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