Skip to content

fix(server): give preemption victim selection a total order - #1301

Merged
inureyes merged 3 commits into
mainfrom
fix/issue-1293-eviction-victim-total-order
Aug 22, 2026
Merged

fix(server): give preemption victim selection a total order#1301
inureyes merged 3 commits into
mainfrom
fix/issue-1293-eviction-victim-total-order

Conversation

@inureyes

Copy link
Copy Markdown
Member

Summary

BatchScheduler::select_eviction_victim decides which in-flight sequence gets preempted when the batch is full and a higher-priority request is waiting. Both policy arms selected through a key that was not total, so tied sequences were separated only by HashMap iteration order: LongestFirst used max_by_key(|seq| seq.generated_tokens.len()) and LowestPriority used min_by on priority then length, both over ActiveBatch::iter_sequences, which is HashMap::values. RandomState seeds each map instance, so the victim was not reproducible from the batch state.

Unlike the four latent sites fixed in #1291 / PR #1299, this tie is reachable through ordinary state. The LongestFirst key is generated_tokens.len(), a small integer that sequences admitted together and decoding in lockstep share routinely, and both PreemptionPolicy::LongestFirst and RequestPriority::Normal are the shipped defaults, so the reachable arm is the one that ships on. The harm is reproducibility rather than correctness: any tied-longest sequence does satisfy the stated policy, nothing is corrupted and nothing crashes, but two identically configured workers under identical load preempt different requests, and an operator investigating why a given request was preempted (and replayed with duplicate streamed tokens, per try_evict_for_preemption) cannot reproduce the choice.

What changed

  • src/server/batch/scheduler.rs: the policy is extracted into select_eviction_victim_from(sequences, policy), a pub(crate) free function, and BatchScheduler::select_eviction_victim becomes a one-line wrapper over it. The structured.is_none() filter moves inside the function so no caller can lose it. The free-function form was chosen over a method on ActiveBatch for the reason the issue gives: ActiveBatch currently knows nothing about PreemptionPolicy, and a method would make active.rs depend on crate::server::config.
  • Both arms now append seq_id as the component that makes the order total, using as_u64() because SequenceId derives only Debug, Clone, Copy, PartialEq, Eq, Hash and is not Ord. This is the same trap PromptCacheKeyDigest posed in refactor(server): give the four LRU eviction min_by_key calls a total order #1291, and it is resolved the same way rather than by widening the type's public trait surface for two call sites. LongestFirst keys on (generated_tokens.len(), Reverse(seq_id.as_u64())); LowestPriority appends .then_with(|| a.seq_id.as_u64().cmp(&b.seq_id.as_u64())). The two directions differ because max_by_key returns the LAST maximum while min_by returns the FIRST minimum, so opposite-facing tie components are needed to break the same direction. Each arm carries a comment naming seq_id as the total-order component so a later simplification pass does not strip it.
  • src/server/batch/scheduler_tests.rs: eviction_selects_longest_first_by_default and eviction_selects_lowest_priority_then_longest now call select_eviction_victim_from instead of reimplementing the selection expression. Neither expected victim changed, because neither test constructs a tie. Two determinism tests and an empty-batch test are added.
  • docs/code-guidelines.md: this instance is added to the list in the "HashMap Iteration Order" section, which becomes eight. The max_by_key correction at the "order-sensitive consumers" bullet and the enforcement note were already in place from PR docs(guidelines): correct the max_by_key tie direction and record the decline's cost #1294 and are not redone; the only enforcement-section edits are to keep the measured catch rates arithmetically consistent with an eight-item list, scoped to what docs(guidelines): record the HashMap-iteration-order defect class and decide whether a static check can catch it #1287 actually measured.

Tie-break direction: smallest seq_id, anti-starvation argument accepted

The issue asked for an explicit decision and it is smallest seq_id wins, with the anti-starvation argument accepted as stated.

