services/crunch-dispatch-slurm/crunch-dispatch-slurm_test.go

   1 package main
   2
   3 import (
   4         "bytes"
   5         "fmt"
   6         "git.curoverse.com/arvados.git/sdk/go/arvados"
   7         "git.curoverse.com/arvados.git/sdk/go/arvadosclient"
   8         "git.curoverse.com/arvados.git/sdk/go/arvadostest"
   9         "git.curoverse.com/arvados.git/sdk/go/dispatch"
  10         "io"
  11         "io/ioutil"
  12         "log"
  13         "net/http"
  14         "net/http/httptest"
  15         "os"
  16         "os/exec"
  17         "strings"
  18         "testing"
  19         "time"
  20
  21         . "gopkg.in/check.v1"
  22 )
  23
  24 // Gocheck boilerplate
  25 func Test(t *testing.T) {
  26         TestingT(t)
  27 }
  28
  29 var _ = Suite(&TestSuite{})
  30 var _ = Suite(&MockArvadosServerSuite{})
  31
  32 type TestSuite struct{}
  33 type MockArvadosServerSuite struct{}
  34
  35 var initialArgs []string
  36
  37 func (s *TestSuite) SetUpSuite(c *C) {
  38         initialArgs = os.Args
  39 }
  40
  41 func (s *TestSuite) TearDownSuite(c *C) {
  42 }
  43
  44 func (s *TestSuite) SetUpTest(c *C) {
  45         args := []string{"crunch-dispatch-slurm"}
  46         os.Args = args
  47
  48         arvadostest.StartAPI()
  49         os.Setenv("ARVADOS_API_TOKEN", arvadostest.Dispatch1Token)
  50 }
  51
  52 func (s *TestSuite) TearDownTest(c *C) {
  53         os.Args = initialArgs
  54         arvadostest.StopAPI()
  55 }
  56
  57 func (s *MockArvadosServerSuite) TearDownTest(c *C) {
  58         arvadostest.ResetEnv()
  59 }
  60
  61 func (s *TestSuite) TestIntegrationNormal(c *C) {
  62         container := s.integrationTest(c, func() *exec.Cmd { return exec.Command("echo", "zzzzz-dz642-queuedcontainer") },
  63                 []string(nil),
  64                 func(dispatcher *dispatch.Dispatcher, container arvados.Container) {
  65                         dispatcher.UpdateState(container.UUID, dispatch.Running)
  66                         time.Sleep(3 * time.Second)
  67                         dispatcher.UpdateState(container.UUID, dispatch.Complete)
  68                 })
  69         c.Check(container.State, Equals, arvados.ContainerStateComplete)
  70 }
  71
  72 func (s *TestSuite) TestIntegrationCancel(c *C) {
  73
  74         // Override sbatchCmd
  75         var scancelCmdLine []string
  76         defer func(orig func(arvados.Container) *exec.Cmd) {
  77                 scancelCmd = orig
  78         }(scancelCmd)
  79         scancelCmd = func(container arvados.Container) *exec.Cmd {
  80                 scancelCmdLine = scancelFunc(container).Args
  81                 return exec.Command("echo")
  82         }
  83
  84         container := s.integrationTest(c,
  85                 func() *exec.Cmd { return exec.Command("echo", "zzzzz-dz642-queuedcontainer") },
  86                 []string(nil),
  87                 func(dispatcher *dispatch.Dispatcher, container arvados.Container) {
  88                         dispatcher.UpdateState(container.UUID, dispatch.Running)
  89                         time.Sleep(1 * time.Second)
  90                         dispatcher.Arv.Update("containers", container.UUID,
  91                                 arvadosclient.Dict{
  92                                         "container": arvadosclient.Dict{"priority": 0}},
  93                                 nil)
  94                 })
  95         c.Check(container.State, Equals, arvados.ContainerStateCancelled)
  96         c.Check(scancelCmdLine, DeepEquals, []string{"scancel", "--name=zzzzz-dz642-queuedcontainer"})
  97 }
  98
  99 func (s *TestSuite) TestIntegrationMissingFromSqueue(c *C) {
 100         container := s.integrationTest(c, func() *exec.Cmd { return exec.Command("echo") }, []string{"sbatch", "--share",
 101                 fmt.Sprintf("--job-name=%s", "zzzzz-dz642-queuedcontainer"),
 102                 fmt.Sprintf("--mem-per-cpu=%d", 2862),
 103                 fmt.Sprintf("--cpus-per-task=%d", 4)},
 104                 func(dispatcher *dispatch.Dispatcher, container arvados.Container) {
 105                         dispatcher.UpdateState(container.UUID, dispatch.Running)
 106                         time.Sleep(3 * time.Second)
 107                         dispatcher.UpdateState(container.UUID, dispatch.Complete)
 108                 })
 109         c.Check(container.State, Equals, arvados.ContainerStateCancelled)
 110 }
 111
 112 func (s *TestSuite) integrationTest(c *C,
 113         newSqueueCmd func() *exec.Cmd,
 114         sbatchCmdComps []string,
 115         runContainer func(*dispatch.Dispatcher, arvados.Container)) arvados.Container {
 116         arvadostest.ResetEnv()
 117
 118         arv, err := arvadosclient.MakeArvadosClient()
 119         c.Assert(err, IsNil)
 120
 121         var sbatchCmdLine []string
 122
 123         // Override sbatchCmd
 124         defer func(orig func(arvados.Container) *exec.Cmd) {
 125                 sbatchCmd = orig
 126         }(sbatchCmd)
 127         sbatchCmd = func(container arvados.Container) *exec.Cmd {
 128                 sbatchCmdLine = sbatchFunc(container).Args
 129                 return exec.Command("sh")
 130         }
 131
 132         // Override squeueCmd
 133         defer func(orig func() *exec.Cmd) {
 134                 squeueCmd = orig
 135         }(squeueCmd)
 136         squeueCmd = newSqueueCmd
 137
 138         // There should be one queued container
 139         params := arvadosclient.Dict{
 140                 "filters": [][]string{{"state", "=", "Queued"}},
 141         }
 142         var containers arvados.ContainerList
 143         err = arv.List("containers", params, &containers)
 144         c.Check(err, IsNil)
 145         c.Check(len(containers.Items), Equals, 1)
 146
 147         theConfig.CrunchRunCommand = []string{"echo"}
 148
 149         dispatcher := dispatch.Dispatcher{
 150                 Arv:        arv,
 151                 PollPeriod: time.Duration(1) * time.Second,
 152                 RunContainer: func(dispatcher *dispatch.Dispatcher,
 153                         container arvados.Container,
 154                         status chan arvados.Container) {
 155                         go runContainer(dispatcher, container)
 156                         run(dispatcher, container, status)
 157                         dispatcher.Stop()
 158                 },
 159         }
 160
 161         squeueUpdater.StartMonitor(time.Duration(500) * time.Millisecond)
 162
 163         err = dispatcher.Run()
 164         c.Assert(err, IsNil)
 165
 166         squeueUpdater.Done()
 167
 168         c.Check(sbatchCmdLine, DeepEquals, sbatchCmdComps)
 169
 170         // There should be no queued containers now
 171         err = arv.List("containers", params, &containers)
 172         c.Check(err, IsNil)
 173         c.Check(len(containers.Items), Equals, 0)
 174
 175         // Previously "Queued" container should now be in "Complete" state
 176         var container arvados.Container
 177         err = arv.Get("containers", "zzzzz-dz642-queuedcontainer", nil, &container)
 178         c.Check(err, IsNil)
 179         return container
 180 }
 181
 182 func (s *MockArvadosServerSuite) TestAPIErrorGettingContainers(c *C) {
 183         apiStubResponses := make(map[string]arvadostest.StubResponse)
 184         apiStubResponses["/arvados/v1/api_client_authorizations/current"] = arvadostest.StubResponse{200, `{"uuid":"` + arvadostest.Dispatch1AuthUUID + `"}`}
 185         apiStubResponses["/arvados/v1/containers"] = arvadostest.StubResponse{500, string(`{}`)}
 186
 187         testWithServerStub(c, apiStubResponses, "echo", "Error getting list of containers")
 188 }
 189
 190 func testWithServerStub(c *C, apiStubResponses map[string]arvadostest.StubResponse, crunchCmd string, expected string) {
 191         apiStub := arvadostest.ServerStub{apiStubResponses}
 192
 193         api := httptest.NewServer(&apiStub)
 194         defer api.Close()
 195
 196         arv := &arvadosclient.ArvadosClient{
 197                 Scheme:    "http",
 198                 ApiServer: api.URL[7:],
 199                 ApiToken:  "abc123",
 200                 Client:    &http.Client{Transport: &http.Transport{}},
 201                 Retries:   0,
 202         }
 203
 204         buf := bytes.NewBuffer(nil)
 205         log.SetOutput(io.MultiWriter(buf, os.Stderr))
 206         defer log.SetOutput(os.Stderr)
 207
 208         theConfig.CrunchRunCommand = []string{crunchCmd}
 209
 210         dispatcher := dispatch.Dispatcher{
 211                 Arv:        arv,
 212                 PollPeriod: time.Duration(1) * time.Second,
 213                 RunContainer: func(dispatcher *dispatch.Dispatcher,
 214                         container arvados.Container,
 215                         status chan arvados.Container) {
 216                         go func() {
 217                                 time.Sleep(1 * time.Second)
 218                                 dispatcher.UpdateState(container.UUID, dispatch.Running)
 219                                 dispatcher.UpdateState(container.UUID, dispatch.Complete)
 220                         }()
 221                         run(dispatcher, container, status)
 222                         dispatcher.Stop()
 223                 },
 224         }
 225
 226         go func() {
 227                 for i := 0; i < 80 && !strings.Contains(buf.String(), expected); i++ {
 228                         time.Sleep(100 * time.Millisecond)
 229                 }
 230                 dispatcher.Stop()
 231         }()
 232
 233         err := dispatcher.Run()
 234         c.Assert(err, IsNil)
 235
 236         c.Check(buf.String(), Matches, `(?ms).*`+expected+`.*`)
 237 }
 238
 239 func (s *MockArvadosServerSuite) TestNoSuchConfigFile(c *C) {
 240         var config Config
 241         err := readConfig(&config, "/nosuchdir89j7879/8hjwr7ojgyy7")
 242         c.Assert(err, NotNil)
 243 }
 244
 245 func (s *MockArvadosServerSuite) TestBadSbatchArgsConfig(c *C) {
 246         var config Config
 247
 248         tmpfile, err := ioutil.TempFile(os.TempDir(), "config")
 249         c.Check(err, IsNil)
 250         defer os.Remove(tmpfile.Name())
 251
 252         _, err = tmpfile.Write([]byte(`{"SbatchArguments": "oops this is not a string array"}`))
 253         c.Check(err, IsNil)
 254
 255         err = readConfig(&config, tmpfile.Name())
 256         c.Assert(err, NotNil)
 257 }
 258
 259 func (s *MockArvadosServerSuite) TestNoSuchArgInConfigIgnored(c *C) {
 260         var config Config
 261
 262         tmpfile, err := ioutil.TempFile(os.TempDir(), "config")
 263         c.Check(err, IsNil)
 264         defer os.Remove(tmpfile.Name())
 265
 266         _, err = tmpfile.Write([]byte(`{"NoSuchArg": "Nobody loves me, not one tiny hunk."}`))
 267         c.Check(err, IsNil)
 268
 269         err = readConfig(&config, tmpfile.Name())
 270         c.Assert(err, IsNil)
 271         c.Check(0, Equals, len(config.SbatchArguments))
 272 }
 273
 274 func (s *MockArvadosServerSuite) TestReadConfig(c *C) {
 275         var config Config
 276
 277         tmpfile, err := ioutil.TempFile(os.TempDir(), "config")
 278         c.Check(err, IsNil)
 279         defer os.Remove(tmpfile.Name())
 280
 281         args := []string{"--arg1=v1", "--arg2", "--arg3=v3"}
 282         argsS := `{"SbatchArguments": ["--arg1=v1",  "--arg2", "--arg3=v3"]}`
 283         _, err = tmpfile.Write([]byte(argsS))
 284         c.Check(err, IsNil)
 285
 286         err = readConfig(&config, tmpfile.Name())
 287         c.Assert(err, IsNil)
 288         c.Check(3, Equals, len(config.SbatchArguments))
 289         c.Check(args, DeepEquals, config.SbatchArguments)
 290 }
 291
 292 func (s *MockArvadosServerSuite) TestSbatchFuncWithNoConfigArgs(c *C) {
 293         testSbatchFuncWithArgs(c, nil)
 294 }
 295
 296 func (s *MockArvadosServerSuite) TestSbatchFuncWithEmptyConfigArgs(c *C) {
 297         testSbatchFuncWithArgs(c, []string{})
 298 }
 299
 300 func (s *MockArvadosServerSuite) TestSbatchFuncWithConfigArgs(c *C) {
 301         testSbatchFuncWithArgs(c, []string{"--arg1=v1", "--arg2"})
 302 }
 303
 304 func testSbatchFuncWithArgs(c *C, args []string) {
 305         theConfig.SbatchArguments = append(theConfig.SbatchArguments, args...)
 306
 307         container := arvados.Container{UUID: "123", RuntimeConstraints: arvados.RuntimeConstraints{RAM: 250000000, VCPUs: 2}}
 308         sbatchCmd := sbatchFunc(container)
 309
 310         var expected []string
 311         expected = append(expected, "sbatch", "--share")
 312         expected = append(expected, theConfig.SbatchArguments...)
 313         expected = append(expected, "--job-name=123", "--mem-per-cpu=120", "--cpus-per-task=2")
 314
 315         c.Check(sbatchCmd.Args, DeepEquals, expected)
 316 }
 317
 318 func (s *MockArvadosServerSuite) TestSbatchPartition(c *C) {
 319         theConfig.SbatchArguments = nil
 320         container := arvados.Container{UUID: "123", RuntimeConstraints: arvados.RuntimeConstraints{RAM: 250000000, VCPUs: 1}, SchedulingParameters: arvados.SchedulingParameters{Partitions: []string{"blurb", "b2"}}}
 321         sbatchCmd := sbatchFunc(container)
 322
 323         var expected []string
 324         expected = append(expected, "sbatch", "--share")
 325         expected = append(expected, "--job-name=123", "--mem-per-cpu=239", "--cpus-per-task=1", "--partition=blurb,b2")
 326
 327         c.Check(sbatchCmd.Args, DeepEquals, expected)
 328 }