9684: update arvados_model -> recursive_stringify to convert ":foo" to "foo"
[arvados.git] / services / crunch-dispatch-slurm / crunch-dispatch-slurm.go
index f718fbcdcea3fd5c00ab8763240ee3056f098a53..b11963c7040ac7e1ad0d5ce42e983c54f1bb5624 100644 (file)
@@ -3,10 +3,13 @@ package main
 // Dispatcher service for Crunch that submits containers to the slurm queue.
 
 import (
+       "encoding/json"
        "flag"
        "fmt"
+       "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"
@@ -16,6 +19,18 @@ import (
        "time"
 )
 
+// Config used by crunch-dispatch-slurm
+type Config struct {
+       SbatchArguments []string
+       PollPeriod      arvados.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() {
        err := doMain()
        if err != nil {
@@ -24,26 +39,38 @@ func main() {
 }
 
 var (
-       crunchRunCommand *string
-       squeueUpdater    Squeue
+       config        Config
+       squeueUpdater Squeue
 )
 
+const defaultConfigPath = "/etc/arvados/crunch-dispatch-slurm/config.json"
+
 func doMain() error {
        flags := flag.NewFlagSet("crunch-dispatch-slurm", flag.ExitOnError)
+       flags.Usage = func() { usage(flags) }
 
-       pollInterval := flags.Int(
-               "poll-interval",
-               10,
-               "Interval in seconds to poll for queued containers")
-
-       crunchRunCommand = flags.String(
-               "crunch-run-command",
-               "/usr/bin/crunch-run",
-               "Crunch command to run container")
+       configPath := flags.String(
+               "config",
+               defaultConfigPath,
+               "`path` to json configuration file")
 
        // Parse args; omit the first arg which is the command name
        flags.Parse(os.Args[1:])
 
+       err := readConfig(&config, *configPath)
+       if err != nil {
+               log.Printf("Error reading configuration: %v", err)
+               return err
+       }
+
+       if config.CrunchRunCommand == nil {
+               config.CrunchRunCommand = []string{"crunch-run"}
+       }
+
+       if config.PollPeriod == 0 {
+               config.PollPeriod = arvados.Duration(10 * time.Second)
+       }
+
        arv, err := arvadosclient.MakeArvadosClient()
        if err != nil {
                log.Printf("Error making Arvados client: %v", err)
@@ -51,13 +78,13 @@ func doMain() error {
        }
        arv.Retries = 25
 
-       squeueUpdater.StartMonitor(time.Duration(*pollInterval) * time.Second)
+       squeueUpdater.StartMonitor(time.Duration(config.PollPeriod))
        defer squeueUpdater.Done()
 
        dispatcher := dispatch.Dispatcher{
                Arv:            arv,
                RunContainer:   run,
-               PollInterval:   time.Duration(*pollInterval) * time.Second,
+               PollInterval:   time.Duration(config.PollPeriod),
                DoneProcessing: make(chan struct{})}
 
        err = dispatcher.RunDispatcher()
@@ -69,17 +96,21 @@ func doMain() error {
 }
 
 // sbatchCmd
-func sbatchFunc(container dispatch.Container) *exec.Cmd {
-       memPerCPU := math.Ceil((float64(container.RuntimeConstraints["ram"])) / (float64(container.RuntimeConstraints["vcpus"] * 1048576)))
-       return exec.Command("sbatch", "--share", "--parsable",
-               fmt.Sprintf("--job-name=%s", container.UUID),
-               fmt.Sprintf("--mem-per-cpu=%d", int(memPerCPU)),
-               fmt.Sprintf("--cpus-per-task=%d", int(container.RuntimeConstraints["vcpus"])),
-               fmt.Sprintf("--priority=%d", container.Priority))
+func sbatchFunc(container arvados.Container) *exec.Cmd {
+       memPerCPU := math.Ceil(float64(container.RuntimeConstraints.RAM) / (float64(container.RuntimeConstraints.VCPUs) * 1048576))
+
+       var sbatchArgs []string
+       sbatchArgs = append(sbatchArgs, "--share")
+       sbatchArgs = append(sbatchArgs, config.SbatchArguments...)
+       sbatchArgs = append(sbatchArgs, fmt.Sprintf("--job-name=%s", container.UUID))
+       sbatchArgs = append(sbatchArgs, fmt.Sprintf("--mem-per-cpu=%d", int(memPerCPU)))
+       sbatchArgs = append(sbatchArgs, fmt.Sprintf("--cpus-per-task=%d", container.RuntimeConstraints.VCPUs))
+
+       return exec.Command("sbatch", sbatchArgs...)
 }
 
 // scancelCmd
-func scancelFunc(container dispatch.Container) *exec.Cmd {
+func scancelFunc(container arvados.Container) *exec.Cmd {
        return exec.Command("scancel", "--name="+container.UUID)
 }
 
@@ -89,9 +120,7 @@ var scancelCmd = scancelFunc
 
 // Submit job to slurm using sbatch.
 func submit(dispatcher *dispatch.Dispatcher,
-       container dispatch.Container, crunchRunCommand string) (jobid string, submitErr error) {
-       submitErr = nil
-
+       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.
@@ -154,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()
@@ -166,13 +195,11 @@ func submit(dispatcher *dispatch.Dispatcher,
        close(stderrChan)
 
        if err != nil {
-               submitErr = fmt.Errorf("Container submission failed %v: %v %v", cmd.Args, err, stderrmsg)
+               submitErr = fmt.Errorf("Container submission failed: %v: %v (stderr: %q)", cmd.Args, err, stderrmsg)
                return
        }
 
-       // If everything worked out, got the jobid on stdout
-       jobid = strings.TrimSpace(string(stdoutMsg))
-
+       log.Printf("sbatch succeeded: %s", strings.TrimSpace(string(stdoutMsg)))
        return
 }
 
@@ -181,7 +208,7 @@ func submit(dispatcher *dispatch.Dispatcher,
 //
 // If the container is marked as Running, check if it is in the slurm queue.
 // If not, mark it as Cancelled.
-func monitorSubmitOrCancel(dispatcher *dispatch.Dispatcher, container dispatch.Container, monitorDone *bool) {
+func monitorSubmitOrCancel(dispatcher *dispatch.Dispatcher, container arvados.Container, monitorDone *bool) {
        submitted := false
        for !*monitorDone {
                if squeueUpdater.CheckSqueue(container.UUID) {
@@ -193,7 +220,7 @@ func monitorSubmitOrCancel(dispatcher *dispatch.Dispatcher, container dispatch.C
 
                        log.Printf("About to submit queued container %v", container.UUID)
 
-                       if _, err := submit(dispatcher, container, *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
@@ -207,13 +234,13 @@ func monitorSubmitOrCancel(dispatcher *dispatch.Dispatcher, container dispatch.C
                        // release it back to the Queue, if it is Running then
                        // clean up the record.
 
-                       var con dispatch.Container
+                       var con arvados.Container
                        err := dispatcher.Arv.Get("containers", container.UUID, nil, &con)
                        if err != nil {
                                log.Printf("Error getting final container state: %v", err)
                        }
 
-                       var st string
+                       var st arvados.ContainerState
                        switch con.State {
                        case dispatch.Locked:
                                st = dispatch.Queued
@@ -236,8 +263,8 @@ func monitorSubmitOrCancel(dispatcher *dispatch.Dispatcher, container dispatch.C
 // Monitor status updates.  If the priority changes to zero, cancel the
 // container using scancel.
 func run(dispatcher *dispatch.Dispatcher,
-       container dispatch.Container,
-       status chan dispatch.Container) {
+       container arvados.Container,
+       status chan arvados.Container) {
 
        log.Printf("Monitoring container %v started", container.UUID)
        defer log.Printf("Monitoring container %v finished", container.UUID)
@@ -271,3 +298,18 @@ func run(dispatcher *dispatch.Dispatcher,
        }
        monitorDone = true
 }
+
+func readConfig(dst interface{}, path string) error {
+       if buf, err := ioutil.ReadFile(path); err != nil && os.IsNotExist(err) {
+               if path == defaultConfigPath {
+                       log.Printf("Config not specified. Continue with default configuration.")
+               } else {
+                       return fmt.Errorf("Config file not found %q: %v", path, err)
+               }
+       } else if err != nil {
+               return fmt.Errorf("Error reading config %q: %v", path, err)
+       } else if err = json.Unmarshal(buf, dst); err != nil {
+               return fmt.Errorf("Error decoding config %q: %v", path, err)
+       }
+       return nil
+}