services/crunch-dispatch-slurm/slurm.go

   1 // Copyright (C) The Arvados Authors. All rights reserved.
   2 //
   3 // SPDX-License-Identifier: AGPL-3.0
   4
   5 package main
   6
   7 import (
   8         "fmt"
   9         "io"
  10         "log"
  11         "os/exec"
  12         "strings"
  13 )
  14
  15 type Slurm interface {
  16         Cancel(name string) error
  17         Renice(name string, nice int) error
  18         QueueCommand(args []string) *exec.Cmd
  19         Batch(script io.Reader, args []string) error
  20 }
  21
  22 type slurmCLI struct{}
  23
  24 func (scli *slurmCLI) Batch(script io.Reader, args []string) error {
  25         return scli.run(script, "sbatch", args)
  26 }
  27
  28 func (scli *slurmCLI) Cancel(name string) error {
  29         for _, args := range [][]string{
  30                 // If the slurm job hasn't started yet, remove it from
  31                 // the queue.
  32                 {"--state=pending"},
  33                 // If the slurm job has started, send SIGTERM. If we
  34                 // cancel a running job without a --signal argument,
  35                 // slurm will send SIGTERM and then (after some
  36                 // site-configured interval) SIGKILL. This would kill
  37                 // crunch-run without stopping the container, which we
  38                 // don't want.
  39                 {"--batch", "--signal=TERM", "--state=running"},
  40                 {"--batch", "--signal=TERM", "--state=suspended"},
  41         } {
  42                 err := scli.run(nil, "scancel", append([]string{"--name=" + name}, args...))
  43                 if err != nil {
  44                         // scancel exits 0 if no job matches the given
  45                         // name and state. Any error from scancel here
  46                         // really indicates something is wrong.
  47                         return err
  48                 }
  49         }
  50         return nil
  51 }
  52
  53 func (scli *slurmCLI) QueueCommand(args []string) *exec.Cmd {
  54         return exec.Command("squeue", args...)
  55 }
  56
  57 func (scli *slurmCLI) Renice(name string, nice int) error {
  58         return scli.run(nil, "scontrol", []string{"update", "JobName=" + name, fmt.Sprintf("Nice=%d", nice)})
  59 }
  60
  61 func (scli *slurmCLI) run(stdin io.Reader, prog string, args []string) error {
  62         cmd := exec.Command(prog, args...)
  63         cmd.Stdin = stdin
  64         out, err := cmd.CombinedOutput()
  65         outTrim := strings.TrimSpace(string(out))
  66         if err != nil || len(out) > 0 {
  67                 log.Printf("%q %q: %q", cmd.Path, cmd.Args, outTrim)
  68         }
  69         if err != nil {
  70                 err = fmt.Errorf("%s: %s (%q)", cmd.Path, err, outTrim)
  71         }
  72         return err
  73 }