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         Batch(script io.Reader, args []string) error
  17         Cancel(name string) error
  18         QueueCommand(args []string) *exec.Cmd
  19         Release(name string) error
  20         Renice(name string, nice int64) error
  21 }
  22
  23 type slurmCLI struct{}
  24
  25 func (scli *slurmCLI) Batch(script io.Reader, args []string) error {
  26         return scli.run(script, "sbatch", args)
  27 }
  28
  29 func (scli *slurmCLI) Cancel(name string) error {
  30         for _, args := range [][]string{
  31                 // If the slurm job hasn't started yet, remove it from
  32                 // the queue.
  33                 {"--state=pending"},
  34                 // If the slurm job has started, send SIGTERM. If we
  35                 // cancel a running job without a --signal argument,
  36                 // slurm will send SIGTERM and then (after some
  37                 // site-configured interval) SIGKILL. This would kill
  38                 // crunch-run without stopping the container, which we
  39                 // don't want.
  40                 {"--batch", "--signal=TERM", "--state=running"},
  41                 {"--batch", "--signal=TERM", "--state=suspended"},
  42         } {
  43                 err := scli.run(nil, "scancel", append([]string{"--name=" + name}, args...))
  44                 if err != nil {
  45                         // scancel exits 0 if no job matches the given
  46                         // name and state. Any error from scancel here
  47                         // really indicates something is wrong.
  48                         return err
  49                 }
  50         }
  51         return nil
  52 }
  53
  54 func (scli *slurmCLI) QueueCommand(args []string) *exec.Cmd {
  55         return exec.Command("squeue", args...)
  56 }
  57
  58 func (scli *slurmCLI) Release(name string) error {
  59         return scli.run(nil, "scontrol", []string{"release", "Name=" + name})
  60 }
  61
  62 func (scli *slurmCLI) Renice(name string, nice int64) error {
  63         return scli.run(nil, "scontrol", []string{"update", "JobName=" + name, fmt.Sprintf("Nice=%d", nice)})
  64 }
  65
  66 func (scli *slurmCLI) run(stdin io.Reader, prog string, args []string) error {
  67         cmd := exec.Command(prog, args...)
  68         cmd.Stdin = stdin
  69         out, err := cmd.CombinedOutput()
  70         outTrim := strings.TrimSpace(string(out))
  71         if err != nil || len(out) > 0 {
  72                 log.Printf("%q %q: %q", cmd.Path, cmd.Args, outTrim)
  73         }
  74         if err != nil {
  75                 err = fmt.Errorf("%s: %s (%q)", cmd.Path, err, outTrim)
  76         }
  77         return err
  78 }