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.
21 "git.curoverse.com/arvados.git/lib/config"
22 "git.curoverse.com/arvados.git/lib/dispatchcloud"
23 "git.curoverse.com/arvados.git/sdk/go/arvados"
24 "git.curoverse.com/arvados.git/sdk/go/arvadosclient"
25 "git.curoverse.com/arvados.git/sdk/go/dispatch"
26 "github.com/coreos/go-systemd/daemon"
27 "github.com/ghodss/yaml"
28 "github.com/sirupsen/logrus"
31 type logger interface {
33 Fatalf(string, ...interface{})
36 const initialNiceValue int64 = 10000
42 type Dispatcher struct {
44 logger logrus.FieldLogger
45 cluster *arvados.Cluster
46 sqCheck *SqueueChecker
53 logger := logrus.StandardLogger()
54 if os.Getenv("DEBUG") != "" {
55 logger.SetLevel(logrus.DebugLevel)
57 logger.Formatter = &logrus.JSONFormatter{
58 TimestampFormat: "2006-01-02T15:04:05.000000000Z07:00",
60 disp := &Dispatcher{logger: logger}
61 err := disp.Run(os.Args[0], os.Args[1:])
63 logrus.Fatalf("%s", err)
67 func (disp *Dispatcher) Run(prog string, args []string) error {
68 if err := disp.configure(prog, args); err != nil {
75 // configure() loads config files. Tests skip this.
76 func (disp *Dispatcher) configure(prog string, args []string) error {
77 if disp.logger == nil {
78 disp.logger = logrus.StandardLogger()
80 flags := flag.NewFlagSet(prog, flag.ExitOnError)
81 flags.Usage = func() { usage(flags) }
83 loader := config.NewLoader(nil, disp.logger)
84 loader.SetupFlags(flags)
86 dumpConfig := flag.Bool(
89 "write current configuration to stdout and exit")
90 getVersion := flags.Bool(
93 "Print version information and exit.")
95 args = loader.MungeLegacyConfigArgs(logrus.StandardLogger(), args, "-legacy-crunch-dispatch-slurm-config")
97 // Parse args; omit the first arg which is the command name
98 err := flags.Parse(args)
100 if err == flag.ErrHelp {
104 // Print version information if requested
106 fmt.Printf("crunch-dispatch-slurm %s\n", version)
110 disp.logger.Printf("crunch-dispatch-slurm %s started", version)
112 cfg, err := loader.Load()
117 if disp.cluster, err = cfg.GetCluster(""); err != nil {
118 return fmt.Errorf("config error: %s", err)
121 disp.Client.APIHost = disp.cluster.Services.Controller.ExternalURL.Host
122 disp.Client.AuthToken = disp.cluster.SystemRootToken
123 disp.Client.Insecure = disp.cluster.TLS.Insecure
125 if disp.Client.APIHost != "" || disp.Client.AuthToken != "" {
126 // Copy real configs into env vars so [a]
127 // MakeArvadosClient() uses them, and [b] they get
128 // propagated to crunch-run via SLURM.
129 os.Setenv("ARVADOS_API_HOST", disp.Client.APIHost)
130 os.Setenv("ARVADOS_API_TOKEN", disp.Client.AuthToken)
131 os.Setenv("ARVADOS_API_HOST_INSECURE", "")
132 if disp.Client.Insecure {
133 os.Setenv("ARVADOS_API_HOST_INSECURE", "1")
136 if length(disp.cluster.Containers.SLURM.KeepServices) > 0 {
137 for _, svc := range disp.cluster.Containers.SLURM.KeepServices {
138 for k, _ := range svc.InternalURLs {
144 os.Setenv("ARVADOS_KEEP_SERVICES", ks)
145 os.Setenv("ARVADOS_EXTERNAL_CLIENT", "")
147 disp.logger.Warnf("Client credentials missing from config, so falling back on environment variables (deprecated).")
151 out, err := yaml.Marshal(cfg)
155 _, err = os.Stdout.Write(out)
164 // setup() initializes private fields after configure().
165 func (disp *Dispatcher) setup() {
166 arv, err := arvadosclient.MakeArvadosClient()
168 disp.logger.Fatalf("Error making Arvados client: %v", err)
172 disp.slurm = NewSlurmCLI()
173 disp.sqCheck = &SqueueChecker{
175 Period: time.Duration(disp.cluster.Containers.CloudVMs.PollInterval),
176 PrioritySpread: disp.cluster.Containers.SLURM.PrioritySpread,
179 disp.Dispatcher = &dispatch.Dispatcher{
182 BatchSize: disp.cluster.API.MaxItemsPerResponse,
183 RunContainer: disp.runContainer,
184 PollPeriod: time.Duration(disp.cluster.Containers.CloudVMs.PollInterval),
185 MinRetryPeriod: time.Duration(disp.cluster.Containers.MinRetryPeriod),
189 func (disp *Dispatcher) run() error {
190 defer disp.sqCheck.Stop()
192 if disp.cluster != nil && len(disp.cluster.InstanceTypes) > 0 {
193 go SlurmNodeTypeFeatureKludge(disp.cluster)
196 if _, err := daemon.SdNotify(false, "READY=1"); err != nil {
197 log.Printf("Error notifying init daemon: %v", err)
199 go disp.checkSqueueForOrphans()
200 return disp.Dispatcher.Run(context.Background())
203 var containerUuidPattern = regexp.MustCompile(`^[a-z0-9]{5}-dz642-[a-z0-9]{15}$`)
205 // Check the next squeue report, and invoke TrackContainer for all the
206 // containers in the report. This gives us a chance to cancel slurm
207 // jobs started by a previous dispatch process that never released
208 // their slurm allocations even though their container states are
209 // Cancelled or Complete. See https://dev.arvados.org/issues/10979
210 func (disp *Dispatcher) checkSqueueForOrphans() {
211 for _, uuid := range disp.sqCheck.All() {
212 if !containerUuidPattern.MatchString(uuid) {
215 err := disp.TrackContainer(uuid)
217 log.Printf("checkSqueueForOrphans: TrackContainer(%s): %s", uuid, err)
222 func (disp *Dispatcher) slurmConstraintArgs(container arvados.Container) []string {
223 mem := int64(math.Ceil(float64(container.RuntimeConstraints.RAM+
224 container.RuntimeConstraints.KeepCacheRAM+
225 int64(disp.cluster.Containers.ReserveExtraRAM)) / float64(1048576)))
227 disk := dispatchcloud.EstimateScratchSpace(&container)
228 disk = int64(math.Ceil(float64(disk) / float64(1048576)))
230 fmt.Sprintf("--mem=%d", mem),
231 fmt.Sprintf("--cpus-per-task=%d", container.RuntimeConstraints.VCPUs),
232 fmt.Sprintf("--tmp=%d", disk),
236 func (disp *Dispatcher) sbatchArgs(container arvados.Container) ([]string, error) {
238 args = append(args, disp.cluster.Containers.SLURM.SbatchArgumentsList...)
239 args = append(args, "--job-name="+container.UUID, fmt.Sprintf("--nice=%d", initialNiceValue), "--no-requeue")
241 if disp.cluster == nil {
242 // no instance types configured
243 args = append(args, disp.slurmConstraintArgs(container)...)
244 } else if it, err := dispatchcloud.ChooseInstanceType(disp.cluster, &container); err == dispatchcloud.ErrInstanceTypesNotConfigured {
246 args = append(args, disp.slurmConstraintArgs(container)...)
247 } else if err != nil {
250 // use instancetype constraint instead of slurm mem/cpu/tmp specs
251 args = append(args, "--constraint=instancetype="+it.Name)
254 if len(container.SchedulingParameters.Partitions) > 0 {
255 args = append(args, "--partition="+strings.Join(container.SchedulingParameters.Partitions, ","))
261 func (disp *Dispatcher) submit(container arvados.Container, crunchRunCommand []string) error {
262 // append() here avoids modifying crunchRunCommand's
263 // underlying array, which is shared with other goroutines.
264 crArgs := append([]string(nil), crunchRunCommand...)
265 crArgs = append(crArgs, container.UUID)
266 crScript := strings.NewReader(execScript(crArgs))
268 sbArgs, err := disp.sbatchArgs(container)
272 log.Printf("running sbatch %+q", sbArgs)
273 return disp.slurm.Batch(crScript, sbArgs)
276 // Submit a container to the slurm queue (or resume monitoring if it's
277 // already in the queue). Cancel the slurm job if the container's
278 // priority changes to zero or its state indicates it's no longer
280 func (disp *Dispatcher) runContainer(_ *dispatch.Dispatcher, ctr arvados.Container, status <-chan arvados.Container) {
281 ctx, cancel := context.WithCancel(context.Background())
284 if ctr.State == dispatch.Locked && !disp.sqCheck.HasUUID(ctr.UUID) {
285 log.Printf("Submitting container %s to slurm", ctr.UUID)
286 cmd := []string{disp.cluster.Containers.CrunchRunCommand}
287 cmd = append(cmd, disp.cluster.Containers.CrunchRunArgumentsList...)
288 if err := disp.submit(ctr, cmd); err != nil {
290 if err, ok := err.(dispatchcloud.ConstraintsNotSatisfiableError); ok {
291 var logBuf bytes.Buffer
292 fmt.Fprintf(&logBuf, "cannot run container %s: %s\n", ctr.UUID, err)
293 if len(err.AvailableTypes) == 0 {
294 fmt.Fprint(&logBuf, "No instance types are configured.\n")
296 fmt.Fprint(&logBuf, "Available instance types:\n")
297 for _, t := range err.AvailableTypes {
299 "Type %q: %d VCPUs, %d RAM, %d Scratch, %f Price\n",
300 t.Name, t.VCPUs, t.RAM, t.Scratch, t.Price,
304 text = logBuf.String()
305 disp.UpdateState(ctr.UUID, dispatch.Cancelled)
307 text = fmt.Sprintf("Error submitting container %s to slurm: %s", ctr.UUID, err)
311 lr := arvadosclient.Dict{"log": arvadosclient.Dict{
312 "object_uuid": ctr.UUID,
313 "event_type": "dispatch",
314 "properties": map[string]string{"text": text}}}
315 disp.Arv.Create("logs", lr, nil)
317 disp.Unlock(ctr.UUID)
322 log.Printf("Start monitoring container %v in state %q", ctr.UUID, ctr.State)
323 defer log.Printf("Done monitoring container %s", ctr.UUID)
325 // If the container disappears from the slurm queue, there is
326 // no point in waiting for further dispatch updates: just
327 // clean up and return.
328 go func(uuid string) {
329 for ctx.Err() == nil && disp.sqCheck.HasUUID(uuid) {
337 // Disappeared from squeue
338 if err := disp.Arv.Get("containers", ctr.UUID, nil, &ctr); err != nil {
339 log.Printf("error getting final container state for %s: %s", ctr.UUID, err)
342 case dispatch.Running:
343 disp.UpdateState(ctr.UUID, dispatch.Cancelled)
344 case dispatch.Locked:
345 disp.Unlock(ctr.UUID)
348 case updated, ok := <-status:
350 log.Printf("container %s is done: cancel slurm job", ctr.UUID)
352 } else if updated.Priority == 0 {
353 log.Printf("container %s has state %q, priority %d: cancel slurm job", ctr.UUID, updated.State, updated.Priority)
356 p := int64(updated.Priority)
359 // user-assigned priority. If
360 // ctrs have equal priority,
361 // run the older one first.
362 p = int64(p)<<50 - (updated.CreatedAt.UnixNano() >> 14)
364 disp.sqCheck.SetPriority(ctr.UUID, p)
369 func (disp *Dispatcher) scancel(ctr arvados.Container) {
370 err := disp.slurm.Cancel(ctr.UUID)
372 log.Printf("scancel: %s", err)
373 time.Sleep(time.Second)
374 } else if disp.sqCheck.HasUUID(ctr.UUID) {
375 log.Printf("container %s is still in squeue after scancel", ctr.UUID)
376 time.Sleep(time.Second)