SequenceId is a u64 newtype handed out by CachePool from an AtomicU64 advanced with fetch_add(1) per allocation, so smallest-id-wins reads as "evict the sequence admitted earliest", which needs no new field to be intelligible.

The decisive part is what preemption does to the id. try_evict_for_preemption reallocates its victim under a fresh, therefore higher, id, so a large id marks a sequence that was preempted recently, not one that arrived late. Smallest-id-wins therefore rotates preemption onto sequences that have not been hit yet, whereas largest-id-wins would make the request that was just replayed the preferred next victim.

SequenceInfo::created_at was considered and rejected. It survives preemption (it is absent from the reset block) and so gives true arrival order, but for exactly that reason it points the wrong way for starvation: smallest-created_at-wins would keep selecting the same oldest request, because nothing about being preempted moves the field. It is also an Instant and not total on its own, so seq_id would have had to sit behind it anyway.

Both arms resolve a tie to the same direction, so the choice does not vary by policy.

Regression test, and its failure against the unfixed code

Two tests, both built per docs/code-guidelines.md "Testing the fix": a fresh ActiveBatch inside each iteration (RandomState seeds per map instance, so probing one batch repeatedly measures nothing), 64 iterations, and an actual tie constructed on every axis the arm compares. Ids are inserted out of ascending order so a selector that returns the first- or last-inserted sequence cannot pass by coincidence. Each test records a mismatch per iteration rather than panicking on the first, so a failure reports the count out of 64.

The tests were written against the extracted-but-unfixed function and run before the tie-break was added:

running 5 tests
test server::batch::scheduler::tests::eviction_returns_none_for_an_empty_batch ... ok
test server::batch::scheduler::tests::eviction_selects_lowest_priority_then_longest ... ok
test server::batch::scheduler::tests::eviction_selects_longest_first_by_default ... ok
test server::batch::scheduler::tests::eviction_longest_first_tie_resolves_to_smallest_seq_id ... FAILED
test server::batch::scheduler::tests::eviction_lowest_priority_tie_resolves_to_smallest_seq_id ... FAILED

---- eviction_longest_first_tie_resolves_to_smallest_seq_id stdout ----
LongestFirst resolved a fully tied batch to something other than seq_id 2 on 58 of 64 freshly built batches (first Some((0, Some(5)))); the tie is falling through to HashMap order

---- eviction_lowest_priority_tie_resolves_to_smallest_seq_id stdout ----
LowestPriority resolved a batch tied on priority and length to something other than seq_id 2 on 39 of 64 freshly built batches (first Some((1, Some(14)))); the tie is falling through to HashMap order

test result: FAILED. 3 passed; 2 failed; 0 ignored; 0 measured; 5895 filtered out

So 58 of 64 for LongestFirst (first mismatch at iteration 0) and 39 of 64 for LowestPriority (first at iteration 1). Both are well under 100 percent, which is the reason the loop exists: a single-shot version of the LowestPriority test would have passed about 39 percent of the time. After the fix, all 64 iterations of both tests agree.

No git stash was used to produce this. The tests were run against the extraction commit before the tie-break was applied, and git stash list is empty.

Notes on the issue

Everything the issue asserted checked out against the code, with two points worth recording:

The structured.is_none() filter is now in the tested code path, but no test constructs a structured: Some(...) sequence: StructuredOutputConstraint has no public test constructor, and building one through build_json_schema_constraint needs a real HuggingFace tokenizer with a vocab plus an llguidance grammar compilation, which does not belong in a scheduler policy unit test. Covering the exclusion directly would need a test-only constructor on that type, which is out of scope here.

