package main
import (
+ "bytes"
+ "fmt"
+ "git.curoverse.com/arvados.git/sdk/go/arvados"
"git.curoverse.com/arvados.git/sdk/go/arvadosclient"
"git.curoverse.com/arvados.git/sdk/go/arvadostest"
"git.curoverse.com/arvados.git/sdk/go/dispatch"
-
- "bytes"
- "fmt"
+ "io"
"log"
- "math"
"net/http"
"net/http/httptest"
"os"
func (s *TestSuite) SetUpSuite(c *C) {
initialArgs = os.Args
- arvadostest.StartAPI()
}
func (s *TestSuite) TearDownSuite(c *C) {
- arvadostest.StopAPI()
}
func (s *TestSuite) SetUpTest(c *C) {
args := []string{"crunch-dispatch-slurm"}
os.Args = args
+ arvadostest.StartAPI()
os.Setenv("ARVADOS_API_TOKEN", arvadostest.Dispatch1Token)
}
func (s *TestSuite) TearDownTest(c *C) {
- arvadostest.ResetEnv()
os.Args = initialArgs
+ arvadostest.StopAPI()
}
func (s *MockArvadosServerSuite) TearDownTest(c *C) {
arvadostest.ResetEnv()
}
-func (s *TestSuite) TestIntegration(c *C) {
+func (s *TestSuite) TestIntegrationNormal(c *C) {
+ container := s.integrationTest(c, func() *exec.Cmd { return exec.Command("echo", "zzzzz-dz642-queuedcontainer") },
+ []string(nil),
+ func(dispatcher *dispatch.Dispatcher, container arvados.Container) {
+ dispatcher.UpdateState(container.UUID, dispatch.Running)
+ time.Sleep(3 * time.Second)
+ dispatcher.UpdateState(container.UUID, dispatch.Complete)
+ })
+ c.Check(container.State, Equals, arvados.ContainerStateComplete)
+}
+
+func (s *TestSuite) TestIntegrationCancel(c *C) {
+
+ // Override sbatchCmd
+ var scancelCmdLine []string
+ defer func(orig func(arvados.Container) *exec.Cmd) {
+ scancelCmd = orig
+ }(scancelCmd)
+ scancelCmd = func(container arvados.Container) *exec.Cmd {
+ scancelCmdLine = scancelFunc(container).Args
+ return exec.Command("echo")
+ }
+
+ container := s.integrationTest(c, func() *exec.Cmd { return exec.Command("echo", "zzzzz-dz642-queuedcontainer") },
+ []string(nil),
+ func(dispatcher *dispatch.Dispatcher, container arvados.Container) {
+ dispatcher.UpdateState(container.UUID, dispatch.Running)
+ time.Sleep(1 * time.Second)
+ dispatcher.Arv.Update("containers", container.UUID,
+ arvadosclient.Dict{
+ "container": arvadosclient.Dict{"priority": 0}},
+ nil)
+ })
+ c.Check(container.State, Equals, arvados.ContainerStateCancelled)
+ c.Check(scancelCmdLine, DeepEquals, []string{"scancel", "--name=zzzzz-dz642-queuedcontainer"})
+}
+
+func (s *TestSuite) TestIntegrationMissingFromSqueue(c *C) {
+ container := s.integrationTest(c, func() *exec.Cmd { return exec.Command("echo") }, []string{"sbatch", "--share", "--parsable",
+ fmt.Sprintf("--job-name=%s", "zzzzz-dz642-queuedcontainer"),
+ fmt.Sprintf("--mem-per-cpu=%d", 2862),
+ fmt.Sprintf("--cpus-per-task=%d", 4),
+ fmt.Sprintf("--priority=%d", 1)},
+ func(dispatcher *dispatch.Dispatcher, container arvados.Container) {
+ dispatcher.UpdateState(container.UUID, dispatch.Running)
+ time.Sleep(3 * time.Second)
+ dispatcher.UpdateState(container.UUID, dispatch.Complete)
+ })
+ c.Check(container.State, Equals, arvados.ContainerStateCancelled)
+}
+
+func (s *TestSuite) integrationTest(c *C,
+ newSqueueCmd func() *exec.Cmd,
+ sbatchCmdComps []string,
+ runContainer func(*dispatch.Dispatcher, arvados.Container)) arvados.Container {
+ arvadostest.ResetEnv()
+
arv, err := arvadosclient.MakeArvadosClient()
c.Assert(err, IsNil)
var sbatchCmdLine []string
- var striggerCmdLine []string
// Override sbatchCmd
- defer func(orig func(dispatch.Container) *exec.Cmd) {
+ defer func(orig func(arvados.Container) *exec.Cmd) {
sbatchCmd = orig
}(sbatchCmd)
- sbatchCmd = func(container dispatch.Container) *exec.Cmd {
+ sbatchCmd = func(container arvados.Container) *exec.Cmd {
sbatchCmdLine = sbatchFunc(container).Args
return exec.Command("sh")
}
- // Override striggerCmd
- defer func(orig func(jobid, containerUUID, finishCommand,
- apiHost, apiToken, apiInsecure string) *exec.Cmd) {
- striggerCmd = orig
- }(striggerCmd)
- striggerCmd = func(jobid, containerUUID, finishCommand, apiHost, apiToken, apiInsecure string) *exec.Cmd {
- striggerCmdLine = striggerFunc(jobid, containerUUID, finishCommand,
- apiHost, apiToken, apiInsecure).Args
- go func() {
- time.Sleep(5 * time.Second)
- arv.Update("containers", containerUUID,
- arvadosclient.Dict{
- "container": arvadosclient.Dict{"state": dispatch.Complete}},
- nil)
- }()
- return exec.Command("echo", striggerCmdLine...)
- }
-
// Override squeueCmd
defer func(orig func() *exec.Cmd) {
squeueCmd = orig
}(squeueCmd)
- squeueCmd = func() *exec.Cmd {
- return exec.Command("echo")
- }
+ squeueCmd = newSqueueCmd
// There should be no queued containers now
params := arvadosclient.Dict{
- "filters": [][]string{[]string{"state", "=", "Queued"}},
+ "filters": [][]string{{"state", "=", "Queued"}},
}
- var containers dispatch.ContainerList
+ var containers arvados.ContainerList
err = arv.List("containers", params, &containers)
c.Check(err, IsNil)
c.Check(len(containers.Items), Equals, 1)
echo := "echo"
crunchRunCommand = &echo
- finishCmd := "/usr/bin/crunch-finish-slurm.sh"
- finishCommand = &finishCmd
doneProcessing := make(chan struct{})
dispatcher := dispatch.Dispatcher{
Arv: arv,
PollInterval: time.Duration(1) * time.Second,
RunContainer: func(dispatcher *dispatch.Dispatcher,
- container dispatch.Container,
- status chan dispatch.Container) {
- go func() {
- time.Sleep(1)
- dispatcher.UpdateState(container.UUID, dispatch.Running)
- dispatcher.UpdateState(container.UUID, dispatch.Complete)
- }()
+ container arvados.Container,
+ status chan arvados.Container) {
+ go runContainer(dispatcher, container)
run(dispatcher, container, status)
doneProcessing <- struct{}{}
},
DoneProcessing: doneProcessing}
+ squeueUpdater.StartMonitor(time.Duration(500) * time.Millisecond)
+
err = dispatcher.RunDispatcher()
c.Assert(err, IsNil)
- item := containers.Items[0]
- sbatchCmdComps := []string{"sbatch", "--share", "--parsable",
- fmt.Sprintf("--job-name=%s", item.UUID),
- fmt.Sprintf("--mem-per-cpu=%d", int(math.Ceil(float64(item.RuntimeConstraints["ram"])/float64(item.RuntimeConstraints["vcpus"]*1048576)))),
- fmt.Sprintf("--cpus-per-task=%d", int(item.RuntimeConstraints["vcpus"])),
- fmt.Sprintf("--priority=%d", item.Priority)}
- c.Check(sbatchCmdLine, DeepEquals, sbatchCmdComps)
+ squeueUpdater.Done()
- c.Check(striggerCmdLine, DeepEquals, []string{"strigger", "--set", "--jobid=zzzzz-dz642-queuedcontainer", "--fini",
- "--program=/usr/bin/crunch-finish-slurm.sh " + os.Getenv("ARVADOS_API_HOST") + " " + arvadostest.Dispatch1Token + " 1 zzzzz-dz642-queuedcontainer"})
+ c.Check(sbatchCmdLine, DeepEquals, sbatchCmdComps)
// There should be no queued containers now
err = arv.List("containers", params, &containers)
c.Check(len(containers.Items), Equals, 0)
// Previously "Queued" container should now be in "Complete" state
- var container dispatch.Container
+ var container arvados.Container
err = arv.Get("containers", "zzzzz-dz642-queuedcontainer", nil, &container)
c.Check(err, IsNil)
- c.Check(container.State, Equals, "Complete")
+ return container
}
func (s *MockArvadosServerSuite) Test_APIErrorGettingContainers(c *C) {
}
buf := bytes.NewBuffer(nil)
- log.SetOutput(buf)
+ log.SetOutput(io.MultiWriter(buf, os.Stderr))
defer log.SetOutput(os.Stderr)
crunchRunCommand = &crunchCmd
- finishCmd := "/usr/bin/crunch-finish-slurm.sh"
- finishCommand = &finishCmd
doneProcessing := make(chan struct{})
dispatcher := dispatch.Dispatcher{
Arv: arv,
PollInterval: time.Duration(1) * time.Second,
RunContainer: func(dispatcher *dispatch.Dispatcher,
- container dispatch.Container,
- status chan dispatch.Container) {
+ container arvados.Container,
+ status chan arvados.Container) {
go func() {
- time.Sleep(1)
+ time.Sleep(1 * time.Second)
dispatcher.UpdateState(container.UUID, dispatch.Running)
dispatcher.UpdateState(container.UUID, dispatch.Complete)
}()