Skip to content

fix(distributed): eviction candidate sorts are stable, so HashMap order still picks which tied sequence is evicted #1286

Description

@inureyes

Problem / Background

Three eviction paths under src/distributed/ build a candidate list from a HashMap, sort it, and then consume a prefix. Sorting looks like it makes the result deterministic and does not, because slice::sort_by_key is a stable sort: entries that compare equal keep their input order, and the input order here is HashMap iteration order, which RandomState randomizes per map instance. So among candidates that tie on the sort key, which ones land in the consumed prefix varies run to run.

This is the subtle member of the class fixed by #1265, #1267, #1276 and #1277. The other four had no sort at all, so the defect was visible on inspection. These three DO sort, which is exactly why they read as safe.

Current Behavior

Instance 1, the most consequential: src/distributed/tensor_parallel/cache_manager.rs. allocations is a HashMap<SequenceId, ShardedCacheAllocation> (src/distributed/tensor_parallel/cache_manager.rs:430). select_eviction_candidates (src/distributed/tensor_parallel/cache_manager.rs:710-721) does let mut entries: Vec<_> = self.allocations.values().collect(); at line 711, then entries.sort_by_key(|a| a.last_accessed) at line 715 for EvictionPolicy::LRU, or entries.sort_by_key(|a| a.current_offset) at line 718 for EvictionPolicy::LeastTokens. Its doc comment (lines 707-709) says it returns "ALL sequence IDs in eviction priority order (first = evict first). Callers are responsible for selecting a subset sufficient to relieve pressure."

The caller, check_pressure (src/distributed/tensor_parallel/cache_manager.rs:639-673), does exactly that: it iterates the candidates and breaks out as soon as projected_used <= target_bytes (lines 656-657). That break makes it a prefix consumer, so ties in the sort key decide which sequences actually lose their cache. The resulting EvictionSignal.sequence_ids is then applied verbatim by apply_eviction (src/distributed/tensor_parallel/cache_manager.rs:679-688), which removes each listed allocation. Ties are not exotic under LeastTokens, where current_offset is a token count (usize, src/distributed/tensor_parallel/cache_manager.rs:167) that many concurrent sequences share.

Instance 2: src/distributed/pipeline/cache_manager.rs. The same shape at select_eviction_candidates (src/distributed/pipeline/cache_manager.rs:696-709): self.allocations.values().collect() at line 697, then sort_by_key on last_accessed (line 701, PreemptionPolicy::LRU), current_offset (line 704, Shortest) or std::cmp::Reverse(current_offset) (line 707, Longest). allocations is a HashMap<SequenceId, StageCacheAllocation> at src/distributed/pipeline/cache_manager.rs:459. The result is published as PreemptionSignal.sequence_ids (assigned at src/distributed/pipeline/cache_manager.rs:686, field declared at line 355), documented as eviction priority order.

Inside this repository the only consumer reads .len() (src/distributed/pipeline/cache_manager.rs:381, the Display impl), so there is no prefix consumer yet and the live impact today is that a documented priority ordering is not reproducible. It is worth fixing with the other two rather than waiting for a consumer to make it matter.

Instance 3: src/distributed/request_tracker.rs. evict_if_needed (src/distributed/request_tracker.rs:328-347) collects terminal requests with inner.requests.iter().filter(..).map(..).collect() (lines 334-339) from a HashMap<String, RequestLifecycle> (src/distributed/request_tracker.rs:233), sorts with completed.sort_by_key(|(_, t)| *t) on created_at (line 341), then consumes a prefix explicitly: for (key, _) in completed.into_iter().take(to_remove) (line 345). Two requests that finish within the same Instant tick tie, and hash order picks which is dropped.

Impact

No corruption and no crash: the eviction still relieves the pressure it set out to relieve, and the counts are right. What is lost is reproducibility of WHICH sequence was sacrificed, which is exactly the thing an operator needs when a request is preempted and they are trying to work out why. Instance 1 is the one with live behavioral consequences today.

Proposed Solution

Maintainer's choice, but the shape is the same in all three: give the sort a total order so ties cannot fall back to hash order. The tie-break key is the unique id already carried by each entry, sequence_id (SequenceId = u64, src/distributed/tensor_parallel/cache_manager.rs:48 and src/distributed/pipeline/cache_manager.rs:44) for instances 1 and 2, and the String request key for instance 3.

Either of these forms is acceptable, since what matters is that the sort key is a total order, not which sort function is called:

// Explicit comparator form.
entries.sort_by(|a, b| a.last_accessed.cmp(&b.last_accessed).then_with(|| a.sequence_id.cmp(&b.sequence_id)));

// Tuple-key form (still sort_by_key, but the key is now total).
entries.sort_by_key(|a| (a.last_accessed, a.sequence_id));
entries.sort_by_key(|a| (std::cmp::Reverse(a.current_offset), a.sequence_id));

Apply this to every policy arm: both arms in instance 1 (LRU, LeastTokens), all three arms in instance 2 (LRU, Shortest, Longest), and the created_at sort in instance 3, where the tuple becomes (created_at, key).

Note that sort_unstable_by would NOT be a fix on its own; it would merely replace one arbitrary tie order with another.

Scope

