From ab668d80bf7af40330e536da978ea21efbdb2297 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Fri, 4 Sep 2026 03:59:39 +0000 Subject: [PATCH 001/102] feat(kv-router): host-supplied inputs for SelectionService parity with the frontend scheduler First half of bringing the standalone selection core to input parity with the frontend KvRouter so one engine can serve both. - SelectionHostHooks on SelectionServiceBuilder: prefill-load estimator, overloaded-worker exclusion set, hard-availability inclusion set, shared KV cache, and LoRA worker filter, applied to every partition scheduler the core creates. The standalone service leaves them unset, so its behavior is unchanged; an embedding host feeds the same state it passes to KvScheduler. - Shared-cache hits are queried concurrently with the indexer and reach the scorer and custom policies; select requests accept the frontend's structured session_context (parent_session_id, session_final, kv_hints, input_trigger), with the flat session_id kept as the legacy form. - LoRA-aware candidate narrowing as a runtime-free LoraWorkerFilter trait with the frontend's exact semantics: never widens the allow-set, keeps a pinned worker inside it, falls back to the caller's set when narrowing empties it. - selection_id -> partition index so prefill_complete, free_reservation, and add_output_block stop scanning every partition; a lazy sweep drops ids whose booking was released out of band. - Advisory select through select_without_admission (no queue, no booking) with the chosen worker's projected load and prefill_busy / decode_busy against the conditional-disagg thresholds; admitted responses gain the same two fields. Tests: providers steer selection through the core; shared-cache hits and every session field reach the picker; LoRA helper matrix and injected-filter core test; reservation index population, indexed-first ordering, removal, and sweep; advisory and admitted responses with and without thresholds. Signed-off-by: Ishan Dhanani --- .../ext-proc/src/selector.rs | 1 + .../router/standalone-selection.md | 65 +- lib/kv-router/src/lib.rs | 1 + lib/kv-router/src/scheduling/local.rs | 5 + lib/kv-router/src/scheduling/lora_filter.rs | 136 +++ lib/kv-router/src/scheduling/mod.rs | 2 + .../src/services/selection/catalog.rs | 28 + .../src/services/selection/core/mod.rs | 823 +++++++++++++++++- lib/kv-router/src/services/selection/mod.rs | 5 +- .../src/services/selection/service.rs | 13 +- lib/kv-router/src/services/selection/types.rs | 149 +++- 11 files changed, 1177 insertions(+), 51 deletions(-) create mode 100644 lib/kv-router/src/scheduling/lora_filter.rs diff --git a/deploy/inference-gateway/ext-proc/src/selector.rs b/deploy/inference-gateway/ext-proc/src/selector.rs index f4e0e56adacb..1927256f56eb 100644 --- a/deploy/inference-gateway/ext-proc/src/selector.rs +++ b/deploy/inference-gateway/ext-proc/src/selector.rs @@ -291,6 +291,7 @@ impl Selector { router_config_override: None, expected_output_tokens: req.expected_output_tokens, session_id: None, + session_context: None, priority_jump: req.priority_jump, strict_priority: req.strict_priority, affinity_target: None, diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-selection.md b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-selection.md index 592200d351d4..3ee257af4f6b 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-selection.md +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-selection.md @@ -181,12 +181,49 @@ globally unique `selection_id`, or allow the service to generate one: "cpu": 96, "disk": 128 }, - "effective_prefill_tokens": 384 + "effective_prefill_tokens": 384, + "potential_decode_blocks": 212, + "decode_busy": false } ``` `select` returns the same selection fields but omits `sequence_hashes`, `isl_tokens`, and -`track_prefill_tokens`. `selection_id` is omitted when absent. All `overlap` +`track_prefill_tokens`. `selection_id` is omitted when absent. + +`potential_decode_blocks` is the scheduler's projection of KV blocks on the +chosen worker once this request is decoding, including the request's own +blocks. `decode_busy` compares it against the worker's `total_kv_blocks` at +`conditional_disagg_decode_busy_threshold` and is omitted when either the +threshold or the capacity is unknown. + +### Advisory selection + +`POST /select` accepts `"advisory": true` to select from current scheduler +state without queue admission. The request never waits in the router queue, +and the response adds the chosen worker's projected load: + +```json +{ + "worker_id": 1, + "dp_rank": 0, + "potential_decode_blocks": 212, + "decode_busy": false, + "worker_load": { + "active_prefill_tokens": 2048, + "prefill_token_capacity": 8192, + "total_kv_blocks": 4096, + "prefill_busy": false + } +} +``` + +`prefill_busy` compares `active_prefill_tokens` against +`prefill_token_capacity` at `conditional_disagg_prefill_busy_threshold` and is +omitted when that threshold is unset. This is the probe a disaggregation +coordinator uses to decide whether to bypass remote prefill: one advisory +`select` against the prefill pool answers "would the prefill worker I'd get be +busy?" without a separate load read. `advisory` is not accepted on +`select_and_reserve`. All `overlap` values are matched token counts. `gpu`, `cpu`, and `disk` use the cumulative Mooncake tier semantics documented in the standalone indexer's [per-instance tier breakdown](standalone-indexer.md#per-instance-tier-breakdown). @@ -232,6 +269,30 @@ The selection service does not persist, replicate, or expire `session_id` bindings. It is not part of the selection response and is not retained by the pending-selection cache, so a `POST /reservations` replay does not carry it. +### `session_context` + +Both endpoints also accept an optional `session_context` object that carries +the full session metadata the frontend hands to worker selection. When it is +present, the flat `session_id` field is ignored. + +```json +{ + "token_ids": [1, 2, 3, 4], + "session_context": { + "session_id": "child-session", + "parent_session_id": "root-session", + "session_final": false, + "kv_hints": { "evict_session": false }, + "input_trigger": "tool_result" + } +} +``` + +Only `session_id` is required inside the object. `input_trigger` is one of +`user_message`, `tool_result`, or `other`. A custom policy reads the values +through `WorkerSelectionContext::session_context()`; the built-in selector +ignores them, exactly as it ignores `session_id`. + ## Ray Select-Then-Reserve Flow Ray can keep model invocation separate from selector admission: diff --git a/lib/kv-router/src/lib.rs b/lib/kv-router/src/lib.rs index 362a2f621f81..9c69a7a4d522 100644 --- a/lib/kv-router/src/lib.rs +++ b/lib/kv-router/src/lib.rs @@ -70,6 +70,7 @@ pub use queue::SchedulerQueue; pub use radix_tree::RadixTree; pub use scheduling::LocalScheduler; pub use scheduling::PrefillLoadEstimator; +pub use scheduling::lora_filter::LoraWorkerFilter; pub use scheduling::policy::{FcfsPolicy, RouterSchedulingPolicy, SchedulingPolicy, WsptPolicy}; pub use scheduling::{ KvSchedulerError, PotentialLoad, SchedulingRequest, SchedulingResponse, SessionContext, diff --git a/lib/kv-router/src/scheduling/local.rs b/lib/kv-router/src/scheduling/local.rs index f480ffa9b8a7..f8fd08151528 100644 --- a/lib/kv-router/src/scheduling/local.rs +++ b/lib/kv-router/src/scheduling/local.rs @@ -566,6 +566,11 @@ where Ok(()) } + /// Whether `request_id` currently holds a booking on any worker. + pub fn has_request(&self, request_id: &str) -> bool { + self.slots.request_worker(&request_id.to_string()).is_some() + } + #[doc(hidden)] pub fn booking_cleanup(&self) -> SchedulerBookingCleanup { self.queue.booking_cleanup() diff --git a/lib/kv-router/src/scheduling/lora_filter.rs b/lib/kv-router/src/scheduling/lora_filter.rs new file mode 100644 index 000000000000..7f550442a403 --- /dev/null +++ b/lib/kv-router/src/scheduling/lora_filter.rs @@ -0,0 +1,136 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! LoRA-aware candidate narrowing shared by every selection host. + +use std::collections::HashSet; + +use crate::protocols::{WorkerId, WorkerWithDpRank}; + +/// Host-owned knowledge of which workers can serve a LoRA adapter. +/// +/// The frontend implements this over its LoRA routing table and state tracker; +/// a runtime-free host may back it with a static map. Implementations return a +/// subset of `available`, ordered arbitrarily. Returning `available` unchanged +/// means "no LoRA-specific preference". +pub trait LoraWorkerFilter: Send + Sync { + fn filter_worker_ids_for_lora(&self, lora_name: &str, available: &[WorkerId]) -> Vec; +} + +/// Narrow `allowed_worker_ids` to the workers `filter` reports for `lora_name`, +/// strictly within the existing candidate universe. +/// +/// - Without a LoRA name the caller's allow-set is returned untouched. +/// - The universe is the caller's allow-set when present, else `all_workers`. +/// - A pinned worker inside the universe is always retained, even if the filter +/// dropped it, so the pin still wins for cache correctness. A pin outside the +/// universe is not re-added: the caller's constraint stands. +/// - An empty narrowing falls back to the caller's allow-set rather than +/// producing an unroutable request. +pub fn narrow_allowed_worker_ids_by_lora( + filter: &dyn LoraWorkerFilter, + lora_name: Option<&str>, + allowed_worker_ids: Option>, + pinned_worker: Option<&WorkerWithDpRank>, + all_workers: impl FnOnce() -> Vec, +) -> Option> { + let Some(lora_name) = lora_name else { + return allowed_worker_ids; + }; + let base: Vec = match &allowed_worker_ids { + Some(allowed) => allowed.iter().copied().collect(), + None => all_workers(), + }; + if base.is_empty() { + return allowed_worker_ids; + } + let mut narrowed: HashSet = filter + .filter_worker_ids_for_lora(lora_name, &base) + .into_iter() + .collect(); + if let Some(pinned) = pinned_worker + && base.contains(&pinned.worker_id) + { + narrowed.insert(pinned.worker_id); + } + if narrowed.is_empty() { + return allowed_worker_ids; + } + Some(narrowed) +} + +#[cfg(test)] +mod tests { + use super::*; + + struct OnlyWorker(WorkerId); + + impl LoraWorkerFilter for OnlyWorker { + fn filter_worker_ids_for_lora( + &self, + _lora_name: &str, + available: &[WorkerId], + ) -> Vec { + available + .iter() + .copied() + .filter(|id| *id == self.0) + .collect() + } + } + + fn narrow( + filter: &dyn LoraWorkerFilter, + lora_name: Option<&str>, + allowed: Option>, + pinned: Option, + ) -> Option> { + narrow_allowed_worker_ids_by_lora(filter, lora_name, allowed, pinned.as_ref(), || { + vec![1, 2, 3] + }) + } + + #[test] + fn no_lora_leaves_allow_set_untouched() { + let allowed = Some(HashSet::from([1, 2])); + assert_eq!(narrow(&OnlyWorker(3), None, allowed.clone(), None), allowed); + assert_eq!(narrow(&OnlyWorker(3), None, None, None), None); + } + + #[test] + fn narrows_within_all_workers_or_allow_set() { + assert_eq!( + narrow(&OnlyWorker(2), Some("a"), None, None), + Some(HashSet::from([2])) + ); + // Filter target outside the caller's allow-set: nothing survives, so + // the caller's set is preserved rather than widened. + let allowed = Some(HashSet::from([1, 3])); + assert_eq!( + narrow(&OnlyWorker(2), Some("a"), allowed.clone(), None), + allowed + ); + } + + #[test] + fn pinned_worker_is_retained_only_inside_the_universe() { + let pinned = WorkerWithDpRank::new(1, 0); + assert_eq!( + narrow(&OnlyWorker(2), Some("a"), None, Some(pinned)), + Some(HashSet::from([1, 2])) + ); + let allowed = Some(HashSet::from([2, 3])); + assert_eq!( + narrow(&OnlyWorker(2), Some("a"), allowed, Some(pinned)), + Some(HashSet::from([2])) + ); + } + + #[test] + fn empty_universe_is_a_noop() { + assert_eq!( + narrow_allowed_worker_ids_by_lora(&OnlyWorker(1), Some("a"), None, None, Vec::new), + None + ); + } +} diff --git a/lib/kv-router/src/scheduling/mod.rs b/lib/kv-router/src/scheduling/mod.rs index 3543367a4112..da58b6653d21 100644 --- a/lib/kv-router/src/scheduling/mod.rs +++ b/lib/kv-router/src/scheduling/mod.rs @@ -4,6 +4,7 @@ pub mod config; mod filter; mod local; +pub mod lora_filter; pub mod overlap; pub mod overlap_refresh; pub mod policy; @@ -19,6 +20,7 @@ mod worker_selection_config; mod types; pub use filter::*; pub use local::LocalScheduler; +pub use lora_filter::{LoraWorkerFilter, narrow_allowed_worker_ids_by_lora}; pub use overlap::{ CacheHitEstimates, OverlapAnalysis, OverlapScoresResponse, OverlapSignals, SelectedWorkerTierSnapshot, SharedCacheOverlapScore, WorkerOverlapScore, diff --git a/lib/kv-router/src/services/selection/catalog.rs b/lib/kv-router/src/services/selection/catalog.rs index 11fc0010d72a..45a88de98099 100644 --- a/lib/kv-router/src/services/selection/catalog.rs +++ b/lib/kv-router/src/services/selection/catalog.rs @@ -99,6 +99,34 @@ impl WorkerCatalog { }) } + pub(super) fn schedulable_worker_ids_for_key(&self, key: &RoutingPartitionId) -> Vec { + self.workers + .read() + .values() + .filter(|record| { + record.lifecycle == WorkerLifecycle::Schedulable + && record.model_name == key.model_name + && record.routing_group == key.routing_group + }) + .map(|record| record.worker_id) + .collect() + } + + /// `total_kv_blocks` published by a schedulable worker in `key`'s partition. + pub(super) fn total_kv_blocks( + &self, + worker_id: WorkerId, + key: &RoutingPartitionId, + ) -> Option { + let workers = self.workers.read(); + let record = workers.get(&worker_id)?; + (record.lifecycle == WorkerLifecycle::Schedulable + && record.model_name == key.model_name + && record.routing_group == key.routing_group) + .then_some(record.total_kv_blocks) + .flatten() + } + pub(super) fn scheduler_configs_for_key( &self, key: &RoutingPartitionId, diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index 4e1eca2af1aa..e8059258e7c5 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -12,17 +12,18 @@ use tokio::sync::{mpsc, watch}; use tokio_util::sync::CancellationToken; use crate::identity::RoutingPartitionId; -use crate::indexer::TieredMatchDetails; +use crate::indexer::{SharedKvCache, TieredMatchDetails}; use crate::protocols::{ - ActiveSequenceEvent, LocalBlockHash, PrefillLoadHint, RoutingConstraints, WorkerAffinityTarget, - WorkerId, WorkerWithDpRank, + ActiveSequenceEvent, LocalBlockHash, PrefillLoadHint, RoutingConstraints, SharedCacheHits, + WorkerAffinityTarget, WorkerId, WorkerWithDpRank, }; use crate::scheduling::config::RouterConfigOverride; use crate::scheduling::selector::WorkerSelectionPolicy; use crate::scheduling::{ - KvSchedulerError, LocalScheduler, OverlapAnalysis, OverlapSignals, PotentialLoad, ScheduleMode, - ScheduleRequest, SessionContext, TieredOverlapRefresher, effective_prefill_tokens, - prefill_load_hint_from_effective_tokens, + KvSchedulerError, LocalScheduler, LoraWorkerFilter, OverlapAnalysis, OverlapSignals, + OverloadedWorkerProvider, PotentialLoad, PrefillLoadEstimator, ScheduleMode, ScheduleRequest, + SessionContext, TieredOverlapRefresher, WorkerAvailabilityProvider, effective_prefill_tokens, + narrow_allowed_worker_ids_by_lora, prefill_load_hint_from_effective_tokens, }; use crate::sequences::{ ActiveSequencesMultiWorker, ReplicaWorkerPolicy, SequenceError, SequenceRequest, @@ -43,8 +44,8 @@ use super::pending::{PendingSelection, SelectionCache, SelectionCacheConfig}; use super::types::{ ModelLoadResponse, OverlapScoresRequest, OverlapScoresResponse, PotentialLoadsRequest, ReadyResponse, ReservationRequest, ReservationResponse, SelectAndReserveRequest, SelectRequest, - SelectResponse, SelectionWorkerConfig, WorkerCatalogRecord, WorkerLifecycle, - WorkerPatchRequest, WorkerRequest, + SelectResponse, SelectionWorkerConfig, SelectionWorkerLoad, WorkerCatalogRecord, + WorkerLifecycle, WorkerPatchRequest, WorkerRequest, }; use crate::WorkerSelectionPolicyFactory; use crate::WorkerType; @@ -71,6 +72,7 @@ struct PreparedSelectionInputs { sequence_hashes: Vec, isl_tokens: usize, overlap: OverlapSignals, + shared_cache_hits: Option, } struct SelectionOperation { @@ -82,11 +84,13 @@ struct SelectionOperation { priority_jump: f64, strict_priority: u32, policy_class: Option, - session_id: Option, + session_context: Option, affinity_target: Option, pinned_worker: Option, allowed_worker_ids: Option>, routing_constraints: RoutingConstraints, + /// Skip queue admission and return the chosen worker's load snapshot. + advisory: bool, } /// Resolved inputs for booking a reservation, shared by the cached and explicit @@ -102,6 +106,48 @@ struct ReservationBooking { lora_name: Option, } +/// Host-supplied hooks applied to every partition scheduler the core creates. +/// +/// A host that owns request transport (for example the frontend `KvRouter`) +/// uses these to feed overload shedding and hard availability from its own +/// client state, to attach a prefill-duration estimator, and to consult a +/// shared (cross-worker) KV cache during selection. The standalone HTTP +/// service leaves them unset. +#[derive(Clone, Default)] +pub struct SelectionHostHooks { + pub prefill_load_estimator: Option>, + pub overloaded_worker_provider: Option, + pub available_worker_provider: Option, + /// Queried alongside the indexer for prompts that carry `token_ids`; a + /// failed lookup is logged and selection proceeds without shared hits. + pub shared_cache: Option>, + /// Narrows candidates to the workers that can serve the request's LoRA + /// adapter, strictly within the caller's `allowed_worker_ids`. + pub lora_worker_filter: Option>, +} + +impl std::fmt::Debug for SelectionHostHooks { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("SelectionHostHooks") + .field( + "prefill_load_estimator", + &self.prefill_load_estimator.is_some(), + ) + .field("shared_cache", &self.shared_cache.is_some()) + .field("lora_worker_filter", &self.lora_worker_filter.is_some()) + .field( + "overloaded_worker_provider", + &self.overloaded_worker_provider.is_some(), + ) + .field( + "available_worker_provider", + &self.available_worker_provider.is_some(), + ) + .finish() + } +} + #[derive(Debug, Clone)] pub struct SelectionServiceConfig { pub port: u16, @@ -113,12 +159,29 @@ pub struct SelectionServiceConfig { pub selection_cache: SelectionCacheConfig, } +type SelectionEntries = RwLock>>>>; + +/// `selection_id` -> partition that holds its booking. +/// +/// Lifecycle calls (`prefill_complete`, `free`, `add_output_block`) arrive with +/// only a selection id, so without this index every call scans every partition +/// scheduler. Selection ids are caller-controlled strings, so this stays on the +/// standard hasher. Entries are removed on `free`, on a `RequestNotFound` from +/// the indexed scheduler, and by a periodic sweep that drops ids whose booking +/// expired underneath them. +type ReservationIndex = RwLock>; + pub struct SelectionCore { catalog: WorkerCatalog, - entries: RwLock>>>>, + entries: Arc, + reservation_index: Arc, + /// Sweep task is started lazily from the first `ensure_entry`, which always + /// runs inside the host runtime; construction itself may not. + reservation_sweep_started: OnceCell<()>, indexer_registry: Arc, kv_router_config: crate::config::KvRouterConfig, worker_selection_policy_factory: Option, + host_hooks: SelectionHostHooks, worker_type: WorkerType, cancel_token: CancellationToken, replica_config: Option, @@ -161,6 +224,7 @@ impl SelectionCore { cancel_token, None, None, + SelectionHostHooks::default(), WorkerType::Aggregated, true, cache_config, @@ -175,6 +239,7 @@ impl SelectionCore { cancel_token: CancellationToken, replica_config: Option, worker_selection_policy_factory: Option, + host_hooks: SelectionHostHooks, worker_type: WorkerType, signal_indexer_ready: bool, cache_config: SelectionCacheConfig, @@ -190,10 +255,13 @@ impl SelectionCore { } Self { catalog: WorkerCatalog::default(), - entries: RwLock::new(HashMap::new()), + entries: Arc::new(RwLock::new(HashMap::new())), + reservation_index: Arc::new(RwLock::new(HashMap::new())), + reservation_sweep_started: OnceCell::new(), indexer_registry, kv_router_config, worker_selection_policy_factory, + host_hooks, worker_type, cancel_token, replica_config, @@ -419,6 +487,13 @@ impl SelectionCore { .ok_or_else(|| SelectionError::BadRequest("block_size is required".to_string()))?; let is_eagle = record.is_eagle.unwrap_or(false); let key = record.key(); + self.reservation_sweep_started.get_or_init(|| { + spawn_reservation_index_sweep( + Arc::clone(&self.entries), + Arc::clone(&self.reservation_index), + self.cancel_token.child_token(), + ); + }); let entry_cell = { self.entries.read().get(&key).cloned() }; let entry_cell = entry_cell.unwrap_or_else(|| { @@ -473,11 +548,12 @@ impl SelectionCore { profile, block_size, selector, - None, + self.host_hooks.prefill_load_estimator.clone(), Some(overlap_refresh), - None, - // Standalone selection has no router Client snapshot. - None, + // Standalone selection has no router Client snapshot, so + // these stay `None` unless an embedding host injects them. + self.host_hooks.overloaded_worker_provider.clone(), + self.host_hooks.available_worker_provider.clone(), self.kv_router_config.router_queue_recheck_interval(), self.kv_router_config.router_track_prefill_tokens, self.cancel_token.child_token(), @@ -614,9 +690,10 @@ impl SelectionCore { pub async fn select_with_policy_class( &self, - req: SelectRequest, + mut req: SelectRequest, policy_class: Option, ) -> Result { + let session_context = req.take_session_context(); self.schedule_selection( SelectionOperation { key: RoutingPartitionId::new(req.model_name, req.routing_group), @@ -627,11 +704,12 @@ impl SelectionCore { priority_jump: req.priority_jump.unwrap_or_default(), strict_priority: req.strict_priority.unwrap_or(0), policy_class, - session_id: req.session_id, + session_context, affinity_target: req.affinity_target, pinned_worker: req.pinned_worker, allowed_worker_ids: req.allowed_worker_ids, routing_constraints: req.routing_constraints, + advisory: req.advisory, }, false, ) @@ -647,9 +725,10 @@ impl SelectionCore { pub async fn select_and_reserve_with_policy_class( &self, - req: SelectAndReserveRequest, + mut req: SelectAndReserveRequest, policy_class: Option, ) -> Result { + let session_context = req.take_session_context(); let selection_id = req .selection_id .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); @@ -663,11 +742,12 @@ impl SelectionCore { priority_jump: req.priority_jump.unwrap_or_default(), strict_priority: req.strict_priority.unwrap_or(0), policy_class, - session_id: req.session_id, + session_context, affinity_target: req.affinity_target, pinned_worker: req.pinned_worker, allowed_worker_ids: req.allowed_worker_ids, routing_constraints: req.routing_constraints, + advisory: false, }, true, ) @@ -688,13 +768,19 @@ impl SelectionCore { priority_jump, strict_priority, policy_class, - session_id, + session_context, affinity_target, pinned_worker, allowed_worker_ids, routing_constraints, + advisory, } = operation; self.ensure_running()?; + if advisory && book { + return Err(SelectionError::BadRequest( + "advisory selection cannot book a reservation".to_string(), + )); + } let entry = self.ready_entry(&key)?; let PreparedSelectionInputs { @@ -702,12 +788,14 @@ impl SelectionCore { sequence_hashes, isl_tokens, overlap, + shared_cache_hits, } = self .prepare_selection_inputs( &entry, &prompt, self.kv_router_config .assume_kv_reuse(router_config_override.as_ref()), + true, ) .await?; let mode = if book { @@ -737,6 +825,16 @@ impl SelectionCore { track_prefill_tokens, ) }); + let allowed_worker_ids = match self.host_hooks.lora_worker_filter.as_deref() { + Some(filter) => narrow_allowed_worker_ids_by_lora( + filter, + prompt.lora_name.as_deref(), + allowed_worker_ids, + pinned_worker.as_ref(), + || self.catalog.schedulable_worker_ids_for_key(&key), + ), + None => allowed_worker_ids, + }; let response_sequence_hashes = book.then(|| sequence_hashes.iter().map(|hash| *hash as i64).collect()); let response_isl_tokens = book.then_some(isl_tokens); @@ -754,21 +852,34 @@ impl SelectionCore { priority_jump, strict_priority, policy_class, - session_context: session_id - .map(|session_id| SessionContext::new(session_id, None, None, None)), + session_context, expected_output_tokens, affinity_target, pinned_worker, allowed_worker_ids, routing_constraints, - shared_cache_hits: None, + shared_cache_hits, }; - let response = tokio::select! { + let (response, advisory_load) = tokio::select! { biased; _ = self.cancel_token.cancelled() => { return Err(SelectionError::Scheduler(KvSchedulerError::SubscriberShutdown)); } - result = entry.scheduler.schedule_request(schedule_request) => result?, + result = async { + if advisory { + entry + .scheduler + .select_without_admission(schedule_request) + .await + .map(|advisory| (advisory.response, Some(advisory.selected_worker_load))) + } else { + entry + .scheduler + .schedule_request(schedule_request) + .await + .map(|response| (response, None)) + } + } => result?, }; let endpoint = self .catalog @@ -785,6 +896,33 @@ impl SelectionCore { ); let effective_prefill = effective_prefill_tokens(isl_tokens, response.cached_tokens); + let potential_decode_blocks = response.potential_decode_blocks as u64; + let total_kv_blocks = advisory_load + .and_then(|load| load.total_kv_blocks.map(|blocks| blocks as u64)) + .or_else(|| { + self.catalog + .total_kv_blocks(response.best_worker.worker_id, &key) + }); + let decode_busy = self + .kv_router_config + .conditional_disagg_decode_busy_threshold + .zip(total_kv_blocks) + .map(|(threshold, total_kv_blocks)| { + potential_decode_blocks as f64 > threshold * total_kv_blocks as f64 + }); + let worker_load = advisory_load.map(|load| SelectionWorkerLoad { + active_prefill_tokens: load.active_prefill_tokens, + prefill_token_capacity: load.prefill_token_capacity, + total_kv_blocks, + prefill_busy: self + .kv_router_config + .conditional_disagg_prefill_busy_threshold + .map(|threshold| load.prefill_load_exceeds(threshold)), + }); + + if book && let Some(selection_id) = selection_id.as_deref() { + self.record_reservation(selection_id, &key); + } if let Some((cache_id, sequence_hashes, lora_name, track_prefill_tokens)) = cached_inputs { self.selection_cache.insert( @@ -816,6 +954,9 @@ impl SelectionCore { block_size: entry.block_size, overlap, effective_prefill_tokens: effective_prefill, + potential_decode_blocks, + decode_busy, + worker_load, }) } @@ -1007,6 +1148,7 @@ impl SelectionCore { lora_name, }) .await?; + self.record_reservation(&selection_id, &key); Ok(ReservationResponse { selection_id, @@ -1018,32 +1160,65 @@ impl SelectionCore { }) } + fn record_reservation(&self, selection_id: &str, key: &RoutingPartitionId) { + self.reservation_index + .write() + .insert(selection_id.to_string(), key.clone()); + } + + fn forget_reservation(&self, selection_id: &str) { + self.reservation_index.write().remove(selection_id); + } + + /// Entries to try for a lifecycle call on `selection_id`: the indexed + /// partition first, then every other initialized partition. The fallback + /// covers bookings mirrored from replica peers, which never pass through + /// this core's booking paths. + fn lifecycle_entries(&self, selection_id: &str) -> Vec> { + let indexed = self + .reservation_index + .read() + .get(selection_id) + .cloned() + .and_then(|key| self.entry(&key)); + let mut entries = self.initialized_entries(); + if let Some(indexed) = indexed { + entries.retain(|entry| !Arc::ptr_eq(entry, &indexed)); + entries.insert(0, indexed); + } + entries + } + pub async fn prefill_complete(&self, selection_id: &str) -> Result<(), SelectionError> { - let entries = self.initialized_entries(); - for entry in entries { + for entry in self.lifecycle_entries(selection_id) { match entry.scheduler.mark_prefill_completed(selection_id).await { Ok(()) => return Ok(()), Err(SequenceError::RequestNotFound { .. }) => continue, Err(error) => return Err(error.into()), } } + self.forget_reservation(selection_id); Err(SelectionError::NotFound(format!( "reservation {selection_id} not found" ))) } pub async fn free_reservation(&self, selection_id: &str) -> Result<(), SelectionError> { - let entries = self.initialized_entries(); - for entry in entries { - match entry.scheduler.free(selection_id).await { - Ok(()) => return Ok(()), - Err(SequenceError::RequestNotFound { .. }) => continue, - Err(error) => return Err(error.into()), + let result = async { + for entry in self.lifecycle_entries(selection_id) { + match entry.scheduler.free(selection_id).await { + Ok(()) => return Ok(()), + Err(SequenceError::RequestNotFound { .. }) => continue, + Err(error) => return Err(error.into()), + } } + Err(SelectionError::NotFound(format!( + "reservation {selection_id} not found" + ))) } - Err(SelectionError::NotFound(format!( - "reservation {selection_id} not found" - ))) + .await; + self.forget_reservation(selection_id); + result } pub fn add_output_block( @@ -1059,8 +1234,7 @@ impl SelectionCore { )); } - let entries = self.initialized_entries(); - for entry in entries { + for entry in self.lifecycle_entries(selection_id) { match entry .scheduler .add_output_block(selection_id, decay_fraction) @@ -1070,6 +1244,7 @@ impl SelectionCore { Err(error) => return Err(error.into()), } } + self.forget_reservation(selection_id); Err(SelectionError::NotFound(format!( "reservation {selection_id} not found" ))) @@ -1117,6 +1292,7 @@ impl SelectionCore { &req.prompt, self.kv_router_config .assume_kv_reuse(req.router_config_override.as_ref()), + false, ) .await?; let track_prefill_tokens = req @@ -1161,11 +1337,16 @@ impl SelectionCore { ) } + /// Normalize the prompt and gather cache signals. The indexer lookup and + /// the optional shared-cache lookup run concurrently; the shared cache is + /// consulted only when `query_shared_cache` is set, a shared cache is + /// attached, and the prompt carries raw `token_ids`. async fn prepare_selection_inputs( &self, entry: &SelectionEntry, prompt: &PromptRequest, assume_kv_reuse: bool, + query_shared_cache: bool, ) -> Result { let normalized = prompt.normalize_for_selection( entry.is_eagle, @@ -1175,15 +1356,36 @@ impl SelectionCore { assume_kv_reuse, }, )?; - let tiered = if normalized.block_hashes.is_empty() { - TieredMatchDetails::default() - } else { - entry - .indexer - .find_tiered_matches(normalized.block_hashes.clone()) + let indexer_lookup = async { + if normalized.block_hashes.is_empty() { + Ok(TieredMatchDetails::default()) + } else { + entry + .indexer + .find_tiered_matches(normalized.block_hashes.clone()) + .await + .map_err(|error| SelectionError::Internal(error.to_string())) + } + }; + let shared_cache = query_shared_cache + .then_some(self.host_hooks.shared_cache.as_deref()) + .flatten() + .zip(prompt.token_ids.as_deref()); + let shared_cache_lookup = async { + let (shared_cache, tokens) = shared_cache?; + match shared_cache + .check_blocks(tokens, entry.block_size, prompt.cache_namespace.as_deref()) .await - .map_err(|error| SelectionError::Internal(error.to_string()))? + { + Ok(hits) => Some(hits), + Err(error) => { + tracing::warn!(%error, "Shared cache query failed, ignoring"); + None + } + } }; + let (tiered, shared_cache_hits) = tokio::join!(indexer_lookup, shared_cache_lookup); + let tiered = tiered?; let overlap = OverlapAnalysis::new(&self.kv_router_config, entry.block_size, &tiered).signals(); drop(tiered); @@ -1192,6 +1394,7 @@ impl SelectionCore { sequence_hashes: normalized.sequence_hashes, isl_tokens: normalized.isl_tokens, overlap, + shared_cache_hits, }) } @@ -1209,6 +1412,65 @@ impl SelectionCore { } } +/// Drop index entries whose booking no longer exists in its partition +/// scheduler (expired by the periodic force-expiry, or freed through a path +/// that bypassed this core). Returns the number of entries removed. +fn sweep_reservation_index(entries: &SelectionEntries, index: &ReservationIndex) -> usize { + let snapshot: Vec<(String, RoutingPartitionId)> = index + .read() + .iter() + .map(|(id, key)| (id.clone(), key.clone())) + .collect(); + if snapshot.is_empty() { + return 0; + } + let stale: Vec = { + let entries = entries.read(); + snapshot + .into_iter() + .filter(|(id, key)| { + entries + .get(key) + .and_then(|cell| cell.get()) + .is_none_or(|entry| !entry.scheduler.has_request(id)) + }) + .map(|(id, _)| id) + .collect() + }; + if stale.is_empty() { + return 0; + } + let mut index = index.write(); + stale + .iter() + .filter(|id| index.remove(id.as_str()).is_some()) + .count() +} + +fn spawn_reservation_index_sweep( + entries: Arc, + index: Arc, + cancel_token: CancellationToken, +) { + let period = crate::sequences::active_request_expiry_duration(); + tokio::spawn(async move { + let mut interval = tokio::time::interval(period); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + interval.tick().await; + loop { + tokio::select! { + _ = cancel_token.cancelled() => break, + _ = interval.tick() => { + let removed = sweep_reservation_index(&entries, &index); + if removed > 0 { + tracing::debug!(removed, "Swept stale selection reservation index entries"); + } + } + } + } + }); +} + fn tracking_scope(entry: &SelectionEntry) -> TrackingHashScope<'_> { TrackingHashScope { partition: entry.key.as_ref(), @@ -1293,10 +1555,12 @@ mod tests { priority_jump: None, strict_priority: None, session_id: None, + session_context: None, affinity_target: None, pinned_worker: None, allowed_worker_ids: None, routing_constraints: RoutingConstraints::default(), + advisory: false, } } @@ -1311,6 +1575,7 @@ mod tests { priority_jump: None, strict_priority: None, session_id: None, + session_context: None, affinity_target: None, pinned_worker: None, allowed_worker_ids: None, @@ -1374,6 +1639,7 @@ mod tests { CancellationToken::new(), None, None, + SelectionHostHooks::default(), worker_type, true, SelectionCacheConfig::default(), @@ -1392,6 +1658,307 @@ mod tests { } } + fn core_with_hooks(hooks: SelectionHostHooks) -> SelectionCore { + core_with_hooks_and_policy(hooks, None) + } + + fn core_with_hooks_and_policy( + hooks: SelectionHostHooks, + policy_factory: Option, + ) -> SelectionCore { + let config = test_config(false); + let tracking_hash = Arc::new( + TrackingHashContext::from_config(&config).expect("valid tracking hash configuration"), + ); + SelectionCore::new_inner( + config, + 1, + CancellationToken::new(), + None, + policy_factory, + hooks, + WorkerType::Aggregated, + true, + SelectionCacheConfig::default(), + tracking_hash, + ) + } + + /// Picker that records what worker selection saw and always takes row 0. + struct CapturingPicker { + observed: Arc>>, + } + + #[derive(Debug, Clone)] + struct SelectionObservation { + session_context: Option, + shared_beyond_device_blocks: Vec, + } + + impl crate::scheduling::selector::WorkerPicker for CapturingPicker { + fn required_worker_inputs(&self) -> crate::scheduling::selector::WorkerInputs { + crate::scheduling::selector::WorkerInputs::CACHE + } + + fn pick( + &mut self, + context: &crate::scheduling::selector::WorkerSelectionContext<'_>, + input: crate::scheduling::selector::WorkerInputView<'_>, + ) -> Result { + self.observed.lock().push(SelectionObservation { + session_context: context.session_context().cloned(), + shared_beyond_device_blocks: input + .cache() + .expect("CACHE inputs requested") + .iter() + .map(|cache| cache.shared_beyond_device_blocks()) + .collect(), + }); + Ok(0) + } + } + + fn capturing_policy_factory() -> ( + WorkerSelectionPolicyFactory, + Arc>>, + ) { + let observed = Arc::new(parking_lot::Mutex::new(Vec::new())); + let factory_observed = Arc::clone(&observed); + let factory: WorkerSelectionPolicyFactory = + Arc::new(move |config, worker_type, _partition| { + WorkerSelectionPolicy::new( + config.clone(), + worker_type.as_str(), + Vec::new(), + Box::new(CapturingPicker { + observed: Arc::clone(&factory_observed), + }), + ) + }); + (factory, observed) + } + + type SharedCacheCalls = Arc, u32, Option)>>>; + + /// Shared cache that reports every block as a hit and records each query. + struct RecordingSharedCache { + calls: SharedCacheCalls, + } + + #[async_trait::async_trait] + impl SharedKvCache for RecordingSharedCache { + async fn check_blocks( + &self, + tokens: &[u32], + block_size: u32, + cache_namespace: Option<&str>, + ) -> Result { + self.calls.lock().push(( + tokens.to_vec(), + block_size, + cache_namespace.map(str::to_string), + )); + let blocks = (tokens.len() / block_size as usize) as u32; + Ok(SharedCacheHits::from_hits(&vec![true; blocks as usize])) + } + } + + struct OnlyWorkerForLora { + worker_id: WorkerId, + } + + impl LoraWorkerFilter for OnlyWorkerForLora { + fn filter_worker_ids_for_lora( + &self, + _lora_name: &str, + available: &[WorkerId], + ) -> Vec { + available + .iter() + .copied() + .filter(|id| *id == self.worker_id) + .collect() + } + } + + #[tokio::test] + async fn injected_lora_filter_narrows_candidates() { + let core = core_with_hooks(SelectionHostHooks { + lora_worker_filter: Some(Arc::new(OnlyWorkerForLora { worker_id: 2 })), + ..SelectionHostHooks::default() + }); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + core.upsert_worker(worker(2)).await.expect("worker upsert"); + + // LoRA request: only the filter's worker is eligible. + for _ in 0..4 { + let mut request = select_request(); + request.prompt.lora_name = Some("adapter-a".to_string()); + let response = core.select(request).await.expect("select"); + assert_eq!(response.worker_id, 2); + } + + // The filter never widens the caller's allow-set: an allow-set that + // excludes the filter's worker is preserved as-is. + let mut request = select_request(); + request.prompt.lora_name = Some("adapter-a".to_string()); + request.allowed_worker_ids = Some(HashSet::from([1])); + let response = core.select(request).await.expect("select"); + assert_eq!(response.worker_id, 1); + + // A pinned worker inside the universe survives the filter. + let mut request = select_request(); + request.prompt.lora_name = Some("adapter-a".to_string()); + request.pinned_worker = Some(WorkerWithDpRank::new(1, 0)); + let response = core.select(request).await.expect("select"); + assert_eq!(response.worker_id, 1); + } + + #[tokio::test] + async fn shared_cache_hits_reach_worker_selection() { + let calls: SharedCacheCalls = Arc::new(parking_lot::Mutex::new(Vec::new())); + let (factory, observed) = capturing_policy_factory(); + let core = core_with_hooks_and_policy( + SelectionHostHooks { + prefill_load_estimator: None, + overloaded_worker_provider: None, + available_worker_provider: None, + shared_cache: Some(Arc::new(RecordingSharedCache { + calls: Arc::clone(&calls), + })), + lora_worker_filter: None, + }, + Some(factory), + ); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + + let mut request = select_request(); + request.prompt.cache_namespace = Some("tenant-a".to_string()); + core.select(request).await.expect("select"); + + assert_eq!( + calls.lock().as_slice(), + &[(vec![1, 2, 3, 4], 4, Some("tenant-a".to_string()))] + ); + let observations = observed.lock().clone(); + assert_eq!(observations.len(), 1); + // One block in the prompt, no device overlap, so the whole prompt is a + // shared-cache hit beyond the device prefix. + assert_eq!(observations[0].shared_beyond_device_blocks, vec![1]); + + // Load projection does not consult the shared cache. + core.potential_loads(PotentialLoadsRequest { + model_name: "model".to_string(), + routing_group: "default".to_string(), + prompt: prompt(), + router_config_override: None, + }) + .await + .expect("potential loads"); + assert_eq!(calls.lock().len(), 1); + + // Prompts without raw tokens cannot be checked against the shared cache. + let mut request = select_request(); + request.prompt = PromptRequest { + token_ids: None, + block_hashes: Some(vec![11]), + sequence_hashes: Some(vec![101]), + isl_tokens: Some(4), + ..PromptRequest::default() + }; + core.select(request).await.expect("select"); + assert_eq!(calls.lock().len(), 1); + assert_eq!(observed.lock()[1].shared_beyond_device_blocks, vec![0]); + } + + #[tokio::test] + async fn session_context_reaches_worker_selection() { + use super::super::types::{ + SelectionKvHints, SelectionSessionContext, + }; + + let (factory, observed) = capturing_policy_factory(); + let core = core_with_hooks_and_policy(SelectionHostHooks::default(), Some(factory)); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + + let mut request = select_request(); + request.session_id = Some("ignored-legacy".to_string()); + request.session_context = Some(SelectionSessionContext { + session_id: "child-session".to_string(), + parent_session_id: Some("root-session".to_string()), + session_final: Some(true), + kv_hints: Some(SelectionKvHints { + evict_session: true, + }), + }); + core.select(request).await.expect("select"); + + let mut request = reserve_request("legacy-session-reservation"); + request.session_id = Some("legacy-only".to_string()); + core.select_and_reserve(request) + .await + .expect("select and reserve"); + + let observations = observed.lock(); + let context = observations[0] + .session_context + .as_ref() + .expect("structured session context"); + assert_eq!(context.session_id(), "child-session"); + assert_eq!(context.parent_session_id(), Some("root-session")); + assert_eq!(context.session_final(), Some(true)); + assert!(context.kv_hints().expect("kv hints").evict_session()); + + let legacy = observations[1] + .session_context + .as_ref() + .expect("legacy session context"); + assert_eq!(legacy.session_id(), "legacy-only"); + assert_eq!(legacy.parent_session_id(), None); + } + + #[tokio::test] + async fn injected_availability_provider_restricts_selection() { + let available: Arc>>>> = + Arc::new(parking_lot::Mutex::new(None)); + let provider_state = Arc::clone(&available); + let core = core_with_hooks(SelectionHostHooks { + prefill_load_estimator: None, + overloaded_worker_provider: None, + available_worker_provider: Some(Arc::new(move || provider_state.lock().clone())), + shared_cache: None, + lora_worker_filter: None, + }); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + core.upsert_worker(worker(2)).await.expect("worker upsert"); + + for only in [1, 2] { + *available.lock() = Some(Arc::new(HashSet::from([only]))); + for _ in 0..4 { + let response = core.select(select_request()).await.expect("select"); + assert_eq!(response.worker_id, only); + } + } + } + + #[tokio::test] + async fn injected_overload_provider_excludes_worker() { + let core = core_with_hooks(SelectionHostHooks { + prefill_load_estimator: None, + overloaded_worker_provider: Some(Arc::new(|| Some(HashSet::from([1])))), + available_worker_provider: None, + shared_cache: None, + lora_worker_filter: None, + }); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + core.upsert_worker(worker(2)).await.expect("worker upsert"); + + for _ in 0..4 { + let response = core.select(select_request()).await.expect("select"); + assert_eq!(response.worker_id, 2); + } + } + #[test] fn shutdown_keeps_parent_alive() { let parent = CancellationToken::new(); @@ -1648,6 +2215,171 @@ mod tests { )); } + #[tokio::test] + async fn advisory_select_reports_worker_load_and_busy_evaluation() { + let mut config = test_config(false); + config.conditional_disagg_prefill_busy_threshold = Some(0.5); + config.conditional_disagg_decode_busy_threshold = Some(0.0); + let core = SelectionCore::try_new_local( + config, + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ).expect("valid test config"); + let mut request = worker(1); + request.total_kv_blocks = Some(1000); + core.upsert_worker(request).await.expect("worker upsert"); + + // Admitted (queued) select: decode evaluation comes from the catalog's + // total_kv_blocks; no load snapshot is taken. + let response = core.select(select_request()).await.expect("select"); + assert!(response.potential_decode_blocks > 0); + assert_eq!( + response.decode_busy, + Some(true), + "threshold 0.0 is always exceeded" + ); + assert!(response.worker_load.is_none()); + + // Advisory select: same decode evaluation plus the projected load. + let mut request = select_request(); + request.advisory = true; + let response = core.select(request).await.expect("advisory select"); + assert!(response.potential_decode_blocks > 0); + assert_eq!(response.decode_busy, Some(true)); + let load = response.worker_load.expect("advisory load"); + assert_eq!(load.total_kv_blocks, Some(1000)); + assert_eq!(load.prefill_token_capacity, 1024); + assert_eq!(load.active_prefill_tokens, 0); + assert_eq!(load.prefill_busy, Some(false)); + + // Advisory selection does not book. + assert!(core.reservation_index.read().is_empty()); + assert_eq!( + core.loads(Some("model"), Some("default"))[0].loads[0].potential_prefill_tokens, + 0 + ); + } + + #[tokio::test] + async fn busy_evaluation_is_absent_without_thresholds_or_capacity() { + let core = SelectionCore::try_new_local( + test_config(false), + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ).expect("valid test config"); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + let mut request = select_request(); + request.advisory = true; + let response = core.select(request).await.expect("advisory select"); + assert_eq!(response.decode_busy, None); + let load = response.worker_load.expect("advisory load"); + assert_eq!(load.total_kv_blocks, None); + assert_eq!(load.prefill_busy, None); + } + + #[tokio::test] + async fn reservation_index_tracks_bookings_until_freed() { + let core = SelectionCore::try_new_local( + test_config(false), + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ).expect("valid test config"); + for (worker_id, routing_group) in [(1, "group-a"), (2, "group-b")] { + let mut request = worker(worker_id); + request.routing_group = routing_group.to_string(); + core.upsert_worker(request).await.expect("worker upsert"); + } + let key_b = RoutingPartitionId::new("model", "group-b"); + + // select_and_reserve records the booking's partition. + let mut request = reserve_request("booked"); + request.routing_group = "group-b".to_string(); + core.select_and_reserve(request).await.expect("reserve"); + assert_eq!(core.reservation_index.read().get("booked"), Some(&key_b)); + assert_eq!( + core.lifecycle_entries("booked")[0].key, + key_b, + "indexed partition is tried first" + ); + + // The explicit reservation path records too. + let mut request = select_request(); + request.routing_group = "group-b".to_string(); + request.selection_id = Some("cached".to_string()); + core.select(request).await.expect("select"); + assert!(core.reservation_index.read().get("cached").is_none()); + core.create_reservation(ReservationRequest { + model_name: "model".to_string(), + routing_group: "group-b".to_string(), + selection_id: "cached".to_string(), + worker_id: None, + dp_rank: None, + prompt: PromptRequest::default(), + router_config_override: None, + expected_output_tokens: None, + effective_prefill_tokens: None, + track_prefill_tokens: None, + }) + .await + .expect("cached reservation"); + assert_eq!(core.reservation_index.read().get("cached"), Some(&key_b)); + + // Lifecycle calls still resolve, and free drops the index entry. + core.prefill_complete("booked") + .await + .expect("prefill complete"); + core.free_reservation("booked").await.expect("free"); + assert!(core.reservation_index.read().get("booked").is_none()); + core.free_reservation("cached").await.expect("free"); + assert!(core.reservation_index.read().is_empty()); + + // Unknown ids fall back to the full scan and stay unindexed. + assert!(matches!( + core.prefill_complete("never-booked").await, + Err(SelectionError::NotFound(_)) + )); + assert!(core.reservation_index.read().is_empty()); + } + + #[tokio::test] + async fn reservation_index_sweep_drops_bookings_released_out_of_band() { + let core = SelectionCore::try_new_local( + test_config(false), + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ).expect("valid test config"); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + core.select_and_reserve(reserve_request("live")) + .await + .expect("reserve live"); + core.select_and_reserve(reserve_request("stale")) + .await + .expect("reserve stale"); + assert_eq!(core.reservation_index.read().len(), 2); + assert_eq!( + sweep_reservation_index(&core.entries, &core.reservation_index), + 0 + ); + + // Release directly through the scheduler, as force-expiry would. + let entry = core + .entry(&RoutingPartitionId::new("model", "default")) + .expect("entry"); + entry.scheduler.free("stale").await.expect("scheduler free"); + + assert_eq!( + sweep_reservation_index(&core.entries, &core.reservation_index), + 1 + ); + let index = core.reservation_index.read(); + assert_eq!(index.len(), 1); + assert!(index.contains_key("live")); + } + #[tokio::test(flavor = "current_thread")] async fn queued_selection_returns_refreshed_overlap_snapshot() { let mut config = test_config(false); @@ -1726,6 +2458,7 @@ mod tests { priority_jump: None, strict_priority: None, session_id: None, + session_context: None, affinity_target: None, pinned_worker: None, allowed_worker_ids: None, diff --git a/lib/kv-router/src/services/selection/mod.rs b/lib/kv-router/src/services/selection/mod.rs index 2f24a0b77ad3..7000f65c2984 100644 --- a/lib/kv-router/src/services/selection/mod.rs +++ b/lib/kv-router/src/services/selection/mod.rs @@ -21,7 +21,7 @@ mod tests; pub use crate::WorkerSelectionPolicyFactory; pub use crate::services::common::replica_sync::ReplicaPeerError; -pub use core::{SelectionCore, SelectionServiceConfig}; +pub use core::{SelectionCore, SelectionHostHooks, SelectionServiceConfig}; pub use error::SelectionError; pub use input::PromptRequest; pub use pending::SelectionCacheConfig; @@ -38,7 +38,8 @@ pub use service::{ pub use types::{ ModelLoadResponse, OutputBlockRequest, OverlapScoresRequest, OverlapScoresResponse, PotentialLoadsRequest, ReadyResponse, ReservationRequest, ReservationResponse, - SelectAndReserveRequest, SelectRequest, SelectResponse, SelectionWorkerConfig, + SelectAndReserveRequest, SelectRequest, SelectResponse, SelectionInputTrigger, + SelectionKvHints, SelectionSessionContext, SelectionWorkerConfig, SelectionWorkerLoad, SharedCacheOverlapScore, WorkerCatalogRecord, WorkerLifecycle, WorkerOverlapScore, WorkerPatchRequest, WorkerRequest, }; diff --git a/lib/kv-router/src/services/selection/service.rs b/lib/kv-router/src/services/selection/service.rs index cfa8ac0d34b4..122790230fba 100644 --- a/lib/kv-router/src/services/selection/service.rs +++ b/lib/kv-router/src/services/selection/service.rs @@ -13,7 +13,7 @@ use crate::services::common::replica_sync::{ }; use crate::tracking_hash::TrackingHashContext; -use super::core::{SelectionCore, SelectionServiceConfig}; +use super::core::{SelectionCore, SelectionHostHooks, SelectionServiceConfig}; use super::error::SelectionError; use super::pending::SelectionCacheConfig; use super::policy_registry::WorkerSelectionPolicyRegistry; @@ -33,6 +33,7 @@ pub struct SelectionServiceBuilder { selection_cache: SelectionCacheConfig, worker_type: WorkerType, worker_selection_policy_registry: WorkerSelectionPolicyRegistry, + host_hooks: SelectionHostHooks, } /// Warn when a host does not construct workers for explicitly configured policy roles. @@ -70,9 +71,18 @@ impl SelectionServiceBuilder { selection_cache: SelectionCacheConfig::default(), worker_type, worker_selection_policy_registry, + host_hooks: SelectionHostHooks::default(), } } + /// Attach host-owned hooks (prefill-load estimator, overload exclusion + /// set, hard-availability set). Applied to every partition scheduler this + /// service creates. + pub fn host_hooks(mut self, hooks: SelectionHostHooks) -> Self { + self.host_hooks = hooks; + self + } + pub fn indexer_threads(mut self, indexer_threads: usize) -> Self { self.indexer_threads = indexer_threads; self @@ -123,6 +133,7 @@ impl SelectionServiceBuilder { cancel_token.clone(), replica_config, worker_selection_policy_factory, + self.host_hooks, self.worker_type, false, self.selection_cache, diff --git a/lib/kv-router/src/services/selection/types.rs b/lib/kv-router/src/services/selection/types.rs index cecd91588362..aceb62516b4e 100644 --- a/lib/kv-router/src/services/selection/types.rs +++ b/lib/kv-router/src/services/selection/types.rs @@ -10,9 +10,11 @@ use crate::protocols::{ DpRank, KvTransferEnforcement, RoutingConstraints, WorkerAffinityTarget, WorkerConfigLike, WorkerId, WorkerWithDpRank, }; -use crate::scheduling::PotentialLoad; use crate::scheduling::config::RouterConfigOverride; pub use crate::scheduling::{OverlapScoresResponse, SharedCacheOverlapScore, WorkerOverlapScore}; +use crate::scheduling::{ + PotentialLoad, SessionContext, WorkerSelectionKvHints, +}; use crate::services::overlap::MooncakeOverlapSummary; use super::input::PromptRequest; @@ -376,12 +378,28 @@ pub struct SelectRequest { pub expected_output_tokens: Option, pub priority_jump: Option, pub strict_priority: Option, + /// Legacy session identity. Ignored when `session_context` is present. pub session_id: Option, + pub session_context: Option, pub affinity_target: Option, pub pinned_worker: Option, pub allowed_worker_ids: Option>, #[serde(default)] pub routing_constraints: RoutingConstraints, + /// Select from current scheduler state without queue admission. + /// + /// The response then carries the chosen worker's `worker_load` snapshot + /// and `prefill_busy` evaluation. The request never waits in the router + /// queue; when the queue would have rejected it the call fails the same + /// way an admitted selection does. Rejected on `select_and_reserve`. + #[serde(default)] + pub advisory: bool, +} + +impl SelectRequest { + pub(super) fn take_session_context(&mut self) -> Option { + resolve_session_context(self.session_context.take(), self.session_id.take()) + } } #[derive(Debug, Deserialize)] @@ -397,7 +415,9 @@ pub struct SelectAndReserveRequest { pub expected_output_tokens: Option, pub priority_jump: Option, pub strict_priority: Option, + /// Legacy session identity. Ignored when `session_context` is present. pub session_id: Option, + pub session_context: Option, pub affinity_target: Option, pub pinned_worker: Option, pub allowed_worker_ids: Option>, @@ -405,6 +425,60 @@ pub struct SelectAndReserveRequest { pub routing_constraints: RoutingConstraints, } +impl SelectAndReserveRequest { + pub(super) fn take_session_context(&mut self) -> Option { + resolve_session_context(self.session_context.take(), self.session_id.take()) + } +} + +/// Session metadata handed to worker selection, mirroring the frontend's +/// agent-context extension field for field. +/// +/// `session_context` supersedes the flat `session_id`: when both are present +/// the structured form wins. +#[derive(Debug, Clone, PartialEq, Eq, Deserialize)] +pub struct SelectionSessionContext { + pub session_id: String, + pub parent_session_id: Option, + pub session_final: Option, + pub kv_hints: Option, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)] +pub struct SelectionKvHints { + #[serde(default)] + pub evict_session: bool, +} + +impl From for SessionContext { + fn from(context: SelectionSessionContext) -> Self { + // Exhaustive so a new wire field must be mapped here. + let SelectionSessionContext { + session_id, + parent_session_id, + session_final, + kv_hints, + } = context; + SessionContext::new( + session_id, + parent_session_id, + session_final, + kv_hints.map(|SelectionKvHints { evict_session }| { + WorkerSelectionKvHints::new(evict_session) + }), + ) + } +} + +fn resolve_session_context( + session_context: Option, + session_id: Option, +) -> Option { + session_context.map(SessionContext::from).or_else(|| { + session_id.map(|session_id| SessionContext::new(session_id, None, None, None)) + }) +} + /// Booking request: replay the selection cached under `selection_id`, or book /// self-contained with `worker_id`. The replay books exactly what `select` captured; /// request fields other than the ids and model/routing-group are ignored. @@ -474,6 +548,33 @@ pub struct SelectResponse { pub block_size: u32, pub overlap: MooncakeOverlapSummary, pub effective_prefill_tokens: usize, + /// Projected KV blocks on the chosen worker once this request decodes, + /// including its own blocks: the scheduler's `potential_decode_blocks`. + pub potential_decode_blocks: u64, + /// `potential_decode_blocks` against the chosen worker's `total_kv_blocks` + /// at `conditional_disagg_decode_busy_threshold`. Absent when either the + /// threshold or the worker's capacity is unknown. + #[serde(skip_serializing_if = "Option::is_none")] + pub decode_busy: Option, + /// Chosen worker's load at selection time. Present only for advisory + /// selections (`SelectRequest::advisory`). + #[serde(skip_serializing_if = "Option::is_none")] + pub worker_load: Option, +} + +/// Load snapshot of the chosen worker, as the scheduler projected it for this +/// request. Mirrors the frontend's advisory selection load. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +pub struct SelectionWorkerLoad { + pub active_prefill_tokens: usize, + pub prefill_token_capacity: usize, + #[serde(skip_serializing_if = "Option::is_none")] + pub total_kv_blocks: Option, + /// `active_prefill_tokens` against `prefill_token_capacity` at + /// `conditional_disagg_prefill_busy_threshold`. Absent when the threshold + /// is unset. + #[serde(skip_serializing_if = "Option::is_none")] + pub prefill_busy: Option, } #[derive(Debug, Serialize)] @@ -501,3 +602,49 @@ pub struct ModelLoadResponse { pub pending_count: usize, pub pending_isl_tokens: usize, } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn select_request_deserializes_structured_session_context() { + let mut request: SelectRequest = serde_json::from_value(serde_json::json!({ + "token_ids": [1, 2, 3, 4], + "session_id": "legacy", + "session_context": { + "session_id": "child", + "parent_session_id": "root", + "session_final": false, + "kv_hints": { "evict_session": true } + } + })) + .expect("valid select request"); + + let context = request + .take_session_context() + .expect("structured context wins over legacy session_id"); + assert_eq!(context.session_id(), "child"); + assert_eq!(context.parent_session_id(), Some("root")); + assert_eq!(context.session_final(), Some(false)); + assert!(context.kv_hints().expect("kv hints").evict_session()); + } + + #[test] + fn select_request_falls_back_to_legacy_session_id() { + let mut request: SelectRequest = serde_json::from_value(serde_json::json!({ + "token_ids": [1, 2, 3, 4], + "session_id": "legacy" + })) + .expect("valid select request"); + let context = request.take_session_context().expect("legacy context"); + assert_eq!(context.session_id(), "legacy"); + assert_eq!(context.parent_session_id(), None); + assert_eq!(context.kv_hints(), None); + + let mut request: SelectAndReserveRequest = + serde_json::from_value(serde_json::json!({ "token_ids": [1, 2, 3, 4] })) + .expect("valid reserve request"); + assert!(request.take_session_context().is_none()); + } +} From 888548e387e778baa5b3df123af9a1547899ea4e Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Wed, 9 Sep 2026 01:49:47 -0700 Subject: [PATCH 002/102] fix(router): preserve current session input context contract Signed-off-by: Ishan Dhanani --- .../src/services/selection/core/mod.rs | 25 ++++++------ lib/kv-router/src/services/selection/mod.rs | 5 +-- lib/kv-router/src/services/selection/types.rs | 39 +++++++++++-------- 3 files changed, 39 insertions(+), 30 deletions(-) diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index e8059258e7c5..e541d53749bd 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -1873,9 +1873,7 @@ mod tests { #[tokio::test] async fn session_context_reaches_worker_selection() { - use super::super::types::{ - SelectionKvHints, SelectionSessionContext, - }; + use super::super::types::SelectionSessionContext; let (factory, observed) = capturing_policy_factory(); let core = core_with_hooks_and_policy(SelectionHostHooks::default(), Some(factory)); @@ -1887,9 +1885,7 @@ mod tests { session_id: "child-session".to_string(), parent_session_id: Some("root-session".to_string()), session_final: Some(true), - kv_hints: Some(SelectionKvHints { - evict_session: true, - }), + input_trigger: Some(super::super::types::SelectionInputTrigger::ToolResult), }); core.select(request).await.expect("select"); @@ -1907,7 +1903,10 @@ mod tests { assert_eq!(context.session_id(), "child-session"); assert_eq!(context.parent_session_id(), Some("root-session")); assert_eq!(context.session_final(), Some(true)); - assert!(context.kv_hints().expect("kv hints").evict_session()); + assert_eq!( + context.input_trigger(), + Some(crate::scheduling::WorkerSelectionInputTrigger::ToolResult) + ); let legacy = observations[1] .session_context @@ -2225,7 +2224,8 @@ mod tests { 1, CancellationToken::new(), SelectionCacheConfig::default(), - ).expect("valid test config"); + ) + .expect("valid test config"); let mut request = worker(1); request.total_kv_blocks = Some(1000); core.upsert_worker(request).await.expect("worker upsert"); @@ -2268,7 +2268,8 @@ mod tests { 1, CancellationToken::new(), SelectionCacheConfig::default(), - ).expect("valid test config"); + ) + .expect("valid test config"); core.upsert_worker(worker(1)).await.expect("worker upsert"); let mut request = select_request(); request.advisory = true; @@ -2286,7 +2287,8 @@ mod tests { 1, CancellationToken::new(), SelectionCacheConfig::default(), - ).expect("valid test config"); + ) + .expect("valid test config"); for (worker_id, routing_group) in [(1, "group-a"), (2, "group-b")] { let mut request = worker(worker_id); request.routing_group = routing_group.to_string(); @@ -2351,7 +2353,8 @@ mod tests { 1, CancellationToken::new(), SelectionCacheConfig::default(), - ).expect("valid test config"); + ) + .expect("valid test config"); core.upsert_worker(worker(1)).await.expect("worker upsert"); core.select_and_reserve(reserve_request("live")) .await diff --git a/lib/kv-router/src/services/selection/mod.rs b/lib/kv-router/src/services/selection/mod.rs index 7000f65c2984..ae344ccf2e62 100644 --- a/lib/kv-router/src/services/selection/mod.rs +++ b/lib/kv-router/src/services/selection/mod.rs @@ -39,7 +39,6 @@ pub use types::{ ModelLoadResponse, OutputBlockRequest, OverlapScoresRequest, OverlapScoresResponse, PotentialLoadsRequest, ReadyResponse, ReservationRequest, ReservationResponse, SelectAndReserveRequest, SelectRequest, SelectResponse, SelectionInputTrigger, - SelectionKvHints, SelectionSessionContext, SelectionWorkerConfig, SelectionWorkerLoad, - SharedCacheOverlapScore, WorkerCatalogRecord, WorkerLifecycle, WorkerOverlapScore, - WorkerPatchRequest, WorkerRequest, + SelectionSessionContext, SelectionWorkerConfig, SelectionWorkerLoad, SharedCacheOverlapScore, + WorkerCatalogRecord, WorkerLifecycle, WorkerOverlapScore, WorkerPatchRequest, WorkerRequest, }; diff --git a/lib/kv-router/src/services/selection/types.rs b/lib/kv-router/src/services/selection/types.rs index aceb62516b4e..4baaad5244f8 100644 --- a/lib/kv-router/src/services/selection/types.rs +++ b/lib/kv-router/src/services/selection/types.rs @@ -12,9 +12,7 @@ use crate::protocols::{ }; use crate::scheduling::config::RouterConfigOverride; pub use crate::scheduling::{OverlapScoresResponse, SharedCacheOverlapScore, WorkerOverlapScore}; -use crate::scheduling::{ - PotentialLoad, SessionContext, WorkerSelectionKvHints, -}; +use crate::scheduling::{PotentialLoad, SessionContext, WorkerSelectionInputTrigger}; use crate::services::overlap::MooncakeOverlapSummary; use super::input::PromptRequest; @@ -439,15 +437,20 @@ impl SelectAndReserveRequest { #[derive(Debug, Clone, PartialEq, Eq, Deserialize)] pub struct SelectionSessionContext { pub session_id: String, + #[serde(default)] pub parent_session_id: Option, + #[serde(default)] pub session_final: Option, - pub kv_hints: Option, + #[serde(default)] + pub input_trigger: Option, } #[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)] -pub struct SelectionKvHints { - #[serde(default)] - pub evict_session: bool, +#[serde(rename_all = "snake_case")] +pub enum SelectionInputTrigger { + UserMessage, + ToolResult, + Other, } impl From for SessionContext { @@ -457,14 +460,16 @@ impl From for SessionContext { session_id, parent_session_id, session_final, - kv_hints, + input_trigger, } = context; SessionContext::new( session_id, parent_session_id, session_final, - kv_hints.map(|SelectionKvHints { evict_session }| { - WorkerSelectionKvHints::new(evict_session) + input_trigger.map(|trigger| match trigger { + SelectionInputTrigger::UserMessage => WorkerSelectionInputTrigger::UserMessage, + SelectionInputTrigger::ToolResult => WorkerSelectionInputTrigger::ToolResult, + SelectionInputTrigger::Other => WorkerSelectionInputTrigger::Other, }), ) } @@ -474,9 +479,9 @@ fn resolve_session_context( session_context: Option, session_id: Option, ) -> Option { - session_context.map(SessionContext::from).or_else(|| { - session_id.map(|session_id| SessionContext::new(session_id, None, None, None)) - }) + session_context + .map(SessionContext::from) + .or_else(|| session_id.map(|session_id| SessionContext::new(session_id, None, None, None))) } /// Booking request: replay the selection cached under `selection_id`, or book @@ -616,7 +621,7 @@ mod tests { "session_id": "child", "parent_session_id": "root", "session_final": false, - "kv_hints": { "evict_session": true } + "input_trigger": "user_message" } })) .expect("valid select request"); @@ -627,7 +632,10 @@ mod tests { assert_eq!(context.session_id(), "child"); assert_eq!(context.parent_session_id(), Some("root")); assert_eq!(context.session_final(), Some(false)); - assert!(context.kv_hints().expect("kv hints").evict_session()); + assert_eq!( + context.input_trigger(), + Some(WorkerSelectionInputTrigger::UserMessage) + ); } #[test] @@ -640,7 +648,6 @@ mod tests { let context = request.take_session_context().expect("legacy context"); assert_eq!(context.session_id(), "legacy"); assert_eq!(context.parent_session_id(), None); - assert_eq!(context.kv_hints(), None); let mut request: SelectAndReserveRequest = serde_json::from_value(serde_json::json!({ "token_ids": [1, 2, 3, 4] })) From 042dc2956507294d40c2d488cb8457d3fcc31b79 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Fri, 4 Sep 2026 03:59:39 +0000 Subject: [PATCH 003/102] feat(kv-router): approximate, remote, and hint-aware indexing in SelectionService Second half of input parity: the standalone core indexes the way the frontend KvRouter does. - IndexerPolicy::from_router_config resolves the frontend's shape rules: event-driven primary, event-driven plus a short-TTL predict-on-route side indexer when router_predicted_ttl_secs is set, approximate TTL primary when use_kv_events=false. Bookings (select_and_reserve, cached replay, explicit reservations with tokens or hashes) are recorded into the approximate indexer; query-only selects record nothing. Previously a service with use_kv_events=false had an empty index and no cache awareness. - Indexer::Remote serves the primary index from a standalone indexer over HTTP (POST /query_tiered_by_hash returns lossless TieredMatchDetails). With a remote primary the core requires no kv_events endpoints and registers no ZMQ listeners. Exposed as SelectionServiceBuilder::remote_indexer, --remote-indexer-url, and remote_indexer_url= on the Python constructor. - Router hints: catalog records carry router_hint_worker_type and per-rank source control endpoints; a booking retains the matched chain when the partition has a hint-capable worker and SelectResponse::router_hint names a same-role source that beats the target's own cached prefix. Tests: backend policy matrix; approximate primary and side-indexer merge for 1 and 2 threads; bookings credited through all three paths; /query_tiered_by_hash round trip; remote indexer against a live standalone indexer including offline mapping; hint produced for a cold-worker booking and absent for warm-worker and query-only selections. Signed-off-by: Ishan Dhanani --- .../router/standalone-indexer.md | 26 + .../router/standalone-selection.md | 57 +- lib/bindings/python/rust/llm/kv.rs | 12 +- lib/kv-router/src/services/indexer/backend.rs | 999 +++++++++++++++++- .../src/services/indexer/registry.rs | 23 +- lib/kv-router/src/services/indexer/server.rs | 179 +++- .../src/services/selection/catalog.rs | 14 + .../src/services/selection/core/mod.rs | 497 ++++++++- .../src/services/selection/pending.rs | 9 +- .../src/services/selection/service.rs | 29 +- lib/kv-router/src/services/selection/tests.rs | 6 +- lib/kv-router/src/services/selection/types.rs | 62 +- 12 files changed, 1853 insertions(+), 60 deletions(-) diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-indexer.md b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-indexer.md index 3483286a242c..90d9c51156f3 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-indexer.md +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-indexer.md @@ -399,6 +399,32 @@ after they have been computed. Callers must precompute these hashes with the int and omit `cache_salt` from `/query_by_hash`. Use `/query` when the indexer should compute salted hashes from tokens server-side. +### `POST /query_tiered_by_hash` — Lossless tiered matches for a remote-primary selector + +```bash +curl -X POST http://localhost:8090/query_tiered_by_hash \ + -H 'Content-Type: application/json' \ + -d '{"block_hashes": [123456, 789012], "model_name": "llama-3-8b", "routing_group": "default"}' +``` + +Returns the indexer's tiered match details in their wire form rather than the +Mooncake score summary, so a caller can run its own overlap analysis: + +```json +{ + "block_size": 16, + "tiered": { + "device": { "scores": [[{"worker_id": 7, "dp_rank": 0}, 2]], "frequencies": [] }, + "lower_tier": [["HostPinned", { "hits": [[{"worker_id": 7, "dp_rank": 0}, 1]] }]] + } +} +``` + +Scores and hits are in blocks. This is the endpoint a +[standalone selection service](standalone-selection.md) configured with +`--remote-indexer-url` uses for its primary lookups. Both `model_name` and +`routing_group` are required; an unknown partition returns `404`. + ### Per-instance tier breakdown Each entry in `instances` is keyed by `instance_id` (as a string) and reports prefix reach across the device, host-pinned, and disk storage tiers: diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-selection.md b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-selection.md index 3ee257af4f6b..8c1b5e8eb1a9 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-selection.md +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-selection.md @@ -65,7 +65,8 @@ APIs. Those bindings should wrap `SelectionService` rather than construct |------|---------|-------------| | `--port` | `8092` | HTTP server port. | | `--threads` | `4` | KV indexer worker threads. | -| `--indexer-peers` | none | Comma-separated HTTP URLs used for startup KV recovery through `/dump`. | +| `--indexer-peers` | none | Comma-separated HTTP URLs used for startup KV recovery through `/dump`. Ignored with `--remote-indexer-url`. | +| `--remote-indexer-url` | none | Base URL of a [standalone indexer](standalone-indexer.md) that serves the primary KV index. The selector then does not subscribe to worker KV events; workers publish to the indexer instead. Requires `use_kv_events=true`. | | `--replica-sync-port` | none | Local ZMQ PUB port for active-load lifecycle events. The selector binds `tcp://*:` internally. | | `--replica-sync-peers` | none | Comma-separated ZMQ PUB endpoints for selector peers. Requires `--replica-sync-port`. | | `--selection-cache-ttl-secs` | `120` | Seconds an unclaimed pending selection lives before eviction. | @@ -79,6 +80,29 @@ APIs. Those bindings should wrap `SelectionService` rather than construct Router scheduling behavior continues to use the standard Dynamo router environment configuration. +### KV Indexing Modes + +The service resolves its indexer shape from the same router configuration the +frontend uses, so the two index the same way: + +| Configuration | Primary indexer | Routing decisions | +|---|---|---| +| `use_kv_events=true` (default) | Event-driven from worker ZMQ events | Not recorded | +| `use_kv_events=true`, `DYN_ROUTER_PREDICTED_TTL_SECS` set | Event-driven | Recorded into a short-TTL side indexer merged into device scores by per-worker max | +| `use_kv_events=false` | Approximate: no events, entries expire after `router_ttl_secs` | Recorded into the primary | +| `--remote-indexer-url` (or `SelectionServiceBuilder::remote_indexer`) | Served by a standalone indexer, queried through `POST /query_tiered_by_hash` | Recorded into the side indexer when `DYN_ROUTER_PREDICTED_TTL_SECS` is set, else not recorded | + +With a remote primary, worker records do not need `kv_events_endpoint(s)` to +become schedulable, no ZMQ listener is started in the selector, and `/dump` +returns no primary events (peers recover from the indexer instead). A remote +lookup failure is reported as an offline indexer for that selection. + +A routing decision is recorded when a reservation is booked (`select_and_reserve`, +or `select` followed by `POST /reservations`), never for a query-only `select`. +Hash-only reservations that supply `sequence_hashes` without `block_hashes` or +`token_ids` are not recorded. `router_approximate_cache_policy=lru` is not yet +supported here and falls back to TTL retention with a warning. + The standalone expiry guard measures absolute age from admission; output progress does not refresh it. Periodic cleanup therefore reclaims stale state approximately five to six minutes after admission by default. The embedded `KvRouter` uses the same `300`-second value as its shared @@ -117,6 +141,19 @@ previous partition and moved to the new one, which can leave the previous partition not ready. Assign a unique ID to every live worker across the entire service. +A worker that can consume router hints registers `router_hint_worker_type` +(its backend role, used to match hint sources to targets) and, when it can also +serve as a source, `router_hint_source_control_endpoints` keyed by global DP +rank: + +```json +{ + "worker_id": 1, + "router_hint_worker_type": "decode", + "router_hint_source_control_endpoints": { "0": "tcp://worker:5600", "1": "tcp://worker:5601" } +} +``` + `PATCH /workers/{worker_id}` updates supplied fields, `DELETE /workers/{worker_id}` removes the worker, and `GET /workers` lists catalog state. `model_name` and `routing_group` scope selection, indexer, and load state; @@ -196,6 +233,24 @@ blocks. `decode_busy` compares it against the worker's `total_kv_blocks` at `conditional_disagg_decode_busy_threshold` and is omitted when either the threshold or the capacity is unknown. +A booking response may also carry a `router_hint` when another worker of the +same `router_hint_worker_type` holds a longer cached prefix than the chosen +worker and advertises a source control endpoint for the matching DP rank: + +```json +{ + "router_hint": { + "source_control_endpoint": "tcp://worker-1:5600", + "block_hashes": [8713492873, 1928374650] + } +} +``` + +`block_hashes` are root-aligned external sequence hashes; entry `i` is request +block `i`, and the target decides which suffix to fetch from the source. Hints +require a local event-driven primary indexer (no approximate or remote primary) +and are never attached to query-only `select` responses. + ### Advisory selection `POST /select` accepts `"advisory": true` to select from current scheduler diff --git a/lib/bindings/python/rust/llm/kv.rs b/lib/bindings/python/rust/llm/kv.rs index 91c86ce80878..9f48c153469e 100644 --- a/lib/bindings/python/rust/llm/kv.rs +++ b/lib/bindings/python/rust/llm/kv.rs @@ -364,6 +364,11 @@ struct SelectServiceCli { #[arg(long, value_delimiter = ',')] indexer_peers: Vec, + /// Base URL of a standalone indexer that serves the primary KV index; this + /// selector then does not subscribe to worker KV events itself + #[arg(long)] + remote_indexer_url: Option, + /// Local ZMQ PUB port for active-load replica events #[arg(long, value_parser = parse_nonzero_port)] replica_sync_port: Option, @@ -552,6 +557,7 @@ where port: cli.port, threads: cli.threads, indexer_peers: cli.indexer_peers, + remote_indexer_url: cli.remote_indexer_url, replica_sync_port: cli.replica_sync_port, replica_sync_peers: cli.replica_sync_peers, kv_router_config, @@ -642,7 +648,7 @@ pub(crate) struct SelectionService { impl SelectionService { /// Create a selection service. `indexer_threads` sizes the KV indexer pool. #[new] - #[pyo3(signature = (*, indexer_threads = 4, indexer_peers = None, replica_sync_port = None, replica_sync_peers = None, selection_cache = None))] + #[pyo3(signature = (*, indexer_threads = 4, indexer_peers = None, replica_sync_port = None, replica_sync_peers = None, selection_cache = None, remote_indexer_url = None))] fn new( py: Python<'_>, indexer_threads: usize, @@ -650,6 +656,7 @@ impl SelectionService { replica_sync_port: Option, replica_sync_peers: Option>, selection_cache: Option, + remote_indexer_url: Option, ) -> PyResult { let replica_sync_peers = replica_sync_peers.unwrap_or_default(); if replica_sync_port.is_none() && !replica_sync_peers.is_empty() { @@ -672,6 +679,9 @@ impl SelectionService { if let Some(port) = replica_sync_port { builder = builder.replica_sync(port, replica_sync_peers); } + if let Some(url) = remote_indexer_url { + builder = builder.remote_indexer(url); + } let inner = py .allow_threads(|| pyo3_async_runtimes::tokio::get_runtime().block_on(builder.build())) .map_err(to_pyerr)?; diff --git a/lib/kv-router/src/services/indexer/backend.rs b/lib/kv-router/src/services/indexer/backend.rs index 65cb770596ee..6918a23c5ad3 100644 --- a/lib/kv-router/src/services/indexer/backend.rs +++ b/lib/kv-router/src/services/indexer/backend.rs @@ -2,18 +2,363 @@ // SPDX-License-Identifier: Apache-2.0 use std::sync::Arc; +use std::time::Duration; -use anyhow::Result; +use anyhow::{Context, Result}; use async_trait::async_trait; +use serde::{Deserialize, Serialize}; use tokio_util::sync::CancellationToken; use crate::ConcurrentRadixTreeCompressed; use crate::ThreadPoolIndexer; +use crate::approx::PruneConfig; +use crate::config::{ApproximateCachePolicyKind, KvRouterConfig}; +use crate::identity::RoutingPartitionId; use crate::indexer::{ - KvIndexer, KvIndexerInterface, KvIndexerMetrics, KvRouterError, LowerTierIndexers, - MatchDetails, TieredMatchDetails, TieredMatchProvider, query_lower_tiers, + ApproximateRetentionConfig, KvIndexer, KvIndexerInterface, KvIndexerMetrics, KvRouterError, + LowerTierIndexers, LowerTierQueryOptions, MatchDetails, RoutingDecisionHashes, SyncIndexer, + TieredMatchDetails, TieredMatchProvider, WireTieredMatchDetails, + query_lower_tiers_with_options, }; -use crate::protocols::{KvCacheEventData, LocalBlockHash, OverlapScores, RouterEvent, WorkerId}; +use crate::protocols::{ + KvCacheEventData, LocalBlockHash, OverlapScores, RouterEvent, WorkerId, WorkerWithDpRank, +}; + +/// How the device-tier primary is populated. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum PrimaryRetention { + /// Engine KV events are the only writer. + EventDriven, + /// No engine events: routing decisions are written into the primary and + /// expire after the TTL. This is the frontend's `use_kv_events=false` + /// approximate mode. + ApproximateTtl(Duration), + /// The primary lives in a standalone indexer service reached over HTTP. + /// Workers publish their KV events to that service, not to this process. + Remote(Arc), +} + +/// Connection to a standalone indexer service, shared by every partition's +/// [`RemoteIndexerClient`]. +#[derive(Debug)] +pub struct RemoteIndexerTransport { + base_url: String, + client: reqwest::Client, +} + +impl PartialEq for RemoteIndexerTransport { + fn eq(&self, other: &Self) -> bool { + self.base_url == other.base_url + } +} + +impl Eq for RemoteIndexerTransport {} + +impl RemoteIndexerTransport { + pub const DEFAULT_TIMEOUT: Duration = Duration::from_secs(5); + + pub fn new(base_url: impl Into) -> Result { + Self::new_with_timeout(base_url, Self::DEFAULT_TIMEOUT) + } + + pub fn new_with_timeout(base_url: impl Into, timeout: Duration) -> Result { + let base_url = base_url.into().trim_end_matches('/').to_string(); + if !(base_url.starts_with("http://") || base_url.starts_with("https://")) { + anyhow::bail!("remote indexer URL must start with http:// or https://: {base_url}"); + } + let client = reqwest::Client::builder() + .timeout(timeout) + .build() + .context("failed to build remote indexer HTTP client")?; + Ok(Self { base_url, client }) + } + + pub fn base_url(&self) -> &str { + &self.base_url + } +} + +/// Request body of the standalone indexer's `POST /query_tiered_by_hash`. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct TieredQueryByHashRequest { + pub block_hashes: Vec, + pub model_name: String, + pub routing_group: String, +} + +/// Response body of the standalone indexer's `POST /query_tiered_by_hash`: +/// the lossless tiered match shape a remote primary needs, as opposed to the +/// Mooncake score summary `/query_by_hash` returns. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct TieredQueryResponse { + pub block_size: u32, + pub tiered: WireTieredMatchDetails, +} + +/// Per-partition view of a remote primary indexer. +#[derive(Debug)] +pub struct RemoteIndexerClient { + transport: Arc, + key: RoutingPartitionId, +} + +impl RemoteIndexerClient { + pub fn new(transport: Arc, key: RoutingPartitionId) -> Self { + Self { transport, key } + } + + pub async fn find_tiered_matches( + &self, + sequence: &[LocalBlockHash], + ) -> std::result::Result { + if sequence.is_empty() { + return Ok(TieredMatchDetails::default()); + } + let request = TieredQueryByHashRequest { + block_hashes: sequence.iter().map(|hash| hash.0 as i64).collect(), + model_name: self.key.model_name.clone(), + routing_group: self.key.routing_group.clone(), + }; + let url = format!("{}/query_tiered_by_hash", self.transport.base_url); + let response = self + .transport + .client + .post(&url) + .json(&request) + .send() + .await + .map_err(|error| { + tracing::warn!(%error, url, "Remote indexer query failed"); + KvRouterError::IndexerOffline + })?; + let status = response.status(); + if !status.is_success() { + tracing::warn!(%status, url, "Remote indexer query returned an error status"); + return Err(KvRouterError::IndexerOffline); + } + let body: TieredQueryResponse = response.json().await.map_err(|error| { + tracing::warn!(%error, url, "Remote indexer query returned an invalid body"); + KvRouterError::IndexerOffline + })?; + Ok(body.tiered.into()) + } +} + +/// Indexer shape derived from router configuration, mirroring the frontend +/// `KvRouter` indexer resolution so a selection service built from the same +/// `KvRouterConfig` indexes the same way. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct IndexerPolicy { + pub primary: PrimaryRetention, + /// TTL of the predict-on-route side indexer layered over an event-driven + /// primary (`router_predicted_ttl_secs`). `None` disables it. + pub side_ttl: Option, +} + +impl Default for IndexerPolicy { + fn default() -> Self { + Self::event_driven() + } +} + +impl IndexerPolicy { + pub fn event_driven() -> Self { + Self { + primary: PrimaryRetention::EventDriven, + side_ttl: None, + } + } + + /// Whether the primary is served by a remote indexer service. Such a + /// service does not listen for worker KV events itself. + pub fn is_remote(&self) -> bool { + matches!(self.primary, PrimaryRetention::Remote(_)) + } + + /// Replace an event-driven primary with a remote one at `base_url`. + /// + /// An approximate primary cannot be remote: the standalone indexer has no + /// routing-decision write path, so `use_kv_events=false` is rejected. + pub fn with_remote_indexer(mut self, base_url: impl Into) -> Result { + match self.primary { + PrimaryRetention::EventDriven | PrimaryRetention::Remote(_) => {} + PrimaryRetention::ApproximateTtl(_) => anyhow::bail!( + "a remote indexer requires use_kv_events=true; the standalone indexer does not record routing decisions" + ), + } + self.primary = PrimaryRetention::Remote(Arc::new(RemoteIndexerTransport::new(base_url)?)); + Ok(self) + } + + /// Resolve the indexer shape for `config`. + /// + /// Rejects the same combinations the frontend rejects. Capacity-bounded LRU + /// retention needs per-request acquire/release fencing the service does not + /// yet have, so `router_approximate_cache_policy=lru` falls back to TTL + /// with a warning instead of failing. + pub fn from_router_config(config: &KvRouterConfig) -> Result { + if config.use_kv_events + && config.router_approximate_cache_policy == ApproximateCachePolicyKind::Lru + { + anyhow::bail!( + "router_approximate_cache_policy=lru requires use_kv_events=false; the local side indexer is TTL-only" + ); + } + if config.router_predicted_ttl_secs.is_some() && !config.use_kv_events { + anyhow::bail!( + "router_predicted_ttl_secs requires use_kv_events=true; \ + do not combine a primary approximate indexer with a side approximate indexer" + ); + } + if config.use_kv_events { + return Ok(Self { + primary: PrimaryRetention::EventDriven, + side_ttl: config + .router_predicted_ttl_secs + .map(Duration::from_secs_f64), + }); + } + if config.router_approximate_cache_policy == ApproximateCachePolicyKind::Lru { + tracing::warn!( + "router_approximate_cache_policy=lru is not supported by the selection service yet; falling back to TTL retention" + ); + } + Ok(Self { + primary: PrimaryRetention::ApproximateTtl(Duration::from_secs_f64( + config.router_ttl_secs, + )), + side_ttl: None, + }) + } +} + +/// Predict-on-route side indexer: a short-TTL approximate index that routing +/// decisions populate so a worker gets cache credit for a prompt before the +/// engine's first KV event for it arrives. Always local, never dumped or +/// replayed, and never used to seed lower-tier lookups. +#[derive(Clone)] +pub enum SideIndexer { + Single(KvIndexer), + Concurrent(Arc>), +} + +impl SideIndexer { + pub fn new( + ttl: Duration, + block_size: u32, + num_threads: usize, + metrics: Arc, + ) -> Self { + let prune_config = Some(PruneConfig { ttl }); + if num_threads > 1 { + return Self::Concurrent(Arc::new(ThreadPoolIndexer::new_with_metrics_and_pruning( + ConcurrentRadixTreeCompressed::new(), + num_threads, + block_size, + Some(metrics), + prune_config, + ))); + } + Self::Single(KvIndexer::new_with_pruning( + CancellationToken::new(), + block_size, + metrics, + prune_config, + )) + } + + async fn find_matches( + &self, + sequence: &[LocalBlockHash], + ) -> std::result::Result { + match self { + Self::Single(indexer) => indexer.find_matches(sequence.to_vec()).await, + Self::Concurrent(indexer) => Ok(indexer.backend().find_matches(sequence, false)), + } + } + + async fn record_routing_decision( + &self, + worker: WorkerWithDpRank, + hashes: RoutingDecisionHashes, + ) -> std::result::Result<(), KvRouterError> { + match self { + Self::Single(indexer) => { + indexer + .process_routing_decision_with_hashes( + worker, + hashes.local_hashes, + hashes.sequence_hashes, + ) + .await + } + Self::Concurrent(indexer) => { + indexer + .process_routing_decision_hash_slices( + worker, + &hashes.local_hashes, + &hashes.sequence_hashes, + ) + .await + } + } + } + + async fn remove_worker(&self, worker_id: WorkerId) { + match self { + Self::Single(indexer) => indexer.remove_worker(worker_id).await, + Self::Concurrent(indexer) => indexer.remove_worker(worker_id).await, + } + } + + async fn remove_worker_dp_rank(&self, worker_id: WorkerId, dp_rank: u32) { + match self { + Self::Single(indexer) => indexer.remove_worker_dp_rank(worker_id, dp_rank).await, + Self::Concurrent(indexer) => indexer.remove_worker_dp_rank(worker_id, dp_rank).await, + } + } +} + +/// Merge side-indexer scores into the primary's device match details by +/// per-worker max. Side-only workers gain a score with no paired +/// `last_matched_hashes` entry, so the result must not seed lower-tier +/// continuations; callers run `query_lower_tiers` on the primary-only details +/// first. +fn merge_side_scores(mut primary: MatchDetails, side: OverlapScores) -> MatchDetails { + for (worker, side_score) in side.scores { + primary + .overlap_scores + .scores + .entry(worker) + .and_modify(|score| { + if side_score > *score { + *score = side_score; + } + }) + .or_insert(side_score); + } + primary +} + +async fn merge_side_or_warn( + side: Option<&SideIndexer>, + primary: MatchDetails, + sequence: &[LocalBlockHash], +) -> MatchDetails { + let Some(side) = side else { + return primary; + }; + match side.find_matches(sequence).await { + Ok(side_scores) => merge_side_scores(primary, side_scores), + Err(error) => { + tracing::warn!( + %error, + "predict-on-route side indexer query failed; using primary only" + ); + primary + } + } +} /// Block-content indexer wrapping a primary device-tier backend plus a /// per-tier registry of lower-tier indexers (host-pinned, disk, …). @@ -21,19 +366,111 @@ use crate::protocols::{KvCacheEventData, LocalBlockHash, OverlapScores, RouterEv /// The service owns a `lower_tier: LowerTierIndexers` so tier-tagged events /// get routed alongside device events and tier-aware queries can return /// per-tier hit counts. +/// +/// `approx` is the optional predict-on-route side indexer, queried alongside +/// the primary and merged by per-worker max. `primary_records_routing_decisions` +/// marks an approximate primary (no engine events) that routing decisions +/// populate directly. At most one of the two is active. #[derive(Clone)] pub enum Indexer { Single { primary: KvIndexer, lower_tier: LowerTierIndexers, + approx: Option, + primary_records_routing_decisions: bool, }, Concurrent { primary: Arc>, lower_tier: LowerTierIndexers, + approx: Option, + primary_records_routing_decisions: bool, + }, + /// Primary served by a standalone indexer over HTTP. Lower tiers are the + /// remote service's concern and arrive in its tiered response; only the + /// side indexer is local. + Remote { + primary: Arc, + approx: Option, }, } impl Indexer { + fn approx(&self) -> Option<&SideIndexer> { + match self { + Self::Single { approx, .. } + | Self::Concurrent { approx, .. } + | Self::Remote { approx, .. } => approx.as_ref(), + } + } + + pub fn is_remote(&self) -> bool { + matches!(self, Self::Remote { .. }) + } + + /// Whether booked routing decisions should be written into this indexer. + pub fn records_routing_decisions(&self) -> bool { + match self { + Self::Single { + approx, + primary_records_routing_decisions, + .. + } + | Self::Concurrent { + approx, + primary_records_routing_decisions, + .. + } => approx.is_some() || *primary_records_routing_decisions, + Self::Remote { approx, .. } => approx.is_some(), + } + } + + /// Router-hint chain retention needs a local primary whose hashes come + /// only from engine events. + pub fn supports_router_hint_chain_retention(&self) -> bool { + !self.is_remote() && !self.records_routing_decisions() + } + + /// Record a booked routing decision. Writes to the side indexer when one + /// is attached, else to an approximate primary, else is a no-op. + pub async fn record_routing_decision( + &self, + worker: WorkerWithDpRank, + hashes: RoutingDecisionHashes, + ) -> std::result::Result<(), KvRouterError> { + if let Some(side) = self.approx() { + return side.record_routing_decision(worker, hashes).await; + } + match self { + Self::Single { + primary, + primary_records_routing_decisions: true, + .. + } => { + primary + .process_routing_decision_with_hashes( + worker, + hashes.local_hashes, + hashes.sequence_hashes, + ) + .await + } + Self::Concurrent { + primary, + primary_records_routing_decisions: true, + .. + } => { + primary + .process_routing_decision_hash_slices( + worker, + &hashes.local_hashes, + &hashes.sequence_hashes, + ) + .await + } + Self::Single { .. } | Self::Concurrent { .. } | Self::Remote { .. } => Ok(()), + } + } + /// Apply an event without tier dispatch — kept for callers that have /// already determined this is a device-tier event. Most callers should /// use [`Self::apply_event_routed`]. @@ -41,6 +478,9 @@ impl Indexer { match self { Indexer::Single { primary, .. } => primary.apply_event(event).await, Indexer::Concurrent { primary, .. } => primary.apply_event(event).await, + Indexer::Remote { .. } => { + tracing::trace!("Dropping KV event: the primary indexer is remote"); + } } } @@ -56,6 +496,7 @@ impl Indexer { Self::Single { lower_tier, .. } | Self::Concurrent { lower_tier, .. } => { lower_tier.record_unsupported_residency_event(&event); } + Self::Remote { .. } => {} } return Ok(()); } @@ -65,6 +506,7 @@ impl Indexer { Indexer::Single { primary, lower_tier, + .. } => { if is_clear { let mut reset_error = None; @@ -95,6 +537,7 @@ impl Indexer { Indexer::Concurrent { primary, lower_tier, + .. } => { if is_clear { let mut reset_error = None; @@ -122,15 +565,22 @@ impl Indexer { .await; } } + Indexer::Remote { .. } => { + tracing::trace!("Dropping KV event: the primary indexer is remote"); + } } Ok(()) } pub async fn remove_worker(&self, worker_id: WorkerId) { + if let Some(side) = self.approx() { + side.remove_worker(worker_id).await; + } match self { Indexer::Single { primary, lower_tier, + .. } => { for indexer in lower_tier.all() { indexer.remove_worker(worker_id).await; @@ -140,20 +590,26 @@ impl Indexer { Indexer::Concurrent { primary, lower_tier, + .. } => { for indexer in lower_tier.all() { indexer.remove_worker(worker_id).await; } primary.remove_worker(worker_id).await; } + Indexer::Remote { .. } => {} } } pub async fn remove_worker_dp_rank(&self, worker_id: WorkerId, dp_rank: u32) { + if let Some(side) = self.approx() { + side.remove_worker_dp_rank(worker_id, dp_rank).await; + } match self { Indexer::Single { primary, lower_tier, + .. } => { for indexer in lower_tier.all() { indexer.remove_worker_dp_rank(worker_id, dp_rank).await; @@ -163,59 +619,102 @@ impl Indexer { Indexer::Concurrent { primary, lower_tier, + .. } => { for indexer in lower_tier.all() { indexer.remove_worker_dp_rank(worker_id, dp_rank).await; } primary.remove_worker_dp_rank(worker_id, dp_rank).await; } + Indexer::Remote { .. } => {} } } - /// Device-tier overlap scores. Existing flat-shape callers continue to use - /// this; tier-aware callers should use [`Self::find_tiered_matches`]. + /// Device-tier overlap scores, including side-indexer credit. Existing + /// flat-shape callers continue to use this; tier-aware callers should use + /// [`Self::find_tiered_matches`]. pub async fn find_matches(&self, hashes: Vec) -> Result { - match self { - Indexer::Single { primary, .. } => { - primary.find_matches(hashes).await.map_err(Into::into) + let primary = match self { + Indexer::Single { primary, .. } => primary.find_matches(hashes.clone()).await?, + Indexer::Concurrent { primary, .. } => primary.find_matches(hashes.clone()).await?, + Indexer::Remote { primary, .. } => { + primary + .find_tiered_matches(&hashes) + .await? + .device + .overlap_scores } - Indexer::Concurrent { primary, .. } => { - primary.find_matches(hashes).await.map_err(Into::into) - } - } + }; + let Some(side) = self.approx() else { + return Ok(primary); + }; + let mut merged = MatchDetails::new(); + merged.overlap_scores = primary; + Ok(merge_side_or_warn(Some(side), merged, &hashes) + .await + .overlap_scores) } /// Device match details + per-tier hits, suitable for building the /// Mooncake-RFC-shape per-instance breakdown. + /// + /// Lower-tier continuations are seeded from confirmed primary matches only; + /// side-indexer scores are merged into the device tier afterwards. pub async fn find_tiered_matches( &self, sequence: Vec, ) -> std::result::Result { - match self { + self.find_tiered_matches_with_options(sequence, LowerTierQueryOptions::default()) + .await + } + + /// [`Self::find_tiered_matches`] with lower-tier query options. Router-hint + /// chain retention is honored only by a local event-driven primary (see + /// [`Self::supports_router_hint_chain_retention`]); other shapes ignore it. + pub async fn find_tiered_matches_with_options( + &self, + sequence: Vec, + options: LowerTierQueryOptions, + ) -> std::result::Result { + let options = LowerTierQueryOptions { + retain_router_hint_chain: options.retain_router_hint_chain + && self.supports_router_hint_chain_retention(), + }; + let (device, lower_tier) = match self { Indexer::Single { primary, lower_tier, + .. } => { - let device = primary.find_match_details(sequence.clone()).await?; - let lt = query_lower_tiers(lower_tier, &sequence, &device); - Ok(TieredMatchDetails { - device, - lower_tier: lt, - }) + let device = primary + .find_match_details_with_options( + sequence.clone(), + options.retain_router_hint_chain, + ) + .await?; + let lt = query_lower_tiers_with_options(lower_tier, &sequence, &device, options); + (device, lt) } Indexer::Concurrent { primary, lower_tier, + .. } => { - let device: MatchDetails = - primary.backend().find_match_details_impl(&sequence, false); - let lt = query_lower_tiers(lower_tier, &sequence, &device); - Ok(TieredMatchDetails { - device, - lower_tier: lt, - }) + let device: MatchDetails = primary.backend().find_match_details_impl_with_options( + &sequence, + false, + options.retain_router_hint_chain, + ); + let lt = query_lower_tiers_with_options(lower_tier, &sequence, &device, options); + (device, lt) } - } + Indexer::Remote { primary, .. } => { + let tiered = primary.find_tiered_matches(&sequence).await?; + (tiered.device, tiered.lower_tier) + } + }; + let device = merge_side_or_warn(self.approx(), device, &sequence).await; + Ok(TieredMatchDetails { device, lower_tier }) } /// Dump every event tracked by this indexer in a form replayable by a @@ -231,6 +730,7 @@ impl Indexer { Indexer::Single { primary, lower_tier, + .. } => ( primary.dump_events().await.map_err(anyhow::Error::from)?, lower_tier.entries(), @@ -238,10 +738,14 @@ impl Indexer { Indexer::Concurrent { primary, lower_tier, + .. } => ( primary.dump_events().await.map_err(anyhow::Error::from)?, lower_tier.entries(), ), + // The remote service owns the events; a peer recovers from it, not + // from this process. + Indexer::Remote { .. } => return Ok(Vec::new()), }; let mut out = primary_events; @@ -448,6 +952,7 @@ mod tests { if let Indexer::Single { primary, lower_tier, + .. } = &source { let _ = primary.flush().await; @@ -482,6 +987,7 @@ mod tests { if let Indexer::Single { primary, lower_tier, + .. } = &replayed { let _ = primary.flush().await; @@ -530,7 +1036,9 @@ mod tests { lower_tier.get_or_create(StorageTier::HostPinned), lower_tier.get_or_create(StorageTier::Disk), ), - Indexer::Concurrent { .. } => unreachable!("test creates the single indexer"), + Indexer::Concurrent { .. } | Indexer::Remote { .. } => { + unreachable!("test creates the single indexer") + } }; indexer @@ -565,6 +1073,384 @@ mod tests { assert!(healthy_tier.dump_events().await.unwrap().is_empty()); } + fn policy_indexer(num_threads: usize, policy: IndexerPolicy) -> Indexer { + create_indexer_with_policy( + &RoutingPartitionId::new("model", "default"), + 4, + num_threads, + Arc::new(KvIndexerMetrics::new_unregistered()), + &policy, + ) + } + + async fn flush_side(approx: &Option) { + match approx { + Some(SideIndexer::Single(side)) => { + let _ = side.flush().await; + } + Some(SideIndexer::Concurrent(side)) => { + let _ = side.flush().await; + } + None => {} + } + } + + /// Wait until every locally enqueued event and routing decision is applied. + async fn flush(indexer: &Indexer) { + match indexer { + Indexer::Single { + primary, approx, .. + } => { + let _ = primary.flush().await; + flush_side(approx).await; + } + Indexer::Concurrent { + primary, approx, .. + } => { + let _ = primary.flush().await; + flush_side(approx).await; + } + Indexer::Remote { approx, .. } => flush_side(approx).await, + } + } + + #[tokio::test] + async fn event_driven_policy_ignores_routing_decisions() { + let indexer = policy_indexer(1, IndexerPolicy::event_driven()); + assert!(!indexer.records_routing_decisions()); + assert!(indexer.supports_router_hint_chain_retention()); + let worker = WorkerWithDpRank::new(7, 0); + indexer + .record_routing_decision( + worker, + RoutingDecisionHashes::from_local_hashes(vec![LocalBlockHash(11)]), + ) + .await + .unwrap(); + flush(&indexer).await; + let tiered = indexer + .find_tiered_matches(vec![LocalBlockHash(11)]) + .await + .unwrap(); + assert!(tiered.device.overlap_scores.scores.is_empty()); + } + + #[tokio::test] + async fn approximate_primary_records_routing_decisions() { + for num_threads in [1, 2] { + let indexer = policy_indexer( + num_threads, + IndexerPolicy { + primary: PrimaryRetention::ApproximateTtl(Duration::from_secs(60)), + side_ttl: None, + }, + ); + assert!(indexer.records_routing_decisions()); + assert!(!indexer.supports_router_hint_chain_retention()); + let worker = WorkerWithDpRank::new(7, 0); + indexer + .record_routing_decision( + worker, + RoutingDecisionHashes::from_local_hashes(vec![ + LocalBlockHash(11), + LocalBlockHash(12), + ]), + ) + .await + .unwrap(); + flush(&indexer).await; + + let sequence = vec![LocalBlockHash(11), LocalBlockHash(12), LocalBlockHash(13)]; + let tiered = indexer.find_tiered_matches(sequence.clone()).await.unwrap(); + assert_eq!( + tiered.device.overlap_scores.scores.get(&worker).copied(), + Some(2), + "{num_threads} thread(s): approximate primary should credit the routed prefix" + ); + let flat = indexer.find_matches(sequence).await.unwrap(); + assert_eq!(flat.scores.get(&worker).copied(), Some(2)); + + // The recorded prefix is dumped like any other primary state, so a + // recovering peer inherits it. + assert!(!indexer.dump_events().await.unwrap().is_empty()); + } + } + + #[tokio::test] + async fn side_indexer_merges_predicted_scores_by_worker_max() { + for num_threads in [1, 2] { + let indexer = policy_indexer( + num_threads, + IndexerPolicy { + primary: PrimaryRetention::EventDriven, + side_ttl: Some(Duration::from_secs(60)), + }, + ); + assert!(indexer.records_routing_decisions()); + let confirmed = WorkerWithDpRank::new(7, 0); + let predicted = WorkerWithDpRank::new(8, 0); + + // Engine event: worker 7 confirmed holds [11, 12]. + indexer + .apply_event_routed(store_event( + confirmed.worker_id, + confirmed.dp_rank, + 1, + &[], + &[11, 12], + StorageTier::Device, + )) + .await + .unwrap(); + // Routing decisions: worker 8 was just routed [11, 12, 13]; worker 7 + // was routed [11] (shorter than what the engine confirmed). + indexer + .record_routing_decision( + predicted, + RoutingDecisionHashes::from_local_hashes(vec![ + LocalBlockHash(11), + LocalBlockHash(12), + LocalBlockHash(13), + ]), + ) + .await + .unwrap(); + indexer + .record_routing_decision( + confirmed, + RoutingDecisionHashes::from_local_hashes(vec![LocalBlockHash(11)]), + ) + .await + .unwrap(); + flush(&indexer).await; + + let sequence = vec![LocalBlockHash(11), LocalBlockHash(12), LocalBlockHash(13)]; + let tiered = indexer.find_tiered_matches(sequence.clone()).await.unwrap(); + let scores = &tiered.device.overlap_scores.scores; + assert_eq!( + scores.get(&predicted).copied(), + Some(3), + "{num_threads} thread(s)" + ); + assert_eq!( + scores.get(&confirmed).copied(), + Some(2), + "{num_threads} thread(s): primary wins when it saw the longer prefix" + ); + // Side scores never seed lower tiers and are never dumped. + assert!(tiered.lower_tier.values().all(|tier| tier.hits.is_empty())); + let dumped = indexer.dump_events().await.unwrap(); + assert!( + dumped + .iter() + .all(|event| event.worker_id == confirmed.worker_id) + ); + + indexer.remove_worker(predicted.worker_id).await; + flush(&indexer).await; + let tiered = indexer.find_tiered_matches(sequence).await.unwrap(); + assert!(!tiered.device.overlap_scores.scores.contains_key(&predicted)); + } + } + + #[tokio::test] + async fn remote_indexer_queries_the_standalone_service_and_merges_side_scores() { + use crate::services::indexer::registry::WorkerRegistry; + use crate::services::indexer::server::spawn_test_indexer_server; + + let key = RoutingPartitionId::new("model", "default"); + let served = Arc::new(WorkerRegistry::new(1)); + let served_indexer = served.get_or_create_indexer(key.clone(), 4); + let confirmed = WorkerWithDpRank::new(7, 0); + served_indexer + .apply_event_routed(store_event(7, 0, 1, &[], &[11, 12], StorageTier::Device)) + .await + .unwrap(); + served_indexer + .apply_event_routed(store_event( + 7, + 0, + 2, + &[11, 12], + &[13], + StorageTier::HostPinned, + )) + .await + .unwrap(); + flush(&served_indexer).await; + if let Indexer::Single { lower_tier, .. } = &served_indexer { + for inner in lower_tier.all() { + let _ = inner.dump_events().await.unwrap(); + } + } + let (base_url, server) = spawn_test_indexer_server(served).await; + + let policy = IndexerPolicy { + primary: PrimaryRetention::EventDriven, + side_ttl: Some(Duration::from_secs(60)), + } + .with_remote_indexer(base_url.clone()) + .unwrap(); + assert!(policy.is_remote()); + let indexer = create_indexer_with_policy( + &key, + 4, + 1, + Arc::new(KvIndexerMetrics::new_unregistered()), + &policy, + ); + assert!(indexer.is_remote()); + assert!(indexer.records_routing_decisions()); + assert!(!indexer.supports_router_hint_chain_retention()); + + // Local events are dropped: the remote service owns the primary. + indexer + .apply_event_routed(store_event( + 9, + 0, + 1, + &[], + &[11, 12, 13], + StorageTier::Device, + )) + .await + .unwrap(); + assert!(indexer.dump_events().await.unwrap().is_empty()); + + // A routing decision lands in the local side indexer only. + let predicted = WorkerWithDpRank::new(8, 0); + indexer + .record_routing_decision( + predicted, + RoutingDecisionHashes::from_local_hashes(vec![LocalBlockHash(11)]), + ) + .await + .unwrap(); + flush(&indexer).await; + + let sequence = vec![LocalBlockHash(11), LocalBlockHash(12), LocalBlockHash(13)]; + let tiered = indexer.find_tiered_matches(sequence.clone()).await.unwrap(); + let scores = &tiered.device.overlap_scores.scores; + assert_eq!( + scores.get(&confirmed).copied(), + Some(2), + "remote device match" + ); + assert_eq!(scores.get(&predicted).copied(), Some(1), "local side match"); + assert!(!scores.contains_key(&WorkerWithDpRank::new(9, 0))); + assert_eq!( + tiered + .lower_tier + .get(&StorageTier::HostPinned) + .and_then(|d| d.hits.get(&confirmed).copied()), + Some(1), + "remote lower tiers are carried through" + ); + let flat = indexer.find_matches(sequence).await.unwrap(); + assert_eq!(flat.scores.get(&confirmed).copied(), Some(2)); + assert_eq!(flat.scores.get(&predicted).copied(), Some(1)); + + // Unknown partitions and an unreachable service degrade to offline. + let other = create_indexer_with_policy( + &RoutingPartitionId::new("other", "default"), + 4, + 1, + Arc::new(KvIndexerMetrics::new_unregistered()), + &policy, + ); + assert!(matches!( + other.find_tiered_matches(vec![LocalBlockHash(11)]).await, + Err(KvRouterError::IndexerOffline) + )); + server.abort(); + let _ = server.await; + assert!(matches!( + indexer.find_tiered_matches(vec![LocalBlockHash(11)]).await, + Err(KvRouterError::IndexerOffline) + )); + } + + #[test] + fn remote_indexer_rejects_approximate_primary_and_bad_urls() { + let approximate = IndexerPolicy { + primary: PrimaryRetention::ApproximateTtl(Duration::from_secs(1)), + side_ttl: None, + }; + assert!( + approximate + .with_remote_indexer("http://indexer:8091") + .is_err() + ); + assert!( + IndexerPolicy::event_driven() + .with_remote_indexer("indexer:8091") + .is_err() + ); + let policy = IndexerPolicy::event_driven() + .with_remote_indexer("http://indexer:8091/") + .unwrap(); + let PrimaryRetention::Remote(transport) = &policy.primary else { + panic!("expected a remote primary"); + }; + assert_eq!(transport.base_url(), "http://indexer:8091"); + } + + #[test] + fn indexer_policy_mirrors_frontend_resolution() { + let event_driven = KvRouterConfig::default(); + assert_eq!( + IndexerPolicy::from_router_config(&event_driven).unwrap(), + IndexerPolicy::event_driven() + ); + + let side = KvRouterConfig { + router_predicted_ttl_secs: Some(2.5), + ..Default::default() + }; + assert_eq!( + IndexerPolicy::from_router_config(&side).unwrap().side_ttl, + Some(Duration::from_secs_f64(2.5)) + ); + + let approximate = KvRouterConfig { + use_kv_events: false, + router_ttl_secs: 30.0, + ..Default::default() + }; + assert_eq!( + IndexerPolicy::from_router_config(&approximate).unwrap(), + IndexerPolicy { + primary: PrimaryRetention::ApproximateTtl(Duration::from_secs(30)), + side_ttl: None, + } + ); + + let lru_without_events = KvRouterConfig { + use_kv_events: false, + router_approximate_cache_policy: ApproximateCachePolicyKind::Lru, + ..Default::default() + }; + assert!(matches!( + IndexerPolicy::from_router_config(&lru_without_events) + .unwrap() + .primary, + PrimaryRetention::ApproximateTtl(_) + )); + + let lru_with_events = KvRouterConfig { + router_approximate_cache_policy: ApproximateCachePolicyKind::Lru, + ..Default::default() + }; + assert!(IndexerPolicy::from_router_config(&lru_with_events).is_err()); + + let side_without_events = KvRouterConfig { + use_kv_events: false, + router_predicted_ttl_secs: Some(1.0), + ..Default::default() + }; + assert!(IndexerPolicy::from_router_config(&side_without_events).is_err()); + } + #[cfg(feature = "metrics")] #[tokio::test] async fn acknowledged_clear_records_primary_event_for_both_indexers() { @@ -611,25 +1497,64 @@ pub fn create_indexer_with_metrics( num_threads: usize, metrics: Arc, ) -> Indexer { + create_indexer_with_policy( + &RoutingPartitionId::new("default", "default"), + block_size, + num_threads, + metrics, + &IndexerPolicy::event_driven(), + ) +} + +pub fn create_indexer_with_policy( + key: &RoutingPartitionId, + block_size: u32, + num_threads: usize, + metrics: Arc, + policy: &IndexerPolicy, +) -> Indexer { + let approx = policy + .side_ttl + .map(|ttl| SideIndexer::new(ttl, block_size, num_threads, Arc::clone(&metrics))); + let retention = match &policy.primary { + PrimaryRetention::EventDriven => None, + PrimaryRetention::ApproximateTtl(ttl) => { + Some(ApproximateRetentionConfig::Ttl(PruneConfig { ttl: *ttl })) + } + PrimaryRetention::Remote(transport) => { + return Indexer::Remote { + primary: Arc::new(RemoteIndexerClient::new(Arc::clone(transport), key.clone())), + approx, + }; + } + }; + let primary_records_routing_decisions = retention.is_some(); if num_threads > 1 { Indexer::Concurrent { - primary: Arc::new(ThreadPoolIndexer::new_with_metrics( - ConcurrentRadixTreeCompressed::new(), - num_threads, - block_size, - Some(metrics), - )), + primary: Arc::new( + ThreadPoolIndexer::new_with_metrics_and_approximate_retention( + ConcurrentRadixTreeCompressed::new(), + num_threads, + block_size, + Some(metrics), + retention, + ), + ), lower_tier: LowerTierIndexers::new(num_threads, block_size), + approx, + primary_records_routing_decisions, } } else { Indexer::Single { - primary: KvIndexer::new_with_pruning( + primary: KvIndexer::new_with_approximate_retention( CancellationToken::new(), block_size, metrics, - None, + retention, ), lower_tier: LowerTierIndexers::new(1, block_size), + approx, + primary_records_routing_decisions, } } } diff --git a/lib/kv-router/src/services/indexer/registry.rs b/lib/kv-router/src/services/indexer/registry.rs index ce96d84b5221..e8f7f51a1074 100644 --- a/lib/kv-router/src/services/indexer/registry.rs +++ b/lib/kv-router/src/services/indexer/registry.rs @@ -18,7 +18,7 @@ use crate::identity::RoutingPartitionId; use crate::indexer::KvIndexerMetrics; use crate::protocols::WorkerId; -use super::backend::{Indexer, create_indexer_with_metrics}; +use super::backend::{Indexer, IndexerPolicy, create_indexer_with_policy}; use super::listener::spawn_zmq_listener; pub struct IndexerEntry { @@ -315,6 +315,9 @@ pub struct WorkerRegistry { watermarks: DashMap<(WorkerId, u32), Arc>, num_threads: usize, indexer_metrics: Arc, + /// Shape of every indexer this registry creates. Set before the first + /// partition is created; later changes affect only new partitions. + indexer_policy: Mutex, ready_tx: watch::Sender, ready_rx: watch::Receiver, root_cancel_token: CancellationToken, @@ -361,12 +364,22 @@ impl WorkerRegistry { watermarks: DashMap::new(), num_threads, indexer_metrics, + indexer_policy: Mutex::new(IndexerPolicy::event_driven()), ready_tx, ready_rx, root_cancel_token, } } + /// Set the shape used for indexers created from now on. + pub fn set_indexer_policy(&self, policy: IndexerPolicy) { + *self.indexer_policy.lock() = policy; + } + + pub fn indexer_policy(&self) -> IndexerPolicy { + self.indexer_policy.lock().clone() + } + pub fn signal_ready(&self) { let _ = self.ready_tx.send(true); } @@ -437,10 +450,12 @@ impl WorkerRegistry { "Creating new indexer" ); IndexerEntry { - indexer: create_indexer_with_metrics( + indexer: create_indexer_with_policy( + &key, block_size, self.num_threads, self.indexer_metrics.clone(), + &self.indexer_policy.lock(), ), block_size, } @@ -743,10 +758,12 @@ impl WorkerRegistry { "Creating indexer from recovery dump" ); IndexerEntry { - indexer: create_indexer_with_metrics( + indexer: create_indexer_with_policy( + &key, block_size, self.num_threads, self.indexer_metrics.clone(), + &self.indexer_policy.lock(), ), block_size, } diff --git a/lib/kv-router/src/services/indexer/server.rs b/lib/kv-router/src/services/indexer/server.rs index 185efae4bbff..1f1483d39332 100644 --- a/lib/kv-router/src/services/indexer/server.rs +++ b/lib/kv-router/src/services/indexer/server.rs @@ -17,11 +17,11 @@ use tokio_util::sync::CancellationToken; use crate::identity::{RoutingPartitionId, default_routing_group}; #[cfg(feature = "metrics")] use crate::indexer::KvIndexerMetrics; -use crate::indexer::TieredMatchDetails; +use crate::indexer::{TieredMatchDetails, WireTieredMatchDetails}; use crate::protocols::{BlockHashOptions, LocalBlockHash, WorkerId, compute_block_hash_for_seq}; use crate::services::overlap::{MooncakeOverlapSummary, build_mooncake_overlap_summaries}; -use super::backend::Indexer; +use super::backend::{Indexer, TieredQueryByHashRequest, TieredQueryResponse}; use super::registry::{ListenerControlError, WorkerRegistry}; /// We need to fit one million tokens as JSON text, this should do it. @@ -387,6 +387,58 @@ async fn query_by_hash( resp } +/// Lossless tiered lookup for a remote-primary selection service. Unlike +/// `/query_by_hash`, which renders the Mooncake score summary, this returns +/// the wire form of `TieredMatchDetails` so the caller can run its own +/// overlap analysis and side-indexer merge. +async fn query_tiered_by_hash( + State(state): State>, + Json(req): Json, +) -> Response { + let model = req.model_name.clone(); + let key = RoutingPartitionId::new(req.model_name, req.routing_group); + let Some(ie) = state.registry.get_indexer(&key) else { + let mut resp = ( + StatusCode::NOT_FOUND, + Json(serde_json::json!({ + "error": format!( + "no indexer for model={} routing_group={}", + key.model_name, key.routing_group + ) + })), + ) + .into_response(); + resp.extensions_mut().insert(AccessLogModel(model)); + return resp; + }; + let block_size = ie.block_size; + let indexer = ie.indexer.clone(); + drop(ie); + + let block_hashes: Vec = req + .block_hashes + .iter() + .map(|h| LocalBlockHash(*h as u64)) + .collect(); + let mut resp = match indexer.find_tiered_matches(block_hashes).await { + Ok(tiered) => ( + StatusCode::OK, + Json(serde_json::json!(TieredQueryResponse { + block_size, + tiered: WireTieredMatchDetails::from(&tiered), + })), + ) + .into_response(), + Err(e) => ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": e.to_string()})), + ) + .into_response(), + }; + resp.extensions_mut().insert(AccessLogModel(model)); + resp +} + #[derive(Deserialize)] struct ListenerControlRequest { instance_id: WorkerId, @@ -563,6 +615,7 @@ fn build_router(state: Arc, test_endpoints: bool) -> Router { post(query).layer(DefaultBodyLimit::max(QUERY_REQUEST_BODY_LIMIT_BYTES)), ) .route("/query_by_hash", post(query_by_hash)) + .route("/query_tiered_by_hash", post(query_tiered_by_hash)) .route("/dump", get(dump_events)) .route("/register_peer", post(register_peer)) .route("/deregister_peer", post(deregister_peer)) @@ -603,9 +656,32 @@ fn build_router(state: Arc, test_endpoints: bool) -> Router { router } +/// Serve a standalone indexer for `registry` on an ephemeral port. Returns the +/// base URL and the server task; abort the task to stop serving. +#[cfg(test)] +pub(crate) async fn spawn_test_indexer_server( + registry: Arc, +) -> (String, tokio::task::JoinHandle<()>) { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind test indexer"); + let base_url = format!("http://{}", listener.local_addr().expect("local addr")); + let app = create_router(Arc::new(AppState { + registry, + access_log_sink: None, + #[cfg(feature = "metrics")] + prom_registry: prometheus::Registry::new(), + })); + let task = tokio::spawn(async move { + let _ = axum::serve(listener, app).await; + }); + (base_url, task) +} + #[cfg(test)] mod tests { use super::*; + use crate::indexer::KvIndexerInterface; use axum::body::Body; use axum::http::{Request, StatusCode, header}; use tower::ServiceExt; @@ -626,6 +702,105 @@ mod tests { body } + #[tokio::test] + async fn query_tiered_by_hash_returns_wire_tiered_matches() { + use super::super::backend::test_util::store_event; + use crate::protocols::{StorageTier, WorkerWithDpRank}; + + let registry = Arc::new(WorkerRegistry::new(1)); + let key = RoutingPartitionId::new("model", "default"); + let indexer = registry.get_or_create_indexer(key, 4); + let worker = WorkerWithDpRank::new(7, 0); + indexer + .apply_event_routed(store_event(7, 0, 1, &[], &[11, 12], StorageTier::Device)) + .await + .unwrap(); + indexer + .apply_event_routed(store_event( + 7, + 0, + 2, + &[11, 12], + &[13], + StorageTier::HostPinned, + )) + .await + .unwrap(); + if let Indexer::Single { + primary, + lower_tier, + .. + } = &indexer + { + let _ = primary.flush().await; + for inner in lower_tier.all() { + let _ = inner.dump_events().await.unwrap(); + } + } + let app = create_router(Arc::new(AppState { + registry, + access_log_sink: None, + #[cfg(feature = "metrics")] + prom_registry: prometheus::Registry::new(), + })); + + let body = serde_json::to_vec(&TieredQueryByHashRequest { + block_hashes: vec![11, 12, 13], + model_name: "model".to_string(), + routing_group: "default".to_string(), + }) + .unwrap(); + let response = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/query_tiered_by_hash") + .header(header::CONTENT_TYPE, "application/json") + .body(Body::from(body)) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let bytes = axum::body::to_bytes(response.into_body(), usize::MAX) + .await + .unwrap(); + let parsed: TieredQueryResponse = serde_json::from_slice(&bytes).unwrap(); + assert_eq!(parsed.block_size, 4); + let tiered: TieredMatchDetails = parsed.tiered.into(); + assert_eq!( + tiered.device.overlap_scores.scores.get(&worker).copied(), + Some(2) + ); + assert_eq!( + tiered + .lower_tier + .get(&crate::protocols::StorageTier::HostPinned) + .and_then(|d| d.hits.get(&worker).copied()), + Some(1) + ); + + let missing = serde_json::to_vec(&TieredQueryByHashRequest { + block_hashes: vec![11], + model_name: "other".to_string(), + routing_group: "default".to_string(), + }) + .unwrap(); + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/query_tiered_by_hash") + .header(header::CONTENT_TYPE, "application/json") + .body(Body::from(missing)) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::NOT_FOUND); + } + #[tokio::test] async fn query_rejects_request_bodies_over_limit() { let app = create_router(Arc::new(AppState { diff --git a/lib/kv-router/src/services/selection/catalog.rs b/lib/kv-router/src/services/selection/catalog.rs index 45a88de98099..39b3babee241 100644 --- a/lib/kv-router/src/services/selection/catalog.rs +++ b/lib/kv-router/src/services/selection/catalog.rs @@ -127,6 +127,20 @@ impl WorkerCatalog { .flatten() } + /// Whether any schedulable worker in `key`'s partition can consume router + /// hints. Worker-level metadata, so one representative rank suffices. + pub(super) fn has_router_hint_capable_workers(&self, key: &RoutingPartitionId) -> bool { + self.workers.read().values().any(|record| { + record.lifecycle == WorkerLifecycle::Schedulable + && record.model_name == key.model_name + && record.routing_group == key.routing_group + && record + .router_hint_worker_type + .as_deref() + .is_some_and(|worker_type| !worker_type.is_empty()) + }) + } + pub(super) fn scheduler_configs_for_key( &self, key: &RoutingPartitionId, diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index e541d53749bd..a5104ce65b35 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -12,11 +12,14 @@ use tokio::sync::{mpsc, watch}; use tokio_util::sync::CancellationToken; use crate::identity::RoutingPartitionId; -use crate::indexer::{SharedKvCache, TieredMatchDetails}; +use crate::indexer::{ + LowerTierQueryOptions, RoutingDecisionHashes, SharedKvCache, TieredMatchDetails, +}; use crate::protocols::{ ActiveSequenceEvent, LocalBlockHash, PrefillLoadHint, RoutingConstraints, SharedCacheHits, - WorkerAffinityTarget, WorkerId, WorkerWithDpRank, + WorkerAffinityTarget, WorkerConfigLike, WorkerId, WorkerWithDpRank, }; +use crate::router_hint::{RouterHint, RouterHintCandidateSource, RouterHintRootCandidates}; use crate::scheduling::config::RouterConfigOverride; use crate::scheduling::selector::WorkerSelectionPolicy; use crate::scheduling::{ @@ -31,7 +34,7 @@ use crate::sequences::{ use crate::services::common::replica_sync::{ ReplicaSyncConfig, ScopedReplicaEvent, ScopedSequencePublisher, setup_scoped_replica_sync, }; -use crate::services::indexer::backend::Indexer; +use crate::services::indexer::backend::{Indexer, IndexerPolicy}; use crate::services::indexer::recovery; use crate::services::indexer::registry::WorkerRegistry; use crate::services::overlap::MooncakeOverlapSummary; @@ -73,6 +76,7 @@ struct PreparedSelectionInputs { isl_tokens: usize, overlap: OverlapSignals, shared_cache_hits: Option, + router_hint_candidates: Option, } struct SelectionOperation { @@ -104,6 +108,8 @@ struct ReservationBooking { expected_output_tokens: Option, track_prefill_tokens: bool, lora_name: Option, + /// Public block hashes to record into an approximate indexer once booked. + routing_hashes: Option>, } /// Host-supplied hooks applied to every partition scheduler the core creates. @@ -153,6 +159,10 @@ pub struct SelectionServiceConfig { pub port: u16, pub threads: usize, pub indexer_peers: Vec, + /// Base URL of a standalone indexer that serves the primary KV index. + /// When set, this service does not listen for worker KV events and + /// `indexer_peers` recovery is skipped. + pub remote_indexer_url: Option, pub replica_sync_port: Option, pub replica_sync_peers: Vec, pub kv_router_config: crate::config::KvRouterConfig, @@ -178,6 +188,10 @@ pub struct SelectionCore { /// Sweep task is started lazily from the first `ensure_entry`, which always /// runs inside the host runtime; construction itself may not. reservation_sweep_started: OnceCell<()>, + /// Whether this core subscribes to worker KV events itself. False when + /// events are disabled or when the primary indexer is a remote service + /// that workers publish to directly. + listens_for_kv_events: bool, indexer_registry: Arc, kv_router_config: crate::config::KvRouterConfig, worker_selection_policy_factory: Option, @@ -218,6 +232,7 @@ impl SelectionCore { .validate_config() .map_err(anyhow::Error::msg)?; let tracking_hash = Arc::new(TrackingHashContext::from_config(&kv_router_config)?); + let indexer_policy = IndexerPolicy::from_router_config(&kv_router_config)?; Ok(Self::new_inner( kv_router_config, indexer_threads, @@ -229,6 +244,7 @@ impl SelectionCore { true, cache_config, tracking_hash, + indexer_policy, )) } @@ -244,12 +260,15 @@ impl SelectionCore { signal_indexer_ready: bool, cache_config: SelectionCacheConfig, tracking_hash: Arc, + indexer_policy: IndexerPolicy, ) -> Self { let cancel_token = cancel_token.child_token(); let indexer_registry = Arc::new(WorkerRegistry::new_with_cancel_token( indexer_threads, cancel_token.clone(), )); + let listens_for_kv_events = kv_router_config.use_kv_events && !indexer_policy.is_remote(); + indexer_registry.set_indexer_policy(indexer_policy); if signal_indexer_ready { indexer_registry.signal_ready(); } @@ -258,6 +277,7 @@ impl SelectionCore { entries: Arc::new(RwLock::new(HashMap::new())), reservation_index: Arc::new(RwLock::new(HashMap::new())), reservation_sweep_started: OnceCell::new(), + listens_for_kv_events, indexer_registry, kv_router_config, worker_selection_policy_factory, @@ -431,8 +451,8 @@ impl SelectionCore { .kv_router_config .queueing_enabled(Some(&record.model_name)) .map_err(|error| SelectionError::BadRequest(error.to_string()))?; - let reasons = record - .missing_schedulable_metadata(queueing_enabled, self.kv_router_config.use_kv_events); + let reasons = + record.missing_schedulable_metadata(queueing_enabled, self.listens_for_kv_events); if !reasons.is_empty() { let updated = self .catalog @@ -445,7 +465,7 @@ impl SelectionCore { if let Err(error) = self.ensure_entry(&record) { return self.mark_incomplete_after_reconcile_error(worker_id, record.key(), error); } - if self.kv_router_config.use_kv_events + if self.listens_for_kv_events && let Err(error) = self.register_indexer_listeners(&record).await { self.cleanup_indexer_registration(&record).await; @@ -629,7 +649,7 @@ impl SelectionCore { } async fn cleanup_indexer_registration(&self, record: &WorkerCatalogRecord) { - if self.kv_router_config.use_kv_events { + if self.listens_for_kv_events { if let Err(error) = self .indexer_registry .deregister(record.worker_id, &record.model_name, &record.routing_group) @@ -783,12 +803,19 @@ impl SelectionCore { } let entry = self.ready_entry(&key)?; + // Router hints are attached to bookings only, and only when a worker in + // this partition can consume them and the indexer can retain the + // matched chain (local, event-driven, no approximate writes). + let retain_router_hint_chain = book + && entry.indexer.supports_router_hint_chain_retention() + && self.catalog.has_router_hint_capable_workers(&key); let PreparedSelectionInputs { block_hashes, sequence_hashes, isl_tokens, overlap, shared_cache_hits, + router_hint_candidates, } = self .prepare_selection_inputs( &entry, @@ -796,6 +823,7 @@ impl SelectionCore { self.kv_router_config .assume_kv_reuse(router_config_override.as_ref()), true, + retain_router_hint_chain, ) .await?; let mode = if book { @@ -839,14 +867,19 @@ impl SelectionCore { book.then(|| sequence_hashes.iter().map(|hash| *hash as i64).collect()); let response_isl_tokens = book.then_some(isl_tokens); let response_track_prefill_tokens = book.then_some(track_prefill_tokens); + // Bookings (now, or later via the pending-selection cache) are recorded + // into an approximate indexer; keep the public hashes for that. + let routing_hashes = (entry.indexer.records_routing_decisions() + && (book || cached_inputs.is_some())) + .then(|| block_hashes.clone()); let schedule_request = ScheduleRequest { mode, token_seq: Some(sequence_hashes), block_hashes: Some(block_hashes), isl_tokens, overlap, - kv_transfer_candidates: None, - retain_kv_transfer_chain: false, + router_hint_candidates, + retain_router_hint_chain, router_config_override, lora_name: prompt.lora_name, priority_jump, @@ -910,6 +943,16 @@ impl SelectionCore { .map(|(threshold, total_kv_blocks)| { potential_decode_blocks as f64 > threshold * total_kv_blocks as f64 }); + let router_hint = if retain_router_hint_chain { + router_hint_for_selection( + &self.catalog.scheduler_configs_for_key(&key), + response.best_worker, + response.target_cached_prefix_blocks, + response.router_hint_candidates.as_ref(), + ) + } else { + None + }; let worker_load = advisory_load.map(|load| SelectionWorkerLoad { active_prefill_tokens: load.active_prefill_tokens, prefill_token_capacity: load.prefill_token_capacity, @@ -922,6 +965,10 @@ impl SelectionCore { if book && let Some(selection_id) = selection_id.as_deref() { self.record_reservation(selection_id, &key); + if let Some(hashes) = routing_hashes.clone() { + self.record_routing_decision(&entry, response.best_worker, hashes) + .await; + } } if let Some((cache_id, sequence_hashes, lora_name, track_prefill_tokens)) = cached_inputs { @@ -936,6 +983,7 @@ impl SelectionCore { expected_output_tokens, track_prefill_tokens, lora_name, + routing_hashes, }, Instant::now(), ); @@ -957,6 +1005,7 @@ impl SelectionCore { potential_decode_blocks, decode_busy, worker_load, + router_hint, }) } @@ -1015,6 +1064,7 @@ impl SelectionCore { expected_output_tokens: pending.expected_output_tokens, track_prefill_tokens, lora_name: pending.lora_name.clone(), + routing_hashes: pending.routing_hashes.clone(), }, ) .await @@ -1096,6 +1146,16 @@ impl SelectionCore { .and_then(|cfg| cfg.track_prefill_tokens) .unwrap_or(self.kv_router_config.router_track_prefill_tokens) }); + // Hash-only reservations (sequence hashes without block hashes) carry + // nothing an indexer can key on; recording is skipped for them. + let can_record = entry.indexer.records_routing_decisions() + && (req.prompt.token_ids.is_some() || req.prompt.block_hashes.is_some()); + let routing_hashes = can_record + .then(|| { + req.prompt + .block_hashes_for_indexer(entry.block_size, entry.is_eagle) + }) + .transpose()?; self.finalize_reservation( entry, @@ -1109,6 +1169,7 @@ impl SelectionCore { expected_output_tokens: req.expected_output_tokens, track_prefill_tokens, lora_name: req.prompt.lora_name, + routing_hashes, }, ) .await @@ -1132,6 +1193,7 @@ impl SelectionCore { expected_output_tokens, track_prefill_tokens, lora_name, + routing_hashes, } = booking; // Strict booking: never lazily recreate a worker/rank removed since the @@ -1149,6 +1211,9 @@ impl SelectionCore { }) .await?; self.record_reservation(&selection_id, &key); + if let Some(hashes) = routing_hashes { + self.record_routing_decision(&entry, worker, hashes).await; + } Ok(ReservationResponse { selection_id, @@ -1170,6 +1235,36 @@ impl SelectionCore { self.reservation_index.write().remove(selection_id); } + /// Record a booked routing decision into the partition's approximate + /// indexer (side or primary). The booking already landed, so a failure + /// here only costs predicted cache credit; it is logged, not returned. + async fn record_routing_decision( + &self, + entry: &SelectionEntry, + worker: WorkerWithDpRank, + block_hashes: Vec, + ) { + if block_hashes.is_empty() { + return; + } + if let Err(error) = entry + .indexer + .record_routing_decision( + worker, + RoutingDecisionHashes::from_local_hashes(block_hashes), + ) + .await + { + tracing::warn!( + %error, + key = %entry.key, + worker_id = worker.worker_id, + dp_rank = worker.dp_rank, + "Failed to record routing decision into approximate indexer" + ); + } + } + /// Entries to try for a lifecycle call on `selection_id`: the indexed /// partition first, then every other initialized partition. The fallback /// covers bookings mirrored from replica peers, which never pass through @@ -1293,6 +1388,7 @@ impl SelectionCore { self.kv_router_config .assume_kv_reuse(req.router_config_override.as_ref()), false, + false, ) .await?; let track_prefill_tokens = req @@ -1347,6 +1443,7 @@ impl SelectionCore { prompt: &PromptRequest, assume_kv_reuse: bool, query_shared_cache: bool, + retain_router_hint_chain: bool, ) -> Result { let normalized = prompt.normalize_for_selection( entry.is_eagle, @@ -1362,7 +1459,12 @@ impl SelectionCore { } else { entry .indexer - .find_tiered_matches(normalized.block_hashes.clone()) + .find_tiered_matches_with_options( + normalized.block_hashes.clone(), + LowerTierQueryOptions { + retain_router_hint_chain, + }, + ) .await .map_err(|error| SelectionError::Internal(error.to_string())) } @@ -1388,6 +1490,9 @@ impl SelectionCore { let tiered = tiered?; let overlap = OverlapAnalysis::new(&self.kv_router_config, entry.block_size, &tiered).signals(); + let router_hint_candidates = retain_router_hint_chain + .then(|| tiered.router_hint_root_candidates().cloned()) + .flatten(); drop(tiered); Ok(PreparedSelectionInputs { block_hashes: normalized.block_hashes, @@ -1395,6 +1500,7 @@ impl SelectionCore { isl_tokens: normalized.isl_tokens, overlap, shared_cache_hits, + router_hint_candidates, }) } @@ -1471,6 +1577,69 @@ fn spawn_reservation_index_sweep( }); } +/// Pick the best router-hint source for `target`: a same-role worker (or +/// cache owner) holding a longer root-aligned prefix than the target's own +/// `target_cached_prefix_blocks`, with a non-empty control endpoint. Mirrors the +/// frontend `KvRouter::router_hint_for_selection`. +fn router_hint_for_selection( + configs: &HashMap, + target: WorkerWithDpRank, + target_cached_prefix_blocks: u32, + candidates: Option<&RouterHintRootCandidates>, +) -> Option { + let candidates = candidates?; + let target_config = configs.get(&target.worker_id)?; + let target_metadata = target_config.router_hint_metadata_for_dp_rank(target.dp_rank)?; + + let prefix_blocks_to_beat = usize::try_from(target_cached_prefix_blocks).unwrap_or(usize::MAX); + let (source, block_hashes) = + candidates.best_source(prefix_blocks_to_beat, |source| match source { + RouterHintCandidateSource::Worker(worker) => { + worker != target + && configs.get(&worker.worker_id).is_some_and(|config| { + config + .router_hint_metadata_for_dp_rank(worker.dp_rank) + .is_some_and(|source_metadata| { + source_metadata.worker_type == target_metadata.worker_type + && source_metadata + .source_control_endpoint + .is_some_and(|endpoint| !endpoint.is_empty()) + }) + }) + } + RouterHintCandidateSource::CacheOwner(owner) => candidates + .routing_snapshot + .as_ref() + .and_then(|snapshot| snapshot.router_hint_source(owner)) + .is_some_and(|source| { + source.attached_worker != Some(target) + && source.metadata.worker_type == target_metadata.worker_type + && !source.metadata.source_control_endpoint.is_empty() + }), + })?; + let source_control_endpoint = match source { + RouterHintCandidateSource::Worker(worker) => configs + .get(&worker.worker_id)? + .router_hint_metadata_for_dp_rank(worker.dp_rank)? + .source_control_endpoint? + .to_string(), + RouterHintCandidateSource::CacheOwner(owner) => candidates + .routing_snapshot + .as_ref()? + .router_hint_source(owner)? + .metadata + .source_control_endpoint + .clone(), + }; + if block_hashes.is_empty() { + return None; + } + Some(RouterHint { + source_control_endpoint, + block_hashes, + }) +} + fn tracking_scope(entry: &SelectionEntry) -> TrackingHashScope<'_> { TrackingHashScope { partition: entry.key.as_ref(), @@ -1520,6 +1689,8 @@ mod tests { kv_transfer_domain: None, kv_transfer_enforcement: None, kv_transfer_preferred_weight: None, + router_hint_worker_type: None, + router_hint_source_control_endpoints: HashMap::new(), } } @@ -1644,6 +1815,7 @@ mod tests { true, SelectionCacheConfig::default(), tracking_hash, + IndexerPolicy::from_router_config(&test_config(false)).expect("indexer policy"), ); core.upsert_worker(worker(1)).await.expect("worker upsert"); @@ -1670,6 +1842,7 @@ mod tests { let tracking_hash = Arc::new( TrackingHashContext::from_config(&config).expect("valid tracking hash configuration"), ); + let indexer_policy = IndexerPolicy::from_router_config(&config).expect("indexer policy"); SelectionCore::new_inner( config, 1, @@ -1681,9 +1854,313 @@ mod tests { true, SelectionCacheConfig::default(), tracking_hash, + indexer_policy, ) } + async fn wait_for_overlap( + core: &SelectionCore, + request: impl Fn() -> SelectRequest, + ) -> SelectResponse { + tokio::time::timeout(Duration::from_secs(2), async { + loop { + let response = core.select(request()).await.expect("select"); + if response.overlap.longest_matched > 0 { + return response; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("approximate indexer never credited the booked prompt") + } + + #[tokio::test] + async fn bookings_populate_the_approximate_primary_without_kv_events() { + // use_kv_events=false: the primary is approximate and bookings feed it. + let core = SelectionCore::new_local( + test_config(false), + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + core.upsert_worker(worker(2)).await.expect("worker upsert"); + + // Query-only selection records nothing. + let first = core.select(select_request()).await.expect("select"); + assert_eq!(first.overlap.longest_matched, 0); + + // select_and_reserve records the routed prefix for the chosen worker. + let booked = core + .select_and_reserve(reserve_request("booked")) + .await + .expect("reserve"); + let credited = wait_for_overlap(&core, || { + let mut request = select_request(); + request.allowed_worker_ids = Some(HashSet::from([booked.worker_id])); + request + }) + .await; + assert_eq!(credited.worker_id, booked.worker_id); + assert_eq!(credited.overlap.longest_matched, 4); + + // The cached replay path records too, on a different prompt. + let prompt_b = || PromptRequest { + token_ids: Some(vec![5, 6, 7, 8]), + ..PromptRequest::default() + }; + let mut request = select_request(); + request.prompt = prompt_b(); + request.selection_id = Some("cached".to_string()); + request.allowed_worker_ids = Some(HashSet::from([2])); + core.select(request).await.expect("select"); + core.create_reservation(ReservationRequest { + model_name: "model".to_string(), + routing_group: "default".to_string(), + selection_id: "cached".to_string(), + worker_id: None, + dp_rank: None, + prompt: PromptRequest::default(), + router_config_override: None, + expected_output_tokens: None, + effective_prefill_tokens: None, + track_prefill_tokens: None, + }) + .await + .expect("cached reservation"); + let credited = wait_for_overlap(&core, || { + let mut request = select_request(); + request.prompt = prompt_b(); + request.allowed_worker_ids = Some(HashSet::from([2])); + request + }) + .await; + assert_eq!(credited.worker_id, 2); + + // And the explicit reservation form. + let prompt_c = || PromptRequest { + token_ids: Some(vec![9, 10, 11, 12]), + ..PromptRequest::default() + }; + core.create_reservation(ReservationRequest { + model_name: "model".to_string(), + routing_group: "default".to_string(), + selection_id: "explicit".to_string(), + worker_id: Some(1), + dp_rank: None, + prompt: prompt_c(), + router_config_override: None, + expected_output_tokens: None, + effective_prefill_tokens: None, + track_prefill_tokens: None, + }) + .await + .expect("explicit reservation"); + let credited = wait_for_overlap(&core, || { + let mut request = select_request(); + request.prompt = prompt_c(); + request.allowed_worker_ids = Some(HashSet::from([1])); + request + }) + .await; + assert_eq!(credited.worker_id, 1); + } + + #[tokio::test] + async fn remote_indexer_serves_selection_without_local_kv_listeners() { + use crate::indexer::KvIndexerInterface; + use crate::protocols::{BlockHashOptions, StorageTier, compute_block_hash_for_seq}; + use crate::services::indexer::registry::WorkerRegistry; + use crate::services::indexer::server::spawn_test_indexer_server; + + // The standalone indexer holds worker 2's cache for the test prompt. + let key = RoutingPartitionId::new("model", "default"); + let served = Arc::new(WorkerRegistry::new(1)); + let served_indexer = served.get_or_create_indexer(key.clone(), 4); + let hashes: Vec = + compute_block_hash_for_seq(&[1, 2, 3, 4], 4, BlockHashOptions::default()) + .into_iter() + .map(|hash| hash.0) + .collect(); + served_indexer + .apply_event_routed(store_event(2, 0, 1, &[], &hashes, StorageTier::Device)) + .await + .unwrap(); + if let Indexer::Single { primary, .. } = &served_indexer { + let _ = primary.flush().await; + } + let (base_url, server) = spawn_test_indexer_server(served).await; + + // use_kv_events=true, but the primary is remote: workers need no + // kv_events endpoints and no ZMQ listener is started here. + let config = test_config(true); + let tracking_hash = Arc::new( + TrackingHashContext::from_config(&config).expect("valid tracking hash configuration"), + ); + let indexer_policy = IndexerPolicy::from_router_config(&config) + .expect("indexer policy") + .with_remote_indexer(base_url) + .expect("remote policy"); + let core = SelectionCore::new_inner( + config, + 1, + CancellationToken::new(), + None, + None, + SelectionHostHooks::default(), + WorkerType::Aggregated, + true, + SelectionCacheConfig::default(), + tracking_hash, + indexer_policy, + ); + assert!(!core.listens_for_kv_events); + for worker_id in [1, 2] { + let record = core + .upsert_worker(worker(worker_id)) + .await + .expect("worker upsert"); + assert_eq!(record.lifecycle, WorkerLifecycle::Schedulable, "{record:?}"); + } + + let response = core.select(select_request()).await.expect("select"); + assert_eq!( + response.worker_id, 2, + "remote cache credit steers selection" + ); + assert_eq!(response.overlap.longest_matched, 4); + + core.delete_worker(2).await.expect("delete worker"); + server.abort(); + } + + #[tokio::test] + async fn booked_selection_attaches_router_hint_from_a_better_source() { + use crate::indexer::KvIndexerInterface; + use crate::protocols::{BlockHashOptions, StorageTier, compute_block_hash_for_seq}; + + let core = SelectionCore::new_local( + test_config(true), + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ); + for worker_id in [1, 2] { + let mut request = worker_with_kv_events(worker_id); + request.router_hint_worker_type = Some("decode".to_string()); + request.router_hint_source_control_endpoints = + HashMap::from([(0, format!("tcp://worker-{worker_id}:9000"))]); + core.upsert_worker(request).await.expect("worker upsert"); + } + let key = RoutingPartitionId::new("model", "default"); + let entry = core.entry(&key).expect("entry"); + assert!(entry.indexer.supports_router_hint_chain_retention()); + + // Worker 1 holds both blocks of the prompt; worker 2 holds nothing. + let tokens: Vec = (1..=8).collect(); + let hashes: Vec = compute_block_hash_for_seq(&tokens, 4, BlockHashOptions::default()) + .into_iter() + .map(|hash| hash.0) + .collect(); + assert_eq!(hashes.len(), 2); + entry + .indexer + .apply_event_routed(store_event(1, 0, 1, &[], &hashes, StorageTier::Device)) + .await + .unwrap(); + if let Indexer::Single { primary, .. } = &entry.indexer { + let _ = primary.flush().await; + } + let prompt = || PromptRequest { + token_ids: Some(tokens.clone()), + ..PromptRequest::default() + }; + + // Booking on worker 2: worker 1 is a same-role source with a longer prefix. + let mut request = reserve_request("to-worker-2"); + request.prompt = prompt(); + request.pinned_worker = Some(WorkerWithDpRank::new(2, 0)); + let response = core.select_and_reserve(request).await.expect("reserve"); + let hint = response.router_hint.expect("router hint for worker 2"); + assert_eq!(hint.source_control_endpoint, "tcp://worker-1:9000"); + assert_eq!(hint.block_hashes.len(), 2); + + // Booking on worker 1 itself: nothing holds a longer prefix. + let mut request = reserve_request("to-worker-1"); + request.prompt = prompt(); + request.pinned_worker = Some(WorkerWithDpRank::new(1, 0)); + let response = core.select_and_reserve(request).await.expect("reserve"); + assert!(response.router_hint.is_none()); + + // Query-only selections never carry a hint. + let mut request = select_request(); + request.prompt = prompt(); + request.pinned_worker = Some(WorkerWithDpRank::new(2, 0)); + let response = core.select(request).await.expect("select"); + assert!(response.router_hint.is_none()); + } + + #[tokio::test] + async fn router_hint_needs_capable_workers() { + use crate::indexer::KvIndexerInterface; + use crate::protocols::{BlockHashOptions, StorageTier, compute_block_hash_for_seq}; + + let core = SelectionCore::new_local( + test_config(true), + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ); + for worker_id in [1, 2] { + core.upsert_worker(worker_with_kv_events(worker_id)) + .await + .expect("worker upsert"); + } + let entry = core + .entry(&RoutingPartitionId::new("model", "default")) + .expect("entry"); + let hashes: Vec = + compute_block_hash_for_seq(&[1, 2, 3, 4], 4, BlockHashOptions::default()) + .into_iter() + .map(|hash| hash.0) + .collect(); + entry + .indexer + .apply_event_routed(store_event(1, 0, 1, &[], &hashes, StorageTier::Device)) + .await + .unwrap(); + if let Indexer::Single { primary, .. } = &entry.indexer { + let _ = primary.flush().await; + } + let mut request = reserve_request("plain"); + request.pinned_worker = Some(WorkerWithDpRank::new(2, 0)); + let response = core.select_and_reserve(request).await.expect("reserve"); + assert!(response.router_hint.is_none()); + } + + #[tokio::test] + async fn event_driven_indexer_does_not_record_bookings() { + let core = SelectionCore::new_local( + test_config(true), + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ); + core.upsert_worker(worker_with_kv_events(1)) + .await + .expect("worker upsert"); + core.select_and_reserve(reserve_request("booked")) + .await + .expect("reserve"); + core.free_reservation("booked").await.expect("free"); + for _ in 0..3 { + let response = core.select(select_request()).await.expect("select"); + assert_eq!(response.overlap.longest_matched, 0); + tokio::task::yield_now().await; + } + } + /// Picker that records what worker selection saw and always takes row 0. struct CapturingPicker { observed: Arc>>, diff --git a/lib/kv-router/src/services/selection/pending.rs b/lib/kv-router/src/services/selection/pending.rs index 7008d251cdaa..7a3a70b1d2a5 100644 --- a/lib/kv-router/src/services/selection/pending.rs +++ b/lib/kv-router/src/services/selection/pending.rs @@ -13,7 +13,7 @@ use dynamo_tokens::SequenceHash; use parking_lot::Mutex; use crate::identity::RoutingPartitionId; -use crate::protocols::WorkerWithDpRank; +use crate::protocols::{LocalBlockHash, WorkerWithDpRank}; /// How long a pending selection lives before it is evicted. const SELECTION_CACHE_TTL: Duration = Duration::from_secs(120); @@ -56,6 +56,9 @@ pub(super) struct PendingSelection { pub expected_output_tokens: Option, pub track_prefill_tokens: bool, pub lora_name: Option, + /// Public block hashes retained only when the partition indexer records + /// routing decisions, so the replayed booking can be recorded too. + pub routing_hashes: Option>, } struct Entry { @@ -71,6 +74,9 @@ type CacheKey = (RoutingPartitionId, String); /// Approximate resident bytes: sequence hashes plus id and scope strings. fn entry_bytes(key: &CacheKey, selection: &PendingSelection) -> usize { selection.sequence_hashes.len() * std::mem::size_of::() + + selection.routing_hashes.as_ref().map_or(0, |hashes| { + hashes.len() * std::mem::size_of::() + }) + key.1.len() + key.0.model_name.len() + key.0.routing_group.len() @@ -267,6 +273,7 @@ mod tests { expected_output_tokens: Some(16), track_prefill_tokens: true, lora_name: None, + routing_hashes: None, } } diff --git a/lib/kv-router/src/services/selection/service.rs b/lib/kv-router/src/services/selection/service.rs index 122790230fba..574c0c4d4ec0 100644 --- a/lib/kv-router/src/services/selection/service.rs +++ b/lib/kv-router/src/services/selection/service.rs @@ -11,6 +11,7 @@ use crate::scheduling::PotentialLoad; use crate::services::common::replica_sync::{ PeerManager, ReplicaPeerError, ReplicaSyncRuntime, setup_replica_sync, }; +use crate::services::indexer::backend::IndexerPolicy; use crate::tracking_hash::TrackingHashContext; use super::core::{SelectionCore, SelectionHostHooks, SelectionServiceConfig}; @@ -34,6 +35,7 @@ pub struct SelectionServiceBuilder { worker_type: WorkerType, worker_selection_policy_registry: WorkerSelectionPolicyRegistry, host_hooks: SelectionHostHooks, + remote_indexer_url: Option, } /// Warn when a host does not construct workers for explicitly configured policy roles. @@ -72,9 +74,18 @@ impl SelectionServiceBuilder { worker_type, worker_selection_policy_registry, host_hooks: SelectionHostHooks::default(), + remote_indexer_url: None, } } + /// Serve the primary KV index from a standalone indexer at `base_url` + /// instead of subscribing to worker KV events in this process. Requires + /// `use_kv_events=true`; `indexer_peers` recovery is skipped. + pub fn remote_indexer(mut self, base_url: impl Into) -> Self { + self.remote_indexer_url = Some(base_url.into()); + self + } + /// Attach host-owned hooks (prefill-load estimator, overload exclusion /// set, hard-availability set). Applied to every partition scheduler this /// service creates. @@ -112,6 +123,18 @@ impl SelectionServiceBuilder { .worker_selection_policy_registry .resolve_for_worker_type(&self.kv_router_config, self.worker_type)?; let tracking_hash = Arc::new(TrackingHashContext::from_config(&self.kv_router_config)?); + let mut indexer_policy = IndexerPolicy::from_router_config(&self.kv_router_config)?; + if let Some(base_url) = &self.remote_indexer_url { + indexer_policy = indexer_policy.with_remote_indexer(base_url.clone())?; + if !self.indexer_peers.is_empty() { + tracing::warn!( + remote_indexer_url = %base_url, + "indexer_peers are ignored when the primary indexer is remote" + ); + } + tracing::info!(remote_indexer_url = %base_url, "Using remote KV indexer"); + } + let recover_from_peers = !self.indexer_peers.is_empty() && !indexer_policy.is_remote(); let cancel_token = CancellationToken::new(); let mut startup_guard = StartupGuard::new(cancel_token.clone()); let replica_runtime = setup_replica_sync( @@ -138,9 +161,10 @@ impl SelectionServiceBuilder { false, self.selection_cache, tracking_hash, + indexer_policy, )); - if !self.indexer_peers.is_empty() { + if recover_from_peers { match core.recover_indexer_from_peers(&self.indexer_peers).await { Ok(true) => tracing::info!("Selection indexer recovery completed"), Ok(false) => { @@ -198,6 +222,9 @@ impl SelectionServiceConfig { if let Some(port) = self.replica_sync_port { builder = builder.replica_sync(port, self.replica_sync_peers.clone()); } + if let Some(url) = &self.remote_indexer_url { + builder = builder.remote_indexer(url.clone()); + } builder } } diff --git a/lib/kv-router/src/services/selection/tests.rs b/lib/kv-router/src/services/selection/tests.rs index 1a9824442765..8d194cc98aa0 100644 --- a/lib/kv-router/src/services/selection/tests.rs +++ b/lib/kv-router/src/services/selection/tests.rs @@ -1759,11 +1759,12 @@ async fn cached_booking_honors_prefill_tracking() { assert_eq!(loads[0]["loads"][0]["potential_prefill_tokens"], 0); // A select-time override is captured and replayed: this booking tracks - // prefill load despite the config default. + // prefill load despite the config default. A distinct prompt keeps the + // approximate indexer (populated by req-1's booking) from crediting it. let select_response = post( app.clone(), "/select", - r#"{"model_name":"model","selection_id":"req-2","token_ids":[1,2,3,4],"router_config_override":{"track_prefill_tokens":true}}"#, + r#"{"model_name":"model","selection_id":"req-2","token_ids":[5,6,7,8],"router_config_override":{"track_prefill_tokens":true}}"#, ) .await; assert_eq!(select_response.status(), StatusCode::OK); @@ -1796,6 +1797,7 @@ async fn selector_replica_sync_propagates_request_lifecycle() { port: 8092, threads: 1, indexer_peers: Vec::new(), + remote_indexer_url: None, replica_sync_port: Some(port_a), replica_sync_peers: Vec::new(), kv_router_config: test_config(), diff --git a/lib/kv-router/src/services/selection/types.rs b/lib/kv-router/src/services/selection/types.rs index 4baaad5244f8..78f633a1a1d8 100644 --- a/lib/kv-router/src/services/selection/types.rs +++ b/lib/kv-router/src/services/selection/types.rs @@ -7,9 +7,10 @@ use serde::{Deserialize, Serialize}; use crate::identity::{RoutingPartitionId, default_routing_group}; use crate::protocols::{ - DpRank, KvTransferEnforcement, RoutingConstraints, WorkerAffinityTarget, WorkerConfigLike, - WorkerId, WorkerWithDpRank, + DpRank, KvTransferEnforcement, RouterHintWorkerMetadata, RoutingConstraints, + WorkerAffinityTarget, WorkerConfigLike, WorkerId, WorkerWithDpRank, }; +use crate::router_hint::RouterHint; use crate::scheduling::config::RouterConfigOverride; pub use crate::scheduling::{OverlapScoresResponse, SharedCacheOverlapScore, WorkerOverlapScore}; use crate::scheduling::{PotentialLoad, SessionContext, WorkerSelectionInputTrigger}; @@ -49,6 +50,13 @@ pub struct SelectionWorkerConfig { pub kv_transfer_domain: Option, pub kv_transfer_enforcement: Option, pub kv_transfer_preferred_weight: Option, + /// Backend role used to match router-hint sources to targets. Presence + /// means the worker can consume router hints. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub router_hint_worker_type: Option, + /// Per-global-DP-rank KV control endpoints a hint target fetches from. + #[serde(default, skip_serializing_if = "HashMap::is_empty")] + pub router_hint_source_control_endpoints: HashMap, } impl WorkerConfigLike for SelectionWorkerConfig { @@ -91,6 +99,24 @@ impl WorkerConfigLike for SelectionWorkerConfig { fn kv_transfer_preferred_weight(&self) -> Option { self.kv_transfer_preferred_weight } + + fn router_hint_metadata_for_dp_rank( + &self, + dp_rank: DpRank, + ) -> Option> { + let worker_type = self.router_hint_worker_type.as_deref()?; + if worker_type.is_empty() { + return None; + } + Some(RouterHintWorkerMetadata { + worker_type, + source_control_endpoint: self + .router_hint_source_control_endpoints + .get(&dp_rank) + .map(String::as_str) + .filter(|endpoint| !endpoint.is_empty()), + }) + } } #[derive(Debug, Clone, Serialize)] @@ -118,6 +144,10 @@ pub struct WorkerCatalogRecord { pub kv_transfer_domain: Option, pub kv_transfer_enforcement: Option, pub kv_transfer_preferred_weight: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub router_hint_worker_type: Option, + #[serde(default, skip_serializing_if = "HashMap::is_empty")] + pub router_hint_source_control_endpoints: HashMap, #[serde(default, skip_serializing_if = "Vec::is_empty")] pub not_schedulable_reasons: Vec, } @@ -145,6 +175,8 @@ impl WorkerCatalogRecord { kv_transfer_domain: req.kv_transfer_domain, kv_transfer_enforcement: req.kv_transfer_enforcement, kv_transfer_preferred_weight: req.kv_transfer_preferred_weight, + router_hint_worker_type: req.router_hint_worker_type, + router_hint_source_control_endpoints: req.router_hint_source_control_endpoints, not_schedulable_reasons: Vec::new(), } } @@ -181,6 +213,8 @@ impl WorkerCatalogRecord { kv_transfer_domain: self.kv_transfer_domain.clone(), kv_transfer_enforcement: self.kv_transfer_enforcement, kv_transfer_preferred_weight: self.kv_transfer_preferred_weight, + router_hint_worker_type: self.router_hint_worker_type.clone(), + router_hint_source_control_endpoints: self.router_hint_source_control_endpoints.clone(), }) } @@ -255,6 +289,8 @@ impl Default for WorkerRequest { kv_transfer_domain: None, kv_transfer_enforcement: None, kv_transfer_preferred_weight: None, + router_hint_worker_type: None, + router_hint_source_control_endpoints: HashMap::new(), } } } @@ -285,6 +321,13 @@ pub struct WorkerRequest { pub kv_transfer_domain: Option, pub kv_transfer_enforcement: Option, pub kv_transfer_preferred_weight: Option, + /// Backend role for router-hint source/target matching. Set it to mark the + /// worker as able to consume router hints. + #[serde(default)] + pub router_hint_worker_type: Option, + /// Per-global-DP-rank KV control endpoints this worker can serve hints from. + #[serde(default)] + pub router_hint_source_control_endpoints: HashMap, } #[derive(Debug, Deserialize)] @@ -305,6 +348,10 @@ pub struct WorkerPatchRequest { pub kv_transfer_domain: Option, pub kv_transfer_enforcement: Option, pub kv_transfer_preferred_weight: Option, + #[serde(default)] + pub router_hint_worker_type: Option, + #[serde(default)] + pub router_hint_source_control_endpoints: Option>, } impl WorkerCatalogRecord { @@ -360,6 +407,12 @@ impl WorkerCatalogRecord { if patch.kv_transfer_preferred_weight.is_some() { self.kv_transfer_preferred_weight = patch.kv_transfer_preferred_weight; } + if patch.router_hint_worker_type.is_some() { + self.router_hint_worker_type = patch.router_hint_worker_type; + } + if let Some(endpoints) = patch.router_hint_source_control_endpoints { + self.router_hint_source_control_endpoints = endpoints; + } } } @@ -565,6 +618,11 @@ pub struct SelectResponse { /// selections (`SelectRequest::advisory`). #[serde(skip_serializing_if = "Option::is_none")] pub worker_load: Option, + /// Source the chosen worker can fetch a longer cached prefix from. Present + /// only for bookings when the partition has router-hint-capable workers, + /// the indexer can retain the matched chain, and a better source exists. + #[serde(skip_serializing_if = "Option::is_none")] + pub router_hint: Option, } /// Load snapshot of the chosen worker, as the scheduler projected it for this From e7ef305f5f4ad91c50b526de5ae0815031228436 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Wed, 9 Sep 2026 01:52:52 -0700 Subject: [PATCH 004/102] fix(router): use typed KV transfer hints in standalone selection Signed-off-by: Ishan Dhanani --- lib/kv-router/src/services/indexer/backend.rs | 18 +-- .../src/services/selection/core/mod.rs | 103 ++++++++++-------- lib/kv-router/src/services/selection/types.rs | 12 +- 3 files changed, 75 insertions(+), 58 deletions(-) diff --git a/lib/kv-router/src/services/indexer/backend.rs b/lib/kv-router/src/services/indexer/backend.rs index 6918a23c5ad3..51f46075975c 100644 --- a/lib/kv-router/src/services/indexer/backend.rs +++ b/lib/kv-router/src/services/indexer/backend.rs @@ -426,7 +426,7 @@ impl Indexer { /// Router-hint chain retention needs a local primary whose hashes come /// only from engine events. - pub fn supports_router_hint_chain_retention(&self) -> bool { + pub fn supports_kv_transfer_chain_retention(&self) -> bool { !self.is_remote() && !self.records_routing_decisions() } @@ -670,15 +670,15 @@ impl Indexer { /// [`Self::find_tiered_matches`] with lower-tier query options. Router-hint /// chain retention is honored only by a local event-driven primary (see - /// [`Self::supports_router_hint_chain_retention`]); other shapes ignore it. + /// [`Self::supports_kv_transfer_chain_retention`]); other shapes ignore it. pub async fn find_tiered_matches_with_options( &self, sequence: Vec, options: LowerTierQueryOptions, ) -> std::result::Result { let options = LowerTierQueryOptions { - retain_router_hint_chain: options.retain_router_hint_chain - && self.supports_router_hint_chain_retention(), + retain_kv_transfer_chain: options.retain_kv_transfer_chain + && self.supports_kv_transfer_chain_retention(), }; let (device, lower_tier) = match self { Indexer::Single { @@ -689,7 +689,7 @@ impl Indexer { let device = primary .find_match_details_with_options( sequence.clone(), - options.retain_router_hint_chain, + options.retain_kv_transfer_chain, ) .await?; let lt = query_lower_tiers_with_options(lower_tier, &sequence, &device, options); @@ -703,7 +703,7 @@ impl Indexer { let device: MatchDetails = primary.backend().find_match_details_impl_with_options( &sequence, false, - options.retain_router_hint_chain, + options.retain_kv_transfer_chain, ); let lt = query_lower_tiers_with_options(lower_tier, &sequence, &device, options); (device, lt) @@ -1118,7 +1118,7 @@ mod tests { async fn event_driven_policy_ignores_routing_decisions() { let indexer = policy_indexer(1, IndexerPolicy::event_driven()); assert!(!indexer.records_routing_decisions()); - assert!(indexer.supports_router_hint_chain_retention()); + assert!(indexer.supports_kv_transfer_chain_retention()); let worker = WorkerWithDpRank::new(7, 0); indexer .record_routing_decision( @@ -1146,7 +1146,7 @@ mod tests { }, ); assert!(indexer.records_routing_decisions()); - assert!(!indexer.supports_router_hint_chain_retention()); + assert!(!indexer.supports_kv_transfer_chain_retention()); let worker = WorkerWithDpRank::new(7, 0); indexer .record_routing_decision( @@ -1301,7 +1301,7 @@ mod tests { ); assert!(indexer.is_remote()); assert!(indexer.records_routing_decisions()); - assert!(!indexer.supports_router_hint_chain_retention()); + assert!(!indexer.supports_kv_transfer_chain_retention()); // Local events are dropped: the remote service owns the primary. indexer diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index a5104ce65b35..1e31d60fa287 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -15,11 +15,13 @@ use crate::identity::RoutingPartitionId; use crate::indexer::{ LowerTierQueryOptions, RoutingDecisionHashes, SharedKvCache, TieredMatchDetails, }; +use crate::kv_hints::{ + KvHint, KvHintAction, KvSourceLocationsPayload, KvTransferCandidateSource, KvTransferCandidates, +}; use crate::protocols::{ ActiveSequenceEvent, LocalBlockHash, PrefillLoadHint, RoutingConstraints, SharedCacheHits, WorkerAffinityTarget, WorkerConfigLike, WorkerId, WorkerWithDpRank, }; -use crate::router_hint::{RouterHint, RouterHintCandidateSource, RouterHintRootCandidates}; use crate::scheduling::config::RouterConfigOverride; use crate::scheduling::selector::WorkerSelectionPolicy; use crate::scheduling::{ @@ -76,7 +78,7 @@ struct PreparedSelectionInputs { isl_tokens: usize, overlap: OverlapSignals, shared_cache_hits: Option, - router_hint_candidates: Option, + kv_transfer_candidates: Option, } struct SelectionOperation { @@ -806,8 +808,8 @@ impl SelectionCore { // Router hints are attached to bookings only, and only when a worker in // this partition can consume them and the indexer can retain the // matched chain (local, event-driven, no approximate writes). - let retain_router_hint_chain = book - && entry.indexer.supports_router_hint_chain_retention() + let retain_kv_transfer_chain = book + && entry.indexer.supports_kv_transfer_chain_retention() && self.catalog.has_router_hint_capable_workers(&key); let PreparedSelectionInputs { block_hashes, @@ -815,7 +817,7 @@ impl SelectionCore { isl_tokens, overlap, shared_cache_hits, - router_hint_candidates, + kv_transfer_candidates, } = self .prepare_selection_inputs( &entry, @@ -823,7 +825,7 @@ impl SelectionCore { self.kv_router_config .assume_kv_reuse(router_config_override.as_ref()), true, - retain_router_hint_chain, + retain_kv_transfer_chain, ) .await?; let mode = if book { @@ -878,8 +880,8 @@ impl SelectionCore { block_hashes: Some(block_hashes), isl_tokens, overlap, - router_hint_candidates, - retain_router_hint_chain, + kv_transfer_candidates, + retain_kv_transfer_chain, router_config_override, lora_name: prompt.lora_name, priority_jump, @@ -943,13 +945,19 @@ impl SelectionCore { .map(|(threshold, total_kv_blocks)| { potential_decode_blocks as f64 > threshold * total_kv_blocks as f64 }); - let router_hint = if retain_router_hint_chain { - router_hint_for_selection( + let kv_hint = if retain_kv_transfer_chain { + transfer_hint_for_selection( &self.catalog.scheduler_configs_for_key(&key), response.best_worker, response.target_cached_prefix_blocks, - response.router_hint_candidates.as_ref(), + response.kv_transfer_candidates.as_ref(), ) + .map(|payload| { + KvHint::new( + selection_id.as_deref().unwrap_or_default(), + vec![KvHintAction::fetch("a1", payload)], + ) + }) } else { None }; @@ -1005,7 +1013,7 @@ impl SelectionCore { potential_decode_blocks, decode_busy, worker_load, - router_hint, + kv_hint, }) } @@ -1443,7 +1451,7 @@ impl SelectionCore { prompt: &PromptRequest, assume_kv_reuse: bool, query_shared_cache: bool, - retain_router_hint_chain: bool, + retain_kv_transfer_chain: bool, ) -> Result { let normalized = prompt.normalize_for_selection( entry.is_eagle, @@ -1462,7 +1470,7 @@ impl SelectionCore { .find_tiered_matches_with_options( normalized.block_hashes.clone(), LowerTierQueryOptions { - retain_router_hint_chain, + retain_kv_transfer_chain, }, ) .await @@ -1490,8 +1498,8 @@ impl SelectionCore { let tiered = tiered?; let overlap = OverlapAnalysis::new(&self.kv_router_config, entry.block_size, &tiered).signals(); - let router_hint_candidates = retain_router_hint_chain - .then(|| tiered.router_hint_root_candidates().cloned()) + let kv_transfer_candidates = retain_kv_transfer_chain + .then(|| tiered.kv_transfer_candidates().cloned()) .flatten(); drop(tiered); Ok(PreparedSelectionInputs { @@ -1500,7 +1508,7 @@ impl SelectionCore { isl_tokens: normalized.isl_tokens, overlap, shared_cache_hits, - router_hint_candidates, + kv_transfer_candidates, }) } @@ -1581,24 +1589,24 @@ fn spawn_reservation_index_sweep( /// cache owner) holding a longer root-aligned prefix than the target's own /// `target_cached_prefix_blocks`, with a non-empty control endpoint. Mirrors the /// frontend `KvRouter::router_hint_for_selection`. -fn router_hint_for_selection( +fn transfer_hint_for_selection( configs: &HashMap, target: WorkerWithDpRank, target_cached_prefix_blocks: u32, - candidates: Option<&RouterHintRootCandidates>, -) -> Option { + candidates: Option<&KvTransferCandidates>, +) -> Option { let candidates = candidates?; let target_config = configs.get(&target.worker_id)?; - let target_metadata = target_config.router_hint_metadata_for_dp_rank(target.dp_rank)?; + let target_metadata = target_config.kv_hint_transfer_metadata_for_dp_rank(target.dp_rank)?; let prefix_blocks_to_beat = usize::try_from(target_cached_prefix_blocks).unwrap_or(usize::MAX); let (source, block_hashes) = candidates.best_source(prefix_blocks_to_beat, |source| match source { - RouterHintCandidateSource::Worker(worker) => { + KvTransferCandidateSource::Worker(worker) => { worker != target && configs.get(&worker.worker_id).is_some_and(|config| { config - .router_hint_metadata_for_dp_rank(worker.dp_rank) + .kv_hint_transfer_metadata_for_dp_rank(worker.dp_rank) .is_some_and(|source_metadata| { source_metadata.worker_type == target_metadata.worker_type && source_metadata @@ -1607,7 +1615,7 @@ fn router_hint_for_selection( }) }) } - RouterHintCandidateSource::CacheOwner(owner) => candidates + KvTransferCandidateSource::CacheOwner(owner) => candidates .routing_snapshot .as_ref() .and_then(|snapshot| snapshot.router_hint_source(owner)) @@ -1618,12 +1626,12 @@ fn router_hint_for_selection( }), })?; let source_control_endpoint = match source { - RouterHintCandidateSource::Worker(worker) => configs + KvTransferCandidateSource::Worker(worker) => configs .get(&worker.worker_id)? - .router_hint_metadata_for_dp_rank(worker.dp_rank)? + .kv_hint_transfer_metadata_for_dp_rank(worker.dp_rank)? .source_control_endpoint? .to_string(), - RouterHintCandidateSource::CacheOwner(owner) => candidates + KvTransferCandidateSource::CacheOwner(owner) => candidates .routing_snapshot .as_ref()? .router_hint_source(owner)? @@ -1634,7 +1642,7 @@ fn router_hint_for_selection( if block_hashes.is_empty() { return None; } - Some(RouterHint { + Some(KvSourceLocationsPayload { source_control_endpoint, block_hashes, }) @@ -1878,12 +1886,13 @@ mod tests { #[tokio::test] async fn bookings_populate_the_approximate_primary_without_kv_events() { // use_kv_events=false: the primary is approximate and bookings feed it. - let core = SelectionCore::new_local( + let core = SelectionCore::try_new_local( test_config(false), 1, CancellationToken::new(), SelectionCacheConfig::default(), - ); + ) + .expect("valid test config"); core.upsert_worker(worker(1)).await.expect("worker upsert"); core.upsert_worker(worker(2)).await.expect("worker upsert"); @@ -2040,12 +2049,13 @@ mod tests { use crate::indexer::KvIndexerInterface; use crate::protocols::{BlockHashOptions, StorageTier, compute_block_hash_for_seq}; - let core = SelectionCore::new_local( + let core = SelectionCore::try_new_local( test_config(true), 1, CancellationToken::new(), SelectionCacheConfig::default(), - ); + ) + .expect("valid test config"); for worker_id in [1, 2] { let mut request = worker_with_kv_events(worker_id); request.router_hint_worker_type = Some("decode".to_string()); @@ -2055,7 +2065,7 @@ mod tests { } let key = RoutingPartitionId::new("model", "default"); let entry = core.entry(&key).expect("entry"); - assert!(entry.indexer.supports_router_hint_chain_retention()); + assert!(entry.indexer.supports_kv_transfer_chain_retention()); // Worker 1 holds both blocks of the prompt; worker 2 holds nothing. let tokens: Vec = (1..=8).collect(); @@ -2082,23 +2092,28 @@ mod tests { request.prompt = prompt(); request.pinned_worker = Some(WorkerWithDpRank::new(2, 0)); let response = core.select_and_reserve(request).await.expect("reserve"); - let hint = response.router_hint.expect("router hint for worker 2"); - assert_eq!(hint.source_control_endpoint, "tcp://worker-1:9000"); - assert_eq!(hint.block_hashes.len(), 2); + let hint = response.kv_hint.expect("router hint for worker 2"); + assert_eq!(hint.message_id, "to-worker-2"); + assert_eq!(hint.actions[0].action_type, "kv.fetch"); + let payload: KvSourceLocationsPayload = + serde_json::from_value(serde_json::to_value(&hint.actions[0].payload).unwrap()) + .unwrap(); + assert_eq!(payload.source_control_endpoint, "tcp://worker-1:9000"); + assert_eq!(payload.block_hashes.len(), 2); // Booking on worker 1 itself: nothing holds a longer prefix. let mut request = reserve_request("to-worker-1"); request.prompt = prompt(); request.pinned_worker = Some(WorkerWithDpRank::new(1, 0)); let response = core.select_and_reserve(request).await.expect("reserve"); - assert!(response.router_hint.is_none()); + assert!(response.kv_hint.is_none()); // Query-only selections never carry a hint. let mut request = select_request(); request.prompt = prompt(); request.pinned_worker = Some(WorkerWithDpRank::new(2, 0)); let response = core.select(request).await.expect("select"); - assert!(response.router_hint.is_none()); + assert!(response.kv_hint.is_none()); } #[tokio::test] @@ -2106,12 +2121,13 @@ mod tests { use crate::indexer::KvIndexerInterface; use crate::protocols::{BlockHashOptions, StorageTier, compute_block_hash_for_seq}; - let core = SelectionCore::new_local( + let core = SelectionCore::try_new_local( test_config(true), 1, CancellationToken::new(), SelectionCacheConfig::default(), - ); + ) + .expect("valid test config"); for worker_id in [1, 2] { core.upsert_worker(worker_with_kv_events(worker_id)) .await @@ -2136,17 +2152,18 @@ mod tests { let mut request = reserve_request("plain"); request.pinned_worker = Some(WorkerWithDpRank::new(2, 0)); let response = core.select_and_reserve(request).await.expect("reserve"); - assert!(response.router_hint.is_none()); + assert!(response.kv_hint.is_none()); } #[tokio::test] async fn event_driven_indexer_does_not_record_bookings() { - let core = SelectionCore::new_local( + let core = SelectionCore::try_new_local( test_config(true), 1, CancellationToken::new(), SelectionCacheConfig::default(), - ); + ) + .expect("valid test config"); core.upsert_worker(worker_with_kv_events(1)) .await .expect("worker upsert"); diff --git a/lib/kv-router/src/services/selection/types.rs b/lib/kv-router/src/services/selection/types.rs index 78f633a1a1d8..05b2f0896853 100644 --- a/lib/kv-router/src/services/selection/types.rs +++ b/lib/kv-router/src/services/selection/types.rs @@ -6,11 +6,11 @@ use std::collections::{HashMap, HashSet}; use serde::{Deserialize, Serialize}; use crate::identity::{RoutingPartitionId, default_routing_group}; +use crate::kv_hints::KvHint; use crate::protocols::{ - DpRank, KvTransferEnforcement, RouterHintWorkerMetadata, RoutingConstraints, + DpRank, KvHintTransferWorkerMetadata, KvTransferEnforcement, RoutingConstraints, WorkerAffinityTarget, WorkerConfigLike, WorkerId, WorkerWithDpRank, }; -use crate::router_hint::RouterHint; use crate::scheduling::config::RouterConfigOverride; pub use crate::scheduling::{OverlapScoresResponse, SharedCacheOverlapScore, WorkerOverlapScore}; use crate::scheduling::{PotentialLoad, SessionContext, WorkerSelectionInputTrigger}; @@ -100,15 +100,15 @@ impl WorkerConfigLike for SelectionWorkerConfig { self.kv_transfer_preferred_weight } - fn router_hint_metadata_for_dp_rank( + fn kv_hint_transfer_metadata_for_dp_rank( &self, dp_rank: DpRank, - ) -> Option> { + ) -> Option> { let worker_type = self.router_hint_worker_type.as_deref()?; if worker_type.is_empty() { return None; } - Some(RouterHintWorkerMetadata { + Some(KvHintTransferWorkerMetadata { worker_type, source_control_endpoint: self .router_hint_source_control_endpoints @@ -622,7 +622,7 @@ pub struct SelectResponse { /// only for bookings when the partition has router-hint-capable workers, /// the indexer can retain the matched chain, and a better source exists. #[serde(skip_serializing_if = "Option::is_none")] - pub router_hint: Option, + pub kv_hint: Option, } /// Load snapshot of the chosen worker, as the scheduler projected it for this From 8dfeecbcc5305219333dd881af32e80c378fdc69 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Wed, 9 Sep 2026 12:05:27 -0700 Subject: [PATCH 005/102] fix(bindings): compile selection service tests with remote indexing Signed-off-by: Ishan Dhanani --- lib/bindings/python/rust/llm/kv.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/lib/bindings/python/rust/llm/kv.rs b/lib/bindings/python/rust/llm/kv.rs index 9f48c153469e..913d32991c70 100644 --- a/lib/bindings/python/rust/llm/kv.rs +++ b/lib/bindings/python/rust/llm/kv.rs @@ -970,7 +970,8 @@ mod selection_service_lifecycle_tests { #[test] fn idempotent_shutdown() { let service = - Python::with_gil(|py| SelectionService::new(py, 1, None, None, None, None)).unwrap(); + Python::with_gil(|py| SelectionService::new(py, 1, None, None, None, None, None)) + .unwrap(); Python::with_gil(|py| { service.shutdown(py); service.shutdown(py); From 2d774263042d70df455816e49b306ee7ace43e12 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Thu, 3 Sep 2026 07:47:47 +0000 Subject: [PATCH 006/102] feat(router): re-host KvRouter scheduling on an embedded selection partition Phase 1a of the runtime-optional frontend roadmap. Behind KvRouterConfig::router_embedded_selection (--router-embedded-selection / DYN_ROUTER_EMBEDDED_SELECTION, default off), the frontend KvRouter runs its scheduling and active-sequence accounting on one partition of an in-process SelectionService instead of the runtime-bound KvScheduler: - RouterScheduler enum in KvRouter dispatches statically to either backend; every router method (admitted and advisory scheduling, bookings, lifecycle, potential loads) goes through it. Request leases fence with the partition scheduler's booking_cleanup exactly as with KvScheduler. - kv_router/embedded.rs builds the service with the registry-resolved policy (or Dynamo's default scorer/picker under the router's metric label), the router's overload/availability providers and prefill estimator as host hooks, and the router's own indexer as the dequeue-time overlap refresh source. A catalog feeder mirrors the runtime-config watch into the partition's worker catalog (dp layout, capacities, taints, topology, KV-transfer, router-hint metadata). - The router keeps its indexer and request transport: it computes overlap against its indexer and hands the partition a fully formed ScheduleRequest. Selection-service additions used by the host: SelectionPartition handle, ensure_partition (create the scheduler before workers register), external_kv_events (host feeds KV state; no ZMQ listeners), a policy-factory override on the builder, and OverlapRefreshSource (partition indexer, external, or disabled). Not carried over in embedded mode: NATS active-sequence replica sync (router_replica_sync warns), per-class queue metrics, and scheduler load publication to the routing-load monitor. Tests: routing_host::embedded_selection_books_and_releases_like_the_runtime_scheduler mirrors the runtime preview/plan/abort/track/free flow on the embedded path; embedded::worker_request_mirrors_runtime_config covers the catalog mapping. All existing kv_router tests pass on the runtime path. Signed-off-by: Ishan Dhanani --- .../configuration/groups/kv_router_args.py | 14 + .../components/frontend-configuration.mdx | 9 + lib/bindings/python/rust/llm/entrypoint.rs | 4 +- lib/bindings/python/src/dynamo/_core.pyi | 3 + lib/kv-router/src/scheduling/config.rs | 14 + .../src/services/common/replica_sync.rs | 2 +- lib/kv-router/src/services/indexer/backend.rs | 9 + .../src/services/selection/core/mod.rs | 175 ++++++- lib/kv-router/src/services/selection/mod.rs | 7 +- .../src/services/selection/service.rs | 52 ++- lib/llm/Cargo.toml | 2 +- lib/llm/src/kv_router.rs | 231 +++++---- lib/llm/src/kv_router/embedded.rs | 442 ++++++++++++++++++ lib/llm/src/kv_router/routing_host/tests.rs | 114 +++++ 14 files changed, 977 insertions(+), 101 deletions(-) create mode 100644 lib/llm/src/kv_router/embedded.rs diff --git a/components/src/dynamo/common/configuration/groups/kv_router_args.py b/components/src/dynamo/common/configuration/groups/kv_router_args.py index 1a43c56a900a..87091c36b2cc 100644 --- a/components/src/dynamo/common/configuration/groups/kv_router_args.py +++ b/components/src/dynamo/common/configuration/groups/kv_router_args.py @@ -39,6 +39,7 @@ "router_temperature", "use_kv_events", "router_replica_sync", + "router_embedded_selection", "router_track_active_blocks", "router_track_output_blocks", "router_assume_kv_reuse", @@ -200,6 +201,7 @@ class KvRouterConfigBase(ConfigBase): router_temperature: float use_kv_events: bool router_replica_sync: bool + router_embedded_selection: bool = False router_track_active_blocks: bool router_track_output_blocks: bool router_assume_kv_reuse: bool @@ -423,6 +425,18 @@ def add_arguments(self, parser) -> None: "the Runtime event plane." ), ) + add_negatable_bool_argument( + g, + flag_name="--router-embedded-selection", + env_var="DYN_ROUTER_EMBEDDED_SELECTION", + default=False, + dest="router_embedded_selection", + help=( + "KV Router: Run scheduling and active-sequence accounting on an " + "embedded selection-service partition instead of the runtime-bound " + "scheduler. Experimental; the KV indexer and request transport are unchanged." + ), + ) add_negatable_bool_argument( g, flag_name="--router-track-active-blocks", diff --git a/docs/fern/pages/reference/components/frontend-configuration.mdx b/docs/fern/pages/reference/components/frontend-configuration.mdx index 729a74ecbde3..e2d1615de025 100644 --- a/docs/fern/pages/reference/components/frontend-configuration.mdx +++ b/docs/fern/pages/reference/components/frontend-configuration.mdx @@ -213,6 +213,15 @@ router. See [Router Guide](../../developer-guide/knowledge-base/modular-componen Environment variable: `DYN_ROUTER_REPLICA_SYNC` + + **Experimental.** Run the router's scheduling and active-sequence accounting on an embedded + selection-service partition (the same core as the standalone selection service) instead of the + runtime-bound scheduler. The KV indexer, request transport, and worker discovery are unchanged; + `--router-replica-sync` is not applied to the embedded partition. + + Environment variable: `DYN_ROUTER_EMBEDDED_SELECTION` + + Hash function for router-derived active-sequence identities. `keyed-xxh3-v1` is experimental and requires both a tracking key file and key ID. diff --git a/lib/bindings/python/rust/llm/entrypoint.rs b/lib/bindings/python/rust/llm/entrypoint.rs index 3634bdb8f9fe..c378d1aa1bba 100644 --- a/lib/bindings/python/rust/llm/entrypoint.rs +++ b/lib/bindings/python/rust/llm/entrypoint.rs @@ -242,7 +242,7 @@ impl AicPerfConfig { #[pymethods] impl KvRouterConfig { #[new] - #[pyo3(signature = (overlap_score_weight=None, host_cache_hit_weight=0.75, disk_cache_hit_weight=0.25, router_temperature=0.0, use_kv_events=true, *, router_replica_sync=false, router_track_active_blocks=true, router_track_output_blocks=false, router_assume_kv_reuse=true, router_track_prefill_tokens=true, router_prefill_load_model="none", router_ttl_secs=120.0, router_approximate_cache_policy="ttl", router_queue_threshold=None, router_event_threads=4, router_queue_policy="fcfs", use_remote_indexer=false, serve_indexer=false, shared_cache_multiplier=0.0, shared_cache_type="none", router_predicted_ttl_secs=None, conditional_disagg_enabled=false, conditional_disagg_policy="isl_bounding", conditional_disagg_eff_isl_threshold=2048, conditional_disagg_eff_isl_ratio_threshold=0.7, conditional_disagg_prefill_busy_threshold=None, conditional_disagg_decode_busy_threshold=None, overlap_score_credit=1.0, overlap_score_credit_decay=0.0, prefill_load_scale=1.0, decode_active_request_weight=0.0, router_policy_config=None, router_prefill_policy=None, router_decode_policy=None, router_tracking_hash="public-xxh3-v1", router_tracking_key_file=None, router_tracking_key_id=None))] + #[pyo3(signature = (overlap_score_weight=None, host_cache_hit_weight=0.75, disk_cache_hit_weight=0.25, router_temperature=0.0, use_kv_events=true, *, router_replica_sync=false, router_embedded_selection=false, router_track_active_blocks=true, router_track_output_blocks=false, router_assume_kv_reuse=true, router_track_prefill_tokens=true, router_prefill_load_model="none", router_ttl_secs=120.0, router_approximate_cache_policy="ttl", router_queue_threshold=None, router_event_threads=4, router_queue_policy="fcfs", use_remote_indexer=false, serve_indexer=false, shared_cache_multiplier=0.0, shared_cache_type="none", router_predicted_ttl_secs=None, conditional_disagg_enabled=false, conditional_disagg_policy="isl_bounding", conditional_disagg_eff_isl_threshold=2048, conditional_disagg_eff_isl_ratio_threshold=0.7, conditional_disagg_prefill_busy_threshold=None, conditional_disagg_decode_busy_threshold=None, overlap_score_credit=1.0, overlap_score_credit_decay=0.0, prefill_load_scale=1.0, decode_active_request_weight=0.0, router_policy_config=None, router_prefill_policy=None, router_decode_policy=None, router_tracking_hash="public-xxh3-v1", router_tracking_key_file=None, router_tracking_key_id=None))] #[allow(clippy::too_many_arguments)] fn new( overlap_score_weight: Option, @@ -251,6 +251,7 @@ impl KvRouterConfig { router_temperature: f64, use_kv_events: bool, router_replica_sync: bool, + router_embedded_selection: bool, router_track_active_blocks: bool, router_track_output_blocks: bool, router_assume_kv_reuse: bool, @@ -301,6 +302,7 @@ impl KvRouterConfig { router_temperature, use_kv_events, router_replica_sync, + router_embedded_selection, router_track_active_blocks, router_track_output_blocks, router_assume_kv_reuse, diff --git a/lib/bindings/python/src/dynamo/_core.pyi b/lib/bindings/python/src/dynamo/_core.pyi index 8e11a3a45dfa..1e192f8f72c7 100644 --- a/lib/bindings/python/src/dynamo/_core.pyi +++ b/lib/bindings/python/src/dynamo/_core.pyi @@ -1784,6 +1784,7 @@ class KvRouterConfig: use_kv_events: bool = True, *, router_replica_sync: bool = False, + router_embedded_selection: bool = False, router_track_active_blocks: bool = True, router_track_output_blocks: bool = False, router_assume_kv_reuse: bool = True, @@ -1829,6 +1830,8 @@ class KvRouterConfig: router_temperature: Temperature for normalized worker sampling via softmax (default: 0.0) use_kv_events: Whether to use KV events from workers (default: True) router_replica_sync: Enable replica synchronization (default: False) + router_embedded_selection: Run scheduling on an embedded selection-service + partition instead of the runtime-bound scheduler (default: False). Experimental. router_track_active_blocks: Track active blocks for load balancing (default: True) router_track_output_blocks: Track output blocks during generation (default: False). When enabled, the router adds placeholder blocks as tokens are generated diff --git a/lib/kv-router/src/scheduling/config.rs b/lib/kv-router/src/scheduling/config.rs index 48ca83d332d8..ba7d913dc527 100644 --- a/lib/kv-router/src/scheduling/config.rs +++ b/lib/kv-router/src/scheduling/config.rs @@ -276,6 +276,9 @@ fn kv_router_config_from_lookup( if let Some(value) = parse_bool(&get_env, "DYN_ROUTER_TRACK_PREFILL_TOKENS") { config.router_track_prefill_tokens = value; } + if let Some(value) = parse_bool(&get_env, "DYN_ROUTER_EMBEDDED_SELECTION") { + config.router_embedded_selection = value; + } if let Some(value) = get_env("DYN_ROUTER_TRACKING_HASH") { config.router_tracking_hash = value.parse()?; } @@ -666,6 +669,8 @@ struct KvRouterConfigSerde { router_track_output_blocks: bool, router_assume_kv_reuse: bool, router_track_prefill_tokens: bool, + #[serde(default)] + router_embedded_selection: bool, router_tracking_hash: TrackingHashAlgorithm, router_tracking_key_file: Option, router_tracking_key_id: Option, @@ -712,6 +717,7 @@ impl Default for KvRouterConfigSerde { router_track_output_blocks: config.router_track_output_blocks, router_assume_kv_reuse: config.router_assume_kv_reuse, router_track_prefill_tokens: config.router_track_prefill_tokens, + router_embedded_selection: config.router_embedded_selection, router_tracking_hash: config.router_tracking_hash, router_tracking_key_file: config.router_tracking_key_file, router_tracking_key_id: config.router_tracking_key_id, @@ -797,6 +803,12 @@ pub struct KvRouterConfig { #[serde(default = "default_track_prefill_tokens")] pub router_track_prefill_tokens: bool, + /// Run the frontend router's scheduling and active-sequence accounting on + /// an embedded selection-service partition instead of the runtime-bound + /// scheduler (default: false). Frontend-local; workers ignore it. + #[serde(default, skip_serializing_if = "is_default")] + pub router_embedded_selection: bool, + /// Hash algorithm used for router-derived active-sequence identities. #[serde(default, skip_serializing_if = "is_default")] pub router_tracking_hash: TrackingHashAlgorithm, @@ -959,6 +971,7 @@ impl Default for KvRouterConfig { router_track_output_blocks: false, router_assume_kv_reuse: true, router_track_prefill_tokens: default_track_prefill_tokens(), + router_embedded_selection: false, router_tracking_hash: TrackingHashAlgorithm::default(), router_tracking_key_file: None, router_tracking_key_id: None, @@ -1023,6 +1036,7 @@ impl TryFrom for KvRouterConfig { router_track_output_blocks: compat.router_track_output_blocks, router_assume_kv_reuse: compat.router_assume_kv_reuse, router_track_prefill_tokens: compat.router_track_prefill_tokens, + router_embedded_selection: compat.router_embedded_selection, router_tracking_hash: compat.router_tracking_hash, router_tracking_key_file: compat.router_tracking_key_file, router_tracking_key_id: compat.router_tracking_key_id, diff --git a/lib/kv-router/src/services/common/replica_sync.rs b/lib/kv-router/src/services/common/replica_sync.rs index efe97053ed4a..4d33bc879e2a 100644 --- a/lib/kv-router/src/services/common/replica_sync.rs +++ b/lib/kv-router/src/services/common/replica_sync.rs @@ -118,7 +118,7 @@ pub(crate) struct ScopedReplicaSync { } #[derive(Clone)] -pub(crate) struct ScopedSequencePublisher { +pub struct ScopedSequencePublisher { replica: Option, } diff --git a/lib/kv-router/src/services/indexer/backend.rs b/lib/kv-router/src/services/indexer/backend.rs index 51f46075975c..73e27a80d3c7 100644 --- a/lib/kv-router/src/services/indexer/backend.rs +++ b/lib/kv-router/src/services/indexer/backend.rs @@ -768,6 +768,15 @@ impl TieredMatchProvider for Indexer { ) -> std::result::Result { self.find_tiered_matches(sequence.to_vec()).await } + + async fn find_tiered_matches_with_options( + &self, + sequence: &[LocalBlockHash], + options: LowerTierQueryOptions, + ) -> std::result::Result { + self.find_tiered_matches_with_options(sequence.to_vec(), options) + .await + } } #[cfg(test)] diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index 1e31d60fa287..5ef463644ee7 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -13,7 +13,8 @@ use tokio_util::sync::CancellationToken; use crate::identity::RoutingPartitionId; use crate::indexer::{ - LowerTierQueryOptions, RoutingDecisionHashes, SharedKvCache, TieredMatchDetails, + KvRouterError, LowerTierQueryOptions, RoutingDecisionHashes, SharedKvCache, TieredMatchDetails, + TieredMatchProvider, }; use crate::kv_hints::{ KvHint, KvHintAction, KvSourceLocationsPayload, KvTransferCandidateSource, KvTransferCandidates, @@ -55,13 +56,78 @@ use super::types::{ use crate::WorkerSelectionPolicyFactory; use crate::WorkerType; -type SelectionScheduler = LocalScheduler< +/// Source of dequeue-time overlap refreshes for a partition scheduler: the +/// partition's own indexer, or one the embedding host supplies when it keeps +/// its own index (the frontend `KvRouter` in embedded mode). +#[derive(Clone)] +pub enum RefreshProvider { + Local(Arc), + External(Arc), +} + +#[async_trait::async_trait] +impl TieredMatchProvider for RefreshProvider { + async fn find_tiered_matches( + &self, + sequence: &[LocalBlockHash], + ) -> Result { + match self { + Self::Local(indexer) => indexer.find_tiered_matches(sequence.to_vec()).await, + Self::External(provider) => provider.find_tiered_matches(sequence).await, + } + } + + async fn find_tiered_matches_with_options( + &self, + sequence: &[LocalBlockHash], + options: LowerTierQueryOptions, + ) -> Result { + match self { + Self::Local(indexer) => { + indexer + .find_tiered_matches_with_options(sequence.to_vec(), options) + .await + } + Self::External(provider) => { + provider + .find_tiered_matches_with_options(sequence, options) + .await + } + } + } +} + +/// The scheduler type every partition runs. +pub type SelectionScheduler = LocalScheduler< ScopedSequencePublisher, SelectionWorkerConfig, WorkerSelectionPolicy, - TieredOverlapRefresher, + TieredOverlapRefresher, >; +/// Handle to one partition's scheduler and indexer for an embedding host that +/// drives scheduling directly (bypassing the request-shaped `select` API). +#[derive(Clone)] +pub struct SelectionPartition(Arc); + +impl SelectionPartition { + pub fn key(&self) -> &RoutingPartitionId { + &self.0.key + } + + pub fn block_size(&self) -> u32 { + self.0.block_size + } + + pub fn scheduler(&self) -> &SelectionScheduler { + &self.0.scheduler + } + + pub fn indexer(&self) -> &Indexer { + &self.0.indexer + } +} + struct SelectionEntry { key: RoutingPartitionId, block_size: u32, @@ -132,6 +198,33 @@ pub struct SelectionHostHooks { /// Narrows candidates to the workers that can serve the request's LoRA /// adapter, strictly within the caller's `allowed_worker_ids`. pub lora_worker_filter: Option>, + /// Where dequeue-time overlap refresh reads from. An embedding host that + /// keeps its own index and drives the scheduler directly points this at + /// that index so queued requests are re-scored against the index that + /// scored them, or disables refresh when its index cannot serve it. + pub overlap_refresh: OverlapRefreshSource, +} + +/// Source of dequeue-time overlap refreshes for partition schedulers. +#[derive(Clone, Default)] +pub enum OverlapRefreshSource { + /// The partition's own indexer (standalone service default). + #[default] + PartitionIndexer, + /// A host-owned index. + External(Arc), + /// No refresh: queued requests keep their admission-time scores. + Disabled, +} + +impl std::fmt::Debug for OverlapRefreshSource { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str(match self { + Self::PartitionIndexer => "PartitionIndexer", + Self::External(_) => "External", + Self::Disabled => "Disabled", + }) + } } impl std::fmt::Debug for SelectionHostHooks { @@ -144,6 +237,7 @@ impl std::fmt::Debug for SelectionHostHooks { ) .field("shared_cache", &self.shared_cache.is_some()) .field("lora_worker_filter", &self.lora_worker_filter.is_some()) + .field("overlap_refresh", &self.overlap_refresh) .field( "overloaded_worker_provider", &self.overloaded_worker_provider.is_some(), @@ -191,8 +285,8 @@ pub struct SelectionCore { /// runs inside the host runtime; construction itself may not. reservation_sweep_started: OnceCell<()>, /// Whether this core subscribes to worker KV events itself. False when - /// events are disabled or when the primary indexer is a remote service - /// that workers publish to directly. + /// events are disabled, when the primary indexer is a remote service that + /// workers publish to directly, or when the embedding host feeds events. listens_for_kv_events: bool, indexer_registry: Arc, kv_router_config: crate::config::KvRouterConfig, @@ -247,9 +341,16 @@ impl SelectionCore { cache_config, tracking_hash, indexer_policy, + false, )) } + /// Scheduler and indexer handle for `key`, once a worker has been upserted + /// into that partition. + pub fn partition(&self, key: &RoutingPartitionId) -> Option { + self.entry(key).map(SelectionPartition) + } + #[allow(clippy::too_many_arguments)] pub(super) fn new_inner( kv_router_config: crate::config::KvRouterConfig, @@ -263,13 +364,15 @@ impl SelectionCore { cache_config: SelectionCacheConfig, tracking_hash: Arc, indexer_policy: IndexerPolicy, + external_kv_events: bool, ) -> Self { let cancel_token = cancel_token.child_token(); let indexer_registry = Arc::new(WorkerRegistry::new_with_cancel_token( indexer_threads, cancel_token.clone(), )); - let listens_for_kv_events = kv_router_config.use_kv_events && !indexer_policy.is_remote(); + let listens_for_kv_events = + kv_router_config.use_kv_events && !indexer_policy.is_remote() && !external_kv_events; indexer_registry.set_indexer_policy(indexer_policy); if signal_indexer_ready { indexer_registry.signal_ready(); @@ -507,8 +610,35 @@ impl SelectionCore { let block_size = record .block_size .ok_or_else(|| SelectionError::BadRequest("block_size is required".to_string()))?; - let is_eagle = record.is_eagle.unwrap_or(false); - let key = record.key(); + self.ensure_entry_for(record.key(), block_size, record.is_eagle.unwrap_or(false)) + } + + /// Create the partition scheduler and indexer for `key` before any worker + /// registers, so an embedding host can hold the scheduler handle from + /// construction. Idempotent; a later worker with a different block size or + /// eagle setting is rejected at reconciliation. + pub fn ensure_partition( + &self, + key: RoutingPartitionId, + block_size: u32, + is_eagle: bool, + ) -> Result { + self.ensure_running()?; + if block_size == 0 { + return Err(SelectionError::BadRequest( + "block_size must be greater than 0".to_string(), + )); + } + self.ensure_entry_for(key, block_size, is_eagle) + .map(SelectionPartition) + } + + fn ensure_entry_for( + &self, + key: RoutingPartitionId, + block_size: u32, + is_eagle: bool, + ) -> Result, SelectionError> { self.reservation_sweep_started.get_or_init(|| { spawn_reservation_index_sweep( Arc::clone(&self.entries), @@ -551,11 +681,22 @@ impl SelectionCore { let indexer = self .indexer_registry .get_or_create_indexer(key.clone(), block_size); - let overlap_refresh = Arc::new(TieredOverlapRefresher::new( - indexer.clone(), - self.kv_router_config.clone(), - block_size, - )); + let refresh_provider = match &self.host_hooks.overlap_refresh { + OverlapRefreshSource::PartitionIndexer => { + Some(RefreshProvider::Local(Arc::new(indexer.clone()))) + } + OverlapRefreshSource::External(provider) => { + Some(RefreshProvider::External(Arc::clone(provider))) + } + OverlapRefreshSource::Disabled => None, + }; + let overlap_refresh = refresh_provider.map(|provider| { + Arc::new(TieredOverlapRefresher::new( + provider, + self.kv_router_config.clone(), + block_size, + )) + }); let selector = self.worker_selection_policy_factory.as_ref().map_or_else( || WorkerSelectionPolicy::default(self.kv_router_config.clone(), worker_label), |factory| factory(&self.kv_router_config, self.worker_type, key.as_ref()), @@ -571,7 +712,7 @@ impl SelectionCore { block_size, selector, self.host_hooks.prefill_load_estimator.clone(), - Some(overlap_refresh), + overlap_refresh, // Standalone selection has no router Client snapshot, so // these stay `None` unless an embedding host injects them. self.host_hooks.overloaded_worker_provider.clone(), @@ -1824,6 +1965,7 @@ mod tests { SelectionCacheConfig::default(), tracking_hash, IndexerPolicy::from_router_config(&test_config(false)).expect("indexer policy"), + false, ); core.upsert_worker(worker(1)).await.expect("worker upsert"); @@ -1863,6 +2005,7 @@ mod tests { SelectionCacheConfig::default(), tracking_hash, indexer_policy, + false, ) } @@ -2023,6 +2166,7 @@ mod tests { SelectionCacheConfig::default(), tracking_hash, indexer_policy, + false, ); assert!(!core.listens_for_kv_events); for worker_id in [1, 2] { @@ -2321,6 +2465,7 @@ mod tests { calls: Arc::clone(&calls), })), lora_worker_filter: None, + overlap_refresh: OverlapRefreshSource::default(), }, Some(factory), ); @@ -2421,6 +2566,7 @@ mod tests { available_worker_provider: Some(Arc::new(move || provider_state.lock().clone())), shared_cache: None, lora_worker_filter: None, + overlap_refresh: OverlapRefreshSource::default(), }); core.upsert_worker(worker(1)).await.expect("worker upsert"); core.upsert_worker(worker(2)).await.expect("worker upsert"); @@ -2442,6 +2588,7 @@ mod tests { available_worker_provider: None, shared_cache: None, lora_worker_filter: None, + overlap_refresh: OverlapRefreshSource::default(), }); core.upsert_worker(worker(1)).await.expect("worker upsert"); core.upsert_worker(worker(2)).await.expect("worker upsert"); diff --git a/lib/kv-router/src/services/selection/mod.rs b/lib/kv-router/src/services/selection/mod.rs index ae344ccf2e62..cd60810fcbb5 100644 --- a/lib/kv-router/src/services/selection/mod.rs +++ b/lib/kv-router/src/services/selection/mod.rs @@ -20,8 +20,11 @@ mod types; mod tests; pub use crate::WorkerSelectionPolicyFactory; -pub use crate::services::common::replica_sync::ReplicaPeerError; -pub use core::{SelectionCore, SelectionHostHooks, SelectionServiceConfig}; +pub use crate::services::common::replica_sync::{ReplicaPeerError, ScopedSequencePublisher}; +pub use core::{ + OverlapRefreshSource, RefreshProvider, SelectionCore, SelectionHostHooks, SelectionPartition, + SelectionScheduler, SelectionServiceConfig, +}; pub use error::SelectionError; pub use input::PromptRequest; pub use pending::SelectionCacheConfig; diff --git a/lib/kv-router/src/services/selection/service.rs b/lib/kv-router/src/services/selection/service.rs index 574c0c4d4ec0..17a96e7c77e9 100644 --- a/lib/kv-router/src/services/selection/service.rs +++ b/lib/kv-router/src/services/selection/service.rs @@ -5,7 +5,9 @@ use std::sync::Arc; use tokio_util::sync::CancellationToken; +use crate::WorkerSelectionPolicyFactory; use crate::config::KvRouterConfig; +use crate::identity::RoutingPartitionId; use crate::protocols::WorkerId; use crate::scheduling::PotentialLoad; use crate::services::common::replica_sync::{ @@ -14,7 +16,7 @@ use crate::services::common::replica_sync::{ use crate::services::indexer::backend::IndexerPolicy; use crate::tracking_hash::TrackingHashContext; -use super::core::{SelectionCore, SelectionHostHooks, SelectionServiceConfig}; +use super::core::{SelectionCore, SelectionHostHooks, SelectionPartition, SelectionServiceConfig}; use super::error::SelectionError; use super::pending::SelectionCacheConfig; use super::policy_registry::WorkerSelectionPolicyRegistry; @@ -36,6 +38,8 @@ pub struct SelectionServiceBuilder { worker_selection_policy_registry: WorkerSelectionPolicyRegistry, host_hooks: SelectionHostHooks, remote_indexer_url: Option, + worker_selection_policy_factory: Option, + external_kv_events: bool, } /// Warn when a host does not construct workers for explicitly configured policy roles. @@ -75,9 +79,28 @@ impl SelectionServiceBuilder { worker_selection_policy_registry, host_hooks: SelectionHostHooks::default(), remote_indexer_url: None, + worker_selection_policy_factory: None, + external_kv_events: false, } } + /// Use `factory` for every partition's worker-selection policy instead of + /// resolving one from router configuration through the registry. + pub fn worker_selection_policy_factory( + mut self, + factory: WorkerSelectionPolicyFactory, + ) -> Self { + self.worker_selection_policy_factory = Some(factory); + self + } + + /// The embedding host feeds KV state itself: workers become schedulable + /// without `kv_events_endpoint(s)` and no ZMQ listener is started here. + pub fn external_kv_events(mut self) -> Self { + self.external_kv_events = true; + self + } + /// Serve the primary KV index from a standalone indexer at `base_url` /// instead of subscribing to worker KV events in this process. Requires /// `use_kv_events=true`; `indexer_peers` recovery is skipped. @@ -119,9 +142,12 @@ impl SelectionServiceBuilder { self.kv_router_config .validate_config() .map_err(anyhow::Error::msg)?; - let worker_selection_policy_factory = self - .worker_selection_policy_registry - .resolve_for_worker_type(&self.kv_router_config, self.worker_type)?; + let worker_selection_policy_factory = match self.worker_selection_policy_factory { + Some(factory) => Some(factory), + None => self + .worker_selection_policy_registry + .resolve_for_worker_type(&self.kv_router_config, self.worker_type)?, + }; let tracking_hash = Arc::new(TrackingHashContext::from_config(&self.kv_router_config)?); let mut indexer_policy = IndexerPolicy::from_router_config(&self.kv_router_config)?; if let Some(base_url) = &self.remote_indexer_url { @@ -162,6 +188,7 @@ impl SelectionServiceBuilder { self.selection_cache, tracking_hash, indexer_policy, + self.external_kv_events, )); if recover_from_peers { @@ -294,6 +321,23 @@ impl SelectionService { self.core.upsert_worker(req).await } + /// Scheduler and indexer handle for `key`, once a worker has been upserted + /// into that partition. Embedding hosts that keep their own index drive + /// the partition scheduler through this handle. + pub fn partition(&self, key: &RoutingPartitionId) -> Option { + self.core.partition(key) + } + + /// Create (or fetch) the partition for `key` ahead of worker registration. + pub fn ensure_partition( + &self, + key: RoutingPartitionId, + block_size: u32, + is_eagle: bool, + ) -> Result { + self.core.ensure_partition(key, block_size, is_eagle) + } + /// The port this service uses for replica synchronization, if enabled. pub fn replica_sync_port(&self) -> Option { self.replica_sync_port diff --git a/lib/llm/Cargo.toml b/lib/llm/Cargo.toml index 623375c2355c..19532648ffd3 100644 --- a/lib/llm/Cargo.toml +++ b/lib/llm/Cargo.toml @@ -72,7 +72,7 @@ required-features = ["media-ffmpeg"] [dependencies] # repo -dynamo-kv-router = { workspace = true, features = ["metrics", "runtime-protocols"] } +dynamo-kv-router = { workspace = true, features = ["metrics", "runtime-protocols", "standalone-selection"] } dynamo-memory = { workspace = true } dynamo-mocker = { workspace = true, features = ["zmq-events"] } dynamo-rl = { workspace = true } diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index ff9c69d5e321..c89db06fb122 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -55,6 +55,7 @@ pub use dynamo_kv_router::protocols; pub use dynamo_kv_router::scheduling; pub use dynamo_kv_router::selector; +pub(crate) mod embedded; pub mod encoder_router; pub mod indexer; pub mod metrics; @@ -534,7 +535,7 @@ where Sel: dynamo_kv_router::selector::WorkerSelector, { indexer: Indexer, - scheduler: KvScheduler>, + scheduler: RouterScheduler, required_worker_inputs: dynamo_kv_router::selector::WorkerInputs, workers_with_configs: RuntimeConfigWatch, block_size: u32, @@ -560,6 +561,39 @@ where teardown_task_guard: Option, } +/// The scheduler behind a `KvRouter`: the runtime-bound `KvScheduler`, or one +/// partition of an embedded selection service +/// (`KvRouterConfig::router_embedded_selection`). Both expose the same +/// surface; every router method dispatches statically through this enum. +enum RouterScheduler +where + Sel: dynamo_kv_router::selector::WorkerSelector, +{ + Runtime(KvScheduler>), + Embedded(embedded::EmbeddedSelection), +} + +macro_rules! dispatch_scheduler { + ($self:expr, $scheduler:ident => $body:expr) => { + match &$self.scheduler { + RouterScheduler::Runtime($scheduler) => $body, + RouterScheduler::Embedded($scheduler) => $body, + } + }; +} + +impl RouterScheduler +where + Sel: dynamo_kv_router::selector::WorkerSelector + Send + 'static, +{ + fn booking_cleanup(&self) -> scheduler::SchedulerBookingCleanup { + match self { + Self::Runtime(scheduler) => scheduler.booking_cleanup(), + Self::Embedded(scheduler) => scheduler.booking_cleanup(), + } + } +} + fn resolve_tracking_model_name( algorithm: TrackingHashAlgorithm, model_name: Option<&str>, @@ -773,27 +807,61 @@ where let available_worker_provider: WorkerAvailabilityProvider = Arc::new(move || client_for_availability.available_instance_ids()); - let scheduler = KvScheduler::start_with_shared_request_leases( - endpoint.clone(), - block_size, - workers_with_configs.clone(), - selector, - &kv_router_config, - prefill_load_estimator.clone(), - overlap_scores_refresh, - Some(overloaded_worker_provider), - Some(available_worker_provider), - model_name.as_deref(), - metric_worker_type, - scheduler_load, - cancellation_token.child_token(), - ) - .await?; + let scheduler = if kv_router_config.router_embedded_selection { + // The router keeps its indexer and transport; scheduling and + // active-sequence accounting run on a selection-service partition. + let overlap_refresh: Option> = + indexer + .supports_overlap_refresh() + .then(|| Arc::new(indexer.clone()) as Arc<_>); + drop(overlap_scores_refresh); + RouterScheduler::Embedded( + embedded::EmbeddedSelection::start( + embedded::EmbeddedSelectionArgs { + kv_router_config: kv_router_config.clone(), + worker_role, + metric_worker_type, + model_name: model_name.clone(), + block_size, + is_eagle, + prefill_load_estimator: prefill_load_estimator.clone(), + overloaded_worker_provider, + available_worker_provider, + overlap_refresh, + }, + workers_with_configs.clone(), + cancellation_token.child_token(), + ) + .await?, + ) + } else { + RouterScheduler::Runtime( + KvScheduler::start_with_shared_request_leases( + endpoint.clone(), + block_size, + workers_with_configs.clone(), + selector, + &kv_router_config, + prefill_load_estimator.clone(), + overlap_scores_refresh, + Some(overloaded_worker_provider), + Some(available_worker_provider), + model_name.as_deref(), + metric_worker_type, + scheduler_load, + cancellation_token.child_token(), + ) + .await?, + ) + }; let request_leases = request_lease::RequestLeaseManager::new( scheduler.booking_cleanup(), cancellation_token.child_token(), ); - if !scheduler.set_replica_request_lease_observer(Arc::new(request_leases.clone())) { + if let RouterScheduler::Runtime(runtime_scheduler) = &scheduler + && !runtime_scheduler + .set_replica_request_lease_observer(Arc::new(request_leases.clone())) + { return Err(anyhow::anyhow!( "request lease observer is already installed for this router" )); @@ -1600,7 +1668,8 @@ where let seq_hash_elapsed = start.elapsed(); let is_admitted_routing = matches!(admission, FindBestMatchAdmission::WithAdmission { .. }); - let supports_overlap_refresh = self.scheduler.supports_overlap_refresh(); + let supports_overlap_refresh = + dispatch_scheduler!(self, scheduler => scheduler.supports_overlap_refresh()); let retain_block_hashes = supports_overlap_refresh || return_routing_hashes; let has_transfer_capable_workers = self.has_transfer_capable_workers(); let should_prepare_transfer_hint = is_admitted_routing && has_transfer_capable_workers; @@ -1691,41 +1760,43 @@ where shared_cache_hits, }; let (response, attempt, selected_worker_load) = match admission { - FindBestMatchAdmission::WithAdmission { .. } => match self - .scheduler + FindBestMatchAdmission::WithAdmission { .. } => { + match dispatch_scheduler!(self, scheduler => scheduler .schedule_request_admitted(schedule_request) .instrument(tracing::info_span!("kv_router.schedule")) - .await - { - Ok(admitted) => (admitted.response, admitted.attempt, None), - Err(KvSchedulerError::QueueRejected(rejection)) => { - return Ok(FindBestMatchInnerOutcome::WithAdmission( - AdmittedFindBestMatchOutcome { - outcome: FindBestMatchOutcome::QueueRejected { rejection }, - attempt: AdmissionAttempt::Untracked, - }, - )); + .await) + { + Ok(admitted) => (admitted.response, admitted.attempt, None), + Err(KvSchedulerError::QueueRejected(rejection)) => { + return Ok(FindBestMatchInnerOutcome::WithAdmission( + AdmittedFindBestMatchOutcome { + outcome: FindBestMatchOutcome::QueueRejected { rejection }, + attempt: AdmissionAttempt::Untracked, + }, + )); + } + Err(error) => return Err(map_scheduler_error(error)), } - Err(error) => return Err(map_scheduler_error(error)), - }, - FindBestMatchAdmission::WithoutAdmission => match self - .scheduler + } + FindBestMatchAdmission::WithoutAdmission => { + match dispatch_scheduler!(self, scheduler => scheduler .select_without_admission(schedule_request) .instrument(tracing::info_span!("kv_router.select_without_admission")) - .await - { - Ok(advisory) => ( - advisory.response, - AdmissionAttempt::Untracked, - Some(advisory.selected_worker_load), - ), - Err(KvSchedulerError::QueueRejected(rejection)) => { - return Ok(FindBestMatchInnerOutcome::WithoutAdmission( - FindBestMatchAdvisoryOutcome::QueueRejected { rejection }, - )); + .await) + { + Ok(advisory) => ( + advisory.response, + AdmissionAttempt::Untracked, + Some(advisory.selected_worker_load), + ), + Err(KvSchedulerError::QueueRejected(rejection)) => { + return Ok(FindBestMatchInnerOutcome::WithoutAdmission( + FindBestMatchAdvisoryOutcome::QueueRejected { rejection }, + )); + } + Err(error) => return Err(map_scheduler_error(error)), } - Err(error) => return Err(map_scheduler_error(error)), - }, + } }; let kv_hint = if is_admitted_routing { self.kv_hint_for_selection( @@ -1855,7 +1926,7 @@ where /// Register externally-provided workers in the slot tracker. pub fn register_workers(&self, worker_ids: &HashSet) { - self.scheduler.register_workers(worker_ids); + dispatch_scheduler!(self, scheduler => scheduler.register_workers(worker_ids)); } #[allow(clippy::too_many_arguments)] @@ -1895,18 +1966,18 @@ where let prefill_load_hint = self.prefill_load_hint_for(isl_tokens, cached_tokens, track_prefill_tokens); - let admission = self - .scheduler - .add_request_admitted(SequenceRequest { - request_id: request_id.clone(), - token_sequence: maybe_seq_hashes, - track_prefill_tokens, - expected_output_tokens, - prefill_load_hint, - worker, - lora_name, - }) - .await; + let sequence_request = SequenceRequest { + request_id: request_id.clone(), + token_sequence: maybe_seq_hashes, + track_prefill_tokens, + expected_output_tokens, + prefill_load_hint, + worker, + lora_name, + }; + let admission = dispatch_scheduler!(self, scheduler => scheduler + .add_request_admitted(sequence_request) + .await); let attempt_id = match admission { Ok(attempt_id) => attempt_id, Err(error) => { @@ -1927,7 +1998,7 @@ where } pub async fn mark_prefill_completed(&self, request_id: &str) -> Result<(), SequenceError> { - self.scheduler.mark_prefill_completed(request_id).await?; + dispatch_scheduler!(self, scheduler => scheduler.mark_prefill_completed(request_id).await)?; self.request_leases.touch_request(request_id); Ok(()) } @@ -1936,7 +2007,7 @@ where if self.request_leases.finish_request(request_id).await { return Ok(()); } - self.scheduler.free(request_id).await + dispatch_scheduler!(self, scheduler => scheduler.free(request_id).await) } /// Release a booking only if it still belongs to `worker`. @@ -1948,7 +2019,7 @@ where request_id: &str, worker: WorkerWithDpRank, ) -> Result<(), SequenceError> { - self.scheduler.free_if_worker(request_id, worker).await + dispatch_scheduler!(self, scheduler => scheduler.free_if_worker(request_id, worker).await) } #[doc(hidden)] @@ -1964,19 +2035,19 @@ where &self, booking: &scheduler::SchedulerBookingDescriptor, ) -> Result<(), KvSchedulerError> { - self.scheduler + dispatch_scheduler!(self, scheduler => scheduler .mark_prefill_completed_if_booking(booking) - .await + .await) } /// Number of requests currently parked in the scheduler queue. pub fn pending_count(&self) -> usize { - self.scheduler.pending_count() + dispatch_scheduler!(self, scheduler => scheduler.pending_count()) } /// Sum of ISL tokens for requests currently parked in the scheduler queue. pub fn pending_isl_tokens(&self) -> usize { - self.scheduler.pending_isl_tokens() + dispatch_scheduler!(self, scheduler => scheduler.pending_isl_tokens()) } fn prefill_load_hint_for( @@ -2019,7 +2090,7 @@ where /// Get the worker type for this router ("prefill" or "decode"). /// Used for Prometheus metric labeling. pub fn worker_type(&self) -> &'static str { - self.scheduler.worker_type() + dispatch_scheduler!(self, scheduler => scheduler.worker_type()) } /// Return the worker's unique global DP rank when it owns exactly one rank. @@ -2034,7 +2105,7 @@ where request_id: &str, decay_fraction: Option, ) -> Result<(), SequenceError> { - self.scheduler.add_output_block(request_id, decay_fraction) + dispatch_scheduler!(self, scheduler => scheduler.add_output_block(request_id, decay_fraction)) } pub(crate) async fn enqueue_output_block_if_booking( @@ -2042,9 +2113,9 @@ where booking: &scheduler::SchedulerBookingDescriptor, decay_fraction: Option, ) -> Result<(), KvSchedulerError> { - self.scheduler + dispatch_scheduler!(self, scheduler => scheduler .enqueue_output_block_if_booking(booking, decay_fraction) - .await + .await) } pub fn block_size(&self) -> u32 { @@ -2156,12 +2227,16 @@ where let tiered_matches = self.indexer.find_matches_by_tier(block_hashes).await?; let cache_hit_estimates = self.cache_hit_estimates_from_tiered_matches(&tiered_matches); - Ok(self.scheduler.get_potential_loads( - maybe_seq_hashes, - isl_tokens, - cache_hit_estimates.cached_tokens.into_iter().collect(), - track_prefill_tokens, - )) + let effective_cached_tokens: HashMap = + cache_hit_estimates.cached_tokens.into_iter().collect(); + Ok( + dispatch_scheduler!(self, scheduler => scheduler.get_potential_loads( + maybe_seq_hashes, + isl_tokens, + effective_cached_tokens, + track_prefill_tokens, + )), + ) } /// Return per-worker KV overlap by storage tier. diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs new file mode 100644 index 000000000000..c658b3ac3e87 --- /dev/null +++ b/lib/llm/src/kv_router/embedded.rs @@ -0,0 +1,442 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Embedded selection backend for the frontend `KvRouter`. +//! +//! When `KvRouterConfig::router_embedded_selection` is set, the router runs its +//! scheduling and active-sequence accounting on one partition of an in-process +//! `SelectionService` (the same core the standalone selection service and the +//! EPP use) instead of the runtime-bound `KvScheduler`. The router keeps its own +//! KV indexer and request transport: it computes overlap against that indexer +//! and hands the partition scheduler a fully formed `ScheduleRequest`, exactly +//! as it does with `KvScheduler`. The partition's worker catalog is fed from the +//! runtime-config watch, so discovery remains the source of truth for +//! membership while the scheduler and its replica sync become runtime-free. + +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; + +use anyhow::{Context, Result}; +use dynamo_kv_router::WorkerType; +use dynamo_kv_router::config::KvRouterConfig; +use dynamo_kv_router::identity::RoutingPartitionId; +use dynamo_kv_router::indexer::TieredMatchProvider; +use dynamo_kv_router::protocols::{WorkerConfigLike, WorkerId, WorkerWithDpRank}; +use dynamo_kv_router::scheduling::queue::{SchedulerBookingCleanup, SchedulerBookingDescriptor}; +use dynamo_kv_router::scheduling::selector::WorkerSelectionPolicy; +use dynamo_kv_router::scheduling::{ + AdmittedSchedulingResponse, AdvisorySchedulingResponse, AttemptId, KvSchedulerError, + OverloadedWorkerProvider, PotentialLoad, ScheduleRequest, WorkerAvailabilityProvider, +}; +use dynamo_kv_router::sequences::{SequenceError, SequenceRequest}; +use dynamo_kv_router::services::selection::{ + OverlapRefreshSource, SelectionHostHooks, SelectionPartition, SelectionService, + SelectionServiceBuilder, WorkerRequest, WorkerSelectionPolicyRegistry, +}; +use dynamo_kv_router::{DEFAULT_ROUTING_GROUP, PrefillLoadEstimator, WorkerSelectionPolicyFactory}; +use dynamo_tokens::SequenceHash; +use tokio_util::sync::CancellationToken; + +use crate::discovery::RuntimeConfigWatch; +use crate::local_model::runtime_config::ModelRuntimeConfig; + +/// Inputs the embedded backend needs from the router at construction. +pub(crate) struct EmbeddedSelectionArgs { + pub kv_router_config: KvRouterConfig, + pub worker_role: Option, + pub metric_worker_type: &'static str, + pub model_name: Option, + pub block_size: u32, + pub is_eagle: bool, + pub prefill_load_estimator: Option>, + pub overloaded_worker_provider: OverloadedWorkerProvider, + pub available_worker_provider: WorkerAvailabilityProvider, + /// The router's own indexer, used for dequeue-time overlap refresh so + /// queued requests are re-scored against the index that scored them. + pub overlap_refresh: Option>, +} + +/// One `SelectionService` partition driven directly by the router. +pub(crate) struct EmbeddedSelection { + /// Keeps the service (listeners, sweep, replica sync) alive for as long as + /// the router holds the partition. + _service: Arc, + partition: SelectionPartition, + worker_type: &'static str, +} + +impl EmbeddedSelection { + pub(crate) async fn start( + args: EmbeddedSelectionArgs, + workers_with_configs: RuntimeConfigWatch, + cancellation_token: CancellationToken, + ) -> Result { + let worker_type = args.worker_role.unwrap_or(WorkerType::Aggregated); + let key = RoutingPartitionId::new( + args.model_name + .clone() + .unwrap_or_else(|| "default".to_string()), + DEFAULT_ROUTING_GROUP, + ); + + // Same selector the runtime path builds: the registry-resolved policy + // when router configuration names one, else Dynamo's default scorer + // and picker under the router's metric worker label. + let label = args.metric_worker_type; + let policy_factory: WorkerSelectionPolicyFactory = + match WorkerSelectionPolicyRegistry::default() + .resolve_for_worker_type(&args.kv_router_config, worker_type)? + { + Some(factory) => factory, + None => Arc::new(move |config: &KvRouterConfig, _worker_type, _partition| { + WorkerSelectionPolicy::default(config.clone(), label) + }), + }; + + if args.kv_router_config.router_replica_sync { + tracing::warn!( + "router_replica_sync is not carried over to the embedded selection partition; \ + active-load replica sync must be configured on the selection service" + ); + } + + let service = SelectionServiceBuilder::new( + args.kv_router_config.clone(), + worker_type, + WorkerSelectionPolicyRegistry::default(), + ) + .worker_selection_policy_factory(policy_factory) + .indexer_threads(1) + .external_kv_events() + .host_hooks(SelectionHostHooks { + prefill_load_estimator: args.prefill_load_estimator, + overloaded_worker_provider: Some(args.overloaded_worker_provider), + available_worker_provider: Some(args.available_worker_provider), + shared_cache: None, + lora_worker_filter: None, + overlap_refresh: match args.overlap_refresh { + Some(provider) => OverlapRefreshSource::External(provider), + None => OverlapRefreshSource::Disabled, + }, + }) + .build() + .await + .context("failed to start embedded selection service")?; + let service = Arc::new(service); + let partition = service + .ensure_partition(key.clone(), args.block_size, args.is_eagle) + .context("failed to create embedded selection partition")?; + + let feeder = CatalogFeeder { + service: Arc::clone(&service), + key, + block_size: args.block_size, + is_eagle: args.is_eagle, + known: HashMap::new(), + }; + feeder + .run(workers_with_configs, cancellation_token.child_token()) + .await; + + tracing::info!( + worker_type = %worker_type, + "KvRouter scheduling on embedded selection partition" + ); + Ok(Self { + _service: service, + partition, + worker_type: args.metric_worker_type, + }) + } + + /// Membership is catalog-driven (runtime-config watch); explicit worker + /// registration is a no-op here. + pub(crate) fn register_workers(&self, _worker_ids: &HashSet) {} + + pub(crate) async fn schedule_request_admitted( + &self, + request: ScheduleRequest, + ) -> Result { + self.partition + .scheduler() + .schedule_request_admitted(request) + .await + } + + pub(crate) async fn select_without_admission( + &self, + request: ScheduleRequest, + ) -> Result { + self.partition + .scheduler() + .select_without_admission(request) + .await + } + + pub(crate) async fn add_request_admitted( + &self, + req: SequenceRequest, + ) -> Result { + self.partition.scheduler().add_request_admitted(req).await + } + + pub(crate) async fn mark_prefill_completed( + &self, + request_id: &str, + ) -> Result<(), SequenceError> { + self.partition + .scheduler() + .mark_prefill_completed(request_id) + .await + } + + pub(crate) async fn free(&self, request_id: &str) -> Result<(), SequenceError> { + self.partition.scheduler().free(request_id).await + } + + pub(crate) async fn free_if_worker( + &self, + request_id: &str, + worker: WorkerWithDpRank, + ) -> Result<(), SequenceError> { + self.partition + .scheduler() + .free_if_worker(request_id, worker) + .await + } + + pub(crate) fn booking_cleanup(&self) -> SchedulerBookingCleanup { + self.partition.scheduler().booking_cleanup() + } + + pub(crate) async fn mark_prefill_completed_if_booking( + &self, + booking: &SchedulerBookingDescriptor, + ) -> Result<(), KvSchedulerError> { + self.partition + .scheduler() + .mark_prefill_completed_if_booking(booking) + .await + } + + pub(crate) fn pending_count(&self) -> usize { + self.partition.scheduler().pending_count() + } + + pub(crate) fn pending_isl_tokens(&self) -> usize { + self.partition.scheduler().pending_isl_tokens() + } + + pub(crate) fn worker_type(&self) -> &'static str { + self.worker_type + } + + pub(crate) fn add_output_block( + &self, + request_id: &str, + decay_fraction: Option, + ) -> Result<(), SequenceError> { + self.partition + .scheduler() + .add_output_block(request_id, decay_fraction) + } + + pub(crate) async fn enqueue_output_block_if_booking( + &self, + booking: &SchedulerBookingDescriptor, + decay_fraction: Option, + ) -> Result<(), KvSchedulerError> { + self.partition + .scheduler() + .enqueue_output_block_if_booking(booking, decay_fraction) + .await + } + + pub(crate) fn get_potential_loads( + &self, + token_seq: Option>, + isl_tokens: usize, + effective_cached_tokens: HashMap, + track_prefill_tokens: bool, + ) -> Vec { + self.partition.scheduler().get_potential_loads( + token_seq, + isl_tokens, + effective_cached_tokens, + track_prefill_tokens, + ) + } + + pub(crate) fn supports_overlap_refresh(&self) -> bool { + self.partition.scheduler().supports_overlap_refresh() + } +} + +/// Mirrors the runtime-config watch into the partition's worker catalog. +struct CatalogFeeder { + service: Arc, + key: RoutingPartitionId, + block_size: u32, + is_eagle: bool, + known: HashMap, +} + +impl CatalogFeeder { + /// Apply the current snapshot, then keep applying changes until cancelled. + async fn run(mut self, mut watch: RuntimeConfigWatch, cancel: CancellationToken) { + let snapshot = watch.borrow_and_update().clone(); + self.apply(snapshot).await; + tokio::spawn(async move { + loop { + tokio::select! { + _ = cancel.cancelled() => break, + changed = watch.changed() => { + if changed.is_err() { + break; + } + } + } + let snapshot = watch.borrow_and_update().clone(); + self.apply(snapshot).await; + } + }); + } + + async fn apply(&mut self, snapshot: HashMap) { + let removed: Vec = self + .known + .keys() + .copied() + .filter(|worker_id| !snapshot.contains_key(worker_id)) + .collect(); + for worker_id in removed { + self.known.remove(&worker_id); + if let Err(error) = self.service.delete_worker(worker_id).await { + tracing::debug!(worker_id, %error, "embedded selection: worker delete skipped"); + } + } + for (worker_id, config) in snapshot { + if self.known.get(&worker_id) == Some(&config) { + continue; + } + let request = worker_request_from_runtime_config( + worker_id, + &config, + &self.key, + self.block_size, + self.is_eagle, + ); + match self.service.upsert_worker(request).await { + Ok(record) => { + if !record.not_schedulable_reasons.is_empty() { + tracing::warn!( + worker_id, + reasons = ?record.not_schedulable_reasons, + "embedded selection: worker is not schedulable" + ); + } + self.known.insert(worker_id, config); + } + Err(error) => { + tracing::warn!(worker_id, %error, "embedded selection: worker upsert failed"); + } + } + } + } +} + +/// Build the catalog record for a discovered worker. The endpoint is a +/// placeholder: dispatch stays on the router's request transport, keyed by +/// worker id. +pub(crate) fn worker_request_from_runtime_config( + worker_id: WorkerId, + config: &ModelRuntimeConfig, + key: &RoutingPartitionId, + block_size: u32, + is_eagle: bool, +) -> WorkerRequest { + let dp_start = config.data_parallel_start_rank(); + let dp_size = config.data_parallel_size(); + let mut router_hint_worker_type = None; + let mut router_hint_source_control_endpoints = HashMap::new(); + for dp_rank in dp_start..dp_start.saturating_add(dp_size) { + if let Some(metadata) = config.router_hint_metadata_for_dp_rank(dp_rank) { + router_hint_worker_type.get_or_insert_with(|| metadata.worker_type.to_string()); + if let Some(endpoint) = metadata.source_control_endpoint { + router_hint_source_control_endpoints.insert(dp_rank, endpoint.to_string()); + } + } + } + WorkerRequest { + worker_id, + model_name: key.model_name.clone(), + routing_group: key.routing_group.clone(), + endpoint: Some(format!("dyn://{worker_id}")), + block_size: Some(block_size), + data_parallel_start_rank: Some(dp_start), + data_parallel_size: Some(dp_size), + max_num_batched_tokens: config.max_num_batched_tokens, + total_kv_blocks: config.total_kv_blocks, + stable_routing_id: config.stable_routing_id.clone(), + is_eagle: Some(is_eagle), + taints: config.taints().clone(), + topology_domains: config.topology_domains.clone(), + kv_transfer_domain: config.kv_transfer_domain.clone(), + kv_transfer_enforcement: config.kv_transfer_enforcement, + kv_transfer_preferred_weight: config.kv_transfer_preferred_weight, + router_hint_worker_type, + router_hint_source_control_endpoints, + ..WorkerRequest::default() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn worker_request_mirrors_runtime_config() { + let mut config = ModelRuntimeConfig { + data_parallel_start_rank: 2, + data_parallel_size: 2, + max_num_batched_tokens: Some(8192), + total_kv_blocks: Some(4096), + stable_routing_id: Some("worker-0".to_string()), + ..ModelRuntimeConfig::default() + }; + config.taints.insert("gpu=h100".to_string()); + config.runtime_data.insert( + dynamo_kv_router::router_hint::ROUTER_HINT_RUNTIME_CAPABILITY_KEY.to_string(), + serde_json::Value::Bool(true), + ); + config.runtime_data.insert( + dynamo_kv_router::router_hint::ROUTER_HINT_WORKER_TYPE_RUNTIME_KEY.to_string(), + serde_json::Value::String("decode".to_string()), + ); + config.runtime_data.insert( + dynamo_kv_router::router_hint::ROUTER_HINT_SOURCE_CONTROL_ENDPOINTS_RUNTIME_KEY + .to_string(), + serde_json::json!({"2": "tcp://w:9002", "3": "tcp://w:9003"}), + ); + let key = RoutingPartitionId::new("model", DEFAULT_ROUTING_GROUP); + let request = worker_request_from_runtime_config(7, &config, &key, 16, false); + assert_eq!(request.worker_id, 7); + assert_eq!(request.model_name, "model"); + assert_eq!(request.endpoint.as_deref(), Some("dyn://7")); + assert_eq!(request.block_size, Some(16)); + assert_eq!(request.data_parallel_start_rank, Some(2)); + assert_eq!(request.data_parallel_size, Some(2)); + assert_eq!(request.max_num_batched_tokens, Some(8192)); + assert_eq!(request.total_kv_blocks, Some(4096)); + assert_eq!(request.stable_routing_id.as_deref(), Some("worker-0")); + assert!(request.taints.contains("gpu=h100")); + assert_eq!(request.router_hint_worker_type.as_deref(), Some("decode")); + assert_eq!( + request.router_hint_source_control_endpoints, + HashMap::from([ + (2, "tcp://w:9002".to_string()), + (3, "tcp://w:9003".to_string()) + ]) + ); + } +} diff --git a/lib/llm/src/kv_router/routing_host/tests.rs b/lib/llm/src/kv_router/routing_host/tests.rs index bbe5e2315d2b..92537a90ae31 100644 --- a/lib/llm/src/kv_router/routing_host/tests.rs +++ b/lib/llm/src/kv_router/routing_host/tests.rs @@ -1154,6 +1154,27 @@ async fn router_with_workers( async fn router_with_worker_configs( session_affinity_ttl: Option, workers: HashMap, +) -> (RoutingHost, Runtime) { + router_with_worker_configs_in_mode(session_affinity_ttl, workers, false).await +} + +/// Same router, scheduling on an embedded selection-service partition. +async fn embedded_router_with_workers( + session_affinity_ttl: Option, + worker_ids: &[u64], +) -> (RoutingHost, Runtime) { + let workers = worker_ids + .iter() + .copied() + .map(|worker_id| (worker_id, ModelRuntimeConfig::default())) + .collect(); + router_with_worker_configs_in_mode(session_affinity_ttl, workers, true).await +} + +async fn router_with_worker_configs_in_mode( + session_affinity_ttl: Option, + workers: HashMap, + embedded_selection: bool, ) -> (RoutingHost, Runtime) { let runtime = Runtime::from_current().unwrap(); let distributed = DistributedRuntime::new(runtime.clone(), DistributedConfig::process_local()) @@ -1172,6 +1193,7 @@ async fn router_with_worker_configs( skip_initial_worker_wait: true, use_kv_events: false, router_track_active_blocks: false, + router_embedded_selection: embedded_selection, ..Default::default() }; let chooser = KvRouter::new( @@ -1405,6 +1427,98 @@ async fn track_request( (request, selection, guard) } +/// The embedded selection partition books and releases exactly like the +/// runtime scheduler: preview does not book, plan books, abort releases, and +/// a tracked request frees on guard completion. +#[tokio::test] +#[serial_test::serial] +async fn embedded_selection_books_and_releases_like_the_runtime_scheduler() { + let (router, runtime) = embedded_router_with_workers(None, &[7, 9]).await; + let request = Context::new(request()); + + let preview = router + .preview_kv_route(&request, RequestPhase::Decode) + .await + .expect("decode preview should select one worker"); + let loads = router + .kv_router() + .get_potential_loads(&[], None, None, None, None) + .await + .unwrap(); + assert!( + loads.iter().all(|load| load.active_requests == 0), + "a preview must not book: {loads:?}" + ); + + let plan = router + .plan_kv_route_from_preview(&request, preview) + .await + .expect("decode plan should admit one request"); + let planned_worker = plan.signals().worker.worker_id; + assert!([7, 9].contains(&planned_worker)); + let loads = router + .kv_router() + .get_potential_loads(&[], None, None, None, None) + .await + .unwrap(); + assert_eq!( + loads + .iter() + .find(|load| load.worker_id == planned_worker && load.dp_rank == 0) + .expect("selected worker must be reported") + .active_requests, + 1 + ); + plan.abort().await; + let loads = router + .kv_router() + .get_potential_loads(&[], None, None, None, None) + .await + .unwrap(); + assert!( + loads.iter().all(|load| load.active_requests == 0), + "abandoned plans must release their embedded reservation: {loads:?}" + ); + + // A tracked request books on admission and frees when its guard finishes. + let (_request, selection, mut guard) = track_request(&router, false).await; + let loads = router + .kv_router() + .get_potential_loads(&[], None, None, None, None) + .await + .unwrap(); + assert_eq!( + loads + .iter() + .find(|load| load.worker_id == selection.worker.worker_id) + .expect("tracked worker must be reported") + .active_requests, + 1 + ); + guard.abort().await; + let released = tokio::time::timeout(Duration::from_secs(2), async { + loop { + let loads = router + .kv_router() + .get_potential_loads(&[], None, None, None, None) + .await + .unwrap(); + if loads.iter().all(|load| load.active_requests == 0) { + return; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await; + assert!( + released.is_ok(), + "guard drop must release the embedded booking" + ); + + drop(router); + runtime.shutdown(); +} + #[tokio::test] #[serial_test::serial] async fn route_plan_from_preview_holds_and_releases_the_decode_reservation() { From fc8a67c33b15b17623cb2afa905065ac495b7051 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Thu, 3 Sep 2026 13:21:47 +0000 Subject: [PATCH 007/102] fix(router): feed scheduler load from the embedded selection partition KvWorkerMonitor's prefill-token overload detection reads the scheduler's own load snapshots through SchedulerLoadSender. The embedded selection partition never published them, so --active-prefill-tokens-threshold could not trigger 529 under --router-embedded-selection (tests/router test_mocker_kv_router_overload_529[prefill-tokens] failed only on that path; decode-blocks worked because it uses worker-published metrics). Add SelectionHostHooks.scheduler_load_sink (a SchedulerLoadSink) that the partition's ScopedSequencePublisher forwards publish_scheduler_load(_batch) into; the embedded host bridges it to the router's SchedulerLoadSender. tests/router overload_529 (agg + disagg, prefill-tokens + decode-blocks) now pass on the embedded path. Signed-off-by: Ishan Dhanani (cherry picked from commit beb6ddc50869c3a4a5d3902c7e0b43011437d6f9) --- .../src/services/common/replica_sync.rs | 39 ++++++++++++++++++- .../src/services/selection/core/mod.rs | 14 ++++++- lib/kv-router/src/services/selection/mod.rs | 4 +- lib/llm/src/kv_router.rs | 1 + lib/llm/src/kv_router/embedded.rs | 18 +++++++++ 5 files changed, 71 insertions(+), 5 deletions(-) diff --git a/lib/kv-router/src/services/common/replica_sync.rs b/lib/kv-router/src/services/common/replica_sync.rs index 4d33bc879e2a..f57dbb2a1b45 100644 --- a/lib/kv-router/src/services/common/replica_sync.rs +++ b/lib/kv-router/src/services/common/replica_sync.rs @@ -117,9 +117,24 @@ pub(crate) struct ScopedReplicaSync { pub channel: Option<(mpsc::Sender, ChannelSequenceSubscriber)>, } +/// Receives the scheduler-owned load snapshots a partition's active-sequence +/// tracker publishes (per worker: active decode blocks, active prefill +/// tokens). An embedding host uses it to drive overload detection from the +/// same numbers the scheduler books against. +pub trait SchedulerLoadSink: Send + Sync { + fn publish(&self, snapshot: crate::sequences::SchedulerLoadSnapshot); + + fn publish_batch(&self, snapshots: Vec) { + for snapshot in snapshots { + self.publish(snapshot); + } + } +} + #[derive(Clone)] pub struct ScopedSequencePublisher { replica: Option, + load_sink: Option>, } #[derive(Clone)] @@ -132,7 +147,16 @@ struct ScopedReplicaPublisher { impl ScopedSequencePublisher { pub(crate) fn disabled() -> Self { - Self { replica: None } + Self { + replica: None, + load_sink: None, + } + } + + /// Forward scheduler load snapshots to `sink` in addition to replica sync. + pub(crate) fn with_load_sink(mut self, sink: Option>) -> Self { + self.load_sink = sink; + self } pub(crate) fn enabled( @@ -148,6 +172,7 @@ impl ScopedSequencePublisher { tx, cancel_token, }), + load_sink: None, } } } @@ -186,7 +211,17 @@ impl SequencePublisher for ScopedSequencePublisher { } } - fn publish_scheduler_load(&self, _load: crate::sequences::SchedulerLoadSnapshot) {} + fn publish_scheduler_load(&self, load: crate::sequences::SchedulerLoadSnapshot) { + if let Some(sink) = &self.load_sink { + sink.publish(load); + } + } + + fn publish_scheduler_load_batch(&self, loads: Vec) { + if let Some(sink) = &self.load_sink { + sink.publish_batch(loads); + } + } fn observe_load( &self, diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index 5ef463644ee7..a2d458e2fac8 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -35,7 +35,8 @@ use crate::sequences::{ ActiveSequencesMultiWorker, ReplicaWorkerPolicy, SequenceError, SequenceRequest, }; use crate::services::common::replica_sync::{ - ReplicaSyncConfig, ScopedReplicaEvent, ScopedSequencePublisher, setup_scoped_replica_sync, + ReplicaSyncConfig, SchedulerLoadSink, ScopedReplicaEvent, ScopedSequencePublisher, + setup_scoped_replica_sync, }; use crate::services::indexer::backend::{Indexer, IndexerPolicy}; use crate::services::indexer::recovery; @@ -203,6 +204,10 @@ pub struct SelectionHostHooks { /// that index so queued requests are re-scored against the index that /// scored them, or disables refresh when its index cannot serve it. pub overlap_refresh: OverlapRefreshSource, + /// Receives each partition's scheduler-owned load snapshots (active decode + /// blocks and prefill tokens per worker) so the host can run overload + /// detection on the numbers the scheduler books against. + pub scheduler_load_sink: Option>, } /// Source of dequeue-time overlap refreshes for partition schedulers. @@ -662,7 +667,9 @@ impl SelectionCore { setup_scoped_replica_sync(self.replica_config.as_ref(), &key, block_size); let worker_label = self.worker_type.as_str(); let slots = Arc::new(ActiveSequencesMultiWorker::new_with_replica_worker_policy( - scoped_replica_sync.publisher, + scoped_replica_sync + .publisher + .with_load_sink(self.host_hooks.scheduler_load_sink.clone()), block_size as usize, HashMap::new(), scoped_replica_sync.enabled, @@ -2466,6 +2473,7 @@ mod tests { })), lora_worker_filter: None, overlap_refresh: OverlapRefreshSource::default(), + scheduler_load_sink: None, }, Some(factory), ); @@ -2567,6 +2575,7 @@ mod tests { shared_cache: None, lora_worker_filter: None, overlap_refresh: OverlapRefreshSource::default(), + scheduler_load_sink: None, }); core.upsert_worker(worker(1)).await.expect("worker upsert"); core.upsert_worker(worker(2)).await.expect("worker upsert"); @@ -2589,6 +2598,7 @@ mod tests { shared_cache: None, lora_worker_filter: None, overlap_refresh: OverlapRefreshSource::default(), + scheduler_load_sink: None, }); core.upsert_worker(worker(1)).await.expect("worker upsert"); core.upsert_worker(worker(2)).await.expect("worker upsert"); diff --git a/lib/kv-router/src/services/selection/mod.rs b/lib/kv-router/src/services/selection/mod.rs index cd60810fcbb5..19756ed4f485 100644 --- a/lib/kv-router/src/services/selection/mod.rs +++ b/lib/kv-router/src/services/selection/mod.rs @@ -20,7 +20,9 @@ mod types; mod tests; pub use crate::WorkerSelectionPolicyFactory; -pub use crate::services::common::replica_sync::{ReplicaPeerError, ScopedSequencePublisher}; +pub use crate::services::common::replica_sync::{ + ReplicaPeerError, SchedulerLoadSink, ScopedSequencePublisher, +}; pub use core::{ OverlapRefreshSource, RefreshProvider, SelectionCore, SelectionHostHooks, SelectionPartition, SelectionScheduler, SelectionServiceConfig, diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index c89db06fb122..7e9ab927ff52 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -828,6 +828,7 @@ where overloaded_worker_provider, available_worker_provider, overlap_refresh, + scheduler_load, }, workers_with_configs.clone(), cancellation_token.child_token(), diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs index c658b3ac3e87..e073f358234d 100644 --- a/lib/llm/src/kv_router/embedded.rs +++ b/lib/llm/src/kv_router/embedded.rs @@ -54,6 +54,9 @@ pub(crate) struct EmbeddedSelectionArgs { /// The router's own indexer, used for dequeue-time overlap refresh so /// queued requests are re-scored against the index that scored them. pub overlap_refresh: Option>, + /// Scheduler-owned load snapshots for the worker monitor's overload + /// detection, the same feed the runtime scheduler publishes. + pub scheduler_load: crate::kv_router::routing_load::SchedulerLoadSender, } /// One `SelectionService` partition driven directly by the router. @@ -65,6 +68,20 @@ pub(crate) struct EmbeddedSelection { worker_type: &'static str, } +/// Bridges the partition's scheduler load snapshots to the router's +/// `SchedulerLoadSender`, which feeds `KvWorkerMonitor`. +struct SenderLoadSink(crate::kv_router::routing_load::SchedulerLoadSender); + +impl dynamo_kv_router::services::selection::SchedulerLoadSink for SenderLoadSink { + fn publish(&self, snapshot: dynamo_kv_router::sequences::SchedulerLoadSnapshot) { + self.0.publish(snapshot); + } + + fn publish_batch(&self, snapshots: Vec) { + self.0.publish_batch(snapshots); + } +} + impl EmbeddedSelection { pub(crate) async fn start( args: EmbeddedSelectionArgs, @@ -118,6 +135,7 @@ impl EmbeddedSelection { Some(provider) => OverlapRefreshSource::External(provider), None => OverlapRefreshSource::Disabled, }, + scheduler_load_sink: Some(Arc::new(SenderLoadSink(args.scheduler_load))), }) .build() .await From ad7e918ca2a124f688571fb9b2e256b6845d1f59 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Thu, 3 Sep 2026 13:41:52 +0000 Subject: [PATCH 008/102] feat(router): make embedded selection the default scheduler Phase 1 exit gates now pass: the frontend A/B benchmark (+6% req/s, lower p50/p99), offline KV replay byte-parity against main on four rows, and tests/router e2e green on both scheduler paths. Flip router_embedded_selection to default true (serde default/skip updated, PyO3 constructor, .pyi, --router-embedded-selection default, docs). The runtime-bound KvScheduler path stays reachable through --no-router-embedded-selection for one release as a rollback and logs a deprecation warning at startup; deleting it is the follow-up. A caller-injected custom WorkerSelector type (not DefaultWorkerSelector) cannot run on the partition, so KvRouter::new keeps the runtime scheduler for it with a warning instead of silently changing selection semantics. Signed-off-by: Ishan Dhanani (cherry picked from commit 568d06c49d) --- .../configuration/groups/kv_router_args.py | 9 +++++---- .../components/frontend-configuration.mdx | 11 ++++++----- lib/bindings/python/rust/llm/entrypoint.rs | 2 +- lib/bindings/python/src/dynamo/_core.pyi | 5 +++-- lib/kv-router/src/scheduling/config.rs | 19 ++++++++++++++----- lib/llm/src/kv_router.rs | 18 +++++++++++++++++- 6 files changed, 46 insertions(+), 18 deletions(-) diff --git a/components/src/dynamo/common/configuration/groups/kv_router_args.py b/components/src/dynamo/common/configuration/groups/kv_router_args.py index 87091c36b2cc..550867216884 100644 --- a/components/src/dynamo/common/configuration/groups/kv_router_args.py +++ b/components/src/dynamo/common/configuration/groups/kv_router_args.py @@ -201,7 +201,7 @@ class KvRouterConfigBase(ConfigBase): router_temperature: float use_kv_events: bool router_replica_sync: bool - router_embedded_selection: bool = False + router_embedded_selection: bool = True router_track_active_blocks: bool router_track_output_blocks: bool router_assume_kv_reuse: bool @@ -429,12 +429,13 @@ def add_arguments(self, parser) -> None: g, flag_name="--router-embedded-selection", env_var="DYN_ROUTER_EMBEDDED_SELECTION", - default=False, + default=True, dest="router_embedded_selection", help=( "KV Router: Run scheduling and active-sequence accounting on an " - "embedded selection-service partition instead of the runtime-bound " - "scheduler. Experimental; the KV indexer and request transport are unchanged." + "embedded selection-service partition (default). " + "--no-router-embedded-selection selects the deprecated runtime-bound " + "scheduler, kept for one release as a rollback." ), ) add_negatable_bool_argument( diff --git a/docs/fern/pages/reference/components/frontend-configuration.mdx b/docs/fern/pages/reference/components/frontend-configuration.mdx index e2d1615de025..0ad51fa491b6 100644 --- a/docs/fern/pages/reference/components/frontend-configuration.mdx +++ b/docs/fern/pages/reference/components/frontend-configuration.mdx @@ -213,11 +213,12 @@ router. See [Router Guide](../../developer-guide/knowledge-base/modular-componen Environment variable: `DYN_ROUTER_REPLICA_SYNC` - - **Experimental.** Run the router's scheduling and active-sequence accounting on an embedded - selection-service partition (the same core as the standalone selection service) instead of the - runtime-bound scheduler. The KV indexer, request transport, and worker discovery are unchanged; - `--router-replica-sync` is not applied to the embedded partition. + + Run the router's scheduling and active-sequence accounting on an embedded selection-service + partition (the same core as the standalone selection service). The KV indexer, request + transport, and worker discovery are unchanged; `--router-replica-sync` is not applied to the + embedded partition. `--no-router-embedded-selection` selects the runtime-bound scheduler, which + is deprecated and kept for one release as a rollback; it logs a warning at startup. Environment variable: `DYN_ROUTER_EMBEDDED_SELECTION` diff --git a/lib/bindings/python/rust/llm/entrypoint.rs b/lib/bindings/python/rust/llm/entrypoint.rs index c378d1aa1bba..d6c17e300f94 100644 --- a/lib/bindings/python/rust/llm/entrypoint.rs +++ b/lib/bindings/python/rust/llm/entrypoint.rs @@ -242,7 +242,7 @@ impl AicPerfConfig { #[pymethods] impl KvRouterConfig { #[new] - #[pyo3(signature = (overlap_score_weight=None, host_cache_hit_weight=0.75, disk_cache_hit_weight=0.25, router_temperature=0.0, use_kv_events=true, *, router_replica_sync=false, router_embedded_selection=false, router_track_active_blocks=true, router_track_output_blocks=false, router_assume_kv_reuse=true, router_track_prefill_tokens=true, router_prefill_load_model="none", router_ttl_secs=120.0, router_approximate_cache_policy="ttl", router_queue_threshold=None, router_event_threads=4, router_queue_policy="fcfs", use_remote_indexer=false, serve_indexer=false, shared_cache_multiplier=0.0, shared_cache_type="none", router_predicted_ttl_secs=None, conditional_disagg_enabled=false, conditional_disagg_policy="isl_bounding", conditional_disagg_eff_isl_threshold=2048, conditional_disagg_eff_isl_ratio_threshold=0.7, conditional_disagg_prefill_busy_threshold=None, conditional_disagg_decode_busy_threshold=None, overlap_score_credit=1.0, overlap_score_credit_decay=0.0, prefill_load_scale=1.0, decode_active_request_weight=0.0, router_policy_config=None, router_prefill_policy=None, router_decode_policy=None, router_tracking_hash="public-xxh3-v1", router_tracking_key_file=None, router_tracking_key_id=None))] + #[pyo3(signature = (overlap_score_weight=None, host_cache_hit_weight=0.75, disk_cache_hit_weight=0.25, router_temperature=0.0, use_kv_events=true, *, router_replica_sync=false, router_embedded_selection=true, router_track_active_blocks=true, router_track_output_blocks=false, router_assume_kv_reuse=true, router_track_prefill_tokens=true, router_prefill_load_model="none", router_ttl_secs=120.0, router_approximate_cache_policy="ttl", router_queue_threshold=None, router_event_threads=4, router_queue_policy="fcfs", use_remote_indexer=false, serve_indexer=false, shared_cache_multiplier=0.0, shared_cache_type="none", router_predicted_ttl_secs=None, conditional_disagg_enabled=false, conditional_disagg_policy="isl_bounding", conditional_disagg_eff_isl_threshold=2048, conditional_disagg_eff_isl_ratio_threshold=0.7, conditional_disagg_prefill_busy_threshold=None, conditional_disagg_decode_busy_threshold=None, overlap_score_credit=1.0, overlap_score_credit_decay=0.0, prefill_load_scale=1.0, decode_active_request_weight=0.0, router_policy_config=None, router_prefill_policy=None, router_decode_policy=None, router_tracking_hash="public-xxh3-v1", router_tracking_key_file=None, router_tracking_key_id=None))] #[allow(clippy::too_many_arguments)] fn new( overlap_score_weight: Option, diff --git a/lib/bindings/python/src/dynamo/_core.pyi b/lib/bindings/python/src/dynamo/_core.pyi index 1e192f8f72c7..b7a2a962a7d7 100644 --- a/lib/bindings/python/src/dynamo/_core.pyi +++ b/lib/bindings/python/src/dynamo/_core.pyi @@ -1784,7 +1784,7 @@ class KvRouterConfig: use_kv_events: bool = True, *, router_replica_sync: bool = False, - router_embedded_selection: bool = False, + router_embedded_selection: bool = True, router_track_active_blocks: bool = True, router_track_output_blocks: bool = False, router_assume_kv_reuse: bool = True, @@ -1831,7 +1831,8 @@ class KvRouterConfig: use_kv_events: Whether to use KV events from workers (default: True) router_replica_sync: Enable replica synchronization (default: False) router_embedded_selection: Run scheduling on an embedded selection-service - partition instead of the runtime-bound scheduler (default: False). Experimental. + partition (default: True). False selects the deprecated runtime-bound + scheduler, kept for one release as a rollback. router_track_active_blocks: Track active blocks for load balancing (default: True) router_track_output_blocks: Track output blocks during generation (default: False). When enabled, the router adds placeholder blocks as tokens are generated diff --git a/lib/kv-router/src/scheduling/config.rs b/lib/kv-router/src/scheduling/config.rs index ba7d913dc527..4f6ffe4a34e9 100644 --- a/lib/kv-router/src/scheduling/config.rs +++ b/lib/kv-router/src/scheduling/config.rs @@ -95,6 +95,14 @@ const fn default_decode_active_request_weight() -> f64 { // Default-valued post-v1.3 fields are omitted so v1.3 frontends can read v1.4 MDCs during // rolling upgrades. Non-default values still serialize and fail closed on the older frontend. // TODO(v1.5): Remove these compatibility skips when v1.3 falls outside the N-1 window. +fn default_true() -> bool { + true +} + +fn is_true(value: &bool) -> bool { + *value +} + fn is_default(value: &T) -> bool { value == &T::default() } @@ -669,7 +677,7 @@ struct KvRouterConfigSerde { router_track_output_blocks: bool, router_assume_kv_reuse: bool, router_track_prefill_tokens: bool, - #[serde(default)] + #[serde(default = "default_true")] router_embedded_selection: bool, router_tracking_hash: TrackingHashAlgorithm, router_tracking_key_file: Option, @@ -804,9 +812,10 @@ pub struct KvRouterConfig { pub router_track_prefill_tokens: bool, /// Run the frontend router's scheduling and active-sequence accounting on - /// an embedded selection-service partition instead of the runtime-bound - /// scheduler (default: false). Frontend-local; workers ignore it. - #[serde(default, skip_serializing_if = "is_default")] + /// an embedded selection-service partition (default: true). `false` + /// selects the deprecated runtime-bound scheduler, kept for one release as + /// a rollback. Frontend-local; workers ignore it. + #[serde(default = "default_true", skip_serializing_if = "is_true")] pub router_embedded_selection: bool, /// Hash algorithm used for router-derived active-sequence identities. @@ -971,7 +980,7 @@ impl Default for KvRouterConfig { router_track_output_blocks: false, router_assume_kv_reuse: true, router_track_prefill_tokens: default_track_prefill_tokens(), - router_embedded_selection: false, + router_embedded_selection: true, router_tracking_hash: TrackingHashAlgorithm::default(), router_tracking_key_file: None, router_tracking_key_id: None, diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index 7e9ab927ff52..2da26383eb98 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -807,7 +807,18 @@ where let available_worker_provider: WorkerAvailabilityProvider = Arc::new(move || client_for_availability.available_instance_ids()); - let scheduler = if kv_router_config.router_embedded_selection { + // The embedded partition builds its own selection policy (the default + // scorer/picker or a registry-resolved policy); a caller-injected custom + // `WorkerSelector` type can only run on the runtime scheduler. + let custom_selector = std::any::TypeId::of::() + != std::any::TypeId::of::(); + if kv_router_config.router_embedded_selection && custom_selector { + tracing::warn!( + selector = std::any::type_name::(), + "custom WorkerSelector is not supported on the embedded selection partition; using the runtime scheduler" + ); + } + let scheduler = if kv_router_config.router_embedded_selection && !custom_selector { // The router keeps its indexer and transport; scheduling and // active-sequence accounting run on a selection-service partition. let overlap_refresh: Option> = @@ -836,6 +847,11 @@ where .await?, ) } else { + tracing::warn!( + "the runtime-bound KV router scheduler (--no-router-embedded-selection) is \ + deprecated and kept for one release as a rollback; the embedded selection \ + partition is the default" + ); RouterScheduler::Runtime( KvScheduler::start_with_shared_request_leases( endpoint.clone(), From 616bcefbf1007183090f668a1889e1e2f7182a9a Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Thu, 3 Sep 2026 14:11:50 +0000 Subject: [PATCH 009/102] feat(router): replica sync for the embedded selection partition --router-replica-sync previously logged that it was not carried over to the embedded partition. Add SelectionHostHooks.replica_sync, a per-partition factory returning HostReplicaChannels (outbound events, inbound peer events, this replica's id); when the service runs no ZMQ peer mesh the partition's publisher forwards to the host's outbound channel and its replica consumer reads the host's inbound channel. The embedded host builds those channels over the runtime event plane (ACTIVE_SEQUENCES_SUBJECT, same wire formats and subscriber as the runtime scheduler used), so multi-replica frontends keep active-sequence sync on the default scheduler path. Signed-off-by: Ishan Dhanani (cherry picked from commit d58689324514c144a8ff2f0059e121955d93d66c) --- .../src/services/common/replica_sync.rs | 56 ++++++++++++++++ .../src/services/selection/core/mod.rs | 24 +++++-- lib/kv-router/src/services/selection/mod.rs | 3 +- .../src/services/slot_tracker/registry.rs | 2 +- lib/llm/src/kv_router.rs | 2 + lib/llm/src/kv_router/embedded.rs | 30 +++++++-- lib/llm/src/kv_router/sequence.rs | 64 +++++++++++++++++++ 7 files changed, 169 insertions(+), 12 deletions(-) diff --git a/lib/kv-router/src/services/common/replica_sync.rs b/lib/kv-router/src/services/common/replica_sync.rs index f57dbb2a1b45..0a9e23bebd5f 100644 --- a/lib/kv-router/src/services/common/replica_sync.rs +++ b/lib/kv-router/src/services/common/replica_sync.rs @@ -131,9 +131,28 @@ pub trait SchedulerLoadSink: Send + Sync { } } +/// Replica-sync plumbing an embedding host supplies for one partition when it +/// carries active-sequence events over its own transport (for example the +/// Dynamo runtime event plane) instead of the service's ZMQ peer mesh: +/// events the partition emits go to `outbound`; events from peer replicas +/// arrive on `inbound_rx` (the host also holds `inbound_tx`). +pub struct HostReplicaChannels { + pub outbound: mpsc::Sender, + pub inbound_tx: mpsc::Sender, + pub inbound_rx: mpsc::Receiver, + /// This replica's id; events carrying it are ignored on receipt. + pub process_id: u64, +} + +/// Per-partition factory for [`HostReplicaChannels`]; `None` disables replica +/// sync for that partition. +pub type HostReplicaSyncFactory = + Arc Option + Send + Sync>; + #[derive(Clone)] pub struct ScopedSequencePublisher { replica: Option, + host_tx: Option>, load_sink: Option>, } @@ -149,6 +168,16 @@ impl ScopedSequencePublisher { pub(crate) fn disabled() -> Self { Self { replica: None, + host_tx: None, + load_sink: None, + } + } + + /// Publish through an embedding host's outbound channel. + pub(crate) fn host(outbound: mpsc::Sender) -> Self { + Self { + replica: None, + host_tx: Some(outbound), load_sink: None, } } @@ -172,6 +201,7 @@ impl ScopedSequencePublisher { tx, cancel_token, }), + host_tx: None, load_sink: None, } } @@ -179,6 +209,19 @@ impl ScopedSequencePublisher { impl SequencePublisher for ScopedSequencePublisher { fn enqueue_event(&self, event: ActiveSequenceEvent) -> Result<()> { + if let Some(host_tx) = &self.host_tx { + return match host_tx.try_send(event) { + Ok(()) => Ok(()), + Err(mpsc::error::TrySendError::Full(event)) => Err(anyhow::Error::new( + SequencePublishQueueError::full(event, host_tx.max_capacity()), + ) + .context("host replica publisher")), + Err(mpsc::error::TrySendError::Closed(event)) => Err(anyhow::Error::new( + SequencePublishQueueError::closed(event, host_tx.max_capacity(), true), + ) + .context("host replica publisher")), + }; + } let Some(replica) = &self.replica else { return Ok(()); }; @@ -299,8 +342,21 @@ pub(crate) fn setup_scoped_replica_sync( config: Option<&ReplicaSyncConfig>, partition: &RoutingPartitionId, block_size: u32, + host: Option, ) -> ScopedReplicaSync { let Some(config) = config else { + // No peer mesh: an embedding host may still carry replica events. + if let Some(host) = host { + return ScopedReplicaSync { + publisher: ScopedSequencePublisher::host(host.outbound), + enabled: true, + process_id: host.process_id, + channel: Some(( + host.inbound_tx, + ChannelSequenceSubscriber::new(host.inbound_rx), + )), + }; + } return ScopedReplicaSync { publisher: ScopedSequencePublisher::disabled(), enabled: false, diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index a2d458e2fac8..a0c98ab98373 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -35,8 +35,8 @@ use crate::sequences::{ ActiveSequencesMultiWorker, ReplicaWorkerPolicy, SequenceError, SequenceRequest, }; use crate::services::common::replica_sync::{ - ReplicaSyncConfig, SchedulerLoadSink, ScopedReplicaEvent, ScopedSequencePublisher, - setup_scoped_replica_sync, + HostReplicaSyncFactory, ReplicaSyncConfig, SchedulerLoadSink, ScopedReplicaEvent, + ScopedSequencePublisher, setup_scoped_replica_sync, }; use crate::services::indexer::backend::{Indexer, IndexerPolicy}; use crate::services::indexer::recovery; @@ -208,6 +208,10 @@ pub struct SelectionHostHooks { /// blocks and prefill tokens per worker) so the host can run overload /// detection on the numbers the scheduler books against. pub scheduler_load_sink: Option>, + /// Replica sync carried by the host's own transport for partitions this + /// service does not mesh itself (ignored when the service runs its own + /// ZMQ replica sync). + pub replica_sync: Option, } /// Source of dequeue-time overlap refreshes for partition schedulers. @@ -663,8 +667,17 @@ impl SelectionCore { let entry = entry_cell .get_or_try_init(|| -> Result, SelectionError> { let (workers_tx, workers_rx) = watch::channel(HashMap::new()); - let scoped_replica_sync = - setup_scoped_replica_sync(self.replica_config.as_ref(), &key, block_size); + let host_replica = self + .host_hooks + .replica_sync + .as_ref() + .and_then(|factory| factory(&key)); + let scoped_replica_sync = setup_scoped_replica_sync( + self.replica_config.as_ref(), + &key, + block_size, + host_replica, + ); let worker_label = self.worker_type.as_str(); let slots = Arc::new(ActiveSequencesMultiWorker::new_with_replica_worker_policy( scoped_replica_sync @@ -2474,6 +2487,7 @@ mod tests { lora_worker_filter: None, overlap_refresh: OverlapRefreshSource::default(), scheduler_load_sink: None, + replica_sync: None, }, Some(factory), ); @@ -2576,6 +2590,7 @@ mod tests { lora_worker_filter: None, overlap_refresh: OverlapRefreshSource::default(), scheduler_load_sink: None, + replica_sync: None, }); core.upsert_worker(worker(1)).await.expect("worker upsert"); core.upsert_worker(worker(2)).await.expect("worker upsert"); @@ -2599,6 +2614,7 @@ mod tests { lora_worker_filter: None, overlap_refresh: OverlapRefreshSource::default(), scheduler_load_sink: None, + replica_sync: None, }); core.upsert_worker(worker(1)).await.expect("worker upsert"); core.upsert_worker(worker(2)).await.expect("worker upsert"); diff --git a/lib/kv-router/src/services/selection/mod.rs b/lib/kv-router/src/services/selection/mod.rs index 19756ed4f485..9a96ef28437e 100644 --- a/lib/kv-router/src/services/selection/mod.rs +++ b/lib/kv-router/src/services/selection/mod.rs @@ -21,7 +21,8 @@ mod tests; pub use crate::WorkerSelectionPolicyFactory; pub use crate::services::common::replica_sync::{ - ReplicaPeerError, SchedulerLoadSink, ScopedSequencePublisher, + HostReplicaChannels, HostReplicaSyncFactory, ReplicaPeerError, SchedulerLoadSink, + ScopedSequencePublisher, }; pub use core::{ OverlapRefreshSource, RefreshProvider, SelectionCore, SelectionHostHooks, SelectionPartition, diff --git a/lib/kv-router/src/services/slot_tracker/registry.rs b/lib/kv-router/src/services/slot_tracker/registry.rs index 6aa627e4c9ec..8a39bc77bff2 100644 --- a/lib/kv-router/src/services/slot_tracker/registry.rs +++ b/lib/kv-router/src/services/slot_tracker/registry.rs @@ -105,7 +105,7 @@ impl TrackerEntry { replica_config: Option<&ReplicaSyncConfig>, ) -> Arc { let cancel_token = root_cancel_token.child_token(); - let scoped_replica_sync = setup_scoped_replica_sync(replica_config, key, block_size); + let scoped_replica_sync = setup_scoped_replica_sync(replica_config, key, block_size, None); let tracker = Arc::new(ActiveSequencesMultiWorker::new_with_replica_worker_policy( scoped_replica_sync.publisher, block_size as usize, diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index 2da26383eb98..2607e539ead7 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -840,6 +840,8 @@ where available_worker_provider, overlap_refresh, scheduler_load, + endpoint: endpoint.clone(), + router_id: endpoint.drt().discovery().instance_id(), }, workers_with_configs.clone(), cancellation_token.child_token(), diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs index e073f358234d..9d9869489ff2 100644 --- a/lib/llm/src/kv_router/embedded.rs +++ b/lib/llm/src/kv_router/embedded.rs @@ -57,6 +57,11 @@ pub(crate) struct EmbeddedSelectionArgs { /// Scheduler-owned load snapshots for the worker monitor's overload /// detection, the same feed the runtime scheduler publishes. pub scheduler_load: crate::kv_router::routing_load::SchedulerLoadSender, + /// Endpoint whose event plane carries replica sync when + /// `router_replica_sync` is set. + pub endpoint: dynamo_runtime::component::Endpoint, + /// This router replica's id (ignored on receipt of its own events). + pub router_id: u64, } /// One `SelectionService` partition driven directly by the router. @@ -110,12 +115,24 @@ impl EmbeddedSelection { }), }; - if args.kv_router_config.router_replica_sync { - tracing::warn!( - "router_replica_sync is not carried over to the embedded selection partition; \ - active-load replica sync must be configured on the selection service" - ); - } + // Replica sync rides the runtime event plane, as the runtime scheduler's + // did; the partition publishes and consumes through host channels. + let replica_sync: Option = + if args.kv_router_config.router_replica_sync { + let channels = crate::kv_router::sequence::host_replica_channels( + &args.endpoint, + args.router_id, + cancellation_token.child_token(), + ) + .await + .context("start replica sync for the embedded selection partition")?; + let slot = std::sync::Mutex::new(Some(channels)); + Some(Arc::new(move |_partition| { + slot.lock().ok().and_then(|mut s| s.take()) + })) + } else { + None + }; let service = SelectionServiceBuilder::new( args.kv_router_config.clone(), @@ -136,6 +153,7 @@ impl EmbeddedSelection { None => OverlapRefreshSource::Disabled, }, scheduler_load_sink: Some(Arc::new(SenderLoadSink(args.scheduler_load))), + replica_sync, }) .build() .await diff --git a/lib/llm/src/kv_router/sequence.rs b/lib/llm/src/kv_router/sequence.rs index 977746b78c2a..9e8295503778 100644 --- a/lib/llm/src/kv_router/sequence.rs +++ b/lib/llm/src/kv_router/sequence.rs @@ -516,6 +516,70 @@ impl SequenceSubscriber for RuntimeSequenceSubscriber { /// Type alias for the runtime-wired multi-worker sequence tracker. pub type ActiveSequencesMulti = ActiveSequencesMultiWorker; +/// Replica-sync channels for an embedded selection partition, carried over the +/// runtime event plane exactly like the runtime scheduler's replica sync: +/// outbound events are published on `ACTIVE_SEQUENCES_SUBJECT`, peer events +/// are subscribed from it and forwarded to the partition. +pub(crate) async fn host_replica_channels( + endpoint: &Endpoint, + router_id: u64, + cancellation_token: CancellationToken, +) -> Result { + let transport_kind = endpoint.drt().default_event_transport_kind(); + let (outbound, outbound_rx) = mpsc::channel(REPLICA_EVENT_CHANNEL_CAPACITY); + let event_publisher = EventPublisher::for_endpoint_with_transport( + endpoint, + ACTIVE_SEQUENCES_SUBJECT, + transport_kind, + ) + .await?; + match active_sequence_event_wire_format(transport_kind) { + ActiveSequenceEventWireFormat::Singleton => { + tokio::spawn(run_replica_singleton_publisher( + event_publisher, + outbound_rx, + cancellation_token.clone(), + )); + } + ActiveSequenceEventWireFormat::Batch => { + tokio::spawn(run_replica_batch_publisher( + event_publisher, + outbound_rx, + cancellation_token.clone(), + )); + } + } + + let (inbound_tx, inbound_rx) = mpsc::channel(REPLICA_EVENT_CHANNEL_CAPACITY); + let mut subscriber = RuntimeSequenceSubscriber::for_endpoint(endpoint).await?; + let forward_tx = inbound_tx.clone(); + tokio::spawn(async move { + loop { + let next = tokio::select! { + _ = cancellation_token.cancelled() => break, + next = subscriber.next_event() => next, + }; + match next { + Some(Ok(event)) => { + if forward_tx.send(event).await.is_err() { + break; + } + } + Some(Err(error)) => { + tracing::warn!(%error, "replica-sync subscriber error; continuing"); + } + None => break, + } + } + }); + Ok(dynamo_kv_router::services::selection::HostReplicaChannels { + outbound, + inbound_tx, + inbound_rx, + process_id: router_id, + }) +} + /// Convenience async constructor that creates the event-plane publishers/subscribers /// and returns an `Arc` with replica sync already running. pub async fn create_multi_worker_sequences( From f417e1cf5097264d1e02ed77922d9f86da8c19f1 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Thu, 3 Sep 2026 14:24:13 +0000 Subject: [PATCH 010/102] feat(router): resolve linked worker-selection policies on the embedded partition Install the process-wide WorkerSelectionPolicyRegistry (the linked custom policy catalog the Python bindings register under --features custom-policy) through dynamo_llm::kv_router::install_worker_selection_policy_registry; embedded partitions and the static frontend resolve KvRouterConfig policy instances against it instead of the built-in default registry. A KvRouter built with Sel = WorkerSelectionPolicy therefore runs on the embedded partition too; only an arbitrary caller-injected WorkerSelector type (unit test fakes) still selects the runtime scheduler. Signed-off-by: Ishan Dhanani (cherry picked from commit fc22bac2ca2f1e4df0fcfdbfbe249e7372027442) --- lib/bindings/python/rust/lib.rs | 2 ++ lib/llm/src/kv_router.rs | 18 +++++++++++----- lib/llm/src/kv_router/embedded.rs | 34 ++++++++++++++++++++++--------- 3 files changed, 39 insertions(+), 15 deletions(-) diff --git a/lib/bindings/python/rust/lib.rs b/lib/bindings/python/rust/lib.rs index e2b66152d130..9def2af65cf0 100644 --- a/lib/bindings/python/rust/lib.rs +++ b/lib/bindings/python/rust/lib.rs @@ -335,6 +335,8 @@ fn register_core_with_custom_worker_selection_policy(m: &Bound<'_, PyModule>) -> dynamo_worker_selection_policy_catalog::register(&mut registry) .map_err(|error| PyRuntimeError::new_err(error.to_string()))?; + // Embedded selection partitions resolve linked policies from the same registry. + dynamo_llm::kv_router::install_worker_selection_policy_registry(registry.clone()); WORKER_SELECTION_POLICY_REGISTRY .set(registry) .map_err(|_| { diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index 2607e539ead7..349a43550a5c 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -73,6 +73,7 @@ pub mod shared_cache; pub use dynamo_kv_router::scheduling::{ OverlapScoresResponse, SharedCacheOverlapScore, WorkerOverlapScore, }; +pub use embedded::{install_worker_selection_policy_registry, worker_selection_policy_registry}; pub use encoder_router::EncoderRouter; pub use indexer::{Indexer, ServedIndexerHandle, ServedIndexerMode, ensure_served_indexer_service}; pub use prefill_router::PrefillRouter; @@ -807,11 +808,18 @@ where let available_worker_provider: WorkerAvailabilityProvider = Arc::new(move || client_for_availability.available_instance_ids()); - // The embedded partition builds its own selection policy (the default - // scorer/picker or a registry-resolved policy); a caller-injected custom - // `WorkerSelector` type can only run on the runtime scheduler. - let custom_selector = std::any::TypeId::of::() - != std::any::TypeId::of::(); + // The embedded partition builds its own selection policy: the default + // scorer/picker, or the registry-resolved policy the router config names + // (the same resolution `WorkerSelectionPolicy` callers perform). Only an + // arbitrary caller-injected `WorkerSelector` type still needs the runtime + // scheduler. + let sel = std::any::TypeId::of::(); + let custom_selector = sel + != std::any::TypeId::of::() + && sel + != std::any::TypeId::of::< + dynamo_kv_router::scheduling::selector::WorkerSelectionPolicy, + >(); if kv_router_config.router_embedded_selection && custom_selector { tracing::warn!( selector = std::any::type_name::(), diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs index 9d9869489ff2..7fa89602657a 100644 --- a/lib/llm/src/kv_router/embedded.rs +++ b/lib/llm/src/kv_router/embedded.rs @@ -73,6 +73,21 @@ pub(crate) struct EmbeddedSelection { worker_type: &'static str, } +static INSTALLED_POLICY_REGISTRY: std::sync::OnceLock = + std::sync::OnceLock::new(); + +/// Install the process-wide worker-selection policy registry (linked custom +/// policies) that embedded selection partitions resolve `KvRouterConfig` +/// policy instances against. Returns `false` if one is already installed. +pub fn install_worker_selection_policy_registry(registry: WorkerSelectionPolicyRegistry) -> bool { + INSTALLED_POLICY_REGISTRY.set(registry).is_ok() +} + +/// The installed registry, or the built-in default. +pub fn worker_selection_policy_registry() -> WorkerSelectionPolicyRegistry { + INSTALLED_POLICY_REGISTRY.get().cloned().unwrap_or_default() +} + /// Bridges the partition's scheduler load snapshots to the router's /// `SchedulerLoadSender`, which feeds `KvWorkerMonitor`. struct SenderLoadSink(crate::kv_router::routing_load::SchedulerLoadSender); @@ -105,15 +120,14 @@ impl EmbeddedSelection { // when router configuration names one, else Dynamo's default scorer // and picker under the router's metric worker label. let label = args.metric_worker_type; - let policy_factory: WorkerSelectionPolicyFactory = - match WorkerSelectionPolicyRegistry::default() - .resolve_for_worker_type(&args.kv_router_config, worker_type)? - { - Some(factory) => factory, - None => Arc::new(move |config: &KvRouterConfig, _worker_type, _partition| { - WorkerSelectionPolicy::default(config.clone(), label) - }), - }; + let policy_factory: WorkerSelectionPolicyFactory = match worker_selection_policy_registry() + .resolve_for_worker_type(&args.kv_router_config, worker_type)? + { + Some(factory) => factory, + None => Arc::new(move |config: &KvRouterConfig, _worker_type, _partition| { + WorkerSelectionPolicy::default(config.clone(), label) + }), + }; // Replica sync rides the runtime event plane, as the runtime scheduler's // did; the partition publishes and consumes through host channels. @@ -137,7 +151,7 @@ impl EmbeddedSelection { let service = SelectionServiceBuilder::new( args.kv_router_config.clone(), worker_type, - WorkerSelectionPolicyRegistry::default(), + worker_selection_policy_registry(), ) .worker_selection_policy_factory(policy_factory) .indexer_threads(1) From b2644f824d7d81b8928b81bdcb1b81551df065e6 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Wed, 9 Sep 2026 11:07:37 -0700 Subject: [PATCH 011/102] fix(router): project current KV transfer metadata into embedded selection Signed-off-by: Ishan Dhanani --- lib/llm/src/kv_router/embedded.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs index 7fa89602657a..abac908a10ba 100644 --- a/lib/llm/src/kv_router/embedded.rs +++ b/lib/llm/src/kv_router/embedded.rs @@ -410,7 +410,7 @@ pub(crate) fn worker_request_from_runtime_config( let mut router_hint_worker_type = None; let mut router_hint_source_control_endpoints = HashMap::new(); for dp_rank in dp_start..dp_start.saturating_add(dp_size) { - if let Some(metadata) = config.router_hint_metadata_for_dp_rank(dp_rank) { + if let Some(metadata) = config.kv_hint_transfer_metadata_for_dp_rank(dp_rank) { router_hint_worker_type.get_or_insert_with(|| metadata.worker_type.to_string()); if let Some(endpoint) = metadata.source_control_endpoint { router_hint_source_control_endpoints.insert(dp_rank, endpoint.to_string()); @@ -456,15 +456,15 @@ mod tests { }; config.taints.insert("gpu=h100".to_string()); config.runtime_data.insert( - dynamo_kv_router::router_hint::ROUTER_HINT_RUNTIME_CAPABILITY_KEY.to_string(), + dynamo_kv_router::kv_hints::KV_HINT_TRANSFER_CAPABILITY_KEY.to_string(), serde_json::Value::Bool(true), ); config.runtime_data.insert( - dynamo_kv_router::router_hint::ROUTER_HINT_WORKER_TYPE_RUNTIME_KEY.to_string(), + dynamo_kv_router::kv_hints::KV_HINT_TRANSFER_WORKER_TYPE_RUNTIME_KEY.to_string(), serde_json::Value::String("decode".to_string()), ); config.runtime_data.insert( - dynamo_kv_router::router_hint::ROUTER_HINT_SOURCE_CONTROL_ENDPOINTS_RUNTIME_KEY + dynamo_kv_router::kv_hints::KV_HINT_TRANSFER_SOURCE_CONTROL_ENDPOINTS_RUNTIME_KEY .to_string(), serde_json::json!({"2": "tcp://w:9002", "3": "tcp://w:9003"}), ); From 14aad43ab60596e51c3d6696371ea6d9fb512597 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Thu, 3 Sep 2026 15:03:08 +0000 Subject: [PATCH 012/102] refactor(router): delete the runtime-bound KV router scheduler The frontend router now always schedules on an embedded selection-service partition. Remove `KvScheduler` (kv_router/scheduler.rs keeps only its runtime-free re-exports), the runtime sequence publisher and its constructors, the direct-ZMQ active-sequence ingress, the two-variant `RouterScheduler` dispatch, and the LoRA load estimator's uncalled scheduler polling. `KvRouter` keeps its type parameter for API compatibility (`PhantomData Sel>`). A caller-injected `WorkerSelector` that is neither the default selector nor a registry `WorkerSelectionPolicy` now runs inside the partition: `WorkerSelectionPolicy::delegated` (new `DelegatedWorkerSelector` trait in dynamo-kv-router) hands the selector the scheduling request and a config-erased view of the eligible workers (`ErasedWorkerConfig`), which the embedded host projects back onto `ModelRuntimeConfig`. The unit tests that inject fake selectors run unchanged on that path. The per-worker `router_worker_registered` gauge is now maintained by the embedded catalog feeder. `router_embedded_selection` is deprecated and ignored; `false` logs a warning. Four sequence-level tests that built the runtime sequences directly are removed; their behavior is covered by the replica-lifecycle and session-affinity e2e tests on the embedded path. Signed-off-by: Ishan Dhanani (cherry picked from commit 71cea67bf3d76c4b6dce4fc8d5b26cbafe9bf77e) --- .../configuration/groups/kv_router_args.py | 6 +- .../components/frontend-configuration.mdx | 9 +- lib/bindings/python/src/dynamo/_core.pyi | 5 +- lib/kv-router/src/scheduling/config.rs | 7 +- lib/kv-router/src/scheduling/selector/mod.rs | 6 +- .../src/scheduling/selector/policy.rs | 122 ++- lib/llm/src/kv_router.rs | 226 ++--- lib/llm/src/kv_router/embedded.rs | 159 +++- lib/llm/src/kv_router/metrics.rs | 120 --- lib/llm/src/kv_router/routing_host/tests.rs | 6 +- lib/llm/src/kv_router/routing_load.rs | 11 +- lib/llm/src/kv_router/scheduler.rs | 708 +--------------- lib/llm/src/kv_router/sequence.rs | 774 +----------------- lib/llm/src/kv_router/sequence/direct_zmq.rs | 290 ------- lib/llm/src/lora/load_estimator.rs | 29 +- 15 files changed, 381 insertions(+), 2097 deletions(-) delete mode 100644 lib/llm/src/kv_router/sequence/direct_zmq.rs diff --git a/components/src/dynamo/common/configuration/groups/kv_router_args.py b/components/src/dynamo/common/configuration/groups/kv_router_args.py index 550867216884..011db193a48c 100644 --- a/components/src/dynamo/common/configuration/groups/kv_router_args.py +++ b/components/src/dynamo/common/configuration/groups/kv_router_args.py @@ -432,10 +432,8 @@ def add_arguments(self, parser) -> None: default=True, dest="router_embedded_selection", help=( - "KV Router: Run scheduling and active-sequence accounting on an " - "embedded selection-service partition (default). " - "--no-router-embedded-selection selects the deprecated runtime-bound " - "scheduler, kept for one release as a rollback." + "Deprecated and ignored: the KV Router always schedules on an embedded " + "selection-service partition. --no-router-embedded-selection only logs a warning." ), ) add_negatable_bool_argument( diff --git a/docs/fern/pages/reference/components/frontend-configuration.mdx b/docs/fern/pages/reference/components/frontend-configuration.mdx index 0ad51fa491b6..8899b3222ec7 100644 --- a/docs/fern/pages/reference/components/frontend-configuration.mdx +++ b/docs/fern/pages/reference/components/frontend-configuration.mdx @@ -214,11 +214,10 @@ router. See [Router Guide](../../developer-guide/knowledge-base/modular-componen - Run the router's scheduling and active-sequence accounting on an embedded selection-service - partition (the same core as the standalone selection service). The KV indexer, request - transport, and worker discovery are unchanged; `--router-replica-sync` is not applied to the - embedded partition. `--no-router-embedded-selection` selects the runtime-bound scheduler, which - is deprecated and kept for one release as a rollback; it logs a warning at startup. + Deprecated and ignored. The router's scheduling and active-sequence accounting always run on an + embedded selection-service partition (the same core as the standalone selection service); the + runtime-bound scheduler was removed. `--router-replica-sync` and overload detection apply to + the embedded partition. `--no-router-embedded-selection` only logs a warning. Environment variable: `DYN_ROUTER_EMBEDDED_SELECTION` diff --git a/lib/bindings/python/src/dynamo/_core.pyi b/lib/bindings/python/src/dynamo/_core.pyi index b7a2a962a7d7..ffeb381961b4 100644 --- a/lib/bindings/python/src/dynamo/_core.pyi +++ b/lib/bindings/python/src/dynamo/_core.pyi @@ -1830,9 +1830,8 @@ class KvRouterConfig: router_temperature: Temperature for normalized worker sampling via softmax (default: 0.0) use_kv_events: Whether to use KV events from workers (default: True) router_replica_sync: Enable replica synchronization (default: False) - router_embedded_selection: Run scheduling on an embedded selection-service - partition (default: True). False selects the deprecated runtime-bound - scheduler, kept for one release as a rollback. + router_embedded_selection: Deprecated and ignored; scheduling always runs on + an embedded selection-service partition. False only logs a warning. router_track_active_blocks: Track active blocks for load balancing (default: True) router_track_output_blocks: Track output blocks during generation (default: False). When enabled, the router adds placeholder blocks as tokens are generated diff --git a/lib/kv-router/src/scheduling/config.rs b/lib/kv-router/src/scheduling/config.rs index 4f6ffe4a34e9..590368968eee 100644 --- a/lib/kv-router/src/scheduling/config.rs +++ b/lib/kv-router/src/scheduling/config.rs @@ -811,10 +811,9 @@ pub struct KvRouterConfig { #[serde(default = "default_track_prefill_tokens")] pub router_track_prefill_tokens: bool, - /// Run the frontend router's scheduling and active-sequence accounting on - /// an embedded selection-service partition (default: true). `false` - /// selects the deprecated runtime-bound scheduler, kept for one release as - /// a rollback. Frontend-local; workers ignore it. + /// Deprecated and ignored: the frontend router always schedules on an + /// embedded selection-service partition. Kept so existing configurations + /// and wire records still parse; `false` logs a warning. #[serde(default = "default_true", skip_serializing_if = "is_true")] pub router_embedded_selection: bool, diff --git a/lib/kv-router/src/scheduling/selector/mod.rs b/lib/kv-router/src/scheduling/selector/mod.rs index 1b81bc78cb74..30d28398a76e 100644 --- a/lib/kv-router/src/scheduling/selector/mod.rs +++ b/lib/kv-router/src/scheduling/selector/mod.rs @@ -10,9 +10,9 @@ pub use default::DefaultWorkerSelector; use default::{DefaultWorkerPicker, DefaultWorkerScorer}; pub use policy::{ - ScoredWorkerCandidate, WorkerCacheInput, WorkerCandidate, WorkerFilter, WorkerInputView, - WorkerInputs, WorkerLoadInput, WorkerPicker, WorkerScorer, WorkerSelectionContext, - WorkerSelectionPolicy, + DelegatedWorkerSelector, ErasedWorkerConfig, ScoredWorkerCandidate, WorkerCacheInput, + WorkerCandidate, WorkerFilter, WorkerInputView, WorkerInputs, WorkerLoadInput, WorkerPicker, + WorkerScorer, WorkerSelectionContext, WorkerSelectionPolicy, }; use default::{pick_default_worker, selection_weights}; diff --git a/lib/kv-router/src/scheduling/selector/policy.rs b/lib/kv-router/src/scheduling/selector/policy.rs index 78a5cbf6e2ab..0dcd5f5ca7dd 100644 --- a/lib/kv-router/src/scheduling/selector/policy.rs +++ b/lib/kv-router/src/scheduling/selector/policy.rs @@ -344,6 +344,99 @@ pub(super) enum WorkerSelectionPolicyState { Default(DefaultWorkerPicker), /// Policy-local state owned and called serially by one scheduler queue actor. Custom(RefCell), + /// A caller-owned selector that sees the request and a config-erased view of + /// the eligible workers. Lets a host run its own `WorkerSelector` on a + /// partition whose worker-config type it does not know. + Delegated(Box), +} + +/// Worker configuration reduced to the [`WorkerConfigLike`] surface, so a +/// selector written against one host's config type can run on any partition. +#[derive(Debug, Clone, Default, PartialEq)] +pub struct ErasedWorkerConfig { + pub data_parallel_start_rank: u32, + pub data_parallel_size: u32, + pub max_num_batched_tokens: Option, + pub total_kv_blocks: Option, + pub native_offloading_capacity_tokens: Option, + pub taints: std::collections::HashSet, + pub stable_routing_id: Option, + pub topology_domains: Option>, + pub kv_transfer_domain: Option, + pub kv_transfer_enforcement: Option, + pub kv_transfer_preferred_weight: Option, +} + +impl ErasedWorkerConfig { + pub fn from_config_like(config: &C) -> Self { + Self { + data_parallel_start_rank: config.data_parallel_start_rank(), + data_parallel_size: config.data_parallel_size(), + max_num_batched_tokens: config.max_num_batched_tokens(), + total_kv_blocks: config.total_kv_blocks(), + native_offloading_capacity_tokens: config.native_offloading_capacity_tokens(), + taints: config.taints().clone(), + stable_routing_id: config.stable_routing_id().map(str::to_owned), + topology_domains: config.topology_domains().cloned(), + kv_transfer_domain: config.kv_transfer_domain().map(str::to_owned), + kv_transfer_enforcement: config.kv_transfer_enforcement(), + kv_transfer_preferred_weight: config.kv_transfer_preferred_weight(), + } + } +} + +impl WorkerConfigLike for ErasedWorkerConfig { + fn data_parallel_start_rank(&self) -> u32 { + self.data_parallel_start_rank + } + fn data_parallel_size(&self) -> u32 { + self.data_parallel_size + } + fn max_num_batched_tokens(&self) -> Option { + self.max_num_batched_tokens + } + fn total_kv_blocks(&self) -> Option { + self.total_kv_blocks + } + fn native_offloading_capacity_tokens(&self) -> Option { + self.native_offloading_capacity_tokens + } + fn taints(&self) -> &std::collections::HashSet { + &self.taints + } + fn stable_routing_id(&self) -> Option<&str> { + self.stable_routing_id.as_deref() + } + fn topology_domains(&self) -> Option<&HashMap> { + self.topology_domains.as_ref() + } + fn kv_transfer_domain(&self) -> Option<&str> { + self.kv_transfer_domain.as_deref() + } + fn kv_transfer_enforcement(&self) -> Option { + self.kv_transfer_enforcement + } + fn kv_transfer_preferred_weight(&self) -> Option { + self.kv_transfer_preferred_weight + } +} + +/// A caller-owned selector a partition policy delegates to. Receives the +/// scheduling request and a config-erased view of the eligible workers. +pub trait DelegatedWorkerSelector: Send + Sync { + fn required_worker_inputs(&self) -> WorkerInputs; + + fn uses_exclusive_affinity_target(&self) -> bool { + false + } + + fn select_worker( + &self, + workers: &HashMap, + request: &SchedulingRequest, + eligibility: RoutingEligibility<'_>, + block_size: u32, + ) -> Result; } pub(super) enum WorkerSelectionPolicyStateRef<'a> { @@ -387,6 +480,19 @@ impl WorkerSelectionPolicy { Self::new_with_filters(kv_router_config, worker_label, Vec::new(), scorers, picker) } + /// Run `selector` for every selection on this partition. + pub fn delegated( + kv_router_config: KvRouterConfig, + worker_label: &'static str, + selector: Box, + ) -> Self { + Self { + kv_router_config, + worker_label, + state: WorkerSelectionPolicyState::Delegated(selector), + } + } + /// Build a custom policy from ordered filters, additive scorers, and one picker. /// /// `worker_label` identifies the worker pool in routing logs. A typed policy factory normally @@ -590,7 +696,13 @@ pub(super) fn collect_custom_candidates( impl WorkerSelector for WorkerSelectionPolicy { fn uses_exclusive_affinity_target(&self) -> bool { - matches!(&self.state, WorkerSelectionPolicyState::Default(_)) + match &self.state { + WorkerSelectionPolicyState::Default(_) => true, + WorkerSelectionPolicyState::Custom(_) => false, + WorkerSelectionPolicyState::Delegated(selector) => { + selector.uses_exclusive_affinity_target() + } + } } fn required_worker_inputs(&self) -> WorkerInputs { @@ -600,6 +712,7 @@ impl WorkerSelector for WorkerSelectionPolicy { let state = state.borrow(); state.filter_inputs | state.scorer_picker_inputs } + WorkerSelectionPolicyState::Delegated(selector) => selector.required_worker_inputs(), } } @@ -616,6 +729,13 @@ impl WorkerSelector for WorkerSelectionPolicy { WorkerSelectionPolicyState::Custom(state) => { WorkerSelectionPolicyStateRef::Custom(state) } + WorkerSelectionPolicyState::Delegated(selector) => { + let erased: HashMap = workers + .iter() + .map(|(id, config)| (*id, ErasedWorkerConfig::from_config_like(config))) + .collect(); + return selector.select_worker(&erased, request, eligibility, block_size); + } }; select_worker_with_policy( &self.kv_router_config, diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index 349a43550a5c..225c51a4b4f6 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -28,8 +28,8 @@ use dynamo_kv_router::{ }, scheduling::{ AdmissionAttempt, AttemptId, CacheHitEstimates, OverlapAnalysis, OverloadedWorkerProvider, - ScheduleMode, ScheduleRequest, TieredOverlapRefresher, WorkerAvailabilityProvider, - effective_prefill_tokens, overlap::cache_hit_estimates_from_tiered_matches, + ScheduleMode, ScheduleRequest, WorkerAvailabilityProvider, effective_prefill_tokens, + overlap::cache_hit_estimates_from_tiered_matches, }, selector::WorkerInputs, }; @@ -85,7 +85,7 @@ pub use routing_load::{ use crate::{ discovery::{KvSourceMembershipWatch, RuntimeConfigWatch}, kv_router::{ - scheduler::{DefaultWorkerSelector, KvScheduler, PotentialLoad}, + scheduler::{DefaultWorkerSelector, PotentialLoad}, sequence::{SequenceError, SequenceRequest}, }, local_model::runtime_config::ModelRuntimeConfig, @@ -536,7 +536,10 @@ where Sel: dynamo_kv_router::selector::WorkerSelector, { indexer: Indexer, - scheduler: RouterScheduler, + scheduler: embedded::EmbeddedSelection, + /// `Sel` only names the selector type the caller injected; the partition + /// owns it. `fn() -> Sel` keeps the router `Send + Sync` regardless of `Sel`. + _selector: std::marker::PhantomData Sel>, required_worker_inputs: dynamo_kv_router::selector::WorkerInputs, workers_with_configs: RuntimeConfigWatch, block_size: u32, @@ -562,39 +565,6 @@ where teardown_task_guard: Option, } -/// The scheduler behind a `KvRouter`: the runtime-bound `KvScheduler`, or one -/// partition of an embedded selection service -/// (`KvRouterConfig::router_embedded_selection`). Both expose the same -/// surface; every router method dispatches statically through this enum. -enum RouterScheduler -where - Sel: dynamo_kv_router::selector::WorkerSelector, -{ - Runtime(KvScheduler>), - Embedded(embedded::EmbeddedSelection), -} - -macro_rules! dispatch_scheduler { - ($self:expr, $scheduler:ident => $body:expr) => { - match &$self.scheduler { - RouterScheduler::Runtime($scheduler) => $body, - RouterScheduler::Embedded($scheduler) => $body, - } - }; -} - -impl RouterScheduler -where - Sel: dynamo_kv_router::selector::WorkerSelector + Send + 'static, -{ - fn booking_cleanup(&self) -> scheduler::SchedulerBookingCleanup { - match self { - Self::Runtime(scheduler) => scheduler.booking_cleanup(), - Self::Embedded(scheduler) => scheduler.booking_cleanup(), - } - } -} - fn resolve_tracking_model_name( algorithm: TrackingHashAlgorithm, model_name: Option<&str>, @@ -793,13 +763,6 @@ where ); } - let overlap_scores_refresh = indexer.supports_overlap_refresh().then(|| { - Arc::new(TieredOverlapRefresher::new( - indexer.clone(), - kv_router_config.clone(), - block_size, - )) - }); let client_for_overload = client.clone(); let overloaded_worker_provider: OverloadedWorkerProvider = Arc::new(move || client_for_overload.overloaded_instance_ids()); @@ -808,11 +771,16 @@ where let available_worker_provider: WorkerAvailabilityProvider = Arc::new(move || client_for_availability.available_instance_ids()); - // The embedded partition builds its own selection policy: the default - // scorer/picker, or the registry-resolved policy the router config names - // (the same resolution `WorkerSelectionPolicy` callers perform). Only an - // arbitrary caller-injected `WorkerSelector` type still needs the runtime - // scheduler. + if !kv_router_config.router_embedded_selection { + tracing::warn!( + "router_embedded_selection=false is ignored: the runtime-bound KV router scheduler \ + was removed and scheduling always runs on the embedded selection partition" + ); + } + // Scheduling and active-sequence accounting run on a selection-service + // partition. A caller-injected `WorkerSelector` type (anything other than + // the default selector or a registry `WorkerSelectionPolicy`) runs inside + // the partition through a delegating policy. let sel = std::any::TypeId::of::(); let custom_selector = sel != std::any::TypeId::of::() @@ -820,79 +788,37 @@ where != std::any::TypeId::of::< dynamo_kv_router::scheduling::selector::WorkerSelectionPolicy, >(); - if kv_router_config.router_embedded_selection && custom_selector { - tracing::warn!( - selector = std::any::type_name::(), - "custom WorkerSelector is not supported on the embedded selection partition; using the runtime scheduler" - ); - } - let scheduler = if kv_router_config.router_embedded_selection && !custom_selector { - // The router keeps its indexer and transport; scheduling and - // active-sequence accounting run on a selection-service partition. - let overlap_refresh: Option> = - indexer - .supports_overlap_refresh() - .then(|| Arc::new(indexer.clone()) as Arc<_>); - drop(overlap_scores_refresh); - RouterScheduler::Embedded( - embedded::EmbeddedSelection::start( - embedded::EmbeddedSelectionArgs { - kv_router_config: kv_router_config.clone(), - worker_role, - metric_worker_type, - model_name: model_name.clone(), - block_size, - is_eagle, - prefill_load_estimator: prefill_load_estimator.clone(), - overloaded_worker_provider, - available_worker_provider, - overlap_refresh, - scheduler_load, - endpoint: endpoint.clone(), - router_id: endpoint.drt().discovery().instance_id(), - }, - workers_with_configs.clone(), - cancellation_token.child_token(), - ) - .await?, - ) - } else { - tracing::warn!( - "the runtime-bound KV router scheduler (--no-router-embedded-selection) is \ - deprecated and kept for one release as a rollback; the embedded selection \ - partition is the default" - ); - RouterScheduler::Runtime( - KvScheduler::start_with_shared_request_leases( - endpoint.clone(), - block_size, - workers_with_configs.clone(), - selector, - &kv_router_config, - prefill_load_estimator.clone(), - overlap_scores_refresh, - Some(overloaded_worker_provider), - Some(available_worker_provider), - model_name.as_deref(), - metric_worker_type, - scheduler_load, - cancellation_token.child_token(), - ) - .await?, - ) - }; + let policy_override = custom_selector + .then(|| embedded::delegating_policy_factory(selector, metric_worker_type)); + let overlap_refresh: Option> = + indexer + .supports_overlap_refresh() + .then(|| Arc::new(indexer.clone()) as Arc<_>); + let scheduler = embedded::EmbeddedSelection::start( + embedded::EmbeddedSelectionArgs { + kv_router_config: kv_router_config.clone(), + worker_role, + metric_worker_type, + model_name: model_name.clone(), + block_size, + is_eagle, + prefill_load_estimator: prefill_load_estimator.clone(), + overloaded_worker_provider, + available_worker_provider, + overlap_refresh, + scheduler_load, + endpoint: endpoint.clone(), + router_id: endpoint.drt().discovery().instance_id(), + policy_override, + }, + workers_with_configs.clone(), + cancellation_token.child_token(), + ) + .await?; let request_leases = request_lease::RequestLeaseManager::new( scheduler.booking_cleanup(), cancellation_token.child_token(), ); - if let RouterScheduler::Runtime(runtime_scheduler) = &scheduler - && !runtime_scheduler - .set_replica_request_lease_observer(Arc::new(request_leases.clone())) - { - return Err(anyhow::anyhow!( - "request lease observer is already installed for this router" - )); - } // Start KV event subscription if needed — skip when using a remote indexer. let kv_event_subscription = if cache_required && kv_event_source_requirement.should_subscribe(&kv_router_config) @@ -950,6 +876,7 @@ where Ok(Self { indexer, scheduler, + _selector: std::marker::PhantomData, required_worker_inputs, workers_with_configs, block_size, @@ -1695,8 +1622,7 @@ where let seq_hash_elapsed = start.elapsed(); let is_admitted_routing = matches!(admission, FindBestMatchAdmission::WithAdmission { .. }); - let supports_overlap_refresh = - dispatch_scheduler!(self, scheduler => scheduler.supports_overlap_refresh()); + let supports_overlap_refresh = self.scheduler.supports_overlap_refresh(); let retain_block_hashes = supports_overlap_refresh || return_routing_hashes; let has_transfer_capable_workers = self.has_transfer_capable_workers(); let should_prepare_transfer_hint = is_admitted_routing && has_transfer_capable_workers; @@ -1788,10 +1714,11 @@ where }; let (response, attempt, selected_worker_load) = match admission { FindBestMatchAdmission::WithAdmission { .. } => { - match dispatch_scheduler!(self, scheduler => scheduler - .schedule_request_admitted(schedule_request) - .instrument(tracing::info_span!("kv_router.schedule")) - .await) + match self + .scheduler + .schedule_request_admitted(schedule_request) + .instrument(tracing::info_span!("kv_router.schedule")) + .await { Ok(admitted) => (admitted.response, admitted.attempt, None), Err(KvSchedulerError::QueueRejected(rejection)) => { @@ -1806,10 +1733,11 @@ where } } FindBestMatchAdmission::WithoutAdmission => { - match dispatch_scheduler!(self, scheduler => scheduler - .select_without_admission(schedule_request) - .instrument(tracing::info_span!("kv_router.select_without_admission")) - .await) + match self + .scheduler + .select_without_admission(schedule_request) + .instrument(tracing::info_span!("kv_router.select_without_admission")) + .await { Ok(advisory) => ( advisory.response, @@ -1953,7 +1881,7 @@ where /// Register externally-provided workers in the slot tracker. pub fn register_workers(&self, worker_ids: &HashSet) { - dispatch_scheduler!(self, scheduler => scheduler.register_workers(worker_ids)); + self.scheduler.register_workers(worker_ids); } #[allow(clippy::too_many_arguments)] @@ -2002,9 +1930,7 @@ where worker, lora_name, }; - let admission = dispatch_scheduler!(self, scheduler => scheduler - .add_request_admitted(sequence_request) - .await); + let admission = self.scheduler.add_request_admitted(sequence_request).await; let attempt_id = match admission { Ok(attempt_id) => attempt_id, Err(error) => { @@ -2025,7 +1951,7 @@ where } pub async fn mark_prefill_completed(&self, request_id: &str) -> Result<(), SequenceError> { - dispatch_scheduler!(self, scheduler => scheduler.mark_prefill_completed(request_id).await)?; + self.scheduler.mark_prefill_completed(request_id).await?; self.request_leases.touch_request(request_id); Ok(()) } @@ -2034,7 +1960,7 @@ where if self.request_leases.finish_request(request_id).await { return Ok(()); } - dispatch_scheduler!(self, scheduler => scheduler.free(request_id).await) + self.scheduler.free(request_id).await } /// Release a booking only if it still belongs to `worker`. @@ -2046,7 +1972,7 @@ where request_id: &str, worker: WorkerWithDpRank, ) -> Result<(), SequenceError> { - dispatch_scheduler!(self, scheduler => scheduler.free_if_worker(request_id, worker).await) + self.scheduler.free_if_worker(request_id, worker).await } #[doc(hidden)] @@ -2062,19 +1988,19 @@ where &self, booking: &scheduler::SchedulerBookingDescriptor, ) -> Result<(), KvSchedulerError> { - dispatch_scheduler!(self, scheduler => scheduler + self.scheduler .mark_prefill_completed_if_booking(booking) - .await) + .await } /// Number of requests currently parked in the scheduler queue. pub fn pending_count(&self) -> usize { - dispatch_scheduler!(self, scheduler => scheduler.pending_count()) + self.scheduler.pending_count() } /// Sum of ISL tokens for requests currently parked in the scheduler queue. pub fn pending_isl_tokens(&self) -> usize { - dispatch_scheduler!(self, scheduler => scheduler.pending_isl_tokens()) + self.scheduler.pending_isl_tokens() } fn prefill_load_hint_for( @@ -2117,7 +2043,7 @@ where /// Get the worker type for this router ("prefill" or "decode"). /// Used for Prometheus metric labeling. pub fn worker_type(&self) -> &'static str { - dispatch_scheduler!(self, scheduler => scheduler.worker_type()) + self.scheduler.worker_type() } /// Return the worker's unique global DP rank when it owns exactly one rank. @@ -2132,7 +2058,7 @@ where request_id: &str, decay_fraction: Option, ) -> Result<(), SequenceError> { - dispatch_scheduler!(self, scheduler => scheduler.add_output_block(request_id, decay_fraction)) + self.scheduler.add_output_block(request_id, decay_fraction) } pub(crate) async fn enqueue_output_block_if_booking( @@ -2140,9 +2066,9 @@ where booking: &scheduler::SchedulerBookingDescriptor, decay_fraction: Option, ) -> Result<(), KvSchedulerError> { - dispatch_scheduler!(self, scheduler => scheduler + self.scheduler .enqueue_output_block_if_booking(booking, decay_fraction) - .await) + .await } pub fn block_size(&self) -> u32 { @@ -2256,14 +2182,12 @@ where let effective_cached_tokens: HashMap = cache_hit_estimates.cached_tokens.into_iter().collect(); - Ok( - dispatch_scheduler!(self, scheduler => scheduler.get_potential_loads( - maybe_seq_hashes, - isl_tokens, - effective_cached_tokens, - track_prefill_tokens, - )), - ) + Ok(self.scheduler.get_potential_loads( + maybe_seq_hashes, + isl_tokens, + effective_cached_tokens, + track_prefill_tokens, + )) } /// Return per-worker KV overlap by storage tier. diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs index abac908a10ba..993a7dca3466 100644 --- a/lib/llm/src/kv_router/embedded.rs +++ b/lib/llm/src/kv_router/embedded.rs @@ -62,6 +62,127 @@ pub(crate) struct EmbeddedSelectionArgs { pub endpoint: dynamo_runtime::component::Endpoint, /// This router replica's id (ignored on receipt of its own events). pub router_id: u64, + /// Use this policy instead of resolving one from the registry (a + /// caller-injected `WorkerSelector`, wrapped by [`delegating_policy_factory`]). + pub policy_override: Option, +} + +/// Runs a caller-owned `WorkerSelector` inside a partition +/// policy: the partition's worker configs are projected onto +/// `ModelRuntimeConfig` per selection. Serialized through a mutex because the +/// selector is only required to be `Send`. +struct DelegatingSelector { + inner: std::sync::Mutex, + inputs: dynamo_kv_router::scheduling::selector::WorkerInputs, + exclusive_affinity: bool, +} + +impl dynamo_kv_router::scheduling::selector::DelegatedWorkerSelector + for DelegatingSelector +where + Sel: dynamo_kv_router::selector::WorkerSelector + Send + 'static, +{ + fn required_worker_inputs(&self) -> dynamo_kv_router::scheduling::selector::WorkerInputs { + self.inputs + } + + fn uses_exclusive_affinity_target(&self) -> bool { + self.exclusive_affinity + } + + fn select_worker( + &self, + workers: &HashMap, + request: &dynamo_kv_router::scheduling::SchedulingRequest, + eligibility: dynamo_kv_router::scheduling::RoutingEligibility<'_>, + block_size: u32, + ) -> std::result::Result< + dynamo_kv_router::protocols::WorkerSelectionResult, + dynamo_kv_router::scheduling::KvSchedulerError, + > { + let projected: HashMap = workers + .iter() + .map(|(id, config)| { + ( + *id, + ModelRuntimeConfig { + data_parallel_start_rank: config.data_parallel_start_rank, + data_parallel_size: config.data_parallel_size, + max_num_batched_tokens: config.max_num_batched_tokens, + total_kv_blocks: config.total_kv_blocks, + taints: config.taints.clone(), + stable_routing_id: config.stable_routing_id.clone(), + topology_domains: config.topology_domains.clone().unwrap_or_default(), + kv_transfer_domain: config.kv_transfer_domain.clone(), + kv_transfer_enforcement: config.kv_transfer_enforcement, + kv_transfer_preferred_weight: config.kv_transfer_preferred_weight, + ..Default::default() + }, + ) + }) + .collect(); + let inner = self + .inner + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + inner.select_worker( + dynamo_kv_router::scheduling::selector::WorkerSelectionInput::Configured { + workers: &projected, + request, + eligibility, + block_size, + }, + ) + } +} + +struct SharedDelegate(Arc); + +impl dynamo_kv_router::scheduling::selector::DelegatedWorkerSelector for SharedDelegate { + fn required_worker_inputs(&self) -> dynamo_kv_router::scheduling::selector::WorkerInputs { + self.0.required_worker_inputs() + } + fn uses_exclusive_affinity_target(&self) -> bool { + self.0.uses_exclusive_affinity_target() + } + fn select_worker( + &self, + workers: &HashMap, + request: &dynamo_kv_router::scheduling::SchedulingRequest, + eligibility: dynamo_kv_router::scheduling::RoutingEligibility<'_>, + block_size: u32, + ) -> std::result::Result< + dynamo_kv_router::protocols::WorkerSelectionResult, + dynamo_kv_router::scheduling::KvSchedulerError, + > { + self.0 + .select_worker(workers, request, eligibility, block_size) + } +} + +/// A policy factory that runs `selector` on every partition of the router. +pub(crate) fn delegating_policy_factory( + selector: Sel, + label: &'static str, +) -> WorkerSelectionPolicyFactory +where + Sel: dynamo_kv_router::selector::WorkerSelector + Send + 'static, +{ + let inputs = selector.required_worker_inputs(); + let exclusive_affinity = selector.uses_exclusive_affinity_target(); + let shared: Arc = + Arc::new(DelegatingSelector { + inner: std::sync::Mutex::new(selector), + inputs, + exclusive_affinity, + }); + Arc::new(move |config: &KvRouterConfig, _worker_type, _partition| { + WorkerSelectionPolicy::delegated( + config.clone(), + label, + Box::new(SharedDelegate(Arc::clone(&shared))), + ) + }) } /// One `SelectionService` partition driven directly by the router. @@ -120,13 +241,16 @@ impl EmbeddedSelection { // when router configuration names one, else Dynamo's default scorer // and picker under the router's metric worker label. let label = args.metric_worker_type; - let policy_factory: WorkerSelectionPolicyFactory = match worker_selection_policy_registry() - .resolve_for_worker_type(&args.kv_router_config, worker_type)? - { + let policy_factory: WorkerSelectionPolicyFactory = match args.policy_override { Some(factory) => factory, - None => Arc::new(move |config: &KvRouterConfig, _worker_type, _partition| { - WorkerSelectionPolicy::default(config.clone(), label) - }), + None => match worker_selection_policy_registry() + .resolve_for_worker_type(&args.kv_router_config, worker_type)? + { + Some(factory) => factory, + None => Arc::new(move |config: &KvRouterConfig, _worker_type, _partition| { + WorkerSelectionPolicy::default(config.clone(), label) + }), + }, }; // Replica sync rides the runtime event plane, as the runtime scheduler's @@ -183,6 +307,10 @@ impl EmbeddedSelection { block_size: args.block_size, is_eagle: args.is_eagle, known: HashMap::new(), + status_metrics: super::metrics::RouterWorkerStatusMetrics::from_component( + args.endpoint.component(), + ), + worker_label: args.metric_worker_type, }; feeder .run(workers_with_configs, cancellation_token.child_token()) @@ -329,6 +457,14 @@ struct CatalogFeeder { block_size: u32, is_eagle: bool, known: HashMap, + /// `router_worker_registered` gauge per worker/dp_rank. + status_metrics: Arc, + worker_label: &'static str, +} + +fn dp_ranks(config: &ModelRuntimeConfig) -> std::ops::Range { + let start = config.data_parallel_start_rank; + start..start.saturating_add(config.data_parallel_size.max(1)) } impl CatalogFeeder { @@ -360,7 +496,12 @@ impl CatalogFeeder { .filter(|worker_id| !snapshot.contains_key(worker_id)) .collect(); for worker_id in removed { - self.known.remove(&worker_id); + if let Some(previous) = self.known.remove(&worker_id) { + for dp_rank in dp_ranks(&previous) { + self.status_metrics + .remove_worker(worker_id, dp_rank, self.worker_label); + } + } if let Err(error) = self.service.delete_worker(worker_id).await { tracing::debug!(worker_id, %error, "embedded selection: worker delete skipped"); } @@ -385,6 +526,10 @@ impl CatalogFeeder { "embedded selection: worker is not schedulable" ); } + for dp_rank in dp_ranks(&config) { + self.status_metrics + .set_registered(worker_id, dp_rank, self.worker_label); + } self.known.insert(worker_id, config); } Err(error) => { diff --git a/lib/llm/src/kv_router/metrics.rs b/lib/llm/src/kv_router/metrics.rs index cbce00c50b49..245f29a862bd 100644 --- a/lib/llm/src/kv_router/metrics.rs +++ b/lib/llm/src/kv_router/metrics.rs @@ -282,126 +282,6 @@ impl KvZmqIngressMetrics { } } -// --------------------------------------------------------------------------- -// Direct-ZMQ active-sequence ingress metrics -// --------------------------------------------------------------------------- - -pub(crate) struct ActiveSequenceZmqIngressMetrics { - tracked_sources: IntGauge, - sequence_gaps_total: IntCounter, - out_of_order_total: IntCounter, - reconnects_total: IntCounter, - replacements_total: IntCounter, - envelope_decode_errors_total: IntCounter, - payload_decode_errors_total: IntCounter, - identity_errors_total: IntCounter, - forced_aborts_total: IntCounter, -} - -static ACTIVE_SEQUENCE_ZMQ_INGRESS_METRICS: OnceLock> = - OnceLock::new(); - -impl ActiveSequenceZmqIngressMetrics { - pub(crate) fn from_component(component: &Component) -> Arc { - ACTIVE_SEQUENCE_ZMQ_INGRESS_METRICS - .get_or_init(|| { - let metrics = component.metrics(); - let counter = |name, help| { - metrics - .create_intcounter(name, help, &[]) - .unwrap_or_else(|error| panic!("failed to create {name}: {error}")) - }; - Arc::new(Self { - tracked_sources: metrics - .create_intgauge( - "router_active_sequence_zmq_ingress_sources", - "Number of tracked direct-ZMQ active-sequence sources", - &[], - ) - .expect("failed to create router_active_sequence_zmq_ingress_sources"), - sequence_gaps_total: counter( - "router_active_sequence_zmq_ingress_sequence_gaps_total", - "Missing direct-ZMQ active-sequence envelopes inferred from sequence gaps", - ), - out_of_order_total: counter( - "router_active_sequence_zmq_ingress_out_of_order_total", - "Non-increasing direct-ZMQ active-sequence envelope sequences", - ), - reconnects_total: counter( - "router_active_sequence_zmq_ingress_reconnects_total", - "Direct-ZMQ active-sequence source reconnect attempts", - ), - replacements_total: counter( - "router_active_sequence_zmq_ingress_replacements_total", - "Direct-ZMQ active-sequence source task replacements", - ), - envelope_decode_errors_total: counter( - "router_active_sequence_zmq_ingress_envelope_decode_errors_total", - "Direct-ZMQ active-sequence envelope decode errors", - ), - payload_decode_errors_total: counter( - "router_active_sequence_zmq_ingress_payload_decode_errors_total", - "Direct-ZMQ active-sequence payload decode errors", - ), - identity_errors_total: counter( - "router_active_sequence_zmq_ingress_identity_errors_total", - "Direct-ZMQ active-sequence frame and envelope identity mismatches", - ), - forced_aborts_total: counter( - "router_active_sequence_zmq_ingress_forced_aborts_total", - "Direct-ZMQ active-sequence source tasks aborted after join timeout", - ), - }) - }) - .clone() - } - - pub(crate) fn source_started(&self) { - self.tracked_sources.inc(); - } - - pub(crate) fn source_stopped(&self) { - self.tracked_sources.dec(); - } - - pub(crate) fn record_gap(&self, missing: u64) { - self.sequence_gaps_total.inc_by(missing); - } - - pub(crate) fn record_out_of_order(&self) { - self.out_of_order_total.inc(); - } - - pub(crate) fn record_reconnect(&self) { - self.reconnects_total.inc(); - } - - pub(crate) fn record_replacement(&self) { - self.replacements_total.inc(); - } - - pub(crate) fn record_envelope_decode_error(&self) { - self.envelope_decode_errors_total.inc(); - } - - pub(crate) fn record_payload_decode_error(&self) { - self.payload_decode_errors_total.inc(); - } - - pub(crate) fn record_identity_error(&self) { - self.identity_errors_total.inc(); - } - - pub(crate) fn record_forced_abort(&self) { - self.forced_aborts_total.inc(); - } -} - -// --------------------------------------------------------------------------- -// Router worker status metrics (component-scoped gauges) -// --------------------------------------------------------------------------- - -/// Component-scoped router gauges for worker discovery. pub(crate) struct RouterWorkerStatusMetrics { pub registered: IntGaugeVec, pub kv_event_source_mismatch_workers: IntGaugeVec, diff --git a/lib/llm/src/kv_router/routing_host/tests.rs b/lib/llm/src/kv_router/routing_host/tests.rs index 92537a90ae31..161c2f9b2e62 100644 --- a/lib/llm/src/kv_router/routing_host/tests.rs +++ b/lib/llm/src/kv_router/routing_host/tests.rs @@ -1171,10 +1171,12 @@ async fn embedded_router_with_workers( router_with_worker_configs_in_mode(session_affinity_ttl, workers, true).await } +/// `_embedded_selection` is kept for call-site stability: the embedded +/// partition is the only scheduler now. async fn router_with_worker_configs_in_mode( session_affinity_ttl: Option, workers: HashMap, - embedded_selection: bool, + _embedded_selection: bool, ) -> (RoutingHost, Runtime) { let runtime = Runtime::from_current().unwrap(); let distributed = DistributedRuntime::new(runtime.clone(), DistributedConfig::process_local()) @@ -1193,7 +1195,7 @@ async fn router_with_worker_configs_in_mode( skip_initial_worker_wait: true, use_kv_events: false, router_track_active_blocks: false, - router_embedded_selection: embedded_selection, + router_embedded_selection: true, ..Default::default() }; let chooser = KvRouter::new( diff --git a/lib/llm/src/kv_router/routing_load.rs b/lib/llm/src/kv_router/routing_load.rs index c6497aeb89e9..2c59c365f929 100644 --- a/lib/llm/src/kv_router/routing_load.rs +++ b/lib/llm/src/kv_router/routing_load.rs @@ -197,27 +197,21 @@ impl SchedulerLoadShared { pub struct SchedulerLoadSender { wake_tx: Option>, shared: Arc, - source: RouterLoadSource, cancellation_token: CancellationToken, } impl SchedulerLoadSender { pub(crate) fn disabled( - source: RouterLoadSource, + _source: RouterLoadSource, cancellation_token: CancellationToken, ) -> Self { Self { wake_tx: None, shared: Arc::new(SchedulerLoadShared::new(SCHEDULER_LOAD_CHANNEL_CAPACITY)), - source, cancellation_token, } } - pub(crate) const fn metric_label(&self) -> &'static str { - self.source.metric_label() - } - pub fn publish(&self, snapshot: SchedulerLoadSnapshot) { self.try_publish(SchedulerLoadCommand::Single(snapshot)); } @@ -291,7 +285,7 @@ pub(crate) fn scheduler_load_channel( } fn scheduler_load_channel_with_capacity( - source: RouterLoadSource, + _source: RouterLoadSource, cancellation_token: CancellationToken, capacity: usize, ) -> (SchedulerLoadSender, SchedulerLoadReceiver) { @@ -301,7 +295,6 @@ fn scheduler_load_channel_with_capacity( SchedulerLoadSender { wake_tx: Some(wake_tx), shared: shared.clone(), - source, cancellation_token, }, SchedulerLoadReceiver { wake_rx, shared }, diff --git a/lib/llm/src/kv_router/scheduler.rs b/lib/llm/src/kv_router/scheduler.rs index 838108af5dcf..97fe82fb85aa 100644 --- a/lib/llm/src/kv_router/scheduler.rs +++ b/lib/llm/src/kv_router/scheduler.rs @@ -1,7 +1,10 @@ // SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use dynamo_kv_router::protocols::{LocalBlockHash, SharedCacheHits}; +//! Scheduling types shared by the router. The router's scheduler is one +//! partition of an embedded selection service (`kv_router::embedded`); these +//! re-exports are the runtime-free scheduling vocabulary it speaks. + pub use dynamo_kv_router::scheduling::overlap_refresh::{ NoopOverlapScoresRefresh, OverlapScoresRefresh, RefreshedOverlap, }; @@ -15,706 +18,3 @@ pub use dynamo_kv_router::scheduling::{ WorkerAvailabilityProvider, }; pub use dynamo_kv_router::selector::DefaultWorkerSelector; -use dynamo_kv_router::selector::WorkerSelector as WorkerSelectorTrait; - -use super::SchedulerLoadSender; -use super::metrics::{ROUTER_QUEUE_METRICS, RouterQueueMetricHandles, RouterRequestMetrics}; -use super::sequence::{ - DeferredReplicaRequestLeaseObserver, RuntimeSequencePublisher, SequenceError, SequenceRequest, - create_multi_worker_sequences_with_observer, -}; -use crate::discovery::RuntimeConfigWatch; -use crate::local_model::runtime_config::ModelRuntimeConfig; -use crate::protocols::common::timing::WORKER_TYPE_PREFILL; -use anyhow::Result; -use dynamo_kv_router::{ - PrefillLoadEstimator, - config::{KvRouterConfig, RouterConfigOverride}, - multi_worker_sequence::ReplicaRequestLeaseObserver, - protocols::{RoutingConstraints, WorkerId, WorkerWithDpRank}, -}; -use dynamo_runtime::component::Endpoint; -use dynamo_runtime::traits::DistributedRuntimeProvider; -use dynamo_tokens::SequenceHash; -use std::collections::{HashMap, HashSet}; -use std::sync::Arc; -use std::time::Duration; -use tokio_util::sync::CancellationToken; - -pub struct KvScheduler -where - Sel: WorkerSelectorTrait, - RF: OverlapScoresRefresh, -{ - inner: Arc>, - replica_request_lease_observer: Option>, - queue_metrics: Vec, - queue_metric_indices: HashMap, -} - -impl KvScheduler -where - Sel: WorkerSelectorTrait + Send + 'static, - RF: OverlapScoresRefresh + Send + Sync + 'static, -{ - /// Start the scheduler, optionally wiring an [`OverlapScoresRefresh`] into the queue so - /// long-waiting requests can be re-scored at dequeue time. - #[expect(clippy::too_many_arguments)] - pub async fn start( - endpoint: Endpoint, - block_size: u32, - workers_with_configs: RuntimeConfigWatch, - selector: Sel, - kv_router_config: &KvRouterConfig, - prefill_load_estimator: Option>, - overlap_scores_refresh: Option>, - overloaded_worker_provider: Option, - available_worker_provider: Option, - model_name: Option<&str>, - worker_type: &'static str, - scheduler_load: SchedulerLoadSender, - cancellation_token: CancellationToken, - ) -> Result { - Self::start_inner( - endpoint, - block_size, - workers_with_configs, - selector, - kv_router_config, - prefill_load_estimator, - overlap_scores_refresh, - overloaded_worker_provider, - available_worker_provider, - model_name, - worker_type, - scheduler_load, - cancellation_token, - false, - ) - .await - } - - /// Start the embedded router scheduler whose request lifecycle is owned by - /// `RequestLeaseManager` rather than the standalone slot-tracker expiry. - #[expect(clippy::too_many_arguments)] - pub(crate) async fn start_with_shared_request_leases( - endpoint: Endpoint, - block_size: u32, - workers_with_configs: RuntimeConfigWatch, - selector: Sel, - kv_router_config: &KvRouterConfig, - prefill_load_estimator: Option>, - overlap_scores_refresh: Option>, - overloaded_worker_provider: Option, - available_worker_provider: Option, - model_name: Option<&str>, - worker_type: &'static str, - scheduler_load: SchedulerLoadSender, - cancellation_token: CancellationToken, - ) -> Result { - Self::start_inner( - endpoint, - block_size, - workers_with_configs, - selector, - kv_router_config, - prefill_load_estimator, - overlap_scores_refresh, - overloaded_worker_provider, - available_worker_provider, - model_name, - worker_type, - scheduler_load, - cancellation_token, - true, - ) - .await - } - - #[expect(clippy::too_many_arguments)] - async fn start_inner( - endpoint: Endpoint, - block_size: u32, - workers_with_configs: RuntimeConfigWatch, - selector: Sel, - kv_router_config: &KvRouterConfig, - prefill_load_estimator: Option>, - overlap_scores_refresh: Option>, - overloaded_worker_provider: Option, - available_worker_provider: Option, - model_name: Option<&str>, - worker_type: &'static str, - scheduler_load: SchedulerLoadSender, - cancellation_token: CancellationToken, - use_shared_request_leases: bool, - ) -> Result { - let initial_workers: HashMap = - workers_with_configs.borrow().clone(); - - let router_id = endpoint.drt().discovery().instance_id(); - let replica_request_lease_observer = use_shared_request_leases - .then(|| Arc::new(DeferredReplicaRequestLeaseObserver::default())); - let observer = replica_request_lease_observer - .as_ref() - .map(|observer| Arc::clone(observer) as Arc); - let slots = create_multi_worker_sequences_with_observer( - endpoint, - block_size as usize, - initial_workers, - kv_router_config.router_replica_sync, - router_id, - scheduler_load, - observer, - cancellation_token.child_token(), - ) - .await - .map_err(|e| KvSchedulerError::InitFailed(e.to_string()))?; - - // `skip_initial_worker_wait` only controls boot-time blocking. The scheduler must - // keep monitoring runtime config changes so late-joining workers enter routing. - let watch_worker_configs = true; - - let profile = kv_router_config - .policy_profile(model_name) - .map_err(|error| KvSchedulerError::InitFailed(error.to_string()))?; - let queue_recheck_interval = kv_router_config.router_queue_recheck_interval(); - let metric_model = model_name.unwrap_or("unknown"); - let queue_metrics = profile - .classes() - .iter() - .map(|class| ROUTER_QUEUE_METRICS.handles(metric_model, worker_type, &class.name)) - .collect::>(); - let queue_metric_indices = profile - .classes() - .iter() - .enumerate() - .map(|(index, class)| (class.name.clone(), index)) - .collect(); - - let inner = Arc::new(LocalScheduler::new( - slots, - workers_with_configs.clone(), - profile, - block_size, - selector, - prefill_load_estimator, - overlap_scores_refresh, - overloaded_worker_provider, - available_worker_provider, - queue_recheck_interval, - kv_router_config.router_track_prefill_tokens, - cancellation_token.child_token(), - worker_type, - watch_worker_configs, - )); - if worker_type == WORKER_TYPE_PREFILL { - let locality_observer: NonMaxOverlapSelectionObserver = - Arc::new(move |request_id, selection| { - let overlap_blocks_lost = selection.overlap_blocks_lost(); - if let Some(metrics) = RouterRequestMetrics::get() { - metrics.observe_non_max_overlap_selection(worker_type, overlap_blocks_lost); - } - tracing::debug!( - request_id, - worker_type, - selected_worker_id = selection.selected_worker.worker_id, - selected_dp_rank = selection.selected_worker.dp_rank, - selected_overlap_blocks = selection.selected_overlap_blocks, - highest_overlap_worker_id = selection.highest_overlap_worker.worker_id, - highest_overlap_dp_rank = selection.highest_overlap_worker.dp_rank, - highest_overlap_blocks = selection.highest_overlap_blocks, - overlap_blocks_lost, - "Router selected a worker with lower KV cache overlap" - ); - }); - if !inner.set_non_max_overlap_selection_observer(locality_observer) { - return Err(KvSchedulerError::InitFailed( - "non-max-overlap observer is already installed".to_string(), - )); - } - } - - let metrics_scheduler = Arc::clone(&inner); - let background_metrics = queue_metrics.clone(); - let metrics_cancel_token = cancellation_token.child_token(); - let mut queue_updates = inner.subscribe_queue_updates(); - tokio::spawn(async move { - let mut recheck_interval = tokio::time::interval(Duration::from_secs(60)); - update_queue_metrics(&background_metrics, |class_index| { - metrics_scheduler.class_queue_stats(class_index) - }); - - loop { - tokio::select! { - _ = metrics_cancel_token.cancelled() => break, - result = queue_updates.changed() => { - if result.is_err() { - break; - } - update_queue_metrics(&background_metrics, |class_index| { - metrics_scheduler.class_queue_stats(class_index) - }); - } - _ = recheck_interval.tick() => { - update_queue_metrics(&background_metrics, |class_index| { - metrics_scheduler.class_queue_stats(class_index) - }); - } - } - } - }); - - Ok(Self { - inner, - replica_request_lease_observer, - queue_metrics, - queue_metric_indices, - }) - } - - pub async fn schedule_request( - &self, - request: ScheduleRequest, - ) -> Result { - let response = self.inner.schedule_request(request).await; - self.observe_schedule_result(&response); - response - } - - pub(crate) async fn schedule_request_admitted( - &self, - request: ScheduleRequest, - ) -> Result { - let response = self.inner.schedule_request_admitted(request).await; - self.observe_schedule_result(&response); - response - } - - #[expect(clippy::too_many_arguments)] - pub async fn schedule( - &self, - maybe_request_id: Option, - isl_tokens: usize, - token_seq: Option>, - tier_overlap_blocks: TierOverlapBlocks, - effective_overlap_blocks: HashMap, - effective_cached_tokens: HashMap, - router_config_override: Option<&RouterConfigOverride>, - update_states: bool, - lora_name: Option, - priority_jump: f64, - strict_priority: u32, - expected_output_tokens: Option, - pinned_worker: Option, - allowed_worker_ids: Option>, - routing_constraints: RoutingConstraints, - shared_cache_hits: Option, - ) -> Result { - self.schedule_with_block_hashes( - maybe_request_id, - isl_tokens, - token_seq, - None, - tier_overlap_blocks, - effective_overlap_blocks, - effective_cached_tokens, - router_config_override, - update_states, - lora_name, - priority_jump, - strict_priority, - expected_output_tokens, - pinned_worker, - allowed_worker_ids, - routing_constraints, - shared_cache_hits, - ) - .await - } - - /// Like [`schedule`](Self::schedule) but forwards the block hashes used to compute the - /// initial overlap scores. Required to enable dequeue-time overlap refresh; ignored if - /// the scheduler was not constructed with an [`OverlapScoresRefresh`]. - #[expect(clippy::too_many_arguments)] - pub async fn schedule_with_block_hashes( - &self, - maybe_request_id: Option, - isl_tokens: usize, - token_seq: Option>, - block_hashes: Option>, - tier_overlap_blocks: TierOverlapBlocks, - effective_overlap_blocks: HashMap, - effective_cached_tokens: HashMap, - router_config_override: Option<&RouterConfigOverride>, - update_states: bool, - lora_name: Option, - priority_jump: f64, - strict_priority: u32, - expected_output_tokens: Option, - pinned_worker: Option, - allowed_worker_ids: Option>, - routing_constraints: RoutingConstraints, - shared_cache_hits: Option, - ) -> Result { - self.schedule_with_policy_class_and_block_hashes( - maybe_request_id, - isl_tokens, - token_seq, - block_hashes, - tier_overlap_blocks, - effective_overlap_blocks, - effective_cached_tokens, - router_config_override, - update_states, - lora_name, - priority_jump, - strict_priority, - None, - expected_output_tokens, - pinned_worker, - allowed_worker_ids, - routing_constraints, - shared_cache_hits, - ) - .await - } - - #[expect(clippy::too_many_arguments)] - pub async fn schedule_with_policy_class_and_block_hashes( - &self, - maybe_request_id: Option, - isl_tokens: usize, - token_seq: Option>, - block_hashes: Option>, - tier_overlap_blocks: TierOverlapBlocks, - effective_overlap_blocks: HashMap, - effective_cached_tokens: HashMap, - router_config_override: Option<&RouterConfigOverride>, - update_states: bool, - lora_name: Option, - priority_jump: f64, - strict_priority: u32, - policy_class: Option, - expected_output_tokens: Option, - pinned_worker: Option, - allowed_worker_ids: Option>, - routing_constraints: RoutingConstraints, - shared_cache_hits: Option, - ) -> Result { - let response = self - .inner - .schedule_with_policy_class_and_block_hashes( - maybe_request_id, - isl_tokens, - token_seq, - block_hashes, - tier_overlap_blocks, - effective_overlap_blocks, - effective_cached_tokens, - router_config_override, - update_states, - lora_name, - priority_jump, - strict_priority, - policy_class, - expected_output_tokens, - pinned_worker, - allowed_worker_ids, - routing_constraints, - shared_cache_hits, - ) - .await; - self.observe_schedule_result(&response); - response - } - - fn observe_schedule_result(&self, response: &Result) { - if let Err(KvSchedulerError::QueueRejected(rejection)) = response - && let Some(metrics) = self - .queue_metric_indices - .get(&rejection.policy_class) - .and_then(|index| self.queue_metrics.get(*index)) - { - match rejection.limit_kind { - dynamo_kv_router::scheduling::QueueLimitKind::Requests => { - metrics.request_limit_rejections.inc(); - } - dynamo_kv_router::scheduling::QueueLimitKind::RawIslTokens => { - metrics.raw_isl_limit_rejections.inc(); - } - dynamo_kv_router::scheduling::QueueLimitKind::CachedTokens => { - metrics.cached_token_limit_rejections.inc(); - } - } - } - self.update_queue_metrics(); - } - - /// Select a worker from current scheduler state without queue admission or booking. - pub async fn select_without_admission( - &self, - request: ScheduleRequest, - ) -> Result { - self.inner.select_without_admission(request).await - } - - pub fn register_workers(&self, worker_ids: &HashSet) { - self.inner.register_workers(worker_ids); - } - - pub async fn add_request(&self, req: SequenceRequest) -> Result<(), SequenceError> { - self.inner.add_request(req).await - } - - pub async fn add_request_admitted( - &self, - req: SequenceRequest, - ) -> Result { - self.inner.add_request_admitted(req).await - } - - pub async fn mark_prefill_completed(&self, request_id: &str) -> Result<(), SequenceError> { - self.inner.mark_prefill_completed(request_id).await?; - self.update_queue_metrics(); - Ok(()) - } - - pub async fn free(&self, request_id: &str) -> Result<(), SequenceError> { - self.inner.free(request_id).await?; - self.update_queue_metrics(); - Ok(()) - } - - /// Release `request_id` only if it is still booked on `worker`. - pub async fn free_if_worker( - &self, - request_id: &str, - worker: WorkerWithDpRank, - ) -> Result<(), SequenceError> { - self.inner.free_if_worker(request_id, worker).await?; - self.update_queue_metrics(); - Ok(()) - } - - pub(crate) fn booking_cleanup(&self) -> SchedulerBookingCleanup { - self.inner.booking_cleanup() - } - - pub(crate) fn set_replica_request_lease_observer( - &self, - observer: Arc, - ) -> bool { - self.replica_request_lease_observer - .as_ref() - .is_some_and(|deferred| deferred.install(observer)) - } - - pub(crate) async fn mark_prefill_completed_if_booking( - &self, - booking: &SchedulerBookingDescriptor, - ) -> Result<(), KvSchedulerError> { - self.inner - .mark_prefill_completed_if_booking(booking) - .await?; - self.update_queue_metrics(); - Ok(()) - } - - pub fn pending_count(&self) -> usize { - self.inner.pending_count() - } - - pub fn pending_isl_tokens(&self) -> usize { - self.inner.pending_isl_tokens() - } - - pub fn worker_type(&self) -> &'static str { - self.inner.worker_type() - } - - pub fn add_output_block( - &self, - request_id: &str, - decay_fraction: Option, - ) -> Result<(), SequenceError> { - self.inner.add_output_block(request_id, decay_fraction) - } - - pub(crate) async fn enqueue_output_block_if_booking( - &self, - booking: &SchedulerBookingDescriptor, - decay_fraction: Option, - ) -> Result<(), KvSchedulerError> { - self.inner - .enqueue_output_block_if_booking(booking, decay_fraction) - .await - } - - pub fn get_potential_loads( - &self, - token_seq: Option>, - isl_tokens: usize, - effective_cached_tokens: HashMap, - track_prefill_tokens: bool, - ) -> Vec { - self.inner.get_potential_loads( - token_seq, - isl_tokens, - effective_cached_tokens, - track_prefill_tokens, - ) - } - - pub fn get_active_lora_counts(&self) -> HashMap { - self.inner.get_active_lora_counts() - } - - pub fn supports_overlap_refresh(&self) -> bool { - self.inner.supports_overlap_refresh() - } - - fn update_queue_metrics(&self) { - update_queue_metrics(&self.queue_metrics, |class_index| { - self.inner.class_queue_stats(class_index) - }); - } -} - -fn update_queue_metrics( - handles: &[RouterQueueMetricHandles], - mut stats_for: impl FnMut(usize) -> Option, -) { - for (class_index, handles) in handles.iter().enumerate() { - let Some(stats) = stats_for(class_index) else { - debug_assert!( - false, - "missing queue counters for policy class {class_index}" - ); - continue; - }; - handles.pending_requests.set(stats.pending_count as i64); - handles - .pending_isl_tokens - .set(stats.pending_isl_tokens as i64); - handles - .pending_cached_tokens - .set(stats.pending_cached_tokens as i64); - } -} - -#[cfg(test)] -mod tests { - use super::*; - use dynamo_runtime::component::Component; - use dynamo_runtime::{DistributedRuntime, Runtime, distributed::DistributedConfig}; - use tokio::sync::watch; - - async fn make_test_component(name: &str) -> Component { - let runtime = Runtime::from_current().unwrap(); - let drt = DistributedRuntime::new(runtime, DistributedConfig::process_local()) - .await - .unwrap(); - let namespace = drt.namespace(format!("test-ns-{name}")).unwrap(); - namespace - .component(format!("test-component-{name}")) - .unwrap() - } - - #[test] - fn queue_metrics_are_updated_by_class_index() { - let handles = ["latency", "bulk"] - .map(|class| ROUTER_QUEUE_METRICS.handles("index-test", "decode", class)); - let stats = [ - dynamo_kv_router::queue::ClassQueueStats { - pending_count: 2, - pending_isl_tokens: 128, - pending_cached_tokens: 64, - }, - dynamo_kv_router::queue::ClassQueueStats { - pending_count: 3, - pending_isl_tokens: 384, - pending_cached_tokens: 192, - }, - ]; - - update_queue_metrics(&handles, |class_index| stats.get(class_index).copied()); - - for (handles, stats) in handles.iter().zip(stats) { - assert_eq!(handles.pending_requests.get(), stats.pending_count as i64); - assert_eq!( - handles.pending_isl_tokens.get(), - stats.pending_isl_tokens as i64 - ); - assert_eq!( - handles.pending_cached_tokens.get(), - stats.pending_cached_tokens as i64 - ); - } - } - - #[tokio::test] - async fn skip_initial_worker_wait_still_monitors_worker_config_updates() { - let component = make_test_component("skip-initial-worker-watch").await; - let mut workers = HashMap::new(); - workers.insert(0, ModelRuntimeConfig::default()); - let (cfg_tx, cfg_rx) = watch::channel(workers); - let config = KvRouterConfig { - skip_initial_worker_wait: true, - use_kv_events: false, - router_track_active_blocks: false, - ..Default::default() - }; - let cancellation_token = CancellationToken::new(); - - let scheduler = KvScheduler::start( - component.endpoint("generate"), - 64, - cfg_rx, - DefaultWorkerSelector::new(Some(config.clone()), "decode"), - &config, - None, - None::>, - None, - None, - Some("test-model"), - "decode", - super::super::routing_load::scheduler_load_channel( - super::super::RouterLoadSource::Decode, - cancellation_token.child_token(), - ) - .0, - cancellation_token.clone(), - ) - .await - .unwrap(); - - assert_eq!( - scheduler - .get_potential_loads(None, 64, HashMap::new(), true) - .len(), - 1 - ); - - let mut updated_workers = HashMap::new(); - updated_workers.insert(0, ModelRuntimeConfig::default()); - updated_workers.insert(1, ModelRuntimeConfig::default()); - cfg_tx.send(updated_workers).unwrap(); - - tokio::time::timeout(Duration::from_secs(1), async { - loop { - if scheduler - .get_potential_loads(None, 64, HashMap::new(), true) - .iter() - .any(|load| load.worker_id == 1) - { - break; - } - tokio::task::yield_now().await; - } - }) - .await - .unwrap(); - - cancellation_token.cancel(); - } -} diff --git a/lib/llm/src/kv_router/sequence.rs b/lib/llm/src/kv_router/sequence.rs index 9e8295503778..e7164b085cda 100644 --- a/lib/llm/src/kv_router/sequence.rs +++ b/lib/llm/src/kv_router/sequence.rs @@ -7,8 +7,6 @@ //! implementations that wire the runtime-agnostic business logic (in `dynamo_kv_router`) //! to the configured event transport and Prometheus metrics. -mod direct_zmq; - pub use dynamo_kv_router::multi_worker_sequence::{ ActiveSequencesMultiWorker, ReplicaRequestLeaseObserver, SchedulerLoadSnapshot, SequenceError, SequencePublishQueueError, SequencePublisher, SequenceRequest, SequenceSubscriber, @@ -17,7 +15,6 @@ use dynamo_kv_router::protocols::{ ActiveSequenceEvent, ActiveSequenceEventBatch, MAX_REPLICA_BATCH_DURATION, MAX_REPLICA_BATCH_EVENTS, WorkerWithDpRank, }; -use dynamo_kv_router::scheduling::queue::SchedulerBookingDescriptor; pub use dynamo_kv_router::sequence::{ActiveSequences, RequestId}; use anyhow::Result; @@ -26,20 +23,15 @@ use dynamo_runtime::traits::DistributedRuntimeProvider; use dynamo_runtime::transports::event_plane::{ EventPublisher, EventSubscriber, EventTransportKind, TypedEventSubscriber, }; -use parking_lot::Mutex; -use std::collections::{HashMap, VecDeque}; +use std::collections::VecDeque; use std::future::Future; -use std::sync::Arc; use std::task::{Context, Poll}; use tokio::sync::mpsc; use tokio::time::Instant; use tokio_util::sync::CancellationToken; -use super::metrics::{RouterWorkerStatusMetrics, WORKER_LOAD_METRICS}; -use crate::kv_router::{ACTIVE_SEQUENCES_SUBJECT, SchedulerLoadSender}; -use crate::local_model::runtime_config::ModelRuntimeConfig; +use crate::kv_router::ACTIVE_SEQUENCES_SUBJECT; #[cfg(test)] -use dynamo_kv_router::protocols::PrefillLoadHint; #[cfg(test)] use dynamo_runtime::transports::event_plane::MsgpackCodec; @@ -47,108 +39,8 @@ use dynamo_runtime::transports::event_plane::MsgpackCodec; // if the queue is full, the newest event is dropped without blocking the local mutation. const REPLICA_EVENT_CHANNEL_CAPACITY: usize = 100_000; -enum DeferredReplicaLeaseEvent { - Admitted(SchedulerBookingDescriptor), - Progressed(SchedulerBookingDescriptor), - Completed(SchedulerBookingDescriptor), -} - -enum DeferredReplicaLeaseState { - Buffering(Vec), - Installing { - target: Arc, - pending: Vec, - }, - Installed(Arc), -} - -impl Default for DeferredReplicaLeaseState { - fn default() -> Self { - Self::Buffering(Vec::new()) - } -} - -/// Buffers the construction-time event window until `KvRouter` installs its -/// request lease manager. After installation, calls pass through directly. -#[derive(Default)] -pub(crate) struct DeferredReplicaRequestLeaseObserver { - state: Mutex, -} - -impl DeferredReplicaRequestLeaseObserver { - pub(crate) fn install(&self, target: Arc) -> bool { - let mut pending = { - let mut state = self.state.lock(); - let DeferredReplicaLeaseState::Buffering(pending) = &mut *state else { - return false; - }; - let pending = std::mem::take(pending); - *state = DeferredReplicaLeaseState::Installing { - target: Arc::clone(&target), - pending: Vec::new(), - }; - pending - }; - - loop { - for event in pending { - Self::notify(&target, event); - } - let mut state = self.state.lock(); - let DeferredReplicaLeaseState::Installing { - target: installing_target, - pending: queued, - } = &mut *state - else { - unreachable!("observer installation state changed unexpectedly"); - }; - if queued.is_empty() { - *state = DeferredReplicaLeaseState::Installed(Arc::clone(installing_target)); - return true; - } - pending = std::mem::take(queued); - } - } - - fn observe(&self, event: DeferredReplicaLeaseEvent) { - let target = { - let mut state = self.state.lock(); - match &mut *state { - DeferredReplicaLeaseState::Buffering(pending) - | DeferredReplicaLeaseState::Installing { pending, .. } => { - pending.push(event); - return; - } - DeferredReplicaLeaseState::Installed(target) => Arc::clone(target), - } - }; - Self::notify(&target, event); - } - - fn notify(target: &Arc, event: DeferredReplicaLeaseEvent) { - match event { - DeferredReplicaLeaseEvent::Admitted(booking) => target.admitted(booking), - DeferredReplicaLeaseEvent::Progressed(booking) => target.progressed(&booking), - DeferredReplicaLeaseEvent::Completed(booking) => target.completed(&booking), - } - } -} - -impl ReplicaRequestLeaseObserver for DeferredReplicaRequestLeaseObserver { - fn admitted(&self, booking: SchedulerBookingDescriptor) { - self.observe(DeferredReplicaLeaseEvent::Admitted(booking)); - } - - fn progressed(&self, booking: &SchedulerBookingDescriptor) { - self.observe(DeferredReplicaLeaseEvent::Progressed(booking.clone())); - } - - fn completed(&self, booking: &SchedulerBookingDescriptor) { - self.observe(DeferredReplicaLeaseEvent::Completed(booking.clone())); - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +/// How active-sequence events are framed on the wire for a transport. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] enum ActiveSequenceEventWireFormat { Singleton, Batch, @@ -253,6 +145,7 @@ impl ActiveSequenceEventPublisher { } } +#[cfg(test)] fn active_sequence_event_channel( enabled: bool, capacity: usize, @@ -265,56 +158,7 @@ fn active_sequence_event_channel( .then(|| ActiveSequenceEventPublisher::channel(capacity, cancellation_token.child_token())) } -/// Concrete [`SequencePublisher`] backed by the runtime event plane and Prometheus gauges. -pub struct RuntimeSequencePublisher { - event_sender: Option, - scheduler_load: SchedulerLoadSender, - worker_status_metrics: Arc, -} - -impl SequencePublisher for RuntimeSequencePublisher { - fn enqueue_event(&self, event: ActiveSequenceEvent) -> anyhow::Result<()> { - let Some(event_sender) = &self.event_sender else { - return Ok(()); - }; - event_sender.enqueue(event) - } - - fn publish_scheduler_load(&self, snapshot: SchedulerLoadSnapshot) { - self.scheduler_load.publish(snapshot); - } - - fn publish_scheduler_load_batch(&self, snapshots: Vec) { - self.scheduler_load.publish_batch(snapshots); - } - - fn observe_load( - &self, - worker: &WorkerWithDpRank, - worker_type: &str, - blocks: usize, - tokens: usize, - ) { - WORKER_LOAD_METRICS.observe( - worker.worker_id, - worker.dp_rank, - worker_type, - blocks, - tokens, - ); - } - - fn observe_worker_registered(&self, worker: &WorkerWithDpRank, worker_type: &str) { - self.worker_status_metrics - .set_registered(worker.worker_id, worker.dp_rank, worker_type); - } - - fn observe_worker_removed(&self, worker: &WorkerWithDpRank, worker_type: &str) { - self.worker_status_metrics - .remove_worker(worker.worker_id, worker.dp_rank, worker_type); - } -} - +/// One event per message, for transports that carry singletons. trait SingletonEventPublisher: Send + Sync { fn publish_event( &self, @@ -513,9 +357,6 @@ impl SequenceSubscriber for RuntimeSequenceSubscriber { } } -/// Type alias for the runtime-wired multi-worker sequence tracker. -pub type ActiveSequencesMulti = ActiveSequencesMultiWorker; - /// Replica-sync channels for an embedded selection partition, carried over the /// runtime event plane exactly like the runtime scheduler's replica sync: /// outbound events are published on `ACTIVE_SEQUENCES_SUBJECT`, peer events @@ -580,179 +421,13 @@ pub(crate) async fn host_replica_channels( }) } -/// Convenience async constructor that creates the event-plane publishers/subscribers -/// and returns an `Arc` with replica sync already running. -pub async fn create_multi_worker_sequences( - endpoint: Endpoint, - block_size: usize, - workers_with_configs: HashMap, - replica_sync: bool, - router_id: u64, - scheduler_load: SchedulerLoadSender, - cancellation_token: CancellationToken, -) -> Result> { - create_multi_worker_sequences_with_observer( - endpoint, - block_size, - workers_with_configs, - replica_sync, - router_id, - scheduler_load, - None, - cancellation_token, - ) - .await -} - -#[expect(clippy::too_many_arguments)] -pub(crate) async fn create_multi_worker_sequences_with_observer( - endpoint: Endpoint, - block_size: usize, - workers_with_configs: HashMap, - replica_sync: bool, - router_id: u64, - scheduler_load: SchedulerLoadSender, - replica_request_lease_observer: Option>, - cancellation_token: CancellationToken, -) -> Result> { - let worker_type = scheduler_load.metric_label(); - let transport_kind = endpoint.drt().default_event_transport_kind(); - let event_sender = if let Some((event_sender, event_rx)) = active_sequence_event_channel( - replica_sync, - REPLICA_EVENT_CHANNEL_CAPACITY, - &cancellation_token, - ) { - let publisher_cancellation_token = event_sender.cancellation_token.clone(); - let event_publisher = EventPublisher::for_endpoint_with_transport( - &endpoint, - ACTIVE_SEQUENCES_SUBJECT, - transport_kind, - ) - .await?; - match active_sequence_event_wire_format(transport_kind) { - ActiveSequenceEventWireFormat::Singleton => { - tokio::spawn(run_replica_singleton_publisher( - event_publisher, - event_rx, - publisher_cancellation_token, - )); - } - ActiveSequenceEventWireFormat::Batch => { - tokio::spawn(run_replica_batch_publisher( - event_publisher, - event_rx, - publisher_cancellation_token, - )); - } - } - Some(event_sender) - } else { - None - }; - let worker_status_metrics = RouterWorkerStatusMetrics::from_component(endpoint.component()); - - let publisher = RuntimeSequencePublisher { - event_sender, - scheduler_load, - worker_status_metrics, - }; - - let dp_range: HashMap = workers_with_configs - .into_iter() - .map(|(id, config)| { - ( - id, - (config.data_parallel_start_rank, config.data_parallel_size), - ) - }) - .collect(); - - let uses_shared_lease_manager = replica_request_lease_observer.is_some(); - let multi_worker = if uses_shared_lease_manager { - // KvRouter's request lease manager owns the single CLOCK reaper for scheduler - // and approximate-LRU cleanup. Cache-retention TTL is unrelated. - ActiveSequencesMultiWorker::new_without_expiry( - publisher, - block_size, - dp_range, - replica_sync, - router_id, - worker_type, - ) - } else { - ActiveSequencesMultiWorker::new( - publisher, - block_size, - dp_range, - replica_sync, - router_id, - worker_type, - ) - }; - - let arc = Arc::new(multi_worker); - if let Some(observer) = replica_request_lease_observer - && !arc.set_replica_request_lease_observer(observer) - { - anyhow::bail!("replica request lease observer is already installed"); - } - if !uses_shared_lease_manager { - arc.start_periodic_force_expiry_across_all_workers(cancellation_token.child_token()); - } - - // Worker-origin completion marks are consumed even when router-to-router replica sync is - // disabled. The tracker filters all other remote lifecycle events in that mode. - let direct_config = direct_zmq::DirectZmqSequenceConfig::from_env(); - let ingress_result = if direct_config.should_use_direct(transport_kind) { - direct_zmq::start( - endpoint, - arc.clone(), - direct_config.rcvhwm, - cancellation_token.child_token(), - ) - .await - .map(|_task| ()) - } else { - match RuntimeSequenceSubscriber::for_endpoint(&endpoint).await { - Ok(subscriber) => { - arc.start_replica_sync(subscriber, cancellation_token.child_token()); - Ok(()) - } - Err(error) => Err(error), - } - }; - if let Err(error) = ingress_result { - tracing::warn!( - %error, - "active-sequence event ingress unavailable; continuing with response-side cleanup" - ); - } - - Ok(arc) -} - #[cfg(test)] mod tests { use super::*; use dynamo_kv_router::protocols::ActiveSequenceEventData; - use dynamo_runtime::{DistributedRuntime, Runtime, distributed::DistributedConfig}; + use std::sync::Arc; use tokio::time::Instant; - fn scheduler_load_sender() -> SchedulerLoadSender { - super::super::routing_load::scheduler_load_channel( - super::super::RouterLoadSource::Decode, - CancellationToken::new(), - ) - .0 - } - - fn tracking_hint(tokens: usize) -> Option { - Some(PrefillLoadHint { - initial_effective_prefill_tokens: tokens, - expected_prefill_duration: None, - }) - } - fn free_event(request_id: impl Into) -> ActiveSequenceEvent { ActiveSequenceEvent { request_id: request_id.into(), @@ -1041,439 +716,4 @@ mod tests { .is_err() ); } - - #[tokio::test] - async fn active_sequence_replica_sync_isolated_by_endpoint() -> Result<()> { - let runtime = Runtime::from_current()?; - let distributed = - DistributedRuntime::new(runtime, DistributedConfig::process_local()).await?; - let namespace = distributed.namespace(format!( - "active-sequence-endpoint-isolation-{}", - uuid::Uuid::new_v4() - ))?; - let component = namespace.component("workers")?; - let endpoint_a = component.endpoint("generate-a"); - let endpoint_b = component.endpoint("generate-b"); - let workers = HashMap::from([(0, ModelRuntimeConfig::new())]); - - let cancel = CancellationToken::new(); - let sequences_a = create_multi_worker_sequences( - endpoint_a.clone(), - 4, - workers.clone(), - true, - 1, - scheduler_load_sender(), - cancel.child_token(), - ) - .await?; - let sequences_a_peer = create_multi_worker_sequences( - endpoint_a, - 4, - workers.clone(), - true, - 3, - scheduler_load_sender(), - cancel.child_token(), - ) - .await?; - let sequences_b = create_multi_worker_sequences( - endpoint_b, - 4, - workers, - true, - 2, - scheduler_load_sender(), - cancel.child_token(), - ) - .await?; - - let worker = WorkerWithDpRank::new(0, 0); - tokio::time::timeout(tokio::time::Duration::from_secs(5), async { - for request_index in 0..100 { - if sequences_a_peer.active_blocks()[&worker] > 0 { - break; - } - - sequences_a.add_request( - SequenceRequest { - request_id: format!("endpoint-a-request-{request_index}"), - token_sequence: Some(vec![1, 2, 3, 4]), - track_prefill_tokens: true, - expected_output_tokens: None, - prefill_load_hint: tracking_hint(4), - worker, - lora_name: None, - }, - Instant::now(), - )?; - tokio::time::sleep(tokio::time::Duration::from_millis(50)).await; - } - - anyhow::ensure!(sequences_a_peer.active_blocks()[&worker] > 0); - Ok::<_, anyhow::Error>(()) - }) - .await??; - assert!(sequences_a.active_blocks()[&worker] > 0); - assert!(sequences_a_peer.active_blocks()[&worker] > 0); - let leaked_to_b = tokio::time::timeout(tokio::time::Duration::from_millis(250), async { - loop { - if sequences_b.active_blocks()[&worker] > 0 { - break; - } - tokio::task::yield_now().await; - } - }) - .await; - assert!( - leaked_to_b.is_err(), - "endpoint B received endpoint A sequence state" - ); - assert_eq!(sequences_b.active_blocks()[&worker], 0); - cancel.cancel(); - Ok(()) - } - - #[tokio::test] - async fn worker_completion_ingress_runs_without_router_replica_sync() -> Result<()> { - let runtime = Runtime::from_current()?; - let distributed = - DistributedRuntime::new(runtime, DistributedConfig::process_local()).await?; - let endpoint = distributed - .namespace(format!( - "worker-completion-ingress-{}", - uuid::Uuid::new_v4() - ))? - .component("workers")? - .endpoint("generate"); - let worker_id = 42; - let worker = WorkerWithDpRank::new(worker_id, 0); - let cancel = CancellationToken::new(); - let sequences = create_multi_worker_sequences( - endpoint.clone(), - 4, - HashMap::from([(worker_id, ModelRuntimeConfig::new())]), - false, - 99, - scheduler_load_sender(), - cancel.child_token(), - ) - .await?; - let request_id = "worker-origin-mark".to_string(); - sequences.add_request( - SequenceRequest { - request_id: request_id.clone(), - token_sequence: Some(vec![1, 2, 3]), - track_prefill_tokens: true, - expected_output_tokens: None, - prefill_load_hint: tracking_hint(12), - worker, - lora_name: None, - }, - Instant::now(), - )?; - - let publisher = ActiveSequenceEventPublisher::for_endpoint(&endpoint, 16).await?; - tokio::time::timeout(tokio::time::Duration::from_secs(5), async { - loop { - publisher.mark_prefill_completed(request_id.clone(), worker_id, 0)?; - if sequences.active_tokens(Instant::now())[&worker] == 0 { - break; - } - tokio::time::sleep(tokio::time::Duration::from_millis(50)).await; - } - Ok::<_, anyhow::Error>(()) - }) - .await??; - - drop(publisher); - cancel.cancel(); - Ok(()) - } - - #[tokio::test] - #[ignore] - async fn test_multi_worker_cross_instance_sync() -> Result<()> { - dynamo_runtime::logging::init(); - - let block_size = 4; - - let runtime = Runtime::from_current()?; - let distributed = DistributedRuntime::from_settings(runtime.clone()).await?; - - let namespace = distributed.namespace("test_cross_instance_sync")?; - let endpoint = namespace.component("sequences")?.endpoint("generate"); - - let mut workers_with_configs = HashMap::new(); - - let mut config_worker_0 = crate::local_model::runtime_config::ModelRuntimeConfig::new(); - config_worker_0.data_parallel_size = 2; - workers_with_configs.insert(0, config_worker_0); - - let config_worker_1 = crate::local_model::runtime_config::ModelRuntimeConfig::new(); - workers_with_configs.insert(1, config_worker_1); - - let seq_manager_1 = create_multi_worker_sequences( - endpoint.clone(), - block_size, - workers_with_configs.clone(), - true, - 1, - scheduler_load_sender(), - CancellationToken::new(), - ) - .await?; - let seq_manager_2 = create_multi_worker_sequences( - endpoint, - block_size, - workers_with_configs, - true, - 2, - scheduler_load_sender(), - CancellationToken::new(), - ) - .await?; - - tokio::time::sleep(tokio::time::Duration::from_millis(200)).await; - let decay_now = Instant::now(); - - seq_manager_1.add_request( - SequenceRequest { - request_id: "request_0".to_string(), - token_sequence: Some(vec![0, 1, 2]), - track_prefill_tokens: true, - expected_output_tokens: None, - prefill_load_hint: tracking_hint(12), - worker: WorkerWithDpRank::new(0, 0), - lora_name: None, - }, - decay_now, - )?; - - seq_manager_1.add_request( - SequenceRequest { - request_id: "request_1".to_string(), - token_sequence: Some(vec![3, 4]), - track_prefill_tokens: true, - expected_output_tokens: None, - prefill_load_hint: tracking_hint(8), - worker: WorkerWithDpRank::new(0, 1), - lora_name: None, - }, - decay_now, - )?; - - seq_manager_2.add_request( - SequenceRequest { - request_id: "request_2".to_string(), - token_sequence: Some(vec![0, 1, 2, 3]), - track_prefill_tokens: true, - expected_output_tokens: None, - prefill_load_hint: tracking_hint(16), - worker: WorkerWithDpRank::new(1, 0), - lora_name: None, - }, - decay_now, - )?; - - tokio::time::sleep(tokio::time::Duration::from_millis(300)).await; - - let blocks_phase1 = seq_manager_1.active_blocks(); - let tokens_phase1 = seq_manager_1.active_tokens(Instant::now()); - - let worker_0_dp0 = WorkerWithDpRank::new(0, 0); - let worker_0_dp1 = WorkerWithDpRank::new(0, 1); - let worker_1_dp0 = WorkerWithDpRank::new(1, 0); - - assert_eq!( - blocks_phase1[&worker_0_dp0], 3, - "Worker 0 dp_rank 0 should have 3 active blocks (from request_0)" - ); - assert_eq!( - blocks_phase1[&worker_0_dp1], 2, - "Worker 0 dp_rank 1 should have 2 active blocks (from request_1)" - ); - assert_eq!( - blocks_phase1[&worker_1_dp0], 4, - "Worker 1 dp_rank 0 should have 4 active blocks (from request_2 added by seq_manager_2)" - ); - assert_eq!( - tokens_phase1[&worker_0_dp0], 12, - "Worker 0 dp_rank 0 should have 12 active tokens" - ); - assert_eq!( - tokens_phase1[&worker_0_dp1], 8, - "Worker 0 dp_rank 1 should have 8 active tokens" - ); - assert_eq!( - tokens_phase1[&worker_1_dp0], 16, - "Worker 1 dp_rank 0 should have 16 active tokens (from request_2 added by seq_manager_2)" - ); - - seq_manager_1.free(&"request_2".to_string(), Instant::now())?; - - seq_manager_2.free(&"request_0".to_string(), Instant::now())?; - seq_manager_2.free(&"request_1".to_string(), Instant::now())?; - - tokio::time::sleep(tokio::time::Duration::from_millis(300)).await; - - let blocks_phase2 = seq_manager_2.active_blocks(); - let tokens_phase2 = seq_manager_2.active_tokens(Instant::now()); - - let all_workers = vec![ - WorkerWithDpRank::new(0, 0), - WorkerWithDpRank::new(0, 1), - WorkerWithDpRank::new(1, 0), - ]; - - for worker in all_workers { - assert_eq!( - blocks_phase2[&worker], 0, - "Worker (id={}, dp_rank={}) should have 0 active blocks after all requests freed", - worker.worker_id, worker.dp_rank - ); - assert_eq!( - tokens_phase2[&worker], 0, - "Worker (id={}, dp_rank={}) should have 0 active tokens after all requests freed", - worker.worker_id, worker.dp_rank - ); - } - - Ok(()) - } - - #[tokio::test] - #[ignore] - async fn test_multi_worker_no_token_sequence_sync() -> Result<()> { - dynamo_runtime::logging::init(); - - let block_size = 4; - - let runtime = Runtime::from_current()?; - let distributed = DistributedRuntime::from_settings(runtime.clone()).await?; - - let namespace = distributed.namespace("test_no_token_seq_sync")?; - let endpoint = namespace.component("sequences")?.endpoint("generate"); - - let mut workers_with_configs = HashMap::new(); - workers_with_configs.insert( - 0, - crate::local_model::runtime_config::ModelRuntimeConfig::new(), - ); - workers_with_configs.insert( - 1, - crate::local_model::runtime_config::ModelRuntimeConfig::new(), - ); - workers_with_configs.insert( - 2, - crate::local_model::runtime_config::ModelRuntimeConfig::new(), - ); - - let seq_manager_1 = create_multi_worker_sequences( - endpoint.clone(), - block_size, - workers_with_configs.clone(), - true, - 1, - scheduler_load_sender(), - CancellationToken::new(), - ) - .await?; - let seq_manager_2 = create_multi_worker_sequences( - endpoint, - block_size, - workers_with_configs, - true, - 2, - scheduler_load_sender(), - CancellationToken::new(), - ) - .await?; - - tokio::time::sleep(tokio::time::Duration::from_millis(200)).await; - let decay_now = Instant::now(); - - seq_manager_1.add_request( - SequenceRequest { - request_id: "request_0".to_string(), - token_sequence: None, - track_prefill_tokens: true, - expected_output_tokens: None, - prefill_load_hint: tracking_hint(12), - worker: WorkerWithDpRank::from_worker_id(0), - lora_name: None, - }, - decay_now, - )?; - - seq_manager_1.add_request( - SequenceRequest { - request_id: "request_1".to_string(), - token_sequence: None, - track_prefill_tokens: true, - expected_output_tokens: None, - prefill_load_hint: tracking_hint(8), - worker: WorkerWithDpRank::from_worker_id(1), - lora_name: None, - }, - decay_now, - )?; - - seq_manager_2.add_request( - SequenceRequest { - request_id: "request_2".to_string(), - token_sequence: None, - track_prefill_tokens: true, - expected_output_tokens: None, - prefill_load_hint: tracking_hint(16), - worker: WorkerWithDpRank::from_worker_id(2), - lora_name: None, - }, - decay_now, - )?; - - tokio::time::sleep(tokio::time::Duration::from_millis(300)).await; - - let tokens_phase1 = seq_manager_1.active_tokens(Instant::now()); - - let worker_0 = WorkerWithDpRank::from_worker_id(0); - let worker_1 = WorkerWithDpRank::from_worker_id(1); - let worker_2 = WorkerWithDpRank::from_worker_id(2); - - assert_eq!( - tokens_phase1[&worker_0], 12, - "Worker 0 should have 12 active tokens" - ); - assert_eq!( - tokens_phase1[&worker_1], 8, - "Worker 1 should have 8 active tokens" - ); - assert_eq!( - tokens_phase1[&worker_2], 16, - "Worker 2 should have 16 active tokens (from request_2 added by seq_manager_2)" - ); - - seq_manager_1.mark_prefill_completed(&"request_2".to_string(), Instant::now())?; - seq_manager_1.free(&"request_2".to_string(), Instant::now())?; - - seq_manager_2.mark_prefill_completed(&"request_0".to_string(), Instant::now())?; - seq_manager_2.mark_prefill_completed(&"request_1".to_string(), Instant::now())?; - seq_manager_2.free(&"request_0".to_string(), Instant::now())?; - seq_manager_2.free(&"request_1".to_string(), Instant::now())?; - - tokio::time::sleep(tokio::time::Duration::from_millis(300)).await; - - let tokens_phase2 = seq_manager_2.active_tokens(Instant::now()); - - for worker_id in 0..=2 { - let worker = WorkerWithDpRank::from_worker_id(worker_id); - assert_eq!( - tokens_phase2[&worker], 0, - "Worker {} should have 0 active tokens after all requests freed", - worker_id - ); - } - - Ok(()) - } } diff --git a/lib/llm/src/kv_router/sequence/direct_zmq.rs b/lib/llm/src/kv_router/sequence/direct_zmq.rs deleted file mode 100644 index 1c458f7d54a1..000000000000 --- a/lib/llm/src/kv_router/sequence/direct_zmq.rs +++ /dev/null @@ -1,290 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -use std::{ffi::OsString, sync::Arc}; - -use anyhow::Result; -use dynamo_kv_router::{ - ActiveSequencesMultiWorker, SequencePublisher, - protocols::{ActiveSequenceEventBatch, MAX_REPLICA_BATCH_EVENTS}, -}; -use dynamo_runtime::{ - component::Endpoint, - config::parse_bool_opt, - discovery::EventTransportKind, - transports::event_plane::{Codec, uses_direct_zmq}, -}; -use tokio::task::JoinHandle; -use tokio_util::sync::CancellationToken; - -use crate::{ - direct_zmq_fan_in::{ContinuityMode, FanInEvent, start_direct_zmq_fan_in}, - kv_router::{ACTIVE_SEQUENCES_SUBJECT, metrics::ActiveSequenceZmqIngressMetrics}, -}; - -const DIRECT_ZMQ_ENV: &str = "DYN_ROUTER_ACTIVE_SEQUENCE_DIRECT_ZMQ"; -const RCVHWM_ENV: &str = "DYN_ROUTER_ACTIVE_SEQUENCE_ZMQ_RCVHWM"; -const DEFAULT_RCVHWM: i32 = 1024; - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(super) struct DirectZmqSequenceConfig { - enabled: bool, - pub(super) rcvhwm: i32, -} - -impl DirectZmqSequenceConfig { - pub(super) fn from_env() -> Self { - Self::from_lookup(|key| std::env::var_os(key)) - } - - fn from_lookup(mut get_env: impl FnMut(&str) -> Option) -> Self { - let enabled = match get_env(DIRECT_ZMQ_ENV) { - None => true, - Some(raw) => { - let value = raw.to_string_lossy(); - match parse_bool_opt(&value) { - Some(enabled) => enabled, - None => { - tracing::warn!( - env = DIRECT_ZMQ_ENV, - %value, - "invalid direct-ZMQ active-sequence setting; using enabled default" - ); - true - } - } - } - }; - - let rcvhwm = match get_env(RCVHWM_ENV) { - None => DEFAULT_RCVHWM, - Some(raw) => match raw.to_string_lossy().trim().parse::() { - Ok(value) if value > 0 => value, - _ => { - tracing::warn!( - env = RCVHWM_ENV, - value = %raw.to_string_lossy(), - default = DEFAULT_RCVHWM, - "invalid direct-ZMQ active-sequence receive HWM; using default" - ); - DEFAULT_RCVHWM - } - }, - }; - - Self { enabled, rcvhwm } - } - - pub(super) fn should_use_direct(self, transport_kind: EventTransportKind) -> bool { - self.should_use_direct_for_topology(transport_kind, uses_direct_zmq(transport_kind)) - } - - fn should_use_direct_for_topology( - self, - transport_kind: EventTransportKind, - direct_zmq_topology: bool, - ) -> bool { - self.enabled && transport_kind == EventTransportKind::Zmq && direct_zmq_topology - } -} - -pub(super) async fn start( - endpoint: Endpoint, - tracker: Arc>, - rcvhwm: i32, - cancellation_token: CancellationToken, -) -> Result> { - let metrics = ActiveSequenceZmqIngressMetrics::from_component(endpoint.component()); - let handler_metrics = metrics.clone(); - let handler = move |envelope: dynamo_runtime::transports::event_plane::ValidatedEnvelope| { - let codec = Codec::default(); - let batch = codec - .decode_payload::(&envelope.payload) - .inspect_err(|_| { - handler_metrics.record_payload_decode_error(); - })?; - if batch.events.len() > MAX_REPLICA_BATCH_EVENTS { - handler_metrics.record_payload_decode_error(); - anyhow::bail!( - "active-sequence batch contains {} events; maximum is {}", - batch.events.len(), - MAX_REPLICA_BATCH_EVENTS - ); - } - tracker.apply_replica_batch(batch.events); - Ok(()) - }; - let observer = - move |observation: crate::direct_zmq_fan_in::FanInObservation| match observation.event { - FanInEvent::SourceStarted => metrics.source_started(), - FanInEvent::SourceStopped => metrics.source_stopped(), - FanInEvent::Disconnected | FanInEvent::DiscoveryReset => {} - FanInEvent::Reconnect => metrics.record_reconnect(), - FanInEvent::Replacement => metrics.record_replacement(), - FanInEvent::EnvelopeDecodeError => metrics.record_envelope_decode_error(), - FanInEvent::IdentityMismatch => metrics.record_identity_error(), - FanInEvent::SequenceGap { missing } => metrics.record_gap(missing), - FanInEvent::OutOfOrder => metrics.record_out_of_order(), - FanInEvent::ForcedAbort => metrics.record_forced_abort(), - }; - - start_direct_zmq_fan_in( - endpoint, - ACTIVE_SEQUENCES_SUBJECT, - rcvhwm, - None, - ContinuityMode::TrackFromZero, - cancellation_token, - handler, - observer, - ) - .await -} - -#[cfg(test)] -mod tests { - use std::{collections::HashMap, time::Duration}; - - use dynamo_kv_router::{ - NoopSequencePublisher, - protocols::{ActiveSequenceEvent, ActiveSequenceEventData, WorkerWithDpRank}, - }; - use dynamo_runtime::{ - DistributedRuntime, Runtime, distributed::DistributedConfig, - transports::event_plane::EventPublisher, - }; - - use super::*; - - fn lookup(direct: Option<&str>, rcvhwm: Option<&str>) -> impl FnMut(&str) -> Option { - let direct = direct.map(OsString::from); - let rcvhwm = rcvhwm.map(OsString::from); - move |key| match key { - DIRECT_ZMQ_ENV => direct.clone(), - RCVHWM_ENV => rcvhwm.clone(), - _ => None, - } - } - - #[test] - fn parses_direct_zmq_configuration() { - assert_eq!( - DirectZmqSequenceConfig::from_lookup(lookup(None, None)), - DirectZmqSequenceConfig { - enabled: true, - rcvhwm: DEFAULT_RCVHWM, - } - ); - assert_eq!( - DirectZmqSequenceConfig::from_lookup(lookup(Some("false"), Some("37"))), - DirectZmqSequenceConfig { - enabled: false, - rcvhwm: 37, - } - ); - for invalid in ["0", "-1", "not-a-number", "2147483648"] { - assert_eq!( - DirectZmqSequenceConfig::from_lookup(lookup(None, Some(invalid))).rcvhwm, - DEFAULT_RCVHWM - ); - } - assert!(DirectZmqSequenceConfig::from_lookup(lookup(Some("invalid"), None)).enabled); - } - - #[test] - fn selects_only_direct_zmq_and_honors_rollback() { - let enabled = DirectZmqSequenceConfig::from_lookup(lookup(None, None)); - let rollback = DirectZmqSequenceConfig::from_lookup(lookup(Some("0"), None)); - - assert!(enabled.should_use_direct_for_topology(EventTransportKind::Zmq, true)); - assert!(!enabled.should_use_direct_for_topology(EventTransportKind::Zmq, false)); - assert!(!enabled.should_use_direct_for_topology(EventTransportKind::Nats, false)); - assert!(!rollback.should_use_direct_for_topology(EventTransportKind::Zmq, true)); - assert!(!rollback.should_use_direct_for_topology(EventTransportKind::Zmq, false)); - assert!(!rollback.should_use_direct_for_topology(EventTransportKind::Nats, false)); - } - - fn add_event(request_id: impl Into) -> ActiveSequenceEvent { - ActiveSequenceEvent { - request_id: request_id.into(), - worker: WorkerWithDpRank::new(0, 0), - data: ActiveSequenceEventData::AddRequest { - token_sequence: Some(vec![1]), - track_prefill_tokens: false, - expected_output_tokens: None, - prefill_load_hint: None, - }, - router_id: 99, - lora_name: None, - } - } - - #[tokio::test] - #[serial_test::serial] - async fn direct_zmq_batch_ingress_applies_after_publisher_recreation() -> Result<()> { - let runtime = Runtime::from_current()?; - let distributed = - DistributedRuntime::new(runtime, DistributedConfig::process_local()).await?; - let endpoint = distributed - .namespace(format!( - "active-sequence-direct-zmq-{}", - uuid::Uuid::new_v4() - ))? - .component("frontend")? - .endpoint("generate"); - let tracker = Arc::new(ActiveSequencesMultiWorker::new( - NoopSequencePublisher, - 4, - HashMap::from([(0, (0, 1))]), - true, - 1, - "test", - )); - let worker = WorkerWithDpRank::new(0, 0); - let cancel = CancellationToken::new(); - let supervisor = start( - endpoint.clone(), - tracker.clone(), - DEFAULT_RCVHWM, - cancel.clone(), - ) - .await?; - - for (index, request_id) in ["before-recreation", "after-recreation"] - .into_iter() - .enumerate() - { - let publisher = EventPublisher::for_endpoint_with_transport( - &endpoint, - ACTIVE_SEQUENCES_SUBJECT, - EventTransportKind::Zmq, - ) - .await?; - tokio::time::timeout(Duration::from_secs(5), async { - loop { - publisher - .publish(&ActiveSequenceEventBatch { - events: vec![add_event(request_id)], - }) - .await - .unwrap(); - if tracker.active_request_counts()[&worker] == index + 1 { - break; - } - tokio::time::sleep(Duration::from_millis(20)).await; - } - }) - .await - .expect("direct-ZMQ source should apply a warmed publisher batch"); - drop(publisher); - } - - cancel.cancel(); - tokio::time::timeout(Duration::from_secs(5), supervisor) - .await - .expect("fan-in supervisor should stop after cancellation") - .expect("fan-in supervisor should not panic"); - distributed.shutdown(); - Ok(()) - } -} diff --git a/lib/llm/src/lora/load_estimator.rs b/lib/llm/src/lora/load_estimator.rs index d956c507423c..9830d33c0a87 100644 --- a/lib/llm/src/lora/load_estimator.rs +++ b/lib/llm/src/lora/load_estimator.rs @@ -17,10 +17,9 @@ use std::time::{Duration, Instant}; use dashmap::DashMap; use dynamo_kv_router::protocols::{ActiveSequenceEvent, ActiveSequenceEventData}; -use dynamo_runtime::component::{Component, Endpoint}; +use dynamo_runtime::component::Endpoint; use dynamo_runtime::traits::DistributedRuntimeProvider; -use crate::kv_router::scheduler::KvScheduler; use crate::kv_router::sequence::{RuntimeSequenceSubscriber, SequenceSubscriber}; use crate::lora::config::PredictorType; use crate::lora::predictor::{EmaPredictor, LoadPredictor}; @@ -454,31 +453,6 @@ impl LoadEstimator { .clear(); } - pub fn start_polling( - self: Arc, - scheduler: Arc, - component: Component, - ) -> tokio::task::JoinHandle<()> { - let cancel_token = component.drt().child_token(); - tokio::spawn(async move { - let mut interval = tokio::time::interval(self.config.read().poll_interval); - tracing::info!("Started LORA load polling"); - - loop { - tokio::select! { - _ = cancel_token.cancelled() => { - tracing::debug!("LORA load polling task cancelled"); - break; - } - _ = interval.tick() => { - let lora_counts = scheduler.get_active_lora_counts(); - self.update_from_counts(lora_counts); - } - } - } - }) - } - pub fn start_event_subscription( self: Arc, endpoint: Endpoint, @@ -632,6 +606,7 @@ impl LoadEstimator { /// — net delta 0) is invisible, and sub-interval oscillation is lost. /// Event-based mode (`handle_event` / `increment_load`) gives accurate /// arrival rates; prefer it when arrival precision matters. + #[cfg(test)] fn update_from_counts(&self, lora_counts: HashMap) { let now = Instant::now(); let cfg = self.config.read(); From 2c3a1e6f8a51ff0b2abf81c968dc7309d03d4367 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Wed, 9 Sep 2026 11:08:50 -0700 Subject: [PATCH 013/102] style(router): remove duplicate test configuration attribute Signed-off-by: Ishan Dhanani --- lib/llm/src/kv_router/sequence.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/lib/llm/src/kv_router/sequence.rs b/lib/llm/src/kv_router/sequence.rs index e7164b085cda..a61b3bc831a4 100644 --- a/lib/llm/src/kv_router/sequence.rs +++ b/lib/llm/src/kv_router/sequence.rs @@ -32,7 +32,6 @@ use tokio_util::sync::CancellationToken; use crate::kv_router::ACTIVE_SEQUENCES_SUBJECT; #[cfg(test)] -#[cfg(test)] use dynamo_runtime::transports::event_plane::MsgpackCodec; // Match the existing standalone replica-sync queue. Lifecycle callers enqueue without awaiting; From 81cc6b68d26e980d1dd29a4fd3ff155722d922b8 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Thu, 3 Sep 2026 20:40:21 +0000 Subject: [PATCH 014/102] refactor(router): remove the ignored router_embedded_selection flag The embedded selection partition is the only scheduler, so the flag has been a no-op since the runtime scheduler was deleted. It never shipped in a release, so no wire compatibility shim is needed: drop it from KvRouterConfig (Rust and serde), the DYN_ROUTER_EMBEDDED_SELECTION environment override, the Python argument group, the bindings signature, and the frontend configuration reference. Signed-off-by: Ishan Dhanani --- .../configuration/groups/kv_router_args.py | 13 ----------- .../components/frontend-configuration.mdx | 9 -------- lib/bindings/python/rust/llm/entrypoint.rs | 4 +--- lib/bindings/python/src/dynamo/_core.pyi | 3 --- lib/kv-router/src/scheduling/config.rs | 22 ------------------- lib/llm/src/kv_router.rs | 6 ----- lib/llm/src/kv_router/embedded.rs | 17 +++++++------- lib/llm/src/kv_router/routing_host/tests.rs | 1 - 8 files changed, 9 insertions(+), 66 deletions(-) diff --git a/components/src/dynamo/common/configuration/groups/kv_router_args.py b/components/src/dynamo/common/configuration/groups/kv_router_args.py index 011db193a48c..1a43c56a900a 100644 --- a/components/src/dynamo/common/configuration/groups/kv_router_args.py +++ b/components/src/dynamo/common/configuration/groups/kv_router_args.py @@ -39,7 +39,6 @@ "router_temperature", "use_kv_events", "router_replica_sync", - "router_embedded_selection", "router_track_active_blocks", "router_track_output_blocks", "router_assume_kv_reuse", @@ -201,7 +200,6 @@ class KvRouterConfigBase(ConfigBase): router_temperature: float use_kv_events: bool router_replica_sync: bool - router_embedded_selection: bool = True router_track_active_blocks: bool router_track_output_blocks: bool router_assume_kv_reuse: bool @@ -425,17 +423,6 @@ def add_arguments(self, parser) -> None: "the Runtime event plane." ), ) - add_negatable_bool_argument( - g, - flag_name="--router-embedded-selection", - env_var="DYN_ROUTER_EMBEDDED_SELECTION", - default=True, - dest="router_embedded_selection", - help=( - "Deprecated and ignored: the KV Router always schedules on an embedded " - "selection-service partition. --no-router-embedded-selection only logs a warning." - ), - ) add_negatable_bool_argument( g, flag_name="--router-track-active-blocks", diff --git a/docs/fern/pages/reference/components/frontend-configuration.mdx b/docs/fern/pages/reference/components/frontend-configuration.mdx index 8899b3222ec7..729a74ecbde3 100644 --- a/docs/fern/pages/reference/components/frontend-configuration.mdx +++ b/docs/fern/pages/reference/components/frontend-configuration.mdx @@ -213,15 +213,6 @@ router. See [Router Guide](../../developer-guide/knowledge-base/modular-componen Environment variable: `DYN_ROUTER_REPLICA_SYNC` - - Deprecated and ignored. The router's scheduling and active-sequence accounting always run on an - embedded selection-service partition (the same core as the standalone selection service); the - runtime-bound scheduler was removed. `--router-replica-sync` and overload detection apply to - the embedded partition. `--no-router-embedded-selection` only logs a warning. - - Environment variable: `DYN_ROUTER_EMBEDDED_SELECTION` - - Hash function for router-derived active-sequence identities. `keyed-xxh3-v1` is experimental and requires both a tracking key file and key ID. diff --git a/lib/bindings/python/rust/llm/entrypoint.rs b/lib/bindings/python/rust/llm/entrypoint.rs index d6c17e300f94..3634bdb8f9fe 100644 --- a/lib/bindings/python/rust/llm/entrypoint.rs +++ b/lib/bindings/python/rust/llm/entrypoint.rs @@ -242,7 +242,7 @@ impl AicPerfConfig { #[pymethods] impl KvRouterConfig { #[new] - #[pyo3(signature = (overlap_score_weight=None, host_cache_hit_weight=0.75, disk_cache_hit_weight=0.25, router_temperature=0.0, use_kv_events=true, *, router_replica_sync=false, router_embedded_selection=true, router_track_active_blocks=true, router_track_output_blocks=false, router_assume_kv_reuse=true, router_track_prefill_tokens=true, router_prefill_load_model="none", router_ttl_secs=120.0, router_approximate_cache_policy="ttl", router_queue_threshold=None, router_event_threads=4, router_queue_policy="fcfs", use_remote_indexer=false, serve_indexer=false, shared_cache_multiplier=0.0, shared_cache_type="none", router_predicted_ttl_secs=None, conditional_disagg_enabled=false, conditional_disagg_policy="isl_bounding", conditional_disagg_eff_isl_threshold=2048, conditional_disagg_eff_isl_ratio_threshold=0.7, conditional_disagg_prefill_busy_threshold=None, conditional_disagg_decode_busy_threshold=None, overlap_score_credit=1.0, overlap_score_credit_decay=0.0, prefill_load_scale=1.0, decode_active_request_weight=0.0, router_policy_config=None, router_prefill_policy=None, router_decode_policy=None, router_tracking_hash="public-xxh3-v1", router_tracking_key_file=None, router_tracking_key_id=None))] + #[pyo3(signature = (overlap_score_weight=None, host_cache_hit_weight=0.75, disk_cache_hit_weight=0.25, router_temperature=0.0, use_kv_events=true, *, router_replica_sync=false, router_track_active_blocks=true, router_track_output_blocks=false, router_assume_kv_reuse=true, router_track_prefill_tokens=true, router_prefill_load_model="none", router_ttl_secs=120.0, router_approximate_cache_policy="ttl", router_queue_threshold=None, router_event_threads=4, router_queue_policy="fcfs", use_remote_indexer=false, serve_indexer=false, shared_cache_multiplier=0.0, shared_cache_type="none", router_predicted_ttl_secs=None, conditional_disagg_enabled=false, conditional_disagg_policy="isl_bounding", conditional_disagg_eff_isl_threshold=2048, conditional_disagg_eff_isl_ratio_threshold=0.7, conditional_disagg_prefill_busy_threshold=None, conditional_disagg_decode_busy_threshold=None, overlap_score_credit=1.0, overlap_score_credit_decay=0.0, prefill_load_scale=1.0, decode_active_request_weight=0.0, router_policy_config=None, router_prefill_policy=None, router_decode_policy=None, router_tracking_hash="public-xxh3-v1", router_tracking_key_file=None, router_tracking_key_id=None))] #[allow(clippy::too_many_arguments)] fn new( overlap_score_weight: Option, @@ -251,7 +251,6 @@ impl KvRouterConfig { router_temperature: f64, use_kv_events: bool, router_replica_sync: bool, - router_embedded_selection: bool, router_track_active_blocks: bool, router_track_output_blocks: bool, router_assume_kv_reuse: bool, @@ -302,7 +301,6 @@ impl KvRouterConfig { router_temperature, use_kv_events, router_replica_sync, - router_embedded_selection, router_track_active_blocks, router_track_output_blocks, router_assume_kv_reuse, diff --git a/lib/bindings/python/src/dynamo/_core.pyi b/lib/bindings/python/src/dynamo/_core.pyi index ffeb381961b4..8e11a3a45dfa 100644 --- a/lib/bindings/python/src/dynamo/_core.pyi +++ b/lib/bindings/python/src/dynamo/_core.pyi @@ -1784,7 +1784,6 @@ class KvRouterConfig: use_kv_events: bool = True, *, router_replica_sync: bool = False, - router_embedded_selection: bool = True, router_track_active_blocks: bool = True, router_track_output_blocks: bool = False, router_assume_kv_reuse: bool = True, @@ -1830,8 +1829,6 @@ class KvRouterConfig: router_temperature: Temperature for normalized worker sampling via softmax (default: 0.0) use_kv_events: Whether to use KV events from workers (default: True) router_replica_sync: Enable replica synchronization (default: False) - router_embedded_selection: Deprecated and ignored; scheduling always runs on - an embedded selection-service partition. False only logs a warning. router_track_active_blocks: Track active blocks for load balancing (default: True) router_track_output_blocks: Track output blocks during generation (default: False). When enabled, the router adds placeholder blocks as tokens are generated diff --git a/lib/kv-router/src/scheduling/config.rs b/lib/kv-router/src/scheduling/config.rs index 590368968eee..48ca83d332d8 100644 --- a/lib/kv-router/src/scheduling/config.rs +++ b/lib/kv-router/src/scheduling/config.rs @@ -95,14 +95,6 @@ const fn default_decode_active_request_weight() -> f64 { // Default-valued post-v1.3 fields are omitted so v1.3 frontends can read v1.4 MDCs during // rolling upgrades. Non-default values still serialize and fail closed on the older frontend. // TODO(v1.5): Remove these compatibility skips when v1.3 falls outside the N-1 window. -fn default_true() -> bool { - true -} - -fn is_true(value: &bool) -> bool { - *value -} - fn is_default(value: &T) -> bool { value == &T::default() } @@ -284,9 +276,6 @@ fn kv_router_config_from_lookup( if let Some(value) = parse_bool(&get_env, "DYN_ROUTER_TRACK_PREFILL_TOKENS") { config.router_track_prefill_tokens = value; } - if let Some(value) = parse_bool(&get_env, "DYN_ROUTER_EMBEDDED_SELECTION") { - config.router_embedded_selection = value; - } if let Some(value) = get_env("DYN_ROUTER_TRACKING_HASH") { config.router_tracking_hash = value.parse()?; } @@ -677,8 +666,6 @@ struct KvRouterConfigSerde { router_track_output_blocks: bool, router_assume_kv_reuse: bool, router_track_prefill_tokens: bool, - #[serde(default = "default_true")] - router_embedded_selection: bool, router_tracking_hash: TrackingHashAlgorithm, router_tracking_key_file: Option, router_tracking_key_id: Option, @@ -725,7 +712,6 @@ impl Default for KvRouterConfigSerde { router_track_output_blocks: config.router_track_output_blocks, router_assume_kv_reuse: config.router_assume_kv_reuse, router_track_prefill_tokens: config.router_track_prefill_tokens, - router_embedded_selection: config.router_embedded_selection, router_tracking_hash: config.router_tracking_hash, router_tracking_key_file: config.router_tracking_key_file, router_tracking_key_id: config.router_tracking_key_id, @@ -811,12 +797,6 @@ pub struct KvRouterConfig { #[serde(default = "default_track_prefill_tokens")] pub router_track_prefill_tokens: bool, - /// Deprecated and ignored: the frontend router always schedules on an - /// embedded selection-service partition. Kept so existing configurations - /// and wire records still parse; `false` logs a warning. - #[serde(default = "default_true", skip_serializing_if = "is_true")] - pub router_embedded_selection: bool, - /// Hash algorithm used for router-derived active-sequence identities. #[serde(default, skip_serializing_if = "is_default")] pub router_tracking_hash: TrackingHashAlgorithm, @@ -979,7 +959,6 @@ impl Default for KvRouterConfig { router_track_output_blocks: false, router_assume_kv_reuse: true, router_track_prefill_tokens: default_track_prefill_tokens(), - router_embedded_selection: true, router_tracking_hash: TrackingHashAlgorithm::default(), router_tracking_key_file: None, router_tracking_key_id: None, @@ -1044,7 +1023,6 @@ impl TryFrom for KvRouterConfig { router_track_output_blocks: compat.router_track_output_blocks, router_assume_kv_reuse: compat.router_assume_kv_reuse, router_track_prefill_tokens: compat.router_track_prefill_tokens, - router_embedded_selection: compat.router_embedded_selection, router_tracking_hash: compat.router_tracking_hash, router_tracking_key_file: compat.router_tracking_key_file, router_tracking_key_id: compat.router_tracking_key_id, diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index 225c51a4b4f6..4abf34b7aa3f 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -771,12 +771,6 @@ where let available_worker_provider: WorkerAvailabilityProvider = Arc::new(move || client_for_availability.available_instance_ids()); - if !kv_router_config.router_embedded_selection { - tracing::warn!( - "router_embedded_selection=false is ignored: the runtime-bound KV router scheduler \ - was removed and scheduling always runs on the embedded selection partition" - ); - } // Scheduling and active-sequence accounting run on a selection-service // partition. A caller-injected `WorkerSelector` type (anything other than // the default selector or a registry `WorkerSelectionPolicy`) runs inside diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs index 993a7dca3466..c055970088af 100644 --- a/lib/llm/src/kv_router/embedded.rs +++ b/lib/llm/src/kv_router/embedded.rs @@ -3,15 +3,14 @@ //! Embedded selection backend for the frontend `KvRouter`. //! -//! When `KvRouterConfig::router_embedded_selection` is set, the router runs its -//! scheduling and active-sequence accounting on one partition of an in-process -//! `SelectionService` (the same core the standalone selection service and the -//! EPP use) instead of the runtime-bound `KvScheduler`. The router keeps its own -//! KV indexer and request transport: it computes overlap against that indexer -//! and hands the partition scheduler a fully formed `ScheduleRequest`, exactly -//! as it does with `KvScheduler`. The partition's worker catalog is fed from the -//! runtime-config watch, so discovery remains the source of truth for -//! membership while the scheduler and its replica sync become runtime-free. +//! The router runs its scheduling and active-sequence accounting on one +//! partition of an in-process `SelectionService` (the same core the standalone +//! selection service and the EPP use). The router keeps its own KV indexer and +//! request transport: it computes overlap against that indexer and hands the +//! partition scheduler a fully formed `ScheduleRequest`. The partition's worker +//! catalog is fed from the runtime-config watch, so discovery remains the source +//! of truth for membership while the scheduler and its replica sync are +//! runtime-free. use std::collections::{HashMap, HashSet}; use std::sync::Arc; diff --git a/lib/llm/src/kv_router/routing_host/tests.rs b/lib/llm/src/kv_router/routing_host/tests.rs index 161c2f9b2e62..dc58656381ab 100644 --- a/lib/llm/src/kv_router/routing_host/tests.rs +++ b/lib/llm/src/kv_router/routing_host/tests.rs @@ -1195,7 +1195,6 @@ async fn router_with_worker_configs_in_mode( skip_initial_worker_wait: true, use_kv_events: false, router_track_active_blocks: false, - router_embedded_selection: true, ..Default::default() }; let chooser = KvRouter::new( From 6c2e7d4d28c0c8c4d3e8c5fe3c09e0e520958092 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Thu, 3 Sep 2026 20:59:21 +0000 Subject: [PATCH 015/102] refactor(router): drop the selector type parameter from the frontend router KvRouter, RoutingHost, RoutePlan, RequestGuard, PrefillRouter, ManagedKvRouter, ModelWatcher, and PreprocessedRouting were generic over a caller-injected WorkerSelector type that only ever named the policy the embedded partition would run. Replace the parameter with SelectionPolicySource: Registry (the linked policy the router config names, else Dynamo's default) or Factory (an explicit WorkerSelectionPolicyFactory). The router resolves it once at construction and instantiates one policy to learn the worker inputs it needs, so shared-cache and indexer gating are unchanged. ModelManager constructors take the policy source in place of a selector (kv_chooser_for_with_policy[_and_client], managed_kv_router_for_with_policy); HttpFrontend, ModelWatcher, and PrefillRouter carry it instead of a per-model selector factory. The bindings' custom-policy feature no longer needs a separate router type. The delegating-selector adapter in the embedded backend is removed with its only caller. Unit tests that injected fake selectors now inject picker or filter policies through the public worker-selection API; the overload test becomes a filtered-workers test because a policy cannot report overload directly. Signed-off-by: Ishan Dhanani --- lib/bindings/python/rust/llm/kv.rs | 74 ++-- lib/llm/src/discovery/model_manager.rs | 116 +++--- lib/llm/src/discovery/watcher.rs | 62 +--- lib/llm/src/entrypoint/input.rs | 2 +- lib/llm/src/entrypoint/input/common.rs | 52 +-- lib/llm/src/entrypoint/input/http.rs | 78 ++--- lib/llm/src/kv_router.rs | 329 +++++++++--------- lib/llm/src/kv_router/embedded.rs | 146 +------- .../kv_router/prefill_router/activation.rs | 59 ++-- .../src/kv_router/prefill_router/admission.rs | 39 +-- .../prefill_router/conditional_bypass.rs | 16 +- lib/llm/src/kv_router/prefill_router/mod.rs | 62 +--- lib/llm/src/kv_router/prefill_router/query.rs | 7 +- lib/llm/src/kv_router/routing_host.rs | 77 ++-- lib/llm/src/kv_router/routing_host/builtin.rs | 7 +- lib/llm/src/kv_router/routing_host/kv.rs | 21 +- .../kv_router/routing_host/kv_selection.rs | 7 +- .../kv_router/routing_host/request_guard.rs | 51 +-- lib/llm/src/kv_router/routing_host/tests.rs | 45 +-- lib/llm/src/kv_router/routing_load.rs | 29 +- 20 files changed, 457 insertions(+), 822 deletions(-) diff --git a/lib/bindings/python/rust/llm/kv.rs b/lib/bindings/python/rust/llm/kv.rs index 913d32991c70..de7d09c6c251 100644 --- a/lib/bindings/python/rust/llm/kv.rs +++ b/lib/bindings/python/rust/llm/kv.rs @@ -20,8 +20,6 @@ use crate::Endpoint; feature = "select-service" ))] use clap::Parser; -#[cfg(feature = "custom-policy")] -use dynamo_kv_router::WorkerSelectionPolicy; use dynamo_kv_router::WorkerSelectionPolicyFactory; #[cfg(feature = "select-service")] use dynamo_kv_router::config::try_kv_router_config_from_dynamo_env; @@ -43,20 +41,15 @@ use dynamo_kv_router::services::selection::{ use dynamo_kv_router::services::slot_tracker::{self, SlotTrackerConfig}; #[cfg(feature = "select-service")] use dynamo_kv_router::{TrackingHashAlgorithm, WorkerType}; +use llm_rs::kv_router::SelectionPolicySource; use rs::pipeline::{AsyncEngine, SingleIn}; use rs::protocols::annotated::Annotated as RsAnnotated; use tracing; use llm_rs::discovery::LoadThresholdConfig as RsLoadThresholdConfig; use llm_rs::kv_router::RoutingHost; -#[cfg(not(feature = "custom-policy"))] type RsRoutingHost = RoutingHost; -#[cfg(feature = "custom-policy")] -type RsRoutingHost = RoutingHost; -#[cfg(not(feature = "custom-policy"))] type RsManagedKvRouter = llm_rs::kv_router::ManagedKvRouter; -#[cfg(feature = "custom-policy")] -type RsManagedKvRouter = llm_rs::kv_router::ManagedKvRouter; use llm_rs::kv_router::publisher::{KvEventSourceConfig, create_stored_blocks}; use llm_rs::protocols::common::timing::RequestTracker; use llm_rs::protocols::common::{OutputOptions, SamplingOptions, StopConditions}; @@ -2015,14 +2008,33 @@ async fn create_kv_router_from_endpoint( .await?; #[cfg(not(feature = "custom-policy"))] + let selection_policy = SelectionPolicySource::Registry; + #[cfg(feature = "custom-policy")] + let selection_policy = match worker_selection_policy_factory { + None => SelectionPolicySource::Registry, + // The policy sees the card's typed role and display name, which can + // differ from the metric role and the routing model name. + Some(factory) => { + let policy_worker_role = policy_worker_role + .expect("a configured worker-selection policy waits for a typed model card above"); + let policy_model_name = policy_model_name.unwrap_or_default(); + SelectionPolicySource::Factory(Arc::new(move |config, _worker_type, _partition| { + factory( + config, + policy_worker_role, + dynamo_kv_router::RoutingPartitionRef::new( + &policy_model_name, + dynamo_kv_router::DEFAULT_ROUTING_GROUP, + ), + ) + })) + } + }; let kv_router = model_manager - .kv_chooser_for_with_selector_and_client( + .kv_chooser_for_with_policy_and_client( client, block_size as u32, - dynamo_kv_router::DefaultWorkerSelector::new( - kv_router_config.clone(), - metric_worker_type, - ), + selection_policy, kv_router_config, prefill_load_estimator, worker_role, @@ -2034,42 +2046,6 @@ async fn create_kv_router_from_endpoint( ) .await?; - #[cfg(feature = "custom-policy")] - let kv_router = { - let effective_config = kv_router_config.clone().unwrap_or_default(); - let selector = worker_selection_policy_factory.map_or_else( - || WorkerSelectionPolicy::default(effective_config.clone(), metric_worker_type), - |factory| { - let policy_worker_role = policy_worker_role.expect( - "a configured worker-selection policy waits for a typed model card above", - ); - factory( - &effective_config, - policy_worker_role, - dynamo_kv_router::RoutingPartitionRef::new( - policy_model_name.as_deref().unwrap_or_default(), - dynamo_kv_router::DEFAULT_ROUTING_GROUP, - ), - ) - }, - ); - model_manager - .kv_chooser_for_with_selector_and_client( - client, - block_size as u32, - selector, - kv_router_config, - prefill_load_estimator, - worker_role, - metric_worker_type, - model_name, - enable_eagle, - load_context.scheduler_load_sender(), - load_context.cancellation_token(), - ) - .await? - }; - Ok(llm_rs::kv_router::ManagedKvRouter::new( load_context, kv_router, diff --git a/lib/llm/src/discovery/model_manager.rs b/lib/llm/src/discovery/model_manager.rs index c62914fc99f3..56f50877b472 100644 --- a/lib/llm/src/discovery/model_manager.rs +++ b/lib/llm/src/discovery/model_manager.rs @@ -15,7 +15,7 @@ use dynamo_kv_router::{ PrefillLoadEstimator, config::KvRouterConfig, protocols::{KvTransferEnforcement, RoutingConstraints, WorkerId, WorkerWithDpRank}, - selector::{WorkerInputs, WorkerSelector}, + selector::WorkerInputs, }; use tokio_util::sync::CancellationToken; @@ -35,7 +35,7 @@ use dynamo_runtime::{ use crate::{ kv_router::{ - KvEventSourceRequirement, KvRouter, router_endpoint_id, scheduler::DefaultWorkerSelector, + KvEventSourceRequirement, KvRouter, SelectionPolicySource, router_endpoint_id, shared_cache::HicacheSharedKvCache, }, local_model::runtime_config::{ @@ -1897,11 +1897,10 @@ impl ModelManager { model_name: Option, is_eagle: bool, ) -> anyhow::Result> { - let selector = DefaultWorkerSelector::new(kv_router_config.clone(), metric_worker_type); - self.kv_chooser_for_with_selector( + self.kv_chooser_for_with_policy( endpoint, kv_cache_block_size, - selector, + SelectionPolicySource::Registry, kv_router_config, prefill_load_estimator, worker_role, @@ -1912,23 +1911,20 @@ impl ModelManager { .await } - /// Construct a KV chooser with a selector resolved by the router host at startup. + /// Construct a KV chooser whose partition runs the given selection policy. #[allow(clippy::too_many_arguments)] - pub async fn kv_chooser_for_with_selector( + pub async fn kv_chooser_for_with_policy( &self, endpoint: &Endpoint, kv_cache_block_size: u32, - selector: Sel, + policy: SelectionPolicySource, kv_router_config: Option, prefill_load_estimator: Option>, worker_role: Option, metric_worker_type: &'static str, model_name: Option, is_eagle: bool, - ) -> anyhow::Result>> - where - Sel: WorkerSelector + Send + 'static, - { + ) -> anyhow::Result> { let client = endpoint.client().await?; let source = crate::kv_router::RouterLoadSource::from_worker_role_or_metric( worker_role, @@ -1937,10 +1933,10 @@ impl ModelManager { let parent_token = endpoint.component().drt().child_token(); let scheduler_load = crate::kv_router::SchedulerLoadSender::disabled(source, parent_token.child_token()); - self.kv_chooser_for_with_selector_and_client( + self.kv_chooser_for_with_policy_and_client( client, kv_cache_block_size, - selector, + policy, kv_router_config, prefill_load_estimator, worker_role, @@ -1989,11 +1985,10 @@ impl ModelManager { model_name: Option, is_eagle: bool, ) -> anyhow::Result { - let selector = DefaultWorkerSelector::new(kv_router_config.clone(), metric_worker_type); - self.managed_kv_router_for_with_selector( + self.managed_kv_router_for_with_policy( endpoint, kv_cache_block_size, - selector, + SelectionPolicySource::Registry, kv_router_config, prefill_load_estimator, worker_role, @@ -2004,23 +1999,20 @@ impl ModelManager { .await } - /// Construct a managed KV router with a selector resolved by the routing host at startup. + /// Construct a managed KV router whose partition runs the given selection policy. #[allow(clippy::too_many_arguments)] - pub async fn managed_kv_router_for_with_selector( + pub async fn managed_kv_router_for_with_policy( &self, endpoint: &Endpoint, kv_cache_block_size: u32, - selector: Sel, + policy: SelectionPolicySource, kv_router_config: Option, prefill_load_estimator: Option>, worker_role: Option, metric_worker_type: &'static str, model_name: Option, is_eagle: bool, - ) -> anyhow::Result> - where - Sel: WorkerSelector + Send + 'static, - { + ) -> anyhow::Result { let client = endpoint.client().await?; let source = crate::kv_router::RouterLoadSource::from_worker_role_or_metric( worker_role, @@ -2035,10 +2027,10 @@ impl ModelManager { ) .await?; let router = self - .kv_chooser_for_with_selector_and_client( + .kv_chooser_for_with_policy_and_client( client, kv_cache_block_size, - selector, + policy, kv_router_config, prefill_load_estimator, worker_role, @@ -2053,11 +2045,11 @@ impl ModelManager { } #[allow(clippy::too_many_arguments)] - pub async fn kv_chooser_for_with_selector_and_client( + pub async fn kv_chooser_for_with_policy_and_client( &self, client: Client, kv_cache_block_size: u32, - selector: Sel, + policy: SelectionPolicySource, kv_router_config: Option, prefill_load_estimator: Option>, worker_role: Option, @@ -2066,10 +2058,7 @@ impl ModelManager { is_eagle: bool, scheduler_load: crate::kv_router::SchedulerLoadSender, cancellation_token: CancellationToken, - ) -> anyhow::Result>> - where - Sel: WorkerSelector + Send + 'static, - { + ) -> anyhow::Result> { let endpoint = client.endpoint.clone(); let lora_domain = self.lora_domain(&endpoint.id()); @@ -2097,39 +2086,44 @@ impl ModelManager { // Get of create runtime config watcher for this endpoint let workers_with_configs = self.get_or_create_runtime_config_watcher(&endpoint).await?; - // A selector that does not consume cache input must not create a shared-cache client or + let effective_kv_router_config = kv_router_config.clone().unwrap_or_default(); + let worker_type = worker_role.unwrap_or(WorkerType::Aggregated); + let policy = + policy.resolve(&effective_kv_router_config, worker_type, metric_worker_type)?; + let required_worker_inputs = crate::kv_router::policy_worker_inputs( + &policy, + &effective_kv_router_config, + worker_type, + model_name.as_deref(), + ); + // A policy that does not consume cache input must not create a shared-cache client or // subscribe to its updates. - let shared_cache: Option> = if selector - .required_worker_inputs() - .contains(WorkerInputs::CACHE) - { - match kv_router_config - .as_ref() - .map(|c| c.shared_cache_type) - .unwrap_or_default() - { - dynamo_kv_router::SharedCacheType::None => None, - dynamo_kv_router::SharedCacheType::Hicache => { - let worker_component_name = &endpoint.id().component; - tracing::info!( - worker_component = worker_component_name, - "Using HiCache shared KV cache" - ); - Some(Box::new( - self.hicache_cache_for(&endpoint, workers_with_configs.clone()), - )) + let shared_cache: Option> = + if required_worker_inputs.contains(WorkerInputs::CACHE) { + match kv_router_config + .as_ref() + .map(|c| c.shared_cache_type) + .unwrap_or_default() + { + dynamo_kv_router::SharedCacheType::None => None, + dynamo_kv_router::SharedCacheType::Hicache => { + let worker_component_name = &endpoint.id().component; + tracing::info!( + worker_component = worker_component_name, + "Using HiCache shared KV cache" + ); + Some(Box::new( + self.hicache_cache_for(&endpoint, workers_with_configs.clone()), + )) + } } - } - } else { - None - }; + } else { + None + }; - let effective_kv_router_config = kv_router_config.clone().unwrap_or_default(); let kv_event_source_requirement = KvEventSourceRequirement::derive(worker_role, &effective_kv_router_config); - let cache_required = selector - .required_worker_inputs() - .contains(WorkerInputs::CACHE) + let cache_required = required_worker_inputs.contains(WorkerInputs::CACHE) || effective_kv_router_config.serve_indexer || matches!( kv_event_source_requirement, @@ -2153,7 +2147,7 @@ impl ModelManager { workers_with_configs, kv_source_membership, kv_cache_block_size, - selector, + SelectionPolicySource::Factory(policy), kv_router_config, prefill_load_estimator, worker_role, diff --git a/lib/llm/src/discovery/watcher.rs b/lib/llm/src/discovery/watcher.rs index d65417d303ec..3d1ebfaa826e 100644 --- a/lib/llm/src/discovery/watcher.rs +++ b/lib/llm/src/discovery/watcher.rs @@ -10,10 +10,7 @@ use tokio::sync::{Notify, mpsc::Sender}; use anyhow::Context as _; use async_trait::async_trait; -use dynamo_kv_router::{ - DEFAULT_ROUTING_GROUP, PrefillLoadEstimator, RoutingPartitionRef, - selector::{DefaultWorkerSelector, WorkerSelector}, -}; +use dynamo_kv_router::PrefillLoadEstimator; use dynamo_runtime::{ DistributedRuntime, discovery::{DiscoveryInstance, DiscoveryQuery, DiscoveryStream, ModelCardInstanceId}, @@ -32,11 +29,9 @@ use crate::{ entrypoint::{self, ChatEngineFactoryCallback, RouterConfig}, http::service::metrics::Metrics, kv_router::{ - EncoderRouter, PrefillRouter, RouterLoadSource, RoutingLoadContext, WorkerSelectorFactory, - }, - local_model::runtime_config::{ - ModelRuntimeConfig, TokenizerBackend, VLLM_INFERENCE_V1_GENERATE_CAPABILITY, + EncoderRouter, PrefillRouter, RouterLoadSource, RoutingLoadContext, SelectionPolicySource, }, + local_model::runtime_config::{TokenizerBackend, VLLM_INFERENCE_V1_GENERATE_CAPABILITY}, model_card::ModelDeploymentCard, model_type::{ModelInput, ModelType}, preprocessor::{ @@ -183,10 +178,7 @@ pub enum ModelUpdate { Removed(ModelDeploymentCard), } -pub struct ModelWatcher -where - Sel: WorkerSelector, -{ +pub struct ModelWatcher { manager: Arc, drt: DistributedRuntime, router_config: RouterConfig, @@ -210,8 +202,8 @@ where /// Worker capabilities accepted by the frontend's engine-native Generate routes. /// Keep raw pipelines out of default-off and backend-mismatched paths. generate_engine_capabilities: Vec<&'static str>, - worker_selector_factory: WorkerSelectorFactory, - /// Custom selector dispatch cannot infer whether an untyped legacy card is decode or aggregated. + selection_policy: SelectionPolicySource, + /// Custom policy dispatch cannot infer whether an untyped legacy card is decode or aggregated. require_typed_worker_role: bool, } @@ -290,7 +282,7 @@ fn removed_model_cards( .collect() } -impl ModelWatcher { +impl ModelWatcher { #[allow(clippy::too_many_arguments)] pub fn new( runtime: DistributedRuntime, @@ -302,7 +294,7 @@ impl ModelWatcher { prefill_load_estimator: Option>, metrics: Arc, ) -> ModelWatcher { - Self::new_with_worker_selector_factory( + Self::new_with_selection_policy( runtime, model_manager, router_config, @@ -312,22 +304,14 @@ impl ModelWatcher { prefill_load_estimator, metrics, false, - Arc::new(|config, worker_type, _partition| { - DefaultWorkerSelector::new( - Some(config.clone()), - worker_type.default_selector_label(), - ) - }), + SelectionPolicySource::Registry, ) } } -impl ModelWatcher -where - Sel: WorkerSelector + Send + 'static, -{ +impl ModelWatcher { #[allow(clippy::too_many_arguments)] - pub(crate) fn new_with_worker_selector_factory( + pub(crate) fn new_with_selection_policy( runtime: DistributedRuntime, model_manager: Arc, router_config: RouterConfig, @@ -337,7 +321,7 @@ where prefill_load_estimator: Option>, metrics: Arc, require_typed_worker_role: bool, - worker_selector_factory: WorkerSelectorFactory, + selection_policy: SelectionPolicySource, ) -> Self { Self { manager: model_manager, @@ -355,7 +339,7 @@ where tokenizer_backend: None, tokenizer_fallback_enabled: None, generate_engine_capabilities: Vec::new(), - worker_selector_factory, + selection_policy, require_typed_worker_role, } } @@ -593,21 +577,16 @@ where // need the shared chooser in KV mode. let kv_chooser = if router_config.router_mode == RouterMode::KV && needs_preprocessed_routing { - let selector = (self.worker_selector_factory)( - &router_config.kv_router_config, - effective_worker_type(card.worker_type, card.model_type), - RoutingPartitionRef::new(&card.display_name, DEFAULT_ROUTING_GROUP), - ); let mut chooser = self .manager - .kv_chooser_for_with_selector_and_client( + .kv_chooser_for_with_policy_and_client( load_context .as_ref() .expect("routing load context must exist") .client() .clone(), card.kv_cache_block_size, - selector, + self.selection_policy.clone(), Some(router_config.kv_router_config.clone()), self.prefill_load_estimator.clone(), card.worker_type, @@ -644,13 +623,13 @@ where // Fallback only: a prefill worker that declares its own // `router_config` overrides this at activation time. - Some(PrefillRouter::new_with_selector_factory( + Some(PrefillRouter::new_with_selection_policy( None, self.manager.clone(), router_config.router_mode, card.kv_cache_block_size, Some(prefill_config), - self.worker_selector_factory.clone(), + self.selection_policy.clone(), self.prefill_load_estimator.clone(), router_config.session_affinity_ttl_secs, router_config.session_affinity_mode, @@ -683,7 +662,7 @@ where let preprocessed_routing = if needs_preprocessed_routing { Some( - entrypoint::input::build_preprocessed_routing_with_selector( + entrypoint::input::build_preprocessed_routing_with_session_affinity_mode( &client, self.manager.clone(), router_config.router_mode, @@ -1032,10 +1011,7 @@ where } #[async_trait] -impl ControllerHost for ModelWatcher -where - Sel: WorkerSelector + Send + 'static, -{ +impl ControllerHost for ModelWatcher { type Prepared = PreparedWorkerSet; fn normalize( diff --git a/lib/llm/src/entrypoint/input.rs b/lib/llm/src/entrypoint/input.rs index 45e11be46682..088931fda151 100644 --- a/lib/llm/src/entrypoint/input.rs +++ b/lib/llm/src/entrypoint/input.rs @@ -14,7 +14,7 @@ use std::{ }; mod common; -pub(crate) use common::build_preprocessed_routing_with_selector; +pub(crate) use common::build_preprocessed_routing_with_session_affinity_mode; pub use common::{PreprocessedRouting, build_preprocessed_routing}; pub mod endpoint; pub mod grpc; diff --git a/lib/llm/src/entrypoint/input/common.rs b/lib/llm/src/entrypoint/input/common.rs index 2e7f983010bc..0685dcae9269 100644 --- a/lib/llm/src/entrypoint/input/common.rs +++ b/lib/llm/src/entrypoint/input/common.rs @@ -34,10 +34,7 @@ use crate::{ }, }; -use dynamo_kv_router::{ - config::min_initial_workers_from_env, - selector::{DefaultWorkerSelector, WorkerSelector}, -}; +use dynamo_kv_router::config::min_initial_workers_from_env; use dynamo_runtime::{ DistributedRuntime, component::Client, @@ -52,13 +49,10 @@ use std::sync::Arc; type LlmPushRouter = PushRouter>; #[derive(Clone)] -pub struct PreprocessedRouting -where - Sel: WorkerSelector + Send + 'static, -{ +pub struct PreprocessedRouting { backend_engine: ServiceEngine, ManyOut>>, - prefill_router: Arc>, + prefill_router: Arc, encoder_router: Arc, } @@ -108,14 +102,11 @@ async fn wait_for_min_initial_workers( } } -fn router_client( +fn router_client( client: &Client, router_mode: RouterMode, - chooser: Option<&Arc>>, -) -> anyhow::Result -where - Sel: WorkerSelector + Send + 'static, -{ + chooser: Option<&Arc>, +) -> anyhow::Result { if router_mode == RouterMode::KV { let Some(chooser) = chooser else { anyhow::bail!("RouterMode::KV requires KVRouter to not be null"); @@ -160,19 +151,16 @@ fn validate_router_mode_for_lora( } #[expect(clippy::too_many_arguments)] -fn preprocessed_backend_engine( +fn preprocessed_backend_engine( router: LlmPushRouter, router_mode: RouterMode, - chooser: Option>>, + chooser: Option>, model_manager: &Arc, endpoint_id: &dynamo_runtime::protocols::EndpointId, affinity: Option, session_affinity_mode: SessionAffinityMode, load_context: Arc, -) -> anyhow::Result>> -where - Sel: WorkerSelector + Send + 'static, -{ +) -> anyhow::Result> { // Reject LoRA + unsupported-mode combinations up front (single source of truth, shared with // the fail-fast check in `build_preprocessed_routing`). After this, the Direct and advanced // arms below are only reached with LoRA serving disabled. @@ -199,7 +187,7 @@ where let lora = model_manager .lora_filter_for(endpoint_id) .map(|filter| (filter, model_manager.lora_load_estimator_for(endpoint_id))); - Arc::new(RoutingHost::::new_builtin_with_capabilities( + Arc::new(RoutingHost::new_builtin_with_capabilities( router, load_context, affinity, @@ -224,7 +212,7 @@ pub async fn build_preprocessed_routing( enable_multimodal_cache_indexer: bool, session_affinity_ttl_secs: Option, ) -> anyhow::Result { - build_preprocessed_routing_with_selector( + build_preprocessed_routing_with_session_affinity_mode( client, model_manager, router_mode, @@ -240,21 +228,18 @@ pub async fn build_preprocessed_routing( } #[allow(clippy::too_many_arguments)] -pub(crate) async fn build_preprocessed_routing_with_selector( +pub(crate) async fn build_preprocessed_routing_with_session_affinity_mode( client: &Client, model_manager: Arc, router_mode: RouterMode, load_context: Arc, - chooser: Option>>, - prefill_chooser: Option>>, + chooser: Option>, + prefill_chooser: Option>, encoder_chooser: Option>, enable_multimodal_cache_indexer: bool, session_affinity_ttl_secs: Option, session_affinity_mode: SessionAffinityMode, -) -> anyhow::Result> -where - Sel: WorkerSelector + Send + 'static, -{ +) -> anyhow::Result { // Fail fast on an unsupported LoRA + router-mode combination BEFORE waiting for the initial // worker set, so a misconfiguration surfaces immediately at startup rather than after the // (possibly long) DYN_ROUTER_MIN_INITIAL_WORKERS wait. @@ -302,7 +287,7 @@ where RouterRequestMetrics::from_component(client.endpoint.component()); let prefill_router = prefill_chooser.unwrap_or_else(|| { - PrefillRouter::::disabled_with_selector( + PrefillRouter::disabled_with_session_affinity_mode( model_manager.clone(), router_mode, session_affinity_ttl_secs, @@ -483,10 +468,7 @@ where .link_terminal(frontend)?) } -impl PreprocessedRouting -where - Sel: WorkerSelector + Send + 'static, -{ +impl PreprocessedRouting { /// The normal way to build an inference pipeline. Connect this directly to HTTP layer. pub fn build_pipeline( &self, diff --git a/lib/llm/src/entrypoint/input/http.rs b/lib/llm/src/entrypoint/input/http.rs index 6a23771bbd66..df6c07b39332 100644 --- a/lib/llm/src/entrypoint/input/http.rs +++ b/lib/llm/src/entrypoint/input/http.rs @@ -13,8 +13,8 @@ use crate::{ FrontendRouteExtension, service_v2::{self, HttpService}, }, - kv_router::WorkerSelectorFactory, - local_model::runtime_config::{ModelRuntimeConfig, TokenizerBackend}, + kv_router::SelectionPolicySource, + local_model::runtime_config::TokenizerBackend, namespace::NamespaceFilter, types::openai::{ chat_completions::{NvCreateChatCompletionRequest, NvCreateChatCompletionStreamResponse}, @@ -22,20 +22,20 @@ use crate::{ }, }; use dynamo_kv_router::{ - KvRouterConfig, RoutingPartitionRef, WorkerSelectionPolicy, WorkerType, - selector::{DefaultWorkerSelector, WorkerSelector}, + KvRouterConfig, RoutingPartitionRef, WorkerSelectionPolicy, WorkerSelectionPolicyFactory, + WorkerType, }; use dynamo_runtime::DistributedRuntime; use dynamo_runtime::metrics::MetricsHierarchy; /// Dynamo's complete discovery-backed HTTP frontend. /// -/// The default frontend uses [`DefaultWorkerSelector`]. A statically linked external crate can -/// replace only worker selection with [`Self::worker_selection_policy_factory`]. +/// The default frontend resolves worker selection from its configuration. A statically linked +/// external crate can replace only worker selection with [`Self::worker_selection_policy_factory`]. #[derive(Default)] pub struct HttpFrontend { frontend_route_extensions: Vec, - worker_selection_policy_factory: Option>, + worker_selection_policy_factory: Option, } impl HttpFrontend { @@ -83,33 +83,19 @@ impl HttpFrontend { super::initialize_input(&distributed_runtime, &engine_config).await; - match self.worker_selection_policy_factory { - Some(factory) => { - run_with_worker_selector_factory( - distributed_runtime, - engine_config, - self.frontend_route_extensions, - true, - factory, - ) - .await - } - None => { - run_with_worker_selector_factory( - distributed_runtime, - engine_config, - self.frontend_route_extensions, - false, - Arc::new(|config, worker_type, _partition| { - DefaultWorkerSelector::new( - Some(config.clone()), - worker_type.default_selector_label(), - ) - }), - ) - .await - } - } + let (require_typed_worker_role, selection_policy) = + match self.worker_selection_policy_factory { + Some(factory) => (true, SelectionPolicySource::Factory(factory)), + None => (false, SelectionPolicySource::Registry), + }; + run_with_selection_policy( + distributed_runtime, + engine_config, + self.frontend_route_extensions, + require_typed_worker_role, + selection_policy, + ) + .await } } @@ -135,16 +121,13 @@ pub async fn run_with_frontend_route_extensions( .await } -async fn run_with_worker_selector_factory( +async fn run_with_selection_policy( distributed_runtime: DistributedRuntime, engine_config: EngineConfig, frontend_route_extensions: Vec, require_typed_worker_role: bool, - worker_selector_factory: WorkerSelectorFactory, -) -> anyhow::Result<()> -where - Sel: WorkerSelector + Send + 'static, -{ + selection_policy: SelectionPolicySource, +) -> anyhow::Result<()> { let local_model = engine_config.local_model(); let mut http_service_builder = match (local_model.tls_cert_path(), local_model.tls_key_path()) { (Some(tls_cert_path), Some(tls_key_path)) => { @@ -232,7 +215,7 @@ where model.runtime_config().tokenizer_fallback_enabled, generate_engine_capabilities, require_typed_worker_role, - worker_selector_factory.clone(), + selection_policy.clone(), ) .await?; http_service @@ -304,7 +287,7 @@ fn enable_in_process_model_endpoints(http_service: &HttpService) -> anyhow::Resu /// Spawns a task that watches for new models in store, /// and registers them with the ModelManager so that the HTTP service can use them. #[allow(clippy::too_many_arguments)] -async fn run_watcher( +async fn run_watcher( runtime: DistributedRuntime, model_manager: Arc, router_config: RouterConfig, @@ -320,11 +303,8 @@ async fn run_watcher( tokenizer_fallback_enabled: Option, generate_engine_capabilities: Vec<&'static str>, require_typed_worker_role: bool, - worker_selector_factory: WorkerSelectorFactory, -) -> anyhow::Result<()> -where - Sel: WorkerSelector + Send + 'static, -{ + selection_policy: SelectionPolicySource, +) -> anyhow::Result<()> { // Start the LoRA allocation controller when LoRA serving is enabled. The // controller itself is additionally gated on the allocation config // (DYN_LORA_ALLOCATION_ENABLED) inside `start_lora_controller`. @@ -333,7 +313,7 @@ where let _controller_handle = model_manager.start_lora_controller(cancel_token); } - let mut watch_obj = ModelWatcher::new_with_worker_selector_factory( + let mut watch_obj = ModelWatcher::new_with_selection_policy( runtime.clone(), model_manager, router_config, @@ -343,7 +323,7 @@ where prefill_load_estimator, metrics.clone(), require_typed_worker_role, - worker_selector_factory, + selection_policy, ); watch_obj.set_local_model_path(local_model_path); watch_obj.set_tokenizer_backend(tokenizer_backend); diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index 4abf34b7aa3f..14fb21efc081 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -12,6 +12,7 @@ use anyhow::Result; use dynamo_kv_router::{ DEFAULT_ROUTING_GROUP, KvSchedulerError, PrefillLoadEstimator, RoutingPartitionRef, SharedKvCache, TrackingHashAlgorithm, TrackingHashContext, TrackingHashScope, + WorkerSelectionPolicy, WorkerSelectionPolicyFactory, config::{KvRouterConfig, RouterConfigOverride, min_initial_workers_from_env}, indexer::{ ApproximateLruIncarnation, ApproximateLruStats, KvRouterError, RoutingDecisionHashes, @@ -85,7 +86,7 @@ pub use routing_load::{ use crate::{ discovery::{KvSourceMembershipWatch, RuntimeConfigWatch}, kv_router::{ - scheduler::{DefaultWorkerSelector, PotentialLoad}, + scheduler::PotentialLoad, sequence::{SequenceError, SequenceRequest}, }, local_model::runtime_config::ModelRuntimeConfig, @@ -95,8 +96,58 @@ use route_lookup::{ TieredLookupOptions, TieredLookupResult, query_tiered_matches, split_retained_block_hashes, }; -pub(crate) type WorkerSelectorFactory = - Arc Fn(&KvRouterConfig, WorkerType, RoutingPartitionRef<'a>) -> Sel + Send + Sync>; +/// Where a `KvRouter` gets the worker-selection policy its partition runs. +#[derive(Clone, Default)] +pub enum SelectionPolicySource { + /// The linked policy `KvRouterConfig` names for the worker role, resolved + /// from the installed [`worker_selection_policy_registry`]; Dynamo's + /// default scorer and picker when it names none. + #[default] + Registry, + /// This factory, called once per routing partition. + Factory(WorkerSelectionPolicyFactory), +} + +impl SelectionPolicySource { + /// Resolve to the factory the partition will call. `label` is the worker + /// pool name the default policy logs under. + pub fn resolve( + &self, + config: &KvRouterConfig, + worker_type: WorkerType, + label: &'static str, + ) -> Result { + match self { + Self::Factory(factory) => Ok(factory.clone()), + Self::Registry => Ok( + match worker_selection_policy_registry() + .resolve_for_worker_type(config, worker_type)? + { + Some(factory) => factory, + None => Arc::new(move |config: &KvRouterConfig, _worker_type, _partition| { + WorkerSelectionPolicy::default(config.clone(), label) + }), + }, + ), + } + } +} + +/// The optional worker inputs one instance of `factory`'s policy consumes. +pub(crate) fn policy_worker_inputs( + factory: &WorkerSelectionPolicyFactory, + config: &KvRouterConfig, + worker_type: WorkerType, + model_name: Option<&str>, +) -> WorkerInputs { + let partition = RoutingPartitionRef::new( + model_name.unwrap_or(embedded::DEFAULT_MODEL_NAME), + DEFAULT_ROUTING_GROUP, + ); + dynamo_kv_router::selector::WorkerSelector::::required_worker_inputs( + &factory(config, worker_type, partition), + ) +} #[derive(Clone, Copy)] struct ApproximateLruRankRegistration { @@ -531,15 +582,9 @@ pub fn router_discovery_query(namespace: String, component: String) -> Discovery /// A KvRouter only decides which worker you should use. It doesn't send you there. /// TODO: Rename this to indicate it only selects a worker, it does not route. -pub struct KvRouter -where - Sel: dynamo_kv_router::selector::WorkerSelector, -{ +pub struct KvRouter { indexer: Indexer, scheduler: embedded::EmbeddedSelection, - /// `Sel` only names the selector type the caller injected; the partition - /// owns it. `fn() -> Sel` keeps the router `Send + Sync` regardless of `Sel`. - _selector: std::marker::PhantomData Sel>, required_worker_inputs: dynamo_kv_router::selector::WorkerInputs, workers_with_configs: RuntimeConfigWatch, block_size: u32, @@ -580,10 +625,7 @@ fn resolve_tracking_model_name( Ok(model_name.unwrap_or_default().to_owned()) } -impl KvRouter -where - Sel: dynamo_kv_router::selector::WorkerSelector + Send + 'static, -{ +impl KvRouter { #[allow(clippy::too_many_arguments)] pub async fn new( endpoint: Endpoint, @@ -591,7 +633,7 @@ where workers_with_configs: RuntimeConfigWatch, kv_source_membership: Option, block_size: u32, - selector: Sel, + policy: SelectionPolicySource, kv_router_config: Option, prefill_load_estimator: Option>, metric_worker_type: &'static str, @@ -606,7 +648,7 @@ where workers_with_configs, kv_source_membership, block_size, - selector, + policy, kv_router_config, prefill_load_estimator, None, @@ -626,7 +668,7 @@ where workers_with_configs: RuntimeConfigWatch, kv_source_membership: Option, block_size: u32, - selector: Sel, + policy: SelectionPolicySource, kv_router_config: Option, prefill_load_estimator: Option>, worker_role: Option, @@ -646,7 +688,7 @@ where workers_with_configs, kv_source_membership, block_size, - selector, + policy, kv_router_config, prefill_load_estimator, worker_role, @@ -668,7 +710,7 @@ where workers_with_configs: RuntimeConfigWatch, kv_source_membership: Option, block_size: u32, - selector: Sel, + policy: SelectionPolicySource, kv_router_config: Option, prefill_load_estimator: Option>, worker_role: Option, @@ -680,7 +722,16 @@ where scheduler_load: SchedulerLoadSender, parent_token: CancellationToken, ) -> Result { - let required_worker_inputs = selector.required_worker_inputs(); + let kv_router_config = kv_router_config.unwrap_or_default(); + kv_router_config.validate().map_err(anyhow::Error::msg)?; + let worker_type = worker_role.unwrap_or(WorkerType::Aggregated); + let policy_factory = policy.resolve(&kv_router_config, worker_type, metric_worker_type)?; + let required_worker_inputs = policy_worker_inputs( + &policy_factory, + &kv_router_config, + worker_type, + model_name.as_deref(), + ); // ModelManager gates client construction as well, but preserve the capability boundary for // direct KvRouter callers. let shared_cache = if required_worker_inputs.contains(WorkerInputs::CACHE) { @@ -688,8 +739,6 @@ where } else { None }; - let kv_router_config = kv_router_config.unwrap_or_default(); - kv_router_config.validate().map_err(anyhow::Error::msg)?; let tracking_hash = TrackingHashContext::from_config(&kv_router_config)?; let tracking_model_name = resolve_tracking_model_name(tracking_hash.algorithm(), model_name.as_deref())?; @@ -771,19 +820,6 @@ where let available_worker_provider: WorkerAvailabilityProvider = Arc::new(move || client_for_availability.available_instance_ids()); - // Scheduling and active-sequence accounting run on a selection-service - // partition. A caller-injected `WorkerSelector` type (anything other than - // the default selector or a registry `WorkerSelectionPolicy`) runs inside - // the partition through a delegating policy. - let sel = std::any::TypeId::of::(); - let custom_selector = sel - != std::any::TypeId::of::() - && sel - != std::any::TypeId::of::< - dynamo_kv_router::scheduling::selector::WorkerSelectionPolicy, - >(); - let policy_override = custom_selector - .then(|| embedded::delegating_policy_factory(selector, metric_worker_type)); let overlap_refresh: Option> = indexer .supports_overlap_refresh() @@ -803,7 +839,7 @@ where scheduler_load, endpoint: endpoint.clone(), router_id: endpoint.drt().discovery().instance_id(), - policy_override, + policy_factory, }, workers_with_configs.clone(), cancellation_token.child_token(), @@ -870,7 +906,6 @@ where Ok(Self { indexer, scheduler, - _selector: std::marker::PhantomData, required_worker_inputs, workers_with_configs, block_size, @@ -2262,11 +2297,7 @@ where // NOTE: KVRouter works like a PushRouter, // but without the reverse proxy functionality, but based on the RouterRequest contract #[async_trait] -impl AsyncEngine, ManyOut>, Error> - for KvRouter -where - Sel: dynamo_kv_router::selector::WorkerSelector + Send + 'static, -{ +impl AsyncEngine, ManyOut>, Error> for KvRouter { async fn generate( &self, request: SingleIn, @@ -2382,10 +2413,7 @@ where } } -impl Drop for KvRouter -where - Sel: dynamo_kv_router::selector::WorkerSelector, -{ +impl Drop for KvRouter { fn drop(&mut self) { tracing::info!("Dropping KvRouter - cancelling background tasks"); self.cancellation_token.cancel(); @@ -2399,7 +2427,7 @@ mod tests { use async_trait::async_trait; use dynamo_kv_router::{ - WorkerSelectionInput, + WorkerSelectionPolicyError, identity::{ CacheOwnerId, CacheSemanticsId, DcId, IdentitySource, IndexerDomainId, PoolId, RoutingScopeId, StableDpSlotId, @@ -2416,6 +2444,9 @@ mod tests { use crate::kv_router::scheduler::KvSchedulerError; use crate::local_model::runtime_config::ModelRuntimeConfig; + use dynamo_kv_router::selector::{ + WorkerCandidate, WorkerFilter, WorkerInputView, WorkerPicker, WorkerSelectionContext, + }; #[test] fn all_filtered_workers_map_to_unavailable() { @@ -2630,70 +2661,90 @@ mod tests { } } - struct InspectingSelector { - expected_hits: Option, - selected_worker: WorkerWithDpRank, + /// Picks a fixed worker after checking the shared-cache hits the router + /// projected into the candidate table. + struct FixedPicker { + expected_shared_blocks: Option, + worker: WorkerWithDpRank, } - impl dynamo_kv_router::selector::WorkerSelector for InspectingSelector { + impl WorkerPicker for FixedPicker { fn required_worker_inputs(&self) -> WorkerInputs { WorkerInputs::CACHE | WorkerInputs::LOAD } - fn select_worker( - &self, - input: WorkerSelectionInput<'_, ModelRuntimeConfig>, - ) -> Result { - let (_workers, request, _eligibility, block_size) = input.into_configured()?; - let observed_hits = request - .shared_cache_hits - .as_ref() - .map(|hits| hits.total_hits); - assert_eq!(observed_hits, self.expected_hits); - - Ok(dynamo_kv_router::protocols::WorkerSelectionResult { - worker: self.selected_worker, - required_blocks: request.isl_tokens.div_ceil(block_size as usize) as u64, - effective_overlap_blocks: 0.0, - cached_tokens: 0, - potential_decode_blocks: request - .worker_load_for(self.selected_worker) - .potential_decode_blocks() - .saturating_add(request.isl_tokens.div_ceil(block_size as usize)), - }) + fn pick( + &mut self, + _context: &WorkerSelectionContext<'_>, + input: WorkerInputView<'_>, + ) -> Result { + let row = input + .candidates() + .iter() + .position(|candidate| candidate.worker() == self.worker) + .ok_or_else(|| WorkerSelectionPolicyError::failed("fixed worker not eligible"))?; + let shared = input + .cache() + .map(|cache| cache[row].shared_beyond_device_blocks()) + .filter(|blocks| *blocks > 0); + assert_eq!(shared, self.expected_shared_blocks); + Ok(row) } } - struct OverloadedSelector; - - impl dynamo_kv_router::selector::WorkerSelector for OverloadedSelector { - fn required_worker_inputs(&self) -> WorkerInputs { - WorkerInputs::NONE - } + struct RejectAll; - fn select_worker( - &self, - _input: WorkerSelectionInput<'_, ModelRuntimeConfig>, - ) -> Result { - Err(KvSchedulerError::AllEligibleWorkersOverloaded) + impl WorkerFilter for RejectAll { + fn keep( + &mut self, + _context: &WorkerSelectionContext<'_>, + _candidate: &WorkerCandidate, + ) -> Result { + Ok(false) } } - struct LoadOnlySelector; + struct LoadOnlyPicker; - impl dynamo_kv_router::selector::WorkerSelector for LoadOnlySelector { + impl WorkerPicker for LoadOnlyPicker { fn required_worker_inputs(&self) -> WorkerInputs { WorkerInputs::LOAD } - fn select_worker( - &self, - _input: WorkerSelectionInput<'_, ModelRuntimeConfig>, - ) -> Result { + fn pick( + &mut self, + _context: &WorkerSelectionContext<'_>, + _input: WorkerInputView<'_>, + ) -> Result { unreachable!("capability construction test does not select a worker") } } + fn fixed_policy( + expected_shared_blocks: Option, + worker: WorkerWithDpRank, + ) -> SelectionPolicySource { + SelectionPolicySource::Factory(Arc::new(move |config: &KvRouterConfig, _, _| { + WorkerSelectionPolicy::new( + config.clone(), + "decode", + Vec::new(), + Box::new(FixedPicker { + expected_shared_blocks, + worker, + }), + ) + })) + } + + fn picker_policy( + picker: impl Fn() -> Box + Send + Sync + 'static, + ) -> SelectionPolicySource { + SelectionPolicySource::Factory(Arc::new(move |config: &KvRouterConfig, _, _| { + WorkerSelectionPolicy::new(config.clone(), "decode", Vec::new(), picker()) + })) + } + async fn make_test_component(name: &str) -> dynamo_runtime::component::Component { let runtime = Runtime::from_current().unwrap(); let drt = DistributedRuntime::new(runtime, DistributedConfig::process_local()) @@ -2722,7 +2773,7 @@ mod tests { workers, None, 16, - DefaultWorkerSelector::new(Some(config.clone()), "decode"), + SelectionPolicySource::Registry, Some(config), None, worker_role, @@ -2771,7 +2822,7 @@ mod tests { workers, None, 16, - LoadOnlySelector, + picker_policy(|| Box::new(LoadOnlyPicker)), Some(config), None, Some(WorkerType::Prefill), @@ -2798,15 +2849,10 @@ mod tests { } async fn make_test_router_with_workers( - selector: impl dynamo_kv_router::selector::WorkerSelector - + Send - + Sync - + 'static, + policy: SelectionPolicySource, shared_cache: Option>, workers: HashMap, - ) -> KvRouter< - impl dynamo_kv_router::selector::WorkerSelector + Send + Sync + 'static, - > { + ) -> KvRouter { let component = make_test_component("shared-cache-router").await; let endpoint = component.endpoint("backend"); let client = endpoint.client().await.unwrap(); @@ -2828,7 +2874,7 @@ mod tests { rx, None, 2, - selector, + policy, Some(config), None, None, @@ -2843,18 +2889,13 @@ mod tests { } async fn make_test_router( - selector: impl dynamo_kv_router::selector::WorkerSelector - + Send - + Sync - + 'static, + policy: SelectionPolicySource, shared_cache: Option>, - ) -> KvRouter< - impl dynamo_kv_router::selector::WorkerSelector + Send + Sync + 'static, - > { + ) -> KvRouter { let mut workers = HashMap::new(); workers.insert(0, ModelRuntimeConfig::default()); workers.insert(1, ModelRuntimeConfig::default()); - make_test_router_with_workers(selector, shared_cache, workers).await + make_test_router_with_workers(policy, shared_cache, workers).await } fn transfer_hint_runtime_config(endpoint: Option<&str>) -> ModelRuntimeConfig { @@ -2935,10 +2976,7 @@ mod tests { ]), ); let router = make_test_router_with_workers( - InspectingSelector { - expected_hits: None, - selected_worker: WorkerWithDpRank::new(7, 0), - }, + fixed_policy(None, WorkerWithDpRank::new(7, 0)), None, workers, ) @@ -2987,10 +3025,7 @@ mod tests { ); workers.insert(8, transfer_hint_runtime_config(source_endpoint)); let router = make_test_router_with_workers( - InspectingSelector { - expected_hits: None, - selected_worker: WorkerWithDpRank::new(7, 0), - }, + fixed_policy(None, WorkerWithDpRank::new(7, 0)), None, workers, ) @@ -3026,10 +3061,7 @@ mod tests { transfer_hint_runtime_config_with_worker_type(Some("tcp://127.0.0.1:23281"), "decode"), ); let router = make_test_router_with_workers( - InspectingSelector { - expected_hits: None, - selected_worker: WorkerWithDpRank::new(7, 0), - }, + fixed_policy(None, WorkerWithDpRank::new(7, 0)), None, workers, ) @@ -3069,10 +3101,7 @@ mod tests { transfer_hint_runtime_config_with_worker_type(Some("tcp://127.0.0.1:23283"), "decode"), ); let router = make_test_router_with_workers( - InspectingSelector { - expected_hits: None, - selected_worker: WorkerWithDpRank::new(7, 0), - }, + fixed_policy(None, WorkerWithDpRank::new(7, 0)), None, workers, ) @@ -3128,15 +3157,7 @@ mod tests { transfer_hint_runtime_config(Some("tcp://stale-worker-endpoint:23280")); stale_source_config.kv_event_source_mode = Some("state_agent_v2".to_string()); workers.insert(8, stale_source_config); - let router = make_test_router_with_workers( - InspectingSelector { - expected_hits: None, - selected_worker: target, - }, - None, - workers, - ) - .await; + let router = make_test_router_with_workers(fixed_policy(None, target), None, workers).await; let owner = router_hint_cache_owner(); let owner_key = ResidencyOwner::cache_owner(owner).compact_key(); let candidates = KvTransferCandidates { @@ -3176,10 +3197,7 @@ mod tests { #[tokio::test] async fn test_find_best_match_passes_shared_cache_hits_to_scheduler() { let router = make_test_router( - InspectingSelector { - expected_hits: Some(2), - selected_worker: WorkerWithDpRank::from_worker_id(1), - }, + fixed_policy(Some(2), WorkerWithDpRank::from_worker_id(1)), Some(Box::new(FakeSharedCache { #[allow(clippy::single_range_in_vec_init)] hits: Some(dynamo_kv_router::protocols::SharedCacheHits::from_ranges( @@ -3215,10 +3233,7 @@ mod tests { #[tokio::test] async fn test_find_best_match_ignores_shared_cache_errors() { let router = make_test_router( - InspectingSelector { - expected_hits: None, - selected_worker: WorkerWithDpRank::from_worker_id(0), - }, + fixed_policy(None, WorkerWithDpRank::from_worker_id(0)), Some(Box::new(FakeSharedCache { hits: None, should_error: true, @@ -3249,8 +3264,17 @@ mod tests { } #[tokio::test] - async fn test_find_best_match_maps_overload_to_resource_exhausted() { - let router = make_test_router(OverloadedSelector, None).await; + async fn test_find_best_match_maps_filtered_workers_to_unavailable() { + let policy = SelectionPolicySource::Factory(Arc::new(|config: &KvRouterConfig, _, _| { + WorkerSelectionPolicy::new_with_filters( + config.clone(), + "decode", + vec![Box::new(RejectAll)], + Vec::new(), + Box::new(LoadOnlyPicker), + ) + })); + let router = make_test_router(policy, None).await; let err = router .find_best_match( @@ -3272,22 +3296,15 @@ mod tests { assert!(dynamo_runtime::error::match_error_chain( err.as_ref(), - &[dynamo_runtime::error::ErrorType::ResourceExhausted], + &[dynamo_runtime::error::ErrorType::Unavailable], &[] )); - assert!( - err.to_string() - .contains("all eligible workers are overloaded") - ); } #[tokio::test] async fn test_find_best_match_details_returns_routing_hashes_when_requested() { let router = make_test_router( - InspectingSelector { - expected_hits: None, - selected_worker: WorkerWithDpRank::from_worker_id(0), - }, + fixed_policy(None, WorkerWithDpRank::from_worker_id(0)), None, ) .await; @@ -3339,10 +3356,7 @@ mod tests { #[tokio::test] async fn test_find_best_match_details_omits_routing_hashes_when_not_requested() { let router = make_test_router( - InspectingSelector { - expected_hits: None, - selected_worker: WorkerWithDpRank::from_worker_id(0), - }, + fixed_policy(None, WorkerWithDpRank::from_worker_id(0)), None, ) .await; @@ -3376,10 +3390,7 @@ mod tests { #[tokio::test] async fn test_get_overlap_scores_returns_tiered_rows_and_shared_hits() { let router = make_test_router( - InspectingSelector { - expected_hits: None, - selected_worker: WorkerWithDpRank::from_worker_id(0), - }, + fixed_policy(None, WorkerWithDpRank::from_worker_id(0)), Some(Box::new(FakeSharedCache { #[allow(clippy::single_range_in_vec_init)] hits: Some(dynamo_kv_router::protocols::SharedCacheHits::from_ranges( diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs index c055970088af..4a74e39018ea 100644 --- a/lib/llm/src/kv_router/embedded.rs +++ b/lib/llm/src/kv_router/embedded.rs @@ -22,7 +22,6 @@ use dynamo_kv_router::identity::RoutingPartitionId; use dynamo_kv_router::indexer::TieredMatchProvider; use dynamo_kv_router::protocols::{WorkerConfigLike, WorkerId, WorkerWithDpRank}; use dynamo_kv_router::scheduling::queue::{SchedulerBookingCleanup, SchedulerBookingDescriptor}; -use dynamo_kv_router::scheduling::selector::WorkerSelectionPolicy; use dynamo_kv_router::scheduling::{ AdmittedSchedulingResponse, AdvisorySchedulingResponse, AttemptId, KvSchedulerError, OverloadedWorkerProvider, PotentialLoad, ScheduleRequest, WorkerAvailabilityProvider, @@ -61,128 +60,13 @@ pub(crate) struct EmbeddedSelectionArgs { pub endpoint: dynamo_runtime::component::Endpoint, /// This router replica's id (ignored on receipt of its own events). pub router_id: u64, - /// Use this policy instead of resolving one from the registry (a - /// caller-injected `WorkerSelector`, wrapped by [`delegating_policy_factory`]). - pub policy_override: Option, + /// Builds the partition's worker-selection policy (see + /// `SelectionPolicySource::resolve`). + pub policy_factory: WorkerSelectionPolicyFactory, } -/// Runs a caller-owned `WorkerSelector` inside a partition -/// policy: the partition's worker configs are projected onto -/// `ModelRuntimeConfig` per selection. Serialized through a mutex because the -/// selector is only required to be `Send`. -struct DelegatingSelector { - inner: std::sync::Mutex, - inputs: dynamo_kv_router::scheduling::selector::WorkerInputs, - exclusive_affinity: bool, -} - -impl dynamo_kv_router::scheduling::selector::DelegatedWorkerSelector - for DelegatingSelector -where - Sel: dynamo_kv_router::selector::WorkerSelector + Send + 'static, -{ - fn required_worker_inputs(&self) -> dynamo_kv_router::scheduling::selector::WorkerInputs { - self.inputs - } - - fn uses_exclusive_affinity_target(&self) -> bool { - self.exclusive_affinity - } - - fn select_worker( - &self, - workers: &HashMap, - request: &dynamo_kv_router::scheduling::SchedulingRequest, - eligibility: dynamo_kv_router::scheduling::RoutingEligibility<'_>, - block_size: u32, - ) -> std::result::Result< - dynamo_kv_router::protocols::WorkerSelectionResult, - dynamo_kv_router::scheduling::KvSchedulerError, - > { - let projected: HashMap = workers - .iter() - .map(|(id, config)| { - ( - *id, - ModelRuntimeConfig { - data_parallel_start_rank: config.data_parallel_start_rank, - data_parallel_size: config.data_parallel_size, - max_num_batched_tokens: config.max_num_batched_tokens, - total_kv_blocks: config.total_kv_blocks, - taints: config.taints.clone(), - stable_routing_id: config.stable_routing_id.clone(), - topology_domains: config.topology_domains.clone().unwrap_or_default(), - kv_transfer_domain: config.kv_transfer_domain.clone(), - kv_transfer_enforcement: config.kv_transfer_enforcement, - kv_transfer_preferred_weight: config.kv_transfer_preferred_weight, - ..Default::default() - }, - ) - }) - .collect(); - let inner = self - .inner - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - inner.select_worker( - dynamo_kv_router::scheduling::selector::WorkerSelectionInput::Configured { - workers: &projected, - request, - eligibility, - block_size, - }, - ) - } -} - -struct SharedDelegate(Arc); - -impl dynamo_kv_router::scheduling::selector::DelegatedWorkerSelector for SharedDelegate { - fn required_worker_inputs(&self) -> dynamo_kv_router::scheduling::selector::WorkerInputs { - self.0.required_worker_inputs() - } - fn uses_exclusive_affinity_target(&self) -> bool { - self.0.uses_exclusive_affinity_target() - } - fn select_worker( - &self, - workers: &HashMap, - request: &dynamo_kv_router::scheduling::SchedulingRequest, - eligibility: dynamo_kv_router::scheduling::RoutingEligibility<'_>, - block_size: u32, - ) -> std::result::Result< - dynamo_kv_router::protocols::WorkerSelectionResult, - dynamo_kv_router::scheduling::KvSchedulerError, - > { - self.0 - .select_worker(workers, request, eligibility, block_size) - } -} - -/// A policy factory that runs `selector` on every partition of the router. -pub(crate) fn delegating_policy_factory( - selector: Sel, - label: &'static str, -) -> WorkerSelectionPolicyFactory -where - Sel: dynamo_kv_router::selector::WorkerSelector + Send + 'static, -{ - let inputs = selector.required_worker_inputs(); - let exclusive_affinity = selector.uses_exclusive_affinity_target(); - let shared: Arc = - Arc::new(DelegatingSelector { - inner: std::sync::Mutex::new(selector), - inputs, - exclusive_affinity, - }); - Arc::new(move |config: &KvRouterConfig, _worker_type, _partition| { - WorkerSelectionPolicy::delegated( - config.clone(), - label, - Box::new(SharedDelegate(Arc::clone(&shared))), - ) - }) -} +/// Partition model name when the router has none. +pub(crate) const DEFAULT_MODEL_NAME: &str = "default"; /// One `SelectionService` partition driven directly by the router. pub(crate) struct EmbeddedSelection { @@ -232,26 +116,10 @@ impl EmbeddedSelection { let key = RoutingPartitionId::new( args.model_name .clone() - .unwrap_or_else(|| "default".to_string()), + .unwrap_or_else(|| DEFAULT_MODEL_NAME.to_string()), DEFAULT_ROUTING_GROUP, ); - // Same selector the runtime path builds: the registry-resolved policy - // when router configuration names one, else Dynamo's default scorer - // and picker under the router's metric worker label. - let label = args.metric_worker_type; - let policy_factory: WorkerSelectionPolicyFactory = match args.policy_override { - Some(factory) => factory, - None => match worker_selection_policy_registry() - .resolve_for_worker_type(&args.kv_router_config, worker_type)? - { - Some(factory) => factory, - None => Arc::new(move |config: &KvRouterConfig, _worker_type, _partition| { - WorkerSelectionPolicy::default(config.clone(), label) - }), - }, - }; - // Replica sync rides the runtime event plane, as the runtime scheduler's // did; the partition publishes and consumes through host channels. let replica_sync: Option = @@ -276,7 +144,7 @@ impl EmbeddedSelection { worker_type, worker_selection_policy_registry(), ) - .worker_selection_policy_factory(policy_factory) + .worker_selection_policy_factory(args.policy_factory) .indexer_threads(1) .external_kv_events() .host_hooks(SelectionHostHooks { diff --git a/lib/llm/src/kv_router/prefill_router/activation.rs b/lib/llm/src/kv_router/prefill_router/activation.rs index 3bf3e80b3b48..da0779032f6c 100644 --- a/lib/llm/src/kv_router/prefill_router/activation.rs +++ b/lib/llm/src/kv_router/prefill_router/activation.rs @@ -8,10 +8,8 @@ use anyhow::{Context as _, Result}; use tokio::sync::{oneshot, watch}; use dynamo_kv_router::{ - DEFAULT_ROUTING_GROUP, PrefillLoadEstimator, RoutingPartitionRef, - conditional_disagg::make_conditional_disagg_policy, + PrefillLoadEstimator, conditional_disagg::make_conditional_disagg_policy, config::KvRouterConfig, - selector::{DefaultWorkerSelector, WorkerSelector}, }; use dynamo_runtime::{ component::Endpoint, @@ -24,8 +22,7 @@ use dynamo_runtime::{ use super::{PrefillBinding, PrefillBuildContext, PrefillLifecycleState, PrefillRouter}; use crate::{ discovery::{LoadThresholdHandle, ModelManager}, - kv_router::{RouterLoadSource, RoutingHost, RoutingLoadContext, WorkerSelectorFactory}, - local_model::runtime_config::ModelRuntimeConfig, + kv_router::{RouterLoadSource, RoutingHost, RoutingLoadContext, SelectionPolicySource}, model_card::ModelDeploymentCard, protocols::common::{ llm_backend::{LLMEngineOutput, PreprocessedRequest}, @@ -102,14 +99,14 @@ fn resolve_advertisement_from_cards( Ok(advertisement) } -impl PrefillRouter { +impl PrefillRouter { /// Create a disabled prefill router that will never activate (passthrough only) pub fn disabled( model_manager: Arc, decode_router_mode: RouterMode, session_affinity_ttl_secs: Option, ) -> Arc { - Self::disabled_with_selector( + Self::disabled_with_session_affinity_mode( model_manager, decode_router_mode, session_affinity_ttl_secs, @@ -135,18 +132,13 @@ impl PrefillRouter { load_thresholds: LoadThresholdHandle, parent_token: tokio_util::sync::CancellationToken, ) -> Arc { - Self::new_with_selector_factory( + Self::new_with_selection_policy( Some(activation_rx), model_manager, decode_router_mode, kv_cache_block_size, kv_router_config, - Arc::new(|config, worker_type, _partition| { - DefaultWorkerSelector::new( - Some(config.clone()), - worker_type.default_selector_label(), - ) - }), + SelectionPolicySource::Registry, prefill_load_estimator, session_affinity_ttl_secs, session_affinity_mode, @@ -159,11 +151,8 @@ impl PrefillRouter { } } -impl PrefillRouter -where - Sel: WorkerSelector + Send + 'static, -{ - pub(crate) fn disabled_with_selector( +impl PrefillRouter { + pub(crate) fn disabled_with_session_affinity_mode( model_manager: Arc, decode_router_mode: RouterMode, session_affinity_ttl_secs: Option, @@ -174,7 +163,7 @@ where target: parking_lot::Mutex::new(None), target_tx: None, decode_routing_host: std::sync::OnceLock::new(), - worker_selector_factory: None, + selection_policy: None, model_manager, cancel_token: tokio_util::sync::CancellationToken::new(), decode_router_mode, @@ -194,13 +183,13 @@ where } #[expect(clippy::too_many_arguments)] - pub(crate) fn new_with_selector_factory( + pub(crate) fn new_with_selection_policy( activation_rx: Option>, model_manager: Arc, decode_router_mode: RouterMode, kv_cache_block_size: u32, kv_router_config: Option, - worker_selector_factory: WorkerSelectorFactory, + selection_policy: SelectionPolicySource, prefill_load_estimator: Option>, session_affinity_ttl_secs: Option, session_affinity_mode: SessionAffinityMode, @@ -226,7 +215,7 @@ where target: parking_lot::Mutex::new(None), target_tx: Some(target_tx), decode_routing_host: std::sync::OnceLock::new(), - worker_selector_factory: Some(worker_selector_factory), + selection_policy: Some(selection_policy), model_manager: model_manager.clone(), cancel_token: cancel_token.clone(), decode_router_mode, @@ -291,11 +280,11 @@ where } async fn build_binding( - context: &PrefillBuildContext, + context: &PrefillBuildContext, endpoint: Endpoint, kv_cache_block_size: u32, kv_router_config: Option, - ) -> Result> { + ) -> Result { let endpoint_id = endpoint.id(); let client = endpoint.client().await?; @@ -359,18 +348,12 @@ where let router = if prefill_router_mode.is_kv_routing() { // Create KV chooser using the endpoint (this is a prefill router) - let effective_kv_router_config = prefill_kv_config.clone().unwrap_or_default(); - let selector = (context.worker_selector_factory)( - &effective_kv_router_config, - crate::worker_type::WorkerType::Prefill, - RoutingPartitionRef::new(&context.model_name, DEFAULT_ROUTING_GROUP), - ); let kv_chooser = context .model_manager - .kv_chooser_for_with_selector_and_client( + .kv_chooser_for_with_policy_and_client( client.clone(), prefill_block_size, - selector, + context.selection_policy.clone(), prefill_kv_config, context.prefill_load_estimator.clone(), Some(crate::worker_type::WorkerType::Prefill), @@ -414,7 +397,7 @@ where ) .await?; - Arc::new(RoutingHost::::new_builtin_with_coordinator( + Arc::new(RoutingHost::new_builtin_with_coordinator( push_router, load_context.clone(), affinity, @@ -438,7 +421,7 @@ where /// wrong strategy because of one transient discovery miss — a downgrade that /// would persist until the binding was rebuilt. async fn resolve_prefill_advertisement( - context: &PrefillBuildContext, + context: &PrefillBuildContext, endpoint: &Endpoint, ) -> Result { let endpoint_id = endpoint.id(); @@ -519,10 +502,10 @@ where let build_context = PrefillBuildContext { model_manager: router_ref.model_manager.clone(), decode_router_mode: router_ref.decode_router_mode, - worker_selector_factory: router_ref - .worker_selector_factory + selection_policy: router_ref + .selection_policy .clone() - .expect("enabled prefill router has a worker selector factory"), + .expect("enabled prefill router has a selection policy"), prefill_load_estimator: router_ref.prefill_load_estimator.clone(), session_affinity_ttl: router_ref.session_affinity_ttl, session_affinity_mode: router_ref.session_affinity_mode, diff --git a/lib/llm/src/kv_router/prefill_router/admission.rs b/lib/llm/src/kv_router/prefill_router/admission.rs index d9bea2a96939..77c73a168b6a 100644 --- a/lib/llm/src/kv_router/prefill_router/admission.rs +++ b/lib/llm/src/kv_router/prefill_router/admission.rs @@ -8,26 +8,18 @@ use futures::StreamExt; use tokio::sync::OwnedSemaphorePermit; use tracing::Instrument; -use dynamo_kv_router::selector::WorkerSelector; - use dynamo_runtime::{ pipeline::ManyOut, protocols::{annotated::Annotated, maybe_error::MaybeError}, }; use super::{PrefillCompletion, PrefillError, PrefillRouter}; -use crate::{ - local_model::runtime_config::ModelRuntimeConfig, - protocols::common::{ - llm_backend::{FinishReason, LLMEngineOutput}, - timing::RequestTracker, - }, +use crate::protocols::common::{ + llm_backend::{FinishReason, LLMEngineOutput}, + timing::RequestTracker, }; -impl PrefillRouter -where - Sel: WorkerSelector + Send + 'static, -{ +impl PrefillRouter { pub(super) async fn consume_prefill_stream( mut prefill_response: ManyOut>, tracker: Option>, @@ -170,7 +162,6 @@ where #[cfg(test)] mod tests { - use dynamo_kv_router::selector::DefaultWorkerSelector; use futures::stream; use serde_json::json; @@ -216,13 +207,9 @@ mod tests { let task_guard: dynamo_runtime::engine::EngineContextGuard = teardown.clone(); drop(teardown); - PrefillRouter::::consume_prefill_stream( - response, - None, - Some(task_guard), - ) - .await - .unwrap(); + PrefillRouter::consume_prefill_stream(response, None, Some(task_guard)) + .await + .unwrap(); assert!(teardown_weak.upgrade().is_some()); release_tx.send(()).unwrap(); @@ -238,7 +225,7 @@ mod tests { #[tokio::test] async fn first_output_error_does_not_record_prefill_complete() { let tracker = Arc::new(RequestTracker::new()); - let result = PrefillRouter::::consume_prefill_stream( + let result = PrefillRouter::consume_prefill_stream( prefill_stream(vec![Annotated::from_error("prefill failed")]), Some(tracker.clone()), None, @@ -256,7 +243,7 @@ mod tests { #[tokio::test] async fn later_output_error_is_propagated_after_prefill_arrival() { let tracker = Arc::new(RequestTracker::new()); - let result = PrefillRouter::::consume_prefill_stream( + let result = PrefillRouter::consume_prefill_stream( prefill_stream(vec![ valid_prefill_output(), Annotated::from_error("prefill stream failed"), @@ -287,7 +274,7 @@ mod tests { finish_reason: Some(finish_reason.clone()), ..Default::default() }; - let result = PrefillRouter::::consume_prefill_stream( + let result = PrefillRouter::consume_prefill_stream( prefill_stream(vec![Annotated::from_data(output)]), None, None, @@ -313,7 +300,7 @@ mod tests { disaggregated_params: Some(json!({"ctx_request_id": 42})), ..Default::default() }; - let result = PrefillRouter::::consume_prefill_stream( + let result = PrefillRouter::consume_prefill_stream( prefill_stream(vec![Annotated::from_data(output)]), None, None, @@ -334,7 +321,7 @@ mod tests { disaggregated_params: Some(json!({"ctx_request_id": 42})), ..Default::default() }; - let result = PrefillRouter::::consume_prefill_stream( + let result = PrefillRouter::consume_prefill_stream( prefill_stream(vec![Annotated::from_data(output)]), None, None, @@ -359,7 +346,7 @@ mod tests { disaggregated_params: Some(disaggregated_params), ..Default::default() }; - let result = PrefillRouter::::consume_prefill_stream( + let result = PrefillRouter::consume_prefill_stream( prefill_stream(vec![Annotated::from_data(output)]), None, None, diff --git a/lib/llm/src/kv_router/prefill_router/conditional_bypass.rs b/lib/llm/src/kv_router/prefill_router/conditional_bypass.rs index 147e2d6a7a02..619f2638c3ef 100644 --- a/lib/llm/src/kv_router/prefill_router/conditional_bypass.rs +++ b/lib/llm/src/kv_router/prefill_router/conditional_bypass.rs @@ -3,21 +3,16 @@ use anyhow::Result; use dynamo_kv_router::conditional_disagg::ConditionalDisaggDecisionInput; -use dynamo_kv_router::selector::WorkerSelector; use dynamo_runtime::pipeline::SingleIn; use super::PrefillRouter; use crate::kv_router::routing_host::{RoutePlan, RoutePlanSignals, is_cancelled}; -use crate::local_model::runtime_config::ModelRuntimeConfig; use crate::protocols::common::{llm_backend::PreprocessedRequest, timing::RequestPhase}; /// An admitted decode route selected by `RoutingHost` together with the /// conditional-disagg policy's diagnostic signals. -pub(super) struct ConditionalDisaggDecodeDecision -where - Sel: WorkerSelector + Send + 'static, -{ - pub plan: RoutePlan, +pub(super) struct ConditionalDisaggDecodeDecision { + pub plan: RoutePlan, pub overlap_tokens: usize, pub net_new_tokens: usize, } @@ -30,17 +25,14 @@ fn decode_gate_allows_bypass( policy_says_bypass && (!decode_gate_configured || matches!(decode_busy, Some(false))) } -impl PrefillRouter -where - Sel: WorkerSelector + Send + 'static, -{ +impl PrefillRouter { /// Preview one decode route, then admit it only when the topology policy /// chooses local decode. pub(super) async fn plan_conditional_disagg_decode( &self, request: &SingleIn, request_id: &str, - ) -> Result>> { + ) -> Result> { // Conditional disagg chooses a cache-hot decode worker, so it only // applies to a KV-routed decode set. if !self.decode_router_mode.is_kv_routing() { diff --git a/lib/llm/src/kv_router/prefill_router/mod.rs b/lib/llm/src/kv_router/prefill_router/mod.rs index 3159a8ec5e12..65a4846d4ef9 100644 --- a/lib/llm/src/kv_router/prefill_router/mod.rs +++ b/lib/llm/src/kv_router/prefill_router/mod.rs @@ -12,12 +12,8 @@ use tokio_util::sync::CancellationToken; use uuid::Uuid; use dynamo_kv_router::{ - PrefillLoadEstimator, - conditional_disagg::ConditionalDisaggPolicy, - config::RouterConfigOverride, - protocols::RoutingConstraints, - scheduling::QueueRejection, - selector::{DefaultWorkerSelector, WorkerSelector}, + PrefillLoadEstimator, conditional_disagg::ConditionalDisaggPolicy, + config::RouterConfigOverride, protocols::RoutingConstraints, scheduling::QueueRejection, }; use dynamo_runtime::{ pipeline::{ @@ -30,8 +26,7 @@ use futures::stream::{self, StreamExt}; use crate::{ discovery::ModelManager, - kv_router::{RoutingHost, WorkerSelectorFactory}, - local_model::runtime_config::ModelRuntimeConfig, + kv_router::{RoutingHost, SelectionPolicySource}, protocols::common::{ extensions::{SESSION_AFFINITY_CONTEXT_KEY, SessionAffinityId}, llm_backend::{LLMEngineOutput, PreprocessedRequest}, @@ -216,17 +211,14 @@ pub(crate) const BYPASS_REMOTE_PREFILL_ANNOTATION: &str = "x-bypass-remote-prefi /// Client-visible logprobs should not be placed in `disaggregated_params`, /// which is an engine-owned KV handoff contract rather than a public response /// channel. -pub struct PrefillRouter -where - Sel: WorkerSelector + Send + 'static, -{ - binding: ArcSwapOption>, +pub struct PrefillRouter { + binding: ArcSwapOption, target: Mutex>, target_tx: Option>>, /// Decode routing owns conditional-disagg planning and dispatch. This is /// installed after the frontend constructs its one decode `RoutingHost`. - decode_routing_host: OnceLock>>, - worker_selector_factory: Option>, + decode_routing_host: OnceLock>, + selection_policy: Option, model_manager: Arc, cancel_token: CancellationToken, /// Mode of the decode set that owns this router. Governs decode-side @@ -253,26 +245,20 @@ where activation_task_state: Arc<()>, } -struct PrefillBinding -where - Sel: WorkerSelector + Send + 'static, -{ +struct PrefillBinding { endpoint_id: EndpointId, - router: Arc>, + router: Arc, /// Resolved at activation from the prefill card. Lives here rather than on /// `PrefillRouter` because it is unknowable until a target is discovered, /// and changes when the binding is rebuilt. prefill_router_mode: RouterMode, } -struct PrefillBuildContext -where - Sel: WorkerSelector + Send + 'static, -{ +struct PrefillBuildContext { model_manager: Arc, /// Fallback mode for the prefill hop when the prefill card advertises none. decode_router_mode: RouterMode, - worker_selector_factory: WorkerSelectorFactory, + selection_policy: SelectionPolicySource, prefill_load_estimator: Option>, session_affinity_ttl: Option, session_affinity_mode: SessionAffinityMode, @@ -286,19 +272,13 @@ pub(crate) trait PrefillRouterLifecycle: Send + Sync { fn set_target(&self, target: Option); } -impl PrefillRouterLifecycle for PrefillRouter -where - Sel: WorkerSelector + Send + 'static, -{ +impl PrefillRouterLifecycle for PrefillRouter { fn set_target(&self, target: Option) { self.set_target(target); } } -impl Drop for PrefillRouter -where - Sel: WorkerSelector + Send + 'static, -{ +impl Drop for PrefillRouter { fn drop(&mut self) { tracing::debug!("Dropping PrefillRouter, cancelling background activation task"); self.cancel_token.cancel(); @@ -306,15 +286,13 @@ where } #[async_trait] -impl +impl Operator< SingleIn, ManyOut>, SingleIn, ManyOut>, - > for PrefillRouter -where - Sel: WorkerSelector + Send + 'static, + > for PrefillRouter { async fn generate( &self, @@ -580,18 +558,12 @@ where } } -impl PrefillRouter -where - Sel: WorkerSelector + Send + 'static, -{ +impl PrefillRouter { pub(crate) fn conditional_disagg_enabled(&self) -> bool { self.conditional_disagg_policy.is_enabled() } - pub(crate) fn set_decode_routing_host( - &self, - routing_host: Arc>, - ) -> Result<()> { + pub(crate) fn set_decode_routing_host(&self, routing_host: Arc) -> Result<()> { match self.decode_routing_host.set(routing_host) { Ok(()) => Ok(()), Err(routing_host) diff --git a/lib/llm/src/kv_router/prefill_router/query.rs b/lib/llm/src/kv_router/prefill_router/query.rs index 86a707f66027..bf3e3ff2f703 100644 --- a/lib/llm/src/kv_router/prefill_router/query.rs +++ b/lib/llm/src/kv_router/prefill_router/query.rs @@ -10,11 +10,9 @@ use dynamo_kv_router::{ AdmissionAttempt, queue::{SchedulerBookingCleanup, SchedulerBookingDescriptor}, }, - selector::WorkerSelector, }; use super::{PrefillError, PrefillLifecycleState, PrefillQueryOutcome, PrefillRouter}; -use crate::local_model::runtime_config::ModelRuntimeConfig; /// A prefill booking that owns cleanup of the exact scheduler attempt it admitted. /// @@ -67,10 +65,7 @@ impl Drop for PrefillReservation { } } -impl PrefillRouter -where - Sel: WorkerSelector + Send + 'static, -{ +impl PrefillRouter { /// Select a prefill worker and reserve it when KV routing is enabled. /// /// If this future is dropped while queued, the scheduler retracts its diff --git a/lib/llm/src/kv_router/routing_host.rs b/lib/llm/src/kv_router/routing_host.rs index 9aea9b89237c..40268cf5defe 100644 --- a/lib/llm/src/kv_router/routing_host.rs +++ b/lib/llm/src/kv_router/routing_host.rs @@ -28,11 +28,7 @@ use futures::stream::{self, StreamExt}; use tracing::Instrument; use crate::{ - kv_router::{ - KvRouter, metrics::RouterRequestMetrics, scheduler::DefaultWorkerSelector, - to_worker_selection_session_context, - }, - local_model::runtime_config::ModelRuntimeConfig, + kv_router::{KvRouter, metrics::RouterRequestMetrics, to_worker_selection_session_context}, lora::{LoadEstimator, LoraFilter}, preprocessor::PreprocessedRequest, protocols::common::{ @@ -75,14 +71,11 @@ fn route_target(worker: WorkerWithDpRank) -> AffinityTarget { AffinityTarget::new(worker.worker_id, Some(worker.dp_rank)) } -fn monitor_response_stream( +fn monitor_response_stream( mut response_stream: ManyOut>, context: Arc, - mut guard: RequestGuard, -) -> impl futures::Stream> + Send -where - Sel: WorkerSelector + Send + 'static, -{ + mut guard: RequestGuard, +) -> impl futures::Stream> + Send { async_stream::stream! { // Keep one cancellation future alive for the whole response stream. Calling // `stopped()` for every item repeatedly clones and polls a watch receiver. @@ -183,13 +176,10 @@ where } } -fn into_monitored_response( +fn into_monitored_response( response_stream: ManyOut>, - guard: RequestGuard, -) -> ManyOut> -where - Sel: WorkerSelector + Send + 'static, -{ + guard: RequestGuard, +) -> ManyOut> { let stream_context = response_stream.context(); let wrapped_stream = Box::pin(monitor_response_stream( response_stream, @@ -199,11 +189,8 @@ where ResponseStream::new(wrapped_stream, stream_context) } -enum RoutingPolicy -where - Sel: WorkerSelector + Send + 'static, -{ - Kv(Arc>), +enum RoutingPolicy { + Kv(Arc), Builtin(BuiltinWorkerSelector), Direct, DeviceAwareWeighted, @@ -241,12 +228,9 @@ struct DeviceAwareTelemetry { /// [`PushRouter`] owns discovery, fault detection, and transport. [`KvRouter`] /// owns optional KV candidate state. `RoutingHost` owns the common request /// lifecycle regardless of which policy selected the worker. -pub struct RoutingHost -where - Sel: WorkerSelector + Send + 'static, -{ +pub struct RoutingHost { inner: PushRouter>, - policy: RoutingPolicy, + policy: RoutingPolicy, request_metrics: Arc, affinity: Option, session_affinity_mode: SessionAffinityMode, @@ -260,13 +244,10 @@ where } /// An admitted KV route awaiting dispatch. -pub(crate) struct RoutePlan -where - Sel: WorkerSelector + Send + 'static, -{ +pub(crate) struct RoutePlan { signals: RoutePlanSignals, selection: WorkerSelection, - cleanup: KvRequestCleanup, + cleanup: KvRequestCleanup, affinity: Option, /// Carried forward from the [`RoutePreview`] this plan was admitted from, so /// preview, admission and dispatch draw on one budget instead of three. @@ -311,10 +292,7 @@ impl RoutePlanSignals { } } -impl RoutePlan -where - Sel: WorkerSelector + Send + 'static, -{ +impl RoutePlan { pub(crate) fn signals(&self) -> RoutePlanSignals { self.signals } @@ -334,15 +312,12 @@ where /// /// This alias remains supported through the Dynamo 1.x series. It may be /// removed only in a 2.0.0 (or later) breaking release. -pub type KvPushRouter = RoutingHost; +pub type KvPushRouter = RoutingHost; -impl RoutingHost -where - Sel: WorkerSelector + Send + 'static, -{ +impl RoutingHost { pub fn new( inner: PushRouter>, - kv_router: Arc>, + kv_router: Arc, session_affinity_ttl: Option, ) -> Result { let affinity = session_affinity_ttl @@ -359,7 +334,7 @@ where pub fn new_with_load_context( inner: PushRouter>, - kv_router: Arc>, + kv_router: Arc, load_context: Arc, session_affinity_ttl: Option, session_affinity_mode: SessionAffinityMode, @@ -379,7 +354,7 @@ where pub(crate) fn new_with_coordinator( inner: PushRouter>, - kv_router: Arc>, + kv_router: Arc, affinity: Option, session_affinity_mode: SessionAffinityMode, ) -> Self { @@ -394,7 +369,7 @@ where pub(crate) fn new_with_load_context_and_coordinator( inner: PushRouter>, - kv_router: Arc>, + kv_router: Arc, load_context: Arc, affinity: Option, session_affinity_mode: SessionAffinityMode, @@ -410,7 +385,7 @@ where fn new_with_optional_load_context_and_coordinator( inner: PushRouter>, - kv_router: Arc>, + kv_router: Arc, load_context: Option>, affinity: Option, session_affinity_mode: SessionAffinityMode, @@ -542,12 +517,12 @@ where } /// The active KV-aware data plane. - pub fn kv_router(&self) -> &Arc> { + pub fn kv_router(&self) -> &Arc { self.kv_router_if_enabled() .expect("routing host has no KV capability") } - pub(crate) fn kv_router_if_enabled(&self) -> Option<&Arc>> { + pub(crate) fn kv_router_if_enabled(&self) -> Option<&Arc> { match &self.policy { RoutingPolicy::Kv(chooser) => Some(chooser), RoutingPolicy::Builtin(_) @@ -720,10 +695,8 @@ where } #[async_trait] -impl AsyncEngine, ManyOut>, Error> - for RoutingHost -where - Sel: WorkerSelector + Send + 'static, +impl AsyncEngine, ManyOut>, Error> + for RoutingHost { /// Generate a request through the selected routing plane. /// diff --git a/lib/llm/src/kv_router/routing_host/builtin.rs b/lib/llm/src/kv_router/routing_host/builtin.rs index 28b511f93c1c..3b91e0c699e3 100644 --- a/lib/llm/src/kv_router/routing_host/builtin.rs +++ b/lib/llm/src/kv_router/routing_host/builtin.rs @@ -95,10 +95,7 @@ fn selection(worker_id: u64) -> WorkerSelectionResult { use super::*; -impl RoutingHost -where - Sel: WorkerSelector + Send + 'static, -{ +impl RoutingHost { fn select_lora_target( &self, request: &PreprocessedRequest, @@ -359,7 +356,7 @@ where let uses_occupancy = self .required_worker_inputs() .contains(WorkerInputs::OCCUPANCY); - let mut guard: RequestGuard = RequestGuard::new_builtin( + let mut guard: RequestGuard = RequestGuard::new_builtin( self.request_metrics.clone(), initial_worker, occupancy_reservation, diff --git a/lib/llm/src/kv_router/routing_host/kv.rs b/lib/llm/src/kv_router/routing_host/kv.rs index c6c3c505e5ab..337a0218e41b 100644 --- a/lib/llm/src/kv_router/routing_host/kv.rs +++ b/lib/llm/src/kv_router/routing_host/kv.rs @@ -4,10 +4,7 @@ use super::*; use crate::kv_router::{FindBestMatchAdmission, routing_host::kv_selection::SelectionOutcome}; -impl RoutingHost -where - Sel: WorkerSelector + Send + 'static, -{ +impl RoutingHost { #[allow(clippy::too_many_arguments)] async fn select_request_outcome( &self, @@ -163,7 +160,7 @@ where &self, request: &SingleIn, preview: RoutePreview, - ) -> Result, Error> { + ) -> Result { // Inherited, not restarted: this stage continues the route the preview // opened. let budget = preview.budget; @@ -221,7 +218,7 @@ where pub(crate) async fn dispatch_kv_plan( &self, request: SingleIn, - plan: RoutePlan, + plan: RoutePlan, ) -> Result>, Error> { let RoutePlan { mut selection, @@ -303,7 +300,7 @@ where phase: RequestPhase, is_query_only: bool, budget: &CleanupBudget, - ) -> Result, Error> { + ) -> Result { self.track_selection_with_cleanup(request, selection, phase, is_query_only, None, budget) .await } @@ -312,9 +309,9 @@ where &self, request: &SingleIn, selection: &mut WorkerSelection, - cleanup: KvRequestCleanup, + cleanup: KvRequestCleanup, budget: &CleanupBudget, - ) -> Result, Error> { + ) -> Result { let phase = request .tracker .as_ref() @@ -330,9 +327,9 @@ where selection: &mut WorkerSelection, phase: RequestPhase, is_query_only: bool, - cleanup: Option>, + cleanup: Option, budget: &CleanupBudget, - ) -> Result, Error> { + ) -> Result { let context_id = request.context().id().to_string(); let staged_kv = StagedKv::for_request(request.content()); let request_context = request.context().clone(); @@ -448,7 +445,7 @@ where &self, request: SingleIn, selection: WorkerSelection, - mut guard: RequestGuard, + mut guard: RequestGuard, budget: &CleanupBudget, ) -> Result>, Error> { let context_id = request.context().id().to_string(); diff --git a/lib/llm/src/kv_router/routing_host/kv_selection.rs b/lib/llm/src/kv_router/routing_host/kv_selection.rs index 18911cbfcd34..4c71f2dbd393 100644 --- a/lib/llm/src/kv_router/routing_host/kv_selection.rs +++ b/lib/llm/src/kv_router/routing_host/kv_selection.rs @@ -11,7 +11,6 @@ use dynamo_kv_router::{ BlockExtraInfo, RoutingConstraints, WorkerAffinityTarget, WorkerId, WorkerWithDpRank, }, scheduling::{AdmissionAttempt, AdvisoryWorkerLoad, QueueRejection, RoutingEligibility}, - selector::WorkerSelector, }; use dynamo_runtime::{dynamo_nvtx_range, pipeline::Error}; @@ -20,7 +19,6 @@ use crate::{ FindBestMatchAdmission, FindBestMatchInnerOutcome, FindBestMatchOutcome, routing_host::RoutingHost, }, - local_model::runtime_config::ModelRuntimeConfig, preprocessor::PreprocessedRequest, protocols::{ TokenIdType, @@ -100,10 +98,7 @@ struct BestMatchArgs<'a> { admission: FindBestMatchAdmission, } -impl RoutingHost -where - Sel: WorkerSelector + Send + 'static, -{ +impl RoutingHost { async fn select_best_match(&self, args: BestMatchArgs<'_>) -> Result { let outcome = self .kv_router() diff --git a/lib/llm/src/kv_router/routing_host/request_guard.rs b/lib/llm/src/kv_router/routing_host/request_guard.rs index a28ff9eacd63..8728650ef56a 100644 --- a/lib/llm/src/kv_router/routing_host/request_guard.rs +++ b/lib/llm/src/kv_router/routing_host/request_guard.rs @@ -10,9 +10,8 @@ use crate::{ metrics::RouterRequestMetrics, prefill_router::BYPASS_REMOTE_PREFILL_ANNOTATION, request_lease::RequestAttemptLease, - scheduler::{DefaultWorkerSelector, SchedulerBookingDescriptor}, + scheduler::SchedulerBookingDescriptor, }, - local_model::runtime_config::ModelRuntimeConfig, lora::LoadEstimator, preprocessor::PreprocessedRequest, protocols::common::{ @@ -28,7 +27,6 @@ use dynamo_kv_router::{ compute_next_seq_hash, }, scheduling::AdmissionAttempt, - selector::WorkerSelector, }; use dynamo_runtime::{ error::DynamoError, @@ -393,23 +391,17 @@ struct OutputBlockTracker { /// Owns the shared attempt-scoped scheduler and approximate-LRU lifecycle after /// a KV worker is selected. -pub(super) struct KvRequestCleanup -where - Sel: WorkerSelector + Send + 'static, -{ - chooser: Arc>, +pub(super) struct KvRequestCleanup { + chooser: Arc, context_id: String, worker: WorkerWithDpRank, approximate_lru: Option, lifecycle: Option, } -impl KvRequestCleanup -where - Sel: WorkerSelector + Send + 'static, -{ +impl KvRequestCleanup { pub(super) fn new( - chooser: Arc>, + chooser: Arc, context_id: String, worker: WorkerWithDpRank, attempt: AdmissionAttempt, @@ -458,11 +450,8 @@ where } /// Policy-specific state released by the host's common request lifecycle. -enum RequestCleanup -where - Sel: WorkerSelector + Send + 'static, -{ - Kv(KvRequestCleanup), +enum RequestCleanup { + Kv(KvRequestCleanup), Stateless { worker_id: u64, }, @@ -472,10 +461,7 @@ where }, } -impl RequestCleanup -where - Sel: WorkerSelector + Send + 'static, -{ +impl RequestCleanup { fn worker_id(&self) -> u64 { match self { Self::Kv(cleanup) => cleanup.worker.worker_id, @@ -569,11 +555,8 @@ impl OutputBlockTracker { /// /// Session-affinity lifetime is separate: `AffinityAcquire` and /// `AffinityLease` own binding commit, release, and invalidation. -pub(super) struct RequestGuard -where - Sel: WorkerSelector + Send + 'static, -{ - cleanup: RequestCleanup, +pub(super) struct RequestGuard { + cleanup: RequestCleanup, observability: RequestObservability, output_blocks: OutputBlockTracker, approximate_lru: Option, @@ -584,12 +567,9 @@ where _lora_load: Option, } -impl RequestGuard -where - Sel: WorkerSelector + Send + 'static, -{ +impl RequestGuard { pub(super) fn new_kv( - chooser: Arc>, + chooser: Arc, request_metrics: Arc, context_id: String, worker: WorkerWithDpRank, @@ -605,7 +585,7 @@ where pub(super) fn new_kv_with_cleanup( request_metrics: Arc, - cleanup: KvRequestCleanup, + cleanup: KvRequestCleanup, request: &PreprocessedRequest, ) -> Self { let chooser = &cleanup.chooser; @@ -820,10 +800,7 @@ where } } -impl Drop for RequestGuard -where - Sel: WorkerSelector + Send + 'static, -{ +impl Drop for RequestGuard { fn drop(&mut self) { self.observability .record_metrics(self.record_itl_at_completion); diff --git a/lib/llm/src/kv_router/routing_host/tests.rs b/lib/llm/src/kv_router/routing_host/tests.rs index dc58656381ab..413690c6a432 100644 --- a/lib/llm/src/kv_router/routing_host/tests.rs +++ b/lib/llm/src/kv_router/routing_host/tests.rs @@ -10,10 +10,9 @@ use std::{ time::{Duration, Instant}, }; -use dynamo_kv_router::{ - DefaultWorkerSelector, WorkerSelectionPolicy, config::KvRouterConfig, - protocols::RoutingConstraints, -}; +use dynamo_kv_router::{config::KvRouterConfig, protocols::RoutingConstraints}; + +use crate::kv_router::SelectionPolicySource; use dynamo_runtime::{ DistributedRuntime, Runtime, component::{Client, Instance}, @@ -94,7 +93,7 @@ fn classify_response_item_separates_terminal_failures_from_healthy_frames() { fn selector_state_remains_owned_by_the_scheduler_actor() { fn assert_send_sync() {} - assert_send_sync::>(); + assert_send_sync::(); } #[test] @@ -126,8 +125,7 @@ async fn builtin_host_constructs_only_declared_capabilities() { let inner = PushRouter::from_client(client.clone(), RouterMode::RoundRobin) .await .unwrap(); - let host = - RoutingHost::::new_builtin(inner, load_context.clone()).unwrap(); + let host = RoutingHost::new_builtin(inner, load_context.clone()).unwrap(); assert_eq!(host.required_worker_inputs(), WorkerInputs::NONE); assert!(host.hosted_occupancy.is_none()); @@ -139,7 +137,7 @@ async fn builtin_host_constructs_only_declared_capabilities() { let inner = PushRouter::from_client(client, RouterMode::PowerOfTwoChoices) .await .unwrap(); - let host = RoutingHost::::new_builtin(inner, load_context).unwrap(); + let host = RoutingHost::new_builtin(inner, load_context).unwrap(); let RoutingPolicy::Builtin(selector) = &host.policy else { unreachable!() }; @@ -154,7 +152,7 @@ async fn builtin_host_constructs_only_declared_capabilities() { assert_eq!(selection.worker_id, 1); assert_eq!(selection.occupancy, 1); assert_eq!(host.inner.occupancy_for_test(1), 1); - let mut guard: RequestGuard = RequestGuard::new_builtin( + let mut guard: RequestGuard = RequestGuard::new_builtin( Arc::clone(&host.request_metrics), selection.worker_id, Some(selection.reservation), @@ -188,7 +186,7 @@ async fn builtin_occupancy_selection_uses_all_selectable_workers() { .unwrap(); client.override_discovered_instances(vec![1, 2]); client.override_instance_avail(vec![1, 2]); - let host = RoutingHost::::new_builtin(inner, load_context).unwrap(); + let host = RoutingHost::new_builtin(inner, load_context).unwrap(); let RoutingPolicy::Builtin(selector) = &host.policy else { unreachable!() }; @@ -229,7 +227,7 @@ async fn builtin_direct_without_worker_is_invalid_argument() { .await .unwrap(); let affinity = AffinityCoordinator::new(Duration::from_secs(10)).unwrap(); - let host = RoutingHost::::new_builtin_with_coordinator( + let host = RoutingHost::new_builtin_with_coordinator( inner, load_context, Some(affinity), @@ -276,7 +274,7 @@ async fn builtin_direct_uses_bound_soft_affinity_as_exact_target() { .await .unwrap(); let affinity = AffinityCoordinator::new(Duration::from_secs(10)).unwrap(); - let host = RoutingHost::::new_builtin_with_coordinator( + let host = RoutingHost::new_builtin_with_coordinator( inner, load_context, Some(affinity.clone()), @@ -362,7 +360,7 @@ async fn builtin_hard_affinity_ignores_local_inhibition() { .await .unwrap(); let affinity = AffinityCoordinator::new(Duration::from_secs(10)).unwrap(); - let host = RoutingHost::::new_builtin_with_coordinator( + let host = RoutingHost::new_builtin_with_coordinator( inner, load_context, Some(affinity.clone()), @@ -441,7 +439,7 @@ async fn builtin_lora_keeps_separate_selection_and_cleanup() { ); let filter = Arc::new(LoraFilter::new(routing_table, LoraStateTracker::new())); let estimator = Arc::new(LoadEstimator::new()); - let host = RoutingHost::::new_builtin_with_capabilities( + let host = RoutingHost::new_builtin_with_capabilities( inner, load_context, None, @@ -519,7 +517,7 @@ async fn builtin_affinity_uses_common_host_for_every_policy() { .await .unwrap(); let affinity = AffinityCoordinator::new(Duration::from_secs(10)).unwrap(); - let host = RoutingHost::::new_builtin_with_coordinator( + let host = RoutingHost::new_builtin_with_coordinator( inner, load_context, Some(affinity.clone()), @@ -577,7 +575,7 @@ async fn builtin_hard_affinity_ignores_overload_while_soft_affinity_falls_back() let inner = PushRouter::from_client(client, RouterMode::RoundRobin) .await .unwrap(); - let host = RoutingHost::::new_builtin(inner, load_context).unwrap(); + let host = RoutingHost::new_builtin(inner, load_context).unwrap(); let request = Context::new(request()); let hard = host @@ -626,7 +624,7 @@ async fn builtin_direct_fallback_stays_disabled_for_affinity() { .await .unwrap(); let affinity = AffinityCoordinator::new(Duration::from_secs(10)).unwrap(); - let host = RoutingHost::::new_builtin_with_coordinator( + let host = RoutingHost::new_builtin_with_coordinator( inner, load_context, Some(affinity.clone()), @@ -1203,7 +1201,7 @@ async fn router_with_worker_configs_in_mode( workers, None, 16, - DefaultWorkerSelector::new(Some(config.clone()), "decode"), + SelectionPolicySource::Registry, Some(config), None, "decode", @@ -1854,13 +1852,8 @@ async fn router_request_counters_follow_admission_and_completion_lifecycle() { assert_eq!(metrics.requests_started_total().get(), started_before + 3); assert_eq!(metrics.requests_total.get(), completed_before + 1); - let mut builtin_guard = RequestGuard::::new_builtin( - Arc::clone(&metrics), - 7, - None, - None, - &request(), - ); + let mut builtin_guard = + RequestGuard::new_builtin(Arc::clone(&metrics), 7, None, None, &request()); assert_eq!(metrics.requests_started_total().get(), started_before + 4); builtin_guard.abort().await; drop(builtin_guard); @@ -2455,7 +2448,7 @@ async fn two_worker_migration_harness( workers, None, 16, - DefaultWorkerSelector::new(Some(config.clone()), "decode"), + SelectionPolicySource::Registry, Some(config), None, None, diff --git a/lib/llm/src/kv_router/routing_load.rs b/lib/llm/src/kv_router/routing_load.rs index 2c59c365f929..f15e29c2e990 100644 --- a/lib/llm/src/kv_router/routing_load.rs +++ b/lib/llm/src/kv_router/routing_load.rs @@ -17,7 +17,6 @@ use dynamo_runtime::pipeline::WorkerLoadMonitor; use crate::discovery::{KvWorkerMonitor, LoadThresholdHandle}; use crate::kv_router::KvRouter; -use crate::local_model::runtime_config::ModelRuntimeConfig; use crate::protocols::common::timing::{WORKER_TYPE_DECODE, WORKER_TYPE_PREFILL}; use crate::worker_type::WorkerType; @@ -388,18 +387,12 @@ impl Drop for RoutingLoadContext { } /// Standalone KV selection surface plus the context that owns its load tasks. -pub struct ManagedKvRouter -where - Sel: dynamo_kv_router::selector::WorkerSelector, -{ +pub struct ManagedKvRouter { load_context: Arc, - router: Arc>, + router: Arc, } -impl Clone for ManagedKvRouter -where - Sel: dynamo_kv_router::selector::WorkerSelector, -{ +impl Clone for ManagedKvRouter { fn clone(&self) -> Self { Self { load_context: self.load_context.clone(), @@ -408,22 +401,16 @@ where } } -impl std::ops::Deref for ManagedKvRouter -where - Sel: dynamo_kv_router::selector::WorkerSelector, -{ - type Target = KvRouter; +impl std::ops::Deref for ManagedKvRouter { + type Target = KvRouter; fn deref(&self) -> &Self::Target { &self.router } } -impl ManagedKvRouter -where - Sel: dynamo_kv_router::selector::WorkerSelector, -{ - pub fn new(load_context: Arc, router: Arc>) -> Self { +impl ManagedKvRouter { + pub fn new(load_context: Arc, router: Arc) -> Self { Self { load_context, router, @@ -434,7 +421,7 @@ where &self.load_context } - pub fn router(&self) -> &Arc> { + pub fn router(&self) -> &Arc { &self.router } } From f9f1259b0f0174a2299d485bfa1e176504bb0428 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Thu, 3 Sep 2026 21:02:20 +0000 Subject: [PATCH 016/102] refactor(kv-router): remove the delegated worker-selector adapter DelegatedWorkerSelector, ErasedWorkerConfig, and the Delegated policy state existed only so the frontend could run an arbitrary caller-injected WorkerSelector type on a partition. That path is gone with the selector type parameter; custom policies use WorkerSelectionPolicy's filter, scorer, and picker API like every other host. Signed-off-by: Ishan Dhanani --- lib/kv-router/src/scheduling/selector/mod.rs | 6 +- .../src/scheduling/selector/policy.rs | 122 +----------------- 2 files changed, 4 insertions(+), 124 deletions(-) diff --git a/lib/kv-router/src/scheduling/selector/mod.rs b/lib/kv-router/src/scheduling/selector/mod.rs index 30d28398a76e..1b81bc78cb74 100644 --- a/lib/kv-router/src/scheduling/selector/mod.rs +++ b/lib/kv-router/src/scheduling/selector/mod.rs @@ -10,9 +10,9 @@ pub use default::DefaultWorkerSelector; use default::{DefaultWorkerPicker, DefaultWorkerScorer}; pub use policy::{ - DelegatedWorkerSelector, ErasedWorkerConfig, ScoredWorkerCandidate, WorkerCacheInput, - WorkerCandidate, WorkerFilter, WorkerInputView, WorkerInputs, WorkerLoadInput, WorkerPicker, - WorkerScorer, WorkerSelectionContext, WorkerSelectionPolicy, + ScoredWorkerCandidate, WorkerCacheInput, WorkerCandidate, WorkerFilter, WorkerInputView, + WorkerInputs, WorkerLoadInput, WorkerPicker, WorkerScorer, WorkerSelectionContext, + WorkerSelectionPolicy, }; use default::{pick_default_worker, selection_weights}; diff --git a/lib/kv-router/src/scheduling/selector/policy.rs b/lib/kv-router/src/scheduling/selector/policy.rs index 0dcd5f5ca7dd..78a5cbf6e2ab 100644 --- a/lib/kv-router/src/scheduling/selector/policy.rs +++ b/lib/kv-router/src/scheduling/selector/policy.rs @@ -344,99 +344,6 @@ pub(super) enum WorkerSelectionPolicyState { Default(DefaultWorkerPicker), /// Policy-local state owned and called serially by one scheduler queue actor. Custom(RefCell), - /// A caller-owned selector that sees the request and a config-erased view of - /// the eligible workers. Lets a host run its own `WorkerSelector` on a - /// partition whose worker-config type it does not know. - Delegated(Box), -} - -/// Worker configuration reduced to the [`WorkerConfigLike`] surface, so a -/// selector written against one host's config type can run on any partition. -#[derive(Debug, Clone, Default, PartialEq)] -pub struct ErasedWorkerConfig { - pub data_parallel_start_rank: u32, - pub data_parallel_size: u32, - pub max_num_batched_tokens: Option, - pub total_kv_blocks: Option, - pub native_offloading_capacity_tokens: Option, - pub taints: std::collections::HashSet, - pub stable_routing_id: Option, - pub topology_domains: Option>, - pub kv_transfer_domain: Option, - pub kv_transfer_enforcement: Option, - pub kv_transfer_preferred_weight: Option, -} - -impl ErasedWorkerConfig { - pub fn from_config_like(config: &C) -> Self { - Self { - data_parallel_start_rank: config.data_parallel_start_rank(), - data_parallel_size: config.data_parallel_size(), - max_num_batched_tokens: config.max_num_batched_tokens(), - total_kv_blocks: config.total_kv_blocks(), - native_offloading_capacity_tokens: config.native_offloading_capacity_tokens(), - taints: config.taints().clone(), - stable_routing_id: config.stable_routing_id().map(str::to_owned), - topology_domains: config.topology_domains().cloned(), - kv_transfer_domain: config.kv_transfer_domain().map(str::to_owned), - kv_transfer_enforcement: config.kv_transfer_enforcement(), - kv_transfer_preferred_weight: config.kv_transfer_preferred_weight(), - } - } -} - -impl WorkerConfigLike for ErasedWorkerConfig { - fn data_parallel_start_rank(&self) -> u32 { - self.data_parallel_start_rank - } - fn data_parallel_size(&self) -> u32 { - self.data_parallel_size - } - fn max_num_batched_tokens(&self) -> Option { - self.max_num_batched_tokens - } - fn total_kv_blocks(&self) -> Option { - self.total_kv_blocks - } - fn native_offloading_capacity_tokens(&self) -> Option { - self.native_offloading_capacity_tokens - } - fn taints(&self) -> &std::collections::HashSet { - &self.taints - } - fn stable_routing_id(&self) -> Option<&str> { - self.stable_routing_id.as_deref() - } - fn topology_domains(&self) -> Option<&HashMap> { - self.topology_domains.as_ref() - } - fn kv_transfer_domain(&self) -> Option<&str> { - self.kv_transfer_domain.as_deref() - } - fn kv_transfer_enforcement(&self) -> Option { - self.kv_transfer_enforcement - } - fn kv_transfer_preferred_weight(&self) -> Option { - self.kv_transfer_preferred_weight - } -} - -/// A caller-owned selector a partition policy delegates to. Receives the -/// scheduling request and a config-erased view of the eligible workers. -pub trait DelegatedWorkerSelector: Send + Sync { - fn required_worker_inputs(&self) -> WorkerInputs; - - fn uses_exclusive_affinity_target(&self) -> bool { - false - } - - fn select_worker( - &self, - workers: &HashMap, - request: &SchedulingRequest, - eligibility: RoutingEligibility<'_>, - block_size: u32, - ) -> Result; } pub(super) enum WorkerSelectionPolicyStateRef<'a> { @@ -480,19 +387,6 @@ impl WorkerSelectionPolicy { Self::new_with_filters(kv_router_config, worker_label, Vec::new(), scorers, picker) } - /// Run `selector` for every selection on this partition. - pub fn delegated( - kv_router_config: KvRouterConfig, - worker_label: &'static str, - selector: Box, - ) -> Self { - Self { - kv_router_config, - worker_label, - state: WorkerSelectionPolicyState::Delegated(selector), - } - } - /// Build a custom policy from ordered filters, additive scorers, and one picker. /// /// `worker_label` identifies the worker pool in routing logs. A typed policy factory normally @@ -696,13 +590,7 @@ pub(super) fn collect_custom_candidates( impl WorkerSelector for WorkerSelectionPolicy { fn uses_exclusive_affinity_target(&self) -> bool { - match &self.state { - WorkerSelectionPolicyState::Default(_) => true, - WorkerSelectionPolicyState::Custom(_) => false, - WorkerSelectionPolicyState::Delegated(selector) => { - selector.uses_exclusive_affinity_target() - } - } + matches!(&self.state, WorkerSelectionPolicyState::Default(_)) } fn required_worker_inputs(&self) -> WorkerInputs { @@ -712,7 +600,6 @@ impl WorkerSelector for WorkerSelectionPolicy { let state = state.borrow(); state.filter_inputs | state.scorer_picker_inputs } - WorkerSelectionPolicyState::Delegated(selector) => selector.required_worker_inputs(), } } @@ -729,13 +616,6 @@ impl WorkerSelector for WorkerSelectionPolicy { WorkerSelectionPolicyState::Custom(state) => { WorkerSelectionPolicyStateRef::Custom(state) } - WorkerSelectionPolicyState::Delegated(selector) => { - let erased: HashMap = workers - .iter() - .map(|(id, config)| (*id, ErasedWorkerConfig::from_config_like(config))) - .collect(); - return selector.select_worker(&erased, request, eligibility, block_size); - } }; select_worker_with_policy( &self.kv_router_config, From e7713c81ed925fb941338bd32d3ce2049b54ea73 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Thu, 3 Sep 2026 21:12:51 +0000 Subject: [PATCH 017/102] refactor(kv-router): group host-supplied selection inputs into SelectionHost SelectionHostHooks was a flat bag of eight optional fields. Regroup it by purpose as SelectionHost { load, cache, eligibility, telemetry, replication }, each group defaulting to the standalone service's behavior so a host sets only the groups it owns. The builder method becomes SelectionServiceBuilder::host. No behavior change; the frontend's embedded partition sets the same values. Signed-off-by: Ishan Dhanani --- .../src/services/selection/core/mod.rs | 202 ++++++++++-------- lib/kv-router/src/services/selection/mod.rs | 5 +- .../src/services/selection/service.rs | 16 +- lib/llm/src/kv_router/embedded.rs | 35 +-- 4 files changed, 143 insertions(+), 115 deletions(-) diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index a0c98ab98373..ebc8b3f09d5d 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -181,37 +181,64 @@ struct ReservationBooking { routing_hashes: Option>, } -/// Host-supplied hooks applied to every partition scheduler the core creates. -/// -/// A host that owns request transport (for example the frontend `KvRouter`) -/// uses these to feed overload shedding and hard availability from its own -/// client state, to attach a prefill-duration estimator, and to consult a -/// shared (cross-worker) KV cache during selection. The standalone HTTP -/// service leaves them unset. +/// What an embedding host supplies to every partition the core creates, +/// grouped by purpose. Each group defaults to the standalone service's +/// behavior, so a host overrides only the groups it owns: the frontend +/// `KvRouter` feeds load and availability from its request client, points +/// overlap refresh at its own index, and carries replica sync on its own +/// transport. +#[derive(Clone, Default)] +pub struct SelectionHost { + pub load: HostLoad, + pub cache: HostCache, + pub eligibility: HostEligibility, + pub telemetry: HostTelemetry, + pub replication: HostReplication, +} + +/// Load signals the host knows and the partition scheduler does not. +#[derive(Clone, Default)] +pub struct HostLoad { + pub prefill_estimator: Option>, + /// Workers to shed from selection (the host's overload detector). + pub overloaded_workers: Option, + /// Workers the host can currently reach; others are never selected. + pub available_workers: Option, +} + +/// KV-cache knowledge outside the partition's own index. #[derive(Clone, Default)] -pub struct SelectionHostHooks { - pub prefill_load_estimator: Option>, - pub overloaded_worker_provider: Option, - pub available_worker_provider: Option, +pub struct HostCache { /// Queried alongside the indexer for prompts that carry `token_ids`; a /// failed lookup is logged and selection proceeds without shared hits. - pub shared_cache: Option>, + pub shared: Option>, + /// Where dequeue-time overlap refresh reads from. A host that keeps its + /// own index points this at that index so queued requests are re-scored + /// against the index that scored them. + pub overlap_refresh: OverlapRefreshSource, +} + +/// Host-owned narrowing of the candidate set. +#[derive(Clone, Default)] +pub struct HostEligibility { /// Narrows candidates to the workers that can serve the request's LoRA /// adapter, strictly within the caller's `allowed_worker_ids`. pub lora_worker_filter: Option>, - /// Where dequeue-time overlap refresh reads from. An embedding host that - /// keeps its own index and drives the scheduler directly points this at - /// that index so queued requests are re-scored against the index that - /// scored them, or disables refresh when its index cannot serve it. - pub overlap_refresh: OverlapRefreshSource, +} + +/// Scheduler state the host consumes. +#[derive(Clone, Default)] +pub struct HostTelemetry { /// Receives each partition's scheduler-owned load snapshots (active decode - /// blocks and prefill tokens per worker) so the host can run overload - /// detection on the numbers the scheduler books against. - pub scheduler_load_sink: Option>, - /// Replica sync carried by the host's own transport for partitions this - /// service does not mesh itself (ignored when the service runs its own - /// ZMQ replica sync). - pub replica_sync: Option, + /// blocks and prefill tokens per worker) for metrics and overload detection. + pub scheduler_load: Option>, +} + +/// Replica-sync transport the host carries for partitions this core does not +/// mesh itself (ignored when the service runs its own ZMQ replica sync). +#[derive(Clone, Default)] +pub struct HostReplication { + pub channels: Option, } /// Source of dequeue-time overlap refreshes for partition schedulers. @@ -236,25 +263,24 @@ impl std::fmt::Debug for OverlapRefreshSource { } } -impl std::fmt::Debug for SelectionHostHooks { +impl std::fmt::Debug for SelectionHost { fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { formatter - .debug_struct("SelectionHostHooks") + .debug_struct("SelectionHost") + .field("prefill_estimator", &self.load.prefill_estimator.is_some()) .field( - "prefill_load_estimator", - &self.prefill_load_estimator.is_some(), + "overloaded_workers", + &self.load.overloaded_workers.is_some(), ) - .field("shared_cache", &self.shared_cache.is_some()) - .field("lora_worker_filter", &self.lora_worker_filter.is_some()) - .field("overlap_refresh", &self.overlap_refresh) + .field("available_workers", &self.load.available_workers.is_some()) + .field("shared_cache", &self.cache.shared.is_some()) + .field("overlap_refresh", &self.cache.overlap_refresh) .field( - "overloaded_worker_provider", - &self.overloaded_worker_provider.is_some(), - ) - .field( - "available_worker_provider", - &self.available_worker_provider.is_some(), + "lora_worker_filter", + &self.eligibility.lora_worker_filter.is_some(), ) + .field("scheduler_load", &self.telemetry.scheduler_load.is_some()) + .field("replication", &self.replication.channels.is_some()) .finish() } } @@ -300,7 +326,7 @@ pub struct SelectionCore { indexer_registry: Arc, kv_router_config: crate::config::KvRouterConfig, worker_selection_policy_factory: Option, - host_hooks: SelectionHostHooks, + host: SelectionHost, worker_type: WorkerType, cancel_token: CancellationToken, replica_config: Option, @@ -344,7 +370,7 @@ impl SelectionCore { cancel_token, None, None, - SelectionHostHooks::default(), + SelectionHost::default(), WorkerType::Aggregated, true, cache_config, @@ -367,7 +393,7 @@ impl SelectionCore { cancel_token: CancellationToken, replica_config: Option, worker_selection_policy_factory: Option, - host_hooks: SelectionHostHooks, + host: SelectionHost, worker_type: WorkerType, signal_indexer_ready: bool, cache_config: SelectionCacheConfig, @@ -395,7 +421,7 @@ impl SelectionCore { indexer_registry, kv_router_config, worker_selection_policy_factory, - host_hooks, + host, worker_type, cancel_token, replica_config, @@ -668,8 +694,9 @@ impl SelectionCore { .get_or_try_init(|| -> Result, SelectionError> { let (workers_tx, workers_rx) = watch::channel(HashMap::new()); let host_replica = self - .host_hooks - .replica_sync + .host + .replication + .channels .as_ref() .and_then(|factory| factory(&key)); let scoped_replica_sync = setup_scoped_replica_sync( @@ -682,7 +709,7 @@ impl SelectionCore { let slots = Arc::new(ActiveSequencesMultiWorker::new_with_replica_worker_policy( scoped_replica_sync .publisher - .with_load_sink(self.host_hooks.scheduler_load_sink.clone()), + .with_load_sink(self.host.telemetry.scheduler_load.clone()), block_size as usize, HashMap::new(), scoped_replica_sync.enabled, @@ -701,7 +728,7 @@ impl SelectionCore { let indexer = self .indexer_registry .get_or_create_indexer(key.clone(), block_size); - let refresh_provider = match &self.host_hooks.overlap_refresh { + let refresh_provider = match &self.host.cache.overlap_refresh { OverlapRefreshSource::PartitionIndexer => { Some(RefreshProvider::Local(Arc::new(indexer.clone()))) } @@ -731,12 +758,12 @@ impl SelectionCore { profile, block_size, selector, - self.host_hooks.prefill_load_estimator.clone(), + self.host.load.prefill_estimator.clone(), overlap_refresh, // Standalone selection has no router Client snapshot, so // these stay `None` unless an embedding host injects them. - self.host_hooks.overloaded_worker_provider.clone(), - self.host_hooks.available_worker_provider.clone(), + self.host.load.overloaded_workers.clone(), + self.host.load.available_workers.clone(), self.kv_router_config.router_queue_recheck_interval(), self.kv_router_config.router_track_prefill_tokens, self.cancel_token.child_token(), @@ -1016,7 +1043,7 @@ impl SelectionCore { track_prefill_tokens, ) }); - let allowed_worker_ids = match self.host_hooks.lora_worker_filter.as_deref() { + let allowed_worker_ids = match self.host.eligibility.lora_worker_filter.as_deref() { Some(filter) => narrow_allowed_worker_ids_by_lora( filter, prompt.lora_name.as_deref(), @@ -1639,7 +1666,7 @@ impl SelectionCore { } }; let shared_cache = query_shared_cache - .then_some(self.host_hooks.shared_cache.as_deref()) + .then_some(self.host.cache.shared.as_deref()) .flatten() .zip(prompt.token_ids.as_deref()); let shared_cache_lookup = async { @@ -1979,7 +2006,7 @@ mod tests { CancellationToken::new(), None, None, - SelectionHostHooks::default(), + SelectionHost::default(), worker_type, true, SelectionCacheConfig::default(), @@ -2000,12 +2027,12 @@ mod tests { } } - fn core_with_hooks(hooks: SelectionHostHooks) -> SelectionCore { - core_with_hooks_and_policy(hooks, None) + fn core_with_host(host: SelectionHost) -> SelectionCore { + core_with_host_and_policy(host, None) } - fn core_with_hooks_and_policy( - hooks: SelectionHostHooks, + fn core_with_host_and_policy( + host: SelectionHost, policy_factory: Option, ) -> SelectionCore { let config = test_config(false); @@ -2019,7 +2046,7 @@ mod tests { CancellationToken::new(), None, policy_factory, - hooks, + host, WorkerType::Aggregated, true, SelectionCacheConfig::default(), @@ -2180,7 +2207,7 @@ mod tests { CancellationToken::new(), None, None, - SelectionHostHooks::default(), + SelectionHost::default(), WorkerType::Aggregated, true, SelectionCacheConfig::default(), @@ -2441,9 +2468,11 @@ mod tests { #[tokio::test] async fn injected_lora_filter_narrows_candidates() { - let core = core_with_hooks(SelectionHostHooks { - lora_worker_filter: Some(Arc::new(OnlyWorkerForLora { worker_id: 2 })), - ..SelectionHostHooks::default() + let core = core_with_host(SelectionHost { + eligibility: HostEligibility { + lora_worker_filter: Some(Arc::new(OnlyWorkerForLora { worker_id: 2 })), + }, + ..SelectionHost::default() }); core.upsert_worker(worker(1)).await.expect("worker upsert"); core.upsert_worker(worker(2)).await.expect("worker upsert"); @@ -2476,18 +2505,15 @@ mod tests { async fn shared_cache_hits_reach_worker_selection() { let calls: SharedCacheCalls = Arc::new(parking_lot::Mutex::new(Vec::new())); let (factory, observed) = capturing_policy_factory(); - let core = core_with_hooks_and_policy( - SelectionHostHooks { - prefill_load_estimator: None, - overloaded_worker_provider: None, - available_worker_provider: None, - shared_cache: Some(Arc::new(RecordingSharedCache { - calls: Arc::clone(&calls), - })), - lora_worker_filter: None, - overlap_refresh: OverlapRefreshSource::default(), - scheduler_load_sink: None, - replica_sync: None, + let core = core_with_host_and_policy( + SelectionHost { + cache: HostCache { + shared: Some(Arc::new(RecordingSharedCache { + calls: Arc::clone(&calls), + })), + ..HostCache::default() + }, + ..SelectionHost::default() }, Some(factory), ); @@ -2537,7 +2563,7 @@ mod tests { use super::super::types::SelectionSessionContext; let (factory, observed) = capturing_policy_factory(); - let core = core_with_hooks_and_policy(SelectionHostHooks::default(), Some(factory)); + let core = core_with_host_and_policy(SelectionHost::default(), Some(factory)); core.upsert_worker(worker(1)).await.expect("worker upsert"); let mut request = select_request(); @@ -2582,15 +2608,12 @@ mod tests { let available: Arc>>>> = Arc::new(parking_lot::Mutex::new(None)); let provider_state = Arc::clone(&available); - let core = core_with_hooks(SelectionHostHooks { - prefill_load_estimator: None, - overloaded_worker_provider: None, - available_worker_provider: Some(Arc::new(move || provider_state.lock().clone())), - shared_cache: None, - lora_worker_filter: None, - overlap_refresh: OverlapRefreshSource::default(), - scheduler_load_sink: None, - replica_sync: None, + let core = core_with_host(SelectionHost { + load: HostLoad { + available_workers: Some(Arc::new(move || provider_state.lock().clone())), + ..HostLoad::default() + }, + ..SelectionHost::default() }); core.upsert_worker(worker(1)).await.expect("worker upsert"); core.upsert_worker(worker(2)).await.expect("worker upsert"); @@ -2606,15 +2629,12 @@ mod tests { #[tokio::test] async fn injected_overload_provider_excludes_worker() { - let core = core_with_hooks(SelectionHostHooks { - prefill_load_estimator: None, - overloaded_worker_provider: Some(Arc::new(|| Some(HashSet::from([1])))), - available_worker_provider: None, - shared_cache: None, - lora_worker_filter: None, - overlap_refresh: OverlapRefreshSource::default(), - scheduler_load_sink: None, - replica_sync: None, + let core = core_with_host(SelectionHost { + load: HostLoad { + overloaded_workers: Some(Arc::new(|| Some(HashSet::from([1])))), + ..HostLoad::default() + }, + ..SelectionHost::default() }); core.upsert_worker(worker(1)).await.expect("worker upsert"); core.upsert_worker(worker(2)).await.expect("worker upsert"); diff --git a/lib/kv-router/src/services/selection/mod.rs b/lib/kv-router/src/services/selection/mod.rs index 9a96ef28437e..c5e9247c19f8 100644 --- a/lib/kv-router/src/services/selection/mod.rs +++ b/lib/kv-router/src/services/selection/mod.rs @@ -25,8 +25,9 @@ pub use crate::services::common::replica_sync::{ ScopedSequencePublisher, }; pub use core::{ - OverlapRefreshSource, RefreshProvider, SelectionCore, SelectionHostHooks, SelectionPartition, - SelectionScheduler, SelectionServiceConfig, + HostCache, HostEligibility, HostLoad, HostReplication, HostTelemetry, OverlapRefreshSource, + RefreshProvider, SelectionCore, SelectionHost, SelectionPartition, SelectionScheduler, + SelectionServiceConfig, }; pub use error::SelectionError; pub use input::PromptRequest; diff --git a/lib/kv-router/src/services/selection/service.rs b/lib/kv-router/src/services/selection/service.rs index 17a96e7c77e9..ce24dacfe0e4 100644 --- a/lib/kv-router/src/services/selection/service.rs +++ b/lib/kv-router/src/services/selection/service.rs @@ -16,7 +16,7 @@ use crate::services::common::replica_sync::{ use crate::services::indexer::backend::IndexerPolicy; use crate::tracking_hash::TrackingHashContext; -use super::core::{SelectionCore, SelectionHostHooks, SelectionPartition, SelectionServiceConfig}; +use super::core::{SelectionCore, SelectionHost, SelectionPartition, SelectionServiceConfig}; use super::error::SelectionError; use super::pending::SelectionCacheConfig; use super::policy_registry::WorkerSelectionPolicyRegistry; @@ -36,7 +36,7 @@ pub struct SelectionServiceBuilder { selection_cache: SelectionCacheConfig, worker_type: WorkerType, worker_selection_policy_registry: WorkerSelectionPolicyRegistry, - host_hooks: SelectionHostHooks, + host: SelectionHost, remote_indexer_url: Option, worker_selection_policy_factory: Option, external_kv_events: bool, @@ -77,7 +77,7 @@ impl SelectionServiceBuilder { selection_cache: SelectionCacheConfig::default(), worker_type, worker_selection_policy_registry, - host_hooks: SelectionHostHooks::default(), + host: SelectionHost::default(), remote_indexer_url: None, worker_selection_policy_factory: None, external_kv_events: false, @@ -109,11 +109,9 @@ impl SelectionServiceBuilder { self } - /// Attach host-owned hooks (prefill-load estimator, overload exclusion - /// set, hard-availability set). Applied to every partition scheduler this - /// service creates. - pub fn host_hooks(mut self, hooks: SelectionHostHooks) -> Self { - self.host_hooks = hooks; + /// What the embedding host supplies to every partition this service creates. + pub fn host(mut self, host: SelectionHost) -> Self { + self.host = host; self } @@ -182,7 +180,7 @@ impl SelectionServiceBuilder { cancel_token.clone(), replica_config, worker_selection_policy_factory, - self.host_hooks, + self.host, self.worker_type, false, self.selection_cache, diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs index 4a74e39018ea..408b571393f1 100644 --- a/lib/llm/src/kv_router/embedded.rs +++ b/lib/llm/src/kv_router/embedded.rs @@ -28,8 +28,9 @@ use dynamo_kv_router::scheduling::{ }; use dynamo_kv_router::sequences::{SequenceError, SequenceRequest}; use dynamo_kv_router::services::selection::{ - OverlapRefreshSource, SelectionHostHooks, SelectionPartition, SelectionService, - SelectionServiceBuilder, WorkerRequest, WorkerSelectionPolicyRegistry, + HostCache, HostEligibility, HostLoad, HostReplication, HostTelemetry, OverlapRefreshSource, + SelectionHost, SelectionPartition, SelectionService, SelectionServiceBuilder, WorkerRequest, + WorkerSelectionPolicyRegistry, }; use dynamo_kv_router::{DEFAULT_ROUTING_GROUP, PrefillLoadEstimator, WorkerSelectionPolicyFactory}; use dynamo_tokens::SequenceHash; @@ -147,18 +148,26 @@ impl EmbeddedSelection { .worker_selection_policy_factory(args.policy_factory) .indexer_threads(1) .external_kv_events() - .host_hooks(SelectionHostHooks { - prefill_load_estimator: args.prefill_load_estimator, - overloaded_worker_provider: Some(args.overloaded_worker_provider), - available_worker_provider: Some(args.available_worker_provider), - shared_cache: None, - lora_worker_filter: None, - overlap_refresh: match args.overlap_refresh { - Some(provider) => OverlapRefreshSource::External(provider), - None => OverlapRefreshSource::Disabled, + .host(SelectionHost { + load: HostLoad { + prefill_estimator: args.prefill_load_estimator, + overloaded_workers: Some(args.overloaded_worker_provider), + available_workers: Some(args.available_worker_provider), + }, + cache: HostCache { + shared: None, + overlap_refresh: match args.overlap_refresh { + Some(provider) => OverlapRefreshSource::External(provider), + None => OverlapRefreshSource::Disabled, + }, + }, + eligibility: HostEligibility::default(), + telemetry: HostTelemetry { + scheduler_load: Some(Arc::new(SenderLoadSink(args.scheduler_load))), + }, + replication: HostReplication { + channels: replica_sync, }, - scheduler_load_sink: Some(Arc::new(SenderLoadSink(args.scheduler_load))), - replica_sync, }) .build() .await From 9cb836bccd906ed51d90ab525273fcc7efea2805 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Wed, 9 Sep 2026 11:10:44 -0700 Subject: [PATCH 018/102] fix(router): update callers after selector type erasure Signed-off-by: Ishan Dhanani --- lib/llm/src/kv_router/prefill_router/query.rs | 17 +++++------------ .../src/kv_router/routing_host/request_guard.rs | 7 ++----- lib/llm/src/kv_router/routing_host/tests.rs | 4 ++-- 3 files changed, 9 insertions(+), 19 deletions(-) diff --git a/lib/llm/src/kv_router/prefill_router/query.rs b/lib/llm/src/kv_router/prefill_router/query.rs index bf3e3ff2f703..6673f232a9dd 100644 --- a/lib/llm/src/kv_router/prefill_router/query.rs +++ b/lib/llm/src/kv_router/prefill_router/query.rs @@ -241,8 +241,9 @@ mod tests { time::Duration, }; + use crate::local_model::runtime_config::ModelRuntimeConfig; use async_trait::async_trait; - use dynamo_kv_router::{config::KvRouterConfig, selector::DefaultWorkerSelector}; + use dynamo_kv_router::config::KvRouterConfig; use dynamo_runtime::{ DistributedRuntime, Runtime, component::Instance, @@ -483,12 +484,7 @@ mod tests { (shared, prefill, worker_runtimes, workers) } - async fn tracked_binding( - label: &str, - ) -> ( - Arc>, - Arc>, - ) { + async fn tracked_binding(label: &str) -> (Arc, Arc) { let runtime = Runtime::from_current().unwrap(); let distributed = DistributedRuntime::new(runtime, DistributedConfig::process_local()) .await @@ -524,7 +520,7 @@ mod tests { workers_rx, None, 16, - DefaultWorkerSelector::new(Some(config.clone()), "prefill"), + crate::kv_router::SelectionPolicySource::Registry, Some(config), None, Some(WorkerType::Prefill), @@ -549,10 +545,7 @@ mod tests { (binding, chooser) } - async fn tracked_prefill_router() -> ( - Arc>, - Arc>, - ) { + async fn tracked_prefill_router() -> (Arc, Arc) { let (binding, chooser) = tracked_binding("primary").await; let router = PrefillRouter::disabled(Arc::new(ModelManager::new()), RouterMode::KV, None); router.binding.store(Some(binding)); diff --git a/lib/llm/src/kv_router/routing_host/request_guard.rs b/lib/llm/src/kv_router/routing_host/request_guard.rs index 8728650ef56a..12e26b3906f5 100644 --- a/lib/llm/src/kv_router/routing_host/request_guard.rs +++ b/lib/llm/src/kv_router/routing_host/request_guard.rs @@ -5,11 +5,8 @@ use std::{collections::HashMap, sync::Arc}; use crate::{ kv_router::{ - KvRouter, - indexer::ApproximateRequestLease, - metrics::RouterRequestMetrics, - prefill_router::BYPASS_REMOTE_PREFILL_ANNOTATION, - request_lease::RequestAttemptLease, + KvRouter, indexer::ApproximateRequestLease, metrics::RouterRequestMetrics, + prefill_router::BYPASS_REMOTE_PREFILL_ANNOTATION, request_lease::RequestAttemptLease, scheduler::SchedulerBookingDescriptor, }, lora::LoadEstimator, diff --git a/lib/llm/src/kv_router/routing_host/tests.rs b/lib/llm/src/kv_router/routing_host/tests.rs index 413690c6a432..d8399eb9a616 100644 --- a/lib/llm/src/kv_router/routing_host/tests.rs +++ b/lib/llm/src/kv_router/routing_host/tests.rs @@ -1305,7 +1305,7 @@ async fn router_with_recorded_dispatch_and_affinity( workers, None, 16, - DefaultWorkerSelector::new(Some(config.clone()), "decode"), + crate::kv_router::SelectionPolicySource::Registry, Some(config), None, "decode", @@ -2681,7 +2681,7 @@ async fn builtin_host_with_recorded_dispatch( ) .await .unwrap(); - let host = RoutingHost::::new_builtin(inner, load_context).unwrap(); + let host = RoutingHost::new_builtin(inner, load_context).unwrap(); (host, dispatch, worker_id, runtime) } From cd1a793ff968c99dd7e17703d1fc7ae3fc247738 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Thu, 3 Sep 2026 09:18:00 +0000 Subject: [PATCH 019/102] feat(epp): data-parallel worker registration in the standalone reflector Phase 4 of the runtime-optional frontend roadmap: the EPP pod reflector is the Kubernetes catalog feeder, and it registered every pod as a single data-parallel rank (the hardcoded data_parallel_size: 1 in selector.rs). - EppStandaloneConfig gains DYN_EPP_DATA_PARALLEL_SIZE (default 1) and DYN_EPP_KV_EVENT_PORT_STRIDE (default 1). Rank r publishes KV events on kv_event_port + r * stride; the reflector emits one KV-event endpoint per rank (RawWorker.kv_events_endpoints) and the registration carries data_parallel_size, which the selector passes through with data_parallel_start_rank 0. Tests: strided per-rank endpoints for a 3-rank pod; existing reflector, adapter, and selector tests updated to the per-rank shape (144 pass). Signed-off-by: Ishan Dhanani --- .../ext-proc/src/epp_standalone_config.rs | 10 + .../ext-proc/src/pod_discovery.rs | 203 ++++++++++++++---- .../ext-proc/src/selector.rs | 12 +- .../ext-proc/src/topology_adapter.rs | 13 +- .../kv-aware-routing/vanilla-vllm-onramp.mdx | 2 + 5 files changed, 194 insertions(+), 46 deletions(-) diff --git a/deploy/inference-gateway/ext-proc/src/epp_standalone_config.rs b/deploy/inference-gateway/ext-proc/src/epp_standalone_config.rs index cd2c48cc2402..874695f70db8 100644 --- a/deploy/inference-gateway/ext-proc/src/epp_standalone_config.rs +++ b/deploy/inference-gateway/ext-proc/src/epp_standalone_config.rs @@ -140,6 +140,13 @@ pub struct EppStandaloneConfig { /// KV-cache block size; MUST equal the inference engine block size. #[validate(range(min = 1, message = "DYN_KV_CACHE_BLOCK_SIZE must be >= 1"))] pub block_size: u32, + /// Data-parallel ranks per worker pod. Each rank publishes KV events on + /// `kv_event_port + rank * kv_event_port_stride`. + #[validate(range(min = 1, message = "DYN_EPP_DATA_PARALLEL_SIZE must be >= 1"))] + pub data_parallel_size: u32, + /// Port distance between consecutive data-parallel ranks' KV event ports. + #[validate(range(min = 1, message = "DYN_EPP_KV_EVENT_PORT_STRIDE must be >= 1"))] + pub kv_event_port_stride: u16, /// KV zmq event port. #[validate(range(min = 1))] pub kv_event_port: u16, @@ -216,6 +223,9 @@ impl EppStandaloneConfig { )? .unwrap_or(DEFAULT_TOKENIZER_MAX_RESPONSE_BYTES), block_size: opt_parse::(get, "DYN_KV_CACHE_BLOCK_SIZE")?.unwrap_or(0), + data_parallel_size: opt_parse::(get, "DYN_EPP_DATA_PARALLEL_SIZE")?.unwrap_or(1), + kv_event_port_stride: opt_parse::(get, "DYN_EPP_KV_EVENT_PORT_STRIDE")? + .unwrap_or(1), kv_event_port: opt_parse::(get, "DYN_EPP_KV_EVENT_PORT")? .unwrap_or(DEFAULT_KV_EVENT_PORT), replay_port: opt_parse::(get, "DYN_EPP_KV_EVENT_REPLAY_PORT")?, diff --git a/deploy/inference-gateway/ext-proc/src/pod_discovery.rs b/deploy/inference-gateway/ext-proc/src/pod_discovery.rs index 9acbea298e1a..d0f425968e15 100644 --- a/deploy/inference-gateway/ext-proc/src/pod_discovery.rs +++ b/deploy/inference-gateway/ext-proc/src/pod_discovery.rs @@ -32,8 +32,9 @@ pub struct RawWorker { pub pod_ip: String, /// OpenAI HTTP inference endpoint, `http://:`. pub http_endpoint: String, - /// Inference engine KV-event ZMQ PUB endpoint, `tcp://:`. - pub kv_events_endpoint: String, + /// Inference engine KV-event ZMQ PUB endpoints by global data-parallel rank, + /// `tcp://:`. + pub kv_events_endpoints: HashMap, /// Optional ZMQ REQ endpoint for live-stream gap replay. pub replay_endpoint: Option, } @@ -102,6 +103,11 @@ impl PodDiscovery { let kv_event_port = cfg.kv_event_port; let replay_port = cfg.replay_port; + let kv_ports = KvEventPorts { + base: kv_event_port, + stride: cfg.kv_event_port_stride, + data_parallel_size: cfg.data_parallel_size, + }; let index: Arc> = Arc::new(RwLock::new(WorkerIndex::new())); @@ -178,7 +184,7 @@ impl PodDiscovery { Delta::Rebuild => rebuild_index( &store, pool_rx.borrow().as_ref(), - kv_event_port, + kv_ports, replay_port, &index_task, ), @@ -186,7 +192,7 @@ impl PodDiscovery { &index_task, &pod, pool_rx.borrow().as_ref(), - kv_event_port, + kv_ports, replay_port, ), Delta::Remove(pod) => remove_pod(&index_task, &pod), @@ -339,14 +345,14 @@ fn upsert_pod( index: &RwLock, pod: &Pod, pool: Option<&PoolState>, - kv_event_port: u16, + kv_ports: KvEventPorts, replay_port: Option, ) -> bool { let Some(worker_id) = pod_worker_id(pod) else { return false; }; let entry = pool - .and_then(|pool| raw_worker_from_pod(pod, pool, kv_event_port, replay_port)) + .and_then(|pool| raw_worker_from_pod(pod, pool, kv_ports, replay_port)) .map(WorkerEntry::from_raw); let mut index = index.write().unwrap(); match entry { @@ -377,14 +383,14 @@ fn remove_pod(index: &RwLock, pod: &Pod) -> bool { fn rebuild_index( store: &kube::runtime::reflector::Store, pool: Option<&PoolState>, - kv_event_port: u16, + kv_ports: KvEventPorts, replay_port: Option, index: &RwLock, ) -> bool { let mut fresh = WorkerIndex::new(); if let Some(pool) = pool { for pod in store.state().iter() { - if let Some(worker) = raw_worker_from_pod(pod, pool, kv_event_port, replay_port) { + if let Some(worker) = raw_worker_from_pod(pod, pool, kv_ports, replay_port) { fresh.insert(worker.worker_id, WorkerEntry::from_raw(worker)); } } @@ -400,10 +406,31 @@ fn rebuild_index( /// Build a [`RawWorker`] from a pod, or `None` if it is not `Ready`, not /// pool-selected, or lacks an IP/name. Pure function — unit-testable. +/// Per-pod KV-event port layout: rank `r` publishes on `base + r * stride`. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct KvEventPorts { + pub base: u16, + pub stride: u16, + pub data_parallel_size: u32, +} + +impl KvEventPorts { + fn endpoints(self, ip: IpAddr) -> HashMap { + (0..self.data_parallel_size.max(1)) + .map(|rank| { + let port = self + .base + .saturating_add(self.stride.saturating_mul(rank as u16)); + (rank, format!("tcp://{}", SocketAddr::new(ip, port))) + }) + .collect() + } +} + fn raw_worker_from_pod( pod: &Pod, pool: &PoolState, - kv_event_port: u16, + kv_ports: KvEventPorts, replay_port: Option, ) -> Option { if !pod_is_ready(pod) || !pod_matches(pod, &pool.match_labels) { @@ -418,7 +445,7 @@ fn raw_worker_from_pod( pod_name: pod_name.to_string(), pod_ip: pod_ip.to_string(), http_endpoint: format!("http://{}", SocketAddr::new(ip, pool.target_port)), - kv_events_endpoint: format!("tcp://{}", SocketAddr::new(ip, kv_event_port)), + kv_events_endpoints: kv_ports.endpoints(ip), replay_endpoint: replay_port.map(|p| format!("tcp://{}", SocketAddr::new(ip, p))), }) } @@ -430,6 +457,42 @@ mod tests { use k8s_openapi::apimachinery::pkg::apis::meta::v1::Time; use kube::api::ObjectMeta; + fn single_rank(port: u16) -> KvEventPorts { + KvEventPorts { + base: port, + stride: 1, + data_parallel_size: 1, + } + } + + #[test] + fn data_parallel_ranks_publish_on_strided_ports() { + let w = raw_worker_from_pod( + &pod( + "vllm-0", + Some("10.0.0.1"), + Some(true), + &[("app", "vllm-qwen")], + ), + &pool(), + KvEventPorts { + base: 5557, + stride: 2, + data_parallel_size: 3, + }, + None, + ) + .expect("ready, selected pod should map"); + assert_eq!( + w.kv_events_endpoints, + HashMap::from([ + (0, "tcp://10.0.0.1:5557".to_string()), + (1, "tcp://10.0.0.1:5559".to_string()), + (2, "tcp://10.0.0.1:5561".to_string()), + ]) + ); + } + fn pool() -> PoolState { PoolState { match_labels: BTreeMap::from([("app".to_string(), "vllm-qwen".to_string())]), @@ -474,13 +537,16 @@ mod tests { &[("app", "vllm-qwen")], ), &pool(), - 5557, + single_rank(5557), Some(5560), ) .expect("ready, selected pod should map"); assert_eq!(w.worker_id, hash_pod_name("vllm-0")); assert_eq!(w.http_endpoint, "http://10.0.0.1:8000"); - assert_eq!(w.kv_events_endpoint, "tcp://10.0.0.1:5557"); + assert_eq!( + w.kv_events_endpoints, + HashMap::from([(0, "tcp://10.0.0.1:5557".to_string())]) + ); assert_eq!(w.replay_endpoint.as_deref(), Some("tcp://10.0.0.1:5560")); } @@ -494,13 +560,16 @@ mod tests { &[("app", "vllm-qwen")], ), &pool(), - 5557, + single_rank(5557), Some(5560), ) .expect("ready, selected IPv6 pod should map"); // SocketAddr brackets the IPv6 host so host and port are unambiguous. assert_eq!(w.http_endpoint, "http://[fd00::10]:8000"); - assert_eq!(w.kv_events_endpoint, "tcp://[fd00::10]:5557"); + assert_eq!( + w.kv_events_endpoints, + HashMap::from([(0, "tcp://[fd00::10]:5557".to_string())]) + ); assert_eq!(w.replay_endpoint.as_deref(), Some("tcp://[fd00::10]:5560")); } @@ -515,7 +584,7 @@ mod tests { &[("app", "vllm-qwen")] ), &pool(), - 5557, + single_rank(5557), None, ) .is_none() @@ -533,7 +602,7 @@ mod tests { &[("app", "something-else")] ), &pool(), - 5557, + single_rank(5557), None, ) .is_none() @@ -551,7 +620,7 @@ mod tests { &[("app", "vllm-qwen")] ), &pool(), - 5557, + single_rank(5557), None, ) .is_none() @@ -567,7 +636,7 @@ mod tests { &[("app", "vllm-qwen")], ); p.metadata.deletion_timestamp = Some(Time(k8s_openapi::chrono::Utc::now())); - assert!(raw_worker_from_pod(&p, &pool(), 5557, None).is_none()); + assert!(raw_worker_from_pod(&p, &pool(), single_rank(5557), None).is_none()); } #[test] @@ -576,7 +645,7 @@ mod tests { raw_worker_from_pod( &pod("vllm-0", None, Some(true), &[("app", "vllm-qwen")]), &pool(), - 5557, + single_rank(5557), None, ) .is_none() @@ -617,14 +686,14 @@ mod tests { assert!(rebuild_index( &store, Some(&pool()), - 5557, + single_rank(5557), Some(5560), &index )); assert!(!rebuild_index( &store, Some(&pool()), - 5557, + single_rank(5557), Some(5560), &index )); @@ -648,7 +717,13 @@ mod tests { &[("app", "vllm-qwen")], )]); let index = RwLock::new(WorkerIndex::new()); - assert!(!rebuild_index(&store, None, 5557, None, &index)); + assert!(!rebuild_index( + &store, + None, + single_rank(5557), + None, + &index + )); assert!(index.read().unwrap().is_empty()); } @@ -664,8 +739,20 @@ mod tests { ); // Ready + selected -> inserted, with a pre-stripped endpoint. - assert!(upsert_pod(&index, &ready, Some(&pool()), 5557, None)); - assert!(!upsert_pod(&index, &ready, Some(&pool()), 5557, None)); + assert!(upsert_pod( + &index, + &ready, + Some(&pool()), + single_rank(5557), + None + )); + assert!(!upsert_pod( + &index, + &ready, + Some(&pool()), + single_rank(5557), + None + )); assert_eq!( index.read().unwrap().get(&id).map(|e| e.endpoint.as_str()), Some("10.0.0.1:8000") @@ -678,12 +765,30 @@ mod tests { Some(false), &[("app", "vllm-qwen")], ); - assert!(upsert_pod(&index, ¬_ready, Some(&pool()), 5557, None)); - assert!(!upsert_pod(&index, ¬_ready, Some(&pool()), 5557, None)); + assert!(upsert_pod( + &index, + ¬_ready, + Some(&pool()), + single_rank(5557), + None + )); + assert!(!upsert_pod( + &index, + ¬_ready, + Some(&pool()), + single_rank(5557), + None + )); assert!(!index.read().unwrap().contains_key(&id)); // Re-add, then a Delete removes it. - assert!(upsert_pod(&index, &ready, Some(&pool()), 5557, None)); + assert!(upsert_pod( + &index, + &ready, + Some(&pool()), + single_rank(5557), + None + )); assert!(index.read().unwrap().contains_key(&id)); assert!(remove_pod(&index, &ready)); assert!(!remove_pod(&index, &ready)); @@ -691,7 +796,13 @@ mod tests { // An unrelated namespace pod does not change the derived worker index. let unselected = pod("other-0", Some("10.0.0.2"), Some(true), &[("app", "other")]); - assert!(!upsert_pod(&index, &unselected, Some(&pool()), 5557, None)); + assert!(!upsert_pod( + &index, + &unselected, + Some(&pool()), + single_rank(5557), + None + )); } #[test] @@ -707,11 +818,17 @@ mod tests { &[("app", "vllm-qwen")], ); - assert!(upsert_pod(&index, &ready, Some(&pool()), 5557, None)); + assert!(upsert_pod( + &index, + &ready, + Some(&pool()), + single_rank(5557), + None + )); assert!(index.read().unwrap().contains_key(&id)); - assert!(upsert_pod(&index, &ready, None, 5557, None)); - assert!(!upsert_pod(&index, &ready, None, 5557, None)); + assert!(upsert_pod(&index, &ready, None, single_rank(5557), None)); + assert!(!upsert_pod(&index, &ready, None, single_rank(5557), None)); assert!(!index.read().unwrap().contains_key(&id)); } @@ -740,7 +857,13 @@ mod tests { writer.apply_watcher_event(&watcher::Event::InitApply(vllm_1.clone())); writer.apply_watcher_event(&watcher::Event::InitDone); let index = RwLock::new(WorkerIndex::new()); - assert!(rebuild_index(&store, Some(&pool()), 5557, None, &index)); + assert!(rebuild_index( + &store, + Some(&pool()), + single_rank(5557), + None, + &index + )); assert_eq!(index.read().unwrap().len(), 2); // During the next LIST, vllm-1 has disappeared. The reflector buffers @@ -769,7 +892,7 @@ mod tests { assert!(rebuild_index( &store, Some(&updated_pool), - 5557, + single_rank(5557), None, &index )); @@ -804,7 +927,7 @@ mod tests { &[("app", "vllm-qwen")], ), Some(&pool()), - 5557, + single_rank(5557), None, ); upsert_pod( @@ -816,7 +939,7 @@ mod tests { &[("app", "vllm-qwen")], ), Some(&pool()), - 5557, + single_rank(5557), None, ); assert_eq!(index.read().unwrap().len(), 2); @@ -828,7 +951,13 @@ mod tests { Some(true), &[("app", "vllm-qwen")], )]); - assert!(rebuild_index(&store, Some(&pool()), 5557, None, &index)); + assert!(rebuild_index( + &store, + Some(&pool()), + single_rank(5557), + None, + &index + )); let index = index.read().unwrap(); assert_eq!(index.len(), 1); @@ -848,7 +977,7 @@ mod tests { .map_or(endpoint, |(ip, _)| ip) .to_string(), http_endpoint: format!("http://{endpoint}"), - kv_events_endpoint: format!("tcp://{endpoint}"), + kv_events_endpoints: HashMap::from([(0, format!("tcp://{endpoint}"))]), replay_endpoint: None, } } diff --git a/deploy/inference-gateway/ext-proc/src/selector.rs b/deploy/inference-gateway/ext-proc/src/selector.rs index 1927256f56eb..380187c1c420 100644 --- a/deploy/inference-gateway/ext-proc/src/selector.rs +++ b/deploy/inference-gateway/ext-proc/src/selector.rs @@ -35,6 +35,9 @@ pub struct WorkerRegistration { pub model_name: String, pub endpoint: String, pub block_size: u32, + /// Data-parallel ranks served by this worker; `kv_events_endpoints` carries + /// one KV-event endpoint per rank. + pub data_parallel_size: u32, pub kv_events_endpoints: HashMap, pub replay_endpoint: Option, pub total_kv_blocks: Option, @@ -197,9 +200,8 @@ impl Selector { routing_group: DEFAULT_ROUTING_GROUP.to_string(), endpoint: Some(reg.endpoint.clone()), block_size: Some(reg.block_size), - // Data parallel size is not yet implemented. - // Default support for single rank DP. - data_parallel_size: Some(1), + data_parallel_start_rank: Some(0), + data_parallel_size: Some(reg.data_parallel_size.max(1)), kv_events_endpoints: reg.kv_events_endpoints.clone(), replay_endpoint: reg.replay_endpoint.clone(), total_kv_blocks: reg.total_kv_blocks, @@ -450,6 +452,8 @@ models: tokenizer_max_response_bytes: 16 * 1024 * 1024, tokenization_timeout_ms: 5_000, block_size: 16, + data_parallel_size: 1, + kv_event_port_stride: 1, kv_event_port: 5557, replay_port: None, total_kv_blocks: None, @@ -467,6 +471,7 @@ models: model_name: "test-model".to_string(), endpoint: "http://10.0.0.1:8000".to_string(), block_size: 0, + data_parallel_size: 1, kv_events_endpoints: HashMap::new(), replay_endpoint: None, total_kv_blocks: None, @@ -488,6 +493,7 @@ models: model_name: "test-model".to_string(), endpoint: format!("http://10.0.0.{worker_id}:8000"), block_size: 16, + data_parallel_size: 1, kv_events_endpoints: HashMap::from([( 0u32, format!("tcp://127.0.0.1:{}", 45_000 + worker_id), diff --git a/deploy/inference-gateway/ext-proc/src/topology_adapter.rs b/deploy/inference-gateway/ext-proc/src/topology_adapter.rs index 9fac749192d5..c95bef0f80d9 100644 --- a/deploy/inference-gateway/ext-proc/src/topology_adapter.rs +++ b/deploy/inference-gateway/ext-proc/src/topology_adapter.rs @@ -7,7 +7,6 @@ //! resolved endpoints and configured defaults, then passes the desired worker //! set to the [`Selector`]. -use std::collections::HashMap; use std::sync::Arc; use tokio_util::sync::CancellationToken; @@ -104,15 +103,14 @@ async fn reconcile_once( } fn build_registration(w: RawWorker, defaults: &RegistrationDefaults) -> WorkerRegistration { - let mut kv_events_endpoints = HashMap::new(); - kv_events_endpoints.insert(0u32, w.kv_events_endpoint); - + let data_parallel_size = (w.kv_events_endpoints.len() as u32).max(1); WorkerRegistration { worker_id: w.worker_id, model_name: defaults.model_name.clone(), endpoint: w.http_endpoint, block_size: defaults.block_size, - kv_events_endpoints, + data_parallel_size, + kv_events_endpoints: w.kv_events_endpoints, replay_endpoint: w.replay_endpoint, total_kv_blocks: defaults.total_kv_blocks, max_num_batched_tokens: defaults.max_num_batched_tokens, @@ -121,6 +119,7 @@ fn build_registration(w: RawWorker, defaults: &RegistrationDefaults) -> WorkerRe #[cfg(test)] mod tests { + use std::collections::HashMap; use std::time::Duration; use super::*; @@ -138,6 +137,8 @@ mod tests { tokenizer_max_response_bytes: 16 * 1024 * 1024, tokenization_timeout_ms: 5_000, block_size: 16, + data_parallel_size: 1, + kv_event_port_stride: 1, kv_event_port: 5557, replay_port: None, total_kv_blocks: Some(1000), @@ -161,7 +162,7 @@ mod tests { pod_name: format!("vllm-{id}"), pod_ip: ip.to_string(), http_endpoint: format!("http://{ip}:8000"), - kv_events_endpoint: format!("tcp://{ip}:5557"), + kv_events_endpoints: HashMap::from([(0, format!("tcp://{ip}:5557"))]), replay_endpoint: None, } } diff --git a/docs/fern/pages/kubernetes/kv-aware-routing/vanilla-vllm-onramp.mdx b/docs/fern/pages/kubernetes/kv-aware-routing/vanilla-vllm-onramp.mdx index ad6f55d959dd..7faccae2c1a1 100644 --- a/docs/fern/pages/kubernetes/kv-aware-routing/vanilla-vllm-onramp.mdx +++ b/docs/fern/pages/kubernetes/kv-aware-routing/vanilla-vllm-onramp.mdx @@ -497,6 +497,8 @@ For a single EPP replica, set `spec.replicas: 1` and remove `DYN_EPP_PEER_SERVIC | `DYN_KV_CACHE_BLOCK_SIZE` | Yes | Must equal vLLM `--block-size`. | | `DYN_EPP_KV_EVENT_PORT` | No | Per-pod vLLM KV event port; defaults to `5557`. | | `DYN_EPP_KV_EVENT_REPLAY_PORT` | No | Per-pod replay port for KV event gap recovery. | +| `DYN_EPP_DATA_PARALLEL_SIZE` | No | Data-parallel ranks per pod; defaults to `1`. Rank `r` publishes KV events on `DYN_EPP_KV_EVENT_PORT + r * DYN_EPP_KV_EVENT_PORT_STRIDE`. | +| `DYN_EPP_KV_EVENT_PORT_STRIDE` | No | Port distance between consecutive ranks' KV event ports; defaults to `1`. | | `DYN_EPP_MAX_NUM_BATCHED_TOKENS` | Policy-dependent | Must equal vLLM `--max-num-batched-tokens`; required when router queue thresholds are enabled. | | `DYN_EPP_TOTAL_KV_BLOCKS` | No | Optional per-worker total KV block hint. | | `DYN_EPP_SELECTION_INDEXER_THREADS` | No | KV indexer thread count; defaults to `4`. | From d140d9c704ac8a60093b2425e49edaa0d18eaecd Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Thu, 3 Sep 2026 21:23:38 +0000 Subject: [PATCH 020/102] feat(kv-router): WorkerCatalogSource and CatalogReconciler for worker membership A host owns discovery and exposes it as a WorkerCatalogSource: a stream of complete desired membership snapshots. CatalogReconciler applies each snapshot to a SelectionCore: it upserts new or changed records, retries records the core reported not schedulable, deletes workers that left (including ones that never converged), rejects duplicate ids before mutating anything, and clears the catalog when the source closes so selection fails closed. A CatalogObserver lets a host keep per-worker metrics in step. This is the reconcile loop the EPP standalone selector and the frontend's embedded partition each carried; both switch to it in the next commits. SelectionService exposes its core for hosts that drive the catalog directly. Signed-off-by: Ishan Dhanani --- .../src/services/selection/membership.rs | 243 ++++++++++++++++++ lib/kv-router/src/services/selection/mod.rs | 2 + .../src/services/selection/service.rs | 5 + lib/kv-router/src/services/selection/types.rs | 2 +- 4 files changed, 251 insertions(+), 1 deletion(-) create mode 100644 lib/kv-router/src/services/selection/membership.rs diff --git a/lib/kv-router/src/services/selection/membership.rs b/lib/kv-router/src/services/selection/membership.rs new file mode 100644 index 000000000000..b485418a2e16 --- /dev/null +++ b/lib/kv-router/src/services/selection/membership.rs @@ -0,0 +1,243 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Worker membership for a selection core. +//! +//! A host owns discovery (runtime discovery, a Kubernetes pod reflector, a +//! workers file) and exposes it as a [`WorkerCatalogSource`]: a stream of +//! complete desired snapshots. [`CatalogReconciler`] turns each snapshot into +//! catalog upserts and deletes, retrying workers the core has not yet accepted +//! as schedulable and removing workers that left. + +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; + +use async_trait::async_trait; +use tokio_util::sync::CancellationToken; + +use super::core::SelectionCore; +use super::error::SelectionError; +use super::types::{WorkerCatalogRecord, WorkerLifecycle, WorkerRequest}; +use crate::protocols::WorkerId; + +/// Desired worker membership, delivered as complete snapshots rather than +/// deltas so a missed change can never leave a stale worker behind. +#[async_trait] +pub trait WorkerCatalogSource: Send { + /// The next desired snapshot, or `None` once the source is closed. The + /// first call returns the current membership without waiting. + async fn next_snapshot(&mut self) -> Option>; +} + +/// Notified after the reconciler applies a catalog change (metrics hooks). +pub trait CatalogObserver: Send + Sync { + fn upserted(&self, record: &WorkerCatalogRecord); + fn removed(&self, record: &WorkerCatalogRecord); +} + +/// Applies desired snapshots to a core's worker catalog. +pub struct CatalogReconciler { + core: Arc, + observer: Option>, + /// Requests the core reported `Schedulable`. An identical desired record + /// skips its upsert; anything else is retried on every snapshot. + converged: HashMap, + /// Every worker id this reconciler introduced, including ones that never + /// became schedulable, so stale deletion covers partial upserts. + tracked: HashSet, +} + +impl CatalogReconciler { + pub fn new(core: Arc) -> Self { + Self { + core, + observer: None, + converged: HashMap::new(), + tracked: HashSet::new(), + } + } + + pub fn with_observer(mut self, observer: Arc) -> Self { + self.observer = Some(observer); + self + } + + /// Apply one desired snapshot. A snapshot with duplicate worker ids is + /// rejected before any catalog mutation. A failed upsert or delete stops + /// the pass; the next snapshot retries it. + pub async fn apply(&mut self, desired: Vec) -> Result<(), SelectionError> { + let mut by_id: HashMap = HashMap::with_capacity(desired.len()); + for request in desired { + let worker_id = request.worker_id; + if by_id.insert(worker_id, request).is_some() { + return Err(SelectionError::BadRequest(format!( + "duplicate worker_id {worker_id} in membership snapshot" + ))); + } + } + + for (worker_id, request) in &by_id { + // Track before the upsert so a partially applied record is still + // deleted when the worker leaves the desired set. + self.tracked.insert(*worker_id); + if self.converged.get(worker_id) == Some(request) { + continue; + } + let record = self.core.upsert_worker(request.clone()).await?; + if let Some(observer) = &self.observer { + observer.upserted(&record); + } + if record.lifecycle == WorkerLifecycle::Schedulable { + self.converged.insert(*worker_id, request.clone()); + } else { + self.converged.remove(worker_id); + tracing::warn!( + worker_id, + lifecycle = ?record.lifecycle, + reasons = ?record.not_schedulable_reasons, + "worker upserted but not schedulable; retrying on the next membership snapshot" + ); + } + } + + let stale: Vec = self + .tracked + .iter() + .copied() + .filter(|worker_id| !by_id.contains_key(worker_id)) + .collect(); + for worker_id in stale { + match self.core.delete_worker(worker_id).await { + Ok(record) => { + if let Some(observer) = &self.observer { + observer.removed(&record); + } + } + Err(SelectionError::NotFound(_)) => {} + Err(error) => return Err(error), + } + self.tracked.remove(&worker_id); + self.converged.remove(&worker_id); + } + Ok(()) + } + + /// Apply every snapshot `source` yields until it closes or `cancel` fires. + /// The catalog keeps its last state when the source closes; a source that + /// wants selection to fail closed yields an empty snapshot before closing. + pub async fn run(mut self, mut source: impl WorkerCatalogSource, cancel: CancellationToken) { + loop { + let snapshot = tokio::select! { + _ = cancel.cancelled() => return, + snapshot = source.next_snapshot() => snapshot, + }; + let Some(snapshot) = snapshot else { + tracing::debug!("worker membership source closed"); + return; + }; + if let Err(error) = self.apply(snapshot).await { + tracing::warn!(%error, "membership reconcile failed; retrying on the next snapshot"); + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::config::KvRouterConfig; + use crate::services::selection::SelectionCacheConfig; + + fn core() -> Arc { + Arc::new(SelectionCore::new_local( + KvRouterConfig { + use_kv_events: false, + router_queue_threshold: None, + ..Default::default() + }, + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + )) + } + + fn schedulable(worker_id: WorkerId) -> WorkerRequest { + WorkerRequest { + worker_id, + endpoint: Some(format!("http://10.0.0.{worker_id}:8000")), + block_size: Some(16), + ..WorkerRequest::default() + } + } + + /// `block_size = 0` fails the schedulable-metadata check regardless of + /// router or KV-event configuration. + fn incomplete(worker_id: WorkerId) -> WorkerRequest { + WorkerRequest { + block_size: Some(0), + ..schedulable(worker_id) + } + } + + fn lifecycle(core: &SelectionCore, worker_id: WorkerId) -> Option { + core.list_workers(None, None) + .into_iter() + .find(|record| record.worker_id == worker_id) + .map(|record| record.lifecycle) + } + + #[tokio::test] + async fn incomplete_worker_is_retried_and_deleted_when_it_leaves() { + let core = core(); + let mut reconciler = CatalogReconciler::new(Arc::clone(&core)); + + reconciler.apply(vec![incomplete(1)]).await.expect("apply"); + assert!(reconciler.converged.is_empty()); + assert_eq!(lifecycle(&core, 1), Some(WorkerLifecycle::Incomplete)); + + // The same snapshot re-upserts rather than skipping the unconverged worker. + reconciler.apply(vec![incomplete(1)]).await.expect("apply"); + assert!(reconciler.converged.is_empty()); + assert!(reconciler.tracked.contains(&1)); + + reconciler.apply(Vec::new()).await.expect("apply"); + assert_eq!(lifecycle(&core, 1), Some(WorkerLifecycle::Unschedulable)); + assert!(reconciler.tracked.is_empty()); + } + + #[tokio::test] + async fn schedulable_worker_converges_and_changed_record_reupserts() { + let core = core(); + let mut reconciler = CatalogReconciler::new(Arc::clone(&core)); + + reconciler.apply(vec![schedulable(1)]).await.expect("apply"); + assert_eq!(lifecycle(&core, 1), Some(WorkerLifecycle::Schedulable)); + assert!(reconciler.converged.contains_key(&1)); + + let mut moved = schedulable(1); + moved.endpoint = Some("http://10.0.0.9:8000".to_string()); + reconciler.apply(vec![moved]).await.expect("apply"); + let record = core + .list_workers(None, None) + .into_iter() + .find(|record| record.worker_id == 1) + .expect("record"); + assert_eq!(record.endpoint.as_deref(), Some("http://10.0.0.9:8000")); + assert_eq!(record.lifecycle, WorkerLifecycle::Schedulable); + } + + #[tokio::test] + async fn duplicate_ids_are_rejected_before_any_mutation() { + let core = core(); + let mut reconciler = CatalogReconciler::new(Arc::clone(&core)); + + let error = reconciler + .apply(vec![schedulable(1), schedulable(1)]) + .await + .expect_err("duplicates are rejected"); + assert!( + matches!(error, SelectionError::BadRequest(message) if message.contains("duplicate worker_id 1")) + ); + assert!(core.list_workers(None, None).is_empty()); + } +} diff --git a/lib/kv-router/src/services/selection/mod.rs b/lib/kv-router/src/services/selection/mod.rs index c5e9247c19f8..41b60e53ea8c 100644 --- a/lib/kv-router/src/services/selection/mod.rs +++ b/lib/kv-router/src/services/selection/mod.rs @@ -10,6 +10,7 @@ mod catalog; mod core; mod error; mod input; +mod membership; mod pending; mod policy_registry; mod server; @@ -31,6 +32,7 @@ pub use core::{ }; pub use error::SelectionError; pub use input::PromptRequest; +pub use membership::{CatalogObserver, CatalogReconciler, WorkerCatalogSource}; pub use pending::SelectionCacheConfig; pub use policy_registry::{ DYN_ROUTER_DECODE_POLICY, DYN_ROUTER_PREFILL_POLICY, DYN_ROUTER_WORKER_SELECTION_POLICY, diff --git a/lib/kv-router/src/services/selection/service.rs b/lib/kv-router/src/services/selection/service.rs index ce24dacfe0e4..58046c613d00 100644 --- a/lib/kv-router/src/services/selection/service.rs +++ b/lib/kv-router/src/services/selection/service.rs @@ -319,6 +319,11 @@ impl SelectionService { self.core.upsert_worker(req).await } + /// The core this service wraps, for hosts that drive its catalog directly. + pub fn core(&self) -> &Arc { + &self.core + } + /// Scheduler and indexer handle for `key`, once a worker has been upserted /// into that partition. Embedding hosts that keep their own index drive /// the partition scheduler through this handle. diff --git a/lib/kv-router/src/services/selection/types.rs b/lib/kv-router/src/services/selection/types.rs index 05b2f0896853..aeef0590f26e 100644 --- a/lib/kv-router/src/services/selection/types.rs +++ b/lib/kv-router/src/services/selection/types.rs @@ -295,7 +295,7 @@ impl Default for WorkerRequest { } } -#[derive(Debug, Deserialize)] +#[derive(Debug, Clone, PartialEq, Deserialize)] pub struct WorkerRequest { pub worker_id: WorkerId, #[serde(default = "default_model_name")] From 2efe0002c1d634e8b8cf507c778c487a2fddb850 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Thu, 3 Sep 2026 21:55:44 +0000 Subject: [PATCH 021/102] refactor(router): feed the embedded partition's catalog through CatalogReconciler The runtime-config watch becomes a WorkerCatalogSource (RuntimeDiscoverySource) and the partition's catalog is driven by the shared CatalogReconciler instead of the router-local CatalogFeeder. The per-worker router_worker_registered gauge is kept through a CatalogObserver. Behavior is unchanged except that a worker the core reports not schedulable is now retried on every membership snapshot rather than only when its runtime config changes. Signed-off-by: Ishan Dhanani --- lib/llm/src/kv_router/embedded.rs | 156 +++++++++++++----------------- 1 file changed, 68 insertions(+), 88 deletions(-) diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs index 408b571393f1..1811875e1383 100644 --- a/lib/llm/src/kv_router/embedded.rs +++ b/lib/llm/src/kv_router/embedded.rs @@ -28,8 +28,9 @@ use dynamo_kv_router::scheduling::{ }; use dynamo_kv_router::sequences::{SequenceError, SequenceRequest}; use dynamo_kv_router::services::selection::{ - HostCache, HostEligibility, HostLoad, HostReplication, HostTelemetry, OverlapRefreshSource, - SelectionHost, SelectionPartition, SelectionService, SelectionServiceBuilder, WorkerRequest, + CatalogObserver, CatalogReconciler, HostCache, HostEligibility, HostLoad, HostReplication, + HostTelemetry, OverlapRefreshSource, SelectionHost, SelectionPartition, SelectionService, + SelectionServiceBuilder, WorkerCatalogRecord, WorkerCatalogSource, WorkerRequest, WorkerSelectionPolicyRegistry, }; use dynamo_kv_router::{DEFAULT_ROUTING_GROUP, PrefillLoadEstimator, WorkerSelectionPolicyFactory}; @@ -177,20 +178,28 @@ impl EmbeddedSelection { .ensure_partition(key.clone(), args.block_size, args.is_eagle) .context("failed to create embedded selection partition")?; - let feeder = CatalogFeeder { - service: Arc::clone(&service), + let mut source = RuntimeDiscoverySource { + watch: workers_with_configs, key, block_size: args.block_size, is_eagle: args.is_eagle, - known: HashMap::new(), - status_metrics: super::metrics::RouterWorkerStatusMetrics::from_component( - args.endpoint.component(), - ), - worker_label: args.metric_worker_type, + primed: false, }; - feeder - .run(workers_with_configs, cancellation_token.child_token()) - .await; + let mut reconciler = CatalogReconciler::new(Arc::clone(service.core())).with_observer( + Arc::new(RegisteredGauge { + metrics: super::metrics::RouterWorkerStatusMetrics::from_component( + args.endpoint.component(), + ), + worker_label: args.metric_worker_type, + }), + ); + // The current membership is in the catalog before the router serves. + if let Some(snapshot) = source.next_snapshot().await + && let Err(error) = reconciler.apply(snapshot).await + { + tracing::warn!(%error, "embedded selection: initial membership reconcile failed"); + } + tokio::spawn(reconciler.run(source, cancellation_token.child_token())); tracing::info!( worker_type = %worker_type, @@ -326,92 +335,63 @@ impl EmbeddedSelection { } } -/// Mirrors the runtime-config watch into the partition's worker catalog. -struct CatalogFeeder { - service: Arc, +/// The runtime-config watch as worker membership for the partition. +struct RuntimeDiscoverySource { + watch: RuntimeConfigWatch, key: RoutingPartitionId, block_size: u32, is_eagle: bool, - known: HashMap, - /// `router_worker_registered` gauge per worker/dp_rank. - status_metrics: Arc, + primed: bool, +} + +#[async_trait::async_trait] +impl WorkerCatalogSource for RuntimeDiscoverySource { + async fn next_snapshot(&mut self) -> Option> { + if self.primed { + self.watch.changed().await.ok()?; + } + self.primed = true; + let snapshot = self + .watch + .borrow_and_update() + .iter() + .map(|(worker_id, config)| { + worker_request_from_runtime_config( + *worker_id, + config, + &self.key, + self.block_size, + self.is_eagle, + ) + }) + .collect(); + Some(snapshot) + } +} + +/// Keeps the per-worker `router_worker_registered` gauge in step with the catalog. +struct RegisteredGauge { + metrics: Arc, worker_label: &'static str, } -fn dp_ranks(config: &ModelRuntimeConfig) -> std::ops::Range { - let start = config.data_parallel_start_rank; - start..start.saturating_add(config.data_parallel_size.max(1)) +fn dp_ranks(record: &WorkerCatalogRecord) -> std::ops::Range { + let start = record.data_parallel_start_rank.unwrap_or(0); + start..start.saturating_add(record.data_parallel_size.unwrap_or(1).max(1)) } -impl CatalogFeeder { - /// Apply the current snapshot, then keep applying changes until cancelled. - async fn run(mut self, mut watch: RuntimeConfigWatch, cancel: CancellationToken) { - let snapshot = watch.borrow_and_update().clone(); - self.apply(snapshot).await; - tokio::spawn(async move { - loop { - tokio::select! { - _ = cancel.cancelled() => break, - changed = watch.changed() => { - if changed.is_err() { - break; - } - } - } - let snapshot = watch.borrow_and_update().clone(); - self.apply(snapshot).await; - } - }); +impl CatalogObserver for RegisteredGauge { + fn upserted(&self, record: &WorkerCatalogRecord) { + for dp_rank in dp_ranks(record) { + self.metrics + .set_registered(record.worker_id, dp_rank, self.worker_label); + } } - async fn apply(&mut self, snapshot: HashMap) { - let removed: Vec = self - .known - .keys() - .copied() - .filter(|worker_id| !snapshot.contains_key(worker_id)) - .collect(); - for worker_id in removed { - if let Some(previous) = self.known.remove(&worker_id) { - for dp_rank in dp_ranks(&previous) { - self.status_metrics - .remove_worker(worker_id, dp_rank, self.worker_label); - } - } - if let Err(error) = self.service.delete_worker(worker_id).await { - tracing::debug!(worker_id, %error, "embedded selection: worker delete skipped"); - } - } - for (worker_id, config) in snapshot { - if self.known.get(&worker_id) == Some(&config) { - continue; - } - let request = worker_request_from_runtime_config( - worker_id, - &config, - &self.key, - self.block_size, - self.is_eagle, - ); - match self.service.upsert_worker(request).await { - Ok(record) => { - if !record.not_schedulable_reasons.is_empty() { - tracing::warn!( - worker_id, - reasons = ?record.not_schedulable_reasons, - "embedded selection: worker is not schedulable" - ); - } - for dp_rank in dp_ranks(&config) { - self.status_metrics - .set_registered(worker_id, dp_rank, self.worker_label); - } - self.known.insert(worker_id, config); - } - Err(error) => { - tracing::warn!(worker_id, %error, "embedded selection: worker upsert failed"); - } - } + fn removed(&self, record: &WorkerCatalogRecord) { + for dp_rank in dp_ranks(record) { + self.metrics + .remove_worker(record.worker_id, dp_rank, self.worker_label); } } } From dab4d2b81970febf39416c3c02e90f7f10a59a0d Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Thu, 3 Sep 2026 21:58:10 +0000 Subject: [PATCH 022/102] refactor(epp): expose the pod reflector as a WorkerCatalogSource The standalone EPP's topology adapter yields the reflector's Ready pods as WorkerRequest snapshots (PodReflectorSource) and lets CatalogReconciler own the desired-versus-actual diff. Selector::reconcile, WorkerRegistration, and the EPP-local reconcile state are removed; their tests moved to dynamo-kv-router with the reconciler. A closed reflector channel still yields one empty snapshot so selection fails closed, as before. Signed-off-by: Ishan Dhanani --- Cargo.lock | 1 + deploy/inference-gateway/ext-proc/Cargo.toml | 1 + deploy/inference-gateway/ext-proc/src/lib.rs | 2 +- .../ext-proc/src/selector.rs | 255 +++--------------- .../ext-proc/src/topology_adapter.rs | 135 +++++----- .../src/services/selection/README.md | 4 + 6 files changed, 116 insertions(+), 282 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 895b6658bc6d..51883a61456d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2542,6 +2542,7 @@ name = "dynamo-ext-proc" version = "1.5.0" dependencies = [ "anyhow", + "async-trait", "axum 0.8.4", "bytes", "dashmap", diff --git a/deploy/inference-gateway/ext-proc/Cargo.toml b/deploy/inference-gateway/ext-proc/Cargo.toml index b4eb1c410153..de1d14186130 100644 --- a/deploy/inference-gateway/ext-proc/Cargo.toml +++ b/deploy/inference-gateway/ext-proc/Cargo.toml @@ -16,6 +16,7 @@ name = "dynamo-ext-proc" path = "src/bin/dynamo-ext-proc.rs" [dependencies] +async-trait = { workspace = true } dynamo-kv-router = { workspace = true, features = [ "metrics", "runtime-protocols", diff --git a/deploy/inference-gateway/ext-proc/src/lib.rs b/deploy/inference-gateway/ext-proc/src/lib.rs index 38543d1f764d..2fc951d887b4 100644 --- a/deploy/inference-gateway/ext-proc/src/lib.rs +++ b/deploy/inference-gateway/ext-proc/src/lib.rs @@ -39,7 +39,7 @@ pub use picker::{Endpoint, EndpointPicker, PickResult, RequestInfo, ResponseUsag pub use pod_discovery::{PodDiscovery, RawWorker}; pub use render_http::RenderError; pub use runner::run; -pub use selector::{OverlapSummary, SelectRequest, SelectResponse, Selector, WorkerRegistration}; +pub use selector::{OverlapSummary, SelectRequest, SelectResponse, Selector}; pub use server::ExtProcServer; pub use sglang_renderer_client::SglangRendererClient; pub use topology_adapter::{RegistrationDefaults, TopologyAdapter}; diff --git a/deploy/inference-gateway/ext-proc/src/selector.rs b/deploy/inference-gateway/ext-proc/src/selector.rs index 380187c1c420..74d165a4d3ab 100644 --- a/deploy/inference-gateway/ext-proc/src/selector.rs +++ b/deploy/inference-gateway/ext-proc/src/selector.rs @@ -3,47 +3,28 @@ //! In-process worker selection for standalone mode. //! -//! Wraps [`SelectionService`] and defines its worker registration and selection -//! types. Optionally synchronizes active load across EPP replicas through -//! [`crate::peer_discovery`]. +//! Wraps [`SelectionService`] and defines its selection types; worker +//! membership is fed by [`crate::topology_adapter`]. Optionally synchronizes +//! active load across EPP replicas through [`crate::peer_discovery`]. -use std::collections::{HashMap, HashSet}; +use std::collections::HashSet; use std::sync::Arc; use anyhow::{Context, Result, anyhow}; -use dynamo_kv_router::WorkerType; use dynamo_kv_router::config::{KvRouterConfig, try_kv_router_config_from_dynamo_env}; use dynamo_kv_router::protocols::RoutingConstraints; use dynamo_kv_router::services::selection::{ - PromptRequest, SelectAndReserveRequest as CoreSelectAndReserveRequest, SelectionError, - SelectionService, SelectionServiceBuilder, WorkerLifecycle, WorkerRequest as CoreWorkerRequest, - WorkerSelectionPolicyRegistry, warn_for_unserved_worker_selection_policies, + PromptRequest, SelectAndReserveRequest as CoreSelectAndReserveRequest, SelectionCore, + SelectionError, SelectionService, SelectionServiceBuilder, WorkerSelectionPolicyRegistry, + warn_for_unserved_worker_selection_policies, }; +use dynamo_kv_router::{DEFAULT_ROUTING_GROUP, WorkerType}; use kube::Client; -use tokio::sync::Mutex; use tokio_util::sync::CancellationToken; use crate::epp_standalone_config::EppStandaloneConfig; -const DEFAULT_ROUTING_GROUP: &str = "default"; - -/// A worker the EPP registers into the selector. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct WorkerRegistration { - pub worker_id: u64, - pub model_name: String, - pub endpoint: String, - pub block_size: u32, - /// Data-parallel ranks served by this worker; `kv_events_endpoints` carries - /// one KV-event endpoint per rank. - pub data_parallel_size: u32, - pub kv_events_endpoints: HashMap, - pub replay_endpoint: Option, - pub total_kv_blocks: Option, - pub max_num_batched_tokens: Option, -} - /// A worker-selection request. #[derive(Debug, Clone)] pub struct SelectRequest { @@ -86,19 +67,6 @@ pub struct Selector { /// Cancels the peer-discovery watch on drop. The `SelectionService`'s own /// `Drop` tears down its core + replica-sync tasks. cancel: CancellationToken, - reconcile_state: Mutex, -} - -/// Local bookkeeping for desired-state reconciliation. -#[derive(Default)] -struct ReconcileState { - /// Registrations confirmed schedulable by the service. Identical desired - /// entries can skip upsert without preventing `Incomplete` retries. - converged: HashMap, - /// Every worker ID this reconciler may have introduced into the service, - /// including `Incomplete` and partially failed upserts. Stale detection must - /// cover this set rather than only the schedulable convergence cache. - tracked_worker_ids: HashSet, } impl Selector { @@ -168,11 +136,7 @@ impl Selector { "Initialized in-process selection service" ); - Ok(Self { - service, - cancel, - reconcile_state: Mutex::new(ReconcileState::default()), - }) + Ok(Self { service, cancel }) } fn validate_queueing_worker_capacity( @@ -193,85 +157,9 @@ impl Selector { Ok(()) } - fn worker_request(reg: &WorkerRegistration) -> CoreWorkerRequest { - CoreWorkerRequest { - worker_id: reg.worker_id, - model_name: reg.model_name.clone(), - routing_group: DEFAULT_ROUTING_GROUP.to_string(), - endpoint: Some(reg.endpoint.clone()), - block_size: Some(reg.block_size), - data_parallel_start_rank: Some(0), - data_parallel_size: Some(reg.data_parallel_size.max(1)), - kv_events_endpoints: reg.kv_events_endpoints.clone(), - replay_endpoint: reg.replay_endpoint.clone(), - total_kv_blocks: reg.total_kv_blocks, - max_num_batched_tokens: reg.max_num_batched_tokens, - ..Default::default() - } - } - - pub async fn reconcile(&self, registrations: &[WorkerRegistration]) -> Result<()> { - // Derive the bookkeeping key from the registration so callers cannot - // provide a map key that disagrees with the ID sent to SelectionService. - // Reject duplicates before mutating either local or service state. - let mut desired = HashMap::with_capacity(registrations.len()); - for registration in registrations { - if desired - .insert(registration.worker_id, registration) - .is_some() - { - anyhow::bail!("duplicate worker_id {}", registration.worker_id); - } - } - - let mut state = self.reconcile_state.lock().await; - - // Upsert new or changed workers. - for (&worker_id, ®) in &desired { - // Track before calling the service so a partially applied upsert can - // still be deleted if the worker later leaves the desired snapshot. - state.tracked_worker_ids.insert(worker_id); - if state.converged.get(&worker_id) == Some(reg) { - continue; - } - let record = self - .service - .upsert_worker(Self::worker_request(reg)) - .await - .map_err(|e| anyhow!("upsert_worker failed: {e}"))?; - // Only cache the registration once the core reports the worker Schedulable. - if record.lifecycle == WorkerLifecycle::Schedulable { - state.converged.insert(worker_id, reg.clone()); - } else { - state.converged.remove(&worker_id); - tracing::warn!( - worker_id, - lifecycle = ?record.lifecycle, - reasons = ?record.not_schedulable_reasons, - "Worker upserted but not schedulable; leaving uncached to retry on the next reconcile" - ); - } - } - - // Delete every worker this reconciler may have introduced that is no - // longer desired, including workers that never became schedulable. - let stale: Vec = state - .tracked_worker_ids - .iter() - .copied() - .filter(|id| !desired.contains_key(id)) - .collect(); - for worker_id in stale { - match self.service.delete_worker(worker_id).await { - // A worker that was never registered is not an error (idempotent). - Ok(_) | Err(SelectionError::NotFound(_)) => {} - Err(e) => return Err(anyhow!("delete_worker failed: {e}")), - } - state.tracked_worker_ids.remove(&worker_id); - state.converged.remove(&worker_id); - } - - Ok(()) + /// The selection core whose catalog the topology adapter feeds. + pub fn core(&self) -> &Arc { + self.service.core() } /// Select a worker for a prompt and book its load in one operation. Takes the @@ -360,9 +248,12 @@ impl Drop for Selector { #[cfg(test)] mod tests { use std::sync::Mutex; + use std::collections::HashMap; use std::sync::atomic::{AtomicUsize, Ordering}; - use dynamo_kv_router::services::selection::WorkerSelectionPolicyParameters; + use dynamo_kv_router::services::selection::{ + CatalogReconciler, WorkerRequest, WorkerSelectionPolicyParameters, + }; use dynamo_kv_router::{ WorkerInputView, WorkerPicker, WorkerSelectionContext, WorkerSelectionPolicy, WorkerSelectionPolicyError, WorkerSelectionPolicyFactory, @@ -465,17 +356,10 @@ models: /// A registration the core marks `Incomplete`: `block_size = 0` fails the /// schedulable-metadata check independent of router/kv-event config, so the /// upsert returns `Ok` with a non-`Schedulable` lifecycle. - fn incomplete_registration(worker_id: u64) -> WorkerRegistration { - WorkerRegistration { - worker_id, - model_name: "test-model".to_string(), - endpoint: "http://10.0.0.1:8000".to_string(), - block_size: 0, - data_parallel_size: 1, - kv_events_endpoints: HashMap::new(), - replay_endpoint: None, - total_kv_blocks: None, - max_num_batched_tokens: None, + fn incomplete_registration(worker_id: u64) -> WorkerRequest { + WorkerRequest { + block_size: Some(0), + ..schedulable_registration(worker_id) } } @@ -487,23 +371,28 @@ models: /// `WorkerRegistry` register tests in `lib/kv-router`), so no network infra is /// required. Queueing is disabled under the default router policy, so /// `max_num_batched_tokens` is not required here. - fn schedulable_registration(worker_id: u64) -> WorkerRegistration { - WorkerRegistration { + fn schedulable_registration(worker_id: u64) -> WorkerRequest { + WorkerRequest { worker_id, model_name: "test-model".to_string(), - endpoint: format!("http://10.0.0.{worker_id}:8000"), - block_size: 16, - data_parallel_size: 1, + endpoint: Some(format!("http://10.0.0.{worker_id}:8000")), + block_size: Some(16), + data_parallel_start_rank: Some(0), + data_parallel_size: Some(1), kv_events_endpoints: HashMap::from([( 0u32, format!("tcp://127.0.0.1:{}", 45_000 + worker_id), )]), - replay_endpoint: None, - total_kv_blocks: None, - max_num_batched_tokens: None, + ..Default::default() } } + async fn register(selector: &Selector, workers: Vec) { + CatalogReconciler::new(Arc::clone(selector.core())) + .apply(workers) + .await + .expect("reconcile should succeed"); + } /// A selection request keyed by `reservation_id` for the schedulable worker's /// model. The prompt is long enough to book non-trivial prefill load. fn select_request(reservation_id: &str) -> SelectRequest { @@ -525,10 +414,7 @@ models: let selector = Selector::new(&test_config(), WorkerSelectionPolicyRegistry::default()) .await .expect("selector should build"); - selector - .reconcile(&[schedulable_registration(1)]) - .await - .expect("reconcile should succeed"); + register(&selector, vec![schedulable_registration(1)]).await; assert!( selector.any_ready().await, "a complete worker must be schedulable in-process" @@ -601,10 +487,7 @@ worker_selection: .expect("custom selection service should build"); assert_eq!(provider_calls.load(Ordering::Relaxed), 1); assert_eq!(factory_calls.load(Ordering::Relaxed), 0); - selector - .reconcile(&[schedulable_registration(1)]) - .await - .expect("worker should register"); + register(&selector, vec![schedulable_registration(1)]).await; assert_eq!(factory_calls.load(Ordering::Relaxed), 1); let response = selector @@ -742,10 +625,7 @@ worker_selection: let selector = Selector::new(&test_config(), WorkerSelectionPolicyRegistry::default()) .await .expect("selector should build"); - selector - .reconcile(&[incomplete_registration(1)]) - .await - .expect("reconcile should succeed"); + register(&selector, vec![incomplete_registration(1)]).await; assert!( !selector.any_ready().await, "an incomplete worker must not be schedulable" @@ -778,69 +658,6 @@ worker_selection: .expect("prefill-complete of an unknown id is a no-op"); } - #[tokio::test] - async fn incomplete_worker_is_not_cached_as_reconciled() { - let selector = Selector::new(&test_config(), WorkerSelectionPolicyRegistry::default()) - .await - .expect("selector should build"); - - let desired = [incomplete_registration(1)]; - selector - .reconcile(&desired) - .await - .expect("reconcile should succeed"); - - // The worker came back Incomplete, so it must NOT be recorded as - // reconciled — otherwise the identical next snapshot would skip the - // re-upsert and the worker would stay silently unconverged. - assert!( - selector.reconcile_state.lock().await.converged.is_empty(), - "Incomplete worker must not be cached as reconciled" - ); - assert!( - !selector.any_ready().await, - "an Incomplete worker must not be schedulable" - ); - - // A second identical reconcile must re-attempt the upsert (cache miss), - // not skip it, so the worker keeps getting a chance to converge. - selector - .reconcile(&desired) - .await - .expect("second reconcile should succeed"); - assert!( - selector.reconcile_state.lock().await.converged.is_empty(), - "Incomplete worker must still be uncached after a repeat reconcile" - ); - } - - #[tokio::test] - async fn stale_incomplete_worker_is_deleted() { - let selector = Selector::new(&test_config(), WorkerSelectionPolicyRegistry::default()) - .await - .expect("selector should build"); - - selector - .reconcile(&[incomplete_registration(1)]) - .await - .expect("incomplete worker should be tracked"); - selector - .reconcile(&[]) - .await - .expect("stale incomplete worker should be deleted"); - - let worker = selector - .service - .list_workers(None, None) - .into_iter() - .find(|worker| worker.worker_id == 1) - .expect("SelectionService retains the terminal catalog record"); - assert_eq!(worker.lifecycle, WorkerLifecycle::Unschedulable); - let state = selector.reconcile_state.lock().await; - assert!(state.tracked_worker_ids.is_empty()); - assert!(state.converged.is_empty()); - } - #[tokio::test] async fn selector_without_peer_replication_disables_replica_sync() { let selector = Selector::new(&test_config(), WorkerSelectionPolicyRegistry::default()) diff --git a/deploy/inference-gateway/ext-proc/src/topology_adapter.rs b/deploy/inference-gateway/ext-proc/src/topology_adapter.rs index c95bef0f80d9..7bbd69ad71ec 100644 --- a/deploy/inference-gateway/ext-proc/src/topology_adapter.rs +++ b/deploy/inference-gateway/ext-proc/src/topology_adapter.rs @@ -1,19 +1,25 @@ // SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! Reconciles discovered vLLM pods with the selection-service worker catalog. +//! Feeds discovered vLLM pods into the selection-service worker catalog. //! -//! The adapter converts each ready pod into a [`WorkerRegistration`] using its -//! resolved endpoints and configured defaults, then passes the desired worker -//! set to the [`Selector`]. +//! The pod reflector is exposed as a [`WorkerCatalogSource`]: each ready pod +//! becomes a [`WorkerRequest`] from its resolved endpoints and the configured +//! defaults, and a [`CatalogReconciler`] keeps the catalog in step. use std::sync::Arc; +use async_trait::async_trait; +use dynamo_kv_router::DEFAULT_ROUTING_GROUP; +use dynamo_kv_router::services::selection::{ + CatalogReconciler, WorkerCatalogSource, WorkerRequest, +}; +use tokio::sync::watch; use tokio_util::sync::CancellationToken; use crate::epp_standalone_config::EppStandaloneConfig; use crate::pod_discovery::{PodDiscovery, RawWorker}; -use crate::selector::{Selector, WorkerRegistration}; +use crate::selector::Selector; #[derive(Debug, Clone)] pub struct RegistrationDefaults { @@ -34,6 +40,39 @@ impl RegistrationDefaults { } } +/// The pod reflector as worker membership: every `Ready`, pool-selected pod. +/// When the reflector stops, the source yields one empty snapshot so selection +/// fails closed rather than routing to pods nobody is watching. +struct PodReflectorSource { + reflector: PodDiscovery, + changes: watch::Receiver, + defaults: RegistrationDefaults, + primed: bool, + closed: bool, +} + +#[async_trait] +impl WorkerCatalogSource for PodReflectorSource { + async fn next_snapshot(&mut self) -> Option> { + if self.closed { + return None; + } + if self.primed && self.changes.changed().await.is_err() { + tracing::warn!("reflector change channel closed; clearing selector topology"); + self.closed = true; + return Some(Vec::new()); + } + self.primed = true; + Some( + self.reflector + .ready_workers() + .into_iter() + .map(|worker| worker_request(worker, &self.defaults)) + .collect(), + ) + } +} + /// Background task that keeps the selector catalog in sync with the reflector. /// Dropping the adapter cancels the task so it stops promptly and releases its /// `Selector`/`PodDiscovery` handles. @@ -48,32 +87,16 @@ impl TopologyAdapter { defaults: RegistrationDefaults, ) -> Self { let cancel = CancellationToken::new(); - let cancel_child = cancel.clone(); - tokio::spawn(async move { - let mut pod_changes = reflector.subscribe_changes(); - loop { - reconcile_once(&reflector, selector.as_ref(), &defaults).await; - tokio::select! { - _ = cancel_child.cancelled() => break, - // Re-reconcile on a pod change. Exit if the sender drops - // (reflector gone). - changed = pod_changes.changed() => { - if changed.is_err() { - tracing::warn!( - "Reflector change channel closed; clearing selector topology" - ); - if let Err(e) = selector.reconcile(&[]).await { - tracing::warn!( - error = %e, - "Failed to clear selector topology after reflector stopped" - ); - } - break; - } - } - } - } - }); + let source = PodReflectorSource { + changes: reflector.subscribe_changes(), + reflector, + defaults, + primed: false, + closed: false, + }; + tokio::spawn( + CatalogReconciler::new(Arc::clone(selector.core())).run(source, cancel.child_token()), + ); Self { cancel } } } @@ -84,36 +107,20 @@ impl Drop for TopologyAdapter { } } -/// Run one reconcile pass: build the desired catalog from the Ready pods and -/// hand it to the selector, which owns the actual-vs-desired diff. -async fn reconcile_once( - reflector: &PodDiscovery, - selector: &Selector, - defaults: &RegistrationDefaults, -) { - let desired: Vec = reflector - .ready_workers() - .into_iter() - .map(|w| build_registration(w, defaults)) - .collect(); - - if let Err(e) = selector.reconcile(&desired).await { - tracing::warn!(error = %e, "Selector reconcile failed; will retry on next change"); - } -} - -fn build_registration(w: RawWorker, defaults: &RegistrationDefaults) -> WorkerRegistration { - let data_parallel_size = (w.kv_events_endpoints.len() as u32).max(1); - WorkerRegistration { +fn worker_request(w: RawWorker, defaults: &RegistrationDefaults) -> WorkerRequest { + WorkerRequest { worker_id: w.worker_id, model_name: defaults.model_name.clone(), - endpoint: w.http_endpoint, - block_size: defaults.block_size, - data_parallel_size, + routing_group: DEFAULT_ROUTING_GROUP.to_string(), + endpoint: Some(w.http_endpoint), + block_size: Some(defaults.block_size), + data_parallel_start_rank: Some(0), + data_parallel_size: Some((w.kv_events_endpoints.len() as u32).max(1)), kv_events_endpoints: w.kv_events_endpoints, replay_endpoint: w.replay_endpoint, total_kv_blocks: defaults.total_kv_blocks, max_num_batched_tokens: defaults.max_num_batched_tokens, + ..Default::default() } } @@ -169,16 +176,20 @@ mod tests { #[test] fn registration_maps_env_and_endpoints() { - let reg = build_registration(worker(7, "10.0.0.1"), &defaults()); - assert_eq!(reg.worker_id, 7); - assert_eq!(reg.model_name, "Qwen/Qwen3-0.6B"); - assert_eq!(reg.endpoint, "http://10.0.0.1:8000"); - assert_eq!(reg.block_size, 16); + let mut raw = worker(7, "10.0.0.1"); + raw.kv_events_endpoints + .insert(1, "tcp://10.0.0.1:5558".to_string()); + let request = worker_request(raw, &defaults()); + assert_eq!(request.worker_id, 7); + assert_eq!(request.model_name, "Qwen/Qwen3-0.6B"); + assert_eq!(request.endpoint.as_deref(), Some("http://10.0.0.1:8000")); + assert_eq!(request.block_size, Some(16)); + assert_eq!(request.data_parallel_size, Some(2)); assert_eq!( - reg.kv_events_endpoints.get(&0).unwrap(), + request.kv_events_endpoints.get(&0).unwrap(), "tcp://10.0.0.1:5557" ); - assert_eq!(reg.total_kv_blocks, Some(1000)); + assert_eq!(request.total_kv_blocks, Some(1000)); } #[tokio::test] diff --git a/lib/kv-router/src/services/selection/README.md b/lib/kv-router/src/services/selection/README.md index e89c98e402c4..3c81c0560552 100644 --- a/lib/kv-router/src/services/selection/README.md +++ b/lib/kv-router/src/services/selection/README.md @@ -12,6 +12,10 @@ and active-sequence accounting. Keep these implementation invariants explicit: - Mooncake overlap fields are raw matched-token observability. Effective prefill tokens use the scheduler's weighted cache credit and are not derived from `longest_matched`. +- Worker membership arrives as complete snapshots through a + `WorkerCatalogSource`; `CatalogReconciler` owns the desired-versus-actual + diff (retry of not-yet-schedulable records, stale deletion). Hosts do not + call `upsert_worker`/`delete_worker` for discovery-driven membership. - Selector replicas synchronize admission, prefill-complete, and free events. - **NOTE:** Output-block updates remain local. They are deliberately excluded from replica sync because their frequency would consume disproportionate From 945b4aea2f4f41914dd78fbb729911b727115d1b Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Wed, 9 Sep 2026 11:12:33 -0700 Subject: [PATCH 023/102] fix(router): update membership tests to reconciler API Signed-off-by: Ishan Dhanani --- .../ext-proc/src/selector.rs | 21 +++++------------ .../src/services/selection/membership.rs | 23 +++++++++++-------- 2 files changed, 19 insertions(+), 25 deletions(-) diff --git a/deploy/inference-gateway/ext-proc/src/selector.rs b/deploy/inference-gateway/ext-proc/src/selector.rs index 74d165a4d3ab..688c71b3cf62 100644 --- a/deploy/inference-gateway/ext-proc/src/selector.rs +++ b/deploy/inference-gateway/ext-proc/src/selector.rs @@ -247,8 +247,8 @@ impl Drop for Selector { #[cfg(test)] mod tests { - use std::sync::Mutex; use std::collections::HashMap; + use std::sync::Mutex; use std::sync::atomic::{AtomicUsize, Ordering}; use dynamo_kv_router::services::selection::{ @@ -718,8 +718,8 @@ worker_selection: .expect("selector should build"); let duplicate = incomplete_registration(1); - let error = selector - .reconcile(&[duplicate.clone(), duplicate]) + let error = CatalogReconciler::new(Arc::clone(selector.core())) + .apply(vec![duplicate.clone(), duplicate]) .await .expect_err("duplicate IDs must be rejected"); assert!(error.to_string().contains("duplicate worker_id 1")); @@ -779,10 +779,7 @@ worker_selection: ) .await .expect("selector should build"); - selector - .reconcile(&[schedulable_registration(1)]) - .await - .expect("worker should register"); + register(&selector, vec![schedulable_registration(1)]).await; let mut req = select_request("res-eot"); req.expected_output_tokens = Some(128); @@ -819,10 +816,7 @@ worker_selection: .expect("selector should build"); let mut reg = schedulable_registration(1); reg.model_name = "threshold-free-model".to_string(); - selector - .reconcile(&[reg]) - .await - .expect("worker should register"); + register(&selector, vec![reg]).await; // A policy-family name resolves to that family's bucketed class. let mut family_req = select_request("res-family"); @@ -869,10 +863,7 @@ worker_selection: let selector = Selector::new(&test_config(), WorkerSelectionPolicyRegistry::default()) .await .expect("selector should build"); - selector - .reconcile(&[schedulable_registration(1)]) - .await - .expect("worker should register"); + register(&selector, vec![schedulable_registration(1)]).await; let mut synthetic_req = select_request("res-synthetic"); synthetic_req.policy_class = Some("anything".to_string()); selector diff --git a/lib/kv-router/src/services/selection/membership.rs b/lib/kv-router/src/services/selection/membership.rs index b485418a2e16..c5eed8d2a128 100644 --- a/lib/kv-router/src/services/selection/membership.rs +++ b/lib/kv-router/src/services/selection/membership.rs @@ -149,16 +149,19 @@ mod tests { use crate::services::selection::SelectionCacheConfig; fn core() -> Arc { - Arc::new(SelectionCore::new_local( - KvRouterConfig { - use_kv_events: false, - router_queue_threshold: None, - ..Default::default() - }, - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - )) + Arc::new( + SelectionCore::try_new_local( + KvRouterConfig { + use_kv_events: false, + router_queue_threshold: None, + ..Default::default() + }, + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ) + .expect("valid test config"), + ) } fn schedulable(worker_id: WorkerId) -> WorkerRequest { From 1990241e32c555987009770021099c937f80e33b Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Thu, 3 Sep 2026 22:13:17 +0000 Subject: [PATCH 024/102] refactor(kv-router): HostCache.index names where a partition's KV index comes from Replace three knobs that together described the same thing (the builder's external_kv_events() and remote_indexer(), and SelectionHost's OverlapRefreshSource) with one KvIndexSource on HostCache: - Owned: the core owns the index and subscribes to worker KV events over ZMQ (the standalone default; workers need KV-event endpoints). - External { refresh }: the host owns and feeds the index; workers need no KV-event metadata; refresh optionally lets the core re-score queued requests against the host's index (the frontend's embedded partition). - Remote(url): a standalone indexer serves the primary index. The listener decision, schedulability metadata, and dequeue-time refresh all derive from this one value. SelectionServiceBuilder::kv_index sets it; SelectionServiceConfig and the Python bindings map remote_indexer_url onto Remote. No behavior change. Signed-off-by: Ishan Dhanani --- lib/bindings/python/rust/llm/kv.rs | 10 +-- .../src/services/selection/core/mod.rs | 83 ++++++++++--------- lib/kv-router/src/services/selection/mod.rs | 2 +- .../src/services/selection/service.rs | 29 ++----- lib/llm/src/kv_router/embedded.rs | 8 +- 5 files changed, 60 insertions(+), 72 deletions(-) diff --git a/lib/bindings/python/rust/llm/kv.rs b/lib/bindings/python/rust/llm/kv.rs index de7d09c6c251..7e81d03790e9 100644 --- a/lib/bindings/python/rust/llm/kv.rs +++ b/lib/bindings/python/rust/llm/kv.rs @@ -31,10 +31,10 @@ use dynamo_kv_router::scheduling::AdmissionAttempt; use dynamo_kv_router::services::indexer::{self, IndexerConfig}; #[cfg(feature = "select-service")] use dynamo_kv_router::services::selection::{ - self, OverlapScoresRequest, PotentialLoadsRequest, ReservationRequest, SelectAndReserveRequest, - SelectRequest, SelectionCacheConfig as RsSelectionCacheConfig, SelectionError, - SelectionService as RustSelectionService, SelectionServiceBuilder, SelectionServiceConfig, - WorkerPatchRequest, WorkerRequest, WorkerSelectionPolicyRegistry, + self, KvIndexSource, OverlapScoresRequest, PotentialLoadsRequest, ReservationRequest, + SelectAndReserveRequest, SelectRequest, SelectionCacheConfig as RsSelectionCacheConfig, + SelectionError, SelectionService as RustSelectionService, SelectionServiceBuilder, + SelectionServiceConfig, WorkerPatchRequest, WorkerRequest, WorkerSelectionPolicyRegistry, warn_for_unserved_worker_selection_policies, }; #[cfg(feature = "slot-tracker")] @@ -673,7 +673,7 @@ impl SelectionService { builder = builder.replica_sync(port, replica_sync_peers); } if let Some(url) = remote_indexer_url { - builder = builder.remote_indexer(url); + builder = builder.kv_index(KvIndexSource::Remote(url)); } let inner = py .allow_threads(|| pyo3_async_runtimes::tokio::get_runtime().block_on(builder.build())) diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index ebc8b3f09d5d..3df89340fa6a 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -206,16 +206,44 @@ pub struct HostLoad { pub available_workers: Option, } -/// KV-cache knowledge outside the partition's own index. +/// Where a partition's KV knowledge comes from. #[derive(Clone, Default)] pub struct HostCache { /// Queried alongside the indexer for prompts that carry `token_ids`; a /// failed lookup is logged and selection proceeds without shared hits. pub shared: Option>, - /// Where dequeue-time overlap refresh reads from. A host that keeps its - /// own index points this at that index so queued requests are re-scored - /// against the index that scored them. - pub overlap_refresh: OverlapRefreshSource, + pub index: KvIndexSource, +} + +/// Where a partition's KV index lives and who feeds it. +#[derive(Clone, Default)] +pub enum KvIndexSource { + /// The core owns the index and subscribes to each worker rank's KV events + /// over ZMQ; workers need KV-event endpoints to be schedulable. + #[default] + Owned, + /// The host owns and feeds the index. Workers need no KV-event metadata. + /// `refresh` lets the core re-score queued requests against that index at + /// dequeue; `None` keeps their admission-time scores. + External { + refresh: Option>, + }, + /// A standalone indexer at this base URL serves the primary index; this + /// core does not subscribe to worker KV events. + Remote(String), +} + +impl std::fmt::Debug for KvIndexSource { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Owned => formatter.write_str("Owned"), + Self::External { refresh } => formatter + .debug_struct("External") + .field("refresh", &refresh.is_some()) + .finish(), + Self::Remote(url) => formatter.debug_tuple("Remote").field(url).finish(), + } + } } /// Host-owned narrowing of the candidate set. @@ -241,28 +269,6 @@ pub struct HostReplication { pub channels: Option, } -/// Source of dequeue-time overlap refreshes for partition schedulers. -#[derive(Clone, Default)] -pub enum OverlapRefreshSource { - /// The partition's own indexer (standalone service default). - #[default] - PartitionIndexer, - /// A host-owned index. - External(Arc), - /// No refresh: queued requests keep their admission-time scores. - Disabled, -} - -impl std::fmt::Debug for OverlapRefreshSource { - fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - formatter.write_str(match self { - Self::PartitionIndexer => "PartitionIndexer", - Self::External(_) => "External", - Self::Disabled => "Disabled", - }) - } -} - impl std::fmt::Debug for SelectionHost { fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { formatter @@ -274,7 +280,7 @@ impl std::fmt::Debug for SelectionHost { ) .field("available_workers", &self.load.available_workers.is_some()) .field("shared_cache", &self.cache.shared.is_some()) - .field("overlap_refresh", &self.cache.overlap_refresh) + .field("index", &self.cache.index) .field( "lora_worker_filter", &self.eligibility.lora_worker_filter.is_some(), @@ -376,7 +382,6 @@ impl SelectionCore { cache_config, tracking_hash, indexer_policy, - false, )) } @@ -399,15 +404,15 @@ impl SelectionCore { cache_config: SelectionCacheConfig, tracking_hash: Arc, indexer_policy: IndexerPolicy, - external_kv_events: bool, ) -> Self { let cancel_token = cancel_token.child_token(); let indexer_registry = Arc::new(WorkerRegistry::new_with_cancel_token( indexer_threads, cancel_token.clone(), )); - let listens_for_kv_events = - kv_router_config.use_kv_events && !indexer_policy.is_remote() && !external_kv_events; + let listens_for_kv_events = kv_router_config.use_kv_events + && !indexer_policy.is_remote() + && !matches!(host.cache.index, KvIndexSource::External { .. }); indexer_registry.set_indexer_policy(indexer_policy); if signal_indexer_ready { indexer_registry.signal_ready(); @@ -728,14 +733,13 @@ impl SelectionCore { let indexer = self .indexer_registry .get_or_create_indexer(key.clone(), block_size); - let refresh_provider = match &self.host.cache.overlap_refresh { - OverlapRefreshSource::PartitionIndexer => { + let refresh_provider = match &self.host.cache.index { + KvIndexSource::Owned | KvIndexSource::Remote(_) => { Some(RefreshProvider::Local(Arc::new(indexer.clone()))) } - OverlapRefreshSource::External(provider) => { - Some(RefreshProvider::External(Arc::clone(provider))) - } - OverlapRefreshSource::Disabled => None, + KvIndexSource::External { refresh } => refresh + .as_ref() + .map(|p| RefreshProvider::External(Arc::clone(p))), }; let overlap_refresh = refresh_provider.map(|provider| { Arc::new(TieredOverlapRefresher::new( @@ -2012,7 +2016,6 @@ mod tests { SelectionCacheConfig::default(), tracking_hash, IndexerPolicy::from_router_config(&test_config(false)).expect("indexer policy"), - false, ); core.upsert_worker(worker(1)).await.expect("worker upsert"); @@ -2052,7 +2055,6 @@ mod tests { SelectionCacheConfig::default(), tracking_hash, indexer_policy, - false, ) } @@ -2213,7 +2215,6 @@ mod tests { SelectionCacheConfig::default(), tracking_hash, indexer_policy, - false, ); assert!(!core.listens_for_kv_events); for worker_id in [1, 2] { diff --git a/lib/kv-router/src/services/selection/mod.rs b/lib/kv-router/src/services/selection/mod.rs index 41b60e53ea8c..cac678fb3bec 100644 --- a/lib/kv-router/src/services/selection/mod.rs +++ b/lib/kv-router/src/services/selection/mod.rs @@ -26,7 +26,7 @@ pub use crate::services::common::replica_sync::{ ScopedSequencePublisher, }; pub use core::{ - HostCache, HostEligibility, HostLoad, HostReplication, HostTelemetry, OverlapRefreshSource, + HostCache, HostEligibility, HostLoad, HostReplication, HostTelemetry, KvIndexSource, RefreshProvider, SelectionCore, SelectionHost, SelectionPartition, SelectionScheduler, SelectionServiceConfig, }; diff --git a/lib/kv-router/src/services/selection/service.rs b/lib/kv-router/src/services/selection/service.rs index 58046c613d00..49f134bec223 100644 --- a/lib/kv-router/src/services/selection/service.rs +++ b/lib/kv-router/src/services/selection/service.rs @@ -16,7 +16,9 @@ use crate::services::common::replica_sync::{ use crate::services::indexer::backend::IndexerPolicy; use crate::tracking_hash::TrackingHashContext; -use super::core::{SelectionCore, SelectionHost, SelectionPartition, SelectionServiceConfig}; +use super::core::{ + KvIndexSource, SelectionCore, SelectionHost, SelectionPartition, SelectionServiceConfig, +}; use super::error::SelectionError; use super::pending::SelectionCacheConfig; use super::policy_registry::WorkerSelectionPolicyRegistry; @@ -37,9 +39,7 @@ pub struct SelectionServiceBuilder { worker_type: WorkerType, worker_selection_policy_registry: WorkerSelectionPolicyRegistry, host: SelectionHost, - remote_indexer_url: Option, worker_selection_policy_factory: Option, - external_kv_events: bool, } /// Warn when a host does not construct workers for explicitly configured policy roles. @@ -78,9 +78,7 @@ impl SelectionServiceBuilder { worker_type, worker_selection_policy_registry, host: SelectionHost::default(), - remote_indexer_url: None, worker_selection_policy_factory: None, - external_kv_events: false, } } @@ -94,18 +92,10 @@ impl SelectionServiceBuilder { self } - /// The embedding host feeds KV state itself: workers become schedulable - /// without `kv_events_endpoint(s)` and no ZMQ listener is started here. - pub fn external_kv_events(mut self) -> Self { - self.external_kv_events = true; - self - } - - /// Serve the primary KV index from a standalone indexer at `base_url` - /// instead of subscribing to worker KV events in this process. Requires - /// `use_kv_events=true`; `indexer_peers` recovery is skipped. - pub fn remote_indexer(mut self, base_url: impl Into) -> Self { - self.remote_indexer_url = Some(base_url.into()); + /// Where each partition's KV index comes from (see [`KvIndexSource`]). + /// Shorthand for setting `host.cache.index`. + pub fn kv_index(mut self, source: KvIndexSource) -> Self { + self.host.cache.index = source; self } @@ -148,7 +138,7 @@ impl SelectionServiceBuilder { }; let tracking_hash = Arc::new(TrackingHashContext::from_config(&self.kv_router_config)?); let mut indexer_policy = IndexerPolicy::from_router_config(&self.kv_router_config)?; - if let Some(base_url) = &self.remote_indexer_url { + if let KvIndexSource::Remote(base_url) = &self.host.cache.index { indexer_policy = indexer_policy.with_remote_indexer(base_url.clone())?; if !self.indexer_peers.is_empty() { tracing::warn!( @@ -186,7 +176,6 @@ impl SelectionServiceBuilder { self.selection_cache, tracking_hash, indexer_policy, - self.external_kv_events, )); if recover_from_peers { @@ -248,7 +237,7 @@ impl SelectionServiceConfig { builder = builder.replica_sync(port, self.replica_sync_peers.clone()); } if let Some(url) = &self.remote_indexer_url { - builder = builder.remote_indexer(url.clone()); + builder = builder.kv_index(KvIndexSource::Remote(url.clone())); } builder } diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs index 1811875e1383..183f0b52049c 100644 --- a/lib/llm/src/kv_router/embedded.rs +++ b/lib/llm/src/kv_router/embedded.rs @@ -29,7 +29,7 @@ use dynamo_kv_router::scheduling::{ use dynamo_kv_router::sequences::{SequenceError, SequenceRequest}; use dynamo_kv_router::services::selection::{ CatalogObserver, CatalogReconciler, HostCache, HostEligibility, HostLoad, HostReplication, - HostTelemetry, OverlapRefreshSource, SelectionHost, SelectionPartition, SelectionService, + HostTelemetry, KvIndexSource, SelectionHost, SelectionPartition, SelectionService, SelectionServiceBuilder, WorkerCatalogRecord, WorkerCatalogSource, WorkerRequest, WorkerSelectionPolicyRegistry, }; @@ -148,7 +148,6 @@ impl EmbeddedSelection { ) .worker_selection_policy_factory(args.policy_factory) .indexer_threads(1) - .external_kv_events() .host(SelectionHost { load: HostLoad { prefill_estimator: args.prefill_load_estimator, @@ -157,9 +156,8 @@ impl EmbeddedSelection { }, cache: HostCache { shared: None, - overlap_refresh: match args.overlap_refresh { - Some(provider) => OverlapRefreshSource::External(provider), - None => OverlapRefreshSource::Disabled, + index: KvIndexSource::External { + refresh: args.overlap_refresh, }, }, eligibility: HostEligibility::default(), From 8772ed0565db0e19340de9d4bb7053f7cbafa2c9 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Thu, 3 Sep 2026 23:32:48 +0000 Subject: [PATCH 025/102] refactor(kv-router): one Indexer type for the frontend and the selection service lib/llm and lib/kv-router each carried an Indexer enum with the same Single/Concurrent/Remote shape (primary, lower tiers, predict-on-route side overlay, primary_records_routing_decisions) and diverging method sets. Fold the frontend's richer implementation into dynamo_kv_router::services::indexer: - lookup.rs: the borrowed/owned HashInput lookup pipeline (primary, lower tiers seeded from confirmed primary matches, side overlay merged by per-worker max), find_matches_by_tier[_ref][_with_options], find_primary_*. - recording.rs: routing-decision recording targets (side overlay, approximate primary, remote primary) and the per-request ApproximateRequestLease. - backend.rs: the merged enum gains None, residency projection setters, try_apply_event (the backpressure-aware admit with a completion barrier on Cleared that recovery relies on), reset_worker_dp_rank_and_wait, and the approximate-LRU capacity and stats API. apply_event_routed and remove_worker[_dp_rank] stay as the standalone listener's path. Remote primaries go behind RemotePrimary so the HTTP standalone client here and the request-plane RemoteIndexer in lib/llm plug into the same enum. dump_events returns KvRouterError like the rest of the surface. - The frontend keeps only construction (indexer::build from KvRouterConfig with component metrics and the request-plane remote), the served-indexer service, and recovery. Its 16 indexer tests move with the code. No behavior change: both hosts run the exact code they ran before for their own paths; TieredMatchProvider keeps the capability-gated variant. Signed-off-by: Ishan Dhanani --- lib/kv-router/src/services/indexer/backend.rs | 518 ++++--- lib/kv-router/src/services/indexer/lookup.rs | 1190 +++++++++++++++ lib/kv-router/src/services/indexer/mod.rs | 2 + .../src/services}/indexer/recording.rs | 105 +- lib/llm/src/kv_router.rs | 2 +- lib/llm/src/kv_router/indexer/lookup.rs | 430 ------ lib/llm/src/kv_router/indexer/mod.rs | 1307 ++--------------- .../kv_router/indexer/recovery/state_agent.rs | 4 +- .../indexer/recovery/worker_query.rs | 4 +- lib/llm/src/kv_router/indexer/remote.rs | 54 +- lib/llm/src/kv_router/indexer/side.rs | 125 -- 11 files changed, 1771 insertions(+), 1970 deletions(-) create mode 100644 lib/kv-router/src/services/indexer/lookup.rs rename lib/{llm/src/kv_router => kv-router/src/services}/indexer/recording.rs (71%) delete mode 100644 lib/llm/src/kv_router/indexer/lookup.rs delete mode 100644 lib/llm/src/kv_router/indexer/side.rs diff --git a/lib/kv-router/src/services/indexer/backend.rs b/lib/kv-router/src/services/indexer/backend.rs index 73e27a80d3c7..92f12c4019eb 100644 --- a/lib/kv-router/src/services/indexer/backend.rs +++ b/lib/kv-router/src/services/indexer/backend.rs @@ -15,15 +15,18 @@ use crate::approx::PruneConfig; use crate::config::{ApproximateCachePolicyKind, KvRouterConfig}; use crate::identity::RoutingPartitionId; use crate::indexer::{ - ApproximateRetentionConfig, KvIndexer, KvIndexerInterface, KvIndexerMetrics, KvRouterError, - LowerTierIndexers, LowerTierQueryOptions, MatchDetails, RoutingDecisionHashes, SyncIndexer, - TieredMatchDetails, TieredMatchProvider, WireTieredMatchDetails, - query_lower_tiers_with_options, + ApproximateLruIncarnation, ApproximateLruStats, ApproximateRetentionConfig, KvIndexer, + KvIndexerInterface, KvIndexerMetrics, KvRouterError, LowerTierIndexers, LowerTierQueryOptions, + MatchDetails, RoutingDecisionHashes, SyncIndexer, TieredMatchDetails, TieredMatchProvider, + WireTieredMatchDetails, record_unsupported_residency_event, }; use crate::protocols::{ - KvCacheEventData, LocalBlockHash, OverlapScores, RouterEvent, WorkerId, WorkerWithDpRank, + DpRank, KvCacheEventData, LocalBlockHash, OverlapScores, ResidencyProjection, + ResidencyRoutingSnapshot, RouterEvent, WorkerId, WorkerWithDpRank, }; +use super::lookup::{HashInput, merge_side_or_warn}; + /// How the device-tier primary is populated. #[derive(Debug, Clone, PartialEq, Eq)] pub enum PrimaryRetention { @@ -243,11 +246,13 @@ pub enum SideIndexer { } impl SideIndexer { + /// `cancel` stops the single-threaded indexer's pruning task. pub fn new( ttl: Duration, block_size: u32, num_threads: usize, metrics: Arc, + cancel: CancellationToken, ) -> Self { let prune_config = Some(PruneConfig { ttl }); if num_threads > 1 { @@ -260,24 +265,26 @@ impl SideIndexer { ))); } Self::Single(KvIndexer::new_with_pruning( - CancellationToken::new(), + cancel, block_size, metrics, prune_config, )) } - async fn find_matches( + pub(super) async fn find_matches_input( &self, - sequence: &[LocalBlockHash], + sequence: HashInput<'_>, ) -> std::result::Result { match self { - Self::Single(indexer) => indexer.find_matches(sequence.to_vec()).await, - Self::Concurrent(indexer) => Ok(indexer.backend().find_matches(sequence, false)), + Self::Single(indexer) => indexer.find_matches(sequence.as_slice().to_vec()).await, + Self::Concurrent(indexer) => { + Ok(indexer.backend().find_matches(sequence.as_slice(), false)) + } } } - async fn record_routing_decision( + pub(super) async fn record_routing_decision( &self, worker: WorkerWithDpRank, hashes: RoutingDecisionHashes, @@ -304,6 +311,25 @@ impl SideIndexer { } } + async fn reset_worker_dp_rank_and_wait( + &self, + worker_id: WorkerId, + dp_rank: DpRank, + ) -> std::result::Result<(), KvRouterError> { + match self { + Self::Single(indexer) => { + indexer + .reset_worker_dp_rank_and_wait(worker_id, dp_rank) + .await + } + Self::Concurrent(indexer) => { + indexer + .reset_worker_dp_rank_and_wait(worker_id, dp_rank) + .await + } + } + } + async fn remove_worker(&self, worker_id: WorkerId) { match self { Self::Single(indexer) => indexer.remove_worker(worker_id).await, @@ -319,53 +345,54 @@ impl SideIndexer { } } -/// Merge side-indexer scores into the primary's device match details by -/// per-worker max. Side-only workers gain a score with no paired -/// `last_matched_hashes` entry, so the result must not seed lower-tier -/// continuations; callers run `query_lower_tiers` on the primary-only details -/// first. -fn merge_side_scores(mut primary: MatchDetails, side: OverlapScores) -> MatchDetails { - for (worker, side_score) in side.scores { - primary - .overlap_scores - .scores - .entry(worker) - .and_modify(|score| { - if side_score > *score { - *score = side_score; - } - }) - .or_insert(side_score); - } - primary +/// A primary index served by another process. The HTTP standalone indexer +/// implements it here; a request-plane one implements it in its own crate. +#[async_trait] +pub trait RemotePrimary: Send + Sync { + async fn find_matches_by_tier( + &self, + block_hashes: Vec, + device_only: bool, + ) -> Result; + + /// Write a booked routing decision into an approximate remote primary. + async fn record_routing_decision( + &self, + worker: WorkerWithDpRank, + hashes: RoutingDecisionHashes, + ) -> Result<()>; + + /// Whether the remote primary is fed by engine KV events (a side indexer + /// may overlay it) rather than by routing decisions. + fn use_kv_events(&self) -> bool; } -async fn merge_side_or_warn( - side: Option<&SideIndexer>, - primary: MatchDetails, - sequence: &[LocalBlockHash], -) -> MatchDetails { - let Some(side) = side else { - return primary; - }; - match side.find_matches(sequence).await { - Ok(side_scores) => merge_side_scores(primary, side_scores), - Err(error) => { - tracing::warn!( - %error, - "predict-on-route side indexer query failed; using primary only" - ); - primary - } +#[async_trait] +impl RemotePrimary for RemoteIndexerClient { + async fn find_matches_by_tier( + &self, + block_hashes: Vec, + _device_only: bool, + ) -> Result { + Ok(self.find_tiered_matches(&block_hashes).await?) + } + + async fn record_routing_decision( + &self, + _worker: WorkerWithDpRank, + _hashes: RoutingDecisionHashes, + ) -> Result<()> { + anyhow::bail!("the standalone indexer does not record routing decisions") + } + + fn use_kv_events(&self) -> bool { + true } } -/// Block-content indexer wrapping a primary device-tier backend plus a -/// per-tier registry of lower-tier indexers (host-pinned, disk, …). -/// -/// The service owns a `lower_tier: LowerTierIndexers` so tier-tagged events -/// get routed alongside device events and tier-aware queries can return -/// per-tier hit counts. +/// Block-content indexer for one routing partition: a primary device-tier +/// backend plus a per-tier registry of lower-tier indexers (host-pinned, +/// disk, …), or a remote primary, or nothing when routing is load-only. /// /// `approx` is the optional predict-on-route side indexer, queried alongside /// the primary and merged by per-worker max. `primary_records_routing_decisions` @@ -385,13 +412,16 @@ pub enum Indexer { approx: Option, primary_records_routing_decisions: bool, }, - /// Primary served by a standalone indexer over HTTP. Lower tiers are the - /// remote service's concern and arrive in its tiered response; only the - /// side indexer is local. + /// Primary served by another process. Lower tiers are that process's + /// concern and arrive in its tiered response; only the side indexer is + /// local. Remote { - primary: Arc, + primary: Arc, approx: Option, + primary_records_routing_decisions: bool, }, + /// No KV awareness. + None, } impl Indexer { @@ -400,6 +430,7 @@ impl Indexer { Self::Single { approx, .. } | Self::Concurrent { approx, .. } | Self::Remote { approx, .. } => approx.as_ref(), + Self::None => None, } } @@ -407,67 +438,44 @@ impl Indexer { matches!(self, Self::Remote { .. }) } - /// Whether booked routing decisions should be written into this indexer. - pub fn records_routing_decisions(&self) -> bool { - match self { - Self::Single { - approx, - primary_records_routing_decisions, - .. - } - | Self::Concurrent { - approx, - primary_records_routing_decisions, - .. - } => approx.is_some() || *primary_records_routing_decisions, - Self::Remote { approx, .. } => approx.is_some(), - } + /// Dequeue-time overlap refresh needs a local primary. + pub fn supports_overlap_refresh(&self) -> bool { + matches!(self, Self::Single { .. } | Self::Concurrent { .. }) } /// Router-hint chain retention needs a local primary whose hashes come /// only from engine events. pub fn supports_kv_transfer_chain_retention(&self) -> bool { - !self.is_remote() && !self.records_routing_decisions() + matches!( + self, + Self::Single { + approx: None, + primary_records_routing_decisions: false, + .. + } | Self::Concurrent { + approx: None, + primary_records_routing_decisions: false, + .. + } + ) } - /// Record a booked routing decision. Writes to the side indexer when one - /// is attached, else to an approximate primary, else is a no-op. - pub async fn record_routing_decision( - &self, - worker: WorkerWithDpRank, - hashes: RoutingDecisionHashes, - ) -> std::result::Result<(), KvRouterError> { - if let Some(side) = self.approx() { - return side.record_routing_decision(worker, hashes).await; - } + /// Publish a control-plane residency projection for subsequent lookups. + pub fn set_residency_projection(&self, projection: ResidencyProjection) { match self { - Self::Single { - primary, - primary_records_routing_decisions: true, - .. - } => { - primary - .process_routing_decision_with_hashes( - worker, - hashes.local_hashes, - hashes.sequence_hashes, - ) - .await + Self::Single { lower_tier, .. } | Self::Concurrent { lower_tier, .. } => { + lower_tier.set_residency_projection(projection) } - Self::Concurrent { - primary, - primary_records_routing_decisions: true, - .. - } => { - primary - .process_routing_decision_hash_slices( - worker, - &hashes.local_hashes, - &hashes.sequence_hashes, - ) - .await + Self::Remote { .. } | Self::None => {} + } + } + + pub fn set_residency_routing_snapshot(&self, snapshot: ResidencyRoutingSnapshot) { + match self { + Self::Single { lower_tier, .. } | Self::Concurrent { lower_tier, .. } => { + lower_tier.set_residency_routing_snapshot(snapshot) } - Self::Single { .. } | Self::Concurrent { .. } | Self::Remote { .. } => Ok(()), + Self::Remote { .. } | Self::None => {} } } @@ -478,8 +486,8 @@ impl Indexer { match self { Indexer::Single { primary, .. } => primary.apply_event(event).await, Indexer::Concurrent { primary, .. } => primary.apply_event(event).await, - Indexer::Remote { .. } => { - tracing::trace!("Dropping KV event: the primary indexer is remote"); + Indexer::Remote { .. } | Indexer::None => { + tracing::trace!("Dropping KV event: no local primary indexer"); } } } @@ -487,8 +495,12 @@ impl Indexer { /// Apply an event, routing to the device-tier primary when /// `event.storage_tier.is_gpu()` and to the appropriate lower-tier /// indexer otherwise. `Cleared` events fan out to their applicable physical - /// indexes according to the event's reset scope. - pub async fn apply_event_routed(&self, event: RouterEvent) -> Result<(), KvRouterError> { + /// indexes according to the event's reset scope. Failures are logged and + /// the first one is returned after every tier was attempted. + pub async fn apply_event_routed( + &self, + event: RouterEvent, + ) -> std::result::Result<(), KvRouterError> { let targets_primary = match event.targets_primary() { Ok(targets_primary) => targets_primary, Err(_) => { @@ -496,7 +508,7 @@ impl Indexer { Self::Single { lower_tier, .. } | Self::Concurrent { lower_tier, .. } => { lower_tier.record_unsupported_residency_event(&event); } - Self::Remote { .. } => {} + Self::Remote { .. } | Self::None => {} } return Ok(()); } @@ -565,13 +577,149 @@ impl Indexer { .await; } } - Indexer::Remote { .. } => { - tracing::trace!("Dropping KV event: the primary indexer is remote"); + Indexer::Remote { .. } | Indexer::None => { + tracing::trace!("Dropping KV event: no local primary indexer"); } } Ok(()) } + /// Admit one event with backpressure: a `Cleared` event completes only once + /// every local tier has applied it (a completion barrier for recovery), a + /// device event waits for the primary's queue to accept it, and a lower-tier + /// event is enqueued. Unlike [`Self::apply_event_routed`], the first failure + /// stops the fan-out. + pub async fn try_apply_event( + &self, + event: RouterEvent, + ) -> std::result::Result<(), KvRouterError> { + let targets_primary = match event.targets_primary() { + Ok(targets_primary) => targets_primary, + Err(_) => { + match self { + Self::Single { lower_tier, .. } | Self::Concurrent { lower_tier, .. } => { + lower_tier.record_unsupported_residency_event(&event); + } + Self::Remote { .. } | Self::None => { + record_unsupported_residency_event(None, &event); + } + } + return Ok(()); + } + }; + let is_clear = matches!(&event.event.data, KvCacheEventData::Cleared); + match self { + Self::Single { + primary, + lower_tier, + .. + } => { + if is_clear { + if targets_primary { + primary + .reset_worker_dp_rank_and_wait(event.worker_id, event.event.dp_rank) + .await?; + } + for indexer in lower_tier.all() { + indexer.apply_event_and_wait(event.clone()).await?; + } + } else if targets_primary { + primary + .event_sender() + .send(event) + .await + .map_err(|_| KvRouterError::IndexerOffline)?; + } else { + lower_tier + .get_or_create(event.storage_tier) + .enqueue_event(event)?; + } + } + Self::Concurrent { + primary, + lower_tier, + .. + } => { + if is_clear { + if targets_primary { + primary.apply_event_and_wait(event.clone()).await?; + } + for indexer in lower_tier.all() { + indexer.apply_event_and_wait(event.clone()).await?; + } + } else if targets_primary { + primary.enqueue_event(event)?; + } else { + lower_tier + .get_or_create(event.storage_tier) + .enqueue_event(event)?; + } + } + Self::Remote { .. } | Self::None => {} + } + Ok(()) + } + + /// Cold-reset one logical rank and wait until every local tier (and the + /// side indexer) has completed the removal. + pub async fn reset_worker_dp_rank_and_wait( + &self, + worker_id: WorkerId, + dp_rank: DpRank, + ) -> std::result::Result<(), KvRouterError> { + match self { + Self::Single { + primary, + lower_tier, + approx, + .. + } => { + primary + .reset_worker_dp_rank_and_wait(worker_id, dp_rank) + .await?; + for indexer in lower_tier.all() { + indexer + .reset_worker_dp_rank_and_wait(worker_id, dp_rank) + .await?; + } + if let Some(approx) = approx { + approx + .reset_worker_dp_rank_and_wait(worker_id, dp_rank) + .await?; + } + } + Self::Concurrent { + primary, + lower_tier, + approx, + .. + } => { + primary + .reset_worker_dp_rank_and_wait(worker_id, dp_rank) + .await?; + for indexer in lower_tier.all() { + indexer + .reset_worker_dp_rank_and_wait(worker_id, dp_rank) + .await?; + } + if let Some(approx) = approx { + approx + .reset_worker_dp_rank_and_wait(worker_id, dp_rank) + .await?; + } + } + Self::Remote { approx, .. } => { + if let Some(approx) = approx { + approx + .reset_worker_dp_rank_and_wait(worker_id, dp_rank) + .await?; + } + } + Self::None => {} + } + Ok(()) + } + pub async fn remove_worker(&self, worker_id: WorkerId) { if let Some(side) = self.approx() { side.remove_worker(worker_id).await; @@ -597,7 +745,7 @@ impl Indexer { } primary.remove_worker(worker_id).await; } - Indexer::Remote { .. } => {} + Indexer::Remote { .. } | Indexer::None => {} } } @@ -626,40 +774,39 @@ impl Indexer { } primary.remove_worker_dp_rank(worker_id, dp_rank).await; } - Indexer::Remote { .. } => {} + Indexer::Remote { .. } | Indexer::None => {} } } - /// Device-tier overlap scores, including side-indexer credit. Existing - /// flat-shape callers continue to use this; tier-aware callers should use - /// [`Self::find_tiered_matches`]. + /// Device-tier overlap scores, including side-indexer credit. Tier-aware + /// callers should use [`Self::find_tiered_matches`]. pub async fn find_matches(&self, hashes: Vec) -> Result { let primary = match self { Indexer::Single { primary, .. } => primary.find_matches(hashes.clone()).await?, Indexer::Concurrent { primary, .. } => primary.find_matches(hashes.clone()).await?, Indexer::Remote { primary, .. } => { primary - .find_tiered_matches(&hashes) + .find_matches_by_tier(hashes.clone(), true) .await? .device .overlap_scores } + Indexer::None => return Ok(OverlapScores::default()), }; let Some(side) = self.approx() else { return Ok(primary); }; let mut merged = MatchDetails::new(); merged.overlap_scores = primary; - Ok(merge_side_or_warn(Some(side), merged, &hashes) - .await - .overlap_scores) + Ok( + merge_side_or_warn(Some(side), merged, HashInput::Owned(hashes)) + .await + .overlap_scores, + ) } /// Device match details + per-tier hits, suitable for building the /// Mooncake-RFC-shape per-instance breakdown. - /// - /// Lower-tier continuations are seeded from confirmed primary matches only; - /// side-indexer scores are merged into the device tier afterwards. pub async fn find_tiered_matches( &self, sequence: Vec, @@ -680,77 +827,39 @@ impl Indexer { retain_kv_transfer_chain: options.retain_kv_transfer_chain && self.supports_kv_transfer_chain_retention(), }; - let (device, lower_tier) = match self { - Indexer::Single { - primary, - lower_tier, - .. - } => { - let device = primary - .find_match_details_with_options( - sequence.clone(), - options.retain_kv_transfer_chain, - ) - .await?; - let lt = query_lower_tiers_with_options(lower_tier, &sequence, &device, options); - (device, lt) - } - Indexer::Concurrent { - primary, - lower_tier, - .. - } => { - let device: MatchDetails = primary.backend().find_match_details_impl_with_options( - &sequence, - false, - options.retain_kv_transfer_chain, - ); - let lt = query_lower_tiers_with_options(lower_tier, &sequence, &device, options); - (device, lt) - } - Indexer::Remote { primary, .. } => { - let tiered = primary.find_tiered_matches(&sequence).await?; - (tiered.device, tiered.lower_tier) - } - }; - let device = merge_side_or_warn(self.approx(), device, &sequence).await; - Ok(TieredMatchDetails { device, lower_tier }) + self.find_matches_by_tier_with_options(sequence, options) + .await } /// Dump every event tracked by this indexer in a form replayable by a - /// peer's [`Self::apply_event_routed`]. - /// - /// Returns the primary device-tier dump first, followed by every allocated - /// lower-tier indexer's dump with each event's `storage_tier` retagged to - /// the tier it lives in. The `LowerTierIndexer::dump_events` synthetic - /// events default `storage_tier` to `Device`, so without retagging a peer - /// would replay them into the wrong slot. - pub async fn dump_events(&self) -> Result> { + /// peer's [`Self::apply_event_routed`]: the primary device-tier dump first, + /// then every allocated lower-tier indexer's dump with `storage_tier` + /// retagged to the tier it lives in. + pub async fn dump_events(&self) -> std::result::Result, KvRouterError> { let (primary_events, lower_tier_entries) = match self { Indexer::Single { primary, lower_tier, .. - } => ( - primary.dump_events().await.map_err(anyhow::Error::from)?, - lower_tier.entries(), - ), + } => (primary.dump_events().await?, lower_tier.entries()), Indexer::Concurrent { primary, lower_tier, .. - } => ( - primary.dump_events().await.map_err(anyhow::Error::from)?, - lower_tier.entries(), - ), + } => (primary.dump_events().await?, lower_tier.entries()), // The remote service owns the events; a peer recovers from it, not // from this process. Indexer::Remote { .. } => return Ok(Vec::new()), + Indexer::None => { + return Err(KvRouterError::Unsupported( + "event dumping requires a KV indexer".to_string(), + )); + } }; let mut out = primary_events; for (tier, indexer) in lower_tier_entries { - let events = indexer.dump_events().await.map_err(anyhow::Error::from)?; + let events = indexer.dump_events().await?; for mut event in events { event.storage_tier = tier; out.push(event); @@ -758,6 +867,41 @@ impl Indexer { } Ok(out) } + + pub fn uses_approximate_lru(&self) -> bool { + match self { + Self::Single { primary, .. } => primary.approximate_lru_enabled(), + Self::Concurrent { primary, .. } => primary.approximate_lru_enabled(), + Self::Remote { .. } | Self::None => false, + } + } + + pub fn set_approximate_lru_capacity_now( + &self, + worker: WorkerWithDpRank, + incarnation: ApproximateLruIncarnation, + capacity: Option, + ) -> std::result::Result<(), KvRouterError> { + match self { + Self::Single { primary, .. } => { + primary.set_approximate_lru_capacity_now(worker, incarnation, capacity) + } + Self::Concurrent { primary, .. } => { + primary.set_approximate_lru_capacity_now(worker, incarnation, capacity) + } + Self::Remote { .. } | Self::None => Ok(()), + } + } + + pub async fn approximate_lru_stats( + &self, + ) -> std::result::Result { + match self { + Self::Single { primary, .. } => primary.approximate_lru_stats().await, + Self::Concurrent { primary, .. } => primary.approximate_lru_stats().await, + Self::Remote { .. } | Self::None => Ok(ApproximateLruStats::default()), + } + } } #[async_trait] @@ -1045,7 +1189,7 @@ mod tests { lower_tier.get_or_create(StorageTier::HostPinned), lower_tier.get_or_create(StorageTier::Disk), ), - Indexer::Concurrent { .. } | Indexer::Remote { .. } => { + Indexer::Concurrent { .. } | Indexer::Remote { .. } | Indexer::None => { unreachable!("test creates the single indexer") } }; @@ -1120,6 +1264,7 @@ mod tests { flush_side(approx).await; } Indexer::Remote { approx, .. } => flush_side(approx).await, + Indexer::None => {} } } @@ -1522,9 +1667,15 @@ pub fn create_indexer_with_policy( metrics: Arc, policy: &IndexerPolicy, ) -> Indexer { - let approx = policy - .side_ttl - .map(|ttl| SideIndexer::new(ttl, block_size, num_threads, Arc::clone(&metrics))); + let approx = policy.side_ttl.map(|ttl| { + SideIndexer::new( + ttl, + block_size, + num_threads, + Arc::clone(&metrics), + CancellationToken::new(), + ) + }); let retention = match &policy.primary { PrimaryRetention::EventDriven => None, PrimaryRetention::ApproximateTtl(ttl) => { @@ -1534,6 +1685,7 @@ pub fn create_indexer_with_policy( return Indexer::Remote { primary: Arc::new(RemoteIndexerClient::new(Arc::clone(transport), key.clone())), approx, + primary_records_routing_decisions: false, }; } }; diff --git a/lib/kv-router/src/services/indexer/lookup.rs b/lib/kv-router/src/services/indexer/lookup.rs new file mode 100644 index 000000000000..1ab24bc3bc35 --- /dev/null +++ b/lib/kv-router/src/services/indexer/lookup.rs @@ -0,0 +1,1190 @@ +// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Match lookups over an [`Indexer`]: primary device tier, lower tiers seeded +//! from confirmed primary matches, and the predict-on-route side overlay. + +use crate::indexer::{ + KvIndexer, KvRouterError, LowerTierIndexers, LowerTierQueryOptions, MatchDetails, + TieredMatchDetails, query_lower_tiers_with_options, +}; +use crate::protocols::{LocalBlockHash, OverlapScores}; +use crate::{ConcurrentRadixTreeCompressed, ThreadPoolIndexer}; + +use super::backend::{Indexer, RemotePrimary, SideIndexer}; + +/// Block hashes for one lookup, borrowed when the caller keeps them and owned +/// when it does not, so a borrowed query is copied only at the one boundary +/// that needs an owned `Vec`. +pub enum HashInput<'a> { + Borrowed(&'a [LocalBlockHash]), + Owned(Vec), +} + +impl<'a> HashInput<'a> { + pub fn as_slice(&self) -> &[LocalBlockHash] { + match self { + Self::Borrowed(hashes) => hashes, + Self::Owned(hashes) => hashes.as_slice(), + } + } + + fn clone_for_boundary(&self) -> Vec { + self.as_slice().to_vec() + } + + fn into_owned_at_boundary(self) -> Vec { + match self { + Self::Borrowed(hashes) => hashes.to_vec(), + Self::Owned(hashes) => hashes, + } + } +} + +pub(super) struct LookupPipeline<'a> { + primary: PrimaryLookup<'a>, + lower_tier: Option<&'a LowerTierIndexers>, + side: Option<&'a SideIndexer>, +} + +enum PrimaryLookup<'a> { + Single(&'a KvIndexer), + Concurrent(&'a ThreadPoolIndexer), + Remote(&'a dyn RemotePrimary), + None, +} + +impl Indexer { + pub(super) fn lookup_pipeline(&self) -> LookupPipeline<'_> { + match self { + Self::Single { + primary, + lower_tier, + approx, + .. + } => LookupPipeline { + primary: PrimaryLookup::Single(primary), + lower_tier: Some(lower_tier), + side: approx.as_ref(), + }, + Self::Concurrent { + primary, + lower_tier, + approx, + .. + } => LookupPipeline { + primary: PrimaryLookup::Concurrent(primary.as_ref()), + lower_tier: Some(lower_tier), + side: approx.as_ref(), + }, + Self::Remote { + primary, approx, .. + } => LookupPipeline { + primary: PrimaryLookup::Remote(primary.as_ref()), + lower_tier: None, + side: approx.as_ref(), + }, + Self::None => LookupPipeline { + primary: PrimaryLookup::None, + lower_tier: None, + side: None, + }, + } + } + + /// Device-tier match details with the side overlay merged in. + pub async fn find_match_details( + &self, + sequence: Vec, + ) -> Result { + self.lookup_pipeline() + .find_match_details(HashInput::Owned(sequence)) + .await + } + + /// Device-tier match details from the primary only (no side overlay). + pub async fn find_primary_match_details( + &self, + sequence: Vec, + ) -> Result { + self.lookup_pipeline() + .find_primary_match_details(HashInput::Owned(sequence)) + .await + } + + /// Device match details plus per-tier hits. Router-hint chain retention is + /// honored as requested; see [`Indexer::find_tiered_matches_with_options`] + /// for the capability-gated variant. + pub async fn find_matches_by_tier( + &self, + sequence: Vec, + ) -> Result { + self.find_matches_by_tier_with_options(sequence, LowerTierQueryOptions::default()) + .await + } + + pub async fn find_matches_by_tier_with_options( + &self, + sequence: Vec, + lower_tier_options: LowerTierQueryOptions, + ) -> Result { + self.lookup_pipeline() + .find_matches_by_tier(HashInput::Owned(sequence), lower_tier_options) + .await + } + + pub async fn find_matches_by_tier_ref( + &self, + sequence: &[LocalBlockHash], + ) -> Result { + self.find_matches_by_tier_ref_with_options(sequence, LowerTierQueryOptions::default()) + .await + } + + pub async fn find_matches_by_tier_ref_with_options( + &self, + sequence: &[LocalBlockHash], + lower_tier_options: LowerTierQueryOptions, + ) -> Result { + self.lookup_pipeline() + .find_matches_by_tier(HashInput::Borrowed(sequence), lower_tier_options) + .await + } + + /// Tiered matches from the primary and lower tiers only (no side overlay). + pub async fn find_primary_matches_by_tier( + &self, + sequence: Vec, + ) -> Result { + self.lookup_pipeline() + .find_primary_matches_by_tier(HashInput::Owned(sequence)) + .await + } +} + +impl<'a> LookupPipeline<'a> { + async fn find_match_details( + &self, + sequence: HashInput<'_>, + ) -> Result { + if self.side.is_none() { + return self.find_primary_match_details(sequence).await; + } + + let primary_details = self.primary.find_match_details_retained(&sequence).await?; + Ok(merge_side_or_warn(self.side, primary_details, sequence).await) + } + + async fn find_primary_match_details( + &self, + sequence: HashInput<'_>, + ) -> Result { + self.primary.find_match_details(sequence).await + } + + async fn find_matches_by_tier( + &self, + sequence: HashInput<'_>, + lower_tier_options: LowerTierQueryOptions, + ) -> Result { + match self.primary { + PrimaryLookup::Single(_) | PrimaryLookup::Concurrent(_) => { + let Some(lower_tier) = self.lower_tier else { + return Ok(TieredMatchDetails::default()); + }; + // Seed lower-tier continuations from confirmed primary matches + // only. Predict-on-route side scores are unconfirmed; using + // them as lower-tier anchors would over-credit host/disk cache + // hits and break the score/hash lockstep `query_lower_tiers` + // expects. + let primary_device = self + .primary + .find_match_details_retained_with_options( + &sequence, + lower_tier_options.retain_kv_transfer_chain, + ) + .await?; + let lt = query_lower_tiers_with_options( + lower_tier, + sequence.as_slice(), + &primary_device, + lower_tier_options, + ); + let device = merge_side_or_warn(self.side, primary_device, sequence).await; + + Ok(TieredMatchDetails { + device, + lower_tier: lt, + }) + } + PrimaryLookup::Remote(primary) => { + if lower_tier_options.retain_kv_transfer_chain { + tracing::warn!( + "router_hint chain retention is not supported with remote primary indexer; proceeding without router hints" + ); + } + let Some(side) = self.side else { + return primary + .find_matches_by_tier(sequence.into_owned_at_boundary(), false) + .await + .map_err(|e| { + tracing::warn!(error = %e, "Remote indexer tiered query failed"); + KvRouterError::IndexerOffline + }); + }; + let mut tiered = primary + .find_matches_by_tier(sequence.clone_for_boundary(), false) + .await + .map_err(|e| { + tracing::warn!(error = %e, "Remote indexer tiered query failed"); + KvRouterError::IndexerOffline + })?; + tiered.device = merge_side_or_warn(Some(side), tiered.device, sequence).await; + Ok(tiered) + } + PrimaryLookup::None => Ok(TieredMatchDetails::default()), + } + } + + async fn find_primary_matches_by_tier( + &self, + sequence: HashInput<'_>, + ) -> Result { + match self.primary { + PrimaryLookup::Single(_) | PrimaryLookup::Concurrent(_) => { + let Some(lower_tier) = self.lower_tier else { + return Ok(TieredMatchDetails::default()); + }; + let device = self.primary.find_match_details_retained(&sequence).await?; + let lt = query_lower_tiers_with_options( + lower_tier, + sequence.as_slice(), + &device, + LowerTierQueryOptions::default(), + ); + Ok(TieredMatchDetails { + device, + lower_tier: lt, + }) + } + PrimaryLookup::Remote(primary) => primary + .find_matches_by_tier(sequence.into_owned_at_boundary(), false) + .await + .map_err(|e| { + tracing::warn!(error = %e, "Remote indexer tiered query failed"); + KvRouterError::IndexerOffline + }), + PrimaryLookup::None => Ok(TieredMatchDetails::default()), + } + } +} + +impl<'a> PrimaryLookup<'a> { + async fn find_match_details( + &self, + sequence: HashInput<'_>, + ) -> Result { + let primary_details = match self { + Self::Single(primary) => { + primary + .find_match_details(sequence.into_owned_at_boundary()) + .await? + } + Self::Concurrent(primary) => primary + .backend() + .find_match_details_impl(sequence.as_slice(), false), + Self::Remote(primary) => { + let tiered = primary + .find_matches_by_tier(sequence.into_owned_at_boundary(), true) + .await + .map_err(|e| { + tracing::warn!(error = %e, "Remote indexer query failed"); + KvRouterError::IndexerOffline + })?; + tiered.device + } + Self::None => return Ok(MatchDetails::new()), + }; + + Ok(primary_details) + } + + async fn find_match_details_retained( + &self, + sequence: &HashInput<'_>, + ) -> Result { + self.find_match_details_retained_with_options(sequence, false) + .await + } + + async fn find_match_details_retained_with_options( + &self, + sequence: &HashInput<'_>, + retain_kv_transfer_chain: bool, + ) -> Result { + let primary_details = match self { + Self::Single(primary) => { + primary + .find_match_details_with_options( + sequence.clone_for_boundary(), + retain_kv_transfer_chain, + ) + .await? + } + Self::Concurrent(primary) => primary.backend().find_match_details_impl_with_options( + sequence.as_slice(), + false, + retain_kv_transfer_chain, + ), + Self::Remote(primary) => { + let tiered = primary + .find_matches_by_tier(sequence.clone_for_boundary(), true) + .await + .map_err(|e| { + tracing::warn!(error = %e, "Remote indexer query failed"); + KvRouterError::IndexerOffline + })?; + tiered.device + } + Self::None => return Ok(MatchDetails::new()), + }; + + Ok(primary_details) + } +} + +/// Merge a side-indexer's `OverlapScores` into the primary's `MatchDetails` +/// by per-worker max. The side indexer covers the window before the engine's +/// first KV event arrives; `last_matched_hashes`, `frequencies`, and +/// `tree_sizes` stay the primary's. +/// +/// The result no longer satisfies the `scores` <-> `last_matched_hashes` +/// lockstep (side-only workers gain a score with no paired hash), so it must +/// not seed `query_lower_tiers`; callers run the lower-tier query against the +/// primary-only details first. +fn merge_overlap_scores(mut primary: MatchDetails, side: OverlapScores) -> MatchDetails { + for (worker, side_score) in side.scores { + primary + .overlap_scores + .scores + .entry(worker) + .and_modify(|s| { + if side_score > *s { + *s = side_score; + } + }) + .or_insert(side_score); + } + primary +} + +/// Query the side indexer (if present) and merge its scores into `primary`. +/// On query error, warn and return `primary` unchanged so the caller still has +/// a usable scheduling signal. +pub(super) async fn merge_side_or_warn( + side: Option<&SideIndexer>, + primary: MatchDetails, + sequence: HashInput<'_>, +) -> MatchDetails { + let Some(side) = side else { + return primary; + }; + match side.find_matches_input(sequence).await { + Ok(side_scores) => merge_overlap_scores(primary, side_scores), + Err(error) => { + tracing::warn!( + error = %error, + "predict-on-route side indexer query failed; using primary only" + ); + primary + } + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + use std::time::Duration; + + use tokio_util::sync::CancellationToken; + + use crate::indexer::LowerTierIndexers; + use crate::services::indexer::backend::Indexer; + use crate::services::indexer::backend::test_util::store_event; + use crate::{ + ConcurrentRadixTreeCompressed, ThreadPoolIndexer, + approx::PruneConfig, + indexer::{KvIndexer, KvIndexerInterface, KvIndexerMetrics, RoutingDecisionHashes}, + protocols::{ + BlockHashOptions, LocalBlockHash, StorageTier, TokensWithHashes, WorkerWithDpRank, + compute_block_hash_for_seq, compute_seq_hash_for_block, + }, + }; + + fn make_test_indexer() -> Indexer { + Indexer::Single { + primary: KvIndexer::new( + CancellationToken::new(), + 4, + Arc::new(KvIndexerMetrics::new_unregistered()), + ), + lower_tier: LowerTierIndexers::new(1, 4), + approx: None, + primary_records_routing_decisions: false, + } + } + + fn make_test_concurrent_indexer() -> Indexer { + Indexer::Concurrent { + primary: Arc::new(ThreadPoolIndexer::new( + ConcurrentRadixTreeCompressed::new(), + 2, + 4, + )), + lower_tier: LowerTierIndexers::new(2, 4), + approx: None, + primary_records_routing_decisions: false, + } + } + + fn make_test_concurrent_approx_indexer() -> Indexer { + Indexer::Concurrent { + primary: Arc::new(ThreadPoolIndexer::new_with_pruning( + ConcurrentRadixTreeCompressed::new(), + 2, + 4, + PruneConfig { + ttl: Duration::from_secs(60), + }, + )), + lower_tier: LowerTierIndexers::new(2, 4), + approx: None, + primary_records_routing_decisions: true, + } + } + + #[test] + fn overlap_refresh_is_limited_to_local_indexers() { + assert!(make_test_indexer().supports_overlap_refresh()); + assert!(make_test_concurrent_indexer().supports_overlap_refresh()); + assert!(!Indexer::None.supports_overlap_refresh()); + } + + #[test] + fn router_hint_chain_retention_requires_event_driven_primary() { + assert!(make_test_indexer().supports_kv_transfer_chain_retention()); + assert!(make_test_concurrent_indexer().supports_kv_transfer_chain_retention()); + assert!(!make_test_concurrent_approx_indexer().supports_kv_transfer_chain_retention()); + assert!(!Indexer::None.supports_kv_transfer_chain_retention()); + } + + async fn flush_indexer(indexer: &Indexer) { + match indexer { + Indexer::Single { + primary, + lower_tier, + .. + } => { + let _ = primary.flush().await; + for indexer in lower_tier.all() { + let _ = indexer.dump_events().await.unwrap(); + } + } + Indexer::Concurrent { + primary, + lower_tier, + .. + } => { + primary.flush().await; + for indexer in lower_tier.all() { + let _ = indexer.dump_events().await.unwrap(); + } + } + Indexer::Remote { .. } | Indexer::None => {} + } + } + + async fn assert_rank_reset_is_acknowledged(indexer: Indexer) { + let reset_rank = WorkerWithDpRank::new(7, 0); + let retained_rank = WorkerWithDpRank::new(7, 1); + + for dp_rank in [reset_rank.dp_rank, retained_rank.dp_rank] { + indexer + .try_apply_event(store_event(7, dp_rank, 1, &[], &[41], StorageTier::Device)) + .await + .unwrap(); + indexer + .try_apply_event(store_event( + 7, + dp_rank, + 2, + &[41], + &[42], + StorageTier::HostPinned, + )) + .await + .unwrap(); + } + flush_indexer(&indexer).await; + + indexer + .reset_worker_dp_rank_and_wait(reset_rank.worker_id, reset_rank.dp_rank) + .await + .unwrap(); + + let matches = indexer + .find_matches_by_tier(vec![LocalBlockHash(41), LocalBlockHash(42)]) + .await + .unwrap(); + assert!( + !matches + .device + .overlap_scores + .scores + .contains_key(&reset_rank) + ); + assert_eq!( + matches.device.overlap_scores.scores.get(&retained_rank), + Some(&1) + ); + let host_hits = &matches + .lower_tier + .get(&StorageTier::HostPinned) + .unwrap() + .hits; + assert!(!host_hits.contains_key(&reset_rank)); + assert_eq!(host_hits.get(&retained_rank), Some(&1)); + } + + #[tokio::test] + async fn single_thread_rank_reset_waits_for_all_local_tiers() { + assert_rank_reset_is_acknowledged(make_test_indexer()).await; + } + + #[tokio::test] + async fn concurrent_rank_reset_waits_for_all_local_tiers() { + assert_rank_reset_is_acknowledged(make_test_concurrent_indexer()).await; + } + + #[tokio::test] + async fn tiered_query_chains_device_host_and_disk() { + let indexer = make_test_indexer(); + let worker = WorkerWithDpRank::new(7, 0); + + indexer + .try_apply_event(store_event(7, 0, 1, &[], &[11, 12], StorageTier::Device)) + .await + .unwrap(); + indexer + .try_apply_event(store_event( + 7, + 0, + 2, + &[11, 12], + &[13], + StorageTier::HostPinned, + )) + .await + .unwrap(); + indexer + .try_apply_event(store_event( + 7, + 0, + 3, + &[11, 12, 13], + &[14], + StorageTier::Disk, + )) + .await + .unwrap(); + flush_indexer(&indexer).await; + + let matches = indexer + .find_matches_by_tier(vec![ + LocalBlockHash(11), + LocalBlockHash(12), + LocalBlockHash(13), + LocalBlockHash(14), + ]) + .await + .unwrap(); + + assert_eq!(matches.device.overlap_scores.scores.get(&worker), Some(&2)); + assert_eq!( + matches + .lower_tier + .get(&StorageTier::HostPinned) + .and_then(|tier| tier.hits.get(&worker)), + Some(&1) + ); + assert_eq!( + matches + .lower_tier + .get(&StorageTier::Disk) + .and_then(|tier| tier.hits.get(&worker)), + Some(&1) + ); + } + + #[tokio::test] + async fn router_dump_includes_all_allocated_physical_tiers() { + let indexer = make_test_indexer(); + for (event_id, tier, block) in [ + (1, StorageTier::Device, 11), + (2, StorageTier::HostPinned, 12), + (3, StorageTier::Disk, 13), + ] { + indexer + .try_apply_event(store_event(7, 0, event_id, &[], &[block], tier)) + .await + .unwrap(); + } + flush_indexer(&indexer).await; + + let events = indexer.dump_events().await.unwrap(); + for tier in [ + StorageTier::Device, + StorageTier::HostPinned, + StorageTier::Disk, + ] { + assert!(events.iter().any(|event| event.storage_tier == tier)); + } + } + + #[tokio::test] + async fn tiered_query_seeds_lower_tier_only_workers_without_affecting_device_scores() { + let indexer = make_test_indexer(); + let device_worker = WorkerWithDpRank::new(10, 0); + let host_only_worker = WorkerWithDpRank::new(20, 0); + let disk_only_worker = WorkerWithDpRank::new(30, 0); + + indexer + .try_apply_event(store_event(10, 0, 1, &[], &[21], StorageTier::Device)) + .await + .unwrap(); + indexer + .try_apply_event(store_event(20, 0, 2, &[], &[21], StorageTier::HostPinned)) + .await + .unwrap(); + indexer + .try_apply_event(store_event(30, 0, 3, &[], &[21], StorageTier::Disk)) + .await + .unwrap(); + flush_indexer(&indexer).await; + + let matches = indexer + .find_matches_by_tier(vec![LocalBlockHash(21)]) + .await + .unwrap(); + + assert_eq!( + matches.device.overlap_scores.scores.get(&device_worker), + Some(&1) + ); + assert!( + !matches + .device + .overlap_scores + .scores + .contains_key(&host_only_worker) + ); + assert!( + !matches + .device + .overlap_scores + .scores + .contains_key(&disk_only_worker) + ); + + assert_eq!( + matches + .lower_tier + .get(&StorageTier::HostPinned) + .and_then(|tier| tier.hits.get(&host_only_worker)), + Some(&1) + ); + assert_eq!( + matches + .lower_tier + .get(&StorageTier::Disk) + .and_then(|tier| tier.hits.get(&disk_only_worker)), + Some(&1) + ); + } + + #[tokio::test] + async fn tiered_query_only_seeds_matching_root_workers() { + let indexer = make_test_indexer(); + let matching_host_worker = WorkerWithDpRank::new(20, 0); + let nonmatching_host_worker = WorkerWithDpRank::new(21, 0); + + indexer + .try_apply_event(store_event(20, 0, 1, &[], &[31], StorageTier::HostPinned)) + .await + .unwrap(); + indexer + .try_apply_event(store_event(21, 0, 2, &[], &[32], StorageTier::HostPinned)) + .await + .unwrap(); + flush_indexer(&indexer).await; + + let matches = indexer + .find_matches_by_tier(vec![LocalBlockHash(31)]) + .await + .unwrap(); + + assert_eq!( + matches + .lower_tier + .get(&StorageTier::HostPinned) + .and_then(|tier| tier.hits.get(&matching_host_worker)), + Some(&1) + ); + assert!( + !matches + .lower_tier + .get(&StorageTier::HostPinned) + .is_some_and(|tier| tier.hits.contains_key(&nonmatching_host_worker)) + ); + } + + #[tokio::test] + async fn concurrent_tiered_query_chains_device_and_lower_tier_matches() { + let indexer = make_test_concurrent_indexer(); + let worker = WorkerWithDpRank::new(7, 0); + + indexer + .try_apply_event(store_event(7, 0, 1, &[], &[11, 12], StorageTier::Device)) + .await + .unwrap(); + indexer + .try_apply_event(store_event( + 7, + 0, + 2, + &[11, 12], + &[13], + StorageTier::HostPinned, + )) + .await + .unwrap(); + flush_indexer(&indexer).await; + + let matches = indexer + .find_matches_by_tier(vec![ + LocalBlockHash(11), + LocalBlockHash(12), + LocalBlockHash(13), + ]) + .await + .unwrap(); + + assert_eq!(matches.device.overlap_scores.scores.get(&worker), Some(&2)); + assert_eq!( + matches + .lower_tier + .get(&StorageTier::HostPinned) + .and_then(|tier| tier.hits.get(&worker)), + Some(&1) + ); + } + + #[tokio::test] + async fn concurrent_records_hashed_routing_decision() { + let indexer = make_test_concurrent_approx_indexer(); + assert!(indexer.records_routing_decisions()); + + let worker = WorkerWithDpRank::new(7, 0); + let tokens = vec![1, 2, 3, 4]; + let block_hashes = compute_block_hash_for_seq(&tokens, 4, BlockHashOptions::default()); + let sequence_hashes = compute_seq_hash_for_block(&block_hashes); + + indexer + .record_hashed_routing_decision(worker, block_hashes.clone(), sequence_hashes) + .await + .unwrap(); + flush_indexer(&indexer).await; + + let matches = indexer.find_matches_by_tier(block_hashes).await.unwrap(); + assert_eq!(matches.device.overlap_scores.scores.get(&worker), Some(&1)); + } + + #[tokio::test] + async fn concurrent_records_precomputed_routing_hashes() { + let indexer = make_test_concurrent_approx_indexer(); + assert!(indexer.records_routing_decisions()); + + let worker = WorkerWithDpRank::new(7, 0); + let local_hashes = vec![LocalBlockHash(91), LocalBlockHash(92)]; + let sequence_hashes = compute_seq_hash_for_block(&local_hashes); + indexer + .record_routing_decision_hashes( + worker, + RoutingDecisionHashes { + local_hashes: local_hashes.clone(), + sequence_hashes, + }, + ) + .await + .unwrap(); + flush_indexer(&indexer).await; + + let matches = indexer.find_matches_by_tier(local_hashes).await.unwrap(); + assert_eq!(matches.device.overlap_scores.scores.get(&worker), Some(&2)); + } + + #[tokio::test] + async fn event_driven_primary_without_side_skips_route_recording() { + let indexer = make_test_indexer(); + assert!(!indexer.records_routing_decisions()); + + let worker = WorkerWithDpRank::new(7, 0); + let tokens = vec![1, 2, 3, 4]; + let block_hashes = compute_block_hash_for_seq(&tokens, 4, BlockHashOptions::default()); + let mut tokens_with_hashes = TokensWithHashes::new(tokens, 4); + + indexer + .process_routing_decision_for_request(&mut tokens_with_hashes, worker) + .await + .unwrap(); + flush_indexer(&indexer).await; + + let matches = indexer.find_matches_by_tier(block_hashes).await.unwrap(); + assert!( + !matches.device.overlap_scores.scores.contains_key(&worker), + "event-driven primary without side overlay should rely on KV events, not route-time writes" + ); + } + + #[tokio::test] + async fn side_only_worker_scored_but_not_used_as_lower_tier_anchor() { + // Build an Indexer::Concurrent with a real side indexer so + // `record_hashed_routing_decision` populates only the side path. + let primary = Arc::new(ThreadPoolIndexer::new( + ConcurrentRadixTreeCompressed::new(), + 2, + 4, + )); + // PruneConfig is required to enable routing-decision recording on the + // side indexer; without it the routing-decision path is a no-op. + let side = Arc::new(ThreadPoolIndexer::new_with_pruning( + ConcurrentRadixTreeCompressed::new(), + 1, + 4, + PruneConfig { + ttl: Duration::from_secs(60), + }, + )); + let side_for_flush = side.clone(); + let indexer = Indexer::Concurrent { + primary, + lower_tier: LowerTierIndexers::new(2, 4), + approx: Some(super::SideIndexer::Concurrent(side)), + primary_records_routing_decisions: false, + }; + assert!(indexer.records_routing_decisions()); + + let primary_worker = WorkerWithDpRank::new(10, 0); + let side_only_worker = WorkerWithDpRank::new(20, 0); + + // Primary sees blocks [11, 12, 13] on Device for primary_worker; + // extension block [14] on HostPinned for primary_worker. + indexer + .try_apply_event(store_event( + 10, + 0, + 1, + &[], + &[11, 12, 13], + StorageTier::Device, + )) + .await + .unwrap(); + indexer + .try_apply_event(store_event( + 10, + 0, + 2, + &[11, 12, 13], + &[14], + StorageTier::HostPinned, + )) + .await + .unwrap(); + // Crucially, also give side_only_worker a HostPinned extension at + // block 14 anchored on the same prefix [11, 12, 13]. If the lower + // tier were seeded from the side-merged device score, the host walk + // would find this and credit a hit; with the reorder it should not. + indexer + .try_apply_event(store_event( + 20, + 0, + 3, + &[11, 12, 13], + &[14], + StorageTier::HostPinned, + )) + .await + .unwrap(); + + // Side-only: route a decision so the side indexer learns + // side_only_worker for the same device prefix. Primary never sees it. + let block_hashes: Vec = + [11, 12, 13].iter().copied().map(LocalBlockHash).collect(); + let sequence_hashes = compute_seq_hash_for_block(&block_hashes); + indexer + .record_hashed_routing_decision(side_only_worker, block_hashes.clone(), sequence_hashes) + .await + .unwrap(); + + flush_indexer(&indexer).await; + side_for_flush.flush().await; + + let matches = indexer + .find_matches_by_tier(vec![ + LocalBlockHash(11), + LocalBlockHash(12), + LocalBlockHash(13), + LocalBlockHash(14), + ]) + .await + .unwrap(); + + // Merge worked: both workers carry device scores. + assert_eq!( + matches + .device + .overlap_scores + .scores + .get(&primary_worker) + .copied(), + Some(3) + ); + assert_eq!( + matches + .device + .overlap_scores + .scores + .get(&side_only_worker) + .copied(), + Some(3), + "side-only worker should appear in merged device scores" + ); + + // Reorder enforced: lower-tier was seeded from primary only. + // primary_worker still extends into HostPinned via its own device + // anchor. side_only_worker's HostPinned extension exists in the + // host tier, but because the side score wasn't used as a device + // anchor, the host walk does not start for it and its host hit is + // not credited. + let host = matches + .lower_tier + .get(&StorageTier::HostPinned) + .expect("host-pinned tier should have been allocated"); + assert_eq!(host.hits.get(&primary_worker).copied(), Some(1)); + assert_eq!( + host.hits.get(&side_only_worker).copied().unwrap_or(0), + 0, + "side-only worker's host extension must not be credited \ + when lower-tier seeding is primary-only" + ); + assert!( + !host.next_continuations.contains_key(&side_only_worker), + "side-only worker must not appear in lower-tier continuations" + ); + } + + #[tokio::test] + async fn borrowed_tiered_lookup_matches_owned_with_lower_tier_and_side_overlay() { + let primary = Arc::new(ThreadPoolIndexer::new( + ConcurrentRadixTreeCompressed::new(), + 2, + 4, + )); + let side = Arc::new(ThreadPoolIndexer::new_with_pruning( + ConcurrentRadixTreeCompressed::new(), + 1, + 4, + PruneConfig { + ttl: Duration::from_secs(60), + }, + )); + let side_for_flush = side.clone(); + let indexer = Indexer::Concurrent { + primary, + lower_tier: LowerTierIndexers::new(2, 4), + approx: Some(super::SideIndexer::Concurrent(side)), + primary_records_routing_decisions: false, + }; + + let primary_worker = WorkerWithDpRank::new(10, 0); + let side_worker = WorkerWithDpRank::new(20, 0); + indexer + .try_apply_event(store_event(10, 0, 1, &[], &[11, 12], StorageTier::Device)) + .await + .unwrap(); + indexer + .try_apply_event(store_event( + 10, + 0, + 2, + &[11, 12], + &[13], + StorageTier::HostPinned, + )) + .await + .unwrap(); + + let side_hashes = vec![LocalBlockHash(11), LocalBlockHash(12), LocalBlockHash(13)]; + indexer + .record_routing_decision_hashes( + side_worker, + RoutingDecisionHashes { + local_hashes: side_hashes.clone(), + sequence_hashes: compute_seq_hash_for_block(&side_hashes), + }, + ) + .await + .unwrap(); + flush_indexer(&indexer).await; + side_for_flush.flush().await; + + let query = vec![LocalBlockHash(11), LocalBlockHash(12), LocalBlockHash(13)]; + let borrowed = indexer.find_matches_by_tier_ref(&query).await.unwrap(); + let owned = indexer.find_matches_by_tier(query).await.unwrap(); + + assert_eq!( + borrowed.device.overlap_scores.scores, + owned.device.overlap_scores.scores + ); + assert_eq!( + borrowed + .lower_tier + .get(&StorageTier::HostPinned) + .map(|tier| &tier.hits), + owned + .lower_tier + .get(&StorageTier::HostPinned) + .map(|tier| &tier.hits) + ); + assert_eq!( + borrowed + .device + .overlap_scores + .scores + .get(&primary_worker) + .copied(), + Some(2) + ); + assert_eq!( + borrowed + .device + .overlap_scores + .scores + .get(&side_worker) + .copied(), + Some(3) + ); + } + + #[tokio::test] + async fn concurrent_tiered_query_seeds_lower_tier_only_workers_without_affecting_device_scores() + { + let indexer = make_test_concurrent_indexer(); + let device_worker = WorkerWithDpRank::new(10, 0); + let host_only_worker = WorkerWithDpRank::new(20, 0); + let disk_only_worker = WorkerWithDpRank::new(30, 0); + + indexer + .try_apply_event(store_event(10, 0, 1, &[], &[21], StorageTier::Device)) + .await + .unwrap(); + indexer + .try_apply_event(store_event(20, 0, 2, &[], &[21], StorageTier::HostPinned)) + .await + .unwrap(); + indexer + .try_apply_event(store_event(30, 0, 3, &[], &[21], StorageTier::Disk)) + .await + .unwrap(); + flush_indexer(&indexer).await; + + let matches = indexer + .find_matches_by_tier(vec![LocalBlockHash(21)]) + .await + .unwrap(); + + assert_eq!( + matches.device.overlap_scores.scores.get(&device_worker), + Some(&1) + ); + assert!( + !matches + .device + .overlap_scores + .scores + .contains_key(&host_only_worker) + ); + assert!( + !matches + .device + .overlap_scores + .scores + .contains_key(&disk_only_worker) + ); + + assert_eq!( + matches + .lower_tier + .get(&StorageTier::HostPinned) + .and_then(|tier| tier.hits.get(&host_only_worker)), + Some(&1) + ); + assert_eq!( + matches + .lower_tier + .get(&StorageTier::Disk) + .and_then(|tier| tier.hits.get(&disk_only_worker)), + Some(&1) + ); + } + + /// Regression test: when a worker has blocks in both device and lower-tier + /// storage (e.g. same prefix stored on GPU and offloaded to host), the + /// Concurrent indexer doesn't return last_matched_hashes. Without the fix, + /// query_lower_tiers would re-query that worker from root in the lower tier, + /// double-counting overlap blocks and producing cached_tokens > ISL. + #[tokio::test] + async fn concurrent_tiered_query_does_not_double_count_device_and_lower_tier_overlap() { + let indexer = make_test_concurrent_indexer(); + let worker = WorkerWithDpRank::new(7, 0); + + // Device owns the prefix block; host-pinned extends it by one block. + indexer + .try_apply_event(store_event(7, 0, 1, &[], &[41], StorageTier::Device)) + .await + .unwrap(); + indexer + .try_apply_event(store_event(7, 0, 2, &[41], &[42], StorageTier::HostPinned)) + .await + .unwrap(); + flush_indexer(&indexer).await; + + let matches = indexer + .find_matches_by_tier(vec![LocalBlockHash(41), LocalBlockHash(42)]) + .await + .unwrap(); + + assert_eq!(matches.device.overlap_scores.scores.get(&worker), Some(&1)); + + let host_hits = matches + .lower_tier + .get(&StorageTier::HostPinned) + .and_then(|tier| tier.hits.get(&worker).copied()) + .unwrap_or(0); + assert_eq!( + host_hits, 1, + "lower-tier should extend the device prefix without double-counting it" + ); + } +} diff --git a/lib/kv-router/src/services/indexer/mod.rs b/lib/kv-router/src/services/indexer/mod.rs index 532366d882b1..2f4838391091 100644 --- a/lib/kv-router/src/services/indexer/mod.rs +++ b/lib/kv-router/src/services/indexer/mod.rs @@ -32,7 +32,9 @@ pub mod backend; pub mod listener; pub mod logging; +pub mod lookup; pub mod metrics; +pub mod recording; pub mod recovery; pub mod registry; pub mod server; diff --git a/lib/llm/src/kv_router/indexer/recording.rs b/lib/kv-router/src/services/indexer/recording.rs similarity index 71% rename from lib/llm/src/kv_router/indexer/recording.rs rename to lib/kv-router/src/services/indexer/recording.rs index 596938f71681..cb9508a91b74 100644 --- a/lib/llm/src/kv_router/indexer/recording.rs +++ b/lib/kv-router/src/services/indexer/recording.rs @@ -1,36 +1,39 @@ -// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use dynamo_kv_router::{ - ConcurrentRadixTreeCompressed, - indexer::{ - ApproximateAcquireMode, ApproximateLruBlock, ApproximateLruIncarnation, - ApproximateLruLease, ApproximateLruReleaseAck, KvIndexer, KvIndexerInterface, - KvRouterError, RoutingDecisionHashes, ThreadPoolIndexer, - }, - protocols::{LocalBlockHash, TokensWithHashes, WorkerWithDpRank}, - scheduling::AttemptId, -}; -use dynamo_tokens::SequenceHash; +//! Routing-decision recording: which physical index a booked routing decision +//! is written into, and the per-request approximate-LRU lease. + use std::sync::{ Arc, atomic::{AtomicU8, Ordering}, }; -use super::{Indexer, SideIndexer, remote::RemoteIndexer}; +use dynamo_tokens::SequenceHash; + +use crate::indexer::{ + ApproximateAcquireMode, ApproximateLruBlock, ApproximateLruIncarnation, ApproximateLruLease, + ApproximateLruReleaseAck, KvIndexer, KvIndexerInterface, KvRouterError, RoutingDecisionHashes, +}; +use crate::protocols::{LocalBlockHash, TokensWithHashes, WorkerWithDpRank}; +use crate::scheduling::AttemptId; +use crate::{ConcurrentRadixTreeCompressed, ThreadPoolIndexer}; + +use super::backend::{Indexer, RemotePrimary, SideIndexer}; const MODE_INACTIVE: u8 = 0; const MODE_LRU: u8 = 1; const MODE_TTL_FALLBACK: u8 = 2; +/// One request's approximate-LRU lease on its primary indexer. #[derive(Clone)] -pub(crate) struct ApproximateRequestLease { +pub struct ApproximateRequestLease { lease: ApproximateLruLease, mode: Arc, } impl ApproximateRequestLease { - pub(crate) async fn acquire( + pub async fn acquire( &mut self, hashes: RoutingDecisionHashes, private_blocks: usize, @@ -56,7 +59,7 @@ impl ApproximateRequestLease { Ok(mode) } - pub(crate) fn materialize( + pub fn materialize( &self, parent_hash: Option, blocks: Vec, @@ -70,37 +73,37 @@ impl ApproximateRequestLease { .materialize(parent_hash, blocks, start_position, private_blocks) } - pub(crate) fn begin_finish(&self) -> Result, KvRouterError> { + pub fn begin_finish(&self) -> Result, KvRouterError> { self.lease.begin_finish() } - pub(crate) fn release_now(&self) { + pub fn release_now(&self) { self.lease.release_now(); } - pub(crate) fn is_active_lru(&self) -> bool { + pub fn is_active_lru(&self) -> bool { self.mode.load(Ordering::Acquire) == MODE_LRU } } #[derive(Clone, Copy)] -pub(super) enum RouteRecordingTarget<'a> { +enum RouteRecordingTarget<'a> { Disabled, PrimaryLocal(&'a KvIndexer), PrimaryConcurrent(&'a ThreadPoolIndexer), - PrimaryRemote(&'a RemoteIndexer), + PrimaryRemote(&'a dyn RemotePrimary), SideOverlay(&'a SideIndexer), } impl Indexer { - pub(crate) fn begin_approximate_lru_request( + pub fn begin_approximate_lru_request( &self, worker: WorkerWithDpRank, incarnation: ApproximateLruIncarnation, attempt_id: AttemptId, ) -> Option { let lease = match self { - Self::KvIndexer { primary, .. } => { + Self::Single { primary, .. } => { primary.begin_approximate_lru_request(worker, incarnation, attempt_id)? } Self::Concurrent { primary, .. } => { @@ -114,13 +117,14 @@ impl Indexer { }) } - pub(crate) fn records_routing_decisions(&self) -> bool { + /// Whether booked routing decisions should be written into this indexer. + pub fn records_routing_decisions(&self) -> bool { !matches!(self.recording_target(), RouteRecordingTarget::Disabled) } - pub(super) fn recording_target(&self) -> RouteRecordingTarget<'_> { + fn recording_target(&self) -> RouteRecordingTarget<'_> { match self { - Self::KvIndexer { + Self::Single { approx: Some(side), .. } | Self::Concurrent { @@ -137,7 +141,7 @@ impl Indexer { ); RouteRecordingTarget::SideOverlay(side) } - Self::KvIndexer { + Self::Single { primary, primary_records_routing_decisions: true, .. @@ -152,40 +156,51 @@ impl Indexer { primary_records_routing_decisions: true, .. } => RouteRecordingTarget::PrimaryRemote(primary.as_ref()), - Self::KvIndexer { .. } | Self::Concurrent { .. } | Self::Remote { .. } | Self::None => { + Self::Single { .. } | Self::Concurrent { .. } | Self::Remote { .. } | Self::None => { RouteRecordingTarget::Disabled } } } - pub(crate) async fn record_hashed_routing_decision( + /// Record a booked routing decision. Writes to the side indexer when one + /// is attached, else to an approximate primary, else is a no-op. + pub async fn record_routing_decision( &self, worker: WorkerWithDpRank, - local_hashes: Vec, - sequence_hashes: Vec, + hashes: RoutingDecisionHashes, ) -> Result<(), KvRouterError> { self.recording_target() - .record_routing_hashes( - worker, - RoutingDecisionHashes { - local_hashes, - sequence_hashes, - }, - ) + .record_routing_hashes(worker, hashes) .await } - pub(crate) async fn record_routing_decision_hashes( + /// Same as [`Self::record_routing_decision`]; kept for callers that name + /// the hashes explicitly. + pub async fn record_routing_decision_hashes( &self, worker: WorkerWithDpRank, hashes: RoutingDecisionHashes, ) -> Result<(), KvRouterError> { - self.recording_target() - .record_routing_hashes(worker, hashes) - .await + self.record_routing_decision(worker, hashes).await + } + + pub async fn record_hashed_routing_decision( + &self, + worker: WorkerWithDpRank, + local_hashes: Vec, + sequence_hashes: Vec, + ) -> Result<(), KvRouterError> { + self.record_routing_decision( + worker, + RoutingDecisionHashes { + local_hashes, + sequence_hashes, + }, + ) + .await } - pub(crate) async fn process_routing_decision_for_request( + pub async fn process_routing_decision_for_request( &self, tokens_with_hashes: &mut TokensWithHashes, worker: WorkerWithDpRank, @@ -241,13 +256,13 @@ impl<'a> RouteRecordingTarget<'a> { .await } Self::PrimaryRemote(primary) => primary - .record_hashed_routing_decision(worker, hashes.local_hashes, hashes.sequence_hashes) + .record_routing_decision(worker, hashes) .await .map_err(|error| { tracing::warn!(error = %error, "Remote indexer write failed"); KvRouterError::IndexerDroppedRequest }), - Self::SideOverlay(side) => side.process_routing_decision_hashes(worker, hashes).await, + Self::SideOverlay(side) => side.record_routing_decision(worker, hashes).await, } } } diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index 14fb21efc081..323b9b2d77be 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -758,7 +758,7 @@ impl KvRouter { let min_initial_workers = min_initial_workers_from_env()?; let indexer = if cache_required { - Indexer::new( + indexer::build( component, &kv_router_config, block_size, diff --git a/lib/llm/src/kv_router/indexer/lookup.rs b/lib/llm/src/kv_router/indexer/lookup.rs deleted file mode 100644 index 0b8aa7b19124..000000000000 --- a/lib/llm/src/kv_router/indexer/lookup.rs +++ /dev/null @@ -1,430 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -use dynamo_kv_router::{ - ConcurrentRadixTreeCompressed, - indexer::{ - KvIndexer, KvRouterError, LowerTierIndexers, LowerTierQueryOptions, MatchDetails, - ThreadPoolIndexer, TieredMatchProvider, query_lower_tiers_with_options, - }, - protocols::{LocalBlockHash, OverlapScores}, -}; -use dynamo_runtime::pipeline::async_trait; - -use super::{Indexer, SideIndexer, TieredMatchDetails, remote::RemoteIndexer}; - -pub(super) enum HashInput<'a> { - Borrowed(&'a [LocalBlockHash]), - Owned(Vec), -} - -impl<'a> HashInput<'a> { - pub(super) fn as_slice(&self) -> &[LocalBlockHash] { - match self { - Self::Borrowed(sequence) => sequence, - Self::Owned(sequence) => sequence, - } - } - - pub(super) fn clone_for_boundary(&self) -> Vec { - self.as_slice().to_vec() - } - - pub(super) fn into_owned_at_boundary(self) -> Vec { - match self { - Self::Borrowed(sequence) => sequence.to_vec(), - Self::Owned(sequence) => sequence, - } - } -} - -#[derive(Clone, Copy)] -pub(super) enum PrimaryLookup<'a> { - KvIndexer(&'a KvIndexer), - Concurrent(&'a ThreadPoolIndexer), - Remote(&'a RemoteIndexer), - None, -} - -#[derive(Clone, Copy)] -pub(super) struct LookupPipeline<'a> { - primary: PrimaryLookup<'a>, - lower_tier: Option<&'a LowerTierIndexers>, - side: Option<&'a SideIndexer>, -} - -impl Indexer { - pub(super) fn lookup_pipeline(&self) -> LookupPipeline<'_> { - match self { - Self::KvIndexer { - primary, - lower_tier, - approx, - .. - } => LookupPipeline { - primary: PrimaryLookup::KvIndexer(primary), - lower_tier: Some(lower_tier), - side: approx.as_ref(), - }, - Self::Concurrent { - primary, - lower_tier, - approx, - .. - } => LookupPipeline { - primary: PrimaryLookup::Concurrent(primary.as_ref()), - lower_tier: Some(lower_tier), - side: approx.as_ref(), - }, - Self::Remote { - primary, approx, .. - } => LookupPipeline { - primary: PrimaryLookup::Remote(primary.as_ref()), - lower_tier: None, - side: approx.as_ref(), - }, - Self::None => LookupPipeline { - primary: PrimaryLookup::None, - lower_tier: None, - side: None, - }, - } - } - - #[cfg(test)] - pub(crate) async fn find_match_details( - &self, - sequence: Vec, - ) -> Result { - self.lookup_pipeline() - .find_match_details(HashInput::Owned(sequence)) - .await - } - - pub(crate) async fn find_primary_match_details( - &self, - sequence: Vec, - ) -> Result { - self.lookup_pipeline() - .find_primary_match_details(HashInput::Owned(sequence)) - .await - } - - pub(crate) async fn find_matches_by_tier( - &self, - sequence: Vec, - ) -> Result { - self.find_matches_by_tier_with_options(sequence, LowerTierQueryOptions::default()) - .await - } - - pub(crate) async fn find_matches_by_tier_with_options( - &self, - sequence: Vec, - lower_tier_options: LowerTierQueryOptions, - ) -> Result { - self.lookup_pipeline() - .find_matches_by_tier(HashInput::Owned(sequence), lower_tier_options) - .await - } - - pub(crate) async fn find_matches_by_tier_ref( - &self, - sequence: &[LocalBlockHash], - ) -> Result { - self.find_matches_by_tier_ref_with_options(sequence, LowerTierQueryOptions::default()) - .await - } - - pub(crate) async fn find_matches_by_tier_ref_with_options( - &self, - sequence: &[LocalBlockHash], - lower_tier_options: LowerTierQueryOptions, - ) -> Result { - self.lookup_pipeline() - .find_matches_by_tier(HashInput::Borrowed(sequence), lower_tier_options) - .await - } - - pub(crate) async fn find_primary_matches_by_tier( - &self, - sequence: Vec, - ) -> Result { - self.lookup_pipeline() - .find_primary_matches_by_tier(HashInput::Owned(sequence)) - .await - } -} - -#[async_trait] -impl TieredMatchProvider for Indexer { - async fn find_tiered_matches( - &self, - sequence: &[LocalBlockHash], - ) -> Result { - self.find_matches_by_tier_ref(sequence).await - } - - async fn find_tiered_matches_with_options( - &self, - sequence: &[LocalBlockHash], - options: LowerTierQueryOptions, - ) -> Result { - self.find_matches_by_tier_ref_with_options(sequence, options) - .await - } -} - -impl<'a> LookupPipeline<'a> { - #[cfg(test)] - async fn find_match_details( - &self, - sequence: HashInput<'_>, - ) -> Result { - if self.side.is_none() { - return self.find_primary_match_details(sequence).await; - } - - let primary_details = self.primary.find_match_details_retained(&sequence).await?; - Ok(merge_side_or_warn(self.side, primary_details, sequence).await) - } - - async fn find_primary_match_details( - &self, - sequence: HashInput<'_>, - ) -> Result { - self.primary.find_match_details(sequence).await - } - - async fn find_matches_by_tier( - &self, - sequence: HashInput<'_>, - lower_tier_options: LowerTierQueryOptions, - ) -> Result { - match self.primary { - PrimaryLookup::KvIndexer(_) | PrimaryLookup::Concurrent(_) => { - let Some(lower_tier) = self.lower_tier else { - return Ok(TieredMatchDetails::default()); - }; - // Seed lower-tier continuations from confirmed primary matches - // only. Predict-on-route side scores are unconfirmed; using - // them as lower-tier anchors would over-credit host/disk cache - // hits and break the score/hash lockstep `query_lower_tiers` - // expects. - let primary_device = self - .primary - .find_match_details_retained_with_options( - &sequence, - lower_tier_options.retain_kv_transfer_chain, - ) - .await?; - let lt = query_lower_tiers_with_options( - lower_tier, - sequence.as_slice(), - &primary_device, - lower_tier_options, - ); - let device = merge_side_or_warn(self.side, primary_device, sequence).await; - - Ok(TieredMatchDetails { - device, - lower_tier: lt, - }) - } - PrimaryLookup::Remote(primary) => { - if lower_tier_options.retain_kv_transfer_chain { - tracing::warn!( - "KV transfer chain retention is not supported with remote primary indexer; proceeding without KV transfer hints" - ); - } - let Some(side) = self.side else { - return primary - .find_matches_by_tier(sequence.into_owned_at_boundary(), false) - .await - .map_err(|e| { - tracing::warn!(error = %e, "Remote indexer tiered query failed"); - KvRouterError::IndexerOffline - }); - }; - let mut tiered = primary - .find_matches_by_tier(sequence.clone_for_boundary(), false) - .await - .map_err(|e| { - tracing::warn!(error = %e, "Remote indexer tiered query failed"); - KvRouterError::IndexerOffline - })?; - tiered.device = merge_side_or_warn(Some(side), tiered.device, sequence).await; - Ok(tiered) - } - PrimaryLookup::None => Ok(TieredMatchDetails::default()), - } - } - - async fn find_primary_matches_by_tier( - &self, - sequence: HashInput<'_>, - ) -> Result { - match self.primary { - PrimaryLookup::KvIndexer(_) | PrimaryLookup::Concurrent(_) => { - let Some(lower_tier) = self.lower_tier else { - return Ok(TieredMatchDetails::default()); - }; - let device = self.primary.find_match_details_retained(&sequence).await?; - let lt = query_lower_tiers_with_options( - lower_tier, - sequence.as_slice(), - &device, - LowerTierQueryOptions::default(), - ); - Ok(TieredMatchDetails { - device, - lower_tier: lt, - }) - } - PrimaryLookup::Remote(primary) => primary - .find_matches_by_tier(sequence.into_owned_at_boundary(), false) - .await - .map_err(|e| { - tracing::warn!(error = %e, "Remote indexer tiered query failed"); - KvRouterError::IndexerOffline - }), - PrimaryLookup::None => Ok(TieredMatchDetails::default()), - } - } -} - -impl<'a> PrimaryLookup<'a> { - async fn find_match_details( - &self, - sequence: HashInput<'_>, - ) -> Result { - let primary_details = match self { - Self::KvIndexer(primary) => { - primary - .find_match_details(sequence.into_owned_at_boundary()) - .await? - } - Self::Concurrent(primary) => primary - .backend() - .find_match_details_impl(sequence.as_slice(), false), - Self::Remote(primary) => { - let tiered = primary - .find_matches_by_tier(sequence.into_owned_at_boundary(), true) - .await - .map_err(|e| { - tracing::warn!(error = %e, "Remote indexer query failed"); - KvRouterError::IndexerOffline - })?; - tiered.device - } - Self::None => return Ok(MatchDetails::new()), - }; - - Ok(primary_details) - } - - async fn find_match_details_retained( - &self, - sequence: &HashInput<'_>, - ) -> Result { - self.find_match_details_retained_with_options(sequence, false) - .await - } - - async fn find_match_details_retained_with_options( - &self, - sequence: &HashInput<'_>, - retain_kv_transfer_chain: bool, - ) -> Result { - let primary_details = match self { - Self::KvIndexer(primary) => { - primary - .find_match_details_with_options( - sequence.clone_for_boundary(), - retain_kv_transfer_chain, - ) - .await? - } - Self::Concurrent(primary) => primary.backend().find_match_details_impl_with_options( - sequence.as_slice(), - false, - retain_kv_transfer_chain, - ), - Self::Remote(primary) => { - let tiered = primary - .find_matches_by_tier(sequence.clone_for_boundary(), true) - .await - .map_err(|e| { - tracing::warn!(error = %e, "Remote indexer query failed"); - KvRouterError::IndexerOffline - })?; - tiered.device - } - Self::None => return Ok(MatchDetails::new()), - }; - - Ok(primary_details) - } -} - -/// Merge a side-indexer's `OverlapScores` into the primary's `MatchDetails` -/// by taking the per-worker max overlap. The side indexer covers the window -/// before the engine's first KV event arrives; for workers it knows about, -/// we use whichever indexer saw the longer prefix. `last_matched_hashes`, -/// `frequencies`, and `tree_sizes` come from the primary -- the side -/// indexer's short-TTL view isn't meaningful for those signals. -/// -/// IMPORTANT: the returned `MatchDetails` is no longer guaranteed to satisfy -/// `overlap_scores.scores` <-> `last_matched_hashes` lockstep. Side-only -/// workers gain a score with no paired hash by design. The result is safe -/// for scheduling / cache-hit signal but MUST NOT be used to seed -/// `query_lower_tiers`, which assumes the lockstep invariant. The local -/// arm of `find_matches_by_tier` enforces this by running the lower-tier -/// query against primary-only `MatchDetails` before merging side scores. -fn merge_overlap_scores(mut primary: MatchDetails, side: OverlapScores) -> MatchDetails { - for (worker, side_score) in side.scores { - primary - .overlap_scores - .scores - .entry(worker) - .and_modify(|s| { - if side_score > *s { - *s = side_score; - } - }) - .or_insert(side_score); - } - primary -} - -/// Query the predict-on-route side indexer (if present) and merge its scores -/// into the primary device match details. Side scores never feed lower-tier -/// or shared-cache scoring. On query error, log a warning and return `primary` -/// unchanged so the caller still has a usable scheduling signal. See -/// [`merge_overlap_scores`] for the lockstep caveat on the returned shape. -/// -/// NOTE: when this merged `MatchDetails` is combined with lower-tier hits -/// seeded from the primary-only anchor (e.g. in `find_matches_by_tier`), the -/// total cached-token signal can in theory overcount: the device score is -/// raised by the side indexer but the lower-tier walk used the lower primary -/// depth. Accepted as edge for now since side scores are short-TTL -/// approximations and the overcount is bounded and rare in practice. -async fn merge_side_or_warn( - side: Option<&SideIndexer>, - primary: MatchDetails, - sequence: HashInput<'_>, -) -> MatchDetails { - let Some(side) = side else { - return primary; - }; - match side.find_matches_input(sequence).await { - Ok(side_scores) => merge_overlap_scores(primary, side_scores), - Err(error) => { - tracing::warn!( - error = %error, - "predict-on-route side indexer query failed; using primary only" - ); - primary - } - } -} diff --git a/lib/llm/src/kv_router/indexer/mod.rs b/lib/llm/src/kv_router/indexer/mod.rs index e463f22d7edc..11a42212673a 100644 --- a/lib/llm/src/kv_router/indexer/mod.rs +++ b/lib/llm/src/kv_router/indexer/mod.rs @@ -1,6 +1,10 @@ // SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +//! The frontend's KV index: the shared `dynamo_kv_router` [`Indexer`] built +//! from router configuration with runtime-backed metrics and, when configured, +//! a request-plane remote primary. + use std::{sync::Arc, time::Duration}; use anyhow::Result; @@ -9,13 +13,8 @@ use dynamo_kv_router::{ approx::PruneConfig, config::{ApproximateCachePolicyKind, KvRouterConfig}, indexer::{ - ApproximateLruIncarnation, ApproximateLruStats, ApproximateRetentionConfig, KvIndexer, - KvIndexerInterface, KvIndexerMetrics, KvRouterError, LowerTierIndexers, ThreadPoolIndexer, - record_unsupported_residency_event, - }, - protocols::{ - DpRank, KvCacheEventData, ResidencyProjection, ResidencyRoutingSnapshot, RouterEvent, - WorkerId, + ApproximateRetentionConfig, KvIndexer, KvIndexerMetrics, LowerTierIndexers, + ThreadPoolIndexer, }, }; @@ -24,23 +23,20 @@ use dynamo_kv_router::{ pub(crate) use dynamo_kv_router::indexer::TieredMatchDetails; #[allow(unused_imports)] pub(crate) use dynamo_kv_router::indexer::WireTieredMatchDetails; +pub use dynamo_kv_router::services::indexer::backend::{Indexer, RemotePrimary, SideIndexer}; +pub(crate) use dynamo_kv_router::services::indexer::recording::ApproximateRequestLease; use dynamo_runtime::component::Component; use tokio_util::sync::CancellationToken; mod embedding_cache; -mod lookup; -mod recording; mod recovery; pub mod remote; -mod side; pub use self::embedding_cache::{ EmbeddingCacheIndexer, preprocessed_multimodal_cache_keys, try_build_cache_indexer, }; -pub(crate) use self::recording::ApproximateRequestLease; use self::remote::RemoteIndexer; pub use self::remote::{ServedIndexerHandle, ServedIndexerMode, ensure_served_indexer_service}; -pub use self::side::SideIndexer; #[cfg(feature = "ckf-diagnostics")] pub(crate) use recovery::WorkerQueryHealthSnapshot; pub(crate) use recovery::{ @@ -53,37 +49,6 @@ pub(crate) use recovery::{ start_subscriber, start_worker_kv_query_endpoint, start_worker_kv_query_endpoint_with_status, }; -/// `approx` is the optional predict-on-route side indexer. It is always local -/// to this router, even when the primary indexer is served or consumed -/// remotely. Routing decisions populate it with a short TTL; engine KV events -/// go to the primary only. `find_match_details` queries both and returns the -/// per-worker max overlap. Keeping this separate from the primary avoids the -/// sequence-hash mismatch problem: vLLM/SGLang salt their hashes with -/// cryptographic digests the router can't reproduce, so writing -/// router-computed hashes into the primary would key the same block under two -/// hashes and pollute the tree. -#[derive(Clone)] -pub enum Indexer { - KvIndexer { - primary: KvIndexer, - lower_tier: LowerTierIndexers, - approx: Option, - primary_records_routing_decisions: bool, - }, - Concurrent { - primary: Arc>, - lower_tier: LowerTierIndexers, - approx: Option, - primary_records_routing_decisions: bool, - }, - Remote { - primary: Arc, - approx: Option, - primary_records_routing_decisions: bool, - }, - None, -} - #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum ResolvedApproximatePrimaryPolicy { Disabled, @@ -116,159 +81,88 @@ fn resolve_approximate_primary_policy( Ok(ResolvedApproximatePrimaryPolicy::Lru) } -async fn dump_local_events( - mut events: Vec, - lower_tiers: &LowerTierIndexers, -) -> Result, KvRouterError> { - for (tier, indexer) in lower_tiers.entries() { - events.extend(indexer.dump_events().await?.into_iter().map(|mut event| { - event.storage_tier = tier; - event - })); - } - Ok(events) -} - -impl Indexer { - /// Publish a control-plane projection snapshot for subsequent lookups. - /// - /// Discovery and attachment reconciliation stay in lib/llm; router-core - /// only consumes this immutable resolved view. - pub fn set_residency_projection(&self, projection: ResidencyProjection) { - match self { - Self::KvIndexer { lower_tier, .. } | Self::Concurrent { lower_tier, .. } => { - lower_tier.set_residency_projection(projection) - } - Self::Remote { .. } | Self::None => {} - } - } - - pub fn set_residency_routing_snapshot(&self, snapshot: ResidencyRoutingSnapshot) { - match self { - Self::KvIndexer { lower_tier, .. } | Self::Concurrent { lower_tier, .. } => { - lower_tier.set_residency_routing_snapshot(snapshot) - } - Self::Remote { .. } | Self::None => {} - } +/// Build the frontend router's index for `kv_router_config`. +pub(crate) async fn build( + component: &Component, + kv_router_config: &KvRouterConfig, + block_size: u32, + model_name: Option<&str>, + cancellation_token: CancellationToken, +) -> Result { + let approximate_policy = resolve_approximate_primary_policy(kv_router_config)?; + if approximate_policy == ResolvedApproximatePrimaryPolicy::Disabled { + return Ok(Indexer::None); + } + + if approximate_policy == ResolvedApproximatePrimaryPolicy::TtlRemoteFallback { + tracing::warn!( + use_remote_indexer = kv_router_config.use_remote_indexer, + serve_indexer = kv_router_config.serve_indexer, + "Approximate LRU requires a router-local primary indexer; falling back to TTL" + ); } - pub(crate) fn supports_overlap_refresh(&self) -> bool { - matches!(self, Self::KvIndexer { .. } | Self::Concurrent { .. }) + if kv_router_config.router_predicted_ttl_secs.is_some() && !kv_router_config.use_kv_events { + anyhow::bail!( + "router_predicted_ttl_secs requires use_kv_events=true; \ + do not combine a primary approximate indexer with a side approximate indexer" + ); } - - pub(crate) fn supports_kv_transfer_chain_retention(&self) -> bool { - matches!( - self, - Self::KvIndexer { - approx: None, - primary_records_routing_decisions: false, - .. - } | Self::Concurrent { - approx: None, - primary_records_routing_decisions: false, - .. - } - ) + if kv_router_config.use_remote_indexer { + let model_name = model_name + .ok_or_else(|| { + anyhow::anyhow!("model_name is required when use_remote_indexer is configured") + })? + .to_string(); + let indexer_component_name = component.name(); + tracing::info!( + indexer_component = %indexer_component_name, + model_name, + "Using remote KV indexer" + ); + let remote = + RemoteIndexer::new(component, model_name, kv_router_config.use_kv_events).await?; + let approx = predict_on_route_side_indexer( + component, + kv_router_config, + block_size, + cancellation_token.child_token(), + ); + return Ok(Indexer::Remote { + primary: Arc::new(remote) as Arc, + approx, + primary_records_routing_decisions: !kv_router_config.use_kv_events, + }); } - pub async fn new( - component: &Component, - kv_router_config: &KvRouterConfig, - block_size: u32, - model_name: Option<&str>, - cancellation_token: CancellationToken, - ) -> Result { - let approximate_policy = resolve_approximate_primary_policy(kv_router_config)?; - if approximate_policy == ResolvedApproximatePrimaryPolicy::Disabled { - return Ok(Self::None); - } - - if approximate_policy == ResolvedApproximatePrimaryPolicy::TtlRemoteFallback { - tracing::warn!( - use_remote_indexer = kv_router_config.use_remote_indexer, - serve_indexer = kv_router_config.serve_indexer, - "Approximate LRU requires a router-local primary indexer; falling back to TTL" - ); - } - - if kv_router_config.router_predicted_ttl_secs.is_some() && !kv_router_config.use_kv_events { - anyhow::bail!( - "router_predicted_ttl_secs requires use_kv_events=true; \ - do not combine a primary approximate indexer with a side approximate indexer" - ); - } - if kv_router_config.use_remote_indexer { - let model_name = model_name - .ok_or_else(|| { - anyhow::anyhow!("model_name is required when use_remote_indexer is configured") - })? - .to_string(); - let indexer_component_name = component.name(); + if !kv_router_config.use_kv_events { + let kv_indexer_metrics = KvIndexerMetrics::from_component(component); + let prune_config = PruneConfig { + ttl: Duration::from_secs_f64(kv_router_config.router_ttl_secs), + }; + let retention = if approximate_policy == ResolvedApproximatePrimaryPolicy::Lru { tracing::info!( - indexer_component = %indexer_component_name, - model_name, - "Using remote KV indexer" + "Starting local primary approximate indexer with capacity-bounded LRU retention" ); - let remote = - RemoteIndexer::new(component, model_name, kv_router_config.use_kv_events).await?; - let approx = SideIndexer::new_predict_on_route( - component, - kv_router_config, - block_size, - cancellation_token.child_token(), - ); - return Ok(Self::Remote { - primary: Arc::new(remote), - approx, - primary_records_routing_decisions: !kv_router_config.use_kv_events, - }); - } - - if !kv_router_config.use_kv_events { - let kv_indexer_metrics = KvIndexerMetrics::from_component(component); - let prune_config = PruneConfig { - ttl: Duration::from_secs_f64(kv_router_config.router_ttl_secs), - }; - let retention = if approximate_policy == ResolvedApproximatePrimaryPolicy::Lru { - tracing::info!( - "Starting local primary approximate indexer with capacity-bounded LRU retention" - ); - ApproximateRetentionConfig::Lru { - fallback_ttl: prune_config, - } - } else { - ApproximateRetentionConfig::Ttl(prune_config) - }; - if kv_router_config.router_event_threads > 1 { - return Ok(Self::Concurrent { - primary: Arc::new( - ThreadPoolIndexer::new_with_metrics_and_approximate_retention( - ConcurrentRadixTreeCompressed::new(), - kv_router_config.router_event_threads as usize, - block_size, - Some(kv_indexer_metrics.clone()), - Some(retention), - ), - ), - lower_tier: LowerTierIndexers::new_with_metrics( + ApproximateRetentionConfig::Lru { + fallback_ttl: prune_config, + } + } else { + ApproximateRetentionConfig::Ttl(prune_config) + }; + if kv_router_config.router_event_threads > 1 { + return Ok(Indexer::Concurrent { + primary: Arc::new( + ThreadPoolIndexer::new_with_metrics_and_approximate_retention( + ConcurrentRadixTreeCompressed::new(), kv_router_config.router_event_threads as usize, block_size, - Some(kv_indexer_metrics), + Some(kv_indexer_metrics.clone()), + Some(retention), ), - approx: None, - primary_records_routing_decisions: true, - }); - } - - return Ok(Self::KvIndexer { - primary: KvIndexer::new_with_approximate_retention( - cancellation_token.child_token(), - block_size, - kv_indexer_metrics.clone(), - Some(retention), ), lower_tier: LowerTierIndexers::new_with_metrics( - 1, + kv_router_config.router_event_threads as usize, block_size, Some(kv_indexer_metrics), ), @@ -277,240 +171,83 @@ impl Indexer { }); } - let approx = SideIndexer::new_predict_on_route( - component, - kv_router_config, - block_size, - cancellation_token.child_token(), - ); - - if kv_router_config.router_event_threads > 1 { - let kv_indexer_metrics = KvIndexerMetrics::from_component(component); - return Ok(Self::Concurrent { - primary: Arc::new(ThreadPoolIndexer::new_with_metrics( - ConcurrentRadixTreeCompressed::new(), - kv_router_config.router_event_threads as usize, - block_size, - Some(kv_indexer_metrics.clone()), - )), - lower_tier: LowerTierIndexers::new_with_metrics( - kv_router_config.router_event_threads as usize, - block_size, - Some(kv_indexer_metrics), - ), - approx, - primary_records_routing_decisions: false, - }); - } - - let kv_indexer_metrics = KvIndexerMetrics::from_component(component); - Ok(Self::KvIndexer { - primary: KvIndexer::new_with_pruning( + return Ok(Indexer::Single { + primary: KvIndexer::new_with_approximate_retention( cancellation_token.child_token(), block_size, kv_indexer_metrics.clone(), - None, + Some(retention), ), lower_tier: LowerTierIndexers::new_with_metrics( 1, block_size, Some(kv_indexer_metrics), ), - approx, - primary_records_routing_decisions: false, - }) - } - - pub(crate) async fn dump_events(&self) -> Result, KvRouterError> { - match self { - Self::KvIndexer { - primary, - lower_tier, - .. - } => dump_local_events(primary.dump_events().await?, lower_tier).await, - Self::Concurrent { - primary, - lower_tier, - .. - } => dump_local_events(primary.dump_events().await?, lower_tier).await, - Self::Remote { .. } => Ok(Vec::new()), - Self::None => Err(KvRouterError::Unsupported( - "event dumping requires a KV indexer".to_string(), - )), - } - } - - pub(crate) async fn try_apply_event(&self, event: RouterEvent) -> Result<(), KvRouterError> { - let targets_primary = match event.targets_primary() { - Ok(targets_primary) => targets_primary, - Err(_) => { - match self { - Self::KvIndexer { lower_tier, .. } | Self::Concurrent { lower_tier, .. } => { - lower_tier.record_unsupported_residency_event(&event); - } - Self::Remote { .. } | Self::None => { - record_unsupported_residency_event(None, &event); - } - } - return Ok(()); - } - }; - let is_clear = matches!(&event.event.data, KvCacheEventData::Cleared); - match self { - Self::KvIndexer { - primary, - lower_tier, - .. - } => { - if is_clear { - if targets_primary { - primary - .reset_worker_dp_rank_and_wait(event.worker_id, event.event.dp_rank) - .await?; - } - - for indexer in lower_tier.all() { - indexer.apply_event_and_wait(event.clone()).await?; - } - } else if targets_primary { - primary - .event_sender() - .send(event) - .await - .map_err(|_| KvRouterError::IndexerOffline)?; - } else { - lower_tier - .get_or_create(event.storage_tier) - .enqueue_event(event)?; - } - } - Self::Concurrent { - primary, - lower_tier, - .. - } => { - if is_clear { - if targets_primary { - primary.apply_event_and_wait(event.clone()).await?; - } - - for indexer in lower_tier.all() { - indexer.apply_event_and_wait(event.clone()).await?; - } - } else if targets_primary { - primary.enqueue_event(event)?; - } else { - lower_tier - .get_or_create(event.storage_tier) - .enqueue_event(event)?; - } - } - Self::Remote { .. } | Self::None => {} - } - Ok(()) - } - - #[cfg(test)] - pub(crate) async fn apply_event(&self, event: RouterEvent) { - if let Err(error) = self.try_apply_event(event).await { - tracing::error!(%error, "Failed to enqueue KV event"); - } + approx: None, + primary_records_routing_decisions: true, + }); } - /// Cold-reset one logical rank and wait until all local index tiers have completed the removal. - /// - /// NOTE: Unlike ordinary event application, rank removal is an infallible lane operation. - /// Its FIFO completion must be visible before source activation or clearing a pending reset. - pub(crate) async fn reset_worker_dp_rank_and_wait( - &self, - worker_id: WorkerId, - dp_rank: DpRank, - ) -> Result<(), KvRouterError> { - match self { - Self::KvIndexer { - primary, - lower_tier, - approx, - .. - } => { - primary - .reset_worker_dp_rank_and_wait(worker_id, dp_rank) - .await?; - for indexer in lower_tier.all() { - indexer - .reset_worker_dp_rank_and_wait(worker_id, dp_rank) - .await?; - } - if let Some(approx) = approx { - approx - .reset_worker_dp_rank_and_wait(worker_id, dp_rank) - .await?; - } - } - Self::Concurrent { - primary, - lower_tier, - approx, - .. - } => { - primary - .reset_worker_dp_rank_and_wait(worker_id, dp_rank) - .await?; - for indexer in lower_tier.all() { - indexer - .reset_worker_dp_rank_and_wait(worker_id, dp_rank) - .await?; - } - if let Some(approx) = approx { - approx - .reset_worker_dp_rank_and_wait(worker_id, dp_rank) - .await?; - } - } - Self::Remote { approx, .. } => { - if let Some(approx) = approx { - approx - .reset_worker_dp_rank_and_wait(worker_id, dp_rank) - .await?; - } - } - Self::None => {} - } - Ok(()) - } + let approx = predict_on_route_side_indexer( + component, + kv_router_config, + block_size, + cancellation_token.child_token(), + ); - pub(crate) fn uses_approximate_lru(&self) -> bool { - match self { - Self::KvIndexer { primary, .. } => primary.approximate_lru_enabled(), - Self::Concurrent { primary, .. } => primary.approximate_lru_enabled(), - Self::Remote { .. } | Self::None => false, - } + if kv_router_config.router_event_threads > 1 { + let kv_indexer_metrics = KvIndexerMetrics::from_component(component); + return Ok(Indexer::Concurrent { + primary: Arc::new(ThreadPoolIndexer::new_with_metrics( + ConcurrentRadixTreeCompressed::new(), + kv_router_config.router_event_threads as usize, + block_size, + Some(kv_indexer_metrics.clone()), + )), + lower_tier: LowerTierIndexers::new_with_metrics( + kv_router_config.router_event_threads as usize, + block_size, + Some(kv_indexer_metrics), + ), + approx, + primary_records_routing_decisions: false, + }); } - pub(crate) fn set_approximate_lru_capacity_now( - &self, - worker: dynamo_kv_router::protocols::WorkerWithDpRank, - incarnation: ApproximateLruIncarnation, - capacity: Option, - ) -> Result<(), KvRouterError> { - match self { - Self::KvIndexer { primary, .. } => { - primary.set_approximate_lru_capacity_now(worker, incarnation, capacity) - } - Self::Concurrent { primary, .. } => { - primary.set_approximate_lru_capacity_now(worker, incarnation, capacity) - } - Self::Remote { .. } | Self::None => Ok(()), - } - } + let kv_indexer_metrics = KvIndexerMetrics::from_component(component); + Ok(Indexer::Single { + primary: KvIndexer::new_with_pruning( + cancellation_token.child_token(), + block_size, + kv_indexer_metrics.clone(), + None, + ), + lower_tier: LowerTierIndexers::new_with_metrics(1, block_size, Some(kv_indexer_metrics)), + approx, + primary_records_routing_decisions: false, + }) +} - pub(crate) async fn approximate_lru_stats(&self) -> Result { - match self { - Self::KvIndexer { primary, .. } => primary.approximate_lru_stats().await, - Self::Concurrent { primary, .. } => primary.approximate_lru_stats().await, - Self::Remote { .. } | Self::None => Ok(ApproximateLruStats::default()), - } - } +/// Predict-on-route side indexer for `router_predicted_ttl_secs`, with the +/// component's indexer metrics. +fn predict_on_route_side_indexer( + component: &Component, + kv_router_config: &KvRouterConfig, + block_size: u32, + cancellation_token: CancellationToken, +) -> Option { + let ttl_secs = kv_router_config.router_predicted_ttl_secs?; + tracing::info!( + ttl_secs, + "Starting predict-on-route side indexer (short-TTL approximate)" + ); + Some(SideIndexer::new( + Duration::from_secs_f64(ttl_secs), + block_size, + kv_router_config.router_event_threads as usize, + KvIndexerMetrics::from_component(component), + cancellation_token, + )) } #[cfg(test)] @@ -569,767 +306,3 @@ pub(super) mod test_util { ) } } - -#[cfg(test)] -mod tests { - use std::sync::Arc; - use std::time::Duration; - - use tokio_util::sync::CancellationToken; - - use super::test_util::store_event; - use super::{Indexer, LowerTierIndexers}; - use dynamo_kv_router::{ - ConcurrentRadixTreeCompressed, ThreadPoolIndexer, - approx::PruneConfig, - indexer::{KvIndexer, KvIndexerInterface, KvIndexerMetrics, RoutingDecisionHashes}, - protocols::{ - BlockHashOptions, LocalBlockHash, StorageTier, TokensWithHashes, WorkerWithDpRank, - compute_block_hash_for_seq, compute_seq_hash_for_block, - }, - }; - - fn make_test_indexer() -> Indexer { - Indexer::KvIndexer { - primary: KvIndexer::new( - CancellationToken::new(), - 4, - Arc::new(KvIndexerMetrics::new_unregistered()), - ), - lower_tier: LowerTierIndexers::new(1, 4), - approx: None, - primary_records_routing_decisions: false, - } - } - - fn make_test_concurrent_indexer() -> Indexer { - Indexer::Concurrent { - primary: Arc::new(ThreadPoolIndexer::new( - ConcurrentRadixTreeCompressed::new(), - 2, - 4, - )), - lower_tier: LowerTierIndexers::new(2, 4), - approx: None, - primary_records_routing_decisions: false, - } - } - - fn make_test_concurrent_approx_indexer() -> Indexer { - Indexer::Concurrent { - primary: Arc::new(ThreadPoolIndexer::new_with_pruning( - ConcurrentRadixTreeCompressed::new(), - 2, - 4, - PruneConfig { - ttl: Duration::from_secs(60), - }, - )), - lower_tier: LowerTierIndexers::new(2, 4), - approx: None, - primary_records_routing_decisions: true, - } - } - - #[test] - fn overlap_refresh_is_limited_to_local_indexers() { - assert!(make_test_indexer().supports_overlap_refresh()); - assert!(make_test_concurrent_indexer().supports_overlap_refresh()); - assert!(!Indexer::None.supports_overlap_refresh()); - } - - #[test] - fn kv_transfer_chain_retention_requires_event_driven_primary() { - assert!(make_test_indexer().supports_kv_transfer_chain_retention()); - assert!(make_test_concurrent_indexer().supports_kv_transfer_chain_retention()); - assert!(!make_test_concurrent_approx_indexer().supports_kv_transfer_chain_retention()); - assert!(!Indexer::None.supports_kv_transfer_chain_retention()); - } - - async fn flush_indexer(indexer: &Indexer) { - match indexer { - Indexer::KvIndexer { - primary, - lower_tier, - .. - } => { - let _ = primary.flush().await; - for indexer in lower_tier.all() { - let _ = indexer.dump_events().await.unwrap(); - } - } - Indexer::Concurrent { - primary, - lower_tier, - .. - } => { - primary.flush().await; - for indexer in lower_tier.all() { - let _ = indexer.dump_events().await.unwrap(); - } - } - Indexer::Remote { .. } | Indexer::None => {} - } - } - - async fn assert_rank_reset_is_acknowledged(indexer: Indexer) { - let reset_rank = WorkerWithDpRank::new(7, 0); - let retained_rank = WorkerWithDpRank::new(7, 1); - - for dp_rank in [reset_rank.dp_rank, retained_rank.dp_rank] { - indexer - .apply_event(store_event(7, dp_rank, 1, &[], &[41], StorageTier::Device)) - .await; - indexer - .apply_event(store_event( - 7, - dp_rank, - 2, - &[41], - &[42], - StorageTier::HostPinned, - )) - .await; - } - flush_indexer(&indexer).await; - - indexer - .reset_worker_dp_rank_and_wait(reset_rank.worker_id, reset_rank.dp_rank) - .await - .unwrap(); - - let matches = indexer - .find_matches_by_tier(vec![LocalBlockHash(41), LocalBlockHash(42)]) - .await - .unwrap(); - assert!( - !matches - .device - .overlap_scores - .scores - .contains_key(&reset_rank) - ); - assert_eq!( - matches.device.overlap_scores.scores.get(&retained_rank), - Some(&1) - ); - let host_hits = &matches - .lower_tier - .get(&StorageTier::HostPinned) - .unwrap() - .hits; - assert!(!host_hits.contains_key(&reset_rank)); - assert_eq!(host_hits.get(&retained_rank), Some(&1)); - } - - #[tokio::test] - async fn single_thread_rank_reset_waits_for_all_local_tiers() { - assert_rank_reset_is_acknowledged(make_test_indexer()).await; - } - - #[tokio::test] - async fn concurrent_rank_reset_waits_for_all_local_tiers() { - assert_rank_reset_is_acknowledged(make_test_concurrent_indexer()).await; - } - - #[tokio::test] - async fn tiered_query_chains_device_host_and_disk() { - let indexer = make_test_indexer(); - let worker = WorkerWithDpRank::new(7, 0); - - indexer - .apply_event(store_event(7, 0, 1, &[], &[11, 12], StorageTier::Device)) - .await; - indexer - .apply_event(store_event( - 7, - 0, - 2, - &[11, 12], - &[13], - StorageTier::HostPinned, - )) - .await; - indexer - .apply_event(store_event( - 7, - 0, - 3, - &[11, 12, 13], - &[14], - StorageTier::Disk, - )) - .await; - flush_indexer(&indexer).await; - - let matches = indexer - .find_matches_by_tier(vec![ - LocalBlockHash(11), - LocalBlockHash(12), - LocalBlockHash(13), - LocalBlockHash(14), - ]) - .await - .unwrap(); - - assert_eq!(matches.device.overlap_scores.scores.get(&worker), Some(&2)); - assert_eq!( - matches - .lower_tier - .get(&StorageTier::HostPinned) - .and_then(|tier| tier.hits.get(&worker)), - Some(&1) - ); - assert_eq!( - matches - .lower_tier - .get(&StorageTier::Disk) - .and_then(|tier| tier.hits.get(&worker)), - Some(&1) - ); - } - - #[tokio::test] - async fn router_dump_includes_all_allocated_physical_tiers() { - let indexer = make_test_indexer(); - for (event_id, tier, block) in [ - (1, StorageTier::Device, 11), - (2, StorageTier::HostPinned, 12), - (3, StorageTier::Disk, 13), - ] { - indexer - .apply_event(store_event(7, 0, event_id, &[], &[block], tier)) - .await; - } - flush_indexer(&indexer).await; - - let events = indexer.dump_events().await.unwrap(); - for tier in [ - StorageTier::Device, - StorageTier::HostPinned, - StorageTier::Disk, - ] { - assert!(events.iter().any(|event| event.storage_tier == tier)); - } - } - - #[tokio::test] - async fn tiered_query_seeds_lower_tier_only_workers_without_affecting_device_scores() { - let indexer = make_test_indexer(); - let device_worker = WorkerWithDpRank::new(10, 0); - let host_only_worker = WorkerWithDpRank::new(20, 0); - let disk_only_worker = WorkerWithDpRank::new(30, 0); - - indexer - .apply_event(store_event(10, 0, 1, &[], &[21], StorageTier::Device)) - .await; - indexer - .apply_event(store_event(20, 0, 2, &[], &[21], StorageTier::HostPinned)) - .await; - indexer - .apply_event(store_event(30, 0, 3, &[], &[21], StorageTier::Disk)) - .await; - flush_indexer(&indexer).await; - - let matches = indexer - .find_matches_by_tier(vec![LocalBlockHash(21)]) - .await - .unwrap(); - - assert_eq!( - matches.device.overlap_scores.scores.get(&device_worker), - Some(&1) - ); - assert!( - !matches - .device - .overlap_scores - .scores - .contains_key(&host_only_worker) - ); - assert!( - !matches - .device - .overlap_scores - .scores - .contains_key(&disk_only_worker) - ); - - assert_eq!( - matches - .lower_tier - .get(&StorageTier::HostPinned) - .and_then(|tier| tier.hits.get(&host_only_worker)), - Some(&1) - ); - assert_eq!( - matches - .lower_tier - .get(&StorageTier::Disk) - .and_then(|tier| tier.hits.get(&disk_only_worker)), - Some(&1) - ); - } - - #[tokio::test] - async fn tiered_query_only_seeds_matching_root_workers() { - let indexer = make_test_indexer(); - let matching_host_worker = WorkerWithDpRank::new(20, 0); - let nonmatching_host_worker = WorkerWithDpRank::new(21, 0); - - indexer - .apply_event(store_event(20, 0, 1, &[], &[31], StorageTier::HostPinned)) - .await; - indexer - .apply_event(store_event(21, 0, 2, &[], &[32], StorageTier::HostPinned)) - .await; - flush_indexer(&indexer).await; - - let matches = indexer - .find_matches_by_tier(vec![LocalBlockHash(31)]) - .await - .unwrap(); - - assert_eq!( - matches - .lower_tier - .get(&StorageTier::HostPinned) - .and_then(|tier| tier.hits.get(&matching_host_worker)), - Some(&1) - ); - assert!( - !matches - .lower_tier - .get(&StorageTier::HostPinned) - .is_some_and(|tier| tier.hits.contains_key(&nonmatching_host_worker)) - ); - } - - #[tokio::test] - async fn concurrent_tiered_query_chains_device_and_lower_tier_matches() { - let indexer = make_test_concurrent_indexer(); - let worker = WorkerWithDpRank::new(7, 0); - - indexer - .apply_event(store_event(7, 0, 1, &[], &[11, 12], StorageTier::Device)) - .await; - indexer - .apply_event(store_event( - 7, - 0, - 2, - &[11, 12], - &[13], - StorageTier::HostPinned, - )) - .await; - flush_indexer(&indexer).await; - - let matches = indexer - .find_matches_by_tier(vec![ - LocalBlockHash(11), - LocalBlockHash(12), - LocalBlockHash(13), - ]) - .await - .unwrap(); - - assert_eq!(matches.device.overlap_scores.scores.get(&worker), Some(&2)); - assert_eq!( - matches - .lower_tier - .get(&StorageTier::HostPinned) - .and_then(|tier| tier.hits.get(&worker)), - Some(&1) - ); - } - - #[tokio::test] - async fn concurrent_records_hashed_routing_decision() { - let indexer = make_test_concurrent_approx_indexer(); - assert!(indexer.records_routing_decisions()); - - let worker = WorkerWithDpRank::new(7, 0); - let tokens = vec![1, 2, 3, 4]; - let block_hashes = compute_block_hash_for_seq(&tokens, 4, BlockHashOptions::default()); - let sequence_hashes = compute_seq_hash_for_block(&block_hashes); - - indexer - .record_hashed_routing_decision(worker, block_hashes.clone(), sequence_hashes) - .await - .unwrap(); - flush_indexer(&indexer).await; - - let matches = indexer.find_matches_by_tier(block_hashes).await.unwrap(); - assert_eq!(matches.device.overlap_scores.scores.get(&worker), Some(&1)); - } - - #[tokio::test] - async fn concurrent_records_precomputed_routing_hashes() { - let indexer = make_test_concurrent_approx_indexer(); - assert!(indexer.records_routing_decisions()); - - let worker = WorkerWithDpRank::new(7, 0); - let local_hashes = vec![LocalBlockHash(91), LocalBlockHash(92)]; - let sequence_hashes = compute_seq_hash_for_block(&local_hashes); - indexer - .record_routing_decision_hashes( - worker, - RoutingDecisionHashes { - local_hashes: local_hashes.clone(), - sequence_hashes, - }, - ) - .await - .unwrap(); - flush_indexer(&indexer).await; - - let matches = indexer.find_matches_by_tier(local_hashes).await.unwrap(); - assert_eq!(matches.device.overlap_scores.scores.get(&worker), Some(&2)); - } - - #[tokio::test] - async fn event_driven_primary_without_side_skips_route_recording() { - let indexer = make_test_indexer(); - assert!(!indexer.records_routing_decisions()); - - let worker = WorkerWithDpRank::new(7, 0); - let tokens = vec![1, 2, 3, 4]; - let block_hashes = compute_block_hash_for_seq(&tokens, 4, BlockHashOptions::default()); - let mut tokens_with_hashes = TokensWithHashes::new(tokens, 4); - - indexer - .process_routing_decision_for_request(&mut tokens_with_hashes, worker) - .await - .unwrap(); - flush_indexer(&indexer).await; - - let matches = indexer.find_matches_by_tier(block_hashes).await.unwrap(); - assert!( - !matches.device.overlap_scores.scores.contains_key(&worker), - "event-driven primary without side overlay should rely on KV events, not route-time writes" - ); - } - - #[tokio::test] - async fn side_only_worker_scored_but_not_used_as_lower_tier_anchor() { - // Build an Indexer::Concurrent with a real side indexer so - // `record_hashed_routing_decision` populates only the side path. - let primary = Arc::new(ThreadPoolIndexer::new( - ConcurrentRadixTreeCompressed::new(), - 2, - 4, - )); - // PruneConfig is required to enable routing-decision recording on the - // side indexer; without it the routing-decision path is a no-op. - let side = Arc::new(ThreadPoolIndexer::new_with_pruning( - ConcurrentRadixTreeCompressed::new(), - 1, - 4, - PruneConfig { - ttl: Duration::from_secs(60), - }, - )); - let side_for_flush = side.clone(); - let indexer = Indexer::Concurrent { - primary, - lower_tier: LowerTierIndexers::new(2, 4), - approx: Some(super::SideIndexer::Concurrent(side)), - primary_records_routing_decisions: false, - }; - assert!(indexer.records_routing_decisions()); - - let primary_worker = WorkerWithDpRank::new(10, 0); - let side_only_worker = WorkerWithDpRank::new(20, 0); - - // Primary sees blocks [11, 12, 13] on Device for primary_worker; - // extension block [14] on HostPinned for primary_worker. - indexer - .apply_event(store_event( - 10, - 0, - 1, - &[], - &[11, 12, 13], - StorageTier::Device, - )) - .await; - indexer - .apply_event(store_event( - 10, - 0, - 2, - &[11, 12, 13], - &[14], - StorageTier::HostPinned, - )) - .await; - // Crucially, also give side_only_worker a HostPinned extension at - // block 14 anchored on the same prefix [11, 12, 13]. If the lower - // tier were seeded from the side-merged device score, the host walk - // would find this and credit a hit; with the reorder it should not. - indexer - .apply_event(store_event( - 20, - 0, - 3, - &[11, 12, 13], - &[14], - StorageTier::HostPinned, - )) - .await; - - // Side-only: route a decision so the side indexer learns - // side_only_worker for the same device prefix. Primary never sees it. - let block_hashes: Vec = - [11, 12, 13].iter().copied().map(LocalBlockHash).collect(); - let sequence_hashes = compute_seq_hash_for_block(&block_hashes); - indexer - .record_hashed_routing_decision(side_only_worker, block_hashes.clone(), sequence_hashes) - .await - .unwrap(); - - flush_indexer(&indexer).await; - side_for_flush.flush().await; - - let matches = indexer - .find_matches_by_tier(vec![ - LocalBlockHash(11), - LocalBlockHash(12), - LocalBlockHash(13), - LocalBlockHash(14), - ]) - .await - .unwrap(); - - // Merge worked: both workers carry device scores. - assert_eq!( - matches - .device - .overlap_scores - .scores - .get(&primary_worker) - .copied(), - Some(3) - ); - assert_eq!( - matches - .device - .overlap_scores - .scores - .get(&side_only_worker) - .copied(), - Some(3), - "side-only worker should appear in merged device scores" - ); - - // Reorder enforced: lower-tier was seeded from primary only. - // primary_worker still extends into HostPinned via its own device - // anchor. side_only_worker's HostPinned extension exists in the - // host tier, but because the side score wasn't used as a device - // anchor, the host walk does not start for it and its host hit is - // not credited. - let host = matches - .lower_tier - .get(&StorageTier::HostPinned) - .expect("host-pinned tier should have been allocated"); - assert_eq!(host.hits.get(&primary_worker).copied(), Some(1)); - assert_eq!( - host.hits.get(&side_only_worker).copied().unwrap_or(0), - 0, - "side-only worker's host extension must not be credited \ - when lower-tier seeding is primary-only" - ); - assert!( - !host.next_continuations.contains_key(&side_only_worker), - "side-only worker must not appear in lower-tier continuations" - ); - } - - #[tokio::test] - async fn borrowed_tiered_lookup_matches_owned_with_lower_tier_and_side_overlay() { - let primary = Arc::new(ThreadPoolIndexer::new( - ConcurrentRadixTreeCompressed::new(), - 2, - 4, - )); - let side = Arc::new(ThreadPoolIndexer::new_with_pruning( - ConcurrentRadixTreeCompressed::new(), - 1, - 4, - PruneConfig { - ttl: Duration::from_secs(60), - }, - )); - let side_for_flush = side.clone(); - let indexer = Indexer::Concurrent { - primary, - lower_tier: LowerTierIndexers::new(2, 4), - approx: Some(super::SideIndexer::Concurrent(side)), - primary_records_routing_decisions: false, - }; - - let primary_worker = WorkerWithDpRank::new(10, 0); - let side_worker = WorkerWithDpRank::new(20, 0); - indexer - .apply_event(store_event(10, 0, 1, &[], &[11, 12], StorageTier::Device)) - .await; - indexer - .apply_event(store_event( - 10, - 0, - 2, - &[11, 12], - &[13], - StorageTier::HostPinned, - )) - .await; - - let side_hashes = vec![LocalBlockHash(11), LocalBlockHash(12), LocalBlockHash(13)]; - indexer - .record_routing_decision_hashes( - side_worker, - RoutingDecisionHashes { - local_hashes: side_hashes.clone(), - sequence_hashes: compute_seq_hash_for_block(&side_hashes), - }, - ) - .await - .unwrap(); - flush_indexer(&indexer).await; - side_for_flush.flush().await; - - let query = vec![LocalBlockHash(11), LocalBlockHash(12), LocalBlockHash(13)]; - let borrowed = indexer.find_matches_by_tier_ref(&query).await.unwrap(); - let owned = indexer.find_matches_by_tier(query).await.unwrap(); - - assert_eq!( - borrowed.device.overlap_scores.scores, - owned.device.overlap_scores.scores - ); - assert_eq!( - borrowed - .lower_tier - .get(&StorageTier::HostPinned) - .map(|tier| &tier.hits), - owned - .lower_tier - .get(&StorageTier::HostPinned) - .map(|tier| &tier.hits) - ); - assert_eq!( - borrowed - .device - .overlap_scores - .scores - .get(&primary_worker) - .copied(), - Some(2) - ); - assert_eq!( - borrowed - .device - .overlap_scores - .scores - .get(&side_worker) - .copied(), - Some(3) - ); - } - - #[tokio::test] - async fn concurrent_tiered_query_seeds_lower_tier_only_workers_without_affecting_device_scores() - { - let indexer = make_test_concurrent_indexer(); - let device_worker = WorkerWithDpRank::new(10, 0); - let host_only_worker = WorkerWithDpRank::new(20, 0); - let disk_only_worker = WorkerWithDpRank::new(30, 0); - - indexer - .apply_event(store_event(10, 0, 1, &[], &[21], StorageTier::Device)) - .await; - indexer - .apply_event(store_event(20, 0, 2, &[], &[21], StorageTier::HostPinned)) - .await; - indexer - .apply_event(store_event(30, 0, 3, &[], &[21], StorageTier::Disk)) - .await; - flush_indexer(&indexer).await; - - let matches = indexer - .find_matches_by_tier(vec![LocalBlockHash(21)]) - .await - .unwrap(); - - assert_eq!( - matches.device.overlap_scores.scores.get(&device_worker), - Some(&1) - ); - assert!( - !matches - .device - .overlap_scores - .scores - .contains_key(&host_only_worker) - ); - assert!( - !matches - .device - .overlap_scores - .scores - .contains_key(&disk_only_worker) - ); - - assert_eq!( - matches - .lower_tier - .get(&StorageTier::HostPinned) - .and_then(|tier| tier.hits.get(&host_only_worker)), - Some(&1) - ); - assert_eq!( - matches - .lower_tier - .get(&StorageTier::Disk) - .and_then(|tier| tier.hits.get(&disk_only_worker)), - Some(&1) - ); - } - - /// Regression test: when a worker has blocks in both device and lower-tier - /// storage (e.g. same prefix stored on GPU and offloaded to host), the - /// Concurrent indexer doesn't return last_matched_hashes. Without the fix, - /// query_lower_tiers would re-query that worker from root in the lower tier, - /// double-counting overlap blocks and producing cached_tokens > ISL. - #[tokio::test] - async fn concurrent_tiered_query_does_not_double_count_device_and_lower_tier_overlap() { - let indexer = make_test_concurrent_indexer(); - let worker = WorkerWithDpRank::new(7, 0); - - // Device owns the prefix block; host-pinned extends it by one block. - indexer - .apply_event(store_event(7, 0, 1, &[], &[41], StorageTier::Device)) - .await; - indexer - .apply_event(store_event(7, 0, 2, &[41], &[42], StorageTier::HostPinned)) - .await; - flush_indexer(&indexer).await; - - let matches = indexer - .find_matches_by_tier(vec![LocalBlockHash(41), LocalBlockHash(42)]) - .await - .unwrap(); - - assert_eq!(matches.device.overlap_scores.scores.get(&worker), Some(&1)); - - let host_hits = matches - .lower_tier - .get(&StorageTier::HostPinned) - .and_then(|tier| tier.hits.get(&worker).copied()) - .unwrap_or(0); - assert_eq!( - host_hits, 1, - "lower-tier should extend the device prefix without double-counting it" - ); - } -} diff --git a/lib/llm/src/kv_router/indexer/recovery/state_agent.rs b/lib/llm/src/kv_router/indexer/recovery/state_agent.rs index a825183c3c26..fd495617c663 100644 --- a/lib/llm/src/kv_router/indexer/recovery/state_agent.rs +++ b/lib/llm/src/kv_router/indexer/recovery/state_agent.rs @@ -1855,7 +1855,7 @@ mod tests { ); ( primary.clone(), - Indexer::KvIndexer { + Indexer::Single { primary, lower_tier: LowerTierIndexers::new(1, 4), approx: None, @@ -2091,7 +2091,7 @@ mod tests { Arc::new(KvIndexerMetrics::new_unregistered()), ); let lower_tiers = LowerTierIndexers::new(1, 4); - let indexer = Indexer::KvIndexer { + let indexer = Indexer::Single { primary, lower_tier: lower_tiers.clone(), approx: None, diff --git a/lib/llm/src/kv_router/indexer/recovery/worker_query.rs b/lib/llm/src/kv_router/indexer/recovery/worker_query.rs index cfdb4a712691..6d19039c28dd 100644 --- a/lib/llm/src/kv_router/indexer/recovery/worker_query.rs +++ b/lib/llm/src/kv_router/indexer/recovery/worker_query.rs @@ -1371,7 +1371,7 @@ mod tests { let indexer = KvIndexer::new(CancellationToken::new(), 4, metrics); ( indexer.clone(), - Indexer::KvIndexer { + Indexer::Single { primary: indexer, lower_tier: LowerTierIndexers::new(1, 4), approx: None, @@ -1835,7 +1835,7 @@ mod tests { let (_tx, rx) = watch::channel(initial.clone()); let (_primary, indexer) = indexer(); let lower_tier = match &indexer { - Indexer::KvIndexer { lower_tier, .. } => lower_tier.clone(), + Indexer::Single { lower_tier, .. } => lower_tier.clone(), _ => unreachable!(), }; let client = diff --git a/lib/llm/src/kv_router/indexer/remote.rs b/lib/llm/src/kv_router/indexer/remote.rs index 59e1c2e31e84..9a6475c4b6c4 100644 --- a/lib/llm/src/kv_router/indexer/remote.rs +++ b/lib/llm/src/kv_router/indexer/remote.rs @@ -198,8 +198,28 @@ impl RemoteIndexer { Ok(()) } +} - pub(super) fn use_kv_events(&self) -> bool { +#[async_trait::async_trait] +impl dynamo_kv_router::services::indexer::backend::RemotePrimary for RemoteIndexer { + async fn find_matches_by_tier( + &self, + block_hashes: Vec, + device_only: bool, + ) -> Result { + RemoteIndexer::find_matches_by_tier(self, block_hashes, device_only).await + } + + async fn record_routing_decision( + &self, + worker: WorkerWithDpRank, + hashes: dynamo_kv_router::indexer::RoutingDecisionHashes, + ) -> Result<()> { + self.record_hashed_routing_decision(worker, hashes.local_hashes, hashes.sequence_hashes) + .await + } + + fn use_kv_events(&self) -> bool { self.use_kv_events } } @@ -577,14 +597,14 @@ mod tests { .unwrap(); let _ = side.flush().await; - let indexer = Indexer::KvIndexer { + let indexer = Indexer::Single { primary: KvIndexer::new( CancellationToken::new(), 4, Arc::new(KvIndexerMetrics::new_unregistered()), ), lower_tier: LowerTierIndexers::new(1, 4), - approx: Some(SideIndexer::KvIndexer(side)), + approx: Some(SideIndexer::Single(side)), primary_records_routing_decisions: false, }; @@ -629,7 +649,7 @@ mod tests { #[tokio::test] async fn query_engine_returns_tiered_scores_with_lower_tier() { let worker = WorkerWithDpRank::new(7, 0); - let indexer = Indexer::KvIndexer { + let indexer = Indexer::Single { primary: KvIndexer::new( CancellationToken::new(), 4, @@ -642,10 +662,11 @@ mod tests { // Worker owns [11, 12] on device and [11, 12, 13] on host-pinned. indexer - .apply_event(store_event(7, 0, 1, &[], &[11, 12], StorageTier::Device)) - .await; + .try_apply_event(store_event(7, 0, 1, &[], &[11, 12], StorageTier::Device)) + .await + .unwrap(); indexer - .apply_event(store_event( + .try_apply_event(store_event( 7, 0, 2, @@ -653,9 +674,10 @@ mod tests { &[13], StorageTier::HostPinned, )) - .await; + .await + .unwrap(); - let Indexer::KvIndexer { + let Indexer::Single { primary, lower_tier, .. @@ -728,7 +750,7 @@ mod tests { #[tokio::test] async fn query_engine_device_only_skips_lower_tiers() { let worker = WorkerWithDpRank::new(7, 0); - let indexer = Indexer::KvIndexer { + let indexer = Indexer::Single { primary: KvIndexer::new( CancellationToken::new(), 4, @@ -740,10 +762,11 @@ mod tests { }; indexer - .apply_event(store_event(7, 0, 1, &[], &[11, 12], StorageTier::Device)) - .await; + .try_apply_event(store_event(7, 0, 1, &[], &[11, 12], StorageTier::Device)) + .await + .unwrap(); indexer - .apply_event(store_event( + .try_apply_event(store_event( 7, 0, 2, @@ -751,9 +774,10 @@ mod tests { &[13], StorageTier::HostPinned, )) - .await; + .await + .unwrap(); - let Indexer::KvIndexer { + let Indexer::Single { primary, lower_tier, .. diff --git a/lib/llm/src/kv_router/indexer/side.rs b/lib/llm/src/kv_router/indexer/side.rs deleted file mode 100644 index c4e2218048d3..000000000000 --- a/lib/llm/src/kv_router/indexer/side.rs +++ /dev/null @@ -1,125 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -use std::{sync::Arc, time::Duration}; - -use anyhow::Result; -use dynamo_kv_router::{ - ConcurrentRadixTreeCompressed, - approx::PruneConfig, - config::KvRouterConfig, - indexer::{ - KvIndexer, KvIndexerInterface, KvIndexerMetrics, KvRouterError, RoutingDecisionHashes, - SyncIndexer, ThreadPoolIndexer, - }, - protocols::{DpRank, OverlapScores, WorkerId, WorkerWithDpRank}, -}; -use dynamo_runtime::component::Component; -use tokio_util::sync::CancellationToken; - -use super::lookup::HashInput; - -#[derive(Clone)] -pub enum SideIndexer { - KvIndexer(KvIndexer), - Concurrent(Arc>), -} - -impl SideIndexer { - pub(super) fn new_predict_on_route( - component: &Component, - kv_router_config: &KvRouterConfig, - block_size: u32, - cancellation_token: CancellationToken, - ) -> Option { - let ttl_secs = kv_router_config.router_predicted_ttl_secs?; - let prune_config = Some(PruneConfig { - ttl: Duration::from_secs_f64(ttl_secs), - }); - let metrics = KvIndexerMetrics::from_component(component); - tracing::info!( - ttl_secs, - "Starting predict-on-route side indexer (short-TTL approximate)" - ); - if kv_router_config.router_event_threads > 1 { - return Some(Self::Concurrent(Arc::new( - ThreadPoolIndexer::new_with_metrics_and_pruning( - ConcurrentRadixTreeCompressed::new(), - kv_router_config.router_event_threads as usize, - block_size, - Some(metrics), - prune_config, - ), - ))); - } - - Some(Self::KvIndexer(KvIndexer::new_with_pruning( - cancellation_token, - block_size, - metrics, - prune_config, - ))) - } - - pub(super) async fn find_matches_input( - &self, - sequence: HashInput<'_>, - ) -> Result { - match self { - Self::KvIndexer(indexer) => { - indexer - .find_matches(sequence.into_owned_at_boundary()) - .await - } - Self::Concurrent(indexer) => { - Ok(indexer.backend().find_matches(sequence.as_slice(), false)) - } - } - } - - pub(super) async fn process_routing_decision_hashes( - &self, - worker: WorkerWithDpRank, - hashes: RoutingDecisionHashes, - ) -> Result<(), KvRouterError> { - match self { - Self::KvIndexer(indexer) => { - indexer - .process_routing_decision_with_hashes( - worker, - hashes.local_hashes, - hashes.sequence_hashes, - ) - .await - } - Self::Concurrent(indexer) => { - indexer - .process_routing_decision_hash_slices( - worker, - &hashes.local_hashes, - &hashes.sequence_hashes, - ) - .await - } - } - } - - pub(super) async fn reset_worker_dp_rank_and_wait( - &self, - worker_id: WorkerId, - dp_rank: DpRank, - ) -> Result<(), KvRouterError> { - match self { - Self::KvIndexer(indexer) => { - indexer - .reset_worker_dp_rank_and_wait(worker_id, dp_rank) - .await - } - Self::Concurrent(indexer) => { - indexer - .reset_worker_dp_rank_and_wait(worker_id, dp_rank) - .await - } - } - } -} From d17aa47dbf761489340ab7aaa30d2aa6a4db6657 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Thu, 3 Sep 2026 23:50:21 +0000 Subject: [PATCH 026/102] refactor(kv-router): the frontend's index is the partition's index KvIndexSource::External { refresh } gave the core only a refresh handle onto an index the host kept elsewhere. Replace it with Provided(Arc): the host builds and feeds the index, and the core holds it as every partition's index, reading it for selection and dequeue-time refresh (when the index supports refresh) and never writing to it. The frontend's embedded partition now shares the KvRouter's Indexer object instead of a TieredMatchProvider view of it. RefreshProvider loses its External variant. Signed-off-by: Ishan Dhanani --- .../src/services/selection/core/mod.rs | 60 +++++++++---------- lib/llm/src/kv_router.rs | 6 +- lib/llm/src/kv_router/embedded.rs | 11 ++-- 3 files changed, 32 insertions(+), 45 deletions(-) diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index 3df89340fa6a..ec1327477363 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -58,12 +58,10 @@ use crate::WorkerSelectionPolicyFactory; use crate::WorkerType; /// Source of dequeue-time overlap refreshes for a partition scheduler: the -/// partition's own indexer, or one the embedding host supplies when it keeps -/// its own index (the frontend `KvRouter` in embedded mode). +/// partition's index. #[derive(Clone)] pub enum RefreshProvider { Local(Arc), - External(Arc), } #[async_trait::async_trait] @@ -74,7 +72,6 @@ impl TieredMatchProvider for RefreshProvider { ) -> Result { match self { Self::Local(indexer) => indexer.find_tiered_matches(sequence.to_vec()).await, - Self::External(provider) => provider.find_tiered_matches(sequence).await, } } @@ -89,11 +86,6 @@ impl TieredMatchProvider for RefreshProvider { .find_tiered_matches_with_options(sequence.to_vec(), options) .await } - Self::External(provider) => { - provider - .find_tiered_matches_with_options(sequence, options) - .await - } } } } @@ -215,19 +207,18 @@ pub struct HostCache { pub index: KvIndexSource, } -/// Where a partition's KV index lives and who feeds it. +/// Where a partition's KV index comes from and who feeds it. #[derive(Clone, Default)] pub enum KvIndexSource { - /// The core owns the index and subscribes to each worker rank's KV events - /// over ZMQ; workers need KV-event endpoints to be schedulable. + /// The core builds the index and subscribes to each worker rank's KV + /// events over ZMQ; workers need KV-event endpoints to be schedulable. #[default] Owned, - /// The host owns and feeds the index. Workers need no KV-event metadata. - /// `refresh` lets the core re-score queued requests against that index at - /// dequeue; `None` keeps their admission-time scores. - External { - refresh: Option>, - }, + /// The host built this index and feeds it (its own event ingress and + /// recovery). Every partition uses it; the core reads it for selection + /// and dequeue-time refresh but never writes to it, and workers need no + /// KV-event metadata. + Provided(Arc), /// A standalone indexer at this base URL serves the primary index; this /// core does not subscribe to worker KV events. Remote(String), @@ -237,9 +228,10 @@ impl std::fmt::Debug for KvIndexSource { fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { Self::Owned => formatter.write_str("Owned"), - Self::External { refresh } => formatter - .debug_struct("External") - .field("refresh", &refresh.is_some()) + Self::Provided(indexer) => formatter + .debug_struct("Provided") + .field("remote", &indexer.is_remote()) + .field("overlap_refresh", &indexer.supports_overlap_refresh()) .finish(), Self::Remote(url) => formatter.debug_tuple("Remote").field(url).finish(), } @@ -412,7 +404,7 @@ impl SelectionCore { )); let listens_for_kv_events = kv_router_config.use_kv_events && !indexer_policy.is_remote() - && !matches!(host.cache.index, KvIndexSource::External { .. }); + && !matches!(host.cache.index, KvIndexSource::Provided(_)); indexer_registry.set_indexer_policy(indexer_policy); if signal_indexer_ready { indexer_registry.signal_ready(); @@ -730,20 +722,19 @@ impl SelectionCore { self.cancel_token.child_token(), ); - let indexer = self - .indexer_registry - .get_or_create_indexer(key.clone(), block_size); - let refresh_provider = match &self.host.cache.index { - KvIndexSource::Owned | KvIndexSource::Remote(_) => { - Some(RefreshProvider::Local(Arc::new(indexer.clone()))) + let (indexer, refreshes) = match &self.host.cache.index { + KvIndexSource::Provided(indexer) => { + (Indexer::clone(indexer), indexer.supports_overlap_refresh()) } - KvIndexSource::External { refresh } => refresh - .as_ref() - .map(|p| RefreshProvider::External(Arc::clone(p))), + KvIndexSource::Owned | KvIndexSource::Remote(_) => ( + self.indexer_registry + .get_or_create_indexer(key.clone(), block_size), + true, + ), }; - let overlap_refresh = refresh_provider.map(|provider| { + let overlap_refresh = refreshes.then(|| { Arc::new(TieredOverlapRefresher::new( - provider, + RefreshProvider::Local(Arc::new(indexer.clone())), self.kv_router_config.clone(), block_size, )) @@ -843,6 +834,9 @@ impl SelectionCore { } async fn cleanup_indexer_registration(&self, record: &WorkerCatalogRecord) { + if matches!(self.host.cache.index, KvIndexSource::Provided(_)) { + return; + } if self.listens_for_kv_events { if let Err(error) = self .indexer_registry diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index 323b9b2d77be..a94f923e486f 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -820,10 +820,6 @@ impl KvRouter { let available_worker_provider: WorkerAvailabilityProvider = Arc::new(move || client_for_availability.available_instance_ids()); - let overlap_refresh: Option> = - indexer - .supports_overlap_refresh() - .then(|| Arc::new(indexer.clone()) as Arc<_>); let scheduler = embedded::EmbeddedSelection::start( embedded::EmbeddedSelectionArgs { kv_router_config: kv_router_config.clone(), @@ -835,7 +831,7 @@ impl KvRouter { prefill_load_estimator: prefill_load_estimator.clone(), overloaded_worker_provider, available_worker_provider, - overlap_refresh, + indexer: indexer.clone(), scheduler_load, endpoint: endpoint.clone(), router_id: endpoint.drt().discovery().instance_id(), diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs index 183f0b52049c..d393fe38d3e3 100644 --- a/lib/llm/src/kv_router/embedded.rs +++ b/lib/llm/src/kv_router/embedded.rs @@ -19,7 +19,6 @@ use anyhow::{Context, Result}; use dynamo_kv_router::WorkerType; use dynamo_kv_router::config::KvRouterConfig; use dynamo_kv_router::identity::RoutingPartitionId; -use dynamo_kv_router::indexer::TieredMatchProvider; use dynamo_kv_router::protocols::{WorkerConfigLike, WorkerId, WorkerWithDpRank}; use dynamo_kv_router::scheduling::queue::{SchedulerBookingCleanup, SchedulerBookingDescriptor}; use dynamo_kv_router::scheduling::{ @@ -51,9 +50,9 @@ pub(crate) struct EmbeddedSelectionArgs { pub prefill_load_estimator: Option>, pub overloaded_worker_provider: OverloadedWorkerProvider, pub available_worker_provider: WorkerAvailabilityProvider, - /// The router's own indexer, used for dequeue-time overlap refresh so - /// queued requests are re-scored against the index that scored them. - pub overlap_refresh: Option>, + /// The router's index, built and fed by the router; the partition reads + /// it for selection and dequeue-time refresh. + pub indexer: super::indexer::Indexer, /// Scheduler-owned load snapshots for the worker monitor's overload /// detection, the same feed the runtime scheduler publishes. pub scheduler_load: crate::kv_router::routing_load::SchedulerLoadSender, @@ -156,9 +155,7 @@ impl EmbeddedSelection { }, cache: HostCache { shared: None, - index: KvIndexSource::External { - refresh: args.overlap_refresh, - }, + index: KvIndexSource::Provided(Arc::new(args.indexer)), }, eligibility: HostEligibility::default(), telemetry: HostTelemetry { From e2633414670a10161d72144cdc1134aa238dfd44 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Fri, 4 Sep 2026 00:02:51 +0000 Subject: [PATCH 027/102] feat(kv-router): KvEventIngress feeds core-owned indexes KvIndexSource::Owned carries the ingress that feeds the core-built index. The trait decides what metadata a worker needs to be schedulable, attaches its ranks when the worker becomes schedulable, and detaches them when it leaves. ZmqDirectIngress is the existing standalone behavior moved out of SelectionCore: endpoint validation, one ZMQ listener per rank on the registry, replay endpoint, deregistration on delete. The core keeps the lifecycle decisions (when to attach, how to mark Incomplete) and no longer knows about ZMQ endpoints. Signed-off-by: Ishan Dhanani --- .../src/services/selection/README.md | 5 + .../src/services/selection/core/mod.rs | 89 +++++--------- .../src/services/selection/ingress.rs | 109 ++++++++++++++++++ lib/kv-router/src/services/selection/mod.rs | 2 + lib/kv-router/src/services/selection/types.rs | 18 +-- 5 files changed, 144 insertions(+), 79 deletions(-) create mode 100644 lib/kv-router/src/services/selection/ingress.rs diff --git a/lib/kv-router/src/services/selection/README.md b/lib/kv-router/src/services/selection/README.md index 3c81c0560552..d011b9fc61fa 100644 --- a/lib/kv-router/src/services/selection/README.md +++ b/lib/kv-router/src/services/selection/README.md @@ -16,6 +16,11 @@ and active-sequence accounting. Keep these implementation invariants explicit: `WorkerCatalogSource`; `CatalogReconciler` owns the desired-versus-actual diff (retry of not-yet-schedulable records, stale deletion). Hosts do not call `upsert_worker`/`delete_worker` for discovery-driven membership. +- A partition's KV index comes from `HostCache.index: KvIndexSource`: `Owned` + (the core builds it and a `KvEventIngress`, by default `ZmqDirectIngress`, + feeds it), `Provided` (the host built and feeds it; the core only reads), + or `Remote` (a standalone indexer serves it). The `Indexer` type itself is + shared with the frontend (`services::indexer::backend`). - Selector replicas synchronize admission, prefill-complete, and free events. - **NOTE:** Output-block updates remain local. They are deliberately excluded from replica sync because their frequency would consume disproportionate diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index ec1327477363..0bb8c7bc4d3e 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -46,6 +46,7 @@ use crate::tracking_hash::{TrackingHashContext, TrackingHashScope}; use super::catalog::WorkerCatalog; use super::error::SelectionError; +use super::ingress::{KvEventIngress, ZmqDirectIngress}; use super::input::{PromptRequest, TrackingHashInput}; use super::pending::{PendingSelection, SelectionCache, SelectionCacheConfig}; use super::types::{ @@ -208,12 +209,12 @@ pub struct HostCache { } /// Where a partition's KV index comes from and who feeds it. -#[derive(Clone, Default)] +#[derive(Clone)] pub enum KvIndexSource { - /// The core builds the index and subscribes to each worker rank's KV - /// events over ZMQ; workers need KV-event endpoints to be schedulable. - #[default] - Owned, + /// The core builds the index and the ingress feeds it with worker KV + /// events; the ingress decides what metadata a worker needs to be + /// schedulable. Defaults to [`ZmqDirectIngress`]. + Owned(Arc), /// The host built this index and feeds it (its own event ingress and /// recovery). Every partition uses it; the core reads it for selection /// and dequeue-time refresh but never writes to it, and workers need no @@ -224,10 +225,16 @@ pub enum KvIndexSource { Remote(String), } +impl Default for KvIndexSource { + fn default() -> Self { + Self::Owned(Arc::new(ZmqDirectIngress)) + } +} + impl std::fmt::Debug for KvIndexSource { fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { - Self::Owned => formatter.write_str("Owned"), + Self::Owned(_) => formatter.write_str("Owned"), Self::Provided(indexer) => formatter .debug_struct("Provided") .field("remote", &indexer.is_remote()) @@ -588,8 +595,10 @@ impl SelectionCore { .kv_router_config .queueing_enabled(Some(&record.model_name)) .map_err(|error| SelectionError::BadRequest(error.to_string()))?; - let reasons = - record.missing_schedulable_metadata(queueing_enabled, self.listens_for_kv_events); + let mut reasons = record.missing_schedulable_metadata(queueing_enabled); + if let Some(ingress) = self.ingress() { + reasons.extend(ingress.missing_metadata(&record)); + } if !reasons.is_empty() { let updated = self .catalog @@ -602,8 +611,8 @@ impl SelectionCore { if let Err(error) = self.ensure_entry(&record) { return self.mark_incomplete_after_reconcile_error(worker_id, record.key(), error); } - if self.listens_for_kv_events - && let Err(error) = self.register_indexer_listeners(&record).await + if let Some(ingress) = self.ingress() + && let Err(error) = ingress.attach(&self.indexer_registry, &record).await { self.cleanup_indexer_registration(&record).await; return self.mark_incomplete_after_reconcile_error(worker_id, record.key(), error); @@ -726,7 +735,7 @@ impl SelectionCore { KvIndexSource::Provided(indexer) => { (Indexer::clone(indexer), indexer.supports_overlap_refresh()) } - KvIndexSource::Owned | KvIndexSource::Remote(_) => ( + KvIndexSource::Owned(_) | KvIndexSource::Remote(_) => ( self.indexer_registry .get_or_create_indexer(key.clone(), block_size), true, @@ -791,64 +800,20 @@ impl SelectionCore { Ok(entry) } - async fn register_indexer_listeners( - &self, - record: &WorkerCatalogRecord, - ) -> Result<(), SelectionError> { - let block_size = record - .block_size - .ok_or_else(|| SelectionError::BadRequest("block_size is required".to_string()))?; - let mut endpoints: Vec<_> = record.listener_endpoints().into_iter().collect(); - endpoints.sort_by_key(|(dp_rank, _)| *dp_rank); - for (dp_rank, endpoint) in endpoints { - crate::services::common::zmq::validate_endpoint(&endpoint).map_err(|error| { - SelectionError::BadRequest(format!( - "invalid kv_events endpoint for worker {} dp_rank {dp_rank}: {error}", - record.worker_id - )) - })?; - if let Some(replay_endpoint) = record.replay_endpoint.as_deref() { - crate::services::common::zmq::validate_endpoint(replay_endpoint).map_err( - |error| { - SelectionError::BadRequest(format!( - "invalid replay endpoint for worker {} dp_rank {dp_rank}: {error}", - record.worker_id - )) - }, - )?; - } - self.indexer_registry - .register( - record.worker_id, - endpoint, - dp_rank, - record.model_name.clone(), - record.routing_group.clone(), - block_size, - record.replay_endpoint.clone(), - ) - .await - .map_err(|error| SelectionError::BadRequest(error.to_string()))?; + /// The ingress feeding core-owned indexes, when this core listens for KV events. + fn ingress(&self) -> Option<&dyn KvEventIngress> { + match &self.host.cache.index { + KvIndexSource::Owned(ingress) if self.listens_for_kv_events => Some(ingress.as_ref()), + _ => None, } - Ok(()) } async fn cleanup_indexer_registration(&self, record: &WorkerCatalogRecord) { if matches!(self.host.cache.index, KvIndexSource::Provided(_)) { return; } - if self.listens_for_kv_events { - if let Err(error) = self - .indexer_registry - .deregister(record.worker_id, &record.model_name, &record.routing_group) - .await - { - tracing::debug!( - worker_id = record.worker_id, - error = %error, - "indexer deregistration skipped or failed" - ); - } + if let Some(ingress) = self.ingress() { + ingress.detach(&self.indexer_registry, record).await; return; } diff --git a/lib/kv-router/src/services/selection/ingress.rs b/lib/kv-router/src/services/selection/ingress.rs new file mode 100644 index 000000000000..bb0ae8ab81dc --- /dev/null +++ b/lib/kv-router/src/services/selection/ingress.rs @@ -0,0 +1,109 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! How a core-owned KV index is fed with worker KV events. + +use async_trait::async_trait; + +use super::error::SelectionError; +use super::types::WorkerCatalogRecord; +use crate::services::indexer::registry::WorkerRegistry; + +/// Feeds a partition's core-owned index ([`super::KvIndexSource::Owned`]). +/// +/// The registry holds the partition indexes and the per-rank listener +/// lifecycle; an ingress attaches and detaches workers on it. A host that feeds +/// the index itself uses [`super::KvIndexSource::Provided`] instead. +#[async_trait] +pub trait KvEventIngress: Send + Sync { + /// Metadata `record` must carry before this ingress can feed its ranks. + /// A non-empty result keeps the worker `Incomplete`. + fn missing_metadata(&self, record: &WorkerCatalogRecord) -> Vec; + + /// Start feeding every rank of `record` into its partition's index. + async fn attach( + &self, + registry: &WorkerRegistry, + record: &WorkerCatalogRecord, + ) -> Result<(), SelectionError>; + + /// Stop feeding `record` and drop its ranks from the index. + async fn detach(&self, registry: &WorkerRegistry, record: &WorkerCatalogRecord); +} + +/// One ZMQ SUB socket per worker rank on the engine's KV-event publisher, with +/// gap replay from the engine's replay endpoint when it advertises one. +#[derive(Debug, Default, Clone, Copy)] +pub struct ZmqDirectIngress; + +#[async_trait] +impl KvEventIngress for ZmqDirectIngress { + fn missing_metadata(&self, record: &WorkerCatalogRecord) -> Vec { + let endpoints = record.listener_endpoints(); + record + .dp_ranks() + .filter(|rank| { + endpoints + .get(rank) + .is_none_or(|endpoint| endpoint.is_empty()) + }) + .map(|rank| format!("kv_events endpoint is required for dp_rank {rank}")) + .collect() + } + + async fn attach( + &self, + registry: &WorkerRegistry, + record: &WorkerCatalogRecord, + ) -> Result<(), SelectionError> { + let block_size = record + .block_size + .ok_or_else(|| SelectionError::BadRequest("block_size is required".to_string()))?; + let mut endpoints: Vec<_> = record.listener_endpoints().into_iter().collect(); + endpoints.sort_by_key(|(dp_rank, _)| *dp_rank); + for (dp_rank, endpoint) in endpoints { + crate::services::common::zmq::validate_endpoint(&endpoint).map_err(|error| { + SelectionError::BadRequest(format!( + "invalid kv_events endpoint for worker {} dp_rank {dp_rank}: {error}", + record.worker_id + )) + })?; + if let Some(replay_endpoint) = record.replay_endpoint.as_deref() { + crate::services::common::zmq::validate_endpoint(replay_endpoint).map_err( + |error| { + SelectionError::BadRequest(format!( + "invalid replay endpoint for worker {} dp_rank {dp_rank}: {error}", + record.worker_id + )) + }, + )?; + } + registry + .register( + record.worker_id, + endpoint, + dp_rank, + record.model_name.clone(), + record.routing_group.clone(), + block_size, + record.replay_endpoint.clone(), + ) + .await + .map_err(|error| SelectionError::BadRequest(error.to_string()))?; + } + Ok(()) + } + + async fn detach(&self, registry: &WorkerRegistry, record: &WorkerCatalogRecord) { + if let Err(error) = registry + .deregister(record.worker_id, &record.model_name, &record.routing_group) + .await + { + tracing::debug!( + worker_id = record.worker_id, + error = %error, + "indexer deregistration skipped or failed" + ); + } + } +} diff --git a/lib/kv-router/src/services/selection/mod.rs b/lib/kv-router/src/services/selection/mod.rs index cac678fb3bec..63ea5504c309 100644 --- a/lib/kv-router/src/services/selection/mod.rs +++ b/lib/kv-router/src/services/selection/mod.rs @@ -9,6 +9,7 @@ mod catalog; mod core; mod error; +mod ingress; mod input; mod membership; mod pending; @@ -31,6 +32,7 @@ pub use core::{ SelectionServiceConfig, }; pub use error::SelectionError; +pub use ingress::{KvEventIngress, ZmqDirectIngress}; pub use input::PromptRequest; pub use membership::{CatalogObserver, CatalogReconciler, WorkerCatalogSource}; pub use pending::SelectionCacheConfig; diff --git a/lib/kv-router/src/services/selection/types.rs b/lib/kv-router/src/services/selection/types.rs index aeef0590f26e..846739b34fb7 100644 --- a/lib/kv-router/src/services/selection/types.rs +++ b/lib/kv-router/src/services/selection/types.rs @@ -229,11 +229,7 @@ impl WorkerCatalogRecord { } } - pub(super) fn missing_schedulable_metadata( - &self, - queueing_enabled: bool, - kv_events_enabled: bool, - ) -> Vec { + pub(super) fn missing_schedulable_metadata(&self, queueing_enabled: bool) -> Vec { let mut missing = Vec::new(); if self.endpoint.as_deref().is_none_or(str::is_empty) { @@ -249,18 +245,6 @@ impl WorkerCatalogRecord { missing .push("max_num_batched_tokens is required while queueing is enabled".to_string()); } - if kv_events_enabled { - let endpoints = self.listener_endpoints(); - for rank in self.dp_ranks() { - if endpoints - .get(&rank) - .is_none_or(|endpoint| endpoint.is_empty()) - { - missing.push(format!("kv_events endpoint is required for dp_rank {rank}")); - } - } - } - missing } } From 9aa0971d66a30ffb6ef9f71b44dc94dbb56b7178 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Fri, 4 Sep 2026 01:25:26 +0000 Subject: [PATCH 028/102] feat(router): RuntimeIngress feeds the partition's index for the frontend KvEventIngress now builds the partition's index (open) as well as feeding it, so KvIndexSource::Owned covers both hosts and Provided goes away: - ZmqDirectIngress opens the registry's index and, on detach, deregisters the worker's listeners or, when there are none (events off or a remote primary), drops the worker from the index: the two behaviors SelectionCore used to branch on itself. - RuntimeIngress (lib/llm/src/kv_router/indexer/ingress.rs) builds the frontend's index from the router config, starts the event-plane subscription with per-rank recovery and the served-indexer endpoint, and hands the index to the partition. Its per-worker hooks are no-ops because membership comes from the KV-source watch. KvRouter::new shrinks to constructing it; task guard and shutdown go through it. The core no longer special-cases a host-fed index anywhere: every Owned index is opened by and detached through its ingress. Signed-off-by: Ishan Dhanani --- .../src/services/indexer/registry.rs | 5 + .../src/services/selection/README.md | 8 +- .../src/services/selection/core/mod.rs | 42 ++--- .../src/services/selection/ingress.rs | 73 ++++++-- lib/llm/src/kv_router.rs | 101 +++-------- lib/llm/src/kv_router/embedded.rs | 14 +- lib/llm/src/kv_router/indexer/ingress.rs | 157 ++++++++++++++++++ lib/llm/src/kv_router/indexer/mod.rs | 2 + 8 files changed, 267 insertions(+), 135 deletions(-) create mode 100644 lib/llm/src/kv_router/indexer/ingress.rs diff --git a/lib/kv-router/src/services/indexer/registry.rs b/lib/kv-router/src/services/indexer/registry.rs index e8f7f51a1074..710bee3f6ca0 100644 --- a/lib/kv-router/src/services/indexer/registry.rs +++ b/lib/kv-router/src/services/indexer/registry.rs @@ -676,6 +676,11 @@ impl WorkerRegistry { Ok(()) } + /// Whether `worker_id` has listeners registered here. + pub fn has_worker(&self, worker_id: WorkerId) -> bool { + self.workers.contains_key(&worker_id) + } + pub fn list(&self) -> Vec { self.list_filtered(None, None) } diff --git a/lib/kv-router/src/services/selection/README.md b/lib/kv-router/src/services/selection/README.md index d011b9fc61fa..083ca127ad12 100644 --- a/lib/kv-router/src/services/selection/README.md +++ b/lib/kv-router/src/services/selection/README.md @@ -17,10 +17,10 @@ and active-sequence accounting. Keep these implementation invariants explicit: diff (retry of not-yet-schedulable records, stale deletion). Hosts do not call `upsert_worker`/`delete_worker` for discovery-driven membership. - A partition's KV index comes from `HostCache.index: KvIndexSource`: `Owned` - (the core builds it and a `KvEventIngress`, by default `ZmqDirectIngress`, - feeds it), `Provided` (the host built and feeds it; the core only reads), - or `Remote` (a standalone indexer serves it). The `Indexer` type itself is - shared with the frontend (`services::indexer::backend`). + (a `KvEventIngress` builds and feeds it: `ZmqDirectIngress` here, the + runtime event plane in the frontend) or `Remote` (a standalone indexer + serves it). The `Indexer` type itself is shared with the frontend + (`services::indexer::backend`). - Selector replicas synchronize admission, prefill-complete, and free events. - **NOTE:** Output-block updates remain local. They are deliberately excluded from replica sync because their frequency would consume disproportionate diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index 0bb8c7bc4d3e..6a65147a8e6e 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -211,15 +211,11 @@ pub struct HostCache { /// Where a partition's KV index comes from and who feeds it. #[derive(Clone)] pub enum KvIndexSource { - /// The core builds the index and the ingress feeds it with worker KV - /// events; the ingress decides what metadata a worker needs to be - /// schedulable. Defaults to [`ZmqDirectIngress`]. + /// The ingress builds each partition's index and feeds it with worker KV + /// events; it also decides what metadata a worker needs to be schedulable + /// and what happens to the index when a worker leaves. Defaults to + /// [`ZmqDirectIngress`]; the frontend supplies its runtime-backed ingress. Owned(Arc), - /// The host built this index and feeds it (its own event ingress and - /// recovery). Every partition uses it; the core reads it for selection - /// and dequeue-time refresh but never writes to it, and workers need no - /// KV-event metadata. - Provided(Arc), /// A standalone indexer at this base URL serves the primary index; this /// core does not subscribe to worker KV events. Remote(String), @@ -235,11 +231,6 @@ impl std::fmt::Debug for KvIndexSource { fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { Self::Owned(_) => formatter.write_str("Owned"), - Self::Provided(indexer) => formatter - .debug_struct("Provided") - .field("remote", &indexer.is_remote()) - .field("overlap_refresh", &indexer.supports_overlap_refresh()) - .finish(), Self::Remote(url) => formatter.debug_tuple("Remote").field(url).finish(), } } @@ -409,9 +400,7 @@ impl SelectionCore { indexer_threads, cancel_token.clone(), )); - let listens_for_kv_events = kv_router_config.use_kv_events - && !indexer_policy.is_remote() - && !matches!(host.cache.index, KvIndexSource::Provided(_)); + let listens_for_kv_events = kv_router_config.use_kv_events && !indexer_policy.is_remote(); indexer_registry.set_indexer_policy(indexer_policy); if signal_indexer_ready { indexer_registry.signal_ready(); @@ -731,17 +720,15 @@ impl SelectionCore { self.cancel_token.child_token(), ); - let (indexer, refreshes) = match &self.host.cache.index { - KvIndexSource::Provided(indexer) => { - (Indexer::clone(indexer), indexer.supports_overlap_refresh()) + let indexer = match &self.host.cache.index { + KvIndexSource::Owned(ingress) => { + ingress.open(&self.indexer_registry, &key, block_size) } - KvIndexSource::Owned(_) | KvIndexSource::Remote(_) => ( - self.indexer_registry - .get_or_create_indexer(key.clone(), block_size), - true, - ), + KvIndexSource::Remote(_) => self + .indexer_registry + .get_or_create_indexer(key.clone(), block_size), }; - let overlap_refresh = refreshes.then(|| { + let overlap_refresh = indexer.supports_overlap_refresh().then(|| { Arc::new(TieredOverlapRefresher::new( RefreshProvider::Local(Arc::new(indexer.clone())), self.kv_router_config.clone(), @@ -809,10 +796,7 @@ impl SelectionCore { } async fn cleanup_indexer_registration(&self, record: &WorkerCatalogRecord) { - if matches!(self.host.cache.index, KvIndexSource::Provided(_)) { - return; - } - if let Some(ingress) = self.ingress() { + if let KvIndexSource::Owned(ingress) = &self.host.cache.index { ingress.detach(&self.indexer_registry, record).await; return; } diff --git a/lib/kv-router/src/services/selection/ingress.rs b/lib/kv-router/src/services/selection/ingress.rs index bb0ae8ab81dc..ac97f9c02807 100644 --- a/lib/kv-router/src/services/selection/ingress.rs +++ b/lib/kv-router/src/services/selection/ingress.rs @@ -7,28 +7,47 @@ use async_trait::async_trait; use super::error::SelectionError; use super::types::WorkerCatalogRecord; +use crate::identity::RoutingPartitionId; +use crate::services::indexer::backend::Indexer; use crate::services::indexer::registry::WorkerRegistry; -/// Feeds a partition's core-owned index ([`super::KvIndexSource::Owned`]). +/// Builds and feeds a partition's index ([`super::KvIndexSource::Owned`]). /// -/// The registry holds the partition indexes and the per-rank listener -/// lifecycle; an ingress attaches and detaches workers on it. A host that feeds -/// the index itself uses [`super::KvIndexSource::Provided`] instead. +/// The registry holds the standalone indexes and the per-rank listener +/// lifecycle; the ZMQ ingress works on it. A host with its own event transport +/// (the frontend over the Dynamo runtime) builds and feeds the index itself +/// and treats the per-worker hooks as membership it already tracks. #[async_trait] pub trait KvEventIngress: Send + Sync { + /// The index for `key`, created on first use. Called once per partition, + /// before any worker attaches. + fn open(&self, registry: &WorkerRegistry, key: &RoutingPartitionId, block_size: u32) + -> Indexer; + /// Metadata `record` must carry before this ingress can feed its ranks. - /// A non-empty result keeps the worker `Incomplete`. - fn missing_metadata(&self, record: &WorkerCatalogRecord) -> Vec; + /// A non-empty result keeps the worker `Incomplete`. Consulted only when + /// the core subscribes to worker KV events. + fn missing_metadata(&self, record: &WorkerCatalogRecord) -> Vec { + let _ = record; + Vec::new() + } /// Start feeding every rank of `record` into its partition's index. + /// Called only when the core subscribes to worker KV events. async fn attach( &self, registry: &WorkerRegistry, record: &WorkerCatalogRecord, - ) -> Result<(), SelectionError>; + ) -> Result<(), SelectionError> { + let _ = (registry, record); + Ok(()) + } - /// Stop feeding `record` and drop its ranks from the index. - async fn detach(&self, registry: &WorkerRegistry, record: &WorkerCatalogRecord); + /// `record` left the catalog: stop feeding it and drop what the index + /// holds for it. Called for every worker, subscribed or not. + async fn detach(&self, registry: &WorkerRegistry, record: &WorkerCatalogRecord) { + let _ = (registry, record); + } } /// One ZMQ SUB socket per worker rank on the engine's KV-event publisher, with @@ -38,6 +57,15 @@ pub struct ZmqDirectIngress; #[async_trait] impl KvEventIngress for ZmqDirectIngress { + fn open( + &self, + registry: &WorkerRegistry, + key: &RoutingPartitionId, + block_size: u32, + ) -> Indexer { + registry.get_or_create_indexer(key.clone(), block_size) + } + fn missing_metadata(&self, record: &WorkerCatalogRecord) -> Vec { let endpoints = record.listener_endpoints(); record @@ -95,15 +123,26 @@ impl KvEventIngress for ZmqDirectIngress { } async fn detach(&self, registry: &WorkerRegistry, record: &WorkerCatalogRecord) { - if let Err(error) = registry - .deregister(record.worker_id, &record.model_name, &record.routing_group) - .await + if registry.has_worker(record.worker_id) { + if let Err(error) = registry + .deregister(record.worker_id, &record.model_name, &record.routing_group) + .await + { + tracing::debug!( + worker_id = record.worker_id, + error = %error, + "indexer deregistration skipped or failed" + ); + } + return; + } + // No listeners (events disabled or a remote primary): drop what an + // approximate or side index recorded for the worker. + if let Some(indexer) = registry + .get_indexer(&record.key()) + .map(|entry| entry.indexer.clone()) { - tracing::debug!( - worker_id = record.worker_id, - error = %error, - "indexer deregistration skipped or failed" - ); + indexer.remove_worker(record.worker_id).await; } } } diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index a94f923e486f..4622d012d660 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -593,12 +593,11 @@ pub struct KvRouter { cancellation_token: CancellationToken, client: Client, is_eagle: bool, - kv_event_subscription: Option, + ingress: Arc, tracking_hash: TrackingHashContext, tracking_model_name: String, approximate_lru_ranks: ApproximateLruRanks, request_leases: request_lease::RequestLeaseManager, - _served_indexer_handle: Option, /// Optional external shared KV cache pool. When present, `find_best_match` /// queries it in parallel with the indexer and factors shared hits into scoring. shared_cache: Option>, @@ -757,18 +756,20 @@ impl KvRouter { let cancellation_guard = cancellation_token.clone().drop_guard(); let min_initial_workers = min_initial_workers_from_env()?; - let indexer = if cache_required { - indexer::build( - component, - &kv_router_config, - block_size, - model_name.as_deref(), - cancellation_token.child_token(), - ) - .await? - } else { - Indexer::None - }; + let ingress = indexer::RuntimeIngress::start(indexer::RuntimeIngressArgs { + endpoint: &endpoint, + kv_router_config: &kv_router_config, + block_size, + model_name: model_name.as_deref(), + worker_role, + metric_worker_type, + cache_required, + kv_event_source_requirement, + kv_source_membership, + cancellation_token: cancellation_token.clone(), + }) + .await?; + let indexer = ingress.indexer().clone(); let approximate_lru_metrics = metrics::ApproximateLruMetrics::from_component(component); let configured_policy = kv_router_config.router_approximate_cache_policy.to_string(); let effective_policy = if kv_router_config.overlap_score_credit <= 0.0 { @@ -831,7 +832,8 @@ impl KvRouter { prefill_load_estimator: prefill_load_estimator.clone(), overloaded_worker_provider, available_worker_provider, - indexer: indexer.clone(), + ingress: Arc::clone(&ingress) + as Arc, scheduler_load, endpoint: endpoint.clone(), router_id: endpoint.drt().discovery().instance_id(), @@ -845,58 +847,6 @@ impl KvRouter { scheduler.booking_cleanup(), cancellation_token.child_token(), ); - // Start KV event subscription if needed — skip when using a remote indexer. - let kv_event_subscription = if cache_required - && kv_event_source_requirement.should_subscribe(&kv_router_config) - { - let membership_watch = kv_source_membership.ok_or_else(|| { - anyhow::anyhow!( - "KV source membership watch is required when local KV event subscription is enabled" - ) - })?; - Some( - indexer::start_subscriber( - endpoint.clone(), - indexer.clone(), - membership_watch, - block_size, - model_name.clone().unwrap_or_else(|| "unknown".to_string()), - worker_role, - kv_event_source_requirement, - metric_worker_type, - cancellation_token.child_token(), - ) - .await?, - ) - } else { - tracing::info!( - requirement = %kv_event_source_requirement, - cache_required, - "Skipping KV event subscription (use_kv_events={}, overlap_score_credit={}, use_remote_indexer={})", - kv_router_config.use_kv_events, - kv_router_config.overlap_score_credit, - kv_router_config.use_remote_indexer, - ); - None - }; - - let served_indexer_handle = if kv_router_config.serve_indexer { - let model_name = model_name.clone().ok_or_else(|| { - anyhow::anyhow!("model_name is required when serve_indexer is configured") - })?; - Some( - ensure_served_indexer_service( - component.clone(), - ServedIndexerMode::from_use_kv_events(kv_router_config.use_kv_events), - model_name, - indexer.clone(), - ) - .await?, - ) - } else { - None - }; - tracing::info!("KV Routing initialized"); let cancellation_token = cancellation_guard.disarm(); Ok(Self { @@ -910,12 +860,11 @@ impl KvRouter { cancellation_token, client, is_eagle, - kv_event_subscription, + ingress, tracking_hash, tracking_model_name, approximate_lru_ranks, request_leases, - _served_indexer_handle: served_indexer_handle, shared_cache, lora_filter, endpoint_registration: None, @@ -934,9 +883,7 @@ impl KvRouter { &mut self, task_guard: dynamo_runtime::engine::EngineContextGuard, ) { - if let Some(subscription) = self.kv_event_subscription.as_mut() { - subscription.set_task_guard(task_guard.clone()); - } + self.ingress.set_task_guard(task_guard.clone()); self.teardown_task_guard = Some(task_guard); } @@ -958,11 +905,9 @@ impl KvRouter { } /// Cancel background work and wait for KV event ingestion to stop. - pub async fn shutdown(mut self) { + pub async fn shutdown(self) { self.cancellation_token.cancel(); - if let Some(subscription) = self.kv_event_subscription.take() { - subscription.shutdown().await; - } + self.ingress.shutdown().await; } pub fn is_eagle(&self) -> bool { @@ -2787,7 +2732,7 @@ mod tests { let router = make_router_without_membership(Some(WorkerType::Decode)) .await .expect("ordinary decode must not require KV source membership"); - assert!(router.kv_event_subscription.is_none()); + assert!(!router.ingress.has_subscription()); let error = make_router_without_membership(None) .await @@ -2836,7 +2781,7 @@ mod tests { assert_eq!(router.required_worker_inputs(), WorkerInputs::LOAD); assert!(matches!(router.indexer, Indexer::None)); - assert!(router.kv_event_subscription.is_none()); + assert!(!router.ingress.has_subscription()); assert!(router.shared_cache.is_none()); assert!(matches!( router.dump_events().await, diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs index d393fe38d3e3..381d43ffe24c 100644 --- a/lib/llm/src/kv_router/embedded.rs +++ b/lib/llm/src/kv_router/embedded.rs @@ -28,9 +28,9 @@ use dynamo_kv_router::scheduling::{ use dynamo_kv_router::sequences::{SequenceError, SequenceRequest}; use dynamo_kv_router::services::selection::{ CatalogObserver, CatalogReconciler, HostCache, HostEligibility, HostLoad, HostReplication, - HostTelemetry, KvIndexSource, SelectionHost, SelectionPartition, SelectionService, - SelectionServiceBuilder, WorkerCatalogRecord, WorkerCatalogSource, WorkerRequest, - WorkerSelectionPolicyRegistry, + HostTelemetry, KvEventIngress, KvIndexSource, SelectionHost, SelectionPartition, + SelectionService, SelectionServiceBuilder, WorkerCatalogRecord, WorkerCatalogSource, + WorkerRequest, WorkerSelectionPolicyRegistry, }; use dynamo_kv_router::{DEFAULT_ROUTING_GROUP, PrefillLoadEstimator, WorkerSelectionPolicyFactory}; use dynamo_tokens::SequenceHash; @@ -50,9 +50,9 @@ pub(crate) struct EmbeddedSelectionArgs { pub prefill_load_estimator: Option>, pub overloaded_worker_provider: OverloadedWorkerProvider, pub available_worker_provider: WorkerAvailabilityProvider, - /// The router's index, built and fed by the router; the partition reads - /// it for selection and dequeue-time refresh. - pub indexer: super::indexer::Indexer, + /// Builds and feeds the router's index from the runtime; the partition + /// takes its index from it. + pub ingress: Arc, /// Scheduler-owned load snapshots for the worker monitor's overload /// detection, the same feed the runtime scheduler publishes. pub scheduler_load: crate::kv_router::routing_load::SchedulerLoadSender, @@ -155,7 +155,7 @@ impl EmbeddedSelection { }, cache: HostCache { shared: None, - index: KvIndexSource::Provided(Arc::new(args.indexer)), + index: KvIndexSource::Owned(args.ingress), }, eligibility: HostEligibility::default(), telemetry: HostTelemetry { diff --git a/lib/llm/src/kv_router/indexer/ingress.rs b/lib/llm/src/kv_router/indexer/ingress.rs new file mode 100644 index 000000000000..6abdd5aa4fd4 --- /dev/null +++ b/lib/llm/src/kv_router/indexer/ingress.rs @@ -0,0 +1,157 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! The frontend's KV-event ingress: builds the router's index and feeds it from +//! the Dynamo runtime (event plane subscription with per-rank recovery, and the +//! served-indexer endpoint when configured). Worker membership comes from the +//! KV-source membership watch, so the per-worker catalog hooks are no-ops. + +use std::sync::Arc; + +use anyhow::Result; +use dynamo_kv_router::WorkerType; +use dynamo_kv_router::config::KvRouterConfig; +use dynamo_kv_router::identity::RoutingPartitionId; +use dynamo_kv_router::services::indexer::registry::WorkerRegistry; +use dynamo_kv_router::services::selection::KvEventIngress; +use dynamo_runtime::component::Endpoint; +use tokio_util::sync::CancellationToken; + +use super::{ + Indexer, KvEventSubscriptionHandle, ServedIndexerMode, ensure_served_indexer_service, + start_subscriber, +}; +use crate::discovery::KvSourceMembershipWatch; +use crate::kv_router::KvEventSourceRequirement; + +pub(crate) struct RuntimeIngressArgs<'a> { + pub endpoint: &'a Endpoint, + pub kv_router_config: &'a KvRouterConfig, + pub block_size: u32, + pub model_name: Option<&'a str>, + pub worker_role: Option, + pub metric_worker_type: &'static str, + /// Whether any consumer needs KV knowledge; without it the index is `None`. + pub cache_required: bool, + pub kv_event_source_requirement: KvEventSourceRequirement, + pub kv_source_membership: Option, + pub cancellation_token: CancellationToken, +} + +pub(crate) struct RuntimeIngress { + indexer: Indexer, + subscription: parking_lot::Mutex>, + _served: Option, +} + +impl RuntimeIngress { + pub(crate) async fn start(args: RuntimeIngressArgs<'_>) -> Result> { + let component = args.endpoint.component(); + let indexer = if args.cache_required { + super::build( + component, + args.kv_router_config, + args.block_size, + args.model_name, + args.cancellation_token.child_token(), + ) + .await? + } else { + Indexer::None + }; + + let subscription = if args.cache_required + && args + .kv_event_source_requirement + .should_subscribe(args.kv_router_config) + { + let membership_watch = args.kv_source_membership.ok_or_else(|| { + anyhow::anyhow!( + "KV source membership watch is required when local KV event subscription is enabled" + ) + })?; + Some( + start_subscriber( + args.endpoint.clone(), + indexer.clone(), + membership_watch, + args.block_size, + args.model_name.unwrap_or("unknown").to_string(), + args.worker_role, + args.kv_event_source_requirement, + args.metric_worker_type, + args.cancellation_token.child_token(), + ) + .await?, + ) + } else { + tracing::info!( + requirement = %args.kv_event_source_requirement, + cache_required = args.cache_required, + "Skipping KV event subscription (use_kv_events={}, overlap_score_credit={}, use_remote_indexer={})", + args.kv_router_config.use_kv_events, + args.kv_router_config.overlap_score_credit, + args.kv_router_config.use_remote_indexer, + ); + None + }; + + let served = if args.kv_router_config.serve_indexer { + let model_name = args.model_name.ok_or_else(|| { + anyhow::anyhow!("model_name is required when serve_indexer is configured") + })?; + Some( + ensure_served_indexer_service( + component.clone(), + ServedIndexerMode::from_use_kv_events(args.kv_router_config.use_kv_events), + model_name.to_string(), + indexer.clone(), + ) + .await?, + ) + } else { + None + }; + + Ok(Arc::new(Self { + indexer, + subscription: parking_lot::Mutex::new(subscription), + _served: served, + })) + } + + pub(crate) fn indexer(&self) -> &Indexer { + &self.indexer + } + + #[cfg(test)] + pub(crate) fn has_subscription(&self) -> bool { + self.subscription.lock().is_some() + } + + pub(crate) fn set_task_guard(&self, task_guard: dynamo_runtime::engine::EngineContextGuard) { + if let Some(subscription) = self.subscription.lock().as_mut() { + subscription.set_task_guard(task_guard); + } + } + + /// Stop event ingestion and wait for it to drain. + pub(crate) async fn shutdown(&self) { + let subscription = self.subscription.lock().take(); + if let Some(subscription) = subscription { + subscription.shutdown().await; + } + } +} + +#[async_trait::async_trait] +impl KvEventIngress for RuntimeIngress { + fn open( + &self, + _registry: &WorkerRegistry, + _key: &RoutingPartitionId, + _block_size: u32, + ) -> Indexer { + self.indexer.clone() + } +} diff --git a/lib/llm/src/kv_router/indexer/mod.rs b/lib/llm/src/kv_router/indexer/mod.rs index 11a42212673a..e54d3d45596d 100644 --- a/lib/llm/src/kv_router/indexer/mod.rs +++ b/lib/llm/src/kv_router/indexer/mod.rs @@ -26,9 +26,11 @@ pub(crate) use dynamo_kv_router::indexer::WireTieredMatchDetails; pub use dynamo_kv_router::services::indexer::backend::{Indexer, RemotePrimary, SideIndexer}; pub(crate) use dynamo_kv_router::services::indexer::recording::ApproximateRequestLease; use dynamo_runtime::component::Component; +pub(crate) use ingress::{RuntimeIngress, RuntimeIngressArgs}; use tokio_util::sync::CancellationToken; mod embedding_cache; +mod ingress; mod recovery; pub mod remote; From f413da4bd432bf6c1db31f3dd4fc2a05ccfa4da2 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Fri, 4 Sep 2026 00:31:16 +0000 Subject: [PATCH 029/102] refactor(kv-router): session-affinity table shared by every selection host Move the session table out of lib/llm/src/session_affinity into dynamo_kv_router::services::selection::affinity: SessionAffinity (session id -> bound worker/rank, versioned, TTL reaper, entry and id limits), the AffinityInitialization / AffinityLease lifecycle (commit, rebind, invalidate, release on drop), hard and soft validation, and replica apply with the same outcome table. Replication becomes an AffinityReplicaSink the host installs; the table stays transport-free. The frontend AffinityCoordinator is now a thin wrapper: it converts the pipeline RouteTarget at the boundary, abandons an initialization wait when the request context stops, holds a lease for the life of a response stream, and keeps its runtime event-plane replica sync (same subject and wire format) behind the sink. SessionAffinityMode moves with the table; lib/llm re-exports it. All 37 session-affinity tests and the routing-host affinity tests run unchanged against the wrapper. Signed-off-by: Ishan Dhanani --- .../src/services/selection/README.md | 5 + .../src/services/selection/affinity.rs | 847 +++++++++++++++++ lib/kv-router/src/services/selection/mod.rs | 1 + lib/llm/src/kv_router/routing_host/tests.rs | 32 +- lib/llm/src/session_affinity/coordinator.rs | 865 +++--------------- lib/llm/src/session_affinity/mod.rs | 35 +- lib/llm/src/session_affinity/replica_sync.rs | 62 +- lib/llm/src/session_affinity/tests.rs | 35 +- 8 files changed, 1082 insertions(+), 800 deletions(-) create mode 100644 lib/kv-router/src/services/selection/affinity.rs diff --git a/lib/kv-router/src/services/selection/README.md b/lib/kv-router/src/services/selection/README.md index 083ca127ad12..ef56760ef9ef 100644 --- a/lib/kv-router/src/services/selection/README.md +++ b/lib/kv-router/src/services/selection/README.md @@ -21,6 +21,11 @@ and active-sequence accounting. Keep these implementation invariants explicit: runtime event plane in the frontend) or `Remote` (a standalone indexer serves it). The `Indexer` type itself is shared with the frontend (`services::indexer::backend`). +- Session affinity (`services::selection::affinity::SessionAffinity`) is the + one session table: versioned bindings, TTL reaper, initialization/lease + lifecycle, hard and soft validation, replica apply. Hosts own the instance + and the replication transport (`AffinityReplicaSink`); the frontend wraps it + for pipeline streams and the runtime event plane. - Selector replicas synchronize admission, prefill-complete, and free events. - **NOTE:** Output-block updates remain local. They are deliberately excluded from replica sync because their frequency would consume disproportionate diff --git a/lib/kv-router/src/services/selection/affinity.rs b/lib/kv-router/src/services/selection/affinity.rs new file mode 100644 index 000000000000..388a8c8d8d8c --- /dev/null +++ b/lib/kv-router/src/services/selection/affinity.rs @@ -0,0 +1,847 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Session affinity: pin a session id to the worker (and optionally the DP +//! rank) that served it, for as long as requests keep arriving within the TTL. +//! +//! [`SessionAffinity`] is the table. A host acquires a session before +//! selecting a worker: a new session yields an [`AffinityInitialization`] the +//! host commits with the worker it dispatched to; a bound session yields the +//! target and an [`AffinityLease`] that keeps the binding alive until the +//! response stream ends. Bindings are versioned so replicas can exchange them +//! through an [`AffinityReplicaSink`]; the transport is the host's. + +use std::pin::Pin; +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; +use std::sync::{Arc, OnceLock, Weak}; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use dashmap::{DashMap, mapref::entry::Entry}; +use serde::{Deserialize, Serialize}; +use tokio::sync::Notify; +use tokio::sync::futures::OwnedNotified; +use tokio::time::Instant; +use tokio_util::sync::CancellationToken; + +pub use crate::protocols::WorkerAffinityTarget as AffinityTarget; + +pub const MAX_SESSION_AFFINITY_TTL_SECS: u64 = 31_536_000; +pub const MAX_SESSION_AFFINITY_ENTRIES: usize = 65_536; +pub const MAX_SESSION_AFFINITY_ID_BYTES: usize = 256; + +/// How a bound session treats a dispatch that landed elsewhere. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum SessionAffinityMode { + /// The binding is exact: dispatching to another worker or rank is an error. + #[default] + Hard, + /// The binding follows the dispatch: the session rebinds to where it ran. + Soft, +} + +impl std::str::FromStr for SessionAffinityMode { + type Err = String; + + fn from_str(value: &str) -> Result { + match value { + "hard" => Ok(Self::Hard), + "soft" => Ok(Self::Soft), + _ => Err(format!( + "invalid session affinity mode {value:?}; expected 'hard' or 'soft'" + )), + } + } +} + +#[derive(Debug, thiserror::Error)] +pub enum AffinityError { + #[error("{0}")] + InvalidArgument(String), + #[error("{0}")] + ResourceExhausted(String), + #[error("session affinity table dropped")] + Dropped, +} + +/// Ordering for replicated bindings: higher wins, ties broken by writer. +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)] +pub struct AffinityVersion { + pub sequence: u64, + pub writer_id: u64, +} + +/// Receives every binding this table publishes for its replicas. +pub trait AffinityReplicaSink: Send + Sync { + fn publish(&self, session_id: &str, target: AffinityTarget, version: AffinityVersion); +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ReplicaApplyOutcome { + Inserted, + Refreshed, + ReplacedExpired, + ReplacedNewer, + IgnoredInitializing, + IgnoredConflict, + RejectedSessionId, + RejectedCapacity, +} + +enum AffinityEntry { + Initializing { + revision: u64, + notify: Arc, + }, + Bound { + target: AffinityTarget, + revision: u64, + version: AffinityVersion, + active_leases: usize, + idle_deadline: Instant, + }, +} + +struct Inner { + entries: DashMap, + ttl: Duration, + max_entries: usize, + max_session_id_bytes: usize, + entry_count: AtomicUsize, + next_revision: AtomicU64, + next_sequence: AtomicU64, + writer_id: AtomicU64, + cancel: CancellationToken, + replica: OnceLock>, + reaper_started: Arc, + waiter_observed: Arc, +} + +impl Drop for Inner { + fn drop(&mut self) { + self.cancel.cancel(); + } +} + +/// The session table. Cheap to clone; the last clone stops the reaper. +#[derive(Clone)] +pub struct SessionAffinity { + inner: Arc, +} + +/// A non-owning handle for replica appliers, so a sink task never keeps the +/// table alive. +#[derive(Clone)] +pub struct WeakSessionAffinity { + inner: Weak, +} + +impl WeakSessionAffinity { + pub fn upgrade(&self) -> Option { + self.inner.upgrade().map(|inner| SessionAffinity { inner }) + } +} + +/// One step of acquiring a session. +pub enum AcquireStep { + /// New (or expired) session: select a worker, then `commit` it. + Initialize(AffinityInitialization), + /// Bound session: route to `target`; the lease releases on drop. + Bound { + target: AffinityTarget, + lease: AffinityLease, + }, + /// Another request is initializing this session; await, then retry. + Wait(Pin>), +} + +impl SessionAffinity { + pub fn new(ttl: Duration) -> Result { + Self::new_with_limits( + ttl, + MAX_SESSION_AFFINITY_ENTRIES, + MAX_SESSION_AFFINITY_ID_BYTES, + ) + } + + pub fn new_with_limits( + ttl: Duration, + max_entries: usize, + max_session_id_bytes: usize, + ) -> Result { + if !(Duration::from_secs(1)..=Duration::from_secs(MAX_SESSION_AFFINITY_TTL_SECS)) + .contains(&ttl) + { + return Err(AffinityError::InvalidArgument(format!( + "session affinity TTL must be between 1 and {MAX_SESSION_AFFINITY_TTL_SECS} seconds" + ))); + } + let inner = Arc::new(Inner { + entries: DashMap::new(), + ttl, + max_entries, + max_session_id_bytes, + entry_count: AtomicUsize::new(0), + next_revision: AtomicU64::new(1), + next_sequence: AtomicU64::new( + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_nanos() as u64, + ), + writer_id: AtomicU64::new(0), + cancel: CancellationToken::new(), + replica: OnceLock::new(), + reaper_started: Arc::new(Notify::new()), + waiter_observed: Arc::new(Notify::new()), + }); + Self::spawn_reaper(&inner); + tracing::info!( + ttl_secs = ttl.as_secs(), + max_entries, + "session affinity enabled" + ); + Ok(Self { inner }) + } + + fn spawn_reaper(inner: &Arc) { + let weak = Arc::downgrade(inner); + let cancel = inner.cancel.clone(); + let period = inner.ttl.min(Duration::from_secs(30)); + let reaper_started = inner.reaper_started.clone(); + tokio::spawn(async move { + reaper_started.notify_one(); + loop { + tokio::select! { + _ = cancel.cancelled() => return, + _ = tokio::time::sleep(period) => {} + } + let Some(inner) = weak.upgrade() else { + return; + }; + let now = Instant::now(); + let mut removed = 0; + inner.entries.retain(|_, entry| { + let retain = !matches!( + entry, + AffinityEntry::Bound { + active_leases: 0, + idle_deadline, + .. + } if *idle_deadline <= now + ); + removed += usize::from(!retain); + retain + }); + inner.entry_count.fetch_sub(removed, Ordering::Relaxed); + } + }); + } + + pub fn downgrade(&self) -> WeakSessionAffinity { + WeakSessionAffinity { + inner: Arc::downgrade(&self.inner), + } + } + + /// Install the replica sink and this replica's writer id. Returns `false` + /// when a sink is already installed. + pub fn enable_replication(&self, writer_id: u64, sink: Arc) -> bool { + self.inner.writer_id.store(writer_id, Ordering::Relaxed); + self.inner.replica.set(sink).is_ok() + } + + pub fn ttl(&self) -> Duration { + self.inner.ttl + } + + /// Acquire `session_id` for a request. `requested_target` is an explicit + /// pin the request carries; a bound session must agree with it. + pub fn try_acquire( + &self, + session_id: &str, + requested_target: Option, + ) -> Result { + self.validate_session_id(session_id)?; + let now = Instant::now(); + match self.inner.entries.entry(session_id.to_string()) { + Entry::Vacant(entry) => { + self.reserve_entry()?; + tracing::debug!( + session_id, + "session affinity miss: new session, pinning after worker selection" + ); + let revision = self.inner.next_revision.fetch_add(1, Ordering::Relaxed); + let notify = Arc::new(Notify::new()); + entry.insert(AffinityEntry::Initializing { + revision, + notify: notify.clone(), + }); + Ok(AcquireStep::Initialize(AffinityInitialization { + table: Arc::downgrade(&self.inner), + session_id: session_id.to_string(), + revision, + notify, + requested_target, + active: true, + })) + } + Entry::Occupied(mut entry) => match entry.get_mut() { + AffinityEntry::Initializing { notify, .. } => { + self.inner.waiter_observed.notify_one(); + // Register before releasing the entry so a commit or a + // dropped initialization between the two cannot be missed. + let mut notified = Box::pin(notify.clone().notified_owned()); + notified.as_mut().enable(); + Ok(AcquireStep::Wait(notified)) + } + AffinityEntry::Bound { + active_leases, + idle_deadline, + .. + } if *active_leases == 0 && *idle_deadline <= now => { + tracing::debug!( + session_id, + "session affinity miss: pin expired (idle past TTL), re-selecting worker" + ); + let revision = self.inner.next_revision.fetch_add(1, Ordering::Relaxed); + let notify = Arc::new(Notify::new()); + *entry.get_mut() = AffinityEntry::Initializing { + revision, + notify: notify.clone(), + }; + Ok(AcquireStep::Initialize(AffinityInitialization { + table: Arc::downgrade(&self.inner), + session_id: session_id.to_string(), + revision, + notify, + requested_target, + active: true, + })) + } + AffinityEntry::Bound { + target, + revision, + version, + active_leases, + .. + } => { + validate_bound_target(session_id, *target, requested_target)?; + tracing::debug!( + session_id, + worker_id = target.worker_id, + dp_rank = ?target.dp_rank, + active_leases = *active_leases + 1, + "session affinity hit: reusing pinned worker" + ); + *active_leases += 1; + Ok(AcquireStep::Bound { + target: *target, + lease: AffinityLease { + table: Arc::downgrade(&self.inner), + session_id: session_id.to_string(), + revision: *revision, + version: *version, + active: true, + }, + }) + } + }, + } + } + + /// Acquire, waiting out another request's initialization. Hosts that need + /// to abandon the wait on their own signal drive [`Self::try_acquire`]. + pub async fn acquire( + &self, + session_id: &str, + requested_target: Option, + ) -> Result { + loop { + match self.try_acquire(session_id, requested_target)? { + AcquireStep::Wait(notified) => notified.await, + step => return Ok(step), + } + } + } + + /// The bound target, without taking a lease. Read-only. + pub fn query_target( + &self, + session_id: &str, + requested_target: Option, + ) -> Result, AffinityError> { + self.validate_session_id(session_id)?; + let Some(entry) = self.inner.entries.get(session_id) else { + return Ok(None); + }; + let AffinityEntry::Bound { + target, + active_leases, + idle_deadline, + .. + } = entry.value() + else { + return Ok(None); + }; + if *active_leases == 0 && *idle_deadline <= Instant::now() { + return Ok(None); + } + validate_bound_target(session_id, *target, requested_target)?; + tracing::debug!( + session_id, + worker_id = target.worker_id, + dp_rank = ?target.dp_rank, + "session affinity hit: reusing pinned worker" + ); + Ok(Some(*target)) + } + + /// Advance the local sequence past one seen from a replica. + pub fn observe_replica_sequence(&self, sequence: u64) { + self.inner.observe_replica_sequence(sequence); + } + + /// Apply a binding published by a replica. + pub fn apply_replica_update( + &self, + session_id: String, + target: AffinityTarget, + version: AffinityVersion, + ) -> ReplicaApplyOutcome { + self.inner.apply_replica_update(session_id, target, version) + } + + pub fn entry_count(&self) -> usize { + self.inner.entry_count.load(Ordering::Relaxed) + } + + /// Cancelled when the last handle drops. + pub fn cancellation_token(&self) -> CancellationToken { + self.inner.cancel.clone() + } + + /// Test hook: the reaper task has started. + pub async fn wait_for_reaper(&self) { + self.inner.reaper_started.notified().await; + } + + /// Test hook: a request observed an initializing session and is waiting. + pub async fn wait_for_initializing_waiter(&self) { + self.inner.waiter_observed.notified().await; + } + + /// Test hook: expire an idle bound session now. + pub fn expire_for_test(&self, session_id: &str) { + let Some(mut entry) = self.inner.entries.get_mut(session_id) else { + panic!("session affinity entry missing"); + }; + let AffinityEntry::Bound { + active_leases, + idle_deadline, + .. + } = entry.value_mut() + else { + panic!("session affinity entry is not bound"); + }; + assert_eq!(*active_leases, 0); + *idle_deadline = Instant::now(); + } + + pub fn next_version(&self) -> AffinityVersion { + self.inner.next_version() + } + + fn validate_session_id(&self, session_id: &str) -> Result<(), AffinityError> { + if session_id.len() > self.inner.max_session_id_bytes { + return Err(AffinityError::InvalidArgument(format!( + "session affinity ID must not exceed {} bytes", + self.inner.max_session_id_bytes + ))); + } + Ok(()) + } + + fn reserve_entry(&self) -> Result<(), AffinityError> { + self.inner.reserve_entry().then_some(()).ok_or_else(|| { + AffinityError::ResourceExhausted("session affinity entry limit reached".to_string()) + }) + } +} + +impl Inner { + fn reserve_entry(&self) -> bool { + self.entry_count + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |count| { + (count < self.max_entries).then_some(count + 1) + }) + .is_ok() + } + + fn publish_replica_update( + &self, + session_id: &str, + target: AffinityTarget, + version: AffinityVersion, + ) { + if let Some(replica) = self.replica.get() { + replica.publish(session_id, target, version); + } + } + + fn next_version(&self) -> AffinityVersion { + AffinityVersion { + sequence: self.next_sequence.fetch_add(1, Ordering::Relaxed), + writer_id: self.writer_id.load(Ordering::Relaxed), + } + } + + fn observe_replica_sequence(&self, sequence: u64) { + self.next_sequence + .fetch_max(sequence.saturating_add(1), Ordering::Relaxed); + } + + fn apply_replica_update( + &self, + session_id: String, + target: AffinityTarget, + version: AffinityVersion, + ) -> ReplicaApplyOutcome { + if session_id.len() > self.max_session_id_bytes { + return ReplicaApplyOutcome::RejectedSessionId; + } + self.observe_replica_sequence(version.sequence); + + let now = Instant::now(); + match self.entries.entry(session_id) { + Entry::Vacant(entry) => { + if !self.reserve_entry() { + return ReplicaApplyOutcome::RejectedCapacity; + } + let revision = self.next_revision.fetch_add(1, Ordering::Relaxed); + entry.insert(AffinityEntry::Bound { + target, + revision, + version, + active_leases: 0, + idle_deadline: now + self.ttl, + }); + ReplicaApplyOutcome::Inserted + } + Entry::Occupied(mut entry) => match entry.get_mut() { + AffinityEntry::Initializing { .. } => ReplicaApplyOutcome::IgnoredInitializing, + AffinityEntry::Bound { + active_leases, + idle_deadline, + .. + } if *active_leases == 0 && *idle_deadline <= now => { + let revision = self.next_revision.fetch_add(1, Ordering::Relaxed); + *entry.get_mut() = AffinityEntry::Bound { + target, + revision, + version, + active_leases: 0, + idle_deadline: now + self.ttl, + }; + ReplicaApplyOutcome::ReplacedExpired + } + AffinityEntry::Bound { + target: existing, + version: existing_version, + idle_deadline, + .. + } if *existing == target && version >= *existing_version => { + *existing_version = version; + *idle_deadline = now + self.ttl; + ReplicaApplyOutcome::Refreshed + } + AffinityEntry::Bound { + target: existing, + version: existing_version, + idle_deadline, + .. + } if version > *existing_version => { + *existing = target; + *existing_version = version; + *idle_deadline = now + self.ttl; + ReplicaApplyOutcome::ReplacedNewer + } + AffinityEntry::Bound { .. } => ReplicaApplyOutcome::IgnoredConflict, + }, + } + } +} + +/// A session this request is the first to bind. Dropping it uncommitted +/// releases the slot and wakes waiters so they re-acquire. +pub struct AffinityInitialization { + table: Weak, + session_id: String, + revision: u64, + notify: Arc, + requested_target: Option, + active: bool, +} + +impl AffinityInitialization { + pub fn session_id(&self) -> &str { + &self.session_id + } + + /// Bind the session to the worker the request was dispatched to. + pub fn commit(mut self, target: AffinityTarget) -> Result { + validate_bound_target(&self.session_id, target, self.requested_target)?; + let Some(inner) = self.table.upgrade() else { + return Err(AffinityError::Dropped); + }; + let Some(mut entry) = inner.entries.get_mut(&self.session_id) else { + return Err(AffinityError::InvalidArgument( + "session affinity initialization was cancelled".to_string(), + )); + }; + if !matches!( + entry.value(), + AffinityEntry::Initializing { revision, .. } if *revision == self.revision + ) { + return Err(AffinityError::InvalidArgument( + "session affinity initialization changed".to_string(), + )); + } + let version = inner.next_version(); + *entry = AffinityEntry::Bound { + target, + revision: self.revision, + version, + active_leases: 1, + idle_deadline: Instant::now() + inner.ttl, + }; + drop(entry); + self.active = false; + self.notify.notify_waiters(); + Ok(AffinityLease { + table: Arc::downgrade(&inner), + session_id: self.session_id.clone(), + revision: self.revision, + version, + active: true, + }) + } +} + +impl Drop for AffinityInitialization { + fn drop(&mut self) { + if !self.active { + return; + } + let Some(inner) = self.table.upgrade() else { + return; + }; + let removed = inner.entries.remove_if(&self.session_id, |_, entry| { + matches!( + entry, + AffinityEntry::Initializing { revision, .. } if *revision == self.revision + ) + }); + if removed.is_some() { + inner.entry_count.fetch_sub(1, Ordering::Relaxed); + } + self.notify.notify_waiters(); + } +} + +/// A live use of a bound session. Dropping it releases the use and refreshes +/// the idle deadline. +pub struct AffinityLease { + table: Weak, + session_id: String, + revision: u64, + version: AffinityVersion, + active: bool, +} + +impl AffinityLease { + pub fn session_id(&self) -> &str { + &self.session_id + } + + /// Publish the binding to replicas. + pub fn publish(&self, target: AffinityTarget) { + if let Some(inner) = self.table.upgrade() { + inner.publish_replica_update(&self.session_id, target, self.version); + } + } + + /// Move the binding from `expected` to `target` if nothing changed it since + /// this lease was taken. + pub fn rebind(&mut self, expected: AffinityTarget, target: AffinityTarget) -> bool { + let Some(inner) = self.table.upgrade() else { + return false; + }; + let Some(mut entry) = inner.entries.get_mut(&self.session_id) else { + return false; + }; + let AffinityEntry::Bound { + target: current, + revision, + version, + .. + } = entry.value_mut() + else { + return false; + }; + if *revision != self.revision || *version != self.version || *current != expected { + return false; + } + let next_version = inner.next_version(); + *current = target; + *version = next_version; + self.version = next_version; + true + } + + /// Soft mode: where a bound session that was dispatched to `dispatched` + /// should now point. A worker-only binding stays worker-only; a ranked + /// binding follows the dispatched rank when there is one. + pub fn rebound_target(bound: AffinityTarget, dispatched: AffinityTarget) -> AffinityTarget { + match bound.dp_rank { + None => AffinityTarget::new(dispatched.worker_id, None), + Some(_) => dispatched.dp_rank.map_or(bound, |dp_rank| { + AffinityTarget::new(dispatched.worker_id, Some(dp_rank)) + }), + } + } + + fn release(&mut self) { + if !self.active { + return; + } + self.active = false; + let Some(inner) = self.table.upgrade() else { + return; + }; + let (target, version) = { + let Some(mut entry) = inner.entries.get_mut(&self.session_id) else { + return; + }; + let AffinityEntry::Bound { + target, + revision, + version, + active_leases, + idle_deadline, + } = entry.value_mut() + else { + return; + }; + if *revision != self.revision || *active_leases == 0 { + return; + } + *active_leases -= 1; + if *version != self.version { + return; + } + *idle_deadline = Instant::now() + inner.ttl; + (*target, *version) + }; + inner.publish_replica_update(&self.session_id, target, version); + } + + /// Drop the binding this lease holds (the bound worker is gone). A binding + /// already replaced by a newer version is left alone and this use released. + pub fn invalidate(&mut self) { + if !self.active { + return; + } + let Some(inner) = self.table.upgrade() else { + self.active = false; + return; + }; + let removed = inner.entries.remove_if(&self.session_id, |_, entry| { + matches!( + entry, + AffinityEntry::Bound { revision, version, .. } + if *revision == self.revision && *version == self.version + ) + }); + match removed { + Some((_, AffinityEntry::Bound { target, .. })) => { + tracing::debug!( + session_id = %self.session_id, + worker_id = target.worker_id, + dp_rank = ?target.dp_rank, + "invalidated current session affinity binding" + ); + self.active = false; + inner.entry_count.fetch_sub(1, Ordering::Relaxed); + } + Some((_, AffinityEntry::Initializing { .. })) => { + unreachable!("bound lease removed an initializing entry") + } + None => self.release(), + } + } +} + +impl Drop for AffinityLease { + fn drop(&mut self) { + self.release(); + } +} + +/// A bound session must agree with the target the request explicitly asked for. +pub fn validate_bound_target( + session_id: &str, + bound: AffinityTarget, + requested: Option, +) -> Result<(), AffinityError> { + let Some(requested) = requested else { + return Ok(()); + }; + if bound.worker_id != requested.worker_id { + return Err(AffinityError::InvalidArgument(format!( + "session {session_id} is bound to worker {}, not {}", + bound.worker_id, requested.worker_id + ))); + } + match (bound.dp_rank, requested.dp_rank) { + (Some(bound), Some(requested)) if bound != requested => { + Err(AffinityError::InvalidArgument(format!( + "session {session_id} is bound to DP rank {bound}, not {requested}" + ))) + } + (None, Some(requested)) => Err(AffinityError::InvalidArgument(format!( + "session {session_id} has worker-only affinity and cannot add DP rank {requested}" + ))), + _ => Ok(()), + } +} + +/// Hard mode: a bound session must have been dispatched where it is bound. +pub fn validate_dispatch_target( + session_id: &str, + bound: AffinityTarget, + dispatched: AffinityTarget, +) -> Result<(), AffinityError> { + if bound.worker_id != dispatched.worker_id { + return Err(AffinityError::InvalidArgument(format!( + "session {session_id} is bound to worker {}, not {}", + bound.worker_id, dispatched.worker_id + ))); + } + if let Some(bound_rank) = bound.dp_rank { + match dispatched.dp_rank { + Some(dispatched_rank) if dispatched_rank == bound_rank => {} + Some(dispatched_rank) => { + return Err(AffinityError::InvalidArgument(format!( + "session {session_id} is bound to DP rank {bound_rank}, not {dispatched_rank}" + ))); + } + None => { + return Err(AffinityError::InvalidArgument(format!( + "session {session_id} is bound to DP rank {bound_rank}, but dispatch did not select a DP rank" + ))); + } + } + } + Ok(()) +} diff --git a/lib/kv-router/src/services/selection/mod.rs b/lib/kv-router/src/services/selection/mod.rs index 63ea5504c309..15ab516d6edd 100644 --- a/lib/kv-router/src/services/selection/mod.rs +++ b/lib/kv-router/src/services/selection/mod.rs @@ -6,6 +6,7 @@ //! The service owns worker selection and reservation state, but never forwards //! model requests and never owns model responses. +pub mod affinity; mod catalog; mod core; mod error; diff --git a/lib/llm/src/kv_router/routing_host/tests.rs b/lib/llm/src/kv_router/routing_host/tests.rs index d8399eb9a616..e31eb37335d0 100644 --- a/lib/llm/src/kv_router/routing_host/tests.rs +++ b/lib/llm/src/kv_router/routing_host/tests.rs @@ -376,7 +376,9 @@ async fn builtin_hard_affinity_ignores_local_inhibition() { }; drop( initializer - .commit(AffinityTarget::worker(worker_id)) + .commit(crate::session_affinity::to_table(AffinityTarget::worker( + worker_id, + ))) .unwrap(), ); @@ -649,7 +651,9 @@ async fn builtin_direct_fallback_stays_disabled_for_affinity() { }; drop( initializer - .commit(AffinityTarget::worker(stale_worker)) + .commit(crate::session_affinity::to_table(AffinityTarget::worker( + stale_worker, + ))) .unwrap(), ); assert!( @@ -1877,7 +1881,11 @@ async fn session_affinity_post_selection_failures_preserve_binding() { else { panic!("first request must initialize"); }; - drop(initializer.commit(original_target).unwrap()); + drop( + initializer + .commit(crate::session_affinity::to_table(original_target)) + .unwrap(), + ); let operation = Some(affinity.acquire(&session_id, None).await.unwrap()); drop(operation); @@ -1915,7 +1923,11 @@ async fn session_affinity_existing_selection_cancellation_preserves_binding_with else { panic!("first request must initialize"); }; - drop(initializer.commit(original_target).unwrap()); + drop( + initializer + .commit(crate::session_affinity::to_table(original_target)) + .unwrap(), + ); let controller = Controller::new("cancelled-selection-request".to_string()); controller.stop(); @@ -1980,7 +1992,11 @@ async fn bind_affinity_target( else { panic!("first request must initialize"); }; - drop(initializer.commit(target).unwrap()); + drop( + initializer + .commit(crate::session_affinity::to_table(target)) + .unwrap(), + ); } #[tokio::test] @@ -2148,7 +2164,11 @@ async fn migration_exclusion_preserves_hard_affinity_without_widening_or_escapin else { panic!("first request must initialize"); }; - drop(initializer.commit(original_target).unwrap()); + drop( + initializer + .commit(crate::session_affinity::to_table(original_target)) + .unwrap(), + ); let mut retry_input = request(); retry_input.routing_mut().allowed_worker_ids = Some(HashSet::from([7, 8])); diff --git a/lib/llm/src/session_affinity/coordinator.rs b/lib/llm/src/session_affinity/coordinator.rs index ab2266f5b3a1..2fe9ff943528 100644 --- a/lib/llm/src/session_affinity/coordinator.rs +++ b/lib/llm/src/session_affinity/coordinator.rs @@ -1,32 +1,33 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +//! The frontend's view of session affinity: the shared `dynamo_kv_router` +//! [`SessionAffinity`] table plus what only the frontend knows about, namely +//! request contexts (to abandon a wait when the client cancels), pipeline +//! response streams (a lease lives until the stream ends), and the runtime +//! event plane that replicates bindings between frontends. + use std::{ pin::Pin, - sync::{ - Arc, OnceLock, Weak, - atomic::{AtomicU64, AtomicUsize, Ordering}, - }, + sync::{Arc, OnceLock}, task::{Context, Poll}, - time::{Duration, SystemTime, UNIX_EPOCH}, + time::Duration, }; -use dashmap::{DashMap, mapref::entry::Entry}; +use dynamo_kv_router::services::selection::affinity::{ + AcquireStep, AffinityError, AffinityInitialization, AffinityLease, SessionAffinity, + validate_dispatch_target, +}; use dynamo_runtime::{ engine::{AsyncEngineContext, AsyncEngineContextProvider}, error::{DynamoError, ErrorType}, pipeline::{Error, ManyOut, ResponseStream}, }; use futures::Stream; -use tokio::{sync::Notify, time::Instant}; -use tokio_util::sync::CancellationToken; #[cfg(test)] use super::replica_sync::SessionAffinityUpdate; -use super::{ - LlmResponse, MAX_SESSION_AFFINITY_ENTRIES, MAX_SESSION_AFFINITY_ID_BYTES, - MAX_SESSION_AFFINITY_TTL_SECS, SessionAffinityMode, replica_sync::ReplicaSyncRuntime, -}; +use super::{LlmResponse, SessionAffinityMode, replica_sync::ReplicaSyncRuntime}; use crate::{ preprocessor::PreprocessedRequest, protocols::common::{ @@ -35,165 +36,73 @@ use crate::{ }, }; +#[cfg(test)] +pub(super) use dynamo_kv_router::services::selection::affinity::ReplicaApplyOutcome; + +/// The pipeline's routing target; the table's `AffinityTarget` has the same +/// shape and the two convert at this boundary. pub type AffinityTarget = dynamo_runtime::pipeline::RouteTarget; -#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)] -pub(super) struct AffinityVersion { - pub(super) sequence: u64, - pub(super) writer_id: u64, +type TableTarget = dynamo_kv_router::services::selection::affinity::AffinityTarget; + +/// The table-side form of a pipeline routing target. +pub(crate) fn to_table(target: AffinityTarget) -> TableTarget { + TableTarget::new(target.worker_id, target.dp_rank) } -enum AffinityEntry { - Initializing { - revision: u64, - notify: Arc, - }, - Bound { - target: AffinityTarget, - revision: u64, - version: AffinityVersion, - active_leases: usize, - idle_deadline: Instant, - }, +fn from_table(target: TableTarget) -> AffinityTarget { + AffinityTarget::new(target.worker_id, target.dp_rank) } -pub(super) struct AffinityCoordinatorInner { - entries: DashMap, - ttl: Duration, - max_entries: usize, - max_session_id_bytes: usize, - entry_count: AtomicUsize, - next_revision: AtomicU64, - next_sequence: AtomicU64, - pub(super) writer_id: AtomicU64, - cancel: CancellationToken, +struct Inner { + table: SessionAffinity, replica: OnceLock, - #[cfg(test)] - reaper_started: Arc, - #[cfg(test)] - waiter_observed: Arc, } -impl Drop for AffinityCoordinatorInner { +impl Drop for Inner { fn drop(&mut self) { - self.cancel.cancel(); if let Some(replica) = self.replica.get_mut() { replica.shutdown_now(); } } } -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub(super) enum ReplicaApplyOutcome { - Inserted, - Refreshed, - ReplacedExpired, - ReplacedNewer, - IgnoredInitializing, - IgnoredConflict, - RejectedSessionId, - RejectedCapacity, -} - #[derive(Clone)] pub struct AffinityCoordinator { - inner: Arc, + inner: Arc, } impl AffinityCoordinator { pub fn new(ttl: Duration) -> Result { - Self::new_with_limits( - ttl, - MAX_SESSION_AFFINITY_ENTRIES, - MAX_SESSION_AFFINITY_ID_BYTES, - ) + Ok(Self::wrap( + SessionAffinity::new(ttl).map_err(affinity_error)?, + )) } - fn new_with_limits( - ttl: Duration, - max_entries: usize, - max_session_id_bytes: usize, - ) -> Result { - if !(Duration::from_secs(1)..=Duration::from_secs(MAX_SESSION_AFFINITY_TTL_SECS)) - .contains(&ttl) - { - return Err(invalid_argument(format!( - "session affinity TTL must be between 1 and {MAX_SESSION_AFFINITY_TTL_SECS} seconds" - ))); + fn wrap(table: SessionAffinity) -> Self { + Self { + inner: Arc::new(Inner { + table, + replica: OnceLock::new(), + }), } - let inner = Arc::new(AffinityCoordinatorInner { - entries: DashMap::new(), - ttl, - max_entries, - max_session_id_bytes, - entry_count: AtomicUsize::new(0), - next_revision: AtomicU64::new(1), - next_sequence: AtomicU64::new( - SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_default() - .as_nanos() as u64, - ), - writer_id: AtomicU64::new(0), - cancel: CancellationToken::new(), - replica: OnceLock::new(), - #[cfg(test)] - reaper_started: Arc::new(Notify::new()), - #[cfg(test)] - waiter_observed: Arc::new(Notify::new()), - }); - Self::spawn_reaper(&inner); - tracing::info!( - ttl_secs = ttl.as_secs(), - max_entries, - "session affinity enabled" - ); - Ok(Self { inner }) - } - - fn spawn_reaper(inner: &Arc) { - let weak = Arc::downgrade(inner); - let cancel = inner.cancel.clone(); - let period = inner.ttl.min(Duration::from_secs(30)); - #[cfg(test)] - let reaper_started = inner.reaper_started.clone(); - tokio::spawn(async move { - #[cfg(test)] - reaper_started.notify_one(); - loop { - tokio::select! { - _ = cancel.cancelled() => return, - _ = tokio::time::sleep(period) => {} - } - let Some(inner) = weak.upgrade() else { - return; - }; - let now = Instant::now(); - let mut removed = 0; - inner.entries.retain(|_, entry| { - let retain = !matches!( - entry, - AffinityEntry::Bound { - active_leases: 0, - idle_deadline, - .. - } if *idle_deadline <= now - ); - removed += usize::from(!retain); - retain - }); - inner.entry_count.fetch_sub(removed, Ordering::Relaxed); - } - }); } pub(crate) async fn enable_replica_sync( &self, client: dynamo_runtime::component::Client, ) -> Result<(), Error> { - let replica = - ReplicaSyncRuntime::start(client, Arc::downgrade(&self.inner), &self.inner.cancel) - .await?; + let (replica, router_id) = + ReplicaSyncRuntime::start(client, self.inner.table.downgrade()).await?; + if !self + .inner + .table + .enable_replication(router_id, replica.sink()) + { + return Err(anyhow::anyhow!( + "session affinity replica sync already enabled" + )); + } self.inner .replica .set(replica) @@ -231,102 +140,33 @@ impl AffinityCoordinator { requested_target: Option, request_context: Option<&dyn AsyncEngineContext>, ) -> Result { - self.validate_session_id(session_id)?; - let session_id = session_id.as_str().to_string(); - + let requested = requested_target.map(to_table); loop { - let now = Instant::now(); - match self.inner.entries.entry(session_id.clone()) { - Entry::Vacant(entry) => { - self.reserve_entry()?; - tracing::debug!( - session_id = %session_id, - "session affinity miss: new session, pinning after worker selection" - ); - return Ok(AffinityAcquire::Initialize(entry.insert_initializing( - &self.inner, - session_id, - requested_target, - ))); + match self + .inner + .table + .try_acquire(session_id.as_str(), requested) + .map_err(affinity_error)? + { + AcquireStep::Initialize(initialization) => { + return Ok(AffinityAcquire::Initialize(initialization)); + } + AcquireStep::Bound { target, lease } => { + return Ok(AffinityAcquire::Bound { + target: from_table(target), + lease, + }); } - Entry::Occupied(mut entry) => match entry.get_mut() { - AffinityEntry::Initializing { notify, .. } => { - #[cfg(test)] - self.inner.waiter_observed.notify_one(); - let notified = notify.clone().notified_owned(); - tokio::pin!(notified); - notified.as_mut().enable(); - drop(entry); - if let Some(context) = request_context { - tokio::select! { - biased; - _ = context.stopped() => { - return Err(cancelled(context.id())); - } - _ = context.killed() => { - return Err(cancelled(context.id())); - } - _ = notified => {} - } - } else { - notified.await; + AcquireStep::Wait(notified) => match request_context { + Some(context) => { + tokio::select! { + biased; + _ = context.stopped() => return Err(cancelled(context.id())), + _ = context.killed() => return Err(cancelled(context.id())), + _ = notified => {} } } - AffinityEntry::Bound { - target: _, - revision, - active_leases, - idle_deadline, - .. - } if *active_leases == 0 && *idle_deadline <= now => { - tracing::debug!( - session_id = %session_id, - "session affinity miss: pin expired (idle past TTL), re-selecting worker" - ); - let revision = self.inner.next_revision.fetch_add(1, Ordering::Relaxed); - let notify = Arc::new(Notify::new()); - *entry.get_mut() = AffinityEntry::Initializing { - revision, - notify: notify.clone(), - }; - drop(entry); - return Ok(AffinityAcquire::Initialize(AffinityInitialization { - coordinator: Arc::downgrade(&self.inner), - session_id, - revision, - notify, - requested_target, - active: true, - })); - } - AffinityEntry::Bound { - target, - revision, - version, - active_leases, - .. - } => { - validate_bound_target(&session_id, *target, requested_target)?; - tracing::debug!( - session_id = %session_id, - worker_id = target.worker_id, - dp_rank = ?target.dp_rank, - active_leases = *active_leases + 1, - "session affinity hit: reusing pinned worker" - ); - *active_leases += 1; - let lease = AffinityLease { - coordinator: Arc::downgrade(&self.inner), - session_id, - revision: *revision, - version: *version, - active: true, - }; - return Ok(AffinityAcquire::Bound { - target: *target, - lease, - }); - } + None => notified.await, }, } } @@ -337,73 +177,48 @@ impl AffinityCoordinator { session_id: &SessionAffinityId, requested_target: Option, ) -> Result, Error> { - self.validate_session_id(session_id)?; - let Some(entry) = self.inner.entries.get(session_id.as_str()) else { - return Ok(None); - }; - let AffinityEntry::Bound { - target, - active_leases, - idle_deadline, - .. - } = entry.value() - else { - return Ok(None); - }; - if *active_leases == 0 && *idle_deadline <= Instant::now() { - return Ok(None); - } - validate_bound_target(session_id.as_str(), *target, requested_target)?; - tracing::debug!( - session_id = %session_id.as_str(), - worker_id = target.worker_id, - dp_rank = ?target.dp_rank, - "session affinity hit: reusing pinned worker" - ); - - Ok(Some(*target)) + self.inner + .table + .query_target(session_id.as_str(), requested_target.map(to_table)) + .map(|target| target.map(from_table)) + .map_err(affinity_error) } #[cfg(test)] pub(super) fn entry_count(&self) -> usize { - self.inner.entry_count.load(Ordering::Relaxed) + self.inner.table.entry_count() } #[cfg(test)] - pub(super) fn cancellation_token(&self) -> CancellationToken { - self.inner.cancel.clone() + pub(super) fn cancellation_token(&self) -> tokio_util::sync::CancellationToken { + self.inner.table.cancellation_token() } #[cfg(test)] pub(super) async fn wait_for_reaper(&self) { - self.inner.reaper_started.notified().await; + self.inner.table.wait_for_reaper().await; } #[cfg(test)] pub(crate) async fn wait_for_initializing_waiter(&self) { - self.inner.waiter_observed.notified().await; + self.inner.table.wait_for_initializing_waiter().await; } #[cfg(test)] pub(super) fn expire_for_test(&self, session_id: &SessionAffinityId) { - let Some(mut entry) = self.inner.entries.get_mut(session_id.as_str()) else { - panic!("session affinity entry missing"); - }; - let AffinityEntry::Bound { - active_leases, - idle_deadline, - .. - } = entry.value_mut() - else { - panic!("session affinity entry is not bound"); - }; - assert_eq!(*active_leases, 0); - *idle_deadline = Instant::now(); + self.inner.table.expire_for_test(session_id.as_str()); } #[cfg(test)] pub(super) fn with_test_limits(max_entries: usize, max_session_id_bytes: usize) -> Self { - Self::new_with_limits(Duration::from_secs(10), max_entries, max_session_id_bytes).unwrap() + Self::wrap( + SessionAffinity::new_with_limits( + Duration::from_secs(10), + max_entries, + max_session_id_bytes, + ) + .unwrap(), + ) } #[cfg(test)] @@ -412,8 +227,13 @@ impl AffinityCoordinator { router_id: u64, capacity: usize, ) -> tokio::sync::mpsc::Receiver { - self.inner.writer_id.store(router_id, Ordering::Relaxed); let (replica, rx) = ReplicaSyncRuntime::for_test(capacity); + assert!( + self.inner + .table + .enable_replication(router_id, replica.sink()), + "session affinity test replica already enabled" + ); self.inner .replica .set(replica) @@ -422,13 +242,17 @@ impl AffinityCoordinator { } #[cfg(test)] - pub(super) fn downgrade_for_test(&self) -> Weak { - Arc::downgrade(&self.inner) + pub(super) fn next_version_for_test( + &self, + ) -> dynamo_kv_router::services::selection::affinity::AffinityVersion { + self.inner.table.next_version() } #[cfg(test)] - pub(super) fn next_version_for_test(&self) -> AffinityVersion { - self.inner.next_version() + pub(super) fn table_for_test( + &self, + ) -> dynamo_kv_router::services::selection::affinity::WeakSessionAffinity { + self.inner.table.downgrade() } #[cfg(test)] @@ -437,14 +261,7 @@ impl AffinityCoordinator { session_id: impl Into, target: AffinityTarget, ) -> ReplicaApplyOutcome { - self.inner.apply_replica_update( - session_id.into(), - target, - AffinityVersion { - sequence: 0, - writer_id: 0, - }, - ) + self.apply_versioned_replica_update_for_test(session_id, target, 0, 0) } #[cfg(test)] @@ -455,168 +272,16 @@ impl AffinityCoordinator { sequence: u64, writer_id: u64, ) -> ReplicaApplyOutcome { - self.inner.apply_replica_update( + use dynamo_kv_router::services::selection::affinity::AffinityVersion; + self.inner.table.apply_replica_update( session_id.into(), - target, + to_table(target), AffinityVersion { sequence, writer_id, }, ) } - - fn validate_session_id(&self, session_id: &SessionAffinityId) -> Result<(), Error> { - if session_id.as_str().len() > self.inner.max_session_id_bytes { - return Err(invalid_argument(format!( - "session affinity ID must not exceed {} bytes", - self.inner.max_session_id_bytes - ))); - } - Ok(()) - } - - fn reserve_entry(&self) -> Result<(), Error> { - self.inner - .reserve_entry() - .then_some(()) - .ok_or_else(|| resource_exhausted("session affinity entry limit reached")) - } -} - -impl AffinityCoordinatorInner { - fn reserve_entry(&self) -> bool { - self.entry_count - .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |count| { - (count < self.max_entries).then_some(count + 1) - }) - .is_ok() - } - - fn publish_replica_update( - &self, - session_id: &str, - target: AffinityTarget, - version: AffinityVersion, - ) { - if let Some(replica) = self.replica.get() { - replica.publish(session_id, target, version); - } - } - - fn next_version(&self) -> AffinityVersion { - AffinityVersion { - sequence: self.next_sequence.fetch_add(1, Ordering::Relaxed), - writer_id: self.writer_id.load(Ordering::Relaxed), - } - } - - pub(super) fn observe_replica_sequence(&self, sequence: u64) { - self.next_sequence - .fetch_max(sequence.saturating_add(1), Ordering::Relaxed); - } - - pub(super) fn apply_replica_update( - &self, - session_id: String, - target: AffinityTarget, - version: AffinityVersion, - ) -> ReplicaApplyOutcome { - if session_id.len() > self.max_session_id_bytes { - return ReplicaApplyOutcome::RejectedSessionId; - } - self.observe_replica_sequence(version.sequence); - - let now = Instant::now(); - match self.entries.entry(session_id) { - Entry::Vacant(entry) => { - if !self.reserve_entry() { - return ReplicaApplyOutcome::RejectedCapacity; - } - let revision = self.next_revision.fetch_add(1, Ordering::Relaxed); - entry.insert(AffinityEntry::Bound { - target, - revision, - version, - active_leases: 0, - idle_deadline: now + self.ttl, - }); - ReplicaApplyOutcome::Inserted - } - Entry::Occupied(mut entry) => match entry.get_mut() { - AffinityEntry::Initializing { .. } => ReplicaApplyOutcome::IgnoredInitializing, - AffinityEntry::Bound { - active_leases, - idle_deadline, - .. - } if *active_leases == 0 && *idle_deadline <= now => { - let revision = self.next_revision.fetch_add(1, Ordering::Relaxed); - *entry.get_mut() = AffinityEntry::Bound { - target, - revision, - version, - active_leases: 0, - idle_deadline: now + self.ttl, - }; - ReplicaApplyOutcome::ReplacedExpired - } - AffinityEntry::Bound { - target: existing, - version: existing_version, - idle_deadline, - .. - } if *existing == target && version >= *existing_version => { - *existing_version = version; - *idle_deadline = now + self.ttl; - ReplicaApplyOutcome::Refreshed - } - AffinityEntry::Bound { - target: existing, - version: existing_version, - idle_deadline, - .. - } if version > *existing_version => { - *existing = target; - *existing_version = version; - *idle_deadline = now + self.ttl; - ReplicaApplyOutcome::ReplacedNewer - } - AffinityEntry::Bound { .. } => ReplicaApplyOutcome::IgnoredConflict, - }, - } - } -} - -trait VacantEntryExt { - fn insert_initializing( - self, - inner: &Arc, - session_id: String, - requested_target: Option, - ) -> AffinityInitialization; -} - -impl<'a> VacantEntryExt for dashmap::mapref::entry::VacantEntry<'a, String, AffinityEntry> { - fn insert_initializing( - self, - inner: &Arc, - session_id: String, - requested_target: Option, - ) -> AffinityInitialization { - let revision = inner.next_revision.fetch_add(1, Ordering::Relaxed); - let notify = Arc::new(Notify::new()); - self.insert(AffinityEntry::Initializing { - revision, - notify: notify.clone(), - }); - AffinityInitialization { - coordinator: Arc::downgrade(inner), - session_id, - revision, - notify, - requested_target, - active: true, - } - } } pub(crate) enum AffinityAcquire { @@ -635,39 +300,38 @@ impl AffinityAcquire { } } + /// Bind (or confirm) the session to where the request was dispatched and + /// hold the lease until `stream` ends. pub(crate) fn into_stream( self, dispatched_target: AffinityTarget, stream: ManyOut, mode: SessionAffinityMode, ) -> Result, Error> { + let dispatched = to_table(dispatched_target); match self { Self::Initialize(initialization) => { - let lease = initialization.commit(dispatched_target)?; - lease.publish(dispatched_target); - Ok(lease.into_stream(stream)) + let lease = initialization.commit(dispatched).map_err(affinity_error)?; + lease.publish(dispatched); + Ok(tracked_stream(lease, stream)) } Self::Bound { target, mut lease } => { + let bound = to_table(target); if mode == SessionAffinityMode::Soft { - let rebound_target = match target.dp_rank { - None => AffinityTarget::worker(dispatched_target.worker_id), - Some(_) => dispatched_target.dp_rank.map_or(target, |dp_rank| { - AffinityTarget::new(dispatched_target.worker_id, Some(dp_rank)) - }), - }; - if target == rebound_target { - lease.publish(target); - } else if lease.rebind(target, rebound_target) { - lease.publish(rebound_target); + let rebound = AffinityLease::rebound_target(bound, dispatched); + if bound == rebound { + lease.publish(bound); + } else if lease.rebind(bound, rebound) { + lease.publish(rebound); } - return Ok(lease.into_stream(stream)); + return Ok(tracked_stream(lease, stream)); } - if let Err(error) = validate_dispatch_target("session", target, dispatched_target) { + if let Err(error) = validate_dispatch_target("session", bound, dispatched) { lease.invalidate(); - return Err(error); + return Err(affinity_error(error)); } - lease.publish(target); - Ok(lease.into_stream(stream)) + lease.publish(bound); + Ok(tracked_stream(lease, stream)) } } } @@ -679,199 +343,18 @@ impl AffinityAcquire { } } -pub(crate) struct AffinityInitialization { - coordinator: Weak, - session_id: String, - revision: u64, - notify: Arc, - requested_target: Option, - active: bool, -} - -impl AffinityInitialization { - pub(crate) fn commit(mut self, target: AffinityTarget) -> Result { - validate_bound_target(&self.session_id, target, self.requested_target)?; - let Some(inner) = self.coordinator.upgrade() else { - return Err(anyhow::anyhow!("session affinity coordinator dropped")); - }; - let Some(mut entry) = inner.entries.get_mut(&self.session_id) else { - return Err(invalid_argument( - "session affinity initialization was cancelled", - )); - }; - if !matches!( - entry.value(), - AffinityEntry::Initializing { revision, .. } if *revision == self.revision - ) { - return Err(invalid_argument("session affinity initialization changed")); - } - let version = inner.next_version(); - *entry = AffinityEntry::Bound { - target, - revision: self.revision, - version, - active_leases: 1, - idle_deadline: Instant::now() + inner.ttl, - }; - drop(entry); - self.active = false; - self.notify.notify_waiters(); - Ok(AffinityLease { - coordinator: Arc::downgrade(&inner), - session_id: self.session_id.clone(), - revision: self.revision, - version, - active: true, - }) - } -} - -impl Drop for AffinityInitialization { - fn drop(&mut self) { - if !self.active { - return; - } - let Some(inner) = self.coordinator.upgrade() else { - return; - }; - let removed = inner.entries.remove_if(&self.session_id, |_, entry| { - matches!( - entry, - AffinityEntry::Initializing { revision, .. } if *revision == self.revision - ) - }); - if removed.is_some() { - inner.entry_count.fetch_sub(1, Ordering::Relaxed); - } - self.notify.notify_waiters(); - } -} - -pub(crate) struct AffinityLease { - coordinator: Weak, - session_id: String, - revision: u64, - version: AffinityVersion, - active: bool, -} - -impl AffinityLease { - fn publish(&self, target: AffinityTarget) { - if let Some(inner) = self.coordinator.upgrade() { - inner.publish_replica_update(&self.session_id, target, self.version); - } - } - - fn rebind(&mut self, expected: AffinityTarget, target: AffinityTarget) -> bool { - let Some(inner) = self.coordinator.upgrade() else { - return false; - }; - let Some(mut entry) = inner.entries.get_mut(&self.session_id) else { - return false; - }; - let AffinityEntry::Bound { - target: current, - revision, - version, - .. - } = entry.value_mut() - else { - return false; - }; - if *revision != self.revision || *version != self.version || *current != expected { - return false; - } - let next_version = inner.next_version(); - *current = target; - *version = next_version; - self.version = next_version; - true - } - - pub(crate) fn into_stream(self, stream: ManyOut) -> ManyOut { - let context = stream.context(); - ResponseStream::new( - Box::pin(AffinityTrackedStream { - stream, - lease: Some(self), - }), - context, - ) - } - - fn release(&mut self) { - if !self.active { - return; - } - self.active = false; - let Some(inner) = self.coordinator.upgrade() else { - return; - }; - let (target, version) = { - let Some(mut entry) = inner.entries.get_mut(&self.session_id) else { - return; - }; - let AffinityEntry::Bound { - target, - revision, - version, - active_leases, - idle_deadline, - } = entry.value_mut() - else { - return; - }; - if *revision != self.revision || *active_leases == 0 { - return; - } - *active_leases -= 1; - if *version != self.version { - return; - } - *idle_deadline = Instant::now() + inner.ttl; - (*target, *version) - }; - inner.publish_replica_update(&self.session_id, target, version); - } - - fn invalidate(&mut self) { - if !self.active { - return; - } - let Some(inner) = self.coordinator.upgrade() else { - self.active = false; - return; - }; - let removed = inner.entries.remove_if(&self.session_id, |_, entry| { - matches!( - entry, - AffinityEntry::Bound { revision, version, .. } - if *revision == self.revision && *version == self.version - ) - }); - match removed { - Some((_, AffinityEntry::Bound { target, .. })) => { - tracing::debug!( - session_id = %self.session_id, - worker_id = target.worker_id, - dp_rank = ?target.dp_rank, - "invalidated current session affinity binding" - ); - self.active = false; - inner.entry_count.fetch_sub(1, Ordering::Relaxed); - } - Some((_, AffinityEntry::Initializing { .. })) => { - unreachable!("bound lease removed an initializing entry") - } - None => self.release(), - } - } -} - -impl Drop for AffinityLease { - fn drop(&mut self) { - self.release(); - } +pub(super) fn tracked_stream( + lease: AffinityLease, + stream: ManyOut, +) -> ManyOut { + let context = stream.context(); + ResponseStream::new( + Box::pin(AffinityTrackedStream { + stream, + lease: Some(lease), + }), + context, + ) } struct AffinityTrackedStream { @@ -927,62 +410,16 @@ pub fn explicit_target( Ok(worker_id.map(|worker_id| AffinityTarget { worker_id, dp_rank })) } -fn validate_bound_target( - session_id: &str, - bound: AffinityTarget, - requested: Option, -) -> Result<(), Error> { - let Some(requested) = requested else { - return Ok(()); - }; - if bound.worker_id != requested.worker_id { - return Err(invalid_argument(format!( - "session {session_id} is bound to worker {}, not {}", - bound.worker_id, requested.worker_id - ))); - } - match (bound.dp_rank, requested.dp_rank) { - (Some(bound), Some(requested)) if bound != requested => Err(invalid_argument(format!( - "session {session_id} is bound to DP rank {bound}, not {requested}" - ))), - (None, Some(requested)) => Err(invalid_argument(format!( - "session {session_id} has worker-only affinity and cannot add DP rank {requested}" - ))), - _ => Ok(()), - } -} - -/// Validates that a request was dispatched within an existing session binding. -/// -/// Unlike an explicit requested target, a dispatch target may add a DP rank to a worker-only -/// binding because load-aware scheduling chooses that rank for this request only. -fn validate_dispatch_target( - session_id: &str, - bound: AffinityTarget, - dispatched: AffinityTarget, -) -> Result<(), Error> { - if bound.worker_id != dispatched.worker_id { - return Err(invalid_argument(format!( - "session {session_id} is bound to worker {}, not {}", - bound.worker_id, dispatched.worker_id - ))); - } - if let Some(bound_rank) = bound.dp_rank { - match dispatched.dp_rank { - Some(dispatched_rank) if dispatched_rank == bound_rank => {} - Some(dispatched_rank) => { - return Err(invalid_argument(format!( - "session {session_id} is bound to DP rank {bound_rank}, not {dispatched_rank}" - ))); - } - None => { - return Err(invalid_argument(format!( - "session {session_id} is bound to DP rank {bound_rank}, but dispatch did not select a DP rank" - ))); - } - } +fn affinity_error(error: AffinityError) -> Error { + match error { + AffinityError::InvalidArgument(message) => invalid_argument(message), + AffinityError::ResourceExhausted(message) => DynamoError::builder() + .error_type(ErrorType::ResourceExhausted) + .message(message) + .build() + .into(), + AffinityError::Dropped => anyhow::anyhow!("session affinity coordinator dropped"), } - Ok(()) } pub(crate) fn invalid_argument(message: impl Into) -> Error { @@ -993,14 +430,6 @@ pub(crate) fn invalid_argument(message: impl Into) -> Error { .into() } -fn resource_exhausted(message: impl Into) -> Error { - DynamoError::builder() - .error_type(ErrorType::ResourceExhausted) - .message(message.into()) - .build() - .into() -} - fn cancelled(context_id: &str) -> Error { DynamoError::builder() .error_type(ErrorType::Cancelled) diff --git a/lib/llm/src/session_affinity/mod.rs b/lib/llm/src/session_affinity/mod.rs index 1a570ba08b7c..5b0341c26423 100644 --- a/lib/llm/src/session_affinity/mod.rs +++ b/lib/llm/src/session_affinity/mod.rs @@ -4,39 +4,18 @@ mod coordinator; mod replica_sync; -use std::{str::FromStr, time::Duration}; +use std::time::Duration; use dynamo_runtime::{component::Client, pipeline::Error}; -use serde::{Deserialize, Serialize}; +#[cfg(test)] +pub(crate) use coordinator::to_table; pub(crate) use coordinator::{AffinityAcquire, affinity_id, invalid_argument}; pub use coordinator::{AffinityCoordinator, AffinityTarget, explicit_target}; - -#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "lowercase")] -pub enum SessionAffinityMode { - #[default] - Hard, - Soft, -} - -impl FromStr for SessionAffinityMode { - type Err = String; - - fn from_str(value: &str) -> Result { - match value { - "hard" => Ok(Self::Hard), - "soft" => Ok(Self::Soft), - _ => Err(format!( - "invalid session affinity mode {value:?}; expected 'hard' or 'soft'" - )), - } - } -} - -pub const MAX_SESSION_AFFINITY_TTL_SECS: u64 = 31_536_000; -pub const MAX_SESSION_AFFINITY_ENTRIES: usize = 65_536; -pub const MAX_SESSION_AFFINITY_ID_BYTES: usize = 256; +pub use dynamo_kv_router::services::selection::affinity::{ + MAX_SESSION_AFFINITY_ENTRIES, MAX_SESSION_AFFINITY_ID_BYTES, MAX_SESSION_AFFINITY_TTL_SECS, + SessionAffinityMode, +}; pub type LlmResponse = crate::types::Annotated; diff --git a/lib/llm/src/session_affinity/replica_sync.rs b/lib/llm/src/session_affinity/replica_sync.rs index 6afb7c027cf8..06ed7f4c9e68 100644 --- a/lib/llm/src/session_affinity/replica_sync.rs +++ b/lib/llm/src/session_affinity/replica_sync.rs @@ -1,7 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use std::sync::Weak; +use std::sync::Arc; use anyhow::{Context, Result}; use dynamo_runtime::{ @@ -19,10 +19,12 @@ use tokio::{ }; use tokio_util::sync::CancellationToken; -use super::coordinator::{AffinityCoordinatorInner, AffinityTarget, AffinityVersion}; use crate::direct_zmq_fan_in::{ ContinuityMode, FanInEvent, FanInObservation, start_direct_zmq_fan_in, }; +use dynamo_kv_router::services::selection::affinity::{ + AffinityReplicaSink, AffinityTarget, AffinityVersion, WeakSessionAffinity, +}; pub(super) const SESSION_AFFINITY_SUBJECT: &str = "session_affinity_events"; const OUTBOUND_CHANNEL_CAPACITY: usize = 4_096; @@ -42,7 +44,7 @@ struct ReplicaUpdateSender { tx: mpsc::Sender, } -impl ReplicaUpdateSender { +impl AffinityReplicaSink for ReplicaUpdateSender { fn publish(&self, session_id: &str, target: AffinityTarget, version: AffinityVersion) { let update = SessionAffinityUpdate { session_id: session_id.to_string(), @@ -66,7 +68,7 @@ impl ReplicaUpdateSender { struct ReplicaUpdateApplier { local_publisher_id: u64, discovered_instances: watch::Receiver>, - coordinator: Weak, + table: WeakSessionAffinity, } impl ReplicaUpdateApplier { @@ -75,10 +77,10 @@ impl ReplicaUpdateApplier { return true; } - let Some(coordinator) = self.coordinator.upgrade() else { + let Some(table) = self.table.upgrade() else { return false; }; - coordinator.observe_replica_sequence(update.sequence); + table.observe_replica_sequence(update.sequence); if !self .discovered_instances .borrow() @@ -87,11 +89,8 @@ impl ReplicaUpdateApplier { { return true; } - let target = AffinityTarget { - worker_id: update.worker_id, - dp_rank: update.dp_rank, - }; - let outcome = coordinator.apply_replica_update( + let target = AffinityTarget::new(update.worker_id, update.dp_rank); + let outcome = table.apply_replica_update( update.session_id, target, AffinityVersion { @@ -99,7 +98,7 @@ impl ReplicaUpdateApplier { writer_id: update.writer_id, }, ); - drop(coordinator); + drop(table); tracing::trace!( worker_id = target.worker_id, dp_rank = ?target.dp_rank, @@ -118,11 +117,8 @@ pub(super) struct ReplicaSyncRuntime { } impl ReplicaSyncRuntime { - pub(super) async fn start( - client: Client, - coordinator: Weak, - parent_cancel: &CancellationToken, - ) -> Result { + /// Returns the runtime and this replica's writer id (its discovery instance id). + pub(super) async fn start(client: Client, table: WeakSessionAffinity) -> Result<(Self, u64)> { let endpoint = &client.endpoint; let router_id = endpoint.drt().discovery().instance_id(); let transport_kind = endpoint.drt().default_event_transport_kind(); @@ -134,18 +130,13 @@ impl ReplicaSyncRuntime { .await .context("create session affinity event publisher")?; let publisher_id = publisher.publisher_id(); - if let Some(inner) = coordinator.upgrade() { - inner - .writer_id - .store(router_id, std::sync::atomic::Ordering::Relaxed); - } let applier = ReplicaUpdateApplier { local_publisher_id: publisher_id, discovered_instances: client.instance_source.as_ref().clone(), - coordinator, + table, }; - let cancel = parent_cancel.child_token(); + let cancel = CancellationToken::new(); let subscriber_task = if should_use_direct_sync(transport_kind, uses_direct_zmq(transport_kind)) { let codec = Codec::default(); @@ -244,14 +235,23 @@ impl ReplicaSyncRuntime { } }); - Ok(Self { - sender, - cancel, - publisher_task: Some(publisher_task), - subscriber_task: Some(subscriber_task), - }) + Ok(( + Self { + sender, + cancel, + publisher_task: Some(publisher_task), + subscriber_task: Some(subscriber_task), + }, + router_id, + )) + } + + /// The sink the table publishes into. + pub(super) fn sink(&self) -> Arc { + Arc::new(self.sender.clone()) } + #[cfg(test)] pub(super) fn publish( &self, session_id: &str, @@ -353,7 +353,7 @@ mod tests { let applier = ReplicaUpdateApplier { local_publisher_id: 7, discovered_instances, - coordinator: coordinator.downgrade_for_test(), + table: coordinator.table_for_test(), }; assert!(applier.apply( diff --git a/lib/llm/src/session_affinity/tests.rs b/lib/llm/src/session_affinity/tests.rs index cfb0fd02d44d..00762006e4c1 100644 --- a/lib/llm/src/session_affinity/tests.rs +++ b/lib/llm/src/session_affinity/tests.rs @@ -14,7 +14,8 @@ use futures::{StreamExt, stream}; use super::SessionAffinityMode::{Hard, Soft}; use super::{ AffinityAcquire, AffinityCoordinator, AffinityTarget, LlmResponse, affinity_id, - coordinator::ReplicaApplyOutcome, explicit_target, + coordinator::{ReplicaApplyOutcome, tracked_stream}, + explicit_target, to_table, }; use crate::{ preprocessor::PreprocessedRequest, @@ -156,7 +157,7 @@ async fn session_affinity_initialization_is_atomic() { coordinator.wait_for_initializing_waiter().await; assert!(!waiter.is_finished()); - let first_lease = first.commit(target(7, Some(0))).unwrap(); + let first_lease = first.commit(to_table(target(7, Some(0)))).unwrap(); let second = waiter.await.unwrap().unwrap(); let AffinityAcquire::Bound { target: second_target, @@ -233,7 +234,7 @@ async fn session_affinity_validates_worker_and_rank_contract() { else { panic!("first request must initialize"); }; - drop(initializer.commit(target(7, None)).unwrap()); + drop(initializer.commit(to_table(target(7, None))).unwrap()); assert!( coordinator @@ -263,7 +264,7 @@ async fn session_affinity_failed_bound_operation_invalidates_binding() { else { panic!("first request must initialize"); }; - drop(initializer.commit(target(7, Some(0))).unwrap()); + drop(initializer.commit(to_table(target(7, Some(0)))).unwrap()); let operation = coordinator.acquire(&session_id(), None).await.unwrap(); assert_eq!(operation.target(), Some(target(7, Some(0)))); @@ -285,9 +286,9 @@ async fn session_affinity_stream_drop_refreshes_idle_ttl() { else { panic!("first request must initialize"); }; - let lease = initializer.commit(target(7, Some(0))).unwrap(); + let lease = initializer.commit(to_table(target(7, Some(0)))).unwrap(); tokio::time::advance(Duration::from_secs(9)).await; - let mut stream = lease.into_stream(response_stream(1)); + let mut stream = tracked_stream(lease, response_stream(1)); assert!(stream.next().await.is_some()); drop(stream); @@ -302,9 +303,9 @@ async fn session_affinity_empty_stream_refreshes_idle_ttl() { else { panic!("first request must initialize"); }; - let lease = initializer.commit(target(7, Some(0))).unwrap(); + let lease = initializer.commit(to_table(target(7, Some(0)))).unwrap(); tokio::time::advance(Duration::from_secs(9)).await; - let mut stream = lease.into_stream(response_stream(0)); + let mut stream = tracked_stream(lease, response_stream(0)); assert!(stream.next().await.is_none()); assert_binding_expires_after_refreshed_ttl(&coordinator).await; @@ -318,7 +319,7 @@ async fn session_affinity_cancelled_stream_refreshes_idle_ttl() { else { panic!("first request must initialize"); }; - drop(initializer.commit(target(7, Some(0))).unwrap()); + drop(initializer.commit(to_table(target(7, Some(0)))).unwrap()); tokio::time::advance(Duration::from_secs(9)).await; let AffinityAcquire::Bound { @@ -329,7 +330,7 @@ async fn session_affinity_cancelled_stream_refreshes_idle_ttl() { panic!("continuation must acquire the existing binding"); }; assert_eq!(bound_target, target(7, Some(0))); - let mut stream = lease.into_stream(cancelled_response_stream()); + let mut stream = tracked_stream(lease, cancelled_response_stream()); assert!(stream.next().await.is_none()); assert_binding_expires_after_refreshed_ttl(&coordinator).await; @@ -356,9 +357,9 @@ async fn session_affinity_error_stream_refreshes_idle_ttl() { else { panic!("first request must initialize"); }; - let lease = initializer.commit(target(7, Some(0))).unwrap(); + let lease = initializer.commit(to_table(target(7, Some(0)))).unwrap(); tokio::time::advance(Duration::from_secs(9)).await; - let mut stream = lease.into_stream(error_response_stream()); + let mut stream = tracked_stream(lease, error_response_stream()); assert!(stream.next().await.unwrap().is_err()); assert!(stream.next().await.is_none()); @@ -373,9 +374,9 @@ async fn session_affinity_stream_eof_refreshes_idle_ttl() { else { panic!("first request must initialize"); }; - let lease = initializer.commit(target(7, Some(0))).unwrap(); + let lease = initializer.commit(to_table(target(7, Some(0)))).unwrap(); tokio::time::advance(Duration::from_secs(9)).await; - let mut stream = lease.into_stream(response_stream(1)); + let mut stream = tracked_stream(lease, response_stream(1)); while stream.next().await.is_some() {} assert_binding_expires_after_refreshed_ttl(&coordinator).await; @@ -389,7 +390,7 @@ async fn session_affinity_bound_lease_drop_refreshes_idle_ttl() { else { panic!("first request must initialize"); }; - drop(initializer.commit(target(7, Some(0))).unwrap()); + drop(initializer.commit(to_table(target(7, Some(0)))).unwrap()); tokio::time::advance(Duration::from_secs(9)).await; let AffinityAcquire::Bound { lease, .. } = @@ -425,7 +426,7 @@ async fn session_affinity_query_is_read_only() { else { panic!("first request must initialize"); }; - drop(initializer.commit(target(7, Some(0))).unwrap()); + drop(initializer.commit(to_table(target(7, Some(0)))).unwrap()); assert_eq!( coordinator.query_target(&session_id(), None).unwrap(), Some(target(7, Some(0))) @@ -444,7 +445,7 @@ async fn session_affinity_reaper_removes_idle_entries_and_stops_on_drop() { else { panic!("first request must initialize"); }; - drop(initializer.commit(target(7, Some(0))).unwrap()); + drop(initializer.commit(to_table(target(7, Some(0)))).unwrap()); coordinator.wait_for_reaper().await; tokio::time::advance(Duration::from_secs(10)).await; From 12d061056a9c1d3ccbbe119b9ef265cd353a560c Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Fri, 4 Sep 2026 01:46:32 +0000 Subject: [PATCH 030/102] feat(kv-router): session affinity for request-shaped selection hosts SelectionServiceBuilder::session_affinity(ttl) gives the core its own SessionAffinity table. A booked selection whose request carries a session id (session_id or session_context) and no explicit pin or affinity target acquires the session first: a bound session steers selection to its worker (exclusive for the default selector) and a new one is bound to the worker the scheduler chose. The lease is held with the booking and released by free_reservation; a bound worker that was not selected (it left) is invalidated and the session rebound. Advisory selects read the binding. Bindings replicate over the existing ZMQ peer mesh on their own topic (dynamo.session-affinity.v1) so peers that do not subscribe never see them; the active-sequence topic and wire format are untouched. The publisher multiplexes both channels onto one socket; PeerManager subscribes to the second topic only when a handler is given. SelectionServiceConfig carries the TTL for the standalone binary. Signed-off-by: Ishan Dhanani --- .../src/services/common/replica_sync.rs | 211 +++++++++++---- lib/kv-router/src/services/common/zmq.rs | 8 +- .../src/services/selection/affinity.rs | 18 +- .../src/services/selection/core/mod.rs | 247 +++++++++++++++++- .../src/services/selection/service.rs | 23 +- lib/kv-router/src/services/selection/tests.rs | 1 + 6 files changed, 453 insertions(+), 55 deletions(-) diff --git a/lib/kv-router/src/services/common/replica_sync.rs b/lib/kv-router/src/services/common/replica_sync.rs index 0a9e23bebd5f..cc294d2e7630 100644 --- a/lib/kv-router/src/services/common/replica_sync.rs +++ b/lib/kv-router/src/services/common/replica_sync.rs @@ -15,11 +15,62 @@ use tokio_util::sync::CancellationToken; use crate::identity::{RoutingPartitionId, RoutingPartitionRef}; use crate::protocols::{ActiveSequenceEvent, WorkerWithDpRank}; use crate::sequences::{SequencePublishQueueError, SequencePublisher, SequenceSubscriber}; -use crate::services::common::zmq::{create_bound_pub_socket, create_sub_socket, validate_endpoint}; +use crate::services::common::zmq::{ + create_bound_pub_socket, create_sub_socket_topics, validate_endpoint, +}; +use crate::services::selection::affinity::{AffinityReplicaSink, AffinityTarget, AffinityVersion}; pub(crate) const REPLICA_EVENT_CHANNEL_CAPACITY: usize = 100_000; const PEER_COMMAND_CHANNEL_CAPACITY: usize = 64; const REPLICA_TOPIC: &[u8] = b"dynamo.slot-tracker.v1"; +/// Session-affinity bindings ride the same mesh on their own topic, so peers +/// that do not subscribe to it never see them. +const AFFINITY_TOPIC: &[u8] = b"dynamo.session-affinity.v1"; +const AFFINITY_EVENT_CHANNEL_CAPACITY: usize = 4_096; + +/// One replicated session binding. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub(crate) struct AffinityBindingEvent { + pub session_id: String, + pub worker_id: u64, + pub dp_rank: Option, + pub sequence: u64, + pub writer_id: u64, +} + +impl AffinityBindingEvent { + pub(crate) fn target(&self) -> AffinityTarget { + AffinityTarget::new(self.worker_id, self.dp_rank) + } + + pub(crate) fn version(&self) -> AffinityVersion { + AffinityVersion { + sequence: self.sequence, + writer_id: self.writer_id, + } + } +} + +/// Publishes bindings from a [`super::super::selection::affinity::SessionAffinity`] +/// into the mesh. Best effort: a full channel drops the update. +struct AffinityMeshSink { + tx: mpsc::Sender, +} + +impl AffinityReplicaSink for AffinityMeshSink { + fn publish(&self, session_id: &str, target: AffinityTarget, version: AffinityVersion) { + let update = AffinityBindingEvent { + session_id: session_id.to_string(), + worker_id: target.worker_id, + dp_rank: target.dp_rank, + sequence: version.sequence, + writer_id: version.writer_id, + }; + if let Err(error) = self.tx.try_send(update) { + tracing::trace!(%error, "dropping best-effort session affinity replica update"); + } + } +} #[derive(Debug, Clone, Serialize, Deserialize)] pub(crate) struct ScopedReplicaEvent { @@ -49,6 +100,7 @@ pub(crate) type ReplicaEventSender = mpsc::Sender; pub(crate) struct ReplicaSyncConfig { process_id: u64, outbound_tx: ReplicaEventSender, + affinity_tx: Option>, cancel_token: CancellationToken, } @@ -96,15 +148,28 @@ impl ReplicaSyncConfig { Self { process_id, outbound_tx, + affinity_tx: None, cancel_token, } } - #[cfg(feature = "standalone-slot-tracker")] + fn with_affinity_publisher(mut self, tx: mpsc::Sender) -> Self { + self.affinity_tx = Some(tx); + self + } + pub(crate) fn process_id(&self) -> u64 { self.process_id } + /// Sink that publishes session bindings into the mesh, when this runtime + /// carries them. + pub(crate) fn affinity_sink(&self) -> Option> { + self.affinity_tx + .clone() + .map(|tx| Arc::new(AffinityMeshSink { tx }) as Arc) + } + pub(crate) fn is_self_event(&self, event: &ActiveSequenceEvent) -> bool { event.router_id == self.process_id } @@ -329,10 +394,11 @@ pub(crate) fn setup_replica_sync( let bind_endpoint = replica_sync_bind_endpoint(port)?; let process_id = generate_process_id(); - let (outbound_tx, publisher_task) = + let (outbound_tx, affinity_tx, publisher_task) = start_replica_publisher(&bind_endpoint, cancel_token.clone())?; Ok(Some(ReplicaSyncRuntime { - config: ReplicaSyncConfig::new(process_id, outbound_tx, cancel_token.clone()), + config: ReplicaSyncConfig::new(process_id, outbound_tx, cancel_token.clone()) + .with_affinity_publisher(affinity_tx), cancel_token, publisher_task: Mutex::new(Some(publisher_task)), })) @@ -382,53 +448,63 @@ pub(crate) fn setup_scoped_replica_sync( pub(crate) fn start_replica_publisher( bind_endpoint: &str, cancel_token: CancellationToken, -) -> Result<(ReplicaEventSender, JoinHandle<()>)> { +) -> Result<( + ReplicaEventSender, + mpsc::Sender, + JoinHandle<()>, +)> { validate_endpoint(bind_endpoint)?; let mut socket = create_bound_pub_socket(bind_endpoint) .with_context(|| format!("failed to bind replica publisher to `{bind_endpoint}`"))?; let (tx, mut rx) = mpsc::channel::(REPLICA_EVENT_CHANNEL_CAPACITY); + let (affinity_tx, mut affinity_rx) = + mpsc::channel::(AFFINITY_EVENT_CHANNEL_CAPACITY); let task = tokio::spawn(async move { loop { - let event = tokio::select! { + let frames = tokio::select! { _ = cancel_token.cancelled() => break, event = rx.recv() => { let Some(event) = event else { break; }; - event + let partition = event.partition_ref(); + match rmp_serde::to_vec_named(&event) { + Ok(payload) => vec![REPLICA_TOPIC.to_vec(), payload], + Err(error) => { + tracing::error!( + model_name = %partition.model_name, + routing_group = %partition.routing_group, + request_id = %event.event.request_id, + "Failed to encode active-sequence replica event: {error}" + ); + continue; + } + } } - }; - - let request_id = event.event.request_id.clone(); - let partition = event.partition_ref(); - let payload = match rmp_serde::to_vec_named(&event) { - Ok(payload) => payload, - Err(error) => { - tracing::error!( - model_name = %partition.model_name, - routing_group = %partition.routing_group, - request_id = %request_id, - "Failed to encode active-sequence replica event: {error}" - ); - continue; + binding = affinity_rx.recv() => { + let Some(binding) = binding else { + break; + }; + match rmp_serde::to_vec_named(&binding) { + Ok(payload) => vec![AFFINITY_TOPIC.to_vec(), payload], + Err(error) => { + tracing::error!( + session_id = %binding.session_id, + "Failed to encode session affinity replica event: {error}" + ); + continue; + } + } } }; - if let Err(error) = socket - .send_multipart(vec![REPLICA_TOPIC.to_vec(), payload]) - .await - { - tracing::error!( - model_name = %partition.model_name, - routing_group = %partition.routing_group, - request_id = %request_id, - "Failed to publish active-sequence replica event: {error}" - ); + if let Err(error) = socket.send_multipart(frames).await { + tracing::error!("Failed to publish replica event: {error}"); } } }); - Ok((tx, task)) + Ok((tx, affinity_tx, task)) } #[derive(Debug, thiserror::Error)] @@ -466,6 +542,7 @@ enum PeerCommand { } impl PeerManager { + #[cfg_attr(not(feature = "standalone-slot-tracker"), allow(dead_code))] pub(crate) fn start( initial_peers: Vec, cancel_token: CancellationToken, @@ -474,7 +551,32 @@ impl PeerManager { where F: Fn(ScopedReplicaEvent) + Send + Sync + 'static, { - let mut socket = create_sub_socket(REPLICA_TOPIC)?; + Self::start_with_affinity( + initial_peers, + cancel_token, + handle_event, + None::, + ) + } + + /// Like [`Self::start`], also subscribing to session-affinity bindings + /// when `handle_affinity` is given. + pub(crate) fn start_with_affinity( + initial_peers: Vec, + cancel_token: CancellationToken, + handle_event: F, + handle_affinity: Option, + ) -> Result + where + F: Fn(ScopedReplicaEvent) + Send + Sync + 'static, + A: Fn(AffinityBindingEvent) + Send + Sync + 'static, + { + let topics: &[&[u8]] = if handle_affinity.is_some() { + &[REPLICA_TOPIC, AFFINITY_TOPIC] + } else { + &[REPLICA_TOPIC] + }; + let mut socket = create_sub_socket_topics(topics)?; let mut configured_peers = HashSet::new(); for endpoint in initial_peers { validate_endpoint(&endpoint) @@ -502,7 +604,7 @@ impl PeerManager { } message = socket.recv_multipart() => { match message { - Ok(frames) => handle_replica_message(&handle_event, frames), + Ok(frames) => handle_replica_message(&handle_event, handle_affinity.as_ref(), frames), Err(error) => { tracing::error!("Failed to receive active-sequence replica event: {error}"); } @@ -615,31 +717,40 @@ fn handle_peer_command( } } -fn handle_replica_message( +fn handle_replica_message( handle_event: &F, + handle_affinity: Option<&A>, frames: crate::services::common::zmq::MultipartMessage, ) where F: Fn(ScopedReplicaEvent), + A: Fn(AffinityBindingEvent), { let [topic, payload] = frames.as_slice() else { tracing::debug!( frame_count = frames.len(), - "Dropping malformed active-sequence replica message" + "Dropping malformed replica message" ); return; }; - if topic.as_slice() != REPLICA_TOPIC { - tracing::debug!("Dropping active-sequence replica message with unexpected topic"); - return; + match topic.as_slice() { + REPLICA_TOPIC => match rmp_serde::from_slice::(payload) { + Ok(event) => handle_event(event), + Err(error) => { + tracing::debug!("Dropping malformed active-sequence replica payload: {error}"); + } + }, + AFFINITY_TOPIC => match ( + handle_affinity, + rmp_serde::from_slice::(payload), + ) { + (Some(handle_affinity), Ok(event)) => handle_affinity(event), + (None, _) => {} + (_, Err(error)) => { + tracing::debug!("Dropping malformed session affinity replica payload: {error}"); + } + }, + _ => tracing::debug!("Dropping replica message with unexpected topic"), } - let event: ScopedReplicaEvent = match rmp_serde::from_slice(payload) { - Ok(event) => event, - Err(error) => { - tracing::debug!("Dropping malformed active-sequence replica payload: {error}"); - return; - } - }; - handle_event(event); } #[cfg(test)] @@ -786,7 +897,7 @@ mod tests { async fn dynamic_peer_registration_controls_delivery() { let endpoint = reserve_tcp_endpoint(); let cancel_token = CancellationToken::new(); - let (outbound, publisher_task) = + let (outbound, _affinity, publisher_task) = start_replica_publisher(&endpoint, cancel_token.child_token()).expect("publisher"); let (received_tx, mut received_rx) = mpsc::channel(16); let manager = PeerManager::start(Vec::new(), cancel_token.child_token(), move |event| { @@ -834,9 +945,9 @@ mod tests { let endpoint_a = reserve_tcp_endpoint(); let endpoint_b = reserve_tcp_endpoint(); let cancel_token = CancellationToken::new(); - let (outbound_a, publisher_a) = + let (outbound_a, _affinity_a, publisher_a) = start_replica_publisher(&endpoint_a, cancel_token.child_token()).unwrap(); - let (outbound_b, publisher_b) = + let (outbound_b, _affinity_b, publisher_b) = start_replica_publisher(&endpoint_b, cancel_token.child_token()).unwrap(); let registry_a = Arc::new(SlotTrackerRegistry::new_with_replica_sync( cancel_token.clone(), diff --git a/lib/kv-router/src/services/common/zmq.rs b/lib/kv-router/src/services/common/zmq.rs index 406be80b5103..5d5d821e8d8b 100644 --- a/lib/kv-router/src/services/common/zmq.rs +++ b/lib/kv-router/src/services/common/zmq.rs @@ -207,9 +207,15 @@ where } pub(crate) fn create_sub_socket(topic: &[u8]) -> Result { + create_sub_socket_topics(&[topic]) +} + +pub(crate) fn create_sub_socket_topics(topics: &[&[u8]]) -> Result { build_socket(zmq::SUB, |socket| { configure_receive_socket(socket)?; - socket.set_subscribe(topic)?; + for topic in topics { + socket.set_subscribe(topic)?; + } Ok(()) }) } diff --git a/lib/kv-router/src/services/selection/affinity.rs b/lib/kv-router/src/services/selection/affinity.rs index 388a8c8d8d8c..f78b3639994d 100644 --- a/lib/kv-router/src/services/selection/affinity.rs +++ b/lib/kv-router/src/services/selection/affinity.rs @@ -142,6 +142,17 @@ impl WeakSessionAffinity { } } +/// A session acquired for one request. +pub enum Acquired { + /// New (or expired) session: select a worker, then `commit` it. + Initialize(AffinityInitialization), + /// Bound session: route to `target`; the lease releases on drop. + Bound { + target: AffinityTarget, + lease: AffinityLease, + }, +} + /// One step of acquiring a session. pub enum AcquireStep { /// New (or expired) session: select a worker, then `commit` it. @@ -356,11 +367,14 @@ impl SessionAffinity { &self, session_id: &str, requested_target: Option, - ) -> Result { + ) -> Result { loop { match self.try_acquire(session_id, requested_target)? { AcquireStep::Wait(notified) => notified.await, - step => return Ok(step), + AcquireStep::Initialize(init) => return Ok(Acquired::Initialize(init)), + AcquireStep::Bound { target, lease } => { + return Ok(Acquired::Bound { target, lease }); + } } } } diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index 6a65147a8e6e..53bb8fa7fb9a 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -3,7 +3,7 @@ use std::collections::{HashMap, HashSet}; use std::sync::Arc; -use std::time::Instant; +use std::time::{Duration, Instant}; use dynamo_tokens::SequenceHash; use once_cell::sync::OnceCell; @@ -44,6 +44,7 @@ use crate::services::indexer::registry::WorkerRegistry; use crate::services::overlap::MooncakeOverlapSummary; use crate::tracking_hash::{TrackingHashContext, TrackingHashScope}; +use super::affinity::{Acquired, AffinityInitialization, AffinityLease, SessionAffinity}; use super::catalog::WorkerCatalog; use super::error::SelectionError; use super::ingress::{KvEventIngress, ZmqDirectIngress}; @@ -57,6 +58,7 @@ use super::types::{ }; use crate::WorkerSelectionPolicyFactory; use crate::WorkerType; +use crate::services::common::replica_sync::AffinityBindingEvent; /// Source of dequeue-time overlap refreshes for a partition scheduler: the /// partition's index. @@ -294,6 +296,8 @@ pub struct SelectionServiceConfig { pub replica_sync_peers: Vec, pub kv_router_config: crate::config::KvRouterConfig, pub selection_cache: SelectionCacheConfig, + /// Session stickiness TTL; `None` disables session affinity. + pub session_affinity_ttl: Option, } type SelectionEntries = RwLock>>>>; @@ -330,6 +334,30 @@ pub struct SelectionCore { /// `create_reservation` can replay them without re-sending the prompt. selection_cache: SelectionCache, tracking_hash: Arc, + /// Session stickiness for request-shaped hosts; `None` when not configured. + session_affinity: Option, + /// Leases held by booked selections, released with the booking. + affinity_leases: parking_lot::Mutex>, +} + +/// What a booked selection holds on its session until the worker is known. +enum AffinityHold { + Initialize(AffinityInitialization), + Bound { + target: SessionTarget, + lease: AffinityLease, + }, +} + +type SessionTarget = super::affinity::AffinityTarget; + +fn affinity_error(error: super::affinity::AffinityError) -> SelectionError { + use super::affinity::AffinityError; + match error { + AffinityError::InvalidArgument(message) => SelectionError::BadRequest(message), + AffinityError::ResourceExhausted(message) => SelectionError::NotReady(message), + AffinityError::Dropped => SelectionError::Internal(error.to_string()), + } } impl SelectionCore { @@ -372,6 +400,7 @@ impl SelectionCore { cache_config, tracking_hash, indexer_policy, + None, )) } @@ -394,8 +423,28 @@ impl SelectionCore { cache_config: SelectionCacheConfig, tracking_hash: Arc, indexer_policy: IndexerPolicy, + session_affinity_ttl: Option, ) -> Self { let cancel_token = cancel_token.child_token(); + let session_affinity = match session_affinity_ttl { + Some(ttl) => match SessionAffinity::new(ttl) { + Ok(table) => { + if let Some(sink) = replica_config.as_ref().and_then(|c| c.affinity_sink()) { + let process_id = replica_config + .as_ref() + .map(|c| c.process_id()) + .unwrap_or_default(); + table.enable_replication(process_id, sink); + } + Some(table) + } + Err(error) => { + tracing::warn!(%error, "session affinity disabled"); + None + } + }, + None => None, + }; let indexer_registry = Arc::new(WorkerRegistry::new_with_cancel_token( indexer_threads, cancel_token.clone(), @@ -420,6 +469,8 @@ impl SelectionCore { replica_config, selection_cache: SelectionCache::new(&cache_config), tracking_hash, + session_affinity, + affinity_leases: parking_lot::Mutex::new(HashMap::new()), } } @@ -940,6 +991,40 @@ impl SelectionCore { } let entry = self.ready_entry(&key)?; + + // Session stickiness: a bound session steers selection (exclusive for + // the default selector); a new session is bound to the worker chosen. + // An explicit affinity target or pin from the caller wins. + let session_id = self + .session_affinity + .as_ref() + .and(session_context.as_ref()) + .filter(|_| affinity_target.is_none() && pinned_worker.is_none()) + .map(|context| context.session_id().to_string()); + let mut affinity_hold = None; + let affinity_target = match (session_id.as_deref(), self.session_affinity.as_ref()) { + (Some(session_id), Some(table)) if book => { + match table + .acquire(session_id, None) + .await + .map_err(affinity_error)? + { + Acquired::Initialize(init) => { + affinity_hold = Some(AffinityHold::Initialize(init)); + None + } + Acquired::Bound { target, lease } => { + affinity_hold = Some(AffinityHold::Bound { target, lease }); + Some(WorkerAffinityTarget::new(target.worker_id, target.dp_rank)) + } + } + } + (Some(session_id), Some(table)) => table + .query_target(session_id, None) + .map_err(affinity_error)? + .map(|target| WorkerAffinityTarget::new(target.worker_id, target.dp_rank)), + _ => affinity_target, + }; // Router hints are attached to bookings only, and only when a worker in // this partition can consume them and the indexer can retain the // matched chain (local, event-driven, no approximate writes). @@ -1051,6 +1136,13 @@ impl SelectionCore { } } => result?, }; + if let (Some(hold), Some(session_id), Some(selection_id)) = ( + affinity_hold.take(), + session_id.as_deref(), + selection_id.as_deref(), + ) { + self.bind_session(hold, session_id, selection_id, response.best_worker); + } let endpoint = self .catalog .schedulable_endpoint(response.best_worker.worker_id, &key) @@ -1376,6 +1468,77 @@ impl SelectionCore { fn forget_reservation(&self, selection_id: &str) { self.reservation_index.write().remove(selection_id); + self.affinity_leases.lock().remove(selection_id); + } + + /// Bind (or confirm) `session_id` to the worker a booking landed on and + /// keep the lease with the booking. A bound session whose worker was not + /// selected (it left) is invalidated and rebound. + fn bind_session( + &self, + hold: AffinityHold, + session_id: &str, + selection_id: &str, + worker: WorkerWithDpRank, + ) { + let Some(table) = self.session_affinity.as_ref() else { + return; + }; + let selected = SessionTarget::new(worker.worker_id, Some(worker.dp_rank)); + let lease = match hold { + AffinityHold::Initialize(init) => init.commit(selected).ok(), + AffinityHold::Bound { target, lease } + if target.worker_id == selected.worker_id + && target.dp_rank.is_none_or(|rank| rank == worker.dp_rank) => + { + lease.publish(target); + return self.hold_lease(selection_id, lease); + } + AffinityHold::Bound { mut lease, .. } => { + lease.invalidate(); + match table.try_acquire(session_id, None) { + Ok(super::affinity::AcquireStep::Initialize(init)) => { + init.commit(selected).ok() + } + _ => None, + } + } + }; + if let Some(lease) = lease { + lease.publish(selected); + self.hold_lease(selection_id, lease); + } + } + + fn hold_lease(&self, selection_id: &str, lease: AffinityLease) { + self.affinity_leases + .lock() + .insert(selection_id.to_string(), lease); + } + + /// Apply a session binding a replica published. + pub(crate) fn dispatch_affinity_event(&self, event: AffinityBindingEvent) { + let Some(table) = self.session_affinity.as_ref() else { + return; + }; + if self + .replica_config + .as_ref() + .is_some_and(|config| config.process_id() == event.writer_id) + { + return; + } + table.observe_replica_sequence(event.sequence); + if self.catalog.get(event.worker_id).is_none() { + return; + } + let (target, version, worker_id) = (event.target(), event.version(), event.worker_id); + let outcome = table.apply_replica_update(event.session_id, target, version); + tracing::trace!( + worker_id, + ?outcome, + "applied session affinity replica update" + ); } /// Record a booked routing decision into the partition's approximate @@ -1959,6 +2122,7 @@ mod tests { SelectionCacheConfig::default(), tracking_hash, IndexerPolicy::from_router_config(&test_config(false)).expect("indexer policy"), + None, ); core.upsert_worker(worker(1)).await.expect("worker upsert"); @@ -1998,6 +2162,7 @@ mod tests { SelectionCacheConfig::default(), tracking_hash, indexer_policy, + None, ) } @@ -2158,6 +2323,7 @@ mod tests { SelectionCacheConfig::default(), tracking_hash, indexer_policy, + None, ); assert!(!core.listens_for_kv_events); for worker_id in [1, 2] { @@ -3123,4 +3289,83 @@ mod tests { assert_eq!(response.overlap.disk, 8); assert_eq!(response.overlap.dp, HashMap::from([("0".to_string(), 8)])); } + + fn core_with_session_affinity() -> SelectionCore { + let config = test_config(false); + let tracking_hash = Arc::new( + TrackingHashContext::from_config(&config).expect("valid tracking hash configuration"), + ); + let indexer_policy = IndexerPolicy::from_router_config(&config).expect("indexer policy"); + SelectionCore::new_inner( + config, + 1, + CancellationToken::new(), + None, + None, + SelectionHost::default(), + WorkerType::Aggregated, + true, + SelectionCacheConfig::default(), + tracking_hash, + indexer_policy, + Some(Duration::from_secs(10)), + ) + } + + fn session_reservation(selection_id: &str, session_id: &str) -> SelectAndReserveRequest { + let mut request = reserve_request(selection_id); + request.session_id = Some(session_id.to_string()); + request + } + + #[tokio::test] + async fn session_stays_on_its_first_worker_across_bookings() { + let core = core_with_session_affinity(); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + core.upsert_worker(worker(2)).await.expect("worker upsert"); + + let first = core + .select_and_reserve(session_reservation("r1", "chat-a")) + .await + .expect("first booking"); + for index in 0..4 { + let response = core + .select_and_reserve(session_reservation(&format!("r-{index}"), "chat-a")) + .await + .expect("booking"); + assert_eq!(response.worker_id, first.worker_id, "session must stay put"); + core.free_reservation(&format!("r-{index}")) + .await + .expect("free"); + } + // A read-only select sees the binding too. + let mut advisory = select_request(); + advisory.session_id = Some("chat-a".to_string()); + let response = core.select(advisory).await.expect("select"); + assert_eq!(response.worker_id, first.worker_id); + core.free_reservation("r1").await.expect("free"); + } + + #[tokio::test] + async fn replicated_binding_steers_a_new_session_and_frees_with_the_booking() { + let core = core_with_session_affinity(); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + core.upsert_worker(worker(2)).await.expect("worker upsert"); + + core.dispatch_affinity_event(AffinityBindingEvent { + session_id: "chat-b".to_string(), + worker_id: 2, + dp_rank: Some(0), + sequence: 1, + writer_id: 99, + }); + let response = core + .select_and_reserve(session_reservation("r2", "chat-b")) + .await + .expect("booking"); + assert_eq!(response.worker_id, 2); + assert!(core.affinity_leases.lock().contains_key("r2")); + core.free_reservation("r2").await.expect("free"); + assert!(core.affinity_leases.lock().is_empty()); + } } diff --git a/lib/kv-router/src/services/selection/service.rs b/lib/kv-router/src/services/selection/service.rs index 49f134bec223..2781cb620f2a 100644 --- a/lib/kv-router/src/services/selection/service.rs +++ b/lib/kv-router/src/services/selection/service.rs @@ -40,6 +40,7 @@ pub struct SelectionServiceBuilder { worker_selection_policy_registry: WorkerSelectionPolicyRegistry, host: SelectionHost, worker_selection_policy_factory: Option, + session_affinity_ttl: Option, } /// Warn when a host does not construct workers for explicitly configured policy roles. @@ -79,6 +80,7 @@ impl SelectionServiceBuilder { worker_selection_policy_registry, host: SelectionHost::default(), worker_selection_policy_factory: None, + session_affinity_ttl: None, } } @@ -115,6 +117,13 @@ impl SelectionServiceBuilder { self } + /// Pin each session id to the worker that served it for `ttl` after its + /// last request. Bindings replicate over the replica mesh when enabled. + pub fn session_affinity(mut self, ttl: std::time::Duration) -> Self { + self.session_affinity_ttl = Some(ttl); + self + } + pub fn replica_sync(mut self, port: u16, peers: Vec) -> Self { self.replica_sync_port = Some(port); self.replica_sync_peers = peers; @@ -176,6 +185,7 @@ impl SelectionServiceBuilder { self.selection_cache, tracking_hash, indexer_policy, + self.session_affinity_ttl, )); if recover_from_peers { @@ -195,7 +205,8 @@ impl SelectionServiceBuilder { let peer_manager = if replica_runtime.is_some() { let weak_core = Arc::downgrade(&core); - Some(PeerManager::start( + let affinity_core = Arc::downgrade(&core); + Some(PeerManager::start_with_affinity( self.replica_sync_peers, cancel_token.child_token(), move |event| { @@ -203,6 +214,13 @@ impl SelectionServiceBuilder { core.dispatch_replica_event(event); } }, + self.session_affinity_ttl.map(|_| { + move |event| { + if let Some(core) = affinity_core.upgrade() { + core.dispatch_affinity_event(event); + } + } + }), )?) } else { None @@ -236,6 +254,9 @@ impl SelectionServiceConfig { if let Some(port) = self.replica_sync_port { builder = builder.replica_sync(port, self.replica_sync_peers.clone()); } + if let Some(ttl) = self.session_affinity_ttl { + builder = builder.session_affinity(ttl); + } if let Some(url) = &self.remote_indexer_url { builder = builder.kv_index(KvIndexSource::Remote(url.clone())); } diff --git a/lib/kv-router/src/services/selection/tests.rs b/lib/kv-router/src/services/selection/tests.rs index 8d194cc98aa0..bd1c37a43ea4 100644 --- a/lib/kv-router/src/services/selection/tests.rs +++ b/lib/kv-router/src/services/selection/tests.rs @@ -1798,6 +1798,7 @@ async fn selector_replica_sync_propagates_request_lifecycle() { threads: 1, indexer_peers: Vec::new(), remote_indexer_url: None, + session_affinity_ttl: None, replica_sync_port: Some(port_a), replica_sync_peers: Vec::new(), kv_router_config: test_config(), From f1d266826653603b614fe399b4f110031b2654e3 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Fri, 4 Sep 2026 01:46:32 +0000 Subject: [PATCH 031/102] feat(epp): session affinity in the standalone EPP DYN_EPP_SESSION_AFFINITY_TTL_SECS turns on the selector's session table. The EPP passes x-dynamo-session-id, the same header the frontend reads, as the request's session id, so repeat requests of a session land on the worker that served it and peer EPPs learn the binding over the replica mesh. Hard mode only: the EPP routes by headers and never observes where a request ran. The standalone selection binary and the Python SelectionService gain the same TTL setting (--session-affinity-ttl-secs, session_affinity_ttl_secs). Signed-off-by: Ishan Dhanani --- .../inference-gateway/ext-proc/src/epp_router.rs | 4 ++++ .../ext-proc/src/epp_standalone_config.rs | 5 +++++ deploy/inference-gateway/ext-proc/src/selector.rs | 10 +++++++++- .../ext-proc/src/topology_adapter.rs | 1 + .../router/standalone-selection.md | 1 + .../kv-aware-routing/vanilla-vllm-onramp.mdx | 1 + lib/bindings/python/rust/llm/kv.rs | 14 +++++++++++++- 7 files changed, 34 insertions(+), 2 deletions(-) diff --git a/deploy/inference-gateway/ext-proc/src/epp_router.rs b/deploy/inference-gateway/ext-proc/src/epp_router.rs index 2c423e590c0d..68bdd2eec971 100644 --- a/deploy/inference-gateway/ext-proc/src/epp_router.rs +++ b/deploy/inference-gateway/ext-proc/src/epp_router.rs @@ -70,6 +70,9 @@ impl RenderClient { } } +/// Session id the standalone EPP pins to a worker when session affinity is on. +const HEADER_SESSION_ID: &str = "x-dynamo-session-id"; + /// Standalone endpoint picker backed by the standalone selection service. pub struct EppRouter { renderer: RenderClient, @@ -341,6 +344,7 @@ impl EndpointPicker for EppRouter { model_name: self.model_name.clone(), reservation_id: reservation_id.clone(), token_ids: tokens, + session_id: first_header(&req.headers, HEADER_SESSION_ID).map(str::to_owned), // `None` on the ordinary path: the selector schedules over its // catalog; `Some` only carries an Envoy subset constraint. allowed_worker_ids: allowed, diff --git a/deploy/inference-gateway/ext-proc/src/epp_standalone_config.rs b/deploy/inference-gateway/ext-proc/src/epp_standalone_config.rs index 874695f70db8..1f2582df5644 100644 --- a/deploy/inference-gateway/ext-proc/src/epp_standalone_config.rs +++ b/deploy/inference-gateway/ext-proc/src/epp_standalone_config.rs @@ -170,6 +170,10 @@ pub struct EppStandaloneConfig { /// throughput throttle). Excess requests are shed with a 503, not queued. #[validate(range(min = 1, message = "DYN_EPP_MAX_INFLIGHT_REQUESTS must be >= 1"))] pub max_inflight_requests: usize, + /// Pin each `x-dynamo-session-id` to the worker that served it for this + /// long after its last request (`DYN_EPP_SESSION_AFFINITY_TTL_SECS`). + /// `None` disables session affinity. + pub session_affinity_ttl_secs: Option, } impl EppStandaloneConfig { @@ -233,6 +237,7 @@ impl EppStandaloneConfig { max_num_batched_tokens: opt_parse::(get, "DYN_EPP_MAX_NUM_BATCHED_TOKENS")?, max_inflight_requests: opt_parse::(get, "DYN_EPP_MAX_INFLIGHT_REQUESTS")? .unwrap_or(DEFAULT_MAX_INFLIGHT_REQUESTS), + session_affinity_ttl_secs: opt_parse::(get, "DYN_EPP_SESSION_AFFINITY_TTL_SECS")?, }) } diff --git a/deploy/inference-gateway/ext-proc/src/selector.rs b/deploy/inference-gateway/ext-proc/src/selector.rs index 688c71b3cf62..2cee3a923861 100644 --- a/deploy/inference-gateway/ext-proc/src/selector.rs +++ b/deploy/inference-gateway/ext-proc/src/selector.rs @@ -38,6 +38,9 @@ pub struct SelectRequest { pub strict_priority: Option, pub expected_output_tokens: Option, pub policy_class: Option, + /// Session to pin (`x-dynamo-session-id`); the selector binds it to the + /// chosen worker when session affinity is enabled. + pub session_id: Option, } /// Observability overlap summary (matched token counts). @@ -112,6 +115,9 @@ impl Selector { if let Some(peer_replication) = peer_replication { builder = builder.replica_sync(peer_replication.sync_port, Vec::new()); } + if let Some(ttl) = cfg.session_affinity_ttl_secs { + builder = builder.session_affinity(std::time::Duration::from_secs_f64(ttl)); + } let service = Arc::new( builder .build() @@ -180,7 +186,7 @@ impl Selector { }, router_config_override: None, expected_output_tokens: req.expected_output_tokens, - session_id: None, + session_id: req.session_id, session_context: None, priority_jump: req.priority_jump, strict_priority: req.strict_priority, @@ -350,6 +356,7 @@ models: total_kv_blocks: None, max_num_batched_tokens: Some(8192), max_inflight_requests: 1024, + session_affinity_ttl_secs: None, } } @@ -405,6 +412,7 @@ models: strict_priority: None, expected_output_tokens: None, policy_class: None, + session_id: None, } } diff --git a/deploy/inference-gateway/ext-proc/src/topology_adapter.rs b/deploy/inference-gateway/ext-proc/src/topology_adapter.rs index 7bbd69ad71ec..d81708f6a238 100644 --- a/deploy/inference-gateway/ext-proc/src/topology_adapter.rs +++ b/deploy/inference-gateway/ext-proc/src/topology_adapter.rs @@ -151,6 +151,7 @@ mod tests { total_kv_blocks: Some(1000), max_num_batched_tokens: Some(8192), max_inflight_requests: 1024, + session_affinity_ttl_secs: None, } } diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-selection.md b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-selection.md index 8c1b5e8eb1a9..71d13f270f71 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-selection.md +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-selection.md @@ -69,6 +69,7 @@ APIs. Those bindings should wrap `SelectionService` rather than construct | `--remote-indexer-url` | none | Base URL of a [standalone indexer](standalone-indexer.md) that serves the primary KV index. The selector then does not subscribe to worker KV events; workers publish to the indexer instead. Requires `use_kv_events=true`. | | `--replica-sync-port` | none | Local ZMQ PUB port for active-load lifecycle events. The selector binds `tcp://*:` internally. | | `--replica-sync-peers` | none | Comma-separated ZMQ PUB endpoints for selector peers. Requires `--replica-sync-port`. | +| `--session-affinity-ttl-secs` | none | Pin each request `session_id` to the worker that served it for this long after its last request. Bindings replicate to peers over the replica mesh (`dynamo.session-affinity.v1` topic). | | `--selection-cache-ttl-secs` | `120` | Seconds an unclaimed pending selection lives before eviction. | | `--selection-cache-max-entries` | `4096` | Maximum resident pending selections, evicting oldest first. | | `--selection-cache-max-bytes` | `268435456` | Approximate byte budget across resident pending selections. | diff --git a/docs/fern/pages/kubernetes/kv-aware-routing/vanilla-vllm-onramp.mdx b/docs/fern/pages/kubernetes/kv-aware-routing/vanilla-vllm-onramp.mdx index 7faccae2c1a1..5e23deadd132 100644 --- a/docs/fern/pages/kubernetes/kv-aware-routing/vanilla-vllm-onramp.mdx +++ b/docs/fern/pages/kubernetes/kv-aware-routing/vanilla-vllm-onramp.mdx @@ -504,6 +504,7 @@ For a single EPP replica, set `spec.replicas: 1` and remove `DYN_EPP_PEER_SERVIC | `DYN_EPP_SELECTION_INDEXER_THREADS` | No | KV indexer thread count; defaults to `4`. | | `DYN_EPP_TOKENIZATION_TIMEOUT_MS` | No | Render request timeout; defaults to `5000`. | | `DYN_EPP_TOKENIZER_MAX_RESPONSE_BYTES` | No | Maximum render response size; defaults to `16777216`. | +| `DYN_EPP_SESSION_AFFINITY_TTL_SECS` | No | Pin each `x-dynamo-session-id` to the worker that served it for this many seconds after its last request; the binding replicates to peer EPPs over the replica mesh. Unset disables session affinity. | | `DYN_EPP_MAX_INFLIGHT_REQUESTS` | No | Concurrent EPP request guardrail; defaults to `1024`, and excess requests receive `503`. | | `DYN_EPP_PEER_SERVICE` | No | Existing EPP Service used to discover sibling replicas. Enables replica sync. | | `DYN_EPP_REPLICA_SYNC_PORT` | No | ZMQ port every replica binds and dials for replica sync; defaults to `9092`. All replicas selected by `DYN_EPP_PEER_SERVICE` must use the same value. | diff --git a/lib/bindings/python/rust/llm/kv.rs b/lib/bindings/python/rust/llm/kv.rs index 7e81d03790e9..b4ce9c621265 100644 --- a/lib/bindings/python/rust/llm/kv.rs +++ b/lib/bindings/python/rust/llm/kv.rs @@ -370,6 +370,11 @@ struct SelectServiceCli { #[arg(long, value_delimiter = ',', requires = "replica_sync_port")] replica_sync_peers: Vec, + /// Pin each session id (`session_id` on the request) to the worker that + /// served it for this many seconds after its last request + #[arg(long)] + session_affinity_ttl_secs: Option, + /// Seconds an unclaimed pending selection lives before eviction #[arg(long)] selection_cache_ttl_secs: Option, @@ -553,6 +558,9 @@ where remote_indexer_url: cli.remote_indexer_url, replica_sync_port: cli.replica_sync_port, replica_sync_peers: cli.replica_sync_peers, + session_affinity_ttl: cli + .session_affinity_ttl_secs + .map(std::time::Duration::from_secs_f64), kv_router_config, selection_cache: selection_cache_config_from_overrides( cli.selection_cache_ttl_secs, @@ -641,7 +649,7 @@ pub(crate) struct SelectionService { impl SelectionService { /// Create a selection service. `indexer_threads` sizes the KV indexer pool. #[new] - #[pyo3(signature = (*, indexer_threads = 4, indexer_peers = None, replica_sync_port = None, replica_sync_peers = None, selection_cache = None, remote_indexer_url = None))] + #[pyo3(signature = (*, indexer_threads = 4, indexer_peers = None, replica_sync_port = None, replica_sync_peers = None, selection_cache = None, remote_indexer_url = None, session_affinity_ttl_secs = None))] fn new( py: Python<'_>, indexer_threads: usize, @@ -650,6 +658,7 @@ impl SelectionService { replica_sync_peers: Option>, selection_cache: Option, remote_indexer_url: Option, + session_affinity_ttl_secs: Option, ) -> PyResult { let replica_sync_peers = replica_sync_peers.unwrap_or_default(); if replica_sync_port.is_none() && !replica_sync_peers.is_empty() { @@ -675,6 +684,9 @@ impl SelectionService { if let Some(url) = remote_indexer_url { builder = builder.kv_index(KvIndexSource::Remote(url)); } + if let Some(ttl) = session_affinity_ttl_secs { + builder = builder.session_affinity(std::time::Duration::from_secs_f64(ttl)); + } let inner = py .allow_threads(|| pyo3_async_runtimes::tokio::get_runtime().block_on(builder.build())) .map_err(to_pyerr)?; From f8432b356f1d0bdb227049e775330ebc7e236a6d Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Wed, 9 Sep 2026 01:26:38 -0700 Subject: [PATCH 032/102] docs(router): clarify partition session affinity ownership Signed-off-by: Ishan Dhanani --- .../router/standalone-selection.md | 31 ++++--------------- .../src/services/selection/README.md | 13 +++++--- 2 files changed, 14 insertions(+), 30 deletions(-) diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-selection.md b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-selection.md index 71d13f270f71..7dfbda2fd44d 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-selection.md +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-selection.md @@ -304,26 +304,11 @@ hash producer with the same algorithm, key, and key ID as the selector. ### `session_id` -Both `POST /select` and `POST /select_and_reserve` accept an optional -`session_id` string. It defaults to absent. Under the built-in selector, -omitting it does not change selection; a custom policy that reads the field can -select differently depending on whether it is present. The selector carries the -value through scheduling and exposes it to worker-selection policy as -`WorkerSelectionContext::session_id()`, so a custom picker or scorer can -implement session affinity by preferring the worker a session used previously. - -> [!NOTE] -> `session_id` is an input to policy, not an affinity mechanism in itself. The -> built-in selector ignores it, so it changes the chosen worker only when you -> supply a custom picker or scorer that reads it. See -> [Write Custom Routing Strategies](custom-worker-selection.mdx). -> It is also distinct from the frontend's own session affinity, which binds -> sessions from request headers rather than from this API; see -> [Configuration and Tuning](configuration-and-tuning.md). - -The selection service does not persist, replicate, or expire `session_id` -bindings. It is not part of the selection response and is not retained by the -pending-selection cache, so a `POST /reservations` replay does not carry it. +Both `POST /select` and `POST /select_and_reserve` accept an optional `session_id` string. With `--session-affinity-ttl-secs` enabled, each model and routing group has its own binding table. `/select_and_reserve` binds a new session to the selected worker and holds its lease until the reservation is released or expires. The idle TTL starts when the last lease is released. `/select` can use an existing binding but does not create one. + +Bindings replicate over the configured replica mesh. An explicit `affinity_target` or `pinned_worker` takes precedence over the session binding. Without a configured TTL, `session_id` is only policy input: custom policies can read it through `WorkerSelectionContext::session_id()`, while the built-in selector ignores it. See [Write Custom Routing Strategies](custom-worker-selection.mdx). + +The pending-selection cache does not retain session metadata, so a later `POST /reservations` does not create a session binding. Use `/select_and_reserve` when the service should manage affinity. The frontend uses the same table implementation for request-header affinity; see [Configuration and Tuning](configuration-and-tuning.md). ### `session_context` @@ -338,16 +323,12 @@ present, the flat `session_id` field is ignored. "session_id": "child-session", "parent_session_id": "root-session", "session_final": false, - "kv_hints": { "evict_session": false }, "input_trigger": "tool_result" } } ``` -Only `session_id` is required inside the object. `input_trigger` is one of -`user_message`, `tool_result`, or `other`. A custom policy reads the values -through `WorkerSelectionContext::session_context()`; the built-in selector -ignores them, exactly as it ignores `session_id`. +Only `session_id` is required inside the object. `input_trigger` is one of `user_message`, `tool_result`, or `other`. The affinity table uses this session ID when a TTL is configured. Custom policies can read the remaining values through `WorkerSelectionContext::session_context()`. ## Ray Select-Then-Reserve Flow diff --git a/lib/kv-router/src/services/selection/README.md b/lib/kv-router/src/services/selection/README.md index ef56760ef9ef..b6a4fbfa98eb 100644 --- a/lib/kv-router/src/services/selection/README.md +++ b/lib/kv-router/src/services/selection/README.md @@ -1,3 +1,8 @@ + + # Dynamo Selection Service The public deployment and HTTP API contract is documented in @@ -21,11 +26,9 @@ and active-sequence accounting. Keep these implementation invariants explicit: runtime event plane in the frontend) or `Remote` (a standalone indexer serves it). The `Indexer` type itself is shared with the frontend (`services::indexer::backend`). -- Session affinity (`services::selection::affinity::SessionAffinity`) is the - one session table: versioned bindings, TTL reaper, initialization/lease - lifecycle, hard and soft validation, replica apply. Hosts own the instance - and the replication transport (`AffinityReplicaSink`); the frontend wraps it - for pipeline streams and the runtime event plane. +- Each partition owns its `SessionAffinity` table, including versioned bindings, idle TTL, and lease lifecycle. A reservation owns its affinity lease, so every reservation removal releases both. Frontend routing hosts share the partition table and one coordinator for stream leases and runtime replication. +- Valid worker metadata updates preserve live bookings on surviving ranks and KV state from unchanged event sources. Catalog commits and ingress changes are serialized; partition policy factories can initialize independently. +- The frontend request lease manager owns expiry for embedded partitions. Standalone partitions use periodic request expiry. - Selector replicas synchronize admission, prefill-complete, and free events. - **NOTE:** Output-block updates remain local. They are deliberately excluded from replica sync because their frequency would consume disproportionate From cc4d75647091069fa5f51339c9f0cdfd71fdc197 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Wed, 9 Sep 2026 01:40:42 -0700 Subject: [PATCH 033/102] fix(router): preserve partition and reservation lifecycle in selection core Signed-off-by: Ishan Dhanani --- .../ext-proc/src/selector.rs | 18 +- lib/kv-router/src/sequences/multi_worker.rs | 8 +- .../src/services/common/replica_sync.rs | 18 +- .../src/services/indexer/registry.rs | 13 +- .../src/services/selection/affinity.rs | 15 +- .../src/services/selection/catalog.rs | 34 +- .../src/services/selection/core/mod.rs | 784 +++++++++++++----- .../src/services/selection/ingress.rs | 66 +- lib/kv-router/src/services/selection/mod.rs | 3 +- .../src/services/selection/service.rs | 3 + lib/llm/src/discovery/model_manager.rs | 6 +- lib/llm/src/entrypoint/input/common.rs | 14 +- lib/llm/src/kv_router.rs | 48 +- lib/llm/src/kv_router/embedded.rs | 27 +- lib/llm/src/kv_router/routing_host.rs | 4 +- lib/llm/src/kv_router/routing_load.rs | 23 +- lib/llm/src/kv_router/sequence.rs | 151 ++++ lib/llm/src/lora/filter.rs | 6 + lib/llm/src/session_affinity/coordinator.rs | 37 +- lib/llm/src/session_affinity/replica_sync.rs | 11 +- 20 files changed, 912 insertions(+), 377 deletions(-) diff --git a/deploy/inference-gateway/ext-proc/src/selector.rs b/deploy/inference-gateway/ext-proc/src/selector.rs index 2cee3a923861..a3c81bfebcdc 100644 --- a/deploy/inference-gateway/ext-proc/src/selector.rs +++ b/deploy/inference-gateway/ext-proc/src/selector.rs @@ -116,7 +116,10 @@ impl Selector { builder = builder.replica_sync(peer_replication.sync_port, Vec::new()); } if let Some(ttl) = cfg.session_affinity_ttl_secs { - builder = builder.session_affinity(std::time::Duration::from_secs_f64(ttl)); + builder = builder.session_affinity( + std::time::Duration::try_from_secs_f64(ttl) + .context("invalid session affinity TTL")?, + ); } let service = Arc::new( builder @@ -879,4 +882,17 @@ worker_selection: .await .expect("synthetic profile must ignore any policy_class value"); } + #[tokio::test] + async fn invalid_affinity_ttl_returns_configuration_error() { + for ttl in [-1.0, 0.0, 0.5, f64::NAN, f64::INFINITY] { + let mut cfg = test_config(); + cfg.session_affinity_ttl_secs = Some(ttl); + assert!( + Selector::new(&cfg, WorkerSelectionPolicyRegistry::default()) + .await + .is_err(), + "TTL={ttl}" + ); + } + } } diff --git a/lib/kv-router/src/sequences/multi_worker.rs b/lib/kv-router/src/sequences/multi_worker.rs index 7b63b7f24a7e..c115b82d1412 100644 --- a/lib/kv-router/src/sequences/multi_worker.rs +++ b/lib/kv-router/src/sequences/multi_worker.rs @@ -281,9 +281,9 @@ pub enum ReplicaWorkerPolicy { } #[derive(Debug, Clone, Copy)] -struct SequenceTrackerOptions { - replica_worker_policy: ReplicaWorkerPolicy, - expiry_duration: Option, +pub(crate) struct SequenceTrackerOptions { + pub(crate) replica_worker_policy: ReplicaWorkerPolicy, + pub(crate) expiry_duration: Option, } /// Errors that can occur during sequence management operations. @@ -443,7 +443,7 @@ impl ActiveSequencesMultiWorker

{ } /// Builds a tracker from resolved replica-admission and expiry policies. - fn new_with_options( + pub(crate) fn new_with_options( publisher: P, block_size: usize, dp_range: HashMap, diff --git a/lib/kv-router/src/services/common/replica_sync.rs b/lib/kv-router/src/services/common/replica_sync.rs index cc294d2e7630..a47408e63de3 100644 --- a/lib/kv-router/src/services/common/replica_sync.rs +++ b/lib/kv-router/src/services/common/replica_sync.rs @@ -31,6 +31,8 @@ const AFFINITY_EVENT_CHANNEL_CAPACITY: usize = 4_096; /// One replicated session binding. #[derive(Debug, Clone, Serialize, Deserialize)] pub(crate) struct AffinityBindingEvent { + #[serde(flatten)] + pub partition: RoutingPartitionId, pub session_id: String, pub worker_id: u64, pub dp_rank: Option, @@ -54,12 +56,14 @@ impl AffinityBindingEvent { /// Publishes bindings from a [`super::super::selection::affinity::SessionAffinity`] /// into the mesh. Best effort: a full channel drops the update. struct AffinityMeshSink { + partition: RoutingPartitionId, tx: mpsc::Sender, } impl AffinityReplicaSink for AffinityMeshSink { fn publish(&self, session_id: &str, target: AffinityTarget, version: AffinityVersion) { let update = AffinityBindingEvent { + partition: self.partition.clone(), session_id: session_id.to_string(), worker_id: target.worker_id, dp_rank: target.dp_rank, @@ -164,10 +168,16 @@ impl ReplicaSyncConfig { /// Sink that publishes session bindings into the mesh, when this runtime /// carries them. - pub(crate) fn affinity_sink(&self) -> Option> { - self.affinity_tx - .clone() - .map(|tx| Arc::new(AffinityMeshSink { tx }) as Arc) + pub(crate) fn affinity_sink( + &self, + partition: &RoutingPartitionId, + ) -> Option> { + self.affinity_tx.clone().map(|tx| { + Arc::new(AffinityMeshSink { + tx, + partition: partition.clone(), + }) as Arc + }) } pub(crate) fn is_self_event(&self, event: &ActiveSequenceEvent) -> bool { diff --git a/lib/kv-router/src/services/indexer/registry.rs b/lib/kv-router/src/services/indexer/registry.rs index 710bee3f6ca0..d8a31dee2238 100644 --- a/lib/kv-router/src/services/indexer/registry.rs +++ b/lib/kv-router/src/services/indexer/registry.rs @@ -311,6 +311,7 @@ pub struct WorkerEntry { pub struct WorkerRegistry { workers: DashMap, indexers: DashMap, + retain_empty_partitions: bool, peers: DashMap, watermarks: DashMap<(WorkerId, u32), Arc>, num_threads: usize, @@ -343,6 +344,7 @@ impl WorkerRegistry { Self::new_inner(num_threads, indexer_metrics, CancellationToken::new()) } + #[cfg(feature = "metrics")] pub(super) fn new_with_indexer_metrics_and_cancel_token( num_threads: usize, indexer_metrics: Arc, @@ -351,6 +353,12 @@ impl WorkerRegistry { Self::new_inner(num_threads, indexer_metrics, root_cancel_token) } + /// Selection partitions retain their index and overlap refresher across scale-to-zero. + pub(crate) fn retain_empty_partitions(mut self) -> Self { + self.retain_empty_partitions = true; + self + } + fn new_inner( num_threads: usize, indexer_metrics: Arc, @@ -360,6 +368,7 @@ impl WorkerRegistry { Self { workers: DashMap::new(), indexers: DashMap::new(), + retain_empty_partitions: false, peers: DashMap::new(), watermarks: DashMap::new(), num_threads, @@ -830,7 +839,9 @@ impl WorkerRegistry { } fn maybe_remove_indexer(&self, key: &RoutingPartitionId) { - if self.workers.iter().any(|entry| entry.value().key == *key) { + if self.retain_empty_partitions + || self.workers.iter().any(|entry| entry.value().key == *key) + { return; } diff --git a/lib/kv-router/src/services/selection/affinity.rs b/lib/kv-router/src/services/selection/affinity.rs index f78b3639994d..0f1ac4ef4a34 100644 --- a/lib/kv-router/src/services/selection/affinity.rs +++ b/lib/kv-router/src/services/selection/affinity.rs @@ -175,11 +175,7 @@ impl SessionAffinity { ) } - pub fn new_with_limits( - ttl: Duration, - max_entries: usize, - max_session_id_bytes: usize, - ) -> Result { + pub(crate) fn validate_ttl(ttl: Duration) -> Result<(), AffinityError> { if !(Duration::from_secs(1)..=Duration::from_secs(MAX_SESSION_AFFINITY_TTL_SECS)) .contains(&ttl) { @@ -187,6 +183,15 @@ impl SessionAffinity { "session affinity TTL must be between 1 and {MAX_SESSION_AFFINITY_TTL_SECS} seconds" ))); } + Ok(()) + } + + pub fn new_with_limits( + ttl: Duration, + max_entries: usize, + max_session_id_bytes: usize, + ) -> Result { + Self::validate_ttl(ttl)?; let inner = Arc::new(Inner { entries: DashMap::new(), ttl, diff --git a/lib/kv-router/src/services/selection/catalog.rs b/lib/kv-router/src/services/selection/catalog.rs index 39b3babee241..805911b154fd 100644 --- a/lib/kv-router/src/services/selection/catalog.rs +++ b/lib/kv-router/src/services/selection/catalog.rs @@ -8,10 +8,7 @@ use parking_lot::RwLock; use crate::identity::RoutingPartitionId; use crate::protocols::{WorkerId, WorkerWithDpRank}; -use super::error::SelectionError; -use super::types::{ - SelectionWorkerConfig, WorkerCatalogRecord, WorkerLifecycle, WorkerPatchRequest, WorkerRequest, -}; +use super::types::{SelectionWorkerConfig, WorkerCatalogRecord, WorkerLifecycle}; #[derive(Debug, Default)] pub(super) struct WorkerCatalog { @@ -19,33 +16,8 @@ pub(super) struct WorkerCatalog { } impl WorkerCatalog { - pub(super) fn upsert( - &self, - req: WorkerRequest, - ) -> (Option, WorkerCatalogRecord) { - let mut workers = self.workers.write(); - let previous = workers.get(&req.worker_id).cloned(); - let record = WorkerCatalogRecord::new(req); - workers.insert(record.worker_id, record.clone()); - (previous, record) - } - - pub(super) fn patch( - &self, - worker_id: WorkerId, - patch: WorkerPatchRequest, - ) -> Result<(WorkerCatalogRecord, WorkerCatalogRecord), SelectionError> { - let mut workers = self.workers.write(); - let Some(record) = workers.get_mut(&worker_id) else { - return Err(SelectionError::NotFound(format!( - "worker {worker_id} not found" - ))); - }; - let previous = record.clone(); - record.apply_patch(patch); - record.lifecycle = WorkerLifecycle::Incomplete; - record.not_schedulable_reasons.clear(); - Ok((previous, record.clone())) + pub(super) fn replace(&self, record: WorkerCatalogRecord) { + self.workers.write().insert(record.worker_id, record); } pub(super) fn get(&self, worker_id: WorkerId) -> Option { diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index 53bb8fa7fb9a..7a45463557bc 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -13,8 +13,7 @@ use tokio_util::sync::CancellationToken; use crate::identity::RoutingPartitionId; use crate::indexer::{ - KvRouterError, LowerTierQueryOptions, RoutingDecisionHashes, SharedKvCache, TieredMatchDetails, - TieredMatchProvider, + LowerTierQueryOptions, RoutingDecisionHashes, SharedKvCache, TieredMatchDetails, }; use crate::kv_hints::{ KvHint, KvHintAction, KvSourceLocationsPayload, KvTransferCandidateSource, KvTransferCandidates, @@ -32,7 +31,8 @@ use crate::scheduling::{ narrow_allowed_worker_ids_by_lora, prefill_load_hint_from_effective_tokens, }; use crate::sequences::{ - ActiveSequencesMultiWorker, ReplicaWorkerPolicy, SequenceError, SequenceRequest, + ActiveSequencesMultiWorker, ReplicaRequestLeaseObserver, ReplicaWorkerPolicy, SequenceError, + SequenceRequest, SequenceTrackerOptions, active_request_expiry_duration, }; use crate::services::common::replica_sync::{ HostReplicaSyncFactory, ReplicaSyncConfig, SchedulerLoadSink, ScopedReplicaEvent, @@ -44,7 +44,7 @@ use crate::services::indexer::registry::WorkerRegistry; use crate::services::overlap::MooncakeOverlapSummary; use crate::tracking_hash::{TrackingHashContext, TrackingHashScope}; -use super::affinity::{Acquired, AffinityInitialization, AffinityLease, SessionAffinity}; +use super::affinity::{Acquired, AffinityLease, SessionAffinity}; use super::catalog::WorkerCatalog; use super::error::SelectionError; use super::ingress::{KvEventIngress, ZmqDirectIngress}; @@ -60,45 +60,12 @@ use crate::WorkerSelectionPolicyFactory; use crate::WorkerType; use crate::services::common::replica_sync::AffinityBindingEvent; -/// Source of dequeue-time overlap refreshes for a partition scheduler: the -/// partition's index. -#[derive(Clone)] -pub enum RefreshProvider { - Local(Arc), -} - -#[async_trait::async_trait] -impl TieredMatchProvider for RefreshProvider { - async fn find_tiered_matches( - &self, - sequence: &[LocalBlockHash], - ) -> Result { - match self { - Self::Local(indexer) => indexer.find_tiered_matches(sequence.to_vec()).await, - } - } - - async fn find_tiered_matches_with_options( - &self, - sequence: &[LocalBlockHash], - options: LowerTierQueryOptions, - ) -> Result { - match self { - Self::Local(indexer) => { - indexer - .find_tiered_matches_with_options(sequence.to_vec(), options) - .await - } - } - } -} - /// The scheduler type every partition runs. pub type SelectionScheduler = LocalScheduler< ScopedSequencePublisher, SelectionWorkerConfig, WorkerSelectionPolicy, - TieredOverlapRefresher, + TieredOverlapRefresher, >; /// Handle to one partition's scheduler and indexer for an embedding host that @@ -122,6 +89,11 @@ impl SelectionPartition { pub fn indexer(&self) -> &Indexer { &self.0.indexer } + + /// Return this partition's affinity table, initialized with its serving TTL. + pub fn session_affinity(&self, ttl: Duration) -> Result { + self.0.session_affinity(ttl).cloned() + } } struct SelectionEntry { @@ -132,6 +104,8 @@ struct SelectionEntry { workers_tx: watch::Sender>, scheduler: SelectionScheduler, replica_tx: Option>, + affinity: OnceCell, + replica_config: Option, } struct PreparedSelectionInputs { @@ -143,6 +117,29 @@ struct PreparedSelectionInputs { kv_transfer_candidates: Option, } +impl SelectionEntry { + fn session_affinity(&self, ttl: Duration) -> Result<&SessionAffinity, SelectionError> { + let table = self + .affinity + .get_or_try_init(|| -> Result<_, SelectionError> { + let table = SessionAffinity::new(ttl).map_err(affinity_error)?; + if let Some(config) = &self.replica_config + && let Some(sink) = config.affinity_sink(&self.key) + { + table.enable_replication(config.process_id(), sink); + } + Ok(table) + })?; + if table.ttl() != ttl { + return Err(SelectionError::Conflict(format!( + "session affinity TTL mismatch for {}", + self.key + ))); + } + Ok(table) + } +} + struct SelectionOperation { key: RoutingPartitionId, selection_id: Option, @@ -259,6 +256,8 @@ pub struct HostTelemetry { #[derive(Clone, Default)] pub struct HostReplication { pub channels: Option, + /// Owns request expiry when supplied; the partition then expires only through lifecycle events. + pub request_leases: Option>, } impl std::fmt::Debug for SelectionHost { @@ -310,10 +309,17 @@ type SelectionEntries = RwLock>; +type ReservationIndex = RwLock>; + +struct Reservation { + partition: RoutingPartitionId, + _affinity_lease: Option, +} pub struct SelectionCore { catalog: WorkerCatalog, + /// Serializes catalog commits and the corresponding ingress changes. Never held by selection. + catalog_updates: tokio::sync::Mutex<()>, entries: Arc, reservation_index: Arc, /// Sweep task is started lazily from the first `ensure_entry`, which always @@ -334,19 +340,7 @@ pub struct SelectionCore { /// `create_reservation` can replay them without re-sending the prompt. selection_cache: SelectionCache, tracking_hash: Arc, - /// Session stickiness for request-shaped hosts; `None` when not configured. - session_affinity: Option, - /// Leases held by booked selections, released with the booking. - affinity_leases: parking_lot::Mutex>, -} - -/// What a booked selection holds on its session until the worker is known. -enum AffinityHold { - Initialize(AffinityInitialization), - Bound { - target: SessionTarget, - lease: AffinityLease, - }, + session_affinity_ttl: Option, } type SessionTarget = super::affinity::AffinityTarget; @@ -426,29 +420,10 @@ impl SelectionCore { session_affinity_ttl: Option, ) -> Self { let cancel_token = cancel_token.child_token(); - let session_affinity = match session_affinity_ttl { - Some(ttl) => match SessionAffinity::new(ttl) { - Ok(table) => { - if let Some(sink) = replica_config.as_ref().and_then(|c| c.affinity_sink()) { - let process_id = replica_config - .as_ref() - .map(|c| c.process_id()) - .unwrap_or_default(); - table.enable_replication(process_id, sink); - } - Some(table) - } - Err(error) => { - tracing::warn!(%error, "session affinity disabled"); - None - } - }, - None => None, - }; - let indexer_registry = Arc::new(WorkerRegistry::new_with_cancel_token( - indexer_threads, - cancel_token.clone(), - )); + let indexer_registry = Arc::new( + WorkerRegistry::new_with_cancel_token(indexer_threads, cancel_token.clone()) + .retain_empty_partitions(), + ); let listens_for_kv_events = kv_router_config.use_kv_events && !indexer_policy.is_remote(); indexer_registry.set_indexer_policy(indexer_policy); if signal_indexer_ready { @@ -456,6 +431,7 @@ impl SelectionCore { } Self { catalog: WorkerCatalog::default(), + catalog_updates: tokio::sync::Mutex::new(()), entries: Arc::new(RwLock::new(HashMap::new())), reservation_index: Arc::new(RwLock::new(HashMap::new())), reservation_sweep_started: OnceCell::new(), @@ -469,8 +445,7 @@ impl SelectionCore { replica_config, selection_cache: SelectionCache::new(&cache_config), tracking_hash, - session_affinity, - affinity_leases: parking_lot::Mutex::new(HashMap::new()), + session_affinity_ttl, } } @@ -554,8 +529,14 @@ impl SelectionCore { req: WorkerRequest, ) -> Result { self.ensure_running()?; - let (previous, record) = self.catalog.upsert(req); - self.reconcile_worker(record.worker_id, previous).await + let mut record = WorkerCatalogRecord::new(req); + // Partition policy factories may construct independently. Only committing + // membership and reconciling ingress needs the catalog mutation lock. + self.prepare_worker(&mut record)?; + let _update = self.catalog_updates.lock().await; + self.ensure_running()?; + let previous = self.catalog.get(record.worker_id); + self.reconcile_worker(record, previous).await } pub async fn patch_worker( @@ -563,16 +544,23 @@ impl SelectionCore { worker_id: WorkerId, patch: WorkerPatchRequest, ) -> Result { + let _update = self.catalog_updates.lock().await; self.ensure_running()?; - let (previous, record) = self.catalog.patch(worker_id, patch)?; - self.reconcile_worker(record.worker_id, Some(previous)) - .await + let previous = self + .catalog + .get(worker_id) + .ok_or_else(|| SelectionError::NotFound(format!("worker {worker_id} not found")))?; + let mut record = previous.clone(); + record.apply_patch(patch); + self.prepare_worker(&mut record)?; + self.reconcile_worker(record, Some(previous)).await } pub async fn delete_worker( &self, worker_id: WorkerId, ) -> Result { + let _update = self.catalog_updates.lock().await; let Some(previous) = self.catalog.get(worker_id) else { return Err(SelectionError::NotFound(format!( "worker {worker_id} not found" @@ -609,79 +597,67 @@ impl SelectionCore { } } - async fn reconcile_worker( - &self, - worker_id: WorkerId, - previous: Option, - ) -> Result { - let Some(record) = self.catalog.get(worker_id) else { - return Err(SelectionError::NotFound(format!( - "worker {worker_id} not found" - ))); - }; - - if previous - .as_ref() - .is_some_and(|record| record.lifecycle == WorkerLifecycle::Schedulable) - { - self.catalog - .set_lifecycle(worker_id, WorkerLifecycle::Draining, Vec::new()); - self.publish_scheduler_config(&previous.as_ref().expect("checked").key())?; - self.cleanup_indexer_registration(previous.as_ref().expect("checked")) - .await; - } - + fn prepare_worker(&self, record: &mut WorkerCatalogRecord) -> Result<(), SelectionError> { let queueing_enabled = self .kv_router_config .queueing_enabled(Some(&record.model_name)) .map_err(|error| SelectionError::BadRequest(error.to_string()))?; - let mut reasons = record.missing_schedulable_metadata(queueing_enabled); + record.not_schedulable_reasons = record.missing_schedulable_metadata(queueing_enabled); if let Some(ingress) = self.ingress() { - reasons.extend(ingress.missing_metadata(&record)); - } - if !reasons.is_empty() { - let updated = self - .catalog - .set_lifecycle(worker_id, WorkerLifecycle::Incomplete, reasons) - .ok_or_else(|| SelectionError::NotFound(format!("worker {worker_id} not found")))?; - self.publish_scheduler_config(&updated.key())?; - return Ok(updated); + record + .not_schedulable_reasons + .extend(ingress.missing_metadata(record)); } - - if let Err(error) = self.ensure_entry(&record) { - return self.mark_incomplete_after_reconcile_error(worker_id, record.key(), error); - } - if let Some(ingress) = self.ingress() - && let Err(error) = ingress.attach(&self.indexer_registry, &record).await + if record.not_schedulable_reasons.is_empty() + && let Err(error) = self.ensure_entry(record) { - self.cleanup_indexer_registration(&record).await; - return self.mark_incomplete_after_reconcile_error(worker_id, record.key(), error); + record + .not_schedulable_reasons + .push(format!("reconciliation failed: {error}")); } - let updated = self - .catalog - .set_lifecycle(worker_id, WorkerLifecycle::Schedulable, Vec::new()) - .ok_or_else(|| SelectionError::NotFound(format!("worker {worker_id} not found")))?; - self.publish_scheduler_config(&updated.key())?; - Ok(updated) + Ok(()) } - fn mark_incomplete_after_reconcile_error( + async fn reconcile_worker( &self, - worker_id: WorkerId, - key: RoutingPartitionId, - error: SelectionError, + mut record: WorkerCatalogRecord, + previous: Option, ) -> Result { - let updated = self - .catalog - .set_lifecycle( - worker_id, - WorkerLifecycle::Incomplete, - vec![format!("reconciliation failed: {error}")], - ) - .ok_or_else(|| SelectionError::NotFound(format!("worker {worker_id} not found")))?; - self.publish_scheduler_config(&key)?; - Ok(updated) + let previous = previous.filter(|old| old.lifecycle == WorkerLifecycle::Schedulable); + let previous = if let Some(old) = previous.as_ref() + && (old.key() != record.key() || !record.not_schedulable_reasons.is_empty()) + { + self.catalog + .set_lifecycle(old.worker_id, WorkerLifecycle::Draining, Vec::new()); + self.publish_scheduler_config(&old.key())?; + self.cleanup_indexer_registration(old).await; + None + } else { + previous + }; + + if record.not_schedulable_reasons.is_empty() + && let Some(ingress) = self.ingress() + && let Err(error) = ingress + .reconcile(&self.indexer_registry, previous.as_ref(), &record) + .await + { + self.cleanup_indexer_registration(&record).await; + record + .not_schedulable_reasons + .push(format!("reconciliation failed: {error}")); + } + record.lifecycle = if record.not_schedulable_reasons.is_empty() { + WorkerLifecycle::Schedulable + } else { + WorkerLifecycle::Incomplete + }; + // Readers see only committed metadata. A valid capacity/topology update preserves + // live bookings on ranks present in both the old and new snapshots. + self.catalog.replace(record.clone()); + self.publish_scheduler_config(&record.key())?; + Ok(record) } fn ensure_entry( @@ -752,7 +728,7 @@ impl SelectionCore { host_replica, ); let worker_label = self.worker_type.as_str(); - let slots = Arc::new(ActiveSequencesMultiWorker::new_with_replica_worker_policy( + let slots = Arc::new(ActiveSequencesMultiWorker::new_with_options( scoped_replica_sync .publisher .with_load_sink(self.host.telemetry.scheduler_load.clone()), @@ -761,15 +737,28 @@ impl SelectionCore { scoped_replica_sync.enabled, scoped_replica_sync.process_id, worker_label, - ReplicaWorkerPolicy::RequireRegistered, + SequenceTrackerOptions { + replica_worker_policy: ReplicaWorkerPolicy::RequireRegistered, + expiry_duration: self + .host + .replication + .request_leases + .is_none() + .then(active_request_expiry_duration), + }, )); + if let Some(observer) = &self.host.replication.request_leases { + slots.set_replica_request_lease_observer(Arc::clone(observer)); + } let replica_tx = scoped_replica_sync.channel.map(|(replica_tx, subscriber)| { slots.start_replica_sync(subscriber, self.cancel_token.child_token()); replica_tx }); - slots.start_periodic_force_expiry_across_all_workers( - self.cancel_token.child_token(), - ); + if self.host.replication.request_leases.is_none() { + slots.start_periodic_force_expiry_across_all_workers( + self.cancel_token.child_token(), + ); + } let indexer = match &self.host.cache.index { KvIndexSource::Owned(ingress) => { @@ -781,7 +770,7 @@ impl SelectionCore { }; let overlap_refresh = indexer.supports_overlap_refresh().then(|| { Arc::new(TieredOverlapRefresher::new( - RefreshProvider::Local(Arc::new(indexer.clone())), + indexer.clone(), self.kv_router_config.clone(), block_size, )) @@ -820,6 +809,8 @@ impl SelectionCore { workers_tx, scheduler, replica_tx, + affinity: OnceCell::new(), + replica_config: self.replica_config.clone(), })) })? .clone(); @@ -835,6 +826,9 @@ impl SelectionCore { entry.is_eagle ))); } + if let Some(ttl) = self.session_affinity_ttl { + entry.session_affinity(ttl)?; + } Ok(entry) } @@ -995,26 +989,24 @@ impl SelectionCore { // Session stickiness: a bound session steers selection (exclusive for // the default selector); a new session is bound to the worker chosen. // An explicit affinity target or pin from the caller wins. - let session_id = self - .session_affinity - .as_ref() + let table = entry.affinity.get(); + let session_id = table .and(session_context.as_ref()) .filter(|_| affinity_target.is_none() && pinned_worker.is_none()) .map(|context| context.session_id().to_string()); let mut affinity_hold = None; - let affinity_target = match (session_id.as_deref(), self.session_affinity.as_ref()) { + let affinity_target = match (session_id.as_deref(), table) { (Some(session_id), Some(table)) if book => { - match table - .acquire(session_id, None) - .await - .map_err(affinity_error)? - { + match tokio::select! { + _ = self.cancel_token.cancelled() => return Err(SelectionError::Scheduler(KvSchedulerError::SubscriberShutdown)), + result = table.acquire(session_id, None) => result.map_err(affinity_error)?, + } { Acquired::Initialize(init) => { - affinity_hold = Some(AffinityHold::Initialize(init)); + affinity_hold = Some(Acquired::Initialize(init)); None } Acquired::Bound { target, lease } => { - affinity_hold = Some(AffinityHold::Bound { target, lease }); + affinity_hold = Some(Acquired::Bound { target, lease }); Some(WorkerAffinityTarget::new(target.worker_id, target.dp_rank)) } } @@ -1136,13 +1128,6 @@ impl SelectionCore { } } => result?, }; - if let (Some(hold), Some(session_id), Some(selection_id)) = ( - affinity_hold.take(), - session_id.as_deref(), - selection_id.as_deref(), - ) { - self.bind_session(hold, session_id, selection_id, response.best_worker); - } let endpoint = self .catalog .schedulable_endpoint(response.best_worker.worker_id, &key) @@ -1199,7 +1184,13 @@ impl SelectionCore { }); if book && let Some(selection_id) = selection_id.as_deref() { - self.record_reservation(selection_id, &key); + let lease = match (affinity_hold, session_id.as_deref(), table) { + (Some(hold), Some(session_id), Some(table)) => { + Self::bind_session(table, hold, session_id, response.best_worker) + } + _ => None, + }; + self.record_reservation(selection_id, &key, lease); if let Some(hashes) = routing_hashes.clone() { self.record_routing_decision(&entry, response.best_worker, hashes) .await; @@ -1445,7 +1436,7 @@ impl SelectionCore { lora_name, }) .await?; - self.record_reservation(&selection_id, &key); + self.record_reservation(&selection_id, &key, None); if let Some(hashes) = routing_hashes { self.record_routing_decision(&entry, worker, hashes).await; } @@ -1460,41 +1451,45 @@ impl SelectionCore { }) } - fn record_reservation(&self, selection_id: &str, key: &RoutingPartitionId) { - self.reservation_index - .write() - .insert(selection_id.to_string(), key.clone()); + fn record_reservation( + &self, + selection_id: &str, + key: &RoutingPartitionId, + lease: Option, + ) { + self.reservation_index.write().insert( + selection_id.to_string(), + Reservation { + partition: key.clone(), + _affinity_lease: lease, + }, + ); } fn forget_reservation(&self, selection_id: &str) { self.reservation_index.write().remove(selection_id); - self.affinity_leases.lock().remove(selection_id); } /// Bind (or confirm) `session_id` to the worker a booking landed on and /// keep the lease with the booking. A bound session whose worker was not /// selected (it left) is invalidated and rebound. fn bind_session( - &self, - hold: AffinityHold, + table: &SessionAffinity, + hold: Acquired, session_id: &str, - selection_id: &str, worker: WorkerWithDpRank, - ) { - let Some(table) = self.session_affinity.as_ref() else { - return; - }; + ) -> Option { let selected = SessionTarget::new(worker.worker_id, Some(worker.dp_rank)); let lease = match hold { - AffinityHold::Initialize(init) => init.commit(selected).ok(), - AffinityHold::Bound { target, lease } + Acquired::Initialize(init) => init.commit(selected).ok(), + Acquired::Bound { target, lease } if target.worker_id == selected.worker_id && target.dp_rank.is_none_or(|rank| rank == worker.dp_rank) => { lease.publish(target); - return self.hold_lease(selection_id, lease); + return Some(lease); } - AffinityHold::Bound { mut lease, .. } => { + Acquired::Bound { mut lease, .. } => { lease.invalidate(); match table.try_acquire(session_id, None) { Ok(super::affinity::AcquireStep::Initialize(init)) => { @@ -1504,21 +1499,18 @@ impl SelectionCore { } } }; - if let Some(lease) = lease { + if let Some(lease) = &lease { lease.publish(selected); - self.hold_lease(selection_id, lease); } - } - - fn hold_lease(&self, selection_id: &str, lease: AffinityLease) { - self.affinity_leases - .lock() - .insert(selection_id.to_string(), lease); + lease } /// Apply a session binding a replica published. pub(crate) fn dispatch_affinity_event(&self, event: AffinityBindingEvent) { - let Some(table) = self.session_affinity.as_ref() else { + let Some(entry) = self.entry(&event.partition) else { + return; + }; + let Some(table) = entry.affinity.get() else { return; }; if self @@ -1529,7 +1521,11 @@ impl SelectionCore { return; } table.observe_replica_sequence(event.sequence); - if self.catalog.get(event.worker_id).is_none() { + if self + .catalog + .get(event.worker_id) + .is_none_or(|record| record.key() != event.partition) + { return; } let (target, version, worker_id) = (event.target(), event.version(), event.worker_id); @@ -1580,7 +1576,7 @@ impl SelectionCore { .reservation_index .read() .get(selection_id) - .cloned() + .map(|reservation| reservation.partition.clone()) .and_then(|key| self.entry(&key)); let mut entries = self.initialized_entries(); if let Some(indexed) = indexed { @@ -1828,35 +1824,16 @@ impl SelectionCore { /// scheduler (expired by the periodic force-expiry, or freed through a path /// that bypassed this core). Returns the number of entries removed. fn sweep_reservation_index(entries: &SelectionEntries, index: &ReservationIndex) -> usize { - let snapshot: Vec<(String, RoutingPartitionId)> = index - .read() - .iter() - .map(|(id, key)| (id.clone(), key.clone())) - .collect(); - if snapshot.is_empty() { - return 0; - } - let stale: Vec = { - let entries = entries.read(); - snapshot - .into_iter() - .filter(|(id, key)| { - entries - .get(key) - .and_then(|cell| cell.get()) - .is_none_or(|entry| !entry.scheduler.has_request(id)) - }) - .map(|(id, _)| id) - .collect() - }; - if stale.is_empty() { - return 0; - } + let entries = entries.read(); let mut index = index.write(); - stale - .iter() - .filter(|id| index.remove(id.as_str()).is_some()) - .count() + let before = index.len(); + index.retain(|id, reservation| { + entries + .get(&reservation.partition) + .and_then(|cell| cell.get()) + .is_some_and(|entry| entry.scheduler.has_request(id)) + }); + before - index.len() } fn spawn_reservation_index_sweep( @@ -3097,7 +3074,13 @@ mod tests { let mut request = reserve_request("booked"); request.routing_group = "group-b".to_string(); core.select_and_reserve(request).await.expect("reserve"); - assert_eq!(core.reservation_index.read().get("booked"), Some(&key_b)); + assert_eq!( + core.reservation_index + .read() + .get("booked") + .map(|r| &r.partition), + Some(&key_b) + ); assert_eq!( core.lifecycle_entries("booked")[0].key, key_b, @@ -3124,7 +3107,13 @@ mod tests { }) .await .expect("cached reservation"); - assert_eq!(core.reservation_index.read().get("cached"), Some(&key_b)); + assert_eq!( + core.reservation_index + .read() + .get("cached") + .map(|r| &r.partition), + Some(&key_b) + ); // Lifecycle calls still resolve, and free drops the index entry. core.prefill_complete("booked") @@ -3353,6 +3342,7 @@ mod tests { core.upsert_worker(worker(2)).await.expect("worker upsert"); core.dispatch_affinity_event(AffinityBindingEvent { + partition: RoutingPartitionId::new("model", "default"), session_id: "chat-b".to_string(), worker_id: 2, dp_rank: Some(0), @@ -3364,8 +3354,348 @@ mod tests { .await .expect("booking"); assert_eq!(response.worker_id, 2); - assert!(core.affinity_leases.lock().contains_key("r2")); + assert!( + core.reservation_index + .read() + .get("r2") + .unwrap() + ._affinity_lease + .is_some() + ); core.free_reservation("r2").await.expect("free"); - assert!(core.affinity_leases.lock().is_empty()); + assert!(core.reservation_index.read().is_empty()); + } + + #[tokio::test] + async fn expired_booking_releases_affinity_lease() { + let core = core_with_session_affinity(); + core.upsert_worker(worker(1)).await.unwrap(); + core.select_and_reserve(session_reservation("abandoned", "session")) + .await + .unwrap(); + let key = RoutingPartitionId::new("model", "default"); + core.entry(&key) + .unwrap() + .scheduler + .free("abandoned") + .await + .unwrap(); + assert_eq!( + sweep_reservation_index(&core.entries, &core.reservation_index), + 1 + ); + tokio::time::pause(); + tokio::time::advance(Duration::from_secs(11)).await; + assert_eq!( + core.entry(&RoutingPartitionId::new("model", "default")) + .unwrap() + .affinity + .get() + .unwrap() + .query_target("session", None) + .unwrap(), + None + ); + } + + #[tokio::test] + async fn partition_sessions_keep_independent_bindings() { + let core = core_with_session_affinity(); + core.upsert_worker(worker(1)).await.unwrap(); + let first = core + .select_and_reserve(session_reservation("first", "shared-session")) + .await + .unwrap(); + core.upsert_worker(worker(2)).await.unwrap(); + let mut other = worker(3); + other.routing_group = "other".to_string(); + core.upsert_worker(other).await.unwrap(); + let mut request = session_reservation("other", "shared-session"); + request.routing_group = "other".to_string(); + assert_eq!(core.select_and_reserve(request).await.unwrap().worker_id, 3); + assert_eq!( + core.entry(&RoutingPartitionId::new("model", "default")) + .unwrap() + .affinity + .get() + .unwrap() + .query_target("shared-session", None) + .unwrap() + .unwrap() + .worker_id, + first.worker_id + ); + let again = core + .select_and_reserve(session_reservation("again", "shared-session")) + .await + .unwrap(); + assert_eq!(again.worker_id, first.worker_id); + } + + #[tokio::test] + async fn rejoined_worker_feeds_partition_index() { + use crate::protocols::{BlockHashOptions, compute_block_hash_for_seq}; + for update_only in [false, true] { + let core = SelectionCore::try_new_local( + test_config(true), + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ) + .unwrap(); + let request = worker_with_kv_events(1); + core.upsert_worker(request.clone()).await.unwrap(); + let key = RoutingPartitionId::new("model", "default"); + let partition = core.partition(&key).unwrap(); + let tokens: Vec = (1..=8).collect(); + let hashes: Vec = + compute_block_hash_for_seq(&tokens, 4, BlockHashOptions::default()) + .into_iter() + .map(|hash| hash.0) + .collect(); + let indexer = core + .indexer_registry + .get_indexer(&key) + .unwrap() + .indexer + .clone(); + indexer + .apply_event_routed(store_event(1, 0, 1, &[], &hashes, StorageTier::Device)) + .await + .unwrap(); + indexer.dump_events().await.unwrap(); + if !update_only { + core.delete_worker(1).await.unwrap(); + } + let mut request = request; + request.total_kv_blocks = Some(2048); + core.upsert_worker(request).await.unwrap(); + let retained = partition.indexer().dump_events().await.unwrap(); + if update_only { + assert!( + !retained.is_empty(), + "metadata update cleared cached blocks" + ); + } else { + assert!(retained.is_empty(), "removed worker retained cached blocks"); + let current = core + .indexer_registry + .get_indexer(&key) + .unwrap() + .indexer + .clone(); + current + .apply_event_routed(store_event(1, 0, 1, &[], &hashes, StorageTier::Device)) + .await + .unwrap(); + current.dump_events().await.unwrap(); + } + let mut request = select_request(); + request.prompt = PromptRequest { + token_ids: Some(tokens), + ..PromptRequest::default() + }; + assert_eq!( + core.select(request).await.unwrap().overlap.longest_matched, + 8, + "update_only={update_only}" + ); + assert_eq!( + partition.indexer().dump_events().await.unwrap().len(), + indexer.dump_events().await.unwrap().len() + ); + } + } + #[tokio::test] + async fn metadata_update_preserves_live_booking() { + struct SuspendedDetach(tokio::sync::Notify); + #[async_trait::async_trait] + impl KvEventIngress for SuspendedDetach { + fn open( + &self, + registry: &WorkerRegistry, + key: &RoutingPartitionId, + block_size: u32, + ) -> Indexer { + registry.get_or_create_indexer(key.clone(), block_size) + } + async fn detach(&self, _registry: &WorkerRegistry, _record: &WorkerCatalogRecord) { + self.0.notify_one(); + std::future::pending().await + } + } + let ingress = Arc::new(SuspendedDetach(tokio::sync::Notify::new())); + let core = core_with_host(SelectionHost { + cache: HostCache { + index: KvIndexSource::Owned(ingress.clone()), + shared: None, + }, + ..SelectionHost::default() + }); + core.upsert_worker(worker(1)).await.unwrap(); + core.select_and_reserve(reserve_request("live")) + .await + .unwrap(); + let entry = core + .entry(&RoutingPartitionId::new("model", "default")) + .unwrap(); + let mut updated = worker(1); + updated.max_num_batched_tokens = Some(2048); + tokio::select! { + result = core.upsert_worker(updated) => { result.unwrap(); }, + _ = ingress.0.notified() => { + // A capacity update must not temporarily withdraw this worker. + assert!(entry.scheduler.has_request("live")); + panic!("capacity update detached the worker"); + } + } + assert!(entry.scheduler.has_request("live")); + core.free_reservation("live").await.unwrap(); + } + #[tokio::test(start_paused = true)] + async fn host_lease_manager_owns_expiry() { + use crate::scheduling::queue::SchedulerBookingDescriptor; + struct HostLeases; + impl ReplicaRequestLeaseObserver for HostLeases { + fn admitted(&self, _: SchedulerBookingDescriptor) {} + fn progressed(&self, _: &SchedulerBookingDescriptor) {} + fn completed(&self, _: &SchedulerBookingDescriptor) {} + } + for host_owned in [false, true] { + let core = core_with_host(SelectionHost { + replication: HostReplication { + request_leases: host_owned + .then(|| Arc::new(HostLeases) as Arc), + ..HostReplication::default() + }, + ..SelectionHost::default() + }); + core.upsert_worker(worker(1)).await.unwrap(); + core.select_and_reserve(reserve_request("live")) + .await + .unwrap(); + tokio::task::yield_now().await; + tokio::time::advance(active_request_expiry_duration() * 3).await; + tokio::task::yield_now().await; + let entry = core + .entry(&RoutingPartitionId::new("model", "default")) + .unwrap(); + assert_eq!(entry.scheduler.has_request("live"), host_owned); + if host_owned { + core.free_reservation("live").await.unwrap(); + assert!(!entry.scheduler.has_request("live")); + } + } + } + + #[tokio::test] + async fn catalog_updates_commit_in_order_without_exposing_candidates() { + struct PausedUpdate { + entered: tokio::sync::Notify, + release: tokio::sync::Notify, + } + #[async_trait::async_trait] + impl KvEventIngress for PausedUpdate { + fn open( + &self, + registry: &WorkerRegistry, + key: &RoutingPartitionId, + block_size: u32, + ) -> Indexer { + registry.get_or_create_indexer(key.clone(), block_size) + } + async fn reconcile( + &self, + _: &WorkerRegistry, + previous: Option<&WorkerCatalogRecord>, + _: &WorkerCatalogRecord, + ) -> Result<(), SelectionError> { + if previous.is_some() { + self.entered.notify_one(); + self.release.notified().await; + } + Ok(()) + } + } + let ingress = Arc::new(PausedUpdate { + entered: tokio::sync::Notify::new(), + release: tokio::sync::Notify::new(), + }); + let mut core = core_with_host(SelectionHost { + cache: HostCache { + index: KvIndexSource::Owned(ingress.clone()), + shared: None, + }, + ..SelectionHost::default() + }); + core.listens_for_kv_events = true; + core.upsert_worker(worker(1)).await.unwrap(); + core.select_and_reserve(reserve_request("live")) + .await + .unwrap(); + let mut updated = worker(1); + updated.max_num_batched_tokens = Some(2048); + let update = core.upsert_worker(updated); + tokio::pin!(update); + tokio::select! { + _ = &mut update => panic!("update should pause before committing"), + _ = ingress.entered.notified() => {} + } + let committed = core.catalog.get(1).unwrap(); + assert_eq!(committed.lifecycle, WorkerLifecycle::Schedulable); + assert_eq!(committed.max_num_batched_tokens, Some(1024)); + let entry = core.entry(&committed.key()).unwrap(); + assert!(entry.scheduler.has_request("live")); + let deletion = core.delete_worker(1); + tokio::pin!(deletion); + assert!( + std::future::Future::poll( + deletion.as_mut(), + &mut std::task::Context::from_waker(std::task::Waker::noop()) + ) + .is_pending() + ); + ingress.release.notify_one(); + assert_eq!(update.await.unwrap().max_num_batched_tokens, Some(2048)); + assert!(entry.scheduler.has_request("live")); + assert_eq!( + deletion.await.unwrap().lifecycle, + WorkerLifecycle::Unschedulable + ); + tokio::time::timeout(Duration::from_secs(1), async { + while entry.scheduler.has_request("live") { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + } + + #[tokio::test] + async fn affinity_configuration_rejects_invalid_or_conflicting_ttl() { + use super::super::service::SelectionServiceBuilder; + for ttl in [ + Duration::ZERO, + Duration::from_secs(super::super::affinity::MAX_SESSION_AFFINITY_TTL_SECS + 1), + ] { + let result = SelectionServiceBuilder::new( + test_config(false), + WorkerType::Aggregated, + Default::default(), + ) + .session_affinity(ttl) + .build() + .await; + assert!(result.is_err()); + } + let core = core_with_session_affinity(); + core.upsert_worker(worker(1)).await.unwrap(); + let partition = core + .partition(&RoutingPartitionId::new("model", "default")) + .unwrap(); + assert!(matches!( + partition.session_affinity(Duration::from_secs(20)), + Err(SelectionError::Conflict(_)) + )); } } diff --git a/lib/kv-router/src/services/selection/ingress.rs b/lib/kv-router/src/services/selection/ingress.rs index ac97f9c02807..ac67a1efa1f1 100644 --- a/lib/kv-router/src/services/selection/ingress.rs +++ b/lib/kv-router/src/services/selection/ingress.rs @@ -32,14 +32,16 @@ pub trait KvEventIngress: Send + Sync { Vec::new() } - /// Start feeding every rank of `record` into its partition's index. + /// Reconcile event sources, retaining unchanged ranks and their cached state. + /// `previous` is the last schedulable record in the same partition. /// Called only when the core subscribes to worker KV events. - async fn attach( + async fn reconcile( &self, registry: &WorkerRegistry, + previous: Option<&WorkerCatalogRecord>, record: &WorkerCatalogRecord, ) -> Result<(), SelectionError> { - let _ = (registry, record); + let _ = (registry, previous, record); Ok(()) } @@ -79,33 +81,61 @@ impl KvEventIngress for ZmqDirectIngress { .collect() } - async fn attach( + async fn reconcile( &self, registry: &WorkerRegistry, + previous: Option<&WorkerCatalogRecord>, record: &WorkerCatalogRecord, ) -> Result<(), SelectionError> { let block_size = record .block_size .ok_or_else(|| SelectionError::BadRequest("block_size is required".to_string()))?; - let mut endpoints: Vec<_> = record.listener_endpoints().into_iter().collect(); - endpoints.sort_by_key(|(dp_rank, _)| *dp_rank); - for (dp_rank, endpoint) in endpoints { - crate::services::common::zmq::validate_endpoint(&endpoint).map_err(|error| { + let endpoints = record.listener_endpoints(); + for (rank, endpoint) in &endpoints { + crate::services::common::zmq::validate_endpoint(endpoint).map_err(|error| { SelectionError::BadRequest(format!( - "invalid kv_events endpoint for worker {} dp_rank {dp_rank}: {error}", + "invalid kv_events endpoint for worker {} dp_rank {rank}: {error}", record.worker_id )) })?; - if let Some(replay_endpoint) = record.replay_endpoint.as_deref() { - crate::services::common::zmq::validate_endpoint(replay_endpoint).map_err( - |error| { - SelectionError::BadRequest(format!( - "invalid replay endpoint for worker {} dp_rank {dp_rank}: {error}", - record.worker_id - )) - }, - )?; + } + if let Some(endpoint) = record.replay_endpoint.as_deref() { + crate::services::common::zmq::validate_endpoint(endpoint).map_err(|error| { + SelectionError::BadRequest(format!( + "invalid replay endpoint for worker {}: {error}", + record.worker_id + )) + })?; + } + let ranks: std::collections::HashSet<_> = record.dp_ranks().collect(); + let mut retained = std::collections::HashSet::new(); + if let Some(previous) = previous { + let old_endpoints = previous.listener_endpoints(); + for rank in previous.dp_ranks() { + if ranks.contains(&rank) + && old_endpoints.get(&rank) == endpoints.get(&rank) + && previous.replay_endpoint == record.replay_endpoint + { + retained.insert(rank); + } else { + registry + .deregister_dp_rank( + record.worker_id, + rank, + &previous.model_name, + &previous.routing_group, + ) + .await + .map_err(|error| SelectionError::Internal(error.to_string()))?; + } } + } + let mut endpoints: Vec<_> = endpoints + .into_iter() + .filter(|(rank, _)| ranks.contains(rank) && !retained.contains(rank)) + .collect(); + endpoints.sort_by_key(|(rank, _)| *rank); + for (dp_rank, endpoint) in endpoints { registry .register( record.worker_id, diff --git a/lib/kv-router/src/services/selection/mod.rs b/lib/kv-router/src/services/selection/mod.rs index 15ab516d6edd..48504e82bafb 100644 --- a/lib/kv-router/src/services/selection/mod.rs +++ b/lib/kv-router/src/services/selection/mod.rs @@ -29,8 +29,7 @@ pub use crate::services::common::replica_sync::{ }; pub use core::{ HostCache, HostEligibility, HostLoad, HostReplication, HostTelemetry, KvIndexSource, - RefreshProvider, SelectionCore, SelectionHost, SelectionPartition, SelectionScheduler, - SelectionServiceConfig, + SelectionCore, SelectionHost, SelectionPartition, SelectionScheduler, SelectionServiceConfig, }; pub use error::SelectionError; pub use ingress::{KvEventIngress, ZmqDirectIngress}; diff --git a/lib/kv-router/src/services/selection/service.rs b/lib/kv-router/src/services/selection/service.rs index 2781cb620f2a..3925dba782d3 100644 --- a/lib/kv-router/src/services/selection/service.rs +++ b/lib/kv-router/src/services/selection/service.rs @@ -136,6 +136,9 @@ impl SelectionServiceBuilder { } pub async fn build(self) -> anyhow::Result { + if let Some(ttl) = self.session_affinity_ttl { + super::affinity::SessionAffinity::validate_ttl(ttl)?; + } self.kv_router_config .validate_config() .map_err(anyhow::Error::msg)?; diff --git a/lib/llm/src/discovery/model_manager.rs b/lib/llm/src/discovery/model_manager.rs index 56f50877b472..674f292a83b5 100644 --- a/lib/llm/src/discovery/model_manager.rs +++ b/lib/llm/src/discovery/model_manager.rs @@ -1926,13 +1926,9 @@ impl ModelManager { is_eagle: bool, ) -> anyhow::Result> { let client = endpoint.client().await?; - let source = crate::kv_router::RouterLoadSource::from_worker_role_or_metric( - worker_role, - metric_worker_type, - ); let parent_token = endpoint.component().drt().child_token(); let scheduler_load = - crate::kv_router::SchedulerLoadSender::disabled(source, parent_token.child_token()); + crate::kv_router::SchedulerLoadSender::disabled(parent_token.child_token()); self.kv_chooser_for_with_policy_and_client( client, kv_cache_block_size, diff --git a/lib/llm/src/entrypoint/input/common.rs b/lib/llm/src/entrypoint/input/common.rs index 0685dcae9269..2f9d2ed9b072 100644 --- a/lib/llm/src/entrypoint/input/common.rs +++ b/lib/llm/src/entrypoint/input/common.rs @@ -254,11 +254,15 @@ pub(crate) async fn build_preprocessed_routing_with_session_affinity_mode( wait_for_min_initial_workers(&router_client, min_initial_workers).await?; let endpoint_id = router_client.endpoint.id(); - let affinity = create_affinity_coordinator( - session_affinity_ttl_secs.map(Duration::from_secs), - router_client.clone(), - ) - .await?; + let ttl = session_affinity_ttl_secs.map(Duration::from_secs); + let affinity = match (ttl, chooser.as_ref()) { + (Some(ttl), Some(chooser)) => { + let affinity = chooser.affinity_coordinator(ttl)?; + affinity.enable_replica_sync(router_client.clone()).await?; + Some(affinity) + } + _ => create_affinity_coordinator(ttl, router_client.clone()).await?, + }; let embedding_cache_indexer = if enable_multimodal_cache_indexer && matches!(router_mode, RouterMode::DeviceAwareWeighted) diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index 4622d012d660..dd7dec912514 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -677,9 +677,8 @@ impl KvRouter { shared_cache: Option>, lora_filter: Option>, ) -> Result { - let source = RouterLoadSource::from_worker_role_or_metric(worker_role, metric_worker_type); let parent_token = endpoint.component().drt().child_token(); - let scheduler_load = SchedulerLoadSender::disabled(source, parent_token.child_token()); + let scheduler_load = SchedulerLoadSender::disabled(parent_token.child_token()); Self::new_with_worker_role_and_scheduler_load( endpoint, @@ -821,6 +820,7 @@ impl KvRouter { let available_worker_provider: WorkerAvailabilityProvider = Arc::new(move || client_for_availability.available_instance_ids()); + let replica_leases = Arc::new(sequence::DeferredReplicaRequestLeaseObserver::default()); let scheduler = embedded::EmbeddedSelection::start( embedded::EmbeddedSelectionArgs { kv_router_config: kv_router_config.clone(), @@ -838,6 +838,7 @@ impl KvRouter { endpoint: endpoint.clone(), router_id: endpoint.drt().discovery().instance_id(), policy_factory, + request_leases: replica_leases.clone(), }, workers_with_configs.clone(), cancellation_token.child_token(), @@ -847,6 +848,7 @@ impl KvRouter { scheduler.booking_cleanup(), cancellation_token.child_token(), ); + replica_leases.install(Arc::new(request_leases.clone())); tracing::info!("KV Routing initialized"); let cancellation_token = cancellation_guard.disarm(); Ok(Self { @@ -1229,35 +1231,16 @@ impl KvRouter { allowed_worker_ids: Option>, pinned_worker: Option<&WorkerWithDpRank>, ) -> Option> { - let (Some(filter), Some(lora_name)) = (self.lora_filter.as_ref(), lora_name) else { + let Some(filter) = self.lora_filter.as_ref() else { return allowed_worker_ids; }; - // Base candidate universe: explicit allow-set if present, else all current workers. - let base: Vec = match &allowed_worker_ids { - Some(allowed) => allowed.iter().copied().collect(), - None => self.workers_with_configs.borrow().keys().copied().collect(), - }; - if base.is_empty() { - return allowed_worker_ids; - } - let mut narrowed: HashSet = filter - .filter_worker_ids_for_lora(Some(lora_name), &base) - .into_iter() - .collect(); - // Retain a pinned worker only if it is already within the candidate universe — never - // widen the caller's `allowed_worker_ids` (KV-cache / EPP / migration invariants depend - // on that set). If the filter excluded an in-universe pinned worker, re-add it so the - // pin still wins for cache correctness; if the pin is outside the universe, honor the - // caller's constraint and drop it. - if let Some(p) = pinned_worker - && base.contains(&p.worker_id) - { - narrowed.insert(p.worker_id); - } - if narrowed.is_empty() { - return allowed_worker_ids; - } - Some(narrowed) + dynamo_kv_router::scheduling::narrow_allowed_worker_ids_by_lora( + filter.as_ref(), + lora_name, + allowed_worker_ids, + pinned_worker, + || self.workers_with_configs.borrow().keys().copied().collect(), + ) } /// Give these tokens, find the worker with the best weighted cache hit. @@ -1946,6 +1929,13 @@ impl KvRouter { } #[doc(hidden)] + pub(crate) fn affinity_coordinator( + &self, + ttl: std::time::Duration, + ) -> anyhow::Result { + self.scheduler.affinity_coordinator(ttl) + } + pub(crate) fn booking_cleanup(&self) -> scheduler::SchedulerBookingCleanup { self.scheduler.booking_cleanup() } diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs index 381d43ffe24c..65b497d85a82 100644 --- a/lib/llm/src/kv_router/embedded.rs +++ b/lib/llm/src/kv_router/embedded.rs @@ -3,14 +3,10 @@ //! Embedded selection backend for the frontend `KvRouter`. //! -//! The router runs its scheduling and active-sequence accounting on one -//! partition of an in-process `SelectionService` (the same core the standalone -//! selection service and the EPP use). The router keeps its own KV indexer and -//! request transport: it computes overlap against that indexer and hands the -//! partition scheduler a fully formed `ScheduleRequest`. The partition's worker -//! catalog is fed from the runtime-config watch, so discovery remains the source -//! of truth for membership while the scheduler and its replica sync are -//! runtime-free. +//! The runtime-config watch feeds one selection partition's worker catalog. +//! Runtime ingress feeds the partition's index, and the frontend prepares overlap +//! and request constraints before scheduling directly on that partition. The +//! frontend retains transport, stream leases, and request-expiry ownership. use std::collections::{HashMap, HashSet}; use std::sync::Arc; @@ -64,6 +60,7 @@ pub(crate) struct EmbeddedSelectionArgs { /// Builds the partition's worker-selection policy (see /// `SelectionPolicySource::resolve`). pub policy_factory: WorkerSelectionPolicyFactory, + pub request_leases: Arc, } /// Partition model name when the router has none. @@ -75,6 +72,7 @@ pub(crate) struct EmbeddedSelection { /// the router holds the partition. _service: Arc, partition: SelectionPartition, + affinity: std::sync::OnceLock, worker_type: &'static str, } @@ -163,6 +161,7 @@ impl EmbeddedSelection { }, replication: HostReplication { channels: replica_sync, + request_leases: Some(args.request_leases), }, }) .build() @@ -203,10 +202,22 @@ impl EmbeddedSelection { Ok(Self { _service: service, partition, + affinity: std::sync::OnceLock::new(), worker_type: args.metric_worker_type, }) } + pub(crate) fn affinity_coordinator( + &self, + ttl: std::time::Duration, + ) -> Result { + let table = self.partition.session_affinity(ttl)?; + Ok(self + .affinity + .get_or_init(|| crate::session_affinity::AffinityCoordinator::wrap(table)) + .clone()) + } + /// Membership is catalog-driven (runtime-config watch); explicit worker /// registration is a no-op here. pub(crate) fn register_workers(&self, _worker_ids: &HashSet) {} diff --git a/lib/llm/src/kv_router/routing_host.rs b/lib/llm/src/kv_router/routing_host.rs index 40268cf5defe..6a3537fdd19e 100644 --- a/lib/llm/src/kv_router/routing_host.rs +++ b/lib/llm/src/kv_router/routing_host.rs @@ -321,7 +321,7 @@ impl RoutingHost { session_affinity_ttl: Option, ) -> Result { let affinity = session_affinity_ttl - .map(AffinityCoordinator::new) + .map(|ttl| kv_router.affinity_coordinator(ttl)) .transpose()?; Ok(Self::new_with_coordinator( @@ -340,7 +340,7 @@ impl RoutingHost { session_affinity_mode: SessionAffinityMode, ) -> Result { let affinity = session_affinity_ttl - .map(AffinityCoordinator::new) + .map(|ttl| kv_router.affinity_coordinator(ttl)) .transpose()?; Ok(Self::new_with_load_context_and_coordinator( diff --git a/lib/llm/src/kv_router/routing_load.rs b/lib/llm/src/kv_router/routing_load.rs index f15e29c2e990..182b314af3d7 100644 --- a/lib/llm/src/kv_router/routing_load.rs +++ b/lib/llm/src/kv_router/routing_load.rs @@ -200,10 +200,7 @@ pub struct SchedulerLoadSender { } impl SchedulerLoadSender { - pub(crate) fn disabled( - _source: RouterLoadSource, - cancellation_token: CancellationToken, - ) -> Self { + pub(crate) fn disabled(cancellation_token: CancellationToken) -> Self { Self { wake_tx: None, shared: Arc::new(SchedulerLoadShared::new(SCHEDULER_LOAD_CHANNEL_CAPACITY)), @@ -273,18 +270,12 @@ impl SchedulerLoadReceiver { } pub(crate) fn scheduler_load_channel( - source: RouterLoadSource, cancellation_token: CancellationToken, ) -> (SchedulerLoadSender, SchedulerLoadReceiver) { - scheduler_load_channel_with_capacity( - source, - cancellation_token, - SCHEDULER_LOAD_CHANNEL_CAPACITY, - ) + scheduler_load_channel_with_capacity(cancellation_token, SCHEDULER_LOAD_CHANNEL_CAPACITY) } fn scheduler_load_channel_with_capacity( - _source: RouterLoadSource, cancellation_token: CancellationToken, capacity: usize, ) -> (SchedulerLoadSender, SchedulerLoadReceiver) { @@ -326,7 +317,7 @@ impl RoutingLoadContext { let cancellation_token = parent_token.child_token(); let (scheduler_load, monitor) = if source.monitors_sequence_load() { let (scheduler_load, scheduler_load_rx) = - scheduler_load_channel(source, cancellation_token.child_token()); + scheduler_load_channel(cancellation_token.child_token()); let monitor = KvWorkerMonitor::new( client.clone(), source, @@ -339,7 +330,7 @@ impl RoutingLoadContext { (scheduler_load, Some(monitor)) } else { ( - SchedulerLoadSender::disabled(source, cancellation_token.child_token()), + SchedulerLoadSender::disabled(cancellation_token.child_token()), None, ) }; @@ -442,8 +433,7 @@ mod tests { #[tokio::test] async fn saturated_channel_coalesces_batch_and_later_absolute_state_converges() { let token = CancellationToken::new(); - let (sender, mut receiver) = - scheduler_load_channel_with_capacity(RouterLoadSource::Decode, token, 1); + let (sender, mut receiver) = scheduler_load_channel_with_capacity(token, 1); sender.publish(snapshot(1, 90)); sender.publish_batch(vec![snapshot(1, 80), snapshot(2, 70)]); @@ -469,8 +459,7 @@ mod tests { #[tokio::test] async fn saturated_channel_preserves_queued_updates_before_coalesced_updates() { let token = CancellationToken::new(); - let (sender, mut receiver) = - scheduler_load_channel_with_capacity(RouterLoadSource::Decode, token, 2); + let (sender, mut receiver) = scheduler_load_channel_with_capacity(token, 2); sender.publish(snapshot(2, 20)); sender.publish(snapshot(1, 10)); diff --git a/lib/llm/src/kv_router/sequence.rs b/lib/llm/src/kv_router/sequence.rs index a61b3bc831a4..c0f6b9390a02 100644 --- a/lib/llm/src/kv_router/sequence.rs +++ b/lib/llm/src/kv_router/sequence.rs @@ -18,13 +18,16 @@ use dynamo_kv_router::protocols::{ pub use dynamo_kv_router::sequence::{ActiveSequences, RequestId}; use anyhow::Result; +use dynamo_kv_router::scheduling::queue::SchedulerBookingDescriptor; use dynamo_runtime::component::Endpoint; use dynamo_runtime::traits::DistributedRuntimeProvider; use dynamo_runtime::transports::event_plane::{ EventPublisher, EventSubscriber, EventTransportKind, TypedEventSubscriber, }; +use parking_lot::Mutex; use std::collections::VecDeque; use std::future::Future; +use std::sync::Arc; use std::task::{Context, Poll}; use tokio::sync::mpsc; use tokio::time::Instant; @@ -38,6 +41,107 @@ use dynamo_runtime::transports::event_plane::MsgpackCodec; // if the queue is full, the newest event is dropped without blocking the local mutation. const REPLICA_EVENT_CHANNEL_CAPACITY: usize = 100_000; +enum DeferredReplicaLeaseEvent { + Admitted(SchedulerBookingDescriptor), + Progressed(SchedulerBookingDescriptor), + Completed(SchedulerBookingDescriptor), +} + +enum DeferredReplicaLeaseState { + Buffering(Vec), + Installing { + target: Arc, + pending: Vec, + }, + Installed(Arc), +} + +impl Default for DeferredReplicaLeaseState { + fn default() -> Self { + Self::Buffering(Vec::new()) + } +} + +/// Buffers the construction-time event window until `KvRouter` installs its +/// request lease manager. After installation, calls pass through directly. +#[derive(Default)] +pub(crate) struct DeferredReplicaRequestLeaseObserver { + state: Mutex, +} + +impl DeferredReplicaRequestLeaseObserver { + pub(crate) fn install(&self, target: Arc) -> bool { + let mut pending = { + let mut state = self.state.lock(); + let DeferredReplicaLeaseState::Buffering(pending) = &mut *state else { + return false; + }; + let pending = std::mem::take(pending); + *state = DeferredReplicaLeaseState::Installing { + target: Arc::clone(&target), + pending: Vec::new(), + }; + pending + }; + + loop { + for event in pending { + Self::notify(&target, event); + } + let mut state = self.state.lock(); + let DeferredReplicaLeaseState::Installing { + target: installing_target, + pending: queued, + } = &mut *state + else { + unreachable!("observer installation state changed unexpectedly"); + }; + if queued.is_empty() { + *state = DeferredReplicaLeaseState::Installed(Arc::clone(installing_target)); + return true; + } + pending = std::mem::take(queued); + } + } + + fn observe(&self, event: DeferredReplicaLeaseEvent) { + let target = { + let mut state = self.state.lock(); + match &mut *state { + DeferredReplicaLeaseState::Buffering(pending) + | DeferredReplicaLeaseState::Installing { pending, .. } => { + pending.push(event); + return; + } + DeferredReplicaLeaseState::Installed(target) => Arc::clone(target), + } + }; + Self::notify(&target, event); + } + + fn notify(target: &Arc, event: DeferredReplicaLeaseEvent) { + match event { + DeferredReplicaLeaseEvent::Admitted(booking) => target.admitted(booking), + DeferredReplicaLeaseEvent::Progressed(booking) => target.progressed(&booking), + DeferredReplicaLeaseEvent::Completed(booking) => target.completed(&booking), + } + } +} + +impl ReplicaRequestLeaseObserver for DeferredReplicaRequestLeaseObserver { + fn admitted(&self, booking: SchedulerBookingDescriptor) { + self.observe(DeferredReplicaLeaseEvent::Admitted(booking)); + } + + fn progressed(&self, booking: &SchedulerBookingDescriptor) { + self.observe(DeferredReplicaLeaseEvent::Progressed(booking.clone())); + } + + fn completed(&self, booking: &SchedulerBookingDescriptor) { + self.observe(DeferredReplicaLeaseEvent::Completed(booking.clone())); + } +} + /// How active-sequence events are framed on the wire for a transport. #[derive(Clone, Copy, Debug, PartialEq, Eq)] enum ActiveSequenceEventWireFormat { @@ -715,4 +819,51 @@ mod tests { .is_err() ); } + #[test] + fn deferred_replica_observer_flushes_completed_startup_attempts_in_order() { + use dynamo_kv_router::sequences::NoopSequencePublisher; + use std::collections::HashMap; + #[derive(Default)] + struct Observer(Mutex>); + impl ReplicaRequestLeaseObserver for Observer { + fn admitted(&self, booking: SchedulerBookingDescriptor) { + self.0.lock().push(("add", booking)); + } + fn progressed(&self, booking: &SchedulerBookingDescriptor) { + self.0.lock().push(("progress", booking.clone())); + } + fn completed(&self, booking: &SchedulerBookingDescriptor) { + self.0.lock().push(("free", booking.clone())); + } + } + let slots = ActiveSequencesMultiWorker::new_without_expiry( + NoopSequencePublisher, + 4, + HashMap::from([(1, (0, 1))]), + true, + 0, + "test", + ); + let deferred = Arc::new(DeferredReplicaRequestLeaseObserver::default()); + assert!(slots.set_replica_request_lease_observer(deferred.clone())); + slots.apply_replica_batch(vec![add_event("startup"), free_event("startup")]); + let observer = Arc::new(Observer::default()); + assert!(deferred.install(observer.clone())); + { + let events = observer.0.lock(); + assert_eq!(events.len(), 2); + let booking = &events[0].1; + assert_eq!(booking.request_id, "startup"); + assert_eq!( + *events, + vec![("add", booking.clone()), ("free", booking.clone())] + ); + } + slots.apply_replica_batch(vec![add_event("after-install")]); + assert_eq!( + observer.0.lock().last().unwrap().1.request_id, + "after-install" + ); + assert!(!deferred.install(observer)); + } } diff --git a/lib/llm/src/lora/filter.rs b/lib/llm/src/lora/filter.rs index 6294c520ab12..8c697be2cb47 100644 --- a/lib/llm/src/lora/filter.rs +++ b/lib/llm/src/lora/filter.rs @@ -23,6 +23,12 @@ pub struct LoraFilter { state_tracker: LoraStateTracker, } +impl dynamo_kv_router::scheduling::LoraWorkerFilter for LoraFilter { + fn filter_worker_ids_for_lora(&self, lora_name: &str, available: &[WorkerId]) -> Vec { + self.filter_worker_ids_for_lora(Some(lora_name), available) + } +} + impl LoraFilter { pub fn new(routing_table: LoraRoutingTable, state_tracker: LoraStateTracker) -> Self { Self { diff --git a/lib/llm/src/session_affinity/coordinator.rs b/lib/llm/src/session_affinity/coordinator.rs index 2fe9ff943528..d8d5b69ef918 100644 --- a/lib/llm/src/session_affinity/coordinator.rs +++ b/lib/llm/src/session_affinity/coordinator.rs @@ -9,7 +9,7 @@ use std::{ pin::Pin, - sync::{Arc, OnceLock}, + sync::Arc, task::{Context, Poll}, time::Duration, }; @@ -56,7 +56,7 @@ fn from_table(target: TableTarget) -> AffinityTarget { struct Inner { table: SessionAffinity, - replica: OnceLock, + replica: tokio::sync::OnceCell, } impl Drop for Inner { @@ -79,11 +79,11 @@ impl AffinityCoordinator { )) } - fn wrap(table: SessionAffinity) -> Self { + pub(crate) fn wrap(table: SessionAffinity) -> Self { Self { inner: Arc::new(Inner { table, - replica: OnceLock::new(), + replica: tokio::sync::OnceCell::new(), }), } } @@ -92,21 +92,24 @@ impl AffinityCoordinator { &self, client: dynamo_runtime::component::Client, ) -> Result<(), Error> { - let (replica, router_id) = - ReplicaSyncRuntime::start(client, self.inner.table.downgrade()).await?; - if !self - .inner - .table - .enable_replication(router_id, replica.sink()) - { - return Err(anyhow::anyhow!( - "session affinity replica sync already enabled" - )); - } self.inner .replica - .set(replica) - .map_err(|_| anyhow::anyhow!("session affinity replica sync already enabled")) + .get_or_try_init(|| async { + let (replica, router_id) = + ReplicaSyncRuntime::start(client, self.inner.table.downgrade()).await?; + if !self + .inner + .table + .enable_replication(router_id, replica.sink()) + { + return Err(anyhow::anyhow!( + "session affinity replica sync already enabled" + )); + } + Ok(replica) + }) + .await?; + Ok(()) } /// Take a slot without cancelling on the request context. diff --git a/lib/llm/src/session_affinity/replica_sync.rs b/lib/llm/src/session_affinity/replica_sync.rs index 06ed7f4c9e68..bf3ef835a37d 100644 --- a/lib/llm/src/session_affinity/replica_sync.rs +++ b/lib/llm/src/session_affinity/replica_sync.rs @@ -394,10 +394,19 @@ mod tests { let client = endpoint.client().await.unwrap(); let original = AffinityCoordinator::new(Duration::from_secs(10)).unwrap(); - original.enable_replica_sync(client.clone()).await.unwrap(); + let shared = original.clone(); + let (first, second) = tokio::join!( + original.enable_replica_sync(client.clone()), + shared.enable_replica_sync(client.clone()), + ); + first.unwrap(); + second.unwrap(); wait_for_registration_count(&drt, &query, 1).await; drop(original); + shared.enable_replica_sync(client.clone()).await.unwrap(); + wait_for_registration_count(&drt, &query, 1).await; + drop(shared); wait_for_registration_count(&drt, &query, 0).await; let replacement = AffinityCoordinator::new(Duration::from_secs(10)).unwrap(); From b91015f183bc159e1fbf6968f2cc967cbbc0b23f Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Wed, 9 Sep 2026 11:48:40 -0700 Subject: [PATCH 034/102] fix(bindings): reject invalid selection affinity TTLs Signed-off-by: Ishan Dhanani --- lib/bindings/python/rust/llm/kv.rs | 24 ++++++++-- lib/bindings/python/src/dynamo/_core.pyi | 8 +++- lib/bindings/python/tests/test_kv_bindings.py | 46 +++++++++++++++++++ 3 files changed, 74 insertions(+), 4 deletions(-) diff --git a/lib/bindings/python/rust/llm/kv.rs b/lib/bindings/python/rust/llm/kv.rs index b4ce9c621265..1fdc028be3d6 100644 --- a/lib/bindings/python/rust/llm/kv.rs +++ b/lib/bindings/python/rust/llm/kv.rs @@ -560,7 +560,9 @@ where replica_sync_peers: cli.replica_sync_peers, session_affinity_ttl: cli .session_affinity_ttl_secs - .map(std::time::Duration::from_secs_f64), + .map(session_affinity_ttl_from_secs) + .transpose() + .map_err(anyhow::Error::msg)?, kv_router_config, selection_cache: selection_cache_config_from_overrides( cli.selection_cache_ttl_secs, @@ -637,6 +639,18 @@ impl SelectionCacheConfig { } } +#[cfg(feature = "select-service")] +fn session_affinity_ttl_from_secs(ttl: f64) -> Result { + use selection::affinity::MAX_SESSION_AFFINITY_TTL_SECS; + + if !(1.0..=MAX_SESSION_AFFINITY_TTL_SECS as f64).contains(&ttl) { + return Err(format!( + "session_affinity_ttl_secs must be between 1 and {MAX_SESSION_AFFINITY_TTL_SECS}" + )); + } + Ok(Duration::from_secs_f64(ttl)) +} + /// In-process handle to a managed Dynamo `SelectionService`. #[cfg(feature = "select-service")] #[pyclass] @@ -650,6 +664,7 @@ impl SelectionService { /// Create a selection service. `indexer_threads` sizes the KV indexer pool. #[new] #[pyo3(signature = (*, indexer_threads = 4, indexer_peers = None, replica_sync_port = None, replica_sync_peers = None, selection_cache = None, remote_indexer_url = None, session_affinity_ttl_secs = None))] + #[allow(clippy::too_many_arguments)] fn new( py: Python<'_>, indexer_threads: usize, @@ -685,7 +700,9 @@ impl SelectionService { builder = builder.kv_index(KvIndexSource::Remote(url)); } if let Some(ttl) = session_affinity_ttl_secs { - builder = builder.session_affinity(std::time::Duration::from_secs_f64(ttl)); + builder = builder.session_affinity( + session_affinity_ttl_from_secs(ttl).map_err(PyValueError::new_err)?, + ); } let inner = py .allow_threads(|| pyo3_async_runtimes::tokio::get_runtime().block_on(builder.build())) @@ -974,8 +991,9 @@ mod selection_service_lifecycle_tests { #[test] fn idempotent_shutdown() { + pyo3::prepare_freethreaded_python(); let service = - Python::with_gil(|py| SelectionService::new(py, 1, None, None, None, None, None)) + Python::with_gil(|py| SelectionService::new(py, 1, None, None, None, None, None, None)) .unwrap(); Python::with_gil(|py| { service.shutdown(py); diff --git a/lib/bindings/python/src/dynamo/_core.pyi b/lib/bindings/python/src/dynamo/_core.pyi index 8e11a3a45dfa..afb6da94d4ec 100644 --- a/lib/bindings/python/src/dynamo/_core.pyi +++ b/lib/bindings/python/src/dynamo/_core.pyi @@ -748,8 +748,14 @@ class SelectionService: replica_sync_port: Optional[int] = None, replica_sync_peers: Optional[list[str]] = None, selection_cache: Optional[SelectionCacheConfig] = None, + remote_indexer_url: Optional[str] = None, + session_affinity_ttl_secs: Optional[float] = None, ) -> None: - """Create a selection service. `indexer_threads` sizes the KV indexer pool.""" + """Create a selection service. `indexer_threads` sizes the KV indexer pool. + + `session_affinity_ttl_secs` enables session affinity with an idle TTL + between 1 and 31536000 seconds. + """ ... def shutdown(self) -> None: diff --git a/lib/bindings/python/tests/test_kv_bindings.py b/lib/bindings/python/tests/test_kv_bindings.py index 573f7c462c9a..72e85018f3c5 100644 --- a/lib/bindings/python/tests/test_kv_bindings.py +++ b/lib/bindings/python/tests/test_kv_bindings.py @@ -15,6 +15,8 @@ import json +import subprocess +import sys import threading import pytest @@ -201,6 +203,50 @@ def worker(worker_id, prepopulate_worker_ids: bool = False): ), f"Expected {expected_blocks_after_removal} block events after removal, got {len(blocks_after_removal)}" +@pytest.mark.skipif( + SelectionService is None, + reason="SelectionService requires the select-service Cargo feature", +) +@pytest.mark.parametrize( + "ttl", [-1, 0, 0.5, float("nan"), float("inf"), 31536001, 1e300] +) +def test_selection_service_rejects_invalid_affinity_ttl(ttl): + with pytest.raises(ValueError, match="session_affinity_ttl_secs must be between"): + SelectionService(session_affinity_ttl_secs=ttl) + + +@pytest.mark.skipif( + SelectionService is None, + reason="SelectionService requires the select-service Cargo feature", +) +@pytest.mark.parametrize("ttl", ["-1", "NaN", "31536001"]) +def test_selection_service_cli_rejects_invalid_affinity_ttl(ttl): + result = subprocess.run( + [ + sys.executable, + "-m", + "dynamo.select_service", + f"--session-affinity-ttl-secs={ttl}", + ], + capture_output=True, + text=True, + timeout=15, + ) + assert result.returncode != 0 + assert "session_affinity_ttl_secs must be between" in result.stderr + assert "panicked" not in result.stderr + + +@pytest.mark.skipif( + SelectionService is None, + reason="SelectionService requires the select-service Cargo feature", +) +@pytest.mark.parametrize("ttl", [1.0, 1.5, 31536000.0]) +def test_selection_service_accepts_valid_affinity_ttl(ttl): + service = SelectionService(indexer_threads=1, session_affinity_ttl_secs=ttl) + service.shutdown() + + @pytest.mark.skipif( SelectionService is None, reason="SelectionService requires the select-service Cargo feature", From efb4c7e7bd9a8800a0e6b4cdee5b5bcbd8894df5 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Wed, 9 Sep 2026 23:04:48 -0700 Subject: [PATCH 035/102] fix(router): drop the reservation-index guard before the partition lookup lifecycle_entries held the reservation_index read guard while entry() took entries.read(); the sweep takes entries.read() then reservation_index.write(). With a partition creation queued on entries.write(), parking_lot's writer priority turned that into a three-way deadlock that wedged every select. Bind the indexed partition first so the guard drops before entries is touched, and add a regression test that reproduces the interleaving. Signed-off-by: Ishan Dhanani --- .../src/services/selection/core/mod.rs | 71 ++++++++++++++++++- 1 file changed, 68 insertions(+), 3 deletions(-) diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index 7a45463557bc..c6f62b6d21d5 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -321,6 +321,7 @@ pub struct SelectionCore { /// Serializes catalog commits and the corresponding ingress changes. Never held by selection. catalog_updates: tokio::sync::Mutex<()>, entries: Arc, + /// Lock order: `entries` before `reservation_index`, never nested the other way. reservation_index: Arc, /// Sweep task is started lazily from the first `ensure_entry`, which always /// runs inside the host runtime; construction itself may not. @@ -1572,12 +1573,15 @@ impl SelectionCore { /// covers bookings mirrored from replica peers, which never pass through /// this core's booking paths. fn lifecycle_entries(&self, selection_id: &str) -> Vec> { - let indexed = self + // Release the index guard before taking `entries`: the sweep takes + // `entries` then `reservation_index`, so nesting here would invert the + // lock order. + let indexed_partition = self .reservation_index .read() .get(selection_id) - .map(|reservation| reservation.partition.clone()) - .and_then(|key| self.entry(&key)); + .map(|reservation| reservation.partition.clone()); + let indexed = indexed_partition.and_then(|key| self.entry(&key)); let mut entries = self.initialized_entries(); if let Some(indexed) = indexed { entries.retain(|entry| !Arc::ptr_eq(entry, &indexed)); @@ -1941,6 +1945,8 @@ mod tests { use super::*; use crate::protocols::StorageTier; use crate::services::indexer::backend::test_util::store_event; + use std::sync::atomic::{AtomicBool, Ordering}; + use std::thread::sleep; use std::time::Duration; fn test_config(use_kv_events: bool) -> crate::config::KvRouterConfig { @@ -3169,6 +3175,65 @@ mod tests { assert!(index.contains_key("live")); } + #[tokio::test(flavor = "current_thread")] + async fn lifecycle_lookup_does_not_nest_reservation_index_inside_entries() { + // Three parties: a sweep holding `entries` and wanting `reservation_index`, + // a lifecycle call, and a partition creation queued on `entries.write()`. + let core = Arc::new( + SelectionCore::try_new_local( + test_config(false), + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ) + .expect("valid test config"), + ); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + core.select_and_reserve(reserve_request("live")) + .await + .expect("reserve live"); + let deadline = Instant::now() + Duration::from_secs(5); + + let sweep_entries = core.entries.read(); + let writer = { + let core = Arc::clone(&core); + std::thread::spawn(move || { + core.entries + .write() + .entry(RoutingPartitionId::new("other", "default")) + .or_insert_with(|| Arc::new(OnceCell::new())); + }) + }; + while core.entries.try_read().is_some() { + assert!(Instant::now() < deadline, "partition writer never queued"); + std::thread::yield_now(); + } + let lifecycle_started = Arc::new(AtomicBool::new(false)); + let lifecycle = { + let core = Arc::clone(&core); + let started = Arc::clone(&lifecycle_started); + std::thread::spawn(move || { + started.store(true, Ordering::Release); + core.lifecycle_entries("live").len() + }) + }; + while !lifecycle_started.load(Ordering::Acquire) { + assert!(Instant::now() < deadline, "lifecycle thread never started"); + std::thread::yield_now(); + } + sleep(Duration::from_millis(50)); + + // With the index guard held across `entries.read()` this never acquires. + drop( + core.reservation_index + .try_write_for(Duration::from_secs(2)) + .expect("reservation index must not be held by a blocked lifecycle call"), + ); + drop(sweep_entries); + writer.join().expect("partition writer"); + assert_eq!(lifecycle.join().expect("lifecycle lookup"), 1); + } + #[tokio::test(flavor = "current_thread")] async fn queued_selection_returns_refreshed_overlap_snapshot() { let mut config = test_config(false); From 533ff450aa597ff8cfed965e9ce0716c9fb41ca8 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Wed, 9 Sep 2026 23:20:06 -0700 Subject: [PATCH 036/102] fix(router): consume worker completion marks without replica sync The embedded partition only opened its active-sequence ingress when router_replica_sync was set, so worker-origin MarkPrefillCompleted events were never received in the default configuration. Always build the inbound channel and gate only the outbound publisher on the flag; a failed inbound subscribe degrades to response-side cleanup with a warning, as before. Restore the direct-ZMQ active-sequence ingress (default on for brokerless ZMQ) with its receive-HWM knob and continuity metrics, forwarding decoded events into the partition's inbound channel; a full channel drops the event and counts it. Signed-off-by: Ishan Dhanani --- .../src/services/common/replica_sync.rs | 32 +- lib/llm/src/kv_router/embedded.rs | 32 +- lib/llm/src/kv_router/metrics.rs | 129 ++++++++ lib/llm/src/kv_router/sequence.rs | 175 ++++++++--- lib/llm/src/kv_router/sequence/direct_zmq.rs | 283 ++++++++++++++++++ 5 files changed, 586 insertions(+), 65 deletions(-) create mode 100644 lib/llm/src/kv_router/sequence/direct_zmq.rs diff --git a/lib/kv-router/src/services/common/replica_sync.rs b/lib/kv-router/src/services/common/replica_sync.rs index a47408e63de3..e9412ff4eb36 100644 --- a/lib/kv-router/src/services/common/replica_sync.rs +++ b/lib/kv-router/src/services/common/replica_sync.rs @@ -211,8 +211,11 @@ pub trait SchedulerLoadSink: Send + Sync { /// Dynamo runtime event plane) instead of the service's ZMQ peer mesh: /// events the partition emits go to `outbound`; events from peer replicas /// arrive on `inbound_rx` (the host also holds `inbound_tx`). +/// +/// `outbound: None` is ingress-only: the partition publishes nothing and the +/// tracker keeps only worker-origin completion marks from `inbound_rx`. pub struct HostReplicaChannels { - pub outbound: mpsc::Sender, + pub outbound: Option>, pub inbound_tx: mpsc::Sender, pub inbound_rx: mpsc::Receiver, /// This replica's id; events carrying it are ignored on receipt. @@ -423,9 +426,13 @@ pub(crate) fn setup_scoped_replica_sync( let Some(config) = config else { // No peer mesh: an embedding host may still carry replica events. if let Some(host) = host { + let enabled = host.outbound.is_some(); return ScopedReplicaSync { - publisher: ScopedSequencePublisher::host(host.outbound), - enabled: true, + publisher: host + .outbound + .map(ScopedSequencePublisher::host) + .unwrap_or_else(ScopedSequencePublisher::disabled), + enabled, process_id: host.process_id, channel: Some(( host.inbound_tx, @@ -1041,4 +1048,23 @@ mod tests { drop(listener); endpoint } + + #[test] + fn ingress_only_host_channels_keep_replica_sync_disabled() { + let (inbound_tx, inbound_rx) = mpsc::channel(1); + let scoped = setup_scoped_replica_sync( + None, + &RoutingPartitionId::new("model", "default"), + 16, + Some(HostReplicaChannels { + outbound: None, + inbound_tx, + inbound_rx, + process_id: 7, + }), + ); + assert!(!scoped.enabled); + assert_eq!(scoped.process_id, 7); + assert!(scoped.channel.is_some()); + } } diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs index 65b497d85a82..c163c01cefcb 100644 --- a/lib/llm/src/kv_router/embedded.rs +++ b/lib/llm/src/kv_router/embedded.rs @@ -119,24 +119,22 @@ impl EmbeddedSelection { DEFAULT_ROUTING_GROUP, ); - // Replica sync rides the runtime event plane, as the runtime scheduler's - // did; the partition publishes and consumes through host channels. + // Replica sync rides the runtime event plane. Worker-origin completion + // marks are consumed even when router-to-router replica sync is + // disabled; only publishing is gated. + let channels = crate::kv_router::sequence::host_replica_channels( + &args.endpoint, + args.router_id, + args.kv_router_config.router_replica_sync, + cancellation_token.child_token(), + ) + .await + .context("start replica sync for the embedded selection partition")?; + let slot = std::sync::Mutex::new(Some(channels)); let replica_sync: Option = - if args.kv_router_config.router_replica_sync { - let channels = crate::kv_router::sequence::host_replica_channels( - &args.endpoint, - args.router_id, - cancellation_token.child_token(), - ) - .await - .context("start replica sync for the embedded selection partition")?; - let slot = std::sync::Mutex::new(Some(channels)); - Some(Arc::new(move |_partition| { - slot.lock().ok().and_then(|mut s| s.take()) - })) - } else { - None - }; + Some(Arc::new(move |_partition| { + slot.lock().ok().and_then(|mut s| s.take()) + })); let service = SelectionServiceBuilder::new( args.kv_router_config.clone(), diff --git a/lib/llm/src/kv_router/metrics.rs b/lib/llm/src/kv_router/metrics.rs index 245f29a862bd..a927bc6eaa32 100644 --- a/lib/llm/src/kv_router/metrics.rs +++ b/lib/llm/src/kv_router/metrics.rs @@ -282,6 +282,135 @@ impl KvZmqIngressMetrics { } } +// --------------------------------------------------------------------------- +// Direct-ZMQ active-sequence ingress metrics +// --------------------------------------------------------------------------- + +pub(crate) struct ActiveSequenceZmqIngressMetrics { + tracked_sources: IntGauge, + sequence_gaps_total: IntCounter, + out_of_order_total: IntCounter, + reconnects_total: IntCounter, + replacements_total: IntCounter, + envelope_decode_errors_total: IntCounter, + payload_decode_errors_total: IntCounter, + identity_errors_total: IntCounter, + forced_aborts_total: IntCounter, + dropped_events_total: IntCounter, +} + +static ACTIVE_SEQUENCE_ZMQ_INGRESS_METRICS: OnceLock> = + OnceLock::new(); + +impl ActiveSequenceZmqIngressMetrics { + pub(crate) fn from_component(component: &Component) -> Arc { + ACTIVE_SEQUENCE_ZMQ_INGRESS_METRICS + .get_or_init(|| { + let metrics = component.metrics(); + let counter = |name, help| { + metrics + .create_intcounter(name, help, &[]) + .unwrap_or_else(|error| panic!("failed to create {name}: {error}")) + }; + Arc::new(Self { + tracked_sources: metrics + .create_intgauge( + "router_active_sequence_zmq_ingress_sources", + "Number of tracked direct-ZMQ active-sequence sources", + &[], + ) + .expect("failed to create router_active_sequence_zmq_ingress_sources"), + sequence_gaps_total: counter( + "router_active_sequence_zmq_ingress_sequence_gaps_total", + "Missing direct-ZMQ active-sequence envelopes inferred from sequence gaps", + ), + out_of_order_total: counter( + "router_active_sequence_zmq_ingress_out_of_order_total", + "Non-increasing direct-ZMQ active-sequence envelope sequences", + ), + reconnects_total: counter( + "router_active_sequence_zmq_ingress_reconnects_total", + "Direct-ZMQ active-sequence source reconnect attempts", + ), + replacements_total: counter( + "router_active_sequence_zmq_ingress_replacements_total", + "Direct-ZMQ active-sequence source task replacements", + ), + envelope_decode_errors_total: counter( + "router_active_sequence_zmq_ingress_envelope_decode_errors_total", + "Direct-ZMQ active-sequence envelope decode errors", + ), + payload_decode_errors_total: counter( + "router_active_sequence_zmq_ingress_payload_decode_errors_total", + "Direct-ZMQ active-sequence payload decode errors", + ), + identity_errors_total: counter( + "router_active_sequence_zmq_ingress_identity_errors_total", + "Direct-ZMQ active-sequence frame and envelope identity mismatches", + ), + forced_aborts_total: counter( + "router_active_sequence_zmq_ingress_forced_aborts_total", + "Direct-ZMQ active-sequence source tasks aborted after join timeout", + ), + dropped_events_total: counter( + "router_active_sequence_zmq_ingress_dropped_events_total", + "Direct-ZMQ active-sequence events dropped because the partition's inbound channel was full or closed", + ), + }) + }) + .clone() + } + + pub(crate) fn source_started(&self) { + self.tracked_sources.inc(); + } + + pub(crate) fn source_stopped(&self) { + self.tracked_sources.dec(); + } + + pub(crate) fn record_gap(&self, missing: u64) { + self.sequence_gaps_total.inc_by(missing); + } + + pub(crate) fn record_out_of_order(&self) { + self.out_of_order_total.inc(); + } + + pub(crate) fn record_reconnect(&self) { + self.reconnects_total.inc(); + } + + pub(crate) fn record_replacement(&self) { + self.replacements_total.inc(); + } + + pub(crate) fn record_envelope_decode_error(&self) { + self.envelope_decode_errors_total.inc(); + } + + pub(crate) fn record_payload_decode_error(&self) { + self.payload_decode_errors_total.inc(); + } + + pub(crate) fn record_identity_error(&self) { + self.identity_errors_total.inc(); + } + + pub(crate) fn record_forced_abort(&self) { + self.forced_aborts_total.inc(); + } + + pub(crate) fn record_dropped_event(&self) { + self.dropped_events_total.inc(); + } +} + +// --------------------------------------------------------------------------- +// Router worker status metrics (component-scoped gauges) +// --------------------------------------------------------------------------- + +/// Component-scoped router gauges for worker discovery. pub(crate) struct RouterWorkerStatusMetrics { pub registered: IntGaugeVec, pub kv_event_source_mismatch_workers: IntGaugeVec, diff --git a/lib/llm/src/kv_router/sequence.rs b/lib/llm/src/kv_router/sequence.rs index c0f6b9390a02..692e91187e5c 100644 --- a/lib/llm/src/kv_router/sequence.rs +++ b/lib/llm/src/kv_router/sequence.rs @@ -33,6 +33,8 @@ use tokio::sync::mpsc; use tokio::time::Instant; use tokio_util::sync::CancellationToken; +mod direct_zmq; + use crate::kv_router::ACTIVE_SEQUENCES_SUBJECT; #[cfg(test)] use dynamo_runtime::transports::event_plane::MsgpackCodec; @@ -462,60 +464,74 @@ impl SequenceSubscriber for RuntimeSequenceSubscriber { /// Replica-sync channels for an embedded selection partition, carried over the /// runtime event plane exactly like the runtime scheduler's replica sync: -/// outbound events are published on `ACTIVE_SEQUENCES_SUBJECT`, peer events -/// are subscribed from it and forwarded to the partition. +/// peer and worker events on `ACTIVE_SEQUENCES_SUBJECT` are always forwarded +/// to the partition (worker-origin completion marks are needed even without +/// router-to-router sync); outbound events are published only when +/// `publishes_outbound` is set. pub(crate) async fn host_replica_channels( endpoint: &Endpoint, router_id: u64, + publishes_outbound: bool, cancellation_token: CancellationToken, ) -> Result { let transport_kind = endpoint.drt().default_event_transport_kind(); - let (outbound, outbound_rx) = mpsc::channel(REPLICA_EVENT_CHANNEL_CAPACITY); - let event_publisher = EventPublisher::for_endpoint_with_transport( - endpoint, - ACTIVE_SEQUENCES_SUBJECT, - transport_kind, - ) - .await?; - match active_sequence_event_wire_format(transport_kind) { - ActiveSequenceEventWireFormat::Singleton => { - tokio::spawn(run_replica_singleton_publisher( - event_publisher, - outbound_rx, - cancellation_token.clone(), - )); - } - ActiveSequenceEventWireFormat::Batch => { - tokio::spawn(run_replica_batch_publisher( - event_publisher, - outbound_rx, - cancellation_token.clone(), - )); + let outbound = if publishes_outbound { + let (outbound, outbound_rx) = mpsc::channel(REPLICA_EVENT_CHANNEL_CAPACITY); + let event_publisher = EventPublisher::for_endpoint_with_transport( + endpoint, + ACTIVE_SEQUENCES_SUBJECT, + transport_kind, + ) + .await?; + match active_sequence_event_wire_format(transport_kind) { + ActiveSequenceEventWireFormat::Singleton => { + tokio::spawn(run_replica_singleton_publisher( + event_publisher, + outbound_rx, + cancellation_token.clone(), + )); + } + ActiveSequenceEventWireFormat::Batch => { + tokio::spawn(run_replica_batch_publisher( + event_publisher, + outbound_rx, + cancellation_token.clone(), + )); + } } - } + Some(outbound) + } else { + None + }; let (inbound_tx, inbound_rx) = mpsc::channel(REPLICA_EVENT_CHANNEL_CAPACITY); - let mut subscriber = RuntimeSequenceSubscriber::for_endpoint(endpoint).await?; - let forward_tx = inbound_tx.clone(); - tokio::spawn(async move { - loop { - let next = tokio::select! { - _ = cancellation_token.cancelled() => break, - next = subscriber.next_event() => next, - }; - match next { - Some(Ok(event)) => { - if forward_tx.send(event).await.is_err() { - break; - } - } - Some(Err(error)) => { - tracing::warn!(%error, "replica-sync subscriber error; continuing"); - } - None => break, - } - } - }); + let direct = direct_zmq::DirectZmqSequenceConfig::from_env(); + let ingress_result = if direct.should_use_direct(transport_kind) { + direct_zmq::start( + endpoint.clone(), + inbound_tx.clone(), + direct.rcvhwm, + cancellation_token, + ) + .await + .map(|_supervisor| ()) + } else { + RuntimeSequenceSubscriber::for_endpoint(endpoint) + .await + .map(|subscriber| { + tokio::spawn(forward_replica_events( + subscriber, + inbound_tx.clone(), + cancellation_token, + )); + }) + }; + if let Err(error) = ingress_result { + tracing::warn!( + %error, + "active-sequence event ingress unavailable; continuing with response-side cleanup" + ); + } Ok(dynamo_kv_router::services::selection::HostReplicaChannels { outbound, inbound_tx, @@ -524,6 +540,30 @@ pub(crate) async fn host_replica_channels( }) } +async fn forward_replica_events( + mut subscriber: RuntimeSequenceSubscriber, + forward_tx: mpsc::Sender, + cancellation_token: CancellationToken, +) { + loop { + let next = tokio::select! { + _ = cancellation_token.cancelled() => break, + next = subscriber.next_event() => next, + }; + match next { + Some(Ok(event)) => { + if forward_tx.send(event).await.is_err() { + break; + } + } + Some(Err(error)) => { + tracing::warn!(%error, "replica-sync subscriber error; continuing"); + } + None => break, + } + } +} + #[cfg(test)] mod tests { use super::*; @@ -866,4 +906,49 @@ mod tests { ); assert!(!deferred.install(observer)); } + + #[tokio::test] + #[serial_test::serial] + async fn worker_completion_ingress_runs_without_router_replica_sync() -> Result<()> { + let runtime = dynamo_runtime::Runtime::from_current()?; + let distributed = dynamo_runtime::DistributedRuntime::new( + runtime, + dynamo_runtime::distributed::DistributedConfig::process_local(), + ) + .await?; + let endpoint = distributed + .namespace(format!( + "worker-completion-ingress-{}", + uuid::Uuid::new_v4() + ))? + .component("workers")? + .endpoint("generate"); + let cancel = CancellationToken::new(); + let mut channels = + host_replica_channels(&endpoint, 99, false, cancel.child_token()).await?; + assert!(channels.outbound.is_none()); + + let publisher = ActiveSequenceEventPublisher::for_endpoint(&endpoint, 16).await?; + let received = tokio::time::timeout(std::time::Duration::from_secs(5), async { + loop { + publisher.mark_prefill_completed("worker-origin-mark".to_string(), 42, 0)?; + tokio::select! { + event = channels.inbound_rx.recv() => return Ok::<_, anyhow::Error>(event), + _ = tokio::time::sleep(std::time::Duration::from_millis(50)) => {} + } + } + }) + .await?? + .expect("inbound channel stays open"); + assert_eq!(received.request_id, "worker-origin-mark"); + assert!(matches!( + received.data, + ActiveSequenceEventData::MarkPrefillCompleted + )); + + drop(publisher); + cancel.cancel(); + distributed.shutdown(); + Ok(()) + } } diff --git a/lib/llm/src/kv_router/sequence/direct_zmq.rs b/lib/llm/src/kv_router/sequence/direct_zmq.rs new file mode 100644 index 000000000000..77d22e088e1c --- /dev/null +++ b/lib/llm/src/kv_router/sequence/direct_zmq.rs @@ -0,0 +1,283 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::ffi::OsString; + +use anyhow::Result; +use dynamo_kv_router::protocols::{ + ActiveSequenceEvent, ActiveSequenceEventBatch, MAX_REPLICA_BATCH_EVENTS, +}; +use dynamo_runtime::{ + component::Endpoint, + config::parse_bool_opt, + discovery::EventTransportKind, + transports::event_plane::{Codec, uses_direct_zmq}, +}; +use tokio::sync::mpsc; +use tokio::task::JoinHandle; +use tokio_util::sync::CancellationToken; + +use crate::{ + direct_zmq_fan_in::{ContinuityMode, FanInEvent, start_direct_zmq_fan_in}, + kv_router::{ACTIVE_SEQUENCES_SUBJECT, metrics::ActiveSequenceZmqIngressMetrics}, +}; + +const DIRECT_ZMQ_ENV: &str = "DYN_ROUTER_ACTIVE_SEQUENCE_DIRECT_ZMQ"; +const RCVHWM_ENV: &str = "DYN_ROUTER_ACTIVE_SEQUENCE_ZMQ_RCVHWM"; +const DEFAULT_RCVHWM: i32 = 1024; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) struct DirectZmqSequenceConfig { + enabled: bool, + pub(super) rcvhwm: i32, +} + +impl DirectZmqSequenceConfig { + pub(super) fn from_env() -> Self { + Self::from_lookup(|key| std::env::var_os(key)) + } + + fn from_lookup(mut get_env: impl FnMut(&str) -> Option) -> Self { + let enabled = match get_env(DIRECT_ZMQ_ENV) { + None => true, + Some(raw) => { + let value = raw.to_string_lossy(); + match parse_bool_opt(&value) { + Some(enabled) => enabled, + None => { + tracing::warn!( + env = DIRECT_ZMQ_ENV, + %value, + "invalid direct-ZMQ active-sequence setting; using enabled default" + ); + true + } + } + } + }; + + let rcvhwm = match get_env(RCVHWM_ENV) { + None => DEFAULT_RCVHWM, + Some(raw) => match raw.to_string_lossy().trim().parse::() { + Ok(value) if value > 0 => value, + _ => { + tracing::warn!( + env = RCVHWM_ENV, + value = %raw.to_string_lossy(), + default = DEFAULT_RCVHWM, + "invalid direct-ZMQ active-sequence receive HWM; using default" + ); + DEFAULT_RCVHWM + } + }, + }; + + Self { enabled, rcvhwm } + } + + pub(super) fn should_use_direct(self, transport_kind: EventTransportKind) -> bool { + self.should_use_direct_for_topology(transport_kind, uses_direct_zmq(transport_kind)) + } + + fn should_use_direct_for_topology( + self, + transport_kind: EventTransportKind, + direct_zmq_topology: bool, + ) -> bool { + self.enabled && transport_kind == EventTransportKind::Zmq && direct_zmq_topology + } +} + +/// Direct-ZMQ ingress for active-sequence batches: one source per publisher, +/// continuity tracked from zero, events forwarded into the partition's inbound +/// replica channel. A full channel drops the event and counts it. +pub(super) async fn start( + endpoint: Endpoint, + inbound_tx: mpsc::Sender, + rcvhwm: i32, + cancellation_token: CancellationToken, +) -> Result> { + let metrics = ActiveSequenceZmqIngressMetrics::from_component(endpoint.component()); + let handler_metrics = metrics.clone(); + let handler = move |envelope: dynamo_runtime::transports::event_plane::ValidatedEnvelope| { + let codec = Codec::default(); + let batch = codec + .decode_payload::(&envelope.payload) + .inspect_err(|_| { + handler_metrics.record_payload_decode_error(); + })?; + if batch.events.len() > MAX_REPLICA_BATCH_EVENTS { + handler_metrics.record_payload_decode_error(); + anyhow::bail!( + "active-sequence batch contains {} events; maximum is {}", + batch.events.len(), + MAX_REPLICA_BATCH_EVENTS + ); + } + for event in batch.events { + if inbound_tx.try_send(event).is_err() { + handler_metrics.record_dropped_event(); + } + } + Ok(()) + }; + let observer = + move |observation: crate::direct_zmq_fan_in::FanInObservation| match observation.event { + FanInEvent::SourceStarted => metrics.source_started(), + FanInEvent::SourceStopped => metrics.source_stopped(), + FanInEvent::Disconnected | FanInEvent::DiscoveryReset => {} + FanInEvent::Reconnect => metrics.record_reconnect(), + FanInEvent::Replacement => metrics.record_replacement(), + FanInEvent::EnvelopeDecodeError => metrics.record_envelope_decode_error(), + FanInEvent::IdentityMismatch => metrics.record_identity_error(), + FanInEvent::SequenceGap { missing } => metrics.record_gap(missing), + FanInEvent::OutOfOrder => metrics.record_out_of_order(), + FanInEvent::ForcedAbort => metrics.record_forced_abort(), + }; + + start_direct_zmq_fan_in( + endpoint, + ACTIVE_SEQUENCES_SUBJECT, + rcvhwm, + None, + ContinuityMode::TrackFromZero, + cancellation_token, + handler, + observer, + ) + .await +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use dynamo_kv_router::protocols::{ActiveSequenceEventData, WorkerWithDpRank}; + use dynamo_runtime::{ + DistributedRuntime, Runtime, distributed::DistributedConfig, + transports::event_plane::EventPublisher, + }; + + use super::*; + + fn lookup(direct: Option<&str>, rcvhwm: Option<&str>) -> impl FnMut(&str) -> Option { + let direct = direct.map(OsString::from); + let rcvhwm = rcvhwm.map(OsString::from); + move |key| match key { + DIRECT_ZMQ_ENV => direct.clone(), + RCVHWM_ENV => rcvhwm.clone(), + _ => None, + } + } + + #[test] + fn parses_direct_zmq_configuration() { + assert_eq!( + DirectZmqSequenceConfig::from_lookup(lookup(None, None)), + DirectZmqSequenceConfig { + enabled: true, + rcvhwm: DEFAULT_RCVHWM, + } + ); + assert_eq!( + DirectZmqSequenceConfig::from_lookup(lookup(Some("false"), Some("37"))), + DirectZmqSequenceConfig { + enabled: false, + rcvhwm: 37, + } + ); + for invalid in ["0", "-1", "not-a-number", "2147483648"] { + assert_eq!( + DirectZmqSequenceConfig::from_lookup(lookup(None, Some(invalid))).rcvhwm, + DEFAULT_RCVHWM + ); + } + assert!(DirectZmqSequenceConfig::from_lookup(lookup(Some("invalid"), None)).enabled); + } + + #[test] + fn selects_only_direct_zmq_and_honors_rollback() { + let enabled = DirectZmqSequenceConfig::from_lookup(lookup(None, None)); + let rollback = DirectZmqSequenceConfig::from_lookup(lookup(Some("0"), None)); + + assert!(enabled.should_use_direct_for_topology(EventTransportKind::Zmq, true)); + assert!(!enabled.should_use_direct_for_topology(EventTransportKind::Zmq, false)); + assert!(!enabled.should_use_direct_for_topology(EventTransportKind::Nats, false)); + assert!(!rollback.should_use_direct_for_topology(EventTransportKind::Zmq, true)); + assert!(!rollback.should_use_direct_for_topology(EventTransportKind::Zmq, false)); + assert!(!rollback.should_use_direct_for_topology(EventTransportKind::Nats, false)); + } + + fn add_event(request_id: impl Into) -> ActiveSequenceEvent { + ActiveSequenceEvent { + request_id: request_id.into(), + worker: WorkerWithDpRank::new(0, 0), + data: ActiveSequenceEventData::AddRequest { + token_sequence: Some(vec![1]), + track_prefill_tokens: false, + expected_output_tokens: None, + prefill_load_hint: None, + }, + router_id: 99, + lora_name: None, + } + } + + #[tokio::test] + #[serial_test::serial] + async fn direct_zmq_batch_ingress_applies_after_publisher_recreation() -> Result<()> { + let runtime = Runtime::from_current()?; + let distributed = + DistributedRuntime::new(runtime, DistributedConfig::process_local()).await?; + let endpoint = distributed + .namespace(format!( + "active-sequence-direct-zmq-{}", + uuid::Uuid::new_v4() + ))? + .component("frontend")? + .endpoint("generate"); + let (inbound_tx, mut inbound_rx) = mpsc::channel(16); + let cancel = CancellationToken::new(); + let supervisor = + start(endpoint.clone(), inbound_tx, DEFAULT_RCVHWM, cancel.clone()).await?; + + for request_id in ["before-recreation", "after-recreation"] { + let publisher = EventPublisher::for_endpoint_with_transport( + &endpoint, + ACTIVE_SEQUENCES_SUBJECT, + EventTransportKind::Zmq, + ) + .await?; + tokio::time::timeout(Duration::from_secs(5), async { + loop { + publisher + .publish(&ActiveSequenceEventBatch { + events: vec![add_event(request_id)], + }) + .await + .unwrap(); + tokio::select! { + received = inbound_rx.recv() => { + // Republished duplicates of the previous id may still be queued. + if received.map(|event| event.request_id).as_deref() == Some(request_id) { + break; + } + } + _ = tokio::time::sleep(Duration::from_millis(20)) => {} + } + } + }) + .await + .expect("direct-ZMQ source should apply a warmed publisher batch"); + drop(publisher); + } + + cancel.cancel(); + tokio::time::timeout(Duration::from_secs(5), supervisor) + .await + .expect("fan-in supervisor should stop after cancellation") + .expect("fan-in supervisor should not panic"); + distributed.shutdown(); + Ok(()) + } +} From ab4c95962866fbdac26986fae15616c6fe960c40 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Wed, 9 Sep 2026 23:27:37 -0700 Subject: [PATCH 037/102] fix(kv-router): compile the slot tracker without the selection feature replica_sync.rs imported selection::affinity unconditionally while the selection module is gated on standalone-selection, so the isolated standalone-slot-tracker build failed. Gate the affinity mesh plumbing on the selection feature and check both isolated feature builds in pre-merge, where workspace feature unification otherwise hides them. Signed-off-by: Ishan Dhanani --- .github/workflows/pre-merge.yml | 7 +++++++ lib/kv-router/src/services/common/replica_sync.rs | 6 ++++++ 2 files changed, 13 insertions(+) diff --git a/.github/workflows/pre-merge.yml b/.github/workflows/pre-merge.yml index 32f9c5820608..120975ff2d65 100644 --- a/.github/workflows/pre-merge.yml +++ b/.github/workflows/pre-merge.yml @@ -435,6 +435,13 @@ jobs: working-directory: ${{ matrix.dir }} # NOTE: --all-targets doesn't run doc tests run: cargo test --locked --all-targets + - name: Check isolated kv-router feature builds + if: matrix.dir == '.' + working-directory: ${{ matrix.dir }} + # Workspace feature unification hides these; check each on its own. + run: | + cargo check --locked -p dynamo-kv-router --no-default-features --features standalone-slot-tracker + cargo check --locked -p dynamo-kv-router --no-default-features --features standalone-indexer - name: Run canonical replay feature tests if: matrix.dir == '.' working-directory: ${{ matrix.dir }} diff --git a/lib/kv-router/src/services/common/replica_sync.rs b/lib/kv-router/src/services/common/replica_sync.rs index e9412ff4eb36..d01ed8cda842 100644 --- a/lib/kv-router/src/services/common/replica_sync.rs +++ b/lib/kv-router/src/services/common/replica_sync.rs @@ -18,6 +18,7 @@ use crate::sequences::{SequencePublishQueueError, SequencePublisher, SequenceSub use crate::services::common::zmq::{ create_bound_pub_socket, create_sub_socket_topics, validate_endpoint, }; +#[cfg(feature = "standalone-selection")] use crate::services::selection::affinity::{AffinityReplicaSink, AffinityTarget, AffinityVersion}; pub(crate) const REPLICA_EVENT_CHANNEL_CAPACITY: usize = 100_000; @@ -40,6 +41,7 @@ pub(crate) struct AffinityBindingEvent { pub writer_id: u64, } +#[cfg(feature = "standalone-selection")] impl AffinityBindingEvent { pub(crate) fn target(&self) -> AffinityTarget { AffinityTarget::new(self.worker_id, self.dp_rank) @@ -53,6 +55,7 @@ impl AffinityBindingEvent { } } +#[cfg(feature = "standalone-selection")] /// Publishes bindings from a [`super::super::selection::affinity::SessionAffinity`] /// into the mesh. Best effort: a full channel drops the update. struct AffinityMeshSink { @@ -60,6 +63,7 @@ struct AffinityMeshSink { tx: mpsc::Sender, } +#[cfg(feature = "standalone-selection")] impl AffinityReplicaSink for AffinityMeshSink { fn publish(&self, session_id: &str, target: AffinityTarget, version: AffinityVersion) { let update = AffinityBindingEvent { @@ -166,6 +170,7 @@ impl ReplicaSyncConfig { self.process_id } + #[cfg(feature = "standalone-selection")] /// Sink that publishes session bindings into the mesh, when this runtime /// carries them. pub(crate) fn affinity_sink( @@ -222,6 +227,7 @@ pub struct HostReplicaChannels { pub process_id: u64, } +#[cfg(feature = "standalone-selection")] /// Per-partition factory for [`HostReplicaChannels`]; `None` disables replica /// sync for that partition. pub type HostReplicaSyncFactory = From 0c0096d780d32cfcc589a80ea6914495346bbd64 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Wed, 9 Sep 2026 23:31:06 -0700 Subject: [PATCH 038/102] fix(router): default unreported worker capacity for discovered workers The shared catalog marks a worker Incomplete when queueing is enabled and it reports no max_num_batched_tokens, so runtime-discovered workers that omit the field were held out of routing. The deleted scheduler never gated on it and the queue already falls back to DEFAULT_MAX_BATCHED_TOKENS; apply that default when building the catalog record. The standalone and EPP gates are unchanged. Signed-off-by: Ishan Dhanani --- lib/llm/src/kv_router/embedded.rs | 23 +++++++++++++++++++++-- 1 file changed, 21 insertions(+), 2 deletions(-) diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs index c163c01cefcb..abfdfffaa8c2 100644 --- a/lib/llm/src/kv_router/embedded.rs +++ b/lib/llm/src/kv_router/embedded.rs @@ -16,7 +16,9 @@ use dynamo_kv_router::WorkerType; use dynamo_kv_router::config::KvRouterConfig; use dynamo_kv_router::identity::RoutingPartitionId; use dynamo_kv_router::protocols::{WorkerConfigLike, WorkerId, WorkerWithDpRank}; -use dynamo_kv_router::scheduling::queue::{SchedulerBookingCleanup, SchedulerBookingDescriptor}; +use dynamo_kv_router::scheduling::queue::{ + DEFAULT_MAX_BATCHED_TOKENS, SchedulerBookingCleanup, SchedulerBookingDescriptor, +}; use dynamo_kv_router::scheduling::{ AdmittedSchedulingResponse, AdvisorySchedulingResponse, AttemptId, KvSchedulerError, OverloadedWorkerProvider, PotentialLoad, ScheduleRequest, WorkerAvailabilityProvider, @@ -430,7 +432,13 @@ pub(crate) fn worker_request_from_runtime_config( block_size: Some(block_size), data_parallel_start_rank: Some(dp_start), data_parallel_size: Some(dp_size), - max_num_batched_tokens: config.max_num_batched_tokens, + // Default unreported capacity so the catalog's queueing gate does not + // mark the worker Incomplete; the scheduler applies the same fallback. + max_num_batched_tokens: Some( + config + .max_num_batched_tokens + .unwrap_or(DEFAULT_MAX_BATCHED_TOKENS), + ), total_kv_blocks: config.total_kv_blocks, stable_routing_id: config.stable_routing_id.clone(), is_eagle: Some(is_eagle), @@ -494,4 +502,15 @@ mod tests { ]) ); } + + #[test] + fn worker_request_defaults_unreported_capacity() { + let key = RoutingPartitionId::new("model", DEFAULT_ROUTING_GROUP); + let request = + worker_request_from_runtime_config(7, &ModelRuntimeConfig::default(), &key, 16, false); + assert_eq!( + request.max_num_batched_tokens, + Some(DEFAULT_MAX_BATCHED_TOKENS) + ); + } } From 76c437417737151a546174913364310a4ebd9dbf Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Wed, 9 Sep 2026 23:39:14 -0700 Subject: [PATCH 039/102] fix(router): restore per-worker load gauges on the embedded partition The partition publisher's observe_load hook was empty, so the frontend's dynamo_frontend_worker_active_decode_blocks and _prefill_tokens gauges had no writer. Route the hook through the existing SchedulerLoadSink and have the router's sink set the gauges with its own worker-type label, matching the worker monitor's cleanup path. Signed-off-by: Ishan Dhanani --- .../src/services/common/replica_sync.rs | 14 ++++- lib/llm/src/kv_router/embedded.rs | 61 +++++++++++++++++-- 2 files changed, 67 insertions(+), 8 deletions(-) diff --git a/lib/kv-router/src/services/common/replica_sync.rs b/lib/kv-router/src/services/common/replica_sync.rs index d01ed8cda842..5c7d0961f623 100644 --- a/lib/kv-router/src/services/common/replica_sync.rs +++ b/lib/kv-router/src/services/common/replica_sync.rs @@ -204,6 +204,11 @@ pub(crate) struct ScopedReplicaSync { pub trait SchedulerLoadSink: Send + Sync { fn publish(&self, snapshot: crate::sequences::SchedulerLoadSnapshot); + /// Per-worker load after any local mutation, including output blocks, + /// which are never published as shared scheduler load. The sink owns the + /// metric label so it matches the host's cleanup path. + fn observe_local_load(&self, _worker: &WorkerWithDpRank, _blocks: usize, _tokens: usize) {} + fn publish_batch(&self, snapshots: Vec) { for snapshot in snapshots { self.publish(snapshot); @@ -352,11 +357,14 @@ impl SequencePublisher for ScopedSequencePublisher { fn observe_load( &self, - _worker: &WorkerWithDpRank, + worker: &WorkerWithDpRank, _worker_type: &str, - _blocks: usize, - _tokens: usize, + blocks: usize, + tokens: usize, ) { + if let Some(sink) = &self.load_sink { + sink.observe_local_load(worker, blocks, tokens); + } } } diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs index abfdfffaa8c2..426a50d89225 100644 --- a/lib/llm/src/kv_router/embedded.rs +++ b/lib/llm/src/kv_router/embedded.rs @@ -35,6 +35,7 @@ use dynamo_tokens::SequenceHash; use tokio_util::sync::CancellationToken; use crate::discovery::RuntimeConfigWatch; +use crate::kv_router::metrics::WORKER_LOAD_METRICS; use crate::local_model::runtime_config::ModelRuntimeConfig; /// Inputs the embedded backend needs from the router at construction. @@ -94,16 +95,30 @@ pub fn worker_selection_policy_registry() -> WorkerSelectionPolicyRegistry { } /// Bridges the partition's scheduler load snapshots to the router's -/// `SchedulerLoadSender`, which feeds `KvWorkerMonitor`. -struct SenderLoadSink(crate::kv_router::routing_load::SchedulerLoadSender); +/// `SchedulerLoadSender`, which feeds `KvWorkerMonitor`, and its per-worker +/// load to the frontend gauges. +struct SenderLoadSink { + sender: crate::kv_router::routing_load::SchedulerLoadSender, + worker_type: &'static str, +} impl dynamo_kv_router::services::selection::SchedulerLoadSink for SenderLoadSink { fn publish(&self, snapshot: dynamo_kv_router::sequences::SchedulerLoadSnapshot) { - self.0.publish(snapshot); + self.sender.publish(snapshot); } fn publish_batch(&self, snapshots: Vec) { - self.0.publish_batch(snapshots); + self.sender.publish_batch(snapshots); + } + + fn observe_local_load(&self, worker: &WorkerWithDpRank, blocks: usize, tokens: usize) { + WORKER_LOAD_METRICS.observe( + worker.worker_id, + worker.dp_rank, + self.worker_type, + blocks, + tokens, + ); } } @@ -157,7 +172,10 @@ impl EmbeddedSelection { }, eligibility: HostEligibility::default(), telemetry: HostTelemetry { - scheduler_load: Some(Arc::new(SenderLoadSink(args.scheduler_load))), + scheduler_load: Some(Arc::new(SenderLoadSink { + sender: args.scheduler_load, + worker_type: args.metric_worker_type, + })), }, replication: HostReplication { channels: replica_sync, @@ -456,6 +474,7 @@ pub(crate) fn worker_request_from_runtime_config( #[cfg(test)] mod tests { use super::*; + use dynamo_kv_router::services::selection::SchedulerLoadSink; #[test] fn worker_request_mirrors_runtime_config() { @@ -513,4 +532,36 @@ mod tests { Some(DEFAULT_MAX_BATCHED_TOKENS) ); } + + #[test] + fn local_load_observation_sets_worker_gauges() { + let sink = SenderLoadSink { + sender: crate::kv_router::routing_load::SchedulerLoadSender::disabled( + CancellationToken::new(), + ), + worker_type: "decode", + }; + sink.observe_local_load(&WorkerWithDpRank::new(3, 1), 5, 7); + let labels = ["3", "1", "decode"]; + assert_eq!( + WORKER_LOAD_METRICS + .active_decode_blocks + .with_label_values(&labels) + .get(), + 5 + ); + assert_eq!( + WORKER_LOAD_METRICS + .active_prefill_tokens + .with_label_values(&labels) + .get(), + 7 + ); + let _ = WORKER_LOAD_METRICS + .active_decode_blocks + .remove_label_values(&labels); + let _ = WORKER_LOAD_METRICS + .active_prefill_tokens + .remove_label_values(&labels); + } } From e5fe3f00a16c133d50a83eec81a7aef69ee80faf Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Thu, 10 Sep 2026 00:02:13 -0700 Subject: [PATCH 040/102] fix(router): restore router queue metrics and the prefill overlap observer Deleting KvScheduler removed the only producer of the per-class queue gauges, the queue rejection counters, and the prefill-only non-max-overlap observer. Build the handles from the same policy profile the partition scheduler resolved, refresh them on schedule results and on the scheduler's update watch, and install the observer on prefill partitions. Signed-off-by: Ishan Dhanani --- lib/llm/src/kv_router/embedded.rs | 184 ++++++++++++++++++++++++++++-- 1 file changed, 177 insertions(+), 7 deletions(-) diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs index 426a50d89225..8ee4435a2379 100644 --- a/lib/llm/src/kv_router/embedded.rs +++ b/lib/llm/src/kv_router/embedded.rs @@ -17,11 +17,13 @@ use dynamo_kv_router::config::KvRouterConfig; use dynamo_kv_router::identity::RoutingPartitionId; use dynamo_kv_router::protocols::{WorkerConfigLike, WorkerId, WorkerWithDpRank}; use dynamo_kv_router::scheduling::queue::{ - DEFAULT_MAX_BATCHED_TOKENS, SchedulerBookingCleanup, SchedulerBookingDescriptor, + ClassQueueStats, DEFAULT_MAX_BATCHED_TOKENS, SchedulerBookingCleanup, + SchedulerBookingDescriptor, }; use dynamo_kv_router::scheduling::{ AdmittedSchedulingResponse, AdvisorySchedulingResponse, AttemptId, KvSchedulerError, - OverloadedWorkerProvider, PotentialLoad, ScheduleRequest, WorkerAvailabilityProvider, + NonMaxOverlapSelectionObserver, OverloadedWorkerProvider, PotentialLoad, QueueLimitKind, + QueueRejection, ScheduleRequest, WorkerAvailabilityProvider, }; use dynamo_kv_router::sequences::{SequenceError, SequenceRequest}; use dynamo_kv_router::services::selection::{ @@ -35,7 +37,9 @@ use dynamo_tokens::SequenceHash; use tokio_util::sync::CancellationToken; use crate::discovery::RuntimeConfigWatch; -use crate::kv_router::metrics::WORKER_LOAD_METRICS; +use crate::kv_router::metrics::{ + ROUTER_QUEUE_METRICS, RouterQueueMetricHandles, RouterRequestMetrics, WORKER_LOAD_METRICS, +}; use crate::local_model::runtime_config::ModelRuntimeConfig; /// Inputs the embedded backend needs from the router at construction. @@ -66,6 +70,94 @@ pub(crate) struct EmbeddedSelectionArgs { pub request_leases: Arc, } +fn record_queue_rejection( + per_class: &[RouterQueueMetricHandles], + indices: &HashMap, + rejection: &QueueRejection, +) { + let Some(handles) = indices + .get(&rejection.policy_class) + .and_then(|index| per_class.get(*index)) + else { + return; + }; + match rejection.limit_kind { + QueueLimitKind::Requests => handles.request_limit_rejections.inc(), + QueueLimitKind::RawIslTokens => handles.raw_isl_limit_rejections.inc(), + QueueLimitKind::CachedTokens => handles.cached_token_limit_rejections.inc(), + } +} + +fn update_queue_metrics( + per_class: &[RouterQueueMetricHandles], + mut stats_for: impl FnMut(usize) -> Option, +) { + for (class_index, handles) in per_class.iter().enumerate() { + let Some(stats) = stats_for(class_index) else { + debug_assert!( + false, + "missing queue counters for policy class {class_index}" + ); + continue; + }; + handles.pending_requests.set(stats.pending_count as i64); + handles + .pending_isl_tokens + .set(stats.pending_isl_tokens as i64); + handles + .pending_cached_tokens + .set(stats.pending_cached_tokens as i64); + } +} + +/// Refresh queue gauges on worker-config changes and at least once a minute. +fn spawn_queue_metrics_updater( + partition: SelectionPartition, + handles: Vec, + cancellation_token: CancellationToken, +) { + let mut queue_updates = partition.scheduler().subscribe_queue_updates(); + tokio::spawn(async move { + let period = std::time::Duration::from_secs(60); + let mut recheck = tokio::time::interval_at(tokio::time::Instant::now() + period, period); + loop { + update_queue_metrics(&handles, |index| { + partition.scheduler().class_queue_stats(index) + }); + tokio::select! { + _ = cancellation_token.cancelled() => break, + changed = queue_updates.changed() => { + if changed.is_err() { + break; + } + } + _ = recheck.tick() => {} + } + } + }); +} + +fn non_max_overlap_observer(worker_type: &'static str) -> NonMaxOverlapSelectionObserver { + Arc::new(move |request_id, selection| { + let overlap_blocks_lost = selection.overlap_blocks_lost(); + if let Some(metrics) = RouterRequestMetrics::get() { + metrics.observe_non_max_overlap_selection(worker_type, overlap_blocks_lost); + } + tracing::debug!( + request_id, + worker_type, + selected_worker_id = selection.selected_worker.worker_id, + selected_dp_rank = selection.selected_worker.dp_rank, + selected_overlap_blocks = selection.selected_overlap_blocks, + highest_overlap_worker_id = selection.highest_overlap_worker.worker_id, + highest_overlap_dp_rank = selection.highest_overlap_worker.dp_rank, + highest_overlap_blocks = selection.highest_overlap_blocks, + overlap_blocks_lost, + "Router selected a worker with lower KV cache overlap" + ); + }) +} + /// Partition model name when the router has none. pub(crate) const DEFAULT_MODEL_NAME: &str = "default"; @@ -77,6 +169,10 @@ pub(crate) struct EmbeddedSelection { partition: SelectionPartition, affinity: std::sync::OnceLock, worker_type: &'static str, + /// Queue gauges and rejection counters per policy class, index-aligned with + /// the scheduler's `class_queue_stats`. + queue_metrics: Vec, + queue_metric_indices: HashMap, } static INSTALLED_POLICY_REGISTRY: std::sync::OnceLock = @@ -190,6 +286,39 @@ impl EmbeddedSelection { .ensure_partition(key.clone(), args.block_size, args.is_eagle) .context("failed to create embedded selection partition")?; + // Same profile the partition scheduler resolved, so class indices align. + let profile = args + .kv_router_config + .policy_profile(Some(&key.model_name)) + .context("failed to resolve the embedded selection policy profile")?; + let queue_metrics: Vec<_> = profile + .classes() + .iter() + .map(|class| { + ROUTER_QUEUE_METRICS.handles(&key.model_name, args.metric_worker_type, &class.name) + }) + .collect(); + let queue_metric_indices = profile + .classes() + .iter() + .enumerate() + .map(|(index, class)| (class.name.clone(), index)) + .collect(); + spawn_queue_metrics_updater( + partition.clone(), + queue_metrics.clone(), + cancellation_token.child_token(), + ); + if worker_type == WorkerType::Prefill + && !partition + .scheduler() + .set_non_max_overlap_selection_observer(non_max_overlap_observer( + args.metric_worker_type, + )) + { + anyhow::bail!("non-max-overlap observer is already installed"); + } + let mut source = RuntimeDiscoverySource { watch: workers_with_configs, key, @@ -222,9 +351,22 @@ impl EmbeddedSelection { partition, affinity: std::sync::OnceLock::new(), worker_type: args.metric_worker_type, + queue_metrics, + queue_metric_indices, }) } + /// Queue changes are not observable through the scheduler's update watch, + /// so every schedule result refreshes the gauges. + fn observe_schedule_result(&self, result: &Result) { + if let Err(KvSchedulerError::QueueRejected(rejection)) = result { + record_queue_rejection(&self.queue_metrics, &self.queue_metric_indices, rejection); + } + update_queue_metrics(&self.queue_metrics, |index| { + self.partition.scheduler().class_queue_stats(index) + }); + } + pub(crate) fn affinity_coordinator( &self, ttl: std::time::Duration, @@ -244,20 +386,26 @@ impl EmbeddedSelection { &self, request: ScheduleRequest, ) -> Result { - self.partition + let result = self + .partition .scheduler() .schedule_request_admitted(request) - .await + .await; + self.observe_schedule_result(&result); + result } pub(crate) async fn select_without_admission( &self, request: ScheduleRequest, ) -> Result { - self.partition + let result = self + .partition .scheduler() .select_without_admission(request) - .await + .await; + self.observe_schedule_result(&result); + result } pub(crate) async fn add_request_admitted( @@ -564,4 +712,26 @@ mod tests { .active_prefill_tokens .remove_label_values(&labels); } + + #[test] + fn queue_rejections_count_against_their_policy_class() { + let per_class = vec![ + ROUTER_QUEUE_METRICS.handles("m-reject", "decode", "interactive"), + ROUTER_QUEUE_METRICS.handles("m-reject", "decode", "batch"), + ]; + let indices = HashMap::from([("interactive".to_string(), 0), ("batch".to_string(), 1)]); + record_queue_rejection( + &per_class, + &indices, + &QueueRejection { + policy_class: "batch".to_string(), + limit_kind: QueueLimitKind::RawIslTokens, + current: 9, + limit: 8, + }, + ); + assert_eq!(per_class[1].raw_isl_limit_rejections.get(), 1); + assert_eq!(per_class[0].raw_isl_limit_rejections.get(), 0); + assert_eq!(per_class[1].request_limit_rejections.get(), 0); + } } From ca0a0c6a1d9b4f074044f037f8d8295261cc46b2 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Thu, 10 Sep 2026 00:04:40 -0700 Subject: [PATCH 041/102] fix(router): route without affinity when the session table is full A full session-affinity table is a router-side limit, not a client fault. Booked selections now proceed unpinned and log at debug instead of failing with NotReady; invalid session ids still map to BadRequest. Signed-off-by: Ishan Dhanani --- .../src/services/selection/core/mod.rs | 60 ++++++++++++++++--- 1 file changed, 52 insertions(+), 8 deletions(-) diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index c6f62b6d21d5..cd76ad6638c8 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -44,7 +44,7 @@ use crate::services::indexer::registry::WorkerRegistry; use crate::services::overlap::MooncakeOverlapSummary; use crate::tracking_hash::{TrackingHashContext, TrackingHashScope}; -use super::affinity::{Acquired, AffinityLease, SessionAffinity}; +use super::affinity::{Acquired, AffinityError, AffinityLease, SessionAffinity}; use super::catalog::WorkerCatalog; use super::error::SelectionError; use super::ingress::{KvEventIngress, ZmqDirectIngress}; @@ -346,8 +346,7 @@ pub struct SelectionCore { type SessionTarget = super::affinity::AffinityTarget; -fn affinity_error(error: super::affinity::AffinityError) -> SelectionError { - use super::affinity::AffinityError; +fn affinity_error(error: AffinityError) -> SelectionError { match error { AffinityError::InvalidArgument(message) => SelectionError::BadRequest(message), AffinityError::ResourceExhausted(message) => SelectionError::NotReady(message), @@ -998,18 +997,29 @@ impl SelectionCore { let mut affinity_hold = None; let affinity_target = match (session_id.as_deref(), table) { (Some(session_id), Some(table)) if book => { - match tokio::select! { + let acquire_result = tokio::select! { _ = self.cancel_token.cancelled() => return Err(SelectionError::Scheduler(KvSchedulerError::SubscriberShutdown)), - result = table.acquire(session_id, None) => result.map_err(affinity_error)?, - } { - Acquired::Initialize(init) => { + result = table.acquire(session_id, None) => result, + }; + match acquire_result { + Ok(Acquired::Initialize(init)) => { affinity_hold = Some(Acquired::Initialize(init)); None } - Acquired::Bound { target, lease } => { + Ok(Acquired::Bound { target, lease }) => { affinity_hold = Some(Acquired::Bound { target, lease }); Some(WorkerAffinityTarget::new(target.worker_id, target.dp_rank)) } + // A full table is a router-side limit, not a client fault: + // route this request unpinned instead of failing it. + Err(AffinityError::ResourceExhausted(_)) => { + tracing::debug!( + session_id, + "affinity table full; routing without session affinity" + ); + None + } + Err(error) => return Err(affinity_error(error)), } } (Some(session_id), Some(table)) => table @@ -2696,6 +2706,40 @@ mod tests { assert_eq!(legacy.parent_session_id(), None); } + #[tokio::test] + async fn full_affinity_table_routes_without_pinning() { + let core = SelectionCore::try_new_local( + test_config(false), + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ) + .expect("valid test config"); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + let entry = core + .entry(&RoutingPartitionId::new("model", "default")) + .expect("entry"); + let table = SessionAffinity::new_with_limits(Duration::from_secs(60), 1, 256) + .expect("affinity table"); + assert!(entry.affinity.set(table).is_ok()); + + for (selection_id, session_id) in [("first", "s1"), ("second", "s2")] { + let mut request = reserve_request(selection_id); + request.session_id = Some(session_id.to_string()); + core.select_and_reserve(request) + .await + .expect("a full affinity table must not fail selection"); + } + assert_eq!( + core.reservation_index + .read() + .values() + .filter(|r| r._affinity_lease.is_some()) + .count(), + 1 + ); + } + #[tokio::test] async fn injected_availability_provider_restricts_selection() { let available: Arc>>>> = From 9493e06e2b091a75e4c7d6577e98b9fa1dcf61c3 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Thu, 10 Sep 2026 00:04:40 -0700 Subject: [PATCH 042/102] fix(router): report remote indexer outages as not-ready A remote indexer failure is a dependency outage, not a router bug: map IndexerOffline to NotReady (503) instead of Internal (500) so clients retry. Bound the connect phase at two seconds so a blackholed indexer fails fast, and carry the partition on the error-status log. Signed-off-by: Ishan Dhanani --- lib/kv-router/src/services/indexer/backend.rs | 19 ++++++++- .../src/services/selection/core/mod.rs | 42 ++++++++++++++++++- 2 files changed, 58 insertions(+), 3 deletions(-) diff --git a/lib/kv-router/src/services/indexer/backend.rs b/lib/kv-router/src/services/indexer/backend.rs index 92f12c4019eb..352bcee20a13 100644 --- a/lib/kv-router/src/services/indexer/backend.rs +++ b/lib/kv-router/src/services/indexer/backend.rs @@ -59,6 +59,8 @@ impl Eq for RemoteIndexerTransport {} impl RemoteIndexerTransport { pub const DEFAULT_TIMEOUT: Duration = Duration::from_secs(5); + /// A blackholed indexer fails fast instead of consuming the whole request budget. + const CONNECT_TIMEOUT: Duration = Duration::from_secs(2); pub fn new(base_url: impl Into) -> Result { Self::new_with_timeout(base_url, Self::DEFAULT_TIMEOUT) @@ -70,6 +72,7 @@ impl RemoteIndexerTransport { anyhow::bail!("remote indexer URL must start with http:// or https://: {base_url}"); } let client = reqwest::Client::builder() + .connect_timeout(Self::CONNECT_TIMEOUT.min(timeout)) .timeout(timeout) .build() .context("failed to build remote indexer HTTP client")?; @@ -131,12 +134,24 @@ impl RemoteIndexerClient { .send() .await .map_err(|error| { - tracing::warn!(%error, url, "Remote indexer query failed"); + tracing::warn!( + %error, + url, + is_timeout = error.is_timeout(), + is_connect = error.is_connect(), + "Remote indexer query failed" + ); KvRouterError::IndexerOffline })?; let status = response.status(); if !status.is_success() { - tracing::warn!(%status, url, "Remote indexer query returned an error status"); + tracing::warn!( + %status, + url, + model_name = %self.key.model_name, + routing_group = %self.key.routing_group, + "Remote indexer query returned an error status" + ); return Err(KvRouterError::IndexerOffline); } let body: TieredQueryResponse = response.json().await.map_err(|error| { diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index cd76ad6638c8..cd467888d548 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -58,6 +58,7 @@ use super::types::{ }; use crate::WorkerSelectionPolicyFactory; use crate::WorkerType; +use crate::indexer::KvRouterError; use crate::services::common::replica_sync::AffinityBindingEvent; /// The scheduler type every partition runs. @@ -1782,7 +1783,12 @@ impl SelectionCore { }, ) .await - .map_err(|error| SelectionError::Internal(error.to_string())) + .map_err(|error| match error { + KvRouterError::IndexerOffline => { + SelectionError::NotReady(error.to_string()) + } + other => SelectionError::Internal(other.to_string()), + }) } }; let shared_cache = query_shared_cache @@ -2269,6 +2275,40 @@ mod tests { assert_eq!(credited.worker_id, 1); } + #[tokio::test] + async fn unreachable_remote_indexer_is_reported_not_ready() { + let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind"); + let base_url = format!("http://{}", listener.local_addr().expect("addr")); + drop(listener); + let config = test_config(true); + let tracking_hash = Arc::new( + TrackingHashContext::from_config(&config).expect("valid tracking hash configuration"), + ); + let indexer_policy = IndexerPolicy::from_router_config(&config) + .expect("indexer policy") + .with_remote_indexer(base_url) + .expect("remote policy"); + let core = SelectionCore::new_inner( + config, + 1, + CancellationToken::new(), + None, + None, + SelectionHost::default(), + WorkerType::Aggregated, + true, + SelectionCacheConfig::default(), + tracking_hash, + indexer_policy, + None, + ); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + assert!(matches!( + core.select(select_request()).await, + Err(SelectionError::NotReady(_)) + )); + } + #[tokio::test] async fn remote_indexer_serves_selection_without_local_kv_listeners() { use crate::indexer::KvIndexerInterface; From 17c73cc2b3a646d8ab17a7f103c8c6d079a7d457 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Thu, 10 Sep 2026 00:04:40 -0700 Subject: [PATCH 043/102] fix(epp): surface selection client errors as 400 Selector::select_and_reserve now returns the typed SelectionError so the picker maps BadRequest (for example an over-long x-dynamo-session-id) to a 400 instead of folding every failure into a 503 routing error. Signed-off-by: Ishan Dhanani --- .../ext-proc/src/epp_router.rs | 5 +++- .../ext-proc/src/selector.rs | 25 ++++++++++++++++--- 2 files changed, 26 insertions(+), 4 deletions(-) diff --git a/deploy/inference-gateway/ext-proc/src/epp_router.rs b/deploy/inference-gateway/ext-proc/src/epp_router.rs index 68bdd2eec971..e3fc72baba9a 100644 --- a/deploy/inference-gateway/ext-proc/src/epp_router.rs +++ b/deploy/inference-gateway/ext-proc/src/epp_router.rs @@ -26,7 +26,7 @@ use std::time::Duration; use anyhow::Result; use tokio::sync::Semaphore; -use dynamo_kv_router::services::selection::WorkerSelectionPolicyRegistry; +use dynamo_kv_router::services::selection::{SelectionError, WorkerSelectionPolicyRegistry}; use dynamo_llm::http::service::metadata::extract_metadata_from_header_pairs; use dynamo_llm::protocols::common::extensions::{ AgentHints, HEADER_REQUEST_PRIORITY, HEADER_REQUEST_STRICT_PRIORITY, resolve_request_priority, @@ -359,6 +359,9 @@ impl EndpointPicker for EppRouter { let resp = match self.selector.select_and_reserve(select_req).await { Ok(resp) => resp, + Err(SelectionError::BadRequest(message)) => { + return Err(PickError::InvalidRequest(message)); + } Err(e) => return Err(PickError::RoutingFailed(e.to_string())), }; diff --git a/deploy/inference-gateway/ext-proc/src/selector.rs b/deploy/inference-gateway/ext-proc/src/selector.rs index a3c81bfebcdc..cd0a3f44ca05 100644 --- a/deploy/inference-gateway/ext-proc/src/selector.rs +++ b/deploy/inference-gateway/ext-proc/src/selector.rs @@ -174,7 +174,10 @@ impl Selector { /// Select a worker for a prompt and book its load in one operation. Takes the /// request by value so per-request fields are moved into the core request /// rather than cloned on the hot path. - pub async fn select_and_reserve(&self, req: SelectRequest) -> Result { + pub async fn select_and_reserve( + &self, + req: SelectRequest, + ) -> std::result::Result { let reservation_id = req.reservation_id; let core_req = CoreSelectAndReserveRequest { model_name: req.model_name, @@ -201,8 +204,7 @@ impl Selector { let resp = self .service .select_and_reserve_with_policy_class(core_req, req.policy_class) - .await - .map_err(|e| anyhow!("select_and_reserve failed: {e}"))?; + .await?; Ok(SelectResponse { reservation_id, worker_id: resp.worker_id, @@ -256,6 +258,7 @@ impl Drop for Selector { #[cfg(test)] mod tests { + use dynamo_kv_router::services::selection::affinity::MAX_SESSION_AFFINITY_ID_BYTES; use std::collections::HashMap; use std::sync::Mutex; use std::sync::atomic::{AtomicUsize, Ordering}; @@ -433,6 +436,22 @@ models: selector } + #[tokio::test] + async fn over_long_session_id_is_a_bad_request() { + let mut cfg = test_config(); + cfg.session_affinity_ttl_secs = Some(60.0); + let selector = Selector::new(&cfg, WorkerSelectionPolicyRegistry::default()) + .await + .expect("selector should build"); + register(&selector, vec![schedulable_registration(1)]).await; + let mut request = select_request("res-long-session"); + request.session_id = Some("s".repeat(MAX_SESSION_AFFINITY_ID_BYTES + 1)); + assert!(matches!( + selector.select_and_reserve(request).await, + Err(SelectionError::BadRequest(_)) + )); + } + #[tokio::test] async fn registered_policy_runs_through_reservation() { struct FirstEligiblePicker; From 20c307dd86c116b1f96e9e2ebadcf6d910c5103e Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Thu, 10 Sep 2026 00:09:23 -0700 Subject: [PATCH 044/102] chore(router): gate affinity test hooks and drop dead re-exports Put the session-affinity test hooks and their Notify fields behind a testing feature so release builds carry no panicking test-only API and no notify under the shard lock; lib/llm enables it as a dev-dependency. Trim the kv_router scheduler shim to the two names still imported through it, and delete the LoadEstimator polling remnants. Signed-off-by: Ishan Dhanani --- lib/kv-router/Cargo.toml | 2 ++ .../src/services/selection/affinity.rs | 18 +++++++++++++----- lib/llm/Cargo.toml | 1 + lib/llm/src/discovery/model_manager.rs | 1 - lib/llm/src/kv_router/scheduler.rs | 11 +---------- lib/llm/src/lora/load_estimator.rs | 15 +++------------ 6 files changed, 20 insertions(+), 28 deletions(-) diff --git a/lib/kv-router/Cargo.toml b/lib/kv-router/Cargo.toml index 2a2e7bc0761f..154d575a3ea8 100644 --- a/lib/kv-router/Cargo.toml +++ b/lib/kv-router/Cargo.toml @@ -21,6 +21,8 @@ profile = [] standalone-indexer = ["dep:axum", "dep:reqwest", "dep:zmq", "dep:chrono", "dep:tracing-appender"] standalone-slot-tracker = ["dep:axum", "dep:zmq"] standalone-selection = ["standalone-indexer", "dep:once_cell"] +# Test-only hooks consumed by dependent crates' tests; never enable in release. +testing = [] [dependencies] # repo diff --git a/lib/kv-router/src/services/selection/affinity.rs b/lib/kv-router/src/services/selection/affinity.rs index 0f1ac4ef4a34..b4a630ddfeef 100644 --- a/lib/kv-router/src/services/selection/affinity.rs +++ b/lib/kv-router/src/services/selection/affinity.rs @@ -113,8 +113,10 @@ struct Inner { writer_id: AtomicU64, cancel: CancellationToken, replica: OnceLock>, + #[cfg(any(test, feature = "testing"))] reaper_started: Arc, - waiter_observed: Arc, + #[cfg(any(test, feature = "testing"))] + waiter_observed: Notify, } impl Drop for Inner { @@ -208,8 +210,10 @@ impl SessionAffinity { writer_id: AtomicU64::new(0), cancel: CancellationToken::new(), replica: OnceLock::new(), + #[cfg(any(test, feature = "testing"))] reaper_started: Arc::new(Notify::new()), - waiter_observed: Arc::new(Notify::new()), + #[cfg(any(test, feature = "testing"))] + waiter_observed: Notify::new(), }); Self::spawn_reaper(&inner); tracing::info!( @@ -224,8 +228,10 @@ impl SessionAffinity { let weak = Arc::downgrade(inner); let cancel = inner.cancel.clone(); let period = inner.ttl.min(Duration::from_secs(30)); + #[cfg(any(test, feature = "testing"))] let reaper_started = inner.reaper_started.clone(); tokio::spawn(async move { + #[cfg(any(test, feature = "testing"))] reaper_started.notify_one(); loop { tokio::select! { @@ -304,6 +310,7 @@ impl SessionAffinity { } Entry::Occupied(mut entry) => match entry.get_mut() { AffinityEntry::Initializing { notify, .. } => { + #[cfg(any(test, feature = "testing"))] self.inner.waiter_observed.notify_one(); // Register before releasing the entry so a commit or a // dropped initialization between the two cannot be missed. @@ -431,6 +438,7 @@ impl SessionAffinity { self.inner.apply_replica_update(session_id, target, version) } + #[cfg(any(test, feature = "testing"))] pub fn entry_count(&self) -> usize { self.inner.entry_count.load(Ordering::Relaxed) } @@ -440,17 +448,17 @@ impl SessionAffinity { self.inner.cancel.clone() } - /// Test hook: the reaper task has started. + #[cfg(any(test, feature = "testing"))] pub async fn wait_for_reaper(&self) { self.inner.reaper_started.notified().await; } - /// Test hook: a request observed an initializing session and is waiting. + #[cfg(any(test, feature = "testing"))] pub async fn wait_for_initializing_waiter(&self) { self.inner.waiter_observed.notified().await; } - /// Test hook: expire an idle bound session now. + #[cfg(any(test, feature = "testing"))] pub fn expire_for_test(&self, session_id: &str) { let Some(mut entry) = self.inner.entries.get_mut(session_id) else { panic!("session affinity entry missing"); diff --git a/lib/llm/Cargo.toml b/lib/llm/Cargo.toml index 23efc4733b0e..5874cd391ac9 100644 --- a/lib/llm/Cargo.toml +++ b/lib/llm/Cargo.toml @@ -229,6 +229,7 @@ llm-multimodal = { version = "=1.9.0", optional = true } llm-tokenizer = { version = "=1.6.0", optional = true } [dev-dependencies] +dynamo-kv-router = { workspace = true, features = ["testing"] } dynamo-runtime = { workspace = true, features = ["testing"] } tokio = { workspace = true, features = ["test-util"] } approx = "0.5" diff --git a/lib/llm/src/discovery/model_manager.rs b/lib/llm/src/discovery/model_manager.rs index 2aa72a04b043..0cd43f314e4f 100644 --- a/lib/llm/src/discovery/model_manager.rs +++ b/lib/llm/src/discovery/model_manager.rs @@ -2275,7 +2275,6 @@ impl ModelManager { buckets_per_second: config.buckets_per_second, predictor_type: config.predictor_type, ema_alpha: config.ema_alpha, - ..Default::default() }); let domain_cancel = cancel_token.child_token(); *domain.controller_cancel.lock() = Some(domain_cancel.clone()); diff --git a/lib/llm/src/kv_router/scheduler.rs b/lib/llm/src/kv_router/scheduler.rs index 97fe82fb85aa..e8892968ae37 100644 --- a/lib/llm/src/kv_router/scheduler.rs +++ b/lib/llm/src/kv_router/scheduler.rs @@ -5,16 +5,7 @@ //! partition of an embedded selection service (`kv_router::embedded`); these //! re-exports are the runtime-free scheduling vocabulary it speaks. -pub use dynamo_kv_router::scheduling::overlap_refresh::{ - NoopOverlapScoresRefresh, OverlapScoresRefresh, RefreshedOverlap, -}; pub use dynamo_kv_router::scheduling::queue::{ SchedulerBookingCleanup, SchedulerBookingDescriptor, }; -pub use dynamo_kv_router::scheduling::{ - AdmittedSchedulingResponse, AdvisorySchedulingResponse, AttemptId, KvSchedulerError, - LocalScheduler, NonMaxOverlapSelectionObserver, OverloadedWorkerProvider, PotentialLoad, - ScheduleRequest, SchedulingRequest, SchedulingResponse, TierOverlapBlocks, - WorkerAvailabilityProvider, -}; -pub use dynamo_kv_router::selector::DefaultWorkerSelector; +pub use dynamo_kv_router::scheduling::{KvSchedulerError, PotentialLoad}; diff --git a/lib/llm/src/lora/load_estimator.rs b/lib/llm/src/lora/load_estimator.rs index 9830d33c0a87..8ca17c011b64 100644 --- a/lib/llm/src/lora/load_estimator.rs +++ b/lib/llm/src/lora/load_estimator.rs @@ -4,7 +4,7 @@ //! LORA Load Estimator //! //! Tracks LORA adapter usage over time to estimate load for allocation decisions. -//! Supports single-router (polling) and multi-router (event-based) modes. +//! Event-based: routers report arrivals and completions as they happen. //! //! The primary load signal is **arrival count in a sliding window**, tracked by //! a lock-free [`BucketedRateCounter`] per LoRA. An optional [`LoadPredictor`] @@ -230,7 +230,6 @@ impl std::fmt::Debug for LoraLoadData { #[derive(Debug, Clone)] pub struct LoadEstimatorConfig { - pub poll_interval: Duration, /// Sliding window size for request-rate calculation. pub rate_window: Duration, pub buckets_per_second: u64, @@ -241,7 +240,6 @@ pub struct LoadEstimatorConfig { impl Default for LoadEstimatorConfig { fn default() -> Self { Self { - poll_interval: Duration::from_secs(5), rate_window: Duration::from_secs(30), buckets_per_second: 1, predictor_type: PredictorType::Ema, @@ -597,15 +595,8 @@ impl LoadEstimator { } } - /// Update active counts from a polled snapshot. - /// - /// **Polling-mode caveat**: arrivals are approximated as the per-poll - /// delta `max(0, current - prev)`, since worker snapshots do not expose - /// request-start events. This is a *lower bound* on real arrivals: - /// in-interval churn (e.g., 10 requests finishing while 10 new ones start - /// — net delta 0) is invisible, and sub-interval oscillation is lost. - /// Event-based mode (`handle_event` / `increment_load`) gives accurate - /// arrival rates; prefer it when arrival precision matters. + /// Update active counts from a snapshot; arrivals are approximated as the + /// positive delta since the previous snapshot. #[cfg(test)] fn update_from_counts(&self, lora_counts: HashMap) { let now = Instant::now(); From f81f949832ff685417023fb0fdb2538d8881815d Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Thu, 10 Sep 2026 00:18:20 -0700 Subject: [PATCH 045/102] fix(router): reject cache_salt on tiered indexer queries and skip no-op scheduler publishes /query_tiered_by_hash now rejects a cache_salt with 400 like /query_by_hash instead of silently ignoring it. publish_scheduler_config only wakes the scheduler when the worker map actually changed. Signed-off-by: Ishan Dhanani --- lib/kv-router/src/services/indexer/server.rs | 51 ++++++++++++++++++- .../src/services/selection/core/mod.rs | 14 +++-- 2 files changed, 61 insertions(+), 4 deletions(-) diff --git a/lib/kv-router/src/services/indexer/server.rs b/lib/kv-router/src/services/indexer/server.rs index 1f1483d39332..c77092160269 100644 --- a/lib/kv-router/src/services/indexer/server.rs +++ b/lib/kv-router/src/services/indexer/server.rs @@ -122,6 +122,16 @@ struct QueryRequest { cache_salt: Option, } +/// Server-side view of `TieredQueryByHashRequest` that also captures the +/// `cache_salt` some `/query_by_hash` callers send, so it can be rejected here +/// as well instead of being silently ignored. +#[derive(Deserialize)] +struct TieredQueryByHashBody { + #[serde(flatten)] + inner: TieredQueryByHashRequest, + cache_salt: Option, +} + #[derive(Deserialize)] struct QueryByHashRequest { block_hashes: Vec, @@ -393,9 +403,23 @@ async fn query_by_hash( /// overlap analysis and side-indexer merge. async fn query_tiered_by_hash( State(state): State>, - Json(req): Json, + Json(TieredQueryByHashBody { + inner: req, + cache_salt, + }): Json, ) -> Response { let model = req.model_name.clone(); + if cache_salt.is_some() { + let mut resp = ( + StatusCode::BAD_REQUEST, + Json(serde_json::json!({ + "error": "cache_salt is not accepted by /query_tiered_by_hash; block_hashes must already include the intended cache salt" + })), + ) + .into_response(); + resp.extensions_mut().insert(AccessLogModel(model)); + return resp; + } let key = RoutingPartitionId::new(req.model_name, req.routing_group); let Some(ie) = state.registry.get_indexer(&key) else { let mut resp = ( @@ -702,6 +726,31 @@ mod tests { body } + #[tokio::test] + async fn query_tiered_by_hash_rejects_cache_salt() { + let app = create_router(Arc::new(AppState { + registry: Arc::new(WorkerRegistry::new(1)), + access_log_sink: None, + #[cfg(feature = "metrics")] + prom_registry: prometheus::Registry::new(), + })); + + let salted = app + .oneshot( + Request::builder() + .method("POST") + .uri("/query_tiered_by_hash") + .header(header::CONTENT_TYPE, "application/json") + .body(Body::from( + r#"{"block_hashes":[1],"model_name":"model","routing_group":"default","cache_salt":"s"}"#, + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(salted.status(), StatusCode::BAD_REQUEST); + } + #[tokio::test] async fn query_tiered_by_hash_returns_wire_tiered_matches() { use super::super::backend::test_util::store_event; diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index cd467888d548..6da1e347db00 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -862,9 +862,17 @@ impl SelectionCore { return Ok(()); }; let workers = self.catalog.scheduler_configs_for_key(key); - entry.workers_tx.send(workers).map_err(|_| { - SelectionError::Internal(format!("scheduler worker watch closed for {key}")) - }) + // Lifecycle transitions between non-schedulable states publish the same + // map; skipping them saves the scheduler a wake and a full map clone. + entry.workers_tx.send_if_modified(|current| { + if *current == workers { + false + } else { + *current = workers; + true + } + }); + Ok(()) } fn ready_entry(&self, key: &RoutingPartitionId) -> Result, SelectionError> { From 774115241be9a53a7d77d1b19f4bdcd598d6a800 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Thu, 10 Sep 2026 00:25:40 -0700 Subject: [PATCH 046/102] fix(router): reject unsupported data-parallel shapes and correct standalone docs The EPP accepts only DYN_EPP_DATA_PARALLEL_SIZE=1 because it cannot convey a selected rank to the pod, and the standalone service marks a multi-rank worker Incomplete when it sets replay_endpoint (the replay protocol carries no rank). Docs now show the kv_hint envelope, the real builder method, the lru and advisory behavior, and replication requiring a replica mesh. Signed-off-by: Ishan Dhanani --- .../ext-proc/src/epp_standalone_config.rs | 25 +++++++++-- .../router/standalone-indexer.md | 2 +- .../router/standalone-selection.md | 45 +++++++++++++------ .../kv-aware-routing/vanilla-vllm-onramp.mdx | 4 +- lib/kv-router/src/services/indexer/backend.rs | 7 ++- .../src/services/selection/core/mod.rs | 31 ++++++++++--- .../src/services/selection/ingress.rs | 9 +++- lib/kv-router/src/services/selection/types.rs | 3 +- 8 files changed, 94 insertions(+), 32 deletions(-) diff --git a/deploy/inference-gateway/ext-proc/src/epp_standalone_config.rs b/deploy/inference-gateway/ext-proc/src/epp_standalone_config.rs index 1f2582df5644..24e31f1cbf7f 100644 --- a/deploy/inference-gateway/ext-proc/src/epp_standalone_config.rs +++ b/deploy/inference-gateway/ext-proc/src/epp_standalone_config.rs @@ -140,9 +140,13 @@ pub struct EppStandaloneConfig { /// KV-cache block size; MUST equal the inference engine block size. #[validate(range(min = 1, message = "DYN_KV_CACHE_BLOCK_SIZE must be >= 1"))] pub block_size: u32, - /// Data-parallel ranks per worker pod. Each rank publishes KV events on - /// `kv_event_port + rank * kv_event_port_stride`. - #[validate(range(min = 1, message = "DYN_EPP_DATA_PARALLEL_SIZE must be >= 1"))] + /// Data-parallel ranks per worker pod. Only `1` is accepted: the EPP routes + /// by worker and cannot convey the selected rank to the pod. + #[validate(range( + min = 1, + max = 1, + message = "DYN_EPP_DATA_PARALLEL_SIZE must be 1; multi-rank standalone serving is unsupported" + ))] pub data_parallel_size: u32, /// Port distance between consecutive data-parallel ranks' KV event ports. #[validate(range(min = 1, message = "DYN_EPP_KV_EVENT_PORT_STRIDE must be >= 1"))] @@ -557,6 +561,21 @@ mod tests { ); } + #[test] + fn data_parallel_size_above_one_fails() { + assert!( + parse_cfg(&[ + ("DYN_EPP_INFERENCE_POOL_NAME", "vllm-qwen-pool"), + ("POD_NAMESPACE", "inference"), + ("DYN_MODEL_NAME", "Qwen/Qwen3-0.6B"), + ("DYN_EPP_TOKENIZER_SERVICE_URL", "http://vllm-render:8000"), + ("DYN_EPP_TOKENIZER_PROTOCOL", "vllm-render"), + ("DYN_EPP_DATA_PARALLEL_SIZE", "2"), + ]) + .is_err() + ); + } + #[test] fn replay_port_is_optional_and_uses_the_kv_event_name() { let cfg = parse_cfg(&[ diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-indexer.md b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-indexer.md index 90d9c51156f3..090bd8a82cd2 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-indexer.md +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-indexer.md @@ -415,7 +415,7 @@ Mooncake score summary, so a caller can run its own overlap analysis: "block_size": 16, "tiered": { "device": { "scores": [[{"worker_id": 7, "dp_rank": 0}, 2]], "frequencies": [] }, - "lower_tier": [["HostPinned", { "hits": [[{"worker_id": 7, "dp_rank": 0}, 1]] }]] + "lower_tier": [["host_pinned", { "hits": [[{"worker_id": 7, "dp_rank": 0}, 1]] }]] } } ``` diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-selection.md b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-selection.md index 7dfbda2fd44d..91b108654c19 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-selection.md +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-selection.md @@ -69,7 +69,7 @@ APIs. Those bindings should wrap `SelectionService` rather than construct | `--remote-indexer-url` | none | Base URL of a [standalone indexer](standalone-indexer.md) that serves the primary KV index. The selector then does not subscribe to worker KV events; workers publish to the indexer instead. Requires `use_kv_events=true`. | | `--replica-sync-port` | none | Local ZMQ PUB port for active-load lifecycle events. The selector binds `tcp://*:` internally. | | `--replica-sync-peers` | none | Comma-separated ZMQ PUB endpoints for selector peers. Requires `--replica-sync-port`. | -| `--session-affinity-ttl-secs` | none | Pin each request `session_id` to the worker that served it for this long after its last request. Bindings replicate to peers over the replica mesh (`dynamo.session-affinity.v1` topic). | +| `--session-affinity-ttl-secs` | none | Pin each request `session_id` to the worker that served it for this long after its last request (1 to 31,536,000 seconds). When `--replica-sync-port` is configured, bindings replicate to peers over the replica mesh (`dynamo.session-affinity.v1` topic). | | `--selection-cache-ttl-secs` | `120` | Seconds an unclaimed pending selection lives before eviction. | | `--selection-cache-max-entries` | `4096` | Maximum resident pending selections, evicting oldest first. | | `--selection-cache-max-bytes` | `268435456` | Approximate byte budget across resident pending selections. | @@ -91,7 +91,7 @@ frontend uses, so the two index the same way: | `use_kv_events=true` (default) | Event-driven from worker ZMQ events | Not recorded | | `use_kv_events=true`, `DYN_ROUTER_PREDICTED_TTL_SECS` set | Event-driven | Recorded into a short-TTL side indexer merged into device scores by per-worker max | | `use_kv_events=false` | Approximate: no events, entries expire after `router_ttl_secs` | Recorded into the primary | -| `--remote-indexer-url` (or `SelectionServiceBuilder::remote_indexer`) | Served by a standalone indexer, queried through `POST /query_tiered_by_hash` | Recorded into the side indexer when `DYN_ROUTER_PREDICTED_TTL_SECS` is set, else not recorded | +| `--remote-indexer-url` (or `SelectionServiceBuilder::kv_index(KvIndexSource::Remote(..))`) | Served by a standalone indexer, queried through `POST /query_tiered_by_hash` | Recorded into the side indexer when `DYN_ROUTER_PREDICTED_TTL_SECS` is set, else not recorded | With a remote primary, worker records do not need `kv_events_endpoint(s)` to become schedulable, no ZMQ listener is started in the selector, and `/dump` @@ -101,8 +101,9 @@ lookup failure is reported as an offline indexer for that selection. A routing decision is recorded when a reservation is booked (`select_and_reserve`, or `select` followed by `POST /reservations`), never for a query-only `select`. Hash-only reservations that supply `sequence_hashes` without `block_hashes` or -`token_ids` are not recorded. `router_approximate_cache_policy=lru` is not yet -supported here and falls back to TTL retention with a warning. +`token_ids` are not recorded. `router_approximate_cache_policy=lru` is rejected at startup when +`use_kv_events=true`; with `use_kv_events=false` it falls back to TTL retention +with a warning. The standalone expiry guard measures absolute age from admission; output progress does not refresh it. Periodic cleanup therefore reclaims stale state approximately five to six minutes after @@ -130,11 +131,14 @@ Content-Type: application/json "kv_events_endpoints": { "0": "tcp://worker:5557", "1": "tcp://worker:5558" - }, - "replay_endpoint": "tcp://worker:5560" + } } ``` +`replay_endpoint` (KV event gap recovery) is accepted only when +`data_parallel_size` is 1; a multi-rank worker that sets it stays `incomplete` +because the replay protocol carries no rank. + `worker_id` is service-wide, not scoped by model or routing group. `POST /workers` is an upsert and returns `201`: reusing an existing ID replaces its catalog record. If the model or routing group changes, the worker is removed from the @@ -234,15 +238,28 @@ blocks. `decode_busy` compares it against the worker's `total_kv_blocks` at `conditional_disagg_decode_busy_threshold` and is omitted when either the threshold or the capacity is unknown. -A booking response may also carry a `router_hint` when another worker of the -same `router_hint_worker_type` holds a longer cached prefix than the chosen -worker and advertises a source control endpoint for the matching DP rank: +A booking response may also carry a `kv_hint` when another worker of the same +`router_hint_worker_type` (the worker capability key is `router_hint`) holds a +longer cached prefix than the chosen worker and advertises a source control +endpoint for the matching DP rank. The hint is a versioned action envelope +whose `message_id` is the `selection_id`: ```json { - "router_hint": { - "source_control_endpoint": "tcp://worker-1:5600", - "block_hashes": [8713492873, 1928374650] + "kv_hint": { + "protocol_version": "0.1", + "message_id": "select-123", + "actions": [ + { + "action_id": "a1", + "action_type": "kv.fetch", + "action_version": "1.0", + "payload": { + "source_control_endpoint": "tcp://worker-1:5600", + "block_hashes": [8713492873, 1928374650] + } + } + ] } } ``` @@ -278,8 +295,8 @@ and the response adds the chosen worker's projected load: omitted when that threshold is unset. This is the probe a disaggregation coordinator uses to decide whether to bypass remote prefill: one advisory `select` against the prefill pool answers "would the prefill worker I'd get be -busy?" without a separate load read. `advisory` is not accepted on -`select_and_reserve`. All `overlap` +busy?" without a separate load read. `advisory` is ignored by +`select_and_reserve`, which always books. All `overlap` values are matched token counts. `gpu`, `cpu`, and `disk` use the cumulative Mooncake tier semantics documented in the standalone indexer's [per-instance tier breakdown](standalone-indexer.md#per-instance-tier-breakdown). diff --git a/docs/fern/pages/kubernetes/kv-aware-routing/vanilla-vllm-onramp.mdx b/docs/fern/pages/kubernetes/kv-aware-routing/vanilla-vllm-onramp.mdx index 5e23deadd132..ebe56b753c2a 100644 --- a/docs/fern/pages/kubernetes/kv-aware-routing/vanilla-vllm-onramp.mdx +++ b/docs/fern/pages/kubernetes/kv-aware-routing/vanilla-vllm-onramp.mdx @@ -497,8 +497,8 @@ For a single EPP replica, set `spec.replicas: 1` and remove `DYN_EPP_PEER_SERVIC | `DYN_KV_CACHE_BLOCK_SIZE` | Yes | Must equal vLLM `--block-size`. | | `DYN_EPP_KV_EVENT_PORT` | No | Per-pod vLLM KV event port; defaults to `5557`. | | `DYN_EPP_KV_EVENT_REPLAY_PORT` | No | Per-pod replay port for KV event gap recovery. | -| `DYN_EPP_DATA_PARALLEL_SIZE` | No | Data-parallel ranks per pod; defaults to `1`. Rank `r` publishes KV events on `DYN_EPP_KV_EVENT_PORT + r * DYN_EPP_KV_EVENT_PORT_STRIDE`. | -| `DYN_EPP_KV_EVENT_PORT_STRIDE` | No | Port distance between consecutive ranks' KV event ports; defaults to `1`. | +| `DYN_EPP_DATA_PARALLEL_SIZE` | No | Data-parallel ranks per pod; defaults to `1`, and `1` is the only accepted value because the EPP cannot convey a selected rank to the pod. | +| `DYN_EPP_KV_EVENT_PORT_STRIDE` | No | Port distance between consecutive ranks' KV event ports; defaults to `1`. Unused while `DYN_EPP_DATA_PARALLEL_SIZE` is fixed at `1`. | | `DYN_EPP_MAX_NUM_BATCHED_TOKENS` | Policy-dependent | Must equal vLLM `--max-num-batched-tokens`; required when router queue thresholds are enabled. | | `DYN_EPP_TOTAL_KV_BLOCKS` | No | Optional per-worker total KV block hint. | | `DYN_EPP_SELECTION_INDEXER_THREADS` | No | KV indexer thread count; defaults to `4`. | diff --git a/lib/kv-router/src/services/indexer/backend.rs b/lib/kv-router/src/services/indexer/backend.rs index 352bcee20a13..65f2774616cd 100644 --- a/lib/kv-router/src/services/indexer/backend.rs +++ b/lib/kv-router/src/services/indexer/backend.rs @@ -210,10 +210,9 @@ impl IndexerPolicy { /// Resolve the indexer shape for `config`. /// - /// Rejects the same combinations the frontend rejects. Capacity-bounded LRU - /// retention needs per-request acquire/release fencing the service does not - /// yet have, so `router_approximate_cache_policy=lru` falls back to TTL - /// with a warning instead of failing. + /// Rejects the same combinations the frontend rejects. The local side indexer + /// is TTL-only, so `router_approximate_cache_policy=lru` is rejected with + /// `use_kv_events=true` and falls back to TTL with a warning otherwise. pub fn from_router_config(config: &KvRouterConfig) -> Result { if config.use_kv_events && config.router_approximate_cache_policy == ApproximateCachePolicyKind::Lru diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index 6da1e347db00..d54702897df5 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -987,11 +987,6 @@ impl SelectionCore { advisory, } = operation; self.ensure_running()?; - if advisory && book { - return Err(SelectionError::BadRequest( - "advisory selection cannot book a reservation".to_string(), - )); - } let entry = self.ready_entry(&key)?; @@ -2847,6 +2842,32 @@ mod tests { assert!(!parent.is_cancelled()); } + #[tokio::test] + async fn multi_rank_worker_with_replay_endpoint_is_incomplete() { + let core = SelectionCore::try_new_local( + test_config(true), + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ) + .expect("valid test config"); + + let record = core + .upsert_worker(WorkerRequest { + data_parallel_size: Some(2), + kv_events_endpoints: HashMap::from([ + (0, "tcp://127.0.0.1:5557".to_string()), + (1, "tcp://127.0.0.1:5558".to_string()), + ]), + replay_endpoint: Some("tcp://127.0.0.1:5600".to_string()), + ..worker(1) + }) + .await + .expect("worker upsert"); + assert_eq!(record.lifecycle, WorkerLifecycle::Incomplete, "{record:?}"); + assert!(!core.indexer_registry.has_worker(1)); + } + #[tokio::test] async fn shutdown_cancels_listeners() { let parent = CancellationToken::new(); diff --git a/lib/kv-router/src/services/selection/ingress.rs b/lib/kv-router/src/services/selection/ingress.rs index ac67a1efa1f1..a54b0e945212 100644 --- a/lib/kv-router/src/services/selection/ingress.rs +++ b/lib/kv-router/src/services/selection/ingress.rs @@ -70,7 +70,7 @@ impl KvEventIngress for ZmqDirectIngress { fn missing_metadata(&self, record: &WorkerCatalogRecord) -> Vec { let endpoints = record.listener_endpoints(); - record + let mut missing: Vec = record .dp_ranks() .filter(|rank| { endpoints @@ -78,7 +78,12 @@ impl KvEventIngress for ZmqDirectIngress { .is_none_or(|endpoint| endpoint.is_empty()) }) .map(|rank| format!("kv_events endpoint is required for dp_rank {rank}")) - .collect() + .collect(); + // The replay protocol carries no rank, and each rank opens its own socket. + if record.dp_size() > 1 && record.replay_endpoint.is_some() { + missing.push("replay_endpoint is supported only for data_parallel_size 1".to_string()); + } + missing } async fn reconcile( diff --git a/lib/kv-router/src/services/selection/types.rs b/lib/kv-router/src/services/selection/types.rs index 846739b34fb7..a36eeaa1de62 100644 --- a/lib/kv-router/src/services/selection/types.rs +++ b/lib/kv-router/src/services/selection/types.rs @@ -426,7 +426,8 @@ pub struct SelectRequest { /// The response then carries the chosen worker's `worker_load` snapshot /// and `prefill_busy` evaluation. The request never waits in the router /// queue; when the queue would have rejected it the call fails the same - /// way an admitted selection does. Rejected on `select_and_reserve`. + /// way an admitted selection does. Ignored on `select_and_reserve`, which + /// always books. #[serde(default)] pub advisory: bool, } From 960500e46965f7d916f856e0ae2ba0b29131768f Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Thu, 10 Sep 2026 00:30:50 -0700 Subject: [PATCH 047/102] fix(router): free core bookings the caller never consumes Core bookings now carry the queue lifecycle lease, so a selection future dropped after the actor booked (or while still queued) releases its booking instead of leaking it. A worker that drained while the request waited in the queue also releases the booking before the selection fails. Signed-off-by: Ishan Dhanani --- .../src/services/selection/core/mod.rs | 155 +++++++++++++++--- 1 file changed, 132 insertions(+), 23 deletions(-) diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index d54702897df5..abd871fce9a7 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -23,12 +23,14 @@ use crate::protocols::{ WorkerAffinityTarget, WorkerConfigLike, WorkerId, WorkerWithDpRank, }; use crate::scheduling::config::RouterConfigOverride; +use crate::scheduling::queue::SchedulerBookingDescriptor; use crate::scheduling::selector::WorkerSelectionPolicy; use crate::scheduling::{ - KvSchedulerError, LocalScheduler, LoraWorkerFilter, OverlapAnalysis, OverlapSignals, - OverloadedWorkerProvider, PotentialLoad, PrefillLoadEstimator, ScheduleMode, ScheduleRequest, - SessionContext, TieredOverlapRefresher, WorkerAvailabilityProvider, effective_prefill_tokens, - narrow_allowed_worker_ids_by_lora, prefill_load_hint_from_effective_tokens, + AdmissionAttempt, KvSchedulerError, LocalScheduler, LoraWorkerFilter, OverlapAnalysis, + OverlapSignals, OverloadedWorkerProvider, PotentialLoad, PrefillLoadEstimator, ScheduleMode, + ScheduleRequest, SessionContext, TieredOverlapRefresher, WorkerAvailabilityProvider, + effective_prefill_tokens, narrow_allowed_worker_ids_by_lora, + prefill_load_hint_from_effective_tokens, }; use crate::sequences::{ ActiveSequencesMultiWorker, ReplicaRequestLeaseObserver, ReplicaWorkerPolicy, SequenceError, @@ -1055,8 +1057,10 @@ impl SelectionCore { retain_kv_transfer_chain, ) .await?; + // The queue lease frees a booking whose response is never consumed + // (the caller dropped this future after the actor booked). let mode = if book { - ScheduleMode::Tracked { + ScheduleMode::TrackedWithLifecycle { request_id: selection_id.clone().ok_or_else(|| { SelectionError::Internal( "booked selection did not include a selection ID".to_string(), @@ -1122,7 +1126,7 @@ impl SelectionCore { routing_constraints, shared_cache_hits, }; - let (response, advisory_load) = tokio::select! { + let (response, advisory_load, attempt) = tokio::select! { biased; _ = self.cancel_token.cancelled() => { return Err(SelectionError::Scheduler(KvSchedulerError::SubscriberShutdown)); @@ -1133,25 +1137,45 @@ impl SelectionCore { .scheduler .select_without_admission(schedule_request) .await - .map(|advisory| (advisory.response, Some(advisory.selected_worker_load))) + .map(|advisory| { + ( + advisory.response, + Some(advisory.selected_worker_load), + AdmissionAttempt::Untracked, + ) + }) } else { entry .scheduler - .schedule_request(schedule_request) + .schedule_request_admitted(schedule_request) .await - .map(|response| (response, None)) + .map(|admitted| (admitted.response, None, admitted.attempt)) } } => result?, }; - let endpoint = self + let Some(endpoint) = self .catalog .schedulable_endpoint(response.best_worker.worker_id, &key) - .ok_or_else(|| { - SelectionError::Internal(format!( - "selected worker {} is no longer schedulable", - response.best_worker.worker_id - )) - })?; + else { + // The worker drained while this request was queued; release the + // booking the actor just made for it. + if let (AdmissionAttempt::Tracked(attempt_id), Some(request_id)) = + (attempt, selection_id.as_deref()) + { + entry + .scheduler + .booking_cleanup() + .enqueue(SchedulerBookingDescriptor { + request_id: request_id.to_string(), + worker: response.best_worker, + attempt_id, + }); + } + return Err(SelectionError::Internal(format!( + "selected worker {} is no longer schedulable", + response.best_worker.worker_id + ))); + }; let overlap = MooncakeOverlapSummary::from_selected_worker_tiers( &response.selected_worker_tiers, entry.block_size, @@ -2062,17 +2086,21 @@ mod tests { } } - async fn wait_for_pending_selection(core: &SelectionCore) { - tokio::time::timeout(Duration::from_secs(1), async { - loop { - if core.loads(Some("model"), Some("default"))[0].pending_count == 1 { - return; - } + async fn wait_until(what: &str, mut condition: impl FnMut() -> bool) { + tokio::time::timeout(Duration::from_secs(2), async { + while !condition() { tokio::task::yield_now().await; } }) .await - .expect("selection did not queue"); + .unwrap_or_else(|_| panic!("timed out waiting for {what}")); + } + + async fn wait_for_pending_selection(core: &SelectionCore) { + wait_until("pending selection", || { + core.loads(Some("model"), Some("default"))[0].pending_count == 1 + }) + .await; } fn assert_shutdown_error(error: SelectionError) { @@ -3048,6 +3076,87 @@ mod tests { )); } + #[tokio::test] + async fn booking_is_freed_when_selected_worker_drained_while_queued() { + let mut config = test_config(false); + config.router_queue_threshold = Some(0.0); + let core = Arc::new( + SelectionCore::try_new_local( + config, + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ) + .expect("valid test config"), + ); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + let key = RoutingPartitionId::new("model", "default"); + let entry = core.entry(&key).expect("entry"); + core.select_and_reserve(reserve_request("res-a")) + .await + .expect("initial reservation"); + + let queued_core = core.clone(); + let queued = tokio::spawn(async move { + queued_core + .select_and_reserve(reserve_request("queued")) + .await + }); + wait_for_pending_selection(&core).await; + // First half of `delete_worker`: the worker drains while the request waits. + core.catalog + .set_lifecycle(1, WorkerLifecycle::Draining, Vec::new()); + core.free_reservation("res-a").await.expect("free res-a"); + + let err = tokio::time::timeout(Duration::from_secs(2), queued) + .await + .expect("queued selection timed out") + .expect("task panicked") + .expect_err("drained worker is not schedulable"); + assert!( + matches!(&err, SelectionError::Internal(m) if m.contains("no longer schedulable")), + "{err:?}" + ); + wait_until("booking release", || !entry.scheduler.has_request("queued")).await; + } + + #[tokio::test] + async fn dropped_selection_future_frees_its_booking() { + let core = SelectionCore::try_new_local( + test_config(false), + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ) + .expect("valid test config"); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + let key = RoutingPartitionId::new("model", "default"); + let entry = core.entry(&key).expect("entry"); + + // Drive the selection by hand so the actor's response is delivered but + // never consumed: poll until the booking exists, then drop the future. + let mut selection = Box::pin(core.select_and_reserve(reserve_request("dropped"))); + let mut context = std::task::Context::from_waker(std::task::Waker::noop()); + wait_until("scheduler booking", || { + if entry.scheduler.has_request("dropped") { + return true; + } + assert!( + selection.as_mut().poll(&mut context).is_pending(), + "selection completed before the booking was observed" + ); + false + }) + .await; + drop(selection); + + wait_until("booking release", || { + !entry.scheduler.has_request("dropped") + }) + .await; + assert!(core.reservation_index.read().get("dropped").is_none()); + } + #[tokio::test] async fn lifecycle_operations_find_reservation_in_later_entry() { let mut config = test_config(false); From 3003cb8ae0046f4a943895a8861dacfa84206a13 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Thu, 10 Sep 2026 00:30:51 -0700 Subject: [PATCH 048/102] fix(router): hold the rebound session lease on concurrent affinity failover When two bookings for the same session both find its worker gone, the second one now keeps the lease of the binding the first one created instead of discarding it and leaving the reservation without a session hold. Signed-off-by: Ishan Dhanani --- .../src/services/selection/core/mod.rs | 49 ++++++++++++++++--- 1 file changed, 43 insertions(+), 6 deletions(-) diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index abd871fce9a7..3020c647eeda 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -46,7 +46,7 @@ use crate::services::indexer::registry::WorkerRegistry; use crate::services::overlap::MooncakeOverlapSummary; use crate::tracking_hash::{TrackingHashContext, TrackingHashScope}; -use super::affinity::{Acquired, AffinityError, AffinityLease, SessionAffinity}; +use super::affinity::{AcquireStep, Acquired, AffinityError, AffinityLease, SessionAffinity}; use super::catalog::WorkerCatalog; use super::error::SelectionError; use super::ingress::{KvEventIngress, ZmqDirectIngress}; @@ -1511,7 +1511,8 @@ impl SelectionCore { /// Bind (or confirm) `session_id` to the worker a booking landed on and /// keep the lease with the booking. A bound session whose worker was not - /// selected (it left) is invalidated and rebound. + /// selected (it left) is invalidated and rebound; when a concurrent request + /// already rebound it, this booking holds that binding's lease instead. fn bind_session( table: &SessionAffinity, hold: Acquired, @@ -1531,10 +1532,9 @@ impl SelectionCore { Acquired::Bound { mut lease, .. } => { lease.invalidate(); match table.try_acquire(session_id, None) { - Ok(super::affinity::AcquireStep::Initialize(init)) => { - init.commit(selected).ok() - } - _ => None, + Ok(AcquireStep::Initialize(init)) => init.commit(selected).ok(), + Ok(AcquireStep::Bound { lease, .. }) => return Some(lease), + Ok(AcquireStep::Wait(_)) | Err(_) => None, } } }; @@ -2811,6 +2811,43 @@ mod tests { ); } + #[tokio::test] + async fn concurrent_failover_keeps_the_winning_binding_lease() { + let table = SessionAffinity::new(Duration::from_secs(60)).expect("affinity table"); + let departed = WorkerWithDpRank::new(1, 0); + let Acquired::Initialize(init) = table.acquire("s", None).await.expect("acquire") else { + panic!("fresh session must initialize"); + }; + drop(init.commit(SessionTarget::new(1, Some(0))).expect("commit")); + let first = table.acquire("s", None).await.expect("first acquire"); + let second = table.acquire("s", None).await.expect("second acquire"); + assert!( + matches!(first, Acquired::Bound { target, .. } if target.worker_id == departed.worker_id) + ); + + // Worker 1 left; both bookings landed elsewhere. The first rebinds the + // session, the second must hold that binding rather than end up leaseless. + let winner = SelectionCore::bind_session(&table, first, "s", WorkerWithDpRank::new(2, 0)) + .expect("first failover rebinds"); + let follower = + SelectionCore::bind_session(&table, second, "s", WorkerWithDpRank::new(3, 0)) + .expect("second failover holds the rebound session"); + let bound = table + .query_target("s", None) + .expect("query") + .expect("bound"); + assert_eq!(bound.worker_id, 2); + drop(winner); + drop(follower); + assert_eq!( + table + .query_target("s", None) + .expect("query") + .map(|t| t.worker_id), + Some(2) + ); + } + #[tokio::test] async fn injected_availability_provider_restricts_selection() { let available: Arc>>>> = From 92ea625801391e3c817dbf42a56c502431e8b70e Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Thu, 10 Sep 2026 00:30:51 -0700 Subject: [PATCH 049/102] fix(router): re-register KV listeners a cancelled worker update removed Worker reconciliation retained a rank whenever its catalog metadata was unchanged, so a listener lost to an update cancelled mid-deregistration was never recreated. Retention now requires the listener to exist. Signed-off-by: Ishan Dhanani --- .../src/services/indexer/registry.rs | 6 ++++ .../src/services/selection/core/mod.rs | 28 +++++++++++++++++++ .../src/services/selection/ingress.rs | 5 ++++ 3 files changed, 39 insertions(+) diff --git a/lib/kv-router/src/services/indexer/registry.rs b/lib/kv-router/src/services/indexer/registry.rs index d8a31dee2238..5c99a5e58713 100644 --- a/lib/kv-router/src/services/indexer/registry.rs +++ b/lib/kv-router/src/services/indexer/registry.rs @@ -690,6 +690,12 @@ impl WorkerRegistry { self.workers.contains_key(&worker_id) } + pub fn has_listener(&self, worker_id: WorkerId, dp_rank: u32) -> bool { + self.workers + .get(&worker_id) + .is_some_and(|entry| entry.listeners.contains_key(&dp_rank)) + } + pub fn list(&self) -> Vec { self.list_filtered(None, None) } diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index 3020c647eeda..6c17cc3c117a 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -2933,6 +2933,34 @@ mod tests { assert!(!core.indexer_registry.has_worker(1)); } + #[tokio::test] + async fn reupsert_recreates_a_listener_lost_to_a_cancelled_update() { + let core = SelectionCore::try_new_local( + test_config(true), + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ) + .expect("valid test config"); + core.upsert_worker(worker_with_kv_events(1)) + .await + .expect("worker upsert"); + // An endpoint update cancelled after deregistration leaves the catalog + // record schedulable with no listener behind it. + core.indexer_registry + .deregister_dp_rank(1, 0, "model", "default") + .await + .expect("deregister"); + assert!(!core.indexer_registry.has_listener(1, 0)); + + let record = core + .upsert_worker(worker_with_kv_events(1)) + .await + .expect("worker re-upsert"); + assert_eq!(record.lifecycle, WorkerLifecycle::Schedulable, "{record:?}"); + assert!(core.indexer_registry.has_listener(1, 0)); + } + #[tokio::test] async fn shutdown_cancels_listeners() { let parent = CancellationToken::new(); diff --git a/lib/kv-router/src/services/selection/ingress.rs b/lib/kv-router/src/services/selection/ingress.rs index a54b0e945212..9575847c7f85 100644 --- a/lib/kv-router/src/services/selection/ingress.rs +++ b/lib/kv-router/src/services/selection/ingress.rs @@ -117,6 +117,11 @@ impl KvEventIngress for ZmqDirectIngress { if let Some(previous) = previous { let old_endpoints = previous.listener_endpoints(); for rank in previous.dp_ranks() { + // A listener can be missing for a schedulable record when an + // earlier update was cancelled mid-deregistration; re-register it. + if !registry.has_listener(record.worker_id, rank) { + continue; + } if ranks.contains(&rank) && old_endpoints.get(&rank) == endpoints.get(&rank) && previous.replay_endpoint == record.replay_endpoint From b4a81a5d4080d4e998e8e403e997e28cbdd07f04 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Thu, 10 Sep 2026 09:51:37 -0700 Subject: [PATCH 050/102] fix(router): purge index blocks for a rank whose listener is already gone A cancelled worker update can remove a rank's KV listener before the index purge runs. Reconciliation now treats a rank without a listener as not retained, purges its indexed blocks, and re-registers it when it is still wanted; the concurrent-failover test and bind_session doc are trimmed. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Ishan Dhanani --- .../src/services/indexer/registry.rs | 22 ++++++++ .../src/services/selection/core/mod.rs | 51 ++++++++++++------- .../src/services/selection/ingress.rs | 17 ++++--- 3 files changed, 66 insertions(+), 24 deletions(-) diff --git a/lib/kv-router/src/services/indexer/registry.rs b/lib/kv-router/src/services/indexer/registry.rs index 5c99a5e58713..754144cc3c00 100644 --- a/lib/kv-router/src/services/indexer/registry.rs +++ b/lib/kv-router/src/services/indexer/registry.rs @@ -696,6 +696,19 @@ impl WorkerRegistry { .is_some_and(|entry| entry.listeners.contains_key(&dp_rank)) } + /// Drop what a rank's listener fed into the index after the listener itself + /// is gone (its deregistration was cancelled before reaching this step). + pub async fn remove_dp_rank_blocks( + &self, + worker_id: WorkerId, + dp_rank: u32, + key: &RoutingPartitionId, + ) { + if let Some(ie) = self.indexers.get(key) { + ie.indexer.remove_worker_dp_rank(worker_id, dp_rank).await; + } + } + pub fn list(&self) -> Vec { self.list_filtered(None, None) } @@ -813,6 +826,15 @@ impl WorkerRegistry { .collect() } + /// The state a `deregister_dp_rank` cancelled after removing the listener + /// leaves behind: no listener, index entries untouched. + #[cfg(test)] + pub(crate) fn forget_listener(&self, instance_id: WorkerId, dp_rank: u32) { + if let Some(mut entry) = self.workers.get_mut(&instance_id) { + entry.listeners.remove(&dp_rank); + } + } + #[cfg(test)] pub(crate) fn listener_cancelled(&self, instance_id: WorkerId, dp_rank: u32) -> Option { self.workers.get(&instance_id).and_then(|entry| { diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index 6c17cc3c117a..05218ff46609 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -1511,8 +1511,7 @@ impl SelectionCore { /// Bind (or confirm) `session_id` to the worker a booking landed on and /// keep the lease with the booking. A bound session whose worker was not - /// selected (it left) is invalidated and rebound; when a concurrent request - /// already rebound it, this booking holds that binding's lease instead. + /// selected (it left) is invalidated and rebound. fn bind_session( table: &SessionAffinity, hold: Acquired, @@ -2814,16 +2813,13 @@ mod tests { #[tokio::test] async fn concurrent_failover_keeps_the_winning_binding_lease() { let table = SessionAffinity::new(Duration::from_secs(60)).expect("affinity table"); - let departed = WorkerWithDpRank::new(1, 0); let Acquired::Initialize(init) = table.acquire("s", None).await.expect("acquire") else { panic!("fresh session must initialize"); }; drop(init.commit(SessionTarget::new(1, Some(0))).expect("commit")); let first = table.acquire("s", None).await.expect("first acquire"); let second = table.acquire("s", None).await.expect("second acquire"); - assert!( - matches!(first, Acquired::Bound { target, .. } if target.worker_id == departed.worker_id) - ); + assert!(matches!(first, Acquired::Bound { target, .. } if target.worker_id == 1)); // Worker 1 left; both bookings landed elsewhere. The first rebinds the // session, the second must hold that binding rather than end up leaseless. @@ -2832,11 +2828,6 @@ mod tests { let follower = SelectionCore::bind_session(&table, second, "s", WorkerWithDpRank::new(3, 0)) .expect("second failover holds the rebound session"); - let bound = table - .query_target("s", None) - .expect("query") - .expect("bound"); - assert_eq!(bound.worker_id, 2); drop(winner); drop(follower); assert_eq!( @@ -2935,6 +2926,8 @@ mod tests { #[tokio::test] async fn reupsert_recreates_a_listener_lost_to_a_cancelled_update() { + use crate::indexer::KvIndexerInterface; + let core = SelectionCore::try_new_local( test_config(true), 1, @@ -2945,13 +2938,21 @@ mod tests { core.upsert_worker(worker_with_kv_events(1)) .await .expect("worker upsert"); - // An endpoint update cancelled after deregistration leaves the catalog - // record schedulable with no listener behind it. - core.indexer_registry - .deregister_dp_rank(1, 0, "model", "default") + let entry = core + .entry(&RoutingPartitionId::new("model", "default")) + .expect("entry"); + let hashes = [11u64, 12]; + entry + .indexer + .apply_event_routed(store_event(1, 0, 1, &[], &hashes, StorageTier::Device)) .await - .expect("deregister"); - assert!(!core.indexer_registry.has_listener(1, 0)); + .unwrap(); + if let Indexer::Single { primary, .. } = &entry.indexer { + let _ = primary.flush().await; + } + // An endpoint update cancelled right after removing the listener leaves + // the catalog record schedulable, the listener gone, and its blocks indexed. + core.indexer_registry.forget_listener(1, 0); let record = core .upsert_worker(worker_with_kv_events(1)) @@ -2959,6 +2960,22 @@ mod tests { .expect("worker re-upsert"); assert_eq!(record.lifecycle, WorkerLifecycle::Schedulable, "{record:?}"); assert!(core.indexer_registry.has_listener(1, 0)); + // The purge is queued to the indexer thread; wait for it to land. + tokio::time::timeout(Duration::from_secs(2), async { + loop { + let overlap = entry + .indexer + .find_matches(hashes.iter().copied().map(LocalBlockHash).collect()) + .await + .expect("find matches"); + if overlap.scores.is_empty() { + return; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("stale blocks survived the re-upsert"); } #[tokio::test] diff --git a/lib/kv-router/src/services/selection/ingress.rs b/lib/kv-router/src/services/selection/ingress.rs index 9575847c7f85..73bc3c9e3360 100644 --- a/lib/kv-router/src/services/selection/ingress.rs +++ b/lib/kv-router/src/services/selection/ingress.rs @@ -117,17 +117,16 @@ impl KvEventIngress for ZmqDirectIngress { if let Some(previous) = previous { let old_endpoints = previous.listener_endpoints(); for rank in previous.dp_ranks() { - // A listener can be missing for a schedulable record when an - // earlier update was cancelled mid-deregistration; re-register it. - if !registry.has_listener(record.worker_id, rank) { - continue; - } - if ranks.contains(&rank) + // A schedulable record can lack its listener when an earlier + // update was cancelled mid-deregistration. + let listener_present = registry.has_listener(record.worker_id, rank); + if listener_present + && ranks.contains(&rank) && old_endpoints.get(&rank) == endpoints.get(&rank) && previous.replay_endpoint == record.replay_endpoint { retained.insert(rank); - } else { + } else if listener_present { registry .deregister_dp_rank( record.worker_id, @@ -137,6 +136,10 @@ impl KvEventIngress for ZmqDirectIngress { ) .await .map_err(|error| SelectionError::Internal(error.to_string()))?; + } else { + registry + .remove_dp_rank_blocks(record.worker_id, rank, &previous.key()) + .await; } } } From e2e4e618d843238cb3b2e831409a199e7c737852 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Thu, 10 Sep 2026 11:26:39 -0700 Subject: [PATCH 051/102] feat(router): guard core bookings with the lifecycle lease and claim selection ids Core bookings keep the queue's RequestLifecycleLease armed until the reservation is installed, so a caller dropping select_and_reserve or a two-phase reservation anywhere after the scheduler booked frees the booking instead of leaking it. The selection id is claimed in the reservation index before booking: a live or in-flight id is rejected with 409 up front instead of failing as a scheduler duplicate after queueing, and lifecycle calls for an in-flight id return 404 rather than mutating the booking being made. Lifecycle calls now resolve through the index to one partition and one scheduler booking descriptor; the id-only scan over every partition is gone. Bookings mirrored from replica peers are indexed by an observer that wraps the host's, so they stay reachable and the index stays exact when peers free or re-admit them. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Ishan Dhanani --- lib/kv-router/src/scheduling/local.rs | 109 ++- lib/kv-router/src/scheduling/queue.rs | 34 +- lib/kv-router/src/sequences/multi_worker.rs | 44 +- .../src/services/selection/core/mod.rs | 647 ++++++++++++++---- lib/llm/src/kv_router.rs | 1 + lib/llm/src/kv_router/embedded.rs | 1 + 6 files changed, 699 insertions(+), 137 deletions(-) diff --git a/lib/kv-router/src/scheduling/local.rs b/lib/kv-router/src/scheduling/local.rs index f8fd08151528..d86c8b149dfc 100644 --- a/lib/kv-router/src/scheduling/local.rs +++ b/lib/kv-router/src/scheduling/local.rs @@ -17,7 +17,8 @@ use super::overlap_refresh::{NoopOverlapScoresRefresh, OverlapScoresRefresh}; use super::policy_config::PolicyProfile; use super::prefill_load::PrefillLoadEstimator; use super::queue::{ - ClassQueueStats, SchedulerBookingCleanup, SchedulerBookingDescriptor, SchedulerQueue, + ClassQueueStats, RequestLifecycleLease, SchedulerBookingCleanup, SchedulerBookingDescriptor, + SchedulerQueue, }; use super::selector::{DefaultWorkerSelector, WorkerSelector}; use super::types::{ @@ -30,8 +31,8 @@ use crate::protocols::RoutingConstraints; use crate::protocols::{LocalBlockHash, WorkerConfigLike, WorkerId, WorkerWithDpRank}; use crate::sequences::topology::WorkerDpRange; use crate::sequences::{ - ActiveSequencesMultiWorker, PrefillTokenDeltas, SequenceError, SequencePublisher, - SequenceRequest, + ActiveSequencesMultiWorker, LifecycleMutationOutcome, PrefillTokenDeltas, SequenceError, + SequencePublisher, SequenceRequest, }; use dynamo_tokens::SequenceHash; @@ -290,6 +291,27 @@ where &self, request: ScheduleRequest, ) -> Result { + let (admitted, lease) = self.schedule_request_with_lease(request).await?; + if let Some(lease) = lease { + let _ = lease.commit(); + } + Ok(admitted) + } + + /// Schedule a request and keep its lifecycle lease armed: dropping the + /// lease frees the booking, `commit` hands it to a longer-lived owner. + /// The lease is `None` unless the mode is `TrackedWithLifecycle`. + #[doc(hidden)] + pub async fn schedule_request_with_lease( + &self, + request: ScheduleRequest, + ) -> Result< + ( + AdmittedSchedulingResponse, + Option>, + ), + KvSchedulerError, + > { let tracked = request.mode.is_tracked(); let (resp_tx, resp_rx) = tokio::sync::oneshot::channel(); let (attempt_tx, attempt_rx) = tokio::sync::oneshot::channel(); @@ -298,7 +320,7 @@ where .new_request_lifecycle_lease(request.mode.lifecycle_request_id()); let (request, block_hashes) = self.make_scheduling_request(request, Some(resp_tx)); - let mut lifecycle_lease = self + let lifecycle_lease = self .queue .enqueue_admitted_with_block_hashes_and_lease( request, @@ -320,10 +342,10 @@ where } else { AdmissionAttempt::Untracked }; - if let Some(lease) = lifecycle_lease.as_mut() { - lease.disarm(); - } - Ok(AdmittedSchedulingResponse { response, attempt }) + Ok(( + AdmittedSchedulingResponse { response, attempt }, + lifecycle_lease, + )) } /// Select a worker from current scheduler state without queue admission or booking. @@ -505,11 +527,23 @@ where /// Book a request only when its worker is already registered, so a request /// racing worker removal cannot lazily recreate the removed worker/rank. - pub async fn add_request_if_registered( + /// The returned lease guards the booking: dropping it frees the booking, + /// `commit` hands it over. + #[doc(hidden)] + pub fn add_request_if_registered_guarded( &self, req: SequenceRequest, - ) -> Result<(), SequenceError> { - self.slots.add_request_if_registered(req, Instant::now()) + ) -> Result, SequenceError> { + let request_id = req.request_id.clone(); + let worker = req.worker; + let attempt_id = self + .slots + .add_request_if_registered_admitted(req, Instant::now())?; + Ok(self.queue.lease_for_booking(SchedulerBookingDescriptor { + request_id, + worker, + attempt_id, + })) } pub async fn mark_prefill_completed(&self, request_id: &str) -> Result<(), SequenceError> { @@ -571,21 +605,56 @@ where self.slots.request_worker(&request_id.to_string()).is_some() } + #[doc(hidden)] + pub fn has_booking(&self, booking: &SchedulerBookingDescriptor) -> bool { + self.slots.has_booking(booking) + } + + /// Release a booking only if `booking` still describes it; `NoChange` when + /// the id was freed or rebooked since the descriptor was taken. + #[doc(hidden)] + pub async fn free_if_booking( + &self, + booking: &SchedulerBookingDescriptor, + ) -> Result { + let outcome = self.slots.free_if_booking( + &booking.request_id, + booking.worker, + booking.attempt_id, + Instant::now(), + )?; + if outcome.is_applied() { + self.queue.update_worker(booking.worker).await; + } + Ok(outcome) + } + #[doc(hidden)] pub fn booking_cleanup(&self) -> SchedulerBookingCleanup { self.queue.booking_cleanup() } + /// `NoChange` when the booking no longer matches or its prefill was + /// already marked complete. #[doc(hidden)] pub async fn mark_prefill_completed_if_booking( &self, booking: &SchedulerBookingDescriptor, - ) -> Result<(), KvSchedulerError> { + ) -> Result { self.queue .mark_prefill_completed_if_booking(booking.clone()) .await } + /// Republish the ordered prefill-completion event while `booking` is live. + #[doc(hidden)] + pub fn publish_prefill_completed_if_booking( + &self, + booking: &SchedulerBookingDescriptor, + ) -> bool { + self.slots.publish_prefill_completed_if_booking(booking) + } + pub fn pending_count(&self) -> usize { self.queue.pending_count() } @@ -630,6 +699,22 @@ where .await } + /// `add_output_block_if_booking` applied inline, like `add_output_block`, + /// for callers that cannot await. + #[doc(hidden)] + pub fn add_output_block_if_booking_sync( + &self, + booking: &SchedulerBookingDescriptor, + decay_fraction: Option, + ) -> Result { + self.slots.add_output_block_if_booking( + &booking.request_id, + booking.worker, + booking.attempt_id, + decay_fraction, + ) + } + #[doc(hidden)] pub async fn enqueue_output_block_if_booking( &self, diff --git a/lib/kv-router/src/scheduling/queue.rs b/lib/kv-router/src/scheduling/queue.rs index 6d64eec2848a..c5f317db20d0 100644 --- a/lib/kv-router/src/scheduling/queue.rs +++ b/lib/kv-router/src/scheduling/queue.rs @@ -363,7 +363,8 @@ impl SchedulerBookingCleanup { } /// Single-owner cleanup lease for one scheduler-tracked request. -pub(crate) struct RequestLifecycleLease { +#[doc(hidden)] +pub struct RequestLifecycleLease { cleanup: Arc, actor_tx: mpsc::Sender, transfer: Option>, @@ -384,6 +385,20 @@ impl RequestLifecycleLease { transfer.disarm(); } } + + /// Hand the booking to its long-term owner: the lease stops guarding it. + #[must_use] + pub fn commit(mut self) -> Option { + let booking = match self.transfer.as_ref().map(|transfer| transfer.state.lock()) { + Some(state) => match &*state { + AdmissionLifecycleState::Booking(booking) => Some(booking.clone()), + _ => None, + }, + None => None, + }; + self.disarm(); + booking + } } impl Drop for RequestLifecycleLease { @@ -723,6 +738,20 @@ impl< } } + /// An armed lease for a booking made outside the admission actor. + pub(crate) fn lease_for_booking( + &self, + booking: SchedulerBookingDescriptor, + ) -> Box { + let transfer = AdmissionLifecycleTransfer::new(booking.request_id.clone()); + transfer.arm_booking(booking); + Box::new(RequestLifecycleLease { + cleanup: Arc::clone(&self.cleanup), + actor_tx: self.admission_tx.clone(), + transfer: Some(Arc::new(transfer)), + }) + } + /// Select a worker from current scheduler state without entering admission. /// /// This is for advisory policy probes that must not wait in the router @@ -758,7 +787,7 @@ impl< pub(crate) async fn mark_prefill_completed_if_booking( &self, booking: SchedulerBookingDescriptor, - ) -> Result<(), KvSchedulerError> { + ) -> Result { let (ack_tx, ack_rx) = oneshot::channel(); self.admission_tx .send(AdmissionCommand::MarkPrefillCompleted { booking, ack_tx }) @@ -767,7 +796,6 @@ impl< ack_rx .await .map_err(|_| KvSchedulerError::SubscriberShutdown)? - .map(|_| ()) .map_err(|error| KvSchedulerError::BookingFailed(error.to_string())) } diff --git a/lib/kv-router/src/sequences/multi_worker.rs b/lib/kv-router/src/sequences/multi_worker.rs index c115b82d1412..8032072108d6 100644 --- a/lib/kv-router/src/sequences/multi_worker.rs +++ b/lib/kv-router/src/sequences/multi_worker.rs @@ -787,7 +787,16 @@ impl ActiveSequencesMultiWorker

{ req: SequenceRequest, decay_now: Instant, ) -> Result<(), SequenceError> { - self.add_request_impl(req, decay_now, false).map(|_| ()) + self.add_request_if_registered_admitted(req, decay_now) + .map(|_| ()) + } + + pub fn add_request_if_registered_admitted( + &self, + req: SequenceRequest, + decay_now: Instant, + ) -> Result { + self.add_request_impl(req, decay_now, false) } fn add_request_impl( @@ -819,6 +828,39 @@ impl ActiveSequencesMultiWorker

{ self.request_index.worker_for(request_id) } + pub(crate) fn has_booking(&self, booking: &SchedulerBookingDescriptor) -> bool { + self.request_index.booking_for(&booking.request_id) + == Some(RequestBooking { + worker: booking.worker, + attempt_id: booking.attempt_id, + }) + } + + /// Republish the ordered completion event for `booking` while it is still + /// the live booking; serialized with free/rebook through the worker's + /// sequence lock. + pub(crate) fn publish_prefill_completed_if_booking( + &self, + booking: &SchedulerBookingDescriptor, + ) -> bool { + let table = self.workers.read(); + let Some(&idx) = table.index.get(&booking.worker) else { + return false; + }; + let _seq = table.slots[idx].sequences.write(); + if !self.has_booking(booking) { + return false; + } + self.enqueue_publish_event(ActiveSequenceEvent { + request_id: booking.request_id.clone(), + worker: booking.worker, + data: ActiveSequenceEventData::MarkPrefillCompleted, + router_id: self.router_id, + lora_name: self.request_index.lora_for(&booking.request_id), + }); + true + } + /// Free all blocks associated with a request. /// /// Note: This operation is idempotent. Calling it multiple times for the same request diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index 05218ff46609..a6f71f2c5922 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -18,6 +18,8 @@ use crate::indexer::{ use crate::kv_hints::{ KvHint, KvHintAction, KvSourceLocationsPayload, KvTransferCandidateSource, KvTransferCandidates, }; +#[cfg(test)] +use crate::protocols::ActiveSequenceEventData; use crate::protocols::{ ActiveSequenceEvent, LocalBlockHash, PrefillLoadHint, RoutingConstraints, SharedCacheHits, WorkerAffinityTarget, WorkerConfigLike, WorkerId, WorkerWithDpRank, @@ -26,16 +28,17 @@ use crate::scheduling::config::RouterConfigOverride; use crate::scheduling::queue::SchedulerBookingDescriptor; use crate::scheduling::selector::WorkerSelectionPolicy; use crate::scheduling::{ - AdmissionAttempt, KvSchedulerError, LocalScheduler, LoraWorkerFilter, OverlapAnalysis, - OverlapSignals, OverloadedWorkerProvider, PotentialLoad, PrefillLoadEstimator, ScheduleMode, - ScheduleRequest, SessionContext, TieredOverlapRefresher, WorkerAvailabilityProvider, - effective_prefill_tokens, narrow_allowed_worker_ids_by_lora, - prefill_load_hint_from_effective_tokens, + KvSchedulerError, LocalScheduler, LoraWorkerFilter, OverlapAnalysis, OverlapSignals, + OverloadedWorkerProvider, PotentialLoad, PrefillLoadEstimator, ScheduleMode, ScheduleRequest, + SessionContext, TieredOverlapRefresher, WorkerAvailabilityProvider, effective_prefill_tokens, + narrow_allowed_worker_ids_by_lora, prefill_load_hint_from_effective_tokens, }; use crate::sequences::{ - ActiveSequencesMultiWorker, ReplicaRequestLeaseObserver, ReplicaWorkerPolicy, SequenceError, - SequenceRequest, SequenceTrackerOptions, active_request_expiry_duration, + ActiveSequencesMultiWorker, LifecycleMutationOutcome, ReplicaRequestLeaseObserver, + ReplicaWorkerPolicy, SequenceRequest, SequenceTrackerOptions, active_request_expiry_duration, }; +#[cfg(test)] +use crate::services::common::replica_sync::HostReplicaChannels; use crate::services::common::replica_sync::{ HostReplicaSyncFactory, ReplicaSyncConfig, SchedulerLoadSink, ScopedReplicaEvent, ScopedSequencePublisher, setup_scoped_replica_sync, @@ -304,21 +307,120 @@ pub struct SelectionServiceConfig { type SelectionEntries = RwLock>>>>; -/// `selection_id` -> partition that holds its booking. +/// `selection_id` -> the booking it holds. /// /// Lifecycle calls (`prefill_complete`, `free`, `add_output_block`) arrive with -/// only a selection id, so without this index every call scans every partition -/// scheduler. Selection ids are caller-controlled strings, so this stays on the -/// standard hasher. Entries are removed on `free`, on a `RequestNotFound` from -/// the indexed scheduler, and by a periodic sweep that drops ids whose booking -/// expired underneath them. +/// only a selection id; the index resolves them to one partition and one +/// scheduler booking, so they never touch a booking made by a later request +/// that reused the id. An id is claimed here before its booking is made and +/// installed once the booking is final; a claim (`booking == None`) rejects a +/// concurrent booking of the same id and is invisible to lifecycle calls. +/// Bookings mirrored from replica peers are indexed by +/// [`ReservationIndexObserver`]. type ReservationIndex = RwLock>; struct Reservation { partition: RoutingPartitionId, + booking: Option, _affinity_lease: Option, } +/// Exclusive ownership of a selection id while its booking is in flight. +/// Dropping the claim without `install` releases the id. +struct ReservationClaim<'a> { + index: &'a ReservationIndex, + selection_id: String, + armed: bool, +} + +impl ReservationClaim<'_> { + fn install(mut self, reservation: Reservation) { + self.armed = false; + self.index + .write() + .insert(std::mem::take(&mut self.selection_id), reservation); + } +} + +impl Drop for ReservationClaim<'_> { + fn drop(&mut self) { + if !self.armed { + return; + } + let mut index = self.index.write(); + if index + .get(&self.selection_id) + .is_some_and(|reservation| reservation.booking.is_none()) + { + index.remove(&self.selection_id); + } + } +} + +/// Keeps the reservation index exact for bookings replicated from peers, which +/// never pass through this core's booking paths, and forwards to the host's +/// own observer. +struct ReservationIndexObserver { + index: Arc, + partition: RoutingPartitionId, + host: Option>, +} + +impl ReplicaRequestLeaseObserver for ReservationIndexObserver { + fn admitted(&self, booking: SchedulerBookingDescriptor) { + { + let mut index = self.index.write(); + // A peer only admits an id this partition's scheduler does not hold, + // so an existing row for it is a stale booking (expired, not yet + // swept) or a claim whose local booking will now fail; the mirror + // replaces both. A row from another partition is left alone. + let own_row = index + .get(&booking.request_id) + .is_none_or(|reservation| reservation.partition == self.partition); + if own_row { + index.insert( + booking.request_id.clone(), + Reservation { + partition: self.partition.clone(), + booking: Some(booking.clone()), + _affinity_lease: None, + }, + ); + } + } + if let Some(host) = &self.host { + host.admitted(booking); + } + } + + fn progressed(&self, booking: &SchedulerBookingDescriptor) { + if let Some(host) = &self.host { + host.progressed(booking); + } + } + + fn completed(&self, booking: &SchedulerBookingDescriptor) { + forget_reservation_if(&self.index, &self.partition, booking); + if let Some(host) = &self.host { + host.completed(booking); + } + } +} + +/// Remove the index entry for `booking` only if it still describes it. +fn forget_reservation_if( + index: &ReservationIndex, + partition: &RoutingPartitionId, + booking: &SchedulerBookingDescriptor, +) { + let mut index = index.write(); + if index.get(&booking.request_id).is_some_and(|reservation| { + reservation.partition == *partition && reservation.booking.as_ref() == Some(booking) + }) { + index.remove(&booking.request_id); + } +} + pub struct SelectionCore { catalog: WorkerCatalog, /// Serializes catalog commits and the corresponding ingress changes. Never held by selection. @@ -750,9 +852,11 @@ impl SelectionCore { .then(active_request_expiry_duration), }, )); - if let Some(observer) = &self.host.replication.request_leases { - slots.set_replica_request_lease_observer(Arc::clone(observer)); - } + slots.set_replica_request_lease_observer(Arc::new(ReservationIndexObserver { + index: Arc::clone(&self.reservation_index), + partition: key.clone(), + host: self.host.replication.request_leases.clone(), + })); let replica_tx = scoped_replica_sync.channel.map(|(replica_tx, subscriber)| { slots.start_replica_sync(subscriber, self.cancel_token.child_token()); replica_tx @@ -991,6 +1095,10 @@ impl SelectionCore { self.ensure_running()?; let entry = self.ready_entry(&key)?; + let claim = match selection_id.as_deref() { + Some(selection_id) if book => Some(self.claim_reservation(selection_id, &key)?), + _ => None, + }; // Session stickiness: a bound session steers selection (exclusive for // the default selector); a new session is bound to the worker chosen. @@ -1126,7 +1234,9 @@ impl SelectionCore { routing_constraints, shared_cache_hits, }; - let (response, advisory_load, attempt) = tokio::select! { + // `lease` guards the booking until it is installed below: any early + // return or drop before then frees it. + let (response, advisory_load, lease) = tokio::select! { biased; _ = self.cancel_token.cancelled() => { return Err(SelectionError::Scheduler(KvSchedulerError::SubscriberShutdown)); @@ -1138,18 +1248,14 @@ impl SelectionCore { .select_without_admission(schedule_request) .await .map(|advisory| { - ( - advisory.response, - Some(advisory.selected_worker_load), - AdmissionAttempt::Untracked, - ) + (advisory.response, Some(advisory.selected_worker_load), None) }) } else { entry .scheduler - .schedule_request_admitted(schedule_request) + .schedule_request_with_lease(schedule_request) .await - .map(|admitted| (admitted.response, None, admitted.attempt)) + .map(|(admitted, lease)| (admitted.response, None, lease)) } } => result?, }; @@ -1157,20 +1263,6 @@ impl SelectionCore { .catalog .schedulable_endpoint(response.best_worker.worker_id, &key) else { - // The worker drained while this request was queued; release the - // booking the actor just made for it. - if let (AdmissionAttempt::Tracked(attempt_id), Some(request_id)) = - (attempt, selection_id.as_deref()) - { - entry - .scheduler - .booking_cleanup() - .enqueue(SchedulerBookingDescriptor { - request_id: request_id.to_string(), - worker: response.best_worker, - attempt_id, - }); - } return Err(SelectionError::Internal(format!( "selected worker {} is no longer schedulable", response.best_worker.worker_id @@ -1222,18 +1314,30 @@ impl SelectionCore { .map(|threshold| load.prefill_load_exceeds(threshold)), }); - if book && let Some(selection_id) = selection_id.as_deref() { - let lease = match (affinity_hold, session_id.as_deref(), table) { - (Some(hold), Some(session_id), Some(table)) => { - Self::bind_session(table, hold, session_id, response.best_worker) - } - _ => None, + if let Some(claim) = claim { + let Some(lease) = lease else { + return Err(SelectionError::Internal( + "booked selection has no lifecycle lease".to_string(), + )); }; - self.record_reservation(selection_id, &key, lease); if let Some(hashes) = routing_hashes.clone() { self.record_routing_decision(&entry, response.best_worker, hashes) .await; } + // Nothing awaits from here to `install`: the binding, the booking + // handover and the index entry land together. + let booking = lease.commit().ok_or_else(missing_booking)?; + let affinity_lease = match (affinity_hold, session_id.as_deref(), table) { + (Some(hold), Some(session_id), Some(table)) => { + Self::bind_session(table, hold, session_id, response.best_worker) + } + _ => None, + }; + claim.install(Reservation { + partition: key.clone(), + booking: Some(booking), + _affinity_lease: affinity_lease, + }); } if let Some((cache_id, sequence_hashes, lora_name, track_prefill_tokens)) = cached_inputs { @@ -1461,11 +1565,13 @@ impl SelectionCore { routing_hashes, } = booking; + let claim = self.claim_reservation(&selection_id, &key)?; // Strict booking: never lazily recreate a worker/rank removed since the - // reservation was resolved. - entry + // reservation was resolved. The lease frees the booking if this future + // is dropped before `install`. + let lease = entry .scheduler - .add_request_if_registered(SequenceRequest { + .add_request_if_registered_guarded(SequenceRequest { request_id: selection_id.clone(), token_sequence: Some(sequence_hashes), track_prefill_tokens, @@ -1473,12 +1579,15 @@ impl SelectionCore { prefill_load_hint, worker, lora_name, - }) - .await?; - self.record_reservation(&selection_id, &key, None); + })?; if let Some(hashes) = routing_hashes { self.record_routing_decision(&entry, worker, hashes).await; } + claim.install(Reservation { + partition: key.clone(), + booking: Some(lease.commit().ok_or_else(missing_booking)?), + _affinity_lease: None, + }); Ok(ReservationResponse { selection_id, @@ -1490,23 +1599,31 @@ impl SelectionCore { }) } - fn record_reservation( + /// Take `selection_id` for a booking about to be made in `key`. + fn claim_reservation( &self, selection_id: &str, key: &RoutingPartitionId, - lease: Option, - ) { - self.reservation_index.write().insert( + ) -> Result, SelectionError> { + let mut index = self.reservation_index.write(); + if index.contains_key(selection_id) { + return Err(SelectionError::Conflict(format!( + "selection {selection_id} is already reserved or being reserved" + ))); + } + index.insert( selection_id.to_string(), Reservation { partition: key.clone(), - _affinity_lease: lease, + booking: None, + _affinity_lease: None, }, ); - } - - fn forget_reservation(&self, selection_id: &str) { - self.reservation_index.write().remove(selection_id); + Ok(ReservationClaim { + index: &self.reservation_index, + selection_id: selection_id.to_string(), + armed: true, + }) } /// Bind (or confirm) `session_id` to the worker a booking landed on and @@ -1605,58 +1722,63 @@ impl SelectionCore { } } - /// Entries to try for a lifecycle call on `selection_id`: the indexed - /// partition first, then every other initialized partition. The fallback - /// covers bookings mirrored from replica peers, which never pass through - /// this core's booking paths. - fn lifecycle_entries(&self, selection_id: &str) -> Vec> { + /// The partition and booking a lifecycle call on `selection_id` may touch. + /// `None` for unknown ids and for ids whose booking is still in flight. + fn indexed_booking( + &self, + selection_id: &str, + ) -> Option<(Arc, SchedulerBookingDescriptor)> { // Release the index guard before taking `entries`: the sweep takes // `entries` then `reservation_index`, so nesting here would invert the // lock order. - let indexed_partition = self - .reservation_index - .read() - .get(selection_id) - .map(|reservation| reservation.partition.clone()); - let indexed = indexed_partition.and_then(|key| self.entry(&key)); - let mut entries = self.initialized_entries(); - if let Some(indexed) = indexed { - entries.retain(|entry| !Arc::ptr_eq(entry, &indexed)); - entries.insert(0, indexed); - } - entries + let (partition, booking) = { + let index = self.reservation_index.read(); + let reservation = index.get(selection_id)?; + (reservation.partition.clone(), reservation.booking.clone()?) + }; + Some((self.entry(&partition)?, booking)) + } + + fn reservation_not_found(selection_id: &str) -> SelectionError { + SelectionError::NotFound(format!("reservation {selection_id} not found")) } pub async fn prefill_complete(&self, selection_id: &str) -> Result<(), SelectionError> { - for entry in self.lifecycle_entries(selection_id) { - match entry.scheduler.mark_prefill_completed(selection_id).await { - Ok(()) => return Ok(()), - Err(SequenceError::RequestNotFound { .. }) => continue, - Err(error) => return Err(error.into()), + let Some((entry, booking)) = self.indexed_booking(selection_id) else { + return Err(Self::reservation_not_found(selection_id)); + }; + match entry + .scheduler + .mark_prefill_completed_if_booking(&booking) + .await? + { + LifecycleMutationOutcome::Applied => Ok(()), + // Already marked: still publish the ordered completion so peers + // that missed the first event converge. + LifecycleMutationOutcome::NoChange + if entry + .scheduler + .publish_prefill_completed_if_booking(&booking) => + { + Ok(()) + } + LifecycleMutationOutcome::NoChange => { + forget_reservation_if(&self.reservation_index, &entry.key, &booking); + Err(Self::reservation_not_found(selection_id)) } } - self.forget_reservation(selection_id); - Err(SelectionError::NotFound(format!( - "reservation {selection_id} not found" - ))) } pub async fn free_reservation(&self, selection_id: &str) -> Result<(), SelectionError> { - let result = async { - for entry in self.lifecycle_entries(selection_id) { - match entry.scheduler.free(selection_id).await { - Ok(()) => return Ok(()), - Err(SequenceError::RequestNotFound { .. }) => continue, - Err(error) => return Err(error.into()), - } - } - Err(SelectionError::NotFound(format!( - "reservation {selection_id} not found" - ))) + let Some((entry, booking)) = self.indexed_booking(selection_id) else { + return Err(Self::reservation_not_found(selection_id)); + }; + let outcome = entry.scheduler.free_if_booking(&booking).await?; + forget_reservation_if(&self.reservation_index, &entry.key, &booking); + match outcome { + LifecycleMutationOutcome::Applied => Ok(()), + LifecycleMutationOutcome::NoChange => Err(Self::reservation_not_found(selection_id)), } - .await; - self.forget_reservation(selection_id); - result } pub fn add_output_block( @@ -1672,20 +1794,19 @@ impl SelectionCore { )); } - for entry in self.lifecycle_entries(selection_id) { - match entry - .scheduler - .add_output_block(selection_id, decay_fraction) - { - Ok(()) => return Ok(()), - Err(SequenceError::RequestNotFound { .. }) => continue, - Err(error) => return Err(error.into()), + let Some((entry, booking)) = self.indexed_booking(selection_id) else { + return Err(Self::reservation_not_found(selection_id)); + }; + match entry + .scheduler + .add_output_block_if_booking_sync(&booking, decay_fraction)? + { + LifecycleMutationOutcome::Applied => Ok(()), + LifecycleMutationOutcome::NoChange => { + forget_reservation_if(&self.reservation_index, &entry.key, &booking); + Err(Self::reservation_not_found(selection_id)) } } - self.forget_reservation(selection_id); - Err(SelectionError::NotFound(format!( - "reservation {selection_id} not found" - ))) } pub fn loads( @@ -1868,20 +1989,28 @@ impl SelectionCore { /// Drop index entries whose booking no longer exists in its partition /// scheduler (expired by the periodic force-expiry, or freed through a path -/// that bypassed this core). Returns the number of entries removed. +/// that bypassed this core). Claims are left for their owner. Returns the +/// number of entries removed. fn sweep_reservation_index(entries: &SelectionEntries, index: &ReservationIndex) -> usize { let entries = entries.read(); let mut index = index.write(); let before = index.len(); - index.retain(|id, reservation| { + index.retain(|_, reservation| { + let Some(booking) = &reservation.booking else { + return true; + }; entries .get(&reservation.partition) .and_then(|cell| cell.get()) - .is_some_and(|entry| entry.scheduler.has_request(id)) + .is_some_and(|entry| entry.scheduler.has_booking(booking)) }); before - index.len() } +fn missing_booking() -> SelectionError { + SelectionError::Internal("booking lease holds no booking".to_string()) +} + fn spawn_reservation_index_sweep( entries: Arc, index: Arc, @@ -3230,6 +3359,13 @@ mod tests { false }) .await; + assert!( + core.reservation_index + .read() + .get("dropped") + .is_some_and(|reservation| reservation.booking.is_none()), + "the id is claimed while its booking is in flight" + ); drop(selection); wait_until("booking release", || { @@ -3240,7 +3376,274 @@ mod tests { } #[tokio::test] - async fn lifecycle_operations_find_reservation_in_later_entry() { + async fn same_selection_id_in_two_partitions_is_a_conflict() { + let core = SelectionCore::try_new_local( + test_config(false), + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ) + .expect("valid test config"); + for (worker_id, routing_group) in [(1, "group-a"), (2, "group-b")] { + let mut request = worker(worker_id); + request.routing_group = routing_group.to_string(); + core.upsert_worker(request).await.expect("worker upsert"); + } + let mut first = reserve_request("shared"); + first.routing_group = "group-a".to_string(); + core.select_and_reserve(first).await.expect("first booking"); + + let mut second = reserve_request("shared"); + second.routing_group = "group-b".to_string(); + let err = core + .select_and_reserve(second) + .await + .expect_err("a live id cannot be booked again"); + assert!(matches!(err, SelectionError::Conflict(_)), "{err:?}"); + + let (entry, _) = core + .indexed_booking("shared") + .expect("first booking indexed"); + assert_eq!(entry.key.routing_group, "group-a"); + assert!(entry.scheduler.has_request("shared")); + assert!( + !core + .entry(&RoutingPartitionId::new("model", "group-b")) + .expect("entry") + .scheduler + .has_request("shared") + ); + } + + #[tokio::test] + async fn explicit_reservation_of_a_live_id_is_a_conflict() { + let core = SelectionCore::try_new_local( + test_config(false), + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ) + .expect("valid test config"); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + core.select_and_reserve(reserve_request("taken")) + .await + .expect("booking"); + let err = core + .create_reservation(ReservationRequest { + model_name: "model".to_string(), + routing_group: "default".to_string(), + selection_id: "taken".to_string(), + worker_id: Some(1), + dp_rank: None, + prompt: prompt(), + router_config_override: None, + expected_output_tokens: None, + effective_prefill_tokens: None, + track_prefill_tokens: None, + }) + .await + .expect_err("explicit booking of a live id"); + assert!(matches!(err, SelectionError::Conflict(_)), "{err:?}"); + } + + #[test] + fn index_observer_replaces_stale_and_claimed_rows_of_its_partition() { + use crate::scheduling::AttemptId; + let index = Arc::new(RwLock::new(HashMap::new())); + let partition = RoutingPartitionId::new("model", "default"); + let observer = ReservationIndexObserver { + index: Arc::clone(&index), + partition: partition.clone(), + host: None, + }; + let booking = |attempt: u64| SchedulerBookingDescriptor { + request_id: "shared".to_string(), + worker: WorkerWithDpRank::new(1, 0), + attempt_id: AttemptId::new(attempt), + }; + let row = |partition: &RoutingPartitionId, booking| Reservation { + partition: partition.clone(), + booking, + _affinity_lease: None, + }; + + // A stale row (its booking expired before the sweep ran) yields to the mirror. + index + .write() + .insert("shared".to_string(), row(&partition, Some(booking(1)))); + observer.admitted(booking(2)); + assert_eq!(index.read()["shared"].booking, Some(booking(2))); + + // A completion for the replaced booking leaves the live mirror alone. + observer.completed(&booking(1)); + assert!(index.read().contains_key("shared")); + observer.completed(&booking(2)); + assert!(!index.read().contains_key("shared")); + + // A claim whose local booking is about to fail also yields, and dropping + // that claim keeps the mirror. + let claim = ReservationClaim { + index: &index, + selection_id: "shared".to_string(), + armed: true, + }; + index + .write() + .insert("shared".to_string(), row(&partition, None)); + observer.admitted(booking(3)); + drop(claim); + assert_eq!(index.read()["shared"].booking, Some(booking(3))); + + // Another partition's row is never replaced. + let other = RoutingPartitionId::new("model", "other"); + index + .write() + .insert("shared".to_string(), row(&other, Some(booking(4)))); + observer.admitted(booking(5)); + assert_eq!(index.read()["shared"].booking, Some(booking(4))); + } + + #[tokio::test] + async fn free_of_an_in_flight_reservation_is_not_found() { + let mut config = test_config(false); + config.router_queue_threshold = Some(0.0); + let core = Arc::new( + SelectionCore::try_new_local( + config, + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ) + .expect("valid test config"), + ); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + core.select_and_reserve(reserve_request("res-a")) + .await + .expect("initial reservation"); + let queued_core = core.clone(); + let queued = tokio::spawn(async move { + queued_core + .select_and_reserve(reserve_request("queued")) + .await + }); + wait_for_pending_selection(&core).await; + + // A free racing the in-flight booking neither frees nor evicts it. + let err = core + .free_reservation("queued") + .await + .expect_err("in-flight id is not a reservation yet"); + assert!(matches!(err, SelectionError::NotFound(_)), "{err:?}"); + assert!(core.reservation_index.read().contains_key("queued")); + + core.free_reservation("res-a").await.expect("free res-a"); + tokio::time::timeout(Duration::from_secs(2), queued) + .await + .expect("queued selection timed out") + .expect("task panicked") + .expect("queued selection books"); + core.free_reservation("queued").await.expect("free queued"); + assert!(core.reservation_index.read().is_empty()); + } + + #[tokio::test] + async fn prefill_complete_is_idempotent_for_a_live_booking() { + let mut config = test_config(false); + config.router_track_prefill_tokens = true; + let core = SelectionCore::try_new_local( + config, + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ) + .expect("valid test config"); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + core.select_and_reserve(reserve_request("live")) + .await + .expect("booking"); + core.prefill_complete("live").await.expect("first mark"); + core.prefill_complete("live") + .await + .expect("a repeated mark on a live booking is not an error"); + let (entry, _) = core.indexed_booking("live").expect("still indexed"); + assert!(entry.scheduler.has_request("live")); + core.free_reservation("live").await.expect("free"); + } + + #[tokio::test] + async fn mirrored_replica_bookings_are_indexed_until_freed() { + let (outbound_tx, _outbound_rx) = mpsc::channel(16); + let (inbound_tx, inbound_rx) = mpsc::channel(16); + let channels = parking_lot::Mutex::new(Some(HostReplicaChannels { + outbound: Some(outbound_tx), + inbound_tx: inbound_tx.clone(), + inbound_rx, + process_id: 7, + })); + let core = core_with_host(SelectionHost { + replication: HostReplication { + channels: Some(Arc::new(move |_| channels.lock().take())), + ..HostReplication::default() + }, + ..SelectionHost::default() + }); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + // Replica events for a worker the scheduler has not registered yet are + // dropped; the first local booking registers it. + core.select_and_reserve(reserve_request("warm")) + .await + .expect("warm booking"); + core.free_reservation("warm").await.expect("free warm"); + let entry = core + .entry(&RoutingPartitionId::new("model", "default")) + .expect("entry"); + let peer_event = |request_id: &str, data| ActiveSequenceEvent { + request_id: request_id.to_string(), + worker: WorkerWithDpRank::new(1, 0), + data, + router_id: 99, + lora_name: None, + }; + let add = |request_id: &str| { + peer_event( + request_id, + ActiveSequenceEventData::AddRequest { + token_sequence: Some(vec![1, 2]), + track_prefill_tokens: false, + expected_output_tokens: None, + prefill_load_hint: None, + }, + ) + }; + + inbound_tx.send(add("peer-a")).await.expect("send"); + inbound_tx.send(add("peer-b")).await.expect("send"); + wait_until("mirrored bookings indexed", || { + core.indexed_booking("peer-a").is_some() && core.indexed_booking("peer-b").is_some() + }) + .await; + + // A lifecycle call on a mirrored booking resolves through the index. + core.free_reservation("peer-a") + .await + .expect("free mirrored"); + assert!(!entry.scheduler.has_request("peer-a")); + assert!(core.indexed_booking("peer-a").is_none()); + + // The peer freeing its own booking removes the mirror. + inbound_tx + .send(peer_event("peer-b", ActiveSequenceEventData::Free)) + .await + .expect("send"); + wait_until("mirror removed", || { + core.indexed_booking("peer-b").is_none() + }) + .await; + assert!(!entry.scheduler.has_request("peer-b")); + } + + #[tokio::test] + async fn lifecycle_operations_resolve_through_the_index() { let mut config = test_config(false); config.router_track_prefill_tokens = true; let core = SelectionCore::try_new_local( @@ -3392,9 +3795,11 @@ mod tests { Some(&key_b) ); assert_eq!( - core.lifecycle_entries("booked")[0].key, - key_b, - "indexed partition is tried first" + core.indexed_booking("booked") + .expect("indexed booking") + .0 + .key, + key_b ); // The explicit reservation path records too. @@ -3518,7 +3923,7 @@ mod tests { let started = Arc::clone(&lifecycle_started); std::thread::spawn(move || { started.store(true, Ordering::Release); - core.lifecycle_entries("live").len() + core.indexed_booking("live").is_some() }) }; while !lifecycle_started.load(Ordering::Acquire) { @@ -3535,7 +3940,7 @@ mod tests { ); drop(sweep_entries); writer.join().expect("partition writer"); - assert_eq!(lifecycle.join().expect("lifecycle lookup"), 1); + assert!(lifecycle.join().expect("lifecycle lookup")); } #[tokio::test(flavor = "current_thread")] diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index dd7dec912514..7e14ea934f5d 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -1951,6 +1951,7 @@ impl KvRouter { self.scheduler .mark_prefill_completed_if_booking(booking) .await + .map(|_| ()) } /// Number of requests currently parked in the scheduler queue. diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs index 8ee4435a2379..c22ea1c4dedb 100644 --- a/lib/llm/src/kv_router/embedded.rs +++ b/lib/llm/src/kv_router/embedded.rs @@ -452,6 +452,7 @@ impl EmbeddedSelection { .scheduler() .mark_prefill_completed_if_booking(booking) .await + .map(|_| ()) } pub(crate) fn pending_count(&self) -> usize { From fe816fc6a2717a954cb3bbfa23d47000d2349d90 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Thu, 10 Sep 2026 12:21:56 -0700 Subject: [PATCH 052/102] feat(router): share the session affinity commit and configure its mode SessionAffinity gains a config (TTL, Hard/Soft mode, limits) and a commit step that binds a held session to the worker a request was dispatched to, with the same Hard and Soft semantics the frontend applies in into_stream. The standalone service and EPP now commit through it, defaulting to Hard: a booking that lands away from a live bound worker is rejected with 400 and the stale binding dropped, instead of silently rebinding. A bound worker the catalog no longer schedules, including one that departs while the request is queued, re-initializes the session on the newly selected worker. The affinity commit runs before the routing decision is recorded and before the booking lease is committed, so a rejection leaves no booking, index entry, or approximate-cache credit behind. SelectionServiceBuilder exposes session_affinity_mode for Soft. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Ishan Dhanani --- .../src/services/selection/affinity.rs | 174 ++++++--- .../src/services/selection/catalog.rs | 45 +-- .../src/services/selection/core/mod.rs | 369 +++++++++++++----- .../src/services/selection/service.rs | 13 +- lib/llm/src/kv_router/embedded.rs | 4 +- lib/llm/src/session_affinity/coordinator.rs | 6 +- 6 files changed, 432 insertions(+), 179 deletions(-) diff --git a/lib/kv-router/src/services/selection/affinity.rs b/lib/kv-router/src/services/selection/affinity.rs index b4a630ddfeef..9b9225e290d4 100644 --- a/lib/kv-router/src/services/selection/affinity.rs +++ b/lib/kv-router/src/services/selection/affinity.rs @@ -102,9 +102,36 @@ enum AffinityEntry { }, } +/// How a table binds sessions; see [`SessionAffinityMode`] for the +/// dispatch-time semantics of `mode`. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct SessionAffinityConfig { + pub ttl: Duration, + pub mode: SessionAffinityMode, + pub max_entries: usize, + pub max_session_id_bytes: usize, +} + +impl SessionAffinityConfig { + pub fn new(ttl: Duration) -> Self { + Self { + ttl, + mode: SessionAffinityMode::default(), + max_entries: MAX_SESSION_AFFINITY_ENTRIES, + max_session_id_bytes: MAX_SESSION_AFFINITY_ID_BYTES, + } + } + + pub fn with_mode(mut self, mode: SessionAffinityMode) -> Self { + self.mode = mode; + self + } +} + struct Inner { entries: DashMap, ttl: Duration, + mode: SessionAffinityMode, max_entries: usize, max_session_id_bytes: usize, entry_count: AtomicUsize, @@ -144,9 +171,10 @@ impl WeakSessionAffinity { } } -/// A session acquired for one request. -pub enum Acquired { - /// New (or expired) session: select a worker, then `commit` it. +/// A session held for one request until [`SessionAffinity::commit`] binds it +/// to the worker the request was dispatched to. +pub enum Hold { + /// New (or expired) session: select a worker, then commit it. Initialize(AffinityInitialization), /// Bound session: route to `target`; the lease releases on drop. Bound { @@ -155,26 +183,26 @@ pub enum Acquired { }, } +impl Hold { + /// The worker a bound session steers to; `None` while initializing. + pub fn target(&self) -> Option { + match self { + Self::Initialize(_) => None, + Self::Bound { target, .. } => Some(*target), + } + } +} + /// One step of acquiring a session. pub enum AcquireStep { - /// New (or expired) session: select a worker, then `commit` it. - Initialize(AffinityInitialization), - /// Bound session: route to `target`; the lease releases on drop. - Bound { - target: AffinityTarget, - lease: AffinityLease, - }, + Held(Hold), /// Another request is initializing this session; await, then retry. Wait(Pin>), } impl SessionAffinity { pub fn new(ttl: Duration) -> Result { - Self::new_with_limits( - ttl, - MAX_SESSION_AFFINITY_ENTRIES, - MAX_SESSION_AFFINITY_ID_BYTES, - ) + Self::with_config(SessionAffinityConfig::new(ttl)) } pub(crate) fn validate_ttl(ttl: Duration) -> Result<(), AffinityError> { @@ -193,12 +221,21 @@ impl SessionAffinity { max_entries: usize, max_session_id_bytes: usize, ) -> Result { - Self::validate_ttl(ttl)?; - let inner = Arc::new(Inner { - entries: DashMap::new(), - ttl, + Self::with_config(SessionAffinityConfig { max_entries, max_session_id_bytes, + ..SessionAffinityConfig::new(ttl) + }) + } + + pub fn with_config(config: SessionAffinityConfig) -> Result { + Self::validate_ttl(config.ttl)?; + let inner = Arc::new(Inner { + entries: DashMap::new(), + ttl: config.ttl, + mode: config.mode, + max_entries: config.max_entries, + max_session_id_bytes: config.max_session_id_bytes, entry_count: AtomicUsize::new(0), next_revision: AtomicU64::new(1), next_sequence: AtomicU64::new( @@ -217,8 +254,9 @@ impl SessionAffinity { }); Self::spawn_reaper(&inner); tracing::info!( - ttl_secs = ttl.as_secs(), - max_entries, + ttl_secs = config.ttl.as_secs(), + mode = ?config.mode, + max_entries = config.max_entries, "session affinity enabled" ); Ok(Self { inner }) @@ -277,6 +315,10 @@ impl SessionAffinity { self.inner.ttl } + pub fn mode(&self) -> SessionAffinityMode { + self.inner.mode + } + /// Acquire `session_id` for a request. `requested_target` is an explicit /// pin the request carries; a bound session must agree with it. pub fn try_acquire( @@ -299,14 +341,16 @@ impl SessionAffinity { revision, notify: notify.clone(), }); - Ok(AcquireStep::Initialize(AffinityInitialization { - table: Arc::downgrade(&self.inner), - session_id: session_id.to_string(), - revision, - notify, - requested_target, - active: true, - })) + Ok(AcquireStep::Held(Hold::Initialize( + AffinityInitialization { + table: Arc::downgrade(&self.inner), + session_id: session_id.to_string(), + revision, + notify, + requested_target, + active: true, + }, + ))) } Entry::Occupied(mut entry) => match entry.get_mut() { AffinityEntry::Initializing { notify, .. } => { @@ -333,14 +377,16 @@ impl SessionAffinity { revision, notify: notify.clone(), }; - Ok(AcquireStep::Initialize(AffinityInitialization { - table: Arc::downgrade(&self.inner), - session_id: session_id.to_string(), - revision, - notify, - requested_target, - active: true, - })) + Ok(AcquireStep::Held(Hold::Initialize( + AffinityInitialization { + table: Arc::downgrade(&self.inner), + session_id: session_id.to_string(), + revision, + notify, + requested_target, + active: true, + }, + ))) } AffinityEntry::Bound { target, @@ -358,7 +404,7 @@ impl SessionAffinity { "session affinity hit: reusing pinned worker" ); *active_leases += 1; - Ok(AcquireStep::Bound { + Ok(AcquireStep::Held(Hold::Bound { target: *target, lease: AffinityLease { table: Arc::downgrade(&self.inner), @@ -367,7 +413,7 @@ impl SessionAffinity { version: *version, active: true, }, - }) + })) } }, } @@ -379,14 +425,46 @@ impl SessionAffinity { &self, session_id: &str, requested_target: Option, - ) -> Result { + ) -> Result { loop { match self.try_acquire(session_id, requested_target)? { AcquireStep::Wait(notified) => notified.await, - AcquireStep::Initialize(init) => return Ok(Acquired::Initialize(init)), - AcquireStep::Bound { target, lease } => { - return Ok(Acquired::Bound { target, lease }); + AcquireStep::Held(hold) => return Ok(hold), + } + } + } + + /// Bind the held session to `dispatched`. A `Bound` session dispatched + /// elsewhere follows the dispatch in `Soft` mode and is invalidated with + /// an error in `Hard` mode, so its next request re-initializes. + pub fn commit( + &self, + hold: Hold, + dispatched: AffinityTarget, + ) -> Result { + match hold { + Hold::Initialize(initialization) => { + let lease = initialization.commit(dispatched)?; + lease.publish(dispatched); + Ok(lease) + } + Hold::Bound { target, mut lease } => { + if self.inner.mode == SessionAffinityMode::Soft { + let rebound = AffinityLease::rebound_target(target, dispatched); + if target == rebound { + lease.publish(target); + } else if lease.rebind(target, rebound) { + lease.publish(rebound); + } + return Ok(lease); } + if let Err(error) = validate_dispatch_target(lease.session_id(), target, dispatched) + { + lease.invalidate(); + return Err(error); + } + lease.publish(target); + Ok(lease) } } } @@ -791,11 +869,19 @@ impl AffinityLease { ) }); match removed { - Some((_, AffinityEntry::Bound { target, .. })) => { + Some(( + _, + AffinityEntry::Bound { + target, + active_leases, + .. + }, + )) => { tracing::debug!( session_id = %self.session_id, worker_id = target.worker_id, dp_rank = ?target.dp_rank, + other_holders = active_leases.saturating_sub(1), "invalidated current session affinity binding" ); self.active = false; diff --git a/lib/kv-router/src/services/selection/catalog.rs b/lib/kv-router/src/services/selection/catalog.rs index 805911b154fd..52910a7b44b2 100644 --- a/lib/kv-router/src/services/selection/catalog.rs +++ b/lib/kv-router/src/services/selection/catalog.rs @@ -64,22 +64,17 @@ impl WorkerCatalog { } pub(super) fn has_schedulable_for_key(&self, key: &RoutingPartitionId) -> bool { - self.workers.read().values().any(|record| { - record.lifecycle == WorkerLifecycle::Schedulable - && record.model_name == key.model_name - && record.routing_group == key.routing_group - }) + self.workers + .read() + .values() + .any(|record| schedulable_in(record, key)) } pub(super) fn schedulable_worker_ids_for_key(&self, key: &RoutingPartitionId) -> Vec { self.workers .read() .values() - .filter(|record| { - record.lifecycle == WorkerLifecycle::Schedulable - && record.model_name == key.model_name - && record.routing_group == key.routing_group - }) + .filter(|record| schedulable_in(record, key)) .map(|record| record.worker_id) .collect() } @@ -120,11 +115,7 @@ impl WorkerCatalog { self.workers .read() .values() - .filter(|record| { - record.lifecycle == WorkerLifecycle::Schedulable - && record.model_name == key.model_name - && record.routing_group == key.routing_group - }) + .filter(|record| schedulable_in(record, key)) .filter_map(|record| { record .scheduler_config() @@ -141,6 +132,13 @@ impl WorkerCatalog { .count() } + pub(super) fn is_schedulable(&self, worker_id: WorkerId, key: &RoutingPartitionId) -> bool { + self.workers + .read() + .get(&worker_id) + .is_some_and(|record| schedulable_in(record, key)) + } + pub(super) fn schedulable_endpoint( &self, worker_id: WorkerId, @@ -148,10 +146,7 @@ impl WorkerCatalog { ) -> Option { let workers = self.workers.read(); let record = workers.get(&worker_id)?; - if record.lifecycle != WorkerLifecycle::Schedulable - || record.model_name != key.model_name - || record.routing_group != key.routing_group - { + if !schedulable_in(record, key) { return None; } record.endpoint.clone() @@ -164,13 +159,15 @@ impl WorkerCatalog { ) -> Option { let workers = self.workers.read(); let record = workers.get(&worker.worker_id)?; - if record.lifecycle != WorkerLifecycle::Schedulable - || record.model_name != key.model_name - || record.routing_group != key.routing_group - || !record.dp_ranks().any(|rank| rank == worker.dp_rank) - { + if !schedulable_in(record, key) || !record.dp_ranks().any(|rank| rank == worker.dp_rank) { return None; } record.endpoint.clone() } } + +fn schedulable_in(record: &WorkerCatalogRecord, key: &RoutingPartitionId) -> bool { + record.lifecycle == WorkerLifecycle::Schedulable + && record.model_name == key.model_name + && record.routing_group == key.routing_group +} diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index a6f71f2c5922..7cc6f0b239fb 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -49,7 +49,7 @@ use crate::services::indexer::registry::WorkerRegistry; use crate::services::overlap::MooncakeOverlapSummary; use crate::tracking_hash::{TrackingHashContext, TrackingHashScope}; -use super::affinity::{AcquireStep, Acquired, AffinityError, AffinityLease, SessionAffinity}; +use super::affinity::{AffinityError, AffinityLease, Hold, SessionAffinity, SessionAffinityConfig}; use super::catalog::WorkerCatalog; use super::error::SelectionError; use super::ingress::{KvEventIngress, ZmqDirectIngress}; @@ -96,9 +96,12 @@ impl SelectionPartition { &self.0.indexer } - /// Return this partition's affinity table, initialized with its serving TTL. - pub fn session_affinity(&self, ttl: Duration) -> Result { - self.0.session_affinity(ttl).cloned() + /// Return this partition's affinity table, initialized with `config`. + pub fn session_affinity( + &self, + config: SessionAffinityConfig, + ) -> Result { + self.0.session_affinity(config).cloned() } } @@ -124,11 +127,14 @@ struct PreparedSelectionInputs { } impl SelectionEntry { - fn session_affinity(&self, ttl: Duration) -> Result<&SessionAffinity, SelectionError> { + fn session_affinity( + &self, + config: SessionAffinityConfig, + ) -> Result<&SessionAffinity, SelectionError> { let table = self .affinity .get_or_try_init(|| -> Result<_, SelectionError> { - let table = SessionAffinity::new(ttl).map_err(affinity_error)?; + let table = SessionAffinity::with_config(config).map_err(affinity_error)?; if let Some(config) = &self.replica_config && let Some(sink) = config.affinity_sink(&self.key) { @@ -136,7 +142,7 @@ impl SelectionEntry { } Ok(table) })?; - if table.ttl() != ttl { + if table.ttl() != config.ttl { return Err(SelectionError::Conflict(format!( "session affinity TTL mismatch for {}", self.key @@ -446,11 +452,9 @@ pub struct SelectionCore { /// `create_reservation` can replay them without re-sending the prompt. selection_cache: SelectionCache, tracking_hash: Arc, - session_affinity_ttl: Option, + session_affinity: Option, } -type SessionTarget = super::affinity::AffinityTarget; - fn affinity_error(error: AffinityError) -> SelectionError { match error { AffinityError::InvalidArgument(message) => SelectionError::BadRequest(message), @@ -522,7 +526,7 @@ impl SelectionCore { cache_config: SelectionCacheConfig, tracking_hash: Arc, indexer_policy: IndexerPolicy, - session_affinity_ttl: Option, + session_affinity: Option, ) -> Self { let cancel_token = cancel_token.child_token(); let indexer_registry = Arc::new( @@ -550,7 +554,7 @@ impl SelectionCore { replica_config, selection_cache: SelectionCache::new(&cache_config), tracking_hash, - session_affinity_ttl, + session_affinity, } } @@ -933,8 +937,8 @@ impl SelectionCore { entry.is_eagle ))); } - if let Some(ttl) = self.session_affinity_ttl { - entry.session_affinity(ttl)?; + if let Some(config) = self.session_affinity { + entry.session_affinity(config)?; } Ok(entry) } @@ -1111,30 +1115,8 @@ impl SelectionCore { let mut affinity_hold = None; let affinity_target = match (session_id.as_deref(), table) { (Some(session_id), Some(table)) if book => { - let acquire_result = tokio::select! { - _ = self.cancel_token.cancelled() => return Err(SelectionError::Scheduler(KvSchedulerError::SubscriberShutdown)), - result = table.acquire(session_id, None) => result, - }; - match acquire_result { - Ok(Acquired::Initialize(init)) => { - affinity_hold = Some(Acquired::Initialize(init)); - None - } - Ok(Acquired::Bound { target, lease }) => { - affinity_hold = Some(Acquired::Bound { target, lease }); - Some(WorkerAffinityTarget::new(target.worker_id, target.dp_rank)) - } - // A full table is a router-side limit, not a client fault: - // route this request unpinned instead of failing it. - Err(AffinityError::ResourceExhausted(_)) => { - tracing::debug!( - session_id, - "affinity table full; routing without session affinity" - ); - None - } - Err(error) => return Err(affinity_error(error)), - } + affinity_hold = self.hold_session(table, session_id, &key).await?; + affinity_hold.as_ref().and_then(Hold::target) } (Some(session_id), Some(table)) => table .query_target(session_id, None) @@ -1320,22 +1302,22 @@ impl SelectionCore { "booked selection has no lifecycle lease".to_string(), )); }; + // A rejected affinity commit returns while the lease is still armed, + // so the booking is freed and nothing below is recorded. + let affinity_lease = match (affinity_hold, session_id.as_deref(), table) { + (Some(hold), Some(session_id), Some(table)) => Some( + self.commit_session(table, hold, session_id, response.best_worker, &key) + .await?, + ), + _ => None, + }; if let Some(hashes) = routing_hashes.clone() { self.record_routing_decision(&entry, response.best_worker, hashes) .await; } - // Nothing awaits from here to `install`: the binding, the booking - // handover and the index entry land together. - let booking = lease.commit().ok_or_else(missing_booking)?; - let affinity_lease = match (affinity_hold, session_id.as_deref(), table) { - (Some(hold), Some(session_id), Some(table)) => { - Self::bind_session(table, hold, session_id, response.best_worker) - } - _ => None, - }; claim.install(Reservation { partition: key.clone(), - booking: Some(booking), + booking: Some(lease.commit().ok_or_else(missing_booking)?), _affinity_lease: affinity_lease, }); } @@ -1626,38 +1608,74 @@ impl SelectionCore { }) } - /// Bind (or confirm) `session_id` to the worker a booking landed on and - /// keep the lease with the booking. A bound session whose worker was not - /// selected (it left) is invalidated and rebound. - fn bind_session( + /// Hold `session_id` for a booking, re-initializing a binding whose worker + /// this partition can no longer schedule. `None` when the table is full: a + /// router-side limit, not a client fault, so the request routes unpinned. + async fn hold_session( + &self, table: &SessionAffinity, - hold: Acquired, session_id: &str, - worker: WorkerWithDpRank, - ) -> Option { - let selected = SessionTarget::new(worker.worker_id, Some(worker.dp_rank)); - let lease = match hold { - Acquired::Initialize(init) => init.commit(selected).ok(), - Acquired::Bound { target, lease } - if target.worker_id == selected.worker_id - && target.dp_rank.is_none_or(|rank| rank == worker.dp_rank) => - { - lease.publish(target); - return Some(lease); + key: &RoutingPartitionId, + ) -> Result, SelectionError> { + loop { + let acquired = tokio::select! { + _ = self.cancel_token.cancelled() => { + return Err(SelectionError::Scheduler(KvSchedulerError::SubscriberShutdown)); + } + result = table.acquire(session_id, None) => result, + }; + match acquired { + Ok(Hold::Bound { target, mut lease }) + if !self.catalog.is_schedulable(target.worker_id, key) => + { + tracing::debug!( + session_id, + worker_id = target.worker_id, + "session affinity target is not schedulable; re-initializing" + ); + lease.invalidate(); + } + Ok(hold) => return Ok(Some(hold)), + Err(AffinityError::ResourceExhausted(_)) => { + tracing::debug!( + session_id, + "affinity table full; routing without session affinity" + ); + return Ok(None); + } + Err(error) => return Err(affinity_error(error)), } - Acquired::Bound { mut lease, .. } => { - lease.invalidate(); - match table.try_acquire(session_id, None) { - Ok(AcquireStep::Initialize(init)) => init.commit(selected).ok(), - Ok(AcquireStep::Bound { lease, .. }) => return Some(lease), - Ok(AcquireStep::Wait(_)) | Err(_) => None, + } + } + + /// Bind the held session to `dispatched`. A `Hard` rejection whose bound + /// worker departed after [`Self::hold_session`] checked it is not a client + /// fault: the binding is re-initialized on the dispatched worker instead. + async fn commit_session( + &self, + table: &SessionAffinity, + hold: Hold, + session_id: &str, + dispatched: WorkerWithDpRank, + key: &RoutingPartitionId, + ) -> Result { + let bound = hold.target(); + let dispatched = WorkerAffinityTarget::new(dispatched.worker_id, Some(dispatched.dp_rank)); + match table.commit(hold, dispatched) { + Ok(lease) => Ok(lease), + Err(error) => { + let departed = + bound.is_some_and(|target| !self.catalog.is_schedulable(target.worker_id, key)); + if !departed { + return Err(affinity_error(error)); + } + // `commit` invalidated the stale binding; the next hold initializes. + match self.hold_session(table, session_id, key).await? { + Some(hold) => table.commit(hold, dispatched).map_err(affinity_error), + None => Err(affinity_error(error)), } } - }; - if let Some(lease) = &lease { - lease.publish(selected); } - lease } /// Apply a session binding a replica published. @@ -2113,6 +2131,7 @@ impl Drop for SelectionCore { #[cfg(test)] mod tests { + use super::super::affinity::SessionAffinityMode; use super::*; use crate::protocols::StorageTier; use crate::services::indexer::backend::test_util::store_event; @@ -2939,35 +2958,6 @@ mod tests { ); } - #[tokio::test] - async fn concurrent_failover_keeps_the_winning_binding_lease() { - let table = SessionAffinity::new(Duration::from_secs(60)).expect("affinity table"); - let Acquired::Initialize(init) = table.acquire("s", None).await.expect("acquire") else { - panic!("fresh session must initialize"); - }; - drop(init.commit(SessionTarget::new(1, Some(0))).expect("commit")); - let first = table.acquire("s", None).await.expect("first acquire"); - let second = table.acquire("s", None).await.expect("second acquire"); - assert!(matches!(first, Acquired::Bound { target, .. } if target.worker_id == 1)); - - // Worker 1 left; both bookings landed elsewhere. The first rebinds the - // session, the second must hold that binding rather than end up leaseless. - let winner = SelectionCore::bind_session(&table, first, "s", WorkerWithDpRank::new(2, 0)) - .expect("first failover rebinds"); - let follower = - SelectionCore::bind_session(&table, second, "s", WorkerWithDpRank::new(3, 0)) - .expect("second failover holds the rebound session"); - drop(winner); - drop(follower); - assert_eq!( - table - .query_target("s", None) - .expect("query") - .map(|t| t.worker_id), - Some(2) - ); - } - #[tokio::test] async fn injected_availability_provider_restricts_selection() { let available: Arc>>>> = @@ -4053,7 +4043,7 @@ mod tests { assert_eq!(response.overlap.dp, HashMap::from([("0".to_string(), 8)])); } - fn core_with_session_affinity() -> SelectionCore { + fn core_with_session_affinity_mode(mode: SessionAffinityMode) -> SelectionCore { let config = test_config(false); let tracking_hash = Arc::new( TrackingHashContext::from_config(&config).expect("valid tracking hash configuration"), @@ -4071,16 +4061,183 @@ mod tests { SelectionCacheConfig::default(), tracking_hash, indexer_policy, - Some(Duration::from_secs(10)), + Some(SessionAffinityConfig::new(Duration::from_secs(10)).with_mode(mode)), ) } + fn core_with_session_affinity() -> SelectionCore { + core_with_session_affinity_mode(SessionAffinityMode::Hard) + } + + fn bound_worker(core: &SelectionCore, session_id: &str) -> Option { + core.entry(&RoutingPartitionId::new("model", "default")) + .and_then(|entry| entry.affinity.get().cloned()) + .and_then(|table| table.query_target(session_id, None).expect("query")) + .map(|target| target.worker_id) + } + fn session_reservation(selection_id: &str, session_id: &str) -> SelectAndReserveRequest { let mut request = reserve_request(selection_id); request.session_id = Some(session_id.to_string()); request } + #[tokio::test] + async fn departed_session_worker_reinitializes_the_session() { + let core = core_with_session_affinity(); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + core.upsert_worker(worker(2)).await.expect("worker upsert"); + let first = core + .select_and_reserve(session_reservation("r1", "s")) + .await + .expect("first booking"); + core.free_reservation("r1").await.expect("free"); + core.delete_worker(first.worker_id).await.expect("delete"); + + let second = core + .select_and_reserve(session_reservation("r2", "s")) + .await + .expect("session must move off a departed worker"); + assert_ne!(second.worker_id, first.worker_id); + assert_eq!(bound_worker(&core, "s"), Some(second.worker_id)); + } + + #[tokio::test] + async fn session_worker_departing_after_the_hold_reinitializes_the_session() { + let core = core_with_session_affinity(); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + core.upsert_worker(worker(2)).await.expect("worker upsert"); + let first = core + .select_and_reserve(session_reservation("r1", "s")) + .await + .expect("first booking"); + core.free_reservation("r1").await.expect("free"); + let key = RoutingPartitionId::new("model", "default"); + + // One poll takes the hold (the bound worker still passes the check) + // and hands the request to the scheduler actor, which has not run yet. + let mut second = Box::pin(core.select_and_reserve(session_reservation("r2", "s"))); + let mut context = std::task::Context::from_waker(std::task::Waker::noop()); + assert!(second.as_mut().poll(&mut context).is_pending()); + core.catalog + .set_lifecycle(first.worker_id, WorkerLifecycle::Draining, Vec::new()); + core.publish_scheduler_config(&key).expect("publish"); + + let second = second + .await + .expect("a departure after the hold is not a client fault"); + assert_ne!(second.worker_id, first.worker_id); + assert_eq!(bound_worker(&core, "s"), Some(second.worker_id)); + } + + #[tokio::test] + async fn concurrent_holds_on_a_departed_worker_both_land_on_the_replacement() { + let table = SessionAffinity::new(Duration::from_secs(60)).expect("affinity table"); + let departed = WorkerAffinityTarget::new(1, Some(0)); + let replacement = WorkerAffinityTarget::new(2, Some(0)); + let Hold::Initialize(init) = table.acquire("s", None).await.expect("acquire") else { + panic!("fresh session must initialize"); + }; + drop( + table + .commit(Hold::Initialize(init), departed) + .expect("bind"), + ); + let ( + Hold::Bound { + lease: mut first, .. + }, + Hold::Bound { + lease: mut second, .. + }, + ) = ( + table.acquire("s", None).await.expect("first hold"), + table.acquire("s", None).await.expect("second hold"), + ) + else { + panic!("both requests hold the departed binding"); + }; + + // Both requests notice the departure, as `hold_session` does: the first + // invalidation drops the binding, the second is a no-op release. + first.invalidate(); + second.invalidate(); + let leases = [ + table.commit( + table.acquire("s", None).await.expect("re-acquire"), + replacement, + ), + table.commit( + table.acquire("s", None).await.expect("re-acquire"), + replacement, + ), + ]; + for lease in leases { + drop(lease.expect("both requests bind to the replacement")); + } + assert_eq!( + table.query_target("s", None).expect("query"), + Some(replacement) + ); + } + + #[tokio::test] + async fn hard_mode_rejects_dispatch_away_from_a_live_binding() { + let core = core_with_session_affinity(); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + core.upsert_worker(worker(2)).await.expect("worker upsert"); + let first = core + .select_and_reserve(session_reservation("r1", "s")) + .await + .expect("first booking"); + core.free_reservation("r1").await.expect("free"); + let other = if first.worker_id == 1 { 2 } else { 1 }; + + // Steering cannot reach the bound worker, so selection lands elsewhere. + let mut request = session_reservation("r2", "s"); + request.allowed_worker_ids = Some(HashSet::from([other])); + let err = core + .select_and_reserve(request) + .await + .expect_err("hard affinity rejects a dispatch away from the binding"); + assert!(matches!(err, SelectionError::BadRequest(_)), "{err:?}"); + let entry = core + .entry(&RoutingPartitionId::new("model", "default")) + .expect("entry"); + wait_until("rejected booking release", || { + !entry.scheduler.has_request("r2") + }) + .await; + assert!(core.reservation_index.read().is_empty()); + assert_eq!(bound_worker(&core, "s"), None, "stale binding is dropped"); + + // The session's next request re-initializes. + let mut request = session_reservation("r3", "s"); + request.allowed_worker_ids = Some(HashSet::from([other])); + let third = core.select_and_reserve(request).await.expect("rebind"); + assert_eq!(third.worker_id, other); + assert_eq!(bound_worker(&core, "s"), Some(other)); + } + + #[tokio::test] + async fn soft_mode_follows_the_dispatch() { + let core = core_with_session_affinity_mode(SessionAffinityMode::Soft); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + core.upsert_worker(worker(2)).await.expect("worker upsert"); + let first = core + .select_and_reserve(session_reservation("r1", "s")) + .await + .expect("first booking"); + core.free_reservation("r1").await.expect("free"); + let other = if first.worker_id == 1 { 2 } else { 1 }; + + let mut request = session_reservation("r2", "s"); + request.allowed_worker_ids = Some(HashSet::from([other])); + let second = core.select_and_reserve(request).await.expect("soft rebind"); + assert_eq!(second.worker_id, other); + assert_eq!(bound_worker(&core, "s"), Some(other)); + } + #[tokio::test] async fn session_stays_on_its_first_worker_across_bookings() { let core = core_with_session_affinity(); @@ -4468,7 +4625,7 @@ mod tests { .partition(&RoutingPartitionId::new("model", "default")) .unwrap(); assert!(matches!( - partition.session_affinity(Duration::from_secs(20)), + partition.session_affinity(SessionAffinityConfig::new(Duration::from_secs(20))), Err(SelectionError::Conflict(_)) )); } diff --git a/lib/kv-router/src/services/selection/service.rs b/lib/kv-router/src/services/selection/service.rs index 3925dba782d3..0cb4ab1f4b03 100644 --- a/lib/kv-router/src/services/selection/service.rs +++ b/lib/kv-router/src/services/selection/service.rs @@ -16,6 +16,7 @@ use crate::services::common::replica_sync::{ use crate::services::indexer::backend::IndexerPolicy; use crate::tracking_hash::TrackingHashContext; +use super::affinity::{SessionAffinityConfig, SessionAffinityMode}; use super::core::{ KvIndexSource, SelectionCore, SelectionHost, SelectionPartition, SelectionServiceConfig, }; @@ -41,6 +42,7 @@ pub struct SelectionServiceBuilder { host: SelectionHost, worker_selection_policy_factory: Option, session_affinity_ttl: Option, + session_affinity_mode: SessionAffinityMode, } /// Warn when a host does not construct workers for explicitly configured policy roles. @@ -81,6 +83,7 @@ impl SelectionServiceBuilder { host: SelectionHost::default(), worker_selection_policy_factory: None, session_affinity_ttl: None, + session_affinity_mode: SessionAffinityMode::default(), } } @@ -124,6 +127,13 @@ impl SelectionServiceBuilder { self } + /// How a bound session treats a booking that lands on another worker: + /// `Hard` (default) rejects it with 400, `Soft` moves the session. + pub fn session_affinity_mode(mut self, mode: SessionAffinityMode) -> Self { + self.session_affinity_mode = mode; + self + } + pub fn replica_sync(mut self, port: u16, peers: Vec) -> Self { self.replica_sync_port = Some(port); self.replica_sync_peers = peers; @@ -188,7 +198,8 @@ impl SelectionServiceBuilder { self.selection_cache, tracking_hash, indexer_policy, - self.session_affinity_ttl, + self.session_affinity_ttl + .map(|ttl| SessionAffinityConfig::new(ttl).with_mode(self.session_affinity_mode)), )); if recover_from_peers { diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs index c22ea1c4dedb..4b9de67863de 100644 --- a/lib/llm/src/kv_router/embedded.rs +++ b/lib/llm/src/kv_router/embedded.rs @@ -371,7 +371,9 @@ impl EmbeddedSelection { &self, ttl: std::time::Duration, ) -> Result { - let table = self.partition.session_affinity(ttl)?; + let table = self.partition.session_affinity( + dynamo_kv_router::services::selection::affinity::SessionAffinityConfig::new(ttl), + )?; Ok(self .affinity .get_or_init(|| crate::session_affinity::AffinityCoordinator::wrap(table)) diff --git a/lib/llm/src/session_affinity/coordinator.rs b/lib/llm/src/session_affinity/coordinator.rs index d8d5b69ef918..9f9d80c27c93 100644 --- a/lib/llm/src/session_affinity/coordinator.rs +++ b/lib/llm/src/session_affinity/coordinator.rs @@ -15,7 +15,7 @@ use std::{ }; use dynamo_kv_router::services::selection::affinity::{ - AcquireStep, AffinityError, AffinityInitialization, AffinityLease, SessionAffinity, + AcquireStep, AffinityError, AffinityInitialization, AffinityLease, Hold, SessionAffinity, validate_dispatch_target, }; use dynamo_runtime::{ @@ -151,10 +151,10 @@ impl AffinityCoordinator { .try_acquire(session_id.as_str(), requested) .map_err(affinity_error)? { - AcquireStep::Initialize(initialization) => { + AcquireStep::Held(Hold::Initialize(initialization)) => { return Ok(AffinityAcquire::Initialize(initialization)); } - AcquireStep::Bound { target, lease } => { + AcquireStep::Held(Hold::Bound { target, lease }) => { return Ok(AffinityAcquire::Bound { target: from_table(target), lease, From 9740a38ec1ecdb8ba95bc0454b9cdef7a998a6e2 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Thu, 10 Sep 2026 12:50:58 -0700 Subject: [PATCH 053/102] refactor(router): route every selection through one core operation SelectionCore::run_selection takes a SelectionOperation (borrowed prompt view, independent session, steering, and admission inputs) and returns a typed SelectionOutcome; the standalone select and select_and_reserve handlers are adapters that build the operation and shape the wire response, so their bodies, including the queue-rejection 503, are unchanged. SelectionAdmission names the three admission shapes (Query, Book, Advisory) instead of a pair of booleans. The two-phase path now caches the session a select steered by and holds it before the replay books, then commits it after, so create_reservation binds the session like select_and_reserve does. Under Hard mode a replay whose session moved to another worker returns 400 and drops the unusable cached selection. The affinity commit never waits on another request's initialization while a booking is held. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Ishan Dhanani --- .../src/services/selection/core/mod.rs | 490 ++++++++++++------ .../src/services/selection/core/operation.rs | 100 ++++ lib/kv-router/src/services/selection/input.rs | 61 ++- lib/kv-router/src/services/selection/mod.rs | 7 +- .../src/services/selection/pending.rs | 4 + lib/kv-router/src/services/selection/tests.rs | 21 + 6 files changed, 518 insertions(+), 165 deletions(-) create mode 100644 lib/kv-router/src/services/selection/core/operation.rs diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index 7cc6f0b239fb..f074bba43185 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -1,7 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use std::collections::{HashMap, HashSet}; +use std::collections::HashMap; use std::sync::Arc; use std::time::{Duration, Instant}; @@ -21,10 +21,9 @@ use crate::kv_hints::{ #[cfg(test)] use crate::protocols::ActiveSequenceEventData; use crate::protocols::{ - ActiveSequenceEvent, LocalBlockHash, PrefillLoadHint, RoutingConstraints, SharedCacheHits, - WorkerAffinityTarget, WorkerConfigLike, WorkerId, WorkerWithDpRank, + ActiveSequenceEvent, LocalBlockHash, PrefillLoadHint, SharedCacheHits, WorkerAffinityTarget, + WorkerConfigLike, WorkerId, WorkerWithDpRank, }; -use crate::scheduling::config::RouterConfigOverride; use crate::scheduling::queue::SchedulerBookingDescriptor; use crate::scheduling::selector::WorkerSelectionPolicy; use crate::scheduling::{ @@ -49,11 +48,19 @@ use crate::services::indexer::registry::WorkerRegistry; use crate::services::overlap::MooncakeOverlapSummary; use crate::tracking_hash::{TrackingHashContext, TrackingHashScope}; -use super::affinity::{AffinityError, AffinityLease, Hold, SessionAffinity, SessionAffinityConfig}; +mod operation; + +pub use operation::{ + Selected, SelectionAdmission, SelectionOperation, SelectionOutcome, SessionBinding, +}; + +use super::affinity::{ + AcquireStep, AffinityError, AffinityLease, Hold, SessionAffinity, SessionAffinityConfig, +}; use super::catalog::WorkerCatalog; use super::error::SelectionError; use super::ingress::{KvEventIngress, ZmqDirectIngress}; -use super::input::{PromptRequest, TrackingHashInput}; +use super::input::{PromptView, TrackingHashInput}; use super::pending::{PendingSelection, SelectionCache, SelectionCacheConfig}; use super::types::{ ModelLoadResponse, OverlapScoresRequest, OverlapScoresResponse, PotentialLoadsRequest, @@ -152,24 +159,6 @@ impl SelectionEntry { } } -struct SelectionOperation { - key: RoutingPartitionId, - selection_id: Option, - prompt: PromptRequest, - router_config_override: Option, - expected_output_tokens: Option, - priority_jump: f64, - strict_priority: u32, - policy_class: Option, - session_context: Option, - affinity_target: Option, - pinned_worker: Option, - allowed_worker_ids: Option>, - routing_constraints: RoutingConstraints, - /// Skip queue admission and return the chosen worker's load snapshot. - advisory: bool, -} - /// Resolved inputs for booking a reservation, shared by the cached and explicit /// `create_reservation` paths. struct ReservationBooking { @@ -183,6 +172,7 @@ struct ReservationBooking { lora_name: Option, /// Public block hashes to record into an approximate indexer once booked. routing_hashes: Option>, + session_id: Option, } /// What an embedding host supplies to every partition the core creates, @@ -1015,26 +1005,36 @@ impl SelectionCore { policy_class: Option, ) -> Result { let session_context = req.take_session_context(); - self.schedule_selection( - SelectionOperation { + let is_steerable = req.affinity_target.is_none() && req.pinned_worker.is_none(); + let session = session_binding(session_context.as_ref(), is_steerable, false); + let admission = if req.advisory { + SelectionAdmission::Advisory { + request_id: req.selection_id.clone(), + } + } else { + SelectionAdmission::Query { + request_id: req.selection_id.clone(), + } + }; + let selected = self + .select_or_reject(SelectionOperation { key: RoutingPartitionId::new(req.model_name, req.routing_group), - selection_id: req.selection_id, - prompt: req.prompt, + prompt: req.prompt.view(), router_config_override: req.router_config_override, expected_output_tokens: req.expected_output_tokens, priority_jump: req.priority_jump.unwrap_or_default(), strict_priority: req.strict_priority.unwrap_or(0), policy_class, session_context, + session, affinity_target: req.affinity_target, pinned_worker: req.pinned_worker, allowed_worker_ids: req.allowed_worker_ids, routing_constraints: req.routing_constraints, - advisory: req.advisory, - }, - false, - ) - .await + admission, + }) + .await?; + Ok(self.select_response(selected, req.selection_id)) } pub async fn select_and_reserve( @@ -1053,36 +1053,109 @@ impl SelectionCore { let selection_id = req .selection_id .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); - self.schedule_selection( - SelectionOperation { + let is_steerable = req.affinity_target.is_none() && req.pinned_worker.is_none(); + let session = session_binding(session_context.as_ref(), is_steerable, true); + let selected = self + .select_or_reject(SelectionOperation { key: RoutingPartitionId::new(req.model_name, req.routing_group), - selection_id: Some(selection_id), - prompt: req.prompt, + prompt: req.prompt.view(), router_config_override: req.router_config_override, expected_output_tokens: req.expected_output_tokens, priority_jump: req.priority_jump.unwrap_or_default(), strict_priority: req.strict_priority.unwrap_or(0), policy_class, session_context, + session, affinity_target: req.affinity_target, pinned_worker: req.pinned_worker, allowed_worker_ids: req.allowed_worker_ids, routing_constraints: req.routing_constraints, - advisory: false, - }, - true, - ) - .await + admission: SelectionAdmission::Book { + selection_id: selection_id.clone(), + }, + }) + .await?; + Ok(self.select_response(selected, Some(selection_id))) } - async fn schedule_selection( + /// The wire shape of a selection; a queue rejection keeps its 503 body. + async fn select_or_reject( &self, - operation: SelectionOperation, - book: bool, - ) -> Result { - let SelectionOperation { + operation: SelectionOperation<'_>, + ) -> Result { + match self.run_selection(operation).await? { + SelectionOutcome::Selected(selected) => Ok(selected), + SelectionOutcome::QueueRejected { rejection } => Err(SelectionError::Scheduler( + KvSchedulerError::QueueRejected(rejection), + )), + } + } + + fn select_response(&self, selected: Selected, selection_id: Option) -> SelectResponse { + let Selected { key, + response, + advisory_load, + total_kv_blocks, + endpoint, + block_size, + isl_tokens, + sequence_hashes, + track_prefill_tokens, + effective_prefill_tokens, + kv_hint, + } = selected; + let booked = sequence_hashes.is_some(); + let potential_decode_blocks = response.potential_decode_blocks as u64; + let decode_busy = self + .kv_router_config + .conditional_disagg_decode_busy_threshold + .zip(total_kv_blocks) + .map(|(threshold, total_kv_blocks)| { + potential_decode_blocks as f64 > threshold * total_kv_blocks as f64 + }); + let worker_load = advisory_load.map(|load| SelectionWorkerLoad { + active_prefill_tokens: load.active_prefill_tokens, + prefill_token_capacity: load.prefill_token_capacity, + total_kv_blocks, + prefill_busy: self + .kv_router_config + .conditional_disagg_prefill_busy_threshold + .map(|threshold| load.prefill_load_exceeds(threshold)), + }); + SelectResponse { selection_id, + sequence_hashes: sequence_hashes + .map(|hashes| hashes.into_iter().map(|hash| hash as i64).collect()), + isl_tokens: booked.then_some(isl_tokens), + track_prefill_tokens: booked.then_some(track_prefill_tokens), + model_name: key.model_name, + routing_group: key.routing_group, + worker_id: response.best_worker.worker_id, + dp_rank: response.best_worker.dp_rank, + endpoint, + block_size, + overlap: MooncakeOverlapSummary::from_selected_worker_tiers( + &response.selected_worker_tiers, + block_size, + ), + effective_prefill_tokens, + potential_decode_blocks, + decode_busy, + worker_load, + kv_hint, + } + } + + /// Run one selection: resolve the session, look the prompt up, schedule, + /// and (for `Book`) install the reservation. Every host's selection goes + /// through here. + pub async fn run_selection( + &self, + operation: SelectionOperation<'_>, + ) -> Result { + let SelectionOperation { + key, prompt, router_config_override, expected_output_tokens, @@ -1090,35 +1163,38 @@ impl SelectionCore { strict_priority, policy_class, session_context, + session, affinity_target, pinned_worker, allowed_worker_ids, routing_constraints, - advisory, + admission, } = operation; self.ensure_running()?; let entry = self.ready_entry(&key)?; - let claim = match selection_id.as_deref() { - Some(selection_id) if book => Some(self.claim_reservation(selection_id, &key)?), + let book = admission.is_booking(); + let claim = match &admission { + SelectionAdmission::Book { selection_id } => { + Some(self.claim_reservation(selection_id, &key)?) + } _ => None, }; // Session stickiness: a bound session steers selection (exclusive for - // the default selector); a new session is bound to the worker chosen. - // An explicit affinity target or pin from the caller wins. + // the default selector); a new session is bound to the worker booked. let table = entry.affinity.get(); - let session_id = table - .and(session_context.as_ref()) - .filter(|_| affinity_target.is_none() && pinned_worker.is_none()) - .map(|context| context.session_id().to_string()); let mut affinity_hold = None; - let affinity_target = match (session_id.as_deref(), table) { - (Some(session_id), Some(table)) if book => { + let managed_session = match (&session, table) { + (SessionBinding::Managed { session_id }, Some(table)) => Some((table, session_id)), + _ => None, + }; + let affinity_target = match (&session, table) { + (SessionBinding::Managed { session_id }, Some(table)) => { affinity_hold = self.hold_session(table, session_id, &key).await?; affinity_hold.as_ref().and_then(Hold::target) } - (Some(session_id), Some(table)) => table + (SessionBinding::Query { session_id }, Some(table)) => table .query_target(session_id, None) .map_err(affinity_error)? .map(|target| WorkerAffinityTarget::new(target.worker_id, target.dp_rank)), @@ -1149,52 +1225,57 @@ impl SelectionCore { .await?; // The queue lease frees a booking whose response is never consumed // (the caller dropped this future after the actor booked). - let mode = if book { - ScheduleMode::TrackedWithLifecycle { - request_id: selection_id.clone().ok_or_else(|| { - SelectionError::Internal( - "booked selection did not include a selection ID".to_string(), - ) - })?, - } - } else { - ScheduleMode::QueryOnly { + let mode = match &admission { + SelectionAdmission::Book { selection_id } => ScheduleMode::TrackedWithLifecycle { request_id: selection_id.clone(), - } + }, + SelectionAdmission::Query { request_id } => ScheduleMode::QueryOnly { + request_id: request_id.clone(), + }, + SelectionAdmission::Advisory { request_id } => ScheduleMode::QueryOnly { + request_id: request_id.clone(), + }, }; let track_prefill_tokens = router_config_override .as_ref() .and_then(|cfg| cfg.track_prefill_tokens) .unwrap_or(self.kv_router_config.router_track_prefill_tokens); - // `select` (book == false) with a selection_id caches the booking inputs - // so a follow-up `create_reservation` can replay them by that id. - let cached_inputs = (!book).then(|| selection_id.clone()).flatten().map(|id| { - ( - id, + // An unbooked selection with an id caches the booking inputs so a + // follow-up `create_reservation` can replay them by that id. + let cached_inputs = match &admission { + SelectionAdmission::Query { + request_id: Some(id), + } + | SelectionAdmission::Advisory { + request_id: Some(id), + } => Some(( + id.clone(), sequence_hashes.clone(), - prompt.lora_name.clone(), + prompt.lora_name.map(str::to_string), track_prefill_tokens, - ) - }); + match &session { + SessionBinding::Query { session_id } => Some(session_id.clone()), + _ => None, + }, + )), + _ => None, + }; let allowed_worker_ids = match self.host.eligibility.lora_worker_filter.as_deref() { Some(filter) => narrow_allowed_worker_ids_by_lora( filter, - prompt.lora_name.as_deref(), + prompt.lora_name, allowed_worker_ids, pinned_worker.as_ref(), || self.catalog.schedulable_worker_ids_for_key(&key), ), None => allowed_worker_ids, }; - let response_sequence_hashes = - book.then(|| sequence_hashes.iter().map(|hash| *hash as i64).collect()); - let response_isl_tokens = book.then_some(isl_tokens); - let response_track_prefill_tokens = book.then_some(track_prefill_tokens); // Bookings (now, or later via the pending-selection cache) are recorded // into an approximate indexer; keep the public hashes for that. let routing_hashes = (entry.indexer.records_routing_decisions() && (book || cached_inputs.is_some())) .then(|| block_hashes.clone()); + let booked_sequence_hashes = book.then(|| sequence_hashes.clone()); let schedule_request = ScheduleRequest { mode, token_seq: Some(sequence_hashes), @@ -1204,7 +1285,7 @@ impl SelectionCore { kv_transfer_candidates, retain_kv_transfer_chain, router_config_override, - lora_name: prompt.lora_name, + lora_name: prompt.lora_name.map(str::to_string), priority_jump, strict_priority, policy_class, @@ -1218,13 +1299,13 @@ impl SelectionCore { }; // `lease` guards the booking until it is installed below: any early // return or drop before then frees it. - let (response, advisory_load, lease) = tokio::select! { + let scheduled = tokio::select! { biased; _ = self.cancel_token.cancelled() => { return Err(SelectionError::Scheduler(KvSchedulerError::SubscriberShutdown)); } result = async { - if advisory { + if matches!(admission, SelectionAdmission::Advisory { .. }) { entry .scheduler .select_without_admission(schedule_request) @@ -1239,7 +1320,14 @@ impl SelectionCore { .await .map(|(admitted, lease)| (admitted.response, None, lease)) } - } => result?, + } => result, + }; + let (response, advisory_load, lease) = match scheduled { + Ok(scheduled) => scheduled, + Err(KvSchedulerError::QueueRejected(rejection)) => { + return Ok(SelectionOutcome::QueueRejected { rejection }); + } + Err(error) => return Err(error.into()), }; let Some(endpoint) = self .catalog @@ -1250,26 +1338,13 @@ impl SelectionCore { response.best_worker.worker_id ))); }; - let overlap = MooncakeOverlapSummary::from_selected_worker_tiers( - &response.selected_worker_tiers, - entry.block_size, - ); - let effective_prefill = effective_prefill_tokens(isl_tokens, response.cached_tokens); - let potential_decode_blocks = response.potential_decode_blocks as u64; let total_kv_blocks = advisory_load .and_then(|load| load.total_kv_blocks.map(|blocks| blocks as u64)) .or_else(|| { self.catalog .total_kv_blocks(response.best_worker.worker_id, &key) }); - let decode_busy = self - .kv_router_config - .conditional_disagg_decode_busy_threshold - .zip(total_kv_blocks) - .map(|(threshold, total_kv_blocks)| { - potential_decode_blocks as f64 > threshold * total_kv_blocks as f64 - }); let kv_hint = if retain_kv_transfer_chain { transfer_hint_for_selection( &self.catalog.scheduler_configs_for_key(&key), @@ -1279,23 +1354,17 @@ impl SelectionCore { ) .map(|payload| { KvHint::new( - selection_id.as_deref().unwrap_or_default(), + admission.request_id().unwrap_or_default(), vec![KvHintAction::fetch("a1", payload)], ) }) } else { None }; - let worker_load = advisory_load.map(|load| SelectionWorkerLoad { - active_prefill_tokens: load.active_prefill_tokens, - prefill_token_capacity: load.prefill_token_capacity, - total_kv_blocks, - prefill_busy: self - .kv_router_config - .conditional_disagg_prefill_busy_threshold - .map(|threshold| load.prefill_load_exceeds(threshold)), - }); + // The routing hashes go to exactly one of: the booking recorded now, or + // the cached inputs a later replay records. + let mut routing_hashes = routing_hashes; if let Some(claim) = claim { let Some(lease) = lease else { return Err(SelectionError::Internal( @@ -1304,14 +1373,13 @@ impl SelectionCore { }; // A rejected affinity commit returns while the lease is still armed, // so the booking is freed and nothing below is recorded. - let affinity_lease = match (affinity_hold, session_id.as_deref(), table) { - (Some(hold), Some(session_id), Some(table)) => Some( - self.commit_session(table, hold, session_id, response.best_worker, &key) - .await?, - ), + let affinity_lease = match (affinity_hold, managed_session) { + (Some(hold), Some((table, session_id))) => { + self.commit_session(table, hold, session_id, response.best_worker, &key)? + } _ => None, }; - if let Some(hashes) = routing_hashes.clone() { + if let Some(hashes) = routing_hashes.take() { self.record_routing_decision(&entry, response.best_worker, hashes) .await; } @@ -1322,7 +1390,9 @@ impl SelectionCore { }); } - if let Some((cache_id, sequence_hashes, lora_name, track_prefill_tokens)) = cached_inputs { + if let Some((cache_id, sequence_hashes, lora_name, track_prefill_tokens, session_id)) = + cached_inputs + { self.selection_cache.insert( cache_id, PendingSelection { @@ -1335,29 +1405,24 @@ impl SelectionCore { track_prefill_tokens, lora_name, routing_hashes, + session_id, }, Instant::now(), ); } - - Ok(SelectResponse { - selection_id, - sequence_hashes: response_sequence_hashes, - isl_tokens: response_isl_tokens, - track_prefill_tokens: response_track_prefill_tokens, - model_name: key.model_name, - routing_group: key.routing_group, - worker_id: response.best_worker.worker_id, - dp_rank: response.best_worker.dp_rank, + Ok(SelectionOutcome::Selected(Selected { + key, + response, + advisory_load, + total_kv_blocks, endpoint, block_size: entry.block_size, - overlap, + isl_tokens, + sequence_hashes: booked_sequence_hashes, + track_prefill_tokens, effective_prefill_tokens: effective_prefill, - potential_decode_blocks, - decode_busy, - worker_load, kv_hint, - }) + })) } pub async fn create_reservation( @@ -1388,7 +1453,15 @@ impl SelectionCore { req.selection_id ))); }; - let response = self.book_cached_selection(pending, &req).await?; + let response = match self.book_cached_selection(pending, &req).await { + Ok(response) => response, + // The session moved: the cached worker can never be booked for it. + Err(SelectionError::BadRequest(message)) => { + self.selection_cache.discard(&key, &req.selection_id); + return Err(SelectionError::BadRequest(message)); + } + Err(error) => return Err(error), + }; self.selection_cache .remove(&key, &req.selection_id, generation); Ok(response) @@ -1416,6 +1489,7 @@ impl SelectionCore { track_prefill_tokens, lora_name: pending.lora_name.clone(), routing_hashes: pending.routing_hashes.clone(), + session_id: pending.session_id.clone(), }, ) .await @@ -1470,7 +1544,7 @@ impl SelectionCore { req: ReservationRequest, ) -> Result { let entry = self.ready_entry(&key)?; - let normalized = req.prompt.normalize_for_reservation( + let normalized = req.prompt.view().normalize_for_reservation( entry.is_eagle, TrackingHashInput { context: &self.tracking_hash, @@ -1504,6 +1578,7 @@ impl SelectionCore { let routing_hashes = can_record .then(|| { req.prompt + .view() .block_hashes_for_indexer(entry.block_size, entry.is_eagle) }) .transpose()?; @@ -1521,6 +1596,7 @@ impl SelectionCore { track_prefill_tokens, lora_name: req.prompt.lora_name, routing_hashes, + session_id: None, }, ) .await @@ -1545,9 +1621,19 @@ impl SelectionCore { track_prefill_tokens, lora_name, routing_hashes, + session_id, } = booking; let claim = self.claim_reservation(&selection_id, &key)?; + // Hold the session before booking: holding after would wait on an + // initializer that may itself be waiting for this worker's capacity. + let session = match (session_id.as_deref(), entry.affinity.get()) { + (Some(session_id), Some(table)) => self + .hold_session(table, session_id, &key) + .await? + .map(|hold| (table, session_id, hold)), + _ => None, + }; // Strict booking: never lazily recreate a worker/rank removed since the // reservation was resolved. The lease frees the booking if this future // is dropped before `install`. @@ -1562,13 +1648,19 @@ impl SelectionCore { worker, lora_name, })?; + let affinity_lease = match session { + Some((table, session_id, hold)) => { + self.commit_session(table, hold, session_id, worker, &key)? + } + None => None, + }; if let Some(hashes) = routing_hashes { self.record_routing_decision(&entry, worker, hashes).await; } claim.install(Reservation { partition: key.clone(), booking: Some(lease.commit().ok_or_else(missing_booking)?), - _affinity_lease: None, + _affinity_lease: affinity_lease, }); Ok(ReservationResponse { @@ -1650,29 +1742,36 @@ impl SelectionCore { /// Bind the held session to `dispatched`. A `Hard` rejection whose bound /// worker departed after [`Self::hold_session`] checked it is not a client - /// fault: the binding is re-initialized on the dispatched worker instead. - async fn commit_session( + /// fault: the session is re-initialized on the dispatched worker instead. + /// This runs with the booking held, so it never waits on another request's + /// initialization (which may itself be queued behind this booking); when + /// one is in flight the booking stands unbound, `None`. + fn commit_session( &self, table: &SessionAffinity, hold: Hold, session_id: &str, dispatched: WorkerWithDpRank, key: &RoutingPartitionId, - ) -> Result { + ) -> Result, SelectionError> { let bound = hold.target(); let dispatched = WorkerAffinityTarget::new(dispatched.worker_id, Some(dispatched.dp_rank)); match table.commit(hold, dispatched) { - Ok(lease) => Ok(lease), + Ok(lease) => Ok(Some(lease)), Err(error) => { let departed = bound.is_some_and(|target| !self.catalog.is_schedulable(target.worker_id, key)); if !departed { return Err(affinity_error(error)); } - // `commit` invalidated the stale binding; the next hold initializes. - match self.hold_session(table, session_id, key).await? { - Some(hold) => table.commit(hold, dispatched).map_err(affinity_error), - None => Err(affinity_error(error)), + // `commit` invalidated the stale binding. + match table.try_acquire(session_id, None) { + Ok(AcquireStep::Held(hold)) => table + .commit(hold, dispatched) + .map(Some) + .map_err(affinity_error), + Ok(AcquireStep::Wait(_)) | Err(AffinityError::ResourceExhausted(_)) => Ok(None), + Err(error) => Err(affinity_error(error)), } } } @@ -1866,7 +1965,7 @@ impl SelectionCore { let prepared = self .prepare_selection_inputs( &entry, - &req.prompt, + &req.prompt.view(), self.kv_router_config .assume_kv_reuse(req.router_config_override.as_ref()), false, @@ -1894,6 +1993,7 @@ impl SelectionCore { let entry = self.ready_entry(&key)?; let block_hashes = req .prompt + .view() .block_hashes_for_indexer(entry.block_size, entry.is_eagle)?; let num_blocks = block_hashes.len(); let tiered = entry @@ -1922,7 +2022,7 @@ impl SelectionCore { async fn prepare_selection_inputs( &self, entry: &SelectionEntry, - prompt: &PromptRequest, + prompt: &PromptView<'_>, assume_kv_reuse: bool, query_shared_cache: bool, retain_kv_transfer_chain: bool, @@ -1959,11 +2059,11 @@ impl SelectionCore { let shared_cache = query_shared_cache .then_some(self.host.cache.shared.as_deref()) .flatten() - .zip(prompt.token_ids.as_deref()); + .zip(prompt.token_ids); let shared_cache_lookup = async { let (shared_cache, tokens) = shared_cache?; match shared_cache - .check_blocks(tokens, entry.block_size, prompt.cache_namespace.as_deref()) + .check_blocks(tokens, entry.block_size, prompt.cache_namespace) .await { Ok(hits) => Some(hits), @@ -2025,6 +2125,24 @@ fn sweep_reservation_index(entries: &SelectionEntries, index: &ReservationIndex) before - index.len() } +/// The core's use of the session table for a request: a session is steered +/// only when the caller did not pin or target a worker explicitly. +fn session_binding( + context: Option<&SessionContext>, + is_steerable: bool, + is_booking: bool, +) -> SessionBinding { + match context.filter(|_| is_steerable) { + Some(context) if is_booking => SessionBinding::Managed { + session_id: context.session_id().to_string(), + }, + Some(context) => SessionBinding::Query { + session_id: context.session_id().to_string(), + }, + None => SessionBinding::None, + } +} + fn missing_booking() -> SelectionError { SelectionError::Internal("booking lease holds no booking".to_string()) } @@ -2132,9 +2250,11 @@ impl Drop for SelectionCore { #[cfg(test)] mod tests { use super::super::affinity::SessionAffinityMode; + use super::super::input::PromptRequest; use super::*; - use crate::protocols::StorageTier; + use crate::protocols::{RoutingConstraints, StorageTier}; use crate::services::indexer::backend::test_util::store_event; + use std::collections::HashSet; use std::sync::atomic::{AtomicBool, Ordering}; use std::thread::sleep; use std::time::Duration; @@ -4181,6 +4301,80 @@ mod tests { ); } + fn replay_reservation(selection_id: &str) -> ReservationRequest { + ReservationRequest { + model_name: "model".to_string(), + routing_group: "default".to_string(), + selection_id: selection_id.to_string(), + worker_id: None, + dp_rank: None, + prompt: PromptRequest::default(), + router_config_override: None, + expected_output_tokens: None, + effective_prefill_tokens: None, + track_prefill_tokens: None, + } + } + + #[tokio::test] + async fn two_phase_reservation_binds_the_session() { + let core = core_with_session_affinity(); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + core.upsert_worker(worker(2)).await.expect("worker upsert"); + let mut request = select_request(); + request.selection_id = Some("pending".to_string()); + request.session_id = Some("s".to_string()); + let selected = core.select(request).await.expect("select"); + assert_eq!(bound_worker(&core, "s"), None, "select alone binds nothing"); + + core.create_reservation(replay_reservation("pending")) + .await + .expect("replay booking"); + assert_eq!(bound_worker(&core, "s"), Some(selected.worker_id)); + core.free_reservation("pending").await.expect("free"); + + // The binding steers the session's next selection. + let mut request = select_request(); + request.session_id = Some("s".to_string()); + assert_eq!( + core.select(request).await.expect("select").worker_id, + selected.worker_id + ); + } + + #[tokio::test] + async fn two_phase_replay_rejects_a_worker_the_session_left() { + let core = core_with_session_affinity(); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + core.upsert_worker(worker(2)).await.expect("worker upsert"); + let mut request = select_request(); + request.selection_id = Some("pending".to_string()); + request.session_id = Some("s".to_string()); + let cached = core.select(request).await.expect("select"); + // Another request binds the session elsewhere before the replay. + let other = if cached.worker_id == 1 { 2 } else { 1 }; + let mut request = session_reservation("r1", "s"); + request.allowed_worker_ids = Some(HashSet::from([other])); + core.select_and_reserve(request) + .await + .expect("bind elsewhere"); + core.free_reservation("r1").await.expect("free"); + + let err = core + .create_reservation(replay_reservation("pending")) + .await + .expect_err("hard affinity rejects the stale cached worker"); + assert!(matches!(err, SelectionError::BadRequest(_)), "{err:?}"); + let entry = core + .entry(&RoutingPartitionId::new("model", "default")) + .expect("entry"); + wait_until("rejected booking release", || { + !entry.scheduler.has_request("pending") + }) + .await; + assert!(core.reservation_index.read().is_empty()); + } + #[tokio::test] async fn hard_mode_rejects_dispatch_away_from_a_live_binding() { let core = core_with_session_affinity(); diff --git a/lib/kv-router/src/services/selection/core/operation.rs b/lib/kv-router/src/services/selection/core/operation.rs new file mode 100644 index 000000000000..4ba513281061 --- /dev/null +++ b/lib/kv-router/src/services/selection/core/operation.rs @@ -0,0 +1,100 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! The one selection operation every host runs. Wire handlers and embedding +//! hosts build a [`SelectionOperation`] and consume a [`SelectionOutcome`]. + +use std::collections::HashSet; + +use dynamo_tokens::SequenceHash; + +use crate::identity::RoutingPartitionId; +use crate::kv_hints::KvHint; +use crate::protocols::{RoutingConstraints, WorkerAffinityTarget, WorkerId, WorkerWithDpRank}; +use crate::scheduling::config::RouterConfigOverride; +use crate::scheduling::{AdvisoryWorkerLoad, QueueRejection, SchedulingResponse, SessionContext}; + +use super::super::input::PromptView; + +/// Every input to one selection. Fields are independent: `session_context` +/// feeds worker selection, `session` says what the core does with the session +/// table, `affinity_target` and `pinned_worker` are explicit steering. +pub struct SelectionOperation<'a> { + pub key: RoutingPartitionId, + pub prompt: PromptView<'a>, + pub router_config_override: Option, + pub expected_output_tokens: Option, + pub priority_jump: f64, + pub strict_priority: u32, + pub policy_class: Option, + pub session_context: Option, + pub session: SessionBinding, + pub affinity_target: Option, + pub pinned_worker: Option, + pub allowed_worker_ids: Option>, + pub routing_constraints: RoutingConstraints, + pub admission: SelectionAdmission, +} + +pub enum SelectionAdmission { + /// Queue admission without a booking; a `request_id` caches the inputs for + /// a later `create_reservation` replay. + Query { request_id: Option }, + /// Queue admission with a booking recorded under `selection_id`. + Book { selection_id: String }, + /// Skip queue admission and report the chosen worker's load; a + /// `request_id` caches the inputs like `Query`. + Advisory { request_id: Option }, +} + +impl SelectionAdmission { + pub fn request_id(&self) -> Option<&str> { + match self { + Self::Query { request_id } | Self::Advisory { request_id } => request_id.as_deref(), + Self::Book { selection_id } => Some(selection_id), + } + } + + pub fn is_booking(&self) -> bool { + matches!(self, Self::Book { .. }) + } +} + +/// What the core does with the partition's session table for this request. +pub enum SessionBinding { + None, + /// Hold the session, steer to its worker, and bind it to the worker booked. + Managed { + session_id: String, + }, + /// Steer to the session's worker without holding or binding it. + Query { + session_id: String, + }, +} + +// Every caller matches this immediately; boxing the common variant to shrink +// the rare one would put an allocation on the hot path. +#[allow(clippy::large_enum_variant)] +pub enum SelectionOutcome { + Selected(Selected), + QueueRejected { rejection: QueueRejection }, +} + +/// A completed selection. For `Book`, the reservation is already installed. +#[must_use] +pub struct Selected { + pub key: RoutingPartitionId, + pub response: SchedulingResponse, + pub advisory_load: Option, + /// The chosen worker's KV capacity as known when it was selected. + pub total_kv_blocks: Option, + pub endpoint: String, + pub block_size: u32, + pub isl_tokens: usize, + /// The hashes the booking tracks; `Book` only. + pub sequence_hashes: Option>, + pub track_prefill_tokens: bool, + pub effective_prefill_tokens: usize, + pub kv_hint: Option, +} diff --git a/lib/kv-router/src/services/selection/input.rs b/lib/kv-router/src/services/selection/input.rs index aaf688a4bf26..5e7780275434 100644 --- a/lib/kv-router/src/services/selection/input.rs +++ b/lib/kv-router/src/services/selection/input.rs @@ -41,6 +41,40 @@ pub struct PromptRequest { } impl PromptRequest { + /// Borrow the request as the shape selection consumes; hosts that already + /// hold the pieces build a [`PromptView`] directly instead of allocating. + pub fn view(&self) -> PromptView<'_> { + PromptView { + token_ids: self.token_ids.as_deref(), + mm_routing_info: self.mm_routing_info.as_ref(), + block_mm_infos: self.block_mm_infos.as_deref(), + block_hashes: self.block_hashes.as_deref(), + sequence_hashes: self.sequence_hashes.as_deref(), + isl_tokens: self.isl_tokens, + lora_name: self.lora_name.as_deref(), + cache_namespace: self.cache_namespace.as_deref(), + is_eagle: self.is_eagle, + } + } +} + +/// A prompt borrowed for one selection: the same fields as [`PromptRequest`] +/// with the same precedence (multimodal routing tokens, then raw tokens, then +/// the hash-only trio). +#[derive(Debug, Clone, Copy)] +pub struct PromptView<'a> { + pub token_ids: Option<&'a [u32]>, + pub mm_routing_info: Option<&'a MmRoutingInfoRequest>, + pub block_mm_infos: Option<&'a [Option]>, + pub block_hashes: Option<&'a [i64]>, + pub sequence_hashes: Option<&'a [i64]>, + pub isl_tokens: Option, + pub lora_name: Option<&'a str>, + pub cache_namespace: Option<&'a str>, + pub is_eagle: Option, +} + +impl PromptView<'_> { pub(super) fn normalize_for_selection( &self, default_is_eagle: bool, @@ -49,18 +83,18 @@ impl PromptRequest { if let Some((token_ids, block_mm_infos)) = self.routing_tokens_and_mm_infos() { return Ok(normalize_tokens_for_selection( token_ids, - self.lora_name.as_deref(), - self.cache_namespace.as_deref(), + self.lora_name, + self.cache_namespace, block_mm_infos, self.is_eagle.unwrap_or(default_is_eagle), tracking, )); } - let block_hashes = self.block_hashes.as_ref().ok_or_else(|| { + let block_hashes = self.block_hashes.ok_or_else(|| { SelectionError::BadRequest("block_hashes is required without token_ids".to_string()) })?; - let sequence_hashes = self.sequence_hashes.as_ref().ok_or_else(|| { + let sequence_hashes = self.sequence_hashes.ok_or_else(|| { SelectionError::BadRequest("sequence_hashes is required without token_ids".to_string()) })?; let isl_tokens = self.isl_tokens.ok_or_else(|| { @@ -77,15 +111,15 @@ impl PromptRequest { if let Some((token_ids, block_mm_infos)) = self.routing_tokens_and_mm_infos() { return Ok(normalize_tokens_for_reservation( token_ids, - self.lora_name.as_deref(), - self.cache_namespace.as_deref(), + self.lora_name, + self.cache_namespace, block_mm_infos, self.is_eagle.unwrap_or(default_is_eagle), tracking, )); } - let sequence_hashes = self.sequence_hashes.as_ref().ok_or_else(|| { + let sequence_hashes = self.sequence_hashes.ok_or_else(|| { SelectionError::BadRequest("sequence_hashes is required without token_ids".to_string()) })?; if self.isl_tokens.is_none() { @@ -110,17 +144,17 @@ impl PromptRequest { block_size, BlockHashOptions { block_mm_infos, - lora_name: self.lora_name.as_deref(), - cache_namespace: self.cache_namespace.as_deref(), + lora_name: self.lora_name, + cache_namespace: self.cache_namespace, is_eagle: Some(self.is_eagle.unwrap_or(default_is_eagle)), }, )); } - let block_hashes = self.block_hashes.as_ref().ok_or_else(|| { + let block_hashes = self.block_hashes.ok_or_else(|| { SelectionError::BadRequest("block_hashes is required without token_ids".to_string()) })?; - let sequence_hashes = self.sequence_hashes.as_ref().ok_or_else(|| { + let sequence_hashes = self.sequence_hashes.ok_or_else(|| { SelectionError::BadRequest("sequence_hashes is required without token_ids".to_string()) })?; let isl_tokens = self.isl_tokens.ok_or_else(|| { @@ -130,7 +164,7 @@ impl PromptRequest { } fn routing_tokens_and_mm_infos(&self) -> Option> { - if let Some(mm_routing_info) = &self.mm_routing_info + if let Some(mm_routing_info) = self.mm_routing_info && !mm_routing_info.routing_token_ids.is_empty() { return Some(( @@ -140,8 +174,7 @@ impl PromptRequest { } self.token_ids - .as_deref() - .map(|token_ids| (token_ids, self.block_mm_infos.as_deref())) + .map(|token_ids| (token_ids, self.block_mm_infos)) } } diff --git a/lib/kv-router/src/services/selection/mod.rs b/lib/kv-router/src/services/selection/mod.rs index 48504e82bafb..099bcdd172d5 100644 --- a/lib/kv-router/src/services/selection/mod.rs +++ b/lib/kv-router/src/services/selection/mod.rs @@ -28,12 +28,13 @@ pub use crate::services::common::replica_sync::{ ScopedSequencePublisher, }; pub use core::{ - HostCache, HostEligibility, HostLoad, HostReplication, HostTelemetry, KvIndexSource, - SelectionCore, SelectionHost, SelectionPartition, SelectionScheduler, SelectionServiceConfig, + HostCache, HostEligibility, HostLoad, HostReplication, HostTelemetry, KvIndexSource, Selected, + SelectionAdmission, SelectionCore, SelectionHost, SelectionOperation, SelectionOutcome, + SelectionPartition, SelectionScheduler, SelectionServiceConfig, SessionBinding, }; pub use error::SelectionError; pub use ingress::{KvEventIngress, ZmqDirectIngress}; -pub use input::PromptRequest; +pub use input::{PromptRequest, PromptView}; pub use membership::{CatalogObserver, CatalogReconciler, WorkerCatalogSource}; pub use pending::SelectionCacheConfig; pub use policy_registry::{ diff --git a/lib/kv-router/src/services/selection/pending.rs b/lib/kv-router/src/services/selection/pending.rs index 7a3a70b1d2a5..c1ce85de6b28 100644 --- a/lib/kv-router/src/services/selection/pending.rs +++ b/lib/kv-router/src/services/selection/pending.rs @@ -59,6 +59,8 @@ pub(super) struct PendingSelection { /// Public block hashes retained only when the partition indexer records /// routing decisions, so the replayed booking can be recorded too. pub routing_hashes: Option>, + /// The session the selection steered by, bound when the booking replays. + pub session_id: Option, } struct Entry { @@ -81,6 +83,7 @@ fn entry_bytes(key: &CacheKey, selection: &PendingSelection) -> usize { + key.0.model_name.len() + key.0.routing_group.len() + selection.lora_name.as_deref().map_or(0, str::len) + + selection.session_id.as_deref().map_or(0, str::len) } struct State { @@ -274,6 +277,7 @@ mod tests { track_prefill_tokens: true, lora_name: None, routing_hashes: None, + session_id: None, } } diff --git a/lib/kv-router/src/services/selection/tests.rs b/lib/kv-router/src/services/selection/tests.rs index bd1c37a43ea4..3a89f3bab511 100644 --- a/lib/kv-router/src/services/selection/tests.rs +++ b/lib/kv-router/src/services/selection/tests.rs @@ -127,6 +127,7 @@ fn normalize_prompt(request: &PromptRequest) -> super::input::NormalizedPrompt { let config = test_config(); let context = TrackingHashContext::from_config(&config).unwrap(); request + .view() .normalize_for_selection( false, TrackingHashInput { @@ -566,6 +567,7 @@ fn keyed_prompt_tracking_leaves_indexer_hashes_public() { .unwrap(); let normalized = request + .view() .normalize_for_selection( false, TrackingHashInput { @@ -604,6 +606,7 @@ fn disabled_kv_reuse_keeps_public_indexer_hashes_and_randomizes_tracking() { .unwrap(); let normalize = || { request + .view() .normalize_for_selection( false, TrackingHashInput { @@ -650,6 +653,7 @@ fn keyed_reservation_hashes_directly_from_tokens() { }; let normalized = request + .view() .normalize_for_reservation( false, TrackingHashInput { @@ -697,6 +701,7 @@ fn keyed_hash_only_inputs_remain_trusted_for_selection_and_reservation() { }; let selection = request + .view() .normalize_for_selection( false, TrackingHashInput { @@ -707,6 +712,7 @@ fn keyed_hash_only_inputs_remain_trusted_for_selection_and_reservation() { ) .unwrap(); let reservation = request + .view() .normalize_for_reservation( false, TrackingHashInput { @@ -750,6 +756,7 @@ fn randomized_reservation_uses_canonical_complete_block_count() { })) .unwrap(); let normalized = request + .view() .normalize_for_reservation( false, TrackingHashInput { @@ -2005,3 +2012,17 @@ async fn hash_path_validation_returns_bad_request() { .await; assert_eq!(response.status(), StatusCode::BAD_REQUEST); } + +#[test] +fn empty_raw_tokens_win_over_supplied_hashes() { + let request: PromptRequest = serde_json::from_value(serde_json::json!({ + "token_ids": [], + "block_hashes": [11, 12], + "sequence_hashes": [21, 22], + "isl_tokens": 8 + })) + .unwrap(); + let normalized = normalize_prompt(&request); + assert!(normalized.block_hashes.is_empty()); + assert_eq!(normalized.isl_tokens, 0); +} From d6b1c8f6cc6d2822454124e5f62590ce074aca70 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Thu, 10 Sep 2026 13:50:28 -0700 Subject: [PATCH 054/102] perf(router): trim the core hint and readiness hot paths and carry the KV event source mode Worker registrations carry kv_event_source_mode, and the core excludes state_agent_v2 workers as router-hint sources the way the frontend already does. The hint path borrows the partition's published worker snapshot instead of copying every catalog config per request, and readiness no longer counts the whole catalog on the success path. A criterion bench sweeps catalog size and hint capability for select_and_reserve plus free. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Ishan Dhanani --- .../router/standalone-selection.md | 4 + lib/kv-router/Cargo.toml | 5 + lib/kv-router/benches/selection_core.rs | 175 ++++++++++++++++++ .../src/services/selection/core/mod.rs | 162 ++++++++-------- lib/kv-router/src/services/selection/types.rs | 16 ++ lib/llm/src/kv_router/embedded.rs | 1 + 6 files changed, 283 insertions(+), 80 deletions(-) create mode 100644 lib/kv-router/benches/selection_core.rs diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-selection.md b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-selection.md index 91b108654c19..84f0d1d2aee6 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-selection.md +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-selection.md @@ -159,6 +159,10 @@ rank: } ``` +A worker whose KV events come from a state agent sets +`"kv_event_source_mode": "state_agent_v2"`; it can receive hints but is never +chosen as a hint source. + `PATCH /workers/{worker_id}` updates supplied fields, `DELETE /workers/{worker_id}` removes the worker, and `GET /workers` lists catalog state. `model_name` and `routing_group` scope selection, indexer, and load state; diff --git a/lib/kv-router/Cargo.toml b/lib/kv-router/Cargo.toml index 154d575a3ea8..deeb18758059 100644 --- a/lib/kv-router/Cargo.toml +++ b/lib/kv-router/Cargo.toml @@ -98,3 +98,8 @@ harness = false [[bench]] name = "worker_selection" harness = false + +[[bench]] +name = "selection_core" +harness = false +required-features = ["standalone-selection"] diff --git a/lib/kv-router/benches/selection_core.rs b/lib/kv-router/benches/selection_core.rs new file mode 100644 index 000000000000..b472c21b1a99 --- /dev/null +++ b/lib/kv-router/benches/selection_core.rs @@ -0,0 +1,175 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Per-request cost of the shared selection core: one booked selection and its +//! release, swept over catalog size and router-hint capability. + +use std::collections::HashMap; +use std::time::Duration; + +use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; +use dynamo_kv_router::KvRouterConfig; +use dynamo_kv_router::identity::RoutingPartitionId; +use dynamo_kv_router::indexer::KvIndexerInterface; +use dynamo_kv_router::protocols::{ + BlockHashOptions, ExternalSequenceBlockHash, KvCacheEvent, KvCacheEventData, KvCacheStoreData, + KvCacheStoredBlockData, RouterEvent, RoutingConstraints, StorageTier, WorkerWithDpRank, + compute_block_hash_for_seq, compute_seq_hash_for_block, +}; +use dynamo_kv_router::services::indexer::backend::Indexer; +use dynamo_kv_router::services::selection::{ + PromptRequest, SelectAndReserveRequest, SelectionCacheConfig, SelectionCore, WorkerRequest, +}; +use tokio_util::sync::CancellationToken; + +const BLOCK_SIZE: u32 = 16; + +fn prompt_tokens() -> Vec { + (1..=256).collect() +} + +/// `hints` makes worker 1 hold the whole prompt so every booking on another +/// worker computes a fetch hint from it. +fn core_with_workers( + runtime: &tokio::runtime::Runtime, + workers: u64, + hints: bool, +) -> SelectionCore { + let config = KvRouterConfig { + use_kv_events: true, + router_queue_threshold: None, + ..Default::default() + }; + let core = SelectionCore::try_new_local( + config, + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ) + .expect("core"); + runtime.block_on(async { + for worker_id in 1..=workers { + let mut request = WorkerRequest { + worker_id, + endpoint: Some(format!("http://worker-{worker_id}:8000")), + kv_events_endpoint: Some(format!("tcp://127.0.0.1:{}", 40_000 + worker_id)), + block_size: Some(BLOCK_SIZE), + max_num_batched_tokens: Some(8192), + ..WorkerRequest::default() + }; + if hints { + request.router_hint_worker_type = Some("decode".to_string()); + request.router_hint_source_control_endpoints = + HashMap::from([(0, format!("tcp://worker-{worker_id}:9000"))]); + } + core.upsert_worker(request).await.expect("upsert"); + } + if hints { + seed_prefix_on_worker_one(&core).await; + } + }); + core +} + +async fn seed_prefix_on_worker_one(core: &SelectionCore) { + let key = RoutingPartitionId::new("default", "default"); + let local_hashes = + compute_block_hash_for_seq(&prompt_tokens(), BLOCK_SIZE, BlockHashOptions::default()); + let sequence_hashes = compute_seq_hash_for_block(&local_hashes); + let blocks = local_hashes + .iter() + .zip(sequence_hashes.iter()) + .map(|(&tokens_hash, &sequence_hash)| KvCacheStoredBlockData { + block_hash: ExternalSequenceBlockHash(sequence_hash), + tokens_hash, + mm_extra_info: None, + }) + .collect(); + let indexer = core.partition(&key).expect("partition").indexer().clone(); + indexer + .apply_event_routed(RouterEvent::with_storage_tier( + 1, + KvCacheEvent { + event_id: 1, + data: KvCacheEventData::Stored(KvCacheStoreData { + parent_hash: None, + start_position: None, + blocks, + }), + dp_rank: 0, + }, + StorageTier::Device, + )) + .await + .expect("seed index"); + if let Indexer::Single { primary, .. } = &indexer { + primary.flush().await; + } + let mut request = reserve_request("seed".to_string()); + request.pinned_worker = Some(WorkerWithDpRank::new(2, 0)); + let response = core.select_and_reserve(request).await.expect("reserve"); + core.free_reservation("seed").await.expect("free"); + assert!(response.kv_hint.is_some(), "seeded prefix produced no hint"); +} + +fn reserve_request(selection_id: String) -> SelectAndReserveRequest { + SelectAndReserveRequest { + model_name: "default".to_string(), + routing_group: "default".to_string(), + selection_id: Some(selection_id), + prompt: PromptRequest { + token_ids: Some(prompt_tokens()), + ..PromptRequest::default() + }, + router_config_override: None, + expected_output_tokens: None, + priority_jump: None, + strict_priority: None, + session_id: None, + session_context: None, + affinity_target: None, + pinned_worker: None, + allowed_worker_ids: None, + routing_constraints: RoutingConstraints::default(), + } +} + +fn bench_select_and_reserve(c: &mut Criterion) { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("runtime"); + let mut group = c.benchmark_group("selection_core/select_and_reserve_then_free"); + group.measurement_time(Duration::from_secs(5)); + group.throughput(Throughput::Elements(1)); + for &workers in &[8u64, 64, 512] { + for hints in [false, true] { + let core = core_with_workers(&runtime, workers, hints); + let mut sequence = 0u64; + group.bench_with_input( + BenchmarkId::new(if hints { "hints" } else { "no_hints" }, workers), + &core, + |b, core| { + b.iter(|| { + sequence += 1; + let id = format!("bench-{sequence}"); + // `no_hints` sweeps scoring across all workers; `hints` + // pins off the seeded worker so every booking computes a hint. + let mut request = reserve_request(id.clone()); + if hints { + request.pinned_worker = Some(WorkerWithDpRank::new(2, 0)); + } + runtime.block_on(async { + core.select_and_reserve(request).await.expect("reserve"); + core.free_reservation(&id).await.expect("free"); + }); + }); + }, + ); + } + } + group.finish(); +} + +criterion_group!(benches, bench_select_and_reserve); +criterion_main!(benches); diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index f074bba43185..b0aa620fd1bf 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -976,25 +976,24 @@ impl SelectionCore { } fn ready_entry(&self, key: &RoutingPartitionId) -> Result, SelectionError> { - if self.catalog.schedulable_count() == 0 { - return Err(SelectionError::NotReady( - "no schedulable workers are available".to_string(), - )); - } - let Some(entry) = self.entry(key) else { - return Err(SelectionError::NotReady(format!( - "no schedulable workers for {key}" - ))); + return Err(self.not_ready(key)); }; if !self.catalog.has_schedulable_for_key(key) { - return Err(SelectionError::NotReady(format!( - "no schedulable workers for {key}" - ))); + return Err(self.not_ready(key)); } Ok(entry) } + /// Only the failure path pays for the catalog-wide count. + fn not_ready(&self, key: &RoutingPartitionId) -> SelectionError { + if self.catalog.schedulable_count() == 0 { + SelectionError::NotReady("no schedulable workers are available".to_string()) + } else { + SelectionError::NotReady(format!("no schedulable workers for {key}")) + } + } + pub async fn select(&self, req: SelectRequest) -> Result { self.select_with_policy_class(req, None).await } @@ -1345,9 +1344,9 @@ impl SelectionCore { self.catalog .total_kv_blocks(response.best_worker.worker_id, &key) }); - let kv_hint = if retain_kv_transfer_chain { + let kv_hint = if retain_kv_transfer_chain && response.kv_transfer_candidates.is_some() { transfer_hint_for_selection( - &self.catalog.scheduler_configs_for_key(&key), + &entry.workers_tx.borrow(), response.best_worker, response.target_cached_prefix_blocks, response.kv_transfer_candidates.as_ref(), @@ -2174,7 +2173,7 @@ fn spawn_reservation_index_sweep( /// Pick the best router-hint source for `target`: a same-role worker (or /// cache owner) holding a longer root-aligned prefix than the target's own /// `target_cached_prefix_blocks`, with a non-empty control endpoint. Mirrors the -/// frontend `KvRouter::router_hint_for_selection`. +/// frontend `KvRouter::transfer_hint_for_selection`. fn transfer_hint_for_selection( configs: &HashMap, target: WorkerWithDpRank, @@ -2191,14 +2190,15 @@ fn transfer_hint_for_selection( KvTransferCandidateSource::Worker(worker) => { worker != target && configs.get(&worker.worker_id).is_some_and(|config| { - config - .kv_hint_transfer_metadata_for_dp_rank(worker.dp_rank) - .is_some_and(|source_metadata| { - source_metadata.worker_type == target_metadata.worker_type - && source_metadata - .source_control_endpoint - .is_some_and(|endpoint| !endpoint.is_empty()) - }) + config.kv_event_source_mode.as_deref() != Some("state_agent_v2") + && config + .kv_hint_transfer_metadata_for_dp_rank(worker.dp_rank) + .is_some_and(|source_metadata| { + source_metadata.worker_type == target_metadata.worker_type + && source_metadata + .source_control_endpoint + .is_some_and(|endpoint| !endpoint.is_empty()) + }) }) } KvTransferCandidateSource::CacheOwner(owner) => candidates @@ -2290,6 +2290,7 @@ mod tests { kv_transfer_preferred_weight: None, router_hint_worker_type: None, router_hint_source_control_endpoints: HashMap::new(), + kv_event_source_mode: None, } } @@ -2676,10 +2677,13 @@ mod tests { server.abort(); } - #[tokio::test] - async fn booked_selection_attaches_router_hint_from_a_better_source() { + /// Two event-driven workers; worker 1 holds every block of an 8-token + /// prompt. `configure` adjusts each registration before upsert. + async fn hint_fixture( + configure: impl Fn(&mut WorkerRequest), + ) -> (SelectionCore, Arc, Vec) { use crate::indexer::KvIndexerInterface; - use crate::protocols::{BlockHashOptions, StorageTier, compute_block_hash_for_seq}; + use crate::protocols::{BlockHashOptions, compute_block_hash_for_seq}; let core = SelectionCore::try_new_local( test_config(true), @@ -2690,22 +2694,17 @@ mod tests { .expect("valid test config"); for worker_id in [1, 2] { let mut request = worker_with_kv_events(worker_id); - request.router_hint_worker_type = Some("decode".to_string()); - request.router_hint_source_control_endpoints = - HashMap::from([(0, format!("tcp://worker-{worker_id}:9000"))]); + configure(&mut request); core.upsert_worker(request).await.expect("worker upsert"); } - let key = RoutingPartitionId::new("model", "default"); - let entry = core.entry(&key).expect("entry"); - assert!(entry.indexer.supports_kv_transfer_chain_retention()); - - // Worker 1 holds both blocks of the prompt; worker 2 holds nothing. + let entry = core + .entry(&RoutingPartitionId::new("model", "default")) + .expect("entry"); let tokens: Vec = (1..=8).collect(); let hashes: Vec = compute_block_hash_for_seq(&tokens, 4, BlockHashOptions::default()) .into_iter() .map(|hash| hash.0) .collect(); - assert_eq!(hashes.len(), 2); entry .indexer .apply_event_routed(store_event(1, 0, 1, &[], &hashes, StorageTier::Device)) @@ -2714,16 +2713,37 @@ mod tests { if let Indexer::Single { primary, .. } = &entry.indexer { let _ = primary.flush().await; } - let prompt = || PromptRequest { - token_ids: Some(tokens.clone()), + (core, entry, tokens) + } + + fn hint_capable(request: &mut WorkerRequest) { + request.router_hint_worker_type = Some("decode".to_string()); + request.router_hint_source_control_endpoints = + HashMap::from([(0, format!("tcp://worker-{}:9000", request.worker_id))]); + } + + async fn reserve_pinned( + core: &SelectionCore, + selection_id: &str, + tokens: &[u32], + worker_id: WorkerId, + ) -> SelectResponse { + let mut request = reserve_request(selection_id); + request.prompt = PromptRequest { + token_ids: Some(tokens.to_vec()), ..PromptRequest::default() }; + request.pinned_worker = Some(WorkerWithDpRank::new(worker_id, 0)); + core.select_and_reserve(request).await.expect("reserve") + } + + #[tokio::test] + async fn booked_selection_attaches_router_hint_from_a_better_source() { + let (core, entry, tokens) = hint_fixture(hint_capable).await; + assert!(entry.indexer.supports_kv_transfer_chain_retention()); // Booking on worker 2: worker 1 is a same-role source with a longer prefix. - let mut request = reserve_request("to-worker-2"); - request.prompt = prompt(); - request.pinned_worker = Some(WorkerWithDpRank::new(2, 0)); - let response = core.select_and_reserve(request).await.expect("reserve"); + let response = reserve_pinned(&core, "to-worker-2", &tokens, 2).await; let hint = response.kv_hint.expect("router hint for worker 2"); assert_eq!(hint.message_id, "to-worker-2"); assert_eq!(hint.actions[0].action_type, "kv.fetch"); @@ -2734,56 +2754,38 @@ mod tests { assert_eq!(payload.block_hashes.len(), 2); // Booking on worker 1 itself: nothing holds a longer prefix. - let mut request = reserve_request("to-worker-1"); - request.prompt = prompt(); - request.pinned_worker = Some(WorkerWithDpRank::new(1, 0)); - let response = core.select_and_reserve(request).await.expect("reserve"); + let response = reserve_pinned(&core, "to-worker-1", &tokens, 1).await; assert!(response.kv_hint.is_none()); // Query-only selections never carry a hint. let mut request = select_request(); - request.prompt = prompt(); + request.prompt = PromptRequest { + token_ids: Some(tokens.clone()), + ..PromptRequest::default() + }; request.pinned_worker = Some(WorkerWithDpRank::new(2, 0)); let response = core.select(request).await.expect("select"); assert!(response.kv_hint.is_none()); } #[tokio::test] - async fn router_hint_needs_capable_workers() { - use crate::indexer::KvIndexerInterface; - use crate::protocols::{BlockHashOptions, StorageTier, compute_block_hash_for_seq}; + async fn state_agent_workers_are_not_router_hint_sources() { + let (core, _entry, tokens) = hint_fixture(|request| { + hint_capable(request); + if request.worker_id == 1 { + request.kv_event_source_mode = Some("state_agent_v2".to_string()); + } + }) + .await; + // Worker 1 holds the prefix but reports through a state agent. + let response = reserve_pinned(&core, "to-worker-2", &tokens, 2).await; + assert!(response.kv_hint.is_none()); + } - let core = SelectionCore::try_new_local( - test_config(true), - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"); - for worker_id in [1, 2] { - core.upsert_worker(worker_with_kv_events(worker_id)) - .await - .expect("worker upsert"); - } - let entry = core - .entry(&RoutingPartitionId::new("model", "default")) - .expect("entry"); - let hashes: Vec = - compute_block_hash_for_seq(&[1, 2, 3, 4], 4, BlockHashOptions::default()) - .into_iter() - .map(|hash| hash.0) - .collect(); - entry - .indexer - .apply_event_routed(store_event(1, 0, 1, &[], &hashes, StorageTier::Device)) - .await - .unwrap(); - if let Indexer::Single { primary, .. } = &entry.indexer { - let _ = primary.flush().await; - } - let mut request = reserve_request("plain"); - request.pinned_worker = Some(WorkerWithDpRank::new(2, 0)); - let response = core.select_and_reserve(request).await.expect("reserve"); + #[tokio::test] + async fn router_hint_needs_capable_workers() { + let (core, _entry, tokens) = hint_fixture(|_| {}).await; + let response = reserve_pinned(&core, "plain", &tokens, 2).await; assert!(response.kv_hint.is_none()); } diff --git a/lib/kv-router/src/services/selection/types.rs b/lib/kv-router/src/services/selection/types.rs index a36eeaa1de62..712ffcd5d0b9 100644 --- a/lib/kv-router/src/services/selection/types.rs +++ b/lib/kv-router/src/services/selection/types.rs @@ -57,6 +57,10 @@ pub struct SelectionWorkerConfig { /// Per-global-DP-rank KV control endpoints a hint target fetches from. #[serde(default, skip_serializing_if = "HashMap::is_empty")] pub router_hint_source_control_endpoints: HashMap, + /// How the worker publishes KV events; a `state_agent_v2` worker is never + /// a router-hint source. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub kv_event_source_mode: Option, } impl WorkerConfigLike for SelectionWorkerConfig { @@ -148,6 +152,8 @@ pub struct WorkerCatalogRecord { pub router_hint_worker_type: Option, #[serde(default, skip_serializing_if = "HashMap::is_empty")] pub router_hint_source_control_endpoints: HashMap, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub kv_event_source_mode: Option, #[serde(default, skip_serializing_if = "Vec::is_empty")] pub not_schedulable_reasons: Vec, } @@ -177,6 +183,7 @@ impl WorkerCatalogRecord { kv_transfer_preferred_weight: req.kv_transfer_preferred_weight, router_hint_worker_type: req.router_hint_worker_type, router_hint_source_control_endpoints: req.router_hint_source_control_endpoints, + kv_event_source_mode: req.kv_event_source_mode, not_schedulable_reasons: Vec::new(), } } @@ -215,6 +222,7 @@ impl WorkerCatalogRecord { kv_transfer_preferred_weight: self.kv_transfer_preferred_weight, router_hint_worker_type: self.router_hint_worker_type.clone(), router_hint_source_control_endpoints: self.router_hint_source_control_endpoints.clone(), + kv_event_source_mode: self.kv_event_source_mode.clone(), }) } @@ -275,6 +283,7 @@ impl Default for WorkerRequest { kv_transfer_preferred_weight: None, router_hint_worker_type: None, router_hint_source_control_endpoints: HashMap::new(), + kv_event_source_mode: None, } } } @@ -312,6 +321,8 @@ pub struct WorkerRequest { /// Per-global-DP-rank KV control endpoints this worker can serve hints from. #[serde(default)] pub router_hint_source_control_endpoints: HashMap, + #[serde(default)] + pub kv_event_source_mode: Option, } #[derive(Debug, Deserialize)] @@ -336,6 +347,8 @@ pub struct WorkerPatchRequest { pub router_hint_worker_type: Option, #[serde(default)] pub router_hint_source_control_endpoints: Option>, + #[serde(default)] + pub kv_event_source_mode: Option, } impl WorkerCatalogRecord { @@ -397,6 +410,9 @@ impl WorkerCatalogRecord { if let Some(endpoints) = patch.router_hint_source_control_endpoints { self.router_hint_source_control_endpoints = endpoints; } + if patch.kv_event_source_mode.is_some() { + self.kv_event_source_mode = patch.kv_event_source_mode; + } } } diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs index 4b9de67863de..25087ea6c962 100644 --- a/lib/llm/src/kv_router/embedded.rs +++ b/lib/llm/src/kv_router/embedded.rs @@ -618,6 +618,7 @@ pub(crate) fn worker_request_from_runtime_config( kv_transfer_preferred_weight: config.kv_transfer_preferred_weight, router_hint_worker_type, router_hint_source_control_endpoints, + kv_event_source_mode: config.kv_event_source_mode.clone(), ..WorkerRequest::default() } } From 6616406aae1e2b6f6080758b8939bce4f3c167c2 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Thu, 10 Sep 2026 14:42:56 -0700 Subject: [PATCH 055/102] feat(router): wire the frontend to the shared selection core without cutover The core gains a Lease admission that books with the queue lifecycle lease and hands it back armed, plus operation-level gates for active-block tracking, returned routing hashes, and replay caching, so the standalone handlers keep today's behavior while an embedding host can express its own. Untracked selections no longer compute tracking hashes they would discard. LoRA narrowing falls back to the partition's published worker snapshot, the set the scheduler actually scores. Indexer lookup failures keep their type in SelectionError::Indexer with unchanged HTTP status and kind. The frontend passes its shared KV cache and LoRA filter into the embedded core's host hooks and gains KvRouter::select_core, the core-backed twin of find_best_match_details_with_policy_class_inner. Two differential tests run both procedures on independent routers over the same corpus; the tracked trace is frozen under tests/data/selection_golden and regenerated only with SELECTION_GOLDEN_UPDATE=1. Production traffic still takes the old path. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Ishan Dhanani --- .../src/services/selection/catalog.rs | 9 - .../src/services/selection/core/mod.rs | 88 +-- .../src/services/selection/core/operation.rs | 36 +- lib/kv-router/src/services/selection/error.rs | 9 + lib/kv-router/src/services/selection/input.rs | 28 +- lib/kv-router/src/services/selection/tests.rs | 20 +- lib/llm/src/discovery/model_manager.rs | 4 +- lib/llm/src/kv_router.rs | 520 +++++++++++++++++- lib/llm/src/kv_router/embedded.rs | 45 +- .../selection_golden/aggregated_tracked.json | 37 ++ 10 files changed, 708 insertions(+), 88 deletions(-) create mode 100644 lib/llm/tests/data/selection_golden/aggregated_tracked.json diff --git a/lib/kv-router/src/services/selection/catalog.rs b/lib/kv-router/src/services/selection/catalog.rs index 52910a7b44b2..bfe49a50bda3 100644 --- a/lib/kv-router/src/services/selection/catalog.rs +++ b/lib/kv-router/src/services/selection/catalog.rs @@ -70,15 +70,6 @@ impl WorkerCatalog { .any(|record| schedulable_in(record, key)) } - pub(super) fn schedulable_worker_ids_for_key(&self, key: &RoutingPartitionId) -> Vec { - self.workers - .read() - .values() - .filter(|record| schedulable_in(record, key)) - .map(|record| record.worker_id) - .collect() - } - /// `total_kv_blocks` published by a schedulable worker in `key`'s partition. pub(super) fn total_kv_blocks( &self, diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index b0aa620fd1bf..84b568abea6c 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -70,7 +70,6 @@ use super::types::{ }; use crate::WorkerSelectionPolicyFactory; use crate::WorkerType; -use crate::indexer::KvRouterError; use crate::services::common::replica_sync::AffinityBindingEvent; /// The scheduler type every partition runs. @@ -1031,6 +1030,9 @@ impl SelectionCore { allowed_worker_ids: req.allowed_worker_ids, routing_constraints: req.routing_constraints, admission, + track_active_blocks: true, + return_routing_hashes: false, + replay_id: req.selection_id.clone(), }) .await?; Ok(self.select_response(selected, req.selection_id)) @@ -1072,6 +1074,9 @@ impl SelectionCore { admission: SelectionAdmission::Book { selection_id: selection_id.clone(), }, + track_active_blocks: true, + return_routing_hashes: false, + replay_id: None, }) .await?; Ok(self.select_response(selected, Some(selection_id))) @@ -1103,6 +1108,8 @@ impl SelectionCore { track_prefill_tokens, effective_prefill_tokens, kv_hint, + routing_hashes: _, + lease: _, } = selected; let booked = sequence_hashes.is_some(); let potential_decode_blocks = response.potential_decode_blocks as u64; @@ -1168,6 +1175,9 @@ impl SelectionCore { allowed_worker_ids, routing_constraints, admission, + track_active_blocks, + return_routing_hashes, + replay_id, } = operation; self.ensure_running()?; @@ -1185,6 +1195,11 @@ impl SelectionCore { let table = entry.affinity.get(); let mut affinity_hold = None; let managed_session = match (&session, table) { + (SessionBinding::Managed { .. }, _) if claim.is_none() => { + return Err(SelectionError::Internal( + "a managed session binding requires Book admission".to_string(), + )); + } (SessionBinding::Managed { session_id }, Some(table)) => Some((table, session_id)), _ => None, }; @@ -1216,8 +1231,10 @@ impl SelectionCore { .prepare_selection_inputs( &entry, &prompt, - self.kv_router_config - .assume_kv_reuse(router_config_override.as_ref()), + track_active_blocks.then(|| { + self.kv_router_config + .assume_kv_reuse(router_config_override.as_ref()) + }), true, retain_kv_transfer_chain, ) @@ -1225,7 +1242,10 @@ impl SelectionCore { // The queue lease frees a booking whose response is never consumed // (the caller dropped this future after the actor booked). let mode = match &admission { - SelectionAdmission::Book { selection_id } => ScheduleMode::TrackedWithLifecycle { + SelectionAdmission::Book { selection_id } + | SelectionAdmission::Lease { + request_id: selection_id, + } => ScheduleMode::TrackedWithLifecycle { request_id: selection_id.clone(), }, SelectionAdmission::Query { request_id } => ScheduleMode::QueryOnly { @@ -1241,14 +1261,9 @@ impl SelectionCore { .unwrap_or(self.kv_router_config.router_track_prefill_tokens); // An unbooked selection with an id caches the booking inputs so a // follow-up `create_reservation` can replay them by that id. - let cached_inputs = match &admission { - SelectionAdmission::Query { - request_id: Some(id), - } - | SelectionAdmission::Advisory { - request_id: Some(id), - } => Some(( - id.clone(), + let cached_inputs = replay_id.filter(|_| !book).map(|id| { + ( + id, sequence_hashes.clone(), prompt.lora_name.map(str::to_string), track_prefill_tokens, @@ -1256,28 +1271,28 @@ impl SelectionCore { SessionBinding::Query { session_id } => Some(session_id.clone()), _ => None, }, - )), - _ => None, - }; + ) + }); let allowed_worker_ids = match self.host.eligibility.lora_worker_filter.as_deref() { Some(filter) => narrow_allowed_worker_ids_by_lora( filter, prompt.lora_name, allowed_worker_ids, pinned_worker.as_ref(), - || self.catalog.schedulable_worker_ids_for_key(&key), + || entry.workers_tx.borrow().keys().copied().collect(), ), None => allowed_worker_ids, }; // Bookings (now, or later via the pending-selection cache) are recorded // into an approximate indexer; keep the public hashes for that. let routing_hashes = (entry.indexer.records_routing_decisions() - && (book || cached_inputs.is_some())) + && (claim.is_some() || cached_inputs.is_some())) .then(|| block_hashes.clone()); let booked_sequence_hashes = book.then(|| sequence_hashes.clone()); + let returned_routing_hashes = return_routing_hashes.then(|| block_hashes.clone()); let schedule_request = ScheduleRequest { mode, - token_seq: Some(sequence_hashes), + token_seq: track_active_blocks.then_some(sequence_hashes), block_hashes: Some(block_hashes), isl_tokens, overlap, @@ -1363,8 +1378,10 @@ impl SelectionCore { // The routing hashes go to exactly one of: the booking recorded now, or // the cached inputs a later replay records. + // The routing hashes go to exactly one of: the reservation recorded + // now, or the replay cache a later reservation records from. let mut routing_hashes = routing_hashes; - if let Some(claim) = claim { + let lease = if let Some(claim) = claim { let Some(lease) = lease else { return Err(SelectionError::Internal( "booked selection has no lifecycle lease".to_string(), @@ -1387,7 +1404,10 @@ impl SelectionCore { booking: Some(lease.commit().ok_or_else(missing_booking)?), _affinity_lease: affinity_lease, }); - } + None + } else { + lease + }; if let Some((cache_id, sequence_hashes, lora_name, track_prefill_tokens, session_id)) = cached_inputs @@ -1421,6 +1441,8 @@ impl SelectionCore { track_prefill_tokens, effective_prefill_tokens: effective_prefill, kv_hint, + routing_hashes: returned_routing_hashes, + lease, })) } @@ -1965,8 +1987,10 @@ impl SelectionCore { .prepare_selection_inputs( &entry, &req.prompt.view(), - self.kv_router_config - .assume_kv_reuse(req.router_config_override.as_ref()), + Some( + self.kv_router_config + .assume_kv_reuse(req.router_config_override.as_ref()), + ), false, false, ) @@ -2018,21 +2042,24 @@ impl SelectionCore { /// the optional shared-cache lookup run concurrently; the shared cache is /// consulted only when `query_shared_cache` is set, a shared cache is /// attached, and the prompt carries raw `token_ids`. + /// `tracking_assume_kv_reuse` is `None` when the caller does not track + /// active blocks. async fn prepare_selection_inputs( &self, entry: &SelectionEntry, prompt: &PromptView<'_>, - assume_kv_reuse: bool, + tracking_assume_kv_reuse: Option, query_shared_cache: bool, retain_kv_transfer_chain: bool, ) -> Result { let normalized = prompt.normalize_for_selection( + entry.block_size, entry.is_eagle, - TrackingHashInput { + tracking_assume_kv_reuse.map(|assume_kv_reuse| TrackingHashInput { context: &self.tracking_hash, scope: tracking_scope(entry), assume_kv_reuse, - }, + }), )?; let indexer_lookup = async { if normalized.block_hashes.is_empty() { @@ -2047,12 +2074,7 @@ impl SelectionCore { }, ) .await - .map_err(|error| match error { - KvRouterError::IndexerOffline => { - SelectionError::NotReady(error.to_string()) - } - other => SelectionError::Internal(other.to_string()), - }) + .map_err(SelectionError::Indexer) } }; let shared_cache = query_shared_cache @@ -2604,7 +2626,9 @@ mod tests { core.upsert_worker(worker(1)).await.expect("worker upsert"); assert!(matches!( core.select(select_request()).await, - Err(SelectionError::NotReady(_)) + Err(SelectionError::Indexer( + crate::indexer::KvRouterError::IndexerOffline + )) )); } diff --git a/lib/kv-router/src/services/selection/core/operation.rs b/lib/kv-router/src/services/selection/core/operation.rs index 4ba513281061..17ccd1865bc2 100644 --- a/lib/kv-router/src/services/selection/core/operation.rs +++ b/lib/kv-router/src/services/selection/core/operation.rs @@ -10,8 +10,11 @@ use dynamo_tokens::SequenceHash; use crate::identity::RoutingPartitionId; use crate::kv_hints::KvHint; -use crate::protocols::{RoutingConstraints, WorkerAffinityTarget, WorkerId, WorkerWithDpRank}; +use crate::protocols::{ + LocalBlockHash, RoutingConstraints, WorkerAffinityTarget, WorkerId, WorkerWithDpRank, +}; use crate::scheduling::config::RouterConfigOverride; +use crate::scheduling::queue::RequestLifecycleLease; use crate::scheduling::{AdvisoryWorkerLoad, QueueRejection, SchedulingResponse, SessionContext}; use super::super::input::PromptView; @@ -34,16 +37,26 @@ pub struct SelectionOperation<'a> { pub allowed_worker_ids: Option>, pub routing_constraints: RoutingConstraints, pub admission: SelectionAdmission, + /// Send the prompt's tracking hashes to the scheduler so the booking's + /// blocks count as active on its worker. + pub track_active_blocks: bool, + /// Return the prompt's public block hashes for the host to record. + pub return_routing_hashes: bool, + /// Cache the inputs under this id for a later `create_reservation` replay + /// (unbooked admissions only). + pub replay_id: Option, } pub enum SelectionAdmission { - /// Queue admission without a booking; a `request_id` caches the inputs for - /// a later `create_reservation` replay. + /// Queue admission without a booking. Query { request_id: Option }, - /// Queue admission with a booking recorded under `selection_id`. + /// Queue admission with a booking the core records as a reservation under + /// `selection_id` and releases through its lifecycle calls. Book { selection_id: String }, - /// Skip queue admission and report the chosen worker's load; a - /// `request_id` caches the inputs like `Query`. + /// Queue admission with a booking the caller owns: its lifecycle lease is + /// returned armed in [`Selected::lease`]. + Lease { request_id: String }, + /// Skip queue admission and report the chosen worker's load. Advisory { request_id: Option }, } @@ -51,12 +64,15 @@ impl SelectionAdmission { pub fn request_id(&self) -> Option<&str> { match self { Self::Query { request_id } | Self::Advisory { request_id } => request_id.as_deref(), - Self::Book { selection_id } => Some(selection_id), + Self::Book { selection_id } + | Self::Lease { + request_id: selection_id, + } => Some(selection_id), } } pub fn is_booking(&self) -> bool { - matches!(self, Self::Book { .. }) + matches!(self, Self::Book { .. } | Self::Lease { .. }) } } @@ -97,4 +113,8 @@ pub struct Selected { pub track_prefill_tokens: bool, pub effective_prefill_tokens: usize, pub kv_hint: Option, + pub routing_hashes: Option>, + /// The booking's lifecycle lease; `Lease` admission only. Dropping it frees + /// the booking, `commit` hands it to the caller's own cleanup. + pub lease: Option>, } diff --git a/lib/kv-router/src/services/selection/error.rs b/lib/kv-router/src/services/selection/error.rs index 0c5f4f36b123..0cce437cfffa 100644 --- a/lib/kv-router/src/services/selection/error.rs +++ b/lib/kv-router/src/services/selection/error.rs @@ -5,6 +5,7 @@ use axum::Json; use axum::http::StatusCode; use axum::response::{IntoResponse, Response}; +use crate::indexer::KvRouterError; use crate::scheduling::KvSchedulerError; use crate::sequences::SequenceError; @@ -24,6 +25,10 @@ pub enum SelectionError { Scheduler(#[from] KvSchedulerError), #[error(transparent)] Sequence(#[from] SequenceError), + /// The KV index lookup failed; an offline remote index is not ready, + /// anything else is internal. + #[error(transparent)] + Indexer(#[from] KvRouterError), } impl SelectionError { @@ -36,6 +41,8 @@ impl SelectionError { Self::Internal(_) => StatusCode::INTERNAL_SERVER_ERROR, Self::Scheduler(error) => scheduler_error_status(error), Self::Sequence(error) => sequence_error_status(error), + Self::Indexer(KvRouterError::IndexerOffline) => StatusCode::SERVICE_UNAVAILABLE, + Self::Indexer(_) => StatusCode::INTERNAL_SERVER_ERROR, } } @@ -55,6 +62,8 @@ impl SelectionError { Self::Internal(_) => "internal", Self::Scheduler(_) => "scheduler", Self::Sequence(_) => "sequence", + Self::Indexer(KvRouterError::IndexerOffline) => "not_ready", + Self::Indexer(_) => "internal", } } } diff --git a/lib/kv-router/src/services/selection/input.rs b/lib/kv-router/src/services/selection/input.rs index 5e7780275434..55ada4e9b26b 100644 --- a/lib/kv-router/src/services/selection/input.rs +++ b/lib/kv-router/src/services/selection/input.rs @@ -75,14 +75,18 @@ pub struct PromptView<'a> { } impl PromptView<'_> { + /// `tracking` is `None` when the caller does not track active blocks; the + /// tracking hashes are then left empty instead of computed and discarded. pub(super) fn normalize_for_selection( &self, + block_size: u32, default_is_eagle: bool, - tracking: TrackingHashInput<'_>, + tracking: Option>, ) -> Result { if let Some((token_ids, block_mm_infos)) = self.routing_tokens_and_mm_infos() { return Ok(normalize_tokens_for_selection( token_ids, + block_size, self.lora_name, self.cache_namespace, block_mm_infos, @@ -180,11 +184,12 @@ impl PromptView<'_> { fn normalize_tokens_for_selection( token_ids: &[u32], + block_size: u32, lora_name: Option<&str>, cache_namespace: Option<&str>, block_mm_infos: Option<&[Option]>, is_eagle: bool, - tracking: TrackingHashInput<'_>, + tracking: Option>, ) -> NormalizedPrompt { let hash_options = BlockHashOptions { block_mm_infos, @@ -192,15 +197,16 @@ fn normalize_tokens_for_selection( cache_namespace, is_eagle: Some(is_eagle), }; - let block_hashes = - compute_block_hash_for_seq(token_ids, tracking.scope.block_size, hash_options); - let sequence_hashes = tracking.context.compute_sequence_hashes_for_tracking( - tracking.scope, - token_ids, - hash_options, - tracking.assume_kv_reuse, - Some(&block_hashes), - ); + let block_hashes = compute_block_hash_for_seq(token_ids, block_size, hash_options); + let sequence_hashes = tracking.map_or_else(Vec::new, |tracking| { + tracking.context.compute_sequence_hashes_for_tracking( + tracking.scope, + token_ids, + hash_options, + tracking.assume_kv_reuse, + Some(&block_hashes), + ) + }); NormalizedPrompt { block_hashes, sequence_hashes, diff --git a/lib/kv-router/src/services/selection/tests.rs b/lib/kv-router/src/services/selection/tests.rs index 3a89f3bab511..ebdc009b0c79 100644 --- a/lib/kv-router/src/services/selection/tests.rs +++ b/lib/kv-router/src/services/selection/tests.rs @@ -129,15 +129,16 @@ fn normalize_prompt(request: &PromptRequest) -> super::input::NormalizedPrompt { request .view() .normalize_for_selection( + 4, false, - TrackingHashInput { + Some(TrackingHashInput { context: &context, scope: TrackingHashScope { partition: RoutingPartitionRef::new("model", "default"), block_size: 4, }, assume_kv_reuse: true, - }, + }), ) .expect("normalize prompt") } @@ -569,15 +570,16 @@ fn keyed_prompt_tracking_leaves_indexer_hashes_public() { let normalized = request .view() .normalize_for_selection( + 4, false, - TrackingHashInput { + Some(TrackingHashInput { context: &context, scope: TrackingHashScope { partition: RoutingPartitionRef::new("model", "default"), block_size: 4, }, assume_kv_reuse: true, - }, + }), ) .unwrap(); let public_blocks = compute_block_hash_for_seq( @@ -608,15 +610,16 @@ fn disabled_kv_reuse_keeps_public_indexer_hashes_and_randomizes_tracking() { request .view() .normalize_for_selection( + 4, false, - TrackingHashInput { + Some(TrackingHashInput { context: &context, scope: TrackingHashScope { partition: RoutingPartitionRef::new("model", "default"), block_size: 4, }, assume_kv_reuse: false, - }, + }), ) .unwrap() }; @@ -703,12 +706,13 @@ fn keyed_hash_only_inputs_remain_trusted_for_selection_and_reservation() { let selection = request .view() .normalize_for_selection( + 4, false, - TrackingHashInput { + Some(TrackingHashInput { context: &context, scope, assume_kv_reuse: true, - }, + }), ) .unwrap(); let reservation = request diff --git a/lib/llm/src/discovery/model_manager.rs b/lib/llm/src/discovery/model_manager.rs index 0cd43f314e4f..af91e9d944a7 100644 --- a/lib/llm/src/discovery/model_manager.rs +++ b/lib/llm/src/discovery/model_manager.rs @@ -2119,7 +2119,7 @@ impl ModelManager { ); // A policy that does not consume cache input must not create a shared-cache client or // subscribe to its updates. - let shared_cache: Option> = + let shared_cache: Option> = if required_worker_inputs.contains(WorkerInputs::CACHE) { match kv_router_config .as_ref() @@ -2133,7 +2133,7 @@ impl ModelManager { worker_component = worker_component_name, "Using HiCache shared KV cache" ); - Some(Box::new( + Some(Arc::new( self.hicache_cache_for(&endpoint, workers_with_configs.clone()), )) } diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index 7e14ea934f5d..1c49d4ecb5fb 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -33,6 +33,10 @@ use dynamo_kv_router::{ overlap::cache_hit_estimates_from_tiered_matches, }, selector::WorkerInputs, + services::selection::{ + PromptView, Selected, SelectionAdmission, SelectionError, SelectionOperation, + SelectionOutcome, SessionBinding, + }, }; use dynamo_runtime::{ CancellationToken, @@ -600,7 +604,7 @@ pub struct KvRouter { request_leases: request_lease::RequestLeaseManager, /// Optional external shared KV cache pool. When present, `find_best_match` /// queries it in parallel with the indexer and factors shared hits into scoring. - shared_cache: Option>, + shared_cache: Option>, /// Optional LoRA filter. When present (LoRA serving enabled), candidate workers are /// narrowed to the LoRA's allocated/loaded replicas inside `find_best_match_details`, /// covering both the decode and prefill routers (both built via `kv_chooser_for`). @@ -638,7 +642,7 @@ impl KvRouter { metric_worker_type: &'static str, model_name: Option, is_eagle: bool, - shared_cache: Option>, + shared_cache: Option>, lora_filter: Option>, ) -> Result { Self::new_with_worker_role( @@ -674,7 +678,7 @@ impl KvRouter { metric_worker_type: &'static str, model_name: Option, is_eagle: bool, - shared_cache: Option>, + shared_cache: Option>, lora_filter: Option>, ) -> Result { let parent_token = endpoint.component().drt().child_token(); @@ -715,7 +719,7 @@ impl KvRouter { metric_worker_type: &'static str, model_name: Option, is_eagle: bool, - shared_cache: Option>, + shared_cache: Option>, lora_filter: Option>, scheduler_load: SchedulerLoadSender, parent_token: CancellationToken, @@ -832,6 +836,10 @@ impl KvRouter { prefill_load_estimator: prefill_load_estimator.clone(), overloaded_worker_provider, available_worker_provider, + shared_cache: shared_cache.clone(), + lora_worker_filter: lora_filter.clone().map(|filter| { + filter as Arc + }), ingress: Arc::clone(&ingress) as Arc, scheduler_load, @@ -1503,6 +1511,154 @@ impl KvRouter { } } + /// The shared core's selection for the same inputs as + /// `find_best_match_details_with_policy_class_inner`. + #[allow(clippy::too_many_arguments)] + #[cfg_attr(not(test), expect(dead_code))] + async fn select_core( + &self, + context_id: Option<&str>, + tokens: &[u32], + block_mm_infos: Option<&[Option]>, + router_config_override: Option<&RouterConfigOverride>, + update_states: bool, + return_routing_hashes: bool, + lora_name: Option<&str>, + cache_namespace: Option<&str>, + priority_jump: f64, + strict_priority: u32, + policy_class: Option, + session_context: Option, + expected_output_tokens: Option, + affinity_target: Option, + pinned_worker: Option, + allowed_worker_ids: Option>, + routing_constraints: RoutingConstraints, + admission: FindBestMatchAdmission, + ) -> anyhow::Result { + if update_states && context_id.is_none() { + anyhow::bail!("context_id must be provided if update_states is true"); + } + let core_admission = match admission { + FindBestMatchAdmission::WithAdmission { .. } if update_states => { + SelectionAdmission::Lease { + request_id: context_id.expect("validated above").to_string(), + } + } + FindBestMatchAdmission::WithAdmission { .. } => SelectionAdmission::Query { + request_id: context_id.map(str::to_string), + }, + FindBestMatchAdmission::WithoutAdmission => SelectionAdmission::Advisory { + request_id: context_id.map(str::to_string), + }, + }; + let outcome = self + .scheduler + .run_selection(SelectionOperation { + key: self.scheduler.partition_key().clone(), + prompt: PromptView { + token_ids: Some(tokens), + mm_routing_info: None, + block_mm_infos, + block_hashes: None, + sequence_hashes: None, + isl_tokens: None, + lora_name, + cache_namespace, + is_eagle: Some(self.is_eagle), + }, + router_config_override: router_config_override.cloned(), + expected_output_tokens, + priority_jump, + strict_priority, + policy_class, + session_context, + session: SessionBinding::None, + affinity_target, + pinned_worker, + allowed_worker_ids, + routing_constraints, + admission: core_admission, + track_active_blocks: self.kv_router_config.router_track_active_blocks, + return_routing_hashes, + replay_id: None, + }) + .await; + let selected = match outcome { + Ok(SelectionOutcome::Selected(selected)) => selected, + Ok(SelectionOutcome::QueueRejected { rejection }) => { + return Ok(match admission { + FindBestMatchAdmission::WithAdmission { .. } => { + FindBestMatchInnerOutcome::WithAdmission(AdmittedFindBestMatchOutcome { + outcome: FindBestMatchOutcome::QueueRejected { rejection }, + attempt: AdmissionAttempt::Untracked, + }) + } + FindBestMatchAdmission::WithoutAdmission => { + FindBestMatchInnerOutcome::WithoutAdmission( + FindBestMatchAdvisoryOutcome::QueueRejected { rejection }, + ) + } + }); + } + Err(SelectionError::Scheduler(error)) => return Err(map_scheduler_error(error)), + Err(SelectionError::Indexer(error)) => return Err(error.into()), + // The partition has no schedulable worker: the scheduler's own answer. + Err(SelectionError::NotReady(_)) => { + return Err(map_scheduler_error(KvSchedulerError::NoEndpoints)); + } + Err(error) => return Err(error.into()), + }; + let Selected { + response, + advisory_load, + kv_hint, + routing_hashes, + lease, + .. + } = selected; + let routing_hashes = routing_hashes.map(RoutingDecisionHashes::from_local_hashes); + let overlap_blocks = response.effective_overlap_blocks.round() as u32; + Ok(match admission { + FindBestMatchAdmission::WithAdmission { .. } => { + // The caller owns cleanup through the attempt identity. + let attempt = match lease { + Some(lease) => { + let booking = lease + .commit() + .ok_or_else(|| anyhow::anyhow!("booking lease holds no booking"))?; + AdmissionAttempt::Tracked(booking.attempt_id) + } + None => AdmissionAttempt::Untracked, + }; + FindBestMatchInnerOutcome::WithAdmission(AdmittedFindBestMatchOutcome { + outcome: FindBestMatchOutcome::Routed { + worker: response.best_worker, + overlap_blocks, + effective_overlap_blocks: response.effective_overlap_blocks, + cached_tokens: response.cached_tokens, + potential_decode_blocks: response.potential_decode_blocks as u64, + routing_hashes, + kv_hint, + }, + attempt, + }) + } + FindBestMatchAdmission::WithoutAdmission => { + FindBestMatchInnerOutcome::WithoutAdmission(FindBestMatchAdvisoryOutcome::Routed { + worker: response.best_worker, + overlap_blocks, + effective_overlap_blocks: response.effective_overlap_blocks, + cached_tokens: response.cached_tokens, + potential_decode_blocks: response.potential_decode_blocks as u64, + selected_worker_load: advisory_load + .expect("without-admission selection returns advisory load"), + routing_hashes, + }) + } + }) + } + #[allow(clippy::too_many_arguments)] async fn find_best_match_details_with_policy_class_inner( &self, @@ -2761,7 +2917,7 @@ mod tests { "prefill", None, false, - Some(Box::new(FakeSharedCache { + Some(Arc::new(FakeSharedCache { hits: None, should_error: false, })), @@ -2780,9 +2936,353 @@ mod tests { )); } + /// Three default-config workers under the registry policy, with + /// `router_track_active_blocks` on so bookings send tracking hashes. + async fn differential_router(name: &str) -> KvRouter { + let component = make_test_component(name).await; + let endpoint = component.endpoint("backend"); + let client = endpoint.client().await.unwrap(); + let workers = (0..3) + .map(|worker_id| (worker_id, ModelRuntimeConfig::default())) + .collect(); + let (_tx, rx) = watch::channel(workers); + // Prefill load decays with wall time and would let near-ties flip + // between two procedures run microseconds apart. + let config = KvRouterConfig { + use_kv_events: false, + router_track_prefill_tokens: false, + skip_initial_worker_wait: true, + ..Default::default() + }; + let router = KvRouter::new_with_worker_role( + endpoint, + client, + rx, + None, + 2, + SelectionPolicySource::Registry, + Some(config), + None, + None, + "decode", + None, + false, + None, + None, + ) + .await + .unwrap(); + // Worker 1 has served prefix A and worker 2 prefix B, so every corpus + // prompt has one unambiguous best worker. + for (worker_id, tokens) in [(1u64, golden_prefix(1)), (2, golden_prefix(2))] { + router + .record_routing_decision( + TokensWithHashes::new(tokens.clone(), 2), + WorkerWithDpRank::from_worker_id(worker_id), + ) + .await + .unwrap(); + // The approximate index applies recordings asynchronously; wait + // until a query sees the whole prefix. + tokio::time::timeout(std::time::Duration::from_secs(5), async { + loop { + let (worker, overlap) = router + .find_best_match( + None, + &tokens, + None, + None, + false, + None, + None, + 0.0, + 0, + None, + None, + RoutingConstraints::default(), + ) + .await + .unwrap(); + if worker.worker_id == worker_id && overlap == 8 { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("seeded prefix became visible"); + } + router + } + + fn golden_prefix(seed: u32) -> Vec { + (0..16).map(|i| seed * 100 + i).collect() + } + + /// Prompts sharing 4, 2, 8 (then continuing past it), 8, and 8 blocks + /// with a seeded prefix. Partial matches come first so they are scored + /// against the seeded index alone, before bookings add decode load. + fn golden_corpus() -> Vec> { + let a = golden_prefix(1); + let b = golden_prefix(2); + vec![ + a[..8].to_vec(), + b[..4].iter().copied().chain(900..912).collect(), + a.iter().copied().chain(1_000..1_008).collect(), + a, + b, + ] + } + + #[derive(Debug, PartialEq, serde::Serialize, serde::Deserialize)] + struct GoldenDecision { + worker_id: WorkerId, + dp_rank: u32, + overlap_blocks: u32, + cached_tokens: usize, + potential_decode_blocks: u64, + } + + impl GoldenDecision { + fn from_outcome(outcome: &FindBestMatchOutcome) -> Self { + match outcome { + FindBestMatchOutcome::Routed { + worker, + overlap_blocks, + cached_tokens, + potential_decode_blocks, + .. + } => Self { + worker_id: worker.worker_id, + dp_rank: worker.dp_rank, + overlap_blocks: *overlap_blocks, + cached_tokens: *cached_tokens, + potential_decode_blocks: *potential_decode_blocks, + }, + FindBestMatchOutcome::QueueRejected { rejection } => { + panic!("golden corpus must not be queue rejected: {rejection:?}") + } + } + } + } + + const GOLDEN_PATH: &str = concat!( + env!("CARGO_MANIFEST_DIR"), + "/tests/data/selection_golden/aggregated_tracked.json" + ); + + /// Both selection procedures, run on independent routers over the same + /// corpus with tracked bookings, must agree with each other and with the + /// frozen trace recorded from the frontend procedure. Regenerate the trace + /// with `SELECTION_GOLDEN_UPDATE=1` only for an intended behavior change. + #[tokio::test] + async fn core_selection_matches_frontend_golden_trace() { + let frontend = differential_router("golden-frontend").await; + let core = differential_router("golden-core").await; + let mut frontend_trace = Vec::new(); + let mut core_trace = Vec::new(); + for (index, tokens) in golden_corpus().iter().enumerate() { + let context_id = format!("golden-{index}"); + let admission = FindBestMatchAdmission::WithAdmission { + track_lifecycle: true, + }; + let old = frontend + .find_best_match_details_with_policy_class_inner( + Some(&context_id), + tokens, + None, + None, + true, + true, + None, + None, + 0.0, + 0, + None, + None, + None, + None, + None, + None, + RoutingConstraints::default(), + admission, + ) + .await + .unwrap(); + let new = core + .select_core( + Some(&context_id), + tokens, + None, + None, + true, + true, + None, + None, + 0.0, + 0, + None, + None, + None, + None, + None, + None, + RoutingConstraints::default(), + admission, + ) + .await + .unwrap(); + let ( + FindBestMatchInnerOutcome::WithAdmission(old), + FindBestMatchInnerOutcome::WithAdmission(new), + ) = (old, new) + else { + panic!("admitted routing returned advisory outcome"); + }; + assert!(matches!(old.attempt, AdmissionAttempt::Tracked(_))); + assert!(matches!(new.attempt, AdmissionAttempt::Tracked(_))); + match (&old.outcome, &new.outcome) { + ( + FindBestMatchOutcome::Routed { + routing_hashes: old_hashes, + .. + }, + FindBestMatchOutcome::Routed { + routing_hashes: new_hashes, + .. + }, + ) => assert_eq!( + old_hashes.as_ref().map(|h| h.local_hashes.clone()), + new_hashes.as_ref().map(|h| h.local_hashes.clone()) + ), + _ => panic!("golden corpus must route"), + } + frontend_trace.push(GoldenDecision::from_outcome(&old.outcome)); + core_trace.push(GoldenDecision::from_outcome(&new.outcome)); + } + assert_eq!( + core_trace, frontend_trace, + "core and frontend procedures diverged" + ); + + if std::env::var_os("SELECTION_GOLDEN_UPDATE").is_some() { + std::fs::write( + GOLDEN_PATH, + serde_json::to_string_pretty(&frontend_trace).unwrap() + "\n", + ) + .unwrap(); + } + let golden: Vec = + serde_json::from_str(&std::fs::read_to_string(GOLDEN_PATH).unwrap()).unwrap(); + assert_eq!( + frontend_trace, golden, + "frontend procedure drifted from its frozen trace" + ); + } + + /// Query-only and advisory selections mutate nothing, so one router serves + /// both procedures. + #[tokio::test] + async fn core_selection_matches_frontend_for_unbooked_requests() { + let router = differential_router("golden-unbooked").await; + for tokens in golden_corpus() { + for admission in [ + FindBestMatchAdmission::WithAdmission { + track_lifecycle: false, + }, + FindBestMatchAdmission::WithoutAdmission, + ] { + let old = router + .find_best_match_details_with_policy_class_inner( + None, + &tokens, + None, + None, + false, + false, + None, + None, + 0.0, + 0, + None, + None, + None, + None, + None, + None, + RoutingConstraints::default(), + admission, + ) + .await + .unwrap(); + let new = router + .select_core( + None, + &tokens, + None, + None, + false, + false, + None, + None, + 0.0, + 0, + None, + None, + None, + None, + None, + None, + RoutingConstraints::default(), + admission, + ) + .await + .unwrap(); + match (old, new) { + ( + FindBestMatchInnerOutcome::WithAdmission(old), + FindBestMatchInnerOutcome::WithAdmission(new), + ) => { + assert!(matches!(old.attempt, AdmissionAttempt::Untracked)); + assert!(matches!(new.attempt, AdmissionAttempt::Untracked)); + assert_eq!( + GoldenDecision::from_outcome(&new.outcome), + GoldenDecision::from_outcome(&old.outcome) + ); + } + ( + FindBestMatchInnerOutcome::WithoutAdmission( + FindBestMatchAdvisoryOutcome::Routed { + worker: old_worker, + overlap_blocks: old_overlap, + selected_worker_load: old_load, + .. + }, + ), + FindBestMatchInnerOutcome::WithoutAdmission( + FindBestMatchAdvisoryOutcome::Routed { + worker: new_worker, + overlap_blocks: new_overlap, + selected_worker_load: new_load, + .. + }, + ), + ) => { + assert_eq!((new_worker, new_overlap), (old_worker, old_overlap)); + assert_eq!( + new_load.active_prefill_tokens, + old_load.active_prefill_tokens + ); + } + _ => panic!("procedures returned different outcome shapes"), + } + } + } + } + async fn make_test_router_with_workers( policy: SelectionPolicySource, - shared_cache: Option>, + shared_cache: Option>, workers: HashMap, ) -> KvRouter { let component = make_test_component("shared-cache-router").await; @@ -2822,7 +3322,7 @@ mod tests { async fn make_test_router( policy: SelectionPolicySource, - shared_cache: Option>, + shared_cache: Option>, ) -> KvRouter { let mut workers = HashMap::new(); workers.insert(0, ModelRuntimeConfig::default()); @@ -3130,7 +3630,7 @@ mod tests { async fn test_find_best_match_passes_shared_cache_hits_to_scheduler() { let router = make_test_router( fixed_policy(Some(2), WorkerWithDpRank::from_worker_id(1)), - Some(Box::new(FakeSharedCache { + Some(Arc::new(FakeSharedCache { #[allow(clippy::single_range_in_vec_init)] hits: Some(dynamo_kv_router::protocols::SharedCacheHits::from_ranges( vec![0..2], @@ -3166,7 +3666,7 @@ mod tests { async fn test_find_best_match_ignores_shared_cache_errors() { let router = make_test_router( fixed_policy(None, WorkerWithDpRank::from_worker_id(0)), - Some(Box::new(FakeSharedCache { + Some(Arc::new(FakeSharedCache { hits: None, should_error: true, })), @@ -3323,7 +3823,7 @@ mod tests { async fn test_get_overlap_scores_returns_tiered_rows_and_shared_hits() { let router = make_test_router( fixed_policy(None, WorkerWithDpRank::from_worker_id(0)), - Some(Box::new(FakeSharedCache { + Some(Arc::new(FakeSharedCache { #[allow(clippy::single_range_in_vec_init)] hits: Some(dynamo_kv_router::protocols::SharedCacheHits::from_ranges( vec![0..2], diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs index 25087ea6c962..d840df160617 100644 --- a/lib/llm/src/kv_router/embedded.rs +++ b/lib/llm/src/kv_router/embedded.rs @@ -28,9 +28,10 @@ use dynamo_kv_router::scheduling::{ use dynamo_kv_router::sequences::{SequenceError, SequenceRequest}; use dynamo_kv_router::services::selection::{ CatalogObserver, CatalogReconciler, HostCache, HostEligibility, HostLoad, HostReplication, - HostTelemetry, KvEventIngress, KvIndexSource, SelectionHost, SelectionPartition, - SelectionService, SelectionServiceBuilder, WorkerCatalogRecord, WorkerCatalogSource, - WorkerRequest, WorkerSelectionPolicyRegistry, + HostTelemetry, KvEventIngress, KvIndexSource, SelectionError, SelectionHost, + SelectionOperation, SelectionOutcome, SelectionPartition, SelectionService, + SelectionServiceBuilder, WorkerCatalogRecord, WorkerCatalogSource, WorkerRequest, + WorkerSelectionPolicyRegistry, }; use dynamo_kv_router::{DEFAULT_ROUTING_GROUP, PrefillLoadEstimator, WorkerSelectionPolicyFactory}; use dynamo_tokens::SequenceHash; @@ -53,6 +54,8 @@ pub(crate) struct EmbeddedSelectionArgs { pub prefill_load_estimator: Option>, pub overloaded_worker_provider: OverloadedWorkerProvider, pub available_worker_provider: WorkerAvailabilityProvider, + pub shared_cache: Option>, + pub lora_worker_filter: Option>, /// Builds and feeds the router's index from the runtime; the partition /// takes its index from it. pub ingress: Arc, @@ -165,7 +168,7 @@ pub(crate) const DEFAULT_MODEL_NAME: &str = "default"; pub(crate) struct EmbeddedSelection { /// Keeps the service (listeners, sweep, replica sync) alive for as long as /// the router holds the partition. - _service: Arc, + service: Arc, partition: SelectionPartition, affinity: std::sync::OnceLock, worker_type: &'static str, @@ -263,10 +266,12 @@ impl EmbeddedSelection { available_workers: Some(args.available_worker_provider), }, cache: HostCache { - shared: None, + shared: args.shared_cache, index: KvIndexSource::Owned(args.ingress), }, - eligibility: HostEligibility::default(), + eligibility: HostEligibility { + lora_worker_filter: args.lora_worker_filter, + }, telemetry: HostTelemetry { scheduler_load: Some(Arc::new(SenderLoadSink { sender: args.scheduler_load, @@ -347,7 +352,7 @@ impl EmbeddedSelection { "KvRouter scheduling on embedded selection partition" ); Ok(Self { - _service: service, + service, partition, affinity: std::sync::OnceLock::new(), worker_type: args.metric_worker_type, @@ -359,7 +364,14 @@ impl EmbeddedSelection { /// Queue changes are not observable through the scheduler's update watch, /// so every schedule result refreshes the gauges. fn observe_schedule_result(&self, result: &Result) { - if let Err(KvSchedulerError::QueueRejected(rejection)) = result { + self.observe_queue(match result { + Err(KvSchedulerError::QueueRejected(rejection)) => Some(rejection), + _ => None, + }); + } + + fn observe_queue(&self, rejection: Option<&QueueRejection>) { + if let Some(rejection) = rejection { record_queue_rejection(&self.queue_metrics, &self.queue_metric_indices, rejection); } update_queue_metrics(&self.queue_metrics, |index| { @@ -384,6 +396,23 @@ impl EmbeddedSelection { /// registration is a no-op here. pub(crate) fn register_workers(&self, _worker_ids: &HashSet) {} + pub(crate) fn partition_key(&self) -> &RoutingPartitionId { + self.partition.key() + } + + /// Run one selection through the shared core. + pub(crate) async fn run_selection( + &self, + operation: SelectionOperation<'_>, + ) -> Result { + let outcome = self.service.core().run_selection(operation).await; + self.observe_queue(match &outcome { + Ok(SelectionOutcome::QueueRejected { rejection }) => Some(rejection), + _ => None, + }); + outcome + } + pub(crate) async fn schedule_request_admitted( &self, request: ScheduleRequest, diff --git a/lib/llm/tests/data/selection_golden/aggregated_tracked.json b/lib/llm/tests/data/selection_golden/aggregated_tracked.json new file mode 100644 index 000000000000..d9906ad0e01b --- /dev/null +++ b/lib/llm/tests/data/selection_golden/aggregated_tracked.json @@ -0,0 +1,37 @@ +[ + { + "worker_id": 1, + "dp_rank": 0, + "overlap_blocks": 4, + "cached_tokens": 8, + "potential_decode_blocks": 4 + }, + { + "worker_id": 2, + "dp_rank": 0, + "overlap_blocks": 2, + "cached_tokens": 4, + "potential_decode_blocks": 8 + }, + { + "worker_id": 1, + "dp_rank": 0, + "overlap_blocks": 8, + "cached_tokens": 16, + "potential_decode_blocks": 12 + }, + { + "worker_id": 1, + "dp_rank": 0, + "overlap_blocks": 8, + "cached_tokens": 16, + "potential_decode_blocks": 12 + }, + { + "worker_id": 2, + "dp_rank": 0, + "overlap_blocks": 8, + "cached_tokens": 16, + "potential_decode_blocks": 14 + } +] From c6702146d4df9b15ef2d227215a18b11ebbff89a Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Thu, 10 Sep 2026 16:38:10 -0700 Subject: [PATCH 056/102] refactor(router): route the frontend through the shared selection core The frontend KvRouter ran its own copy of the selection pipeline: index lookup, shared-cache lookup, LoRA narrowing, transfer-hint computation and scheduler booking. It now builds one SelectionOperation with Lease admission and consumes the core's SelectionRun; the booking's lifecycle lease rides back in Selected.lease and is committed into the request's cleanup (KvRequestCleanup, prefill ReservationRelease, or the bindings' enroll path). Dropping an uncommitted lease frees the booking. - SelectionRun carries LookupTimings alongside the result so a host can account hashing, indexer and shared-cache time even when scheduling fails; the timings are recorded before the indexer error propagates and include the elapsed lookup wall time so scheduling is not overcharged. - Selected.endpoint is optional: a Lease admission whose worker left the catalog while the request queued still returns the selection, so the embedded host dispatches by worker id and the transport's connect failure drives migration as before. Wire admissions keep their 500. - The transfer-hint retention warning moves into the core next to the decision it describes. - The frontend hint helpers, route_lookup.rs, the without-admission path and lora_filter are deleted; the three frontend hint tests are ported onto transfer_hint_for_selection in the core. - A golden trace test and a dropped-lease test cover the cutover. Behavior changes: query-only selections attach no transfer hints; find_best_match_details_without_admission is removed; a request cancelled while queued no longer counts a completed shared-cache lookup failure. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Ishan Dhanani --- lib/bindings/python/rust/llm/kv.rs | 17 +- .../src/services/selection/core/mod.rs | 367 ++++- .../src/services/selection/core/operation.rs | 32 +- lib/kv-router/src/services/selection/input.rs | 32 +- lib/kv-router/src/services/selection/mod.rs | 7 +- lib/llm/src/kv_router.rs | 1191 +++-------------- lib/llm/src/kv_router/embedded.rs | 67 +- lib/llm/src/kv_router/prefill_router/query.rs | 17 +- lib/llm/src/kv_router/route_lookup.rs | 238 ---- lib/llm/src/kv_router/routing_host/kv.rs | 14 +- .../kv_router/routing_host/kv_selection.rs | 11 +- .../kv_router/routing_host/request_guard.rs | 44 +- lib/llm/src/kv_router/routing_host/tests.rs | 12 +- 13 files changed, 608 insertions(+), 1441 deletions(-) delete mode 100644 lib/llm/src/kv_router/route_lookup.rs diff --git a/lib/bindings/python/rust/llm/kv.rs b/lib/bindings/python/rust/llm/kv.rs index 1fdc028be3d6..1cea5c26df84 100644 --- a/lib/bindings/python/rust/llm/kv.rs +++ b/lib/bindings/python/rust/llm/kv.rs @@ -26,7 +26,6 @@ use dynamo_kv_router::config::try_kv_router_config_from_dynamo_env; use dynamo_kv_router::config::{KvRouterConfig, RouterConfigOverride}; use dynamo_kv_router::protocols::compute_block_hash_for_seq; use dynamo_kv_router::protocols::*; -use dynamo_kv_router::scheduling::AdmissionAttempt; #[cfg(feature = "kv-indexer")] use dynamo_kv_router::services::indexer::{self, IndexerConfig}; #[cfg(feature = "select-service")] @@ -2522,7 +2521,7 @@ impl KvRouter { ) .await .map_err(to_pyerr)?; - let (outcome, attempt) = admitted.into_parts(); + let (outcome, lease) = admitted.into_parts(); let (best_worker, overlap_blocks) = match outcome { llm_rs::kv_router::FindBestMatchOutcome::Routed { worker, @@ -2558,19 +2557,9 @@ impl KvRouter { None }; - if let Some(request_id) = request_id { - let AdmissionAttempt::Tracked(attempt_id) = attempt else { - return Err(PyErr::new::( - "tracked admission returned no attempt identity", - )); - }; + if let Some(lease) = lease { chooser - .enroll_public_request_attempt( - request_id, - best_worker, - attempt_id, - routing_decision, - ) + .enroll_public_request_attempt(lease, routing_decision) .await .map_err(to_pyerr)?; } else if let Some(tokens_with_hashes) = routing_decision { diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index 84b568abea6c..caeae7b686a2 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -10,6 +10,7 @@ use once_cell::sync::OnceCell; use parking_lot::RwLock; use tokio::sync::{mpsc, watch}; use tokio_util::sync::CancellationToken; +use tracing::Instrument; use crate::identity::RoutingPartitionId; use crate::indexer::{ @@ -51,7 +52,8 @@ use crate::tracking_hash::{TrackingHashContext, TrackingHashScope}; mod operation; pub use operation::{ - Selected, SelectionAdmission, SelectionOperation, SelectionOutcome, SessionBinding, + LookupTimings, Selected, SelectionAdmission, SelectionOperation, SelectionOutcome, + SelectionRun, SessionBinding, }; use super::affinity::{ @@ -1014,7 +1016,7 @@ impl SelectionCore { request_id: req.selection_id.clone(), } }; - let selected = self + let (selected, endpoint) = self .select_or_reject(SelectionOperation { key: RoutingPartitionId::new(req.model_name, req.routing_group), prompt: req.prompt.view(), @@ -1035,7 +1037,7 @@ impl SelectionCore { replay_id: req.selection_id.clone(), }) .await?; - Ok(self.select_response(selected, req.selection_id)) + Ok(self.select_response(selected, endpoint, req.selection_id)) } pub async fn select_and_reserve( @@ -1056,7 +1058,7 @@ impl SelectionCore { .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); let is_steerable = req.affinity_target.is_none() && req.pinned_worker.is_none(); let session = session_binding(session_context.as_ref(), is_steerable, true); - let selected = self + let (selected, endpoint) = self .select_or_reject(SelectionOperation { key: RoutingPartitionId::new(req.model_name, req.routing_group), prompt: req.prompt.view(), @@ -1079,29 +1081,44 @@ impl SelectionCore { replay_id: None, }) .await?; - Ok(self.select_response(selected, Some(selection_id))) + Ok(self.select_response(selected, endpoint, Some(selection_id))) } /// The wire shape of a selection; a queue rejection keeps its 503 body. async fn select_or_reject( &self, operation: SelectionOperation<'_>, - ) -> Result { - match self.run_selection(operation).await? { - SelectionOutcome::Selected(selected) => Ok(selected), - SelectionOutcome::QueueRejected { rejection } => Err(SelectionError::Scheduler( - KvSchedulerError::QueueRejected(rejection), - )), - } + ) -> Result<(Selected, String), SelectionError> { + let mut selected = match self.run_selection(operation).await.result? { + SelectionOutcome::Selected(selected) => selected, + SelectionOutcome::QueueRejected { rejection } => { + return Err(SelectionError::Scheduler(KvSchedulerError::QueueRejected( + rejection, + ))); + } + }; + // Wire admissions never book without an endpoint; see `Selected::endpoint`. + let endpoint = selected.endpoint.take().ok_or_else(|| { + SelectionError::Internal(format!( + "selected worker {} is no longer schedulable", + selected.response.best_worker.worker_id + )) + })?; + Ok((selected, endpoint)) } - fn select_response(&self, selected: Selected, selection_id: Option) -> SelectResponse { + fn select_response( + &self, + selected: Selected, + endpoint: String, + selection_id: Option, + ) -> SelectResponse { let Selected { key, response, advisory_load, total_kv_blocks, - endpoint, + endpoint: _, block_size, isl_tokens, sequence_hashes, @@ -1109,6 +1126,7 @@ impl SelectionCore { effective_prefill_tokens, kv_hint, routing_hashes: _, + shared_cache_hits: _, lease: _, } = selected; let booked = sequence_hashes.is_some(); @@ -1156,9 +1174,16 @@ impl SelectionCore { /// Run one selection: resolve the session, look the prompt up, schedule, /// and (for `Book`) install the reservation. Every host's selection goes /// through here. - pub async fn run_selection( + pub async fn run_selection(&self, operation: SelectionOperation<'_>) -> SelectionRun { + let mut lookup = None; + let result = self.run_selection_inner(operation, &mut lookup).await; + SelectionRun { result, lookup } + } + + async fn run_selection_inner( &self, operation: SelectionOperation<'_>, + lookup: &mut Option, ) -> Result { let SelectionOperation { key, @@ -1217,9 +1242,17 @@ impl SelectionCore { // Router hints are attached to bookings only, and only when a worker in // this partition can consume them and the indexer can retain the // matched chain (local, event-driven, no approximate writes). - let retain_kv_transfer_chain = book - && entry.indexer.supports_kv_transfer_chain_retention() - && self.catalog.has_router_hint_capable_workers(&key); + let hint_capable_workers = book && self.catalog.has_router_hint_capable_workers(&key); + let retain_kv_transfer_chain = + hint_capable_workers && entry.indexer.supports_kv_transfer_chain_retention(); + if hint_capable_workers && !retain_kv_transfer_chain { + static WARN_ONCE: std::sync::Once = std::sync::Once::new(); + WARN_ONCE.call_once(|| { + tracing::warn!( + "router hints need a local event-driven indexer with no approximate side indexer; workers advertise hint capability but no hints will be attached" + ); + }); + } let PreparedSelectionInputs { block_hashes, sequence_hashes, @@ -1237,8 +1270,21 @@ impl SelectionCore { }), true, retain_kv_transfer_chain, + lookup, ) .await?; + if tracing::enabled!(tracing::Level::DEBUG) { + let local_hashes: Vec = block_hashes.iter().map(|hash| hash.0).collect(); + tracing::debug!( + request_id = admission.request_id().unwrap_or_default(), + isl_tokens, + block_size = entry.block_size, + num_blocks = local_hashes.len(), + ?local_hashes, + "[ROUTING_INPUT] request local hashes" + ); + } + let host_shared_cache_hits = shared_cache_hits.clone(); // The queue lease frees a booking whose response is never consumed // (the caller dropped this future after the actor booked). let mode = match &admission { @@ -1323,6 +1369,7 @@ impl SelectionCore { entry .scheduler .select_without_admission(schedule_request) + .instrument(tracing::info_span!("kv_router.select_without_admission")) .await .map(|advisory| { (advisory.response, Some(advisory.selected_worker_load), None) @@ -1331,6 +1378,7 @@ impl SelectionCore { entry .scheduler .schedule_request_with_lease(schedule_request) + .instrument(tracing::info_span!("kv_router.schedule")) .await .map(|(admitted, lease)| (admitted.response, None, lease)) } @@ -1343,15 +1391,15 @@ impl SelectionCore { } Err(error) => return Err(error.into()), }; - let Some(endpoint) = self + let endpoint = self .catalog - .schedulable_endpoint(response.best_worker.worker_id, &key) - else { + .schedulable_endpoint(response.best_worker.worker_id, &key); + if endpoint.is_none() && !matches!(admission, SelectionAdmission::Lease { .. }) { return Err(SelectionError::Internal(format!( "selected worker {} is no longer schedulable", response.best_worker.worker_id ))); - }; + } let effective_prefill = effective_prefill_tokens(isl_tokens, response.cached_tokens); let total_kv_blocks = advisory_load .and_then(|load| load.total_kv_blocks.map(|blocks| blocks as u64)) @@ -1376,8 +1424,6 @@ impl SelectionCore { None }; - // The routing hashes go to exactly one of: the booking recorded now, or - // the cached inputs a later replay records. // The routing hashes go to exactly one of: the reservation recorded // now, or the replay cache a later reservation records from. let mut routing_hashes = routing_hashes; @@ -1442,6 +1488,7 @@ impl SelectionCore { effective_prefill_tokens: effective_prefill, kv_hint, routing_hashes: returned_routing_hashes, + shared_cache_hits: host_shared_cache_hits, lease, })) } @@ -1993,6 +2040,7 @@ impl SelectionCore { ), false, false, + &mut None, ) .await?; let track_prefill_tokens = req @@ -2051,6 +2099,7 @@ impl SelectionCore { tracking_assume_kv_reuse: Option, query_shared_cache: bool, retain_kv_transfer_chain: bool, + lookup: &mut Option, ) -> Result { let normalized = prompt.normalize_for_selection( entry.block_size, @@ -2062,7 +2111,8 @@ impl SelectionCore { }), )?; let indexer_lookup = async { - if normalized.block_hashes.is_empty() { + let started = Instant::now(); + let tiered = if normalized.block_hashes.is_empty() { Ok(TieredMatchDetails::default()) } else { entry @@ -2073,28 +2123,49 @@ impl SelectionCore { retain_kv_transfer_chain, }, ) + .instrument(tracing::info_span!("kv_router.find_matches")) .await .map_err(SelectionError::Indexer) - } + }; + (tiered, started.elapsed()) }; let shared_cache = query_shared_cache .then_some(self.host.cache.shared.as_deref()) .flatten() .zip(prompt.token_ids); + // `(hits, duration, failed)`; all `None`/`false` without a shared cache. let shared_cache_lookup = async { - let (shared_cache, tokens) = shared_cache?; - match shared_cache + let Some((shared_cache, tokens)) = shared_cache else { + return (None, None, false); + }; + let started = Instant::now(); + let result = shared_cache .check_blocks(tokens, entry.block_size, prompt.cache_namespace) - .await - { - Ok(hits) => Some(hits), + .instrument(tracing::info_span!("kv_router.shared_cache_check")) + .await; + let elapsed = started.elapsed(); + match result { + Ok(hits) => (Some(hits), Some(elapsed), false), Err(error) => { tracing::warn!(%error, "Shared cache query failed, ignoring"); - None + (None, Some(elapsed), true) } } }; - let (tiered, shared_cache_hits) = tokio::join!(indexer_lookup, shared_cache_lookup); + let lookups_started = Instant::now(); + let ((tiered, indexer), (shared_cache_hits, shared_cache, shared_cache_error)) = + tokio::join!(indexer_lookup, shared_cache_lookup); + let timings = LookupTimings { + block_hashing: normalized.block_hashing, + seq_hashing: normalized.seq_hashing, + lookups: lookups_started.elapsed(), + indexer, + shared_cache, + shared_cache_error, + }; + // Recorded before `?` so a shared-cache failure is still accounted + // when the index lookup fails. + *lookup = Some(timings); let tiered = tiered?; let overlap = OverlapAnalysis::new(&self.kv_router_config, entry.block_size, &tiered).signals(); @@ -3467,6 +3538,82 @@ mod tests { wait_until("booking release", || !entry.scheduler.has_request("queued")).await; } + /// The embedded host dispatches by worker id, so a worker that drains + /// while the request queues still comes back selected: the transport + /// reports the departure and migration takes over. The endpoint is the + /// only thing missing. + #[tokio::test] + async fn lease_admission_keeps_a_selection_whose_worker_drained_while_queued() { + let mut config = test_config(false); + config.router_queue_threshold = Some(0.0); + let core = Arc::new( + SelectionCore::try_new_local( + config, + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ) + .expect("valid test config"), + ); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + let key = RoutingPartitionId::new("model", "default"); + let entry = core.entry(&key).expect("entry"); + core.select_and_reserve(reserve_request("res-a")) + .await + .expect("initial reservation"); + + let queued_core = core.clone(); + let queued = tokio::spawn(async move { + let req = reserve_request("queued"); + let run = queued_core + .run_selection(SelectionOperation { + key: RoutingPartitionId::new("model", "default"), + prompt: req.prompt.view(), + router_config_override: None, + expected_output_tokens: None, + priority_jump: 0.0, + strict_priority: 0, + policy_class: None, + session_context: None, + session: SessionBinding::None, + affinity_target: None, + pinned_worker: None, + allowed_worker_ids: None, + routing_constraints: RoutingConstraints::default(), + admission: SelectionAdmission::Lease { + request_id: "queued".to_string(), + }, + track_active_blocks: false, + return_routing_hashes: false, + replay_id: None, + }) + .await; + match run.result { + Ok(SelectionOutcome::Selected(selected)) => ( + selected.response.best_worker.worker_id, + selected.endpoint, + selected.lease.is_some(), + ), + Ok(SelectionOutcome::QueueRejected { .. }) => panic!("queue rejected"), + Err(error) => panic!("lease selection failed: {error:?}"), + } + }); + wait_for_pending_selection(&core).await; + core.catalog + .set_lifecycle(1, WorkerLifecycle::Draining, Vec::new()); + core.free_reservation("res-a").await.expect("free res-a"); + + let (worker_id, endpoint, leased) = tokio::time::timeout(Duration::from_secs(2), queued) + .await + .expect("queued selection timed out") + .expect("task panicked"); + assert_eq!(worker_id, 1); + assert_eq!(endpoint, None); + assert!(leased, "lease admission returns the booking's lease"); + // The lease was dropped with the selection, so the booking is gone. + wait_until("booking release", || !entry.scheduler.has_request("queued")).await; + } + #[tokio::test] async fn dropped_selection_future_frees_its_booking() { let core = SelectionCore::try_new_local( @@ -4849,4 +4996,158 @@ mod tests { Err(SelectionError::Conflict(_)) )); } + + fn hint_config(worker_type: &str, endpoints: &[(u32, &str)]) -> SelectionWorkerConfig { + SelectionWorkerConfig { + endpoint: "http://worker:8000".to_string(), + data_parallel_start_rank: 0, + data_parallel_size: endpoints.len().max(1) as u32, + max_num_batched_tokens: None, + total_kv_blocks: None, + stable_routing_id: None, + is_eagle: None, + taints: HashSet::new(), + topology_domains: HashMap::new(), + kv_transfer_domain: None, + kv_transfer_enforcement: None, + kv_transfer_preferred_weight: None, + router_hint_worker_type: Some(worker_type.to_string()), + router_hint_source_control_endpoints: endpoints + .iter() + .map(|(rank, endpoint)| (*rank, endpoint.to_string())) + .collect(), + kv_event_source_mode: None, + } + } + + fn hint_candidates( + hashes: &[u64], + owners: Vec<(KvTransferCandidateSource, usize)>, + ) -> KvTransferCandidates { + KvTransferCandidates { + block_hashes: hashes + .iter() + .map(|h| crate::protocols::ExternalSequenceBlockHash(*h)) + .collect(), + owner_prefix_blocks: owners, + routing_snapshot: None, + } + } + + #[test] + fn hint_source_endpoint_follows_the_source_dp_rank() { + let configs = HashMap::from([( + 7, + hint_config( + "prefill", + &[(0, "tcp://127.0.0.1:23280"), (1, "tcp://127.0.0.1:23281")], + ), + )]); + let candidates = + hint_candidates(&[101, 102], vec![(WorkerWithDpRank::new(7, 1).into(), 2)]); + let hint = transfer_hint_for_selection( + &configs, + WorkerWithDpRank::new(7, 0), + 0, + Some(&candidates), + ); + assert_eq!( + hint.map(|payload| payload.source_control_endpoint), + Some("tcp://127.0.0.1:23281".to_string()) + ); + } + + #[test] + fn hint_source_must_share_the_target_worker_type() { + let configs = HashMap::from([ + (7, hint_config("prefill", &[(0, "tcp://127.0.0.1:23280")])), + (8, hint_config("prefill", &[(0, "tcp://127.0.0.1:23281")])), + (9, hint_config("decode", &[(0, "tcp://127.0.0.1:23282")])), + (10, hint_config("decode", &[(0, "tcp://127.0.0.1:23283")])), + ]); + let candidates = hint_candidates( + &[101, 102, 103], + vec![ + (WorkerWithDpRank::new(8, 0).into(), 2), + (WorkerWithDpRank::new(9, 0).into(), 3), + ], + ); + // The longer decode prefix is skipped for a prefill target. + let prefill = transfer_hint_for_selection( + &configs, + WorkerWithDpRank::new(7, 0), + 0, + Some(&candidates), + ) + .expect("prefill hint"); + assert_eq!(prefill.source_control_endpoint, "tcp://127.0.0.1:23281"); + assert_eq!(prefill.block_hashes.len(), 2); + let decode = transfer_hint_for_selection( + &configs, + WorkerWithDpRank::new(10, 0), + 0, + Some(&candidates), + ) + .expect("decode hint"); + assert_eq!(decode.source_control_endpoint, "tcp://127.0.0.1:23282"); + assert_eq!(decode.block_hashes.len(), 3); + } + + #[test] + fn hint_resolves_a_persistent_cache_owner_over_a_state_agent_worker() { + use crate::identity::{ + CacheOwnerId, CacheSemanticsId, DcId, IdentitySource, IndexerDomainId, PoolId, + RoutingScopeId, StableDpSlotId, + }; + use crate::protocols::{ + ResidencyOwner, ResidencyProjection, ResidencyRoutingSnapshot, RouterHintSourceMetadata, + }; + let mut stale = hint_config("prefill", &[(0, "tcp://stale-worker:23280")]); + stale.kv_event_source_mode = Some("state_agent_v2".to_string()); + let configs = HashMap::from([(7, hint_config("prefill", &[])), (8, stale)]); + let owner = CacheOwnerId::new( + PoolId::new( + IndexerDomainId::new( + CacheSemanticsId::new([1; 16], IdentitySource::Explicit), + RoutingScopeId::new([2; 16], IdentitySource::Explicit), + ), + DcId::new(3), + ), + StableDpSlotId::new([4; 16], IdentitySource::Explicit), + ); + let mut candidates = hint_candidates( + &[101, 102], + vec![ + ( + KvTransferCandidateSource::Worker(WorkerWithDpRank::new(8, 0)), + 2, + ), + ( + KvTransferCandidateSource::CacheOwner( + ResidencyOwner::cache_owner(owner).compact_key(), + ), + 2, + ), + ], + ); + candidates.routing_snapshot = Some(Arc::new(ResidencyRoutingSnapshot::new( + ResidencyProjection::default(), + [( + owner, + RouterHintSourceMetadata { + source_control_endpoint: "tcp://persistent-owner:23280".to_string(), + worker_type: "prefill".to_string(), + }, + None, + )], + ))); + let hint = transfer_hint_for_selection( + &configs, + WorkerWithDpRank::new(7, 0), + 0, + Some(&candidates), + ) + .expect("hint"); + assert_eq!(hint.source_control_endpoint, "tcp://persistent-owner:23280"); + } } diff --git a/lib/kv-router/src/services/selection/core/operation.rs b/lib/kv-router/src/services/selection/core/operation.rs index 17ccd1865bc2..8d0a47baca82 100644 --- a/lib/kv-router/src/services/selection/core/operation.rs +++ b/lib/kv-router/src/services/selection/core/operation.rs @@ -5,18 +5,21 @@ //! hosts build a [`SelectionOperation`] and consume a [`SelectionOutcome`]. use std::collections::HashSet; +use std::time::Duration; use dynamo_tokens::SequenceHash; use crate::identity::RoutingPartitionId; use crate::kv_hints::KvHint; use crate::protocols::{ - LocalBlockHash, RoutingConstraints, WorkerAffinityTarget, WorkerId, WorkerWithDpRank, + LocalBlockHash, RoutingConstraints, SharedCacheHits, WorkerAffinityTarget, WorkerId, + WorkerWithDpRank, }; use crate::scheduling::config::RouterConfigOverride; use crate::scheduling::queue::RequestLifecycleLease; use crate::scheduling::{AdvisoryWorkerLoad, QueueRejection, SchedulingResponse, SessionContext}; +use super::super::error::SelectionError; use super::super::input::PromptView; /// Every input to one selection. Fields are independent: `session_context` @@ -89,6 +92,27 @@ pub enum SessionBinding { }, } +/// One selection's result together with how long its lookups took. The +/// timings ride with errors too, so a host can account a shared-cache failure +/// even when scheduling then fails. +pub struct SelectionRun { + pub result: Result, + /// `None` when selection failed before the lookups started. + pub lookup: Option, +} + +#[derive(Debug, Clone, Copy, Default)] +pub struct LookupTimings { + pub block_hashing: Duration, + pub seq_hashing: Duration, + /// Wall time of the index and shared-cache lookups together. + pub lookups: Duration, + pub indexer: Duration, + /// `None` when no shared cache was queried. + pub shared_cache: Option, + pub shared_cache_error: bool, +} + // Every caller matches this immediately; boxing the common variant to shrink // the rare one would put an allocation on the hot path. #[allow(clippy::large_enum_variant)] @@ -105,7 +129,10 @@ pub struct Selected { pub advisory_load: Option, /// The chosen worker's KV capacity as known when it was selected. pub total_kv_blocks: Option, - pub endpoint: String, + /// `None` only for `Lease` admission when the worker left the catalog + /// after it was booked: the host dispatches by worker id and lets the + /// transport report the departure, which is what drives migration. + pub endpoint: Option, pub block_size: u32, pub isl_tokens: usize, /// The hashes the booking tracks; `Book` only. @@ -114,6 +141,7 @@ pub struct Selected { pub effective_prefill_tokens: usize, pub kv_hint: Option, pub routing_hashes: Option>, + pub shared_cache_hits: Option, /// The booking's lifecycle lease; `Lease` admission only. Dropping it frees /// the booking, `commit` hands it to the caller's own cleanup. pub lease: Option>, diff --git a/lib/kv-router/src/services/selection/input.rs b/lib/kv-router/src/services/selection/input.rs index 55ada4e9b26b..048622692f05 100644 --- a/lib/kv-router/src/services/selection/input.rs +++ b/lib/kv-router/src/services/selection/input.rs @@ -1,6 +1,8 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +use std::time::{Duration, Instant}; + use dynamo_tokens::SequenceHash; use serde::Deserialize; @@ -197,20 +199,27 @@ fn normalize_tokens_for_selection( cache_namespace, is_eagle: Some(is_eagle), }; - let block_hashes = compute_block_hash_for_seq(token_ids, block_size, hash_options); - let sequence_hashes = tracking.map_or_else(Vec::new, |tracking| { - tracking.context.compute_sequence_hashes_for_tracking( - tracking.scope, - token_ids, - hash_options, - tracking.assume_kv_reuse, - Some(&block_hashes), - ) + let started = Instant::now(); + let block_hashes = tracing::info_span!("kv_router.compute_block_hashes") + .in_scope(|| compute_block_hash_for_seq(token_ids, block_size, hash_options)); + let block_hashing = started.elapsed(); + let sequence_hashes = tracing::info_span!("kv_router.compute_seq_hashes").in_scope(|| { + tracking.map_or_else(Vec::new, |tracking| { + tracking.context.compute_sequence_hashes_for_tracking( + tracking.scope, + token_ids, + hash_options, + tracking.assume_kv_reuse, + Some(&block_hashes), + ) + }) }); NormalizedPrompt { block_hashes, sequence_hashes, isl_tokens: token_ids.len(), + block_hashing, + seq_hashing: started.elapsed().saturating_sub(block_hashing), } } @@ -265,6 +274,8 @@ fn normalize_hashes( .collect(), sequence_hashes: signed_sequence_hashes(sequence_hashes), isl_tokens, + block_hashing: Duration::ZERO, + seq_hashing: Duration::ZERO, }) } @@ -276,6 +287,9 @@ pub(super) struct NormalizedPrompt { pub(super) block_hashes: Vec, pub(super) sequence_hashes: Vec, pub(super) isl_tokens: usize, + /// Zero for hash-only inputs. + pub(super) block_hashing: Duration, + pub(super) seq_hashing: Duration, } pub(super) struct NormalizedReservation { diff --git a/lib/kv-router/src/services/selection/mod.rs b/lib/kv-router/src/services/selection/mod.rs index 099bcdd172d5..dbf0a2cbb841 100644 --- a/lib/kv-router/src/services/selection/mod.rs +++ b/lib/kv-router/src/services/selection/mod.rs @@ -28,9 +28,10 @@ pub use crate::services::common::replica_sync::{ ScopedSequencePublisher, }; pub use core::{ - HostCache, HostEligibility, HostLoad, HostReplication, HostTelemetry, KvIndexSource, Selected, - SelectionAdmission, SelectionCore, SelectionHost, SelectionOperation, SelectionOutcome, - SelectionPartition, SelectionScheduler, SelectionServiceConfig, SessionBinding, + HostCache, HostEligibility, HostLoad, HostReplication, HostTelemetry, KvIndexSource, + LookupTimings, Selected, SelectionAdmission, SelectionCore, SelectionHost, SelectionOperation, + SelectionOutcome, SelectionPartition, SelectionRun, SelectionScheduler, SelectionServiceConfig, + SessionBinding, }; pub use error::SelectionError; pub use ingress::{KvEventIngress, ZmqDirectIngress}; diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index 1c49d4ecb5fb..95b3c6d1cca7 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -17,25 +17,22 @@ use dynamo_kv_router::{ indexer::{ ApproximateLruIncarnation, ApproximateLruStats, KvRouterError, RoutingDecisionHashes, }, - kv_hints::{ - KvHint, KvHintAction, KvSourceLocationsPayload, KvTransferCandidateSource, - KvTransferCandidates, - }, + kv_hints::KvHint, protocols::KV_EVENT_SUBJECT, protocols::{ - BlockExtraInfo, BlockHashOptions, LocalBlockHash, PrefillLoadHint, RouterEvent, - RouterRequest, RouterResponse, RoutingConstraints, TokensWithHashes, WorkerConfigLike, - WorkerId, WorkerWithDpRank, compute_block_hash_for_seq, + BlockExtraInfo, BlockHashOptions, PrefillLoadHint, RouterEvent, RouterRequest, + RouterResponse, RoutingConstraints, TokensWithHashes, WorkerConfigLike, WorkerId, + WorkerWithDpRank, compute_block_hash_for_seq, }, scheduling::{ - AdmissionAttempt, AttemptId, CacheHitEstimates, OverlapAnalysis, OverloadedWorkerProvider, - ScheduleMode, ScheduleRequest, WorkerAvailabilityProvider, effective_prefill_tokens, - overlap::cache_hit_estimates_from_tiered_matches, + CacheHitEstimates, OverlapAnalysis, OverloadedWorkerProvider, WorkerAvailabilityProvider, + effective_prefill_tokens, overlap::cache_hit_estimates_from_tiered_matches, + queue::RequestLifecycleLease, }, selector::WorkerInputs, services::selection::{ PromptView, Selected, SelectionAdmission, SelectionError, SelectionOperation, - SelectionOutcome, SessionBinding, + SelectionOutcome, SelectionRun, SessionBinding, }, }; use dynamo_runtime::{ @@ -52,7 +49,6 @@ use dynamo_runtime::{ traits::DistributedRuntimeProvider, }; use futures::stream; -use tracing::Instrument; // Re-export from dynamo-kv-router crate pub use dynamo_kv_router::approx; @@ -68,7 +64,6 @@ pub(crate) mod metrics_subscriber; pub mod prefill_router; pub mod publisher; mod request_lease; -mod route_lookup; mod routing_host; pub(crate) mod routing_load; pub mod scheduler; @@ -96,9 +91,6 @@ use crate::{ local_model::runtime_config::ModelRuntimeConfig, worker_type::WorkerType, }; -use route_lookup::{ - TieredLookupOptions, TieredLookupResult, query_tiered_matches, split_retained_block_hashes, -}; /// Where a `KvRouter` gets the worker-selection policy its partition runs. #[derive(Clone, Default)] @@ -436,20 +428,23 @@ pub enum FindBestMatchAdvisoryOutcome { #[derive(Debug, Clone, Copy)] pub(super) enum FindBestMatchAdmission { - WithAdmission { track_lifecycle: bool }, + WithAdmission, WithoutAdmission, } +/// A routed outcome with the booking's lifecycle lease, when the request was +/// booked. Dropping the lease frees the booking; the caller commits it into +/// its own cleanup. #[doc(hidden)] pub struct AdmittedFindBestMatchOutcome { pub(super) outcome: FindBestMatchOutcome, - pub(super) attempt: AdmissionAttempt, + pub(super) lease: Option>, } impl AdmittedFindBestMatchOutcome { #[doc(hidden)] - pub fn into_parts(self) -> (FindBestMatchOutcome, AdmissionAttempt) { - (self.outcome, self.attempt) + pub fn into_parts(self) -> (FindBestMatchOutcome, Option>) { + (self.outcome, self.lease) } } @@ -541,28 +536,6 @@ fn cancelled_error(context_id: &str) -> anyhow::Error { .into() } -fn log_routing_input_hashes( - request_id: Option<&str>, - block_size: u32, - tokens: &[u32], - local_hashes: &[LocalBlockHash], -) { - if !tracing::enabled!(tracing::Level::DEBUG) { - return; - } - - let local_hash_ids: Vec = local_hashes.iter().map(|hash| hash.0).collect(); - - tracing::debug!( - request_id = request_id.unwrap_or(""), - isl_tokens = tokens.len(), - block_size, - num_blocks = local_hashes.len(), - local_hashes = ?local_hash_ids, - "[ROUTING_INPUT] request local hashes" - ); -} - // for router discovery registration pub const KV_ROUTER_ENDPOINT: &str = "router-discovery"; @@ -605,10 +578,6 @@ pub struct KvRouter { /// Optional external shared KV cache pool. When present, `find_best_match` /// queries it in parallel with the indexer and factors shared hits into scoring. shared_cache: Option>, - /// Optional LoRA filter. When present (LoRA serving enabled), candidate workers are - /// narrowed to the LoRA's allocated/loaded replicas inside `find_best_match_details`, - /// covering both the decode and prefill routers (both built via `kv_chooser_for`). - lora_filter: Option>, endpoint_registration: Option, teardown_task_guard: Option, } @@ -837,7 +806,7 @@ impl KvRouter { overloaded_worker_provider, available_worker_provider, shared_cache: shared_cache.clone(), - lora_worker_filter: lora_filter.clone().map(|filter| { + lora_worker_filter: lora_filter.map(|filter| { filter as Arc }), ingress: Arc::clone(&ingress) @@ -876,7 +845,6 @@ impl KvRouter { approximate_lru_ranks, request_leases, shared_cache, - lora_filter, endpoint_registration: None, teardown_task_guard: None, }) @@ -1007,107 +975,6 @@ impl KvRouter { cache_hit_for_worker(cache_hit_estimates, worker) } - fn has_transfer_capable_workers(&self) -> bool { - // TRANSFER capability is worker-level metadata. Check one - // representative DP rank here so the coarse request-path gate does not - // scale with data_parallel_size. Follow-up: cache this from the runtime - // config watch if the per-worker scan shows up in large-fleet routing - // benchmarks. - self.workers_with_configs.borrow().values().any(|config| { - config - .kv_hint_transfer_metadata_for_dp_rank(config.data_parallel_start_rank()) - .is_some() - }) - } - - fn transfer_hint_for_selection( - &self, - target: WorkerWithDpRank, - target_cached_prefix_blocks: u32, - candidates: Option<&KvTransferCandidates>, - ) -> Option { - let candidates = candidates?; - - let (block_hashes, source_control_endpoint) = { - let configs = self.workers_with_configs.borrow(); - let target_config = configs.get(&target.worker_id)?; - let target_metadata = - target_config.kv_hint_transfer_metadata_for_dp_rank(target.dp_rank)?; - - let prefix_blocks_to_beat = - usize::try_from(target_cached_prefix_blocks).unwrap_or(usize::MAX); - let (source, block_hashes) = - candidates.best_source(prefix_blocks_to_beat, |source| match source { - KvTransferCandidateSource::Worker(worker) => { - worker != target - && configs.get(&worker.worker_id).is_some_and(|config| { - config.kv_event_source_mode.as_deref() != Some("state_agent_v2") - && config - .kv_hint_transfer_metadata_for_dp_rank(worker.dp_rank) - .is_some_and(|source_metadata| { - source_metadata.worker_type - == target_metadata.worker_type - && source_metadata - .source_control_endpoint - .is_some_and(|endpoint| !endpoint.is_empty()) - }) - }) - } - KvTransferCandidateSource::CacheOwner(owner) => candidates - .routing_snapshot - .as_ref() - .and_then(|snapshot| snapshot.router_hint_source(owner)) - .is_some_and(|source| { - source.attached_worker != Some(target) - && source.metadata.worker_type == target_metadata.worker_type - && !source.metadata.source_control_endpoint.is_empty() - }), - })?; - let source_control_endpoint = match source { - KvTransferCandidateSource::Worker(worker) => configs - .get(&worker.worker_id)? - .kv_hint_transfer_metadata_for_dp_rank(worker.dp_rank)? - .source_control_endpoint? - .to_string(), - KvTransferCandidateSource::CacheOwner(owner) => candidates - .routing_snapshot - .as_ref()? - .router_hint_source(owner)? - .metadata - .source_control_endpoint - .clone(), - }; - (block_hashes, source_control_endpoint) - }; - - if block_hashes.is_empty() { - return None; - } - - Some(KvSourceLocationsPayload { - source_control_endpoint, - block_hashes, - }) - } - - fn kv_hint_for_selection( - &self, - message_id: Option<&str>, - target: WorkerWithDpRank, - target_cached_prefix_blocks: u32, - transfer_candidates: Option<&KvTransferCandidates>, - ) -> Option { - let payload = self.transfer_hint_for_selection( - target, - target_cached_prefix_blocks, - transfer_candidates, - )?; - Some(KvHint::new( - message_id.unwrap_or_default(), - vec![KvHintAction::fetch("a1", payload)], - )) - } - pub async fn record_routing_decision( &self, mut tokens_with_hashes: TokensWithHashes, @@ -1145,21 +1012,22 @@ impl KvRouter { #[doc(hidden)] pub async fn enroll_public_request_attempt( &self, - request_id: String, - worker: WorkerWithDpRank, - attempt_id: AttemptId, + lease: Box, routing_decision: Option, ) -> Result<(), KvRouterError> { + // Nothing awaits between taking the booking over and registering it. + let booking = lease.commit().ok_or_else(|| { + KvRouterError::Unsupported("booking lease holds no booking".to_string()) + })?; + let worker = booking.worker; let lru_registration = self.approximate_lru_rank_registration(worker); let approximate_lru = lru_registration.and_then(|registration| { - self.indexer - .begin_approximate_lru_request(worker, registration.incarnation, attempt_id) + self.indexer.begin_approximate_lru_request( + worker, + registration.incarnation, + booking.attempt_id, + ) }); - let booking = scheduler::SchedulerBookingDescriptor { - request_id, - worker, - attempt_id, - }; let enrollment = self .request_leases .register_detached(booking, approximate_lru.clone()); @@ -1223,34 +1091,6 @@ impl KvRouter { .await } - /// Narrow the candidate workers to this LoRA's allocated/loaded replicas, staying strictly - /// within the existing candidate universe (never widening). Returns the (possibly narrowed) - /// `allowed_worker_ids` to pass to the scheduler. - /// - /// - No filter (LoRA serving disabled) or base-model request (`lora_name` is `None`): - /// returns `allowed_worker_ids` unchanged. - /// - Pinned worker: KV-cache correctness wins — it is always retained even if not in the - /// LoRA replica set (the worker lazy-loads the adapter). - /// - If narrowing would exclude every candidate, falls back to the original set so the - /// request stays routable (lazy-load path) rather than failing. - fn narrow_allowed_by_lora( - &self, - lora_name: Option<&str>, - allowed_worker_ids: Option>, - pinned_worker: Option<&WorkerWithDpRank>, - ) -> Option> { - let Some(filter) = self.lora_filter.as_ref() else { - return allowed_worker_ids; - }; - dynamo_kv_router::scheduling::narrow_allowed_worker_ids_by_lora( - filter.as_ref(), - lora_name, - allowed_worker_ids, - pinned_worker, - || self.workers_with_configs.borrow().keys().copied().collect(), - ) - } - /// Give these tokens, find the worker with the best weighted cache hit. /// Returns the full match details for the selected worker. /// @@ -1340,9 +1180,7 @@ impl KvRouter { pinned_worker, allowed_worker_ids, routing_constraints, - FindBestMatchAdmission::WithAdmission { - track_lifecycle: true, - }, + FindBestMatchAdmission::WithAdmission, ) .await? { @@ -1393,14 +1231,8 @@ impl KvRouter { routing_constraints, ) .await?; - if let ( - Some(request_id), - FindBestMatchOutcome::Routed { worker, .. }, - AdmissionAttempt::Tracked(attempt_id), - ) = (context_id, &admitted.outcome, admitted.attempt) - { - self.enroll_public_request_attempt(request_id.to_string(), *worker, attempt_id, None) - .await?; + if let Some(lease) = admitted.lease { + self.enroll_public_request_attempt(lease, None).await?; } Ok(admitted.outcome) } @@ -1448,9 +1280,7 @@ impl KvRouter { pinned_worker, allowed_worker_ids, routing_constraints, - FindBestMatchAdmission::WithAdmission { - track_lifecycle: false, - }, + FindBestMatchAdmission::WithAdmission, ) .await? { @@ -1462,12 +1292,13 @@ impl KvRouter { } #[allow(clippy::too_many_arguments)] - pub async fn find_best_match_details_without_admission( + async fn find_best_match_details_with_policy_class_inner( &self, context_id: Option<&str>, tokens: &[u32], block_mm_infos: Option<&[Option]>, router_config_override: Option<&RouterConfigOverride>, + update_states: bool, return_routing_hashes: bool, lora_name: Option, cache_namespace: Option, @@ -1480,79 +1311,28 @@ impl KvRouter { pinned_worker: Option, allowed_worker_ids: Option>, routing_constraints: RoutingConstraints, - ) -> anyhow::Result { - match self - .find_best_match_details_with_policy_class_inner( - context_id, - tokens, - block_mm_infos, - router_config_override, - false, - return_routing_hashes, - lora_name, - cache_namespace, - priority_jump, - strict_priority, - policy_class, - session_context, - expected_output_tokens, - affinity_target, - pinned_worker, - allowed_worker_ids, - routing_constraints, - FindBestMatchAdmission::WithoutAdmission, - ) - .await? - { - FindBestMatchInnerOutcome::WithoutAdmission(outcome) => Ok(outcome), - FindBestMatchInnerOutcome::WithAdmission(_) => { - unreachable!("without-admission routing returned admitted outcome") - } - } - } - - /// The shared core's selection for the same inputs as - /// `find_best_match_details_with_policy_class_inner`. - #[allow(clippy::too_many_arguments)] - #[cfg_attr(not(test), expect(dead_code))] - async fn select_core( - &self, - context_id: Option<&str>, - tokens: &[u32], - block_mm_infos: Option<&[Option]>, - router_config_override: Option<&RouterConfigOverride>, - update_states: bool, - return_routing_hashes: bool, - lora_name: Option<&str>, - cache_namespace: Option<&str>, - priority_jump: f64, - strict_priority: u32, - policy_class: Option, - session_context: Option, - expected_output_tokens: Option, - affinity_target: Option, - pinned_worker: Option, - allowed_worker_ids: Option>, - routing_constraints: RoutingConstraints, admission: FindBestMatchAdmission, ) -> anyhow::Result { + let start = Instant::now(); if update_states && context_id.is_none() { anyhow::bail!("context_id must be provided if update_states is true"); } + let is_admitted_routing = matches!(admission, FindBestMatchAdmission::WithAdmission); let core_admission = match admission { - FindBestMatchAdmission::WithAdmission { .. } if update_states => { - SelectionAdmission::Lease { - request_id: context_id.expect("validated above").to_string(), - } - } - FindBestMatchAdmission::WithAdmission { .. } => SelectionAdmission::Query { + FindBestMatchAdmission::WithAdmission if update_states => SelectionAdmission::Lease { + request_id: context_id.expect("validated above").to_string(), + }, + FindBestMatchAdmission::WithAdmission => SelectionAdmission::Query { request_id: context_id.map(str::to_string), }, FindBestMatchAdmission::WithoutAdmission => SelectionAdmission::Advisory { request_id: context_id.map(str::to_string), }, }; - let outcome = self + let SelectionRun { + result: outcome, + lookup, + } = self .scheduler .run_selection(SelectionOperation { key: self.scheduler.partition_key().clone(), @@ -1563,8 +1343,8 @@ impl KvRouter { block_hashes: None, sequence_hashes: None, isl_tokens: None, - lora_name, - cache_namespace, + lora_name: lora_name.as_deref(), + cache_namespace: cache_namespace.as_deref(), is_eagle: Some(self.is_eagle), }, router_config_override: router_config_override.cloned(), @@ -1584,14 +1364,19 @@ impl KvRouter { replay_id: None, }) .await; + if lookup.is_some_and(|lookup| lookup.shared_cache_error) + && let Some(m) = metrics::RoutingOverheadMetrics::get() + { + m.inc_shared_cache_errors(); + } let selected = match outcome { Ok(SelectionOutcome::Selected(selected)) => selected, Ok(SelectionOutcome::QueueRejected { rejection }) => { return Ok(match admission { - FindBestMatchAdmission::WithAdmission { .. } => { + FindBestMatchAdmission::WithAdmission => { FindBestMatchInnerOutcome::WithAdmission(AdmittedFindBestMatchOutcome { outcome: FindBestMatchOutcome::QueueRejected { rejection }, - attempt: AdmissionAttempt::Untracked, + lease: None, }) } FindBestMatchAdmission::WithoutAdmission => { @@ -1612,25 +1397,49 @@ impl KvRouter { let Selected { response, advisory_load, + isl_tokens, kv_hint, routing_hashes, + shared_cache_hits, lease, .. } = selected; + if update_states && is_admitted_routing && lease.is_none() { + anyhow::bail!("booked selection returned no lifecycle lease"); + } let routing_hashes = routing_hashes.map(RoutingDecisionHashes::from_local_hashes); let overlap_blocks = response.effective_overlap_blocks.round() as u32; + + // Routing metrics stay scoped to requests admitted by this call. + if is_admitted_routing { + let total_elapsed = start.elapsed(); + if let (Some(lookup), Some(m)) = (lookup, metrics::RoutingOverheadMetrics::get()) { + let hash_elapsed = lookup.block_hashing; + let seq_hash_elapsed = hash_elapsed + lookup.seq_hashing; + let find_matches_elapsed = seq_hash_elapsed + lookup.lookups; + m.observe( + hash_elapsed, + seq_hash_elapsed, + lookup.indexer, + lookup.shared_cache, + find_matches_elapsed, + total_elapsed, + ); + } + if let (Some(hits), Some(m)) = (shared_cache_hits, metrics::RouterRequestMetrics::get()) + { + let num_blocks = isl_tokens / self.block_size as usize; + if num_blocks > 0 { + m.shared_cache_hit_rate + .observe(hits.total_hits as f64 / num_blocks as f64); + } + m.shared_cache_beyond_blocks + .observe(hits.hits_beyond(overlap_blocks) as f64); + } + } + Ok(match admission { - FindBestMatchAdmission::WithAdmission { .. } => { - // The caller owns cleanup through the attempt identity. - let attempt = match lease { - Some(lease) => { - let booking = lease - .commit() - .ok_or_else(|| anyhow::anyhow!("booking lease holds no booking"))?; - AdmissionAttempt::Tracked(booking.attempt_id) - } - None => AdmissionAttempt::Untracked, - }; + FindBestMatchAdmission::WithAdmission => { FindBestMatchInnerOutcome::WithAdmission(AdmittedFindBestMatchOutcome { outcome: FindBestMatchOutcome::Routed { worker: response.best_worker, @@ -1641,7 +1450,7 @@ impl KvRouter { routing_hashes, kv_hint, }, - attempt, + lease, }) } FindBestMatchAdmission::WithoutAdmission => { @@ -1659,289 +1468,6 @@ impl KvRouter { }) } - #[allow(clippy::too_many_arguments)] - async fn find_best_match_details_with_policy_class_inner( - &self, - context_id: Option<&str>, - tokens: &[u32], - block_mm_infos: Option<&[Option]>, - router_config_override: Option<&RouterConfigOverride>, - update_states: bool, - return_routing_hashes: bool, - lora_name: Option, - cache_namespace: Option, - priority_jump: f64, - strict_priority: u32, - policy_class: Option, - session_context: Option, - expected_output_tokens: Option, - affinity_target: Option, - pinned_worker: Option, - allowed_worker_ids: Option>, - routing_constraints: RoutingConstraints, - admission: FindBestMatchAdmission, - ) -> anyhow::Result { - let start = Instant::now(); - - if update_states && context_id.is_none() { - anyhow::bail!("context_id must be provided if update_states is true"); - } - let mode = match admission { - FindBestMatchAdmission::WithAdmission { track_lifecycle } - if update_states && track_lifecycle => - { - ScheduleMode::TrackedWithLifecycle { - request_id: context_id.expect("validated above").to_string(), - } - } - FindBestMatchAdmission::WithAdmission { .. } if update_states => { - ScheduleMode::Tracked { - request_id: context_id.expect("validated above").to_string(), - } - } - FindBestMatchAdmission::WithAdmission { .. } - | FindBestMatchAdmission::WithoutAdmission => ScheduleMode::QueryOnly { - request_id: context_id.map(str::to_string), - }, - }; - let isl_tokens = tokens.len(); - let hash_options = BlockHashOptions { - block_mm_infos, - lora_name: lora_name.as_deref(), - cache_namespace: cache_namespace.as_deref(), - is_eagle: Some(self.is_eagle), - }; - - let block_hashes = tracing::info_span!("kv_router.compute_block_hashes") - .in_scope(|| compute_block_hash_for_seq(tokens, self.block_size, hash_options)); - log_routing_input_hashes(context_id, self.block_size, tokens, &block_hashes); - let hash_elapsed = start.elapsed(); - // Compute seq_hashes only if scheduler needs it for active blocks tracking - let maybe_seq_hashes = tracing::info_span!("kv_router.compute_seq_hashes").in_scope(|| { - self.kv_router_config - .compute_seq_hashes_for_tracking_with_context( - &self.tracking_hash, - self.tracking_hash_scope(), - tokens, - router_config_override, - hash_options, - Some(&block_hashes), - ) - }); - let seq_hash_elapsed = start.elapsed(); - - let is_admitted_routing = matches!(admission, FindBestMatchAdmission::WithAdmission { .. }); - let supports_overlap_refresh = self.scheduler.supports_overlap_refresh(); - let retain_block_hashes = supports_overlap_refresh || return_routing_hashes; - let has_transfer_capable_workers = self.has_transfer_capable_workers(); - let should_prepare_transfer_hint = is_admitted_routing && has_transfer_capable_workers; - let retain_kv_transfer_chain = - should_prepare_transfer_hint && self.indexer.supports_kv_transfer_chain_retention(); - if should_prepare_transfer_hint && !retain_kv_transfer_chain { - static WARN_ONCE: std::sync::Once = std::sync::Once::new(); - WARN_ONCE.call_once(|| { - tracing::warn!( - "TRANSFER hint chain retention requires a local event-driven indexer with no approximate side indexer and no remote-recorded routing decisions; proceeding without transfer hints" - ); - }); - } - - let TieredLookupResult { - tiered_matches, - shared_cache_hits, - indexer_duration, - shared_cache_duration, - retained_block_hashes, - } = query_tiered_matches( - &self.indexer, - self.shared_cache.as_deref(), - tokens, - self.block_size, - block_hashes, - TieredLookupOptions { - cache_namespace: cache_namespace.as_deref(), - retain_block_hashes, - retain_kv_transfer_chain, - }, - ) - .await?; - - let (block_hashes_for_refresh, routing_block_hashes) = retained_block_hashes - .map(|block_hashes| { - split_retained_block_hashes( - block_hashes, - supports_overlap_refresh, - return_routing_hashes, - ) - }) - .unwrap_or((None, None)); - - let overlap = - OverlapAnalysis::new(&self.kv_router_config, self.block_size, &tiered_matches) - .signals(); - let kv_transfer_candidates = retain_kv_transfer_chain - .then(|| tiered_matches.kv_transfer_candidates().cloned()) - .flatten(); - drop(tiered_matches); - let find_matches_elapsed = start.elapsed(); - - // Capture shared cache info for metrics before moving into schedule(). - // Clone the hits so we can compute `hits_beyond(overlap_blocks)` after - // scheduling returns, since `overlap_blocks` isn't known until then. - let num_blocks = isl_tokens / self.block_size as usize; - let sc_hits_for_metrics = shared_cache_hits.clone(); - - // LoRA-aware candidate narrowing: restrict to this LoRA's allocated/loaded replicas, - // strictly within the existing candidate universe (never widening). Covers both the - // decode and prefill routers, since both flow through this method. - let allowed_worker_ids = self.narrow_allowed_by_lora( - lora_name.as_deref(), - allowed_worker_ids, - pinned_worker.as_ref(), - ); - - let schedule_request = ScheduleRequest { - mode, - token_seq: maybe_seq_hashes, - block_hashes: block_hashes_for_refresh, - isl_tokens, - overlap, - kv_transfer_candidates, - retain_kv_transfer_chain, - router_config_override: router_config_override.cloned(), - lora_name, - priority_jump, - strict_priority, - policy_class, - session_context, - expected_output_tokens, - affinity_target, - pinned_worker, - allowed_worker_ids, - routing_constraints, - shared_cache_hits, - }; - let (response, attempt, selected_worker_load) = match admission { - FindBestMatchAdmission::WithAdmission { .. } => { - match self - .scheduler - .schedule_request_admitted(schedule_request) - .instrument(tracing::info_span!("kv_router.schedule")) - .await - { - Ok(admitted) => (admitted.response, admitted.attempt, None), - Err(KvSchedulerError::QueueRejected(rejection)) => { - return Ok(FindBestMatchInnerOutcome::WithAdmission( - AdmittedFindBestMatchOutcome { - outcome: FindBestMatchOutcome::QueueRejected { rejection }, - attempt: AdmissionAttempt::Untracked, - }, - )); - } - Err(error) => return Err(map_scheduler_error(error)), - } - } - FindBestMatchAdmission::WithoutAdmission => { - match self - .scheduler - .select_without_admission(schedule_request) - .instrument(tracing::info_span!("kv_router.select_without_admission")) - .await - { - Ok(advisory) => ( - advisory.response, - AdmissionAttempt::Untracked, - Some(advisory.selected_worker_load), - ), - Err(KvSchedulerError::QueueRejected(rejection)) => { - return Ok(FindBestMatchInnerOutcome::WithoutAdmission( - FindBestMatchAdvisoryOutcome::QueueRejected { rejection }, - )); - } - Err(error) => return Err(map_scheduler_error(error)), - } - } - }; - let kv_hint = if is_admitted_routing { - self.kv_hint_for_selection( - context_id, - response.best_worker, - response.target_cached_prefix_blocks, - response.kv_transfer_candidates.as_ref(), - ) - } else { - None - }; - - let total_elapsed = start.elapsed(); - let routing_hashes = routing_block_hashes.map(RoutingDecisionHashes::from_local_hashes); - - // Keep existing routing metrics scoped to requests admitted into the scheduler by this call. - if is_admitted_routing && let Some(m) = metrics::RoutingOverheadMetrics::get() { - m.observe( - hash_elapsed, - seq_hash_elapsed, - indexer_duration, - shared_cache_duration, - find_matches_elapsed, - total_elapsed, - ); - } - - // Observe per-request shared cache metrics. - if is_admitted_routing - && let Some(hits) = sc_hits_for_metrics - && let Some(m) = metrics::RouterRequestMetrics::get() - { - if num_blocks > 0 { - m.shared_cache_hit_rate - .observe(hits.total_hits as f64 / num_blocks as f64); - } - let beyond = hits.hits_beyond(response.effective_overlap_blocks.round() as u32); - m.shared_cache_beyond_blocks.observe(beyond as f64); - } - - #[cfg(feature = "bench")] - tracing::info!( - isl_tokens, - hash_us = hash_elapsed.as_micros() as u64, - seq_hash_us = (seq_hash_elapsed - hash_elapsed).as_micros() as u64, - find_matches_us = (find_matches_elapsed - seq_hash_elapsed).as_micros() as u64, - schedule_us = (total_elapsed - find_matches_elapsed).as_micros() as u64, - total_us = total_elapsed.as_micros() as u64, - "find_best_match completed" - ); - - match admission { - FindBestMatchAdmission::WithAdmission { .. } => Ok( - FindBestMatchInnerOutcome::WithAdmission(AdmittedFindBestMatchOutcome { - outcome: FindBestMatchOutcome::Routed { - worker: response.best_worker, - overlap_blocks: response.effective_overlap_blocks.round() as u32, - effective_overlap_blocks: response.effective_overlap_blocks, - cached_tokens: response.cached_tokens, - potential_decode_blocks: response.potential_decode_blocks as u64, - routing_hashes, - kv_hint, - }, - attempt, - }), - ), - FindBestMatchAdmission::WithoutAdmission => Ok( - FindBestMatchInnerOutcome::WithoutAdmission(FindBestMatchAdvisoryOutcome::Routed { - worker: response.best_worker, - overlap_blocks: response.effective_overlap_blocks.round() as u32, - effective_overlap_blocks: response.effective_overlap_blocks, - cached_tokens: response.cached_tokens, - potential_decode_blocks: response.potential_decode_blocks as u64, - selected_worker_load: selected_worker_load - .expect("without-admission selection returns advisory load"), - routing_hashes, - }), - ), - } - } - /// Give these tokens, find the worker with the best match in its KV cache. /// Returns the best worker (with dp_rank) and approximate effective overlap in blocks. #[allow(clippy::too_many_arguments)] @@ -2516,16 +2042,8 @@ mod tests { use async_trait::async_trait; use dynamo_kv_router::{ WorkerSelectionPolicyError, - identity::{ - CacheOwnerId, CacheSemanticsId, DcId, IdentitySource, IndexerDomainId, PoolId, - RoutingScopeId, StableDpSlotId, - }, indexer::{LowerTierMatchDetails, MatchDetails}, - protocols::{ - ExternalSequenceBlockHash, OverlapScores, ResidencyOwner, ResidencyProjection, - ResidencyRoutingSnapshot, RouterHintSourceMetadata, StorageTier, - compute_seq_hash_for_block, - }, + protocols::{OverlapScores, StorageTier, compute_seq_hash_for_block}, }; use dynamo_runtime::{DistributedRuntime, Runtime, distributed::DistributedConfig}; use tokio::sync::watch; @@ -3071,22 +2589,17 @@ mod tests { "/tests/data/selection_golden/aggregated_tracked.json" ); - /// Both selection procedures, run on independent routers over the same - /// corpus with tracked bookings, must agree with each other and with the - /// frozen trace recorded from the frontend procedure. Regenerate the trace - /// with `SELECTION_GOLDEN_UPDATE=1` only for an intended behavior change. + /// The selection procedure, run over the corpus with tracked bookings, must + /// match the trace frozen from the frontend's pre-unification procedure. + /// Regenerate the trace with `SELECTION_GOLDEN_UPDATE=1` only for an + /// intended behavior change. #[tokio::test] - async fn core_selection_matches_frontend_golden_trace() { - let frontend = differential_router("golden-frontend").await; - let core = differential_router("golden-core").await; - let mut frontend_trace = Vec::new(); - let mut core_trace = Vec::new(); + async fn selection_matches_frozen_frontend_trace() { + let router = differential_router("golden").await; + let mut trace = Vec::new(); for (index, tokens) in golden_corpus().iter().enumerate() { let context_id = format!("golden-{index}"); - let admission = FindBestMatchAdmission::WithAdmission { - track_lifecycle: true, - }; - let old = frontend + let FindBestMatchInnerOutcome::WithAdmission(admitted) = router .find_best_match_details_with_policy_class_inner( Some(&context_id), tokens, @@ -3105,179 +2618,89 @@ mod tests { None, None, RoutingConstraints::default(), - admission, + FindBestMatchAdmission::WithAdmission, ) .await - .unwrap(); - let new = core - .select_core( - Some(&context_id), - tokens, - None, - None, - true, - true, - None, - None, - 0.0, - 0, - None, - None, - None, - None, - None, - None, - RoutingConstraints::default(), - admission, - ) - .await - .unwrap(); - let ( - FindBestMatchInnerOutcome::WithAdmission(old), - FindBestMatchInnerOutcome::WithAdmission(new), - ) = (old, new) + .unwrap() else { panic!("admitted routing returned advisory outcome"); }; - assert!(matches!(old.attempt, AdmissionAttempt::Tracked(_))); - assert!(matches!(new.attempt, AdmissionAttempt::Tracked(_))); - match (&old.outcome, &new.outcome) { - ( - FindBestMatchOutcome::Routed { - routing_hashes: old_hashes, - .. - }, - FindBestMatchOutcome::Routed { - routing_hashes: new_hashes, - .. - }, - ) => assert_eq!( - old_hashes.as_ref().map(|h| h.local_hashes.clone()), - new_hashes.as_ref().map(|h| h.local_hashes.clone()) - ), - _ => panic!("golden corpus must route"), - } - frontend_trace.push(GoldenDecision::from_outcome(&old.outcome)); - core_trace.push(GoldenDecision::from_outcome(&new.outcome)); + let FindBestMatchOutcome::Routed { routing_hashes, .. } = &admitted.outcome else { + panic!("golden corpus must route"); + }; + assert!(routing_hashes.is_some(), "routing hashes were requested"); + trace.push(GoldenDecision::from_outcome(&admitted.outcome)); + // Keep the booking: later rows are scored against its load. + let booking = admitted + .lease + .and_then(|lease| lease.commit()) + .expect("tracked selection carries its booking lease"); + assert_eq!(booking.request_id, context_id); } - assert_eq!( - core_trace, frontend_trace, - "core and frontend procedures diverged" - ); if std::env::var_os("SELECTION_GOLDEN_UPDATE").is_some() { std::fs::write( GOLDEN_PATH, - serde_json::to_string_pretty(&frontend_trace).unwrap() + "\n", + serde_json::to_string_pretty(&trace).unwrap() + "\n", ) .unwrap(); } let golden: Vec = serde_json::from_str(&std::fs::read_to_string(GOLDEN_PATH).unwrap()).unwrap(); assert_eq!( - frontend_trace, golden, - "frontend procedure drifted from its frozen trace" + trace, golden, + "selection drifted from the frozen frontend trace" ); } - /// Query-only and advisory selections mutate nothing, so one router serves - /// both procedures. + /// A dropped tracked selection frees its booking; a committed one keeps it. #[tokio::test] - async fn core_selection_matches_frontend_for_unbooked_requests() { - let router = differential_router("golden-unbooked").await; - for tokens in golden_corpus() { - for admission in [ - FindBestMatchAdmission::WithAdmission { - track_lifecycle: false, - }, - FindBestMatchAdmission::WithoutAdmission, - ] { - let old = router - .find_best_match_details_with_policy_class_inner( - None, - &tokens, - None, - None, - false, - false, - None, - None, - 0.0, - 0, - None, - None, - None, - None, - None, - None, - RoutingConstraints::default(), - admission, - ) - .await - .unwrap(); - let new = router - .select_core( - None, - &tokens, - None, - None, - false, - false, - None, - None, - 0.0, - 0, - None, - None, - None, - None, - None, - None, - RoutingConstraints::default(), - admission, - ) - .await - .unwrap(); - match (old, new) { - ( - FindBestMatchInnerOutcome::WithAdmission(old), - FindBestMatchInnerOutcome::WithAdmission(new), - ) => { - assert!(matches!(old.attempt, AdmissionAttempt::Untracked)); - assert!(matches!(new.attempt, AdmissionAttempt::Untracked)); - assert_eq!( - GoldenDecision::from_outcome(&new.outcome), - GoldenDecision::from_outcome(&old.outcome) - ); - } - ( - FindBestMatchInnerOutcome::WithoutAdmission( - FindBestMatchAdvisoryOutcome::Routed { - worker: old_worker, - overlap_blocks: old_overlap, - selected_worker_load: old_load, - .. - }, - ), - FindBestMatchInnerOutcome::WithoutAdmission( - FindBestMatchAdvisoryOutcome::Routed { - worker: new_worker, - overlap_blocks: new_overlap, - selected_worker_load: new_load, - .. - }, - ), - ) => { - assert_eq!((new_worker, new_overlap), (old_worker, old_overlap)); - assert_eq!( - new_load.active_prefill_tokens, - old_load.active_prefill_tokens - ); - } - _ => panic!("procedures returned different outcome shapes"), - } + async fn dropped_selection_lease_frees_the_booking() { + let router = differential_router("lease").await; + let prompt = golden_prefix(1); + let select = |context_id: &'static str| { + router.find_best_match_details_with_policy_class_inner( + Some(context_id), + &prompt, + None, + None, + true, + false, + None, + None, + 0.0, + 0, + None, + None, + None, + None, + None, + None, + RoutingConstraints::default(), + FindBestMatchAdmission::WithAdmission, + ) + }; + let FindBestMatchInnerOutcome::WithAdmission(dropped) = select("dropped").await.unwrap() + else { + panic!("admitted routing returned advisory outcome"); + }; + drop(dropped); + let FindBestMatchInnerOutcome::WithAdmission(kept) = select("kept").await.unwrap() else { + panic!("admitted routing returned advisory outcome"); + }; + let _booking = kept + .lease + .and_then(|lease| lease.commit()) + .expect("booking"); + + tokio::time::timeout(std::time::Duration::from_secs(2), async { + while router.scheduler.has_request("dropped") { + tokio::task::yield_now().await; } - } + }) + .await + .expect("dropped selection released its booking"); + assert!(router.scheduler.has_request("kept")); } async fn make_test_router_with_workers( @@ -3330,302 +2753,6 @@ mod tests { make_test_router_with_workers(policy, shared_cache, workers).await } - fn transfer_hint_runtime_config(endpoint: Option<&str>) -> ModelRuntimeConfig { - transfer_hint_runtime_config_with_worker_type(endpoint, "prefill") - } - - fn transfer_hint_runtime_config_with_worker_type( - endpoint: Option<&str>, - worker_type: &str, - ) -> ModelRuntimeConfig { - let mut runtime_config = ModelRuntimeConfig::default(); - runtime_config.runtime_data.insert( - dynamo_kv_router::kv_hints::KV_HINT_TRANSFER_CAPABILITY_KEY.to_string(), - serde_json::Value::Bool(true), - ); - runtime_config.runtime_data.insert( - dynamo_kv_router::kv_hints::KV_HINT_TRANSFER_WORKER_TYPE_RUNTIME_KEY.to_string(), - serde_json::Value::String(worker_type.to_string()), - ); - if let Some(endpoint) = endpoint { - let mut endpoints = serde_json::Map::new(); - endpoints.insert( - "0".to_string(), - serde_json::Value::String(endpoint.to_string()), - ); - runtime_config.runtime_data.insert( - dynamo_kv_router::kv_hints::KV_HINT_TRANSFER_SOURCE_CONTROL_ENDPOINTS_RUNTIME_KEY - .to_string(), - serde_json::Value::Object(endpoints), - ); - } - runtime_config - } - - fn transfer_hint_runtime_config_with_dp_endpoints( - endpoints: &[(u32, &str)], - ) -> ModelRuntimeConfig { - let mut runtime_config = transfer_hint_runtime_config(None); - runtime_config.runtime_data.insert( - dynamo_kv_router::kv_hints::KV_HINT_TRANSFER_SOURCE_CONTROL_ENDPOINTS_RUNTIME_KEY - .to_string(), - serde_json::Value::Object( - endpoints - .iter() - .map(|(dp_rank, endpoint)| { - ( - dp_rank.to_string(), - serde_json::Value::String(endpoint.to_string()), - ) - }) - .collect(), - ), - ); - runtime_config - } - - fn router_hint_cache_owner() -> CacheOwnerId { - CacheOwnerId::new( - PoolId::new( - IndexerDomainId::new( - CacheSemanticsId::new([1; 16], IdentitySource::Explicit), - RoutingScopeId::new([2; 16], IdentitySource::Explicit), - ), - DcId::new(3), - ), - StableDpSlotId::new([4; 16], IdentitySource::Explicit), - ) - } - - #[tokio::test] - async fn kv_hint_composer_wraps_transfer_from_another_dp_rank() { - let mut workers = HashMap::new(); - workers.insert( - 7, - transfer_hint_runtime_config_with_dp_endpoints(&[ - (0, "tcp://127.0.0.1:23280"), - (1, "tcp://127.0.0.1:23281"), - ]), - ); - let router = make_test_router_with_workers( - fixed_policy(None, WorkerWithDpRank::new(7, 0)), - None, - workers, - ) - .await; - let candidates = KvTransferCandidates { - block_hashes: vec![ - ExternalSequenceBlockHash(101), - ExternalSequenceBlockHash(102), - ], - owner_prefix_blocks: vec![(WorkerWithDpRank::new(7, 1).into(), 2)], - routing_snapshot: None, - }; - - let hint = router.kv_hint_for_selection( - Some("msg-123"), - WorkerWithDpRank::new(7, 0), - 0, - Some(&candidates), - ); - - assert_eq!( - hint, - Some(KvHint::new( - "msg-123", - vec![KvHintAction::fetch( - "a1", - KvSourceLocationsPayload { - source_control_endpoint: "tcp://127.0.0.1:23281".to_string(), - block_hashes: vec![ - ExternalSequenceBlockHash(101), - ExternalSequenceBlockHash(102), - ], - }, - )], - )) - ); - } - - #[tokio::test] - async fn transfer_hint_skips_sources_without_usable_endpoint() { - for source_endpoint in [None, Some("")] { - let mut workers = HashMap::new(); - workers.insert( - 7, - transfer_hint_runtime_config(Some("tcp://127.0.0.1:23280")), - ); - workers.insert(8, transfer_hint_runtime_config(source_endpoint)); - let router = make_test_router_with_workers( - fixed_policy(None, WorkerWithDpRank::new(7, 0)), - None, - workers, - ) - .await; - let candidates = KvTransferCandidates { - block_hashes: vec![ - ExternalSequenceBlockHash(101), - ExternalSequenceBlockHash(102), - ], - owner_prefix_blocks: vec![(WorkerWithDpRank::new(8, 0).into(), 2)], - routing_snapshot: None, - }; - - let hint = router.transfer_hint_for_selection( - WorkerWithDpRank::new(7, 0), - 0, - Some(&candidates), - ); - - assert_eq!(hint, None); - } - } - - #[tokio::test] - async fn transfer_hint_skips_sources_with_different_worker_type() { - let mut workers = HashMap::new(); - workers.insert( - 7, - transfer_hint_runtime_config_with_worker_type(Some("tcp://127.0.0.1:23280"), "prefill"), - ); - workers.insert( - 8, - transfer_hint_runtime_config_with_worker_type(Some("tcp://127.0.0.1:23281"), "decode"), - ); - let router = make_test_router_with_workers( - fixed_policy(None, WorkerWithDpRank::new(7, 0)), - None, - workers, - ) - .await; - let candidates = KvTransferCandidates { - block_hashes: vec![ - ExternalSequenceBlockHash(101), - ExternalSequenceBlockHash(102), - ], - owner_prefix_blocks: vec![(WorkerWithDpRank::new(8, 0).into(), 2)], - routing_snapshot: None, - }; - - let hint = - router.transfer_hint_for_selection(WorkerWithDpRank::new(7, 0), 0, Some(&candidates)); - - assert_eq!(hint, None); - } - - #[tokio::test] - async fn transfer_hint_selects_source_with_matching_worker_type() { - let mut workers = HashMap::new(); - workers.insert( - 7, - transfer_hint_runtime_config_with_worker_type(Some("tcp://127.0.0.1:23280"), "prefill"), - ); - workers.insert( - 8, - transfer_hint_runtime_config_with_worker_type(Some("tcp://127.0.0.1:23281"), "prefill"), - ); - workers.insert( - 9, - transfer_hint_runtime_config_with_worker_type(Some("tcp://127.0.0.1:23282"), "decode"), - ); - workers.insert( - 10, - transfer_hint_runtime_config_with_worker_type(Some("tcp://127.0.0.1:23283"), "decode"), - ); - let router = make_test_router_with_workers( - fixed_policy(None, WorkerWithDpRank::new(7, 0)), - None, - workers, - ) - .await; - let candidates = KvTransferCandidates { - block_hashes: vec![ - ExternalSequenceBlockHash(101), - ExternalSequenceBlockHash(102), - ExternalSequenceBlockHash(103), - ], - owner_prefix_blocks: vec![ - (WorkerWithDpRank::new(8, 0).into(), 2), - (WorkerWithDpRank::new(9, 0).into(), 3), - ], - routing_snapshot: None, - }; - - let prefill_hint = - router.transfer_hint_for_selection(WorkerWithDpRank::new(7, 0), 0, Some(&candidates)); - assert_eq!( - prefill_hint, - Some(KvSourceLocationsPayload { - source_control_endpoint: "tcp://127.0.0.1:23281".to_string(), - block_hashes: vec![ - ExternalSequenceBlockHash(101), - ExternalSequenceBlockHash(102), - ], - }) - ); - - let decode_hint = - router.transfer_hint_for_selection(WorkerWithDpRank::new(10, 0), 0, Some(&candidates)); - assert_eq!( - decode_hint, - Some(KvSourceLocationsPayload { - source_control_endpoint: "tcp://127.0.0.1:23282".to_string(), - block_hashes: vec![ - ExternalSequenceBlockHash(101), - ExternalSequenceBlockHash(102), - ExternalSequenceBlockHash(103), - ], - }) - ); - } - - #[tokio::test] - async fn kv_transfer_resolves_persistent_owner_without_state_agent_fallback() { - let target = WorkerWithDpRank::new(7, 0); - let stale_source = WorkerWithDpRank::new(8, 0); - let mut workers = HashMap::new(); - workers.insert(7, transfer_hint_runtime_config(None)); - let mut stale_source_config = - transfer_hint_runtime_config(Some("tcp://stale-worker-endpoint:23280")); - stale_source_config.kv_event_source_mode = Some("state_agent_v2".to_string()); - workers.insert(8, stale_source_config); - let router = make_test_router_with_workers(fixed_policy(None, target), None, workers).await; - let owner = router_hint_cache_owner(); - let owner_key = ResidencyOwner::cache_owner(owner).compact_key(); - let candidates = KvTransferCandidates { - block_hashes: vec![ - ExternalSequenceBlockHash(101), - ExternalSequenceBlockHash(102), - ], - owner_prefix_blocks: vec![ - (KvTransferCandidateSource::Worker(stale_source), 2), - (KvTransferCandidateSource::CacheOwner(owner_key), 2), - ], - routing_snapshot: Some(Arc::new(ResidencyRoutingSnapshot::new( - ResidencyProjection::default(), - [( - owner, - RouterHintSourceMetadata { - source_control_endpoint: "tcp://persistent-owner:23280".to_string(), - worker_type: "prefill".to_string(), - }, - None, - )], - ))), - }; - - assert_eq!( - router.transfer_hint_for_selection(target, 0, Some(&candidates)), - Some(KvSourceLocationsPayload { - source_control_endpoint: "tcp://persistent-owner:23280".to_string(), - block_hashes: vec![ - ExternalSequenceBlockHash(101), - ExternalSequenceBlockHash(102), - ], - }) - ); - } - #[tokio::test] async fn test_find_best_match_passes_shared_cache_hits_to_scheduler() { let router = make_test_router( diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs index d840df160617..bc1315fdba75 100644 --- a/lib/llm/src/kv_router/embedded.rs +++ b/lib/llm/src/kv_router/embedded.rs @@ -21,17 +21,15 @@ use dynamo_kv_router::scheduling::queue::{ SchedulerBookingDescriptor, }; use dynamo_kv_router::scheduling::{ - AdmittedSchedulingResponse, AdvisorySchedulingResponse, AttemptId, KvSchedulerError, - NonMaxOverlapSelectionObserver, OverloadedWorkerProvider, PotentialLoad, QueueLimitKind, - QueueRejection, ScheduleRequest, WorkerAvailabilityProvider, + AttemptId, KvSchedulerError, NonMaxOverlapSelectionObserver, OverloadedWorkerProvider, + PotentialLoad, QueueLimitKind, QueueRejection, WorkerAvailabilityProvider, }; use dynamo_kv_router::sequences::{SequenceError, SequenceRequest}; use dynamo_kv_router::services::selection::{ CatalogObserver, CatalogReconciler, HostCache, HostEligibility, HostLoad, HostReplication, - HostTelemetry, KvEventIngress, KvIndexSource, SelectionError, SelectionHost, - SelectionOperation, SelectionOutcome, SelectionPartition, SelectionService, - SelectionServiceBuilder, WorkerCatalogRecord, WorkerCatalogSource, WorkerRequest, - WorkerSelectionPolicyRegistry, + HostTelemetry, KvEventIngress, KvIndexSource, SelectionHost, SelectionOperation, + SelectionOutcome, SelectionPartition, SelectionRun, SelectionService, SelectionServiceBuilder, + WorkerCatalogRecord, WorkerCatalogSource, WorkerRequest, WorkerSelectionPolicyRegistry, }; use dynamo_kv_router::{DEFAULT_ROUTING_GROUP, PrefillLoadEstimator, WorkerSelectionPolicyFactory}; use dynamo_tokens::SequenceHash; @@ -361,15 +359,6 @@ impl EmbeddedSelection { }) } - /// Queue changes are not observable through the scheduler's update watch, - /// so every schedule result refreshes the gauges. - fn observe_schedule_result(&self, result: &Result) { - self.observe_queue(match result { - Err(KvSchedulerError::QueueRejected(rejection)) => Some(rejection), - _ => None, - }); - } - fn observe_queue(&self, rejection: Option<&QueueRejection>) { if let Some(rejection) = rejection { record_queue_rejection(&self.queue_metrics, &self.queue_metric_indices, rejection); @@ -401,42 +390,13 @@ impl EmbeddedSelection { } /// Run one selection through the shared core. - pub(crate) async fn run_selection( - &self, - operation: SelectionOperation<'_>, - ) -> Result { - let outcome = self.service.core().run_selection(operation).await; - self.observe_queue(match &outcome { + pub(crate) async fn run_selection(&self, operation: SelectionOperation<'_>) -> SelectionRun { + let run = self.service.core().run_selection(operation).await; + self.observe_queue(match &run.result { Ok(SelectionOutcome::QueueRejected { rejection }) => Some(rejection), _ => None, }); - outcome - } - - pub(crate) async fn schedule_request_admitted( - &self, - request: ScheduleRequest, - ) -> Result { - let result = self - .partition - .scheduler() - .schedule_request_admitted(request) - .await; - self.observe_schedule_result(&result); - result - } - - pub(crate) async fn select_without_admission( - &self, - request: ScheduleRequest, - ) -> Result { - let result = self - .partition - .scheduler() - .select_without_admission(request) - .await; - self.observe_schedule_result(&result); - result + run } pub(crate) async fn add_request_admitted( @@ -486,6 +446,11 @@ impl EmbeddedSelection { .map(|_| ()) } + #[cfg(test)] + pub(crate) fn has_request(&self, request_id: &str) -> bool { + self.partition.scheduler().has_request(request_id) + } + pub(crate) fn pending_count(&self) -> usize { self.partition.scheduler().pending_count() } @@ -533,10 +498,6 @@ impl EmbeddedSelection { track_prefill_tokens, ) } - - pub(crate) fn supports_overlap_refresh(&self) -> bool { - self.partition.scheduler().supports_overlap_refresh() - } } /// The runtime-config watch as worker membership for the partition. diff --git a/lib/llm/src/kv_router/prefill_router/query.rs b/lib/llm/src/kv_router/prefill_router/query.rs index 6673f232a9dd..8e467d65a0cb 100644 --- a/lib/llm/src/kv_router/prefill_router/query.rs +++ b/lib/llm/src/kv_router/prefill_router/query.rs @@ -6,10 +6,7 @@ use std::collections::HashSet; use anyhow::Result; use dynamo_kv_router::{ protocols::{BlockExtraInfo, RoutingConstraints, WorkerId, WorkerWithDpRank}, - scheduling::{ - AdmissionAttempt, - queue::{SchedulerBookingCleanup, SchedulerBookingDescriptor}, - }, + scheduling::queue::{SchedulerBookingCleanup, SchedulerBookingDescriptor}, }; use super::{PrefillError, PrefillLifecycleState, PrefillQueryOutcome, PrefillRouter}; @@ -128,22 +125,18 @@ impl PrefillRouter { routing_constraints, ) .await?; - let (outcome, attempt) = admitted.into_parts(); + let (outcome, lease) = admitted.into_parts(); match outcome { crate::kv_router::FindBestMatchOutcome::Routed { worker, .. } => { - let AdmissionAttempt::Tracked(attempt_id) = attempt else { - anyhow::bail!("prefill reservation admission did not return a tracked attempt"); + let Some(booking) = lease.and_then(|lease| lease.commit()) else { + anyhow::bail!("prefill reservation admission did not return a booking"); }; Ok(PrefillReservation { worker, dp_rank: Some(worker.dp_rank), release: ReservationRelease::Kv { cleanup: chooser.booking_cleanup(), - booking: Some(SchedulerBookingDescriptor { - request_id: reservation_id.to_string(), - worker, - attempt_id, - }), + booking: Some(booking), }, }) } diff --git a/lib/llm/src/kv_router/route_lookup.rs b/lib/llm/src/kv_router/route_lookup.rs deleted file mode 100644 index 021915552a1d..000000000000 --- a/lib/llm/src/kv_router/route_lookup.rs +++ /dev/null @@ -1,238 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -use std::{ - future::Future, - time::{Duration, Instant}, -}; - -use dynamo_kv_router::{ - SharedKvCache, - indexer::{KvRouterError, LowerTierQueryOptions}, - protocols::{LocalBlockHash, SharedCacheHits}, -}; -use tracing::Instrument; - -use super::{Indexer, indexer::TieredMatchDetails, metrics}; - -pub(super) struct TieredLookupOptions<'a> { - pub(super) cache_namespace: Option<&'a str>, - pub(super) retain_block_hashes: bool, - pub(super) retain_kv_transfer_chain: bool, -} - -pub(super) struct TieredLookupResult { - pub(super) tiered_matches: TieredMatchDetails, - pub(super) shared_cache_hits: Option, - pub(super) indexer_duration: Duration, - pub(super) shared_cache_duration: Option, - pub(super) retained_block_hashes: Option>, -} - -pub(super) fn split_retained_block_hashes( - block_hashes: Vec, - supports_overlap_refresh: bool, - return_routing_hashes: bool, -) -> (Option>, Option>) { - debug_assert!(supports_overlap_refresh || return_routing_hashes); - - match (supports_overlap_refresh, return_routing_hashes) { - (true, true) => (Some(block_hashes.clone()), Some(block_hashes)), - (true, false) => (Some(block_hashes), None), - (false, true) => (None, Some(block_hashes)), - (false, false) => (None, None), - } -} - -pub(super) async fn query_tiered_matches( - indexer: &Indexer, - shared_cache: Option<&dyn SharedKvCache>, - tokens: &[u32], - block_size: u32, - block_hashes: Vec, - options: TieredLookupOptions<'_>, -) -> Result { - let retain_kv_transfer_chain = - options.retain_kv_transfer_chain && indexer.supports_kv_transfer_chain_retention(); - if options.retain_kv_transfer_chain && !retain_kv_transfer_chain { - tracing::warn!( - "KV transfer chain retention is not supported by this indexer configuration; KV transfer hints will not be attached" - ); - } - let lower_tier_options = LowerTierQueryOptions { - retain_kv_transfer_chain, - }; - if options.retain_block_hashes { - let (tiered_matches, shared_cache_hits, indexer_duration, shared_cache_duration) = - query_retained( - indexer, - shared_cache, - tokens, - block_size, - &block_hashes, - options.cache_namespace, - lower_tier_options, - ) - .await?; - - return Ok(TieredLookupResult { - tiered_matches, - shared_cache_hits, - indexer_duration, - shared_cache_duration, - retained_block_hashes: Some(block_hashes), - }); - } - - let (tiered_matches, shared_cache_hits, indexer_duration, shared_cache_duration) = query_owned( - indexer, - shared_cache, - tokens, - block_size, - block_hashes, - options.cache_namespace, - lower_tier_options, - ) - .await?; - - Ok(TieredLookupResult { - tiered_matches, - shared_cache_hits, - indexer_duration, - shared_cache_duration, - retained_block_hashes: None, - }) -} - -async fn query_retained( - indexer: &Indexer, - shared_cache: Option<&dyn SharedKvCache>, - tokens: &[u32], - block_size: u32, - block_hashes: &[LocalBlockHash], - cache_namespace: Option<&str>, - lower_tier_options: LowerTierQueryOptions, -) -> Result< - ( - TieredMatchDetails, - Option, - Duration, - Option, - ), - KvRouterError, -> { - let Some(shared_cache) = shared_cache else { - let t = Instant::now(); - let tiered = indexer - .find_matches_by_tier_ref_with_options(block_hashes, lower_tier_options) - .instrument(tracing::info_span!("kv_router.find_matches")) - .await?; - return Ok((tiered, None, t.elapsed(), None)); - }; - - let indexer_fut = indexer - .find_matches_by_tier_ref_with_options(block_hashes, lower_tier_options) - .instrument(tracing::info_span!("kv_router.find_matches")); - join_indexer_and_shared_cache( - indexer_fut, - shared_cache, - tokens, - block_size, - cache_namespace, - ) - .await -} - -async fn query_owned( - indexer: &Indexer, - shared_cache: Option<&dyn SharedKvCache>, - tokens: &[u32], - block_size: u32, - block_hashes: Vec, - cache_namespace: Option<&str>, - lower_tier_options: LowerTierQueryOptions, -) -> Result< - ( - TieredMatchDetails, - Option, - Duration, - Option, - ), - KvRouterError, -> { - let Some(shared_cache) = shared_cache else { - let t = Instant::now(); - let tiered = indexer - .find_matches_by_tier_with_options(block_hashes, lower_tier_options) - .instrument(tracing::info_span!("kv_router.find_matches")) - .await?; - return Ok((tiered, None, t.elapsed(), None)); - }; - - let indexer_fut = indexer - .find_matches_by_tier_with_options(block_hashes, lower_tier_options) - .instrument(tracing::info_span!("kv_router.find_matches")); - join_indexer_and_shared_cache( - indexer_fut, - shared_cache, - tokens, - block_size, - cache_namespace, - ) - .await -} - -async fn join_indexer_and_shared_cache( - indexer_fut: I, - shared_cache: &dyn SharedKvCache, - tokens: &[u32], - block_size: u32, - cache_namespace: Option<&str>, -) -> Result< - ( - TieredMatchDetails, - Option, - Duration, - Option, - ), - KvRouterError, -> -where - I: Future>, -{ - let shared_fut = shared_cache - .check_blocks(tokens, block_size, cache_namespace) - .instrument(tracing::info_span!("kv_router.shared_cache_check")); - - let indexer_timed = async { - let t = Instant::now(); - let r = indexer_fut.await; - (r, t.elapsed()) - }; - let shared_timed = async { - let t = Instant::now(); - let r = shared_fut.await; - (r, t.elapsed()) - }; - - let ((indexer_result, indexer_duration), (shared_result, shared_cache_duration)) = - tokio::join!(indexer_timed, shared_timed); - let tiered_matches = indexer_result?; - let shared_cache_hits = match shared_result { - Ok(hits) => Some(hits), - Err(e) => { - tracing::warn!(error = %e, "Shared cache query failed, ignoring"); - if let Some(m) = metrics::RoutingOverheadMetrics::get() { - m.inc_shared_cache_errors(); - } - None - } - }; - - Ok(( - tiered_matches, - shared_cache_hits, - indexer_duration, - Some(shared_cache_duration), - )) -} diff --git a/lib/llm/src/kv_router/routing_host/kv.rs b/lib/llm/src/kv_router/routing_host/kv.rs index 337a0218e41b..90659dd2e56b 100644 --- a/lib/llm/src/kv_router/routing_host/kv.rs +++ b/lib/llm/src/kv_router/routing_host/kv.rs @@ -74,9 +74,7 @@ impl RoutingHost { is_query_only, affinity_target, None, - FindBestMatchAdmission::WithAdmission { - track_lifecycle: true, - }, + FindBestMatchAdmission::WithAdmission, budget, ) .await? @@ -179,7 +177,7 @@ impl RoutingHost { let phase_label = phase.to_string(); let route_guard = StageGuard::new(STAGE_ROUTE, &phase_label); let planned_worker = preview.signals.worker; - let (selection, affinity) = self + let (mut selection, affinity) = self .select_with_session_affinity(request, phase, false, &budget, |target| { let budget = &budget; async move { @@ -189,9 +187,7 @@ impl RoutingHost { false, target, Some(planned_worker), - FindBestMatchAdmission::WithAdmission { - track_lifecycle: true, - }, + FindBestMatchAdmission::WithAdmission, budget, ) .await? @@ -207,7 +203,7 @@ impl RoutingHost { Arc::clone(self.kv_router()), request.context().id().to_string(), selection.worker, - selection.attempt, + selection.lease.take(), ), selection, affinity, @@ -346,7 +342,7 @@ impl RoutingHost { self.request_metrics.clone(), context_id.clone(), selected_worker, - selection.attempt, + selection.lease.take(), request, ), }; diff --git a/lib/llm/src/kv_router/routing_host/kv_selection.rs b/lib/llm/src/kv_router/routing_host/kv_selection.rs index 4c71f2dbd393..980b47943c31 100644 --- a/lib/llm/src/kv_router/routing_host/kv_selection.rs +++ b/lib/llm/src/kv_router/routing_host/kv_selection.rs @@ -10,7 +10,9 @@ use dynamo_kv_router::{ protocols::{ BlockExtraInfo, RoutingConstraints, WorkerAffinityTarget, WorkerId, WorkerWithDpRank, }, - scheduling::{AdmissionAttempt, AdvisoryWorkerLoad, QueueRejection, RoutingEligibility}, + scheduling::{ + AdvisoryWorkerLoad, QueueRejection, RoutingEligibility, queue::RequestLifecycleLease, + }, }; use dynamo_runtime::{dynamo_nvtx_range, pipeline::Error}; @@ -29,7 +31,8 @@ use crate::{ pub(super) struct WorkerSelection { pub(super) worker: WorkerWithDpRank, - pub(super) attempt: AdmissionAttempt, + /// The booking's lifecycle lease until the request's cleanup takes it over. + pub(super) lease: Option>, pub(super) overlap_amount: u32, pub(super) effective_overlap_blocks: f64, pub(super) cached_tokens: usize, @@ -135,7 +138,7 @@ impl RoutingHost { kv_hint, } => Ok(SelectionOutcome::Routed(WorkerSelection { worker, - attempt: admitted.attempt, + lease: admitted.lease, overlap_amount: overlap_blocks, effective_overlap_blocks, cached_tokens, @@ -159,7 +162,7 @@ impl RoutingHost { routing_hashes, } => Ok(SelectionOutcome::Routed(WorkerSelection { worker, - attempt: AdmissionAttempt::Untracked, + lease: None, overlap_amount: overlap_blocks, effective_overlap_blocks, cached_tokens, diff --git a/lib/llm/src/kv_router/routing_host/request_guard.rs b/lib/llm/src/kv_router/routing_host/request_guard.rs index 12e26b3906f5..6d96d676b051 100644 --- a/lib/llm/src/kv_router/routing_host/request_guard.rs +++ b/lib/llm/src/kv_router/routing_host/request_guard.rs @@ -7,7 +7,6 @@ use crate::{ kv_router::{ KvRouter, indexer::ApproximateRequestLease, metrics::RouterRequestMetrics, prefill_router::BYPASS_REMOTE_PREFILL_ANNOTATION, request_lease::RequestAttemptLease, - scheduler::SchedulerBookingDescriptor, }, lora::LoadEstimator, preprocessor::PreprocessedRequest, @@ -23,7 +22,7 @@ use dynamo_kv_router::{ BlockExtraInfo, BlockHashOptions, WorkerWithDpRank, compute_block_hash_for_seq, compute_next_seq_hash, }, - scheduling::AdmissionAttempt, + scheduling::queue::RequestLifecycleLease, }; use dynamo_runtime::{ error::DynamoError, @@ -397,35 +396,28 @@ pub(super) struct KvRequestCleanup { } impl KvRequestCleanup { + /// Takes the booking over from `lease`: from here the request's cleanup + /// owns it. pub(super) fn new( chooser: Arc, context_id: String, worker: WorkerWithDpRank, - attempt: AdmissionAttempt, + lease: Option>, ) -> Self { - let attempt_id = match attempt { - AdmissionAttempt::Untracked => None, - AdmissionAttempt::Tracked(attempt_id) => Some(attempt_id), - }; - let approximate_lru = attempt_id - .and_then(|_| chooser.approximate_lru_rank_registration(worker)) - .and_then(|registration| { - chooser.indexer().begin_approximate_lru_request( - worker, - registration.incarnation, - attempt_id?, - ) - }); - let lifecycle = attempt_id.map(|attempt_id| { - chooser.request_lease_manager().register_local( - SchedulerBookingDescriptor { - request_id: context_id.clone(), - worker, - attempt_id, - }, - approximate_lru.clone(), + let booking = lease.and_then(|lease| lease.commit()); + let approximate_lru = booking.as_ref().and_then(|booking| { + let registration = chooser.approximate_lru_rank_registration(worker)?; + chooser.indexer().begin_approximate_lru_request( + worker, + registration.incarnation, + booking.attempt_id, ) }); + let lifecycle = booking.map(|booking| { + chooser + .request_lease_manager() + .register_local(booking, approximate_lru.clone()) + }); Self { chooser, context_id, @@ -570,12 +562,12 @@ impl RequestGuard { request_metrics: Arc, context_id: String, worker: WorkerWithDpRank, - attempt: AdmissionAttempt, + lease: Option>, request: &PreprocessedRequest, ) -> Self { Self::new_kv_with_cleanup( request_metrics, - KvRequestCleanup::new(chooser, context_id, worker, attempt), + KvRequestCleanup::new(chooser, context_id, worker, lease), request, ) } diff --git a/lib/llm/src/kv_router/routing_host/tests.rs b/lib/llm/src/kv_router/routing_host/tests.rs index e31eb37335d0..96787c1cdeab 100644 --- a/lib/llm/src/kv_router/routing_host/tests.rs +++ b/lib/llm/src/kv_router/routing_host/tests.rs @@ -696,7 +696,7 @@ async fn terminal_item_does_not_skip_transport_eof() { Arc::clone(&router.request_metrics), "terminal-drain".to_string(), WorkerWithDpRank::from_worker_id(0), - dynamo_kv_router::scheduling::AdmissionAttempt::Untracked, + None, &request(), ); let monitored = monitor_response_stream(source, context, guard); @@ -765,7 +765,7 @@ async fn shutdown_cancellation_drains_trailing_engine_shutdown_error() { Arc::clone(&router.request_metrics), "shutdown-drain".to_string(), WorkerWithDpRank::from_worker_id(0), - dynamo_kv_router::scheduling::AdmissionAttempt::Untracked, + None, &request(), ); let monitored = monitor_response_stream(source, context, guard); @@ -811,7 +811,7 @@ async fn client_cancellation_still_ends_stream_without_draining() { Arc::clone(&router.request_metrics), "client-cancelled-drain".to_string(), WorkerWithDpRank::from_worker_id(0), - dynamo_kv_router::scheduling::AdmissionAttempt::Untracked, + None, &request(), ); let monitored = monitor_response_stream(source, context, guard); @@ -849,7 +849,7 @@ async fn drain_without_trailing_error_gives_up_at_the_deadline() { Arc::clone(&router.request_metrics), "shutdown-drain-deadline".to_string(), WorkerWithDpRank::from_worker_id(0), - dynamo_kv_router::scheduling::AdmissionAttempt::Untracked, + None, &request(), ); let monitored = monitor_response_stream(source, context, guard); @@ -904,7 +904,7 @@ async fn trailing_error_within_the_drain_window_still_reaches_migration() { Arc::clone(&router.request_metrics), "drain-window-armed".to_string(), WorkerWithDpRank::from_worker_id(0), - dynamo_kv_router::scheduling::AdmissionAttempt::Untracked, + None, &request(), ); let monitored = monitor_response_stream(source, context, guard); @@ -950,7 +950,7 @@ async fn always_ready_terminals_cannot_starve_the_drain_deadline() { Arc::clone(&router.request_metrics), "starvation-guard".to_string(), WorkerWithDpRank::from_worker_id(0), - dynamo_kv_router::scheduling::AdmissionAttempt::Untracked, + None, &request(), ); let monitored = monitor_response_stream(source, context, guard); From 1172f06f30864c8a548e3635857aa143ca26dacd Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Thu, 10 Sep 2026 17:01:32 -0700 Subject: [PATCH 057/102] refactor(router): commit session affinity through the shared table The frontend's affinity coordinator duplicated the commit path that the shared SessionAffinity table already owns: initialize-then-bind, the Soft rebind CAS, and the Hard validate-and-invalidate. The coordinator now returns the table's Hold and commits through SessionAffinity::commit; the frontend AffinityAcquire enum and its copy of that logic are deleted and Hold::invalidate moves next to the table. The Hard/Soft mode is configured on the table when the coordinator is built (AffinityCoordinator::new(ttl, mode), KvRouter::affinity_coordinator, create_affinity_coordinator). RoutingHost derives its mode from the coordinator instead of taking a second copy, so the constructors that accept a coordinator lose their mode parameter and the two can no longer disagree. A partition rejects a second host requesting a different TTL or mode for the same table. The three dispatch-site matches collapse into RoutingHost::bind_affinity. Behavior changes: Hard-mode dispatch mismatch errors now name the actual session id instead of the literal "session"; KvPushRouter::new on a partition another host configured Soft fails at construction. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Ishan Dhanani --- .../src/services/selection/affinity.rs | 8 ++ .../src/services/selection/core/mod.rs | 19 ++- lib/llm/src/entrypoint/input/common.rs | 9 +- lib/llm/src/kv_router.rs | 3 +- lib/llm/src/kv_router/embedded.rs | 4 +- .../kv_router/prefill_router/activation.rs | 18 ++- lib/llm/src/kv_router/prefill_router/query.rs | 8 +- lib/llm/src/kv_router/routing_host.rs | 87 +++++------- lib/llm/src/kv_router/routing_host/builtin.rs | 10 +- lib/llm/src/kv_router/routing_host/kv.rs | 14 +- .../kv_router/routing_host/request_guard.rs | 2 +- lib/llm/src/kv_router/routing_host/tests.rs | 81 +++++------ lib/llm/src/session_affinity/coordinator.rs | 105 ++++---------- lib/llm/src/session_affinity/mod.rs | 6 +- lib/llm/src/session_affinity/replica_sync.rs | 18 ++- lib/llm/src/session_affinity/tests.rs | 131 +++++++++--------- 16 files changed, 232 insertions(+), 291 deletions(-) diff --git a/lib/kv-router/src/services/selection/affinity.rs b/lib/kv-router/src/services/selection/affinity.rs index 9b9225e290d4..c7a1c3a77ab8 100644 --- a/lib/kv-router/src/services/selection/affinity.rs +++ b/lib/kv-router/src/services/selection/affinity.rs @@ -191,6 +191,14 @@ impl Hold { Self::Bound { target, .. } => Some(*target), } } + + /// `AffinityLease::invalidate` for a bound session; an initializing hold + /// just releases. + pub fn invalidate(self) { + if let Self::Bound { mut lease, .. } = self { + lease.invalidate(); + } + } } /// One step of acquiring a session. diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index caeae7b686a2..19915e33167c 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -150,10 +150,14 @@ impl SelectionEntry { } Ok(table) })?; - if table.ttl() != config.ttl { + if table.ttl() != config.ttl || table.mode() != config.mode { return Err(SelectionError::Conflict(format!( - "session affinity TTL mismatch for {}", - self.key + "session affinity config mismatch for {}: existing=({:?}, {:?}) requested=({:?}, {:?})", + self.key, + table.ttl(), + table.mode(), + config.ttl, + config.mode ))); } Ok(table) @@ -4970,7 +4974,7 @@ mod tests { } #[tokio::test] - async fn affinity_configuration_rejects_invalid_or_conflicting_ttl() { + async fn affinity_configuration_rejects_invalid_or_conflicting_config() { use super::super::service::SelectionServiceBuilder; for ttl in [ Duration::ZERO, @@ -4995,6 +4999,13 @@ mod tests { partition.session_affinity(SessionAffinityConfig::new(Duration::from_secs(20))), Err(SelectionError::Conflict(_)) )); + assert!(matches!( + partition.session_affinity( + SessionAffinityConfig::new(Duration::from_secs(10)) + .with_mode(SessionAffinityMode::Soft) + ), + Err(SelectionError::Conflict(_)) + )); } fn hint_config(worker_type: &str, endpoints: &[(u32, &str)]) -> SelectionWorkerConfig { diff --git a/lib/llm/src/entrypoint/input/common.rs b/lib/llm/src/entrypoint/input/common.rs index 2f9d2ed9b072..048438914355 100644 --- a/lib/llm/src/entrypoint/input/common.rs +++ b/lib/llm/src/entrypoint/input/common.rs @@ -150,7 +150,6 @@ fn validate_router_mode_for_lora( } } -#[expect(clippy::too_many_arguments)] fn preprocessed_backend_engine( router: LlmPushRouter, router_mode: RouterMode, @@ -158,7 +157,6 @@ fn preprocessed_backend_engine( model_manager: &Arc, endpoint_id: &dynamo_runtime::protocols::EndpointId, affinity: Option, - session_affinity_mode: SessionAffinityMode, load_context: Arc, ) -> anyhow::Result> { // Reject LoRA + unsupported-mode combinations up front (single source of truth, shared with @@ -180,7 +178,6 @@ fn preprocessed_backend_engine( chooser, load_context, affinity, - session_affinity_mode, )) } _ => { @@ -191,7 +188,6 @@ fn preprocessed_backend_engine( router, load_context, affinity, - session_affinity_mode, lora, )?) } @@ -257,11 +253,11 @@ pub(crate) async fn build_preprocessed_routing_with_session_affinity_mode( let ttl = session_affinity_ttl_secs.map(Duration::from_secs); let affinity = match (ttl, chooser.as_ref()) { (Some(ttl), Some(chooser)) => { - let affinity = chooser.affinity_coordinator(ttl)?; + let affinity = chooser.affinity_coordinator(ttl, session_affinity_mode)?; affinity.enable_replica_sync(router_client.clone()).await?; Some(affinity) } - _ => create_affinity_coordinator(ttl, router_client.clone()).await?, + _ => create_affinity_coordinator(ttl, session_affinity_mode, router_client.clone()).await?, }; let embedding_cache_indexer = if enable_multimodal_cache_indexer @@ -306,7 +302,6 @@ pub(crate) async fn build_preprocessed_routing_with_session_affinity_mode( &model_manager, &endpoint_id, affinity, - session_affinity_mode, load_context, )?; if router_mode.is_kv_routing() && prefill_router.conditional_disagg_enabled() { diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index 95b3c6d1cca7..09c211d9c0cc 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -1614,8 +1614,9 @@ impl KvRouter { pub(crate) fn affinity_coordinator( &self, ttl: std::time::Duration, + mode: crate::session_affinity::SessionAffinityMode, ) -> anyhow::Result { - self.scheduler.affinity_coordinator(ttl) + self.scheduler.affinity_coordinator(ttl, mode) } pub(crate) fn booking_cleanup(&self) -> scheduler::SchedulerBookingCleanup { diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs index bc1315fdba75..c665f0d36a88 100644 --- a/lib/llm/src/kv_router/embedded.rs +++ b/lib/llm/src/kv_router/embedded.rs @@ -371,9 +371,11 @@ impl EmbeddedSelection { pub(crate) fn affinity_coordinator( &self, ttl: std::time::Duration, + mode: crate::session_affinity::SessionAffinityMode, ) -> Result { let table = self.partition.session_affinity( - dynamo_kv_router::services::selection::affinity::SessionAffinityConfig::new(ttl), + dynamo_kv_router::services::selection::affinity::SessionAffinityConfig::new(ttl) + .with_mode(mode), )?; Ok(self .affinity diff --git a/lib/llm/src/kv_router/prefill_router/activation.rs b/lib/llm/src/kv_router/prefill_router/activation.rs index da0779032f6c..e5f7b330e612 100644 --- a/lib/llm/src/kv_router/prefill_router/activation.rs +++ b/lib/llm/src/kv_router/prefill_router/activation.rs @@ -365,8 +365,12 @@ impl PrefillRouter { ) .await?; - let affinity = - create_affinity_coordinator(prefill_session_affinity_ttl, client.clone()).await?; + let affinity = create_affinity_coordinator( + prefill_session_affinity_ttl, + context.session_affinity_mode, + client.clone(), + ) + .await?; // Build the PushRouter for prefill with KV mode using the shared client let push_router = PushRouter::>::from_client_with_monitor( @@ -381,11 +385,14 @@ impl PrefillRouter { kv_chooser, load_context.clone(), affinity, - context.session_affinity_mode, )) } else { - let affinity = - create_affinity_coordinator(prefill_session_affinity_ttl, client.clone()).await?; + let affinity = create_affinity_coordinator( + prefill_session_affinity_ttl, + context.session_affinity_mode, + client.clone(), + ) + .await?; // Create the transport and discovery layer for the builtin policy. // Note: Per-worker metrics (active_prefill_tokens, active_decode_blocks) are only @@ -401,7 +408,6 @@ impl PrefillRouter { push_router, load_context.clone(), affinity, - context.session_affinity_mode, )?) }; diff --git a/lib/llm/src/kv_router/prefill_router/query.rs b/lib/llm/src/kv_router/prefill_router/query.rs index 8e467d65a0cb..8c40384e5bb7 100644 --- a/lib/llm/src/kv_router/prefill_router/query.rs +++ b/lib/llm/src/kv_router/prefill_router/query.rs @@ -449,13 +449,7 @@ mod tests { .await .unwrap(); let shared = Arc::new( - RoutingHost::new_builtin_with_coordinator( - push_router, - load_context, - None, - crate::session_affinity::SessionAffinityMode::Hard, - ) - .unwrap(), + RoutingHost::new_builtin_with_coordinator(push_router, load_context, None).unwrap(), ); let prefill = PrefillRouter::disabled(Arc::new(ModelManager::new()), mode, None); prefill.binding.store(Some(Arc::new( diff --git a/lib/llm/src/kv_router/routing_host.rs b/lib/llm/src/kv_router/routing_host.rs index 6a3537fdd19e..8c95432a3db3 100644 --- a/lib/llm/src/kv_router/routing_host.rs +++ b/lib/llm/src/kv_router/routing_host.rs @@ -38,8 +38,8 @@ use crate::{ timing::{RequestPhase, RoutingData, WORKER_TYPE_DECODE, WORKER_TYPE_PREFILL}, }, session_affinity::{ - AffinityAcquire, AffinityCoordinator, AffinityTarget, SessionAffinityMode, affinity_id, - explicit_target, invalid_argument, + AffinityCoordinator, AffinityTarget, Hold, SessionAffinityMode, affinity_id, + explicit_target, from_table, invalid_argument, }, }; @@ -248,7 +248,7 @@ pub(crate) struct RoutePlan { signals: RoutePlanSignals, selection: WorkerSelection, cleanup: KvRequestCleanup, - affinity: Option, + affinity: Option, /// Carried forward from the [`RoutePreview`] this plan was admitted from, so /// preview, admission and dispatch draw on one budget instead of three. budget: CleanupBudget, @@ -314,6 +314,12 @@ impl RoutePlan { /// removed only in a 2.0.0 (or later) breaking release. pub type KvPushRouter = RoutingHost; +/// The host steers and retries by the same mode its table commits with; a +/// host without session affinity never reads it. +fn affinity_mode(affinity: Option<&AffinityCoordinator>) -> SessionAffinityMode { + affinity.map(AffinityCoordinator::mode).unwrap_or_default() +} + impl RoutingHost { pub fn new( inner: PushRouter>, @@ -321,15 +327,10 @@ impl RoutingHost { session_affinity_ttl: Option, ) -> Result { let affinity = session_affinity_ttl - .map(|ttl| kv_router.affinity_coordinator(ttl)) + .map(|ttl| kv_router.affinity_coordinator(ttl, SessionAffinityMode::Hard)) .transpose()?; - Ok(Self::new_with_coordinator( - inner, - kv_router, - affinity, - SessionAffinityMode::Hard, - )) + Ok(Self::new_with_coordinator(inner, kv_router, affinity)) } pub fn new_with_load_context( @@ -340,7 +341,7 @@ impl RoutingHost { session_affinity_mode: SessionAffinityMode, ) -> Result { let affinity = session_affinity_ttl - .map(|ttl| kv_router.affinity_coordinator(ttl)) + .map(|ttl| kv_router.affinity_coordinator(ttl, session_affinity_mode)) .transpose()?; Ok(Self::new_with_load_context_and_coordinator( @@ -348,7 +349,6 @@ impl RoutingHost { kv_router, load_context, affinity, - session_affinity_mode, )) } @@ -356,15 +356,8 @@ impl RoutingHost { inner: PushRouter>, kv_router: Arc, affinity: Option, - session_affinity_mode: SessionAffinityMode, ) -> Self { - Self::new_with_optional_load_context_and_coordinator( - inner, - kv_router, - None, - affinity, - session_affinity_mode, - ) + Self::new_with_optional_load_context_and_coordinator(inner, kv_router, None, affinity) } pub(crate) fn new_with_load_context_and_coordinator( @@ -372,14 +365,12 @@ impl RoutingHost { kv_router: Arc, load_context: Arc, affinity: Option, - session_affinity_mode: SessionAffinityMode, ) -> Self { Self::new_with_optional_load_context_and_coordinator( inner, kv_router, Some(load_context), affinity, - session_affinity_mode, ) } @@ -388,7 +379,6 @@ impl RoutingHost { kv_router: Arc, load_context: Option>, affinity: Option, - session_affinity_mode: SessionAffinityMode, ) -> Self { // Eagerly register router request metrics (as zeros) so they are // scrapeable before any requests arrive. Both the frontend pipeline @@ -400,8 +390,8 @@ impl RoutingHost { inner, policy: RoutingPolicy::Kv(kv_router), request_metrics, + session_affinity_mode: affinity_mode(affinity.as_ref()), affinity, - session_affinity_mode, hosted_occupancy: None, lora: None, routing_context: load_context, @@ -413,35 +403,21 @@ impl RoutingHost { inner: PushRouter>, load_context: Arc, ) -> Result { - Self::new_builtin_with_capabilities( - inner, - load_context, - None, - SessionAffinityMode::Hard, - None, - ) + Self::new_builtin_with_capabilities(inner, load_context, None, None) } pub(crate) fn new_builtin_with_coordinator( inner: PushRouter>, load_context: Arc, affinity: Option, - session_affinity_mode: SessionAffinityMode, ) -> Result { - Self::new_builtin_with_capabilities( - inner, - load_context, - affinity, - session_affinity_mode, - None, - ) + Self::new_builtin_with_capabilities(inner, load_context, affinity, None) } pub(crate) fn new_builtin_with_capabilities( inner: PushRouter>, load_context: Arc, affinity: Option, - session_affinity_mode: SessionAffinityMode, lora: Option<(Arc, Arc)>, ) -> Result { if affinity.is_some() && lora.is_some() { @@ -488,8 +464,8 @@ impl RoutingHost { inner, policy, request_metrics, + session_affinity_mode: affinity_mode(affinity.as_ref()), affinity, - session_affinity_mode, hosted_occupancy, lora: lora .zip(lora_selector) @@ -551,6 +527,20 @@ impl RoutingHost { } } + /// Commit a held session to the dispatched worker; a request without a + /// session passes its stream through. + fn bind_affinity( + &self, + hold: Option, + dispatched_target: AffinityTarget, + stream: ManyOut>, + ) -> Result>, Error> { + let (Some(hold), Some(affinity)) = (hold, self.affinity.as_ref()) else { + return Ok(stream); + }; + affinity.commit_to_stream(hold, dispatched_target, stream) + } + fn affinity_target_is_valid(&self, target: AffinityTarget) -> bool { if !self.inner.client.is_instance_discovered(target.worker_id) { return false; @@ -589,7 +579,7 @@ impl RoutingHost { phase: RequestPhase, staged_kv: StagedKv, budget: &CleanupBudget, - ) -> Result { + ) -> Result { match DispatchCancellation::for_request(phase, staged_kv) { DispatchCancellation::CancelWhenStopped => { affinity @@ -616,7 +606,7 @@ impl RoutingHost { is_query_only: bool, budget: &CleanupBudget, mut select: Select, - ) -> Result<(T, Option), Error> + ) -> Result<(T, Option), Error> where Select: FnMut(Option) -> SelectionFuture, SelectionFuture: Future>, @@ -646,7 +636,7 @@ impl RoutingHost { budget, ) .await?; - let target = operation.target(); + let target = operation.target().map(from_table); match select(target).await { Ok(selection) => Ok((selection, Some(operation))), Err(error) if is_cancelled(&error) => Err(error), @@ -667,7 +657,7 @@ impl RoutingHost { budget, ) .await?; - let selection = select(retry.target()).await?; + let selection = select(retry.target().map(from_table)).await?; Ok((selection, Some(retry))) } Err(error) => Err(error), @@ -819,12 +809,7 @@ impl AsyncEngine, ManyOut { - operation.into_stream(selected_target, stream, self.session_affinity_mode) - } - None => Ok(stream), - } + self.bind_affinity(operation, selected_target, stream) } } diff --git a/lib/llm/src/kv_router/routing_host/builtin.rs b/lib/llm/src/kv_router/routing_host/builtin.rs index 3b91e0c699e3..17faf4beeb23 100644 --- a/lib/llm/src/kv_router/routing_host/builtin.rs +++ b/lib/llm/src/kv_router/routing_host/builtin.rs @@ -343,7 +343,7 @@ impl RoutingHost { device_aware_telemetry, } = selection; let soft_affinity_target = if self.session_affinity_mode == SessionAffinityMode::Soft { - operation.as_ref().and_then(AffinityAcquire::target) + operation.as_ref().and_then(Hold::target).map(from_table) } else { None }; @@ -518,13 +518,7 @@ impl RoutingHost { } guard.mark_dispatched(); let stream = into_monitored_response(response_stream, guard); - match operation { - Some(operation) => Ok(( - metadata, - operation.into_stream(target, stream, self.session_affinity_mode)?, - )), - None => Ok((metadata, stream)), - } + Ok((metadata, self.bind_affinity(operation, target, stream)?)) } } diff --git a/lib/llm/src/kv_router/routing_host/kv.rs b/lib/llm/src/kv_router/routing_host/kv.rs index 90659dd2e56b..83b86ef753eb 100644 --- a/lib/llm/src/kv_router/routing_host/kv.rs +++ b/lib/llm/src/kv_router/routing_host/kv.rs @@ -87,7 +87,7 @@ impl RoutingHost { phase: RequestPhase, is_query_only: bool, budget: &CleanupBudget, - ) -> Result<(WorkerSelection, Option), Error> { + ) -> Result<(WorkerSelection, Option), Error> { self.select_with_session_affinity(request, phase, is_query_only, budget, |target| { self.select_request(request, phase, is_query_only, target, budget) }) @@ -246,12 +246,7 @@ impl RoutingHost { return Err(error); } }; - match affinity { - Some(affinity) => { - affinity.into_stream(selected_target, stream, self.session_affinity_mode) - } - None => Ok(stream), - } + self.bind_affinity(affinity, selected_target, stream) } pub(crate) async fn prefill_worker_busy( @@ -600,12 +595,9 @@ impl RoutingHost { return Err(error); } }; - let Some(operation) = operation else { - return Ok((metadata, stream)); - }; Ok(( metadata, - operation.into_stream(selected_target, stream, self.session_affinity_mode)?, + self.bind_affinity(operation, selected_target, stream)?, )) } } diff --git a/lib/llm/src/kv_router/routing_host/request_guard.rs b/lib/llm/src/kv_router/routing_host/request_guard.rs index 6d96d676b051..f59dce786977 100644 --- a/lib/llm/src/kv_router/routing_host/request_guard.rs +++ b/lib/llm/src/kv_router/routing_host/request_guard.rs @@ -542,7 +542,7 @@ impl OutputBlockTracker { /// Coordinates scheduler cleanup, observability, and streamed load tracking. /// -/// Session-affinity lifetime is separate: `AffinityAcquire` and +/// Session-affinity lifetime is separate: the affinity `Hold` and /// `AffinityLease` own binding commit, release, and invalidation. pub(super) struct RequestGuard { cleanup: RequestCleanup, diff --git a/lib/llm/src/kv_router/routing_host/tests.rs b/lib/llm/src/kv_router/routing_host/tests.rs index 96787c1cdeab..3ee32c1b4c4c 100644 --- a/lib/llm/src/kv_router/routing_host/tests.rs +++ b/lib/llm/src/kv_router/routing_host/tests.rs @@ -226,14 +226,13 @@ async fn builtin_direct_without_worker_is_invalid_argument() { let inner = PushRouter::from_client(client, RouterMode::Direct) .await .unwrap(); - let affinity = AffinityCoordinator::new(Duration::from_secs(10)).unwrap(); - let host = RoutingHost::new_builtin_with_coordinator( - inner, - load_context, - Some(affinity), + let affinity = AffinityCoordinator::new( + Duration::from_secs(10), crate::session_affinity::SessionAffinityMode::Hard, ) .unwrap(); + let host = + RoutingHost::new_builtin_with_coordinator(inner, load_context, Some(affinity)).unwrap(); let error = host .generate(affinity_request("direct-unbound", None)) @@ -273,14 +272,14 @@ async fn builtin_direct_uses_bound_soft_affinity_as_exact_target() { ) .await .unwrap(); - let affinity = AffinityCoordinator::new(Duration::from_secs(10)).unwrap(); - let host = RoutingHost::new_builtin_with_coordinator( - inner, - load_context, - Some(affinity.clone()), + let affinity = AffinityCoordinator::new( + Duration::from_secs(10), crate::session_affinity::SessionAffinityMode::Soft, ) .unwrap(); + let host = + RoutingHost::new_builtin_with_coordinator(inner, load_context, Some(affinity.clone())) + .unwrap(); let session_id = SessionAffinityId::new("direct-soft-bound"); bind_affinity_target(&host, &session_id, AffinityTarget::worker(worker_id)).await; @@ -359,19 +358,17 @@ async fn builtin_hard_affinity_ignores_local_inhibition() { ) .await .unwrap(); - let affinity = AffinityCoordinator::new(Duration::from_secs(10)).unwrap(); - let host = RoutingHost::new_builtin_with_coordinator( - inner, - load_context, - Some(affinity.clone()), + let affinity = AffinityCoordinator::new( + Duration::from_secs(10), crate::session_affinity::SessionAffinityMode::Hard, ) .unwrap(); + let host = + RoutingHost::new_builtin_with_coordinator(inner, load_context, Some(affinity.clone())) + .unwrap(); let session_id = SessionAffinityId::new("local-inhibition"); - let AffinityAcquire::Initialize(initializer) = - affinity.acquire(&session_id, None).await.unwrap() - else { + let Hold::Initialize(initializer) = affinity.acquire(&session_id, None).await.unwrap() else { panic!("new affinity session must initialize"); }; drop( @@ -445,7 +442,6 @@ async fn builtin_lora_keeps_separate_selection_and_cleanup() { inner, load_context, None, - crate::session_affinity::SessionAffinityMode::Hard, Some((filter, Arc::clone(&estimator))), ) .unwrap(); @@ -518,14 +514,14 @@ async fn builtin_affinity_uses_common_host_for_every_policy() { ) .await .unwrap(); - let affinity = AffinityCoordinator::new(Duration::from_secs(10)).unwrap(); - let host = RoutingHost::new_builtin_with_coordinator( - inner, - load_context, - Some(affinity.clone()), + let affinity = AffinityCoordinator::new( + Duration::from_secs(10), crate::session_affinity::SessionAffinityMode::Hard, ) .unwrap(); + let host = + RoutingHost::new_builtin_with_coordinator(inner, load_context, Some(affinity.clone())) + .unwrap(); let session_id = format!("session-{index}"); let affinity_id = SessionAffinityId::new(session_id.clone()); let explicit_worker = (mode == RouterMode::Direct).then_some(worker_id); @@ -625,14 +621,14 @@ async fn builtin_direct_fallback_stays_disabled_for_affinity() { ) .await .unwrap(); - let affinity = AffinityCoordinator::new(Duration::from_secs(10)).unwrap(); - let host = RoutingHost::new_builtin_with_coordinator( - inner, - load_context, - Some(affinity.clone()), + let affinity = AffinityCoordinator::new( + Duration::from_secs(10), crate::session_affinity::SessionAffinityMode::Hard, ) .unwrap(); + let host = + RoutingHost::new_builtin_with_coordinator(inner, load_context, Some(affinity.clone())) + .unwrap(); let mut standalone = request(); standalone.routing_mut().backend_instance_id = Some(stale_worker); @@ -644,9 +640,7 @@ async fn builtin_direct_fallback_stays_disabled_for_affinity() { ); let session_id = SessionAffinityId::new("direct-affinity"); - let AffinityAcquire::Initialize(initializer) = - affinity.acquire(&session_id, None).await.unwrap() - else { + let Hold::Initialize(initializer) = affinity.acquire(&session_id, None).await.unwrap() else { panic!("new affinity session must initialize"); }; drop( @@ -1666,7 +1660,7 @@ async fn route_preview_does_not_acquire_session_affinity() { .await .expect("preview must not leave affinity initialization pending") .unwrap(); - assert!(matches!(acquisition, AffinityAcquire::Initialize(_))); + assert!(matches!(acquisition, Hold::Initialize(_))); drop(acquisition); drop(router); @@ -1753,7 +1747,7 @@ async fn aborted_route_plan_drops_pending_affinity_initialization() { .await .expect("abandoned plan must not leave affinity initialization pending") .unwrap(); - assert!(matches!(acquisition, AffinityAcquire::Initialize(_))); + assert!(matches!(acquisition, Hold::Initialize(_))); drop(acquisition); drop(router); @@ -1876,9 +1870,7 @@ async fn session_affinity_post_selection_failures_preserve_binding() { worker_id: 7, dp_rank: Some(0), }; - let AffinityAcquire::Initialize(initializer) = - affinity.acquire(&session_id, None).await.unwrap() - else { + let Hold::Initialize(initializer) = affinity.acquire(&session_id, None).await.unwrap() else { panic!("first request must initialize"); }; drop( @@ -1913,7 +1905,7 @@ async fn session_affinity_existing_selection_cancellation_preserves_binding_with worker_id: 7, dp_rank: Some(0), }; - let AffinityAcquire::Initialize(initializer) = router + let Hold::Initialize(initializer) = router .affinity .as_ref() .unwrap() @@ -1960,7 +1952,7 @@ async fn session_affinity_existing_selection_cancellation_preserves_binding_with Some(original_target) ); - let AffinityAcquire::Bound { target, lease } = router + let Hold::Bound { target, lease } = router .affinity .as_ref() .unwrap() @@ -1970,7 +1962,7 @@ async fn session_affinity_existing_selection_cancellation_preserves_binding_with else { panic!("cancellation must preserve the existing binding"); }; - assert_eq!(target, original_target); + assert_eq!(target, crate::session_affinity::to_table(original_target)); drop(lease); drop(router); @@ -1982,7 +1974,7 @@ async fn bind_affinity_target( session_id: &SessionAffinityId, target: AffinityTarget, ) { - let AffinityAcquire::Initialize(initializer) = router + let Hold::Initialize(initializer) = router .affinity .as_ref() .unwrap() @@ -2087,7 +2079,7 @@ async fn stale_affinity_rank_recovers_within_request() { .await .unwrap(); assert_eq!(selection.worker, WorkerWithDpRank::new(7, 0)); - assert!(matches!(operation, Some(AffinityAcquire::Initialize(_)))); + assert!(matches!(operation, Some(Hold::Initialize(_)))); router.kv_router().free(request.id()).await.unwrap(); drop(operation); @@ -2154,7 +2146,7 @@ async fn migration_exclusion_preserves_hard_affinity_without_widening_or_escapin worker_id: 7, dp_rank: Some(0), }; - let AffinityAcquire::Initialize(initializer) = router + let Hold::Initialize(initializer) = router .affinity .as_ref() .unwrap() @@ -2913,8 +2905,7 @@ async fn kv_stopped_decode_request_survives_a_contended_session_affinity_wait() // leg has to wait rather than take the slot immediately. let router = Arc::new(router); let coordinator = router.affinity.as_ref().unwrap().clone(); - let AffinityAcquire::Initialize(holder) = coordinator.acquire(&session_id, None).await.unwrap() - else { + let Hold::Initialize(holder) = coordinator.acquire(&session_id, None).await.unwrap() else { panic!("the first acquisition must initialize the session"); }; diff --git a/lib/llm/src/session_affinity/coordinator.rs b/lib/llm/src/session_affinity/coordinator.rs index 9f9d80c27c93..c06dd089f6ed 100644 --- a/lib/llm/src/session_affinity/coordinator.rs +++ b/lib/llm/src/session_affinity/coordinator.rs @@ -15,8 +15,7 @@ use std::{ }; use dynamo_kv_router::services::selection::affinity::{ - AcquireStep, AffinityError, AffinityInitialization, AffinityLease, Hold, SessionAffinity, - validate_dispatch_target, + AcquireStep, AffinityError, AffinityLease, Hold, SessionAffinity, SessionAffinityConfig, }; use dynamo_runtime::{ engine::{AsyncEngineContext, AsyncEngineContextProvider}, @@ -50,7 +49,7 @@ pub(crate) fn to_table(target: AffinityTarget) -> TableTarget { TableTarget::new(target.worker_id, target.dp_rank) } -fn from_table(target: TableTarget) -> AffinityTarget { +pub(crate) fn from_table(target: TableTarget) -> AffinityTarget { AffinityTarget::new(target.worker_id, target.dp_rank) } @@ -73,9 +72,10 @@ pub struct AffinityCoordinator { } impl AffinityCoordinator { - pub fn new(ttl: Duration) -> Result { + pub fn new(ttl: Duration, mode: SessionAffinityMode) -> Result { Ok(Self::wrap( - SessionAffinity::new(ttl).map_err(affinity_error)?, + SessionAffinity::with_config(SessionAffinityConfig::new(ttl).with_mode(mode)) + .map_err(affinity_error)?, )) } @@ -123,7 +123,7 @@ impl AffinityCoordinator { &self, session_id: &SessionAffinityId, requested_target: Option, - ) -> Result { + ) -> Result { self.acquire_inner(session_id, requested_target, None).await } @@ -132,7 +132,7 @@ impl AffinityCoordinator { session_id: &SessionAffinityId, requested_target: Option, request_context: &dyn AsyncEngineContext, - ) -> Result { + ) -> Result { self.acquire_inner(session_id, requested_target, Some(request_context)) .await } @@ -142,7 +142,7 @@ impl AffinityCoordinator { session_id: &SessionAffinityId, requested_target: Option, request_context: Option<&dyn AsyncEngineContext>, - ) -> Result { + ) -> Result { let requested = requested_target.map(to_table); loop { match self @@ -151,15 +151,7 @@ impl AffinityCoordinator { .try_acquire(session_id.as_str(), requested) .map_err(affinity_error)? { - AcquireStep::Held(Hold::Initialize(initialization)) => { - return Ok(AffinityAcquire::Initialize(initialization)); - } - AcquireStep::Held(Hold::Bound { target, lease }) => { - return Ok(AffinityAcquire::Bound { - target: from_table(target), - lease, - }); - } + AcquireStep::Held(hold) => return Ok(hold), AcquireStep::Wait(notified) => match request_context { Some(context) => { tokio::select! { @@ -187,6 +179,26 @@ impl AffinityCoordinator { .map_err(affinity_error) } + /// The mode the table binds sessions with. + pub(crate) fn mode(&self) -> SessionAffinityMode { + self.inner.table.mode() + } + + /// `SessionAffinity::commit`, holding the lease until `stream` ends. + pub(crate) fn commit_to_stream( + &self, + hold: Hold, + dispatched_target: AffinityTarget, + stream: ManyOut, + ) -> Result, Error> { + let lease = self + .inner + .table + .commit(hold, to_table(dispatched_target)) + .map_err(affinity_error)?; + Ok(tracked_stream(lease, stream)) + } + #[cfg(test)] pub(super) fn entry_count(&self) -> usize { self.inner.table.entry_count() @@ -287,65 +299,6 @@ impl AffinityCoordinator { } } -pub(crate) enum AffinityAcquire { - Initialize(AffinityInitialization), - Bound { - target: AffinityTarget, - lease: AffinityLease, - }, -} - -impl AffinityAcquire { - pub(crate) fn target(&self) -> Option { - match self { - Self::Initialize(_) => None, - Self::Bound { target, .. } => Some(*target), - } - } - - /// Bind (or confirm) the session to where the request was dispatched and - /// hold the lease until `stream` ends. - pub(crate) fn into_stream( - self, - dispatched_target: AffinityTarget, - stream: ManyOut, - mode: SessionAffinityMode, - ) -> Result, Error> { - let dispatched = to_table(dispatched_target); - match self { - Self::Initialize(initialization) => { - let lease = initialization.commit(dispatched).map_err(affinity_error)?; - lease.publish(dispatched); - Ok(tracked_stream(lease, stream)) - } - Self::Bound { target, mut lease } => { - let bound = to_table(target); - if mode == SessionAffinityMode::Soft { - let rebound = AffinityLease::rebound_target(bound, dispatched); - if bound == rebound { - lease.publish(bound); - } else if lease.rebind(bound, rebound) { - lease.publish(rebound); - } - return Ok(tracked_stream(lease, stream)); - } - if let Err(error) = validate_dispatch_target("session", bound, dispatched) { - lease.invalidate(); - return Err(affinity_error(error)); - } - lease.publish(bound); - Ok(tracked_stream(lease, stream)) - } - } - } - - pub(crate) fn invalidate(self) { - if let Self::Bound { mut lease, .. } = self { - lease.invalidate(); - } - } -} - pub(super) fn tracked_stream( lease: AffinityLease, stream: ManyOut, diff --git a/lib/llm/src/session_affinity/mod.rs b/lib/llm/src/session_affinity/mod.rs index 5b0341c26423..18108044eb8b 100644 --- a/lib/llm/src/session_affinity/mod.rs +++ b/lib/llm/src/session_affinity/mod.rs @@ -10,8 +10,9 @@ use dynamo_runtime::{component::Client, pipeline::Error}; #[cfg(test)] pub(crate) use coordinator::to_table; -pub(crate) use coordinator::{AffinityAcquire, affinity_id, invalid_argument}; pub use coordinator::{AffinityCoordinator, AffinityTarget, explicit_target}; +pub(crate) use coordinator::{affinity_id, from_table, invalid_argument}; +pub(crate) use dynamo_kv_router::services::selection::affinity::Hold; pub use dynamo_kv_router::services::selection::affinity::{ MAX_SESSION_AFFINITY_ENTRIES, MAX_SESSION_AFFINITY_ID_BYTES, MAX_SESSION_AFFINITY_TTL_SECS, SessionAffinityMode, @@ -22,12 +23,13 @@ pub type LlmResponse = pub(crate) async fn create_affinity_coordinator( ttl: Option, + mode: SessionAffinityMode, client: Client, ) -> Result, Error> { let Some(ttl) = ttl else { return Ok(None); }; - let coordinator = AffinityCoordinator::new(ttl)?; + let coordinator = AffinityCoordinator::new(ttl, mode)?; coordinator.enable_replica_sync(client).await?; Ok(Some(coordinator)) } diff --git a/lib/llm/src/session_affinity/replica_sync.rs b/lib/llm/src/session_affinity/replica_sync.rs index bf3ef835a37d..39539db3b3e4 100644 --- a/lib/llm/src/session_affinity/replica_sync.rs +++ b/lib/llm/src/session_affinity/replica_sync.rs @@ -346,7 +346,11 @@ mod tests { #[tokio::test] async fn rejected_worker_update_still_advances_replica_clock() { - let coordinator = AffinityCoordinator::new(Duration::from_secs(10)).unwrap(); + let coordinator = AffinityCoordinator::new( + Duration::from_secs(10), + crate::session_affinity::SessionAffinityMode::Hard, + ) + .unwrap(); let baseline = coordinator.next_version_for_test(); let sequence = baseline.sequence.saturating_add(10); let (_tx, discovered_instances) = watch::channel(Vec::new()); @@ -393,7 +397,11 @@ mod tests { )); let client = endpoint.client().await.unwrap(); - let original = AffinityCoordinator::new(Duration::from_secs(10)).unwrap(); + let original = AffinityCoordinator::new( + Duration::from_secs(10), + crate::session_affinity::SessionAffinityMode::Hard, + ) + .unwrap(); let shared = original.clone(); let (first, second) = tokio::join!( original.enable_replica_sync(client.clone()), @@ -409,7 +417,11 @@ mod tests { drop(shared); wait_for_registration_count(&drt, &query, 0).await; - let replacement = AffinityCoordinator::new(Duration::from_secs(10)).unwrap(); + let replacement = AffinityCoordinator::new( + Duration::from_secs(10), + crate::session_affinity::SessionAffinityMode::Hard, + ) + .unwrap(); replacement.enable_replica_sync(client).await.unwrap(); wait_for_registration_count(&drt, &query, 1).await; diff --git a/lib/llm/src/session_affinity/tests.rs b/lib/llm/src/session_affinity/tests.rs index 00762006e4c1..20681ac8430f 100644 --- a/lib/llm/src/session_affinity/tests.rs +++ b/lib/llm/src/session_affinity/tests.rs @@ -13,7 +13,7 @@ use futures::{StreamExt, stream}; use super::SessionAffinityMode::{Hard, Soft}; use super::{ - AffinityAcquire, AffinityCoordinator, AffinityTarget, LlmResponse, affinity_id, + AffinityCoordinator, AffinityTarget, Hold, LlmResponse, affinity_id, coordinator::{ReplicaApplyOutcome, tracked_stream}, explicit_target, to_table, }; @@ -37,7 +37,11 @@ fn target(worker_id: u64, dp_rank: Option) -> AffinityTarget { } fn coordinator() -> AffinityCoordinator { - AffinityCoordinator::new(Duration::from_secs(10)).unwrap() + AffinityCoordinator::new(Duration::from_secs(10), Hard).unwrap() +} + +fn soft_coordinator() -> AffinityCoordinator { + AffinityCoordinator::new(Duration::from_secs(10), Soft).unwrap() } fn response_stream(items: usize) -> dynamo_runtime::pipeline::ManyOut { @@ -63,8 +67,8 @@ fn cancelled_response_stream() -> dynamo_runtime::pipeline::ManyOut async fn bind(coordinator: &AffinityCoordinator, target: AffinityTarget) { let operation = coordinator.acquire(&session_id(), None).await.unwrap(); - let mut stream = operation - .into_stream(target, response_stream(1), Hard) + let mut stream = coordinator + .commit_to_stream(operation, target, response_stream(1)) .unwrap(); while stream.next().await.is_some() {} } @@ -148,7 +152,7 @@ fn session_affinity_context_type_errors_are_preserved() { async fn session_affinity_initialization_is_atomic() { let coordinator = coordinator(); let first = coordinator.acquire(&session_id(), None).await.unwrap(); - let AffinityAcquire::Initialize(first) = first else { + let Hold::Initialize(first) = first else { panic!("first request must initialize"); }; @@ -159,14 +163,14 @@ async fn session_affinity_initialization_is_atomic() { let first_lease = first.commit(to_table(target(7, Some(0)))).unwrap(); let second = waiter.await.unwrap().unwrap(); - let AffinityAcquire::Bound { + let Hold::Bound { target: second_target, lease: second_lease, } = second else { panic!("waiter must acquire the committed binding"); }; - assert_eq!(second_target, target(7, Some(0))); + assert_eq!(second_target, to_table(target(7, Some(0)))); drop(first_lease); drop(second_lease); } @@ -175,7 +179,7 @@ async fn session_affinity_initialization_is_atomic() { async fn session_affinity_initializer_cancellation_wakes_waiter() { let coordinator = coordinator(); let first = coordinator.acquire(&session_id(), None).await.unwrap(); - let AffinityAcquire::Initialize(first) = first else { + let Hold::Initialize(first) = first else { panic!("first request must initialize"); }; @@ -185,12 +189,12 @@ async fn session_affinity_initializer_cancellation_wakes_waiter() { drop(first); let next = waiter.await.unwrap().unwrap(); - assert!(matches!(&next, AffinityAcquire::Initialize(_))); + assert!(matches!(&next, Hold::Initialize(_))); drop(next); assert_eq!(coordinator.entry_count(), 0); assert!(matches!( coordinator.acquire(&session_id(), None).await.unwrap(), - AffinityAcquire::Initialize(_) + Hold::Initialize(_) )); } @@ -198,7 +202,7 @@ async fn session_affinity_initializer_cancellation_wakes_waiter() { async fn session_affinity_wait_stops_when_request_is_cancelled() { let coordinator = coordinator(); let first = coordinator.acquire(&session_id(), None).await.unwrap(); - let AffinityAcquire::Initialize(first) = first else { + let Hold::Initialize(first) = first else { panic!("first request must initialize"); }; @@ -227,7 +231,7 @@ async fn session_affinity_wait_stops_when_request_is_cancelled() { #[tokio::test(start_paused = true)] async fn session_affinity_validates_worker_and_rank_contract() { let coordinator = coordinator(); - let AffinityAcquire::Initialize(initializer) = coordinator + let Hold::Initialize(initializer) = coordinator .acquire(&session_id(), Some(target(7, None))) .await .unwrap() @@ -259,30 +263,28 @@ async fn session_affinity_validates_worker_and_rank_contract() { #[tokio::test(start_paused = true)] async fn session_affinity_failed_bound_operation_invalidates_binding() { let coordinator = coordinator(); - let AffinityAcquire::Initialize(initializer) = - coordinator.acquire(&session_id(), None).await.unwrap() + let Hold::Initialize(initializer) = coordinator.acquire(&session_id(), None).await.unwrap() else { panic!("first request must initialize"); }; drop(initializer.commit(to_table(target(7, Some(0)))).unwrap()); let operation = coordinator.acquire(&session_id(), None).await.unwrap(); - assert_eq!(operation.target(), Some(target(7, Some(0)))); + assert_eq!(operation.target(), Some(to_table(target(7, Some(0))))); operation.invalidate(); assert_eq!(coordinator.query_target(&session_id(), None).unwrap(), None); assert_eq!(coordinator.entry_count(), 0); assert!(matches!( coordinator.acquire(&session_id(), None).await.unwrap(), - AffinityAcquire::Initialize(_) + Hold::Initialize(_) )); } #[tokio::test(start_paused = true)] async fn session_affinity_stream_drop_refreshes_idle_ttl() { let coordinator = coordinator(); - let AffinityAcquire::Initialize(initializer) = - coordinator.acquire(&session_id(), None).await.unwrap() + let Hold::Initialize(initializer) = coordinator.acquire(&session_id(), None).await.unwrap() else { panic!("first request must initialize"); }; @@ -298,8 +300,7 @@ async fn session_affinity_stream_drop_refreshes_idle_ttl() { #[tokio::test(start_paused = true)] async fn session_affinity_empty_stream_refreshes_idle_ttl() { let coordinator = coordinator(); - let AffinityAcquire::Initialize(initializer) = - coordinator.acquire(&session_id(), None).await.unwrap() + let Hold::Initialize(initializer) = coordinator.acquire(&session_id(), None).await.unwrap() else { panic!("first request must initialize"); }; @@ -314,22 +315,21 @@ async fn session_affinity_empty_stream_refreshes_idle_ttl() { #[tokio::test(start_paused = true)] async fn session_affinity_cancelled_stream_refreshes_idle_ttl() { let coordinator = coordinator(); - let AffinityAcquire::Initialize(initializer) = - coordinator.acquire(&session_id(), None).await.unwrap() + let Hold::Initialize(initializer) = coordinator.acquire(&session_id(), None).await.unwrap() else { panic!("first request must initialize"); }; drop(initializer.commit(to_table(target(7, Some(0)))).unwrap()); tokio::time::advance(Duration::from_secs(9)).await; - let AffinityAcquire::Bound { + let Hold::Bound { target: bound_target, lease, } = coordinator.acquire(&session_id(), None).await.unwrap() else { panic!("continuation must acquire the existing binding"); }; - assert_eq!(bound_target, target(7, Some(0))); + assert_eq!(bound_target, to_table(target(7, Some(0)))); let mut stream = tracked_stream(lease, cancelled_response_stream()); assert!(stream.next().await.is_none()); @@ -340,8 +340,8 @@ async fn session_affinity_cancelled_stream_refreshes_idle_ttl() { async fn session_affinity_committed_binding_survives_cancelled_stream_until_ttl() { let coordinator = coordinator(); let operation = coordinator.acquire(&session_id(), None).await.unwrap(); - let mut stream = operation - .into_stream(target(7, Some(0)), cancelled_response_stream(), Hard) + let mut stream = coordinator + .commit_to_stream(operation, target(7, Some(0)), cancelled_response_stream()) .unwrap(); tokio::time::advance(Duration::from_secs(9)).await; assert!(stream.next().await.is_none()); @@ -352,8 +352,7 @@ async fn session_affinity_committed_binding_survives_cancelled_stream_until_ttl( #[tokio::test(start_paused = true)] async fn session_affinity_error_stream_refreshes_idle_ttl() { let coordinator = coordinator(); - let AffinityAcquire::Initialize(initializer) = - coordinator.acquire(&session_id(), None).await.unwrap() + let Hold::Initialize(initializer) = coordinator.acquire(&session_id(), None).await.unwrap() else { panic!("first request must initialize"); }; @@ -369,8 +368,7 @@ async fn session_affinity_error_stream_refreshes_idle_ttl() { #[tokio::test(start_paused = true)] async fn session_affinity_stream_eof_refreshes_idle_ttl() { let coordinator = coordinator(); - let AffinityAcquire::Initialize(initializer) = - coordinator.acquire(&session_id(), None).await.unwrap() + let Hold::Initialize(initializer) = coordinator.acquire(&session_id(), None).await.unwrap() else { panic!("first request must initialize"); }; @@ -385,17 +383,14 @@ async fn session_affinity_stream_eof_refreshes_idle_ttl() { #[tokio::test(start_paused = true)] async fn session_affinity_bound_lease_drop_refreshes_idle_ttl() { let coordinator = coordinator(); - let AffinityAcquire::Initialize(initializer) = - coordinator.acquire(&session_id(), None).await.unwrap() + let Hold::Initialize(initializer) = coordinator.acquire(&session_id(), None).await.unwrap() else { panic!("first request must initialize"); }; drop(initializer.commit(to_table(target(7, Some(0)))).unwrap()); tokio::time::advance(Duration::from_secs(9)).await; - let AffinityAcquire::Bound { lease, .. } = - coordinator.acquire(&session_id(), None).await.unwrap() - else { + let Hold::Bound { lease, .. } = coordinator.acquire(&session_id(), None).await.unwrap() else { panic!("continuation must acquire the binding"); }; tokio::time::advance(Duration::from_secs(2)).await; @@ -421,8 +416,7 @@ async fn session_affinity_query_is_read_only() { drop(initializing); assert_eq!(coordinator.entry_count(), 0); - let AffinityAcquire::Initialize(initializer) = - coordinator.acquire(&session_id(), None).await.unwrap() + let Hold::Initialize(initializer) = coordinator.acquire(&session_id(), None).await.unwrap() else { panic!("first request must initialize"); }; @@ -440,8 +434,7 @@ async fn session_affinity_query_is_read_only() { async fn session_affinity_reaper_removes_idle_entries_and_stops_on_drop() { let coordinator = coordinator(); let cancellation = coordinator.cancellation_token(); - let AffinityAcquire::Initialize(initializer) = - coordinator.acquire(&session_id(), None).await.unwrap() + let Hold::Initialize(initializer) = coordinator.acquire(&session_id(), None).await.unwrap() else { panic!("first request must initialize"); }; @@ -462,7 +455,7 @@ fn session_affinity_rejects_invalid_ttl_before_starting_reaper() { Duration::ZERO, Duration::from_secs(super::MAX_SESSION_AFFINITY_TTL_SECS + 1), ] { - let Err(error) = AffinityCoordinator::new(ttl) else { + let Err(error) = AffinityCoordinator::new(ttl, Hard) else { panic!("invalid TTL must fail coordinator construction"); }; assert!(dynamo_runtime::error::match_error_chain( @@ -504,7 +497,7 @@ async fn session_affinity_enforces_id_and_entry_limits() { assert_eq!(coordinator.entry_count(), 0); assert!(matches!( coordinator.acquire(&second_id, None).await.unwrap(), - AffinityAcquire::Initialize(_) + Hold::Initialize(_) )); } @@ -607,8 +600,8 @@ async fn session_affinity_publishes_after_dispatch_and_lease_completion() { let mut updates = coordinator.enable_test_replica(99, 4); let selected_target = target(7, Some(0)); let operation = coordinator.acquire(&session_id(), None).await.unwrap(); - let stream = operation - .into_stream(selected_target, response_stream(1), Hard) + let stream = coordinator + .commit_to_stream(operation, selected_target, response_stream(1)) .unwrap(); let after_dispatch = updates.recv().await.unwrap(); @@ -638,9 +631,7 @@ async fn session_affinity_republishes_the_stored_replica_version() { ReplicaApplyOutcome::Inserted ); - let AffinityAcquire::Bound { lease, .. } = - coordinator.acquire(&session_id(), None).await.unwrap() - else { + let Hold::Bound { lease, .. } = coordinator.acquire(&session_id(), None).await.unwrap() else { panic!("replicated binding must be acquired"); }; drop(lease); @@ -657,14 +648,14 @@ async fn session_affinity_worker_only_binding_allows_ranked_dispatch_without_nar let initialization = coordinator.acquire(&session_id(), None).await.unwrap(); drop( - initialization - .into_stream(worker_binding, response_stream(1), Hard) + coordinator + .commit_to_stream(initialization, worker_binding, response_stream(1)) .unwrap(), ); let continuation = coordinator.acquire(&session_id(), None).await.unwrap(); - let stream = continuation - .into_stream(target(7, Some(3)), response_stream(1), Hard) + let stream = coordinator + .commit_to_stream(continuation, target(7, Some(3)), response_stream(1)) .expect("worker-only affinity must allow the scheduler to select a DP rank"); assert_eq!( @@ -676,13 +667,13 @@ async fn session_affinity_worker_only_binding_allows_ranked_dispatch_without_nar #[tokio::test(start_paused = true)] async fn soft_worker_only_binding_stays_worker_scoped_after_ranked_dispatch() { - let coordinator = coordinator(); + let coordinator = soft_coordinator(); let worker_binding = target(7, None); bind(&coordinator, worker_binding).await; let continuation = coordinator.acquire(&session_id(), None).await.unwrap(); - let mut stream = continuation - .into_stream(target(7, Some(3)), response_stream(1), Soft) + let mut stream = coordinator + .commit_to_stream(continuation, target(7, Some(3)), response_stream(1)) .unwrap(); while stream.next().await.is_some() {} @@ -694,13 +685,13 @@ async fn soft_worker_only_binding_stays_worker_scoped_after_ranked_dispatch() { #[tokio::test(start_paused = true)] async fn soft_ranked_binding_is_not_widened_by_worker_only_dispatch() { - let coordinator = coordinator(); + let coordinator = soft_coordinator(); let ranked_binding = target(7, Some(2)); bind(&coordinator, ranked_binding).await; let continuation = coordinator.acquire(&session_id(), None).await.unwrap(); - let mut stream = continuation - .into_stream(target(8, None), response_stream(1), Soft) + let mut stream = coordinator + .commit_to_stream(continuation, target(8, None), response_stream(1)) .unwrap(); while stream.next().await.is_some() {} @@ -717,15 +708,15 @@ async fn session_affinity_ranked_binding_rejects_mismatched_rank_dispatch() { let initialization = coordinator.acquire(&session_id(), None).await.unwrap(); drop( - initialization - .into_stream(binding, response_stream(1), Hard) + coordinator + .commit_to_stream(initialization, binding, response_stream(1)) .unwrap(), ); let continuation = coordinator.acquire(&session_id(), None).await.unwrap(); assert!( - continuation - .into_stream(target(7, Some(3)), response_stream(1), Hard) + coordinator + .commit_to_stream(continuation, target(7, Some(3)), response_stream(1)) .is_err() ); assert_eq!(coordinator.query_target(&session_id(), None).unwrap(), None); @@ -733,7 +724,7 @@ async fn session_affinity_ranked_binding_rejects_mismatched_rank_dispatch() { #[tokio::test(start_paused = true)] async fn soft_affinity_rebinds_after_dispatch() { - let coordinator = coordinator(); + let coordinator = soft_coordinator(); let original = target(7, Some(0)); let replacement = target(8, Some(1)); bind(&coordinator, original).await; @@ -746,8 +737,8 @@ async fn soft_affinity_rebinds_after_dispatch() { ); let successful_attempt = coordinator.acquire(&session_id(), None).await.unwrap(); - let stream = successful_attempt - .into_stream(replacement, response_stream(1), Soft) + let stream = coordinator + .commit_to_stream(successful_attempt, replacement, response_stream(1)) .unwrap(); assert_eq!( coordinator.query_target(&session_id(), None).unwrap(), @@ -758,7 +749,7 @@ async fn soft_affinity_rebinds_after_dispatch() { #[tokio::test(start_paused = true)] async fn concurrent_soft_rebind_uses_observed_version_cas() { - let coordinator = coordinator(); + let coordinator = soft_coordinator(); let original = target(7, Some(0)); let winner = target(8, Some(0)); let stale = target(9, Some(0)); @@ -766,8 +757,12 @@ async fn concurrent_soft_rebind_uses_observed_version_cas() { let first = coordinator.acquire(&session_id(), None).await.unwrap(); let second = coordinator.acquire(&session_id(), None).await.unwrap(); - let mut first_stream = first.into_stream(winner, response_stream(1), Soft).unwrap(); - let mut second_stream = second.into_stream(stale, response_stream(1), Soft).unwrap(); + let mut first_stream = coordinator + .commit_to_stream(first, winner, response_stream(1)) + .unwrap(); + let mut second_stream = coordinator + .commit_to_stream(second, stale, response_stream(1)) + .unwrap(); while first_stream.next().await.is_some() {} while second_stream.next().await.is_some() {} assert_eq!( @@ -844,8 +839,8 @@ async fn session_affinity_completion_restores_expired_remote_binding() { let replica = coordinator(); let replicated_target = target(7, Some(0)); let operation = origin.acquire(&session_id(), None).await.unwrap(); - let stream = operation - .into_stream(replicated_target, response_stream(1), Hard) + let stream = origin + .commit_to_stream(operation, replicated_target, response_stream(1)) .unwrap(); let after_dispatch = updates.recv().await.unwrap(); From fabac4b3228921115bbacd7e1650d6cd4110235f Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Thu, 10 Sep 2026 17:16:20 -0700 Subject: [PATCH 058/102] refactor(router): split the selection core into one module per concern core/mod.rs held six responsibilities in 2,347 non-test lines. Pure moves, no logic change: mod.rs keeps SelectionCore, the host types, partition and entry construction, readiness and replica dispatch; run.rs holds one selection end to end from the wire adapters through run_selection and the lookups to the scheduler call and the post-scheduler transaction; reservations.rs holds the selection-id index, the replay path and the lifecycle calls; workers.rs holds catalog membership; queries.rs the read-only load and overlap queries; hint.rs the transfer hint; tests.rs the former inline test module. Children inherit the parent's imports; items used across siblings are pub(super). Co-Authored-By: Claude Fable 5.1 Signed-off-by: Ishan Dhanani --- .../src/services/selection/core/hint.rs | 70 + .../src/services/selection/core/mod.rs | 4686 +---------------- .../src/services/selection/core/queries.rs | 113 + .../services/selection/core/reservations.rs | 553 ++ .../src/services/selection/core/run.rs | 740 +++ .../src/services/selection/core/tests.rs | 2803 ++++++++++ .../src/services/selection/core/workers.rs | 365 ++ 7 files changed, 4687 insertions(+), 4643 deletions(-) create mode 100644 lib/kv-router/src/services/selection/core/hint.rs create mode 100644 lib/kv-router/src/services/selection/core/queries.rs create mode 100644 lib/kv-router/src/services/selection/core/reservations.rs create mode 100644 lib/kv-router/src/services/selection/core/run.rs create mode 100644 lib/kv-router/src/services/selection/core/tests.rs create mode 100644 lib/kv-router/src/services/selection/core/workers.rs diff --git a/lib/kv-router/src/services/selection/core/hint.rs b/lib/kv-router/src/services/selection/core/hint.rs new file mode 100644 index 000000000000..a5e7f0779211 --- /dev/null +++ b/lib/kv-router/src/services/selection/core/hint.rs @@ -0,0 +1,70 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! The KV transfer hint attached to a booked selection. + +use super::*; + +/// Pick the best router-hint source for `target`: a same-role worker (or +/// cache owner) holding a longer root-aligned prefix than the target's own +/// `target_cached_prefix_blocks`, with a non-empty control endpoint. Mirrors the +/// frontend `KvRouter::transfer_hint_for_selection`. +pub(super) fn transfer_hint_for_selection( + configs: &HashMap, + target: WorkerWithDpRank, + target_cached_prefix_blocks: u32, + candidates: Option<&KvTransferCandidates>, +) -> Option { + let candidates = candidates?; + let target_config = configs.get(&target.worker_id)?; + let target_metadata = target_config.kv_hint_transfer_metadata_for_dp_rank(target.dp_rank)?; + + let prefix_blocks_to_beat = usize::try_from(target_cached_prefix_blocks).unwrap_or(usize::MAX); + let (source, block_hashes) = + candidates.best_source(prefix_blocks_to_beat, |source| match source { + KvTransferCandidateSource::Worker(worker) => { + worker != target + && configs.get(&worker.worker_id).is_some_and(|config| { + config.kv_event_source_mode.as_deref() != Some("state_agent_v2") + && config + .kv_hint_transfer_metadata_for_dp_rank(worker.dp_rank) + .is_some_and(|source_metadata| { + source_metadata.worker_type == target_metadata.worker_type + && source_metadata + .source_control_endpoint + .is_some_and(|endpoint| !endpoint.is_empty()) + }) + }) + } + KvTransferCandidateSource::CacheOwner(owner) => candidates + .routing_snapshot + .as_ref() + .and_then(|snapshot| snapshot.router_hint_source(owner)) + .is_some_and(|source| { + source.attached_worker != Some(target) + && source.metadata.worker_type == target_metadata.worker_type + && !source.metadata.source_control_endpoint.is_empty() + }), + })?; + let source_control_endpoint = match source { + KvTransferCandidateSource::Worker(worker) => configs + .get(&worker.worker_id)? + .kv_hint_transfer_metadata_for_dp_rank(worker.dp_rank)? + .source_control_endpoint? + .to_string(), + KvTransferCandidateSource::CacheOwner(owner) => candidates + .routing_snapshot + .as_ref()? + .router_hint_source(owner)? + .metadata + .source_control_endpoint + .clone(), + }; + if block_hashes.is_empty() { + return None; + } + Some(KvSourceLocationsPayload { + source_control_endpoint, + block_hashes, + }) +} diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index 19915e33167c..f287c0859636 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -1,6 +1,10 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +//! The shared selection core: one `SelectionCore` per process holding the +//! partitions, their catalog, and the reservation index. Selection itself is +//! in `run`, reservations in `reservations`, worker membership in `workers`. + use std::collections::HashMap; use std::sync::Arc; use std::time::{Duration, Instant}; @@ -19,8 +23,6 @@ use crate::indexer::{ use crate::kv_hints::{ KvHint, KvHintAction, KvSourceLocationsPayload, KvTransferCandidateSource, KvTransferCandidates, }; -#[cfg(test)] -use crate::protocols::ActiveSequenceEventData; use crate::protocols::{ ActiveSequenceEvent, LocalBlockHash, PrefillLoadHint, SharedCacheHits, WorkerAffinityTarget, WorkerConfigLike, WorkerId, WorkerWithDpRank, @@ -37,8 +39,6 @@ use crate::sequences::{ ActiveSequencesMultiWorker, LifecycleMutationOutcome, ReplicaRequestLeaseObserver, ReplicaWorkerPolicy, SequenceRequest, SequenceTrackerOptions, active_request_expiry_duration, }; -#[cfg(test)] -use crate::services::common::replica_sync::HostReplicaChannels; use crate::services::common::replica_sync::{ HostReplicaSyncFactory, ReplicaSyncConfig, SchedulerLoadSink, ScopedReplicaEvent, ScopedSequencePublisher, setup_scoped_replica_sync, @@ -49,7 +49,16 @@ use crate::services::indexer::registry::WorkerRegistry; use crate::services::overlap::MooncakeOverlapSummary; use crate::tracking_hash::{TrackingHashContext, TrackingHashScope}; +mod hint; mod operation; +mod queries; +mod reservations; +mod run; +#[cfg(test)] +mod tests; +mod workers; + +use reservations::ReservationIndex; pub use operation::{ LookupTimings, Selected, SelectionAdmission, SelectionOperation, SelectionOutcome, @@ -125,15 +134,6 @@ struct SelectionEntry { replica_config: Option, } -struct PreparedSelectionInputs { - block_hashes: Vec, - sequence_hashes: Vec, - isl_tokens: usize, - overlap: OverlapSignals, - shared_cache_hits: Option, - kv_transfer_candidates: Option, -} - impl SelectionEntry { fn session_affinity( &self, @@ -164,22 +164,6 @@ impl SelectionEntry { } } -/// Resolved inputs for booking a reservation, shared by the cached and explicit -/// `create_reservation` paths. -struct ReservationBooking { - key: RoutingPartitionId, - selection_id: String, - worker: WorkerWithDpRank, - sequence_hashes: Vec, - prefill_load_hint: Option, - expected_output_tokens: Option, - track_prefill_tokens: bool, - lora_name: Option, - /// Public block hashes to record into an approximate indexer once booked. - routing_hashes: Option>, - session_id: Option, -} - /// What an embedding host supplies to every partition the core creates, /// grouped by purpose. Each group defaults to the standalone service's /// behavior, so a host overrides only the groups it owns: the frontend @@ -308,120 +292,6 @@ pub struct SelectionServiceConfig { type SelectionEntries = RwLock>>>>; -/// `selection_id` -> the booking it holds. -/// -/// Lifecycle calls (`prefill_complete`, `free`, `add_output_block`) arrive with -/// only a selection id; the index resolves them to one partition and one -/// scheduler booking, so they never touch a booking made by a later request -/// that reused the id. An id is claimed here before its booking is made and -/// installed once the booking is final; a claim (`booking == None`) rejects a -/// concurrent booking of the same id and is invisible to lifecycle calls. -/// Bookings mirrored from replica peers are indexed by -/// [`ReservationIndexObserver`]. -type ReservationIndex = RwLock>; - -struct Reservation { - partition: RoutingPartitionId, - booking: Option, - _affinity_lease: Option, -} - -/// Exclusive ownership of a selection id while its booking is in flight. -/// Dropping the claim without `install` releases the id. -struct ReservationClaim<'a> { - index: &'a ReservationIndex, - selection_id: String, - armed: bool, -} - -impl ReservationClaim<'_> { - fn install(mut self, reservation: Reservation) { - self.armed = false; - self.index - .write() - .insert(std::mem::take(&mut self.selection_id), reservation); - } -} - -impl Drop for ReservationClaim<'_> { - fn drop(&mut self) { - if !self.armed { - return; - } - let mut index = self.index.write(); - if index - .get(&self.selection_id) - .is_some_and(|reservation| reservation.booking.is_none()) - { - index.remove(&self.selection_id); - } - } -} - -/// Keeps the reservation index exact for bookings replicated from peers, which -/// never pass through this core's booking paths, and forwards to the host's -/// own observer. -struct ReservationIndexObserver { - index: Arc, - partition: RoutingPartitionId, - host: Option>, -} - -impl ReplicaRequestLeaseObserver for ReservationIndexObserver { - fn admitted(&self, booking: SchedulerBookingDescriptor) { - { - let mut index = self.index.write(); - // A peer only admits an id this partition's scheduler does not hold, - // so an existing row for it is a stale booking (expired, not yet - // swept) or a claim whose local booking will now fail; the mirror - // replaces both. A row from another partition is left alone. - let own_row = index - .get(&booking.request_id) - .is_none_or(|reservation| reservation.partition == self.partition); - if own_row { - index.insert( - booking.request_id.clone(), - Reservation { - partition: self.partition.clone(), - booking: Some(booking.clone()), - _affinity_lease: None, - }, - ); - } - } - if let Some(host) = &self.host { - host.admitted(booking); - } - } - - fn progressed(&self, booking: &SchedulerBookingDescriptor) { - if let Some(host) = &self.host { - host.progressed(booking); - } - } - - fn completed(&self, booking: &SchedulerBookingDescriptor) { - forget_reservation_if(&self.index, &self.partition, booking); - if let Some(host) = &self.host { - host.completed(booking); - } - } -} - -/// Remove the index entry for `booking` only if it still describes it. -fn forget_reservation_if( - index: &ReservationIndex, - partition: &RoutingPartitionId, - booking: &SchedulerBookingDescriptor, -) { - let mut index = index.write(); - if index.get(&booking.request_id).is_some_and(|reservation| { - reservation.partition == *partition && reservation.booking.as_ref() == Some(booking) - }) { - index.remove(&booking.request_id); - } -} - pub struct SelectionCore { catalog: WorkerCatalog, /// Serializes catalog commits and the corresponding ingress changes. Never held by selection. @@ -628,356 +498,36 @@ impl SelectionCore { } } - pub async fn upsert_worker( - &self, - req: WorkerRequest, - ) -> Result { - self.ensure_running()?; - let mut record = WorkerCatalogRecord::new(req); - // Partition policy factories may construct independently. Only committing - // membership and reconciling ingress needs the catalog mutation lock. - self.prepare_worker(&mut record)?; - let _update = self.catalog_updates.lock().await; - self.ensure_running()?; - let previous = self.catalog.get(record.worker_id); - self.reconcile_worker(record, previous).await - } - - pub async fn patch_worker( - &self, - worker_id: WorkerId, - patch: WorkerPatchRequest, - ) -> Result { - let _update = self.catalog_updates.lock().await; - self.ensure_running()?; - let previous = self - .catalog - .get(worker_id) - .ok_or_else(|| SelectionError::NotFound(format!("worker {worker_id} not found")))?; - let mut record = previous.clone(); - record.apply_patch(patch); - self.prepare_worker(&mut record)?; - self.reconcile_worker(record, Some(previous)).await - } - - pub async fn delete_worker( - &self, - worker_id: WorkerId, - ) -> Result { - let _update = self.catalog_updates.lock().await; - let Some(previous) = self.catalog.get(worker_id) else { - return Err(SelectionError::NotFound(format!( - "worker {worker_id} not found" - ))); + /// Apply a session binding a replica published. + pub(crate) fn dispatch_affinity_event(&self, event: AffinityBindingEvent) { + let Some(entry) = self.entry(&event.partition) else { + return; }; - let key = previous.key(); - self.catalog - .set_lifecycle(worker_id, WorkerLifecycle::Draining, Vec::new()); - self.publish_scheduler_config(&key)?; - self.cleanup_indexer_registration(&previous).await; - let record = self - .catalog - .set_lifecycle(worker_id, WorkerLifecycle::Unschedulable, Vec::new()) - .ok_or_else(|| SelectionError::NotFound(format!("worker {worker_id} not found")))?; - self.publish_scheduler_config(&key)?; - Ok(record) - } - - pub fn list_workers( - &self, - model_name: Option<&str>, - routing_group: Option<&str>, - ) -> Vec { - self.catalog.list(model_name, routing_group) - } - - pub fn ready(&self) -> ReadyResponse { - let schedulable_workers = self.catalog.schedulable_count(); - let workers = self.catalog.list(None, None); - ReadyResponse { - ready: !self.cancel_token.is_cancelled() && schedulable_workers > 0, - schedulable_workers, - workers, - } - } - - fn prepare_worker(&self, record: &mut WorkerCatalogRecord) -> Result<(), SelectionError> { - let queueing_enabled = self - .kv_router_config - .queueing_enabled(Some(&record.model_name)) - .map_err(|error| SelectionError::BadRequest(error.to_string()))?; - record.not_schedulable_reasons = record.missing_schedulable_metadata(queueing_enabled); - if let Some(ingress) = self.ingress() { - record - .not_schedulable_reasons - .extend(ingress.missing_metadata(record)); - } - if record.not_schedulable_reasons.is_empty() - && let Err(error) = self.ensure_entry(record) - { - record - .not_schedulable_reasons - .push(format!("reconciliation failed: {error}")); - } - - Ok(()) - } - - async fn reconcile_worker( - &self, - mut record: WorkerCatalogRecord, - previous: Option, - ) -> Result { - let previous = previous.filter(|old| old.lifecycle == WorkerLifecycle::Schedulable); - let previous = if let Some(old) = previous.as_ref() - && (old.key() != record.key() || !record.not_schedulable_reasons.is_empty()) - { - self.catalog - .set_lifecycle(old.worker_id, WorkerLifecycle::Draining, Vec::new()); - self.publish_scheduler_config(&old.key())?; - self.cleanup_indexer_registration(old).await; - None - } else { - previous + let Some(table) = entry.affinity.get() else { + return; }; - - if record.not_schedulable_reasons.is_empty() - && let Some(ingress) = self.ingress() - && let Err(error) = ingress - .reconcile(&self.indexer_registry, previous.as_ref(), &record) - .await + if self + .replica_config + .as_ref() + .is_some_and(|config| config.process_id() == event.writer_id) { - self.cleanup_indexer_registration(&record).await; - record - .not_schedulable_reasons - .push(format!("reconciliation failed: {error}")); - } - record.lifecycle = if record.not_schedulable_reasons.is_empty() { - WorkerLifecycle::Schedulable - } else { - WorkerLifecycle::Incomplete - }; - // Readers see only committed metadata. A valid capacity/topology update preserves - // live bookings on ranks present in both the old and new snapshots. - self.catalog.replace(record.clone()); - self.publish_scheduler_config(&record.key())?; - Ok(record) - } - - fn ensure_entry( - &self, - record: &WorkerCatalogRecord, - ) -> Result, SelectionError> { - let block_size = record - .block_size - .ok_or_else(|| SelectionError::BadRequest("block_size is required".to_string()))?; - self.ensure_entry_for(record.key(), block_size, record.is_eagle.unwrap_or(false)) - } - - /// Create the partition scheduler and indexer for `key` before any worker - /// registers, so an embedding host can hold the scheduler handle from - /// construction. Idempotent; a later worker with a different block size or - /// eagle setting is rejected at reconciliation. - pub fn ensure_partition( - &self, - key: RoutingPartitionId, - block_size: u32, - is_eagle: bool, - ) -> Result { - self.ensure_running()?; - if block_size == 0 { - return Err(SelectionError::BadRequest( - "block_size must be greater than 0".to_string(), - )); - } - self.ensure_entry_for(key, block_size, is_eagle) - .map(SelectionPartition) - } - - fn ensure_entry_for( - &self, - key: RoutingPartitionId, - block_size: u32, - is_eagle: bool, - ) -> Result, SelectionError> { - self.reservation_sweep_started.get_or_init(|| { - spawn_reservation_index_sweep( - Arc::clone(&self.entries), - Arc::clone(&self.reservation_index), - self.cancel_token.child_token(), - ); - }); - - let entry_cell = { self.entries.read().get(&key).cloned() }; - let entry_cell = entry_cell.unwrap_or_else(|| { - self.entries - .write() - .entry(key.clone()) - .or_insert_with(|| Arc::new(OnceCell::new())) - .clone() - }); - let entry = entry_cell - .get_or_try_init(|| -> Result, SelectionError> { - let (workers_tx, workers_rx) = watch::channel(HashMap::new()); - let host_replica = self - .host - .replication - .channels - .as_ref() - .and_then(|factory| factory(&key)); - let scoped_replica_sync = setup_scoped_replica_sync( - self.replica_config.as_ref(), - &key, - block_size, - host_replica, - ); - let worker_label = self.worker_type.as_str(); - let slots = Arc::new(ActiveSequencesMultiWorker::new_with_options( - scoped_replica_sync - .publisher - .with_load_sink(self.host.telemetry.scheduler_load.clone()), - block_size as usize, - HashMap::new(), - scoped_replica_sync.enabled, - scoped_replica_sync.process_id, - worker_label, - SequenceTrackerOptions { - replica_worker_policy: ReplicaWorkerPolicy::RequireRegistered, - expiry_duration: self - .host - .replication - .request_leases - .is_none() - .then(active_request_expiry_duration), - }, - )); - slots.set_replica_request_lease_observer(Arc::new(ReservationIndexObserver { - index: Arc::clone(&self.reservation_index), - partition: key.clone(), - host: self.host.replication.request_leases.clone(), - })); - let replica_tx = scoped_replica_sync.channel.map(|(replica_tx, subscriber)| { - slots.start_replica_sync(subscriber, self.cancel_token.child_token()); - replica_tx - }); - if self.host.replication.request_leases.is_none() { - slots.start_periodic_force_expiry_across_all_workers( - self.cancel_token.child_token(), - ); - } - - let indexer = match &self.host.cache.index { - KvIndexSource::Owned(ingress) => { - ingress.open(&self.indexer_registry, &key, block_size) - } - KvIndexSource::Remote(_) => self - .indexer_registry - .get_or_create_indexer(key.clone(), block_size), - }; - let overlap_refresh = indexer.supports_overlap_refresh().then(|| { - Arc::new(TieredOverlapRefresher::new( - indexer.clone(), - self.kv_router_config.clone(), - block_size, - )) - }); - let selector = self.worker_selection_policy_factory.as_ref().map_or_else( - || WorkerSelectionPolicy::default(self.kv_router_config.clone(), worker_label), - |factory| factory(&self.kv_router_config, self.worker_type, key.as_ref()), - ); - let profile = self - .kv_router_config - .policy_profile(Some(&key.model_name)) - .map_err(|error| SelectionError::BadRequest(error.to_string()))?; - let scheduler = LocalScheduler::new( - slots, - workers_rx, - profile, - block_size, - selector, - self.host.load.prefill_estimator.clone(), - overlap_refresh, - // Standalone selection has no router Client snapshot, so - // these stay `None` unless an embedding host injects them. - self.host.load.overloaded_workers.clone(), - self.host.load.available_workers.clone(), - self.kv_router_config.router_queue_recheck_interval(), - self.kv_router_config.router_track_prefill_tokens, - self.cancel_token.child_token(), - worker_label, - true, - ); - Ok(Arc::new(SelectionEntry { - key: key.clone(), - block_size, - is_eagle, - indexer, - workers_tx, - scheduler, - replica_tx, - affinity: OnceCell::new(), - replica_config: self.replica_config.clone(), - })) - })? - .clone(); - if entry.block_size != block_size { - return Err(SelectionError::Conflict(format!( - "block_size mismatch for {key}: existing={} requested={block_size}", - entry.block_size - ))); - } - if entry.is_eagle != is_eagle { - return Err(SelectionError::Conflict(format!( - "is_eagle mismatch for {key}: existing={} requested={is_eagle}", - entry.is_eagle - ))); - } - if let Some(config) = self.session_affinity { - entry.session_affinity(config)?; - } - Ok(entry) - } - - /// The ingress feeding core-owned indexes, when this core listens for KV events. - fn ingress(&self) -> Option<&dyn KvEventIngress> { - match &self.host.cache.index { - KvIndexSource::Owned(ingress) if self.listens_for_kv_events => Some(ingress.as_ref()), - _ => None, - } - } - - async fn cleanup_indexer_registration(&self, record: &WorkerCatalogRecord) { - if let KvIndexSource::Owned(ingress) = &self.host.cache.index { - ingress.detach(&self.indexer_registry, record).await; return; } - - let key = record.key(); - let indexer = self - .indexer_registry - .get_indexer(&key) - .map(|entry| entry.indexer.clone()); - if let Some(indexer) = indexer { - indexer.remove_worker(record.worker_id).await; + table.observe_replica_sequence(event.sequence); + if self + .catalog + .get(event.worker_id) + .is_none_or(|record| record.key() != event.partition) + { + return; } - } - - fn publish_scheduler_config(&self, key: &RoutingPartitionId) -> Result<(), SelectionError> { - let Some(entry) = self.entry(key) else { - return Ok(()); - }; - let workers = self.catalog.scheduler_configs_for_key(key); - // Lifecycle transitions between non-schedulable states publish the same - // map; skipping them saves the scheduler a wake and a full map clone. - entry.workers_tx.send_if_modified(|current| { - if *current == workers { - false - } else { - *current = workers; - true - } - }); - Ok(()) + let (target, version, worker_id) = (event.target(), event.version(), event.worker_id); + let outcome = table.apply_replica_update(event.session_id, target, version); + tracing::trace!( + worker_id, + ?outcome, + "applied session affinity replica update" + ); } fn ready_entry(&self, key: &RoutingPartitionId) -> Result, SelectionError> { @@ -998,4167 +548,17 @@ impl SelectionCore { SelectionError::NotReady(format!("no schedulable workers for {key}")) } } +} - pub async fn select(&self, req: SelectRequest) -> Result { - self.select_with_policy_class(req, None).await - } - - pub async fn select_with_policy_class( - &self, - mut req: SelectRequest, - policy_class: Option, - ) -> Result { - let session_context = req.take_session_context(); - let is_steerable = req.affinity_target.is_none() && req.pinned_worker.is_none(); - let session = session_binding(session_context.as_ref(), is_steerable, false); - let admission = if req.advisory { - SelectionAdmission::Advisory { - request_id: req.selection_id.clone(), - } - } else { - SelectionAdmission::Query { - request_id: req.selection_id.clone(), - } - }; - let (selected, endpoint) = self - .select_or_reject(SelectionOperation { - key: RoutingPartitionId::new(req.model_name, req.routing_group), - prompt: req.prompt.view(), - router_config_override: req.router_config_override, - expected_output_tokens: req.expected_output_tokens, - priority_jump: req.priority_jump.unwrap_or_default(), - strict_priority: req.strict_priority.unwrap_or(0), - policy_class, - session_context, - session, - affinity_target: req.affinity_target, - pinned_worker: req.pinned_worker, - allowed_worker_ids: req.allowed_worker_ids, - routing_constraints: req.routing_constraints, - admission, - track_active_blocks: true, - return_routing_hashes: false, - replay_id: req.selection_id.clone(), - }) - .await?; - Ok(self.select_response(selected, endpoint, req.selection_id)) - } - - pub async fn select_and_reserve( - &self, - req: SelectAndReserveRequest, - ) -> Result { - self.select_and_reserve_with_policy_class(req, None).await - } - - pub async fn select_and_reserve_with_policy_class( - &self, - mut req: SelectAndReserveRequest, - policy_class: Option, - ) -> Result { - let session_context = req.take_session_context(); - let selection_id = req - .selection_id - .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); - let is_steerable = req.affinity_target.is_none() && req.pinned_worker.is_none(); - let session = session_binding(session_context.as_ref(), is_steerable, true); - let (selected, endpoint) = self - .select_or_reject(SelectionOperation { - key: RoutingPartitionId::new(req.model_name, req.routing_group), - prompt: req.prompt.view(), - router_config_override: req.router_config_override, - expected_output_tokens: req.expected_output_tokens, - priority_jump: req.priority_jump.unwrap_or_default(), - strict_priority: req.strict_priority.unwrap_or(0), - policy_class, - session_context, - session, - affinity_target: req.affinity_target, - pinned_worker: req.pinned_worker, - allowed_worker_ids: req.allowed_worker_ids, - routing_constraints: req.routing_constraints, - admission: SelectionAdmission::Book { - selection_id: selection_id.clone(), - }, - track_active_blocks: true, - return_routing_hashes: false, - replay_id: None, - }) - .await?; - Ok(self.select_response(selected, endpoint, Some(selection_id))) - } - - /// The wire shape of a selection; a queue rejection keeps its 503 body. - async fn select_or_reject( - &self, - operation: SelectionOperation<'_>, - ) -> Result<(Selected, String), SelectionError> { - let mut selected = match self.run_selection(operation).await.result? { - SelectionOutcome::Selected(selected) => selected, - SelectionOutcome::QueueRejected { rejection } => { - return Err(SelectionError::Scheduler(KvSchedulerError::QueueRejected( - rejection, - ))); - } - }; - // Wire admissions never book without an endpoint; see `Selected::endpoint`. - let endpoint = selected.endpoint.take().ok_or_else(|| { - SelectionError::Internal(format!( - "selected worker {} is no longer schedulable", - selected.response.best_worker.worker_id - )) - })?; - Ok((selected, endpoint)) - } - - fn select_response( - &self, - selected: Selected, - endpoint: String, - selection_id: Option, - ) -> SelectResponse { - let Selected { - key, - response, - advisory_load, - total_kv_blocks, - endpoint: _, - block_size, - isl_tokens, - sequence_hashes, - track_prefill_tokens, - effective_prefill_tokens, - kv_hint, - routing_hashes: _, - shared_cache_hits: _, - lease: _, - } = selected; - let booked = sequence_hashes.is_some(); - let potential_decode_blocks = response.potential_decode_blocks as u64; - let decode_busy = self - .kv_router_config - .conditional_disagg_decode_busy_threshold - .zip(total_kv_blocks) - .map(|(threshold, total_kv_blocks)| { - potential_decode_blocks as f64 > threshold * total_kv_blocks as f64 - }); - let worker_load = advisory_load.map(|load| SelectionWorkerLoad { - active_prefill_tokens: load.active_prefill_tokens, - prefill_token_capacity: load.prefill_token_capacity, - total_kv_blocks, - prefill_busy: self - .kv_router_config - .conditional_disagg_prefill_busy_threshold - .map(|threshold| load.prefill_load_exceeds(threshold)), - }); - SelectResponse { - selection_id, - sequence_hashes: sequence_hashes - .map(|hashes| hashes.into_iter().map(|hash| hash as i64).collect()), - isl_tokens: booked.then_some(isl_tokens), - track_prefill_tokens: booked.then_some(track_prefill_tokens), - model_name: key.model_name, - routing_group: key.routing_group, - worker_id: response.best_worker.worker_id, - dp_rank: response.best_worker.dp_rank, - endpoint, - block_size, - overlap: MooncakeOverlapSummary::from_selected_worker_tiers( - &response.selected_worker_tiers, - block_size, - ), - effective_prefill_tokens, - potential_decode_blocks, - decode_busy, - worker_load, - kv_hint, - } - } - - /// Run one selection: resolve the session, look the prompt up, schedule, - /// and (for `Book`) install the reservation. Every host's selection goes - /// through here. - pub async fn run_selection(&self, operation: SelectionOperation<'_>) -> SelectionRun { - let mut lookup = None; - let result = self.run_selection_inner(operation, &mut lookup).await; - SelectionRun { result, lookup } - } - - async fn run_selection_inner( - &self, - operation: SelectionOperation<'_>, - lookup: &mut Option, - ) -> Result { - let SelectionOperation { - key, - prompt, - router_config_override, - expected_output_tokens, - priority_jump, - strict_priority, - policy_class, - session_context, - session, - affinity_target, - pinned_worker, - allowed_worker_ids, - routing_constraints, - admission, - track_active_blocks, - return_routing_hashes, - replay_id, - } = operation; - self.ensure_running()?; - - let entry = self.ready_entry(&key)?; - let book = admission.is_booking(); - let claim = match &admission { - SelectionAdmission::Book { selection_id } => { - Some(self.claim_reservation(selection_id, &key)?) - } - _ => None, - }; - - // Session stickiness: a bound session steers selection (exclusive for - // the default selector); a new session is bound to the worker booked. - let table = entry.affinity.get(); - let mut affinity_hold = None; - let managed_session = match (&session, table) { - (SessionBinding::Managed { .. }, _) if claim.is_none() => { - return Err(SelectionError::Internal( - "a managed session binding requires Book admission".to_string(), - )); - } - (SessionBinding::Managed { session_id }, Some(table)) => Some((table, session_id)), - _ => None, - }; - let affinity_target = match (&session, table) { - (SessionBinding::Managed { session_id }, Some(table)) => { - affinity_hold = self.hold_session(table, session_id, &key).await?; - affinity_hold.as_ref().and_then(Hold::target) - } - (SessionBinding::Query { session_id }, Some(table)) => table - .query_target(session_id, None) - .map_err(affinity_error)? - .map(|target| WorkerAffinityTarget::new(target.worker_id, target.dp_rank)), - _ => affinity_target, - }; - // Router hints are attached to bookings only, and only when a worker in - // this partition can consume them and the indexer can retain the - // matched chain (local, event-driven, no approximate writes). - let hint_capable_workers = book && self.catalog.has_router_hint_capable_workers(&key); - let retain_kv_transfer_chain = - hint_capable_workers && entry.indexer.supports_kv_transfer_chain_retention(); - if hint_capable_workers && !retain_kv_transfer_chain { - static WARN_ONCE: std::sync::Once = std::sync::Once::new(); - WARN_ONCE.call_once(|| { - tracing::warn!( - "router hints need a local event-driven indexer with no approximate side indexer; workers advertise hint capability but no hints will be attached" - ); - }); - } - let PreparedSelectionInputs { - block_hashes, - sequence_hashes, - isl_tokens, - overlap, - shared_cache_hits, - kv_transfer_candidates, - } = self - .prepare_selection_inputs( - &entry, - &prompt, - track_active_blocks.then(|| { - self.kv_router_config - .assume_kv_reuse(router_config_override.as_ref()) - }), - true, - retain_kv_transfer_chain, - lookup, - ) - .await?; - if tracing::enabled!(tracing::Level::DEBUG) { - let local_hashes: Vec = block_hashes.iter().map(|hash| hash.0).collect(); - tracing::debug!( - request_id = admission.request_id().unwrap_or_default(), - isl_tokens, - block_size = entry.block_size, - num_blocks = local_hashes.len(), - ?local_hashes, - "[ROUTING_INPUT] request local hashes" - ); - } - let host_shared_cache_hits = shared_cache_hits.clone(); - // The queue lease frees a booking whose response is never consumed - // (the caller dropped this future after the actor booked). - let mode = match &admission { - SelectionAdmission::Book { selection_id } - | SelectionAdmission::Lease { - request_id: selection_id, - } => ScheduleMode::TrackedWithLifecycle { - request_id: selection_id.clone(), - }, - SelectionAdmission::Query { request_id } => ScheduleMode::QueryOnly { - request_id: request_id.clone(), - }, - SelectionAdmission::Advisory { request_id } => ScheduleMode::QueryOnly { - request_id: request_id.clone(), - }, - }; - let track_prefill_tokens = router_config_override - .as_ref() - .and_then(|cfg| cfg.track_prefill_tokens) - .unwrap_or(self.kv_router_config.router_track_prefill_tokens); - // An unbooked selection with an id caches the booking inputs so a - // follow-up `create_reservation` can replay them by that id. - let cached_inputs = replay_id.filter(|_| !book).map(|id| { - ( - id, - sequence_hashes.clone(), - prompt.lora_name.map(str::to_string), - track_prefill_tokens, - match &session { - SessionBinding::Query { session_id } => Some(session_id.clone()), - _ => None, - }, - ) - }); - let allowed_worker_ids = match self.host.eligibility.lora_worker_filter.as_deref() { - Some(filter) => narrow_allowed_worker_ids_by_lora( - filter, - prompt.lora_name, - allowed_worker_ids, - pinned_worker.as_ref(), - || entry.workers_tx.borrow().keys().copied().collect(), - ), - None => allowed_worker_ids, - }; - // Bookings (now, or later via the pending-selection cache) are recorded - // into an approximate indexer; keep the public hashes for that. - let routing_hashes = (entry.indexer.records_routing_decisions() - && (claim.is_some() || cached_inputs.is_some())) - .then(|| block_hashes.clone()); - let booked_sequence_hashes = book.then(|| sequence_hashes.clone()); - let returned_routing_hashes = return_routing_hashes.then(|| block_hashes.clone()); - let schedule_request = ScheduleRequest { - mode, - token_seq: track_active_blocks.then_some(sequence_hashes), - block_hashes: Some(block_hashes), - isl_tokens, - overlap, - kv_transfer_candidates, - retain_kv_transfer_chain, - router_config_override, - lora_name: prompt.lora_name.map(str::to_string), - priority_jump, - strict_priority, - policy_class, - session_context, - expected_output_tokens, - affinity_target, - pinned_worker, - allowed_worker_ids, - routing_constraints, - shared_cache_hits, - }; - // `lease` guards the booking until it is installed below: any early - // return or drop before then frees it. - let scheduled = tokio::select! { - biased; - _ = self.cancel_token.cancelled() => { - return Err(SelectionError::Scheduler(KvSchedulerError::SubscriberShutdown)); - } - result = async { - if matches!(admission, SelectionAdmission::Advisory { .. }) { - entry - .scheduler - .select_without_admission(schedule_request) - .instrument(tracing::info_span!("kv_router.select_without_admission")) - .await - .map(|advisory| { - (advisory.response, Some(advisory.selected_worker_load), None) - }) - } else { - entry - .scheduler - .schedule_request_with_lease(schedule_request) - .instrument(tracing::info_span!("kv_router.schedule")) - .await - .map(|(admitted, lease)| (admitted.response, None, lease)) - } - } => result, - }; - let (response, advisory_load, lease) = match scheduled { - Ok(scheduled) => scheduled, - Err(KvSchedulerError::QueueRejected(rejection)) => { - return Ok(SelectionOutcome::QueueRejected { rejection }); - } - Err(error) => return Err(error.into()), - }; - let endpoint = self - .catalog - .schedulable_endpoint(response.best_worker.worker_id, &key); - if endpoint.is_none() && !matches!(admission, SelectionAdmission::Lease { .. }) { - return Err(SelectionError::Internal(format!( - "selected worker {} is no longer schedulable", - response.best_worker.worker_id - ))); - } - let effective_prefill = effective_prefill_tokens(isl_tokens, response.cached_tokens); - let total_kv_blocks = advisory_load - .and_then(|load| load.total_kv_blocks.map(|blocks| blocks as u64)) - .or_else(|| { - self.catalog - .total_kv_blocks(response.best_worker.worker_id, &key) - }); - let kv_hint = if retain_kv_transfer_chain && response.kv_transfer_candidates.is_some() { - transfer_hint_for_selection( - &entry.workers_tx.borrow(), - response.best_worker, - response.target_cached_prefix_blocks, - response.kv_transfer_candidates.as_ref(), - ) - .map(|payload| { - KvHint::new( - admission.request_id().unwrap_or_default(), - vec![KvHintAction::fetch("a1", payload)], - ) - }) - } else { - None - }; - - // The routing hashes go to exactly one of: the reservation recorded - // now, or the replay cache a later reservation records from. - let mut routing_hashes = routing_hashes; - let lease = if let Some(claim) = claim { - let Some(lease) = lease else { - return Err(SelectionError::Internal( - "booked selection has no lifecycle lease".to_string(), - )); - }; - // A rejected affinity commit returns while the lease is still armed, - // so the booking is freed and nothing below is recorded. - let affinity_lease = match (affinity_hold, managed_session) { - (Some(hold), Some((table, session_id))) => { - self.commit_session(table, hold, session_id, response.best_worker, &key)? - } - _ => None, - }; - if let Some(hashes) = routing_hashes.take() { - self.record_routing_decision(&entry, response.best_worker, hashes) - .await; - } - claim.install(Reservation { - partition: key.clone(), - booking: Some(lease.commit().ok_or_else(missing_booking)?), - _affinity_lease: affinity_lease, - }); - None - } else { - lease - }; - - if let Some((cache_id, sequence_hashes, lora_name, track_prefill_tokens, session_id)) = - cached_inputs - { - self.selection_cache.insert( - cache_id, - PendingSelection { - key: key.clone(), - worker: response.best_worker, - sequence_hashes, - isl_tokens, - effective_prefill_tokens: effective_prefill, - expected_output_tokens, - track_prefill_tokens, - lora_name, - routing_hashes, - session_id, - }, - Instant::now(), - ); - } - Ok(SelectionOutcome::Selected(Selected { - key, - response, - advisory_load, - total_kv_blocks, - endpoint, - block_size: entry.block_size, - isl_tokens, - sequence_hashes: booked_sequence_hashes, - track_prefill_tokens, - effective_prefill_tokens: effective_prefill, - kv_hint, - routing_hashes: returned_routing_hashes, - shared_cache_hits: host_shared_cache_hits, - lease, - })) - } - - pub async fn create_reservation( - &self, - req: ReservationRequest, - ) -> Result { - self.ensure_running()?; - - let key = RoutingPartitionId::new(req.model_name.clone(), req.routing_group.clone()); - - // Explicit form: book on the given worker under selection_id, discarding - // any cached selection for the id so a later replay can't book stale state. - if let Some(worker_id) = req.worker_id { - self.selection_cache.discard(&key, &req.selection_id); - return self.reserve_explicit(key, worker_id, req).await; - } - - // Replay form: peek, book, and consume only once the booking lands. A - // failure leaves the entry for a retry; concurrent replays of the same - // id collide at the scheduler, so they can't double-book. - let Some((pending, generation)) = - self.selection_cache - .peek(&key, &req.selection_id, Instant::now()) - else { - return Err(SelectionError::NotFound(format!( - "no pending selection {} for {key} (expired, already used, \ - or never selected)", - req.selection_id - ))); - }; - let response = match self.book_cached_selection(pending, &req).await { - Ok(response) => response, - // The session moved: the cached worker can never be booked for it. - Err(SelectionError::BadRequest(message)) => { - self.selection_cache.discard(&key, &req.selection_id); - return Err(SelectionError::BadRequest(message)); - } - Err(error) => return Err(error), - }; - self.selection_cache - .remove(&key, &req.selection_id, generation); - Ok(response) - } - - /// Book a reservation replaying what the matching `select` captured; request - /// fields other than the ids are ignored. - async fn book_cached_selection( - &self, - pending: Arc, - req: &ReservationRequest, - ) -> Result { - let (entry, endpoint, prefill_load_hint) = self.resolve_cached_booking(&pending)?; - let track_prefill_tokens = pending.track_prefill_tokens; - self.finalize_reservation( - entry, - endpoint, - ReservationBooking { - key: pending.key.clone(), - selection_id: req.selection_id.clone(), - worker: pending.worker, - sequence_hashes: pending.sequence_hashes.clone(), - prefill_load_hint: track_prefill_tokens.then_some(prefill_load_hint), - expected_output_tokens: pending.expected_output_tokens, - track_prefill_tokens, - lora_name: pending.lora_name.clone(), - routing_hashes: pending.routing_hashes.clone(), - session_id: pending.session_id.clone(), - }, - ) - .await - } - - /// Resolve everything a cached booking needs (ready entry, schedulable - /// endpoint, prefill hint), so the only fallible step left in - /// `finalize_reservation` is the scheduler call. - fn resolve_cached_booking( - &self, - pending: &PendingSelection, - ) -> Result<(Arc, String, PrefillLoadHint), SelectionError> { - let entry = self.ready_entry(&pending.key)?; - // Validate the full worker/rank against current topology; a rank a PATCH - // removed during the window is rejected (the entry stays for a retry). - let endpoint = self - .catalog - .schedulable_worker_endpoint(pending.worker, &pending.key) - .ok_or_else(|| { - SelectionError::NotFound(format!( - "schedulable worker {} (dp_rank {}) not found for {}", - pending.worker.worker_id, pending.worker.dp_rank, pending.key - )) - })?; - let prefill_load_hint = prefill_load_hint_from_effective_tokens( - pending.isl_tokens, - pending.effective_prefill_tokens, - ) - .map_err(|error| SelectionError::BadRequest(error.to_string()))?; - Ok((entry, endpoint, prefill_load_hint)) - } - - fn schedulable_endpoint( - &self, - worker_id: WorkerId, - key: &RoutingPartitionId, - ) -> Result { - self.catalog - .schedulable_endpoint(worker_id, key) - .ok_or_else(|| { - SelectionError::NotFound(format!( - "schedulable worker {worker_id} not found for {key}" - )) - }) - } - - /// Book a reservation from a self-contained request (explicit worker_id and prompt). - async fn reserve_explicit( - &self, - key: RoutingPartitionId, - worker_id: WorkerId, - req: ReservationRequest, - ) -> Result { - let entry = self.ready_entry(&key)?; - let normalized = req.prompt.view().normalize_for_reservation( - entry.is_eagle, - TrackingHashInput { - context: &self.tracking_hash, - scope: tracking_scope(&entry), - assume_kv_reuse: self - .kv_router_config - .assume_kv_reuse(req.router_config_override.as_ref()), - }, - )?; - let prefill_load_hint = req - .effective_prefill_tokens - .map(|tokens| { - prefill_load_hint_from_effective_tokens(normalized.isl_tokens, tokens) - .map_err(|error| SelectionError::BadRequest(error.to_string())) - }) - .transpose()?; - let worker = WorkerWithDpRank::new(worker_id, req.dp_rank.unwrap_or(0)); - let endpoint = self.schedulable_endpoint(worker.worker_id, &key)?; - let track_prefill_tokens = req.track_prefill_tokens.unwrap_or_else(|| { - req.effective_prefill_tokens.is_some() - || req - .router_config_override - .as_ref() - .and_then(|cfg| cfg.track_prefill_tokens) - .unwrap_or(self.kv_router_config.router_track_prefill_tokens) - }); - // Hash-only reservations (sequence hashes without block hashes) carry - // nothing an indexer can key on; recording is skipped for them. - let can_record = entry.indexer.records_routing_decisions() - && (req.prompt.token_ids.is_some() || req.prompt.block_hashes.is_some()); - let routing_hashes = can_record - .then(|| { - req.prompt - .view() - .block_hashes_for_indexer(entry.block_size, entry.is_eagle) - }) - .transpose()?; - - self.finalize_reservation( - entry, - endpoint, - ReservationBooking { - key, - selection_id: req.selection_id, - worker, - sequence_hashes: normalized.sequence_hashes, - prefill_load_hint, - expected_output_tokens: req.expected_output_tokens, - track_prefill_tokens, - lora_name: req.prompt.lora_name, - routing_hashes, - session_id: None, - }, - ) - .await - } - - /// Register the booking with the scheduler. All fallible resolution happens - /// in the caller; the scheduler add here is the last step that can fail, and - /// the cached path leaves its selection in place (to retry) if it does. - async fn finalize_reservation( - &self, - entry: Arc, - endpoint: String, - booking: ReservationBooking, - ) -> Result { - let ReservationBooking { - key, - selection_id, - worker, - sequence_hashes, - prefill_load_hint, - expected_output_tokens, - track_prefill_tokens, - lora_name, - routing_hashes, - session_id, - } = booking; - - let claim = self.claim_reservation(&selection_id, &key)?; - // Hold the session before booking: holding after would wait on an - // initializer that may itself be waiting for this worker's capacity. - let session = match (session_id.as_deref(), entry.affinity.get()) { - (Some(session_id), Some(table)) => self - .hold_session(table, session_id, &key) - .await? - .map(|hold| (table, session_id, hold)), - _ => None, - }; - // Strict booking: never lazily recreate a worker/rank removed since the - // reservation was resolved. The lease frees the booking if this future - // is dropped before `install`. - let lease = entry - .scheduler - .add_request_if_registered_guarded(SequenceRequest { - request_id: selection_id.clone(), - token_sequence: Some(sequence_hashes), - track_prefill_tokens, - expected_output_tokens, - prefill_load_hint, - worker, - lora_name, - })?; - let affinity_lease = match session { - Some((table, session_id, hold)) => { - self.commit_session(table, hold, session_id, worker, &key)? - } - None => None, - }; - if let Some(hashes) = routing_hashes { - self.record_routing_decision(&entry, worker, hashes).await; - } - claim.install(Reservation { - partition: key.clone(), - booking: Some(lease.commit().ok_or_else(missing_booking)?), - _affinity_lease: affinity_lease, - }); - - Ok(ReservationResponse { - selection_id, - model_name: key.model_name, - routing_group: key.routing_group, - worker_id: worker.worker_id, - dp_rank: worker.dp_rank, - endpoint, - }) - } - - /// Take `selection_id` for a booking about to be made in `key`. - fn claim_reservation( - &self, - selection_id: &str, - key: &RoutingPartitionId, - ) -> Result, SelectionError> { - let mut index = self.reservation_index.write(); - if index.contains_key(selection_id) { - return Err(SelectionError::Conflict(format!( - "selection {selection_id} is already reserved or being reserved" - ))); - } - index.insert( - selection_id.to_string(), - Reservation { - partition: key.clone(), - booking: None, - _affinity_lease: None, - }, - ); - Ok(ReservationClaim { - index: &self.reservation_index, - selection_id: selection_id.to_string(), - armed: true, - }) +fn tracking_scope(entry: &SelectionEntry) -> TrackingHashScope<'_> { + TrackingHashScope { + partition: entry.key.as_ref(), + block_size: entry.block_size, } - - /// Hold `session_id` for a booking, re-initializing a binding whose worker - /// this partition can no longer schedule. `None` when the table is full: a - /// router-side limit, not a client fault, so the request routes unpinned. - async fn hold_session( - &self, - table: &SessionAffinity, - session_id: &str, - key: &RoutingPartitionId, - ) -> Result, SelectionError> { - loop { - let acquired = tokio::select! { - _ = self.cancel_token.cancelled() => { - return Err(SelectionError::Scheduler(KvSchedulerError::SubscriberShutdown)); - } - result = table.acquire(session_id, None) => result, - }; - match acquired { - Ok(Hold::Bound { target, mut lease }) - if !self.catalog.is_schedulable(target.worker_id, key) => - { - tracing::debug!( - session_id, - worker_id = target.worker_id, - "session affinity target is not schedulable; re-initializing" - ); - lease.invalidate(); - } - Ok(hold) => return Ok(Some(hold)), - Err(AffinityError::ResourceExhausted(_)) => { - tracing::debug!( - session_id, - "affinity table full; routing without session affinity" - ); - return Ok(None); - } - Err(error) => return Err(affinity_error(error)), - } - } - } - - /// Bind the held session to `dispatched`. A `Hard` rejection whose bound - /// worker departed after [`Self::hold_session`] checked it is not a client - /// fault: the session is re-initialized on the dispatched worker instead. - /// This runs with the booking held, so it never waits on another request's - /// initialization (which may itself be queued behind this booking); when - /// one is in flight the booking stands unbound, `None`. - fn commit_session( - &self, - table: &SessionAffinity, - hold: Hold, - session_id: &str, - dispatched: WorkerWithDpRank, - key: &RoutingPartitionId, - ) -> Result, SelectionError> { - let bound = hold.target(); - let dispatched = WorkerAffinityTarget::new(dispatched.worker_id, Some(dispatched.dp_rank)); - match table.commit(hold, dispatched) { - Ok(lease) => Ok(Some(lease)), - Err(error) => { - let departed = - bound.is_some_and(|target| !self.catalog.is_schedulable(target.worker_id, key)); - if !departed { - return Err(affinity_error(error)); - } - // `commit` invalidated the stale binding. - match table.try_acquire(session_id, None) { - Ok(AcquireStep::Held(hold)) => table - .commit(hold, dispatched) - .map(Some) - .map_err(affinity_error), - Ok(AcquireStep::Wait(_)) | Err(AffinityError::ResourceExhausted(_)) => Ok(None), - Err(error) => Err(affinity_error(error)), - } - } - } - } - - /// Apply a session binding a replica published. - pub(crate) fn dispatch_affinity_event(&self, event: AffinityBindingEvent) { - let Some(entry) = self.entry(&event.partition) else { - return; - }; - let Some(table) = entry.affinity.get() else { - return; - }; - if self - .replica_config - .as_ref() - .is_some_and(|config| config.process_id() == event.writer_id) - { - return; - } - table.observe_replica_sequence(event.sequence); - if self - .catalog - .get(event.worker_id) - .is_none_or(|record| record.key() != event.partition) - { - return; - } - let (target, version, worker_id) = (event.target(), event.version(), event.worker_id); - let outcome = table.apply_replica_update(event.session_id, target, version); - tracing::trace!( - worker_id, - ?outcome, - "applied session affinity replica update" - ); - } - - /// Record a booked routing decision into the partition's approximate - /// indexer (side or primary). The booking already landed, so a failure - /// here only costs predicted cache credit; it is logged, not returned. - async fn record_routing_decision( - &self, - entry: &SelectionEntry, - worker: WorkerWithDpRank, - block_hashes: Vec, - ) { - if block_hashes.is_empty() { - return; - } - if let Err(error) = entry - .indexer - .record_routing_decision( - worker, - RoutingDecisionHashes::from_local_hashes(block_hashes), - ) - .await - { - tracing::warn!( - %error, - key = %entry.key, - worker_id = worker.worker_id, - dp_rank = worker.dp_rank, - "Failed to record routing decision into approximate indexer" - ); - } - } - - /// The partition and booking a lifecycle call on `selection_id` may touch. - /// `None` for unknown ids and for ids whose booking is still in flight. - fn indexed_booking( - &self, - selection_id: &str, - ) -> Option<(Arc, SchedulerBookingDescriptor)> { - // Release the index guard before taking `entries`: the sweep takes - // `entries` then `reservation_index`, so nesting here would invert the - // lock order. - let (partition, booking) = { - let index = self.reservation_index.read(); - let reservation = index.get(selection_id)?; - (reservation.partition.clone(), reservation.booking.clone()?) - }; - Some((self.entry(&partition)?, booking)) - } - - fn reservation_not_found(selection_id: &str) -> SelectionError { - SelectionError::NotFound(format!("reservation {selection_id} not found")) - } - - pub async fn prefill_complete(&self, selection_id: &str) -> Result<(), SelectionError> { - let Some((entry, booking)) = self.indexed_booking(selection_id) else { - return Err(Self::reservation_not_found(selection_id)); - }; - match entry - .scheduler - .mark_prefill_completed_if_booking(&booking) - .await? - { - LifecycleMutationOutcome::Applied => Ok(()), - // Already marked: still publish the ordered completion so peers - // that missed the first event converge. - LifecycleMutationOutcome::NoChange - if entry - .scheduler - .publish_prefill_completed_if_booking(&booking) => - { - Ok(()) - } - LifecycleMutationOutcome::NoChange => { - forget_reservation_if(&self.reservation_index, &entry.key, &booking); - Err(Self::reservation_not_found(selection_id)) - } - } - } - - pub async fn free_reservation(&self, selection_id: &str) -> Result<(), SelectionError> { - let Some((entry, booking)) = self.indexed_booking(selection_id) else { - return Err(Self::reservation_not_found(selection_id)); - }; - let outcome = entry.scheduler.free_if_booking(&booking).await?; - forget_reservation_if(&self.reservation_index, &entry.key, &booking); - match outcome { - LifecycleMutationOutcome::Applied => Ok(()), - LifecycleMutationOutcome::NoChange => Err(Self::reservation_not_found(selection_id)), - } - } - - pub fn add_output_block( - &self, - selection_id: &str, - decay_fraction: Option, - ) -> Result<(), SelectionError> { - if let Some(frac) = decay_fraction - && !(0.0..=1.0).contains(&frac) - { - return Err(SelectionError::BadRequest( - "decay_fraction must be between 0.0 and 1.0".to_string(), - )); - } - - let Some((entry, booking)) = self.indexed_booking(selection_id) else { - return Err(Self::reservation_not_found(selection_id)); - }; - match entry - .scheduler - .add_output_block_if_booking_sync(&booking, decay_fraction)? - { - LifecycleMutationOutcome::Applied => Ok(()), - LifecycleMutationOutcome::NoChange => { - forget_reservation_if(&self.reservation_index, &entry.key, &booking); - Err(Self::reservation_not_found(selection_id)) - } - } - } - - pub fn loads( - &self, - model_name: Option<&str>, - routing_group: Option<&str>, - ) -> Vec { - let entries = self.initialized_entries(); - let mut loads = Vec::new(); - for entry in entries { - if model_name.is_some_and(|model_name| entry.key.model_name != model_name) - || routing_group - .is_some_and(|routing_group| entry.key.routing_group != routing_group) - { - continue; - } - loads.push(ModelLoadResponse { - model_name: entry.key.model_name.clone(), - routing_group: entry.key.routing_group.clone(), - loads: entry - .scheduler - .get_potential_loads(None, 0, HashMap::new(), false), - pending_count: entry.scheduler.pending_count(), - pending_isl_tokens: entry.scheduler.pending_isl_tokens(), - }); - } - loads.sort_by(|a, b| { - (&a.model_name, &a.routing_group).cmp(&(&b.model_name, &b.routing_group)) - }); - loads - } - - pub async fn potential_loads( - &self, - req: PotentialLoadsRequest, - ) -> Result, SelectionError> { - let key = RoutingPartitionId::new(req.model_name.clone(), req.routing_group.clone()); - let entry = self.ready_entry(&key)?; - let prepared = self - .prepare_selection_inputs( - &entry, - &req.prompt.view(), - Some( - self.kv_router_config - .assume_kv_reuse(req.router_config_override.as_ref()), - ), - false, - false, - &mut None, - ) - .await?; - let track_prefill_tokens = req - .router_config_override - .as_ref() - .and_then(|cfg| cfg.track_prefill_tokens) - .unwrap_or(self.kv_router_config.router_track_prefill_tokens); - Ok(entry.scheduler.get_potential_loads( - Some(prepared.sequence_hashes), - prepared.isl_tokens, - prepared.overlap.effective_cached_tokens, - track_prefill_tokens, - )) - } - - pub async fn overlap_scores( - &self, - req: OverlapScoresRequest, - ) -> Result { - let key = RoutingPartitionId::new(req.model_name.clone(), req.routing_group.clone()); - let entry = self.ready_entry(&key)?; - let block_hashes = req - .prompt - .view() - .block_hashes_for_indexer(entry.block_size, entry.is_eagle)?; - let num_blocks = block_hashes.len(); - let tiered = entry - .indexer - .find_tiered_matches(block_hashes) - .await - .map_err(|error| SelectionError::Internal(error.to_string()))?; - let schedulable_workers = self.schedulable_worker_ranks(&key); - Ok( - OverlapAnalysis::new(&self.kv_router_config, entry.block_size, &tiered) - .scores_response( - req.router_config_override.as_ref(), - num_blocks, - schedulable_workers, - false, - None, - None, - ), - ) - } - - /// Normalize the prompt and gather cache signals. The indexer lookup and - /// the optional shared-cache lookup run concurrently; the shared cache is - /// consulted only when `query_shared_cache` is set, a shared cache is - /// attached, and the prompt carries raw `token_ids`. - /// `tracking_assume_kv_reuse` is `None` when the caller does not track - /// active blocks. - async fn prepare_selection_inputs( - &self, - entry: &SelectionEntry, - prompt: &PromptView<'_>, - tracking_assume_kv_reuse: Option, - query_shared_cache: bool, - retain_kv_transfer_chain: bool, - lookup: &mut Option, - ) -> Result { - let normalized = prompt.normalize_for_selection( - entry.block_size, - entry.is_eagle, - tracking_assume_kv_reuse.map(|assume_kv_reuse| TrackingHashInput { - context: &self.tracking_hash, - scope: tracking_scope(entry), - assume_kv_reuse, - }), - )?; - let indexer_lookup = async { - let started = Instant::now(); - let tiered = if normalized.block_hashes.is_empty() { - Ok(TieredMatchDetails::default()) - } else { - entry - .indexer - .find_tiered_matches_with_options( - normalized.block_hashes.clone(), - LowerTierQueryOptions { - retain_kv_transfer_chain, - }, - ) - .instrument(tracing::info_span!("kv_router.find_matches")) - .await - .map_err(SelectionError::Indexer) - }; - (tiered, started.elapsed()) - }; - let shared_cache = query_shared_cache - .then_some(self.host.cache.shared.as_deref()) - .flatten() - .zip(prompt.token_ids); - // `(hits, duration, failed)`; all `None`/`false` without a shared cache. - let shared_cache_lookup = async { - let Some((shared_cache, tokens)) = shared_cache else { - return (None, None, false); - }; - let started = Instant::now(); - let result = shared_cache - .check_blocks(tokens, entry.block_size, prompt.cache_namespace) - .instrument(tracing::info_span!("kv_router.shared_cache_check")) - .await; - let elapsed = started.elapsed(); - match result { - Ok(hits) => (Some(hits), Some(elapsed), false), - Err(error) => { - tracing::warn!(%error, "Shared cache query failed, ignoring"); - (None, Some(elapsed), true) - } - } - }; - let lookups_started = Instant::now(); - let ((tiered, indexer), (shared_cache_hits, shared_cache, shared_cache_error)) = - tokio::join!(indexer_lookup, shared_cache_lookup); - let timings = LookupTimings { - block_hashing: normalized.block_hashing, - seq_hashing: normalized.seq_hashing, - lookups: lookups_started.elapsed(), - indexer, - shared_cache, - shared_cache_error, - }; - // Recorded before `?` so a shared-cache failure is still accounted - // when the index lookup fails. - *lookup = Some(timings); - let tiered = tiered?; - let overlap = - OverlapAnalysis::new(&self.kv_router_config, entry.block_size, &tiered).signals(); - let kv_transfer_candidates = retain_kv_transfer_chain - .then(|| tiered.kv_transfer_candidates().cloned()) - .flatten(); - drop(tiered); - Ok(PreparedSelectionInputs { - block_hashes: normalized.block_hashes, - sequence_hashes: normalized.sequence_hashes, - isl_tokens: normalized.isl_tokens, - overlap, - shared_cache_hits, - kv_transfer_candidates, - }) - } - - fn schedulable_worker_ranks(&self, key: &RoutingPartitionId) -> Vec { - let configs = self.catalog.scheduler_configs_for_key(key); - let mut workers = Vec::new(); - for (worker_id, config) in configs { - let start = config.data_parallel_start_rank; - let end = start.saturating_add(config.data_parallel_size); - for dp_rank in start..end { - workers.push(WorkerWithDpRank::new(worker_id, dp_rank)); - } - } - workers - } -} - -/// Drop index entries whose booking no longer exists in its partition -/// scheduler (expired by the periodic force-expiry, or freed through a path -/// that bypassed this core). Claims are left for their owner. Returns the -/// number of entries removed. -fn sweep_reservation_index(entries: &SelectionEntries, index: &ReservationIndex) -> usize { - let entries = entries.read(); - let mut index = index.write(); - let before = index.len(); - index.retain(|_, reservation| { - let Some(booking) = &reservation.booking else { - return true; - }; - entries - .get(&reservation.partition) - .and_then(|cell| cell.get()) - .is_some_and(|entry| entry.scheduler.has_booking(booking)) - }); - before - index.len() -} - -/// The core's use of the session table for a request: a session is steered -/// only when the caller did not pin or target a worker explicitly. -fn session_binding( - context: Option<&SessionContext>, - is_steerable: bool, - is_booking: bool, -) -> SessionBinding { - match context.filter(|_| is_steerable) { - Some(context) if is_booking => SessionBinding::Managed { - session_id: context.session_id().to_string(), - }, - Some(context) => SessionBinding::Query { - session_id: context.session_id().to_string(), - }, - None => SessionBinding::None, - } -} - -fn missing_booking() -> SelectionError { - SelectionError::Internal("booking lease holds no booking".to_string()) -} - -fn spawn_reservation_index_sweep( - entries: Arc, - index: Arc, - cancel_token: CancellationToken, -) { - let period = crate::sequences::active_request_expiry_duration(); - tokio::spawn(async move { - let mut interval = tokio::time::interval(period); - interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); - interval.tick().await; - loop { - tokio::select! { - _ = cancel_token.cancelled() => break, - _ = interval.tick() => { - let removed = sweep_reservation_index(&entries, &index); - if removed > 0 { - tracing::debug!(removed, "Swept stale selection reservation index entries"); - } - } - } - } - }); -} - -/// Pick the best router-hint source for `target`: a same-role worker (or -/// cache owner) holding a longer root-aligned prefix than the target's own -/// `target_cached_prefix_blocks`, with a non-empty control endpoint. Mirrors the -/// frontend `KvRouter::transfer_hint_for_selection`. -fn transfer_hint_for_selection( - configs: &HashMap, - target: WorkerWithDpRank, - target_cached_prefix_blocks: u32, - candidates: Option<&KvTransferCandidates>, -) -> Option { - let candidates = candidates?; - let target_config = configs.get(&target.worker_id)?; - let target_metadata = target_config.kv_hint_transfer_metadata_for_dp_rank(target.dp_rank)?; - - let prefix_blocks_to_beat = usize::try_from(target_cached_prefix_blocks).unwrap_or(usize::MAX); - let (source, block_hashes) = - candidates.best_source(prefix_blocks_to_beat, |source| match source { - KvTransferCandidateSource::Worker(worker) => { - worker != target - && configs.get(&worker.worker_id).is_some_and(|config| { - config.kv_event_source_mode.as_deref() != Some("state_agent_v2") - && config - .kv_hint_transfer_metadata_for_dp_rank(worker.dp_rank) - .is_some_and(|source_metadata| { - source_metadata.worker_type == target_metadata.worker_type - && source_metadata - .source_control_endpoint - .is_some_and(|endpoint| !endpoint.is_empty()) - }) - }) - } - KvTransferCandidateSource::CacheOwner(owner) => candidates - .routing_snapshot - .as_ref() - .and_then(|snapshot| snapshot.router_hint_source(owner)) - .is_some_and(|source| { - source.attached_worker != Some(target) - && source.metadata.worker_type == target_metadata.worker_type - && !source.metadata.source_control_endpoint.is_empty() - }), - })?; - let source_control_endpoint = match source { - KvTransferCandidateSource::Worker(worker) => configs - .get(&worker.worker_id)? - .kv_hint_transfer_metadata_for_dp_rank(worker.dp_rank)? - .source_control_endpoint? - .to_string(), - KvTransferCandidateSource::CacheOwner(owner) => candidates - .routing_snapshot - .as_ref()? - .router_hint_source(owner)? - .metadata - .source_control_endpoint - .clone(), - }; - if block_hashes.is_empty() { - return None; - } - Some(KvSourceLocationsPayload { - source_control_endpoint, - block_hashes, - }) -} - -fn tracking_scope(entry: &SelectionEntry) -> TrackingHashScope<'_> { - TrackingHashScope { - partition: entry.key.as_ref(), - block_size: entry.block_size, - } -} +} impl Drop for SelectionCore { fn drop(&mut self) { self.shutdown(); } } - -#[cfg(test)] -mod tests { - use super::super::affinity::SessionAffinityMode; - use super::super::input::PromptRequest; - use super::*; - use crate::protocols::{RoutingConstraints, StorageTier}; - use crate::services::indexer::backend::test_util::store_event; - use std::collections::HashSet; - use std::sync::atomic::{AtomicBool, Ordering}; - use std::thread::sleep; - use std::time::Duration; - - fn test_config(use_kv_events: bool) -> crate::config::KvRouterConfig { - crate::config::KvRouterConfig { - use_kv_events, - router_queue_threshold: None, - ..Default::default() - } - } - - fn worker(worker_id: WorkerId) -> WorkerRequest { - WorkerRequest { - worker_id, - model_name: "model".to_string(), - routing_group: "default".to_string(), - endpoint: Some(format!("http://worker-{worker_id}:8000")), - kv_events_endpoint: None, - kv_events_endpoints: HashMap::new(), - replay_endpoint: None, - block_size: Some(4), - data_parallel_start_rank: None, - data_parallel_size: None, - max_num_batched_tokens: Some(1024), - total_kv_blocks: None, - stable_routing_id: None, - is_eagle: None, - taints: HashSet::new(), - topology_domains: HashMap::new(), - kv_transfer_domain: None, - kv_transfer_enforcement: None, - kv_transfer_preferred_weight: None, - router_hint_worker_type: None, - router_hint_source_control_endpoints: HashMap::new(), - kv_event_source_mode: None, - } - } - - fn worker_with_kv_events(worker_id: WorkerId) -> WorkerRequest { - WorkerRequest { - kv_events_endpoint: Some("tcp://127.0.0.1:5557".to_string()), - ..worker(worker_id) - } - } - - fn prompt() -> PromptRequest { - PromptRequest { - token_ids: Some(vec![1, 2, 3, 4]), - mm_routing_info: None, - block_mm_infos: None, - block_hashes: None, - sequence_hashes: None, - isl_tokens: None, - lora_name: None, - cache_namespace: None, - is_eagle: None, - } - } - - fn select_request() -> SelectRequest { - SelectRequest { - model_name: "model".to_string(), - routing_group: "default".to_string(), - selection_id: None, - prompt: prompt(), - router_config_override: None, - expected_output_tokens: None, - priority_jump: None, - strict_priority: None, - session_id: None, - session_context: None, - affinity_target: None, - pinned_worker: None, - allowed_worker_ids: None, - routing_constraints: RoutingConstraints::default(), - advisory: false, - } - } - - fn reserve_request(selection_id: &str) -> SelectAndReserveRequest { - SelectAndReserveRequest { - model_name: "model".to_string(), - routing_group: "default".to_string(), - selection_id: Some(selection_id.to_string()), - prompt: prompt(), - router_config_override: None, - expected_output_tokens: None, - priority_jump: None, - strict_priority: None, - session_id: None, - session_context: None, - affinity_target: None, - pinned_worker: None, - allowed_worker_ids: None, - routing_constraints: RoutingConstraints::default(), - } - } - - async fn wait_until(what: &str, mut condition: impl FnMut() -> bool) { - tokio::time::timeout(Duration::from_secs(2), async { - while !condition() { - tokio::task::yield_now().await; - } - }) - .await - .unwrap_or_else(|_| panic!("timed out waiting for {what}")); - } - - async fn wait_for_pending_selection(core: &SelectionCore) { - wait_until("pending selection", || { - core.loads(Some("model"), Some("default"))[0].pending_count == 1 - }) - .await; - } - - fn assert_shutdown_error(error: SelectionError) { - assert!(matches!( - error, - SelectionError::NotReady(message) - if message == "selection service is shutting down" - )); - } - - #[test] - fn parent_cancel_cancels_core() { - let parent = CancellationToken::new(); - let core = SelectionCore::try_new_local( - test_config(false), - 1, - parent.clone(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"); - - assert!(!core.cancel_token.is_cancelled()); - parent.cancel(); - assert!(core.cancel_token.is_cancelled()); - } - - #[tokio::test] - async fn selection_setup_uses_worker_type_label() { - for (worker_type, expected_label) in [ - (WorkerType::Prefill, "prefill"), - (WorkerType::Decode, "decode"), - (WorkerType::Encode, "encode"), - (WorkerType::Aggregated, "aggregated"), - ] { - let config = test_config(false); - let tracking_hash = Arc::new( - TrackingHashContext::from_config(&config) - .expect("valid tracking hash configuration"), - ); - let core = SelectionCore::new_inner( - config, - 1, - CancellationToken::new(), - None, - None, - SelectionHost::default(), - worker_type, - true, - SelectionCacheConfig::default(), - tracking_hash, - IndexerPolicy::from_router_config(&test_config(false)).expect("indexer policy"), - None, - ); - - core.upsert_worker(worker(1)).await.expect("worker upsert"); - let entry = core - .entry(&RoutingPartitionId::new("model", "default")) - .expect("selection entry"); - assert_eq!( - entry.scheduler.worker_type(), - expected_label, - "{worker_type}" - ); - } - } - - fn core_with_host(host: SelectionHost) -> SelectionCore { - core_with_host_and_policy(host, None) - } - - fn core_with_host_and_policy( - host: SelectionHost, - policy_factory: Option, - ) -> SelectionCore { - let config = test_config(false); - let tracking_hash = Arc::new( - TrackingHashContext::from_config(&config).expect("valid tracking hash configuration"), - ); - let indexer_policy = IndexerPolicy::from_router_config(&config).expect("indexer policy"); - SelectionCore::new_inner( - config, - 1, - CancellationToken::new(), - None, - policy_factory, - host, - WorkerType::Aggregated, - true, - SelectionCacheConfig::default(), - tracking_hash, - indexer_policy, - None, - ) - } - - async fn wait_for_overlap( - core: &SelectionCore, - request: impl Fn() -> SelectRequest, - ) -> SelectResponse { - tokio::time::timeout(Duration::from_secs(2), async { - loop { - let response = core.select(request()).await.expect("select"); - if response.overlap.longest_matched > 0 { - return response; - } - tokio::task::yield_now().await; - } - }) - .await - .expect("approximate indexer never credited the booked prompt") - } - - #[tokio::test] - async fn bookings_populate_the_approximate_primary_without_kv_events() { - // use_kv_events=false: the primary is approximate and bookings feed it. - let core = SelectionCore::try_new_local( - test_config(false), - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"); - core.upsert_worker(worker(1)).await.expect("worker upsert"); - core.upsert_worker(worker(2)).await.expect("worker upsert"); - - // Query-only selection records nothing. - let first = core.select(select_request()).await.expect("select"); - assert_eq!(first.overlap.longest_matched, 0); - - // select_and_reserve records the routed prefix for the chosen worker. - let booked = core - .select_and_reserve(reserve_request("booked")) - .await - .expect("reserve"); - let credited = wait_for_overlap(&core, || { - let mut request = select_request(); - request.allowed_worker_ids = Some(HashSet::from([booked.worker_id])); - request - }) - .await; - assert_eq!(credited.worker_id, booked.worker_id); - assert_eq!(credited.overlap.longest_matched, 4); - - // The cached replay path records too, on a different prompt. - let prompt_b = || PromptRequest { - token_ids: Some(vec![5, 6, 7, 8]), - ..PromptRequest::default() - }; - let mut request = select_request(); - request.prompt = prompt_b(); - request.selection_id = Some("cached".to_string()); - request.allowed_worker_ids = Some(HashSet::from([2])); - core.select(request).await.expect("select"); - core.create_reservation(ReservationRequest { - model_name: "model".to_string(), - routing_group: "default".to_string(), - selection_id: "cached".to_string(), - worker_id: None, - dp_rank: None, - prompt: PromptRequest::default(), - router_config_override: None, - expected_output_tokens: None, - effective_prefill_tokens: None, - track_prefill_tokens: None, - }) - .await - .expect("cached reservation"); - let credited = wait_for_overlap(&core, || { - let mut request = select_request(); - request.prompt = prompt_b(); - request.allowed_worker_ids = Some(HashSet::from([2])); - request - }) - .await; - assert_eq!(credited.worker_id, 2); - - // And the explicit reservation form. - let prompt_c = || PromptRequest { - token_ids: Some(vec![9, 10, 11, 12]), - ..PromptRequest::default() - }; - core.create_reservation(ReservationRequest { - model_name: "model".to_string(), - routing_group: "default".to_string(), - selection_id: "explicit".to_string(), - worker_id: Some(1), - dp_rank: None, - prompt: prompt_c(), - router_config_override: None, - expected_output_tokens: None, - effective_prefill_tokens: None, - track_prefill_tokens: None, - }) - .await - .expect("explicit reservation"); - let credited = wait_for_overlap(&core, || { - let mut request = select_request(); - request.prompt = prompt_c(); - request.allowed_worker_ids = Some(HashSet::from([1])); - request - }) - .await; - assert_eq!(credited.worker_id, 1); - } - - #[tokio::test] - async fn unreachable_remote_indexer_is_reported_not_ready() { - let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind"); - let base_url = format!("http://{}", listener.local_addr().expect("addr")); - drop(listener); - let config = test_config(true); - let tracking_hash = Arc::new( - TrackingHashContext::from_config(&config).expect("valid tracking hash configuration"), - ); - let indexer_policy = IndexerPolicy::from_router_config(&config) - .expect("indexer policy") - .with_remote_indexer(base_url) - .expect("remote policy"); - let core = SelectionCore::new_inner( - config, - 1, - CancellationToken::new(), - None, - None, - SelectionHost::default(), - WorkerType::Aggregated, - true, - SelectionCacheConfig::default(), - tracking_hash, - indexer_policy, - None, - ); - core.upsert_worker(worker(1)).await.expect("worker upsert"); - assert!(matches!( - core.select(select_request()).await, - Err(SelectionError::Indexer( - crate::indexer::KvRouterError::IndexerOffline - )) - )); - } - - #[tokio::test] - async fn remote_indexer_serves_selection_without_local_kv_listeners() { - use crate::indexer::KvIndexerInterface; - use crate::protocols::{BlockHashOptions, StorageTier, compute_block_hash_for_seq}; - use crate::services::indexer::registry::WorkerRegistry; - use crate::services::indexer::server::spawn_test_indexer_server; - - // The standalone indexer holds worker 2's cache for the test prompt. - let key = RoutingPartitionId::new("model", "default"); - let served = Arc::new(WorkerRegistry::new(1)); - let served_indexer = served.get_or_create_indexer(key.clone(), 4); - let hashes: Vec = - compute_block_hash_for_seq(&[1, 2, 3, 4], 4, BlockHashOptions::default()) - .into_iter() - .map(|hash| hash.0) - .collect(); - served_indexer - .apply_event_routed(store_event(2, 0, 1, &[], &hashes, StorageTier::Device)) - .await - .unwrap(); - if let Indexer::Single { primary, .. } = &served_indexer { - let _ = primary.flush().await; - } - let (base_url, server) = spawn_test_indexer_server(served).await; - - // use_kv_events=true, but the primary is remote: workers need no - // kv_events endpoints and no ZMQ listener is started here. - let config = test_config(true); - let tracking_hash = Arc::new( - TrackingHashContext::from_config(&config).expect("valid tracking hash configuration"), - ); - let indexer_policy = IndexerPolicy::from_router_config(&config) - .expect("indexer policy") - .with_remote_indexer(base_url) - .expect("remote policy"); - let core = SelectionCore::new_inner( - config, - 1, - CancellationToken::new(), - None, - None, - SelectionHost::default(), - WorkerType::Aggregated, - true, - SelectionCacheConfig::default(), - tracking_hash, - indexer_policy, - None, - ); - assert!(!core.listens_for_kv_events); - for worker_id in [1, 2] { - let record = core - .upsert_worker(worker(worker_id)) - .await - .expect("worker upsert"); - assert_eq!(record.lifecycle, WorkerLifecycle::Schedulable, "{record:?}"); - } - - let response = core.select(select_request()).await.expect("select"); - assert_eq!( - response.worker_id, 2, - "remote cache credit steers selection" - ); - assert_eq!(response.overlap.longest_matched, 4); - - core.delete_worker(2).await.expect("delete worker"); - server.abort(); - } - - /// Two event-driven workers; worker 1 holds every block of an 8-token - /// prompt. `configure` adjusts each registration before upsert. - async fn hint_fixture( - configure: impl Fn(&mut WorkerRequest), - ) -> (SelectionCore, Arc, Vec) { - use crate::indexer::KvIndexerInterface; - use crate::protocols::{BlockHashOptions, compute_block_hash_for_seq}; - - let core = SelectionCore::try_new_local( - test_config(true), - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"); - for worker_id in [1, 2] { - let mut request = worker_with_kv_events(worker_id); - configure(&mut request); - core.upsert_worker(request).await.expect("worker upsert"); - } - let entry = core - .entry(&RoutingPartitionId::new("model", "default")) - .expect("entry"); - let tokens: Vec = (1..=8).collect(); - let hashes: Vec = compute_block_hash_for_seq(&tokens, 4, BlockHashOptions::default()) - .into_iter() - .map(|hash| hash.0) - .collect(); - entry - .indexer - .apply_event_routed(store_event(1, 0, 1, &[], &hashes, StorageTier::Device)) - .await - .unwrap(); - if let Indexer::Single { primary, .. } = &entry.indexer { - let _ = primary.flush().await; - } - (core, entry, tokens) - } - - fn hint_capable(request: &mut WorkerRequest) { - request.router_hint_worker_type = Some("decode".to_string()); - request.router_hint_source_control_endpoints = - HashMap::from([(0, format!("tcp://worker-{}:9000", request.worker_id))]); - } - - async fn reserve_pinned( - core: &SelectionCore, - selection_id: &str, - tokens: &[u32], - worker_id: WorkerId, - ) -> SelectResponse { - let mut request = reserve_request(selection_id); - request.prompt = PromptRequest { - token_ids: Some(tokens.to_vec()), - ..PromptRequest::default() - }; - request.pinned_worker = Some(WorkerWithDpRank::new(worker_id, 0)); - core.select_and_reserve(request).await.expect("reserve") - } - - #[tokio::test] - async fn booked_selection_attaches_router_hint_from_a_better_source() { - let (core, entry, tokens) = hint_fixture(hint_capable).await; - assert!(entry.indexer.supports_kv_transfer_chain_retention()); - - // Booking on worker 2: worker 1 is a same-role source with a longer prefix. - let response = reserve_pinned(&core, "to-worker-2", &tokens, 2).await; - let hint = response.kv_hint.expect("router hint for worker 2"); - assert_eq!(hint.message_id, "to-worker-2"); - assert_eq!(hint.actions[0].action_type, "kv.fetch"); - let payload: KvSourceLocationsPayload = - serde_json::from_value(serde_json::to_value(&hint.actions[0].payload).unwrap()) - .unwrap(); - assert_eq!(payload.source_control_endpoint, "tcp://worker-1:9000"); - assert_eq!(payload.block_hashes.len(), 2); - - // Booking on worker 1 itself: nothing holds a longer prefix. - let response = reserve_pinned(&core, "to-worker-1", &tokens, 1).await; - assert!(response.kv_hint.is_none()); - - // Query-only selections never carry a hint. - let mut request = select_request(); - request.prompt = PromptRequest { - token_ids: Some(tokens.clone()), - ..PromptRequest::default() - }; - request.pinned_worker = Some(WorkerWithDpRank::new(2, 0)); - let response = core.select(request).await.expect("select"); - assert!(response.kv_hint.is_none()); - } - - #[tokio::test] - async fn state_agent_workers_are_not_router_hint_sources() { - let (core, _entry, tokens) = hint_fixture(|request| { - hint_capable(request); - if request.worker_id == 1 { - request.kv_event_source_mode = Some("state_agent_v2".to_string()); - } - }) - .await; - // Worker 1 holds the prefix but reports through a state agent. - let response = reserve_pinned(&core, "to-worker-2", &tokens, 2).await; - assert!(response.kv_hint.is_none()); - } - - #[tokio::test] - async fn router_hint_needs_capable_workers() { - let (core, _entry, tokens) = hint_fixture(|_| {}).await; - let response = reserve_pinned(&core, "plain", &tokens, 2).await; - assert!(response.kv_hint.is_none()); - } - - #[tokio::test] - async fn event_driven_indexer_does_not_record_bookings() { - let core = SelectionCore::try_new_local( - test_config(true), - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"); - core.upsert_worker(worker_with_kv_events(1)) - .await - .expect("worker upsert"); - core.select_and_reserve(reserve_request("booked")) - .await - .expect("reserve"); - core.free_reservation("booked").await.expect("free"); - for _ in 0..3 { - let response = core.select(select_request()).await.expect("select"); - assert_eq!(response.overlap.longest_matched, 0); - tokio::task::yield_now().await; - } - } - - /// Picker that records what worker selection saw and always takes row 0. - struct CapturingPicker { - observed: Arc>>, - } - - #[derive(Debug, Clone)] - struct SelectionObservation { - session_context: Option, - shared_beyond_device_blocks: Vec, - } - - impl crate::scheduling::selector::WorkerPicker for CapturingPicker { - fn required_worker_inputs(&self) -> crate::scheduling::selector::WorkerInputs { - crate::scheduling::selector::WorkerInputs::CACHE - } - - fn pick( - &mut self, - context: &crate::scheduling::selector::WorkerSelectionContext<'_>, - input: crate::scheduling::selector::WorkerInputView<'_>, - ) -> Result { - self.observed.lock().push(SelectionObservation { - session_context: context.session_context().cloned(), - shared_beyond_device_blocks: input - .cache() - .expect("CACHE inputs requested") - .iter() - .map(|cache| cache.shared_beyond_device_blocks()) - .collect(), - }); - Ok(0) - } - } - - fn capturing_policy_factory() -> ( - WorkerSelectionPolicyFactory, - Arc>>, - ) { - let observed = Arc::new(parking_lot::Mutex::new(Vec::new())); - let factory_observed = Arc::clone(&observed); - let factory: WorkerSelectionPolicyFactory = - Arc::new(move |config, worker_type, _partition| { - WorkerSelectionPolicy::new( - config.clone(), - worker_type.as_str(), - Vec::new(), - Box::new(CapturingPicker { - observed: Arc::clone(&factory_observed), - }), - ) - }); - (factory, observed) - } - - type SharedCacheCalls = Arc, u32, Option)>>>; - - /// Shared cache that reports every block as a hit and records each query. - struct RecordingSharedCache { - calls: SharedCacheCalls, - } - - #[async_trait::async_trait] - impl SharedKvCache for RecordingSharedCache { - async fn check_blocks( - &self, - tokens: &[u32], - block_size: u32, - cache_namespace: Option<&str>, - ) -> Result { - self.calls.lock().push(( - tokens.to_vec(), - block_size, - cache_namespace.map(str::to_string), - )); - let blocks = (tokens.len() / block_size as usize) as u32; - Ok(SharedCacheHits::from_hits(&vec![true; blocks as usize])) - } - } - - struct OnlyWorkerForLora { - worker_id: WorkerId, - } - - impl LoraWorkerFilter for OnlyWorkerForLora { - fn filter_worker_ids_for_lora( - &self, - _lora_name: &str, - available: &[WorkerId], - ) -> Vec { - available - .iter() - .copied() - .filter(|id| *id == self.worker_id) - .collect() - } - } - - #[tokio::test] - async fn injected_lora_filter_narrows_candidates() { - let core = core_with_host(SelectionHost { - eligibility: HostEligibility { - lora_worker_filter: Some(Arc::new(OnlyWorkerForLora { worker_id: 2 })), - }, - ..SelectionHost::default() - }); - core.upsert_worker(worker(1)).await.expect("worker upsert"); - core.upsert_worker(worker(2)).await.expect("worker upsert"); - - // LoRA request: only the filter's worker is eligible. - for _ in 0..4 { - let mut request = select_request(); - request.prompt.lora_name = Some("adapter-a".to_string()); - let response = core.select(request).await.expect("select"); - assert_eq!(response.worker_id, 2); - } - - // The filter never widens the caller's allow-set: an allow-set that - // excludes the filter's worker is preserved as-is. - let mut request = select_request(); - request.prompt.lora_name = Some("adapter-a".to_string()); - request.allowed_worker_ids = Some(HashSet::from([1])); - let response = core.select(request).await.expect("select"); - assert_eq!(response.worker_id, 1); - - // A pinned worker inside the universe survives the filter. - let mut request = select_request(); - request.prompt.lora_name = Some("adapter-a".to_string()); - request.pinned_worker = Some(WorkerWithDpRank::new(1, 0)); - let response = core.select(request).await.expect("select"); - assert_eq!(response.worker_id, 1); - } - - #[tokio::test] - async fn shared_cache_hits_reach_worker_selection() { - let calls: SharedCacheCalls = Arc::new(parking_lot::Mutex::new(Vec::new())); - let (factory, observed) = capturing_policy_factory(); - let core = core_with_host_and_policy( - SelectionHost { - cache: HostCache { - shared: Some(Arc::new(RecordingSharedCache { - calls: Arc::clone(&calls), - })), - ..HostCache::default() - }, - ..SelectionHost::default() - }, - Some(factory), - ); - core.upsert_worker(worker(1)).await.expect("worker upsert"); - - let mut request = select_request(); - request.prompt.cache_namespace = Some("tenant-a".to_string()); - core.select(request).await.expect("select"); - - assert_eq!( - calls.lock().as_slice(), - &[(vec![1, 2, 3, 4], 4, Some("tenant-a".to_string()))] - ); - let observations = observed.lock().clone(); - assert_eq!(observations.len(), 1); - // One block in the prompt, no device overlap, so the whole prompt is a - // shared-cache hit beyond the device prefix. - assert_eq!(observations[0].shared_beyond_device_blocks, vec![1]); - - // Load projection does not consult the shared cache. - core.potential_loads(PotentialLoadsRequest { - model_name: "model".to_string(), - routing_group: "default".to_string(), - prompt: prompt(), - router_config_override: None, - }) - .await - .expect("potential loads"); - assert_eq!(calls.lock().len(), 1); - - // Prompts without raw tokens cannot be checked against the shared cache. - let mut request = select_request(); - request.prompt = PromptRequest { - token_ids: None, - block_hashes: Some(vec![11]), - sequence_hashes: Some(vec![101]), - isl_tokens: Some(4), - ..PromptRequest::default() - }; - core.select(request).await.expect("select"); - assert_eq!(calls.lock().len(), 1); - assert_eq!(observed.lock()[1].shared_beyond_device_blocks, vec![0]); - } - - #[tokio::test] - async fn session_context_reaches_worker_selection() { - use super::super::types::SelectionSessionContext; - - let (factory, observed) = capturing_policy_factory(); - let core = core_with_host_and_policy(SelectionHost::default(), Some(factory)); - core.upsert_worker(worker(1)).await.expect("worker upsert"); - - let mut request = select_request(); - request.session_id = Some("ignored-legacy".to_string()); - request.session_context = Some(SelectionSessionContext { - session_id: "child-session".to_string(), - parent_session_id: Some("root-session".to_string()), - session_final: Some(true), - input_trigger: Some(super::super::types::SelectionInputTrigger::ToolResult), - }); - core.select(request).await.expect("select"); - - let mut request = reserve_request("legacy-session-reservation"); - request.session_id = Some("legacy-only".to_string()); - core.select_and_reserve(request) - .await - .expect("select and reserve"); - - let observations = observed.lock(); - let context = observations[0] - .session_context - .as_ref() - .expect("structured session context"); - assert_eq!(context.session_id(), "child-session"); - assert_eq!(context.parent_session_id(), Some("root-session")); - assert_eq!(context.session_final(), Some(true)); - assert_eq!( - context.input_trigger(), - Some(crate::scheduling::WorkerSelectionInputTrigger::ToolResult) - ); - - let legacy = observations[1] - .session_context - .as_ref() - .expect("legacy session context"); - assert_eq!(legacy.session_id(), "legacy-only"); - assert_eq!(legacy.parent_session_id(), None); - } - - #[tokio::test] - async fn full_affinity_table_routes_without_pinning() { - let core = SelectionCore::try_new_local( - test_config(false), - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"); - core.upsert_worker(worker(1)).await.expect("worker upsert"); - let entry = core - .entry(&RoutingPartitionId::new("model", "default")) - .expect("entry"); - let table = SessionAffinity::new_with_limits(Duration::from_secs(60), 1, 256) - .expect("affinity table"); - assert!(entry.affinity.set(table).is_ok()); - - for (selection_id, session_id) in [("first", "s1"), ("second", "s2")] { - let mut request = reserve_request(selection_id); - request.session_id = Some(session_id.to_string()); - core.select_and_reserve(request) - .await - .expect("a full affinity table must not fail selection"); - } - assert_eq!( - core.reservation_index - .read() - .values() - .filter(|r| r._affinity_lease.is_some()) - .count(), - 1 - ); - } - - #[tokio::test] - async fn injected_availability_provider_restricts_selection() { - let available: Arc>>>> = - Arc::new(parking_lot::Mutex::new(None)); - let provider_state = Arc::clone(&available); - let core = core_with_host(SelectionHost { - load: HostLoad { - available_workers: Some(Arc::new(move || provider_state.lock().clone())), - ..HostLoad::default() - }, - ..SelectionHost::default() - }); - core.upsert_worker(worker(1)).await.expect("worker upsert"); - core.upsert_worker(worker(2)).await.expect("worker upsert"); - - for only in [1, 2] { - *available.lock() = Some(Arc::new(HashSet::from([only]))); - for _ in 0..4 { - let response = core.select(select_request()).await.expect("select"); - assert_eq!(response.worker_id, only); - } - } - } - - #[tokio::test] - async fn injected_overload_provider_excludes_worker() { - let core = core_with_host(SelectionHost { - load: HostLoad { - overloaded_workers: Some(Arc::new(|| Some(HashSet::from([1])))), - ..HostLoad::default() - }, - ..SelectionHost::default() - }); - core.upsert_worker(worker(1)).await.expect("worker upsert"); - core.upsert_worker(worker(2)).await.expect("worker upsert"); - - for _ in 0..4 { - let response = core.select(select_request()).await.expect("select"); - assert_eq!(response.worker_id, 2); - } - } - - #[test] - fn shutdown_keeps_parent_alive() { - let parent = CancellationToken::new(); - let core = SelectionCore::try_new_local( - test_config(false), - 1, - parent.clone(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"); - - core.shutdown(); - - assert!(core.cancel_token.is_cancelled()); - assert!(!parent.is_cancelled()); - } - - #[tokio::test] - async fn multi_rank_worker_with_replay_endpoint_is_incomplete() { - let core = SelectionCore::try_new_local( - test_config(true), - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"); - - let record = core - .upsert_worker(WorkerRequest { - data_parallel_size: Some(2), - kv_events_endpoints: HashMap::from([ - (0, "tcp://127.0.0.1:5557".to_string()), - (1, "tcp://127.0.0.1:5558".to_string()), - ]), - replay_endpoint: Some("tcp://127.0.0.1:5600".to_string()), - ..worker(1) - }) - .await - .expect("worker upsert"); - assert_eq!(record.lifecycle, WorkerLifecycle::Incomplete, "{record:?}"); - assert!(!core.indexer_registry.has_worker(1)); - } - - #[tokio::test] - async fn reupsert_recreates_a_listener_lost_to_a_cancelled_update() { - use crate::indexer::KvIndexerInterface; - - let core = SelectionCore::try_new_local( - test_config(true), - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"); - core.upsert_worker(worker_with_kv_events(1)) - .await - .expect("worker upsert"); - let entry = core - .entry(&RoutingPartitionId::new("model", "default")) - .expect("entry"); - let hashes = [11u64, 12]; - entry - .indexer - .apply_event_routed(store_event(1, 0, 1, &[], &hashes, StorageTier::Device)) - .await - .unwrap(); - if let Indexer::Single { primary, .. } = &entry.indexer { - let _ = primary.flush().await; - } - // An endpoint update cancelled right after removing the listener leaves - // the catalog record schedulable, the listener gone, and its blocks indexed. - core.indexer_registry.forget_listener(1, 0); - - let record = core - .upsert_worker(worker_with_kv_events(1)) - .await - .expect("worker re-upsert"); - assert_eq!(record.lifecycle, WorkerLifecycle::Schedulable, "{record:?}"); - assert!(core.indexer_registry.has_listener(1, 0)); - // The purge is queued to the indexer thread; wait for it to land. - tokio::time::timeout(Duration::from_secs(2), async { - loop { - let overlap = entry - .indexer - .find_matches(hashes.iter().copied().map(LocalBlockHash).collect()) - .await - .expect("find matches"); - if overlap.scores.is_empty() { - return; - } - tokio::task::yield_now().await; - } - }) - .await - .expect("stale blocks survived the re-upsert"); - } - - #[tokio::test] - async fn shutdown_cancels_listeners() { - let parent = CancellationToken::new(); - let core = SelectionCore::try_new_local( - test_config(true), - 1, - parent, - SelectionCacheConfig::default(), - ) - .expect("valid test config"); - - let record = core - .upsert_worker(worker_with_kv_events(1)) - .await - .expect("worker upsert"); - assert_eq!(record.lifecycle, WorkerLifecycle::Schedulable); - assert_eq!(core.indexer_registry.listener_cancelled(1, 0), Some(false)); - - core.shutdown(); - assert_eq!(core.indexer_registry.listener_cancelled(1, 0), Some(true)); - } - - #[tokio::test] - async fn upsert_moves_global_worker_id_between_routing_groups() { - let core = SelectionCore::try_new_local( - test_config(true), - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"); - let mut group_a = worker_with_kv_events(1); - group_a.routing_group = "group-a".to_string(); - core.upsert_worker(group_a).await.expect("group A upsert"); - assert_eq!( - core.indexer_registry - .list_filtered(Some("model"), Some("group-a")) - .len(), - 1 - ); - - let mut group_b = worker_with_kv_events(1); - group_b.routing_group = "group-b".to_string(); - core.upsert_worker(group_b).await.expect("group B upsert"); - - assert!(core.list_workers(Some("model"), Some("group-a")).is_empty()); - assert_eq!(core.list_workers(Some("model"), Some("group-b")).len(), 1); - assert!( - core.indexer_registry - .list_filtered(Some("model"), Some("group-a")) - .is_empty() - ); - assert_eq!( - core.indexer_registry - .list_filtered(Some("model"), Some("group-b")) - .len(), - 1 - ); - - let mut select_a = select_request(); - select_a.routing_group = "group-a".to_string(); - assert!(matches!( - core.select(select_a).await, - Err(SelectionError::NotReady(_)) - )); - let mut select_b = select_request(); - select_b.routing_group = "group-b".to_string(); - assert_eq!(core.select(select_b).await.unwrap().worker_id, 1); - - core.delete_worker(1).await.expect("delete group B worker"); - assert!( - core.indexer_registry - .list_filtered(Some("model"), Some("group-b")) - .is_empty() - ); - } - - #[tokio::test] - async fn shutdown_reports_not_ready_and_rejects_new_work() { - let core = SelectionCore::try_new_local( - test_config(false), - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"); - core.upsert_worker(worker(1)).await.expect("worker upsert"); - assert!(core.ready().ready); - - core.shutdown(); - - let ready = core.ready(); - assert!(!ready.ready); - assert_eq!(ready.schedulable_workers, 1); - - let upsert_error = core - .upsert_worker(worker(2)) - .await - .expect_err("upsert should fail after shutdown"); - assert_shutdown_error(upsert_error); - - let patch = serde_json::from_value(serde_json::json!({ - "endpoint": "http://worker-1:9000" - })) - .expect("worker patch"); - let patch_error = core - .patch_worker(1, patch) - .await - .expect_err("patch should fail after shutdown"); - assert_shutdown_error(patch_error); - - let select_error = core - .select(select_request()) - .await - .expect_err("selection should fail after shutdown"); - assert_shutdown_error(select_error); - - let reservation_error = core - .create_reservation(ReservationRequest { - model_name: "model".to_string(), - routing_group: "default".to_string(), - selection_id: "res-after-shutdown".to_string(), - worker_id: Some(1), - dp_rank: None, - prompt: prompt(), - router_config_override: None, - expected_output_tokens: None, - effective_prefill_tokens: None, - track_prefill_tokens: None, - }) - .await - .expect_err("reservation should fail after shutdown"); - assert_shutdown_error(reservation_error); - - assert_eq!(core.list_workers(None, None).len(), 1); - assert_eq!(core.loads(None, None).len(), 1); - let deleted = core - .delete_worker(1) - .await - .expect("delete should remain available after shutdown"); - assert_eq!(deleted.lifecycle, WorkerLifecycle::Unschedulable); - } - - #[tokio::test] - async fn queued_selection_errors_on_shutdown() { - let mut config = test_config(false); - config.router_queue_threshold = Some(0.0); - let core = Arc::new( - SelectionCore::try_new_local( - config, - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"), - ); - - let record = core.upsert_worker(worker(1)).await.expect("worker upsert"); - assert_eq!(record.lifecycle, WorkerLifecycle::Schedulable); - core.select_and_reserve(reserve_request("res-a")) - .await - .expect("initial reservation"); - - let queued_core = core.clone(); - let queued = tokio::spawn(async move { queued_core.select(select_request()).await }); - wait_for_pending_selection(&core).await; - - core.shutdown(); - let err = tokio::time::timeout(Duration::from_secs(1), queued) - .await - .expect("queued selection timed out") - .expect("queued selection task panicked") - .expect_err("queued selection should fail"); - - assert!(matches!( - err, - SelectionError::Scheduler(KvSchedulerError::SubscriberShutdown) - )); - } - - #[tokio::test] - async fn booking_is_freed_when_selected_worker_drained_while_queued() { - let mut config = test_config(false); - config.router_queue_threshold = Some(0.0); - let core = Arc::new( - SelectionCore::try_new_local( - config, - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"), - ); - core.upsert_worker(worker(1)).await.expect("worker upsert"); - let key = RoutingPartitionId::new("model", "default"); - let entry = core.entry(&key).expect("entry"); - core.select_and_reserve(reserve_request("res-a")) - .await - .expect("initial reservation"); - - let queued_core = core.clone(); - let queued = tokio::spawn(async move { - queued_core - .select_and_reserve(reserve_request("queued")) - .await - }); - wait_for_pending_selection(&core).await; - // First half of `delete_worker`: the worker drains while the request waits. - core.catalog - .set_lifecycle(1, WorkerLifecycle::Draining, Vec::new()); - core.free_reservation("res-a").await.expect("free res-a"); - - let err = tokio::time::timeout(Duration::from_secs(2), queued) - .await - .expect("queued selection timed out") - .expect("task panicked") - .expect_err("drained worker is not schedulable"); - assert!( - matches!(&err, SelectionError::Internal(m) if m.contains("no longer schedulable")), - "{err:?}" - ); - wait_until("booking release", || !entry.scheduler.has_request("queued")).await; - } - - /// The embedded host dispatches by worker id, so a worker that drains - /// while the request queues still comes back selected: the transport - /// reports the departure and migration takes over. The endpoint is the - /// only thing missing. - #[tokio::test] - async fn lease_admission_keeps_a_selection_whose_worker_drained_while_queued() { - let mut config = test_config(false); - config.router_queue_threshold = Some(0.0); - let core = Arc::new( - SelectionCore::try_new_local( - config, - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"), - ); - core.upsert_worker(worker(1)).await.expect("worker upsert"); - let key = RoutingPartitionId::new("model", "default"); - let entry = core.entry(&key).expect("entry"); - core.select_and_reserve(reserve_request("res-a")) - .await - .expect("initial reservation"); - - let queued_core = core.clone(); - let queued = tokio::spawn(async move { - let req = reserve_request("queued"); - let run = queued_core - .run_selection(SelectionOperation { - key: RoutingPartitionId::new("model", "default"), - prompt: req.prompt.view(), - router_config_override: None, - expected_output_tokens: None, - priority_jump: 0.0, - strict_priority: 0, - policy_class: None, - session_context: None, - session: SessionBinding::None, - affinity_target: None, - pinned_worker: None, - allowed_worker_ids: None, - routing_constraints: RoutingConstraints::default(), - admission: SelectionAdmission::Lease { - request_id: "queued".to_string(), - }, - track_active_blocks: false, - return_routing_hashes: false, - replay_id: None, - }) - .await; - match run.result { - Ok(SelectionOutcome::Selected(selected)) => ( - selected.response.best_worker.worker_id, - selected.endpoint, - selected.lease.is_some(), - ), - Ok(SelectionOutcome::QueueRejected { .. }) => panic!("queue rejected"), - Err(error) => panic!("lease selection failed: {error:?}"), - } - }); - wait_for_pending_selection(&core).await; - core.catalog - .set_lifecycle(1, WorkerLifecycle::Draining, Vec::new()); - core.free_reservation("res-a").await.expect("free res-a"); - - let (worker_id, endpoint, leased) = tokio::time::timeout(Duration::from_secs(2), queued) - .await - .expect("queued selection timed out") - .expect("task panicked"); - assert_eq!(worker_id, 1); - assert_eq!(endpoint, None); - assert!(leased, "lease admission returns the booking's lease"); - // The lease was dropped with the selection, so the booking is gone. - wait_until("booking release", || !entry.scheduler.has_request("queued")).await; - } - - #[tokio::test] - async fn dropped_selection_future_frees_its_booking() { - let core = SelectionCore::try_new_local( - test_config(false), - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"); - core.upsert_worker(worker(1)).await.expect("worker upsert"); - let key = RoutingPartitionId::new("model", "default"); - let entry = core.entry(&key).expect("entry"); - - // Drive the selection by hand so the actor's response is delivered but - // never consumed: poll until the booking exists, then drop the future. - let mut selection = Box::pin(core.select_and_reserve(reserve_request("dropped"))); - let mut context = std::task::Context::from_waker(std::task::Waker::noop()); - wait_until("scheduler booking", || { - if entry.scheduler.has_request("dropped") { - return true; - } - assert!( - selection.as_mut().poll(&mut context).is_pending(), - "selection completed before the booking was observed" - ); - false - }) - .await; - assert!( - core.reservation_index - .read() - .get("dropped") - .is_some_and(|reservation| reservation.booking.is_none()), - "the id is claimed while its booking is in flight" - ); - drop(selection); - - wait_until("booking release", || { - !entry.scheduler.has_request("dropped") - }) - .await; - assert!(core.reservation_index.read().get("dropped").is_none()); - } - - #[tokio::test] - async fn same_selection_id_in_two_partitions_is_a_conflict() { - let core = SelectionCore::try_new_local( - test_config(false), - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"); - for (worker_id, routing_group) in [(1, "group-a"), (2, "group-b")] { - let mut request = worker(worker_id); - request.routing_group = routing_group.to_string(); - core.upsert_worker(request).await.expect("worker upsert"); - } - let mut first = reserve_request("shared"); - first.routing_group = "group-a".to_string(); - core.select_and_reserve(first).await.expect("first booking"); - - let mut second = reserve_request("shared"); - second.routing_group = "group-b".to_string(); - let err = core - .select_and_reserve(second) - .await - .expect_err("a live id cannot be booked again"); - assert!(matches!(err, SelectionError::Conflict(_)), "{err:?}"); - - let (entry, _) = core - .indexed_booking("shared") - .expect("first booking indexed"); - assert_eq!(entry.key.routing_group, "group-a"); - assert!(entry.scheduler.has_request("shared")); - assert!( - !core - .entry(&RoutingPartitionId::new("model", "group-b")) - .expect("entry") - .scheduler - .has_request("shared") - ); - } - - #[tokio::test] - async fn explicit_reservation_of_a_live_id_is_a_conflict() { - let core = SelectionCore::try_new_local( - test_config(false), - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"); - core.upsert_worker(worker(1)).await.expect("worker upsert"); - core.select_and_reserve(reserve_request("taken")) - .await - .expect("booking"); - let err = core - .create_reservation(ReservationRequest { - model_name: "model".to_string(), - routing_group: "default".to_string(), - selection_id: "taken".to_string(), - worker_id: Some(1), - dp_rank: None, - prompt: prompt(), - router_config_override: None, - expected_output_tokens: None, - effective_prefill_tokens: None, - track_prefill_tokens: None, - }) - .await - .expect_err("explicit booking of a live id"); - assert!(matches!(err, SelectionError::Conflict(_)), "{err:?}"); - } - - #[test] - fn index_observer_replaces_stale_and_claimed_rows_of_its_partition() { - use crate::scheduling::AttemptId; - let index = Arc::new(RwLock::new(HashMap::new())); - let partition = RoutingPartitionId::new("model", "default"); - let observer = ReservationIndexObserver { - index: Arc::clone(&index), - partition: partition.clone(), - host: None, - }; - let booking = |attempt: u64| SchedulerBookingDescriptor { - request_id: "shared".to_string(), - worker: WorkerWithDpRank::new(1, 0), - attempt_id: AttemptId::new(attempt), - }; - let row = |partition: &RoutingPartitionId, booking| Reservation { - partition: partition.clone(), - booking, - _affinity_lease: None, - }; - - // A stale row (its booking expired before the sweep ran) yields to the mirror. - index - .write() - .insert("shared".to_string(), row(&partition, Some(booking(1)))); - observer.admitted(booking(2)); - assert_eq!(index.read()["shared"].booking, Some(booking(2))); - - // A completion for the replaced booking leaves the live mirror alone. - observer.completed(&booking(1)); - assert!(index.read().contains_key("shared")); - observer.completed(&booking(2)); - assert!(!index.read().contains_key("shared")); - - // A claim whose local booking is about to fail also yields, and dropping - // that claim keeps the mirror. - let claim = ReservationClaim { - index: &index, - selection_id: "shared".to_string(), - armed: true, - }; - index - .write() - .insert("shared".to_string(), row(&partition, None)); - observer.admitted(booking(3)); - drop(claim); - assert_eq!(index.read()["shared"].booking, Some(booking(3))); - - // Another partition's row is never replaced. - let other = RoutingPartitionId::new("model", "other"); - index - .write() - .insert("shared".to_string(), row(&other, Some(booking(4)))); - observer.admitted(booking(5)); - assert_eq!(index.read()["shared"].booking, Some(booking(4))); - } - - #[tokio::test] - async fn free_of_an_in_flight_reservation_is_not_found() { - let mut config = test_config(false); - config.router_queue_threshold = Some(0.0); - let core = Arc::new( - SelectionCore::try_new_local( - config, - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"), - ); - core.upsert_worker(worker(1)).await.expect("worker upsert"); - core.select_and_reserve(reserve_request("res-a")) - .await - .expect("initial reservation"); - let queued_core = core.clone(); - let queued = tokio::spawn(async move { - queued_core - .select_and_reserve(reserve_request("queued")) - .await - }); - wait_for_pending_selection(&core).await; - - // A free racing the in-flight booking neither frees nor evicts it. - let err = core - .free_reservation("queued") - .await - .expect_err("in-flight id is not a reservation yet"); - assert!(matches!(err, SelectionError::NotFound(_)), "{err:?}"); - assert!(core.reservation_index.read().contains_key("queued")); - - core.free_reservation("res-a").await.expect("free res-a"); - tokio::time::timeout(Duration::from_secs(2), queued) - .await - .expect("queued selection timed out") - .expect("task panicked") - .expect("queued selection books"); - core.free_reservation("queued").await.expect("free queued"); - assert!(core.reservation_index.read().is_empty()); - } - - #[tokio::test] - async fn prefill_complete_is_idempotent_for_a_live_booking() { - let mut config = test_config(false); - config.router_track_prefill_tokens = true; - let core = SelectionCore::try_new_local( - config, - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"); - core.upsert_worker(worker(1)).await.expect("worker upsert"); - core.select_and_reserve(reserve_request("live")) - .await - .expect("booking"); - core.prefill_complete("live").await.expect("first mark"); - core.prefill_complete("live") - .await - .expect("a repeated mark on a live booking is not an error"); - let (entry, _) = core.indexed_booking("live").expect("still indexed"); - assert!(entry.scheduler.has_request("live")); - core.free_reservation("live").await.expect("free"); - } - - #[tokio::test] - async fn mirrored_replica_bookings_are_indexed_until_freed() { - let (outbound_tx, _outbound_rx) = mpsc::channel(16); - let (inbound_tx, inbound_rx) = mpsc::channel(16); - let channels = parking_lot::Mutex::new(Some(HostReplicaChannels { - outbound: Some(outbound_tx), - inbound_tx: inbound_tx.clone(), - inbound_rx, - process_id: 7, - })); - let core = core_with_host(SelectionHost { - replication: HostReplication { - channels: Some(Arc::new(move |_| channels.lock().take())), - ..HostReplication::default() - }, - ..SelectionHost::default() - }); - core.upsert_worker(worker(1)).await.expect("worker upsert"); - // Replica events for a worker the scheduler has not registered yet are - // dropped; the first local booking registers it. - core.select_and_reserve(reserve_request("warm")) - .await - .expect("warm booking"); - core.free_reservation("warm").await.expect("free warm"); - let entry = core - .entry(&RoutingPartitionId::new("model", "default")) - .expect("entry"); - let peer_event = |request_id: &str, data| ActiveSequenceEvent { - request_id: request_id.to_string(), - worker: WorkerWithDpRank::new(1, 0), - data, - router_id: 99, - lora_name: None, - }; - let add = |request_id: &str| { - peer_event( - request_id, - ActiveSequenceEventData::AddRequest { - token_sequence: Some(vec![1, 2]), - track_prefill_tokens: false, - expected_output_tokens: None, - prefill_load_hint: None, - }, - ) - }; - - inbound_tx.send(add("peer-a")).await.expect("send"); - inbound_tx.send(add("peer-b")).await.expect("send"); - wait_until("mirrored bookings indexed", || { - core.indexed_booking("peer-a").is_some() && core.indexed_booking("peer-b").is_some() - }) - .await; - - // A lifecycle call on a mirrored booking resolves through the index. - core.free_reservation("peer-a") - .await - .expect("free mirrored"); - assert!(!entry.scheduler.has_request("peer-a")); - assert!(core.indexed_booking("peer-a").is_none()); - - // The peer freeing its own booking removes the mirror. - inbound_tx - .send(peer_event("peer-b", ActiveSequenceEventData::Free)) - .await - .expect("send"); - wait_until("mirror removed", || { - core.indexed_booking("peer-b").is_none() - }) - .await; - assert!(!entry.scheduler.has_request("peer-b")); - } - - #[tokio::test] - async fn lifecycle_operations_resolve_through_the_index() { - let mut config = test_config(false); - config.router_track_prefill_tokens = true; - let core = SelectionCore::try_new_local( - config, - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"); - - for (worker_id, routing_group) in [(1, "group-a"), (2, "group-b")] { - let mut request = worker(worker_id); - request.routing_group = routing_group.to_string(); - core.upsert_worker(request).await.expect("worker upsert"); - } - - let entries = core.initialized_entries(); - assert_eq!(entries.len(), 2); - let target = &entries[1]; - let target_group = target.key.routing_group.clone(); - let target_worker = *target - .workers_tx - .borrow() - .keys() - .next() - .expect("target worker"); - - let mut request = reserve_request("later-entry-reservation"); - request.routing_group = target_group.clone(); - core.select_and_reserve(request) - .await - .expect("reserve in later entry"); - - let load = || { - core.loads(Some("model"), Some(&target_group))[0] - .loads - .iter() - .find(|load| load.worker_id == target_worker) - .expect("target load") - .potential_prefill_tokens - }; - assert_eq!(load(), 4); - - core.prefill_complete("later-entry-reservation") - .await - .expect("complete prefill in later entry"); - assert_eq!(load(), 0); - - core.free_reservation("later-entry-reservation") - .await - .expect("free reservation in later entry"); - assert!(matches!( - core.add_output_block("later-entry-reservation", None), - Err(SelectionError::NotFound(_)) - )); - } - - #[tokio::test] - async fn advisory_select_reports_worker_load_and_busy_evaluation() { - let mut config = test_config(false); - config.conditional_disagg_prefill_busy_threshold = Some(0.5); - config.conditional_disagg_decode_busy_threshold = Some(0.0); - let core = SelectionCore::try_new_local( - config, - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"); - let mut request = worker(1); - request.total_kv_blocks = Some(1000); - core.upsert_worker(request).await.expect("worker upsert"); - - // Admitted (queued) select: decode evaluation comes from the catalog's - // total_kv_blocks; no load snapshot is taken. - let response = core.select(select_request()).await.expect("select"); - assert!(response.potential_decode_blocks > 0); - assert_eq!( - response.decode_busy, - Some(true), - "threshold 0.0 is always exceeded" - ); - assert!(response.worker_load.is_none()); - - // Advisory select: same decode evaluation plus the projected load. - let mut request = select_request(); - request.advisory = true; - let response = core.select(request).await.expect("advisory select"); - assert!(response.potential_decode_blocks > 0); - assert_eq!(response.decode_busy, Some(true)); - let load = response.worker_load.expect("advisory load"); - assert_eq!(load.total_kv_blocks, Some(1000)); - assert_eq!(load.prefill_token_capacity, 1024); - assert_eq!(load.active_prefill_tokens, 0); - assert_eq!(load.prefill_busy, Some(false)); - - // Advisory selection does not book. - assert!(core.reservation_index.read().is_empty()); - assert_eq!( - core.loads(Some("model"), Some("default"))[0].loads[0].potential_prefill_tokens, - 0 - ); - } - - #[tokio::test] - async fn busy_evaluation_is_absent_without_thresholds_or_capacity() { - let core = SelectionCore::try_new_local( - test_config(false), - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"); - core.upsert_worker(worker(1)).await.expect("worker upsert"); - let mut request = select_request(); - request.advisory = true; - let response = core.select(request).await.expect("advisory select"); - assert_eq!(response.decode_busy, None); - let load = response.worker_load.expect("advisory load"); - assert_eq!(load.total_kv_blocks, None); - assert_eq!(load.prefill_busy, None); - } - - #[tokio::test] - async fn reservation_index_tracks_bookings_until_freed() { - let core = SelectionCore::try_new_local( - test_config(false), - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"); - for (worker_id, routing_group) in [(1, "group-a"), (2, "group-b")] { - let mut request = worker(worker_id); - request.routing_group = routing_group.to_string(); - core.upsert_worker(request).await.expect("worker upsert"); - } - let key_b = RoutingPartitionId::new("model", "group-b"); - - // select_and_reserve records the booking's partition. - let mut request = reserve_request("booked"); - request.routing_group = "group-b".to_string(); - core.select_and_reserve(request).await.expect("reserve"); - assert_eq!( - core.reservation_index - .read() - .get("booked") - .map(|r| &r.partition), - Some(&key_b) - ); - assert_eq!( - core.indexed_booking("booked") - .expect("indexed booking") - .0 - .key, - key_b - ); - - // The explicit reservation path records too. - let mut request = select_request(); - request.routing_group = "group-b".to_string(); - request.selection_id = Some("cached".to_string()); - core.select(request).await.expect("select"); - assert!(core.reservation_index.read().get("cached").is_none()); - core.create_reservation(ReservationRequest { - model_name: "model".to_string(), - routing_group: "group-b".to_string(), - selection_id: "cached".to_string(), - worker_id: None, - dp_rank: None, - prompt: PromptRequest::default(), - router_config_override: None, - expected_output_tokens: None, - effective_prefill_tokens: None, - track_prefill_tokens: None, - }) - .await - .expect("cached reservation"); - assert_eq!( - core.reservation_index - .read() - .get("cached") - .map(|r| &r.partition), - Some(&key_b) - ); - - // Lifecycle calls still resolve, and free drops the index entry. - core.prefill_complete("booked") - .await - .expect("prefill complete"); - core.free_reservation("booked").await.expect("free"); - assert!(core.reservation_index.read().get("booked").is_none()); - core.free_reservation("cached").await.expect("free"); - assert!(core.reservation_index.read().is_empty()); - - // Unknown ids fall back to the full scan and stay unindexed. - assert!(matches!( - core.prefill_complete("never-booked").await, - Err(SelectionError::NotFound(_)) - )); - assert!(core.reservation_index.read().is_empty()); - } - - #[tokio::test] - async fn reservation_index_sweep_drops_bookings_released_out_of_band() { - let core = SelectionCore::try_new_local( - test_config(false), - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"); - core.upsert_worker(worker(1)).await.expect("worker upsert"); - core.select_and_reserve(reserve_request("live")) - .await - .expect("reserve live"); - core.select_and_reserve(reserve_request("stale")) - .await - .expect("reserve stale"); - assert_eq!(core.reservation_index.read().len(), 2); - assert_eq!( - sweep_reservation_index(&core.entries, &core.reservation_index), - 0 - ); - - // Release directly through the scheduler, as force-expiry would. - let entry = core - .entry(&RoutingPartitionId::new("model", "default")) - .expect("entry"); - entry.scheduler.free("stale").await.expect("scheduler free"); - - assert_eq!( - sweep_reservation_index(&core.entries, &core.reservation_index), - 1 - ); - let index = core.reservation_index.read(); - assert_eq!(index.len(), 1); - assert!(index.contains_key("live")); - } - - #[tokio::test(flavor = "current_thread")] - async fn lifecycle_lookup_does_not_nest_reservation_index_inside_entries() { - // Three parties: a sweep holding `entries` and wanting `reservation_index`, - // a lifecycle call, and a partition creation queued on `entries.write()`. - let core = Arc::new( - SelectionCore::try_new_local( - test_config(false), - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"), - ); - core.upsert_worker(worker(1)).await.expect("worker upsert"); - core.select_and_reserve(reserve_request("live")) - .await - .expect("reserve live"); - let deadline = Instant::now() + Duration::from_secs(5); - - let sweep_entries = core.entries.read(); - let writer = { - let core = Arc::clone(&core); - std::thread::spawn(move || { - core.entries - .write() - .entry(RoutingPartitionId::new("other", "default")) - .or_insert_with(|| Arc::new(OnceCell::new())); - }) - }; - while core.entries.try_read().is_some() { - assert!(Instant::now() < deadline, "partition writer never queued"); - std::thread::yield_now(); - } - let lifecycle_started = Arc::new(AtomicBool::new(false)); - let lifecycle = { - let core = Arc::clone(&core); - let started = Arc::clone(&lifecycle_started); - std::thread::spawn(move || { - started.store(true, Ordering::Release); - core.indexed_booking("live").is_some() - }) - }; - while !lifecycle_started.load(Ordering::Acquire) { - assert!(Instant::now() < deadline, "lifecycle thread never started"); - std::thread::yield_now(); - } - sleep(Duration::from_millis(50)); - - // With the index guard held across `entries.read()` this never acquires. - drop( - core.reservation_index - .try_write_for(Duration::from_secs(2)) - .expect("reservation index must not be held by a blocked lifecycle call"), - ); - drop(sweep_entries); - writer.join().expect("partition writer"); - assert!(lifecycle.join().expect("lifecycle lookup")); - } - - #[tokio::test(flavor = "current_thread")] - async fn queued_selection_returns_refreshed_overlap_snapshot() { - let mut config = test_config(false); - config.router_queue_threshold = Some(0.0); - let core = Arc::new( - SelectionCore::try_new_local( - config, - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"), - ); - - for worker_id in [1, 2] { - let mut request = worker(worker_id); - request.max_num_batched_tokens = Some(8); - core.upsert_worker(request).await.expect("worker upsert"); - } - let key = RoutingPartitionId::new("model", "default"); - let entry = core.entry(&key).expect("entry"); - entry - .indexer - .apply_event_routed(store_event(1, 0, 1, &[], &[11], StorageTier::Device)) - .await - .unwrap(); - entry.indexer.dump_events().await.expect("flush indexer"); - - for worker_id in [1, 2] { - core.create_reservation(ReservationRequest { - model_name: "model".to_string(), - routing_group: "default".to_string(), - selection_id: format!("occupy-{worker_id}"), - worker_id: Some(worker_id), - dp_rank: Some(0), - prompt: PromptRequest { - token_ids: None, - mm_routing_info: None, - block_mm_infos: None, - block_hashes: None, - sequence_hashes: Some(vec![1, 2]), - isl_tokens: Some(8), - lora_name: None, - cache_namespace: None, - is_eagle: None, - }, - router_config_override: None, - expected_output_tokens: None, - effective_prefill_tokens: Some(8), - track_prefill_tokens: None, - }) - .await - .expect("occupy worker"); - } - - let queued_core = Arc::clone(&core); - let queued = tokio::spawn(async move { - queued_core - .select_and_reserve(SelectAndReserveRequest { - model_name: "model".to_string(), - routing_group: "default".to_string(), - selection_id: Some("refresh-selection".to_string()), - prompt: PromptRequest { - token_ids: None, - mm_routing_info: None, - block_mm_infos: None, - block_hashes: Some(vec![11, 12]), - sequence_hashes: Some(vec![101, 102]), - isl_tokens: Some(8), - lora_name: None, - cache_namespace: None, - is_eagle: None, - }, - router_config_override: None, - expected_output_tokens: None, - priority_jump: None, - strict_priority: None, - session_id: None, - session_context: None, - affinity_target: None, - pinned_worker: None, - allowed_worker_ids: None, - routing_constraints: RoutingConstraints::default(), - }) - .await - }); - wait_for_pending_selection(&core).await; - - entry - .indexer - .apply_event_routed(store_event(2, 0, 1, &[], &[11, 12], StorageTier::Device)) - .await - .unwrap(); - entry.indexer.dump_events().await.expect("flush indexer"); - // Freeze time only after async setup so background timers cannot expire the fixture. - tokio::time::pause(); - tokio::time::advance(Duration::from_secs(11)).await; - core.free_reservation("occupy-2") - .await - .expect("release worker 2"); - - let response = queued.await.expect("selection task").expect("selection"); - assert_eq!(response.worker_id, 2); - assert_eq!(response.effective_prefill_tokens, 0); - assert_eq!(response.overlap.gpu, 8); - assert_eq!(response.overlap.cpu, 8); - assert_eq!(response.overlap.disk, 8); - assert_eq!(response.overlap.dp, HashMap::from([("0".to_string(), 8)])); - } - - fn core_with_session_affinity_mode(mode: SessionAffinityMode) -> SelectionCore { - let config = test_config(false); - let tracking_hash = Arc::new( - TrackingHashContext::from_config(&config).expect("valid tracking hash configuration"), - ); - let indexer_policy = IndexerPolicy::from_router_config(&config).expect("indexer policy"); - SelectionCore::new_inner( - config, - 1, - CancellationToken::new(), - None, - None, - SelectionHost::default(), - WorkerType::Aggregated, - true, - SelectionCacheConfig::default(), - tracking_hash, - indexer_policy, - Some(SessionAffinityConfig::new(Duration::from_secs(10)).with_mode(mode)), - ) - } - - fn core_with_session_affinity() -> SelectionCore { - core_with_session_affinity_mode(SessionAffinityMode::Hard) - } - - fn bound_worker(core: &SelectionCore, session_id: &str) -> Option { - core.entry(&RoutingPartitionId::new("model", "default")) - .and_then(|entry| entry.affinity.get().cloned()) - .and_then(|table| table.query_target(session_id, None).expect("query")) - .map(|target| target.worker_id) - } - - fn session_reservation(selection_id: &str, session_id: &str) -> SelectAndReserveRequest { - let mut request = reserve_request(selection_id); - request.session_id = Some(session_id.to_string()); - request - } - - #[tokio::test] - async fn departed_session_worker_reinitializes_the_session() { - let core = core_with_session_affinity(); - core.upsert_worker(worker(1)).await.expect("worker upsert"); - core.upsert_worker(worker(2)).await.expect("worker upsert"); - let first = core - .select_and_reserve(session_reservation("r1", "s")) - .await - .expect("first booking"); - core.free_reservation("r1").await.expect("free"); - core.delete_worker(first.worker_id).await.expect("delete"); - - let second = core - .select_and_reserve(session_reservation("r2", "s")) - .await - .expect("session must move off a departed worker"); - assert_ne!(second.worker_id, first.worker_id); - assert_eq!(bound_worker(&core, "s"), Some(second.worker_id)); - } - - #[tokio::test] - async fn session_worker_departing_after_the_hold_reinitializes_the_session() { - let core = core_with_session_affinity(); - core.upsert_worker(worker(1)).await.expect("worker upsert"); - core.upsert_worker(worker(2)).await.expect("worker upsert"); - let first = core - .select_and_reserve(session_reservation("r1", "s")) - .await - .expect("first booking"); - core.free_reservation("r1").await.expect("free"); - let key = RoutingPartitionId::new("model", "default"); - - // One poll takes the hold (the bound worker still passes the check) - // and hands the request to the scheduler actor, which has not run yet. - let mut second = Box::pin(core.select_and_reserve(session_reservation("r2", "s"))); - let mut context = std::task::Context::from_waker(std::task::Waker::noop()); - assert!(second.as_mut().poll(&mut context).is_pending()); - core.catalog - .set_lifecycle(first.worker_id, WorkerLifecycle::Draining, Vec::new()); - core.publish_scheduler_config(&key).expect("publish"); - - let second = second - .await - .expect("a departure after the hold is not a client fault"); - assert_ne!(second.worker_id, first.worker_id); - assert_eq!(bound_worker(&core, "s"), Some(second.worker_id)); - } - - #[tokio::test] - async fn concurrent_holds_on_a_departed_worker_both_land_on_the_replacement() { - let table = SessionAffinity::new(Duration::from_secs(60)).expect("affinity table"); - let departed = WorkerAffinityTarget::new(1, Some(0)); - let replacement = WorkerAffinityTarget::new(2, Some(0)); - let Hold::Initialize(init) = table.acquire("s", None).await.expect("acquire") else { - panic!("fresh session must initialize"); - }; - drop( - table - .commit(Hold::Initialize(init), departed) - .expect("bind"), - ); - let ( - Hold::Bound { - lease: mut first, .. - }, - Hold::Bound { - lease: mut second, .. - }, - ) = ( - table.acquire("s", None).await.expect("first hold"), - table.acquire("s", None).await.expect("second hold"), - ) - else { - panic!("both requests hold the departed binding"); - }; - - // Both requests notice the departure, as `hold_session` does: the first - // invalidation drops the binding, the second is a no-op release. - first.invalidate(); - second.invalidate(); - let leases = [ - table.commit( - table.acquire("s", None).await.expect("re-acquire"), - replacement, - ), - table.commit( - table.acquire("s", None).await.expect("re-acquire"), - replacement, - ), - ]; - for lease in leases { - drop(lease.expect("both requests bind to the replacement")); - } - assert_eq!( - table.query_target("s", None).expect("query"), - Some(replacement) - ); - } - - fn replay_reservation(selection_id: &str) -> ReservationRequest { - ReservationRequest { - model_name: "model".to_string(), - routing_group: "default".to_string(), - selection_id: selection_id.to_string(), - worker_id: None, - dp_rank: None, - prompt: PromptRequest::default(), - router_config_override: None, - expected_output_tokens: None, - effective_prefill_tokens: None, - track_prefill_tokens: None, - } - } - - #[tokio::test] - async fn two_phase_reservation_binds_the_session() { - let core = core_with_session_affinity(); - core.upsert_worker(worker(1)).await.expect("worker upsert"); - core.upsert_worker(worker(2)).await.expect("worker upsert"); - let mut request = select_request(); - request.selection_id = Some("pending".to_string()); - request.session_id = Some("s".to_string()); - let selected = core.select(request).await.expect("select"); - assert_eq!(bound_worker(&core, "s"), None, "select alone binds nothing"); - - core.create_reservation(replay_reservation("pending")) - .await - .expect("replay booking"); - assert_eq!(bound_worker(&core, "s"), Some(selected.worker_id)); - core.free_reservation("pending").await.expect("free"); - - // The binding steers the session's next selection. - let mut request = select_request(); - request.session_id = Some("s".to_string()); - assert_eq!( - core.select(request).await.expect("select").worker_id, - selected.worker_id - ); - } - - #[tokio::test] - async fn two_phase_replay_rejects_a_worker_the_session_left() { - let core = core_with_session_affinity(); - core.upsert_worker(worker(1)).await.expect("worker upsert"); - core.upsert_worker(worker(2)).await.expect("worker upsert"); - let mut request = select_request(); - request.selection_id = Some("pending".to_string()); - request.session_id = Some("s".to_string()); - let cached = core.select(request).await.expect("select"); - // Another request binds the session elsewhere before the replay. - let other = if cached.worker_id == 1 { 2 } else { 1 }; - let mut request = session_reservation("r1", "s"); - request.allowed_worker_ids = Some(HashSet::from([other])); - core.select_and_reserve(request) - .await - .expect("bind elsewhere"); - core.free_reservation("r1").await.expect("free"); - - let err = core - .create_reservation(replay_reservation("pending")) - .await - .expect_err("hard affinity rejects the stale cached worker"); - assert!(matches!(err, SelectionError::BadRequest(_)), "{err:?}"); - let entry = core - .entry(&RoutingPartitionId::new("model", "default")) - .expect("entry"); - wait_until("rejected booking release", || { - !entry.scheduler.has_request("pending") - }) - .await; - assert!(core.reservation_index.read().is_empty()); - } - - #[tokio::test] - async fn hard_mode_rejects_dispatch_away_from_a_live_binding() { - let core = core_with_session_affinity(); - core.upsert_worker(worker(1)).await.expect("worker upsert"); - core.upsert_worker(worker(2)).await.expect("worker upsert"); - let first = core - .select_and_reserve(session_reservation("r1", "s")) - .await - .expect("first booking"); - core.free_reservation("r1").await.expect("free"); - let other = if first.worker_id == 1 { 2 } else { 1 }; - - // Steering cannot reach the bound worker, so selection lands elsewhere. - let mut request = session_reservation("r2", "s"); - request.allowed_worker_ids = Some(HashSet::from([other])); - let err = core - .select_and_reserve(request) - .await - .expect_err("hard affinity rejects a dispatch away from the binding"); - assert!(matches!(err, SelectionError::BadRequest(_)), "{err:?}"); - let entry = core - .entry(&RoutingPartitionId::new("model", "default")) - .expect("entry"); - wait_until("rejected booking release", || { - !entry.scheduler.has_request("r2") - }) - .await; - assert!(core.reservation_index.read().is_empty()); - assert_eq!(bound_worker(&core, "s"), None, "stale binding is dropped"); - - // The session's next request re-initializes. - let mut request = session_reservation("r3", "s"); - request.allowed_worker_ids = Some(HashSet::from([other])); - let third = core.select_and_reserve(request).await.expect("rebind"); - assert_eq!(third.worker_id, other); - assert_eq!(bound_worker(&core, "s"), Some(other)); - } - - #[tokio::test] - async fn soft_mode_follows_the_dispatch() { - let core = core_with_session_affinity_mode(SessionAffinityMode::Soft); - core.upsert_worker(worker(1)).await.expect("worker upsert"); - core.upsert_worker(worker(2)).await.expect("worker upsert"); - let first = core - .select_and_reserve(session_reservation("r1", "s")) - .await - .expect("first booking"); - core.free_reservation("r1").await.expect("free"); - let other = if first.worker_id == 1 { 2 } else { 1 }; - - let mut request = session_reservation("r2", "s"); - request.allowed_worker_ids = Some(HashSet::from([other])); - let second = core.select_and_reserve(request).await.expect("soft rebind"); - assert_eq!(second.worker_id, other); - assert_eq!(bound_worker(&core, "s"), Some(other)); - } - - #[tokio::test] - async fn session_stays_on_its_first_worker_across_bookings() { - let core = core_with_session_affinity(); - core.upsert_worker(worker(1)).await.expect("worker upsert"); - core.upsert_worker(worker(2)).await.expect("worker upsert"); - - let first = core - .select_and_reserve(session_reservation("r1", "chat-a")) - .await - .expect("first booking"); - for index in 0..4 { - let response = core - .select_and_reserve(session_reservation(&format!("r-{index}"), "chat-a")) - .await - .expect("booking"); - assert_eq!(response.worker_id, first.worker_id, "session must stay put"); - core.free_reservation(&format!("r-{index}")) - .await - .expect("free"); - } - // A read-only select sees the binding too. - let mut advisory = select_request(); - advisory.session_id = Some("chat-a".to_string()); - let response = core.select(advisory).await.expect("select"); - assert_eq!(response.worker_id, first.worker_id); - core.free_reservation("r1").await.expect("free"); - } - - #[tokio::test] - async fn replicated_binding_steers_a_new_session_and_frees_with_the_booking() { - let core = core_with_session_affinity(); - core.upsert_worker(worker(1)).await.expect("worker upsert"); - core.upsert_worker(worker(2)).await.expect("worker upsert"); - - core.dispatch_affinity_event(AffinityBindingEvent { - partition: RoutingPartitionId::new("model", "default"), - session_id: "chat-b".to_string(), - worker_id: 2, - dp_rank: Some(0), - sequence: 1, - writer_id: 99, - }); - let response = core - .select_and_reserve(session_reservation("r2", "chat-b")) - .await - .expect("booking"); - assert_eq!(response.worker_id, 2); - assert!( - core.reservation_index - .read() - .get("r2") - .unwrap() - ._affinity_lease - .is_some() - ); - core.free_reservation("r2").await.expect("free"); - assert!(core.reservation_index.read().is_empty()); - } - - #[tokio::test] - async fn expired_booking_releases_affinity_lease() { - let core = core_with_session_affinity(); - core.upsert_worker(worker(1)).await.unwrap(); - core.select_and_reserve(session_reservation("abandoned", "session")) - .await - .unwrap(); - let key = RoutingPartitionId::new("model", "default"); - core.entry(&key) - .unwrap() - .scheduler - .free("abandoned") - .await - .unwrap(); - assert_eq!( - sweep_reservation_index(&core.entries, &core.reservation_index), - 1 - ); - tokio::time::pause(); - tokio::time::advance(Duration::from_secs(11)).await; - assert_eq!( - core.entry(&RoutingPartitionId::new("model", "default")) - .unwrap() - .affinity - .get() - .unwrap() - .query_target("session", None) - .unwrap(), - None - ); - } - - #[tokio::test] - async fn partition_sessions_keep_independent_bindings() { - let core = core_with_session_affinity(); - core.upsert_worker(worker(1)).await.unwrap(); - let first = core - .select_and_reserve(session_reservation("first", "shared-session")) - .await - .unwrap(); - core.upsert_worker(worker(2)).await.unwrap(); - let mut other = worker(3); - other.routing_group = "other".to_string(); - core.upsert_worker(other).await.unwrap(); - let mut request = session_reservation("other", "shared-session"); - request.routing_group = "other".to_string(); - assert_eq!(core.select_and_reserve(request).await.unwrap().worker_id, 3); - assert_eq!( - core.entry(&RoutingPartitionId::new("model", "default")) - .unwrap() - .affinity - .get() - .unwrap() - .query_target("shared-session", None) - .unwrap() - .unwrap() - .worker_id, - first.worker_id - ); - let again = core - .select_and_reserve(session_reservation("again", "shared-session")) - .await - .unwrap(); - assert_eq!(again.worker_id, first.worker_id); - } - - #[tokio::test] - async fn rejoined_worker_feeds_partition_index() { - use crate::protocols::{BlockHashOptions, compute_block_hash_for_seq}; - for update_only in [false, true] { - let core = SelectionCore::try_new_local( - test_config(true), - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .unwrap(); - let request = worker_with_kv_events(1); - core.upsert_worker(request.clone()).await.unwrap(); - let key = RoutingPartitionId::new("model", "default"); - let partition = core.partition(&key).unwrap(); - let tokens: Vec = (1..=8).collect(); - let hashes: Vec = - compute_block_hash_for_seq(&tokens, 4, BlockHashOptions::default()) - .into_iter() - .map(|hash| hash.0) - .collect(); - let indexer = core - .indexer_registry - .get_indexer(&key) - .unwrap() - .indexer - .clone(); - indexer - .apply_event_routed(store_event(1, 0, 1, &[], &hashes, StorageTier::Device)) - .await - .unwrap(); - indexer.dump_events().await.unwrap(); - if !update_only { - core.delete_worker(1).await.unwrap(); - } - let mut request = request; - request.total_kv_blocks = Some(2048); - core.upsert_worker(request).await.unwrap(); - let retained = partition.indexer().dump_events().await.unwrap(); - if update_only { - assert!( - !retained.is_empty(), - "metadata update cleared cached blocks" - ); - } else { - assert!(retained.is_empty(), "removed worker retained cached blocks"); - let current = core - .indexer_registry - .get_indexer(&key) - .unwrap() - .indexer - .clone(); - current - .apply_event_routed(store_event(1, 0, 1, &[], &hashes, StorageTier::Device)) - .await - .unwrap(); - current.dump_events().await.unwrap(); - } - let mut request = select_request(); - request.prompt = PromptRequest { - token_ids: Some(tokens), - ..PromptRequest::default() - }; - assert_eq!( - core.select(request).await.unwrap().overlap.longest_matched, - 8, - "update_only={update_only}" - ); - assert_eq!( - partition.indexer().dump_events().await.unwrap().len(), - indexer.dump_events().await.unwrap().len() - ); - } - } - #[tokio::test] - async fn metadata_update_preserves_live_booking() { - struct SuspendedDetach(tokio::sync::Notify); - #[async_trait::async_trait] - impl KvEventIngress for SuspendedDetach { - fn open( - &self, - registry: &WorkerRegistry, - key: &RoutingPartitionId, - block_size: u32, - ) -> Indexer { - registry.get_or_create_indexer(key.clone(), block_size) - } - async fn detach(&self, _registry: &WorkerRegistry, _record: &WorkerCatalogRecord) { - self.0.notify_one(); - std::future::pending().await - } - } - let ingress = Arc::new(SuspendedDetach(tokio::sync::Notify::new())); - let core = core_with_host(SelectionHost { - cache: HostCache { - index: KvIndexSource::Owned(ingress.clone()), - shared: None, - }, - ..SelectionHost::default() - }); - core.upsert_worker(worker(1)).await.unwrap(); - core.select_and_reserve(reserve_request("live")) - .await - .unwrap(); - let entry = core - .entry(&RoutingPartitionId::new("model", "default")) - .unwrap(); - let mut updated = worker(1); - updated.max_num_batched_tokens = Some(2048); - tokio::select! { - result = core.upsert_worker(updated) => { result.unwrap(); }, - _ = ingress.0.notified() => { - // A capacity update must not temporarily withdraw this worker. - assert!(entry.scheduler.has_request("live")); - panic!("capacity update detached the worker"); - } - } - assert!(entry.scheduler.has_request("live")); - core.free_reservation("live").await.unwrap(); - } - #[tokio::test(start_paused = true)] - async fn host_lease_manager_owns_expiry() { - use crate::scheduling::queue::SchedulerBookingDescriptor; - struct HostLeases; - impl ReplicaRequestLeaseObserver for HostLeases { - fn admitted(&self, _: SchedulerBookingDescriptor) {} - fn progressed(&self, _: &SchedulerBookingDescriptor) {} - fn completed(&self, _: &SchedulerBookingDescriptor) {} - } - for host_owned in [false, true] { - let core = core_with_host(SelectionHost { - replication: HostReplication { - request_leases: host_owned - .then(|| Arc::new(HostLeases) as Arc), - ..HostReplication::default() - }, - ..SelectionHost::default() - }); - core.upsert_worker(worker(1)).await.unwrap(); - core.select_and_reserve(reserve_request("live")) - .await - .unwrap(); - tokio::task::yield_now().await; - tokio::time::advance(active_request_expiry_duration() * 3).await; - tokio::task::yield_now().await; - let entry = core - .entry(&RoutingPartitionId::new("model", "default")) - .unwrap(); - assert_eq!(entry.scheduler.has_request("live"), host_owned); - if host_owned { - core.free_reservation("live").await.unwrap(); - assert!(!entry.scheduler.has_request("live")); - } - } - } - - #[tokio::test] - async fn catalog_updates_commit_in_order_without_exposing_candidates() { - struct PausedUpdate { - entered: tokio::sync::Notify, - release: tokio::sync::Notify, - } - #[async_trait::async_trait] - impl KvEventIngress for PausedUpdate { - fn open( - &self, - registry: &WorkerRegistry, - key: &RoutingPartitionId, - block_size: u32, - ) -> Indexer { - registry.get_or_create_indexer(key.clone(), block_size) - } - async fn reconcile( - &self, - _: &WorkerRegistry, - previous: Option<&WorkerCatalogRecord>, - _: &WorkerCatalogRecord, - ) -> Result<(), SelectionError> { - if previous.is_some() { - self.entered.notify_one(); - self.release.notified().await; - } - Ok(()) - } - } - let ingress = Arc::new(PausedUpdate { - entered: tokio::sync::Notify::new(), - release: tokio::sync::Notify::new(), - }); - let mut core = core_with_host(SelectionHost { - cache: HostCache { - index: KvIndexSource::Owned(ingress.clone()), - shared: None, - }, - ..SelectionHost::default() - }); - core.listens_for_kv_events = true; - core.upsert_worker(worker(1)).await.unwrap(); - core.select_and_reserve(reserve_request("live")) - .await - .unwrap(); - let mut updated = worker(1); - updated.max_num_batched_tokens = Some(2048); - let update = core.upsert_worker(updated); - tokio::pin!(update); - tokio::select! { - _ = &mut update => panic!("update should pause before committing"), - _ = ingress.entered.notified() => {} - } - let committed = core.catalog.get(1).unwrap(); - assert_eq!(committed.lifecycle, WorkerLifecycle::Schedulable); - assert_eq!(committed.max_num_batched_tokens, Some(1024)); - let entry = core.entry(&committed.key()).unwrap(); - assert!(entry.scheduler.has_request("live")); - let deletion = core.delete_worker(1); - tokio::pin!(deletion); - assert!( - std::future::Future::poll( - deletion.as_mut(), - &mut std::task::Context::from_waker(std::task::Waker::noop()) - ) - .is_pending() - ); - ingress.release.notify_one(); - assert_eq!(update.await.unwrap().max_num_batched_tokens, Some(2048)); - assert!(entry.scheduler.has_request("live")); - assert_eq!( - deletion.await.unwrap().lifecycle, - WorkerLifecycle::Unschedulable - ); - tokio::time::timeout(Duration::from_secs(1), async { - while entry.scheduler.has_request("live") { - tokio::task::yield_now().await; - } - }) - .await - .unwrap(); - } - - #[tokio::test] - async fn affinity_configuration_rejects_invalid_or_conflicting_config() { - use super::super::service::SelectionServiceBuilder; - for ttl in [ - Duration::ZERO, - Duration::from_secs(super::super::affinity::MAX_SESSION_AFFINITY_TTL_SECS + 1), - ] { - let result = SelectionServiceBuilder::new( - test_config(false), - WorkerType::Aggregated, - Default::default(), - ) - .session_affinity(ttl) - .build() - .await; - assert!(result.is_err()); - } - let core = core_with_session_affinity(); - core.upsert_worker(worker(1)).await.unwrap(); - let partition = core - .partition(&RoutingPartitionId::new("model", "default")) - .unwrap(); - assert!(matches!( - partition.session_affinity(SessionAffinityConfig::new(Duration::from_secs(20))), - Err(SelectionError::Conflict(_)) - )); - assert!(matches!( - partition.session_affinity( - SessionAffinityConfig::new(Duration::from_secs(10)) - .with_mode(SessionAffinityMode::Soft) - ), - Err(SelectionError::Conflict(_)) - )); - } - - fn hint_config(worker_type: &str, endpoints: &[(u32, &str)]) -> SelectionWorkerConfig { - SelectionWorkerConfig { - endpoint: "http://worker:8000".to_string(), - data_parallel_start_rank: 0, - data_parallel_size: endpoints.len().max(1) as u32, - max_num_batched_tokens: None, - total_kv_blocks: None, - stable_routing_id: None, - is_eagle: None, - taints: HashSet::new(), - topology_domains: HashMap::new(), - kv_transfer_domain: None, - kv_transfer_enforcement: None, - kv_transfer_preferred_weight: None, - router_hint_worker_type: Some(worker_type.to_string()), - router_hint_source_control_endpoints: endpoints - .iter() - .map(|(rank, endpoint)| (*rank, endpoint.to_string())) - .collect(), - kv_event_source_mode: None, - } - } - - fn hint_candidates( - hashes: &[u64], - owners: Vec<(KvTransferCandidateSource, usize)>, - ) -> KvTransferCandidates { - KvTransferCandidates { - block_hashes: hashes - .iter() - .map(|h| crate::protocols::ExternalSequenceBlockHash(*h)) - .collect(), - owner_prefix_blocks: owners, - routing_snapshot: None, - } - } - - #[test] - fn hint_source_endpoint_follows_the_source_dp_rank() { - let configs = HashMap::from([( - 7, - hint_config( - "prefill", - &[(0, "tcp://127.0.0.1:23280"), (1, "tcp://127.0.0.1:23281")], - ), - )]); - let candidates = - hint_candidates(&[101, 102], vec![(WorkerWithDpRank::new(7, 1).into(), 2)]); - let hint = transfer_hint_for_selection( - &configs, - WorkerWithDpRank::new(7, 0), - 0, - Some(&candidates), - ); - assert_eq!( - hint.map(|payload| payload.source_control_endpoint), - Some("tcp://127.0.0.1:23281".to_string()) - ); - } - - #[test] - fn hint_source_must_share_the_target_worker_type() { - let configs = HashMap::from([ - (7, hint_config("prefill", &[(0, "tcp://127.0.0.1:23280")])), - (8, hint_config("prefill", &[(0, "tcp://127.0.0.1:23281")])), - (9, hint_config("decode", &[(0, "tcp://127.0.0.1:23282")])), - (10, hint_config("decode", &[(0, "tcp://127.0.0.1:23283")])), - ]); - let candidates = hint_candidates( - &[101, 102, 103], - vec![ - (WorkerWithDpRank::new(8, 0).into(), 2), - (WorkerWithDpRank::new(9, 0).into(), 3), - ], - ); - // The longer decode prefix is skipped for a prefill target. - let prefill = transfer_hint_for_selection( - &configs, - WorkerWithDpRank::new(7, 0), - 0, - Some(&candidates), - ) - .expect("prefill hint"); - assert_eq!(prefill.source_control_endpoint, "tcp://127.0.0.1:23281"); - assert_eq!(prefill.block_hashes.len(), 2); - let decode = transfer_hint_for_selection( - &configs, - WorkerWithDpRank::new(10, 0), - 0, - Some(&candidates), - ) - .expect("decode hint"); - assert_eq!(decode.source_control_endpoint, "tcp://127.0.0.1:23282"); - assert_eq!(decode.block_hashes.len(), 3); - } - - #[test] - fn hint_resolves_a_persistent_cache_owner_over_a_state_agent_worker() { - use crate::identity::{ - CacheOwnerId, CacheSemanticsId, DcId, IdentitySource, IndexerDomainId, PoolId, - RoutingScopeId, StableDpSlotId, - }; - use crate::protocols::{ - ResidencyOwner, ResidencyProjection, ResidencyRoutingSnapshot, RouterHintSourceMetadata, - }; - let mut stale = hint_config("prefill", &[(0, "tcp://stale-worker:23280")]); - stale.kv_event_source_mode = Some("state_agent_v2".to_string()); - let configs = HashMap::from([(7, hint_config("prefill", &[])), (8, stale)]); - let owner = CacheOwnerId::new( - PoolId::new( - IndexerDomainId::new( - CacheSemanticsId::new([1; 16], IdentitySource::Explicit), - RoutingScopeId::new([2; 16], IdentitySource::Explicit), - ), - DcId::new(3), - ), - StableDpSlotId::new([4; 16], IdentitySource::Explicit), - ); - let mut candidates = hint_candidates( - &[101, 102], - vec![ - ( - KvTransferCandidateSource::Worker(WorkerWithDpRank::new(8, 0)), - 2, - ), - ( - KvTransferCandidateSource::CacheOwner( - ResidencyOwner::cache_owner(owner).compact_key(), - ), - 2, - ), - ], - ); - candidates.routing_snapshot = Some(Arc::new(ResidencyRoutingSnapshot::new( - ResidencyProjection::default(), - [( - owner, - RouterHintSourceMetadata { - source_control_endpoint: "tcp://persistent-owner:23280".to_string(), - worker_type: "prefill".to_string(), - }, - None, - )], - ))); - let hint = transfer_hint_for_selection( - &configs, - WorkerWithDpRank::new(7, 0), - 0, - Some(&candidates), - ) - .expect("hint"); - assert_eq!(hint.source_control_endpoint, "tcp://persistent-owner:23280"); - } -} diff --git a/lib/kv-router/src/services/selection/core/queries.rs b/lib/kv-router/src/services/selection/core/queries.rs new file mode 100644 index 000000000000..66d862bd08e6 --- /dev/null +++ b/lib/kv-router/src/services/selection/core/queries.rs @@ -0,0 +1,113 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Read-only partition queries: loads, potential loads and overlap scores. + +use super::*; + +impl SelectionCore { + pub fn loads( + &self, + model_name: Option<&str>, + routing_group: Option<&str>, + ) -> Vec { + let entries = self.initialized_entries(); + let mut loads = Vec::new(); + for entry in entries { + if model_name.is_some_and(|model_name| entry.key.model_name != model_name) + || routing_group + .is_some_and(|routing_group| entry.key.routing_group != routing_group) + { + continue; + } + loads.push(ModelLoadResponse { + model_name: entry.key.model_name.clone(), + routing_group: entry.key.routing_group.clone(), + loads: entry + .scheduler + .get_potential_loads(None, 0, HashMap::new(), false), + pending_count: entry.scheduler.pending_count(), + pending_isl_tokens: entry.scheduler.pending_isl_tokens(), + }); + } + loads.sort_by(|a, b| { + (&a.model_name, &a.routing_group).cmp(&(&b.model_name, &b.routing_group)) + }); + loads + } + + pub async fn potential_loads( + &self, + req: PotentialLoadsRequest, + ) -> Result, SelectionError> { + let key = RoutingPartitionId::new(req.model_name.clone(), req.routing_group.clone()); + let entry = self.ready_entry(&key)?; + let prepared = self + .prepare_selection_inputs( + &entry, + &req.prompt.view(), + Some( + self.kv_router_config + .assume_kv_reuse(req.router_config_override.as_ref()), + ), + false, + false, + &mut None, + ) + .await?; + let track_prefill_tokens = req + .router_config_override + .as_ref() + .and_then(|cfg| cfg.track_prefill_tokens) + .unwrap_or(self.kv_router_config.router_track_prefill_tokens); + Ok(entry.scheduler.get_potential_loads( + Some(prepared.sequence_hashes), + prepared.isl_tokens, + prepared.overlap.effective_cached_tokens, + track_prefill_tokens, + )) + } + + pub async fn overlap_scores( + &self, + req: OverlapScoresRequest, + ) -> Result { + let key = RoutingPartitionId::new(req.model_name.clone(), req.routing_group.clone()); + let entry = self.ready_entry(&key)?; + let block_hashes = req + .prompt + .view() + .block_hashes_for_indexer(entry.block_size, entry.is_eagle)?; + let num_blocks = block_hashes.len(); + let tiered = entry + .indexer + .find_tiered_matches(block_hashes) + .await + .map_err(|error| SelectionError::Internal(error.to_string()))?; + let schedulable_workers = self.schedulable_worker_ranks(&key); + Ok( + OverlapAnalysis::new(&self.kv_router_config, entry.block_size, &tiered) + .scores_response( + req.router_config_override.as_ref(), + num_blocks, + schedulable_workers, + false, + None, + None, + ), + ) + } + + fn schedulable_worker_ranks(&self, key: &RoutingPartitionId) -> Vec { + let configs = self.catalog.scheduler_configs_for_key(key); + let mut workers = Vec::new(); + for (worker_id, config) in configs { + let start = config.data_parallel_start_rank; + let end = start.saturating_add(config.data_parallel_size); + for dp_rank in start..end { + workers.push(WorkerWithDpRank::new(worker_id, dp_rank)); + } + } + workers + } +} diff --git a/lib/kv-router/src/services/selection/core/reservations.rs b/lib/kv-router/src/services/selection/core/reservations.rs new file mode 100644 index 000000000000..92c6fd84302a --- /dev/null +++ b/lib/kv-router/src/services/selection/core/reservations.rs @@ -0,0 +1,553 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Reservations: the selection-id index over scheduler bookings, the +//! `create_reservation` replay path, and the lifecycle calls that advance or +//! release them. + +use super::*; + +/// Resolved inputs for booking a reservation, shared by the cached and explicit +/// `create_reservation` paths. +struct ReservationBooking { + key: RoutingPartitionId, + selection_id: String, + worker: WorkerWithDpRank, + sequence_hashes: Vec, + prefill_load_hint: Option, + expected_output_tokens: Option, + track_prefill_tokens: bool, + lora_name: Option, + /// Public block hashes to record into an approximate indexer once booked. + routing_hashes: Option>, + session_id: Option, +} + +/// `selection_id` -> the booking it holds. +/// +/// Lifecycle calls (`prefill_complete`, `free`, `add_output_block`) arrive with +/// only a selection id; the index resolves them to one partition and one +/// scheduler booking, so they never touch a booking made by a later request +/// that reused the id. An id is claimed here before its booking is made and +/// installed once the booking is final; a claim (`booking == None`) rejects a +/// concurrent booking of the same id and is invisible to lifecycle calls. +/// Bookings mirrored from replica peers are indexed by +/// [`ReservationIndexObserver`]. +pub(super) type ReservationIndex = RwLock>; + +pub(super) struct Reservation { + pub(super) partition: RoutingPartitionId, + pub(super) booking: Option, + pub(super) _affinity_lease: Option, +} + +/// Exclusive ownership of a selection id while its booking is in flight. +/// Dropping the claim without `install` releases the id. +pub(super) struct ReservationClaim<'a> { + pub(super) index: &'a ReservationIndex, + pub(super) selection_id: String, + pub(super) armed: bool, +} + +impl ReservationClaim<'_> { + pub(super) fn install(mut self, reservation: Reservation) { + self.armed = false; + self.index + .write() + .insert(std::mem::take(&mut self.selection_id), reservation); + } +} + +impl Drop for ReservationClaim<'_> { + fn drop(&mut self) { + if !self.armed { + return; + } + let mut index = self.index.write(); + if index + .get(&self.selection_id) + .is_some_and(|reservation| reservation.booking.is_none()) + { + index.remove(&self.selection_id); + } + } +} + +/// Keeps the reservation index exact for bookings replicated from peers, which +/// never pass through this core's booking paths, and forwards to the host's +/// own observer. +pub(super) struct ReservationIndexObserver { + pub(super) index: Arc, + pub(super) partition: RoutingPartitionId, + pub(super) host: Option>, +} + +impl ReplicaRequestLeaseObserver for ReservationIndexObserver { + fn admitted(&self, booking: SchedulerBookingDescriptor) { + { + let mut index = self.index.write(); + // A peer only admits an id this partition's scheduler does not hold, + // so an existing row for it is a stale booking (expired, not yet + // swept) or a claim whose local booking will now fail; the mirror + // replaces both. A row from another partition is left alone. + let own_row = index + .get(&booking.request_id) + .is_none_or(|reservation| reservation.partition == self.partition); + if own_row { + index.insert( + booking.request_id.clone(), + Reservation { + partition: self.partition.clone(), + booking: Some(booking.clone()), + _affinity_lease: None, + }, + ); + } + } + if let Some(host) = &self.host { + host.admitted(booking); + } + } + + fn progressed(&self, booking: &SchedulerBookingDescriptor) { + if let Some(host) = &self.host { + host.progressed(booking); + } + } + + fn completed(&self, booking: &SchedulerBookingDescriptor) { + forget_reservation_if(&self.index, &self.partition, booking); + if let Some(host) = &self.host { + host.completed(booking); + } + } +} + +/// Remove the index entry for `booking` only if it still describes it. +fn forget_reservation_if( + index: &ReservationIndex, + partition: &RoutingPartitionId, + booking: &SchedulerBookingDescriptor, +) { + let mut index = index.write(); + if index.get(&booking.request_id).is_some_and(|reservation| { + reservation.partition == *partition && reservation.booking.as_ref() == Some(booking) + }) { + index.remove(&booking.request_id); + } +} + +impl SelectionCore { + pub async fn create_reservation( + &self, + req: ReservationRequest, + ) -> Result { + self.ensure_running()?; + + let key = RoutingPartitionId::new(req.model_name.clone(), req.routing_group.clone()); + + // Explicit form: book on the given worker under selection_id, discarding + // any cached selection for the id so a later replay can't book stale state. + if let Some(worker_id) = req.worker_id { + self.selection_cache.discard(&key, &req.selection_id); + return self.reserve_explicit(key, worker_id, req).await; + } + + // Replay form: peek, book, and consume only once the booking lands. A + // failure leaves the entry for a retry; concurrent replays of the same + // id collide at the scheduler, so they can't double-book. + let Some((pending, generation)) = + self.selection_cache + .peek(&key, &req.selection_id, Instant::now()) + else { + return Err(SelectionError::NotFound(format!( + "no pending selection {} for {key} (expired, already used, \ + or never selected)", + req.selection_id + ))); + }; + let response = match self.book_cached_selection(pending, &req).await { + Ok(response) => response, + // The session moved: the cached worker can never be booked for it. + Err(SelectionError::BadRequest(message)) => { + self.selection_cache.discard(&key, &req.selection_id); + return Err(SelectionError::BadRequest(message)); + } + Err(error) => return Err(error), + }; + self.selection_cache + .remove(&key, &req.selection_id, generation); + Ok(response) + } + + /// Book a reservation replaying what the matching `select` captured; request + /// fields other than the ids are ignored. + async fn book_cached_selection( + &self, + pending: Arc, + req: &ReservationRequest, + ) -> Result { + let (entry, endpoint, prefill_load_hint) = self.resolve_cached_booking(&pending)?; + let track_prefill_tokens = pending.track_prefill_tokens; + self.finalize_reservation( + entry, + endpoint, + ReservationBooking { + key: pending.key.clone(), + selection_id: req.selection_id.clone(), + worker: pending.worker, + sequence_hashes: pending.sequence_hashes.clone(), + prefill_load_hint: track_prefill_tokens.then_some(prefill_load_hint), + expected_output_tokens: pending.expected_output_tokens, + track_prefill_tokens, + lora_name: pending.lora_name.clone(), + routing_hashes: pending.routing_hashes.clone(), + session_id: pending.session_id.clone(), + }, + ) + .await + } + + /// Resolve everything a cached booking needs (ready entry, schedulable + /// endpoint, prefill hint), so the only fallible step left in + /// `finalize_reservation` is the scheduler call. + fn resolve_cached_booking( + &self, + pending: &PendingSelection, + ) -> Result<(Arc, String, PrefillLoadHint), SelectionError> { + let entry = self.ready_entry(&pending.key)?; + // Validate the full worker/rank against current topology; a rank a PATCH + // removed during the window is rejected (the entry stays for a retry). + let endpoint = self + .catalog + .schedulable_worker_endpoint(pending.worker, &pending.key) + .ok_or_else(|| { + SelectionError::NotFound(format!( + "schedulable worker {} (dp_rank {}) not found for {}", + pending.worker.worker_id, pending.worker.dp_rank, pending.key + )) + })?; + let prefill_load_hint = prefill_load_hint_from_effective_tokens( + pending.isl_tokens, + pending.effective_prefill_tokens, + ) + .map_err(|error| SelectionError::BadRequest(error.to_string()))?; + Ok((entry, endpoint, prefill_load_hint)) + } + + fn schedulable_endpoint( + &self, + worker_id: WorkerId, + key: &RoutingPartitionId, + ) -> Result { + self.catalog + .schedulable_endpoint(worker_id, key) + .ok_or_else(|| { + SelectionError::NotFound(format!( + "schedulable worker {worker_id} not found for {key}" + )) + }) + } + + /// Book a reservation from a self-contained request (explicit worker_id and prompt). + async fn reserve_explicit( + &self, + key: RoutingPartitionId, + worker_id: WorkerId, + req: ReservationRequest, + ) -> Result { + let entry = self.ready_entry(&key)?; + let normalized = req.prompt.view().normalize_for_reservation( + entry.is_eagle, + TrackingHashInput { + context: &self.tracking_hash, + scope: tracking_scope(&entry), + assume_kv_reuse: self + .kv_router_config + .assume_kv_reuse(req.router_config_override.as_ref()), + }, + )?; + let prefill_load_hint = req + .effective_prefill_tokens + .map(|tokens| { + prefill_load_hint_from_effective_tokens(normalized.isl_tokens, tokens) + .map_err(|error| SelectionError::BadRequest(error.to_string())) + }) + .transpose()?; + let worker = WorkerWithDpRank::new(worker_id, req.dp_rank.unwrap_or(0)); + let endpoint = self.schedulable_endpoint(worker.worker_id, &key)?; + let track_prefill_tokens = req.track_prefill_tokens.unwrap_or_else(|| { + req.effective_prefill_tokens.is_some() + || req + .router_config_override + .as_ref() + .and_then(|cfg| cfg.track_prefill_tokens) + .unwrap_or(self.kv_router_config.router_track_prefill_tokens) + }); + // Hash-only reservations (sequence hashes without block hashes) carry + // nothing an indexer can key on; recording is skipped for them. + let can_record = entry.indexer.records_routing_decisions() + && (req.prompt.token_ids.is_some() || req.prompt.block_hashes.is_some()); + let routing_hashes = can_record + .then(|| { + req.prompt + .view() + .block_hashes_for_indexer(entry.block_size, entry.is_eagle) + }) + .transpose()?; + + self.finalize_reservation( + entry, + endpoint, + ReservationBooking { + key, + selection_id: req.selection_id, + worker, + sequence_hashes: normalized.sequence_hashes, + prefill_load_hint, + expected_output_tokens: req.expected_output_tokens, + track_prefill_tokens, + lora_name: req.prompt.lora_name, + routing_hashes, + session_id: None, + }, + ) + .await + } + + /// Register the booking with the scheduler. All fallible resolution happens + /// in the caller; the scheduler add here is the last step that can fail, and + /// the cached path leaves its selection in place (to retry) if it does. + async fn finalize_reservation( + &self, + entry: Arc, + endpoint: String, + booking: ReservationBooking, + ) -> Result { + let ReservationBooking { + key, + selection_id, + worker, + sequence_hashes, + prefill_load_hint, + expected_output_tokens, + track_prefill_tokens, + lora_name, + routing_hashes, + session_id, + } = booking; + + let claim = self.claim_reservation(&selection_id, &key)?; + // Hold the session before booking: holding after would wait on an + // initializer that may itself be waiting for this worker's capacity. + let session = match (session_id.as_deref(), entry.affinity.get()) { + (Some(session_id), Some(table)) => self + .hold_session(table, session_id, &key) + .await? + .map(|hold| (table, session_id, hold)), + _ => None, + }; + // Strict booking: never lazily recreate a worker/rank removed since the + // reservation was resolved. The lease frees the booking if this future + // is dropped before `install`. + let lease = entry + .scheduler + .add_request_if_registered_guarded(SequenceRequest { + request_id: selection_id.clone(), + token_sequence: Some(sequence_hashes), + track_prefill_tokens, + expected_output_tokens, + prefill_load_hint, + worker, + lora_name, + })?; + let affinity_lease = match session { + Some((table, session_id, hold)) => { + self.commit_session(table, hold, session_id, worker, &key)? + } + None => None, + }; + if let Some(hashes) = routing_hashes { + self.record_routing_decision(&entry, worker, hashes).await; + } + claim.install(Reservation { + partition: key.clone(), + booking: Some(lease.commit().ok_or_else(missing_booking)?), + _affinity_lease: affinity_lease, + }); + + Ok(ReservationResponse { + selection_id, + model_name: key.model_name, + routing_group: key.routing_group, + worker_id: worker.worker_id, + dp_rank: worker.dp_rank, + endpoint, + }) + } + + /// Take `selection_id` for a booking about to be made in `key`. + pub(super) fn claim_reservation( + &self, + selection_id: &str, + key: &RoutingPartitionId, + ) -> Result, SelectionError> { + let mut index = self.reservation_index.write(); + if index.contains_key(selection_id) { + return Err(SelectionError::Conflict(format!( + "selection {selection_id} is already reserved or being reserved" + ))); + } + index.insert( + selection_id.to_string(), + Reservation { + partition: key.clone(), + booking: None, + _affinity_lease: None, + }, + ); + Ok(ReservationClaim { + index: &self.reservation_index, + selection_id: selection_id.to_string(), + armed: true, + }) + } + + /// The partition and booking a lifecycle call on `selection_id` may touch. + /// `None` for unknown ids and for ids whose booking is still in flight. + pub(super) fn indexed_booking( + &self, + selection_id: &str, + ) -> Option<(Arc, SchedulerBookingDescriptor)> { + // Release the index guard before taking `entries`: the sweep takes + // `entries` then `reservation_index`, so nesting here would invert the + // lock order. + let (partition, booking) = { + let index = self.reservation_index.read(); + let reservation = index.get(selection_id)?; + (reservation.partition.clone(), reservation.booking.clone()?) + }; + Some((self.entry(&partition)?, booking)) + } + + fn reservation_not_found(selection_id: &str) -> SelectionError { + SelectionError::NotFound(format!("reservation {selection_id} not found")) + } + + pub async fn prefill_complete(&self, selection_id: &str) -> Result<(), SelectionError> { + let Some((entry, booking)) = self.indexed_booking(selection_id) else { + return Err(Self::reservation_not_found(selection_id)); + }; + match entry + .scheduler + .mark_prefill_completed_if_booking(&booking) + .await? + { + LifecycleMutationOutcome::Applied => Ok(()), + // Already marked: still publish the ordered completion so peers + // that missed the first event converge. + LifecycleMutationOutcome::NoChange + if entry + .scheduler + .publish_prefill_completed_if_booking(&booking) => + { + Ok(()) + } + LifecycleMutationOutcome::NoChange => { + forget_reservation_if(&self.reservation_index, &entry.key, &booking); + Err(Self::reservation_not_found(selection_id)) + } + } + } + + pub async fn free_reservation(&self, selection_id: &str) -> Result<(), SelectionError> { + let Some((entry, booking)) = self.indexed_booking(selection_id) else { + return Err(Self::reservation_not_found(selection_id)); + }; + let outcome = entry.scheduler.free_if_booking(&booking).await?; + forget_reservation_if(&self.reservation_index, &entry.key, &booking); + match outcome { + LifecycleMutationOutcome::Applied => Ok(()), + LifecycleMutationOutcome::NoChange => Err(Self::reservation_not_found(selection_id)), + } + } + + pub fn add_output_block( + &self, + selection_id: &str, + decay_fraction: Option, + ) -> Result<(), SelectionError> { + if let Some(frac) = decay_fraction + && !(0.0..=1.0).contains(&frac) + { + return Err(SelectionError::BadRequest( + "decay_fraction must be between 0.0 and 1.0".to_string(), + )); + } + + let Some((entry, booking)) = self.indexed_booking(selection_id) else { + return Err(Self::reservation_not_found(selection_id)); + }; + match entry + .scheduler + .add_output_block_if_booking_sync(&booking, decay_fraction)? + { + LifecycleMutationOutcome::Applied => Ok(()), + LifecycleMutationOutcome::NoChange => { + forget_reservation_if(&self.reservation_index, &entry.key, &booking); + Err(Self::reservation_not_found(selection_id)) + } + } + } +} + +/// Drop index entries whose booking no longer exists in its partition +/// scheduler (expired by the periodic force-expiry, or freed through a path +/// that bypassed this core). Claims are left for their owner. Returns the +/// number of entries removed. +pub(super) fn sweep_reservation_index( + entries: &SelectionEntries, + index: &ReservationIndex, +) -> usize { + let entries = entries.read(); + let mut index = index.write(); + let before = index.len(); + index.retain(|_, reservation| { + let Some(booking) = &reservation.booking else { + return true; + }; + entries + .get(&reservation.partition) + .and_then(|cell| cell.get()) + .is_some_and(|entry| entry.scheduler.has_booking(booking)) + }); + before - index.len() +} + +pub(super) fn missing_booking() -> SelectionError { + SelectionError::Internal("booking lease holds no booking".to_string()) +} + +pub(super) fn spawn_reservation_index_sweep( + entries: Arc, + index: Arc, + cancel_token: CancellationToken, +) { + let period = crate::sequences::active_request_expiry_duration(); + tokio::spawn(async move { + let mut interval = tokio::time::interval(period); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + interval.tick().await; + loop { + tokio::select! { + _ = cancel_token.cancelled() => break, + _ = interval.tick() => { + let removed = sweep_reservation_index(&entries, &index); + if removed > 0 { + tracing::debug!(removed, "Swept stale selection reservation index entries"); + } + } + } + } + }); +} diff --git a/lib/kv-router/src/services/selection/core/run.rs b/lib/kv-router/src/services/selection/core/run.rs new file mode 100644 index 000000000000..649d4ca4d945 --- /dev/null +++ b/lib/kv-router/src/services/selection/core/run.rs @@ -0,0 +1,740 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! One selection, end to end: the wire adapters, `run_selection`, the +//! lookups that feed the scheduler, and the session binding around it. + +use super::hint::transfer_hint_for_selection; +use super::reservations::{Reservation, missing_booking}; +use super::*; + +pub(super) struct PreparedSelectionInputs { + pub(super) block_hashes: Vec, + pub(super) sequence_hashes: Vec, + pub(super) isl_tokens: usize, + pub(super) overlap: OverlapSignals, + pub(super) shared_cache_hits: Option, + pub(super) kv_transfer_candidates: Option, +} + +impl SelectionCore { + pub async fn select(&self, req: SelectRequest) -> Result { + self.select_with_policy_class(req, None).await + } + + pub async fn select_with_policy_class( + &self, + mut req: SelectRequest, + policy_class: Option, + ) -> Result { + let session_context = req.take_session_context(); + let is_steerable = req.affinity_target.is_none() && req.pinned_worker.is_none(); + let session = session_binding(session_context.as_ref(), is_steerable, false); + let admission = if req.advisory { + SelectionAdmission::Advisory { + request_id: req.selection_id.clone(), + } + } else { + SelectionAdmission::Query { + request_id: req.selection_id.clone(), + } + }; + let (selected, endpoint) = self + .select_or_reject(SelectionOperation { + key: RoutingPartitionId::new(req.model_name, req.routing_group), + prompt: req.prompt.view(), + router_config_override: req.router_config_override, + expected_output_tokens: req.expected_output_tokens, + priority_jump: req.priority_jump.unwrap_or_default(), + strict_priority: req.strict_priority.unwrap_or(0), + policy_class, + session_context, + session, + affinity_target: req.affinity_target, + pinned_worker: req.pinned_worker, + allowed_worker_ids: req.allowed_worker_ids, + routing_constraints: req.routing_constraints, + admission, + track_active_blocks: true, + return_routing_hashes: false, + replay_id: req.selection_id.clone(), + }) + .await?; + Ok(self.select_response(selected, endpoint, req.selection_id)) + } + + pub async fn select_and_reserve( + &self, + req: SelectAndReserveRequest, + ) -> Result { + self.select_and_reserve_with_policy_class(req, None).await + } + + pub async fn select_and_reserve_with_policy_class( + &self, + mut req: SelectAndReserveRequest, + policy_class: Option, + ) -> Result { + let session_context = req.take_session_context(); + let selection_id = req + .selection_id + .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); + let is_steerable = req.affinity_target.is_none() && req.pinned_worker.is_none(); + let session = session_binding(session_context.as_ref(), is_steerable, true); + let (selected, endpoint) = self + .select_or_reject(SelectionOperation { + key: RoutingPartitionId::new(req.model_name, req.routing_group), + prompt: req.prompt.view(), + router_config_override: req.router_config_override, + expected_output_tokens: req.expected_output_tokens, + priority_jump: req.priority_jump.unwrap_or_default(), + strict_priority: req.strict_priority.unwrap_or(0), + policy_class, + session_context, + session, + affinity_target: req.affinity_target, + pinned_worker: req.pinned_worker, + allowed_worker_ids: req.allowed_worker_ids, + routing_constraints: req.routing_constraints, + admission: SelectionAdmission::Book { + selection_id: selection_id.clone(), + }, + track_active_blocks: true, + return_routing_hashes: false, + replay_id: None, + }) + .await?; + Ok(self.select_response(selected, endpoint, Some(selection_id))) + } + + /// The wire shape of a selection; a queue rejection keeps its 503 body. + async fn select_or_reject( + &self, + operation: SelectionOperation<'_>, + ) -> Result<(Selected, String), SelectionError> { + let mut selected = match self.run_selection(operation).await.result? { + SelectionOutcome::Selected(selected) => selected, + SelectionOutcome::QueueRejected { rejection } => { + return Err(SelectionError::Scheduler(KvSchedulerError::QueueRejected( + rejection, + ))); + } + }; + // Wire admissions never book without an endpoint; see `Selected::endpoint`. + let endpoint = selected.endpoint.take().ok_or_else(|| { + SelectionError::Internal(format!( + "selected worker {} is no longer schedulable", + selected.response.best_worker.worker_id + )) + })?; + Ok((selected, endpoint)) + } + + fn select_response( + &self, + selected: Selected, + endpoint: String, + selection_id: Option, + ) -> SelectResponse { + let Selected { + key, + response, + advisory_load, + total_kv_blocks, + endpoint: _, + block_size, + isl_tokens, + sequence_hashes, + track_prefill_tokens, + effective_prefill_tokens, + kv_hint, + routing_hashes: _, + shared_cache_hits: _, + lease: _, + } = selected; + let booked = sequence_hashes.is_some(); + let potential_decode_blocks = response.potential_decode_blocks as u64; + let decode_busy = self + .kv_router_config + .conditional_disagg_decode_busy_threshold + .zip(total_kv_blocks) + .map(|(threshold, total_kv_blocks)| { + potential_decode_blocks as f64 > threshold * total_kv_blocks as f64 + }); + let worker_load = advisory_load.map(|load| SelectionWorkerLoad { + active_prefill_tokens: load.active_prefill_tokens, + prefill_token_capacity: load.prefill_token_capacity, + total_kv_blocks, + prefill_busy: self + .kv_router_config + .conditional_disagg_prefill_busy_threshold + .map(|threshold| load.prefill_load_exceeds(threshold)), + }); + SelectResponse { + selection_id, + sequence_hashes: sequence_hashes + .map(|hashes| hashes.into_iter().map(|hash| hash as i64).collect()), + isl_tokens: booked.then_some(isl_tokens), + track_prefill_tokens: booked.then_some(track_prefill_tokens), + model_name: key.model_name, + routing_group: key.routing_group, + worker_id: response.best_worker.worker_id, + dp_rank: response.best_worker.dp_rank, + endpoint, + block_size, + overlap: MooncakeOverlapSummary::from_selected_worker_tiers( + &response.selected_worker_tiers, + block_size, + ), + effective_prefill_tokens, + potential_decode_blocks, + decode_busy, + worker_load, + kv_hint, + } + } + + /// Run one selection: resolve the session, look the prompt up, schedule, + /// and (for `Book`) install the reservation. Every host's selection goes + /// through here. + pub async fn run_selection(&self, operation: SelectionOperation<'_>) -> SelectionRun { + let mut lookup = None; + let result = self.run_selection_inner(operation, &mut lookup).await; + SelectionRun { result, lookup } + } + + async fn run_selection_inner( + &self, + operation: SelectionOperation<'_>, + lookup: &mut Option, + ) -> Result { + let SelectionOperation { + key, + prompt, + router_config_override, + expected_output_tokens, + priority_jump, + strict_priority, + policy_class, + session_context, + session, + affinity_target, + pinned_worker, + allowed_worker_ids, + routing_constraints, + admission, + track_active_blocks, + return_routing_hashes, + replay_id, + } = operation; + self.ensure_running()?; + + let entry = self.ready_entry(&key)?; + let book = admission.is_booking(); + let claim = match &admission { + SelectionAdmission::Book { selection_id } => { + Some(self.claim_reservation(selection_id, &key)?) + } + _ => None, + }; + + // Session stickiness: a bound session steers selection (exclusive for + // the default selector); a new session is bound to the worker booked. + let table = entry.affinity.get(); + let mut affinity_hold = None; + let managed_session = match (&session, table) { + (SessionBinding::Managed { .. }, _) if claim.is_none() => { + return Err(SelectionError::Internal( + "a managed session binding requires Book admission".to_string(), + )); + } + (SessionBinding::Managed { session_id }, Some(table)) => Some((table, session_id)), + _ => None, + }; + let affinity_target = match (&session, table) { + (SessionBinding::Managed { session_id }, Some(table)) => { + affinity_hold = self.hold_session(table, session_id, &key).await?; + affinity_hold.as_ref().and_then(Hold::target) + } + (SessionBinding::Query { session_id }, Some(table)) => table + .query_target(session_id, None) + .map_err(affinity_error)? + .map(|target| WorkerAffinityTarget::new(target.worker_id, target.dp_rank)), + _ => affinity_target, + }; + // Router hints are attached to bookings only, and only when a worker in + // this partition can consume them and the indexer can retain the + // matched chain (local, event-driven, no approximate writes). + let hint_capable_workers = book && self.catalog.has_router_hint_capable_workers(&key); + let retain_kv_transfer_chain = + hint_capable_workers && entry.indexer.supports_kv_transfer_chain_retention(); + if hint_capable_workers && !retain_kv_transfer_chain { + static WARN_ONCE: std::sync::Once = std::sync::Once::new(); + WARN_ONCE.call_once(|| { + tracing::warn!( + "router hints need a local event-driven indexer with no approximate side indexer; workers advertise hint capability but no hints will be attached" + ); + }); + } + let PreparedSelectionInputs { + block_hashes, + sequence_hashes, + isl_tokens, + overlap, + shared_cache_hits, + kv_transfer_candidates, + } = self + .prepare_selection_inputs( + &entry, + &prompt, + track_active_blocks.then(|| { + self.kv_router_config + .assume_kv_reuse(router_config_override.as_ref()) + }), + true, + retain_kv_transfer_chain, + lookup, + ) + .await?; + if tracing::enabled!(tracing::Level::DEBUG) { + let local_hashes: Vec = block_hashes.iter().map(|hash| hash.0).collect(); + tracing::debug!( + request_id = admission.request_id().unwrap_or_default(), + isl_tokens, + block_size = entry.block_size, + num_blocks = local_hashes.len(), + ?local_hashes, + "[ROUTING_INPUT] request local hashes" + ); + } + let host_shared_cache_hits = shared_cache_hits.clone(); + // The queue lease frees a booking whose response is never consumed + // (the caller dropped this future after the actor booked). + let mode = match &admission { + SelectionAdmission::Book { selection_id } + | SelectionAdmission::Lease { + request_id: selection_id, + } => ScheduleMode::TrackedWithLifecycle { + request_id: selection_id.clone(), + }, + SelectionAdmission::Query { request_id } => ScheduleMode::QueryOnly { + request_id: request_id.clone(), + }, + SelectionAdmission::Advisory { request_id } => ScheduleMode::QueryOnly { + request_id: request_id.clone(), + }, + }; + let track_prefill_tokens = router_config_override + .as_ref() + .and_then(|cfg| cfg.track_prefill_tokens) + .unwrap_or(self.kv_router_config.router_track_prefill_tokens); + // An unbooked selection with an id caches the booking inputs so a + // follow-up `create_reservation` can replay them by that id. + let cached_inputs = replay_id.filter(|_| !book).map(|id| { + ( + id, + sequence_hashes.clone(), + prompt.lora_name.map(str::to_string), + track_prefill_tokens, + match &session { + SessionBinding::Query { session_id } => Some(session_id.clone()), + _ => None, + }, + ) + }); + let allowed_worker_ids = match self.host.eligibility.lora_worker_filter.as_deref() { + Some(filter) => narrow_allowed_worker_ids_by_lora( + filter, + prompt.lora_name, + allowed_worker_ids, + pinned_worker.as_ref(), + || entry.workers_tx.borrow().keys().copied().collect(), + ), + None => allowed_worker_ids, + }; + // Bookings (now, or later via the pending-selection cache) are recorded + // into an approximate indexer; keep the public hashes for that. + let routing_hashes = (entry.indexer.records_routing_decisions() + && (claim.is_some() || cached_inputs.is_some())) + .then(|| block_hashes.clone()); + let booked_sequence_hashes = book.then(|| sequence_hashes.clone()); + let returned_routing_hashes = return_routing_hashes.then(|| block_hashes.clone()); + let schedule_request = ScheduleRequest { + mode, + token_seq: track_active_blocks.then_some(sequence_hashes), + block_hashes: Some(block_hashes), + isl_tokens, + overlap, + kv_transfer_candidates, + retain_kv_transfer_chain, + router_config_override, + lora_name: prompt.lora_name.map(str::to_string), + priority_jump, + strict_priority, + policy_class, + session_context, + expected_output_tokens, + affinity_target, + pinned_worker, + allowed_worker_ids, + routing_constraints, + shared_cache_hits, + }; + // `lease` guards the booking until it is installed below: any early + // return or drop before then frees it. + let scheduled = tokio::select! { + biased; + _ = self.cancel_token.cancelled() => { + return Err(SelectionError::Scheduler(KvSchedulerError::SubscriberShutdown)); + } + result = async { + if matches!(admission, SelectionAdmission::Advisory { .. }) { + entry + .scheduler + .select_without_admission(schedule_request) + .instrument(tracing::info_span!("kv_router.select_without_admission")) + .await + .map(|advisory| { + (advisory.response, Some(advisory.selected_worker_load), None) + }) + } else { + entry + .scheduler + .schedule_request_with_lease(schedule_request) + .instrument(tracing::info_span!("kv_router.schedule")) + .await + .map(|(admitted, lease)| (admitted.response, None, lease)) + } + } => result, + }; + let (response, advisory_load, lease) = match scheduled { + Ok(scheduled) => scheduled, + Err(KvSchedulerError::QueueRejected(rejection)) => { + return Ok(SelectionOutcome::QueueRejected { rejection }); + } + Err(error) => return Err(error.into()), + }; + let endpoint = self + .catalog + .schedulable_endpoint(response.best_worker.worker_id, &key); + if endpoint.is_none() && !matches!(admission, SelectionAdmission::Lease { .. }) { + return Err(SelectionError::Internal(format!( + "selected worker {} is no longer schedulable", + response.best_worker.worker_id + ))); + } + let effective_prefill = effective_prefill_tokens(isl_tokens, response.cached_tokens); + let total_kv_blocks = advisory_load + .and_then(|load| load.total_kv_blocks.map(|blocks| blocks as u64)) + .or_else(|| { + self.catalog + .total_kv_blocks(response.best_worker.worker_id, &key) + }); + let kv_hint = if retain_kv_transfer_chain && response.kv_transfer_candidates.is_some() { + transfer_hint_for_selection( + &entry.workers_tx.borrow(), + response.best_worker, + response.target_cached_prefix_blocks, + response.kv_transfer_candidates.as_ref(), + ) + .map(|payload| { + KvHint::new( + admission.request_id().unwrap_or_default(), + vec![KvHintAction::fetch("a1", payload)], + ) + }) + } else { + None + }; + + // The routing hashes go to exactly one of: the reservation recorded + // now, or the replay cache a later reservation records from. + let mut routing_hashes = routing_hashes; + let lease = if let Some(claim) = claim { + let Some(lease) = lease else { + return Err(SelectionError::Internal( + "booked selection has no lifecycle lease".to_string(), + )); + }; + // A rejected affinity commit returns while the lease is still armed, + // so the booking is freed and nothing below is recorded. + let affinity_lease = match (affinity_hold, managed_session) { + (Some(hold), Some((table, session_id))) => { + self.commit_session(table, hold, session_id, response.best_worker, &key)? + } + _ => None, + }; + if let Some(hashes) = routing_hashes.take() { + self.record_routing_decision(&entry, response.best_worker, hashes) + .await; + } + claim.install(Reservation { + partition: key.clone(), + booking: Some(lease.commit().ok_or_else(missing_booking)?), + _affinity_lease: affinity_lease, + }); + None + } else { + lease + }; + + if let Some((cache_id, sequence_hashes, lora_name, track_prefill_tokens, session_id)) = + cached_inputs + { + self.selection_cache.insert( + cache_id, + PendingSelection { + key: key.clone(), + worker: response.best_worker, + sequence_hashes, + isl_tokens, + effective_prefill_tokens: effective_prefill, + expected_output_tokens, + track_prefill_tokens, + lora_name, + routing_hashes, + session_id, + }, + Instant::now(), + ); + } + Ok(SelectionOutcome::Selected(Selected { + key, + response, + advisory_load, + total_kv_blocks, + endpoint, + block_size: entry.block_size, + isl_tokens, + sequence_hashes: booked_sequence_hashes, + track_prefill_tokens, + effective_prefill_tokens: effective_prefill, + kv_hint, + routing_hashes: returned_routing_hashes, + shared_cache_hits: host_shared_cache_hits, + lease, + })) + } + + /// Hold `session_id` for a booking, re-initializing a binding whose worker + /// this partition can no longer schedule. `None` when the table is full: a + /// router-side limit, not a client fault, so the request routes unpinned. + pub(super) async fn hold_session( + &self, + table: &SessionAffinity, + session_id: &str, + key: &RoutingPartitionId, + ) -> Result, SelectionError> { + loop { + let acquired = tokio::select! { + _ = self.cancel_token.cancelled() => { + return Err(SelectionError::Scheduler(KvSchedulerError::SubscriberShutdown)); + } + result = table.acquire(session_id, None) => result, + }; + match acquired { + Ok(Hold::Bound { target, mut lease }) + if !self.catalog.is_schedulable(target.worker_id, key) => + { + tracing::debug!( + session_id, + worker_id = target.worker_id, + "session affinity target is not schedulable; re-initializing" + ); + lease.invalidate(); + } + Ok(hold) => return Ok(Some(hold)), + Err(AffinityError::ResourceExhausted(_)) => { + tracing::debug!( + session_id, + "affinity table full; routing without session affinity" + ); + return Ok(None); + } + Err(error) => return Err(affinity_error(error)), + } + } + } + + /// Bind the held session to `dispatched`. A `Hard` rejection whose bound + /// worker departed after [`Self::hold_session`] checked it is not a client + /// fault: the session is re-initialized on the dispatched worker instead. + /// This runs with the booking held, so it never waits on another request's + /// initialization (which may itself be queued behind this booking); when + /// one is in flight the booking stands unbound, `None`. + pub(super) fn commit_session( + &self, + table: &SessionAffinity, + hold: Hold, + session_id: &str, + dispatched: WorkerWithDpRank, + key: &RoutingPartitionId, + ) -> Result, SelectionError> { + let bound = hold.target(); + let dispatched = WorkerAffinityTarget::new(dispatched.worker_id, Some(dispatched.dp_rank)); + match table.commit(hold, dispatched) { + Ok(lease) => Ok(Some(lease)), + Err(error) => { + let departed = + bound.is_some_and(|target| !self.catalog.is_schedulable(target.worker_id, key)); + if !departed { + return Err(affinity_error(error)); + } + // `commit` invalidated the stale binding. + match table.try_acquire(session_id, None) { + Ok(AcquireStep::Held(hold)) => table + .commit(hold, dispatched) + .map(Some) + .map_err(affinity_error), + Ok(AcquireStep::Wait(_)) | Err(AffinityError::ResourceExhausted(_)) => Ok(None), + Err(error) => Err(affinity_error(error)), + } + } + } + } + + /// Record a booked routing decision into the partition's approximate + /// indexer (side or primary). The booking already landed, so a failure + /// here only costs predicted cache credit; it is logged, not returned. + pub(super) async fn record_routing_decision( + &self, + entry: &SelectionEntry, + worker: WorkerWithDpRank, + block_hashes: Vec, + ) { + if block_hashes.is_empty() { + return; + } + if let Err(error) = entry + .indexer + .record_routing_decision( + worker, + RoutingDecisionHashes::from_local_hashes(block_hashes), + ) + .await + { + tracing::warn!( + %error, + key = %entry.key, + worker_id = worker.worker_id, + dp_rank = worker.dp_rank, + "Failed to record routing decision into approximate indexer" + ); + } + } + + /// Normalize the prompt and gather cache signals. The indexer lookup and + /// the optional shared-cache lookup run concurrently; the shared cache is + /// consulted only when `query_shared_cache` is set, a shared cache is + /// attached, and the prompt carries raw `token_ids`. + /// `tracking_assume_kv_reuse` is `None` when the caller does not track + /// active blocks. + pub(super) async fn prepare_selection_inputs( + &self, + entry: &SelectionEntry, + prompt: &PromptView<'_>, + tracking_assume_kv_reuse: Option, + query_shared_cache: bool, + retain_kv_transfer_chain: bool, + lookup: &mut Option, + ) -> Result { + let normalized = prompt.normalize_for_selection( + entry.block_size, + entry.is_eagle, + tracking_assume_kv_reuse.map(|assume_kv_reuse| TrackingHashInput { + context: &self.tracking_hash, + scope: tracking_scope(entry), + assume_kv_reuse, + }), + )?; + let indexer_lookup = async { + let started = Instant::now(); + let tiered = if normalized.block_hashes.is_empty() { + Ok(TieredMatchDetails::default()) + } else { + entry + .indexer + .find_tiered_matches_with_options( + normalized.block_hashes.clone(), + LowerTierQueryOptions { + retain_kv_transfer_chain, + }, + ) + .instrument(tracing::info_span!("kv_router.find_matches")) + .await + .map_err(SelectionError::Indexer) + }; + (tiered, started.elapsed()) + }; + let shared_cache = query_shared_cache + .then_some(self.host.cache.shared.as_deref()) + .flatten() + .zip(prompt.token_ids); + // `(hits, duration, failed)`; all `None`/`false` without a shared cache. + let shared_cache_lookup = async { + let Some((shared_cache, tokens)) = shared_cache else { + return (None, None, false); + }; + let started = Instant::now(); + let result = shared_cache + .check_blocks(tokens, entry.block_size, prompt.cache_namespace) + .instrument(tracing::info_span!("kv_router.shared_cache_check")) + .await; + let elapsed = started.elapsed(); + match result { + Ok(hits) => (Some(hits), Some(elapsed), false), + Err(error) => { + tracing::warn!(%error, "Shared cache query failed, ignoring"); + (None, Some(elapsed), true) + } + } + }; + let lookups_started = Instant::now(); + let ((tiered, indexer), (shared_cache_hits, shared_cache, shared_cache_error)) = + tokio::join!(indexer_lookup, shared_cache_lookup); + let timings = LookupTimings { + block_hashing: normalized.block_hashing, + seq_hashing: normalized.seq_hashing, + lookups: lookups_started.elapsed(), + indexer, + shared_cache, + shared_cache_error, + }; + // Recorded before `?` so a shared-cache failure is still accounted + // when the index lookup fails. + *lookup = Some(timings); + let tiered = tiered?; + let overlap = + OverlapAnalysis::new(&self.kv_router_config, entry.block_size, &tiered).signals(); + let kv_transfer_candidates = retain_kv_transfer_chain + .then(|| tiered.kv_transfer_candidates().cloned()) + .flatten(); + drop(tiered); + Ok(PreparedSelectionInputs { + block_hashes: normalized.block_hashes, + sequence_hashes: normalized.sequence_hashes, + isl_tokens: normalized.isl_tokens, + overlap, + shared_cache_hits, + kv_transfer_candidates, + }) + } +} + +/// The core's use of the session table for a request: a session is steered +/// only when the caller did not pin or target a worker explicitly. +fn session_binding( + context: Option<&SessionContext>, + is_steerable: bool, + is_booking: bool, +) -> SessionBinding { + match context.filter(|_| is_steerable) { + Some(context) if is_booking => SessionBinding::Managed { + session_id: context.session_id().to_string(), + }, + Some(context) => SessionBinding::Query { + session_id: context.session_id().to_string(), + }, + None => SessionBinding::None, + } +} diff --git a/lib/kv-router/src/services/selection/core/tests.rs b/lib/kv-router/src/services/selection/core/tests.rs new file mode 100644 index 000000000000..03aacee52f8e --- /dev/null +++ b/lib/kv-router/src/services/selection/core/tests.rs @@ -0,0 +1,2803 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use super::super::affinity::SessionAffinityMode; +use super::super::input::PromptRequest; +use super::hint::transfer_hint_for_selection; +use super::reservations::{ + Reservation, ReservationClaim, ReservationIndexObserver, sweep_reservation_index, +}; +use super::*; +use crate::protocols::ActiveSequenceEventData; +use crate::protocols::{RoutingConstraints, StorageTier}; +use crate::services::common::replica_sync::HostReplicaChannels; +use crate::services::indexer::backend::test_util::store_event; +use std::collections::HashSet; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::thread::sleep; +use std::time::Duration; + +fn test_config(use_kv_events: bool) -> crate::config::KvRouterConfig { + crate::config::KvRouterConfig { + use_kv_events, + router_queue_threshold: None, + ..Default::default() + } +} + +fn worker(worker_id: WorkerId) -> WorkerRequest { + WorkerRequest { + worker_id, + model_name: "model".to_string(), + routing_group: "default".to_string(), + endpoint: Some(format!("http://worker-{worker_id}:8000")), + kv_events_endpoint: None, + kv_events_endpoints: HashMap::new(), + replay_endpoint: None, + block_size: Some(4), + data_parallel_start_rank: None, + data_parallel_size: None, + max_num_batched_tokens: Some(1024), + total_kv_blocks: None, + stable_routing_id: None, + is_eagle: None, + taints: HashSet::new(), + topology_domains: HashMap::new(), + kv_transfer_domain: None, + kv_transfer_enforcement: None, + kv_transfer_preferred_weight: None, + router_hint_worker_type: None, + router_hint_source_control_endpoints: HashMap::new(), + kv_event_source_mode: None, + } +} + +fn worker_with_kv_events(worker_id: WorkerId) -> WorkerRequest { + WorkerRequest { + kv_events_endpoint: Some("tcp://127.0.0.1:5557".to_string()), + ..worker(worker_id) + } +} + +fn prompt() -> PromptRequest { + PromptRequest { + token_ids: Some(vec![1, 2, 3, 4]), + mm_routing_info: None, + block_mm_infos: None, + block_hashes: None, + sequence_hashes: None, + isl_tokens: None, + lora_name: None, + cache_namespace: None, + is_eagle: None, + } +} + +fn select_request() -> SelectRequest { + SelectRequest { + model_name: "model".to_string(), + routing_group: "default".to_string(), + selection_id: None, + prompt: prompt(), + router_config_override: None, + expected_output_tokens: None, + priority_jump: None, + strict_priority: None, + session_id: None, + session_context: None, + affinity_target: None, + pinned_worker: None, + allowed_worker_ids: None, + routing_constraints: RoutingConstraints::default(), + advisory: false, + } +} + +fn reserve_request(selection_id: &str) -> SelectAndReserveRequest { + SelectAndReserveRequest { + model_name: "model".to_string(), + routing_group: "default".to_string(), + selection_id: Some(selection_id.to_string()), + prompt: prompt(), + router_config_override: None, + expected_output_tokens: None, + priority_jump: None, + strict_priority: None, + session_id: None, + session_context: None, + affinity_target: None, + pinned_worker: None, + allowed_worker_ids: None, + routing_constraints: RoutingConstraints::default(), + } +} + +async fn wait_until(what: &str, mut condition: impl FnMut() -> bool) { + tokio::time::timeout(Duration::from_secs(2), async { + while !condition() { + tokio::task::yield_now().await; + } + }) + .await + .unwrap_or_else(|_| panic!("timed out waiting for {what}")); +} + +async fn wait_for_pending_selection(core: &SelectionCore) { + wait_until("pending selection", || { + core.loads(Some("model"), Some("default"))[0].pending_count == 1 + }) + .await; +} + +fn assert_shutdown_error(error: SelectionError) { + assert!(matches!( + error, + SelectionError::NotReady(message) + if message == "selection service is shutting down" + )); +} + +#[test] +fn parent_cancel_cancels_core() { + let parent = CancellationToken::new(); + let core = SelectionCore::try_new_local( + test_config(false), + 1, + parent.clone(), + SelectionCacheConfig::default(), + ) + .expect("valid test config"); + + assert!(!core.cancel_token.is_cancelled()); + parent.cancel(); + assert!(core.cancel_token.is_cancelled()); +} + +#[tokio::test] +async fn selection_setup_uses_worker_type_label() { + for (worker_type, expected_label) in [ + (WorkerType::Prefill, "prefill"), + (WorkerType::Decode, "decode"), + (WorkerType::Encode, "encode"), + (WorkerType::Aggregated, "aggregated"), + ] { + let config = test_config(false); + let tracking_hash = Arc::new( + TrackingHashContext::from_config(&config).expect("valid tracking hash configuration"), + ); + let core = SelectionCore::new_inner( + config, + 1, + CancellationToken::new(), + None, + None, + SelectionHost::default(), + worker_type, + true, + SelectionCacheConfig::default(), + tracking_hash, + IndexerPolicy::from_router_config(&test_config(false)).expect("indexer policy"), + None, + ); + + core.upsert_worker(worker(1)).await.expect("worker upsert"); + let entry = core + .entry(&RoutingPartitionId::new("model", "default")) + .expect("selection entry"); + assert_eq!( + entry.scheduler.worker_type(), + expected_label, + "{worker_type}" + ); + } +} + +fn core_with_host(host: SelectionHost) -> SelectionCore { + core_with_host_and_policy(host, None) +} + +fn core_with_host_and_policy( + host: SelectionHost, + policy_factory: Option, +) -> SelectionCore { + let config = test_config(false); + let tracking_hash = Arc::new( + TrackingHashContext::from_config(&config).expect("valid tracking hash configuration"), + ); + let indexer_policy = IndexerPolicy::from_router_config(&config).expect("indexer policy"); + SelectionCore::new_inner( + config, + 1, + CancellationToken::new(), + None, + policy_factory, + host, + WorkerType::Aggregated, + true, + SelectionCacheConfig::default(), + tracking_hash, + indexer_policy, + None, + ) +} + +async fn wait_for_overlap( + core: &SelectionCore, + request: impl Fn() -> SelectRequest, +) -> SelectResponse { + tokio::time::timeout(Duration::from_secs(2), async { + loop { + let response = core.select(request()).await.expect("select"); + if response.overlap.longest_matched > 0 { + return response; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("approximate indexer never credited the booked prompt") +} + +#[tokio::test] +async fn bookings_populate_the_approximate_primary_without_kv_events() { + // use_kv_events=false: the primary is approximate and bookings feed it. + let core = SelectionCore::try_new_local( + test_config(false), + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ) + .expect("valid test config"); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + core.upsert_worker(worker(2)).await.expect("worker upsert"); + + // Query-only selection records nothing. + let first = core.select(select_request()).await.expect("select"); + assert_eq!(first.overlap.longest_matched, 0); + + // select_and_reserve records the routed prefix for the chosen worker. + let booked = core + .select_and_reserve(reserve_request("booked")) + .await + .expect("reserve"); + let credited = wait_for_overlap(&core, || { + let mut request = select_request(); + request.allowed_worker_ids = Some(HashSet::from([booked.worker_id])); + request + }) + .await; + assert_eq!(credited.worker_id, booked.worker_id); + assert_eq!(credited.overlap.longest_matched, 4); + + // The cached replay path records too, on a different prompt. + let prompt_b = || PromptRequest { + token_ids: Some(vec![5, 6, 7, 8]), + ..PromptRequest::default() + }; + let mut request = select_request(); + request.prompt = prompt_b(); + request.selection_id = Some("cached".to_string()); + request.allowed_worker_ids = Some(HashSet::from([2])); + core.select(request).await.expect("select"); + core.create_reservation(ReservationRequest { + model_name: "model".to_string(), + routing_group: "default".to_string(), + selection_id: "cached".to_string(), + worker_id: None, + dp_rank: None, + prompt: PromptRequest::default(), + router_config_override: None, + expected_output_tokens: None, + effective_prefill_tokens: None, + track_prefill_tokens: None, + }) + .await + .expect("cached reservation"); + let credited = wait_for_overlap(&core, || { + let mut request = select_request(); + request.prompt = prompt_b(); + request.allowed_worker_ids = Some(HashSet::from([2])); + request + }) + .await; + assert_eq!(credited.worker_id, 2); + + // And the explicit reservation form. + let prompt_c = || PromptRequest { + token_ids: Some(vec![9, 10, 11, 12]), + ..PromptRequest::default() + }; + core.create_reservation(ReservationRequest { + model_name: "model".to_string(), + routing_group: "default".to_string(), + selection_id: "explicit".to_string(), + worker_id: Some(1), + dp_rank: None, + prompt: prompt_c(), + router_config_override: None, + expected_output_tokens: None, + effective_prefill_tokens: None, + track_prefill_tokens: None, + }) + .await + .expect("explicit reservation"); + let credited = wait_for_overlap(&core, || { + let mut request = select_request(); + request.prompt = prompt_c(); + request.allowed_worker_ids = Some(HashSet::from([1])); + request + }) + .await; + assert_eq!(credited.worker_id, 1); +} + +#[tokio::test] +async fn unreachable_remote_indexer_is_reported_not_ready() { + let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind"); + let base_url = format!("http://{}", listener.local_addr().expect("addr")); + drop(listener); + let config = test_config(true); + let tracking_hash = Arc::new( + TrackingHashContext::from_config(&config).expect("valid tracking hash configuration"), + ); + let indexer_policy = IndexerPolicy::from_router_config(&config) + .expect("indexer policy") + .with_remote_indexer(base_url) + .expect("remote policy"); + let core = SelectionCore::new_inner( + config, + 1, + CancellationToken::new(), + None, + None, + SelectionHost::default(), + WorkerType::Aggregated, + true, + SelectionCacheConfig::default(), + tracking_hash, + indexer_policy, + None, + ); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + assert!(matches!( + core.select(select_request()).await, + Err(SelectionError::Indexer( + crate::indexer::KvRouterError::IndexerOffline + )) + )); +} + +#[tokio::test] +async fn remote_indexer_serves_selection_without_local_kv_listeners() { + use crate::indexer::KvIndexerInterface; + use crate::protocols::{BlockHashOptions, StorageTier, compute_block_hash_for_seq}; + use crate::services::indexer::registry::WorkerRegistry; + use crate::services::indexer::server::spawn_test_indexer_server; + + // The standalone indexer holds worker 2's cache for the test prompt. + let key = RoutingPartitionId::new("model", "default"); + let served = Arc::new(WorkerRegistry::new(1)); + let served_indexer = served.get_or_create_indexer(key.clone(), 4); + let hashes: Vec = + compute_block_hash_for_seq(&[1, 2, 3, 4], 4, BlockHashOptions::default()) + .into_iter() + .map(|hash| hash.0) + .collect(); + served_indexer + .apply_event_routed(store_event(2, 0, 1, &[], &hashes, StorageTier::Device)) + .await + .unwrap(); + if let Indexer::Single { primary, .. } = &served_indexer { + let _ = primary.flush().await; + } + let (base_url, server) = spawn_test_indexer_server(served).await; + + // use_kv_events=true, but the primary is remote: workers need no + // kv_events endpoints and no ZMQ listener is started here. + let config = test_config(true); + let tracking_hash = Arc::new( + TrackingHashContext::from_config(&config).expect("valid tracking hash configuration"), + ); + let indexer_policy = IndexerPolicy::from_router_config(&config) + .expect("indexer policy") + .with_remote_indexer(base_url) + .expect("remote policy"); + let core = SelectionCore::new_inner( + config, + 1, + CancellationToken::new(), + None, + None, + SelectionHost::default(), + WorkerType::Aggregated, + true, + SelectionCacheConfig::default(), + tracking_hash, + indexer_policy, + None, + ); + assert!(!core.listens_for_kv_events); + for worker_id in [1, 2] { + let record = core + .upsert_worker(worker(worker_id)) + .await + .expect("worker upsert"); + assert_eq!(record.lifecycle, WorkerLifecycle::Schedulable, "{record:?}"); + } + + let response = core.select(select_request()).await.expect("select"); + assert_eq!( + response.worker_id, 2, + "remote cache credit steers selection" + ); + assert_eq!(response.overlap.longest_matched, 4); + + core.delete_worker(2).await.expect("delete worker"); + server.abort(); +} + +/// Two event-driven workers; worker 1 holds every block of an 8-token +/// prompt. `configure` adjusts each registration before upsert. +async fn hint_fixture( + configure: impl Fn(&mut WorkerRequest), +) -> (SelectionCore, Arc, Vec) { + use crate::indexer::KvIndexerInterface; + use crate::protocols::{BlockHashOptions, compute_block_hash_for_seq}; + + let core = SelectionCore::try_new_local( + test_config(true), + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ) + .expect("valid test config"); + for worker_id in [1, 2] { + let mut request = worker_with_kv_events(worker_id); + configure(&mut request); + core.upsert_worker(request).await.expect("worker upsert"); + } + let entry = core + .entry(&RoutingPartitionId::new("model", "default")) + .expect("entry"); + let tokens: Vec = (1..=8).collect(); + let hashes: Vec = compute_block_hash_for_seq(&tokens, 4, BlockHashOptions::default()) + .into_iter() + .map(|hash| hash.0) + .collect(); + entry + .indexer + .apply_event_routed(store_event(1, 0, 1, &[], &hashes, StorageTier::Device)) + .await + .unwrap(); + if let Indexer::Single { primary, .. } = &entry.indexer { + let _ = primary.flush().await; + } + (core, entry, tokens) +} + +fn hint_capable(request: &mut WorkerRequest) { + request.router_hint_worker_type = Some("decode".to_string()); + request.router_hint_source_control_endpoints = + HashMap::from([(0, format!("tcp://worker-{}:9000", request.worker_id))]); +} + +async fn reserve_pinned( + core: &SelectionCore, + selection_id: &str, + tokens: &[u32], + worker_id: WorkerId, +) -> SelectResponse { + let mut request = reserve_request(selection_id); + request.prompt = PromptRequest { + token_ids: Some(tokens.to_vec()), + ..PromptRequest::default() + }; + request.pinned_worker = Some(WorkerWithDpRank::new(worker_id, 0)); + core.select_and_reserve(request).await.expect("reserve") +} + +#[tokio::test] +async fn booked_selection_attaches_router_hint_from_a_better_source() { + let (core, entry, tokens) = hint_fixture(hint_capable).await; + assert!(entry.indexer.supports_kv_transfer_chain_retention()); + + // Booking on worker 2: worker 1 is a same-role source with a longer prefix. + let response = reserve_pinned(&core, "to-worker-2", &tokens, 2).await; + let hint = response.kv_hint.expect("router hint for worker 2"); + assert_eq!(hint.message_id, "to-worker-2"); + assert_eq!(hint.actions[0].action_type, "kv.fetch"); + let payload: KvSourceLocationsPayload = + serde_json::from_value(serde_json::to_value(&hint.actions[0].payload).unwrap()).unwrap(); + assert_eq!(payload.source_control_endpoint, "tcp://worker-1:9000"); + assert_eq!(payload.block_hashes.len(), 2); + + // Booking on worker 1 itself: nothing holds a longer prefix. + let response = reserve_pinned(&core, "to-worker-1", &tokens, 1).await; + assert!(response.kv_hint.is_none()); + + // Query-only selections never carry a hint. + let mut request = select_request(); + request.prompt = PromptRequest { + token_ids: Some(tokens.clone()), + ..PromptRequest::default() + }; + request.pinned_worker = Some(WorkerWithDpRank::new(2, 0)); + let response = core.select(request).await.expect("select"); + assert!(response.kv_hint.is_none()); +} + +#[tokio::test] +async fn state_agent_workers_are_not_router_hint_sources() { + let (core, _entry, tokens) = hint_fixture(|request| { + hint_capable(request); + if request.worker_id == 1 { + request.kv_event_source_mode = Some("state_agent_v2".to_string()); + } + }) + .await; + // Worker 1 holds the prefix but reports through a state agent. + let response = reserve_pinned(&core, "to-worker-2", &tokens, 2).await; + assert!(response.kv_hint.is_none()); +} + +#[tokio::test] +async fn router_hint_needs_capable_workers() { + let (core, _entry, tokens) = hint_fixture(|_| {}).await; + let response = reserve_pinned(&core, "plain", &tokens, 2).await; + assert!(response.kv_hint.is_none()); +} + +#[tokio::test] +async fn event_driven_indexer_does_not_record_bookings() { + let core = SelectionCore::try_new_local( + test_config(true), + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ) + .expect("valid test config"); + core.upsert_worker(worker_with_kv_events(1)) + .await + .expect("worker upsert"); + core.select_and_reserve(reserve_request("booked")) + .await + .expect("reserve"); + core.free_reservation("booked").await.expect("free"); + for _ in 0..3 { + let response = core.select(select_request()).await.expect("select"); + assert_eq!(response.overlap.longest_matched, 0); + tokio::task::yield_now().await; + } +} + +/// Picker that records what worker selection saw and always takes row 0. +struct CapturingPicker { + observed: Arc>>, +} + +#[derive(Debug, Clone)] +struct SelectionObservation { + session_context: Option, + shared_beyond_device_blocks: Vec, +} + +impl crate::scheduling::selector::WorkerPicker for CapturingPicker { + fn required_worker_inputs(&self) -> crate::scheduling::selector::WorkerInputs { + crate::scheduling::selector::WorkerInputs::CACHE + } + + fn pick( + &mut self, + context: &crate::scheduling::selector::WorkerSelectionContext<'_>, + input: crate::scheduling::selector::WorkerInputView<'_>, + ) -> Result { + self.observed.lock().push(SelectionObservation { + session_context: context.session_context().cloned(), + shared_beyond_device_blocks: input + .cache() + .expect("CACHE inputs requested") + .iter() + .map(|cache| cache.shared_beyond_device_blocks()) + .collect(), + }); + Ok(0) + } +} + +fn capturing_policy_factory() -> ( + WorkerSelectionPolicyFactory, + Arc>>, +) { + let observed = Arc::new(parking_lot::Mutex::new(Vec::new())); + let factory_observed = Arc::clone(&observed); + let factory: WorkerSelectionPolicyFactory = Arc::new(move |config, worker_type, _partition| { + WorkerSelectionPolicy::new( + config.clone(), + worker_type.as_str(), + Vec::new(), + Box::new(CapturingPicker { + observed: Arc::clone(&factory_observed), + }), + ) + }); + (factory, observed) +} + +type SharedCacheCalls = Arc, u32, Option)>>>; + +/// Shared cache that reports every block as a hit and records each query. +struct RecordingSharedCache { + calls: SharedCacheCalls, +} + +#[async_trait::async_trait] +impl SharedKvCache for RecordingSharedCache { + async fn check_blocks( + &self, + tokens: &[u32], + block_size: u32, + cache_namespace: Option<&str>, + ) -> Result { + self.calls.lock().push(( + tokens.to_vec(), + block_size, + cache_namespace.map(str::to_string), + )); + let blocks = (tokens.len() / block_size as usize) as u32; + Ok(SharedCacheHits::from_hits(&vec![true; blocks as usize])) + } +} + +struct OnlyWorkerForLora { + worker_id: WorkerId, +} + +impl LoraWorkerFilter for OnlyWorkerForLora { + fn filter_worker_ids_for_lora( + &self, + _lora_name: &str, + available: &[WorkerId], + ) -> Vec { + available + .iter() + .copied() + .filter(|id| *id == self.worker_id) + .collect() + } +} + +#[tokio::test] +async fn injected_lora_filter_narrows_candidates() { + let core = core_with_host(SelectionHost { + eligibility: HostEligibility { + lora_worker_filter: Some(Arc::new(OnlyWorkerForLora { worker_id: 2 })), + }, + ..SelectionHost::default() + }); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + core.upsert_worker(worker(2)).await.expect("worker upsert"); + + // LoRA request: only the filter's worker is eligible. + for _ in 0..4 { + let mut request = select_request(); + request.prompt.lora_name = Some("adapter-a".to_string()); + let response = core.select(request).await.expect("select"); + assert_eq!(response.worker_id, 2); + } + + // The filter never widens the caller's allow-set: an allow-set that + // excludes the filter's worker is preserved as-is. + let mut request = select_request(); + request.prompt.lora_name = Some("adapter-a".to_string()); + request.allowed_worker_ids = Some(HashSet::from([1])); + let response = core.select(request).await.expect("select"); + assert_eq!(response.worker_id, 1); + + // A pinned worker inside the universe survives the filter. + let mut request = select_request(); + request.prompt.lora_name = Some("adapter-a".to_string()); + request.pinned_worker = Some(WorkerWithDpRank::new(1, 0)); + let response = core.select(request).await.expect("select"); + assert_eq!(response.worker_id, 1); +} + +#[tokio::test] +async fn shared_cache_hits_reach_worker_selection() { + let calls: SharedCacheCalls = Arc::new(parking_lot::Mutex::new(Vec::new())); + let (factory, observed) = capturing_policy_factory(); + let core = core_with_host_and_policy( + SelectionHost { + cache: HostCache { + shared: Some(Arc::new(RecordingSharedCache { + calls: Arc::clone(&calls), + })), + ..HostCache::default() + }, + ..SelectionHost::default() + }, + Some(factory), + ); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + + let mut request = select_request(); + request.prompt.cache_namespace = Some("tenant-a".to_string()); + core.select(request).await.expect("select"); + + assert_eq!( + calls.lock().as_slice(), + &[(vec![1, 2, 3, 4], 4, Some("tenant-a".to_string()))] + ); + let observations = observed.lock().clone(); + assert_eq!(observations.len(), 1); + // One block in the prompt, no device overlap, so the whole prompt is a + // shared-cache hit beyond the device prefix. + assert_eq!(observations[0].shared_beyond_device_blocks, vec![1]); + + // Load projection does not consult the shared cache. + core.potential_loads(PotentialLoadsRequest { + model_name: "model".to_string(), + routing_group: "default".to_string(), + prompt: prompt(), + router_config_override: None, + }) + .await + .expect("potential loads"); + assert_eq!(calls.lock().len(), 1); + + // Prompts without raw tokens cannot be checked against the shared cache. + let mut request = select_request(); + request.prompt = PromptRequest { + token_ids: None, + block_hashes: Some(vec![11]), + sequence_hashes: Some(vec![101]), + isl_tokens: Some(4), + ..PromptRequest::default() + }; + core.select(request).await.expect("select"); + assert_eq!(calls.lock().len(), 1); + assert_eq!(observed.lock()[1].shared_beyond_device_blocks, vec![0]); +} + +#[tokio::test] +async fn session_context_reaches_worker_selection() { + use super::super::types::SelectionSessionContext; + + let (factory, observed) = capturing_policy_factory(); + let core = core_with_host_and_policy(SelectionHost::default(), Some(factory)); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + + let mut request = select_request(); + request.session_id = Some("ignored-legacy".to_string()); + request.session_context = Some(SelectionSessionContext { + session_id: "child-session".to_string(), + parent_session_id: Some("root-session".to_string()), + session_final: Some(true), + input_trigger: Some(super::super::types::SelectionInputTrigger::ToolResult), + }); + core.select(request).await.expect("select"); + + let mut request = reserve_request("legacy-session-reservation"); + request.session_id = Some("legacy-only".to_string()); + core.select_and_reserve(request) + .await + .expect("select and reserve"); + + let observations = observed.lock(); + let context = observations[0] + .session_context + .as_ref() + .expect("structured session context"); + assert_eq!(context.session_id(), "child-session"); + assert_eq!(context.parent_session_id(), Some("root-session")); + assert_eq!(context.session_final(), Some(true)); + assert_eq!( + context.input_trigger(), + Some(crate::scheduling::WorkerSelectionInputTrigger::ToolResult) + ); + + let legacy = observations[1] + .session_context + .as_ref() + .expect("legacy session context"); + assert_eq!(legacy.session_id(), "legacy-only"); + assert_eq!(legacy.parent_session_id(), None); +} + +#[tokio::test] +async fn full_affinity_table_routes_without_pinning() { + let core = SelectionCore::try_new_local( + test_config(false), + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ) + .expect("valid test config"); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + let entry = core + .entry(&RoutingPartitionId::new("model", "default")) + .expect("entry"); + let table = + SessionAffinity::new_with_limits(Duration::from_secs(60), 1, 256).expect("affinity table"); + assert!(entry.affinity.set(table).is_ok()); + + for (selection_id, session_id) in [("first", "s1"), ("second", "s2")] { + let mut request = reserve_request(selection_id); + request.session_id = Some(session_id.to_string()); + core.select_and_reserve(request) + .await + .expect("a full affinity table must not fail selection"); + } + assert_eq!( + core.reservation_index + .read() + .values() + .filter(|r| r._affinity_lease.is_some()) + .count(), + 1 + ); +} + +#[tokio::test] +async fn injected_availability_provider_restricts_selection() { + let available: Arc>>>> = + Arc::new(parking_lot::Mutex::new(None)); + let provider_state = Arc::clone(&available); + let core = core_with_host(SelectionHost { + load: HostLoad { + available_workers: Some(Arc::new(move || provider_state.lock().clone())), + ..HostLoad::default() + }, + ..SelectionHost::default() + }); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + core.upsert_worker(worker(2)).await.expect("worker upsert"); + + for only in [1, 2] { + *available.lock() = Some(Arc::new(HashSet::from([only]))); + for _ in 0..4 { + let response = core.select(select_request()).await.expect("select"); + assert_eq!(response.worker_id, only); + } + } +} + +#[tokio::test] +async fn injected_overload_provider_excludes_worker() { + let core = core_with_host(SelectionHost { + load: HostLoad { + overloaded_workers: Some(Arc::new(|| Some(HashSet::from([1])))), + ..HostLoad::default() + }, + ..SelectionHost::default() + }); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + core.upsert_worker(worker(2)).await.expect("worker upsert"); + + for _ in 0..4 { + let response = core.select(select_request()).await.expect("select"); + assert_eq!(response.worker_id, 2); + } +} + +#[test] +fn shutdown_keeps_parent_alive() { + let parent = CancellationToken::new(); + let core = SelectionCore::try_new_local( + test_config(false), + 1, + parent.clone(), + SelectionCacheConfig::default(), + ) + .expect("valid test config"); + + core.shutdown(); + + assert!(core.cancel_token.is_cancelled()); + assert!(!parent.is_cancelled()); +} + +#[tokio::test] +async fn multi_rank_worker_with_replay_endpoint_is_incomplete() { + let core = SelectionCore::try_new_local( + test_config(true), + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ) + .expect("valid test config"); + + let record = core + .upsert_worker(WorkerRequest { + data_parallel_size: Some(2), + kv_events_endpoints: HashMap::from([ + (0, "tcp://127.0.0.1:5557".to_string()), + (1, "tcp://127.0.0.1:5558".to_string()), + ]), + replay_endpoint: Some("tcp://127.0.0.1:5600".to_string()), + ..worker(1) + }) + .await + .expect("worker upsert"); + assert_eq!(record.lifecycle, WorkerLifecycle::Incomplete, "{record:?}"); + assert!(!core.indexer_registry.has_worker(1)); +} + +#[tokio::test] +async fn reupsert_recreates_a_listener_lost_to_a_cancelled_update() { + use crate::indexer::KvIndexerInterface; + + let core = SelectionCore::try_new_local( + test_config(true), + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ) + .expect("valid test config"); + core.upsert_worker(worker_with_kv_events(1)) + .await + .expect("worker upsert"); + let entry = core + .entry(&RoutingPartitionId::new("model", "default")) + .expect("entry"); + let hashes = [11u64, 12]; + entry + .indexer + .apply_event_routed(store_event(1, 0, 1, &[], &hashes, StorageTier::Device)) + .await + .unwrap(); + if let Indexer::Single { primary, .. } = &entry.indexer { + let _ = primary.flush().await; + } + // An endpoint update cancelled right after removing the listener leaves + // the catalog record schedulable, the listener gone, and its blocks indexed. + core.indexer_registry.forget_listener(1, 0); + + let record = core + .upsert_worker(worker_with_kv_events(1)) + .await + .expect("worker re-upsert"); + assert_eq!(record.lifecycle, WorkerLifecycle::Schedulable, "{record:?}"); + assert!(core.indexer_registry.has_listener(1, 0)); + // The purge is queued to the indexer thread; wait for it to land. + tokio::time::timeout(Duration::from_secs(2), async { + loop { + let overlap = entry + .indexer + .find_matches(hashes.iter().copied().map(LocalBlockHash).collect()) + .await + .expect("find matches"); + if overlap.scores.is_empty() { + return; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("stale blocks survived the re-upsert"); +} + +#[tokio::test] +async fn shutdown_cancels_listeners() { + let parent = CancellationToken::new(); + let core = SelectionCore::try_new_local( + test_config(true), + 1, + parent, + SelectionCacheConfig::default(), + ) + .expect("valid test config"); + + let record = core + .upsert_worker(worker_with_kv_events(1)) + .await + .expect("worker upsert"); + assert_eq!(record.lifecycle, WorkerLifecycle::Schedulable); + assert_eq!(core.indexer_registry.listener_cancelled(1, 0), Some(false)); + + core.shutdown(); + assert_eq!(core.indexer_registry.listener_cancelled(1, 0), Some(true)); +} + +#[tokio::test] +async fn upsert_moves_global_worker_id_between_routing_groups() { + let core = SelectionCore::try_new_local( + test_config(true), + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ) + .expect("valid test config"); + let mut group_a = worker_with_kv_events(1); + group_a.routing_group = "group-a".to_string(); + core.upsert_worker(group_a).await.expect("group A upsert"); + assert_eq!( + core.indexer_registry + .list_filtered(Some("model"), Some("group-a")) + .len(), + 1 + ); + + let mut group_b = worker_with_kv_events(1); + group_b.routing_group = "group-b".to_string(); + core.upsert_worker(group_b).await.expect("group B upsert"); + + assert!(core.list_workers(Some("model"), Some("group-a")).is_empty()); + assert_eq!(core.list_workers(Some("model"), Some("group-b")).len(), 1); + assert!( + core.indexer_registry + .list_filtered(Some("model"), Some("group-a")) + .is_empty() + ); + assert_eq!( + core.indexer_registry + .list_filtered(Some("model"), Some("group-b")) + .len(), + 1 + ); + + let mut select_a = select_request(); + select_a.routing_group = "group-a".to_string(); + assert!(matches!( + core.select(select_a).await, + Err(SelectionError::NotReady(_)) + )); + let mut select_b = select_request(); + select_b.routing_group = "group-b".to_string(); + assert_eq!(core.select(select_b).await.unwrap().worker_id, 1); + + core.delete_worker(1).await.expect("delete group B worker"); + assert!( + core.indexer_registry + .list_filtered(Some("model"), Some("group-b")) + .is_empty() + ); +} + +#[tokio::test] +async fn shutdown_reports_not_ready_and_rejects_new_work() { + let core = SelectionCore::try_new_local( + test_config(false), + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ) + .expect("valid test config"); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + assert!(core.ready().ready); + + core.shutdown(); + + let ready = core.ready(); + assert!(!ready.ready); + assert_eq!(ready.schedulable_workers, 1); + + let upsert_error = core + .upsert_worker(worker(2)) + .await + .expect_err("upsert should fail after shutdown"); + assert_shutdown_error(upsert_error); + + let patch = serde_json::from_value(serde_json::json!({ + "endpoint": "http://worker-1:9000" + })) + .expect("worker patch"); + let patch_error = core + .patch_worker(1, patch) + .await + .expect_err("patch should fail after shutdown"); + assert_shutdown_error(patch_error); + + let select_error = core + .select(select_request()) + .await + .expect_err("selection should fail after shutdown"); + assert_shutdown_error(select_error); + + let reservation_error = core + .create_reservation(ReservationRequest { + model_name: "model".to_string(), + routing_group: "default".to_string(), + selection_id: "res-after-shutdown".to_string(), + worker_id: Some(1), + dp_rank: None, + prompt: prompt(), + router_config_override: None, + expected_output_tokens: None, + effective_prefill_tokens: None, + track_prefill_tokens: None, + }) + .await + .expect_err("reservation should fail after shutdown"); + assert_shutdown_error(reservation_error); + + assert_eq!(core.list_workers(None, None).len(), 1); + assert_eq!(core.loads(None, None).len(), 1); + let deleted = core + .delete_worker(1) + .await + .expect("delete should remain available after shutdown"); + assert_eq!(deleted.lifecycle, WorkerLifecycle::Unschedulable); +} + +#[tokio::test] +async fn queued_selection_errors_on_shutdown() { + let mut config = test_config(false); + config.router_queue_threshold = Some(0.0); + let core = Arc::new( + SelectionCore::try_new_local( + config, + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ) + .expect("valid test config"), + ); + + let record = core.upsert_worker(worker(1)).await.expect("worker upsert"); + assert_eq!(record.lifecycle, WorkerLifecycle::Schedulable); + core.select_and_reserve(reserve_request("res-a")) + .await + .expect("initial reservation"); + + let queued_core = core.clone(); + let queued = tokio::spawn(async move { queued_core.select(select_request()).await }); + wait_for_pending_selection(&core).await; + + core.shutdown(); + let err = tokio::time::timeout(Duration::from_secs(1), queued) + .await + .expect("queued selection timed out") + .expect("queued selection task panicked") + .expect_err("queued selection should fail"); + + assert!(matches!( + err, + SelectionError::Scheduler(KvSchedulerError::SubscriberShutdown) + )); +} + +#[tokio::test] +async fn booking_is_freed_when_selected_worker_drained_while_queued() { + let mut config = test_config(false); + config.router_queue_threshold = Some(0.0); + let core = Arc::new( + SelectionCore::try_new_local( + config, + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ) + .expect("valid test config"), + ); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + let key = RoutingPartitionId::new("model", "default"); + let entry = core.entry(&key).expect("entry"); + core.select_and_reserve(reserve_request("res-a")) + .await + .expect("initial reservation"); + + let queued_core = core.clone(); + let queued = tokio::spawn(async move { + queued_core + .select_and_reserve(reserve_request("queued")) + .await + }); + wait_for_pending_selection(&core).await; + // First half of `delete_worker`: the worker drains while the request waits. + core.catalog + .set_lifecycle(1, WorkerLifecycle::Draining, Vec::new()); + core.free_reservation("res-a").await.expect("free res-a"); + + let err = tokio::time::timeout(Duration::from_secs(2), queued) + .await + .expect("queued selection timed out") + .expect("task panicked") + .expect_err("drained worker is not schedulable"); + assert!( + matches!(&err, SelectionError::Internal(m) if m.contains("no longer schedulable")), + "{err:?}" + ); + wait_until("booking release", || !entry.scheduler.has_request("queued")).await; +} + +/// The embedded host dispatches by worker id, so a worker that drains +/// while the request queues still comes back selected: the transport +/// reports the departure and migration takes over. The endpoint is the +/// only thing missing. +#[tokio::test] +async fn lease_admission_keeps_a_selection_whose_worker_drained_while_queued() { + let mut config = test_config(false); + config.router_queue_threshold = Some(0.0); + let core = Arc::new( + SelectionCore::try_new_local( + config, + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ) + .expect("valid test config"), + ); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + let key = RoutingPartitionId::new("model", "default"); + let entry = core.entry(&key).expect("entry"); + core.select_and_reserve(reserve_request("res-a")) + .await + .expect("initial reservation"); + + let queued_core = core.clone(); + let queued = tokio::spawn(async move { + let req = reserve_request("queued"); + let run = queued_core + .run_selection(SelectionOperation { + key: RoutingPartitionId::new("model", "default"), + prompt: req.prompt.view(), + router_config_override: None, + expected_output_tokens: None, + priority_jump: 0.0, + strict_priority: 0, + policy_class: None, + session_context: None, + session: SessionBinding::None, + affinity_target: None, + pinned_worker: None, + allowed_worker_ids: None, + routing_constraints: RoutingConstraints::default(), + admission: SelectionAdmission::Lease { + request_id: "queued".to_string(), + }, + track_active_blocks: false, + return_routing_hashes: false, + replay_id: None, + }) + .await; + match run.result { + Ok(SelectionOutcome::Selected(selected)) => ( + selected.response.best_worker.worker_id, + selected.endpoint, + selected.lease.is_some(), + ), + Ok(SelectionOutcome::QueueRejected { .. }) => panic!("queue rejected"), + Err(error) => panic!("lease selection failed: {error:?}"), + } + }); + wait_for_pending_selection(&core).await; + core.catalog + .set_lifecycle(1, WorkerLifecycle::Draining, Vec::new()); + core.free_reservation("res-a").await.expect("free res-a"); + + let (worker_id, endpoint, leased) = tokio::time::timeout(Duration::from_secs(2), queued) + .await + .expect("queued selection timed out") + .expect("task panicked"); + assert_eq!(worker_id, 1); + assert_eq!(endpoint, None); + assert!(leased, "lease admission returns the booking's lease"); + // The lease was dropped with the selection, so the booking is gone. + wait_until("booking release", || !entry.scheduler.has_request("queued")).await; +} + +#[tokio::test] +async fn dropped_selection_future_frees_its_booking() { + let core = SelectionCore::try_new_local( + test_config(false), + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ) + .expect("valid test config"); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + let key = RoutingPartitionId::new("model", "default"); + let entry = core.entry(&key).expect("entry"); + + // Drive the selection by hand so the actor's response is delivered but + // never consumed: poll until the booking exists, then drop the future. + let mut selection = Box::pin(core.select_and_reserve(reserve_request("dropped"))); + let mut context = std::task::Context::from_waker(std::task::Waker::noop()); + wait_until("scheduler booking", || { + if entry.scheduler.has_request("dropped") { + return true; + } + assert!( + selection.as_mut().poll(&mut context).is_pending(), + "selection completed before the booking was observed" + ); + false + }) + .await; + assert!( + core.reservation_index + .read() + .get("dropped") + .is_some_and(|reservation| reservation.booking.is_none()), + "the id is claimed while its booking is in flight" + ); + drop(selection); + + wait_until("booking release", || { + !entry.scheduler.has_request("dropped") + }) + .await; + assert!(core.reservation_index.read().get("dropped").is_none()); +} + +#[tokio::test] +async fn same_selection_id_in_two_partitions_is_a_conflict() { + let core = SelectionCore::try_new_local( + test_config(false), + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ) + .expect("valid test config"); + for (worker_id, routing_group) in [(1, "group-a"), (2, "group-b")] { + let mut request = worker(worker_id); + request.routing_group = routing_group.to_string(); + core.upsert_worker(request).await.expect("worker upsert"); + } + let mut first = reserve_request("shared"); + first.routing_group = "group-a".to_string(); + core.select_and_reserve(first).await.expect("first booking"); + + let mut second = reserve_request("shared"); + second.routing_group = "group-b".to_string(); + let err = core + .select_and_reserve(second) + .await + .expect_err("a live id cannot be booked again"); + assert!(matches!(err, SelectionError::Conflict(_)), "{err:?}"); + + let (entry, _) = core + .indexed_booking("shared") + .expect("first booking indexed"); + assert_eq!(entry.key.routing_group, "group-a"); + assert!(entry.scheduler.has_request("shared")); + assert!( + !core + .entry(&RoutingPartitionId::new("model", "group-b")) + .expect("entry") + .scheduler + .has_request("shared") + ); +} + +#[tokio::test] +async fn explicit_reservation_of_a_live_id_is_a_conflict() { + let core = SelectionCore::try_new_local( + test_config(false), + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ) + .expect("valid test config"); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + core.select_and_reserve(reserve_request("taken")) + .await + .expect("booking"); + let err = core + .create_reservation(ReservationRequest { + model_name: "model".to_string(), + routing_group: "default".to_string(), + selection_id: "taken".to_string(), + worker_id: Some(1), + dp_rank: None, + prompt: prompt(), + router_config_override: None, + expected_output_tokens: None, + effective_prefill_tokens: None, + track_prefill_tokens: None, + }) + .await + .expect_err("explicit booking of a live id"); + assert!(matches!(err, SelectionError::Conflict(_)), "{err:?}"); +} + +#[test] +fn index_observer_replaces_stale_and_claimed_rows_of_its_partition() { + use crate::scheduling::AttemptId; + let index = Arc::new(RwLock::new(HashMap::new())); + let partition = RoutingPartitionId::new("model", "default"); + let observer = ReservationIndexObserver { + index: Arc::clone(&index), + partition: partition.clone(), + host: None, + }; + let booking = |attempt: u64| SchedulerBookingDescriptor { + request_id: "shared".to_string(), + worker: WorkerWithDpRank::new(1, 0), + attempt_id: AttemptId::new(attempt), + }; + let row = |partition: &RoutingPartitionId, booking| Reservation { + partition: partition.clone(), + booking, + _affinity_lease: None, + }; + + // A stale row (its booking expired before the sweep ran) yields to the mirror. + index + .write() + .insert("shared".to_string(), row(&partition, Some(booking(1)))); + observer.admitted(booking(2)); + assert_eq!(index.read()["shared"].booking, Some(booking(2))); + + // A completion for the replaced booking leaves the live mirror alone. + observer.completed(&booking(1)); + assert!(index.read().contains_key("shared")); + observer.completed(&booking(2)); + assert!(!index.read().contains_key("shared")); + + // A claim whose local booking is about to fail also yields, and dropping + // that claim keeps the mirror. + let claim = ReservationClaim { + index: &index, + selection_id: "shared".to_string(), + armed: true, + }; + index + .write() + .insert("shared".to_string(), row(&partition, None)); + observer.admitted(booking(3)); + drop(claim); + assert_eq!(index.read()["shared"].booking, Some(booking(3))); + + // Another partition's row is never replaced. + let other = RoutingPartitionId::new("model", "other"); + index + .write() + .insert("shared".to_string(), row(&other, Some(booking(4)))); + observer.admitted(booking(5)); + assert_eq!(index.read()["shared"].booking, Some(booking(4))); +} + +#[tokio::test] +async fn free_of_an_in_flight_reservation_is_not_found() { + let mut config = test_config(false); + config.router_queue_threshold = Some(0.0); + let core = Arc::new( + SelectionCore::try_new_local( + config, + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ) + .expect("valid test config"), + ); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + core.select_and_reserve(reserve_request("res-a")) + .await + .expect("initial reservation"); + let queued_core = core.clone(); + let queued = tokio::spawn(async move { + queued_core + .select_and_reserve(reserve_request("queued")) + .await + }); + wait_for_pending_selection(&core).await; + + // A free racing the in-flight booking neither frees nor evicts it. + let err = core + .free_reservation("queued") + .await + .expect_err("in-flight id is not a reservation yet"); + assert!(matches!(err, SelectionError::NotFound(_)), "{err:?}"); + assert!(core.reservation_index.read().contains_key("queued")); + + core.free_reservation("res-a").await.expect("free res-a"); + tokio::time::timeout(Duration::from_secs(2), queued) + .await + .expect("queued selection timed out") + .expect("task panicked") + .expect("queued selection books"); + core.free_reservation("queued").await.expect("free queued"); + assert!(core.reservation_index.read().is_empty()); +} + +#[tokio::test] +async fn prefill_complete_is_idempotent_for_a_live_booking() { + let mut config = test_config(false); + config.router_track_prefill_tokens = true; + let core = SelectionCore::try_new_local( + config, + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ) + .expect("valid test config"); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + core.select_and_reserve(reserve_request("live")) + .await + .expect("booking"); + core.prefill_complete("live").await.expect("first mark"); + core.prefill_complete("live") + .await + .expect("a repeated mark on a live booking is not an error"); + let (entry, _) = core.indexed_booking("live").expect("still indexed"); + assert!(entry.scheduler.has_request("live")); + core.free_reservation("live").await.expect("free"); +} + +#[tokio::test] +async fn mirrored_replica_bookings_are_indexed_until_freed() { + let (outbound_tx, _outbound_rx) = mpsc::channel(16); + let (inbound_tx, inbound_rx) = mpsc::channel(16); + let channels = parking_lot::Mutex::new(Some(HostReplicaChannels { + outbound: Some(outbound_tx), + inbound_tx: inbound_tx.clone(), + inbound_rx, + process_id: 7, + })); + let core = core_with_host(SelectionHost { + replication: HostReplication { + channels: Some(Arc::new(move |_| channels.lock().take())), + ..HostReplication::default() + }, + ..SelectionHost::default() + }); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + // Replica events for a worker the scheduler has not registered yet are + // dropped; the first local booking registers it. + core.select_and_reserve(reserve_request("warm")) + .await + .expect("warm booking"); + core.free_reservation("warm").await.expect("free warm"); + let entry = core + .entry(&RoutingPartitionId::new("model", "default")) + .expect("entry"); + let peer_event = |request_id: &str, data| ActiveSequenceEvent { + request_id: request_id.to_string(), + worker: WorkerWithDpRank::new(1, 0), + data, + router_id: 99, + lora_name: None, + }; + let add = |request_id: &str| { + peer_event( + request_id, + ActiveSequenceEventData::AddRequest { + token_sequence: Some(vec![1, 2]), + track_prefill_tokens: false, + expected_output_tokens: None, + prefill_load_hint: None, + }, + ) + }; + + inbound_tx.send(add("peer-a")).await.expect("send"); + inbound_tx.send(add("peer-b")).await.expect("send"); + wait_until("mirrored bookings indexed", || { + core.indexed_booking("peer-a").is_some() && core.indexed_booking("peer-b").is_some() + }) + .await; + + // A lifecycle call on a mirrored booking resolves through the index. + core.free_reservation("peer-a") + .await + .expect("free mirrored"); + assert!(!entry.scheduler.has_request("peer-a")); + assert!(core.indexed_booking("peer-a").is_none()); + + // The peer freeing its own booking removes the mirror. + inbound_tx + .send(peer_event("peer-b", ActiveSequenceEventData::Free)) + .await + .expect("send"); + wait_until("mirror removed", || { + core.indexed_booking("peer-b").is_none() + }) + .await; + assert!(!entry.scheduler.has_request("peer-b")); +} + +#[tokio::test] +async fn lifecycle_operations_resolve_through_the_index() { + let mut config = test_config(false); + config.router_track_prefill_tokens = true; + let core = SelectionCore::try_new_local( + config, + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ) + .expect("valid test config"); + + for (worker_id, routing_group) in [(1, "group-a"), (2, "group-b")] { + let mut request = worker(worker_id); + request.routing_group = routing_group.to_string(); + core.upsert_worker(request).await.expect("worker upsert"); + } + + let entries = core.initialized_entries(); + assert_eq!(entries.len(), 2); + let target = &entries[1]; + let target_group = target.key.routing_group.clone(); + let target_worker = *target + .workers_tx + .borrow() + .keys() + .next() + .expect("target worker"); + + let mut request = reserve_request("later-entry-reservation"); + request.routing_group = target_group.clone(); + core.select_and_reserve(request) + .await + .expect("reserve in later entry"); + + let load = || { + core.loads(Some("model"), Some(&target_group))[0] + .loads + .iter() + .find(|load| load.worker_id == target_worker) + .expect("target load") + .potential_prefill_tokens + }; + assert_eq!(load(), 4); + + core.prefill_complete("later-entry-reservation") + .await + .expect("complete prefill in later entry"); + assert_eq!(load(), 0); + + core.free_reservation("later-entry-reservation") + .await + .expect("free reservation in later entry"); + assert!(matches!( + core.add_output_block("later-entry-reservation", None), + Err(SelectionError::NotFound(_)) + )); +} + +#[tokio::test] +async fn advisory_select_reports_worker_load_and_busy_evaluation() { + let mut config = test_config(false); + config.conditional_disagg_prefill_busy_threshold = Some(0.5); + config.conditional_disagg_decode_busy_threshold = Some(0.0); + let core = SelectionCore::try_new_local( + config, + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ) + .expect("valid test config"); + let mut request = worker(1); + request.total_kv_blocks = Some(1000); + core.upsert_worker(request).await.expect("worker upsert"); + + // Admitted (queued) select: decode evaluation comes from the catalog's + // total_kv_blocks; no load snapshot is taken. + let response = core.select(select_request()).await.expect("select"); + assert!(response.potential_decode_blocks > 0); + assert_eq!( + response.decode_busy, + Some(true), + "threshold 0.0 is always exceeded" + ); + assert!(response.worker_load.is_none()); + + // Advisory select: same decode evaluation plus the projected load. + let mut request = select_request(); + request.advisory = true; + let response = core.select(request).await.expect("advisory select"); + assert!(response.potential_decode_blocks > 0); + assert_eq!(response.decode_busy, Some(true)); + let load = response.worker_load.expect("advisory load"); + assert_eq!(load.total_kv_blocks, Some(1000)); + assert_eq!(load.prefill_token_capacity, 1024); + assert_eq!(load.active_prefill_tokens, 0); + assert_eq!(load.prefill_busy, Some(false)); + + // Advisory selection does not book. + assert!(core.reservation_index.read().is_empty()); + assert_eq!( + core.loads(Some("model"), Some("default"))[0].loads[0].potential_prefill_tokens, + 0 + ); +} + +#[tokio::test] +async fn busy_evaluation_is_absent_without_thresholds_or_capacity() { + let core = SelectionCore::try_new_local( + test_config(false), + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ) + .expect("valid test config"); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + let mut request = select_request(); + request.advisory = true; + let response = core.select(request).await.expect("advisory select"); + assert_eq!(response.decode_busy, None); + let load = response.worker_load.expect("advisory load"); + assert_eq!(load.total_kv_blocks, None); + assert_eq!(load.prefill_busy, None); +} + +#[tokio::test] +async fn reservation_index_tracks_bookings_until_freed() { + let core = SelectionCore::try_new_local( + test_config(false), + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ) + .expect("valid test config"); + for (worker_id, routing_group) in [(1, "group-a"), (2, "group-b")] { + let mut request = worker(worker_id); + request.routing_group = routing_group.to_string(); + core.upsert_worker(request).await.expect("worker upsert"); + } + let key_b = RoutingPartitionId::new("model", "group-b"); + + // select_and_reserve records the booking's partition. + let mut request = reserve_request("booked"); + request.routing_group = "group-b".to_string(); + core.select_and_reserve(request).await.expect("reserve"); + assert_eq!( + core.reservation_index + .read() + .get("booked") + .map(|r| &r.partition), + Some(&key_b) + ); + assert_eq!( + core.indexed_booking("booked") + .expect("indexed booking") + .0 + .key, + key_b + ); + + // The explicit reservation path records too. + let mut request = select_request(); + request.routing_group = "group-b".to_string(); + request.selection_id = Some("cached".to_string()); + core.select(request).await.expect("select"); + assert!(core.reservation_index.read().get("cached").is_none()); + core.create_reservation(ReservationRequest { + model_name: "model".to_string(), + routing_group: "group-b".to_string(), + selection_id: "cached".to_string(), + worker_id: None, + dp_rank: None, + prompt: PromptRequest::default(), + router_config_override: None, + expected_output_tokens: None, + effective_prefill_tokens: None, + track_prefill_tokens: None, + }) + .await + .expect("cached reservation"); + assert_eq!( + core.reservation_index + .read() + .get("cached") + .map(|r| &r.partition), + Some(&key_b) + ); + + // Lifecycle calls still resolve, and free drops the index entry. + core.prefill_complete("booked") + .await + .expect("prefill complete"); + core.free_reservation("booked").await.expect("free"); + assert!(core.reservation_index.read().get("booked").is_none()); + core.free_reservation("cached").await.expect("free"); + assert!(core.reservation_index.read().is_empty()); + + // Unknown ids fall back to the full scan and stay unindexed. + assert!(matches!( + core.prefill_complete("never-booked").await, + Err(SelectionError::NotFound(_)) + )); + assert!(core.reservation_index.read().is_empty()); +} + +#[tokio::test] +async fn reservation_index_sweep_drops_bookings_released_out_of_band() { + let core = SelectionCore::try_new_local( + test_config(false), + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ) + .expect("valid test config"); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + core.select_and_reserve(reserve_request("live")) + .await + .expect("reserve live"); + core.select_and_reserve(reserve_request("stale")) + .await + .expect("reserve stale"); + assert_eq!(core.reservation_index.read().len(), 2); + assert_eq!( + sweep_reservation_index(&core.entries, &core.reservation_index), + 0 + ); + + // Release directly through the scheduler, as force-expiry would. + let entry = core + .entry(&RoutingPartitionId::new("model", "default")) + .expect("entry"); + entry.scheduler.free("stale").await.expect("scheduler free"); + + assert_eq!( + sweep_reservation_index(&core.entries, &core.reservation_index), + 1 + ); + let index = core.reservation_index.read(); + assert_eq!(index.len(), 1); + assert!(index.contains_key("live")); +} + +#[tokio::test(flavor = "current_thread")] +async fn lifecycle_lookup_does_not_nest_reservation_index_inside_entries() { + // Three parties: a sweep holding `entries` and wanting `reservation_index`, + // a lifecycle call, and a partition creation queued on `entries.write()`. + let core = Arc::new( + SelectionCore::try_new_local( + test_config(false), + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ) + .expect("valid test config"), + ); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + core.select_and_reserve(reserve_request("live")) + .await + .expect("reserve live"); + let deadline = Instant::now() + Duration::from_secs(5); + + let sweep_entries = core.entries.read(); + let writer = { + let core = Arc::clone(&core); + std::thread::spawn(move || { + core.entries + .write() + .entry(RoutingPartitionId::new("other", "default")) + .or_insert_with(|| Arc::new(OnceCell::new())); + }) + }; + while core.entries.try_read().is_some() { + assert!(Instant::now() < deadline, "partition writer never queued"); + std::thread::yield_now(); + } + let lifecycle_started = Arc::new(AtomicBool::new(false)); + let lifecycle = { + let core = Arc::clone(&core); + let started = Arc::clone(&lifecycle_started); + std::thread::spawn(move || { + started.store(true, Ordering::Release); + core.indexed_booking("live").is_some() + }) + }; + while !lifecycle_started.load(Ordering::Acquire) { + assert!(Instant::now() < deadline, "lifecycle thread never started"); + std::thread::yield_now(); + } + sleep(Duration::from_millis(50)); + + // With the index guard held across `entries.read()` this never acquires. + drop( + core.reservation_index + .try_write_for(Duration::from_secs(2)) + .expect("reservation index must not be held by a blocked lifecycle call"), + ); + drop(sweep_entries); + writer.join().expect("partition writer"); + assert!(lifecycle.join().expect("lifecycle lookup")); +} + +#[tokio::test(flavor = "current_thread")] +async fn queued_selection_returns_refreshed_overlap_snapshot() { + let mut config = test_config(false); + config.router_queue_threshold = Some(0.0); + let core = Arc::new( + SelectionCore::try_new_local( + config, + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ) + .expect("valid test config"), + ); + + for worker_id in [1, 2] { + let mut request = worker(worker_id); + request.max_num_batched_tokens = Some(8); + core.upsert_worker(request).await.expect("worker upsert"); + } + let key = RoutingPartitionId::new("model", "default"); + let entry = core.entry(&key).expect("entry"); + entry + .indexer + .apply_event_routed(store_event(1, 0, 1, &[], &[11], StorageTier::Device)) + .await + .unwrap(); + entry.indexer.dump_events().await.expect("flush indexer"); + + for worker_id in [1, 2] { + core.create_reservation(ReservationRequest { + model_name: "model".to_string(), + routing_group: "default".to_string(), + selection_id: format!("occupy-{worker_id}"), + worker_id: Some(worker_id), + dp_rank: Some(0), + prompt: PromptRequest { + token_ids: None, + mm_routing_info: None, + block_mm_infos: None, + block_hashes: None, + sequence_hashes: Some(vec![1, 2]), + isl_tokens: Some(8), + lora_name: None, + cache_namespace: None, + is_eagle: None, + }, + router_config_override: None, + expected_output_tokens: None, + effective_prefill_tokens: Some(8), + track_prefill_tokens: None, + }) + .await + .expect("occupy worker"); + } + + let queued_core = Arc::clone(&core); + let queued = tokio::spawn(async move { + queued_core + .select_and_reserve(SelectAndReserveRequest { + model_name: "model".to_string(), + routing_group: "default".to_string(), + selection_id: Some("refresh-selection".to_string()), + prompt: PromptRequest { + token_ids: None, + mm_routing_info: None, + block_mm_infos: None, + block_hashes: Some(vec![11, 12]), + sequence_hashes: Some(vec![101, 102]), + isl_tokens: Some(8), + lora_name: None, + cache_namespace: None, + is_eagle: None, + }, + router_config_override: None, + expected_output_tokens: None, + priority_jump: None, + strict_priority: None, + session_id: None, + session_context: None, + affinity_target: None, + pinned_worker: None, + allowed_worker_ids: None, + routing_constraints: RoutingConstraints::default(), + }) + .await + }); + wait_for_pending_selection(&core).await; + + entry + .indexer + .apply_event_routed(store_event(2, 0, 1, &[], &[11, 12], StorageTier::Device)) + .await + .unwrap(); + entry.indexer.dump_events().await.expect("flush indexer"); + // Freeze time only after async setup so background timers cannot expire the fixture. + tokio::time::pause(); + tokio::time::advance(Duration::from_secs(11)).await; + core.free_reservation("occupy-2") + .await + .expect("release worker 2"); + + let response = queued.await.expect("selection task").expect("selection"); + assert_eq!(response.worker_id, 2); + assert_eq!(response.effective_prefill_tokens, 0); + assert_eq!(response.overlap.gpu, 8); + assert_eq!(response.overlap.cpu, 8); + assert_eq!(response.overlap.disk, 8); + assert_eq!(response.overlap.dp, HashMap::from([("0".to_string(), 8)])); +} + +fn core_with_session_affinity_mode(mode: SessionAffinityMode) -> SelectionCore { + let config = test_config(false); + let tracking_hash = Arc::new( + TrackingHashContext::from_config(&config).expect("valid tracking hash configuration"), + ); + let indexer_policy = IndexerPolicy::from_router_config(&config).expect("indexer policy"); + SelectionCore::new_inner( + config, + 1, + CancellationToken::new(), + None, + None, + SelectionHost::default(), + WorkerType::Aggregated, + true, + SelectionCacheConfig::default(), + tracking_hash, + indexer_policy, + Some(SessionAffinityConfig::new(Duration::from_secs(10)).with_mode(mode)), + ) +} + +fn core_with_session_affinity() -> SelectionCore { + core_with_session_affinity_mode(SessionAffinityMode::Hard) +} + +fn bound_worker(core: &SelectionCore, session_id: &str) -> Option { + core.entry(&RoutingPartitionId::new("model", "default")) + .and_then(|entry| entry.affinity.get().cloned()) + .and_then(|table| table.query_target(session_id, None).expect("query")) + .map(|target| target.worker_id) +} + +fn session_reservation(selection_id: &str, session_id: &str) -> SelectAndReserveRequest { + let mut request = reserve_request(selection_id); + request.session_id = Some(session_id.to_string()); + request +} + +#[tokio::test] +async fn departed_session_worker_reinitializes_the_session() { + let core = core_with_session_affinity(); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + core.upsert_worker(worker(2)).await.expect("worker upsert"); + let first = core + .select_and_reserve(session_reservation("r1", "s")) + .await + .expect("first booking"); + core.free_reservation("r1").await.expect("free"); + core.delete_worker(first.worker_id).await.expect("delete"); + + let second = core + .select_and_reserve(session_reservation("r2", "s")) + .await + .expect("session must move off a departed worker"); + assert_ne!(second.worker_id, first.worker_id); + assert_eq!(bound_worker(&core, "s"), Some(second.worker_id)); +} + +#[tokio::test] +async fn session_worker_departing_after_the_hold_reinitializes_the_session() { + let core = core_with_session_affinity(); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + core.upsert_worker(worker(2)).await.expect("worker upsert"); + let first = core + .select_and_reserve(session_reservation("r1", "s")) + .await + .expect("first booking"); + core.free_reservation("r1").await.expect("free"); + let key = RoutingPartitionId::new("model", "default"); + + // One poll takes the hold (the bound worker still passes the check) + // and hands the request to the scheduler actor, which has not run yet. + let mut second = Box::pin(core.select_and_reserve(session_reservation("r2", "s"))); + let mut context = std::task::Context::from_waker(std::task::Waker::noop()); + assert!(second.as_mut().poll(&mut context).is_pending()); + core.catalog + .set_lifecycle(first.worker_id, WorkerLifecycle::Draining, Vec::new()); + core.publish_scheduler_config(&key).expect("publish"); + + let second = second + .await + .expect("a departure after the hold is not a client fault"); + assert_ne!(second.worker_id, first.worker_id); + assert_eq!(bound_worker(&core, "s"), Some(second.worker_id)); +} + +#[tokio::test] +async fn concurrent_holds_on_a_departed_worker_both_land_on_the_replacement() { + let table = SessionAffinity::new(Duration::from_secs(60)).expect("affinity table"); + let departed = WorkerAffinityTarget::new(1, Some(0)); + let replacement = WorkerAffinityTarget::new(2, Some(0)); + let Hold::Initialize(init) = table.acquire("s", None).await.expect("acquire") else { + panic!("fresh session must initialize"); + }; + drop( + table + .commit(Hold::Initialize(init), departed) + .expect("bind"), + ); + let ( + Hold::Bound { + lease: mut first, .. + }, + Hold::Bound { + lease: mut second, .. + }, + ) = ( + table.acquire("s", None).await.expect("first hold"), + table.acquire("s", None).await.expect("second hold"), + ) + else { + panic!("both requests hold the departed binding"); + }; + + // Both requests notice the departure, as `hold_session` does: the first + // invalidation drops the binding, the second is a no-op release. + first.invalidate(); + second.invalidate(); + let leases = [ + table.commit( + table.acquire("s", None).await.expect("re-acquire"), + replacement, + ), + table.commit( + table.acquire("s", None).await.expect("re-acquire"), + replacement, + ), + ]; + for lease in leases { + drop(lease.expect("both requests bind to the replacement")); + } + assert_eq!( + table.query_target("s", None).expect("query"), + Some(replacement) + ); +} + +fn replay_reservation(selection_id: &str) -> ReservationRequest { + ReservationRequest { + model_name: "model".to_string(), + routing_group: "default".to_string(), + selection_id: selection_id.to_string(), + worker_id: None, + dp_rank: None, + prompt: PromptRequest::default(), + router_config_override: None, + expected_output_tokens: None, + effective_prefill_tokens: None, + track_prefill_tokens: None, + } +} + +#[tokio::test] +async fn two_phase_reservation_binds_the_session() { + let core = core_with_session_affinity(); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + core.upsert_worker(worker(2)).await.expect("worker upsert"); + let mut request = select_request(); + request.selection_id = Some("pending".to_string()); + request.session_id = Some("s".to_string()); + let selected = core.select(request).await.expect("select"); + assert_eq!(bound_worker(&core, "s"), None, "select alone binds nothing"); + + core.create_reservation(replay_reservation("pending")) + .await + .expect("replay booking"); + assert_eq!(bound_worker(&core, "s"), Some(selected.worker_id)); + core.free_reservation("pending").await.expect("free"); + + // The binding steers the session's next selection. + let mut request = select_request(); + request.session_id = Some("s".to_string()); + assert_eq!( + core.select(request).await.expect("select").worker_id, + selected.worker_id + ); +} + +#[tokio::test] +async fn two_phase_replay_rejects_a_worker_the_session_left() { + let core = core_with_session_affinity(); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + core.upsert_worker(worker(2)).await.expect("worker upsert"); + let mut request = select_request(); + request.selection_id = Some("pending".to_string()); + request.session_id = Some("s".to_string()); + let cached = core.select(request).await.expect("select"); + // Another request binds the session elsewhere before the replay. + let other = if cached.worker_id == 1 { 2 } else { 1 }; + let mut request = session_reservation("r1", "s"); + request.allowed_worker_ids = Some(HashSet::from([other])); + core.select_and_reserve(request) + .await + .expect("bind elsewhere"); + core.free_reservation("r1").await.expect("free"); + + let err = core + .create_reservation(replay_reservation("pending")) + .await + .expect_err("hard affinity rejects the stale cached worker"); + assert!(matches!(err, SelectionError::BadRequest(_)), "{err:?}"); + let entry = core + .entry(&RoutingPartitionId::new("model", "default")) + .expect("entry"); + wait_until("rejected booking release", || { + !entry.scheduler.has_request("pending") + }) + .await; + assert!(core.reservation_index.read().is_empty()); +} + +#[tokio::test] +async fn hard_mode_rejects_dispatch_away_from_a_live_binding() { + let core = core_with_session_affinity(); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + core.upsert_worker(worker(2)).await.expect("worker upsert"); + let first = core + .select_and_reserve(session_reservation("r1", "s")) + .await + .expect("first booking"); + core.free_reservation("r1").await.expect("free"); + let other = if first.worker_id == 1 { 2 } else { 1 }; + + // Steering cannot reach the bound worker, so selection lands elsewhere. + let mut request = session_reservation("r2", "s"); + request.allowed_worker_ids = Some(HashSet::from([other])); + let err = core + .select_and_reserve(request) + .await + .expect_err("hard affinity rejects a dispatch away from the binding"); + assert!(matches!(err, SelectionError::BadRequest(_)), "{err:?}"); + let entry = core + .entry(&RoutingPartitionId::new("model", "default")) + .expect("entry"); + wait_until("rejected booking release", || { + !entry.scheduler.has_request("r2") + }) + .await; + assert!(core.reservation_index.read().is_empty()); + assert_eq!(bound_worker(&core, "s"), None, "stale binding is dropped"); + + // The session's next request re-initializes. + let mut request = session_reservation("r3", "s"); + request.allowed_worker_ids = Some(HashSet::from([other])); + let third = core.select_and_reserve(request).await.expect("rebind"); + assert_eq!(third.worker_id, other); + assert_eq!(bound_worker(&core, "s"), Some(other)); +} + +#[tokio::test] +async fn soft_mode_follows_the_dispatch() { + let core = core_with_session_affinity_mode(SessionAffinityMode::Soft); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + core.upsert_worker(worker(2)).await.expect("worker upsert"); + let first = core + .select_and_reserve(session_reservation("r1", "s")) + .await + .expect("first booking"); + core.free_reservation("r1").await.expect("free"); + let other = if first.worker_id == 1 { 2 } else { 1 }; + + let mut request = session_reservation("r2", "s"); + request.allowed_worker_ids = Some(HashSet::from([other])); + let second = core.select_and_reserve(request).await.expect("soft rebind"); + assert_eq!(second.worker_id, other); + assert_eq!(bound_worker(&core, "s"), Some(other)); +} + +#[tokio::test] +async fn session_stays_on_its_first_worker_across_bookings() { + let core = core_with_session_affinity(); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + core.upsert_worker(worker(2)).await.expect("worker upsert"); + + let first = core + .select_and_reserve(session_reservation("r1", "chat-a")) + .await + .expect("first booking"); + for index in 0..4 { + let response = core + .select_and_reserve(session_reservation(&format!("r-{index}"), "chat-a")) + .await + .expect("booking"); + assert_eq!(response.worker_id, first.worker_id, "session must stay put"); + core.free_reservation(&format!("r-{index}")) + .await + .expect("free"); + } + // A read-only select sees the binding too. + let mut advisory = select_request(); + advisory.session_id = Some("chat-a".to_string()); + let response = core.select(advisory).await.expect("select"); + assert_eq!(response.worker_id, first.worker_id); + core.free_reservation("r1").await.expect("free"); +} + +#[tokio::test] +async fn replicated_binding_steers_a_new_session_and_frees_with_the_booking() { + let core = core_with_session_affinity(); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + core.upsert_worker(worker(2)).await.expect("worker upsert"); + + core.dispatch_affinity_event(AffinityBindingEvent { + partition: RoutingPartitionId::new("model", "default"), + session_id: "chat-b".to_string(), + worker_id: 2, + dp_rank: Some(0), + sequence: 1, + writer_id: 99, + }); + let response = core + .select_and_reserve(session_reservation("r2", "chat-b")) + .await + .expect("booking"); + assert_eq!(response.worker_id, 2); + assert!( + core.reservation_index + .read() + .get("r2") + .unwrap() + ._affinity_lease + .is_some() + ); + core.free_reservation("r2").await.expect("free"); + assert!(core.reservation_index.read().is_empty()); +} + +#[tokio::test] +async fn expired_booking_releases_affinity_lease() { + let core = core_with_session_affinity(); + core.upsert_worker(worker(1)).await.unwrap(); + core.select_and_reserve(session_reservation("abandoned", "session")) + .await + .unwrap(); + let key = RoutingPartitionId::new("model", "default"); + core.entry(&key) + .unwrap() + .scheduler + .free("abandoned") + .await + .unwrap(); + assert_eq!( + sweep_reservation_index(&core.entries, &core.reservation_index), + 1 + ); + tokio::time::pause(); + tokio::time::advance(Duration::from_secs(11)).await; + assert_eq!( + core.entry(&RoutingPartitionId::new("model", "default")) + .unwrap() + .affinity + .get() + .unwrap() + .query_target("session", None) + .unwrap(), + None + ); +} + +#[tokio::test] +async fn partition_sessions_keep_independent_bindings() { + let core = core_with_session_affinity(); + core.upsert_worker(worker(1)).await.unwrap(); + let first = core + .select_and_reserve(session_reservation("first", "shared-session")) + .await + .unwrap(); + core.upsert_worker(worker(2)).await.unwrap(); + let mut other = worker(3); + other.routing_group = "other".to_string(); + core.upsert_worker(other).await.unwrap(); + let mut request = session_reservation("other", "shared-session"); + request.routing_group = "other".to_string(); + assert_eq!(core.select_and_reserve(request).await.unwrap().worker_id, 3); + assert_eq!( + core.entry(&RoutingPartitionId::new("model", "default")) + .unwrap() + .affinity + .get() + .unwrap() + .query_target("shared-session", None) + .unwrap() + .unwrap() + .worker_id, + first.worker_id + ); + let again = core + .select_and_reserve(session_reservation("again", "shared-session")) + .await + .unwrap(); + assert_eq!(again.worker_id, first.worker_id); +} + +#[tokio::test] +async fn rejoined_worker_feeds_partition_index() { + use crate::protocols::{BlockHashOptions, compute_block_hash_for_seq}; + for update_only in [false, true] { + let core = SelectionCore::try_new_local( + test_config(true), + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ) + .unwrap(); + let request = worker_with_kv_events(1); + core.upsert_worker(request.clone()).await.unwrap(); + let key = RoutingPartitionId::new("model", "default"); + let partition = core.partition(&key).unwrap(); + let tokens: Vec = (1..=8).collect(); + let hashes: Vec = compute_block_hash_for_seq(&tokens, 4, BlockHashOptions::default()) + .into_iter() + .map(|hash| hash.0) + .collect(); + let indexer = core + .indexer_registry + .get_indexer(&key) + .unwrap() + .indexer + .clone(); + indexer + .apply_event_routed(store_event(1, 0, 1, &[], &hashes, StorageTier::Device)) + .await + .unwrap(); + indexer.dump_events().await.unwrap(); + if !update_only { + core.delete_worker(1).await.unwrap(); + } + let mut request = request; + request.total_kv_blocks = Some(2048); + core.upsert_worker(request).await.unwrap(); + let retained = partition.indexer().dump_events().await.unwrap(); + if update_only { + assert!( + !retained.is_empty(), + "metadata update cleared cached blocks" + ); + } else { + assert!(retained.is_empty(), "removed worker retained cached blocks"); + let current = core + .indexer_registry + .get_indexer(&key) + .unwrap() + .indexer + .clone(); + current + .apply_event_routed(store_event(1, 0, 1, &[], &hashes, StorageTier::Device)) + .await + .unwrap(); + current.dump_events().await.unwrap(); + } + let mut request = select_request(); + request.prompt = PromptRequest { + token_ids: Some(tokens), + ..PromptRequest::default() + }; + assert_eq!( + core.select(request).await.unwrap().overlap.longest_matched, + 8, + "update_only={update_only}" + ); + assert_eq!( + partition.indexer().dump_events().await.unwrap().len(), + indexer.dump_events().await.unwrap().len() + ); + } +} +#[tokio::test] +async fn metadata_update_preserves_live_booking() { + struct SuspendedDetach(tokio::sync::Notify); + #[async_trait::async_trait] + impl KvEventIngress for SuspendedDetach { + fn open( + &self, + registry: &WorkerRegistry, + key: &RoutingPartitionId, + block_size: u32, + ) -> Indexer { + registry.get_or_create_indexer(key.clone(), block_size) + } + async fn detach(&self, _registry: &WorkerRegistry, _record: &WorkerCatalogRecord) { + self.0.notify_one(); + std::future::pending().await + } + } + let ingress = Arc::new(SuspendedDetach(tokio::sync::Notify::new())); + let core = core_with_host(SelectionHost { + cache: HostCache { + index: KvIndexSource::Owned(ingress.clone()), + shared: None, + }, + ..SelectionHost::default() + }); + core.upsert_worker(worker(1)).await.unwrap(); + core.select_and_reserve(reserve_request("live")) + .await + .unwrap(); + let entry = core + .entry(&RoutingPartitionId::new("model", "default")) + .unwrap(); + let mut updated = worker(1); + updated.max_num_batched_tokens = Some(2048); + tokio::select! { + result = core.upsert_worker(updated) => { result.unwrap(); }, + _ = ingress.0.notified() => { + // A capacity update must not temporarily withdraw this worker. + assert!(entry.scheduler.has_request("live")); + panic!("capacity update detached the worker"); + } + } + assert!(entry.scheduler.has_request("live")); + core.free_reservation("live").await.unwrap(); +} +#[tokio::test(start_paused = true)] +async fn host_lease_manager_owns_expiry() { + use crate::scheduling::queue::SchedulerBookingDescriptor; + struct HostLeases; + impl ReplicaRequestLeaseObserver for HostLeases { + fn admitted(&self, _: SchedulerBookingDescriptor) {} + fn progressed(&self, _: &SchedulerBookingDescriptor) {} + fn completed(&self, _: &SchedulerBookingDescriptor) {} + } + for host_owned in [false, true] { + let core = core_with_host(SelectionHost { + replication: HostReplication { + request_leases: host_owned + .then(|| Arc::new(HostLeases) as Arc), + ..HostReplication::default() + }, + ..SelectionHost::default() + }); + core.upsert_worker(worker(1)).await.unwrap(); + core.select_and_reserve(reserve_request("live")) + .await + .unwrap(); + tokio::task::yield_now().await; + tokio::time::advance(active_request_expiry_duration() * 3).await; + tokio::task::yield_now().await; + let entry = core + .entry(&RoutingPartitionId::new("model", "default")) + .unwrap(); + assert_eq!(entry.scheduler.has_request("live"), host_owned); + if host_owned { + core.free_reservation("live").await.unwrap(); + assert!(!entry.scheduler.has_request("live")); + } + } +} + +#[tokio::test] +async fn catalog_updates_commit_in_order_without_exposing_candidates() { + struct PausedUpdate { + entered: tokio::sync::Notify, + release: tokio::sync::Notify, + } + #[async_trait::async_trait] + impl KvEventIngress for PausedUpdate { + fn open( + &self, + registry: &WorkerRegistry, + key: &RoutingPartitionId, + block_size: u32, + ) -> Indexer { + registry.get_or_create_indexer(key.clone(), block_size) + } + async fn reconcile( + &self, + _: &WorkerRegistry, + previous: Option<&WorkerCatalogRecord>, + _: &WorkerCatalogRecord, + ) -> Result<(), SelectionError> { + if previous.is_some() { + self.entered.notify_one(); + self.release.notified().await; + } + Ok(()) + } + } + let ingress = Arc::new(PausedUpdate { + entered: tokio::sync::Notify::new(), + release: tokio::sync::Notify::new(), + }); + let mut core = core_with_host(SelectionHost { + cache: HostCache { + index: KvIndexSource::Owned(ingress.clone()), + shared: None, + }, + ..SelectionHost::default() + }); + core.listens_for_kv_events = true; + core.upsert_worker(worker(1)).await.unwrap(); + core.select_and_reserve(reserve_request("live")) + .await + .unwrap(); + let mut updated = worker(1); + updated.max_num_batched_tokens = Some(2048); + let update = core.upsert_worker(updated); + tokio::pin!(update); + tokio::select! { + _ = &mut update => panic!("update should pause before committing"), + _ = ingress.entered.notified() => {} + } + let committed = core.catalog.get(1).unwrap(); + assert_eq!(committed.lifecycle, WorkerLifecycle::Schedulable); + assert_eq!(committed.max_num_batched_tokens, Some(1024)); + let entry = core.entry(&committed.key()).unwrap(); + assert!(entry.scheduler.has_request("live")); + let deletion = core.delete_worker(1); + tokio::pin!(deletion); + assert!( + std::future::Future::poll( + deletion.as_mut(), + &mut std::task::Context::from_waker(std::task::Waker::noop()) + ) + .is_pending() + ); + ingress.release.notify_one(); + assert_eq!(update.await.unwrap().max_num_batched_tokens, Some(2048)); + assert!(entry.scheduler.has_request("live")); + assert_eq!( + deletion.await.unwrap().lifecycle, + WorkerLifecycle::Unschedulable + ); + tokio::time::timeout(Duration::from_secs(1), async { + while entry.scheduler.has_request("live") { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); +} + +#[tokio::test] +async fn affinity_configuration_rejects_invalid_or_conflicting_config() { + use super::super::service::SelectionServiceBuilder; + for ttl in [ + Duration::ZERO, + Duration::from_secs(super::super::affinity::MAX_SESSION_AFFINITY_TTL_SECS + 1), + ] { + let result = SelectionServiceBuilder::new( + test_config(false), + WorkerType::Aggregated, + Default::default(), + ) + .session_affinity(ttl) + .build() + .await; + assert!(result.is_err()); + } + let core = core_with_session_affinity(); + core.upsert_worker(worker(1)).await.unwrap(); + let partition = core + .partition(&RoutingPartitionId::new("model", "default")) + .unwrap(); + assert!(matches!( + partition.session_affinity(SessionAffinityConfig::new(Duration::from_secs(20))), + Err(SelectionError::Conflict(_)) + )); + assert!(matches!( + partition.session_affinity( + SessionAffinityConfig::new(Duration::from_secs(10)) + .with_mode(SessionAffinityMode::Soft) + ), + Err(SelectionError::Conflict(_)) + )); +} + +fn hint_config(worker_type: &str, endpoints: &[(u32, &str)]) -> SelectionWorkerConfig { + SelectionWorkerConfig { + endpoint: "http://worker:8000".to_string(), + data_parallel_start_rank: 0, + data_parallel_size: endpoints.len().max(1) as u32, + max_num_batched_tokens: None, + total_kv_blocks: None, + stable_routing_id: None, + is_eagle: None, + taints: HashSet::new(), + topology_domains: HashMap::new(), + kv_transfer_domain: None, + kv_transfer_enforcement: None, + kv_transfer_preferred_weight: None, + router_hint_worker_type: Some(worker_type.to_string()), + router_hint_source_control_endpoints: endpoints + .iter() + .map(|(rank, endpoint)| (*rank, endpoint.to_string())) + .collect(), + kv_event_source_mode: None, + } +} + +fn hint_candidates( + hashes: &[u64], + owners: Vec<(KvTransferCandidateSource, usize)>, +) -> KvTransferCandidates { + KvTransferCandidates { + block_hashes: hashes + .iter() + .map(|h| crate::protocols::ExternalSequenceBlockHash(*h)) + .collect(), + owner_prefix_blocks: owners, + routing_snapshot: None, + } +} + +#[test] +fn hint_source_endpoint_follows_the_source_dp_rank() { + let configs = HashMap::from([( + 7, + hint_config( + "prefill", + &[(0, "tcp://127.0.0.1:23280"), (1, "tcp://127.0.0.1:23281")], + ), + )]); + let candidates = hint_candidates(&[101, 102], vec![(WorkerWithDpRank::new(7, 1).into(), 2)]); + let hint = + transfer_hint_for_selection(&configs, WorkerWithDpRank::new(7, 0), 0, Some(&candidates)); + assert_eq!( + hint.map(|payload| payload.source_control_endpoint), + Some("tcp://127.0.0.1:23281".to_string()) + ); +} + +#[test] +fn hint_source_must_share_the_target_worker_type() { + let configs = HashMap::from([ + (7, hint_config("prefill", &[(0, "tcp://127.0.0.1:23280")])), + (8, hint_config("prefill", &[(0, "tcp://127.0.0.1:23281")])), + (9, hint_config("decode", &[(0, "tcp://127.0.0.1:23282")])), + (10, hint_config("decode", &[(0, "tcp://127.0.0.1:23283")])), + ]); + let candidates = hint_candidates( + &[101, 102, 103], + vec![ + (WorkerWithDpRank::new(8, 0).into(), 2), + (WorkerWithDpRank::new(9, 0).into(), 3), + ], + ); + // The longer decode prefix is skipped for a prefill target. + let prefill = + transfer_hint_for_selection(&configs, WorkerWithDpRank::new(7, 0), 0, Some(&candidates)) + .expect("prefill hint"); + assert_eq!(prefill.source_control_endpoint, "tcp://127.0.0.1:23281"); + assert_eq!(prefill.block_hashes.len(), 2); + let decode = + transfer_hint_for_selection(&configs, WorkerWithDpRank::new(10, 0), 0, Some(&candidates)) + .expect("decode hint"); + assert_eq!(decode.source_control_endpoint, "tcp://127.0.0.1:23282"); + assert_eq!(decode.block_hashes.len(), 3); +} + +#[test] +fn hint_resolves_a_persistent_cache_owner_over_a_state_agent_worker() { + use crate::identity::{ + CacheOwnerId, CacheSemanticsId, DcId, IdentitySource, IndexerDomainId, PoolId, + RoutingScopeId, StableDpSlotId, + }; + use crate::protocols::{ + ResidencyOwner, ResidencyProjection, ResidencyRoutingSnapshot, RouterHintSourceMetadata, + }; + let mut stale = hint_config("prefill", &[(0, "tcp://stale-worker:23280")]); + stale.kv_event_source_mode = Some("state_agent_v2".to_string()); + let configs = HashMap::from([(7, hint_config("prefill", &[])), (8, stale)]); + let owner = CacheOwnerId::new( + PoolId::new( + IndexerDomainId::new( + CacheSemanticsId::new([1; 16], IdentitySource::Explicit), + RoutingScopeId::new([2; 16], IdentitySource::Explicit), + ), + DcId::new(3), + ), + StableDpSlotId::new([4; 16], IdentitySource::Explicit), + ); + let mut candidates = hint_candidates( + &[101, 102], + vec![ + ( + KvTransferCandidateSource::Worker(WorkerWithDpRank::new(8, 0)), + 2, + ), + ( + KvTransferCandidateSource::CacheOwner( + ResidencyOwner::cache_owner(owner).compact_key(), + ), + 2, + ), + ], + ); + candidates.routing_snapshot = Some(Arc::new(ResidencyRoutingSnapshot::new( + ResidencyProjection::default(), + [( + owner, + RouterHintSourceMetadata { + source_control_endpoint: "tcp://persistent-owner:23280".to_string(), + worker_type: "prefill".to_string(), + }, + None, + )], + ))); + let hint = + transfer_hint_for_selection(&configs, WorkerWithDpRank::new(7, 0), 0, Some(&candidates)) + .expect("hint"); + assert_eq!(hint.source_control_endpoint, "tcp://persistent-owner:23280"); +} diff --git a/lib/kv-router/src/services/selection/core/workers.rs b/lib/kv-router/src/services/selection/core/workers.rs new file mode 100644 index 000000000000..deeb0317b6b1 --- /dev/null +++ b/lib/kv-router/src/services/selection/core/workers.rs @@ -0,0 +1,365 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Worker membership: catalog upserts, patches and deletes, partition +//! creation, indexer registration and scheduler config publication. + +use super::reservations::{ReservationIndexObserver, spawn_reservation_index_sweep}; +use super::*; + +impl SelectionCore { + pub async fn upsert_worker( + &self, + req: WorkerRequest, + ) -> Result { + self.ensure_running()?; + let mut record = WorkerCatalogRecord::new(req); + // Partition policy factories may construct independently. Only committing + // membership and reconciling ingress needs the catalog mutation lock. + self.prepare_worker(&mut record)?; + let _update = self.catalog_updates.lock().await; + self.ensure_running()?; + let previous = self.catalog.get(record.worker_id); + self.reconcile_worker(record, previous).await + } + + pub async fn patch_worker( + &self, + worker_id: WorkerId, + patch: WorkerPatchRequest, + ) -> Result { + let _update = self.catalog_updates.lock().await; + self.ensure_running()?; + let previous = self + .catalog + .get(worker_id) + .ok_or_else(|| SelectionError::NotFound(format!("worker {worker_id} not found")))?; + let mut record = previous.clone(); + record.apply_patch(patch); + self.prepare_worker(&mut record)?; + self.reconcile_worker(record, Some(previous)).await + } + + pub async fn delete_worker( + &self, + worker_id: WorkerId, + ) -> Result { + let _update = self.catalog_updates.lock().await; + let Some(previous) = self.catalog.get(worker_id) else { + return Err(SelectionError::NotFound(format!( + "worker {worker_id} not found" + ))); + }; + let key = previous.key(); + self.catalog + .set_lifecycle(worker_id, WorkerLifecycle::Draining, Vec::new()); + self.publish_scheduler_config(&key)?; + self.cleanup_indexer_registration(&previous).await; + let record = self + .catalog + .set_lifecycle(worker_id, WorkerLifecycle::Unschedulable, Vec::new()) + .ok_or_else(|| SelectionError::NotFound(format!("worker {worker_id} not found")))?; + self.publish_scheduler_config(&key)?; + Ok(record) + } + + pub fn list_workers( + &self, + model_name: Option<&str>, + routing_group: Option<&str>, + ) -> Vec { + self.catalog.list(model_name, routing_group) + } + + pub fn ready(&self) -> ReadyResponse { + let schedulable_workers = self.catalog.schedulable_count(); + let workers = self.catalog.list(None, None); + ReadyResponse { + ready: !self.cancel_token.is_cancelled() && schedulable_workers > 0, + schedulable_workers, + workers, + } + } + + fn prepare_worker(&self, record: &mut WorkerCatalogRecord) -> Result<(), SelectionError> { + let queueing_enabled = self + .kv_router_config + .queueing_enabled(Some(&record.model_name)) + .map_err(|error| SelectionError::BadRequest(error.to_string()))?; + record.not_schedulable_reasons = record.missing_schedulable_metadata(queueing_enabled); + if let Some(ingress) = self.ingress() { + record + .not_schedulable_reasons + .extend(ingress.missing_metadata(record)); + } + if record.not_schedulable_reasons.is_empty() + && let Err(error) = self.ensure_entry(record) + { + record + .not_schedulable_reasons + .push(format!("reconciliation failed: {error}")); + } + + Ok(()) + } + + async fn reconcile_worker( + &self, + mut record: WorkerCatalogRecord, + previous: Option, + ) -> Result { + let previous = previous.filter(|old| old.lifecycle == WorkerLifecycle::Schedulable); + let previous = if let Some(old) = previous.as_ref() + && (old.key() != record.key() || !record.not_schedulable_reasons.is_empty()) + { + self.catalog + .set_lifecycle(old.worker_id, WorkerLifecycle::Draining, Vec::new()); + self.publish_scheduler_config(&old.key())?; + self.cleanup_indexer_registration(old).await; + None + } else { + previous + }; + + if record.not_schedulable_reasons.is_empty() + && let Some(ingress) = self.ingress() + && let Err(error) = ingress + .reconcile(&self.indexer_registry, previous.as_ref(), &record) + .await + { + self.cleanup_indexer_registration(&record).await; + record + .not_schedulable_reasons + .push(format!("reconciliation failed: {error}")); + } + record.lifecycle = if record.not_schedulable_reasons.is_empty() { + WorkerLifecycle::Schedulable + } else { + WorkerLifecycle::Incomplete + }; + // Readers see only committed metadata. A valid capacity/topology update preserves + // live bookings on ranks present in both the old and new snapshots. + self.catalog.replace(record.clone()); + self.publish_scheduler_config(&record.key())?; + Ok(record) + } + + fn ensure_entry( + &self, + record: &WorkerCatalogRecord, + ) -> Result, SelectionError> { + let block_size = record + .block_size + .ok_or_else(|| SelectionError::BadRequest("block_size is required".to_string()))?; + self.ensure_entry_for(record.key(), block_size, record.is_eagle.unwrap_or(false)) + } + + /// Create the partition scheduler and indexer for `key` before any worker + /// registers, so an embedding host can hold the scheduler handle from + /// construction. Idempotent; a later worker with a different block size or + /// eagle setting is rejected at reconciliation. + pub fn ensure_partition( + &self, + key: RoutingPartitionId, + block_size: u32, + is_eagle: bool, + ) -> Result { + self.ensure_running()?; + if block_size == 0 { + return Err(SelectionError::BadRequest( + "block_size must be greater than 0".to_string(), + )); + } + self.ensure_entry_for(key, block_size, is_eagle) + .map(SelectionPartition) + } + + fn ensure_entry_for( + &self, + key: RoutingPartitionId, + block_size: u32, + is_eagle: bool, + ) -> Result, SelectionError> { + self.reservation_sweep_started.get_or_init(|| { + spawn_reservation_index_sweep( + Arc::clone(&self.entries), + Arc::clone(&self.reservation_index), + self.cancel_token.child_token(), + ); + }); + + let entry_cell = { self.entries.read().get(&key).cloned() }; + let entry_cell = entry_cell.unwrap_or_else(|| { + self.entries + .write() + .entry(key.clone()) + .or_insert_with(|| Arc::new(OnceCell::new())) + .clone() + }); + let entry = entry_cell + .get_or_try_init(|| -> Result, SelectionError> { + let (workers_tx, workers_rx) = watch::channel(HashMap::new()); + let host_replica = self + .host + .replication + .channels + .as_ref() + .and_then(|factory| factory(&key)); + let scoped_replica_sync = setup_scoped_replica_sync( + self.replica_config.as_ref(), + &key, + block_size, + host_replica, + ); + let worker_label = self.worker_type.as_str(); + let slots = Arc::new(ActiveSequencesMultiWorker::new_with_options( + scoped_replica_sync + .publisher + .with_load_sink(self.host.telemetry.scheduler_load.clone()), + block_size as usize, + HashMap::new(), + scoped_replica_sync.enabled, + scoped_replica_sync.process_id, + worker_label, + SequenceTrackerOptions { + replica_worker_policy: ReplicaWorkerPolicy::RequireRegistered, + expiry_duration: self + .host + .replication + .request_leases + .is_none() + .then(active_request_expiry_duration), + }, + )); + slots.set_replica_request_lease_observer(Arc::new(ReservationIndexObserver { + index: Arc::clone(&self.reservation_index), + partition: key.clone(), + host: self.host.replication.request_leases.clone(), + })); + let replica_tx = scoped_replica_sync.channel.map(|(replica_tx, subscriber)| { + slots.start_replica_sync(subscriber, self.cancel_token.child_token()); + replica_tx + }); + if self.host.replication.request_leases.is_none() { + slots.start_periodic_force_expiry_across_all_workers( + self.cancel_token.child_token(), + ); + } + + let indexer = match &self.host.cache.index { + KvIndexSource::Owned(ingress) => { + ingress.open(&self.indexer_registry, &key, block_size) + } + KvIndexSource::Remote(_) => self + .indexer_registry + .get_or_create_indexer(key.clone(), block_size), + }; + let overlap_refresh = indexer.supports_overlap_refresh().then(|| { + Arc::new(TieredOverlapRefresher::new( + indexer.clone(), + self.kv_router_config.clone(), + block_size, + )) + }); + let selector = self.worker_selection_policy_factory.as_ref().map_or_else( + || WorkerSelectionPolicy::default(self.kv_router_config.clone(), worker_label), + |factory| factory(&self.kv_router_config, self.worker_type, key.as_ref()), + ); + let profile = self + .kv_router_config + .policy_profile(Some(&key.model_name)) + .map_err(|error| SelectionError::BadRequest(error.to_string()))?; + let scheduler = LocalScheduler::new( + slots, + workers_rx, + profile, + block_size, + selector, + self.host.load.prefill_estimator.clone(), + overlap_refresh, + // Standalone selection has no router Client snapshot, so + // these stay `None` unless an embedding host injects them. + self.host.load.overloaded_workers.clone(), + self.host.load.available_workers.clone(), + self.kv_router_config.router_queue_recheck_interval(), + self.kv_router_config.router_track_prefill_tokens, + self.cancel_token.child_token(), + worker_label, + true, + ); + Ok(Arc::new(SelectionEntry { + key: key.clone(), + block_size, + is_eagle, + indexer, + workers_tx, + scheduler, + replica_tx, + affinity: OnceCell::new(), + replica_config: self.replica_config.clone(), + })) + })? + .clone(); + if entry.block_size != block_size { + return Err(SelectionError::Conflict(format!( + "block_size mismatch for {key}: existing={} requested={block_size}", + entry.block_size + ))); + } + if entry.is_eagle != is_eagle { + return Err(SelectionError::Conflict(format!( + "is_eagle mismatch for {key}: existing={} requested={is_eagle}", + entry.is_eagle + ))); + } + if let Some(config) = self.session_affinity { + entry.session_affinity(config)?; + } + Ok(entry) + } + + /// The ingress feeding core-owned indexes, when this core listens for KV events. + fn ingress(&self) -> Option<&dyn KvEventIngress> { + match &self.host.cache.index { + KvIndexSource::Owned(ingress) if self.listens_for_kv_events => Some(ingress.as_ref()), + _ => None, + } + } + + async fn cleanup_indexer_registration(&self, record: &WorkerCatalogRecord) { + if let KvIndexSource::Owned(ingress) = &self.host.cache.index { + ingress.detach(&self.indexer_registry, record).await; + return; + } + + let key = record.key(); + let indexer = self + .indexer_registry + .get_indexer(&key) + .map(|entry| entry.indexer.clone()); + if let Some(indexer) = indexer { + indexer.remove_worker(record.worker_id).await; + } + } + + pub(super) fn publish_scheduler_config( + &self, + key: &RoutingPartitionId, + ) -> Result<(), SelectionError> { + let Some(entry) = self.entry(key) else { + return Ok(()); + }; + let workers = self.catalog.scheduler_configs_for_key(key); + // Lifecycle transitions between non-schedulable states publish the same + // map; skipping them saves the scheduler a wake and a full map clone. + entry.workers_tx.send_if_modified(|current| { + if *current == workers { + false + } else { + *current = workers; + true + } + }); + Ok(()) + } +} From 9c4b3347f75496338683b34ed09164956b7cbcda Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Thu, 10 Sep 2026 17:25:25 -0700 Subject: [PATCH 059/102] chore(router): delete the selection paths the core cutover left unused SelectionServiceBuilder::session_affinity_mode had no caller (the standalone service keeps the default mode; the frontend sets the mode on its table directly); LocalScheduler and SchedulerQueue::supports_overlap_refresh lost their only caller with the frontend pipeline; SessionAffinity::new_with_limits was test-only and is inlined as with_config; the free_if_worker wrappers on KvRouter and EmbeddedSelection had no caller. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Ishan Dhanani --- lib/kv-router/src/scheduling/local.rs | 4 ---- lib/kv-router/src/scheduling/queue.rs | 4 ---- lib/kv-router/src/services/selection/affinity.rs | 12 ------------ lib/kv-router/src/services/selection/core/tests.rs | 8 ++++++-- lib/kv-router/src/services/selection/service.rs | 14 ++------------ lib/llm/src/kv_router.rs | 12 ------------ lib/llm/src/kv_router/embedded.rs | 11 ----------- lib/llm/src/session_affinity/coordinator.rs | 6 +++--- 8 files changed, 11 insertions(+), 60 deletions(-) diff --git a/lib/kv-router/src/scheduling/local.rs b/lib/kv-router/src/scheduling/local.rs index d86c8b149dfc..6eb5f8c12e18 100644 --- a/lib/kv-router/src/scheduling/local.rs +++ b/lib/kv-router/src/scheduling/local.rs @@ -667,10 +667,6 @@ where self.queue.class_queue_stats(class_index) } - pub fn supports_overlap_refresh(&self) -> bool { - self.queue.supports_overlap_refresh() - } - pub fn worker_type(&self) -> &'static str { self.worker_type } diff --git a/lib/kv-router/src/scheduling/queue.rs b/lib/kv-router/src/scheduling/queue.rs index c5f317db20d0..e7996c039c81 100644 --- a/lib/kv-router/src/scheduling/queue.rs +++ b/lib/kv-router/src/scheduling/queue.rs @@ -873,10 +873,6 @@ impl< }) } - pub fn supports_overlap_refresh(&self) -> bool { - self.supports_overlap_refresh - } - fn prepare_block_hashes_for_refresh( &self, block_hashes: Option>, diff --git a/lib/kv-router/src/services/selection/affinity.rs b/lib/kv-router/src/services/selection/affinity.rs index c7a1c3a77ab8..e291b829351f 100644 --- a/lib/kv-router/src/services/selection/affinity.rs +++ b/lib/kv-router/src/services/selection/affinity.rs @@ -224,18 +224,6 @@ impl SessionAffinity { Ok(()) } - pub fn new_with_limits( - ttl: Duration, - max_entries: usize, - max_session_id_bytes: usize, - ) -> Result { - Self::with_config(SessionAffinityConfig { - max_entries, - max_session_id_bytes, - ..SessionAffinityConfig::new(ttl) - }) - } - pub fn with_config(config: SessionAffinityConfig) -> Result { Self::validate_ttl(config.ttl)?; let inner = Arc::new(Inner { diff --git a/lib/kv-router/src/services/selection/core/tests.rs b/lib/kv-router/src/services/selection/core/tests.rs index 03aacee52f8e..93e559f81e6a 100644 --- a/lib/kv-router/src/services/selection/core/tests.rs +++ b/lib/kv-router/src/services/selection/core/tests.rs @@ -816,8 +816,12 @@ async fn full_affinity_table_routes_without_pinning() { let entry = core .entry(&RoutingPartitionId::new("model", "default")) .expect("entry"); - let table = - SessionAffinity::new_with_limits(Duration::from_secs(60), 1, 256).expect("affinity table"); + let table = SessionAffinity::with_config(SessionAffinityConfig { + max_entries: 1, + max_session_id_bytes: 256, + ..SessionAffinityConfig::new(Duration::from_secs(60)) + }) + .expect("affinity table"); assert!(entry.affinity.set(table).is_ok()); for (selection_id, session_id) in [("first", "s1"), ("second", "s2")] { diff --git a/lib/kv-router/src/services/selection/service.rs b/lib/kv-router/src/services/selection/service.rs index 0cb4ab1f4b03..f9790b52fafc 100644 --- a/lib/kv-router/src/services/selection/service.rs +++ b/lib/kv-router/src/services/selection/service.rs @@ -16,7 +16,7 @@ use crate::services::common::replica_sync::{ use crate::services::indexer::backend::IndexerPolicy; use crate::tracking_hash::TrackingHashContext; -use super::affinity::{SessionAffinityConfig, SessionAffinityMode}; +use super::affinity::SessionAffinityConfig; use super::core::{ KvIndexSource, SelectionCore, SelectionHost, SelectionPartition, SelectionServiceConfig, }; @@ -42,7 +42,6 @@ pub struct SelectionServiceBuilder { host: SelectionHost, worker_selection_policy_factory: Option, session_affinity_ttl: Option, - session_affinity_mode: SessionAffinityMode, } /// Warn when a host does not construct workers for explicitly configured policy roles. @@ -83,7 +82,6 @@ impl SelectionServiceBuilder { host: SelectionHost::default(), worker_selection_policy_factory: None, session_affinity_ttl: None, - session_affinity_mode: SessionAffinityMode::default(), } } @@ -127,13 +125,6 @@ impl SelectionServiceBuilder { self } - /// How a bound session treats a booking that lands on another worker: - /// `Hard` (default) rejects it with 400, `Soft` moves the session. - pub fn session_affinity_mode(mut self, mode: SessionAffinityMode) -> Self { - self.session_affinity_mode = mode; - self - } - pub fn replica_sync(mut self, port: u16, peers: Vec) -> Self { self.replica_sync_port = Some(port); self.replica_sync_peers = peers; @@ -198,8 +189,7 @@ impl SelectionServiceBuilder { self.selection_cache, tracking_hash, indexer_policy, - self.session_affinity_ttl - .map(|ttl| SessionAffinityConfig::new(ttl).with_mode(self.session_affinity_mode)), + self.session_affinity_ttl.map(SessionAffinityConfig::new), )); if recover_from_peers { diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index 09c211d9c0cc..272cdbfe4f7d 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -1598,18 +1598,6 @@ impl KvRouter { self.scheduler.free(request_id).await } - /// Release a booking only if it still belongs to `worker`. - /// - /// An ownership mismatch is a harmless no-op, which makes this safe for - /// delayed cleanup that captured the worker when it acquired the booking. - pub async fn free_if_worker( - &self, - request_id: &str, - worker: WorkerWithDpRank, - ) -> Result<(), SequenceError> { - self.scheduler.free_if_worker(request_id, worker).await - } - #[doc(hidden)] pub(crate) fn affinity_coordinator( &self, diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs index c665f0d36a88..bb8f15340ee3 100644 --- a/lib/llm/src/kv_router/embedded.rs +++ b/lib/llm/src/kv_router/embedded.rs @@ -422,17 +422,6 @@ impl EmbeddedSelection { self.partition.scheduler().free(request_id).await } - pub(crate) async fn free_if_worker( - &self, - request_id: &str, - worker: WorkerWithDpRank, - ) -> Result<(), SequenceError> { - self.partition - .scheduler() - .free_if_worker(request_id, worker) - .await - } - pub(crate) fn booking_cleanup(&self) -> SchedulerBookingCleanup { self.partition.scheduler().booking_cleanup() } diff --git a/lib/llm/src/session_affinity/coordinator.rs b/lib/llm/src/session_affinity/coordinator.rs index c06dd089f6ed..c18574a4678b 100644 --- a/lib/llm/src/session_affinity/coordinator.rs +++ b/lib/llm/src/session_affinity/coordinator.rs @@ -227,11 +227,11 @@ impl AffinityCoordinator { #[cfg(test)] pub(super) fn with_test_limits(max_entries: usize, max_session_id_bytes: usize) -> Self { Self::wrap( - SessionAffinity::new_with_limits( - Duration::from_secs(10), + SessionAffinity::with_config(SessionAffinityConfig { max_entries, max_session_id_bytes, - ) + ..SessionAffinityConfig::new(Duration::from_secs(10)) + }) .unwrap(), ) } From 89785a3e7b2c52e161efabf3629612751804f030 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Thu, 10 Sep 2026 17:42:51 -0700 Subject: [PATCH 060/102] refactor(router): start replica ingress after the lease manager is installed The embedded partition reports request-lifecycle events to the router's RequestLeaseManager, which can only be built once the partition exists. The inbound replica ingress used to start with the partition, so a buffering observer replayed early peer events into the manager after installation. host_replica_channels now starts only the outbound publisher and hands back the inbound leg as a ReplicaIngress. EmbeddedSelection::start returns it wrapped as PendingReplicaIngress, whose start installs the manager and then runs the ingress, so the ordering is fixed by the types rather than by call order in KvRouter::new. The buffering DeferredReplicaRequestLeaseObserver and its state machine are replaced by LateBoundLeaseObserver, which logs and drops an event that arrives before the manager exists. Behavior change: a peer replica event published while the router is still reconciling its initial worker membership is no longer mirrored locally; replica state was already best-effort across the subscription start. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Ishan Dhanani --- lib/llm/src/kv_router.rs | 9 +- lib/llm/src/kv_router/embedded.rs | 48 ++++-- lib/llm/src/kv_router/sequence.rs | 257 ++++++++++++++---------------- 3 files changed, 158 insertions(+), 156 deletions(-) diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index 272cdbfe4f7d..fe2a528fdad7 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -793,8 +793,7 @@ impl KvRouter { let available_worker_provider: WorkerAvailabilityProvider = Arc::new(move || client_for_availability.available_instance_ids()); - let replica_leases = Arc::new(sequence::DeferredReplicaRequestLeaseObserver::default()); - let scheduler = embedded::EmbeddedSelection::start( + let (scheduler, replica_ingress) = embedded::EmbeddedSelection::start( embedded::EmbeddedSelectionArgs { kv_router_config: kv_router_config.clone(), worker_role, @@ -815,7 +814,6 @@ impl KvRouter { endpoint: endpoint.clone(), router_id: endpoint.drt().discovery().instance_id(), policy_factory, - request_leases: replica_leases.clone(), }, workers_with_configs.clone(), cancellation_token.child_token(), @@ -825,7 +823,10 @@ impl KvRouter { scheduler.booking_cleanup(), cancellation_token.child_token(), ); - replica_leases.install(Arc::new(request_leases.clone())); + // Inbound lifecycle events start only now that their consumer exists. + replica_ingress + .start(Arc::new(request_leases.clone())) + .await; tracing::info!("KV Routing initialized"); let cancellation_token = cancellation_guard.disarm(); Ok(Self { diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs index bb8f15340ee3..78a1b0b9005b 100644 --- a/lib/llm/src/kv_router/embedded.rs +++ b/lib/llm/src/kv_router/embedded.rs @@ -68,7 +68,6 @@ pub(crate) struct EmbeddedSelectionArgs { /// Builds the partition's worker-selection policy (see /// `SelectionPolicySource::resolve`). pub policy_factory: WorkerSelectionPolicyFactory, - pub request_leases: Arc, } fn record_queue_rejection( @@ -176,6 +175,24 @@ pub(crate) struct EmbeddedSelection { queue_metric_indices: HashMap, } +/// The partition's inbound replica ingress, not yet running. Starting it +/// installs the router's lease manager first, so no lifecycle event reaches +/// the partition before its consumer exists. +pub(crate) struct PendingReplicaIngress { + ingress: crate::kv_router::sequence::ReplicaIngress, + leases: Arc, +} + +impl PendingReplicaIngress { + pub(crate) async fn start( + self, + leases: Arc, + ) { + self.leases.install(leases); + self.ingress.start().await; + } +} + static INSTALLED_POLICY_REGISTRY: std::sync::OnceLock = std::sync::OnceLock::new(); @@ -224,7 +241,7 @@ impl EmbeddedSelection { args: EmbeddedSelectionArgs, workers_with_configs: RuntimeConfigWatch, cancellation_token: CancellationToken, - ) -> Result { + ) -> Result<(Self, PendingReplicaIngress)> { let worker_type = args.worker_role.unwrap_or(WorkerType::Aggregated); let key = RoutingPartitionId::new( args.model_name @@ -236,7 +253,7 @@ impl EmbeddedSelection { // Replica sync rides the runtime event plane. Worker-origin completion // marks are consumed even when router-to-router replica sync is // disabled; only publishing is gated. - let channels = crate::kv_router::sequence::host_replica_channels( + let (channels, replica_ingress) = crate::kv_router::sequence::host_replica_channels( &args.endpoint, args.router_id, args.kv_router_config.router_replica_sync, @@ -244,6 +261,7 @@ impl EmbeddedSelection { ) .await .context("start replica sync for the embedded selection partition")?; + let leases = Arc::new(crate::kv_router::sequence::LateBoundLeaseObserver::default()); let slot = std::sync::Mutex::new(Some(channels)); let replica_sync: Option = Some(Arc::new(move |_partition| { @@ -278,7 +296,7 @@ impl EmbeddedSelection { }, replication: HostReplication { channels: replica_sync, - request_leases: Some(args.request_leases), + request_leases: Some(leases.clone()), }, }) .build() @@ -349,14 +367,20 @@ impl EmbeddedSelection { worker_type = %worker_type, "KvRouter scheduling on embedded selection partition" ); - Ok(Self { - service, - partition, - affinity: std::sync::OnceLock::new(), - worker_type: args.metric_worker_type, - queue_metrics, - queue_metric_indices, - }) + Ok(( + Self { + service, + partition, + affinity: std::sync::OnceLock::new(), + worker_type: args.metric_worker_type, + queue_metrics, + queue_metric_indices, + }, + PendingReplicaIngress { + ingress: replica_ingress, + leases, + }, + )) } fn observe_queue(&self, rejection: Option<&QueueRejection>) { diff --git a/lib/llm/src/kv_router/sequence.rs b/lib/llm/src/kv_router/sequence.rs index 692e91187e5c..f8a260cc1a05 100644 --- a/lib/llm/src/kv_router/sequence.rs +++ b/lib/llm/src/kv_router/sequence.rs @@ -24,7 +24,6 @@ use dynamo_runtime::traits::DistributedRuntimeProvider; use dynamo_runtime::transports::event_plane::{ EventPublisher, EventSubscriber, EventTransportKind, TypedEventSubscriber, }; -use parking_lot::Mutex; use std::collections::VecDeque; use std::future::Future; use std::sync::Arc; @@ -43,104 +42,53 @@ use dynamo_runtime::transports::event_plane::MsgpackCodec; // if the queue is full, the newest event is dropped without blocking the local mutation. const REPLICA_EVENT_CHANNEL_CAPACITY: usize = 100_000; -enum DeferredReplicaLeaseEvent { - Admitted(SchedulerBookingDescriptor), - Progressed(SchedulerBookingDescriptor), - Completed(SchedulerBookingDescriptor), -} - -enum DeferredReplicaLeaseState { - Buffering(Vec), - Installing { - target: Arc, - pending: Vec, - }, - Installed(Arc), -} - -impl Default for DeferredReplicaLeaseState { - fn default() -> Self { - Self::Buffering(Vec::new()) - } -} - -/// Buffers the construction-time event window until `KvRouter` installs its -/// request lease manager. After installation, calls pass through directly. +/// The partition's lease observer until the router's `RequestLeaseManager` +/// exists; an event before `install` is a wiring bug and is dropped loudly. #[derive(Default)] -pub(crate) struct DeferredReplicaRequestLeaseObserver { - state: Mutex, +pub(crate) struct LateBoundLeaseObserver { + target: std::sync::OnceLock>, } -impl DeferredReplicaRequestLeaseObserver { - pub(crate) fn install(&self, target: Arc) -> bool { - let mut pending = { - let mut state = self.state.lock(); - let DeferredReplicaLeaseState::Buffering(pending) = &mut *state else { - return false; - }; - let pending = std::mem::take(pending); - *state = DeferredReplicaLeaseState::Installing { - target: Arc::clone(&target), - pending: Vec::new(), - }; - pending - }; - - loop { - for event in pending { - Self::notify(&target, event); - } - let mut state = self.state.lock(); - let DeferredReplicaLeaseState::Installing { - target: installing_target, - pending: queued, - } = &mut *state - else { - unreachable!("observer installation state changed unexpectedly"); - }; - if queued.is_empty() { - *state = DeferredReplicaLeaseState::Installed(Arc::clone(installing_target)); - return true; - } - pending = std::mem::take(queued); - } - } - - fn observe(&self, event: DeferredReplicaLeaseEvent) { - let target = { - let mut state = self.state.lock(); - match &mut *state { - DeferredReplicaLeaseState::Buffering(pending) - | DeferredReplicaLeaseState::Installing { pending, .. } => { - pending.push(event); - return; - } - DeferredReplicaLeaseState::Installed(target) => Arc::clone(target), - } - }; - Self::notify(&target, event); +impl LateBoundLeaseObserver { + pub(crate) fn install(&self, target: Arc) { + let installed = self.target.set(target).is_ok(); + debug_assert!(installed, "request lease manager installed twice"); } - fn notify(target: &Arc, event: DeferredReplicaLeaseEvent) { - match event { - DeferredReplicaLeaseEvent::Admitted(booking) => target.admitted(booking), - DeferredReplicaLeaseEvent::Progressed(booking) => target.progressed(&booking), - DeferredReplicaLeaseEvent::Completed(booking) => target.completed(&booking), + fn target( + &self, + event: &str, + booking: &SchedulerBookingDescriptor, + ) -> Option<&Arc> { + let target = self.target.get(); + if target.is_none() { + tracing::error!( + request_id = %booking.request_id, + event, + "request lifecycle event before the lease manager was installed; the manager will not track this mirrored booking" + ); } + target } } -impl ReplicaRequestLeaseObserver for DeferredReplicaRequestLeaseObserver { +impl ReplicaRequestLeaseObserver for LateBoundLeaseObserver { fn admitted(&self, booking: SchedulerBookingDescriptor) { - self.observe(DeferredReplicaLeaseEvent::Admitted(booking)); + if let Some(target) = self.target("admitted", &booking) { + target.admitted(booking); + } } fn progressed(&self, booking: &SchedulerBookingDescriptor) { - self.observe(DeferredReplicaLeaseEvent::Progressed(booking.clone())); + if let Some(target) = self.target("progressed", booking) { + target.progressed(booking); + } } fn completed(&self, booking: &SchedulerBookingDescriptor) { - self.observe(DeferredReplicaLeaseEvent::Completed(booking.clone())); + if let Some(target) = self.target("completed", booking) { + target.completed(booking); + } } } @@ -467,13 +415,18 @@ impl SequenceSubscriber for RuntimeSequenceSubscriber { /// peer and worker events on `ACTIVE_SEQUENCES_SUBJECT` are always forwarded /// to the partition (worker-origin completion marks are needed even without /// router-to-router sync); outbound events are published only when -/// `publishes_outbound` is set. +/// `publishes_outbound` is set. The outbound publisher runs from here; the +/// inbound leg runs when the returned [`ReplicaIngress`] is started, so the +/// caller can install every consumer of lifecycle events first. pub(crate) async fn host_replica_channels( endpoint: &Endpoint, router_id: u64, publishes_outbound: bool, cancellation_token: CancellationToken, -) -> Result { +) -> Result<( + dynamo_kv_router::services::selection::HostReplicaChannels, + ReplicaIngress, +)> { let transport_kind = endpoint.drt().default_event_transport_kind(); let outbound = if publishes_outbound { let (outbound, outbound_rx) = mpsc::channel(REPLICA_EVENT_CHANNEL_CAPACITY); @@ -505,39 +458,62 @@ pub(crate) async fn host_replica_channels( }; let (inbound_tx, inbound_rx) = mpsc::channel(REPLICA_EVENT_CHANNEL_CAPACITY); - let direct = direct_zmq::DirectZmqSequenceConfig::from_env(); - let ingress_result = if direct.should_use_direct(transport_kind) { - direct_zmq::start( - endpoint.clone(), - inbound_tx.clone(), - direct.rcvhwm, - cancellation_token, - ) - .await - .map(|_supervisor| ()) - } else { - RuntimeSequenceSubscriber::for_endpoint(endpoint) - .await - .map(|subscriber| { - tokio::spawn(forward_replica_events( - subscriber, - inbound_tx.clone(), - cancellation_token, - )); - }) + let ingress = ReplicaIngress { + endpoint: endpoint.clone(), + inbound_tx: inbound_tx.clone(), + cancellation_token, }; - if let Err(error) = ingress_result { - tracing::warn!( - %error, - "active-sequence event ingress unavailable; continuing with response-side cleanup" - ); + Ok(( + dynamo_kv_router::services::selection::HostReplicaChannels { + outbound, + inbound_tx, + inbound_rx, + process_id: router_id, + }, + ingress, + )) +} + +/// The inbound leg of [`host_replica_channels`]: peer replica events and +/// worker-origin completion marks feed `inbound_tx` from the direct-ZMQ +/// fan-in or the runtime subscriber, whichever the transport selects. +pub(crate) struct ReplicaIngress { + endpoint: Endpoint, + inbound_tx: mpsc::Sender, + cancellation_token: CancellationToken, +} + +impl ReplicaIngress { + pub(crate) async fn start(self) { + let transport_kind = self.endpoint.drt().default_event_transport_kind(); + let direct = direct_zmq::DirectZmqSequenceConfig::from_env(); + let ingress_result = if direct.should_use_direct(transport_kind) { + direct_zmq::start( + self.endpoint, + self.inbound_tx, + direct.rcvhwm, + self.cancellation_token, + ) + .await + .map(|_supervisor| ()) + } else { + RuntimeSequenceSubscriber::for_endpoint(&self.endpoint) + .await + .map(|subscriber| { + tokio::spawn(forward_replica_events( + subscriber, + self.inbound_tx, + self.cancellation_token, + )); + }) + }; + if let Err(error) = ingress_result { + tracing::warn!( + %error, + "active-sequence event ingress unavailable; continuing with response-side cleanup" + ); + } } - Ok(dynamo_kv_router::services::selection::HostReplicaChannels { - outbound, - inbound_tx, - inbound_rx, - process_id: router_id, - }) } async fn forward_replica_events( @@ -568,6 +544,7 @@ async fn forward_replica_events( mod tests { use super::*; use dynamo_kv_router::protocols::ActiveSequenceEventData; + use parking_lot::Mutex; use std::sync::Arc; use tokio::time::Instant; @@ -860,20 +837,20 @@ mod tests { ); } #[test] - fn deferred_replica_observer_flushes_completed_startup_attempts_in_order() { + fn late_bound_observer_drops_before_install_and_forwards_after() { use dynamo_kv_router::sequences::NoopSequencePublisher; use std::collections::HashMap; #[derive(Default)] - struct Observer(Mutex>); + struct Observer(Mutex>); impl ReplicaRequestLeaseObserver for Observer { fn admitted(&self, booking: SchedulerBookingDescriptor) { - self.0.lock().push(("add", booking)); + self.0.lock().push(("add", booking.request_id)); } fn progressed(&self, booking: &SchedulerBookingDescriptor) { - self.0.lock().push(("progress", booking.clone())); + self.0.lock().push(("progress", booking.request_id.clone())); } fn completed(&self, booking: &SchedulerBookingDescriptor) { - self.0.lock().push(("free", booking.clone())); + self.0.lock().push(("free", booking.request_id.clone())); } } let slots = ActiveSequencesMultiWorker::new_without_expiry( @@ -884,27 +861,26 @@ mod tests { 0, "test", ); - let deferred = Arc::new(DeferredReplicaRequestLeaseObserver::default()); - assert!(slots.set_replica_request_lease_observer(deferred.clone())); - slots.apply_replica_batch(vec![add_event("startup"), free_event("startup")]); + let late = Arc::new(LateBoundLeaseObserver::default()); + assert!(slots.set_replica_request_lease_observer(late.clone())); + // Dropped: no target installed. + slots.apply_replica_batch(vec![ + add_event("before-install"), + free_event("before-install"), + ]); let observer = Arc::new(Observer::default()); - assert!(deferred.install(observer.clone())); - { - let events = observer.0.lock(); - assert_eq!(events.len(), 2); - let booking = &events[0].1; - assert_eq!(booking.request_id, "startup"); - assert_eq!( - *events, - vec![("add", booking.clone()), ("free", booking.clone())] - ); - } - slots.apply_replica_batch(vec![add_event("after-install")]); + late.install(observer.clone()); + slots.apply_replica_batch(vec![ + add_event("after-install"), + free_event("after-install"), + ]); assert_eq!( - observer.0.lock().last().unwrap().1.request_id, - "after-install" + *observer.0.lock(), + vec![ + ("add", "after-install".to_string()), + ("free", "after-install".to_string()) + ] ); - assert!(!deferred.install(observer)); } #[tokio::test] @@ -924,9 +900,10 @@ mod tests { .component("workers")? .endpoint("generate"); let cancel = CancellationToken::new(); - let mut channels = + let (mut channels, ingress) = host_replica_channels(&endpoint, 99, false, cancel.child_token()).await?; assert!(channels.outbound.is_none()); + ingress.start().await; let publisher = ActiveSequenceEventPublisher::for_endpoint(&endpoint, 16).await?; let received = tokio::time::timeout(std::time::Duration::from_secs(5), async { From 25aea16655b9d014e1e10693f541c30f760c1e7f Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Thu, 10 Sep 2026 19:07:12 -0700 Subject: [PATCH 061/102] fix(router): hold the routed selection future once on the cleanup path With the frontend routing through the core, every disaggregated request segfaulted the frontend in a debug build: a tokio worker stack overflowed in the prologue of await_with_cleanup_policy, the wrapper the routing host passes its selection future through. The wrapper kept the operation by value three times (its argument, the move into cancel_on_stop, and a second pin in the other arm), so its state was three copies of an already large future, nested under the prefill router's request future. The wrapper now pins the operation once and lends it to both arms (cancel_on_stop takes Pin<&mut _>), and the embedded host boxes the core's selection future so the frontend's request future carries a pointer. A test keeps the KV selection future under 16 KiB (6.4 KB today in debug). Co-Authored-By: Claude Fable 5.1 Signed-off-by: Ishan Dhanani --- lib/llm/src/kv_router/embedded.rs | 3 +- .../kv_router/routing_host/cancellation.rs | 36 +++++++++++++------ lib/llm/src/kv_router/routing_host/tests.rs | 21 +++++++++++ 3 files changed, 48 insertions(+), 12 deletions(-) diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs index 78a1b0b9005b..0ce88dac16e4 100644 --- a/lib/llm/src/kv_router/embedded.rs +++ b/lib/llm/src/kv_router/embedded.rs @@ -417,7 +417,8 @@ impl EmbeddedSelection { /// Run one selection through the shared core. pub(crate) async fn run_selection(&self, operation: SelectionOperation<'_>) -> SelectionRun { - let run = self.service.core().run_selection(operation).await; + // Keep the selection state out of the frontend's nested request future. + let run = Box::pin(self.service.core().run_selection(operation)).await; self.observe_queue(match &run.result { Ok(SelectionOutcome::QueueRejected { rejection }) => Some(rejection), _ => None, diff --git a/lib/llm/src/kv_router/routing_host/cancellation.rs b/lib/llm/src/kv_router/routing_host/cancellation.rs index af115de5fd1a..302021290814 100644 --- a/lib/llm/src/kv_router/routing_host/cancellation.rs +++ b/lib/llm/src/kv_router/routing_host/cancellation.rs @@ -1,7 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use std::{future::Future, sync::Mutex, time::Duration}; +use std::{future::Future, pin::Pin, sync::Mutex, time::Duration}; // `tokio::time::Instant` rather than the std type so the budget is measured on // the same clock as the `tokio::time::timeout` that consumes it. @@ -108,10 +108,12 @@ pub(super) async fn await_with_cleanup_policy( budget: &CleanupBudget, operation: impl Future, ) -> Result { + // Pinned once here and lent to both arms, so the wrapper's state holds the + // operation a single time rather than once per arm. + tokio::pin!(operation); match DispatchCancellation::for_request(phase, staged_kv) { DispatchCancellation::CancelWhenStopped => cancel_on_stop(context, operation).await, DispatchCancellation::DispatchWhenStopped => { - tokio::pin!(operation); tokio::select! { biased; @@ -165,9 +167,8 @@ pub(super) fn cancelled_error(context_id: &str) -> Error { pub(super) async fn cancel_on_stop( context: &dyn AsyncEngineContext, - operation: impl Future, + mut operation: Pin<&mut impl Future>, ) -> Result { - tokio::pin!(operation); tokio::select! { biased; @@ -196,7 +197,11 @@ mod tests { pipeline::{AsyncEngineContext, context::Controller}, }; - use super::{CLEANUP_DISPATCH_TIMEOUT, CleanupBudget, cancel_on_stop}; + use super::{ + CLEANUP_DISPATCH_TIMEOUT, CleanupBudget, StagedKv, await_with_cleanup_policy, + cancel_on_stop, + }; + use crate::protocols::common::timing::RequestPhase; struct PendingUntilDropped(Arc); @@ -247,9 +252,18 @@ mod tests { context.stop(); let dropped = Arc::new(AtomicBool::new(false)); - let error = cancel_on_stop(&context, PendingUntilDropped(dropped.clone())) - .await - .unwrap_err(); + // Through the wrapper, which owns the operation: it is gone once the + // cancellation error is returned. + let error = await_with_cleanup_policy( + &context, + RequestPhase::Aggregated, + StagedKv::Absent, + "test", + &CleanupBudget::default(), + PendingUntilDropped(dropped.clone()), + ) + .await + .unwrap_err(); let error = error .downcast_ref::() @@ -263,9 +277,9 @@ mod tests { let context = Controller::new("completed-request".to_string()); context.stop(); - let result = cancel_on_stop(&context, std::future::ready(42)) - .await - .unwrap(); + let ready = std::future::ready(42); + tokio::pin!(ready); + let result = cancel_on_stop(&context, ready).await.unwrap(); assert_eq!(result, 42); } diff --git a/lib/llm/src/kv_router/routing_host/tests.rs b/lib/llm/src/kv_router/routing_host/tests.rs index 3ee32c1b4c4c..c17601746993 100644 --- a/lib/llm/src/kv_router/routing_host/tests.rs +++ b/lib/llm/src/kv_router/routing_host/tests.rs @@ -968,6 +968,27 @@ async fn always_ready_terminals_cannot_starve_the_drain_deadline() { runtime.shutdown(); } +/// The selection future is nested inside every request future (deeper still +/// on the disaggregated path). A debug build once overflowed a worker stack +/// when this grew and `await_with_cleanup_policy` held it by value; keep it +/// small enough that nesting stays cheap. +#[tokio::test] +#[serial_test::serial] +async fn kv_selection_future_stays_small() { + let (router, _runtime) = embedded_router_with_workers(None, &[1]).await; + let request = Context::with_id_and_metadata( + request(), + "selection-future-size".to_string(), + Default::default(), + ); + let budget = CleanupBudget::default(); + let future = router.select_with_affinity(&request, RequestPhase::Aggregated, false, &budget); + let size = std::mem::size_of_val(&future); + drop(future); + eprintln!("kv selection future size: {size} bytes"); + assert!(size < 16 * 1024, "kv selection future is {size} bytes"); +} + /// Transport EOF ends the drain and releases the booking. #[tokio::test] #[serial_test::serial] From e22fab5ee7974cabfafd99828f178ae435a2987b Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Fri, 11 Sep 2026 12:58:09 -0700 Subject: [PATCH 062/102] fix(router): apply the review round on the selection core Eleven review items from the current head, each a small local change: - docs: the two-phase `/reservations` contract binds the session; the earlier paragraph said the cache dropped session metadata. - indexer registry: clone the `Indexer` handle and drop the DashMap shard guard before awaiting the rank removal. - indexer server: serialize `TieredQueryResponse` directly instead of through a `serde_json::Value`, and give `/query_by_hash` and `/query_tiered_by_hash` the same 8 MiB body limit as `/query`. - reservation index: drop removed entries (and their session bindings, which take a shard lock and publish) after the index write lock is released, in both `forget_reservation_if` and the sweep. - selection run: copy `sequence_hashes` into the result only for `Book` admissions; the frontend's `Lease` path never reads them. - LoRA filter: intersect the filter's output with the candidate set so a filter cannot widen the caller's allowlist; module made private and the helper `pub(crate)`. - session affinity: store the replica writer id only after the sink installs, so a rejected second call cannot overwrite it. - bindings test: pytest timeout marker on the subprocess test. - delete the `register_workers` chain (`EmbeddedSelection`, `KvRouter`, `PrefillRouter`) and its two EPP call sites: the terminal method was a no-op since membership became catalog-driven. EPP comments that described the old upsert behavior now describe what the scheduler does with an unknown id. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Ishan Dhanani --- deploy/inference-gateway/ext-proc/src/epp.rs | 32 +++++++------------ .../router/standalone-selection.md | 2 +- lib/bindings/python/tests/test_kv_bindings.py | 1 + lib/kv-router/src/lib.rs | 2 +- lib/kv-router/src/scheduling/lora_filter.rs | 30 +++++++++++++++++ lib/kv-router/src/scheduling/mod.rs | 5 +-- .../src/services/indexer/registry.rs | 6 ++-- lib/kv-router/src/services/indexer/server.rs | 14 +++++--- .../src/services/selection/affinity.rs | 8 +++-- .../services/selection/core/reservations.rs | 32 ++++++++++++++----- .../src/services/selection/core/run.rs | 5 ++- lib/llm/src/kv_router.rs | 5 --- lib/llm/src/kv_router/embedded.rs | 6 +--- lib/llm/src/kv_router/prefill_router/query.rs | 8 ----- 14 files changed, 96 insertions(+), 60 deletions(-) diff --git a/deploy/inference-gateway/ext-proc/src/epp.rs b/deploy/inference-gateway/ext-proc/src/epp.rs index 7f6eb80d0988..a3e9367f29df 100644 --- a/deploy/inference-gateway/ext-proc/src/epp.rs +++ b/deploy/inference-gateway/ext-proc/src/epp.rs @@ -475,10 +475,6 @@ impl Router { allowed_worker_ids: Option>, routing_constraints: RoutingConstraints, ) -> Result { - if let Some(ref ids) = allowed_worker_ids { - self.prefill_router.register_workers(ids); - } - self.prefill_router .reserve_prefill_worker( reservation_id, @@ -519,10 +515,6 @@ impl Router { allowed_worker_ids: Option>, routing_constraints: RoutingConstraints, ) -> Result<(WorkerWithDpRank, u32)> { - if let Some(ref ids) = allowed_worker_ids { - self.decode_router.register_workers(ids); - } - let config_override = decode_router_config_override(is_disaggregated); let outcome = self @@ -926,8 +918,9 @@ fn indexed_endpoint_address(endpoint: &Endpoint) -> Option { /// The mode is exclusive, and so are the identities. A worker process picks one /// `KubeDiscoveryTarget` from its own mode, so under container discovery /// nothing registers under the bare pod identity — emitting it there would -/// invent a worker that `register_workers` upserts at zero load and zero KV -/// overlap, making it the most attractive candidate the scheduler sees. +/// put a worker id in `allowed_worker_ids` that no backend registered. The +/// scheduler filters unknown ids out, so the pod contributes nothing, and a +/// subset made only of such ids selects no worker at all. /// `"main"` hashes to the pod identity (`hash_container_name`), so a pod whose /// main container is Ready still contributes that id through the container /// path; one whose main container is *not* Ready correctly contributes nothing @@ -1384,11 +1377,9 @@ impl EndpointPicker for Router { // Only pod discovery registers a worker under its pod identity. // Under container discovery each engine container registers under // its own (`KubeDiscoveryTarget::Container`), so a hand-built pod - // hash names a worker present in no registry: `register_workers` - // would upsert it at zero load and zero KV overlap, making it the - // scheduler's most attractive candidate, and the reverse lookup - // below would then fail to match and silently forward to - // `endpoints[0]`. The index is the one place that knows which + // hash names a worker present in no registry, which the scheduler + // filters out, leaving the subset short one candidate or empty. + // The index is the one place that knows which // identity scheme is in effect (see `pod_worker_ids`), so ask it. let wm: Vec<(u64, &Endpoint)> = { let index = read_index(&self.worker_index); @@ -2183,9 +2174,9 @@ mod tests { /// Under pod discovery a worker registers under its pod identity alone, so /// a pod's ready sidecars must contribute no worker ids. Emitting them - /// would invent workers no backend registered under: they miss - /// `register_workers`' discovery lookup, default to `(0, 1)` with no load - /// and no KV overlap, and so look maximally attractive to the scheduler. + /// would name workers no backend registered under; the scheduler filters + /// them out, so the pod's real worker would be the subset's only survivor + /// only by accident of the id set. #[test] fn pod_worker_ids_ignores_containers_under_pod_discovery() { let pod = pod_mode_worker_pod(); @@ -2251,7 +2242,7 @@ mod tests { } /// End of the chain that made this matter: the index feeds - /// `subset_to_worker_ids` -> `allowed_worker_ids` -> `register_workers`, + /// `subset_to_worker_ids` -> `allowed_worker_ids` -> the scheduler, /// so one backend pod must contribute exactly one worker id under pod /// discovery rather than one per ready container. #[test] @@ -2486,8 +2477,7 @@ mod tests { /// An externally supplied endpoint must resolve to the identity the /// reflector actually holds. Under container discovery that is the engine /// container's id, never `hash_pod_name` -- deriving the latter names a - /// worker no registry contains, which `register_workers` then upserts at - /// zero load as the scheduler's most attractive candidate. + /// worker no registry contains. #[test] fn external_endpoint_resolves_to_the_indexed_container_identity() { let mut index = WorkerEndpointIndex::new(true); diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-selection.md b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-selection.md index 84f0d1d2aee6..def9ef969bcf 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-selection.md +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-selection.md @@ -329,7 +329,7 @@ Both `POST /select` and `POST /select_and_reserve` accept an optional `session_i Bindings replicate over the configured replica mesh. An explicit `affinity_target` or `pinned_worker` takes precedence over the session binding. Without a configured TTL, `session_id` is only policy input: custom policies can read it through `WorkerSelectionContext::session_id()`, while the built-in selector ignores it. See [Write Custom Routing Strategies](custom-worker-selection.mdx). -The pending-selection cache does not retain session metadata, so a later `POST /reservations` does not create a session binding. Use `/select_and_reserve` when the service should manage affinity. The frontend uses the same table implementation for request-header affinity; see [Configuration and Tuning](configuration-and-tuning.md). +The pending-selection cache keeps the session metadata, so a later `POST /reservations` for that selection binds the session to the booked worker, the same as `/select_and_reserve` does. The frontend uses the same table implementation for request-header affinity; see [Configuration and Tuning](configuration-and-tuning.md). ### `session_context` diff --git a/lib/bindings/python/tests/test_kv_bindings.py b/lib/bindings/python/tests/test_kv_bindings.py index 72e85018f3c5..bb9bca53222e 100644 --- a/lib/bindings/python/tests/test_kv_bindings.py +++ b/lib/bindings/python/tests/test_kv_bindings.py @@ -220,6 +220,7 @@ def test_selection_service_rejects_invalid_affinity_ttl(ttl): reason="SelectionService requires the select-service Cargo feature", ) @pytest.mark.parametrize("ttl", ["-1", "NaN", "31536001"]) +@pytest.mark.timeout(60) # subprocess wait is capped at 15 s per case def test_selection_service_cli_rejects_invalid_affinity_ttl(ttl): result = subprocess.run( [ diff --git a/lib/kv-router/src/lib.rs b/lib/kv-router/src/lib.rs index 9c69a7a4d522..7f16fe045417 100644 --- a/lib/kv-router/src/lib.rs +++ b/lib/kv-router/src/lib.rs @@ -69,8 +69,8 @@ pub use protocols::{ pub use queue::SchedulerQueue; pub use radix_tree::RadixTree; pub use scheduling::LocalScheduler; +pub use scheduling::LoraWorkerFilter; pub use scheduling::PrefillLoadEstimator; -pub use scheduling::lora_filter::LoraWorkerFilter; pub use scheduling::policy::{FcfsPolicy, RouterSchedulingPolicy, SchedulingPolicy, WsptPolicy}; pub use scheduling::{ KvSchedulerError, PotentialLoad, SchedulingRequest, SchedulingResponse, SessionContext, diff --git a/lib/kv-router/src/scheduling/lora_filter.rs b/lib/kv-router/src/scheduling/lora_filter.rs index 7f550442a403..69f6c483d872 100644 --- a/lib/kv-router/src/scheduling/lora_filter.rs +++ b/lib/kv-router/src/scheduling/lora_filter.rs @@ -44,9 +44,11 @@ pub fn narrow_allowed_worker_ids_by_lora( if base.is_empty() { return allowed_worker_ids; } + // A filter may only narrow: ids it returns from outside `base` are dropped. let mut narrowed: HashSet = filter .filter_worker_ids_for_lora(lora_name, &base) .into_iter() + .filter(|id| base.contains(id)) .collect(); if let Some(pinned) = pinned_worker && base.contains(&pinned.worker_id) @@ -126,6 +128,34 @@ mod tests { ); } + /// A filter that answers with ids outside the caller's universe. + struct Widening(Vec); + + impl LoraWorkerFilter for Widening { + fn filter_worker_ids_for_lora( + &self, + _lora_name: &str, + _available: &[WorkerId], + ) -> Vec { + self.0.clone() + } + } + + #[test] + fn filter_output_outside_the_universe_is_dropped() { + // 9 is in neither the allow-set nor the worker list: it must not appear. + let allowed = Some(HashSet::from([1, 2])); + assert_eq!( + narrow(&Widening(vec![2, 9]), Some("a"), allowed, None), + Some(HashSet::from([2])) + ); + assert_eq!( + narrow(&Widening(vec![9]), Some("a"), None, None), + None, + "only out-of-universe ids: nothing survives, so the allow-set is unchanged" + ); + } + #[test] fn empty_universe_is_a_noop() { assert_eq!( diff --git a/lib/kv-router/src/scheduling/mod.rs b/lib/kv-router/src/scheduling/mod.rs index da58b6653d21..c48797e7293d 100644 --- a/lib/kv-router/src/scheduling/mod.rs +++ b/lib/kv-router/src/scheduling/mod.rs @@ -4,7 +4,7 @@ pub mod config; mod filter; mod local; -pub mod lora_filter; +mod lora_filter; pub mod overlap; pub mod overlap_refresh; pub mod policy; @@ -20,7 +20,8 @@ mod worker_selection_config; mod types; pub use filter::*; pub use local::LocalScheduler; -pub use lora_filter::{LoraWorkerFilter, narrow_allowed_worker_ids_by_lora}; +pub use lora_filter::LoraWorkerFilter; +pub(crate) use lora_filter::narrow_allowed_worker_ids_by_lora; pub use overlap::{ CacheHitEstimates, OverlapAnalysis, OverlapScoresResponse, OverlapSignals, SelectedWorkerTierSnapshot, SharedCacheOverlapScore, WorkerOverlapScore, diff --git a/lib/kv-router/src/services/indexer/registry.rs b/lib/kv-router/src/services/indexer/registry.rs index 754144cc3c00..0a8dbb464d85 100644 --- a/lib/kv-router/src/services/indexer/registry.rs +++ b/lib/kv-router/src/services/indexer/registry.rs @@ -704,8 +704,10 @@ impl WorkerRegistry { dp_rank: u32, key: &RoutingPartitionId, ) { - if let Some(ie) = self.indexers.get(key) { - ie.indexer.remove_worker_dp_rank(worker_id, dp_rank).await; + // Clone the handle so the shard guard drops before the await. + let indexer = self.indexers.get(key).map(|ie| ie.indexer.clone()); + if let Some(indexer) = indexer { + indexer.remove_worker_dp_rank(worker_id, dp_rank).await; } } diff --git a/lib/kv-router/src/services/indexer/server.rs b/lib/kv-router/src/services/indexer/server.rs index c77092160269..4b25c21f979e 100644 --- a/lib/kv-router/src/services/indexer/server.rs +++ b/lib/kv-router/src/services/indexer/server.rs @@ -447,10 +447,10 @@ async fn query_tiered_by_hash( let mut resp = match indexer.find_tiered_matches(block_hashes).await { Ok(tiered) => ( StatusCode::OK, - Json(serde_json::json!(TieredQueryResponse { + Json(TieredQueryResponse { block_size, tiered: WireTieredMatchDetails::from(&tiered), - })), + }), ) .into_response(), Err(e) => ( @@ -638,8 +638,14 @@ fn build_router(state: Arc, test_endpoints: bool) -> Router { "/query", post(query).layer(DefaultBodyLimit::max(QUERY_REQUEST_BODY_LIMIT_BYTES)), ) - .route("/query_by_hash", post(query_by_hash)) - .route("/query_tiered_by_hash", post(query_tiered_by_hash)) + .route( + "/query_by_hash", + post(query_by_hash).layer(DefaultBodyLimit::max(QUERY_REQUEST_BODY_LIMIT_BYTES)), + ) + .route( + "/query_tiered_by_hash", + post(query_tiered_by_hash).layer(DefaultBodyLimit::max(QUERY_REQUEST_BODY_LIMIT_BYTES)), + ) .route("/dump", get(dump_events)) .route("/register_peer", post(register_peer)) .route("/deregister_peer", post(deregister_peer)) diff --git a/lib/kv-router/src/services/selection/affinity.rs b/lib/kv-router/src/services/selection/affinity.rs index e291b829351f..cde62dfa850b 100644 --- a/lib/kv-router/src/services/selection/affinity.rs +++ b/lib/kv-router/src/services/selection/affinity.rs @@ -303,8 +303,12 @@ impl SessionAffinity { /// Install the replica sink and this replica's writer id. Returns `false` /// when a sink is already installed. pub fn enable_replication(&self, writer_id: u64, sink: Arc) -> bool { - self.inner.writer_id.store(writer_id, Ordering::Relaxed); - self.inner.replica.set(sink).is_ok() + // Record the id only for the sink that actually got installed. + let installed = self.inner.replica.set(sink).is_ok(); + if installed { + self.inner.writer_id.store(writer_id, Ordering::Relaxed); + } + installed } pub fn ttl(&self) -> Duration { diff --git a/lib/kv-router/src/services/selection/core/reservations.rs b/lib/kv-router/src/services/selection/core/reservations.rs index 92c6fd84302a..87debde19db5 100644 --- a/lib/kv-router/src/services/selection/core/reservations.rs +++ b/lib/kv-router/src/services/selection/core/reservations.rs @@ -130,11 +130,16 @@ fn forget_reservation_if( booking: &SchedulerBookingDescriptor, ) { let mut index = index.write(); - if index.get(&booking.request_id).is_some_and(|reservation| { - reservation.partition == *partition && reservation.booking.as_ref() == Some(booking) - }) { - index.remove(&booking.request_id); - } + let removed = index + .get(&booking.request_id) + .is_some_and(|reservation| { + reservation.partition == *partition && reservation.booking.as_ref() == Some(booking) + }) + .then(|| index.remove(&booking.request_id)); + drop(index); + // The reservation's session binding releases on drop (shard lock, replica + // publish); keep that work outside the index write lock. + drop(removed); } impl SelectionCore { @@ -512,16 +517,27 @@ pub(super) fn sweep_reservation_index( let entries = entries.read(); let mut index = index.write(); let before = index.len(); + let mut removed = Vec::new(); index.retain(|_, reservation| { let Some(booking) = &reservation.booking else { return true; }; - entries + let live = entries .get(&reservation.partition) .and_then(|cell| cell.get()) - .is_some_and(|entry| entry.scheduler.has_booking(booking)) + .is_some_and(|entry| entry.scheduler.has_booking(booking)); + if !live { + removed.push(reservation._affinity_lease.take()); + } + live }); - before - index.len() + let swept = before - index.len(); + drop(index); + drop(entries); + // Session bindings release on drop (shard lock, replica publish); do that + // after both locks are gone. + drop(removed); + swept } pub(super) fn missing_booking() -> SelectionError { diff --git a/lib/kv-router/src/services/selection/core/run.rs b/lib/kv-router/src/services/selection/core/run.rs index 649d4ca4d945..24c1ee1b028d 100644 --- a/lib/kv-router/src/services/selection/core/run.rs +++ b/lib/kv-router/src/services/selection/core/run.rs @@ -357,7 +357,10 @@ impl SelectionCore { let routing_hashes = (entry.indexer.records_routing_decisions() && (claim.is_some() || cached_inputs.is_some())) .then(|| block_hashes.clone()); - let booked_sequence_hashes = book.then(|| sequence_hashes.clone()); + // Only `Book` callers read the hashes back (the reservation response); + // the frontend's `Lease` path discards them. + let booked_sequence_hashes = + matches!(admission, SelectionAdmission::Book { .. }).then(|| sequence_hashes.clone()); let returned_routing_hashes = return_routing_hashes.then(|| block_hashes.clone()); let schedule_request = ScheduleRequest { mode, diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index fe2a528fdad7..d2f7d1af22ba 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -1515,11 +1515,6 @@ impl KvRouter { } } - /// Register externally-provided workers in the slot tracker. - pub fn register_workers(&self, worker_ids: &HashSet) { - self.scheduler.register_workers(worker_ids); - } - #[allow(clippy::too_many_arguments)] pub async fn add_request( &self, diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs index 0ce88dac16e4..a3d1ec74875f 100644 --- a/lib/llm/src/kv_router/embedded.rs +++ b/lib/llm/src/kv_router/embedded.rs @@ -8,7 +8,7 @@ //! and request constraints before scheduling directly on that partition. The //! frontend retains transport, stream leases, and request-expiry ownership. -use std::collections::{HashMap, HashSet}; +use std::collections::HashMap; use std::sync::Arc; use anyhow::{Context, Result}; @@ -407,10 +407,6 @@ impl EmbeddedSelection { .clone()) } - /// Membership is catalog-driven (runtime-config watch); explicit worker - /// registration is a no-op here. - pub(crate) fn register_workers(&self, _worker_ids: &HashSet) {} - pub(crate) fn partition_key(&self) -> &RoutingPartitionId { self.partition.key() } diff --git a/lib/llm/src/kv_router/prefill_router/query.rs b/lib/llm/src/kv_router/prefill_router/query.rs index 8c40384e5bb7..bf84c8c15bba 100644 --- a/lib/llm/src/kv_router/prefill_router/query.rs +++ b/lib/llm/src/kv_router/prefill_router/query.rs @@ -216,14 +216,6 @@ impl PrefillRouter { } } } - - pub fn register_workers(&self, worker_ids: &HashSet) { - if let Some(binding) = self.binding.load_full() - && let Some(kv_router) = binding.router.kv_router_if_enabled() - { - kv_router.register_workers(worker_ids); - } - } } #[cfg(test)] From 317eb30e8639550dbc7692723387a6dc2ae97ef5 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Fri, 11 Sep 2026 14:34:15 -0700 Subject: [PATCH 063/102] chore(router): clear the small self-review findings Forty-six review findings from the 2026-09-10 pass, each a local change with no behavior change intended: - indexer: tiered and flat lookups borrow the hash slice instead of cloning it (`find_tiered_matches_ref_with_options`, used by the selection core and the tiered match provider); `find_matches` checks for a side indexer before copying; the indexer policy is cloned out of its mutex before the indexer build closure; the unused `indexer_policy()` getter and the `record_routing_decision_hashes` forwarder are gone; `TieredQueryByHashRequest.routing_group` gets the serde default the sibling request has; `try_apply_event` traces the dropped-event arm like its siblings; the merge-overcount note is back on `merge_overlap_scores`. - selection core: `KV_HINT_FETCH_ACTION_ID` names the hint action; `dispatch_affinity_event` traces its silent returns; the advisory doc no longer promises queue-rejection parity; the per-partition "session affinity enabled" line is `debug!`; log messages start in sentence case; inline `std::` paths fold into imports; `DEFAULT_MODEL_NAME` is exported once and reused by the frontend host. - sequences: `request_worker`, `worker_for`, `booking_for` take `&str`, so `has_request` stops allocating per probe. - replica sync: publish failures log the partition and request or session id; `SchedulerLoadSnapshot` is imported once. - frontend: `impl Drop for Inner` in the affinity coordinator was a duplicate of `ReplicaSyncRuntime::drop`; the replica-enable error names its trigger; the `RegisteredGauge` clamp, the dead registry argument and `..WorkerRequest::default()` are gone from the embedded host; `#[cfg(test)]` helpers move into their test modules; the inherent `find_matches_by_tier` on the remote indexer is renamed so it no longer shadows the trait method; two adjacent `impl ModelWatcher` blocks are one; AGENTS.md names `embedded.rs`. - ext-proc: `HEADER_DYNAMO_SESSION_ID` is shared from `dynamo_llm` instead of duplicated; `DYN_EPP_SESSION_AFFINITY_TTL_SECS` is validated and its bound documented; `Selector::core()` is crate-private; the TTL test asserts the error text. - bindings: the TTL bound uses `MAX_SESSION_AFFINITY_TTL_SECS`; the `.pyi` documents `remote_indexer_url`; timeout markers on two tests. - `reqwest` in dynamo-kv-router enables `rustls-tls` like the workspace. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Ishan Dhanani --- .../ext-proc/src/epp_router.rs | 6 +- .../ext-proc/src/epp_standalone_config.rs | 5 ++ .../ext-proc/src/pod_discovery.rs | 4 +- .../ext-proc/src/selector.rs | 12 +-- .../ext-proc/src/topology_adapter.rs | 2 +- .../kv-aware-routing/vanilla-vllm-onramp.mdx | 2 +- lib/bindings/python/rust/llm/entrypoint.rs | 11 ++- lib/bindings/python/rust/llm/kv.rs | 16 ++-- lib/bindings/python/src/dynamo/_core.pyi | 4 + lib/bindings/python/tests/test_kv_bindings.py | 2 + lib/kv-router/Cargo.toml | 2 +- lib/kv-router/src/scheduling/local.rs | 2 +- lib/kv-router/src/sequences/multi_worker.rs | 76 ++++------------- lib/kv-router/src/sequences/request_maps.rs | 8 +- .../src/services/common/replica_sync.rs | 26 ++++-- lib/kv-router/src/services/indexer/backend.rs | 63 ++++++++++---- lib/kv-router/src/services/indexer/lookup.rs | 13 ++- .../src/services/indexer/recording.rs | 10 --- .../src/services/indexer/registry.rs | 10 +-- .../src/services/selection/README.md | 13 ++- .../src/services/selection/affinity.rs | 12 +-- .../src/services/selection/core/mod.rs | 15 +++- .../src/services/selection/core/run.rs | 13 +-- .../src/services/selection/core/tests.rs | 2 +- lib/kv-router/src/services/selection/mod.rs | 11 +-- .../src/services/selection/service.rs | 5 +- lib/kv-router/src/services/selection/types.rs | 10 ++- lib/llm/src/discovery/watcher.rs | 2 - lib/llm/src/kv_router.rs | 14 +-- lib/llm/src/kv_router/AGENTS.md | 2 +- lib/llm/src/kv_router/embedded.rs | 39 +++++---- lib/llm/src/kv_router/indexer/ingress.rs | 8 +- lib/llm/src/kv_router/indexer/remote.rs | 4 +- lib/llm/src/kv_router/sequence.rs | 18 ---- lib/llm/src/lora/load_estimator.rs | 85 ++++++++++--------- lib/llm/src/protocols/agents.rs | 2 +- lib/llm/src/session_affinity/coordinator.rs | 10 +-- 37 files changed, 275 insertions(+), 264 deletions(-) diff --git a/deploy/inference-gateway/ext-proc/src/epp_router.rs b/deploy/inference-gateway/ext-proc/src/epp_router.rs index e3fc72baba9a..8bdcde4c1404 100644 --- a/deploy/inference-gateway/ext-proc/src/epp_router.rs +++ b/deploy/inference-gateway/ext-proc/src/epp_router.rs @@ -28,6 +28,7 @@ use tokio::sync::Semaphore; use dynamo_kv_router::services::selection::{SelectionError, WorkerSelectionPolicyRegistry}; use dynamo_llm::http::service::metadata::extract_metadata_from_header_pairs; +use dynamo_llm::protocols::agents::HEADER_DYNAMO_SESSION_ID; use dynamo_llm::protocols::common::extensions::{ AgentHints, HEADER_REQUEST_PRIORITY, HEADER_REQUEST_STRICT_PRIORITY, resolve_request_priority, }; @@ -70,9 +71,6 @@ impl RenderClient { } } -/// Session id the standalone EPP pins to a worker when session affinity is on. -const HEADER_SESSION_ID: &str = "x-dynamo-session-id"; - /// Standalone endpoint picker backed by the standalone selection service. pub struct EppRouter { renderer: RenderClient, @@ -344,7 +342,7 @@ impl EndpointPicker for EppRouter { model_name: self.model_name.clone(), reservation_id: reservation_id.clone(), token_ids: tokens, - session_id: first_header(&req.headers, HEADER_SESSION_ID).map(str::to_owned), + session_id: first_header(&req.headers, HEADER_DYNAMO_SESSION_ID).map(str::to_owned), // `None` on the ordinary path: the selector schedules over its // catalog; `Some` only carries an Envoy subset constraint. allowed_worker_ids: allowed, diff --git a/deploy/inference-gateway/ext-proc/src/epp_standalone_config.rs b/deploy/inference-gateway/ext-proc/src/epp_standalone_config.rs index 24e31f1cbf7f..93b7120fb5db 100644 --- a/deploy/inference-gateway/ext-proc/src/epp_standalone_config.rs +++ b/deploy/inference-gateway/ext-proc/src/epp_standalone_config.rs @@ -177,6 +177,11 @@ pub struct EppStandaloneConfig { /// Pin each `x-dynamo-session-id` to the worker that served it for this /// long after its last request (`DYN_EPP_SESSION_AFFINITY_TTL_SECS`). /// `None` disables session affinity. + #[validate(range( + min = 1.0, + max = 31536000.0, + message = "DYN_EPP_SESSION_AFFINITY_TTL_SECS must be between 1 and 31536000 seconds" + ))] pub session_affinity_ttl_secs: Option, } diff --git a/deploy/inference-gateway/ext-proc/src/pod_discovery.rs b/deploy/inference-gateway/ext-proc/src/pod_discovery.rs index d0f425968e15..1920c1f2f5b4 100644 --- a/deploy/inference-gateway/ext-proc/src/pod_discovery.rs +++ b/deploy/inference-gateway/ext-proc/src/pod_discovery.rs @@ -404,8 +404,6 @@ fn rebuild_index( } } -/// Build a [`RawWorker`] from a pod, or `None` if it is not `Ready`, not -/// pool-selected, or lacks an IP/name. Pure function — unit-testable. /// Per-pod KV-event port layout: rank `r` publishes on `base + r * stride`. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) struct KvEventPorts { @@ -427,6 +425,8 @@ impl KvEventPorts { } } +/// Build a [`RawWorker`] from a pod, or `None` if it is not `Ready`, not +/// pool-selected, or lacks an IP/name. Pure function — unit-testable. fn raw_worker_from_pod( pod: &Pod, pool: &PoolState, diff --git a/deploy/inference-gateway/ext-proc/src/selector.rs b/deploy/inference-gateway/ext-proc/src/selector.rs index cd0a3f44ca05..e00c9499e005 100644 --- a/deploy/inference-gateway/ext-proc/src/selector.rs +++ b/deploy/inference-gateway/ext-proc/src/selector.rs @@ -167,7 +167,7 @@ impl Selector { } /// The selection core whose catalog the topology adapter feeds. - pub fn core(&self) -> &Arc { + pub(crate) fn core(&self) -> &Arc { self.service.core() } @@ -906,11 +906,13 @@ worker_selection: for ttl in [-1.0, 0.0, 0.5, f64::NAN, f64::INFINITY] { let mut cfg = test_config(); cfg.session_affinity_ttl_secs = Some(ttl); + let error = Selector::new(&cfg, WorkerSelectionPolicyRegistry::default()) + .await + .err() + .unwrap_or_else(|| panic!("TTL={ttl} must be rejected")); assert!( - Selector::new(&cfg, WorkerSelectionPolicyRegistry::default()) - .await - .is_err(), - "TTL={ttl}" + format!("{error:#}").contains("session affinity TTL"), + "TTL={ttl}: {error:#}" ); } } diff --git a/deploy/inference-gateway/ext-proc/src/topology_adapter.rs b/deploy/inference-gateway/ext-proc/src/topology_adapter.rs index d81708f6a238..34a9fb48f644 100644 --- a/deploy/inference-gateway/ext-proc/src/topology_adapter.rs +++ b/deploy/inference-gateway/ext-proc/src/topology_adapter.rs @@ -58,7 +58,7 @@ impl WorkerCatalogSource for PodReflectorSource { return None; } if self.primed && self.changes.changed().await.is_err() { - tracing::warn!("reflector change channel closed; clearing selector topology"); + tracing::warn!("Reflector change channel closed; clearing selector topology"); self.closed = true; return Some(Vec::new()); } diff --git a/docs/fern/pages/kubernetes/kv-aware-routing/vanilla-vllm-onramp.mdx b/docs/fern/pages/kubernetes/kv-aware-routing/vanilla-vllm-onramp.mdx index ebe56b753c2a..df7a559662f3 100644 --- a/docs/fern/pages/kubernetes/kv-aware-routing/vanilla-vllm-onramp.mdx +++ b/docs/fern/pages/kubernetes/kv-aware-routing/vanilla-vllm-onramp.mdx @@ -504,7 +504,7 @@ For a single EPP replica, set `spec.replicas: 1` and remove `DYN_EPP_PEER_SERVIC | `DYN_EPP_SELECTION_INDEXER_THREADS` | No | KV indexer thread count; defaults to `4`. | | `DYN_EPP_TOKENIZATION_TIMEOUT_MS` | No | Render request timeout; defaults to `5000`. | | `DYN_EPP_TOKENIZER_MAX_RESPONSE_BYTES` | No | Maximum render response size; defaults to `16777216`. | -| `DYN_EPP_SESSION_AFFINITY_TTL_SECS` | No | Pin each `x-dynamo-session-id` to the worker that served it for this many seconds after its last request; the binding replicates to peer EPPs over the replica mesh. Unset disables session affinity. | +| `DYN_EPP_SESSION_AFFINITY_TTL_SECS` | No | Pin each `x-dynamo-session-id` to the worker that served it for this many seconds (1 to 31,536,000) after its last request; the binding replicates to peer EPPs over the replica mesh. Unset disables session affinity. | | `DYN_EPP_MAX_INFLIGHT_REQUESTS` | No | Concurrent EPP request guardrail; defaults to `1024`, and excess requests receive `503`. | | `DYN_EPP_PEER_SERVICE` | No | Existing EPP Service used to discover sibling replicas. Enables replica sync. | | `DYN_EPP_REPLICA_SYNC_PORT` | No | ZMQ port every replica binds and dials for replica sync; defaults to `9092`. All replicas selected by `DYN_EPP_PEER_SERVICE` must use the same value. | diff --git a/lib/bindings/python/rust/llm/entrypoint.rs b/lib/bindings/python/rust/llm/entrypoint.rs index 3634bdb8f9fe..ad63a0913f1c 100644 --- a/lib/bindings/python/rust/llm/entrypoint.rs +++ b/lib/bindings/python/rust/llm/entrypoint.rs @@ -32,6 +32,7 @@ use dynamo_llm::local_model::{LocalModel, LocalModelBuilder}; use dynamo_llm::mocker::make_mocker_engine; use dynamo_llm::model_card::ModelDeploymentCard as RsModelDeploymentCard; use dynamo_llm::reasoning_field::ReasoningField; +use dynamo_llm::session_affinity::MAX_SESSION_AFFINITY_TTL_SECS; use dynamo_llm::session_affinity::SessionAffinityMode as RsSessionAffinityMode; use dynamo_llm::types::openai::chat_completions::OpenAIChatCompletionsStreamingEngine; use dynamo_mocker::common::perf_model::PerfModel; @@ -507,10 +508,12 @@ impl RouterConfig { ); }); } - if session_affinity_ttl_secs.is_some_and(|ttl| !(1..=31_536_000).contains(&ttl)) { - return Err(PyValueError::new_err( - "session_affinity_ttl_secs must be between 1 and 31536000", - )); + if session_affinity_ttl_secs + .is_some_and(|ttl| !(1..=MAX_SESSION_AFFINITY_TTL_SECS).contains(&ttl)) + { + return Err(PyValueError::new_err(format!( + "session_affinity_ttl_secs must be between 1 and {MAX_SESSION_AFFINITY_TTL_SECS}" + ))); } RsLoadThresholdConfig { active_decode_blocks_threshold, diff --git a/lib/bindings/python/rust/llm/kv.rs b/lib/bindings/python/rust/llm/kv.rs index 1cea5c26df84..1936666281a8 100644 --- a/lib/bindings/python/rust/llm/kv.rs +++ b/lib/bindings/python/rust/llm/kv.rs @@ -52,7 +52,9 @@ type RsManagedKvRouter = llm_rs::kv_router::ManagedKvRouter; use llm_rs::kv_router::publisher::{KvEventSourceConfig, create_stored_blocks}; use llm_rs::protocols::common::timing::RequestTracker; use llm_rs::protocols::common::{OutputOptions, SamplingOptions, StopConditions}; -use llm_rs::session_affinity::SessionAffinityMode as RsSessionAffinityMode; +use llm_rs::session_affinity::{ + MAX_SESSION_AFFINITY_TTL_SECS, SessionAffinityMode as RsSessionAffinityMode, +}; use super::aic_callback::create_aic_prefill_load_estimator; use super::entrypoint::AicPerfConfig; @@ -640,8 +642,6 @@ impl SelectionCacheConfig { #[cfg(feature = "select-service")] fn session_affinity_ttl_from_secs(ttl: f64) -> Result { - use selection::affinity::MAX_SESSION_AFFINITY_TTL_SECS; - if !(1.0..=MAX_SESSION_AFFINITY_TTL_SECS as f64).contains(&ttl) { return Err(format!( "session_affinity_ttl_secs must be between 1 and {MAX_SESSION_AFFINITY_TTL_SECS}" @@ -2230,10 +2230,12 @@ impl KvRouter { load_threshold_config: Option<&LoadThresholdConfig>, session_affinity_mode: &str, ) -> PyResult { - if session_affinity_ttl_secs.is_some_and(|ttl| !(1..=31_536_000).contains(&ttl)) { - return Err(PyValueError::new_err( - "session_affinity_ttl_secs must be between 1 and 31536000", - )); + if session_affinity_ttl_secs + .is_some_and(|ttl| !(1..=MAX_SESSION_AFFINITY_TTL_SECS).contains(&ttl)) + { + return Err(PyValueError::new_err(format!( + "session_affinity_ttl_secs must be between 1 and {MAX_SESSION_AFFINITY_TTL_SECS}" + ))); } let session_affinity_mode = session_affinity_mode .parse::() diff --git a/lib/bindings/python/src/dynamo/_core.pyi b/lib/bindings/python/src/dynamo/_core.pyi index ee2caab10b56..32091729f281 100644 --- a/lib/bindings/python/src/dynamo/_core.pyi +++ b/lib/bindings/python/src/dynamo/_core.pyi @@ -753,6 +753,10 @@ class SelectionService: ) -> None: """Create a selection service. `indexer_threads` sizes the KV indexer pool. + `remote_indexer_url` is the base URL of a standalone indexer that serves + the primary KV index; when set, this service does not subscribe to worker + KV events itself. + `session_affinity_ttl_secs` enables session affinity with an idle TTL between 1 and 31536000 seconds. """ diff --git a/lib/bindings/python/tests/test_kv_bindings.py b/lib/bindings/python/tests/test_kv_bindings.py index bb9bca53222e..e7643d8baf61 100644 --- a/lib/bindings/python/tests/test_kv_bindings.py +++ b/lib/bindings/python/tests/test_kv_bindings.py @@ -210,6 +210,7 @@ def worker(worker_id, prepopulate_worker_ids: bool = False): @pytest.mark.parametrize( "ttl", [-1, 0, 0.5, float("nan"), float("inf"), 31536001, 1e300] ) +@pytest.mark.timeout(5) def test_selection_service_rejects_invalid_affinity_ttl(ttl): with pytest.raises(ValueError, match="session_affinity_ttl_secs must be between"): SelectionService(session_affinity_ttl_secs=ttl) @@ -243,6 +244,7 @@ def test_selection_service_cli_rejects_invalid_affinity_ttl(ttl): reason="SelectionService requires the select-service Cargo feature", ) @pytest.mark.parametrize("ttl", [1.0, 1.5, 31536000.0]) +@pytest.mark.timeout(5) def test_selection_service_accepts_valid_affinity_ttl(ttl): service = SelectionService(indexer_threads=1, session_affinity_ttl_secs=ttl) service.shutdown() diff --git a/lib/kv-router/Cargo.toml b/lib/kv-router/Cargo.toml index deeb18758059..605f24de3233 100644 --- a/lib/kv-router/Cargo.toml +++ b/lib/kv-router/Cargo.toml @@ -75,7 +75,7 @@ axum = { version = "=0.8.4", default-features = false, features = [ "tracing", ], optional = true } chrono = { workspace = true, optional = true } -reqwest = { version = "0.12.24", default-features = false, features = ["json"], optional = true } +reqwest = { version = "0.12.24", default-features = false, features = ["json", "rustls-tls"], optional = true } tracing-appender = { version = "0.2", optional = true } [dev-dependencies] diff --git a/lib/kv-router/src/scheduling/local.rs b/lib/kv-router/src/scheduling/local.rs index 6eb5f8c12e18..ff685a5bb111 100644 --- a/lib/kv-router/src/scheduling/local.rs +++ b/lib/kv-router/src/scheduling/local.rs @@ -602,7 +602,7 @@ where /// Whether `request_id` currently holds a booking on any worker. pub fn has_request(&self, request_id: &str) -> bool { - self.slots.request_worker(&request_id.to_string()).is_some() + self.slots.request_worker(request_id).is_some() } #[doc(hidden)] diff --git a/lib/kv-router/src/sequences/multi_worker.rs b/lib/kv-router/src/sequences/multi_worker.rs index 8032072108d6..c46f1c829cb7 100644 --- a/lib/kv-router/src/sequences/multi_worker.rs +++ b/lib/kv-router/src/sequences/multi_worker.rs @@ -824,7 +824,7 @@ impl ActiveSequencesMultiWorker

{ Ok(attempt_id) } - pub(crate) fn request_worker(&self, request_id: &RequestId) -> Option { + pub(crate) fn request_worker(&self, request_id: &str) -> Option { self.request_index.worker_for(request_id) } @@ -2491,10 +2491,7 @@ mod tests { ); assert_eq!(active_request_count(&sequences, worker_a), 1); assert_eq!(active_request_count(&sequences, worker_b), 0); - assert_eq!( - sequences.request_worker(&"req-1".to_string()), - Some(worker_a) - ); + assert_eq!(sequences.request_worker("req-1"), Some(worker_a)); } #[tokio::test] @@ -2523,10 +2520,7 @@ mod tests { .unwrap(); assert_eq!(active_request_count(&sequences, booked_worker), 1); - assert_eq!( - sequences.request_worker(&"req-1".to_string()), - Some(booked_worker) - ); + assert_eq!(sequences.request_worker("req-1"), Some(booked_worker)); } #[test] @@ -3078,14 +3072,8 @@ mod tests { assert_eq!(batches[0][0].active_prefill_tokens, 0); assert_eq!(sequences.remote_state_update_count(), 1); assert_eq!(sequences.prompt_registry.cleanup_attempts(), 1); - assert_eq!( - sequences.request_index.worker_for(&"req-2".to_string()), - Some(worker) - ); - assert_eq!( - sequences.request_index.worker_for(&"req-1".to_string()), - None - ); + assert_eq!(sequences.request_index.worker_for("req-2"), Some(worker)); + assert_eq!(sequences.request_index.worker_for("req-1"), None); assert_eq!( sequences.active_request_counts().get(&worker).copied(), Some(1) @@ -3112,14 +3100,8 @@ mod tests { assert_eq!(batches[0][0].active_prefill_tokens, 0); assert_eq!(sequences.remote_state_update_count(), 1); assert_eq!(sequences.prompt_registry.cleanup_attempts(), 1); - assert_eq!( - sequences.request_index.worker_for(&"req-2".to_string()), - Some(worker) - ); - assert_eq!( - sequences.request_index.worker_for(&"req-1".to_string()), - None - ); + assert_eq!(sequences.request_index.worker_for("req-2"), Some(worker)); + assert_eq!(sequences.request_index.worker_for("req-1"), None); } #[test] @@ -3182,11 +3164,11 @@ mod tests { "source B should apply while source A remains blocked" ); assert_eq!( - sequences.request_index.worker_for(&"source-a".to_string()), + sequences.request_index.worker_for("source-a"), Some(worker_a) ); assert_eq!( - sequences.request_index.worker_for(&"source-b".to_string()), + sequences.request_index.worker_for("source-b"), Some(worker_b) ); } @@ -3239,10 +3221,7 @@ mod tests { sequences.apply_replica_batch(vec![replica_add("req-1", worker_b, vec![4, 5, 6])]); - assert_eq!( - sequences.request_index.worker_for(&"req-1".to_string()), - Some(worker_a) - ); + assert_eq!(sequences.request_index.worker_for("req-1"), Some(worker_a)); assert_eq!(sequences.active_blocks()[&worker_a], 3); assert!(!sequences.active_blocks().contains_key(&worker_b)); assert_eq!(sequences.num_workers(), 1); @@ -3293,10 +3272,7 @@ mod tests { sequences.prompt_registry.cleanup_attempts(), cleanup_count + 1 ); - assert_eq!( - sequences.request_index.worker_for(&"req-1".to_string()), - None - ); + assert_eq!(sequences.request_index.worker_for("req-1"), None); } #[tokio::test(start_paused = true)] @@ -3332,10 +3308,7 @@ mod tests { .copied(), Some(0) ); - assert_eq!( - sequences.request_index.worker_for(&"req-1".to_string()), - Some(worker) - ); + assert_eq!(sequences.request_index.worker_for("req-1"), Some(worker)); } #[tokio::test(start_paused = true)] @@ -3414,10 +3387,7 @@ mod tests { let batches = publisher.load_batches(); assert_eq!(batches.len(), 1); assert_eq!(batches[0].len(), 1); - assert_eq!( - sequences.request_index.worker_for(&"req-1".to_string()), - Some(worker) - ); + assert_eq!(sequences.request_index.worker_for("req-1"), Some(worker)); } #[tokio::test(start_paused = true)] @@ -3439,10 +3409,7 @@ mod tests { let batches = publisher.load_batches(); assert_eq!(batches.len(), 1); assert_eq!(batches[0].len(), 1); - assert_eq!( - sequences.request_index.worker_for(&"req-1".to_string()), - Some(worker) - ); + assert_eq!(sequences.request_index.worker_for("req-1"), Some(worker)); } #[tokio::test] @@ -3703,10 +3670,7 @@ mod tests { .unwrap(); assert_eq!(sequences.num_workers(), 1); - assert_eq!( - sequences.request_index.worker_for(&"req-1".to_string()), - Some(worker) - ); + assert_eq!(sequences.request_index.worker_for("req-1"), Some(worker)); assert!(!sequences.prompt_registry.is_block_index_empty()); assert_eq!(sequences.active_blocks().get(&worker).copied(), Some(3)); } @@ -3735,10 +3699,7 @@ mod tests { .unwrap(); assert_eq!(sequences.num_workers(), 0); - assert_eq!( - sequences.request_index.worker_for(&"req-1".to_string()), - None - ); + assert_eq!(sequences.request_index.worker_for("req-1"), None); assert!(sequences.prompt_registry.is_block_index_empty()); } @@ -3767,10 +3728,7 @@ mod tests { .unwrap(); assert_eq!(sequences.num_workers(), 0); - assert_eq!( - sequences.request_index.worker_for(&"req-1".to_string()), - None - ); + assert_eq!(sequences.request_index.worker_for("req-1"), None); } #[test] diff --git a/lib/kv-router/src/sequences/request_maps.rs b/lib/kv-router/src/sequences/request_maps.rs index 114fe256cd87..583c12a2fbea 100644 --- a/lib/kv-router/src/sequences/request_maps.rs +++ b/lib/kv-router/src/sequences/request_maps.rs @@ -60,11 +60,11 @@ impl RequestIndex { } } - pub(super) fn worker_for(&self, request_id: &RequestId) -> Option { + pub(super) fn worker_for(&self, request_id: &str) -> Option { self.booking_for(request_id).map(|booking| booking.worker) } - pub(super) fn booking_for(&self, request_id: &RequestId) -> Option { + pub(super) fn booking_for(&self, request_id: &str) -> Option { self.request_to_booking.get(request_id).map(|entry| *entry) } @@ -172,7 +172,7 @@ mod tests { index.try_insert_request("req-1".to_string(), WorkerWithDpRank::new(2, 0), None), Err(worker) ); - assert_eq!(index.worker_for(&"req-1".to_string()), Some(worker)); + assert_eq!(index.worker_for("req-1"), Some(worker)); assert_eq!( index.lora_for(&"req-1".to_string()), Some("adapter".to_string()) @@ -222,7 +222,7 @@ mod tests { let mut removed = index.remove_worker_requests(worker_a); removed.sort(); assert_eq!(removed, vec!["req-a".to_string(), "req-c".to_string()]); - assert_eq!(index.worker_for(&"req-b".to_string()), Some(worker_b)); + assert_eq!(index.worker_for("req-b"), Some(worker_b)); assert_eq!( index.active_lora_counts(), HashMap::from([("adapter-b".to_string(), 1)]) diff --git a/lib/kv-router/src/services/common/replica_sync.rs b/lib/kv-router/src/services/common/replica_sync.rs index 5c7d0961f623..b425239a0de9 100644 --- a/lib/kv-router/src/services/common/replica_sync.rs +++ b/lib/kv-router/src/services/common/replica_sync.rs @@ -14,7 +14,9 @@ use tokio_util::sync::CancellationToken; use crate::identity::{RoutingPartitionId, RoutingPartitionRef}; use crate::protocols::{ActiveSequenceEvent, WorkerWithDpRank}; -use crate::sequences::{SequencePublishQueueError, SequencePublisher, SequenceSubscriber}; +use crate::sequences::{ + SchedulerLoadSnapshot, SequencePublishQueueError, SequencePublisher, SequenceSubscriber, +}; use crate::services::common::zmq::{ create_bound_pub_socket, create_sub_socket_topics, validate_endpoint, }; @@ -202,14 +204,14 @@ pub(crate) struct ScopedReplicaSync { /// tokens). An embedding host uses it to drive overload detection from the /// same numbers the scheduler books against. pub trait SchedulerLoadSink: Send + Sync { - fn publish(&self, snapshot: crate::sequences::SchedulerLoadSnapshot); + fn publish(&self, snapshot: SchedulerLoadSnapshot); /// Per-worker load after any local mutation, including output blocks, /// which are never published as shared scheduler load. The sink owns the /// metric label so it matches the host's cleanup path. fn observe_local_load(&self, _worker: &WorkerWithDpRank, _blocks: usize, _tokens: usize) {} - fn publish_batch(&self, snapshots: Vec) { + fn publish_batch(&self, snapshots: Vec) { for snapshot in snapshots { self.publish(snapshot); } @@ -343,13 +345,13 @@ impl SequencePublisher for ScopedSequencePublisher { } } - fn publish_scheduler_load(&self, load: crate::sequences::SchedulerLoadSnapshot) { + fn publish_scheduler_load(&self, load: SchedulerLoadSnapshot) { if let Some(sink) = &self.load_sink { sink.publish(load); } } - fn publish_scheduler_load_batch(&self, loads: Vec) { + fn publish_scheduler_load_batch(&self, loads: Vec) { if let Some(sink) = &self.load_sink { sink.publish_batch(loads); } @@ -493,7 +495,7 @@ pub(crate) fn start_replica_publisher( let task = tokio::spawn(async move { loop { - let frames = tokio::select! { + let (frames, label) = tokio::select! { _ = cancel_token.cancelled() => break, event = rx.recv() => { let Some(event) = event else { @@ -501,7 +503,10 @@ pub(crate) fn start_replica_publisher( }; let partition = event.partition_ref(); match rmp_serde::to_vec_named(&event) { - Ok(payload) => vec![REPLICA_TOPIC.to_vec(), payload], + Ok(payload) => ( + vec![REPLICA_TOPIC.to_vec(), payload], + format!("{partition} request_id={}", event.event.request_id), + ), Err(error) => { tracing::error!( model_name = %partition.model_name, @@ -518,7 +523,10 @@ pub(crate) fn start_replica_publisher( break; }; match rmp_serde::to_vec_named(&binding) { - Ok(payload) => vec![AFFINITY_TOPIC.to_vec(), payload], + Ok(payload) => ( + vec![AFFINITY_TOPIC.to_vec(), payload], + format!("session_id={}", binding.session_id), + ), Err(error) => { tracing::error!( session_id = %binding.session_id, @@ -530,7 +538,7 @@ pub(crate) fn start_replica_publisher( } }; if let Err(error) = socket.send_multipart(frames).await { - tracing::error!("Failed to publish replica event: {error}"); + tracing::error!(event = %label, "Failed to publish replica event: {error}"); } } }); diff --git a/lib/kv-router/src/services/indexer/backend.rs b/lib/kv-router/src/services/indexer/backend.rs index 65f2774616cd..296c8be0b609 100644 --- a/lib/kv-router/src/services/indexer/backend.rs +++ b/lib/kv-router/src/services/indexer/backend.rs @@ -13,7 +13,7 @@ use crate::ConcurrentRadixTreeCompressed; use crate::ThreadPoolIndexer; use crate::approx::PruneConfig; use crate::config::{ApproximateCachePolicyKind, KvRouterConfig}; -use crate::identity::RoutingPartitionId; +use crate::identity::{RoutingPartitionId, default_routing_group}; use crate::indexer::{ ApproximateLruIncarnation, ApproximateLruStats, ApproximateRetentionConfig, KvIndexer, KvIndexerInterface, KvIndexerMetrics, KvRouterError, LowerTierIndexers, LowerTierQueryOptions, @@ -89,6 +89,7 @@ impl RemoteIndexerTransport { pub struct TieredQueryByHashRequest { pub block_hashes: Vec, pub model_name: String, + #[serde(default = "default_routing_group")] pub routing_group: String, } @@ -291,7 +292,11 @@ impl SideIndexer { sequence: HashInput<'_>, ) -> std::result::Result { match self { - Self::Single(indexer) => indexer.find_matches(sequence.as_slice().to_vec()).await, + Self::Single(indexer) => { + indexer + .find_matches(sequence.into_owned_at_boundary()) + .await + } Self::Concurrent(indexer) => { Ok(indexer.backend().find_matches(sequence.as_slice(), false)) } @@ -669,7 +674,9 @@ impl Indexer { .enqueue_event(event)?; } } - Self::Remote { .. } | Self::None => {} + Self::Remote { .. } | Self::None => { + tracing::trace!("Dropping KV event: no local primary indexer"); + } } Ok(()) } @@ -795,21 +802,10 @@ impl Indexer { /// Device-tier overlap scores, including side-indexer credit. Tier-aware /// callers should use [`Self::find_tiered_matches`]. pub async fn find_matches(&self, hashes: Vec) -> Result { - let primary = match self { - Indexer::Single { primary, .. } => primary.find_matches(hashes.clone()).await?, - Indexer::Concurrent { primary, .. } => primary.find_matches(hashes.clone()).await?, - Indexer::Remote { primary, .. } => { - primary - .find_matches_by_tier(hashes.clone(), true) - .await? - .device - .overlap_scores - } - Indexer::None => return Ok(OverlapScores::default()), - }; let Some(side) = self.approx() else { - return Ok(primary); + return self.find_primary_matches(hashes).await; }; + let primary = self.find_primary_matches(hashes.clone()).await?; let mut merged = MatchDetails::new(); merged.overlap_scores = primary; Ok( @@ -819,6 +815,22 @@ impl Indexer { ) } + /// Device-tier overlap scores from the primary only (no side-indexer credit). + async fn find_primary_matches(&self, hashes: Vec) -> Result { + Ok(match self { + Indexer::Single { primary, .. } => primary.find_matches(hashes).await?, + Indexer::Concurrent { primary, .. } => primary.find_matches(hashes).await?, + Indexer::Remote { primary, .. } => { + primary + .find_matches_by_tier(hashes, true) + .await? + .device + .overlap_scores + } + Indexer::None => OverlapScores::default(), + }) + } + /// Device match details + per-tier hits, suitable for building the /// Mooncake-RFC-shape per-instance breakdown. pub async fn find_tiered_matches( @@ -845,6 +857,20 @@ impl Indexer { .await } + /// Borrowed variant of [`Self::find_tiered_matches_with_options`]. + pub async fn find_tiered_matches_ref_with_options( + &self, + sequence: &[LocalBlockHash], + options: LowerTierQueryOptions, + ) -> std::result::Result { + let options = LowerTierQueryOptions { + retain_kv_transfer_chain: options.retain_kv_transfer_chain + && self.supports_kv_transfer_chain_retention(), + }; + self.find_matches_by_tier_ref_with_options(sequence, options) + .await + } + /// Dump every event tracked by this indexer in a form replayable by a /// peer's [`Self::apply_event_routed`]: the primary device-tier dump first, /// then every allocated lower-tier indexer's dump with `storage_tier` @@ -924,7 +950,8 @@ impl TieredMatchProvider for Indexer { &self, sequence: &[LocalBlockHash], ) -> std::result::Result { - self.find_tiered_matches(sequence.to_vec()).await + self.find_tiered_matches_ref_with_options(sequence, LowerTierQueryOptions::default()) + .await } async fn find_tiered_matches_with_options( @@ -932,7 +959,7 @@ impl TieredMatchProvider for Indexer { sequence: &[LocalBlockHash], options: LowerTierQueryOptions, ) -> std::result::Result { - self.find_tiered_matches_with_options(sequence.to_vec(), options) + self.find_tiered_matches_ref_with_options(sequence, options) .await } } diff --git a/lib/kv-router/src/services/indexer/lookup.rs b/lib/kv-router/src/services/indexer/lookup.rs index 1ab24bc3bc35..fca609242417 100644 --- a/lib/kv-router/src/services/indexer/lookup.rs +++ b/lib/kv-router/src/services/indexer/lookup.rs @@ -33,7 +33,7 @@ impl<'a> HashInput<'a> { self.as_slice().to_vec() } - fn into_owned_at_boundary(self) -> Vec { + pub(super) fn into_owned_at_boundary(self) -> Vec { match self { Self::Borrowed(hashes) => hashes.to_vec(), Self::Owned(hashes) => hashes, @@ -362,6 +362,13 @@ impl<'a> PrimaryLookup<'a> { /// lockstep (side-only workers gain a score with no paired hash), so it must /// not seed `query_lower_tiers`; callers run the lower-tier query against the /// primary-only details first. +/// +/// NOTE: when this merged `MatchDetails` is combined with lower-tier hits +/// seeded from the primary-only anchor (e.g. in `find_matches_by_tier`), the +/// total cached-token signal can overcount: the device score is raised by the +/// side indexer but the lower-tier walk used the lower primary depth. Accepted +/// since side scores are short-TTL approximations and the overcount is bounded +/// and rare in practice. fn merge_overlap_scores(mut primary: MatchDetails, side: OverlapScores) -> MatchDetails { for (worker, side_score) in side.scores { primary @@ -818,7 +825,7 @@ mod tests { let local_hashes = vec![LocalBlockHash(91), LocalBlockHash(92)]; let sequence_hashes = compute_seq_hash_for_block(&local_hashes); indexer - .record_routing_decision_hashes( + .record_routing_decision( worker, RoutingDecisionHashes { local_hashes: local_hashes.clone(), @@ -1037,7 +1044,7 @@ mod tests { let side_hashes = vec![LocalBlockHash(11), LocalBlockHash(12), LocalBlockHash(13)]; indexer - .record_routing_decision_hashes( + .record_routing_decision( side_worker, RoutingDecisionHashes { local_hashes: side_hashes.clone(), diff --git a/lib/kv-router/src/services/indexer/recording.rs b/lib/kv-router/src/services/indexer/recording.rs index cb9508a91b74..dc80abcfcc0a 100644 --- a/lib/kv-router/src/services/indexer/recording.rs +++ b/lib/kv-router/src/services/indexer/recording.rs @@ -174,16 +174,6 @@ impl Indexer { .await } - /// Same as [`Self::record_routing_decision`]; kept for callers that name - /// the hashes explicitly. - pub async fn record_routing_decision_hashes( - &self, - worker: WorkerWithDpRank, - hashes: RoutingDecisionHashes, - ) -> Result<(), KvRouterError> { - self.record_routing_decision(worker, hashes).await - } - pub async fn record_hashed_routing_decision( &self, worker: WorkerWithDpRank, diff --git a/lib/kv-router/src/services/indexer/registry.rs b/lib/kv-router/src/services/indexer/registry.rs index 0a8dbb464d85..b1e29a454e6d 100644 --- a/lib/kv-router/src/services/indexer/registry.rs +++ b/lib/kv-router/src/services/indexer/registry.rs @@ -385,10 +385,6 @@ impl WorkerRegistry { *self.indexer_policy.lock() = policy; } - pub fn indexer_policy(&self) -> IndexerPolicy { - self.indexer_policy.lock().clone() - } - pub fn signal_ready(&self) { let _ = self.ready_tx.send(true); } @@ -451,6 +447,7 @@ impl WorkerRegistry { } } + let indexer_policy = self.indexer_policy.lock().clone(); let indexer_entry = self.indexers.entry(key.clone()).or_insert_with(|| { tracing::info!( model_name = %key.model_name, @@ -464,7 +461,7 @@ impl WorkerRegistry { block_size, self.num_threads, self.indexer_metrics.clone(), - &self.indexer_policy.lock(), + &indexer_policy, ), block_size, } @@ -785,6 +782,7 @@ impl WorkerRegistry { } pub fn get_or_create_indexer(&self, key: RoutingPartitionId, block_size: u32) -> Indexer { + let indexer_policy = self.indexer_policy.lock().clone(); let entry = self.indexers.entry(key.clone()).or_insert_with(|| { tracing::info!( model_name = %key.model_name, @@ -798,7 +796,7 @@ impl WorkerRegistry { block_size, self.num_threads, self.indexer_metrics.clone(), - &self.indexer_policy.lock(), + &indexer_policy, ), block_size, } diff --git a/lib/kv-router/src/services/selection/README.md b/lib/kv-router/src/services/selection/README.md index b6a4fbfa98eb..195159579f89 100644 --- a/lib/kv-router/src/services/selection/README.md +++ b/lib/kv-router/src/services/selection/README.md @@ -26,9 +26,16 @@ and active-sequence accounting. Keep these implementation invariants explicit: runtime event plane in the frontend) or `Remote` (a standalone indexer serves it). The `Indexer` type itself is shared with the frontend (`services::indexer::backend`). -- Each partition owns its `SessionAffinity` table, including versioned bindings, idle TTL, and lease lifecycle. A reservation owns its affinity lease, so every reservation removal releases both. Frontend routing hosts share the partition table and one coordinator for stream leases and runtime replication. -- Valid worker metadata updates preserve live bookings on surviving ranks and KV state from unchanged event sources. Catalog commits and ingress changes are serialized; partition policy factories can initialize independently. -- The frontend request lease manager owns expiry for embedded partitions. Standalone partitions use periodic request expiry. +- Each partition owns its `SessionAffinity` table, including versioned + bindings, idle TTL, and lease lifecycle. A reservation owns its affinity + lease, so every reservation removal releases both. Frontend routing hosts + share the partition table and one coordinator for stream leases and runtime + replication. +- Valid worker metadata updates preserve live bookings on surviving ranks and + KV state from unchanged event sources. Catalog commits and ingress changes + are serialized; partition policy factories can initialize independently. +- The frontend request lease manager owns expiry for embedded partitions. + Standalone partitions use periodic request expiry. - Selector replicas synchronize admission, prefill-complete, and free events. - **NOTE:** Output-block updates remain local. They are deliberately excluded from replica sync because their frequency would consume disproportionate diff --git a/lib/kv-router/src/services/selection/affinity.rs b/lib/kv-router/src/services/selection/affinity.rs index cde62dfa850b..3f72b6dab116 100644 --- a/lib/kv-router/src/services/selection/affinity.rs +++ b/lib/kv-router/src/services/selection/affinity.rs @@ -249,11 +249,11 @@ impl SessionAffinity { waiter_observed: Notify::new(), }); Self::spawn_reaper(&inner); - tracing::info!( + tracing::debug!( ttl_secs = config.ttl.as_secs(), mode = ?config.mode, max_entries = config.max_entries, - "session affinity enabled" + "Session affinity enabled" ); Ok(Self { inner }) } @@ -333,7 +333,7 @@ impl SessionAffinity { self.reserve_entry()?; tracing::debug!( session_id, - "session affinity miss: new session, pinning after worker selection" + "Session affinity miss: new session, pinning after worker selection" ); let revision = self.inner.next_revision.fetch_add(1, Ordering::Relaxed); let notify = Arc::new(Notify::new()); @@ -369,7 +369,7 @@ impl SessionAffinity { } if *active_leases == 0 && *idle_deadline <= now => { tracing::debug!( session_id, - "session affinity miss: pin expired (idle past TTL), re-selecting worker" + "Session affinity miss: pin expired (idle past TTL), re-selecting worker" ); let revision = self.inner.next_revision.fetch_add(1, Ordering::Relaxed); let notify = Arc::new(Notify::new()); @@ -401,7 +401,7 @@ impl SessionAffinity { worker_id = target.worker_id, dp_rank = ?target.dp_rank, active_leases = *active_leases + 1, - "session affinity hit: reusing pinned worker" + "Session affinity hit: reusing pinned worker" ); *active_leases += 1; Ok(AcquireStep::Held(Hold::Bound { @@ -496,7 +496,7 @@ impl SessionAffinity { session_id, worker_id = target.worker_id, dp_rank = ?target.dp_rank, - "session affinity hit: reusing pinned worker" + "Session affinity hit: reusing pinned worker" ); Ok(Some(*target)) } diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index f287c0859636..a5d788b23997 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -501,9 +501,17 @@ impl SelectionCore { /// Apply a session binding a replica published. pub(crate) fn dispatch_affinity_event(&self, event: AffinityBindingEvent) { let Some(entry) = self.entry(&event.partition) else { + tracing::trace!( + key = %event.partition, + "Dropping session affinity replica update for unknown selector entry" + ); return; }; let Some(table) = entry.affinity.get() else { + tracing::trace!( + key = %event.partition, + "Dropping session affinity replica update: no affinity table" + ); return; }; if self @@ -519,6 +527,11 @@ impl SelectionCore { .get(event.worker_id) .is_none_or(|record| record.key() != event.partition) { + tracing::trace!( + key = %event.partition, + worker_id = event.worker_id, + "Dropping session affinity replica update: worker not in partition" + ); return; } let (target, version, worker_id) = (event.target(), event.version(), event.worker_id); @@ -526,7 +539,7 @@ impl SelectionCore { tracing::trace!( worker_id, ?outcome, - "applied session affinity replica update" + "Applied session affinity replica update" ); } diff --git a/lib/kv-router/src/services/selection/core/run.rs b/lib/kv-router/src/services/selection/core/run.rs index 24c1ee1b028d..7b93bd59265a 100644 --- a/lib/kv-router/src/services/selection/core/run.rs +++ b/lib/kv-router/src/services/selection/core/run.rs @@ -8,6 +8,9 @@ use super::hint::transfer_hint_for_selection; use super::reservations::{Reservation, missing_booking}; use super::*; +/// Action id of the single `kv.fetch` action a selection's KV hint carries. +const KV_HINT_FETCH_ACTION_ID: &str = "a1"; + pub(super) struct PreparedSelectionInputs { pub(super) block_hashes: Vec, pub(super) sequence_hashes: Vec, @@ -443,7 +446,7 @@ impl SelectionCore { .map(|payload| { KvHint::new( admission.request_id().unwrap_or_default(), - vec![KvHintAction::fetch("a1", payload)], + vec![KvHintAction::fetch(KV_HINT_FETCH_ACTION_ID, payload)], ) }) } else { @@ -542,7 +545,7 @@ impl SelectionCore { tracing::debug!( session_id, worker_id = target.worker_id, - "session affinity target is not schedulable; re-initializing" + "Session affinity target is not schedulable; re-initializing" ); lease.invalidate(); } @@ -550,7 +553,7 @@ impl SelectionCore { Err(AffinityError::ResourceExhausted(_)) => { tracing::debug!( session_id, - "affinity table full; routing without session affinity" + "Affinity table full; routing without session affinity" ); return Ok(None); } @@ -657,8 +660,8 @@ impl SelectionCore { } else { entry .indexer - .find_tiered_matches_with_options( - normalized.block_hashes.clone(), + .find_tiered_matches_ref_with_options( + &normalized.block_hashes, LowerTierQueryOptions { retain_kv_transfer_chain, }, diff --git a/lib/kv-router/src/services/selection/core/tests.rs b/lib/kv-router/src/services/selection/core/tests.rs index 93e559f81e6a..8d3c568695a4 100644 --- a/lib/kv-router/src/services/selection/core/tests.rs +++ b/lib/kv-router/src/services/selection/core/tests.rs @@ -2508,13 +2508,13 @@ async fn metadata_update_preserves_live_booking() { result = core.upsert_worker(updated) => { result.unwrap(); }, _ = ingress.0.notified() => { // A capacity update must not temporarily withdraw this worker. - assert!(entry.scheduler.has_request("live")); panic!("capacity update detached the worker"); } } assert!(entry.scheduler.has_request("live")); core.free_reservation("live").await.unwrap(); } + #[tokio::test(start_paused = true)] async fn host_lease_manager_owns_expiry() { use crate::scheduling::queue::SchedulerBookingDescriptor; diff --git a/lib/kv-router/src/services/selection/mod.rs b/lib/kv-router/src/services/selection/mod.rs index dbf0a2cbb841..49f8f754b15c 100644 --- a/lib/kv-router/src/services/selection/mod.rs +++ b/lib/kv-router/src/services/selection/mod.rs @@ -49,9 +49,10 @@ pub use service::{ SelectionService, SelectionServiceBuilder, warn_for_unserved_worker_selection_policies, }; pub use types::{ - ModelLoadResponse, OutputBlockRequest, OverlapScoresRequest, OverlapScoresResponse, - PotentialLoadsRequest, ReadyResponse, ReservationRequest, ReservationResponse, - SelectAndReserveRequest, SelectRequest, SelectResponse, SelectionInputTrigger, - SelectionSessionContext, SelectionWorkerConfig, SelectionWorkerLoad, SharedCacheOverlapScore, - WorkerCatalogRecord, WorkerLifecycle, WorkerOverlapScore, WorkerPatchRequest, WorkerRequest, + DEFAULT_MODEL_NAME, ModelLoadResponse, OutputBlockRequest, OverlapScoresRequest, + OverlapScoresResponse, PotentialLoadsRequest, ReadyResponse, ReservationRequest, + ReservationResponse, SelectAndReserveRequest, SelectRequest, SelectResponse, + SelectionInputTrigger, SelectionSessionContext, SelectionWorkerConfig, SelectionWorkerLoad, + SharedCacheOverlapScore, WorkerCatalogRecord, WorkerLifecycle, WorkerOverlapScore, + WorkerPatchRequest, WorkerRequest, }; diff --git a/lib/kv-router/src/services/selection/service.rs b/lib/kv-router/src/services/selection/service.rs index f9790b52fafc..280016a2b91c 100644 --- a/lib/kv-router/src/services/selection/service.rs +++ b/lib/kv-router/src/services/selection/service.rs @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 use std::sync::Arc; +use std::time::Duration; use tokio_util::sync::CancellationToken; @@ -41,7 +42,7 @@ pub struct SelectionServiceBuilder { worker_selection_policy_registry: WorkerSelectionPolicyRegistry, host: SelectionHost, worker_selection_policy_factory: Option, - session_affinity_ttl: Option, + session_affinity_ttl: Option, } /// Warn when a host does not construct workers for explicitly configured policy roles. @@ -120,7 +121,7 @@ impl SelectionServiceBuilder { /// Pin each session id to the worker that served it for `ttl` after its /// last request. Bindings replicate over the replica mesh when enabled. - pub fn session_affinity(mut self, ttl: std::time::Duration) -> Self { + pub fn session_affinity(mut self, ttl: Duration) -> Self { self.session_affinity_ttl = Some(ttl); self } diff --git a/lib/kv-router/src/services/selection/types.rs b/lib/kv-router/src/services/selection/types.rs index 712ffcd5d0b9..9b77f3dfd532 100644 --- a/lib/kv-router/src/services/selection/types.rs +++ b/lib/kv-router/src/services/selection/types.rs @@ -18,7 +18,7 @@ use crate::services::overlap::MooncakeOverlapSummary; use super::input::PromptRequest; -const DEFAULT_MODEL_NAME: &str = "default"; +pub const DEFAULT_MODEL_NAME: &str = "default"; pub(super) const REQUEST_BODY_LIMIT_BYTES: usize = 8 * 1024 * 1024; fn default_model_name() -> String { @@ -441,9 +441,11 @@ pub struct SelectRequest { /// /// The response then carries the chosen worker's `worker_load` snapshot /// and `prefill_busy` evaluation. The request never waits in the router - /// queue; when the queue would have rejected it the call fails the same - /// way an admitted selection does. Ignored on `select_and_reserve`, which - /// always books. + /// queue and is not subject to its admission checks, so an advisory + /// selection can succeed where an admitted one would have been rejected. + /// A `selection_id` still caches the booking inputs for a follow-up + /// `create_reservation`. Ignored on `select_and_reserve`, which always + /// books. #[serde(default)] pub advisory: bool, } diff --git a/lib/llm/src/discovery/watcher.rs b/lib/llm/src/discovery/watcher.rs index 0fb628a18b3b..feb987cc2f67 100644 --- a/lib/llm/src/discovery/watcher.rs +++ b/lib/llm/src/discovery/watcher.rs @@ -285,9 +285,7 @@ impl ModelWatcher { SelectionPolicySource::Registry, ) } -} -impl ModelWatcher { #[allow(clippy::too_many_arguments)] pub(crate) fn new_with_selection_policy( runtime: DistributedRuntime, diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index d2f7d1af22ba..1097be333db4 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -137,7 +137,7 @@ pub(crate) fn policy_worker_inputs( model_name: Option<&str>, ) -> WorkerInputs { let partition = RoutingPartitionRef::new( - model_name.unwrap_or(embedded::DEFAULT_MODEL_NAME), + model_name.unwrap_or(dynamo_kv_router::services::selection::DEFAULT_MODEL_NAME), DEFAULT_ROUTING_GROUP, ); dynamo_kv_router::selector::WorkerSelector::::required_worker_inputs( @@ -1087,9 +1087,7 @@ impl KvRouter { hashes: RoutingDecisionHashes, worker: WorkerWithDpRank, ) -> Result<(), KvRouterError> { - self.indexer - .record_routing_decision_hashes(worker, hashes) - .await + self.indexer.record_routing_decision(worker, hashes).await } /// Give these tokens, find the worker with the best weighted cache hit. @@ -1594,7 +1592,6 @@ impl KvRouter { self.scheduler.free(request_id).await } - #[doc(hidden)] pub(crate) fn affinity_coordinator( &self, ttl: std::time::Duration, @@ -2047,6 +2044,13 @@ mod tests { .expect("filtered workers should produce a DynamoError"); assert_eq!(dynamo_error.error_type(), ErrorType::Unavailable); + + let error = map_scheduler_error(KvSchedulerError::AllEligibleWorkersOverloaded); + let dynamo_error = error + .downcast_ref::() + .expect("overloaded workers should produce a DynamoError"); + + assert_eq!(dynamo_error.error_type(), ErrorType::ResourceExhausted); } #[test] diff --git a/lib/llm/src/kv_router/AGENTS.md b/lib/llm/src/kv_router/AGENTS.md index bbecd4aab726..2f860e620198 100644 --- a/lib/llm/src/kv_router/AGENTS.md +++ b/lib/llm/src/kv_router/AGENTS.md @@ -3,7 +3,7 @@ ## Module map - `kv_router.rs` estimates how much of a request each worker already has cached, then asks the scheduler to choose a worker. It selects a worker but does not send the request. -- `scheduler.rs` connects discovered workers and their current load to the scheduling code in `lib/kv-router`. +- `embedded.rs` connects discovered workers and their current load to the scheduling code in `lib/kv-router`; `scheduler.rs` only re-exports from it. - `indexer/` and `route_lookup.rs` track which KV-cache blocks each worker holds and look for reusable blocks before selection. - `publisher/` receives KV-cache events and worker metrics from inference engines. - `routing_host.rs` owns selection, dispatch, and response cleanup. `routing_host/kv_selection.rs` chooses KV-routed workers, `routing_host/request_guard.rs` tracks progress and cleanup, and `routing_host/cancellation.rs` stops unfinished work when the client cancels. diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs index a3d1ec74875f..9e5965705972 100644 --- a/lib/llm/src/kv_router/embedded.rs +++ b/lib/llm/src/kv_router/embedded.rs @@ -9,7 +9,9 @@ //! frontend retains transport, stream leases, and request-expiry ownership. use std::collections::HashMap; -use std::sync::Arc; +use std::ops::Range; +use std::sync::{Arc, OnceLock}; +use std::time::Duration; use anyhow::{Context, Result}; use dynamo_kv_router::WorkerType; @@ -26,10 +28,11 @@ use dynamo_kv_router::scheduling::{ }; use dynamo_kv_router::sequences::{SequenceError, SequenceRequest}; use dynamo_kv_router::services::selection::{ - CatalogObserver, CatalogReconciler, HostCache, HostEligibility, HostLoad, HostReplication, - HostTelemetry, KvEventIngress, KvIndexSource, SelectionHost, SelectionOperation, - SelectionOutcome, SelectionPartition, SelectionRun, SelectionService, SelectionServiceBuilder, - WorkerCatalogRecord, WorkerCatalogSource, WorkerRequest, WorkerSelectionPolicyRegistry, + CatalogObserver, CatalogReconciler, DEFAULT_MODEL_NAME, HostCache, HostEligibility, HostLoad, + HostReplication, HostTelemetry, KvEventIngress, KvIndexSource, SelectionHost, + SelectionOperation, SelectionOutcome, SelectionPartition, SelectionRun, SelectionService, + SelectionServiceBuilder, WorkerCatalogRecord, WorkerCatalogSource, WorkerRequest, + WorkerSelectionPolicyRegistry, }; use dynamo_kv_router::{DEFAULT_ROUTING_GROUP, PrefillLoadEstimator, WorkerSelectionPolicyFactory}; use dynamo_tokens::SequenceHash; @@ -118,7 +121,7 @@ fn spawn_queue_metrics_updater( ) { let mut queue_updates = partition.scheduler().subscribe_queue_updates(); tokio::spawn(async move { - let period = std::time::Duration::from_secs(60); + let period = Duration::from_secs(60); let mut recheck = tokio::time::interval_at(tokio::time::Instant::now() + period, period); loop { update_queue_metrics(&handles, |index| { @@ -158,16 +161,13 @@ fn non_max_overlap_observer(worker_type: &'static str) -> NonMaxOverlapSelection }) } -/// Partition model name when the router has none. -pub(crate) const DEFAULT_MODEL_NAME: &str = "default"; - /// One `SelectionService` partition driven directly by the router. pub(crate) struct EmbeddedSelection { /// Keeps the service (listeners, sweep, replica sync) alive for as long as /// the router holds the partition. service: Arc, partition: SelectionPartition, - affinity: std::sync::OnceLock, + affinity: OnceLock, worker_type: &'static str, /// Queue gauges and rejection counters per policy class, index-aligned with /// the scheduler's `class_queue_stats`. @@ -193,8 +193,7 @@ impl PendingReplicaIngress { } } -static INSTALLED_POLICY_REGISTRY: std::sync::OnceLock = - std::sync::OnceLock::new(); +static INSTALLED_POLICY_REGISTRY: OnceLock = OnceLock::new(); /// Install the process-wide worker-selection policy registry (linked custom /// policies) that embedded selection partitions resolve `KvRouterConfig` @@ -268,10 +267,12 @@ impl EmbeddedSelection { slot.lock().ok().and_then(|mut s| s.take()) })); + // The registry is only consulted when no factory is set; `policy_factory` + // always is, so the builder never reads it. let service = SelectionServiceBuilder::new( args.kv_router_config.clone(), worker_type, - worker_selection_policy_registry(), + WorkerSelectionPolicyRegistry::default(), ) .worker_selection_policy_factory(args.policy_factory) .indexer_threads(1) @@ -371,7 +372,7 @@ impl EmbeddedSelection { Self { service, partition, - affinity: std::sync::OnceLock::new(), + affinity: OnceLock::new(), worker_type: args.metric_worker_type, queue_metrics, queue_metric_indices, @@ -394,7 +395,7 @@ impl EmbeddedSelection { pub(crate) fn affinity_coordinator( &self, - ttl: std::time::Duration, + ttl: Duration, mode: crate::session_affinity::SessionAffinityMode, ) -> Result { let table = self.partition.session_affinity( @@ -552,9 +553,9 @@ struct RegisteredGauge { worker_label: &'static str, } -fn dp_ranks(record: &WorkerCatalogRecord) -> std::ops::Range { +fn dp_ranks(record: &WorkerCatalogRecord) -> Range { let start = record.data_parallel_start_rank.unwrap_or(0); - start..start.saturating_add(record.data_parallel_size.unwrap_or(1).max(1)) + start..start.saturating_add(record.data_parallel_size.unwrap_or(1)) } impl CatalogObserver for RegisteredGauge { @@ -600,6 +601,9 @@ pub(crate) fn worker_request_from_runtime_config( model_name: key.model_name.clone(), routing_group: key.routing_group.clone(), endpoint: Some(format!("dyn://{worker_id}")), + kv_events_endpoint: None, + kv_events_endpoints: HashMap::new(), + replay_endpoint: None, block_size: Some(block_size), data_parallel_start_rank: Some(dp_start), data_parallel_size: Some(dp_size), @@ -621,7 +625,6 @@ pub(crate) fn worker_request_from_runtime_config( router_hint_worker_type, router_hint_source_control_endpoints, kv_event_source_mode: config.kv_event_source_mode.clone(), - ..WorkerRequest::default() } } diff --git a/lib/llm/src/kv_router/indexer/ingress.rs b/lib/llm/src/kv_router/indexer/ingress.rs index 6abdd5aa4fd4..bfe40424cf9a 100644 --- a/lib/llm/src/kv_router/indexer/ingress.rs +++ b/lib/llm/src/kv_router/indexer/ingress.rs @@ -88,10 +88,10 @@ impl RuntimeIngress { tracing::info!( requirement = %args.kv_event_source_requirement, cache_required = args.cache_required, - "Skipping KV event subscription (use_kv_events={}, overlap_score_credit={}, use_remote_indexer={})", - args.kv_router_config.use_kv_events, - args.kv_router_config.overlap_score_credit, - args.kv_router_config.use_remote_indexer, + use_kv_events = args.kv_router_config.use_kv_events, + overlap_score_credit = args.kv_router_config.overlap_score_credit, + use_remote_indexer = args.kv_router_config.use_remote_indexer, + "Skipping KV event subscription" ); None }; diff --git a/lib/llm/src/kv_router/indexer/remote.rs b/lib/llm/src/kv_router/indexer/remote.rs index 9a6475c4b6c4..07376493a8a5 100644 --- a/lib/llm/src/kv_router/indexer/remote.rs +++ b/lib/llm/src/kv_router/indexer/remote.rs @@ -85,7 +85,7 @@ impl RemoteIndexer { }) } - pub(super) async fn find_matches_by_tier( + pub(super) async fn fetch_matches_by_tier( &self, block_hashes: Vec, device_only: bool, @@ -207,7 +207,7 @@ impl dynamo_kv_router::services::indexer::backend::RemotePrimary for RemoteIndex block_hashes: Vec, device_only: bool, ) -> Result { - RemoteIndexer::find_matches_by_tier(self, block_hashes, device_only).await + self.fetch_matches_by_tier(block_hashes, device_only).await } async fn record_routing_decision( diff --git a/lib/llm/src/kv_router/sequence.rs b/lib/llm/src/kv_router/sequence.rs index f8a260cc1a05..75b8b727f9cb 100644 --- a/lib/llm/src/kv_router/sequence.rs +++ b/lib/llm/src/kv_router/sequence.rs @@ -198,19 +198,6 @@ impl ActiveSequenceEventPublisher { } } -#[cfg(test)] -fn active_sequence_event_channel( - enabled: bool, - capacity: usize, - cancellation_token: &CancellationToken, -) -> Option<( - ActiveSequenceEventPublisher, - mpsc::Receiver, -)> { - enabled - .then(|| ActiveSequenceEventPublisher::channel(capacity, cancellation_token.child_token())) -} - /// One event per message, for transports that carry singletons. trait SingletonEventPublisher: Send + Sync { fn publish_event( @@ -658,11 +645,6 @@ mod tests { assert_eq!(event_rx.try_recv().unwrap().request_id, "accepted"); } - #[test] - fn active_sequence_publish_channel_is_absent_when_replica_sync_disabled() { - assert!(active_sequence_event_channel(false, 1, &CancellationToken::new()).is_none()); - } - #[test] fn active_sequence_publish_sender_classifies_closed_queue_by_cancellation() { let cancellation_token = CancellationToken::new(); diff --git a/lib/llm/src/lora/load_estimator.rs b/lib/llm/src/lora/load_estimator.rs index 8ca17c011b64..a136edd65e41 100644 --- a/lib/llm/src/lora/load_estimator.rs +++ b/lib/llm/src/lora/load_estimator.rs @@ -595,48 +595,6 @@ impl LoadEstimator { } } - /// Update active counts from a snapshot; arrivals are approximated as the - /// positive delta since the previous snapshot. - #[cfg(test)] - fn update_from_counts(&self, lora_counts: HashMap) { - let now = Instant::now(); - let cfg = self.config.read(); - let num_buckets = cfg.num_buckets(); - let bucket_duration = cfg.bucket_duration(); - drop(cfg); - - for (lora_name, count) in &lora_counts { - self.data - .entry(lora_name.clone()) - .and_modify(|data| { - let prev = data.active_count.load(Ordering::Relaxed); - data.active_count.store(*count, Ordering::Relaxed); - // Record only the delta (new arrivals since the last poll) to - // avoid double-counting sustained requests. A request active - // for N ticks should contribute 1 arrival, not N. - let arrivals = count.saturating_sub(prev) as u64; - if arrivals > 0 { - data.rate_counter.record_count(arrivals, now); - } - }) - .or_insert_with(|| { - let counter = BucketedRateCounter::new(num_buckets, bucket_duration, now); - // First observation: the entire count represents new arrivals. - counter.record_count(*count as u64, now); - LoraLoadData { - active_count: AtomicUsize::new(*count), - rate_counter: counter, - } - }); - } - - for entry in self.data.iter() { - if !lora_counts.contains_key(entry.key()) { - entry.value().active_count.store(0, Ordering::Relaxed); - } - } - } - /// Get current load using arrival count in the sliding window. pub fn get_current_load(&self) -> HashMap { self.get_current_load_at(Instant::now()) @@ -755,6 +713,49 @@ mod tests { use crate::kv_router::ACTIVE_SEQUENCES_SUBJECT; + impl LoadEstimator { + /// Update active counts from a snapshot; arrivals are approximated as the + /// positive delta since the previous snapshot. + fn update_from_counts(&self, lora_counts: HashMap) { + let now = Instant::now(); + let cfg = self.config.read(); + let num_buckets = cfg.num_buckets(); + let bucket_duration = cfg.bucket_duration(); + drop(cfg); + + for (lora_name, count) in &lora_counts { + self.data + .entry(lora_name.clone()) + .and_modify(|data| { + let prev = data.active_count.load(Ordering::Relaxed); + data.active_count.store(*count, Ordering::Relaxed); + // Record only the delta (new arrivals since the last poll) to + // avoid double-counting sustained requests. A request active + // for N ticks should contribute 1 arrival, not N. + let arrivals = count.saturating_sub(prev) as u64; + if arrivals > 0 { + data.rate_counter.record_count(arrivals, now); + } + }) + .or_insert_with(|| { + let counter = BucketedRateCounter::new(num_buckets, bucket_duration, now); + // First observation: the entire count represents new arrivals. + counter.record_count(*count as u64, now); + LoraLoadData { + active_count: AtomicUsize::new(*count), + rate_counter: counter, + } + }); + } + + for entry in self.data.iter() { + if !lora_counts.contains_key(entry.key()) { + entry.value().active_count.store(0, Ordering::Relaxed); + } + } + } + } + fn lora_event( request_id: &str, lora_name: Option<&str>, diff --git a/lib/llm/src/protocols/agents.rs b/lib/llm/src/protocols/agents.rs index 9ee27c87ceeb..998ca96dfbdf 100644 --- a/lib/llm/src/protocols/agents.rs +++ b/lib/llm/src/protocols/agents.rs @@ -16,7 +16,7 @@ pub(crate) const HEADER_CODEX_PARENT_THREAD_ID: &str = "x-codex-parent-thread-id pub(crate) const HEADER_CODEX_TURN_METADATA: &str = "x-codex-turn-metadata"; pub(crate) const HEADER_OPENCODE_SESSION_ID: &str = "x-session-id"; pub(crate) const HEADER_OPENCODE_PARENT_SESSION_ID: &str = "x-parent-session-id"; -pub(crate) const HEADER_DYNAMO_SESSION_ID: &str = "x-dynamo-session-id"; +pub const HEADER_DYNAMO_SESSION_ID: &str = "x-dynamo-session-id"; pub(crate) const HEADER_DYNAMO_PARENT_SESSION_ID: &str = "x-dynamo-parent-session-id"; pub(crate) const HEADER_DYNAMO_SESSION_FINAL: &str = "x-dynamo-session-final"; diff --git a/lib/llm/src/session_affinity/coordinator.rs b/lib/llm/src/session_affinity/coordinator.rs index c18574a4678b..675bd34bb401 100644 --- a/lib/llm/src/session_affinity/coordinator.rs +++ b/lib/llm/src/session_affinity/coordinator.rs @@ -58,14 +58,6 @@ struct Inner { replica: tokio::sync::OnceCell, } -impl Drop for Inner { - fn drop(&mut self) { - if let Some(replica) = self.replica.get_mut() { - replica.shutdown_now(); - } - } -} - #[derive(Clone)] pub struct AffinityCoordinator { inner: Arc, @@ -103,7 +95,7 @@ impl AffinityCoordinator { .enable_replication(router_id, replica.sink()) { return Err(anyhow::anyhow!( - "session affinity replica sync already enabled" + "session affinity table already has a replica sink installed" )); } Ok(replica) From ea037f3b854f43753057bbf0459c2748b22e8265 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Fri, 11 Sep 2026 16:09:11 -0700 Subject: [PATCH 064/102] test(router): pin the booking lifecycle invariants State the five invariants a booking holds through both booking paths (`run_selection_inner` and `finalize_reservation`) as a module doc on `core/run.rs`, and add the tests that were missing for them: - `lease_admission_installs_no_index_row_and_records_nothing`: a `Lease` admission leaves the reservation index empty, returns no booked hashes, and records nothing into an approximate primary (checked after a real indexer flush). - `dropped_book_selection_during_routing_record_frees_booking_and_claim`: dropping a `Book` run while its routing decision is being recorded frees the booking and releases the selection-id claim. - `enroll_public_request_attempt_cancelled_during_routing_update_frees_booking` (frontend): a public enrollment cancelled during the routing update frees the booking through the detached enrollment's drop. Each test fails when the guarded code is broken (claim taken for `Lease`, install moved ahead of the record await, detached enrollment left disarmed). `Indexer::flush` is exposed to the selection core's tests for the first check. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Ishan Dhanani --- lib/kv-router/src/services/indexer/backend.rs | 8 + .../src/services/selection/core/run.rs | 32 +++ .../src/services/selection/core/tests.rs | 190 ++++++++++++++++++ lib/llm/src/kv_router.rs | 100 +++++++++ 4 files changed, 330 insertions(+) diff --git a/lib/kv-router/src/services/indexer/backend.rs b/lib/kv-router/src/services/indexer/backend.rs index 296c8be0b609..22fad6388a42 100644 --- a/lib/kv-router/src/services/indexer/backend.rs +++ b/lib/kv-router/src/services/indexer/backend.rs @@ -1309,6 +1309,14 @@ mod tests { } } + impl Indexer { + /// Test-only: wait until every locally enqueued event and routing + /// decision is applied. + pub(crate) async fn flush(&self) { + flush(self).await + } + } + #[tokio::test] async fn event_driven_policy_ignores_routing_decisions() { let indexer = policy_indexer(1, IndexerPolicy::event_driven()); diff --git a/lib/kv-router/src/services/selection/core/run.rs b/lib/kv-router/src/services/selection/core/run.rs index 7b93bd59265a..8b9ee42a5238 100644 --- a/lib/kv-router/src/services/selection/core/run.rs +++ b/lib/kv-router/src/services/selection/core/run.rs @@ -3,6 +3,38 @@ //! One selection, end to end: the wire adapters, `run_selection`, the //! lookups that feed the scheduler, and the session binding around it. +//! +//! # Booking lifecycle invariants +//! +//! A booking is `(request_id, worker, attempt_id)` in the partition scheduler. +//! Both booking paths, `run_selection_inner` here and `finalize_reservation` +//! in `core/reservations.rs`, hold it through these invariants; the named +//! tests in `core/tests.rs` pin each one. +//! +//! - A booking is released exactly once, by whichever owner holds it when the +//! request ends (`dropped_selection_future_frees_its_booking`, +//! `reservation_index_tracks_bookings_until_freed`). +//! - Release ownership moves in one direction: the scheduler actor books, hands +//! the armed lease to this run, and the run hands it to the index row for +//! `Book` (`claim.install`) or to the host for `Lease` (`Selected::lease`) +//! (`dropped_selection_future_frees_its_booking`, +//! `lease_admission_installs_no_index_row_and_records_nothing`). +//! - The session hold is taken before the booking is made +//! (`session_worker_departing_after_the_hold_reinitializes_the_session`) and +//! the binding is committed while the booking is held. The commit never +//! waits: `commit_session` is a synchronous fn that re-binds only through +//! `try_acquire`, so when the table cannot bind at commit time (`try_acquire` +//! returns `Wait`, or the table is full) it returns `Ok(None)` and the +//! request routes unpinned rather than waiting on a request queued behind +//! this booking. A rejected commit frees the booking +//! (`two_phase_replay_rejects_a_worker_the_session_left`). +//! - A `Lease` admission installs no index row and records no routing hashes +//! in the core; the host owns both +//! (`lease_admission_installs_no_index_row_and_records_nothing`). +//! - Dropping the run at any await after the booking lands and before the +//! index row is installed frees the booking and releases the id claim +//! (`dropped_selection_future_frees_its_booking`, +//! `dropped_book_selection_during_routing_record_frees_booking_and_claim`). use super::hint::transfer_hint_for_selection; use super::reservations::{Reservation, missing_booking}; diff --git a/lib/kv-router/src/services/selection/core/tests.rs b/lib/kv-router/src/services/selection/core/tests.rs index 8d3c568695a4..f74600679177 100644 --- a/lib/kv-router/src/services/selection/core/tests.rs +++ b/lib/kv-router/src/services/selection/core/tests.rs @@ -1280,6 +1280,76 @@ async fn lease_admission_keeps_a_selection_whose_worker_drained_while_queued() { wait_until("booking release", || !entry.scheduler.has_request("queued")).await; } +/// A `Lease` admission hands the booking to the host: no index row, no +/// booked hashes, and nothing recorded into the approximate primary that +/// `Book` feeds. +#[tokio::test] +async fn lease_admission_installs_no_index_row_and_records_nothing() { + let core = SelectionCore::try_new_local( + test_config(false), + 1, + CancellationToken::new(), + SelectionCacheConfig::default(), + ) + .expect("valid test config"); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + let key = RoutingPartitionId::new("model", "default"); + let entry = core.entry(&key).expect("entry"); + + let req = reserve_request("leased"); + let run = core + .run_selection(SelectionOperation { + key: key.clone(), + prompt: req.prompt.view(), + router_config_override: None, + expected_output_tokens: None, + priority_jump: 0.0, + strict_priority: 0, + policy_class: None, + session_context: None, + session: SessionBinding::None, + affinity_target: None, + pinned_worker: None, + allowed_worker_ids: None, + routing_constraints: RoutingConstraints::default(), + admission: SelectionAdmission::Lease { + request_id: "leased".to_string(), + }, + track_active_blocks: true, + return_routing_hashes: false, + replay_id: None, + }) + .await; + let Ok(SelectionOutcome::Selected(selected)) = run.result else { + panic!("lease selection failed"); + }; + assert!(entry.scheduler.has_request("leased")); + assert!( + core.reservation_index.read().is_empty(), + "a lease admission installs no index row" + ); + assert!( + selected.lease.is_some(), + "the booking's lease goes to the host" + ); + assert!(selected.sequence_hashes.is_none()); + assert!(selected.routing_hashes.is_none()); + // Nothing was recorded for the leased prompt: once the indexer has + // applied everything enqueued so far, a query restricted to the booked + // worker sees no cached blocks. + entry.indexer.flush().await; + let mut probe = select_request(); + probe.allowed_worker_ids = Some(HashSet::from([1])); + let probe = core.select(probe).await.expect("select"); + assert_eq!( + probe.overlap.longest_matched, 0, + "a lease admission records nothing into the indexer" + ); + + drop(selected); + wait_until("booking release", || !entry.scheduler.has_request("leased")).await; +} + #[tokio::test] async fn dropped_selection_future_frees_its_booking() { let core = SelectionCore::try_new_local( @@ -1324,6 +1394,126 @@ async fn dropped_selection_future_frees_its_booking() { assert!(core.reservation_index.read().get("dropped").is_none()); } +/// The routing-decision record is an await between the booking (and the +/// session commit) and the index install; the armed lease, not lock +/// discipline, covers it. Dropping the run there frees the booking and +/// releases the id claim. +#[tokio::test] +async fn dropped_book_selection_during_routing_record_frees_booking_and_claim() { + use crate::indexer::TieredMatchDetails; + use crate::services::indexer::backend::RemotePrimary; + struct PausedRecord { + entered: AtomicBool, + release: tokio::sync::Notify, + } + #[async_trait::async_trait] + impl RemotePrimary for PausedRecord { + async fn find_matches_by_tier( + &self, + _: Vec, + _: bool, + ) -> anyhow::Result { + Ok(TieredMatchDetails::default()) + } + async fn record_routing_decision( + &self, + _: WorkerWithDpRank, + _: RoutingDecisionHashes, + ) -> anyhow::Result<()> { + self.entered.store(true, Ordering::Release); + self.release.notified().await; + Ok(()) + } + fn use_kv_events(&self) -> bool { + false + } + } + struct PausedRecordIngress(Arc); + #[async_trait::async_trait] + impl KvEventIngress for PausedRecordIngress { + fn open(&self, _: &WorkerRegistry, _: &RoutingPartitionId, _: u32) -> Indexer { + Indexer::Remote { + primary: self.0.clone(), + approx: None, + primary_records_routing_decisions: true, + } + } + } + let record = Arc::new(PausedRecord { + entered: AtomicBool::new(false), + release: tokio::sync::Notify::new(), + }); + let config = test_config(false); + let tracking_hash = Arc::new( + TrackingHashContext::from_config(&config).expect("valid tracking hash configuration"), + ); + let indexer_policy = IndexerPolicy::from_router_config(&config).expect("indexer policy"); + let core = SelectionCore::new_inner( + config, + 1, + CancellationToken::new(), + None, + None, + SelectionHost { + cache: HostCache { + index: KvIndexSource::Owned(Arc::new(PausedRecordIngress(record.clone()))), + shared: None, + }, + ..SelectionHost::default() + }, + WorkerType::Aggregated, + true, + SelectionCacheConfig::default(), + tracking_hash, + indexer_policy, + Some( + SessionAffinityConfig::new(Duration::from_secs(10)) + .with_mode(SessionAffinityMode::Hard), + ), + ); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + let key = RoutingPartitionId::new("model", "default"); + let entry = core.entry(&key).expect("entry"); + + // Poll by hand until the booking exists and the record is awaited. + let mut selection = Box::pin(core.select_and_reserve(session_reservation("recording", "s"))); + let mut context = std::task::Context::from_waker(std::task::Waker::noop()); + wait_until("routing record entered", || { + if record.entered.load(Ordering::Acquire) { + return true; + } + assert!( + selection.as_mut().poll(&mut context).is_pending(), + "selection completed before the routing record was entered" + ); + false + }) + .await; + assert!(entry.scheduler.has_request("recording")); + assert!( + core.reservation_index + .read() + .get("recording") + .is_some_and(|reservation| reservation.booking.is_none()), + "the id is still a claim while the record is in flight" + ); + assert_eq!( + bound_worker(&core, "s"), + Some(1), + "the session commits before the record" + ); + drop(selection); + + wait_until("booking release", || { + !entry.scheduler.has_request("recording") + }) + .await; + assert!( + core.reservation_index.read().get("recording").is_none(), + "the claim is released with the booking" + ); +} + #[tokio::test] async fn same_selection_id_in_two_partitions_is_a_conflict() { let core = SelectionCore::try_new_local( diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index 1097be333db4..b6f8fd5b89a1 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -2692,6 +2692,106 @@ mod tests { assert!(router.scheduler.has_request("kept")); } + /// A public enrollment cancelled while its routing update is in flight + /// frees the booking through the detached enrollment's drop. + #[tokio::test] + async fn enroll_public_request_attempt_cancelled_during_routing_update_frees_booking() { + use std::sync::atomic::{AtomicBool, Ordering}; + + use dynamo_kv_router::indexer::TieredMatchDetails; + use dynamo_kv_router::protocols::LocalBlockHash; + use dynamo_kv_router::services::indexer::backend::RemotePrimary; + + /// Records park until `release` is notified. + struct PausedRecord { + entered: AtomicBool, + release: tokio::sync::Notify, + } + #[async_trait] + impl RemotePrimary for PausedRecord { + async fn find_matches_by_tier( + &self, + _: Vec, + _: bool, + ) -> anyhow::Result { + Ok(TieredMatchDetails::default()) + } + async fn record_routing_decision( + &self, + _: WorkerWithDpRank, + _: RoutingDecisionHashes, + ) -> anyhow::Result<()> { + self.entered.store(true, Ordering::Release); + self.release.notified().await; + Ok(()) + } + fn use_kv_events(&self) -> bool { + false + } + } + + let mut router = differential_router("enroll").await; + let record = Arc::new(PausedRecord { + entered: AtomicBool::new(false), + release: tokio::sync::Notify::new(), + }); + router.indexer = Indexer::Remote { + primary: record.clone(), + approx: None, + primary_records_routing_decisions: true, + }; + let router = router; + let prompt = golden_prefix(1); + let outcome = router + .find_best_match_details_with_policy_class_inner( + Some("cancelled"), + &prompt, + None, + None, + true, + false, + None, + None, + 0.0, + 0, + None, + None, + None, + None, + None, + None, + RoutingConstraints::default(), + FindBestMatchAdmission::WithAdmission, + ) + .await + .unwrap(); + let FindBestMatchInnerOutcome::WithAdmission(admitted) = outcome else { + panic!("admitted routing returned advisory outcome"); + }; + let lease = admitted + .lease + .expect("tracked selection carries its booking lease"); + + // Park the routing update after `register_detached` has run. + let mut enroll = Box::pin( + router.enroll_public_request_attempt(lease, Some(TokensWithHashes::new(prompt, 2))), + ); + assert!(futures::poll!(&mut enroll).is_pending()); + assert!( + record.entered.load(Ordering::Acquire), + "routing update is in flight" + ); + assert!(router.scheduler.has_request("cancelled")); + drop(enroll); + tokio::time::timeout(std::time::Duration::from_secs(2), async { + while router.scheduler.has_request("cancelled") { + tokio::task::yield_now().await; + } + }) + .await + .expect("cancelled enrollment released its booking"); + } + async fn make_test_router_with_workers( policy: SelectionPolicySource, shared_cache: Option>, From c95307e99eff684a3df3e99307ca93259e8cf281 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Fri, 11 Sep 2026 16:34:49 -0700 Subject: [PATCH 065/102] fix(router): gate the LoRA narrowing helper behind standalone-selection The helper became crate-private in 317eb30e86, but its only caller is the selection core, which only builds with the `standalone-selection` feature. Without the feature the re-export and the function were dead code and the default-feature clippy run failed. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Ishan Dhanani --- lib/kv-router/src/scheduling/lora_filter.rs | 1 + lib/kv-router/src/scheduling/mod.rs | 1 + 2 files changed, 2 insertions(+) diff --git a/lib/kv-router/src/scheduling/lora_filter.rs b/lib/kv-router/src/scheduling/lora_filter.rs index 69f6c483d872..32416c3de2bf 100644 --- a/lib/kv-router/src/scheduling/lora_filter.rs +++ b/lib/kv-router/src/scheduling/lora_filter.rs @@ -27,6 +27,7 @@ pub trait LoraWorkerFilter: Send + Sync { /// universe is not re-added: the caller's constraint stands. /// - An empty narrowing falls back to the caller's allow-set rather than /// producing an unroutable request. +#[cfg_attr(not(feature = "standalone-selection"), allow(dead_code))] pub fn narrow_allowed_worker_ids_by_lora( filter: &dyn LoraWorkerFilter, lora_name: Option<&str>, diff --git a/lib/kv-router/src/scheduling/mod.rs b/lib/kv-router/src/scheduling/mod.rs index c48797e7293d..82eab3885b66 100644 --- a/lib/kv-router/src/scheduling/mod.rs +++ b/lib/kv-router/src/scheduling/mod.rs @@ -21,6 +21,7 @@ mod types; pub use filter::*; pub use local::LocalScheduler; pub use lora_filter::LoraWorkerFilter; +#[cfg(feature = "standalone-selection")] pub(crate) use lora_filter::narrow_allowed_worker_ids_by_lora; pub use overlap::{ CacheHitEstimates, OverlapAnalysis, OverlapScoresResponse, OverlapSignals, From 1f91e531b0ea244da805f254887371861f71a73e Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Fri, 11 Sep 2026 16:34:49 -0700 Subject: [PATCH 066/102] refactor(router): hand bookings to hosts as one BookingHandle Every host that received a booking from the scheduler got a `Box` and a `commit()` that might return `None`, and each host re-wrapped it: `Selected.lease`, `AdmittedFindBestMatchOutcome`, `WorkerSelection`, `KvRequestCleanup::new`, `enroll_public_request_attempt`, and `PrefillReservation` with its own release enum and `Drop`. `BookingHandle` is the one owner type: `descriptor()`, `commit()` to hand the booking to a longer-lived owner, `release()` for an acknowledged free, and `Drop` that frees an armed handle through the scheduler's cleanup queue. The scheduler builds it right after the internal lifecycle lease commits, with no await in between, so the exactly-once release invariant is unchanged; `RequestLifecycleLease` is now crate-private and covers only the queued window. `PrefillReservation` keeps its behavior with the handle as its field and loses its `Drop` impl and `ReservationRelease` enum; the unreachable "lease holds no booking" error paths are gone. `SelectionOutcome` in the frontend host carries the handle inline (about 100 bytes) rather than boxing it per selection; the enum is built once per request and consumed immediately, so the size-variance lint is allowed there. A unit test pins the handle's contract: drop while armed frees, commit leaves the booking to its owner, release frees before returning. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Ishan Dhanani --- lib/bindings/python/rust/llm/kv.rs | 6 +- lib/kv-router/src/scheduling/local.rs | 40 +++-- lib/kv-router/src/scheduling/queue.rs | 145 ++++++++++++++++-- .../src/services/selection/core/operation.rs | 12 +- .../services/selection/core/reservations.rs | 10 +- .../src/services/selection/core/run.rs | 30 ++-- .../src/services/selection/core/tests.rs | 6 +- lib/llm/src/kv_router.rs | 61 ++++---- lib/llm/src/kv_router/prefill_router/query.rs | 39 +---- lib/llm/src/kv_router/routing_host/kv.rs | 4 +- .../kv_router/routing_host/kv_selection.rs | 14 +- .../kv_router/routing_host/request_guard.rs | 12 +- 12 files changed, 228 insertions(+), 151 deletions(-) diff --git a/lib/bindings/python/rust/llm/kv.rs b/lib/bindings/python/rust/llm/kv.rs index 1936666281a8..d9bf1e1fd97f 100644 --- a/lib/bindings/python/rust/llm/kv.rs +++ b/lib/bindings/python/rust/llm/kv.rs @@ -2523,7 +2523,7 @@ impl KvRouter { ) .await .map_err(to_pyerr)?; - let (outcome, lease) = admitted.into_parts(); + let (outcome, booking) = admitted.into_parts(); let (best_worker, overlap_blocks) = match outcome { llm_rs::kv_router::FindBestMatchOutcome::Routed { worker, @@ -2559,9 +2559,9 @@ impl KvRouter { None }; - if let Some(lease) = lease { + if let Some(booking) = booking { chooser - .enroll_public_request_attempt(lease, routing_decision) + .enroll_public_request_attempt(booking, routing_decision) .await .map_err(to_pyerr)?; } else if let Some(tokens_with_hashes) = routing_decision { diff --git a/lib/kv-router/src/scheduling/local.rs b/lib/kv-router/src/scheduling/local.rs index ff685a5bb111..061b6f58d7f4 100644 --- a/lib/kv-router/src/scheduling/local.rs +++ b/lib/kv-router/src/scheduling/local.rs @@ -17,7 +17,7 @@ use super::overlap_refresh::{NoopOverlapScoresRefresh, OverlapScoresRefresh}; use super::policy_config::PolicyProfile; use super::prefill_load::PrefillLoadEstimator; use super::queue::{ - ClassQueueStats, RequestLifecycleLease, SchedulerBookingCleanup, SchedulerBookingDescriptor, + BookingHandle, ClassQueueStats, SchedulerBookingCleanup, SchedulerBookingDescriptor, SchedulerQueue, }; use super::selector::{DefaultWorkerSelector, WorkerSelector}; @@ -291,27 +291,21 @@ where &self, request: ScheduleRequest, ) -> Result { - let (admitted, lease) = self.schedule_request_with_lease(request).await?; - if let Some(lease) = lease { - let _ = lease.commit(); + let (admitted, booking) = self.schedule_request_with_booking(request).await?; + if let Some(booking) = booking { + let _ = booking.commit(); } Ok(admitted) } - /// Schedule a request and keep its lifecycle lease armed: dropping the - /// lease frees the booking, `commit` hands it to a longer-lived owner. - /// The lease is `None` unless the mode is `TrackedWithLifecycle`. + /// Schedule a request and return an armed handle for its booking: dropping + /// the handle frees the booking, `commit` hands it to a longer-lived owner. + /// The handle is `None` unless the mode is `TrackedWithLifecycle`. #[doc(hidden)] - pub async fn schedule_request_with_lease( + pub async fn schedule_request_with_booking( &self, request: ScheduleRequest, - ) -> Result< - ( - AdmittedSchedulingResponse, - Option>, - ), - KvSchedulerError, - > { + ) -> Result<(AdmittedSchedulingResponse, Option), KvSchedulerError> { let tracked = request.mode.is_tracked(); let (resp_tx, resp_rx) = tokio::sync::oneshot::channel(); let (attempt_tx, attempt_rx) = tokio::sync::oneshot::channel(); @@ -342,10 +336,12 @@ where } else { AdmissionAttempt::Untracked }; - Ok(( - AdmittedSchedulingResponse { response, attempt }, - lifecycle_lease, - )) + // No await between `commit()` and the handle build: the booking is + // always guarded by exactly one of the lease and the handle. + let booking = lifecycle_lease + .and_then(|lease| lease.commit()) + .map(|booking| self.queue.booking_handle(booking)); + Ok((AdmittedSchedulingResponse { response, attempt }, booking)) } /// Select a worker from current scheduler state without queue admission or booking. @@ -527,19 +523,19 @@ where /// Book a request only when its worker is already registered, so a request /// racing worker removal cannot lazily recreate the removed worker/rank. - /// The returned lease guards the booking: dropping it frees the booking, + /// The returned handle guards the booking: dropping it frees the booking, /// `commit` hands it over. #[doc(hidden)] pub fn add_request_if_registered_guarded( &self, req: SequenceRequest, - ) -> Result, SequenceError> { + ) -> Result { let request_id = req.request_id.clone(); let worker = req.worker; let attempt_id = self .slots .add_request_if_registered_admitted(req, Instant::now())?; - Ok(self.queue.lease_for_booking(SchedulerBookingDescriptor { + Ok(self.queue.booking_handle(SchedulerBookingDescriptor { request_id, worker, attempt_id, diff --git a/lib/kv-router/src/scheduling/queue.rs b/lib/kv-router/src/scheduling/queue.rs index e7996c039c81..01effd45c202 100644 --- a/lib/kv-router/src/scheduling/queue.rs +++ b/lib/kv-router/src/scheduling/queue.rs @@ -362,9 +362,58 @@ impl SchedulerBookingCleanup { } } -/// Single-owner cleanup lease for one scheduler-tracked request. +/// A booking whose release this handle owns until `commit` hands it over. +/// Dropping an armed handle frees the booking through the scheduler's cleanup queue. #[doc(hidden)] -pub struct RequestLifecycleLease { +#[must_use] +pub struct BookingHandle { + booking: SchedulerBookingDescriptor, + cleanup: SchedulerBookingCleanup, + armed: bool, +} + +impl std::fmt::Debug for BookingHandle { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("BookingHandle") + .field("booking", &self.booking) + .field("armed", &self.armed) + .finish() + } +} + +impl BookingHandle { + pub fn descriptor(&self) -> &SchedulerBookingDescriptor { + &self.booking + } + + /// Hand the booking to a longer-lived owner; the handle stops guarding it. + #[must_use] + pub fn commit(mut self) -> SchedulerBookingDescriptor { + self.armed = false; + self.booking.clone() + } + + /// Free the booking now and wait for the scheduler to acknowledge it. + pub async fn release(mut self) -> Result<(), SequenceError> { + self.armed = false; + self.cleanup + .enqueue_acknowledged(self.booking.clone()) + .wait() + .await + } +} + +impl Drop for BookingHandle { + fn drop(&mut self) { + if self.armed { + self.cleanup.enqueue(self.booking.clone()); + } + } +} + +/// Single-owner cleanup lease for one scheduler-tracked request. +pub(crate) struct RequestLifecycleLease { cleanup: Arc, actor_tx: mpsc::Sender, transfer: Option>, @@ -738,18 +787,13 @@ impl< } } - /// An armed lease for a booking made outside the admission actor. - pub(crate) fn lease_for_booking( - &self, - booking: SchedulerBookingDescriptor, - ) -> Box { - let transfer = AdmissionLifecycleTransfer::new(booking.request_id.clone()); - transfer.arm_booking(booking); - Box::new(RequestLifecycleLease { - cleanup: Arc::clone(&self.cleanup), - actor_tx: self.admission_tx.clone(), - transfer: Some(Arc::new(transfer)), - }) + /// An armed handle for an existing booking. + pub(crate) fn booking_handle(&self, booking: SchedulerBookingDescriptor) -> BookingHandle { + BookingHandle { + booking, + cleanup: self.booking_cleanup(), + armed: true, + } } /// Select a worker from current scheduler state without entering admission. @@ -1728,6 +1772,79 @@ mod tests { assert_eq!(cleanup, booking); } + fn book_directly( + slots: &ActiveSequencesMultiWorker, + request_id: &str, + ) -> SchedulerBookingDescriptor { + let worker = WorkerWithDpRank::new(0, 0); + let attempt_id = slots + .add_request_admitted( + crate::sequences::SequenceRequest { + request_id: request_id.to_string(), + token_sequence: None, + track_prefill_tokens: true, + expected_output_tokens: None, + prefill_load_hint: None, + worker, + lora_name: None, + }, + Instant::now(), + ) + .expect("worker 0 is registered"); + SchedulerBookingDescriptor { + request_id: request_id.to_string(), + worker, + attempt_id, + } + } + + async fn wait_freed( + slots: &ActiveSequencesMultiWorker, + booking: &SchedulerBookingDescriptor, + ) { + tokio::time::timeout(Duration::from_secs(2), async { + while slots.has_booking(booking) { + tokio::task::yield_now().await; + } + }) + .await + .expect("booking was freed"); + } + + /// The handle frees its booking exactly when it is dropped armed; `commit` + /// hands the booking over untouched and `release` frees it with an ack. + #[tokio::test(flavor = "multi_thread")] + async fn booking_handle_frees_only_while_armed() { + let (queue, slots) = make_queue(1, 16, 512, Some(0.0)); + + let armed = book_directly(&slots, "armed"); + drop(queue.booking_handle(armed.clone())); + wait_freed(&slots, &armed).await; + + let committed = book_directly(&slots, "committed"); + let descriptor = queue.booking_handle(committed.clone()).commit(); + assert_eq!(descriptor, committed); + + // The cleanup queue drains in order, so an acknowledged release + // enqueued after the commit proves the commit enqueued nothing. + let released = book_directly(&slots, "released"); + queue + .booking_handle(released.clone()) + .release() + .await + .expect("acknowledged release"); + assert!( + !slots.has_booking(&released), + "release frees before it returns" + ); + assert!( + slots.has_booking(&committed), + "a committed handle leaves the booking to its new owner" + ); + slots.free(&committed.request_id, decay_now()).unwrap(); + slots.assert_completely_drained(decay_now()); + } + struct DropResponseOnLoadPublisher { response_rx: Arc>>, } diff --git a/lib/kv-router/src/services/selection/core/operation.rs b/lib/kv-router/src/services/selection/core/operation.rs index 8d0a47baca82..bf8929b78c99 100644 --- a/lib/kv-router/src/services/selection/core/operation.rs +++ b/lib/kv-router/src/services/selection/core/operation.rs @@ -16,7 +16,7 @@ use crate::protocols::{ WorkerWithDpRank, }; use crate::scheduling::config::RouterConfigOverride; -use crate::scheduling::queue::RequestLifecycleLease; +use crate::scheduling::queue::BookingHandle; use crate::scheduling::{AdvisoryWorkerLoad, QueueRejection, SchedulingResponse, SessionContext}; use super::super::error::SelectionError; @@ -56,8 +56,8 @@ pub enum SelectionAdmission { /// Queue admission with a booking the core records as a reservation under /// `selection_id` and releases through its lifecycle calls. Book { selection_id: String }, - /// Queue admission with a booking the caller owns: its lifecycle lease is - /// returned armed in [`Selected::lease`]. + /// Queue admission with a booking the caller owns: its booking handle is + /// returned armed in [`Selected::booking`]. Lease { request_id: String }, /// Skip queue admission and report the chosen worker's load. Advisory { request_id: Option }, @@ -142,7 +142,7 @@ pub struct Selected { pub kv_hint: Option, pub routing_hashes: Option>, pub shared_cache_hits: Option, - /// The booking's lifecycle lease; `Lease` admission only. Dropping it frees - /// the booking, `commit` hands it to the caller's own cleanup. - pub lease: Option>, + /// The booking's handle; `Lease` admission only. Dropping it frees the + /// booking, `commit` hands it to the caller's own cleanup. + pub booking: Option, } diff --git a/lib/kv-router/src/services/selection/core/reservations.rs b/lib/kv-router/src/services/selection/core/reservations.rs index 87debde19db5..3cc069eae657 100644 --- a/lib/kv-router/src/services/selection/core/reservations.rs +++ b/lib/kv-router/src/services/selection/core/reservations.rs @@ -353,9 +353,9 @@ impl SelectionCore { _ => None, }; // Strict booking: never lazily recreate a worker/rank removed since the - // reservation was resolved. The lease frees the booking if this future + // reservation was resolved. The handle frees the booking if this future // is dropped before `install`. - let lease = entry + let booking = entry .scheduler .add_request_if_registered_guarded(SequenceRequest { request_id: selection_id.clone(), @@ -377,7 +377,7 @@ impl SelectionCore { } claim.install(Reservation { partition: key.clone(), - booking: Some(lease.commit().ok_or_else(missing_booking)?), + booking: Some(booking.commit()), _affinity_lease: affinity_lease, }); @@ -540,10 +540,6 @@ pub(super) fn sweep_reservation_index( swept } -pub(super) fn missing_booking() -> SelectionError { - SelectionError::Internal("booking lease holds no booking".to_string()) -} - pub(super) fn spawn_reservation_index_sweep( entries: Arc, index: Arc, diff --git a/lib/kv-router/src/services/selection/core/run.rs b/lib/kv-router/src/services/selection/core/run.rs index 8b9ee42a5238..0b733438fa02 100644 --- a/lib/kv-router/src/services/selection/core/run.rs +++ b/lib/kv-router/src/services/selection/core/run.rs @@ -15,8 +15,8 @@ //! request ends (`dropped_selection_future_frees_its_booking`, //! `reservation_index_tracks_bookings_until_freed`). //! - Release ownership moves in one direction: the scheduler actor books, hands -//! the armed lease to this run, and the run hands it to the index row for -//! `Book` (`claim.install`) or to the host for `Lease` (`Selected::lease`) +//! the armed handle to this run, and the run hands it to the index row for +//! `Book` (`claim.install`) or to the host for `Lease` (`Selected::booking`) //! (`dropped_selection_future_frees_its_booking`, //! `lease_admission_installs_no_index_row_and_records_nothing`). //! - The session hold is taken before the booking is made @@ -37,7 +37,7 @@ //! `dropped_book_selection_during_routing_record_frees_booking_and_claim`). use super::hint::transfer_hint_for_selection; -use super::reservations::{Reservation, missing_booking}; +use super::reservations::Reservation; use super::*; /// Action id of the single `kv.fetch` action a selection's KV hint carries. @@ -185,7 +185,7 @@ impl SelectionCore { kv_hint, routing_hashes: _, shared_cache_hits: _, - lease: _, + booking: _, } = selected; let booked = sequence_hashes.is_some(); let potential_decode_blocks = response.potential_decode_blocks as u64; @@ -418,7 +418,7 @@ impl SelectionCore { routing_constraints, shared_cache_hits, }; - // `lease` guards the booking until it is installed below: any early + // `booking` guards the booking until it is installed below: any early // return or drop before then frees it. let scheduled = tokio::select! { biased; @@ -438,14 +438,14 @@ impl SelectionCore { } else { entry .scheduler - .schedule_request_with_lease(schedule_request) + .schedule_request_with_booking(schedule_request) .instrument(tracing::info_span!("kv_router.schedule")) .await - .map(|(admitted, lease)| (admitted.response, None, lease)) + .map(|(admitted, booking)| (admitted.response, None, booking)) } } => result, }; - let (response, advisory_load, lease) = match scheduled { + let (response, advisory_load, booking) = match scheduled { Ok(scheduled) => scheduled, Err(KvSchedulerError::QueueRejected(rejection)) => { return Ok(SelectionOutcome::QueueRejected { rejection }); @@ -488,13 +488,13 @@ impl SelectionCore { // The routing hashes go to exactly one of: the reservation recorded // now, or the replay cache a later reservation records from. let mut routing_hashes = routing_hashes; - let lease = if let Some(claim) = claim { - let Some(lease) = lease else { + let booking = if let Some(claim) = claim { + let Some(booking) = booking else { return Err(SelectionError::Internal( - "booked selection has no lifecycle lease".to_string(), + "booked selection has no booking handle".to_string(), )); }; - // A rejected affinity commit returns while the lease is still armed, + // A rejected affinity commit returns while the handle is still armed, // so the booking is freed and nothing below is recorded. let affinity_lease = match (affinity_hold, managed_session) { (Some(hold), Some((table, session_id))) => { @@ -508,12 +508,12 @@ impl SelectionCore { } claim.install(Reservation { partition: key.clone(), - booking: Some(lease.commit().ok_or_else(missing_booking)?), + booking: Some(booking.commit()), _affinity_lease: affinity_lease, }); None } else { - lease + booking }; if let Some((cache_id, sequence_hashes, lora_name, track_prefill_tokens, session_id)) = @@ -550,7 +550,7 @@ impl SelectionCore { kv_hint, routing_hashes: returned_routing_hashes, shared_cache_hits: host_shared_cache_hits, - lease, + booking, })) } diff --git a/lib/kv-router/src/services/selection/core/tests.rs b/lib/kv-router/src/services/selection/core/tests.rs index f74600679177..3386aecaf0ee 100644 --- a/lib/kv-router/src/services/selection/core/tests.rs +++ b/lib/kv-router/src/services/selection/core/tests.rs @@ -1258,7 +1258,7 @@ async fn lease_admission_keeps_a_selection_whose_worker_drained_while_queued() { Ok(SelectionOutcome::Selected(selected)) => ( selected.response.best_worker.worker_id, selected.endpoint, - selected.lease.is_some(), + selected.booking.is_some(), ), Ok(SelectionOutcome::QueueRejected { .. }) => panic!("queue rejected"), Err(error) => panic!("lease selection failed: {error:?}"), @@ -1329,8 +1329,8 @@ async fn lease_admission_installs_no_index_row_and_records_nothing() { "a lease admission installs no index row" ); assert!( - selected.lease.is_some(), - "the booking's lease goes to the host" + selected.booking.is_some(), + "the booking's handle goes to the host" ); assert!(selected.sequence_hashes.is_none()); assert!(selected.routing_hashes.is_none()); diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index b6f8fd5b89a1..65665c4871a5 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -27,7 +27,7 @@ use dynamo_kv_router::{ scheduling::{ CacheHitEstimates, OverlapAnalysis, OverloadedWorkerProvider, WorkerAvailabilityProvider, effective_prefill_tokens, overlap::cache_hit_estimates_from_tiered_matches, - queue::RequestLifecycleLease, + queue::BookingHandle, }, selector::WorkerInputs, services::selection::{ @@ -432,19 +432,19 @@ pub(super) enum FindBestMatchAdmission { WithoutAdmission, } -/// A routed outcome with the booking's lifecycle lease, when the request was -/// booked. Dropping the lease frees the booking; the caller commits it into -/// its own cleanup. +/// A routed outcome with the booking's handle, when the request was booked. +/// Dropping the handle frees the booking; the caller commits it into its own +/// cleanup. #[doc(hidden)] pub struct AdmittedFindBestMatchOutcome { pub(super) outcome: FindBestMatchOutcome, - pub(super) lease: Option>, + pub(super) booking: Option, } impl AdmittedFindBestMatchOutcome { #[doc(hidden)] - pub fn into_parts(self) -> (FindBestMatchOutcome, Option>) { - (self.outcome, self.lease) + pub fn into_parts(self) -> (FindBestMatchOutcome, Option) { + (self.outcome, self.booking) } } @@ -1013,13 +1013,11 @@ impl KvRouter { #[doc(hidden)] pub async fn enroll_public_request_attempt( &self, - lease: Box, + booking: BookingHandle, routing_decision: Option, ) -> Result<(), KvRouterError> { // Nothing awaits between taking the booking over and registering it. - let booking = lease.commit().ok_or_else(|| { - KvRouterError::Unsupported("booking lease holds no booking".to_string()) - })?; + let booking = booking.commit(); let worker = booking.worker; let lru_registration = self.approximate_lru_rank_registration(worker); let approximate_lru = lru_registration.and_then(|registration| { @@ -1230,14 +1228,14 @@ impl KvRouter { routing_constraints, ) .await?; - if let Some(lease) = admitted.lease { - self.enroll_public_request_attempt(lease, None).await?; + if let Some(booking) = admitted.booking { + self.enroll_public_request_attempt(booking, None).await?; } Ok(admitted.outcome) } - /// Return the admitted routing wrapper without enrolling it in a detached - /// lifecycle lease. Internal bindings use this to attach optional LRU state + /// Return the admitted routing wrapper without enrolling its booking handle + /// in a detached lease. Internal bindings use this to attach optional LRU state /// before installing the one shared request lease. #[doc(hidden)] #[allow(clippy::too_many_arguments)] @@ -1375,7 +1373,7 @@ impl KvRouter { FindBestMatchAdmission::WithAdmission => { FindBestMatchInnerOutcome::WithAdmission(AdmittedFindBestMatchOutcome { outcome: FindBestMatchOutcome::QueueRejected { rejection }, - lease: None, + booking: None, }) } FindBestMatchAdmission::WithoutAdmission => { @@ -1400,11 +1398,11 @@ impl KvRouter { kv_hint, routing_hashes, shared_cache_hits, - lease, + booking, .. } = selected; - if update_states && is_admitted_routing && lease.is_none() { - anyhow::bail!("booked selection returned no lifecycle lease"); + if update_states && is_admitted_routing && booking.is_none() { + anyhow::bail!("booked selection returned no booking handle"); } let routing_hashes = routing_hashes.map(RoutingDecisionHashes::from_local_hashes); let overlap_blocks = response.effective_overlap_blocks.round() as u32; @@ -1449,7 +1447,7 @@ impl KvRouter { routing_hashes, kv_hint, }, - lease, + booking, }) } FindBestMatchAdmission::WithoutAdmission => { @@ -1600,10 +1598,6 @@ impl KvRouter { self.scheduler.affinity_coordinator(ttl, mode) } - pub(crate) fn booking_cleanup(&self) -> scheduler::SchedulerBookingCleanup { - self.scheduler.booking_cleanup() - } - pub(crate) fn request_lease_manager(&self) -> &request_lease::RequestLeaseManager { &self.request_leases } @@ -2621,9 +2615,9 @@ mod tests { trace.push(GoldenDecision::from_outcome(&admitted.outcome)); // Keep the booking: later rows are scored against its load. let booking = admitted - .lease - .and_then(|lease| lease.commit()) - .expect("tracked selection carries its booking lease"); + .booking + .map(BookingHandle::commit) + .expect("tracked selection carries its booking handle"); assert_eq!(booking.request_id, context_id); } @@ -2677,10 +2671,7 @@ mod tests { let FindBestMatchInnerOutcome::WithAdmission(kept) = select("kept").await.unwrap() else { panic!("admitted routing returned advisory outcome"); }; - let _booking = kept - .lease - .and_then(|lease| lease.commit()) - .expect("booking"); + let _booking = kept.booking.map(BookingHandle::commit).expect("booking"); tokio::time::timeout(std::time::Duration::from_secs(2), async { while router.scheduler.has_request("dropped") { @@ -2768,13 +2759,13 @@ mod tests { let FindBestMatchInnerOutcome::WithAdmission(admitted) = outcome else { panic!("admitted routing returned advisory outcome"); }; - let lease = admitted - .lease - .expect("tracked selection carries its booking lease"); + let booking = admitted + .booking + .expect("tracked selection carries its booking handle"); // Park the routing update after `register_detached` has run. let mut enroll = Box::pin( - router.enroll_public_request_attempt(lease, Some(TokensWithHashes::new(prompt, 2))), + router.enroll_public_request_attempt(booking, Some(TokensWithHashes::new(prompt, 2))), ); assert!(futures::poll!(&mut enroll).is_pending()); assert!( diff --git a/lib/llm/src/kv_router/prefill_router/query.rs b/lib/llm/src/kv_router/prefill_router/query.rs index bf84c8c15bba..39d1201a3e6f 100644 --- a/lib/llm/src/kv_router/prefill_router/query.rs +++ b/lib/llm/src/kv_router/prefill_router/query.rs @@ -6,7 +6,7 @@ use std::collections::HashSet; use anyhow::Result; use dynamo_kv_router::{ protocols::{BlockExtraInfo, RoutingConstraints, WorkerId, WorkerWithDpRank}, - scheduling::queue::{SchedulerBookingCleanup, SchedulerBookingDescriptor}, + scheduling::queue::BookingHandle, }; use super::{PrefillError, PrefillLifecycleState, PrefillQueryOutcome, PrefillRouter}; @@ -21,15 +21,7 @@ use super::{PrefillError, PrefillLifecycleState, PrefillQueryOutcome, PrefillRou pub struct PrefillReservation { worker: WorkerWithDpRank, dp_rank: Option, - release: ReservationRelease, -} - -enum ReservationRelease { - Kv { - cleanup: SchedulerBookingCleanup, - booking: Option, - }, - None, + booking: Option, } impl PrefillReservation { @@ -43,25 +35,13 @@ impl PrefillReservation { /// Release this booking and wait for the scheduler to acknowledge it. pub async fn release(mut self) -> Result<()> { - if let ReservationRelease::Kv { cleanup, booking } = &mut self.release - && let Some(booking) = booking.take() - { - cleanup.enqueue_acknowledged(booking).wait().await?; + if let Some(booking) = self.booking.take() { + booking.release().await?; } Ok(()) } } -impl Drop for PrefillReservation { - fn drop(&mut self) { - if let ReservationRelease::Kv { cleanup, booking } = &mut self.release - && let Some(booking) = booking.take() - { - cleanup.enqueue(booking); - } - } -} - impl PrefillRouter { /// Select a prefill worker and reserve it when KV routing is enabled. /// @@ -103,7 +83,7 @@ impl PrefillRouter { return Ok(PrefillReservation { worker: WorkerWithDpRank::new(worker_id, 0), dp_rank: None, - release: ReservationRelease::None, + booking: None, }); }; let admitted = chooser @@ -125,19 +105,16 @@ impl PrefillRouter { routing_constraints, ) .await?; - let (outcome, lease) = admitted.into_parts(); + let (outcome, booking) = admitted.into_parts(); match outcome { crate::kv_router::FindBestMatchOutcome::Routed { worker, .. } => { - let Some(booking) = lease.and_then(|lease| lease.commit()) else { + let Some(booking) = booking else { anyhow::bail!("prefill reservation admission did not return a booking"); }; Ok(PrefillReservation { worker, dp_rank: Some(worker.dp_rank), - release: ReservationRelease::Kv { - cleanup: chooser.booking_cleanup(), - booking: Some(booking), - }, + booking: Some(booking), }) } crate::kv_router::FindBestMatchOutcome::QueueRejected { rejection } => { diff --git a/lib/llm/src/kv_router/routing_host/kv.rs b/lib/llm/src/kv_router/routing_host/kv.rs index 83b86ef753eb..7c96a640c2d9 100644 --- a/lib/llm/src/kv_router/routing_host/kv.rs +++ b/lib/llm/src/kv_router/routing_host/kv.rs @@ -203,7 +203,7 @@ impl RoutingHost { Arc::clone(self.kv_router()), request.context().id().to_string(), selection.worker, - selection.lease.take(), + selection.booking.take(), ), selection, affinity, @@ -337,7 +337,7 @@ impl RoutingHost { self.request_metrics.clone(), context_id.clone(), selected_worker, - selection.lease.take(), + selection.booking.take(), request, ), }; diff --git a/lib/llm/src/kv_router/routing_host/kv_selection.rs b/lib/llm/src/kv_router/routing_host/kv_selection.rs index 980b47943c31..a0a680cbedc5 100644 --- a/lib/llm/src/kv_router/routing_host/kv_selection.rs +++ b/lib/llm/src/kv_router/routing_host/kv_selection.rs @@ -10,9 +10,7 @@ use dynamo_kv_router::{ protocols::{ BlockExtraInfo, RoutingConstraints, WorkerAffinityTarget, WorkerId, WorkerWithDpRank, }, - scheduling::{ - AdvisoryWorkerLoad, QueueRejection, RoutingEligibility, queue::RequestLifecycleLease, - }, + scheduling::{AdvisoryWorkerLoad, QueueRejection, RoutingEligibility, queue::BookingHandle}, }; use dynamo_runtime::{dynamo_nvtx_range, pipeline::Error}; @@ -31,8 +29,8 @@ use crate::{ pub(super) struct WorkerSelection { pub(super) worker: WorkerWithDpRank, - /// The booking's lifecycle lease until the request's cleanup takes it over. - pub(super) lease: Option>, + /// The booking's handle until the request's cleanup takes it over. + pub(super) booking: Option, pub(super) overlap_amount: u32, pub(super) effective_overlap_blocks: f64, pub(super) cached_tokens: usize, @@ -42,6 +40,8 @@ pub(super) struct WorkerSelection { pub(super) kv_hint: Option, } +// Transient return value; `Routed` is moved into `WorkerSelection` right away. +#[allow(clippy::large_enum_variant)] pub(super) enum SelectionOutcome { Routed(WorkerSelection), QueueRejected(QueueRejection), @@ -138,7 +138,7 @@ impl RoutingHost { kv_hint, } => Ok(SelectionOutcome::Routed(WorkerSelection { worker, - lease: admitted.lease, + booking: admitted.booking, overlap_amount: overlap_blocks, effective_overlap_blocks, cached_tokens, @@ -162,7 +162,7 @@ impl RoutingHost { routing_hashes, } => Ok(SelectionOutcome::Routed(WorkerSelection { worker, - lease: None, + booking: None, overlap_amount: overlap_blocks, effective_overlap_blocks, cached_tokens, diff --git a/lib/llm/src/kv_router/routing_host/request_guard.rs b/lib/llm/src/kv_router/routing_host/request_guard.rs index f59dce786977..ad98cfc2e151 100644 --- a/lib/llm/src/kv_router/routing_host/request_guard.rs +++ b/lib/llm/src/kv_router/routing_host/request_guard.rs @@ -22,7 +22,7 @@ use dynamo_kv_router::{ BlockExtraInfo, BlockHashOptions, WorkerWithDpRank, compute_block_hash_for_seq, compute_next_seq_hash, }, - scheduling::queue::RequestLifecycleLease, + scheduling::queue::BookingHandle, }; use dynamo_runtime::{ error::DynamoError, @@ -396,15 +396,15 @@ pub(super) struct KvRequestCleanup { } impl KvRequestCleanup { - /// Takes the booking over from `lease`: from here the request's cleanup + /// Takes the booking over from `booking`: from here the request's cleanup /// owns it. pub(super) fn new( chooser: Arc, context_id: String, worker: WorkerWithDpRank, - lease: Option>, + booking: Option, ) -> Self { - let booking = lease.and_then(|lease| lease.commit()); + let booking = booking.map(BookingHandle::commit); let approximate_lru = booking.as_ref().and_then(|booking| { let registration = chooser.approximate_lru_rank_registration(worker)?; chooser.indexer().begin_approximate_lru_request( @@ -562,12 +562,12 @@ impl RequestGuard { request_metrics: Arc, context_id: String, worker: WorkerWithDpRank, - lease: Option>, + booking: Option, request: &PreprocessedRequest, ) -> Self { Self::new_kv_with_cleanup( request_metrics, - KvRequestCleanup::new(chooser, context_id, worker, lease), + KvRequestCleanup::new(chooser, context_id, worker, booking), request, ) } From bbd8412a91f437898451309090d33c46e104786f Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Fri, 11 Sep 2026 16:50:27 -0700 Subject: [PATCH 067/102] refactor(router): wire the request lease manager directly as the partition observer Peer-mirrored booking events reached the frontend's `RequestLeaseManager` through `LateBoundLeaseObserver`, a `OnceLock` shim installed when the replica ingress started. The shim existed only because the manager needed the scheduler's cleanup handle while the scheduler needed the manager as its observer. The manager now holds the cleanup handle in a `OnceLock` set right after the partition scheduler exists, so it is built first and passed to the selection core directly. The replica ingress still starts after the handle is set. Release paths that run before the handle is set log an error and skip the scheduler, which a test pins along with the normal acknowledged release; the ordering itself is fixed at the one construction site. The shim and its test are gone. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Ishan Dhanani --- lib/llm/src/kv_router.rs | 64 ++++++++++++++--- lib/llm/src/kv_router/embedded.rs | 19 ++--- lib/llm/src/kv_router/request_lease.rs | 52 +++++++++++--- lib/llm/src/kv_router/sequence.rs | 99 -------------------------- 4 files changed, 106 insertions(+), 128 deletions(-) diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index 65665c4871a5..2bf99ffad271 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -793,6 +793,10 @@ impl KvRouter { let available_worker_provider: WorkerAvailabilityProvider = Arc::new(move || client_for_availability.available_instance_ids()); + // The manager is the partition's lease observer, so it exists before + // the scheduler; its scheduler cleanup is set once the scheduler does. + let request_leases = + request_lease::RequestLeaseManager::new(cancellation_token.child_token()); let (scheduler, replica_ingress) = embedded::EmbeddedSelection::start( embedded::EmbeddedSelectionArgs { kv_router_config: kv_router_config.clone(), @@ -816,17 +820,14 @@ impl KvRouter { policy_factory, }, workers_with_configs.clone(), + Some(Arc::new(request_leases.clone())), cancellation_token.child_token(), ) .await?; - let request_leases = request_lease::RequestLeaseManager::new( - scheduler.booking_cleanup(), - cancellation_token.child_token(), - ); - // Inbound lifecycle events start only now that their consumer exists. - replica_ingress - .start(Arc::new(request_leases.clone())) - .await; + request_leases.set_scheduler(scheduler.booking_cleanup()); + // Inbound lifecycle events start only now that their consumer can + // release bookings. + replica_ingress.start().await; tracing::info!("KV Routing initialized"); let cancellation_token = cancellation_guard.disarm(); Ok(Self { @@ -2683,6 +2684,53 @@ mod tests { assert!(router.scheduler.has_request("kept")); } + /// A lease released before `set_scheduler` is a wiring bug: it logs and + /// leaves the booking alone. After `set_scheduler`, the release reaches + /// the scheduler's booking cleanup. + #[tokio::test] + async fn request_lease_manager_releases_bookings_only_once_scheduler_is_set() { + use dynamo_kv_router::scheduling::queue::SchedulerBookingDescriptor; + + let router = differential_router("lease-manager").await; + let worker = WorkerWithDpRank::from_worker_id(1); + let scheduler = &router.scheduler; + let book = |request_id: &'static str| async move { + let attempt_id = scheduler + .add_request_admitted(SequenceRequest { + request_id: request_id.to_string(), + token_sequence: None, + track_prefill_tokens: false, + expected_output_tokens: None, + prefill_load_hint: None, + worker, + lora_name: None, + }) + .await + .expect("booking"); + SchedulerBookingDescriptor { + request_id: request_id.to_string(), + worker, + attempt_id, + } + }; + let cancel = CancellationToken::new(); + let manager = request_lease::RequestLeaseManager::new(cancel.child_token()); + + let before = manager.register_local(book("before").await, None); + drop(before); + assert!( + router.scheduler.has_request("before"), + "no scheduler to release through yet" + ); + + manager.set_scheduler(router.scheduler.booking_cleanup()); + let after = manager.register_local(book("after").await, None); + after.finish().await; + assert!(!router.scheduler.has_request("after")); + assert!(router.scheduler.has_request("before")); + cancel.cancel(); + } + /// A public enrollment cancelled while its routing update is in flight /// frees the booking through the detached enrollment's drop. #[tokio::test] diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs index 9e5965705972..6ecfc33da4af 100644 --- a/lib/llm/src/kv_router/embedded.rs +++ b/lib/llm/src/kv_router/embedded.rs @@ -175,20 +175,16 @@ pub(crate) struct EmbeddedSelection { queue_metric_indices: HashMap, } -/// The partition's inbound replica ingress, not yet running. Starting it -/// installs the router's lease manager first, so no lifecycle event reaches -/// the partition before its consumer exists. +/// The partition's inbound replica ingress, not yet running. The caller starts +/// it only after the lease manager it passed as `request_leases` has its +/// scheduler set, so no lifecycle event reaches the manager before it can +/// release the booking. pub(crate) struct PendingReplicaIngress { ingress: crate::kv_router::sequence::ReplicaIngress, - leases: Arc, } impl PendingReplicaIngress { - pub(crate) async fn start( - self, - leases: Arc, - ) { - self.leases.install(leases); + pub(crate) async fn start(self) { self.ingress.start().await; } } @@ -239,6 +235,7 @@ impl EmbeddedSelection { pub(crate) async fn start( args: EmbeddedSelectionArgs, workers_with_configs: RuntimeConfigWatch, + request_leases: Option>, cancellation_token: CancellationToken, ) -> Result<(Self, PendingReplicaIngress)> { let worker_type = args.worker_role.unwrap_or(WorkerType::Aggregated); @@ -260,7 +257,6 @@ impl EmbeddedSelection { ) .await .context("start replica sync for the embedded selection partition")?; - let leases = Arc::new(crate::kv_router::sequence::LateBoundLeaseObserver::default()); let slot = std::sync::Mutex::new(Some(channels)); let replica_sync: Option = Some(Arc::new(move |_partition| { @@ -297,7 +293,7 @@ impl EmbeddedSelection { }, replication: HostReplication { channels: replica_sync, - request_leases: Some(leases.clone()), + request_leases, }, }) .build() @@ -379,7 +375,6 @@ impl EmbeddedSelection { }, PendingReplicaIngress { ingress: replica_ingress, - leases, }, )) } diff --git a/lib/llm/src/kv_router/request_lease.rs b/lib/llm/src/kv_router/request_lease.rs index 9eaa56b647cc..df0291c75eed 100644 --- a/lib/llm/src/kv_router/request_lease.rs +++ b/lib/llm/src/kv_router/request_lease.rs @@ -4,7 +4,7 @@ use std::{ collections::HashMap, sync::{ - Arc, Weak, + Arc, OnceLock, Weak, atomic::{AtomicU8, Ordering}, }, time::Duration, @@ -97,7 +97,9 @@ impl RequestLeaseRecord { struct RequestLeaseManagerInner { active: Mutex, - scheduler: SchedulerBookingCleanup, + /// Set once the scheduler exists (`set_scheduler`); the manager is built + /// first because the scheduler takes it as its lease observer. + scheduler: OnceLock, } #[derive(Default)] @@ -168,8 +170,29 @@ impl RequestLeaseManagerInner { } } + /// The scheduler cleanup, or `None` (logged) for an event that arrives + /// before `set_scheduler`; that is a wiring bug and the booking is not + /// released here. + fn scheduler( + &self, + event: &str, + booking: &SchedulerBookingDescriptor, + ) -> Option<&SchedulerBookingCleanup> { + let scheduler = self.scheduler.get(); + if scheduler.is_none() { + tracing::error!( + request_id = %booking.request_id, + event, + "request lifecycle event before the lease manager's scheduler was set; the booking will not be released" + ); + } + scheduler + } + fn enqueue_completion(&self, record: &RequestLeaseRecord) { - self.scheduler.enqueue(record.booking.clone()); + if let Some(scheduler) = self.scheduler("completed", &record.booking) { + scheduler.enqueue(record.booking.clone()); + } if let Some(approximate_lru) = &record.approximate_lru { approximate_lru.release_now(); } @@ -179,7 +202,9 @@ impl RequestLeaseManagerInner { // NOTE: Request-liveness expiry is deliberately isolated to this router. // Local and mirrored scheduler copies expire independently, and only an // explicit lifecycle completion publishes `Free` to peer routers. - self.scheduler.enqueue_expired(record.booking.clone()); + if let Some(scheduler) = self.scheduler("expired", &record.booking) { + scheduler.enqueue_expired(record.booking.clone()); + } if let Some(approximate_lru) = &record.approximate_lru { approximate_lru.release_now(); } @@ -210,10 +235,10 @@ pub(crate) struct RequestLeaseManager { } impl RequestLeaseManager { - pub(crate) fn new(scheduler: SchedulerBookingCleanup, cancellation: CancellationToken) -> Self { + pub(crate) fn new(cancellation: CancellationToken) -> Self { let inner = Arc::new(RequestLeaseManagerInner { active: Mutex::new(ActiveRequestLeases::default()), - scheduler, + scheduler: OnceLock::new(), }); start_reaper( Arc::downgrade(&inner), @@ -223,6 +248,13 @@ impl RequestLeaseManager { Self { inner } } + /// Install the scheduler cleanup. Must run before any lifecycle event can + /// reach this manager (the replica ingress starts after it). + pub(crate) fn set_scheduler(&self, scheduler: SchedulerBookingCleanup) { + let installed = self.inner.scheduler.set(scheduler).is_ok(); + debug_assert!(installed, "request lease manager scheduler set twice"); + } + pub(crate) fn register_local( &self, booking: SchedulerBookingDescriptor, @@ -302,15 +334,17 @@ impl RequestLeaseManager { // the finishing future therefore cannot strand either cleanup. let scheduler_ack = self .inner - .scheduler - .enqueue_acknowledged(record.booking.clone()); + .scheduler("finished", &record.booking) + .map(|scheduler| scheduler.enqueue_acknowledged(record.booking.clone())); let lru_ack = record .approximate_lru .as_ref() .map(ApproximateRequestLease::begin_finish) .transpose(); - if let Err(error) = scheduler_ack.wait().await { + if let Some(scheduler_ack) = scheduler_ack + && let Err(error) = scheduler_ack.wait().await + { tracing::warn!( request_id = %record.booking.request_id, worker = ?record.booking.worker, diff --git a/lib/llm/src/kv_router/sequence.rs b/lib/llm/src/kv_router/sequence.rs index 75b8b727f9cb..0b5795bfb2b2 100644 --- a/lib/llm/src/kv_router/sequence.rs +++ b/lib/llm/src/kv_router/sequence.rs @@ -18,7 +18,6 @@ use dynamo_kv_router::protocols::{ pub use dynamo_kv_router::sequence::{ActiveSequences, RequestId}; use anyhow::Result; -use dynamo_kv_router::scheduling::queue::SchedulerBookingDescriptor; use dynamo_runtime::component::Endpoint; use dynamo_runtime::traits::DistributedRuntimeProvider; use dynamo_runtime::transports::event_plane::{ @@ -26,7 +25,6 @@ use dynamo_runtime::transports::event_plane::{ }; use std::collections::VecDeque; use std::future::Future; -use std::sync::Arc; use std::task::{Context, Poll}; use tokio::sync::mpsc; use tokio::time::Instant; @@ -42,56 +40,6 @@ use dynamo_runtime::transports::event_plane::MsgpackCodec; // if the queue is full, the newest event is dropped without blocking the local mutation. const REPLICA_EVENT_CHANNEL_CAPACITY: usize = 100_000; -/// The partition's lease observer until the router's `RequestLeaseManager` -/// exists; an event before `install` is a wiring bug and is dropped loudly. -#[derive(Default)] -pub(crate) struct LateBoundLeaseObserver { - target: std::sync::OnceLock>, -} - -impl LateBoundLeaseObserver { - pub(crate) fn install(&self, target: Arc) { - let installed = self.target.set(target).is_ok(); - debug_assert!(installed, "request lease manager installed twice"); - } - - fn target( - &self, - event: &str, - booking: &SchedulerBookingDescriptor, - ) -> Option<&Arc> { - let target = self.target.get(); - if target.is_none() { - tracing::error!( - request_id = %booking.request_id, - event, - "request lifecycle event before the lease manager was installed; the manager will not track this mirrored booking" - ); - } - target - } -} - -impl ReplicaRequestLeaseObserver for LateBoundLeaseObserver { - fn admitted(&self, booking: SchedulerBookingDescriptor) { - if let Some(target) = self.target("admitted", &booking) { - target.admitted(booking); - } - } - - fn progressed(&self, booking: &SchedulerBookingDescriptor) { - if let Some(target) = self.target("progressed", booking) { - target.progressed(booking); - } - } - - fn completed(&self, booking: &SchedulerBookingDescriptor) { - if let Some(target) = self.target("completed", booking) { - target.completed(booking); - } - } -} - /// How active-sequence events are framed on the wire for a transport. #[derive(Clone, Copy, Debug, PartialEq, Eq)] enum ActiveSequenceEventWireFormat { @@ -531,7 +479,6 @@ async fn forward_replica_events( mod tests { use super::*; use dynamo_kv_router::protocols::ActiveSequenceEventData; - use parking_lot::Mutex; use std::sync::Arc; use tokio::time::Instant; @@ -818,52 +765,6 @@ mod tests { .is_err() ); } - #[test] - fn late_bound_observer_drops_before_install_and_forwards_after() { - use dynamo_kv_router::sequences::NoopSequencePublisher; - use std::collections::HashMap; - #[derive(Default)] - struct Observer(Mutex>); - impl ReplicaRequestLeaseObserver for Observer { - fn admitted(&self, booking: SchedulerBookingDescriptor) { - self.0.lock().push(("add", booking.request_id)); - } - fn progressed(&self, booking: &SchedulerBookingDescriptor) { - self.0.lock().push(("progress", booking.request_id.clone())); - } - fn completed(&self, booking: &SchedulerBookingDescriptor) { - self.0.lock().push(("free", booking.request_id.clone())); - } - } - let slots = ActiveSequencesMultiWorker::new_without_expiry( - NoopSequencePublisher, - 4, - HashMap::from([(1, (0, 1))]), - true, - 0, - "test", - ); - let late = Arc::new(LateBoundLeaseObserver::default()); - assert!(slots.set_replica_request_lease_observer(late.clone())); - // Dropped: no target installed. - slots.apply_replica_batch(vec![ - add_event("before-install"), - free_event("before-install"), - ]); - let observer = Arc::new(Observer::default()); - late.install(observer.clone()); - slots.apply_replica_batch(vec![ - add_event("after-install"), - free_event("after-install"), - ]); - assert_eq!( - *observer.0.lock(), - vec![ - ("add", "after-install".to_string()), - ("free", "after-install".to_string()) - ] - ); - } #[tokio::test] #[serial_test::serial] From 5c2eb04504208c5193c6c7bc6b3f60f686e7a0c2 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Sat, 12 Sep 2026 13:11:12 -0700 Subject: [PATCH 068/102] chore(router): cleanup round one, drop test-only shims and duplicate tests Targeted at shrinking the diff without changing behavior or losing a check. Every deleted test names the test that still makes its assertions. Production: - `KvRouter::find_best_match`, `find_best_match_details`, `find_best_match_details_with_lifecycle`, `PrefillRouter::query_prefill_worker` and `PrefillQueryOutcome` had only test callers; the tests now go through `find_best_match_details_with_policy_class` and `reserve_prefill_worker`. `get_cache_hit_estimate` loses a layer whose other branch had no caller. `KvRouter::shutdown`, both `add_output_block` wrappers, `RuntimeIngress::shutdown`, `router_discovery_query`, dead consts and dead re-exports are gone; the `kv_router/scheduler.rs` shim is deleted and its six uses import from `dynamo_kv_router::scheduling`. `RoutePlan::signals` and `RoutePreview::signals` become a `pub(crate)` field. `AffinityCoordinator::acquire` calls the table directly instead of through a second copy of the wait loop. - kv-router: `SelectionService::{partition, cancelled}`, `SelectionPartition::block_size`, `AffinityInitialization::session_id`, `SessionAffinity::new` and four unused re-exports are gone; the six test-only hooks on `SessionAffinity` sit in one gated impl block; an identity map in `core/run.rs` is dropped. Tests: - `core/tests.rs` builds cores through `local_core`, `core_with` and `replay_reservation` instead of repeating the constructor literal twenty-eight times. - Deleted with survivors: `session_stays_on_its_first_worker_across_bookings` (`partition_sessions_keep_independent_bindings`, `two_phase_reservation_binds_the_session`), two blocks of `injected_lora_filter_narrows_candidates` (`lora_filter.rs` unit tests), the two `injected_*_provider` tests (merged into one), `embedded_selection_books_and_releases_like_the_runtime_scheduler` (route preview and plan tests; its held-booking assertion now lives in the request counters test), `dropped_selection_lease_frees_the_booking` (`booking_handle_frees_only_while_armed`, `dropped_selection_future_frees_its_booking`), `test_find_best_match_details_omits_routing_hashes_when_not_requested` (`lease_admission_installs_no_index_row_and_records_nothing`). `weighted_cache_hit_estimates_include_lower_tiers` moves next to the function it tests in `scheduling/overlap.rs`. - Frontend test fixtures merge into `make_router` and `tracked_router`; the golden-trace seeding runs only for the test that reads it. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Ishan Dhanani --- lib/kv-router/src/scheduling/overlap.rs | 29 + .../src/services/selection/affinity.rs | 48 +- .../src/services/selection/core/mod.rs | 4 - .../src/services/selection/core/run.rs | 3 +- .../src/services/selection/core/tests.rs | 614 +++++------------ lib/kv-router/src/services/selection/mod.rs | 10 +- .../src/services/selection/service.rs | 11 - lib/llm/src/kv_router.rs | 620 ++++-------------- lib/llm/src/kv_router/AGENTS.md | 2 +- lib/llm/src/kv_router/embedded.rs | 10 - lib/llm/src/kv_router/indexer/ingress.rs | 8 - .../kv_router/indexer/recovery/subscriber.rs | 1 + .../kv_router/prefill_router/activation.rs | 21 +- .../prefill_router/conditional_bypass.rs | 2 +- lib/llm/src/kv_router/prefill_router/mod.rs | 15 +- lib/llm/src/kv_router/prefill_router/query.rs | 117 +--- lib/llm/src/kv_router/request_lease.rs | 10 +- lib/llm/src/kv_router/routing_host.rs | 12 +- lib/llm/src/kv_router/routing_host/tests.rs | 151 +---- lib/llm/src/kv_router/scheduler.rs | 11 - lib/llm/src/session_affinity/coordinator.rs | 35 +- lib/llm/src/session_affinity/mod.rs | 3 +- 22 files changed, 421 insertions(+), 1316 deletions(-) delete mode 100644 lib/llm/src/kv_router/scheduler.rs diff --git a/lib/kv-router/src/scheduling/overlap.rs b/lib/kv-router/src/scheduling/overlap.rs index a6aed097188e..217980b09aa8 100644 --- a/lib/kv-router/src/scheduling/overlap.rs +++ b/lib/kv-router/src/scheduling/overlap.rs @@ -434,6 +434,35 @@ mod tests { assert_eq!(tiers.disk[&worker], 9); } + /// Default tier weights; `worker_2` has lower-tier hits but no device row. + #[test] + fn weighted_cache_hit_estimates_include_lower_tiers() { + let worker_1 = WorkerWithDpRank::new(1, 0); + let worker_2 = WorkerWithDpRank::new(2, 0); + let mut device = OverlapScores::new(); + device.scores.insert(worker_1, 2); + let mut host = LowerTierMatchDetails::default(); + host.hits.insert(worker_1, 1); + host.hits.insert(worker_2, 1); + let mut disk = LowerTierMatchDetails::default(); + disk.hits.insert(worker_1, 2); + let tiered = TieredMatchDetails { + device: MatchDetails { + overlap_scores: device, + ..Default::default() + }, + lower_tier: HashMap::from([(StorageTier::HostPinned, host), (StorageTier::Disk, disk)]), + }; + + let estimates = + cache_hit_estimates_from_tiered_matches(&KvRouterConfig::default(), 16, &tiered); + + assert_eq!(estimates.effective_overlap_blocks[&worker_1], 3.25); + assert_eq!(estimates.cached_tokens[&worker_1], 52); + assert_eq!(estimates.effective_overlap_blocks[&worker_2], 0.75); + assert_eq!(estimates.cached_tokens[&worker_2], 12); + } + #[test] fn score_response_uses_explicit_block_count_and_shared_credit() { let warm = WorkerWithDpRank::new(7, 1); diff --git a/lib/kv-router/src/services/selection/affinity.rs b/lib/kv-router/src/services/selection/affinity.rs index 3f72b6dab116..f466884dd1fe 100644 --- a/lib/kv-router/src/services/selection/affinity.rs +++ b/lib/kv-router/src/services/selection/affinity.rs @@ -209,10 +209,6 @@ pub enum AcquireStep { } impl SessionAffinity { - pub fn new(ttl: Duration) -> Result { - Self::with_config(SessionAffinityConfig::new(ttl)) - } - pub(crate) fn validate_ttl(ttl: Duration) -> Result<(), AffinityError> { if !(Duration::from_secs(1)..=Duration::from_secs(MAX_SESSION_AFFINITY_TTL_SECS)) .contains(&ttl) @@ -516,7 +512,26 @@ impl SessionAffinity { self.inner.apply_replica_update(session_id, target, version) } - #[cfg(any(test, feature = "testing"))] + fn validate_session_id(&self, session_id: &str) -> Result<(), AffinityError> { + if session_id.len() > self.inner.max_session_id_bytes { + return Err(AffinityError::InvalidArgument(format!( + "session affinity ID must not exceed {} bytes", + self.inner.max_session_id_bytes + ))); + } + Ok(()) + } + + fn reserve_entry(&self) -> Result<(), AffinityError> { + self.inner.reserve_entry().then_some(()).ok_or_else(|| { + AffinityError::ResourceExhausted("session affinity entry limit reached".to_string()) + }) + } +} + +/// Test hooks, also consumed by dependent crates' tests via `testing`. +#[cfg(any(test, feature = "testing"))] +impl SessionAffinity { pub fn entry_count(&self) -> usize { self.inner.entry_count.load(Ordering::Relaxed) } @@ -526,17 +541,14 @@ impl SessionAffinity { self.inner.cancel.clone() } - #[cfg(any(test, feature = "testing"))] pub async fn wait_for_reaper(&self) { self.inner.reaper_started.notified().await; } - #[cfg(any(test, feature = "testing"))] pub async fn wait_for_initializing_waiter(&self) { self.inner.waiter_observed.notified().await; } - #[cfg(any(test, feature = "testing"))] pub fn expire_for_test(&self, session_id: &str) { let Some(mut entry) = self.inner.entries.get_mut(session_id) else { panic!("session affinity entry missing"); @@ -556,22 +568,6 @@ impl SessionAffinity { pub fn next_version(&self) -> AffinityVersion { self.inner.next_version() } - - fn validate_session_id(&self, session_id: &str) -> Result<(), AffinityError> { - if session_id.len() > self.inner.max_session_id_bytes { - return Err(AffinityError::InvalidArgument(format!( - "session affinity ID must not exceed {} bytes", - self.inner.max_session_id_bytes - ))); - } - Ok(()) - } - - fn reserve_entry(&self) -> Result<(), AffinityError> { - self.inner.reserve_entry().then_some(()).ok_or_else(|| { - AffinityError::ResourceExhausted("session affinity entry limit reached".to_string()) - }) - } } impl Inner { @@ -689,10 +685,6 @@ pub struct AffinityInitialization { } impl AffinityInitialization { - pub fn session_id(&self) -> &str { - &self.session_id - } - /// Bind the session to the worker the request was dispatched to. pub fn commit(mut self, target: AffinityTarget) -> Result { validate_bound_target(&self.session_id, target, self.requested_target)?; diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index def4b6153b69..ef4f0fe292be 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -101,10 +101,6 @@ impl SelectionPartition { &self.0.key } - pub fn block_size(&self) -> u32 { - self.0.block_size - } - pub fn scheduler(&self) -> &SelectionScheduler { &self.0.scheduler } diff --git a/lib/kv-router/src/services/selection/core/run.rs b/lib/kv-router/src/services/selection/core/run.rs index 0b733438fa02..df4d2bbf10b6 100644 --- a/lib/kv-router/src/services/selection/core/run.rs +++ b/lib/kv-router/src/services/selection/core/run.rs @@ -293,8 +293,7 @@ impl SelectionCore { } (SessionBinding::Query { session_id }, Some(table)) => table .query_target(session_id, None) - .map_err(affinity_error)? - .map(|target| WorkerAffinityTarget::new(target.worker_id, target.dp_rank)), + .map_err(affinity_error)?, _ => affinity_target, }; // Router hints are attached to bookings only, and only when a worker in diff --git a/lib/kv-router/src/services/selection/core/tests.rs b/lib/kv-router/src/services/selection/core/tests.rs index 8e4ee2a414a7..43035f5da945 100644 --- a/lib/kv-router/src/services/selection/core/tests.rs +++ b/lib/kv-router/src/services/selection/core/tests.rs @@ -25,6 +25,92 @@ fn test_config(use_kv_events: bool) -> crate::config::KvRouterConfig { } } +fn local_core(config: crate::config::KvRouterConfig) -> SelectionCore { + local_core_with(config, 1, CancellationToken::new()) +} + +fn local_core_with( + config: crate::config::KvRouterConfig, + indexer_threads: usize, + cancel_token: CancellationToken, +) -> SelectionCore { + SelectionCore::try_new_local( + config, + indexer_threads, + cancel_token, + SelectionCacheConfig::default(), + ) + .expect("valid test config") +} + +/// `new_inner` with the test defaults; `remote_indexer` swaps the primary for +/// the standalone indexer at that URL. +fn core_with( + config: crate::config::KvRouterConfig, + host: SelectionHost, + policy_factory: Option, + worker_type: WorkerType, + remote_indexer: Option, + affinity: Option, +) -> SelectionCore { + let tracking_hash = Arc::new( + TrackingHashContext::from_config(&config).expect("valid tracking hash configuration"), + ); + let mut indexer_policy = IndexerPolicy::from_router_config(&config).expect("indexer policy"); + if let Some(base_url) = remote_indexer { + indexer_policy = indexer_policy + .with_remote_indexer(base_url) + .expect("remote policy"); + } + SelectionCore::new_inner( + config, + 1, + CancellationToken::new(), + None, + policy_factory, + host, + worker_type, + true, + SelectionCacheConfig::default(), + tracking_hash, + indexer_policy, + affinity, + ) +} + +fn core_with_host(host: SelectionHost) -> SelectionCore { + core_with_host_and_policy(host, None) +} + +fn core_with_host_and_policy( + host: SelectionHost, + policy_factory: Option, +) -> SelectionCore { + core_with( + test_config(false), + host, + policy_factory, + WorkerType::Aggregated, + None, + None, + ) +} + +fn replay_reservation(selection_id: &str) -> ReservationRequest { + ReservationRequest { + model_name: "model".to_string(), + routing_group: "default".to_string(), + selection_id: selection_id.to_string(), + worker_id: None, + dp_rank: None, + prompt: PromptRequest::default(), + router_config_override: None, + expected_output_tokens: None, + effective_prefill_tokens: None, + track_prefill_tokens: None, + } +} + fn worker(worker_id: WorkerId) -> WorkerRequest { WorkerRequest { worker_id, @@ -140,13 +226,7 @@ fn assert_shutdown_error(error: SelectionError) { #[test] fn parent_cancel_cancels_core() { let parent = CancellationToken::new(); - let core = SelectionCore::try_new_local( - test_config(false), - 1, - parent.clone(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"); + let core = local_core_with(test_config(false), 1, parent.clone()); assert!(!core.cancel_token.is_cancelled()); parent.cancel(); @@ -161,22 +241,12 @@ async fn selection_setup_uses_worker_type_label() { (WorkerType::Encode, "encode"), (WorkerType::Aggregated, "aggregated"), ] { - let config = test_config(false); - let tracking_hash = Arc::new( - TrackingHashContext::from_config(&config).expect("valid tracking hash configuration"), - ); - let core = SelectionCore::new_inner( - config, - 1, - CancellationToken::new(), - None, - None, + let core = core_with( + test_config(false), SelectionHost::default(), + None, worker_type, - true, - SelectionCacheConfig::default(), - tracking_hash, - IndexerPolicy::from_router_config(&test_config(false)).expect("indexer policy"), + None, None, ); @@ -192,35 +262,6 @@ async fn selection_setup_uses_worker_type_label() { } } -fn core_with_host(host: SelectionHost) -> SelectionCore { - core_with_host_and_policy(host, None) -} - -fn core_with_host_and_policy( - host: SelectionHost, - policy_factory: Option, -) -> SelectionCore { - let config = test_config(false); - let tracking_hash = Arc::new( - TrackingHashContext::from_config(&config).expect("valid tracking hash configuration"), - ); - let indexer_policy = IndexerPolicy::from_router_config(&config).expect("indexer policy"); - SelectionCore::new_inner( - config, - 1, - CancellationToken::new(), - None, - policy_factory, - host, - WorkerType::Aggregated, - true, - SelectionCacheConfig::default(), - tracking_hash, - indexer_policy, - None, - ) -} - async fn wait_for_overlap( core: &SelectionCore, request: impl Fn() -> SelectRequest, @@ -241,13 +282,7 @@ async fn wait_for_overlap( #[tokio::test] async fn bookings_populate_the_approximate_primary_without_kv_events() { // use_kv_events=false: the primary is approximate and bookings feed it. - let core = SelectionCore::try_new_local( - test_config(false), - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"); + let core = local_core(test_config(false)); core.upsert_worker(worker(1)).await.expect("worker upsert"); core.upsert_worker(worker(2)).await.expect("worker upsert"); @@ -279,20 +314,9 @@ async fn bookings_populate_the_approximate_primary_without_kv_events() { request.selection_id = Some("cached".to_string()); request.allowed_worker_ids = Some(HashSet::from([2])); core.select(request).await.expect("select"); - core.create_reservation(ReservationRequest { - model_name: "model".to_string(), - routing_group: "default".to_string(), - selection_id: "cached".to_string(), - worker_id: None, - dp_rank: None, - prompt: PromptRequest::default(), - router_config_override: None, - expected_output_tokens: None, - effective_prefill_tokens: None, - track_prefill_tokens: None, - }) - .await - .expect("cached reservation"); + core.create_reservation(replay_reservation("cached")) + .await + .expect("cached reservation"); let credited = wait_for_overlap(&core, || { let mut request = select_request(); request.prompt = prompt_b(); @@ -308,16 +332,9 @@ async fn bookings_populate_the_approximate_primary_without_kv_events() { ..PromptRequest::default() }; core.create_reservation(ReservationRequest { - model_name: "model".to_string(), - routing_group: "default".to_string(), - selection_id: "explicit".to_string(), worker_id: Some(1), - dp_rank: None, prompt: prompt_c(), - router_config_override: None, - expected_output_tokens: None, - effective_prefill_tokens: None, - track_prefill_tokens: None, + ..replay_reservation("explicit") }) .await .expect("explicit reservation"); @@ -336,26 +353,12 @@ async fn unreachable_remote_indexer_is_reported_not_ready() { let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind"); let base_url = format!("http://{}", listener.local_addr().expect("addr")); drop(listener); - let config = test_config(true); - let tracking_hash = Arc::new( - TrackingHashContext::from_config(&config).expect("valid tracking hash configuration"), - ); - let indexer_policy = IndexerPolicy::from_router_config(&config) - .expect("indexer policy") - .with_remote_indexer(base_url) - .expect("remote policy"); - let core = SelectionCore::new_inner( - config, - 1, - CancellationToken::new(), - None, - None, + let core = core_with( + test_config(true), SelectionHost::default(), + None, WorkerType::Aggregated, - true, - SelectionCacheConfig::default(), - tracking_hash, - indexer_policy, + Some(base_url), None, ); core.upsert_worker(worker(1)).await.expect("worker upsert"); @@ -394,26 +397,12 @@ async fn remote_indexer_serves_selection_without_local_kv_listeners() { // use_kv_events=true, but the primary is remote: workers need no // kv_events endpoints and no ZMQ listener is started here. - let config = test_config(true); - let tracking_hash = Arc::new( - TrackingHashContext::from_config(&config).expect("valid tracking hash configuration"), - ); - let indexer_policy = IndexerPolicy::from_router_config(&config) - .expect("indexer policy") - .with_remote_indexer(base_url) - .expect("remote policy"); - let core = SelectionCore::new_inner( - config, - 1, - CancellationToken::new(), - None, - None, + let core = core_with( + test_config(true), SelectionHost::default(), + None, WorkerType::Aggregated, - true, - SelectionCacheConfig::default(), - tracking_hash, - indexer_policy, + Some(base_url), None, ); assert!(!core.listens_for_kv_events); @@ -444,13 +433,7 @@ async fn hint_fixture( use crate::indexer::KvIndexerInterface; use crate::protocols::{BlockHashOptions, compute_block_hash_for_seq}; - let core = SelectionCore::try_new_local( - test_config(true), - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"); + let core = local_core(test_config(true)); for worker_id in [1, 2] { let mut request = worker_with_kv_events(worker_id); configure(&mut request); @@ -549,13 +532,7 @@ async fn router_hint_needs_capable_workers() { #[tokio::test] async fn event_driven_indexer_does_not_record_bookings() { - let core = SelectionCore::try_new_local( - test_config(true), - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"); + let core = local_core(test_config(true)); core.upsert_worker(worker_with_kv_events(1)) .await .expect("worker upsert"); @@ -677,28 +654,14 @@ async fn injected_lora_filter_narrows_candidates() { core.upsert_worker(worker(1)).await.expect("worker upsert"); core.upsert_worker(worker(2)).await.expect("worker upsert"); - // LoRA request: only the filter's worker is eligible. + // LoRA request: only the filter's worker is eligible. Allow-set and + // pinned-worker interplay is pinned by the `lora_filter` unit tests. for _ in 0..4 { let mut request = select_request(); request.prompt.lora_name = Some("adapter-a".to_string()); let response = core.select(request).await.expect("select"); assert_eq!(response.worker_id, 2); } - - // The filter never widens the caller's allow-set: an allow-set that - // excludes the filter's worker is preserved as-is. - let mut request = select_request(); - request.prompt.lora_name = Some("adapter-a".to_string()); - request.allowed_worker_ids = Some(HashSet::from([1])); - let response = core.select(request).await.expect("select"); - assert_eq!(response.worker_id, 1); - - // A pinned worker inside the universe survives the filter. - let mut request = select_request(); - request.prompt.lora_name = Some("adapter-a".to_string()); - request.pinned_worker = Some(WorkerWithDpRank::new(1, 0)); - let response = core.select(request).await.expect("select"); - assert_eq!(response.worker_id, 1); } #[tokio::test] @@ -805,13 +768,7 @@ async fn session_context_reaches_worker_selection() { #[tokio::test] async fn full_affinity_table_routes_without_pinning() { - let core = SelectionCore::try_new_local( - test_config(false), - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"); + let core = local_core(test_config(false)); core.upsert_worker(worker(1)).await.expect("worker upsert"); let entry = core .entry(&RoutingPartitionId::new("model", "default")) @@ -841,22 +798,34 @@ async fn full_affinity_table_routes_without_pinning() { ); } +/// Both `HostLoad` providers reach the scheduler; their semantics are pinned +/// by the `scheduling::queue` unit tests. #[tokio::test] -async fn injected_availability_provider_restricts_selection() { +async fn injected_load_providers_restrict_selection() { let available: Arc>>>> = Arc::new(parking_lot::Mutex::new(None)); let provider_state = Arc::clone(&available); let core = core_with_host(SelectionHost { load: HostLoad { available_workers: Some(Arc::new(move || provider_state.lock().clone())), + overloaded_workers: Some(Arc::new(|| Some(HashSet::from([1])))), ..HostLoad::default() }, ..SelectionHost::default() }); - core.upsert_worker(worker(1)).await.expect("worker upsert"); - core.upsert_worker(worker(2)).await.expect("worker upsert"); + for worker_id in [1, 2, 3] { + core.upsert_worker(worker(worker_id)) + .await + .expect("worker upsert"); + } - for only in [1, 2] { + // The overloaded worker is excluded while availability is unrestricted. + for _ in 0..4 { + let response = core.select(select_request()).await.expect("select"); + assert_ne!(response.worker_id, 1); + } + // The availability provider narrows the remaining candidates. + for only in [2, 3] { *available.lock() = Some(Arc::new(HashSet::from([only]))); for _ in 0..4 { let response = core.select(select_request()).await.expect("select"); @@ -865,34 +834,10 @@ async fn injected_availability_provider_restricts_selection() { } } -#[tokio::test] -async fn injected_overload_provider_excludes_worker() { - let core = core_with_host(SelectionHost { - load: HostLoad { - overloaded_workers: Some(Arc::new(|| Some(HashSet::from([1])))), - ..HostLoad::default() - }, - ..SelectionHost::default() - }); - core.upsert_worker(worker(1)).await.expect("worker upsert"); - core.upsert_worker(worker(2)).await.expect("worker upsert"); - - for _ in 0..4 { - let response = core.select(select_request()).await.expect("select"); - assert_eq!(response.worker_id, 2); - } -} - #[tokio::test] async fn shutdown_cancels_listeners_but_keeps_parent_alive() { let parent = CancellationToken::new(); - let core = SelectionCore::try_new_local( - test_config(true), - 1, - parent.clone(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"); + let core = local_core_with(test_config(true), 1, parent.clone()); let record = core .upsert_worker(worker_with_kv_events(1)) @@ -912,13 +857,7 @@ async fn shutdown_cancels_listeners_but_keeps_parent_alive() { #[case(2)] #[tokio::test] async fn selection_sees_cache_after_last_worker_replacement(#[case] indexer_threads: usize) { - let core = SelectionCore::try_new_local( - test_config(true), - indexer_threads, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"); + let core = local_core_with(test_config(true), indexer_threads, CancellationToken::new()); let key = RoutingPartitionId::new("model", "default"); let request = || { let mut request = select_request(); @@ -972,13 +911,7 @@ async fn selection_sees_cache_after_last_worker_replacement(#[case] indexer_thre #[tokio::test] async fn multi_rank_worker_with_replay_endpoint_is_incomplete() { - let core = SelectionCore::try_new_local( - test_config(true), - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"); + let core = local_core(test_config(true)); let record = core .upsert_worker(WorkerRequest { @@ -1000,13 +933,7 @@ async fn multi_rank_worker_with_replay_endpoint_is_incomplete() { async fn reupsert_recreates_a_listener_lost_to_a_cancelled_update() { use crate::indexer::KvIndexerInterface; - let core = SelectionCore::try_new_local( - test_config(true), - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"); + let core = local_core(test_config(true)); core.upsert_worker(worker_with_kv_events(1)) .await .expect("worker upsert"); @@ -1052,13 +979,7 @@ async fn reupsert_recreates_a_listener_lost_to_a_cancelled_update() { #[tokio::test] async fn upsert_moves_global_worker_id_between_routing_groups() { - let core = SelectionCore::try_new_local( - test_config(true), - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"); + let core = local_core(test_config(true)); let mut group_a = worker_with_kv_events(1); group_a.routing_group = "group-a".to_string(); core.upsert_worker(group_a).await.expect("group A upsert"); @@ -1107,13 +1028,7 @@ async fn upsert_moves_global_worker_id_between_routing_groups() { #[tokio::test] async fn shutdown_reports_not_ready_and_rejects_new_work() { - let core = SelectionCore::try_new_local( - test_config(false), - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"); + let core = local_core(test_config(false)); core.upsert_worker(worker(1)).await.expect("worker upsert"); assert!(core.ready().ready); @@ -1147,16 +1062,9 @@ async fn shutdown_reports_not_ready_and_rejects_new_work() { let reservation_error = core .create_reservation(ReservationRequest { - model_name: "model".to_string(), - routing_group: "default".to_string(), - selection_id: "res-after-shutdown".to_string(), worker_id: Some(1), - dp_rank: None, prompt: prompt(), - router_config_override: None, - expected_output_tokens: None, - effective_prefill_tokens: None, - track_prefill_tokens: None, + ..replay_reservation("res-after-shutdown") }) .await .expect_err("reservation should fail after shutdown"); @@ -1175,15 +1083,7 @@ async fn shutdown_reports_not_ready_and_rejects_new_work() { async fn queued_selection_errors_on_shutdown() { let mut config = test_config(false); config.router_queue_threshold = Some(0.0); - let core = Arc::new( - SelectionCore::try_new_local( - config, - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"), - ); + let core = Arc::new(local_core(config)); let record = core.upsert_worker(worker(1)).await.expect("worker upsert"); assert_eq!(record.lifecycle, WorkerLifecycle::Schedulable); @@ -1212,15 +1112,7 @@ async fn queued_selection_errors_on_shutdown() { async fn booking_is_freed_when_selected_worker_drained_while_queued() { let mut config = test_config(false); config.router_queue_threshold = Some(0.0); - let core = Arc::new( - SelectionCore::try_new_local( - config, - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"), - ); + let core = Arc::new(local_core(config)); core.upsert_worker(worker(1)).await.expect("worker upsert"); let key = RoutingPartitionId::new("model", "default"); let entry = core.entry(&key).expect("entry"); @@ -1260,15 +1152,7 @@ async fn booking_is_freed_when_selected_worker_drained_while_queued() { async fn lease_admission_keeps_a_selection_whose_worker_drained_while_queued() { let mut config = test_config(false); config.router_queue_threshold = Some(0.0); - let core = Arc::new( - SelectionCore::try_new_local( - config, - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"), - ); + let core = Arc::new(local_core(config)); core.upsert_worker(worker(1)).await.expect("worker upsert"); let key = RoutingPartitionId::new("model", "default"); let entry = core.entry(&key).expect("entry"); @@ -1333,13 +1217,7 @@ async fn lease_admission_keeps_a_selection_whose_worker_drained_while_queued() { /// `Book` feeds. #[tokio::test] async fn lease_admission_installs_no_index_row_and_records_nothing() { - let core = SelectionCore::try_new_local( - test_config(false), - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"); + let core = local_core(test_config(false)); core.upsert_worker(worker(1)).await.expect("worker upsert"); let key = RoutingPartitionId::new("model", "default"); let entry = core.entry(&key).expect("entry"); @@ -1400,13 +1278,7 @@ async fn lease_admission_installs_no_index_row_and_records_nothing() { #[tokio::test] async fn dropped_selection_future_frees_its_booking() { - let core = SelectionCore::try_new_local( - test_config(false), - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"); + let core = local_core(test_config(false)); core.upsert_worker(worker(1)).await.expect("worker upsert"); let key = RoutingPartitionId::new("model", "default"); let entry = core.entry(&key).expect("entry"); @@ -1491,17 +1363,8 @@ async fn dropped_book_selection_during_routing_record_frees_booking_and_claim() entered: AtomicBool::new(false), release: tokio::sync::Notify::new(), }); - let config = test_config(false); - let tracking_hash = Arc::new( - TrackingHashContext::from_config(&config).expect("valid tracking hash configuration"), - ); - let indexer_policy = IndexerPolicy::from_router_config(&config).expect("indexer policy"); - let core = SelectionCore::new_inner( - config, - 1, - CancellationToken::new(), - None, - None, + let core = core_with( + test_config(false), SelectionHost { cache: HostCache { index: KvIndexSource::Owned(Arc::new(PausedRecordIngress(record.clone()))), @@ -1509,11 +1372,9 @@ async fn dropped_book_selection_during_routing_record_frees_booking_and_claim() }, ..SelectionHost::default() }, + None, WorkerType::Aggregated, - true, - SelectionCacheConfig::default(), - tracking_hash, - indexer_policy, + None, Some( SessionAffinityConfig::new(Duration::from_secs(10)) .with_mode(SessionAffinityMode::Hard), @@ -1564,13 +1425,7 @@ async fn dropped_book_selection_during_routing_record_frees_booking_and_claim() #[tokio::test] async fn same_selection_id_in_two_partitions_is_a_conflict() { - let core = SelectionCore::try_new_local( - test_config(false), - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"); + let core = local_core(test_config(false)); for (worker_id, routing_group) in [(1, "group-a"), (2, "group-b")] { let mut request = worker(worker_id); request.routing_group = routing_group.to_string(); @@ -1604,29 +1459,16 @@ async fn same_selection_id_in_two_partitions_is_a_conflict() { #[tokio::test] async fn explicit_reservation_of_a_live_id_is_a_conflict() { - let core = SelectionCore::try_new_local( - test_config(false), - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"); + let core = local_core(test_config(false)); core.upsert_worker(worker(1)).await.expect("worker upsert"); core.select_and_reserve(reserve_request("taken")) .await .expect("booking"); let err = core .create_reservation(ReservationRequest { - model_name: "model".to_string(), - routing_group: "default".to_string(), - selection_id: "taken".to_string(), worker_id: Some(1), - dp_rank: None, prompt: prompt(), - router_config_override: None, - expected_output_tokens: None, - effective_prefill_tokens: None, - track_prefill_tokens: None, + ..replay_reservation("taken") }) .await .expect_err("explicit booking of a live id"); @@ -1694,15 +1536,7 @@ fn index_observer_replaces_stale_and_claimed_rows_of_its_partition() { async fn free_of_an_in_flight_reservation_is_not_found() { let mut config = test_config(false); config.router_queue_threshold = Some(0.0); - let core = Arc::new( - SelectionCore::try_new_local( - config, - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"), - ); + let core = Arc::new(local_core(config)); core.upsert_worker(worker(1)).await.expect("worker upsert"); core.select_and_reserve(reserve_request("res-a")) .await @@ -1737,13 +1571,7 @@ async fn free_of_an_in_flight_reservation_is_not_found() { async fn prefill_complete_is_idempotent_for_a_live_booking() { let mut config = test_config(false); config.router_track_prefill_tokens = true; - let core = SelectionCore::try_new_local( - config, - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"); + let core = local_core(config); core.upsert_worker(worker(1)).await.expect("worker upsert"); core.select_and_reserve(reserve_request("live")) .await @@ -1833,13 +1661,7 @@ async fn mirrored_replica_bookings_are_indexed_until_freed() { async fn lifecycle_operations_resolve_through_the_index() { let mut config = test_config(false); config.router_track_prefill_tokens = true; - let core = SelectionCore::try_new_local( - config, - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"); + let core = local_core(config); for (worker_id, routing_group) in [(1, "group-a"), (2, "group-b")] { let mut request = worker(worker_id); @@ -1893,13 +1715,7 @@ async fn advisory_select_reports_worker_load_and_busy_evaluation() { let mut config = test_config(false); config.conditional_disagg_prefill_busy_threshold = Some(0.5); config.conditional_disagg_decode_busy_threshold = Some(0.0); - let core = SelectionCore::try_new_local( - config, - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"); + let core = local_core(config); let mut request = worker(1); request.total_kv_blocks = Some(1000); core.upsert_worker(request).await.expect("worker upsert"); @@ -1937,13 +1753,7 @@ async fn advisory_select_reports_worker_load_and_busy_evaluation() { #[tokio::test] async fn busy_evaluation_is_absent_without_thresholds_or_capacity() { - let core = SelectionCore::try_new_local( - test_config(false), - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"); + let core = local_core(test_config(false)); core.upsert_worker(worker(1)).await.expect("worker upsert"); let mut request = select_request(); request.advisory = true; @@ -1956,13 +1766,7 @@ async fn busy_evaluation_is_absent_without_thresholds_or_capacity() { #[tokio::test] async fn reservation_index_tracks_bookings_until_freed() { - let core = SelectionCore::try_new_local( - test_config(false), - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"); + let core = local_core(test_config(false)); for (worker_id, routing_group) in [(1, "group-a"), (2, "group-b")] { let mut request = worker(worker_id); request.routing_group = routing_group.to_string(); @@ -1996,16 +1800,8 @@ async fn reservation_index_tracks_bookings_until_freed() { core.select(request).await.expect("select"); assert!(core.reservation_index.read().get("cached").is_none()); core.create_reservation(ReservationRequest { - model_name: "model".to_string(), routing_group: "group-b".to_string(), - selection_id: "cached".to_string(), - worker_id: None, - dp_rank: None, - prompt: PromptRequest::default(), - router_config_override: None, - expected_output_tokens: None, - effective_prefill_tokens: None, - track_prefill_tokens: None, + ..replay_reservation("cached") }) .await .expect("cached reservation"); @@ -2036,13 +1832,7 @@ async fn reservation_index_tracks_bookings_until_freed() { #[tokio::test] async fn reservation_index_sweep_drops_bookings_released_out_of_band() { - let core = SelectionCore::try_new_local( - test_config(false), - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"); + let core = local_core(test_config(false)); core.upsert_worker(worker(1)).await.expect("worker upsert"); core.select_and_reserve(reserve_request("live")) .await @@ -2075,15 +1865,7 @@ async fn reservation_index_sweep_drops_bookings_released_out_of_band() { async fn lifecycle_lookup_does_not_nest_reservation_index_inside_entries() { // Three parties: a sweep holding `entries` and wanting `reservation_index`, // a lifecycle call, and a partition creation queued on `entries.write()`. - let core = Arc::new( - SelectionCore::try_new_local( - test_config(false), - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"), - ); + let core = Arc::new(local_core(test_config(false))); core.upsert_worker(worker(1)).await.expect("worker upsert"); core.select_and_reserve(reserve_request("live")) .await @@ -2134,15 +1916,7 @@ async fn lifecycle_lookup_does_not_nest_reservation_index_inside_entries() { async fn queued_selection_returns_refreshed_overlap_snapshot() { let mut config = test_config(false); config.router_queue_threshold = Some(0.0); - let core = Arc::new( - SelectionCore::try_new_local( - config, - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .expect("valid test config"), - ); + let core = Arc::new(local_core(config)); for worker_id in [1, 2] { let mut request = worker(worker_id); @@ -2160,9 +1934,6 @@ async fn queued_selection_returns_refreshed_overlap_snapshot() { for worker_id in [1, 2] { core.create_reservation(ReservationRequest { - model_name: "model".to_string(), - routing_group: "default".to_string(), - selection_id: format!("occupy-{worker_id}"), worker_id: Some(worker_id), dp_rank: Some(0), prompt: PromptRequest { @@ -2176,10 +1947,8 @@ async fn queued_selection_returns_refreshed_overlap_snapshot() { cache_namespace: None, is_eagle: None, }, - router_config_override: None, - expected_output_tokens: None, effective_prefill_tokens: Some(8), - track_prefill_tokens: None, + ..replay_reservation(&format!("occupy-{worker_id}")) }) .await .expect("occupy worker"); @@ -2241,23 +2010,12 @@ async fn queued_selection_returns_refreshed_overlap_snapshot() { } fn core_with_session_affinity_mode(mode: SessionAffinityMode) -> SelectionCore { - let config = test_config(false); - let tracking_hash = Arc::new( - TrackingHashContext::from_config(&config).expect("valid tracking hash configuration"), - ); - let indexer_policy = IndexerPolicy::from_router_config(&config).expect("indexer policy"); - SelectionCore::new_inner( - config, - 1, - CancellationToken::new(), - None, - None, + core_with( + test_config(false), SelectionHost::default(), + None, WorkerType::Aggregated, - true, - SelectionCacheConfig::default(), - tracking_hash, - indexer_policy, + None, Some(SessionAffinityConfig::new(Duration::from_secs(10)).with_mode(mode)), ) } @@ -2329,7 +2087,8 @@ async fn session_worker_departing_after_the_hold_reinitializes_the_session() { #[tokio::test] async fn concurrent_holds_on_a_departed_worker_both_land_on_the_replacement() { - let table = SessionAffinity::new(Duration::from_secs(60)).expect("affinity table"); + let table = SessionAffinity::with_config(SessionAffinityConfig::new(Duration::from_secs(60))) + .expect("affinity table"); let departed = WorkerAffinityTarget::new(1, Some(0)); let replacement = WorkerAffinityTarget::new(2, Some(0)); let Hold::Initialize(init) = table.acquire("s", None).await.expect("acquire") else { @@ -2378,21 +2137,6 @@ async fn concurrent_holds_on_a_departed_worker_both_land_on_the_replacement() { ); } -fn replay_reservation(selection_id: &str) -> ReservationRequest { - ReservationRequest { - model_name: "model".to_string(), - routing_group: "default".to_string(), - selection_id: selection_id.to_string(), - worker_id: None, - dp_rank: None, - prompt: PromptRequest::default(), - router_config_override: None, - expected_output_tokens: None, - effective_prefill_tokens: None, - track_prefill_tokens: None, - } -} - #[tokio::test] async fn two_phase_reservation_binds_the_session() { let core = core_with_session_affinity(); @@ -2509,34 +2253,6 @@ async fn soft_mode_follows_the_dispatch() { assert_eq!(bound_worker(&core, "s"), Some(other)); } -#[tokio::test] -async fn session_stays_on_its_first_worker_across_bookings() { - let core = core_with_session_affinity(); - core.upsert_worker(worker(1)).await.expect("worker upsert"); - core.upsert_worker(worker(2)).await.expect("worker upsert"); - - let first = core - .select_and_reserve(session_reservation("r1", "chat-a")) - .await - .expect("first booking"); - for index in 0..4 { - let response = core - .select_and_reserve(session_reservation(&format!("r-{index}"), "chat-a")) - .await - .expect("booking"); - assert_eq!(response.worker_id, first.worker_id, "session must stay put"); - core.free_reservation(&format!("r-{index}")) - .await - .expect("free"); - } - // A read-only select sees the binding too. - let mut advisory = select_request(); - advisory.session_id = Some("chat-a".to_string()); - let response = core.select(advisory).await.expect("select"); - assert_eq!(response.worker_id, first.worker_id); - core.free_reservation("r1").await.expect("free"); -} - #[tokio::test] async fn replicated_binding_steers_a_new_session_and_frees_with_the_booking() { let core = core_with_session_affinity(); @@ -2638,13 +2354,7 @@ async fn partition_sessions_keep_independent_bindings() { async fn rejoined_worker_feeds_partition_index() { use crate::protocols::{BlockHashOptions, compute_block_hash_for_seq}; for update_only in [false, true] { - let core = SelectionCore::try_new_local( - test_config(true), - 1, - CancellationToken::new(), - SelectionCacheConfig::default(), - ) - .unwrap(); + let core = local_core(test_config(true)); let request = worker_with_kv_events(1); core.upsert_worker(request.clone()).await.unwrap(); let key = RoutingPartitionId::new("model", "default"); diff --git a/lib/kv-router/src/services/selection/mod.rs b/lib/kv-router/src/services/selection/mod.rs index 49f8f754b15c..f60883cdbd18 100644 --- a/lib/kv-router/src/services/selection/mod.rs +++ b/lib/kv-router/src/services/selection/mod.rs @@ -25,16 +25,14 @@ mod tests; pub use crate::WorkerSelectionPolicyFactory; pub use crate::services::common::replica_sync::{ HostReplicaChannels, HostReplicaSyncFactory, ReplicaPeerError, SchedulerLoadSink, - ScopedSequencePublisher, }; pub use core::{ - HostCache, HostEligibility, HostLoad, HostReplication, HostTelemetry, KvIndexSource, - LookupTimings, Selected, SelectionAdmission, SelectionCore, SelectionHost, SelectionOperation, - SelectionOutcome, SelectionPartition, SelectionRun, SelectionScheduler, SelectionServiceConfig, - SessionBinding, + HostCache, HostEligibility, HostLoad, HostReplication, HostTelemetry, KvIndexSource, Selected, + SelectionAdmission, SelectionCore, SelectionHost, SelectionOperation, SelectionOutcome, + SelectionPartition, SelectionRun, SelectionServiceConfig, SessionBinding, }; pub use error::SelectionError; -pub use ingress::{KvEventIngress, ZmqDirectIngress}; +pub use ingress::KvEventIngress; pub use input::{PromptRequest, PromptView}; pub use membership::{CatalogObserver, CatalogReconciler, WorkerCatalogSource}; pub use pending::SelectionCacheConfig; diff --git a/lib/kv-router/src/services/selection/service.rs b/lib/kv-router/src/services/selection/service.rs index 280016a2b91c..a84f46a87924 100644 --- a/lib/kv-router/src/services/selection/service.rs +++ b/lib/kv-router/src/services/selection/service.rs @@ -339,13 +339,6 @@ impl SelectionService { &self.core } - /// Scheduler and indexer handle for `key`, once a worker has been upserted - /// into that partition. Embedding hosts that keep their own index drive - /// the partition scheduler through this handle. - pub fn partition(&self, key: &RoutingPartitionId) -> Option { - self.core.partition(key) - } - /// Create (or fetch) the partition for `key` ahead of worker registration. pub fn ensure_partition( &self, @@ -496,10 +489,6 @@ impl SelectionService { self.core.recover_indexer_from_peers(peers).await } - pub async fn cancelled(&self) { - self.cancel_token.cancelled().await; - } - pub async fn shutdown(&self) { self.cancel_token.cancel(); self.core.shutdown(); diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index 2bf99ffad271..6c8268fad9f6 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -25,9 +25,10 @@ use dynamo_kv_router::{ WorkerWithDpRank, compute_block_hash_for_seq, }, scheduling::{ - CacheHitEstimates, OverlapAnalysis, OverloadedWorkerProvider, WorkerAvailabilityProvider, - effective_prefill_tokens, overlap::cache_hit_estimates_from_tiered_matches, - queue::BookingHandle, + CacheHitEstimates, OverlapAnalysis, OverloadedWorkerProvider, PotentialLoad, + WorkerAvailabilityProvider, effective_prefill_tokens, + overlap::cache_hit_estimates_from_tiered_matches, + queue::{BookingHandle, SchedulerBookingDescriptor}, }, selector::WorkerInputs, services::selection::{ @@ -38,7 +39,6 @@ use dynamo_kv_router::{ use dynamo_runtime::{ CancellationToken, component::{Client, Endpoint}, - discovery::DiscoveryQuery, error::{DynamoError, ErrorType}, pipeline::{ AsyncEngine, AsyncEngineContextProvider, Error, ManyOut, ResponseStream, SingleIn, @@ -51,10 +51,8 @@ use dynamo_runtime::{ use futures::stream; // Re-export from dynamo-kv-router crate -pub use dynamo_kv_router::approx; pub use dynamo_kv_router::protocols; pub use dynamo_kv_router::scheduling; -pub use dynamo_kv_router::selector; pub(crate) mod embedded; pub mod encoder_router; @@ -66,16 +64,13 @@ pub mod publisher; mod request_lease; mod routing_host; pub(crate) mod routing_load; -pub mod scheduler; pub mod sequence; pub mod shared_cache; -pub use dynamo_kv_router::scheduling::{ - OverlapScoresResponse, SharedCacheOverlapScore, WorkerOverlapScore, -}; +pub use dynamo_kv_router::scheduling::OverlapScoresResponse; pub use embedded::{install_worker_selection_policy_registry, worker_selection_policy_registry}; pub use encoder_router::EncoderRouter; -pub use indexer::{Indexer, ServedIndexerHandle, ServedIndexerMode, ensure_served_indexer_service}; +pub use indexer::Indexer; pub use prefill_router::PrefillRouter; pub use routing_host::{KvPushRouter, RoutingHost}; pub use routing_load::{ @@ -84,10 +79,7 @@ pub use routing_load::{ use crate::{ discovery::{KvSourceMembershipWatch, RuntimeConfigWatch}, - kv_router::{ - scheduler::PotentialLoad, - sequence::{SequenceError, SequenceRequest}, - }, + kv_router::sequence::{SequenceError, SequenceRequest}, local_model::runtime_config::ModelRuntimeConfig, worker_type::WorkerType, }; @@ -477,24 +469,13 @@ fn cache_hit_for_worker( } } -// [gluo TODO] shouldn't need to be public -// this should be discovered from the component - -// for metric scraping (pull-based) -pub const KV_METRICS_ENDPOINT: &str = "load_metrics"; - // for metric publishing (push-based) pub const KV_METRICS_SUBJECT: &str = "kv_metrics"; pub const MULTIMODAL_EMBEDDING_CACHE_SUBJECT: &str = "multimodal_embedding_cache"; // for inter-router comms -pub const PREFILL_SUBJECT: &str = "prefill_events"; pub const ACTIVE_SEQUENCES_SUBJECT: &str = "active_sequences_events"; -// for radix tree snapshot storage -pub const RADIX_STATE_BUCKET: &str = "radix-bucket"; -pub const RADIX_STATE_FILE: &str = "radix-state"; - // for worker-local kvindexer query pub const WORKER_KV_INDEXER_BUFFER_SIZE: usize = 1024; // store 1024 most recent events in worker buffer @@ -548,15 +529,6 @@ pub fn router_endpoint_id(namespace: String, component: String) -> EndpointId { } } -/// Creates a DiscoveryQuery for the KV router in the given namespace. -pub fn router_discovery_query(namespace: String, component: String) -> DiscoveryQuery { - DiscoveryQuery::Endpoint { - namespace, - component, - endpoint: KV_ROUTER_ENDPOINT.to_string(), - } -} - /// A KvRouter only decides which worker you should use. It doesn't send you there. /// TODO: Rename this to indicate it only selects a worker, it does not route. pub struct KvRouter { @@ -884,12 +856,6 @@ impl KvRouter { self.required_worker_inputs } - /// Cancel background work and wait for KV event ingestion to stop. - pub async fn shutdown(self) { - self.cancellation_token.cancel(); - self.ingress.shutdown().await; - } - pub fn is_eagle(&self) -> bool { self.is_eagle } @@ -1096,99 +1062,6 @@ impl KvRouter { /// that exact worker/rank. /// /// When `allowed_worker_ids` is Some, only workers in that set are considered for selection. - #[allow(clippy::too_many_arguments)] - pub async fn find_best_match_details( - &self, - context_id: Option<&str>, - tokens: &[u32], - block_mm_infos: Option<&[Option]>, - router_config_override: Option<&RouterConfigOverride>, - update_states: bool, - return_routing_hashes: bool, - lora_name: Option, - cache_namespace: Option, - priority_jump: f64, - strict_priority: u32, - expected_output_tokens: Option, - pinned_worker: Option, - allowed_worker_ids: Option>, - routing_constraints: RoutingConstraints, - ) -> anyhow::Result { - self.find_best_match_details_with_policy_class( - context_id, - tokens, - block_mm_infos, - router_config_override, - update_states, - return_routing_hashes, - lora_name, - cache_namespace, - priority_jump, - strict_priority, - None, - None, - expected_output_tokens, - pinned_worker, - allowed_worker_ids, - routing_constraints, - ) - .await - } - - /// Admit a best-match request with scheduler-owned cancellation cleanup. - /// - /// If the future is dropped before it returns, the scheduler retracts any - /// pending or admitted booking. After success, the returned outcome carries - /// the tracked attempt identity required for exact lifecycle cleanup. - #[allow(clippy::too_many_arguments)] - pub(crate) async fn find_best_match_details_with_lifecycle( - &self, - context_id: Option<&str>, - tokens: &[u32], - block_mm_infos: Option<&[Option]>, - router_config_override: Option<&RouterConfigOverride>, - update_states: bool, - return_routing_hashes: bool, - lora_name: Option, - cache_namespace: Option, - priority_jump: f64, - strict_priority: u32, - policy_class: Option, - expected_output_tokens: Option, - pinned_worker: Option, - allowed_worker_ids: Option>, - routing_constraints: RoutingConstraints, - ) -> anyhow::Result { - match self - .find_best_match_details_with_policy_class_inner( - context_id, - tokens, - block_mm_infos, - router_config_override, - update_states, - return_routing_hashes, - lora_name, - cache_namespace, - priority_jump, - strict_priority, - policy_class, - None, - expected_output_tokens, - None, - pinned_worker, - allowed_worker_ids, - routing_constraints, - FindBestMatchAdmission::WithAdmission, - ) - .await? - { - FindBestMatchInnerOutcome::WithAdmission(outcome) => Ok(outcome), - FindBestMatchInnerOutcome::WithoutAdmission(_) => { - unreachable!("with-admission routing returned advisory outcome") - } - } - } - #[allow(clippy::too_many_arguments)] pub async fn find_best_match_details_with_policy_class( &self, @@ -1466,52 +1339,6 @@ impl KvRouter { }) } - /// Give these tokens, find the worker with the best match in its KV cache. - /// Returns the best worker (with dp_rank) and approximate effective overlap in blocks. - #[allow(clippy::too_many_arguments)] - pub async fn find_best_match( - &self, - context_id: Option<&str>, - tokens: &[u32], - block_mm_infos: Option<&[Option]>, - router_config_override: Option<&RouterConfigOverride>, - update_states: bool, - lora_name: Option, - cache_namespace: Option, - priority_jump: f64, - strict_priority: u32, - expected_output_tokens: Option, - allowed_worker_ids: Option>, - routing_constraints: RoutingConstraints, - ) -> anyhow::Result<(WorkerWithDpRank, u32)> { - let result = self - .find_best_match_details( - context_id, - tokens, - block_mm_infos, - router_config_override, - update_states, - false, - lora_name, - cache_namespace, - priority_jump, - strict_priority, - expected_output_tokens, - None, - allowed_worker_ids, - routing_constraints, - ) - .await?; - match result { - FindBestMatchOutcome::Routed { - worker, - overlap_blocks, - .. - } => Ok((worker, overlap_blocks)), - FindBestMatchOutcome::QueueRejected { rejection } => Err(rejection.into()), - } - } - #[allow(clippy::too_many_arguments)] pub async fn add_request( &self, @@ -1568,7 +1395,7 @@ impl KvRouter { }; self.request_leases .register_detached( - scheduler::SchedulerBookingDescriptor { + SchedulerBookingDescriptor { request_id, worker, attempt_id, @@ -1605,7 +1432,7 @@ impl KvRouter { pub(crate) async fn mark_prefill_completed_if_booking( &self, - booking: &scheduler::SchedulerBookingDescriptor, + booking: &SchedulerBookingDescriptor, ) -> Result<(), KvSchedulerError> { self.scheduler .mark_prefill_completed_if_booking(booking) @@ -1673,17 +1500,9 @@ impl KvRouter { (config.data_parallel_size == 1).then_some(config.data_parallel_start_rank) } - pub fn add_output_block( - &self, - request_id: &str, - decay_fraction: Option, - ) -> Result<(), SequenceError> { - self.scheduler.add_output_block(request_id, decay_fraction) - } - pub(crate) async fn enqueue_output_block_if_booking( &self, - booking: &scheduler::SchedulerBookingDescriptor, + booking: &SchedulerBookingDescriptor, decay_fraction: Option, ) -> Result<(), KvSchedulerError> { self.scheduler @@ -1719,27 +1538,6 @@ impl KvRouter { lora_name: Option<&str>, cache_namespace: Option<&str>, ) -> Result { - self.get_cache_hit_estimate_with_hashes( - tokens, - block_mm_infos, - worker, - lora_name, - cache_namespace, - false, - ) - .await - .map(|(estimate, _)| estimate) - } - - pub(crate) async fn get_cache_hit_estimate_with_hashes( - &self, - tokens: &[u32], - block_mm_infos: Option<&[Option]>, - worker: WorkerWithDpRank, - lora_name: Option<&str>, - cache_namespace: Option<&str>, - return_routing_hashes: bool, - ) -> Result<(WorkerCacheHitEstimate, Option), KvRouterError> { let block_hashes = compute_block_hash_for_seq( tokens, self.block_size, @@ -1750,20 +1548,9 @@ impl KvRouter { is_eagle: Some(self.is_eagle), }, ); - let (tiered_matches, routing_hashes) = if return_routing_hashes { - let tiered_matches = self.indexer.find_matches_by_tier_ref(&block_hashes).await?; - ( - tiered_matches, - Some(RoutingDecisionHashes::from_local_hashes(block_hashes)), - ) - } else { - (self.indexer.find_matches_by_tier(block_hashes).await?, None) - }; + let tiered_matches = self.indexer.find_matches_by_tier(block_hashes).await?; let cache_hit_estimates = self.cache_hit_estimates_from_tiered_matches(&tiered_matches); - Ok(( - self.cache_hit_for_worker(&cache_hit_estimates, worker), - routing_hashes, - )) + Ok(self.cache_hit_for_worker(&cache_hit_estimates, worker)) } /// Get potential prefill and decode loads for all workers @@ -2017,15 +1804,10 @@ mod tests { use std::collections::{HashMap, HashSet}; use async_trait::async_trait; - use dynamo_kv_router::{ - WorkerSelectionPolicyError, - indexer::{LowerTierMatchDetails, MatchDetails}, - protocols::{OverlapScores, StorageTier, compute_seq_hash_for_block}, - }; + use dynamo_kv_router::{WorkerSelectionPolicyError, protocols::compute_seq_hash_for_block}; use dynamo_runtime::{DistributedRuntime, Runtime, distributed::DistributedConfig}; use tokio::sync::watch; - use crate::kv_router::scheduler::KvSchedulerError; use crate::local_model::runtime_config::ModelRuntimeConfig; use dynamo_kv_router::selector::{ WorkerCandidate, WorkerFilter, WorkerInputView, WorkerPicker, WorkerSelectionContext, @@ -2189,47 +1971,6 @@ mod tests { } } - #[test] - fn weighted_cache_hit_estimates_include_lower_tiers() { - let worker_1 = WorkerWithDpRank::new(1, 0); - let worker_2 = WorkerWithDpRank::new(2, 0); - let mut device_overlap_scores = OverlapScores::new(); - device_overlap_scores.scores.insert(worker_1, 2); - let mut host_match_details = LowerTierMatchDetails::default(); - host_match_details.hits.insert(worker_1, 1); - host_match_details.hits.insert(worker_2, 1); - let mut disk_match_details = LowerTierMatchDetails::default(); - disk_match_details.hits.insert(worker_1, 2); - - let tiered_matches = indexer::TieredMatchDetails { - device: MatchDetails { - overlap_scores: device_overlap_scores, - ..Default::default() - }, - lower_tier: HashMap::from([ - (StorageTier::HostPinned, host_match_details), - (StorageTier::Disk, disk_match_details), - ]), - }; - - let estimates = cache_hit_estimates_from_tiered_matches( - &KvRouterConfig::default(), - 16, - &tiered_matches, - ); - - assert_eq!( - estimates.effective_overlap_blocks.get(&worker_1), - Some(&3.25) - ); - assert_eq!(estimates.cached_tokens.get(&worker_1), Some(&52)); - assert_eq!( - estimates.effective_overlap_blocks.get(&worker_2), - Some(&0.75) - ); - assert_eq!(estimates.cached_tokens.get(&worker_2), Some(&12)); - } - struct FakeSharedCache { hits: Option, should_error: bool, @@ -2438,44 +2179,86 @@ mod tests { )); } - /// Three default-config workers under the registry policy, with - /// `router_track_active_blocks` on so bookings send tracking hashes. - async fn differential_router(name: &str) -> KvRouter { + /// One block-size-2 router over `workers`, no KV event subscription. + async fn make_router( + name: &str, + workers: HashMap, + policy: SelectionPolicySource, + shared_cache: Option>, + config: KvRouterConfig, + ) -> KvRouter { let component = make_test_component(name).await; let endpoint = component.endpoint("backend"); let client = endpoint.client().await.unwrap(); - let workers = (0..3) - .map(|worker_id| (worker_id, ModelRuntimeConfig::default())) - .collect(); let (_tx, rx) = watch::channel(workers); - // Prefill load decays with wall time and would let near-ties flip - // between two procedures run microseconds apart. - let config = KvRouterConfig { - use_kv_events: false, - router_track_prefill_tokens: false, - skip_initial_worker_wait: true, - ..Default::default() - }; - let router = KvRouter::new_with_worker_role( + KvRouter::new_with_worker_role( endpoint, client, rx, None, 2, - SelectionPolicySource::Registry, + policy, Some(config), None, None, "decode", None, false, - None, + shared_cache, None, ) .await - .unwrap(); - // Worker 1 has served prefix A and worker 2 prefix B, so every corpus - // prompt has one unambiguous best worker. + .unwrap() + } + + /// Three default-config workers under the registry policy, with + /// `router_track_active_blocks` on so bookings send tracking hashes. + async fn tracked_router(name: &str) -> KvRouter { + // Prefill load decays with wall time and would let near-ties flip + // between two runs microseconds apart. + let config = KvRouterConfig { + use_kv_events: false, + router_track_prefill_tokens: false, + skip_initial_worker_wait: true, + ..Default::default() + }; + let workers = (0..3) + .map(|worker_id| (worker_id, ModelRuntimeConfig::default())) + .collect(); + make_router(name, workers, SelectionPolicySource::Registry, None, config).await + } + + /// Advisory best-match query with default routing arguments. + async fn find_best_match( + router: &KvRouter, + tokens: &[u32], + return_routing_hashes: bool, + ) -> anyhow::Result { + router + .find_best_match_details_with_policy_class( + None, + tokens, + None, + None, + false, + return_routing_hashes, + None, + None, + 0.0, + 0, + None, + None, + None, + None, + None, + RoutingConstraints::default(), + ) + .await + } + + /// Worker 1 has served prefix A and worker 2 prefix B, so every corpus + /// prompt has one unambiguous best worker. + async fn seed_golden_prefixes(router: &KvRouter) { for (worker_id, tokens) in [(1u64, golden_prefix(1)), (2, golden_prefix(2))] { router .record_routing_decision( @@ -2488,24 +2271,15 @@ mod tests { // until a query sees the whole prefix. tokio::time::timeout(std::time::Duration::from_secs(5), async { loop { - let (worker, overlap) = router - .find_best_match( - None, - &tokens, - None, - None, - false, - None, - None, - 0.0, - 0, - None, - None, - RoutingConstraints::default(), - ) - .await - .unwrap(); - if worker.worker_id == worker_id && overlap == 8 { + let FindBestMatchOutcome::Routed { + worker, + overlap_blocks, + .. + } = find_best_match(router, &tokens, false).await.unwrap() + else { + panic!("seeding query must route"); + }; + if worker.worker_id == worker_id && overlap_blocks == 8 { break; } tokio::task::yield_now().await; @@ -2514,7 +2288,6 @@ mod tests { .await .expect("seeded prefix became visible"); } - router } fn golden_prefix(seed: u32) -> Vec { @@ -2579,7 +2352,8 @@ mod tests { /// intended behavior change. #[tokio::test] async fn selection_matches_frozen_frontend_trace() { - let router = differential_router("golden").await; + let router = tracked_router("golden").await; + seed_golden_prefixes(&router).await; let mut trace = Vec::new(); for (index, tokens) in golden_corpus().iter().enumerate() { let context_id = format!("golden-{index}"); @@ -2637,53 +2411,6 @@ mod tests { ); } - /// A dropped tracked selection frees its booking; a committed one keeps it. - #[tokio::test] - async fn dropped_selection_lease_frees_the_booking() { - let router = differential_router("lease").await; - let prompt = golden_prefix(1); - let select = |context_id: &'static str| { - router.find_best_match_details_with_policy_class_inner( - Some(context_id), - &prompt, - None, - None, - true, - false, - None, - None, - 0.0, - 0, - None, - None, - None, - None, - None, - None, - RoutingConstraints::default(), - FindBestMatchAdmission::WithAdmission, - ) - }; - let FindBestMatchInnerOutcome::WithAdmission(dropped) = select("dropped").await.unwrap() - else { - panic!("admitted routing returned advisory outcome"); - }; - drop(dropped); - let FindBestMatchInnerOutcome::WithAdmission(kept) = select("kept").await.unwrap() else { - panic!("admitted routing returned advisory outcome"); - }; - let _booking = kept.booking.map(BookingHandle::commit).expect("booking"); - - tokio::time::timeout(std::time::Duration::from_secs(2), async { - while router.scheduler.has_request("dropped") { - tokio::task::yield_now().await; - } - }) - .await - .expect("dropped selection released its booking"); - assert!(router.scheduler.has_request("kept")); - } - /// A lease released before `set_scheduler` is a wiring bug: it logs and /// leaves the booking alone. After `set_scheduler`, the release reaches /// the scheduler's booking cleanup. @@ -2691,7 +2418,7 @@ mod tests { async fn request_lease_manager_releases_bookings_only_once_scheduler_is_set() { use dynamo_kv_router::scheduling::queue::SchedulerBookingDescriptor; - let router = differential_router("lease-manager").await; + let router = tracked_router("lease-manager").await; let worker = WorkerWithDpRank::from_worker_id(1); let scheduler = &router.scheduler; let book = |request_id: &'static str| async move { @@ -2769,7 +2496,7 @@ mod tests { } } - let mut router = differential_router("enroll").await; + let mut router = tracked_router("enroll").await; let record = Arc::new(PausedRecord { entered: AtomicBool::new(false), release: tokio::sync::Notify::new(), @@ -2831,16 +2558,11 @@ mod tests { .expect("cancelled enrollment released its booking"); } - async fn make_test_router_with_workers( + /// Two default-config workers, deterministic scoring, untracked bookings. + async fn make_test_router( policy: SelectionPolicySource, shared_cache: Option>, - workers: HashMap, ) -> KvRouter { - let component = make_test_component("shared-cache-router").await; - let endpoint = component.endpoint("backend"); - let client = endpoint.client().await.unwrap(); - let (_tx, rx) = watch::channel(workers); - let config = KvRouterConfig { overlap_score_credit: 0.0, router_temperature: 0.0, @@ -2850,35 +2572,11 @@ mod tests { skip_initial_worker_wait: true, ..Default::default() }; - - KvRouter::new_with_worker_role( - endpoint, - client, - rx, - None, - 2, - policy, - Some(config), - None, - None, - "decode", - None, - false, - shared_cache, - None, - ) - .await - .unwrap() - } - - async fn make_test_router( - policy: SelectionPolicySource, - shared_cache: Option>, - ) -> KvRouter { - let mut workers = HashMap::new(); - workers.insert(0, ModelRuntimeConfig::default()); - workers.insert(1, ModelRuntimeConfig::default()); - make_test_router_with_workers(policy, shared_cache, workers).await + let workers = HashMap::from([ + (0, ModelRuntimeConfig::default()), + (1, ModelRuntimeConfig::default()), + ]); + make_router("shared-cache-router", workers, policy, shared_cache, config).await } #[tokio::test] @@ -2895,26 +2593,19 @@ mod tests { ) .await; - let (worker, overlap) = router - .find_best_match( - None, - &[11, 12, 21, 22], - None, - None, - false, - None, - None, - 0.0, - 0, - None, - None, - RoutingConstraints::default(), - ) + let FindBestMatchOutcome::Routed { + worker, + overlap_blocks, + .. + } = find_best_match(&router, &[11, 12, 21, 22], false) .await - .unwrap(); + .unwrap() + else { + panic!("expected routed outcome"); + }; assert_eq!(worker, WorkerWithDpRank::from_worker_id(1)); - assert_eq!(overlap, 0); + assert_eq!(overlap_blocks, 0); } #[tokio::test] @@ -2928,26 +2619,19 @@ mod tests { ) .await; - let (worker, overlap) = router - .find_best_match( - None, - &[11, 12, 21, 22], - None, - None, - false, - None, - None, - 0.0, - 0, - None, - None, - RoutingConstraints::default(), - ) + let FindBestMatchOutcome::Routed { + worker, + overlap_blocks, + .. + } = find_best_match(&router, &[11, 12, 21, 22], false) .await - .unwrap(); + .unwrap() + else { + panic!("expected routed outcome"); + }; assert_eq!(worker, WorkerWithDpRank::from_worker_id(0)); - assert_eq!(overlap, 0); + assert_eq!(overlap_blocks, 0); } #[tokio::test] @@ -2963,23 +2647,9 @@ mod tests { })); let router = make_test_router(policy, None).await; - let err = router - .find_best_match( - None, - &[11, 12], - None, - None, - false, - None, - None, - 0.0, - 0, - None, - None, - RoutingConstraints::default(), - ) - .await - .unwrap_err(); + let Err(err) = find_best_match(&router, &[11, 12], false).await else { + panic!("filtered workers must not route"); + }; assert!(dynamo_runtime::error::match_error_chain( err.as_ref(), @@ -2997,25 +2667,7 @@ mod tests { .await; let tokens = [11, 12, 21, 22]; - let outcome = router - .find_best_match_details( - None, - &tokens, - None, - None, - false, - true, - None, - None, - 0.0, - 0, - None, - None, - None, - RoutingConstraints::default(), - ) - .await - .unwrap(); + let outcome = find_best_match(&router, &tokens, true).await.unwrap(); let FindBestMatchOutcome::Routed { routing_hashes: Some(hashes), @@ -3040,40 +2692,6 @@ mod tests { assert_eq!(hashes.sequence_hashes, expected_sequence); } - #[tokio::test] - async fn test_find_best_match_details_omits_routing_hashes_when_not_requested() { - let router = make_test_router( - fixed_policy(None, WorkerWithDpRank::from_worker_id(0)), - None, - ) - .await; - - let outcome = router - .find_best_match_details( - None, - &[11, 12, 21, 22], - None, - None, - false, - false, - None, - None, - 0.0, - 0, - None, - None, - None, - RoutingConstraints::default(), - ) - .await - .unwrap(); - - let FindBestMatchOutcome::Routed { routing_hashes, .. } = outcome else { - panic!("expected routed outcome"); - }; - assert!(routing_hashes.is_none()); - } - #[tokio::test] async fn test_get_overlap_scores_returns_tiered_rows_and_shared_hits() { let router = make_test_router( diff --git a/lib/llm/src/kv_router/AGENTS.md b/lib/llm/src/kv_router/AGENTS.md index 2f860e620198..279e6958cac1 100644 --- a/lib/llm/src/kv_router/AGENTS.md +++ b/lib/llm/src/kv_router/AGENTS.md @@ -3,7 +3,7 @@ ## Module map - `kv_router.rs` estimates how much of a request each worker already has cached, then asks the scheduler to choose a worker. It selects a worker but does not send the request. -- `embedded.rs` connects discovered workers and their current load to the scheduling code in `lib/kv-router`; `scheduler.rs` only re-exports from it. +- `embedded.rs` connects discovered workers and their current load to the scheduling code in `lib/kv-router`. - `indexer/` and `route_lookup.rs` track which KV-cache blocks each worker holds and look for reusable blocks before selection. - `publisher/` receives KV-cache events and worker metrics from inference engines. - `routing_host.rs` owns selection, dispatch, and response cleanup. `routing_host/kv_selection.rs` chooses KV-routed workers, `routing_host/request_guard.rs` tracks progress and cleanup, and `routing_host/cancellation.rs` stops unfinished work when the client cancels. diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs index 6ecfc33da4af..6f2e155d3a63 100644 --- a/lib/llm/src/kv_router/embedded.rs +++ b/lib/llm/src/kv_router/embedded.rs @@ -471,16 +471,6 @@ impl EmbeddedSelection { self.worker_type } - pub(crate) fn add_output_block( - &self, - request_id: &str, - decay_fraction: Option, - ) -> Result<(), SequenceError> { - self.partition - .scheduler() - .add_output_block(request_id, decay_fraction) - } - pub(crate) async fn enqueue_output_block_if_booking( &self, booking: &SchedulerBookingDescriptor, diff --git a/lib/llm/src/kv_router/indexer/ingress.rs b/lib/llm/src/kv_router/indexer/ingress.rs index bfe40424cf9a..19750738b35a 100644 --- a/lib/llm/src/kv_router/indexer/ingress.rs +++ b/lib/llm/src/kv_router/indexer/ingress.rs @@ -134,14 +134,6 @@ impl RuntimeIngress { subscription.set_task_guard(task_guard); } } - - /// Stop event ingestion and wait for it to drain. - pub(crate) async fn shutdown(&self) { - let subscription = self.subscription.lock().take(); - if let Some(subscription) = subscription { - subscription.shutdown().await; - } - } } #[async_trait::async_trait] diff --git a/lib/llm/src/kv_router/indexer/recovery/subscriber.rs b/lib/llm/src/kv_router/indexer/recovery/subscriber.rs index 9dad1493a5d5..da9b213beeee 100644 --- a/lib/llm/src/kv_router/indexer/recovery/subscriber.rs +++ b/lib/llm/src/kv_router/indexer/recovery/subscriber.rs @@ -379,6 +379,7 @@ impl KvEventSubscriptionHandle { self.task_guard = Some(task_guard); } + #[cfg(test)] pub(crate) async fn shutdown(mut self) { self.cancel.cancel(); for completion in self.completions.drain(..) { diff --git a/lib/llm/src/kv_router/prefill_router/activation.rs b/lib/llm/src/kv_router/prefill_router/activation.rs index 7a2440eda900..8fbf22f7051d 100644 --- a/lib/llm/src/kv_router/prefill_router/activation.rs +++ b/lib/llm/src/kv_router/prefill_router/activation.rs @@ -628,11 +628,6 @@ impl PrefillRouter { } } - /// Whether the inner router has initialized. - pub fn is_activated(&self) -> bool { - self.binding.load().is_some() - } - pub(super) fn lifecycle_state(&self) -> PrefillLifecycleState { PrefillLifecycleState::from_atomic(self.lifecycle.load(Ordering::Acquire)) } @@ -947,8 +942,9 @@ mod tests { admissions.send_replace(vec![ids[1]]); tokio::time::timeout(Duration::from_secs(5), async { loop { - let outcome = router - .query_prefill_worker( + let reservation = router + .reserve_prefill_worker( + "admitted-probe", &[1; 128], None, None, @@ -956,16 +952,13 @@ mod tests { 0.0, 0, None, + None, RoutingConstraints::default(), ) .await - .unwrap(); - let crate::kv_router::prefill_router::PrefillQueryOutcome::Routed { - worker_id, .. - } = outcome - else { - panic!("prefill query must remain routable"); - }; + .expect("prefill reservation must remain routable"); + let worker_id = reservation.worker_id(); + reservation.release().await.unwrap(); assert!( ids[..2].contains(&worker_id), "rejected prefill worker became selectable" diff --git a/lib/llm/src/kv_router/prefill_router/conditional_bypass.rs b/lib/llm/src/kv_router/prefill_router/conditional_bypass.rs index 619f2638c3ef..da162d9077b3 100644 --- a/lib/llm/src/kv_router/prefill_router/conditional_bypass.rs +++ b/lib/llm/src/kv_router/prefill_router/conditional_bypass.rs @@ -67,7 +67,7 @@ impl PrefillRouter { let preview = decode_host .preview_kv_route(request, RequestPhase::Decode) .await?; - let signals = preview.signals(); + let signals = preview.signals; let mut input = ConditionalDisaggDecisionInput::new(routing_token_ids.len(), signals.cached_tokens); if self.conditional_disagg_policy.needs_prefill_worker_busy() { diff --git a/lib/llm/src/kv_router/prefill_router/mod.rs b/lib/llm/src/kv_router/prefill_router/mod.rs index 4d165b01c118..df0b9a4466a3 100644 --- a/lib/llm/src/kv_router/prefill_router/mod.rs +++ b/lib/llm/src/kv_router/prefill_router/mod.rs @@ -13,7 +13,7 @@ use uuid::Uuid; use dynamo_kv_router::{ PrefillLoadEstimator, conditional_disagg::ConditionalDisaggPolicy, - config::RouterConfigOverride, protocols::RoutingConstraints, scheduling::QueueRejection, + config::RouterConfigOverride, protocols::RoutingConstraints, }; use dynamo_runtime::{ pipeline::{ @@ -156,17 +156,6 @@ struct PreparedPrefill { topology_constraints: Option, } -/// Advisory prefill worker selection result. -pub enum PrefillQueryOutcome { - Routed { - worker_id: u64, - dp_rank: Option, - }, - QueueRejected { - rejection: QueueRejection, - }, -} - enum PrefillCompletion { Handoff { result: PrefillResult, @@ -339,7 +328,7 @@ impl .await { Ok(Some(decision)) => { - let signals = decision.plan.signals(); + let signals = decision.plan.signals; tracing::info!( request_id = %request_id, worker_id = signals.worker.worker_id, diff --git a/lib/llm/src/kv_router/prefill_router/query.rs b/lib/llm/src/kv_router/prefill_router/query.rs index 4d943dc614cb..92f6d6659e6d 100644 --- a/lib/llm/src/kv_router/prefill_router/query.rs +++ b/lib/llm/src/kv_router/prefill_router/query.rs @@ -9,7 +9,7 @@ use dynamo_kv_router::{ scheduling::queue::BookingHandle, }; -use super::{PrefillError, PrefillLifecycleState, PrefillQueryOutcome, PrefillRouter}; +use super::{PrefillError, PrefillLifecycleState, PrefillRouter}; /// A prefill booking that owns cleanup of the exact scheduler attempt it admitted. /// @@ -87,7 +87,7 @@ impl PrefillRouter { }); }; let admitted = chooser - .find_best_match_details_with_lifecycle( + .find_best_match_details_with_policy_class_admitted( Some(reservation_id), token_ids, block_mm_infos, @@ -101,6 +101,7 @@ impl PrefillRouter { policy_class, None, None, + None, allowed_worker_ids, routing_constraints, ) @@ -128,71 +129,6 @@ impl PrefillRouter { } } } - - /// Query the best prefill worker without executing a request. - /// - /// This query is advisory and does not book scheduler or occupancy state; - /// concurrent callers may observe the same worker. - #[expect(clippy::too_many_arguments)] - pub async fn query_prefill_worker( - &self, - token_ids: &[u32], - block_mm_infos: Option<&[Option]>, - lora_name: Option, - cache_namespace: Option, - priority_jump: f64, - strict_priority: u32, - allowed_worker_ids: Option>, - routing_constraints: RoutingConstraints, - ) -> Result { - if self.lifecycle_state() != PrefillLifecycleState::Active { - return Err(anyhow::anyhow!(PrefillError::NotActivated)); - } - let binding = self - .binding - .load_full() - .ok_or_else(|| anyhow::anyhow!(PrefillError::NotActivated))?; - - let router = &binding.router; - let Some(kv_router) = router.kv_router_if_enabled() else { - let worker_id = router - .peek_next_worker() - .ok_or_else(|| anyhow::anyhow!("No workers available for prefill"))?; - return Ok(PrefillQueryOutcome::Routed { - worker_id, - dp_rank: None, - }); - }; - let outcome = kv_router - .find_best_match_details( - None, - token_ids, - block_mm_infos, - None, - false, - false, - lora_name, - cache_namespace, - priority_jump, - strict_priority, - None, - None, - allowed_worker_ids, - routing_constraints, - ) - .await?; - match outcome { - crate::kv_router::FindBestMatchOutcome::Routed { worker, .. } => { - Ok(PrefillQueryOutcome::Routed { - worker_id: worker.worker_id, - dp_rank: Some(worker.dp_rank), - }) - } - crate::kv_router::FindBestMatchOutcome::QueueRejected { rejection } => { - Ok(PrefillQueryOutcome::QueueRejected { rejection }) - } - } - } } #[cfg(test)] @@ -219,14 +155,14 @@ mod tests { storage::kv, traits::DistributedRuntimeProvider, }; - use futures::{StreamExt, future::join_all}; + use futures::StreamExt; use tokio::sync::watch; use super::*; use crate::{ discovery::ModelManager, kv_router::prefill_router::PrefillBinding, - kv_router::{KvPushRouter, KvRouter, RouterLoadSource, RoutingHost, RoutingLoadContext}, + kv_router::{KvRouter, RouterLoadSource, RoutingHost, RoutingLoadContext}, protocols::common::{ FinishReason, llm_backend::LLMEngineOutput, preprocessor::PreprocessedRequest, }, @@ -320,27 +256,16 @@ mod tests { .unwrap() } - async fn query_worker(router: &PrefillRouter) -> u64 { - match router - .query_prefill_worker( - &[1, 2, 3], - None, - None, - None, - 0.0, - 0, - None, - RoutingConstraints::default(), - ) - .await + /// The advisory builtin-host peek `reserve_prefill_worker`'s non-KV arm + /// routes through. + fn query_worker(router: &PrefillRouter) -> u64 { + router + .binding + .load_full() .unwrap() - { - PrefillQueryOutcome::Routed { worker_id, dp_rank } => { - assert_eq!(dp_rank, None); - worker_id - } - PrefillQueryOutcome::QueueRejected { .. } => panic!("RR query cannot queue"), - } + .router + .peek_next_worker() + .expect("RR query cannot fail") } async fn shared_router( @@ -503,7 +428,7 @@ mod tests { let binding = Arc::new(PrefillBinding { target_id: crate::discovery::WorkerSetTargetId::Legacy(endpoint_id.clone()), endpoint_id, - router: Arc::new(KvPushRouter::new(push_router, chooser.clone(), None).unwrap()), + router: Arc::new(RoutingHost::new(push_router, chooser.clone(), None).unwrap()), prefill_router_mode: RouterMode::KV, }); (binding, chooser) @@ -620,7 +545,9 @@ mod tests { assert_eq!(reservation.dp_rank(), None); reservation.release().await.unwrap(); - let concurrent_peeks = join_all((0..16).map(|_| query_worker(&prefill_router))).await; + let concurrent_peeks = (0..16) + .map(|_| query_worker(&prefill_router)) + .collect::>(); assert!( concurrent_peeks .iter() @@ -628,8 +555,8 @@ mod tests { ); for expected_worker in &expected_workers { - assert_eq!(query_worker(&prefill_router).await, *expected_worker); - assert_eq!(query_worker(&prefill_router).await, *expected_worker); + assert_eq!(query_worker(&prefill_router), *expected_worker); + assert_eq!(query_worker(&prefill_router), *expected_worker); let mut stream = shared_router .generate(Context::new(request())) .await @@ -665,7 +592,9 @@ mod tests { ) .await; - let _ = join_all((0..16).map(|_| query_worker(&prefill))).await; + for _ in 0..16 { + query_worker(&prefill); + } assert_eq!( workers .iter() diff --git a/lib/llm/src/kv_router/request_lease.rs b/lib/llm/src/kv_router/request_lease.rs index df0291c75eed..abaa928ec176 100644 --- a/lib/llm/src/kv_router/request_lease.rs +++ b/lib/llm/src/kv_router/request_lease.rs @@ -12,15 +12,15 @@ use std::{ use dynamo_kv_router::{ multi_worker_sequence::{ReplicaRequestLeaseObserver, active_request_expiry_duration}, - scheduling::AttemptId, + scheduling::{ + AttemptId, + queue::{SchedulerBookingCleanup, SchedulerBookingDescriptor}, + }, }; use parking_lot::Mutex; use tokio_util::sync::CancellationToken; -use super::{ - indexer::ApproximateRequestLease, - scheduler::{SchedulerBookingCleanup, SchedulerBookingDescriptor}, -}; +use super::indexer::ApproximateRequestLease; const LEASE_QUIET: u8 = 0; const LEASE_TOUCHED: u8 = 1; diff --git a/lib/llm/src/kv_router/routing_host.rs b/lib/llm/src/kv_router/routing_host.rs index 8c95432a3db3..0310e80017a5 100644 --- a/lib/llm/src/kv_router/routing_host.rs +++ b/lib/llm/src/kv_router/routing_host.rs @@ -245,7 +245,7 @@ pub struct RoutingHost { /// An admitted KV route awaiting dispatch. pub(crate) struct RoutePlan { - signals: RoutePlanSignals, + pub(crate) signals: RoutePlanSignals, selection: WorkerSelection, cleanup: KvRequestCleanup, affinity: Option, @@ -258,7 +258,7 @@ pub(crate) struct RoutePlan { pub(crate) struct RoutePreview { request_id: String, phase: RequestPhase, - signals: RoutePlanSignals, + pub(crate) signals: RoutePlanSignals, /// Starts here because the conditional route's first stage is the preview. budget: CleanupBudget, } @@ -273,10 +273,6 @@ pub(crate) struct RoutePlanSignals { } impl RoutePreview { - pub(crate) fn signals(&self) -> RoutePlanSignals { - self.signals - } - /// Starts the budget's clock and reports what is left, so a test can follow /// one budget across the real preview/plan/dispatch chain. #[cfg(test)] @@ -293,10 +289,6 @@ impl RoutePlanSignals { } impl RoutePlan { - pub(crate) fn signals(&self) -> RoutePlanSignals { - self.signals - } - #[cfg(test)] pub(crate) fn cleanup_budget_remaining(&self) -> std::time::Duration { self.budget.remaining() diff --git a/lib/llm/src/kv_router/routing_host/tests.rs b/lib/llm/src/kv_router/routing_host/tests.rs index c17601746993..cc01353373e3 100644 --- a/lib/llm/src/kv_router/routing_host/tests.rs +++ b/lib/llm/src/kv_router/routing_host/tests.rs @@ -975,7 +975,7 @@ async fn always_ready_terminals_cannot_starve_the_drain_deadline() { #[tokio::test] #[serial_test::serial] async fn kv_selection_future_stays_small() { - let (router, _runtime) = embedded_router_with_workers(None, &[1]).await; + let (router, _runtime) = router_with_workers(None, &[1]).await; let request = Context::with_id_and_metadata( request(), "selection-future-size".to_string(), @@ -1171,29 +1171,6 @@ async fn router_with_workers( async fn router_with_worker_configs( session_affinity_ttl: Option, workers: HashMap, -) -> (RoutingHost, Runtime) { - router_with_worker_configs_in_mode(session_affinity_ttl, workers, false).await -} - -/// Same router, scheduling on an embedded selection-service partition. -async fn embedded_router_with_workers( - session_affinity_ttl: Option, - worker_ids: &[u64], -) -> (RoutingHost, Runtime) { - let workers = worker_ids - .iter() - .copied() - .map(|worker_id| (worker_id, ModelRuntimeConfig::default())) - .collect(); - router_with_worker_configs_in_mode(session_affinity_ttl, workers, true).await -} - -/// `_embedded_selection` is kept for call-site stability: the embedded -/// partition is the only scheduler now. -async fn router_with_worker_configs_in_mode( - session_affinity_ttl: Option, - workers: HashMap, - _embedded_selection: bool, ) -> (RoutingHost, Runtime) { let runtime = Runtime::from_current().unwrap(); let distributed = DistributedRuntime::new(runtime.clone(), DistributedConfig::process_local()) @@ -1445,98 +1422,6 @@ async fn track_request( (request, selection, guard) } -/// The embedded selection partition books and releases exactly like the -/// runtime scheduler: preview does not book, plan books, abort releases, and -/// a tracked request frees on guard completion. -#[tokio::test] -#[serial_test::serial] -async fn embedded_selection_books_and_releases_like_the_runtime_scheduler() { - let (router, runtime) = embedded_router_with_workers(None, &[7, 9]).await; - let request = Context::new(request()); - - let preview = router - .preview_kv_route(&request, RequestPhase::Decode) - .await - .expect("decode preview should select one worker"); - let loads = router - .kv_router() - .get_potential_loads(&[], None, None, None, None) - .await - .unwrap(); - assert!( - loads.iter().all(|load| load.active_requests == 0), - "a preview must not book: {loads:?}" - ); - - let plan = router - .plan_kv_route_from_preview(&request, preview) - .await - .expect("decode plan should admit one request"); - let planned_worker = plan.signals().worker.worker_id; - assert!([7, 9].contains(&planned_worker)); - let loads = router - .kv_router() - .get_potential_loads(&[], None, None, None, None) - .await - .unwrap(); - assert_eq!( - loads - .iter() - .find(|load| load.worker_id == planned_worker && load.dp_rank == 0) - .expect("selected worker must be reported") - .active_requests, - 1 - ); - plan.abort().await; - let loads = router - .kv_router() - .get_potential_loads(&[], None, None, None, None) - .await - .unwrap(); - assert!( - loads.iter().all(|load| load.active_requests == 0), - "abandoned plans must release their embedded reservation: {loads:?}" - ); - - // A tracked request books on admission and frees when its guard finishes. - let (_request, selection, mut guard) = track_request(&router, false).await; - let loads = router - .kv_router() - .get_potential_loads(&[], None, None, None, None) - .await - .unwrap(); - assert_eq!( - loads - .iter() - .find(|load| load.worker_id == selection.worker.worker_id) - .expect("tracked worker must be reported") - .active_requests, - 1 - ); - guard.abort().await; - let released = tokio::time::timeout(Duration::from_secs(2), async { - loop { - let loads = router - .kv_router() - .get_potential_loads(&[], None, None, None, None) - .await - .unwrap(); - if loads.iter().all(|load| load.active_requests == 0) { - return; - } - tokio::time::sleep(Duration::from_millis(10)).await; - } - }) - .await; - assert!( - released.is_ok(), - "guard drop must release the embedded booking" - ); - - drop(router); - runtime.shutdown(); -} - #[tokio::test] #[serial_test::serial] async fn route_plan_from_preview_holds_and_releases_the_decode_reservation() { @@ -1552,7 +1437,7 @@ async fn route_plan_from_preview_holds_and_releases_the_decode_reservation() { .plan_kv_route_from_preview(&request, preview) .await .expect("decode plan should admit one request"); - assert_eq!(plan.signals().worker.worker_id, 7); + assert_eq!(plan.signals.worker.worker_id, 7); assert_eq!( router.request_metrics.requests_started_total().get(), requests_started_before, @@ -1603,7 +1488,7 @@ async fn route_preview_does_not_admit_a_request() { .preview_kv_route(&request, RequestPhase::Decode) .await .expect("decode preview should select one request"); - assert_eq!(preview.signals().worker.worker_id, 7); + assert_eq!(preview.signals.worker.worker_id, 7); let loads = router .kv_router() .get_potential_loads(&[], None, None, None, None) @@ -1628,13 +1513,13 @@ async fn route_plan_from_preview_admits_the_previewed_worker() { .preview_kv_route(&request, RequestPhase::Decode) .await .unwrap(); - let previewed_worker = preview.signals().worker; + let previewed_worker = preview.signals.worker; let plan = router .plan_kv_route_from_preview(&request, preview) .await .unwrap(); - assert_eq!(plan.signals().worker, previewed_worker); + assert_eq!(plan.signals.worker, previewed_worker); let loads = router .kv_router() .get_potential_loads(&[], None, None, None, None) @@ -1813,16 +1698,40 @@ async fn router_request_counters_follow_admission_and_completion_lifecycle() { drop(query_guard); assert_eq!(metrics.requests_started_total().get(), started_before); - let (_, _, mut cancelled_guard) = track_request(&router, false).await; + let (_, selection, mut cancelled_guard) = track_request(&router, false).await; assert_eq!(metrics.requests_started_total().get(), started_before + 1); assert_eq!(metrics.requests_total.get(), completed_before); + // The booking is held from admission until the guard finishes. + let loads = router + .kv_router() + .get_potential_loads(&[], None, None, None, None) + .await + .unwrap(); + assert_eq!( + loads + .iter() + .find(|load| load.worker_id == selection.worker.worker_id) + .expect("tracked worker must be reported") + .active_requests, + 1, + "a tracked request stays booked while its guard is live: {loads:?}" + ); // Admission remains counted even when the request aborts before dispatch. cancelled_guard.abort().await; drop(cancelled_guard); assert_eq!(metrics.requests_started_total().get(), started_before + 1); assert_eq!(metrics.requests_total.get(), completed_before); + let loads = router + .kv_router() + .get_potential_loads(&[], None, None, None, None) + .await + .unwrap(); + assert!( + loads.iter().all(|load| load.active_requests == 0), + "an aborted guard frees its booking: {loads:?}" + ); let mut failed_input = request(); failed_input.migration_state = Some(Default::default()); diff --git a/lib/llm/src/kv_router/scheduler.rs b/lib/llm/src/kv_router/scheduler.rs deleted file mode 100644 index e8892968ae37..000000000000 --- a/lib/llm/src/kv_router/scheduler.rs +++ /dev/null @@ -1,11 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -//! Scheduling types shared by the router. The router's scheduler is one -//! partition of an embedded selection service (`kv_router::embedded`); these -//! re-exports are the runtime-free scheduling vocabulary it speaks. - -pub use dynamo_kv_router::scheduling::queue::{ - SchedulerBookingCleanup, SchedulerBookingDescriptor, -}; -pub use dynamo_kv_router::scheduling::{KvSchedulerError, PotentialLoad}; diff --git a/lib/llm/src/session_affinity/coordinator.rs b/lib/llm/src/session_affinity/coordinator.rs index 675bd34bb401..8b016caf6146 100644 --- a/lib/llm/src/session_affinity/coordinator.rs +++ b/lib/llm/src/session_affinity/coordinator.rs @@ -116,24 +116,18 @@ impl AffinityCoordinator { session_id: &SessionAffinityId, requested_target: Option, ) -> Result { - self.acquire_inner(session_id, requested_target, None).await - } - - pub(crate) async fn acquire_with_context( - &self, - session_id: &SessionAffinityId, - requested_target: Option, - request_context: &dyn AsyncEngineContext, - ) -> Result { - self.acquire_inner(session_id, requested_target, Some(request_context)) + self.inner + .table + .acquire(session_id.as_str(), requested_target.map(to_table)) .await + .map_err(affinity_error) } - async fn acquire_inner( + pub(crate) async fn acquire_with_context( &self, session_id: &SessionAffinityId, requested_target: Option, - request_context: Option<&dyn AsyncEngineContext>, + context: &dyn AsyncEngineContext, ) -> Result { let requested = requested_target.map(to_table); loop { @@ -144,17 +138,14 @@ impl AffinityCoordinator { .map_err(affinity_error)? { AcquireStep::Held(hold) => return Ok(hold), - AcquireStep::Wait(notified) => match request_context { - Some(context) => { - tokio::select! { - biased; - _ = context.stopped() => return Err(cancelled(context.id())), - _ = context.killed() => return Err(cancelled(context.id())), - _ = notified => {} - } + AcquireStep::Wait(notified) => { + tokio::select! { + biased; + _ = context.stopped() => return Err(cancelled(context.id())), + _ = context.killed() => return Err(cancelled(context.id())), + _ = notified => {} } - None => notified.await, - }, + } } } } diff --git a/lib/llm/src/session_affinity/mod.rs b/lib/llm/src/session_affinity/mod.rs index 18108044eb8b..b0dd066ec75f 100644 --- a/lib/llm/src/session_affinity/mod.rs +++ b/lib/llm/src/session_affinity/mod.rs @@ -14,8 +14,7 @@ pub use coordinator::{AffinityCoordinator, AffinityTarget, explicit_target}; pub(crate) use coordinator::{affinity_id, from_table, invalid_argument}; pub(crate) use dynamo_kv_router::services::selection::affinity::Hold; pub use dynamo_kv_router::services::selection::affinity::{ - MAX_SESSION_AFFINITY_ENTRIES, MAX_SESSION_AFFINITY_ID_BYTES, MAX_SESSION_AFFINITY_TTL_SECS, - SessionAffinityMode, + MAX_SESSION_AFFINITY_TTL_SECS, SessionAffinityMode, }; pub type LlmResponse = From 63f7ac4998bd864752fbd6f080386255f4b862bf Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Sat, 12 Sep 2026 13:50:34 -0700 Subject: [PATCH 069/102] chore(router): cleanup round two, one result shape and fewer wrappers Behavior-neutral collapses in the frontend host and the selection core. - One selection result shape in the frontend: `AdmittedFindBestMatchOutcome` carries `advisory_load`, set exactly for without-admission selection, and `FindBestMatchAdvisoryOutcome` / `FindBestMatchInnerOutcome` are gone. The Python-visible `FindBestMatchOutcome`, `into_parts` and the admitted entry point are unchanged. A debug assertion pins the advisory/admitted shape the core guarantees. - `EmbeddedSelection` exposes its `scheduler()` instead of ten forwarding methods; the replica ingress is returned directly instead of through a wrapper; `WorkerCatalogRecord::dp_ranks` replaces a frontend copy; `update_queue_metrics` reads the scheduler it is given. - kv-router: `SelectionService::ensure_partition` forwarder removed (the frontend calls the core); the "remove worker from index" fallback is one helper in `ingress.rs` used by both the ZMQ detach and the core's registration cleanup; `schedule_request_with_booking` is crate-private. - Test `metadata_update_preserves_live_booking` is deleted; its survivor `catalog_updates_commit_in_order_without_exposing_candidates` now counts detaches (none during the update, one on deletion). Co-Authored-By: Claude Fable 5.1 Signed-off-by: Ishan Dhanani --- lib/kv-router/src/scheduling/local.rs | 3 +- .../src/services/selection/core/hint.rs | 4 +- .../src/services/selection/core/tests.rs | 64 ++---- .../src/services/selection/core/workers.rs | 11 +- .../src/services/selection/ingress.rs | 23 +- lib/kv-router/src/services/selection/mod.rs | 2 +- .../src/services/selection/service.rs | 15 +- lib/kv-router/src/services/selection/types.rs | 2 +- lib/llm/src/kv_router.rs | 206 +++++++----------- lib/llm/src/kv_router/embedded.rs | 147 ++----------- .../kv_router/routing_host/kv_selection.rs | 79 +++---- 11 files changed, 172 insertions(+), 384 deletions(-) diff --git a/lib/kv-router/src/scheduling/local.rs b/lib/kv-router/src/scheduling/local.rs index 061b6f58d7f4..00494c6478f6 100644 --- a/lib/kv-router/src/scheduling/local.rs +++ b/lib/kv-router/src/scheduling/local.rs @@ -301,8 +301,7 @@ where /// Schedule a request and return an armed handle for its booking: dropping /// the handle frees the booking, `commit` hands it to a longer-lived owner. /// The handle is `None` unless the mode is `TrackedWithLifecycle`. - #[doc(hidden)] - pub async fn schedule_request_with_booking( + pub(crate) async fn schedule_request_with_booking( &self, request: ScheduleRequest, ) -> Result<(AdmittedSchedulingResponse, Option), KvSchedulerError> { diff --git a/lib/kv-router/src/services/selection/core/hint.rs b/lib/kv-router/src/services/selection/core/hint.rs index a5e7f0779211..dd382c6ce721 100644 --- a/lib/kv-router/src/services/selection/core/hint.rs +++ b/lib/kv-router/src/services/selection/core/hint.rs @@ -7,8 +7,8 @@ use super::*; /// Pick the best router-hint source for `target`: a same-role worker (or /// cache owner) holding a longer root-aligned prefix than the target's own -/// `target_cached_prefix_blocks`, with a non-empty control endpoint. Mirrors the -/// frontend `KvRouter::transfer_hint_for_selection`. +/// `target_cached_prefix_blocks`, with a non-empty control endpoint. Both hosts +/// reach it through `run_selection`. pub(super) fn transfer_hint_for_selection( configs: &HashMap, target: WorkerWithDpRank, diff --git a/lib/kv-router/src/services/selection/core/tests.rs b/lib/kv-router/src/services/selection/core/tests.rs index 43035f5da945..b4d88f5f8018 100644 --- a/lib/kv-router/src/services/selection/core/tests.rs +++ b/lib/kv-router/src/services/selection/core/tests.rs @@ -13,7 +13,7 @@ use crate::protocols::{RoutingConstraints, StorageTier}; use crate::services::common::replica_sync::HostReplicaChannels; use crate::services::indexer::backend::test_util::store_event; use std::collections::HashSet; -use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::thread::sleep; use std::time::Duration; @@ -2417,51 +2417,6 @@ async fn rejoined_worker_feeds_partition_index() { ); } } -#[tokio::test] -async fn metadata_update_preserves_live_booking() { - struct SuspendedDetach(tokio::sync::Notify); - #[async_trait::async_trait] - impl KvEventIngress for SuspendedDetach { - fn open( - &self, - registry: &WorkerRegistry, - key: &RoutingPartitionId, - block_size: u32, - ) -> Indexer { - registry.get_or_create_indexer(key.clone(), block_size) - } - async fn detach(&self, _registry: &WorkerRegistry, _record: &WorkerCatalogRecord) { - self.0.notify_one(); - std::future::pending().await - } - } - let ingress = Arc::new(SuspendedDetach(tokio::sync::Notify::new())); - let core = core_with_host(SelectionHost { - cache: HostCache { - index: KvIndexSource::Owned(ingress.clone()), - shared: None, - }, - ..SelectionHost::default() - }); - core.upsert_worker(worker(1)).await.unwrap(); - core.select_and_reserve(reserve_request("live")) - .await - .unwrap(); - let entry = core - .entry(&RoutingPartitionId::new("model", "default")) - .unwrap(); - let mut updated = worker(1); - updated.max_num_batched_tokens = Some(2048); - tokio::select! { - result = core.upsert_worker(updated) => { result.unwrap(); }, - _ = ingress.0.notified() => { - // A capacity update must not temporarily withdraw this worker. - panic!("capacity update detached the worker"); - } - } - assert!(entry.scheduler.has_request("live")); - core.free_reservation("live").await.unwrap(); -} #[tokio::test(start_paused = true)] async fn host_lease_manager_owns_expiry() { @@ -2504,6 +2459,7 @@ async fn catalog_updates_commit_in_order_without_exposing_candidates() { struct PausedUpdate { entered: tokio::sync::Notify, release: tokio::sync::Notify, + detaches: AtomicUsize, } #[async_trait::async_trait] impl KvEventIngress for PausedUpdate { @@ -2527,10 +2483,14 @@ async fn catalog_updates_commit_in_order_without_exposing_candidates() { } Ok(()) } + async fn detach(&self, _: &WorkerRegistry, _: &WorkerCatalogRecord) { + self.detaches.fetch_add(1, Ordering::SeqCst); + } } let ingress = Arc::new(PausedUpdate { entered: tokio::sync::Notify::new(), release: tokio::sync::Notify::new(), + detaches: AtomicUsize::new(0), }); let mut core = core_with_host(SelectionHost { cache: HostCache { @@ -2568,11 +2528,23 @@ async fn catalog_updates_commit_in_order_without_exposing_candidates() { ); ingress.release.notify_one(); assert_eq!(update.await.unwrap().max_num_batched_tokens, Some(2048)); + // A capacity update must not temporarily withdraw this worker; only the + // pending deletion detaches it. + assert_eq!( + ingress.detaches.load(Ordering::SeqCst), + 0, + "capacity update detached the worker" + ); assert!(entry.scheduler.has_request("live")); assert_eq!( deletion.await.unwrap().lifecycle, WorkerLifecycle::Unschedulable ); + assert_eq!( + ingress.detaches.load(Ordering::SeqCst), + 1, + "deletion detaches the worker" + ); tokio::time::timeout(Duration::from_secs(1), async { while entry.scheduler.has_request("live") { tokio::task::yield_now().await; diff --git a/lib/kv-router/src/services/selection/core/workers.rs b/lib/kv-router/src/services/selection/core/workers.rs index deeb0317b6b1..335498773f3f 100644 --- a/lib/kv-router/src/services/selection/core/workers.rs +++ b/lib/kv-router/src/services/selection/core/workers.rs @@ -6,6 +6,7 @@ use super::reservations::{ReservationIndexObserver, spawn_reservation_index_sweep}; use super::*; +use crate::services::selection::ingress::remove_worker_from_index; impl SelectionCore { pub async fn upsert_worker( @@ -331,15 +332,7 @@ impl SelectionCore { ingress.detach(&self.indexer_registry, record).await; return; } - - let key = record.key(); - let indexer = self - .indexer_registry - .get_indexer(&key) - .map(|entry| entry.indexer.clone()); - if let Some(indexer) = indexer { - indexer.remove_worker(record.worker_id).await; - } + remove_worker_from_index(&self.indexer_registry, record).await; } pub(super) fn publish_scheduler_config( diff --git a/lib/kv-router/src/services/selection/ingress.rs b/lib/kv-router/src/services/selection/ingress.rs index 73bc3c9e3360..272eef76c6eb 100644 --- a/lib/kv-router/src/services/selection/ingress.rs +++ b/lib/kv-router/src/services/selection/ingress.rs @@ -179,13 +179,20 @@ impl KvEventIngress for ZmqDirectIngress { } return; } - // No listeners (events disabled or a remote primary): drop what an - // approximate or side index recorded for the worker. - if let Some(indexer) = registry - .get_indexer(&record.key()) - .map(|entry| entry.indexer.clone()) - { - indexer.remove_worker(record.worker_id).await; - } + remove_worker_from_index(registry, record).await; + } +} + +/// No listeners for `record` (events disabled or a remote primary): drop what +/// an approximate or side index recorded for the worker. +pub(super) async fn remove_worker_from_index( + registry: &WorkerRegistry, + record: &WorkerCatalogRecord, +) { + if let Some(indexer) = registry + .get_indexer(&record.key()) + .map(|entry| entry.indexer.clone()) + { + indexer.remove_worker(record.worker_id).await; } } diff --git a/lib/kv-router/src/services/selection/mod.rs b/lib/kv-router/src/services/selection/mod.rs index f60883cdbd18..909a08e245a8 100644 --- a/lib/kv-router/src/services/selection/mod.rs +++ b/lib/kv-router/src/services/selection/mod.rs @@ -29,7 +29,7 @@ pub use crate::services::common::replica_sync::{ pub use core::{ HostCache, HostEligibility, HostLoad, HostReplication, HostTelemetry, KvIndexSource, Selected, SelectionAdmission, SelectionCore, SelectionHost, SelectionOperation, SelectionOutcome, - SelectionPartition, SelectionRun, SelectionServiceConfig, SessionBinding, + SelectionPartition, SelectionRun, SelectionScheduler, SelectionServiceConfig, SessionBinding, }; pub use error::SelectionError; pub use ingress::KvEventIngress; diff --git a/lib/kv-router/src/services/selection/service.rs b/lib/kv-router/src/services/selection/service.rs index a84f46a87924..87932f187619 100644 --- a/lib/kv-router/src/services/selection/service.rs +++ b/lib/kv-router/src/services/selection/service.rs @@ -8,7 +8,6 @@ use tokio_util::sync::CancellationToken; use crate::WorkerSelectionPolicyFactory; use crate::config::KvRouterConfig; -use crate::identity::RoutingPartitionId; use crate::protocols::WorkerId; use crate::scheduling::PotentialLoad; use crate::services::common::replica_sync::{ @@ -18,9 +17,7 @@ use crate::services::indexer::backend::IndexerPolicy; use crate::tracking_hash::TrackingHashContext; use super::affinity::SessionAffinityConfig; -use super::core::{ - KvIndexSource, SelectionCore, SelectionHost, SelectionPartition, SelectionServiceConfig, -}; +use super::core::{KvIndexSource, SelectionCore, SelectionHost, SelectionServiceConfig}; use super::error::SelectionError; use super::pending::SelectionCacheConfig; use super::policy_registry::WorkerSelectionPolicyRegistry; @@ -339,16 +336,6 @@ impl SelectionService { &self.core } - /// Create (or fetch) the partition for `key` ahead of worker registration. - pub fn ensure_partition( - &self, - key: RoutingPartitionId, - block_size: u32, - is_eagle: bool, - ) -> Result { - self.core.ensure_partition(key, block_size, is_eagle) - } - /// The port this service uses for replica synchronization, if enabled. pub fn replica_sync_port(&self) -> Option { self.replica_sync_port diff --git a/lib/kv-router/src/services/selection/types.rs b/lib/kv-router/src/services/selection/types.rs index 9b77f3dfd532..22f479e12c98 100644 --- a/lib/kv-router/src/services/selection/types.rs +++ b/lib/kv-router/src/services/selection/types.rs @@ -200,7 +200,7 @@ impl WorkerCatalogRecord { self.data_parallel_size.unwrap_or(1) } - pub(super) fn dp_ranks(&self) -> impl Iterator { + pub fn dp_ranks(&self) -> impl Iterator { let start = self.dp_start(); let size = self.dp_size(); start..start.saturating_add(size) diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index 6c8268fad9f6..e78718848fa3 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -399,25 +399,6 @@ pub enum FindBestMatchOutcome { }, } -/// For probes that return best-match routing decisions plus selected-worker -/// scheduler-load snapshots, without admitting the request into scheduler state. -/// `FindBestMatchInnerOutcome` keeps this advisory shape internal so admitted -/// routing can keep using `FindBestMatchOutcome` unchanged. -pub enum FindBestMatchAdvisoryOutcome { - Routed { - worker: WorkerWithDpRank, - overlap_blocks: u32, - effective_overlap_blocks: f64, - cached_tokens: usize, - potential_decode_blocks: u64, - selected_worker_load: scheduling::AdvisoryWorkerLoad, - routing_hashes: Option, - }, - QueueRejected { - rejection: scheduling::QueueRejection, - }, -} - #[derive(Debug, Clone, Copy)] pub(super) enum FindBestMatchAdmission { WithAdmission, @@ -431,6 +412,9 @@ pub(super) enum FindBestMatchAdmission { pub struct AdmittedFindBestMatchOutcome { pub(super) outcome: FindBestMatchOutcome, pub(super) booking: Option, + /// The selected worker's scheduler-load snapshot; set only by advisory + /// probes (`FindBestMatchAdmission::WithoutAdmission`), which never book. + pub(super) advisory_load: Option, } impl AdmittedFindBestMatchOutcome { @@ -440,11 +424,6 @@ impl AdmittedFindBestMatchOutcome { } } -pub(super) enum FindBestMatchInnerOutcome { - WithAdmission(AdmittedFindBestMatchOutcome), - WithoutAdmission(FindBestMatchAdvisoryOutcome), -} - #[derive(Debug, Clone, Copy, Default)] pub(crate) struct WorkerCacheHitEstimate { pub effective_overlap_blocks: f64, @@ -533,7 +512,7 @@ pub fn router_endpoint_id(namespace: String, component: String) -> EndpointId { /// TODO: Rename this to indicate it only selects a worker, it does not route. pub struct KvRouter { indexer: Indexer, - scheduler: embedded::EmbeddedSelection, + selection: embedded::EmbeddedSelection, required_worker_inputs: dynamo_kv_router::selector::WorkerInputs, workers_with_configs: RuntimeConfigWatch, block_size: u32, @@ -769,7 +748,7 @@ impl KvRouter { // the scheduler; its scheduler cleanup is set once the scheduler does. let request_leases = request_lease::RequestLeaseManager::new(cancellation_token.child_token()); - let (scheduler, replica_ingress) = embedded::EmbeddedSelection::start( + let (selection, replica_ingress) = embedded::EmbeddedSelection::start( embedded::EmbeddedSelectionArgs { kv_router_config: kv_router_config.clone(), worker_role, @@ -796,7 +775,7 @@ impl KvRouter { cancellation_token.child_token(), ) .await?; - request_leases.set_scheduler(scheduler.booking_cleanup()); + request_leases.set_scheduler(selection.scheduler().booking_cleanup()); // Inbound lifecycle events start only now that their consumer can // release bookings. replica_ingress.start().await; @@ -804,7 +783,7 @@ impl KvRouter { let cancellation_token = cancellation_guard.disarm(); Ok(Self { indexer, - scheduler, + selection, required_worker_inputs, workers_with_configs, block_size, @@ -1132,34 +1111,27 @@ impl KvRouter { allowed_worker_ids: Option>, routing_constraints: RoutingConstraints, ) -> anyhow::Result { - match self - .find_best_match_details_with_policy_class_inner( - context_id, - tokens, - block_mm_infos, - router_config_override, - update_states, - return_routing_hashes, - lora_name, - cache_namespace, - priority_jump, - strict_priority, - policy_class, - session_context, - expected_output_tokens, - None, - pinned_worker, - allowed_worker_ids, - routing_constraints, - FindBestMatchAdmission::WithAdmission, - ) - .await? - { - FindBestMatchInnerOutcome::WithAdmission(admitted) => Ok(admitted), - FindBestMatchInnerOutcome::WithoutAdmission(_) => { - unreachable!("with-admission routing returned advisory outcome") - } - } + self.find_best_match_details_with_policy_class_inner( + context_id, + tokens, + block_mm_infos, + router_config_override, + update_states, + return_routing_hashes, + lora_name, + cache_namespace, + priority_jump, + strict_priority, + policy_class, + session_context, + expected_output_tokens, + None, + pinned_worker, + allowed_worker_ids, + routing_constraints, + FindBestMatchAdmission::WithAdmission, + ) + .await } #[allow(clippy::too_many_arguments)] @@ -1183,7 +1155,7 @@ impl KvRouter { allowed_worker_ids: Option>, routing_constraints: RoutingConstraints, admission: FindBestMatchAdmission, - ) -> anyhow::Result { + ) -> anyhow::Result { let start = Instant::now(); if update_states && context_id.is_none() { anyhow::bail!("context_id must be provided if update_states is true"); @@ -1204,9 +1176,9 @@ impl KvRouter { result: outcome, lookup, } = self - .scheduler + .selection .run_selection(SelectionOperation { - key: self.scheduler.partition_key().clone(), + key: self.selection.partition_key().clone(), prompt: PromptView { token_ids: Some(tokens), mm_routing_info: None, @@ -1243,18 +1215,10 @@ impl KvRouter { let selected = match outcome { Ok(SelectionOutcome::Selected(selected)) => selected, Ok(SelectionOutcome::QueueRejected { rejection }) => { - return Ok(match admission { - FindBestMatchAdmission::WithAdmission => { - FindBestMatchInnerOutcome::WithAdmission(AdmittedFindBestMatchOutcome { - outcome: FindBestMatchOutcome::QueueRejected { rejection }, - booking: None, - }) - } - FindBestMatchAdmission::WithoutAdmission => { - FindBestMatchInnerOutcome::WithoutAdmission( - FindBestMatchAdvisoryOutcome::QueueRejected { rejection }, - ) - } + return Ok(AdmittedFindBestMatchOutcome { + outcome: FindBestMatchOutcome::QueueRejected { rejection }, + booking: None, + advisory_load: None, }); } Err(SelectionError::Scheduler(error)) => return Err(map_scheduler_error(error)), @@ -1309,33 +1273,25 @@ impl KvRouter { } } - Ok(match admission { - FindBestMatchAdmission::WithAdmission => { - FindBestMatchInnerOutcome::WithAdmission(AdmittedFindBestMatchOutcome { - outcome: FindBestMatchOutcome::Routed { - worker: response.best_worker, - overlap_blocks, - effective_overlap_blocks: response.effective_overlap_blocks, - cached_tokens: response.cached_tokens, - potential_decode_blocks: response.potential_decode_blocks as u64, - routing_hashes, - kv_hint, - }, - booking, - }) - } - FindBestMatchAdmission::WithoutAdmission => { - FindBestMatchInnerOutcome::WithoutAdmission(FindBestMatchAdvisoryOutcome::Routed { - worker: response.best_worker, - overlap_blocks, - effective_overlap_blocks: response.effective_overlap_blocks, - cached_tokens: response.cached_tokens, - potential_decode_blocks: response.potential_decode_blocks as u64, - selected_worker_load: advisory_load - .expect("without-admission selection returns advisory load"), - routing_hashes, - }) - } + // Advisory selections carry no booking or hint and always a load + // snapshot; admitted ones the reverse (`SelectionCore::select_or_reject`). + debug_assert_eq!( + advisory_load.is_some(), + !is_admitted_routing, + "advisory load is set exactly for without-admission selection" + ); + Ok(AdmittedFindBestMatchOutcome { + outcome: FindBestMatchOutcome::Routed { + worker: response.best_worker, + overlap_blocks, + effective_overlap_blocks: response.effective_overlap_blocks, + cached_tokens: response.cached_tokens, + potential_decode_blocks: response.potential_decode_blocks as u64, + routing_hashes, + kv_hint, + }, + booking, + advisory_load, }) } @@ -1385,7 +1341,11 @@ impl KvRouter { worker, lora_name, }; - let admission = self.scheduler.add_request_admitted(sequence_request).await; + let admission = self + .selection + .scheduler() + .add_request_admitted(sequence_request) + .await; let attempt_id = match admission { Ok(attempt_id) => attempt_id, Err(error) => { @@ -1406,7 +1366,10 @@ impl KvRouter { } pub async fn mark_prefill_completed(&self, request_id: &str) -> Result<(), SequenceError> { - self.scheduler.mark_prefill_completed(request_id).await?; + self.selection + .scheduler() + .mark_prefill_completed(request_id) + .await?; self.request_leases.touch_request(request_id); Ok(()) } @@ -1415,7 +1378,7 @@ impl KvRouter { if self.request_leases.finish_request(request_id).await { return Ok(()); } - self.scheduler.free(request_id).await + self.selection.scheduler().free(request_id).await } pub(crate) fn affinity_coordinator( @@ -1423,7 +1386,7 @@ impl KvRouter { ttl: std::time::Duration, mode: crate::session_affinity::SessionAffinityMode, ) -> anyhow::Result { - self.scheduler.affinity_coordinator(ttl, mode) + self.selection.affinity_coordinator(ttl, mode) } pub(crate) fn request_lease_manager(&self) -> &request_lease::RequestLeaseManager { @@ -1434,7 +1397,8 @@ impl KvRouter { &self, booking: &SchedulerBookingDescriptor, ) -> Result<(), KvSchedulerError> { - self.scheduler + self.selection + .scheduler() .mark_prefill_completed_if_booking(booking) .await .map(|_| ()) @@ -1442,12 +1406,12 @@ impl KvRouter { /// Number of requests currently parked in the scheduler queue. pub fn pending_count(&self) -> usize { - self.scheduler.pending_count() + self.selection.scheduler().pending_count() } /// Sum of ISL tokens for requests currently parked in the scheduler queue. pub fn pending_isl_tokens(&self) -> usize { - self.scheduler.pending_isl_tokens() + self.selection.scheduler().pending_isl_tokens() } fn prefill_load_hint_for( @@ -1490,7 +1454,7 @@ impl KvRouter { /// Get the worker type for this router ("prefill" or "decode"). /// Used for Prometheus metric labeling. pub fn worker_type(&self) -> &'static str { - self.scheduler.worker_type() + self.selection.worker_type() } /// Return the worker's unique global DP rank when it owns exactly one rank. @@ -1505,7 +1469,8 @@ impl KvRouter { booking: &SchedulerBookingDescriptor, decay_fraction: Option, ) -> Result<(), KvSchedulerError> { - self.scheduler + self.selection + .scheduler() .enqueue_output_block_if_booking(booking, decay_fraction) .await } @@ -1589,7 +1554,7 @@ impl KvRouter { let effective_cached_tokens: HashMap = cache_hit_estimates.cached_tokens.into_iter().collect(); - Ok(self.scheduler.get_potential_loads( + Ok(self.selection.scheduler().get_potential_loads( maybe_seq_hashes, isl_tokens, effective_cached_tokens, @@ -2357,7 +2322,7 @@ mod tests { let mut trace = Vec::new(); for (index, tokens) in golden_corpus().iter().enumerate() { let context_id = format!("golden-{index}"); - let FindBestMatchInnerOutcome::WithAdmission(admitted) = router + let admitted = router .find_best_match_details_with_policy_class_inner( Some(&context_id), tokens, @@ -2379,10 +2344,7 @@ mod tests { FindBestMatchAdmission::WithAdmission, ) .await - .unwrap() - else { - panic!("admitted routing returned advisory outcome"); - }; + .unwrap(); let FindBestMatchOutcome::Routed { routing_hashes, .. } = &admitted.outcome else { panic!("golden corpus must route"); }; @@ -2420,7 +2382,7 @@ mod tests { let router = tracked_router("lease-manager").await; let worker = WorkerWithDpRank::from_worker_id(1); - let scheduler = &router.scheduler; + let scheduler = router.selection.scheduler(); let book = |request_id: &'static str| async move { let attempt_id = scheduler .add_request_admitted(SequenceRequest { @@ -2446,15 +2408,15 @@ mod tests { let before = manager.register_local(book("before").await, None); drop(before); assert!( - router.scheduler.has_request("before"), + router.selection.scheduler().has_request("before"), "no scheduler to release through yet" ); - manager.set_scheduler(router.scheduler.booking_cleanup()); + manager.set_scheduler(router.selection.scheduler().booking_cleanup()); let after = manager.register_local(book("after").await, None); after.finish().await; - assert!(!router.scheduler.has_request("after")); - assert!(router.scheduler.has_request("before")); + assert!(!router.selection.scheduler().has_request("after")); + assert!(router.selection.scheduler().has_request("before")); cancel.cancel(); } @@ -2531,9 +2493,7 @@ mod tests { ) .await .unwrap(); - let FindBestMatchInnerOutcome::WithAdmission(admitted) = outcome else { - panic!("admitted routing returned advisory outcome"); - }; + let admitted = outcome; let booking = admitted .booking .expect("tracked selection carries its booking handle"); @@ -2547,10 +2507,10 @@ mod tests { record.entered.load(Ordering::Acquire), "routing update is in flight" ); - assert!(router.scheduler.has_request("cancelled")); + assert!(router.selection.scheduler().has_request("cancelled")); drop(enroll); tokio::time::timeout(std::time::Duration::from_secs(2), async { - while router.scheduler.has_request("cancelled") { + while router.selection.scheduler().has_request("cancelled") { tokio::task::yield_now().await; } }) diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs index 6f2e155d3a63..5d014155ec0d 100644 --- a/lib/llm/src/kv_router/embedded.rs +++ b/lib/llm/src/kv_router/embedded.rs @@ -9,7 +9,6 @@ //! frontend retains transport, stream leases, and request-expiry ownership. use std::collections::HashMap; -use std::ops::Range; use std::sync::{Arc, OnceLock}; use std::time::Duration; @@ -18,24 +17,19 @@ use dynamo_kv_router::WorkerType; use dynamo_kv_router::config::KvRouterConfig; use dynamo_kv_router::identity::RoutingPartitionId; use dynamo_kv_router::protocols::{WorkerConfigLike, WorkerId, WorkerWithDpRank}; -use dynamo_kv_router::scheduling::queue::{ - ClassQueueStats, DEFAULT_MAX_BATCHED_TOKENS, SchedulerBookingCleanup, - SchedulerBookingDescriptor, -}; +use dynamo_kv_router::scheduling::queue::DEFAULT_MAX_BATCHED_TOKENS; use dynamo_kv_router::scheduling::{ - AttemptId, KvSchedulerError, NonMaxOverlapSelectionObserver, OverloadedWorkerProvider, - PotentialLoad, QueueLimitKind, QueueRejection, WorkerAvailabilityProvider, + NonMaxOverlapSelectionObserver, OverloadedWorkerProvider, QueueLimitKind, QueueRejection, + WorkerAvailabilityProvider, }; -use dynamo_kv_router::sequences::{SequenceError, SequenceRequest}; use dynamo_kv_router::services::selection::{ CatalogObserver, CatalogReconciler, DEFAULT_MODEL_NAME, HostCache, HostEligibility, HostLoad, HostReplication, HostTelemetry, KvEventIngress, KvIndexSource, SelectionHost, - SelectionOperation, SelectionOutcome, SelectionPartition, SelectionRun, SelectionService, - SelectionServiceBuilder, WorkerCatalogRecord, WorkerCatalogSource, WorkerRequest, - WorkerSelectionPolicyRegistry, + SelectionOperation, SelectionOutcome, SelectionPartition, SelectionRun, SelectionScheduler, + SelectionService, SelectionServiceBuilder, WorkerCatalogRecord, WorkerCatalogSource, + WorkerRequest, WorkerSelectionPolicyRegistry, }; use dynamo_kv_router::{DEFAULT_ROUTING_GROUP, PrefillLoadEstimator, WorkerSelectionPolicyFactory}; -use dynamo_tokens::SequenceHash; use tokio_util::sync::CancellationToken; use crate::discovery::RuntimeConfigWatch; @@ -91,12 +85,9 @@ fn record_queue_rejection( } } -fn update_queue_metrics( - per_class: &[RouterQueueMetricHandles], - mut stats_for: impl FnMut(usize) -> Option, -) { +fn update_queue_metrics(per_class: &[RouterQueueMetricHandles], scheduler: &SelectionScheduler) { for (class_index, handles) in per_class.iter().enumerate() { - let Some(stats) = stats_for(class_index) else { + let Some(stats) = scheduler.class_queue_stats(class_index) else { debug_assert!( false, "missing queue counters for policy class {class_index}" @@ -124,9 +115,7 @@ fn spawn_queue_metrics_updater( let period = Duration::from_secs(60); let mut recheck = tokio::time::interval_at(tokio::time::Instant::now() + period, period); loop { - update_queue_metrics(&handles, |index| { - partition.scheduler().class_queue_stats(index) - }); + update_queue_metrics(&handles, partition.scheduler()); tokio::select! { _ = cancellation_token.cancelled() => break, changed = queue_updates.changed() => { @@ -175,20 +164,6 @@ pub(crate) struct EmbeddedSelection { queue_metric_indices: HashMap, } -/// The partition's inbound replica ingress, not yet running. The caller starts -/// it only after the lease manager it passed as `request_leases` has its -/// scheduler set, so no lifecycle event reaches the manager before it can -/// release the booking. -pub(crate) struct PendingReplicaIngress { - ingress: crate::kv_router::sequence::ReplicaIngress, -} - -impl PendingReplicaIngress { - pub(crate) async fn start(self) { - self.ingress.start().await; - } -} - static INSTALLED_POLICY_REGISTRY: OnceLock = OnceLock::new(); /// Install the process-wide worker-selection policy registry (linked custom @@ -232,12 +207,16 @@ impl dynamo_kv_router::services::selection::SchedulerLoadSink for SenderLoadSink } impl EmbeddedSelection { + /// Returns the partition's inbound replica ingress, not yet running. The + /// caller starts it only after the lease manager it passed as + /// `request_leases` has its scheduler set, so no lifecycle event reaches + /// the manager before it can release the booking. pub(crate) async fn start( args: EmbeddedSelectionArgs, workers_with_configs: RuntimeConfigWatch, request_leases: Option>, cancellation_token: CancellationToken, - ) -> Result<(Self, PendingReplicaIngress)> { + ) -> Result<(Self, crate::kv_router::sequence::ReplicaIngress)> { let worker_type = args.worker_role.unwrap_or(WorkerType::Aggregated); let key = RoutingPartitionId::new( args.model_name @@ -301,6 +280,7 @@ impl EmbeddedSelection { .context("failed to start embedded selection service")?; let service = Arc::new(service); let partition = service + .core() .ensure_partition(key.clone(), args.block_size, args.is_eagle) .context("failed to create embedded selection partition")?; @@ -373,9 +353,7 @@ impl EmbeddedSelection { queue_metrics, queue_metric_indices, }, - PendingReplicaIngress { - ingress: replica_ingress, - }, + replica_ingress, )) } @@ -383,9 +361,7 @@ impl EmbeddedSelection { if let Some(rejection) = rejection { record_queue_rejection(&self.queue_metrics, &self.queue_metric_indices, rejection); } - update_queue_metrics(&self.queue_metrics, |index| { - self.partition.scheduler().class_queue_stats(index) - }); + update_queue_metrics(&self.queue_metrics, self.partition.scheduler()); } pub(crate) fn affinity_coordinator( @@ -418,83 +394,13 @@ impl EmbeddedSelection { run } - pub(crate) async fn add_request_admitted( - &self, - req: SequenceRequest, - ) -> Result { - self.partition.scheduler().add_request_admitted(req).await - } - - pub(crate) async fn mark_prefill_completed( - &self, - request_id: &str, - ) -> Result<(), SequenceError> { - self.partition - .scheduler() - .mark_prefill_completed(request_id) - .await - } - - pub(crate) async fn free(&self, request_id: &str) -> Result<(), SequenceError> { - self.partition.scheduler().free(request_id).await - } - - pub(crate) fn booking_cleanup(&self) -> SchedulerBookingCleanup { - self.partition.scheduler().booking_cleanup() - } - - pub(crate) async fn mark_prefill_completed_if_booking( - &self, - booking: &SchedulerBookingDescriptor, - ) -> Result<(), KvSchedulerError> { - self.partition - .scheduler() - .mark_prefill_completed_if_booking(booking) - .await - .map(|_| ()) - } - - #[cfg(test)] - pub(crate) fn has_request(&self, request_id: &str) -> bool { - self.partition.scheduler().has_request(request_id) - } - - pub(crate) fn pending_count(&self) -> usize { - self.partition.scheduler().pending_count() - } - - pub(crate) fn pending_isl_tokens(&self) -> usize { - self.partition.scheduler().pending_isl_tokens() - } - pub(crate) fn worker_type(&self) -> &'static str { self.worker_type } - pub(crate) async fn enqueue_output_block_if_booking( - &self, - booking: &SchedulerBookingDescriptor, - decay_fraction: Option, - ) -> Result<(), KvSchedulerError> { - self.partition - .scheduler() - .enqueue_output_block_if_booking(booking, decay_fraction) - .await - } - - pub(crate) fn get_potential_loads( - &self, - token_seq: Option>, - isl_tokens: usize, - effective_cached_tokens: HashMap, - track_prefill_tokens: bool, - ) -> Vec { - self.partition.scheduler().get_potential_loads( - token_seq, - isl_tokens, - effective_cached_tokens, - track_prefill_tokens, - ) + /// The partition's scheduler, for bookings and load queries. + pub(crate) fn scheduler(&self) -> &SelectionScheduler { + self.partition.scheduler() } } @@ -538,21 +444,16 @@ struct RegisteredGauge { worker_label: &'static str, } -fn dp_ranks(record: &WorkerCatalogRecord) -> Range { - let start = record.data_parallel_start_rank.unwrap_or(0); - start..start.saturating_add(record.data_parallel_size.unwrap_or(1)) -} - impl CatalogObserver for RegisteredGauge { fn upserted(&self, record: &WorkerCatalogRecord) { - for dp_rank in dp_ranks(record) { + for dp_rank in record.dp_ranks() { self.metrics .set_registered(record.worker_id, dp_rank, self.worker_label); } } fn removed(&self, record: &WorkerCatalogRecord) { - for dp_rank in dp_ranks(record) { + for dp_rank in record.dp_ranks() { self.metrics .remove_worker(record.worker_id, dp_rank, self.worker_label); } @@ -586,9 +487,6 @@ pub(crate) fn worker_request_from_runtime_config( model_name: key.model_name.clone(), routing_group: key.routing_group.clone(), endpoint: Some(format!("dyn://{worker_id}")), - kv_events_endpoint: None, - kv_events_endpoints: HashMap::new(), - replay_endpoint: None, block_size: Some(block_size), data_parallel_start_rank: Some(dp_start), data_parallel_size: Some(dp_size), @@ -610,6 +508,7 @@ pub(crate) fn worker_request_from_runtime_config( router_hint_worker_type, router_hint_source_control_endpoints, kv_event_source_mode: config.kv_event_source_mode.clone(), + ..WorkerRequest::default() } } diff --git a/lib/llm/src/kv_router/routing_host/kv_selection.rs b/lib/llm/src/kv_router/routing_host/kv_selection.rs index a0a680cbedc5..fb76a5387b78 100644 --- a/lib/llm/src/kv_router/routing_host/kv_selection.rs +++ b/lib/llm/src/kv_router/routing_host/kv_selection.rs @@ -15,10 +15,7 @@ use dynamo_kv_router::{ use dynamo_runtime::{dynamo_nvtx_range, pipeline::Error}; use crate::{ - kv_router::{ - FindBestMatchAdmission, FindBestMatchInnerOutcome, FindBestMatchOutcome, - routing_host::RoutingHost, - }, + kv_router::{FindBestMatchAdmission, FindBestMatchOutcome, routing_host::RoutingHost}, preprocessor::PreprocessedRequest, protocols::{ TokenIdType, @@ -103,7 +100,7 @@ struct BestMatchArgs<'a> { impl RoutingHost { async fn select_best_match(&self, args: BestMatchArgs<'_>) -> Result { - let outcome = self + let admitted = self .kv_router() .find_best_match_details_with_policy_class_inner( Some(args.context_id), @@ -126,55 +123,29 @@ impl RoutingHost { args.admission, ) .await?; - match outcome { - FindBestMatchInnerOutcome::WithAdmission(admitted) => match admitted.outcome { - FindBestMatchOutcome::Routed { - worker, - overlap_blocks, - effective_overlap_blocks, - cached_tokens, - potential_decode_blocks, - routing_hashes, - kv_hint, - } => Ok(SelectionOutcome::Routed(WorkerSelection { - worker, - booking: admitted.booking, - overlap_amount: overlap_blocks, - effective_overlap_blocks, - cached_tokens, - potential_decode_blocks, - selected_worker_load: None, - routing_hashes, - kv_hint, - })), - FindBestMatchOutcome::QueueRejected { rejection } => { - Ok(SelectionOutcome::QueueRejected(rejection)) - } - }, - FindBestMatchInnerOutcome::WithoutAdmission(outcome) => match outcome { - crate::kv_router::FindBestMatchAdvisoryOutcome::Routed { - worker, - overlap_blocks, - effective_overlap_blocks, - cached_tokens, - potential_decode_blocks, - selected_worker_load, - routing_hashes, - } => Ok(SelectionOutcome::Routed(WorkerSelection { - worker, - booking: None, - overlap_amount: overlap_blocks, - effective_overlap_blocks, - cached_tokens, - potential_decode_blocks, - selected_worker_load: Some(selected_worker_load), - routing_hashes, - kv_hint: None, - })), - crate::kv_router::FindBestMatchAdvisoryOutcome::QueueRejected { rejection } => { - Ok(SelectionOutcome::QueueRejected(rejection)) - } - }, + match admitted.outcome { + FindBestMatchOutcome::Routed { + worker, + overlap_blocks, + effective_overlap_blocks, + cached_tokens, + potential_decode_blocks, + routing_hashes, + kv_hint, + } => Ok(SelectionOutcome::Routed(WorkerSelection { + worker, + booking: admitted.booking, + overlap_amount: overlap_blocks, + effective_overlap_blocks, + cached_tokens, + potential_decode_blocks, + selected_worker_load: admitted.advisory_load, + routing_hashes, + kv_hint, + })), + FindBestMatchOutcome::QueueRejected { rejection } => { + Ok(SelectionOutcome::QueueRejected(rejection)) + } } } From 99f8574f50b0c9e226efcae44e08bc08b1750601 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Sat, 12 Sep 2026 14:43:51 -0700 Subject: [PATCH 070/102] chore(router): cleanup round three, verbatim moves and shorter tests Last behavior-neutral round. Where the branch had renamed or reworded code it only moved from `lib/llm`, main's text is restored so the hunks read as moves: `PrimaryLookup`/`SideIndexer` variant names, the `record_routing_decision_hashes` name and method order, the `merge_overlap_scores` doc, the affinity table's field and log wording, the sequence module layout, the runtime ingress body (destructured once), the recording SPDX year. Dead or duplicated pieces: the unused `Debug` impls on `SelectionHost` and `KvIndexSource`, `affinity_mode()` and `Selector::core()` accessors, `BookingHandle::descriptor`, `Indexer::is_remote` outside tests, a `publish_scheduler_config` result nobody could receive, and the third copy of the session-affinity TTL range check in the bindings. Tests keep every assertion and lose repeated literals: `default_key`, `saturated_core`, `lease_operation`, `bound_session`, `bound_worker` (which now requires a configured affinity table, so `None` means unbound), `builtin_host_with_affinity`, `normalize_with`, one `single_rank(5557)` binding per EPP discovery test, and `make_router` taking the role and label the two open-coded constructors passed. Doc comments that restated the code are gone; three that stated a rule are back: which crate implements the LoRA filter and why the pin wins, that an unregistered worker deletes idempotently, and that lookups live in `services/indexer/lookup.rs`. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Ishan Dhanani --- deploy/inference-gateway/ext-proc/src/epp.rs | 5 +- .../ext-proc/src/pod_discovery.rs | 112 ++----- .../ext-proc/src/selector.rs | 15 +- .../ext-proc/src/topology_adapter.rs | 3 +- lib/bindings/python/rust/llm/entrypoint.rs | 9 +- lib/bindings/python/rust/llm/kv.rs | 16 +- lib/kv-router/src/scheduling/lora_filter.rs | 18 +- lib/kv-router/src/scheduling/queue.rs | 4 - .../src/services/common/replica_sync.rs | 4 - lib/kv-router/src/services/indexer/backend.rs | 44 +-- lib/kv-router/src/services/indexer/lookup.rs | 69 ++-- .../src/services/indexer/recording.rs | 34 +- .../src/services/selection/README.md | 6 +- .../src/services/selection/affinity.rs | 51 ++- .../src/services/selection/core/hint.rs | 3 +- .../src/services/selection/core/mod.rs | 44 +-- .../src/services/selection/core/run.rs | 5 +- .../src/services/selection/core/tests.rs | 313 ++++++------------ .../src/services/selection/core/workers.rs | 16 +- .../src/services/selection/membership.rs | 1 + lib/kv-router/src/services/selection/tests.rs | 83 ++--- lib/kv-router/src/services/selection/types.rs | 5 +- lib/llm/src/discovery/model_manager.rs | 48 +-- lib/llm/src/discovery/watcher.rs | 2 +- lib/llm/src/kv_router.rs | 138 ++++---- lib/llm/src/kv_router/AGENTS.md | 2 +- lib/llm/src/kv_router/embedded.rs | 2 +- lib/llm/src/kv_router/indexer/ingress.rs | 62 ++-- lib/llm/src/kv_router/indexer/remote.rs | 66 ++-- lib/llm/src/kv_router/prefill_router/query.rs | 3 +- lib/llm/src/kv_router/routing_host.rs | 16 +- .../kv_router/routing_host/request_guard.rs | 2 - lib/llm/src/kv_router/routing_host/tests.rs | 90 ++--- lib/llm/src/kv_router/sequence.rs | 36 +- lib/llm/src/session_affinity/coordinator.rs | 2 - 35 files changed, 507 insertions(+), 822 deletions(-) diff --git a/deploy/inference-gateway/ext-proc/src/epp.rs b/deploy/inference-gateway/ext-proc/src/epp.rs index a3e9367f29df..34daa35fab9f 100644 --- a/deploy/inference-gateway/ext-proc/src/epp.rs +++ b/deploy/inference-gateway/ext-proc/src/epp.rs @@ -2173,10 +2173,7 @@ mod tests { } /// Under pod discovery a worker registers under its pod identity alone, so - /// a pod's ready sidecars must contribute no worker ids. Emitting them - /// would name workers no backend registered under; the scheduler filters - /// them out, so the pod's real worker would be the subset's only survivor - /// only by accident of the id set. + /// a pod's ready sidecars must contribute no worker ids. #[test] fn pod_worker_ids_ignores_containers_under_pod_discovery() { let pod = pod_mode_worker_pod(); diff --git a/deploy/inference-gateway/ext-proc/src/pod_discovery.rs b/deploy/inference-gateway/ext-proc/src/pod_discovery.rs index 1920c1f2f5b4..bdd652c1a417 100644 --- a/deploy/inference-gateway/ext-proc/src/pod_discovery.rs +++ b/deploy/inference-gateway/ext-proc/src/pod_discovery.rs @@ -666,6 +666,7 @@ mod tests { #[test] fn rebuild_index_keeps_only_ready_selected_pods() { + let kp = single_rank(5557); let store = store_from_pods(vec![ pod( "vllm-0", @@ -683,17 +684,11 @@ mod tests { ]); let index = RwLock::new(WorkerIndex::new()); - assert!(rebuild_index( - &store, - Some(&pool()), - single_rank(5557), - Some(5560), - &index - )); + assert!(rebuild_index(&store, Some(&pool()), kp, Some(5560), &index)); assert!(!rebuild_index( &store, Some(&pool()), - single_rank(5557), + kp, Some(5560), &index )); @@ -710,6 +705,7 @@ mod tests { #[test] fn rebuild_index_is_empty_without_pool() { + let kp = single_rank(5557); let store = store_from_pods(vec![pod( "vllm-0", Some("10.0.0.1"), @@ -717,18 +713,13 @@ mod tests { &[("app", "vllm-qwen")], )]); let index = RwLock::new(WorkerIndex::new()); - assert!(!rebuild_index( - &store, - None, - single_rank(5557), - None, - &index - )); + assert!(!rebuild_index(&store, None, kp, None, &index)); assert!(index.read().unwrap().is_empty()); } #[test] fn upsert_and_remove_pod_mutate_index_incrementally() { + let kp = single_rank(5557); let index = RwLock::new(WorkerIndex::new()); let id = hash_pod_name("vllm-0"); let ready = pod( @@ -739,20 +730,8 @@ mod tests { ); // Ready + selected -> inserted, with a pre-stripped endpoint. - assert!(upsert_pod( - &index, - &ready, - Some(&pool()), - single_rank(5557), - None - )); - assert!(!upsert_pod( - &index, - &ready, - Some(&pool()), - single_rank(5557), - None - )); + assert!(upsert_pod(&index, &ready, Some(&pool()), kp, None)); + assert!(!upsert_pod(&index, &ready, Some(&pool()), kp, None)); assert_eq!( index.read().unwrap().get(&id).map(|e| e.endpoint.as_str()), Some("10.0.0.1:8000") @@ -765,30 +744,12 @@ mod tests { Some(false), &[("app", "vllm-qwen")], ); - assert!(upsert_pod( - &index, - ¬_ready, - Some(&pool()), - single_rank(5557), - None - )); - assert!(!upsert_pod( - &index, - ¬_ready, - Some(&pool()), - single_rank(5557), - None - )); + assert!(upsert_pod(&index, ¬_ready, Some(&pool()), kp, None)); + assert!(!upsert_pod(&index, ¬_ready, Some(&pool()), kp, None)); assert!(!index.read().unwrap().contains_key(&id)); // Re-add, then a Delete removes it. - assert!(upsert_pod( - &index, - &ready, - Some(&pool()), - single_rank(5557), - None - )); + assert!(upsert_pod(&index, &ready, Some(&pool()), kp, None)); assert!(index.read().unwrap().contains_key(&id)); assert!(remove_pod(&index, &ready)); assert!(!remove_pod(&index, &ready)); @@ -796,17 +757,12 @@ mod tests { // An unrelated namespace pod does not change the derived worker index. let unselected = pod("other-0", Some("10.0.0.2"), Some(true), &[("app", "other")]); - assert!(!upsert_pod( - &index, - &unselected, - Some(&pool()), - single_rank(5557), - None - )); + assert!(!upsert_pod(&index, &unselected, Some(&pool()), kp, None)); } #[test] fn upsert_pod_without_pool_drops_entry() { + let kp = single_rank(5557); // A `None` pool (unresolved or deleted) means nothing is routable, so an // upsert must evict any existing entry rather than leave stale routing. let index = RwLock::new(WorkerIndex::new()); @@ -818,22 +774,17 @@ mod tests { &[("app", "vllm-qwen")], ); - assert!(upsert_pod( - &index, - &ready, - Some(&pool()), - single_rank(5557), - None - )); + assert!(upsert_pod(&index, &ready, Some(&pool()), kp, None)); assert!(index.read().unwrap().contains_key(&id)); - assert!(upsert_pod(&index, &ready, None, single_rank(5557), None)); - assert!(!upsert_pod(&index, &ready, None, single_rank(5557), None)); + assert!(upsert_pod(&index, &ready, None, kp, None)); + assert!(!upsert_pod(&index, &ready, None, kp, None)); assert!(!index.read().unwrap().contains_key(&id)); } #[test] fn pool_edit_during_relist_rebuilds_at_init_done_from_completed_store() { + let kp = single_rank(5557); use kube::runtime::watcher; let vllm_0 = pod( @@ -857,13 +808,7 @@ mod tests { writer.apply_watcher_event(&watcher::Event::InitApply(vllm_1.clone())); writer.apply_watcher_event(&watcher::Event::InitDone); let index = RwLock::new(WorkerIndex::new()); - assert!(rebuild_index( - &store, - Some(&pool()), - single_rank(5557), - None, - &index - )); + assert!(rebuild_index(&store, Some(&pool()), kp, None, &index)); assert_eq!(index.read().unwrap().len(), 2); // During the next LIST, vllm-1 has disappeared. The reflector buffers @@ -889,13 +834,7 @@ mod tests { // InitDone makes the staged list live. Its single rebuild uses both the // completed one-pod Store and the latest PoolState. writer.apply_watcher_event(&watcher::Event::InitDone); - assert!(rebuild_index( - &store, - Some(&updated_pool), - single_rank(5557), - None, - &index - )); + assert!(rebuild_index(&store, Some(&updated_pool), kp, None, &index)); let index = index.read().unwrap(); assert_eq!(index.len(), 1); assert_eq!( @@ -912,6 +851,7 @@ mod tests { #[test] fn rebuild_index_drops_workers_absent_from_the_store() { + let kp = single_rank(5557); // A relist arrives as a fresh snapshot with no `Delete` events for pods // that disappeared during the disconnect, so the `InitDone`/pool-change // rebuild must *replace* the index, not merge into it — otherwise dead @@ -927,7 +867,7 @@ mod tests { &[("app", "vllm-qwen")], ), Some(&pool()), - single_rank(5557), + kp, None, ); upsert_pod( @@ -939,7 +879,7 @@ mod tests { &[("app", "vllm-qwen")], ), Some(&pool()), - single_rank(5557), + kp, None, ); assert_eq!(index.read().unwrap().len(), 2); @@ -951,13 +891,7 @@ mod tests { Some(true), &[("app", "vllm-qwen")], )]); - assert!(rebuild_index( - &store, - Some(&pool()), - single_rank(5557), - None, - &index - )); + assert!(rebuild_index(&store, Some(&pool()), kp, None, &index)); let index = index.read().unwrap(); assert_eq!(index.len(), 1); diff --git a/deploy/inference-gateway/ext-proc/src/selector.rs b/deploy/inference-gateway/ext-proc/src/selector.rs index e00c9499e005..8bdb113e53e0 100644 --- a/deploy/inference-gateway/ext-proc/src/selector.rs +++ b/deploy/inference-gateway/ext-proc/src/selector.rs @@ -15,8 +15,8 @@ use anyhow::{Context, Result, anyhow}; use dynamo_kv_router::config::{KvRouterConfig, try_kv_router_config_from_dynamo_env}; use dynamo_kv_router::protocols::RoutingConstraints; use dynamo_kv_router::services::selection::{ - PromptRequest, SelectAndReserveRequest as CoreSelectAndReserveRequest, SelectionCore, - SelectionError, SelectionService, SelectionServiceBuilder, WorkerSelectionPolicyRegistry, + PromptRequest, SelectAndReserveRequest as CoreSelectAndReserveRequest, SelectionError, + SelectionService, SelectionServiceBuilder, WorkerSelectionPolicyRegistry, warn_for_unserved_worker_selection_policies, }; use dynamo_kv_router::{DEFAULT_ROUTING_GROUP, WorkerType}; @@ -66,7 +66,7 @@ pub struct SelectResponse { /// In-process runtime-free selector wrapping a [`SelectionService`]. pub struct Selector { - service: Arc, + pub(crate) service: Arc, /// Cancels the peer-discovery watch on drop. The `SelectionService`'s own /// `Drop` tears down its core + replica-sync tasks. cancel: CancellationToken, @@ -166,11 +166,6 @@ impl Selector { Ok(()) } - /// The selection core whose catalog the topology adapter feeds. - pub(crate) fn core(&self) -> &Arc { - self.service.core() - } - /// Select a worker for a prompt and book its load in one operation. Takes the /// request by value so per-request fields are moved into the core request /// rather than cloned on the hot path. @@ -401,7 +396,7 @@ models: } async fn register(selector: &Selector, workers: Vec) { - CatalogReconciler::new(Arc::clone(selector.core())) + CatalogReconciler::new(Arc::clone(selector.service.core())) .apply(workers) .await .expect("reconcile should succeed"); @@ -748,7 +743,7 @@ worker_selection: .expect("selector should build"); let duplicate = incomplete_registration(1); - let error = CatalogReconciler::new(Arc::clone(selector.core())) + let error = CatalogReconciler::new(Arc::clone(selector.service.core())) .apply(vec![duplicate.clone(), duplicate]) .await .expect_err("duplicate IDs must be rejected"); diff --git a/deploy/inference-gateway/ext-proc/src/topology_adapter.rs b/deploy/inference-gateway/ext-proc/src/topology_adapter.rs index 34a9fb48f644..244381551259 100644 --- a/deploy/inference-gateway/ext-proc/src/topology_adapter.rs +++ b/deploy/inference-gateway/ext-proc/src/topology_adapter.rs @@ -95,7 +95,8 @@ impl TopologyAdapter { closed: false, }; tokio::spawn( - CatalogReconciler::new(Arc::clone(selector.core())).run(source, cancel.child_token()), + CatalogReconciler::new(Arc::clone(selector.service.core())) + .run(source, cancel.child_token()), ); Self { cancel } } diff --git a/lib/bindings/python/rust/llm/entrypoint.rs b/lib/bindings/python/rust/llm/entrypoint.rs index ad63a0913f1c..7adaf99b4a41 100644 --- a/lib/bindings/python/rust/llm/entrypoint.rs +++ b/lib/bindings/python/rust/llm/entrypoint.rs @@ -32,7 +32,6 @@ use dynamo_llm::local_model::{LocalModel, LocalModelBuilder}; use dynamo_llm::mocker::make_mocker_engine; use dynamo_llm::model_card::ModelDeploymentCard as RsModelDeploymentCard; use dynamo_llm::reasoning_field::ReasoningField; -use dynamo_llm::session_affinity::MAX_SESSION_AFFINITY_TTL_SECS; use dynamo_llm::session_affinity::SessionAffinityMode as RsSessionAffinityMode; use dynamo_llm::types::openai::chat_completions::OpenAIChatCompletionsStreamingEngine; use dynamo_mocker::common::perf_model::PerfModel; @@ -508,13 +507,7 @@ impl RouterConfig { ); }); } - if session_affinity_ttl_secs - .is_some_and(|ttl| !(1..=MAX_SESSION_AFFINITY_TTL_SECS).contains(&ttl)) - { - return Err(PyValueError::new_err(format!( - "session_affinity_ttl_secs must be between 1 and {MAX_SESSION_AFFINITY_TTL_SECS}" - ))); - } + super::kv::check_session_affinity_ttl_secs(session_affinity_ttl_secs)?; RsLoadThresholdConfig { active_decode_blocks_threshold, active_prefill_tokens_threshold, diff --git a/lib/bindings/python/rust/llm/kv.rs b/lib/bindings/python/rust/llm/kv.rs index d9bf1e1fd97f..c0503466d2ef 100644 --- a/lib/bindings/python/rust/llm/kv.rs +++ b/lib/bindings/python/rust/llm/kv.rs @@ -640,7 +640,6 @@ impl SelectionCacheConfig { } } -#[cfg(feature = "select-service")] fn session_affinity_ttl_from_secs(ttl: f64) -> Result { if !(1.0..=MAX_SESSION_AFFINITY_TTL_SECS as f64).contains(&ttl) { return Err(format!( @@ -650,6 +649,13 @@ fn session_affinity_ttl_from_secs(ttl: f64) -> Result { Ok(Duration::from_secs_f64(ttl)) } +/// Range check for a whole-second TTL; `None` passes. +pub(crate) fn check_session_affinity_ttl_secs(ttl: Option) -> PyResult<()> { + ttl.map(|ttl| session_affinity_ttl_from_secs(ttl as f64).map_err(PyValueError::new_err)) + .transpose()?; + Ok(()) +} + /// In-process handle to a managed Dynamo `SelectionService`. #[cfg(feature = "select-service")] #[pyclass] @@ -2230,13 +2236,7 @@ impl KvRouter { load_threshold_config: Option<&LoadThresholdConfig>, session_affinity_mode: &str, ) -> PyResult { - if session_affinity_ttl_secs - .is_some_and(|ttl| !(1..=MAX_SESSION_AFFINITY_TTL_SECS).contains(&ttl)) - { - return Err(PyValueError::new_err(format!( - "session_affinity_ttl_secs must be between 1 and {MAX_SESSION_AFFINITY_TTL_SECS}" - ))); - } + check_session_affinity_ttl_secs(session_affinity_ttl_secs)?; let session_affinity_mode = session_affinity_mode .parse::() .map_err(PyValueError::new_err)?; diff --git a/lib/kv-router/src/scheduling/lora_filter.rs b/lib/kv-router/src/scheduling/lora_filter.rs index 32416c3de2bf..74a3ef76c81a 100644 --- a/lib/kv-router/src/scheduling/lora_filter.rs +++ b/lib/kv-router/src/scheduling/lora_filter.rs @@ -9,10 +9,8 @@ use crate::protocols::{WorkerId, WorkerWithDpRank}; /// Host-owned knowledge of which workers can serve a LoRA adapter. /// -/// The frontend implements this over its LoRA routing table and state tracker; -/// a runtime-free host may back it with a static map. Implementations return a -/// subset of `available`, ordered arbitrarily. Returning `available` unchanged -/// means "no LoRA-specific preference". +/// Implementations return a subset of `available`, ordered arbitrarily. +/// Returning `available` unchanged means "no LoRA-specific preference". pub trait LoraWorkerFilter: Send + Sync { fn filter_worker_ids_for_lora(&self, lora_name: &str, available: &[WorkerId]) -> Vec; } @@ -22,11 +20,13 @@ pub trait LoraWorkerFilter: Send + Sync { /// /// - Without a LoRA name the caller's allow-set is returned untouched. /// - The universe is the caller's allow-set when present, else `all_workers`. -/// - A pinned worker inside the universe is always retained, even if the filter -/// dropped it, so the pin still wins for cache correctness. A pin outside the -/// universe is not re-added: the caller's constraint stands. -/// - An empty narrowing falls back to the caller's allow-set rather than -/// producing an unroutable request. +/// - Pinned worker: KV-cache correctness wins. A pin inside the universe is +/// always retained even if not in the LoRA replica set (the worker lazy-loads +/// the adapter); a pin outside the universe is not re-added, the caller's +/// constraint stands. +/// - If narrowing would exclude every candidate, falls back to the caller's +/// allow-set so the request stays routable (lazy-load path) rather than +/// failing. #[cfg_attr(not(feature = "standalone-selection"), allow(dead_code))] pub fn narrow_allowed_worker_ids_by_lora( filter: &dyn LoraWorkerFilter, diff --git a/lib/kv-router/src/scheduling/queue.rs b/lib/kv-router/src/scheduling/queue.rs index 01effd45c202..7d8a74e6eb38 100644 --- a/lib/kv-router/src/scheduling/queue.rs +++ b/lib/kv-router/src/scheduling/queue.rs @@ -383,10 +383,6 @@ impl std::fmt::Debug for BookingHandle { } impl BookingHandle { - pub fn descriptor(&self) -> &SchedulerBookingDescriptor { - &self.booking - } - /// Hand the booking to a longer-lived owner; the handle stops guarding it. #[must_use] pub fn commit(mut self) -> SchedulerBookingDescriptor { diff --git a/lib/kv-router/src/services/common/replica_sync.rs b/lib/kv-router/src/services/common/replica_sync.rs index b425239a0de9..afe17399d4c9 100644 --- a/lib/kv-router/src/services/common/replica_sync.rs +++ b/lib/kv-router/src/services/common/replica_sync.rs @@ -264,7 +264,6 @@ impl ScopedSequencePublisher { } } - /// Publish through an embedding host's outbound channel. pub(crate) fn host(outbound: mpsc::Sender) -> Self { Self { replica: None, @@ -273,7 +272,6 @@ impl ScopedSequencePublisher { } } - /// Forward scheduler load snapshots to `sink` in addition to replica sync. pub(crate) fn with_load_sink(mut self, sink: Option>) -> Self { self.load_sink = sink; self @@ -598,8 +596,6 @@ impl PeerManager { ) } - /// Like [`Self::start`], also subscribing to session-affinity bindings - /// when `handle_affinity` is given. pub(crate) fn start_with_affinity( initial_peers: Vec, cancel_token: CancellationToken, diff --git a/lib/kv-router/src/services/indexer/backend.rs b/lib/kv-router/src/services/indexer/backend.rs index 22fad6388a42..144910aa956b 100644 --- a/lib/kv-router/src/services/indexer/backend.rs +++ b/lib/kv-router/src/services/indexer/backend.rs @@ -188,8 +188,6 @@ impl IndexerPolicy { } } - /// Whether the primary is served by a remote indexer service. Such a - /// service does not listen for worker KV events itself. pub fn is_remote(&self) -> bool { matches!(self.primary, PrimaryRetention::Remote(_)) } @@ -256,7 +254,7 @@ impl IndexerPolicy { /// replayed, and never used to seed lower-tier lookups. #[derive(Clone)] pub enum SideIndexer { - Single(KvIndexer), + KvIndexer(KvIndexer), Concurrent(Arc>), } @@ -279,7 +277,7 @@ impl SideIndexer { prune_config, ))); } - Self::Single(KvIndexer::new_with_pruning( + Self::KvIndexer(KvIndexer::new_with_pruning( cancel, block_size, metrics, @@ -292,7 +290,7 @@ impl SideIndexer { sequence: HashInput<'_>, ) -> std::result::Result { match self { - Self::Single(indexer) => { + Self::KvIndexer(indexer) => { indexer .find_matches(sequence.into_owned_at_boundary()) .await @@ -309,7 +307,7 @@ impl SideIndexer { hashes: RoutingDecisionHashes, ) -> std::result::Result<(), KvRouterError> { match self { - Self::Single(indexer) => { + Self::KvIndexer(indexer) => { indexer .process_routing_decision_with_hashes( worker, @@ -336,7 +334,7 @@ impl SideIndexer { dp_rank: DpRank, ) -> std::result::Result<(), KvRouterError> { match self { - Self::Single(indexer) => { + Self::KvIndexer(indexer) => { indexer .reset_worker_dp_rank_and_wait(worker_id, dp_rank) .await @@ -351,14 +349,14 @@ impl SideIndexer { async fn remove_worker(&self, worker_id: WorkerId) { match self { - Self::Single(indexer) => indexer.remove_worker(worker_id).await, + Self::KvIndexer(indexer) => indexer.remove_worker(worker_id).await, Self::Concurrent(indexer) => indexer.remove_worker(worker_id).await, } } async fn remove_worker_dp_rank(&self, worker_id: WorkerId, dp_rank: u32) { match self { - Self::Single(indexer) => indexer.remove_worker_dp_rank(worker_id, dp_rank).await, + Self::KvIndexer(indexer) => indexer.remove_worker_dp_rank(worker_id, dp_rank).await, Self::Concurrent(indexer) => indexer.remove_worker_dp_rank(worker_id, dp_rank).await, } } @@ -453,17 +451,15 @@ impl Indexer { } } - pub fn is_remote(&self) -> bool { + #[cfg(test)] + pub(crate) fn is_remote(&self) -> bool { matches!(self, Self::Remote { .. }) } - /// Dequeue-time overlap refresh needs a local primary. pub fn supports_overlap_refresh(&self) -> bool { matches!(self, Self::Single { .. } | Self::Concurrent { .. }) } - /// Router-hint chain retention needs a local primary whose hashes come - /// only from engine events. pub fn supports_kv_transfer_chain_retention(&self) -> bool { matches!( self, @@ -1279,7 +1275,7 @@ mod tests { async fn flush_side(approx: &Option) { match approx { - Some(SideIndexer::Single(side)) => { + Some(SideIndexer::KvIndexer(side)) => { let _ = side.flush().await; } Some(SideIndexer::Concurrent(side)) => { @@ -1310,8 +1306,6 @@ mod tests { } impl Indexer { - /// Test-only: wait until every locally enqueued event and routing - /// decision is applied. pub(crate) async fn flush(&self) { flush(self).await } @@ -1324,7 +1318,7 @@ mod tests { assert!(indexer.supports_kv_transfer_chain_retention()); let worker = WorkerWithDpRank::new(7, 0); indexer - .record_routing_decision( + .record_routing_decision_hashes( worker, RoutingDecisionHashes::from_local_hashes(vec![LocalBlockHash(11)]), ) @@ -1352,7 +1346,7 @@ mod tests { assert!(!indexer.supports_kv_transfer_chain_retention()); let worker = WorkerWithDpRank::new(7, 0); indexer - .record_routing_decision( + .record_routing_decision_hashes( worker, RoutingDecisionHashes::from_local_hashes(vec![ LocalBlockHash(11), @@ -1408,7 +1402,7 @@ mod tests { // Routing decisions: worker 8 was just routed [11, 12, 13]; worker 7 // was routed [11] (shorter than what the engine confirmed). indexer - .record_routing_decision( + .record_routing_decision_hashes( predicted, RoutingDecisionHashes::from_local_hashes(vec![ LocalBlockHash(11), @@ -1419,7 +1413,7 @@ mod tests { .await .unwrap(); indexer - .record_routing_decision( + .record_routing_decision_hashes( confirmed, RoutingDecisionHashes::from_local_hashes(vec![LocalBlockHash(11)]), ) @@ -1495,13 +1489,7 @@ mod tests { .with_remote_indexer(base_url.clone()) .unwrap(); assert!(policy.is_remote()); - let indexer = create_indexer_with_policy( - &key, - 4, - 1, - Arc::new(KvIndexerMetrics::new_unregistered()), - &policy, - ); + let indexer = policy_indexer(1, policy.clone()); assert!(indexer.is_remote()); assert!(indexer.records_routing_decisions()); assert!(!indexer.supports_kv_transfer_chain_retention()); @@ -1523,7 +1511,7 @@ mod tests { // A routing decision lands in the local side indexer only. let predicted = WorkerWithDpRank::new(8, 0); indexer - .record_routing_decision( + .record_routing_decision_hashes( predicted, RoutingDecisionHashes::from_local_hashes(vec![LocalBlockHash(11)]), ) diff --git a/lib/kv-router/src/services/indexer/lookup.rs b/lib/kv-router/src/services/indexer/lookup.rs index fca609242417..ae439b5f1ef0 100644 --- a/lib/kv-router/src/services/indexer/lookup.rs +++ b/lib/kv-router/src/services/indexer/lookup.rs @@ -24,8 +24,8 @@ pub enum HashInput<'a> { impl<'a> HashInput<'a> { pub fn as_slice(&self) -> &[LocalBlockHash] { match self { - Self::Borrowed(hashes) => hashes, - Self::Owned(hashes) => hashes.as_slice(), + Self::Borrowed(sequence) => sequence, + Self::Owned(sequence) => sequence, } } @@ -35,20 +35,22 @@ impl<'a> HashInput<'a> { pub(super) fn into_owned_at_boundary(self) -> Vec { match self { - Self::Borrowed(hashes) => hashes.to_vec(), - Self::Owned(hashes) => hashes, + Self::Borrowed(sequence) => sequence.to_vec(), + Self::Owned(sequence) => sequence, } } } +#[derive(Clone, Copy)] pub(super) struct LookupPipeline<'a> { primary: PrimaryLookup<'a>, lower_tier: Option<&'a LowerTierIndexers>, side: Option<&'a SideIndexer>, } +#[derive(Clone, Copy)] enum PrimaryLookup<'a> { - Single(&'a KvIndexer), + KvIndexer(&'a KvIndexer), Concurrent(&'a ThreadPoolIndexer), Remote(&'a dyn RemotePrimary), None, @@ -63,7 +65,7 @@ impl Indexer { approx, .. } => LookupPipeline { - primary: PrimaryLookup::Single(primary), + primary: PrimaryLookup::KvIndexer(primary), lower_tier: Some(lower_tier), side: approx.as_ref(), }, @@ -188,7 +190,7 @@ impl<'a> LookupPipeline<'a> { lower_tier_options: LowerTierQueryOptions, ) -> Result { match self.primary { - PrimaryLookup::Single(_) | PrimaryLookup::Concurrent(_) => { + PrimaryLookup::KvIndexer(_) | PrimaryLookup::Concurrent(_) => { let Some(lower_tier) = self.lower_tier else { return Ok(TieredMatchDetails::default()); }; @@ -220,7 +222,7 @@ impl<'a> LookupPipeline<'a> { PrimaryLookup::Remote(primary) => { if lower_tier_options.retain_kv_transfer_chain { tracing::warn!( - "router_hint chain retention is not supported with remote primary indexer; proceeding without router hints" + "KV transfer chain retention is not supported with remote primary indexer; proceeding without KV transfer hints" ); } let Some(side) = self.side else { @@ -251,7 +253,7 @@ impl<'a> LookupPipeline<'a> { sequence: HashInput<'_>, ) -> Result { match self.primary { - PrimaryLookup::Single(_) | PrimaryLookup::Concurrent(_) => { + PrimaryLookup::KvIndexer(_) | PrimaryLookup::Concurrent(_) => { let Some(lower_tier) = self.lower_tier else { return Ok(TieredMatchDetails::default()); }; @@ -285,7 +287,7 @@ impl<'a> PrimaryLookup<'a> { sequence: HashInput<'_>, ) -> Result { let primary_details = match self { - Self::Single(primary) => { + Self::KvIndexer(primary) => { primary .find_match_details(sequence.into_owned_at_boundary()) .await? @@ -323,7 +325,7 @@ impl<'a> PrimaryLookup<'a> { retain_kv_transfer_chain: bool, ) -> Result { let primary_details = match self { - Self::Single(primary) => { + Self::KvIndexer(primary) => { primary .find_match_details_with_options( sequence.clone_for_boundary(), @@ -354,21 +356,19 @@ impl<'a> PrimaryLookup<'a> { } /// Merge a side-indexer's `OverlapScores` into the primary's `MatchDetails` -/// by per-worker max. The side indexer covers the window before the engine's -/// first KV event arrives; `last_matched_hashes`, `frequencies`, and -/// `tree_sizes` stay the primary's. +/// by taking the per-worker max overlap. The side indexer covers the window +/// before the engine's first KV event arrives; for workers it knows about, +/// we use whichever indexer saw the longer prefix. `last_matched_hashes`, +/// `frequencies`, and `tree_sizes` come from the primary -- the side +/// indexer's short-TTL view isn't meaningful for those signals. /// -/// The result no longer satisfies the `scores` <-> `last_matched_hashes` -/// lockstep (side-only workers gain a score with no paired hash), so it must -/// not seed `query_lower_tiers`; callers run the lower-tier query against the -/// primary-only details first. -/// -/// NOTE: when this merged `MatchDetails` is combined with lower-tier hits -/// seeded from the primary-only anchor (e.g. in `find_matches_by_tier`), the -/// total cached-token signal can overcount: the device score is raised by the -/// side indexer but the lower-tier walk used the lower primary depth. Accepted -/// since side scores are short-TTL approximations and the overcount is bounded -/// and rare in practice. +/// IMPORTANT: the returned `MatchDetails` is no longer guaranteed to satisfy +/// `overlap_scores.scores` <-> `last_matched_hashes` lockstep. Side-only +/// workers gain a score with no paired hash by design. The result is safe +/// for scheduling / cache-hit signal but MUST NOT be used to seed +/// `query_lower_tiers`, which assumes the lockstep invariant. The local +/// arm of `find_matches_by_tier` enforces this by running the lower-tier +/// query against primary-only `MatchDetails` before merging side scores. fn merge_overlap_scores(mut primary: MatchDetails, side: OverlapScores) -> MatchDetails { for (worker, side_score) in side.scores { primary @@ -385,9 +385,18 @@ fn merge_overlap_scores(mut primary: MatchDetails, side: OverlapScores) -> Match primary } -/// Query the side indexer (if present) and merge its scores into `primary`. -/// On query error, warn and return `primary` unchanged so the caller still has -/// a usable scheduling signal. +/// Query the predict-on-route side indexer (if present) and merge its scores +/// into the primary device match details. Side scores never feed lower-tier +/// or shared-cache scoring. On query error, log a warning and return `primary` +/// unchanged so the caller still has a usable scheduling signal. See +/// [`merge_overlap_scores`] for the lockstep caveat on the returned shape. +/// +/// NOTE: when this merged `MatchDetails` is combined with lower-tier hits +/// seeded from the primary-only anchor (e.g. in `find_matches_by_tier`), the +/// total cached-token signal can in theory overcount: the device score is +/// raised by the side indexer but the lower-tier walk used the lower primary +/// depth. Accepted as edge for now since side scores are short-TTL +/// approximations and the overcount is bounded and rare in practice. pub(super) async fn merge_side_or_warn( side: Option<&SideIndexer>, primary: MatchDetails, @@ -825,7 +834,7 @@ mod tests { let local_hashes = vec![LocalBlockHash(91), LocalBlockHash(92)]; let sequence_hashes = compute_seq_hash_for_block(&local_hashes); indexer - .record_routing_decision( + .record_routing_decision_hashes( worker, RoutingDecisionHashes { local_hashes: local_hashes.clone(), @@ -1044,7 +1053,7 @@ mod tests { let side_hashes = vec![LocalBlockHash(11), LocalBlockHash(12), LocalBlockHash(13)]; indexer - .record_routing_decision( + .record_routing_decision_hashes( side_worker, RoutingDecisionHashes { local_hashes: side_hashes.clone(), diff --git a/lib/kv-router/src/services/indexer/recording.rs b/lib/kv-router/src/services/indexer/recording.rs index dc80abcfcc0a..95e484d0d453 100644 --- a/lib/kv-router/src/services/indexer/recording.rs +++ b/lib/kv-router/src/services/indexer/recording.rs @@ -1,4 +1,4 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 //! Routing-decision recording: which physical index a booked routing decision @@ -117,7 +117,6 @@ impl Indexer { }) } - /// Whether booked routing decisions should be written into this indexer. pub fn records_routing_decisions(&self) -> bool { !matches!(self.recording_target(), RouteRecordingTarget::Disabled) } @@ -162,32 +161,31 @@ impl Indexer { } } - /// Record a booked routing decision. Writes to the side indexer when one - /// is attached, else to an approximate primary, else is a no-op. - pub async fn record_routing_decision( + pub async fn record_hashed_routing_decision( &self, worker: WorkerWithDpRank, - hashes: RoutingDecisionHashes, + local_hashes: Vec, + sequence_hashes: Vec, ) -> Result<(), KvRouterError> { self.recording_target() - .record_routing_hashes(worker, hashes) + .record_routing_hashes( + worker, + RoutingDecisionHashes { + local_hashes, + sequence_hashes, + }, + ) .await } - pub async fn record_hashed_routing_decision( + pub async fn record_routing_decision_hashes( &self, worker: WorkerWithDpRank, - local_hashes: Vec, - sequence_hashes: Vec, + hashes: RoutingDecisionHashes, ) -> Result<(), KvRouterError> { - self.record_routing_decision( - worker, - RoutingDecisionHashes { - local_hashes, - sequence_hashes, - }, - ) - .await + self.recording_target() + .record_routing_hashes(worker, hashes) + .await } pub async fn process_routing_decision_for_request( diff --git a/lib/kv-router/src/services/selection/README.md b/lib/kv-router/src/services/selection/README.md index 195159579f89..0991c675c60e 100644 --- a/lib/kv-router/src/services/selection/README.md +++ b/lib/kv-router/src/services/selection/README.md @@ -17,10 +17,8 @@ and active-sequence accounting. Keep these implementation invariants explicit: - Mooncake overlap fields are raw matched-token observability. Effective prefill tokens use the scheduler's weighted cache credit and are not derived from `longest_matched`. -- Worker membership arrives as complete snapshots through a - `WorkerCatalogSource`; `CatalogReconciler` owns the desired-versus-actual - diff (retry of not-yet-schedulable records, stale deletion). Hosts do not - call `upsert_worker`/`delete_worker` for discovery-driven membership. +- Discovery-driven worker membership goes through `WorkerCatalogSource` and + `CatalogReconciler` (`membership.rs`), not `upsert_worker`/`delete_worker`. - A partition's KV index comes from `HostCache.index: KvIndexSource`: `Owned` (a `KvEventIngress` builds and feeds it: `ZmqDirectIngress` here, the runtime event plane in the frontend) or `Remote` (a standalone indexer diff --git a/lib/kv-router/src/services/selection/affinity.rs b/lib/kv-router/src/services/selection/affinity.rs index f466884dd1fe..270dca637087 100644 --- a/lib/kv-router/src/services/selection/affinity.rs +++ b/lib/kv-router/src/services/selection/affinity.rs @@ -143,7 +143,7 @@ struct Inner { #[cfg(any(test, feature = "testing"))] reaper_started: Arc, #[cfg(any(test, feature = "testing"))] - waiter_observed: Notify, + waiter_observed: Arc, } impl Drop for Inner { @@ -184,7 +184,6 @@ pub enum Hold { } impl Hold { - /// The worker a bound session steers to; `None` while initializing. pub fn target(&self) -> Option { match self { Self::Initialize(_) => None, @@ -192,8 +191,6 @@ impl Hold { } } - /// `AffinityLease::invalidate` for a bound session; an initializing hold - /// just releases. pub fn invalidate(self) { if let Self::Bound { mut lease, .. } = self { lease.invalidate(); @@ -242,7 +239,7 @@ impl SessionAffinity { #[cfg(any(test, feature = "testing"))] reaper_started: Arc::new(Notify::new()), #[cfg(any(test, feature = "testing"))] - waiter_observed: Notify::new(), + waiter_observed: Arc::new(Notify::new()), }); Self::spawn_reaper(&inner); tracing::debug!( @@ -328,8 +325,8 @@ impl SessionAffinity { Entry::Vacant(entry) => { self.reserve_entry()?; tracing::debug!( - session_id, - "Session affinity miss: new session, pinning after worker selection" + session_id = %session_id, + "session affinity miss: new session, pinning after worker selection" ); let revision = self.inner.next_revision.fetch_add(1, Ordering::Relaxed); let notify = Arc::new(Notify::new()); @@ -339,7 +336,7 @@ impl SessionAffinity { }); Ok(AcquireStep::Held(Hold::Initialize( AffinityInitialization { - table: Arc::downgrade(&self.inner), + coordinator: Arc::downgrade(&self.inner), session_id: session_id.to_string(), revision, notify, @@ -364,8 +361,8 @@ impl SessionAffinity { .. } if *active_leases == 0 && *idle_deadline <= now => { tracing::debug!( - session_id, - "Session affinity miss: pin expired (idle past TTL), re-selecting worker" + session_id = %session_id, + "session affinity miss: pin expired (idle past TTL), re-selecting worker" ); let revision = self.inner.next_revision.fetch_add(1, Ordering::Relaxed); let notify = Arc::new(Notify::new()); @@ -375,7 +372,7 @@ impl SessionAffinity { }; Ok(AcquireStep::Held(Hold::Initialize( AffinityInitialization { - table: Arc::downgrade(&self.inner), + coordinator: Arc::downgrade(&self.inner), session_id: session_id.to_string(), revision, notify, @@ -393,7 +390,7 @@ impl SessionAffinity { } => { validate_bound_target(session_id, *target, requested_target)?; tracing::debug!( - session_id, + session_id = %session_id, worker_id = target.worker_id, dp_rank = ?target.dp_rank, active_leases = *active_leases + 1, @@ -403,7 +400,7 @@ impl SessionAffinity { Ok(AcquireStep::Held(Hold::Bound { target: *target, lease: AffinityLease { - table: Arc::downgrade(&self.inner), + coordinator: Arc::downgrade(&self.inner), session_id: session_id.to_string(), revision: *revision, version: *version, @@ -489,7 +486,7 @@ impl SessionAffinity { } validate_bound_target(session_id, *target, requested_target)?; tracing::debug!( - session_id, + session_id = %session_id, worker_id = target.worker_id, dp_rank = ?target.dp_rank, "Session affinity hit: reusing pinned worker" @@ -497,12 +494,10 @@ impl SessionAffinity { Ok(Some(*target)) } - /// Advance the local sequence past one seen from a replica. pub fn observe_replica_sequence(&self, sequence: u64) { self.inner.observe_replica_sequence(sequence); } - /// Apply a binding published by a replica. pub fn apply_replica_update( &self, session_id: String, @@ -536,7 +531,6 @@ impl SessionAffinity { self.inner.entry_count.load(Ordering::Relaxed) } - /// Cancelled when the last handle drops. pub fn cancellation_token(&self) -> CancellationToken { self.inner.cancel.clone() } @@ -676,7 +670,7 @@ impl Inner { /// A session this request is the first to bind. Dropping it uncommitted /// releases the slot and wakes waiters so they re-acquire. pub struct AffinityInitialization { - table: Weak, + coordinator: Weak, session_id: String, revision: u64, notify: Arc, @@ -688,7 +682,7 @@ impl AffinityInitialization { /// Bind the session to the worker the request was dispatched to. pub fn commit(mut self, target: AffinityTarget) -> Result { validate_bound_target(&self.session_id, target, self.requested_target)?; - let Some(inner) = self.table.upgrade() else { + let Some(inner) = self.coordinator.upgrade() else { return Err(AffinityError::Dropped); }; let Some(mut entry) = inner.entries.get_mut(&self.session_id) else { @@ -716,7 +710,7 @@ impl AffinityInitialization { self.active = false; self.notify.notify_waiters(); Ok(AffinityLease { - table: Arc::downgrade(&inner), + coordinator: Arc::downgrade(&inner), session_id: self.session_id.clone(), revision: self.revision, version, @@ -730,7 +724,7 @@ impl Drop for AffinityInitialization { if !self.active { return; } - let Some(inner) = self.table.upgrade() else { + let Some(inner) = self.coordinator.upgrade() else { return; }; let removed = inner.entries.remove_if(&self.session_id, |_, entry| { @@ -749,7 +743,7 @@ impl Drop for AffinityInitialization { /// A live use of a bound session. Dropping it releases the use and refreshes /// the idle deadline. pub struct AffinityLease { - table: Weak, + coordinator: Weak, session_id: String, revision: u64, version: AffinityVersion, @@ -761,17 +755,16 @@ impl AffinityLease { &self.session_id } - /// Publish the binding to replicas. - pub fn publish(&self, target: AffinityTarget) { - if let Some(inner) = self.table.upgrade() { + fn publish(&self, target: AffinityTarget) { + if let Some(inner) = self.coordinator.upgrade() { inner.publish_replica_update(&self.session_id, target, self.version); } } /// Move the binding from `expected` to `target` if nothing changed it since /// this lease was taken. - pub fn rebind(&mut self, expected: AffinityTarget, target: AffinityTarget) -> bool { - let Some(inner) = self.table.upgrade() else { + fn rebind(&mut self, expected: AffinityTarget, target: AffinityTarget) -> bool { + let Some(inner) = self.coordinator.upgrade() else { return false; }; let Some(mut entry) = inner.entries.get_mut(&self.session_id) else { @@ -813,7 +806,7 @@ impl AffinityLease { return; } self.active = false; - let Some(inner) = self.table.upgrade() else { + let Some(inner) = self.coordinator.upgrade() else { return; }; let (target, version) = { @@ -849,7 +842,7 @@ impl AffinityLease { if !self.active { return; } - let Some(inner) = self.table.upgrade() else { + let Some(inner) = self.coordinator.upgrade() else { self.active = false; return; }; diff --git a/lib/kv-router/src/services/selection/core/hint.rs b/lib/kv-router/src/services/selection/core/hint.rs index dd382c6ce721..a01903988943 100644 --- a/lib/kv-router/src/services/selection/core/hint.rs +++ b/lib/kv-router/src/services/selection/core/hint.rs @@ -7,8 +7,7 @@ use super::*; /// Pick the best router-hint source for `target`: a same-role worker (or /// cache owner) holding a longer root-aligned prefix than the target's own -/// `target_cached_prefix_blocks`, with a non-empty control endpoint. Both hosts -/// reach it through `run_selection`. +/// `target_cached_prefix_blocks`, with a non-empty control endpoint. pub(super) fn transfer_hint_for_selection( configs: &HashMap, target: WorkerWithDpRank, diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index ef4f0fe292be..3989139d6cc6 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -83,7 +83,6 @@ use crate::WorkerSelectionPolicyFactory; use crate::WorkerType; use crate::services::common::replica_sync::AffinityBindingEvent; -/// The scheduler type every partition runs. pub type SelectionScheduler = LocalScheduler< ScopedSequencePublisher, SelectionWorkerConfig, @@ -162,10 +161,7 @@ impl SelectionEntry { /// What an embedding host supplies to every partition the core creates, /// grouped by purpose. Each group defaults to the standalone service's -/// behavior, so a host overrides only the groups it owns: the frontend -/// `KvRouter` feeds load and availability from its request client, points -/// overlap refresh at its own index, and carries replica sync on its own -/// transport. +/// behavior, so a host overrides only the groups it owns. #[derive(Clone, Default)] pub struct SelectionHost { pub load: HostLoad, @@ -199,8 +195,7 @@ pub struct HostCache { pub enum KvIndexSource { /// The ingress builds each partition's index and feeds it with worker KV /// events; it also decides what metadata a worker needs to be schedulable - /// and what happens to the index when a worker leaves. Defaults to - /// [`ZmqDirectIngress`]; the frontend supplies its runtime-backed ingress. + /// and what happens to the index when a worker leaves. Owned(Arc), /// A standalone indexer at this base URL serves the primary index; this /// core does not subscribe to worker KV events. @@ -213,15 +208,6 @@ impl Default for KvIndexSource { } } -impl std::fmt::Debug for KvIndexSource { - fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - Self::Owned(_) => formatter.write_str("Owned"), - Self::Remote(url) => formatter.debug_tuple("Remote").field(url).finish(), - } - } -} - /// Host-owned narrowing of the candidate set. #[derive(Clone, Default)] pub struct HostEligibility { @@ -247,36 +233,12 @@ pub struct HostReplication { pub request_leases: Option>, } -impl std::fmt::Debug for SelectionHost { - fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - formatter - .debug_struct("SelectionHost") - .field("prefill_estimator", &self.load.prefill_estimator.is_some()) - .field( - "overloaded_workers", - &self.load.overloaded_workers.is_some(), - ) - .field("available_workers", &self.load.available_workers.is_some()) - .field("shared_cache", &self.cache.shared.is_some()) - .field("index", &self.cache.index) - .field( - "lora_worker_filter", - &self.eligibility.lora_worker_filter.is_some(), - ) - .field("scheduler_load", &self.telemetry.scheduler_load.is_some()) - .field("replication", &self.replication.channels.is_some()) - .finish() - } -} - #[derive(Debug, Clone)] pub struct SelectionServiceConfig { pub port: u16, pub threads: usize, pub indexer_peers: Vec, /// Base URL of a standalone indexer that serves the primary KV index. - /// When set, this service does not listen for worker KV events and - /// `indexer_peers` recovery is skipped. pub remote_indexer_url: Option, pub replica_sync_port: Option, pub replica_sync_peers: Vec, @@ -368,8 +330,6 @@ impl SelectionCore { )) } - /// Scheduler and indexer handle for `key`, once a worker has been upserted - /// into that partition. pub fn partition(&self, key: &RoutingPartitionId) -> Option { self.entry(key).map(SelectionPartition) } diff --git a/lib/kv-router/src/services/selection/core/run.rs b/lib/kv-router/src/services/selection/core/run.rs index df4d2bbf10b6..98dd43c57330 100644 --- a/lib/kv-router/src/services/selection/core/run.rs +++ b/lib/kv-router/src/services/selection/core/run.rs @@ -596,9 +596,6 @@ impl SelectionCore { /// Bind the held session to `dispatched`. A `Hard` rejection whose bound /// worker departed after [`Self::hold_session`] checked it is not a client /// fault: the session is re-initialized on the dispatched worker instead. - /// This runs with the booking held, so it never waits on another request's - /// initialization (which may itself be queued behind this booking); when - /// one is in flight the booking stands unbound, `None`. pub(super) fn commit_session( &self, table: &SessionAffinity, @@ -644,7 +641,7 @@ impl SelectionCore { } if let Err(error) = entry .indexer - .record_routing_decision( + .record_routing_decision_hashes( worker, RoutingDecisionHashes::from_local_hashes(block_hashes), ) diff --git a/lib/kv-router/src/services/selection/core/tests.rs b/lib/kv-router/src/services/selection/core/tests.rs index b4d88f5f8018..557c1a587241 100644 --- a/lib/kv-router/src/services/selection/core/tests.rs +++ b/lib/kv-router/src/services/selection/core/tests.rs @@ -223,6 +223,46 @@ fn assert_shutdown_error(error: SelectionError) { )); } +fn default_key() -> RoutingPartitionId { + RoutingPartitionId::new("model", "default") +} + +/// A core whose queue threshold is zero, so a second booking queues behind +/// the first. +fn saturated_core() -> Arc { + let mut config = test_config(false); + config.router_queue_threshold = Some(0.0); + Arc::new(local_core(config)) +} + +fn lease_operation<'a>( + prompt: PromptView<'a>, + request_id: &str, + track_active_blocks: bool, +) -> SelectionOperation<'a> { + SelectionOperation { + key: default_key(), + prompt, + router_config_override: None, + expected_output_tokens: None, + priority_jump: 0.0, + strict_priority: 0, + policy_class: None, + session_context: None, + session: SessionBinding::None, + affinity_target: None, + pinned_worker: None, + allowed_worker_ids: None, + routing_constraints: RoutingConstraints::default(), + admission: SelectionAdmission::Lease { + request_id: request_id.to_string(), + }, + track_active_blocks, + return_routing_hashes: false, + replay_id: None, + } +} + #[test] fn parent_cancel_cancels_core() { let parent = CancellationToken::new(); @@ -251,9 +291,7 @@ async fn selection_setup_uses_worker_type_label() { ); core.upsert_worker(worker(1)).await.expect("worker upsert"); - let entry = core - .entry(&RoutingPartitionId::new("model", "default")) - .expect("selection entry"); + let entry = core.entry(&default_key()).expect("selection entry"); assert_eq!( entry.scheduler.worker_type(), expected_label, @@ -378,7 +416,7 @@ async fn remote_indexer_serves_selection_without_local_kv_listeners() { use crate::services::indexer::server::spawn_test_indexer_server; // The standalone indexer holds worker 2's cache for the test prompt. - let key = RoutingPartitionId::new("model", "default"); + let key = default_key(); let served = Arc::new(WorkerRegistry::new(1)); let served_indexer = served.get_or_create_indexer(key.clone(), 4); let hashes: Vec = @@ -439,9 +477,7 @@ async fn hint_fixture( configure(&mut request); core.upsert_worker(request).await.expect("worker upsert"); } - let entry = core - .entry(&RoutingPartitionId::new("model", "default")) - .expect("entry"); + let entry = core.entry(&default_key()).expect("entry"); let tokens: Vec = (1..=8).collect(); let hashes: Vec = compute_block_hash_for_seq(&tokens, 4, BlockHashOptions::default()) .into_iter() @@ -770,9 +806,7 @@ async fn session_context_reaches_worker_selection() { async fn full_affinity_table_routes_without_pinning() { let core = local_core(test_config(false)); core.upsert_worker(worker(1)).await.expect("worker upsert"); - let entry = core - .entry(&RoutingPartitionId::new("model", "default")) - .expect("entry"); + let entry = core.entry(&default_key()).expect("entry"); let table = SessionAffinity::with_config(SessionAffinityConfig { max_entries: 1, max_session_id_bytes: 256, @@ -858,7 +892,7 @@ async fn shutdown_cancels_listeners_but_keeps_parent_alive() { #[tokio::test] async fn selection_sees_cache_after_last_worker_replacement(#[case] indexer_threads: usize) { let core = local_core_with(test_config(true), indexer_threads, CancellationToken::new()); - let key = RoutingPartitionId::new("model", "default"); + let key = default_key(); let request = || { let mut request = select_request(); request.prompt.token_ids = None; @@ -937,9 +971,7 @@ async fn reupsert_recreates_a_listener_lost_to_a_cancelled_update() { core.upsert_worker(worker_with_kv_events(1)) .await .expect("worker upsert"); - let entry = core - .entry(&RoutingPartitionId::new("model", "default")) - .expect("entry"); + let entry = core.entry(&default_key()).expect("entry"); let hashes = [11u64, 12]; entry .indexer @@ -983,12 +1015,12 @@ async fn upsert_moves_global_worker_id_between_routing_groups() { let mut group_a = worker_with_kv_events(1); group_a.routing_group = "group-a".to_string(); core.upsert_worker(group_a).await.expect("group A upsert"); - assert_eq!( + let indexed = |group: &str| { core.indexer_registry - .list_filtered(Some("model"), Some("group-a")) - .len(), - 1 - ); + .list_filtered(Some("model"), Some(group)) + .len() + }; + assert_eq!(indexed("group-a"), 1); let mut group_b = worker_with_kv_events(1); group_b.routing_group = "group-b".to_string(); @@ -996,17 +1028,8 @@ async fn upsert_moves_global_worker_id_between_routing_groups() { assert!(core.list_workers(Some("model"), Some("group-a")).is_empty()); assert_eq!(core.list_workers(Some("model"), Some("group-b")).len(), 1); - assert!( - core.indexer_registry - .list_filtered(Some("model"), Some("group-a")) - .is_empty() - ); - assert_eq!( - core.indexer_registry - .list_filtered(Some("model"), Some("group-b")) - .len(), - 1 - ); + assert_eq!(indexed("group-a"), 0); + assert_eq!(indexed("group-b"), 1); let mut select_a = select_request(); select_a.routing_group = "group-a".to_string(); @@ -1019,11 +1042,7 @@ async fn upsert_moves_global_worker_id_between_routing_groups() { assert_eq!(core.select(select_b).await.unwrap().worker_id, 1); core.delete_worker(1).await.expect("delete group B worker"); - assert!( - core.indexer_registry - .list_filtered(Some("model"), Some("group-b")) - .is_empty() - ); + assert_eq!(indexed("group-b"), 0); } #[tokio::test] @@ -1081,9 +1100,7 @@ async fn shutdown_reports_not_ready_and_rejects_new_work() { #[tokio::test] async fn queued_selection_errors_on_shutdown() { - let mut config = test_config(false); - config.router_queue_threshold = Some(0.0); - let core = Arc::new(local_core(config)); + let core = saturated_core(); let record = core.upsert_worker(worker(1)).await.expect("worker upsert"); assert_eq!(record.lifecycle, WorkerLifecycle::Schedulable); @@ -1110,12 +1127,9 @@ async fn queued_selection_errors_on_shutdown() { #[tokio::test] async fn booking_is_freed_when_selected_worker_drained_while_queued() { - let mut config = test_config(false); - config.router_queue_threshold = Some(0.0); - let core = Arc::new(local_core(config)); + let core = saturated_core(); core.upsert_worker(worker(1)).await.expect("worker upsert"); - let key = RoutingPartitionId::new("model", "default"); - let entry = core.entry(&key).expect("entry"); + let entry = core.entry(&default_key()).expect("entry"); core.select_and_reserve(reserve_request("res-a")) .await .expect("initial reservation"); @@ -1150,12 +1164,9 @@ async fn booking_is_freed_when_selected_worker_drained_while_queued() { /// only thing missing. #[tokio::test] async fn lease_admission_keeps_a_selection_whose_worker_drained_while_queued() { - let mut config = test_config(false); - config.router_queue_threshold = Some(0.0); - let core = Arc::new(local_core(config)); + let core = saturated_core(); core.upsert_worker(worker(1)).await.expect("worker upsert"); - let key = RoutingPartitionId::new("model", "default"); - let entry = core.entry(&key).expect("entry"); + let entry = core.entry(&default_key()).expect("entry"); core.select_and_reserve(reserve_request("res-a")) .await .expect("initial reservation"); @@ -1164,27 +1175,7 @@ async fn lease_admission_keeps_a_selection_whose_worker_drained_while_queued() { let queued = tokio::spawn(async move { let req = reserve_request("queued"); let run = queued_core - .run_selection(SelectionOperation { - key: RoutingPartitionId::new("model", "default"), - prompt: req.prompt.view(), - router_config_override: None, - expected_output_tokens: None, - priority_jump: 0.0, - strict_priority: 0, - policy_class: None, - session_context: None, - session: SessionBinding::None, - affinity_target: None, - pinned_worker: None, - allowed_worker_ids: None, - routing_constraints: RoutingConstraints::default(), - admission: SelectionAdmission::Lease { - request_id: "queued".to_string(), - }, - track_active_blocks: false, - return_routing_hashes: false, - replay_id: None, - }) + .run_selection(lease_operation(req.prompt.view(), "queued", false)) .await; match run.result { Ok(SelectionOutcome::Selected(selected)) => ( @@ -1219,32 +1210,12 @@ async fn lease_admission_keeps_a_selection_whose_worker_drained_while_queued() { async fn lease_admission_installs_no_index_row_and_records_nothing() { let core = local_core(test_config(false)); core.upsert_worker(worker(1)).await.expect("worker upsert"); - let key = RoutingPartitionId::new("model", "default"); + let key = default_key(); let entry = core.entry(&key).expect("entry"); let req = reserve_request("leased"); let run = core - .run_selection(SelectionOperation { - key: key.clone(), - prompt: req.prompt.view(), - router_config_override: None, - expected_output_tokens: None, - priority_jump: 0.0, - strict_priority: 0, - policy_class: None, - session_context: None, - session: SessionBinding::None, - affinity_target: None, - pinned_worker: None, - allowed_worker_ids: None, - routing_constraints: RoutingConstraints::default(), - admission: SelectionAdmission::Lease { - request_id: "leased".to_string(), - }, - track_active_blocks: true, - return_routing_hashes: false, - replay_id: None, - }) + .run_selection(lease_operation(req.prompt.view(), "leased", true)) .await; let Ok(SelectionOutcome::Selected(selected)) = run.result else { panic!("lease selection failed"); @@ -1280,7 +1251,7 @@ async fn lease_admission_installs_no_index_row_and_records_nothing() { async fn dropped_selection_future_frees_its_booking() { let core = local_core(test_config(false)); core.upsert_worker(worker(1)).await.expect("worker upsert"); - let key = RoutingPartitionId::new("model", "default"); + let key = default_key(); let entry = core.entry(&key).expect("entry"); // Drive the selection by hand so the actor's response is delivered but @@ -1381,7 +1352,7 @@ async fn dropped_book_selection_during_routing_record_frees_booking_and_claim() ), ); core.upsert_worker(worker(1)).await.expect("worker upsert"); - let key = RoutingPartitionId::new("model", "default"); + let key = default_key(); let entry = core.entry(&key).expect("entry"); // Poll by hand until the booking exists and the record is awaited. @@ -1479,7 +1450,7 @@ async fn explicit_reservation_of_a_live_id_is_a_conflict() { fn index_observer_replaces_stale_and_claimed_rows_of_its_partition() { use crate::scheduling::AttemptId; let index = Arc::new(RwLock::new(HashMap::new())); - let partition = RoutingPartitionId::new("model", "default"); + let partition = default_key(); let observer = ReservationIndexObserver { index: Arc::clone(&index), partition: partition.clone(), @@ -1534,9 +1505,7 @@ fn index_observer_replaces_stale_and_claimed_rows_of_its_partition() { #[tokio::test] async fn free_of_an_in_flight_reservation_is_not_found() { - let mut config = test_config(false); - config.router_queue_threshold = Some(0.0); - let core = Arc::new(local_core(config)); + let core = saturated_core(); core.upsert_worker(worker(1)).await.expect("worker upsert"); core.select_and_reserve(reserve_request("res-a")) .await @@ -1609,9 +1578,7 @@ async fn mirrored_replica_bookings_are_indexed_until_freed() { .await .expect("warm booking"); core.free_reservation("warm").await.expect("free warm"); - let entry = core - .entry(&RoutingPartitionId::new("model", "default")) - .expect("entry"); + let entry = core.entry(&default_key()).expect("entry"); let peer_event = |request_id: &str, data| ActiveSequenceEvent { request_id: request_id.to_string(), worker: WorkerWithDpRank::new(1, 0), @@ -1847,9 +1814,7 @@ async fn reservation_index_sweep_drops_bookings_released_out_of_band() { ); // Release directly through the scheduler, as force-expiry would. - let entry = core - .entry(&RoutingPartitionId::new("model", "default")) - .expect("entry"); + let entry = core.entry(&default_key()).expect("entry"); entry.scheduler.free("stale").await.expect("scheduler free"); assert_eq!( @@ -1914,16 +1879,14 @@ async fn lifecycle_lookup_does_not_nest_reservation_index_inside_entries() { #[tokio::test(flavor = "current_thread")] async fn queued_selection_returns_refreshed_overlap_snapshot() { - let mut config = test_config(false); - config.router_queue_threshold = Some(0.0); - let core = Arc::new(local_core(config)); + let core = saturated_core(); for worker_id in [1, 2] { let mut request = worker(worker_id); request.max_num_batched_tokens = Some(8); core.upsert_worker(request).await.expect("worker upsert"); } - let key = RoutingPartitionId::new("model", "default"); + let key = default_key(); let entry = core.entry(&key).expect("entry"); entry .indexer @@ -1937,15 +1900,9 @@ async fn queued_selection_returns_refreshed_overlap_snapshot() { worker_id: Some(worker_id), dp_rank: Some(0), prompt: PromptRequest { - token_ids: None, - mm_routing_info: None, - block_mm_infos: None, - block_hashes: None, sequence_hashes: Some(vec![1, 2]), isl_tokens: Some(8), - lora_name: None, - cache_namespace: None, - is_eagle: None, + ..PromptRequest::default() }, effective_prefill_tokens: Some(8), ..replay_reservation(&format!("occupy-{worker_id}")) @@ -1958,30 +1915,13 @@ async fn queued_selection_returns_refreshed_overlap_snapshot() { let queued = tokio::spawn(async move { queued_core .select_and_reserve(SelectAndReserveRequest { - model_name: "model".to_string(), - routing_group: "default".to_string(), - selection_id: Some("refresh-selection".to_string()), prompt: PromptRequest { - token_ids: None, - mm_routing_info: None, - block_mm_infos: None, block_hashes: Some(vec![11, 12]), sequence_hashes: Some(vec![101, 102]), isl_tokens: Some(8), - lora_name: None, - cache_namespace: None, - is_eagle: None, + ..PromptRequest::default() }, - router_config_override: None, - expected_output_tokens: None, - priority_jump: None, - strict_priority: None, - session_id: None, - session_context: None, - affinity_target: None, - pinned_worker: None, - allowed_worker_ids: None, - routing_constraints: RoutingConstraints::default(), + ..reserve_request("refresh-selection") }) .await }); @@ -2024,10 +1964,14 @@ fn core_with_session_affinity() -> SelectionCore { core_with_session_affinity_mode(SessionAffinityMode::Hard) } +/// The worker `session_id` is bound to; panics if the partition has no +/// affinity table, so `None` means unbound rather than unconfigured. fn bound_worker(core: &SelectionCore, session_id: &str) -> Option { - core.entry(&RoutingPartitionId::new("model", "default")) - .and_then(|entry| entry.affinity.get().cloned()) - .and_then(|table| table.query_target(session_id, None).expect("query")) + let entry = core.entry(&default_key()).expect("default partition"); + let table = entry.affinity.get().expect("affinity table configured"); + table + .query_target(session_id, None) + .expect("query") .map(|target| target.worker_id) } @@ -2037,9 +1981,9 @@ fn session_reservation(selection_id: &str, session_id: &str) -> SelectAndReserve request } -#[tokio::test] -async fn departed_session_worker_reinitializes_the_session() { - let core = core_with_session_affinity(); +/// Two workers, session `s` bound by booking `r1` (already freed). +async fn bound_session(mode: SessionAffinityMode) -> (SelectionCore, SelectResponse) { + let core = core_with_session_affinity_mode(mode); core.upsert_worker(worker(1)).await.expect("worker upsert"); core.upsert_worker(worker(2)).await.expect("worker upsert"); let first = core @@ -2047,6 +1991,12 @@ async fn departed_session_worker_reinitializes_the_session() { .await .expect("first booking"); core.free_reservation("r1").await.expect("free"); + (core, first) +} + +#[tokio::test] +async fn departed_session_worker_reinitializes_the_session() { + let (core, first) = bound_session(SessionAffinityMode::Hard).await; core.delete_worker(first.worker_id).await.expect("delete"); let second = core @@ -2059,15 +2009,8 @@ async fn departed_session_worker_reinitializes_the_session() { #[tokio::test] async fn session_worker_departing_after_the_hold_reinitializes_the_session() { - let core = core_with_session_affinity(); - core.upsert_worker(worker(1)).await.expect("worker upsert"); - core.upsert_worker(worker(2)).await.expect("worker upsert"); - let first = core - .select_and_reserve(session_reservation("r1", "s")) - .await - .expect("first booking"); - core.free_reservation("r1").await.expect("free"); - let key = RoutingPartitionId::new("model", "default"); + let (core, first) = bound_session(SessionAffinityMode::Hard).await; + let key = default_key(); // One poll takes the hold (the bound worker still passes the check) // and hands the request to the scheduler actor, which has not run yet. @@ -2076,7 +2019,7 @@ async fn session_worker_departing_after_the_hold_reinitializes_the_session() { assert!(second.as_mut().poll(&mut context).is_pending()); core.catalog .set_lifecycle(first.worker_id, WorkerLifecycle::Draining, Vec::new()); - core.publish_scheduler_config(&key).expect("publish"); + core.publish_scheduler_config(&key); let second = second .await @@ -2186,9 +2129,7 @@ async fn two_phase_replay_rejects_a_worker_the_session_left() { .await .expect_err("hard affinity rejects the stale cached worker"); assert!(matches!(err, SelectionError::BadRequest(_)), "{err:?}"); - let entry = core - .entry(&RoutingPartitionId::new("model", "default")) - .expect("entry"); + let entry = core.entry(&default_key()).expect("entry"); wait_until("rejected booking release", || { !entry.scheduler.has_request("pending") }) @@ -2198,14 +2139,7 @@ async fn two_phase_replay_rejects_a_worker_the_session_left() { #[tokio::test] async fn hard_mode_rejects_dispatch_away_from_a_live_binding() { - let core = core_with_session_affinity(); - core.upsert_worker(worker(1)).await.expect("worker upsert"); - core.upsert_worker(worker(2)).await.expect("worker upsert"); - let first = core - .select_and_reserve(session_reservation("r1", "s")) - .await - .expect("first booking"); - core.free_reservation("r1").await.expect("free"); + let (core, first) = bound_session(SessionAffinityMode::Hard).await; let other = if first.worker_id == 1 { 2 } else { 1 }; // Steering cannot reach the bound worker, so selection lands elsewhere. @@ -2216,9 +2150,7 @@ async fn hard_mode_rejects_dispatch_away_from_a_live_binding() { .await .expect_err("hard affinity rejects a dispatch away from the binding"); assert!(matches!(err, SelectionError::BadRequest(_)), "{err:?}"); - let entry = core - .entry(&RoutingPartitionId::new("model", "default")) - .expect("entry"); + let entry = core.entry(&default_key()).expect("entry"); wait_until("rejected booking release", || { !entry.scheduler.has_request("r2") }) @@ -2236,14 +2168,7 @@ async fn hard_mode_rejects_dispatch_away_from_a_live_binding() { #[tokio::test] async fn soft_mode_follows_the_dispatch() { - let core = core_with_session_affinity_mode(SessionAffinityMode::Soft); - core.upsert_worker(worker(1)).await.expect("worker upsert"); - core.upsert_worker(worker(2)).await.expect("worker upsert"); - let first = core - .select_and_reserve(session_reservation("r1", "s")) - .await - .expect("first booking"); - core.free_reservation("r1").await.expect("free"); + let (core, first) = bound_session(SessionAffinityMode::Soft).await; let other = if first.worker_id == 1 { 2 } else { 1 }; let mut request = session_reservation("r2", "s"); @@ -2260,7 +2185,7 @@ async fn replicated_binding_steers_a_new_session_and_frees_with_the_booking() { core.upsert_worker(worker(2)).await.expect("worker upsert"); core.dispatch_affinity_event(AffinityBindingEvent { - partition: RoutingPartitionId::new("model", "default"), + partition: default_key(), session_id: "chat-b".to_string(), worker_id: 2, dp_rank: Some(0), @@ -2291,7 +2216,7 @@ async fn expired_booking_releases_affinity_lease() { core.select_and_reserve(session_reservation("abandoned", "session")) .await .unwrap(); - let key = RoutingPartitionId::new("model", "default"); + let key = default_key(); core.entry(&key) .unwrap() .scheduler @@ -2304,16 +2229,7 @@ async fn expired_booking_releases_affinity_lease() { ); tokio::time::pause(); tokio::time::advance(Duration::from_secs(11)).await; - assert_eq!( - core.entry(&RoutingPartitionId::new("model", "default")) - .unwrap() - .affinity - .get() - .unwrap() - .query_target("session", None) - .unwrap(), - None - ); + assert_eq!(bound_worker(&core, "session"), None); } #[tokio::test] @@ -2331,18 +2247,7 @@ async fn partition_sessions_keep_independent_bindings() { let mut request = session_reservation("other", "shared-session"); request.routing_group = "other".to_string(); assert_eq!(core.select_and_reserve(request).await.unwrap().worker_id, 3); - assert_eq!( - core.entry(&RoutingPartitionId::new("model", "default")) - .unwrap() - .affinity - .get() - .unwrap() - .query_target("shared-session", None) - .unwrap() - .unwrap() - .worker_id, - first.worker_id - ); + assert_eq!(bound_worker(&core, "shared-session"), Some(first.worker_id)); let again = core .select_and_reserve(session_reservation("again", "shared-session")) .await @@ -2357,7 +2262,7 @@ async fn rejoined_worker_feeds_partition_index() { let core = local_core(test_config(true)); let request = worker_with_kv_events(1); core.upsert_worker(request.clone()).await.unwrap(); - let key = RoutingPartitionId::new("model", "default"); + let key = default_key(); let partition = core.partition(&key).unwrap(); let tokens: Vec = (1..=8).collect(); let hashes: Vec = compute_block_hash_for_seq(&tokens, 4, BlockHashOptions::default()) @@ -2443,9 +2348,7 @@ async fn host_lease_manager_owns_expiry() { tokio::task::yield_now().await; tokio::time::advance(active_request_expiry_duration() * 3).await; tokio::task::yield_now().await; - let entry = core - .entry(&RoutingPartitionId::new("model", "default")) - .unwrap(); + let entry = core.entry(&default_key()).unwrap(); assert_eq!(entry.scheduler.has_request("live"), host_owned); if host_owned { core.free_reservation("live").await.unwrap(); @@ -2545,13 +2448,7 @@ async fn catalog_updates_commit_in_order_without_exposing_candidates() { 1, "deletion detaches the worker" ); - tokio::time::timeout(Duration::from_secs(1), async { - while entry.scheduler.has_request("live") { - tokio::task::yield_now().await; - } - }) - .await - .unwrap(); + wait_until("booking release", || !entry.scheduler.has_request("live")).await; } #[tokio::test] @@ -2573,9 +2470,7 @@ async fn affinity_configuration_rejects_invalid_or_conflicting_config() { } let core = core_with_session_affinity(); core.upsert_worker(worker(1)).await.unwrap(); - let partition = core - .partition(&RoutingPartitionId::new("model", "default")) - .unwrap(); + let partition = core.partition(&default_key()).unwrap(); assert!(matches!( partition.session_affinity(SessionAffinityConfig::new(Duration::from_secs(20))), Err(SelectionError::Conflict(_)) diff --git a/lib/kv-router/src/services/selection/core/workers.rs b/lib/kv-router/src/services/selection/core/workers.rs index 335498773f3f..f42a17cb391f 100644 --- a/lib/kv-router/src/services/selection/core/workers.rs +++ b/lib/kv-router/src/services/selection/core/workers.rs @@ -54,13 +54,13 @@ impl SelectionCore { let key = previous.key(); self.catalog .set_lifecycle(worker_id, WorkerLifecycle::Draining, Vec::new()); - self.publish_scheduler_config(&key)?; + self.publish_scheduler_config(&key); self.cleanup_indexer_registration(&previous).await; let record = self .catalog .set_lifecycle(worker_id, WorkerLifecycle::Unschedulable, Vec::new()) .ok_or_else(|| SelectionError::NotFound(format!("worker {worker_id} not found")))?; - self.publish_scheduler_config(&key)?; + self.publish_scheduler_config(&key); Ok(record) } @@ -115,7 +115,7 @@ impl SelectionCore { { self.catalog .set_lifecycle(old.worker_id, WorkerLifecycle::Draining, Vec::new()); - self.publish_scheduler_config(&old.key())?; + self.publish_scheduler_config(&old.key()); self.cleanup_indexer_registration(old).await; None } else { @@ -141,7 +141,7 @@ impl SelectionCore { // Readers see only committed metadata. A valid capacity/topology update preserves // live bookings on ranks present in both the old and new snapshots. self.catalog.replace(record.clone()); - self.publish_scheduler_config(&record.key())?; + self.publish_scheduler_config(&record.key()); Ok(record) } @@ -335,12 +335,9 @@ impl SelectionCore { remove_worker_from_index(&self.indexer_registry, record).await; } - pub(super) fn publish_scheduler_config( - &self, - key: &RoutingPartitionId, - ) -> Result<(), SelectionError> { + pub(super) fn publish_scheduler_config(&self, key: &RoutingPartitionId) { let Some(entry) = self.entry(key) else { - return Ok(()); + return; }; let workers = self.catalog.scheduler_configs_for_key(key); // Lifecycle transitions between non-schedulable states publish the same @@ -353,6 +350,5 @@ impl SelectionCore { true } }); - Ok(()) } } diff --git a/lib/kv-router/src/services/selection/membership.rs b/lib/kv-router/src/services/selection/membership.rs index c5eed8d2a128..10e79b77bf6c 100644 --- a/lib/kv-router/src/services/selection/membership.rs +++ b/lib/kv-router/src/services/selection/membership.rs @@ -113,6 +113,7 @@ impl CatalogReconciler { observer.removed(&record); } } + // A worker that was never registered is not an error (idempotent). Err(SelectionError::NotFound(_)) => {} Err(error) => return Err(error), } diff --git a/lib/kv-router/src/services/selection/tests.rs b/lib/kv-router/src/services/selection/tests.rs index ebdc009b0c79..c6efa387fdeb 100644 --- a/lib/kv-router/src/services/selection/tests.rs +++ b/lib/kv-router/src/services/selection/tests.rs @@ -123,24 +123,30 @@ impl WorkerFilter for RejectWorker { } } +/// Block size 4, partition `model/default`, hashes tracked through `context`. +fn normalize_with( + request: &PromptRequest, + context: &TrackingHashContext, + assume_kv_reuse: bool, +) -> Result { + request.view().normalize_for_selection( + 4, + false, + Some(TrackingHashInput { + context, + scope: TrackingHashScope { + partition: RoutingPartitionRef::new("model", "default"), + block_size: 4, + }, + assume_kv_reuse, + }), + ) +} + fn normalize_prompt(request: &PromptRequest) -> super::input::NormalizedPrompt { let config = test_config(); let context = TrackingHashContext::from_config(&config).unwrap(); - request - .view() - .normalize_for_selection( - 4, - false, - Some(TrackingHashInput { - context: &context, - scope: TrackingHashScope { - partition: RoutingPartitionRef::new("model", "default"), - block_size: 4, - }, - assume_kv_reuse: true, - }), - ) - .expect("normalize prompt") + normalize_with(request, &context, true).expect("normalize prompt") } async fn response_json(response: Response) -> serde_json::Value { @@ -567,21 +573,7 @@ fn keyed_prompt_tracking_leaves_indexer_hashes_public() { })) .unwrap(); - let normalized = request - .view() - .normalize_for_selection( - 4, - false, - Some(TrackingHashInput { - context: &context, - scope: TrackingHashScope { - partition: RoutingPartitionRef::new("model", "default"), - block_size: 4, - }, - assume_kv_reuse: true, - }), - ) - .unwrap(); + let normalized = normalize_with(&request, &context, true).unwrap(); let public_blocks = compute_block_hash_for_seq( &[1, 2, 3, 4, 5, 6, 7, 8], 4, @@ -606,23 +598,7 @@ fn disabled_kv_reuse_keeps_public_indexer_hashes_and_randomizes_tracking() { "token_ids": [1, 2, 3, 4, 5, 6, 7, 8] })) .unwrap(); - let normalize = || { - request - .view() - .normalize_for_selection( - 4, - false, - Some(TrackingHashInput { - context: &context, - scope: TrackingHashScope { - partition: RoutingPartitionRef::new("model", "default"), - block_size: 4, - }, - assume_kv_reuse: false, - }), - ) - .unwrap() - }; + let normalize = || normalize_with(&request, &context, false).unwrap(); let first = normalize(); let second = normalize(); @@ -703,18 +679,7 @@ fn keyed_hash_only_inputs_remain_trusted_for_selection_and_reservation() { block_size: 4, }; - let selection = request - .view() - .normalize_for_selection( - 4, - false, - Some(TrackingHashInput { - context: &context, - scope, - assume_kv_reuse: true, - }), - ) - .unwrap(); + let selection = normalize_with(&request, &context, true).unwrap(); let reservation = request .view() .normalize_for_reservation( diff --git a/lib/kv-router/src/services/selection/types.rs b/lib/kv-router/src/services/selection/types.rs index 22f479e12c98..a2d347df4cd2 100644 --- a/lib/kv-router/src/services/selection/types.rs +++ b/lib/kv-router/src/services/selection/types.rs @@ -485,8 +485,7 @@ impl SelectAndReserveRequest { } } -/// Session metadata handed to worker selection, mirroring the frontend's -/// agent-context extension field for field. +/// Session metadata handed to worker selection. /// /// `session_context` supersedes the flat `session_id`: when both are present /// the structured form wins. @@ -629,7 +628,7 @@ pub struct SelectResponse { } /// Load snapshot of the chosen worker, as the scheduler projected it for this -/// request. Mirrors the frontend's advisory selection load. +/// request. #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] pub struct SelectionWorkerLoad { pub active_prefill_tokens: usize, diff --git a/lib/llm/src/discovery/model_manager.rs b/lib/llm/src/discovery/model_manager.rs index 090757fc6148..5f5c32b3eb4d 100644 --- a/lib/llm/src/discovery/model_manager.rs +++ b/lib/llm/src/discovery/model_manager.rs @@ -1949,7 +1949,7 @@ impl ModelManager { .await } - /// Construct a KV chooser whose partition runs the given selection policy. + /// Construct a KV chooser with a selector resolved by the router host at startup. #[allow(clippy::too_many_arguments)] pub async fn kv_chooser_for_with_policy( &self, @@ -2033,7 +2033,7 @@ impl ModelManager { .await } - /// Construct a managed KV router whose partition runs the given selection policy. + /// Construct a managed KV router with a selector resolved by the routing host at startup. #[allow(clippy::too_many_arguments)] pub async fn managed_kv_router_for_with_policy( &self, @@ -2132,32 +2132,32 @@ impl ModelManager { ); // A policy that does not consume cache input must not create a shared-cache client or // subscribe to its updates. - let shared_cache: Option> = - if required_worker_inputs.contains(WorkerInputs::CACHE) { - match kv_router_config - .as_ref() - .map(|c| c.shared_cache_type) - .unwrap_or_default() - { - dynamo_kv_router::SharedCacheType::None => None, - dynamo_kv_router::SharedCacheType::Hicache => { - let worker_component_name = &endpoint.id().component; - tracing::info!( - worker_component = worker_component_name, - "Using HiCache shared KV cache" - ); - Some(Arc::new( - self.hicache_cache_for(&endpoint, workers_with_configs.clone()), - )) - } + let wants_cache = required_worker_inputs.contains(WorkerInputs::CACHE); + let shared_cache: Option> = if wants_cache { + match kv_router_config + .as_ref() + .map(|c| c.shared_cache_type) + .unwrap_or_default() + { + dynamo_kv_router::SharedCacheType::None => None, + dynamo_kv_router::SharedCacheType::Hicache => { + let worker_component_name = &endpoint.id().component; + tracing::info!( + worker_component = worker_component_name, + "Using HiCache shared KV cache" + ); + Some(Arc::new( + self.hicache_cache_for(&endpoint, workers_with_configs.clone()), + )) } - } else { - None - }; + } + } else { + None + }; let kv_event_source_requirement = KvEventSourceRequirement::derive(worker_role, &effective_kv_router_config); - let cache_required = required_worker_inputs.contains(WorkerInputs::CACHE) + let cache_required = wants_cache || effective_kv_router_config.serve_indexer || matches!( kv_event_source_requirement, diff --git a/lib/llm/src/discovery/watcher.rs b/lib/llm/src/discovery/watcher.rs index f67ee17b4d9c..451bbff61228 100644 --- a/lib/llm/src/discovery/watcher.rs +++ b/lib/llm/src/discovery/watcher.rs @@ -203,7 +203,7 @@ pub struct ModelWatcher { /// Keep raw pipelines out of default-off and backend-mismatched paths. generate_engine_capabilities: Vec<&'static str>, selection_policy: SelectionPolicySource, - /// Custom policy dispatch cannot infer whether an untyped legacy card is decode or aggregated. + /// Custom selector dispatch cannot infer whether an untyped legacy card is decode or aggregated. require_typed_worker_role: bool, } diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index e78718848fa3..6ab8ad82c5d4 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -744,8 +744,6 @@ impl KvRouter { let available_worker_provider: WorkerAvailabilityProvider = Arc::new(move || client_for_availability.available_instance_ids()); - // The manager is the partition's lease observer, so it exists before - // the scheduler; its scheduler cleanup is set once the scheduler does. let request_leases = request_lease::RequestLeaseManager::new(cancellation_token.child_token()); let (selection, replica_ingress) = embedded::EmbeddedSelection::start( @@ -1031,7 +1029,9 @@ impl KvRouter { hashes: RoutingDecisionHashes, worker: WorkerWithDpRank, ) -> Result<(), KvRouterError> { - self.indexer.record_routing_decision(worker, hashes).await + self.indexer + .record_routing_decision_hashes(worker, hashes) + .await } /// Give these tokens, find the worker with the best weighted cache hit. @@ -1332,19 +1332,18 @@ impl KvRouter { let prefill_load_hint = self.prefill_load_hint_for(isl_tokens, cached_tokens, track_prefill_tokens); - let sequence_request = SequenceRequest { - request_id: request_id.clone(), - token_sequence: maybe_seq_hashes, - track_prefill_tokens, - expected_output_tokens, - prefill_load_hint, - worker, - lora_name, - }; let admission = self .selection .scheduler() - .add_request_admitted(sequence_request) + .add_request_admitted(SequenceRequest { + request_id: request_id.clone(), + token_sequence: maybe_seq_hashes, + track_prefill_tokens, + expected_output_tokens, + prefill_load_hint, + worker, + lora_name, + }) .await; let attempt_id = match admission { Ok(attempt_id) => attempt_id, @@ -1552,12 +1551,10 @@ impl KvRouter { let tiered_matches = self.indexer.find_matches_by_tier(block_hashes).await?; let cache_hit_estimates = self.cache_hit_estimates_from_tiered_matches(&tiered_matches); - let effective_cached_tokens: HashMap = - cache_hit_estimates.cached_tokens.into_iter().collect(); Ok(self.selection.scheduler().get_potential_loads( maybe_seq_hashes, isl_tokens, - effective_cached_tokens, + cache_hit_estimates.cached_tokens.into_iter().collect(), track_prefill_tokens, )) } @@ -2053,31 +2050,19 @@ mod tests { } async fn make_router_without_membership(worker_role: Option) -> Result { - let component = make_test_component("role-aware-subscription").await; - let endpoint = component.endpoint("backend"); - let client = endpoint.client().await?; - let (_tx, workers) = watch::channel(HashMap::from([(7, ModelRuntimeConfig::default())])); - let config = KvRouterConfig { - skip_initial_worker_wait: true, - router_event_threads: 1, - ..Default::default() - }; - - KvRouter::new_with_worker_role( - endpoint, - client, - workers, - None, + make_router( + "role-aware-subscription", + HashMap::from([(7, ModelRuntimeConfig::default())]), 16, SelectionPolicySource::Registry, - Some(config), None, worker_role, "decode", - None, - false, - None, - None, + KvRouterConfig { + skip_initial_worker_wait: true, + router_event_threads: 1, + ..Default::default() + }, ) .await } @@ -2102,34 +2087,22 @@ mod tests { #[tokio::test] async fn load_only_selector_skips_cache_inputs() { - let component = make_test_component("load-only-capability").await; - let endpoint = component.endpoint("backend"); - let client = endpoint.client().await.unwrap(); - let (_tx, workers) = watch::channel(HashMap::from([(7, ModelRuntimeConfig::default())])); - let config = KvRouterConfig { - skip_initial_worker_wait: true, - router_event_threads: 1, - ..Default::default() - }; - - let router = KvRouter::new_with_worker_role( - endpoint, - client, - workers, - None, + let router = make_router( + "load-only-capability", + HashMap::from([(7, ModelRuntimeConfig::default())]), 16, picker_policy(|| Box::new(LoadOnlyPicker)), - Some(config), - None, - Some(WorkerType::Prefill), - "prefill", - None, - false, Some(Arc::new(FakeSharedCache { hits: None, should_error: false, })), - None, + Some(WorkerType::Prefill), + "prefill", + KvRouterConfig { + skip_initial_worker_wait: true, + router_event_threads: 1, + ..Default::default() + }, ) .await .unwrap(); @@ -2144,36 +2117,40 @@ mod tests { )); } - /// One block-size-2 router over `workers`, no KV event subscription. + /// One router over `workers`, no KV event subscription; `role` is the + /// worker role and its metric label. + #[allow(clippy::too_many_arguments)] async fn make_router( name: &str, workers: HashMap, + block_size: u32, policy: SelectionPolicySource, shared_cache: Option>, + worker_role: Option, + metric_label: &'static str, config: KvRouterConfig, - ) -> KvRouter { + ) -> Result { let component = make_test_component(name).await; let endpoint = component.endpoint("backend"); - let client = endpoint.client().await.unwrap(); + let client = endpoint.client().await?; let (_tx, rx) = watch::channel(workers); KvRouter::new_with_worker_role( endpoint, client, rx, None, - 2, + block_size, policy, Some(config), None, - None, - "decode", + worker_role, + metric_label, None, false, shared_cache, None, ) .await - .unwrap() } /// Three default-config workers under the registry policy, with @@ -2190,7 +2167,18 @@ mod tests { let workers = (0..3) .map(|worker_id| (worker_id, ModelRuntimeConfig::default())) .collect(); - make_router(name, workers, SelectionPolicySource::Registry, None, config).await + make_router( + name, + workers, + 2, + SelectionPolicySource::Registry, + None, + None, + "decode", + config, + ) + .await + .unwrap() } /// Advisory best-match query with default routing arguments. @@ -2312,9 +2300,8 @@ mod tests { ); /// The selection procedure, run over the corpus with tracked bookings, must - /// match the trace frozen from the frontend's pre-unification procedure. - /// Regenerate the trace with `SELECTION_GOLDEN_UPDATE=1` only for an - /// intended behavior change. + /// match the frozen trace. Regenerate it with `SELECTION_GOLDEN_UPDATE=1` + /// only for an intended behavior change. #[tokio::test] async fn selection_matches_frozen_frontend_trace() { let router = tracked_router("golden").await; @@ -2536,7 +2523,18 @@ mod tests { (0, ModelRuntimeConfig::default()), (1, ModelRuntimeConfig::default()), ]); - make_router("shared-cache-router", workers, policy, shared_cache, config).await + make_router( + "shared-cache-router", + workers, + 2, + policy, + shared_cache, + None, + "decode", + config, + ) + .await + .unwrap() } #[tokio::test] diff --git a/lib/llm/src/kv_router/AGENTS.md b/lib/llm/src/kv_router/AGENTS.md index 279e6958cac1..03a3e5a7b871 100644 --- a/lib/llm/src/kv_router/AGENTS.md +++ b/lib/llm/src/kv_router/AGENTS.md @@ -4,7 +4,7 @@ - `kv_router.rs` estimates how much of a request each worker already has cached, then asks the scheduler to choose a worker. It selects a worker but does not send the request. - `embedded.rs` connects discovered workers and their current load to the scheduling code in `lib/kv-router`. -- `indexer/` and `route_lookup.rs` track which KV-cache blocks each worker holds and look for reusable blocks before selection. +- `indexer/` tracks which KV-cache blocks each worker holds; lookups for reusable blocks run in `lib/kv-router` before selection. - `publisher/` receives KV-cache events and worker metrics from inference engines. - `routing_host.rs` owns selection, dispatch, and response cleanup. `routing_host/kv_selection.rs` chooses KV-routed workers, `routing_host/request_guard.rs` tracks progress and cleanup, and `routing_host/cancellation.rs` stops unfinished work when the client cancels. - `prefill_router/` optionally runs a request on a prefill worker before sending it to decode. diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs index 5d014155ec0d..3ed4f4d1849c 100644 --- a/lib/llm/src/kv_router/embedded.rs +++ b/lib/llm/src/kv_router/embedded.rs @@ -55,7 +55,7 @@ pub(crate) struct EmbeddedSelectionArgs { /// takes its index from it. pub ingress: Arc, /// Scheduler-owned load snapshots for the worker monitor's overload - /// detection, the same feed the runtime scheduler publishes. + /// detection. pub scheduler_load: crate::kv_router::routing_load::SchedulerLoadSender, /// Endpoint whose event plane carries replica sync when /// `router_replica_sync` is set. diff --git a/lib/llm/src/kv_router/indexer/ingress.rs b/lib/llm/src/kv_router/indexer/ingress.rs index 19750738b35a..d42fa4900b0e 100644 --- a/lib/llm/src/kv_router/indexer/ingress.rs +++ b/lib/llm/src/kv_router/indexer/ingress.rs @@ -46,64 +46,74 @@ pub(crate) struct RuntimeIngress { impl RuntimeIngress { pub(crate) async fn start(args: RuntimeIngressArgs<'_>) -> Result> { - let component = args.endpoint.component(); - let indexer = if args.cache_required { + let RuntimeIngressArgs { + endpoint, + kv_router_config, + block_size, + model_name, + worker_role, + metric_worker_type, + cache_required, + kv_event_source_requirement, + kv_source_membership, + cancellation_token, + } = args; + let component = endpoint.component(); + let indexer = if cache_required { super::build( component, - args.kv_router_config, - args.block_size, - args.model_name, - args.cancellation_token.child_token(), + kv_router_config, + block_size, + model_name, + cancellation_token.child_token(), ) .await? } else { Indexer::None }; - let subscription = if args.cache_required - && args - .kv_event_source_requirement - .should_subscribe(args.kv_router_config) + let subscription = if cache_required + && kv_event_source_requirement.should_subscribe(kv_router_config) { - let membership_watch = args.kv_source_membership.ok_or_else(|| { + let membership_watch = kv_source_membership.ok_or_else(|| { anyhow::anyhow!( "KV source membership watch is required when local KV event subscription is enabled" ) })?; Some( start_subscriber( - args.endpoint.clone(), + endpoint.clone(), indexer.clone(), membership_watch, - args.block_size, - args.model_name.unwrap_or("unknown").to_string(), - args.worker_role, - args.kv_event_source_requirement, - args.metric_worker_type, - args.cancellation_token.child_token(), + block_size, + model_name.unwrap_or("unknown").to_string(), + worker_role, + kv_event_source_requirement, + metric_worker_type, + cancellation_token.child_token(), ) .await?, ) } else { tracing::info!( - requirement = %args.kv_event_source_requirement, - cache_required = args.cache_required, - use_kv_events = args.kv_router_config.use_kv_events, - overlap_score_credit = args.kv_router_config.overlap_score_credit, - use_remote_indexer = args.kv_router_config.use_remote_indexer, + requirement = %kv_event_source_requirement, + cache_required = cache_required, + use_kv_events = kv_router_config.use_kv_events, + overlap_score_credit = kv_router_config.overlap_score_credit, + use_remote_indexer = kv_router_config.use_remote_indexer, "Skipping KV event subscription" ); None }; - let served = if args.kv_router_config.serve_indexer { - let model_name = args.model_name.ok_or_else(|| { + let served = if kv_router_config.serve_indexer { + let model_name = model_name.ok_or_else(|| { anyhow::anyhow!("model_name is required when serve_indexer is configured") })?; Some( ensure_served_indexer_service( component.clone(), - ServedIndexerMode::from_use_kv_events(args.kv_router_config.use_kv_events), + ServedIndexerMode::from_use_kv_events(kv_router_config.use_kv_events), model_name.to_string(), indexer.clone(), ) diff --git a/lib/llm/src/kv_router/indexer/remote.rs b/lib/llm/src/kv_router/indexer/remote.rs index 07376493a8a5..e3c239501fcc 100644 --- a/lib/llm/src/kv_router/indexer/remote.rs +++ b/lib/llm/src/kv_router/indexer/remote.rs @@ -85,41 +85,6 @@ impl RemoteIndexer { }) } - pub(super) async fn fetch_matches_by_tier( - &self, - block_hashes: Vec, - device_only: bool, - ) -> Result { - self.validate_topology_if_ready().await.inspect_err(|_| { - self.metrics.increment_query_failures(); - })?; - - let request = IndexerQueryRequest { - model_name: self.model_name.clone(), - block_hashes, - device_only, - }; - let mut stream: ManyOut = self - .query_router - .round_robin(SingleIn::new(request)) - .await - .inspect_err(|_| { - self.metrics.increment_query_failures(); - })?; - - match stream.next().await { - Some(IndexerQueryResponse::TieredScores(wire)) => Ok(wire.into()), - Some(IndexerQueryResponse::Error(msg)) => { - self.metrics.increment_query_failures(); - Err(anyhow::anyhow!("Remote indexer error: {}", msg)) - } - None => { - self.metrics.increment_query_failures(); - Err(anyhow::anyhow!("Remote indexer returned empty response")) - } - } - } - pub(super) async fn record_hashed_routing_decision( &self, worker: WorkerWithDpRank, @@ -207,7 +172,34 @@ impl dynamo_kv_router::services::indexer::backend::RemotePrimary for RemoteIndex block_hashes: Vec, device_only: bool, ) -> Result { - self.fetch_matches_by_tier(block_hashes, device_only).await + self.validate_topology_if_ready().await.inspect_err(|_| { + self.metrics.increment_query_failures(); + })?; + + let request = IndexerQueryRequest { + model_name: self.model_name.clone(), + block_hashes, + device_only, + }; + let mut stream: ManyOut = self + .query_router + .round_robin(SingleIn::new(request)) + .await + .inspect_err(|_| { + self.metrics.increment_query_failures(); + })?; + + match stream.next().await { + Some(IndexerQueryResponse::TieredScores(wire)) => Ok(wire.into()), + Some(IndexerQueryResponse::Error(msg)) => { + self.metrics.increment_query_failures(); + Err(anyhow::anyhow!("Remote indexer error: {}", msg)) + } + None => { + self.metrics.increment_query_failures(); + Err(anyhow::anyhow!("Remote indexer returned empty response")) + } + } } async fn record_routing_decision( @@ -604,7 +596,7 @@ mod tests { Arc::new(KvIndexerMetrics::new_unregistered()), ), lower_tier: LowerTierIndexers::new(1, 4), - approx: Some(SideIndexer::Single(side)), + approx: Some(SideIndexer::KvIndexer(side)), primary_records_routing_decisions: false, }; diff --git a/lib/llm/src/kv_router/prefill_router/query.rs b/lib/llm/src/kv_router/prefill_router/query.rs index 92f6d6659e6d..f225b37a4bc6 100644 --- a/lib/llm/src/kv_router/prefill_router/query.rs +++ b/lib/llm/src/kv_router/prefill_router/query.rs @@ -256,8 +256,7 @@ mod tests { .unwrap() } - /// The advisory builtin-host peek `reserve_prefill_worker`'s non-KV arm - /// routes through. + /// The advisory builtin-host peek used by `reserve_prefill_worker`'s non-KV arm. fn query_worker(router: &PrefillRouter) -> u64 { router .binding diff --git a/lib/llm/src/kv_router/routing_host.rs b/lib/llm/src/kv_router/routing_host.rs index 0310e80017a5..944306d8fefc 100644 --- a/lib/llm/src/kv_router/routing_host.rs +++ b/lib/llm/src/kv_router/routing_host.rs @@ -306,12 +306,6 @@ impl RoutePlan { /// removed only in a 2.0.0 (or later) breaking release. pub type KvPushRouter = RoutingHost; -/// The host steers and retries by the same mode its table commits with; a -/// host without session affinity never reads it. -fn affinity_mode(affinity: Option<&AffinityCoordinator>) -> SessionAffinityMode { - affinity.map(AffinityCoordinator::mode).unwrap_or_default() -} - impl RoutingHost { pub fn new( inner: PushRouter>, @@ -382,7 +376,10 @@ impl RoutingHost { inner, policy: RoutingPolicy::Kv(kv_router), request_metrics, - session_affinity_mode: affinity_mode(affinity.as_ref()), + session_affinity_mode: affinity + .as_ref() + .map(AffinityCoordinator::mode) + .unwrap_or_default(), affinity, hosted_occupancy: None, lora: None, @@ -456,7 +453,10 @@ impl RoutingHost { inner, policy, request_metrics, - session_affinity_mode: affinity_mode(affinity.as_ref()), + session_affinity_mode: affinity + .as_ref() + .map(AffinityCoordinator::mode) + .unwrap_or_default(), affinity, hosted_occupancy, lora: lora diff --git a/lib/llm/src/kv_router/routing_host/request_guard.rs b/lib/llm/src/kv_router/routing_host/request_guard.rs index ad98cfc2e151..71f771b7133f 100644 --- a/lib/llm/src/kv_router/routing_host/request_guard.rs +++ b/lib/llm/src/kv_router/routing_host/request_guard.rs @@ -396,8 +396,6 @@ pub(super) struct KvRequestCleanup { } impl KvRequestCleanup { - /// Takes the booking over from `booking`: from here the request's cleanup - /// owns it. pub(super) fn new( chooser: Arc, context_id: String, diff --git a/lib/llm/src/kv_router/routing_host/tests.rs b/lib/llm/src/kv_router/routing_host/tests.rs index cc01353373e3..d5c811c0ce56 100644 --- a/lib/llm/src/kv_router/routing_host/tests.rs +++ b/lib/llm/src/kv_router/routing_host/tests.rs @@ -226,13 +226,11 @@ async fn builtin_direct_without_worker_is_invalid_argument() { let inner = PushRouter::from_client(client, RouterMode::Direct) .await .unwrap(); - let affinity = AffinityCoordinator::new( - Duration::from_secs(10), + let (host, _) = builtin_host_with_affinity( + inner, + load_context, crate::session_affinity::SessionAffinityMode::Hard, - ) - .unwrap(); - let host = - RoutingHost::new_builtin_with_coordinator(inner, load_context, Some(affinity)).unwrap(); + ); let error = host .generate(affinity_request("direct-unbound", None)) @@ -272,14 +270,11 @@ async fn builtin_direct_uses_bound_soft_affinity_as_exact_target() { ) .await .unwrap(); - let affinity = AffinityCoordinator::new( - Duration::from_secs(10), + let (host, _) = builtin_host_with_affinity( + inner, + load_context, crate::session_affinity::SessionAffinityMode::Soft, - ) - .unwrap(); - let host = - RoutingHost::new_builtin_with_coordinator(inner, load_context, Some(affinity.clone())) - .unwrap(); + ); let session_id = SessionAffinityId::new("direct-soft-bound"); bind_affinity_target(&host, &session_id, AffinityTarget::worker(worker_id)).await; @@ -358,26 +353,14 @@ async fn builtin_hard_affinity_ignores_local_inhibition() { ) .await .unwrap(); - let affinity = AffinityCoordinator::new( - Duration::from_secs(10), + let (host, affinity) = builtin_host_with_affinity( + inner, + load_context, crate::session_affinity::SessionAffinityMode::Hard, - ) - .unwrap(); - let host = - RoutingHost::new_builtin_with_coordinator(inner, load_context, Some(affinity.clone())) - .unwrap(); + ); let session_id = SessionAffinityId::new("local-inhibition"); - let Hold::Initialize(initializer) = affinity.acquire(&session_id, None).await.unwrap() else { - panic!("new affinity session must initialize"); - }; - drop( - initializer - .commit(crate::session_affinity::to_table(AffinityTarget::worker( - worker_id, - ))) - .unwrap(), - ); + bind_affinity_target(&host, &session_id, AffinityTarget::worker(worker_id)).await; client.report_instance_down(worker_id); assert!(client.instance_ids().contains(&worker_id)); @@ -514,14 +497,11 @@ async fn builtin_affinity_uses_common_host_for_every_policy() { ) .await .unwrap(); - let affinity = AffinityCoordinator::new( - Duration::from_secs(10), + let (host, affinity) = builtin_host_with_affinity( + inner, + load_context, crate::session_affinity::SessionAffinityMode::Hard, - ) - .unwrap(); - let host = - RoutingHost::new_builtin_with_coordinator(inner, load_context, Some(affinity.clone())) - .unwrap(); + ); let session_id = format!("session-{index}"); let affinity_id = SessionAffinityId::new(session_id.clone()); let explicit_worker = (mode == RouterMode::Direct).then_some(worker_id); @@ -621,14 +601,11 @@ async fn builtin_direct_fallback_stays_disabled_for_affinity() { ) .await .unwrap(); - let affinity = AffinityCoordinator::new( - Duration::from_secs(10), + let (host, affinity) = builtin_host_with_affinity( + inner, + load_context, crate::session_affinity::SessionAffinityMode::Hard, - ) - .unwrap(); - let host = - RoutingHost::new_builtin_with_coordinator(inner, load_context, Some(affinity.clone())) - .unwrap(); + ); let mut standalone = request(); standalone.routing_mut().backend_instance_id = Some(stale_worker); @@ -640,16 +617,7 @@ async fn builtin_direct_fallback_stays_disabled_for_affinity() { ); let session_id = SessionAffinityId::new("direct-affinity"); - let Hold::Initialize(initializer) = affinity.acquire(&session_id, None).await.unwrap() else { - panic!("new affinity session must initialize"); - }; - drop( - initializer - .commit(crate::session_affinity::to_table(AffinityTarget::worker( - stale_worker, - ))) - .unwrap(), - ); + bind_affinity_target(&host, &session_id, AffinityTarget::worker(stale_worker)).await; assert!( host.generate(affinity_request("direct-affinity", Some(stale_worker))) .await @@ -1899,6 +1867,20 @@ async fn session_affinity_existing_selection_cancellation_preserves_binding_with runtime.shutdown(); } +/// A builtin host over `inner` whose session-affinity coordinator has a 10s +/// TTL and binds in `mode`. +fn builtin_host_with_affinity( + inner: PushRouter>, + load_context: Arc, + mode: crate::session_affinity::SessionAffinityMode, +) -> (RoutingHost, AffinityCoordinator) { + let affinity = AffinityCoordinator::new(Duration::from_secs(10), mode).unwrap(); + let host = + RoutingHost::new_builtin_with_coordinator(inner, load_context, Some(affinity.clone())) + .unwrap(); + (host, affinity) +} + async fn bind_affinity_target( router: &RoutingHost, session_id: &SessionAffinityId, diff --git a/lib/llm/src/kv_router/sequence.rs b/lib/llm/src/kv_router/sequence.rs index 0b5795bfb2b2..055d5ebfea51 100644 --- a/lib/llm/src/kv_router/sequence.rs +++ b/lib/llm/src/kv_router/sequence.rs @@ -7,6 +7,8 @@ //! implementations that wire the runtime-agnostic business logic (in `dynamo_kv_router`) //! to the configured event transport and Prometheus metrics. +mod direct_zmq; + pub use dynamo_kv_router::multi_worker_sequence::{ ActiveSequencesMultiWorker, ReplicaRequestLeaseObserver, SchedulerLoadSnapshot, SequenceError, SequencePublishQueueError, SequencePublisher, SequenceRequest, SequenceSubscriber, @@ -30,8 +32,6 @@ use tokio::sync::mpsc; use tokio::time::Instant; use tokio_util::sync::CancellationToken; -mod direct_zmq; - use crate::kv_router::ACTIVE_SEQUENCES_SUBJECT; #[cfg(test)] use dynamo_runtime::transports::event_plane::MsgpackCodec; @@ -41,7 +41,7 @@ use dynamo_runtime::transports::event_plane::MsgpackCodec; const REPLICA_EVENT_CHANNEL_CAPACITY: usize = 100_000; /// How active-sequence events are framed on the wire for a transport. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] enum ActiveSequenceEventWireFormat { Singleton, Batch, @@ -345,14 +345,11 @@ impl SequenceSubscriber for RuntimeSequenceSubscriber { } } -/// Replica-sync channels for an embedded selection partition, carried over the -/// runtime event plane exactly like the runtime scheduler's replica sync: -/// peer and worker events on `ACTIVE_SEQUENCES_SUBJECT` are always forwarded -/// to the partition (worker-origin completion marks are needed even without -/// router-to-router sync); outbound events are published only when -/// `publishes_outbound` is set. The outbound publisher runs from here; the -/// inbound leg runs when the returned [`ReplicaIngress`] is started, so the -/// caller can install every consumer of lifecycle events first. +/// Replica-sync channels for an embedded selection partition over the runtime +/// event plane. Inbound events on `ACTIVE_SEQUENCES_SUBJECT` are always +/// forwarded; outbound events are published only when `publishes_outbound` is +/// set. The inbound leg runs when the returned [`ReplicaIngress`] is started, +/// so the caller can install every consumer of lifecycle events first. pub(crate) async fn host_replica_channels( endpoint: &Endpoint, router_id: u64, @@ -363,8 +360,9 @@ pub(crate) async fn host_replica_channels( ReplicaIngress, )> { let transport_kind = endpoint.drt().default_event_transport_kind(); - let outbound = if publishes_outbound { - let (outbound, outbound_rx) = mpsc::channel(REPLICA_EVENT_CHANNEL_CAPACITY); + let event_sender = if publishes_outbound { + let (event_sender, event_rx) = mpsc::channel(REPLICA_EVENT_CHANNEL_CAPACITY); + let publisher_cancellation_token = cancellation_token.clone(); let event_publisher = EventPublisher::for_endpoint_with_transport( endpoint, ACTIVE_SEQUENCES_SUBJECT, @@ -375,19 +373,19 @@ pub(crate) async fn host_replica_channels( ActiveSequenceEventWireFormat::Singleton => { tokio::spawn(run_replica_singleton_publisher( event_publisher, - outbound_rx, - cancellation_token.clone(), + event_rx, + publisher_cancellation_token, )); } ActiveSequenceEventWireFormat::Batch => { tokio::spawn(run_replica_batch_publisher( event_publisher, - outbound_rx, - cancellation_token.clone(), + event_rx, + publisher_cancellation_token, )); } } - Some(outbound) + Some(event_sender) } else { None }; @@ -400,7 +398,7 @@ pub(crate) async fn host_replica_channels( }; Ok(( dynamo_kv_router::services::selection::HostReplicaChannels { - outbound, + outbound: event_sender, inbound_tx, inbound_rx, process_id: router_id, diff --git a/lib/llm/src/session_affinity/coordinator.rs b/lib/llm/src/session_affinity/coordinator.rs index 8b016caf6146..bd485cfb870a 100644 --- a/lib/llm/src/session_affinity/coordinator.rs +++ b/lib/llm/src/session_affinity/coordinator.rs @@ -44,7 +44,6 @@ pub type AffinityTarget = dynamo_runtime::pipeline::RouteTarget; type TableTarget = dynamo_kv_router::services::selection::affinity::AffinityTarget; -/// The table-side form of a pipeline routing target. pub(crate) fn to_table(target: AffinityTarget) -> TableTarget { TableTarget::new(target.worker_id, target.dp_rank) } @@ -162,7 +161,6 @@ impl AffinityCoordinator { .map_err(affinity_error) } - /// The mode the table binds sessions with. pub(crate) fn mode(&self) -> SessionAffinityMode { self.inner.table.mode() } From 37cfba9b5a08f4adb9f427b5e9113c353472804c Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Sat, 12 Sep 2026 14:52:57 -0700 Subject: [PATCH 071/102] test(router): restore two pins the cutover dropped Two tests on main had no successor after the frontend moved onto the selection core: - `skip_initial_worker_wait_still_monitors_worker_config_updates`: a router built with `skip_initial_worker_wait` still applies later worker-config watch updates. Rebuilt against `KvRouter` through the test router builder, which now returns the watch sender. - `active_sequence_replica_sync_isolated_by_endpoint`: replica ingresses on different endpoints do not see each other's active-sequence events. Pinned at the `ReplicaIngress` layer, where the endpoint scoping lives. Both fail when the guarded behavior is removed (watch ignored after start; second ingress bound to the first endpoint). Co-Authored-By: Claude Fable 5.1 Signed-off-by: Ishan Dhanani --- lib/llm/src/kv_router.rs | 82 ++++++++++++++++++++++++++++++- lib/llm/src/kv_router/sequence.rs | 72 +++++++++++++++++++++++++++ 2 files changed, 152 insertions(+), 2 deletions(-) diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index 6ab8ad82c5d4..ee351c0c1bee 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -2130,11 +2130,40 @@ mod tests { metric_label: &'static str, config: KvRouterConfig, ) -> Result { + make_router_with_watch( + name, + workers, + block_size, + policy, + shared_cache, + worker_role, + metric_label, + config, + ) + .await + .map(|(router, _tx)| router) + } + + /// [`make_router`] that also hands back the worker-config watch sender. + #[allow(clippy::too_many_arguments)] + async fn make_router_with_watch( + name: &str, + workers: HashMap, + block_size: u32, + policy: SelectionPolicySource, + shared_cache: Option>, + worker_role: Option, + metric_label: &'static str, + config: KvRouterConfig, + ) -> Result<( + KvRouter, + watch::Sender>, + )> { let component = make_test_component(name).await; let endpoint = component.endpoint("backend"); let client = endpoint.client().await?; - let (_tx, rx) = watch::channel(workers); - KvRouter::new_with_worker_role( + let (tx, rx) = watch::channel(workers); + let router = KvRouter::new_with_worker_role( endpoint, client, rx, @@ -2150,7 +2179,56 @@ mod tests { shared_cache, None, ) + .await?; + Ok((router, tx)) + } + + /// Workers that appear on the config watch after construction become + /// routable even when the router did not wait for an initial worker. + #[tokio::test] + async fn skip_initial_worker_wait_still_monitors_worker_config_updates() { + let (router, tx) = make_router_with_watch( + "skip-initial-worker-watch", + HashMap::from([(0, ModelRuntimeConfig::default())]), + 2, + SelectionPolicySource::Registry, + None, + Some(WorkerType::Decode), + "decode", + KvRouterConfig { + skip_initial_worker_wait: true, + use_kv_events: false, + router_track_active_blocks: false, + ..Default::default() + }, + ) .await + .unwrap(); + // Catalog upserts reach the scheduler's slots asynchronously. + async fn wait_for_worker(router: &KvRouter, worker_id: WorkerId) -> Vec { + tokio::time::timeout(std::time::Duration::from_secs(5), async { + loop { + let loads = router + .get_potential_loads(&[1, 2, 3, 4], None, None, None, None) + .await + .unwrap(); + if loads.iter().any(|load| load.worker_id == worker_id) { + return loads; + } + tokio::task::yield_now().await; + } + }) + .await + .unwrap_or_else(|_| panic!("worker {worker_id} never became routable")) + } + assert_eq!(wait_for_worker(&router, 0).await.len(), 1); + + tx.send(HashMap::from([ + (0, ModelRuntimeConfig::default()), + (1, ModelRuntimeConfig::default()), + ])) + .unwrap(); + assert_eq!(wait_for_worker(&router, 1).await.len(), 2); } /// Three default-config workers under the registry policy, with diff --git a/lib/llm/src/kv_router/sequence.rs b/lib/llm/src/kv_router/sequence.rs index 055d5ebfea51..024be03aa84a 100644 --- a/lib/llm/src/kv_router/sequence.rs +++ b/lib/llm/src/kv_router/sequence.rs @@ -809,4 +809,76 @@ mod tests { distributed.shutdown(); Ok(()) } + + /// Replica sync is scoped to the endpoint: an ingress on endpoint B never + /// sees events published on endpoint A, even once both planes are live. + #[tokio::test] + #[serial_test::serial] + async fn active_sequence_replica_sync_isolated_by_endpoint() -> Result<()> { + let runtime = dynamo_runtime::Runtime::from_current()?; + let distributed = dynamo_runtime::DistributedRuntime::new( + runtime, + dynamo_runtime::distributed::DistributedConfig::process_local(), + ) + .await?; + let component = distributed + .namespace(format!( + "active-sequence-endpoint-isolation-{}", + uuid::Uuid::new_v4() + ))? + .component("workers")?; + let endpoint_a = component.endpoint("generate-a"); + let endpoint_b = component.endpoint("generate-b"); + let cancel = CancellationToken::new(); + let (mut channels_a, ingress_a) = + host_replica_channels(&endpoint_a, 1, false, cancel.child_token()).await?; + let (mut channels_b, ingress_b) = + host_replica_channels(&endpoint_b, 2, false, cancel.child_token()).await?; + ingress_a.start().await; + ingress_b.start().await; + + // Publish on both planes until each ingress has received something, so + // B's silence about A cannot be blamed on B's subscription not being up. + let publisher_a = ActiveSequenceEventPublisher::for_endpoint(&endpoint_a, 16).await?; + let publisher_b = ActiveSequenceEventPublisher::for_endpoint(&endpoint_b, 16).await?; + let mut received_a = Vec::new(); + let mut received_b = Vec::new(); + tokio::time::timeout(std::time::Duration::from_secs(5), async { + while received_a.is_empty() || received_b.is_empty() { + publisher_a.mark_prefill_completed("endpoint-a-mark".to_string(), 42, 0)?; + publisher_b.mark_prefill_completed("endpoint-b-mark".to_string(), 42, 0)?; + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + while let Ok(event) = channels_a.inbound_rx.try_recv() { + received_a.push(event.request_id); + } + while let Ok(event) = channels_b.inbound_rx.try_recv() { + received_b.push(event.request_id); + } + } + Ok::<_, anyhow::Error>(()) + }) + .await??; + // Give late A events 250ms of silence to leak into B before judging. + while let Ok(Some(event)) = tokio::time::timeout( + std::time::Duration::from_millis(250), + channels_b.inbound_rx.recv(), + ) + .await + { + received_b.push(event.request_id); + } + assert!( + received_a.iter().all(|id| id == "endpoint-a-mark"), + "endpoint A received endpoint B sequence state: {received_a:?}" + ); + assert!( + received_b.iter().all(|id| id == "endpoint-b-mark"), + "endpoint B received endpoint A sequence state: {received_b:?}" + ); + + drop((publisher_a, publisher_b)); + cancel.cancel(); + distributed.shutdown(); + Ok(()) + } } From 7d637baed3e36e2729d3778b43b49de7c203d8d8 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Sat, 12 Sep 2026 17:04:57 -0700 Subject: [PATCH 072/102] ci: refresh the kvbm bindings lockfile and fix a moved docs link The kvbm bindings workspace has its own Cargo.lock; `cargo metadata --locked` refused it after the merge of main brought in dependencies it did not list. The ModelExpress vLLM refit example moved to `dynamo_vllm_reshard_refit`, which the link checker flagged. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Ishan Dhanani --- .../reinforcement-learning/weight-updates.md | 2 +- lib/bindings/kvbm/Cargo.lock | 25 +++++++++++++++++++ 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/docs/fern/pages/use-cases/reinforcement-learning/weight-updates.md b/docs/fern/pages/use-cases/reinforcement-learning/weight-updates.md index e6e43f253dab..994ea820ec81 100644 --- a/docs/fern/pages/use-cases/reinforcement-learning/weight-updates.md +++ b/docs/fern/pages/use-cases/reinforcement-learning/weight-updates.md @@ -17,7 +17,7 @@ ModelExpress can stage a version from three sources: - **Artifact to inference:** an inference worker prepares a canonical checkpoint from S3, including the current exact-base XOR delta format, before engine installation. - **Inference to inference:** after an inference worker applies a version, a rank-compatible worker can pull that version from it instead of returning to the trainer. -All three refit paths are Experimental. The [ModelExpress RL refit package](https://github.com/ai-dynamo/modelexpress/tree/main/modelexpress_client/python/modelexpress_rl) contains the current source strategies. The [Dynamo vLLM refit example](https://github.com/ai-dynamo/modelexpress/tree/main/examples/rl/dynamo_vllm_refit) validates the full-weight trainer-to-inference path only; it does not qualify the S3 delta path or every backend and topology. ModelExpress startup loading remains a separate boot and scale-out workflow. +All three refit paths are Experimental. The [ModelExpress RL refit package](https://github.com/ai-dynamo/modelexpress/tree/main/modelexpress_client/python/modelexpress_rl) contains the current source strategies. The [Dynamo vLLM refit example](https://github.com/ai-dynamo/modelexpress/tree/main/examples/rl/dynamo_vllm_reshard_refit) validates the full-weight trainer-to-inference path only; it does not qualify the S3 delta path or every backend and topology. ModelExpress startup loading remains a separate boot and scale-out workflow. > [!NOTE] > ModelExpress startup loading and live RL refit have different interfaces. vLLM 0.23 and later use the native `--load-format modelexpress`; `mx` remains a backward-compatible alias. Earlier plugin-based images generally use `mx`, while legacy split-loader images can expose `mx-source` and `mx-target`. Match the loader to the installed vLLM, ModelExpress package, and runtime image instead of inferring it from the refit API. See [Dynamo deprecations](../../reference/general/releases/deprecations.mdx) for the Dynamo-owned loader migration. diff --git a/lib/bindings/kvbm/Cargo.lock b/lib/bindings/kvbm/Cargo.lock index 0a511088ba9c..c2a12b1db2a5 100644 --- a/lib/bindings/kvbm/Cargo.lock +++ b/lib/bindings/kvbm/Cargo.lock @@ -2084,7 +2084,9 @@ dependencies = [ "anyhow", "arc-swap", "async-trait", + "axum", "blake3", + "chrono", "crossbeam-queue", "dashmap", "derive_builder", @@ -2095,9 +2097,11 @@ dependencies = [ "fastrand", "flume", "indexmap 2.14.0", + "once_cell", "ordered-float 4.6.0", "parking_lot", "prometheus", + "reqwest 0.12.28", "rmp-serde", "rustc-hash 2.1.2", "seqlock", @@ -2109,9 +2113,11 @@ dependencies = [ "tokio", "tokio-util", "tracing", + "tracing-appender", "uuid", "xxhash-rust", "zeroize", + "zmq", ] [[package]] @@ -7594,6 +7600,12 @@ version = "2.6.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" +[[package]] +name = "symlink" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7973cce6668464ea31f176d85b13c7ab3bba2cb3b77a2ed26abd7801688010a" + [[package]] name = "syn" version = "2.0.117" @@ -8303,6 +8315,19 @@ dependencies = [ "tracing-core", ] +[[package]] +name = "tracing-appender" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "050686193eb999b4bb3bc2acfa891a13da00f79734704c4b8b4ef1a10b368a3c" +dependencies = [ + "crossbeam-channel", + "symlink", + "thiserror 2.0.18", + "time", + "tracing-subscriber", +] + [[package]] name = "tracing-attributes" version = "0.1.31" From 29a05b4833b57ed738b92a8357135170607d22d0 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Sat, 12 Sep 2026 17:27:26 -0700 Subject: [PATCH 073/102] chore(router): finish the deferred mechanical cleanups - Session-affinity tests build targets through one `table_target` fixture. - Routing-host tests share `potential_loads`, `active_requests_for` and `plan_decode_route` instead of repeating the load lookup and the preview-to-plan stage. - Membership tests assert through the public `lifecycle` accessor and a counting catalog observer instead of private fields. - `validate_selector_worker_role` is `validate_policy_worker_role` and takes the `SelectionPolicySource`; the redundant `require_typed_worker_role` flag threaded through the HTTP entrypoint and the watcher is gone. - `indexer/backend.rs` signatures use `Result` through the imported alias. No behavior change; every test keeps its assertions. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Ishan Dhanani --- lib/kv-router/src/services/indexer/backend.rs | 38 +++--- .../src/services/selection/membership.rs | 60 ++++++++- lib/llm/src/discovery/watcher.rs | 25 ++-- lib/llm/src/entrypoint/input/http.rs | 16 +-- lib/llm/src/kv_router/routing_host/tests.rs | 120 +++++++----------- lib/llm/src/session_affinity/tests.rs | 35 +++-- 6 files changed, 155 insertions(+), 139 deletions(-) diff --git a/lib/kv-router/src/services/indexer/backend.rs b/lib/kv-router/src/services/indexer/backend.rs index 144910aa956b..2c50365b5049 100644 --- a/lib/kv-router/src/services/indexer/backend.rs +++ b/lib/kv-router/src/services/indexer/backend.rs @@ -117,7 +117,7 @@ impl RemoteIndexerClient { pub async fn find_tiered_matches( &self, sequence: &[LocalBlockHash], - ) -> std::result::Result { + ) -> Result { if sequence.is_empty() { return Ok(TieredMatchDetails::default()); } @@ -288,7 +288,7 @@ impl SideIndexer { pub(super) async fn find_matches_input( &self, sequence: HashInput<'_>, - ) -> std::result::Result { + ) -> Result { match self { Self::KvIndexer(indexer) => { indexer @@ -305,7 +305,7 @@ impl SideIndexer { &self, worker: WorkerWithDpRank, hashes: RoutingDecisionHashes, - ) -> std::result::Result<(), KvRouterError> { + ) -> Result<(), KvRouterError> { match self { Self::KvIndexer(indexer) => { indexer @@ -332,7 +332,7 @@ impl SideIndexer { &self, worker_id: WorkerId, dp_rank: DpRank, - ) -> std::result::Result<(), KvRouterError> { + ) -> Result<(), KvRouterError> { match self { Self::KvIndexer(indexer) => { indexer @@ -512,10 +512,7 @@ impl Indexer { /// indexer otherwise. `Cleared` events fan out to their applicable physical /// indexes according to the event's reset scope. Failures are logged and /// the first one is returned after every tier was attempted. - pub async fn apply_event_routed( - &self, - event: RouterEvent, - ) -> std::result::Result<(), KvRouterError> { + pub async fn apply_event_routed(&self, event: RouterEvent) -> Result<(), KvRouterError> { let targets_primary = match event.targets_primary() { Ok(targets_primary) => targets_primary, Err(_) => { @@ -604,10 +601,7 @@ impl Indexer { /// device event waits for the primary's queue to accept it, and a lower-tier /// event is enqueued. Unlike [`Self::apply_event_routed`], the first failure /// stops the fan-out. - pub async fn try_apply_event( - &self, - event: RouterEvent, - ) -> std::result::Result<(), KvRouterError> { + pub async fn try_apply_event(&self, event: RouterEvent) -> Result<(), KvRouterError> { let targets_primary = match event.targets_primary() { Ok(targets_primary) => targets_primary, Err(_) => { @@ -683,7 +677,7 @@ impl Indexer { &self, worker_id: WorkerId, dp_rank: DpRank, - ) -> std::result::Result<(), KvRouterError> { + ) -> Result<(), KvRouterError> { match self { Self::Single { primary, @@ -832,7 +826,7 @@ impl Indexer { pub async fn find_tiered_matches( &self, sequence: Vec, - ) -> std::result::Result { + ) -> Result { self.find_tiered_matches_with_options(sequence, LowerTierQueryOptions::default()) .await } @@ -844,7 +838,7 @@ impl Indexer { &self, sequence: Vec, options: LowerTierQueryOptions, - ) -> std::result::Result { + ) -> Result { let options = LowerTierQueryOptions { retain_kv_transfer_chain: options.retain_kv_transfer_chain && self.supports_kv_transfer_chain_retention(), @@ -858,7 +852,7 @@ impl Indexer { &self, sequence: &[LocalBlockHash], options: LowerTierQueryOptions, - ) -> std::result::Result { + ) -> Result { let options = LowerTierQueryOptions { retain_kv_transfer_chain: options.retain_kv_transfer_chain && self.supports_kv_transfer_chain_retention(), @@ -871,7 +865,7 @@ impl Indexer { /// peer's [`Self::apply_event_routed`]: the primary device-tier dump first, /// then every allocated lower-tier indexer's dump with `storage_tier` /// retagged to the tier it lives in. - pub async fn dump_events(&self) -> std::result::Result, KvRouterError> { + pub async fn dump_events(&self) -> Result, KvRouterError> { let (primary_events, lower_tier_entries) = match self { Indexer::Single { primary, @@ -917,7 +911,7 @@ impl Indexer { worker: WorkerWithDpRank, incarnation: ApproximateLruIncarnation, capacity: Option, - ) -> std::result::Result<(), KvRouterError> { + ) -> Result<(), KvRouterError> { match self { Self::Single { primary, .. } => { primary.set_approximate_lru_capacity_now(worker, incarnation, capacity) @@ -929,9 +923,7 @@ impl Indexer { } } - pub async fn approximate_lru_stats( - &self, - ) -> std::result::Result { + pub async fn approximate_lru_stats(&self) -> Result { match self { Self::Single { primary, .. } => primary.approximate_lru_stats().await, Self::Concurrent { primary, .. } => primary.approximate_lru_stats().await, @@ -945,7 +937,7 @@ impl TieredMatchProvider for Indexer { async fn find_tiered_matches( &self, sequence: &[LocalBlockHash], - ) -> std::result::Result { + ) -> Result { self.find_tiered_matches_ref_with_options(sequence, LowerTierQueryOptions::default()) .await } @@ -954,7 +946,7 @@ impl TieredMatchProvider for Indexer { &self, sequence: &[LocalBlockHash], options: LowerTierQueryOptions, - ) -> std::result::Result { + ) -> Result { self.find_tiered_matches_ref_with_options(sequence, options) .await } diff --git a/lib/kv-router/src/services/selection/membership.rs b/lib/kv-router/src/services/selection/membership.rs index 10e79b77bf6c..09af68b2d589 100644 --- a/lib/kv-router/src/services/selection/membership.rs +++ b/lib/kv-router/src/services/selection/membership.rs @@ -145,6 +145,8 @@ impl CatalogReconciler { #[cfg(test)] mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + use super::*; use crate::config::KvRouterConfig; use crate::services::selection::SelectionCacheConfig; @@ -190,37 +192,81 @@ mod tests { .map(|record| record.lifecycle) } + /// Counts observer callbacks so tests can assert convergence (a skipped + /// upsert) and stale deletion through the public surface. + #[derive(Default)] + struct Counter { + upserted: AtomicUsize, + removed: AtomicUsize, + } + + impl Counter { + fn upserts(&self) -> usize { + self.upserted.load(Ordering::SeqCst) + } + + fn removals(&self) -> usize { + self.removed.load(Ordering::SeqCst) + } + } + + impl CatalogObserver for Counter { + fn upserted(&self, _record: &WorkerCatalogRecord) { + self.upserted.fetch_add(1, Ordering::SeqCst); + } + + fn removed(&self, _record: &WorkerCatalogRecord) { + self.removed.fetch_add(1, Ordering::SeqCst); + } + } + + fn reconciler(core: &Arc) -> (CatalogReconciler, Arc) { + let counter = Arc::new(Counter::default()); + let reconciler = CatalogReconciler::new(Arc::clone(core)) + .with_observer(Arc::clone(&counter) as Arc); + (reconciler, counter) + } + #[tokio::test] async fn incomplete_worker_is_retried_and_deleted_when_it_leaves() { let core = core(); - let mut reconciler = CatalogReconciler::new(Arc::clone(&core)); + let (mut reconciler, counter) = reconciler(&core); reconciler.apply(vec![incomplete(1)]).await.expect("apply"); - assert!(reconciler.converged.is_empty()); + assert_eq!(counter.upserts(), 1); assert_eq!(lifecycle(&core, 1), Some(WorkerLifecycle::Incomplete)); // The same snapshot re-upserts rather than skipping the unconverged worker. reconciler.apply(vec![incomplete(1)]).await.expect("apply"); - assert!(reconciler.converged.is_empty()); - assert!(reconciler.tracked.contains(&1)); + assert_eq!(counter.upserts(), 2); + // The never-schedulable worker was still tracked, so leaving deletes it. reconciler.apply(Vec::new()).await.expect("apply"); assert_eq!(lifecycle(&core, 1), Some(WorkerLifecycle::Unschedulable)); - assert!(reconciler.tracked.is_empty()); + assert_eq!(counter.removals(), 1); + + // Nothing is tracked any more: an empty snapshot deletes nothing. + reconciler.apply(Vec::new()).await.expect("apply"); + assert_eq!(counter.removals(), 1); } #[tokio::test] async fn schedulable_worker_converges_and_changed_record_reupserts() { let core = core(); - let mut reconciler = CatalogReconciler::new(Arc::clone(&core)); + let (mut reconciler, counter) = reconciler(&core); reconciler.apply(vec![schedulable(1)]).await.expect("apply"); assert_eq!(lifecycle(&core, 1), Some(WorkerLifecycle::Schedulable)); - assert!(reconciler.converged.contains_key(&1)); + assert_eq!(counter.upserts(), 1); + + // An identical desired record is converged and skips its upsert. + reconciler.apply(vec![schedulable(1)]).await.expect("apply"); + assert_eq!(counter.upserts(), 1); let mut moved = schedulable(1); moved.endpoint = Some("http://10.0.0.9:8000".to_string()); reconciler.apply(vec![moved]).await.expect("apply"); + assert_eq!(counter.upserts(), 2); let record = core .list_workers(None, None) .into_iter() diff --git a/lib/llm/src/discovery/watcher.rs b/lib/llm/src/discovery/watcher.rs index 451bbff61228..344019b55791 100644 --- a/lib/llm/src/discovery/watcher.rs +++ b/lib/llm/src/discovery/watcher.rs @@ -203,8 +203,6 @@ pub struct ModelWatcher { /// Keep raw pipelines out of default-off and backend-mismatched paths. generate_engine_capabilities: Vec<&'static str>, selection_policy: SelectionPolicySource, - /// Custom selector dispatch cannot infer whether an untyped legacy card is decode or aggregated. - require_typed_worker_role: bool, } pub(crate) struct PreparedWorkerSet { @@ -281,7 +279,6 @@ impl ModelWatcher { chat_engine_factory, prefill_load_estimator, metrics, - false, SelectionPolicySource::Registry, ) } @@ -296,7 +293,6 @@ impl ModelWatcher { chat_engine_factory: Option, prefill_load_estimator: Option>, metrics: Arc, - require_typed_worker_role: bool, selection_policy: SelectionPolicySource, ) -> Self { Self { @@ -316,7 +312,6 @@ impl ModelWatcher { tokenizer_fallback_enabled: None, generate_engine_capabilities: Vec::new(), selection_policy, - require_typed_worker_role, } } @@ -413,7 +408,7 @@ impl ModelWatcher { card.download_config(self.local_model_path.as_deref()) .await?; - validate_selector_worker_role(card, self.require_typed_worker_role)?; + validate_policy_worker_role(card, &self.selection_policy)?; // Use per-worker-set router config if the worker provided one in its MDC, // otherwise fall back to the frontend-level global config. Policy selections @@ -1254,11 +1249,13 @@ fn effective_router_config<'a>( Cow::Owned(effective) } -fn validate_selector_worker_role( +/// A custom policy factory cannot infer whether an untyped legacy card is +/// decode or aggregated, so it requires an explicit `worker_type`. +fn validate_policy_worker_role( card: &ModelDeploymentCard, - require_typed_worker_role: bool, + policy: &SelectionPolicySource, ) -> anyhow::Result<()> { - if require_typed_worker_role && card.worker_type.is_none() { + if matches!(policy, SelectionPolicySource::Factory(_)) && card.worker_type.is_none() { anyhow::bail!( "custom worker-selection policies require model cards with an explicit worker_type" ); @@ -2384,12 +2381,16 @@ mod tests { #[test] fn custom_selector_requires_explicit_worker_type() { + let registry = SelectionPolicySource::Registry; + let factory = SelectionPolicySource::Factory(Arc::new(|_, _, _| { + unreachable!("role validation never constructs the policy") + })); let mut card = ModelDeploymentCard::with_name_only("model"); - assert!(validate_selector_worker_role(&card, false).is_ok()); - assert!(validate_selector_worker_role(&card, true).is_err()); + assert!(validate_policy_worker_role(&card, ®istry).is_ok()); + assert!(validate_policy_worker_role(&card, &factory).is_err()); card.worker_type = Some(WorkerType::Decode); - assert!(validate_selector_worker_role(&card, true).is_ok()); + assert!(validate_policy_worker_role(&card, &factory).is_ok()); } #[test] diff --git a/lib/llm/src/entrypoint/input/http.rs b/lib/llm/src/entrypoint/input/http.rs index 2c446b358e4d..850b19fb8f4b 100644 --- a/lib/llm/src/entrypoint/input/http.rs +++ b/lib/llm/src/entrypoint/input/http.rs @@ -49,7 +49,7 @@ impl HttpFrontend { self } - /// Replace the default worker selector with a statically linked native policy. + /// Replace the registry-resolved worker-selection policy with a statically linked native one. /// /// The factory is called when each decode or prefill worker set is constructed, not per /// request. Workers must advertise an explicit typed role; legacy untyped cards are rejected @@ -84,16 +84,14 @@ impl HttpFrontend { super::initialize_input(&distributed_runtime, &engine_config).await; - let (require_typed_worker_role, selection_policy) = - match self.worker_selection_policy_factory { - Some(factory) => (true, SelectionPolicySource::Factory(factory)), - None => (false, SelectionPolicySource::Registry), - }; + let selection_policy = match self.worker_selection_policy_factory { + Some(factory) => SelectionPolicySource::Factory(factory), + None => SelectionPolicySource::Registry, + }; run_with_selection_policy( distributed_runtime, engine_config, self.frontend_route_extensions, - require_typed_worker_role, selection_policy, ) .await @@ -126,7 +124,6 @@ async fn run_with_selection_policy( distributed_runtime: DistributedRuntime, engine_config: EngineConfig, frontend_route_extensions: Vec, - require_typed_worker_role: bool, selection_policy: SelectionPolicySource, ) -> anyhow::Result<()> { let local_model = engine_config.local_model(); @@ -215,7 +212,6 @@ async fn run_with_selection_policy( model.runtime_config().tokenizer_backend, model.runtime_config().tokenizer_fallback_enabled, generate_engine_capabilities, - require_typed_worker_role, selection_policy.clone(), ) .await?; @@ -305,7 +301,6 @@ async fn run_watcher( tokenizer_backend: Option, tokenizer_fallback_enabled: Option, generate_engine_capabilities: Vec<&'static str>, - require_typed_worker_role: bool, selection_policy: SelectionPolicySource, ) -> anyhow::Result<()> { // Start the LoRA allocation controller when LoRA serving is enabled. The @@ -325,7 +320,6 @@ async fn run_watcher( chat_engine_factory, prefill_load_estimator, metrics.clone(), - require_typed_worker_role, selection_policy, ); watch_obj.set_local_model_path(local_model_path); diff --git a/lib/llm/src/kv_router/routing_host/tests.rs b/lib/llm/src/kv_router/routing_host/tests.rs index d5c811c0ce56..3c547e841b27 100644 --- a/lib/llm/src/kv_router/routing_host/tests.rs +++ b/lib/llm/src/kv_router/routing_host/tests.rs @@ -1390,6 +1390,42 @@ async fn track_request( (request, selection, guard) } +/// Scheduler-reported loads with no prompt, the way the reservation tests read them. +async fn potential_loads(router: &RoutingHost) -> Vec { + router + .kv_router() + .get_potential_loads(&[], None, None, None, None) + .await + .unwrap() +} + +fn active_requests_for( + loads: &[dynamo_kv_router::protocols::PotentialLoad], + worker_id: u64, + dp_rank: u32, +) -> usize { + loads + .iter() + .find(|load| load.worker_id == worker_id && load.dp_rank == dp_rank) + .expect("selected worker must be reported") + .active_requests +} + +/// Preview and admit one decode route, the two stages the plan tests exercise. +async fn plan_decode_route( + router: &RoutingHost, + request: &Context, +) -> RoutePlan { + let preview = router + .preview_kv_route(request, RequestPhase::Decode) + .await + .expect("decode preview should select one request"); + router + .plan_kv_route_from_preview(request, preview) + .await + .expect("decode plan should admit one request") +} + #[tokio::test] #[serial_test::serial] async fn route_plan_from_preview_holds_and_releases_the_decode_reservation() { @@ -1397,40 +1433,18 @@ async fn route_plan_from_preview_holds_and_releases_the_decode_reservation() { let request = Context::new(request()); let requests_started_before = router.request_metrics.requests_started_total().get(); - let preview = router - .preview_kv_route(&request, RequestPhase::Decode) - .await - .expect("decode preview should select one request"); - let plan = router - .plan_kv_route_from_preview(&request, preview) - .await - .expect("decode plan should admit one request"); + let plan = plan_decode_route(&router, &request).await; assert_eq!(plan.signals.worker.worker_id, 7); assert_eq!( router.request_metrics.requests_started_total().get(), requests_started_before, "a topology decision is not a started request" ); - let admitted_loads = router - .kv_router() - .get_potential_loads(&[], None, None, None, None) - .await - .unwrap(); - assert_eq!( - admitted_loads - .iter() - .find(|load| load.worker_id == 7 && load.dp_rank == 0) - .expect("selected worker must be reported") - .active_requests, - 1 - ); + let admitted_loads = potential_loads(&router).await; + assert_eq!(active_requests_for(&admitted_loads, 7, 0), 1); plan.abort().await; - let released_loads = router - .kv_router() - .get_potential_loads(&[], None, None, None, None) - .await - .unwrap(); + let released_loads = potential_loads(&router).await; assert!( released_loads.iter().all(|load| load.active_requests == 0), "abandoned plans must release their scheduler reservation: {released_loads:?}" @@ -1457,11 +1471,7 @@ async fn route_preview_does_not_admit_a_request() { .await .expect("decode preview should select one request"); assert_eq!(preview.signals.worker.worker_id, 7); - let loads = router - .kv_router() - .get_potential_loads(&[], None, None, None, None) - .await - .unwrap(); + let loads = potential_loads(&router).await; assert!(loads.iter().all(|load| load.active_requests == 0)); assert_eq!( router.request_metrics.requests_started_total().get(), @@ -1488,20 +1498,9 @@ async fn route_plan_from_preview_admits_the_previewed_worker() { .await .unwrap(); assert_eq!(plan.signals.worker, previewed_worker); - let loads = router - .kv_router() - .get_potential_loads(&[], None, None, None, None) - .await - .unwrap(); + let loads = potential_loads(&router).await; assert_eq!( - loads - .iter() - .find(|load| { - load.worker_id == previewed_worker.worker_id - && load.dp_rank == previewed_worker.dp_rank - }) - .expect("previewed worker must be reported") - .active_requests, + active_requests_for(&loads, previewed_worker.worker_id, previewed_worker.dp_rank), 1 ); assert_eq!( @@ -1547,25 +1546,14 @@ async fn planned_dispatch_transfers_the_reservation_to_request_cleanup() { let (router, runtime) = router(None).await; let requests_started_before = router.request_metrics.requests_started_total().get(); let request = Context::new(request()); - let preview = router - .preview_kv_route(&request, RequestPhase::Decode) - .await - .unwrap(); - let plan = router - .plan_kv_route_from_preview(&request, preview) - .await - .unwrap(); + let plan = plan_decode_route(&router, &request).await; assert!(router.dispatch_kv_plan(request, plan).await.is_err()); assert_eq!( router.request_metrics.requests_started_total().get(), requests_started_before + 1 ); - let loads = router - .kv_router() - .get_potential_loads(&[], None, None, None, None) - .await - .unwrap(); + let loads = potential_loads(&router).await; assert!(loads.iter().all(|load| load.active_requests == 0)); drop(router); @@ -1580,11 +1568,7 @@ async fn prefill_busy_probe_does_not_admit_a_request() { let requests_started_before = router.request_metrics.requests_started_total().get(); assert!(!router.prefill_worker_busy(&request, 0.5).await.unwrap()); - let loads = router - .kv_router() - .get_potential_loads(&[], None, None, None, None) - .await - .unwrap(); + let loads = potential_loads(&router).await; assert!(loads.iter().all(|load| load.active_requests == 0)); assert_eq!( router.request_metrics.requests_started_total().get(), @@ -1671,11 +1655,7 @@ async fn router_request_counters_follow_admission_and_completion_lifecycle() { assert_eq!(metrics.requests_started_total().get(), started_before + 1); assert_eq!(metrics.requests_total.get(), completed_before); // The booking is held from admission until the guard finishes. - let loads = router - .kv_router() - .get_potential_loads(&[], None, None, None, None) - .await - .unwrap(); + let loads = potential_loads(&router).await; assert_eq!( loads .iter() @@ -1691,11 +1671,7 @@ async fn router_request_counters_follow_admission_and_completion_lifecycle() { drop(cancelled_guard); assert_eq!(metrics.requests_started_total().get(), started_before + 1); assert_eq!(metrics.requests_total.get(), completed_before); - let loads = router - .kv_router() - .get_potential_loads(&[], None, None, None, None) - .await - .unwrap(); + let loads = potential_loads(&router).await; assert!( loads.iter().all(|load| load.active_requests == 0), "an aborted guard frees its booking: {loads:?}" diff --git a/lib/llm/src/session_affinity/tests.rs b/lib/llm/src/session_affinity/tests.rs index 20681ac8430f..0a19f68c90a8 100644 --- a/lib/llm/src/session_affinity/tests.rs +++ b/lib/llm/src/session_affinity/tests.rs @@ -36,6 +36,13 @@ fn target(worker_id: u64, dp_rank: Option) -> AffinityTarget { AffinityTarget { worker_id, dp_rank } } +fn table_target( + worker_id: u64, + dp_rank: Option, +) -> dynamo_kv_router::services::selection::affinity::AffinityTarget { + to_table(target(worker_id, dp_rank)) +} + fn coordinator() -> AffinityCoordinator { AffinityCoordinator::new(Duration::from_secs(10), Hard).unwrap() } @@ -161,7 +168,7 @@ async fn session_affinity_initialization_is_atomic() { coordinator.wait_for_initializing_waiter().await; assert!(!waiter.is_finished()); - let first_lease = first.commit(to_table(target(7, Some(0)))).unwrap(); + let first_lease = first.commit(table_target(7, Some(0))).unwrap(); let second = waiter.await.unwrap().unwrap(); let Hold::Bound { target: second_target, @@ -170,7 +177,7 @@ async fn session_affinity_initialization_is_atomic() { else { panic!("waiter must acquire the committed binding"); }; - assert_eq!(second_target, to_table(target(7, Some(0)))); + assert_eq!(second_target, table_target(7, Some(0))); drop(first_lease); drop(second_lease); } @@ -238,7 +245,7 @@ async fn session_affinity_validates_worker_and_rank_contract() { else { panic!("first request must initialize"); }; - drop(initializer.commit(to_table(target(7, None))).unwrap()); + drop(initializer.commit(table_target(7, None)).unwrap()); assert!( coordinator @@ -267,10 +274,10 @@ async fn session_affinity_failed_bound_operation_invalidates_binding() { else { panic!("first request must initialize"); }; - drop(initializer.commit(to_table(target(7, Some(0)))).unwrap()); + drop(initializer.commit(table_target(7, Some(0))).unwrap()); let operation = coordinator.acquire(&session_id(), None).await.unwrap(); - assert_eq!(operation.target(), Some(to_table(target(7, Some(0))))); + assert_eq!(operation.target(), Some(table_target(7, Some(0)))); operation.invalidate(); assert_eq!(coordinator.query_target(&session_id(), None).unwrap(), None); @@ -288,7 +295,7 @@ async fn session_affinity_stream_drop_refreshes_idle_ttl() { else { panic!("first request must initialize"); }; - let lease = initializer.commit(to_table(target(7, Some(0)))).unwrap(); + let lease = initializer.commit(table_target(7, Some(0))).unwrap(); tokio::time::advance(Duration::from_secs(9)).await; let mut stream = tracked_stream(lease, response_stream(1)); assert!(stream.next().await.is_some()); @@ -304,7 +311,7 @@ async fn session_affinity_empty_stream_refreshes_idle_ttl() { else { panic!("first request must initialize"); }; - let lease = initializer.commit(to_table(target(7, Some(0)))).unwrap(); + let lease = initializer.commit(table_target(7, Some(0))).unwrap(); tokio::time::advance(Duration::from_secs(9)).await; let mut stream = tracked_stream(lease, response_stream(0)); assert!(stream.next().await.is_none()); @@ -319,7 +326,7 @@ async fn session_affinity_cancelled_stream_refreshes_idle_ttl() { else { panic!("first request must initialize"); }; - drop(initializer.commit(to_table(target(7, Some(0)))).unwrap()); + drop(initializer.commit(table_target(7, Some(0))).unwrap()); tokio::time::advance(Duration::from_secs(9)).await; let Hold::Bound { @@ -329,7 +336,7 @@ async fn session_affinity_cancelled_stream_refreshes_idle_ttl() { else { panic!("continuation must acquire the existing binding"); }; - assert_eq!(bound_target, to_table(target(7, Some(0)))); + assert_eq!(bound_target, table_target(7, Some(0))); let mut stream = tracked_stream(lease, cancelled_response_stream()); assert!(stream.next().await.is_none()); @@ -356,7 +363,7 @@ async fn session_affinity_error_stream_refreshes_idle_ttl() { else { panic!("first request must initialize"); }; - let lease = initializer.commit(to_table(target(7, Some(0)))).unwrap(); + let lease = initializer.commit(table_target(7, Some(0))).unwrap(); tokio::time::advance(Duration::from_secs(9)).await; let mut stream = tracked_stream(lease, error_response_stream()); assert!(stream.next().await.unwrap().is_err()); @@ -372,7 +379,7 @@ async fn session_affinity_stream_eof_refreshes_idle_ttl() { else { panic!("first request must initialize"); }; - let lease = initializer.commit(to_table(target(7, Some(0)))).unwrap(); + let lease = initializer.commit(table_target(7, Some(0))).unwrap(); tokio::time::advance(Duration::from_secs(9)).await; let mut stream = tracked_stream(lease, response_stream(1)); while stream.next().await.is_some() {} @@ -387,7 +394,7 @@ async fn session_affinity_bound_lease_drop_refreshes_idle_ttl() { else { panic!("first request must initialize"); }; - drop(initializer.commit(to_table(target(7, Some(0)))).unwrap()); + drop(initializer.commit(table_target(7, Some(0))).unwrap()); tokio::time::advance(Duration::from_secs(9)).await; let Hold::Bound { lease, .. } = coordinator.acquire(&session_id(), None).await.unwrap() else { @@ -420,7 +427,7 @@ async fn session_affinity_query_is_read_only() { else { panic!("first request must initialize"); }; - drop(initializer.commit(to_table(target(7, Some(0)))).unwrap()); + drop(initializer.commit(table_target(7, Some(0))).unwrap()); assert_eq!( coordinator.query_target(&session_id(), None).unwrap(), Some(target(7, Some(0))) @@ -438,7 +445,7 @@ async fn session_affinity_reaper_removes_idle_entries_and_stops_on_drop() { else { panic!("first request must initialize"); }; - drop(initializer.commit(to_table(target(7, Some(0)))).unwrap()); + drop(initializer.commit(table_target(7, Some(0))).unwrap()); coordinator.wait_for_reaper().await; tokio::time::advance(Duration::from_secs(10)).await; From 7b2f6c3faf0783d6e78aa3abdfa86dbb58790ed8 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Sat, 12 Sep 2026 17:42:00 -0700 Subject: [PATCH 074/102] fix(router): track hint capability per partition and clear stale rank gauges Two review items. The core scanned the whole worker catalog on every booking to learn whether any worker in the partition can serve router hints. The pass that publishes the partition's scheduler configuration on every catalog mutation now computes that flag too and stores it on the entry, so a booking reads one atomic. Deleting a worker also removes its catalog record once the partition has published without it, so retired records no longer accumulate; every reader already tolerated a missing record. The frontend's per-rank worker gauges were set from the current record only, so when a worker's data-parallel size shrank the ranks that left kept their last value. The gauge remembers each worker's last rank range and clears the ranks that leave it, and clears everything on removal. Tests: `hint_capability_tracks_catalog_membership` (flag and `kv_hint` follow adds, deletes and re-adds; catalog length after delete) and `registered_gauge_clears_ranks_that_leave_the_worker`. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Ishan Dhanani --- .../src/services/selection/catalog.rs | 32 +++---- .../src/services/selection/core/mod.rs | 5 + .../src/services/selection/core/queries.rs | 2 +- .../src/services/selection/core/run.rs | 2 +- .../src/services/selection/core/tests.rs | 93 ++++++++++++++++++ .../src/services/selection/core/workers.rs | 7 +- .../src/services/selection/membership.rs | 5 +- lib/kv-router/src/services/selection/types.rs | 4 +- lib/llm/src/kv_router/embedded.rs | 95 +++++++++++++++++-- lib/llm/src/kv_router/metrics.rs | 23 +++++ 10 files changed, 239 insertions(+), 29 deletions(-) diff --git a/lib/kv-router/src/services/selection/catalog.rs b/lib/kv-router/src/services/selection/catalog.rs index bfe49a50bda3..b35bf280e7de 100644 --- a/lib/kv-router/src/services/selection/catalog.rs +++ b/lib/kv-router/src/services/selection/catalog.rs @@ -85,34 +85,34 @@ impl WorkerCatalog { .flatten() } - /// Whether any schedulable worker in `key`'s partition can consume router - /// hints. Worker-level metadata, so one representative rank suffices. - pub(super) fn has_router_hint_capable_workers(&self, key: &RoutingPartitionId) -> bool { - self.workers.read().values().any(|record| { - record.lifecycle == WorkerLifecycle::Schedulable - && record.model_name == key.model_name - && record.routing_group == key.routing_group - && record - .router_hint_worker_type - .as_deref() - .is_some_and(|worker_type| !worker_type.is_empty()) - }) + pub(super) fn remove(&self, worker_id: WorkerId) -> Option { + self.workers.write().remove(&worker_id) } - pub(super) fn scheduler_configs_for_key( + /// The schedulable workers in `key`'s partition and whether any of them + /// can consume router hints (worker-level metadata, so one representative + /// rank suffices). One pass over the catalog for both. + pub(super) fn partition_view( &self, key: &RoutingPartitionId, - ) -> HashMap { - self.workers + ) -> (HashMap, bool) { + let mut hint_capable = false; + let workers = self + .workers .read() .values() .filter(|record| schedulable_in(record, key)) .filter_map(|record| { + hint_capable |= record + .router_hint_worker_type + .as_deref() + .is_some_and(|worker_type| !worker_type.is_empty()); record .scheduler_config() .map(|config| (record.worker_id, config)) }) - .collect() + .collect(); + (workers, hint_capable) } pub(super) fn schedulable_count(&self) -> usize { diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index 3989139d6cc6..fb8be575184d 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -7,6 +7,7 @@ use std::collections::HashMap; use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; use std::time::{Duration, Instant}; use dynamo_tokens::SequenceHash; @@ -123,6 +124,10 @@ struct SelectionEntry { is_eagle: bool, indexer: Indexer, workers_tx: watch::Sender>, + /// Some schedulable worker in this partition advertises a router hint + /// worker type. Recomputed by `publish_scheduler_config` on every catalog + /// mutation, so bookings read a bool instead of scanning the catalog. + hint_capable: AtomicBool, scheduler: SelectionScheduler, replica_tx: Option>, affinity: OnceCell, diff --git a/lib/kv-router/src/services/selection/core/queries.rs b/lib/kv-router/src/services/selection/core/queries.rs index 66d862bd08e6..0d5cc4a63fba 100644 --- a/lib/kv-router/src/services/selection/core/queries.rs +++ b/lib/kv-router/src/services/selection/core/queries.rs @@ -99,7 +99,7 @@ impl SelectionCore { } fn schedulable_worker_ranks(&self, key: &RoutingPartitionId) -> Vec { - let configs = self.catalog.scheduler_configs_for_key(key); + let (configs, _) = self.catalog.partition_view(key); let mut workers = Vec::new(); for (worker_id, config) in configs { let start = config.data_parallel_start_rank; diff --git a/lib/kv-router/src/services/selection/core/run.rs b/lib/kv-router/src/services/selection/core/run.rs index 98dd43c57330..8ee43339568d 100644 --- a/lib/kv-router/src/services/selection/core/run.rs +++ b/lib/kv-router/src/services/selection/core/run.rs @@ -299,7 +299,7 @@ impl SelectionCore { // Router hints are attached to bookings only, and only when a worker in // this partition can consume them and the indexer can retain the // matched chain (local, event-driven, no approximate writes). - let hint_capable_workers = book && self.catalog.has_router_hint_capable_workers(&key); + let hint_capable_workers = book && entry.hint_capable.load(Ordering::Acquire); let retain_kv_transfer_chain = hint_capable_workers && entry.indexer.supports_kv_transfer_chain_retention(); if hint_capable_workers && !retain_kv_transfer_chain { diff --git a/lib/kv-router/src/services/selection/core/tests.rs b/lib/kv-router/src/services/selection/core/tests.rs index 557c1a587241..e5ee970283ed 100644 --- a/lib/kv-router/src/services/selection/core/tests.rs +++ b/lib/kv-router/src/services/selection/core/tests.rs @@ -566,6 +566,99 @@ async fn router_hint_needs_capable_workers() { assert!(response.kv_hint.is_none()); } +/// The partition's hint-capability flag follows catalog membership: set by +/// the upsert that adds a capable worker, cleared by the delete that removes +/// the last one, unaffected by other partitions. Deleted workers leave the +/// catalog instead of lingering as `Unschedulable`. +#[tokio::test] +async fn hint_capability_tracks_catalog_membership() { + use crate::indexer::KvIndexerInterface; + use crate::protocols::{BlockHashOptions, compute_block_hash_for_seq}; + + let (core, entry, tokens) = hint_fixture(|_| {}).await; + let hashes: Vec = compute_block_hash_for_seq(&tokens, 4, BlockHashOptions::default()) + .into_iter() + .map(|hash| hash.0) + .collect(); + let seed_worker_1 = || async { + entry + .indexer + .apply_event_routed(store_event(1, 0, 1, &[], &hashes, StorageTier::Device)) + .await + .unwrap(); + if let Indexer::Single { primary, .. } = &entry.indexer { + let _ = primary.flush().await; + } + }; + let capable = |worker_id: WorkerId| { + let mut request = worker_with_kv_events(worker_id); + hint_capable(&mut request); + request + }; + let flag = || entry.hint_capable.load(Ordering::Acquire); + + assert!(!flag()); + let response = reserve_pinned(&core, "before", &tokens, 2).await; + assert!(response.kv_hint.is_none()); + + // A capable worker in another partition does not flip this one. + let mut other = capable(9); + other.routing_group = "group-b".to_string(); + core.upsert_worker(other).await.expect("group-b upsert"); + assert!(!flag()); + + // Re-registering both workers with hint capability enables hints. + core.upsert_worker(capable(1)).await.expect("upsert 1"); + core.upsert_worker(capable(2)).await.expect("upsert 2"); + seed_worker_1().await; + let response = reserve_pinned(&core, "enabled", &tokens, 2).await; + assert!( + response.kv_hint.is_some(), + "capable partition attaches hints" + ); + assert!(flag()); + + // Deleting the capable workers clears the flag and drops their records. + core.delete_worker(1).await.expect("delete 1"); + assert!(flag(), "worker 2 is still capable"); + assert_eq!(core.list_workers(None, None).len(), 2, "1 left the catalog"); + core.delete_worker(2).await.expect("delete 2"); + assert!(!flag()); + assert_eq!( + core.list_workers(None, None).len(), + 1, + "only group-b remains" + ); + assert!(core.list_workers(Some("model"), Some("default")).is_empty()); + + // Non-capable replacements keep the flag clear and hints off. + core.upsert_worker(worker_with_kv_events(1)) + .await + .expect("plain upsert 1"); + core.upsert_worker(worker_with_kv_events(2)) + .await + .expect("plain upsert 2"); + seed_worker_1().await; + assert!(!flag()); + let response = reserve_pinned(&core, "after", &tokens, 2).await; + assert!(response.kv_hint.is_none()); + + // Re-adding capable workers restores hints. + core.upsert_worker(capable(1)) + .await + .expect("upsert 1 again"); + core.upsert_worker(capable(2)) + .await + .expect("upsert 2 again"); + seed_worker_1().await; + assert!(flag()); + let response = reserve_pinned(&core, "restored", &tokens, 2).await; + assert!( + response.kv_hint.is_some(), + "re-added capable worker restores hints" + ); +} + #[tokio::test] async fn event_driven_indexer_does_not_record_bookings() { let core = local_core(test_config(true)); diff --git a/lib/kv-router/src/services/selection/core/workers.rs b/lib/kv-router/src/services/selection/core/workers.rs index f42a17cb391f..297829ff983d 100644 --- a/lib/kv-router/src/services/selection/core/workers.rs +++ b/lib/kv-router/src/services/selection/core/workers.rs @@ -61,6 +61,9 @@ impl SelectionCore { .set_lifecycle(worker_id, WorkerLifecycle::Unschedulable, Vec::new()) .ok_or_else(|| SelectionError::NotFound(format!("worker {worker_id} not found")))?; self.publish_scheduler_config(&key); + // Departed workers do not linger: every reader treats a missing record + // like an unschedulable one, and the catalog stays bounded by live ids. + self.catalog.remove(worker_id); Ok(record) } @@ -294,6 +297,7 @@ impl SelectionCore { is_eagle, indexer, workers_tx, + hint_capable: AtomicBool::new(false), scheduler, replica_tx, affinity: OnceCell::new(), @@ -339,7 +343,8 @@ impl SelectionCore { let Some(entry) = self.entry(key) else { return; }; - let workers = self.catalog.scheduler_configs_for_key(key); + let (workers, hint_capable) = self.catalog.partition_view(key); + entry.hint_capable.store(hint_capable, Ordering::Release); // Lifecycle transitions between non-schedulable states publish the same // map; skipping them saves the scheduler a wake and a full map clone. entry.workers_tx.send_if_modified(|current| { diff --git a/lib/kv-router/src/services/selection/membership.rs b/lib/kv-router/src/services/selection/membership.rs index 09af68b2d589..22b05f86127a 100644 --- a/lib/kv-router/src/services/selection/membership.rs +++ b/lib/kv-router/src/services/selection/membership.rs @@ -240,9 +240,10 @@ mod tests { reconciler.apply(vec![incomplete(1)]).await.expect("apply"); assert_eq!(counter.upserts(), 2); - // The never-schedulable worker was still tracked, so leaving deletes it. + // The never-schedulable worker was still tracked, so leaving deletes it + // and the catalog no longer lists it. reconciler.apply(Vec::new()).await.expect("apply"); - assert_eq!(lifecycle(&core, 1), Some(WorkerLifecycle::Unschedulable)); + assert_eq!(lifecycle(&core, 1), None); assert_eq!(counter.removals(), 1); // Nothing is tracked any more: an empty snapshot deletes nothing. diff --git a/lib/kv-router/src/services/selection/types.rs b/lib/kv-router/src/services/selection/types.rs index a2d347df4cd2..3893f19f6887 100644 --- a/lib/kv-router/src/services/selection/types.rs +++ b/lib/kv-router/src/services/selection/types.rs @@ -159,7 +159,7 @@ pub struct WorkerCatalogRecord { } impl WorkerCatalogRecord { - pub(super) fn new(req: WorkerRequest) -> Self { + pub fn new(req: WorkerRequest) -> Self { Self { worker_id: req.worker_id, model_name: req.model_name, @@ -200,7 +200,7 @@ impl WorkerCatalogRecord { self.data_parallel_size.unwrap_or(1) } - pub fn dp_ranks(&self) -> impl Iterator { + pub fn dp_ranks(&self) -> std::ops::Range { let start = self.dp_start(); let size = self.dp_size(); start..start.saturating_add(size) diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs index 3ed4f4d1849c..b3ee2620ce90 100644 --- a/lib/llm/src/kv_router/embedded.rs +++ b/lib/llm/src/kv_router/embedded.rs @@ -325,12 +325,12 @@ impl EmbeddedSelection { primed: false, }; let mut reconciler = CatalogReconciler::new(Arc::clone(service.core())).with_observer( - Arc::new(RegisteredGauge { - metrics: super::metrics::RouterWorkerStatusMetrics::from_component( + Arc::new(RegisteredGauge::new( + super::metrics::RouterWorkerStatusMetrics::from_component( args.endpoint.component(), ), - worker_label: args.metric_worker_type, - }), + args.metric_worker_type, + )), ); // The current membership is in the catalog before the router serves. if let Some(snapshot) = source.next_snapshot().await @@ -442,18 +442,55 @@ impl WorkerCatalogSource for RuntimeDiscoverySource { struct RegisteredGauge { metrics: Arc, worker_label: &'static str, + /// Last published rank range per worker, so a shrinking or shifted + /// `data_parallel_size` clears the ranks that left instead of stranding + /// them at 1. + ranks: std::sync::Mutex>>, +} + +impl RegisteredGauge { + fn new( + metrics: Arc, + worker_label: &'static str, + ) -> Self { + Self { + metrics, + worker_label, + ranks: std::sync::Mutex::new(HashMap::new()), + } + } } impl CatalogObserver for RegisteredGauge { fn upserted(&self, record: &WorkerCatalogRecord) { - for dp_rank in record.dp_ranks() { + let current = record.dp_ranks(); + let previous = self + .ranks + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .insert(record.worker_id, current.clone()); + for dp_rank in previous + .into_iter() + .flatten() + .filter(|dp_rank| !current.contains(dp_rank)) + { + self.metrics + .remove_worker(record.worker_id, dp_rank, self.worker_label); + } + for dp_rank in current { self.metrics .set_registered(record.worker_id, dp_rank, self.worker_label); } } fn removed(&self, record: &WorkerCatalogRecord) { - for dp_rank in record.dp_ranks() { + let previous = self + .ranks + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .remove(&record.worker_id); + // The record's own ranks plus whatever an earlier upsert published. + for dp_rank in record.dp_ranks().chain(previous.into_iter().flatten()) { self.metrics .remove_worker(record.worker_id, dp_rank, self.worker_label); } @@ -574,6 +611,52 @@ mod tests { ); } + /// A data-parallel shrink clears the gauges of the ranks that left; a + /// removal clears every rank the worker ever published. + #[test] + fn registered_gauge_clears_ranks_that_leave_the_worker() { + let metrics = Arc::new(super::super::metrics::RouterWorkerStatusMetrics::unregistered()); + let gauge = RegisteredGauge::new(Arc::clone(&metrics), "decode"); + let record = |dp_size: u32| { + WorkerCatalogRecord::new(WorkerRequest { + worker_id: 7, + data_parallel_start_rank: Some(0), + data_parallel_size: Some(dp_size), + ..WorkerRequest::default() + }) + }; + // `get_metric_with_label_values` creates the child it looks up; + // `collect` reports only the children that exist. + let registered = |dp_rank: u32| { + use prometheus::core::Collector; + let dp_rank = dp_rank.to_string(); + metrics.registered.collect().into_iter().find_map(|family| { + family + .get_metric() + .iter() + .find(|metric| { + metric + .get_label() + .iter() + .any(|label| label.name() == "dp_rank" && label.value() == dp_rank) + }) + .map(|metric| metric.get_gauge().value() as i64) + }) + }; + + gauge.upserted(&record(4)); + assert!((0..4).all(|dp_rank| registered(dp_rank) == Some(1))); + + gauge.upserted(&record(2)); + assert_eq!(registered(0), Some(1)); + assert_eq!(registered(1), Some(1)); + assert_eq!(registered(2), None, "rank 2 left the worker"); + assert_eq!(registered(3), None, "rank 3 left the worker"); + + gauge.removed(&record(2)); + assert!((0..4).all(|dp_rank| registered(dp_rank).is_none())); + } + #[test] fn local_load_observation_sets_worker_gauges() { let sink = SenderLoadSink { diff --git a/lib/llm/src/kv_router/metrics.rs b/lib/llm/src/kv_router/metrics.rs index a927bc6eaa32..696636ef2cd5 100644 --- a/lib/llm/src/kv_router/metrics.rs +++ b/lib/llm/src/kv_router/metrics.rs @@ -459,6 +459,29 @@ impl RouterWorkerStatusMetrics { .clone() } + /// Gauges on no registry, for observer tests that only read them back. + #[cfg(test)] + pub(crate) fn unregistered() -> Self { + Self { + registered: IntGaugeVec::new( + Opts::new(router::WORKER_REGISTERED, "registered"), + &[ROUTER_WORKER_ID_LABEL, labels::DP_RANK, labels::WORKER_TYPE], + ) + .expect("valid gauge"), + kv_event_source_mismatch_workers: IntGaugeVec::new( + Opts::new(router::KV_EVENT_SOURCE_MISMATCH_WORKERS, "mismatch"), + &[ + labels::MODEL, + labels::WORKER_TYPE, + TARGET_NAMESPACE_LABEL, + TARGET_COMPONENT_LABEL, + TARGET_ENDPOINT_LABEL, + ], + ) + .expect("valid gauge"), + } + } + pub fn set_registered(&self, worker_id: u64, dp_rank: u32, worker_type: &str) { let worker_id = worker_id.to_string(); let dp_rank = dp_rank.to_string(); From 9a18ffb10996782fed9c5f71e01da54a06ebd18e Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Sat, 12 Sep 2026 17:52:28 -0700 Subject: [PATCH 075/102] fix(router): reconcile a catalog snapshot as one batch, retry a failed one Three review items on the catalog reconciler. A snapshot with one bad record used to stop at the first error, leaving stale workers in the catalog and the good records unapplied. `apply` now continues past a failed upsert or delete, still removes the workers the snapshot no longer lists, and returns the first error. A failed snapshot is kept and re-applied after five seconds; a newer snapshot from the source replaces the pending one. An initial snapshot of N workers published the partition's scheduler configuration N times. `upsert_workers` takes the catalog lock once, reconciles every record, and publishes each affected partition once; the draining publish on a partition move stays immediate and the single-worker path is unchanged. Tests: `stale_worker_is_deleted_when_another_upsert_fails`, `failed_snapshot_is_retried_after_the_delay`, `newer_snapshot_supersedes_the_pending_retry`, `initial_snapshot_publishes_once_per_partition`; each fails without its change. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Ishan Dhanani --- .../src/services/selection/core/mod.rs | 11 + .../src/services/selection/core/workers.rs | 77 +++++- .../src/services/selection/membership.rs | 253 ++++++++++++++++-- 3 files changed, 316 insertions(+), 25 deletions(-) diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index fb8be575184d..03f7b41b0f84 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -281,6 +281,13 @@ pub struct SelectionCore { selection_cache: SelectionCache, tracking_hash: Arc, session_affinity: Option, + /// Worker ids whose upsert fails with `Internal` before any catalog + /// mutation, so membership tests can exercise per-worker error paths. + #[cfg(test)] + pub(super) fail_upsert_for: parking_lot::Mutex>, + /// Scheduler-config publishes that changed a partition's worker map. + #[cfg(test)] + pub(super) publish_count: std::sync::atomic::AtomicUsize, } fn affinity_error(error: AffinityError) -> SelectionError { @@ -381,6 +388,10 @@ impl SelectionCore { selection_cache: SelectionCache::new(&cache_config), tracking_hash, session_affinity, + #[cfg(test)] + fail_upsert_for: parking_lot::Mutex::default(), + #[cfg(test)] + publish_count: std::sync::atomic::AtomicUsize::new(0), } } diff --git a/lib/kv-router/src/services/selection/core/workers.rs b/lib/kv-router/src/services/selection/core/workers.rs index 297829ff983d..f744174802cd 100644 --- a/lib/kv-router/src/services/selection/core/workers.rs +++ b/lib/kv-router/src/services/selection/core/workers.rs @@ -4,6 +4,8 @@ //! Worker membership: catalog upserts, patches and deletes, partition //! creation, indexer registration and scheduler config publication. +use std::collections::HashSet; + use super::reservations::{ReservationIndexObserver, spawn_reservation_index_sweep}; use super::*; use crate::services::selection::ingress::remove_worker_from_index; @@ -13,15 +15,55 @@ impl SelectionCore { &self, req: WorkerRequest, ) -> Result { - self.ensure_running()?; - let mut record = WorkerCatalogRecord::new(req); + self.upsert_workers(vec![req]) + .await + .pop() + .expect("one result per request") + } + + /// Upsert one membership snapshot under a single catalog lock, publishing + /// each affected partition's scheduler config once after every record has + /// committed. Results are positional; one failure does not block the rest. + pub async fn upsert_workers( + &self, + requests: Vec, + ) -> Vec> { // Partition policy factories may construct independently. Only committing // membership and reconciling ingress needs the catalog mutation lock. - self.prepare_worker(&mut record)?; + let prepared: Vec> = requests + .into_iter() + .map(|req| { + self.ensure_running()?; + #[cfg(test)] + if self.fail_upsert_for.lock().contains(&req.worker_id) { + return Err(SelectionError::Internal(format!( + "test hook: upsert of worker {} fails", + req.worker_id + ))); + } + let mut record = WorkerCatalogRecord::new(req); + self.prepare_worker(&mut record)?; + Ok(record) + }) + .collect(); let _update = self.catalog_updates.lock().await; - self.ensure_running()?; - let previous = self.catalog.get(record.worker_id); - self.reconcile_worker(record, previous).await + let mut affected: HashSet = HashSet::new(); + let mut results = Vec::with_capacity(prepared.len()); + for record in prepared { + let result = match (record, self.ensure_running()) { + (Ok(record), Ok(())) => { + let previous = self.catalog.get(record.worker_id); + self.reconcile_worker(record, previous, Some(&mut affected)) + .await + } + (Err(error), _) | (_, Err(error)) => Err(error), + }; + results.push(result); + } + for key in &affected { + self.publish_scheduler_config(key); + } + results } pub async fn patch_worker( @@ -38,7 +80,7 @@ impl SelectionCore { let mut record = previous.clone(); record.apply_patch(patch); self.prepare_worker(&mut record)?; - self.reconcile_worker(record, Some(previous)).await + self.reconcile_worker(record, Some(previous), None).await } pub async fn delete_worker( @@ -107,10 +149,16 @@ impl SelectionCore { Ok(()) } + /// Commit `record`. With `deferred`, the final publish for the record's + /// partition is recorded there instead of sent, so a snapshot publishes + /// once per partition; the Draining publish on a partition move or loss + /// of schedulability stays immediate because it must precede the indexer + /// cleanup that follows it. async fn reconcile_worker( &self, mut record: WorkerCatalogRecord, previous: Option, + deferred: Option<&mut HashSet>, ) -> Result { let previous = previous.filter(|old| old.lifecycle == WorkerLifecycle::Schedulable); let previous = if let Some(old) = previous.as_ref() @@ -144,7 +192,12 @@ impl SelectionCore { // Readers see only committed metadata. A valid capacity/topology update preserves // live bookings on ranks present in both the old and new snapshots. self.catalog.replace(record.clone()); - self.publish_scheduler_config(&record.key()); + match deferred { + Some(affected) => { + affected.insert(record.key()); + } + None => self.publish_scheduler_config(&record.key()), + } Ok(record) } @@ -347,7 +400,7 @@ impl SelectionCore { entry.hint_capable.store(hint_capable, Ordering::Release); // Lifecycle transitions between non-schedulable states publish the same // map; skipping them saves the scheduler a wake and a full map clone. - entry.workers_tx.send_if_modified(|current| { + let modified = entry.workers_tx.send_if_modified(|current| { if *current == workers { false } else { @@ -355,5 +408,11 @@ impl SelectionCore { true } }); + #[cfg(test)] + if modified { + self.publish_count.fetch_add(1, Ordering::SeqCst); + } + #[cfg(not(test))] + let _ = modified; } } diff --git a/lib/kv-router/src/services/selection/membership.rs b/lib/kv-router/src/services/selection/membership.rs index 22b05f86127a..afd0ee820485 100644 --- a/lib/kv-router/src/services/selection/membership.rs +++ b/lib/kv-router/src/services/selection/membership.rs @@ -11,6 +11,7 @@ use std::collections::{HashMap, HashSet}; use std::sync::Arc; +use std::time::Duration; use async_trait::async_trait; use tokio_util::sync::CancellationToken; @@ -20,6 +21,9 @@ use super::error::SelectionError; use super::types::{WorkerCatalogRecord, WorkerLifecycle, WorkerRequest}; use crate::protocols::WorkerId; +/// Delay before [`CatalogReconciler::run`] re-applies a snapshot whose pass failed. +const RECONCILE_RETRY_DELAY: Duration = Duration::from_secs(5); + /// Desired worker membership, delivered as complete snapshots rather than /// deltas so a missed change can never leave a stale worker behind. #[async_trait] @@ -45,6 +49,9 @@ pub struct CatalogReconciler { /// Every worker id this reconciler introduced, including ones that never /// became schedulable, so stale deletion covers partial upserts. tracked: HashSet, + /// Passes started by `run`, including retries. + #[cfg(test)] + applies: Arc, } impl CatalogReconciler { @@ -54,6 +61,8 @@ impl CatalogReconciler { observer: None, converged: HashMap::new(), tracked: HashSet::new(), + #[cfg(test)] + applies: Arc::default(), } } @@ -63,8 +72,10 @@ impl CatalogReconciler { } /// Apply one desired snapshot. A snapshot with duplicate worker ids is - /// rejected before any catalog mutation. A failed upsert or delete stops - /// the pass; the next snapshot retries it. + /// rejected before any catalog mutation. A failed upsert or delete does + /// not stop the pass: every other worker is still upserted and every + /// worker absent from the snapshot is still deleted. The first error is + /// returned so the caller retries the snapshot. pub async fn apply(&mut self, desired: Vec) -> Result<(), SelectionError> { let mut by_id: HashMap = HashMap::with_capacity(desired.len()); for request in desired { @@ -76,21 +87,37 @@ impl CatalogReconciler { } } + let mut to_upsert: Vec = Vec::new(); for (worker_id, request) in &by_id { // Track before the upsert so a partially applied record is still // deleted when the worker leaves the desired set. self.tracked.insert(*worker_id); - if self.converged.get(worker_id) == Some(request) { - continue; + if self.converged.get(worker_id) != Some(request) { + to_upsert.push(request.clone()); } - let record = self.core.upsert_worker(request.clone()).await?; + } + // One catalog lock and one scheduler-config publish per partition for + // the whole snapshot, instead of one of each per worker. + let ids: Vec = to_upsert.iter().map(|request| request.worker_id).collect(); + let results = self.core.upsert_workers(to_upsert).await; + let mut first_error: Option = None; + for (worker_id, result) in ids.into_iter().zip(results) { + let record = match result { + Ok(record) => record, + Err(error) => { + tracing::warn!(worker_id, %error, "worker upsert failed; continuing the pass"); + first_error.get_or_insert(error); + continue; + } + }; if let Some(observer) = &self.observer { observer.upserted(&record); } if record.lifecycle == WorkerLifecycle::Schedulable { - self.converged.insert(*worker_id, request.clone()); + let request = by_id.get(&worker_id).expect("upserted id came from by_id"); + self.converged.insert(worker_id, request.clone()); } else { - self.converged.remove(worker_id); + self.converged.remove(&worker_id); tracing::warn!( worker_id, lifecycle = ?record.lifecycle, @@ -115,29 +142,56 @@ impl CatalogReconciler { } // A worker that was never registered is not an error (idempotent). Err(SelectionError::NotFound(_)) => {} - Err(error) => return Err(error), + // Stays tracked so the next pass deletes it. + Err(error) => { + tracing::warn!(worker_id, %error, "stale worker delete failed; continuing the pass"); + first_error.get_or_insert(error); + continue; + } } self.tracked.remove(&worker_id); self.converged.remove(&worker_id); } - Ok(()) + first_error.map_or(Ok(()), Err) } /// Apply every snapshot `source` yields until it closes or `cancel` fires. /// The catalog keeps its last state when the source closes; a source that /// wants selection to fail closed yields an empty snapshot before closing. + /// + /// A snapshot whose pass failed is kept in a single pending slot and + /// re-applied after [`RECONCILE_RETRY_DELAY`]; a newer snapshot from the + /// source supersedes it. Sources block until membership changes, so + /// without this a failed pass would wait for an unrelated change. pub async fn run(mut self, mut source: impl WorkerCatalogSource, cancel: CancellationToken) { + let mut pending: Option> = None; loop { let snapshot = tokio::select! { _ = cancel.cancelled() => return, - snapshot = source.next_snapshot() => snapshot, - }; - let Some(snapshot) = snapshot else { - tracing::debug!("worker membership source closed"); - return; + snapshot = source.next_snapshot() => { + let Some(snapshot) = snapshot else { + tracing::debug!("worker membership source closed"); + return; + }; + snapshot + } + _ = tokio::time::sleep(RECONCILE_RETRY_DELAY), if pending.is_some() => { + pending.take().expect("guarded by pending.is_some()") + } }; - if let Err(error) = self.apply(snapshot).await { - tracing::warn!(%error, "membership reconcile failed; retrying on the next snapshot"); + #[cfg(test)] + self.applies + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); + match self.apply(snapshot.clone()).await { + Ok(()) => pending = None, + Err(error) => { + tracing::warn!( + %error, + retry_in = ?RECONCILE_RETRY_DELAY, + "membership reconcile failed; retrying" + ); + pending = Some(snapshot); + } } } } @@ -277,6 +331,173 @@ mod tests { assert_eq!(record.lifecycle, WorkerLifecycle::Schedulable); } + #[tokio::test] + async fn stale_worker_is_deleted_when_another_upsert_fails() { + let core = core(); + let (mut reconciler, counter) = reconciler(&core); + + reconciler + .apply(vec![schedulable(1), schedulable(2)]) + .await + .expect("apply"); + assert_eq!(counter.upserts(), 2); + + // Worker 2 changes and its upsert fails; worker 1 left the snapshot. + core.fail_upsert_for.lock().insert(2); + let mut changed = schedulable(2); + changed.endpoint = Some("http://10.0.0.9:8000".to_string()); + let error = reconciler + .apply(vec![changed.clone()]) + .await + .expect_err("failed upsert is reported"); + assert!(matches!(error, SelectionError::Internal(_)), "{error}"); + + // The stale worker is still deleted on the pass that observed its absence. + assert_eq!(lifecycle(&core, 1), None); + assert_eq!(counter.removals(), 1); + assert_eq!(reconciler.tracked, HashSet::from([2])); + // The failed worker is still tracked and is retried once the fault clears. + core.fail_upsert_for.lock().clear(); + reconciler.apply(vec![changed]).await.expect("apply"); + assert_eq!(counter.upserts(), 3); + } + + /// Snapshots pushed by the test; closes when the sender drops. + struct ChannelSource(tokio::sync::mpsc::UnboundedReceiver>); + + #[async_trait] + impl WorkerCatalogSource for ChannelSource { + async fn next_snapshot(&mut self) -> Option> { + self.0.recv().await + } + } + + /// Yield until `applies` reaches `expected` without letting the paused + /// clock auto-advance (the test task never parks). + async fn wait_for_applies(applies: &AtomicUsize, expected: usize) { + for _ in 0..1000 { + if applies.load(Ordering::SeqCst) >= expected { + return; + } + tokio::task::yield_now().await; + } + panic!( + "expected {expected} applies, saw {}", + applies.load(Ordering::SeqCst) + ); + } + + #[tokio::test(start_paused = true)] + async fn failed_snapshot_is_retried_after_the_delay() { + let core = core(); + let (reconciler, counter) = reconciler(&core); + let applies = Arc::clone(&reconciler.applies); + let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); + let cancel = CancellationToken::new(); + let task = tokio::spawn(reconciler.run(ChannelSource(rx), cancel.clone())); + + core.fail_upsert_for.lock().insert(1); + tx.send(vec![schedulable(1)]).expect("send"); + wait_for_applies(&applies, 1).await; + assert_eq!(lifecycle(&core, 1), None); + + // The fault clears, but nothing re-applies before the delay elapses. + core.fail_upsert_for.lock().clear(); + tokio::time::advance(RECONCILE_RETRY_DELAY - Duration::from_millis(1)).await; + for _ in 0..10 { + tokio::task::yield_now().await; + } + assert_eq!(applies.load(Ordering::SeqCst), 1); + assert_eq!(lifecycle(&core, 1), None); + + tokio::time::advance(Duration::from_millis(1)).await; + wait_for_applies(&applies, 2).await; + assert_eq!(lifecycle(&core, 1), Some(WorkerLifecycle::Schedulable)); + assert_eq!(counter.upserts(), 1); + + // A successful pass clears the slot: no further retry fires. + tokio::time::advance(RECONCILE_RETRY_DELAY * 2).await; + for _ in 0..10 { + tokio::task::yield_now().await; + } + assert_eq!(applies.load(Ordering::SeqCst), 2); + + cancel.cancel(); + task.await.expect("run exits on cancel"); + } + + #[tokio::test(start_paused = true)] + async fn newer_snapshot_supersedes_the_pending_retry() { + let core = core(); + let (reconciler, counter) = reconciler(&core); + let applies = Arc::clone(&reconciler.applies); + let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); + let cancel = CancellationToken::new(); + let task = tokio::spawn(reconciler.run(ChannelSource(rx), cancel.clone())); + + core.fail_upsert_for.lock().insert(1); + tx.send(vec![schedulable(1)]).expect("send"); + wait_for_applies(&applies, 1).await; + assert_eq!(lifecycle(&core, 1), None); + + // A newer snapshot inside the retry window replaces the pending one. + tokio::time::advance(Duration::from_secs(1)).await; + tx.send(vec![schedulable(2)]).expect("send"); + wait_for_applies(&applies, 2).await; + assert_eq!(lifecycle(&core, 2), Some(WorkerLifecycle::Schedulable)); + + // Even after the fault clears and the window elapses, the superseded + // snapshot is never re-applied: worker 1 stays absent. + core.fail_upsert_for.lock().clear(); + tokio::time::advance(RECONCILE_RETRY_DELAY * 2).await; + for _ in 0..10 { + tokio::task::yield_now().await; + } + assert_eq!(applies.load(Ordering::SeqCst), 2); + assert_eq!(lifecycle(&core, 1), None); + assert_eq!(counter.upserts(), 1); + + cancel.cancel(); + task.await.expect("run exits on cancel"); + } + + #[tokio::test] + async fn initial_snapshot_publishes_once_per_partition() { + let core = core(); + let (mut reconciler, counter) = reconciler(&core); + let publishes = || core.publish_count.load(Ordering::SeqCst); + + reconciler + .apply((1..=4).map(schedulable).collect()) + .await + .expect("apply"); + assert_eq!( + publishes(), + 1, + "one publish for four workers in one partition" + ); + assert_eq!(counter.upserts(), 4, "observer still fires per worker"); + let schedulable_count = core + .list_workers(None, None) + .into_iter() + .filter(|record| record.lifecycle == WorkerLifecycle::Schedulable) + .count(); + assert_eq!(schedulable_count, 4); + + // A second partition in the same snapshot publishes once for itself. + let mut other = schedulable(5); + other.routing_group = "other".to_string(); + reconciler + .apply((1..=4).map(schedulable).chain([other]).collect()) + .await + .expect("apply"); + assert_eq!(publishes(), 2); + + // The single-worker path publishes once per call. + core.upsert_worker(schedulable(6)).await.expect("upsert"); + assert_eq!(publishes(), 3); + } + #[tokio::test] async fn duplicate_ids_are_rejected_before_any_mutation() { let core = core(); From 3cecd2af8d1ae48946c5fe1d3f8a2a0fec38d747 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Sat, 12 Sep 2026 17:55:23 -0700 Subject: [PATCH 076/102] fix(router): read hint capability from the published partition snapshot The partition snapshot the hint builder already borrows carries each worker's hint worker type, so a booking reads capability from that same view instead of a separate flag written on publish. This removes the window between the catalog replace and the publish where the flag could lag, and the atomic itself. The catalog pass goes back to building only the scheduler configuration. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Ishan Dhanani --- .../src/services/selection/catalog.rs | 19 +++++-------------- .../src/services/selection/core/hint.rs | 11 +++++++++++ .../src/services/selection/core/mod.rs | 7 ++----- .../src/services/selection/core/queries.rs | 2 +- .../src/services/selection/core/run.rs | 6 ++++-- .../src/services/selection/core/tests.rs | 2 +- .../src/services/selection/core/workers.rs | 4 +--- 7 files changed, 25 insertions(+), 26 deletions(-) diff --git a/lib/kv-router/src/services/selection/catalog.rs b/lib/kv-router/src/services/selection/catalog.rs index b35bf280e7de..f18692fc5849 100644 --- a/lib/kv-router/src/services/selection/catalog.rs +++ b/lib/kv-router/src/services/selection/catalog.rs @@ -89,30 +89,21 @@ impl WorkerCatalog { self.workers.write().remove(&worker_id) } - /// The schedulable workers in `key`'s partition and whether any of them - /// can consume router hints (worker-level metadata, so one representative - /// rank suffices). One pass over the catalog for both. - pub(super) fn partition_view( + /// The schedulable workers in `key`'s partition, keyed by worker id. + pub(super) fn scheduler_configs_for_key( &self, key: &RoutingPartitionId, - ) -> (HashMap, bool) { - let mut hint_capable = false; - let workers = self - .workers + ) -> HashMap { + self.workers .read() .values() .filter(|record| schedulable_in(record, key)) .filter_map(|record| { - hint_capable |= record - .router_hint_worker_type - .as_deref() - .is_some_and(|worker_type| !worker_type.is_empty()); record .scheduler_config() .map(|config| (record.worker_id, config)) }) - .collect(); - (workers, hint_capable) + .collect() } pub(super) fn schedulable_count(&self) -> usize { diff --git a/lib/kv-router/src/services/selection/core/hint.rs b/lib/kv-router/src/services/selection/core/hint.rs index a01903988943..b3c3685e7114 100644 --- a/lib/kv-router/src/services/selection/core/hint.rs +++ b/lib/kv-router/src/services/selection/core/hint.rs @@ -8,6 +8,17 @@ use super::*; /// Pick the best router-hint source for `target`: a same-role worker (or /// cache owner) holding a longer root-aligned prefix than the target's own /// `target_cached_prefix_blocks`, with a non-empty control endpoint. +/// Whether any worker in the published partition snapshot advertises a router +/// hint worker type (worker-level metadata, so one rank suffices). +pub(super) fn hint_capable_partition(configs: &HashMap) -> bool { + configs.values().any(|config| { + config + .router_hint_worker_type + .as_deref() + .is_some_and(|worker_type| !worker_type.is_empty()) + }) +} + pub(super) fn transfer_hint_for_selection( configs: &HashMap, target: WorkerWithDpRank, diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index 03f7b41b0f84..65e49a900c1b 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -7,7 +7,8 @@ use std::collections::HashMap; use std::sync::Arc; -use std::sync::atomic::{AtomicBool, Ordering}; +#[cfg(test)] +use std::sync::atomic::Ordering; use std::time::{Duration, Instant}; use dynamo_tokens::SequenceHash; @@ -124,10 +125,6 @@ struct SelectionEntry { is_eagle: bool, indexer: Indexer, workers_tx: watch::Sender>, - /// Some schedulable worker in this partition advertises a router hint - /// worker type. Recomputed by `publish_scheduler_config` on every catalog - /// mutation, so bookings read a bool instead of scanning the catalog. - hint_capable: AtomicBool, scheduler: SelectionScheduler, replica_tx: Option>, affinity: OnceCell, diff --git a/lib/kv-router/src/services/selection/core/queries.rs b/lib/kv-router/src/services/selection/core/queries.rs index 0d5cc4a63fba..66d862bd08e6 100644 --- a/lib/kv-router/src/services/selection/core/queries.rs +++ b/lib/kv-router/src/services/selection/core/queries.rs @@ -99,7 +99,7 @@ impl SelectionCore { } fn schedulable_worker_ranks(&self, key: &RoutingPartitionId) -> Vec { - let (configs, _) = self.catalog.partition_view(key); + let configs = self.catalog.scheduler_configs_for_key(key); let mut workers = Vec::new(); for (worker_id, config) in configs { let start = config.data_parallel_start_rank; diff --git a/lib/kv-router/src/services/selection/core/run.rs b/lib/kv-router/src/services/selection/core/run.rs index 8ee43339568d..df2e5b9a6ad4 100644 --- a/lib/kv-router/src/services/selection/core/run.rs +++ b/lib/kv-router/src/services/selection/core/run.rs @@ -36,7 +36,7 @@ //! (`dropped_selection_future_frees_its_booking`, //! `dropped_book_selection_during_routing_record_frees_booking_and_claim`). -use super::hint::transfer_hint_for_selection; +use super::hint::{hint_capable_partition, transfer_hint_for_selection}; use super::reservations::Reservation; use super::*; @@ -299,7 +299,9 @@ impl SelectionCore { // Router hints are attached to bookings only, and only when a worker in // this partition can consume them and the indexer can retain the // matched chain (local, event-driven, no approximate writes). - let hint_capable_workers = book && entry.hint_capable.load(Ordering::Acquire); + // Read from the published partition snapshot, the same view the hint + // builder uses, rather than scanning the catalog per booking. + let hint_capable_workers = book && hint_capable_partition(&entry.workers_tx.borrow()); let retain_kv_transfer_chain = hint_capable_workers && entry.indexer.supports_kv_transfer_chain_retention(); if hint_capable_workers && !retain_kv_transfer_chain { diff --git a/lib/kv-router/src/services/selection/core/tests.rs b/lib/kv-router/src/services/selection/core/tests.rs index e5ee970283ed..e563c62c3f90 100644 --- a/lib/kv-router/src/services/selection/core/tests.rs +++ b/lib/kv-router/src/services/selection/core/tests.rs @@ -595,7 +595,7 @@ async fn hint_capability_tracks_catalog_membership() { hint_capable(&mut request); request }; - let flag = || entry.hint_capable.load(Ordering::Acquire); + let flag = || super::hint::hint_capable_partition(&entry.workers_tx.borrow()); assert!(!flag()); let response = reserve_pinned(&core, "before", &tokens, 2).await; diff --git a/lib/kv-router/src/services/selection/core/workers.rs b/lib/kv-router/src/services/selection/core/workers.rs index f744174802cd..cf4d5f883584 100644 --- a/lib/kv-router/src/services/selection/core/workers.rs +++ b/lib/kv-router/src/services/selection/core/workers.rs @@ -350,7 +350,6 @@ impl SelectionCore { is_eagle, indexer, workers_tx, - hint_capable: AtomicBool::new(false), scheduler, replica_tx, affinity: OnceCell::new(), @@ -396,8 +395,7 @@ impl SelectionCore { let Some(entry) = self.entry(key) else { return; }; - let (workers, hint_capable) = self.catalog.partition_view(key); - entry.hint_capable.store(hint_capable, Ordering::Release); + let workers = self.catalog.scheduler_configs_for_key(key); // Lifecycle transitions between non-schedulable states publish the same // map; skipping them saves the scheduler a wake and a full map clone. let modified = entry.workers_tx.send_if_modified(|current| { From 29b2b9c6324d92d10989806d0807a28ad36011e9 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Sat, 12 Sep 2026 18:08:29 -0700 Subject: [PATCH 077/102] fix(router): construct the worker-selection policy once per partition The frontend built a policy instance only to read which worker inputs it needs, then let the factory run again when the partition scheduler was created, three times per partition on the model-manager path and twice through `KvRouter::new`. `PreparedSelectionPolicy` builds the instance for the embedded partition key once, records its required inputs, and hands that same instance to the first factory call for that partition; later calls (other partitions) go to the underlying factory. The model manager prepares the policy before constructing the router and passes it through as `SelectionPolicySource::Prepared`. Python bindings and EPP keep their `Factory` and `Registry` sources unchanged. Test: a counting factory tags each instance; the router constructs exactly one, reports its inputs, and routes through that instance. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Ishan Dhanani --- lib/llm/src/discovery/model_manager.rs | 14 +- lib/llm/src/discovery/watcher.rs | 6 +- lib/llm/src/kv_router.rs | 179 ++++++++++++++++++++++--- 3 files changed, 170 insertions(+), 29 deletions(-) diff --git a/lib/llm/src/discovery/model_manager.rs b/lib/llm/src/discovery/model_manager.rs index 5f5c32b3eb4d..dbbb903a5fe7 100644 --- a/lib/llm/src/discovery/model_manager.rs +++ b/lib/llm/src/discovery/model_manager.rs @@ -2122,17 +2122,17 @@ impl ModelManager { let effective_kv_router_config = kv_router_config.clone().unwrap_or_default(); let worker_type = worker_role.unwrap_or(WorkerType::Aggregated); - let policy = - policy.resolve(&effective_kv_router_config, worker_type, metric_worker_type)?; - let required_worker_inputs = crate::kv_router::policy_worker_inputs( - &policy, + // One construction for the router's partition: the probed instance is + // the one `KvRouter` hands to the partition scheduler. + let policy = policy.prepare( &effective_kv_router_config, worker_type, + metric_worker_type, model_name.as_deref(), - ); + )?; // A policy that does not consume cache input must not create a shared-cache client or // subscribe to its updates. - let wants_cache = required_worker_inputs.contains(WorkerInputs::CACHE); + let wants_cache = policy.inputs().contains(WorkerInputs::CACHE); let shared_cache: Option> = if wants_cache { match kv_router_config .as_ref() @@ -2181,7 +2181,7 @@ impl ModelManager { workers_with_configs, kv_source_membership, kv_cache_block_size, - SelectionPolicySource::Factory(policy), + SelectionPolicySource::Prepared(policy), kv_router_config, prefill_load_estimator, worker_role, diff --git a/lib/llm/src/discovery/watcher.rs b/lib/llm/src/discovery/watcher.rs index 344019b55791..af6d2417c91b 100644 --- a/lib/llm/src/discovery/watcher.rs +++ b/lib/llm/src/discovery/watcher.rs @@ -1255,7 +1255,11 @@ fn validate_policy_worker_role( card: &ModelDeploymentCard, policy: &SelectionPolicySource, ) -> anyhow::Result<()> { - if matches!(policy, SelectionPolicySource::Factory(_)) && card.worker_type.is_none() { + if matches!( + policy, + SelectionPolicySource::Factory(_) | SelectionPolicySource::Prepared(_) + ) && card.worker_type.is_none() + { anyhow::bail!( "custom worker-selection policies require model cards with an explicit worker_type" ); diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index ee351c0c1bee..e69dc1274bb2 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -10,9 +10,9 @@ use std::{ use anyhow::Result; use dynamo_kv_router::{ - DEFAULT_ROUTING_GROUP, KvSchedulerError, PrefillLoadEstimator, RoutingPartitionRef, - SharedKvCache, TrackingHashAlgorithm, TrackingHashContext, TrackingHashScope, - WorkerSelectionPolicy, WorkerSelectionPolicyFactory, + DEFAULT_ROUTING_GROUP, KvSchedulerError, PrefillLoadEstimator, RoutingPartitionId, + RoutingPartitionRef, SharedKvCache, TrackingHashAlgorithm, TrackingHashContext, + TrackingHashScope, WorkerSelectionPolicy, WorkerSelectionPolicyFactory, config::{KvRouterConfig, RouterConfigOverride, min_initial_workers_from_env}, indexer::{ ApproximateLruIncarnation, ApproximateLruStats, KvRouterError, RoutingDecisionHashes, @@ -94,6 +94,9 @@ pub enum SelectionPolicySource { Registry, /// This factory, called once per routing partition. Factory(WorkerSelectionPolicyFactory), + /// A factory whose instance for the router's own partition is already + /// built and probed; see [`PreparedSelectionPolicy`]. + Prepared(PreparedSelectionPolicy), } impl SelectionPolicySource { @@ -107,6 +110,7 @@ impl SelectionPolicySource { ) -> Result { match self { Self::Factory(factory) => Ok(factory.clone()), + Self::Prepared(prepared) => Ok(prepared.factory.clone()), Self::Registry => Ok( match worker_selection_policy_registry() .resolve_for_worker_type(config, worker_type)? @@ -119,22 +123,76 @@ impl SelectionPolicySource { ), } } + + /// Construct the policy instance for the router's own partition (`model_name` + /// in [`DEFAULT_ROUTING_GROUP`]) once and read its required inputs. An + /// already `Prepared` source is returned as-is, so chained callers never + /// invoke the factory a second time for that partition. + pub(crate) fn prepare( + self, + config: &KvRouterConfig, + worker_type: WorkerType, + label: &'static str, + model_name: Option<&str>, + ) -> Result { + if let Self::Prepared(prepared) = self { + return Ok(prepared); + } + let factory = self.resolve(config, worker_type, label)?; + Ok(PreparedSelectionPolicy::prepare( + factory, + config, + worker_type, + model_name, + )) + } } -/// The optional worker inputs one instance of `factory`'s policy consumes. -pub(crate) fn policy_worker_inputs( - factory: &WorkerSelectionPolicyFactory, - config: &KvRouterConfig, - worker_type: WorkerType, - model_name: Option<&str>, -) -> WorkerInputs { - let partition = RoutingPartitionRef::new( - model_name.unwrap_or(dynamo_kv_router::services::selection::DEFAULT_MODEL_NAME), - DEFAULT_ROUTING_GROUP, - ); - dynamo_kv_router::selector::WorkerSelector::::required_worker_inputs( - &factory(config, worker_type, partition), - ) +/// One constructed policy instance: its required worker inputs, and a factory +/// that hands this instance to the first request for `partition` (the key +/// `embedded::EmbeddedSelection::start` builds) and defers every other call to +/// the wrapped factory. This keeps the one-call-per-partition contract while +/// letting the router read the inputs of the instance that actually serves. +#[derive(Clone)] +pub struct PreparedSelectionPolicy { + factory: WorkerSelectionPolicyFactory, + inputs: WorkerInputs, +} + +impl PreparedSelectionPolicy { + fn prepare( + inner: WorkerSelectionPolicyFactory, + config: &KvRouterConfig, + worker_type: WorkerType, + model_name: Option<&str>, + ) -> Self { + let key = RoutingPartitionId::new( + model_name.unwrap_or(dynamo_kv_router::services::selection::DEFAULT_MODEL_NAME), + DEFAULT_ROUTING_GROUP, + ); + let policy = inner(config, worker_type, key.as_ref()); + let inputs = + dynamo_kv_router::selector::WorkerSelector::::required_worker_inputs( + &policy, + ); + let slot = Arc::new(parking_lot::Mutex::new(Some(policy))); + let factory: WorkerSelectionPolicyFactory = Arc::new( + move |config: &KvRouterConfig, worker_type, partition: RoutingPartitionRef<'_>| { + if partition == key.as_ref() + && let Some(policy) = slot.lock().take() + { + return policy; + } + inner(config, worker_type, partition) + }, + ); + Self { factory, inputs } + } + + /// The optional worker inputs the prepared instance consumes. + pub fn inputs(&self) -> WorkerInputs { + self.inputs + } } #[derive(Clone, Copy)] @@ -647,13 +705,14 @@ impl KvRouter { let kv_router_config = kv_router_config.unwrap_or_default(); kv_router_config.validate().map_err(anyhow::Error::msg)?; let worker_type = worker_role.unwrap_or(WorkerType::Aggregated); - let policy_factory = policy.resolve(&kv_router_config, worker_type, metric_worker_type)?; - let required_worker_inputs = policy_worker_inputs( - &policy_factory, + let prepared = policy.prepare( &kv_router_config, worker_type, + metric_worker_type, model_name.as_deref(), - ); + )?; + let required_worker_inputs = prepared.inputs(); + let policy_factory = prepared.factory; // ModelManager gates client construction as well, but preserve the capability boundary for // direct KvRouter callers. let shared_cache = if required_worker_inputs.contains(WorkerInputs::CACHE) { @@ -2615,6 +2674,84 @@ mod tests { .unwrap() } + /// Picks worker 0 and records which construction it belongs to. Its + /// required inputs differ by tag so the router's probed inputs identify + /// the instance they were read from. + struct TaggedPicker { + tag: usize, + picked_tag: Arc>>, + } + + impl WorkerPicker for TaggedPicker { + fn required_worker_inputs(&self) -> WorkerInputs { + if self.tag == 1 { + WorkerInputs::LOAD + } else { + WorkerInputs::CACHE | WorkerInputs::LOAD + } + } + + fn pick( + &mut self, + _context: &WorkerSelectionContext<'_>, + input: WorkerInputView<'_>, + ) -> Result { + *self.picked_tag.lock() = Some(self.tag); + input + .candidates() + .iter() + .position(|candidate| candidate.worker() == WorkerWithDpRank::from_worker_id(0)) + .ok_or_else(|| WorkerSelectionPolicyError::failed("worker 0 not eligible")) + } + } + + /// The factory runs once for the router's partition, and the instance + /// whose inputs the router read is the instance that serves selections. + #[tokio::test] + async fn policy_factory_runs_once_and_the_probed_instance_serves() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + let constructions = Arc::new(AtomicUsize::new(0)); + let picked_tag = Arc::new(parking_lot::Mutex::new(None)); + let counting = { + let constructions = Arc::clone(&constructions); + let picked_tag = Arc::clone(&picked_tag); + SelectionPolicySource::Factory(Arc::new(move |config: &KvRouterConfig, _, _| { + let tag = constructions.fetch_add(1, Ordering::SeqCst) + 1; + WorkerSelectionPolicy::new( + config.clone(), + "decode", + Vec::new(), + Box::new(TaggedPicker { + tag, + picked_tag: Arc::clone(&picked_tag), + }), + ) + })) + }; + + let router = make_test_router(counting, None).await; + assert_eq!( + constructions.load(Ordering::SeqCst), + 1, + "factory must run once for the router's partition" + ); + assert_eq!(router.required_worker_inputs(), WorkerInputs::LOAD); + + let FindBestMatchOutcome::Routed { worker, .. } = + find_best_match(&router, &[11, 12], false).await.unwrap() + else { + panic!("expected routed outcome"); + }; + assert_eq!(worker, WorkerWithDpRank::from_worker_id(0)); + assert_eq!( + *picked_tag.lock(), + Some(1), + "the probed instance must serve the partition" + ); + assert_eq!(constructions.load(Ordering::SeqCst), 1); + } + #[tokio::test] async fn test_find_best_match_passes_shared_cache_hits_to_scheduler() { let router = make_test_router( From 4cc6f49acf010e3fee393ee5972d00b95e3ba117 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Sat, 12 Sep 2026 18:11:16 -0700 Subject: [PATCH 078/102] fix(router): publish a completed partition move before the batch continues Review follow-ups on the batched reconcile: a worker that finished moving partitions publishes its new partition immediately instead of at the end of the batch, so the destination is routable even if a later worker in the batch blocks in indexer cleanup; the test-only apply counter advances after the pass completes; the publish test checks the published map carries all four workers; the retry delay is stated in the module doc and the standalone-selection registration section. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Ishan Dhanani --- .../router/standalone-selection.md | 5 ++++- .../src/services/selection/core/workers.rs | 15 ++++++++----- .../src/services/selection/membership.rs | 21 +++++++++++++++++-- 3 files changed, 33 insertions(+), 8 deletions(-) diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-selection.md b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-selection.md index 0823c7135cde..aef639e67fcc 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-selection.md +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-selection.md @@ -114,7 +114,10 @@ after the last progress touch. ## Worker Registration Every selector replica must receive the same worker catalog before it serves -selection traffic. Replica traffic never creates workers. +selection traffic. Replica traffic never creates workers. When the catalog comes +from a discovery source (Kubernetes, a workers file), a snapshot whose apply +fails is retried after 5 seconds unless a newer snapshot arrives first, so a +stale worker can persist for that long after a failed pass. ```http POST /workers diff --git a/lib/kv-router/src/services/selection/core/workers.rs b/lib/kv-router/src/services/selection/core/workers.rs index cf4d5f883584..6248b02f4dfa 100644 --- a/lib/kv-router/src/services/selection/core/workers.rs +++ b/lib/kv-router/src/services/selection/core/workers.rs @@ -151,9 +151,11 @@ impl SelectionCore { /// Commit `record`. With `deferred`, the final publish for the record's /// partition is recorded there instead of sent, so a snapshot publishes - /// once per partition; the Draining publish on a partition move or loss - /// of schedulability stays immediate because it must precede the indexer - /// cleanup that follows it. + /// once per partition. Two publishes stay immediate even then: the + /// Draining publish on a partition move or loss of schedulability, which + /// must precede the indexer cleanup that follows it, and the publish of a + /// completed move, so the destination partition becomes routable before + /// the rest of the batch (which may block in indexer cleanup) finishes. async fn reconcile_worker( &self, mut record: WorkerCatalogRecord, @@ -161,6 +163,9 @@ impl SelectionCore { deferred: Option<&mut HashSet>, ) -> Result { let previous = previous.filter(|old| old.lifecycle == WorkerLifecycle::Schedulable); + let moved_partition = previous + .as_ref() + .is_some_and(|old| old.key() != record.key()); let previous = if let Some(old) = previous.as_ref() && (old.key() != record.key() || !record.not_schedulable_reasons.is_empty()) { @@ -193,10 +198,10 @@ impl SelectionCore { // live bookings on ranks present in both the old and new snapshots. self.catalog.replace(record.clone()); match deferred { - Some(affected) => { + Some(affected) if !moved_partition => { affected.insert(record.key()); } - None => self.publish_scheduler_config(&record.key()), + _ => self.publish_scheduler_config(&record.key()), } Ok(record) } diff --git a/lib/kv-router/src/services/selection/membership.rs b/lib/kv-router/src/services/selection/membership.rs index afd0ee820485..1729ba39224c 100644 --- a/lib/kv-router/src/services/selection/membership.rs +++ b/lib/kv-router/src/services/selection/membership.rs @@ -7,7 +7,8 @@ //! workers file) and exposes it as a [`WorkerCatalogSource`]: a stream of //! complete desired snapshots. [`CatalogReconciler`] turns each snapshot into //! catalog upserts and deletes, retrying workers the core has not yet accepted -//! as schedulable and removing workers that left. +//! as schedulable and removing workers that left. A snapshot whose pass fails +//! is re-applied after [`RECONCILE_RETRY_DELAY`] unless a newer one arrives. use std::collections::{HashMap, HashSet}; use std::sync::Arc; @@ -179,10 +180,11 @@ impl CatalogReconciler { pending.take().expect("guarded by pending.is_some()") } }; + let outcome = self.apply(snapshot.clone()).await; #[cfg(test)] self.applies .fetch_add(1, std::sync::atomic::Ordering::SeqCst); - match self.apply(snapshot.clone()).await { + match outcome { Ok(()) => pending = None, Err(error) => { tracing::warn!( @@ -477,6 +479,21 @@ mod tests { "one publish for four workers in one partition" ); assert_eq!(counter.upserts(), 4, "observer still fires per worker"); + // The scheduler applies the published map asynchronously; the one + // publish must have carried all four workers. + let published = || -> usize { + core.loads(None, None) + .iter() + .map(|load| load.loads.len()) + .sum() + }; + for _ in 0..1000 { + if published() == 4 { + break; + } + tokio::task::yield_now().await; + } + assert_eq!(published(), 4, "the one publish carries all four workers"); let schedulable_count = core .list_workers(None, None) .into_iter() From 9e1e92a87bdaf63d6aaddabf52b9f7c96c65734e Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Sat, 12 Sep 2026 18:28:35 -0700 Subject: [PATCH 079/102] fix(router): keep a lease when a session commit lands behind a competing failover Two requests can hold the binding of a worker that then departs. The first re-initializes the session and books a replacement; the second's commit found the session initializing, got `Wait`, and proceeded with no lease on the new binding, so the binding could expire under it or the session's next request could land elsewhere. The affinity table now lets a committing request join an in-flight initialization: `join_initializing` hands out a lease that counts toward the binding's active leases once the initializer commits, and is a no-op if the initializer is dropped. `commit_session` uses it on `Wait`, retrying `try_acquire` once if the initialization finished in between. The commit still never waits, so the ordering invariant in the module doc stays true; its wording now covers the join. Tests: `joined_lease_counts_once_the_competing_initialization_commits`, `joined_lease_is_a_noop_after_the_initializer_drops` (table), and `failover_commit_behind_an_initializing_hold_keeps_a_lease` (core, hand-polled interleaving: both requests bound to the replacement, lease count 2 to 1 to 0, binding survives to its TTL). Each fails without the change. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Ishan Dhanani --- .../src/services/selection/affinity.rs | 161 +++++++++++++++++- .../src/services/selection/core/run.rs | 41 +++-- .../src/services/selection/core/tests.rs | 63 +++++++ 3 files changed, 247 insertions(+), 18 deletions(-) diff --git a/lib/kv-router/src/services/selection/affinity.rs b/lib/kv-router/src/services/selection/affinity.rs index 270dca637087..9a1bc8006bea 100644 --- a/lib/kv-router/src/services/selection/affinity.rs +++ b/lib/kv-router/src/services/selection/affinity.rs @@ -71,6 +71,15 @@ pub struct AffinityVersion { pub writer_id: u64, } +impl AffinityVersion { + /// Version of a lease joined to an initialization that has not committed + /// yet; it matches no stored binding, so such a lease can only release. + pub const PENDING: Self = Self { + sequence: u64::MAX, + writer_id: u64::MAX, + }; +} + /// Receives every binding this table publishes for its replicas. pub trait AffinityReplicaSink: Send + Sync { fn publish(&self, session_id: &str, target: AffinityTarget, version: AffinityVersion); @@ -92,6 +101,9 @@ enum AffinityEntry { Initializing { revision: u64, notify: Arc, + /// Leases handed out by [`SessionAffinity::join_initializing`]; they + /// count as active once the initializer commits. + pending_leases: usize, }, Bound { target: AffinityTarget, @@ -333,6 +345,7 @@ impl SessionAffinity { entry.insert(AffinityEntry::Initializing { revision, notify: notify.clone(), + pending_leases: 0, }); Ok(AcquireStep::Held(Hold::Initialize( AffinityInitialization { @@ -369,6 +382,7 @@ impl SessionAffinity { *entry.get_mut() = AffinityEntry::Initializing { revision, notify: notify.clone(), + pending_leases: 0, }; Ok(AcquireStep::Held(Hold::Initialize( AffinityInitialization { @@ -427,6 +441,31 @@ impl SessionAffinity { } } + /// Take a lease on the binding another request is initializing, so this + /// request's use counts once that commit lands and the binding cannot + /// idle out under it. `None` when the entry is not initializing (the + /// caller falls back to [`Self::try_acquire`]). If the initializer drops + /// uncommitted, the lease releases as a no-op. + pub fn join_initializing(&self, session_id: &str) -> Option { + let mut entry = self.inner.entries.get_mut(session_id)?; + let AffinityEntry::Initializing { + revision, + pending_leases, + .. + } = entry.value_mut() + else { + return None; + }; + *pending_leases += 1; + Some(AffinityLease { + coordinator: Arc::downgrade(&self.inner), + session_id: session_id.to_string(), + revision: *revision, + version: AffinityVersion::PENDING, + active: true, + }) + } + /// Bind the held session to `dispatched`. A `Bound` session dispatched /// elsewhere follows the dispatch in `Soft` mode and is invalidated with /// an error in `Hard` mode, so its next request re-initializes. @@ -562,6 +601,16 @@ impl SessionAffinity { pub fn next_version(&self) -> AffinityVersion { self.inner.next_version() } + + /// Leases on `session_id`: pending joiners while initializing, active + /// leases once bound. `None` when there is no entry. + pub fn lease_count(&self, session_id: &str) -> Option { + let entry = self.inner.entries.get(session_id)?; + Some(match entry.value() { + AffinityEntry::Initializing { pending_leases, .. } => *pending_leases, + AffinityEntry::Bound { active_leases, .. } => *active_leases, + }) + } } impl Inner { @@ -690,10 +739,17 @@ impl AffinityInitialization { "session affinity initialization was cancelled".to_string(), )); }; - if !matches!( - entry.value(), - AffinityEntry::Initializing { revision, .. } if *revision == self.revision - ) { + let AffinityEntry::Initializing { + revision, + pending_leases, + .. + } = entry.value() + else { + return Err(AffinityError::InvalidArgument( + "session affinity initialization changed".to_string(), + )); + }; + if *revision != self.revision { return Err(AffinityError::InvalidArgument( "session affinity initialization changed".to_string(), )); @@ -703,7 +759,8 @@ impl AffinityInitialization { target, revision: self.revision, version, - active_leases: 1, + // Joined leases become live uses of the binding they waited for. + active_leases: 1 + pending_leases, idle_deadline: Instant::now() + inner.ttl, }; drop(entry); @@ -827,7 +884,9 @@ impl AffinityLease { return; } *active_leases -= 1; - if *version != self.version { + // A joined lease was a use of whatever binding the initializer + // committed, so it refreshes that binding's idle deadline. + if *version != self.version && self.version != AffinityVersion::PENDING { return; } *idle_deadline = Instant::now() + inner.ttl; @@ -943,3 +1002,93 @@ pub fn validate_dispatch_target( } Ok(()) } + +#[cfg(test)] +mod tests { + use super::*; + + const TTL: Duration = Duration::from_secs(10); + + fn table() -> SessionAffinity { + SessionAffinity::with_config(SessionAffinityConfig::new(TTL)).expect("affinity table") + } + + fn initialize(table: &SessionAffinity) -> AffinityInitialization { + match table.try_acquire("s", None).expect("acquire") { + AcquireStep::Held(Hold::Initialize(init)) => init, + AcquireStep::Held(Hold::Bound { .. }) => panic!("session is already bound"), + AcquireStep::Wait(_) => panic!("session is being initialized"), + } + } + + #[tokio::test(start_paused = true)] + async fn joined_lease_counts_once_the_competing_initialization_commits() { + let table = table(); + let target = AffinityTarget::new(2, Some(0)); + let init = initialize(&table); + let joined = table + .join_initializing("s") + .expect("join an initializing session"); + assert_eq!(table.lease_count("s"), Some(1)); + assert!( + table.join_initializing("x").is_none(), + "no join without an initializer" + ); + + let lease = init.commit(target).expect("commit"); + assert_eq!(table.lease_count("s"), Some(2)); + assert_eq!(table.query_target("s", None).expect("query"), Some(target)); + + drop(lease); + assert_eq!(table.lease_count("s"), Some(1)); + // The joiner's release is a use of the binding: it refreshes the idle + // deadline, so the binding outlives the initializer's own TTL. + tokio::time::advance(TTL - Duration::from_secs(1)).await; + drop(joined); + assert_eq!(table.lease_count("s"), Some(0)); + tokio::time::advance(Duration::from_secs(2)).await; + assert_eq!( + table.query_target("s", None).expect("query"), + Some(target), + "joined lease release must refresh the idle deadline" + ); + tokio::time::advance(TTL).await; + assert_eq!(table.query_target("s", None).expect("query"), None); + assert!( + matches!( + table.try_acquire("s", None).expect("acquire"), + AcquireStep::Held(Hold::Initialize(_)) + ), + "an idled binding re-initializes" + ); + } + + #[tokio::test(start_paused = true)] + async fn joined_lease_is_a_noop_after_the_initializer_drops() { + let table = table(); + let init = initialize(&table); + let joined = table + .join_initializing("s") + .expect("join an initializing session"); + assert_eq!(table.entry_count(), 1); + + drop(init); + assert_eq!( + table.entry_count(), + 0, + "a dropped initializer frees the slot" + ); + assert_eq!(table.lease_count("s"), None); + drop(joined); + assert_eq!(table.entry_count(), 0); + + // The session is re-initializable and the stale joiner left no count. + let init = initialize(&table); + let lease = init + .commit(AffinityTarget::new(3, Some(0))) + .expect("commit"); + assert_eq!(table.lease_count("s"), Some(1)); + drop(lease); + assert_eq!(table.lease_count("s"), Some(0)); + } +} diff --git a/lib/kv-router/src/services/selection/core/run.rs b/lib/kv-router/src/services/selection/core/run.rs index df2e5b9a6ad4..7a3859d7dea9 100644 --- a/lib/kv-router/src/services/selection/core/run.rs +++ b/lib/kv-router/src/services/selection/core/run.rs @@ -23,10 +23,13 @@ //! (`session_worker_departing_after_the_hold_reinitializes_the_session`) and //! the binding is committed while the booking is held. The commit never //! waits: `commit_session` is a synchronous fn that re-binds only through -//! `try_acquire`, so when the table cannot bind at commit time (`try_acquire` -//! returns `Wait`, or the table is full) it returns `Ok(None)` and the -//! request routes unpinned rather than waiting on a request queued behind -//! this booking. A rejected commit frees the booking +//! `try_acquire`. A commit that finds another request initializing the +//! session joins that initialization instead of waiting: it takes a lease +//! that counts once the initializer commits +//! (`failover_commit_behind_an_initializing_hold_keeps_a_lease`). When the +//! table is full it returns `Ok(None)` and the request routes unpinned +//! rather than waiting on a request queued behind this booking. A rejected +//! commit frees the booking //! (`two_phase_replay_rejects_a_worker_the_session_left`). //! - A `Lease` admission installs no index row and records no routing hashes //! in the core; the host owns both @@ -616,15 +619,29 @@ impl SelectionCore { if !departed { return Err(affinity_error(error)); } - // `commit` invalidated the stale binding. - match table.try_acquire(session_id, None) { - Ok(AcquireStep::Held(hold)) => table - .commit(hold, dispatched) - .map(Some) - .map_err(affinity_error), - Ok(AcquireStep::Wait(_)) | Err(AffinityError::ResourceExhausted(_)) => Ok(None), - Err(error) => Err(affinity_error(error)), + // `commit` invalidated the stale binding. Another request may + // already be initializing the replacement: join it rather + // than wait, so this booking holds a lease on the new + // binding. The join misses only if that initialization + // resolved between the two calls, so retry once. + for _ in 0..2 { + match table.try_acquire(session_id, None) { + Ok(AcquireStep::Held(hold)) => { + return table + .commit(hold, dispatched) + .map(Some) + .map_err(affinity_error); + } + Ok(AcquireStep::Wait(_)) => { + if let Some(lease) = table.join_initializing(session_id) { + return Ok(Some(lease)); + } + } + Err(AffinityError::ResourceExhausted(_)) => return Ok(None), + Err(error) => return Err(affinity_error(error)), + } } + Ok(None) } } } diff --git a/lib/kv-router/src/services/selection/core/tests.rs b/lib/kv-router/src/services/selection/core/tests.rs index e563c62c3f90..2e196ce145de 100644 --- a/lib/kv-router/src/services/selection/core/tests.rs +++ b/lib/kv-router/src/services/selection/core/tests.rs @@ -2173,6 +2173,69 @@ async fn concurrent_holds_on_a_departed_worker_both_land_on_the_replacement() { ); } +/// Leases on session `s` in the default partition's table. +fn lease_count(core: &SelectionCore, session_id: &str) -> Option { + let entry = core.entry(&default_key()).expect("default partition"); + let table = entry.affinity.get().expect("affinity table configured"); + table.lease_count(session_id) +} + +/// r2 holds the departed binding when r3 re-initializes the session; r2's +/// commit then finds r3's initialization in progress and must join it, not +/// route without a lease, so the replacement binding cannot idle out while +/// r2 is still active. +#[tokio::test] +async fn failover_commit_behind_an_initializing_hold_keeps_a_lease() { + let (core, first) = bound_session(SessionAffinityMode::Hard).await; + let key = default_key(); + let replacement = if first.worker_id == 1 { 2 } else { 1 }; + let mut context = std::task::Context::from_waker(std::task::Waker::noop()); + + // r2 takes `Bound{w1}` and is queued in the scheduler actor. + let mut r2 = Box::pin(core.select_and_reserve(session_reservation("r2", "s"))); + assert!(r2.as_mut().poll(&mut context).is_pending()); + core.catalog + .set_lifecycle(first.worker_id, WorkerLifecycle::Draining, Vec::new()); + core.publish_scheduler_config(&key); + + // r3 sees w1 unschedulable, invalidates the binding, and holds + // `Initialize` for as long as it sits in the actor unpolled. + let mut r3 = Box::pin(core.select_and_reserve(session_reservation("r3", "s"))); + assert!(r3.as_mut().poll(&mut context).is_pending()); + assert_eq!(bound_worker(&core, "s"), None, "r3 is initializing"); + assert_eq!(lease_count(&core, "s"), Some(0)); + + // r2's `Hard` commit fails, the worker departed, and `try_acquire` finds + // r3's initialization: r2 joins it. + let r2 = r2 + .await + .expect("a departure after the hold is not a client fault"); + assert_eq!(r2.worker_id, replacement); + assert_eq!(bound_worker(&core, "s"), None, "r3 has not committed yet"); + assert_eq!( + lease_count(&core, "s"), + Some(1), + "r2 must hold a pending lease on r3's initialization" + ); + + let r3 = r3.await.expect("r3 binds the replacement"); + assert_eq!(r3.worker_id, replacement); + assert_eq!(bound_worker(&core, "s"), Some(replacement)); + assert_eq!(lease_count(&core, "s"), Some(2)); + + core.free_reservation("r3").await.expect("free r3"); + wait_until("r3 lease release", || lease_count(&core, "s") == Some(1)).await; + assert_eq!(bound_worker(&core, "s"), Some(replacement)); + + core.free_reservation("r2").await.expect("free r2"); + wait_until("r2 lease release", || lease_count(&core, "s") == Some(0)).await; + // Idle, not gone: the binding lasts until the TTL. + assert_eq!(bound_worker(&core, "s"), Some(replacement)); + let entry = core.entry(&key).expect("default partition"); + entry.affinity.get().expect("table").expire_for_test("s"); + assert_eq!(bound_worker(&core, "s"), None); +} + #[tokio::test] async fn two_phase_reservation_binds_the_session() { let core = core_with_session_affinity(); From 2dfc529eb22fee0b16f9adac78f5221197214133 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Sat, 12 Sep 2026 18:28:35 -0700 Subject: [PATCH 080/102] fix(router): derive the embedded partition key in one place The prepared policy probed a key it computed itself and the embedded host rebuilt the same key independently; a drift between the two would have silently recreated the double construction. `embedded_partition_key` is the one definition both use. A test pins the named-model path and that preparing an already prepared source does not probe again. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Ishan Dhanani --- lib/llm/src/kv_router.rs | 57 +++++++++++++++++++++++++++---- lib/llm/src/kv_router/embedded.rs | 17 +++++---- 2 files changed, 61 insertions(+), 13 deletions(-) diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index e69dc1274bb2..db8d3c5e27de 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -10,9 +10,9 @@ use std::{ use anyhow::Result; use dynamo_kv_router::{ - DEFAULT_ROUTING_GROUP, KvSchedulerError, PrefillLoadEstimator, RoutingPartitionId, - RoutingPartitionRef, SharedKvCache, TrackingHashAlgorithm, TrackingHashContext, - TrackingHashScope, WorkerSelectionPolicy, WorkerSelectionPolicyFactory, + DEFAULT_ROUTING_GROUP, KvSchedulerError, PrefillLoadEstimator, RoutingPartitionRef, + SharedKvCache, TrackingHashAlgorithm, TrackingHashContext, TrackingHashScope, + WorkerSelectionPolicy, WorkerSelectionPolicyFactory, config::{KvRouterConfig, RouterConfigOverride, min_initial_workers_from_env}, indexer::{ ApproximateLruIncarnation, ApproximateLruStats, KvRouterError, RoutingDecisionHashes, @@ -166,10 +166,7 @@ impl PreparedSelectionPolicy { worker_type: WorkerType, model_name: Option<&str>, ) -> Self { - let key = RoutingPartitionId::new( - model_name.unwrap_or(dynamo_kv_router::services::selection::DEFAULT_MODEL_NAME), - DEFAULT_ROUTING_GROUP, - ); + let key = embedded::embedded_partition_key(model_name); let policy = inner(config, worker_type, key.as_ref()); let inputs = dynamo_kv_router::selector::WorkerSelector::::required_worker_inputs( @@ -2705,6 +2702,52 @@ mod tests { } } + /// The prepared wrapper hands its parked instance to the embedded + /// partition key for a named model too, and a `Prepared` source is not + /// probed again. + #[test] + fn prepared_policy_matches_the_named_model_partition_and_is_not_reprepared() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + let constructions = Arc::new(AtomicUsize::new(0)); + let counting: WorkerSelectionPolicyFactory = { + let constructions = Arc::clone(&constructions); + Arc::new(move |config: &KvRouterConfig, _, _| { + constructions.fetch_add(1, Ordering::SeqCst); + WorkerSelectionPolicy::default(config.clone(), "decode") + }) + }; + let config = KvRouterConfig::default(); + let prepared = PreparedSelectionPolicy::prepare( + counting, + &config, + WorkerType::Decode, + Some("named-model"), + ); + assert_eq!(constructions.load(Ordering::SeqCst), 1); + let inputs = prepared.inputs(); + + // The first call for the embedded key takes the parked instance. + let key = embedded::embedded_partition_key(Some("named-model")); + let _served = (prepared.factory)(&config, WorkerType::Decode, key.as_ref()); + assert_eq!( + constructions.load(Ordering::SeqCst), + 1, + "parked instance served" + ); + // Another partition constructs its own. + let other = dynamo_kv_router::RoutingPartitionId::new("other-model", DEFAULT_ROUTING_GROUP); + let _other = (prepared.factory)(&config, WorkerType::Decode, other.as_ref()); + assert_eq!(constructions.load(Ordering::SeqCst), 2); + + // Preparing an already prepared source is a passthrough. + let again = SelectionPolicySource::Prepared(prepared) + .prepare(&config, WorkerType::Decode, "decode", Some("named-model")) + .expect("prepare"); + assert_eq!(again.inputs(), inputs, "inputs carried through unchanged"); + assert_eq!(constructions.load(Ordering::SeqCst), 2, "no second probe"); + } + /// The factory runs once for the router's partition, and the instance /// whose inputs the router read is the instance that serves selections. #[tokio::test] diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs index b3ee2620ce90..6750ace1a8fa 100644 --- a/lib/llm/src/kv_router/embedded.rs +++ b/lib/llm/src/kv_router/embedded.rs @@ -150,6 +150,16 @@ fn non_max_overlap_observer(worker_type: &'static str) -> NonMaxOverlapSelection }) } +/// The partition an embedded router serves: `model_name` (or the default) in +/// the default routing group. The prepared policy probes the same key, so the +/// instance it inspected is the one this partition receives. +pub(crate) fn embedded_partition_key(model_name: Option<&str>) -> RoutingPartitionId { + RoutingPartitionId::new( + model_name.unwrap_or(DEFAULT_MODEL_NAME), + DEFAULT_ROUTING_GROUP, + ) +} + /// One `SelectionService` partition driven directly by the router. pub(crate) struct EmbeddedSelection { /// Keeps the service (listeners, sweep, replica sync) alive for as long as @@ -218,12 +228,7 @@ impl EmbeddedSelection { cancellation_token: CancellationToken, ) -> Result<(Self, crate::kv_router::sequence::ReplicaIngress)> { let worker_type = args.worker_role.unwrap_or(WorkerType::Aggregated); - let key = RoutingPartitionId::new( - args.model_name - .clone() - .unwrap_or_else(|| DEFAULT_MODEL_NAME.to_string()), - DEFAULT_ROUTING_GROUP, - ); + let key = embedded_partition_key(args.model_name.as_deref()); // Replica sync rides the runtime event plane. Worker-origin completion // marks are consumed even when router-to-router replica sync is From d4be24436df63de962088a5a008f4e498130a02d Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Sat, 12 Sep 2026 18:42:33 -0700 Subject: [PATCH 081/102] fix(router): uncount a joined lease released before its initializer commits A joiner that finished while the initializer was still queued was never subtracted from the pending count, so the committed binding over-counted and could never idle out, holding its table slot until the worker left. A release on an initializing entry now decrements the pending count. The join retry in `commit_session` is bounded at three attempts instead of one, and the module doc says when the commit still routes unpinned. Test: `joined_lease_released_before_the_commit_is_not_counted`. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Ishan Dhanani --- .../src/services/selection/affinity.rs | 50 +++++++++++++++++++ .../src/services/selection/core/run.rs | 14 ++++-- 2 files changed, 59 insertions(+), 5 deletions(-) diff --git a/lib/kv-router/src/services/selection/affinity.rs b/lib/kv-router/src/services/selection/affinity.rs index 9a1bc8006bea..786c86548785 100644 --- a/lib/kv-router/src/services/selection/affinity.rs +++ b/lib/kv-router/src/services/selection/affinity.rs @@ -870,6 +870,22 @@ impl AffinityLease { let Some(mut entry) = inner.entries.get_mut(&self.session_id) else { return; }; + // A joined lease released before its initializer commits must not + // be counted by that commit, or the binding never idles out. + if let AffinityEntry::Initializing { + revision, + pending_leases, + .. + } = entry.value_mut() + { + if *revision == self.revision + && self.version == AffinityVersion::PENDING + && *pending_leases > 0 + { + *pending_leases -= 1; + } + return; + } let AffinityEntry::Bound { target, revision, @@ -1021,6 +1037,40 @@ mod tests { } } + #[tokio::test(start_paused = true)] + async fn joined_lease_released_before_the_commit_is_not_counted() { + let table = table(); + let target = AffinityTarget::new(2, Some(0)); + let init = initialize(&table); + let joined = table + .join_initializing("s") + .expect("join an initializing session"); + assert_eq!(table.lease_count("s"), Some(1)); + + // The joiner finishes while the initializer is still queued. + drop(joined); + assert_eq!( + table.lease_count("s"), + Some(0), + "early release uncounts the join" + ); + + let lease = init.commit(target).expect("commit"); + assert_eq!( + table.lease_count("s"), + Some(1), + "commit counts only the initializer" + ); + drop(lease); + assert_eq!(table.lease_count("s"), Some(0)); + tokio::time::advance(TTL + Duration::from_secs(1)).await; + assert_eq!( + table.query_target("s", None).expect("query"), + None, + "a binding with no live leases idles out" + ); + } + #[tokio::test(start_paused = true)] async fn joined_lease_counts_once_the_competing_initialization_commits() { let table = table(); diff --git a/lib/kv-router/src/services/selection/core/run.rs b/lib/kv-router/src/services/selection/core/run.rs index 7a3859d7dea9..ded4432a02da 100644 --- a/lib/kv-router/src/services/selection/core/run.rs +++ b/lib/kv-router/src/services/selection/core/run.rs @@ -26,9 +26,12 @@ //! `try_acquire`. A commit that finds another request initializing the //! session joins that initialization instead of waiting: it takes a lease //! that counts once the initializer commits -//! (`failover_commit_behind_an_initializing_hold_keeps_a_lease`). When the -//! table is full it returns `Ok(None)` and the request routes unpinned -//! rather than waiting on a request queued behind this booking. A rejected +//! (`failover_commit_behind_an_initializing_hold_keeps_a_lease`); a joiner +//! that finishes first is uncounted again +//! (`joined_lease_released_before_the_commit_is_not_counted`). When the +//! table is full, or the join misses three times, it returns `Ok(None)` and +//! the request routes unpinned rather than waiting on a request queued +//! behind this booking. A rejected //! commit frees the booking //! (`two_phase_replay_rejects_a_worker_the_session_left`). //! - A `Lease` admission installs no index row and records no routing hashes @@ -623,8 +626,9 @@ impl SelectionCore { // already be initializing the replacement: join it rather // than wait, so this booking holds a lease on the new // binding. The join misses only if that initialization - // resolved between the two calls, so retry once. - for _ in 0..2 { + // resolved between the two calls, so retry a few times before + // routing unpinned. + for _ in 0..3 { match table.try_acquire(session_id, None) { Ok(AcquireStep::Held(hold)) => { return table From 02165d72b1942bdab070bbeddd29514ae37a1640 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Sat, 12 Sep 2026 18:42:33 -0700 Subject: [PATCH 082/102] feat(router): instrument the active-sequence ingress funnel Peer and worker active-sequence events for a partition pass through one bounded channel and one apply task. A review asked for per-publisher direct apply; before restructuring, expose what would show the funnel as a bottleneck: - `router_active_sequence_ingress_queue_depth` (gauge): backlog left after each drain batch. - `router_active_sequence_ingress_events_applied_total` and `router_active_sequence_ingress_drain_batches_total` (counters), so events per batch is derivable. The existing dropped-events counter already covers channel overflow. The funnel is a bottleneck only if depth trends upward under steady load or drops move; a depth near zero with small batches means the apply task is idle between batches and direct apply would gain nothing. Counters update once per batch, not per event. Test: `host_ingress_observer_counts_applied_events_and_drains_to_zero`. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Ishan Dhanani --- lib/kv-router/src/sequences/multi_worker.rs | 5 + lib/kv-router/src/sequences/replica_sync.rs | 1 + .../src/services/common/replica_sync.rs | 112 +++++++++++++++++- .../src/services/selection/core/tests.rs | 4 + lib/kv-router/src/services/selection/mod.rs | 3 +- lib/llm/src/kv_router/embedded.rs | 9 +- lib/llm/src/kv_router/metrics.rs | 89 ++++++++++++++ lib/llm/src/kv_router/sequence.rs | 1 + 8 files changed, 219 insertions(+), 5 deletions(-) diff --git a/lib/kv-router/src/sequences/multi_worker.rs b/lib/kv-router/src/sequences/multi_worker.rs index c46f1c829cb7..3806888e2f9b 100644 --- a/lib/kv-router/src/sequences/multi_worker.rs +++ b/lib/kv-router/src/sequences/multi_worker.rs @@ -265,6 +265,11 @@ pub trait SequenceSubscriber: Send { ) -> Poll>> { Poll::Pending } + + /// Called once after each drain batch is applied and flushed, with the + /// number of events in the batch. Implementations may sample their own + /// backlog here; the default records nothing. + fn record_drain(&mut self, _applied: usize) {} } // --------------------------------------------------------------------------- diff --git a/lib/kv-router/src/sequences/replica_sync.rs b/lib/kv-router/src/sequences/replica_sync.rs index d56a3fbb1a9c..4aa12be5517a 100644 --- a/lib/kv-router/src/sequences/replica_sync.rs +++ b/lib/kv-router/src/sequences/replica_sync.rs @@ -143,6 +143,7 @@ impl ActiveSequencesMultiWorker

{ } self.flush_replica_batch_effects(&mut effects); + subscriber.record_drain(batch_events); if exit_after_flush { break; diff --git a/lib/kv-router/src/services/common/replica_sync.rs b/lib/kv-router/src/services/common/replica_sync.rs index afe17399d4c9..d3acf8e3ed97 100644 --- a/lib/kv-router/src/services/common/replica_sync.rs +++ b/lib/kv-router/src/services/common/replica_sync.rs @@ -218,6 +218,14 @@ pub trait SchedulerLoadSink: Send + Sync { } } +/// Observes the partition's inbound replica funnel once per drain batch, after +/// the batch is applied and flushed. `queue_depth` is the number of events +/// still waiting in the inbound channel at that moment; `applied` is the batch +/// size. Called from the apply task, so implementations must be cheap. +pub trait ReplicaIngressObserver: Send + Sync { + fn observe_drain(&self, queue_depth: usize, applied: usize); +} + /// Replica-sync plumbing an embedding host supplies for one partition when it /// carries active-sequence events over its own transport (for example the /// Dynamo runtime event plane) instead of the service's ZMQ peer mesh: @@ -232,6 +240,8 @@ pub struct HostReplicaChannels { pub inbound_rx: mpsc::Receiver, /// This replica's id; events carrying it are ignored on receipt. pub process_id: u64, + /// Optional drain-batch observer for the inbound funnel. + pub ingress_observer: Option>, } #[cfg(feature = "standalone-selection")] @@ -370,11 +380,19 @@ impl SequencePublisher for ScopedSequencePublisher { pub(crate) struct ChannelSequenceSubscriber { rx: mpsc::Receiver, + observer: Option>, } impl ChannelSequenceSubscriber { pub(crate) fn new(rx: mpsc::Receiver) -> Self { - Self { rx } + Self { rx, observer: None } + } + + pub(crate) fn with_observer( + rx: mpsc::Receiver, + observer: Option>, + ) -> Self { + Self { rx, observer } } } @@ -389,6 +407,12 @@ impl SequenceSubscriber for ChannelSequenceSubscriber { ) -> Poll>> { self.rx.poll_recv(cx).map(|event| event.map(Ok)) } + + fn record_drain(&mut self, applied: usize) { + if let Some(observer) = &self.observer { + observer.observe_drain(self.rx.len(), applied); + } + } } fn generate_process_id() -> u64 { @@ -450,7 +474,10 @@ pub(crate) fn setup_scoped_replica_sync( process_id: host.process_id, channel: Some(( host.inbound_tx, - ChannelSequenceSubscriber::new(host.inbound_rx), + ChannelSequenceSubscriber::with_observer( + host.inbound_rx, + host.ingress_observer, + ), )), }; } @@ -1079,10 +1106,91 @@ mod tests { inbound_tx, inbound_rx, process_id: 7, + ingress_observer: None, }), ); assert!(!scoped.enabled); assert_eq!(scoped.process_id, 7); assert!(scoped.channel.is_some()); } + + #[derive(Default)] + struct RecordingIngressObserver { + applied: std::sync::atomic::AtomicUsize, + batches: std::sync::atomic::AtomicUsize, + last_depth: std::sync::atomic::AtomicUsize, + } + + impl ReplicaIngressObserver for RecordingIngressObserver { + fn observe_drain(&self, queue_depth: usize, applied: usize) { + use std::sync::atomic::Ordering::Relaxed; + self.applied.fetch_add(applied, Relaxed); + self.batches.fetch_add(1, Relaxed); + self.last_depth.store(queue_depth, Relaxed); + } + } + + #[tokio::test] + async fn host_ingress_observer_counts_applied_events_and_drains_to_zero() { + use crate::sequences::{ActiveSequencesMultiWorker, NoopSequencePublisher}; + use std::sync::atomic::Ordering::Relaxed; + + const N: usize = 1_000; + let observer = Arc::new(RecordingIngressObserver::default()); + let (inbound_tx, inbound_rx) = mpsc::channel(REPLICA_EVENT_CHANNEL_CAPACITY); + let scoped = setup_scoped_replica_sync( + None, + &RoutingPartitionId::new("model", "default"), + 16, + Some(HostReplicaChannels { + outbound: None, + inbound_tx, + inbound_rx, + process_id: 7, + ingress_observer: Some(Arc::clone(&observer) as Arc), + }), + ); + let (inbound_tx, subscriber) = scoped.channel.expect("host channel"); + let tracker = Arc::new(ActiveSequencesMultiWorker::new_without_expiry( + NoopSequencePublisher, + 16, + std::collections::HashMap::from([(1_u64, (0_u32, 1_u32))]), + true, + scoped.process_id, + "test", + )); + let cancel_token = CancellationToken::new(); + tracker.start_replica_sync(subscriber, cancel_token.clone()); + + for index in 0..N { + inbound_tx + .try_send(ActiveSequenceEvent { + request_id: format!("req-{index}"), + worker: WorkerWithDpRank::new(1, 0), + data: ActiveSequenceEventData::AddRequest { + token_sequence: Some(vec![index as u64]), + track_prefill_tokens: false, + expected_output_tokens: None, + prefill_load_hint: None, + }, + router_id: 99, + lora_name: None, + }) + .expect("inbound channel has capacity"); + } + + tokio::time::timeout(std::time::Duration::from_secs(5), async { + while observer.applied.load(Relaxed) < N { + tokio::task::yield_now().await; + } + }) + .await + .expect("all events applied"); + + assert_eq!(observer.applied.load(Relaxed), N); + assert!(observer.batches.load(Relaxed) >= 1); + assert!(observer.batches.load(Relaxed) <= N); + assert_eq!(observer.last_depth.load(Relaxed), 0); + cancel_token.cancel(); + } } diff --git a/lib/kv-router/src/services/selection/core/tests.rs b/lib/kv-router/src/services/selection/core/tests.rs index 2e196ce145de..d2da09d9baa4 100644 --- a/lib/kv-router/src/services/selection/core/tests.rs +++ b/lib/kv-router/src/services/selection/core/tests.rs @@ -1656,6 +1656,7 @@ async fn mirrored_replica_bookings_are_indexed_until_freed() { inbound_tx: inbound_tx.clone(), inbound_rx, process_id: 7, + ingress_observer: None, })); let core = core_with_host(SelectionHost { replication: HostReplication { @@ -2180,6 +2181,9 @@ fn lease_count(core: &SelectionCore, session_id: &str) -> Option { table.lease_count(session_id) } +/// Two workers, so both requests must land on the replacement; with more +/// workers the joiner could book elsewhere while still counting toward the +/// session's binding, which is the same split a plain `Bound` lease has. /// r2 holds the departed binding when r3 re-initializes the session; r2's /// commit then finds r3's initialization in progress and must join it, not /// route without a lease, so the replacement binding cannot idle out while diff --git a/lib/kv-router/src/services/selection/mod.rs b/lib/kv-router/src/services/selection/mod.rs index 909a08e245a8..d10a4bb91313 100644 --- a/lib/kv-router/src/services/selection/mod.rs +++ b/lib/kv-router/src/services/selection/mod.rs @@ -24,7 +24,8 @@ mod tests; pub use crate::WorkerSelectionPolicyFactory; pub use crate::services::common::replica_sync::{ - HostReplicaChannels, HostReplicaSyncFactory, ReplicaPeerError, SchedulerLoadSink, + HostReplicaChannels, HostReplicaSyncFactory, ReplicaIngressObserver, ReplicaPeerError, + SchedulerLoadSink, }; pub use core::{ HostCache, HostEligibility, HostLoad, HostReplication, HostTelemetry, KvIndexSource, Selected, diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs index 6750ace1a8fa..bc7aab5a0565 100644 --- a/lib/llm/src/kv_router/embedded.rs +++ b/lib/llm/src/kv_router/embedded.rs @@ -34,7 +34,8 @@ use tokio_util::sync::CancellationToken; use crate::discovery::RuntimeConfigWatch; use crate::kv_router::metrics::{ - ROUTER_QUEUE_METRICS, RouterQueueMetricHandles, RouterRequestMetrics, WORKER_LOAD_METRICS, + ActiveSequenceIngressMetrics, ROUTER_QUEUE_METRICS, RouterQueueMetricHandles, + RouterRequestMetrics, WORKER_LOAD_METRICS, }; use crate::local_model::runtime_config::ModelRuntimeConfig; @@ -233,7 +234,7 @@ impl EmbeddedSelection { // Replica sync rides the runtime event plane. Worker-origin completion // marks are consumed even when router-to-router replica sync is // disabled; only publishing is gated. - let (channels, replica_ingress) = crate::kv_router::sequence::host_replica_channels( + let (mut channels, replica_ingress) = crate::kv_router::sequence::host_replica_channels( &args.endpoint, args.router_id, args.kv_router_config.router_replica_sync, @@ -241,6 +242,10 @@ impl EmbeddedSelection { ) .await .context("start replica sync for the embedded selection partition")?; + channels.ingress_observer = Some(Arc::new( + ActiveSequenceIngressMetrics::from_component(args.endpoint.component()) + .handles(&key.model_name, &key.routing_group), + )); let slot = std::sync::Mutex::new(Some(channels)); let replica_sync: Option = Some(Arc::new(move |_partition| { diff --git a/lib/llm/src/kv_router/metrics.rs b/lib/llm/src/kv_router/metrics.rs index 696636ef2cd5..f650b0710eb5 100644 --- a/lib/llm/src/kv_router/metrics.rs +++ b/lib/llm/src/kv_router/metrics.rs @@ -406,6 +406,95 @@ impl ActiveSequenceZmqIngressMetrics { } } +// --------------------------------------------------------------------------- +// Active-sequence inbound funnel metrics (per partition) +// --------------------------------------------------------------------------- + +const ROUTING_GROUP_LABEL: &str = "routing_group"; + +/// Instrumentation for the per-partition inbound replica funnel: one channel +/// and one apply task per partition, fed by every publisher's ingress source. +/// Updated once per drain batch from the apply task. +pub(crate) struct ActiveSequenceIngressMetrics { + queue_depth: IntGaugeVec, + events_applied_total: IntCounterVec, + drain_batches_total: IntCounterVec, +} + +/// Per-partition handles implementing the core's drain observer. +pub(crate) struct ActiveSequenceIngressMetricHandles { + queue_depth: IntGauge, + events_applied_total: IntCounter, + drain_batches_total: IntCounter, +} + +static ACTIVE_SEQUENCE_INGRESS_METRICS: OnceLock> = + OnceLock::new(); + +impl ActiveSequenceIngressMetrics { + pub(crate) fn from_component(component: &Component) -> Arc { + ACTIVE_SEQUENCE_INGRESS_METRICS + .get_or_init(|| { + let metrics = component.metrics(); + let labels = [labels::MODEL, ROUTING_GROUP_LABEL]; + Arc::new(Self { + queue_depth: metrics + .create_intgaugevec( + "router_active_sequence_ingress_queue_depth", + "Active-sequence events waiting in the partition's inbound channel, sampled after each drain batch", + &labels, + &[], + ) + .expect("failed to create router_active_sequence_ingress_queue_depth"), + events_applied_total: metrics + .create_intcountervec( + "router_active_sequence_ingress_events_applied_total", + "Active-sequence events applied from the partition's inbound channel", + &labels, + &[], + ) + .expect( + "failed to create router_active_sequence_ingress_events_applied_total", + ), + drain_batches_total: metrics + .create_intcountervec( + "router_active_sequence_ingress_drain_batches_total", + "Drain batches applied from the partition's inbound channel", + &labels, + &[], + ) + .expect( + "failed to create router_active_sequence_ingress_drain_batches_total", + ), + }) + }) + .clone() + } + + pub(crate) fn handles( + &self, + model: &str, + routing_group: &str, + ) -> ActiveSequenceIngressMetricHandles { + let labels = [model, routing_group]; + ActiveSequenceIngressMetricHandles { + queue_depth: self.queue_depth.with_label_values(&labels), + events_applied_total: self.events_applied_total.with_label_values(&labels), + drain_batches_total: self.drain_batches_total.with_label_values(&labels), + } + } +} + +impl dynamo_kv_router::services::selection::ReplicaIngressObserver + for ActiveSequenceIngressMetricHandles +{ + fn observe_drain(&self, queue_depth: usize, applied: usize) { + self.queue_depth.set(queue_depth as i64); + self.events_applied_total.inc_by(applied as u64); + self.drain_batches_total.inc(); + } +} + // --------------------------------------------------------------------------- // Router worker status metrics (component-scoped gauges) // --------------------------------------------------------------------------- diff --git a/lib/llm/src/kv_router/sequence.rs b/lib/llm/src/kv_router/sequence.rs index 024be03aa84a..51bd1de751a3 100644 --- a/lib/llm/src/kv_router/sequence.rs +++ b/lib/llm/src/kv_router/sequence.rs @@ -402,6 +402,7 @@ pub(crate) async fn host_replica_channels( inbound_tx, inbound_rx, process_id: router_id, + ingress_observer: None, }, ingress, )) From 27c03ee2a4d9c3c4cf9881a933624323f1496bab Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Sat, 12 Sep 2026 18:56:23 -0700 Subject: [PATCH 083/102] fix(router): label ingress metrics by worker type and pin the queued backlog Prefill and decode routers of one model run separate ingress channels but would have shared the same metric handles; the worker type is now a label. The ingress observer test asserts the first drain sample shows the queued backlog, not only that depth ends at zero. The join retry in `commit_session` is a named constant of two attempts: one retry covers an initialization that resolved between the two calls, and a booking never waits on a request queued behind it. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Ishan Dhanani --- lib/kv-router/src/services/common/replica_sync.rs | 9 +++++++++ lib/kv-router/src/services/selection/core/run.rs | 12 ++++++++---- lib/llm/src/kv_router/embedded.rs | 7 +++++-- lib/llm/src/kv_router/metrics.rs | 8 ++++++-- 4 files changed, 28 insertions(+), 8 deletions(-) diff --git a/lib/kv-router/src/services/common/replica_sync.rs b/lib/kv-router/src/services/common/replica_sync.rs index d3acf8e3ed97..5ca2dd2d0157 100644 --- a/lib/kv-router/src/services/common/replica_sync.rs +++ b/lib/kv-router/src/services/common/replica_sync.rs @@ -1119,6 +1119,7 @@ mod tests { applied: std::sync::atomic::AtomicUsize, batches: std::sync::atomic::AtomicUsize, last_depth: std::sync::atomic::AtomicUsize, + max_depth: std::sync::atomic::AtomicUsize, } impl ReplicaIngressObserver for RecordingIngressObserver { @@ -1127,6 +1128,7 @@ mod tests { self.applied.fetch_add(applied, Relaxed); self.batches.fetch_add(1, Relaxed); self.last_depth.store(queue_depth, Relaxed); + self.max_depth.fetch_max(queue_depth, Relaxed); } } @@ -1191,6 +1193,13 @@ mod tests { assert!(observer.batches.load(Relaxed) >= 1); assert!(observer.batches.load(Relaxed) <= N); assert_eq!(observer.last_depth.load(Relaxed), 0); + // All N events were queued before the first drain, whose batch is + // capped, so the first sample must show the backlog. + assert!( + observer.max_depth.load(Relaxed) >= N - crate::protocols::MAX_REPLICA_BATCH_EVENTS, + "depth sample reflects the queued backlog: {}", + observer.max_depth.load(Relaxed) + ); cancel_token.cancel(); } } diff --git a/lib/kv-router/src/services/selection/core/run.rs b/lib/kv-router/src/services/selection/core/run.rs index ded4432a02da..8517ced5c673 100644 --- a/lib/kv-router/src/services/selection/core/run.rs +++ b/lib/kv-router/src/services/selection/core/run.rs @@ -29,7 +29,7 @@ //! (`failover_commit_behind_an_initializing_hold_keeps_a_lease`); a joiner //! that finishes first is uncounted again //! (`joined_lease_released_before_the_commit_is_not_counted`). When the -//! table is full, or the join misses three times, it returns `Ok(None)` and +//! table is full, or the join misses twice, it returns `Ok(None)` and //! the request routes unpinned rather than waiting on a request queued //! behind this booking. A rejected //! commit frees the booking @@ -43,6 +43,10 @@ //! `dropped_book_selection_during_routing_record_frees_booking_and_claim`). use super::hint::{hint_capable_partition, transfer_hint_for_selection}; + +/// `try_acquire` then `join_initializing` attempts before a commit that finds +/// the session initializing routes unpinned (see the module doc). +const JOIN_ATTEMPTS: usize = 2; use super::reservations::Reservation; use super::*; @@ -626,9 +630,9 @@ impl SelectionCore { // already be initializing the replacement: join it rather // than wait, so this booking holds a lease on the new // binding. The join misses only if that initialization - // resolved between the two calls, so retry a few times before - // routing unpinned. - for _ in 0..3 { + // resolved between the two calls; one retry covers that, and + // a booking never waits on a request queued behind it. + for _ in 0..JOIN_ATTEMPTS { match table.try_acquire(session_id, None) { Ok(AcquireStep::Held(hold)) => { return table diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs index bc7aab5a0565..95957656bd53 100644 --- a/lib/llm/src/kv_router/embedded.rs +++ b/lib/llm/src/kv_router/embedded.rs @@ -243,8 +243,11 @@ impl EmbeddedSelection { .await .context("start replica sync for the embedded selection partition")?; channels.ingress_observer = Some(Arc::new( - ActiveSequenceIngressMetrics::from_component(args.endpoint.component()) - .handles(&key.model_name, &key.routing_group), + ActiveSequenceIngressMetrics::from_component(args.endpoint.component()).handles( + &key.model_name, + &key.routing_group, + args.metric_worker_type, + ), )); let slot = std::sync::Mutex::new(Some(channels)); let replica_sync: Option = diff --git a/lib/llm/src/kv_router/metrics.rs b/lib/llm/src/kv_router/metrics.rs index f650b0710eb5..2de92fb220bd 100644 --- a/lib/llm/src/kv_router/metrics.rs +++ b/lib/llm/src/kv_router/metrics.rs @@ -436,7 +436,7 @@ impl ActiveSequenceIngressMetrics { ACTIVE_SEQUENCE_INGRESS_METRICS .get_or_init(|| { let metrics = component.metrics(); - let labels = [labels::MODEL, ROUTING_GROUP_LABEL]; + let labels = [labels::MODEL, ROUTING_GROUP_LABEL, labels::WORKER_TYPE]; Arc::new(Self { queue_depth: metrics .create_intgaugevec( @@ -471,12 +471,16 @@ impl ActiveSequenceIngressMetrics { .clone() } + /// One handle set per router partition. `worker_type` separates the + /// prefill and decode routers of the same model, which run separate + /// ingress channels. pub(crate) fn handles( &self, model: &str, routing_group: &str, + worker_type: &str, ) -> ActiveSequenceIngressMetricHandles { - let labels = [model, routing_group]; + let labels = [model, routing_group, worker_type]; ActiveSequenceIngressMetricHandles { queue_depth: self.queue_depth.with_label_values(&labels), events_applied_total: self.events_applied_total.with_label_values(&labels), From 461fff6537d05b63ef10c7791811d081b691fb11 Mon Sep 17 00:00:00 2001 From: Thomas Montfort Date: Sun, 13 Sep 2026 14:04:13 -0700 Subject: [PATCH 084/102] fix(epp): address EPP-specific bugs and config layering issues in selection core MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fix several correctness and layering issues in the EPP standalone path: - Add `session_affinity_ttl_secs` + `session_affinity_mode` to `KvRouterConfig`, read from `DYN_ROUTER_SESSION_AFFINITY_TTL_SECS` / `DYN_ROUTER_SESSION_AFFINITY_MODE`. `SelectionServiceBuilder::build()` auto-consumes these, removing the need for callers to call `session_affinity()` explicitly. - Remove `data_parallel_size`, `kv_event_port_stride`, and `session_affinity_ttl_secs` from `EppStandaloneConfig`. DP > 1 is not yet supported (requires per-rank HTTP inference ports too); the stride was an engine implementation detail. Session affinity config belongs in `DYN_ROUTER_*` env vars, not `DYN_EPP_*`. - Remove `KvEventPorts` struct from `pod_discovery.rs`; always produce a single rank-0 KV-event endpoint `{0: "tcp://ip:kv_event_port"}`. - Make `Selector.service` private; add `pub(crate) fn catalog_reconciler()` to eliminate the `selector.service.core()` layer violation in `topology_adapter.rs` and tests. - Expand session-affinity header detection in `epp_router.rs` to cover Claude Code (`x-claude-code-agent-id` / `x-claude-code-session-id`), Codex (`thread-id`), and OpenCode (`x-session-id`) in addition to `x-dynamo-session-id`. - Map `SelectionError::Conflict` to `PickError::Internal` → HTTP 500. Conflict fires only when the EPP's own UUID collides in the selection index — always an EPP bug, never a client error. - Add `PickError::Internal(String)` variant → 500 for EPP-internal invariant violations. - Move `SessionAffinityMode` from `services::selection::affinity` to `scheduling::config` (re-exported from `affinity`) so `KvRouterConfig` can hold it without a feature-gated import cycle. - Make `session_affinity_header_value` in `dynamo_llm::protocols::agents` `pub`. Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Thomas Montfort --- .../ext-proc/src/epp_router.rs | 21 +++- .../ext-proc/src/epp_standalone_config.rs | 39 ------- .../inference-gateway/ext-proc/src/picker.rs | 4 + .../ext-proc/src/pod_discovery.rs | 106 ++++-------------- .../ext-proc/src/selector.rs | 57 +++++----- .../inference-gateway/ext-proc/src/server.rs | 6 + .../ext-proc/src/topology_adapter.rs | 12 +- lib/kv-router/src/scheduling/config.rs | 61 ++++++++++ .../src/services/selection/affinity.rs | 27 +---- .../src/services/selection/service.rs | 19 +++- lib/llm/src/protocols/agents.rs | 2 +- 11 files changed, 166 insertions(+), 188 deletions(-) diff --git a/deploy/inference-gateway/ext-proc/src/epp_router.rs b/deploy/inference-gateway/ext-proc/src/epp_router.rs index 2217cb420d64..cbfc048449e9 100644 --- a/deploy/inference-gateway/ext-proc/src/epp_router.rs +++ b/deploy/inference-gateway/ext-proc/src/epp_router.rs @@ -28,7 +28,6 @@ use tokio::sync::Semaphore; use dynamo_kv_router::services::selection::{SelectionError, WorkerSelectionPolicyRegistry}; use dynamo_llm::http::service::metadata::extract_metadata_from_header_pairs; -use dynamo_llm::protocols::agents::HEADER_DYNAMO_SESSION_ID; use dynamo_llm::protocols::common::extensions::{ AgentHints, HEADER_REQUEST_PRIORITY, HEADER_REQUEST_STRICT_PRIORITY, resolve_request_priority, }; @@ -226,7 +225,6 @@ struct RoutingNvExt { agent_hints: Option, } -/// Case-insensitive lookup of the first non-empty, trimmed value for `name`. fn first_header<'a>(headers: &'a [(String, String)], name: &str) -> Option<&'a str> { headers .iter() @@ -235,6 +233,17 @@ fn first_header<'a>(headers: &'a [(String, String)], name: &str) -> Option<&'a s .filter(|v| !v.is_empty()) } +fn session_id_from_headers(headers: &[(String, String)]) -> Option { + use axum::http::{HeaderMap, HeaderName, HeaderValue}; + let mut map = HeaderMap::with_capacity(headers.len()); + for (k, v) in headers { + if let (Ok(name), Ok(val)) = (k.parse::(), v.parse::()) { + map.insert(name, val); + } + } + dynamo_llm::protocols::agents::session_affinity_header_value(&map) +} + #[tonic::async_trait] impl EndpointPicker for EppRouter { async fn pick( @@ -342,7 +351,7 @@ impl EndpointPicker for EppRouter { model_name: self.model_name.clone(), reservation_id: reservation_id.clone(), token_ids: tokens, - session_id: first_header(&req.headers, HEADER_DYNAMO_SESSION_ID).map(str::to_owned), + session_id: session_id_from_headers(&req.headers), // `None` on the ordinary path: the selector schedules over its // catalog; `Some` only carries an Envoy subset constraint. allowed_worker_ids: allowed, @@ -360,6 +369,12 @@ impl EndpointPicker for EppRouter { Err(SelectionError::BadRequest(message)) => { return Err(PickError::InvalidRequest(message)); } + // A Conflict arises only when the EPP's own UUID collides in the selection + // index — this is an EPP-internal invariant violation, never caused by client + // input, so 500 is correct (not 400 or 503). + Err(SelectionError::Conflict(msg)) => { + return Err(PickError::Internal(format!("selection id conflict: {msg}"))); + } Err(e) => return Err(PickError::RoutingFailed(e.to_string())), }; diff --git a/deploy/inference-gateway/ext-proc/src/epp_standalone_config.rs b/deploy/inference-gateway/ext-proc/src/epp_standalone_config.rs index 93b7120fb5db..cd2c48cc2402 100644 --- a/deploy/inference-gateway/ext-proc/src/epp_standalone_config.rs +++ b/deploy/inference-gateway/ext-proc/src/epp_standalone_config.rs @@ -140,17 +140,6 @@ pub struct EppStandaloneConfig { /// KV-cache block size; MUST equal the inference engine block size. #[validate(range(min = 1, message = "DYN_KV_CACHE_BLOCK_SIZE must be >= 1"))] pub block_size: u32, - /// Data-parallel ranks per worker pod. Only `1` is accepted: the EPP routes - /// by worker and cannot convey the selected rank to the pod. - #[validate(range( - min = 1, - max = 1, - message = "DYN_EPP_DATA_PARALLEL_SIZE must be 1; multi-rank standalone serving is unsupported" - ))] - pub data_parallel_size: u32, - /// Port distance between consecutive data-parallel ranks' KV event ports. - #[validate(range(min = 1, message = "DYN_EPP_KV_EVENT_PORT_STRIDE must be >= 1"))] - pub kv_event_port_stride: u16, /// KV zmq event port. #[validate(range(min = 1))] pub kv_event_port: u16, @@ -174,15 +163,6 @@ pub struct EppStandaloneConfig { /// throughput throttle). Excess requests are shed with a 503, not queued. #[validate(range(min = 1, message = "DYN_EPP_MAX_INFLIGHT_REQUESTS must be >= 1"))] pub max_inflight_requests: usize, - /// Pin each `x-dynamo-session-id` to the worker that served it for this - /// long after its last request (`DYN_EPP_SESSION_AFFINITY_TTL_SECS`). - /// `None` disables session affinity. - #[validate(range( - min = 1.0, - max = 31536000.0, - message = "DYN_EPP_SESSION_AFFINITY_TTL_SECS must be between 1 and 31536000 seconds" - ))] - pub session_affinity_ttl_secs: Option, } impl EppStandaloneConfig { @@ -236,9 +216,6 @@ impl EppStandaloneConfig { )? .unwrap_or(DEFAULT_TOKENIZER_MAX_RESPONSE_BYTES), block_size: opt_parse::(get, "DYN_KV_CACHE_BLOCK_SIZE")?.unwrap_or(0), - data_parallel_size: opt_parse::(get, "DYN_EPP_DATA_PARALLEL_SIZE")?.unwrap_or(1), - kv_event_port_stride: opt_parse::(get, "DYN_EPP_KV_EVENT_PORT_STRIDE")? - .unwrap_or(1), kv_event_port: opt_parse::(get, "DYN_EPP_KV_EVENT_PORT")? .unwrap_or(DEFAULT_KV_EVENT_PORT), replay_port: opt_parse::(get, "DYN_EPP_KV_EVENT_REPLAY_PORT")?, @@ -246,7 +223,6 @@ impl EppStandaloneConfig { max_num_batched_tokens: opt_parse::(get, "DYN_EPP_MAX_NUM_BATCHED_TOKENS")?, max_inflight_requests: opt_parse::(get, "DYN_EPP_MAX_INFLIGHT_REQUESTS")? .unwrap_or(DEFAULT_MAX_INFLIGHT_REQUESTS), - session_affinity_ttl_secs: opt_parse::(get, "DYN_EPP_SESSION_AFFINITY_TTL_SECS")?, }) } @@ -566,21 +542,6 @@ mod tests { ); } - #[test] - fn data_parallel_size_above_one_fails() { - assert!( - parse_cfg(&[ - ("DYN_EPP_INFERENCE_POOL_NAME", "vllm-qwen-pool"), - ("POD_NAMESPACE", "inference"), - ("DYN_MODEL_NAME", "Qwen/Qwen3-0.6B"), - ("DYN_EPP_TOKENIZER_SERVICE_URL", "http://vllm-render:8000"), - ("DYN_EPP_TOKENIZER_PROTOCOL", "vllm-render"), - ("DYN_EPP_DATA_PARALLEL_SIZE", "2"), - ]) - .is_err() - ); - } - #[test] fn replay_port_is_optional_and_uses_the_kv_event_name() { let cfg = parse_cfg(&[ diff --git a/deploy/inference-gateway/ext-proc/src/picker.rs b/deploy/inference-gateway/ext-proc/src/picker.rs index be638de20417..b66a96c686bc 100644 --- a/deploy/inference-gateway/ext-proc/src/picker.rs +++ b/deploy/inference-gateway/ext-proc/src/picker.rs @@ -177,4 +177,8 @@ pub enum PickError { /// multiplexing means the connection cap does not bound concurrent requests. #[error("endpoint picker overloaded")] Overloaded, + /// An internal EPP invariant was violated; always an EPP bug, never a client + /// error → 500. + #[error("internal error: {0}")] + Internal(String), } diff --git a/deploy/inference-gateway/ext-proc/src/pod_discovery.rs b/deploy/inference-gateway/ext-proc/src/pod_discovery.rs index bdd652c1a417..83bf734277cd 100644 --- a/deploy/inference-gateway/ext-proc/src/pod_discovery.rs +++ b/deploy/inference-gateway/ext-proc/src/pod_discovery.rs @@ -103,11 +103,6 @@ impl PodDiscovery { let kv_event_port = cfg.kv_event_port; let replay_port = cfg.replay_port; - let kv_ports = KvEventPorts { - base: kv_event_port, - stride: cfg.kv_event_port_stride, - data_parallel_size: cfg.data_parallel_size, - }; let index: Arc> = Arc::new(RwLock::new(WorkerIndex::new())); @@ -184,7 +179,7 @@ impl PodDiscovery { Delta::Rebuild => rebuild_index( &store, pool_rx.borrow().as_ref(), - kv_ports, + kv_event_port, replay_port, &index_task, ), @@ -192,7 +187,7 @@ impl PodDiscovery { &index_task, &pod, pool_rx.borrow().as_ref(), - kv_ports, + kv_event_port, replay_port, ), Delta::Remove(pod) => remove_pod(&index_task, &pod), @@ -345,14 +340,14 @@ fn upsert_pod( index: &RwLock, pod: &Pod, pool: Option<&PoolState>, - kv_ports: KvEventPorts, + kv_event_port: u16, replay_port: Option, ) -> bool { let Some(worker_id) = pod_worker_id(pod) else { return false; }; let entry = pool - .and_then(|pool| raw_worker_from_pod(pod, pool, kv_ports, replay_port)) + .and_then(|pool| raw_worker_from_pod(pod, pool, kv_event_port, replay_port)) .map(WorkerEntry::from_raw); let mut index = index.write().unwrap(); match entry { @@ -383,14 +378,14 @@ fn remove_pod(index: &RwLock, pod: &Pod) -> bool { fn rebuild_index( store: &kube::runtime::reflector::Store, pool: Option<&PoolState>, - kv_ports: KvEventPorts, + kv_event_port: u16, replay_port: Option, index: &RwLock, ) -> bool { let mut fresh = WorkerIndex::new(); if let Some(pool) = pool { for pod in store.state().iter() { - if let Some(worker) = raw_worker_from_pod(pod, pool, kv_ports, replay_port) { + if let Some(worker) = raw_worker_from_pod(pod, pool, kv_event_port, replay_port) { fresh.insert(worker.worker_id, WorkerEntry::from_raw(worker)); } } @@ -404,33 +399,12 @@ fn rebuild_index( } } -/// Per-pod KV-event port layout: rank `r` publishes on `base + r * stride`. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(crate) struct KvEventPorts { - pub base: u16, - pub stride: u16, - pub data_parallel_size: u32, -} - -impl KvEventPorts { - fn endpoints(self, ip: IpAddr) -> HashMap { - (0..self.data_parallel_size.max(1)) - .map(|rank| { - let port = self - .base - .saturating_add(self.stride.saturating_mul(rank as u16)); - (rank, format!("tcp://{}", SocketAddr::new(ip, port))) - }) - .collect() - } -} - /// Build a [`RawWorker`] from a pod, or `None` if it is not `Ready`, not /// pool-selected, or lacks an IP/name. Pure function — unit-testable. fn raw_worker_from_pod( pod: &Pod, pool: &PoolState, - kv_ports: KvEventPorts, + kv_event_port: u16, replay_port: Option, ) -> Option { if !pod_is_ready(pod) || !pod_matches(pod, &pool.match_labels) { @@ -445,7 +419,10 @@ fn raw_worker_from_pod( pod_name: pod_name.to_string(), pod_ip: pod_ip.to_string(), http_endpoint: format!("http://{}", SocketAddr::new(ip, pool.target_port)), - kv_events_endpoints: kv_ports.endpoints(ip), + kv_events_endpoints: HashMap::from([( + 0u32, + format!("tcp://{}", SocketAddr::new(ip, kv_event_port)), + )]), replay_endpoint: replay_port.map(|p| format!("tcp://{}", SocketAddr::new(ip, p))), }) } @@ -457,41 +434,6 @@ mod tests { use k8s_openapi::apimachinery::pkg::apis::meta::v1::Time; use kube::api::ObjectMeta; - fn single_rank(port: u16) -> KvEventPorts { - KvEventPorts { - base: port, - stride: 1, - data_parallel_size: 1, - } - } - - #[test] - fn data_parallel_ranks_publish_on_strided_ports() { - let w = raw_worker_from_pod( - &pod( - "vllm-0", - Some("10.0.0.1"), - Some(true), - &[("app", "vllm-qwen")], - ), - &pool(), - KvEventPorts { - base: 5557, - stride: 2, - data_parallel_size: 3, - }, - None, - ) - .expect("ready, selected pod should map"); - assert_eq!( - w.kv_events_endpoints, - HashMap::from([ - (0, "tcp://10.0.0.1:5557".to_string()), - (1, "tcp://10.0.0.1:5559".to_string()), - (2, "tcp://10.0.0.1:5561".to_string()), - ]) - ); - } fn pool() -> PoolState { PoolState { @@ -537,7 +479,7 @@ mod tests { &[("app", "vllm-qwen")], ), &pool(), - single_rank(5557), + 5557u16, Some(5560), ) .expect("ready, selected pod should map"); @@ -560,7 +502,7 @@ mod tests { &[("app", "vllm-qwen")], ), &pool(), - single_rank(5557), + 5557u16, Some(5560), ) .expect("ready, selected IPv6 pod should map"); @@ -584,7 +526,7 @@ mod tests { &[("app", "vllm-qwen")] ), &pool(), - single_rank(5557), + 5557u16, None, ) .is_none() @@ -602,7 +544,7 @@ mod tests { &[("app", "something-else")] ), &pool(), - single_rank(5557), + 5557u16, None, ) .is_none() @@ -620,7 +562,7 @@ mod tests { &[("app", "vllm-qwen")] ), &pool(), - single_rank(5557), + 5557u16, None, ) .is_none() @@ -636,7 +578,7 @@ mod tests { &[("app", "vllm-qwen")], ); p.metadata.deletion_timestamp = Some(Time(k8s_openapi::chrono::Utc::now())); - assert!(raw_worker_from_pod(&p, &pool(), single_rank(5557), None).is_none()); + assert!(raw_worker_from_pod(&p, &pool(), 5557u16, None).is_none()); } #[test] @@ -645,7 +587,7 @@ mod tests { raw_worker_from_pod( &pod("vllm-0", None, Some(true), &[("app", "vllm-qwen")]), &pool(), - single_rank(5557), + 5557u16, None, ) .is_none() @@ -666,7 +608,7 @@ mod tests { #[test] fn rebuild_index_keeps_only_ready_selected_pods() { - let kp = single_rank(5557); + let kp = 5557u16; let store = store_from_pods(vec![ pod( "vllm-0", @@ -705,7 +647,7 @@ mod tests { #[test] fn rebuild_index_is_empty_without_pool() { - let kp = single_rank(5557); + let kp = 5557u16; let store = store_from_pods(vec![pod( "vllm-0", Some("10.0.0.1"), @@ -719,7 +661,7 @@ mod tests { #[test] fn upsert_and_remove_pod_mutate_index_incrementally() { - let kp = single_rank(5557); + let kp = 5557u16; let index = RwLock::new(WorkerIndex::new()); let id = hash_pod_name("vllm-0"); let ready = pod( @@ -762,7 +704,7 @@ mod tests { #[test] fn upsert_pod_without_pool_drops_entry() { - let kp = single_rank(5557); + let kp = 5557u16; // A `None` pool (unresolved or deleted) means nothing is routable, so an // upsert must evict any existing entry rather than leave stale routing. let index = RwLock::new(WorkerIndex::new()); @@ -784,7 +726,7 @@ mod tests { #[test] fn pool_edit_during_relist_rebuilds_at_init_done_from_completed_store() { - let kp = single_rank(5557); + let kp = 5557u16; use kube::runtime::watcher; let vllm_0 = pod( @@ -851,7 +793,7 @@ mod tests { #[test] fn rebuild_index_drops_workers_absent_from_the_store() { - let kp = single_rank(5557); + let kp = 5557u16; // A relist arrives as a fresh snapshot with no `Delete` events for pods // that disappeared during the disconnect, so the `InitDone`/pool-change // rebuild must *replace* the index, not merge into it — otherwise dead diff --git a/deploy/inference-gateway/ext-proc/src/selector.rs b/deploy/inference-gateway/ext-proc/src/selector.rs index 8bdb113e53e0..943f09d9b0d7 100644 --- a/deploy/inference-gateway/ext-proc/src/selector.rs +++ b/deploy/inference-gateway/ext-proc/src/selector.rs @@ -15,8 +15,8 @@ use anyhow::{Context, Result, anyhow}; use dynamo_kv_router::config::{KvRouterConfig, try_kv_router_config_from_dynamo_env}; use dynamo_kv_router::protocols::RoutingConstraints; use dynamo_kv_router::services::selection::{ - PromptRequest, SelectAndReserveRequest as CoreSelectAndReserveRequest, SelectionError, - SelectionService, SelectionServiceBuilder, WorkerSelectionPolicyRegistry, + CatalogReconciler, PromptRequest, SelectAndReserveRequest as CoreSelectAndReserveRequest, + SelectionError, SelectionService, SelectionServiceBuilder, WorkerSelectionPolicyRegistry, warn_for_unserved_worker_selection_policies, }; use dynamo_kv_router::{DEFAULT_ROUTING_GROUP, WorkerType}; @@ -66,7 +66,7 @@ pub struct SelectResponse { /// In-process runtime-free selector wrapping a [`SelectionService`]. pub struct Selector { - pub(crate) service: Arc, + service: Arc, /// Cancels the peer-discovery watch on drop. The `SelectionService`'s own /// `Drop` tears down its core + replica-sync tasks. cancel: CancellationToken, @@ -115,12 +115,6 @@ impl Selector { if let Some(peer_replication) = peer_replication { builder = builder.replica_sync(peer_replication.sync_port, Vec::new()); } - if let Some(ttl) = cfg.session_affinity_ttl_secs { - builder = builder.session_affinity( - std::time::Duration::try_from_secs_f64(ttl) - .context("invalid session affinity TTL")?, - ); - } let service = Arc::new( builder .build() @@ -166,6 +160,10 @@ impl Selector { Ok(()) } + pub(crate) fn catalog_reconciler(&self) -> CatalogReconciler { + CatalogReconciler::new(Arc::clone(self.service.core())) + } + /// Select a worker for a prompt and book its load in one operation. Takes the /// request by value so per-request fields are moved into the core request /// rather than cloned on the hot path. @@ -350,14 +348,11 @@ models: tokenizer_max_response_bytes: 16 * 1024 * 1024, tokenization_timeout_ms: 5_000, block_size: 16, - data_parallel_size: 1, - kv_event_port_stride: 1, kv_event_port: 5557, replay_port: None, total_kv_blocks: None, max_num_batched_tokens: Some(8192), max_inflight_requests: 1024, - session_affinity_ttl_secs: None, } } @@ -396,7 +391,8 @@ models: } async fn register(selector: &Selector, workers: Vec) { - CatalogReconciler::new(Arc::clone(selector.service.core())) + selector + .catalog_reconciler() .apply(workers) .await .expect("reconcile should succeed"); @@ -433,11 +429,16 @@ models: #[tokio::test] async fn over_long_session_id_is_a_bad_request() { - let mut cfg = test_config(); - cfg.session_affinity_ttl_secs = Some(60.0); - let selector = Selector::new(&cfg, WorkerSelectionPolicyRegistry::default()) - .await - .expect("selector should build"); + let selector = Selector::new_with_kv_router_config( + &test_config(), + KvRouterConfig { + session_affinity_ttl_secs: Some(60.0), + ..Default::default() + }, + WorkerSelectionPolicyRegistry::default(), + ) + .await + .expect("selector should build"); register(&selector, vec![schedulable_registration(1)]).await; let mut request = select_request("res-long-session"); request.session_id = Some("s".repeat(MAX_SESSION_AFFINITY_ID_BYTES + 1)); @@ -743,7 +744,8 @@ worker_selection: .expect("selector should build"); let duplicate = incomplete_registration(1); - let error = CatalogReconciler::new(Arc::clone(selector.service.core())) + let error = selector + .catalog_reconciler() .apply(vec![duplicate.clone(), duplicate]) .await .expect_err("duplicate IDs must be rejected"); @@ -899,12 +901,17 @@ worker_selection: #[tokio::test] async fn invalid_affinity_ttl_returns_configuration_error() { for ttl in [-1.0, 0.0, 0.5, f64::NAN, f64::INFINITY] { - let mut cfg = test_config(); - cfg.session_affinity_ttl_secs = Some(ttl); - let error = Selector::new(&cfg, WorkerSelectionPolicyRegistry::default()) - .await - .err() - .unwrap_or_else(|| panic!("TTL={ttl} must be rejected")); + let error = Selector::new_with_kv_router_config( + &test_config(), + KvRouterConfig { + session_affinity_ttl_secs: Some(ttl), + ..Default::default() + }, + WorkerSelectionPolicyRegistry::default(), + ) + .await + .err() + .unwrap_or_else(|| panic!("TTL={ttl} must be rejected")); assert!( format!("{error:#}").contains("session affinity TTL"), "TTL={ttl}: {error:#}" diff --git a/deploy/inference-gateway/ext-proc/src/server.rs b/deploy/inference-gateway/ext-proc/src/server.rs index 91af6cb545e2..240a62af81af 100644 --- a/deploy/inference-gateway/ext-proc/src/server.rs +++ b/deploy/inference-gateway/ext-proc/src/server.rs @@ -970,6 +970,12 @@ impl ExtProcError { status_code: StatusCode::ServiceUnavailable, message: e.to_string(), }, + // Internal EPP invariant violation — always an EPP bug. 500 keeps + // it distinct from client errors (400) and transient failures (503). + PickError::Internal(msg) => Self { + status_code: StatusCode::InternalServerError, + message: msg, + }, } } diff --git a/deploy/inference-gateway/ext-proc/src/topology_adapter.rs b/deploy/inference-gateway/ext-proc/src/topology_adapter.rs index 244381551259..3a9c2d13bcf6 100644 --- a/deploy/inference-gateway/ext-proc/src/topology_adapter.rs +++ b/deploy/inference-gateway/ext-proc/src/topology_adapter.rs @@ -11,9 +11,7 @@ use std::sync::Arc; use async_trait::async_trait; use dynamo_kv_router::DEFAULT_ROUTING_GROUP; -use dynamo_kv_router::services::selection::{ - CatalogReconciler, WorkerCatalogSource, WorkerRequest, -}; +use dynamo_kv_router::services::selection::{WorkerCatalogSource, WorkerRequest}; use tokio::sync::watch; use tokio_util::sync::CancellationToken; @@ -94,10 +92,7 @@ impl TopologyAdapter { primed: false, closed: false, }; - tokio::spawn( - CatalogReconciler::new(Arc::clone(selector.service.core())) - .run(source, cancel.child_token()), - ); + tokio::spawn(selector.catalog_reconciler().run(source, cancel.child_token())); Self { cancel } } } @@ -145,14 +140,11 @@ mod tests { tokenizer_max_response_bytes: 16 * 1024 * 1024, tokenization_timeout_ms: 5_000, block_size: 16, - data_parallel_size: 1, - kv_event_port_stride: 1, kv_event_port: 5557, replay_port: None, total_kv_blocks: Some(1000), max_num_batched_tokens: Some(8192), max_inflight_requests: 1024, - session_affinity_ttl_secs: None, } } diff --git a/lib/kv-router/src/scheduling/config.rs b/lib/kv-router/src/scheduling/config.rs index 48ca83d332d8..090ba82eec6b 100644 --- a/lib/kv-router/src/scheduling/config.rs +++ b/lib/kv-router/src/scheduling/config.rs @@ -200,6 +200,8 @@ fn log_env_config(config: &KvRouterConfig) { disk_cache_hit_weight = config.disk_cache_hit_weight, router_prefill_load_model = %config.router_prefill_load_model, router_approximate_cache_policy = %config.router_approximate_cache_policy, + session_affinity_ttl_secs = ?config.session_affinity_ttl_secs, + session_affinity_mode = ?config.session_affinity_mode, "KvRouterConfig initialized (DYN_* env overrides applied)" ); } @@ -364,6 +366,12 @@ fn kv_router_config_from_lookup( if let Some(value) = get_env("DYN_ROUTER_PREFILL_LOAD_MODEL") { config.router_prefill_load_model = value.parse()?; } + if let Some(value) = parse_f64(&get_env, "DYN_ROUTER_SESSION_AFFINITY_TTL_SECS") { + config.session_affinity_ttl_secs = Some(value); + } + if let Some(value) = get_env("DYN_ROUTER_SESSION_AFFINITY_MODE") { + config.session_affinity_mode = value.parse()?; + } Ok(config) } @@ -449,6 +457,39 @@ impl FromStr for SharedCacheType { } } +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum SessionAffinityMode { + /// The binding is exact: dispatching to another worker is an error. + #[default] + Hard, + /// The binding follows the dispatch: the session rebinds to where it ran. + Soft, +} + +impl fmt::Display for SessionAffinityMode { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Hard => f.write_str("hard"), + Self::Soft => f.write_str("soft"), + } + } +} + +impl FromStr for SessionAffinityMode { + type Err = String; + + fn from_str(value: &str) -> Result { + match value { + "hard" => Ok(Self::Hard), + "soft" => Ok(Self::Soft), + _ => Err(format!( + "invalid session affinity mode {value:?}; expected 'hard' or 'soft'" + )), + } + } +} + #[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "lowercase")] pub enum RouterQueuePolicy { @@ -893,6 +934,15 @@ pub struct KvRouterConfig { /// maximum overlap. pub router_predicted_ttl_secs: Option, + /// Session-affinity TTL in seconds (`DYN_ROUTER_SESSION_AFFINITY_TTL_SECS`). + /// `None` disables session affinity. Must be finite and in [1, 31536000]. + #[serde(skip)] + pub session_affinity_ttl_secs: Option, + + /// Session-affinity binding mode (`DYN_ROUTER_SESSION_AFFINITY_MODE`). + #[serde(skip)] + pub session_affinity_mode: SessionAffinityMode, + /// Enable conditional-disagg bypass. When true, the `PrefillRouter` /// may short-circuit selected requests to prefill+decode on a decode worker. #[serde(default, skip_serializing_if = "is_default")] @@ -979,6 +1029,8 @@ impl Default for KvRouterConfig { shared_cache_multiplier: 0.0, shared_cache_type: SharedCacheType::default(), router_predicted_ttl_secs: None, + session_affinity_ttl_secs: None, + session_affinity_mode: SessionAffinityMode::Hard, conditional_disagg_enabled: false, conditional_disagg_policy: ConditionalDisaggPolicyKind::default(), conditional_disagg_eff_isl_threshold: default_conditional_disagg_eff_isl_threshold(), @@ -1043,6 +1095,8 @@ impl TryFrom for KvRouterConfig { shared_cache_multiplier: compat.shared_cache_multiplier, shared_cache_type: compat.shared_cache_type, router_predicted_ttl_secs: compat.router_predicted_ttl_secs, + session_affinity_ttl_secs: None, + session_affinity_mode: SessionAffinityMode::Hard, conditional_disagg_enabled: compat.conditional_disagg_enabled, conditional_disagg_policy: compat.conditional_disagg_policy, conditional_disagg_eff_isl_threshold: compat.conditional_disagg_eff_isl_threshold, @@ -1408,6 +1462,13 @@ impl KvRouterConfig { if let Some(value) = self.conditional_disagg_decode_busy_threshold { validate_min("conditional_disagg_decode_busy_threshold", value, 0.0)?; } + if let Some(ttl) = self.session_affinity_ttl_secs { + if !ttl.is_finite() || ttl < 1.0 || ttl > 31_536_000.0 { + return Err(format!( + "session affinity TTL must be a finite value between 1 and 31536000 seconds, got {ttl}" + )); + } + } validate_kv_router_config(self) } diff --git a/lib/kv-router/src/services/selection/affinity.rs b/lib/kv-router/src/services/selection/affinity.rs index 786c86548785..845fff7dbd6c 100644 --- a/lib/kv-router/src/services/selection/affinity.rs +++ b/lib/kv-router/src/services/selection/affinity.rs @@ -17,43 +17,18 @@ use std::sync::{Arc, OnceLock, Weak}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; use dashmap::{DashMap, mapref::entry::Entry}; -use serde::{Deserialize, Serialize}; use tokio::sync::Notify; use tokio::sync::futures::OwnedNotified; use tokio::time::Instant; use tokio_util::sync::CancellationToken; pub use crate::protocols::WorkerAffinityTarget as AffinityTarget; +pub use crate::scheduling::config::SessionAffinityMode; pub const MAX_SESSION_AFFINITY_TTL_SECS: u64 = 31_536_000; pub const MAX_SESSION_AFFINITY_ENTRIES: usize = 65_536; pub const MAX_SESSION_AFFINITY_ID_BYTES: usize = 256; -/// How a bound session treats a dispatch that landed elsewhere. -#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "lowercase")] -pub enum SessionAffinityMode { - /// The binding is exact: dispatching to another worker or rank is an error. - #[default] - Hard, - /// The binding follows the dispatch: the session rebinds to where it ran. - Soft, -} - -impl std::str::FromStr for SessionAffinityMode { - type Err = String; - - fn from_str(value: &str) -> Result { - match value { - "hard" => Ok(Self::Hard), - "soft" => Ok(Self::Soft), - _ => Err(format!( - "invalid session affinity mode {value:?}; expected 'hard' or 'soft'" - )), - } - } -} - #[derive(Debug, thiserror::Error)] pub enum AffinityError { #[error("{0}")] diff --git a/lib/kv-router/src/services/selection/service.rs b/lib/kv-router/src/services/selection/service.rs index 87932f187619..a818ed568a59 100644 --- a/lib/kv-router/src/services/selection/service.rs +++ b/lib/kv-router/src/services/selection/service.rs @@ -141,6 +141,21 @@ impl SelectionServiceBuilder { self.kv_router_config .validate_config() .map_err(anyhow::Error::msg)?; + + // Explicit session_affinity() call takes precedence over KvRouterConfig. + let affinity_config: Option = + if let Some(ttl) = self.session_affinity_ttl { + Some(SessionAffinityConfig::new(ttl)) + } else if let Some(secs) = self.kv_router_config.session_affinity_ttl_secs { + let ttl = Duration::try_from_secs_f64(secs) + .map_err(|_| anyhow::anyhow!("session affinity TTL {secs} out of range"))?; + Some( + SessionAffinityConfig::new(ttl) + .with_mode(self.kv_router_config.session_affinity_mode), + ) + } else { + None + }; let worker_selection_policy_factory = match self.worker_selection_policy_factory { Some(factory) => Some(factory), None => self @@ -187,7 +202,7 @@ impl SelectionServiceBuilder { self.selection_cache, tracking_hash, indexer_policy, - self.session_affinity_ttl.map(SessionAffinityConfig::new), + affinity_config, )); if recover_from_peers { @@ -216,7 +231,7 @@ impl SelectionServiceBuilder { core.dispatch_replica_event(event); } }, - self.session_affinity_ttl.map(|_| { + affinity_config.as_ref().map(|_| { move |event| { if let Some(core) = affinity_core.upgrade() { core.dispatch_affinity_event(event); diff --git a/lib/llm/src/protocols/agents.rs b/lib/llm/src/protocols/agents.rs index 998ca96dfbdf..7fd5f45aeee2 100644 --- a/lib/llm/src/protocols/agents.rs +++ b/lib/llm/src/protocols/agents.rs @@ -122,7 +122,7 @@ fn codex_compaction_header_value(headers: &HeaderMap) -> Option .then(|| metadata.compaction.unwrap_or_default()) } -pub(crate) fn session_affinity_header_value(headers: &HeaderMap) -> Option { +pub fn session_affinity_header_value(headers: &HeaderMap) -> Option { if let Some(session_id) = borrowed_header_value(headers, HEADER_DYNAMO_SESSION_ID) { return Some(session_id.to_owned()); } From 41d1e6846d3211ae66100ca6e68b24e454cce43b Mon Sep 17 00:00:00 2001 From: Thomas Montfort Date: Sun, 13 Sep 2026 15:14:14 -0700 Subject: [PATCH 085/102] fix(router): address PR review comments on session affinity - Apply session_affinity_mode in the explicit session_affinity() builder branch; previously only the KvRouterConfig path called .with_mode() - Document SessionAffinityMode::Hard variant to clarify Hard only rejects overloaded-worker dispatch; worker departure re-initializes regardless - Document KvRouterConfig session_affinity_* fields as EPP-only (populated by kv_router_config_from_lookup); Python/DGD configures via RouterConfig - Return NotReady (503) instead of BadRequest (400) when Hard affinity commit fails because the bound worker is currently overloaded Co-Authored-By: Claude Sonnet 4.6 --- lib/kv-router/src/scheduling/config.rs | 11 +++++++---- .../src/services/selection/core/run.rs | 19 +++++++++++++++++-- .../src/services/selection/service.rs | 2 +- 3 files changed, 25 insertions(+), 7 deletions(-) diff --git a/lib/kv-router/src/scheduling/config.rs b/lib/kv-router/src/scheduling/config.rs index 090ba82eec6b..1b06792f35c8 100644 --- a/lib/kv-router/src/scheduling/config.rs +++ b/lib/kv-router/src/scheduling/config.rs @@ -460,7 +460,9 @@ impl FromStr for SharedCacheType { #[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "lowercase")] pub enum SessionAffinityMode { - /// The binding is exact: dispatching to another worker is an error. + /// The binding is exact: dispatching an overloaded worker's session to + /// another worker is rejected (503). Worker departure re-initializes the + /// session on the new worker regardless of mode. #[default] Hard, /// The binding follows the dispatch: the session rebinds to where it ran. @@ -934,12 +936,13 @@ pub struct KvRouterConfig { /// maximum overlap. pub router_predicted_ttl_secs: Option, - /// Session-affinity TTL in seconds (`DYN_ROUTER_SESSION_AFFINITY_TTL_SECS`). - /// `None` disables session affinity. Must be finite and in [1, 31536000]. + /// Session-affinity TTL in seconds. Populated by `kv_router_config_from_lookup()` + /// via `DYN_ROUTER_SESSION_AFFINITY_TTL_SECS` (EPP path only). The Python/DGD + /// path sets session affinity on `RouterConfig` instead. #[serde(skip)] pub session_affinity_ttl_secs: Option, - /// Session-affinity binding mode (`DYN_ROUTER_SESSION_AFFINITY_MODE`). + /// Session-affinity binding mode. EPP path only; see `session_affinity_ttl_secs`. #[serde(skip)] pub session_affinity_mode: SessionAffinityMode, diff --git a/lib/kv-router/src/services/selection/core/run.rs b/lib/kv-router/src/services/selection/core/run.rs index 8517ced5c673..8ae5e79eaa08 100644 --- a/lib/kv-router/src/services/selection/core/run.rs +++ b/lib/kv-router/src/services/selection/core/run.rs @@ -606,8 +606,8 @@ impl SelectionCore { } /// Bind the held session to `dispatched`. A `Hard` rejection whose bound - /// worker departed after [`Self::hold_session`] checked it is not a client - /// fault: the session is re-initialized on the dispatched worker instead. + /// worker departed is not a client fault: the session re-initializes on + /// the dispatched worker. An overloaded bound worker returns `NotReady`. pub(super) fn commit_session( &self, table: &SessionAffinity, @@ -624,6 +624,21 @@ impl SelectionCore { let departed = bound.is_some_and(|target| !self.catalog.is_schedulable(target.worker_id, key)); if !departed { + // If the bound worker is overloaded, the client is not at fault. + let overloaded = bound.is_some_and(|target| { + self.host + .load + .overloaded_workers + .as_ref() + .and_then(|provider| provider()) + .is_some_and(|ids| ids.contains(&target.worker_id)) + }); + if overloaded { + return Err(SelectionError::NotReady(format!( + "session affinity worker {} is overloaded", + bound.unwrap().worker_id, + ))); + } return Err(affinity_error(error)); } // `commit` invalidated the stale binding. Another request may diff --git a/lib/kv-router/src/services/selection/service.rs b/lib/kv-router/src/services/selection/service.rs index a818ed568a59..da8494bafae8 100644 --- a/lib/kv-router/src/services/selection/service.rs +++ b/lib/kv-router/src/services/selection/service.rs @@ -145,7 +145,7 @@ impl SelectionServiceBuilder { // Explicit session_affinity() call takes precedence over KvRouterConfig. let affinity_config: Option = if let Some(ttl) = self.session_affinity_ttl { - Some(SessionAffinityConfig::new(ttl)) + Some(SessionAffinityConfig::new(ttl).with_mode(self.kv_router_config.session_affinity_mode)) } else if let Some(secs) = self.kv_router_config.session_affinity_ttl_secs { let ttl = Duration::try_from_secs_f64(secs) .map_err(|_| anyhow::anyhow!("session affinity TTL {secs} out of range"))?; From 3b660dba4297ef1e024b6e855e1e74f49d68a98b Mon Sep 17 00:00:00 2001 From: Thomas Montfort Date: Sun, 13 Sep 2026 15:36:16 -0700 Subject: [PATCH 086/102] refactor(router): unify session affinity config into KvRouterConfig Move session_affinity_ttl_secs and session_affinity_mode from RouterConfig kwargs into KvRouterConfig on both the Rust Python binding and the Python configuration layer. Both paths now configure session affinity through KvRouterConfig, eliminating the dead fields on the DGD/Python path and making log_env_config on the EPP path reflect the same config object that the selection service builder consumes. - KvRouterConfig::new() binding accepts session_affinity_ttl_secs and session_affinity_mode; validation and parsing happen there - RouterConfig::new() binding drops these as separate kwargs; From<> derives them from kv_router_config.inner - _KV_ROUTER_FIELDS / KvRouterConfigBase gain session_affinity_*; CLI args move from RouterArgGroup to KvRouterArgGroup - _ROUTER_FIELDS / RouterConfigBase drop session_affinity_* Co-Authored-By: Claude Sonnet 4.6 --- .../configuration/groups/kv_router_args.py | 31 ++++++++++++++++++ .../configuration/groups/router_args.py | 31 ------------------ .../configuration/test_kv_router_args.py | 12 +++---- lib/bindings/python/rust/llm/entrypoint.rs | 32 +++++++++++-------- 4 files changed, 54 insertions(+), 52 deletions(-) diff --git a/components/src/dynamo/common/configuration/groups/kv_router_args.py b/components/src/dynamo/common/configuration/groups/kv_router_args.py index 1a43c56a900a..61878d3958a6 100644 --- a/components/src/dynamo/common/configuration/groups/kv_router_args.py +++ b/components/src/dynamo/common/configuration/groups/kv_router_args.py @@ -66,6 +66,8 @@ "conditional_disagg_prefill_busy_threshold", "conditional_disagg_decode_busy_threshold", "router_predicted_ttl_secs", + "session_affinity_ttl_secs", + "session_affinity_mode", ) CONDITIONAL_DISAGG_POLICY_CHOICES: tuple[str, ...] = ( @@ -228,6 +230,8 @@ class KvRouterConfigBase(ConfigBase): conditional_disagg_prefill_busy_threshold: Optional[float] = None conditional_disagg_decode_busy_threshold: Optional[float] = None router_predicted_ttl_secs: Optional[float] = None + session_affinity_ttl_secs: Optional[int] = None + session_affinity_mode: str = "hard" load_aware: bool = False def apply_load_aware_preset(self) -> None: @@ -696,3 +700,30 @@ def add_arguments(self, parser) -> None: ), arg_type=float, ) + add_argument( + g, + flag_name="--router-session-affinity-ttl-secs", + env_var="DYN_ROUTER_SESSION_AFFINITY_TTL_SECS", + default=None, + help=( + "Enable session affinity with this router-local idle TTL in seconds. " + "Bindings synchronize across router replicas on a best-effort basis. " + "Affinity is disabled when this option is omitted. " + "This is independent of KV prediction TTL settings." + ), + arg_type=int, + dest="session_affinity_ttl_secs", + ) + add_argument( + g, + flag_name="--router-session-affinity-mode", + env_var="DYN_ROUTER_SESSION_AFFINITY_MODE", + default="hard", + help=( + "How an existing session binding participates in worker selection. " + "hard makes the binding an exact constraint; soft exposes it as a " + "policy-visible preference." + ), + choices=("hard", "soft"), + dest="session_affinity_mode", + ) diff --git a/components/src/dynamo/common/configuration/groups/router_args.py b/components/src/dynamo/common/configuration/groups/router_args.py index cdf3bf2f50bb..dceecebab86a 100644 --- a/components/src/dynamo/common/configuration/groups/router_args.py +++ b/components/src/dynamo/common/configuration/groups/router_args.py @@ -36,8 +36,6 @@ "active_decode_blocks_threshold", "active_prefill_tokens_threshold", "active_prefill_tokens_threshold_frac", - "session_affinity_ttl_secs", - "session_affinity_mode", ) _ENFORCE_DISAGG_DEPRECATION = ( @@ -77,8 +75,6 @@ class RouterConfigBase(ConfigBase): router_mode: str min_initial_workers: int enforce_disagg: bool - session_affinity_ttl_secs: Optional[int] - session_affinity_mode: str active_decode_blocks_threshold: Optional[float] active_prefill_tokens_threshold: Optional[int] active_prefill_tokens_threshold_frac: Optional[float] @@ -244,33 +240,6 @@ def add_arguments(self, parser) -> None: "device-aware-weighted", ], ) - add_argument( - g, - flag_name="--router-session-affinity-ttl-secs", - env_var="DYN_ROUTER_SESSION_AFFINITY_TTL_SECS", - default=None, - help=( - "Enable session affinity with this router-local idle TTL in seconds. " - "Bindings synchronize across router replicas on a best-effort basis. " - "Affinity is disabled when this option is omitted. " - "This is independent of KV prediction TTL settings." - ), - arg_type=int, - dest="session_affinity_ttl_secs", - ) - add_argument( - g, - flag_name="--router-session-affinity-mode", - env_var="DYN_ROUTER_SESSION_AFFINITY_MODE", - default="hard", - help=( - "How an existing session binding participates in worker selection. " - "hard makes the binding an exact constraint; soft exposes it as a " - "policy-visible preference." - ), - choices=("hard", "soft"), - dest="session_affinity_mode", - ) add_argument( g, flag_name="--active-decode-blocks-threshold", diff --git a/components/src/dynamo/common/tests/configuration/test_kv_router_args.py b/components/src/dynamo/common/tests/configuration/test_kv_router_args.py index 6e0a7cc9e634..551be9fe601a 100644 --- a/components/src/dynamo/common/tests/configuration/test_kv_router_args.py +++ b/components/src/dynamo/common/tests/configuration/test_kv_router_args.py @@ -589,9 +589,9 @@ def test_frontend_rejection_thresholds_default_to_none( "active_decode_blocks_threshold": None, "active_prefill_tokens_threshold": None, "active_prefill_tokens_threshold_frac": None, - "session_affinity_ttl_secs": None, - "session_affinity_mode": "hard", } + assert config.kv_router_kwargs()["session_affinity_ttl_secs"] is None + assert config.kv_router_kwargs()["session_affinity_mode"] == "hard" assert "busy-worker rejection disabled" in caplog.text @@ -718,8 +718,6 @@ def test_all_rejection_thresholds_and_queue_override_are_forwarded( "active_decode_blocks_threshold": 0.5, "active_prefill_tokens_threshold": 1000, "active_prefill_tokens_threshold_frac": 2.0, - "session_affinity_ttl_secs": None, - "session_affinity_mode": "hard", } assert config.kv_router_kwargs()["router_queue_threshold"] == 32.0 @@ -869,7 +867,7 @@ def test_session_affinity_ttl_cli_and_environment(monkeypatch) -> None: config = FrontendConfig.from_cli_args(parser.parse_args([])) config.validate() assert config.session_affinity_ttl_secs is None - assert config.router_kwargs()["session_affinity_ttl_secs"] is None + assert config.kv_router_kwargs()["session_affinity_ttl_secs"] is None monkeypatch.setenv("DYN_ROUTER_SESSION_AFFINITY_TTL_SECS", "600") parser = argparse.ArgumentParser() @@ -877,7 +875,7 @@ def test_session_affinity_ttl_cli_and_environment(monkeypatch) -> None: config = FrontendConfig.from_cli_args(parser.parse_args([])) config.validate() assert config.session_affinity_ttl_secs == 600 - assert config.router_kwargs()["session_affinity_ttl_secs"] == 600 + assert config.kv_router_kwargs()["session_affinity_ttl_secs"] == 600 parser = argparse.ArgumentParser() FrontendArgGroup().add_arguments(parser) @@ -894,7 +892,7 @@ def test_session_affinity_mode_cli_and_environment(monkeypatch) -> None: FrontendArgGroup().add_arguments(parser) config = FrontendConfig.from_cli_args(parser.parse_args([])) assert config.session_affinity_mode == "hard" - assert config.router_kwargs()["session_affinity_mode"] == "hard" + assert config.kv_router_kwargs()["session_affinity_mode"] == "hard" monkeypatch.setenv("DYN_ROUTER_SESSION_AFFINITY_MODE", "soft") parser = argparse.ArgumentParser() diff --git a/lib/bindings/python/rust/llm/entrypoint.rs b/lib/bindings/python/rust/llm/entrypoint.rs index 7adaf99b4a41..b02c3115b1e6 100644 --- a/lib/bindings/python/rust/llm/entrypoint.rs +++ b/lib/bindings/python/rust/llm/entrypoint.rs @@ -242,7 +242,7 @@ impl AicPerfConfig { #[pymethods] impl KvRouterConfig { #[new] - #[pyo3(signature = (overlap_score_weight=None, host_cache_hit_weight=0.75, disk_cache_hit_weight=0.25, router_temperature=0.0, use_kv_events=true, *, router_replica_sync=false, router_track_active_blocks=true, router_track_output_blocks=false, router_assume_kv_reuse=true, router_track_prefill_tokens=true, router_prefill_load_model="none", router_ttl_secs=120.0, router_approximate_cache_policy="ttl", router_queue_threshold=None, router_event_threads=4, router_queue_policy="fcfs", use_remote_indexer=false, serve_indexer=false, shared_cache_multiplier=0.0, shared_cache_type="none", router_predicted_ttl_secs=None, conditional_disagg_enabled=false, conditional_disagg_policy="isl_bounding", conditional_disagg_eff_isl_threshold=2048, conditional_disagg_eff_isl_ratio_threshold=0.7, conditional_disagg_prefill_busy_threshold=None, conditional_disagg_decode_busy_threshold=None, overlap_score_credit=1.0, overlap_score_credit_decay=0.0, prefill_load_scale=1.0, decode_active_request_weight=0.0, router_policy_config=None, router_prefill_policy=None, router_decode_policy=None, router_tracking_hash="public-xxh3-v1", router_tracking_key_file=None, router_tracking_key_id=None))] + #[pyo3(signature = (overlap_score_weight=None, host_cache_hit_weight=0.75, disk_cache_hit_weight=0.25, router_temperature=0.0, use_kv_events=true, *, router_replica_sync=false, router_track_active_blocks=true, router_track_output_blocks=false, router_assume_kv_reuse=true, router_track_prefill_tokens=true, router_prefill_load_model="none", router_ttl_secs=120.0, router_approximate_cache_policy="ttl", router_queue_threshold=None, router_event_threads=4, router_queue_policy="fcfs", use_remote_indexer=false, serve_indexer=false, shared_cache_multiplier=0.0, shared_cache_type="none", router_predicted_ttl_secs=None, conditional_disagg_enabled=false, conditional_disagg_policy="isl_bounding", conditional_disagg_eff_isl_threshold=2048, conditional_disagg_eff_isl_ratio_threshold=0.7, conditional_disagg_prefill_busy_threshold=None, conditional_disagg_decode_busy_threshold=None, overlap_score_credit=1.0, overlap_score_credit_decay=0.0, prefill_load_scale=1.0, decode_active_request_weight=0.0, router_policy_config=None, router_prefill_policy=None, router_decode_policy=None, router_tracking_hash="public-xxh3-v1", router_tracking_key_file=None, router_tracking_key_id=None, session_affinity_ttl_secs=None, session_affinity_mode="hard"))] #[allow(clippy::too_many_arguments)] fn new( overlap_score_weight: Option, @@ -282,6 +282,8 @@ impl KvRouterConfig { router_tracking_hash: &str, router_tracking_key_file: Option, router_tracking_key_id: Option, + session_affinity_ttl_secs: Option, + session_affinity_mode: &str, ) -> PyResult { if let Some(value) = overlap_score_weight { apply_deprecated_overlap_score_weight( @@ -290,6 +292,10 @@ impl KvRouterConfig { &mut prefill_load_scale, ); } + super::kv::check_session_affinity_ttl_secs(session_affinity_ttl_secs)?; + let session_affinity_mode = session_affinity_mode + .parse::() + .map_err(PyValueError::new_err)?; let inner = RsKvRouterConfig { overlap_score_credit, @@ -339,6 +345,8 @@ impl KvRouterConfig { conditional_disagg_prefill_busy_threshold, conditional_disagg_decode_busy_threshold, router_predicted_ttl_secs, + session_affinity_ttl_secs: session_affinity_ttl_secs.map(|t| t as f64), + session_affinity_mode, }; validate_kv_router_config(&inner)?; Ok(KvRouterConfig { inner }) @@ -480,15 +488,13 @@ pub struct RouterConfig { active_prefill_tokens_threshold: Option, /// Threshold for active prefill tokens as fraction of max_num_batched_tokens active_prefill_tokens_threshold_frac: Option, - session_affinity_ttl_secs: Option, - session_affinity_mode: RsSessionAffinityMode, } #[pymethods] impl RouterConfig { #[new] #[allow(clippy::too_many_arguments)] - #[pyo3(signature = (mode, config=None, active_decode_blocks_threshold=None, active_prefill_tokens_threshold=None, active_prefill_tokens_threshold_frac=None, enforce_disagg=false, session_affinity_ttl_secs=None, session_affinity_mode="hard"))] + #[pyo3(signature = (mode, config=None, active_decode_blocks_threshold=None, active_prefill_tokens_threshold=None, active_prefill_tokens_threshold_frac=None, enforce_disagg=false))] pub fn new( mode: RouterMode, config: Option, @@ -496,8 +502,6 @@ impl RouterConfig { active_prefill_tokens_threshold: Option, active_prefill_tokens_threshold_frac: Option, enforce_disagg: bool, - session_affinity_ttl_secs: Option, - session_affinity_mode: &str, ) -> PyResult { if enforce_disagg { static WARN_ONCE: std::sync::Once = std::sync::Once::new(); @@ -507,7 +511,6 @@ impl RouterConfig { ); }); } - super::kv::check_session_affinity_ttl_secs(session_affinity_ttl_secs)?; RsLoadThresholdConfig { active_decode_blocks_threshold, active_prefill_tokens_threshold, @@ -515,23 +518,24 @@ impl RouterConfig { } .validate() .map_err(PyValueError::new_err)?; - let session_affinity_mode = session_affinity_mode - .parse() - .map_err(PyValueError::new_err)?; Ok(Self { router_mode: mode, kv_router_config: config.unwrap_or_default(), active_decode_blocks_threshold, active_prefill_tokens_threshold, active_prefill_tokens_threshold_frac, - session_affinity_ttl_secs, - session_affinity_mode, }) } } impl From for RsRouterConfig { fn from(rc: RouterConfig) -> RsRouterConfig { + let session_affinity_ttl_secs = rc + .kv_router_config + .inner + .session_affinity_ttl_secs + .map(|f| f as u64); + let session_affinity_mode = rc.kv_router_config.inner.session_affinity_mode; RsRouterConfig { router_mode: rc.router_mode.into(), kv_router_config: rc.kv_router_config.inner, @@ -541,8 +545,8 @@ impl From for RsRouterConfig { active_prefill_tokens_threshold_frac: rc.active_prefill_tokens_threshold_frac, }, enforce_disagg: false, - session_affinity_ttl_secs: rc.session_affinity_ttl_secs, - session_affinity_mode: rc.session_affinity_mode, + session_affinity_ttl_secs, + session_affinity_mode, } } } From 6b6d5c2f23e85dd67a168b4153b7f34c5ad5e30e Mon Sep 17 00:00:00 2001 From: Thomas Montfort Date: Sun, 13 Sep 2026 15:37:03 -0700 Subject: [PATCH 087/102] docs(router): restore canonical doc comment on session affinity fields Now that both EPP and Python/DGD paths populate these fields through KvRouterConfig, the EPP-only caveat no longer applies. Co-Authored-By: Claude Sonnet 4.6 --- lib/kv-router/src/scheduling/config.rs | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/lib/kv-router/src/scheduling/config.rs b/lib/kv-router/src/scheduling/config.rs index 1b06792f35c8..f2490ac65174 100644 --- a/lib/kv-router/src/scheduling/config.rs +++ b/lib/kv-router/src/scheduling/config.rs @@ -936,13 +936,12 @@ pub struct KvRouterConfig { /// maximum overlap. pub router_predicted_ttl_secs: Option, - /// Session-affinity TTL in seconds. Populated by `kv_router_config_from_lookup()` - /// via `DYN_ROUTER_SESSION_AFFINITY_TTL_SECS` (EPP path only). The Python/DGD - /// path sets session affinity on `RouterConfig` instead. + /// Session-affinity TTL in seconds (`DYN_ROUTER_SESSION_AFFINITY_TTL_SECS`). + /// `None` disables session affinity. Must be finite and in [1, 31536000]. #[serde(skip)] pub session_affinity_ttl_secs: Option, - /// Session-affinity binding mode. EPP path only; see `session_affinity_ttl_secs`. + /// Session-affinity binding mode (`DYN_ROUTER_SESSION_AFFINITY_MODE`). #[serde(skip)] pub session_affinity_mode: SessionAffinityMode, From fd09b4bf701e275629bb611525f6d8071617c05d Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Tue, 15 Sep 2026 17:34:33 +0000 Subject: [PATCH 088/102] revert(router): back out EPP affinity follow-up Revert ishandhanani/dynamo#56. Its configuration unification stopped forwarding session affinity settings for non-KV routing modes. Restore the pre-merge configuration and EPP behavior while preserving the later merge of main. This reverts commit ea7ab6e6a65759b03bf531d7f2f4c9c322f2fe6c, reversing changes made to 27c03ee2a4d9c3c4cf9881a933624323f1496bab. Signed-off-by: Ishan Dhanani --- .../configuration/groups/kv_router_args.py | 31 ----- .../configuration/groups/router_args.py | 31 +++++ .../configuration/test_kv_router_args.py | 12 +- .../ext-proc/src/epp_router.rs | 21 +--- .../ext-proc/src/epp_standalone_config.rs | 39 +++++++ .../inference-gateway/ext-proc/src/picker.rs | 4 - .../ext-proc/src/pod_discovery.rs | 107 ++++++++++++++---- .../ext-proc/src/selector.rs | 57 ++++------ .../inference-gateway/ext-proc/src/server.rs | 6 - .../ext-proc/src/topology_adapter.rs | 10 +- lib/bindings/python/rust/llm/entrypoint.rs | 32 +++--- lib/kv-router/src/scheduling/config.rs | 63 ----------- .../src/services/selection/affinity.rs | 27 ++++- .../src/services/selection/core/run.rs | 19 +--- .../src/services/selection/service.rs | 22 +--- lib/llm/src/protocols/agents.rs | 2 +- 16 files changed, 240 insertions(+), 243 deletions(-) diff --git a/components/src/dynamo/common/configuration/groups/kv_router_args.py b/components/src/dynamo/common/configuration/groups/kv_router_args.py index 61878d3958a6..1a43c56a900a 100644 --- a/components/src/dynamo/common/configuration/groups/kv_router_args.py +++ b/components/src/dynamo/common/configuration/groups/kv_router_args.py @@ -66,8 +66,6 @@ "conditional_disagg_prefill_busy_threshold", "conditional_disagg_decode_busy_threshold", "router_predicted_ttl_secs", - "session_affinity_ttl_secs", - "session_affinity_mode", ) CONDITIONAL_DISAGG_POLICY_CHOICES: tuple[str, ...] = ( @@ -230,8 +228,6 @@ class KvRouterConfigBase(ConfigBase): conditional_disagg_prefill_busy_threshold: Optional[float] = None conditional_disagg_decode_busy_threshold: Optional[float] = None router_predicted_ttl_secs: Optional[float] = None - session_affinity_ttl_secs: Optional[int] = None - session_affinity_mode: str = "hard" load_aware: bool = False def apply_load_aware_preset(self) -> None: @@ -700,30 +696,3 @@ def add_arguments(self, parser) -> None: ), arg_type=float, ) - add_argument( - g, - flag_name="--router-session-affinity-ttl-secs", - env_var="DYN_ROUTER_SESSION_AFFINITY_TTL_SECS", - default=None, - help=( - "Enable session affinity with this router-local idle TTL in seconds. " - "Bindings synchronize across router replicas on a best-effort basis. " - "Affinity is disabled when this option is omitted. " - "This is independent of KV prediction TTL settings." - ), - arg_type=int, - dest="session_affinity_ttl_secs", - ) - add_argument( - g, - flag_name="--router-session-affinity-mode", - env_var="DYN_ROUTER_SESSION_AFFINITY_MODE", - default="hard", - help=( - "How an existing session binding participates in worker selection. " - "hard makes the binding an exact constraint; soft exposes it as a " - "policy-visible preference." - ), - choices=("hard", "soft"), - dest="session_affinity_mode", - ) diff --git a/components/src/dynamo/common/configuration/groups/router_args.py b/components/src/dynamo/common/configuration/groups/router_args.py index dceecebab86a..cdf3bf2f50bb 100644 --- a/components/src/dynamo/common/configuration/groups/router_args.py +++ b/components/src/dynamo/common/configuration/groups/router_args.py @@ -36,6 +36,8 @@ "active_decode_blocks_threshold", "active_prefill_tokens_threshold", "active_prefill_tokens_threshold_frac", + "session_affinity_ttl_secs", + "session_affinity_mode", ) _ENFORCE_DISAGG_DEPRECATION = ( @@ -75,6 +77,8 @@ class RouterConfigBase(ConfigBase): router_mode: str min_initial_workers: int enforce_disagg: bool + session_affinity_ttl_secs: Optional[int] + session_affinity_mode: str active_decode_blocks_threshold: Optional[float] active_prefill_tokens_threshold: Optional[int] active_prefill_tokens_threshold_frac: Optional[float] @@ -240,6 +244,33 @@ def add_arguments(self, parser) -> None: "device-aware-weighted", ], ) + add_argument( + g, + flag_name="--router-session-affinity-ttl-secs", + env_var="DYN_ROUTER_SESSION_AFFINITY_TTL_SECS", + default=None, + help=( + "Enable session affinity with this router-local idle TTL in seconds. " + "Bindings synchronize across router replicas on a best-effort basis. " + "Affinity is disabled when this option is omitted. " + "This is independent of KV prediction TTL settings." + ), + arg_type=int, + dest="session_affinity_ttl_secs", + ) + add_argument( + g, + flag_name="--router-session-affinity-mode", + env_var="DYN_ROUTER_SESSION_AFFINITY_MODE", + default="hard", + help=( + "How an existing session binding participates in worker selection. " + "hard makes the binding an exact constraint; soft exposes it as a " + "policy-visible preference." + ), + choices=("hard", "soft"), + dest="session_affinity_mode", + ) add_argument( g, flag_name="--active-decode-blocks-threshold", diff --git a/components/src/dynamo/common/tests/configuration/test_kv_router_args.py b/components/src/dynamo/common/tests/configuration/test_kv_router_args.py index 551be9fe601a..6e0a7cc9e634 100644 --- a/components/src/dynamo/common/tests/configuration/test_kv_router_args.py +++ b/components/src/dynamo/common/tests/configuration/test_kv_router_args.py @@ -589,9 +589,9 @@ def test_frontend_rejection_thresholds_default_to_none( "active_decode_blocks_threshold": None, "active_prefill_tokens_threshold": None, "active_prefill_tokens_threshold_frac": None, + "session_affinity_ttl_secs": None, + "session_affinity_mode": "hard", } - assert config.kv_router_kwargs()["session_affinity_ttl_secs"] is None - assert config.kv_router_kwargs()["session_affinity_mode"] == "hard" assert "busy-worker rejection disabled" in caplog.text @@ -718,6 +718,8 @@ def test_all_rejection_thresholds_and_queue_override_are_forwarded( "active_decode_blocks_threshold": 0.5, "active_prefill_tokens_threshold": 1000, "active_prefill_tokens_threshold_frac": 2.0, + "session_affinity_ttl_secs": None, + "session_affinity_mode": "hard", } assert config.kv_router_kwargs()["router_queue_threshold"] == 32.0 @@ -867,7 +869,7 @@ def test_session_affinity_ttl_cli_and_environment(monkeypatch) -> None: config = FrontendConfig.from_cli_args(parser.parse_args([])) config.validate() assert config.session_affinity_ttl_secs is None - assert config.kv_router_kwargs()["session_affinity_ttl_secs"] is None + assert config.router_kwargs()["session_affinity_ttl_secs"] is None monkeypatch.setenv("DYN_ROUTER_SESSION_AFFINITY_TTL_SECS", "600") parser = argparse.ArgumentParser() @@ -875,7 +877,7 @@ def test_session_affinity_ttl_cli_and_environment(monkeypatch) -> None: config = FrontendConfig.from_cli_args(parser.parse_args([])) config.validate() assert config.session_affinity_ttl_secs == 600 - assert config.kv_router_kwargs()["session_affinity_ttl_secs"] == 600 + assert config.router_kwargs()["session_affinity_ttl_secs"] == 600 parser = argparse.ArgumentParser() FrontendArgGroup().add_arguments(parser) @@ -892,7 +894,7 @@ def test_session_affinity_mode_cli_and_environment(monkeypatch) -> None: FrontendArgGroup().add_arguments(parser) config = FrontendConfig.from_cli_args(parser.parse_args([])) assert config.session_affinity_mode == "hard" - assert config.kv_router_kwargs()["session_affinity_mode"] == "hard" + assert config.router_kwargs()["session_affinity_mode"] == "hard" monkeypatch.setenv("DYN_ROUTER_SESSION_AFFINITY_MODE", "soft") parser = argparse.ArgumentParser() diff --git a/deploy/inference-gateway/ext-proc/src/epp_router.rs b/deploy/inference-gateway/ext-proc/src/epp_router.rs index 7f23fb2de1d7..5f192f4d069e 100644 --- a/deploy/inference-gateway/ext-proc/src/epp_router.rs +++ b/deploy/inference-gateway/ext-proc/src/epp_router.rs @@ -28,6 +28,7 @@ use tokio::sync::Semaphore; use dynamo_kv_router::services::selection::{SelectionError, WorkerSelectionPolicyRegistry}; use dynamo_llm::http::service::metadata::extract_metadata_from_header_pairs; +use dynamo_llm::protocols::agents::HEADER_DYNAMO_SESSION_ID; use dynamo_llm::protocols::common::extensions::{ AgentHints, HEADER_REQUEST_PRIORITY, HEADER_REQUEST_STRICT_PRIORITY, resolve_request_priority, }; @@ -252,6 +253,7 @@ struct RoutingNvExt { cache_namespace: Option, } +/// Case-insensitive lookup of the first non-empty, trimmed value for `name`. fn first_header<'a>(headers: &'a [(String, String)], name: &str) -> Option<&'a str> { headers .iter() @@ -260,17 +262,6 @@ fn first_header<'a>(headers: &'a [(String, String)], name: &str) -> Option<&'a s .filter(|v| !v.is_empty()) } -fn session_id_from_headers(headers: &[(String, String)]) -> Option { - use axum::http::{HeaderMap, HeaderName, HeaderValue}; - let mut map = HeaderMap::with_capacity(headers.len()); - for (k, v) in headers { - if let (Ok(name), Ok(val)) = (k.parse::(), v.parse::()) { - map.insert(name, val); - } - } - dynamo_llm::protocols::agents::session_affinity_header_value(&map) -} - #[tonic::async_trait] impl EndpointPicker for EppRouter { async fn pick( @@ -378,7 +369,7 @@ impl EndpointPicker for EppRouter { model_name: self.model_name.clone(), reservation_id: reservation_id.clone(), token_ids: tokens, - session_id: session_id_from_headers(&req.headers), + session_id: first_header(&req.headers, HEADER_DYNAMO_SESSION_ID).map(str::to_owned), // `None` on the ordinary path: the selector schedules over its // catalog; `Some` only carries an Envoy subset constraint. allowed_worker_ids: allowed, @@ -397,12 +388,6 @@ impl EndpointPicker for EppRouter { Err(SelectionError::BadRequest(message)) => { return Err(PickError::InvalidRequest(message)); } - // A Conflict arises only when the EPP's own UUID collides in the selection - // index — this is an EPP-internal invariant violation, never caused by client - // input, so 500 is correct (not 400 or 503). - Err(SelectionError::Conflict(msg)) => { - return Err(PickError::Internal(format!("selection id conflict: {msg}"))); - } Err(e) => return Err(PickError::RoutingFailed(e.to_string())), }; diff --git a/deploy/inference-gateway/ext-proc/src/epp_standalone_config.rs b/deploy/inference-gateway/ext-proc/src/epp_standalone_config.rs index 4fe4c185b24f..d8a6b38610fb 100644 --- a/deploy/inference-gateway/ext-proc/src/epp_standalone_config.rs +++ b/deploy/inference-gateway/ext-proc/src/epp_standalone_config.rs @@ -140,6 +140,17 @@ pub struct EppStandaloneConfig { /// KV-cache block size; MUST equal the inference engine block size. #[validate(range(min = 1, message = "DYN_KV_CACHE_BLOCK_SIZE must be >= 1"))] pub block_size: u32, + /// Data-parallel ranks per worker pod. Only `1` is accepted: the EPP routes + /// by worker and cannot convey the selected rank to the pod. + #[validate(range( + min = 1, + max = 1, + message = "DYN_EPP_DATA_PARALLEL_SIZE must be 1; multi-rank standalone serving is unsupported" + ))] + pub data_parallel_size: u32, + /// Port distance between consecutive data-parallel ranks' KV event ports. + #[validate(range(min = 1, message = "DYN_EPP_KV_EVENT_PORT_STRIDE must be >= 1"))] + pub kv_event_port_stride: u16, /// KV zmq event port. #[validate(range(min = 1))] pub kv_event_port: u16, @@ -163,6 +174,15 @@ pub struct EppStandaloneConfig { /// throughput throttle). Excess requests are shed with a 503, not queued. #[validate(range(min = 1, message = "DYN_EPP_MAX_INFLIGHT_REQUESTS must be >= 1"))] pub max_inflight_requests: usize, + /// Pin each `x-dynamo-session-id` to the worker that served it for this + /// long after its last request (`DYN_EPP_SESSION_AFFINITY_TTL_SECS`). + /// `None` disables session affinity. + #[validate(range( + min = 1.0, + max = 31536000.0, + message = "DYN_EPP_SESSION_AFFINITY_TTL_SECS must be between 1 and 31536000 seconds" + ))] + pub session_affinity_ttl_secs: Option, } impl EppStandaloneConfig { @@ -216,6 +236,9 @@ impl EppStandaloneConfig { )? .unwrap_or(DEFAULT_TOKENIZER_MAX_RESPONSE_BYTES), block_size: opt_parse::(get, "DYN_KV_CACHE_BLOCK_SIZE")?.unwrap_or(0), + data_parallel_size: opt_parse::(get, "DYN_EPP_DATA_PARALLEL_SIZE")?.unwrap_or(1), + kv_event_port_stride: opt_parse::(get, "DYN_EPP_KV_EVENT_PORT_STRIDE")? + .unwrap_or(1), kv_event_port: opt_parse::(get, "DYN_EPP_KV_EVENT_PORT")? .unwrap_or(DEFAULT_KV_EVENT_PORT), replay_port: opt_parse::(get, "DYN_EPP_KV_EVENT_REPLAY_PORT")?, @@ -223,6 +246,7 @@ impl EppStandaloneConfig { max_num_batched_tokens: opt_parse::(get, "DYN_EPP_MAX_NUM_BATCHED_TOKENS")?, max_inflight_requests: opt_parse::(get, "DYN_EPP_MAX_INFLIGHT_REQUESTS")? .unwrap_or(DEFAULT_MAX_INFLIGHT_REQUESTS), + session_affinity_ttl_secs: opt_parse::(get, "DYN_EPP_SESSION_AFFINITY_TTL_SECS")?, }) } @@ -543,6 +567,21 @@ mod tests { ); } + #[test] + fn data_parallel_size_above_one_fails() { + assert!( + parse_cfg(&[ + ("DYN_EPP_INFERENCE_POOL_NAME", "vllm-qwen-pool"), + ("POD_NAMESPACE", "inference"), + ("DYN_MODEL_NAME", "Qwen/Qwen3-0.6B"), + ("DYN_EPP_TOKENIZER_SERVICE_URL", "http://vllm-render:8000"), + ("DYN_EPP_TOKENIZER_PROTOCOL", "vllm-render"), + ("DYN_EPP_DATA_PARALLEL_SIZE", "2"), + ]) + .is_err() + ); + } + #[test] fn replay_port_is_optional_and_uses_the_kv_event_name() { let cfg = parse_cfg(&[ diff --git a/deploy/inference-gateway/ext-proc/src/picker.rs b/deploy/inference-gateway/ext-proc/src/picker.rs index 6708626a9be6..f1b0aec80340 100644 --- a/deploy/inference-gateway/ext-proc/src/picker.rs +++ b/deploy/inference-gateway/ext-proc/src/picker.rs @@ -203,10 +203,6 @@ pub enum PickError { /// multiplexing means the connection cap does not bound concurrent requests. #[error("endpoint picker overloaded")] Overloaded, - /// An internal EPP invariant was violated; always an EPP bug, never a client - /// error → 500. - #[error("internal error: {0}")] - Internal(String), } #[cfg(test)] diff --git a/deploy/inference-gateway/ext-proc/src/pod_discovery.rs b/deploy/inference-gateway/ext-proc/src/pod_discovery.rs index f9c0753213a8..39c49f1c7028 100644 --- a/deploy/inference-gateway/ext-proc/src/pod_discovery.rs +++ b/deploy/inference-gateway/ext-proc/src/pod_discovery.rs @@ -109,6 +109,11 @@ impl PodDiscovery { let kv_event_port = cfg.kv_event_port; let replay_port = cfg.replay_port; + let kv_ports = KvEventPorts { + base: kv_event_port, + stride: cfg.kv_event_port_stride, + data_parallel_size: cfg.data_parallel_size, + }; let index: Arc> = Arc::new(RwLock::new(WorkerIndex::new())); @@ -185,7 +190,7 @@ impl PodDiscovery { Delta::Rebuild => rebuild_index( &store, pool_rx.borrow().as_ref(), - kv_event_port, + kv_ports, replay_port, &index_task, ), @@ -193,7 +198,7 @@ impl PodDiscovery { &index_task, &pod, pool_rx.borrow().as_ref(), - kv_event_port, + kv_ports, replay_port, ), Delta::Remove(pod) => remove_pod(&index_task, &pod), @@ -346,14 +351,14 @@ fn upsert_pod( index: &RwLock, pod: &Pod, pool: Option<&PoolState>, - kv_event_port: u16, + kv_ports: KvEventPorts, replay_port: Option, ) -> bool { let Some(worker_id) = pod_worker_id(pod) else { return false; }; let entry = pool - .and_then(|pool| raw_worker_from_pod(pod, pool, kv_event_port, replay_port)) + .and_then(|pool| raw_worker_from_pod(pod, pool, kv_ports, replay_port)) .map(WorkerEntry::from_raw); let mut index = index.write().unwrap(); match entry { @@ -384,14 +389,14 @@ fn remove_pod(index: &RwLock, pod: &Pod) -> bool { fn rebuild_index( store: &kube::runtime::reflector::Store, pool: Option<&PoolState>, - kv_event_port: u16, + kv_ports: KvEventPorts, replay_port: Option, index: &RwLock, ) -> bool { let mut fresh = WorkerIndex::new(); if let Some(pool) = pool { for pod in store.state().iter() { - if let Some(worker) = raw_worker_from_pod(pod, pool, kv_event_port, replay_port) { + if let Some(worker) = raw_worker_from_pod(pod, pool, kv_ports, replay_port) { fresh.insert(worker.worker_id, WorkerEntry::from_raw(worker)); } } @@ -405,12 +410,33 @@ fn rebuild_index( } } +/// Per-pod KV-event port layout: rank `r` publishes on `base + r * stride`. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct KvEventPorts { + pub base: u16, + pub stride: u16, + pub data_parallel_size: u32, +} + +impl KvEventPorts { + fn endpoints(self, ip: IpAddr) -> HashMap { + (0..self.data_parallel_size.max(1)) + .map(|rank| { + let port = self + .base + .saturating_add(self.stride.saturating_mul(rank as u16)); + (rank, format!("tcp://{}", SocketAddr::new(ip, port))) + }) + .collect() + } +} + /// Build a [`RawWorker`] from a pod, or `None` if it is not `Ready`, not /// pool-selected, or lacks an IP/name. Pure function — unit-testable. fn raw_worker_from_pod( pod: &Pod, pool: &PoolState, - kv_event_port: u16, + kv_ports: KvEventPorts, replay_port: Option, ) -> Option { if !pod_is_ready(pod) || !pod_matches(pod, &pool.match_labels) { @@ -425,10 +451,7 @@ fn raw_worker_from_pod( pod_name: pod_name.to_string(), pod_ip: pod_ip.to_string(), http_endpoint: format!("http://{}", SocketAddr::new(ip, pool.target_port)), - kv_events_endpoints: HashMap::from([( - 0u32, - format!("tcp://{}", SocketAddr::new(ip, kv_event_port)), - )]), + kv_events_endpoints: kv_ports.endpoints(ip), replay_endpoint: replay_port.map(|p| format!("tcp://{}", SocketAddr::new(ip, p))), }) } @@ -440,6 +463,42 @@ mod tests { use k8s_openapi::apimachinery::pkg::apis::meta::v1::Time; use kube::api::ObjectMeta; + fn single_rank(port: u16) -> KvEventPorts { + KvEventPorts { + base: port, + stride: 1, + data_parallel_size: 1, + } + } + + #[test] + fn data_parallel_ranks_publish_on_strided_ports() { + let w = raw_worker_from_pod( + &pod( + "vllm-0", + Some("10.0.0.1"), + Some(true), + &[("app", "vllm-qwen")], + ), + &pool(), + KvEventPorts { + base: 5557, + stride: 2, + data_parallel_size: 3, + }, + None, + ) + .expect("ready, selected pod should map"); + assert_eq!( + w.kv_events_endpoints, + HashMap::from([ + (0, "tcp://10.0.0.1:5557".to_string()), + (1, "tcp://10.0.0.1:5559".to_string()), + (2, "tcp://10.0.0.1:5561".to_string()), + ]) + ); + } + fn pool() -> PoolState { PoolState { match_labels: BTreeMap::from([("app".to_string(), "vllm-qwen".to_string())]), @@ -484,7 +543,7 @@ mod tests { &[("app", "vllm-qwen")], ), &pool(), - 5557u16, + single_rank(5557), Some(5560), ) .expect("ready, selected pod should map"); @@ -507,7 +566,7 @@ mod tests { &[("app", "vllm-qwen")], ), &pool(), - 5557u16, + single_rank(5557), Some(5560), ) .expect("ready, selected IPv6 pod should map"); @@ -531,7 +590,7 @@ mod tests { &[("app", "vllm-qwen")] ), &pool(), - 5557u16, + single_rank(5557), None, ) .is_none() @@ -549,7 +608,7 @@ mod tests { &[("app", "something-else")] ), &pool(), - 5557u16, + single_rank(5557), None, ) .is_none() @@ -567,7 +626,7 @@ mod tests { &[("app", "vllm-qwen")] ), &pool(), - 5557u16, + single_rank(5557), None, ) .is_none() @@ -583,7 +642,7 @@ mod tests { &[("app", "vllm-qwen")], ); p.metadata.deletion_timestamp = Some(Time(k8s_openapi::chrono::Utc::now())); - assert!(raw_worker_from_pod(&p, &pool(), 5557u16, None).is_none()); + assert!(raw_worker_from_pod(&p, &pool(), single_rank(5557), None).is_none()); } #[test] @@ -592,7 +651,7 @@ mod tests { raw_worker_from_pod( &pod("vllm-0", None, Some(true), &[("app", "vllm-qwen")]), &pool(), - 5557u16, + single_rank(5557), None, ) .is_none() @@ -613,7 +672,7 @@ mod tests { #[test] fn rebuild_index_keeps_only_ready_selected_pods() { - let kp = 5557u16; + let kp = single_rank(5557); let store = store_from_pods(vec![ pod( "vllm-0", @@ -652,7 +711,7 @@ mod tests { #[test] fn rebuild_index_is_empty_without_pool() { - let kp = 5557u16; + let kp = single_rank(5557); let store = store_from_pods(vec![pod( "vllm-0", Some("10.0.0.1"), @@ -666,7 +725,7 @@ mod tests { #[test] fn upsert_and_remove_pod_mutate_index_incrementally() { - let kp = 5557u16; + let kp = single_rank(5557); let index = RwLock::new(WorkerIndex::new()); let id = hash_pod_name("vllm-0"); let ready = pod( @@ -709,7 +768,7 @@ mod tests { #[test] fn upsert_pod_without_pool_drops_entry() { - let kp = 5557u16; + let kp = single_rank(5557); // A `None` pool (unresolved or deleted) means nothing is routable, so an // upsert must evict any existing entry rather than leave stale routing. let index = RwLock::new(WorkerIndex::new()); @@ -731,7 +790,7 @@ mod tests { #[test] fn pool_edit_during_relist_rebuilds_at_init_done_from_completed_store() { - let kp = 5557u16; + let kp = single_rank(5557); use kube::runtime::watcher; let vllm_0 = pod( @@ -798,7 +857,7 @@ mod tests { #[test] fn rebuild_index_drops_workers_absent_from_the_store() { - let kp = 5557u16; + let kp = single_rank(5557); // A relist arrives as a fresh snapshot with no `Delete` events for pods // that disappeared during the disconnect, so the `InitDone`/pool-change // rebuild must *replace* the index, not merge into it — otherwise dead diff --git a/deploy/inference-gateway/ext-proc/src/selector.rs b/deploy/inference-gateway/ext-proc/src/selector.rs index 85314a8500fe..d597b0b2e10a 100644 --- a/deploy/inference-gateway/ext-proc/src/selector.rs +++ b/deploy/inference-gateway/ext-proc/src/selector.rs @@ -15,8 +15,8 @@ use anyhow::{Context, Result, anyhow}; use dynamo_kv_router::config::{KvRouterConfig, try_kv_router_config_from_dynamo_env}; use dynamo_kv_router::protocols::RoutingConstraints; use dynamo_kv_router::services::selection::{ - CatalogReconciler, PromptRequest, SelectAndReserveRequest as CoreSelectAndReserveRequest, - SelectionError, SelectionService, SelectionServiceBuilder, WorkerSelectionPolicyRegistry, + PromptRequest, SelectAndReserveRequest as CoreSelectAndReserveRequest, SelectionError, + SelectionService, SelectionServiceBuilder, WorkerSelectionPolicyRegistry, warn_for_unserved_worker_selection_policies, }; use dynamo_kv_router::{DEFAULT_ROUTING_GROUP, WorkerType}; @@ -67,7 +67,7 @@ pub struct SelectResponse { /// In-process runtime-free selector wrapping a [`SelectionService`]. pub struct Selector { - service: Arc, + pub(crate) service: Arc, /// Cancels the peer-discovery watch on drop. The `SelectionService`'s own /// `Drop` tears down its core + replica-sync tasks. cancel: CancellationToken, @@ -116,6 +116,12 @@ impl Selector { if let Some(peer_replication) = peer_replication { builder = builder.replica_sync(peer_replication.sync_port, Vec::new()); } + if let Some(ttl) = cfg.session_affinity_ttl_secs { + builder = builder.session_affinity( + std::time::Duration::try_from_secs_f64(ttl) + .context("invalid session affinity TTL")?, + ); + } let service = Arc::new( builder .build() @@ -161,10 +167,6 @@ impl Selector { Ok(()) } - pub(crate) fn catalog_reconciler(&self) -> CatalogReconciler { - CatalogReconciler::new(Arc::clone(self.service.core())) - } - /// Select a worker for a prompt and book its load in one operation. Takes the /// request by value so per-request fields are moved into the core request /// rather than cloned on the hot path. @@ -350,11 +352,14 @@ models: tokenizer_max_response_bytes: 16 * 1024 * 1024, tokenization_timeout_ms: 5_000, block_size: 16, + data_parallel_size: 1, + kv_event_port_stride: 1, kv_event_port: 5557, replay_port: None, total_kv_blocks: None, max_num_batched_tokens: Some(8192), max_inflight_requests: 1024, + session_affinity_ttl_secs: None, } } @@ -393,8 +398,7 @@ models: } async fn register(selector: &Selector, workers: Vec) { - selector - .catalog_reconciler() + CatalogReconciler::new(Arc::clone(selector.service.core())) .apply(workers) .await .expect("reconcile should succeed"); @@ -432,16 +436,11 @@ models: #[tokio::test] async fn over_long_session_id_is_a_bad_request() { - let selector = Selector::new_with_kv_router_config( - &test_config(), - KvRouterConfig { - session_affinity_ttl_secs: Some(60.0), - ..Default::default() - }, - WorkerSelectionPolicyRegistry::default(), - ) - .await - .expect("selector should build"); + let mut cfg = test_config(); + cfg.session_affinity_ttl_secs = Some(60.0); + let selector = Selector::new(&cfg, WorkerSelectionPolicyRegistry::default()) + .await + .expect("selector should build"); register(&selector, vec![schedulable_registration(1)]).await; let mut request = select_request("res-long-session"); request.session_id = Some("s".repeat(MAX_SESSION_AFFINITY_ID_BYTES + 1)); @@ -747,8 +746,7 @@ worker_selection: .expect("selector should build"); let duplicate = incomplete_registration(1); - let error = selector - .catalog_reconciler() + let error = CatalogReconciler::new(Arc::clone(selector.service.core())) .apply(vec![duplicate.clone(), duplicate]) .await .expect_err("duplicate IDs must be rejected"); @@ -904,17 +902,12 @@ worker_selection: #[tokio::test] async fn invalid_affinity_ttl_returns_configuration_error() { for ttl in [-1.0, 0.0, 0.5, f64::NAN, f64::INFINITY] { - let error = Selector::new_with_kv_router_config( - &test_config(), - KvRouterConfig { - session_affinity_ttl_secs: Some(ttl), - ..Default::default() - }, - WorkerSelectionPolicyRegistry::default(), - ) - .await - .err() - .unwrap_or_else(|| panic!("TTL={ttl} must be rejected")); + let mut cfg = test_config(); + cfg.session_affinity_ttl_secs = Some(ttl); + let error = Selector::new(&cfg, WorkerSelectionPolicyRegistry::default()) + .await + .err() + .unwrap_or_else(|| panic!("TTL={ttl} must be rejected")); assert!( format!("{error:#}").contains("session affinity TTL"), "TTL={ttl}: {error:#}" diff --git a/deploy/inference-gateway/ext-proc/src/server.rs b/deploy/inference-gateway/ext-proc/src/server.rs index 8ab80f781331..e39bfa837389 100644 --- a/deploy/inference-gateway/ext-proc/src/server.rs +++ b/deploy/inference-gateway/ext-proc/src/server.rs @@ -998,12 +998,6 @@ impl ExtProcError { status_code: StatusCode::ServiceUnavailable, message: e.to_string(), }, - // Internal EPP invariant violation — always an EPP bug. 500 keeps - // it distinct from client errors (400) and transient failures (503). - PickError::Internal(msg) => Self { - status_code: StatusCode::InternalServerError, - message: msg, - }, } } diff --git a/deploy/inference-gateway/ext-proc/src/topology_adapter.rs b/deploy/inference-gateway/ext-proc/src/topology_adapter.rs index 4a04c3d13bc8..244381551259 100644 --- a/deploy/inference-gateway/ext-proc/src/topology_adapter.rs +++ b/deploy/inference-gateway/ext-proc/src/topology_adapter.rs @@ -11,7 +11,9 @@ use std::sync::Arc; use async_trait::async_trait; use dynamo_kv_router::DEFAULT_ROUTING_GROUP; -use dynamo_kv_router::services::selection::{WorkerCatalogSource, WorkerRequest}; +use dynamo_kv_router::services::selection::{ + CatalogReconciler, WorkerCatalogSource, WorkerRequest, +}; use tokio::sync::watch; use tokio_util::sync::CancellationToken; @@ -93,8 +95,7 @@ impl TopologyAdapter { closed: false, }; tokio::spawn( - selector - .catalog_reconciler() + CatalogReconciler::new(Arc::clone(selector.service.core())) .run(source, cancel.child_token()), ); Self { cancel } @@ -144,11 +145,14 @@ mod tests { tokenizer_max_response_bytes: 16 * 1024 * 1024, tokenization_timeout_ms: 5_000, block_size: 16, + data_parallel_size: 1, + kv_event_port_stride: 1, kv_event_port: 5557, replay_port: None, total_kv_blocks: Some(1000), max_num_batched_tokens: Some(8192), max_inflight_requests: 1024, + session_affinity_ttl_secs: None, } } diff --git a/lib/bindings/python/rust/llm/entrypoint.rs b/lib/bindings/python/rust/llm/entrypoint.rs index 93fcd66539fc..d8765e9ae45c 100644 --- a/lib/bindings/python/rust/llm/entrypoint.rs +++ b/lib/bindings/python/rust/llm/entrypoint.rs @@ -242,7 +242,7 @@ impl AicPerfConfig { #[pymethods] impl KvRouterConfig { #[new] - #[pyo3(signature = (overlap_score_weight=None, host_cache_hit_weight=0.75, disk_cache_hit_weight=0.25, router_temperature=0.0, use_kv_events=true, *, router_replica_sync=false, router_track_active_blocks=true, router_track_output_blocks=false, router_assume_kv_reuse=true, router_track_prefill_tokens=true, router_prefill_load_model="none", router_ttl_secs=120.0, router_approximate_cache_policy="ttl", router_queue_threshold=None, router_event_threads=4, router_queue_policy="fcfs", use_remote_indexer=false, serve_indexer=false, shared_cache_multiplier=0.0, shared_cache_type="none", router_predicted_ttl_secs=None, conditional_disagg_enabled=false, conditional_disagg_policy="isl_bounding", conditional_disagg_eff_isl_threshold=2048, conditional_disagg_eff_isl_ratio_threshold=0.7, conditional_disagg_prefill_busy_threshold=None, conditional_disagg_decode_busy_threshold=None, overlap_score_credit=1.0, overlap_score_credit_decay=0.0, prefill_load_scale=1.0, decode_active_request_weight=0.0, router_policy_config=None, router_prefill_policy=None, router_decode_policy=None, router_tracking_hash="public-xxh3-v1", router_tracking_key_file=None, router_tracking_key_id=None, session_affinity_ttl_secs=None, session_affinity_mode="hard"))] + #[pyo3(signature = (overlap_score_weight=None, host_cache_hit_weight=0.75, disk_cache_hit_weight=0.25, router_temperature=0.0, use_kv_events=true, *, router_replica_sync=false, router_track_active_blocks=true, router_track_output_blocks=false, router_assume_kv_reuse=true, router_track_prefill_tokens=true, router_prefill_load_model="none", router_ttl_secs=120.0, router_approximate_cache_policy="ttl", router_queue_threshold=None, router_event_threads=4, router_queue_policy="fcfs", use_remote_indexer=false, serve_indexer=false, shared_cache_multiplier=0.0, shared_cache_type="none", router_predicted_ttl_secs=None, conditional_disagg_enabled=false, conditional_disagg_policy="isl_bounding", conditional_disagg_eff_isl_threshold=2048, conditional_disagg_eff_isl_ratio_threshold=0.7, conditional_disagg_prefill_busy_threshold=None, conditional_disagg_decode_busy_threshold=None, overlap_score_credit=1.0, overlap_score_credit_decay=0.0, prefill_load_scale=1.0, decode_active_request_weight=0.0, router_policy_config=None, router_prefill_policy=None, router_decode_policy=None, router_tracking_hash="public-xxh3-v1", router_tracking_key_file=None, router_tracking_key_id=None))] #[allow(clippy::too_many_arguments)] fn new( overlap_score_weight: Option, @@ -282,8 +282,6 @@ impl KvRouterConfig { router_tracking_hash: &str, router_tracking_key_file: Option, router_tracking_key_id: Option, - session_affinity_ttl_secs: Option, - session_affinity_mode: &str, ) -> PyResult { if let Some(value) = overlap_score_weight { apply_deprecated_overlap_score_weight( @@ -292,10 +290,6 @@ impl KvRouterConfig { &mut prefill_load_scale, ); } - super::kv::check_session_affinity_ttl_secs(session_affinity_ttl_secs)?; - let session_affinity_mode = session_affinity_mode - .parse::() - .map_err(PyValueError::new_err)?; let inner = RsKvRouterConfig { overlap_score_credit, @@ -345,8 +339,6 @@ impl KvRouterConfig { conditional_disagg_prefill_busy_threshold, conditional_disagg_decode_busy_threshold, router_predicted_ttl_secs, - session_affinity_ttl_secs: session_affinity_ttl_secs.map(|t| t as f64), - session_affinity_mode, }; validate_kv_router_config(&inner)?; Ok(KvRouterConfig { inner }) @@ -488,13 +480,15 @@ pub struct RouterConfig { active_prefill_tokens_threshold: Option, /// Threshold for active prefill tokens as fraction of max_num_batched_tokens active_prefill_tokens_threshold_frac: Option, + session_affinity_ttl_secs: Option, + session_affinity_mode: RsSessionAffinityMode, } #[pymethods] impl RouterConfig { #[new] #[allow(clippy::too_many_arguments)] - #[pyo3(signature = (mode, config=None, active_decode_blocks_threshold=None, active_prefill_tokens_threshold=None, active_prefill_tokens_threshold_frac=None, enforce_disagg=false))] + #[pyo3(signature = (mode, config=None, active_decode_blocks_threshold=None, active_prefill_tokens_threshold=None, active_prefill_tokens_threshold_frac=None, enforce_disagg=false, session_affinity_ttl_secs=None, session_affinity_mode="hard"))] pub fn new( mode: RouterMode, config: Option, @@ -502,6 +496,8 @@ impl RouterConfig { active_prefill_tokens_threshold: Option, active_prefill_tokens_threshold_frac: Option, enforce_disagg: bool, + session_affinity_ttl_secs: Option, + session_affinity_mode: &str, ) -> PyResult { if enforce_disagg { static WARN_ONCE: std::sync::Once = std::sync::Once::new(); @@ -511,6 +507,7 @@ impl RouterConfig { ); }); } + super::kv::check_session_affinity_ttl_secs(session_affinity_ttl_secs)?; RsLoadThresholdConfig { active_decode_blocks_threshold, active_prefill_tokens_threshold, @@ -518,24 +515,23 @@ impl RouterConfig { } .validate() .map_err(PyValueError::new_err)?; + let session_affinity_mode = session_affinity_mode + .parse() + .map_err(PyValueError::new_err)?; Ok(Self { router_mode: mode, kv_router_config: config.unwrap_or_default(), active_decode_blocks_threshold, active_prefill_tokens_threshold, active_prefill_tokens_threshold_frac, + session_affinity_ttl_secs, + session_affinity_mode, }) } } impl From for RsRouterConfig { fn from(rc: RouterConfig) -> RsRouterConfig { - let session_affinity_ttl_secs = rc - .kv_router_config - .inner - .session_affinity_ttl_secs - .map(|f| f as u64); - let session_affinity_mode = rc.kv_router_config.inner.session_affinity_mode; RsRouterConfig { router_mode: rc.router_mode.into(), kv_router_config: rc.kv_router_config.inner, @@ -545,8 +541,8 @@ impl From for RsRouterConfig { active_prefill_tokens_threshold_frac: rc.active_prefill_tokens_threshold_frac, }, enforce_disagg: false, - session_affinity_ttl_secs, - session_affinity_mode, + session_affinity_ttl_secs: rc.session_affinity_ttl_secs, + session_affinity_mode: rc.session_affinity_mode, } } } diff --git a/lib/kv-router/src/scheduling/config.rs b/lib/kv-router/src/scheduling/config.rs index f2490ac65174..48ca83d332d8 100644 --- a/lib/kv-router/src/scheduling/config.rs +++ b/lib/kv-router/src/scheduling/config.rs @@ -200,8 +200,6 @@ fn log_env_config(config: &KvRouterConfig) { disk_cache_hit_weight = config.disk_cache_hit_weight, router_prefill_load_model = %config.router_prefill_load_model, router_approximate_cache_policy = %config.router_approximate_cache_policy, - session_affinity_ttl_secs = ?config.session_affinity_ttl_secs, - session_affinity_mode = ?config.session_affinity_mode, "KvRouterConfig initialized (DYN_* env overrides applied)" ); } @@ -366,12 +364,6 @@ fn kv_router_config_from_lookup( if let Some(value) = get_env("DYN_ROUTER_PREFILL_LOAD_MODEL") { config.router_prefill_load_model = value.parse()?; } - if let Some(value) = parse_f64(&get_env, "DYN_ROUTER_SESSION_AFFINITY_TTL_SECS") { - config.session_affinity_ttl_secs = Some(value); - } - if let Some(value) = get_env("DYN_ROUTER_SESSION_AFFINITY_MODE") { - config.session_affinity_mode = value.parse()?; - } Ok(config) } @@ -457,41 +449,6 @@ impl FromStr for SharedCacheType { } } -#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "lowercase")] -pub enum SessionAffinityMode { - /// The binding is exact: dispatching an overloaded worker's session to - /// another worker is rejected (503). Worker departure re-initializes the - /// session on the new worker regardless of mode. - #[default] - Hard, - /// The binding follows the dispatch: the session rebinds to where it ran. - Soft, -} - -impl fmt::Display for SessionAffinityMode { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - Self::Hard => f.write_str("hard"), - Self::Soft => f.write_str("soft"), - } - } -} - -impl FromStr for SessionAffinityMode { - type Err = String; - - fn from_str(value: &str) -> Result { - match value { - "hard" => Ok(Self::Hard), - "soft" => Ok(Self::Soft), - _ => Err(format!( - "invalid session affinity mode {value:?}; expected 'hard' or 'soft'" - )), - } - } -} - #[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "lowercase")] pub enum RouterQueuePolicy { @@ -936,15 +893,6 @@ pub struct KvRouterConfig { /// maximum overlap. pub router_predicted_ttl_secs: Option, - /// Session-affinity TTL in seconds (`DYN_ROUTER_SESSION_AFFINITY_TTL_SECS`). - /// `None` disables session affinity. Must be finite and in [1, 31536000]. - #[serde(skip)] - pub session_affinity_ttl_secs: Option, - - /// Session-affinity binding mode (`DYN_ROUTER_SESSION_AFFINITY_MODE`). - #[serde(skip)] - pub session_affinity_mode: SessionAffinityMode, - /// Enable conditional-disagg bypass. When true, the `PrefillRouter` /// may short-circuit selected requests to prefill+decode on a decode worker. #[serde(default, skip_serializing_if = "is_default")] @@ -1031,8 +979,6 @@ impl Default for KvRouterConfig { shared_cache_multiplier: 0.0, shared_cache_type: SharedCacheType::default(), router_predicted_ttl_secs: None, - session_affinity_ttl_secs: None, - session_affinity_mode: SessionAffinityMode::Hard, conditional_disagg_enabled: false, conditional_disagg_policy: ConditionalDisaggPolicyKind::default(), conditional_disagg_eff_isl_threshold: default_conditional_disagg_eff_isl_threshold(), @@ -1097,8 +1043,6 @@ impl TryFrom for KvRouterConfig { shared_cache_multiplier: compat.shared_cache_multiplier, shared_cache_type: compat.shared_cache_type, router_predicted_ttl_secs: compat.router_predicted_ttl_secs, - session_affinity_ttl_secs: None, - session_affinity_mode: SessionAffinityMode::Hard, conditional_disagg_enabled: compat.conditional_disagg_enabled, conditional_disagg_policy: compat.conditional_disagg_policy, conditional_disagg_eff_isl_threshold: compat.conditional_disagg_eff_isl_threshold, @@ -1464,13 +1408,6 @@ impl KvRouterConfig { if let Some(value) = self.conditional_disagg_decode_busy_threshold { validate_min("conditional_disagg_decode_busy_threshold", value, 0.0)?; } - if let Some(ttl) = self.session_affinity_ttl_secs { - if !ttl.is_finite() || ttl < 1.0 || ttl > 31_536_000.0 { - return Err(format!( - "session affinity TTL must be a finite value between 1 and 31536000 seconds, got {ttl}" - )); - } - } validate_kv_router_config(self) } diff --git a/lib/kv-router/src/services/selection/affinity.rs b/lib/kv-router/src/services/selection/affinity.rs index 845fff7dbd6c..786c86548785 100644 --- a/lib/kv-router/src/services/selection/affinity.rs +++ b/lib/kv-router/src/services/selection/affinity.rs @@ -17,18 +17,43 @@ use std::sync::{Arc, OnceLock, Weak}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; use dashmap::{DashMap, mapref::entry::Entry}; +use serde::{Deserialize, Serialize}; use tokio::sync::Notify; use tokio::sync::futures::OwnedNotified; use tokio::time::Instant; use tokio_util::sync::CancellationToken; pub use crate::protocols::WorkerAffinityTarget as AffinityTarget; -pub use crate::scheduling::config::SessionAffinityMode; pub const MAX_SESSION_AFFINITY_TTL_SECS: u64 = 31_536_000; pub const MAX_SESSION_AFFINITY_ENTRIES: usize = 65_536; pub const MAX_SESSION_AFFINITY_ID_BYTES: usize = 256; +/// How a bound session treats a dispatch that landed elsewhere. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum SessionAffinityMode { + /// The binding is exact: dispatching to another worker or rank is an error. + #[default] + Hard, + /// The binding follows the dispatch: the session rebinds to where it ran. + Soft, +} + +impl std::str::FromStr for SessionAffinityMode { + type Err = String; + + fn from_str(value: &str) -> Result { + match value { + "hard" => Ok(Self::Hard), + "soft" => Ok(Self::Soft), + _ => Err(format!( + "invalid session affinity mode {value:?}; expected 'hard' or 'soft'" + )), + } + } +} + #[derive(Debug, thiserror::Error)] pub enum AffinityError { #[error("{0}")] diff --git a/lib/kv-router/src/services/selection/core/run.rs b/lib/kv-router/src/services/selection/core/run.rs index 8ae5e79eaa08..8517ced5c673 100644 --- a/lib/kv-router/src/services/selection/core/run.rs +++ b/lib/kv-router/src/services/selection/core/run.rs @@ -606,8 +606,8 @@ impl SelectionCore { } /// Bind the held session to `dispatched`. A `Hard` rejection whose bound - /// worker departed is not a client fault: the session re-initializes on - /// the dispatched worker. An overloaded bound worker returns `NotReady`. + /// worker departed after [`Self::hold_session`] checked it is not a client + /// fault: the session is re-initialized on the dispatched worker instead. pub(super) fn commit_session( &self, table: &SessionAffinity, @@ -624,21 +624,6 @@ impl SelectionCore { let departed = bound.is_some_and(|target| !self.catalog.is_schedulable(target.worker_id, key)); if !departed { - // If the bound worker is overloaded, the client is not at fault. - let overloaded = bound.is_some_and(|target| { - self.host - .load - .overloaded_workers - .as_ref() - .and_then(|provider| provider()) - .is_some_and(|ids| ids.contains(&target.worker_id)) - }); - if overloaded { - return Err(SelectionError::NotReady(format!( - "session affinity worker {} is overloaded", - bound.unwrap().worker_id, - ))); - } return Err(affinity_error(error)); } // `commit` invalidated the stale binding. Another request may diff --git a/lib/kv-router/src/services/selection/service.rs b/lib/kv-router/src/services/selection/service.rs index 31206c0be4c6..87932f187619 100644 --- a/lib/kv-router/src/services/selection/service.rs +++ b/lib/kv-router/src/services/selection/service.rs @@ -141,24 +141,6 @@ impl SelectionServiceBuilder { self.kv_router_config .validate_config() .map_err(anyhow::Error::msg)?; - - // Explicit session_affinity() call takes precedence over KvRouterConfig. - let affinity_config: Option = - if let Some(ttl) = self.session_affinity_ttl { - Some( - SessionAffinityConfig::new(ttl) - .with_mode(self.kv_router_config.session_affinity_mode), - ) - } else if let Some(secs) = self.kv_router_config.session_affinity_ttl_secs { - let ttl = Duration::try_from_secs_f64(secs) - .map_err(|_| anyhow::anyhow!("session affinity TTL {secs} out of range"))?; - Some( - SessionAffinityConfig::new(ttl) - .with_mode(self.kv_router_config.session_affinity_mode), - ) - } else { - None - }; let worker_selection_policy_factory = match self.worker_selection_policy_factory { Some(factory) => Some(factory), None => self @@ -205,7 +187,7 @@ impl SelectionServiceBuilder { self.selection_cache, tracking_hash, indexer_policy, - affinity_config, + self.session_affinity_ttl.map(SessionAffinityConfig::new), )); if recover_from_peers { @@ -234,7 +216,7 @@ impl SelectionServiceBuilder { core.dispatch_replica_event(event); } }, - affinity_config.as_ref().map(|_| { + self.session_affinity_ttl.map(|_| { move |event| { if let Some(core) = affinity_core.upgrade() { core.dispatch_affinity_event(event); diff --git a/lib/llm/src/protocols/agents.rs b/lib/llm/src/protocols/agents.rs index 7fd5f45aeee2..998ca96dfbdf 100644 --- a/lib/llm/src/protocols/agents.rs +++ b/lib/llm/src/protocols/agents.rs @@ -122,7 +122,7 @@ fn codex_compaction_header_value(headers: &HeaderMap) -> Option .then(|| metadata.compaction.unwrap_or_default()) } -pub fn session_affinity_header_value(headers: &HeaderMap) -> Option { +pub(crate) fn session_affinity_header_value(headers: &HeaderMap) -> Option { if let Some(session_id) = borrowed_header_value(headers, HEADER_DYNAMO_SESSION_ID) { return Some(session_id.to_owned()); } From 71335cf945ea58005332ea51b7f36fe1c95d2664 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Tue, 15 Sep 2026 18:24:04 +0000 Subject: [PATCH 089/102] fix(router): validate session affinity DP ranks during failover Signed-off-by: Ishan Dhanani --- .../src/services/selection/catalog.rs | 17 +++-- .../src/services/selection/core/run.rs | 8 +-- .../src/services/selection/core/tests.rs | 64 +++++++++++++++++++ 3 files changed, 81 insertions(+), 8 deletions(-) diff --git a/lib/kv-router/src/services/selection/catalog.rs b/lib/kv-router/src/services/selection/catalog.rs index f18692fc5849..11ddaab292e0 100644 --- a/lib/kv-router/src/services/selection/catalog.rs +++ b/lib/kv-router/src/services/selection/catalog.rs @@ -6,7 +6,7 @@ use std::collections::HashMap; use parking_lot::RwLock; use crate::identity::RoutingPartitionId; -use crate::protocols::{WorkerId, WorkerWithDpRank}; +use crate::protocols::{WorkerAffinityTarget, WorkerId, WorkerWithDpRank}; use super::types::{SelectionWorkerConfig, WorkerCatalogRecord, WorkerLifecycle}; @@ -114,11 +114,20 @@ impl WorkerCatalog { .count() } - pub(super) fn is_schedulable(&self, worker_id: WorkerId, key: &RoutingPartitionId) -> bool { + pub(super) fn is_schedulable( + &self, + target: WorkerAffinityTarget, + key: &RoutingPartitionId, + ) -> bool { self.workers .read() - .get(&worker_id) - .is_some_and(|record| schedulable_in(record, key)) + .get(&target.worker_id) + .is_some_and(|record| { + schedulable_in(record, key) + && target + .dp_rank + .is_none_or(|rank| record.dp_ranks().contains(&rank)) + }) } pub(super) fn schedulable_endpoint( diff --git a/lib/kv-router/src/services/selection/core/run.rs b/lib/kv-router/src/services/selection/core/run.rs index 8517ced5c673..d1138959f1ef 100644 --- a/lib/kv-router/src/services/selection/core/run.rs +++ b/lib/kv-router/src/services/selection/core/run.rs @@ -565,7 +565,7 @@ impl SelectionCore { })) } - /// Hold `session_id` for a booking, re-initializing a binding whose worker + /// Hold `session_id` for a booking, re-initializing a binding whose target /// this partition can no longer schedule. `None` when the table is full: a /// router-side limit, not a client fault, so the request routes unpinned. pub(super) async fn hold_session( @@ -583,7 +583,7 @@ impl SelectionCore { }; match acquired { Ok(Hold::Bound { target, mut lease }) - if !self.catalog.is_schedulable(target.worker_id, key) => + if !self.catalog.is_schedulable(target, key) => { tracing::debug!( session_id, @@ -606,7 +606,7 @@ impl SelectionCore { } /// Bind the held session to `dispatched`. A `Hard` rejection whose bound - /// worker departed after [`Self::hold_session`] checked it is not a client + /// worker or rank departed after [`Self::hold_session`] checked it is not a client /// fault: the session is re-initialized on the dispatched worker instead. pub(super) fn commit_session( &self, @@ -622,7 +622,7 @@ impl SelectionCore { Ok(lease) => Ok(Some(lease)), Err(error) => { let departed = - bound.is_some_and(|target| !self.catalog.is_schedulable(target.worker_id, key)); + bound.is_some_and(|target| !self.catalog.is_schedulable(target, key)); if !departed { return Err(affinity_error(error)); } diff --git a/lib/kv-router/src/services/selection/core/tests.rs b/lib/kv-router/src/services/selection/core/tests.rs index d2da09d9baa4..a8a9ea1747b0 100644 --- a/lib/kv-router/src/services/selection/core/tests.rs +++ b/lib/kv-router/src/services/selection/core/tests.rs @@ -2101,6 +2101,70 @@ async fn departed_session_worker_reinitializes_the_session() { assert_eq!(bound_worker(&core, "s"), Some(second.worker_id)); } +#[rstest::rstest] +#[case::removed_rank(Some(1))] +#[case::retained_rank(Some(0))] +#[case::worker_only(None)] +#[tokio::test] +async fn session_dp_rank_shrink_rebinds_only_removed_targets( + #[case] dp_rank: Option, + #[values(false, true)] after_hold: bool, +) { + let core = core_with_session_affinity(); + let key = default_key(); + core.upsert_worker(WorkerRequest { + data_parallel_size: Some(2), + ..worker(1) + }) + .await + .expect("worker upsert"); + let entry = core.entry(&key).expect("default partition"); + let table = entry.affinity.get().expect("affinity table"); + let target = WorkerAffinityTarget::new(1, dp_rank); + drop( + table + .commit(table.acquire("s", None).await.expect("hold"), target) + .expect("bind session"), + ); + + let mut request = Box::pin(core.select_and_reserve(session_reservation("r1", "s"))); + let mut context = std::task::Context::from_waker(std::task::Waker::noop()); + if after_hold { + // Take the hold, then shrink before the scheduler actor can run. + assert!(request.as_mut().poll(&mut context).is_pending()); + let mut record = core.catalog.get(1).expect("worker"); + record.data_parallel_size = Some(1); + core.catalog.replace(record); + core.publish_scheduler_config(&key); + } else { + core.patch_worker( + 1, + serde_json::from_value(serde_json::json!({"data_parallel_size": 1})) + .expect("worker patch"), + ) + .await + .expect("shrink worker"); + assert!(request.as_mut().poll(&mut context).is_pending()); + assert_eq!( + table.query_target("s", None).expect("query"), + (dp_rank != Some(1)).then_some(target), + "only a removed target must be invalidated before scheduling", + ); + } + + let response = request.await.expect("rank removal is not a client fault"); + assert_eq!((response.worker_id, response.dp_rank), (1, 0)); + assert_eq!( + table.query_target("s", None).expect("query"), + Some(WorkerAffinityTarget::new(1, dp_rank.map(|_| 0))), + ); + core.free_reservation("r1").await.expect("free booking"); + wait_until("session lease release", || { + lease_count(&core, "s") == Some(0) + }) + .await; +} + #[tokio::test] async fn session_worker_departing_after_the_hold_reinitializes_the_session() { let (core, first) = bound_session(SessionAffinityMode::Hard).await; From 54c2076e08ac4f98d1c97f1d638c3fe1fbd44bbf Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Tue, 15 Sep 2026 18:36:17 +0000 Subject: [PATCH 090/102] perf(router): keep LoRA candidate filtering linear Signed-off-by: Ishan Dhanani --- lib/kv-router/src/scheduling/lora_filter.rs | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/lib/kv-router/src/scheduling/lora_filter.rs b/lib/kv-router/src/scheduling/lora_filter.rs index 74a3ef76c81a..03dba81ad6e8 100644 --- a/lib/kv-router/src/scheduling/lora_filter.rs +++ b/lib/kv-router/src/scheduling/lora_filter.rs @@ -5,6 +5,8 @@ use std::collections::HashSet; +use rustc_hash::FxHashSet; + use crate::protocols::{WorkerId, WorkerWithDpRank}; /// Host-owned knowledge of which workers can serve a LoRA adapter. @@ -46,10 +48,14 @@ pub fn narrow_allowed_worker_ids_by_lora( return allowed_worker_ids; } // A filter may only narrow: ids it returns from outside `base` are dropped. - let mut narrowed: HashSet = filter + let preferred: FxHashSet = filter .filter_worker_ids_for_lora(lora_name, &base) .into_iter() - .filter(|id| base.contains(id)) + .collect(); + let mut narrowed: HashSet = base + .iter() + .copied() + .filter(|id| preferred.contains(id)) .collect(); if let Some(pinned) = pinned_worker && base.contains(&pinned.worker_id) @@ -157,6 +163,14 @@ mod tests { ); } + #[test] + fn all_workers_qualifying_preserves_the_candidate_universe() { + assert_eq!( + narrow(&Widening(vec![3, 2, 1, 2, 9]), Some("a"), None, None), + Some(HashSet::from([1, 2, 3])) + ); + } + #[test] fn empty_universe_is_a_noop() { assert_eq!( From 800070df871fec70a3b5187338eab7927b991797 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Tue, 15 Sep 2026 18:36:17 +0000 Subject: [PATCH 091/102] perf(router): move the committed booking request ID Signed-off-by: Ishan Dhanani --- lib/kv-router/src/scheduling/queue.rs | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/lib/kv-router/src/scheduling/queue.rs b/lib/kv-router/src/scheduling/queue.rs index 7d8a74e6eb38..5d53b2283471 100644 --- a/lib/kv-router/src/scheduling/queue.rs +++ b/lib/kv-router/src/scheduling/queue.rs @@ -387,7 +387,11 @@ impl BookingHandle { #[must_use] pub fn commit(mut self) -> SchedulerBookingDescriptor { self.armed = false; - self.booking.clone() + SchedulerBookingDescriptor { + request_id: std::mem::take(&mut self.booking.request_id), + worker: self.booking.worker, + attempt_id: self.booking.attempt_id, + } } /// Free the booking now and wait for the scheduler to acknowledge it. @@ -1818,8 +1822,11 @@ mod tests { wait_freed(&slots, &armed).await; let committed = book_directly(&slots, "committed"); - let descriptor = queue.booking_handle(committed.clone()).commit(); + let handle = queue.booking_handle(committed.clone()); + let request_id_ptr = handle.booking.request_id.as_ptr(); + let descriptor = handle.commit(); assert_eq!(descriptor, committed); + assert_eq!(descriptor.request_id.as_ptr(), request_id_ptr); // The cleanup queue drains in order, so an acknowledged release // enqueued after the commit proves the commit enqueued nothing. From 928321e803f86390cbf96f5c26147b78e4dbbf9b Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Tue, 15 Sep 2026 18:36:17 +0000 Subject: [PATCH 092/102] docs(router): clarify default lookup retention Signed-off-by: Ishan Dhanani --- lib/kv-router/src/services/indexer/lookup.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/lib/kv-router/src/services/indexer/lookup.rs b/lib/kv-router/src/services/indexer/lookup.rs index ae439b5f1ef0..32d9742f2520 100644 --- a/lib/kv-router/src/services/indexer/lookup.rs +++ b/lib/kv-router/src/services/indexer/lookup.rs @@ -114,9 +114,9 @@ impl Indexer { .await } - /// Device match details plus per-tier hits. Router-hint chain retention is - /// honored as requested; see [`Indexer::find_tiered_matches_with_options`] - /// for the capability-gated variant. + /// Device match details plus per-tier hits, with router-hint chain retention + /// disabled. Use [`Indexer::find_tiered_matches_with_options`] to request + /// retention through the capability-gated variant. pub async fn find_matches_by_tier( &self, sequence: Vec, From 0a477b17629b543292fa1f7125e9d8d77f3d8607 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Tue, 15 Sep 2026 18:44:56 +0000 Subject: [PATCH 093/102] fix(router): preserve affinity ownership across concurrent failover Signed-off-by: Ishan Dhanani --- .../src/services/selection/affinity.rs | 232 +++++++++++------- .../src/services/selection/core/run.rs | 21 +- .../src/services/selection/core/tests.rs | 93 ++++++- 3 files changed, 242 insertions(+), 104 deletions(-) diff --git a/lib/kv-router/src/services/selection/affinity.rs b/lib/kv-router/src/services/selection/affinity.rs index 786c86548785..f4d822fdb619 100644 --- a/lib/kv-router/src/services/selection/affinity.rs +++ b/lib/kv-router/src/services/selection/affinity.rs @@ -71,15 +71,6 @@ pub struct AffinityVersion { pub writer_id: u64, } -impl AffinityVersion { - /// Version of a lease joined to an initialization that has not committed - /// yet; it matches no stored binding, so such a lease can only release. - pub const PENDING: Self = Self { - sequence: u64::MAX, - writer_id: u64::MAX, - }; -} - /// Receives every binding this table publishes for its replicas. pub trait AffinityReplicaSink: Send + Sync { fn publish(&self, session_id: &str, target: AffinityTarget, version: AffinityVersion); @@ -101,9 +92,6 @@ enum AffinityEntry { Initializing { revision: u64, notify: Arc, - /// Leases handed out by [`SessionAffinity::join_initializing`]; they - /// count as active once the initializer commits. - pending_leases: usize, }, Bound { target: AffinityTarget, @@ -345,7 +333,6 @@ impl SessionAffinity { entry.insert(AffinityEntry::Initializing { revision, notify: notify.clone(), - pending_leases: 0, }); Ok(AcquireStep::Held(Hold::Initialize( AffinityInitialization { @@ -382,7 +369,6 @@ impl SessionAffinity { *entry.get_mut() = AffinityEntry::Initializing { revision, notify: notify.clone(), - pending_leases: 0, }; Ok(AcquireStep::Held(Hold::Initialize( AffinityInitialization { @@ -441,29 +427,41 @@ impl SessionAffinity { } } - /// Take a lease on the binding another request is initializing, so this - /// request's use counts once that commit lands and the binding cannot - /// idle out under it. `None` when the entry is not initializing (the - /// caller falls back to [`Self::try_acquire`]). If the initializer drops - /// uncommitted, the lease releases as a no-op. - pub fn join_initializing(&self, session_id: &str) -> Option { + /// Finish another request's initialization with an already selected target. + /// The first successful booking owns a bound lease immediately, independent + /// of whether the original initializer later commits or is cancelled. + /// `None` means the caller must re-acquire the entry's current state. + pub fn join_initializing( + &self, + session_id: &str, + target: AffinityTarget, + ) -> Option { let mut entry = self.inner.entries.get_mut(session_id)?; - let AffinityEntry::Initializing { - revision, - pending_leases, - .. - } = entry.value_mut() - else { + let AffinityEntry::Initializing { revision, notify } = entry.value() else { return None; }; - *pending_leases += 1; - Some(AffinityLease { + let revision = *revision; + let notify = Arc::clone(notify); + let version = self.inner.next_version(); + *entry = AffinityEntry::Bound { + target, + revision, + version, + // Reserve the initializer's use until it commits or cancels. + active_leases: 2, + idle_deadline: Instant::now() + self.inner.ttl, + }; + drop(entry); + notify.notify_waiters(); + let lease = AffinityLease { coordinator: Arc::downgrade(&self.inner), session_id: session_id.to_string(), - revision: *revision, - version: AffinityVersion::PENDING, + revision, + version, active: true, - }) + }; + lease.publish(target); + Some(lease) } /// Bind the held session to `dispatched`. A `Bound` session dispatched @@ -475,11 +473,7 @@ impl SessionAffinity { dispatched: AffinityTarget, ) -> Result { match hold { - Hold::Initialize(initialization) => { - let lease = initialization.commit(dispatched)?; - lease.publish(dispatched); - Ok(lease) - } + Hold::Initialize(initialization) => initialization.commit(dispatched), Hold::Bound { target, mut lease } => { if self.inner.mode == SessionAffinityMode::Soft { let rebound = AffinityLease::rebound_target(target, dispatched); @@ -602,12 +596,11 @@ impl SessionAffinity { self.inner.next_version() } - /// Leases on `session_id`: pending joiners while initializing, active - /// leases once bound. `None` when there is no entry. + /// Active leases on `session_id`; an unfinished initialization has none. pub fn lease_count(&self, session_id: &str) -> Option { let entry = self.inner.entries.get(session_id)?; Some(match entry.value() { - AffinityEntry::Initializing { pending_leases, .. } => *pending_leases, + AffinityEntry::Initializing { .. } => 0, AffinityEntry::Bound { active_leases, .. } => *active_leases, }) } @@ -717,7 +710,7 @@ impl Inner { } /// A session this request is the first to bind. Dropping it uncommitted -/// releases the slot and wakes waiters so they re-acquire. +/// releases the slot and wakes waiters, unless another request already bound it. pub struct AffinityInitialization { coordinator: Weak, session_id: String, @@ -739,12 +732,34 @@ impl AffinityInitialization { "session affinity initialization was cancelled".to_string(), )); }; - let AffinityEntry::Initializing { + if let AffinityEntry::Bound { + target: bound, revision, - pending_leases, + version, .. - } = entry.value() - else { + } = entry.value_mut() + && *revision == self.revision + { + // A successful joiner already owns this binding. A conflicting + // initializer must release its booking without invalidating it. + if inner.mode == SessionAffinityMode::Hard { + validate_dispatch_target(&self.session_id, *bound, target)?; + } + let hold = Hold::Bound { + target: *bound, + lease: AffinityLease { + coordinator: Arc::downgrade(&inner), + session_id: self.session_id.clone(), + revision: self.revision, + version: *version, + active: true, + }, + }; + drop(entry); + self.active = false; + return SessionAffinity { inner }.commit(hold, target); + } + let AffinityEntry::Initializing { revision, .. } = entry.value() else { return Err(AffinityError::InvalidArgument( "session affinity initialization changed".to_string(), )); @@ -759,20 +774,21 @@ impl AffinityInitialization { target, revision: self.revision, version, - // Joined leases become live uses of the binding they waited for. - active_leases: 1 + pending_leases, + active_leases: 1, idle_deadline: Instant::now() + inner.ttl, }; drop(entry); self.active = false; self.notify.notify_waiters(); - Ok(AffinityLease { + let lease = AffinityLease { coordinator: Arc::downgrade(&inner), session_id: self.session_id.clone(), revision: self.revision, version, active: true, - }) + }; + lease.publish(target); + Ok(lease) } } @@ -784,11 +800,22 @@ impl Drop for AffinityInitialization { let Some(inner) = self.coordinator.upgrade() else { return; }; - let removed = inner.entries.remove_if(&self.session_id, |_, entry| { - matches!( - entry, - AffinityEntry::Initializing { revision, .. } if *revision == self.revision - ) + let removed = inner.entries.remove_if_mut(&self.session_id, |_, entry| { + match entry { + AffinityEntry::Initializing { revision, .. } => *revision == self.revision, + AffinityEntry::Bound { + revision, + active_leases, + .. + } => { + if *revision == self.revision { + // A joiner reserved this use. Cancellation never + // invalidates its binding or refreshes a newer version. + *active_leases -= 1; + } + false + } + } }); if removed.is_some() { inner.entry_count.fetch_sub(1, Ordering::Relaxed); @@ -870,22 +897,6 @@ impl AffinityLease { let Some(mut entry) = inner.entries.get_mut(&self.session_id) else { return; }; - // A joined lease released before its initializer commits must not - // be counted by that commit, or the binding never idles out. - if let AffinityEntry::Initializing { - revision, - pending_leases, - .. - } = entry.value_mut() - { - if *revision == self.revision - && self.version == AffinityVersion::PENDING - && *pending_leases > 0 - { - *pending_leases -= 1; - } - return; - } let AffinityEntry::Bound { target, revision, @@ -900,9 +911,7 @@ impl AffinityLease { return; } *active_leases -= 1; - // A joined lease was a use of whatever binding the initializer - // committed, so it refreshes that binding's idle deadline. - if *version != self.version && self.version != AffinityVersion::PENDING { + if *version != self.version { return; } *idle_deadline = Instant::now() + inner.ttl; @@ -1043,18 +1052,21 @@ mod tests { let target = AffinityTarget::new(2, Some(0)); let init = initialize(&table); let joined = table - .join_initializing("s") + .join_initializing("s", target) .expect("join an initializing session"); - assert_eq!(table.lease_count("s"), Some(1)); + assert_eq!(table.lease_count("s"), Some(2)); // The joiner finishes while the initializer is still queued. drop(joined); assert_eq!( table.lease_count("s"), - Some(0), - "early release uncounts the join" + Some(1), + "early release leaves the initializer's use" ); + tokio::time::advance(TTL + Duration::from_secs(1)).await; + assert_eq!(table.query_target("s", None).expect("query"), Some(target)); + let lease = init.commit(target).expect("commit"); assert_eq!( table.lease_count("s"), @@ -1076,12 +1088,18 @@ mod tests { let table = table(); let target = AffinityTarget::new(2, Some(0)); let init = initialize(&table); + let AcquireStep::Wait(waiter) = table.try_acquire("s", None).expect("waiter") else { + panic!("initialization must have a waiter"); + }; let joined = table - .join_initializing("s") + .join_initializing("s", target) .expect("join an initializing session"); - assert_eq!(table.lease_count("s"), Some(1)); + tokio::time::timeout(Duration::from_millis(1), waiter) + .await + .expect("joining must wake existing waiters"); + assert_eq!(table.lease_count("s"), Some(2)); assert!( - table.join_initializing("x").is_none(), + table.join_initializing("x", target).is_none(), "no join without an initializer" ); @@ -1114,25 +1132,25 @@ mod tests { } #[tokio::test(start_paused = true)] - async fn joined_lease_is_a_noop_after_the_initializer_drops() { + async fn joined_binding_survives_the_initializer_dropping() { let table = table(); let init = initialize(&table); + let target = AffinityTarget::new(2, Some(0)); let joined = table - .join_initializing("s") + .join_initializing("s", target) .expect("join an initializing session"); assert_eq!(table.entry_count(), 1); drop(init); - assert_eq!( - table.entry_count(), - 0, - "a dropped initializer frees the slot" - ); - assert_eq!(table.lease_count("s"), None); + assert_eq!(table.entry_count(), 1); + assert_eq!(table.query_target("s", None).expect("query"), Some(target)); + assert_eq!(table.lease_count("s"), Some(1)); drop(joined); - assert_eq!(table.entry_count(), 0); + assert_eq!(table.lease_count("s"), Some(0)); + tokio::time::advance(TTL + Duration::from_secs(1)).await; + assert_eq!(table.query_target("s", None).expect("query"), None); - // The session is re-initializable and the stale joiner left no count. + // Once the completed booking's TTL expires, another target can bind. let init = initialize(&table); let lease = init .commit(AffinityTarget::new(3, Some(0))) @@ -1141,4 +1159,40 @@ mod tests { drop(lease); assert_eq!(table.lease_count("s"), Some(0)); } + + #[tokio::test] + async fn initializer_cannot_change_the_joined_rank_in_hard_mode() { + let table = table(); + let init = initialize(&table); + let target = AffinityTarget::new(2, Some(0)); + let joined = table.join_initializing("s", target).expect("join"); + assert!(init.commit(AffinityTarget::new(2, Some(1))).is_err()); + assert_eq!(table.query_target("s", None).expect("query"), Some(target)); + assert_eq!(table.lease_count("s"), Some(1)); + drop(joined); + assert_eq!(table.lease_count("s"), Some(0)); + } + + #[tokio::test] + async fn soft_initializer_follows_dispatch_and_keeps_worker_only_affinity() { + let table = SessionAffinity::with_config( + SessionAffinityConfig::new(TTL).with_mode(SessionAffinityMode::Soft), + ) + .expect("table"); + let init = initialize(&table); + let joined = table + .join_initializing("s", AffinityTarget::new(2, None)) + .expect("join"); + let lease = init + .commit(AffinityTarget::new(3, Some(0))) + .expect("commit"); + assert_eq!( + table.query_target("s", None).expect("query"), + Some(AffinityTarget::new(3, None)) + ); + assert_eq!(table.lease_count("s"), Some(2)); + drop(lease); + drop(joined); + assert_eq!(table.lease_count("s"), Some(0)); + } } diff --git a/lib/kv-router/src/services/selection/core/run.rs b/lib/kv-router/src/services/selection/core/run.rs index d1138959f1ef..3190e23caf79 100644 --- a/lib/kv-router/src/services/selection/core/run.rs +++ b/lib/kv-router/src/services/selection/core/run.rs @@ -24,10 +24,10 @@ //! the binding is committed while the booking is held. The commit never //! waits: `commit_session` is a synchronous fn that re-binds only through //! `try_acquire`. A commit that finds another request initializing the -//! session joins that initialization instead of waiting: it takes a lease -//! that counts once the initializer commits +//! session completes that initialization with its selected target and takes +//! a bound lease immediately, even if the initializer is later cancelled //! (`failover_commit_behind_an_initializing_hold_keeps_a_lease`); a joiner -//! that finishes first is uncounted again +//! that finishes first releases its lease normally //! (`joined_lease_released_before_the_commit_is_not_counted`). When the //! table is full, or the join misses twice, it returns `Ok(None)` and //! the request routes unpinned rather than waiting on a request queued @@ -42,6 +42,7 @@ //! (`dropped_selection_future_frees_its_booking`, //! `dropped_book_selection_during_routing_record_frees_booking_and_claim`). +use super::super::affinity::{SessionAffinityMode, validate_dispatch_target}; use super::hint::{hint_capable_partition, transfer_hint_for_selection}; /// `try_acquire` then `join_initializing` attempts before a commit that finds @@ -628,20 +629,28 @@ impl SelectionCore { } // `commit` invalidated the stale binding. Another request may // already be initializing the replacement: join it rather - // than wait, so this booking holds a lease on the new - // binding. The join misses only if that initialization + // than wait, binding it to this successful booking's target. + // The join misses only if that initialization // resolved between the two calls; one retry covers that, and // a booking never waits on a request queued behind it. for _ in 0..JOIN_ATTEMPTS { match table.try_acquire(session_id, None) { Ok(AcquireStep::Held(hold)) => { + if table.mode() == SessionAffinityMode::Hard + && let Some(target) = hold.target() + { + // A competing failover already bound a valid + // target; reject a mismatch without erasing it. + validate_dispatch_target(session_id, target, dispatched) + .map_err(affinity_error)?; + } return table .commit(hold, dispatched) .map(Some) .map_err(affinity_error); } Ok(AcquireStep::Wait(_)) => { - if let Some(lease) = table.join_initializing(session_id) { + if let Some(lease) = table.join_initializing(session_id, dispatched) { return Ok(Some(lease)); } } diff --git a/lib/kv-router/src/services/selection/core/tests.rs b/lib/kv-router/src/services/selection/core/tests.rs index a8a9ea1747b0..bc4c3c729aa1 100644 --- a/lib/kv-router/src/services/selection/core/tests.rs +++ b/lib/kv-router/src/services/selection/core/tests.rs @@ -2245,13 +2245,9 @@ fn lease_count(core: &SelectionCore, session_id: &str) -> Option { table.lease_count(session_id) } -/// Two workers, so both requests must land on the replacement; with more -/// workers the joiner could book elsewhere while still counting toward the -/// session's binding, which is the same split a plain `Bound` lease has. /// r2 holds the departed binding when r3 re-initializes the session; r2's -/// commit then finds r3's initialization in progress and must join it, not -/// route without a lease, so the replacement binding cannot idle out while -/// r2 is still active. +/// successful booking binds the replacement immediately. r3 can later join +/// the same target, and the binding stays alive until both requests finish. #[tokio::test] async fn failover_commit_behind_an_initializing_hold_keeps_a_lease() { let (core, first) = bound_session(SessionAffinityMode::Hard).await; @@ -2279,11 +2275,11 @@ async fn failover_commit_behind_an_initializing_hold_keeps_a_lease() { .await .expect("a departure after the hold is not a client fault"); assert_eq!(r2.worker_id, replacement); - assert_eq!(bound_worker(&core, "s"), None, "r3 has not committed yet"); + assert_eq!(bound_worker(&core, "s"), Some(replacement)); assert_eq!( lease_count(&core, "s"), - Some(1), - "r2 must hold a pending lease on r3's initialization" + Some(2), + "r2 and the queued initializer both keep the binding alive" ); let r3 = r3.await.expect("r3 binds the replacement"); @@ -2304,6 +2300,85 @@ async fn failover_commit_behind_an_initializing_hold_keeps_a_lease() { assert_eq!(bound_worker(&core, "s"), None); } +#[rstest::rstest] +#[case::initializer_cancelled(false)] +#[case::initializer_selects_another_worker(true)] +#[tokio::test] +async fn joined_failover_preserves_the_successful_bookings_binding(#[case] commit_other: bool) { + let (core, first) = bound_session(SessionAffinityMode::Hard).await; + let key = default_key(); + let replacement = if first.worker_id == 1 { 2 } else { 1 }; + core.upsert_worker(worker(3)).await.expect("third worker"); + let mut context = std::task::Context::from_waker(std::task::Waker::noop()); + + let mut r2_request = session_reservation("r2", "s"); + r2_request.allowed_worker_ids = Some(HashSet::from([replacement])); + let mut r2 = Box::pin(core.select_and_reserve(r2_request)); + assert!(r2.as_mut().poll(&mut context).is_pending()); + core.catalog + .set_lifecycle(first.worker_id, WorkerLifecycle::Draining, Vec::new()); + core.publish_scheduler_config(&key); + + let mut r3_request = session_reservation("r3", "s"); + r3_request.allowed_worker_ids = Some(HashSet::from([3])); + let mut r3 = Box::pin(core.select_and_reserve(r3_request)); + assert!(r3.as_mut().poll(&mut context).is_pending()); + assert_eq!(bound_worker(&core, "s"), None); + + let r2 = r2.await.expect("successful failover booking"); + assert_eq!(r2.worker_id, replacement); + if commit_other { + assert!( + matches!(r3.await, Err(SelectionError::BadRequest(_))), + "Hard affinity must reject the conflicting initializer" + ); + } else { + drop(r3); + } + assert_eq!(bound_worker(&core, "s"), Some(replacement)); + assert_eq!(lease_count(&core, "s"), Some(1)); + assert!(!core.reservation_index.read().contains_key("r3")); + let entry = core.entry(&key).expect("entry"); + wait_until("initializer booking rollback", || { + !entry.scheduler.has_request("r3") + }) + .await; + core.free_reservation("r2").await.expect("free r2"); + wait_until("successful booking release", || { + lease_count(&core, "s") == Some(0) + }) + .await; +} + +#[tokio::test] +async fn late_failover_mismatch_preserves_the_committed_replacement() { + let (core, first) = bound_session(SessionAffinityMode::Hard).await; + core.upsert_worker(worker(3)).await.expect("third worker"); + let key = default_key(); + let entry = core.entry(&key).expect("entry"); + let table = entry.affinity.get().expect("affinity table"); + let old_hold = core.hold_session(table, "s", &key).await.unwrap().unwrap(); + core.catalog + .set_lifecycle(first.worker_id, WorkerLifecycle::Draining, Vec::new()); + let replacement = if first.worker_id == 1 { 2 } else { 1 }; + let new_hold = core.hold_session(table, "s", &key).await.unwrap().unwrap(); + let lease = core + .commit_session( + table, + new_hold, + "s", + WorkerWithDpRank::new(replacement, 0), + &key, + ) + .expect("bind replacement"); + let result = core.commit_session(table, old_hold, "s", WorkerWithDpRank::new(3, 0), &key); + assert!(matches!(result, Err(SelectionError::BadRequest(_)))); + assert_eq!(bound_worker(&core, "s"), Some(replacement)); + assert_eq!(lease_count(&core, "s"), Some(1)); + drop(lease); + assert_eq!(lease_count(&core, "s"), Some(0)); +} + #[tokio::test] async fn two_phase_reservation_binds_the_session() { let core = core_with_session_affinity(); From 0049ca394000de62f54aaf662f1c8c655a5fe656 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Tue, 15 Sep 2026 18:51:01 +0000 Subject: [PATCH 094/102] perf(router): omit unused metadata from frontend leases Signed-off-by: Ishan Dhanani --- .../src/services/selection/core/operation.rs | 7 ++-- .../src/services/selection/core/run.rs | 36 +++++++++++-------- .../src/services/selection/core/tests.rs | 12 +++++-- 3 files changed, 33 insertions(+), 22 deletions(-) diff --git a/lib/kv-router/src/services/selection/core/operation.rs b/lib/kv-router/src/services/selection/core/operation.rs index bf8929b78c99..2817ee94f1ef 100644 --- a/lib/kv-router/src/services/selection/core/operation.rs +++ b/lib/kv-router/src/services/selection/core/operation.rs @@ -127,11 +127,10 @@ pub struct Selected { pub key: RoutingPartitionId, pub response: SchedulingResponse, pub advisory_load: Option, - /// The chosen worker's KV capacity as known when it was selected. + /// The chosen worker's KV capacity for wire responses; omitted for `Lease`. pub total_kv_blocks: Option, - /// `None` only for `Lease` admission when the worker left the catalog - /// after it was booked: the host dispatches by worker id and lets the - /// transport report the departure, which is what drives migration. + /// The endpoint for wire responses; omitted for `Lease`. The embedding + /// host dispatches by worker id and lets transport report departures. pub endpoint: Option, pub block_size: u32, pub isl_tokens: usize, diff --git a/lib/kv-router/src/services/selection/core/run.rs b/lib/kv-router/src/services/selection/core/run.rs index 3190e23caf79..2054852f813a 100644 --- a/lib/kv-router/src/services/selection/core/run.rs +++ b/lib/kv-router/src/services/selection/core/run.rs @@ -464,22 +464,28 @@ impl SelectionCore { } Err(error) => return Err(error.into()), }; - let endpoint = self - .catalog - .schedulable_endpoint(response.best_worker.worker_id, &key); - if endpoint.is_none() && !matches!(admission, SelectionAdmission::Lease { .. }) { - return Err(SelectionError::Internal(format!( - "selected worker {} is no longer schedulable", - response.best_worker.worker_id - ))); - } + let (endpoint, total_kv_blocks) = if matches!(admission, SelectionAdmission::Lease { .. }) { + // The embedding host dispatches by worker id and discards wire metadata. + (None, None) + } else { + let endpoint = self + .catalog + .schedulable_endpoint(response.best_worker.worker_id, &key) + .ok_or_else(|| { + SelectionError::Internal(format!( + "selected worker {} is no longer schedulable", + response.best_worker.worker_id + )) + })?; + let total_kv_blocks = advisory_load + .and_then(|load| load.total_kv_blocks.map(|blocks| blocks as u64)) + .or_else(|| { + self.catalog + .total_kv_blocks(response.best_worker.worker_id, &key) + }); + (Some(endpoint), total_kv_blocks) + }; let effective_prefill = effective_prefill_tokens(isl_tokens, response.cached_tokens); - let total_kv_blocks = advisory_load - .and_then(|load| load.total_kv_blocks.map(|blocks| blocks as u64)) - .or_else(|| { - self.catalog - .total_kv_blocks(response.best_worker.worker_id, &key) - }); let kv_hint = if retain_kv_transfer_chain && response.kv_transfer_candidates.is_some() { transfer_hint_for_selection( &entry.workers_tx.borrow(), diff --git a/lib/kv-router/src/services/selection/core/tests.rs b/lib/kv-router/src/services/selection/core/tests.rs index bc4c3c729aa1..5597a2861940 100644 --- a/lib/kv-router/src/services/selection/core/tests.rs +++ b/lib/kv-router/src/services/selection/core/tests.rs @@ -1253,8 +1253,7 @@ async fn booking_is_freed_when_selected_worker_drained_while_queued() { /// The embedded host dispatches by worker id, so a worker that drains /// while the request queues still comes back selected: the transport -/// reports the departure and migration takes over. The endpoint is the -/// only thing missing. +/// reports the departure and migration takes over. #[tokio::test] async fn lease_admission_keeps_a_selection_whose_worker_drained_while_queued() { let core = saturated_core(); @@ -1302,7 +1301,12 @@ async fn lease_admission_keeps_a_selection_whose_worker_drained_while_queued() { #[tokio::test] async fn lease_admission_installs_no_index_row_and_records_nothing() { let core = local_core(test_config(false)); - core.upsert_worker(worker(1)).await.expect("worker upsert"); + core.upsert_worker(WorkerRequest { + total_kv_blocks: Some(2048), + ..worker(1) + }) + .await + .expect("worker upsert"); let key = default_key(); let entry = core.entry(&key).expect("entry"); @@ -1324,6 +1328,8 @@ async fn lease_admission_installs_no_index_row_and_records_nothing() { ); assert!(selected.sequence_hashes.is_none()); assert!(selected.routing_hashes.is_none()); + assert!(selected.endpoint.is_none()); + assert!(selected.total_kv_blocks.is_none()); // Nothing was recorded for the leased prompt: once the indexer has // applied everything enqueued so far, a query restricted to the booked // worker sees no cached blocks. From c71f3d3d113d18da9f325774be933781ea523d2a Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Tue, 15 Sep 2026 21:24:55 +0000 Subject: [PATCH 095/102] refactor(router): defer standalone HTTP indexer integration Signed-off-by: Ishan Dhanani --- .../router/standalone-indexer.md | 26 -- .../router/standalone-selection.md | 9 +- lib/bindings/python/rust/llm/kv.rs | 20 +- lib/bindings/python/src/dynamo/_core.pyi | 5 - .../src/indexer/lower_tier_indexers.rs | 29 ++ lib/kv-router/src/services/indexer/backend.rs | 289 +++--------------- .../src/services/indexer/registry.rs | 2 - lib/kv-router/src/services/indexer/server.rs | 231 +------------- .../src/services/selection/README.md | 6 +- .../src/services/selection/core/mod.rs | 7 +- .../src/services/selection/core/tests.rs | 130 ++++---- .../src/services/selection/core/workers.rs | 24 +- .../src/services/selection/service.rs | 17 +- lib/kv-router/src/services/selection/tests.rs | 1 - 14 files changed, 141 insertions(+), 655 deletions(-) diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-indexer.md b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-indexer.md index 09901a69b1d1..264715c360ba 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-indexer.md +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-indexer.md @@ -403,32 +403,6 @@ after they have been computed. Callers must precompute these hashes with the int and omit `cache_salt` from `/query_by_hash`. Use `/query` when the indexer should compute salted hashes from tokens server-side. -### `POST /query_tiered_by_hash` — Lossless tiered matches for a remote-primary selector - -```bash -curl -X POST http://localhost:8090/query_tiered_by_hash \ - -H 'Content-Type: application/json' \ - -d '{"block_hashes": [123456, 789012], "model_name": "llama-3-8b", "routing_group": "default"}' -``` - -Returns the indexer's tiered match details in their wire form rather than the -Mooncake score summary, so a caller can run its own overlap analysis: - -```json -{ - "block_size": 16, - "tiered": { - "device": { "scores": [[{"worker_id": 7, "dp_rank": 0}, 2]], "frequencies": [] }, - "lower_tier": [["host_pinned", { "hits": [[{"worker_id": 7, "dp_rank": 0}, 1]] }]] - } -} -``` - -Scores and hits are in blocks. This is the endpoint a -[standalone selection service](standalone-selection.md) configured with -`--remote-indexer-url` uses for its primary lookups. Both `model_name` and -`routing_group` are required; an unknown partition returns `404`. - ### Per-instance tier breakdown Each entry in `instances` is keyed by `instance_id` (as a string) and reports prefix reach across the device, host-pinned, and disk storage tiers: diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-selection.md b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-selection.md index aef639e67fcc..b0438fe23dbc 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-selection.md +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-selection.md @@ -65,8 +65,7 @@ APIs. Those bindings should wrap `SelectionService` rather than construct |------|---------|-------------| | `--port` | `8092` | HTTP server port. | | `--threads` | `4` | KV indexer worker threads. | -| `--indexer-peers` | none | Comma-separated HTTP URLs used for startup KV recovery through `/dump`. Ignored with `--remote-indexer-url`. | -| `--remote-indexer-url` | none | Base URL of a [standalone indexer](standalone-indexer.md) that serves the primary KV index. The selector then does not subscribe to worker KV events; workers publish to the indexer instead. Requires `use_kv_events=true`. | +| `--indexer-peers` | none | Comma-separated HTTP URLs used for startup KV recovery through `/dump`. | | `--replica-sync-port` | none | Local ZMQ PUB port for active-load lifecycle events. The selector binds `tcp://*:` internally. | | `--replica-sync-peers` | none | Comma-separated ZMQ PUB endpoints for selector peers. Requires `--replica-sync-port`. | | `--session-affinity-ttl-secs` | none | Pin each request `session_id` to the worker that served it for this long after its last request (1 to 31,536,000 seconds). When `--replica-sync-port` is configured, bindings replicate to peers over the replica mesh (`dynamo.session-affinity.v1` topic). | @@ -91,12 +90,6 @@ frontend uses, so the two index the same way: | `use_kv_events=true` (default) | Event-driven from worker ZMQ events | Not recorded | | `use_kv_events=true`, `DYN_ROUTER_PREDICTED_TTL_SECS` set | Event-driven | Recorded into a short-TTL side indexer merged into device scores by per-worker max | | `use_kv_events=false` | Approximate: no events, entries expire after `router_ttl_secs` | Recorded into the primary | -| `--remote-indexer-url` (or `SelectionServiceBuilder::kv_index(KvIndexSource::Remote(..))`) | Served by a standalone indexer, queried through `POST /query_tiered_by_hash` | Recorded into the side indexer when `DYN_ROUTER_PREDICTED_TTL_SECS` is set, else not recorded | - -With a remote primary, worker records do not need `kv_events_endpoint(s)` to -become schedulable, no ZMQ listener is started in the selector, and `/dump` -returns no primary events (peers recover from the indexer instead). A remote -lookup failure is reported as an offline indexer for that selection. A routing decision is recorded when a reservation is booked (`select_and_reserve`, or `select` followed by `POST /reservations`), never for a query-only `select`. diff --git a/lib/bindings/python/rust/llm/kv.rs b/lib/bindings/python/rust/llm/kv.rs index c0503466d2ef..50e9575a608a 100644 --- a/lib/bindings/python/rust/llm/kv.rs +++ b/lib/bindings/python/rust/llm/kv.rs @@ -30,10 +30,10 @@ use dynamo_kv_router::protocols::*; use dynamo_kv_router::services::indexer::{self, IndexerConfig}; #[cfg(feature = "select-service")] use dynamo_kv_router::services::selection::{ - self, KvIndexSource, OverlapScoresRequest, PotentialLoadsRequest, ReservationRequest, - SelectAndReserveRequest, SelectRequest, SelectionCacheConfig as RsSelectionCacheConfig, - SelectionError, SelectionService as RustSelectionService, SelectionServiceBuilder, - SelectionServiceConfig, WorkerPatchRequest, WorkerRequest, WorkerSelectionPolicyRegistry, + self, OverlapScoresRequest, PotentialLoadsRequest, ReservationRequest, SelectAndReserveRequest, + SelectRequest, SelectionCacheConfig as RsSelectionCacheConfig, SelectionError, + SelectionService as RustSelectionService, SelectionServiceBuilder, SelectionServiceConfig, + WorkerPatchRequest, WorkerRequest, WorkerSelectionPolicyRegistry, warn_for_unserved_worker_selection_policies, }; #[cfg(feature = "slot-tracker")] @@ -358,11 +358,6 @@ struct SelectServiceCli { #[arg(long, value_delimiter = ',')] indexer_peers: Vec, - /// Base URL of a standalone indexer that serves the primary KV index; this - /// selector then does not subscribe to worker KV events itself - #[arg(long)] - remote_indexer_url: Option, - /// Local ZMQ PUB port for active-load replica events #[arg(long, value_parser = parse_nonzero_port)] replica_sync_port: Option, @@ -556,7 +551,6 @@ where port: cli.port, threads: cli.threads, indexer_peers: cli.indexer_peers, - remote_indexer_url: cli.remote_indexer_url, replica_sync_port: cli.replica_sync_port, replica_sync_peers: cli.replica_sync_peers, session_affinity_ttl: cli @@ -668,7 +662,7 @@ pub(crate) struct SelectionService { impl SelectionService { /// Create a selection service. `indexer_threads` sizes the KV indexer pool. #[new] - #[pyo3(signature = (*, indexer_threads = 4, indexer_peers = None, replica_sync_port = None, replica_sync_peers = None, selection_cache = None, remote_indexer_url = None, session_affinity_ttl_secs = None))] + #[pyo3(signature = (*, indexer_threads = 4, indexer_peers = None, replica_sync_port = None, replica_sync_peers = None, selection_cache = None, session_affinity_ttl_secs = None))] #[allow(clippy::too_many_arguments)] fn new( py: Python<'_>, @@ -677,7 +671,6 @@ impl SelectionService { replica_sync_port: Option, replica_sync_peers: Option>, selection_cache: Option, - remote_indexer_url: Option, session_affinity_ttl_secs: Option, ) -> PyResult { let replica_sync_peers = replica_sync_peers.unwrap_or_default(); @@ -701,9 +694,6 @@ impl SelectionService { if let Some(port) = replica_sync_port { builder = builder.replica_sync(port, replica_sync_peers); } - if let Some(url) = remote_indexer_url { - builder = builder.kv_index(KvIndexSource::Remote(url)); - } if let Some(ttl) = session_affinity_ttl_secs { builder = builder.session_affinity( session_affinity_ttl_from_secs(ttl).map_err(PyValueError::new_err)?, diff --git a/lib/bindings/python/src/dynamo/_core.pyi b/lib/bindings/python/src/dynamo/_core.pyi index 5f77afbf4c38..9df103a9c146 100644 --- a/lib/bindings/python/src/dynamo/_core.pyi +++ b/lib/bindings/python/src/dynamo/_core.pyi @@ -748,15 +748,10 @@ class SelectionService: replica_sync_port: Optional[int] = None, replica_sync_peers: Optional[list[str]] = None, selection_cache: Optional[SelectionCacheConfig] = None, - remote_indexer_url: Optional[str] = None, session_affinity_ttl_secs: Optional[float] = None, ) -> None: """Create a selection service. `indexer_threads` sizes the KV indexer pool. - `remote_indexer_url` is the base URL of a standalone indexer that serves - the primary KV index; when set, this service does not subscribe to worker - KV events itself. - `session_affinity_ttl_secs` enables session affinity with an idle TTL between 1 and 31536000 seconds. """ diff --git a/lib/kv-router/src/indexer/lower_tier_indexers.rs b/lib/kv-router/src/indexer/lower_tier_indexers.rs index f2d067fcb6a2..a6b0e945ddf5 100644 --- a/lib/kv-router/src/indexer/lower_tier_indexers.rs +++ b/lib/kv-router/src/indexer/lower_tier_indexers.rs @@ -408,6 +408,35 @@ mod tests { }; use crate::test_utils::{router_event, stored_blocks_with_sequence_hashes}; + #[test] + fn tiered_matches_json_round_trip_preserves_device_and_lower_tier_scores() { + let worker = WorkerWithDpRank::new(7, 1); + let mut original = TieredMatchDetails::default(); + original.device.overlap_scores.scores.insert(worker, 2); + original.device.overlap_scores.frequencies = vec![1, 1]; + let mut lower = LowerTierMatchDetails::default(); + lower.hits.insert(worker, 1); + original.lower_tier.insert(StorageTier::HostPinned, lower); + + let json = serde_json::to_vec(&WireTieredMatchDetails::from(&original)).unwrap(); + let wire: WireTieredMatchDetails = serde_json::from_slice(&json).unwrap(); + let restored = TieredMatchDetails::from(wire); + + assert_eq!( + restored.device.overlap_scores.scores, + original.device.overlap_scores.scores + ); + assert_eq!( + restored.device.overlap_scores.frequencies, + original.device.overlap_scores.frequencies + ); + assert_eq!(restored.lower_tier.len(), 1); + assert_eq!( + restored.lower_tier[&StorageTier::HostPinned].hits, + original.lower_tier[&StorageTier::HostPinned].hits + ); + } + fn local_hashes(values: &[u64]) -> Vec { values.iter().copied().map(LocalBlockHash).collect() } diff --git a/lib/kv-router/src/services/indexer/backend.rs b/lib/kv-router/src/services/indexer/backend.rs index 2c50365b5049..9c4bd7dc6210 100644 --- a/lib/kv-router/src/services/indexer/backend.rs +++ b/lib/kv-router/src/services/indexer/backend.rs @@ -4,21 +4,19 @@ use std::sync::Arc; use std::time::Duration; -use anyhow::{Context, Result}; +use anyhow::Result; use async_trait::async_trait; -use serde::{Deserialize, Serialize}; use tokio_util::sync::CancellationToken; use crate::ConcurrentRadixTreeCompressed; use crate::ThreadPoolIndexer; use crate::approx::PruneConfig; use crate::config::{ApproximateCachePolicyKind, KvRouterConfig}; -use crate::identity::{RoutingPartitionId, default_routing_group}; use crate::indexer::{ ApproximateLruIncarnation, ApproximateLruStats, ApproximateRetentionConfig, KvIndexer, KvIndexerInterface, KvIndexerMetrics, KvRouterError, LowerTierIndexers, LowerTierQueryOptions, MatchDetails, RoutingDecisionHashes, SyncIndexer, TieredMatchDetails, TieredMatchProvider, - WireTieredMatchDetails, record_unsupported_residency_event, + record_unsupported_residency_event, }; use crate::protocols::{ DpRank, KvCacheEventData, LocalBlockHash, OverlapScores, ResidencyProjection, @@ -36,131 +34,6 @@ pub enum PrimaryRetention { /// expire after the TTL. This is the frontend's `use_kv_events=false` /// approximate mode. ApproximateTtl(Duration), - /// The primary lives in a standalone indexer service reached over HTTP. - /// Workers publish their KV events to that service, not to this process. - Remote(Arc), -} - -/// Connection to a standalone indexer service, shared by every partition's -/// [`RemoteIndexerClient`]. -#[derive(Debug)] -pub struct RemoteIndexerTransport { - base_url: String, - client: reqwest::Client, -} - -impl PartialEq for RemoteIndexerTransport { - fn eq(&self, other: &Self) -> bool { - self.base_url == other.base_url - } -} - -impl Eq for RemoteIndexerTransport {} - -impl RemoteIndexerTransport { - pub const DEFAULT_TIMEOUT: Duration = Duration::from_secs(5); - /// A blackholed indexer fails fast instead of consuming the whole request budget. - const CONNECT_TIMEOUT: Duration = Duration::from_secs(2); - - pub fn new(base_url: impl Into) -> Result { - Self::new_with_timeout(base_url, Self::DEFAULT_TIMEOUT) - } - - pub fn new_with_timeout(base_url: impl Into, timeout: Duration) -> Result { - let base_url = base_url.into().trim_end_matches('/').to_string(); - if !(base_url.starts_with("http://") || base_url.starts_with("https://")) { - anyhow::bail!("remote indexer URL must start with http:// or https://: {base_url}"); - } - let client = reqwest::Client::builder() - .connect_timeout(Self::CONNECT_TIMEOUT.min(timeout)) - .timeout(timeout) - .build() - .context("failed to build remote indexer HTTP client")?; - Ok(Self { base_url, client }) - } - - pub fn base_url(&self) -> &str { - &self.base_url - } -} - -/// Request body of the standalone indexer's `POST /query_tiered_by_hash`. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct TieredQueryByHashRequest { - pub block_hashes: Vec, - pub model_name: String, - #[serde(default = "default_routing_group")] - pub routing_group: String, -} - -/// Response body of the standalone indexer's `POST /query_tiered_by_hash`: -/// the lossless tiered match shape a remote primary needs, as opposed to the -/// Mooncake score summary `/query_by_hash` returns. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct TieredQueryResponse { - pub block_size: u32, - pub tiered: WireTieredMatchDetails, -} - -/// Per-partition view of a remote primary indexer. -#[derive(Debug)] -pub struct RemoteIndexerClient { - transport: Arc, - key: RoutingPartitionId, -} - -impl RemoteIndexerClient { - pub fn new(transport: Arc, key: RoutingPartitionId) -> Self { - Self { transport, key } - } - - pub async fn find_tiered_matches( - &self, - sequence: &[LocalBlockHash], - ) -> Result { - if sequence.is_empty() { - return Ok(TieredMatchDetails::default()); - } - let request = TieredQueryByHashRequest { - block_hashes: sequence.iter().map(|hash| hash.0 as i64).collect(), - model_name: self.key.model_name.clone(), - routing_group: self.key.routing_group.clone(), - }; - let url = format!("{}/query_tiered_by_hash", self.transport.base_url); - let response = self - .transport - .client - .post(&url) - .json(&request) - .send() - .await - .map_err(|error| { - tracing::warn!( - %error, - url, - is_timeout = error.is_timeout(), - is_connect = error.is_connect(), - "Remote indexer query failed" - ); - KvRouterError::IndexerOffline - })?; - let status = response.status(); - if !status.is_success() { - tracing::warn!( - %status, - url, - model_name = %self.key.model_name, - routing_group = %self.key.routing_group, - "Remote indexer query returned an error status" - ); - return Err(KvRouterError::IndexerOffline); - } - let body: TieredQueryResponse = response.json().await.map_err(|error| { - tracing::warn!(%error, url, "Remote indexer query returned an invalid body"); - KvRouterError::IndexerOffline - })?; - Ok(body.tiered.into()) - } } /// Indexer shape derived from router configuration, mirroring the frontend @@ -188,25 +61,6 @@ impl IndexerPolicy { } } - pub fn is_remote(&self) -> bool { - matches!(self.primary, PrimaryRetention::Remote(_)) - } - - /// Replace an event-driven primary with a remote one at `base_url`. - /// - /// An approximate primary cannot be remote: the standalone indexer has no - /// routing-decision write path, so `use_kv_events=false` is rejected. - pub fn with_remote_indexer(mut self, base_url: impl Into) -> Result { - match self.primary { - PrimaryRetention::EventDriven | PrimaryRetention::Remote(_) => {} - PrimaryRetention::ApproximateTtl(_) => anyhow::bail!( - "a remote indexer requires use_kv_events=true; the standalone indexer does not record routing decisions" - ), - } - self.primary = PrimaryRetention::Remote(Arc::new(RemoteIndexerTransport::new(base_url)?)); - Ok(self) - } - /// Resolve the indexer shape for `config`. /// /// Rejects the same combinations the frontend rejects. The local side indexer @@ -362,8 +216,7 @@ impl SideIndexer { } } -/// A primary index served by another process. The HTTP standalone indexer -/// implements it here; a request-plane one implements it in its own crate. +/// A primary index served by another process, with transport supplied by the host. #[async_trait] pub trait RemotePrimary: Send + Sync { async fn find_matches_by_tier( @@ -384,29 +237,6 @@ pub trait RemotePrimary: Send + Sync { fn use_kv_events(&self) -> bool; } -#[async_trait] -impl RemotePrimary for RemoteIndexerClient { - async fn find_matches_by_tier( - &self, - block_hashes: Vec, - _device_only: bool, - ) -> Result { - Ok(self.find_tiered_matches(&block_hashes).await?) - } - - async fn record_routing_decision( - &self, - _worker: WorkerWithDpRank, - _hashes: RoutingDecisionHashes, - ) -> Result<()> { - anyhow::bail!("the standalone indexer does not record routing decisions") - } - - fn use_kv_events(&self) -> bool { - true - } -} - /// Block-content indexer for one routing partition: a primary device-tier /// backend plus a per-tier registry of lower-tier indexers (host-pinned, /// disk, …), or a remote primary, or nothing when routing is load-only. @@ -451,11 +281,6 @@ impl Indexer { } } - #[cfg(test)] - pub(crate) fn is_remote(&self) -> bool { - matches!(self, Self::Remote { .. }) - } - pub fn supports_overlap_refresh(&self) -> bool { matches!(self, Self::Single { .. } | Self::Concurrent { .. }) } @@ -1257,7 +1082,6 @@ mod tests { fn policy_indexer(num_threads: usize, policy: IndexerPolicy) -> Indexer { create_indexer_with_policy( - &RoutingPartitionId::new("model", "default"), 4, num_threads, Arc::new(KvIndexerMetrics::new_unregistered()), @@ -1443,13 +1267,33 @@ mod tests { } #[tokio::test] - async fn remote_indexer_queries_the_standalone_service_and_merges_side_scores() { - use crate::services::indexer::registry::WorkerRegistry; - use crate::services::indexer::server::spawn_test_indexer_server; + async fn remote_primary_merges_local_side_scores() { + struct TestRemote(Indexer); + + #[async_trait] + impl RemotePrimary for TestRemote { + async fn find_matches_by_tier( + &self, + hashes: Vec, + _device_only: bool, + ) -> Result { + Ok(self.0.find_tiered_matches(hashes).await?) + } + + async fn record_routing_decision( + &self, + _worker: WorkerWithDpRank, + _hashes: RoutingDecisionHashes, + ) -> Result<()> { + anyhow::bail!("event-driven remote primary must not receive routing decisions") + } + + fn use_kv_events(&self) -> bool { + true + } + } - let key = RoutingPartitionId::new("model", "default"); - let served = Arc::new(WorkerRegistry::new(1)); - let served_indexer = served.get_or_create_indexer(key.clone(), 4); + let served_indexer = create_indexer(4, 1); let confirmed = WorkerWithDpRank::new(7, 0); served_indexer .apply_event_routed(store_event(7, 0, 1, &[], &[11, 12], StorageTier::Device)) @@ -1472,17 +1316,17 @@ mod tests { let _ = inner.dump_events().await.unwrap(); } } - let (base_url, server) = spawn_test_indexer_server(served).await; - - let policy = IndexerPolicy { - primary: PrimaryRetention::EventDriven, - side_ttl: Some(Duration::from_secs(60)), - } - .with_remote_indexer(base_url.clone()) - .unwrap(); - assert!(policy.is_remote()); - let indexer = policy_indexer(1, policy.clone()); - assert!(indexer.is_remote()); + let indexer = Indexer::Remote { + primary: Arc::new(TestRemote(served_indexer)), + approx: Some(SideIndexer::new( + Duration::from_secs(60), + 4, + 1, + Arc::new(KvIndexerMetrics::new_unregistered()), + CancellationToken::new(), + )), + primary_records_routing_decisions: false, + }; assert!(indexer.records_routing_decisions()); assert!(!indexer.supports_kv_transfer_chain_retention()); @@ -1532,50 +1376,6 @@ mod tests { let flat = indexer.find_matches(sequence).await.unwrap(); assert_eq!(flat.scores.get(&confirmed).copied(), Some(2)); assert_eq!(flat.scores.get(&predicted).copied(), Some(1)); - - // Unknown partitions and an unreachable service degrade to offline. - let other = create_indexer_with_policy( - &RoutingPartitionId::new("other", "default"), - 4, - 1, - Arc::new(KvIndexerMetrics::new_unregistered()), - &policy, - ); - assert!(matches!( - other.find_tiered_matches(vec![LocalBlockHash(11)]).await, - Err(KvRouterError::IndexerOffline) - )); - server.abort(); - let _ = server.await; - assert!(matches!( - indexer.find_tiered_matches(vec![LocalBlockHash(11)]).await, - Err(KvRouterError::IndexerOffline) - )); - } - - #[test] - fn remote_indexer_rejects_approximate_primary_and_bad_urls() { - let approximate = IndexerPolicy { - primary: PrimaryRetention::ApproximateTtl(Duration::from_secs(1)), - side_ttl: None, - }; - assert!( - approximate - .with_remote_indexer("http://indexer:8091") - .is_err() - ); - assert!( - IndexerPolicy::event_driven() - .with_remote_indexer("indexer:8091") - .is_err() - ); - let policy = IndexerPolicy::event_driven() - .with_remote_indexer("http://indexer:8091/") - .unwrap(); - let PrimaryRetention::Remote(transport) = &policy.primary else { - panic!("expected a remote primary"); - }; - assert_eq!(transport.base_url(), "http://indexer:8091"); } #[test] @@ -1681,7 +1481,6 @@ pub fn create_indexer_with_metrics( metrics: Arc, ) -> Indexer { create_indexer_with_policy( - &RoutingPartitionId::new("default", "default"), block_size, num_threads, metrics, @@ -1690,7 +1489,6 @@ pub fn create_indexer_with_metrics( } pub fn create_indexer_with_policy( - key: &RoutingPartitionId, block_size: u32, num_threads: usize, metrics: Arc, @@ -1710,13 +1508,6 @@ pub fn create_indexer_with_policy( PrimaryRetention::ApproximateTtl(ttl) => { Some(ApproximateRetentionConfig::Ttl(PruneConfig { ttl: *ttl })) } - PrimaryRetention::Remote(transport) => { - return Indexer::Remote { - primary: Arc::new(RemoteIndexerClient::new(Arc::clone(transport), key.clone())), - approx, - primary_records_routing_decisions: false, - }; - } }; let primary_records_routing_decisions = retention.is_some(); if num_threads > 1 { diff --git a/lib/kv-router/src/services/indexer/registry.rs b/lib/kv-router/src/services/indexer/registry.rs index 140a29695636..5fec65edc7b4 100644 --- a/lib/kv-router/src/services/indexer/registry.rs +++ b/lib/kv-router/src/services/indexer/registry.rs @@ -462,7 +462,6 @@ impl WorkerRegistry { ); IndexerEntry { indexer: create_indexer_with_policy( - &key, block_size, self.num_threads, self.indexer_metrics.clone(), @@ -805,7 +804,6 @@ impl WorkerRegistry { ); IndexerEntry { indexer: create_indexer_with_policy( - &key, block_size, self.num_threads, self.indexer_metrics.clone(), diff --git a/lib/kv-router/src/services/indexer/server.rs b/lib/kv-router/src/services/indexer/server.rs index 4b25c21f979e..9f1f9bfa3877 100644 --- a/lib/kv-router/src/services/indexer/server.rs +++ b/lib/kv-router/src/services/indexer/server.rs @@ -17,11 +17,11 @@ use tokio_util::sync::CancellationToken; use crate::identity::{RoutingPartitionId, default_routing_group}; #[cfg(feature = "metrics")] use crate::indexer::KvIndexerMetrics; -use crate::indexer::{TieredMatchDetails, WireTieredMatchDetails}; +use crate::indexer::TieredMatchDetails; use crate::protocols::{BlockHashOptions, LocalBlockHash, WorkerId, compute_block_hash_for_seq}; use crate::services::overlap::{MooncakeOverlapSummary, build_mooncake_overlap_summaries}; -use super::backend::{Indexer, TieredQueryByHashRequest, TieredQueryResponse}; +use super::backend::Indexer; use super::registry::{ListenerControlError, WorkerRegistry}; /// We need to fit one million tokens as JSON text, this should do it. @@ -122,16 +122,6 @@ struct QueryRequest { cache_salt: Option, } -/// Server-side view of `TieredQueryByHashRequest` that also captures the -/// `cache_salt` some `/query_by_hash` callers send, so it can be rejected here -/// as well instead of being silently ignored. -#[derive(Deserialize)] -struct TieredQueryByHashBody { - #[serde(flatten)] - inner: TieredQueryByHashRequest, - cache_salt: Option, -} - #[derive(Deserialize)] struct QueryByHashRequest { block_hashes: Vec, @@ -397,72 +387,6 @@ async fn query_by_hash( resp } -/// Lossless tiered lookup for a remote-primary selection service. Unlike -/// `/query_by_hash`, which renders the Mooncake score summary, this returns -/// the wire form of `TieredMatchDetails` so the caller can run its own -/// overlap analysis and side-indexer merge. -async fn query_tiered_by_hash( - State(state): State>, - Json(TieredQueryByHashBody { - inner: req, - cache_salt, - }): Json, -) -> Response { - let model = req.model_name.clone(); - if cache_salt.is_some() { - let mut resp = ( - StatusCode::BAD_REQUEST, - Json(serde_json::json!({ - "error": "cache_salt is not accepted by /query_tiered_by_hash; block_hashes must already include the intended cache salt" - })), - ) - .into_response(); - resp.extensions_mut().insert(AccessLogModel(model)); - return resp; - } - let key = RoutingPartitionId::new(req.model_name, req.routing_group); - let Some(ie) = state.registry.get_indexer(&key) else { - let mut resp = ( - StatusCode::NOT_FOUND, - Json(serde_json::json!({ - "error": format!( - "no indexer for model={} routing_group={}", - key.model_name, key.routing_group - ) - })), - ) - .into_response(); - resp.extensions_mut().insert(AccessLogModel(model)); - return resp; - }; - let block_size = ie.block_size; - let indexer = ie.indexer.clone(); - drop(ie); - - let block_hashes: Vec = req - .block_hashes - .iter() - .map(|h| LocalBlockHash(*h as u64)) - .collect(); - let mut resp = match indexer.find_tiered_matches(block_hashes).await { - Ok(tiered) => ( - StatusCode::OK, - Json(TieredQueryResponse { - block_size, - tiered: WireTieredMatchDetails::from(&tiered), - }), - ) - .into_response(), - Err(e) => ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": e.to_string()})), - ) - .into_response(), - }; - resp.extensions_mut().insert(AccessLogModel(model)); - resp -} - #[derive(Deserialize)] struct ListenerControlRequest { instance_id: WorkerId, @@ -642,10 +566,6 @@ fn build_router(state: Arc, test_endpoints: bool) -> Router { "/query_by_hash", post(query_by_hash).layer(DefaultBodyLimit::max(QUERY_REQUEST_BODY_LIMIT_BYTES)), ) - .route( - "/query_tiered_by_hash", - post(query_tiered_by_hash).layer(DefaultBodyLimit::max(QUERY_REQUEST_BODY_LIMIT_BYTES)), - ) .route("/dump", get(dump_events)) .route("/register_peer", post(register_peer)) .route("/deregister_peer", post(deregister_peer)) @@ -686,32 +606,9 @@ fn build_router(state: Arc, test_endpoints: bool) -> Router { router } -/// Serve a standalone indexer for `registry` on an ephemeral port. Returns the -/// base URL and the server task; abort the task to stop serving. -#[cfg(test)] -pub(crate) async fn spawn_test_indexer_server( - registry: Arc, -) -> (String, tokio::task::JoinHandle<()>) { - let listener = tokio::net::TcpListener::bind("127.0.0.1:0") - .await - .expect("bind test indexer"); - let base_url = format!("http://{}", listener.local_addr().expect("local addr")); - let app = create_router(Arc::new(AppState { - registry, - access_log_sink: None, - #[cfg(feature = "metrics")] - prom_registry: prometheus::Registry::new(), - })); - let task = tokio::spawn(async move { - let _ = axum::serve(listener, app).await; - }); - (base_url, task) -} - #[cfg(test)] mod tests { use super::*; - use crate::indexer::KvIndexerInterface; use axum::body::Body; use axum::http::{Request, StatusCode, header}; use tower::ServiceExt; @@ -732,130 +629,6 @@ mod tests { body } - #[tokio::test] - async fn query_tiered_by_hash_rejects_cache_salt() { - let app = create_router(Arc::new(AppState { - registry: Arc::new(WorkerRegistry::new(1)), - access_log_sink: None, - #[cfg(feature = "metrics")] - prom_registry: prometheus::Registry::new(), - })); - - let salted = app - .oneshot( - Request::builder() - .method("POST") - .uri("/query_tiered_by_hash") - .header(header::CONTENT_TYPE, "application/json") - .body(Body::from( - r#"{"block_hashes":[1],"model_name":"model","routing_group":"default","cache_salt":"s"}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(salted.status(), StatusCode::BAD_REQUEST); - } - - #[tokio::test] - async fn query_tiered_by_hash_returns_wire_tiered_matches() { - use super::super::backend::test_util::store_event; - use crate::protocols::{StorageTier, WorkerWithDpRank}; - - let registry = Arc::new(WorkerRegistry::new(1)); - let key = RoutingPartitionId::new("model", "default"); - let indexer = registry.get_or_create_indexer(key, 4); - let worker = WorkerWithDpRank::new(7, 0); - indexer - .apply_event_routed(store_event(7, 0, 1, &[], &[11, 12], StorageTier::Device)) - .await - .unwrap(); - indexer - .apply_event_routed(store_event( - 7, - 0, - 2, - &[11, 12], - &[13], - StorageTier::HostPinned, - )) - .await - .unwrap(); - if let Indexer::Single { - primary, - lower_tier, - .. - } = &indexer - { - let _ = primary.flush().await; - for inner in lower_tier.all() { - let _ = inner.dump_events().await.unwrap(); - } - } - let app = create_router(Arc::new(AppState { - registry, - access_log_sink: None, - #[cfg(feature = "metrics")] - prom_registry: prometheus::Registry::new(), - })); - - let body = serde_json::to_vec(&TieredQueryByHashRequest { - block_hashes: vec![11, 12, 13], - model_name: "model".to_string(), - routing_group: "default".to_string(), - }) - .unwrap(); - let response = app - .clone() - .oneshot( - Request::builder() - .method("POST") - .uri("/query_tiered_by_hash") - .header(header::CONTENT_TYPE, "application/json") - .body(Body::from(body)) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(response.status(), StatusCode::OK); - let bytes = axum::body::to_bytes(response.into_body(), usize::MAX) - .await - .unwrap(); - let parsed: TieredQueryResponse = serde_json::from_slice(&bytes).unwrap(); - assert_eq!(parsed.block_size, 4); - let tiered: TieredMatchDetails = parsed.tiered.into(); - assert_eq!( - tiered.device.overlap_scores.scores.get(&worker).copied(), - Some(2) - ); - assert_eq!( - tiered - .lower_tier - .get(&crate::protocols::StorageTier::HostPinned) - .and_then(|d| d.hits.get(&worker).copied()), - Some(1) - ); - - let missing = serde_json::to_vec(&TieredQueryByHashRequest { - block_hashes: vec![11], - model_name: "other".to_string(), - routing_group: "default".to_string(), - }) - .unwrap(); - let response = app - .oneshot( - Request::builder() - .method("POST") - .uri("/query_tiered_by_hash") - .header(header::CONTENT_TYPE, "application/json") - .body(Body::from(missing)) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(response.status(), StatusCode::NOT_FOUND); - } - #[tokio::test] async fn query_rejects_request_bodies_over_limit() { let app = create_router(Arc::new(AppState { diff --git a/lib/kv-router/src/services/selection/README.md b/lib/kv-router/src/services/selection/README.md index 0991c675c60e..485b527b2dfb 100644 --- a/lib/kv-router/src/services/selection/README.md +++ b/lib/kv-router/src/services/selection/README.md @@ -19,11 +19,7 @@ and active-sequence accounting. Keep these implementation invariants explicit: from `longest_matched`. - Discovery-driven worker membership goes through `WorkerCatalogSource` and `CatalogReconciler` (`membership.rs`), not `upsert_worker`/`delete_worker`. -- A partition's KV index comes from `HostCache.index: KvIndexSource`: `Owned` - (a `KvEventIngress` builds and feeds it: `ZmqDirectIngress` here, the - runtime event plane in the frontend) or `Remote` (a standalone indexer - serves it). The `Indexer` type itself is shared with the frontend - (`services::indexer::backend`). +- A partition's KV index comes from `HostCache.index: KvIndexSource::Owned`. Its `KvEventIngress` builds and feeds it: `ZmqDirectIngress` here, `RuntimeIngress` in the frontend. The frontend ingress also supports Dynamo-native remote indexing. The `Indexer` type itself is shared with the frontend (`services::indexer::backend`). - Each partition owns its `SessionAffinity` table, including versioned bindings, idle TTL, and lease lifecycle. A reservation owns its affinity lease, so every reservation removal releases both. Frontend routing hosts diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index 65e49a900c1b..bfb3391bf5ee 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -199,9 +199,6 @@ pub enum KvIndexSource { /// events; it also decides what metadata a worker needs to be schedulable /// and what happens to the index when a worker leaves. Owned(Arc), - /// A standalone indexer at this base URL serves the primary index; this - /// core does not subscribe to worker KV events. - Remote(String), } impl Default for KvIndexSource { @@ -240,8 +237,6 @@ pub struct SelectionServiceConfig { pub port: u16, pub threads: usize, pub indexer_peers: Vec, - /// Base URL of a standalone indexer that serves the primary KV index. - pub remote_indexer_url: Option, pub replica_sync_port: Option, pub replica_sync_peers: Vec, pub kv_router_config: crate::config::KvRouterConfig, @@ -363,7 +358,7 @@ impl SelectionCore { WorkerRegistry::new_with_cancel_token(indexer_threads, cancel_token.clone()) .with_retained_indexers(), ); - let listens_for_kv_events = kv_router_config.use_kv_events && !indexer_policy.is_remote(); + let listens_for_kv_events = kv_router_config.use_kv_events; indexer_registry.set_indexer_policy(indexer_policy); if signal_indexer_ready { indexer_registry.signal_ready(); diff --git a/lib/kv-router/src/services/selection/core/tests.rs b/lib/kv-router/src/services/selection/core/tests.rs index 5597a2861940..853e9d40935f 100644 --- a/lib/kv-router/src/services/selection/core/tests.rs +++ b/lib/kv-router/src/services/selection/core/tests.rs @@ -43,25 +43,18 @@ fn local_core_with( .expect("valid test config") } -/// `new_inner` with the test defaults; `remote_indexer` swaps the primary for -/// the standalone indexer at that URL. +/// `new_inner` with the test defaults. fn core_with( config: crate::config::KvRouterConfig, host: SelectionHost, policy_factory: Option, worker_type: WorkerType, - remote_indexer: Option, affinity: Option, ) -> SelectionCore { let tracking_hash = Arc::new( TrackingHashContext::from_config(&config).expect("valid tracking hash configuration"), ); - let mut indexer_policy = IndexerPolicy::from_router_config(&config).expect("indexer policy"); - if let Some(base_url) = remote_indexer { - indexer_policy = indexer_policy - .with_remote_indexer(base_url) - .expect("remote policy"); - } + let indexer_policy = IndexerPolicy::from_router_config(&config).expect("indexer policy"); SelectionCore::new_inner( config, 1, @@ -92,7 +85,6 @@ fn core_with_host_and_policy( policy_factory, WorkerType::Aggregated, None, - None, ) } @@ -287,7 +279,6 @@ async fn selection_setup_uses_worker_type_label() { None, worker_type, None, - None, ); core.upsert_worker(worker(1)).await.expect("worker upsert"); @@ -388,79 +379,68 @@ async fn bookings_populate_the_approximate_primary_without_kv_events() { #[tokio::test] async fn unreachable_remote_indexer_is_reported_not_ready() { - let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind"); - let base_url = format!("http://{}", listener.local_addr().expect("addr")); - drop(listener); - let core = core_with( - test_config(true), - SelectionHost::default(), - None, - WorkerType::Aggregated, - Some(base_url), - None, - ); - core.upsert_worker(worker(1)).await.expect("worker upsert"); - assert!(matches!( - core.select(select_request()).await, - Err(SelectionError::Indexer( - crate::indexer::KvRouterError::IndexerOffline - )) - )); -} + use crate::indexer::{KvRouterError, TieredMatchDetails}; + use crate::services::indexer::backend::RemotePrimary; -#[tokio::test] -async fn remote_indexer_serves_selection_without_local_kv_listeners() { - use crate::indexer::KvIndexerInterface; - use crate::protocols::{BlockHashOptions, StorageTier, compute_block_hash_for_seq}; - use crate::services::indexer::registry::WorkerRegistry; - use crate::services::indexer::server::spawn_test_indexer_server; + struct OfflineRemote; + #[async_trait::async_trait] + impl RemotePrimary for OfflineRemote { + async fn find_matches_by_tier( + &self, + _: Vec, + _: bool, + ) -> anyhow::Result { + Err(KvRouterError::IndexerOffline.into()) + } - // The standalone indexer holds worker 2's cache for the test prompt. - let key = default_key(); - let served = Arc::new(WorkerRegistry::new(1)); - let served_indexer = served.get_or_create_indexer(key.clone(), 4); - let hashes: Vec = - compute_block_hash_for_seq(&[1, 2, 3, 4], 4, BlockHashOptions::default()) - .into_iter() - .map(|hash| hash.0) - .collect(); - served_indexer - .apply_event_routed(store_event(2, 0, 1, &[], &hashes, StorageTier::Device)) - .await - .unwrap(); - if let Indexer::Single { primary, .. } = &served_indexer { - let _ = primary.flush().await; + async fn record_routing_decision( + &self, + _: WorkerWithDpRank, + _: RoutingDecisionHashes, + ) -> anyhow::Result<()> { + anyhow::bail!("event-driven remote primary must not receive routing decisions") + } + + fn use_kv_events(&self) -> bool { + true + } + } + + struct OfflineIngress; + #[async_trait::async_trait] + impl KvEventIngress for OfflineIngress { + fn open(&self, _: &WorkerRegistry, _: &RoutingPartitionId, _: u32) -> Indexer { + Indexer::Remote { + primary: Arc::new(OfflineRemote), + approx: None, + primary_records_routing_decisions: false, + } + } } - let (base_url, server) = spawn_test_indexer_server(served).await; - // use_kv_events=true, but the primary is remote: workers need no - // kv_events endpoints and no ZMQ listener is started here. let core = core_with( test_config(true), - SelectionHost::default(), + SelectionHost { + cache: HostCache { + index: KvIndexSource::Owned(Arc::new(OfflineIngress)), + shared: None, + }, + ..SelectionHost::default() + }, None, WorkerType::Aggregated, - Some(base_url), None, ); - assert!(!core.listens_for_kv_events); - for worker_id in [1, 2] { - let record = core - .upsert_worker(worker(worker_id)) - .await - .expect("worker upsert"); - assert_eq!(record.lifecycle, WorkerLifecycle::Schedulable, "{record:?}"); - } - - let response = core.select(select_request()).await.expect("select"); - assert_eq!( - response.worker_id, 2, - "remote cache credit steers selection" - ); - assert_eq!(response.overlap.longest_matched, 4); - - core.delete_worker(2).await.expect("delete worker"); - server.abort(); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + let error = core + .select(select_request()) + .await + .expect_err("offline indexer"); + assert!(matches!( + error, + SelectionError::Indexer(KvRouterError::IndexerOffline) + )); + assert_eq!(error.status_code(), 503); } /// Two event-driven workers; worker 1 holds every block of an 8-token @@ -1444,7 +1424,6 @@ async fn dropped_book_selection_during_routing_record_frees_booking_and_claim() }, None, WorkerType::Aggregated, - None, Some( SessionAffinityConfig::new(Duration::from_secs(10)) .with_mode(SessionAffinityMode::Hard), @@ -2055,7 +2034,6 @@ fn core_with_session_affinity_mode(mode: SessionAffinityMode) -> SelectionCore { SelectionHost::default(), None, WorkerType::Aggregated, - None, Some(SessionAffinityConfig::new(Duration::from_secs(10)).with_mode(mode)), ) } diff --git a/lib/kv-router/src/services/selection/core/workers.rs b/lib/kv-router/src/services/selection/core/workers.rs index 6248b02f4dfa..7e04f42b58a3 100644 --- a/lib/kv-router/src/services/selection/core/workers.rs +++ b/lib/kv-router/src/services/selection/core/workers.rs @@ -8,7 +8,6 @@ use std::collections::HashSet; use super::reservations::{ReservationIndexObserver, spawn_reservation_index_sweep}; use super::*; -use crate::services::selection::ingress::remove_worker_from_index; impl SelectionCore { pub async fn upsert_worker( @@ -308,14 +307,8 @@ impl SelectionCore { ); } - let indexer = match &self.host.cache.index { - KvIndexSource::Owned(ingress) => { - ingress.open(&self.indexer_registry, &key, block_size) - } - KvIndexSource::Remote(_) => self - .indexer_registry - .get_or_create_indexer(key.clone(), block_size), - }; + let KvIndexSource::Owned(ingress) = &self.host.cache.index; + let indexer = ingress.open(&self.indexer_registry, &key, block_size); let overlap_refresh = indexer.supports_overlap_refresh().then(|| { Arc::new(TieredOverlapRefresher::new( indexer.clone(), @@ -382,18 +375,13 @@ impl SelectionCore { /// The ingress feeding core-owned indexes, when this core listens for KV events. fn ingress(&self) -> Option<&dyn KvEventIngress> { - match &self.host.cache.index { - KvIndexSource::Owned(ingress) if self.listens_for_kv_events => Some(ingress.as_ref()), - _ => None, - } + let KvIndexSource::Owned(ingress) = &self.host.cache.index; + self.listens_for_kv_events.then_some(ingress.as_ref()) } async fn cleanup_indexer_registration(&self, record: &WorkerCatalogRecord) { - if let KvIndexSource::Owned(ingress) = &self.host.cache.index { - ingress.detach(&self.indexer_registry, record).await; - return; - } - remove_worker_from_index(&self.indexer_registry, record).await; + let KvIndexSource::Owned(ingress) = &self.host.cache.index; + ingress.detach(&self.indexer_registry, record).await; } pub(super) fn publish_scheduler_config(&self, key: &RoutingPartitionId) { diff --git a/lib/kv-router/src/services/selection/service.rs b/lib/kv-router/src/services/selection/service.rs index 87932f187619..3be85bd12052 100644 --- a/lib/kv-router/src/services/selection/service.rs +++ b/lib/kv-router/src/services/selection/service.rs @@ -148,18 +148,8 @@ impl SelectionServiceBuilder { .resolve_for_worker_type(&self.kv_router_config, self.worker_type)?, }; let tracking_hash = Arc::new(TrackingHashContext::from_config(&self.kv_router_config)?); - let mut indexer_policy = IndexerPolicy::from_router_config(&self.kv_router_config)?; - if let KvIndexSource::Remote(base_url) = &self.host.cache.index { - indexer_policy = indexer_policy.with_remote_indexer(base_url.clone())?; - if !self.indexer_peers.is_empty() { - tracing::warn!( - remote_indexer_url = %base_url, - "indexer_peers are ignored when the primary indexer is remote" - ); - } - tracing::info!(remote_indexer_url = %base_url, "Using remote KV indexer"); - } - let recover_from_peers = !self.indexer_peers.is_empty() && !indexer_policy.is_remote(); + let indexer_policy = IndexerPolicy::from_router_config(&self.kv_router_config)?; + let recover_from_peers = !self.indexer_peers.is_empty(); let cancel_token = CancellationToken::new(); let mut startup_guard = StartupGuard::new(cancel_token.clone()); let replica_runtime = setup_replica_sync( @@ -259,9 +249,6 @@ impl SelectionServiceConfig { if let Some(ttl) = self.session_affinity_ttl { builder = builder.session_affinity(ttl); } - if let Some(url) = &self.remote_indexer_url { - builder = builder.kv_index(KvIndexSource::Remote(url.clone())); - } builder } } diff --git a/lib/kv-router/src/services/selection/tests.rs b/lib/kv-router/src/services/selection/tests.rs index c6efa387fdeb..4fb73b48d772 100644 --- a/lib/kv-router/src/services/selection/tests.rs +++ b/lib/kv-router/src/services/selection/tests.rs @@ -1773,7 +1773,6 @@ async fn selector_replica_sync_propagates_request_lifecycle() { port: 8092, threads: 1, indexer_peers: Vec::new(), - remote_indexer_url: None, session_affinity_ttl: None, replica_sync_port: Some(port_a), replica_sync_peers: Vec::new(), From 90557bb9df48a603a691d61c4fae43bb38668aff Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Tue, 15 Sep 2026 22:34:22 +0000 Subject: [PATCH 096/102] test(bindings): update selection service lifecycle constructor call Signed-off-by: Ishan Dhanani --- lib/bindings/python/rust/llm/kv.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lib/bindings/python/rust/llm/kv.rs b/lib/bindings/python/rust/llm/kv.rs index 50e9575a608a..9d69d8db565b 100644 --- a/lib/bindings/python/rust/llm/kv.rs +++ b/lib/bindings/python/rust/llm/kv.rs @@ -988,7 +988,7 @@ mod selection_service_lifecycle_tests { fn idempotent_shutdown() { pyo3::prepare_freethreaded_python(); let service = - Python::with_gil(|py| SelectionService::new(py, 1, None, None, None, None, None, None)) + Python::with_gil(|py| SelectionService::new(py, 1, None, None, None, None, None)) .unwrap(); Python::with_gil(|py| { service.shutdown(py); From 3b40f4df3e85a30ca8ebf49f90e60e0ae0f9f150 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Wed, 16 Sep 2026 00:37:28 +0000 Subject: [PATCH 097/102] fix(router): preserve frontend accounting before discovery Signed-off-by: Ishan Dhanani --- .../src/services/selection/core/mod.rs | 15 +++- .../src/services/selection/core/tests.rs | 75 +++++++++++++++++++ .../src/services/selection/core/workers.rs | 2 +- lib/llm/src/kv_router/embedded.rs | 2 + 4 files changed, 92 insertions(+), 2 deletions(-) diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index bfb3391bf5ee..e351ff49ca1f 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -225,11 +225,24 @@ pub struct HostTelemetry { /// Replica-sync transport the host carries for partitions this core does not /// mesh itself (ignored when the service runs its own ZMQ replica sync). -#[derive(Clone, Default)] +#[derive(Clone)] pub struct HostReplication { pub channels: Option, /// Owns request expiry when supplied; the partition then expires only through lifecycle events. pub request_leases: Option>, + /// Whether peer events may create accounting before catalog discovery. + /// This does not make an undiscovered worker eligible for selection. + pub replica_worker_policy: ReplicaWorkerPolicy, +} + +impl Default for HostReplication { + fn default() -> Self { + Self { + channels: None, + request_leases: None, + replica_worker_policy: ReplicaWorkerPolicy::RequireRegistered, + } + } } #[derive(Debug, Clone)] diff --git a/lib/kv-router/src/services/selection/core/tests.rs b/lib/kv-router/src/services/selection/core/tests.rs index 853e9d40935f..1e1fe833f974 100644 --- a/lib/kv-router/src/services/selection/core/tests.rs +++ b/lib/kv-router/src/services/selection/core/tests.rs @@ -1632,6 +1632,81 @@ async fn prefill_complete_is_idempotent_for_a_live_booking() { core.free_reservation("live").await.expect("free"); } +#[tokio::test] +async fn early_peer_accounting_follows_host_policy_without_granting_eligibility() { + for policy in [ + ReplicaWorkerPolicy::LazyRegister, + ReplicaWorkerPolicy::RequireRegistered, + ] { + let (outbound_tx, _outbound_rx) = mpsc::channel(16); + let (inbound_tx, inbound_rx) = mpsc::channel(16); + let channels = parking_lot::Mutex::new(Some(HostReplicaChannels { + outbound: Some(outbound_tx), + inbound_tx: inbound_tx.clone(), + inbound_rx, + process_id: 7, + ingress_observer: None, + })); + let core = core_with_host(SelectionHost { + replication: HostReplication { + channels: Some(Arc::new(move |_| channels.lock().take())), + replica_worker_policy: policy, + ..HostReplication::default() + }, + ..SelectionHost::default() + }); + core.upsert_worker(worker(1)).await.unwrap(); + core.select_and_reserve(reserve_request("warm")) + .await + .unwrap(); + core.free_reservation("warm").await.unwrap(); + for (request_id, worker_id) in [("early", 2), ("barrier", 1)] { + inbound_tx + .send(ActiveSequenceEvent { + request_id: request_id.to_string(), + worker: WorkerWithDpRank::new(worker_id, 0), + data: ActiveSequenceEventData::AddRequest { + token_sequence: Some(vec![1, 2]), + track_prefill_tokens: false, + expected_output_tokens: None, + prefill_load_hint: None, + }, + router_id: 99, + lora_name: None, + }) + .await + .unwrap(); + } + wait_until("peer batch applied", || { + core.indexed_booking("barrier").is_some() + }) + .await; + assert_eq!( + core.indexed_booking("early").is_some(), + policy == ReplicaWorkerPolicy::LazyRegister + ); + assert!( + !core + .catalog + .is_schedulable(WorkerAffinityTarget::new(2, None), &default_key()) + ); + let mut request = select_request(); + request.allowed_worker_ids = Some(HashSet::from([2])); + assert!( + core.select(request).await.is_err(), + "accounting alone made worker 2 selectable" + ); + core.upsert_worker(worker(2)).await.unwrap(); + core.select_and_reserve(reserve_request("after-discovery")) + .await + .unwrap(); + assert_eq!( + core.indexed_booking("early").is_some(), + policy == ReplicaWorkerPolicy::LazyRegister + ); + } +} + #[tokio::test] async fn mirrored_replica_bookings_are_indexed_until_freed() { let (outbound_tx, _outbound_rx) = mpsc::channel(16); diff --git a/lib/kv-router/src/services/selection/core/workers.rs b/lib/kv-router/src/services/selection/core/workers.rs index 7e04f42b58a3..2d87f3706c84 100644 --- a/lib/kv-router/src/services/selection/core/workers.rs +++ b/lib/kv-router/src/services/selection/core/workers.rs @@ -283,7 +283,7 @@ impl SelectionCore { scoped_replica_sync.process_id, worker_label, SequenceTrackerOptions { - replica_worker_policy: ReplicaWorkerPolicy::RequireRegistered, + replica_worker_policy: self.host.replication.replica_worker_policy, expiry_duration: self .host .replication diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs index 95957656bd53..0bda792bfd89 100644 --- a/lib/llm/src/kv_router/embedded.rs +++ b/lib/llm/src/kv_router/embedded.rs @@ -22,6 +22,7 @@ use dynamo_kv_router::scheduling::{ NonMaxOverlapSelectionObserver, OverloadedWorkerProvider, QueueLimitKind, QueueRejection, WorkerAvailabilityProvider, }; +use dynamo_kv_router::sequences::ReplicaWorkerPolicy; use dynamo_kv_router::services::selection::{ CatalogObserver, CatalogReconciler, DEFAULT_MODEL_NAME, HostCache, HostEligibility, HostLoad, HostReplication, HostTelemetry, KvEventIngress, KvIndexSource, SelectionHost, @@ -286,6 +287,7 @@ impl EmbeddedSelection { replication: HostReplication { channels: replica_sync, request_leases, + replica_worker_policy: ReplicaWorkerPolicy::LazyRegister, }, }) .build() From 9d1859a2cf0c9320edcde5bc68ec578cf5ec6bca Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Wed, 16 Sep 2026 00:37:43 +0000 Subject: [PATCH 098/102] fix(router): fence pending reservation claims by identity Signed-off-by: Ishan Dhanani --- .../services/selection/core/reservations.rs | 41 +++++-- .../src/services/selection/core/run.rs | 7 +- .../src/services/selection/core/tests.rs | 103 +++++++++++++++++- 3 files changed, 130 insertions(+), 21 deletions(-) diff --git a/lib/kv-router/src/services/selection/core/reservations.rs b/lib/kv-router/src/services/selection/core/reservations.rs index 3cc069eae657..7cf0153f9aa7 100644 --- a/lib/kv-router/src/services/selection/core/reservations.rs +++ b/lib/kv-router/src/services/selection/core/reservations.rs @@ -6,6 +6,10 @@ //! release them. use super::*; +use crate::scheduling::queue::BookingHandle; +use std::sync::atomic::{AtomicU64, Ordering}; + +static NEXT_CLAIM_ID: AtomicU64 = AtomicU64::new(1); /// Resolved inputs for booking a reservation, shared by the cached and explicit /// `create_reservation` paths. @@ -38,6 +42,7 @@ pub(super) type ReservationIndex = RwLock>; pub(super) struct Reservation { pub(super) partition: RoutingPartitionId, pub(super) booking: Option, + pub(super) claim_id: Option, pub(super) _affinity_lease: Option, } @@ -46,15 +51,31 @@ pub(super) struct Reservation { pub(super) struct ReservationClaim<'a> { pub(super) index: &'a ReservationIndex, pub(super) selection_id: String, + pub(super) claim_id: u64, pub(super) armed: bool, } impl ReservationClaim<'_> { - pub(super) fn install(mut self, reservation: Reservation) { + pub(super) fn install( + mut self, + booking: BookingHandle, + affinity_lease: Option, + ) -> Result<(), SelectionError> { + let mut index = self.index.write(); + let Some(reservation) = index + .get_mut(&self.selection_id) + .filter(|reservation| reservation.claim_id == Some(self.claim_id)) + else { + return Err(SelectionError::Conflict(format!( + "selection {} reservation claim was replaced", + self.selection_id + ))); + }; + reservation.booking = Some(booking.commit()); + reservation.claim_id = None; + reservation._affinity_lease = affinity_lease; self.armed = false; - self.index - .write() - .insert(std::mem::take(&mut self.selection_id), reservation); + Ok(()) } } @@ -66,7 +87,7 @@ impl Drop for ReservationClaim<'_> { let mut index = self.index.write(); if index .get(&self.selection_id) - .is_some_and(|reservation| reservation.booking.is_none()) + .is_some_and(|reservation| reservation.claim_id == Some(self.claim_id)) { index.remove(&self.selection_id); } @@ -99,6 +120,7 @@ impl ReplicaRequestLeaseObserver for ReservationIndexObserver { Reservation { partition: self.partition.clone(), booking: Some(booking.clone()), + claim_id: None, _affinity_lease: None, }, ); @@ -375,11 +397,7 @@ impl SelectionCore { if let Some(hashes) = routing_hashes { self.record_routing_decision(&entry, worker, hashes).await; } - claim.install(Reservation { - partition: key.clone(), - booking: Some(booking.commit()), - _affinity_lease: affinity_lease, - }); + claim.install(booking, affinity_lease)?; Ok(ReservationResponse { selection_id, @@ -403,17 +421,20 @@ impl SelectionCore { "selection {selection_id} is already reserved or being reserved" ))); } + let claim_id = NEXT_CLAIM_ID.fetch_add(1, Ordering::Relaxed); index.insert( selection_id.to_string(), Reservation { partition: key.clone(), booking: None, + claim_id: Some(claim_id), _affinity_lease: None, }, ); Ok(ReservationClaim { index: &self.reservation_index, selection_id: selection_id.to_string(), + claim_id, armed: true, }) } diff --git a/lib/kv-router/src/services/selection/core/run.rs b/lib/kv-router/src/services/selection/core/run.rs index 2054852f813a..21f2e9857924 100644 --- a/lib/kv-router/src/services/selection/core/run.rs +++ b/lib/kv-router/src/services/selection/core/run.rs @@ -48,7 +48,6 @@ use super::hint::{hint_capable_partition, transfer_hint_for_selection}; /// `try_acquire` then `join_initializing` attempts before a commit that finds /// the session initializing routes unpinned (see the module doc). const JOIN_ATTEMPTS: usize = 2; -use super::reservations::Reservation; use super::*; /// Action id of the single `kv.fetch` action a selection's KV hint carries. @@ -524,11 +523,7 @@ impl SelectionCore { self.record_routing_decision(&entry, response.best_worker, hashes) .await; } - claim.install(Reservation { - partition: key.clone(), - booking: Some(booking.commit()), - _affinity_lease: affinity_lease, - }); + claim.install(booking, affinity_lease)?; None } else { booking diff --git a/lib/kv-router/src/services/selection/core/tests.rs b/lib/kv-router/src/services/selection/core/tests.rs index 1e1fe833f974..ed2653c7b36d 100644 --- a/lib/kv-router/src/services/selection/core/tests.rs +++ b/lib/kv-router/src/services/selection/core/tests.rs @@ -1539,11 +1539,13 @@ fn index_observer_replaces_stale_and_claimed_rows_of_its_partition() { worker: WorkerWithDpRank::new(1, 0), attempt_id: AttemptId::new(attempt), }; - let row = |partition: &RoutingPartitionId, booking| Reservation { - partition: partition.clone(), - booking, - _affinity_lease: None, - }; + let row = + |partition: &RoutingPartitionId, booking: Option| Reservation { + partition: partition.clone(), + claim_id: booking.is_none().then_some(0), + booking, + _affinity_lease: None, + }; // A stale row (its booking expired before the sweep ran) yields to the mirror. index @@ -1563,6 +1565,7 @@ fn index_observer_replaces_stale_and_claimed_rows_of_its_partition() { let claim = ReservationClaim { index: &index, selection_id: "shared".to_string(), + claim_id: 0, armed: true, }; index @@ -1581,6 +1584,96 @@ fn index_observer_replaces_stale_and_claimed_rows_of_its_partition() { assert_eq!(index.read()["shared"].booking, Some(booking(4))); } +#[rstest::rstest] +#[tokio::test] +async fn replaced_claim_cannot_remove_or_overwrite_a_new_reservation( + #[values("drop", "pending", "installed")] replacement: &str, +) { + use crate::scheduling::AttemptId; + let core = local_core(test_config(false)); + let key = default_key(); + core.upsert_worker(worker(1)).await.unwrap(); + let other_key = RoutingPartitionId::new("model", "other"); + core.upsert_worker(WorkerRequest { + routing_group: other_key.routing_group.clone(), + ..worker(2) + }) + .await + .unwrap(); + // Both partitions have registered slots before the controlled race. + for key in [&key, &other_key] { + let mut warm = reserve_request("warm"); + warm.routing_group = key.routing_group.clone(); + core.select_and_reserve(warm).await.unwrap(); + core.free_reservation("warm").await.unwrap(); + } + let old = core.claim_reservation("shared", &key).unwrap(); + let observer = ReservationIndexObserver { + index: Arc::clone(&core.reservation_index), + partition: key.clone(), + host: None, + }; + let peer = SchedulerBookingDescriptor { + request_id: "shared".to_string(), + worker: WorkerWithDpRank::new(1, 0), + attempt_id: AttemptId::new(999), + }; + observer.admitted(peer.clone()); + observer.completed(&peer); + let new = core.claim_reservation("shared", &other_key).unwrap(); + let request = |worker_id| SequenceRequest { + request_id: "shared".to_string(), + worker: WorkerWithDpRank::new(worker_id, 0), + token_sequence: Some(vec![1, 2]), + track_prefill_tokens: false, + expected_output_tokens: None, + prefill_load_hint: None, + lora_name: None, + }; + let old_entry = core.entry(&key).unwrap(); + let new_entry = core.entry(&other_key).unwrap(); + let new_id = new.claim_id; + let mut new = Some(new); + if replacement == "installed" { + let booking = new_entry + .scheduler + .add_request_if_registered_guarded(request(2)) + .unwrap(); + new.take().unwrap().install(booking, None).unwrap(); + } + if replacement != "drop" { + let booking = old_entry + .scheduler + .add_request_if_registered_guarded(request(1)) + .unwrap(); + assert!(matches!( + old.install(booking, None), + Err(SelectionError::Conflict(_)) + )); + wait_until("stale booking released", || { + !old_entry.scheduler.has_request("shared") + }) + .await; + } else { + drop(old); + } + if let Some(new) = new { + assert_eq!( + core.reservation_index.read()["shared"].claim_id, + Some(new_id) + ); + let booking = new_entry + .scheduler + .add_request_if_registered_guarded(request(2)) + .unwrap(); + new.install(booking, None).unwrap(); + } + assert_eq!(core.indexed_booking("shared").unwrap().0.key, other_key); + core.free_reservation("shared").await.unwrap(); + assert!(!new_entry.scheduler.has_request("shared")); + assert!(core.reservation_index.read().is_empty()); +} + #[tokio::test] async fn free_of_an_in_flight_reservation_is_not_found() { let core = saturated_core(); From 6bef1c458d1fc01b0f496c3cd57cd8ba625fec7a Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Wed, 16 Sep 2026 00:37:52 +0000 Subject: [PATCH 099/102] fix(router): yield during repeated affinity invalidation Signed-off-by: Ishan Dhanani --- .../src/services/selection/core/mod.rs | 4 ++ .../src/services/selection/core/run.rs | 13 ++++ .../src/services/selection/core/tests.rs | 68 +++++++++++++++++++ 3 files changed, 85 insertions(+) diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index e351ff49ca1f..bae4224752ff 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -293,6 +293,8 @@ pub struct SelectionCore { /// Scheduler-config publishes that changed a partition's worker map. #[cfg(test)] pub(super) publish_count: std::sync::atomic::AtomicUsize, + #[cfg(test)] + pub(super) after_affinity_invalidation: Option>, } fn affinity_error(error: AffinityError) -> SelectionError { @@ -397,6 +399,8 @@ impl SelectionCore { fail_upsert_for: parking_lot::Mutex::default(), #[cfg(test)] publish_count: std::sync::atomic::AtomicUsize::new(0), + #[cfg(test)] + after_affinity_invalidation: None, } } diff --git a/lib/kv-router/src/services/selection/core/run.rs b/lib/kv-router/src/services/selection/core/run.rs index 21f2e9857924..2df9803a2179 100644 --- a/lib/kv-router/src/services/selection/core/run.rs +++ b/lib/kv-router/src/services/selection/core/run.rs @@ -48,6 +48,7 @@ use super::hint::{hint_capable_partition, transfer_hint_for_selection}; /// `try_acquire` then `join_initializing` attempts before a commit that finds /// the session initializing routes unpinned (see the module doc). const JOIN_ATTEMPTS: usize = 2; +const AFFINITY_INVALIDATIONS_BEFORE_YIELD: usize = 32; use super::*; /// Action id of the single `kv.fetch` action a selection's KV hint carries. @@ -576,6 +577,7 @@ impl SelectionCore { session_id: &str, key: &RoutingPartitionId, ) -> Result, SelectionError> { + let mut invalidations = 0; loop { let acquired = tokio::select! { _ = self.cancel_token.cancelled() => { @@ -593,6 +595,17 @@ impl SelectionCore { "Session affinity target is not schedulable; re-initializing" ); lease.invalidate(); + drop(lease); + #[cfg(test)] + if let Some(hook) = &self.after_affinity_invalidation { + hook(); + } + invalidations += 1; + if invalidations == AFFINITY_INVALIDATIONS_BEFORE_YIELD { + // Replica updates can keep acquire immediately ready. + tokio::task::yield_now().await; + invalidations = 0; + } } Ok(hold) => return Ok(Some(hold)), Err(AffinityError::ResourceExhausted(_)) => { diff --git a/lib/kv-router/src/services/selection/core/tests.rs b/lib/kv-router/src/services/selection/core/tests.rs index ed2653c7b36d..ca167878e3aa 100644 --- a/lib/kv-router/src/services/selection/core/tests.rs +++ b/lib/kv-router/src/services/selection/core/tests.rs @@ -1725,6 +1725,74 @@ async fn prefill_complete_is_idempotent_for_a_live_booking() { core.free_reservation("live").await.expect("free"); } +#[tokio::test] +async fn repeated_affinity_invalidation_yields_and_preserves_new_bindings() { + use super::super::affinity::AffinityVersion; + use std::future::Future; + + for cancel in [false, true] { + let mut core = core_with_session_affinity(); + core.upsert_worker(worker(1)).await.unwrap(); + let key = default_key(); + let entry = core.entry(&key).unwrap(); + let table = entry.affinity.get().unwrap().clone(); + let unavailable = WorkerAffinityTarget::new(2, Some(0)); + table.apply_replica_update( + "s".into(), + unavailable, + AffinityVersion { + sequence: 1, + writer_id: 99, + }, + ); + let invalidations = Arc::new(AtomicUsize::new(0)); + core.after_affinity_invalidation = Some(Arc::new({ + let table = table.clone(); + let invalidations = invalidations.clone(); + move || { + let count = invalidations.fetch_add(1, Ordering::SeqCst) + 1; + // A finite burst also makes a missing yield fail instead of hanging the test. + if count <= 64 { + table.apply_replica_update( + "s".into(), + unavailable, + AffinityVersion { + sequence: count as u64 + 1, + writer_id: 99, + }, + ); + } + } + })); + let mut pending = Box::pin(core.hold_session(&table, "s", &key)); + let mut context = std::task::Context::from_waker(std::task::Waker::noop()); + assert!(pending.as_mut().poll(&mut context).is_pending()); + assert_eq!(invalidations.load(Ordering::SeqCst), 32); + if cancel { + core.cancel_token.cancel(); + assert!(matches!( + pending.await, + Err(SelectionError::Scheduler( + KvSchedulerError::SubscriberShutdown + )) + )); + } else { + let healthy = WorkerAffinityTarget::new(1, Some(0)); + table.apply_replica_update( + "s".into(), + healthy, + AffinityVersion { + sequence: 100, + writer_id: 99, + }, + ); + let hold = pending.await.unwrap().unwrap(); + assert!(matches!(hold, Hold::Bound { target, .. } if target == healthy)); + assert_eq!(invalidations.load(Ordering::SeqCst), 32); + } + } +} + #[tokio::test] async fn early_peer_accounting_follows_host_policy_without_granting_eligibility() { for policy in [ From 8827c6132e5b915042df1e2fe3a363aea324dfdf Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Wed, 16 Sep 2026 22:05:18 +0000 Subject: [PATCH 100/102] fix(router): address selection core review feedback Preserve newer cached selections when an older replay fails. Release matching reservation ownership before cancellable queue progress, and drop displaced affinity leases outside the reservation-index lock. Record explicit multimodal-only reservations, remove unused worker registration helpers, and restore the embedded-router overload regression. Add deterministic coverage for replay replacement, cancelled frees, and lock-held lease cleanup. Signed-off-by: Ishan Dhanani --- lib/kv-router/src/scheduling/local.rs | 30 ++-- lib/kv-router/src/scheduling/queue.rs | 48 ++---- .../services/selection/core/reservations.rs | 24 ++- .../src/services/selection/core/tests.rs | 145 +++++++++++++++++- lib/kv-router/src/services/selection/input.rs | 2 +- lib/llm/src/kv_router.rs | 23 +++ 6 files changed, 204 insertions(+), 68 deletions(-) diff --git a/lib/kv-router/src/scheduling/local.rs b/lib/kv-router/src/scheduling/local.rs index 00494c6478f6..4cb363064621 100644 --- a/lib/kv-router/src/scheduling/local.rs +++ b/lib/kv-router/src/scheduling/local.rs @@ -503,10 +503,6 @@ where .await } - pub fn register_workers(&self, worker_ids: &HashSet) { - self.queue.register_workers(worker_ids); - } - pub async fn add_request(&self, req: SequenceRequest) -> Result<(), SequenceError> { self.slots.add_request(req, Instant::now()) } @@ -611,6 +607,16 @@ where pub async fn free_if_booking( &self, booking: &SchedulerBookingDescriptor, + ) -> Result { + self.free_if_booking_with_cleanup(booking, || {}).await + } + + /// Complete owner cleanup after a successful release (including a stale + /// booking's `NoChange`), before the cancellable queue-progress wait. + pub(crate) async fn free_if_booking_with_cleanup( + &self, + booking: &SchedulerBookingDescriptor, + cleanup: impl FnOnce(), ) -> Result { let outcome = self.slots.free_if_booking( &booking.request_id, @@ -618,6 +624,7 @@ where booking.attempt_id, Instant::now(), )?; + cleanup(); if outcome.is_applied() { self.queue.update_worker(booking.worker).await; } @@ -1570,21 +1577,6 @@ mod tests { cancel_token.cancel(); } - #[tokio::test] - async fn test_register_workers_uses_default_dp_fallback() { - let (scheduler, _slots, _cfg_tx, cancel_token) = - make_scheduler(HashMap::new(), None, false, None); - - scheduler.register_workers(&HashSet::from([42])); - let loads = scheduler.get_potential_loads(None, 64, HashMap::new(), true); - - assert_eq!(loads.len(), 1); - assert_eq!(loads[0].worker_id, 42); - assert_eq!(loads[0].dp_rank, 0); - - cancel_token.cancel(); - } - #[tokio::test] async fn test_worker_watch_updates_slot_ranges() { let mut workers = HashMap::new(); diff --git a/lib/kv-router/src/scheduling/queue.rs b/lib/kv-router/src/scheduling/queue.rs index 5d53b2283471..07a9f5ef6441 100644 --- a/lib/kv-router/src/scheduling/queue.rs +++ b/lib/kv-router/src/scheduling/queue.rs @@ -33,7 +33,6 @@ use crate::protocols::{ LocalBlockHash, PrefillLoadHint, WorkerConfigLike, WorkerId, WorkerSelectionResult, WorkerWithDpRank, }; -use crate::sequences::topology::WorkerDpRange; use crate::sequences::{ ActiveSequencesMultiWorker, LifecycleMutationOutcome, SequenceError, SequencePublisher, SequenceRequest, @@ -513,12 +512,12 @@ pub struct SchedulerQueue< /// Incremented after push, decremented after pop. Lock-free reads via `Relaxed` load. pending_isl_tokens: Arc, class_counters: Arc>, - slots: Arc>, - workers_with_configs: watch::Receiver>, queueing_enabled: bool, supports_overlap_refresh: bool, non_max_overlap_selection_observer: Arc>, - _marker: PhantomData (Sel, RF)>, + #[allow(clippy::type_complexity)] + // Covariant type markers, without ownership or auto-trait bounds. + _marker: PhantomData (P, C, Sel, RF)>, } impl< @@ -620,8 +619,8 @@ impl< pending_count: Arc::clone(&pending_count), pending_isl_tokens: Arc::clone(&pending_isl_tokens), class_counters: Arc::clone(&class_counters), - slots: Arc::clone(&slots), - workers_with_configs: workers_with_configs.clone(), + slots, + workers_with_configs, start_time: Instant::now(), block_size, selector, @@ -639,8 +638,6 @@ impl< pending_count, pending_isl_tokens, class_counters, - slots, - workers_with_configs, queueing_enabled, supports_overlap_refresh: overlap_refresh_after.is_some(), non_max_overlap_selection_observer, @@ -656,29 +653,6 @@ impl< RF: OverlapScoresRefresh + Send + Sync + 'static, > SchedulerQueue { - /// Register externally-provided workers in the slot tracker. - /// - /// Looks up DP rank/size from the discovery watch channel; defaults to - /// `(0, 1)` for workers not yet known to discovery. - pub fn register_workers(&self, worker_ids: &std::collections::HashSet) { - let discovery_workers = self.workers_with_configs.borrow(); - for &worker_id in worker_ids { - let (dp_start, dp_size) = discovery_workers - .get(&worker_id) - .map(|runtime_config| { - ( - runtime_config.data_parallel_start_rank(), - runtime_config.data_parallel_size(), - ) - }) - .unwrap_or((0, 1)); - let range = WorkerDpRange::new(worker_id, dp_start, dp_size); - if let Err(error) = self.slots.upsert_worker(range) { - tracing::warn!(worker_id, %error, "Invalid externally-provided worker topology"); - } - } - } - /// Install the observer for admitted selections that sacrifice KV overlap. /// /// Returns `false` when an observer is already installed. @@ -1729,6 +1703,7 @@ mod tests { use crate::scheduling::OverlapSignals; use crate::scheduling::types::{KvSchedulerError, ScheduleMode}; use crate::scheduling::{RefreshedOverlap, RouterPolicyConfig}; + use crate::sequences::topology::WorkerDpRange; use crate::sequences::{ActiveSequencesMultiWorker, SequencePublisher}; use crate::test_utils::{NoopSequencePublisher, SimpleWorkerConfig}; use crate::{DefaultWorkerSelector, WorkerInputs, WorkerSelector}; @@ -3311,10 +3286,9 @@ policy_classes: )); } - /// Simulates the EPP path: router starts with zero workers (skip_initial_worker_wait), - /// then register_workers lazily injects workers before routing. + /// A queue starting with zero workers can route after slots and configs arrive. #[tokio::test(flavor = "multi_thread")] - async fn test_register_workers_lazy_epp_path() { + async fn test_worker_updates_after_empty_start() { let block_size = 16; let isl = 512; @@ -3330,7 +3304,7 @@ policy_classes: resp, Err(crate::scheduling::types::KvSchedulerError::NoEndpoints) ), - "expected NoEndpoints before register_workers, got {resp:?}" + "expected NoEndpoints before worker updates, got {resp:?}" ); // Lazily register two workers in the slot tracker (EPP supplies pod list) @@ -3372,9 +3346,9 @@ policy_classes: .unwrap(); } - /// Register_workers is additive: calling with a new set does NOT remove old workers. + /// Upserting a new worker preserves existing workers. #[tokio::test(flavor = "multi_thread")] - async fn test_register_workers_additive() { + async fn test_worker_updates_are_additive() { let block_size = 16; let isl = 256; diff --git a/lib/kv-router/src/services/selection/core/reservations.rs b/lib/kv-router/src/services/selection/core/reservations.rs index 7cf0153f9aa7..6074758f3bf4 100644 --- a/lib/kv-router/src/services/selection/core/reservations.rs +++ b/lib/kv-router/src/services/selection/core/reservations.rs @@ -105,7 +105,7 @@ pub(super) struct ReservationIndexObserver { impl ReplicaRequestLeaseObserver for ReservationIndexObserver { fn admitted(&self, booking: SchedulerBookingDescriptor) { - { + let displaced = { let mut index = self.index.write(); // A peer only admits an id this partition's scheduler does not hold, // so an existing row for it is a stale booking (expired, not yet @@ -123,9 +123,13 @@ impl ReplicaRequestLeaseObserver for ReservationIndexObserver { claim_id: None, _affinity_lease: None, }, - ); + ) + } else { + None } - } + }; + // Releasing an affinity lease takes a shard lock and publishes to peers. + drop(displaced); if let Some(host) = &self.host { host.admitted(booking); } @@ -197,7 +201,8 @@ impl SelectionCore { Ok(response) => response, // The session moved: the cached worker can never be booked for it. Err(SelectionError::BadRequest(message)) => { - self.selection_cache.discard(&key, &req.selection_id); + self.selection_cache + .remove(&key, &req.selection_id, generation); return Err(SelectionError::BadRequest(message)); } Err(error) => return Err(error), @@ -314,7 +319,8 @@ impl SelectionCore { // Hash-only reservations (sequence hashes without block hashes) carry // nothing an indexer can key on; recording is skipped for them. let can_record = entry.indexer.records_routing_decisions() - && (req.prompt.token_ids.is_some() || req.prompt.block_hashes.is_some()); + && (req.prompt.view().routing_tokens_and_mm_infos().is_some() + || req.prompt.block_hashes.is_some()); let routing_hashes = can_record .then(|| { req.prompt @@ -490,8 +496,12 @@ impl SelectionCore { let Some((entry, booking)) = self.indexed_booking(selection_id) else { return Err(Self::reservation_not_found(selection_id)); }; - let outcome = entry.scheduler.free_if_booking(&booking).await?; - forget_reservation_if(&self.reservation_index, &entry.key, &booking); + let outcome = entry + .scheduler + .free_if_booking_with_cleanup(&booking, || { + forget_reservation_if(&self.reservation_index, &entry.key, &booking); + }) + .await?; match outcome { LifecycleMutationOutcome::Applied => Ok(()), LifecycleMutationOutcome::NoChange => Err(Self::reservation_not_found(selection_id)), diff --git a/lib/kv-router/src/services/selection/core/tests.rs b/lib/kv-router/src/services/selection/core/tests.rs index ca167878e3aa..408944a166bb 100644 --- a/lib/kv-router/src/services/selection/core/tests.rs +++ b/lib/kv-router/src/services/selection/core/tests.rs @@ -308,8 +308,13 @@ async fn wait_for_overlap( .expect("approximate indexer never credited the booked prompt") } +#[rstest::rstest] +#[case::text(false)] +#[case::multimodal_only(true)] #[tokio::test] -async fn bookings_populate_the_approximate_primary_without_kv_events() { +async fn bookings_populate_the_approximate_primary_without_kv_events( + #[case] multimodal_only: bool, +) { // use_kv_events=false: the primary is approximate and bookings feed it. let core = local_core(test_config(false)); core.upsert_worker(worker(1)).await.expect("worker upsert"); @@ -356,9 +361,13 @@ async fn bookings_populate_the_approximate_primary_without_kv_events() { assert_eq!(credited.worker_id, 2); // And the explicit reservation form. - let prompt_c = || PromptRequest { - token_ids: Some(vec![9, 10, 11, 12]), - ..PromptRequest::default() + let prompt_c = || { + serde_json::from_value::(if multimodal_only { + serde_json::json!({"mm_routing_info": {"routing_token_ids": [9, 10, 11, 12]}}) + } else { + serde_json::json!({"token_ids": [9, 10, 11, 12]}) + }) + .expect("valid wire prompt") }; core.create_reservation(ReservationRequest { worker_id: Some(1), @@ -375,6 +384,7 @@ async fn bookings_populate_the_approximate_primary_without_kv_events() { }) .await; assert_eq!(credited.worker_id, 1); + assert_eq!(credited.overlap.longest_matched, 4); } #[tokio::test] @@ -1584,6 +1594,56 @@ fn index_observer_replaces_stale_and_claimed_rows_of_its_partition() { assert_eq!(index.read()["shared"].booking, Some(booking(4))); } +#[tokio::test] +async fn index_observer_releases_displaced_affinity_after_unlock() { + use super::super::affinity::{AffinityReplicaSink, AffinityVersion}; + use crate::scheduling::AttemptId; + + struct UnlockedIndexSink { + index: Arc, + published: AtomicBool, + } + + impl AffinityReplicaSink for UnlockedIndexSink { + fn publish(&self, _: &str, _: WorkerAffinityTarget, _: AffinityVersion) { + assert!(self.index.try_write().is_some(), "index must be unlocked"); + self.published.store(true, Ordering::Relaxed); + } + } + + let core = core_with_session_affinity(); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + core.select_and_reserve(session_reservation("shared", "s")) + .await + .expect("local booking"); + let (entry, booking) = core.indexed_booking("shared").expect("indexed booking"); + // A scheduler release that bypasses the core leaves a stale indexed lease. + entry.scheduler.free_if_booking(&booking).await.unwrap(); + let table = entry.affinity.get().expect("affinity table"); + let sink = Arc::new(UnlockedIndexSink { + index: Arc::clone(&core.reservation_index), + published: AtomicBool::new(false), + }); + assert!(table.enable_replication(1, sink.clone())); + let replacement = SchedulerBookingDescriptor { + attempt_id: AttemptId::new(999), + ..booking + }; + let observer = ReservationIndexObserver { + index: Arc::clone(&core.reservation_index), + partition: default_key(), + host: None, + }; + observer.admitted(replacement.clone()); + + assert!(sink.published.load(Ordering::Relaxed)); + assert_eq!(table.lease_count("s"), Some(0)); + assert_eq!( + core.reservation_index.read()["shared"].booking, + Some(replacement) + ); +} + #[rstest::rstest] #[tokio::test] async fn replaced_claim_cannot_remove_or_overwrite_a_new_reservation( @@ -2656,6 +2716,83 @@ async fn two_phase_replay_rejects_a_worker_the_session_left() { assert!(core.reservation_index.read().is_empty()); } +#[tokio::test] +async fn failed_replay_preserves_a_newer_cached_selection() { + let core = core_with_session_affinity(); + for id in [1, 2] { + core.upsert_worker(worker(id)).await.expect("worker upsert"); + } + let request = |worker_id| { + let mut request = select_request(); + request.selection_id = Some("pending".to_string()); + request.session_id = Some("s".to_string()); + request.allowed_worker_ids = Some(HashSet::from([worker_id])); + request + }; + core.select(request(1)).await.expect("first select"); + + let entry = core.entry(&default_key()).expect("entry"); + let table = entry.affinity.get().expect("affinity table"); + let initializer = table.acquire("s", None).await.expect("initializer"); + let mut replay = Box::pin(core.create_reservation(replay_reservation("pending"))); + let mut context = std::task::Context::from_waker(std::task::Waker::noop()); + assert!(replay.as_mut().poll(&mut context).is_pending()); + + assert_eq!(core.select(request(2)).await.expect("refresh").worker_id, 2); + let _lease = table + .commit(initializer, WorkerAffinityTarget::new(2, Some(0))) + .expect("bind elsewhere"); + assert!(matches!(replay.await, Err(SelectionError::BadRequest(_)))); + wait_until("rejected booking release", || { + !entry.scheduler.has_request("pending") + }) + .await; + + let replacement = core + .create_reservation(replay_reservation("pending")) + .await + .expect("new selection survives the old replay failure"); + assert_eq!(replacement.worker_id, 2); +} + +#[tokio::test(start_paused = true)] +async fn cancelled_free_releases_reservation_and_affinity() { + let mut config = test_config(false); + config.router_queue_threshold = Some(0.0); + let core = core_with( + config, + SelectionHost::default(), + None, + WorkerType::Aggregated, + Some(SessionAffinityConfig::new(Duration::from_secs(1))), + ); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + core.select_and_reserve(session_reservation("r1", "s")) + .await + .expect("booking"); + let entry = core.entry(&default_key()).expect("entry"); + + // An unpolled free must leave the live booking and its lease untouched. + drop(core.free_reservation("r1")); + assert!(entry.scheduler.has_request("r1")); + assert_eq!(lease_count(&core, "s"), Some(1)); + + let mut freeing = Box::pin(core.free_reservation("r1")); + let mut context = std::task::Context::from_waker(std::task::Waker::noop()); + assert!(freeing.as_mut().poll(&mut context).is_pending()); + assert!(!entry.scheduler.has_request("r1")); + // The queue actor has not acknowledged the release yet. + drop(freeing); + assert!(!core.reservation_index.read().contains_key("r1")); + assert_eq!(lease_count(&core, "s"), Some(0)); + + tokio::time::advance(Duration::from_secs(2)).await; + assert_eq!(bound_worker(&core, "s"), None); + core.select_and_reserve(session_reservation("r1", "s")) + .await + .expect("cancelled cleanup must not prevent ID reuse"); +} + #[tokio::test] async fn hard_mode_rejects_dispatch_away_from_a_live_binding() { let (core, first) = bound_session(SessionAffinityMode::Hard).await; diff --git a/lib/kv-router/src/services/selection/input.rs b/lib/kv-router/src/services/selection/input.rs index 048622692f05..f14702e5a673 100644 --- a/lib/kv-router/src/services/selection/input.rs +++ b/lib/kv-router/src/services/selection/input.rs @@ -169,7 +169,7 @@ impl PromptView<'_> { Ok(normalize_hashes(block_hashes, sequence_hashes, isl_tokens)?.block_hashes) } - fn routing_tokens_and_mm_infos(&self) -> Option> { + pub(super) fn routing_tokens_and_mm_infos(&self) -> Option> { if let Some(mm_routing_info) = self.mm_routing_info && !mm_routing_info.routing_token_ids.is_empty() { diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index db8d3c5e27de..c0302695ff1b 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -2850,6 +2850,29 @@ mod tests { assert_eq!(overlap_blocks, 0); } + #[tokio::test] + async fn test_find_best_match_maps_overload_to_resource_exhausted() { + let router = make_test_router(SelectionPolicySource::Registry, None).await; + router.client.set_overloaded_instances(&[0, 1]); + + let Err(error) = find_best_match(&router, &[11, 12], false).await else { + panic!("overloaded pool must not route"); + }; + assert!(dynamo_runtime::error::match_error_chain( + error.as_ref(), + &[ErrorType::ResourceExhausted], + &[] + )); + assert!( + error + .to_string() + .contains("all eligible workers are overloaded") + ); + + router.client.set_overloaded_instances(&[]); + assert!(find_best_match(&router, &[11, 12], false).await.is_ok()); + } + #[tokio::test] async fn test_find_best_match_maps_filtered_workers_to_unavailable() { let policy = SelectionPolicySource::Factory(Arc::new(|config: &KvRouterConfig, _, _| { From 3335e4bacc3ca55d54aaceb3ca8976d124a04b9b Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Wed, 16 Sep 2026 23:32:56 +0000 Subject: [PATCH 101/102] refactor(router): remove unused indexer lookup wrappers Signed-off-by: Ishan Dhanani --- .../src/indexer/lower_tier_indexers.rs | 46 ++++--------------- lib/kv-router/src/scheduling/queue.rs | 2 +- lib/kv-router/src/services/indexer/lookup.rs | 21 ++++----- 3 files changed, 19 insertions(+), 50 deletions(-) diff --git a/lib/kv-router/src/indexer/lower_tier_indexers.rs b/lib/kv-router/src/indexer/lower_tier_indexers.rs index a6b0e945ddf5..412e97053303 100644 --- a/lib/kv-router/src/indexer/lower_tier_indexers.rs +++ b/lib/kv-router/src/indexer/lower_tier_indexers.rs @@ -207,23 +207,6 @@ pub struct LowerTierQueryOptions { pub retain_kv_transfer_chain: bool, } -/// Walk every allocated lower tier in [`lower_tier_query_order`] and build a -/// per-tier match map seeded from `device_matches`. Per-worker continuations -/// flow forward: a worker that matched N device blocks starts the host walk -/// at block N (anchored on its last device hash), and so on. -pub fn query_lower_tiers( - indexers: &LowerTierIndexers, - sequence: &[LocalBlockHash], - device_matches: &MatchDetails, -) -> HashMap { - query_lower_tiers_with_options( - indexers, - sequence, - device_matches, - LowerTierQueryOptions::default(), - ) -} - fn merge_kv_transfer_tier_candidates( device_candidates: Option<&KvTransferCandidates>, tier_matches: &LowerTierMatchDetails, @@ -286,6 +269,10 @@ fn merge_kv_transfer_tier_candidates( }) } +/// Walk every allocated lower tier in [`lower_tier_query_order`] and build a +/// per-tier match map seeded from `device_matches`. Per-worker continuations +/// flow forward: a worker that matched N device blocks starts the host walk +/// at block N (anchored on its last device hash), and so on. pub fn query_lower_tiers_with_options( indexers: &LowerTierIndexers, sequence: &[LocalBlockHash], @@ -305,24 +292,6 @@ pub fn query_lower_tiers_with_options( ) } -pub fn query_lower_tiers_with_options_and_projection( - indexers: &LowerTierIndexers, - sequence: &[LocalBlockHash], - device_matches: &MatchDetails, - options: LowerTierQueryOptions, - projection: &ResidencyProjection, -) -> HashMap { - query_lower_tiers_with_options_and_snapshot( - indexers, - sequence, - device_matches, - options, - Arc::new(ResidencyRoutingSnapshot::from_projection( - projection.clone(), - )), - ) -} - pub fn query_lower_tiers_with_options_and_snapshot( indexers: &LowerTierIndexers, sequence: &[LocalBlockHash], @@ -563,7 +532,12 @@ mod tests { }; let sequence = vec![LocalBlockHash(1), LocalBlockHash(2)]; - let result = query_lower_tiers(&indexers, &sequence, &device_matches); + let result = query_lower_tiers_with_options( + &indexers, + &sequence, + &device_matches, + LowerTierQueryOptions::default(), + ); assert!(result.is_empty()); } diff --git a/lib/kv-router/src/scheduling/queue.rs b/lib/kv-router/src/scheduling/queue.rs index 07a9f5ef6441..358af306a6fd 100644 --- a/lib/kv-router/src/scheduling/queue.rs +++ b/lib/kv-router/src/scheduling/queue.rs @@ -3307,7 +3307,7 @@ policy_classes: "expected NoEndpoints before worker updates, got {resp:?}" ); - // Lazily register two workers in the slot tracker (EPP supplies pod list) + // Add two workers to the slot tracker, then publish their configs. slots.upsert_worker(WorkerDpRange::new(100, 0, 1)).unwrap(); slots.upsert_worker(WorkerDpRange::new(200, 0, 1)).unwrap(); diff --git a/lib/kv-router/src/services/indexer/lookup.rs b/lib/kv-router/src/services/indexer/lookup.rs index 32d9742f2520..9cb50faf99b9 100644 --- a/lib/kv-router/src/services/indexer/lookup.rs +++ b/lib/kv-router/src/services/indexer/lookup.rs @@ -135,14 +135,6 @@ impl Indexer { .await } - pub async fn find_matches_by_tier_ref( - &self, - sequence: &[LocalBlockHash], - ) -> Result { - self.find_matches_by_tier_ref_with_options(sequence, LowerTierQueryOptions::default()) - .await - } - pub async fn find_matches_by_tier_ref_with_options( &self, sequence: &[LocalBlockHash], @@ -197,7 +189,7 @@ impl<'a> LookupPipeline<'a> { // Seed lower-tier continuations from confirmed primary matches // only. Predict-on-route side scores are unconfirmed; using // them as lower-tier anchors would over-credit host/disk cache - // hits and break the score/hash lockstep `query_lower_tiers` + // hits and break the score/hash lockstep `query_lower_tiers_with_options` // expects. let primary_device = self .primary @@ -366,7 +358,7 @@ impl<'a> PrimaryLookup<'a> { /// `overlap_scores.scores` <-> `last_matched_hashes` lockstep. Side-only /// workers gain a score with no paired hash by design. The result is safe /// for scheduling / cache-hit signal but MUST NOT be used to seed -/// `query_lower_tiers`, which assumes the lockstep invariant. The local +/// `query_lower_tiers_with_options`, which assumes the lockstep invariant. The local /// arm of `find_matches_by_tier` enforces this by running the lower-tier /// query against primary-only `MatchDetails` before merging side scores. fn merge_overlap_scores(mut primary: MatchDetails, side: OverlapScores) -> MatchDetails { @@ -424,7 +416,7 @@ mod tests { use tokio_util::sync::CancellationToken; - use crate::indexer::LowerTierIndexers; + use crate::indexer::{LowerTierIndexers, LowerTierQueryOptions}; use crate::services::indexer::backend::Indexer; use crate::services::indexer::backend::test_util::store_event; use crate::{ @@ -1066,7 +1058,10 @@ mod tests { side_for_flush.flush().await; let query = vec![LocalBlockHash(11), LocalBlockHash(12), LocalBlockHash(13)]; - let borrowed = indexer.find_matches_by_tier_ref(&query).await.unwrap(); + let borrowed = indexer + .find_matches_by_tier_ref_with_options(&query, LowerTierQueryOptions::default()) + .await + .unwrap(); let owned = indexer.find_matches_by_tier(query).await.unwrap(); assert_eq!( @@ -1168,7 +1163,7 @@ mod tests { /// Regression test: when a worker has blocks in both device and lower-tier /// storage (e.g. same prefix stored on GPU and offloaded to host), the /// Concurrent indexer doesn't return last_matched_hashes. Without the fix, - /// query_lower_tiers would re-query that worker from root in the lower tier, + /// query_lower_tiers_with_options would re-query that worker from root in the lower tier, /// double-counting overlap blocks and producing cached_tokens > ISL. #[tokio::test] async fn concurrent_tiered_query_does_not_double_count_device_and_lower_tier_overlap() { From 1b82fe24d7357351dc851f6d8cecd38dbb7b1fd5 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Thu, 17 Sep 2026 04:18:37 +0000 Subject: [PATCH 102/102] perf(router): borrow membership snapshots during reconciliation Avoid cloning the complete worker snapshot on every reconciliation pass. Borrow worker requests while computing changes and retain the owned snapshot for retries, cloning only records that need an upsert. Signed-off-by: Ishan Dhanani --- .../ext-proc/src/selector.rs | 4 +-- .../src/services/selection/membership.rs | 36 +++++++++---------- lib/llm/src/kv_router/embedded.rs | 2 +- 3 files changed, 21 insertions(+), 21 deletions(-) diff --git a/deploy/inference-gateway/ext-proc/src/selector.rs b/deploy/inference-gateway/ext-proc/src/selector.rs index d597b0b2e10a..2a96eec5e32b 100644 --- a/deploy/inference-gateway/ext-proc/src/selector.rs +++ b/deploy/inference-gateway/ext-proc/src/selector.rs @@ -399,7 +399,7 @@ models: async fn register(selector: &Selector, workers: Vec) { CatalogReconciler::new(Arc::clone(selector.service.core())) - .apply(workers) + .apply(&workers) .await .expect("reconcile should succeed"); } @@ -747,7 +747,7 @@ worker_selection: let duplicate = incomplete_registration(1); let error = CatalogReconciler::new(Arc::clone(selector.service.core())) - .apply(vec![duplicate.clone(), duplicate]) + .apply(&[duplicate.clone(), duplicate]) .await .expect_err("duplicate IDs must be rejected"); assert!(error.to_string().contains("duplicate worker_id 1")); diff --git a/lib/kv-router/src/services/selection/membership.rs b/lib/kv-router/src/services/selection/membership.rs index 1729ba39224c..9158303561fa 100644 --- a/lib/kv-router/src/services/selection/membership.rs +++ b/lib/kv-router/src/services/selection/membership.rs @@ -77,8 +77,8 @@ impl CatalogReconciler { /// not stop the pass: every other worker is still upserted and every /// worker absent from the snapshot is still deleted. The first error is /// returned so the caller retries the snapshot. - pub async fn apply(&mut self, desired: Vec) -> Result<(), SelectionError> { - let mut by_id: HashMap = HashMap::with_capacity(desired.len()); + pub async fn apply(&mut self, desired: &[WorkerRequest]) -> Result<(), SelectionError> { + let mut by_id: HashMap = HashMap::with_capacity(desired.len()); for request in desired { let worker_id = request.worker_id; if by_id.insert(worker_id, request).is_some() { @@ -89,7 +89,7 @@ impl CatalogReconciler { } let mut to_upsert: Vec = Vec::new(); - for (worker_id, request) in &by_id { + for (worker_id, &request) in &by_id { // Track before the upsert so a partially applied record is still // deleted when the worker leaves the desired set. self.tracked.insert(*worker_id); @@ -115,7 +115,7 @@ impl CatalogReconciler { observer.upserted(&record); } if record.lifecycle == WorkerLifecycle::Schedulable { - let request = by_id.get(&worker_id).expect("upserted id came from by_id"); + let request = *by_id.get(&worker_id).expect("upserted id came from by_id"); self.converged.insert(worker_id, request.clone()); } else { self.converged.remove(&worker_id); @@ -180,7 +180,7 @@ impl CatalogReconciler { pending.take().expect("guarded by pending.is_some()") } }; - let outcome = self.apply(snapshot.clone()).await; + let outcome = self.apply(&snapshot).await; #[cfg(test)] self.applies .fetch_add(1, std::sync::atomic::Ordering::SeqCst); @@ -288,22 +288,22 @@ mod tests { let core = core(); let (mut reconciler, counter) = reconciler(&core); - reconciler.apply(vec![incomplete(1)]).await.expect("apply"); + reconciler.apply(&[incomplete(1)]).await.expect("apply"); assert_eq!(counter.upserts(), 1); assert_eq!(lifecycle(&core, 1), Some(WorkerLifecycle::Incomplete)); // The same snapshot re-upserts rather than skipping the unconverged worker. - reconciler.apply(vec![incomplete(1)]).await.expect("apply"); + reconciler.apply(&[incomplete(1)]).await.expect("apply"); assert_eq!(counter.upserts(), 2); // The never-schedulable worker was still tracked, so leaving deletes it // and the catalog no longer lists it. - reconciler.apply(Vec::new()).await.expect("apply"); + reconciler.apply(&[]).await.expect("apply"); assert_eq!(lifecycle(&core, 1), None); assert_eq!(counter.removals(), 1); // Nothing is tracked any more: an empty snapshot deletes nothing. - reconciler.apply(Vec::new()).await.expect("apply"); + reconciler.apply(&[]).await.expect("apply"); assert_eq!(counter.removals(), 1); } @@ -312,17 +312,17 @@ mod tests { let core = core(); let (mut reconciler, counter) = reconciler(&core); - reconciler.apply(vec![schedulable(1)]).await.expect("apply"); + reconciler.apply(&[schedulable(1)]).await.expect("apply"); assert_eq!(lifecycle(&core, 1), Some(WorkerLifecycle::Schedulable)); assert_eq!(counter.upserts(), 1); // An identical desired record is converged and skips its upsert. - reconciler.apply(vec![schedulable(1)]).await.expect("apply"); + reconciler.apply(&[schedulable(1)]).await.expect("apply"); assert_eq!(counter.upserts(), 1); let mut moved = schedulable(1); moved.endpoint = Some("http://10.0.0.9:8000".to_string()); - reconciler.apply(vec![moved]).await.expect("apply"); + reconciler.apply(&[moved]).await.expect("apply"); assert_eq!(counter.upserts(), 2); let record = core .list_workers(None, None) @@ -339,7 +339,7 @@ mod tests { let (mut reconciler, counter) = reconciler(&core); reconciler - .apply(vec![schedulable(1), schedulable(2)]) + .apply(&[schedulable(1), schedulable(2)]) .await .expect("apply"); assert_eq!(counter.upserts(), 2); @@ -349,7 +349,7 @@ mod tests { let mut changed = schedulable(2); changed.endpoint = Some("http://10.0.0.9:8000".to_string()); let error = reconciler - .apply(vec![changed.clone()]) + .apply(std::slice::from_ref(&changed)) .await .expect_err("failed upsert is reported"); assert!(matches!(error, SelectionError::Internal(_)), "{error}"); @@ -360,7 +360,7 @@ mod tests { assert_eq!(reconciler.tracked, HashSet::from([2])); // The failed worker is still tracked and is retried once the fault clears. core.fail_upsert_for.lock().clear(); - reconciler.apply(vec![changed]).await.expect("apply"); + reconciler.apply(&[changed]).await.expect("apply"); assert_eq!(counter.upserts(), 3); } @@ -470,7 +470,7 @@ mod tests { let publishes = || core.publish_count.load(Ordering::SeqCst); reconciler - .apply((1..=4).map(schedulable).collect()) + .apply(&(1..=4).map(schedulable).collect::>()) .await .expect("apply"); assert_eq!( @@ -505,7 +505,7 @@ mod tests { let mut other = schedulable(5); other.routing_group = "other".to_string(); reconciler - .apply((1..=4).map(schedulable).chain([other]).collect()) + .apply(&(1..=4).map(schedulable).chain([other]).collect::>()) .await .expect("apply"); assert_eq!(publishes(), 2); @@ -521,7 +521,7 @@ mod tests { let mut reconciler = CatalogReconciler::new(Arc::clone(&core)); let error = reconciler - .apply(vec![schedulable(1), schedulable(1)]) + .apply(&[schedulable(1), schedulable(1)]) .await .expect_err("duplicates are rejected"); assert!( diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs index 0bda792bfd89..e60b18eb4e19 100644 --- a/lib/llm/src/kv_router/embedded.rs +++ b/lib/llm/src/kv_router/embedded.rs @@ -349,7 +349,7 @@ impl EmbeddedSelection { ); // The current membership is in the catalog before the router serves. if let Some(snapshot) = source.next_snapshot().await - && let Err(error) = reconciler.apply(snapshot).await + && let Err(error) = reconciler.apply(&snapshot).await { tracing::warn!(%error, "embedded selection: initial membership reconcile failed"); }