[iris] Cross-variant preemption on fungible TPU reservations#6531
[iris] Cross-variant preemption on fungible TPU reservations#6531rjpower wants to merge 20 commits into
Conversation
On a reserved TPU pool the chips are fungible across slice sizes: freeing a v4-8 makes room for part of a v4-32 from the same reservation. The scheduler only ran same-variant preemption, so a low-priority job on slice size X could indefinitely block a high-priority job needing size Y even when the reservation had the chips. (#6217) Model a fungible reservation as the per-size scale groups that share a quota_pool, tagged with reservation_chips (the physical chip budget, half the GCP vN-SIZE suffix). A read-only chip ledger (reserved_pool.py) reports, per pool, how many chips live and in-flight slices consume and how many remain free. The scheduler gains a cross-variant pass that runs after the same-variant pass, over reserved pools only: a higher-band preemptor on a full pool evicts the minimal set of lowest-band victim slices of any variant in the same pool to cover its chip deficit, skipping victims the same-variant pass already claimed. It emits the victim PREEMPTs plus the worker ids whose slices must drain. The drain decision is re-derived from committed task state each tick, so it is restart-safe and needs no new persisted state. The autoscaler reaps the named workers via an intentional-drain teardown that does not feed the churn detector (a drain is not a preemption signal), and a per-pool cooldown keeps it from re-preempting while the replacement slice is still booting. Demand for a reserved pool is ordered by band every tick so the autoscaler provisions the higher-priority job's replacement before the just-evicted victim re-grabs the chips its own eviction freed. The controller fails drained workers out of the registry post-commit, mirroring the dead-worker teardown, so they stop being schedulable immediately.
Rewrite the drain docstrings to state the caller-facing contract instead of cross-referencing the sibling terminate path and narrating the churn-detector internals. Collapse the four reservation-config rejection tests, which shared an identical write-then-raise body, into one parametrized test.
|
Claude finished @rjpower's task in 11m 58s —— View job Code reviewNo blocking issues found. Checked for bugs and CLAUDE.md/AGENTS.md/TESTING.md compliance across the changed files. I ran a multi-agent pass (2 independent bug/logic reviewers + 2 compliance/test reviewers) and validated every candidate finding against the actual code. The two flagged candidates did not survive validation:
Things I specifically verified as correct:
One non-blocking observation (not a bug): |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f38081bd74
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
Two correctness fixes for cross-variant reserved-pool preemption. A schedule-only mini-tick (submit wake) runs the scheduler but not the autoscaler, yet still commits the schedule's preemptions. Cross-variant preemption emitted there would finalize victims to PENDING with no slice drain to reclaim their reserved chips, so the preemptor never gets capacity. Gate the reserved view on the autoscaler running this tick: schedule passes autoscale_runs through ScheduleInput, and the backend builds the reserved view only when it is set, so the victim PREEMPT and its drain always commit together. Group reserved-pool victims by physical slice instead of by task. The drain tears down a whole slice, so two independent tasks colocated on one multi-VM slice were each counted as freeing the full slice (double counting), and selecting only the lower-band one would still drain the sibling worker, killing an unchecked — possibly higher-band — task. Victims now group on the slice id the autoscaler reports, gate on the highest-priority member, free the slice's chips once, and preempt every task on it. A coscheduled job spanning multiple slices is excluded (its cross-slice cascade can't be drained coherently).
|
🤖 Thanks — both P1s were real. Fixed in d2ec08a. Schedule-only tick committing preemptions without a drain. Confirmed: a submit wake runs Victim grouping by task vs. real slice. Confirmed. Victims now group by the physical slice id the autoscaler reports ( |
_TickInputs carries no run_autoscale field; the schedule phase read it off the inputs and crashed every control tick (AttributeError), which the integration and smoke jobs caught. Thread run_autoscale from _control_tick into _schedule_phase as an explicit keyword argument instead.
…l launches by chips route_demand is reservation-blind: each per-size group of a fungible reservation pool is bounded only by its own static max_slices, so unconstrained planning can request far more chips than the reservation holds (256 v4-8 AND 128 v4-16 = 2048 chips against a 1024-chip pool), and a just-drained preemption victim re-grabs the freed chips before the higher-priority preemptor can use them. The previous demand reorder could not fix this — ordering a skip-and-continue loop over independent per-group budgets imposes no shared cap. Replace the reorder with a real per-quota_pool chip cap applied after per-group planning: trim each fungible pool's total slices_to_add * chips to the chips free against the reservation, admitting groups in band order with head-of-line blocking so an unsatisfied high-priority slice holds freed chips (e.g. accumulating across a multi-tick drain) instead of a lower-priority slice re-grabbing them. The scheduler stamps each demand entry's effective band so the autoscaler can order them. Fixes #6217
…tring cross-refs Make the fungible-pool admission helper take a named _PoolCandidate record instead of a positionally-accessed 4-tuple, and replace the :func:/:class: docstring cross-references in reserved_pool and scaling_group with plain prose describing the contract at hand.
Replace the implicit machinery behind cross-variant reserved-pool preemption with one legible per-tick capacity object and an explicit slice state, so the scheduler's preemption pass and the autoscaler's launch planner reason from the same view and a drain is visible rather than inferred. ReservationLedger (reserved_pool.py) supersedes ReservedPoolView: per fungible pool it buckets chips into live / in-flight / draining and exposes free_chips (launchable now) and incoming_chips (free + being-freed). The planner caps launches against free_chips; the pass asks whether a deficit is already being addressed. DRAINING is now a real slice lifecycle. A cross-variant victim is marked DRAINING and kept counted (its chips show as draining, not free) instead of being detached immediately; refresh reaps it once the cloud reports its VMs gone, cleanly — no FAILED outcome, no churn-detector decay — with a timeout fallback for a stuck delete. This lets the 5-minute RESERVED_DRAIN_COOLDOWN timer go away. The pass no longer re-preempts a pool whose deficit is already covered by free-or-draining chips (the drain window) or by a replacement slice of the preemptor's own variant already booting (the reprovision window) — a direct observation of in-flight recovery rather than a guessed duration. Band-ordered launch admission and DemandEntry.band are kept; the task-binding claim machine is deferred. The dashboard surfaces DRAINING as its own slice badge so an in-flight preemption is visible.
Extract the victim-slice grouping out of run_reserved_pool_preemption into _victim_slices_by_pool so the pass reads as three named stages (build victims, then match preemptors against ledger chips). Drop the unused handle argument from _fold_draining_slice and trim the reservation_ledger docstring's duplicated consumer narration.
Collapse the autoscaler's slice-removal paths onto a single delete -> drain -> cleanup lifecycle: a live slice is marked DRAINING and its VMs terminated, then refresh reaps it once the cloud confirms they are gone (or re-terminates past a timeout). A slice leaves tracking only at the reap, never before the cloud confirms its allocation is gone. The failure-vs-drain asymmetry in operations.py collapses into one drain_slices_for_workers parameterized by a SliceDrainCause enum; ScalingGroup gains drain_slice and loses scale_down (detach-then-terminate). Idle scale-down and worker-liveness/zombie teardown now drain like cross-variant preemption does. Slices the cloud already reports dead (boot-FAILED, unresolvable UNKNOWN) detach directly -- the cleanup step, not a drain. Tracking every terminating slice until the cloud confirms it is gone closes the orphan gap: terminate was fire-and-forget with no steady-state reconciliation, so a lost delete left a live VM pinning reserved chips until the next controller restart.
_fold_draining_slice reaps every drain cause now, not just preemption; detach_slice's docstring described its callers rather than its contract; and autoscale()'s docstring still said dead-worker slices are "terminated" when they drain like the rest.
The two new reserved-pool test files used class-based grouping (pure namespacing with no shared fixtures), against the repo's preference for top-level test functions. Flatten them; this also removes the lone test class that sat among flat functions in test_reserved_preemption.py. Tests added to already-class-organized files keep their file's convention.
Lint-review findings: _VictimSlice.pool was set at construction but never read (victims are already grouped in victims_by_pool keyed by pool). Two docstrings narrated their callers or cross-referenced a sibling pass rather than stating their own contract.
Merge test_reserved_preemption.py into test_reserved_pool.py so the fungible pool's ledger, launch cap, and cross-variant preemption tests live together as flat top-level functions.
A coscheduled job's host candidates were only deduped on the eviction- success path, so a multi-host preemptor's later host candidates could re-consume the pool's incoming/replacement capacity already spent on the same job's first candidate, causing unnecessary victim evictions. Mark the preemptor job handled as soon as its first candidate is processed, regardless of which branch resolves it.
…r-preemption # Conflicts: # lib/iris/src/iris/cluster/backends/k8s/tasks.py # lib/iris/src/iris/cluster/backends/rpc/backend.py # lib/iris/src/iris/cluster/controller/backend.py # lib/iris/src/iris/cluster/controller/controller.py # lib/iris/src/iris/cluster/controller/scheduling/decision.py # lib/iris/src/iris/rpc/config.proto # lib/iris/src/iris/rpc/config_pb2.py # lib/iris/src/iris/rpc/config_pb2.pyi # lib/iris/tests/cluster/backends/test_config.py # lib/iris/tests/cluster/backends/test_scaling_group.py # lib/iris/tests/cluster/controller/conftest.py # lib/iris/tests/cluster/controller/test_autoscaler.py # lib/iris/tests/cluster/controller/test_direct_controller.py # lib/iris/tests/cluster/controller/test_reconcile.py
…) drain branches ScheduleInput.autoscale_runs was never read downstream — the actual cross-variant-preemption gate is the ledger built (or not) before ScheduleInput is constructed. RpcTaskBackend.autoscale's dead_workers and drain_workers branches were identical except for the worker list and drain cause; collapse them and assert the two never co-occur instead of silently letting dead_workers take precedence.
Two simplifications to the fungible reservation pool logic, per review on PR #6531: - Each size's max_slices now defaults to reservation_chips // chips_per_slice (the most that size could ever place against the shared budget) instead of requiring a hand-computed number per size in YAML. An override may only tighten the default, never loosen it — a looser value is rejected at load time. marin.yaml's v4-reserved pool drops its 9 hand-computed values. - _admit_in_band_order no longer holds a pool's remaining chips for an unsatisfied high-priority band unconditionally. It only holds them when relief is actually achievable (free chips plus what's currently draining would cover the band's next slice); otherwise a lower-priority band that fits in what's free now is admitted instead of leaving chips idle for a drain that could never satisfy the high band anyway. - test_rpc_backend.py's call-count test is replaced with one that drives a real scheduling tick (real ControllerDB, ScalingGroups, and a dispatched task) and asserts on the observable ScheduleResult.reserved_drain_workers.
Flagged by the lint-catalog review (ml-impl-narration-docstring): the docstring listed drain_slice's callers and cross-referenced detach_slice by name instead of describing its own behavior.
main (#6799) landed while this branch was in flight and made WorkerAttrsProjection scoped per-backend (owns_scale_group required, and RpcTaskBackend now builds its own from BackendRuntime.db + owns_scale_group instead of taking one via BackendRuntime.worker_attrs). Update the new reserved-pool scheduling test for the merged interface.
Problem
On a fungible reserved TPU pool the per-size scale groups share one physical chip
reservation, but the scheduler and autoscaler never reasoned about that shared
budget together. A low-priority job on slice size X could block a high-priority
job needing size Y, and the cross-variant preemption machinery that did exist
was implicit and hard to follow: a 5-minute cooldown timer, a band head-of-line
launch cap, and an immediate slice detach — three proxies for "a drain is in
flight," none of them legible.
Fixes #6217.
How scheduling, autoscaling & preemption work now
The scheduler and autoscaler stay pure functions of the live slice state. They
keep no cross-tick drain memory; the only durable drain state is the slice's own
DRAININGlifecycle, and a singleReservationLedger, recomputed every tick,is the one capacity view both read.
1. One shared ledger feeds both phases
The control tick is single-threaded:
schedule → reconcile → autoscale → commit.flowchart LR State["live slice state<br/>(per scale group)"] --> Ledger["ReservationLedger<br/>per fungible pool:<br/>live · inflight · draining chips"] subgraph Schedule["schedule phase — pure scheduler"] Pass["cross-variant<br/>preemption pass"] end subgraph Autoscale["autoscale phase — autoscaler"] Plan["launch planner<br/>(caps per-pool launches)"] Drain["mark victims DRAINING<br/>+ terminate their VMs"] Reap["refresh reaps<br/>gone DRAINING slices"] end Ledger -->|"incoming_chips · inflight_slices"| Pass Ledger -->|"free_chips"| Plan Pass -->|"PREEMPT pairs + drain_workers<br/>(only when autoscale runs this tick)"| Drain Drain --> State Reap --> State Plan --> Statefree_chipsis launchable now (excludes draining);incoming_chipsisfree + draining. That split is the whole trick.2. The preemption decision — observe in-flight recovery, no timer
flowchart TD Start["pending high-priority task needs<br/>a slice of variant V on a full reserved pool"] Start --> Q1{"free + draining chips ≥ need?"} Q1 -->|yes| Skip1["skip — free now,<br/>or a drain is already freeing it"] Q1 -->|no| Q2{"a replacement slice of V<br/>already booting?"} Q2 -->|yes| Skip2["skip — reprovision already in flight"] Q2 -->|no| Q3{"enough lower-band victim slices<br/>to cover the deficit?"} Q3 -->|no| NoOp["evict nothing<br/>(never partial-evict)"] Q3 -->|yes| Evict["mark minimal lowest-band victim<br/>slices DRAINING + emit PREEMPT pairs"]The two skip branches replace the old 5-minute cooldown: direct observations of
in-flight recovery rather than a guessed duration.
One teardown lifecycle
DRAININGis not preemption-specific. A slice now leaves tracking only at thereap — never before the cloud confirms its VMs are gone.
stateDiagram-v2 [*] --> REQUESTING REQUESTING --> BOOTING BOOTING --> INITIALIZING INITIALIZING --> READY READY --> DRAINING: drain_slice() — any teardown of a live slice DRAINING --> [*]: refresh reaps once VMs gone (detach_slice) REQUESTING --> [*]: boot-FAILED or unresolvable — detach directly note right of DRAINING Reached by idle scale-down, worker-failure, probe-zombie, and cross-variant preemption alike. Still counts against the reservation pool until reaped; the reap is clean (no FAILED outcome, no churn decay). A stuck delete is force-reaped on a 15-minute timeout. end notedrain_slice()(mark DRAINING + terminate) is the single begin-teardown forevery live slice;
refreshreaps it viadetach_slice()once the cloud reportsit gone. Idle scale-down, worker-liveness teardown, probe-health zombies, and
cross-variant preemption all route through it, differing only in whether the loss
is charged to the group's churn detector — a
SliceDrainCauseenum, not aboolean. Slices the cloud already reports dead (boot-
FAILED, unresolvableUNKNOWN) detach directly: that is the cleanup step, not a drain.This closes an orphan gap. Terminate was fire-and-forget and there is no
steady-state reconciliation (
list_all_slicesruns only at controller boot), soa lost delete left a live VM pinning reserved chips, invisible until the next
restart. Tracking every terminating slice until the cloud confirms it is gone —
plus the timeout re-terminate — is the universal guard. It also collapses the
code:
operations.py'sremove-callable indirection and two near-identicalfunctions become one
drain_slices_for_workers, andScalingGroup.scale_down(detach-then-terminate) is deleted.
End-to-end: preempt → drain → reprovision
The held chips stay reserved for the preemptor across the drain window (victim
DRAINING) and the reprovision window (replacement booting) — no cooldown.
sequenceDiagram participant Sch as Scheduler (pass) participant Led as ReservationLedger participant Auto as Autoscaler participant Cloud as Cloud reservation Note over Sch,Cloud: Pool full; high-priority job needs a bigger slice Sch->>Led: incoming_chips(pool) = 0 Sch->>Auto: PREEMPT victims + drain_workers Auto->>Auto: mark victim slices DRAINING Auto->>Cloud: terminate victim VMs Note over Sch,Cloud: Drain window — victim DRAINING, chips held Sch->>Led: incoming = free(0) + draining(N) ≥ need Note over Sch: skip (drain already covers the deficit) Cloud-->>Auto: victim VMs gone Auto->>Auto: refresh reaps DRAINING slice (clean) Auto->>Cloud: launch replacement once free ≥ need Note over Sch,Cloud: Reprovision window — replacement booting Sch->>Led: inflight_slices(pool, V) ≥ 1 Note over Sch: skip (replacement already booting) Cloud-->>Auto: replacement READY Note over Sch,Cloud: High-priority job places — no cooldown timer involvedKey code
The pass skip — two windows, both observed (
policy.py):The reap is clean — intentional, so no failure signal (
runtime.py):Scope
Band-ordered launch admission and
DemandEntry.bandare retained. Thetask-binding claim machine (provisional
claimed_byslices, which fixes thesame-band race where two equal-priority preemptors contend for freed chips, and
would let a replacement queue against a draining victim) is a deliberate
follow-up; this PR makes the existing machinery sound and legible.