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         . "gopkg.in/check.v1"
  11 )
  12
  13 var _ = Suite(&SqueueSuite{})
  14
  15 type SqueueSuite struct{}
  16
  17 func (s *SqueueSuite) TestReleasePending(c *C) {
  18         uuids := []string{
  19                 "zzzzz-dz642-fake0fake0fake0",
  20                 "zzzzz-dz642-fake1fake1fake1",
  21                 "zzzzz-dz642-fake2fake2fake2",
  22         }
  23         slurm := &slurmFake{
  24                 queue: uuids[0] + " 10000 4294000000 PENDING Resources\n" + uuids[1] + " 10000 4294000111 PENDING Resources\n" + uuids[2] + " 10000 0 PENDING BadConstraints\n",
  25         }
  26         sqc := &SqueueChecker{
  27                 Slurm:  slurm,
  28                 Period: time.Hour,
  29         }
  30         sqc.startOnce.Do(sqc.start)
  31         defer sqc.Stop()
  32
  33         done := make(chan struct{})
  34         go func() {
  35                 for _, u := range uuids {
  36                         sqc.SetPriority(u, 1)
  37                 }
  38                 close(done)
  39         }()
  40         callUntilReady(sqc.check, done)
  41
  42         slurm.didRelease = nil
  43         sqc.check()
  44         c.Check(slurm.didRelease, DeepEquals, []string{uuids[2]})
  45 }
  46
  47 func (s *SqueueSuite) TestReniceAll(c *C) {
  48         uuids := []string{"zzzzz-dz642-fake0fake0fake0", "zzzzz-dz642-fake1fake1fake1", "zzzzz-dz642-fake2fake2fake2"}
  49         for _, test := range []struct {
  50                 spread int64
  51                 squeue string
  52                 want   map[string]int64
  53                 expect [][]string
  54         }{
  55                 {
  56                         spread: 1,
  57                         squeue: uuids[0] + " 10000 4294000000 PENDING Resources\n",
  58                         want:   map[string]int64{uuids[0]: 1},
  59                         expect: [][]string{{uuids[0], "0"}},
  60                 },
  61                 { // fake0 priority is too high
  62                         spread: 1,
  63                         squeue: uuids[0] + " 10000 4294000777 PENDING Resources\n" + uuids[1] + " 10000 4294000444 PENDING Resources\n",
  64                         want:   map[string]int64{uuids[0]: 1, uuids[1]: 999},
  65                         expect: [][]string{{uuids[1], "0"}, {uuids[0], "334"}},
  66                 },
  67                 { // specify spread
  68                         spread: 100,
  69                         squeue: uuids[0] + " 10000 4294000777 PENDING Resources\n" + uuids[1] + " 10000 4294000444 PENDING Resources\n",
  70                         want:   map[string]int64{uuids[0]: 1, uuids[1]: 999},
  71                         expect: [][]string{{uuids[1], "0"}, {uuids[0], "433"}},
  72                 },
  73                 { // ignore fake2 because SetPriority() not called
  74                         spread: 1,
  75                         squeue: uuids[0] + " 10000 4294000000 PENDING Resources\n" + uuids[1] + " 10000 4294000111 PENDING Resources\n" + uuids[2] + " 10000 4294000222 PENDING Resources\n",
  76                         want:   map[string]int64{uuids[0]: 999, uuids[1]: 1},
  77                         expect: [][]string{{uuids[0], "0"}, {uuids[1], "112"}},
  78                 },
  79                 { // ignore fake2 because slurm priority=0
  80                         spread: 1,
  81                         squeue: uuids[0] + " 10000 4294000000 PENDING Resources\n" + uuids[1] + " 10000 4294000111 PENDING Resources\n" + uuids[2] + " 10000 0 PENDING Resources\n",
  82                         want:   map[string]int64{uuids[0]: 999, uuids[1]: 1, uuids[2]: 997},
  83                         expect: [][]string{{uuids[0], "0"}, {uuids[1], "112"}},
  84                 },
  85         } {
  86                 c.Logf("spread=%d squeue=%q want=%v -> expect=%v", test.spread, test.squeue, test.want, test.expect)
  87                 slurm := &slurmFake{
  88                         queue: test.squeue,
  89                 }
  90                 sqc := &SqueueChecker{
  91                         Slurm:          slurm,
  92                         PrioritySpread: test.spread,
  93                         Period:         time.Hour,
  94                 }
  95                 sqc.startOnce.Do(sqc.start)
  96                 sqc.check()
  97                 for uuid, pri := range test.want {
  98                         sqc.SetPriority(uuid, pri)
  99                 }
 100                 sqc.reniceAll()
 101                 c.Check(slurm.didRenice, DeepEquals, test.expect)
 102                 sqc.Stop()
 103         }
 104 }
 105
 106 // If the given UUID isn't in the slurm queue yet, SetPriority()
 107 // should wait for it to appear on the very next poll, then give up.
 108 func (s *SqueueSuite) TestSetPriorityBeforeQueued(c *C) {
 109         uuidGood := "zzzzz-dz642-fake0fake0fake0"
 110         uuidBad := "zzzzz-dz642-fake1fake1fake1"
 111
 112         slurm := &slurmFake{}
 113         sqc := &SqueueChecker{
 114                 Slurm:  slurm,
 115                 Period: time.Hour,
 116         }
 117         sqc.startOnce.Do(sqc.start)
 118         sqc.Stop()
 119         sqc.check()
 120
 121         done := make(chan struct{})
 122         go func() {
 123                 sqc.SetPriority(uuidGood, 123)
 124                 sqc.SetPriority(uuidBad, 345)
 125                 close(done)
 126         }()
 127         c.Check(sqc.queue[uuidGood], IsNil)
 128         c.Check(sqc.queue[uuidBad], IsNil)
 129         timeout := time.NewTimer(time.Second)
 130         defer timeout.Stop()
 131         tick := time.NewTicker(time.Millisecond)
 132         defer tick.Stop()
 133         for {
 134                 select {
 135                 case <-tick.C:
 136                         slurm.queue = uuidGood + " 0 12345 PENDING Resources\n"
 137                         sqc.check()
 138
 139                         // Avoid immediately selecting this case again
 140                         // on the next iteration if check() took
 141                         // longer than one tick.
 142                         select {
 143                         case <-tick.C:
 144                         default:
 145                         }
 146                 case <-timeout.C:
 147                         c.Fatal("timed out")
 148                 case <-done:
 149                         c.Assert(sqc.queue[uuidGood], NotNil)
 150                         c.Check(sqc.queue[uuidGood].wantPriority, Equals, int64(123))
 151                         c.Check(sqc.queue[uuidBad], IsNil)
 152                         return
 153                 }
 154         }
 155 }
 156
 157 func callUntilReady(fn func(), done <-chan struct{}) {
 158         tick := time.NewTicker(time.Millisecond)
 159         defer tick.Stop()
 160         for {
 161                 select {
 162                 case <-done:
 163                         return
 164                 case <-tick.C:
 165                         fn()
 166                 }
 167         }
 168 }