9406: Support passing additional arguments from crunch-dispatch-slurm to crunch-run.

diff --git a/services/crunch-dispatch-slurm/crunch-dispatch-slurm.go b/services/crunch-dispatch-slurm/crunch-dispatch-slurm.go
index 740df55ecd103426150dde9295849cb9b12b9b57..af2c42e9b2faf87d42c47e3e306d58598226640e 100644 (file)
--- a/services/crunch-dispatch-slurm/crunch-dispatch-slurm.go
+++ b/services/crunch-dispatch-slurm/crunch-dispatch-slurm.go
@@ -9,6 +9,7 @@ import (
        "git.curoverse.com/arvados.git/sdk/go/arvados"
        "git.curoverse.com/arvados.git/sdk/go/arvadosclient"
        "git.curoverse.com/arvados.git/sdk/go/dispatch"
+       "io"
        "io/ioutil"
        "log"
        "math"
@@ -20,9 +21,14 @@ import (
 
 // Config used by crunch-dispatch-slurm
 type Config struct {
-       SbatchArguments  []string
-       PollPeriod       *time.Duration
-       CrunchRunCommand *string
+       SbatchArguments []string
+       PollPeriod      *time.Duration
+
+       // crunch-run command to invoke. The container UUID will be
+       // appended. If nil, []string{"crunch-run"} will be used.
+       //
+       // Example: []string{"crunch-run", "--cgroup-parent-subsystem=memory"}
+       CrunchRunCommand []string
 }
 
 func main() {
@@ -52,11 +58,6 @@ func doMain() error {
                10*time.Second,
                "Time duration to poll for queued containers")
 
-       config.CrunchRunCommand = flags.String(
-               "crunch-run-command",
-               "/usr/bin/crunch-run",
-               "Crunch command to run container")
-
        // Parse args; omit the first arg which is the command name
        flags.Parse(os.Args[1:])
 
@@ -66,6 +67,10 @@ func doMain() error {
                return err
        }
 
+       if config.CrunchRunCommand == nil {
+               config.CrunchRunCommand = []string{"crunch-run"}
+       }
+
        arv, err := arvadosclient.MakeArvadosClient()
        if err != nil {
                log.Printf("Error making Arvados client: %v", err)
@@ -115,7 +120,7 @@ var scancelCmd = scancelFunc
 
 // Submit job to slurm using sbatch.
 func submit(dispatcher *dispatch.Dispatcher,
-       container arvados.Container, crunchRunCommand string) (submitErr error) {
+       container arvados.Container, crunchRunCommand []string) (submitErr error) {
        defer func() {
                // If we didn't get as far as submitting a slurm job,
                // unlock the container and return it to the queue.
@@ -178,7 +183,7 @@ func submit(dispatcher *dispatch.Dispatcher,
 
        // Send a tiny script on stdin to execute the crunch-run command
        // slurm actually enforces that this must be a #! script
-       fmt.Fprintf(stdinWriter, "#!/bin/sh\nexec '%s' '%s'\n", crunchRunCommand, container.UUID)
+       io.WriteString(stdinWriter, execScript(append(crunchRunCommand, container.UUID)))
        stdinWriter.Close()
 
        err = cmd.Wait()
@@ -215,7 +220,7 @@ func monitorSubmitOrCancel(dispatcher *dispatch.Dispatcher, container arvados.Co
 
                        log.Printf("About to submit queued container %v", container.UUID)
 
-                       if err := submit(dispatcher, container, *config.CrunchRunCommand); err != nil {
+                       if err := submit(dispatcher, container, config.CrunchRunCommand); err != nil {
                                log.Printf("Error submitting container %s to slurm: %v",
                                        container.UUID, err)
                                // maybe sbatch is broken, put it back to queued

diff --git a/services/crunch-dispatch-slurm/crunch-dispatch-slurm_test.go b/services/crunch-dispatch-slurm/crunch-dispatch-slurm_test.go
index a55929838bc39cd86c27664ccee27c544b6b5cba..6692f7f80e341fecfeec5cab1a3dba6ddb3d27a6 100644 (file)
--- a/services/crunch-dispatch-slurm/crunch-dispatch-slurm_test.go
+++ b/services/crunch-dispatch-slurm/crunch-dispatch-slurm_test.go
@@ -143,8 +143,7 @@ func (s *TestSuite) integrationTest(c *C,
        c.Check(err, IsNil)
        c.Check(len(containers.Items), Equals, 1)
 
-       echo := "echo"
-       config.CrunchRunCommand = &echo
+       config.CrunchRunCommand = []string{"echo"}
 
        doneProcessing := make(chan struct{})
        dispatcher := dispatch.Dispatcher{
@@ -206,7 +205,7 @@ func testWithServerStub(c *C, apiStubResponses map[string]arvadostest.StubRespon
        log.SetOutput(io.MultiWriter(buf, os.Stderr))
        defer log.SetOutput(os.Stderr)
 
-       config.CrunchRunCommand = &crunchCmd
+       config.CrunchRunCommand = []string{crunchCmd}
 
        doneProcessing := make(chan struct{})
        dispatcher := dispatch.Dispatcher{

diff --git a/services/crunch-dispatch-slurm/script.go b/services/crunch-dispatch-slurm/script.go
new file mode 100644 (file)
index 0000000..93ae6b5
--- /dev/null
+++ b/services/crunch-dispatch-slurm/script.go
@@ -0,0 +1,15 @@
+package main
+
+import (
+       "strings"
+)
+
+func execScript(args []string) string {
+       s := "#!/bin/sh\nexec"
+       for _, w := range args {
+               s += ` '`
+               s += strings.Replace(w, `'`, `'\''`, -1)
+               s += `'`
+       }
+       return s + "\n"
+}

diff --git a/services/crunch-dispatch-slurm/script_test.go b/services/crunch-dispatch-slurm/script_test.go
new file mode 100644 (file)
index 0000000..3cb407d
--- /dev/null
+++ b/services/crunch-dispatch-slurm/script_test.go
@@ -0,0 +1,24 @@
+package main
+
+import (
+       . "gopkg.in/check.v1"
+)
+
+var _ = Suite(&ScriptSuite{})
+
+type ScriptSuite struct{}
+
+func (s *ScriptSuite) TestExecScript(c *C) {
+       for _, test := range []struct {
+               args   []string
+               script string
+       }{
+               {nil, `exec`},
+               {[]string{`foo`}, `exec 'foo'`},
+               {[]string{`foo`, `bar baz`}, `exec 'foo' 'bar baz'`},
+               {[]string{`foo"`, "'waz 'qux\n"}, `exec 'foo"' ''\''waz '\''qux` + "\n" + `'`},
+       } {
+               c.Logf("%+v -> %+v", test.args, test.script)
+               c.Check(execScript(test.args), Equals, "#!/bin/sh\n"+test.script+"\n")
+       }
+}