Skip to content

[iris] Cross-variant preemption on fungible TPU reservations#6531

Closed
rjpower wants to merge 20 commits into
mainfrom
weaver/iris-autoscaler-preemption
Closed

[iris] Cross-variant preemption on fungible TPU reservations#6531
rjpower wants to merge 20 commits into
mainfrom
weaver/iris-autoscaler-preemption

Conversation

@rjpower

@rjpower rjpower commented Jun 21, 2026

Copy link
Copy Markdown
Collaborator

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
DRAINING lifecycle, and a single ReservationLedger, 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 --> State
Loading

free_chips is launchable now (excludes draining); incoming_chips is
free + 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"]
Loading

The two skip branches replace the old 5-minute cooldown: direct observations of
in-flight recovery rather than a guessed duration.

One teardown lifecycle

DRAINING is not preemption-specific. A slice now leaves tracking only at the
reap
— 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 note
Loading

drain_slice() (mark DRAINING + terminate) is the single begin-teardown for
every live slice; refresh reaps it via detach_slice() once the cloud reports
it 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 SliceDrainCause enum, not a
boolean. Slices the cloud already reports dead (boot-FAILED, unresolvable
UNKNOWN) 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_slices runs only at controller boot), so
a 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's remove-callable indirection and two near-identical
functions become one drain_slices_for_workers, and ScalingGroup.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 involved
Loading

Key code

The pass skip — two windows, both observed (policy.py):

if incoming[pool] >= need:            # free chips, or chips a drain is freeing
    incoming[pool] -= need
    continue
if pool_replacement[variant] >= 1:    # a replacement of the variant is booting
    pool_replacement[variant] -= 1
    continue
# otherwise evict minimal lowest-band victims to cover the deficit

The reap is clean — intentional, so no failure signal (runtime.py):

gone = status.state == CloudSliceState.FAILED or not status.workers
if gone:
    group.detach_slice(slice_id)      # no mark_slice_failed, no backoff fed
    self._unregister_slice_workers(slice_id)
    ...

Scope

