From 0f3f581b52612db96f9981b3b95d25d0dbd487c9 Mon Sep 17 00:00:00 2001 From: jthomson04 Date: Thu, 10 Sep 2026 20:43:39 -0700 Subject: [PATCH 01/14] fix(kv-router): prevent lookup repair from restoring removed entries (#11891) Signed-off-by: jthomson04 <64760228+jthomson04@users.noreply.github.com> (cherry picked from commit 6b78d18fef0e402a55bc5237f29e159d33fc5d71) --- .../concurrent_radix_tree_compressed/mod.rs | 2 +- .../repair.rs | 2 +- .../concurrent_radix_tree_compressed/tests.rs | 63 +++++++++++++++++++ lib/kv-router/src/lookup_update.rs | 35 +++++++++-- 4 files changed, 96 insertions(+), 6 deletions(-) diff --git a/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/mod.rs b/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/mod.rs index 19b61e046781..7ef4efc217f9 100644 --- a/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/mod.rs +++ b/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/mod.rs @@ -19,7 +19,7 @@ use super::{ MatchDetails, PreBoundEventCounters, SyncIndexer, WorkerLookupStats, WorkerTask, }; use crate::cleanup::{CleanupGuard, CleanupState}; -use crate::lookup_update::update_arc_lookup_for_keys; +use crate::lookup_update::{update_arc_lookup_for_keys, update_existing_arc_lookup_for_keys}; use crate::protocols::*; mod children; diff --git a/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/repair.rs b/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/repair.rs index 1576c6e12c0a..18c654a70068 100644 --- a/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/repair.rs +++ b/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/repair.rs @@ -61,7 +61,7 @@ impl ConcurrentRadixTreeCompressed { let mut changed_entries_total = 0u64; for (&worker, worker_lookup) in lookup.iter_mut() { - let _changed_entries = update_arc_lookup_for_keys( + let _changed_entries = update_existing_arc_lookup_for_keys( worker_lookup, resolved.lookup_hashes_for_worker_repair(worker, hash, direction), resolved, diff --git a/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/tests.rs b/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/tests.rs index f327dfbefbe2..948674449b96 100644 --- a/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/tests.rs +++ b/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/tests.rs @@ -1379,3 +1379,66 @@ fn remove_after_split_and_children_clear_scrubs_lookup() { apply_direct(&index, &mut l2, remove_event(2, 103, 0, w2_leaf)); assert_eq!(worker_lookup_len(&l2, w2), Some(0)); } + +#[test] +fn successful_repair_does_not_restore_scrubbed_other_worker_entries() { + let index = ConcurrentRadixTreeCompressed::new(); + let scrubbed_worker = worker(1); + let repairing_worker = worker(2); + let mut shared_lookup = direct_lookup(); + let mut splitter_lookup = direct_lookup(); + + apply_direct( + &index, + &mut shared_lookup, + make_store_event(1, &[1, 2, 3, 4]), + ); + apply_direct( + &index, + &mut shared_lookup, + make_store_event(2, &[1, 2, 3, 4]), + ); + + // A split from another event thread leaves both local lookups pointing at + // the old prefix node for the [3, 4] suffix. + apply_direct( + &index, + &mut splitter_lookup, + make_store_event(3, &[1, 2, 9]), + ); + let suffix_hashes = remove_hashes_with_parent(&[1, 2], &[3, 4]); + + // A resolve-miss remove scrubs lookup state but cannot update coverage on + // the node it failed to find. Model that boundary directly while keeping + // the resolved live node's coverage intact. + let scrubbed_lookup = shared_lookup.get_mut(&scrubbed_worker).unwrap(); + for hash in &suffix_hashes { + scrubbed_lookup.remove(hash); + } + assert_eq!(worker_lookup_len(&shared_lookup, scrubbed_worker), Some(2)); + assert_direct_score(&index, &[1, 2, 3, 4], scrubbed_worker, 4); + + // A different worker's stale lookup successfully resolves the suffix and + // repairs the event thread's shared lookup map. + let resolved = index + .resolve_lookup( + &mut shared_lookup, + repairing_worker, + suffix_hashes[0], + LookupRepairDirection::TowardTail, + ) + .expect("repairing worker should resolve the split suffix"); + + let repairing_lookup = shared_lookup.get(&repairing_worker).unwrap(); + for hash in &suffix_hashes { + assert!(Arc::ptr_eq(repairing_lookup.get(hash).unwrap(), &resolved)); + } + + let scrubbed_lookup = shared_lookup.get(&scrubbed_worker).unwrap(); + for hash in suffix_hashes { + assert!( + !scrubbed_lookup.contains_key(&hash), + "repair for another worker restored a scrubbed lookup entry" + ); + } +} diff --git a/lib/kv-router/src/lookup_update.rs b/lib/kv-router/src/lookup_update.rs index 593e5146996c..c345403e2f5d 100644 --- a/lib/kv-router/src/lookup_update.rs +++ b/lib/kv-router/src/lookup_update.rs @@ -9,10 +9,9 @@ use rustc_hash::FxHashMap; /// Update an `Arc` lookup map for keys that should point at `node`. /// -/// Duplicate store/repair paths often revisit keys that already resolve to the -/// same node. Skipping those replacements avoids unnecessary hash table writes -/// and `Arc` refcount churn while still repairing stale entries. Returns the -/// number of entries that were inserted or changed. +/// Duplicate store paths often revisit keys that already resolve to the same +/// node. Skipping those replacements avoids unnecessary hash table writes and +/// `Arc` refcount churn. Returns the number of entries inserted or changed. pub(crate) fn update_arc_lookup_for_keys( lookup: &mut FxHashMap>, keys: impl IntoIterator, @@ -39,3 +38,31 @@ where changed } + +/// Update existing `Arc` lookup entries for keys that should point at `node`. +/// +/// Unlike [`update_arc_lookup_for_keys`], this does not insert missing keys. +/// Lookup repair uses absence as meaningful state: a remove can scrub an entry +/// before another worker on the same event thread repairs a stale lookup. +/// Returns the number of existing entries that changed. +pub(crate) fn update_existing_arc_lookup_for_keys( + lookup: &mut FxHashMap>, + keys: impl IntoIterator, + node: &Arc, +) -> usize +where + K: Eq + Hash, +{ + let mut changed = 0; + + for key in keys { + if let Some(entry) = lookup.get_mut(&key) + && !Arc::ptr_eq(entry, node) + { + *entry = Arc::clone(node); + changed += 1; + } + } + + changed +} From 6ac814ea0eb778aa8a5b99158f2534f94809306f Mon Sep 17 00:00:00 2001 From: jthomson04 Date: Wed, 16 Sep 2026 08:56:10 -0700 Subject: [PATCH 02/14] fix(kv-router): release unowned radix branches after eviction (#14878) Signed-off-by: jthomson04 (cherry picked from commit 8f9510112c208a2557e7d2fdc0432a3004348ee8) --- lib/kv-router/src/indexer/README.md | 22 ++- lib/kv-router/src/indexer/radix_tree.rs | 180 +++++++++++++++++++++--- 2 files changed, 182 insertions(+), 20 deletions(-) diff --git a/lib/kv-router/src/indexer/README.md b/lib/kv-router/src/indexer/README.md index 5f3e070b271a..e1a9e756d5e2 100644 --- a/lib/kv-router/src/indexer/README.md +++ b/lib/kv-router/src/indexer/README.md @@ -114,9 +114,25 @@ RadixBlock ├── state.worker_cutoffs: HashMap ├── state.full_edge_workers: HashSet ├── children: HashMap +├── parent: Weak> └── internal: bool ``` +Children and worker lookups hold strong references; parent links are weak to avoid +reference cycles. A child is keyed by its first local block hash. Splits preserve +that key for the prefix and update the parent links of children moved to the suffix. +Cleanup detaches an unowned node only if the parent's slot still points to it, so +an old detached node cannot remove a replacement. + +This single-threaded tree cleans up eagerly during mutations. The shared helper in +[`cleanup.rs`](../cleanup.rs) uses `Arc>` and periodic concurrent sweeps; +it does not support this tree's `Rc>` handles. + +The existing mid-chain removal gap can leave descendant lookup entries alive; see +the `TODO(CORRECTNESS)` in `apply_removed`. Such detached nodes can also retain weak +parent links. A weak link retains the parent's allocation after its value is +dropped. Complete descendant and back-link cleanup remains separate follow-up work. + ### Visual Representation ``` @@ -151,9 +167,9 @@ RadixBlock **remove_blocks(worker, block_hashes)**: 1. For each hash, find node via `lookup[worker][hash]` -2. Remove worker from node's `workers` set -3. If `workers` empty, clear children (cascading cleanup) -4. Remove from `lookup[worker]` +2. Truncate the worker's coverage at the removed block +3. Clear child edges if no worker covers the full node; detach the node if it has no owners +4. Remove newly uncovered hashes from `lookup[worker]` **find_matches(local_hashes, early_exit)**: 1. Start at root with all workers as candidates diff --git a/lib/kv-router/src/indexer/radix_tree.rs b/lib/kv-router/src/indexer/radix_tree.rs index 60ee6a3145c7..49d1dc496b3f 100644 --- a/lib/kv-router/src/indexer/radix_tree.rs +++ b/lib/kv-router/src/indexer/radix_tree.rs @@ -3,7 +3,11 @@ //! Single-threaded compressed radix tree for KV cache routing. -use std::{cell::RefCell, collections::VecDeque, rc::Rc}; +use std::{ + cell::RefCell, + collections::VecDeque, + rc::{Rc, Weak}, +}; use rustc_hash::{FxBuildHasher, FxHashMap, FxHashSet}; @@ -18,6 +22,7 @@ type WorkerLookup = FxHashMap; pub(crate) struct RadixBlock { state: NodeState, children: FxHashMap, + parent: Weak>, /// Once a node has children it is never eligible for leaf extension again. internal: bool, } @@ -27,21 +32,50 @@ impl RadixBlock { Self { state: NodeState::empty(), children: FxHashMap::default(), + parent: Weak::new(), internal: true, } } - fn for_blocks(blocks: &[KvCacheStoredBlockData], worker: WorkerWithDpRank) -> Self { + fn for_blocks( + blocks: &[KvCacheStoredBlockData], + worker: WorkerWithDpRank, + parent: &SharedRadixBlock, + ) -> Self { Self { state: NodeState::for_blocks(blocks, worker), children: FxHashMap::default(), + parent: Rc::downgrade(parent), internal: false, } } - fn clear_children_if_unreachable(&mut self) { - if self.state.full_edge_workers.is_empty() { - self.children.clear(); + fn prune_unreachable(node: &SharedRadixBlock) { + let (parent, key) = { + let mut node_ref = node.borrow_mut(); + debug_assert!(!node_ref.state.edge.is_empty(), "cannot prune the root"); + if node_ref.state.full_edge_workers.is_empty() { + node_ref.children.clear(); + } + if node_ref.state.has_any_workers() { + return; + } + ( + node_ref.parent.upgrade(), + node_ref.state.edge.first().map(|&(key, _)| key), + ) + }; + if let (Some(parent), Some(key)) = (parent, key) { + let mut parent = parent.borrow_mut(); + // A detached node can outlive its old edge through a worker lookup. + // Do not remove a replacement that now uses the same local hash. + if parent + .children + .get(&key) + .is_some_and(|child| Rc::ptr_eq(child, node)) + { + parent.children.remove(&key); + } } } } @@ -339,7 +373,9 @@ impl RadixTree { break; } - let child = Rc::new(RefCell::new(RadixBlock::for_blocks(remaining, worker))); + let child = Rc::new(RefCell::new(RadixBlock::for_blocks( + remaining, worker, &parent, + ))); { let mut parent_ref = parent.borrow_mut(); parent_ref.internal = true; @@ -397,7 +433,7 @@ impl RadixTree { self.update_lookup_for_blocks(worker, &remaining[..match_len], &child); let tail = &remaining[match_len..]; - let tail_node = Rc::new(RefCell::new(RadixBlock::for_blocks(tail, worker))); + let tail_node = Rc::new(RefCell::new(RadixBlock::for_blocks(tail, worker, &child))); { let mut child_ref = child.borrow_mut(); child_ref.internal = true; @@ -511,8 +547,12 @@ impl RadixTree { full_edge_workers: suffix_full, }, children: std::mem::take(&mut node_ref.children), + parent: Rc::downgrade(node), internal: node_ref.internal, })); + for child in suffix.borrow().children.values() { + child.borrow_mut().parent = Rc::downgrade(&suffix); + } node_ref.children.insert(suffix_first_local, suffix.clone()); node_ref.internal = true; suffix @@ -532,6 +572,8 @@ impl RadixTree { lookup.insert(hash, suffix.clone()); } } + drop(suffix_ref); + RadixBlock::prune_unreachable(&suffix); } fn update_lookup_for_blocks( @@ -638,12 +680,11 @@ impl RadixTree { // worker. Otherwise stale descendants can be reused as store parents, // reactivated by restoring only the removed block, or emitted by dumps // without a valid worker-specific parent. - let outcome = node_ref + node_ref .state - .remove_worker_at_pos(worker, min_pos, min_hash); - node_ref.clear_children_if_unreachable(); - outcome + .remove_worker_at_pos(worker, min_pos, min_hash) }; + RadixBlock::prune_unreachable(&node); for stale_hash in outcome.stale_hashes { lookup.remove(&stale_hash); eagerly_removed.insert(stale_hash); @@ -670,9 +711,8 @@ impl RadixTree { if !seen.insert(Rc::as_ptr(&node)) { continue; } - let mut node_ref = node.borrow_mut(); - node_ref.state.drop_worker(worker); - node_ref.clear_children_if_unreachable(); + node.borrow_mut().state.drop_worker(worker); + RadixBlock::prune_unreachable(&node); } if keep_worker { self.lookup.insert(worker_key, FxHashMap::default()); @@ -694,9 +734,8 @@ impl RadixTree { if !seen.insert(Rc::as_ptr(&node)) { continue; } - let mut node_ref = node.borrow_mut(); - node_ref.state.drop_worker(worker); - node_ref.clear_children_if_unreachable(); + node.borrow_mut().state.drop_worker(worker); + RadixBlock::prune_unreachable(&node); } } @@ -833,6 +872,113 @@ mod tests { create_remove_event, create_store_event, make_store_event, snapshot_events, }; + #[test] + fn cache_churn_releases_evicted_branches() { + let mut tree = RadixTree::new(); + tree.apply_event(create_store_event(1, 0, vec![1, 2, 3], None)) + .unwrap(); + for hash in [4, 5, 6, 7, 8] { + for worker in [1, 2] { + tree.apply_event(create_store_event(worker, hash, vec![1, 2, hash], None)) + .unwrap(); + } + let branch = Rc::downgrade( + &tree.lookup[&WorkerWithDpRank::new(1, 0)][&ExternalSequenceBlockHash(hash * 100)], + ); + if hash == 4 { + // Splitting the parent moves its existing children to the suffix. + tree.apply_event(create_store_event(1, hash, vec![1, 99], None)) + .unwrap(); + } + tree.apply_event(create_remove_event(1, hash + 1, vec![hash])) + .unwrap(); + assert!( + branch.upgrade().is_some(), + "another worker still owns the branch" + ); + match hash % 4 { + 0 => tree + .apply_event(create_remove_event(2, hash + 1, vec![hash])) + .unwrap(), + 1 => tree.remove_worker(2), + 2 => tree.remove_worker_dp_rank(2, 0), + _ => tree.clear_all_blocks(2), + } + assert!( + branch.upgrade().is_none(), + "evicted branch is still retained" + ); + assert_eq!(tree.edge_lengths_for_test(), vec![1, 1, 1, 1]); + assert_eq!(tree.current_size(), if hash % 4 == 0 { 6 } else { 4 }); + } + } + + #[test] + fn removing_detached_node_preserves_replacement() { + let mut tree = RadixTree::new(); + for blocks in [vec![1, 2, 3, 4], vec![1, 2, 3, 9]] { + tree.apply_event(create_store_event(2, 0, blocks, None)) + .unwrap(); + } + tree.apply_event(create_store_event(1, 0, vec![1, 2, 3], None)) + .unwrap(); + let old_branch = Rc::downgrade( + &tree.lookup[&WorkerWithDpRank::new(2, 0)][&ExternalSequenceBlockHash(400)], + ); + for worker in [1, 2] { + tree.apply_event(create_remove_event(worker, 1, vec![3])) + .unwrap(); + } + // The existing descendant-invalidation gap leaves worker 2's old node + // in its lookup after the parent has disconnected it. + assert!(old_branch.upgrade().is_some()); + tree.apply_event(create_store_event( + 1, + 2, + vec![3], + Some(ExternalSequenceBlockHash(200)), + )) + .unwrap(); + tree.apply_event(create_store_event( + 1, + 3, + vec![4], + Some(ExternalSequenceBlockHash(300)), + )) + .unwrap(); + tree.apply_event(create_remove_event(2, 3, vec![4])) + .unwrap(); + + assert!(old_branch.upgrade().is_none()); + let scores = tree.find_matches((1..=4).map(LocalBlockHash).collect(), false); + assert_eq!(scores.scores.get(&WorkerWithDpRank::new(1, 0)), Some(&4)); + } + + #[test] + fn split_releases_unowned_suffix() { + let mut tree = RadixTree::new(); + for worker in [1, 2] { + tree.apply_event(create_store_event(worker, 0, vec![1, 2, 3], None)) + .unwrap(); + } + tree.apply_event(create_remove_event(1, 1, vec![2])) + .unwrap(); + tree.apply_event(create_remove_event(2, 1, vec![1])) + .unwrap(); + assert_eq!(tree.edge_lengths_for_test(), vec![3]); + + tree.apply_event(create_store_event( + 1, + 2, + vec![9], + Some(ExternalSequenceBlockHash(100)), + )) + .unwrap(); + + assert_eq!(tree.edge_lengths_for_test(), vec![1, 1]); + assert_eq!(tree.current_size(), 2); + } + #[test] fn rejects_self_referencing_store() { let mut tree = RadixTree::new(); From b5534409d7ce12418affed10af5e3eaa6c3e3a9d Mon Sep 17 00:00:00 2001 From: Yan Ru Pei Date: Fri, 11 Sep 2026 12:39:22 -0700 Subject: [PATCH 03/14] fix(kv-router): make standalone indexer cleanup safe (#14652) Signed-off-by: PeaBrane (cherry picked from commit 6c66d4b2a6345edda51f17f7c150862fee37385d) --- .../router/standalone-indexer.md | 27 ++ .../router/standalone-selection.md | 20 ++ lib/kv-router/src/indexer/lower_tier.rs | 130 ++++----- .../src/services/indexer/listener.rs | 113 +++++++- .../src/services/indexer/registry.rs | 262 ++++++++++++++---- .../src/services/selection/core/mod.rs | 86 ++++-- 6 files changed, 481 insertions(+), 157 deletions(-) diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-indexer.md b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-indexer.md index 3483286a242c..b0624e653124 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 @@ -201,6 +201,10 @@ Returns: Register a ZMQ endpoint for an instance. Each call creates or reuses the indexer for the given `(model_name, routing_group)` pair. Registration is non-blocking: if the worker is not up yet, the listener is accepted in `pending` state and transitions to `active` once the initial ZMQ connection succeeds. +Repeating `/register` for an existing `instance_id` and `dp_rank` returns an error; +it does not replace the listener or reset its sequence watermark. For worker restarts, +see [Worker Restarts](#worker-restarts). + ```bash # Single model, default routing group curl -X POST http://localhost:8090/register \ @@ -490,6 +494,29 @@ recovers history still retained by the engine. Without a replay endpoint, or whe range is no longer retained, restart with a healthy indexer in `--peers` to recover its startup snapshot, or perform another full resynchronization before relying on the rebuilt worker state. +## Worker Restarts + +A ZMQ reconnect does not identify a new worker or cache-event publisher lifetime. +If the publisher restarts with its sequence reset, an existing listener retains its +watermark and discards batches at or below it. Previously indexed ownership can also +remain stale. + +On a worker or publisher restart, call `/unregister` and wait for the response before +calling `/register` with the current worker metadata. For a whole-worker restart, +omit `dp_rank` from `/unregister` to remove all ranks. Apply this lifecycle change to +every indexer replica; peer registration does not propagate worker lifecycle changes. +Serialize lifecycle operations for the same worker. + +This sequence requests removal of old cache ownership and starts a fresh sequence +watermark. It does not certify a complete cache view or provide an atomic boundary +against in-flight old events. Registration success and listener status `active` do +not establish that startup events were received or replayed. Recover missing history +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. + ## Limitations - **Standalone mode is ZMQ only**: Workers must publish KV events via ZMQ PUB sockets. 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 4a740248663e..9bcd9b9da3fe 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 @@ -127,6 +127,26 @@ both default to `"default"` when omitted. `GET /health` is process liveness. `GET /ready` returns `200` only after at least one worker is schedulable, otherwise `503` with lifecycle details. +### Worker Restarts + +When a worker or its cache-event publisher restarts, update its registration on every +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. + +This differs from the standalone indexer's `/register`, which rejects duplicate +registrations and requires `/unregister` first. See +[Indexer Worker Restarts](standalone-indexer.md#worker-restarts) for cleanup and +replay limits. Neither a successful catalog update nor `/ready` certifies a complete +cache view: missed startup events must still be recovered, and expired replay +history requires another source of complete state. + ## Selection API ### `POST /select` diff --git a/lib/kv-router/src/indexer/lower_tier.rs b/lib/kv-router/src/indexer/lower_tier.rs index 8445072149f6..77fe9d7948dd 100644 --- a/lib/kv-router/src/indexer/lower_tier.rs +++ b/lib/kv-router/src/indexer/lower_tier.rs @@ -622,7 +622,7 @@ impl LowerTierIndexer { let indexed_owner = IndexedResidencyOwner::from_exact(owner); let remove_worker_entry = { let Some(owner_state) = worker_blocks.get_mut(&indexed_owner) else { - return Err(KvCacheEventError::BlockNotFound); + return Ok(()); }; if owner_state.owner != owner { return Err(KvCacheEventError::UnsupportedResidencyDomain); @@ -631,7 +631,7 @@ impl LowerTierIndexer { for block_hash in block_hashes { let Some(key) = worker_map.remove(block_hash) else { - return Err(KvCacheEventError::BlockNotFound); + continue; }; self.remove_owner_from_edge(key, indexed_owner); @@ -1488,11 +1488,6 @@ mod tests { self.index.remove_worker(&mut self.worker_blocks, worker_id); } - fn remove_worker_dp_rank(&mut self, worker_id: u64, dp_rank: u32) { - self.index - .remove_worker_dp_rank(&mut self.worker_blocks, worker_id, dp_rank); - } - fn root_workers(&self, local_hash: LocalBlockHash) -> Vec { self.index .root_workers(local_hash, &projection(WorkerWithDpRank::new(7, 0))) @@ -2014,6 +2009,55 @@ mod tests { ); } + #[test] + fn removal_batches_are_idempotent_and_preserve_other_owners() { + let mut index = TestLowerTierIndex::new(); + let removed = WorkerWithDpRank::new(1, 0); + let other_rank = WorkerWithDpRank::new(1, 1); + let other_worker = WorkerWithDpRank::new(2, 0); + for worker in [removed, other_rank, other_worker] { + index + .apply_event(store_event( + worker.worker_id, + worker.dp_rank, + 0, + None, + &[11, 12], + &[101, 102], + )) + .unwrap(); + } + + let batch = remove_event( + 1, + 1, + 0, + vec![ + ExternalSequenceBlockHash(999), + ExternalSequenceBlockHash(101), + ExternalSequenceBlockHash(101), + ExternalSequenceBlockHash(102), + ], + ); + index.apply_event(batch.clone()).unwrap(); + index.apply_event(batch).unwrap(); + index + .apply_event(remove_event(99, 2, 0, vec![ExternalSequenceBlockHash(101)])) + .unwrap(); + + let continuations = [removed, other_rank, other_worker] + .into_iter() + .map(|worker| (worker, LowerTierContinuation::from_root(0))) + .collect::>(); + let hits = index.query_contiguous_hits(&local_hashes(&[11, 12]), &continuations); + assert_eq!(hits.get(&removed), Some(&0)); + assert_eq!(hits.get(&other_rank), Some(&2)); + assert_eq!(hits.get(&other_worker), Some(&2)); + assert!(index.dump_events().iter().all(|event| { + event.worker_id != removed.worker_id || event.event.dp_rank != removed.dp_rank + })); + } + #[test] fn removing_one_owner_preserves_shared_edge_for_other_workers() { let mut index = TestLowerTierIndex::new(); @@ -2095,21 +2139,21 @@ mod tests { } #[test] - fn start_pos_past_end_returns_zero() { + fn exhausted_sequence_returns_zero() { let mut index = TestLowerTierIndex::new(); index .apply_event(store_event(23, 0, 0, Some(1100), &[91], &[901])) .unwrap(); - let query = local_hashes(&[91]); - let mut continuations = FxHashMap::default(); - continuations.insert( - WorkerWithDpRank::new(23, 0), - LowerTierContinuation::new(1, ExternalSequenceBlockHash(1100)), - ); - - let hits = index.query_contiguous_hits(&query, &continuations); - assert_eq!(hits.get(&WorkerWithDpRank::new(23, 0)), Some(&0)); + let worker = WorkerWithDpRank::new(23, 0); + for query in [local_hashes(&[91]), Vec::new()] { + let continuation = + LowerTierContinuation::new(query.len(), ExternalSequenceBlockHash(1100)); + let continuations = FxHashMap::from_iter([(worker, continuation)]); + let details = index.query_match_details(&query, &continuations); + assert_eq!(details.hits.get(&worker), Some(&0)); + assert_eq!(details.next_continuations.get(&worker), Some(&continuation)); + } } #[test] @@ -2201,9 +2245,8 @@ mod tests { .apply_event(store_event(41, 0, 0, Some(3000), &[1], &[301])) .unwrap(); index - .apply_event(store_event(41, 1, 1, Some(4000), &[2], &[401])) + .apply_event(store_event(41, 1, 1, Some(4000), &[1], &[401])) .unwrap(); - index.remove_worker(41); let mut continuations = FxHashMap::default(); continuations.insert( @@ -2215,37 +2258,16 @@ mod tests { LowerTierContinuation::new(0, ExternalSequenceBlockHash(4000)), ); + let before = index.query_contiguous_hits(&local_hashes(&[1]), &continuations); + assert_eq!(before.get(&WorkerWithDpRank::new(41, 0)), Some(&1)); + assert_eq!(before.get(&WorkerWithDpRank::new(41, 1)), Some(&1)); + index.remove_worker(41); + let hits = index.query_contiguous_hits(&local_hashes(&[1]), &continuations); assert_eq!(hits.get(&WorkerWithDpRank::new(41, 0)), Some(&0)); assert_eq!(hits.get(&WorkerWithDpRank::new(41, 1)), Some(&0)); } - #[test] - fn remove_worker_dp_rank_keeps_other_ranks() { - let mut index = TestLowerTierIndex::new(); - index - .apply_event(store_event(43, 0, 0, Some(5000), &[1], &[501])) - .unwrap(); - index - .apply_event(store_event(43, 1, 1, Some(6000), &[2], &[601])) - .unwrap(); - index.remove_worker_dp_rank(43, 0); - - let mut continuations = FxHashMap::default(); - continuations.insert( - WorkerWithDpRank::new(43, 0), - LowerTierContinuation::new(0, ExternalSequenceBlockHash(5000)), - ); - continuations.insert( - WorkerWithDpRank::new(43, 1), - LowerTierContinuation::new(0, ExternalSequenceBlockHash(6000)), - ); - - let hits = index.query_contiguous_hits(&local_hashes(&[2]), &continuations); - assert_eq!(hits.get(&WorkerWithDpRank::new(43, 0)), Some(&0)); - assert_eq!(hits.get(&WorkerWithDpRank::new(43, 1)), Some(&1)); - } - #[test] fn removing_parent_block_keeps_child_continuation_edge() { let mut index = TestLowerTierIndex::new(); @@ -2626,24 +2648,6 @@ mod tests { assert_eq!(details.hits.get(&WorkerWithDpRank::new(103, 0)), Some(&0),); } - /// Empty sequence — every worker should get 0 hits. - #[test] - fn empty_sequence_returns_zero_hits() { - let mut index = TestLowerTierIndex::new(); - index - .apply_event(store_event(111, 0, 0, None, &[1], &[101])) - .unwrap(); - - let mut continuations = FxHashMap::default(); - continuations.insert( - WorkerWithDpRank::new(111, 0), - LowerTierContinuation::from_root(0), - ); - - let details = index.query_match_details(&local_hashes(&[]), &continuations); - assert_eq!(details.hits.get(&WorkerWithDpRank::new(111, 0)), Some(&0)); - } - // --- dump_events tests --- /// Helper: replay dumped events into a fresh indexer and return it. diff --git a/lib/kv-router/src/services/indexer/listener.rs b/lib/kv-router/src/services/indexer/listener.rs index c844c4fbbd80..d2cd5d05996f 100644 --- a/lib/kv-router/src/services/indexer/listener.rs +++ b/lib/kv-router/src/services/indexer/listener.rs @@ -229,25 +229,24 @@ impl ListenerLoop { } } }; - if msg.len() != 4 { - tracing::warn!( - worker_id, - dp_rank, - "Unexpected replay frame count: {}", - msg.len() - ); - break; - } - - // vLLM replay responses include the DEALER identity delimiter plus the - // original PUB topic, so the payload shape is: - // [empty delimiter, topic, sequence number, encoded event batch]. - let payload = msg.get(3).expect("frame count checked above"); + // DEALER strips the ROUTER identity. vLLM includes the PUB topic; + // SGLang sends only the delimiter, sequence number, and payload. + let (seq_bytes, payload) = match msg.as_slice() { + [_, seq, payload] | [_, _, seq, payload] => (seq, payload), + _ => { + tracing::warn!( + worker_id, + dp_rank, + "Unexpected replay frame count: {}", + msg.len() + ); + break; + } + }; if payload.is_empty() { break; } - let seq_bytes = msg.get(2).expect("frame count checked above"); if seq_bytes.len() != 8 { tracing::warn!( worker_id, @@ -555,3 +554,87 @@ async fn connect_replay_socket( } } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::protocols::compute_block_hash_for_seq; + use crate::services::indexer::backend::create_indexer; + use std::time::Duration; + + #[rstest::rstest] + #[case::sglang(false)] + #[case::vllm(true)] + #[tokio::test] + async fn replay_accepts_engine_framing(#[case] include_topic: bool) { + let context = zmq::Context::new(); + let router = context.socket(zmq::ROUTER).unwrap(); + router.set_linger(0).unwrap(); + router.set_rcvtimeo(5000).unwrap(); + router.set_sndtimeo(5000).unwrap(); + router.bind("tcp://127.0.0.1:*").unwrap(); + let endpoint = router.get_last_endpoint().unwrap().unwrap(); + let replay_socket = connect_dealer_socket(&endpoint).unwrap(); + let publisher = context.socket(zmq::PUB).unwrap(); + publisher.bind("tcp://127.0.0.1:*").unwrap(); + let live_socket = + connect_sub_socket(&publisher.get_last_endpoint().unwrap().unwrap()).unwrap(); + + let server = tokio::task::spawn_blocking(move || { + let request = router.recv_multipart(0).unwrap(); + assert_eq!(&request[1..], &[vec![], 0_u64.to_be_bytes().to_vec()]); + for seq in 0_u64..=2 { + let mut frames = vec![request[0].clone(), vec![]]; + if include_topic { + frames.push(b"kv-events".to_vec()); + } + if seq == 2 { + frames.extend([u64::MAX.to_be_bytes().to_vec(), vec![]]); + } else { + let event = ( + "BlockStored", + vec![seq + 100], + if seq == 0 { None } else { Some(100_u64) }, + vec![1_u32 + seq as u32; 4], + 4_usize, + Option::::None, + "GPU", + ); + let payload = rmp_serde::to_vec(&(0.0_f64, vec![event], Some(0_i32))).unwrap(); + frames.extend([seq.to_be_bytes().to_vec(), payload]); + } + router.send_multipart(frames, 0).unwrap(); + } + // Keep the socket open until the client consumes the queued replies. + assert_eq!(router.recv_multipart(0).unwrap()[1], b"done"); + }); + let indexer = create_indexer(4, 1); + let watermark = Arc::new(AtomicU64::new(WATERMARK_UNSET)); + let cancel = CancellationToken::new(); + let mut listener = ListenerLoop::new( + 1, + 0, + 4, + indexer.clone(), + cancel.clone(), + live_socket, + Some(replay_socket.clone()), + watermark.clone(), + ); + let replayed = tokio::time::timeout(Duration::from_secs(5), listener.replay_gap(0, 2)) + .await + .expect("replay completion marker") + .unwrap(); + send_multipart(&replay_socket, vec![b"done".to_vec()]) + .await + .unwrap(); + server.await.unwrap(); + assert_eq!(replayed, 2); + assert_eq!(watermark.load(Ordering::Acquire), 1); + indexer.dump_events().await.expect("flush indexer"); + let hashes = compute_block_hash_for_seq(&[1, 1, 1, 1, 2, 2, 2, 2], 4, Default::default()); + let matches = indexer.find_matches(hashes).await.unwrap(); + assert_eq!(matches.scores.get(&WorkerWithDpRank::new(1, 0)), Some(&2)); + cancel.cancel(); + } +} diff --git a/lib/kv-router/src/services/indexer/registry.rs b/lib/kv-router/src/services/indexer/registry.rs index ce96d84b5221..15fa98476cbd 100644 --- a/lib/kv-router/src/services/indexer/registry.rs +++ b/lib/kv-router/src/services/indexer/registry.rs @@ -311,6 +311,8 @@ pub struct WorkerEntry { pub struct WorkerRegistry { workers: DashMap, indexers: DashMap, + // Serialize indexer claims through worker publication with empty-indexer removal. + indexer_lifecycle: tokio::sync::Mutex<()>, peers: DashMap, watermarks: DashMap<(WorkerId, u32), Arc>, num_threads: usize, @@ -318,6 +320,7 @@ pub struct WorkerRegistry { ready_tx: watch::Sender, ready_rx: watch::Receiver, root_cancel_token: CancellationToken, + retain_empty_indexers: bool, } impl WorkerRegistry { @@ -357,6 +360,7 @@ impl WorkerRegistry { Self { workers: DashMap::new(), indexers: DashMap::new(), + indexer_lifecycle: tokio::sync::Mutex::new(()), peers: DashMap::new(), watermarks: DashMap::new(), num_threads, @@ -364,9 +368,17 @@ impl WorkerRegistry { ready_tx, ready_rx, root_cancel_token, + retain_empty_indexers: false, } } + #[cfg(feature = "standalone-selection")] + pub(crate) fn with_retained_indexers(mut self) -> Self { + // Selection entries and their schedulers retain these indexers across worker updates. + self.retain_empty_indexers = true; + self + } + pub fn signal_ready(&self) { let _ = self.ready_tx.send(true); } @@ -414,6 +426,7 @@ impl WorkerRegistry { replay_endpoint: Option, ) -> Result<()> { let key = RoutingPartitionId::new(model_name, routing_group); + let registration = self.indexer_lifecycle.lock().await; if let Some(entry) = self.workers.get(&instance_id) { if entry.key != key { @@ -486,6 +499,7 @@ impl WorkerRegistry { entry.listeners.insert(dp_rank, record.clone()); } + drop(registration); self.spawn_listener(instance_id, dp_rank, attempt, record); Ok(()) } @@ -521,10 +535,12 @@ impl WorkerRegistry { } } - if let Some(ie) = self.indexers.get(&key) { - ie.indexer.remove_worker(instance_id).await; + // Registration needs this map's write lock while removal can wait on indexer queues. + let indexer = self.indexers.get(&key).map(|entry| entry.indexer.clone()); + if let Some(indexer) = indexer { + indexer.remove_worker(instance_id).await; } - self.maybe_remove_indexer(&key); + self.maybe_remove_indexer(&key).await; Ok(()) } @@ -569,13 +585,17 @@ impl WorkerRegistry { .remove_if(&instance_id, |_, entry| entry.listeners.is_empty()) .is_some(); if actually_removed { - if let Some(ie) = self.indexers.get(&key) { - ie.indexer.remove_worker(instance_id).await; + let indexer = self.indexers.get(&key).map(|entry| entry.indexer.clone()); + if let Some(indexer) = indexer { + indexer.remove_worker(instance_id).await; } - self.maybe_remove_indexer(&key); + self.maybe_remove_indexer(&key).await; + } + } else { + let indexer = self.indexers.get(&key).map(|entry| entry.indexer.clone()); + if let Some(indexer) = indexer { + indexer.remove_worker_dp_rank(instance_id, dp_rank).await; } - } else if let Some(ie) = self.indexers.get(&key) { - ie.indexer.remove_worker_dp_rank(instance_id, dp_rank).await; } Ok(()) @@ -610,10 +630,11 @@ impl WorkerRegistry { } } - if let Some(ie) = self.indexers.get(&key) { - ie.indexer.remove_worker(instance_id).await; + let indexer = self.indexers.get(&key).map(|entry| entry.indexer.clone()); + if let Some(indexer) = indexer { + indexer.remove_worker(instance_id).await; } - self.maybe_remove_indexer(&key); + self.maybe_remove_indexer(&key).await; Ok(()) } @@ -807,7 +828,12 @@ impl WorkerRegistry { ); } - fn maybe_remove_indexer(&self, key: &RoutingPartitionId) { + async fn maybe_remove_indexer(&self, key: &RoutingPartitionId) { + if self.retain_empty_indexers { + return; + } + + let _lifecycle = self.indexer_lifecycle.lock().await; if self.workers.iter().any(|entry| entry.value().key == *key) { return; } @@ -819,41 +845,184 @@ impl WorkerRegistry { #[cfg(test)] mod tests { use super::*; + use crate::protocols::{LocalBlockHash, StorageTier, WorkerWithDpRank}; + use crate::services::indexer::backend::test_util::store_event; + use std::future::Future; use std::sync::atomic::Ordering; + use std::task::{Context, Poll, Waker}; + use std::time::Duration; fn test_registry() -> WorkerRegistry { WorkerRegistry::new(1) } + #[rstest::rstest] + #[case("worker")] + #[case("all_groups")] + #[case("last_rank")] #[tokio::test] - async fn deregister_removes_watermark() { + async fn registration_progresses_while_removal_is_backpressured(#[case] removal: &str) { let registry = test_registry(); - registry.signal_ready(); - + let key = RoutingPartitionId::new("test-model", "default"); registry .register( 1, - "tcp://127.0.0.1:15557".to_string(), + "tcp://127.0.0.1:15557".into(), 0, - "test-model".to_string(), - "default".to_string(), + "test-model".into(), + "default".into(), 1, None, ) .await .unwrap(); + let indexer = registry.indexers.get(&key).unwrap().indexer.clone(); + let Indexer::Single { primary, .. } = indexer else { + unreachable!(); + }; + let sender = primary.remove_worker_sender(); + // Reserve the entire queue so removal must yield without blocking the executor. + let permits = sender.reserve_many(sender.max_capacity()).await.unwrap(); + let removal = async { + match removal { + "worker" => registry.deregister(1, "test-model", "default").await, + "all_groups" => { + registry + .deregister_all_routing_groups(1, "test-model") + .await + } + "last_rank" => { + registry + .deregister_dp_rank(1, 0, "test-model", "default") + .await + } + _ => unreachable!(), + } + }; + tokio::pin!(removal); + assert!(matches!( + removal + .as_mut() + .poll(&mut Context::from_waker(Waker::noop())), + Poll::Pending + )); + // A nonblocking probe makes the regression fail instead of deadlocking the test. + assert!(matches!( + registry.indexers.try_get_mut(&key), + dashmap::try_result::TryResult::Present(_) + )); + registry + .register( + 2, + "tcp://127.0.0.1:15558".into(), + 0, + "test-model".into(), + "default".into(), + 1, + None, + ) + .await + .unwrap(); + drop(permits); + tokio::time::timeout(Duration::from_secs(5), removal) + .await + .expect("removal should finish after queue capacity is released") + .unwrap(); + assert!(registry.workers.contains_key(&2)); + assert!(registry.indexers.contains_key(&key)); + registry.root_cancel_token.cancel(); + } - assert!(registry.watermarks.contains_key(&(1, 0))); - + #[rstest::rstest] + #[case("worker")] + #[case("all_groups")] + #[case("last_rank")] + #[tokio::test] + async fn empty_indexer_removal_waits_for_registration(#[case] removal: &str) { + let registry = test_registry(); + let key = RoutingPartitionId::new("test-model", "default"); registry - .deregister(1, "test-model", "default") + .register( + 1, + "tcp://127.0.0.1:15557".into(), + 0, + "test-model".into(), + "default".into(), + 1, + None, + ) .await .unwrap(); - assert!( - !registry.watermarks.contains_key(&(1, 0)), - "watermark should be removed after deregister" + // Queue registration ahead of pruning, without timing or thread scheduling assumptions. + let lifecycle = registry.indexer_lifecycle.lock().await; + let registration = registry.register( + 2, + "tcp://127.0.0.1:15558".into(), + 0, + "test-model".into(), + "default".into(), + 1, + None, ); + tokio::pin!(registration); + assert!(matches!( + registration + .as_mut() + .poll(&mut Context::from_waker(Waker::noop())), + Poll::Pending + )); + + let removal = async { + match removal { + "worker" => registry.deregister(1, "test-model", "default").await, + "all_groups" => { + registry + .deregister_all_routing_groups(1, "test-model") + .await + } + "last_rank" => { + registry + .deregister_dp_rank(1, 0, "test-model", "default") + .await + } + _ => unreachable!(), + } + }; + tokio::pin!(removal); + assert!(matches!( + removal + .as_mut() + .poll(&mut Context::from_waker(Waker::noop())), + Poll::Pending + )); + assert!(registry.workers.is_empty()); + assert!(!registry.watermarks.contains_key(&(1, 0))); + assert!(registry.get_indexer(&key).is_some()); + drop(lifecycle); + registration.await.unwrap(); + removal.await.unwrap(); + + // Events written by the new listener must remain visible through the query registry. + let listener_indexer = registry.workers.get(&2).unwrap().listeners[&0] + .indexer + .clone(); + listener_indexer + .apply_event_routed(store_event(2, 0, 0, &[], &[11], StorageTier::Device)) + .await + .unwrap(); + listener_indexer.dump_events().await.unwrap(); + let query_indexer = registry + .get_indexer(&key) + .expect("registered indexer") + .indexer + .clone(); + let scores = query_indexer + .find_matches(vec![LocalBlockHash(11)]) + .await + .unwrap(); + assert_eq!(scores.scores.get(&WorkerWithDpRank::new(2, 0)), Some(&1)); + registry.root_cancel_token.cancel(); } #[tokio::test] @@ -969,15 +1138,23 @@ mod tests { .unwrap(); // Simulate that the listener advanced the watermark. - if let Some(wm) = registry.watermarks.get(&(1, 0)) { - wm.store(42, Ordering::Release); - } + registry + .watermarks + .get(&(1, 0)) + .unwrap() + .store(42, Ordering::Release); registry .deregister(1, "test-model", "default") .await .unwrap(); + assert!( + registry + .get_indexer(&RoutingPartitionId::new("test-model", "default")) + .is_none() + ); + registry .register( 1, @@ -1002,37 +1179,6 @@ mod tests { ); } - #[tokio::test] - async fn deregister_all_routing_groups_removes_watermarks() { - let registry = test_registry(); - registry.signal_ready(); - - registry - .register( - 1, - "tcp://127.0.0.1:15562".to_string(), - 0, - "test-model".to_string(), - "default".to_string(), - 1, - None, - ) - .await - .unwrap(); - - assert!(registry.watermarks.contains_key(&(1, 0))); - - registry - .deregister_all_routing_groups(1, "test-model") - .await - .unwrap(); - - assert!( - !registry.watermarks.contains_key(&(1, 0)), - "watermark should be removed after deregister_all_routing_groups" - ); - } - // ── list_filtered tests ─────────────────────────────────────────────────── #[tokio::test] diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index f4c45b1fa7dd..6ee9143eb200 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -203,10 +203,10 @@ impl SelectionCore { tracking_hash: Arc, ) -> Self { let cancel_token = cancel_token.child_token(); - let indexer_registry = Arc::new(WorkerRegistry::new_with_cancel_token( - indexer_threads, - cancel_token.clone(), - )); + let indexer_registry = Arc::new( + WorkerRegistry::new_with_cancel_token(indexer_threads, cancel_token.clone()) + .with_retained_indexers(), + ); if signal_indexer_ready { indexer_registry.signal_ready(); } @@ -1413,41 +1413,85 @@ mod tests { } } - #[test] - fn shutdown_keeps_parent_alive() { + #[tokio::test] + async fn shutdown_cancels_listeners_but_keeps_parent_alive() { let parent = CancellationToken::new(); let core = SelectionCore::new_local( - test_config(false), + test_config(true), 1, parent.clone(), SelectionCacheConfig::default(), ); - core.shutdown(); + 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 shutdown_cancels_listeners() { - let parent = CancellationToken::new(); - let core = SelectionCore::new_local( + async fn selection_sees_cache_after_last_worker_replacement(#[case] indexer_threads: usize) { + let core = SelectionCore::try_new_local( test_config(true), - 1, - parent, + 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 + }; - 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)); + // 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(); - assert_eq!(core.indexer_registry.listener_cancelled(1, 0), Some(true)); } #[tokio::test] From a57c6cefda88537184c7e51c8ca3b06773fcbb29 Mon Sep 17 00:00:00 2001 From: Thach Nguyen Date: Wed, 23 Sep 2026 22:16:43 +0000 Subject: [PATCH 04/14] indexer image: build the standalone indexer from a git archive Carries forward DeepInfra's indexer image recipe: a two-stage Ubuntu 24.04 build that compiles the dynamo._core wheel with kv-indexer-metrics and installs only that wheel into a slim runtime. Ports cd18634d7 (image + build script), ef5b033ac (release build) and f4b4f1d20 (touch archived sources so a sibling branch's newer artifact in the shared cargo cache is never reused). Header comment now lists the three service flavors the one image serves. Claude-Session: https://claude.ai/code/session_01XKZ49SpoymBqDqgGRJYMNa --- container/build-indexer-image.sh | 33 ++++++++++++ container/indexer.Dockerfile | 93 ++++++++++++++++++++++++++++++++ 2 files changed, 126 insertions(+) create mode 100755 container/build-indexer-image.sh create mode 100644 container/indexer.Dockerfile diff --git a/container/build-indexer-image.sh b/container/build-indexer-image.sh new file mode 100755 index 000000000000..4b713baf5105 --- /dev/null +++ b/container/build-indexer-image.sh @@ -0,0 +1,33 @@ +#!/usr/bin/env bash +# Build the standalone KV-cache indexer image from COMMITTED dynamo source. +# +# ./container/build-indexer-image.sh [tag] +# +# Builds a multi-stage image (compile wheel -> slim runtime) using a git-archive +# of HEAD as the build context, so the multi-GB local target/ is never shipped +# to the docker daemon. Default tag: /dynamo-indexer:kvtest-. +# Push separately: docker push +set -euo pipefail + +REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +REGISTRY="${REGISTRY:-localhost:30500}" +SHA="$(git -C "$REPO_ROOT" rev-parse --short=10 HEAD)" +TAG="${1:-${REGISTRY}/dynamo-indexer:kvtest-${SHA}}" + +BUILD_DIR="$(mktemp -d)" +trap 'rm -rf "$BUILD_DIR"' EXIT + +echo "==> git archive ${SHA} -> build context (committed files only)" +git -C "$REPO_ROOT" archive --format=tar HEAD -o "$BUILD_DIR/dynamo-src.tar" +cp "$REPO_ROOT/container/indexer.Dockerfile" "$BUILD_DIR/Dockerfile" + +echo "==> docker build $TAG" +# --ulimit nofile: BuildKit RUN steps otherwise inherit a low fd limit, which a +# big parallel cargo build exhausts ("too many open files"). The Dockerfile also +# caps CARGO_BUILD_JOBS, but raise the ceiling here too. +DOCKER_BUILDKIT=1 docker build \ + --ulimit nofile=1048576:1048576 \ + -f "$BUILD_DIR/Dockerfile" -t "$TAG" "$BUILD_DIR" + +echo "==> done: $TAG" +echo " push with: docker push $TAG" diff --git a/container/indexer.Dockerfile b/container/indexer.Dockerfile new file mode 100644 index 000000000000..ced345c512c6 --- /dev/null +++ b/container/indexer.Dockerfile @@ -0,0 +1,93 @@ +# syntax=docker/dockerfile:1.7 +# ============================================================================= +# Standalone KV-cache indexer image -> `python -m dynamo.indexer` +# +# Serves DeepInfra's per-model kv-indexer services; one image, the flavor is a flag: +# kv-indexer:reality (no flags) real evictions, per-shard trees +# kv-indexer:routing --keep-evictions evictions parked, per-shard trees +# kv-indexer:h24 --h24 flat "stored by anyone" counterfactual +# Each discovers the model's engine pods, subscribes to their ZMQ KV events and +# serves /query /query_by_hash /workers /dump /health (+ /metrics). +# +# Self-contained: compiles the Rust python bindings (dynamo._core) from source, +# then installs ONLY that wheel into a slim runtime. Both stages are ubuntu:24.04, +# so the runtime libstdc++ is identical to what the extension was linked against +# -> no LD_PRELOAD shim needed (that hack was only for the mismatched conda env). +# +# The build context is a `git archive` of the dynamo source (committed files +# only), so the multi-GB local target/ never enters the build. The indexer is +# fully standalone: no etcd / NATS, just HTTP + ZMQ. See build-indexer-image.sh. +# ============================================================================= + +# ---------- builder: compile the wheel ------------------------------------- +FROM ubuntu:24.04 AS builder +ENV DEBIAN_FRONTEND=noninteractive +RUN apt-get update && apt-get install -y --no-install-recommends \ + ca-certificates curl git \ + build-essential pkg-config cmake \ + clang libclang-dev \ + protobuf-compiler libprotobuf-dev \ + libzmq3-dev \ + python3 python3-dev python3-venv \ + && rm -rf /var/lib/apt/lists/* + +# Rust: rustup with NO default toolchain. dynamo's rust-toolchain.toml pins +# the toolchain, which rustup auto-fetches on the first cargo run in /src. +RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs \ + | sh -s -- -y --no-modify-path --default-toolchain none --profile minimal +ENV PATH=/root/.cargo/bin:$PATH + +# maturin in an isolated venv (Ubuntu 24.04 system python is externally-managed). +RUN python3 -m venv /opt/maturin && /opt/maturin/bin/pip install --no-cache-dir maturin +ENV PATH=/opt/maturin/bin:$PATH + +# Source tarball (git archive), auto-extracted by ADD. rust-toolchain.toml lands +# at /src so rustup resolves the pinned toolchain for the whole workspace. +ADD dynamo-src.tar /src +# git archive stamps every file with the commit time, and cargo judges a path +# crate fresh by mtime against the shared /cargo-target cache. A sibling build +# (another branch) can therefore leave a newer artifact that cargo reuses with +# the wrong contents; stamping the sources with the build time prevents that. +RUN find /src -type f -exec touch {} + + +WORKDIR /src/lib/bindings/python +ENV CARGO_TARGET_DIR=/cargo-target +# Cap parallelism: this compiles on a shared ~288-core services node. The default +# (one rustc per core) both starves prod neighbours and exhausts the build's +# file-descriptor limit (hundreds of parallel rustc -> "too many open files"). +# 16 jobs is plenty and well-behaved; build-indexer-image.sh also raises nofile. +ENV CARGO_BUILD_JOBS=16 +# Release build: optimized for real indexer throughput (slower to compile than a +# debug build, but the radix-tree query/apply hot paths need the optimizations). +# kv-indexer-metrics adds the `dynamo.indexer` binary + Prometheus /metrics. +# nixl-sys' build script runs bindgen, so point it at libclang (resolved +# dynamically to survive llvm bumps). Cache mounts keep the crate registry + +# compiled artifacts warm across rebuilds. +RUN --mount=type=cache,target=/root/.cargo/registry \ + --mount=type=cache,target=/cargo-target \ + export LIBCLANG_PATH="$(dirname "$(find /usr/lib/llvm-* -name 'libclang.so*' 2>/dev/null | head -1)")" \ + && maturin build --release --features kv-indexer-metrics --out /wheels + +# ---------- runtime: install just the wheel -------------------------------- +FROM ubuntu:24.04 AS runtime +ENV DEBIAN_FRONTEND=noninteractive +RUN apt-get update && apt-get install -y --no-install-recommends \ + ca-certificates \ + libzmq5 \ + python3 python3-venv \ + && rm -rf /var/lib/apt/lists/* + +# venv with only the indexer wheel and its deps (pydantic, uvloop -> manylinux +# wheels from PyPI, no compiler needed here). +COPY --from=builder /wheels/ /wheels/ +RUN python3 -m venv /opt/venv \ + && /opt/venv/bin/pip install --no-cache-dir /wheels/ai_dynamo_runtime*.whl \ + && rm -rf /wheels +ENV PATH=/opt/venv/bin:$PATH + +# Smoke-test the import + CLI at build time so a broken wheel fails the build. +RUN python -m dynamo.indexer --help >/dev/null + +EXPOSE 8090 +ENTRYPOINT ["python", "-m", "dynamo.indexer"] +CMD ["--port", "8090"] From 6607fcd1475ced398229ef8f961cfbc215ea17ad Mon Sep 17 00:00:00 2001 From: Thach Nguyen Date: Wed, 23 Sep 2026 22:17:30 +0000 Subject: [PATCH 05/14] standalone-indexer: unbounded ZMQ receive queue on SUB sockets A listener blocked in gap recovery can let its SUB pipe reach the default RCVHWM of 1000; with heartbeats on, libzmq 4.3.4 then aborts on `Assertion failed: _input_stopped` (zeromq/libzmq#3596, #3937). Seen in prod as an h24 indexer crash loop while 66 startup TreeDumps applied. Only the indexer's connect_sub_socket changes; the shared receive config used by replica sync, selection and the slot tracker keeps its default. Ports ae7403cbe. Claude-Session: https://claude.ai/code/session_01XKZ49SpoymBqDqgGRJYMNa --- lib/kv-router/src/services/common/zmq.rs | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/lib/kv-router/src/services/common/zmq.rs b/lib/kv-router/src/services/common/zmq.rs index 406be80b5103..8b0a4acc27e0 100644 --- a/lib/kv-router/src/services/common/zmq.rs +++ b/lib/kv-router/src/services/common/zmq.rs @@ -18,6 +18,12 @@ pub(crate) type MultipartMessage = Vec>; pub(crate) type SharedSocket = Arc>; const ZMQ_RCVTIMEOUT_MS: i32 = 100; +// Unbounded on purpose: libzmq 4.3.4 asserts (`_input_stopped`) when a +// heartbeating peer restarts a pipe that HWM-stopped input, so an indexer +// listener stalled in gap recovery must never let its queue fill +// (zeromq/libzmq#3596). +#[cfg(feature = "standalone-indexer")] +const ZMQ_INDEXER_RCVHWM: i32 = 0; const ZMQ_SNDTIMEOUT_MS: i32 = 0; const ZMQ_RECONNECT_IVL_MS: i32 = 100; const ZMQ_RECONNECT_IVL_MAX_MS: i32 = 5000; @@ -217,6 +223,8 @@ pub(crate) fn create_sub_socket(topic: &[u8]) -> Result { #[cfg(feature = "standalone-indexer")] pub(crate) fn connect_sub_socket(endpoint: &str) -> Result { let socket = create_sub_socket(b"")?; + // Must precede connect: HWM applies to pipes created afterwards. + socket.socket().set_rcvhwm(ZMQ_INDEXER_RCVHWM)?; socket.connect(endpoint)?; Ok(Arc::new(Mutex::new(socket))) } @@ -298,3 +306,15 @@ pub(crate) fn validate_endpoint(endpoint: &str) -> Result<()> { )), } } + +#[cfg(all(test, feature = "standalone-indexer"))] +mod tests { + use super::*; + + #[tokio::test] + async fn indexer_sub_socket_receive_queue_is_unbounded() { + let socket = connect_sub_socket("inproc://rcvhwm-test").unwrap(); + let guard = socket.lock().await; + assert_eq!(guard.socket().get_rcvhwm().unwrap(), 0); + } +} From 41d3a5ae4df12c7038fc1d57cd1bf1cf826c8287 Mon Sep 17 00:00:00 2001 From: Thach Nguyen Date: Wed, 23 Sep 2026 22:19:31 +0000 Subject: [PATCH 06/14] standalone-indexer: exit on a block-size mismatch, skip partial-prefix stores The shared normalizer drops every block whose size differs from the configured one, so an indexer started with the wrong --block-size stays empty while its stream looks alive. The standalone listener now checks each preprocessed BlockStored first: - a store of blocks larger than configured is a misconfiguration: log and exit so k8s surfaces a CrashLoopBackOff; - a store smaller than configured is a vLLM partial-prefix entry (prefix_match_unit < cache block, e.g. 32/64/96 of 128 on V4.1): drop it, logging the first three. Ports 959e440f0 and 90409c7f8. Unlike the originals this lives in the standalone listener (new block_size.rs) and leaves convert_event's signature and the lib/llm publisher untouched; the publisher half of 959e440f0 is Dynamo's own worker-side path, which we don't run. Claude-Session: https://claude.ai/code/session_01XKZ49SpoymBqDqgGRJYMNa --- .../src/services/indexer/block_size.rs | 171 ++++++++++++++++++ .../src/services/indexer/listener.rs | 35 +++- lib/kv-router/src/services/indexer/mod.rs | 1 + 3 files changed, 197 insertions(+), 10 deletions(-) create mode 100644 lib/kv-router/src/services/indexer/block_size.rs diff --git a/lib/kv-router/src/services/indexer/block_size.rs b/lib/kv-router/src/services/indexer/block_size.rs new file mode 100644 index 000000000000..784116653209 --- /dev/null +++ b/lib/kv-router/src/services/indexer/block_size.rs @@ -0,0 +1,171 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Block-size check on engine `BlockStored` events before they are indexed. +//! +//! The shared ZMQ normalizer silently drops every block whose size differs +//! from the configured one, so an indexer started with the wrong +//! `--block-size` stays empty while its event stream looks healthy. The +//! standalone indexer treats that as a fatal misconfiguration instead, except +//! for vLLM's partial-prefix entries (see [`StoredBlockSize::PartialPrefix`]). + +use crate::protocols::WorkerWithDpRank; +use crate::zmq_wire::RawKvEvent; + +/// Per-listener gate applied to every preprocessed engine event. +pub(super) struct BlockSizeGuard { + configured_block_size: u32, + partial_prefix_skips: u64, +} + +impl BlockSizeGuard { + pub(super) fn new(configured_block_size: u32) -> Self { + Self { + configured_block_size, + partial_prefix_skips: 0, + } + } + + /// Whether `raw` should be indexed. Partial-prefix stores are dropped + /// (first few logged). A block-size mismatch **exits the process**: every + /// later store would be dropped the same way, and a crash loop is the + /// only signal k8s surfaces for an indexer that looks alive but is empty. + pub(super) fn admit(&mut self, raw: &RawKvEvent, worker: WorkerWithDpRank) -> bool { + match classify_stored_block_size(raw, self.configured_block_size) { + StoredBlockSize::Indexable => true, + StoredBlockSize::PartialPrefix { event_block_size } => { + self.partial_prefix_skips += 1; + if self.partial_prefix_skips <= 3 { + tracing::warn!( + worker_id = worker.worker_id, + dp_rank = worker.dp_rank, + event_block_size, + configured_block_size = self.configured_block_size, + "Skipping sub-block BlockStored: a partial-prefix entry \ + (prefix_match_unit < block_size), unless the indexer's \ + --block-size is larger than the engine's -- then the \ + index stays empty; check configured vs event size" + ); + } + false + } + StoredBlockSize::Mismatch { event_block_size } => { + tracing::error!( + worker_id = worker.worker_id, + dp_rank = worker.dp_rank, + event_block_size, + configured_block_size = self.configured_block_size, + "Block size mismatch: every stored block would be dropped and the \ + index would stay empty. Fix the indexer's --block-size to match \ + the engine's --block-size. Exiting." + ); + std::process::exit(1); + } + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum StoredBlockSize { + /// Not a store, an empty store, or a store of whole configured blocks. + Indexable, + /// A store whose `block_size` is smaller than the configured cache block: + /// vLLM hashes prefixes every `prefix_match_unit` (hash_block_size) + /// tokens, which can be finer than the cache block, and publishes the + /// prompt tail inside a cache block as a sub-block store (32/64/96 tokens + /// for a 128-token block). It carries nothing the index keys on. + PartialPrefix { event_block_size: usize }, + /// A store of blocks larger than configured: the indexer's `--block-size` + /// does not match the engine's, and every stored block would be dropped. + Mismatch { event_block_size: usize }, +} + +pub(super) fn classify_stored_block_size( + raw: &RawKvEvent, + configured_block_size: u32, +) -> StoredBlockSize { + let RawKvEvent::BlockStored { + block_hashes, + block_size, + .. + } = raw + else { + return StoredBlockSize::Indexable; + }; + let event_block_size = *block_size; + let configured = configured_block_size as usize; + if block_hashes.is_empty() || event_block_size == configured { + StoredBlockSize::Indexable + } else if event_block_size > 0 && event_block_size < configured { + StoredBlockSize::PartialPrefix { event_block_size } + } else { + StoredBlockSize::Mismatch { event_block_size } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::zmq_wire::BlockHashValue; + + fn stored(block_size: usize, num_blocks: u64) -> RawKvEvent { + RawKvEvent::BlockStored { + block_hashes: (0..num_blocks).map(BlockHashValue::Unsigned).collect(), + parent_block_hash: None, + token_ids: vec![0; block_size * num_blocks as usize], + block_size, + medium: None, + lora_name: None, + cache_namespace: None, + block_mm_infos: None, + is_eagle: None, + group_idx: None, + kv_cache_spec_kind: None, + kv_cache_spec_sliding_window: None, + locality: None, + ownership: None, + } + } + + #[test] + fn whole_blocks_are_indexable() { + assert_eq!( + classify_stored_block_size(&stored(128, 2), 128), + StoredBlockSize::Indexable + ); + } + + #[test] + fn sub_block_store_is_a_partial_prefix_entry() { + for size in [32, 64, 96] { + assert_eq!( + classify_stored_block_size(&stored(size, 1), 128), + StoredBlockSize::PartialPrefix { + event_block_size: size + } + ); + } + } + + #[test] + fn larger_engine_block_is_a_mismatch() { + assert_eq!( + classify_stored_block_size(&stored(256, 1), 128), + StoredBlockSize::Mismatch { + event_block_size: 256 + } + ); + } + + #[test] + fn empty_store_and_non_store_events_are_indexable() { + assert_eq!( + classify_stored_block_size(&stored(256, 0), 128), + StoredBlockSize::Indexable + ); + assert_eq!( + classify_stored_block_size(&RawKvEvent::Ignored, 128), + StoredBlockSize::Indexable + ); + } +} diff --git a/lib/kv-router/src/services/indexer/listener.rs b/lib/kv-router/src/services/indexer/listener.rs index d2cd5d05996f..76fa8bde4ad9 100644 --- a/lib/kv-router/src/services/indexer/listener.rs +++ b/lib/kv-router/src/services/indexer/listener.rs @@ -12,6 +12,7 @@ use crate::recovery::{CursorObservation, CursorState}; use crate::zmq_wire::{ZmqEventNormalizer, decode_event_batch}; use super::backend::Indexer; +use super::block_size::BlockSizeGuard; use super::registry::ListenerRecord; use crate::services::common::zmq::{ MultipartMessage, SharedSocket, connect_dealer_socket, connect_sub_socket, recv_multipart, @@ -153,6 +154,7 @@ struct ListenerLoop { replay_socket: Option, watermark: Arc, normalizer: ZmqEventNormalizer, + block_size_guard: BlockSizeGuard, messages_processed: u64, } @@ -177,6 +179,7 @@ impl ListenerLoop { replay_socket, watermark, normalizer: ZmqEventNormalizer::new(block_size), + block_size_guard: BlockSizeGuard::new(block_size), messages_processed: 0, } } @@ -267,11 +270,17 @@ impl ListenerLoop { .data_parallel_rank .map_or(dp_rank, |rank| rank.cast_unsigned()); for raw_event in batch.events { - let Some(placement_event) = self.normalizer.normalize( - raw_event, - seq, - WorkerWithDpRank::new(worker_id, effective_dp_rank), - ) else { + let worker = WorkerWithDpRank::new(worker_id, effective_dp_rank); + let Some(raw_event) = self.normalizer.preprocess(raw_event, worker) else { + continue; + }; + if !self.block_size_guard.admit(&raw_event, worker) { + continue; + } + let Some(placement_event) = self + .normalizer + .normalize_preprocessed(raw_event, seq, worker) + else { continue; }; let router_event = placement_event @@ -343,11 +352,17 @@ impl ListenerLoop { .data_parallel_rank .map_or(self.dp_rank, |rank| rank.cast_unsigned()); for raw_event in batch.events { - let Some(placement_event) = self.normalizer.normalize( - raw_event, - seq, - WorkerWithDpRank::new(self.worker_id, effective_dp_rank), - ) else { + let worker = WorkerWithDpRank::new(self.worker_id, effective_dp_rank); + let Some(raw_event) = self.normalizer.preprocess(raw_event, worker) else { + continue; + }; + if !self.block_size_guard.admit(&raw_event, worker) { + continue; + } + let Some(placement_event) = self + .normalizer + .normalize_preprocessed(raw_event, seq, worker) + else { continue; }; let router_event = placement_event diff --git a/lib/kv-router/src/services/indexer/mod.rs b/lib/kv-router/src/services/indexer/mod.rs index 532366d882b1..37c455a91bfd 100644 --- a/lib/kv-router/src/services/indexer/mod.rs +++ b/lib/kv-router/src/services/indexer/mod.rs @@ -30,6 +30,7 @@ //! response struct in [`server`] for the exact semantics. pub mod backend; +mod block_size; pub mod listener; pub mod logging; pub mod metrics; From 76fb460bdeb04b8e53070405f28166f812ca4c8d Mon Sep 17 00:00:00 2001 From: Thach Nguyen Date: Wed, 23 Sep 2026 22:28:37 +0000 Subject: [PATCH 07/14] standalone-indexer: recover gaps over HTTP /kv_recover with timeout, gate and retry Our engines serve their local indexer at GET /kv_recover (HTTP) rather than on a ZMQ ROUTER replay socket. A listener registered with a recover_endpoint now recovers a sequence gap by fetching a WorkerKvQueryResponse for [start, end) and applying it: - Events: apply all under the consumer-assigned (worker, dp_rank), resume at last_event_id; - TreeDump (whole-rank scope): remove_worker_dp_rank, apply all, resume at last_event_id; a domain-scoped dump applies without the reset; - TooNew / InvalidRange / Error / TreeDumpFailed: no-op. Without a recover_endpoint the upstream ZMQ replay path is unchanged. Downloads share one reqwest client and a process-wide gate (--recover-concurrency, default 8), with a total timeout of --recover-timeout-secs (default 120) and 3 attempts with 2/4 s backoff before the gap is given up ("kv_recover request failed; giving up"). v1.5.0 has no HTTP recovery at all (its only HTTP client is peer /dump recovery, hard-coded to 10 s). The settings are injected through WorkerRegistry::with_kv_recover rather than process globals, and /register takes recover_endpoint as a new field next to replay_endpoint instead of aliasing it. register() keeps upstream's signature; register_with_extras() carries our fields. Tests: a /kv_recover response captured from a prod GLM-5.3 engine deserializes and plans under the consumer identity; retry until success; give up after the last attempt; per-attempt timeout. Ports ab351013f and 37e1b1a26. Claude-Session: https://claude.ai/code/session_01XKZ49SpoymBqDqgGRJYMNa --- lib/bindings/python/rust/llm/kv.rs | 14 + .../src/services/indexer/kv_recover.rs | 410 ++++++++++++++++++ .../src/services/indexer/listener.rs | 80 +++- lib/kv-router/src/services/indexer/mod.rs | 32 +- .../src/services/indexer/registry.rs | 66 +++ lib/kv-router/src/services/indexer/server.rs | 47 +- 6 files changed, 632 insertions(+), 17 deletions(-) create mode 100644 lib/kv-router/src/services/indexer/kv_recover.rs diff --git a/lib/bindings/python/rust/llm/kv.rs b/lib/bindings/python/rust/llm/kv.rs index 91c86ce80878..78f9e012fc8b 100644 --- a/lib/bindings/python/rust/llm/kv.rs +++ b/lib/bindings/python/rust/llm/kv.rs @@ -228,6 +228,16 @@ struct KvIndexerCli { /// Use local timezone for access log timestamps (default: UTC) #[arg(long)] access_log_local_time: bool, + + /// Total timeout in seconds for one engine `GET /kv_recover` download + /// (connect + body). A full TreeDump of a large engine is tens of MB. + #[arg(long, default_value_t = indexer::kv_recover::DEFAULT_RECOVER_TIMEOUT_S)] + recover_timeout_secs: u64, + + /// Maximum concurrent `/kv_recover` downloads across all listeners of this + /// indexer. Bounds the load a fleet-wide (re)subscription puts on engines. + #[arg(long, default_value_t = indexer::kv_recover::DEFAULT_RECOVER_CONCURRENCY)] + recover_concurrency: usize, } pub fn run_kv_indexer_cli(args: I) -> anyhow::Result<()> @@ -262,6 +272,10 @@ where access_log: cli.access_log, trace_id_header, access_log_local_time: cli.access_log_local_time, + kv_recover: indexer::kv_recover::KvRecoverSettings { + timeout_s: cli.recover_timeout_secs, + concurrency: cli.recover_concurrency, + }, })) } diff --git a/lib/kv-router/src/services/indexer/kv_recover.rs b/lib/kv-router/src/services/indexer/kv_recover.rs new file mode 100644 index 000000000000..30d0f5230644 --- /dev/null +++ b/lib/kv-router/src/services/indexer/kv_recover.rs @@ -0,0 +1,410 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Per-worker gap recovery over HTTP: `GET /kv_recover`. +//! +//! An alternative to the ZMQ DEALER replay socket for engines that serve their +//! local indexer over HTTP. On a sequence gap the listener fetches a +//! [`WorkerKvQueryResponse`] for `[start, end)` and applies it per +//! [`plan_recovery`]: +//! +//! - `Events` — apply all, resume at `last_event_id`. +//! - `TreeDump` — drop this `(worker, dp_rank)`'s state, apply all, resume at +//! `last_event_id` (dump event ids are synthetic). +//! - anything else — no-op, the watermark stays where it was. +//! +//! Recovery is at-least-once: the resume point is always the response's +//! `last_event_id`, never a count of applied events. Events are applied under +//! the consumer-assigned identity; the worker's self-reported one is ignored. +//! +//! A large engine's TreeDump is tens of MB serialized inside the engine +//! process, and a freshly started indexer asks every engine at once, so +//! downloads share a process-wide gate, get a long total timeout, and are +//! retried with backoff before the gap is given up. + +use std::time::Duration; + +use anyhow::Context; +use tokio::sync::Semaphore; +use tokio_util::sync::CancellationToken; + +use crate::indexer::WorkerKvQueryResponse; +use crate::protocols::{ResetScope, RouterEvent, WorkerId}; + +pub const DEFAULT_RECOVER_TIMEOUT_S: u64 = 120; +pub const DEFAULT_RECOVER_CONCURRENCY: usize = 8; + +const RECOVER_CONNECT_TIMEOUT: Duration = Duration::from_secs(5); +const RECOVER_ATTEMPTS: u32 = 3; +const RECOVER_BACKOFF_BASE: Duration = Duration::from_secs(2); + +#[derive(Debug, Clone, Copy)] +pub struct KvRecoverSettings { + /// Total timeout (connect + body) for one download. + pub timeout_s: u64, + /// Downloads allowed in flight across every listener of this process. + pub concurrency: usize, +} + +impl Default for KvRecoverSettings { + fn default() -> Self { + Self { + timeout_s: DEFAULT_RECOVER_TIMEOUT_S, + concurrency: DEFAULT_RECOVER_CONCURRENCY, + } + } +} + +/// Shared by every listener of a registry: one HTTP client and one gate. +pub struct KvRecoverClient { + http: reqwest::Client, + gate: Semaphore, + backoff_base: Duration, +} + +impl KvRecoverClient { + pub fn new(settings: KvRecoverSettings) -> anyhow::Result { + Self::with_backoff(settings, RECOVER_BACKOFF_BASE) + } + + fn with_backoff(settings: KvRecoverSettings, backoff_base: Duration) -> anyhow::Result { + let http = reqwest::Client::builder() + .connect_timeout(RECOVER_CONNECT_TIMEOUT) + .timeout(Duration::from_secs(settings.timeout_s.max(1))) + .build() + .context("failed to build kv_recover HTTP client")?; + Ok(Self { + http, + gate: Semaphore::new(settings.concurrency.max(1)), + backoff_base, + }) + } + + /// Download the worker's recovery response for `[start_seq, end_seq)`, + /// retrying failed downloads. `None` when every attempt failed (logged as + /// `kv_recover request failed; giving up`) or `cancel` fired. + pub(super) async fn fetch( + &self, + recover_endpoint: &str, + start_seq: u64, + end_seq: u64, + worker_id: WorkerId, + dp_rank: u32, + cancel: &CancellationToken, + ) -> Option { + let url = format!("{}/kv_recover", recover_endpoint.trim_end_matches('/')); + for attempt in 0..RECOVER_ATTEMPTS { + let result = tokio::select! { + _ = cancel.cancelled() => return None, + result = self.fetch_once(&url, start_seq, end_seq) => result, + }; + let error = match result { + Ok(response) => return Some(response), + Err(error) => error, + }; + if attempt + 1 == RECOVER_ATTEMPTS { + tracing::error!( + worker_id, + dp_rank, + attempts = RECOVER_ATTEMPTS, + gap_size = end_seq.saturating_sub(start_seq), + error = %format!("{error:#}"), + "kv_recover request failed; giving up, batches lost" + ); + return None; + } + let delay = self.backoff(attempt); + tracing::warn!( + worker_id, + dp_rank, + attempt = attempt + 1, + retry_in_ms = delay.as_millis() as u64, + error = %format!("{error:#}"), + "kv_recover request failed; retrying" + ); + tokio::select! { + _ = cancel.cancelled() => return None, + _ = tokio::time::sleep(delay) => {} + } + } + None + } + + async fn fetch_once( + &self, + url: &str, + start_seq: u64, + end_seq: u64, + ) -> anyhow::Result { + let _permit = self + .gate + .acquire() + .await + .expect("kv_recover gate is never closed"); + let response = self + .http + .get(url) + .query(&[("start", start_seq), ("end", end_seq)]) + .send() + .await?; + anyhow::ensure!( + response.status().is_success(), + "kv_recover returned status {}", + response.status() + ); + Ok(response.json().await?) + } + + /// 2 s, 4 s, 8 s, ... capped at 16x the base. + fn backoff(&self, attempt: u32) -> Duration { + self.backoff_base * (1u32 << attempt.min(4)) + } +} + +/// What a listener must do with one recovery response. +#[derive(Debug, Default)] +pub(super) struct RecoveryPlan { + /// Drop this `(worker, dp_rank)`'s indexed state before applying. + pub reset_dp_rank: bool, + /// Events to apply, already rewritten to the consumer-assigned identity. + pub events: Vec, + /// New watermark; `None` leaves it unchanged. + pub resume_at: Option, +} + +pub(super) fn plan_recovery( + response: WorkerKvQueryResponse, + worker_id: WorkerId, + dp_rank: u32, +) -> RecoveryPlan { + let reassign = |mut events: Vec| { + for event in &mut events { + event.worker_id = worker_id; + event.event.dp_rank = dp_rank; + } + events + }; + match response { + WorkerKvQueryResponse::Events { + events, + last_event_id, + } => RecoveryPlan { + reset_dp_rank: false, + events: reassign(events), + resume_at: Some(last_event_id), + }, + WorkerKvQueryResponse::TreeDump { + events, + last_event_id, + reset_scope, + } => { + // A domain-scoped snapshot does not describe the whole rank, so it + // must not wipe it; applying it on top stays at-least-once. + let reset_dp_rank = matches!(reset_scope, ResetScope::All); + if !reset_dp_rank { + tracing::warn!( + worker_id, + dp_rank, + ?reset_scope, + "kv_recover TreeDump is domain-scoped; applying without a reset" + ); + } + RecoveryPlan { + reset_dp_rank, + events: reassign(events), + resume_at: Some(last_event_id), + } + } + WorkerKvQueryResponse::StateAgentRecovery { response, .. } => { + plan_recovery(*response, worker_id, dp_rank) + } + WorkerKvQueryResponse::TooNew { + newest_available, .. + } => { + tracing::debug!( + worker_id, + dp_rank, + newest_available, + "kv_recover returned TooNew; consumer is ahead, no-op" + ); + RecoveryPlan::default() + } + other => { + tracing::warn!( + worker_id, + dp_rank, + response = ?other, + "kv_recover returned no recoverable state; no-op" + ); + RecoveryPlan::default() + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::protocols::{KvCacheEventData, StorageTier}; + use crate::services::indexer::backend::test_util::store_event; + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + /// A real `/kv_recover` Events response captured from a prod + /// zai-org/GLM-5.3 engine (2026-09-23), trimmed to two events. + const CAPTURED_EVENTS_RESPONSE: &str = r#"{"Events": {"events": [{"worker_id": 0, "storage_tier": "device", "event": {"event_id": 22636960, "data": {"removed": {"block_hashes": [18405780835976372970]}}, "dp_rank": 0}}, {"worker_id": 0, "storage_tier": "device", "event": {"event_id": 22636961, "data": {"stored": {"parent_hash": 270656378982115943, "start_position": null, "blocks": [{"block_hash": 2733357436102908004, "tokens_hash": 2825904518766229643, "mm_extra_info": null}, {"block_hash": 2451106649480522483, "tokens_hash": 10358785732907640880, "mm_extra_info": null}]}}, "dp_rank": 0}}], "last_event_id": 22637089}}"#; + + #[test] + fn captured_engine_response_plans_events_under_consumer_identity() { + let response: WorkerKvQueryResponse = + serde_json::from_str(CAPTURED_EVENTS_RESPONSE).unwrap(); + let plan = plan_recovery(response, 7, 1); + + assert!(!plan.reset_dp_rank); + assert_eq!(plan.resume_at, Some(22637089)); + assert_eq!(plan.events.len(), 2); + assert!( + plan.events + .iter() + .all(|e| e.worker_id == 7 && e.event.dp_rank == 1) + ); + let KvCacheEventData::Stored(stored) = &plan.events[1].event.data else { + panic!("second captured event is a store"); + }; + assert_eq!(stored.blocks.len(), 2); + } + + #[test] + fn legacy_tree_dump_resets_the_rank() { + let response: WorkerKvQueryResponse = serde_json::from_value(serde_json::json!({ + "TreeDump": { + "events": [store_event(999, 3, 0, &[], &[22], StorageTier::Device)], + "last_event_id": 10 + } + })) + .unwrap(); + let plan = plan_recovery(response, 7, 0); + + assert!(plan.reset_dp_rank); + assert_eq!(plan.resume_at, Some(10)); + assert_eq!(plan.events[0].worker_id, 7); + } + + #[test] + fn too_new_leaves_the_watermark_alone() { + let plan = plan_recovery( + WorkerKvQueryResponse::TooNew { + requested_start: Some(4), + requested_end: Some(9), + newest_available: 3, + }, + 7, + 0, + ); + assert!(!plan.reset_dp_rank); + assert!(plan.events.is_empty()); + assert_eq!(plan.resume_at, None); + } + + /// Serve `/kv_recover` on a local port: the first `failures` requests get + /// a 500, the rest the captured Events body. + async fn flaky_recover_server(failures: usize) -> (String, Arc) { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let endpoint = format!("http://{}", listener.local_addr().unwrap()); + let requests = Arc::new(AtomicUsize::new(0)); + let seen = requests.clone(); + tokio::spawn(async move { + loop { + let (mut stream, _) = listener.accept().await.unwrap(); + let n = seen.fetch_add(1, Ordering::SeqCst); + let mut buf = [0u8; 4096]; + let _ = stream.read(&mut buf).await; + let (status, body) = if n < failures { + ("500 Internal Server Error", "boom") + } else { + ("200 OK", CAPTURED_EVENTS_RESPONSE) + }; + let reply = format!( + "HTTP/1.1 {status}\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}", + body.len() + ); + let _ = stream.write_all(reply.as_bytes()).await; + } + }); + (endpoint, requests) + } + + fn fast_client(timeout_s: u64) -> KvRecoverClient { + KvRecoverClient::with_backoff( + KvRecoverSettings { + timeout_s, + concurrency: 1, + }, + Duration::from_millis(10), + ) + .unwrap() + } + + #[tokio::test] + async fn failed_downloads_are_retried_until_one_succeeds() { + let (endpoint, requests) = flaky_recover_server(2).await; + let response = fast_client(5) + .fetch(&endpoint, 5, 9, 7, 0, &CancellationToken::new()) + .await; + + assert!(matches!( + response, + Some(WorkerKvQueryResponse::Events { .. }) + )); + assert_eq!(requests.load(Ordering::SeqCst), 3); + } + + #[tokio::test] + async fn gap_is_given_up_after_the_last_attempt() { + let (endpoint, requests) = flaky_recover_server(usize::MAX).await; + let response = fast_client(5) + .fetch(&endpoint, 5, 9, 7, 0, &CancellationToken::new()) + .await; + + assert!(response.is_none()); + assert_eq!(requests.load(Ordering::SeqCst), RECOVER_ATTEMPTS as usize); + } + + #[tokio::test] + async fn a_download_slower_than_the_timeout_fails_the_attempt() { + // Accepts and never answers. + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let endpoint = format!("http://{}", listener.local_addr().unwrap()); + let _hold = tokio::spawn(async move { + let mut open = Vec::new(); + loop { + let (stream, _) = listener.accept().await.unwrap(); + open.push(stream); + } + }); + + let started = std::time::Instant::now(); + let response = fast_client(1) + .fetch(&endpoint, 5, 9, 7, 0, &CancellationToken::new()) + .await; + + assert!(response.is_none()); + let elapsed = started.elapsed(); + assert!( + elapsed >= Duration::from_secs(RECOVER_ATTEMPTS as u64) + && elapsed < Duration::from_secs(10), + "each attempt must end at the 1 s timeout, took {elapsed:?}" + ); + } + + #[test] + fn default_settings_match_the_documented_flags() { + let settings = KvRecoverSettings::default(); + assert_eq!(settings.timeout_s, 120); + assert_eq!(settings.concurrency, 8); + let client = KvRecoverClient::new(settings).unwrap(); + assert_eq!(client.backoff(0), Duration::from_secs(2)); + assert_eq!(client.backoff(1), Duration::from_secs(4)); + assert_eq!(client.backoff(9), Duration::from_secs(32)); + } +} diff --git a/lib/kv-router/src/services/indexer/listener.rs b/lib/kv-router/src/services/indexer/listener.rs index 76fa8bde4ad9..8ef8604e5fa3 100644 --- a/lib/kv-router/src/services/indexer/listener.rs +++ b/lib/kv-router/src/services/indexer/listener.rs @@ -13,7 +13,8 @@ use crate::zmq_wire::{ZmqEventNormalizer, decode_event_batch}; use super::backend::Indexer; use super::block_size::BlockSizeGuard; -use super::registry::ListenerRecord; +use super::kv_recover::plan_recovery; +use super::registry::{ListenerRecord, RecoverTarget}; use crate::services::common::zmq::{ MultipartMessage, SharedSocket, connect_dealer_socket, connect_sub_socket, recv_multipart, send_multipart, @@ -152,6 +153,7 @@ struct ListenerLoop { cancel: CancellationToken, live_socket: SharedSocket, replay_socket: Option, + recover: Option, watermark: Arc, normalizer: ZmqEventNormalizer, block_size_guard: BlockSizeGuard, @@ -168,6 +170,7 @@ impl ListenerLoop { cancel: CancellationToken, live_socket: SharedSocket, replay_socket: Option, + recover: Option, watermark: Arc, ) -> Self { Self { @@ -177,6 +180,7 @@ impl ListenerLoop { cancel, live_socket, replay_socket, + recover, watermark, normalizer: ZmqEventNormalizer::new(block_size), block_size_guard: BlockSizeGuard::new(block_size), @@ -306,6 +310,74 @@ impl ListenerLoop { Ok(replayed) } + /// Recover `[start_seq, end_seq)` over HTTP `/kv_recover` when the worker + /// serves it, else over the ZMQ replay socket. + async fn recover_gap(&mut self, start_seq: u64, end_seq: u64) -> Result { + match self.recover.clone() { + Some(target) => self.recover_over_http(&target, start_seq, end_seq).await, + None => self.replay_gap(start_seq, end_seq).await, + } + } + + async fn recover_over_http( + &mut self, + target: &RecoverTarget, + start_seq: u64, + end_seq: u64, + ) -> Result { + let (worker_id, dp_rank) = (self.worker_id, self.dp_rank); + tracing::info!( + worker_id, + dp_rank, + start_seq, + end_seq, + "Requesting recovery from worker via /kv_recover" + ); + let Some(response) = target + .client + .fetch( + &target.endpoint, + start_seq, + end_seq, + worker_id, + dp_rank, + &self.cancel, + ) + .await + else { + return Ok(0); + }; + + let plan = plan_recovery(response, worker_id, dp_rank); + if plan.reset_dp_rank { + self.indexer.remove_worker_dp_rank(worker_id, dp_rank).await; + } + let mut applied = 0u64; + for event in plan.events { + self.indexer + .apply_event_routed(event) + .await + .map_err(|error| { + format!( + "failed to apply recovered event for worker {worker_id} dp_rank {dp_rank}: {error}" + ) + })?; + applied += 1; + } + if let Some(last_event_id) = plan.resume_at { + self.watermark.store(last_event_id, Ordering::Release); + } + tracing::info!( + worker_id, + dp_rank, + applied, + reset = plan.reset_dp_rank, + resume_at = ?plan.resume_at, + "Recovery via /kv_recover complete" + ); + Ok(applied) + } + async fn handle_gap(&mut self, seq: u64) -> Result<(), String> { match self.cursor().observe(seq) { CursorObservation::Initial { got } if got > 0 => { @@ -316,7 +388,7 @@ impl ListenerLoop { got, "Gap detected: expected seq 0, got {got}" ); - self.replay_gap(0, got).await?; + self.recover_gap(0, got).await?; } CursorObservation::Gap { expected, got } => { tracing::warn!( @@ -326,7 +398,7 @@ impl ListenerLoop { got, "Gap detected: expected seq {expected}, got {got}" ); - self.replay_gap(expected, got).await?; + self.recover_gap(expected, got).await?; } CursorObservation::Initial { .. } | CursorObservation::Contiguous { .. } @@ -530,6 +602,7 @@ async fn run_listener( cancel, socket, replay_socket, + record.recover_target(), watermark, ) .run() @@ -634,6 +707,7 @@ mod tests { cancel.clone(), live_socket, Some(replay_socket.clone()), + None, watermark.clone(), ); let replayed = tokio::time::timeout(Duration::from_secs(5), listener.replay_gap(0, 2)) diff --git a/lib/kv-router/src/services/indexer/mod.rs b/lib/kv-router/src/services/indexer/mod.rs index 37c455a91bfd..d96a3c1b8fa2 100644 --- a/lib/kv-router/src/services/indexer/mod.rs +++ b/lib/kv-router/src/services/indexer/mod.rs @@ -31,6 +31,7 @@ pub mod backend; mod block_size; +pub mod kv_recover; pub mod listener; pub mod logging; pub mod metrics; @@ -48,6 +49,7 @@ use tokio_util::sync::CancellationToken; use crate::config::min_initial_workers_from_env; use crate::services::common::zmq::validate_endpoint as validate_zmq_endpoint; use axum::http::header::HeaderName; +use kv_recover::KvRecoverSettings; use logging::AccessLogSink; use registry::WorkerRegistry; use server::{AppState, create_router}; @@ -63,6 +65,8 @@ pub struct IndexerConfig { pub access_log: Option, pub trace_id_header: HeaderName, pub access_log_local_time: bool, + /// Timeout and concurrency of `/kv_recover` gap-recovery downloads. + pub kv_recover: KvRecoverSettings, } pub(super) fn validate_listener_endpoints( @@ -78,6 +82,23 @@ pub(super) fn validate_listener_endpoints( Ok(()) } +/// A `recover_endpoint` is the base URL queried at `/kv_recover`: an +/// absolute `http`/`https` URL with a host. +pub(super) fn validate_recover_endpoint(endpoint: &str) -> anyhow::Result<()> { + let url = reqwest::Url::parse(endpoint) + .map_err(|error| anyhow::anyhow!("invalid recover endpoint `{endpoint}`: {error}"))?; + anyhow::ensure!( + matches!(url.scheme(), "http" | "https"), + "invalid recover endpoint `{endpoint}`: scheme must be http or https, got `{}`", + url.scheme() + ); + anyhow::ensure!( + url.host().is_some(), + "invalid recover endpoint `{endpoint}`: missing host" + ); + Ok(()) +} + pub fn parse_workers(s: &str) -> anyhow::Result> { let mut workers = Vec::new(); @@ -142,7 +163,8 @@ pub async fn run_server(config: IndexerConfig) -> anyhow::Result<()> { "Starting standalone KV cache indexer (HTTP-only mode)" ); - let mut state = AppState::new_with_cancel_token(config.threads, cancel_token.clone())?; + let mut state = + AppState::new_with_cancel_token(config.threads, cancel_token.clone(), config.kv_recover)?; state.access_log_sink = match config.access_log { Some(ref path) => { let s = AccessLogSink::new( @@ -278,6 +300,14 @@ mod tests { validate_zmq_endpoint("ipc:///tmp/dynamo.sock").unwrap(); } + #[test] + fn recover_endpoint_must_be_an_http_url_with_a_host() { + validate_recover_endpoint("http://10.0.0.1:5558").unwrap(); + validate_recover_endpoint("https://engine.local").unwrap(); + assert!(validate_recover_endpoint("tcp://10.0.0.1:5558").is_err()); + assert!(validate_recover_endpoint("10.0.0.1:5558").is_err()); + } + #[test] fn test_validate_zmq_endpoint_rejects_invalid_values() { assert!(validate_zmq_endpoint("tcp://host").is_err()); diff --git a/lib/kv-router/src/services/indexer/registry.rs b/lib/kv-router/src/services/indexer/registry.rs index 15fa98476cbd..8fbd42b4073b 100644 --- a/lib/kv-router/src/services/indexer/registry.rs +++ b/lib/kv-router/src/services/indexer/registry.rs @@ -19,6 +19,7 @@ use crate::indexer::KvIndexerMetrics; use crate::protocols::WorkerId; use super::backend::{Indexer, create_indexer_with_metrics}; +use super::kv_recover::{KvRecoverClient, KvRecoverSettings}; use super::listener::spawn_zmq_listener; pub struct IndexerEntry { @@ -141,9 +142,25 @@ struct ListenerRuntime { generation: u64, } +/// Where and how a listener fetches `/kv_recover` on a gap. +#[derive(Clone)] +pub struct RecoverTarget { + pub endpoint: String, + pub client: Arc, +} + +/// Optional registration fields beyond upstream's `register` arguments. +#[derive(Debug, Clone, Default)] +pub struct ListenerExtras { + /// HTTP base URL of the worker's `GET /kv_recover`. Takes precedence over + /// the ZMQ `replay_endpoint` for gap recovery when both are set. + pub recover_endpoint: Option, +} + pub struct ListenerRecord { endpoint: String, replay_endpoint: Option, + recover: Option, block_size: u32, indexer: Indexer, watermark: Arc, @@ -154,6 +171,7 @@ impl ListenerRecord { fn new( endpoint: String, replay_endpoint: Option, + recover: Option, block_size: u32, indexer: Indexer, watermark: Arc, @@ -161,6 +179,7 @@ impl ListenerRecord { Self { endpoint, replay_endpoint, + recover, block_size, indexer, watermark, @@ -181,6 +200,10 @@ impl ListenerRecord { self.replay_endpoint.as_deref() } + pub(super) fn recover_target(&self) -> Option { + self.recover.clone() + } + pub(super) fn block_size(&self) -> u32 { self.block_size } @@ -321,6 +344,7 @@ pub struct WorkerRegistry { ready_rx: watch::Receiver, root_cancel_token: CancellationToken, retain_empty_indexers: bool, + kv_recover: Arc, } impl WorkerRegistry { @@ -369,9 +393,20 @@ impl WorkerRegistry { ready_rx, root_cancel_token, retain_empty_indexers: false, + kv_recover: Arc::new( + KvRecoverClient::new(KvRecoverSettings::default()) + .expect("default kv_recover client builds"), + ), } } + /// Replace the `/kv_recover` timeout and concurrency gate shared by every + /// listener registered afterwards. + pub fn with_kv_recover(mut self, settings: KvRecoverSettings) -> Result { + self.kv_recover = Arc::new(KvRecoverClient::new(settings)?); + Ok(self) + } + #[cfg(feature = "standalone-selection")] pub(crate) fn with_retained_indexers(mut self) -> Self { // Selection entries and their schedulers retain these indexers across worker updates. @@ -424,6 +459,32 @@ impl WorkerRegistry { routing_group: String, block_size: u32, replay_endpoint: Option, + ) -> Result<()> { + self.register_with_extras( + instance_id, + endpoint, + dp_rank, + model_name, + routing_group, + block_size, + replay_endpoint, + ListenerExtras::default(), + ) + .await + } + + /// [`Self::register`] plus the DeepInfra-only registration fields. + #[expect(clippy::too_many_arguments)] + pub async fn register_with_extras( + &self, + instance_id: WorkerId, + endpoint: String, + dp_rank: u32, + model_name: String, + routing_group: String, + block_size: u32, + replay_endpoint: Option, + extras: ListenerExtras, ) -> Result<()> { let key = RoutingPartitionId::new(model_name, routing_group); let registration = self.indexer_lifecycle.lock().await; @@ -479,9 +540,14 @@ impl WorkerRegistry { .or_insert_with(|| Arc::new(AtomicU64::new(u64::MAX))) .clone(); + let recover = extras.recover_endpoint.map(|endpoint| RecoverTarget { + endpoint, + client: self.kv_recover.clone(), + }); let record = Arc::new(ListenerRecord::new( endpoint, replay_endpoint, + recover, bs, indexer, watermark, diff --git a/lib/kv-router/src/services/indexer/server.rs b/lib/kv-router/src/services/indexer/server.rs index b07ee2f07405..3fee668d3804 100644 --- a/lib/kv-router/src/services/indexer/server.rs +++ b/lib/kv-router/src/services/indexer/server.rs @@ -22,7 +22,8 @@ use crate::protocols::{BlockHashOptions, LocalBlockHash, WorkerId, compute_block use crate::services::overlap::{MooncakeOverlapSummary, build_mooncake_overlap_summaries}; use super::backend::Indexer; -use super::registry::{ListenerControlError, WorkerRegistry}; +use super::kv_recover::KvRecoverSettings; +use super::registry::{ListenerControlError, ListenerExtras, WorkerRegistry}; /// We need to fit one million tokens as JSON text, this should do it. const QUERY_REQUEST_BODY_LIMIT_BYTES: usize = 8 * 1024 * 1024; @@ -46,12 +47,17 @@ pub struct AppState { impl AppState { pub fn new(indexer_threads: usize) -> anyhow::Result { - Self::new_with_cancel_token(indexer_threads, CancellationToken::new()) + Self::new_with_cancel_token( + indexer_threads, + CancellationToken::new(), + KvRecoverSettings::default(), + ) } pub(super) fn new_with_cancel_token( indexer_threads: usize, root_cancel_token: CancellationToken, + kv_recover: KvRecoverSettings, ) -> anyhow::Result { #[cfg(feature = "metrics")] { @@ -59,11 +65,14 @@ impl AppState { super::metrics::register(&prom_registry)?; let indexer_metrics = KvIndexerMetrics::new_registered(&prom_registry)?; Ok(Self { - registry: Arc::new(WorkerRegistry::new_with_indexer_metrics_and_cancel_token( - indexer_threads, - indexer_metrics, - root_cancel_token, - )), + registry: Arc::new( + WorkerRegistry::new_with_indexer_metrics_and_cancel_token( + indexer_threads, + indexer_metrics, + root_cancel_token, + ) + .with_kv_recover(kv_recover)?, + ), access_log_sink: None, prom_registry, }) @@ -71,10 +80,10 @@ impl AppState { #[cfg(not(feature = "metrics"))] Ok(Self { - registry: Arc::new(WorkerRegistry::new_with_cancel_token( - indexer_threads, - root_cancel_token, - )), + registry: Arc::new( + WorkerRegistry::new_with_cancel_token(indexer_threads, root_cancel_token) + .with_kv_recover(kv_recover)?, + ), access_log_sink: None, }) } @@ -94,6 +103,11 @@ struct RegisterRequest { dp_rank: Option, #[serde(default)] replay_endpoint: Option, + /// HTTP base URL of the worker's `GET /kv_recover`, e.g. + /// `http://10.0.0.1:5558`. Used for gap recovery instead of + /// `replay_endpoint` when set. + #[serde(default)] + recover_endpoint: Option, /// Optional per-tenant salt (Mooncake RFC #1403 `additionalsalt`). /// Currently accepted but not yet mixed into hashes — engines apply /// their own salt internally. Plumbed for forward compatibility. @@ -163,7 +177,11 @@ async fn register( ) -> Response { let model = req.model_name.clone(); if let Err(error) = - super::validate_listener_endpoints(&req.endpoint, req.replay_endpoint.as_deref()) + super::validate_listener_endpoints(&req.endpoint, req.replay_endpoint.as_deref()).and( + req.recover_endpoint + .as_deref() + .map_or(Ok(()), super::validate_recover_endpoint), + ) { let mut resp = ( StatusCode::BAD_REQUEST, @@ -176,7 +194,7 @@ async fn register( let resp = match state .registry - .register( + .register_with_extras( req.instance_id, req.endpoint, req.dp_rank.unwrap_or(0), @@ -184,6 +202,9 @@ async fn register( req.routing_group, req.block_size, req.replay_endpoint, + ListenerExtras { + recover_endpoint: req.recover_endpoint, + }, ) .await { From 4e55981b33d1697ae0ad908b38994898c2f140ce Mon Sep 17 00:00:00 2001 From: Thach Nguyen Date: Wed, 23 Sep 2026 22:36:43 +0000 Subject: [PATCH 08/14] standalone-indexer: discover engine pods from Kubernetes With --watch-namespace plus --watch-model-name (or a raw --watch-label), the indexer watches the model's engine pods and keeps the registry in sync: a Ready pod is registered, a deleted / unready pod deregistered. kube::runtime's reflector re-lists on 410 Gone, and every change reconciles the desired Ready set against what is registered, so a missed delete is corrected on the next re-list. - The selector is di/model_name=, the label the backend stamps on every engine pod across engine_hash variants. - Each pod registers under routing_group = its engine_hash label (the configured group if unlabeled), so hash-incompatible engine generations never share a tree during a migration; a changed label re-registers the pod under its new group. - --watch-dp-size N registers one listener per data-parallel rank: dp_rank r on tcp://:, recovering from http://:. A pod whose registration fails part-way is deregistered before the retry. - Configuring discovery in a build without the kube-discovery feature is now an error instead of a warning. Discovery config and the label helpers (discovery.rs) are always compiled; the kube client lives only in the feature-gated pod_watcher.rs. The python bindings' kv-indexer feature enables kube-discovery. Ports cd18634d7 (watcher), the model-name watch of d14e05d6d, the engine_hash grouping of b662110cc, and bb6dcfd5c (DP ranks). Claude-Session: https://claude.ai/code/session_01XKZ49SpoymBqDqgGRJYMNa --- Cargo.lock | 3 + lib/bindings/python/Cargo.lock | 3 + lib/bindings/python/Cargo.toml | 2 +- lib/bindings/python/rust/llm/kv.rs | 76 +++ lib/kv-router/Cargo.toml | 7 + .../src/services/indexer/discovery.rs | 88 +++ lib/kv-router/src/services/indexer/mod.rs | 19 + .../src/services/indexer/pod_watcher.rs | 565 ++++++++++++++++++ 8 files changed, 762 insertions(+), 1 deletion(-) create mode 100644 lib/kv-router/src/services/indexer/discovery.rs create mode 100644 lib/kv-router/src/services/indexer/pod_watcher.rs diff --git a/Cargo.lock b/Cargo.lock index 351821513b76..9d232619f727 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2592,7 +2592,10 @@ dependencies = [ "dynamo-truthy", "fastrand", "flume 0.12.0", + "futures", "indexmap 2.14.0", + "k8s-openapi", + "kube", "once_cell", "ordered-float 4.6.0", "parking_lot", diff --git a/lib/bindings/python/Cargo.lock b/lib/bindings/python/Cargo.lock index cb812629cd5a..c47ee5620a1c 100644 --- a/lib/bindings/python/Cargo.lock +++ b/lib/bindings/python/Cargo.lock @@ -2133,7 +2133,10 @@ dependencies = [ "dynamo-truthy", "fastrand", "flume", + "futures", "indexmap 2.14.0", + "k8s-openapi", + "kube", "once_cell", "ordered-float 4.6.0", "parking_lot", diff --git a/lib/bindings/python/Cargo.toml b/lib/bindings/python/Cargo.toml index ff24d9e9218d..8a8a7d852799 100644 --- a/lib/bindings/python/Cargo.toml +++ b/lib/bindings/python/Cargo.toml @@ -33,7 +33,7 @@ custom-policy = [ "dynamo-kv-router/standalone-selection", ] media-ffmpeg = ["dynamo-llm/media-ffmpeg"] -kv-indexer = ["dep:clap", "dep:tracing-subscriber", "dynamo-kv-router/standalone-indexer"] +kv-indexer = ["dep:clap", "dep:tracing-subscriber", "dynamo-kv-router/standalone-indexer", "dynamo-kv-router/kube-discovery"] kv-indexer-metrics = ["kv-indexer", "dynamo-kv-router/metrics"] # Exposes the optional Relay stats/snapshot methods to Python diagnostic builds. ckf-diagnostics = ["dynamo-llm/ckf-diagnostics"] diff --git a/lib/bindings/python/rust/llm/kv.rs b/lib/bindings/python/rust/llm/kv.rs index 78f9e012fc8b..8e7e4a2181a4 100644 --- a/lib/bindings/python/rust/llm/kv.rs +++ b/lib/bindings/python/rust/llm/kv.rs @@ -238,6 +238,79 @@ struct KvIndexerCli { /// indexer. Bounds the load a fleet-wide (re)subscription puts on engines. #[arg(long, default_value_t = indexer::kv_recover::DEFAULT_RECOVER_CONCURRENCY)] recover_concurrency: usize, + + /// Kubernetes namespace to watch for engine pods. Together with + /// --watch-model-name (or --watch-label) this enables pod auto-discovery: + /// subscribe on Ready, unsubscribe on delete. + #[arg(long)] + watch_namespace: Option, + + /// Raw label selector for this model's engine pods, e.g. + /// "engine_hash=d4b7a85131172ca6". Overrides the selector derived from + /// --watch-model-name. + #[arg(long)] + watch_label: Option, + + /// ZMQ KV-event port the discovered engines publish on. + #[arg(long, default_value_t = 5557)] + watch_zmq_port: u16, + + /// HTTP port serving `GET /kv_recover` on the engines, used for gap + /// recovery at `http://:/kv_recover`. + #[arg(long)] + watch_recover_port: Option, + + /// Data-parallel ranks per engine pod (vLLM --data-parallel-size). Rank r + /// publishes on --watch-zmq-port + r and recovers on + /// --watch-recover-port + r. Leaving this at 1 on a DP engine indexes + /// only rank 0's cache. + #[arg(long, default_value_t = 1)] + watch_dp_size: u32, + + /// Model whose engine pods to discover, e.g. "openai/gpt-oss-120b". The + /// watch uses `di/model_name=` unless --watch-label is + /// given; discovered pods register under this name (default: + /// --model-name). + #[arg(long)] + watch_model_name: Option, +} + +#[cfg(feature = "kv-indexer")] +impl KvIndexerCli { + /// Pod-discovery config, `None` unless --watch-namespace is set. + fn kube_discovery(&self) -> anyhow::Result> { + let Some(namespace) = self.watch_namespace.clone() else { + anyhow::ensure!( + self.watch_label.is_none() && self.watch_model_name.is_none(), + "--watch-label/--watch-model-name require --watch-namespace" + ); + return Ok(None); + }; + let label_selector = match (&self.watch_label, &self.watch_model_name) { + (Some(raw), _) => raw.clone(), + (None, Some(model_name)) => indexer::discovery::model_name_label_selector(model_name), + (None, None) => { + anyhow::bail!("--watch-namespace requires --watch-model-name or --watch-label") + } + }; + let block_size = self + .block_size + .ok_or_else(|| anyhow::anyhow!("--block-size is required with --watch-namespace"))?; + anyhow::ensure!(self.watch_dp_size >= 1, "--watch-dp-size must be at least 1"); + Ok(Some(indexer::discovery::KubeDiscoveryConfig { + namespace, + label_selector, + zmq_port: self.watch_zmq_port, + recover_port: self.watch_recover_port, + dp_size: self.watch_dp_size, + model_name: self + .watch_model_name + .clone() + .unwrap_or_else(|| self.model_name.clone()), + routing_group: self.routing_group.clone(), + block_size, + })) + } } pub fn run_kv_indexer_cli(args: I) -> anyhow::Result<()> @@ -258,6 +331,8 @@ where anyhow::anyhow!("invalid --trace-id-header '{}': {e}", cli.trace_id_header) })?; + let kube_discovery = cli.kube_discovery()?; + init_standalone_logging(); let rt = tokio::runtime::Runtime::new()?; @@ -276,6 +351,7 @@ where timeout_s: cli.recover_timeout_secs, concurrency: cli.recover_concurrency, }, + kube_discovery, })) } diff --git a/lib/kv-router/Cargo.toml b/lib/kv-router/Cargo.toml index b1d1f4642bb9..c3dc1cff2452 100644 --- a/lib/kv-router/Cargo.toml +++ b/lib/kv-router/Cargo.toml @@ -21,6 +21,8 @@ profile = [] standalone-indexer = ["dep:axum", "dep:serde_json", "dep:reqwest", "dep:zmq", "dep:chrono", "dep:tracing-appender"] standalone-slot-tracker = ["dep:axum", "dep:serde_json", "dep:zmq"] standalone-selection = ["standalone-indexer", "dep:once_cell"] +# Kubernetes pod auto-discovery for the standalone indexer (pulls in the kube client). +kube-discovery = ["standalone-indexer", "dep:kube", "dep:k8s-openapi", "dep:futures"] [dependencies] # repo @@ -76,6 +78,11 @@ chrono = { workspace = true, optional = true } reqwest = { version = "0.12.24", default-features = false, features = ["json"], optional = true } tracing-appender = { version = "0.2", optional = true } +# kube-discovery (optional) +kube = { workspace = true, optional = true } +k8s-openapi = { workspace = true, optional = true } +futures = { workspace = true, optional = true } + [dev-dependencies] criterion = "0.5" rstest = "0.18.2" diff --git a/lib/kv-router/src/services/indexer/discovery.rs b/lib/kv-router/src/services/indexer/discovery.rs new file mode 100644 index 000000000000..b34713dc7370 --- /dev/null +++ b/lib/kv-router/src/services/indexer/discovery.rs @@ -0,0 +1,88 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Configuration for Kubernetes pod auto-discovery. +//! +//! Plain data, always compiled, so the CLI can build it without the +//! `kube-discovery` feature; only the feature-gated `pod_watcher` consumes it. + +/// Label the backend stamps on every engine pod with the sanitized model name. +/// Stable across GPU configs and engine versions, unlike `engine_hash`, which +/// fragments per GPU config (B200 vs B300 replicas of one model differ). +pub const MODEL_NAME_LABEL: &str = "di/model_name"; + +/// The backend's `di/model_name` label value for `name`: lowercase, then every +/// char not in `[a-z0-9-]` becomes `-`. Must match the backend's +/// `re.sub('[^a-z0-9-]', '-', name.lower())` exactly, or the selector matches +/// zero pods. +pub fn sanitize_model_name_label(name: &str) -> String { + name.to_lowercase() + .chars() + .map(|c| { + if c.is_ascii_lowercase() || c.is_ascii_digit() || c == '-' { + c + } else { + '-' + } + }) + .collect() +} + +/// Selector matching every engine pod of `model_name`, across engine_hash +/// variants. +pub fn model_name_label_selector(model_name: &str) -> String { + format!( + "{MODEL_NAME_LABEL}={}", + sanitize_model_name_label(model_name) + ) +} + +#[derive(Debug, Clone)] +pub struct KubeDiscoveryConfig { + /// Namespace to watch for engine pods. + pub namespace: String, + /// Selector picking out this model's engine pods: derived from the model + /// name, or supplied raw (e.g. `engine_hash=d4b7a85131172ca6`). + pub label_selector: String, + /// KV-event ZMQ port; data-parallel rank `r` publishes on `zmq_port + r`. + pub zmq_port: u16, + /// Port serving `GET /kv_recover`; rank `r` serves on `recover_port + r`. + pub recover_port: Option, + /// Data-parallel ranks per engine pod, each registered as a listener of + /// the pod's instance. + pub dp_size: u32, + /// Model name discovered pods are registered under. + pub model_name: String, + /// Routing group for pods without an `engine_hash` label. + pub routing_group: String, + /// KV cache block size of the discovered engines. + pub block_size: u32, +} + +#[cfg(test)] +mod tests { + use super::*; + + /// The sanitization contract with the backend, verbatim. + #[test] + fn model_name_label_selector_matches_backend_contract() { + assert_eq!( + model_name_label_selector("openai/gpt-oss-120b"), + "di/model_name=openai-gpt-oss-120b" + ); + assert_eq!( + model_name_label_selector("meta-llama/Llama-3.3-70B"), + "di/model_name=meta-llama-llama-3-3-70b" + ); + assert_eq!( + model_name_label_selector("Qwen/Qwen2.5-72B-Instruct"), + "di/model_name=qwen-qwen2-5-72b-instruct" + ); + } + + #[test] + fn sanitize_replaces_every_disallowed_char() { + assert_eq!(sanitize_model_name_label("A_b.c/Dé9-"), "a-b-c-d-9-"); + assert_eq!(sanitize_model_name_label("gpt-oss-20b"), "gpt-oss-20b"); + } +} diff --git a/lib/kv-router/src/services/indexer/mod.rs b/lib/kv-router/src/services/indexer/mod.rs index d96a3c1b8fa2..a31eec043637 100644 --- a/lib/kv-router/src/services/indexer/mod.rs +++ b/lib/kv-router/src/services/indexer/mod.rs @@ -31,10 +31,13 @@ pub mod backend; mod block_size; +pub mod discovery; pub mod kv_recover; pub mod listener; pub mod logging; pub mod metrics; +#[cfg(feature = "kube-discovery")] +mod pod_watcher; pub mod recovery; pub mod registry; pub mod server; @@ -49,6 +52,7 @@ use tokio_util::sync::CancellationToken; use crate::config::min_initial_workers_from_env; use crate::services::common::zmq::validate_endpoint as validate_zmq_endpoint; use axum::http::header::HeaderName; +use discovery::KubeDiscoveryConfig; use kv_recover::KvRecoverSettings; use logging::AccessLogSink; use registry::WorkerRegistry; @@ -67,6 +71,8 @@ pub struct IndexerConfig { pub access_log_local_time: bool, /// Timeout and concurrency of `/kv_recover` gap-recovery downloads. pub kv_recover: KvRecoverSettings, + /// When set, watch Kubernetes and register/deregister engine pods. + pub kube_discovery: Option, } pub(super) fn validate_listener_endpoints( @@ -260,6 +266,19 @@ async fn run_common( wait_for_min_initial_workers(registry, &cancel_token).await?; registry.signal_ready(); + if let Some(kube_config) = config.kube_discovery.clone() { + #[cfg(feature = "kube-discovery")] + pod_watcher::spawn_pod_watcher(kube_config, registry.clone(), cancel_token.clone()); + #[cfg(not(feature = "kube-discovery"))] + { + let _ = kube_config; + anyhow::bail!( + "pod discovery is configured but this binary was built without the \ + `kube-discovery` feature" + ); + } + } + let app = create_router(state); let listener = TcpListener::bind(("0.0.0.0", config.port)).await?; tracing::info!("HTTP server listening on 0.0.0.0:{}", config.port); diff --git a/lib/kv-router/src/services/indexer/pod_watcher.rs b/lib/kv-router/src/services/indexer/pod_watcher.rs new file mode 100644 index 000000000000..e2b1a8cc8a0c --- /dev/null +++ b/lib/kv-router/src/services/indexer/pod_watcher.rs @@ -0,0 +1,565 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Kubernetes pod auto-discovery for the standalone indexer. +//! +//! Watches engine pods in a namespace (filtered by a label selector) and keeps +//! the [`WorkerRegistry`] in sync with them: +//! +//! - a pod that becomes **Ready** is registered (subscribe to its ZMQ KV-event +//! port), and +//! - a pod that is **deleted** (or stops being Ready) is deregistered. +//! +//! Each pod is registered under `routing_group = `, so +//! engine generations that hash KV blocks differently never share a tree +//! during a migration; `/query` fans out over a model's routing groups. +//! +//! This module is the *only* place that imports the Kubernetes client. It +//! depends on [`WorkerRegistry`] and calls nothing but +//! [`WorkerRegistry::register_with_extras`] / [`WorkerRegistry::deregister`], +//! so the indexer core never learns about Kubernetes. +//! +//! ## How resilience is handled +//! +//! We use the high-level `kube::runtime::watcher` + `reflector` machinery. The +//! reflector maintains an in-memory `Store` of the current pods and the watcher +//! transparently reconnects and **re-lists on `410 Gone`** (the "I fell behind +//! the resourceVersion history" case) with exponential backoff. On every change +//! we recompute the desired set of Ready pods from the Store and reconcile it +//! against what we have registered — so a missed delete during a disconnect is +//! corrected on the next re-list rather than leaking a subscription. +//! +//! ## Required RBAC +//! +//! The indexer's ServiceAccount needs `get`, `list`, and `watch` on `pods` in +//! the target namespace: +//! +//! ```yaml +//! apiVersion: rbac.authorization.k8s.io/v1 +//! kind: Role +//! metadata: +//! name: kv-indexer-pod-watcher +//! namespace: +//! rules: +//! - apiGroups: [""] +//! resources: ["pods"] +//! verbs: ["get", "list", "watch"] +//! --- +//! apiVersion: rbac.authorization.k8s.io/v1 +//! kind: RoleBinding +//! metadata: +//! name: kv-indexer-pod-watcher +//! namespace: +//! subjects: +//! - kind: ServiceAccount +//! name: +//! namespace: +//! roleRef: +//! kind: Role +//! name: kv-indexer-pod-watcher +//! apiGroup: rbac.authorization.k8s.io +//! ``` + +use std::collections::HashMap; +use std::sync::Arc; + +use futures::StreamExt; +use k8s_openapi::api::core::v1::Pod; +use kube::runtime::{WatchStreamExt, reflector, watcher}; +use kube::{Api, Client}; +use tokio::sync::Notify; +use tokio_util::sync::CancellationToken; +use xxhash_rust::xxh3; + +use crate::protocols::WorkerId; + +use super::discovery::KubeDiscoveryConfig; +use super::registry::{ListenerExtras, WorkerRegistry}; + +/// Label the backend stamps on every engine pod with a hash of the exact +/// deployment (image, env, args). Pods with different values may hash KV +/// blocks differently, so each value gets its own tree via `routing_group`. +const ENGINE_HASH_LABEL: &str = "engine_hash"; + +/// A pod we have an active subscription for, keyed by pod name. +struct Subscribed { + instance_id: WorkerId, + ip: String, + routing_group: String, +} + +/// A Ready engine pod as extracted from the watch store. +#[derive(Debug, PartialEq)] +struct DiscoveredPod { + name: String, + ip: String, + /// Value of the pod's `engine_hash` label, when present. + engine_hash: Option, +} + +impl DiscoveredPod { + /// The routing group this pod's tree lives under: its `engine_hash`, so + /// hash-incompatible engine generations never share a tree during + /// migrations. Falls back to the configured group for unlabeled pods. + fn routing_group(&self, fallback: &str) -> String { + self.engine_hash + .clone() + .unwrap_or_else(|| fallback.to_string()) + } +} + +/// Spawn the pod watcher as a detached background task. +/// +/// Returns immediately; the watch loop runs until `cancel` fires. Any fatal +/// error (e.g. the Kubernetes client cannot be created) is logged and the task +/// exits — it never panics or blocks the caller. +pub fn spawn_pod_watcher( + config: KubeDiscoveryConfig, + registry: Arc, + cancel: CancellationToken, +) { + tokio::spawn(async move { + tracing::info!( + namespace = %config.namespace, + label_selector = %config.label_selector, + zmq_port = config.zmq_port, + dp_size = config.dp_size, + model_name = %config.model_name, + block_size = config.block_size, + "Starting Kubernetes pod watcher" + ); + if let Err(error) = run_pod_watcher(config, registry, cancel).await { + tracing::error!(error = %error, "Pod watcher exited with error"); + } + }); +} + +async fn run_pod_watcher( + config: KubeDiscoveryConfig, + registry: Arc, + cancel: CancellationToken, +) -> anyhow::Result<()> { + // Reads in-cluster config (mounted ServiceAccount token) when running as a + // pod, or the local kubeconfig otherwise. + let client = Client::try_default().await?; + let api: Api = Api::namespaced(client, &config.namespace); + let watcher_config = watcher::Config::default().labels(&config.label_selector); + + // The reflector keeps `reader` populated with the current pod set; the + // watcher feeds it and handles reconnects / 410-relist with backoff. + let (reader, writer) = reflector::store(); + let notify = Arc::new(Notify::new()); + + // Drive the watch stream on its own task. Each change (or error) pokes the + // reconcile loop below; the store itself is updated by the reflector. + let stream_notify = notify.clone(); + let stream_cancel = cancel.clone(); + tokio::spawn(async move { + let stream = reflector(writer, watcher(api, watcher_config)) + .default_backoff() + .touched_objects() + .for_each(move |event| { + if let Err(error) = event { + tracing::warn!(error = %error, "Pod watch stream error (auto-retrying)"); + } + stream_notify.notify_one(); + futures::future::ready(()) + }); + tokio::select! { + _ = stream => {} + _ = stream_cancel.cancelled() => {} + } + }); + + let mut subscribed: HashMap = HashMap::new(); + + // Initial sync. `wait_until_ready` returns once the reflector has committed + // its first full LIST of pods, so the Store is readable. We then reconcile + // once — this is what subscribes to engines that already exist at startup + // (and on every restart), rather than only reacting to later changes. + tokio::select! { + _ = cancel.cancelled() => { + tracing::info!("Pod watcher shutting down before initial sync"); + return Ok(()); + } + result = reader.wait_until_ready() => { + result?; + } + } + reconcile(&reader, &mut subscribed, ®istry, &config).await; + tracing::info!(engines = subscribed.len(), "Initial pod sync complete"); + + // Steady state: reconcile on every subsequent change. + loop { + tokio::select! { + _ = cancel.cancelled() => { + tracing::info!("Pod watcher shutting down"); + return Ok(()); + } + _ = notify.notified() => { + reconcile(&reader, &mut subscribed, ®istry, &config).await; + } + } + } +} + +/// Compare the current set of Ready pods against our subscriptions and apply +/// the difference: register newly-ready pods, deregister vanished ones. +async fn reconcile( + store: &reflector::Store, + subscribed: &mut HashMap, + registry: &WorkerRegistry, + config: &KubeDiscoveryConfig, +) { + // Desired state: every Ready, non-terminating pod, keyed by name. + let mut desired: HashMap = HashMap::new(); + for pod in store.state() { + if let Some(discovered) = ready_pod_endpoint(pod.as_ref()) { + desired.insert(discovered.name.clone(), discovered); + } + } + + // Unsubscribe first: pods that are gone, no longer Ready, or whose + // engine_hash changed (the same pod name must re-register under its new + // routing group, and register() rejects a worker that switches keys in place). + let vanished: Vec = subscribed + .iter() + .filter(|(name, entry)| { + desired + .get(*name) + .is_none_or(|pod| pod.routing_group(&config.routing_group) != entry.routing_group) + }) + .map(|(name, _)| name.clone()) + .collect(); + for name in vanished { + let entry = subscribed + .remove(&name) + .expect("name came from subscribed keys"); + match registry + .deregister(entry.instance_id, &config.model_name, &entry.routing_group) + .await + { + Ok(()) => tracing::info!( + pod = %name, + ip = %entry.ip, + instance_id = entry.instance_id, + routing_group = %entry.routing_group, + "Unsubscribed from engine pod" + ), + // Already gone from the registry — an idempotent no-op. + Err(error) => tracing::debug!( + pod = %name, + error = %error, + "Deregister was a no-op" + ), + } + } + + // Subscribe to pods that are newly Ready (or regrouped above). + for (name, pod) in &desired { + if subscribed.contains_key(name) { + continue; // already subscribed — idempotent no-op + } + + let instance_id = instance_id_for(name); + let ip = &pod.ip; + let routing_group = pod.routing_group(&config.routing_group); + + let mut failed: Option<(u32, anyhow::Error)> = None; + for rank in rank_endpoints(ip, config) { + if let Err(error) = registry + .register_with_extras( + instance_id, + rank.endpoint, + rank.dp_rank, + config.model_name.clone(), + routing_group.clone(), + config.block_size, + None, + ListenerExtras { + recover_endpoint: rank.recover_endpoint, + }, + ) + .await + { + failed = Some((rank.dp_rank, error)); + break; + } + } + + match failed { + None => { + tracing::info!( + pod = %name, + ip = %ip, + instance_id, + routing_group = %routing_group, + dp_size = config.dp_size, + "Subscribed to engine pod" + ); + subscribed.insert( + name.clone(), + Subscribed { + instance_id, + ip: ip.clone(), + routing_group, + }, + ); + } + Some((dp_rank, error)) => { + // Not recorded, so it is retried on the next reconcile. Drop + // any ranks that did register: register() rejects a rank that + // is already present, so a partial instance would never + // converge on retry. + tracing::warn!( + pod = %name, + ip = %ip, + dp_rank, + error = %error, + "Failed to register engine pod; will retry" + ); + if let Err(error) = registry + .deregister(instance_id, &config.model_name, &routing_group) + .await + { + tracing::debug!( + pod = %name, + error = %error, + "Cleanup deregister was a no-op" + ); + } + } + } + } +} + +/// One listener to register for a pod: its data-parallel rank and the +/// per-rank event and recovery endpoints. +#[derive(Debug, PartialEq)] +struct RankEndpoint { + dp_rank: u32, + endpoint: String, + recover_endpoint: Option, +} + +/// Endpoints for every data-parallel rank of a pod. vLLM offsets both the ZMQ +/// KV-event port and the `/kv_recover` port by the rank, so rank `r` of a pod +/// at `ip` publishes on `zmq_port + r` and recovers on `recover_port + r`. +fn rank_endpoints(ip: &str, config: &KubeDiscoveryConfig) -> Vec { + (0..config.dp_size.max(1)) + .map(|dp_rank| RankEndpoint { + dp_rank, + endpoint: format!("tcp://{ip}:{}", u32::from(config.zmq_port) + dp_rank), + recover_endpoint: config + .recover_port + .map(|port| format!("http://{ip}:{}", u32::from(port) + dp_rank)), + }) + .collect() +} + +/// Derive a stable `WorkerId` from the pod name so the same pod always maps to +/// the same registry entry across MODIFIED events. +fn instance_id_for(pod_name: &str) -> WorkerId { + xxh3::xxh3_64(pod_name.as_bytes()) +} + +/// Return the pod's name, IP, and `engine_hash` label if it is Ready, has an +/// IP, and is not terminating; otherwise `None`. +fn ready_pod_endpoint(pod: &Pod) -> Option { + // A pod with a deletion timestamp is shutting down — treat it as gone. + if pod.metadata.deletion_timestamp.is_some() { + return None; + } + let name = pod.metadata.name.clone()?; + let status = pod.status.as_ref()?; + let ip = status.pod_ip.clone()?; + + let ready = status.conditions.as_ref().is_some_and(|conditions| { + conditions + .iter() + .any(|c| c.type_ == "Ready" && c.status == "True") + }); + if !ready { + return None; + } + + let engine_hash = pod + .metadata + .labels + .as_ref() + .and_then(|labels| labels.get(ENGINE_HASH_LABEL)) + .cloned(); + + Some(DiscoveredPod { + name, + ip, + engine_hash, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use k8s_openapi::api::core::v1::{PodCondition, PodStatus}; + use k8s_openapi::apimachinery::pkg::apis::meta::v1::{ObjectMeta, Time}; + use k8s_openapi::chrono::Utc; + + fn pod(name: &str, ip: Option<&str>, ready: Option<&str>, terminating: bool) -> Pod { + pod_with_labels(name, ip, ready, terminating, &[]) + } + + fn pod_with_labels( + name: &str, + ip: Option<&str>, + ready: Option<&str>, + terminating: bool, + labels: &[(&str, &str)], + ) -> Pod { + Pod { + metadata: ObjectMeta { + name: Some(name.to_string()), + deletion_timestamp: terminating.then(|| Time(Utc::now())), + labels: (!labels.is_empty()).then(|| { + labels + .iter() + .map(|(k, v)| (k.to_string(), v.to_string())) + .collect() + }), + ..Default::default() + }, + status: Some(PodStatus { + pod_ip: ip.map(String::from), + conditions: ready.map(|status| { + vec![PodCondition { + type_: "Ready".to_string(), + status: status.to_string(), + ..Default::default() + }] + }), + ..Default::default() + }), + ..Default::default() + } + } + + #[test] + fn ready_pod_with_ip_is_subscribable() { + let p = pod("engine-abc", Some("10.0.0.1"), Some("True"), false); + assert_eq!( + ready_pod_endpoint(&p), + Some(DiscoveredPod { + name: "engine-abc".to_string(), + ip: "10.0.0.1".to_string(), + engine_hash: None, + }) + ); + } + + #[test] + fn engine_hash_label_becomes_routing_group() { + let p = pod_with_labels( + "engine-abc", + Some("10.0.0.1"), + Some("True"), + false, + &[("engine_hash", "975d72d83374a79b"), ("engine", "pytrtllm")], + ); + let discovered = ready_pod_endpoint(&p).unwrap(); + assert_eq!(discovered.engine_hash.as_deref(), Some("975d72d83374a79b")); + assert_eq!(discovered.routing_group("default"), "975d72d83374a79b"); + } + + #[test] + fn missing_engine_hash_falls_back_to_configured_routing_group() { + let p = pod_with_labels( + "engine-abc", + Some("10.0.0.1"), + Some("True"), + false, + &[("engine", "pytrtllm")], + ); + let discovered = ready_pod_endpoint(&p).unwrap(); + assert_eq!(discovered.engine_hash, None); + assert_eq!(discovered.routing_group("default"), "default"); + } + + #[test] + fn not_ready_pod_is_skipped() { + let p = pod("engine-abc", Some("10.0.0.1"), Some("False"), false); + assert_eq!(ready_pod_endpoint(&p), None); + } + + #[test] + fn ready_pod_without_ip_is_skipped() { + let p = pod("engine-abc", None, Some("True"), false); + assert_eq!(ready_pod_endpoint(&p), None); + } + + #[test] + fn terminating_pod_is_skipped_even_if_ready() { + let p = pod("engine-abc", Some("10.0.0.1"), Some("True"), true); + assert_eq!(ready_pod_endpoint(&p), None); + } + + #[test] + fn instance_id_is_stable_and_name_specific() { + assert_eq!(instance_id_for("engine-abc"), instance_id_for("engine-abc")); + assert_ne!(instance_id_for("engine-abc"), instance_id_for("engine-xyz")); + } + + fn discovery(dp_size: u32, recover_port: Option) -> KubeDiscoveryConfig { + KubeDiscoveryConfig { + namespace: "deepinfra".to_string(), + label_selector: "di/model_name=m".to_string(), + zmq_port: 5557, + recover_port, + dp_size, + model_name: "m".to_string(), + routing_group: "default".to_string(), + block_size: 128, + } + } + + #[test] + fn single_rank_engine_registers_rank_zero_on_base_ports() { + assert_eq!( + rank_endpoints("10.0.0.1", &discovery(1, Some(5559))), + vec![RankEndpoint { + dp_rank: 0, + endpoint: "tcp://10.0.0.1:5557".to_string(), + recover_endpoint: Some("http://10.0.0.1:5559".to_string()), + }] + ); + } + + #[test] + fn data_parallel_ranks_offset_event_and_recovery_ports() { + assert_eq!( + rank_endpoints("10.0.0.1", &discovery(2, Some(5559))), + vec![ + RankEndpoint { + dp_rank: 0, + endpoint: "tcp://10.0.0.1:5557".to_string(), + recover_endpoint: Some("http://10.0.0.1:5559".to_string()), + }, + RankEndpoint { + dp_rank: 1, + endpoint: "tcp://10.0.0.1:5558".to_string(), + recover_endpoint: Some("http://10.0.0.1:5560".to_string()), + }, + ] + ); + } + + #[test] + fn no_recover_port_means_no_recovery_endpoint_on_any_rank() { + let ranks = rank_endpoints("10.0.0.1", &discovery(2, None)); + assert_eq!(ranks.len(), 2); + assert!(ranks.iter().all(|r| r.recover_endpoint.is_none())); + } + + #[test] + fn zero_dp_size_still_registers_rank_zero() { + assert_eq!(rank_endpoints("10.0.0.1", &discovery(0, None)).len(), 1); + } +} From e7855ccd52d7b184bc592b12ad49c27aa6270cfb Mon Sep 17 00:00:00 2001 From: Thach Nguyen Date: Wed, 23 Sep 2026 22:38:44 +0000 Subject: [PATCH 09/14] standalone-indexer: carry pod_name through /workers and /query Workers are keyed by instance_id (xxh3 of the pod name), which is one way. deepapi maps /query_by_hash matches back to shards by pod name (kv_shards_from_indexer_response), so keep the raw pod name on the worker record when the pod watcher or a /register caller provides it and surface it in /workers entries and in each /query `instances` entry. v1.5.0's shared MooncakeOverlapSummary has no pod_name; the indexer server wraps it in a flattened InstanceMatch rather than changing the shared type. Purely additive to the response. Adds deepapi_contract_tests.rs: the /query_by_hash request/response shape deepapi depends on. Ports b9a2249c4. Claude-Session: https://claude.ai/code/session_01XKZ49SpoymBqDqgGRJYMNa --- .../indexer/deepapi_contract_tests.rs | 138 ++++++++++++++++++ lib/kv-router/src/services/indexer/mod.rs | 2 + .../src/services/indexer/pod_watcher.rs | 1 + .../src/services/indexer/registry.rs | 24 +++ lib/kv-router/src/services/indexer/server.rs | 51 ++++++- 5 files changed, 210 insertions(+), 6 deletions(-) create mode 100644 lib/kv-router/src/services/indexer/deepapi_contract_tests.rs diff --git a/lib/kv-router/src/services/indexer/deepapi_contract_tests.rs b/lib/kv-router/src/services/indexer/deepapi_contract_tests.rs new file mode 100644 index 000000000000..589fa8969207 --- /dev/null +++ b/lib/kv-router/src/services/indexer/deepapi_contract_tests.rs @@ -0,0 +1,138 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! The HTTP contract DeepInfra's deepapi relies on (`query_kv_indexer`, +//! `kv_shards_from_indexer_response` in deepinfra/backend): it POSTs +//! `{"model_name", "tenant_id": "default", "block_hashes": [i64]}` to +//! `/query_by_hash` and reads `instances[*].{pod_name, longest_matched, gpu}`. + +use std::sync::Arc; + +use axum::body::Body; +use axum::http::{Request, StatusCode, header}; +use tower::ServiceExt; + +use crate::protocols::StorageTier; + +use super::backend::Indexer; +use super::backend::test_util::store_event; +use super::registry::{ListenerExtras, WorkerRegistry}; +use super::server::{AppState, create_router}; + +const BLOCK_SIZE: u32 = 4; + +async fn register_pod( + registry: &WorkerRegistry, + instance_id: u64, + port: u16, + model_name: &str, + routing_group: &str, + pod_name: &str, +) { + registry + .register_with_extras( + instance_id, + format!("tcp://127.0.0.1:{port}"), + 0, + model_name.to_string(), + routing_group.to_string(), + BLOCK_SIZE, + None, + ListenerExtras { + recover_endpoint: None, + pod_name: Some(pod_name.to_string()), + }, + ) + .await + .unwrap(); +} + +async fn store_blocks(indexer: &Indexer, worker_id: u64, blocks: &[u64]) { + indexer + .apply_event_routed(store_event( + worker_id, + 0, + 1, + &[], + blocks, + StorageTier::Device, + )) + .await + .unwrap(); + indexer.dump_events().await.expect("flush indexer"); +} + +fn router(registry: Arc) -> axum::Router { + create_router(Arc::new(AppState { + registry, + access_log_sink: None, + #[cfg(feature = "metrics")] + prom_registry: prometheus::Registry::new(), + })) +} + +async fn post_json( + app: &axum::Router, + uri: &str, + body: serde_json::Value, +) -> (StatusCode, serde_json::Value) { + let response = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri(uri) + .header(header::CONTENT_TYPE, "application/json") + .body(Body::from(body.to_string())) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let bytes = axum::body::to_bytes(response.into_body(), usize::MAX) + .await + .unwrap(); + (status, serde_json::from_slice(&bytes).unwrap()) +} + +#[tokio::test] +async fn query_by_hash_names_each_matched_instance_by_pod() { + let registry = Arc::new(WorkerRegistry::new(1)); + registry.signal_ready(); + register_pod(®istry, 7, 15700, "m", "default", "vllm-m-abcd").await; + let indexer = registry + .get_indexer(&crate::identity::RoutingPartitionId::new("m", "default")) + .unwrap() + .indexer + .clone(); + store_blocks(&indexer, 7, &[11, 12]).await; + + let app = router(registry); + let (status, body) = post_json( + &app, + "/query_by_hash", + serde_json::json!({"model_name": "m", "tenant_id": "default", "block_hashes": [11, 12]}), + ) + .await; + + assert_eq!(status, StatusCode::OK, "{body}"); + let instance = &body["instances"]["7"]; + assert_eq!(instance["pod_name"], "vllm-m-abcd"); + assert_eq!(instance["longest_matched"], 2 * BLOCK_SIZE); + assert_eq!(instance["gpu"], 2 * BLOCK_SIZE); + + let workers = app + .oneshot( + Request::builder() + .uri("/workers") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let bytes = axum::body::to_bytes(workers.into_body(), usize::MAX) + .await + .unwrap(); + let workers: serde_json::Value = serde_json::from_slice(&bytes).unwrap(); + assert_eq!(workers[0]["pod_name"], "vllm-m-abcd"); +} diff --git a/lib/kv-router/src/services/indexer/mod.rs b/lib/kv-router/src/services/indexer/mod.rs index a31eec043637..4b2e3f15875b 100644 --- a/lib/kv-router/src/services/indexer/mod.rs +++ b/lib/kv-router/src/services/indexer/mod.rs @@ -31,6 +31,8 @@ pub mod backend; mod block_size; +#[cfg(test)] +mod deepapi_contract_tests; pub mod discovery; pub mod kv_recover; pub mod listener; diff --git a/lib/kv-router/src/services/indexer/pod_watcher.rs b/lib/kv-router/src/services/indexer/pod_watcher.rs index e2b1a8cc8a0c..25792f7b94f8 100644 --- a/lib/kv-router/src/services/indexer/pod_watcher.rs +++ b/lib/kv-router/src/services/indexer/pod_watcher.rs @@ -278,6 +278,7 @@ async fn reconcile( None, ListenerExtras { recover_endpoint: rank.recover_endpoint, + pod_name: Some(name.clone()), }, ) .await diff --git a/lib/kv-router/src/services/indexer/registry.rs b/lib/kv-router/src/services/indexer/registry.rs index 8fbd42b4073b..5ece712a9ab9 100644 --- a/lib/kv-router/src/services/indexer/registry.rs +++ b/lib/kv-router/src/services/indexer/registry.rs @@ -103,6 +103,10 @@ pub struct ListenerInfo { #[derive(Debug, Clone, Serialize)] pub struct WorkerInfo { instance_id: WorkerId, + /// Pod name this worker was registered under, when the registrant knows + /// it. `instance_id` stays the canonical key. + #[serde(skip_serializing_if = "Option::is_none")] + pod_name: Option, source: WorkerSource, status: ListenerStatus, model_name: String, @@ -155,6 +159,8 @@ pub struct ListenerExtras { /// HTTP base URL of the worker's `GET /kv_recover`. Takes precedence over /// the ZMQ `replay_endpoint` for gap recovery when both are set. pub recover_endpoint: Option, + /// Pod name, surfaced in `/workers` and in `/query` instances. + pub pod_name: Option, } pub struct ListenerRecord { @@ -328,6 +334,7 @@ impl ListenerRecord { pub struct WorkerEntry { key: RoutingPartitionId, + pod_name: Option, listeners: HashMap>, } @@ -560,8 +567,12 @@ impl WorkerRegistry { .entry(instance_id) .or_insert_with(|| WorkerEntry { key: key.clone(), + pod_name: None, listeners: HashMap::new(), }); + if extras.pod_name.is_some() { + entry.pod_name = extras.pod_name; + } entry.listeners.insert(dp_rank, record.clone()); } @@ -802,6 +813,7 @@ impl WorkerRegistry { let status = ListenerStatus::aggregate(listeners.values().map(|info| info.status)); Some(WorkerInfo { instance_id: *entry.key(), + pod_name: worker.pod_name.clone(), source: WorkerSource::Zmq, status, model_name: key.model_name.clone(), @@ -821,6 +833,17 @@ impl WorkerRegistry { self.indexers.get(key) } + /// Known pod names by worker id, for enriching query responses. + pub fn pod_names(&self) -> HashMap { + self.workers + .iter() + .filter_map(|entry| { + let pod_name = entry.value().pod_name.clone()?; + Some((*entry.key(), pod_name)) + }) + .collect() + } + pub fn get_or_create_indexer(&self, key: RoutingPartitionId, block_size: u32) -> Indexer { let entry = self.indexers.entry(key.clone()).or_insert_with(|| { tracing::info!( @@ -1331,6 +1354,7 @@ mod tests { 99u64, WorkerEntry { key, + pod_name: None, listeners: HashMap::new(), }, ); diff --git a/lib/kv-router/src/services/indexer/server.rs b/lib/kv-router/src/services/indexer/server.rs index 3fee668d3804..60ea6073a735 100644 --- a/lib/kv-router/src/services/indexer/server.rs +++ b/lib/kv-router/src/services/indexer/server.rs @@ -108,6 +108,9 @@ struct RegisterRequest { /// `replay_endpoint` when set. #[serde(default)] recover_endpoint: Option, + /// Optional pod name, surfaced in `/workers` and `/query` instances. + #[serde(default)] + pod_name: Option, /// Optional per-tenant salt (Mooncake RFC #1403 `additionalsalt`). /// Currently accepted but not yet mixed into hashes — engines apply /// their own salt internally. Plumbed for forward compatibility. @@ -168,7 +171,17 @@ struct ScoreResponse { scores: HashMap>, frequencies: Vec, /// Per-instance tier breakdown (Mooncake RFC #1403 alignment). - instances: HashMap, + instances: HashMap, +} + +/// One `instances` entry: the shared Mooncake summary plus the worker's pod +/// name when known (deepapi maps matches back to shards by pod name). +#[derive(Serialize)] +struct InstanceMatch { + #[serde(skip_serializing_if = "Option::is_none")] + pod_name: Option, + #[serde(flatten)] + summary: MooncakeOverlapSummary, } async fn register( @@ -204,6 +217,7 @@ async fn register( req.replay_endpoint, ListenerExtras { recover_endpoint: req.recover_endpoint, + pod_name: req.pod_name, }, ) .await @@ -289,7 +303,11 @@ async fn list_workers( /// (Mooncake RFC #1403) shapes from a tiered match result. /// /// All token counts are scaled from blocks → tokens via `block_size`. -fn build_score_response(tiered: &TieredMatchDetails, block_size: u32) -> ScoreResponse { +fn build_score_response( + tiered: &TieredMatchDetails, + block_size: u32, + pod_names: &HashMap, +) -> ScoreResponse { // Flat fields (unchanged) come from the device-tier overlap. let device = &tiered.device.overlap_scores; @@ -303,7 +321,13 @@ fn build_score_response(tiered: &TieredMatchDetails, block_size: u32) -> ScoreRe let instances = build_mooncake_overlap_summaries(tiered, block_size, []) .into_iter() - .map(|(worker_id, summary)| (worker_id.to_string(), summary)) + .map(|(worker_id, summary)| { + let instance = InstanceMatch { + pod_name: pod_names.get(&worker_id).cloned(), + summary, + }; + (worker_id.to_string(), instance) + }) .collect(); ScoreResponse { @@ -319,11 +343,14 @@ async fn run_tiered_query( indexer: &Indexer, block_hashes: Vec, block_size: u32, + pod_names: &HashMap, ) -> (StatusCode, Json) { match indexer.find_tiered_matches(block_hashes).await { Ok(tiered) => ( StatusCode::OK, - Json(serde_json::json!(build_score_response(&tiered, block_size))), + Json(serde_json::json!(build_score_response( + &tiered, block_size, pod_names + ))), ), Err(e) => ( StatusCode::INTERNAL_SERVER_ERROR, @@ -362,7 +389,13 @@ async fn query(State(state): State>, Json(req): Json ..Default::default() }, ); - let (status, json) = run_tiered_query(&indexer, block_hashes, block_size).await; + let (status, json) = run_tiered_query( + &indexer, + block_hashes, + block_size, + &state.registry.pod_names(), + ) + .await; let mut resp = (status, json).into_response(); resp.extensions_mut().insert(AccessLogModel(model)); resp @@ -409,7 +442,13 @@ async fn query_by_hash( .iter() .map(|h| LocalBlockHash(*h as u64)) .collect(); - let (status, json) = run_tiered_query(&indexer, block_hashes, block_size).await; + let (status, json) = run_tiered_query( + &indexer, + block_hashes, + block_size, + &state.registry.pod_names(), + ) + .await; let mut resp = (status, json).into_response(); resp.extensions_mut().insert(AccessLogModel(model)); resp From f1ef11d9a71bcc9929ce2456cb53b25f83fac565 Mon Sep 17 00:00:00 2001 From: Thach Nguyen Date: Wed, 23 Sep 2026 22:41:51 +0000 Subject: [PATCH 10/14] standalone-indexer: answer /query from every routing group of the model The pod watcher gives each engine_hash regime its own tree (routing group). A /query or /query_by_hash that names no routing_group now fans out over every tree of model_name and merges the per-tree responses: worker sets are disjoint, so scores and instances union; frequencies sum element-wise. A probe hashed in one regime scores 0 against the others, so routing converges on pods that can reuse the cache during image migrations. A request that names a routing_group still reads only that tree, as upstream does; the legacy tenant_id stays ignored. deepapi sends only tenant_id "default", so it gets the fan-out. A failing tree is skipped while any tree answers (500 only if all fail); 404 only when the model has no tree. The fan-out lives in model_query.rs and replaces upstream's single-tree run_tiered_query. Upstream's query_by_hash_isolates_routing_groups_and_ignores_legacy_tenant_id is renamed and its no-routing_group case now expects the fan-out. New tests: /query_by_hash across two groups excluding another model, /query by tokens, and deepapi's golden block hashes (plain and LoRA-seeded) from deepinfra/utils.kv_block_hashes. Ports the query half of b662110cc. Claude-Session: https://claude.ai/code/session_01XKZ49SpoymBqDqgGRJYMNa --- .../indexer/deepapi_contract_tests.rs | 108 +++++++++++++ lib/kv-router/src/services/indexer/mod.rs | 1 + .../src/services/indexer/model_query.rs | 85 ++++++++++ .../src/services/indexer/registry.rs | 32 ++++ lib/kv-router/src/services/indexer/server.rs | 153 +++++++----------- .../tests/standalone_indexer_http.rs | 6 +- 6 files changed, 291 insertions(+), 94 deletions(-) create mode 100644 lib/kv-router/src/services/indexer/model_query.rs diff --git a/lib/kv-router/src/services/indexer/deepapi_contract_tests.rs b/lib/kv-router/src/services/indexer/deepapi_contract_tests.rs index 589fa8969207..10198bc2041f 100644 --- a/lib/kv-router/src/services/indexer/deepapi_contract_tests.rs +++ b/lib/kv-router/src/services/indexer/deepapi_contract_tests.rs @@ -136,3 +136,111 @@ async fn query_by_hash_names_each_matched_instance_by_pod() { let workers: serde_json::Value = serde_json::from_slice(&bytes).unwrap(); assert_eq!(workers[0]["pod_name"], "vllm-m-abcd"); } + +fn indexer_of(registry: &WorkerRegistry, model_name: &str, routing_group: &str) -> Indexer { + registry + .get_indexer(&crate::identity::RoutingPartitionId::new( + model_name, + routing_group, + )) + .unwrap() + .indexer + .clone() +} + +/// One model under two routing groups (two engine_hash regimes) plus an +/// unrelated model: deepapi's request names only the model (and +/// `tenant_id: "default"`, which matches no tree), and must be answered from +/// both groups of that model and nothing else. +#[tokio::test] +async fn query_by_hash_fans_out_across_routing_groups_of_the_model() { + let registry = Arc::new(WorkerRegistry::new(1)); + registry.signal_ready(); + register_pod(®istry, 7, 15710, "m", "hash-a", "pod-a").await; + register_pod(®istry, 8, 15711, "m", "hash-b", "pod-b").await; + register_pod(®istry, 9, 15712, "other", "hash-a", "pod-c").await; + store_blocks(&indexer_of(®istry, "m", "hash-a"), 7, &[11, 12]).await; + store_blocks(&indexer_of(®istry, "m", "hash-b"), 8, &[11]).await; + store_blocks(&indexer_of(®istry, "other", "hash-a"), 9, &[11, 12]).await; + + let app = router(registry); + let (status, body) = post_json( + &app, + "/query_by_hash", + serde_json::json!({"model_name": "m", "tenant_id": "default", "block_hashes": [11, 12]}), + ) + .await; + + assert_eq!(status, StatusCode::OK, "{body}"); + let instances = body["instances"].as_object().unwrap(); + assert_eq!(instances.len(), 2, "{body}"); + assert_eq!(instances["7"]["pod_name"], "pod-a"); + assert_eq!(instances["7"]["longest_matched"], 2 * BLOCK_SIZE); + assert_eq!(instances["8"]["pod_name"], "pod-b"); + assert_eq!(instances["8"]["longest_matched"], BLOCK_SIZE); + assert_eq!(body["scores"]["7"]["0"], 2 * BLOCK_SIZE); + + let (status, _) = post_json( + &app, + "/query_by_hash", + serde_json::json!({"model_name": "nonexistent", "block_hashes": [11]}), + ) + .await; + assert_eq!(status, StatusCode::NOT_FOUND); +} + +/// deepapi hashes prompts itself (`deepinfra/utils.kv_block_hashes`: XXH3-64 +/// of little-endian u32 tokens, seed 1337, lora seed 1337 + xxh3(name)). These +/// vectors were produced by that function; `/query` must hash identically or +/// every deepapi probe misses. +#[test] +fn token_hashing_matches_deepapi_golden_vectors() { + use crate::protocols::{BlockHashOptions, compute_block_hash_for_seq}; + let tokens: Vec = (1..=8).collect(); + let plain: Vec = compute_block_hash_for_seq(&tokens, 4, BlockHashOptions::default()) + .into_iter() + .map(|h| h.0) + .collect(); + assert_eq!(plain, [14643705804678351452, 16777012769546811212]); + + let lora: Vec = compute_block_hash_for_seq( + &tokens, + 4, + BlockHashOptions { + lora_name: Some("lora_7"), + ..Default::default() + }, + ) + .into_iter() + .map(|h| h.0) + .collect(); + assert_eq!(lora, [6090315765804930700, 3193399369686574168]); +} + +#[tokio::test] +async fn query_by_tokens_fans_out_across_routing_groups() { + use crate::protocols::{BlockHashOptions, compute_block_hash_for_seq}; + let registry = Arc::new(WorkerRegistry::new(1)); + registry.signal_ready(); + register_pod(®istry, 7, 15713, "m", "hash-a", "pod-a").await; + register_pod(®istry, 8, 15714, "m", "hash-b", "pod-b").await; + let token_ids: Vec = (0..8).collect(); + let hashes: Vec = + compute_block_hash_for_seq(&token_ids, BLOCK_SIZE, BlockHashOptions::default()) + .into_iter() + .map(|h| h.0) + .collect(); + store_blocks(&indexer_of(®istry, "m", "hash-a"), 7, &hashes).await; + store_blocks(&indexer_of(®istry, "m", "hash-b"), 8, &hashes[..1]).await; + + let (status, body) = post_json( + &router(registry), + "/query", + serde_json::json!({"model_name": "m", "token_ids": token_ids}), + ) + .await; + + assert_eq!(status, StatusCode::OK, "{body}"); + assert_eq!(body["instances"]["7"]["longest_matched"], 2 * BLOCK_SIZE); + assert_eq!(body["instances"]["8"]["longest_matched"], BLOCK_SIZE); +} diff --git a/lib/kv-router/src/services/indexer/mod.rs b/lib/kv-router/src/services/indexer/mod.rs index 4b2e3f15875b..a2c950dd388c 100644 --- a/lib/kv-router/src/services/indexer/mod.rs +++ b/lib/kv-router/src/services/indexer/mod.rs @@ -38,6 +38,7 @@ pub mod kv_recover; pub mod listener; pub mod logging; pub mod metrics; +mod model_query; #[cfg(feature = "kube-discovery")] mod pod_watcher; pub mod recovery; diff --git a/lib/kv-router/src/services/indexer/model_query.rs b/lib/kv-router/src/services/indexer/model_query.rs new file mode 100644 index 000000000000..2a5a86dda2bd --- /dev/null +++ b/lib/kv-router/src/services/indexer/model_query.rs @@ -0,0 +1,85 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Model-level queries: one `/query` or `/query_by_hash` answered from every +//! routing group's tree of the model, merged into one response. +//! +//! The pod watcher keys each pod's tree by its `engine_hash`, so a model's +//! routing groups are hash regimes, not caller-visible partitions. A probe +//! hashes in one regime and scores 0 against the others, so the merged +//! response converges on the pods that can reuse the cache. A request that +//! names a `routing_group` reads only that tree (upstream behavior); the +//! legacy `tenant_id` is ignored. + +use std::collections::HashMap; + +use axum::http::StatusCode; + +use crate::identity::RoutingPartitionId; +use crate::protocols::{LocalBlockHash, WorkerId}; + +use super::backend::Indexer; +use super::server::{ScoreResponse, build_score_response}; + +pub(super) struct ModelQueryOutcome { + pub status: StatusCode, + pub body: serde_json::Value, +} + +/// Query every tree in `trees` (non-empty; callers 404 before calling) and +/// merge. `hashes_for(block_size)` supplies each tree's probe hashes. A tree +/// that fails is skipped while any tree answers; 500 only if all fail. +pub(super) async fn query_model( + trees: Vec<(RoutingPartitionId, Indexer, u32)>, + mut hashes_for: impl FnMut(u32) -> Vec, + pod_names: &HashMap, +) -> ModelQueryOutcome { + let mut merged: Option = None; + let mut last_error: Option = None; + + for (key, indexer, block_size) in trees { + match indexer.find_tiered_matches(hashes_for(block_size)).await { + Ok(tiered) => { + let response = build_score_response(&tiered, block_size, pod_names); + match merged.as_mut() { + Some(acc) => merge_score_responses(acc, response), + None => merged = Some(response), + } + } + Err(error) => { + tracing::warn!( + model_name = %key.model_name, + routing_group = %key.routing_group, + %error, + "Routing-group tree query failed; continuing fan-out" + ); + last_error = Some(error.to_string()); + } + } + } + + let (status, body) = match merged { + Some(response) => (StatusCode::OK, serde_json::json!(response)), + None => ( + StatusCode::INTERNAL_SERVER_ERROR, + serde_json::json!({ + "error": last_error.unwrap_or_else(|| "no trees queried".to_string()) + }), + ), + }; + ModelQueryOutcome { status, body } +} + +/// Worker sets are disjoint across trees (a worker registers into exactly one +/// routing group), so `scores` and `instances` merge by union; +/// `frequencies[i]` counts workers holding block `i`, so it sums element-wise. +fn merge_score_responses(acc: &mut ScoreResponse, other: ScoreResponse) { + acc.scores.extend(other.scores); + acc.instances.extend(other.instances); + if acc.frequencies.len() < other.frequencies.len() { + acc.frequencies.resize(other.frequencies.len(), 0); + } + for (slot, count) in acc.frequencies.iter_mut().zip(other.frequencies) { + *slot += count; + } +} diff --git a/lib/kv-router/src/services/indexer/registry.rs b/lib/kv-router/src/services/indexer/registry.rs index 5ece712a9ab9..f36a3eebe4c1 100644 --- a/lib/kv-router/src/services/indexer/registry.rs +++ b/lib/kv-router/src/services/indexer/registry.rs @@ -873,6 +873,38 @@ impl WorkerRegistry { entry.indexer.clone() } + /// The trees a query for `model_name` reads: the named routing group's, + /// or every group's when none is named. + pub fn query_trees( + &self, + model_name: &str, + routing_group: Option<&str>, + ) -> Vec<(RoutingPartitionId, Indexer, u32)> { + let Some(routing_group) = routing_group else { + return self.indexers_for_model(model_name); + }; + let key = RoutingPartitionId::new(model_name, routing_group); + self.indexers + .get(&key) + .map(|entry| vec![(key.clone(), entry.indexer.clone(), entry.block_size)]) + .unwrap_or_default() + } + + /// Every routing group's tree of `model_name`, for model-level queries. + pub fn indexers_for_model(&self, model_name: &str) -> Vec<(RoutingPartitionId, Indexer, u32)> { + self.indexers + .iter() + .filter(|entry| entry.key().model_name == model_name) + .map(|entry| { + ( + entry.key().clone(), + entry.value().indexer.clone(), + entry.value().block_size, + ) + }) + .collect() + } + pub fn all_indexers_with_block_size(&self) -> Vec<(RoutingPartitionId, Indexer, u32)> { self.indexers .iter() diff --git a/lib/kv-router/src/services/indexer/server.rs b/lib/kv-router/src/services/indexer/server.rs index 60ea6073a735..8bbbf4543a2a 100644 --- a/lib/kv-router/src/services/indexer/server.rs +++ b/lib/kv-router/src/services/indexer/server.rs @@ -14,15 +14,15 @@ use prometheus::Encoder; use serde::{Deserialize, Serialize}; use tokio_util::sync::CancellationToken; -use crate::identity::{RoutingPartitionId, default_routing_group}; +use crate::identity::default_routing_group; #[cfg(feature = "metrics")] use crate::indexer::KvIndexerMetrics; use crate::indexer::TieredMatchDetails; use crate::protocols::{BlockHashOptions, LocalBlockHash, WorkerId, compute_block_hash_for_seq}; use crate::services::overlap::{MooncakeOverlapSummary, build_mooncake_overlap_summaries}; -use super::backend::Indexer; use super::kv_recover::KvRecoverSettings; +use super::model_query::query_model; use super::registry::{ListenerControlError, ListenerExtras, WorkerRegistry}; /// We need to fit one million tokens as JSON text, this should do it. @@ -134,8 +134,10 @@ struct UnregisterRequest { struct QueryRequest { token_ids: Vec, model_name: String, - #[serde(default = "default_routing_group")] - routing_group: String, + /// Queries only this group's tree when given; absent, the query fans out + /// over every routing group of `model_name` (see `model_query`). + #[serde(default)] + routing_group: Option, #[serde(default = "default_routing_group", rename = "tenant_id")] _tenant_id: String, #[serde(default)] @@ -149,8 +151,10 @@ struct QueryRequest { struct QueryByHashRequest { block_hashes: Vec, model_name: String, - #[serde(default = "default_routing_group")] - routing_group: String, + /// Queries only this group's tree when given; absent, the query fans out + /// over every routing group of `model_name` (see `model_query`). + #[serde(default)] + routing_group: Option, #[serde(default = "default_routing_group", rename = "tenant_id")] _tenant_id: String, /// Invalid for `/query_by_hash`. Callers must precompute `block_hashes` with the intended @@ -167,17 +171,17 @@ struct QueryByHashRequest { /// RFC #1403 (kvcache-ai/Mooncake#1403): /// `{instance_id: {longest_matched, gpu, dp: {rank: count}, cpu, disk}}`. #[derive(Serialize)] -struct ScoreResponse { - scores: HashMap>, - frequencies: Vec, +pub(super) struct ScoreResponse { + pub(super) scores: HashMap>, + pub(super) frequencies: Vec, /// Per-instance tier breakdown (Mooncake RFC #1403 alignment). - instances: HashMap, + pub(super) instances: HashMap, } /// One `instances` entry: the shared Mooncake summary plus the worker's pod /// name when known (deepapi maps matches back to shards by pod name). #[derive(Serialize)] -struct InstanceMatch { +pub(super) struct InstanceMatch { #[serde(skip_serializing_if = "Option::is_none")] pod_name: Option, #[serde(flatten)] @@ -303,7 +307,7 @@ async fn list_workers( /// (Mooncake RFC #1403) shapes from a tiered match result. /// /// All token counts are scaled from blocks → tokens via `block_size`. -fn build_score_response( +pub(super) fn build_score_response( tiered: &TieredMatchDetails, block_size: u32, pod_names: &HashMap, @@ -337,66 +341,49 @@ fn build_score_response( } } -/// Run a tiered query and serialize the result, returning the appropriate -/// HTTP status. Shared between `/query` and `/query_by_hash`. -async fn run_tiered_query( - indexer: &Indexer, - block_hashes: Vec, - block_size: u32, - pod_names: &HashMap, -) -> (StatusCode, Json) { - match indexer.find_tiered_matches(block_hashes).await { - Ok(tiered) => ( - StatusCode::OK, - Json(serde_json::json!(build_score_response( - &tiered, block_size, pod_names - ))), - ), - Err(e) => ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": e.to_string()})), - ), - } -} - async fn query(State(state): State>, Json(req): Json) -> Response { let model = req.model_name.clone(); - let key = RoutingPartitionId::new(req.model_name, req.routing_group); - let Some(ie) = state.registry.get_indexer(&key) else { - let mut resp = ( - StatusCode::NOT_FOUND, - Json(serde_json::json!({ - "error": format!( - "no indexer for model={} routing_group={}", - key.model_name, key.routing_group - ) - })), - ) - .into_response(); - resp.extensions_mut().insert(AccessLogModel(model)); - return resp; - }; - let block_size = ie.block_size; - let indexer = ie.indexer.clone(); - drop(ie); - - let block_hashes = compute_block_hash_for_seq( - &req.token_ids, - block_size, - BlockHashOptions { - lora_name: req.lora_name.as_deref(), - cache_namespace: req.cache_salt.as_deref(), - ..Default::default() + let trees = state + .registry + .query_trees(&req.model_name, req.routing_group.as_deref()); + if trees.is_empty() { + return model_not_found(model); + } + + // Trees of one model normally share a block size; hash once per size. + let mut hashes_by_block_size: HashMap> = HashMap::new(); + let outcome = query_model( + trees, + |block_size| { + hashes_by_block_size + .entry(block_size) + .or_insert_with(|| { + compute_block_hash_for_seq( + &req.token_ids, + block_size, + BlockHashOptions { + lora_name: req.lora_name.as_deref(), + cache_namespace: req.cache_salt.as_deref(), + ..Default::default() + }, + ) + }) + .clone() }, - ); - let (status, json) = run_tiered_query( - &indexer, - block_hashes, - block_size, &state.registry.pod_names(), ) .await; - let mut resp = (status, json).into_response(); + let mut resp = (outcome.status, Json(outcome.body)).into_response(); + resp.extensions_mut().insert(AccessLogModel(model)); + resp +} + +fn model_not_found(model: String) -> Response { + let mut resp = ( + StatusCode::NOT_FOUND, + Json(serde_json::json!({ "error": format!("no indexer for model={model}") })), + ) + .into_response(); resp.extensions_mut().insert(AccessLogModel(model)); resp } @@ -418,38 +405,20 @@ async fn query_by_hash( return resp; } - let key = RoutingPartitionId::new(req.model_name, req.routing_group); - let Some(ie) = state.registry.get_indexer(&key) else { - let mut resp = ( - StatusCode::NOT_FOUND, - Json(serde_json::json!({ - "error": format!( - "no indexer for model={} routing_group={}", - key.model_name, key.routing_group - ) - })), - ) - .into_response(); - resp.extensions_mut().insert(AccessLogModel(model)); - return resp; - }; - let block_size = ie.block_size; - let indexer = ie.indexer.clone(); - drop(ie); + let trees = state + .registry + .query_trees(&req.model_name, req.routing_group.as_deref()); + if trees.is_empty() { + return model_not_found(model); + } let block_hashes: Vec = req .block_hashes .iter() .map(|h| LocalBlockHash(*h as u64)) .collect(); - let (status, json) = run_tiered_query( - &indexer, - block_hashes, - block_size, - &state.registry.pod_names(), - ) - .await; - let mut resp = (status, json).into_response(); + let outcome = query_model(trees, |_| block_hashes.clone(), &state.registry.pod_names()).await; + let mut resp = (outcome.status, Json(outcome.body)).into_response(); resp.extensions_mut().insert(AccessLogModel(model)); resp } diff --git a/lib/kv-router/tests/standalone_indexer_http.rs b/lib/kv-router/tests/standalone_indexer_http.rs index ba5cfc0546bc..bd1a2d8c8150 100644 --- a/lib/kv-router/tests/standalone_indexer_http.rs +++ b/lib/kv-router/tests/standalone_indexer_http.rs @@ -281,7 +281,7 @@ async fn query_by_hash_returns_per_instance_tier_breakdown() { } #[tokio::test] -async fn query_by_hash_isolates_routing_groups_and_ignores_legacy_tenant_id() { +async fn query_by_hash_reads_the_named_routing_group_or_fans_out() { const BLOCK_SIZE: u32 = 4; const MODEL: &str = "test-model"; let registry = registry_with_events( @@ -315,8 +315,10 @@ async fn query_by_hash_isolates_routing_groups_and_ignores_legacy_tenant_id() { .json() .await .expect("legacy tenant-only response"); + // DeepInfra: without a routing_group the query fans out over every group + // of the model (engine_hash regimes); tenant_id stays ignored. assert_eq!(legacy_only["scores"]["7"]["0"], BLOCK_SIZE); - assert!(legacy_only["scores"].get("8").is_none()); + assert_eq!(legacy_only["scores"]["8"]["0"], BLOCK_SIZE); let explicit_group: serde_json::Value = client .post(format!("{base_url}/query_by_hash")) From cc3a18d72d9e18d28150a7a1ef889c6d7f2dbc08 Mon Sep 17 00:00:00 2001 From: Thach Nguyen Date: Wed, 23 Sep 2026 22:44:15 +0000 Subject: [PATCH 11/14] standalone-indexer: --keep-evictions parks evictions, sweeps under memory pressure The kv-indexer:routing flavor measures "infinite memory, per-shard" hit rates: engine Removed events are parked in a per-listener buffer instead of applied (Cleared is dropped), and a store of a parked block cancels its eviction. Every 15 s, once memory use crosses --evict-memory-threshold (default 0.75) of the cgroup limit, evictions older than --evict-retention-secs (default 1800) are replayed into the tree as ordinary Removed events, so the pod degrades toward reality instead of OOMing. A per-minute status line logs memory fraction, pending and queued entries, and listener count. A replayed eviction is applied only if it is still the block's latest; a /kv_recover TreeDump that resets the rank clears the rank's buffer. The flag is a RegistryOptions field (no process globals): only listeners of a keep-evictions registry carry a buffer. Ports the keep-evictions half of d14e05d6d and c9eabbb46; supersedes 6d16bb8a9 (--ignore-evictions / --merge-shards). Claude-Session: https://claude.ai/code/session_01XKZ49SpoymBqDqgGRJYMNa --- lib/bindings/python/rust/llm/kv.rs | 23 + .../src/services/indexer/evictions.rs | 482 ++++++++++++++++++ .../src/services/indexer/listener.rs | 47 +- lib/kv-router/src/services/indexer/mod.rs | 28 +- .../src/services/indexer/registry.rs | 54 +- lib/kv-router/src/services/indexer/server.rs | 11 +- 6 files changed, 630 insertions(+), 15 deletions(-) create mode 100644 lib/kv-router/src/services/indexer/evictions.rs diff --git a/lib/bindings/python/rust/llm/kv.rs b/lib/bindings/python/rust/llm/kv.rs index 8e7e4a2181a4..adf78cce6593 100644 --- a/lib/bindings/python/rust/llm/kv.rs +++ b/lib/bindings/python/rust/llm/kv.rs @@ -239,6 +239,23 @@ struct KvIndexerCli { #[arg(long, default_value_t = indexer::kv_recover::DEFAULT_RECOVER_CONCURRENCY)] recover_concurrency: usize, + /// Park Removed events in a buffer (and drop Cleared) instead of applying + /// them, approximating infinite memory (the kv-indexer:routing flavor). + /// Buffered evictions older than --evict-retention-secs are replayed into + /// the tree once memory use crosses --evict-memory-threshold of the limit. + #[arg(long, default_value_t = false)] + keep_evictions: bool, + + /// Minimum age (seconds) a parked eviction must reach before the memory + /// sweep may apply it. Only meaningful with --keep-evictions. + #[arg(long, default_value_t = 1800)] + evict_retention_secs: u64, + + /// Fraction of the memory limit (0..1] above which the sweep replays aged + /// evictions. Only meaningful with --keep-evictions. + #[arg(long, default_value_t = 0.75)] + evict_memory_threshold: f64, + /// Kubernetes namespace to watch for engine pods. Together with /// --watch-model-name (or --watch-label) this enables pod auto-discovery: /// subscribe on Ready, unsubscribe on delete. @@ -352,6 +369,12 @@ where concurrency: cli.recover_concurrency, }, kube_discovery, + keep_evictions: cli.keep_evictions.then_some( + indexer::evictions::KeepEvictionsConfig { + retention_s: cli.evict_retention_secs, + memory_threshold: cli.evict_memory_threshold, + }, + ), })) } diff --git a/lib/kv-router/src/services/indexer/evictions.rs b/lib/kv-router/src/services/indexer/evictions.rs new file mode 100644 index 000000000000..940a2b998259 --- /dev/null +++ b/lib/kv-router/src/services/indexer/evictions.rs @@ -0,0 +1,482 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! `--keep-evictions` support: park engine eviction events in a per-listener +//! buffer instead of applying them, then replay the aged ones into the tree +//! through the normal `Removed` apply path once the pod's memory crosses a +//! threshold. The radix tree itself is untouched — from its point of view +//! evictions simply arrive late. +//! +//! Correctness hinges on one rule: a buffered eviction is only replayed if it +//! is still the block's *latest* eviction. The map holds the latest eviction +//! time per hash; the FIFO queue entries each carry the time they were added, +//! and an entry whose time no longer matches the map is dead (the block was +//! re-stored, or re-evicted later — the newer queue entry handles it). + +use std::collections::VecDeque; +use std::sync::{Arc, LazyLock}; +use std::time::{Duration, Instant}; + +use rustc_hash::FxHashMap; +use tokio_util::sync::CancellationToken; + +use crate::protocols::{ + ExternalSequenceBlockHash, KvCacheEvent, KvCacheEventData, KvCacheRemoveData, RouterEvent, + StorageTier, +}; + +use super::registry::WorkerRegistry; + +/// How often the cleanup loop re-checks memory pressure. +const SWEEP_INTERVAL: Duration = Duration::from_secs(15); + +/// Max hashes per synthesized `Removed` event, to bound message size on the +/// indexer's internal apply channels. +const REMOVE_CHUNK: usize = 4096; + +/// Monotonic seconds since process start. Immune to wall-clock jumps, which +/// matters because entry age decides what gets removed from the tree. +pub(super) fn monotonic_secs() -> u64 { + static START: LazyLock = LazyLock::new(Instant::now); + START.elapsed().as_secs() +} + +/// `--keep-evictions` settings. +#[derive(Debug, Clone, Copy)] +pub struct KeepEvictionsConfig { + /// Minimum age a parked eviction must reach before the sweep replays it. + pub retention_s: u64, + /// Fraction of the memory limit, in (0, 1], above which the sweep runs. + pub memory_threshold: f64, +} + +impl KeepEvictionsConfig { + pub fn validate(&self) -> anyhow::Result<()> { + anyhow::ensure!( + self.memory_threshold > 0.0 && self.memory_threshold <= 1.0, + "--evict-memory-threshold must be in (0, 1], got {}", + self.memory_threshold + ); + Ok(()) + } +} + +/// Buffer of eviction events not yet applied to the tree, owned by one +/// listener (one `(worker_id, dp_rank)` event stream). +#[derive(Default)] +pub struct PendingEvictions { + /// Latest eviction per sequence hash: when it happened and on which tier. + pending: FxHashMap, + /// Arrival-ordered `(hash, at_secs, tier)`. Only pushed at the back and + /// popped at the front, so timestamps are non-decreasing front to back. + queue: VecDeque<(u64, u64, StorageTier)>, +} + +impl PendingEvictions { + /// Record an eviction batch instead of applying it. + pub fn buffer(&mut self, hashes: &[ExternalSequenceBlockHash], tier: StorageTier) { + self.buffer_at(hashes, tier, monotonic_secs()); + } + + /// Timestamp-explicit form of [`Self::buffer`]. `at` must be + /// non-decreasing across calls (queue order is arrival order). + fn buffer_at(&mut self, hashes: &[ExternalSequenceBlockHash], tier: StorageTier, at: u64) { + for hash in hashes { + self.pending.insert(hash.0, (at, tier)); + self.queue.push_back((hash.0, at, tier)); + } + } + + /// Cancel pending evictions for blocks the engine just re-stored. Their + /// stale queue entries stay behind and are skipped at drain time. + pub fn cancel(&mut self, hashes: impl IntoIterator) { + for hash in hashes { + self.pending.remove(&hash); + } + } + + /// Drop all state (used when a recovery snapshot replaces the worker's + /// tree state wholesale). + pub fn clear(&mut self) { + self.pending.clear(); + self.queue.clear(); + } + + /// Number of evictions currently pending replay. + pub fn len(&self) -> usize { + self.pending.len() + } + + /// Raw queue length, including stale (cancelled/superseded) entries that + /// still occupy memory until a drain pops them. `queue_len - len` is the + /// dead weight from re-stored blocks. + pub fn queue_len(&self) -> usize { + self.queue.len() + } + + pub fn is_empty(&self) -> bool { + self.pending.is_empty() + } + + /// Pop every queue entry with `at <= cutoff_secs` and return the hashes + /// whose latest eviction that entry is, grouped by storage tier. Entries + /// that were cancelled by a re-store, or superseded by a later eviction + /// of the same hash, are skipped — not stopped at. + pub fn drain_older_than( + &mut self, + cutoff_secs: u64, + ) -> FxHashMap> { + let mut doomed: FxHashMap> = + FxHashMap::default(); + while let Some(&(hash, at, _)) = self.queue.front() { + if at > cutoff_secs { + break; + } + self.queue.pop_front(); + // Only replay if this entry is still the hash's latest eviction. + // The tier stored in `pending` is authoritative (same write as + // the timestamp), so read it from there. + if let Some(&(pending_at, tier)) = self.pending.get(&hash) + && pending_at == at + { + self.pending.remove(&hash); + doomed + .entry(tier) + .or_default() + .push(ExternalSequenceBlockHash(hash)); + } + } + doomed + } +} + +/// Fraction of the memory limit currently in use, best-effort: +/// cgroup v2, then cgroup v1, then RSS against total host memory. +/// `None` if no source could be read (never expected on Linux). +pub fn memory_usage_fraction() -> Option { + fn read_u64(path: &str) -> Option { + std::fs::read_to_string(path).ok()?.trim().parse().ok() + } + // "max" in cgroup v2 (or the v1 sentinel ~2^63) means "no limit set"; + // fall through to the host total in that case. + fn read_limit(path: &str) -> Option { + let raw = std::fs::read_to_string(path).ok()?; + let raw = raw.trim(); + if raw == "max" { + return None; + } + let value: u64 = raw.parse().ok()?; + (value < (1 << 62)).then_some(value) + } + + let usage = read_u64("/sys/fs/cgroup/memory.current") + .or_else(|| read_u64("/sys/fs/cgroup/memory/memory.usage_in_bytes")) + .or_else(vm_rss_bytes)?; + let limit = read_limit("/sys/fs/cgroup/memory.max") + .or_else(|| read_limit("/sys/fs/cgroup/memory/memory.limit_in_bytes")) + .or_else(host_total_bytes)?; + (limit > 0).then(|| usage as f64 / limit as f64) +} + +/// Resident set size from `/proc/self/status` (`VmRSS`, reported in kB). +fn vm_rss_bytes() -> Option { + let status = std::fs::read_to_string("/proc/self/status").ok()?; + parse_kb_line(&status, "VmRSS:") +} + +/// Total host memory from `/proc/meminfo` (`MemTotal`, reported in kB). +fn host_total_bytes() -> Option { + let meminfo = std::fs::read_to_string("/proc/meminfo").ok()?; + parse_kb_line(&meminfo, "MemTotal:") +} + +fn parse_kb_line(text: &str, key: &str) -> Option { + let line = text.lines().find(|l| l.starts_with(key))?; + let kb: u64 = line.split_whitespace().nth(1)?.parse().ok()?; + Some(kb * 1024) +} + +/// Spawn the background sweep: every [`SWEEP_INTERVAL`], if memory usage is +/// at or above `memory_threshold` of the limit, drain evictions older than +/// `retention_secs` from every listener's buffer and replay them into that +/// listener's indexer as ordinary `Removed` events. +pub fn spawn_cleanup_loop( + registry: Arc, + retention_secs: u64, + memory_threshold: f64, + cancel: CancellationToken, +) { + tokio::spawn(async move { + let mut ticker = tokio::time::interval(SWEEP_INTERVAL); + ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + let mut tick_count: u64 = 0; + loop { + tokio::select! { + _ = cancel.cancelled() => return, + _ = ticker.tick() => {} + } + tick_count += 1; + + let usage = memory_usage_fraction(); + + // Once a minute, log buffer/memory state unconditionally — this is + // the only visibility into the buffers below the sweep threshold, + // and the pod's log retention is far too short to rely on + // startup-time lines. + if tick_count % 4 == 1 { + let (mut pending, mut queued) = (0usize, 0usize); + let records = registry.listener_records(); + for (_, _, buffer, _) in &records { + let buffer = buffer.lock(); + pending += buffer.len(); + queued += buffer.queue_len(); + } + tracing::info!( + memory_fraction = usage.map(|u| format!("{u:.3}")), + pending_evictions = pending, + queued_entries = queued, + listeners = records.len(), + "keep-evictions status" + ); + } + + let Some(usage) = usage else { + tracing::warn!("keep-evictions sweep: could not read memory usage; skipping"); + continue; + }; + if usage < memory_threshold { + continue; + } + + let cutoff = monotonic_secs().saturating_sub(retention_secs); + let mut replayed = 0usize; + let mut still_pending = 0usize; + for (worker_id, dp_rank, pending_evictions, indexer) in registry.listener_records() { + // Drain under the lock, apply outside it: applies await on the + // indexer channels and must not block the listener's event path. + let (doomed, remaining) = { + let mut buffer = pending_evictions.lock(); + let doomed = buffer.drain_older_than(cutoff); + (doomed, buffer.len()) + }; + still_pending += remaining; + for (tier, hashes) in doomed { + replayed += hashes.len(); + for chunk in hashes.chunks(REMOVE_CHUNK) { + let event = RouterEvent::with_storage_tier( + worker_id, + KvCacheEvent { + // Synthetic id: the standalone apply path does + // gap-tracking on ZMQ seq numbers, not on + // event_id, so 0 is safe here. + event_id: 0, + data: KvCacheEventData::Removed(KvCacheRemoveData { + block_hashes: chunk.to_vec(), + }), + dp_rank, + }, + tier, + ); + if let Err(error) = indexer.apply_event_routed(event).await { + tracing::warn!( + worker_id, + dp_rank, + %error, + "keep-evictions sweep: replayed eviction was rejected" + ); + } + } + } + } + tracing::info!( + usage = format!("{:.1}%", usage * 100.0), + threshold = format!("{:.1}%", memory_threshold * 100.0), + replayed, + still_pending, + retention_secs, + "keep-evictions sweep: replayed aged evictions into the tree" + ); + } + }); +} + +#[cfg(test)] +mod tests { + use super::*; + + fn hashes(values: &[u64]) -> Vec { + values + .iter() + .copied() + .map(ExternalSequenceBlockHash) + .collect() + } + + fn drained_flat(buffer: &mut PendingEvictions, cutoff: u64) -> Vec { + let mut out: Vec = buffer + .drain_older_than(cutoff) + .into_values() + .flatten() + .map(|h| h.0) + .collect(); + out.sort_unstable(); + out + } + + #[test] + fn drain_respects_cutoff() { + let mut buffer = PendingEvictions::default(); + buffer.buffer_at(&hashes(&[1, 2]), StorageTier::Device, 100); + buffer.buffer_at(&hashes(&[3]), StorageTier::Device, 200); + + // Entries newer than the cutoff stay. + assert!(drained_flat(&mut buffer, 99).is_empty()); + assert_eq!(buffer.len(), 3); + + // Cutoff between the batches: only the older batch drains. + assert_eq!(drained_flat(&mut buffer, 150), vec![1, 2]); + assert_eq!(buffer.len(), 1); + + // Cutoff past everything: the rest drains; a second drain is a no-op. + assert_eq!(drained_flat(&mut buffer, 500), vec![3]); + assert!(buffer.is_empty()); + assert!(drained_flat(&mut buffer, 500).is_empty()); + } + + #[test] + fn restore_cancels_pending_eviction() { + let mut buffer = PendingEvictions::default(); + buffer.buffer(&hashes(&[1, 2, 3]), StorageTier::Device); + buffer.cancel([2]); + + assert_eq!(buffer.len(), 2); + let drained = drained_flat(&mut buffer, monotonic_secs() + 1); + assert_eq!(drained, vec![1, 3], "re-stored block 2 must not be removed"); + // The stale queue entry for 2 must not stall the drain of 3 behind it. + assert!(buffer.is_empty()); + } + + #[test] + fn later_re_eviction_is_not_removed_early() { + let mut buffer = PendingEvictions::default(); + // Evict at t=100, re-store, evict again at t=2000. A drain whose + // cutoff covers only the first eviction must remove NOTHING: the + // block's current eviction is too fresh, and the stale first entry + // must not act on the map's newer state. + buffer.buffer_at(&hashes(&[7]), StorageTier::Device, 100); + buffer.cancel([7]); + buffer.buffer_at(&hashes(&[7]), StorageTier::Device, 2000); + + assert!( + drained_flat(&mut buffer, 1000).is_empty(), + "the fresh re-eviction at t=2000 must not be removed by the stale t=100 entry" + ); + assert_eq!(buffer.len(), 1); + + // Once the second eviction ages past the cutoff, it drains exactly once. + assert_eq!(drained_flat(&mut buffer, 2000), vec![7]); + assert!(buffer.is_empty()); + assert!(buffer.queue.is_empty(), "no leftover queue entries"); + } + + #[test] + fn duplicate_eviction_without_restore_drains_once() { + let mut buffer = PendingEvictions::default(); + // Same hash evicted twice with no store in between (shouldn't happen, + // but must not double-remove): only the latest entry replays. + buffer.buffer_at(&hashes(&[9]), StorageTier::Device, 100); + buffer.buffer_at(&hashes(&[9]), StorageTier::Device, 200); + + assert_eq!(buffer.len(), 1); + assert_eq!(drained_flat(&mut buffer, 500), vec![9]); + assert!(buffer.queue.is_empty()); + } + + #[test] + fn drain_groups_by_storage_tier() { + let mut buffer = PendingEvictions::default(); + buffer.buffer(&hashes(&[1]), StorageTier::Device); + buffer.buffer(&hashes(&[2]), StorageTier::HostPinned); + + let drained = buffer.drain_older_than(monotonic_secs() + 1); + assert_eq!(drained.get(&StorageTier::Device).map(Vec::len), Some(1)); + assert_eq!(drained.get(&StorageTier::HostPinned).map(Vec::len), Some(1)); + } + + #[test] + fn clear_drops_everything() { + let mut buffer = PendingEvictions::default(); + buffer.buffer(&hashes(&[1, 2]), StorageTier::Device); + buffer.clear(); + assert!(buffer.is_empty()); + assert!(drained_flat(&mut buffer, monotonic_secs() + 1).is_empty()); + } + + /// End-to-end shape of the sweep: a stored-then-evicted block stays + /// queryable while parked, and disappears once the drained eviction is + /// replayed through the normal `Removed` apply path. + #[tokio::test] + async fn drained_evictions_replay_into_the_tree() { + use crate::protocols::{LocalBlockHash, WorkerWithDpRank, compute_seq_hash_for_block}; + use crate::services::indexer::backend::{create_indexer, test_util::store_event}; + + let worker = WorkerWithDpRank::new(7, 0); + let indexer = create_indexer(4, 1); + indexer + .apply_event_routed(store_event( + worker.worker_id, + worker.dp_rank, + 0, + &[], + &[11, 12], + StorageTier::Device, + )) + .await + .unwrap(); + let _ = indexer.dump_events().await; // flush in-flight applies + + let sequence = vec![LocalBlockHash(11), LocalBlockHash(12)]; + let scores = indexer.find_matches(sequence.clone()).await.unwrap(); + assert_eq!(scores.scores.get(&worker).copied(), Some(2)); + + // The engine evicts the tail block; the buffer parks it, so the tree + // still matches both blocks. + let tail_seq_hash = ExternalSequenceBlockHash(compute_seq_hash_for_block(&sequence)[1]); + let mut buffer = PendingEvictions::default(); + buffer.buffer_at(&[tail_seq_hash], StorageTier::Device, 100); + let scores = indexer.find_matches(sequence.clone()).await.unwrap(); + assert_eq!(scores.scores.get(&worker).copied(), Some(2)); + + // The sweep drains the aged entry and replays it as a Removed event — + // exactly what spawn_cleanup_loop's body does. + for (tier, hashes) in buffer.drain_older_than(1000) { + let event = RouterEvent::with_storage_tier( + worker.worker_id, + KvCacheEvent { + event_id: 0, + data: KvCacheEventData::Removed(KvCacheRemoveData { + block_hashes: hashes, + }), + dp_rank: worker.dp_rank, + }, + tier, + ); + indexer.apply_event_routed(event).await.unwrap(); + } + let _ = indexer.dump_events().await; + + let scores = indexer.find_matches(sequence).await.unwrap(); + assert_eq!( + scores.scores.get(&worker).copied().unwrap_or(0), + 1, + "only the surviving head block should match after the replayed eviction" + ); + assert!(buffer.is_empty()); + } + + #[test] + fn memory_usage_fraction_reads_something_sane() { + let usage = memory_usage_fraction().expect("some memory source must be readable"); + assert!(usage > 0.0 && usage < 10.0, "got {usage}"); + } +} diff --git a/lib/kv-router/src/services/indexer/listener.rs b/lib/kv-router/src/services/indexer/listener.rs index 8ef8604e5fa3..cea935029f7b 100644 --- a/lib/kv-router/src/services/indexer/listener.rs +++ b/lib/kv-router/src/services/indexer/listener.rs @@ -7,14 +7,14 @@ use std::sync::atomic::{AtomicU64, Ordering}; use tokio::sync::watch; use tokio_util::sync::CancellationToken; -use crate::protocols::{WorkerId, WorkerWithDpRank}; +use crate::protocols::{KvCacheEventData, RouterEvent, WorkerId, WorkerWithDpRank}; use crate::recovery::{CursorObservation, CursorState}; use crate::zmq_wire::{ZmqEventNormalizer, decode_event_batch}; use super::backend::Indexer; use super::block_size::BlockSizeGuard; use super::kv_recover::plan_recovery; -use super::registry::{ListenerRecord, RecoverTarget}; +use super::registry::{ListenerRecord, RecoverTarget, SharedPendingEvictions}; use crate::services::common::zmq::{ MultipartMessage, SharedSocket, connect_dealer_socket, connect_sub_socket, recv_multipart, send_multipart, @@ -154,6 +154,7 @@ struct ListenerLoop { live_socket: SharedSocket, replay_socket: Option, recover: Option, + pending_evictions: Option, watermark: Arc, normalizer: ZmqEventNormalizer, block_size_guard: BlockSizeGuard, @@ -171,6 +172,7 @@ impl ListenerLoop { live_socket: SharedSocket, replay_socket: Option, recover: Option, + pending_evictions: Option, watermark: Arc, ) -> Self { Self { @@ -181,6 +183,7 @@ impl ListenerLoop { live_socket, replay_socket, recover, + pending_evictions, watermark, normalizer: ZmqEventNormalizer::new(block_size), block_size_guard: BlockSizeGuard::new(block_size), @@ -188,6 +191,30 @@ impl ListenerLoop { } } + /// `--keep-evictions`: a store cancels any parked eviction of its blocks + /// and applies; an eviction is parked instead of applied; a clear is + /// dropped. `true` when the event must not reach the tree. + fn park_kept_eviction(&self, event: &RouterEvent) -> bool { + let Some(pending) = self.pending_evictions.as_ref() else { + return false; + }; + match &event.event.data { + KvCacheEventData::Stored(data) => { + pending + .lock() + .cancel(data.blocks.iter().map(|b| b.block_hash.0)); + false + } + KvCacheEventData::Removed(data) => { + pending + .lock() + .buffer(&data.block_hashes, event.storage_tier); + true + } + KvCacheEventData::Cleared => true, + } + } + fn cursor(&self) -> CursorState { cursor_from_watermark(self.watermark.load(Ordering::Acquire)) } @@ -290,6 +317,9 @@ impl ListenerLoop { let router_event = placement_event .into_router_event() .expect("local worker placement must convert to router event"); + if self.park_kept_eviction(&router_event) { + continue; + } indexer .apply_event_routed(router_event) .await @@ -350,10 +380,18 @@ impl ListenerLoop { let plan = plan_recovery(response, worker_id, dp_rank); if plan.reset_dp_rank { + // Parked evictions describe the state being replaced; replaying + // them onto the snapshot could remove blocks it says are live. + if let Some(pending) = self.pending_evictions.as_ref() { + pending.lock().clear(); + } self.indexer.remove_worker_dp_rank(worker_id, dp_rank).await; } let mut applied = 0u64; for event in plan.events { + if self.park_kept_eviction(&event) { + continue; + } self.indexer .apply_event_routed(event) .await @@ -440,6 +478,9 @@ impl ListenerLoop { let router_event = placement_event .into_router_event() .expect("local worker placement must convert to router event"); + if self.park_kept_eviction(&router_event) { + continue; + } self.indexer .apply_event_routed(router_event) .await @@ -603,6 +644,7 @@ async fn run_listener( socket, replay_socket, record.recover_target(), + record.pending_evictions(), watermark, ) .run() @@ -708,6 +750,7 @@ mod tests { live_socket, Some(replay_socket.clone()), None, + None, watermark.clone(), ); let replayed = tokio::time::timeout(Duration::from_secs(5), listener.replay_gap(0, 2)) diff --git a/lib/kv-router/src/services/indexer/mod.rs b/lib/kv-router/src/services/indexer/mod.rs index a2c950dd388c..8820ff6254ac 100644 --- a/lib/kv-router/src/services/indexer/mod.rs +++ b/lib/kv-router/src/services/indexer/mod.rs @@ -34,6 +34,7 @@ mod block_size; #[cfg(test)] mod deepapi_contract_tests; pub mod discovery; +pub mod evictions; pub mod kv_recover; pub mod listener; pub mod logging; @@ -56,9 +57,10 @@ use crate::config::min_initial_workers_from_env; use crate::services::common::zmq::validate_endpoint as validate_zmq_endpoint; use axum::http::header::HeaderName; use discovery::KubeDiscoveryConfig; +use evictions::KeepEvictionsConfig; use kv_recover::KvRecoverSettings; use logging::AccessLogSink; -use registry::WorkerRegistry; +use registry::{RegistryOptions, WorkerRegistry}; use server::{AppState, create_router}; pub struct IndexerConfig { @@ -76,6 +78,8 @@ pub struct IndexerConfig { pub kv_recover: KvRecoverSettings, /// When set, watch Kubernetes and register/deregister engine pods. pub kube_discovery: Option, + /// When set, park evictions and replay aged ones under memory pressure. + pub keep_evictions: Option, } pub(super) fn validate_listener_endpoints( @@ -172,8 +176,11 @@ pub async fn run_server(config: IndexerConfig) -> anyhow::Result<()> { "Starting standalone KV cache indexer (HTTP-only mode)" ); - let mut state = - AppState::new_with_cancel_token(config.threads, cancel_token.clone(), config.kv_recover)?; + let options = RegistryOptions { + kv_recover: config.kv_recover, + keep_evictions: config.keep_evictions.is_some(), + }; + let mut state = AppState::new_with_cancel_token(config.threads, cancel_token.clone(), options)?; state.access_log_sink = match config.access_log { Some(ref path) => { let s = AccessLogSink::new( @@ -269,6 +276,21 @@ async fn run_common( wait_for_min_initial_workers(registry, &cancel_token).await?; registry.signal_ready(); + if let Some(keep) = config.keep_evictions { + keep.validate()?; + evictions::spawn_cleanup_loop( + state.registry.clone(), + keep.retention_s, + keep.memory_threshold, + cancel_token.clone(), + ); + tracing::info!( + retention_s = keep.retention_s, + memory_threshold = keep.memory_threshold, + "keep-evictions enabled: parking eviction events, sweeping under memory pressure" + ); + } + if let Some(kube_config) = config.kube_discovery.clone() { #[cfg(feature = "kube-discovery")] pod_watcher::spawn_pod_watcher(kube_config, registry.clone(), cancel_token.clone()); diff --git a/lib/kv-router/src/services/indexer/registry.rs b/lib/kv-router/src/services/indexer/registry.rs index f36a3eebe4c1..e9611915ea59 100644 --- a/lib/kv-router/src/services/indexer/registry.rs +++ b/lib/kv-router/src/services/indexer/registry.rs @@ -19,6 +19,7 @@ use crate::indexer::KvIndexerMetrics; use crate::protocols::WorkerId; use super::backend::{Indexer, create_indexer_with_metrics}; +use super::evictions::PendingEvictions; use super::kv_recover::{KvRecoverClient, KvRecoverSettings}; use super::listener::spawn_zmq_listener; @@ -146,6 +147,14 @@ struct ListenerRuntime { generation: u64, } +/// DeepInfra registry behavior chosen at startup. +#[derive(Debug, Clone, Copy, Default)] +pub struct RegistryOptions { + pub kv_recover: KvRecoverSettings, + /// Park evictions instead of applying them (`--keep-evictions`). + pub keep_evictions: bool, +} + /// Where and how a listener fetches `/kv_recover` on a gap. #[derive(Clone)] pub struct RecoverTarget { @@ -163,10 +172,15 @@ pub struct ListenerExtras { pub pod_name: Option, } +/// Eviction events a `--keep-evictions` listener parks instead of applying. +pub type SharedPendingEvictions = Arc>; + pub struct ListenerRecord { endpoint: String, replay_endpoint: Option, recover: Option, + /// `Some` only under `--keep-evictions`; drained by the evictions sweep. + pending_evictions: Option, block_size: u32, indexer: Indexer, watermark: Arc, @@ -178,6 +192,7 @@ impl ListenerRecord { endpoint: String, replay_endpoint: Option, recover: Option, + pending_evictions: Option, block_size: u32, indexer: Indexer, watermark: Arc, @@ -186,6 +201,7 @@ impl ListenerRecord { endpoint, replay_endpoint, recover, + pending_evictions, block_size, indexer, watermark, @@ -210,6 +226,10 @@ impl ListenerRecord { self.recover.clone() } + pub(super) fn pending_evictions(&self) -> Option { + self.pending_evictions.clone() + } + pub(super) fn block_size(&self) -> u32 { self.block_size } @@ -352,6 +372,7 @@ pub struct WorkerRegistry { root_cancel_token: CancellationToken, retain_empty_indexers: bool, kv_recover: Arc, + keep_evictions: bool, } impl WorkerRegistry { @@ -404,13 +425,14 @@ impl WorkerRegistry { KvRecoverClient::new(KvRecoverSettings::default()) .expect("default kv_recover client builds"), ), + keep_evictions: false, } } - /// Replace the `/kv_recover` timeout and concurrency gate shared by every - /// listener registered afterwards. - pub fn with_kv_recover(mut self, settings: KvRecoverSettings) -> Result { - self.kv_recover = Arc::new(KvRecoverClient::new(settings)?); + /// Apply startup options; call before any registration. + pub fn with_options(mut self, options: RegistryOptions) -> Result { + self.kv_recover = Arc::new(KvRecoverClient::new(options.kv_recover)?); + self.keep_evictions = options.keep_evictions; Ok(self) } @@ -551,10 +573,14 @@ impl WorkerRegistry { endpoint, client: self.kv_recover.clone(), }); + let pending_evictions = self + .keep_evictions + .then(|| Arc::new(Mutex::new(PendingEvictions::default()))); let record = Arc::new(ListenerRecord::new( endpoint, replay_endpoint, recover, + pending_evictions, bs, indexer, watermark, @@ -873,6 +899,26 @@ impl WorkerRegistry { entry.indexer.clone() } + /// Every `--keep-evictions` listener as `(worker_id, dp_rank, buffer, + /// indexer)`, for the evictions sweep. + pub(super) fn listener_records(&self) -> Vec<(WorkerId, u32, SharedPendingEvictions, Indexer)> { + self.workers + .iter() + .flat_map(|entry| { + let worker_id = *entry.key(); + entry + .value() + .listeners + .iter() + .filter_map(|(dp_rank, record)| { + let buffer = record.pending_evictions()?; + Some((worker_id, *dp_rank, buffer, record.indexer())) + }) + .collect::>() + }) + .collect() + } + /// The trees a query for `model_name` reads: the named routing group's, /// or every group's when none is named. pub fn query_trees( diff --git a/lib/kv-router/src/services/indexer/server.rs b/lib/kv-router/src/services/indexer/server.rs index 8bbbf4543a2a..d2a03ccae7f8 100644 --- a/lib/kv-router/src/services/indexer/server.rs +++ b/lib/kv-router/src/services/indexer/server.rs @@ -21,9 +21,8 @@ use crate::indexer::TieredMatchDetails; use crate::protocols::{BlockHashOptions, LocalBlockHash, WorkerId, compute_block_hash_for_seq}; use crate::services::overlap::{MooncakeOverlapSummary, build_mooncake_overlap_summaries}; -use super::kv_recover::KvRecoverSettings; use super::model_query::query_model; -use super::registry::{ListenerControlError, ListenerExtras, WorkerRegistry}; +use super::registry::{ListenerControlError, ListenerExtras, RegistryOptions, WorkerRegistry}; /// We need to fit one million tokens as JSON text, this should do it. const QUERY_REQUEST_BODY_LIMIT_BYTES: usize = 8 * 1024 * 1024; @@ -50,14 +49,14 @@ impl AppState { Self::new_with_cancel_token( indexer_threads, CancellationToken::new(), - KvRecoverSettings::default(), + RegistryOptions::default(), ) } pub(super) fn new_with_cancel_token( indexer_threads: usize, root_cancel_token: CancellationToken, - kv_recover: KvRecoverSettings, + options: RegistryOptions, ) -> anyhow::Result { #[cfg(feature = "metrics")] { @@ -71,7 +70,7 @@ impl AppState { indexer_metrics, root_cancel_token, ) - .with_kv_recover(kv_recover)?, + .with_options(options)?, ), access_log_sink: None, prom_registry, @@ -82,7 +81,7 @@ impl AppState { Ok(Self { registry: Arc::new( WorkerRegistry::new_with_cancel_token(indexer_threads, root_cancel_token) - .with_kv_recover(kv_recover)?, + .with_options(options)?, ), access_log_sink: None, }) From eb1ae310e34f6b3b1062c2cc3fb00d5514bbf888 Mon Sep 17 00:00:00 2001 From: Thach Nguyen Date: Wed, 23 Sep 2026 22:46:20 +0000 Subject: [PATCH 12/14] standalone-indexer: --enable-logging writes kv_audit lines for queries and events Off by default. When on, the indexer logs one kv_audit line per /query and /query_by_hash (model, routing groups queried, status, probe block hashes, full JSON response) and one per engine event it ingests (STORE / EVICT / CLEAR with worker, dp_rank, tier, event id, seq, and the token and sequence block hashes), tagged source=live|replay|recover and written before the --keep-evictions filter so it shows what the engine published. Isolate with RUST_LOG=kv_audit=info. Upstream's --access-log (#10700) records requests but not hashes, responses or events. Ports 5e5f3b8dd. Not ported: the non-blocking stdout writer from d14e05d6d's CLI init; no prod deployment passes --enable-logging. Claude-Session: https://claude.ai/code/session_01XKZ49SpoymBqDqgGRJYMNa --- lib/bindings/python/rust/llm/kv.rs | 7 ++ lib/kv-router/src/services/indexer/audit.rs | 106 ++++++++++++++++++ .../src/services/indexer/listener.rs | 21 +++- lib/kv-router/src/services/indexer/mod.rs | 11 ++ .../src/services/indexer/model_query.rs | 10 +- .../src/services/indexer/registry.rs | 42 +++++-- lib/kv-router/src/services/indexer/server.rs | 24 +++- 7 files changed, 205 insertions(+), 16 deletions(-) create mode 100644 lib/kv-router/src/services/indexer/audit.rs diff --git a/lib/bindings/python/rust/llm/kv.rs b/lib/bindings/python/rust/llm/kv.rs index adf78cce6593..0d75981a6ad1 100644 --- a/lib/bindings/python/rust/llm/kv.rs +++ b/lib/bindings/python/rust/llm/kv.rs @@ -256,6 +256,12 @@ struct KvIndexerCli { #[arg(long, default_value_t = 0.75)] evict_memory_threshold: f64, + /// Emit verbose audit logs on the `kv_audit` tracing target: one line per + /// query (block hashes + full response) and one per store/evict/clear + /// event ingested from the engines. Filter with RUST_LOG=kv_audit=info. + #[arg(long, default_value_t = false)] + enable_logging: bool, + /// Kubernetes namespace to watch for engine pods. Together with /// --watch-model-name (or --watch-label) this enables pod auto-discovery: /// subscribe on Ready, unsubscribe on delete. @@ -369,6 +375,7 @@ where concurrency: cli.recover_concurrency, }, kube_discovery, + audit_log: cli.enable_logging, keep_evictions: cli.keep_evictions.then_some( indexer::evictions::KeepEvictionsConfig { retention_s: cli.evict_retention_secs, diff --git a/lib/kv-router/src/services/indexer/audit.rs b/lib/kv-router/src/services/indexer/audit.rs new file mode 100644 index 000000000000..e342470f9ab4 --- /dev/null +++ b/lib/kv-router/src/services/indexer/audit.rs @@ -0,0 +1,106 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! `--enable-logging`: one `kv_audit` line per query and per ingested engine +//! event, for offline analysis (filter with `RUST_LOG=kv_audit=info`). +//! +//! Event lines are written before the `--keep-evictions` filter, so they show +//! what the engine published. `source` is `live` (ZMQ SUB), `replay` (ZMQ +//! replay socket) or `recover` (`/kv_recover`). + +use crate::protocols::{KvCacheEventData, LocalBlockHash, RouterEvent}; + +fn now_unix_ms() -> u64 { + use std::time::{SystemTime, UNIX_EPOCH}; + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_millis() as u64) + .unwrap_or(0) +} + +pub(super) fn log_event(event: &RouterEvent, seq: u64, source: &'static str) { + let ts_ms = now_unix_ms(); + let worker_id = event.worker_id; + let dp_rank = event.event.dp_rank; + let storage_tier = event.storage_tier; + let event_id = event.event.event_id; + match &event.event.data { + KvCacheEventData::Stored(data) => { + // tokens_hash is the local block hash (what /query computes); + // block_hash is the engine's chained sequence hash. + let token_block_hashes: Vec = + data.blocks.iter().map(|b| b.tokens_hash.0).collect(); + let sequence_block_hashes: Vec = + data.blocks.iter().map(|b| b.block_hash.0).collect(); + tracing::info!( + target: "kv_audit", + kind = "STORE", + ts_ms, + source, + seq, + worker_id, + dp_rank, + storage_tier = ?storage_tier, + event_id, + parent_hash = ?data.parent_hash.map(|h| h.0), + num_blocks = token_block_hashes.len(), + token_block_hashes = ?token_block_hashes, + sequence_block_hashes = ?sequence_block_hashes, + "kv_audit STORE" + ); + } + KvCacheEventData::Removed(data) => { + let sequence_block_hashes: Vec = data.block_hashes.iter().map(|h| h.0).collect(); + tracing::info!( + target: "kv_audit", + kind = "EVICT", + ts_ms, + source, + seq, + worker_id, + dp_rank, + storage_tier = ?storage_tier, + event_id, + num_blocks = sequence_block_hashes.len(), + sequence_block_hashes = ?sequence_block_hashes, + "kv_audit EVICT" + ); + } + KvCacheEventData::Cleared => { + tracing::info!( + target: "kv_audit", + kind = "CLEAR", + ts_ms, + source, + seq, + worker_id, + dp_rank, + storage_tier = ?storage_tier, + event_id, + "kv_audit CLEAR" + ); + } + } +} + +pub(super) fn log_query( + model_name: &str, + queried_groups: &[String], + status: u16, + block_hashes: &[LocalBlockHash], + response: &serde_json::Value, +) { + let block_hashes: Vec = block_hashes.iter().map(|h| h.0).collect(); + tracing::info!( + target: "kv_audit", + kind = "QUERY", + ts_ms = now_unix_ms(), + model_name, + fanout_groups = ?queried_groups, + status, + num_blocks = block_hashes.len(), + block_hashes = ?block_hashes, + response = %response, + "kv_audit QUERY" + ); +} diff --git a/lib/kv-router/src/services/indexer/listener.rs b/lib/kv-router/src/services/indexer/listener.rs index cea935029f7b..7b4e6f0c68f9 100644 --- a/lib/kv-router/src/services/indexer/listener.rs +++ b/lib/kv-router/src/services/indexer/listener.rs @@ -11,6 +11,7 @@ use crate::protocols::{KvCacheEventData, RouterEvent, WorkerId, WorkerWithDpRank use crate::recovery::{CursorObservation, CursorState}; use crate::zmq_wire::{ZmqEventNormalizer, decode_event_batch}; +use super::audit; use super::backend::Indexer; use super::block_size::BlockSizeGuard; use super::kv_recover::plan_recovery; @@ -155,6 +156,7 @@ struct ListenerLoop { replay_socket: Option, recover: Option, pending_evictions: Option, + audit_log: bool, watermark: Arc, normalizer: ZmqEventNormalizer, block_size_guard: BlockSizeGuard, @@ -173,6 +175,7 @@ impl ListenerLoop { replay_socket: Option, recover: Option, pending_evictions: Option, + audit_log: bool, watermark: Arc, ) -> Self { Self { @@ -184,6 +187,7 @@ impl ListenerLoop { replay_socket, recover, pending_evictions, + audit_log, watermark, normalizer: ZmqEventNormalizer::new(block_size), block_size_guard: BlockSizeGuard::new(block_size), @@ -191,6 +195,15 @@ impl ListenerLoop { } } + /// Audit-log `event` (as published) and apply the `--keep-evictions` + /// filter. `true` when the event must not reach the tree. + fn withhold_from_tree(&self, event: &RouterEvent, seq: u64, source: &'static str) -> bool { + if self.audit_log { + audit::log_event(event, seq, source); + } + self.park_kept_eviction(event) + } + /// `--keep-evictions`: a store cancels any parked eviction of its blocks /// and applies; an eviction is parked instead of applied; a clear is /// dropped. `true` when the event must not reach the tree. @@ -317,7 +330,7 @@ impl ListenerLoop { let router_event = placement_event .into_router_event() .expect("local worker placement must convert to router event"); - if self.park_kept_eviction(&router_event) { + if self.withhold_from_tree(&router_event, seq, "replay") { continue; } indexer @@ -389,7 +402,7 @@ impl ListenerLoop { } let mut applied = 0u64; for event in plan.events { - if self.park_kept_eviction(&event) { + if self.withhold_from_tree(&event, event.event.event_id, "recover") { continue; } self.indexer @@ -478,7 +491,7 @@ impl ListenerLoop { let router_event = placement_event .into_router_event() .expect("local worker placement must convert to router event"); - if self.park_kept_eviction(&router_event) { + if self.withhold_from_tree(&router_event, seq, "live") { continue; } self.indexer @@ -645,6 +658,7 @@ async fn run_listener( replay_socket, record.recover_target(), record.pending_evictions(), + record.audit_log(), watermark, ) .run() @@ -751,6 +765,7 @@ mod tests { Some(replay_socket.clone()), None, None, + false, watermark.clone(), ); let replayed = tokio::time::timeout(Duration::from_secs(5), listener.replay_gap(0, 2)) diff --git a/lib/kv-router/src/services/indexer/mod.rs b/lib/kv-router/src/services/indexer/mod.rs index 8820ff6254ac..026f8397ea7b 100644 --- a/lib/kv-router/src/services/indexer/mod.rs +++ b/lib/kv-router/src/services/indexer/mod.rs @@ -29,6 +29,7 @@ //! Tier counts are CUMULATIVE through each tier's walk — see the doc on the //! response struct in [`server`] for the exact semantics. +mod audit; pub mod backend; mod block_size; #[cfg(test)] @@ -80,6 +81,8 @@ pub struct IndexerConfig { pub kube_discovery: Option, /// When set, park evictions and replay aged ones under memory pressure. pub keep_evictions: Option, + /// Emit `kv_audit` lines for every query and ingested event. + pub audit_log: bool, } pub(super) fn validate_listener_endpoints( @@ -179,6 +182,7 @@ pub async fn run_server(config: IndexerConfig) -> anyhow::Result<()> { let options = RegistryOptions { kv_recover: config.kv_recover, keep_evictions: config.keep_evictions.is_some(), + audit_log: config.audit_log, }; let mut state = AppState::new_with_cancel_token(config.threads, cancel_token.clone(), options)?; state.access_log_sink = match config.access_log { @@ -276,6 +280,13 @@ async fn run_common( wait_for_min_initial_workers(registry, &cancel_token).await?; registry.signal_ready(); + if config.audit_log { + tracing::info!( + target: "kv_audit", + "kv_audit logging enabled: queries and store/evict/clear events will be logged" + ); + } + if let Some(keep) = config.keep_evictions { keep.validate()?; evictions::spawn_cleanup_loop( diff --git a/lib/kv-router/src/services/indexer/model_query.rs b/lib/kv-router/src/services/indexer/model_query.rs index 2a5a86dda2bd..1de6fc895ade 100644 --- a/lib/kv-router/src/services/indexer/model_query.rs +++ b/lib/kv-router/src/services/indexer/model_query.rs @@ -24,6 +24,8 @@ use super::server::{ScoreResponse, build_score_response}; pub(super) struct ModelQueryOutcome { pub status: StatusCode, pub body: serde_json::Value, + /// Routing groups whose tree answered. + pub queried_groups: Vec, } /// Query every tree in `trees` (non-empty; callers 404 before calling) and @@ -36,6 +38,7 @@ pub(super) async fn query_model( ) -> ModelQueryOutcome { let mut merged: Option = None; let mut last_error: Option = None; + let mut queried_groups = Vec::with_capacity(trees.len()); for (key, indexer, block_size) in trees { match indexer.find_tiered_matches(hashes_for(block_size)).await { @@ -45,6 +48,7 @@ pub(super) async fn query_model( Some(acc) => merge_score_responses(acc, response), None => merged = Some(response), } + queried_groups.push(key.routing_group); } Err(error) => { tracing::warn!( @@ -67,7 +71,11 @@ pub(super) async fn query_model( }), ), }; - ModelQueryOutcome { status, body } + ModelQueryOutcome { + status, + body, + queried_groups, + } } /// Worker sets are disjoint across trees (a worker registers into exactly one diff --git a/lib/kv-router/src/services/indexer/registry.rs b/lib/kv-router/src/services/indexer/registry.rs index e9611915ea59..dcfedc6e04f3 100644 --- a/lib/kv-router/src/services/indexer/registry.rs +++ b/lib/kv-router/src/services/indexer/registry.rs @@ -153,6 +153,8 @@ pub struct RegistryOptions { pub kv_recover: KvRecoverSettings, /// Park evictions instead of applying them (`--keep-evictions`). pub keep_evictions: bool, + /// Emit `kv_audit` lines for queries and events (`--enable-logging`). + pub audit_log: bool, } /// Where and how a listener fetches `/kv_recover` on a gap. @@ -175,12 +177,18 @@ pub struct ListenerExtras { /// Eviction events a `--keep-evictions` listener parks instead of applying. pub type SharedPendingEvictions = Arc>; -pub struct ListenerRecord { - endpoint: String, - replay_endpoint: Option, +/// DeepInfra per-listener ingest behavior, fixed at registration. +struct ListenerIngest { recover: Option, /// `Some` only under `--keep-evictions`; drained by the evictions sweep. pending_evictions: Option, + audit_log: bool, +} + +pub struct ListenerRecord { + endpoint: String, + replay_endpoint: Option, + ingest: ListenerIngest, block_size: u32, indexer: Indexer, watermark: Arc, @@ -191,8 +199,7 @@ impl ListenerRecord { fn new( endpoint: String, replay_endpoint: Option, - recover: Option, - pending_evictions: Option, + ingest: ListenerIngest, block_size: u32, indexer: Indexer, watermark: Arc, @@ -200,8 +207,7 @@ impl ListenerRecord { Self { endpoint, replay_endpoint, - recover, - pending_evictions, + ingest, block_size, indexer, watermark, @@ -223,11 +229,15 @@ impl ListenerRecord { } pub(super) fn recover_target(&self) -> Option { - self.recover.clone() + self.ingest.recover.clone() } pub(super) fn pending_evictions(&self) -> Option { - self.pending_evictions.clone() + self.ingest.pending_evictions.clone() + } + + pub(super) fn audit_log(&self) -> bool { + self.ingest.audit_log } pub(super) fn block_size(&self) -> u32 { @@ -373,6 +383,7 @@ pub struct WorkerRegistry { retain_empty_indexers: bool, kv_recover: Arc, keep_evictions: bool, + audit_log: bool, } impl WorkerRegistry { @@ -426,6 +437,7 @@ impl WorkerRegistry { .expect("default kv_recover client builds"), ), keep_evictions: false, + audit_log: false, } } @@ -433,9 +445,14 @@ impl WorkerRegistry { pub fn with_options(mut self, options: RegistryOptions) -> Result { self.kv_recover = Arc::new(KvRecoverClient::new(options.kv_recover)?); self.keep_evictions = options.keep_evictions; + self.audit_log = options.audit_log; Ok(self) } + pub fn audit_log_enabled(&self) -> bool { + self.audit_log + } + #[cfg(feature = "standalone-selection")] pub(crate) fn with_retained_indexers(mut self) -> Self { // Selection entries and their schedulers retain these indexers across worker updates. @@ -579,8 +596,11 @@ impl WorkerRegistry { let record = Arc::new(ListenerRecord::new( endpoint, replay_endpoint, - recover, - pending_evictions, + ListenerIngest { + recover, + pending_evictions, + audit_log: self.audit_log, + }, bs, indexer, watermark, diff --git a/lib/kv-router/src/services/indexer/server.rs b/lib/kv-router/src/services/indexer/server.rs index d2a03ccae7f8..67bb0e284d57 100644 --- a/lib/kv-router/src/services/indexer/server.rs +++ b/lib/kv-router/src/services/indexer/server.rs @@ -21,7 +21,8 @@ use crate::indexer::TieredMatchDetails; use crate::protocols::{BlockHashOptions, LocalBlockHash, WorkerId, compute_block_hash_for_seq}; use crate::services::overlap::{MooncakeOverlapSummary, build_mooncake_overlap_summaries}; -use super::model_query::query_model; +use super::audit; +use super::model_query::{ModelQueryOutcome, query_model}; use super::registry::{ListenerControlError, ListenerExtras, RegistryOptions, WorkerRegistry}; /// We need to fit one million tokens as JSON text, this should do it. @@ -372,11 +373,29 @@ async fn query(State(state): State>, Json(req): Json &state.registry.pod_names(), ) .await; + if state.registry.audit_log_enabled() { + let probe = hashes_by_block_size + .values() + .next() + .cloned() + .unwrap_or_default(); + audit_query(&model, &outcome, &probe); + } let mut resp = (outcome.status, Json(outcome.body)).into_response(); resp.extensions_mut().insert(AccessLogModel(model)); resp } +fn audit_query(model_name: &str, outcome: &ModelQueryOutcome, probe: &[LocalBlockHash]) { + audit::log_query( + model_name, + &outcome.queried_groups, + outcome.status.as_u16(), + probe, + &outcome.body, + ); +} + fn model_not_found(model: String) -> Response { let mut resp = ( StatusCode::NOT_FOUND, @@ -417,6 +436,9 @@ async fn query_by_hash( .map(|h| LocalBlockHash(*h as u64)) .collect(); let outcome = query_model(trees, |_| block_hashes.clone(), &state.registry.pod_names()).await; + if state.registry.audit_log_enabled() { + audit_query(&model, &outcome, &block_hashes); + } let mut resp = (outcome.status, Json(outcome.body)).into_response(); resp.extensions_mut().insert(AccessLogModel(model)); resp From ed0c84156f4e4e002d08af8921a1dba6e64e384e Mon Sep 17 00:00:00 2001 From: Thach Nguyen Date: Wed, 23 Sep 2026 22:50:43 +0000 Subject: [PATCH 13/14] standalone-indexer: --h24 flat counterfactual indexer mode A new Indexer::H24 backend answers "how many prefix blocks of this query were stored by ANY worker within the retention horizon", ignoring evictions, clears and worker identity. Two flat maps replace the radix tree: engine sequence hash -> chained prefix hash (so stores naming their parent by engine hash keep chaining) and prefix hash -> last_stored / last_touched. Matches are reported under a synthetic worker 0, so the /query_by_hash response shape is unchanged. - Matched-block ages feed dynamo_kvindexer_h24_age_seconds, whose bucket CDF over queried blocks is the hit-rate-vs-retention curve; queried blocks are counted once per request, not per tree. - An expiry sweep every 10 min drops entries older than --h24-horizon-secs (default 172800), and tightens to half the horizon above 92% of the cgroup memory limit. - Every routing group of a model collapses into one h24 tree, so a block matched in several engine_hash regimes is counted once (the 120% hit-rate incident). - /dump returns nothing in h24 mode; --h24 and --keep-evictions are mutually exclusive. A runtime flag of the one image rather than a separate -h24 branch: every live h24 deployment already passes --h24. Ports 73474e51a and 247ac26ad. Claude-Session: https://claude.ai/code/session_01XKZ49SpoymBqDqgGRJYMNa --- lib/bindings/python/rust/llm/kv.rs | 11 + lib/kv-router/src/services/indexer/backend.rs | 46 +- lib/kv-router/src/services/indexer/h24.rs | 469 ++++++++++++++++++ lib/kv-router/src/services/indexer/metrics.rs | 110 +++- lib/kv-router/src/services/indexer/mod.rs | 20 + .../src/services/indexer/model_query.rs | 12 +- .../src/services/indexer/registry.rs | 53 +- 7 files changed, 703 insertions(+), 18 deletions(-) create mode 100644 lib/kv-router/src/services/indexer/h24.rs diff --git a/lib/bindings/python/rust/llm/kv.rs b/lib/bindings/python/rust/llm/kv.rs index 0d75981a6ad1..e08938ef7086 100644 --- a/lib/bindings/python/rust/llm/kv.rs +++ b/lib/bindings/python/rust/llm/kv.rs @@ -262,6 +262,16 @@ struct KvIndexerCli { #[arg(long, default_value_t = false)] enable_logging: bool, + /// Run as the flat h24 counterfactual (kv-indexer:h24): one map per model + /// of every prefix stored by any worker within --h24-horizon-secs, + /// ignoring evictions. Exclusive with --keep-evictions. + #[arg(long, default_value_t = false)] + h24: bool, + + /// Retention horizon of the h24 expiry sweep, in seconds. + #[arg(long, default_value_t = 172800)] + h24_horizon_secs: u64, + /// Kubernetes namespace to watch for engine pods. Together with /// --watch-model-name (or --watch-label) this enables pod auto-discovery: /// subscribe on Ready, unsubscribe on delete. @@ -376,6 +386,7 @@ where }, kube_discovery, audit_log: cli.enable_logging, + h24_horizon_s: cli.h24.then_some(cli.h24_horizon_secs), keep_evictions: cli.keep_evictions.then_some( indexer::evictions::KeepEvictionsConfig { retention_s: cli.evict_retention_secs, diff --git a/lib/kv-router/src/services/indexer/backend.rs b/lib/kv-router/src/services/indexer/backend.rs index 65cb770596ee..0ffbd1886209 100644 --- a/lib/kv-router/src/services/indexer/backend.rs +++ b/lib/kv-router/src/services/indexer/backend.rs @@ -31,6 +31,9 @@ pub enum Indexer { primary: Arc>, lower_tier: LowerTierIndexers, }, + /// The flat counterfactual indexer (`--h24`): no workers, no evictions, + /// no tiers. See [`super::h24`]. + H24(Arc), } impl Indexer { @@ -41,6 +44,7 @@ impl Indexer { match self { Indexer::Single { primary, .. } => primary.apply_event(event).await, Indexer::Concurrent { primary, .. } => primary.apply_event(event).await, + Indexer::H24(h24) => h24_apply(h24, &event), } } @@ -49,6 +53,12 @@ impl Indexer { /// indexer otherwise. `Cleared` events fan out to their applicable physical /// indexes according to the event's reset scope. pub async fn apply_event_routed(&self, event: RouterEvent) -> Result<(), KvRouterError> { + // Tier-agnostic on purpose: in the counterfactual a store on ANY tier + // means the content existed. + if let Indexer::H24(h24) = self { + h24_apply(h24, &event); + return Ok(()); + } let targets_primary = match event.targets_primary() { Ok(targets_primary) => targets_primary, Err(_) => { @@ -56,6 +66,7 @@ impl Indexer { Self::Single { lower_tier, .. } | Self::Concurrent { lower_tier, .. } => { lower_tier.record_unsupported_residency_event(&event); } + Self::H24(_) => unreachable!("h24 events return above"), } return Ok(()); } @@ -122,6 +133,7 @@ impl Indexer { .await; } } + Indexer::H24(_) => unreachable!("h24 events return above"), } Ok(()) } @@ -146,6 +158,8 @@ impl Indexer { } primary.remove_worker(worker_id).await; } + // The counterfactual keeps content regardless of worker lifecycle. + Indexer::H24(_) => {} } } @@ -169,6 +183,7 @@ impl Indexer { } primary.remove_worker_dp_rank(worker_id, dp_rank).await; } + Indexer::H24(_) => {} } } @@ -182,6 +197,7 @@ impl Indexer { Indexer::Concurrent { primary, .. } => { primary.find_matches(hashes).await.map_err(Into::into) } + Indexer::H24(h24) => Ok(h24.find_matches(&hashes)), } } @@ -215,6 +231,13 @@ impl Indexer { lower_tier: lt, }) } + Indexer::H24(h24) => Ok(TieredMatchDetails { + device: MatchDetails { + overlap_scores: h24.find_matches(&sequence), + ..Default::default() + }, + lower_tier: Default::default(), + }), } } @@ -242,6 +265,12 @@ impl Indexer { primary.dump_events().await.map_err(anyhow::Error::from)?, lower_tier.entries(), ), + // Flat maps cannot reconstruct chains; /dump and peer recovery + // are unsupported in h24 mode. + Indexer::H24(_) => { + tracing::warn!("h24 indexer does not support dump_events; returning empty"); + return Ok(Vec::new()); + } }; let mut out = primary_events; @@ -530,7 +559,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::H24(_) => { + unreachable!("test creates the single indexer") + } }; indexer @@ -598,6 +629,19 @@ mod tests { } } +/// Stores chain into the flat h24 maps; evictions and clears are the +/// counterfactual's no-op and are only counted. +fn h24_apply(h24: &super::h24::H24Indexer, event: &RouterEvent) { + match &event.event.data { + KvCacheEventData::Stored(data) => h24.apply_store(data), + KvCacheEventData::Removed(_) | KvCacheEventData::Cleared => h24.note_ignored_event(), + } +} + +pub fn create_h24_indexer() -> Indexer { + Indexer::H24(Arc::new(super::h24::H24Indexer::new())) +} + pub fn create_indexer(block_size: u32, num_threads: usize) -> Indexer { create_indexer_with_metrics( block_size, diff --git a/lib/kv-router/src/services/indexer/h24.rs b/lib/kv-router/src/services/indexer/h24.rs new file mode 100644 index 000000000000..30ca263a3fa1 --- /dev/null +++ b/lib/kv-router/src/services/indexer/h24.rs @@ -0,0 +1,469 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! `--h24` counterfactual indexer: answers "how many prefix blocks of this +//! query were stored by ANY worker within the retention horizon", ignoring +//! evictions and worker identity entirely. +//! +//! Instead of a radix tree it keeps two flat maps: +//! - `attach`: engine sequence hash -> our chained prefix hash, so store +//! events that name their parent by engine hash can keep chaining. +//! - `seen`: chained prefix hash -> (last_stored, last_touched) timestamps. +//! +//! The chain recurrence is the crate-wide [`compute_next_seq_hash`] (first +//! block's prefix hash equals its tokens hash), so `seen` keys are the same +//! `SequenceHash` values the rest of the codebase computes for a prefix. +//! Queries recompute the chain from the caller's tokens hashes and walk it to +//! the first miss; matched-block ages feed the retention-curve histograms. +//! +//! Matches are reported under a single synthetic worker ([`H24_WORKER_ID`]) +//! so the HTTP response shape stays identical for existing callers. + +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::Duration; + +use parking_lot::Mutex; +use rustc_hash::FxHashMap; +use tokio_util::sync::CancellationToken; + +use crate::protocols::{ + KvCacheStoreData, LocalBlockHash, OverlapScores, WorkerWithDpRank, compute_next_seq_hash, +}; + +use super::backend::Indexer; +use super::evictions::memory_usage_fraction; +use super::metrics; +use super::registry::WorkerRegistry; + +/// Synthetic worker id all h24 matches are reported under. Real worker ids +/// are xxh3 hashes of pod names, so 0 cannot collide with one. +pub const H24_WORKER_ID: u64 = 0; + +/// How often the expiry sweep runs. +const SWEEP_INTERVAL: Duration = Duration::from_secs(600); + +/// Memory fraction above which the sweep tightens the horizon to half. +const MEMORY_GUARD_THRESHOLD: f64 = 0.92; + +/// Wall-clock seconds since the Unix epoch (u32 is fine until 2106). +fn unix_secs() -> u32 { + use std::time::{SystemTime, UNIX_EPOCH}; + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_secs() as u32) + .unwrap_or(0) +} + +#[derive(Clone, Copy)] +struct AttachEntry { + prefix: u64, + last_seen: u32, +} + +#[derive(Clone, Copy)] +struct SeenEntry { + last_stored: u32, + last_touched: u32, +} + +#[derive(Default)] +struct Inner { + /// Engine sequence hash -> our chained prefix hash. Store-side plumbing + /// only: lets a store that names its parent by engine hash keep chaining. + attach: FxHashMap, + /// Chained prefix hash -> timestamps. One entry per unique prefix. + seen: FxHashMap, +} + +pub struct H24Indexer { + inner: Mutex, + parent_misses: AtomicU64, + ignored_events: AtomicU64, +} + +impl H24Indexer { + pub fn new() -> Self { + Self { + inner: Mutex::new(Inner::default()), + parent_misses: AtomicU64::new(0), + ignored_events: AtomicU64::new(0), + } + } + + /// Apply a store event: resolve the parent's prefix hash, chain through + /// the event's blocks, and stamp `seen`/`attach`. A missing parent drops + /// the whole event (every block chains off it) and is counted. + pub fn apply_store(&self, data: &KvCacheStoreData) { + self.apply_store_at(data, unix_secs()); + } + + fn apply_store_at(&self, data: &KvCacheStoreData, now: u32) { + if data.blocks.is_empty() { + return; + } + + let mut inner = self.inner.lock(); + + // None parent = sequence start: the first block's prefix hash equals + // its tokens hash (compute_seq_hash_for_block's recurrence). + let mut parent_prefix: Option = match data.parent_hash { + None => None, + Some(parent) => match inner.attach.get(&parent.0) { + Some(entry) => Some(entry.prefix), + None => { + self.parent_misses.fetch_add(1, Ordering::Relaxed); + metrics::h24_inc_parent_miss(); + return; + } + }, + }; + + for block in &data.blocks { + let prefix = match parent_prefix { + None => block.tokens_hash.0, + Some(parent) => compute_next_seq_hash(parent, block.tokens_hash), + }; + let entry = inner.seen.entry(prefix).or_insert(SeenEntry { + last_stored: now, + last_touched: now, + }); + entry.last_stored = now; + entry.last_touched = entry.last_touched.max(now); + inner.attach.insert( + block.block_hash.0, + AttachEntry { + prefix, + last_seen: now, + }, + ); + parent_prefix = Some(prefix); + } + } + + /// Longest-prefix match: rebuild the chain from the query's tokens hashes + /// and walk `seen` to the first miss. Matched blocks bump `last_touched` + /// and feed the age histograms; the score is reported under the synthetic + /// worker. + pub fn find_matches(&self, sequence: &[LocalBlockHash]) -> OverlapScores { + self.find_matches_at(sequence, unix_secs()) + } + + fn find_matches_at(&self, sequence: &[LocalBlockHash], now: u32) -> OverlapScores { + let mut scores = OverlapScores::new(); + if sequence.is_empty() { + return scores; + } + + let mut inner = self.inner.lock(); + let mut depth = 0u32; + let mut prefix: Option = None; + for hash in sequence { + let next = match prefix { + None => hash.0, + Some(parent) => compute_next_seq_hash(parent, *hash), + }; + let Some(entry) = inner.seen.get_mut(&next) else { + break; + }; + metrics::h24_observe_age("since_stored", now.saturating_sub(entry.last_stored)); + metrics::h24_observe_age("since_touched", now.saturating_sub(entry.last_touched)); + entry.last_touched = entry.last_touched.max(now); + depth += 1; + prefix = Some(next); + } + drop(inner); + + if depth > 0 { + metrics::h24_add_matched_blocks(depth as usize); + scores + .scores + .insert(WorkerWithDpRank::new(H24_WORKER_ID, 0), depth); + } + scores + } + + pub fn note_ignored_event(&self) { + self.ignored_events.fetch_add(1, Ordering::Relaxed); + } + + /// Drop entries not seen (stored or touched) since `cutoff_secs`. + /// Returns (seen_dropped, attach_dropped). + pub fn expire_older_than(&self, cutoff_secs: u32) -> (usize, usize) { + let mut inner = self.inner.lock(); + let seen_before = inner.seen.len(); + let attach_before = inner.attach.len(); + inner + .seen + .retain(|_, e| e.last_stored.max(e.last_touched) >= cutoff_secs); + inner.attach.retain(|_, e| e.last_seen >= cutoff_secs); + ( + seen_before - inner.seen.len(), + attach_before - inner.attach.len(), + ) + } + + /// (seen entries, attach entries). + pub fn sizes(&self) -> (usize, usize) { + let inner = self.inner.lock(); + (inner.seen.len(), inner.attach.len()) + } + + pub fn parent_misses(&self) -> u64 { + self.parent_misses.load(Ordering::Relaxed) + } + + pub fn ignored_events(&self) -> u64 { + self.ignored_events.load(Ordering::Relaxed) + } +} + +impl Default for H24Indexer { + fn default() -> Self { + Self::new() + } +} + +/// Background expiry sweep: every [`SWEEP_INTERVAL`], drop entries older than +/// `horizon_secs` from every h24 indexer in the registry, refresh the size +/// gauges, and tighten to half the horizon if memory is still above +/// [`MEMORY_GUARD_THRESHOLD`] afterwards. +pub fn spawn_expiry_loop( + registry: Arc, + horizon_secs: u64, + cancel: CancellationToken, +) { + tokio::spawn(async move { + let mut ticker = tokio::time::interval(SWEEP_INTERVAL); + ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + loop { + tokio::select! { + _ = cancel.cancelled() => return, + _ = ticker.tick() => {} + } + + let now = unix_secs(); + let cutoff = now.saturating_sub(horizon_secs as u32); + let mut dropped = (0usize, 0usize); + let mut sizes = (0usize, 0usize); + let indexers = registry.all_indexers_with_block_size(); + for (_, indexer, _) in &indexers { + let Indexer::H24(h24) = indexer else { + continue; + }; + let (seen_dropped, attach_dropped) = h24.expire_older_than(cutoff); + dropped.0 += seen_dropped; + dropped.1 += attach_dropped; + let (seen, attach) = h24.sizes(); + sizes.0 += seen; + sizes.1 += attach; + } + + // Memory guard: horizon is the target, memory is the hard limit. + let usage = memory_usage_fraction(); + if usage.is_some_and(|u| u >= MEMORY_GUARD_THRESHOLD) { + let tight_cutoff = now.saturating_sub((horizon_secs / 2) as u32); + let mut guard_dropped = 0usize; + for (_, indexer, _) in &indexers { + if let Indexer::H24(h24) = indexer { + let (s, a) = h24.expire_older_than(tight_cutoff); + guard_dropped += s + a; + } + } + tracing::warn!( + usage = ?usage, + horizon_secs, + guard_dropped, + "h24 memory guard: usage above threshold, expired at half horizon" + ); + } + + metrics::h24_set_sizes(sizes.0, sizes.1); + tracing::info!( + seen_blocks = sizes.0, + attach_entries = sizes.1, + seen_dropped = dropped.0, + attach_dropped = dropped.1, + horizon_secs, + memory_fraction = usage.map(|u| format!("{u:.3}")), + "h24 expiry sweep" + ); + } + }); +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::protocols::{ + ExternalSequenceBlockHash, KvCacheStoredBlockData, compute_seq_hash_for_block, + }; + + /// Store `local_hashes` chained after `prefix_hashes`, with engine ext + /// hashes derived from the crate chain recurrence offset by `ext_salt` + /// (different salts model different hash regimes storing the same tokens). + fn store_data(prefix_hashes: &[u64], local_hashes: &[u64], ext_salt: u64) -> KvCacheStoreData { + let prefix: Vec = + prefix_hashes.iter().copied().map(LocalBlockHash).collect(); + let full: Vec = prefix_hashes + .iter() + .chain(local_hashes.iter()) + .copied() + .map(LocalBlockHash) + .collect(); + let full_seq = compute_seq_hash_for_block(&full); + let parent_hash = compute_seq_hash_for_block(&prefix) + .last() + .map(|h| ExternalSequenceBlockHash(h.wrapping_add(ext_salt))); + let blocks = local_hashes + .iter() + .zip(full_seq[prefix_hashes.len()..].iter()) + .map(|(&local, &seq)| KvCacheStoredBlockData { + block_hash: ExternalSequenceBlockHash(seq.wrapping_add(ext_salt)), + tokens_hash: LocalBlockHash(local), + mm_extra_info: None, + }) + .collect(); + KvCacheStoreData { + parent_hash, + start_position: None, + blocks, + } + } + + fn seq(hashes: &[u64]) -> Vec { + hashes.iter().copied().map(LocalBlockHash).collect() + } + + fn depth(scores: &OverlapScores) -> u32 { + scores + .scores + .get(&WorkerWithDpRank::new(H24_WORKER_ID, 0)) + .copied() + .unwrap_or(0) + } + + #[test] + fn store_then_query_matches_prefix() { + let h24 = H24Indexer::new(); + h24.apply_store_at(&store_data(&[], &[1, 2, 3], 0), 100); + + assert_eq!(depth(&h24.find_matches_at(&seq(&[1, 2, 3]), 200)), 3); + assert_eq!(depth(&h24.find_matches_at(&seq(&[1, 2]), 200)), 2); + // Divergence after block 1 stops the walk. + assert_eq!(depth(&h24.find_matches_at(&seq(&[1, 9, 3]), 200)), 1); + // Different first block = no match at all. + assert_eq!(depth(&h24.find_matches_at(&seq(&[9, 2, 3]), 200)), 0); + assert_eq!(h24.sizes(), (3, 3)); + } + + #[test] + fn chained_store_attaches_via_engine_hash() { + let h24 = H24Indexer::new(); + h24.apply_store_at(&store_data(&[], &[1, 2], 0), 100); + // Second event chains off ext hash of block 2. + h24.apply_store_at(&store_data(&[1, 2], &[3], 0), 150); + + assert_eq!(depth(&h24.find_matches_at(&seq(&[1, 2, 3]), 200)), 3); + assert_eq!(h24.parent_misses(), 0); + } + + #[test] + fn same_tokens_different_ext_regime_dedup() { + let h24 = H24Indexer::new(); + // Two regimes (different engine ext hashes) store identical tokens. + h24.apply_store_at(&store_data(&[], &[1, 2], 0), 100); + h24.apply_store_at(&store_data(&[], &[1, 2], 777), 150); + + let (seen, attach) = h24.sizes(); + assert_eq!(seen, 2, "identical token chains must share seen entries"); + assert_eq!(attach, 4, "each regime keeps its own attach aliases"); + // Both regimes can chain further stores off their own ext hashes. + h24.apply_store_at(&store_data(&[1, 2], &[3], 777), 160); + assert_eq!(depth(&h24.find_matches_at(&seq(&[1, 2, 3]), 200)), 3); + assert_eq!(h24.parent_misses(), 0); + } + + #[test] + fn missing_parent_drops_event_and_counts() { + let h24 = H24Indexer::new(); + h24.apply_store_at(&store_data(&[1, 2], &[3], 0), 100); + + assert_eq!(h24.parent_misses(), 1); + assert_eq!(h24.sizes(), (0, 0)); + assert_eq!(depth(&h24.find_matches_at(&seq(&[1, 2, 3]), 200)), 0); + } + + #[test] + fn restore_refreshes_timestamps_and_expiry_respects_touch() { + let h24 = H24Indexer::new(); + h24.apply_store_at(&store_data(&[], &[1, 2], 0), 100); + + // A query at t=5000 touches block 1 only (divergent second block). + let _ = h24.find_matches_at(&seq(&[1, 9]), 5000); + + // Expire everything not seen since t=1000: block 2's entry dies + // (stored 100, never touched), block 1 survives via its touch. + let (seen_dropped, attach_dropped) = h24.expire_older_than(1000); + assert_eq!(seen_dropped, 1); + assert_eq!(attach_dropped, 2, "attach expires on store recency only"); + assert_eq!(depth(&h24.find_matches_at(&seq(&[1, 2]), 6000)), 1); + + // Re-store refreshes: nothing left to expire at the same cutoff. + h24.apply_store_at(&store_data(&[], &[1, 2], 0), 7000); + assert_eq!(h24.expire_older_than(1000), (0, 0)); + assert_eq!(depth(&h24.find_matches_at(&seq(&[1, 2]), 8000)), 2); + } + + /// Every routing group of a model shares ONE h24 map, and deregistration + /// with a pod's real routing group still finds it. + #[tokio::test] + async fn registry_collapses_routing_groups_into_one_h24_indexer() { + use crate::services::indexer::registry::RegistryOptions; + let registry = WorkerRegistry::new(1) + .with_options(RegistryOptions { + h24: true, + ..Default::default() + }) + .unwrap(); + registry.signal_ready(); + for (instance, group) in [(1u64, "regime-a"), (2u64, "regime-b")] { + registry + .register( + instance, + format!("tcp://127.0.0.1:1576{instance}"), + 0, + "m".to_string(), + group.to_string(), + 4, + None, + ) + .await + .unwrap(); + } + + let trees = registry.indexers_for_model("m"); + assert_eq!(trees.len(), 1, "h24 keeps one tree per model"); + assert_eq!(trees[0].0.routing_group, "h24"); + assert!(matches!(trees[0].1, Indexer::H24(_))); + + registry.deregister(1, "m", "regime-a").await.unwrap(); + registry.deregister(2, "m", "regime-b").await.unwrap(); + assert!(registry.list().is_empty()); + } + + #[test] + fn empty_query_and_empty_store_are_noops() { + let h24 = H24Indexer::new(); + h24.apply_store_at( + &KvCacheStoreData { + parent_hash: None, + start_position: None, + blocks: vec![], + }, + 100, + ); + assert_eq!(h24.sizes(), (0, 0)); + assert!(h24.find_matches_at(&[], 100).scores.is_empty()); + } +} diff --git a/lib/kv-router/src/services/indexer/metrics.rs b/lib/kv-router/src/services/indexer/metrics.rs index ef3950b14bb6..c8f9f63abf3f 100644 --- a/lib/kv-router/src/services/indexer/metrics.rs +++ b/lib/kv-router/src/services/indexer/metrics.rs @@ -10,7 +10,8 @@ use std::time::Instant; use axum::{extract::MatchedPath, http::Request, middleware::Next, response::Response}; #[cfg(feature = "metrics")] use prometheus::{ - HistogramVec, IntCounterVec, IntGauge, IntGaugeVec, Opts, exponential_buckets, histogram_opts, + HistogramVec, IntCounter, IntCounterVec, IntGauge, IntGaugeVec, Opts, exponential_buckets, + histogram_opts, }; #[cfg(feature = "metrics")] @@ -39,6 +40,15 @@ pub struct StandaloneIndexerMetrics { pub models: IntGauge, pub workers: IntGauge, pub listeners: IntGaugeVec, + /// h24 mode: age of each matched block at query time, labeled + /// kind = since_stored | since_touched. The bucket CDF over queried + /// blocks is the hit-rate-vs-retention curve. + pub h24_age_seconds: HistogramVec, + pub h24_queried_blocks: IntCounter, + pub h24_matched_blocks: IntCounter, + pub h24_parent_miss: IntCounter, + pub h24_seen_blocks: IntGauge, + pub h24_attach_entries: IntGauge, } #[cfg(feature = "metrics")] @@ -86,6 +96,43 @@ static METRICS: LazyLock = LazyLock::new(|| Standalone &["status"], ) .expect("valid gauge"), + h24_age_seconds: HistogramVec::new( + histogram_opts!( + format!("{METRICS_PREFIX}_h24_age_seconds"), + "h24: age of each matched block at query time", + vec![ + 60.0, 300.0, 900.0, 1800.0, 3600.0, 10800.0, 21600.0, 43200.0, 86400.0, 129600.0, + 172800.0 + ] + ), + &["kind"], + ) + .expect("valid histogram"), + h24_queried_blocks: IntCounter::new( + format!("{METRICS_PREFIX}_h24_queried_blocks_total"), + "h24: blocks in incoming queries (denominator of the retention curve)", + ) + .expect("valid counter"), + h24_matched_blocks: IntCounter::new( + format!("{METRICS_PREFIX}_h24_matched_blocks_total"), + "h24: matched prefix blocks across queries", + ) + .expect("valid counter"), + h24_parent_miss: IntCounter::new( + format!("{METRICS_PREFIX}_h24_parent_miss_total"), + "h24: store events dropped because the parent hash was unknown", + ) + .expect("valid counter"), + h24_seen_blocks: IntGauge::new( + format!("{METRICS_PREFIX}_h24_seen_blocks"), + "h24: unique prefix blocks currently tracked", + ) + .expect("valid gauge"), + h24_attach_entries: IntGauge::new( + format!("{METRICS_PREFIX}_h24_attach_entries"), + "h24: engine-hash attach aliases currently tracked", + ) + .expect("valid gauge"), }); #[cfg(feature = "metrics")] @@ -97,6 +144,12 @@ pub fn register(registry: &prometheus::Registry) -> Result<(), prometheus::Error registry.register(Box::new(m.models.clone()))?; registry.register(Box::new(m.workers.clone()))?; registry.register(Box::new(m.listeners.clone()))?; + registry.register(Box::new(m.h24_age_seconds.clone()))?; + registry.register(Box::new(m.h24_queried_blocks.clone()))?; + registry.register(Box::new(m.h24_matched_blocks.clone()))?; + registry.register(Box::new(m.h24_parent_miss.clone()))?; + registry.register(Box::new(m.h24_seen_blocks.clone()))?; + registry.register(Box::new(m.h24_attach_entries.clone()))?; Ok(()) } @@ -144,6 +197,50 @@ pub fn set_worker_state(models: usize, workers: usize, listener_counts: [i64; 4] #[cfg(not(feature = "metrics"))] pub fn set_worker_state(_models: usize, _workers: usize, _listener_counts: [i64; 4]) {} +#[cfg(feature = "metrics")] +pub fn h24_observe_age(kind: &str, age_s: u32) { + METRICS + .h24_age_seconds + .with_label_values(&[kind]) + .observe(age_s as f64); +} + +#[cfg(not(feature = "metrics"))] +pub fn h24_observe_age(_kind: &str, _age_s: u32) {} + +#[cfg(feature = "metrics")] +pub fn h24_add_queried_blocks(n: usize) { + METRICS.h24_queried_blocks.inc_by(n as u64); +} + +#[cfg(not(feature = "metrics"))] +pub fn h24_add_queried_blocks(_n: usize) {} + +#[cfg(feature = "metrics")] +pub fn h24_add_matched_blocks(n: usize) { + METRICS.h24_matched_blocks.inc_by(n as u64); +} + +#[cfg(not(feature = "metrics"))] +pub fn h24_add_matched_blocks(_n: usize) {} + +#[cfg(feature = "metrics")] +pub fn h24_inc_parent_miss() { + METRICS.h24_parent_miss.inc(); +} + +#[cfg(not(feature = "metrics"))] +pub fn h24_inc_parent_miss() {} + +#[cfg(feature = "metrics")] +pub fn h24_set_sizes(seen: usize, attach: usize) { + METRICS.h24_seen_blocks.set(seen as i64); + METRICS.h24_attach_entries.set(attach as i64); +} + +#[cfg(not(feature = "metrics"))] +pub fn h24_set_sizes(_seen: usize, _attach: usize) {} + #[cfg(all(test, feature = "metrics"))] mod tests { use super::*; @@ -167,6 +264,11 @@ mod tests { .errors_total .with_label_values(&["/query", "4xx"]) .inc(); + h24_observe_age("since_stored", 60); + h24_add_queried_blocks(4); + h24_add_matched_blocks(2); + h24_inc_parent_miss(); + h24_set_sizes(10, 12); let encoder = prometheus::TextEncoder::new(); let mut buf = Vec::new(); @@ -180,5 +282,11 @@ mod tests { assert!(output.contains("dynamo_kvindexer_workers 2")); assert!(output.contains("dynamo_kvindexer_listeners{status=\"pending\"} 1")); assert!(output.contains("dynamo_kvindexer_listeners{status=\"active\"} 1")); + assert!(output.contains("dynamo_kvindexer_h24_age_seconds")); + assert!(output.contains("dynamo_kvindexer_h24_queried_blocks_total")); + assert!(output.contains("dynamo_kvindexer_h24_matched_blocks_total")); + assert!(output.contains("dynamo_kvindexer_h24_parent_miss_total")); + assert!(output.contains("dynamo_kvindexer_h24_seen_blocks 10")); + assert!(output.contains("dynamo_kvindexer_h24_attach_entries 12")); } } diff --git a/lib/kv-router/src/services/indexer/mod.rs b/lib/kv-router/src/services/indexer/mod.rs index 026f8397ea7b..7a6496f28eb9 100644 --- a/lib/kv-router/src/services/indexer/mod.rs +++ b/lib/kv-router/src/services/indexer/mod.rs @@ -36,6 +36,7 @@ mod block_size; mod deepapi_contract_tests; pub mod discovery; pub mod evictions; +pub mod h24; pub mod kv_recover; pub mod listener; pub mod logging; @@ -83,6 +84,9 @@ pub struct IndexerConfig { pub keep_evictions: Option, /// Emit `kv_audit` lines for every query and ingested event. pub audit_log: bool, + /// Run every indexer as the flat h24 counterfactual with this retention + /// horizon (seconds). Exclusive with `keep_evictions`. + pub h24_horizon_s: Option, } pub(super) fn validate_listener_endpoints( @@ -183,7 +187,15 @@ pub async fn run_server(config: IndexerConfig) -> anyhow::Result<()> { kv_recover: config.kv_recover, keep_evictions: config.keep_evictions.is_some(), audit_log: config.audit_log, + h24: config.h24_horizon_s.is_some(), }; + if let Some(horizon_s) = config.h24_horizon_s { + anyhow::ensure!( + config.keep_evictions.is_none(), + "--h24 and --keep-evictions are mutually exclusive" + ); + anyhow::ensure!(horizon_s > 0, "--h24-horizon-secs must be positive"); + } let mut state = AppState::new_with_cancel_token(config.threads, cancel_token.clone(), options)?; state.access_log_sink = match config.access_log { Some(ref path) => { @@ -287,6 +299,14 @@ async fn run_common( ); } + if let Some(horizon_s) = config.h24_horizon_s { + h24::spawn_expiry_loop(state.registry.clone(), horizon_s, cancel_token.clone()); + tracing::info!( + horizon_s, + "h24 mode enabled: flat counterfactual indexer, evictions ignored" + ); + } + if let Some(keep) = config.keep_evictions { keep.validate()?; evictions::spawn_cleanup_loop( diff --git a/lib/kv-router/src/services/indexer/model_query.rs b/lib/kv-router/src/services/indexer/model_query.rs index 1de6fc895ade..3da5e046a37d 100644 --- a/lib/kv-router/src/services/indexer/model_query.rs +++ b/lib/kv-router/src/services/indexer/model_query.rs @@ -39,9 +39,14 @@ pub(super) async fn query_model( let mut merged: Option = None; let mut last_error: Option = None; let mut queried_groups = Vec::with_capacity(trees.len()); + let mut h24_queried_blocks: Option = None; for (key, indexer, block_size) in trees { - match indexer.find_tiered_matches(hashes_for(block_size)).await { + let hashes = hashes_for(block_size); + if matches!(indexer, Indexer::H24(_)) { + h24_queried_blocks = Some(h24_queried_blocks.unwrap_or(0).max(hashes.len())); + } + match indexer.find_tiered_matches(hashes).await { Ok(tiered) => { let response = build_score_response(&tiered, block_size, pod_names); match merged.as_mut() { @@ -62,6 +67,11 @@ pub(super) async fn query_model( } } + // Once per request, not per tree: the h24 retention curve's denominator. + if let Some(blocks) = h24_queried_blocks { + super::metrics::h24_add_queried_blocks(blocks); + } + let (status, body) = match merged { Some(response) => (StatusCode::OK, serde_json::json!(response)), None => ( diff --git a/lib/kv-router/src/services/indexer/registry.rs b/lib/kv-router/src/services/indexer/registry.rs index dcfedc6e04f3..2beec5144f4d 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, create_h24_indexer, create_indexer_with_metrics}; use super::evictions::PendingEvictions; use super::kv_recover::{KvRecoverClient, KvRecoverSettings}; use super::listener::spawn_zmq_listener; @@ -155,6 +155,8 @@ pub struct RegistryOptions { pub keep_evictions: bool, /// Emit `kv_audit` lines for queries and events (`--enable-logging`). pub audit_log: bool, + /// Every indexer is the flat h24 counterfactual, one per model (`--h24`). + pub h24: bool, } /// Where and how a listener fetches `/kv_recover` on a gap. @@ -384,6 +386,7 @@ pub struct WorkerRegistry { kv_recover: Arc, keep_evictions: bool, audit_log: bool, + h24: bool, } impl WorkerRegistry { @@ -438,6 +441,7 @@ impl WorkerRegistry { ), keep_evictions: false, audit_log: false, + h24: false, } } @@ -446,9 +450,35 @@ impl WorkerRegistry { self.kv_recover = Arc::new(KvRecoverClient::new(options.kv_recover)?); self.keep_evictions = options.keep_evictions; self.audit_log = options.audit_log; + self.h24 = options.h24; Ok(self) } + fn new_indexer(&self, block_size: u32) -> Indexer { + if self.h24 { + create_h24_indexer() + } else { + create_indexer_with_metrics(block_size, self.num_threads, self.indexer_metrics.clone()) + } + } + + /// h24 is worker-less and regime-agnostic ("stored by anyone"), so every + /// routing group of a model shares one flat map. Per-group maps would count + /// a block matched in N regimes N times against a once-per-request + /// denominator (live incident: 120% hit rate). Requires one block size + /// across the model's regimes. + fn partition_key(&self, model_name: String, routing_group: String) -> RoutingPartitionId { + if self.h24 { + RoutingPartitionId::new(model_name, "h24") + } else { + RoutingPartitionId::new(model_name, routing_group) + } + } + + pub fn h24_enabled(&self) -> bool { + self.h24 + } + pub fn audit_log_enabled(&self) -> bool { self.audit_log } @@ -532,7 +562,7 @@ impl WorkerRegistry { replay_endpoint: Option, extras: ListenerExtras, ) -> Result<()> { - let key = RoutingPartitionId::new(model_name, routing_group); + let key = self.partition_key(model_name.to_string(), routing_group.to_string()); let registration = self.indexer_lifecycle.lock().await; if let Some(entry) = self.workers.get(&instance_id) { @@ -557,11 +587,7 @@ impl WorkerRegistry { "Creating new indexer" ); IndexerEntry { - indexer: create_indexer_with_metrics( - block_size, - self.num_threads, - self.indexer_metrics.clone(), - ), + indexer: self.new_indexer(block_size), block_size, } }); @@ -633,7 +659,7 @@ impl WorkerRegistry { model_name: &str, routing_group: &str, ) -> Result<()> { - let key = RoutingPartitionId::new(model_name, routing_group); + let key = self.partition_key(model_name.to_string(), routing_group.to_string()); if let Some(entry) = self.workers.get(&instance_id) { if entry.key != key { @@ -674,7 +700,7 @@ impl WorkerRegistry { model_name: &str, routing_group: &str, ) -> Result<()> { - let key = RoutingPartitionId::new(model_name, routing_group); + let key = self.partition_key(model_name.to_string(), routing_group.to_string()); let (record, remove_worker) = { let mut entry = self @@ -891,6 +917,7 @@ impl WorkerRegistry { } pub fn get_or_create_indexer(&self, key: RoutingPartitionId, block_size: u32) -> Indexer { + let key = self.partition_key(key.model_name, key.routing_group); let entry = self.indexers.entry(key.clone()).or_insert_with(|| { tracing::info!( model_name = %key.model_name, @@ -899,11 +926,7 @@ impl WorkerRegistry { "Creating indexer from recovery dump" ); IndexerEntry { - indexer: create_indexer_with_metrics( - block_size, - self.num_threads, - self.indexer_metrics.clone(), - ), + indexer: self.new_indexer(block_size), block_size, } }); @@ -949,7 +972,7 @@ impl WorkerRegistry { let Some(routing_group) = routing_group else { return self.indexers_for_model(model_name); }; - let key = RoutingPartitionId::new(model_name, routing_group); + let key = self.partition_key(model_name.to_string(), routing_group.to_string()); self.indexers .get(&key) .map(|entry| vec![(key.clone(), entry.indexer.clone(), entry.block_size)]) From 812b1101cba9ce263d61347b03114e1bb280e74b Mon Sep 17 00:00:00 2001 From: Thach Nguyen Date: Wed, 23 Sep 2026 22:52:50 +0000 Subject: [PATCH 14/14] standalone-indexer: listener tests for the DeepInfra ingest paths Covers, through ListenerLoop with wire-format msgpack batches: - "CPU" medium stores land on the HostPinned tier, not the device tree (upstream #10368; the fork this replaces filed them as device); - sub-block partial-prefix stores are skipped without exiting; - --keep-evictions parks an eviction and a re-store cancels it; - a sequence gap is filled from /kv_recover: a TreeDump replaces the rank's state under the consumer identity and sets the watermark to its last_event_id. Claude-Session: https://claude.ai/code/session_01XKZ49SpoymBqDqgGRJYMNa --- .../src/services/indexer/listener.rs | 168 ++++++++++++++++++ 1 file changed, 168 insertions(+) diff --git a/lib/kv-router/src/services/indexer/listener.rs b/lib/kv-router/src/services/indexer/listener.rs index 7b4e6f0c68f9..9d6202e2e47b 100644 --- a/lib/kv-router/src/services/indexer/listener.rs +++ b/lib/kv-router/src/services/indexer/listener.rs @@ -785,3 +785,171 @@ mod tests { cancel.cancel(); } } + +#[cfg(test)] +mod deepinfra_tests { + use super::*; + use crate::protocols::{ + BlockHashOptions, LocalBlockHash, StorageTier, compute_block_hash_for_seq, + }; + use crate::services::indexer::backend::create_indexer; + use crate::services::indexer::evictions::PendingEvictions; + + const BLOCK_SIZE: u32 = 4; + + fn listener(pending: Option) -> (ListenerLoop, Indexer) { + // connect() is lazy, so no publisher is needed. + let live_socket = connect_sub_socket("tcp://127.0.0.1:1").unwrap(); + let indexer = create_indexer(BLOCK_SIZE, 1); + let listener = ListenerLoop::new( + 7, + 0, + BLOCK_SIZE, + indexer.clone(), + CancellationToken::new(), + live_socket, + None, + None, + pending, + false, + Arc::new(AtomicU64::new(WATERMARK_UNSET)), + ); + (listener, indexer) + } + + /// A vLLM batch of `BlockStored` events, positional msgpack as on the wire. + fn stored_batch(tokens: &[u32], block_size: usize, medium: &str) -> Vec { + let hashes: Vec = (0..tokens.len() / block_size) + .map(|i| 1000 + i as u64) + .collect(); + let event = ( + "BlockStored", + hashes, + Option::::None, + tokens.to_vec(), + block_size, + Option::::None, + medium, + ); + rmp_serde::to_vec(&(0.0_f64, vec![event], Some(0_i32))).unwrap() + } + + fn removed_batch(hashes: &[u64]) -> Vec { + let event = ("BlockRemoved", hashes.to_vec(), "GPU"); + rmp_serde::to_vec(&(0.0_f64, vec![event], Some(0_i32))).unwrap() + } + + fn probe(tokens: &[u32]) -> Vec { + compute_block_hash_for_seq(tokens, BLOCK_SIZE, BlockHashOptions::default()) + } + + /// vLLM's CPU offload connector publishes `medium: "CPU"`. The fork this + /// replaces filed those on the device tree; they must land in HostPinned + /// so `/query` reports them as `cpu`, not `gpu` (upstream #10368). + #[tokio::test] + async fn cpu_medium_events_are_indexed_on_the_host_pinned_tier() { + let (mut listener, indexer) = listener(None); + let tokens: Vec = (1..=8).collect(); + listener + .apply_live_batch(0, &stored_batch(&tokens, 4, "CPU")) + .await + .unwrap(); + indexer.dump_events().await.unwrap(); + + let tiered = indexer.find_tiered_matches(probe(&tokens)).await.unwrap(); + let worker = WorkerWithDpRank::new(7, 0); + assert_eq!(tiered.device.overlap_scores.scores.get(&worker), None); + assert_eq!( + tiered + .lower_tier + .get(&StorageTier::HostPinned) + .and_then(|m| m.hits.get(&worker)) + .copied(), + Some(2) + ); + } + + #[tokio::test] + async fn partial_prefix_stores_are_skipped_not_fatal() { + let (mut listener, indexer) = listener(None); + let tokens: Vec = (1..=2).collect(); + listener + .apply_live_batch(0, &stored_batch(&tokens, 2, "GPU")) + .await + .unwrap(); + assert!(indexer.dump_events().await.unwrap().is_empty()); + } + + /// `--keep-evictions`: an eviction is parked, not applied, so the block + /// stays matchable; a re-store cancels the parked eviction. + #[tokio::test] + async fn kept_evictions_are_parked_and_cancelled_by_a_restore() { + let pending: SharedPendingEvictions = + Arc::new(parking_lot::Mutex::new(PendingEvictions::default())); + let (mut listener, indexer) = listener(Some(pending.clone())); + let tokens: Vec = (1..=4).collect(); + listener + .apply_live_batch(0, &stored_batch(&tokens, 4, "GPU")) + .await + .unwrap(); + listener + .apply_live_batch(1, &removed_batch(&[1000])) + .await + .unwrap(); + indexer.dump_events().await.unwrap(); + + let scores = indexer.find_matches(probe(&tokens)).await.unwrap(); + assert_eq!(scores.scores.get(&WorkerWithDpRank::new(7, 0)), Some(&1)); + assert_eq!(pending.lock().len(), 1); + + listener + .apply_live_batch(2, &stored_batch(&tokens, 4, "GPU")) + .await + .unwrap(); + assert_eq!(pending.lock().len(), 0); + } + + /// A gap on a listener with a recover endpoint is filled from + /// `/kv_recover`: a TreeDump replaces the rank's state and moves the + /// watermark to its `last_event_id`. + #[tokio::test] + async fn gap_is_recovered_from_kv_recover_tree_dump() { + use crate::services::indexer::backend::test_util::store_event; + use crate::services::indexer::kv_recover::{KvRecoverClient, KvRecoverSettings}; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + let tokens: Vec = (1..=8).collect(); + let hashes: Vec = probe(&tokens).into_iter().map(|h| h.0).collect(); + let body = serde_json::json!({ + "TreeDump": { + "events": [store_event(999, 3, 0, &[], &hashes, StorageTier::Device)], + "last_event_id": 41 + } + }) + .to_string(); + let server = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let endpoint = format!("http://{}", server.local_addr().unwrap()); + tokio::spawn(async move { + let (mut stream, _) = server.accept().await.unwrap(); + let mut buf = [0u8; 4096]; + let _ = stream.read(&mut buf).await; + let reply = format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}", + body.len() + ); + stream.write_all(reply.as_bytes()).await.unwrap(); + }); + + let (mut listener, indexer) = listener(None); + listener.recover = Some(RecoverTarget { + endpoint, + client: Arc::new(KvRecoverClient::new(KvRecoverSettings::default()).unwrap()), + }); + listener.handle_gap(42).await.unwrap(); + indexer.dump_events().await.unwrap(); + + assert_eq!(listener.watermark.load(Ordering::Acquire), 41); + let scores = indexer.find_matches(probe(&tokens)).await.unwrap(); + assert_eq!(scores.scores.get(&WorkerWithDpRank::new(7, 0)), Some(&2)); + } +}