Test plan

  • cargo fmt --all -- --check (exit 0)
  • cargo check --lib --tests --features cuda (exit 0)
  • cargo clippy --lib --tests --features cuda -- -D warnings (exit 0)
  • cargo test --profile test-fast --features cuda --lib server::batch::scheduler::tests::eviction (exit 0, 5 passed, 0 failed)
  • cargo test --profile test-fast --features cuda --lib server::batch::scheduler (exit 0, 111 passed, 0 failed, 7 ignored)
  • cargo test --profile test-fast --features cuda --lib server::batch (exit 0, 373 passed, 0 failed, 7 ignored)
  • make verify-test-cuda is run separately against this branch.

Closes #1293

`BatchScheduler::select_eviction_victim` picked which in-flight sequence to preempt through a key that was not total, so both policy arms tie-broke on `HashMap` iteration order. `LongestFirst` used `max_by_key(|seq| seq.generated_tokens.len())` and `LowestPriority` used `min_by` on priority then length, both over `ActiveBatch::iter_sequences`, which is `HashMap::values`. `RandomState` seeds each map instance, so the victim was not reproducible from the batch state.

Unlike the four latent sites fixed in #1291, this tie is reachable through ordinary state. The `LongestFirst` key is `generated_tokens.len()`, a small integer that sequences admitted together and decoding in lockstep share routinely, and both `PreemptionPolicy::LongestFirst` and `RequestPriority::Normal` are the shipped defaults. Nothing was corrupted, since any tied-longest sequence does satisfy the policy, but two identically configured workers under identical load preempted different requests, and an operator could not reproduce why a given request was preempted and replayed with duplicate streamed tokens.

Both arms now append `seq_id` (via `as_u64()`, since `SequenceId` is not `Ord`), facing opposite ways because `max_by_key` returns the last maximum while `min_by` returns the first minimum. The direction is smallest-id-wins: `seq_id` is monotonic, and preemption reallocates its victim under a fresh higher id, so a large id marks a recently preempted sequence. Preferring the smallest id rotates preemption onto sequences that have not been hit yet, where largest-id-wins would keep re-picking the request that was just replayed. `created_at` survives preemption and would give true arrival order, but for that same reason it points the other way and would repeatedly select the same oldest request.

The policy is extracted into `select_eviction_victim_from`, a `pub(crate)` free function that both the scheduler method and the tests call, so there is exactly one copy. The two existing tests reimplemented the selection expression inline because `BatchScheduler` needs a real model to construct, and both copies had already drifted by losing the `structured.is_none()` guard; they now call the extracted function, which brings that guard under test. Two new tests build a fresh `ActiveBatch` per iteration over 64 iterations with an actual tie constructed on every axis. Against the unfixed logic they failed on 58 of 64 and 39 of 64 batches respectively.

Refs #1293
The "HashMap Iteration Order" section already discussed #1293 in its enforcement note, where PR #1294 recorded it as the finding both static-check candidates missed, but the catalogue itself still said seven. Add the entry so the list matches the note.

The entry states what the other seven do not: this is the only one whose tie is reachable through ordinary state rather than through a coarsening of the key, since the key is `generated_tokens.len()` and both `PreemptionPolicy::LongestFirst` and `RequestPriority::Normal` are the shipped defaults.

The two counted claims in the enforcement bullets are updated to stay arithmetically consistent with an eight-item list (0 of 8, 3 of 8), and the sentence describing what the prototype checks were run against is scoped to the five fixes that had landed when the measurement ran, so the #1287 record is not silently restated as covering something it did not measure. The `max_by_key` direction correction and the enforcement note itself are unchanged.

Refs #1293
@inureyes inureyes added status:review Under review type:bug Bug fixes, error corrections, or issue resolutions priority:medium Medium priority area:architecture Architecture and code structure changes area:inference Generation, sampling, decoding (incl. speculative, DRY) labels Aug 22, 2026
@inureyes

Copy link
Copy Markdown
Member Author

make verify-test-cuda on GB10 (CUDA sm_121, Linux aarch64) at 478bd6be, run with no other cargo process on the box: 8256 passed, 0 failed, 311 ignored, 101 suites, exit 0. The branch was current with main at gate time.

