X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/a5445a2ae553b0723d8579462bcb48855f71a17c..a19ffad966b25b3869e666f749f7c6da187bef68:/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 5a1ebc54e0..60dc6071b4 100644 --- a/services/crunch-dispatch-slurm/crunch-dispatch-slurm.go +++ b/services/crunch-dispatch-slurm/crunch-dispatch-slurm.go @@ -3,6 +3,7 @@ package main // Dispatcher service for Crunch that submits containers to the slurm queue. import ( + "bytes" "flag" "fmt" "git.curoverse.com/arvados.git/sdk/go/arvados" @@ -10,8 +11,6 @@ import ( "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" "os" @@ -32,6 +31,9 @@ type Config struct { // // Example: []string{"crunch-run", "--cgroup-parent-subsystem=memory"} CrunchRunCommand []string + + // Minimum time between two attempts to run the same container + MinRetryPeriod arvados.Duration } func main() { @@ -42,8 +44,8 @@ func main() { } var ( - theConfig Config - squeueUpdater Squeue + theConfig Config + sqCheck SqueueChecker ) const defaultConfigPath = "/etc/arvados/crunch-dispatch-slurm/crunch-dispatch-slurm.yml" @@ -79,11 +81,11 @@ func doMain() error { // 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", "") + os.Setenv("ARVADOS_API_HOST_INSECURE", "") if theConfig.Client.Insecure { - os.Setenv("ARVADOS_API_INSECURE", "1") + os.Setenv("ARVADOS_API_HOST_INSECURE", "1") } - os.Setenv("ARVADOS_KEEP_SERVICES", "") + 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).") @@ -96,25 +98,21 @@ func doMain() error { } arv.Retries = 25 - squeueUpdater.StartMonitor(time.Duration(theConfig.PollPeriod)) - defer squeueUpdater.Done() + sqCheck = SqueueChecker{Period: time.Duration(theConfig.PollPeriod)} + defer sqCheck.Stop() dispatcher := dispatch.Dispatcher{ Arv: arv, RunContainer: run, - 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) + PollPeriod: time.Duration(theConfig.PollPeriod), + MinRetryPeriod: time.Duration(theConfig.MinRetryPeriod), } - err = dispatcher.RunDispatcher() - if err != nil { - return err + if _, err := daemon.SdNotify(false, "READY=1"); err != nil { + log.Printf("Error notifying init daemon: %v", err) } - return nil + return dispatcher.Run() } // sbatchCmd @@ -127,6 +125,9 @@ func sbatchFunc(container arvados.Container) *exec.Cmd { 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...) } @@ -156,70 +157,31 @@ func submit(dispatcher *dispatch.Dispatcher, } }() - // Create the command and attach to stdin/stdout cmd := sbatchCmd(container) - stdinWriter, stdinerr := cmd.StdinPipe() - if stdinerr != nil { - submitErr = fmt.Errorf("Error creating stdin pipe %v: %q", container.UUID, stdinerr) - return - } - stdoutReader, stdoutErr := cmd.StdoutPipe() - if stdoutErr != nil { - submitErr = fmt.Errorf("Error creating stdout pipe %v: %q", container.UUID, stdoutErr) - return - } + // Send a tiny script on stdin to execute the crunch-run + // command (slurm requires this to be a #! script) + cmd.Stdin = strings.NewReader(execScript(append(crunchRunCommand, container.UUID))) - stderrReader, stderrErr := cmd.StderrPipe() - if stderrErr != nil { - submitErr = fmt.Errorf("Error creating stderr pipe %v: %q", container.UUID, stderrErr) - return - } + var stdout, stderr bytes.Buffer + cmd.Stdout = &stdout + cmd.Stderr = &stderr // Mutex between squeue sync and running sbatch or scancel. - squeueUpdater.SlurmLock.Lock() - defer squeueUpdater.SlurmLock.Unlock() - - err := cmd.Start() - if err != nil { - submitErr = fmt.Errorf("Error starting %v: %v", cmd.Args, err) - return - } - - stdoutChan := make(chan []byte) - go func() { - b, _ := ioutil.ReadAll(stdoutReader) - stdoutReader.Close() - stdoutChan <- b - }() - - stderrChan := make(chan []byte) - go func() { - b, _ := ioutil.ReadAll(stderrReader) - stderrReader.Close() - stderrChan <- b - }() - - // Send a tiny script on stdin to execute the crunch-run command - // slurm actually enforces that this must be a #! script - io.WriteString(stdinWriter, execScript(append(crunchRunCommand, container.UUID))) - stdinWriter.Close() - - err = cmd.Wait() - - stdoutMsg := <-stdoutChan - stderrmsg := <-stderrChan - - close(stdoutChan) - close(stderrChan) - - if err != nil { - submitErr = fmt.Errorf("Container submission failed: %v: %v (stderr: %q)", cmd.Args, err, stderrmsg) - return + sqCheck.L.Lock() + defer sqCheck.L.Unlock() + + log.Printf("exec sbatch %+q", cmd.Args) + err := cmd.Run() + switch err.(type) { + case nil: + log.Printf("sbatch succeeded: %q", strings.TrimSpace(stdout.String())) + return nil + case *exec.ExitError: + return fmt.Errorf("sbatch %+q failed: %v (stderr: %q)", cmd.Args, err, stderr) + default: + return fmt.Errorf("exec failed: %v", err) } - - log.Printf("sbatch succeeded: %s", strings.TrimSpace(string(stdoutMsg))) - return } // If the container is marked as Locked, check if it is already in the slurm @@ -230,7 +192,7 @@ func submit(dispatcher *dispatch.Dispatcher, func monitorSubmitOrCancel(dispatcher *dispatch.Dispatcher, container arvados.Container, monitorDone *bool) { submitted := false for !*monitorDone { - if squeueUpdater.CheckSqueue(container.UUID) { + if sqCheck.HasUUID(container.UUID) { // Found in the queue, so continue monitoring submitted = true } else if container.State == dispatch.Locked && !submitted { @@ -292,27 +254,24 @@ func run(dispatcher *dispatch.Dispatcher, go monitorSubmitOrCancel(dispatcher, container, &monitorDone) for container = range status { - if container.State == dispatch.Locked || container.State == dispatch.Running { - if container.Priority == 0 { - log.Printf("Canceling container %s", container.UUID) - - // Mutex between squeue sync and running sbatch or scancel. - squeueUpdater.SlurmLock.Lock() - err := scancelCmd(container).Run() - squeueUpdater.SlurmLock.Unlock() - - if err != nil { - log.Printf("Error stopping container %s with scancel: %v", - container.UUID, err) - if squeueUpdater.CheckSqueue(container.UUID) { - log.Printf("Container %s is still in squeue after scancel.", - container.UUID) - continue - } - } + if container.Priority == 0 && (container.State == dispatch.Locked || container.State == dispatch.Running) { + log.Printf("Canceling container %s", container.UUID) + // Mutex between squeue sync and running sbatch or scancel. + sqCheck.L.Lock() + cmd := scancelCmd(container) + msg, err := cmd.CombinedOutput() + sqCheck.L.Unlock() - err = dispatcher.UpdateState(container.UUID, dispatch.Cancelled) + if err != nil { + log.Printf("Error stopping container %s with %v %v: %v %v", container.UUID, cmd.Path, cmd.Args, err, string(msg)) + if sqCheck.HasUUID(container.UUID) { + log.Printf("Container %s is still in squeue after scancel.", container.UUID) + continue + } } + + // Ignore errors; if necessary, we'll try again next time + dispatcher.UpdateState(container.UUID, dispatch.Cancelled) } } monitorDone = true