In scope:

  • src/distributed/tensor_parallel/cache_manager.rs: total-order sort in select_eviction_candidates (lines 713-720), both policy arms.
  • src/distributed/pipeline/cache_manager.rs: total-order sort in select_eviction_candidates (lines 699-708), all three policy arms.
  • src/distributed/request_tracker.rs: total-order sort in evict_if_needed (line 341).
  • Regression tests in src/distributed/tensor_parallel/cache_manager_tests.rs, src/distributed/pipeline/cache_manager_tests.rs, and src/distributed/request_tracker_tests.rs.

Out of scope:

  • check_pipeline_pressure (src/distributed/pipeline/cache_manager.rs:870-879) selects across a &[&PipelineCacheManager] slice with max_by, so its input order is caller-controlled and not hash-derived. No change needed.
  • Registry::nodes_at_stage / nodes_at_rank (src/distributed/registry.rs:193-227) also stable-sort a HashMap-derived list on a possibly-tied key (rank, stage), but have no non-test consumer in this repository today, and a tie there means two nodes share a (stage, rank) slot, which is a configuration error rather than normal operation. Leave them alone unless a consumer appears.
  • Changing the eviction policies themselves. The order among tied entries becomes defined, not "better".

Implementation Notes

  • Reuse: the tie-break value is already a field on each entry (ShardedCacheAllocation::sequence_id at src/distributed/tensor_parallel/cache_manager.rs:159, StageCacheAllocation::sequence_id at src/distributed/pipeline/cache_manager.rs:119). No new field, accessor, or id source is needed. SequenceId is u64 and String is Ord, so both tie-breaks compile as-is.
  • Test access to private state: each test file starts with use super::*; and is attached with #[path = "..._tests.rs"] mod tests; inside the module it tests (src/distributed/tensor_parallel/cache_manager.rs:791, src/distributed/pipeline/cache_manager.rs:1071, src/distributed/request_tracker.rs:352). A child module can reach the parent's private items, so tests can call the private select_eviction_candidates and mutate self.allocations / self.inner directly. Do NOT add a public setter or a #[cfg(test)] accessor just to make ties constructible.
  • Forcing ties: last_accessed and created_at are Instant set from Instant::now(), which on Linux has nanosecond resolution and will effectively never tie naturally. Tests must write equal values into the map entries directly rather than hoping for a collision. current_offset ties need no such help; allocate several sequences with the same token count.
  • Constraints: no public API change, no change to EvictionSignal / PreemptionSignal shape, no behavior change for candidate lists that have no ties. Existing tests that assert on eviction ordering must keep passing unchanged.
  • Edge cases: empty allocations (candidate list empty, check_pressure returns None at src/distributed/tensor_parallel/cache_manager.rs:643-645, unchanged); a single candidate (no tie, unchanged); all candidates tied (the full list must now come out in ascending sequence_id order); to_remove larger than the number of terminal requests in instance 3 (take already saturates, unchanged).
  • Doc discrepancy worth fixing while in there: the PreemptionSignal.sequence_ids doc at src/distributed/pipeline/cache_manager.rs:354 reads "in priority order (evict first last)", while select_eviction_candidates at line 695 says "first = evict first". The parenthetical is garbled and contradicts the producer. Correct it to match the producer.

Acceptance Criteria

Verification

cargo fmt --all -- --check
cargo clippy --workspace --all-targets -- -D warnings
cargo test --profile test-fast distributed::tensor_parallel::cache_manager
cargo test --profile test-fast distributed::pipeline::cache_manager
cargo test --profile test-fast distributed::request_tracker
make verify-test-cuda

A pass is: every new determinism test green over its full 32-iteration loop, every pre-existing test in those three modules still green, and clippy silent. Before trusting the tests, revert the three sorts to their current sort_by_key form and confirm each new test fails; a test that still passes against unfixed code is not testing the tie order.

Technical Considerations

Found by a sweep of the remaining .values() and .iter() collection sites under src/distributed/ while implementing #1277 (PR #1284). Deliberately not folded into that PR, because these change eviction behavior rather than routing and deserve their own tests.

The sweep at 4d4ac957 covered every sort_by, sort_by_key and sort() call under src/distributed/. Everything else is already safe: src/distributed/registry.rs sorts on the unique config.id after #1284; src/distributed/pipeline/metrics.rs:374,482,535 sorts on keys that are themselves the unique HashMap key ((src_stage, dst_stage), (stage_index, reason), stage_index); src/distributed/cluster_init.rs:552 sorts distinct SocketAddr strings and dedups; src/distributed/disaggregated/request_router.rs:668 was already fixed and carries an explanatory comment; the rest sort Vec-derived or slice-derived input. src/distributed/routing.rs:250 is a genuine prefix consumer (online[0] after a stable sort that ties on an idle cluster) but its input now arrives already ordered by node id through Registry::all_nodes (src/distributed/registry.rs:143-148) via SchedulerCore::build_candidates (src/distributed/scheduler.rs:413-438), so #1284 already cured it.

Related: #1265 and PR #1266, #1267 and PR #1269, #1276 and PR #1281, #1277 and PR #1284, all instances of the same root-cause class.

Metadata

Metadata

Assignees

No one assigned

    Labels

    area:architectureArchitecture and code structure changesarea:coremlxcel-core: MLX FFI, primitives, KV cache, layerspriority: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