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

   1 // Copyright (C) The Arvados Authors. All rights reserved.
   2 //
   3 // SPDX-License-Identifier: AGPL-3.0
   4
   5 // Dispatcher service for Crunch that submits containers to the slurm queue.
   6 package dispatchslurm
   7
   8 import (
   9         "context"
  10         "fmt"
  11         "log"
  12         "math"
  13         "net/http"
  14         "os"
  15         "regexp"
  16         "strings"
  17         "time"
  18
  19         "git.arvados.org/arvados.git/lib/cmd"
  20         "git.arvados.org/arvados.git/lib/dispatchcloud"
  21         "git.arvados.org/arvados.git/lib/service"
  22         "git.arvados.org/arvados.git/sdk/go/arvados"
  23         "git.arvados.org/arvados.git/sdk/go/arvadosclient"
  24         "git.arvados.org/arvados.git/sdk/go/ctxlog"
  25         "git.arvados.org/arvados.git/sdk/go/dispatch"
  26         "github.com/coreos/go-systemd/daemon"
  27         "github.com/prometheus/client_golang/prometheus"
  28         "github.com/sirupsen/logrus"
  29 )
  30
  31 var Command cmd.Handler = service.Command(arvados.ServiceNameDispatchSLURM, newHandler)
  32
  33 func newHandler(ctx context.Context, cluster *arvados.Cluster, _ string, _ *prometheus.Registry) service.Handler {
  34         logger := ctxlog.FromContext(ctx)
  35         disp := &Dispatcher{logger: logger, cluster: cluster}
  36         if err := disp.configure(); err != nil {
  37                 return service.ErrorHandler(ctx, cluster, err)
  38         }
  39         disp.setup()
  40         go func() {
  41                 disp.err = disp.run()
  42                 close(disp.done)
  43         }()
  44         return disp
  45 }
  46
  47 type logger interface {
  48         dispatch.Logger
  49         Fatalf(string, ...interface{})
  50 }
  51
  52 const initialNiceValue int64 = 10000
  53
  54 type Dispatcher struct {
  55         *dispatch.Dispatcher
  56         logger  logrus.FieldLogger
  57         cluster *arvados.Cluster
  58         sqCheck *SqueueChecker
  59         slurm   Slurm
  60
  61         done chan struct{}
  62         err  error
  63
  64         Client arvados.Client
  65 }
  66
  67 func (disp *Dispatcher) CheckHealth() error {
  68         return disp.err
  69 }
  70
  71 func (disp *Dispatcher) Done() <-chan struct{} {
  72         return disp.done
  73 }
  74
  75 func (disp *Dispatcher) ServeHTTP(w http.ResponseWriter, r *http.Request) {
  76         http.NotFound(w, r)
  77 }
  78
  79 // configure() loads config files. Some tests skip this (see
  80 // StubbedSuite).
  81 func (disp *Dispatcher) configure() error {
  82         if disp.logger == nil {
  83                 disp.logger = logrus.StandardLogger()
  84         }
  85         disp.logger = disp.logger.WithField("ClusterID", disp.cluster.ClusterID)
  86         disp.logger.Printf("crunch-dispatch-slurm %s started", cmd.Version.String())
  87
  88         disp.Client.APIHost = disp.cluster.Services.Controller.ExternalURL.Host
  89         disp.Client.AuthToken = disp.cluster.SystemRootToken
  90         disp.Client.Insecure = disp.cluster.TLS.Insecure
  91
  92         if disp.Client.APIHost != "" || disp.Client.AuthToken != "" {
  93                 // Copy real configs into env vars so [a]
  94                 // MakeArvadosClient() uses them, and [b] they get
  95                 // propagated to crunch-run via SLURM.
  96                 os.Setenv("ARVADOS_API_HOST", disp.Client.APIHost)
  97                 os.Setenv("ARVADOS_API_TOKEN", disp.Client.AuthToken)
  98                 os.Setenv("ARVADOS_API_HOST_INSECURE", "")
  99                 if disp.Client.Insecure {
 100                         os.Setenv("ARVADOS_API_HOST_INSECURE", "1")
 101                 }
 102                 os.Setenv("ARVADOS_EXTERNAL_CLIENT", "")
 103                 for k, v := range disp.cluster.Containers.SLURM.SbatchEnvironmentVariables {
 104                         os.Setenv(k, v)
 105                 }
 106         } else {
 107                 disp.logger.Warnf("Client credentials missing from config, so falling back on environment variables (deprecated).")
 108         }
 109         return nil
 110 }
 111
 112 // setup() initializes private fields after configure().
 113 func (disp *Dispatcher) setup() {
 114         disp.done = make(chan struct{})
 115         arv, err := arvadosclient.MakeArvadosClient()
 116         if err != nil {
 117                 disp.logger.Fatalf("Error making Arvados client: %v", err)
 118         }
 119         arv.Retries = 25
 120
 121         disp.slurm = NewSlurmCLI()
 122         disp.sqCheck = &SqueueChecker{
 123                 Logger:         disp.logger,
 124                 Period:         time.Duration(disp.cluster.Containers.CloudVMs.PollInterval),
 125                 PrioritySpread: disp.cluster.Containers.SLURM.PrioritySpread,
 126                 Slurm:          disp.slurm,
 127         }
 128         disp.Dispatcher = &dispatch.Dispatcher{
 129                 Arv:            arv,
 130                 Logger:         disp.logger,
 131                 BatchSize:      disp.cluster.API.MaxItemsPerResponse,
 132                 RunContainer:   disp.runContainer,
 133                 PollPeriod:     time.Duration(disp.cluster.Containers.CloudVMs.PollInterval),
 134                 MinRetryPeriod: time.Duration(disp.cluster.Containers.MinRetryPeriod),
 135         }
 136 }
 137
 138 func (disp *Dispatcher) run() error {
 139         defer disp.sqCheck.Stop()
 140
 141         if disp.cluster != nil && len(disp.cluster.InstanceTypes) > 0 {
 142                 go SlurmNodeTypeFeatureKludge(disp.cluster)
 143         }
 144
 145         if _, err := daemon.SdNotify(false, "READY=1"); err != nil {
 146                 log.Printf("Error notifying init daemon: %v", err)
 147         }
 148         go disp.checkSqueueForOrphans()
 149         return disp.Dispatcher.Run(context.Background())
 150 }
 151
 152 var containerUuidPattern = regexp.MustCompile(`^[a-z0-9]{5}-dz642-[a-z0-9]{15}$`)
 153
 154 // Check the next squeue report, and invoke TrackContainer for all the
 155 // containers in the report. This gives us a chance to cancel slurm
 156 // jobs started by a previous dispatch process that never released
 157 // their slurm allocations even though their container states are
 158 // Cancelled or Complete. See https://dev.arvados.org/issues/10979
 159 func (disp *Dispatcher) checkSqueueForOrphans() {
 160         for _, uuid := range disp.sqCheck.All() {
 161                 if !containerUuidPattern.MatchString(uuid) || !strings.HasPrefix(uuid, disp.cluster.ClusterID) {
 162                         continue
 163                 }
 164                 err := disp.TrackContainer(uuid)
 165                 if err != nil {
 166                         log.Printf("checkSqueueForOrphans: TrackContainer(%s): %s", uuid, err)
 167                 }
 168         }
 169 }
 170
 171 func (disp *Dispatcher) slurmConstraintArgs(container arvados.Container) []string {
 172         mem := int64(math.Ceil(float64(container.RuntimeConstraints.RAM+
 173                 container.RuntimeConstraints.KeepCacheRAM+
 174                 int64(disp.cluster.Containers.ReserveExtraRAM)) / float64(1048576)))
 175
 176         disk := dispatchcloud.EstimateScratchSpace(&container)
 177         disk = int64(math.Ceil(float64(disk) / float64(1048576)))
 178         return []string{
 179                 fmt.Sprintf("--mem=%d", mem),
 180                 fmt.Sprintf("--cpus-per-task=%d", container.RuntimeConstraints.VCPUs),
 181                 fmt.Sprintf("--tmp=%d", disk),
 182         }
 183 }
 184
 185 func (disp *Dispatcher) sbatchArgs(container arvados.Container) ([]string, error) {
 186         var args []string
 187         args = append(args, disp.cluster.Containers.SLURM.SbatchArgumentsList...)
 188         args = append(args, "--job-name="+container.UUID, fmt.Sprintf("--nice=%d", initialNiceValue), "--no-requeue")
 189
 190         if disp.cluster == nil {
 191                 // no instance types configured
 192                 args = append(args, disp.slurmConstraintArgs(container)...)
 193         } else if it, err := dispatchcloud.ChooseInstanceType(disp.cluster, &container); err == dispatchcloud.ErrInstanceTypesNotConfigured {
 194                 // ditto
 195                 args = append(args, disp.slurmConstraintArgs(container)...)
 196         } else if err != nil {
 197                 return nil, err
 198         } else {
 199                 // use instancetype constraint instead of slurm mem/cpu/tmp specs
 200                 args = append(args, "--constraint=instancetype="+it.Name)
 201         }
 202
 203         if len(container.SchedulingParameters.Partitions) > 0 {
 204                 args = append(args, "--partition="+strings.Join(container.SchedulingParameters.Partitions, ","))
 205         }
 206
 207         return args, nil
 208 }
 209
 210 func (disp *Dispatcher) submit(container arvados.Container, crunchRunCommand []string) error {
 211         // append() here avoids modifying crunchRunCommand's
 212         // underlying array, which is shared with other goroutines.
 213         crArgs := append([]string(nil), crunchRunCommand...)
 214         crArgs = append(crArgs, "--runtime-engine="+disp.cluster.Containers.RuntimeEngine)
 215         crArgs = append(crArgs, container.UUID)
 216         crScript := strings.NewReader(execScript(crArgs))
 217
 218         sbArgs, err := disp.sbatchArgs(container)
 219         if err != nil {
 220                 return err
 221         }
 222         log.Printf("running sbatch %+q", sbArgs)
 223         return disp.slurm.Batch(crScript, sbArgs)
 224 }
 225
 226 // Submit a container to the slurm queue (or resume monitoring if it's
 227 // already in the queue).  Cancel the slurm job if the container's
 228 // priority changes to zero or its state indicates it's no longer
 229 // running.
 230 func (disp *Dispatcher) runContainer(_ *dispatch.Dispatcher, ctr arvados.Container, status <-chan arvados.Container) error {
 231         ctx, cancel := context.WithCancel(context.Background())
 232         defer cancel()
 233
 234         if ctr.State == dispatch.Locked && !disp.sqCheck.HasUUID(ctr.UUID) {
 235                 log.Printf("Submitting container %s to slurm", ctr.UUID)
 236                 cmd := []string{disp.cluster.Containers.CrunchRunCommand}
 237                 cmd = append(cmd, disp.cluster.Containers.CrunchRunArgumentsList...)
 238                 err := disp.submit(ctr, cmd)
 239                 if err != nil {
 240                         return err
 241                 }
 242         }
 243
 244         log.Printf("Start monitoring container %v in state %q", ctr.UUID, ctr.State)
 245         defer log.Printf("Done monitoring container %s", ctr.UUID)
 246
 247         // If the container disappears from the slurm queue, there is
 248         // no point in waiting for further dispatch updates: just
 249         // clean up and return.
 250         go func(uuid string) {
 251                 for ctx.Err() == nil && disp.sqCheck.HasUUID(uuid) {
 252                 }
 253                 cancel()
 254         }(ctr.UUID)
 255
 256         for {
 257                 select {
 258                 case <-ctx.Done():
 259                         // Disappeared from squeue
 260                         if err := disp.Arv.Get("containers", ctr.UUID, nil, &ctr); err != nil {
 261                                 log.Printf("error getting final container state for %s: %s", ctr.UUID, err)
 262                         }
 263                         switch ctr.State {
 264                         case dispatch.Running:
 265                                 disp.UpdateState(ctr.UUID, dispatch.Cancelled)
 266                         case dispatch.Locked:
 267                                 disp.Unlock(ctr.UUID)
 268                         }
 269                         return nil
 270                 case updated, ok := <-status:
 271                         if !ok {
 272                                 log.Printf("container %s is done: cancel slurm job", ctr.UUID)
 273                                 disp.scancel(ctr)
 274                         } else if updated.Priority == 0 {
 275                                 log.Printf("container %s has state %q, priority %d: cancel slurm job", ctr.UUID, updated.State, updated.Priority)
 276                                 disp.scancel(ctr)
 277                         } else {
 278                                 p := int64(updated.Priority)
 279                                 if p <= 1000 {
 280                                         // API is providing
 281                                         // user-assigned priority. If
 282                                         // ctrs have equal priority,
 283                                         // run the older one first.
 284                                         p = int64(p)<<50 - (updated.CreatedAt.UnixNano() >> 14)
 285                                 }
 286                                 disp.sqCheck.SetPriority(ctr.UUID, p)
 287                         }
 288                 }
 289         }
 290 }
 291 func (disp *Dispatcher) scancel(ctr arvados.Container) {
 292         err := disp.slurm.Cancel(ctr.UUID)
 293         if err != nil {
 294                 log.Printf("scancel: %s", err)
 295                 time.Sleep(time.Second)
 296         } else if disp.sqCheck.HasUUID(ctr.UUID) {
 297                 log.Printf("container %s is still in squeue after scancel", ctr.UUID)
 298                 time.Sleep(time.Second)
 299         }
 300 }