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
2026-11-10 10:13:50 +08:00 · 2026-10-02 23:56:47 +00:00
parent b748aa2a5e
commit a25a4fbcfe
2 changed files with 302 additions and 16 deletions
@@ -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
 }
@@ -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)
+	}
+}