Skip to content

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

Description

@inureyes

Problem / Background

BatchScheduler::select_eviction_victim (src/server/batch/scheduler.rs:6803-6826) decides which in-flight sequence gets preempted when the batch is full and a higher-priority request is waiting. Both policy arms tie-break on HashMap iteration order, and unlike the four sites in #1291, the tie here is the common case rather than an unreachable one.

  • PreemptionPolicy::LongestFirst (src/server/batch/scheduler.rs:6807-6811) is self.active_batch.iter_sequences().filter(|seq| seq.structured.is_none()).max_by_key(|seq| seq.generated_tokens.len()). Iterator::max_by_key returns the LAST maximum, so sequences tied on token count are separated by input order.
  • PreemptionPolicy::LowestPriority (src/server/batch/scheduler.rs:6815-6823) is min_by(|a, b| a.priority.cmp(&b.priority).then_with(|| b.generated_tokens.len().cmp(&a.generated_tokens.len()))). Two sequences equal on BOTH priority and generated length fall through to input order, and min_by returns the FIRST minimum.

The input order is HashMap order. ActiveBatch::iter_sequences is self.sequences.values() (src/server/batch/active.rs:120-122) over sequences: HashMap<SequenceId, SequenceInfo> (src/server/batch/active.rs:33). Rust's RandomState seeds each map instance, so that order is arbitrary and varies between processes, per the measurement recorded in docs/code-guidelines.md:128-130.

Why the tie is reachable here and was not in #1291. #1291's four LRU sites tie on std::time::Instant, stamped once per call, so at nanosecond resolution two entries effectively never collide and nothing misbehaves today. The key here is generated_tokens.len(), a small integer. Sequences admitted to the same batch decode in lockstep, one token per sequence per decode step, so a batch that admitted several requests together holds several sequences at an identical token count for as long as they stay together. PreemptionPolicy::LongestFirst is the default (src/server/config.rs:224-231), so the reachable-tie arm is the one that ships on by default. For the LowestPriority arm the tie needs equal RequestPriority as well, and Normal is the default priority (src/server/batch/sequence.rs:64-73), so a batch of unlabeled requests ties on that axis too.

Scope of the harm, stated honestly. Any tied-longest sequence does satisfy "evict the longest", so the policy's stated intent is not violated, no output is corrupted, and there is no crash. What is lost is reproducibility. An operator investigating why a particular request was preempted (and, per try_evict_for_preemption at src/server/batch/scheduler.rs:6658-6669, replayed with duplicate streamed tokens) cannot reproduce the choice from the same batch state, and across a fleet two identically configured workers facing identical load pick different victims. That is the same category of harm #1277 was filed for in node selection, and #1286 was filed as type:bug for the structurally identical tie-break-falls-to-hash-order shape in the distributed cache managers.

Current Behavior

Verified at a0e402ad:

// src/server/batch/scheduler.rs:6803-6826
fn select_eviction_victim(&self) -> Option<SequenceId> {
    match self.preemption_policy {
        PreemptionPolicy::LongestFirst => {
            self.active_batch
                .iter_sequences()
                .filter(|seq| seq.structured.is_none())
                .max_by_key(|seq| seq.generated_tokens.len())
                .map(|seq| seq.seq_id)
        }
        PreemptionPolicy::LowestPriority => {
            self.active_batch
                .iter_sequences()
                .filter(|seq| seq.structured.is_none())
                .min_by(|a, b| {
                    a.priority
                        .cmp(&b.priority)
                        .then_with(|| b.generated_tokens.len().cmp(&a.generated_tokens.len()))
                })
                .map(|seq| seq.seq_id)
        }
    }
}

Called from try_evict_for_preemption at src/server/batch/scheduler.rs:6659, whose returned victim is removed from the batch, has its prompt-cache context and KV caches released, is reset for re-prefill, and is reallocated under a fresh SequenceId (src/server/batch/scheduler.rs:6671-6760).

The existing tests do not cover this function and cannot catch a regression in it. eviction_selects_longest_first_by_default (src/server/batch/scheduler_tests.rs:693-719) and eviction_selects_lowest_priority_then_longest (src/server/batch/scheduler_tests.rs:722-757) each reimplement the selection expression inline against a locally built ActiveBatch instead of calling select_eviction_victim. They have already drifted: neither copy carries the filter(|seq| seq.structured.is_none()) guard that production has. This matters for the fix, because a regression test written in the same shape would pin the copy and pass regardless of what the production code does.

