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         return scli.run(nil, "scancel", []string{"--name=" + name})
  30 }
  31
  32 func (scli *slurmCLI) QueueCommand(args []string) *exec.Cmd {
  33         return exec.Command("squeue", args...)
  34 }
  35
  36 func (scli *slurmCLI) Renice(name string, nice int) error {
  37         return scli.run(nil, "scontrol", []string{"update", "JobName=" + name, fmt.Sprintf("Nice=%d", nice)})
  38 }
  39
  40 func (scli *slurmCLI) run(stdin io.Reader, prog string, args []string) error {
  41         cmd := exec.Command(prog, args...)
  42         cmd.Stdin = stdin
  43         out, err := cmd.CombinedOutput()
  44         outTrim := strings.TrimSpace(string(out))
  45         if err != nil || len(out) > 0 {
  46                 log.Printf("%q %q: %q", cmd.Path, cmd.Args, outTrim)
  47         }
  48         if err != nil {
  49                 err = fmt.Errorf("%s: %s (%q)", cmd.Path, err, outTrim)
  50         }
  51         return err
  52 }