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         "bufio"
   7         "flag"
   8         "fmt"
   9         "git.curoverse.com/arvados.git/sdk/go/arvadosclient"
  10         "git.curoverse.com/arvados.git/sdk/go/dispatch"
  11         "io/ioutil"
  12         "log"
  13         "math"
  14         "os"
  15         "os/exec"
  16         "strings"
  17         "sync"
  18         "time"
  19 )
  20
  21 type Squeue struct {
  22         sync.Mutex
  23         squeueContents []string
  24         SqueueDone     chan struct{}
  25 }
  26
  27 func main() {
  28         err := doMain()
  29         if err != nil {
  30                 log.Fatalf("%q", err)
  31         }
  32 }
  33
  34 var (
  35         crunchRunCommand *string
  36         squeueUpdater    Squeue
  37 )
  38
  39 func doMain() error {
  40         flags := flag.NewFlagSet("crunch-dispatch-slurm", flag.ExitOnError)
  41
  42         pollInterval := flags.Int(
  43                 "poll-interval",
  44                 10,
  45                 "Interval in seconds to poll for queued containers")
  46
  47         crunchRunCommand = flags.String(
  48                 "crunch-run-command",
  49                 "/usr/bin/crunch-run",
  50                 "Crunch command to run container")
  51
  52         // Parse args; omit the first arg which is the command name
  53         flags.Parse(os.Args[1:])
  54
  55         arv, err := arvadosclient.MakeArvadosClient()
  56         if err != nil {
  57                 log.Printf("Error making Arvados client: %v", err)
  58                 return err
  59         }
  60         arv.Retries = 25
  61
  62         dispatcher := dispatch.Dispatcher{
  63                 Arv:            arv,
  64                 RunContainer:   run,
  65                 PollInterval:   time.Duration(*pollInterval) * time.Second,
  66                 DoneProcessing: make(chan struct{})}
  67
  68         squeueUpdater.SqueueDone = make(chan struct{})
  69         go squeueUpdater.SyncSqueue(time.Duration(*pollInterval) * time.Second)
  70
  71         err = dispatcher.RunDispatcher()
  72         if err != nil {
  73                 return err
  74         }
  75
  76         squeueUpdater.SqueueDone <- struct{}{}
  77         close(squeueUpdater.SqueueDone)
  78
  79         return nil
  80 }
  81
  82 // sbatchCmd
  83 func sbatchFunc(container dispatch.Container) *exec.Cmd {
  84         memPerCPU := math.Ceil((float64(container.RuntimeConstraints["ram"])) / (float64(container.RuntimeConstraints["vcpus"] * 1048576)))
  85         return exec.Command("sbatch", "--share", "--parsable",
  86                 fmt.Sprintf("--job-name=%s", container.UUID),
  87                 fmt.Sprintf("--mem-per-cpu=%d", int(memPerCPU)),
  88                 fmt.Sprintf("--cpus-per-task=%d", int(container.RuntimeConstraints["vcpus"])),
  89                 fmt.Sprintf("--priority=%d", container.Priority))
  90 }
  91
  92 // squeueFunc
  93 func squeueFunc() *exec.Cmd {
  94         return exec.Command("squeue", "--format=%j")
  95 }
  96
  97 // Wrap these so that they can be overridden by tests
  98 var sbatchCmd = sbatchFunc
  99 var squeueCmd = squeueFunc
 100
 101 // Submit job to slurm using sbatch.
 102 func submit(dispatcher *dispatch.Dispatcher,
 103         container dispatch.Container, crunchRunCommand string) (jobid string, submitErr error) {
 104         submitErr = nil
 105
 106         defer func() {
 107                 // If we didn't get as far as submitting a slurm job,
 108                 // unlock the container and return it to the queue.
 109                 if submitErr == nil {
 110                         // OK, no cleanup needed
 111                         return
 112                 }
 113                 err := dispatcher.Arv.Update("containers", container.UUID,
 114                         arvadosclient.Dict{
 115                                 "container": arvadosclient.Dict{"state": "Queued"}},
 116                         nil)
 117                 if err != nil {
 118                         log.Printf("Error unlocking container %s: %v", container.UUID, err)
 119                 }
 120         }()
 121
 122         // Create the command and attach to stdin/stdout
 123         cmd := sbatchCmd(container)
 124         stdinWriter, stdinerr := cmd.StdinPipe()
 125         if stdinerr != nil {
 126                 submitErr = fmt.Errorf("Error creating stdin pipe %v: %q", container.UUID, stdinerr)
 127                 return
 128         }
 129
 130         stdoutReader, stdoutErr := cmd.StdoutPipe()
 131         if stdoutErr != nil {
 132                 submitErr = fmt.Errorf("Error creating stdout pipe %v: %q", container.UUID, stdoutErr)
 133                 return
 134         }
 135
 136         stderrReader, stderrErr := cmd.StderrPipe()
 137         if stderrErr != nil {
 138                 submitErr = fmt.Errorf("Error creating stderr pipe %v: %q", container.UUID, stderrErr)
 139                 return
 140         }
 141
 142         err := cmd.Start()
 143         if err != nil {
 144                 submitErr = fmt.Errorf("Error starting %v: %v", cmd.Args, err)
 145                 return
 146         }
 147
 148         stdoutChan := make(chan []byte)
 149         go func() {
 150                 b, _ := ioutil.ReadAll(stdoutReader)
 151                 stdoutReader.Close()
 152                 stdoutChan <- b
 153         }()
 154
 155         stderrChan := make(chan []byte)
 156         go func() {
 157                 b, _ := ioutil.ReadAll(stderrReader)
 158                 stderrReader.Close()
 159                 stderrChan <- b
 160         }()
 161
 162         // Send a tiny script on stdin to execute the crunch-run command
 163         // slurm actually enforces that this must be a #! script
 164         fmt.Fprintf(stdinWriter, "#!/bin/sh\nexec '%s' '%s'\n", crunchRunCommand, container.UUID)
 165         stdinWriter.Close()
 166
 167         err = cmd.Wait()
 168
 169         stdoutMsg := <-stdoutChan
 170         stderrmsg := <-stderrChan
 171
 172         close(stdoutChan)
 173         close(stderrChan)
 174
 175         if err != nil {
 176                 submitErr = fmt.Errorf("Container submission failed %v: %v %v", cmd.Args, err, stderrmsg)
 177                 return
 178         }
 179
 180         // If everything worked out, got the jobid on stdout
 181         jobid = strings.TrimSpace(string(stdoutMsg))
 182
 183         return
 184 }
 185
 186 func (squeue *Squeue) runSqueue() ([]string, error) {
 187         var newSqueueContents []string
 188
 189         cmd := squeueCmd()
 190         sq, err := cmd.StdoutPipe()
 191         if err != nil {
 192                 return nil, err
 193         }
 194         cmd.Start()
 195         scanner := bufio.NewScanner(sq)
 196         for scanner.Scan() {
 197                 newSqueueContents = append(newSqueueContents, scanner.Text())
 198         }
 199         if err := scanner.Err(); err != nil {
 200                 cmd.Wait()
 201                 return nil, err
 202         }
 203
 204         err = cmd.Wait()
 205         if err != nil {
 206                 return nil, err
 207         }
 208
 209         return newSqueueContents, nil
 210 }
 211
 212 func (squeue *Squeue) CheckSqueue(uuid string, check bool) (bool, error) {
 213         if check {
 214                 n, err := squeue.runSqueue()
 215                 if err != nil {
 216                         return false, err
 217                 }
 218                 squeue.Lock()
 219                 squeue.squeueContents = n
 220                 squeue.Unlock()
 221         }
 222
 223         if uuid != "" {
 224                 squeue.Lock()
 225                 defer squeue.Unlock()
 226                 for _, k := range squeue.squeueContents {
 227                         if k == uuid {
 228                                 return true, nil
 229                         }
 230                 }
 231         }
 232         return false, nil
 233 }
 234
 235 func (squeue *Squeue) SyncSqueue(pollInterval time.Duration) {
 236         // TODO: considering using "squeue -i" instead of polling squeue.
 237         ticker := time.NewTicker(pollInterval)
 238         for {
 239                 select {
 240                 case <-squeueUpdater.SqueueDone:
 241                         return
 242                 case <-ticker.C:
 243                         squeue.CheckSqueue("", true)
 244                 }
 245         }
 246 }
 247
 248 // Run or monitor a container.
 249 //
 250 // If the container is marked as Locked, check if it is already in the slurm
 251 // queue.  If not, submit it.
 252 //
 253 // If the container is marked as Running, check if it is in the slurm queue.
 254 // If not, mark it as Cancelled.
 255 //
 256 // Monitor status updates.  If the priority changes to zero, cancel the
 257 // container using scancel.
 258 func run(dispatcher *dispatch.Dispatcher,
 259         container dispatch.Container,
 260         status chan dispatch.Container) {
 261
 262         uuid := container.UUID
 263
 264         if container.State == dispatch.Locked {
 265                 if inQ, err := squeueUpdater.CheckSqueue(container.UUID, true); err != nil {
 266                         // maybe squeue is broken, put it back in the queue
 267                         log.Printf("Error running squeue: %v", err)
 268                         dispatcher.UpdateState(container.UUID, dispatch.Queued)
 269                 } else if !inQ {
 270                         log.Printf("About to submit queued container %v", container.UUID)
 271
 272                         if _, err := submit(dispatcher, container, *crunchRunCommand); err != nil {
 273                                 log.Printf("Error submitting container %s to slurm: %v",
 274                                         container.UUID, err)
 275                                 // maybe sbatch is broken, put it back to queued
 276                                 dispatcher.UpdateState(container.UUID, dispatch.Queued)
 277                         }
 278                 }
 279         }
 280
 281         log.Printf("Monitoring container %v started", uuid)
 282
 283         // periodically check squeue
 284         doneSqueue := make(chan struct{})
 285         go func() {
 286                 squeueUpdater.CheckSqueue(container.UUID, true)
 287                 ticker := time.NewTicker(dispatcher.PollInterval)
 288                 for {
 289                         select {
 290                         case <-ticker.C:
 291                                 if inQ, err := squeueUpdater.CheckSqueue(container.UUID, false); err != nil {
 292                                         log.Printf("Error running squeue: %v", err)
 293                                         // don't cancel, just leave it the way it is
 294                                 } else if !inQ {
 295                                         var con dispatch.Container
 296                                         err := dispatcher.Arv.Get("containers", uuid, nil, &con)
 297                                         if err != nil {
 298                                                 log.Printf("Error getting final container state: %v", err)
 299                                         }
 300
 301                                         var st string
 302                                         switch con.State {
 303                                         case dispatch.Locked:
 304                                                 st = dispatch.Queued
 305                                         case dispatch.Running:
 306                                                 st = dispatch.Cancelled
 307                                         default:
 308                                                 st = ""
 309                                         }
 310
 311                                         if st != "" {
 312                                                 log.Printf("Container %s in state %v but missing from slurm queue, changing to %v.",
 313                                                         uuid, con.State, st)
 314                                                 dispatcher.UpdateState(uuid, st)
 315                                         }
 316                                 }
 317                         case <-doneSqueue:
 318                                 close(doneSqueue)
 319                                 ticker.Stop()
 320                                 return
 321                         }
 322                 }
 323         }()
 324
 325         for container = range status {
 326                 if container.State == dispatch.Locked || container.State == dispatch.Running {
 327                         if container.Priority == 0 {
 328                                 log.Printf("Canceling container %s", container.UUID)
 329
 330                                 err := exec.Command("scancel", "--name="+container.UUID).Run()
 331                                 if err != nil {
 332                                         log.Printf("Error stopping container %s with scancel: %v",
 333                                                 container.UUID, err)
 334                                         if inQ, err := squeueUpdater.CheckSqueue(container.UUID, true); err != nil {
 335                                                 log.Printf("Error running squeue: %v", err)
 336                                                 continue
 337                                         } else if inQ {
 338                                                 log.Printf("Container %s is still in squeue after scancel.",
 339                                                         container.UUID)
 340                                                 continue
 341                                         }
 342                                 }
 343
 344                                 err = dispatcher.UpdateState(container.UUID, dispatch.Cancelled)
 345                         }
 346                 }
 347         }
 348
 349         doneSqueue <- struct{}{}
 350
 351         log.Printf("Monitoring container %v finished", uuid)
 352 }