Proposed Solution

Give both arms a total order by appending the unique seq_id. docs/code-guidelines.md:109-199 ("HashMap Iteration Order") records the accepted forms and the reasoning; follow those rather than inventing a new shape. Two details that will bite otherwise:

  1. SequenceId derives only Debug, Clone, Copy, PartialEq, Eq, Hash (src/lib/mlxcel-core/src/cache.rs:6089-6090), so it is not Ord and cannot go into a comparison as-is. Use seq.seq_id.as_u64() (src/lib/mlxcel-core/src/cache.rs:6102-6104) rather than widening the type's public trait surface for two call sites.
  2. max_by_key returns the LAST maximum while min_by returns the FIRST minimum, so the two arms need opposite-facing tie components to break the same direction. Verified empirically with rustc -O on this box: over [(1,5),(2,5),(3,5)] keyed on the second field, max_by_key yields (3,5) and both min_by_key and min_by yield (1,5).

Concrete forms, with the direction argued below:

// LongestFirst. `max_by_key` returns the LAST maximum, so `Reverse` on the id
// makes the smallest seq_id win the tie. The id component is what makes the key
// a TOTAL order, so hash order cannot reach the result; do not simplify it away.
.max_by_key(|seq| (seq.generated_tokens.len(), std::cmp::Reverse(seq.seq_id.as_u64())))

// LowestPriority. `min_by` returns the FIRST minimum; the trailing seq_id
// comparison makes the order total so that "first" is never decided by hash order.
.min_by(|a, b| {
    a.priority
        .cmp(&b.priority)
        .then_with(|| b.generated_tokens.len().cmp(&a.generated_tokens.len()))
        .then_with(|| a.seq_id.as_u64().cmp(&b.seq_id.as_u64()))
})

Tie-break direction: smallest seq_id, and it is meaningful, not just total. SequenceId is a u64 newtype handed out by CachePool from next_id: AtomicU64 (src/lib/mlxcel-core/src/cache.rs:6315), initialized to 0 (:6349) and advanced by fetch_add(1, Ordering::Relaxed) on every allocation (:6393), and the type's own doc comment states the monotonic contract (:6084-6088). So the id already encodes admission order and "smallest id wins" is "evict the sequence admitted earliest", which is intelligible to an operator without any new field.

There is one nuance worth stating rather than discovering later. Preemption reallocates the victim under a fresh, therefore higher, id (src/server/batch/scheduler.rs:6734-6737), so an id records the most recent admission, not the original arrival. That makes smallest-id-wins the safe direction: a just-preempted sequence carries a young id and is the least likely to be picked again, whereas largest-id-wins would make the sequence that was just replayed the preferred next victim, which risks repeatedly starving the same request. If the maintainer wants true arrival order instead, SequenceInfo::created_at survives preemption (it is absent from the reset block at src/server/batch/scheduler.rs:6725-6732) and is the field to use, but it is an Instant and therefore not total on its own, so seq_id would still have to be appended behind it. Record whichever is chosen.

Make the function testable, because right now it is not. A regression test cannot call select_eviction_victim: it is a private method on BatchScheduler, which needs a real model to construct, and the two existing tests worked around that by copying the expression. Extract the policy into one pure function that both the scheduler and the tests call, for example a pub(crate) free function in src/server/batch/scheduler.rs:

pub(crate) fn select_eviction_victim_from<'a>(
    sequences: impl Iterator<Item = &'a SequenceInfo>,
    policy: PreemptionPolicy,
) -> Option<SequenceId>

with BatchScheduler::select_eviction_victim reduced to select_eviction_victim_from(self.active_batch.iter_sequences(), self.preemption_policy). A method on ActiveBatch taking the policy is the alternative; prefer the free function unless you would rather have active.rs depend on crate::server::config::PreemptionPolicy, since ActiveBatch currently knows nothing about preemption policy. Either way there must be exactly one copy of the selection logic, and the two existing tests at src/server/batch/scheduler_tests.rs:693 and :722 must be rewritten to call it, which also restores the structured.is_none() filter to their coverage.

