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