X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/ac4ee0eb5a2d581b7bb4d26b475f0fbcf64b0630..ffe1fe1c77743147ee82aacdc50edde3672cd748:/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 1d9015d505..5a1ebc54e0 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,7 +22,7 @@ import ( // Config used by crunch-dispatch-slurm type Config struct { - arvados.Client + Client arvados.Client SbatchArguments []string PollPeriod arvados.Duration @@ -36,16 +37,16 @@ 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) @@ -54,33 +55,32 @@ func doMain() error { configPath := flags.String( "config", defaultConfigPath, - "`path` to json configuration file") + "`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 config.PollPeriod == 0 { - config.PollPeriod = arvados.Duration(10 * time.Second) + if theConfig.PollPeriod == 0 { + theConfig.PollPeriod = arvados.Duration(10 * time.Second) } - if config.Client.APIHost != "" || config.Client.AuthToken != "" { + 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", config.Client.APIHost) - os.Setenv("ARVADOS_API_TOKEN", config.Client.AuthToken) + os.Setenv("ARVADOS_API_HOST", theConfig.Client.APIHost) + os.Setenv("ARVADOS_API_TOKEN", theConfig.Client.AuthToken) os.Setenv("ARVADOS_API_INSECURE", "") - if config.Client.Insecure { + if theConfig.Client.Insecure { os.Setenv("ARVADOS_API_INSECURE", "1") } os.Setenv("ARVADOS_KEEP_SERVICES", "") @@ -96,15 +96,19 @@ func doMain() error { } arv.Retries = 25 - squeueUpdater.StartMonitor(time.Duration(config.PollPeriod)) + squeueUpdater.StartMonitor(time.Duration(theConfig.PollPeriod)) defer squeueUpdater.Done() dispatcher := dispatch.Dispatcher{ Arv: arv, RunContainer: run, - PollInterval: time.Duration(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 @@ -119,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)) @@ -146,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) } @@ -238,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 { @@ -258,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) } } } @@ -318,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 }