That is +3 against main, and the diff adds exactly 3 #[test] functions and removes none, so the totals reconcile. The two pre-existing eviction tests were rewritten to call the extracted function rather than reimplement it, which changes what they cover but not their count or their expected victim, since neither constructs a tie.

The log was scanned for process-level aborts as well as failing tests, because a per-suite tally cannot see a teardown crash: no SIGABRT, no terminate called, no error: test failed.

cargo fmt --all -- --check, cargo check --lib --tests --features cuda and cargo clippy --lib --tests --features cuda -- -D warnings are all clean on the same tree.

Bilingual report for the preemption victim total-order fix, recorded before the
merge because this repository carries the `TECHNICAL_REPORTS/.keep-reports`
marker. Force-added since `TECHNICAL_REPORTS/` is gitignored.

Refs #1293
@inureyes inureyes added the status:done Completed label Aug 22, 2026
@inureyes

Copy link
Copy Markdown
Member Author

Merging with the cargo-clippy check still queued, and recording why rather than leaving it implicit.

That job is queued behind the in-progress v0.6.0 release build, which holds the GB10 self-hosted runner. It is queued, not failing, and mergeStateStatus is UNSTABLE rather than BLOCKED, so it is not a required check. Every other check on this PR passed.

Rather than wait on an unrelated release, its exact command was run locally on this tree: cargo clippy -p mlxcel --lib --tests -- -D warnings, default features, exit 0 with no diagnostics. That is the same invocation the job runs, so the pending check has no information left to add. cargo clippy --lib --tests --features cuda -- -D warnings is also exit 0, as is make verify-test-cuda at 8256 passed and 0 failed.

Worth recording as a consequence of the job added in #1285: it runs on the GB10 self-hosted runner to get a warm cache, but GB10 also serves xla-compile and the release build, so a release in flight queues clippy behind it for every open PR. The lint this job exists to catch reproduces at default features and needs no accelerator, and pipeline-parallel-ci.yml already runs the same command on ubuntu-latest, so moving it there would trade a warm cache for independence from the release runner. Worth revisiting.

@inureyes
inureyes merged commit 88f2994 into main Aug 22, 2026
10 checks passed
@inureyes
inureyes deleted the fix/issue-1293-eviction-victim-total-order branch August 22, 2026 08:10
inureyes added a commit that referenced this pull request Aug 22, 2026
The job added in #1285 queued behind the in-progress v0.6.0 release on 2026-08-22
and delayed merging #1301, which raised the question of moving it to
`ubuntu-latest`. Measured instead of assumed, and the answer is to stay.

`src/lib/mlxcel-core/build.rs` builds MLX through cmake unconditionally: the
accelerator features select a backend, they do not decide whether the C++ builds
at all. A GitHub-hosted runner would therefore pay a cold MLX build on every run.
The only `ubuntu-latest` job this repository already has,
`pipeline-parallel-ci.yml`, caches the cargo registry and not `target/`, and it
takes about 28 minutes. On GB10 the same clippy command measured 2m44s cold and
27 to 40 seconds warm off its persistent target directory.

So the trade is a rare queue wait against roughly 20 minutes on every PR, and
the latter would re-create the cost objection that removed clippy from PR-time
CI in the first place (#21, #23). Also worth recording because it was initially
stated the other way round: a queued job consumes nothing, so this job did not
contend with the release. The contention that mattered that day was local
full-suite gates run by hand on the same machine.

Comment only. No workflow behavior changes.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:architecture Architecture and code structure changes area:inference Generation, sampling, decoding (incl. speculative, DRY) priority:medium Medium priority status:done Completed status:review Under review type:bug Bug fixes, error corrections, or issue resolutions

Projects

None yet

Development

Successfully merging this pull request may close these issues.

fix(server): preemption victim selection falls back to HashMap order when sequences tie

1 participant