From: Tom Clegg <tom@curoverse.com>
Date: Tue, 2 Aug 2016 01:42:37 +0000 (-0400)
Subject: 9406: Support passing additional arguments from crunch-dispatch-slurm to crunch-run.
X-Git-Tag: 1.1.0~806^2~2
X-Git-Url: https://git.arvados.org/arvados.git/commitdiff_plain/483ca35ac4348924cfbc187dab4f1b88a272eea7

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 740df55ecd..af2c42e9b2 100644
--- 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 a55929838b..6692f7f80e 100644
--- 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
index 0000000000..93ae6b5eb3
--- /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
index 0000000000..3cb407dbb0
--- /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")
+	}
+}