20831: Adding update_remote_user method to User, wip
[arvados.git] / services / crunch-dispatch-slurm / squeue.go
1 // Copyright (C) The Arvados Authors. All rights reserved.
2 //
3 // SPDX-License-Identifier: AGPL-3.0
4
5 package dispatchslurm
6
7 import (
8         "bytes"
9         "fmt"
10         "sort"
11         "strings"
12         "sync"
13         "time"
14 )
15
16 const slurm15NiceLimit int64 = 10000
17
18 type slurmJob struct {
19         uuid         string
20         wantPriority int64
21         priority     int64 // current slurm priority (incorporates nice value)
22         nice         int64 // current slurm nice value
23         hitNiceLimit bool
24 }
25
26 // SqueueChecker implements asynchronous polling monitor of the SLURM queue
27 // using the command 'squeue'.
28 type SqueueChecker struct {
29         Logger         logger
30         Period         time.Duration
31         PrioritySpread int64
32         Slurm          Slurm
33         queue          map[string]*slurmJob
34         startOnce      sync.Once
35         done           chan struct{}
36         lock           sync.RWMutex
37         notify         sync.Cond
38 }
39
40 // HasUUID checks if a given container UUID is in the slurm queue.
41 // This does not run squeue directly, but instead blocks until woken
42 // up by next successful update of squeue.
43 func (sqc *SqueueChecker) HasUUID(uuid string) bool {
44         sqc.startOnce.Do(sqc.start)
45
46         sqc.lock.RLock()
47         defer sqc.lock.RUnlock()
48
49         // block until next squeue broadcast signaling an update.
50         sqc.notify.Wait()
51         _, exists := sqc.queue[uuid]
52         return exists
53 }
54
55 // SetPriority sets or updates the desired (Arvados) priority for a
56 // container.
57 func (sqc *SqueueChecker) SetPriority(uuid string, want int64) {
58         sqc.startOnce.Do(sqc.start)
59
60         sqc.lock.RLock()
61         job := sqc.queue[uuid]
62         if job == nil {
63                 // Wait in case the slurm job was just submitted and
64                 // will appear in the next squeue update.
65                 sqc.notify.Wait()
66                 job = sqc.queue[uuid]
67         }
68         needUpdate := job != nil && job.wantPriority != want
69         sqc.lock.RUnlock()
70
71         if needUpdate {
72                 sqc.lock.Lock()
73                 job.wantPriority = want
74                 sqc.lock.Unlock()
75         }
76 }
77
78 // adjust slurm job nice values as needed to ensure slurm priority
79 // order matches Arvados priority order.
80 func (sqc *SqueueChecker) reniceAll() {
81         // This is slow (it shells out to scontrol many times) and no
82         // other goroutines update sqc.queue or any of the job fields
83         // we use here, so we don't acquire a lock.
84         jobs := make([]*slurmJob, 0, len(sqc.queue))
85         for _, j := range sqc.queue {
86                 if j.wantPriority == 0 {
87                         // SLURM job with unknown Arvados priority
88                         // (perhaps it's not an Arvados job)
89                         continue
90                 }
91                 if j.priority <= 2*slurm15NiceLimit {
92                         // SLURM <= 15.x implements "hold" by setting
93                         // priority to 0. If we include held jobs
94                         // here, we'll end up trying to push other
95                         // jobs below them using negative priority,
96                         // which won't help anything.
97                         continue
98                 }
99                 jobs = append(jobs, j)
100         }
101
102         sort.Slice(jobs, func(i, j int) bool {
103                 if jobs[i].wantPriority != jobs[j].wantPriority {
104                         return jobs[i].wantPriority > jobs[j].wantPriority
105                 }
106                 // break ties with container uuid --
107                 // otherwise, the ordering would change from
108                 // one interval to the next, and we'd do many
109                 // pointless slurm queue rearrangements.
110                 return jobs[i].uuid > jobs[j].uuid
111         })
112         renice := wantNice(jobs, sqc.PrioritySpread)
113         for i, job := range jobs {
114                 niceNew := renice[i]
115                 if job.hitNiceLimit && niceNew > slurm15NiceLimit {
116                         niceNew = slurm15NiceLimit
117                 }
118                 if niceNew == job.nice {
119                         continue
120                 }
121                 err := sqc.Slurm.Renice(job.uuid, niceNew)
122                 if err != nil && niceNew > slurm15NiceLimit && strings.Contains(err.Error(), "Invalid nice value") {
123                         sqc.Logger.Warnf("container %q clamping nice values at %d, priority order will not be correct -- see https://dev.arvados.org/projects/arvados/wiki/SLURM_integration#Limited-nice-values-SLURM-15", job.uuid, slurm15NiceLimit)
124                         job.hitNiceLimit = true
125                 }
126         }
127 }
128
129 // Stop stops the squeue monitoring goroutine. Do not call HasUUID
130 // after calling Stop.
131 func (sqc *SqueueChecker) Stop() {
132         if sqc.done != nil {
133                 close(sqc.done)
134         }
135 }
136
137 // check gets the names of jobs in the SLURM queue (running and
138 // queued). If it succeeds, it updates sqc.queue and wakes up any
139 // goroutines that are waiting in HasUUID() or All().
140 func (sqc *SqueueChecker) check() {
141         cmd := sqc.Slurm.QueueCommand([]string{"--all", "--noheader", "--format=%j %y %Q %T %r"})
142         stdout, stderr := &bytes.Buffer{}, &bytes.Buffer{}
143         cmd.Stdout, cmd.Stderr = stdout, stderr
144         if err := cmd.Run(); err != nil {
145                 sqc.Logger.Warnf("Error running %q %q: %s %q", cmd.Path, cmd.Args, err, stderr.String())
146                 return
147         }
148
149         lines := strings.Split(stdout.String(), "\n")
150         newq := make(map[string]*slurmJob, len(lines))
151         for _, line := range lines {
152                 if line == "" {
153                         continue
154                 }
155                 var uuid, state, reason string
156                 var n, p int64
157                 if _, err := fmt.Sscan(line, &uuid, &n, &p, &state, &reason); err != nil {
158                         sqc.Logger.Warnf("ignoring unparsed line in squeue output: %q", line)
159                         continue
160                 }
161
162                 // No other goroutines write to jobs' priority or nice
163                 // fields, so we can read and write them without
164                 // locks.
165                 replacing, ok := sqc.queue[uuid]
166                 if !ok {
167                         replacing = &slurmJob{uuid: uuid}
168                 }
169                 replacing.priority = p
170                 replacing.nice = n
171                 newq[uuid] = replacing
172
173                 if state == "PENDING" && ((reason == "BadConstraints" && p <= 2*slurm15NiceLimit) || reason == "launch failed requeued held") && replacing.wantPriority > 0 {
174                         // When using SLURM 14.x or 15.x, our queued
175                         // jobs land in this state when "scontrol
176                         // reconfigure" invalidates their feature
177                         // constraints by clearing all node features.
178                         // They stay in this state even after the
179                         // features reappear, until we run "scontrol
180                         // release {jobid}". Priority is usually 0 in
181                         // this state, but sometimes (due to a race
182                         // with nice adjustments?) it's a small
183                         // positive value.
184                         //
185                         // "scontrol release" is silent and successful
186                         // regardless of whether the features have
187                         // reappeared, so rather than second-guessing
188                         // whether SLURM is ready, we just keep trying
189                         // this until it works.
190                         //
191                         // "launch failed requeued held" seems to be
192                         // another manifestation of this problem,
193                         // resolved the same way.
194                         sqc.Logger.Printf("releasing held job %q (priority=%d, state=%q, reason=%q)", uuid, p, state, reason)
195                         sqc.Slurm.Release(uuid)
196                 } else if state != "RUNNING" && p <= 2*slurm15NiceLimit && replacing.wantPriority > 0 {
197                         sqc.Logger.Warnf("job %q has low priority %d, nice %d, state %q, reason %q", uuid, p, n, state, reason)
198                 }
199         }
200         sqc.lock.Lock()
201         sqc.queue = newq
202         sqc.lock.Unlock()
203         sqc.notify.Broadcast()
204 }
205
206 // Initialize, and start a goroutine to call check() once per
207 // squeue.Period until terminated by calling Stop().
208 func (sqc *SqueueChecker) start() {
209         sqc.notify.L = sqc.lock.RLocker()
210         sqc.done = make(chan struct{})
211         go func() {
212                 ticker := time.NewTicker(sqc.Period)
213                 for {
214                         select {
215                         case <-sqc.done:
216                                 ticker.Stop()
217                                 return
218                         case <-ticker.C:
219                                 sqc.check()
220                                 sqc.reniceAll()
221                                 select {
222                                 case <-ticker.C:
223                                         // If this iteration took
224                                         // longer than sqc.Period,
225                                         // consume the next tick and
226                                         // wait. Otherwise we would
227                                         // starve other goroutines.
228                                 default:
229                                 }
230                         }
231                 }
232         }()
233 }
234
235 // All waits for the next squeue invocation, and returns all job
236 // names reported by squeue.
237 func (sqc *SqueueChecker) All() []string {
238         sqc.startOnce.Do(sqc.start)
239         sqc.lock.RLock()
240         defer sqc.lock.RUnlock()
241         sqc.notify.Wait()
242         var uuids []string
243         for u := range sqc.queue {
244                 uuids = append(uuids, u)
245         }
246         return uuids
247 }