Arvados-DCO-1.1-Signed-off-by: Radhika Chippada <radhika@curoverse.com>
[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         "os/exec"
18         "regexp"
19         "strings"
20         "time"
21
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/coreos/go-systemd/daemon"
27 )
28
29 // Config used by crunch-dispatch-slurm
30 type Config struct {
31         Client arvados.Client
32
33         SbatchArguments []string
34         PollPeriod      arvados.Duration
35
36         // crunch-run command to invoke. The container UUID will be
37         // appended. If nil, []string{"crunch-run"} will be used.
38         //
39         // Example: []string{"crunch-run", "--cgroup-parent-subsystem=memory"}
40         CrunchRunCommand []string
41
42         // Minimum time between two attempts to run the same container
43         MinRetryPeriod arvados.Duration
44 }
45
46 func main() {
47         err := doMain()
48         if err != nil {
49                 log.Fatal(err)
50         }
51 }
52
53 var (
54         theConfig Config
55         sqCheck   = &SqueueChecker{}
56 )
57
58 const defaultConfigPath = "/etc/arvados/crunch-dispatch-slurm/crunch-dispatch-slurm.yml"
59
60 func doMain() error {
61         flags := flag.NewFlagSet("crunch-dispatch-slurm", flag.ExitOnError)
62         flags.Usage = func() { usage(flags) }
63
64         configPath := flags.String(
65                 "config",
66                 defaultConfigPath,
67                 "`path` to JSON or YAML configuration file")
68         dumpConfig := flag.Bool(
69                 "dump-config",
70                 false,
71                 "write current configuration to stdout and exit")
72
73         // Parse args; omit the first arg which is the command name
74         flags.Parse(os.Args[1:])
75
76         err := readConfig(&theConfig, *configPath)
77         if err != nil {
78                 return err
79         }
80
81         if theConfig.CrunchRunCommand == nil {
82                 theConfig.CrunchRunCommand = []string{"crunch-run"}
83         }
84
85         if theConfig.PollPeriod == 0 {
86                 theConfig.PollPeriod = arvados.Duration(10 * time.Second)
87         }
88
89         if theConfig.Client.APIHost != "" || theConfig.Client.AuthToken != "" {
90                 // Copy real configs into env vars so [a]
91                 // MakeArvadosClient() uses them, and [b] they get
92                 // propagated to crunch-run via SLURM.
93                 os.Setenv("ARVADOS_API_HOST", theConfig.Client.APIHost)
94                 os.Setenv("ARVADOS_API_TOKEN", theConfig.Client.AuthToken)
95                 os.Setenv("ARVADOS_API_HOST_INSECURE", "")
96                 if theConfig.Client.Insecure {
97                         os.Setenv("ARVADOS_API_HOST_INSECURE", "1")
98                 }
99                 os.Setenv("ARVADOS_KEEP_SERVICES", strings.Join(theConfig.Client.KeepServiceURIs, " "))
100                 os.Setenv("ARVADOS_EXTERNAL_CLIENT", "")
101         } else {
102                 log.Printf("warning: Client credentials missing from config, so falling back on environment variables (deprecated).")
103         }
104
105         if *dumpConfig {
106                 log.Fatal(config.DumpAndExit(theConfig))
107         }
108
109         arv, err := arvadosclient.MakeArvadosClient()
110         if err != nil {
111                 log.Printf("Error making Arvados client: %v", err)
112                 return err
113         }
114         arv.Retries = 25
115
116         sqCheck = &SqueueChecker{Period: time.Duration(theConfig.PollPeriod)}
117         defer sqCheck.Stop()
118
119         dispatcher := &dispatch.Dispatcher{
120                 Arv:            arv,
121                 RunContainer:   run,
122                 PollPeriod:     time.Duration(theConfig.PollPeriod),
123                 MinRetryPeriod: time.Duration(theConfig.MinRetryPeriod),
124         }
125
126         if _, err := daemon.SdNotify(false, "READY=1"); err != nil {
127                 log.Printf("Error notifying init daemon: %v", err)
128         }
129
130         go checkSqueueForOrphans(dispatcher, sqCheck)
131
132         return dispatcher.Run(context.Background())
133 }
134
135 var containerUuidPattern = regexp.MustCompile(`^[a-z0-9]{5}-dz642-[a-z0-9]{15}$`)
136
137 // Check the next squeue report, and invoke TrackContainer for all the
138 // containers in the report. This gives us a chance to cancel slurm
139 // jobs started by a previous dispatch process that never released
140 // their slurm allocations even though their container states are
141 // Cancelled or Complete. See https://dev.arvados.org/issues/10979
142 func checkSqueueForOrphans(dispatcher *dispatch.Dispatcher, sqCheck *SqueueChecker) {
143         for _, uuid := range sqCheck.All() {
144                 if !containerUuidPattern.MatchString(uuid) {
145                         continue
146                 }
147                 err := dispatcher.TrackContainer(uuid)
148                 if err != nil {
149                         log.Printf("checkSqueueForOrphans: TrackContainer(%s): %s", uuid, err)
150                 }
151         }
152 }
153
154 // sbatchCmd
155 func sbatchFunc(container arvados.Container) *exec.Cmd {
156         mem := int64(math.Ceil(float64(container.RuntimeConstraints.RAM+container.RuntimeConstraints.KeepCacheRAM) / float64(1048576)))
157
158         var disk int64
159         for _, m := range container.Mounts {
160                 if m.Kind == "tmp" {
161                         disk += m.Capacity
162                 }
163         }
164         disk = int64(math.Ceil(float64(disk) / float64(1048576)))
165
166         var sbatchArgs []string
167         sbatchArgs = append(sbatchArgs, theConfig.SbatchArguments...)
168         sbatchArgs = append(sbatchArgs, fmt.Sprintf("--job-name=%s", container.UUID))
169         sbatchArgs = append(sbatchArgs, fmt.Sprintf("--mem=%d", mem))
170         sbatchArgs = append(sbatchArgs, fmt.Sprintf("--cpus-per-task=%d", container.RuntimeConstraints.VCPUs))
171         sbatchArgs = append(sbatchArgs, fmt.Sprintf("--tmp=%d", disk))
172         if len(container.SchedulingParameters.Partitions) > 0 {
173                 sbatchArgs = append(sbatchArgs, fmt.Sprintf("--partition=%s", strings.Join(container.SchedulingParameters.Partitions, ",")))
174         }
175
176         return exec.Command("sbatch", sbatchArgs...)
177 }
178
179 // scancelCmd
180 func scancelFunc(container arvados.Container) *exec.Cmd {
181         return exec.Command("scancel", "--name="+container.UUID)
182 }
183
184 // Wrap these so that they can be overridden by tests
185 var sbatchCmd = sbatchFunc
186 var scancelCmd = scancelFunc
187
188 // Submit job to slurm using sbatch.
189 func submit(dispatcher *dispatch.Dispatcher, container arvados.Container, crunchRunCommand []string) error {
190         cmd := sbatchCmd(container)
191
192         // Send a tiny script on stdin to execute the crunch-run
193         // command (slurm requires this to be a #! script)
194         cmd.Stdin = strings.NewReader(execScript(append(crunchRunCommand, container.UUID)))
195
196         var stdout, stderr bytes.Buffer
197         cmd.Stdout = &stdout
198         cmd.Stderr = &stderr
199
200         // Mutex between squeue sync and running sbatch or scancel.
201         sqCheck.L.Lock()
202         defer sqCheck.L.Unlock()
203
204         log.Printf("exec sbatch %+q", cmd.Args)
205         err := cmd.Run()
206
207         switch err.(type) {
208         case nil:
209                 log.Printf("sbatch succeeded: %q", strings.TrimSpace(stdout.String()))
210                 return nil
211
212         case *exec.ExitError:
213                 dispatcher.Unlock(container.UUID)
214                 return fmt.Errorf("sbatch %+q failed: %v (stderr: %q)", cmd.Args, err, stderr.Bytes())
215
216         default:
217                 dispatcher.Unlock(container.UUID)
218                 return fmt.Errorf("exec failed: %v", err)
219         }
220 }
221
222 // Submit a container to the slurm queue (or resume monitoring if it's
223 // already in the queue).  Cancel the slurm job if the container's
224 // priority changes to zero or its state indicates it's no longer
225 // running.
226 func run(disp *dispatch.Dispatcher, ctr arvados.Container, status <-chan arvados.Container) {
227         ctx, cancel := context.WithCancel(context.Background())
228         defer cancel()
229
230         if ctr.State == dispatch.Locked && !sqCheck.HasUUID(ctr.UUID) {
231                 log.Printf("Submitting container %s to slurm", ctr.UUID)
232                 if err := submit(disp, ctr, theConfig.CrunchRunCommand); err != nil {
233                         text := fmt.Sprintf("Error submitting container %s to slurm: %s", ctr.UUID, err)
234                         log.Print(text)
235
236                         lr := arvadosclient.Dict{"log": arvadosclient.Dict{
237                                 "object_uuid": ctr.UUID,
238                                 "event_type":  "dispatch",
239                                 "properties":  map[string]string{"text": text}}}
240                         disp.Arv.Create("logs", lr, nil)
241
242                         disp.Unlock(ctr.UUID)
243                         return
244                 }
245         }
246
247         log.Printf("Start monitoring container %v in state %q", ctr.UUID, ctr.State)
248         defer log.Printf("Done monitoring container %s", ctr.UUID)
249
250         // If the container disappears from the slurm queue, there is
251         // no point in waiting for further dispatch updates: just
252         // clean up and return.
253         go func(uuid string) {
254                 for ctx.Err() == nil && sqCheck.HasUUID(uuid) {
255                 }
256                 cancel()
257         }(ctr.UUID)
258
259         for {
260                 select {
261                 case <-ctx.Done():
262                         // Disappeared from squeue
263                         if err := disp.Arv.Get("containers", ctr.UUID, nil, &ctr); err != nil {
264                                 log.Printf("Error getting final container state for %s: %s", ctr.UUID, err)
265                         }
266                         switch ctr.State {
267                         case dispatch.Running:
268                                 disp.UpdateState(ctr.UUID, dispatch.Cancelled)
269                         case dispatch.Locked:
270                                 disp.Unlock(ctr.UUID)
271                         }
272                         return
273                 case updated, ok := <-status:
274                         if !ok {
275                                 log.Printf("Dispatcher says container %s is done: cancel slurm job", ctr.UUID)
276                                 scancel(ctr)
277                         } else if updated.Priority == 0 {
278                                 log.Printf("Container %s has state %q, priority %d: cancel slurm job", ctr.UUID, updated.State, updated.Priority)
279                                 scancel(ctr)
280                         }
281                 }
282         }
283 }
284
285 func scancel(ctr arvados.Container) {
286         sqCheck.L.Lock()
287         cmd := scancelCmd(ctr)
288         msg, err := cmd.CombinedOutput()
289         sqCheck.L.Unlock()
290
291         if err != nil {
292                 log.Printf("%q %q: %s %q", cmd.Path, cmd.Args, err, msg)
293                 time.Sleep(time.Second)
294         } else if sqCheck.HasUUID(ctr.UUID) {
295                 log.Printf("container %s is still in squeue after scancel", ctr.UUID)
296                 time.Sleep(time.Second)
297         }
298 }
299
300 func readConfig(dst interface{}, path string) error {
301         err := config.LoadFile(dst, path)
302         if err != nil && os.IsNotExist(err) && path == defaultConfigPath {
303                 log.Printf("Config not specified. Continue with default configuration.")
304                 err = nil
305         }
306         return err
307 }