X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/c3cbac378fb45e3bf996a5d691cd3f205dfb3f90..50128b53da4003912635b03fb27b5be2c5beaca1:/services/crunch-dispatch-slurm/crunch-dispatch-slurm.go diff --git a/services/crunch-dispatch-slurm/crunch-dispatch-slurm.go b/services/crunch-dispatch-slurm/crunch-dispatch-slurm.go index af2c42e9b2..3c4f281912 100644 --- a/services/crunch-dispatch-slurm/crunch-dispatch-slurm.go +++ b/services/crunch-dispatch-slurm/crunch-dispatch-slurm.go @@ -3,12 +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/config" "git.curoverse.com/arvados.git/sdk/go/dispatch" + "github.com/coreos/go-systemd/daemon" "io" "io/ioutil" "log" @@ -21,8 +22,10 @@ import ( // Config used by crunch-dispatch-slurm type Config struct { + Client arvados.Client + SbatchArguments []string - PollPeriod *time.Duration + PollPeriod arvados.Duration // crunch-run command to invoke. The container UUID will be // appended. If nil, []string{"crunch-run"} will be used. @@ -34,41 +37,56 @@ type Config struct { 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") + "`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 config.CrunchRunCommand == nil { - config.CrunchRunCommand = []string{"crunch-run"} + 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_HOST_INSECURE", "") + if theConfig.Client.Insecure { + os.Setenv("ARVADOS_API_HOST_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() @@ -78,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(false, "READY=1"); err != nil { + log.Printf("Error notifying init daemon: %v", err) + } + err = dispatcher.RunDispatcher() if err != nil { return err @@ -101,10 +123,13 @@ 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)) + if container.SchedulingParameters.Partitions != nil { + sbatchArgs = append(sbatchArgs, fmt.Sprintf("--partition=%s", strings.Join(container.SchedulingParameters.Partitions, ","))) + } return exec.Command("sbatch", sbatchArgs...) } @@ -128,10 +153,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) } @@ -161,9 +183,10 @@ func submit(dispatcher *dispatch.Dispatcher, squeueUpdater.SlurmLock.Lock() defer squeueUpdater.SlurmLock.Unlock() + log.Printf("sbatch starting: %+q", cmd.Args) err := cmd.Start() if err != nil { - submitErr = fmt.Errorf("Error starting %v: %v", cmd.Args, err) + submitErr = fmt.Errorf("Error starting sbatch: %v", err) return } @@ -220,11 +243,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 { @@ -240,20 +263,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) } } } @@ -300,16 +323,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 }