services/crunch-dispatch-slurm/crunch-dispatch-slurm.go

   1 package main
   2
   3 // Dispatcher service for Crunch that submits containers to the slurm queue.
   4
   5 import (
   6         "flag"
   7         "fmt"
   8         "io"
   9         "io/ioutil"
  10         "log"
  11         "math"
  12         "os"
  13         "os/exec"
  14         "strings"
  15         "time"
  16
  17         "git.curoverse.com/arvados.git/sdk/go/arvados"
  18         "git.curoverse.com/arvados.git/sdk/go/arvadosclient"
  19         "git.curoverse.com/arvados.git/sdk/go/config"
  20         "git.curoverse.com/arvados.git/sdk/go/dispatch"
  21         "github.com/coreos/go-systemd/daemon"
  22         "github.com/ghodss/yaml"
  23 )
  24
  25 // Config used by crunch-dispatch-slurm
  26 type Config struct {
  27         Client arvados.Client
  28
  29         SbatchArguments []string
  30         PollPeriod      arvados.Duration
  31
  32         // crunch-run command to invoke. The container UUID will be
  33         // appended. If nil, []string{"crunch-run"} will be used.
  34         //
  35         // Example: []string{"crunch-run", "--cgroup-parent-subsystem=memory"}
  36         CrunchRunCommand []string
  37 }
  38
  39 func main() {
  40         err := doMain()
  41         if err != nil {
  42                 log.Fatal(err)
  43         }
  44 }
  45
  46 var (
  47         theConfig     Config
  48         squeueUpdater Squeue
  49 )
  50
  51 const defaultConfigPath = "/etc/arvados/crunch-dispatch-slurm/crunch-dispatch-slurm.yml"
  52
  53 func doMain() error {
  54         flags := flag.NewFlagSet("crunch-dispatch-slurm", flag.ExitOnError)
  55         flags.Usage = func() { usage(flags) }
  56
  57         configPath := flags.String(
  58                 "config",
  59                 defaultConfigPath,
  60                 "`path` to JSON or YAML configuration file")
  61         dumpConfig := flag.Bool(
  62                 "dump-config",
  63                 false,
  64                 "write current configuration to stdout and exit")
  65
  66         // Parse args; omit the first arg which is the command name
  67         flags.Parse(os.Args[1:])
  68
  69         err := readConfig(&theConfig, *configPath)
  70         if err != nil {
  71                 return err
  72         }
  73
  74         if theConfig.CrunchRunCommand == nil {
  75                 theConfig.CrunchRunCommand = []string{"crunch-run"}
  76         }
  77
  78         if theConfig.PollPeriod == 0 {
  79                 theConfig.PollPeriod = arvados.Duration(10 * time.Second)
  80         }
  81
  82         if theConfig.Client.APIHost != "" || theConfig.Client.AuthToken != "" {
  83                 // Copy real configs into env vars so [a]
  84                 // MakeArvadosClient() uses them, and [b] they get
  85                 // propagated to crunch-run via SLURM.
  86                 os.Setenv("ARVADOS_API_HOST", theConfig.Client.APIHost)
  87                 os.Setenv("ARVADOS_API_TOKEN", theConfig.Client.AuthToken)
  88                 os.Setenv("ARVADOS_API_HOST_INSECURE", "")
  89                 if theConfig.Client.Insecure {
  90                         os.Setenv("ARVADOS_API_HOST_INSECURE", "1")
  91                 }
  92                 os.Setenv("ARVADOS_KEEP_SERVICES", strings.Join(theConfig.Client.KeepServiceURIs, " "))
  93                 os.Setenv("ARVADOS_EXTERNAL_CLIENT", "")
  94         } else {
  95                 log.Printf("warning: Client credentials missing from config, so falling back on environment variables (deprecated).")
  96         }
  97
  98         if *dumpConfig {
  99                 y, err := yaml.Marshal(theConfig)
 100                 if err != nil {
 101                         log.Fatal(err)
 102                 }
 103                 os.Stdout.Write(y)
 104                 os.Exit(0)
 105         }
 106
 107         arv, err := arvadosclient.MakeArvadosClient()
 108         if err != nil {
 109                 log.Printf("Error making Arvados client: %v", err)
 110                 return err
 111         }
 112         arv.Retries = 25
 113
 114         squeueUpdater.StartMonitor(time.Duration(theConfig.PollPeriod))
 115         defer squeueUpdater.Done()
 116
 117         dispatcher := dispatch.Dispatcher{
 118                 Arv:            arv,
 119                 RunContainer:   run,
 120                 PollInterval:   time.Duration(theConfig.PollPeriod),
 121                 DoneProcessing: make(chan struct{})}
 122
 123         if _, err := daemon.SdNotify(false, "READY=1"); err != nil {
 124                 log.Printf("Error notifying init daemon: %v", err)
 125         }
 126
 127         err = dispatcher.RunDispatcher()
 128         if err != nil {
 129                 return err
 130         }
 131
 132         return nil
 133 }
 134
 135 // sbatchCmd
 136 func sbatchFunc(container arvados.Container) *exec.Cmd {
 137         memPerCPU := math.Ceil(float64(container.RuntimeConstraints.RAM) / (float64(container.RuntimeConstraints.VCPUs) * 1048576))
 138
 139         var sbatchArgs []string
 140         sbatchArgs = append(sbatchArgs, "--share")
 141         sbatchArgs = append(sbatchArgs, theConfig.SbatchArguments...)
 142         sbatchArgs = append(sbatchArgs, fmt.Sprintf("--job-name=%s", container.UUID))
 143         sbatchArgs = append(sbatchArgs, fmt.Sprintf("--mem-per-cpu=%d", int(memPerCPU)))
 144         sbatchArgs = append(sbatchArgs, fmt.Sprintf("--cpus-per-task=%d", container.RuntimeConstraints.VCPUs))
 145         if container.SchedulingParameters.Partitions != nil {
 146                 sbatchArgs = append(sbatchArgs, fmt.Sprintf("--partition=%s", strings.Join(container.SchedulingParameters.Partitions, ",")))
 147         }
 148
 149         return exec.Command("sbatch", sbatchArgs...)
 150 }
 151
 152 // scancelCmd
 153 func scancelFunc(container arvados.Container) *exec.Cmd {
 154         return exec.Command("scancel", "--name="+container.UUID)
 155 }
 156
 157 // Wrap these so that they can be overridden by tests
 158 var sbatchCmd = sbatchFunc
 159 var scancelCmd = scancelFunc
 160
 161 // Submit job to slurm using sbatch.
 162 func submit(dispatcher *dispatch.Dispatcher,
 163         container arvados.Container, crunchRunCommand []string) (submitErr error) {
 164         defer func() {
 165                 // If we didn't get as far as submitting a slurm job,
 166                 // unlock the container and return it to the queue.
 167                 if submitErr == nil {
 168                         // OK, no cleanup needed
 169                         return
 170                 }
 171                 err := dispatcher.Unlock(container.UUID)
 172                 if err != nil {
 173                         log.Printf("Error unlocking container %s: %v", container.UUID, err)
 174                 }
 175         }()
 176
 177         // Create the command and attach to stdin/stdout
 178         cmd := sbatchCmd(container)
 179         stdinWriter, stdinerr := cmd.StdinPipe()
 180         if stdinerr != nil {
 181                 submitErr = fmt.Errorf("Error creating stdin pipe %v: %q", container.UUID, stdinerr)
 182                 return
 183         }
 184
 185         stdoutReader, stdoutErr := cmd.StdoutPipe()
 186         if stdoutErr != nil {
 187                 submitErr = fmt.Errorf("Error creating stdout pipe %v: %q", container.UUID, stdoutErr)
 188                 return
 189         }
 190
 191         stderrReader, stderrErr := cmd.StderrPipe()
 192         if stderrErr != nil {
 193                 submitErr = fmt.Errorf("Error creating stderr pipe %v: %q", container.UUID, stderrErr)
 194                 return
 195         }
 196
 197         // Mutex between squeue sync and running sbatch or scancel.
 198         squeueUpdater.SlurmLock.Lock()
 199         defer squeueUpdater.SlurmLock.Unlock()
 200
 201         log.Printf("sbatch starting: %+q", cmd.Args)
 202         err := cmd.Start()
 203         if err != nil {
 204                 submitErr = fmt.Errorf("Error starting sbatch: %v", err)
 205                 return
 206         }
 207
 208         stdoutChan := make(chan []byte)
 209         go func() {
 210                 b, _ := ioutil.ReadAll(stdoutReader)
 211                 stdoutReader.Close()
 212                 stdoutChan <- b
 213                 close(stdoutChan)
 214         }()
 215
 216         stderrChan := make(chan []byte)
 217         go func() {
 218                 b, _ := ioutil.ReadAll(stderrReader)
 219                 stderrReader.Close()
 220                 stderrChan <- b
 221                 close(stderrChan)
 222         }()
 223
 224         // Send a tiny script on stdin to execute the crunch-run command
 225         // slurm actually enforces that this must be a #! script
 226         io.WriteString(stdinWriter, execScript(append(crunchRunCommand, container.UUID)))
 227         stdinWriter.Close()
 228
 229         stdoutMsg := <-stdoutChan
 230         stderrmsg := <-stderrChan
 231
 232         err = cmd.Wait()
 233
 234         if err != nil {
 235                 submitErr = fmt.Errorf("Container submission failed: %v: %v (stderr: %q)", cmd.Args, err, stderrmsg)
 236                 return
 237         }
 238
 239         log.Printf("sbatch succeeded: %s", strings.TrimSpace(string(stdoutMsg)))
 240         return
 241 }
 242
 243 // If the container is marked as Locked, check if it is already in the slurm
 244 // queue.  If not, submit it.
 245 //
 246 // If the container is marked as Running, check if it is in the slurm queue.
 247 // If not, mark it as Cancelled.
 248 func monitorSubmitOrCancel(dispatcher *dispatch.Dispatcher, container arvados.Container, monitorDone *bool) {
 249         submitted := false
 250         for !*monitorDone {
 251                 if squeueUpdater.CheckSqueue(container.UUID) {
 252                         // Found in the queue, so continue monitoring
 253                         submitted = true
 254                 } else if container.State == dispatch.Locked && !submitted {
 255                         // Not in queue but in Locked state and we haven't
 256                         // submitted it yet, so submit it.
 257
 258                         log.Printf("About to submit queued container %v", container.UUID)
 259
 260                         if err := submit(dispatcher, container, theConfig.CrunchRunCommand); err != nil {
 261                                 log.Printf("Error submitting container %s to slurm: %v",
 262                                         container.UUID, err)
 263                                 // maybe sbatch is broken, put it back to queued
 264                                 dispatcher.Unlock(container.UUID)
 265                         }
 266                         submitted = true
 267                 } else {
 268                         // Not in queue and we are not going to submit it.
 269                         // Refresh the container state. If it is
 270                         // Complete/Cancelled, do nothing, if it is Locked then
 271                         // release it back to the Queue, if it is Running then
 272                         // clean up the record.
 273
 274                         var con arvados.Container
 275                         err := dispatcher.Arv.Get("containers", container.UUID, nil, &con)
 276                         if err != nil {
 277                                 log.Printf("Error getting final container state: %v", err)
 278                         }
 279
 280                         switch con.State {
 281                         case dispatch.Locked:
 282                                 log.Printf("Container %s in state %v but missing from slurm queue, changing to %v.",
 283                                         container.UUID, con.State, dispatch.Queued)
 284                                 dispatcher.Unlock(container.UUID)
 285                         case dispatch.Running:
 286                                 st := dispatch.Cancelled
 287                                 log.Printf("Container %s in state %v but missing from slurm queue, changing to %v.",
 288                                         container.UUID, con.State, st)
 289                                 dispatcher.UpdateState(container.UUID, st)
 290                         default:
 291                                 // Container state is Queued, Complete or Cancelled so stop monitoring it.
 292                                 return
 293                         }
 294                 }
 295         }
 296 }
 297
 298 // Run or monitor a container.
 299 //
 300 // Monitor status updates.  If the priority changes to zero, cancel the
 301 // container using scancel.
 302 func run(dispatcher *dispatch.Dispatcher,
 303         container arvados.Container,
 304         status chan arvados.Container) {
 305
 306         log.Printf("Monitoring container %v started", container.UUID)
 307         defer log.Printf("Monitoring container %v finished", container.UUID)
 308
 309         monitorDone := false
 310         go monitorSubmitOrCancel(dispatcher, container, &monitorDone)
 311
 312         for container = range status {
 313                 if container.State == dispatch.Locked || container.State == dispatch.Running {
 314                         if container.Priority == 0 {
 315                                 log.Printf("Canceling container %s", container.UUID)
 316
 317                                 // Mutex between squeue sync and running sbatch or scancel.
 318                                 squeueUpdater.SlurmLock.Lock()
 319                                 cmd := scancelCmd(container)
 320                                 msg, err := cmd.CombinedOutput()
 321                                 squeueUpdater.SlurmLock.Unlock()
 322
 323                                 if err != nil {
 324                                         log.Printf("Error stopping container %s with %v %v: %v %v",
 325                                                 container.UUID, cmd.Path, cmd.Args, err, string(msg))
 326                                         if squeueUpdater.CheckSqueue(container.UUID) {
 327                                                 log.Printf("Container %s is still in squeue after scancel.",
 328                                                         container.UUID)
 329                                                 continue
 330                                         }
 331                                 }
 332
 333                                 err = dispatcher.UpdateState(container.UUID, dispatch.Cancelled)
 334                         }
 335                 }
 336         }
 337         monitorDone = true
 338 }
 339
 340 func readConfig(dst interface{}, path string) error {
 341         err := config.LoadFile(dst, path)
 342         if err != nil && os.IsNotExist(err) && path == defaultConfigPath {
 343                 log.Printf("Config not specified. Continue with default configuration.")
 344                 err = nil
 345         }
 346         return err
 347 }