The guideline correction this issue originally asked for is already done. docs/code-guidelines.md:137 used to read "min_by_key and max_by_key. Both return the FIRST extremum", which is wrong for max_by_key. It was corrected in PR #1294 (698f6237) and now states the two directions separately, so read the current line rather than the text quoted in earlier revisions of this issue. The enforcement block at the end of that section also already records this issue as the finding both static-check candidates missed, so there is nothing left to add there either. Nothing about the fix below changes; the distinction simply no longer has to be rediscovered.

Scope

In scope: src/server/batch/scheduler.rs:6803-6826 (both arms, plus the extraction described above and a comment at each site naming seq_id as the total-order component); src/server/batch/scheduler_tests.rs:693-757 (rewrite the two existing tests to call the extracted function) plus the new determinism test; docs/code-guidelines.md only if you judge it worth adding this instance to the list at :184-190, which currently reads "The seven instances" and would become eight. The max_by_key correction at :137 and the enforcement note at :192-201 are already in place from PR #1294, so do not redo them.

Out of scope:

Implementation Notes

Acceptance Criteria

  • Both arms of select_eviction_victim select through a total order, so the victim cannot depend on HashMap iteration order.
  • The selection logic exists in exactly one place, callable from tests, and the two existing tests at src/server/batch/scheduler_tests.rs:693 and :722 call it instead of reimplementing it (which also brings the structured.is_none() filter under test).
  • Each arm carries a comment naming seq_id as the component that makes the order total, per docs/code-guidelines.md:166, so a later simplification pass does not strip it.
  • A decision is recorded in the PR body on the tie-break direction (smallest seq_id, largest, or created_at with seq_id behind it) and why, including whether the anti-starvation argument above was accepted.
  • A regression test constructs an actual tie (equal generated_tokens.len() for the LongestFirst arm; equal priority AND equal length for the LowestPriority arm), rebuilds the ActiveBatch inside each iteration, loops at least 32 times, and asserts the same victim every time. Per docs/code-guidelines.md:168-178 the rebuild is mandatory because RandomState seeds per map instance.
  • That test is demonstrated failing against the unfixed code, with the observed failure count out of the iteration total quoted in the PR body (PR fix(rt-detr-v2): make needs_sanitize independent of HashMap order #1281 recorded 27 of 64, PR fix(lang-bias): preserve YAML bias: order so priority is deterministic #1269 recorded first failures at iterations 0, 0 and 2).
  • docs/code-guidelines.md:137 no longer claims max_by_key returns the first extremum, and this instance is added to the list at :184-190.
  • make verify-test-cuda passes.

Verification

# The new and rewritten tests, before and after the fix.
cargo test --profile test-fast --features cuda eviction -- --test-threads=1

# Confirm the test actually fails against the unfixed code: revert only the two
# arms (keep the extraction and the test), re-run, and record the failure count.
git stash push -- src/server/batch/scheduler.rs

make verify-fmt
make verify-clippy
make verify-test-cuda

A pass is: the determinism test reports the same victim on all 32-plus iterations after the fix, and reports a nonzero failure count before it. Note that the full workspace suite has a pre-existing abort on this CUDA box unrelated to this change; confirm any failure also reproduces on unmodified main before attributing it here.

Technical Considerations

This is the eighth instance of the class recorded in docs/code-guidelines.md:109-199 (filed as #1287, landed as PR #1290), after #1265 (PR #1266), #1267 (PR #1269), #1276 (PR #1281), #1277 (PR #1284) and the three in #1286 (PR #1288), and separate from the four latent ones in #1291. It is the first in the lineage that is reachable through ordinary batch state rather than through a coarsening of the key.

It is also a live demonstration of the blind spot that guideline documents. The static-check candidates evaluated in #1287 (docs/code-guidelines.md:194-199) reported four min_by_key/max_by_key flags on the current tree and nothing else, and missed this site because the receiver is a cross-module accessor method (iter_sequences()) rather than a literal .values() on a map declared in the same file. That is exactly the limitation the guideline describes for #1277, and it is the stated reason the checks were declined rather than adopted. The enforcement section should note that the blind spot has now cost a real finding, so the recorded decision stays honest about what it gave up.

Metadata

Metadata

Assignees

No one assigned

    Labels

    area:architectureArchitecture and code structure changesarea:inferenceGeneration, sampling, decoding (incl. speculative, DRY)priority:mediumMedium prioritystatus:doneCompletedtype:bugBug fixes, error corrections, or issue resolutions

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions