From a25a4fbcfed0c0852bd2ebb3b110f46150cd5e25 Mon Sep 17 00:00:00 2001 From: Chenyi Wang Date: Fri, 2 Oct 2026 23:56:47 +0000 Subject: [PATCH] ateapi: prefer less-loaded workers via power of two choices (#1915) ### Summary With multi-actor worker support (#1836) enabled by default (`--max-actors=1000`), selecting a single worker uniformly at random across all candidates with remaining capacity can stack actors onto busy workers while lighter workers exist, whereas strict global least-loaded selection over the asynchronous watch cache can cause concurrent bursts to herd onto the same worker. This change updates `Scheduler.Schedule` to use **power of two choices**: sample two distinct eligible candidates that satisfy `HasRoom` and pick the less-loaded one, comparing actor-slot utilization (`Allocated.Actors / Capacity.Actors`) first and breaking ties on dominant compute-resource utilization (CPU/memory). ### Testing - Added unit tests covering power-of-two-choices ordering, heterogeneous worker actor-capacity utilization, compute-resource utilization tie-breaking, fallback when an idle worker lacks resource room, and tie-breaking reachability in `cmd/ateapi/internal/scheduling/scheduling_test.go`. - [x] Tests pass - [ ] Appropriate changes to documentation are included in the PR --- cmd/ateapi/internal/scheduling/scheduling.go | 135 +++++++++++-- .../internal/scheduling/scheduling_test.go | 183 +++++++++++++++++- 2 files changed, 302 insertions(+), 16 deletions(-) diff --git a/cmd/ateapi/internal/scheduling/scheduling.go b/cmd/ateapi/internal/scheduling/scheduling.go index d78fa8c90..2eb1af0dc 100644 --- a/cmd/ateapi/internal/scheduling/scheduling.go +++ b/cmd/ateapi/internal/scheduling/scheduling.go @@ -81,8 +81,8 @@ type scheduler struct { // Option configures the Scheduler returned by New. type Option func(*scheduler) -// WithIntn overrides the random source used to pick among equally suitable -// workers. n is always >= 1. +// WithIntn overrides the random source used to sample candidate pairs: when +// n >= 2 candidates exist, Schedule calls intn(n) and then intn(n-1). func WithIntn(intn func(n int) int) Option { return func(s *scheduler) { s.intn = intn } } @@ -96,25 +96,120 @@ func New(source WorkerSource, opts ...Option) Scheduler { return s } -// Schedule filters the current worker fleet to find unassigned candidates matching the given constraints. +// Schedule filters the fleet for eligible candidates with room, samples two at +// random (power of two choices), and returns the less-loaded one, where load is +// the higher of actor-slot and compute-resource utilization. Spreading across +// the warm pool avoids hotspots until autoscaling reclaims idle workers. func (s *scheduler) Schedule(ctx context.Context, constraints Constraints) (*ateapipb.Worker, error) { workers, err := s.source.Workers() if err != nil { return nil, fmt.Errorf("while listing workers: %w", err) } - var candidates []*ateapipb.Worker + want, err := resources.ParseQuantities(constraints.Limits) + if err != nil { + return nil, fmt.Errorf("while parsing actor resource limits: %w", err) + } + + var candidates []candidate for _, worker := range workers { - if s.Applies(worker, constraints) && s.HasRoom(worker, constraints) { - candidates = append(candidates, worker) + if !s.Applies(worker, constraints) { + continue + } + if cand, ok := checkRoom(worker, want); ok { + candidates = append(candidates, cand) } } if len(candidates) == 0 { return nil, ErrNoCapacity } + if len(candidates) == 1 { + return candidates[0].worker, nil + } - return candidates[s.intn(len(candidates))], nil + i := s.intn(len(candidates)) + j := s.intn(len(candidates) - 1) + if j >= i { + j++ + } + if lessLoaded(&candidates[j], &candidates[i]) { + return candidates[j].worker, nil + } + return candidates[i].worker, nil +} + +// candidate pairs an eligible worker with its cached compute utilization. +type candidate struct { + worker *ateapipb.Worker + // resUtil is the worker's dominant compute-resource utilization in [0,1], + // or -1 if not yet computed. + resUtil float64 +} + +func (c *candidate) resourceUtilization() float64 { + if c.resUtil < 0 { + c.resUtil = workerResourceUtilization(c.worker) + } + return c.resUtil +} + +// lessLoaded compares dominant utilization — the higher of actor-slot +// utilization (allocated/capacity) and compute-resource utilization — so either +// dimension can mark a worker as hot. Ties fall back to actor-slot utilization, +// then to remaining actor slots. +func lessLoaded(a, b *candidate) bool { + aAlloc := int64(a.worker.GetStatus().GetAllocated().GetActors()) + bAlloc := int64(b.worker.GetStatus().GetAllocated().GetActors()) + // checkRoom admits only workers with allocated < capacity, so capacity >= 1. + aCap := int64(a.worker.GetStatus().GetCapacity().GetActors()) + bCap := int64(b.worker.GetStatus().GetCapacity().GetActors()) + + aLoad := max(float64(aAlloc)/float64(aCap), a.resourceUtilization()) + bLoad := max(float64(bAlloc)/float64(bCap), b.resourceUtilization()) + if aLoad != bLoad { + return aLoad < bLoad + } + + // Compare slot utilization exactly via cross-multiplication. + if lhs, rhs := aAlloc*bCap, bAlloc*aCap; lhs != rhs { + return lhs < rhs + } + + return aCap-aAlloc > bCap-bAlloc +} + +// workerResourceUtilization returns the highest allocated/capacity ratio across +// compute dimensions, or 1 when capacity is unreported or quantities fail to parse. +func workerResourceUtilization(w *ateapipb.Worker) float64 { + capQ, err := resources.ParseQuantities(w.GetStatus().GetCapacity().GetResources()) + if err != nil || len(capQ) == 0 { + return 1 + } + usedQ, err := resources.ParseQuantities(w.GetStatus().GetAllocated().GetResources()) + if err != nil { + return 1 + } + return quantitiesUtilization(capQ, usedQ) +} + +func quantitiesUtilization(capQ, usedQ resources.Quantities) float64 { + var ( + maxRatio float64 + hasPositive bool + ) + for name, capVal := range capQ { + if capFloat := capVal.AsApproximateFloat64(); capFloat > 0 { + hasPositive = true + if usedVal, ok := usedQ[name]; ok { + maxRatio = max(maxRatio, usedVal.AsApproximateFloat64()/capFloat) + } + } + } + if !hasPositive { + return 1 + } + return maxRatio } func (s *scheduler) Applies(worker *ateapipb.Worker, constraints Constraints) bool { @@ -145,30 +240,42 @@ func (s *scheduler) Applies(worker *ateapipb.Worker, constraints Constraints) bo // having no room: it is the only answer that cannot overcommit a worker whose // true occupancy is unreadable. func (s *scheduler) HasRoom(worker *ateapipb.Worker, constraints Constraints) bool { + want, err := resources.ParseQuantities(constraints.Limits) + if err != nil { + return false + } + _, ok := checkRoom(worker, want) + return ok +} + +func checkRoom(worker *ateapipb.Worker, want resources.Quantities) (candidate, bool) { capacity := worker.GetStatus().GetCapacity() used := worker.GetStatus().GetAllocated() // No per-actor size to compare: every assignment costs one, so a worker at // its limit has no room however small the next actor is. if used.GetActors() >= capacity.GetActors() { - return false + return candidate{}, false } - want, err := resources.ParseQuantities(constraints.Limits) - if err != nil || len(want) == 0 { - return err == nil + if len(want) == 0 { + return candidate{worker: worker, resUtil: -1}, true } free, err := resources.ParseQuantities(capacity.GetResources()) if err != nil { - return false + return candidate{}, false } if free == nil { free = resources.Quantities{} } allocated, err := resources.ParseQuantities(used.GetResources()) if err != nil { - return false + return candidate{}, false } + resUtil := quantitiesUtilization(free, allocated) free.Sub(allocated) - return free.Covers(want) + if !free.Covers(want) { + return candidate{}, false + } + return candidate{worker: worker, resUtil: resUtil}, true } diff --git a/cmd/ateapi/internal/scheduling/scheduling_test.go b/cmd/ateapi/internal/scheduling/scheduling_test.go index 7b1e58612..9cbc5e6ff 100644 --- a/cmd/ateapi/internal/scheduling/scheduling_test.go +++ b/cmd/ateapi/internal/scheduling/scheduling_test.go @@ -35,9 +35,12 @@ func TestSchedule(t *testing.T) { wantPod string // "" means ErrNoCapacity expected }{ { + // Two full workers precede w-free so pair sampling cannot reach w-free + // unless full workers are filtered out. name: "picks the only eligible free worker", fleet: fleet{ - worker("w-busy", "gvisor", "node-a", tierTwo, assigned("demo", "other")), + worker("w-busy-1", "gvisor", "node-a", tierTwo, assigned("demo", "other-1")), + worker("w-busy-2", "gvisor", "node-a", tierTwo, assigned("demo", "other-2")), worker("w-free", "gvisor", "node-a", tierTwo), }, constraints: Constraints{SandboxClass: "gvisor"}, @@ -219,6 +222,109 @@ func TestSchedule(t *testing.T) { constraints: Constraints{SandboxClass: "gvisor", Limits: resources.CPUMemory(1000, 4<<30)}, wantPod: "w-half", }, + { + name: "prefers less-loaded worker when first sample is busier", + fleet: fleet{ + worker("w-busy", "gvisor", "node-a", tierTwo, withMaxActors(4), assigned("demo", "a"), assigned("demo", "b")), + worker("w-idle", "gvisor", "node-b", tierTwo, withMaxActors(4)), + }, + constraints: Constraints{SandboxClass: "gvisor"}, + wantPod: "w-idle", + }, + { + name: "keeps first sample when it is already less loaded than second", + fleet: fleet{ + worker("w-idle", "gvisor", "node-a", tierTwo, withMaxActors(4)), + worker("w-busy", "gvisor", "node-b", tierTwo, withMaxActors(4), assigned("demo", "a")), + }, + constraints: Constraints{SandboxClass: "gvisor"}, + wantPod: "w-idle", + }, + { + name: "compares actor utilization across heterogeneous worker capacities", + fleet: fleet{ + // 1/2 (50%) vs 2/16 (12.5%): w-large has more actors but lower utilization. + worker("w-small", "gvisor", "node-a", tierTwo, withMaxActors(2), assigned("demo", "a")), + worker("w-large", "gvisor", "node-b", tierTwo, withMaxActors(16), assigned("demo", "b"), assigned("demo", "c")), + }, + constraints: Constraints{SandboxClass: "gvisor"}, + wantPod: "w-large", + }, + { + name: "prefers lower compute utilization at equal actor utilization", + fleet: fleet{ + worker("w-heavy", "gvisor", "node-a", tierTwo, withCapacity(4000, 8<<30), withMaxActors(4), + assignedFor("demo", "a", resources.CPUMemory(1000, 6<<30))), + worker("w-light", "gvisor", "node-b", tierTwo, withCapacity(4000, 8<<30), withMaxActors(4), + assignedFor("demo", "b", resources.CPUMemory(1000, 1<<30))), + }, + constraints: Constraints{SandboxClass: "gvisor", Limits: resources.CPUMemory(500, 1<<30)}, + wantPod: "w-light", + }, + { + // With a uniform actor ceiling, slot utilization alone would pick + // w-hot (5/1000 < 6/1000) even though its memory is 95% used. + name: "compute hotspot outweighs one fewer actor", + fleet: fleet{ + worker("w-hot", "gvisor", "node-a", tierTwo, withCapacity(4000, 20<<30), withMaxActors(1000), + assignedFor("demo", "a1", resources.CPUMemory(100, 19<<30)), + assigned("demo", "a2"), assigned("demo", "a3"), assigned("demo", "a4"), assigned("demo", "a5")), + worker("w-cool", "gvisor", "node-b", tierTwo, withCapacity(4000, 20<<30), withMaxActors(1000), + assignedFor("demo", "b1", resources.CPUMemory(100, 2<<30)), + assigned("demo", "b2"), assigned("demo", "b3"), assigned("demo", "b4"), assigned("demo", "b5"), + assigned("demo", "b6")), + }, + constraints: Constraints{SandboxClass: "gvisor", Limits: resources.CPUMemory(100, 512<<20)}, + wantPod: "w-cool", + }, + { + name: "prefers worker with known resource utilization over unreported capacity", + fleet: fleet{ + worker("w-unreported", "gvisor", "node-a", tierTwo, withMaxActors(4), assigned("demo", "a")), + worker("w-known", "gvisor", "node-b", tierTwo, withCapacity(4000, 8<<30), withMaxActors(4), + assignedFor("demo", "b", resources.CPUMemory(1000, 2<<30))), + }, + constraints: Constraints{SandboxClass: "gvisor"}, + wantPod: "w-known", + }, + { + name: "prefers worker with known resource utilization over malformed allocated quantities", + fleet: fleet{ + worker("w-malformed", "gvisor", "node-a", tierTwo, withCapacity(4000, 8<<30), withMaxActors(4), + func(w *ateapipb.Worker) { + w.Status.Allocated = &ateapipb.WorkerResources{ + Actors: 1, + Resources: &ateapipb.Resources{Limits: []*ateapipb.Limits{{Name: resources.ResourceCPU, Quantity: "not-a-quantity"}}}, + } + }), + worker("w-known", "gvisor", "node-b", tierTwo, withCapacity(4000, 8<<30), withMaxActors(4), + assignedFor("demo", "b", resources.CPUMemory(1000, 2<<30))), + }, + constraints: Constraints{SandboxClass: "gvisor"}, + wantPod: "w-known", + }, + { + name: "breaks utilization ties using remaining actor slots", + fleet: fleet{ + // Both workers are at 50% actor-slot utilization (1/2 vs 4/8) and + // equal compute utilization, so w-large wins with 4 free slots vs 1. + worker("w-small", "gvisor", "node-a", tierTwo, withMaxActors(2), assigned("demo", "a")), + worker("w-large", "gvisor", "node-b", tierTwo, withMaxActors(8), + assigned("demo", "b"), assigned("demo", "c"), assigned("demo", "d"), assigned("demo", "e")), + }, + constraints: Constraints{SandboxClass: "gvisor"}, + wantPod: "w-large", + }, + { + name: "skips idle worker that lacks resource room and picks busier worker with room", + fleet: fleet{ + worker("w-idle-small", "gvisor", "node-a", tierTwo, withCapacity(1000, 1<<30), withMaxActors(4)), + worker("w-busy-big", "gvisor", "node-b", tierTwo, withCapacity(4000, 8<<30), withMaxActors(4), + assignedFor("demo", "a", resources.CPUMemory(1000, 2<<30))), + }, + constraints: Constraints{SandboxClass: "gvisor", Limits: resources.CPUMemory(1000, 2<<30)}, + wantPod: "w-busy-big", + }, } for _, tc := range tests { @@ -414,5 +520,78 @@ func TestAppliesIgnoresRoom(t *testing.T) { } } -// firstIntn always picks the first candidate, making Schedule deterministic. +// firstIntn always returns 0, sampling candidates[0] and candidates[1] when >=2 exist. func firstIntn(int) int { return 0 } + +func TestSchedulePowerOfTwoChoicesTieBreaking(t *testing.T) { + f := fleet{ + worker("w-0", "gvisor", "node-a", nil, withMaxActors(4), assigned("demo", "a")), + worker("w-1", "gvisor", "node-b", nil, withMaxActors(4), assigned("demo", "b")), + worker("w-2", "gvisor", "node-c", nil, withMaxActors(4), assigned("demo", "c")), + } + constraints := Constraints{SandboxClass: "gvisor"} + + // On equal utilization, Schedule must preserve the first sampled candidate so + // every tied worker remains reachable via the random source. + for wantIdx, wantPod := range []string{"w-0", "w-1", "w-2"} { + calls := 0 + s := New(f, WithIntn(func(n int) int { + calls++ + if calls == 1 { + return wantIdx + } + return 0 + })) + got, err := s.Schedule(context.Background(), constraints) + if err != nil { + t.Fatalf("Schedule() for index %d error = %v", wantIdx, err) + } + if got.GetWorkerPod() != wantPod { + t.Fatalf("Schedule() for index %d = %q, want %q", wantIdx, got.GetWorkerPod(), wantPod) + } + } +} + +// seqIntn returns vals in order, one per call. +func seqIntn(vals ...int) func(int) int { + i := 0 + return func(int) int { + v := vals[i] + i++ + return v + } +} + +func TestSchedulePowerOfTwoChoicesShiftedSecondSample(t *testing.T) { + // intn yields i=0, then j=1; since j >= i the second sample shifts to + // candidates[2], the only lighter worker. An off-by-one in the shift would + // compare w-0 against w-1 and return a busy worker. + f := fleet{ + worker("w-0", "gvisor", "node-a", nil, withMaxActors(4), assigned("demo", "a"), assigned("demo", "b")), + worker("w-1", "gvisor", "node-b", nil, withMaxActors(4), assigned("demo", "c"), assigned("demo", "d")), + worker("w-2", "gvisor", "node-c", nil, withMaxActors(4)), + } + s := New(f, WithIntn(seqIntn(0, 1))) + got, err := s.Schedule(context.Background(), Constraints{SandboxClass: "gvisor"}) + if err != nil { + t.Fatalf("Schedule() error = %v", err) + } + if got.GetWorkerPod() != "w-2" { + t.Fatalf("Schedule() = %q, want %q", got.GetWorkerPod(), "w-2") + } +} + +func TestScheduleMalformedLimitsIsNotNoCapacity(t *testing.T) { + f := fleet{worker("w-1", "gvisor", "node-a", nil)} + constraints := Constraints{ + SandboxClass: "gvisor", + Limits: &ateapipb.Resources{Limits: []*ateapipb.Limits{{Name: resources.ResourceCPU, Quantity: "not-a-quantity"}}}, + } + _, err := New(f, WithIntn(firstIntn)).Schedule(context.Background(), constraints) + if err == nil { + t.Fatal("Schedule() error = nil, want a parse error") + } + if errors.Is(err, ErrNoCapacity) { + t.Fatalf("Schedule() error = %v, want an error other than ErrNoCapacity", err) + } +}