Band-ordered launch admission and DemandEntry.band are retained. The
task-binding claim machine (provisional claimed_by slices, which fixes the
same-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.

rjpower added 2 commits June 21, 2026 16:49
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.
@rjpower rjpower added the agent-generated Created by automation/agent label Jun 21, 2026
@claude

claude Bot commented Jun 21, 2026

Copy link
Copy Markdown
Contributor

Claude finished @rjpower's task in 11m 58s —— View job


Code review

No 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:

  • Chip-ledger "over-eviction" (false positive). A reviewer suspected that same-variant evictions in a fungible pool aren't credited back into the free ledger the cross-variant pass consumes (policy.py:610), causing over-eviction. Tracing it through: at scheduling-decision time the victim is still counted as consumed and the same-variant preemptor's replacement slice is not yet a pending scale-up, so those chips are reserved for the same-variant preemptor and genuinely unavailable to a cross-variant preemptor. Not crediting them is correct; no over-eviction.
  • Test of a private _order_reserved_demand_by_band (not high-signal). test_reserved_preemption.py imports a module-level _-prefixed function. The Iris TESTING.md rule targets private attributes/state, and the repo already has an established pattern of unit-testing private module functions (e.g. test_scaling_group.py, test_provisioning.py). A senior reviewer wouldn't block on this.

Things I specifically verified as correct:

  • Chip math (reserved_pool.py): chip_count = half the vN-SIZE suffix; consumed = slice_count() * chips_per_slice covers live + in-flight; free = budget - consumed. Conflicting per-pool budgets raise.
  • Eviction selection (policy.py): band gate mirrors the same-variant pass (BATCH never preempts; victims must be strictly lower-band); lowest-band-then-smallest ordering; no-partial-eviction guard (freed < deficit → continue); claimed dedups victims and handled_preemptor_jobs dedups coscheduled preemptors; free[pool] += freed - need correctly prevents a second preemptor reusing the same freed chips.
  • Drain path (operations.py/runtime.py): intentional drain skips the failure/backoff feed (feed_backoff=False), stamps a per-pool cooldown that expires by elapsed-time comparison.
  • Post-commit fail-worker (controller.py): ordering is autoscale-drain → commit (victims → PENDING) → remove drained workers; the fail path re-checks liveness, so a victim whose attempt already closed isn't re-charged or re-terminated.
  • Config validation (config.py): rejects reservation_chips without fungible_reservation/non-reserved capacity/non-positive budgets, and requires the largest configured slice to fit (allowing exact equality).
  • Proto regen: both config_pb2.py and config_pb2.pyi are present and consistent with the .proto change (field 82 reservation_chips). No new persisted schema/migration — matches the description.
  • Tests: the new TestAutoscalerDrainSlices and reserved-pool/preemption tests are behavior-focused (drain tears down + returns siblings, drain doesn't decay detector health, cooldown stamp/expiry) and use public APIs.

One non-blocking observation (not a bug): run_reserved_pool_preemption dedups preemptors by parent job every tick, so a non-coscheduled job with several independent pending tasks each needing its own slice only gets one slice's chips freed per tick. This is conservative (under-preempts) and recovers on later ticks, so it's fine — just noting in case more aggressive multi-slice freeing was intended.

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 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".

Comment thread lib/iris/src/iris/cluster/controller/controller.py Outdated
Comment thread lib/iris/src/iris/cluster/controller/scheduling/policy.py Outdated
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).
@rjpower

rjpower commented Jun 21, 2026

Copy link
Copy Markdown
Collaborator Author

🤖 Thanks — both P1s were real. Fixed in d2ec08a.

Schedule-only tick committing preemptions without a drain. Confirmed: a submit wake runs run_schedule without run_autoscale, and _commit_tick finalizes the schedule's preemptions regardless, so a cross-variant PREEMPT there would strand the victim (PENDING, slice never torn down). Fix: the cross-variant pass is now gated on the autoscaler running this tick. ScheduleInput carries autoscale_runs (set from inputs.run_autoscale), and RpcTaskBackend.schedule builds the reserved view only when it's set — so the victim PREEMPT and its drain always commit on the same tick. Same-variant preemption is unaffected (it needs no drain). Regression test: test_reserved_view_only_built_when_autoscale_runs.

Victim grouping by task vs. real slice. Confirmed. Victims now group by the physical slice id the autoscaler reports (ReservedPoolView.worker_slice), not by task id: the slice's chips are counted once, eligibility gates on the highest-priority member (so a sibling task that's equal/higher band protects the whole slice from eviction), and every task on the slice is preempted. A coscheduled job spanning multiple slices is excluded, since its cross-slice cascade can't be drained coherently. Regression tests: test_slice_with_a_higher_band_task_is_not_evicted, test_slice_chips_counted_once_for_colocated_tasks.

rjpower added 10 commits June 21, 2026 17:47
_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.
@rjpower rjpower requested a review from Helw150 June 22, 2026 16:43
Comment thread lib/iris/src/iris/cluster/controller/scheduling/policy.py
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.
Comment thread lib/iris/src/iris/cluster/controller/autoscaler/planning.py Outdated
Comment thread lib/iris/tests/cluster/backends/test_rpc_backend.py Outdated
rjpower added 6 commits June 30, 2026 22:59
…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.
@rjpower

rjpower commented Jul 1, 2026

Copy link
Copy Markdown
Collaborator Author

🤖 Splitting this into two stacked PRs for a cleaner, lower-risk review: #6818 (reservation ledger + unified drain lifecycle, no behavior change to preemption) and #6819 (the actual cross-variant preemption pass, stacked on #6818). Closing this one in favor of those.

@rjpower rjpower closed this Jul 1, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

agent-generated Created by automation/agent

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Iris scheduler] Higher Queues Should Preempt Across Node Types for Reserved Compute

2 participants