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 package main
   6
   7 // Dispatcher service for Crunch that submits containers to the slurm queue.
   8
   9 import (
  10         "context"
  11         "flag"
  12         "fmt"
  13         "log"
  14         "math"
  15         "os"
  16         "regexp"
  17         "strings"
  18         "time"
  19
  20         "git.curoverse.com/arvados.git/lib/dispatchcloud"
  21         "git.curoverse.com/arvados.git/sdk/go/arvados"
  22         "git.curoverse.com/arvados.git/sdk/go/arvadosclient"
  23         "git.curoverse.com/arvados.git/sdk/go/config"
  24         "git.curoverse.com/arvados.git/sdk/go/dispatch"
  25         "github.com/coreos/go-systemd/daemon"
  26 )
  27
  28 const initialNiceValue int64 = 10000
  29
  30 var (
  31         version           = "dev"
  32         defaultConfigPath = "/etc/arvados/crunch-dispatch-slurm/crunch-dispatch-slurm.yml"
  33 )
  34
  35 type Dispatcher struct {
  36         *dispatch.Dispatcher
  37         cluster *arvados.Cluster
  38         sqCheck *SqueueChecker
  39         slurm   Slurm
  40
  41         Client arvados.Client
  42
  43         SbatchArguments []string
  44         PollPeriod      arvados.Duration
  45         PrioritySpread  int64
  46
  47         // crunch-run command to invoke. The container UUID will be
  48         // appended. If nil, []string{"crunch-run"} will be used.
  49         //
  50         // Example: []string{"crunch-run", "--cgroup-parent-subsystem=memory"}
  51         CrunchRunCommand []string
  52
  53         // Minimum time between two attempts to run the same container
  54         MinRetryPeriod arvados.Duration
  55 }
  56
  57 func main() {
  58         disp := &Dispatcher{}
  59         err := disp.Run(os.Args[0], os.Args[1:])
  60         if err != nil {
  61                 log.Fatal(err)
  62         }
  63 }
  64
  65 func (disp *Dispatcher) Run(prog string, args []string) error {
  66         if err := disp.configure(prog, args); err != nil {
  67                 return err
  68         }
  69         disp.setup()
  70         return disp.run()
  71 }
  72
  73 // configure() loads config files. Tests skip this.
  74 func (disp *Dispatcher) configure(prog string, args []string) error {
  75         flags := flag.NewFlagSet(prog, flag.ExitOnError)
  76         flags.Usage = func() { usage(flags) }
  77
  78         configPath := flags.String(
  79                 "config",
  80                 defaultConfigPath,
  81                 "`path` to JSON or YAML configuration file")
  82         dumpConfig := flag.Bool(
  83                 "dump-config",
  84                 false,
  85                 "write current configuration to stdout and exit")
  86         getVersion := flags.Bool(
  87                 "version",
  88                 false,
  89                 "Print version information and exit.")
  90         // Parse args; omit the first arg which is the command name
  91         flags.Parse(args)
  92
  93         // Print version information if requested
  94         if *getVersion {
  95                 fmt.Printf("crunch-dispatch-slurm %s\n", version)
  96                 return nil
  97         }
  98
  99         log.Printf("crunch-dispatch-slurm %s started", version)
 100
 101         err := disp.readConfig(*configPath)
 102         if err != nil {
 103                 return err
 104         }
 105
 106         if disp.CrunchRunCommand == nil {
 107                 disp.CrunchRunCommand = []string{"crunch-run"}
 108         }
 109
 110         if disp.PollPeriod == 0 {
 111                 disp.PollPeriod = arvados.Duration(10 * time.Second)
 112         }
 113
 114         if disp.Client.APIHost != "" || disp.Client.AuthToken != "" {
 115                 // Copy real configs into env vars so [a]
 116                 // MakeArvadosClient() uses them, and [b] they get
 117                 // propagated to crunch-run via SLURM.
 118                 os.Setenv("ARVADOS_API_HOST", disp.Client.APIHost)
 119                 os.Setenv("ARVADOS_API_TOKEN", disp.Client.AuthToken)
 120                 os.Setenv("ARVADOS_API_HOST_INSECURE", "")
 121                 if disp.Client.Insecure {
 122                         os.Setenv("ARVADOS_API_HOST_INSECURE", "1")
 123                 }
 124                 os.Setenv("ARVADOS_KEEP_SERVICES", strings.Join(disp.Client.KeepServiceURIs, " "))
 125                 os.Setenv("ARVADOS_EXTERNAL_CLIENT", "")
 126         } else {
 127                 log.Printf("warning: Client credentials missing from config, so falling back on environment variables (deprecated).")
 128         }
 129
 130         if *dumpConfig {
 131                 return config.DumpAndExit(disp)
 132         }
 133
 134         siteConfig, err := arvados.GetConfig(arvados.DefaultConfigFile)
 135         if os.IsNotExist(err) {
 136                 log.Printf("warning: no cluster config (%s), proceeding with no node types defined", err)
 137         } else if err != nil {
 138                 return fmt.Errorf("error loading config: %s", err)
 139         } else if disp.cluster, err = siteConfig.GetCluster(""); err != nil {
 140                 return fmt.Errorf("config error: %s", err)
 141         }
 142
 143         return nil
 144 }
 145
 146 // setup() initializes private fields after configure().
 147 func (disp *Dispatcher) setup() {
 148         arv, err := arvadosclient.MakeArvadosClient()
 149         if err != nil {
 150                 log.Fatalf("Error making Arvados client: %v", err)
 151         }
 152         arv.Retries = 25
 153
 154         disp.slurm = &slurmCLI{}
 155         disp.sqCheck = &SqueueChecker{
 156                 Period:         time.Duration(disp.PollPeriod),
 157                 PrioritySpread: disp.PrioritySpread,
 158                 Slurm:          disp.slurm,
 159         }
 160         disp.Dispatcher = &dispatch.Dispatcher{
 161                 Arv:            arv,
 162                 RunContainer:   disp.runContainer,
 163                 PollPeriod:     time.Duration(disp.PollPeriod),
 164                 MinRetryPeriod: time.Duration(disp.MinRetryPeriod),
 165         }
 166 }
 167
 168 func (disp *Dispatcher) run() error {
 169         defer disp.sqCheck.Stop()
 170
 171         if disp.cluster != nil && len(disp.cluster.InstanceTypes) > 0 {
 172                 go dispatchcloud.SlurmNodeTypeFeatureKludge(disp.cluster)
 173         }
 174
 175         if _, err := daemon.SdNotify(false, "READY=1"); err != nil {
 176                 log.Printf("Error notifying init daemon: %v", err)
 177         }
 178         go disp.checkSqueueForOrphans()
 179         return disp.Dispatcher.Run(context.Background())
 180 }
 181
 182 var containerUuidPattern = regexp.MustCompile(`^[a-z0-9]{5}-dz642-[a-z0-9]{15}$`)
 183
 184 // Check the next squeue report, and invoke TrackContainer for all the
 185 // containers in the report. This gives us a chance to cancel slurm
 186 // jobs started by a previous dispatch process that never released
 187 // their slurm allocations even though their container states are
 188 // Cancelled or Complete. See https://dev.arvados.org/issues/10979
 189 func (disp *Dispatcher) checkSqueueForOrphans() {
 190         for _, uuid := range disp.sqCheck.All() {
 191                 if !containerUuidPattern.MatchString(uuid) {
 192                         continue
 193                 }
 194                 err := disp.TrackContainer(uuid)
 195                 if err != nil {
 196                         log.Printf("checkSqueueForOrphans: TrackContainer(%s): %s", uuid, err)
 197                 }
 198         }
 199 }
 200
 201 func (disp *Dispatcher) sbatchArgs(container arvados.Container) ([]string, error) {
 202         mem := int64(math.Ceil(float64(container.RuntimeConstraints.RAM+container.RuntimeConstraints.KeepCacheRAM) / float64(1048576)))
 203
 204         var disk int64
 205         for _, m := range container.Mounts {
 206                 if m.Kind == "tmp" {
 207                         disk += m.Capacity
 208                 }
 209         }
 210         disk = int64(math.Ceil(float64(disk) / float64(1048576)))
 211
 212         var sbatchArgs []string
 213         sbatchArgs = append(sbatchArgs, disp.SbatchArguments...)
 214         sbatchArgs = append(sbatchArgs, fmt.Sprintf("--job-name=%s", container.UUID))
 215         sbatchArgs = append(sbatchArgs, fmt.Sprintf("--mem=%d", mem))
 216         sbatchArgs = append(sbatchArgs, fmt.Sprintf("--cpus-per-task=%d", container.RuntimeConstraints.VCPUs))
 217         sbatchArgs = append(sbatchArgs, fmt.Sprintf("--tmp=%d", disk))
 218         sbatchArgs = append(sbatchArgs, fmt.Sprintf("--nice=%d", initialNiceValue))
 219         if len(container.SchedulingParameters.Partitions) > 0 {
 220                 sbatchArgs = append(sbatchArgs, fmt.Sprintf("--partition=%s", strings.Join(container.SchedulingParameters.Partitions, ",")))
 221         }
 222
 223         if disp.cluster == nil {
 224                 // no instance types configured
 225         } else if it, err := dispatchcloud.ChooseInstanceType(disp.cluster, &container); err == dispatchcloud.ErrInstanceTypesNotConfigured {
 226                 // ditto
 227         } else if err != nil {
 228                 return nil, err
 229         } else {
 230                 sbatchArgs = append(sbatchArgs, "--constraint=instancetype="+it.Name)
 231         }
 232
 233         return sbatchArgs, nil
 234 }
 235
 236 func (disp *Dispatcher) submit(container arvados.Container, crunchRunCommand []string) error {
 237         // append() here avoids modifying crunchRunCommand's
 238         // underlying array, which is shared with other goroutines.
 239         crArgs := append([]string(nil), crunchRunCommand...)
 240         crArgs = append(crArgs, container.UUID)
 241         crScript := strings.NewReader(execScript(crArgs))
 242
 243         disp.sqCheck.L.Lock()
 244         defer disp.sqCheck.L.Unlock()
 245
 246         sbArgs, err := disp.sbatchArgs(container)
 247         if err != nil {
 248                 return err
 249         }
 250         log.Printf("running sbatch %+q", sbArgs)
 251         return disp.slurm.Batch(crScript, sbArgs)
 252 }
 253
 254 // Submit a container to the slurm queue (or resume monitoring if it's
 255 // already in the queue).  Cancel the slurm job if the container's
 256 // priority changes to zero or its state indicates it's no longer
 257 // running.
 258 func (disp *Dispatcher) runContainer(_ *dispatch.Dispatcher, ctr arvados.Container, status <-chan arvados.Container) {
 259         ctx, cancel := context.WithCancel(context.Background())
 260         defer cancel()
 261
 262         if ctr.State == dispatch.Locked && !disp.sqCheck.HasUUID(ctr.UUID) {
 263                 log.Printf("Submitting container %s to slurm", ctr.UUID)
 264                 if err := disp.submit(ctr, disp.CrunchRunCommand); err != nil {
 265                         var text string
 266                         if err == dispatchcloud.ErrConstraintsNotSatisfiable {
 267                                 text = fmt.Sprintf("cannot run container %s: %s", ctr.UUID, err)
 268                                 disp.UpdateState(ctr.UUID, dispatch.Cancelled)
 269                         } else {
 270                                 text = fmt.Sprintf("Error submitting container %s to slurm: %s", ctr.UUID, err)
 271                         }
 272                         log.Print(text)
 273
 274                         lr := arvadosclient.Dict{"log": arvadosclient.Dict{
 275                                 "object_uuid": ctr.UUID,
 276                                 "event_type":  "dispatch",
 277                                 "properties":  map[string]string{"text": text}}}
 278                         disp.Arv.Create("logs", lr, nil)
 279
 280                         disp.Unlock(ctr.UUID)
 281                         return
 282                 }
 283         }
 284
 285         log.Printf("Start monitoring container %v in state %q", ctr.UUID, ctr.State)
 286         defer log.Printf("Done monitoring container %s", ctr.UUID)
 287
 288         // If the container disappears from the slurm queue, there is
 289         // no point in waiting for further dispatch updates: just
 290         // clean up and return.
 291         go func(uuid string) {
 292                 for ctx.Err() == nil && disp.sqCheck.HasUUID(uuid) {
 293                 }
 294                 cancel()
 295         }(ctr.UUID)
 296
 297         for {
 298                 select {
 299                 case <-ctx.Done():
 300                         // Disappeared from squeue
 301                         if err := disp.Arv.Get("containers", ctr.UUID, nil, &ctr); err != nil {
 302                                 log.Printf("Error getting final container state for %s: %s", ctr.UUID, err)
 303                         }
 304                         switch ctr.State {
 305                         case dispatch.Running:
 306                                 disp.UpdateState(ctr.UUID, dispatch.Cancelled)
 307                         case dispatch.Locked:
 308                                 disp.Unlock(ctr.UUID)
 309                         }
 310                         return
 311                 case updated, ok := <-status:
 312                         if !ok {
 313                                 log.Printf("container %s is done: cancel slurm job", ctr.UUID)
 314                                 disp.scancel(ctr)
 315                         } else if updated.Priority == 0 {
 316                                 log.Printf("container %s has state %q, priority %d: cancel slurm job", ctr.UUID, updated.State, updated.Priority)
 317                                 disp.scancel(ctr)
 318                         } else {
 319                                 p := int64(updated.Priority)
 320                                 if p <= 1000 {
 321                                         // API is providing
 322                                         // user-assigned priority. If
 323                                         // ctrs have equal priority,
 324                                         // run the older one first.
 325                                         p = int64(p)<<50 - (updated.CreatedAt.UnixNano() >> 14)
 326                                 }
 327                                 disp.sqCheck.SetPriority(ctr.UUID, p)
 328                         }
 329                 }
 330         }
 331 }
 332 func (disp *Dispatcher) scancel(ctr arvados.Container) {
 333         disp.sqCheck.L.Lock()
 334         err := disp.slurm.Cancel(ctr.UUID)
 335         disp.sqCheck.L.Unlock()
 336
 337         if err != nil {
 338                 log.Printf("scancel: %s", err)
 339                 time.Sleep(time.Second)
 340         } else if disp.sqCheck.HasUUID(ctr.UUID) {
 341                 log.Printf("container %s is still in squeue after scancel", ctr.UUID)
 342                 time.Sleep(time.Second)
 343         }
 344 }
 345
 346 func (disp *Dispatcher) readConfig(path string) error {
 347         err := config.LoadFile(disp, path)
 348         if err != nil && os.IsNotExist(err) && path == defaultConfigPath {
 349                 log.Printf("Config not specified. Continue with default configuration.")
 350                 err = nil
 351         }
 352         return err
 353 }