1 // Copyright (C) The Arvados Authors. All rights reserved.
3 // SPDX-License-Identifier: AGPL-3.0
7 // Dispatcher service for Crunch that submits containers to the slurm queue.
22 "git.curoverse.com/arvados.git/lib/config"
23 "git.curoverse.com/arvados.git/lib/dispatchcloud"
24 "git.curoverse.com/arvados.git/sdk/go/arvados"
25 "git.curoverse.com/arvados.git/sdk/go/arvadosclient"
26 "git.curoverse.com/arvados.git/sdk/go/dispatch"
27 "github.com/coreos/go-systemd/daemon"
28 "github.com/ghodss/yaml"
29 "github.com/sirupsen/logrus"
32 type logger interface {
34 Fatalf(string, ...interface{})
37 const initialNiceValue int64 = 10000
43 type Dispatcher struct {
45 logger logrus.FieldLogger
46 cluster *arvados.Cluster
47 sqCheck *SqueueChecker
54 logger := logrus.StandardLogger()
55 if os.Getenv("DEBUG") != "" {
56 logger.SetLevel(logrus.DebugLevel)
58 logger.Formatter = &logrus.JSONFormatter{
59 TimestampFormat: "2006-01-02T15:04:05.000000000Z07:00",
61 disp := &Dispatcher{logger: logger}
62 err := disp.Run(os.Args[0], os.Args[1:])
64 logrus.Fatalf("%s", err)
68 func (disp *Dispatcher) Run(prog string, args []string) error {
69 if err := disp.configure(prog, args); err != nil {
76 // configure() loads config files. Tests skip this.
77 func (disp *Dispatcher) configure(prog string, args []string) error {
78 if disp.logger == nil {
79 disp.logger = logrus.StandardLogger()
81 flags := flag.NewFlagSet(prog, flag.ExitOnError)
82 flags.Usage = func() { usage(flags) }
84 loader := config.NewLoader(nil, disp.logger)
85 loader.SetupFlags(flags)
87 dumpConfig := flag.Bool(
90 "write current configuration to stdout and exit")
91 getVersion := flags.Bool(
94 "Print version information and exit.")
96 args = loader.MungeLegacyConfigArgs(logrus.StandardLogger(), args, "-legacy-crunch-dispatch-slurm-config")
98 // Parse args; omit the first arg which is the command name
99 err := flags.Parse(args)
101 if err == flag.ErrHelp {
105 // Print version information if requested
107 fmt.Printf("crunch-dispatch-slurm %s\n", version)
111 disp.logger.Printf("crunch-dispatch-slurm %s started", version)
113 cfg, err := loader.Load()
118 if disp.cluster, err = cfg.GetCluster(""); err != nil {
119 return fmt.Errorf("config error: %s", err)
122 disp.Client.APIHost = disp.cluster.Services.Controller.ExternalURL.Host
123 disp.Client.AuthToken = disp.cluster.SystemRootToken
124 disp.Client.Insecure = disp.cluster.TLS.Insecure
126 if disp.Client.APIHost != "" || disp.Client.AuthToken != "" {
127 // Copy real configs into env vars so [a]
128 // MakeArvadosClient() uses them, and [b] they get
129 // propagated to crunch-run via SLURM.
130 os.Setenv("ARVADOS_API_HOST", disp.Client.APIHost)
131 os.Setenv("ARVADOS_API_TOKEN", disp.Client.AuthToken)
132 os.Setenv("ARVADOS_API_HOST_INSECURE", "")
133 if disp.Client.Insecure {
134 os.Setenv("ARVADOS_API_HOST_INSECURE", "1")
137 if len(disp.cluster.Containers.SLURM.KeepServices) > 0 {
138 for _, svc := range disp.cluster.Containers.SLURM.KeepServices {
139 for k, _ := range svc.InternalURLs {
146 os.Setenv("ARVADOS_KEEP_SERVICES", strings.TrimSuffix(ks, " "))
147 os.Setenv("ARVADOS_EXTERNAL_CLIENT", "")
149 disp.logger.Warnf("Client credentials missing from config, so falling back on environment variables (deprecated).")
153 out, err := yaml.Marshal(cfg)
157 _, err = os.Stdout.Write(out)
166 // setup() initializes private fields after configure().
167 func (disp *Dispatcher) setup() {
168 arv, err := arvadosclient.MakeArvadosClient()
170 disp.logger.Fatalf("Error making Arvados client: %v", err)
174 disp.slurm = NewSlurmCLI()
175 disp.sqCheck = &SqueueChecker{
177 Period: time.Duration(disp.cluster.Containers.CloudVMs.PollInterval),
178 PrioritySpread: disp.cluster.Containers.SLURM.PrioritySpread,
181 disp.Dispatcher = &dispatch.Dispatcher{
184 BatchSize: disp.cluster.API.MaxItemsPerResponse,
185 RunContainer: disp.runContainer,
186 PollPeriod: time.Duration(disp.cluster.Containers.CloudVMs.PollInterval),
187 MinRetryPeriod: time.Duration(disp.cluster.Containers.MinRetryPeriod),
191 func (disp *Dispatcher) run() error {
192 defer disp.sqCheck.Stop()
194 if disp.cluster != nil && len(disp.cluster.InstanceTypes) > 0 {
195 go SlurmNodeTypeFeatureKludge(disp.cluster)
198 if _, err := daemon.SdNotify(false, "READY=1"); err != nil {
199 log.Printf("Error notifying init daemon: %v", err)
201 go disp.checkSqueueForOrphans()
202 return disp.Dispatcher.Run(context.Background())
205 var containerUuidPattern = regexp.MustCompile(`^[a-z0-9]{5}-dz642-[a-z0-9]{15}$`)
207 // Check the next squeue report, and invoke TrackContainer for all the
208 // containers in the report. This gives us a chance to cancel slurm
209 // jobs started by a previous dispatch process that never released
210 // their slurm allocations even though their container states are
211 // Cancelled or Complete. See https://dev.arvados.org/issues/10979
212 func (disp *Dispatcher) checkSqueueForOrphans() {
213 for _, uuid := range disp.sqCheck.All() {
214 if !containerUuidPattern.MatchString(uuid) {
217 err := disp.TrackContainer(uuid)
219 log.Printf("checkSqueueForOrphans: TrackContainer(%s): %s", uuid, err)
224 func (disp *Dispatcher) slurmConstraintArgs(container arvados.Container) []string {
225 mem := int64(math.Ceil(float64(container.RuntimeConstraints.RAM+
226 container.RuntimeConstraints.KeepCacheRAM+
227 int64(disp.cluster.Containers.ReserveExtraRAM)) / float64(1048576)))
229 disk := dispatchcloud.EstimateScratchSpace(&container)
230 disk = int64(math.Ceil(float64(disk) / float64(1048576)))
232 fmt.Sprintf("--mem=%d", mem),
233 fmt.Sprintf("--cpus-per-task=%d", container.RuntimeConstraints.VCPUs),
234 fmt.Sprintf("--tmp=%d", disk),
238 func (disp *Dispatcher) sbatchArgs(container arvados.Container) ([]string, error) {
240 args = append(args, disp.cluster.Containers.SLURM.SbatchArgumentsList...)
241 args = append(args, "--job-name="+container.UUID, fmt.Sprintf("--nice=%d", initialNiceValue), "--no-requeue")
243 if disp.cluster == nil {
244 // no instance types configured
245 args = append(args, disp.slurmConstraintArgs(container)...)
246 } else if it, err := dispatchcloud.ChooseInstanceType(disp.cluster, &container); err == dispatchcloud.ErrInstanceTypesNotConfigured {
248 args = append(args, disp.slurmConstraintArgs(container)...)
249 } else if err != nil {
252 // use instancetype constraint instead of slurm mem/cpu/tmp specs
253 args = append(args, "--constraint=instancetype="+it.Name)
256 if len(container.SchedulingParameters.Partitions) > 0 {
257 args = append(args, "--partition="+strings.Join(container.SchedulingParameters.Partitions, ","))
263 func (disp *Dispatcher) submit(container arvados.Container, crunchRunCommand []string) error {
264 // append() here avoids modifying crunchRunCommand's
265 // underlying array, which is shared with other goroutines.
266 crArgs := append([]string(nil), crunchRunCommand...)
267 crArgs = append(crArgs, container.UUID)
268 crScript := strings.NewReader(execScript(crArgs))
270 sbArgs, err := disp.sbatchArgs(container)
274 log.Printf("running sbatch %+q", sbArgs)
275 return disp.slurm.Batch(crScript, sbArgs)
278 // Submit a container to the slurm queue (or resume monitoring if it's
279 // already in the queue). Cancel the slurm job if the container's
280 // priority changes to zero or its state indicates it's no longer
282 func (disp *Dispatcher) runContainer(_ *dispatch.Dispatcher, ctr arvados.Container, status <-chan arvados.Container) {
283 ctx, cancel := context.WithCancel(context.Background())
286 if ctr.State == dispatch.Locked && !disp.sqCheck.HasUUID(ctr.UUID) {
287 log.Printf("Submitting container %s to slurm", ctr.UUID)
288 cmd := []string{disp.cluster.Containers.CrunchRunCommand}
289 cmd = append(cmd, disp.cluster.Containers.CrunchRunArgumentsList...)
290 if err := disp.submit(ctr, cmd); err != nil {
292 if err, ok := err.(dispatchcloud.ConstraintsNotSatisfiableError); ok {
293 var logBuf bytes.Buffer
294 fmt.Fprintf(&logBuf, "cannot run container %s: %s\n", ctr.UUID, err)
295 if len(err.AvailableTypes) == 0 {
296 fmt.Fprint(&logBuf, "No instance types are configured.\n")
298 fmt.Fprint(&logBuf, "Available instance types:\n")
299 for _, t := range err.AvailableTypes {
301 "Type %q: %d VCPUs, %d RAM, %d Scratch, %f Price\n",
302 t.Name, t.VCPUs, t.RAM, t.Scratch, t.Price,
306 text = logBuf.String()
307 disp.UpdateState(ctr.UUID, dispatch.Cancelled)
309 text = fmt.Sprintf("Error submitting container %s to slurm: %s", ctr.UUID, err)
313 lr := arvadosclient.Dict{"log": arvadosclient.Dict{
314 "object_uuid": ctr.UUID,
315 "event_type": "dispatch",
316 "properties": map[string]string{"text": text}}}
317 disp.Arv.Create("logs", lr, nil)
319 disp.Unlock(ctr.UUID)
324 log.Printf("Start monitoring container %v in state %q", ctr.UUID, ctr.State)
325 defer log.Printf("Done monitoring container %s", ctr.UUID)
327 // If the container disappears from the slurm queue, there is
328 // no point in waiting for further dispatch updates: just
329 // clean up and return.
330 go func(uuid string) {
331 for ctx.Err() == nil && disp.sqCheck.HasUUID(uuid) {
339 // Disappeared from squeue
340 if err := disp.Arv.Get("containers", ctr.UUID, nil, &ctr); err != nil {
341 log.Printf("error getting final container state for %s: %s", ctr.UUID, err)
344 case dispatch.Running:
345 disp.UpdateState(ctr.UUID, dispatch.Cancelled)
346 case dispatch.Locked:
347 disp.Unlock(ctr.UUID)
350 case updated, ok := <-status:
352 log.Printf("container %s is done: cancel slurm job", ctr.UUID)
354 } else if updated.Priority == 0 {
355 log.Printf("container %s has state %q, priority %d: cancel slurm job", ctr.UUID, updated.State, updated.Priority)
358 p := int64(updated.Priority)
361 // user-assigned priority. If
362 // ctrs have equal priority,
363 // run the older one first.
364 p = int64(p)<<50 - (updated.CreatedAt.UnixNano() >> 14)
366 disp.sqCheck.SetPriority(ctr.UUID, p)
371 func (disp *Dispatcher) scancel(ctr arvados.Container) {
372 err := disp.slurm.Cancel(ctr.UUID)
374 log.Printf("scancel: %s", err)
375 time.Sleep(time.Second)
376 } else if disp.sqCheck.HasUUID(ctr.UUID) {
377 log.Printf("container %s is still in squeue after scancel", ctr.UUID)
378 time.Sleep(time.Second)