diff --git a/.github/workflows/pre-merge.yml b/.github/workflows/pre-merge.yml index b04b69723b47..9a442fabb52c 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 standalone slot-tracker feature tests if: matrix.dir == '.' working-directory: ${{ matrix.dir }} diff --git a/Cargo.lock b/Cargo.lock index d1c6a3dd6b95..5abb2e966fbc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2561,6 +2561,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/epp.rs b/deploy/inference-gateway/ext-proc/src/epp.rs index 085aa9a4137c..481b58802a28 100644 --- a/deploy/inference-gateway/ext-proc/src/epp.rs +++ b/deploy/inference-gateway/ext-proc/src/epp.rs @@ -483,10 +483,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, @@ -527,10 +523,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 @@ -934,8 +926,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 @@ -1392,11 +1385,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); @@ -2146,10 +2137,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 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. + /// 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(); @@ -2215,7 +2203,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] @@ -2450,8 +2438,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/deploy/inference-gateway/ext-proc/src/epp_router.rs b/deploy/inference-gateway/ext-proc/src/epp_router.rs index 022cd5144183..5f192f4d069e 100644 --- a/deploy/inference-gateway/ext-proc/src/epp_router.rs +++ b/deploy/inference-gateway/ext-proc/src/epp_router.rs @@ -26,8 +26,9 @@ 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::agents::HEADER_DYNAMO_SESSION_ID; use dynamo_llm::protocols::common::extensions::{ AgentHints, HEADER_REQUEST_PRIORITY, HEADER_REQUEST_STRICT_PRIORITY, resolve_request_priority, }; @@ -368,6 +369,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), // `None` on the ordinary path: the selector schedules over its // catalog; `Some` only carries an Envoy subset constraint. allowed_worker_ids: allowed, @@ -383,6 +385,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/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/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/pod_discovery.rs b/deploy/inference-gateway/ext-proc/src/pod_discovery.rs index 5b2c8744a8b6..39c49f1c7028 100644 --- a/deploy/inference-gateway/ext-proc/src/pod_discovery.rs +++ b/deploy/inference-gateway/ext-proc/src/pod_discovery.rs @@ -38,8 +38,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, } @@ -108,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())); @@ -184,7 +190,7 @@ impl PodDiscovery { Delta::Rebuild => rebuild_index( &store, pool_rx.borrow().as_ref(), - kv_event_port, + kv_ports, replay_port, &index_task, ), @@ -192,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), @@ -345,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 { @@ -383,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)); } } @@ -404,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) { @@ -424,7 +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_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))), }) } @@ -436,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())]), @@ -480,13 +543,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")); } @@ -500,13 +566,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")); } @@ -521,7 +590,7 @@ mod tests { &[("app", "vllm-qwen")] ), &pool(), - 5557, + single_rank(5557), None, ) .is_none() @@ -539,7 +608,7 @@ mod tests { &[("app", "something-else")] ), &pool(), - 5557, + single_rank(5557), None, ) .is_none() @@ -557,7 +626,7 @@ mod tests { &[("app", "vllm-qwen")] ), &pool(), - 5557, + single_rank(5557), None, ) .is_none() @@ -573,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(), 5557, None).is_none()); + assert!(raw_worker_from_pod(&p, &pool(), single_rank(5557), None).is_none()); } #[test] @@ -582,7 +651,7 @@ mod tests { raw_worker_from_pod( &pod("vllm-0", None, Some(true), &[("app", "vllm-qwen")]), &pool(), - 5557, + single_rank(5557), None, ) .is_none() @@ -603,6 +672,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", @@ -620,17 +690,11 @@ mod tests { ]); let index = RwLock::new(WorkerIndex::new()); - assert!(rebuild_index( - &store, - Some(&pool()), - 5557, - Some(5560), - &index - )); + assert!(rebuild_index(&store, Some(&pool()), kp, Some(5560), &index)); assert!(!rebuild_index( &store, Some(&pool()), - 5557, + kp, Some(5560), &index )); @@ -647,6 +711,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"), @@ -654,12 +719,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, 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( @@ -670,8 +736,8 @@ 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()), 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") @@ -684,12 +750,12 @@ 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()), 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()), 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)); @@ -697,11 +763,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()), 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()); @@ -713,16 +780,17 @@ mod tests { &[("app", "vllm-qwen")], ); - assert!(upsert_pod(&index, &ready, Some(&pool()), 5557, None)); + assert!(upsert_pod(&index, &ready, Some(&pool()), kp, 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, 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( @@ -746,7 +814,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()), 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 @@ -772,13 +840,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), - 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!( @@ -795,6 +857,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 @@ -810,7 +873,7 @@ mod tests { &[("app", "vllm-qwen")], ), Some(&pool()), - 5557, + kp, None, ); upsert_pod( @@ -822,7 +885,7 @@ mod tests { &[("app", "vllm-qwen")], ), Some(&pool()), - 5557, + kp, None, ); assert_eq!(index.read().unwrap().len(), 2); @@ -834,7 +897,7 @@ mod tests { Some(true), &[("app", "vllm-qwen")], )]); - assert!(rebuild_index(&store, Some(&pool()), 5557, None, &index)); + assert!(rebuild_index(&store, Some(&pool()), kp, None, &index)); let index = index.read().unwrap(); assert_eq!(index.len(), 1); @@ -854,7 +917,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 92d9796700b5..2a96eec5e32b 100644 --- a/deploy/inference-gateway/ext-proc/src/selector.rs +++ b/deploy/inference-gateway/ext-proc/src/selector.rs @@ -3,44 +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, + 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, - 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 { @@ -54,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, pub cache_namespace: Option, } @@ -80,23 +67,10 @@ 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, - 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 { @@ -142,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() @@ -166,11 +146,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( @@ -191,92 +167,13 @@ 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 size is not yet implemented. - // Default support for single rank DP. - data_parallel_size: Some(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(()) - } - /// 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, @@ -292,7 +189,8 @@ 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, affinity_target: None, @@ -303,8 +201,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, @@ -358,10 +255,14 @@ 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}; - 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, @@ -451,27 +352,24 @@ 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, } } /// 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, - 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) } } @@ -483,22 +381,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, + 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.service.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 { @@ -511,6 +415,7 @@ models: strict_priority: None, expected_output_tokens: None, policy_class: None, + session_id: None, cache_namespace: None, } } @@ -521,10 +426,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" @@ -532,6 +434,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; @@ -597,10 +515,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 @@ -738,10 +653,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" @@ -774,69 +686,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()) @@ -897,8 +746,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.service.core())) + .apply(&[duplicate.clone(), duplicate]) .await .expect_err("duplicate IDs must be rejected"); assert!(error.to_string().contains("duplicate worker_id 1")); @@ -958,10 +807,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); @@ -998,10 +844,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"); @@ -1048,10 +891,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 @@ -1059,4 +899,19 @@ 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); + 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/topology_adapter.rs b/deploy/inference-gateway/ext-proc/src/topology_adapter.rs index 9fac749192d5..244381551259 100644 --- a/deploy/inference-gateway/ext-proc/src/topology_adapter.rs +++ b/deploy/inference-gateway/ext-proc/src/topology_adapter.rs @@ -1,20 +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::collections::HashMap; 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 { @@ -35,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. @@ -49,32 +87,17 @@ 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.service.core())) + .run(source, cancel.child_token()), + ); Self { cancel } } } @@ -85,42 +108,26 @@ 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 mut kv_events_endpoints = HashMap::new(); - kv_events_endpoints.insert(0u32, w.kv_events_endpoint); - - 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, - kv_events_endpoints, + 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() } } #[cfg(test)] mod tests { + use std::collections::HashMap; use std::time::Duration; use super::*; @@ -138,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, } } @@ -161,23 +171,27 @@ 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, } } #[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/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 b0624e653124..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 @@ -515,7 +515,9 @@ before relying on cache overlap, subject to the limits in [Gap Detection and Replay](#gap-detection-and-replay). The [standalone selector](standalone-selection.md#worker-restarts) has a different -registration API: updating a schedulable worker reconciles its indexer registration. +registration API: updating a schedulable worker reconciles its indexer registration +per rank, keeping listeners whose endpoints did not change, and a restart at the same +address needs `DELETE /workers/{worker_id}` before the new `POST /workers`. ## Limitations 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 d8ddea359f41..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 @@ -68,6 +68,7 @@ APIs. Those bindings should wrap `SelectionService` rather than construct | `--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). | | `--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. | @@ -79,6 +80,24 @@ 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 | + +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 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 admission by default. The embedded `KvRouter` uses the same `300`-second value as its shared @@ -88,7 +107,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 @@ -105,11 +127,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 @@ -117,6 +142,23 @@ 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" } +} +``` + +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; @@ -131,12 +173,17 @@ When a worker or its cache-event publisher restarts, update its registration on selector replica. A reconnect at the same address does not reset the existing event sequence watermark or invalidate old cache ownership. -With KV events enabled, `POST /workers` for an existing schedulable worker drains it, -removes its indexer registration, and creates new listeners with fresh sequence -watermarks. Supply the complete current worker record because this endpoint replaces -the catalog record. `PATCH /workers/{worker_id}` also reconciles a schedulable worker; -even an unchanged update can rebuild its cache view, so avoid using registration -updates as periodic heartbeats. Serialize lifecycle operations for the same worker. +With KV events enabled, `POST /workers` or `PATCH /workers/{worker_id}` for an +existing schedulable worker reconciles its indexer registration rank by rank: a rank +whose event endpoint and replay endpoint are unchanged keeps its listener, its +sequence watermark, and its indexed blocks, so capacity or topology updates do not +drain the cache view. A rank whose endpoint changed gets a new listener with a fresh +watermark. Because an unchanged update keeps the old watermark, it does not recover a +publisher that restarted at the same address: for that, call +`DELETE /workers/{worker_id}`, wait for the response, then `POST /workers` with the +current record, or register the restarted publisher under a new `kv_events_endpoint`. +Supply the complete current worker record to `POST /workers` because it replaces the +catalog record. Serialize lifecycle operations for the same worker. This differs from the standalone indexer's `/register`, which rejects duplicate registrations and requires `/unregister` first. See @@ -201,12 +248,80 @@ 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. + +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 +{ + "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] + } + } + ] + } +} +``` + +`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 +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 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). @@ -231,26 +346,31 @@ 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 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` + +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, + "input_trigger": "tool_result" + } +} +``` + +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/docs/fern/pages/kubernetes/kv-aware-routing/vanilla-vllm-onramp.mdx b/docs/fern/pages/kubernetes/kv-aware-routing/vanilla-vllm-onramp.mdx index eff44badb87c..ec91a4bc9b05 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 @@ -502,11 +502,14 @@ 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`, 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`. | | `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 (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/kvbm/Cargo.lock b/lib/bindings/kvbm/Cargo.lock index 09e2c27c77aa..cf4d22d260d7 100644 --- a/lib/bindings/kvbm/Cargo.lock +++ b/lib/bindings/kvbm/Cargo.lock @@ -2085,7 +2085,9 @@ dependencies = [ "anyhow", "arc-swap", "async-trait", + "axum", "blake3", + "chrono", "crossbeam-queue", "dashmap", "derive_builder", @@ -2096,9 +2098,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", @@ -2110,9 +2114,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" diff --git a/lib/bindings/python/rust/lib.rs b/lib/bindings/python/rust/lib.rs index 3940208cac31..7ae9539e9fb8 100644 --- a/lib/bindings/python/rust/lib.rs +++ b/lib/bindings/python/rust/lib.rs @@ -460,6 +460,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/bindings/python/rust/llm/entrypoint.rs b/lib/bindings/python/rust/llm/entrypoint.rs index 07a8eee9eeb4..f17608631db7 100644 --- a/lib/bindings/python/rust/llm/entrypoint.rs +++ b/lib/bindings/python/rust/llm/entrypoint.rs @@ -509,11 +509,7 @@ 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", - )); - } + 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 295bc29105f7..3336d017dcbf 100644 --- a/lib/bindings/python/rust/llm/kv.rs +++ b/lib/bindings/python/rust/llm/kv.rs @@ -20,15 +20,12 @@ 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; 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")] @@ -43,24 +40,21 @@ 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}; -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; @@ -372,6 +366,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, @@ -554,6 +553,11 @@ where indexer_peers: cli.indexer_peers, replica_sync_port: cli.replica_sync_port, replica_sync_peers: cli.replica_sync_peers, + session_affinity_ttl: cli + .session_affinity_ttl_secs + .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, @@ -630,6 +634,22 @@ impl SelectionCacheConfig { } } +fn session_affinity_ttl_from_secs(ttl: f64) -> Result { + 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)) +} + +/// 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] @@ -642,7 +662,8 @@ 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, session_affinity_ttl_secs = None))] + #[allow(clippy::too_many_arguments)] fn new( py: Python<'_>, indexer_threads: usize, @@ -650,6 +671,7 @@ impl SelectionService { replica_sync_port: Option, replica_sync_peers: Option>, selection_cache: 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() { @@ -672,6 +694,11 @@ impl SelectionService { if let Some(port) = replica_sync_port { builder = builder.replica_sync(port, replica_sync_peers); } + if let Some(ttl) = session_affinity_ttl_secs { + builder = builder.session_affinity( + session_affinity_ttl_from_secs(ttl).map_err(PyValueError::new_err)?, + ); + } let inner = py .allow_threads(|| crate::bridge_runtime().block_on(builder.build())) .map_err(to_pyerr)?; @@ -959,8 +986,10 @@ 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)).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); @@ -2004,14 +2033,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, @@ -2023,42 +2071,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, @@ -2214,11 +2226,7 @@ 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", - )); - } + check_session_affinity_ttl_secs(session_affinity_ttl_secs)?; let session_affinity_mode = session_affinity_mode .parse::() .map_err(PyValueError::new_err)?; @@ -2505,7 +2513,7 @@ impl KvRouter { ) .await .map_err(to_pyerr)?; - let (outcome, attempt) = admitted.into_parts(); + let (outcome, booking) = admitted.into_parts(); let (best_worker, overlap_blocks) = match outcome { llm_rs::kv_router::FindBestMatchOutcome::Routed { worker, @@ -2541,19 +2549,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(booking) = booking { chooser - .enroll_public_request_attempt( - request_id, - best_worker, - attempt_id, - 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/bindings/python/src/dynamo/_core.pyi b/lib/bindings/python/src/dynamo/_core.pyi index d1d59cb61cec..1f8df6c915e0 100644 --- a/lib/bindings/python/src/dynamo/_core.pyi +++ b/lib/bindings/python/src/dynamo/_core.pyi @@ -748,8 +748,13 @@ class SelectionService: replica_sync_port: Optional[int] = None, replica_sync_peers: Optional[list[str]] = None, selection_cache: Optional[SelectionCacheConfig] = 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..e7643d8baf61 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,53 @@ 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] +) +@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) + + +@pytest.mark.skipif( + SelectionService is None, + 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( + [ + 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]) +@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() + + @pytest.mark.skipif( SelectionService is None, reason="SelectionService requires the select-service Cargo feature", diff --git a/lib/kv-router/Cargo.toml b/lib/kv-router/Cargo.toml index 371e9cda2ee9..a778cc5ce945 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 @@ -73,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] @@ -100,3 +102,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/indexer/lower_tier_indexers.rs b/lib/kv-router/src/indexer/lower_tier_indexers.rs index f2d067fcb6a2..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], @@ -408,6 +377,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() } @@ -534,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/lib.rs b/lib/kv-router/src/lib.rs index 9f7eeba5f3f2..c3a9a92f19ec 100644 --- a/lib/kv-router/src/lib.rs +++ b/lib/kv-router/src/lib.rs @@ -70,6 +70,7 @@ 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::policy::{FcfsPolicy, RouterSchedulingPolicy, SchedulingPolicy, WsptPolicy}; pub use scheduling::{ diff --git a/lib/kv-router/src/scheduling/local.rs b/lib/kv-router/src/scheduling/local.rs index f480ffa9b8a7..4cb363064621 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, + BookingHandle, ClassQueueStats, 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,20 @@ where &self, request: ScheduleRequest, ) -> Result { + let (admitted, booking) = self.schedule_request_with_booking(request).await?; + if let Some(booking) = booking { + let _ = booking.commit(); + } + Ok(admitted) + } + + /// 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`. + pub(crate) async fn schedule_request_with_booking( + &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 +313,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 +335,12 @@ where } else { AdmissionAttempt::Untracked }; - if let Some(lease) = lifecycle_lease.as_mut() { - lease.disarm(); - } - Ok(AdmittedSchedulingResponse { response, attempt }) + // 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. @@ -486,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()) } @@ -505,11 +518,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 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> { - self.slots.add_request_if_registered(req, Instant::now()) + ) -> 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.booking_handle(SchedulerBookingDescriptor { + request_id, + worker, + attempt_id, + })) } pub async fn mark_prefill_completed(&self, request_id: &str) -> Result<(), SequenceError> { @@ -566,21 +591,72 @@ 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).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 { + 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, + booking.worker, + booking.attempt_id, + Instant::now(), + )?; + cleanup(); + 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() } @@ -593,10 +669,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 } @@ -625,6 +697,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, @@ -1489,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/lora_filter.rs b/lib/kv-router/src/scheduling/lora_filter.rs new file mode 100644 index 000000000000..03dba81ad6e8 --- /dev/null +++ b/lib/kv-router/src/scheduling/lora_filter.rs @@ -0,0 +1,181 @@ +// 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 rustc_hash::FxHashSet; + +use crate::protocols::{WorkerId, WorkerWithDpRank}; + +/// Host-owned knowledge of which workers can serve a LoRA adapter. +/// +/// 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`. +/// - 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, + 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; + } + // A filter may only narrow: ids it returns from outside `base` are dropped. + let preferred: FxHashSet = filter + .filter_worker_ids_for_lora(lora_name, &base) + .into_iter() + .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) + { + 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])) + ); + } + + /// 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 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!( + 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..82eab3885b66 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; +mod lora_filter; pub mod overlap; pub mod overlap_refresh; pub mod policy; @@ -19,6 +20,9 @@ mod worker_selection_config; 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, SelectedWorkerTierSnapshot, SharedCacheOverlapScore, WorkerOverlapScore, 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/scheduling/queue.rs b/lib/kv-router/src/scheduling/queue.rs index 6d64eec2848a..358af306a6fd 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, @@ -362,6 +361,56 @@ impl SchedulerBookingCleanup { } } +/// 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)] +#[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 { + /// Hand the booking to a longer-lived owner; the handle stops guarding it. + #[must_use] + pub fn commit(mut self) -> SchedulerBookingDescriptor { + self.armed = false; + 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. + 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, @@ -384,6 +433,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 { @@ -449,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< @@ -556,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, @@ -575,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, @@ -592,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. @@ -723,6 +761,15 @@ impl< } } + /// 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. /// /// This is for advisory policy probes that must not wait in the router @@ -758,7 +805,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 +814,6 @@ impl< ack_rx .await .map_err(|_| KvSchedulerError::SubscriberShutdown)? - .map(|_| ()) .map_err(|error| KvSchedulerError::BookingFailed(error.to_string())) } @@ -845,10 +891,6 @@ impl< }) } - pub fn supports_overlap_refresh(&self) -> bool { - self.supports_overlap_refresh - } - fn prepare_block_hashes_for_refresh( &self, block_hashes: Option>, @@ -1661,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}; @@ -1704,6 +1747,82 @@ 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 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. + 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>>, } @@ -3167,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; @@ -3186,10 +3304,10 @@ 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) + // 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(); @@ -3228,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/sequences/multi_worker.rs b/lib/kv-router/src/sequences/multi_worker.rs index 7b63b7f24a7e..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) {} } // --------------------------------------------------------------------------- @@ -281,9 +286,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 +448,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, @@ -787,7 +792,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( @@ -815,10 +829,43 @@ 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) } + 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 @@ -2449,10 +2496,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] @@ -2481,10 +2525,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] @@ -3036,14 +3077,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) @@ -3070,14 +3105,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] @@ -3140,11 +3169,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) ); } @@ -3197,10 +3226,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); @@ -3251,10 +3277,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)] @@ -3290,10 +3313,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)] @@ -3372,10 +3392,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)] @@ -3397,10 +3414,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] @@ -3661,10 +3675,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)); } @@ -3693,10 +3704,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()); } @@ -3725,10 +3733,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/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/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 efe97053ed4a..5ca2dd2d0157 100644 --- a/lib/kv-router/src/services/common/replica_sync.rs +++ b/lib/kv-router/src/services/common/replica_sync.rs @@ -14,12 +14,73 @@ 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::sequences::{ + SchedulerLoadSnapshot, SequencePublishQueueError, SequencePublisher, SequenceSubscriber, +}; +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; 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 { + #[serde(flatten)] + pub partition: RoutingPartitionId, + pub session_id: String, + pub worker_id: u64, + pub dp_rank: Option, + pub sequence: u64, + pub writer_id: u64, +} + +#[cfg(feature = "standalone-selection")] +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, + } + } +} + +#[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 { + partition: RoutingPartitionId, + tx: mpsc::Sender, +} + +#[cfg(feature = "standalone-selection")] +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, + 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 +110,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 +158,35 @@ 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 } + #[cfg(feature = "standalone-selection")] + /// Sink that publishes session bindings into the mesh, when this runtime + /// carries them. + 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 { event.router_id == self.process_id } @@ -117,9 +199,62 @@ 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: 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); + } + } +} + +/// 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: +/// 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: Option>, + 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, + /// Optional drain-batch observer for the inbound funnel. + pub ingress_observer: Option>, +} + +#[cfg(feature = "standalone-selection")] +/// Per-partition factory for [`HostReplicaChannels`]; `None` disables replica +/// sync for that partition. +pub type HostReplicaSyncFactory = + Arc Option + Send + Sync>; + #[derive(Clone)] -pub(crate) struct ScopedSequencePublisher { +pub struct ScopedSequencePublisher { replica: Option, + host_tx: Option>, + load_sink: Option>, } #[derive(Clone)] @@ -132,7 +267,24 @@ struct ScopedReplicaPublisher { impl ScopedSequencePublisher { pub(crate) fn disabled() -> Self { - Self { replica: None } + Self { + replica: None, + host_tx: None, + load_sink: None, + } + } + + pub(crate) fn host(outbound: mpsc::Sender) -> Self { + Self { + replica: None, + host_tx: Some(outbound), + load_sink: None, + } + } + + pub(crate) fn with_load_sink(mut self, sink: Option>) -> Self { + self.load_sink = sink; + self } pub(crate) fn enabled( @@ -148,12 +300,27 @@ impl ScopedSequencePublisher { tx, cancel_token, }), + host_tx: None, + load_sink: None, } } } 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(()); }; @@ -186,25 +353,46 @@ 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) { + if let Some(sink) = &self.load_sink { + sink.publish_batch(loads); + } + } 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); + } } } 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 } } } @@ -219,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 { @@ -251,10 +445,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)), })) @@ -264,8 +459,28 @@ 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 { + let enabled = host.outbound.is_some(); + return ScopedReplicaSync { + publisher: host + .outbound + .map(ScopedSequencePublisher::host) + .unwrap_or_else(ScopedSequencePublisher::disabled), + enabled, + process_id: host.process_id, + channel: Some(( + host.inbound_tx, + ChannelSequenceSubscriber::with_observer( + host.inbound_rx, + host.ingress_observer, + ), + )), + }; + } return ScopedReplicaSync { publisher: ScopedSequencePublisher::disabled(), enabled: false, @@ -291,53 +506,69 @@ 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, label) = 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], + format!("{partition} request_id={}", event.event.request_id), + ), + 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], + format!("session_id={}", binding.session_id), + ), + 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!(event = %label, "Failed to publish replica event: {error}"); } } }); - Ok((tx, task)) + Ok((tx, affinity_tx, task)) } #[derive(Debug, thiserror::Error)] @@ -375,6 +606,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, @@ -383,7 +615,30 @@ 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::, + ) + } + + 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) @@ -411,7 +666,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}"); } @@ -524,31 +779,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)] @@ -695,7 +959,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| { @@ -743,9 +1007,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(), @@ -829,4 +1093,113 @@ 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, + 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, + max_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); + self.max_depth.fetch_max(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); + // 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/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/indexer/backend.rs b/lib/kv-router/src/services/indexer/backend.rs index 65cb770596ee..978479464bcc 100644 --- a/lib/kv-router/src/services/indexer/backend.rs +++ b/lib/kv-router/src/services/indexer/backend.rs @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 use std::sync::Arc; +use std::time::Duration; use anyhow::Result; use async_trait::async_trait; @@ -9,31 +10,318 @@ use tokio_util::sync::CancellationToken; use crate::ConcurrentRadixTreeCompressed; use crate::ThreadPoolIndexer; +use crate::approx::PruneConfig; +use crate::config::{ApproximateCachePolicyKind, KvRouterConfig}; use crate::indexer::{ - KvIndexer, KvIndexerInterface, KvIndexerMetrics, KvRouterError, LowerTierIndexers, - MatchDetails, TieredMatchDetails, TieredMatchProvider, query_lower_tiers, + ApproximateLruIncarnation, ApproximateLruStats, ApproximateRetentionConfig, KvIndexer, + KvIndexerInterface, KvIndexerMetrics, KvRouterError, LowerTierIndexers, LowerTierQueryOptions, + MatchDetails, RoutingDecisionHashes, SyncIndexer, TieredMatchDetails, TieredMatchProvider, + record_unsupported_residency_event, }; -use crate::protocols::{KvCacheEventData, LocalBlockHash, OverlapScores, RouterEvent, WorkerId}; +use crate::protocols::{ + DpRank, ExternalSequenceBlockHash, KvCacheEventData, LocalBlockHash, OverlapScores, + ResidencyProjection, ResidencyRoutingSnapshot, RouterEvent, WorkerId, WorkerWithDpRank, +}; + +use super::lookup::{HashInput, merge_side_or_warn}; +use super::session_updates::{SessionMutation, SessionUpdateSender}; + +/// 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), +} + +/// 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, + } + } + + /// Resolve the indexer shape for `config`. + /// + /// 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 + { + 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 { + KvIndexer(KvIndexer), + Concurrent(Arc>), +} + +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 { + return Self::Concurrent(Arc::new(ThreadPoolIndexer::new_with_metrics_and_pruning( + ConcurrentRadixTreeCompressed::new(), + num_threads, + block_size, + Some(metrics), + prune_config, + ))); + } + Self::KvIndexer(KvIndexer::new_with_pruning( + cancel, + 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 record_routing_decision( + &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 + } + } + } + + 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 + } + } + } + + async fn remove_worker(&self, worker_id: WorkerId) { + match self { + 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::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, + } + } +} + +/// 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( + &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; +} -/// Block-content indexer wrapping a primary device-tier backend plus a -/// per-tier registry of lower-tier indexers (host-pinned, disk, …). +/// 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. /// -/// 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, + session_updates: Option, }, Concurrent { primary: Arc>, lower_tier: LowerTierIndexers, + approx: Option, + primary_records_routing_decisions: bool, + session_updates: Option, }, + /// 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, + approx: Option, + primary_records_routing_decisions: bool, + }, + /// No KV awareness. + None, } impl Indexer { + fn approx(&self) -> Option<&SideIndexer> { + match self { + Self::Single { approx, .. } + | Self::Concurrent { approx, .. } + | Self::Remote { approx, .. } => approx.as_ref(), + Self::None => None, + } + } + + pub fn supports_overlap_refresh(&self) -> bool { + matches!(self, Self::Single { .. } | Self::Concurrent { .. }) + } + + pub fn supports_kv_transfer_chain_retention(&self) -> bool { + matches!( + self, + Self::Single { + approx: None, + primary_records_routing_decisions: false, + .. + } | Self::Concurrent { + approx: None, + primary_records_routing_decisions: false, + .. + } + ) + } + + /// Publish a control-plane residency projection for subsequent lookups. + pub fn set_residency_projection(&self, projection: ResidencyProjection) { + match self { + Self::Single { 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::Single { lower_tier, .. } | Self::Concurrent { lower_tier, .. } => { + lower_tier.set_residency_routing_snapshot(snapshot) + } + Self::Remote { .. } | Self::None => {} + } + } + /// 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,13 +329,17 @@ impl Indexer { match self { Indexer::Single { primary, .. } => primary.apply_event(event).await, Indexer::Concurrent { primary, .. } => primary.apply_event(event).await, + Indexer::Remote { .. } | Indexer::None => { + tracing::trace!("Dropping KV event: no local primary 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. + /// 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) -> Result<(), KvRouterError> { let targets_primary = match event.targets_primary() { Ok(targets_primary) => targets_primary, @@ -56,6 +348,7 @@ impl Indexer { Self::Single { lower_tier, .. } | Self::Concurrent { lower_tier, .. } => { lower_tier.record_unsupported_residency_event(&event); } + Self::Remote { .. } | Self::None => {} } return Ok(()); } @@ -65,6 +358,7 @@ impl Indexer { Indexer::Single { primary, lower_tier, + .. } => { if is_clear { let mut reset_error = None; @@ -95,6 +389,7 @@ impl Indexer { Indexer::Concurrent { primary, lower_tier, + .. } => { if is_clear { let mut reset_error = None; @@ -122,15 +417,246 @@ impl Indexer { .await; } } + 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) -> 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 session_update = if targets_primary { + match self { + Self::Single { + session_updates, .. + } + | Self::Concurrent { + session_updates, .. + } => session_updates.as_ref().and_then(|sender| { + SessionMutation::from_event(&event).map(|mutation| (sender.clone(), mutation)) + }), + Self::Remote { .. } | Self::None => None, + } + } else { + None + }; + 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 => { + tracing::trace!("Dropping KV event: no local primary indexer"); + } + } + if let Some((sender, update)) = session_update { + sender.enqueue(update)?; + } + Ok(()) + } + + pub fn session_residency_version(&self, worker: WorkerWithDpRank) -> Option { + match self { + Self::Single { + session_updates, .. + } + | Self::Concurrent { + session_updates, .. + } => session_updates + .as_ref() + .map(|sender| sender.residency_version(worker)), + Self::Remote { .. } | Self::None => None, + } + } + + pub fn enqueue_session_match( + &self, + session_id: &str, + worker: WorkerWithDpRank, + matched_hash: ExternalSequenceBlockHash, + residency_version: u64, + ) -> Result<(), KvRouterError> { + match self { + Self::Single { + session_updates, .. + } + | Self::Concurrent { + session_updates, .. + } => match session_updates { + Some(sender) => sender.enqueue(SessionMutation::Matched { + worker, + session_id: session_id.to_owned(), + matched_hash, + residency_version, + }), + None => Ok(()), + }, + Self::Remote { .. } | Self::None => Ok(()), + } + } + + #[cfg(test)] + async fn flush_session_updates(&self) -> Result<(), KvRouterError> { + match self { + Self::Single { + session_updates, .. + } + | Self::Concurrent { + session_updates, .. + } => match session_updates { + Some(updates) => updates.flush().await, + None => Ok(()), + }, + 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, + ) -> 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 => {} + } + let session_updates = match self { + Self::Single { + session_updates, .. + } + | Self::Concurrent { + session_updates, .. + } => session_updates.as_ref(), + Self::Remote { .. } | Self::None => None, + }; + if let Some(session_updates) = session_updates { + session_updates.enqueue(SessionMutation::Cleared { + worker: WorkerWithDpRank::new(worker_id, dp_rank), + })?; + session_updates.flush().await?; } 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 +666,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 { .. } | Indexer::None => {} } } 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,26 +695,47 @@ 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 { .. } | Indexer::None => {} } } - /// 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. 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) - } - Indexer::Concurrent { primary, .. } => { - primary.find_matches(hashes).await.map_err(Into::into) + let Some(side) = self.approx() else { + 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( + merge_side_or_warn(Some(side), merged, HashInput::Owned(hashes)) + .await + .overlap_scores, + ) + } + + /// 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 @@ -190,63 +743,70 @@ impl Indexer { pub async fn find_tiered_matches( &self, sequence: Vec, - ) -> std::result::Result { - 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, - }) - } - 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, - }) - } - } + ) -> Result { + 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_kv_transfer_chain_retention`]); other shapes ignore it. + pub async fn find_tiered_matches_with_options( + &self, + sequence: Vec, + options: LowerTierQueryOptions, + ) -> Result { + let options = LowerTierQueryOptions { + retain_kv_transfer_chain: options.retain_kv_transfer_chain + && self.supports_kv_transfer_chain_retention(), + }; + self.find_matches_by_tier_with_options(sequence, options) + .await + } + + /// Borrowed variant of [`Self::find_tiered_matches_with_options`]. + pub async fn find_tiered_matches_ref_with_options( + &self, + sequence: &[LocalBlockHash], + options: LowerTierQueryOptions, + ) -> 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`]. - /// - /// 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) -> 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); @@ -254,6 +814,39 @@ 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, + ) -> 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) -> 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] @@ -261,8 +854,18 @@ impl TieredMatchProvider for Indexer { async fn find_tiered_matches( &self, sequence: &[LocalBlockHash], - ) -> std::result::Result { - self.find_tiered_matches(sequence.to_vec()).await + ) -> Result { + self.find_tiered_matches_ref_with_options(sequence, LowerTierQueryOptions::default()) + .await + } + + async fn find_tiered_matches_with_options( + &self, + sequence: &[LocalBlockHash], + options: LowerTierQueryOptions, + ) -> Result { + self.find_tiered_matches_ref_with_options(sequence, options) + .await } } @@ -448,6 +1051,7 @@ mod tests { if let Indexer::Single { primary, lower_tier, + .. } = &source { let _ = primary.flush().await; @@ -482,6 +1086,7 @@ mod tests { if let Indexer::Single { primary, lower_tier, + .. } = &replayed { let _ = primary.flush().await; @@ -530,7 +1135,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 { .. } | Indexer::None => { + unreachable!("test creates the single indexer") + } }; indexer @@ -565,12 +1172,366 @@ mod tests { assert!(healthy_tier.dump_events().await.unwrap().is_empty()); } - #[cfg(feature = "metrics")] - #[tokio::test] - async fn acknowledged_clear_records_primary_event_for_both_indexers() { - for num_threads in [1, 2] { - let metrics = Arc::new(KvIndexerMetrics::new_unregistered()); - let indexer = create_indexer_with_metrics(4, num_threads, metrics.clone()); + fn policy_indexer(num_threads: usize, policy: IndexerPolicy) -> Indexer { + create_indexer_with_policy( + 4, + num_threads, + Arc::new(KvIndexerMetrics::new_unregistered()), + &policy, + ) + } + + async fn flush_side(approx: &Option) { + match approx { + Some(SideIndexer::KvIndexer(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, + Indexer::None => {} + } + } + + impl Indexer { + 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()); + assert!(!indexer.records_routing_decisions()); + assert!(indexer.supports_kv_transfer_chain_retention()); + let worker = WorkerWithDpRank::new(7, 0); + indexer + .record_routing_decision_hashes( + 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_kv_transfer_chain_retention()); + let worker = WorkerWithDpRank::new(7, 0); + indexer + .record_routing_decision_hashes( + 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_hashes( + predicted, + RoutingDecisionHashes::from_local_hashes(vec![ + LocalBlockHash(11), + LocalBlockHash(12), + LocalBlockHash(13), + ]), + ) + .await + .unwrap(); + indexer + .record_routing_decision_hashes( + 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_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 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)) + .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 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()); + + // 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_hashes( + 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)); + } + + #[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() { + for num_threads in [1, 2] { + let metrics = Arc::new(KvIndexerMetrics::new_unregistered()); + let indexer = create_indexer_with_metrics(4, num_threads, metrics.clone()); indexer .apply_event_routed(RouterEvent::with_residency_domain( 7, @@ -611,25 +1572,371 @@ pub fn create_indexer_with_metrics( num_threads: usize, metrics: Arc, ) -> Indexer { + create_indexer_with_policy( + block_size, + num_threads, + metrics, + &IndexerPolicy::event_driven(), + ) +} + +pub fn create_indexer_with_policy( + 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), + CancellationToken::new(), + ) + }); + let retention = match &policy.primary { + PrimaryRetention::EventDriven => None, + PrimaryRetention::ApproximateTtl(ttl) => { + Some(ApproximateRetentionConfig::Ttl(PruneConfig { ttl: *ttl })) + } + }; + 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, + session_updates: None, } } 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, + session_updates: None, } } } + +#[cfg(test)] +mod session_tests { + use super::test_util::store_event; + use super::*; + use crate::SessionPrefixIndexer; + use crate::protocols::{KvCacheEvent, KvCacheRemoveData, StorageTier}; + + fn make_session_test_indexer() -> (Indexer, Arc) { + let session_prefix_index = Arc::new(SessionPrefixIndexer::new()); + let primary = KvIndexer::new( + CancellationToken::new(), + 4, + Arc::new(KvIndexerMetrics::new_unregistered()), + ); + let indexer = Indexer::Single { + primary: primary.clone(), + lower_tier: LowerTierIndexers::new(1, 4), + approx: None, + primary_records_routing_decisions: false, + session_updates: Some(SessionUpdateSender::for_legacy( + Arc::clone(&session_prefix_index), + primary, + )), + }; + (indexer, session_prefix_index) + } + + fn make_concurrent_session_test_indexer() -> (Indexer, Arc) { + let session_prefix_index = Arc::new(SessionPrefixIndexer::new()); + let primary = Arc::new(ThreadPoolIndexer::new( + ConcurrentRadixTreeCompressed::new(), + 2, + 4, + )); + let indexer = Indexer::Concurrent { + primary: Arc::clone(&primary), + lower_tier: LowerTierIndexers::new(2, 4), + approx: None, + primary_records_routing_decisions: false, + session_updates: Some(SessionUpdateSender::for_concurrent( + Arc::clone(&session_prefix_index), + primary, + )), + }; + (indexer, session_prefix_index) + } + + pub(crate) fn remove_event( + worker_id: u64, + dp_rank: u32, + event_id: u64, + block_hashes: Vec, + ) -> RouterEvent { + RouterEvent::new( + worker_id, + KvCacheEvent { + event_id, + data: KvCacheEventData::Removed(KvCacheRemoveData { block_hashes }), + dp_rank, + }, + ) + } + + pub(crate) fn clear_event(worker_id: u64, dp_rank: u32, event_id: u64) -> RouterEvent { + RouterEvent::new( + worker_id, + KvCacheEvent { + event_id, + data: KvCacheEventData::Cleared, + dp_rank, + }, + ) + } + + #[tokio::test] + async fn only_attributed_stored_events_update_session_lineage() { + let (indexer, session_prefix_index) = make_session_test_indexer(); + + indexer + .try_apply_event( + store_event(7, 0, 1, &[], &[41], StorageTier::Device).with_session_id("session-1"), + ) + .await + .unwrap(); + indexer + .try_apply_event(store_event(7, 0, 2, &[], &[51], StorageTier::HostPinned)) + .await + .unwrap(); + indexer.flush_session_updates().await.unwrap(); + + assert_eq!( + session_prefix_index + .get_session_block_lineage("session-1", WorkerWithDpRank::new(7, 0), None) + .unwrap(), + vec![vec![ExternalSequenceBlockHash(41)]] + ); + assert_eq!(session_prefix_index.session_count(), 1); + assert_eq!(session_prefix_index.node_count(), 1); + assert!( + session_prefix_index + .get_node_from_hash(ExternalSequenceBlockHash(51)) + .is_none() + ); + } + + #[tokio::test] + async fn physical_events_update_worker_local_session_frontiers_in_order() { + let (indexer, session_prefix_index) = make_session_test_indexer(); + let first_worker = WorkerWithDpRank::new(7, 0); + let second_worker = WorkerWithDpRank::new(8, 0); + let first_store = store_event(7, 0, 1, &[], &[41, 42, 43], StorageTier::Device) + .with_session_id("session-1"); + let block_hashes = match &first_store.event.data { + KvCacheEventData::Stored(stored) => stored + .blocks + .iter() + .map(|block| block.block_hash) + .collect::>(), + _ => unreachable!(), + }; + let second_store = store_event(8, 0, 1, &[], &[41, 42, 43], StorageTier::Device) + .with_session_id("session-1"); + + indexer.try_apply_event(first_store).await.unwrap(); + indexer.try_apply_event(second_store).await.unwrap(); + indexer + .try_apply_event(remove_event(7, 0, 2, block_hashes[1..].to_vec())) + .await + .unwrap(); + indexer.flush_session_updates().await.unwrap(); + + assert_eq!( + session_prefix_index + .get_session_block_lineage("session-1", first_worker, None) + .unwrap(), + vec![vec![block_hashes[0]]] + ); + assert_eq!( + session_prefix_index + .get_session_block_lineage("session-1", second_worker, None) + .unwrap(), + vec![block_hashes.clone()] + ); + + indexer.try_apply_event(clear_event(8, 0, 2)).await.unwrap(); + indexer.flush_session_updates().await.unwrap(); + assert!( + session_prefix_index + .get_session_block_lineage("session-1", second_worker, None) + .unwrap() + .is_empty() + ); + assert_eq!( + session_prefix_index.node_count(), + 3, + "physical removal preserves logical topology" + ); + } + + #[tokio::test] + async fn concurrent_physical_events_precede_session_updates() { + let (indexer, session_prefix_index) = make_concurrent_session_test_indexer(); + let worker = WorkerWithDpRank::new(7, 0); + let store = + store_event(7, 0, 1, &[], &[41, 42], StorageTier::Device).with_session_id("session-1"); + let block_hashes = match &store.event.data { + KvCacheEventData::Stored(stored) => stored + .blocks + .iter() + .map(|block| block.block_hash) + .collect::>(), + _ => unreachable!(), + }; + + indexer.try_apply_event(store).await.unwrap(); + indexer + .try_apply_event(remove_event(7, 0, 2, vec![block_hashes[1]])) + .await + .unwrap(); + indexer.flush_session_updates().await.unwrap(); + + assert_eq!( + session_prefix_index + .get_session_block_lineage("session-1", worker, None) + .unwrap(), + vec![vec![block_hashes[0]]] + ); + } + + #[tokio::test] + async fn route_match_is_ordered_after_pending_remove_and_store() { + let (indexer, session_prefix_index) = make_session_test_indexer(); + let worker = WorkerWithDpRank::new(7, 0); + let initial_store = + store_event(7, 0, 1, &[], &[41], StorageTier::Device).with_session_id("session-old"); + let block_hash = match &initial_store.event.data { + KvCacheEventData::Stored(stored) => stored.blocks[0].block_hash, + _ => unreachable!(), + }; + + indexer.try_apply_event(initial_store).await.unwrap(); + indexer.flush_session_updates().await.unwrap(); + + indexer + .try_apply_event(remove_event(7, 0, 2, vec![block_hash])) + .await + .unwrap(); + indexer + .try_apply_event( + store_event(7, 0, 3, &[], &[41], StorageTier::Device).with_session_id("session-a"), + ) + .await + .unwrap(); + let residency_version = indexer.session_residency_version(worker).unwrap(); + indexer + .enqueue_session_match("session-b", worker, block_hash, residency_version) + .unwrap(); + indexer.flush_session_updates().await.unwrap(); + + for session in ["session-a", "session-b"] { + assert_eq!( + session_prefix_index + .get_session_block_lineage(session, worker, None) + .unwrap(), + vec![vec![block_hash]] + ); + } + } + + #[tokio::test] + async fn stale_route_match_does_not_undo_removal() { + let (indexer, session_prefix_index) = make_session_test_indexer(); + let worker = WorkerWithDpRank::new(7, 0); + let store = + store_event(7, 0, 1, &[], &[41], StorageTier::Device).with_session_id("session-b"); + let block_hash = match &store.event.data { + KvCacheEventData::Stored(stored) => stored.blocks[0].block_hash, + _ => unreachable!(), + }; + + indexer.try_apply_event(store).await.unwrap(); + indexer.flush_session_updates().await.unwrap(); + indexer + .try_apply_event(remove_event(7, 0, 2, vec![block_hash])) + .await + .unwrap(); + // The physical removal is queued, but its ordered session mutation has + // not run yet. A match observed in this window must still be invalidated. + let stale_version = indexer.session_residency_version(worker).unwrap(); + indexer + .enqueue_session_match("session-b", worker, block_hash, stale_version) + .unwrap(); + indexer.flush_session_updates().await.unwrap(); + + assert!( + session_prefix_index + .get_session_block_lineage("session-b", worker, None) + .unwrap() + .is_empty() + ); + } + + #[tokio::test] + async fn rank_reset_waits_for_session_frontier_clear() { + let (indexer, session_prefix_index) = make_session_test_indexer(); + let reset_rank = WorkerWithDpRank::new(7, 0); + let retained_rank = WorkerWithDpRank::new(7, 1); + let block_hash = ExternalSequenceBlockHash(41); + + for rank in [reset_rank, retained_rank] { + indexer + .try_apply_event( + store_event( + rank.worker_id, + rank.dp_rank, + 1, + &[], + &[block_hash.0], + StorageTier::Device, + ) + .with_session_id("session-1"), + ) + .await + .unwrap(); + } + indexer.flush_session_updates().await.unwrap(); + + indexer + .reset_worker_dp_rank_and_wait(reset_rank.worker_id, reset_rank.dp_rank) + .await + .unwrap(); + + assert!( + session_prefix_index + .get_session_block_lineage("session-1", reset_rank, None) + .unwrap() + .is_empty() + ); + assert_eq!( + session_prefix_index + .get_session_block_lineage("session-1", retained_rank, None) + .unwrap(), + vec![vec![block_hash]] + ); + } +} 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..aea7cdf380d0 --- /dev/null +++ b/lib/kv-router/src/services/indexer/lookup.rs @@ -0,0 +1,1215 @@ +// 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(sequence) => sequence, + Self::Owned(sequence) => sequence, + } + } + + 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) struct LookupPipeline<'a> { + primary: PrimaryLookup<'a>, + lower_tier: Option<&'a LowerTierIndexers>, + side: Option<&'a SideIndexer>, +} + +#[derive(Clone, Copy)] +enum PrimaryLookup<'a> { + KvIndexer(&'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::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, + }, + } + } + + /// 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 + } + + pub async fn find_primary_match_details_ref( + &self, + sequence: &[LocalBlockHash], + ) -> Result { + self.lookup_pipeline() + .find_primary_match_details(HashInput::Borrowed(sequence)) + .await + } + + /// 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, + ) -> 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_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::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_with_options` + // 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_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 { + 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. +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, LowerTierQueryOptions}; + 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, + session_updates: None, + } + } + + 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, + session_updates: None, + } + } + + 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, + session_updates: None, + } + } + + #[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, + session_updates: None, + }; + 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, + session_updates: None, + }; + + 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_with_options(&query, LowerTierQueryOptions::default()) + .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_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() { + 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..74e56bc6b3c7 100644 --- a/lib/kv-router/src/services/indexer/mod.rs +++ b/lib/kv-router/src/services/indexer/mod.rs @@ -32,10 +32,14 @@ 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; +#[doc(hidden)] +pub mod session_updates; use std::path::PathBuf; use std::sync::Arc; diff --git a/lib/llm/src/kv_router/indexer/recording.rs b/lib/kv-router/src/services/indexer/recording.rs similarity index 80% rename from lib/llm/src/kv_router/indexer/recording.rs rename to lib/kv-router/src/services/indexer/recording.rs index 596938f71681..95e484d0d453 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-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,13 @@ impl Indexer { }) } - pub(crate) fn records_routing_decisions(&self) -> bool { + 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 +140,7 @@ impl Indexer { ); RouteRecordingTarget::SideOverlay(side) } - Self::KvIndexer { + Self::Single { primary, primary_records_routing_decisions: true, .. @@ -152,13 +155,13 @@ 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( + pub async fn record_hashed_routing_decision( &self, worker: WorkerWithDpRank, local_hashes: Vec, @@ -175,7 +178,7 @@ impl Indexer { .await } - pub(crate) async fn record_routing_decision_hashes( + pub async fn record_routing_decision_hashes( &self, worker: WorkerWithDpRank, hashes: RoutingDecisionHashes, @@ -185,7 +188,7 @@ impl Indexer { .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 +244,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/kv-router/src/services/indexer/registry.rs b/lib/kv-router/src/services/indexer/registry.rs index 15fa98476cbd..5fec65edc7b4 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 { @@ -317,6 +317,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, @@ -343,6 +346,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, @@ -365,6 +369,7 @@ impl WorkerRegistry { watermarks: DashMap::new(), num_threads, indexer_metrics, + indexer_policy: Mutex::new(IndexerPolicy::event_driven()), ready_tx, ready_rx, root_cancel_token, @@ -372,6 +377,11 @@ impl WorkerRegistry { } } + /// Set the shape used for indexers created from now on. + pub fn set_indexer_policy(&self, policy: IndexerPolicy) { + *self.indexer_policy.lock() = policy; + } + #[cfg(feature = "standalone-selection")] pub(crate) fn with_retained_indexers(mut self) -> Self { // Selection entries and their schedulers retain these indexers across worker updates. @@ -442,6 +452,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, @@ -450,10 +461,11 @@ impl WorkerRegistry { "Creating new indexer" ); IndexerEntry { - indexer: create_indexer_with_metrics( + indexer: create_indexer_with_policy( block_size, self.num_threads, self.indexer_metrics.clone(), + &indexer_policy, ), block_size, } @@ -682,6 +694,32 @@ 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 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)) + } + + /// 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, + ) { + // 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; + } + } + pub fn list(&self) -> Vec { self.list_filtered(None, None) } @@ -756,6 +794,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, @@ -764,10 +803,11 @@ impl WorkerRegistry { "Creating indexer from recovery dump" ); IndexerEntry { - indexer: create_indexer_with_metrics( + indexer: create_indexer_with_policy( block_size, self.num_threads, self.indexer_metrics.clone(), + &indexer_policy, ), block_size, } @@ -797,6 +837,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/indexer/server.rs b/lib/kv-router/src/services/indexer/server.rs index 185efae4bbff..9f1f9bfa3877 100644 --- a/lib/kv-router/src/services/indexer/server.rs +++ b/lib/kv-router/src/services/indexer/server.rs @@ -562,7 +562,10 @@ 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_by_hash", + post(query_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/llm/src/kv_router/indexer/session_updates.rs b/lib/kv-router/src/services/indexer/session_updates.rs similarity index 98% rename from lib/llm/src/kv_router/indexer/session_updates.rs rename to lib/kv-router/src/services/indexer/session_updates.rs index 509cd9723c9d..6342a42f0916 100644 --- a/lib/llm/src/kv_router/indexer/session_updates.rs +++ b/lib/kv-router/src/services/indexer/session_updates.rs @@ -9,12 +9,12 @@ use std::{ }, }; -use dashmap::{DashMap, mapref::entry::Entry}; -use dynamo_kv_router::{ +use crate::{ ConcurrentRadixTreeCompressed, SessionPrefixIndexer, indexer::{KvIndexer, KvRouterError, ThreadPoolIndexer}, protocols::{ExternalSequenceBlockHash, KvCacheEventData, RouterEvent, WorkerWithDpRank}, }; +use dashmap::{DashMap, mapref::entry::Entry}; use tokio::sync::{mpsc, oneshot}; #[derive(Clone)] @@ -198,11 +198,11 @@ impl PrimaryBarrier { } impl SessionUpdateSender { - pub(super) fn for_legacy(index: Arc, primary: KvIndexer) -> Self { + pub fn for_legacy(index: Arc, primary: KvIndexer) -> Self { Self::spawn(index, PrimaryBarrier::Legacy(primary)) } - pub(super) fn for_concurrent( + pub fn for_concurrent( index: Arc, primary: Arc>, ) -> Self { diff --git a/lib/kv-router/src/services/selection/README.md b/lib/kv-router/src/services/selection/README.md index e89c98e402c4..485b527b2dfb 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 @@ -12,6 +17,19 @@ 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`. +- 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`. 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 + 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 new file mode 100644 index 000000000000..f4d822fdb619 --- /dev/null +++ b/lib/kv-router/src/services/selection/affinity.rs @@ -0,0 +1,1198 @@ +// 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, + }, +} + +/// 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, + next_revision: AtomicU64, + next_sequence: AtomicU64, + writer_id: AtomicU64, + cancel: CancellationToken, + replica: OnceLock>, + #[cfg(any(test, feature = "testing"))] + reaper_started: Arc, + #[cfg(any(test, feature = "testing"))] + 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 }) + } +} + +/// 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 { + target: AffinityTarget, + lease: AffinityLease, + }, +} + +impl Hold { + pub fn target(&self) -> Option { + match self { + Self::Initialize(_) => None, + Self::Bound { target, .. } => Some(*target), + } + } + + pub fn invalidate(self) { + if let Self::Bound { mut lease, .. } = self { + lease.invalidate(); + } + } +} + +/// One step of acquiring a session. +pub enum AcquireStep { + Held(Hold), + /// Another request is initializing this session; await, then retry. + Wait(Pin>), +} + +impl SessionAffinity { + pub(crate) fn validate_ttl(ttl: Duration) -> Result<(), AffinityError> { + 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" + ))); + } + Ok(()) + } + + 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( + 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(any(test, feature = "testing"))] + reaper_started: Arc::new(Notify::new()), + #[cfg(any(test, feature = "testing"))] + waiter_observed: Arc::new(Notify::new()), + }); + Self::spawn_reaper(&inner); + tracing::debug!( + ttl_secs = config.ttl.as_secs(), + mode = ?config.mode, + max_entries = config.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(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! { + _ = 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 { + // 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 { + 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( + &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_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::Held(Hold::Initialize( + AffinityInitialization { + coordinator: 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, .. } => { + #[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. + 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_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::Held(Hold::Initialize( + AffinityInitialization { + coordinator: 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 = %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::Held(Hold::Bound { + target: *target, + lease: AffinityLease { + coordinator: 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, + AcquireStep::Held(hold) => return Ok(hold), + } + } + } + + /// 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, notify } = entry.value() else { + return None; + }; + 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, + version, + active: true, + }; + lease.publish(target); + Some(lease) + } + + /// 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) => initialization.commit(dispatched), + 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) + } + } + } + + /// 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 = %session_id, + worker_id = target.worker_id, + dp_rank = ?target.dp_rank, + "Session affinity hit: reusing pinned worker" + ); + Ok(Some(*target)) + } + + pub fn observe_replica_sequence(&self, sequence: u64) { + self.inner.observe_replica_sequence(sequence); + } + + pub fn apply_replica_update( + &self, + session_id: String, + target: AffinityTarget, + version: AffinityVersion, + ) -> ReplicaApplyOutcome { + self.inner.apply_replica_update(session_id, target, 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()) + }) + } +} + +/// 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) + } + + pub fn cancellation_token(&self) -> CancellationToken { + self.inner.cancel.clone() + } + + pub async fn wait_for_reaper(&self) { + self.inner.reaper_started.notified().await; + } + + pub async fn wait_for_initializing_waiter(&self) { + self.inner.waiter_observed.notified().await; + } + + 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() + } + + /// 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 { .. } => 0, + AffinityEntry::Bound { active_leases, .. } => *active_leases, + }) + } +} + +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, unless another request already bound it. +pub struct AffinityInitialization { + coordinator: Weak, + session_id: String, + revision: u64, + notify: Arc, + requested_target: Option, + active: bool, +} + +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.coordinator.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 let AffinityEntry::Bound { + target: bound, + revision, + version, + .. + } = 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(), + )); + }; + 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(); + let lease = AffinityLease { + coordinator: Arc::downgrade(&inner), + session_id: self.session_id.clone(), + revision: self.revision, + version, + active: true, + }; + lease.publish(target); + Ok(lease) + } +} + +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_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); + } + self.notify.notify_waiters(); + } +} + +/// A live use of a bound session. Dropping it releases the use and refreshes +/// the idle deadline. +pub struct AffinityLease { + coordinator: Weak, + session_id: String, + revision: u64, + version: AffinityVersion, + active: bool, +} + +impl AffinityLease { + pub fn session_id(&self) -> &str { + &self.session_id + } + + 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. + 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 + } + + /// 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.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); + } + + /// 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.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, + 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; + 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(()) +} + +#[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_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", target) + .expect("join an initializing session"); + 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(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"), + 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(); + 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", target) + .expect("join an initializing session"); + 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", target).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_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", target) + .expect("join an initializing session"); + assert_eq!(table.entry_count(), 1); + + drop(init); + 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.lease_count("s"), Some(0)); + tokio::time::advance(TTL + Duration::from_secs(1)).await; + assert_eq!(table.query_target("s", None).expect("query"), None); + + // Once the completed booking's TTL expires, another target can bind. + 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)); + } + + #[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/catalog.rs b/lib/kv-router/src/services/selection/catalog.rs index 11fc0010d72a..11ddaab292e0 100644 --- a/lib/kv-router/src/services/selection/catalog.rs +++ b/lib/kv-router/src/services/selection/catalog.rs @@ -6,12 +6,9 @@ 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::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 { @@ -92,13 +64,32 @@ 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)) + } + + /// `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 remove(&self, worker_id: WorkerId) -> Option { + self.workers.write().remove(&worker_id) } + /// The schedulable workers in `key`'s partition, keyed by worker id. pub(super) fn scheduler_configs_for_key( &self, key: &RoutingPartitionId, @@ -106,11 +97,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() @@ -127,6 +114,22 @@ impl WorkerCatalog { .count() } + pub(super) fn is_schedulable( + &self, + target: WorkerAffinityTarget, + key: &RoutingPartitionId, + ) -> bool { + self.workers + .read() + .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( &self, worker_id: WorkerId, @@ -134,10 +137,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() @@ -150,13 +150,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/hint.rs b/lib/kv-router/src/services/selection/core/hint.rs new file mode 100644 index 000000000000..b3c3685e7114 --- /dev/null +++ b/lib/kv-router/src/services/selection/core/hint.rs @@ -0,0 +1,80 @@ +// 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. +/// 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, + 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 c32027936eec..bae4224752ff 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -1,61 +1,124 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use std::collections::{HashMap, HashSet}; +//! 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::Instant; +#[cfg(test)] +use std::sync::atomic::Ordering; +use std::time::{Duration, Instant}; use dynamo_tokens::SequenceHash; 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::TieredMatchDetails; +use crate::indexer::{ + LowerTierQueryOptions, RoutingDecisionHashes, SharedKvCache, TieredMatchDetails, +}; +use crate::kv_hints::{ + KvHint, KvHintAction, KvSourceLocationsPayload, KvTransferCandidateSource, KvTransferCandidates, +}; use crate::protocols::{ - ActiveSequenceEvent, LocalBlockHash, PrefillLoadHint, RoutingConstraints, WorkerAffinityTarget, - 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::{ - 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, + ActiveSequencesMultiWorker, LifecycleMutationOutcome, ReplicaRequestLeaseObserver, + ReplicaWorkerPolicy, SequenceRequest, SequenceTrackerOptions, active_request_expiry_duration, }; use crate::services::common::replica_sync::{ - ReplicaSyncConfig, ScopedReplicaEvent, ScopedSequencePublisher, setup_scoped_replica_sync, + HostReplicaSyncFactory, ReplicaSyncConfig, SchedulerLoadSink, 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; 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, + SelectionRun, SessionBinding, +}; + +use super::affinity::{ + AcquireStep, AffinityError, AffinityLease, Hold, SessionAffinity, SessionAffinityConfig, +}; use super::catalog::WorkerCatalog; use super::error::SelectionError; -use super::input::{PromptRequest, TrackingHashInput}; +use super::ingress::{KvEventIngress, ZmqDirectIngress}; +use super::input::{PromptView, TrackingHashInput}; 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; +use crate::services::common::replica_sync::AffinityBindingEvent; -type SelectionScheduler = LocalScheduler< +pub type SelectionScheduler = LocalScheduler< ScopedSequencePublisher, SelectionWorkerConfig, WorkerSelectionPolicy, 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 scheduler(&self) -> &SelectionScheduler { + &self.0.scheduler + } + + pub fn indexer(&self) -> &Indexer { + &self.0.indexer + } + + /// Return this partition's affinity table, initialized with `config`. + pub fn session_affinity( + &self, + config: SessionAffinityConfig, + ) -> Result { + self.0.session_affinity(config).cloned() + } +} + struct SelectionEntry { key: RoutingPartitionId, block_size: u32, @@ -64,42 +127,122 @@ struct SelectionEntry { workers_tx: watch::Sender>, scheduler: SelectionScheduler, replica_tx: Option>, + affinity: OnceCell, + replica_config: Option, } -struct PreparedSelectionInputs { - block_hashes: Vec, - sequence_hashes: Vec, - isl_tokens: usize, - overlap: OverlapSignals, +impl SelectionEntry { + fn session_affinity( + &self, + config: SessionAffinityConfig, + ) -> Result<&SessionAffinity, SelectionError> { + let table = self + .affinity + .get_or_try_init(|| -> Result<_, SelectionError> { + 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) + { + table.enable_replication(config.process_id(), sink); + } + Ok(table) + })?; + if table.ttl() != config.ttl || table.mode() != config.mode { + return Err(SelectionError::Conflict(format!( + "session affinity config mismatch for {}: existing=({:?}, {:?}) requested=({:?}, {:?})", + self.key, + table.ttl(), + table.mode(), + config.ttl, + config.mode + ))); + } + Ok(table) + } } -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_id: Option, - affinity_target: Option, - pinned_worker: Option, - allowed_worker_ids: Option>, - routing_constraints: RoutingConstraints, +/// 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. +#[derive(Clone, Default)] +pub struct SelectionHost { + pub load: HostLoad, + pub cache: HostCache, + pub eligibility: HostEligibility, + pub telemetry: HostTelemetry, + pub replication: HostReplication, } -/// 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, +/// 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, +} + +/// 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>, + pub index: KvIndexSource, +} + +/// Where a partition's KV index comes from and who feeds it. +#[derive(Clone)] +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. + Owned(Arc), +} + +impl Default for KvIndexSource { + fn default() -> Self { + Self::Owned(Arc::new(ZmqDirectIngress)) + } +} + +/// 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>, +} + +/// 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) 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)] +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)] @@ -111,14 +254,30 @@ 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>>>>; + pub struct SelectionCore { catalog: WorkerCatalog, - entries: RwLock>>>>, + /// 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. + reservation_sweep_started: OnceCell<()>, + /// Whether this core subscribes to worker KV events itself. False when + /// 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, worker_selection_policy_factory: Option, + host: SelectionHost, worker_type: WorkerType, cancel_token: CancellationToken, replica_config: Option, @@ -126,6 +285,24 @@ pub struct SelectionCore { /// `create_reservation` can replay them without re-sending the prompt. 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, + #[cfg(test)] + pub(super) after_affinity_invalidation: Option>, +} + +fn affinity_error(error: AffinityError) -> SelectionError { + match error { + AffinityError::InvalidArgument(message) => SelectionError::BadRequest(message), + AffinityError::ResourceExhausted(message) => SelectionError::NotReady(message), + AffinityError::Dropped => SelectionError::Internal(error.to_string()), + } } impl SelectionCore { @@ -155,19 +332,27 @@ 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, cancel_token, None, None, + SelectionHost::default(), WorkerType::Aggregated, true, cache_config, tracking_hash, + indexer_policy, + None, )) } + 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, @@ -175,30 +360,47 @@ impl SelectionCore { cancel_token: CancellationToken, replica_config: Option, worker_selection_policy_factory: Option, + host: SelectionHost, worker_type: WorkerType, signal_indexer_ready: bool, cache_config: SelectionCacheConfig, tracking_hash: Arc, + indexer_policy: IndexerPolicy, + session_affinity: Option, ) -> Self { let cancel_token = cancel_token.child_token(); let indexer_registry = Arc::new( 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_registry.set_indexer_policy(indexer_policy); if signal_indexer_ready { indexer_registry.signal_ready(); } Self { catalog: WorkerCatalog::default(), - entries: RwLock::new(HashMap::new()), + 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(), + listens_for_kv_events, indexer_registry, kv_router_config, worker_selection_policy_factory, + host, worker_type, cancel_token, replica_config, 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), + #[cfg(test)] + after_affinity_invalidation: None, } } @@ -277,935 +479,68 @@ impl SelectionCore { } } - pub async fn upsert_worker( - &self, - req: WorkerRequest, - ) -> Result { - self.ensure_running()?; - let (previous, record) = self.catalog.upsert(req); - self.reconcile_worker(record.worker_id, previous).await - } - - pub async fn patch_worker( - &self, - worker_id: WorkerId, - patch: WorkerPatchRequest, - ) -> Result { - self.ensure_running()?; - let (previous, record) = self.catalog.patch(worker_id, patch)?; - self.reconcile_worker(record.worker_id, Some(previous)) - .await - } - - pub async fn delete_worker( - &self, - worker_id: WorkerId, - ) -> Result { - 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 { + tracing::trace!( + key = %event.partition, + "Dropping session affinity replica update for unknown selector entry" + ); + 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, - } - } - - 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" - ))); + let Some(table) = entry.affinity.get() else { + tracing::trace!( + key = %event.partition, + "Dropping session affinity replica update: no affinity table" + ); + return; }; - - if previous + if self + .replica_config .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; - } - - let queueing_enabled = self - .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); - 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); - } - - 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 - && let Err(error) = self.register_indexer_listeners(&record).await + .is_some_and(|config| config.process_id() == event.writer_id) { - self.cleanup_indexer_registration(&record).await; - return self.mark_incomplete_after_reconcile_error(worker_id, record.key(), error); + return; } - - 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) - } - - fn mark_incomplete_after_reconcile_error( - &self, - worker_id: WorkerId, - key: RoutingPartitionId, - error: SelectionError, - ) -> Result { - let updated = self + table.observe_replica_sequence(event.sequence); + if 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) - } - - 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()))?; - let is_eagle = record.is_eagle.unwrap_or(false); - let key = record.key(); - - 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 scoped_replica_sync = - 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, - block_size as usize, - HashMap::new(), - scoped_replica_sync.enabled, - scoped_replica_sync.process_id, - worker_label, - ReplicaWorkerPolicy::RequireRegistered, - )); - 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(), - ); - - 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 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, - None, - Some(overlap_refresh), - None, - // Standalone selection has no router Client snapshot. - None, - 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, - })) - })? - .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 - ))); - } - 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()))?; - } - Ok(()) - } - - async fn cleanup_indexer_registration(&self, record: &WorkerCatalogRecord) { - if self.kv_router_config.use_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" - ); - } + .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 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; - } - } - - 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); - entry.workers_tx.send(workers).map_err(|_| { - SelectionError::Internal(format!("scheduler worker watch closed for {key}")) - }) + 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> { - 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) } - pub async fn select(&self, req: SelectRequest) -> Result { - self.select_with_policy_class(req, None).await - } - - pub async fn select_with_policy_class( - &self, - req: SelectRequest, - policy_class: Option, - ) -> Result { - self.schedule_selection( - SelectionOperation { - key: RoutingPartitionId::new(req.model_name, req.routing_group), - selection_id: req.selection_id, - prompt: req.prompt, - 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_id: req.session_id, - affinity_target: req.affinity_target, - pinned_worker: req.pinned_worker, - allowed_worker_ids: req.allowed_worker_ids, - routing_constraints: req.routing_constraints, - }, - false, - ) - .await - } - - 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, - req: SelectAndReserveRequest, - policy_class: Option, - ) -> Result { - let selection_id = req - .selection_id - .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); - self.schedule_selection( - SelectionOperation { - key: RoutingPartitionId::new(req.model_name, req.routing_group), - selection_id: Some(selection_id), - prompt: req.prompt, - 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_id: req.session_id, - affinity_target: req.affinity_target, - pinned_worker: req.pinned_worker, - allowed_worker_ids: req.allowed_worker_ids, - routing_constraints: req.routing_constraints, - }, - true, - ) - .await - } - - async fn schedule_selection( - &self, - operation: SelectionOperation, - book: bool, - ) -> Result { - let SelectionOperation { - key, - selection_id, - prompt, - router_config_override, - expected_output_tokens, - priority_jump, - strict_priority, - policy_class, - session_id, - affinity_target, - pinned_worker, - allowed_worker_ids, - routing_constraints, - } = operation; - self.ensure_running()?; - - let entry = self.ready_entry(&key)?; - let PreparedSelectionInputs { - block_hashes, - sequence_hashes, - isl_tokens, - overlap, - } = self - .prepare_selection_inputs( - &entry, - &prompt, - self.kv_router_config - .assume_kv_reuse(router_config_override.as_ref()), - ) - .await?; - let mode = if book { - ScheduleMode::Tracked { - request_id: selection_id.clone().ok_or_else(|| { - SelectionError::Internal( - "booked selection did not include a selection ID".to_string(), - ) - })?, - } - } else { - ScheduleMode::QueryOnly { - request_id: selection_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, - sequence_hashes.clone(), - prompt.lora_name.clone(), - track_prefill_tokens, - ) - }); - 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); - 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_config_override, - lora_name: prompt.lora_name, - priority_jump, - strict_priority, - policy_class, - session_context: session_id - .map(|session_id| SessionContext::new(session_id, None, None, None)), - expected_output_tokens, - affinity_target, - pinned_worker, - allowed_worker_ids, - routing_constraints, - shared_cache_hits: None, - }; - let response = tokio::select! { - biased; - _ = self.cancel_token.cancelled() => { - return Err(SelectionError::Scheduler(KvSchedulerError::SubscriberShutdown)); - } - result = entry.scheduler.schedule_request(schedule_request) => result?, - }; - 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 overlap = MooncakeOverlapSummary::from_selected_worker_tiers( - &response.selected_worker_tiers, - entry.block_size, - ); - - let effective_prefill = effective_prefill_tokens(isl_tokens, response.cached_tokens); - - if let Some((cache_id, sequence_hashes, lora_name, track_prefill_tokens)) = 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, - }, - 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, - endpoint, - block_size: entry.block_size, - overlap, - effective_prefill_tokens: effective_prefill, - }) - } - - 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 = self.book_cached_selection(pending, &req).await?; - 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(), - }, - ) - .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.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) - }); - - 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, - }, - ) - .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, - } = booking; - - // Strict booking: never lazily recreate a worker/rank removed since the - // reservation was resolved. - entry - .scheduler - .add_request_if_registered(SequenceRequest { - request_id: selection_id.clone(), - token_sequence: Some(sequence_hashes), - track_prefill_tokens, - expected_output_tokens, - prefill_load_hint, - worker, - lora_name, - }) - .await?; - - 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, - }) - } - - pub async fn prefill_complete(&self, selection_id: &str) -> Result<(), SelectionError> { - let entries = self.initialized_entries(); - for entry in entries { - match entry.scheduler.mark_prefill_completed(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" - ))) - } - - 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()), - } - } - Err(SelectionError::NotFound(format!( - "reservation {selection_id} not found" - ))) - } - - 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 entries = self.initialized_entries(); - for entry in entries { - match entry - .scheduler - .add_output_block(selection_id, decay_fraction) - { - Ok(()) => return Ok(()), - Err(SequenceError::RequestNotFound { .. }) => continue, - Err(error) => return Err(error.into()), - } - } - Err(SelectionError::NotFound(format!( - "reservation {selection_id} not found" - ))) - } - - 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, - self.kv_router_config - .assume_kv_reuse(req.router_config_override.as_ref()), - ) - .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 - .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, - ), - ) - } - - async fn prepare_selection_inputs( - &self, - entry: &SelectionEntry, - prompt: &PromptRequest, - assume_kv_reuse: bool, - ) -> Result { - let normalized = prompt.normalize_for_selection( - entry.is_eagle, - TrackingHashInput { - context: &self.tracking_hash, - scope: tracking_scope(entry), - assume_kv_reuse, - }, - )?; - let tiered = if normalized.block_hashes.is_empty() { - TieredMatchDetails::default() + /// 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 { - entry - .indexer - .find_tiered_matches(normalized.block_hashes.clone()) - .await - .map_err(|error| SelectionError::Internal(error.to_string()))? - }; - let overlap = - OverlapAnalysis::new(&self.kv_router_config, entry.block_size, &tiered).signals(); - drop(tiered); - Ok(PreparedSelectionInputs { - block_hashes: normalized.block_hashes, - sequence_hashes: normalized.sequence_hashes, - isl_tokens: normalized.isl_tokens, - overlap, - }) - } - - 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)); - } + SelectionError::NotReady(format!("no schedulable workers for {key}")) } - workers } } @@ -1221,582 +556,3 @@ impl Drop for SelectionCore { self.shutdown(); } } - -#[cfg(test)] -mod tests { - use super::*; - use crate::protocols::StorageTier; - use crate::services::indexer::backend::test_util::store_event; - 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, - } - } - - 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, - affinity_target: None, - pinned_worker: None, - allowed_worker_ids: None, - routing_constraints: RoutingConstraints::default(), - } - } - - 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, - affinity_target: None, - pinned_worker: None, - allowed_worker_ids: None, - routing_constraints: RoutingConstraints::default(), - } - } - - 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; - } - tokio::task::yield_now().await; - } - }) - .await - .expect("selection did not queue"); - } - - 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, - worker_type, - true, - SelectionCacheConfig::default(), - tracking_hash, - ); - - 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}" - ); - } - } - - #[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 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!(core.cancel_token.is_cancelled()); - assert!(!parent.is_cancelled()); - assert_eq!(core.indexer_registry.listener_cancelled(1, 0), Some(true)); - } - - #[rstest::rstest] - #[case(1)] - #[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 key = RoutingPartitionId::new("model", "default"); - let request = || { - let mut request = select_request(); - request.prompt.token_ids = None; - request.prompt.block_hashes = Some(vec![11]); - request.prompt.sequence_hashes = Some(vec![101]); - request.prompt.isl_tokens = Some(4); - request - }; - - // Exercise an in-place update, then removing the last worker and adding a new one. - for worker_id in [1, 1, 2] { - if worker_id == 2 { - core.delete_worker(1).await.expect("delete last worker"); - } - core.upsert_worker(worker_with_kv_events(worker_id)) - .await - .expect("worker upsert"); - assert_eq!(core.select(request()).await.unwrap().overlap.gpu, 0); - - // Write through the registry used by listeners, not the selector's saved reference. - let indexer = core - .indexer_registry - .get_indexer(&key) - .unwrap() - .indexer - .clone(); - indexer - .apply_event_routed(store_event( - worker_id, - 0, - 1, - &[], - &[11], - StorageTier::Device, - )) - .await - .unwrap(); - indexer.dump_events().await.expect("flush indexer"); - let selected = core.select(request()).await.expect("select cached worker"); - assert_eq!(selected.worker_id, worker_id); - assert_eq!(selected.overlap.gpu, 4); - } - core.shutdown(); - } - - #[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 lifecycle_operations_find_reservation_in_later_entry() { - 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(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, - 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)])); - } -} 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..2817ee94f1ef --- /dev/null +++ b/lib/kv-router/src/services/selection/core/operation.rs @@ -0,0 +1,147 @@ +// 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 std::time::Duration; + +use dynamo_tokens::SequenceHash; + +use crate::identity::RoutingPartitionId; +use crate::kv_hints::KvHint; +use crate::protocols::{ + LocalBlockHash, RoutingConstraints, SharedCacheHits, WorkerAffinityTarget, WorkerId, + WorkerWithDpRank, +}; +use crate::scheduling::config::RouterConfigOverride; +use crate::scheduling::queue::BookingHandle; +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` +/// 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, + /// 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. + Query { request_id: Option }, + /// 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 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 }, +} + +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 } + | Self::Lease { + request_id: selection_id, + } => Some(selection_id), + } + } + + pub fn is_booking(&self) -> bool { + matches!(self, Self::Book { .. } | Self::Lease { .. }) + } +} + +/// 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, + }, +} + +/// 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)] +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 for wire responses; omitted for `Lease`. + pub total_kv_blocks: Option, + /// 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, + /// 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, + pub routing_hashes: Option>, + pub shared_cache_hits: 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/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..6074758f3bf4 --- /dev/null +++ b/lib/kv-router/src/services/selection/core/reservations.rs @@ -0,0 +1,596 @@ +// 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::*; +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. +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) claim_id: 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) claim_id: u64, + pub(super) armed: bool, +} + +impl ReservationClaim<'_> { + 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; + Ok(()) + } +} + +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.claim_id == Some(self.claim_id)) + { + 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 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 + // 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()), + 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); + } + } + + 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(); + 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 { + 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 + .remove(&key, &req.selection_id, generation); + 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.view().routing_tokens_and_mm_infos().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 handle frees the booking if this future + // is dropped before `install`. + let booking = 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(booking, 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" + ))); + } + 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, + }) + } + + /// 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_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)), + } + } + + 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(); + let mut removed = Vec::new(); + index.retain(|_, reservation| { + let Some(booking) = &reservation.booking else { + return true; + }; + let live = entries + .get(&reservation.partition) + .and_then(|cell| cell.get()) + .is_some_and(|entry| entry.scheduler.has_booking(booking)); + if !live { + removed.push(reservation._affinity_lease.take()); + } + live + }); + 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 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..2df9803a2179 --- /dev/null +++ b/lib/kv-router/src/services/selection/core/run.rs @@ -0,0 +1,824 @@ +// 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. +//! +//! # 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 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 +//! (`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`. A commit that finds another request initializing the +//! 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 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 +//! 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::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 +/// 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. +const KV_HINT_FETCH_ACTION_ID: &str = "a1"; + +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: _, + booking: _, + } = 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)?, + _ => 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). + // 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 { + 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()); + // 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, + 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, + }; + // `booking` 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_booking(schedule_request) + .instrument(tracing::info_span!("kv_router.schedule")) + .await + .map(|(admitted, booking)| (admitted.response, None, booking)) + } + } => result, + }; + let (response, advisory_load, booking) = match scheduled { + Ok(scheduled) => scheduled, + Err(KvSchedulerError::QueueRejected(rejection)) => { + return Ok(SelectionOutcome::QueueRejected { rejection }); + } + Err(error) => return Err(error.into()), + }; + 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 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(KV_HINT_FETCH_ACTION_ID, 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 booking = if let Some(claim) = claim { + let Some(booking) = booking else { + return Err(SelectionError::Internal( + "booked selection has no booking handle".to_string(), + )); + }; + // 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))) => { + 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(booking, affinity_lease)?; + None + } else { + booking + }; + + 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, + booking, + })) + } + + /// 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( + &self, + table: &SessionAffinity, + session_id: &str, + key: &RoutingPartitionId, + ) -> Result, SelectionError> { + let mut invalidations = 0; + 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, key) => + { + tracing::debug!( + session_id, + worker_id = target.worker_id, + "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(_)) => { + 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 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, + 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, key)); + if !departed { + return Err(affinity_error(error)); + } + // `commit` invalidated the stale binding. Another request may + // already be initializing the replacement: join it rather + // 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, dispatched) { + return Ok(Some(lease)); + } + } + Err(AffinityError::ResourceExhausted(_)) => return Ok(None), + Err(error) => return Err(affinity_error(error)), + } + } + Ok(None) + } + } + } + + /// 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_hashes( + 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_ref_with_options( + &normalized.block_hashes, + 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..408944a166bb --- /dev/null +++ b/lib/kv-router/src/services/selection/core/tests.rs @@ -0,0 +1,3278 @@ +// 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, AtomicUsize, 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 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. +fn core_with( + config: crate::config::KvRouterConfig, + host: SelectionHost, + policy_factory: Option, + worker_type: WorkerType, + affinity: Option, +) -> SelectionCore { + 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, + 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, + ) +} + +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, + 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" + )); +} + +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(); + let core = local_core_with(test_config(false), 1, parent.clone()); + + 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 core = core_with( + test_config(false), + SelectionHost::default(), + None, + worker_type, + None, + ); + + core.upsert_worker(worker(1)).await.expect("worker upsert"); + let entry = core.entry(&default_key()).expect("selection entry"); + assert_eq!( + entry.scheduler.worker_type(), + expected_label, + "{worker_type}" + ); + } +} + +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") +} + +#[rstest::rstest] +#[case::text(false)] +#[case::multimodal_only(true)] +#[tokio::test] +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"); + 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(replay_reservation("cached")) + .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 = || { + 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), + prompt: prompt_c(), + ..replay_reservation("explicit") + }) + .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); + assert_eq!(credited.overlap.longest_matched, 4); +} + +#[tokio::test] +async fn unreachable_remote_indexer_is_reported_not_ready() { + use crate::indexer::{KvRouterError, TieredMatchDetails}; + use crate::services::indexer::backend::RemotePrimary; + + 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()) + } + + 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 core = core_with( + test_config(true), + SelectionHost { + cache: HostCache { + index: KvIndexSource::Owned(Arc::new(OfflineIngress)), + shared: None, + }, + ..SelectionHost::default() + }, + None, + WorkerType::Aggregated, + None, + ); + 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 +/// 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 = local_core(test_config(true)); + 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(&default_key()).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()); +} + +/// 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 = || super::hint::hint_capable_partition(&entry.workers_tx.borrow()); + + 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)); + 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. 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); + } +} + +#[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 = local_core(test_config(false)); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + let entry = core.entry(&default_key()).expect("entry"); + 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")] { + 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 + ); +} + +/// Both `HostLoad` providers reach the scheduler; their semantics are pinned +/// by the `scheduling::queue` unit tests. +#[tokio::test] +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() + }); + for worker_id in [1, 2, 3] { + core.upsert_worker(worker(worker_id)) + .await + .expect("worker upsert"); + } + + // 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"); + assert_eq!(response.worker_id, only); + } + } +} + +#[tokio::test] +async fn shutdown_cancels_listeners_but_keeps_parent_alive() { + let parent = CancellationToken::new(); + let core = local_core_with(test_config(true), 1, parent.clone()); + + 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!(core.cancel_token.is_cancelled()); + assert!(!parent.is_cancelled()); + assert_eq!(core.indexer_registry.listener_cancelled(1, 0), Some(true)); +} + +#[rstest::rstest] +#[case(1)] +#[case(2)] +#[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 = default_key(); + let request = || { + let mut request = select_request(); + request.prompt.token_ids = None; + request.prompt.block_hashes = Some(vec![11]); + request.prompt.sequence_hashes = Some(vec![101]); + request.prompt.isl_tokens = Some(4); + request + }; + + // Exercise an in-place update, then removing the last worker and adding a new one. + // An in-place update whose listener endpoints are unchanged keeps the rank's + // listener and index rows; only a replacement starts from an empty cache. + for (worker_id, cached_before_store) in [(1, 0), (1, 4), (2, 0)] { + if worker_id == 2 { + core.delete_worker(1).await.expect("delete last worker"); + } + core.upsert_worker(worker_with_kv_events(worker_id)) + .await + .expect("worker upsert"); + assert_eq!( + core.select(request()).await.unwrap().overlap.gpu, + cached_before_store + ); + + // Write through the registry used by listeners, not the selector's saved reference. + let indexer = core + .indexer_registry + .get_indexer(&key) + .unwrap() + .indexer + .clone(); + indexer + .apply_event_routed(store_event( + worker_id, + 0, + 1, + &[], + &[11], + StorageTier::Device, + )) + .await + .unwrap(); + indexer.dump_events().await.expect("flush indexer"); + let selected = core.select(request()).await.expect("select cached worker"); + assert_eq!(selected.worker_id, worker_id); + assert_eq!(selected.overlap.gpu, 4); + } + core.shutdown(); +} + +#[tokio::test] +async fn multi_rank_worker_with_replay_endpoint_is_incomplete() { + let core = local_core(test_config(true)); + + 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 = local_core(test_config(true)); + core.upsert_worker(worker_with_kv_events(1)) + .await + .expect("worker upsert"); + let entry = core.entry(&default_key()).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 upsert_moves_global_worker_id_between_routing_groups() { + 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"); + let indexed = |group: &str| { + core.indexer_registry + .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(); + 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_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(); + 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_eq!(indexed("group-b"), 0); +} + +#[tokio::test] +async fn shutdown_reports_not_ready_and_rejects_new_work() { + let core = local_core(test_config(false)); + 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 { + worker_id: Some(1), + prompt: prompt(), + ..replay_reservation("res-after-shutdown") + }) + .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 core = saturated_core(); + + 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 core = saturated_core(); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + let entry = core.entry(&default_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. +#[tokio::test] +async fn lease_admission_keeps_a_selection_whose_worker_drained_while_queued() { + let core = saturated_core(); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + let entry = core.entry(&default_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(lease_operation(req.prompt.view(), "queued", false)) + .await; + match run.result { + Ok(SelectionOutcome::Selected(selected)) => ( + selected.response.best_worker.worker_id, + selected.endpoint, + selected.booking.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; +} + +/// 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 = local_core(test_config(false)); + 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"); + + let req = reserve_request("leased"); + let run = core + .run_selection(lease_operation(req.prompt.view(), "leased", true)) + .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.booking.is_some(), + "the booking's handle goes to the host" + ); + 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. + 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 = local_core(test_config(false)); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + let key = default_key(); + 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()); +} + +/// 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 core = core_with( + test_config(false), + SelectionHost { + cache: HostCache { + index: KvIndexSource::Owned(Arc::new(PausedRecordIngress(record.clone()))), + shared: None, + }, + ..SelectionHost::default() + }, + None, + WorkerType::Aggregated, + Some( + SessionAffinityConfig::new(Duration::from_secs(10)) + .with_mode(SessionAffinityMode::Hard), + ), + ); + core.upsert_worker(worker(1)).await.expect("worker upsert"); + let key = default_key(); + 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 = 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(); + 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 = 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 { + worker_id: Some(1), + prompt: prompt(), + ..replay_reservation("taken") + }) + .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 = default_key(); + 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: 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 + .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(), + claim_id: 0, + 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 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( + #[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(); + 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 = local_core(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 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 [ + 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); + 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())), + ..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(&default_key()).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 = local_core(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 = local_core(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 = local_core(test_config(false)); + 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 = 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(); + 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 { + routing_group: "group-b".to_string(), + ..replay_reservation("cached") + }) + .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 = local_core(test_config(false)); + 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(&default_key()).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(local_core(test_config(false))); + 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 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 = default_key(); + 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 { + worker_id: Some(worker_id), + dp_rank: Some(0), + prompt: PromptRequest { + sequence_hashes: Some(vec![1, 2]), + isl_tokens: Some(8), + ..PromptRequest::default() + }, + effective_prefill_tokens: Some(8), + ..replay_reservation(&format!("occupy-{worker_id}")) + }) + .await + .expect("occupy worker"); + } + + let queued_core = Arc::clone(&core); + let queued = tokio::spawn(async move { + queued_core + .select_and_reserve(SelectAndReserveRequest { + prompt: PromptRequest { + block_hashes: Some(vec![11, 12]), + sequence_hashes: Some(vec![101, 102]), + isl_tokens: Some(8), + ..PromptRequest::default() + }, + ..reserve_request("refresh-selection") + }) + .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 { + core_with( + test_config(false), + SelectionHost::default(), + None, + WorkerType::Aggregated, + Some(SessionAffinityConfig::new(Duration::from_secs(10)).with_mode(mode)), + ) +} + +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 { + 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) +} + +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 +} + +/// 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 + .select_and_reserve(session_reservation("r1", "s")) + .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 + .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)); +} + +#[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; + 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. + 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); + + 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::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 { + 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) + ); +} + +/// 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 +/// 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; + 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"), Some(replacement)); + assert_eq!( + lease_count(&core, "s"), + Some(2), + "r2 and the queued initializer both keep the binding alive" + ); + + 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); +} + +#[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(); + 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(&default_key()).expect("entry"); + wait_until("rejected booking release", || { + !entry.scheduler.has_request("pending") + }) + .await; + 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; + 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(&default_key()).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, first) = bound_session(SessionAffinityMode::Soft).await; + 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 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: default_key(), + 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 = default_key(); + 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!(bound_worker(&core, "session"), 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!(bound_worker(&core, "shared-session"), Some(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 = local_core(test_config(true)); + let request = worker_with_kv_events(1); + core.upsert_worker(request.clone()).await.unwrap(); + 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()) + .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(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(&default_key()).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, + detaches: AtomicUsize, + } + #[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(()) + } + 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 { + 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)); + // 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" + ); + wait_until("booking release", || !entry.scheduler.has_request("live")).await; +} + +#[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(&default_key()).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..2d87f3706c84 --- /dev/null +++ b/lib/kv-router/src/services/selection/core/workers.rs @@ -0,0 +1,409 @@ +// 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 std::collections::HashSet; + +use super::reservations::{ReservationIndexObserver, spawn_reservation_index_sweep}; +use super::*; + +impl SelectionCore { + pub async fn upsert_worker( + &self, + req: WorkerRequest, + ) -> Result { + 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. + 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; + 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( + &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), None).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); + // 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) + } + + 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(()) + } + + /// 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. 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, + previous: Option, + 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()) + { + 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()); + match deferred { + Some(affected) if !moved_partition => { + affected.insert(record.key()); + } + _ => 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: self.host.replication.replica_worker_policy, + 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 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(), + 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> { + 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) { + let KvIndexSource::Owned(ingress) = &self.host.cache.index; + ingress.detach(&self.indexer_registry, record).await; + } + + pub(super) fn publish_scheduler_config(&self, key: &RoutingPartitionId) { + let Some(entry) = self.entry(key) else { + return; + }; + 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| { + if *current == workers { + false + } else { + *current = workers; + 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/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/ingress.rs b/lib/kv-router/src/services/selection/ingress.rs new file mode 100644 index 000000000000..272eef76c6eb --- /dev/null +++ b/lib/kv-router/src/services/selection/ingress.rs @@ -0,0 +1,198 @@ +// 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::identity::RoutingPartitionId; +use crate::services::indexer::backend::Indexer; +use crate::services::indexer::registry::WorkerRegistry; + +/// Builds and feeds a partition's index ([`super::KvIndexSource::Owned`]). +/// +/// 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`. Consulted only when + /// the core subscribes to worker KV events. + fn missing_metadata(&self, record: &WorkerCatalogRecord) -> Vec { + let _ = record; + Vec::new() + } + + /// 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 reconcile( + &self, + registry: &WorkerRegistry, + previous: Option<&WorkerCatalogRecord>, + record: &WorkerCatalogRecord, + ) -> Result<(), SelectionError> { + let _ = (registry, previous, record); + Ok(()) + } + + /// `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 +/// 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 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(); + let mut missing: Vec = 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(); + // 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( + &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 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 {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() { + // 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 if listener_present { + registry + .deregister_dp_rank( + record.worker_id, + rank, + &previous.model_name, + &previous.routing_group, + ) + .await + .map_err(|error| SelectionError::Internal(error.to_string()))?; + } else { + registry + .remove_dp_rank_blocks(record.worker_id, rank, &previous.key()) + .await; + } + } + } + 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, + 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 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; + } + 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/input.rs b/lib/kv-router/src/services/selection/input.rs index aaf688a4bf26..f14702e5a673 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; @@ -41,26 +43,64 @@ 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<'_> { + /// `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, - self.lora_name.as_deref(), - self.cache_namespace.as_deref(), + block_size, + 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 +117,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 +150,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(|| { @@ -129,8 +169,8 @@ impl PromptRequest { Ok(normalize_hashes(block_hashes, sequence_hashes, isl_tokens)?.block_hashes) } - fn routing_tokens_and_mm_infos(&self) -> Option> { - if let Some(mm_routing_info) = &self.mm_routing_info + 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() { return Some(( @@ -140,18 +180,18 @@ 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)) } } 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, @@ -159,19 +199,27 @@ 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 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), } } @@ -226,6 +274,8 @@ fn normalize_hashes( .collect(), sequence_hashes: signed_sequence_hashes(sequence_hashes), isl_tokens, + block_hashing: Duration::ZERO, + seq_hashing: Duration::ZERO, }) } @@ -237,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/membership.rs b/lib/kv-router/src/services/selection/membership.rs new file mode 100644 index 000000000000..9158303561fa --- /dev/null +++ b/lib/kv-router/src/services/selection/membership.rs @@ -0,0 +1,532 @@ +// 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. 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; +use std::time::Duration; + +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; + +/// 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] +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, + /// Passes started by `run`, including retries. + #[cfg(test)] + applies: Arc, +} + +impl CatalogReconciler { + pub fn new(core: Arc) -> Self { + Self { + core, + observer: None, + converged: HashMap::new(), + tracked: HashSet::new(), + #[cfg(test)] + applies: Arc::default(), + } + } + + 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 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: &[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() { + return Err(SelectionError::BadRequest(format!( + "duplicate worker_id {worker_id} in membership snapshot" + ))); + } + } + + 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) { + to_upsert.push(request.clone()); + } + } + // 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 { + 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); + 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); + } + } + // A worker that was never registered is not an error (idempotent). + Err(SelectionError::NotFound(_)) => {} + // 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); + } + 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() => { + 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()") + } + }; + let outcome = self.apply(&snapshot).await; + #[cfg(test)] + self.applies + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); + match outcome { + Ok(()) => pending = None, + Err(error) => { + tracing::warn!( + %error, + retry_in = ?RECONCILE_RETRY_DELAY, + "membership reconcile failed; retrying" + ); + pending = Some(snapshot); + } + } + } + } +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use super::*; + use crate::config::KvRouterConfig; + use crate::services::selection::SelectionCacheConfig; + + fn core() -> Arc { + 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 { + 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) + } + + /// 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, counter) = reconciler(&core); + + 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(&[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(&[]).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(&[]).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, counter) = reconciler(&core); + + 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(&[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(&[moved]).await.expect("apply"); + assert_eq!(counter.upserts(), 2); + 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 stale_worker_is_deleted_when_another_upsert_fails() { + let core = core(); + let (mut reconciler, counter) = reconciler(&core); + + reconciler + .apply(&[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(std::slice::from_ref(&changed)) + .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(&[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"); + // 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() + .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(); + let mut reconciler = CatalogReconciler::new(Arc::clone(&core)); + + let error = reconciler + .apply(&[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 2f24a0b77ad3..d10a4bb91313 100644 --- a/lib/kv-router/src/services/selection/mod.rs +++ b/lib/kv-router/src/services/selection/mod.rs @@ -6,10 +6,13 @@ //! 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; +mod ingress; mod input; +mod membership; mod pending; mod policy_registry; mod server; @@ -20,10 +23,19 @@ mod types; mod tests; pub use crate::WorkerSelectionPolicyFactory; -pub use crate::services::common::replica_sync::ReplicaPeerError; -pub use core::{SelectionCore, SelectionServiceConfig}; +pub use crate::services::common::replica_sync::{ + HostReplicaChannels, HostReplicaSyncFactory, ReplicaIngressObserver, ReplicaPeerError, + SchedulerLoadSink, +}; +pub use core::{ + HostCache, HostEligibility, HostLoad, HostReplication, HostTelemetry, KvIndexSource, Selected, + SelectionAdmission, SelectionCore, SelectionHost, SelectionOperation, SelectionOutcome, + SelectionPartition, SelectionRun, SelectionScheduler, SelectionServiceConfig, SessionBinding, +}; pub use error::SelectionError; -pub use input::PromptRequest; +pub use ingress::KvEventIngress; +pub use input::{PromptRequest, PromptView}; +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, @@ -36,9 +48,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, SelectionWorkerConfig, + 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/pending.rs b/lib/kv-router/src/services/selection/pending.rs index 7008d251cdaa..c1ce85de6b28 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,11 @@ 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>, + /// The session the selection steered by, bound when the booking replays. + pub session_id: Option, } struct Entry { @@ -71,10 +76,14 @@ 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() + selection.lora_name.as_deref().map_or(0, str::len) + + selection.session_id.as_deref().map_or(0, str::len) } struct State { @@ -267,6 +276,8 @@ mod tests { expected_output_tokens: Some(16), track_prefill_tokens: true, lora_name: None, + routing_hashes: None, + session_id: None, } } diff --git a/lib/kv-router/src/services/selection/service.rs b/lib/kv-router/src/services/selection/service.rs index cfa8ac0d34b4..3be85bd12052 100644 --- a/lib/kv-router/src/services/selection/service.rs +++ b/lib/kv-router/src/services/selection/service.rs @@ -2,18 +2,22 @@ // SPDX-License-Identifier: Apache-2.0 use std::sync::Arc; +use std::time::Duration; use tokio_util::sync::CancellationToken; +use crate::WorkerSelectionPolicyFactory; use crate::config::KvRouterConfig; use crate::protocols::WorkerId; 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, SelectionServiceConfig}; +use super::affinity::SessionAffinityConfig; +use super::core::{KvIndexSource, SelectionCore, SelectionHost, SelectionServiceConfig}; use super::error::SelectionError; use super::pending::SelectionCacheConfig; use super::policy_registry::WorkerSelectionPolicyRegistry; @@ -33,6 +37,9 @@ pub struct SelectionServiceBuilder { selection_cache: SelectionCacheConfig, worker_type: WorkerType, 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. @@ -70,9 +77,35 @@ impl SelectionServiceBuilder { selection_cache: SelectionCacheConfig::default(), worker_type, worker_selection_policy_registry, + host: SelectionHost::default(), + worker_selection_policy_factory: None, + session_affinity_ttl: None, } } + /// 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 + } + + /// 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 + } + + /// What the embedding host supplies to every partition this service creates. + pub fn host(mut self, host: SelectionHost) -> Self { + self.host = host; + self + } + pub fn indexer_threads(mut self, indexer_threads: usize) -> Self { self.indexer_threads = indexer_threads; self @@ -83,6 +116,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: 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; @@ -95,13 +135,21 @@ 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)?; - 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 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( @@ -123,13 +171,16 @@ impl SelectionServiceBuilder { cancel_token.clone(), replica_config, worker_selection_policy_factory, + self.host, self.worker_type, false, self.selection_cache, tracking_hash, + indexer_policy, + self.session_affinity_ttl.map(SessionAffinityConfig::new), )); - 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) => { @@ -146,7 +197,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| { @@ -154,6 +206,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 @@ -187,6 +246,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); + } builder } } @@ -256,6 +318,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 + } + /// The port this service uses for replica synchronization, if enabled. pub fn replica_sync_port(&self) -> Option { self.replica_sync_port @@ -396,10 +463,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/kv-router/src/services/selection/tests.rs b/lib/kv-router/src/services/selection/tests.rs index 1a9824442765..4fb73b48d772 100644 --- a/lib/kv-router/src/services/selection/tests.rs +++ b/lib/kv-router/src/services/selection/tests.rs @@ -123,22 +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 - .normalize_for_selection( - false, - 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 { @@ -565,19 +573,7 @@ fn keyed_prompt_tracking_leaves_indexer_hashes_public() { })) .unwrap(); - let normalized = request - .normalize_for_selection( - false, - 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, @@ -602,21 +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 - .normalize_for_selection( - false, - 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(); @@ -650,6 +632,7 @@ fn keyed_reservation_hashes_directly_from_tokens() { }; let normalized = request + .view() .normalize_for_reservation( false, TrackingHashInput { @@ -696,17 +679,9 @@ fn keyed_hash_only_inputs_remain_trusted_for_selection_and_reservation() { block_size: 4, }; - let selection = request - .normalize_for_selection( - false, - TrackingHashInput { - context: &context, - scope, - assume_kv_reuse: true, - }, - ) - .unwrap(); + let selection = normalize_with(&request, &context, true).unwrap(); let reservation = request + .view() .normalize_for_reservation( false, TrackingHashInput { @@ -750,6 +725,7 @@ fn randomized_reservation_uses_canonical_complete_block_count() { })) .unwrap(); let normalized = request + .view() .normalize_for_reservation( false, TrackingHashInput { @@ -1759,11 +1735,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 +1773,7 @@ async fn selector_replica_sync_propagates_request_lifecycle() { port: 8092, threads: 1, indexer_peers: Vec::new(), + session_affinity_ttl: None, replica_sync_port: Some(port_a), replica_sync_peers: Vec::new(), kv_router_config: test_config(), @@ -2002,3 +1980,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); +} diff --git a/lib/kv-router/src/services/selection/types.rs b/lib/kv-router/src/services/selection/types.rs index cecd91588362..3893f19f6887 100644 --- a/lib/kv-router/src/services/selection/types.rs +++ b/lib/kv-router/src/services/selection/types.rs @@ -6,18 +6,19 @@ 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, RoutingConstraints, WorkerAffinityTarget, WorkerConfigLike, - WorkerId, WorkerWithDpRank, + DpRank, KvHintTransferWorkerMetadata, 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, WorkerSelectionInputTrigger}; 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 { @@ -49,6 +50,17 @@ 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, + /// 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 { @@ -91,6 +103,24 @@ impl WorkerConfigLike for SelectionWorkerConfig { fn kv_transfer_preferred_weight(&self) -> Option { self.kv_transfer_preferred_weight } + + fn kv_hint_transfer_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(KvHintTransferWorkerMetadata { + 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,12 +148,18 @@ 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 = "Option::is_none")] + pub kv_event_source_mode: Option, #[serde(default, skip_serializing_if = "Vec::is_empty")] pub not_schedulable_reasons: Vec, } 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, @@ -145,6 +181,9 @@ 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, + kv_event_source_mode: req.kv_event_source_mode, not_schedulable_reasons: Vec::new(), } } @@ -161,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) -> std::ops::Range { let start = self.dp_start(); let size = self.dp_size(); start..start.saturating_add(size) @@ -181,6 +220,9 @@ 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(), + kv_event_source_mode: self.kv_event_source_mode.clone(), }) } @@ -195,11 +237,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) { @@ -215,18 +253,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 } } @@ -255,11 +281,14 @@ 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(), + kv_event_source_mode: None, } } } -#[derive(Debug, Deserialize)] +#[derive(Debug, Clone, PartialEq, Deserialize)] pub struct WorkerRequest { pub worker_id: WorkerId, #[serde(default = "default_model_name")] @@ -285,6 +314,15 @@ 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, + #[serde(default)] + pub kv_event_source_mode: Option, } #[derive(Debug, Deserialize)] @@ -305,6 +343,12 @@ 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>, + #[serde(default)] + pub kv_event_source_mode: Option, } impl WorkerCatalogRecord { @@ -360,6 +404,15 @@ 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; + } + if patch.kv_event_source_mode.is_some() { + self.kv_event_source_mode = patch.kv_event_source_mode; + } } } @@ -376,12 +429,31 @@ 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 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, +} + +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 +469,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 +479,66 @@ 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. +/// +/// `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, + #[serde(default)] + pub parent_session_id: Option, + #[serde(default)] + pub session_final: Option, + #[serde(default)] + pub input_trigger: Option, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum SelectionInputTrigger { + UserMessage, + ToolResult, + Other, +} + +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, + input_trigger, + } = context; + SessionContext::new( + session_id, + parent_session_id, + session_final, + input_trigger.map(|trigger| match trigger { + SelectionInputTrigger::UserMessage => WorkerSelectionInputTrigger::UserMessage, + SelectionInputTrigger::ToolResult => WorkerSelectionInputTrigger::ToolResult, + SelectionInputTrigger::Other => WorkerSelectionInputTrigger::Other, + }), + ) + } +} + +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 +608,38 @@ 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, + /// 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 kv_hint: Option, +} + +/// Load snapshot of the chosen worker, as the scheduler projected it for this +/// request. +#[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 +667,51 @@ 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, + "input_trigger": "user_message" + } + })) + .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_eq!( + context.input_trigger(), + Some(WorkerSelectionInputTrigger::UserMessage) + ); + } + + #[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); + + 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()); + } +} 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/Cargo.toml b/lib/llm/Cargo.toml index aad78af22bdd..c847abc86e8f 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 } @@ -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 3b4b0be8ea0e..d206edc630fc 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::{ @@ -2055,11 +2055,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, @@ -2072,33 +2071,26 @@ impl ModelManager { /// 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_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, - metric_worker_type, - ); 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( + crate::kv_router::SchedulerLoadSender::disabled(parent_token.child_token()); + self.kv_chooser_for_with_policy_and_client( client, kv_cache_block_size, - selector, + policy, kv_router_config, prefill_load_estimator, worker_role, @@ -2147,11 +2139,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, @@ -2164,21 +2155,18 @@ impl ModelManager { /// 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_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, @@ -2193,10 +2181,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, @@ -2211,11 +2199,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, @@ -2224,10 +2212,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()); @@ -2255,12 +2240,20 @@ 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); + // 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 shared_cache: Option> = if selector - .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() .map(|c| c.shared_cache_type) @@ -2273,7 +2266,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()), )) } @@ -2282,12 +2275,9 @@ impl ModelManager { 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 = wants_cache || effective_kv_router_config.serve_indexer || effective_kv_router_config.enable_session_prefix_index || matches!( @@ -2312,7 +2302,7 @@ impl ModelManager { workers_with_configs, kv_source_membership, kv_cache_block_size, - selector, + SelectionPolicySource::Prepared(policy), kv_router_config, prefill_load_estimator, worker_role, @@ -2419,7 +2409,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/discovery/watcher.rs b/lib/llm/src/discovery/watcher.rs index 38a068244192..0d7e326f2f03 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,10 +29,10 @@ use crate::{ entrypoint::{self, ChatEngineFactoryCallback, RouterConfig}, http::service::metrics::Metrics, kv_router::{ - EncoderRouter, PrefillRouter, RouterLoadSource, RoutingLoadContext, WorkerSelectorFactory, + EncoderRouter, PrefillRouter, RouterLoadSource, RoutingLoadContext, SelectionPolicySource, }, local_model::runtime_config::{ - ModelRuntimeConfig, TokenizerBackend, VLLM_INFERENCE_V1_GENERATE_CAPABILITY, + TokenizerBackend, VLLM_INFERENCE_V1_GENERATE_CAPABILITY, VLLM_QWEN_VIDEO_PROCESSOR_CONTRACT_RUNTIME_KEY, }, model_card::ModelDeploymentCard, @@ -185,10 +182,7 @@ pub enum ModelUpdate { Removed(ModelDeploymentCard), } -pub struct ModelWatcher -where - Sel: WorkerSelector, -{ +pub struct ModelWatcher { manager: Arc, drt: DistributedRuntime, router_config: RouterConfig, @@ -212,9 +206,7 @@ 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. - require_typed_worker_role: bool, + selection_policy: SelectionPolicySource, } pub(crate) struct PreparedWorkerSet { @@ -271,7 +263,7 @@ fn removed_model_cards( .collect() } -impl ModelWatcher { +impl ModelWatcher { #[allow(clippy::too_many_arguments)] pub fn new( runtime: DistributedRuntime, @@ -283,7 +275,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, @@ -292,23 +284,12 @@ impl ModelWatcher { chat_engine_factory, 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, -{ #[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, @@ -317,8 +298,7 @@ where chat_engine_factory: Option, prefill_load_estimator: Option>, metrics: Arc, - require_typed_worker_role: bool, - worker_selector_factory: WorkerSelectorFactory, + selection_policy: SelectionPolicySource, ) -> Self { Self { manager: model_manager, @@ -336,8 +316,7 @@ where tokenizer_backend: None, tokenizer_fallback_enabled: None, generate_engine_capabilities: Vec::new(), - worker_selector_factory, - require_typed_worker_role, + selection_policy, } } @@ -434,7 +413,7 @@ where 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)?; // Prepare without exact video routing unless the cohort agreed on a contract. if spec.video_contract.is_none() @@ -597,21 +576,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, @@ -648,13 +622,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, @@ -687,7 +661,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, @@ -1045,10 +1019,7 @@ where } #[async_trait] -impl ControllerHost for ModelWatcher -where - Sel: WorkerSelector + Send + 'static, -{ +impl ControllerHost for ModelWatcher { type Prepared = PreparedWorkerSet; fn normalize( @@ -1391,11 +1362,17 @@ 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(_) | SelectionPolicySource::Prepared(_) + ) && card.worker_type.is_none() + { anyhow::bail!( "custom worker-selection policies require model cards with an explicit worker_type" ); @@ -2738,12 +2715,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.rs b/lib/llm/src/entrypoint/input.rs index 6f5d368b693a..42ccd83867e8 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..048438914355 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"); @@ -159,20 +150,15 @@ 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. @@ -192,18 +178,16 @@ where chooser, load_context, affinity, - session_affinity_mode, )) } _ => { 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, - session_affinity_mode, lora, )?) } @@ -224,7 +208,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 +224,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. @@ -269,11 +250,15 @@ where 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, session_affinity_mode)?; + affinity.enable_replica_sync(router_client.clone()).await?; + Some(affinity) + } + _ => create_affinity_coordinator(ttl, session_affinity_mode, router_client.clone()).await?, + }; let embedding_cache_indexer = if enable_multimodal_cache_indexer && matches!(router_mode, RouterMode::DeviceAwareWeighted) @@ -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, @@ -317,7 +302,6 @@ where &model_manager, &endpoint_id, affinity, - session_affinity_mode, load_context, )?; if router_mode.is_kv_routing() && prefill_router.conditional_disagg_enabled() { @@ -483,10 +467,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 ba14540340a4..69fbb1856b6e 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, model_type::ModelType, namespace::NamespaceFilter, types::openai::{ @@ -23,20 +23,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 { @@ -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 @@ -93,33 +93,17 @@ impl HttpFrontend { super::initialize_input(&distributed_runtime, &engine_config).await; - let result = 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 selection_policy = match self.worker_selection_policy_factory { + Some(factory) => SelectionPolicySource::Factory(factory), + None => SelectionPolicySource::Registry, }; + let result = run_with_selection_policy( + distributed_runtime, + engine_config, + self.frontend_route_extensions, + selection_policy, + ) + .await; active_input.release_and_drain().await; @@ -149,16 +133,12 @@ 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(), @@ -261,8 +241,7 @@ where model.runtime_config().tokenizer_backend, model.runtime_config().tokenizer_fallback_enabled, generate_engine_capabilities, - require_typed_worker_role, - worker_selector_factory.clone(), + selection_policy.clone(), ) .await?; http_service @@ -336,7 +315,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, @@ -351,12 +330,8 @@ async fn run_watcher( tokenizer_backend: Option, 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`. @@ -365,7 +340,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, @@ -374,8 +349,7 @@ where chat_engine_factory, 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 15caa51cc4ff..c177e137b511 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -12,33 +12,33 @@ use anyhow::Result; use dynamo_kv_router::{ DEFAULT_ROUTING_GROUP, KvSchedulerError, PrefillLoadEstimator, RoutingPartitionRef, SessionPrefixIndexer, SharedKvCache, TrackingHashAlgorithm, TrackingHashContext, - TrackingHashScope, + TrackingHashScope, WorkerSelectionPolicy, WorkerSelectionPolicyFactory, config::{KvRouterConfig, RouterConfigOverride, min_initial_workers_from_env}, indexer::{ - ApproximateLruIncarnation, ApproximateLruStats, KvRouterError, MatchDetails, - RoutingDecisionHashes, - }, - kv_hints::{ - KvHint, KvHintAction, KvSourceLocationsPayload, KvTransferCandidateSource, - KvTransferCandidates, + ApproximateLruIncarnation, ApproximateLruStats, KvRouterError, RoutingDecisionHashes, }, + kv_hints::KvHint, protocols::KV_EVENT_SUBJECT, protocols::{ - BlockExtraInfo, BlockHashOptions, ExternalSequenceBlockHash, 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, TieredOverlapRefresher, WorkerAvailabilityProvider, - effective_prefill_tokens, overlap::cache_hit_estimates_from_tiered_matches, + CacheHitEstimates, OverlapAnalysis, OverloadedWorkerProvider, PotentialLoad, + WorkerAvailabilityProvider, effective_prefill_tokens, + overlap::cache_hit_estimates_from_tiered_matches, + queue::{BookingHandle, SchedulerBookingDescriptor}, }, selector::WorkerInputs, + services::selection::{ + PromptView, Selected, SelectionAdmission, SelectionError, SelectionOperation, + SelectionOutcome, SelectionRun, SessionBinding, + }, }; use dynamo_runtime::{ CancellationToken, component::{Client, Endpoint}, - discovery::DiscoveryQuery, error::{DynamoError, ErrorType}, pipeline::{ AsyncEngine, AsyncEngineContextProvider, Error, ManyOut, ResponseStream, SingleIn, @@ -49,14 +49,12 @@ use dynamo_runtime::{ traits::DistributedRuntimeProvider, }; use futures::stream; -use tracing::Instrument; // 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; pub mod indexer; pub mod metrics; @@ -64,18 +62,15 @@ 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; 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,19 +79,118 @@ pub use routing_load::{ use crate::{ discovery::{KvSourceMembershipWatch, RuntimeConfigWatch}, - kv_router::{ - scheduler::{DefaultWorkerSelector, KvScheduler, PotentialLoad}, - sequence::{SequenceError, SequenceRequest}, - }, + kv_router::sequence::{SequenceError, SequenceRequest}, local_model::runtime_config::ModelRuntimeConfig, worker_type::WorkerType, }; -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), + /// A factory whose instance for the router's own partition is already + /// built and probed; see [`PreparedSelectionPolicy`]. + Prepared(PreparedSelectionPolicy), +} + +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::Prepared(prepared) => Ok(prepared.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) + }), + }, + ), + } + } + + /// 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, + )) + } +} + +/// 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 = 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( + &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)] struct ApproximateLruRankRegistration { @@ -360,49 +454,31 @@ 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 { track_lifecycle: bool }, + WithAdmission, WithoutAdmission, } +/// 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) attempt: AdmissionAttempt, + 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 { #[doc(hidden)] - pub fn into_parts(self) -> (FindBestMatchOutcome, AdmissionAttempt) { - (self.outcome, self.attempt) + pub fn into_parts(self) -> (FindBestMatchOutcome, Option) { + (self.outcome, self.booking) } } -pub(super) enum FindBestMatchInnerOutcome { - WithAdmission(AdmittedFindBestMatchOutcome), - WithoutAdmission(FindBestMatchAdvisoryOutcome), -} - #[derive(Debug, Clone, Copy, Default)] pub(crate) struct WorkerCacheHitEstimate { pub effective_overlap_blocks: f64, @@ -427,24 +503,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 @@ -486,35 +551,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" - ); -} - -fn matched_hash_for_worker( - match_details: &MatchDetails, - worker: WorkerWithDpRank, -) -> Option { - match_details.last_matched_hashes.get(&worker).copied() -} - // for router discovery registration pub const KV_ROUTER_ENDPOINT: &str = "router-discovery"; @@ -527,23 +563,11 @@ 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 -where - Sel: dynamo_kv_router::selector::WorkerSelector, -{ +pub struct KvRouter { indexer: Indexer, - scheduler: KvScheduler>, + selection: embedded::EmbeddedSelection, required_worker_inputs: dynamo_kv_router::selector::WorkerInputs, workers_with_configs: RuntimeConfigWatch, block_size: u32, @@ -552,19 +576,14 @@ where 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>, - /// 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>, + shared_cache: Option>, endpoint_registration: Option, teardown_task_guard: Option, /// Optional session-aware logical prefix index. @@ -586,10 +605,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, @@ -597,13 +613,13 @@ 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, model_name: Option, is_eagle: bool, - shared_cache: Option>, + shared_cache: Option>, lora_filter: Option>, ) -> Result { Self::new_with_worker_role( @@ -612,7 +628,7 @@ where workers_with_configs, kv_source_membership, block_size, - selector, + policy, kv_router_config, prefill_load_estimator, None, @@ -632,19 +648,18 @@ 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, metric_worker_type: &'static str, model_name: Option, is_eagle: bool, - shared_cache: Option>, + 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, @@ -652,7 +667,7 @@ where workers_with_configs, kv_source_membership, block_size, - selector, + policy, kv_router_config, prefill_load_estimator, worker_role, @@ -674,19 +689,29 @@ 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, 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, ) -> 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 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) { @@ -694,8 +719,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())?; @@ -718,19 +741,21 @@ where .enable_session_prefix_index .then(|| Arc::new(SessionPrefixIndexer::new())); - let indexer = if cache_required { - Indexer::new( - component, - &kv_router_config, - block_size, - model_name.as_deref(), - cancellation_token.child_token(), - session_prefix_index.clone(), - ) - .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(), + session_prefix_index: session_prefix_index.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 { @@ -774,13 +799,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()); @@ -789,88 +807,44 @@ 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, + let request_leases = + request_lease::RequestLeaseManager::new(cancellation_token.child_token()); + let (selection, replica_ingress) = 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, + shared_cache: shared_cache.clone(), + lora_worker_filter: lora_filter.map(|filter| { + filter as Arc + }), + ingress: Arc::clone(&ingress) + as Arc, + scheduler_load, + endpoint: endpoint.clone(), + router_id: endpoint.drt().discovery().instance_id(), + policy_factory, + }, 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, + 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(), - ); - if !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) - { - 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 - }; - + request_leases.set_scheduler(selection.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 { indexer, - scheduler, + selection, required_worker_inputs, workers_with_configs, block_size, @@ -879,14 +853,12 @@ where 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, teardown_task_guard: None, session_prefix_index, @@ -904,9 +876,7 @@ where &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); } @@ -927,14 +897,6 @@ where self.required_worker_inputs } - /// Cancel background work and wait for KV event ingestion to stop. - pub async fn shutdown(mut self) { - self.cancellation_token.cancel(); - if let Some(subscription) = self.kv_event_subscription.take() { - subscription.shutdown().await; - } - } - pub fn is_eagle(&self) -> bool { self.is_eagle } @@ -1022,107 +984,6 @@ where 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, @@ -1160,21 +1021,20 @@ where #[doc(hidden)] pub async fn enroll_public_request_attempt( &self, - request_id: String, - worker: WorkerWithDpRank, - attempt_id: AttemptId, + booking: BookingHandle, routing_decision: Option, ) -> Result<(), KvRouterError> { + // Nothing awaits between taking the booking over and registering it. + 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| { - 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()); @@ -1238,53 +1098,6 @@ where .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), Some(lora_name)) = (self.lora_filter.as_ref(), lora_name) 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) - } - /// Give these tokens, find the worker with the best weighted cache hit. /// Returns the full match details for the selected worker. /// @@ -1292,101 +1105,6 @@ where /// 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 { - track_lifecycle: true, - }, - ) - .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, @@ -1427,20 +1145,14 @@ where 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(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)] @@ -1463,86 +1175,27 @@ where 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 { - track_lifecycle: false, - }, - ) - .await? - { - FindBestMatchInnerOutcome::WithAdmission(admitted) => Ok(admitted), - FindBestMatchInnerOutcome::WithoutAdmission(_) => { - unreachable!("with-admission routing returned advisory outcome") - } - } - } - - #[allow(clippy::too_many_arguments)] - pub async fn find_best_match_details_without_admission( - &self, - context_id: Option<&str>, - tokens: &[u32], - block_mm_infos: Option<&[Option]>, - router_config_override: Option<&RouterConfigOverride>, - 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, - ) -> 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") - } - } + 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)] @@ -1566,57 +1219,12 @@ where 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"); } - 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 is_admitted_routing = matches!(admission, FindBestMatchAdmission::WithAdmission); let session_index_context = if is_admitted_routing { self.session_prefix_index .as_ref() @@ -1625,142 +1233,92 @@ where } else { None }; - let session_block_hashes = session_index_context.as_ref().map(|_| block_hashes.clone()); - 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, + let core_admission = match admission { + FindBestMatchAdmission::WithAdmission if update_states => SelectionAdmission::Lease { + request_id: context_id.expect("validated above").to_string(), }, - ) - .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::WithAdmission => SelectionAdmission::Query { + request_id: context_id.map(str::to_string), }, - 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)), + FindBestMatchAdmission::WithoutAdmission => SelectionAdmission::Advisory { + request_id: context_id.map(str::to_string), }, }; - + let SelectionRun { + result: outcome, + lookup, + } = self + .selection + .run_selection(SelectionOperation { + key: self.selection.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: lora_name.as_deref(), + cache_namespace: cache_namespace.as_deref(), + 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: return_routing_hashes || session_index_context.is_some(), + 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(AdmittedFindBestMatchOutcome { + outcome: FindBestMatchOutcome::QueueRejected { rejection }, + booking: None, + advisory_load: None, + }); + } + 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, + isl_tokens, + kv_hint, + routing_hashes, + shared_cache_hits, + booking, + .. + } = selected; + if update_states && is_admitted_routing && booking.is_none() { + anyhow::bail!("booked selection returned no booking handle"); + } // Indexing failures never affect routing. - if let (Some(session_id), Some(block_hashes)) = ( - session_index_context.as_ref(), - session_block_hashes.as_deref(), - ) && let Some(mut residency_version) = - self.indexer.session_residency_version(response.best_worker) + if let (Some(session_id), Some(block_hashes)) = + (session_index_context.as_ref(), routing_hashes.as_deref()) + && let Some(mut residency_version) = + self.indexer.session_residency_version(response.best_worker) { for attempt in 0..2 { match self @@ -1786,8 +1344,10 @@ where break; } - if let Some(matched_hash) = - matched_hash_for_worker(&match_details, response.best_worker) + if let Some(matched_hash) = match_details + .last_matched_hashes + .get(&response.best_worker) + .copied() && let Err(err) = self.indexer.enqueue_session_match( session_id, response.best_worker, @@ -1807,135 +1367,61 @@ where } } - 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(), - ) + let routing_hashes = if return_routing_hashes { + routing_hashes.map(RoutingDecisionHashes::from_local_hashes) } 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 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); } - 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" + // 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" ); - - 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)] - 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, + Ok(AdmittedFindBestMatchOutcome { + outcome: FindBestMatchOutcome::Routed { + worker: response.best_worker, overlap_blocks, - .. - } => Ok((worker, overlap_blocks)), - FindBestMatchOutcome::QueueRejected { rejection } => Err(rejection.into()), - } - } - - /// Register externally-provided workers in the slot tracker. - pub fn register_workers(&self, worker_ids: &HashSet) { - self.scheduler.register_workers(worker_ids); + 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, + }) } #[allow(clippy::too_many_arguments)] @@ -1976,7 +1462,8 @@ where self.prefill_load_hint_for(isl_tokens, cached_tokens, track_prefill_tokens); let admission = self - .scheduler + .selection + .scheduler() .add_request_admitted(SequenceRequest { request_id: request_id.clone(), token_sequence: maybe_seq_hashes, @@ -1996,7 +1483,7 @@ where }; self.request_leases .register_detached( - scheduler::SchedulerBookingDescriptor { + SchedulerBookingDescriptor { request_id, worker, attempt_id, @@ -2007,7 +1494,10 @@ where } 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(()) } @@ -2016,24 +1506,15 @@ where if self.request_leases.finish_request(request_id).await { return Ok(()); } - self.scheduler.free(request_id).await + self.selection.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( + pub(crate) fn affinity_coordinator( &self, - request_id: &str, - worker: WorkerWithDpRank, - ) -> Result<(), SequenceError> { - self.scheduler.free_if_worker(request_id, worker).await - } - - #[doc(hidden)] - pub(crate) fn booking_cleanup(&self) -> scheduler::SchedulerBookingCleanup { - self.scheduler.booking_cleanup() + ttl: std::time::Duration, + mode: crate::session_affinity::SessionAffinityMode, + ) -> anyhow::Result { + self.selection.affinity_coordinator(ttl, mode) } pub(crate) fn request_lease_manager(&self) -> &request_lease::RequestLeaseManager { @@ -2042,21 +1523,23 @@ where pub(crate) async fn mark_prefill_completed_if_booking( &self, - booking: &scheduler::SchedulerBookingDescriptor, + booking: &SchedulerBookingDescriptor, ) -> Result<(), KvSchedulerError> { - self.scheduler + self.selection + .scheduler() .mark_prefill_completed_if_booking(booking) .await + .map(|_| ()) } /// 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( @@ -2099,7 +1582,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() + self.selection.worker_type() } /// Return the worker's unique global DP rank when it owns exactly one rank. @@ -2109,20 +1592,13 @@ where (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 + self.selection + .scheduler() .enqueue_output_block_if_booking(booking, decay_fraction) .await } @@ -2155,27 +1631,6 @@ where 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, @@ -2186,20 +1641,9 @@ where 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 @@ -2236,7 +1680,7 @@ 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( + Ok(self.selection.scheduler().get_potential_loads( maybe_seq_hashes, isl_tokens, cache_hit_estimates.cached_tokens.into_iter().collect(), @@ -2322,11 +1766,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, @@ -2442,10 +1882,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(); @@ -2458,24 +1895,14 @@ mod tests { use std::collections::{HashMap, HashSet}; use async_trait::async_trait; - use dynamo_kv_router::{ - WorkerSelectionInput, - 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, - }, - }; + 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, + }; #[test] fn all_filtered_workers_map_to_unavailable() { @@ -2485,6 +1912,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] @@ -2628,47 +2062,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, @@ -2690,70 +2083,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()) @@ -2766,31 +2179,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, - DefaultWorkerSelector::new(Some(config.clone()), "decode"), - Some(config), + SelectionPolicySource::Registry, None, worker_role, "decode", - None, - false, - None, - None, + KvRouterConfig { + skip_initial_worker_wait: true, + router_event_threads: 1, + ..Default::default() + }, ) .await } @@ -2800,7 +2201,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 @@ -2815,41 +2216,29 @@ 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, - LoadOnlySelector, - Some(config), - None, - Some(WorkerType::Prefill), - "prefill", - None, - false, - Some(Box::new(FakeSharedCache { + picker_policy(|| Box::new(LoadOnlyPicker)), + 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(); 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, @@ -2857,390 +2246,766 @@ mod tests { )); } - async fn make_test_router_with_workers( - selector: impl dynamo_kv_router::selector::WorkerSelector - + Send - + Sync - + 'static, - shared_cache: Option>, + /// 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, - ) -> KvRouter< - impl dynamo_kv_router::selector::WorkerSelector + Send + Sync + 'static, - > { - 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, - use_kv_events: false, - router_track_active_blocks: false, - shared_cache_multiplier: 0.5, - skip_initial_worker_wait: true, - ..Default::default() - }; + block_size: u32, + policy: SelectionPolicySource, + shared_cache: Option>, + worker_role: Option, + 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) + } - KvRouter::new_with_worker_role( + /// [`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); + let router = KvRouter::new_with_worker_role( endpoint, client, rx, None, - 2, - selector, + block_size, + policy, Some(config), None, - None, - "decode", + worker_role, + metric_label, None, false, shared_cache, None, ) - .await - .unwrap() + .await?; + Ok((router, tx)) } - async fn make_test_router( - selector: impl dynamo_kv_router::selector::WorkerSelector - + Send - + Sync - + 'static, - shared_cache: Option>, - ) -> KvRouter< - impl dynamo_kv_router::selector::WorkerSelector + Send + Sync + 'static, - > { - 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 - } - - 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), - ); + /// 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")) } - 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 + 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); } - 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), + /// 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, + 2, + SelectionPolicySource::Registry, + None, + None, + "decode", + config, ) + .await + .unwrap() } + #[rstest::rstest] + #[case::single(1)] + #[case::concurrent(2)] #[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( - InspectingSelector { - expected_hits: None, - selected_worker: WorkerWithDpRank::new(7, 0), - }, + async fn session_prefix_tracking_survives_shared_core_selection(#[case] threads: u32) { + use dynamo_kv_router::protocols::{ExternalSequenceBlockHash, StorageTier}; + + let router = make_router( + "session-prefix-core", + HashMap::from([(7, ModelRuntimeConfig::default())]), + 2, + SelectionPolicySource::Registry, None, - workers, + Some(WorkerType::Decode), + "decode", + KvRouterConfig { + enable_session_prefix_index: true, + router_event_threads: threads, + router_temperature: 0.0, + skip_initial_worker_wait: true, + ..Default::default() + }, ) - .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), - ); + .await + .unwrap(); + let session_index = router.session_prefix_index.as_ref().unwrap(); + let worker = WorkerWithDpRank::new(7, 0); + let tokens = [11, 12, 13, 14]; + let hashes = compute_block_hash_for_seq(&tokens, 2, BlockHashOptions::default()); + let expected = RoutingDecisionHashes::from_local_hashes(hashes.clone()) + .sequence_hashes + .into_iter() + .map(ExternalSequenceBlockHash) + .collect::>(); + router + .indexer + .try_apply_event( + indexer::test_util::store_event( + 7, + 0, + 1, + &[], + &hashes.iter().map(|hash| hash.0).collect::>(), + StorageTier::Device, + ) + .with_session_id("seed"), + ) + .await + .unwrap(); + tokio::time::timeout(std::time::Duration::from_secs(2), async { + while session_index + .get_session_block_lineage("seed", worker, None) + .unwrap() + .is_empty() + { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); - 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), - ], - }, - )], - )) + for (session, admission, update_states, return_hashes) in [ + ( + "advisory", + FindBestMatchAdmission::WithoutAdmission, + false, + false, + ), + ("query", FindBestMatchAdmission::WithAdmission, false, false), + ("booked", FindBestMatchAdmission::WithAdmission, true, true), + ] { + let selected = router + .find_best_match_details_with_policy_class_inner( + Some(session), + &tokens, + None, + None, + update_states, + return_hashes, + None, + None, + 0.0, + 0, + None, + Some(dynamo_kv_router::SessionContext::new( + session.into(), + None, + None, + None, + )), + None, + None, + None, + None, + RoutingConstraints::default(), + admission, + ) + .await + .unwrap(); + let FindBestMatchOutcome::Routed { routing_hashes, .. } = selected.outcome else { + panic!("session request should route"); + }; + assert_eq!(routing_hashes.is_some(), return_hashes); + } + // The last match observes all earlier queued session updates. + tokio::time::timeout(std::time::Duration::from_secs(2), async { + while session_index + .get_session_block_lineage("booked", worker, None) + .unwrap() + .is_empty() + { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + for session in ["seed", "query", "booked"] { + assert_eq!( + session_index + .get_session_block_lineage(session, worker, None) + .unwrap(), + vec![expected.clone()], + ); + } + assert!( + session_index + .get_session_block_lineage("advisory", worker, None) + .unwrap() + .is_empty() ); + assert_eq!(session_index.session_count(), 3); } - #[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( - InspectingSelector { - expected_hits: None, - selected_worker: WorkerWithDpRank::new(7, 0), - }, + /// 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, - 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), + tokens, + None, + None, + false, + return_routing_hashes, + None, + None, + 0.0, 0, - Some(&candidates), - ); + None, + None, + None, + None, + None, + RoutingConstraints::default(), + ) + .await + } - assert_eq!(hint, None); + /// 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( + 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 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; + } + }) + .await + .expect("seeded prefix became visible"); } } + 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" + ); + + /// The selection procedure, run over the corpus with tracked bookings, must + /// match the frozen trace. Regenerate it with `SELECTION_GOLDEN_UPDATE=1` + /// only for an intended behavior change. #[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"), + async fn selection_matches_frozen_frontend_trace() { + 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}"); + let admitted = router + .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(), + FindBestMatchAdmission::WithAdmission, + ) + .await + .unwrap(); + 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 + .booking + .map(BookingHandle::commit) + .expect("tracked selection carries its booking handle"); + assert_eq!(booking.request_id, context_id); + } + + if std::env::var_os("SELECTION_GOLDEN_UPDATE").is_some() { + std::fs::write( + GOLDEN_PATH, + 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!( + trace, golden, + "selection drifted from the frozen frontend trace" ); - let router = make_test_router_with_workers( - InspectingSelector { - expected_hits: None, - selected_worker: 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, + } + + /// 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 = tracked_router("lease-manager").await; + let worker = WorkerWithDpRank::from_worker_id(1); + let scheduler = router.selection.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 hint = - router.transfer_hint_for_selection(WorkerWithDpRank::new(7, 0), 0, Some(&candidates)); + let before = manager.register_local(book("before").await, None); + drop(before); + assert!( + router.selection.scheduler().has_request("before"), + "no scheduler to release through yet" + ); - assert_eq!(hint, None); + manager.set_scheduler(router.selection.scheduler().booking_cleanup()); + let after = manager.register_local(book("after").await, None); + after.finish().await; + assert!(!router.selection.scheduler().has_request("after")); + assert!(router.selection.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] - 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"), + 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 = tracked_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 admitted = outcome; + 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(booking, Some(TokensWithHashes::new(prompt, 2))), ); - workers.insert( - 10, - transfer_hint_runtime_config_with_worker_type(Some("tcp://127.0.0.1:23283"), "decode"), + assert!(futures::poll!(&mut enroll).is_pending()); + assert!( + record.entered.load(Ordering::Acquire), + "routing update is in flight" ); - let router = make_test_router_with_workers( - InspectingSelector { - expected_hits: None, - selected_worker: WorkerWithDpRank::new(7, 0), - }, - None, + assert!(router.selection.scheduler().has_request("cancelled")); + drop(enroll); + tokio::time::timeout(std::time::Duration::from_secs(2), async { + while router.selection.scheduler().has_request("cancelled") { + tokio::task::yield_now().await; + } + }) + .await + .expect("cancelled enrollment released its booking"); + } + + /// Two default-config workers, deterministic scoring, untracked bookings. + async fn make_test_router( + policy: SelectionPolicySource, + shared_cache: Option>, + ) -> KvRouter { + let config = KvRouterConfig { + overlap_score_credit: 0.0, + router_temperature: 0.0, + use_kv_events: false, + router_track_active_blocks: false, + shared_cache_multiplier: 0.5, + skip_initial_worker_wait: true, + ..Default::default() + }; + let workers = HashMap::from([ + (0, ModelRuntimeConfig::default()), + (1, ModelRuntimeConfig::default()), + ]); + make_router( + "shared-cache-router", workers, + 2, + policy, + shared_cache, + None, + "decode", + config, ) - .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, - }; + .await + .unwrap() + } - 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), - ], + /// 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 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(); - let decode_hint = - router.transfer_hint_for_selection(WorkerWithDpRank::new(10, 0), 0, Some(&candidates)); + // 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!( - decode_hint, - Some(KvSourceLocationsPayload { - source_control_endpoint: "tcp://127.0.0.1:23282".to_string(), - block_hashes: vec![ - ExternalSequenceBlockHash(101), - ExternalSequenceBlockHash(102), - ExternalSequenceBlockHash(103), - ], - }) + 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] - 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( - InspectingSelector { - expected_hits: None, - selected_worker: 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, - )], - ))), + 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!( - 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), - ], - }) + 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( - InspectingSelector { - expected_hits: Some(2), - selected_worker: WorkerWithDpRank::from_worker_id(1), - }, - Some(Box::new(FakeSharedCache { + fixed_policy(Some(2), WorkerWithDpRank::from_worker_id(1)), + Some(Arc::new(FakeSharedCache { #[allow(clippy::single_range_in_vec_init)] hits: Some(dynamo_kv_router::protocols::SharedCacheHits::from_ranges( vec![0..2], @@ -3250,128 +3015,104 @@ 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] 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), - }, - Some(Box::new(FakeSharedCache { + fixed_policy(None, WorkerWithDpRank::from_worker_id(0)), + Some(Arc::new(FakeSharedCache { hits: None, should_error: true, })), ) .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] async fn test_find_best_match_maps_overload_to_resource_exhausted() { - let router = make_test_router(OverloadedSelector, 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 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( - err.as_ref(), - &[dynamo_runtime::error::ErrorType::ResourceExhausted], + error.as_ref(), + &[ErrorType::ResourceExhausted], &[] )); assert!( - err.to_string() + 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, _, _| { + 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(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(), + &[dynamo_runtime::error::ErrorType::Unavailable], + &[] + )); } #[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; 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), @@ -3396,51 +3137,11 @@ 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( - InspectingSelector { - expected_hits: None, - selected_worker: 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( - InspectingSelector { - expected_hits: None, - selected_worker: WorkerWithDpRank::from_worker_id(0), - }, - Some(Box::new(FakeSharedCache { + fixed_policy(None, WorkerWithDpRank::from_worker_id(0)), + 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/AGENTS.md b/lib/llm/src/kv_router/AGENTS.md index bbecd4aab726..03a3e5a7b871 100644 --- a/lib/llm/src/kv_router/AGENTS.md +++ b/lib/llm/src/kv_router/AGENTS.md @@ -3,8 +3,8 @@ ## 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`. -- `indexer/` and `route_lookup.rs` track which KV-cache blocks each worker holds and look for reusable blocks before selection. +- `embedded.rs` connects discovered workers and their current load to the scheduling code in `lib/kv-router`. +- `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 new file mode 100644 index 000000000000..e60b18eb4e19 --- /dev/null +++ b/lib/llm/src/kv_router/embedded.rs @@ -0,0 +1,728 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Embedded selection backend for the frontend `KvRouter`. +//! +//! 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; +use std::sync::{Arc, OnceLock}; +use std::time::Duration; + +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::protocols::{WorkerConfigLike, WorkerId, WorkerWithDpRank}; +use dynamo_kv_router::scheduling::queue::DEFAULT_MAX_BATCHED_TOKENS; +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, + SelectionOperation, SelectionOutcome, SelectionPartition, SelectionRun, SelectionScheduler, + SelectionService, SelectionServiceBuilder, WorkerCatalogRecord, WorkerCatalogSource, + WorkerRequest, WorkerSelectionPolicyRegistry, +}; +use dynamo_kv_router::{DEFAULT_ROUTING_GROUP, PrefillLoadEstimator, WorkerSelectionPolicyFactory}; +use tokio_util::sync::CancellationToken; + +use crate::discovery::RuntimeConfigWatch; +use crate::kv_router::metrics::{ + ActiveSequenceIngressMetrics, 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. +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, + 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, + /// Scheduler-owned load snapshots for the worker monitor's overload + /// detection. + 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, + /// Builds the partition's worker-selection policy (see + /// `SelectionPolicySource::resolve`). + pub policy_factory: WorkerSelectionPolicyFactory, +} + +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], scheduler: &SelectionScheduler) { + for (class_index, handles) in per_class.iter().enumerate() { + let Some(stats) = scheduler.class_queue_stats(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 = Duration::from_secs(60); + let mut recheck = tokio::time::interval_at(tokio::time::Instant::now() + period, period); + loop { + update_queue_metrics(&handles, partition.scheduler()); + 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" + ); + }) +} + +/// 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 + /// the router holds the partition. + service: Arc, + partition: SelectionPartition, + affinity: 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: OnceLock = 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`, 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.sender.publish(snapshot); + } + + fn publish_batch(&self, snapshots: Vec) { + 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, + ); + } +} + +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, crate::kv_router::sequence::ReplicaIngress)> { + let worker_type = args.worker_role.unwrap_or(WorkerType::Aggregated); + 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 + // disabled; only publishing is gated. + let (mut channels, replica_ingress) = 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")?; + channels.ingress_observer = Some(Arc::new( + 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 = + Some(Arc::new(move |_partition| { + 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, + WorkerSelectionPolicyRegistry::default(), + ) + .worker_selection_policy_factory(args.policy_factory) + .indexer_threads(1) + .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: args.shared_cache, + index: KvIndexSource::Owned(args.ingress), + }, + eligibility: HostEligibility { + lora_worker_filter: args.lora_worker_filter, + }, + telemetry: HostTelemetry { + scheduler_load: Some(Arc::new(SenderLoadSink { + sender: args.scheduler_load, + worker_type: args.metric_worker_type, + })), + }, + replication: HostReplication { + channels: replica_sync, + request_leases, + replica_worker_policy: ReplicaWorkerPolicy::LazyRegister, + }, + }) + .build() + .await + .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")?; + + // 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, + block_size: args.block_size, + is_eagle: args.is_eagle, + primed: false, + }; + let mut reconciler = CatalogReconciler::new(Arc::clone(service.core())).with_observer( + Arc::new(RegisteredGauge::new( + super::metrics::RouterWorkerStatusMetrics::from_component( + args.endpoint.component(), + ), + 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, + "KvRouter scheduling on embedded selection partition" + ); + Ok(( + Self { + service, + partition, + affinity: OnceLock::new(), + worker_type: args.metric_worker_type, + queue_metrics, + queue_metric_indices, + }, + replica_ingress, + )) + } + + 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, self.partition.scheduler()); + } + + pub(crate) fn affinity_coordinator( + &self, + ttl: Duration, + mode: crate::session_affinity::SessionAffinityMode, + ) -> Result { + let table = self.partition.session_affinity( + dynamo_kv_router::services::selection::affinity::SessionAffinityConfig::new(ttl) + .with_mode(mode), + )?; + Ok(self + .affinity + .get_or_init(|| crate::session_affinity::AffinityCoordinator::wrap(table)) + .clone()) + } + + 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<'_>) -> SelectionRun { + // 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, + }); + run + } + + pub(crate) fn worker_type(&self) -> &'static str { + self.worker_type + } + + /// The partition's scheduler, for bookings and load queries. + pub(crate) fn scheduler(&self) -> &SelectionScheduler { + self.partition.scheduler() + } +} + +/// The runtime-config watch as worker membership for the partition. +struct RuntimeDiscoverySource { + watch: RuntimeConfigWatch, + key: RoutingPartitionId, + block_size: u32, + is_eagle: bool, + 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, + /// 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) { + 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) { + 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); + } + } +} + +/// 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.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()); + } + } + } + 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), + // 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), + 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, + kv_event_source_mode: config.kv_event_source_mode.clone(), + ..WorkerRequest::default() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use dynamo_kv_router::services::selection::SchedulerLoadSink; + + #[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::kv_hints::KV_HINT_TRANSFER_CAPABILITY_KEY.to_string(), + serde_json::Value::Bool(true), + ); + config.runtime_data.insert( + 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::kv_hints::KV_HINT_TRANSFER_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()) + ]) + ); + } + + #[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) + ); + } + + /// 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 { + 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); + } + + #[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); + } +} 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..987144a41665 --- /dev/null +++ b/lib/llm/src/kv_router/indexer/ingress.rs @@ -0,0 +1,162 @@ +// 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::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_kv_router::{SessionPrefixIndexer, WorkerType}; +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 session_prefix_index: Option>, +} + +pub(crate) struct RuntimeIngress { + indexer: Indexer, + subscription: parking_lot::Mutex>, + _served: Option, +} + +impl RuntimeIngress { + pub(crate) async fn start(args: RuntimeIngressArgs<'_>) -> Result> { + 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, + session_prefix_index, + } = args; + let component = endpoint.component(); + let indexer = if cache_required { + super::build( + component, + kv_router_config, + block_size, + model_name, + cancellation_token.child_token(), + session_prefix_index, + ) + .await? + } else { + Indexer::None + }; + + let 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( + start_subscriber( + endpoint.clone(), + indexer.clone(), + membership_watch, + 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 = %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 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(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); + } + } +} + +#[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/lookup.rs b/lib/llm/src/kv_router/indexer/lookup.rs deleted file mode 100644 index 039837afaded..000000000000 --- a/lib/llm/src/kv_router/indexer/lookup.rs +++ /dev/null @@ -1,439 +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_primary_match_details_ref( - &self, - sequence: &[LocalBlockHash], - ) -> Result { - self.lookup_pipeline() - .find_primary_match_details(HashInput::Borrowed(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 5d0a94eab37a..98d4800ee492 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, ExternalSequenceBlockHash, KvCacheEventData, ResidencyProjection, - ResidencyRoutingSnapshot, RouterEvent, WorkerId, WorkerWithDpRank, + ApproximateRetentionConfig, KvIndexer, KvIndexerMetrics, LowerTierIndexers, + ThreadPoolIndexer, }, }; @@ -24,26 +23,23 @@ 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_kv_router::services::indexer::session_updates::SessionUpdateSender; use dynamo_runtime::component::Component; +pub(crate) use ingress::{RuntimeIngress, RuntimeIngressArgs}; use tokio_util::sync::CancellationToken; mod embedding_cache; -mod lookup; -mod recording; +mod ingress; mod recovery; pub mod remote; -#[doc(hidden)] -pub mod session_updates; -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}; -use self::session_updates::{SessionMutation, SessionUpdateSender}; -pub use self::side::SideIndexer; #[cfg(feature = "ckf-diagnostics")] pub(crate) use recovery::WorkerQueryHealthSnapshot; pub(crate) use recovery::{ @@ -56,39 +52,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, - session_updates: Option, - }, - Concurrent { - primary: Arc>, - lower_tier: LowerTierIndexers, - approx: Option, - primary_records_routing_decisions: bool, - session_updates: Option, - }, - Remote { - primary: Arc, - approx: Option, - primary_records_routing_decisions: bool, - }, - None, -} - #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum ResolvedApproximatePrimaryPolicy { Disabled, @@ -121,515 +84,197 @@ 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 => {} - } - } - - pub(crate) fn supports_overlap_refresh(&self) -> bool { - matches!(self, Self::KvIndexer { .. } | Self::Concurrent { .. }) +/// 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, + session_prefix_index: Option>, +) -> Result { + let approximate_policy = resolve_approximate_primary_policy(kv_router_config)?; + if approximate_policy == ResolvedApproximatePrimaryPolicy::Disabled { + return Ok(Indexer::None); } - 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 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 async fn new( - component: &Component, - kv_router_config: &KvRouterConfig, - block_size: u32, - model_name: Option<&str>, - cancellation_token: CancellationToken, - session_prefix_index: Option>, - ) -> 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(); - 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 = 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 { - let 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), - ), - ); - let session_updates = session_prefix_index - .map(|index| SessionUpdateSender::for_concurrent(index, Arc::clone(&primary))); - return Ok(Self::Concurrent { - primary, - lower_tier: LowerTierIndexers::new_with_metrics( - kv_router_config.router_event_threads as usize, - block_size, - Some(kv_indexer_metrics), - ), - approx: None, - primary_records_routing_decisions: true, - session_updates, - }); - } - - let primary = KvIndexer::new_with_approximate_retention( - cancellation_token.child_token(), - block_size, - kv_indexer_metrics.clone(), - Some(retention), - ); - let session_updates = session_prefix_index - .map(|index| SessionUpdateSender::for_legacy(index, primary.clone())); - return Ok(Self::KvIndexer { - primary, - lower_tier: LowerTierIndexers::new_with_metrics( - 1, - block_size, - Some(kv_indexer_metrics), - ), - approx: None, - primary_records_routing_decisions: true, - session_updates, - }); - } - - let approx = SideIndexer::new_predict_on_route( + 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(); + 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), + 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 { - let kv_indexer_metrics = KvIndexerMetrics::from_component(component); - let primary = Arc::new(ThreadPoolIndexer::new_with_metrics( - ConcurrentRadixTreeCompressed::new(), - kv_router_config.router_event_threads as usize, - block_size, - Some(kv_indexer_metrics.clone()), - )); + let 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), + ), + ); let session_updates = session_prefix_index .map(|index| SessionUpdateSender::for_concurrent(index, Arc::clone(&primary))); - return Ok(Self::Concurrent { + return Ok(Indexer::Concurrent { primary, 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, + approx: None, + primary_records_routing_decisions: true, session_updates, }); } - let kv_indexer_metrics = KvIndexerMetrics::from_component(component); - let primary = KvIndexer::new_with_pruning( + let primary = KvIndexer::new_with_approximate_retention( cancellation_token.child_token(), block_size, kv_indexer_metrics.clone(), - None, + Some(retention), ); let session_updates = session_prefix_index .map(|index| SessionUpdateSender::for_legacy(index, primary.clone())); - Ok(Self::KvIndexer { + return Ok(Indexer::Single { primary, lower_tier: LowerTierIndexers::new_with_metrics( 1, block_size, Some(kv_indexer_metrics), ), - approx, - primary_records_routing_decisions: false, + approx: None, + primary_records_routing_decisions: true, session_updates, - }) - } - - 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 session_update = if targets_primary { - match self { - Self::KvIndexer { - session_updates, .. - } - | Self::Concurrent { - session_updates, .. - } => session_updates.as_ref().and_then(|sender| { - SessionMutation::from_event(&event).map(|mutation| (sender.clone(), mutation)) - }), - Self::Remote { .. } | Self::None => None, - } - } else { - None - }; - 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 => {} - } - if let Some((sender, update)) = session_update { - sender.enqueue(update)?; - } - Ok(()) - } - - pub(crate) fn session_residency_version(&self, worker: WorkerWithDpRank) -> Option { - match self { - Self::KvIndexer { - session_updates, .. - } - | Self::Concurrent { - session_updates, .. - } => session_updates - .as_ref() - .map(|sender| sender.residency_version(worker)), - Self::Remote { .. } | Self::None => None, - } - } - - pub(crate) fn enqueue_session_match( - &self, - session_id: &str, - worker: WorkerWithDpRank, - matched_hash: ExternalSequenceBlockHash, - residency_version: u64, - ) -> Result<(), KvRouterError> { - match self { - Self::KvIndexer { - session_updates, .. - } - | Self::Concurrent { - session_updates, .. - } => match session_updates { - Some(sender) => sender.enqueue(SessionMutation::Matched { - worker, - session_id: session_id.to_owned(), - matched_hash, - residency_version, - }), - None => Ok(()), - }, - 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"); - } - } - - #[cfg(test)] - async fn flush_session_updates(&self) -> Result<(), KvRouterError> { - match self { - Self::KvIndexer { - session_updates, .. - } - | Self::Concurrent { - session_updates, .. - } => match session_updates { - Some(updates) => updates.flush().await, - None => Ok(()), - }, - Self::Remote { .. } | Self::None => Ok(()), - } - } - - /// 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 => {} - } - - let session_updates = match self { - Self::KvIndexer { - session_updates, .. - } - | Self::Concurrent { - session_updates, .. - } => session_updates.as_ref(), - Self::Remote { .. } | Self::None => None, - }; - if let Some(session_updates) = session_updates { - session_updates.enqueue(SessionMutation::Cleared { - worker: WorkerWithDpRank::new(worker_id, dp_rank), - })?; - session_updates.flush().await?; - } - Ok(()) + }); } - 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, - } - } + let approx = predict_on_route_side_indexer( + component, + kv_router_config, + block_size, + cancellation_token.child_token(), + ); - 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(()), - } - } + if kv_router_config.router_event_threads > 1 { + let kv_indexer_metrics = KvIndexerMetrics::from_component(component); + let primary = Arc::new(ThreadPoolIndexer::new_with_metrics( + ConcurrentRadixTreeCompressed::new(), + kv_router_config.router_event_threads as usize, + block_size, + Some(kv_indexer_metrics.clone()), + )); + let session_updates = session_prefix_index + .map(|index| SessionUpdateSender::for_concurrent(index, Arc::clone(&primary))); + return Ok(Indexer::Concurrent { + primary, + 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, + session_updates, + }); + } + + let kv_indexer_metrics = KvIndexerMetrics::from_component(component); + let primary = KvIndexer::new_with_pruning( + cancellation_token.child_token(), + block_size, + kv_indexer_metrics.clone(), + None, + ); + let session_updates = + session_prefix_index.map(|index| SessionUpdateSender::for_legacy(index, primary.clone())); + Ok(Indexer::Single { + primary, + lower_tier: LowerTierIndexers::new_with_metrics(1, block_size, Some(kv_indexer_metrics)), + approx, + primary_records_routing_decisions: false, + session_updates, + }) +} - 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)] pub(super) mod test_util { use dynamo_kv_router::protocols::{ - ExternalSequenceBlockHash, KvCacheEvent, KvCacheEventData, KvCacheRemoveData, - KvCacheStoreData, KvCacheStoredBlockData, LocalBlockHash, RouterEvent, StorageTier, + ExternalSequenceBlockHash, KvCacheEvent, KvCacheEventData, KvCacheStoreData, + KvCacheStoredBlockData, LocalBlockHash, RouterEvent, StorageTier, compute_seq_hash_for_block, }; @@ -680,1065 +325,4 @@ pub(super) mod test_util { storage_tier, ) } - - pub(crate) fn remove_event( - worker_id: u64, - dp_rank: u32, - event_id: u64, - block_hashes: Vec, - ) -> RouterEvent { - RouterEvent::new( - worker_id, - KvCacheEvent { - event_id, - data: KvCacheEventData::Removed(KvCacheRemoveData { block_hashes }), - dp_rank, - }, - ) - } - - pub(crate) fn clear_event(worker_id: u64, dp_rank: u32, event_id: u64) -> RouterEvent { - RouterEvent::new( - worker_id, - KvCacheEvent { - event_id, - data: KvCacheEventData::Cleared, - dp_rank, - }, - ) - } -} - -#[cfg(test)] -mod tests { - use std::sync::Arc; - use std::time::Duration; - - use tokio_util::sync::CancellationToken; - - use super::test_util::{clear_event, remove_event, store_event}; - use super::{Indexer, LowerTierIndexers, SessionUpdateSender}; - use dynamo_kv_router::{ - ConcurrentRadixTreeCompressed, SessionPrefixIndexer, ThreadPoolIndexer, - approx::PruneConfig, - indexer::{KvIndexer, KvIndexerInterface, KvIndexerMetrics, RoutingDecisionHashes}, - protocols::{ - BlockHashOptions, ExternalSequenceBlockHash, KvCacheEventData, 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, - session_updates: None, - } - } - - 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, - session_updates: None, - } - } - - 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, - session_updates: None, - } - } - - fn make_session_test_indexer() -> (Indexer, Arc) { - let session_prefix_index = Arc::new(SessionPrefixIndexer::new()); - let primary = KvIndexer::new( - CancellationToken::new(), - 4, - Arc::new(KvIndexerMetrics::new_unregistered()), - ); - let indexer = Indexer::KvIndexer { - primary: primary.clone(), - lower_tier: LowerTierIndexers::new(1, 4), - approx: None, - primary_records_routing_decisions: false, - session_updates: Some(SessionUpdateSender::for_legacy( - Arc::clone(&session_prefix_index), - primary, - )), - }; - (indexer, session_prefix_index) - } - - fn make_concurrent_session_test_indexer() -> (Indexer, Arc) { - let session_prefix_index = Arc::new(SessionPrefixIndexer::new()); - let primary = Arc::new(ThreadPoolIndexer::new( - ConcurrentRadixTreeCompressed::new(), - 2, - 4, - )); - let indexer = Indexer::Concurrent { - primary: Arc::clone(&primary), - lower_tier: LowerTierIndexers::new(2, 4), - approx: None, - primary_records_routing_decisions: false, - session_updates: Some(SessionUpdateSender::for_concurrent( - Arc::clone(&session_prefix_index), - primary, - )), - }; - (indexer, session_prefix_index) - } - - #[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()); - } - - #[tokio::test] - async fn only_attributed_stored_events_update_session_lineage() { - let (indexer, session_prefix_index) = make_session_test_indexer(); - - indexer - .apply_event( - store_event(7, 0, 1, &[], &[41], StorageTier::Device).with_session_id("session-1"), - ) - .await; - indexer - .apply_event(store_event(7, 0, 2, &[], &[51], StorageTier::HostPinned)) - .await; - indexer.flush_session_updates().await.unwrap(); - - assert_eq!( - session_prefix_index - .get_session_block_lineage("session-1", WorkerWithDpRank::new(7, 0), None) - .unwrap(), - vec![vec![ExternalSequenceBlockHash(41)]] - ); - assert_eq!(session_prefix_index.session_count(), 1); - assert_eq!(session_prefix_index.node_count(), 1); - assert!( - session_prefix_index - .get_node_from_hash(ExternalSequenceBlockHash(51)) - .is_none() - ); - } - - #[tokio::test] - async fn physical_events_update_worker_local_session_frontiers_in_order() { - let (indexer, session_prefix_index) = make_session_test_indexer(); - let first_worker = WorkerWithDpRank::new(7, 0); - let second_worker = WorkerWithDpRank::new(8, 0); - let first_store = store_event(7, 0, 1, &[], &[41, 42, 43], StorageTier::Device) - .with_session_id("session-1"); - let block_hashes = match &first_store.event.data { - KvCacheEventData::Stored(stored) => stored - .blocks - .iter() - .map(|block| block.block_hash) - .collect::>(), - _ => unreachable!(), - }; - let second_store = store_event(8, 0, 1, &[], &[41, 42, 43], StorageTier::Device) - .with_session_id("session-1"); - - indexer.apply_event(first_store).await; - indexer.apply_event(second_store).await; - indexer - .apply_event(remove_event(7, 0, 2, block_hashes[1..].to_vec())) - .await; - indexer.flush_session_updates().await.unwrap(); - - assert_eq!( - session_prefix_index - .get_session_block_lineage("session-1", first_worker, None) - .unwrap(), - vec![vec![block_hashes[0]]] - ); - assert_eq!( - session_prefix_index - .get_session_block_lineage("session-1", second_worker, None) - .unwrap(), - vec![block_hashes.clone()] - ); - - indexer.apply_event(clear_event(8, 0, 2)).await; - indexer.flush_session_updates().await.unwrap(); - assert!( - session_prefix_index - .get_session_block_lineage("session-1", second_worker, None) - .unwrap() - .is_empty() - ); - assert_eq!( - session_prefix_index.node_count(), - 3, - "physical removal preserves logical topology" - ); - } - - #[tokio::test] - async fn concurrent_physical_events_precede_session_updates() { - let (indexer, session_prefix_index) = make_concurrent_session_test_indexer(); - let worker = WorkerWithDpRank::new(7, 0); - let store = - store_event(7, 0, 1, &[], &[41, 42], StorageTier::Device).with_session_id("session-1"); - let block_hashes = match &store.event.data { - KvCacheEventData::Stored(stored) => stored - .blocks - .iter() - .map(|block| block.block_hash) - .collect::>(), - _ => unreachable!(), - }; - - indexer.apply_event(store).await; - indexer - .apply_event(remove_event(7, 0, 2, vec![block_hashes[1]])) - .await; - indexer.flush_session_updates().await.unwrap(); - - assert_eq!( - session_prefix_index - .get_session_block_lineage("session-1", worker, None) - .unwrap(), - vec![vec![block_hashes[0]]] - ); - } - - #[tokio::test] - async fn route_match_is_ordered_after_pending_remove_and_store() { - let (indexer, session_prefix_index) = make_session_test_indexer(); - let worker = WorkerWithDpRank::new(7, 0); - let initial_store = - store_event(7, 0, 1, &[], &[41], StorageTier::Device).with_session_id("session-old"); - let block_hash = match &initial_store.event.data { - KvCacheEventData::Stored(stored) => stored.blocks[0].block_hash, - _ => unreachable!(), - }; - - indexer.apply_event(initial_store).await; - indexer.flush_session_updates().await.unwrap(); - - indexer - .apply_event(remove_event(7, 0, 2, vec![block_hash])) - .await; - indexer - .apply_event( - store_event(7, 0, 3, &[], &[41], StorageTier::Device).with_session_id("session-a"), - ) - .await; - let residency_version = indexer.session_residency_version(worker).unwrap(); - indexer - .enqueue_session_match("session-b", worker, block_hash, residency_version) - .unwrap(); - indexer.flush_session_updates().await.unwrap(); - - for session in ["session-a", "session-b"] { - assert_eq!( - session_prefix_index - .get_session_block_lineage(session, worker, None) - .unwrap(), - vec![vec![block_hash]] - ); - } - } - - #[tokio::test] - async fn stale_route_match_does_not_undo_removal() { - let (indexer, session_prefix_index) = make_session_test_indexer(); - let worker = WorkerWithDpRank::new(7, 0); - let store = - store_event(7, 0, 1, &[], &[41], StorageTier::Device).with_session_id("session-b"); - let block_hash = match &store.event.data { - KvCacheEventData::Stored(stored) => stored.blocks[0].block_hash, - _ => unreachable!(), - }; - - indexer.apply_event(store).await; - indexer.flush_session_updates().await.unwrap(); - indexer - .apply_event(remove_event(7, 0, 2, vec![block_hash])) - .await; - // The physical removal is queued, but its ordered session mutation has - // not run yet. A match observed in this window must still be invalidated. - let stale_version = indexer.session_residency_version(worker).unwrap(); - indexer - .enqueue_session_match("session-b", worker, block_hash, stale_version) - .unwrap(); - indexer.flush_session_updates().await.unwrap(); - - assert!( - session_prefix_index - .get_session_block_lineage("session-b", worker, None) - .unwrap() - .is_empty() - ); - } - - 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 rank_reset_waits_for_session_frontier_clear() { - let (indexer, session_prefix_index) = make_session_test_indexer(); - let reset_rank = WorkerWithDpRank::new(7, 0); - let retained_rank = WorkerWithDpRank::new(7, 1); - let block_hash = ExternalSequenceBlockHash(41); - - for rank in [reset_rank, retained_rank] { - indexer - .apply_event( - store_event( - rank.worker_id, - rank.dp_rank, - 1, - &[], - &[block_hash.0], - StorageTier::Device, - ) - .with_session_id("session-1"), - ) - .await; - } - indexer.flush_session_updates().await.unwrap(); - - indexer - .reset_worker_dp_rank_and_wait(reset_rank.worker_id, reset_rank.dp_rank) - .await - .unwrap(); - - assert!( - session_prefix_index - .get_session_block_lineage("session-1", reset_rank, None) - .unwrap() - .is_empty() - ); - assert_eq!( - session_prefix_index - .get_session_block_lineage("session-1", retained_rank, None) - .unwrap(), - vec![vec![block_hash]] - ); - } - - #[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, - session_updates: None, - }; - 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, - session_updates: None, - }; - - 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 c8a74de889ba..b2aedaabeef7 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, @@ -2092,7 +2092,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/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/indexer/recovery/worker_query.rs b/lib/llm/src/kv_router/indexer/recovery/worker_query.rs index b76499a1d1cb..04422d628e0a 100644 --- a/lib/llm/src/kv_router/indexer/recovery/worker_query.rs +++ b/lib/llm/src/kv_router/indexer/recovery/worker_query.rs @@ -1449,7 +1449,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, @@ -1914,7 +1914,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 efe2e7e03c13..82a66bf9dbb9 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 find_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, @@ -198,8 +163,55 @@ 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 { + 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( + &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,7 +589,7 @@ mod tests { .unwrap(); let _ = side.flush().await; - let indexer = Indexer::KvIndexer { + let indexer = Indexer::Single { primary: KvIndexer::new( CancellationToken::new(), 4, @@ -630,7 +642,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, @@ -644,10 +656,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, @@ -655,9 +668,10 @@ mod tests { &[13], StorageTier::HostPinned, )) - .await; + .await + .unwrap(); - let Indexer::KvIndexer { + let Indexer::Single { primary, lower_tier, .. @@ -730,7 +744,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, @@ -743,10 +757,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, @@ -754,9 +769,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 - } - } - } -} diff --git a/lib/llm/src/kv_router/metrics.rs b/lib/llm/src/kv_router/metrics.rs index 4f6aa98baba6..02c584de6c62 100644 --- a/lib/llm/src/kv_router/metrics.rs +++ b/lib/llm/src/kv_router/metrics.rs @@ -296,6 +296,7 @@ pub(crate) struct ActiveSequenceZmqIngressMetrics { payload_decode_errors_total: IntCounter, identity_errors_total: IntCounter, forced_aborts_total: IntCounter, + dropped_events_total: IntCounter, } static ACTIVE_SEQUENCE_ZMQ_INGRESS_METRICS: OnceLock> = @@ -351,6 +352,10 @@ impl ActiveSequenceZmqIngressMetrics { "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() @@ -395,6 +400,103 @@ impl ActiveSequenceZmqIngressMetrics { pub(crate) fn record_forced_abort(&self) { self.forced_aborts_total.inc(); } + + pub(crate) fn record_dropped_event(&self) { + self.dropped_events_total.inc(); + } +} + +// --------------------------------------------------------------------------- +// 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, labels::WORKER_TYPE]; + 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() + } + + /// 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, worker_type]; + 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(); + } } // --------------------------------------------------------------------------- @@ -450,6 +552,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(); diff --git a/lib/llm/src/kv_router/prefill_router/activation.rs b/lib/llm/src/kv_router/prefill_router/activation.rs index 07f051f69239..8fbf22f7051d 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, WorkerSetTarget}, - 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, target: WorkerSetTarget, kv_cache_block_size: u32, kv_router_config: Option, - ) -> Result> { + ) -> Result { let target_id = target.id(); let endpoint = target.endpoint(); let endpoint_id = endpoint.id(); @@ -369,18 +358,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), @@ -392,8 +375,12 @@ where ) .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( @@ -408,11 +395,14 @@ where 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 @@ -424,11 +414,10 @@ where ) .await?; - Arc::new(RoutingHost::::new_builtin_with_coordinator( + Arc::new(RoutingHost::new_builtin_with_coordinator( push_router, load_context.clone(), affinity, - context.session_affinity_mode, )?) }; @@ -449,7 +438,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(); @@ -531,10 +520,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, @@ -639,11 +628,6 @@ where } } - /// 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)) } @@ -896,13 +880,13 @@ mod tests { admitted_ids, }) }; - let router = PrefillRouter::new_with_selector_factory( + let router = PrefillRouter::new_with_selection_policy( None, Arc::new(ModelManager::new()), RouterMode::RoundRobin, 16, None, - Arc::new(|config, _, _| DefaultWorkerSelector::new(Some(config.clone()), "prefill")), + SelectionPolicySource::Registry, None, None, SessionAffinityMode::Hard, @@ -958,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, @@ -967,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/admission.rs b/lib/llm/src/kv_router/prefill_router/admission.rs index 35cadd547250..f7045328f32b 100644 --- a/lib/llm/src/kv_router/prefill_router/admission.rs +++ b/lib/llm/src/kv_router/prefill_router/admission.rs @@ -7,26 +7,18 @@ use anyhow::Result; use futures::StreamExt; use tokio::sync::OwnedSemaphorePermit; -use dynamo_kv_router::selector::WorkerSelector; - use dynamo_runtime::{ pipeline::ManyOut, protocols::{annotated::Annotated, maybe_error::MaybeError}, }; use super::{PrefillCompletion, PrefillError, PrefillRouter, handoff::PrefillTask}; -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>, @@ -174,7 +166,6 @@ where #[cfg(test)] mod tests { - use dynamo_kv_router::selector::DefaultWorkerSelector; use futures::stream; use serde_json::json; @@ -217,11 +208,9 @@ mod tests { failure })); let response = ResponseStream::new(Box::pin(stream), Arc::new(Controller::default())); - let result = PrefillRouter::::consume_prefill_stream( - response, None, None, - ) - .await - .unwrap(); + let result = PrefillRouter::consume_prefill_stream(response, None, None) + .await + .unwrap(); let PrefillCompletion::Handoff { completion: Some(completion), .. @@ -273,13 +262,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(); @@ -295,7 +280,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, @@ -313,7 +298,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"), @@ -344,7 +329,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, @@ -370,7 +355,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, @@ -391,7 +376,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, @@ -416,7 +401,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..da162d9077b3 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() { @@ -75,7 +67,7 @@ where 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 86951edf06dd..e0a4977bad8f 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, }; use dynamo_runtime::{ error::{ErrorType, match_error_chain}, @@ -31,8 +27,7 @@ use futures::stream::{self, StreamExt}; use crate::{ discovery::{ModelManager, WorkerSetTarget, WorkerSetTargetId}, - 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}, @@ -164,17 +159,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, @@ -220,17 +204,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 @@ -257,27 +238,21 @@ where activation_task_state: Arc<()>, } -struct PrefillBinding -where - Sel: WorkerSelector + Send + 'static, -{ +struct PrefillBinding { target_id: WorkerSetTargetId, 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, @@ -291,19 +266,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(); @@ -311,15 +280,13 @@ where } #[async_trait] -impl +impl Operator< SingleIn, ManyOut>, SingleIn, ManyOut>, - > for PrefillRouter -where - Sel: WorkerSelector + Send + 'static, + > for PrefillRouter { async fn generate( &self, @@ -364,7 +331,7 @@ where .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, @@ -608,18 +575,12 @@ fn independent_prefill_context( Ok(prefill) } -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) @@ -929,7 +890,7 @@ mod tests { #[tokio::test] async fn conditional_disagg_rejects_unknown_decode_target_before_fallback() { - use crate::kv_router::KvRouter; + use crate::{kv_router::KvRouter, local_model::runtime_config::ModelRuntimeConfig}; use dynamo_runtime::{ DistributedRuntime, Runtime, distributed::DistributedConfig, pipeline::PushRouter, }; @@ -959,7 +920,7 @@ mod tests { workers, None, 16, - DefaultWorkerSelector::new(Some(config.clone()), "decode"), + SelectionPolicySource::Registry, Some(config), None, "decode", diff --git a/lib/llm/src/kv_router/prefill_router/query.rs b/lib/llm/src/kv_router/prefill_router/query.rs index b7468dd9e0ab..f225b37a4bc6 100644 --- a/lib/llm/src/kv_router/prefill_router/query.rs +++ b/lib/llm/src/kv_router/prefill_router/query.rs @@ -6,15 +6,10 @@ use std::collections::HashSet; use anyhow::Result; use dynamo_kv_router::{ protocols::{BlockExtraInfo, RoutingConstraints, WorkerId, WorkerWithDpRank}, - scheduling::{ - AdmissionAttempt, - queue::{SchedulerBookingCleanup, SchedulerBookingDescriptor}, - }, - selector::WorkerSelector, + scheduling::queue::BookingHandle, }; -use super::{PrefillError, PrefillLifecycleState, PrefillQueryOutcome, PrefillRouter}; -use crate::local_model::runtime_config::ModelRuntimeConfig; +use super::{PrefillError, PrefillLifecycleState, PrefillRouter}; /// A prefill booking that owns cleanup of the exact scheduler attempt it admitted. /// @@ -26,15 +21,7 @@ use crate::local_model::runtime_config::ModelRuntimeConfig; pub struct PrefillReservation { worker: WorkerWithDpRank, dp_rank: Option, - release: ReservationRelease, -} - -enum ReservationRelease { - Kv { - cleanup: SchedulerBookingCleanup, - booking: Option, - }, - None, + booking: Option, } impl PrefillReservation { @@ -48,29 +35,14 @@ 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 -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 @@ -111,11 +83,11 @@ where return Ok(PrefillReservation { worker: WorkerWithDpRank::new(worker_id, 0), dp_rank: None, - release: ReservationRelease::None, + booking: None, }); }; 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, @@ -129,27 +101,21 @@ where policy_class, None, None, + None, allowed_worker_ids, routing_constraints, ) .await?; - let (outcome, attempt) = admitted.into_parts(); + let (outcome, booking) = 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) = 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(SchedulerBookingDescriptor { - request_id: reservation_id.to_string(), - worker, - attempt_id, - }), - }, + booking: Some(booking), }) } crate::kv_router::FindBestMatchOutcome::QueueRejected { rejection } => { @@ -163,79 +129,6 @@ where } } } - - /// 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 }) - } - } - } - - 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)] @@ -246,8 +139,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, @@ -261,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, }, @@ -362,27 +256,15 @@ 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 used by `reserve_prefill_worker`'s non-KV arm. + 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( @@ -460,13 +342,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( @@ -495,12 +371,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 @@ -536,7 +407,7 @@ mod tests { workers_rx, None, 16, - DefaultWorkerSelector::new(Some(config.clone()), "prefill"), + crate::kv_router::SelectionPolicySource::Registry, Some(config), None, Some(WorkerType::Prefill), @@ -556,16 +427,13 @@ 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) } - 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)); @@ -676,7 +544,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() @@ -684,8 +554,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 @@ -721,7 +591,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 9eaa56b647cc..abaa928ec176 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, @@ -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; @@ -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/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.rs b/lib/llm/src/kv_router/routing_host.rs index df3e075c2baa..8fe847f8571d 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::{ @@ -42,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, }, }; @@ -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,14 +244,11 @@ where } /// An admitted KV route awaiting dispatch. -pub(crate) struct RoutePlan -where - Sel: WorkerSelector + Send + 'static, -{ - signals: RoutePlanSignals, +pub(crate) struct RoutePlan { + pub(crate) signals: RoutePlanSignals, selection: WorkerSelection, - cleanup: KvRequestCleanup, - affinity: Option, + 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. budget: CleanupBudget, @@ -277,7 +258,7 @@ where 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, } @@ -292,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)] @@ -311,14 +288,7 @@ impl RoutePlanSignals { } } -impl RoutePlan -where - Sel: WorkerSelector + Send + 'static, -{ - pub(crate) fn signals(&self) -> RoutePlanSignals { - self.signals - } - +impl RoutePlan { #[cfg(test)] pub(crate) fn cleanup_budget_remaining(&self) -> std::time::Duration { self.budget.remaining() @@ -334,38 +304,30 @@ 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 - .map(AffinityCoordinator::new) + .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( inner: PushRouter>, - kv_router: Arc>, + kv_router: Arc, load_context: Arc, session_affinity_ttl: Option, session_affinity_mode: SessionAffinityMode, ) -> Result { let affinity = session_affinity_ttl - .map(AffinityCoordinator::new) + .map(|ttl| kv_router.affinity_coordinator(ttl, session_affinity_mode)) .transpose()?; Ok(Self::new_with_load_context_and_coordinator( @@ -373,47 +335,36 @@ where kv_router, load_context, affinity, - session_affinity_mode, )) } pub(crate) fn new_with_coordinator( inner: PushRouter>, - kv_router: Arc>, + 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( inner: PushRouter>, - kv_router: Arc>, + 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, ) } 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, ) -> Self { // Eagerly register router request metrics (as zeros) so they are // scrapeable before any requests arrive. Both the frontend pipeline @@ -425,8 +376,11 @@ where inner, policy: RoutingPolicy::Kv(kv_router), request_metrics, + session_affinity_mode: affinity + .as_ref() + .map(AffinityCoordinator::mode) + .unwrap_or_default(), affinity, - session_affinity_mode, hosted_occupancy: None, lora: None, routing_context: load_context, @@ -438,35 +392,21 @@ where 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() { @@ -513,8 +453,11 @@ where inner, policy, request_metrics, + session_affinity_mode: affinity + .as_ref() + .map(AffinityCoordinator::mode) + .unwrap_or_default(), affinity, - session_affinity_mode, hosted_occupancy, lora: lora .zip(lora_selector) @@ -542,12 +485,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(_) @@ -576,6 +519,20 @@ where } } + /// 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) + } + /// Check request-supplied targets only at initial selection, not after a route preview /// or when dispatching: a worker disappearing later is a service failure. fn validate_explicit_worker( @@ -648,7 +605,7 @@ where phase: RequestPhase, staged_kv: StagedKv, budget: &CleanupBudget, - ) -> Result { + ) -> Result { match DispatchCancellation::for_request(phase, staged_kv) { DispatchCancellation::CancelWhenStopped => { affinity @@ -675,7 +632,7 @@ where is_query_only: bool, budget: &CleanupBudget, mut select: Select, - ) -> Result<(T, Option), Error> + ) -> Result<(T, Option), Error> where Select: FnMut(Option) -> SelectionFuture, SelectionFuture: Future>, @@ -705,7 +662,7 @@ where 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), @@ -726,7 +683,7 @@ where 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), @@ -754,10 +711,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. /// @@ -880,12 +835,7 @@ where return Err(error); } }; - match operation { - Some(operation) => { - 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 168adc439dad..5777630e18ae 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, @@ -347,7 +344,7 @@ where 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 }; @@ -360,7 +357,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, @@ -522,13 +519,7 @@ where } 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/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/kv.rs b/lib/llm/src/kv_router/routing_host/kv.rs index 1f6e0de6d68e..650172bac3e9 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, @@ -77,9 +74,7 @@ where is_query_only, affinity_target, None, - FindBestMatchAdmission::WithAdmission { - track_lifecycle: true, - }, + FindBestMatchAdmission::WithAdmission, budget, ) .await? @@ -92,7 +87,7 @@ where phase: RequestPhase, is_query_only: bool, budget: &CleanupBudget, - ) -> Result<(WorkerSelection, Option), Error> { + ) -> Result<(WorkerSelection, Option), Error> { self.validate_explicit_worker(request.content(), phase)?; self.select_with_session_affinity(request, phase, is_query_only, budget, |target| { self.select_request(request, phase, is_query_only, target, budget) @@ -165,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; @@ -184,7 +179,7 @@ where 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 { @@ -194,9 +189,7 @@ where false, target, Some(planned_worker), - FindBestMatchAdmission::WithAdmission { - track_lifecycle: true, - }, + FindBestMatchAdmission::WithAdmission, budget, ) .await? @@ -212,7 +205,7 @@ where Arc::clone(self.kv_router()), request.context().id().to_string(), selection.worker, - selection.attempt, + selection.booking.take(), ), selection, affinity, @@ -223,7 +216,7 @@ where pub(crate) async fn dispatch_kv_plan( &self, request: SingleIn, - plan: RoutePlan, + plan: RoutePlan, ) -> Result>, Error> { let RoutePlan { mut selection, @@ -255,12 +248,7 @@ where 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( @@ -305,7 +293,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 } @@ -314,9 +302,9 @@ where &self, request: &SingleIn, selection: &mut WorkerSelection, - cleanup: KvRequestCleanup, + cleanup: KvRequestCleanup, budget: &CleanupBudget, - ) -> Result, Error> { + ) -> Result { let phase = request .tracker .as_ref() @@ -332,9 +320,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(); @@ -351,7 +339,7 @@ where self.request_metrics.clone(), context_id.clone(), selected_worker, - selection.attempt, + selection.booking.take(), request, ), }; @@ -450,7 +438,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(); @@ -609,12 +597,9 @@ where 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/kv_selection.rs b/lib/llm/src/kv_router/routing_host/kv_selection.rs index 18911cbfcd34..fb76a5387b78 100644 --- a/lib/llm/src/kv_router/routing_host/kv_selection.rs +++ b/lib/llm/src/kv_router/routing_host/kv_selection.rs @@ -10,17 +10,12 @@ use dynamo_kv_router::{ protocols::{ BlockExtraInfo, RoutingConstraints, WorkerAffinityTarget, WorkerId, WorkerWithDpRank, }, - scheduling::{AdmissionAttempt, AdvisoryWorkerLoad, QueueRejection, RoutingEligibility}, - selector::WorkerSelector, + scheduling::{AdvisoryWorkerLoad, QueueRejection, RoutingEligibility, queue::BookingHandle}, }; use dynamo_runtime::{dynamo_nvtx_range, pipeline::Error}; use crate::{ - kv_router::{ - FindBestMatchAdmission, FindBestMatchInnerOutcome, FindBestMatchOutcome, - routing_host::RoutingHost, - }, - local_model::runtime_config::ModelRuntimeConfig, + kv_router::{FindBestMatchAdmission, FindBestMatchOutcome, routing_host::RoutingHost}, preprocessor::PreprocessedRequest, protocols::{ TokenIdType, @@ -31,7 +26,8 @@ use crate::{ pub(super) struct WorkerSelection { pub(super) worker: WorkerWithDpRank, - pub(super) attempt: AdmissionAttempt, + /// 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, @@ -41,6 +37,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), @@ -100,12 +98,9 @@ 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 + let admitted = self .kv_router() .find_best_match_details_with_policy_class_inner( Some(args.context_id), @@ -128,55 +123,29 @@ where 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, - attempt: admitted.attempt, - 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, - attempt: AdmissionAttempt::Untracked, - 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)) + } } } 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 c98bbd11411e..135eec5ab788 100644 --- a/lib/llm/src/kv_router/routing_host/request_guard.rs +++ b/lib/llm/src/kv_router/routing_host/request_guard.rs @@ -5,14 +5,9 @@ use std::{collections::HashMap, sync::Arc}; use crate::{ kv_router::{ - KvRouter, - indexer::ApproximateRequestLease, - metrics::RouterRequestMetrics, - prefill_router::BYPASS_REMOTE_PREFILL_ANNOTATION, - request_lease::RequestAttemptLease, - scheduler::{DefaultWorkerSelector, SchedulerBookingDescriptor}, + KvRouter, indexer::ApproximateRequestLease, metrics::RouterRequestMetrics, + prefill_router::BYPASS_REMOTE_PREFILL_ANNOTATION, request_lease::RequestAttemptLease, }, - local_model::runtime_config::ModelRuntimeConfig, lora::LoadEstimator, preprocessor::PreprocessedRequest, protocols::common::{ @@ -27,8 +22,7 @@ use dynamo_kv_router::{ BlockExtraInfo, BlockHashOptions, WorkerWithDpRank, compute_block_hash_for_seq, compute_next_seq_hash, }, - scheduling::AdmissionAttempt, - selector::WorkerSelector, + scheduling::queue::BookingHandle, }; use dynamo_runtime::{ error::DynamoError, @@ -393,50 +387,35 @@ 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, + booking: 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 = 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( + 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, @@ -458,11 +437,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 +448,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, @@ -567,13 +540,10 @@ 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 -where - Sel: WorkerSelector + Send + 'static, -{ - cleanup: RequestCleanup, +pub(super) struct RequestGuard { + cleanup: RequestCleanup, observability: RequestObservability, output_blocks: OutputBlockTracker, approximate_lru: Option, @@ -584,28 +554,25 @@ 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, - attempt: AdmissionAttempt, + booking: Option, request: &PreprocessedRequest, ) -> Self { Self::new_kv_with_cleanup( request_metrics, - KvRequestCleanup::new(chooser, context_id, worker, attempt), + KvRequestCleanup::new(chooser, context_id, worker, booking), request, ) } pub(super) fn new_kv_with_cleanup( request_metrics: Arc, - cleanup: KvRequestCleanup, + cleanup: KvRequestCleanup, request: &PreprocessedRequest, ) -> Self { let chooser = &cleanup.chooser; @@ -820,10 +787,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 1034d5228fa5..7cae3e4fdab4 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!() }; @@ -228,14 +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)).unwrap(); - let host = RoutingHost::::new_builtin_with_coordinator( + let (host, _) = builtin_host_with_affinity( inner, load_context, - Some(affinity), crate::session_affinity::SessionAffinityMode::Hard, - ) - .unwrap(); + ); let error = host .generate(affinity_request("direct-unbound", None)) @@ -275,14 +270,11 @@ 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, _) = builtin_host_with_affinity( inner, load_context, - Some(affinity.clone()), crate::session_affinity::SessionAffinityMode::Soft, - ) - .unwrap(); + ); let session_id = SessionAffinityId::new("direct-soft-bound"); bind_affinity_target(&host, &session_id, AffinityTarget::worker(worker_id)).await; @@ -361,26 +353,14 @@ 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, affinity) = builtin_host_with_affinity( inner, load_context, - Some(affinity.clone()), crate::session_affinity::SessionAffinityMode::Hard, - ) - .unwrap(); + ); let session_id = SessionAffinityId::new("local-inhibition"); - let AffinityAcquire::Initialize(initializer) = - affinity.acquire(&session_id, None).await.unwrap() - else { - panic!("new affinity session must initialize"); - }; - drop( - initializer - .commit(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)); @@ -441,11 +421,10 @@ 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, - crate::session_affinity::SessionAffinityMode::Hard, Some((filter, Arc::clone(&estimator))), ) .unwrap(); @@ -519,14 +498,11 @@ 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, affinity) = builtin_host_with_affinity( inner, load_context, - Some(affinity.clone()), crate::session_affinity::SessionAffinityMode::Hard, - ) - .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); @@ -578,7 +554,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,14 +602,11 @@ async fn builtin_direct_distinguishes_unknown_requests_from_stale_affinity() { ) .await .unwrap(); - let affinity = AffinityCoordinator::new(Duration::from_secs(10)).unwrap(); - let host = RoutingHost::::new_builtin_with_coordinator( + let (host, affinity) = builtin_host_with_affinity( inner, load_context, - Some(affinity.clone()), crate::session_affinity::SessionAffinityMode::Hard, - ) - .unwrap(); + ); let mut standalone = request(); standalone.routing_mut().backend_instance_id = Some(stale_worker); @@ -646,16 +619,7 @@ async fn builtin_direct_distinguishes_unknown_requests_from_stale_affinity() { assert!(dispatch.worker_ids.lock().unwrap().is_empty()); let session_id = SessionAffinityId::new("direct-affinity"); - let AffinityAcquire::Initialize(initializer) = - affinity.acquire(&session_id, None).await.unwrap() - else { - panic!("new affinity session must initialize"); - }; - drop( - initializer - .commit(AffinityTarget::worker(stale_worker)) - .unwrap(), - ); + bind_affinity_target(&host, &session_id, AffinityTarget::worker(stale_worker)).await; let error = host .generate(affinity_request("direct-affinity", None)) .await @@ -706,7 +670,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); @@ -775,7 +739,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); @@ -821,7 +785,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); @@ -859,7 +823,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); @@ -914,7 +878,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); @@ -960,7 +924,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); @@ -984,6 +948,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) = 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] @@ -1193,7 +1178,7 @@ async fn router_with_worker_configs( workers, None, 16, - DefaultWorkerSelector::new(Some(config.clone()), "decode"), + SelectionPolicySource::Registry, Some(config), None, "decode", @@ -1299,7 +1284,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", @@ -1420,46 +1405,60 @@ 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() { let (router, runtime) = router(None).await; let request = Context::new(request()); - 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"); - assert_eq!(plan.signals().worker.worker_id, 7); + 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(), 0, "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:?}" @@ -1484,12 +1483,8 @@ 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); - let loads = router - .kv_router() - .get_potential_loads(&[], None, None, None, None) - .await - .unwrap(); + assert_eq!(preview.signals.worker.worker_id, 7); + let loads = potential_loads(&router).await; assert!(loads.iter().all(|load| load.active_requests == 0)); assert_eq!(router.request_metrics.requests_started_total.get(), 0); @@ -1506,27 +1501,16 @@ 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); - let loads = router - .kv_router() - .get_potential_loads(&[], None, None, None, None) - .await - .unwrap(); + assert_eq!(plan.signals.worker, previewed_worker); + 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!( @@ -1559,7 +1543,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); @@ -1571,22 +1555,11 @@ async fn route_preview_does_not_acquire_session_affinity() { async fn planned_dispatch_transfers_the_reservation_to_request_cleanup() { let (router, runtime) = router(None).await; 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(), 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); @@ -1600,11 +1573,7 @@ async fn prefill_busy_probe_does_not_admit_a_request() { let request = Context::new(request()); 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(), 0); @@ -1638,7 +1607,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); @@ -1681,16 +1650,32 @@ async fn router_request_counters_follow_admission_and_completion_lifecycle() { drop(query_guard); assert_eq!(metrics.requests_started_total.get(), 0); - 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(), 1); assert_eq!(metrics.requests_total.get(), 0); + // The booking is held from admission until the guard finishes. + let loads = potential_loads(&router).await; + 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(), 1); assert_eq!(metrics.requests_total.get(), 0); + let loads = potential_loads(&router).await; + 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()); @@ -1739,13 +1724,8 @@ async fn router_request_counters_follow_admission_and_completion_lifecycle() { assert_eq!(metrics.requests_started_total.get(), 3); assert_eq!(metrics.requests_total.get(), 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(), 4); builtin_guard.abort().await; drop(builtin_guard); @@ -1775,12 +1755,14 @@ 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(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); @@ -1808,7 +1790,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() @@ -1818,7 +1800,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(); @@ -1851,7 +1837,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() @@ -1861,19 +1847,33 @@ 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); 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, target: AffinityTarget, ) { - let AffinityAcquire::Initialize(initializer) = router + let Hold::Initialize(initializer) = router .affinity .as_ref() .unwrap() @@ -1883,7 +1883,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] @@ -1974,7 +1978,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); @@ -2041,7 +2045,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() @@ -2051,7 +2055,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])); @@ -2351,7 +2359,7 @@ async fn two_worker_migration_harness( workers, None, 16, - DefaultWorkerSelector::new(Some(config.clone()), "decode"), + SelectionPolicySource::Registry, Some(config), None, None, @@ -2584,7 +2592,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) } @@ -2796,8 +2804,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"); }; @@ -2891,7 +2898,7 @@ async fn unknown_explicit_workers_are_rejected_before_builtin_dispatch() { ) .await .unwrap(); - let host = RoutingHost::::new_builtin(inner, load_context).unwrap(); + let host = RoutingHost::new_builtin(inner, load_context).unwrap(); let worker_id = live_worker.wrapping_add(1); for (field, phase) in [ ("backend_instance_id", RequestPhase::Aggregated), @@ -2970,10 +2977,8 @@ async fn unknown_explicit_workers_are_rejected_before_kv_admission() { } assert!(dispatch.worker_ids.lock().unwrap().is_empty()); assert!( - host.kv_router() - .get_potential_loads(&[], None, None, None, None) + potential_loads(&host) .await - .unwrap() .iter() .all(|load| load.active_requests == 0) ); diff --git a/lib/llm/src/kv_router/routing_load.rs b/lib/llm/src/kv_router/routing_load.rs index c6497aeb89e9..182b314af3d7 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; @@ -197,27 +196,18 @@ impl SchedulerLoadShared { pub struct SchedulerLoadSender { wake_tx: Option>, shared: Arc, - source: RouterLoadSource, cancellation_token: CancellationToken, } 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)), - 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)); } @@ -280,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) { @@ -301,7 +285,6 @@ fn scheduler_load_channel_with_capacity( SchedulerLoadSender { wake_tx: Some(wake_tx), shared: shared.clone(), - source, cancellation_token, }, SchedulerLoadReceiver { wake_rx, shared }, @@ -334,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, @@ -347,7 +330,7 @@ impl RoutingLoadContext { (scheduler_load, Some(monitor)) } else { ( - SchedulerLoadSender::disabled(source, cancellation_token.child_token()), + SchedulerLoadSender::disabled(cancellation_token.child_token()), None, ) }; @@ -395,18 +378,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(), @@ -415,22 +392,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, @@ -441,7 +412,7 @@ where &self.load_context } - pub fn router(&self) -> &Arc> { + pub fn router(&self) -> &Arc { &self.router } } @@ -462,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)]); @@ -489,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/scheduler.rs b/lib/llm/src/kv_router/scheduler.rs deleted file mode 100644 index 838108af5dcf..000000000000 --- a/lib/llm/src/kv_router/scheduler.rs +++ /dev/null @@ -1,720 +0,0 @@ -// 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}; -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; -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 977746b78c2a..51bd1de751a3 100644 --- a/lib/llm/src/kv_router/sequence.rs +++ b/lib/llm/src/kv_router/sequence.rs @@ -17,7 +17,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 +25,14 @@ 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; -#[cfg(test)] -use dynamo_kv_router::protocols::PrefillLoadHint; +use crate::kv_router::ACTIVE_SEQUENCES_SUBJECT; #[cfg(test)] use dynamo_runtime::transports::event_plane::MsgpackCodec; @@ -47,107 +40,7 @@ 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(Debug, Clone, Copy, PartialEq, Eq)] enum ActiveSequenceEventWireFormat { Singleton, @@ -253,68 +146,7 @@ impl ActiveSequenceEventPublisher { } } -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())) -} - -/// 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,54 +345,26 @@ impl SequenceSubscriber for RuntimeSequenceSubscriber { } } -/// Type alias for the runtime-wired multi-worker sequence tracker. -pub type ActiveSequencesMulti = ActiveSequencesMultiWorker; - -/// 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, +/// 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, - scheduler_load: SchedulerLoadSender, + publishes_outbound: bool, 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(); +) -> Result<( + dynamo_kv_router::services::selection::HostReplicaChannels, + ReplicaIngress, +)> { 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_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, + endpoint, ACTIVE_SEQUENCES_SUBJECT, transport_kind, ) @@ -585,110 +389,98 @@ pub(crate) async fn create_multi_worker_sequences_with_observer( } else { None }; - let worker_status_metrics = RouterWorkerStatusMetrics::from_component(endpoint.component()); - let publisher = RuntimeSequencePublisher { - event_sender, - scheduler_load, - worker_status_metrics, + let (inbound_tx, inbound_rx) = mpsc::channel(REPLICA_EVENT_CHANNEL_CAPACITY); + let ingress = ReplicaIngress { + endpoint: endpoint.clone(), + inbound_tx: inbound_tx.clone(), + cancellation_token, }; + Ok(( + dynamo_kv_router::services::selection::HostReplicaChannels { + outbound: event_sender, + inbound_tx, + inbound_rx, + process_id: router_id, + ingress_observer: None, + }, + ingress, + )) +} - 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, - ) - }; +/// 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, +} - 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()); +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" + ); + } } +} - // 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(()) +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"); } - Err(error) => Err(error), + None => break, } - }; - 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(), @@ -799,11 +591,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(); @@ -979,102 +766,14 @@ mod tests { } #[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(), + #[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 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-{}", @@ -1082,334 +781,105 @@ mod tests { ))? .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 (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?; - tokio::time::timeout(tokio::time::Duration::from_secs(5), async { + let received = tokio::time::timeout(std::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; + 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)) => {} } - tokio::time::sleep(tokio::time::Duration::from_millis(50)).await; } - Ok::<_, anyhow::Error>(()) }) - .await??; + .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(()) } + /// 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] - #[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(), + #[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 seq_manager_2 = create_multi_worker_sequences( - endpoint, - block_size, - workers_with_configs, - true, - 2, - scheduler_load_sender(), - CancellationToken::new(), + 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?; - - 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 - ); + .await + { + received_b.push(event.request_id); } - - 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!( + received_a.iter().all(|id| id == "endpoint-a-mark"), + "endpoint A received endpoint B sequence state: {received_a:?}" ); - assert_eq!( - tokens_phase1[&worker_2], 16, - "Worker 2 should have 16 active tokens (from request_2 added by seq_manager_2)" + assert!( + received_b.iter().all(|id| id == "endpoint-b-mark"), + "endpoint B received endpoint A sequence state: {received_b:?}" ); - 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 - ); - } - + drop((publisher_a, publisher_b)); + 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 index 1c458f7d54a1..77d22e088e1c 100644 --- a/lib/llm/src/kv_router/sequence/direct_zmq.rs +++ b/lib/llm/src/kv_router/sequence/direct_zmq.rs @@ -1,12 +1,11 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use std::{ffi::OsString, sync::Arc}; +use std::ffi::OsString; use anyhow::Result; -use dynamo_kv_router::{ - ActiveSequencesMultiWorker, SequencePublisher, - protocols::{ActiveSequenceEventBatch, MAX_REPLICA_BATCH_EVENTS}, +use dynamo_kv_router::protocols::{ + ActiveSequenceEvent, ActiveSequenceEventBatch, MAX_REPLICA_BATCH_EVENTS, }; use dynamo_runtime::{ component::Endpoint, @@ -14,6 +13,7 @@ use dynamo_runtime::{ discovery::EventTransportKind, transports::event_plane::{Codec, uses_direct_zmq}, }; +use tokio::sync::mpsc; use tokio::task::JoinHandle; use tokio_util::sync::CancellationToken; @@ -88,9 +88,12 @@ impl DirectZmqSequenceConfig { } } -pub(super) async fn start( +/// 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, - tracker: Arc>, + inbound_tx: mpsc::Sender, rcvhwm: i32, cancellation_token: CancellationToken, ) -> Result> { @@ -111,7 +114,11 @@ pub(super) async fn start( MAX_REPLICA_BATCH_EVENTS ); } - tracker.apply_replica_batch(batch.events); + for event in batch.events { + if inbound_tx.try_send(event).is_err() { + handler_metrics.record_dropped_event(); + } + } Ok(()) }; let observer = @@ -143,12 +150,9 @@ pub(super) async fn start( #[cfg(test)] mod tests { - use std::{collections::HashMap, time::Duration}; + use std::time::Duration; - use dynamo_kv_router::{ - NoopSequencePublisher, - protocols::{ActiveSequenceEvent, ActiveSequenceEventData, WorkerWithDpRank}, - }; + use dynamo_kv_router::protocols::{ActiveSequenceEventData, WorkerWithDpRank}; use dynamo_runtime::{ DistributedRuntime, Runtime, distributed::DistributedConfig, transports::event_plane::EventPublisher, @@ -232,28 +236,12 @@ mod tests { ))? .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 (inbound_tx, mut inbound_rx) = mpsc::channel(16); let cancel = CancellationToken::new(); - let supervisor = start( - endpoint.clone(), - tracker.clone(), - DEFAULT_RCVHWM, - cancel.clone(), - ) - .await?; + let supervisor = + start(endpoint.clone(), inbound_tx, DEFAULT_RCVHWM, cancel.clone()).await?; - for (index, request_id) in ["before-recreation", "after-recreation"] - .into_iter() - .enumerate() - { + for request_id in ["before-recreation", "after-recreation"] { let publisher = EventPublisher::for_endpoint_with_transport( &endpoint, ACTIVE_SEQUENCES_SUBJECT, @@ -268,10 +256,15 @@ mod tests { }) .await .unwrap(); - if tracker.active_request_counts()[&worker] == index + 1 { - break; + 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)) => {} } - tokio::time::sleep(Duration::from_millis(20)).await; } }) .await 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/lora/load_estimator.rs b/lib/llm/src/lora/load_estimator.rs index d956c507423c..a136edd65e41 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`] @@ -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}; @@ -231,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, @@ -242,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, @@ -454,31 +451,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, @@ -623,54 +595,6 @@ 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. - 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()) @@ -789,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 ab2266f5b3a1..bd485cfb870a 100644 --- a/lib/llm/src/session_affinity/coordinator.rs +++ b/lib/llm/src/session_affinity/coordinator.rs @@ -1,32 +1,32 @@ // 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, 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, AffinityLease, Hold, SessionAffinity, SessionAffinityConfig, +}; 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,169 +35,72 @@ 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; -enum AffinityEntry { - Initializing { - revision: u64, - notify: Arc, - }, - Bound { - target: AffinityTarget, - revision: u64, - version: AffinityVersion, - active_leases: usize, - idle_deadline: Instant, - }, +pub(crate) fn to_table(target: AffinityTarget) -> TableTarget { + TableTarget::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, - replica: OnceLock, - #[cfg(test)] - reaper_started: Arc, - #[cfg(test)] - waiter_observed: Arc, +pub(crate) fn from_table(target: TableTarget) -> AffinityTarget { + AffinityTarget::new(target.worker_id, target.dp_rank) } -impl Drop for AffinityCoordinatorInner { - 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, +struct Inner { + table: SessionAffinity, + replica: tokio::sync::OnceCell, } #[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, - ) + pub fn new(ttl: Duration, mode: SessionAffinityMode) -> Result { + Ok(Self::wrap( + SessionAffinity::with_config(SessionAffinityConfig::new(ttl).with_mode(mode)) + .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" - ))); + pub(crate) fn wrap(table: SessionAffinity) -> Self { + Self { + inner: Arc::new(Inner { + table, + replica: tokio::sync::OnceCell::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?; 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 table already has a replica sink installed" + )); + } + Ok(replica) + }) + .await?; + Ok(()) } /// Take a slot without cancelling on the request context. @@ -211,123 +114,37 @@ impl AffinityCoordinator { &self, 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)) + ) -> Result { + 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>, - ) -> Result { - self.validate_session_id(session_id)?; - let session_id = session_id.as_str().to_string(); - + context: &dyn AsyncEngineContext, + ) -> Result { + 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, - ))); - } - 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; - } - } - 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, - }); + match self + .inner + .table + .try_acquire(session_id.as_str(), requested) + .map_err(affinity_error)? + { + AcquireStep::Held(hold) => return Ok(hold), + AcquireStep::Wait(notified) => { + tokio::select! { + biased; + _ = context.stopped() => return Err(cancelled(context.id())), + _ = context.killed() => return Err(cancelled(context.id())), + _ = notified => {} } - }, + } } } } @@ -337,73 +154,67 @@ 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" - ); + self.inner + .table + .query_target(session_id.as_str(), requested_target.map(to_table)) + .map(|target| target.map(from_table)) + .map_err(affinity_error) + } - Ok(Some(*target)) + 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.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::with_config(SessionAffinityConfig { + max_entries, + max_session_id_bytes, + ..SessionAffinityConfig::new(Duration::from_secs(10)) + }) + .unwrap(), + ) } #[cfg(test)] @@ -412,8 +223,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 +238,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 +257,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,423 +268,30 @@ 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 { - 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), - } - } - - pub(crate) fn into_stream( - self, - dispatched_target: AffinityTarget, - stream: ManyOut, - mode: SessionAffinityMode, - ) -> Result, Error> { - match self { - Self::Initialize(initialization) => { - let lease = initialization.commit(dispatched_target)?; - lease.publish(dispatched_target); - Ok(lease.into_stream(stream)) - } - Self::Bound { target, mut lease } => { - 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); - } - return Ok(lease.into_stream(stream)); - } - if let Err(error) = validate_dispatch_target("session", target, dispatched_target) { - lease.invalidate(); - return Err(error); - } - lease.publish(target); - Ok(lease.into_stream(stream)) - } - } - } - - pub(crate) fn invalidate(self) { - if let Self::Bound { mut lease, .. } = self { - lease.invalidate(); - } - } -} - -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 +347,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 +367,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..b0dd066ec75f 100644 --- a/lib/llm/src/session_affinity/mod.rs +++ b/lib/llm/src/session_affinity/mod.rs @@ -4,51 +4,31 @@ 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}; -pub(crate) use coordinator::{AffinityAcquire, affinity_id, invalid_argument}; +#[cfg(test)] +pub(crate) use coordinator::to_table; 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(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_TTL_SECS, SessionAffinityMode, +}; pub type LlmResponse = crate::types::Annotated; 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 6afb7c027cf8..39539db3b3e4 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, @@ -346,14 +346,18 @@ 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()); let applier = ReplicaUpdateApplier { local_publisher_id: 7, discovered_instances, - coordinator: coordinator.downgrade_for_test(), + table: coordinator.table_for_test(), }; assert!(applier.apply( @@ -393,14 +397,31 @@ 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 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()), + 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(); + 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 cfb0fd02d44d..0a19f68c90a8 100644 --- a/lib/llm/src/session_affinity/tests.rs +++ b/lib/llm/src/session_affinity/tests.rs @@ -13,8 +13,9 @@ use futures::{StreamExt, stream}; use super::SessionAffinityMode::{Hard, Soft}; use super::{ - AffinityAcquire, AffinityCoordinator, AffinityTarget, LlmResponse, affinity_id, - coordinator::ReplicaApplyOutcome, explicit_target, + AffinityCoordinator, AffinityTarget, Hold, LlmResponse, affinity_id, + coordinator::{ReplicaApplyOutcome, tracked_stream}, + explicit_target, to_table, }; use crate::{ preprocessor::PreprocessedRequest, @@ -35,8 +36,19 @@ 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)).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 { @@ -62,8 +74,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() {} } @@ -147,7 +159,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"); }; @@ -156,16 +168,16 @@ 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(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, table_target(7, Some(0))); drop(first_lease); drop(second_lease); } @@ -174,7 +186,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"); }; @@ -184,12 +196,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(_) )); } @@ -197,7 +209,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"); }; @@ -226,14 +238,14 @@ 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() else { panic!("first request must initialize"); }; - drop(initializer.commit(target(7, None)).unwrap()); + drop(initializer.commit(table_target(7, None)).unwrap()); assert!( coordinator @@ -258,36 +270,34 @@ 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(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(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); 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"); }; - let lease = initializer.commit(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 = lease.into_stream(response_stream(1)); + let mut stream = tracked_stream(lease, response_stream(1)); assert!(stream.next().await.is_some()); drop(stream); @@ -297,14 +307,13 @@ 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"); }; - let lease = initializer.commit(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 = 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; @@ -313,23 +322,22 @@ 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(target(7, Some(0))).unwrap()); + drop(initializer.commit(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))); - let mut stream = lease.into_stream(cancelled_response_stream()); + assert_eq!(bound_target, table_target(7, Some(0))); + let mut stream = tracked_stream(lease, cancelled_response_stream()); assert!(stream.next().await.is_none()); assert_binding_expires_after_refreshed_ttl(&coordinator).await; @@ -339,8 +347,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()); @@ -351,14 +359,13 @@ 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"); }; - let lease = initializer.commit(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 = 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()); @@ -368,14 +375,13 @@ 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"); }; - let lease = initializer.commit(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 = 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; @@ -384,17 +390,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(target(7, Some(0))).unwrap()); + drop(initializer.commit(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; @@ -420,12 +423,11 @@ 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"); }; - drop(initializer.commit(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))) @@ -439,12 +441,11 @@ 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"); }; - drop(initializer.commit(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; @@ -461,7 +462,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( @@ -503,7 +504,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(_) )); } @@ -606,8 +607,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(); @@ -637,9 +638,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); @@ -656,14 +655,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!( @@ -675,13 +674,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() {} @@ -693,13 +692,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() {} @@ -716,15 +715,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); @@ -732,7 +731,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; @@ -745,8 +744,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(), @@ -757,7 +756,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)); @@ -765,8 +764,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!( @@ -843,8 +846,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(); 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 + } +]