lib/dispatchcloud/scheduler/map_test.go

   1 // Copyright (C) The Arvados Authors. All rights reserved.
   2 //
   3 // SPDX-License-Identifier: AGPL-3.0
   4
   5 package scheduler
   6
   7 import (
   8         "errors"
   9         "fmt"
  10         "time"
  11
  12         "git.curoverse.com/arvados.git/lib/dispatchcloud/container"
  13         "git.curoverse.com/arvados.git/lib/dispatchcloud/test"
  14         "git.curoverse.com/arvados.git/lib/dispatchcloud/worker"
  15         "git.curoverse.com/arvados.git/sdk/go/arvados"
  16         "github.com/Sirupsen/logrus"
  17         check "gopkg.in/check.v1"
  18 )
  19
  20 var (
  21         logger = logrus.StandardLogger()
  22
  23         // arbitrary example instance types
  24         types = func() (r []arvados.InstanceType) {
  25                 for i := 0; i < 16; i++ {
  26                         r = append(r, test.InstanceType(i))
  27                 }
  28                 return
  29         }()
  30
  31         // arbitrary example container UUIDs
  32         uuids = func() (r []string) {
  33                 for i := 0; i < 16; i++ {
  34                         r = append(r, test.ContainerUUID(i))
  35                 }
  36                 return
  37         }()
  38 )
  39
  40 type stubQueue struct {
  41         ents map[string]container.QueueEnt
  42 }
  43
  44 func (q *stubQueue) Entries() (map[string]container.QueueEnt, time.Time) {
  45         return q.ents, time.Now()
  46 }
  47 func (q *stubQueue) Lock(uuid string) error {
  48         return q.setState(uuid, arvados.ContainerStateLocked)
  49 }
  50 func (q *stubQueue) Unlock(uuid string) error {
  51         return q.setState(uuid, arvados.ContainerStateQueued)
  52 }
  53 func (q *stubQueue) Cancel(uuid string) error {
  54         return q.setState(uuid, arvados.ContainerStateCancelled)
  55 }
  56 func (q *stubQueue) Forget(uuid string) {
  57 }
  58 func (q *stubQueue) Get(uuid string) (arvados.Container, bool) {
  59         ent, ok := q.ents[uuid]
  60         return ent.Container, ok
  61 }
  62 func (q *stubQueue) setState(uuid string, state arvados.ContainerState) error {
  63         ent, ok := q.ents[uuid]
  64         if !ok {
  65                 return fmt.Errorf("no such ent: %q", uuid)
  66         }
  67         ent.Container.State = state
  68         q.ents[uuid] = ent
  69         return nil
  70 }
  71
  72 type stubQuotaError struct {
  73         error
  74 }
  75
  76 func (stubQuotaError) IsQuotaError() bool { return true }
  77
  78 type stubPool struct {
  79         notify    <-chan struct{}
  80         unalloc   map[arvados.InstanceType]int // idle+booting+unknown
  81         idle      map[arvados.InstanceType]int
  82         running   map[string]time.Time
  83         atQuota   bool
  84         canCreate int
  85         creates   []arvados.InstanceType
  86         starts    []string
  87         shutdowns int
  88 }
  89
  90 func (p *stubPool) AtQuota() bool                 { return p.atQuota }
  91 func (p *stubPool) Subscribe() <-chan struct{}    { return p.notify }
  92 func (p *stubPool) Unsubscribe(<-chan struct{})   {}
  93 func (p *stubPool) Running() map[string]time.Time { return p.running }
  94 func (p *stubPool) Unallocated() map[arvados.InstanceType]int {
  95         r := map[arvados.InstanceType]int{}
  96         for it, n := range p.unalloc {
  97                 r[it] = n
  98         }
  99         return r
 100 }
 101 func (p *stubPool) Create(it arvados.InstanceType) error {
 102         p.creates = append(p.creates, it)
 103         if p.canCreate < 1 {
 104                 return stubQuotaError{errors.New("quota")}
 105         }
 106         p.canCreate--
 107         p.unalloc[it]++
 108         return nil
 109 }
 110 func (p *stubPool) KillContainer(uuid string) {
 111         p.running[uuid] = time.Now()
 112 }
 113 func (p *stubPool) Shutdown(arvados.InstanceType) bool {
 114         p.shutdowns++
 115         return false
 116 }
 117 func (p *stubPool) Workers() map[worker.State]int {
 118         return map[worker.State]int{
 119                 worker.StateBooting: len(p.unalloc) - len(p.idle),
 120                 worker.StateRunning: len(p.idle) - len(p.running),
 121         }
 122 }
 123 func (p *stubPool) StartContainer(it arvados.InstanceType, ctr arvados.Container) bool {
 124         p.starts = append(p.starts, ctr.UUID)
 125         if p.idle[it] == 0 {
 126                 return false
 127         }
 128         p.idle[it]--
 129         p.unalloc[it]--
 130         p.running[ctr.UUID] = time.Time{}
 131         return true
 132 }
 133
 134 var _ = check.Suite(&SchedulerSuite{})
 135
 136 type SchedulerSuite struct{}
 137
 138 // Map priority=4 container to idle node. Create a new instance for
 139 // the priority=3 container. Don't try to start any priority<3
 140 // containers because priority=3 container didn't start
 141 // immediately. Don't try to create any other nodes after the failed
 142 // create.
 143 func (*SchedulerSuite) TestMapIdle(c *check.C) {
 144         queue := stubQueue{
 145                 ents: map[string]container.QueueEnt{
 146                         uuids[1]: {
 147                                 Container:    arvados.Container{UUID: uuids[1], Priority: 1, State: arvados.ContainerStateQueued},
 148                                 InstanceType: types[1],
 149                         },
 150                         uuids[2]: {
 151                                 Container:    arvados.Container{UUID: uuids[2], Priority: 2, State: arvados.ContainerStateQueued},
 152                                 InstanceType: types[1],
 153                         },
 154                         uuids[3]: {
 155                                 Container:    arvados.Container{UUID: uuids[3], Priority: 3, State: arvados.ContainerStateQueued},
 156                                 InstanceType: types[1],
 157                         },
 158                         uuids[4]: {
 159                                 Container:    arvados.Container{UUID: uuids[4], Priority: 4, State: arvados.ContainerStateQueued},
 160                                 InstanceType: types[1],
 161                         },
 162                 },
 163         }
 164         pool := stubPool{
 165                 unalloc: map[arvados.InstanceType]int{
 166                         types[1]: 1,
 167                         types[2]: 2,
 168                 },
 169                 idle: map[arvados.InstanceType]int{
 170                         types[1]: 1,
 171                         types[2]: 2,
 172                 },
 173                 running:   map[string]time.Time{},
 174                 canCreate: 1,
 175         }
 176         Map(logger, &queue, &pool)
 177         c.Check(pool.creates, check.DeepEquals, []arvados.InstanceType{types[1]})
 178         c.Check(pool.starts, check.DeepEquals, []string{uuids[4], uuids[3]})
 179         c.Check(pool.running, check.HasLen, 1)
 180         for uuid := range pool.running {
 181                 c.Check(uuid, check.Equals, uuids[4])
 182         }
 183 }
 184
 185 // Shutdown some nodes if Create() fails -- and without even calling
 186 // Create(), if AtQuota() is true.
 187 func (*SchedulerSuite) TestMapShutdownAtQuota(c *check.C) {
 188         for quota := 0; quota < 2; quota++ {
 189                 shouldCreate := types[1 : 1+quota]
 190                 queue := stubQueue{
 191                         ents: map[string]container.QueueEnt{
 192                                 uuids[1]: {
 193                                         Container:    arvados.Container{UUID: uuids[1], Priority: 1, State: arvados.ContainerStateQueued},
 194                                         InstanceType: types[1],
 195                                 },
 196                         },
 197                 }
 198                 pool := stubPool{
 199                         atQuota: quota == 0,
 200                         unalloc: map[arvados.InstanceType]int{
 201                                 types[2]: 2,
 202                         },
 203                         idle: map[arvados.InstanceType]int{
 204                                 types[2]: 2,
 205                         },
 206                         running:   map[string]time.Time{},
 207                         creates:   []arvados.InstanceType{},
 208                         starts:    []string{},
 209                         canCreate: 0,
 210                 }
 211                 Map(logger, &queue, &pool)
 212                 c.Check(pool.creates, check.DeepEquals, shouldCreate)
 213                 c.Check(pool.starts, check.DeepEquals, []string{})
 214                 c.Check(pool.shutdowns, check.Not(check.Equals), 0)
 215         }
 216 }
 217
 218 // Start lower-priority containers while waiting for new/existing
 219 // workers to come up for higher-priority containers.
 220 func (*SchedulerSuite) TestMapStartWhileCreating(c *check.C) {
 221         pool := stubPool{
 222                 unalloc: map[arvados.InstanceType]int{
 223                         types[1]: 1,
 224                         types[2]: 1,
 225                 },
 226                 idle: map[arvados.InstanceType]int{
 227                         types[1]: 1,
 228                         types[2]: 1,
 229                 },
 230                 running:   map[string]time.Time{},
 231                 canCreate: 2,
 232         }
 233         queue := stubQueue{
 234                 ents: map[string]container.QueueEnt{
 235                         uuids[1]: {
 236                                 // create a new worker
 237                                 Container:    arvados.Container{UUID: uuids[1], Priority: 1, State: arvados.ContainerStateQueued},
 238                                 InstanceType: types[1],
 239                         },
 240                         uuids[2]: {
 241                                 // tentatively map to unalloc worker
 242                                 Container:    arvados.Container{UUID: uuids[2], Priority: 2, State: arvados.ContainerStateQueued},
 243                                 InstanceType: types[1],
 244                         },
 245                         uuids[3]: {
 246                                 // start now on idle worker
 247                                 Container:    arvados.Container{UUID: uuids[3], Priority: 3, State: arvados.ContainerStateQueued},
 248                                 InstanceType: types[1],
 249                         },
 250                         uuids[4]: {
 251                                 // create a new worker
 252                                 Container:    arvados.Container{UUID: uuids[4], Priority: 4, State: arvados.ContainerStateQueued},
 253                                 InstanceType: types[2],
 254                         },
 255                         uuids[5]: {
 256                                 // tentatively map to unalloc worker
 257                                 Container:    arvados.Container{UUID: uuids[5], Priority: 5, State: arvados.ContainerStateQueued},
 258                                 InstanceType: types[2],
 259                         },
 260                         uuids[6]: {
 261                                 // start now on idle worker
 262                                 Container:    arvados.Container{UUID: uuids[6], Priority: 6, State: arvados.ContainerStateQueued},
 263                                 InstanceType: types[2],
 264                         },
 265                 },
 266         }
 267         Map(logger, &queue, &pool)
 268         c.Check(pool.creates, check.DeepEquals, []arvados.InstanceType{types[2], types[1]})
 269         c.Check(pool.starts, check.DeepEquals, []string{uuids[6], uuids[5], uuids[3], uuids[2]})
 270         running := map[string]bool{}
 271         for uuid, t := range pool.running {
 272                 if t.IsZero() {
 273                         running[uuid] = false
 274                 } else {
 275                         running[uuid] = true
 276                 }
 277         }
 278         c.Check(running, check.DeepEquals, map[string]bool{uuids[3]: false, uuids[6]: false})
 279 }