Merge branch '10041-test-arvados-keep-services'
[arvados.git] / services / crunch-dispatch-slurm / crunch-dispatch-slurm.go
index 740df55ecd103426150dde9295849cb9b12b9b57..aaea51cbf7950d30e2a8fe91aed8843566c98114 100644 (file)
@@ -3,12 +3,14 @@ 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/config"
        "git.curoverse.com/arvados.git/sdk/go/dispatch"
+       "github.com/coreos/go-systemd/daemon"
+       "io"
        "io/ioutil"
        "log"
        "math"
@@ -20,52 +22,73 @@ import (
 
 // Config used by crunch-dispatch-slurm
 type Config struct {
-       SbatchArguments  []string
-       PollPeriod       *time.Duration
-       CrunchRunCommand *string
+       Client arvados.Client
+
+       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 {
-               log.Fatalf("%q", err)
+               log.Fatal(err)
        }
 }
 
 var (
-       config        Config
+       theConfig     Config
        squeueUpdater Squeue
 )
 
-const defaultConfigPath = "/etc/arvados/crunch-dispatch-slurm/config.json"
+const defaultConfigPath = "/etc/arvados/crunch-dispatch-slurm/crunch-dispatch-slurm.yml"
 
 func doMain() error {
        flags := flag.NewFlagSet("crunch-dispatch-slurm", flag.ExitOnError)
+       flags.Usage = func() { usage(flags) }
 
        configPath := flags.String(
                "config",
                defaultConfigPath,
-               "`path` to json configuration file")
-
-       config.PollPeriod = flags.Duration(
-               "poll-interval",
-               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")
+               "`path` to JSON or YAML configuration file")
 
        // Parse args; omit the first arg which is the command name
        flags.Parse(os.Args[1:])
 
-       err := readConfig(&config, *configPath)
+       err := readConfig(&theConfig, *configPath)
        if err != nil {
-               log.Printf("Error reading configuration: %v", err)
                return err
        }
 
+       if theConfig.CrunchRunCommand == nil {
+               theConfig.CrunchRunCommand = []string{"crunch-run"}
+       }
+
+       if theConfig.PollPeriod == 0 {
+               theConfig.PollPeriod = arvados.Duration(10 * time.Second)
+       }
+
+       if theConfig.Client.APIHost != "" || theConfig.Client.AuthToken != "" {
+               // Copy real configs into env vars so [a]
+               // MakeArvadosClient() uses them, and [b] they get
+               // propagated to crunch-run via SLURM.
+               os.Setenv("ARVADOS_API_HOST", theConfig.Client.APIHost)
+               os.Setenv("ARVADOS_API_TOKEN", theConfig.Client.AuthToken)
+               os.Setenv("ARVADOS_API_INSECURE", "")
+               if theConfig.Client.Insecure {
+                       os.Setenv("ARVADOS_API_INSECURE", "1")
+               }
+               os.Setenv("ARVADOS_KEEP_SERVICES", strings.Join(theConfig.Client.KeepServiceURIs, " "))
+               os.Setenv("ARVADOS_EXTERNAL_CLIENT", "")
+       } else {
+               log.Printf("warning: Client credentials missing from config, so falling back on environment variables (deprecated).")
+       }
+
        arv, err := arvadosclient.MakeArvadosClient()
        if err != nil {
                log.Printf("Error making Arvados client: %v", err)
@@ -73,15 +96,19 @@ func doMain() error {
        }
        arv.Retries = 25
 
-       squeueUpdater.StartMonitor(*config.PollPeriod)
+       squeueUpdater.StartMonitor(time.Duration(theConfig.PollPeriod))
        defer squeueUpdater.Done()
 
        dispatcher := dispatch.Dispatcher{
                Arv:            arv,
                RunContainer:   run,
-               PollInterval:   *config.PollPeriod,
+               PollInterval:   time.Duration(theConfig.PollPeriod),
                DoneProcessing: make(chan struct{})}
 
+       if _, err := daemon.SdNotify("READY=1"); err != nil {
+               log.Printf("Error notifying init daemon: %v", err)
+       }
+
        err = dispatcher.RunDispatcher()
        if err != nil {
                return err
@@ -96,7 +123,7 @@ func sbatchFunc(container arvados.Container) *exec.Cmd {
 
        var sbatchArgs []string
        sbatchArgs = append(sbatchArgs, "--share")
-       sbatchArgs = append(sbatchArgs, config.SbatchArguments...)
+       sbatchArgs = append(sbatchArgs, theConfig.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))
@@ -115,7 +142,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.
@@ -123,10 +150,7 @@ func submit(dispatcher *dispatch.Dispatcher,
                        // OK, no cleanup needed
                        return
                }
-               err := dispatcher.Arv.Update("containers", container.UUID,
-                       arvadosclient.Dict{
-                               "container": arvadosclient.Dict{"state": "Queued"}},
-                       nil)
+               err := dispatcher.Unlock(container.UUID)
                if err != nil {
                        log.Printf("Error unlocking container %s: %v", container.UUID, err)
                }
@@ -178,7 +202,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,11 +239,11 @@ 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, theConfig.CrunchRunCommand); err != nil {
                                log.Printf("Error submitting container %s to slurm: %v",
                                        container.UUID, err)
                                // maybe sbatch is broken, put it back to queued
-                               dispatcher.UpdateState(container.UUID, dispatch.Queued)
+                               dispatcher.Unlock(container.UUID)
                        }
                        submitted = true
                } else {
@@ -235,20 +259,20 @@ func monitorSubmitOrCancel(dispatcher *dispatch.Dispatcher, container arvados.Co
                                log.Printf("Error getting final container state: %v", err)
                        }
 
-                       var st arvados.ContainerState
                        switch con.State {
                        case dispatch.Locked:
-                               st = dispatch.Queued
+                               log.Printf("Container %s in state %v but missing from slurm queue, changing to %v.",
+                                       container.UUID, con.State, dispatch.Queued)
+                               dispatcher.Unlock(container.UUID)
                        case dispatch.Running:
-                               st = dispatch.Cancelled
+                               st := dispatch.Cancelled
+                               log.Printf("Container %s in state %v but missing from slurm queue, changing to %v.",
+                                       container.UUID, con.State, st)
+                               dispatcher.UpdateState(container.UUID, st)
                        default:
                                // Container state is Queued, Complete or Cancelled so stop monitoring it.
                                return
                        }
-
-                       log.Printf("Container %s in state %v but missing from slurm queue, changing to %v.",
-                               container.UUID, con.State, st)
-                       dispatcher.UpdateState(container.UUID, st)
                }
        }
 }
@@ -295,16 +319,10 @@ func run(dispatcher *dispatch.Dispatcher,
 }
 
 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)
+       err := config.LoadFile(dst, path)
+       if err != nil && os.IsNotExist(err) && path == defaultConfigPath {
+               log.Printf("Config not specified. Continue with default configuration.")
+               err = nil
        }
-       return nil
+       return err
 }