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