9945: Merge branch 'master' into 9945-make-python-package-dependency-free
[arvados.git] / services / crunch-dispatch-slurm / squeue_test.go
1 // Copyright (C) The Arvados Authors. All rights reserved.
2 //
3 // SPDX-License-Identifier: AGPL-3.0
4
5 package main
6
7 import (
8         "time"
9
10         "github.com/sirupsen/logrus"
11         . "gopkg.in/check.v1"
12 )
13
14 var _ = Suite(&SqueueSuite{})
15
16 type SqueueSuite struct{}
17
18 func (s *SqueueSuite) TestReleasePending(c *C) {
19         uuids := []string{
20                 "zzzzz-dz642-fake0fake0fake0",
21                 "zzzzz-dz642-fake1fake1fake1",
22                 "zzzzz-dz642-fake2fake2fake2",
23         }
24         slurm := &slurmFake{
25                 queue: uuids[0] + " 10000 4294000000 PENDING Resources\n" + uuids[1] + " 10000 4294000111 PENDING Resources\n" + uuids[2] + " 10000 0 PENDING BadConstraints\n",
26         }
27         sqc := &SqueueChecker{
28                 Logger: logrus.StandardLogger(),
29                 Slurm:  slurm,
30                 Period: time.Hour,
31         }
32         sqc.startOnce.Do(sqc.start)
33         defer sqc.Stop()
34
35         done := make(chan struct{})
36         go func() {
37                 for _, u := range uuids {
38                         sqc.SetPriority(u, 1)
39                 }
40                 close(done)
41         }()
42         callUntilReady(sqc.check, done)
43
44         slurm.didRelease = nil
45         sqc.check()
46         c.Check(slurm.didRelease, DeepEquals, []string{uuids[2]})
47 }
48
49 func (s *SqueueSuite) TestReniceAll(c *C) {
50         uuids := []string{"zzzzz-dz642-fake0fake0fake0", "zzzzz-dz642-fake1fake1fake1", "zzzzz-dz642-fake2fake2fake2"}
51         for _, test := range []struct {
52                 spread int64
53                 squeue string
54                 want   map[string]int64
55                 expect [][]string
56         }{
57                 {
58                         spread: 1,
59                         squeue: uuids[0] + " 10000 4294000000 PENDING Resources\n",
60                         want:   map[string]int64{uuids[0]: 1},
61                         expect: [][]string{{uuids[0], "0"}},
62                 },
63                 { // fake0 priority is too high
64                         spread: 1,
65                         squeue: uuids[0] + " 10000 4294000777 PENDING Resources\n" + uuids[1] + " 10000 4294000444 PENDING Resources\n",
66                         want:   map[string]int64{uuids[0]: 1, uuids[1]: 999},
67                         expect: [][]string{{uuids[1], "0"}, {uuids[0], "334"}},
68                 },
69                 { // specify spread
70                         spread: 100,
71                         squeue: uuids[0] + " 10000 4294000777 PENDING Resources\n" + uuids[1] + " 10000 4294000444 PENDING Resources\n",
72                         want:   map[string]int64{uuids[0]: 1, uuids[1]: 999},
73                         expect: [][]string{{uuids[1], "0"}, {uuids[0], "433"}},
74                 },
75                 { // ignore fake2 because SetPriority() not called
76                         spread: 1,
77                         squeue: uuids[0] + " 10000 4294000000 PENDING Resources\n" + uuids[1] + " 10000 4294000111 PENDING Resources\n" + uuids[2] + " 10000 4294000222 PENDING Resources\n",
78                         want:   map[string]int64{uuids[0]: 999, uuids[1]: 1},
79                         expect: [][]string{{uuids[0], "0"}, {uuids[1], "112"}},
80                 },
81                 { // ignore fake2 because slurm priority=0
82                         spread: 1,
83                         squeue: uuids[0] + " 10000 4294000000 PENDING Resources\n" + uuids[1] + " 10000 4294000111 PENDING Resources\n" + uuids[2] + " 10000 0 PENDING Resources\n",
84                         want:   map[string]int64{uuids[0]: 999, uuids[1]: 1, uuids[2]: 997},
85                         expect: [][]string{{uuids[0], "0"}, {uuids[1], "112"}},
86                 },
87         } {
88                 c.Logf("spread=%d squeue=%q want=%v -> expect=%v", test.spread, test.squeue, test.want, test.expect)
89                 slurm := &slurmFake{
90                         queue: test.squeue,
91                 }
92                 sqc := &SqueueChecker{
93                         Logger:         logrus.StandardLogger(),
94                         Slurm:          slurm,
95                         PrioritySpread: test.spread,
96                         Period:         time.Hour,
97                 }
98                 sqc.startOnce.Do(sqc.start)
99                 sqc.check()
100                 for uuid, pri := range test.want {
101                         sqc.SetPriority(uuid, pri)
102                 }
103                 sqc.reniceAll()
104                 c.Check(slurm.didRenice, DeepEquals, test.expect)
105                 sqc.Stop()
106         }
107 }
108
109 // If a limited nice range prevents desired priority adjustments, give
110 // up and clamp nice to 10K.
111 func (s *SqueueSuite) TestReniceInvalidNiceValue(c *C) {
112         uuids := []string{"zzzzz-dz642-fake0fake0fake0", "zzzzz-dz642-fake1fake1fake1", "zzzzz-dz642-fake2fake2fake2"}
113         slurm := &slurmFake{
114                 queue:         uuids[0] + " 0 4294000222 PENDING Resources\n" + uuids[1] + " 0 4294555222 PENDING Resources\n",
115                 rejectNice10K: true,
116         }
117         sqc := &SqueueChecker{
118                 Logger:         logrus.StandardLogger(),
119                 Slurm:          slurm,
120                 PrioritySpread: 1,
121                 Period:         time.Hour,
122         }
123         sqc.startOnce.Do(sqc.start)
124         sqc.check()
125         sqc.SetPriority(uuids[0], 2)
126         sqc.SetPriority(uuids[1], 1)
127
128         // First attempt should renice to 555001, which will fail
129         sqc.reniceAll()
130         c.Check(slurm.didRenice, DeepEquals, [][]string{{uuids[1], "555001"}})
131
132         // Next attempt should renice to 10K, which will succeed
133         sqc.reniceAll()
134         c.Check(slurm.didRenice, DeepEquals, [][]string{{uuids[1], "555001"}, {uuids[1], "10000"}})
135         // ...so we'll change the squeue response to reflect the
136         // updated priority+nice, and make sure sqc sees that...
137         slurm.queue = uuids[0] + " 0 4294000222 PENDING Resources\n" + uuids[1] + " 10000 4294545222 PENDING Resources\n"
138         sqc.check()
139
140         // Next attempt should leave nice alone because it's already
141         // at the 10K limit
142         sqc.reniceAll()
143         c.Check(slurm.didRenice, DeepEquals, [][]string{{uuids[1], "555001"}, {uuids[1], "10000"}})
144
145         // Back to normal if desired nice value falls below 10K
146         slurm.queue = uuids[0] + " 0 4294000222 PENDING Resources\n" + uuids[1] + " 10000 4294000111 PENDING Resources\n"
147         sqc.check()
148         sqc.reniceAll()
149         c.Check(slurm.didRenice, DeepEquals, [][]string{{uuids[1], "555001"}, {uuids[1], "10000"}, {uuids[1], "9890"}})
150
151         sqc.Stop()
152 }
153
154 // If the given UUID isn't in the slurm queue yet, SetPriority()
155 // should wait for it to appear on the very next poll, then give up.
156 func (s *SqueueSuite) TestSetPriorityBeforeQueued(c *C) {
157         uuidGood := "zzzzz-dz642-fake0fake0fake0"
158         uuidBad := "zzzzz-dz642-fake1fake1fake1"
159
160         slurm := &slurmFake{}
161         sqc := &SqueueChecker{
162                 Logger: logrus.StandardLogger(),
163                 Slurm:  slurm,
164                 Period: time.Hour,
165         }
166         sqc.startOnce.Do(sqc.start)
167         sqc.Stop()
168         sqc.check()
169
170         done := make(chan struct{})
171         go func() {
172                 sqc.SetPriority(uuidGood, 123)
173                 sqc.SetPriority(uuidBad, 345)
174                 close(done)
175         }()
176         c.Check(sqc.queue[uuidGood], IsNil)
177         c.Check(sqc.queue[uuidBad], IsNil)
178         timeout := time.NewTimer(time.Second)
179         defer timeout.Stop()
180         tick := time.NewTicker(time.Millisecond)
181         defer tick.Stop()
182         for {
183                 select {
184                 case <-tick.C:
185                         slurm.queue = uuidGood + " 0 12345 PENDING Resources\n"
186                         sqc.check()
187
188                         // Avoid immediately selecting this case again
189                         // on the next iteration if check() took
190                         // longer than one tick.
191                         select {
192                         case <-tick.C:
193                         default:
194                         }
195                 case <-timeout.C:
196                         c.Fatal("timed out")
197                 case <-done:
198                         c.Assert(sqc.queue[uuidGood], NotNil)
199                         c.Check(sqc.queue[uuidGood].wantPriority, Equals, int64(123))
200                         c.Check(sqc.queue[uuidBad], IsNil)
201                         return
202                 }
203         }
204 }
205
206 func callUntilReady(fn func(), done <-chan struct{}) {
207         tick := time.NewTicker(time.Millisecond)
208         defer tick.Stop()
209         for {
210                 select {
211                 case <-done:
212                         return
213                 case <-tick.C:
214                         fn()
215                 }
216         }
217 }