From 1dcd91e7aa446f6c201a6e436d943b1c8fd09f67 Mon Sep 17 00:00:00 2001 From: Akira Lonske Date: Fri, 9 Oct 2026 22:42:24 -0400 Subject: [PATCH 01/12] feat(kv-router): expose per-rank resident block counts to selection plugins Add WorkerInputs::RESIDENT_BLOCKS. A filter, scorer, or picker that declares it reads how many blocks the router's indexer tracks for each worker rank, including cached blocks no running request uses, through WorkerCandidate::resident_blocks() or WorkerInputView::resident_blocks(). The selection core starts a poller for a partition only when its policy declares the input. The poller reads the indexer's existing per-rank lookup sizes off the request path every 100 ms and lends one snapshot to each selection, so find_matches, the event apply path, and candidate rows are unchanged. If the indexer stops answering, the poller withdraws the counts. The frontend rejects RESIDENT_BLOCKS without CACHE, because the counts come from the indexer CACHE starts. Relates to #10697 Signed-off-by: Akira Lonske --- .../router/custom-worker-selection.mdx | 5 +- .../router/custom-policy-example/README.md | 2 + lib/kv-router/src/indexer/kv_indexer.rs | 29 ++- lib/kv-router/src/indexer/mod.rs | 2 + lib/kv-router/src/indexer/radix_tree.rs | 11 +- lib/kv-router/src/indexer/resident_blocks.rs | 68 ++++++ lib/kv-router/src/indexer/thread_pool.rs | 53 +++++ lib/kv-router/src/indexer/types.rs | 6 + .../src/plugins/worker_selection/inputs.rs | 29 +++ lib/kv-router/src/scheduling/selector/mod.rs | 8 + .../src/scheduling/selector/policy.rs | 205 +++++++++++++++++- .../src/scheduling/selector/reference.rs | 1 + lib/kv-router/src/services/indexer/backend.rs | 172 ++++++++++++++- .../src/services/selection/core/tests.rs | 73 +++++++ .../src/services/selection/core/workers.rs | 15 +- lib/llm/src/kv_router.rs | 39 ++++ 16 files changed, 704 insertions(+), 14 deletions(-) create mode 100644 lib/kv-router/src/indexer/resident_blocks.rs diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx index 096da72cd8d4..6e59c93cb56e 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx @@ -370,7 +370,7 @@ Each worker policy is constructed once per routing partition. An empty bundle re Worker identity and ordinary request metadata are always available. Dynamo calculates optional per-worker signals only for groups that a filter, scorer, or picker requests. -Routing cache resources start only if a component requests `CACHE`. `serve_indexer` and `enable_session_prefix_index` also require `CACHE`. +Routing cache resources start only if a component requests `CACHE`. `serve_indexer` and `enable_session_prefix_index` also require `CACHE`, and the frontend rejects `RESIDENT_BLOCKS` without `CACHE`. ### Request Context @@ -446,6 +446,7 @@ If a component needs no optional worker data, return `WorkerInputs::NONE`. Combi | `LOAD` | `decode_cost_blocks()` | Projected active decode footprint, including this request's additional active blocks | | `LOAD` | `active_requests()` | Requests currently active on the worker | | `PREFERRED_TAINT` | `preferred_taint_multiplier()` | Optional cost multiplier from matching preferred routing taints | +| `RESIDENT_BLOCKS` | `resident_blocks()` | Blocks the router's indexer tracks for the worker rank, including cached blocks no running request uses; refreshed at most every 100 ms and not clamped to capacity | The GPU, CPU, and disk overlap accessors return raw tier counts. A missing tier or worker entry is zero; Dynamo does not substitute its weighted effective-overlap estimate. Each custom scorer chooses how to combine the raw counts. @@ -453,6 +454,8 @@ A filter or scorer reads groups through `WorkerCandidate::cache()` or `load()`. `WorkerCandidate::preferred_taint_multiplier()` and `ScoredWorkerCandidate::preferred_taint_multiplier()` return the optional cost multiplier from preferred routing constraints. A filter, scorer, or picker must request `WorkerInputs::PREFERRED_TAINT` before reading the multiplier. Without that declaration, the component receives `None`, even when another component requested the multiplier. Exact hard-pinned requests also do not materialize it. Required routing constraints remain host-enforced eligibility rules. +`WorkerCandidate::resident_blocks()` and `WorkerInputView::resident_blocks(row)` return how many blocks the router's indexer tracks for the worker rank, including cached blocks no running request uses; use `LOAD` for work in flight. The router refreshes the counts off the request path at most every 100 ms, later while the indexer is backlogged, so they miss routing decisions made since the last refresh. Every component sees the same counts during one selection. With KV events, the count is the device-tier blocks the indexer currently tracks from engine events. With approximate routing, it is the blocks routed to the rank within the TTL, or the distinct full blocks the LRU cache policy retains, which leaves out its partial tail blocks and duplicate copies. The count is unweighted and not clamped to `total_kv_blocks`. A rank with no tracked blocks reads `0`, including a rank the indexer has no data for yet, such as a worker whose first KV events have not arrived. A component that did not request `WorkerInputs::RESIDENT_BLOCKS` receives `None`, and so does every component before the first refresh completes, with a remote indexer, or after the indexer stops answering. The frontend rejects a policy that requests `RESIDENT_BLOCKS` without `CACHE`, because the counts come from the indexer that `CACHE` starts. With KV events every router replica sees the same counts; with approximate routing each replica counts only the requests it routed. + ## Link the Policy Into Dynamo Both paths use the same policy crate, catalog, and YAML file. Choose the process that owns worker selection. diff --git a/examples/router/custom-policy-example/README.md b/examples/router/custom-policy-example/README.md index d67868309824..4a4f39d20542 100644 --- a/examples/router/custom-policy-example/README.md +++ b/examples/router/custom-policy-example/README.md @@ -24,6 +24,8 @@ Declare `WorkerInputs::CACHE` to read cache data through `candidate.cache()` (fi Preferred routing taints are optional candidate metadata. A filter, scorer, or picker must request `WorkerInputs::PREFERRED_TAINT` before reading `preferred_taint_multiplier()` from a candidate; otherwise, that component receives `None`, even if another component requested it. Exact hard-pinned requests also do not materialize it. Required routing taints remain Dynamo eligibility rules. +Declare `WorkerInputs::RESIDENT_BLOCKS` to read `resident_blocks()` from a candidate, or `input.resident_blocks(row)` in a picker. It returns how many blocks the router's indexer tracks for that worker rank, including cached blocks no running request uses, refreshed in the background. Without the declaration, before the first refresh, or with a remote indexer, it returns `None`. The frontend also requires `CACHE` when a policy requests `RESIDENT_BLOCKS`. + ## Pick a Starting Point | Crate | Use it for | diff --git a/lib/kv-router/src/indexer/kv_indexer.rs b/lib/kv-router/src/indexer/kv_indexer.rs index ecf97d1f25fa..afe4cdf59f3c 100644 --- a/lib/kv-router/src/indexer/kv_indexer.rs +++ b/lib/kv-router/src/indexer/kv_indexer.rs @@ -15,7 +15,8 @@ use super::{ ApproximateLruLease, ApproximateLruStats, ApproximateLruTask, ApproximateRetentionConfig, ContainsWorkerBlockRequest, DumpRequest, EventKind, FlushRequest, GetWorkersRequest, KvIndexerInterface, KvIndexerMetrics, KvRouterError, MatchDetails, MatchDetailsRequest, - MatchRequest, PreBoundEventCounters, RadixTree, RoutingDecisionRequest, panic_payload_message, + MatchRequest, PreBoundEventCounters, RadixTree, RoutingDecisionRequest, WorkerLookupStats, + WorkerLookupStatsRequest, panic_payload_message, }; use crate::indexer::pruning::{BlockEntry, PruneConfig, WorkerPruneManager}; use crate::protocols::*; @@ -336,6 +337,9 @@ pub struct KvIndexer { remove_worker_dp_rank_tx: mpsc::Sender<(WorkerId, DpRank)>, /// A sender for get workers requests. get_workers_tx: mpsc::Sender, + /// A sender for per-rank block count requests. + #[cfg_attr(not(feature = "standalone-indexer"), allow(dead_code))] + worker_lookup_stats_tx: mpsc::Sender, /// A sender for dump requests. dump_tx: mpsc::Sender, /// A sender for flush requests. @@ -430,6 +434,8 @@ impl KvIndexer { let (remove_worker_dp_rank_tx, remove_worker_dp_rank_rx) = mpsc::channel::<(WorkerId, DpRank)>(16); let (get_workers_tx, get_workers_rx) = mpsc::channel::(16); + let (worker_lookup_stats_tx, worker_lookup_stats_rx) = + mpsc::channel::(16); let (dump_tx, dump_rx) = mpsc::channel::(16); let (flush_tx, flush_rx) = mpsc::channel::(16); let (routing_tx, mut routing_rx) = mpsc::channel::(2048); @@ -460,6 +466,7 @@ impl KvIndexer { let mut remove_worker_rx = remove_worker_rx; let mut remove_worker_dp_rank_rx = remove_worker_dp_rank_rx; let mut get_workers_rx = get_workers_rx; + let mut worker_lookup_stats_rx = worker_lookup_stats_rx; let mut dump_rx = dump_rx; let mut flush_rx = flush_rx; let mut trie = delegate.map_or_else(RadixTree::new, RadixTree::new_with_delegate); @@ -559,6 +566,10 @@ impl KvIndexer { let _ = get_workers_req.resp.send(workers); } + Some(stats_req) = worker_lookup_stats_rx.recv() => { + let _ = stats_req.resp.send(trie.worker_lookup_stats()); + } + Some(dump_req) = dump_rx.recv() => { // NOTE: Dump requests use a separate channel from mutations, so the // actor may observe one while already-accepted mutations are still @@ -656,6 +667,7 @@ impl KvIndexer { remove_worker_tx, remove_worker_dp_rank_tx, get_workers_tx, + worker_lookup_stats_tx, dump_tx, flush_tx, routing_tx, @@ -774,6 +786,21 @@ impl KvIndexer { pub fn get_workers_sender(&self) -> mpsc::Sender { self.get_workers_tx.clone() } + + /// Distinct blocks the tree holds for each worker rank. + #[cfg_attr(not(feature = "standalone-indexer"), allow(dead_code))] + pub(crate) async fn worker_lookup_stats(&self) -> Result { + let (resp_tx, resp_rx) = oneshot::channel(); + self.worker_lookup_stats_tx + .send(WorkerLookupStatsRequest { resp: resp_tx }) + .await + .map_err(|_| KvRouterError::IndexerOffline)?; + // A request sent while the actor drops its receiver stays queued with resp_tx alive. + tokio::select! { + resp = resp_rx => resp.map_err(|_| KvRouterError::IndexerDroppedRequest), + _ = self.worker_lookup_stats_tx.closed() => Err(KvRouterError::IndexerOffline), + } + } } #[async_trait] diff --git a/lib/kv-router/src/indexer/mod.rs b/lib/kv-router/src/indexer/mod.rs index 5a9e98205605..b71734e66a85 100644 --- a/lib/kv-router/src/indexer/mod.rs +++ b/lib/kv-router/src/indexer/mod.rs @@ -69,6 +69,7 @@ mod lower_tier_indexers; mod metrics; #[cfg(feature = "bench")] mod observation; +mod resident_blocks; mod thread_pool; mod traits; mod types; @@ -95,6 +96,7 @@ pub use lower_tier_indexers::*; pub use metrics::*; #[cfg(feature = "bench")] pub use observation::*; +pub(crate) use resident_blocks::*; pub use thread_pool::*; pub use traits::*; pub use types::*; diff --git a/lib/kv-router/src/indexer/radix_tree.rs b/lib/kv-router/src/indexer/radix_tree.rs index 054b0a257271..3dd8044c3525 100644 --- a/lib/kv-router/src/indexer/radix_tree.rs +++ b/lib/kv-router/src/indexer/radix_tree.rs @@ -13,7 +13,7 @@ use std::{ use rustc_hash::{FxBuildHasher, FxHashMap, FxHashSet}; use super::compressed_radix::{NodeState, append_dump_events}; -use super::{EventWarningKind, MatchDetails, PreBoundEventCounters}; +use super::{EventWarningKind, MatchDetails, PreBoundEventCounters, WorkerLookupStats}; use crate::protocols::*; pub(crate) type SharedRadixBlock = Rc>; @@ -860,6 +860,15 @@ impl RadixTree { self.lookup.values().map(FxHashMap::len).sum() } + /// Distinct blocks held by each worker rank. + pub(crate) fn worker_lookup_stats(&self) -> WorkerLookupStats { + WorkerLookupStats::from_worker_block_counts( + self.lookup + .iter() + .map(|(worker, worker_lookup)| (*worker, worker_lookup.len())), + ) + } + #[cfg(test)] pub(crate) fn tree_size_for_worker(&self, worker: WorkerWithDpRank) -> Option { self.lookup.get(&worker).map(FxHashMap::len) diff --git a/lib/kv-router/src/indexer/resident_blocks.rs b/lib/kv-router/src/indexer/resident_blocks.rs new file mode 100644 index 000000000000..49215056575e --- /dev/null +++ b/lib/kv-router/src/indexer/resident_blocks.rs @@ -0,0 +1,68 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Per-rank resident block counts, published off the request path. +//! +//! Event threads own the per-rank block maps, so reading them is a queued round +//! trip. A poller makes that trip and publishes the result here, where worker +//! selection can read it without waiting on the indexer. + +use std::sync::Arc; +use std::time::Duration; + +use arc_swap::ArcSwapOption; +use rustc_hash::FxHashMap; + +use super::WorkerLookupStats; +use crate::protocols::WorkerWithDpRank; + +/// How often routing hosts republish [`ResidentBlockCounts`]. +#[cfg_attr(not(feature = "standalone-selection"), allow(dead_code))] +pub(crate) const RESIDENT_BLOCK_COUNTS_POLL_INTERVAL: Duration = Duration::from_millis(100); + +/// Blocks the indexer tracked for each worker rank when polled. +#[derive(Debug)] +pub(crate) struct ResidentBlockCounts { + counts: FxHashMap, +} + +impl ResidentBlockCounts { + /// Blocks tracked for `worker`; zero when the indexer tracks none for it. + pub(crate) fn get(&self, worker: WorkerWithDpRank) -> u64 { + self.counts.get(&worker).copied().unwrap_or(0) + } +} + +impl From for ResidentBlockCounts { + fn from(stats: WorkerLookupStats) -> Self { + let mut counts = FxHashMap::default(); + for (worker, blocks) in stats.worker_blocks { + *counts.entry(worker).or_default() += blocks as u64; + } + Self { counts } + } +} + +/// The latest published [`ResidentBlockCounts`], shared by the poller and its readers. +#[derive(Clone, Default)] +pub(crate) struct ResidentBlockCountsHandle { + latest: Arc>, +} + +impl ResidentBlockCountsHandle { + /// The latest counts, or `None` before the first poll or after the poller stops. + pub(crate) fn load(&self) -> Option> { + self.latest.load_full() + } + + #[cfg_attr(not(feature = "standalone-indexer"), allow(dead_code))] + pub(crate) fn publish(&self, counts: ResidentBlockCounts) { + self.latest.store(Some(Arc::new(counts))); + } + + /// Withdraw the counts so readers see `None` rather than a frozen snapshot. + #[cfg_attr(not(feature = "standalone-indexer"), allow(dead_code))] + pub(crate) fn clear(&self) { + self.latest.store(None); + } +} diff --git a/lib/kv-router/src/indexer/thread_pool.rs b/lib/kv-router/src/indexer/thread_pool.rs index 135115cd0876..430bfc123943 100644 --- a/lib/kv-router/src/indexer/thread_pool.rs +++ b/lib/kv-router/src/indexer/thread_pool.rs @@ -599,6 +599,38 @@ impl ThreadPoolIndexer { WorkerLookupStats::from_worker_block_counts(worker_blocks) } + /// Like [`Self::worker_lookup_stats`], but fails if an event thread is gone, so a caller + /// never mistakes a dead thread's ranks for empty ones. Each rank lives on one thread. + #[cfg_attr(not(feature = "standalone-indexer"), allow(dead_code))] + pub(crate) async fn try_worker_lookup_stats(&self) -> Result { + let mut receivers = Vec::with_capacity(self.worker_event_channels.len()); + for channel in &self.worker_event_channels { + let (resp_tx, resp_rx) = oneshot::channel(); + channel + .send(WorkerTask::Stats(resp_tx)) + .map_err(|_| KvRouterError::IndexerOffline)?; + receivers.push(resp_rx); + } + let mut worker_blocks = Vec::new(); + for (channel, mut receiver) in self.worker_event_channels.iter().zip(receivers) { + // A lane that dies with this request queued neither answers nor drops it, since our + // own senders keep the queue alive, so check for the dead lane while waiting. + let stats = loop { + match tokio::time::timeout(std::time::Duration::from_millis(50), &mut receiver) + .await + { + Ok(reply) => break reply.map_err(|_| KvRouterError::IndexerDroppedRequest)?, + Err(_) if channel.is_disconnected() => { + return Err(KvRouterError::IndexerOffline); + } + Err(_) => {} + } + }; + worker_blocks.extend(stats.worker_blocks); + } + Ok(WorkerLookupStats { worker_blocks }) + } + pub async fn get_workers(&self) -> Vec { self.worker_lookup_stats() .await @@ -1576,6 +1608,27 @@ mod tests { assert_score(&indexer, &[20], rank1, 1).await; } + #[tokio::test] + async fn strict_stats_fail_when_a_lane_dies_with_the_request_queued() { + let indexer = ThreadPoolIndexer::new(ConcurrentRadixTreeCompressed::new(), 2, 16); + // Keep lane 0 busy so the stats request queues behind its Terminate. + for block in 0..20_000u64 { + indexer + .apply_event(make_store_event_with_dp_rank(7, &[block, block + 1], 0)) + .await; + } + indexer.worker_event_channels[0] + .send(WorkerTask::Terminate) + .unwrap(); + let read = tokio::time::timeout( + std::time::Duration::from_secs(5), + indexer.try_worker_lookup_stats(), + ) + .await + .expect("a dead lane must not hang the read"); + assert!(read.is_err()); + } + #[tokio::test] async fn whole_worker_remove_barriers_every_rank_lane_before_returning() { let indexer = ThreadPoolIndexer::new(ConcurrentRadixTreeCompressed::new(), 2, 16); diff --git a/lib/kv-router/src/indexer/types.rs b/lib/kv-router/src/indexer/types.rs index edf72d2777f8..70bb99e27efc 100644 --- a/lib/kv-router/src/indexer/types.rs +++ b/lib/kv-router/src/indexer/types.rs @@ -570,6 +570,12 @@ pub struct GetWorkersRequest { pub resp: oneshot::Sender>, } +/// A request for the number of blocks tracked for each worker rank +pub(crate) struct WorkerLookupStatsRequest { + /// Channel to send the per-rank block counts + pub resp: oneshot::Sender, +} + #[derive(Debug, Default)] pub struct WorkerLookupStats { pub worker_blocks: Vec<(WorkerWithDpRank, usize)>, diff --git a/lib/kv-router/src/plugins/worker_selection/inputs.rs b/lib/kv-router/src/plugins/worker_selection/inputs.rs index 4ce3dd6aaf4c..c19aecc39625 100644 --- a/lib/kv-router/src/plugins/worker_selection/inputs.rs +++ b/lib/kv-router/src/plugins/worker_selection/inputs.rs @@ -3,6 +3,7 @@ //! Worker input groups, stored snapshots, and component-scoped borrowed views. +use crate::indexer::ResidentBlockCounts; use crate::protocols::{SharedCacheHits, WorkerWithDpRank}; use std::ops::BitOr; @@ -67,6 +68,19 @@ impl<'a> WorkerCandidate<'a> { None } } + + /// Blocks the router's indexer tracks for this worker rank, only when this component declared + /// RESIDENT_BLOCKS: device blocks from KV events, or the router's prediction in approximate + /// modes. Refreshed off the request path, unweighted, and not clamped to capacity; 0 also covers + /// ranks with no data yet. None means access was not requested or the host has no counts. + pub fn resident_blocks(self) -> Option { + if !self.inputs.contains(WorkerInputs::RESIDENT_BLOCKS) { + return None; + } + self.cache_snapshot + .resident_blocks + .map(|counts| counts.get(self.data.worker)) + } } /// Borrowed candidate batch restricted to one scorer's declared inputs. @@ -140,6 +154,9 @@ impl WorkerInputs { pub const LOAD: Self = Self(1 << 1); /// Request preferred-taint routing metadata. pub const PREFERRED_TAINT: Self = Self(1 << 2); + /// Request per-rank counts of blocks the router's indexer tracks (routing predictions in + /// approximate modes). The frontend rejects RESIDENT_BLOCKS without CACHE. + pub const RESIDENT_BLOCKS: Self = Self(1 << 3); /// Request host-owned active-request counts. pub const OCCUPANCY: Self = Self(1 << 5); #[cfg(any(test, feature = "bench"))] @@ -166,6 +183,8 @@ impl BitOr for WorkerInputs { pub(crate) struct CacheSnapshot<'a> { pub(crate) shared_hits: Option<&'a SharedCacheHits>, pub(crate) has_tier_matches: bool, + /// Loaded only when a component declared RESIDENT_BLOCKS. + pub(crate) resident_blocks: Option<&'a ResidentBlockCounts>, } /// Numeric cache row retained in host buffers between selections. @@ -238,6 +257,7 @@ pub struct WorkerInputView<'a> { pub(crate) candidates: &'a [ScoredWorkerCandidate], pub(crate) cache: Option>, pub(crate) load: Option<&'a [WorkerLoadInput]>, + pub(crate) resident_blocks: Option<&'a ResidentBlockCounts>, } impl CandidateData { @@ -369,4 +389,13 @@ impl<'a> WorkerInputView<'a> { pub fn load(self) -> Option<&'a [WorkerLoadInput]> { self.load } + + /// Return [`WorkerCandidate::resident_blocks`] for candidate `row` when the picker requested + /// [`WorkerInputs::RESIDENT_BLOCKS`], or None for an out-of-range row. + pub fn resident_blocks(self, row: usize) -> Option { + let counts = self.resident_blocks?; + self.candidates + .get(row) + .map(|candidate| counts.get(candidate.worker)) + } } diff --git a/lib/kv-router/src/scheduling/selector/mod.rs b/lib/kv-router/src/scheduling/selector/mod.rs index 8935d1e9d499..076a953d2c36 100644 --- a/lib/kv-router/src/scheduling/selector/mod.rs +++ b/lib/kv-router/src/scheduling/selector/mod.rs @@ -27,6 +27,7 @@ use reference::pick_default_worker; use super::filter::{RoutingEligibility, WorkerEligibilityError}; use super::types::{KvSchedulerError, SchedulingRequest, WorkerSelectionPolicyError}; +use crate::indexer::ResidentBlockCounts; use crate::protocols::{WorkerConfigLike, WorkerId, WorkerSelectionResult, WorkerWithDpRank}; /// Low-level selector used by routing hosts. @@ -145,6 +146,7 @@ impl<'a> MaterializedSelectionInput<'a> { has_tier_matches: !request.overlap.tier_overlap_blocks.device.is_empty() || !request.overlap.tier_overlap_blocks.host_pinned.is_empty() || !request.overlap.tier_overlap_blocks.disk.is_empty(), + resident_blocks: None, }, context: WorkerSelectionContext { request, @@ -365,6 +367,7 @@ fn select_worker_with_policy( request: &SchedulingRequest, eligibility: RoutingEligibility<'_>, block_size: u32, + resident_blocks: Option<&ResidentBlockCounts>, ) -> Result { eligibility.validate_pinned_worker_allowed()?; @@ -382,6 +385,7 @@ fn select_worker_with_policy( let mut input = MaterializedSelectionInput::new(request, block_size); input.context.pinned_worker = eligibility.pinned_worker(); + input.cache_snapshot.resident_blocks = resident_blocks; let selected = match state { #[cfg(any(test, feature = "bench"))] WorkerSelectionPolicyStateRef::Reference(kv_router_config, picker) => { @@ -428,6 +432,10 @@ fn select_worker_with_policy( load: picker_inputs .contains(WorkerInputs::LOAD) .then_some(load_inputs.as_slice()), + resident_blocks: input + .cache_snapshot + .resident_blocks + .filter(|_| picker_inputs.contains(WorkerInputs::RESIDENT_BLOCKS)), }; let row = picker.pick(&input.context, picker_input)?; let Some(candidate) = candidates.get(row) else { diff --git a/lib/kv-router/src/scheduling/selector/policy.rs b/lib/kv-router/src/scheduling/selector/policy.rs index ac4bed6fea61..6cc4f0089a9b 100644 --- a/lib/kv-router/src/scheduling/selector/policy.rs +++ b/lib/kv-router/src/scheduling/selector/policy.rs @@ -10,6 +10,7 @@ use super::{ MaterializedSelectionInput, WorkerSelectionInput, WorkerSelector, select_worker_with_policy, }; +use crate::indexer::ResidentBlockCountsHandle; use crate::protocols::{WorkerConfigLike, WorkerId, WorkerSelectionResult}; use crate::scheduling::config::KvRouterConfig; use crate::scheduling::filter::RoutingEligibility; @@ -57,6 +58,7 @@ pub struct WorkerSelectionPolicy { worker_label: &'static str, state: WorkerSelectionPolicyState, exclusive_affinity: bool, + resident_blocks: Option, } impl WorkerSelectionPolicy { @@ -107,6 +109,7 @@ impl WorkerSelectionPolicy { Self { worker_label, exclusive_affinity: false, + resident_blocks: None, state: WorkerSelectionPolicyState::Composed(RefCell::new(ComposedPolicyState { filters, scorers, @@ -130,6 +133,25 @@ impl WorkerSelectionPolicy { self } + /// Read per-rank resident block counts from `counts`. Routing hosts call this only for a + /// policy that declares [`WorkerInputs::RESIDENT_BLOCKS`]. + #[cfg_attr(not(feature = "standalone-selection"), allow(dead_code))] + pub(crate) fn with_resident_block_counts(mut self, counts: ResidentBlockCountsHandle) -> Self { + self.resident_blocks = Some(counts); + self + } + + pub(crate) fn declared_inputs(&self) -> WorkerInputs { + match &self.state { + #[cfg(any(test, feature = "bench"))] + WorkerSelectionPolicyState::Reference(..) => WorkerInputs::CACHE | WorkerInputs::LOAD, + WorkerSelectionPolicyState::Composed(state) => { + let state = state.borrow(); + state.filter_inputs | state.scorer_picker_inputs + } + } + } + /// Construct the native reference implementation for parity tests and benchmarks. /// /// `worker_label` selects the built-in scoring and logging contract. Typed hosts use @@ -140,6 +162,7 @@ impl WorkerSelectionPolicy { Self { worker_label, exclusive_affinity: false, + resident_blocks: None, state: WorkerSelectionPolicyState::Reference(Box::new(kv_router_config), picker), } } @@ -339,14 +362,7 @@ impl WorkerSelector for WorkerSelectionPolicy { } fn required_worker_inputs(&self) -> WorkerInputs { - match &self.state { - #[cfg(any(test, feature = "bench"))] - WorkerSelectionPolicyState::Reference(..) => WorkerInputs::CACHE | WorkerInputs::LOAD, - WorkerSelectionPolicyState::Composed(state) => { - let state = state.borrow(); - state.filter_inputs | state.scorer_picker_inputs - } - } + self.declared_inputs() } #[inline(always)] @@ -364,6 +380,11 @@ impl WorkerSelector for WorkerSelectionPolicy { WorkerSelectionPolicyStateRef::Composed(state) } }; + // Held here so the per-selection snapshot borrows it and stays plain data. + let resident_blocks = self + .resident_blocks + .as_ref() + .and_then(ResidentBlockCountsHandle::load); select_worker_with_policy( self.worker_label, state, @@ -371,6 +392,7 @@ impl WorkerSelector for WorkerSelectionPolicy { request, eligibility, block_size, + resident_blocks.as_deref(), ) } } @@ -656,6 +678,173 @@ mod tests { assert_eq!(workers[&0].taint_reads.get(), 1); } + #[test] + fn resident_blocks_reach_only_declaring_components() { + use crate::indexer::{ResidentBlockCountsHandle, WorkerLookupStats}; + + // Worker 8 holds no tracked blocks, which reads as 0 rather than None. Worker ids differ + // from row indices, so a row-for-worker mix-up fails. + fn held(worker: WorkerWithDpRank) -> u64 { + if worker.worker_id == 3 { 12 } else { 0 } + } + + struct ResidencyProbe(WorkerInputs); + + impl ResidencyProbe { + fn check(&self, candidate: WorkerCandidate<'_>) { + let expected = self + .0 + .contains(WorkerInputs::RESIDENT_BLOCKS) + .then_some(held(candidate.worker())); + assert_eq!(candidate.resident_blocks(), expected); + } + } + + impl WorkerFilter for ResidencyProbe { + fn required_worker_inputs(&self) -> WorkerInputs { + self.0 + } + + fn keep( + &mut self, + _context: &WorkerSelectionContext<'_>, + candidate: WorkerCandidate<'_>, + ) -> Result { + self.check(candidate); + Ok(true) + } + } + + impl WorkerScorer for ResidencyProbe { + fn required_worker_inputs(&self) -> WorkerInputs { + self.0 + } + + fn score( + &mut self, + _context: &WorkerSelectionContext<'_>, + candidates: WorkerCandidates<'_>, + costs: &mut [f64], + ) -> Result<(), WorkerSelectionPolicyError> { + for (candidate, cost) in candidates.iter().zip(costs) { + self.check(candidate); + *cost = 0.0; + } + Ok(()) + } + } + + struct FewestResidentPicker; + + impl WorkerPicker for FewestResidentPicker { + fn required_worker_inputs(&self) -> WorkerInputs { + WorkerInputs::RESIDENT_BLOCKS + } + + fn pick( + &mut self, + _context: &WorkerSelectionContext<'_>, + input: WorkerInputView<'_>, + ) -> Result { + let rows = input.candidates().len(); + assert_eq!(input.resident_blocks(rows), None); + for (row, candidate) in input.candidates().iter().enumerate() { + assert_eq!(input.resident_blocks(row), Some(held(candidate.worker()))); + } + Ok((0..rows) + .min_by_key(|&row| input.resident_blocks(row).unwrap()) + .unwrap()) + } + } + + let workers = HashMap::from([ + (3, TaintedWorkerConfig::default()), + (8, TaintedWorkerConfig::default()), + ]); + let request = base_request(16); + let counts = ResidentBlockCountsHandle::default(); + counts.publish( + WorkerLookupStats::from_worker_block_counts([(WorkerWithDpRank::new(3, 0), 12)]).into(), + ); + let probes = || -> Vec> { + vec![ + Box::new(ResidencyProbe(WorkerInputs::RESIDENT_BLOCKS)), + Box::new(ResidencyProbe(WorkerInputs::NONE)), + ] + }; + let policy = WorkerSelectionPolicy::new_with_filters( + KvRouterConfig::default(), + "test", + probes() + .into_iter() + .map(|probe| probe as Box) + .collect(), + probes() + .into_iter() + .map(|probe| probe as Box) + .collect(), + Box::new(FewestResidentPicker), + ) + .with_resident_block_counts(counts); + + let selected = policy + .select_worker(WorkerSelectionInput::configured( + &workers, + &request, + request.eligibility(), + 16, + )) + .unwrap(); + assert_eq!(selected.worker, WorkerWithDpRank::new(8, 0)); + } + + #[test] + fn resident_blocks_are_none_without_published_counts() { + use crate::indexer::ResidentBlockCountsHandle; + + struct ExpectNoCounts; + + impl WorkerPicker for ExpectNoCounts { + fn required_worker_inputs(&self) -> WorkerInputs { + WorkerInputs::RESIDENT_BLOCKS + } + + fn pick( + &mut self, + _context: &WorkerSelectionContext<'_>, + input: WorkerInputView<'_>, + ) -> Result { + assert_eq!(input.resident_blocks(0), None); + Ok(0) + } + } + + let workers = HashMap::from([(0, TaintedWorkerConfig::default())]); + let request = base_request(16); + let policy = || { + WorkerSelectionPolicy::new( + KvRouterConfig::default(), + "test", + Vec::new(), + Box::new(ExpectNoCounts), + ) + }; + // No host handle, then a handle whose poller has not published yet. + for policy in [ + policy(), + policy().with_resident_block_counts(ResidentBlockCountsHandle::default()), + ] { + policy + .select_worker(WorkerSelectionInput::configured( + &workers, + &request, + request.eligibility(), + 16, + )) + .unwrap(); + } + } + #[test] fn custom_policy_does_not_receive_effective_overlap_as_device_overlap() { struct RawDeviceOverlapPicker; diff --git a/lib/kv-router/src/scheduling/selector/reference.rs b/lib/kv-router/src/scheduling/selector/reference.rs index 41d87b956127..e34a649b7568 100644 --- a/lib/kv-router/src/scheduling/selector/reference.rs +++ b/lib/kv-router/src/scheduling/selector/reference.rs @@ -601,6 +601,7 @@ impl WorkerSelector for DefaultWorkerSelector { request, eligibility, block_size, + None, ) } } diff --git a/lib/kv-router/src/services/indexer/backend.rs b/lib/kv-router/src/services/indexer/backend.rs index 72fba144b224..857609b1d00d 100644 --- a/lib/kv-router/src/services/indexer/backend.rs +++ b/lib/kv-router/src/services/indexer/backend.rs @@ -15,8 +15,8 @@ use crate::config::{ApproximateCachePolicyKind, KvRouterConfig}; use crate::indexer::{ ApproximateLruIncarnation, ApproximateLruStats, ApproximateRetentionConfig, KvIndexer, KvIndexerInterface, KvIndexerMetrics, KvRouterError, LowerTierIndexers, LowerTierQueryOptions, - MatchDetails, RoutingDecisionHashes, SyncIndexer, TieredMatchDetails, TieredMatchProvider, - record_unsupported_residency_event, + MatchDetails, ResidentBlockCounts, ResidentBlockCountsHandle, RoutingDecisionHashes, + SyncIndexer, TieredMatchDetails, TieredMatchProvider, record_unsupported_residency_event, }; use crate::protocols::{ DpRank, ExternalSequenceBlockHash, KvCacheEventData, LocalBlockHash, OverlapScores, @@ -870,6 +870,62 @@ impl Indexer { Self::Remote { .. } | Self::None => Ok(ApproximateLruStats::default()), } } + + /// Distinct blocks the primary holds for each worker rank, or `None` without a + /// local primary. The read queues behind pending events, so keep it off the + /// request path. + #[cfg_attr(not(feature = "standalone-selection"), allow(dead_code))] + async fn resident_block_counts(&self) -> Result, KvRouterError> { + let stats = match self { + Self::Single { primary, .. } => primary.worker_lookup_stats().await?, + Self::Concurrent { primary, .. } => primary.try_worker_lookup_stats().await?, + Self::Remote { .. } | Self::None => return Ok(None), + }; + Ok(Some(stats.into())) + } + + /// Republish [`Self::resident_block_counts`] every `interval` until `cancel` fires or the + /// primary stops answering, then withdraw them. `None` without a local primary. + #[cfg_attr(not(feature = "standalone-selection"), allow(dead_code))] + pub(crate) fn spawn_resident_block_counts_poller( + &self, + interval: Duration, + cancel: CancellationToken, + ) -> Option { + if !matches!(self, Self::Single { .. } | Self::Concurrent { .. }) { + return None; + } + let handle = ResidentBlockCountsHandle::default(); + let published = handle.clone(); + let indexer = self.clone(); + tokio::spawn(async move { + let poll = async { + let mut ticks = tokio::time::interval(interval); + ticks.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + loop { + ticks.tick().await; + match indexer.resident_block_counts().await { + Ok(Some(counts)) => published.publish(counts), + Ok(None) => return, + Err(error) => { + // Debug: the indexer logs its own failure, and an ordinary shutdown + // can stop the indexer before this task's cancel fires. + tracing::debug!(%error, "Stopped publishing resident block counts"); + return; + } + } + } + }; + // Biased, so a cancel that races a read failed by the same shutdown wins. + tokio::select! { + biased; + _ = cancel.cancelled() => {} + _ = poll => {} + } + published.clear(); + }); + Some(handle) + } } #[async_trait] @@ -1298,6 +1354,12 @@ mod tests { ); let flat = indexer.find_matches(sequence).await.unwrap(); assert_eq!(flat.scores.get(&worker).copied(), Some(2)); + let counts = indexer.resident_block_counts().await.unwrap().unwrap(); + assert_eq!( + counts.get(worker), + 2, + "{num_threads} thread(s): routed prefix" + ); // The recorded prefix is dumped like any other primary state, so a // recovering peer inherits it. @@ -1581,6 +1643,112 @@ mod tests { ); } } + + async fn wait_until(mut done: impl FnMut() -> bool) { + tokio::time::timeout(Duration::from_secs(5), async { + while !done() { + tokio::time::sleep(Duration::from_millis(5)).await; + } + }) + .await + .expect("condition not reached"); + } + + #[tokio::test] + async fn resident_block_counts_track_each_rank() { + use crate::protocols::compute_seq_hash_for_block; + use crate::test_utils::make_store_event_with_dp_rank; + + for num_threads in [1, 4] { + let indexer = create_indexer(16, num_threads); + for event in [ + make_store_event_with_dp_rank(1, &[1, 2, 3], 0), + // A repeated store must not count its blocks twice. + make_store_event_with_dp_rank(1, &[1, 2, 3], 0), + make_store_event_with_dp_rank(1, &[1, 2], 1), + make_store_event_with_dp_rank(2, &[1], 0), + ] { + indexer.apply_event(event).await; + } + let chain = [LocalBlockHash(1), LocalBlockHash(2), LocalBlockHash(3)]; + let tail = ExternalSequenceBlockHash(compute_seq_hash_for_block(&chain)[2]); + indexer + .apply_event(crate::test_utils::remove_event(1, 1, 0, vec![tail])) + .await; + flush(&indexer).await; + + let counts = indexer.resident_block_counts().await.unwrap().unwrap(); + let held = |worker_id, dp_rank| counts.get(WorkerWithDpRank::new(worker_id, dp_rank)); + assert_eq!( + [held(1, 0), held(1, 1), held(2, 0), held(3, 0)], + [2, 2, 1, 0], + "{num_threads} indexer thread(s)" + ); + + indexer.remove_worker_dp_rank(1, 1).await; + flush(&indexer).await; + let counts = indexer.resident_block_counts().await.unwrap().unwrap(); + let held = |worker_id, dp_rank| counts.get(WorkerWithDpRank::new(worker_id, dp_rank)); + assert_eq!( + [held(1, 0), held(1, 1)], + [2, 0], + "{num_threads} indexer thread(s)" + ); + } + } + + #[tokio::test] + async fn resident_block_counts_poller_publishes_then_withdraws() { + let indexer = create_indexer(16, 4); + let cancel = CancellationToken::new(); + let handle = indexer + .spawn_resident_block_counts_poller(Duration::from_millis(5), cancel.clone()) + .expect("a local primary has counts to poll"); + let worker = WorkerWithDpRank::new(1, 0); + // Store only after the first publish, so seeing it takes a later poll. + wait_until(|| handle.load().is_some()).await; + indexer + .apply_event(crate::test_utils::make_store_event_with_dp_rank( + 1, + &[1, 2], + 0, + )) + .await; + wait_until(|| handle.load().is_some_and(|counts| counts.get(worker) == 2)).await; + + cancel.cancel(); + wait_until(|| handle.load().is_none()).await; + } + + #[tokio::test] + async fn resident_block_counts_poller_withdraws_when_the_primary_stops() { + for num_threads in [1, 4] { + let indexer = create_indexer(16, num_threads); + let handle = indexer + .spawn_resident_block_counts_poller( + Duration::from_millis(5), + CancellationToken::new(), + ) + .expect("a local primary has counts to poll"); + wait_until(|| handle.load().is_some()).await; + match &indexer { + Indexer::Single { primary, .. } => primary.shutdown(), + Indexer::Concurrent { primary, .. } => primary.shutdown(), + Indexer::Remote { .. } | Indexer::None => unreachable!(), + } + // A dead primary reads as None, not as a frozen or all-zero snapshot. + wait_until(|| handle.load().is_none()).await; + } + } + + #[tokio::test] + async fn resident_block_counts_need_a_local_primary() { + let indexer = Indexer::None; + assert!(indexer.resident_block_counts().await.unwrap().is_none()); + let poller = indexer + .spawn_resident_block_counts_poller(Duration::from_millis(5), CancellationToken::new()); + assert!(poller.is_none()); + } } pub fn create_indexer(block_size: u32, num_threads: usize) -> Indexer { diff --git a/lib/kv-router/src/services/selection/core/tests.rs b/lib/kv-router/src/services/selection/core/tests.rs index bf89fa75c03f..1139372f7ded 100644 --- a/lib/kv-router/src/services/selection/core/tests.rs +++ b/lib/kv-router/src/services/selection/core/tests.rs @@ -734,6 +734,79 @@ fn capturing_policy_factory() -> ( (factory, observed) } +/// Picker that records the resident block count of row 0 and takes it. +struct ResidentBlocksRecorder { + observed: Arc>>>, +} + +impl crate::scheduling::selector::WorkerPicker for ResidentBlocksRecorder { + fn required_worker_inputs(&self) -> crate::scheduling::selector::WorkerInputs { + crate::scheduling::selector::WorkerInputs::CACHE + | crate::scheduling::selector::WorkerInputs::RESIDENT_BLOCKS + } + + fn pick( + &mut self, + _context: &crate::scheduling::selector::WorkerSelectionContext<'_>, + input: crate::scheduling::selector::WorkerInputView<'_>, + ) -> Result { + self.observed.lock().push(input.resident_blocks(0)); + Ok(0) + } +} + +#[tokio::test] +async fn declaring_policy_reads_polled_resident_blocks() { + let observed = Arc::new(parking_lot::Mutex::new(Vec::new())); + let factory_observed = Arc::clone(&observed); + let factory: WorkerSelectionPolicyFactory = Arc::new(move |config, worker_type, _partition| { + WorkerSelectionPolicy::new( + config.clone(), + worker_type.as_str(), + Vec::new(), + Box::new(ResidentBlocksRecorder { + observed: Arc::clone(&factory_observed), + }), + ) + }); + let core = core_with( + test_config(true), + SelectionHost::default(), + Some(factory), + WorkerType::Aggregated, + None, + ); + core.upsert_worker(worker_with_kv_events(1)) + .await + .expect("worker upsert"); + let entry = core.entry(&default_key()).expect("entry"); + entry + .indexer + .apply_event_routed(store_event( + 1, + 0, + 1, + &[], + &[11, 12, 13], + StorageTier::Device, + )) + .await + .unwrap(); + + // The partition's poller publishes in the background, so select until the count arrives. + tokio::time::timeout(Duration::from_secs(5), async { + loop { + core.select(select_request()).await.expect("select"); + if observed.lock().last() == Some(&Some(3)) { + break; + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + }) + .await + .expect("a declaring policy should read the polled count"); +} + type SharedCacheCalls = Arc, u32, Option)>>>; /// Shared cache that reports every block as a hit and records each query. diff --git a/lib/kv-router/src/services/selection/core/workers.rs b/lib/kv-router/src/services/selection/core/workers.rs index abd0bedebd99..ad077e9b0487 100644 --- a/lib/kv-router/src/services/selection/core/workers.rs +++ b/lib/kv-router/src/services/selection/core/workers.rs @@ -8,6 +8,8 @@ use std::collections::HashSet; use super::reservations::{ReservationIndexObserver, spawn_reservation_index_sweep}; use super::*; +use crate::indexer::RESIDENT_BLOCK_COUNTS_POLL_INTERVAL; +use crate::plugins::worker_selection::WorkerInputs; use crate::sequences::topology::MAX_DATA_PARALLEL_RANKS_PER_WORKER; impl SelectionCore { @@ -324,7 +326,7 @@ impl SelectionCore { block_size, )) }); - let selector = (self.worker_selection_policy_factory)( + let mut selector = (self.worker_selection_policy_factory)( &self.kv_router_config, self.worker_type, key.as_ref(), @@ -333,6 +335,17 @@ impl SelectionCore { .kv_router_config .policy_profile(Some(&key.model_name)) .map_err(|error| SelectionError::BadRequest(error.to_string()))?; + // Poll this partition's indexer only for a policy that reads the counts. + if selector + .declared_inputs() + .contains(WorkerInputs::RESIDENT_BLOCKS) + && let Some(counts) = indexer.spawn_resident_block_counts_poller( + RESIDENT_BLOCK_COUNTS_POLL_INTERVAL, + self.cancel_token.child_token(), + ) + { + selector = selector.with_resident_block_counts(counts); + } let scheduler = LocalScheduler::new( slots, workers_rx, diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index de86f332caf1..9a3144b9f674 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -157,6 +157,11 @@ impl SelectionPolicySource { !config.enable_session_prefix_index, "enable_session_prefix_index requires the worker-selection policy to declare WorkerInputs::CACHE" ); + // Resident block counts come from the indexer that CACHE starts. + anyhow::ensure!( + !prepared.inputs().contains(WorkerInputs::RESIDENT_BLOCKS), + "RESIDENT_BLOCKS requires the worker-selection policy to declare WorkerInputs::CACHE" + ); } Ok(prepared) } @@ -2227,6 +2232,22 @@ mod tests { } } + struct DeclaredInputsPicker(WorkerInputs); + + impl WorkerPicker for DeclaredInputsPicker { + fn required_worker_inputs(&self) -> WorkerInputs { + self.0 + } + + fn pick( + &mut self, + _context: &WorkerSelectionContext<'_>, + _input: WorkerInputView<'_>, + ) -> Result { + unreachable!("capability construction test does not select a worker") + } + } + fn fixed_policy( expected_shared_blocks: Option, worker: WorkerWithDpRank, @@ -2459,6 +2480,24 @@ mod tests { } } + #[tokio::test] + async fn resident_blocks_require_declared_cache_input() { + let config = KvRouterConfig::default(); + let error = picker_policy(|| Box::new(DeclaredInputsPicker(WorkerInputs::RESIDENT_BLOCKS))) + .prepare(&config, WorkerType::Aggregated, "decode", Some("model")) + .err() + .expect("without CACHE there is no indexer to count blocks"); + assert!(error.to_string().contains("RESIDENT_BLOCKS")); + assert!(error.to_string().contains("WorkerInputs::CACHE")); + picker_policy(|| { + Box::new(DeclaredInputsPicker( + WorkerInputs::RESIDENT_BLOCKS | WorkerInputs::CACHE, + )) + }) + .prepare(&config, WorkerType::Aggregated, "decode", Some("model")) + .expect("CACHE starts the indexer the counts come from"); + } + #[tokio::test] async fn load_only_selector_skips_cache_inputs() { let router = make_router( From 51b5f93124ea891cf0ed28db1137f9a5ff49367f Mon Sep 17 00:00:00 2001 From: Akira Lonske Date: Sat, 10 Oct 2026 19:28:21 -0400 Subject: [PATCH 02/12] fix(kv-router): withdraw stale resident block counts and detect exited threads - Withdraw the counts while a read takes longer than three poll intervals, and republish once a read completes in time. The poller waits out a slow read instead of queueing more behind the backlog. - Answer count requests ahead of mutations and matches in the single-thread indexer, so steady traffic cannot starve them. - Cancel a per-thread token when a thread-pool event thread exits, and race each stats reply against it. Both stats reads share this and wait on all threads at once, so the lenient read no longer hangs on an exited thread. - Build the count map directly from the indexer's replies. - Borrow the published counts through an arc-swap guard. - Reject RESIDENT_BLOCKS without CACHE whenever a partition's policy is built, not only for the frontend's own partition at startup. - Replace the backlog-based thread-exit test with one that holds the thread until both reads are queued. Signed-off-by: Akira Lonske --- .../router/custom-worker-selection.mdx | 6 +- .../router/custom-policy-example/README.md | 2 +- lib/kv-router/src/indexer/kv_indexer.rs | 31 ++-- lib/kv-router/src/indexer/radix_tree.rs | 13 +- lib/kv-router/src/indexer/resident_blocks.rs | 18 +-- lib/kv-router/src/indexer/thread_pool.rs | 135 ++++++++++------- lib/kv-router/src/indexer/types.rs | 4 +- .../src/plugins/worker_selection/inputs.rs | 4 +- .../src/scheduling/selector/policy.rs | 12 +- lib/kv-router/src/services/indexer/backend.rs | 141 ++++++++++++++---- .../src/services/selection/core/tests.rs | 40 ++++- .../src/services/selection/core/workers.rs | 13 +- 12 files changed, 277 insertions(+), 142 deletions(-) diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx index 6e59c93cb56e..c03b0d26f7a8 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx @@ -370,7 +370,7 @@ Each worker policy is constructed once per routing partition. An empty bundle re Worker identity and ordinary request metadata are always available. Dynamo calculates optional per-worker signals only for groups that a filter, scorer, or picker requests. -Routing cache resources start only if a component requests `CACHE`. `serve_indexer` and `enable_session_prefix_index` also require `CACHE`, and the frontend rejects `RESIDENT_BLOCKS` without `CACHE`. +Routing cache resources start only if a component requests `CACHE`. `serve_indexer`, `enable_session_prefix_index`, and `RESIDENT_BLOCKS` also require `CACHE`. ### Request Context @@ -446,7 +446,7 @@ If a component needs no optional worker data, return `WorkerInputs::NONE`. Combi | `LOAD` | `decode_cost_blocks()` | Projected active decode footprint, including this request's additional active blocks | | `LOAD` | `active_requests()` | Requests currently active on the worker | | `PREFERRED_TAINT` | `preferred_taint_multiplier()` | Optional cost multiplier from matching preferred routing taints | -| `RESIDENT_BLOCKS` | `resident_blocks()` | Blocks the router's indexer tracks for the worker rank, including cached blocks no running request uses; refreshed at most every 100 ms and not clamped to capacity | +| `RESIDENT_BLOCKS` | `resident_blocks()` | Blocks the router's indexer tracks for the worker rank, including cached blocks no running request uses; refreshed every 100 ms and not clamped to capacity | The GPU, CPU, and disk overlap accessors return raw tier counts. A missing tier or worker entry is zero; Dynamo does not substitute its weighted effective-overlap estimate. Each custom scorer chooses how to combine the raw counts. @@ -454,7 +454,7 @@ A filter or scorer reads groups through `WorkerCandidate::cache()` or `load()`. `WorkerCandidate::preferred_taint_multiplier()` and `ScoredWorkerCandidate::preferred_taint_multiplier()` return the optional cost multiplier from preferred routing constraints. A filter, scorer, or picker must request `WorkerInputs::PREFERRED_TAINT` before reading the multiplier. Without that declaration, the component receives `None`, even when another component requested the multiplier. Exact hard-pinned requests also do not materialize it. Required routing constraints remain host-enforced eligibility rules. -`WorkerCandidate::resident_blocks()` and `WorkerInputView::resident_blocks(row)` return how many blocks the router's indexer tracks for the worker rank, including cached blocks no running request uses; use `LOAD` for work in flight. The router refreshes the counts off the request path at most every 100 ms, later while the indexer is backlogged, so they miss routing decisions made since the last refresh. Every component sees the same counts during one selection. With KV events, the count is the device-tier blocks the indexer currently tracks from engine events. With approximate routing, it is the blocks routed to the rank within the TTL, or the distinct full blocks the LRU cache policy retains, which leaves out its partial tail blocks and duplicate copies. The count is unweighted and not clamped to `total_kv_blocks`. A rank with no tracked blocks reads `0`, including a rank the indexer has no data for yet, such as a worker whose first KV events have not arrived. A component that did not request `WorkerInputs::RESIDENT_BLOCKS` receives `None`, and so does every component before the first refresh completes, with a remote indexer, or after the indexer stops answering. The frontend rejects a policy that requests `RESIDENT_BLOCKS` without `CACHE`, because the counts come from the indexer that `CACHE` starts. With KV events every router replica sees the same counts; with approximate routing each replica counts only the requests it routed. +`WorkerCandidate::resident_blocks()` and `WorkerInputView::resident_blocks(row)` return how many blocks the router's indexer tracks for the worker rank, including cached blocks no running request uses; use `LOAD` for work in flight. The router refreshes the counts off the request path every 100 ms, so they miss routing decisions made since the last refresh. If the indexer is backlogged and a refresh takes longer than 300 ms, components receive `None` until a refresh completes in time. Every component sees the same counts during one selection. With KV events, the count is the device-tier blocks the indexer currently tracks from engine events. With approximate routing, it is the blocks routed to the rank within the TTL, or the distinct full blocks the LRU cache policy retains, which leaves out its partial tail blocks and duplicate copies. The count is unweighted and not clamped to `total_kv_blocks`. A rank with no tracked blocks reads `0`, including a rank the indexer has no data for yet, such as a worker whose first KV events have not arrived. A component that did not request `WorkerInputs::RESIDENT_BLOCKS` receives `None`, and so does every component before the first refresh completes, while a refresh is overdue, with a remote indexer, or after the indexer stops answering. A policy that requests `RESIDENT_BLOCKS` must also request `CACHE`, because the counts come from the indexer that `CACHE` starts. The router checks this whenever it builds a policy for a routing partition: the frontend fails to start, and the standalone selection service marks the partition's workers unschedulable with the reason. With KV events every router replica sees the same counts; with approximate routing each replica counts only the requests it routed. ## Link the Policy Into Dynamo diff --git a/examples/router/custom-policy-example/README.md b/examples/router/custom-policy-example/README.md index 4a4f39d20542..d549e578db8b 100644 --- a/examples/router/custom-policy-example/README.md +++ b/examples/router/custom-policy-example/README.md @@ -24,7 +24,7 @@ Declare `WorkerInputs::CACHE` to read cache data through `candidate.cache()` (fi Preferred routing taints are optional candidate metadata. A filter, scorer, or picker must request `WorkerInputs::PREFERRED_TAINT` before reading `preferred_taint_multiplier()` from a candidate; otherwise, that component receives `None`, even if another component requested it. Exact hard-pinned requests also do not materialize it. Required routing taints remain Dynamo eligibility rules. -Declare `WorkerInputs::RESIDENT_BLOCKS` to read `resident_blocks()` from a candidate, or `input.resident_blocks(row)` in a picker. It returns how many blocks the router's indexer tracks for that worker rank, including cached blocks no running request uses, refreshed in the background. Without the declaration, before the first refresh, or with a remote indexer, it returns `None`. The frontend also requires `CACHE` when a policy requests `RESIDENT_BLOCKS`. +Declare `WorkerInputs::RESIDENT_BLOCKS` to read `resident_blocks()` from a candidate, or `input.resident_blocks(row)` in a picker. It returns how many blocks the router's indexer tracks for that worker rank, including cached blocks no running request uses, refreshed in the background. Without the declaration, before the first refresh, while the indexer is backlogged, or with a remote indexer, it returns `None`. A policy that requests `RESIDENT_BLOCKS` must also request `CACHE`. ## Pick a Starting Point diff --git a/lib/kv-router/src/indexer/kv_indexer.rs b/lib/kv-router/src/indexer/kv_indexer.rs index afe4cdf59f3c..88718c1d38e3 100644 --- a/lib/kv-router/src/indexer/kv_indexer.rs +++ b/lib/kv-router/src/indexer/kv_indexer.rs @@ -15,8 +15,8 @@ use super::{ ApproximateLruLease, ApproximateLruStats, ApproximateLruTask, ApproximateRetentionConfig, ContainsWorkerBlockRequest, DumpRequest, EventKind, FlushRequest, GetWorkersRequest, KvIndexerInterface, KvIndexerMetrics, KvRouterError, MatchDetails, MatchDetailsRequest, - MatchRequest, PreBoundEventCounters, RadixTree, RoutingDecisionRequest, WorkerLookupStats, - WorkerLookupStatsRequest, panic_payload_message, + MatchRequest, PreBoundEventCounters, RadixTree, ResidentBlockCounts, + ResidentBlockCountsRequest, RoutingDecisionRequest, panic_payload_message, }; use crate::indexer::pruning::{BlockEntry, PruneConfig, WorkerPruneManager}; use crate::protocols::*; @@ -339,7 +339,7 @@ pub struct KvIndexer { get_workers_tx: mpsc::Sender, /// A sender for per-rank block count requests. #[cfg_attr(not(feature = "standalone-indexer"), allow(dead_code))] - worker_lookup_stats_tx: mpsc::Sender, + resident_block_counts_tx: mpsc::Sender, /// A sender for dump requests. dump_tx: mpsc::Sender, /// A sender for flush requests. @@ -434,8 +434,8 @@ impl KvIndexer { let (remove_worker_dp_rank_tx, remove_worker_dp_rank_rx) = mpsc::channel::<(WorkerId, DpRank)>(16); let (get_workers_tx, get_workers_rx) = mpsc::channel::(16); - let (worker_lookup_stats_tx, worker_lookup_stats_rx) = - mpsc::channel::(16); + let (resident_block_counts_tx, resident_block_counts_rx) = + mpsc::channel::(16); let (dump_tx, dump_rx) = mpsc::channel::(16); let (flush_tx, flush_rx) = mpsc::channel::(16); let (routing_tx, mut routing_rx) = mpsc::channel::(2048); @@ -466,7 +466,7 @@ impl KvIndexer { let mut remove_worker_rx = remove_worker_rx; let mut remove_worker_dp_rank_rx = remove_worker_dp_rank_rx; let mut get_workers_rx = get_workers_rx; - let mut worker_lookup_stats_rx = worker_lookup_stats_rx; + let mut resident_block_counts_rx = resident_block_counts_rx; let mut dump_rx = dump_rx; let mut flush_rx = flush_rx; let mut trie = delegate.map_or_else(RadixTree::new, RadixTree::new_with_delegate); @@ -505,6 +505,11 @@ impl KvIndexer { } } + // Ahead of mutations and matches: under steady traffic, later arms may never run. + Some(req) = resident_block_counts_rx.recv() => { + let _ = req.resp.send(trie.resident_block_counts()); + } + Some(mutation) = mutation_rx.recv() => { apply_mutation(&mut trie, mutation, &counters, &prune_manager); } @@ -566,10 +571,6 @@ impl KvIndexer { let _ = get_workers_req.resp.send(workers); } - Some(stats_req) = worker_lookup_stats_rx.recv() => { - let _ = stats_req.resp.send(trie.worker_lookup_stats()); - } - Some(dump_req) = dump_rx.recv() => { // NOTE: Dump requests use a separate channel from mutations, so the // actor may observe one while already-accepted mutations are still @@ -667,7 +668,7 @@ impl KvIndexer { remove_worker_tx, remove_worker_dp_rank_tx, get_workers_tx, - worker_lookup_stats_tx, + resident_block_counts_tx, dump_tx, flush_tx, routing_tx, @@ -789,16 +790,16 @@ impl KvIndexer { /// Distinct blocks the tree holds for each worker rank. #[cfg_attr(not(feature = "standalone-indexer"), allow(dead_code))] - pub(crate) async fn worker_lookup_stats(&self) -> Result { + pub(crate) async fn resident_block_counts(&self) -> Result { let (resp_tx, resp_rx) = oneshot::channel(); - self.worker_lookup_stats_tx - .send(WorkerLookupStatsRequest { resp: resp_tx }) + self.resident_block_counts_tx + .send(ResidentBlockCountsRequest { resp: resp_tx }) .await .map_err(|_| KvRouterError::IndexerOffline)?; // A request sent while the actor drops its receiver stays queued with resp_tx alive. tokio::select! { resp = resp_rx => resp.map_err(|_| KvRouterError::IndexerDroppedRequest), - _ = self.worker_lookup_stats_tx.closed() => Err(KvRouterError::IndexerOffline), + _ = self.resident_block_counts_tx.closed() => Err(KvRouterError::IndexerOffline), } } } diff --git a/lib/kv-router/src/indexer/radix_tree.rs b/lib/kv-router/src/indexer/radix_tree.rs index 3dd8044c3525..8f8cf9f72bb8 100644 --- a/lib/kv-router/src/indexer/radix_tree.rs +++ b/lib/kv-router/src/indexer/radix_tree.rs @@ -13,7 +13,7 @@ use std::{ use rustc_hash::{FxBuildHasher, FxHashMap, FxHashSet}; use super::compressed_radix::{NodeState, append_dump_events}; -use super::{EventWarningKind, MatchDetails, PreBoundEventCounters, WorkerLookupStats}; +use super::{EventWarningKind, MatchDetails, PreBoundEventCounters, ResidentBlockCounts}; use crate::protocols::*; pub(crate) type SharedRadixBlock = Rc>; @@ -861,12 +861,11 @@ impl RadixTree { } /// Distinct blocks held by each worker rank. - pub(crate) fn worker_lookup_stats(&self) -> WorkerLookupStats { - WorkerLookupStats::from_worker_block_counts( - self.lookup - .iter() - .map(|(worker, worker_lookup)| (*worker, worker_lookup.len())), - ) + pub(crate) fn resident_block_counts(&self) -> ResidentBlockCounts { + self.lookup + .iter() + .map(|(worker, worker_lookup)| (*worker, worker_lookup.len())) + .collect() } #[cfg(test)] diff --git a/lib/kv-router/src/indexer/resident_blocks.rs b/lib/kv-router/src/indexer/resident_blocks.rs index 49215056575e..fd28d77f7151 100644 --- a/lib/kv-router/src/indexer/resident_blocks.rs +++ b/lib/kv-router/src/indexer/resident_blocks.rs @@ -10,10 +10,9 @@ use std::sync::Arc; use std::time::Duration; -use arc_swap::ArcSwapOption; +use arc_swap::{ArcSwapOption, Guard}; use rustc_hash::FxHashMap; -use super::WorkerLookupStats; use crate::protocols::WorkerWithDpRank; /// How often routing hosts republish [`ResidentBlockCounts`]. @@ -33,10 +32,10 @@ impl ResidentBlockCounts { } } -impl From for ResidentBlockCounts { - fn from(stats: WorkerLookupStats) -> Self { +impl FromIterator<(WorkerWithDpRank, usize)> for ResidentBlockCounts { + fn from_iter>(iter: I) -> Self { let mut counts = FxHashMap::default(); - for (worker, blocks) in stats.worker_blocks { + for (worker, blocks) in iter { *counts.entry(worker).or_default() += blocks as u64; } Self { counts } @@ -50,9 +49,10 @@ pub(crate) struct ResidentBlockCountsHandle { } impl ResidentBlockCountsHandle { - /// The latest counts, or `None` before the first poll or after the poller stops. - pub(crate) fn load(&self) -> Option> { - self.latest.load_full() + /// The latest counts, or `None` before the first poll, while the indexer is backlogged, or + /// after the poller stops. Hold the guard only for one selection. + pub(crate) fn load(&self) -> Guard>> { + self.latest.load() } #[cfg_attr(not(feature = "standalone-indexer"), allow(dead_code))] @@ -60,7 +60,7 @@ impl ResidentBlockCountsHandle { self.latest.store(Some(Arc::new(counts))); } - /// Withdraw the counts so readers see `None` rather than a frozen snapshot. + /// Withdraw the counts so readers see `None` rather than stale counts. #[cfg_attr(not(feature = "standalone-indexer"), allow(dead_code))] pub(crate) fn clear(&self) { self.latest.store(None); diff --git a/lib/kv-router/src/indexer/thread_pool.rs b/lib/kv-router/src/indexer/thread_pool.rs index 430bfc123943..e3bb41c60b67 100644 --- a/lib/kv-router/src/indexer/thread_pool.rs +++ b/lib/kv-router/src/indexer/thread_pool.rs @@ -19,8 +19,8 @@ use super::concurrent_radix_tree_compressed::ConcurrentRadixTreeCompressed; use super::{ ApproximateLruClient, ApproximateLruCommandSink, ApproximateLruIncarnation, ApproximateLruLease, ApproximateLruStats, ApproximateLruTask, ApproximateRetentionConfig, - KvIndexerInterface, KvIndexerMetrics, KvRouterError, ShardSizeSnapshot, SyncIndexer, - WorkerLookupStats, WorkerTask, panic_payload_message, + KvIndexerInterface, KvIndexerMetrics, KvRouterError, ResidentBlockCounts, ShardSizeSnapshot, + SyncIndexer, WorkerLookupStats, WorkerTask, panic_payload_message, }; #[cfg(feature = "bench")] use super::{ @@ -83,6 +83,9 @@ pub struct ThreadPoolIndexer { /// Handles to worker threads for joining on shutdown. thread_handles: Mutex>>, + /// Cancelled when the matching worker thread exits, including by panic. + thread_exited: Vec, + /// Approximate-mode TTL pruning manager. None for normal event-driven mode. prune_manager: Option, @@ -267,17 +270,21 @@ impl ThreadPoolIndexer { let backend = Arc::new(backend); let mut worker_event_senders = Vec::new(); let mut thread_handles = Vec::new(); + let mut thread_exited = Vec::new(); let worker_assignments = Arc::new(DashMap::with_hasher(FxBuildHasher)); let worker_assignment_count = Arc::new(AtomicUsize::new(0)); let synthetic_event_id = Arc::new(AtomicU64::new(0)); for worker_idx in 0..num_workers { let (event_sender, event_receiver) = flume::unbounded::(); worker_event_senders.push(event_sender); + let exited = tokio_util::sync::CancellationToken::new(); + thread_exited.push(exited.clone()); let backend = Arc::clone(&backend); let metrics = metrics.clone(); let handle = std::thread::spawn(move || { + let _exited = exited.drop_guard(); // This is observability, not recovery: if the worker panics, log // through tracing and then preserve the panic for join(). let panic_result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { @@ -327,6 +334,7 @@ impl ThreadPoolIndexer { num_workers, kv_block_size, thread_handles: Mutex::new(thread_handles), + thread_exited, prune_manager, prune_pump_cancel, synthetic_event_id, @@ -578,57 +586,46 @@ impl ThreadPoolIndexer { } } - pub(crate) async fn worker_lookup_stats(&self) -> WorkerLookupStats { - let mut receivers = Vec::new(); - for channel in &self.worker_event_channels { + /// Each worker thread's per-rank block counts, or `None` for a thread that has exited. + async fn lane_lookup_stats(&self) -> Vec> { + let lanes = self.worker_event_channels.iter().zip(&self.thread_exited); + futures_util::future::join_all(lanes.map(|(channel, exited)| async move { let (resp_tx, resp_rx) = oneshot::channel(); - if channel.send(WorkerTask::Stats(resp_tx)).is_ok() { - receivers.push(resp_rx); + channel.send(WorkerTask::Stats(resp_tx)).ok()?; + // A thread that exits with this request queued neither answers nor drops it, since + // our own senders keep its queue alive. + tokio::select! { + biased; + stats = resp_rx => stats.ok(), + _ = exited.cancelled() => None, } - } + })) + .await + } + pub(crate) async fn worker_lookup_stats(&self) -> WorkerLookupStats { let mut worker_blocks = BTreeMap::new(); - for receiver in receivers { - if let Ok(stats) = receiver.await { - for (worker, block_count) in stats.worker_blocks { - *worker_blocks.entry(worker).or_insert(0usize) += block_count; - } + for stats in self.lane_lookup_stats().await.into_iter().flatten() { + for (worker, block_count) in stats.worker_blocks { + *worker_blocks.entry(worker).or_insert(0usize) += block_count; } } - WorkerLookupStats::from_worker_block_counts(worker_blocks) } - /// Like [`Self::worker_lookup_stats`], but fails if an event thread is gone, so a caller - /// never mistakes a dead thread's ranks for empty ones. Each rank lives on one thread. + /// Per-rank block counts. Fails if a worker thread has exited, so a caller never mistakes a + /// dead thread's ranks for empty ones. #[cfg_attr(not(feature = "standalone-indexer"), allow(dead_code))] - pub(crate) async fn try_worker_lookup_stats(&self) -> Result { - let mut receivers = Vec::with_capacity(self.worker_event_channels.len()); - for channel in &self.worker_event_channels { - let (resp_tx, resp_rx) = oneshot::channel(); - channel - .send(WorkerTask::Stats(resp_tx)) - .map_err(|_| KvRouterError::IndexerOffline)?; - receivers.push(resp_rx); + pub(crate) async fn resident_block_counts(&self) -> Result { + let lanes = self.lane_lookup_stats().await; + if lanes.iter().any(Option::is_none) { + return Err(KvRouterError::IndexerOffline); } - let mut worker_blocks = Vec::new(); - for (channel, mut receiver) in self.worker_event_channels.iter().zip(receivers) { - // A lane that dies with this request queued neither answers nor drops it, since our - // own senders keep the queue alive, so check for the dead lane while waiting. - let stats = loop { - match tokio::time::timeout(std::time::Duration::from_millis(50), &mut receiver) - .await - { - Ok(reply) => break reply.map_err(|_| KvRouterError::IndexerDroppedRequest)?, - Err(_) if channel.is_disconnected() => { - return Err(KvRouterError::IndexerOffline); - } - Err(_) => {} - } - }; - worker_blocks.extend(stats.worker_blocks); - } - Ok(WorkerLookupStats { worker_blocks }) + Ok(lanes + .into_iter() + .flatten() + .flat_map(|stats| stats.worker_blocks) + .collect()) } pub async fn get_workers(&self) -> Vec { @@ -1608,25 +1605,49 @@ mod tests { assert_score(&indexer, &[20], rank1, 1).await; } - #[tokio::test] - async fn strict_stats_fail_when_a_lane_dies_with_the_request_queued() { - let indexer = ThreadPoolIndexer::new(ConcurrentRadixTreeCompressed::new(), 2, 16); - // Keep lane 0 busy so the stats request queues behind its Terminate. - for block in 0..20_000u64 { - indexer - .apply_event(make_store_event_with_dp_rank(7, &[block, block + 1], 0)) - .await; + /// Holds each worker thread until the gate's sender drops, so a test can queue tasks behind + /// `Terminate`. + struct GatedBackend(flume::Receiver<()>); + + impl SyncIndexer for GatedBackend { + fn worker( + &self, + tasks: flume::Receiver, + _metrics: Option>, + ) -> anyhow::Result<()> { + let _ = self.0.recv(); + while let Ok(task) = tasks.recv() { + if let WorkerTask::Terminate = task { + break; + } + } + Ok(()) } + + fn find_matches(&self, _sequence: &[LocalBlockHash], _early_exit: bool) -> OverlapScores { + OverlapScores::default() + } + } + + #[tokio::test] + async fn stats_reads_fail_when_a_thread_exits_with_them_queued() { + let (open_gate, gate) = flume::bounded::<()>(1); + let indexer = ThreadPoolIndexer::new(GatedBackend(gate), 1, 16); indexer.worker_event_channels[0] .send(WorkerTask::Terminate) .unwrap(); - let read = tokio::time::timeout( - std::time::Duration::from_secs(5), - indexer.try_worker_lookup_stats(), - ) - .await - .expect("a dead lane must not hang the read"); - assert!(read.is_err()); + // Both requests queue behind Terminate before the gate opens, so neither is answered. + let timeout = std::time::Duration::from_secs(5); + let (strict, lenient, ()) = tokio::join!( + tokio::time::timeout(timeout, indexer.resident_block_counts()), + tokio::time::timeout(timeout, indexer.worker_lookup_stats()), + async { + tokio::task::yield_now().await; + drop(open_gate); + }, + ); + assert!(strict.expect("strict read hung").is_err()); + assert_eq!(lenient.expect("lenient read hung").worker_count(), 0); } #[tokio::test] diff --git a/lib/kv-router/src/indexer/types.rs b/lib/kv-router/src/indexer/types.rs index 70bb99e27efc..e29ab2b02154 100644 --- a/lib/kv-router/src/indexer/types.rs +++ b/lib/kv-router/src/indexer/types.rs @@ -571,9 +571,9 @@ pub struct GetWorkersRequest { } /// A request for the number of blocks tracked for each worker rank -pub(crate) struct WorkerLookupStatsRequest { +pub(crate) struct ResidentBlockCountsRequest { /// Channel to send the per-rank block counts - pub resp: oneshot::Sender, + pub resp: oneshot::Sender, } #[derive(Debug, Default)] diff --git a/lib/kv-router/src/plugins/worker_selection/inputs.rs b/lib/kv-router/src/plugins/worker_selection/inputs.rs index c19aecc39625..eb3786e596f5 100644 --- a/lib/kv-router/src/plugins/worker_selection/inputs.rs +++ b/lib/kv-router/src/plugins/worker_selection/inputs.rs @@ -72,7 +72,7 @@ impl<'a> WorkerCandidate<'a> { /// Blocks the router's indexer tracks for this worker rank, only when this component declared /// RESIDENT_BLOCKS: device blocks from KV events, or the router's prediction in approximate /// modes. Refreshed off the request path, unweighted, and not clamped to capacity; 0 also covers - /// ranks with no data yet. None means access was not requested or the host has no counts. + /// ranks with no data yet. None means access was not requested or the host has no fresh counts. pub fn resident_blocks(self) -> Option { if !self.inputs.contains(WorkerInputs::RESIDENT_BLOCKS) { return None; @@ -155,7 +155,7 @@ impl WorkerInputs { /// Request preferred-taint routing metadata. pub const PREFERRED_TAINT: Self = Self(1 << 2); /// Request per-rank counts of blocks the router's indexer tracks (routing predictions in - /// approximate modes). The frontend rejects RESIDENT_BLOCKS without CACHE. + /// approximate modes). Requires CACHE. pub const RESIDENT_BLOCKS: Self = Self(1 << 3); /// Request host-owned active-request counts. pub const OCCUPANCY: Self = Self(1 << 5); diff --git a/lib/kv-router/src/scheduling/selector/policy.rs b/lib/kv-router/src/scheduling/selector/policy.rs index 6cc4f0089a9b..02ccff6b4408 100644 --- a/lib/kv-router/src/scheduling/selector/policy.rs +++ b/lib/kv-router/src/scheduling/selector/policy.rs @@ -384,7 +384,7 @@ impl WorkerSelector for WorkerSelectionPolicy { let resident_blocks = self .resident_blocks .as_ref() - .and_then(ResidentBlockCountsHandle::load); + .map(ResidentBlockCountsHandle::load); select_worker_with_policy( self.worker_label, state, @@ -392,7 +392,9 @@ impl WorkerSelector for WorkerSelectionPolicy { request, eligibility, block_size, - resident_blocks.as_deref(), + resident_blocks + .as_ref() + .and_then(|counts| counts.as_deref()), ) } } @@ -680,7 +682,7 @@ mod tests { #[test] fn resident_blocks_reach_only_declaring_components() { - use crate::indexer::{ResidentBlockCountsHandle, WorkerLookupStats}; + use crate::indexer::ResidentBlockCountsHandle; // Worker 8 holds no tracked blocks, which reads as 0 rather than None. Worker ids differ // from row indices, so a row-for-worker mix-up fails. @@ -763,9 +765,7 @@ mod tests { ]); let request = base_request(16); let counts = ResidentBlockCountsHandle::default(); - counts.publish( - WorkerLookupStats::from_worker_block_counts([(WorkerWithDpRank::new(3, 0), 12)]).into(), - ); + counts.publish([(WorkerWithDpRank::new(3, 0), 12)].into_iter().collect()); let probes = || -> Vec> { vec![ Box::new(ResidencyProbe(WorkerInputs::RESIDENT_BLOCKS)), diff --git a/lib/kv-router/src/services/indexer/backend.rs b/lib/kv-router/src/services/indexer/backend.rs index 857609b1d00d..c28562f512df 100644 --- a/lib/kv-router/src/services/indexer/backend.rs +++ b/lib/kv-router/src/services/indexer/backend.rs @@ -876,16 +876,16 @@ impl Indexer { /// request path. #[cfg_attr(not(feature = "standalone-selection"), allow(dead_code))] async fn resident_block_counts(&self) -> Result, KvRouterError> { - let stats = match self { - Self::Single { primary, .. } => primary.worker_lookup_stats().await?, - Self::Concurrent { primary, .. } => primary.try_worker_lookup_stats().await?, - Self::Remote { .. } | Self::None => return Ok(None), - }; - Ok(Some(stats.into())) + match self { + Self::Single { primary, .. } => primary.resident_block_counts().await.map(Some), + Self::Concurrent { primary, .. } => primary.resident_block_counts().await.map(Some), + Self::Remote { .. } | Self::None => Ok(None), + } } /// Republish [`Self::resident_block_counts`] every `interval` until `cancel` fires or the - /// primary stops answering, then withdraw them. `None` without a local primary. + /// primary stops answering, then withdraw them. A read slower than three intervals also + /// withdraws them until a read answers in time. `None` without a local primary. #[cfg_attr(not(feature = "standalone-selection"), allow(dead_code))] pub(crate) fn spawn_resident_block_counts_poller( &self, @@ -896,36 +896,71 @@ impl Indexer { return None; } let handle = ResidentBlockCountsHandle::default(); - let published = handle.clone(); let indexer = self.clone(); - tokio::spawn(async move { - let poll = async { - let mut ticks = tokio::time::interval(interval); - ticks.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); - loop { - ticks.tick().await; - match indexer.resident_block_counts().await { - Ok(Some(counts)) => published.publish(counts), - Ok(None) => return, - Err(error) => { - // Debug: the indexer logs its own failure, and an ordinary shutdown - // can stop the indexer before this task's cancel fires. - tracing::debug!(%error, "Stopped publishing resident block counts"); - return; - } + let read = move || { + let indexer = indexer.clone(); + async move { indexer.resident_block_counts().await } + }; + tokio::spawn(publish_resident_block_counts( + read, + handle.clone(), + interval, + cancel, + )); + Some(handle) + } +} + +/// The loop behind [`Indexer::spawn_resident_block_counts_poller`], with the read injected for tests. +#[cfg_attr(not(feature = "standalone-selection"), allow(dead_code))] +async fn publish_resident_block_counts( + mut read: F, + published: ResidentBlockCountsHandle, + interval: Duration, + cancel: CancellationToken, +) where + F: FnMut() -> Fut, + Fut: Future, KvRouterError>>, +{ + let deadline = interval * 3; + let poll = async { + let mut ticks = tokio::time::interval(interval); + ticks.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + loop { + ticks.tick().await; + let pending = read(); + tokio::pin!(pending); + let result = match tokio::time::timeout(deadline, &mut pending).await { + Ok(result) => result, + Err(_) => { + // Backlogged: withdraw the counts and let this read finish rather than queue + // more behind the backlog. Only a read that meets the deadline republishes. + published.clear(); + match pending.await { + Ok(Some(_)) => continue, + result => result, } } }; - // Biased, so a cancel that races a read failed by the same shutdown wins. - tokio::select! { - biased; - _ = cancel.cancelled() => {} - _ = poll => {} + match result { + Ok(Some(counts)) => published.publish(counts), + Ok(None) => return, + Err(error) => { + // Debug: the indexer logs its own failure, and an ordinary shutdown + // can stop the indexer before this task's cancel fires. + tracing::debug!(%error, "Stopped publishing resident block counts"); + return; + } } - published.clear(); - }); - Some(handle) + } + }; + // Biased, so a cancel that races a read failed by the same shutdown wins. + tokio::select! { + biased; + _ = cancel.cancelled() => {} + _ = poll => {} } + published.clear(); } #[async_trait] @@ -1714,7 +1749,13 @@ mod tests { 0, )) .await; - wait_until(|| handle.load().is_some_and(|counts| counts.get(worker) == 2)).await; + wait_until(|| { + handle + .load() + .as_ref() + .is_some_and(|counts| counts.get(worker) == 2) + }) + .await; cancel.cancel(); wait_until(|| handle.load().is_none()).await; @@ -1741,6 +1782,42 @@ mod tests { } } + #[tokio::test(start_paused = true)] + async fn resident_block_counts_withdraw_while_a_read_misses_the_deadline() { + let worker = WorkerWithDpRank::new(1, 0); + let published = ResidentBlockCountsHandle::default(); + let mut reads = 0; + let read = move || { + reads += 1; + let blocks = reads; + async move { + // The second read waits behind a backlog, past the 300 ms deadline. + if blocks == 2 { + tokio::time::sleep(Duration::from_secs(1)).await; + } + Ok(Some([(worker, blocks)].into_iter().collect())) + } + }; + tokio::spawn(publish_resident_block_counts( + read, + published.clone(), + Duration::from_millis(100), + CancellationToken::new(), + )); + let held = || published.load().as_ref().map(|counts| counts.get(worker)); + + tokio::time::sleep(Duration::from_millis(50)).await; + assert_eq!(held(), Some(1)); + // The second read starts at 100 ms, so its deadline passes at 400 ms. + tokio::time::sleep(Duration::from_millis(300)).await; + assert_eq!(held(), Some(1), "within the deadline"); + tokio::time::sleep(Duration::from_millis(100)).await; + assert_eq!(held(), None, "past the deadline"); + // The late read ends at 1.1 s and is dropped; the next read answers in time. + tokio::time::sleep(Duration::from_millis(700)).await; + assert_eq!(held(), Some(3)); + } + #[tokio::test] async fn resident_block_counts_need_a_local_primary() { let indexer = Indexer::None; diff --git a/lib/kv-router/src/services/selection/core/tests.rs b/lib/kv-router/src/services/selection/core/tests.rs index 1139372f7ded..d9b1ebf76767 100644 --- a/lib/kv-router/src/services/selection/core/tests.rs +++ b/lib/kv-router/src/services/selection/core/tests.rs @@ -734,15 +734,15 @@ fn capturing_policy_factory() -> ( (factory, observed) } -/// Picker that records the resident block count of row 0 and takes it. +/// Picker that declares `inputs`, records the resident block count of row 0, and takes it. struct ResidentBlocksRecorder { + inputs: crate::scheduling::selector::WorkerInputs, observed: Arc>>>, } impl crate::scheduling::selector::WorkerPicker for ResidentBlocksRecorder { fn required_worker_inputs(&self) -> crate::scheduling::selector::WorkerInputs { - crate::scheduling::selector::WorkerInputs::CACHE - | crate::scheduling::selector::WorkerInputs::RESIDENT_BLOCKS + self.inputs } fn pick( @@ -755,8 +755,9 @@ impl crate::scheduling::selector::WorkerPicker for ResidentBlocksRecorder { } } -#[tokio::test] -async fn declaring_policy_reads_polled_resident_blocks() { +fn resident_blocks_core( + inputs: crate::scheduling::selector::WorkerInputs, +) -> (SelectionCore, Arc>>>) { let observed = Arc::new(parking_lot::Mutex::new(Vec::new())); let factory_observed = Arc::clone(&observed); let factory: WorkerSelectionPolicyFactory = Arc::new(move |config, worker_type, _partition| { @@ -765,6 +766,7 @@ async fn declaring_policy_reads_polled_resident_blocks() { worker_type.as_str(), Vec::new(), Box::new(ResidentBlocksRecorder { + inputs, observed: Arc::clone(&factory_observed), }), ) @@ -776,6 +778,15 @@ async fn declaring_policy_reads_polled_resident_blocks() { WorkerType::Aggregated, None, ); + (core, observed) +} + +#[tokio::test] +async fn declaring_policy_reads_polled_resident_blocks() { + let (core, observed) = resident_blocks_core( + crate::scheduling::selector::WorkerInputs::CACHE + | crate::scheduling::selector::WorkerInputs::RESIDENT_BLOCKS, + ); core.upsert_worker(worker_with_kv_events(1)) .await .expect("worker upsert"); @@ -807,6 +818,25 @@ async fn declaring_policy_reads_polled_resident_blocks() { .expect("a declaring policy should read the polled count"); } +#[tokio::test] +async fn partition_rejects_resident_blocks_without_cache() { + let (core, _) = + resident_blocks_core(crate::scheduling::selector::WorkerInputs::RESIDENT_BLOCKS); + let record = core + .upsert_worker(worker_with_kv_events(1)) + .await + .expect("worker upsert"); + assert_eq!(record.lifecycle, WorkerLifecycle::Incomplete); + assert!( + record + .not_schedulable_reasons + .iter() + .any(|reason| reason.contains("WorkerInputs::CACHE")), + "{:?}", + record.not_schedulable_reasons + ); +} + type SharedCacheCalls = Arc, u32, Option)>>>; /// Shared cache that reports every block as a hit and records each query. diff --git a/lib/kv-router/src/services/selection/core/workers.rs b/lib/kv-router/src/services/selection/core/workers.rs index ad077e9b0487..d4cdf1d897c2 100644 --- a/lib/kv-router/src/services/selection/core/workers.rs +++ b/lib/kv-router/src/services/selection/core/workers.rs @@ -331,14 +331,21 @@ impl SelectionCore { self.worker_type, key.as_ref(), ); + // Checked per partition, since the factory may build a different policy for each. + let inputs = selector.declared_inputs(); + if inputs.contains(WorkerInputs::RESIDENT_BLOCKS) + && !inputs.contains(WorkerInputs::CACHE) + { + return Err(SelectionError::BadRequest(format!( + "RESIDENT_BLOCKS requires the worker-selection policy for {key} to declare WorkerInputs::CACHE" + ))); + } let profile = self .kv_router_config .policy_profile(Some(&key.model_name)) .map_err(|error| SelectionError::BadRequest(error.to_string()))?; // Poll this partition's indexer only for a policy that reads the counts. - if selector - .declared_inputs() - .contains(WorkerInputs::RESIDENT_BLOCKS) + if inputs.contains(WorkerInputs::RESIDENT_BLOCKS) && let Some(counts) = indexer.spawn_resident_block_counts_poller( RESIDENT_BLOCK_COUNTS_POLL_INTERVAL, self.cancel_token.child_token(), From 258d9b51d0080006c10ca57c3d9333d580b54cd1 Mon Sep 17 00:00:00 2001 From: Akira Lonske Date: Sat, 10 Oct 2026 21:32:38 -0400 Subject: [PATCH 03/12] test(kv-router): benchmark the opt-in resident block read Add resident_blocks_read to the worker-selection bench. It compares a scorer that adds each candidate's resident block count with the same scorer without it, at 1,024 and 10,000 workers. A bench-only constructor attaches fixed counts, since the host-side setter is crate-private. Signed-off-by: Akira Lonske --- lib/kv-router/benches/worker_selection.rs | 91 +++++++++++++++++++ .../src/scheduling/selector/policy.rs | 11 +++ 2 files changed, 102 insertions(+) diff --git a/lib/kv-router/benches/worker_selection.rs b/lib/kv-router/benches/worker_selection.rs index 6448eeaf54a7..8c709cc432fa 100644 --- a/lib/kv-router/benches/worker_selection.rs +++ b/lib/kv-router/benches/worker_selection.rs @@ -125,6 +125,42 @@ impl WorkerScorer for IgnoringPreferenceScorer { } } +/// `IgnoringPreferenceScorer` plus each candidate's resident block count. +struct ResidentBlocksScorer; + +impl WorkerScorer for ResidentBlocksScorer { + fn required_worker_inputs(&self) -> WorkerInputs { + WorkerInputs::CACHE | WorkerInputs::LOAD | WorkerInputs::RESIDENT_BLOCKS + } + + fn score( + &mut self, + context: &WorkerSelectionContext<'_>, + candidates: WorkerCandidates<'_>, + costs: &mut [f64], + ) -> Result<(), WorkerSelectionPolicyError> { + for (candidate, cost) in candidates.iter().zip(costs) { + let cache = candidate + .cache() + .ok_or_else(|| WorkerSelectionPolicyError::failed("cache input unavailable"))?; + let load = candidate + .load() + .ok_or_else(|| WorkerSelectionPolicyError::failed("load input unavailable"))?; + // Real policies should treat None (no fresh counts) as neutral; failing keeps the bench honest. + let resident_blocks = candidate.resident_blocks().ok_or_else(|| { + WorkerSelectionPolicyError::failed("resident blocks input unavailable") + })?; + let uncached_blocks = + (context.request_blocks() as f64 - cache.device_overlap_blocks()).max(0.0); + let load_blocks = load.active_prefill_tokens() as f64 / context.block_size() as f64 + + load.decode_cost_blocks() + + load.active_requests() as f64; + *cost = uncached_blocks + load_blocks + resident_blocks as f64; + } + Ok(()) + } +} + struct LowestCostPicker; impl WorkerPicker for LowestCostPicker { @@ -371,6 +407,60 @@ fn unused_preferred_taint_metadata(c: &mut Criterion) { group.finish(); } +fn resident_blocks_read(c: &mut Criterion) { + let config = KvRouterConfig { + router_temperature: 0.0, + ..Default::default() + }; + let mut group = c.benchmark_group("resident_blocks_read"); + group.warm_up_time(Duration::from_secs(2)); + group.measurement_time(Duration::from_secs(5)); + group.sample_size(50); + + for worker_count in [2, 32, 1_024, 10_000] { + let (workers, request) = fixture(worker_count); + group.throughput(Throughput::Elements(worker_count as u64)); + let without = WorkerSelectionPolicy::new( + config.clone(), + "prefill", + vec![Box::new(IgnoringPreferenceScorer)], + Box::new(LowestCostPicker), + ); + let with = WorkerSelectionPolicy::new( + config.clone(), + "prefill", + vec![Box::new(ResidentBlocksScorer)], + Box::new(LowestCostPicker), + ) + .with_fixed_resident_block_counts( + workers + .keys() + .map(|&worker_id| (WorkerWithDpRank::from_worker_id(worker_id), 4_096)), + ); + for (name, policy) in [("without", &without), ("with", &with)] { + group.bench_with_input( + BenchmarkId::new(name, worker_count), + &worker_count, + |b, _| { + b.iter(|| { + black_box( + policy + .select_worker(WorkerSelectionInput::configured( + black_box(&workers), + black_box(&request), + request.eligibility(), + black_box(16), + )) + .unwrap(), + ) + }) + }, + ); + } + } + group.finish(); +} + fn default_policy_wrapper(c: &mut Criterion) { let config = KvRouterConfig { router_temperature: 0.0, @@ -469,6 +559,7 @@ criterion_group!( worker_selection, custom_worker_selection, unused_preferred_taint_metadata, + resident_blocks_read, default_policy_wrapper, worker_projection ); diff --git a/lib/kv-router/src/scheduling/selector/policy.rs b/lib/kv-router/src/scheduling/selector/policy.rs index 02ccff6b4408..ff66b0f29575 100644 --- a/lib/kv-router/src/scheduling/selector/policy.rs +++ b/lib/kv-router/src/scheduling/selector/policy.rs @@ -141,6 +141,17 @@ impl WorkerSelectionPolicy { self } + /// Attach fixed per-rank resident block counts, for benchmarking the read path. + #[cfg(feature = "bench")] + pub fn with_fixed_resident_block_counts( + self, + counts: impl IntoIterator, + ) -> Self { + let handle = ResidentBlockCountsHandle::default(); + handle.publish(counts.into_iter().collect()); + self.with_resident_block_counts(handle) + } + pub(crate) fn declared_inputs(&self) -> WorkerInputs { match &self.state { #[cfg(any(test, feature = "bench"))] From 0df2458ab6c9b87cb869bf6c74be7aadd18fdc98 Mon Sep 17 00:00:00 2001 From: Akira Lonske Date: Sun, 11 Oct 2026 15:35:56 -0400 Subject: [PATCH 04/12] fix(kv-router): pin resident block gates and warn when counts can't be polled - Test that a picker which doesn't declare RESIDENT_BLOCKS reads None even when another component in its policy declares it. - Move the decision to poll into one function, and test that only a declaring policy starts a poller. - Warn when a declaring policy's partition has no local indexer, such as a remote indexer, since its components can only read None. - Docs: a KV event feed that never connects reads 0, and a replica that serves its indexer also counts other replicas' routing decisions. Signed-off-by: Akira Lonske --- .../router/custom-worker-selection.mdx | 2 +- .../src/scheduling/selector/policy.rs | 63 ++++++++++--------- .../src/services/selection/core/tests.rs | 12 ++++ .../src/services/selection/core/workers.rs | 37 ++++++++--- 4 files changed, 77 insertions(+), 37 deletions(-) diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx index c03b0d26f7a8..090da81697f6 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx @@ -454,7 +454,7 @@ A filter or scorer reads groups through `WorkerCandidate::cache()` or `load()`. `WorkerCandidate::preferred_taint_multiplier()` and `ScoredWorkerCandidate::preferred_taint_multiplier()` return the optional cost multiplier from preferred routing constraints. A filter, scorer, or picker must request `WorkerInputs::PREFERRED_TAINT` before reading the multiplier. Without that declaration, the component receives `None`, even when another component requested the multiplier. Exact hard-pinned requests also do not materialize it. Required routing constraints remain host-enforced eligibility rules. -`WorkerCandidate::resident_blocks()` and `WorkerInputView::resident_blocks(row)` return how many blocks the router's indexer tracks for the worker rank, including cached blocks no running request uses; use `LOAD` for work in flight. The router refreshes the counts off the request path every 100 ms, so they miss routing decisions made since the last refresh. If the indexer is backlogged and a refresh takes longer than 300 ms, components receive `None` until a refresh completes in time. Every component sees the same counts during one selection. With KV events, the count is the device-tier blocks the indexer currently tracks from engine events. With approximate routing, it is the blocks routed to the rank within the TTL, or the distinct full blocks the LRU cache policy retains, which leaves out its partial tail blocks and duplicate copies. The count is unweighted and not clamped to `total_kv_blocks`. A rank with no tracked blocks reads `0`, including a rank the indexer has no data for yet, such as a worker whose first KV events have not arrived. A component that did not request `WorkerInputs::RESIDENT_BLOCKS` receives `None`, and so does every component before the first refresh completes, while a refresh is overdue, with a remote indexer, or after the indexer stops answering. A policy that requests `RESIDENT_BLOCKS` must also request `CACHE`, because the counts come from the indexer that `CACHE` starts. The router checks this whenever it builds a policy for a routing partition: the frontend fails to start, and the standalone selection service marks the partition's workers unschedulable with the reason. With KV events every router replica sees the same counts; with approximate routing each replica counts only the requests it routed. +`WorkerCandidate::resident_blocks()` and `WorkerInputView::resident_blocks(row)` return how many blocks the router's indexer tracks for the worker rank, including cached blocks no running request uses; use `LOAD` for work in flight. The router refreshes the counts off the request path every 100 ms, so they miss routing decisions made since the last refresh. If the indexer is backlogged and a refresh takes longer than 300 ms, components receive `None` until a refresh completes in time. Every component sees the same counts during one selection. With KV events, the count is the device-tier blocks the indexer currently tracks from engine events. With approximate routing, it is the blocks routed to the rank within the TTL, or the distinct full blocks the LRU cache policy retains, which leaves out its partial tail blocks and duplicate copies. The count is unweighted and not clamped to `total_kv_blocks`. A rank with no tracked blocks reads `0`, including a rank the indexer has no data for yet, such as a worker whose first KV events have not arrived or never arrive. A component that did not request `WorkerInputs::RESIDENT_BLOCKS` receives `None`, and so does every component before the first refresh completes, while a refresh is overdue, with a remote indexer, or after the indexer stops answering. A policy that requests `RESIDENT_BLOCKS` must also request `CACHE`, because the counts come from the indexer that `CACHE` starts. The router checks this whenever it builds a policy for a routing partition: the frontend fails to start, and the standalone selection service marks the partition's workers unschedulable with the reason. With KV events every router replica sees the same counts. With approximate routing each replica counts only the requests it routed, except a replica that serves its indexer (`serve_indexer`), whose counts also include the routing decisions other replicas send it. ## Link the Policy Into Dynamo diff --git a/lib/kv-router/src/scheduling/selector/policy.rs b/lib/kv-router/src/scheduling/selector/policy.rs index ff66b0f29575..49a4a90bb72a 100644 --- a/lib/kv-router/src/scheduling/selector/policy.rs +++ b/lib/kv-router/src/scheduling/selector/policy.rs @@ -747,11 +747,12 @@ mod tests { } } - struct FewestResidentPicker; + /// Picks the fewest resident blocks; an undeclared picker reads None and takes row 0. + struct ResidencyPicker(WorkerInputs); - impl WorkerPicker for FewestResidentPicker { + impl WorkerPicker for ResidencyPicker { fn required_worker_inputs(&self) -> WorkerInputs { - WorkerInputs::RESIDENT_BLOCKS + self.0 } fn pick( @@ -761,11 +762,13 @@ mod tests { ) -> Result { let rows = input.candidates().len(); assert_eq!(input.resident_blocks(rows), None); + let declared = self.0.contains(WorkerInputs::RESIDENT_BLOCKS); for (row, candidate) in input.candidates().iter().enumerate() { - assert_eq!(input.resident_blocks(row), Some(held(candidate.worker()))); + let expected = declared.then_some(held(candidate.worker())); + assert_eq!(input.resident_blocks(row), expected); } Ok((0..rows) - .min_by_key(|&row| input.resident_blocks(row).unwrap()) + .min_by_key(|&row| input.resident_blocks(row).unwrap_or(0)) .unwrap()) } } @@ -783,30 +786,34 @@ mod tests { Box::new(ResidencyProbe(WorkerInputs::NONE)), ] }; - let policy = WorkerSelectionPolicy::new_with_filters( - KvRouterConfig::default(), - "test", - probes() - .into_iter() - .map(|probe| probe as Box) - .collect(), - probes() - .into_iter() - .map(|probe| probe as Box) - .collect(), - Box::new(FewestResidentPicker), - ) - .with_resident_block_counts(counts); + for picker_inputs in [WorkerInputs::RESIDENT_BLOCKS, WorkerInputs::NONE] { + let policy = WorkerSelectionPolicy::new_with_filters( + KvRouterConfig::default(), + "test", + probes() + .into_iter() + .map(|probe| probe as Box) + .collect(), + probes() + .into_iter() + .map(|probe| probe as Box) + .collect(), + Box::new(ResidencyPicker(picker_inputs)), + ) + .with_resident_block_counts(counts.clone()); - let selected = policy - .select_worker(WorkerSelectionInput::configured( - &workers, - &request, - request.eligibility(), - 16, - )) - .unwrap(); - assert_eq!(selected.worker, WorkerWithDpRank::new(8, 0)); + let selected = policy + .select_worker(WorkerSelectionInput::configured( + &workers, + &request, + request.eligibility(), + 16, + )) + .unwrap(); + if picker_inputs.contains(WorkerInputs::RESIDENT_BLOCKS) { + assert_eq!(selected.worker, WorkerWithDpRank::new(8, 0)); + } + } } #[test] diff --git a/lib/kv-router/src/services/selection/core/tests.rs b/lib/kv-router/src/services/selection/core/tests.rs index d9b1ebf76767..827a49a19a40 100644 --- a/lib/kv-router/src/services/selection/core/tests.rs +++ b/lib/kv-router/src/services/selection/core/tests.rs @@ -818,6 +818,18 @@ async fn declaring_policy_reads_polled_resident_blocks() { .expect("a declaring policy should read the polled count"); } +#[tokio::test] +async fn only_a_declaring_policy_polls_resident_blocks() { + use crate::scheduling::selector::WorkerInputs; + let indexer = crate::services::indexer::backend::create_indexer(16, 1); + let key = default_key(); + let poll = |inputs| { + super::workers::poll_resident_block_counts(&key, inputs, &indexer, CancellationToken::new()) + }; + assert!(poll(WorkerInputs::CACHE | WorkerInputs::LOAD).is_none()); + assert!(poll(WorkerInputs::CACHE | WorkerInputs::RESIDENT_BLOCKS).is_some()); +} + #[tokio::test] async fn partition_rejects_resident_blocks_without_cache() { let (core, _) = diff --git a/lib/kv-router/src/services/selection/core/workers.rs b/lib/kv-router/src/services/selection/core/workers.rs index d4cdf1d897c2..de457db548dd 100644 --- a/lib/kv-router/src/services/selection/core/workers.rs +++ b/lib/kv-router/src/services/selection/core/workers.rs @@ -8,7 +8,7 @@ use std::collections::HashSet; use super::reservations::{ReservationIndexObserver, spawn_reservation_index_sweep}; use super::*; -use crate::indexer::RESIDENT_BLOCK_COUNTS_POLL_INTERVAL; +use crate::indexer::{RESIDENT_BLOCK_COUNTS_POLL_INTERVAL, ResidentBlockCountsHandle}; use crate::plugins::worker_selection::WorkerInputs; use crate::sequences::topology::MAX_DATA_PARALLEL_RANKS_PER_WORKER; @@ -344,13 +344,12 @@ impl SelectionCore { .kv_router_config .policy_profile(Some(&key.model_name)) .map_err(|error| SelectionError::BadRequest(error.to_string()))?; - // Poll this partition's indexer only for a policy that reads the counts. - if inputs.contains(WorkerInputs::RESIDENT_BLOCKS) - && let Some(counts) = indexer.spawn_resident_block_counts_poller( - RESIDENT_BLOCK_COUNTS_POLL_INTERVAL, - self.cancel_token.child_token(), - ) - { + if let Some(counts) = poll_resident_block_counts( + &key, + inputs, + &indexer, + self.cancel_token.child_token(), + ) { selector = selector.with_resident_block_counts(counts); } let scheduler = LocalScheduler::new( @@ -436,3 +435,25 @@ impl SelectionCore { let _ = modified; } } + +/// Polls `indexer` for a policy that reads resident block counts. Warns when the partition has +/// no local indexer, since those components can only ever read `None`. +pub(super) fn poll_resident_block_counts( + key: &RoutingPartitionId, + inputs: WorkerInputs, + indexer: &Indexer, + cancel: CancellationToken, +) -> Option { + if !inputs.contains(WorkerInputs::RESIDENT_BLOCKS) { + return None; + } + let counts = + indexer.spawn_resident_block_counts_poller(RESIDENT_BLOCK_COUNTS_POLL_INTERVAL, cancel); + if counts.is_none() { + tracing::warn!( + partition = %key, + "Worker-selection policy reads resident block counts, but the partition has no local indexer; components will read None" + ); + } + counts +} From 863e5e55e72095f4032f66e01fe86e8fb681f87b Mon Sep 17 00:00:00 2001 From: Akira Lonske Date: Wed, 14 Oct 2026 10:37:28 -0400 Subject: [PATCH 05/12] fix(kv-router): check a partition's policy before starting its tasks A partition whose policy is rejected (RESIDENT_BLOCKS without CACHE, or a bad policy profile) used to fail after the replica sync and expiry tasks had started and the indexer had been opened. Because a failed init reruns on every later upsert, each retry leaked another set of tasks. Build and check the policy first, and test that a rejected partition opens no indexer. Docs: the resident count measures occupancy, not reusable cache. Signed-off-by: Akira Lonske --- .../router/custom-worker-selection.mdx | 2 +- .../src/services/selection/core/tests.rs | 4 ++ .../src/services/selection/core/workers.rs | 38 ++++++++++--------- 3 files changed, 25 insertions(+), 19 deletions(-) diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx index 090da81697f6..5f1303a98c7f 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx @@ -454,7 +454,7 @@ A filter or scorer reads groups through `WorkerCandidate::cache()` or `load()`. `WorkerCandidate::preferred_taint_multiplier()` and `ScoredWorkerCandidate::preferred_taint_multiplier()` return the optional cost multiplier from preferred routing constraints. A filter, scorer, or picker must request `WorkerInputs::PREFERRED_TAINT` before reading the multiplier. Without that declaration, the component receives `None`, even when another component requested the multiplier. Exact hard-pinned requests also do not materialize it. Required routing constraints remain host-enforced eligibility rules. -`WorkerCandidate::resident_blocks()` and `WorkerInputView::resident_blocks(row)` return how many blocks the router's indexer tracks for the worker rank, including cached blocks no running request uses; use `LOAD` for work in flight. The router refreshes the counts off the request path every 100 ms, so they miss routing decisions made since the last refresh. If the indexer is backlogged and a refresh takes longer than 300 ms, components receive `None` until a refresh completes in time. Every component sees the same counts during one selection. With KV events, the count is the device-tier blocks the indexer currently tracks from engine events. With approximate routing, it is the blocks routed to the rank within the TTL, or the distinct full blocks the LRU cache policy retains, which leaves out its partial tail blocks and duplicate copies. The count is unweighted and not clamped to `total_kv_blocks`. A rank with no tracked blocks reads `0`, including a rank the indexer has no data for yet, such as a worker whose first KV events have not arrived or never arrive. A component that did not request `WorkerInputs::RESIDENT_BLOCKS` receives `None`, and so does every component before the first refresh completes, while a refresh is overdue, with a remote indexer, or after the indexer stops answering. A policy that requests `RESIDENT_BLOCKS` must also request `CACHE`, because the counts come from the indexer that `CACHE` starts. The router checks this whenever it builds a policy for a routing partition: the frontend fails to start, and the standalone selection service marks the partition's workers unschedulable with the reason. With KV events every router replica sees the same counts. With approximate routing each replica counts only the requests it routed, except a replica that serves its indexer (`serve_indexer`), whose counts also include the routing decisions other replicas send it. +`WorkerCandidate::resident_blocks()` and `WorkerInputView::resident_blocks(row)` return how many blocks the router's indexer tracks for the worker rank, including cached blocks no running request uses; use `LOAD` for work in flight. The router refreshes the counts off the request path every 100 ms, so they miss routing decisions made since the last refresh. If the indexer is backlogged and a refresh takes longer than 300 ms, components receive `None` until a refresh completes in time. Every component sees the same counts during one selection. With KV events, the count is the device-tier blocks the indexer currently tracks from engine events. With approximate routing, it is the blocks routed to the rank within the TTL, or the distinct full blocks the LRU cache policy retains, which leaves out its partial tail blocks and duplicate copies. The count is unweighted and not clamped to `total_kv_blocks`. It measures occupancy, not reusable cache: it includes blocks that prefix matching can no longer reach, such as descendants of an evicted block, until their own remove events arrive, and lost remove events inflate it until the rank is removed. A rank with no tracked blocks reads `0`, including a rank the indexer has no data for yet, such as a worker whose first KV events have not arrived or never arrive. A component that did not request `WorkerInputs::RESIDENT_BLOCKS` receives `None`, and so does every component before the first refresh completes, while a refresh is overdue, with a remote indexer, or after the indexer stops answering. A policy that requests `RESIDENT_BLOCKS` must also request `CACHE`, because the counts come from the indexer that `CACHE` starts. The router checks this whenever it builds a policy for a routing partition: the frontend fails to start, and the standalone selection service marks the partition's workers unschedulable with the reason. With KV events every router replica sees the same counts. With approximate routing each replica counts only the requests it routed, except a replica that serves its indexer (`serve_indexer`), whose counts also include the routing decisions other replicas send it. ## Link the Policy Into Dynamo diff --git a/lib/kv-router/src/services/selection/core/tests.rs b/lib/kv-router/src/services/selection/core/tests.rs index 827a49a19a40..1f3a4d9c4fab 100644 --- a/lib/kv-router/src/services/selection/core/tests.rs +++ b/lib/kv-router/src/services/selection/core/tests.rs @@ -847,6 +847,10 @@ async fn partition_rejects_resident_blocks_without_cache() { "{:?}", record.not_schedulable_reasons ); + assert!( + core.indexer_registry.get_indexer(&default_key()).is_none(), + "a rejected partition opened an indexer" + ); } type SharedCacheCalls = Arc, u32, Option)>>>; diff --git a/lib/kv-router/src/services/selection/core/workers.rs b/lib/kv-router/src/services/selection/core/workers.rs index de457db548dd..fedf518f700b 100644 --- a/lib/kv-router/src/services/selection/core/workers.rs +++ b/lib/kv-router/src/services/selection/core/workers.rs @@ -269,6 +269,26 @@ impl SelectionCore { }); let entry = entry_cell .get_or_try_init(|| -> Result, SelectionError> { + // Fail before starting tasks or opening the indexer: a failed init reruns on every + // later upsert to this partition. + let mut selector = (self.worker_selection_policy_factory)( + &self.kv_router_config, + self.worker_type, + key.as_ref(), + ); + // Checked per partition, since the factory may build a different policy for each. + let inputs = selector.declared_inputs(); + if inputs.contains(WorkerInputs::RESIDENT_BLOCKS) + && !inputs.contains(WorkerInputs::CACHE) + { + return Err(SelectionError::BadRequest(format!( + "RESIDENT_BLOCKS requires the worker-selection policy for {key} to declare WorkerInputs::CACHE" + ))); + } + let profile = self + .kv_router_config + .policy_profile(Some(&key.model_name)) + .map_err(|error| SelectionError::BadRequest(error.to_string()))?; let (workers_tx, workers_rx) = watch::channel(HashMap::new()); let host_replica = self .host @@ -326,24 +346,6 @@ impl SelectionCore { block_size, )) }); - let mut selector = (self.worker_selection_policy_factory)( - &self.kv_router_config, - self.worker_type, - key.as_ref(), - ); - // Checked per partition, since the factory may build a different policy for each. - let inputs = selector.declared_inputs(); - if inputs.contains(WorkerInputs::RESIDENT_BLOCKS) - && !inputs.contains(WorkerInputs::CACHE) - { - return Err(SelectionError::BadRequest(format!( - "RESIDENT_BLOCKS requires the worker-selection policy for {key} to declare WorkerInputs::CACHE" - ))); - } - let profile = self - .kv_router_config - .policy_profile(Some(&key.model_name)) - .map_err(|error| SelectionError::BadRequest(error.to_string()))?; if let Some(counts) = poll_resident_block_counts( &key, inputs, From 250ad609701154635a695557e7d314a6cfe2420e Mon Sep 17 00:00:00 2001 From: Akira Lonske Date: Wed, 14 Oct 2026 14:14:00 -0400 Subject: [PATCH 06/12] feat(kv-router): keep resident block counts live on the event threads Replace the 100 ms poller with per-rank counters that each event thread updates as it applies events. - Each event thread keeps a cache-line-padded atomic counter for every rank it owns and, after each event it applies, stores the size of the rank's live block map, so the count cannot drift. - Readers find the counters through a copy-on-write map that changes only when a rank appears or is removed. A rank the indexer has not seen, or has removed, reads None; a known rank with no blocks reads 0. - Tracking starts when a partition's policy declares RESIDENT_BLOCKS, and each thread backfills counters for the ranks it already holds. A thread's counters are removed when it exits, including by panic. - Remove the poller, the strict stats read, the single-thread actor's stats request, and the 300 ms deadline. - Benchmarks: resident_blocks_read gains a concurrent-writer arm, resident_blocks_write measures the event-thread cost, and the Mooncake replay gains --track-resident-blocks and --resident-blocks-reader-cpu. Signed-off-by: Akira Lonske --- .../router/custom-worker-selection.mdx | 4 +- .../router/custom-policy-example/README.md | 2 +- lib/bench/kv_router/INDEXER_BENCH.md | 10 + lib/bench/kv_router/mooncake_bench.rs | 39 ++- lib/kv-router/benches/worker_selection.rs | 143 +++++++++-- .../sync_impl.rs | 44 +++- lib/kv-router/src/indexer/kv_indexer.rs | 47 ++-- lib/kv-router/src/indexer/lower_tier.rs | 2 + lib/kv-router/src/indexer/mod.rs | 3 +- lib/kv-router/src/indexer/positional.rs | 2 + lib/kv-router/src/indexer/radix_tree.rs | 50 +++- lib/kv-router/src/indexer/resident_blocks.rs | 210 +++++++++++++--- lib/kv-router/src/indexer/thread_pool.rs | 55 ++-- lib/kv-router/src/indexer/types.rs | 8 +- .../src/plugins/worker_selection/inputs.rs | 13 +- .../src/scheduling/selector/policy.rs | 42 +++- lib/kv-router/src/services/indexer/backend.rs | 237 ++++-------------- .../src/services/selection/core/tests.rs | 16 +- .../src/services/selection/core/workers.rs | 17 +- 19 files changed, 594 insertions(+), 350 deletions(-) diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx index 5f1303a98c7f..85f90712df50 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx @@ -446,7 +446,7 @@ If a component needs no optional worker data, return `WorkerInputs::NONE`. Combi | `LOAD` | `decode_cost_blocks()` | Projected active decode footprint, including this request's additional active blocks | | `LOAD` | `active_requests()` | Requests currently active on the worker | | `PREFERRED_TAINT` | `preferred_taint_multiplier()` | Optional cost multiplier from matching preferred routing taints | -| `RESIDENT_BLOCKS` | `resident_blocks()` | Blocks the router's indexer tracks for the worker rank, including cached blocks no running request uses; refreshed every 100 ms and not clamped to capacity | +| `RESIDENT_BLOCKS` | `resident_blocks()` | Blocks the router's indexer tracks for the worker rank, including cached blocks no running request uses; current as of the last event the indexer applied and not clamped to capacity | The GPU, CPU, and disk overlap accessors return raw tier counts. A missing tier or worker entry is zero; Dynamo does not substitute its weighted effective-overlap estimate. Each custom scorer chooses how to combine the raw counts. @@ -454,7 +454,7 @@ A filter or scorer reads groups through `WorkerCandidate::cache()` or `load()`. `WorkerCandidate::preferred_taint_multiplier()` and `ScoredWorkerCandidate::preferred_taint_multiplier()` return the optional cost multiplier from preferred routing constraints. A filter, scorer, or picker must request `WorkerInputs::PREFERRED_TAINT` before reading the multiplier. Without that declaration, the component receives `None`, even when another component requested the multiplier. Exact hard-pinned requests also do not materialize it. Required routing constraints remain host-enforced eligibility rules. -`WorkerCandidate::resident_blocks()` and `WorkerInputView::resident_blocks(row)` return how many blocks the router's indexer tracks for the worker rank, including cached blocks no running request uses; use `LOAD` for work in flight. The router refreshes the counts off the request path every 100 ms, so they miss routing decisions made since the last refresh. If the indexer is backlogged and a refresh takes longer than 300 ms, components receive `None` until a refresh completes in time. Every component sees the same counts during one selection. With KV events, the count is the device-tier blocks the indexer currently tracks from engine events. With approximate routing, it is the blocks routed to the rank within the TTL, or the distinct full blocks the LRU cache policy retains, which leaves out its partial tail blocks and duplicate copies. The count is unweighted and not clamped to `total_kv_blocks`. It measures occupancy, not reusable cache: it includes blocks that prefix matching can no longer reach, such as descendants of an evicted block, until their own remove events arrive, and lost remove events inflate it until the rank is removed. A rank with no tracked blocks reads `0`, including a rank the indexer has no data for yet, such as a worker whose first KV events have not arrived or never arrive. A component that did not request `WorkerInputs::RESIDENT_BLOCKS` receives `None`, and so does every component before the first refresh completes, while a refresh is overdue, with a remote indexer, or after the indexer stops answering. A policy that requests `RESIDENT_BLOCKS` must also request `CACHE`, because the counts come from the indexer that `CACHE` starts. The router checks this whenever it builds a policy for a routing partition: the frontend fails to start, and the standalone selection service marks the partition's workers unschedulable with the reason. With KV events every router replica sees the same counts. With approximate routing each replica counts only the requests it routed, except a replica that serves its indexer (`serve_indexer`), whose counts also include the routing decisions other replicas send it. +`WorkerCandidate::resident_blocks()` and `WorkerInputView::resident_blocks(row)` return how many blocks the router's indexer holds for the worker rank, including cached blocks no running request uses; use `LOAD` for work in flight. The indexer thread that owns a rank updates its count after every event it applies, so a count is current as of the last applied event, the same freshness as the `CACHE` overlap scores, and can trail events still queued for the indexer. Reads are live, so a count can change between two reads in one selection. With KV events, the count is the device-tier blocks the indexer currently tracks from engine events. With approximate routing, it is the blocks routed to the rank within the TTL, or the distinct full blocks the LRU cache policy retains, which leaves out its partial tail blocks and duplicate copies. The count is unweighted and not clamped to `total_kv_blocks`. It measures occupancy, not reusable cache: it includes blocks that prefix matching can no longer reach, such as descendants of an evicted block, until their own remove events arrive, and lost remove events inflate it until the rank is removed. A rank the indexer has seen but holds no blocks for, such as after a clear, reads `0`. A component that did not request `WorkerInputs::RESIDENT_BLOCKS` receives `None`, and so does every component for a rank the indexer has not seen or has removed, with a remote indexer, or after an indexer thread stops; a worker whose KV events never arrive therefore reads `None` rather than looking empty. A policy that requests `RESIDENT_BLOCKS` must also request `CACHE`, because the counts come from the indexer that `CACHE` starts. The router checks this whenever it builds a policy for a routing partition: the frontend fails to start, and the standalone selection service marks the partition's workers unschedulable with the reason. With KV events every router replica sees the same counts. With approximate routing each replica counts only the requests it routed, except a replica that serves its indexer (`serve_indexer`), whose counts also include the routing decisions other replicas send it. ## Link the Policy Into Dynamo diff --git a/examples/router/custom-policy-example/README.md b/examples/router/custom-policy-example/README.md index d549e578db8b..48040725d943 100644 --- a/examples/router/custom-policy-example/README.md +++ b/examples/router/custom-policy-example/README.md @@ -24,7 +24,7 @@ Declare `WorkerInputs::CACHE` to read cache data through `candidate.cache()` (fi Preferred routing taints are optional candidate metadata. A filter, scorer, or picker must request `WorkerInputs::PREFERRED_TAINT` before reading `preferred_taint_multiplier()` from a candidate; otherwise, that component receives `None`, even if another component requested it. Exact hard-pinned requests also do not materialize it. Required routing taints remain Dynamo eligibility rules. -Declare `WorkerInputs::RESIDENT_BLOCKS` to read `resident_blocks()` from a candidate, or `input.resident_blocks(row)` in a picker. It returns how many blocks the router's indexer tracks for that worker rank, including cached blocks no running request uses, refreshed in the background. Without the declaration, before the first refresh, while the indexer is backlogged, or with a remote indexer, it returns `None`. A policy that requests `RESIDENT_BLOCKS` must also request `CACHE`. +Declare `WorkerInputs::RESIDENT_BLOCKS` to read `resident_blocks()` from a candidate, or `input.resident_blocks(row)` in a picker. It returns how many blocks the router's indexer holds for that worker rank, including cached blocks no running request uses, as of the last event the indexer applied. Without the declaration, with a remote indexer, or for a rank the indexer has not seen or has removed, it returns `None`. A policy that requests `RESIDENT_BLOCKS` must also request `CACHE`. ## Pick a Starting Point diff --git a/lib/bench/kv_router/INDEXER_BENCH.md b/lib/bench/kv_router/INDEXER_BENCH.md index 37020d371f90..59c819d6e40c 100644 --- a/lib/bench/kv_router/INDEXER_BENCH.md +++ b/lib/bench/kv_router/INDEXER_BENCH.md @@ -158,6 +158,16 @@ print(f"runs={len(values)} mean={statistics.fmean(values):.3f} median={statistic PY ``` +### Resident block tracking cost + +To measure what `RESIDENT_BLOCKS` tracking costs the event threads, run the single +CRTC trial with `--track-resident-blocks`, which keeps live per-rank block counters, +and compare `achieved_block_ops_per_sec` against an untracked run. Add +`--resident-blocks-reader-cpu ` to also read every counter in a loop from a +thread pinned to that CPU, as selections do, so the result includes the cache-line +traffic between event threads and readers. Keep that CPU outside the backend and +issuer masks. + ### Sweep Sweep cells use the same completion-correct implementation. Each cell regenerates diff --git a/lib/bench/kv_router/mooncake_bench.rs b/lib/bench/kv_router/mooncake_bench.rs index bbb4022ac1df..a52adbd1de33 100644 --- a/lib/bench/kv_router/mooncake_bench.rs +++ b/lib/bench/kv_router/mooncake_bench.rs @@ -23,6 +23,7 @@ use mooncake_shared::{ }; use sha2::{Digest, Sha256}; use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; #[cfg(target_os = "linux")] const PRE_RUN_QUIESCENCE_MS: u64 = 5_000; @@ -130,6 +131,16 @@ struct Args { #[clap(long, default_value = "1")] benchmark_runs: usize, + /// Keep live per-rank resident block counts on the event threads, as when a + /// selection policy reads `RESIDENT_BLOCKS` (concurrent-radix-tree-compressed only). + #[clap(long)] + track_resident_blocks: bool, + + /// Also read every resident block counter in a loop from a thread pinned to this + /// CPU, as selections do. Implies `--track-resident-blocks`. + #[clap(long)] + resident_blocks_reader_cpu: Option, + /// Indexer backend to benchmark. Defaults to concurrent-radix-tree-compressed /// with `--num-event-workers` event threads. #[clap(subcommand)] @@ -270,6 +281,11 @@ async fn run_open_loop_for_config( let metrics = || Some(Arc::new(KvIndexerMetrics::new_unregistered())); let open_config = parse_open_loop_config(args)?; pin_current_thread_to_cpus(&open_config.backend_cpus)?; + let track_resident_blocks = + args.track_resident_blocks || args.resident_blocks_reader_cpu.is_some(); + if track_resident_blocks && config.kind != MooncakeIndexerKind::ConcurrentRadixTreeCompressed { + anyhow::bail!("resident block tracking applies only to concurrent-radix-tree-compressed"); + } match config.kind { MooncakeIndexerKind::NestedMap => { @@ -288,7 +304,28 @@ async fn run_open_loop_for_config( args.common.block_size, metrics(), )); - run_backend(config.short_name(), indexer, trial, open_config).await + let counts = track_resident_blocks.then(|| indexer.bench_track_resident_blocks()); + let stop = Arc::new(AtomicBool::new(false)); + let reader = counts + .zip(args.resident_blocks_reader_cpu) + .map(|(counts, cpu)| { + let stop = stop.clone(); + std::thread::spawn(move || { + pin_current_thread_to_cpus(&[cpu])?; + while !stop.load(Ordering::Relaxed) { + std::hint::black_box(counts.bench_sum()); + } + std::io::Result::Ok(()) + }) + }); + let result = run_backend(config.short_name(), indexer, trial, open_config).await; + stop.store(true, Ordering::Relaxed); + if let Some(reader) = reader { + reader + .join() + .map_err(|_| anyhow::anyhow!("resident block reader panicked"))??; + } + result } MooncakeIndexerKind::RadixTree | MooncakeIndexerKind::BranchShardedCrtc => { anyhow::bail!( diff --git a/lib/kv-router/benches/worker_selection.rs b/lib/kv-router/benches/worker_selection.rs index 8c709cc432fa..3f1459621e98 100644 --- a/lib/kv-router/benches/worker_selection.rs +++ b/lib/kv-router/benches/worker_selection.rs @@ -6,11 +6,18 @@ //! Run with: `cargo bench -p dynamo-kv-router --bench worker_selection` use std::collections::{HashMap, HashSet}; +use std::sync::atomic::{AtomicBool, Ordering}; use std::time::Duration; use criterion::{BenchmarkId, Criterion, Throughput, black_box, criterion_group, criterion_main}; +use dynamo_kv_router::indexer::{ + KvIndexerInterface, ThreadPoolIndexer, + concurrent_radix_tree_compressed::ConcurrentRadixTreeCompressed, +}; use dynamo_kv_router::protocols::{ - RoutingConstraints, WorkerConfigLike, WorkerId, WorkerWithDpRank, + ExternalSequenceBlockHash, KvCacheEvent, KvCacheEventData, KvCacheStoreData, + KvCacheStoredBlockData, LocalBlockHash, RouterEvent, RoutingConstraints, WorkerConfigLike, + WorkerId, WorkerWithDpRank, compute_seq_hash_for_block, }; use dynamo_kv_router::scheduling::{OverlapSignals, ScheduleMode}; use dynamo_kv_router::{ @@ -426,7 +433,7 @@ fn resident_blocks_read(c: &mut Criterion) { vec![Box::new(IgnoringPreferenceScorer)], Box::new(LowestCostPicker), ); - let with = WorkerSelectionPolicy::new( + let (with, writer) = WorkerSelectionPolicy::new( config.clone(), "prefill", vec![Box::new(ResidentBlocksScorer)], @@ -437,30 +444,123 @@ fn resident_blocks_read(c: &mut Criterion) { .keys() .map(|&worker_id| (WorkerWithDpRank::from_worker_id(worker_id), 4_096)), ); - for (name, policy) in [("without", &without), ("with", &with)] { - group.bench_with_input( - BenchmarkId::new(name, worker_count), - &worker_count, - |b, _| { - b.iter(|| { - black_box( - policy - .select_worker(WorkerSelectionInput::configured( - black_box(&workers), - black_box(&request), - request.eligibility(), - black_box(16), - )) - .unwrap(), - ) - }) - }, - ); + for (name, policy, writing) in [ + ("without", &without, false), + ("with", &with, false), + // Another thread storing to every counter in a loop: the worst case for the + // cross-core misses event-thread writes cause readers. + ("with_writer", &with, true), + ] { + let stop = AtomicBool::new(false); + std::thread::scope(|scope| { + if writing { + scope.spawn(|| { + let mut blocks = 0; + while !stop.load(Ordering::Relaxed) { + writer.store_all(blocks); + blocks += 1; + } + }); + } + group.bench_with_input( + BenchmarkId::new(name, worker_count), + &worker_count, + |b, _| { + b.iter(|| { + black_box( + policy + .select_worker(WorkerSelectionInput::configured( + black_box(&workers), + black_box(&request), + request.eligibility(), + black_box(16), + )) + .unwrap(), + ) + }) + }, + ); + stop.store(true, Ordering::Relaxed); + }); } } group.finish(); } +fn kv_event(worker: WorkerId, event_id: u64, data: KvCacheEventData) -> RouterEvent { + RouterEvent::new( + worker, + KvCacheEvent { + event_id, + dp_rank: 0, + data, + }, + ) +} + +/// Event-thread cost of keeping resident block counts: stores then clears for 128 workers on +/// 8 event threads, untracked and tracked. A concurrent reader belongs in the pinned Mooncake +/// replay (`--resident-blocks-reader-cpu`): unpinned, any busy thread skews this bench. +fn resident_blocks_write(c: &mut Criterion) { + const WORKERS: u64 = 128; + const BLOCKS: u64 = 64; + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let stores: Vec<_> = (0..WORKERS) + .map(|worker| { + let local: Vec<_> = (0..BLOCKS) + .map(|i| LocalBlockHash(worker * BLOCKS + i)) + .collect(); + let blocks = local + .iter() + .zip(compute_seq_hash_for_block(&local)) + .map(|(&tokens_hash, hash)| KvCacheStoredBlockData { + block_hash: ExternalSequenceBlockHash(hash), + tokens_hash, + mm_extra_info: None, + }) + .collect(); + let store = KvCacheStoreData { + parent_hash: None, + start_position: None, + blocks, + }; + kv_event(worker, 0, KvCacheEventData::Stored(store)) + }) + .collect(); + let mut group = c.benchmark_group("resident_blocks_write"); + group.sample_size(20); + group.warm_up_time(Duration::from_secs(1)); + group.measurement_time(Duration::from_secs(3)); + group.throughput(Throughput::Elements(WORKERS * 2)); + for (name, tracked) in [("untracked", false), ("tracked", true)] { + let indexer = ThreadPoolIndexer::new(ConcurrentRadixTreeCompressed::new(), 8, 32); + if tracked { + indexer.bench_track_resident_blocks(); + } + group.bench_function(name, |b| { + b.iter(|| { + runtime.block_on(async { + for store in &stores { + indexer.apply_event(store.clone()).await; + } + indexer.flush().await; + for worker in 0..WORKERS { + indexer + .apply_event(kv_event(worker, 1, KvCacheEventData::Cleared)) + .await; + } + indexer.flush().await; + }) + }) + }); + indexer.shutdown(); + } + group.finish(); +} + fn default_policy_wrapper(c: &mut Criterion) { let config = KvRouterConfig { router_temperature: 0.0, @@ -560,6 +660,7 @@ criterion_group!( custom_worker_selection, unused_preferred_taint_metadata, resident_blocks_read, + resident_blocks_write, default_policy_wrapper, worker_projection ); diff --git a/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/sync_impl.rs b/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/sync_impl.rs index 74c7b6d31d2f..35454b1589c7 100644 --- a/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/sync_impl.rs +++ b/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/sync_impl.rs @@ -4,7 +4,20 @@ use super::*; #[cfg(feature = "bench")] use crate::indexer::WorkerObservationState; -use crate::indexer::{AnchorCapableSyncIndexer, ApproximateLruLane, ApproximateLruTask}; +use crate::indexer::{ + AnchorCapableSyncIndexer, ApproximateLruLane, ApproximateLruTask, ResidentBlockTracker, +}; + +/// Publishes `worker`'s block count once a policy reads resident blocks. +fn track( + resident_blocks: &mut Option, + lookup: &FxHashMap, + worker: WorkerWithDpRank, +) { + if let Some(tracker) = resident_blocks { + tracker.update(worker, lookup.get(&worker).map(WorkerLookup::len)); + } +} // ============================================================================ // SyncIndexer implementation for ConcurrentRadixTreeCompressed @@ -20,6 +33,7 @@ impl SyncIndexer for ConcurrentRadixTreeCompressed { let mut lookup = FxHashMap::default(); let counters = metrics.as_ref().map(|m| m.prebind()); let mut approximate_lru = ApproximateLruLane::default(); + let mut resident_blocks = None; #[cfg(feature = "bench")] let mut observation = WorkerObservationState::default(); @@ -27,7 +41,9 @@ impl SyncIndexer for ConcurrentRadixTreeCompressed { match task { WorkerTask::Event(event) => { let kind = EventKind::of(&event.event.data); + let worker = WorkerWithDpRank::new(event.worker_id, event.event.dp_rank); let result = self.apply_event(&mut lookup, event, counters.as_ref()); + track(&mut resident_blocks, &lookup, worker); if result.is_err() { tracing::warn!("Failed to apply event: {:?}", result.as_ref().err()); } @@ -37,7 +53,9 @@ impl SyncIndexer for ConcurrentRadixTreeCompressed { } WorkerTask::EventWithAck { event, resp } => { let kind = EventKind::of(&event.event.data); + let worker = WorkerWithDpRank::new(event.worker_id, event.event.dp_rank); let result = self.apply_event(&mut lookup, event, counters.as_ref()); + track(&mut resident_blocks, &lookup, worker); let applied = result.is_ok(); if result.is_err() { tracing::warn!("Failed to apply event: {:?}", result.as_ref().err()); @@ -63,7 +81,10 @@ impl SyncIndexer for ConcurrentRadixTreeCompressed { } = output; for event in events { let kind = EventKind::of(&event.event.data); + let worker = + WorkerWithDpRank::new(event.worker_id, event.event.dp_rank); let applied = self.apply_event(&mut lookup, event, counters.as_ref()); + track(&mut resident_blocks, &lookup, worker); if let Some(ref counters) = counters { counters.inc(kind, applied); } @@ -96,7 +117,9 @@ impl SyncIndexer for ConcurrentRadixTreeCompressed { correlation_id, } => { let kind = EventKind::of(&event.event.data); + let worker = WorkerWithDpRank::new(event.worker_id, event.event.dp_rank); let result = self.apply_event(&mut lookup, event, counters.as_ref()); + track(&mut resident_blocks, &lookup, worker); observation.record(correlation_id, result.is_ok()); if result.is_err() { tracing::warn!("Failed to apply event: {:?}", result.as_ref().err()); @@ -125,6 +148,9 @@ impl SyncIndexer for ConcurrentRadixTreeCompressed { WorkerRemovalTarget::WorkerId(worker_id), sweep_tree, ); + if let Some(tracker) = &mut resident_blocks { + tracker.remove(|worker| worker.worker_id == worker_id); + } let _ = resp.send(()); } WorkerTask::RemoveWorkerDpRank { @@ -132,12 +158,16 @@ impl SyncIndexer for ConcurrentRadixTreeCompressed { dp_rank, sweep_tree, } => { - approximate_lru.forget_rank(WorkerWithDpRank::new(worker_id, dp_rank)); + let rank = WorkerWithDpRank::new(worker_id, dp_rank); + approximate_lru.forget_rank(rank); self.erase_worker_coverage( &mut lookup, - WorkerRemovalTarget::DpRank(WorkerWithDpRank::new(worker_id, dp_rank)), + WorkerRemovalTarget::DpRank(rank), sweep_tree, ); + if let Some(tracker) = &mut resident_blocks { + tracker.remove(|worker| worker == rank); + } } WorkerTask::CleanupStaleChildren => { self.run_cleanup_task(); @@ -166,6 +196,14 @@ impl SyncIndexer for ConcurrentRadixTreeCompressed { WorkerTask::Flush(sender) => { let _ = sender.send(()); } + WorkerTask::TrackResidentBlocks(counts) => { + resident_blocks = Some(ResidentBlockTracker::new( + counts, + lookup + .iter() + .map(|(worker, worker_lookup)| (*worker, worker_lookup.len())), + )); + } WorkerTask::Terminate => { break; } diff --git a/lib/kv-router/src/indexer/kv_indexer.rs b/lib/kv-router/src/indexer/kv_indexer.rs index 88718c1d38e3..5d76d3aaf688 100644 --- a/lib/kv-router/src/indexer/kv_indexer.rs +++ b/lib/kv-router/src/indexer/kv_indexer.rs @@ -15,8 +15,8 @@ use super::{ ApproximateLruLease, ApproximateLruStats, ApproximateLruTask, ApproximateRetentionConfig, ContainsWorkerBlockRequest, DumpRequest, EventKind, FlushRequest, GetWorkersRequest, KvIndexerInterface, KvIndexerMetrics, KvRouterError, MatchDetails, MatchDetailsRequest, - MatchRequest, PreBoundEventCounters, RadixTree, ResidentBlockCounts, - ResidentBlockCountsRequest, RoutingDecisionRequest, panic_payload_message, + MatchRequest, PreBoundEventCounters, RadixTree, ResidentBlockCountsHandle, + RoutingDecisionRequest, panic_payload_message, }; use crate::indexer::pruning::{BlockEntry, PruneConfig, WorkerPruneManager}; use crate::protocols::*; @@ -337,9 +337,12 @@ pub struct KvIndexer { remove_worker_dp_rank_tx: mpsc::Sender<(WorkerId, DpRank)>, /// A sender for get workers requests. get_workers_tx: mpsc::Sender, - /// A sender for per-rank block count requests. + /// Starts live per-rank block counts in the actor. #[cfg_attr(not(feature = "standalone-indexer"), allow(dead_code))] - resident_block_counts_tx: mpsc::Sender, + track_resident_blocks_tx: mpsc::Sender, + /// The live counts, once a policy reads them. Shared so every clone tracks into one handle. + #[cfg_attr(not(feature = "standalone-indexer"), allow(dead_code))] + resident_blocks: Arc>, /// A sender for dump requests. dump_tx: mpsc::Sender, /// A sender for flush requests. @@ -434,8 +437,8 @@ impl KvIndexer { let (remove_worker_dp_rank_tx, remove_worker_dp_rank_rx) = mpsc::channel::<(WorkerId, DpRank)>(16); let (get_workers_tx, get_workers_rx) = mpsc::channel::(16); - let (resident_block_counts_tx, resident_block_counts_rx) = - mpsc::channel::(16); + let (track_resident_blocks_tx, track_resident_blocks_rx) = + mpsc::channel::(1); let (dump_tx, dump_rx) = mpsc::channel::(16); let (flush_tx, flush_rx) = mpsc::channel::(16); let (routing_tx, mut routing_rx) = mpsc::channel::(2048); @@ -466,7 +469,7 @@ impl KvIndexer { let mut remove_worker_rx = remove_worker_rx; let mut remove_worker_dp_rank_rx = remove_worker_dp_rank_rx; let mut get_workers_rx = get_workers_rx; - let mut resident_block_counts_rx = resident_block_counts_rx; + let mut track_resident_blocks_rx = track_resident_blocks_rx; let mut dump_rx = dump_rx; let mut flush_rx = flush_rx; let mut trie = delegate.map_or_else(RadixTree::new, RadixTree::new_with_delegate); @@ -505,9 +508,8 @@ impl KvIndexer { } } - // Ahead of mutations and matches: under steady traffic, later arms may never run. - Some(req) = resident_block_counts_rx.recv() => { - let _ = req.resp.send(trie.resident_block_counts()); + Some(counts) = track_resident_blocks_rx.recv() => { + trie.track_resident_blocks(counts); } Some(mutation) = mutation_rx.recv() => { @@ -668,7 +670,8 @@ impl KvIndexer { remove_worker_tx, remove_worker_dp_rank_tx, get_workers_tx, - resident_block_counts_tx, + track_resident_blocks_tx, + resident_blocks: Arc::default(), dump_tx, flush_tx, routing_tx, @@ -788,19 +791,17 @@ impl KvIndexer { self.get_workers_tx.clone() } - /// Distinct blocks the tree holds for each worker rank. + /// Live per-rank block counts, which the actor starts keeping on the first call. #[cfg_attr(not(feature = "standalone-indexer"), allow(dead_code))] - pub(crate) async fn resident_block_counts(&self) -> Result { - let (resp_tx, resp_rx) = oneshot::channel(); - self.resident_block_counts_tx - .send(ResidentBlockCountsRequest { resp: resp_tx }) - .await - .map_err(|_| KvRouterError::IndexerOffline)?; - // A request sent while the actor drops its receiver stays queued with resp_tx alive. - tokio::select! { - resp = resp_rx => resp.map_err(|_| KvRouterError::IndexerDroppedRequest), - _ = self.resident_block_counts_tx.closed() => Err(KvRouterError::IndexerOffline), - } + pub(crate) fn track_resident_blocks(&self) -> ResidentBlockCountsHandle { + self.resident_blocks + .get_or_init(|| { + let counts = ResidentBlockCountsHandle::default(); + // A stopped actor never sees this, and every rank reads None. + let _ = self.track_resident_blocks_tx.try_send(counts.clone()); + counts + }) + .clone() } } diff --git a/lib/kv-router/src/indexer/lower_tier.rs b/lib/kv-router/src/indexer/lower_tier.rs index b485fcdad66a..133a80113a1a 100644 --- a/lib/kv-router/src/indexer/lower_tier.rs +++ b/lib/kv-router/src/indexer/lower_tier.rs @@ -1121,6 +1121,8 @@ impl SyncIndexer for LowerTierIndexer { WorkerTask::Flush(sender) => { let _ = sender.send(()); } + // Only the device-tier primary tracks resident blocks. + WorkerTask::TrackResidentBlocks(_) => {} WorkerTask::CleanupStaleChildren => {} WorkerTask::Terminate => { break; diff --git a/lib/kv-router/src/indexer/mod.rs b/lib/kv-router/src/indexer/mod.rs index b71734e66a85..c34b10b4afc7 100644 --- a/lib/kv-router/src/indexer/mod.rs +++ b/lib/kv-router/src/indexer/mod.rs @@ -96,7 +96,8 @@ pub use lower_tier_indexers::*; pub use metrics::*; #[cfg(feature = "bench")] pub use observation::*; -pub(crate) use resident_blocks::*; +pub use resident_blocks::ResidentBlockCountsHandle; +pub(crate) use resident_blocks::{ResidentBlockCounts, ResidentBlockTracker}; pub use thread_pool::*; pub use traits::*; pub use types::*; diff --git a/lib/kv-router/src/indexer/positional.rs b/lib/kv-router/src/indexer/positional.rs index 274230c7d1c9..cf44591db118 100644 --- a/lib/kv-router/src/indexer/positional.rs +++ b/lib/kv-router/src/indexer/positional.rs @@ -325,6 +325,8 @@ impl SyncIndexer for PositionalIndexer { WorkerTask::Flush(sender) => { let _ = sender.send(()); } + // Only the device-tier primary tracks resident blocks. + WorkerTask::TrackResidentBlocks(_) => {} WorkerTask::Terminate => { break; } diff --git a/lib/kv-router/src/indexer/radix_tree.rs b/lib/kv-router/src/indexer/radix_tree.rs index 8f8cf9f72bb8..4011ff425658 100644 --- a/lib/kv-router/src/indexer/radix_tree.rs +++ b/lib/kv-router/src/indexer/radix_tree.rs @@ -13,7 +13,10 @@ use std::{ use rustc_hash::{FxBuildHasher, FxHashMap, FxHashSet}; use super::compressed_radix::{NodeState, append_dump_events}; -use super::{EventWarningKind, MatchDetails, PreBoundEventCounters, ResidentBlockCounts}; +use super::{ + EventWarningKind, MatchDetails, PreBoundEventCounters, ResidentBlockCountsHandle, + ResidentBlockTracker, +}; use crate::protocols::*; pub(crate) type SharedRadixBlock = Rc>; @@ -85,6 +88,8 @@ pub struct RadixTree { root: SharedRadixBlock, lookup: FxHashMap, lifecycle: super::HashLifecycle, + /// Live per-rank block counts, once a policy reads them. + resident_blocks: Option, } impl Default for RadixTree { @@ -126,6 +131,7 @@ impl RadixTree { root: Rc::new(RefCell::new(RadixBlock::root())), lookup: FxHashMap::default(), lifecycle: super::HashLifecycle::default(), + resident_blocks: None, } } @@ -292,14 +298,18 @@ impl RadixTree { let event_id = event.event.event_id; self.lookup.entry(worker).or_default(); - match event.event.data { + let result = match event.event.data { KvCacheEventData::Stored(store) => self.apply_stored(worker, store, event_id, counters), KvCacheEventData::Removed(remove) => self.apply_removed(worker, remove, event_id), KvCacheEventData::Cleared => { - self.remove_worker_dp_rank(worker.worker_id, worker.dp_rank); + self.drop_rank(worker); Ok(()) } + }; + if let Some(tracker) = &mut self.resident_blocks { + tracker.update(worker, self.lookup.get(&worker).map(WorkerLookup::len)); } + result } fn apply_stored( @@ -736,10 +746,21 @@ impl RadixTree { pub fn remove_worker(&mut self, worker_id: WorkerId) { self.remove_or_clear_worker_blocks(worker_id, false); + if let Some(tracker) = &mut self.resident_blocks { + tracker.remove(|worker| worker.worker_id == worker_id); + } } pub fn remove_worker_dp_rank(&mut self, worker_id: WorkerId, dp_rank: DpRank) { - let worker = WorkerWithDpRank { worker_id, dp_rank }; + let rank = WorkerWithDpRank { worker_id, dp_rank }; + self.drop_rank(rank); + if let Some(tracker) = &mut self.resident_blocks { + tracker.remove(|worker| worker == rank); + } + } + + /// Drops `worker`'s blocks without forgetting the rank, as a clear does. + fn drop_rank(&mut self, worker: WorkerWithDpRank) { let Some(blocks) = self.lookup.remove(&worker) else { return; }; @@ -760,6 +781,13 @@ impl RadixTree { pub fn clear_all_blocks(&mut self, worker_id: WorkerId) { self.remove_or_clear_worker_blocks(worker_id, true); + if let Some(tracker) = &mut self.resident_blocks { + for (worker, blocks) in &self.lookup { + if worker.worker_id == worker_id { + tracker.update(*worker, Some(blocks.len())); + } + } + } } pub fn get_workers(&self) -> Vec { @@ -860,12 +888,14 @@ impl RadixTree { self.lookup.values().map(FxHashMap::len).sum() } - /// Distinct blocks held by each worker rank. - pub(crate) fn resident_block_counts(&self) -> ResidentBlockCounts { - self.lookup - .iter() - .map(|(worker, worker_lookup)| (*worker, worker_lookup.len())) - .collect() + /// Keeps live per-rank block counts in `counts` from now on, starting with the ranks held now. + pub(crate) fn track_resident_blocks(&mut self, counts: ResidentBlockCountsHandle) { + self.resident_blocks = Some(ResidentBlockTracker::new( + counts, + self.lookup + .iter() + .map(|(worker, worker_lookup)| (*worker, worker_lookup.len())), + )); } #[cfg(test)] diff --git a/lib/kv-router/src/indexer/resident_blocks.rs b/lib/kv-router/src/indexer/resident_blocks.rs index fd28d77f7151..f60de91d9bb0 100644 --- a/lib/kv-router/src/indexer/resident_blocks.rs +++ b/lib/kv-router/src/indexer/resident_blocks.rs @@ -1,68 +1,204 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! Per-rank resident block counts, published off the request path. +//! Live per-rank resident block counts. //! -//! Event threads own the per-rank block maps, so reading them is a queued round -//! trip. A poller makes that trip and publishes the result here, where worker -//! selection can read it without waiting on the indexer. +//! Each event thread owns the block maps of its ranks. Once a policy reads resident blocks, the +//! thread keeps one counter per rank it owns and sets it to the rank's block count after each +//! event it applies. Readers find the counters through a copy-on-write map that changes only +//! when a rank appears or is removed. use std::sync::Arc; -use std::time::Duration; +use std::sync::atomic::{AtomicU64, Ordering}; -use arc_swap::{ArcSwapOption, Guard}; +use arc_swap::{ArcSwap, Guard}; +use crossbeam_utils::CachePadded; use rustc_hash::FxHashMap; use crate::protocols::WorkerWithDpRank; -/// How often routing hosts republish [`ResidentBlockCounts`]. -#[cfg_attr(not(feature = "standalone-selection"), allow(dead_code))] -pub(crate) const RESIDENT_BLOCK_COUNTS_POLL_INTERVAL: Duration = Duration::from_millis(100); +/// Padded so counters written by different event threads never share a cache line. +type Counter = CachePadded; -/// Blocks the indexer tracked for each worker rank when polled. -#[derive(Debug)] +fn counter(blocks: usize) -> Arc { + Arc::new(CachePadded::new(AtomicU64::new(blocks as u64))) +} + +/// The counters of every rank an indexer holds. +#[derive(Clone, Debug, Default)] pub(crate) struct ResidentBlockCounts { - counts: FxHashMap, + counters: FxHashMap>, } impl ResidentBlockCounts { - /// Blocks tracked for `worker`; zero when the indexer tracks none for it. - pub(crate) fn get(&self, worker: WorkerWithDpRank) -> u64 { - self.counts.get(&worker).copied().unwrap_or(0) + /// Blocks held for `worker` as of the last event applied to it; `None` for a rank the indexer + /// has not seen or has removed. + pub(crate) fn get(&self, worker: WorkerWithDpRank) -> Option { + self.counters + .get(&worker) + .map(|counter| counter.load(Ordering::Relaxed)) } } -impl FromIterator<(WorkerWithDpRank, usize)> for ResidentBlockCounts { - fn from_iter>(iter: I) -> Self { - let mut counts = FxHashMap::default(); - for (worker, blocks) in iter { - *counts.entry(worker).or_default() += blocks as u64; +/// The counters a policy reads, shared with the indexer's event threads. +#[derive(Clone, Debug, Default)] +pub struct ResidentBlockCountsHandle { + counts: Arc>, +} + +impl ResidentBlockCountsHandle { + /// The counters as of the last rank change. Hold the guard only for one selection. + pub(crate) fn load(&self) -> Guard> { + self.counts.load() + } + + /// Counters fixed at `counts`, for tests and benchmarks. + #[cfg(any(test, feature = "bench"))] + pub(crate) fn fixed(counts: impl IntoIterator) -> Self { + let counters = counts + .into_iter() + .map(|(worker, blocks)| (worker, counter(blocks as usize))) + .collect(); + Self { + counts: Arc::new(ArcSwap::from_pointee(ResidentBlockCounts { counters })), } - Self { counts } + } + + /// Sum of every counter, read the way selections read them. + #[cfg(feature = "bench")] + #[doc(hidden)] + pub fn bench_sum(&self) -> u64 { + self.counts + .load() + .counters + .values() + .map(|counter| counter.load(Ordering::Relaxed)) + .sum() + } + + /// Sets every counter to `blocks`, the way event threads write them. + #[cfg(feature = "bench")] + pub(crate) fn store_all(&self, blocks: u64) { + for counter in self.counts.load().counters.values() { + counter.store(blocks, Ordering::Relaxed); + } + } + + fn publish(&self, added: &[(WorkerWithDpRank, Arc)]) { + self.counts.rcu(|current| { + let mut next = ResidentBlockCounts::clone(current); + next.counters.extend( + added + .iter() + .map(|(worker, counter)| (*worker, counter.clone())), + ); + next + }); + } + + fn unpublish(&self, removed: &[(WorkerWithDpRank, Arc)]) { + self.counts.rcu(|current| { + let mut next = ResidentBlockCounts::clone(current); + for (worker, counter) in removed { + if next + .counters + .get(worker) + .is_some_and(|published| Arc::ptr_eq(published, counter)) + { + next.counters.remove(worker); + } + } + next + }); } } -/// The latest published [`ResidentBlockCounts`], shared by the poller and its readers. -#[derive(Clone, Default)] -pub(crate) struct ResidentBlockCountsHandle { - latest: Arc>, +/// The counters for the ranks one event thread owns. Dropping it, as when the thread exits, +/// removes them, so readers get `None` instead of counts nothing updates anymore. +pub(crate) struct ResidentBlockTracker { + counts: ResidentBlockCountsHandle, + owned: FxHashMap>, } -impl ResidentBlockCountsHandle { - /// The latest counts, or `None` before the first poll, while the indexer is backlogged, or - /// after the poller stops. Hold the guard only for one selection. - pub(crate) fn load(&self) -> Guard>> { - self.latest.load() +impl ResidentBlockTracker { + /// Starts tracking, with counters for the ranks the thread already holds. + pub(crate) fn new( + counts: ResidentBlockCountsHandle, + ranks: impl IntoIterator, + ) -> Self { + let owned: FxHashMap<_, _> = ranks + .into_iter() + .map(|(worker, blocks)| (worker, counter(blocks))) + .collect(); + let added: Vec<_> = owned + .iter() + .map(|(worker, counter)| (*worker, counter.clone())) + .collect(); + counts.publish(&added); + Self { counts, owned } } - #[cfg_attr(not(feature = "standalone-indexer"), allow(dead_code))] - pub(crate) fn publish(&self, counts: ResidentBlockCounts) { - self.latest.store(Some(Arc::new(counts))); + /// Records `worker`'s block count after an event. `None` means the thread holds no blocks + /// for it, as after a clear; a rank already counted then reads 0. + pub(crate) fn update(&mut self, worker: WorkerWithDpRank, blocks: Option) { + if let Some(counter) = self.owned.get(&worker) { + counter.store(blocks.unwrap_or(0) as u64, Ordering::Relaxed); + } else if let Some(blocks) = blocks { + let counter = counter(blocks); + self.owned.insert(worker, counter.clone()); + self.counts.publish(&[(worker, counter)]); + } + } + + /// Stops counting the ranks `removed` selects, once the indexer has removed them. + pub(crate) fn remove(&mut self, removed: impl Fn(WorkerWithDpRank) -> bool) { + let mut gone = Vec::new(); + self.owned.retain(|worker, counter| { + let keep = !removed(*worker); + if !keep { + gone.push((*worker, counter.clone())); + } + keep + }); + if !gone.is_empty() { + self.counts.unpublish(&gone); + } } +} + +impl Drop for ResidentBlockTracker { + fn drop(&mut self) { + let owned: Vec<_> = self.owned.drain().collect(); + self.counts.unpublish(&owned); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn tracker_publishes_counts_and_withdraws_them() { + let counts = ResidentBlockCountsHandle::default(); + let (held, emptied, removed, unseen) = ( + WorkerWithDpRank::new(1, 0), + WorkerWithDpRank::new(1, 1), + WorkerWithDpRank::new(2, 0), + WorkerWithDpRank::new(3, 0), + ); + let mut tracker = ResidentBlockTracker::new(counts.clone(), [(held, 4), (emptied, 2)]); + tracker.update(held, Some(5)); + tracker.update(emptied, None); + tracker.update(removed, Some(7)); + let read = |worker| counts.load().get(worker); + assert_eq!( + [read(held), read(emptied), read(removed), read(unseen)], + [Some(5), Some(0), Some(7), None] + ); - /// Withdraw the counts so readers see `None` rather than stale counts. - #[cfg_attr(not(feature = "standalone-indexer"), allow(dead_code))] - pub(crate) fn clear(&self) { - self.latest.store(None); + tracker.remove(|worker| worker.worker_id == 2); + assert_eq!(read(removed), None); + drop(tracker); + assert_eq!(read(held), None, "an exited thread's ranks read None"); } } diff --git a/lib/kv-router/src/indexer/thread_pool.rs b/lib/kv-router/src/indexer/thread_pool.rs index e3bb41c60b67..e6fdc25e4e41 100644 --- a/lib/kv-router/src/indexer/thread_pool.rs +++ b/lib/kv-router/src/indexer/thread_pool.rs @@ -19,8 +19,8 @@ use super::concurrent_radix_tree_compressed::ConcurrentRadixTreeCompressed; use super::{ ApproximateLruClient, ApproximateLruCommandSink, ApproximateLruIncarnation, ApproximateLruLease, ApproximateLruStats, ApproximateLruTask, ApproximateRetentionConfig, - KvIndexerInterface, KvIndexerMetrics, KvRouterError, ResidentBlockCounts, ShardSizeSnapshot, - SyncIndexer, WorkerLookupStats, WorkerTask, panic_payload_message, + KvIndexerInterface, KvIndexerMetrics, KvRouterError, ResidentBlockCountsHandle, + ShardSizeSnapshot, SyncIndexer, WorkerLookupStats, WorkerTask, panic_payload_message, }; #[cfg(feature = "bench")] use super::{ @@ -86,6 +86,9 @@ pub struct ThreadPoolIndexer { /// Cancelled when the matching worker thread exits, including by panic. thread_exited: Vec, + /// Live per-rank block counts, once a policy reads them. + resident_blocks: std::sync::OnceLock, + /// Approximate-mode TTL pruning manager. None for normal event-driven mode. prune_manager: Option, @@ -335,6 +338,7 @@ impl ThreadPoolIndexer { kv_block_size, thread_handles: Mutex::new(thread_handles), thread_exited, + resident_blocks: std::sync::OnceLock::new(), prune_manager, prune_pump_cancel, synthetic_event_id, @@ -613,19 +617,26 @@ impl ThreadPoolIndexer { WorkerLookupStats::from_worker_block_counts(worker_blocks) } - /// Per-rank block counts. Fails if a worker thread has exited, so a caller never mistakes a - /// dead thread's ranks for empty ones. + /// [`Self::track_resident_blocks`], for benchmarking the event threads' cost of tracking. + #[cfg(feature = "bench")] + #[doc(hidden)] + pub fn bench_track_resident_blocks(&self) -> ResidentBlockCountsHandle { + self.track_resident_blocks() + } + + /// Live per-rank block counts, which every worker thread starts keeping on the first call. #[cfg_attr(not(feature = "standalone-indexer"), allow(dead_code))] - pub(crate) async fn resident_block_counts(&self) -> Result { - let lanes = self.lane_lookup_stats().await; - if lanes.iter().any(Option::is_none) { - return Err(KvRouterError::IndexerOffline); - } - Ok(lanes - .into_iter() - .flatten() - .flat_map(|stats| stats.worker_blocks) - .collect()) + pub(crate) fn track_resident_blocks(&self) -> ResidentBlockCountsHandle { + self.resident_blocks + .get_or_init(|| { + let counts = ResidentBlockCountsHandle::default(); + for channel in &self.worker_event_channels { + // An exited thread never sees this, and its ranks read None. + let _ = channel.send(WorkerTask::TrackResidentBlocks(counts.clone())); + } + counts + }) + .clone() } pub async fn get_workers(&self) -> Vec { @@ -1630,24 +1641,24 @@ mod tests { } #[tokio::test] - async fn stats_reads_fail_when_a_thread_exits_with_them_queued() { + async fn stats_read_skips_a_thread_that_exits_with_it_queued() { let (open_gate, gate) = flume::bounded::<()>(1); let indexer = ThreadPoolIndexer::new(GatedBackend(gate), 1, 16); indexer.worker_event_channels[0] .send(WorkerTask::Terminate) .unwrap(); - // Both requests queue behind Terminate before the gate opens, so neither is answered. - let timeout = std::time::Duration::from_secs(5); - let (strict, lenient, ()) = tokio::join!( - tokio::time::timeout(timeout, indexer.resident_block_counts()), - tokio::time::timeout(timeout, indexer.worker_lookup_stats()), + // The request queues behind Terminate before the gate opens, so it is never answered. + let (read, ()) = tokio::join!( + tokio::time::timeout( + std::time::Duration::from_secs(5), + indexer.worker_lookup_stats() + ), async { tokio::task::yield_now().await; drop(open_gate); }, ); - assert!(strict.expect("strict read hung").is_err()); - assert_eq!(lenient.expect("lenient read hung").worker_count(), 0); + assert_eq!(read.expect("stats read hung").worker_count(), 0); } #[tokio::test] diff --git a/lib/kv-router/src/indexer/types.rs b/lib/kv-router/src/indexer/types.rs index e29ab2b02154..497ddf4dc18c 100644 --- a/lib/kv-router/src/indexer/types.rs +++ b/lib/kv-router/src/indexer/types.rs @@ -570,12 +570,6 @@ pub struct GetWorkersRequest { pub resp: oneshot::Sender>, } -/// A request for the number of blocks tracked for each worker rank -pub(crate) struct ResidentBlockCountsRequest { - /// Channel to send the per-rank block counts - pub resp: oneshot::Sender, -} - #[derive(Debug, Default)] pub struct WorkerLookupStats { pub worker_blocks: Vec<(WorkerWithDpRank, usize)>, @@ -661,6 +655,8 @@ pub enum WorkerTask { resp: oneshot::Sender, }, Flush(oneshot::Sender<()>), + /// Keep live per-rank block counts in these counters from now on. + TrackResidentBlocks(super::ResidentBlockCountsHandle), Terminate, } diff --git a/lib/kv-router/src/plugins/worker_selection/inputs.rs b/lib/kv-router/src/plugins/worker_selection/inputs.rs index eb3786e596f5..d81d6077c22d 100644 --- a/lib/kv-router/src/plugins/worker_selection/inputs.rs +++ b/lib/kv-router/src/plugins/worker_selection/inputs.rs @@ -69,17 +69,18 @@ impl<'a> WorkerCandidate<'a> { } } - /// Blocks the router's indexer tracks for this worker rank, only when this component declared - /// RESIDENT_BLOCKS: device blocks from KV events, or the router's prediction in approximate - /// modes. Refreshed off the request path, unweighted, and not clamped to capacity; 0 also covers - /// ranks with no data yet. None means access was not requested or the host has no fresh counts. + /// Blocks the router's indexer holds for this worker rank as of the last event it applied, only + /// when this component declared RESIDENT_BLOCKS: device blocks from KV events, or the router's + /// prediction in approximate modes. Unweighted and not clamped to capacity. None means access + /// was not requested, the host has no local indexer, or the indexer has not seen the rank or + /// has removed it. pub fn resident_blocks(self) -> Option { if !self.inputs.contains(WorkerInputs::RESIDENT_BLOCKS) { return None; } self.cache_snapshot .resident_blocks - .map(|counts| counts.get(self.data.worker)) + .and_then(|counts| counts.get(self.data.worker)) } } @@ -396,6 +397,6 @@ impl<'a> WorkerInputView<'a> { let counts = self.resident_blocks?; self.candidates .get(row) - .map(|candidate| counts.get(candidate.worker)) + .and_then(|candidate| counts.get(candidate.worker)) } } diff --git a/lib/kv-router/src/scheduling/selector/policy.rs b/lib/kv-router/src/scheduling/selector/policy.rs index 49a4a90bb72a..7943daf1f22a 100644 --- a/lib/kv-router/src/scheduling/selector/policy.rs +++ b/lib/kv-router/src/scheduling/selector/policy.rs @@ -141,15 +141,18 @@ impl WorkerSelectionPolicy { self } - /// Attach fixed per-rank resident block counts, for benchmarking the read path. + /// Attach fixed per-rank resident block counts for benchmarking the read path, with a writer + /// that stores to them the way event threads do. #[cfg(feature = "bench")] pub fn with_fixed_resident_block_counts( self, - counts: impl IntoIterator, - ) -> Self { - let handle = ResidentBlockCountsHandle::default(); - handle.publish(counts.into_iter().collect()); - self.with_resident_block_counts(handle) + counts: impl IntoIterator, + ) -> (Self, ResidentBlockCountsWriter) { + let counts = ResidentBlockCountsHandle::fixed(counts); + ( + self.with_resident_block_counts(counts.clone()), + ResidentBlockCountsWriter(counts), + ) } pub(crate) fn declared_inputs(&self) -> WorkerInputs { @@ -403,13 +406,24 @@ impl WorkerSelector for WorkerSelectionPolicy { request, eligibility, block_size, - resident_blocks - .as_ref() - .and_then(|counts| counts.as_deref()), + resident_blocks.as_deref().map(|counts| &**counts), ) } } +/// Stores to the counters from [`WorkerSelectionPolicy::with_fixed_resident_block_counts`], for +/// benchmarking reads while event threads write. +#[cfg(feature = "bench")] +pub struct ResidentBlockCountsWriter(ResidentBlockCountsHandle); + +#[cfg(feature = "bench")] +impl ResidentBlockCountsWriter { + /// Sets every counter to `blocks`. + pub fn store_all(&self, blocks: u64) { + self.0.store_all(blocks); + } +} + #[cfg(test)] mod tests { use crate::plugins::worker_selection::WorkerInputView; @@ -695,7 +709,7 @@ mod tests { fn resident_blocks_reach_only_declaring_components() { use crate::indexer::ResidentBlockCountsHandle; - // Worker 8 holds no tracked blocks, which reads as 0 rather than None. Worker ids differ + // Worker 8 is tracked with no blocks, which reads as 0 rather than None. Worker ids differ // from row indices, so a row-for-worker mix-up fails. fn held(worker: WorkerWithDpRank) -> u64 { if worker.worker_id == 3 { 12 } else { 0 } @@ -778,8 +792,10 @@ mod tests { (8, TaintedWorkerConfig::default()), ]); let request = base_request(16); - let counts = ResidentBlockCountsHandle::default(); - counts.publish([(WorkerWithDpRank::new(3, 0), 12)].into_iter().collect()); + let counts = ResidentBlockCountsHandle::fixed([ + (WorkerWithDpRank::new(3, 0), 12), + (WorkerWithDpRank::new(8, 0), 0), + ]); let probes = || -> Vec> { vec![ Box::new(ResidencyProbe(WorkerInputs::RESIDENT_BLOCKS)), @@ -847,7 +863,7 @@ mod tests { Box::new(ExpectNoCounts), ) }; - // No host handle, then a handle whose poller has not published yet. + // No host handle, then a handle with no tracked ranks yet. for policy in [ policy(), policy().with_resident_block_counts(ResidentBlockCountsHandle::default()), diff --git a/lib/kv-router/src/services/indexer/backend.rs b/lib/kv-router/src/services/indexer/backend.rs index c28562f512df..06dc5aa6abdb 100644 --- a/lib/kv-router/src/services/indexer/backend.rs +++ b/lib/kv-router/src/services/indexer/backend.rs @@ -15,8 +15,8 @@ use crate::config::{ApproximateCachePolicyKind, KvRouterConfig}; use crate::indexer::{ ApproximateLruIncarnation, ApproximateLruStats, ApproximateRetentionConfig, KvIndexer, KvIndexerInterface, KvIndexerMetrics, KvRouterError, LowerTierIndexers, LowerTierQueryOptions, - MatchDetails, ResidentBlockCounts, ResidentBlockCountsHandle, RoutingDecisionHashes, - SyncIndexer, TieredMatchDetails, TieredMatchProvider, record_unsupported_residency_event, + MatchDetails, ResidentBlockCountsHandle, RoutingDecisionHashes, SyncIndexer, + TieredMatchDetails, TieredMatchProvider, record_unsupported_residency_event, }; use crate::protocols::{ DpRank, ExternalSequenceBlockHash, KvCacheEventData, LocalBlockHash, OverlapScores, @@ -871,96 +871,16 @@ impl Indexer { } } - /// Distinct blocks the primary holds for each worker rank, or `None` without a - /// local primary. The read queues behind pending events, so keep it off the - /// request path. + /// Live per-rank block counts from the local primary, which starts keeping them on the first + /// call; `None` without a local primary. #[cfg_attr(not(feature = "standalone-selection"), allow(dead_code))] - async fn resident_block_counts(&self) -> Result, KvRouterError> { + pub(crate) fn track_resident_blocks(&self) -> Option { match self { - Self::Single { primary, .. } => primary.resident_block_counts().await.map(Some), - Self::Concurrent { primary, .. } => primary.resident_block_counts().await.map(Some), - Self::Remote { .. } | Self::None => Ok(None), - } - } - - /// Republish [`Self::resident_block_counts`] every `interval` until `cancel` fires or the - /// primary stops answering, then withdraw them. A read slower than three intervals also - /// withdraws them until a read answers in time. `None` without a local primary. - #[cfg_attr(not(feature = "standalone-selection"), allow(dead_code))] - pub(crate) fn spawn_resident_block_counts_poller( - &self, - interval: Duration, - cancel: CancellationToken, - ) -> Option { - if !matches!(self, Self::Single { .. } | Self::Concurrent { .. }) { - return None; - } - let handle = ResidentBlockCountsHandle::default(); - let indexer = self.clone(); - let read = move || { - let indexer = indexer.clone(); - async move { indexer.resident_block_counts().await } - }; - tokio::spawn(publish_resident_block_counts( - read, - handle.clone(), - interval, - cancel, - )); - Some(handle) - } -} - -/// The loop behind [`Indexer::spawn_resident_block_counts_poller`], with the read injected for tests. -#[cfg_attr(not(feature = "standalone-selection"), allow(dead_code))] -async fn publish_resident_block_counts( - mut read: F, - published: ResidentBlockCountsHandle, - interval: Duration, - cancel: CancellationToken, -) where - F: FnMut() -> Fut, - Fut: Future, KvRouterError>>, -{ - let deadline = interval * 3; - let poll = async { - let mut ticks = tokio::time::interval(interval); - ticks.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); - loop { - ticks.tick().await; - let pending = read(); - tokio::pin!(pending); - let result = match tokio::time::timeout(deadline, &mut pending).await { - Ok(result) => result, - Err(_) => { - // Backlogged: withdraw the counts and let this read finish rather than queue - // more behind the backlog. Only a read that meets the deadline republishes. - published.clear(); - match pending.await { - Ok(Some(_)) => continue, - result => result, - } - } - }; - match result { - Ok(Some(counts)) => published.publish(counts), - Ok(None) => return, - Err(error) => { - // Debug: the indexer logs its own failure, and an ordinary shutdown - // can stop the indexer before this task's cancel fires. - tracing::debug!(%error, "Stopped publishing resident block counts"); - return; - } - } + Self::Single { primary, .. } => Some(primary.track_resident_blocks()), + Self::Concurrent { primary, .. } => Some(primary.track_resident_blocks()), + Self::Remote { .. } | Self::None => None, } - }; - // Biased, so a cancel that races a read failed by the same shutdown wins. - tokio::select! { - biased; - _ = cancel.cancelled() => {} - _ = poll => {} } - published.clear(); } #[async_trait] @@ -1389,10 +1309,11 @@ mod tests { ); let flat = indexer.find_matches(sequence).await.unwrap(); assert_eq!(flat.scores.get(&worker).copied(), Some(2)); - let counts = indexer.resident_block_counts().await.unwrap().unwrap(); + let counts = indexer.track_resident_blocks().unwrap(); + flush(&indexer).await; assert_eq!( - counts.get(worker), - 2, + counts.load().get(worker), + Some(2), "{num_threads} thread(s): routed prefix" ); @@ -1692,12 +1613,19 @@ mod tests { #[tokio::test] async fn resident_block_counts_track_each_rank() { use crate::protocols::compute_seq_hash_for_block; - use crate::test_utils::make_store_event_with_dp_rank; + use crate::test_utils::{make_clear_event_with_dp_rank, make_store_event_with_dp_rank}; for num_threads in [1, 4] { let indexer = create_indexer(16, num_threads); + // Stored before tracking starts, so the rank is backfilled. + indexer + .apply_event(make_store_event_with_dp_rank(1, &[1, 2, 3], 0)) + .await; + flush(&indexer).await; + let counts = indexer.track_resident_blocks().unwrap(); + let held = + |worker_id, dp_rank| counts.load().get(WorkerWithDpRank::new(worker_id, dp_rank)); for event in [ - make_store_event_with_dp_rank(1, &[1, 2, 3], 0), // A repeated store must not count its blocks twice. make_store_event_with_dp_rank(1, &[1, 2, 3], 0), make_store_event_with_dp_rank(1, &[1, 2], 1), @@ -1711,120 +1639,63 @@ mod tests { .apply_event(crate::test_utils::remove_event(1, 1, 0, vec![tail])) .await; flush(&indexer).await; - - let counts = indexer.resident_block_counts().await.unwrap().unwrap(); - let held = |worker_id, dp_rank| counts.get(WorkerWithDpRank::new(worker_id, dp_rank)); assert_eq!( [held(1, 0), held(1, 1), held(2, 0), held(3, 0)], - [2, 2, 1, 0], + [Some(2), Some(2), Some(1), None], "{num_threads} indexer thread(s)" ); + // A cleared rank stays known and empty; a removed rank is gone. + indexer + .apply_event(make_clear_event_with_dp_rank(2, 0)) + .await; indexer.remove_worker_dp_rank(1, 1).await; flush(&indexer).await; - let counts = indexer.resident_block_counts().await.unwrap().unwrap(); - let held = |worker_id, dp_rank| counts.get(WorkerWithDpRank::new(worker_id, dp_rank)); assert_eq!( - [held(1, 0), held(1, 1)], - [2, 0], + [held(2, 0), held(1, 1)], + [Some(0), None], "{num_threads} indexer thread(s)" ); + // Clones of the indexer share one set of counters. + assert_eq!( + indexer + .clone() + .track_resident_blocks() + .unwrap() + .load() + .get(WorkerWithDpRank::new(1, 0)), + Some(2) + ); } } #[tokio::test] - async fn resident_block_counts_poller_publishes_then_withdraws() { - let indexer = create_indexer(16, 4); - let cancel = CancellationToken::new(); - let handle = indexer - .spawn_resident_block_counts_poller(Duration::from_millis(5), cancel.clone()) - .expect("a local primary has counts to poll"); - let worker = WorkerWithDpRank::new(1, 0); - // Store only after the first publish, so seeing it takes a later poll. - wait_until(|| handle.load().is_some()).await; - indexer - .apply_event(crate::test_utils::make_store_event_with_dp_rank( - 1, - &[1, 2], - 0, - )) - .await; - wait_until(|| { - handle - .load() - .as_ref() - .is_some_and(|counts| counts.get(worker) == 2) - }) - .await; - - cancel.cancel(); - wait_until(|| handle.load().is_none()).await; - } - - #[tokio::test] - async fn resident_block_counts_poller_withdraws_when_the_primary_stops() { + async fn resident_block_counts_withdraw_when_the_primary_stops() { for num_threads in [1, 4] { let indexer = create_indexer(16, num_threads); - let handle = indexer - .spawn_resident_block_counts_poller( - Duration::from_millis(5), - CancellationToken::new(), - ) - .expect("a local primary has counts to poll"); - wait_until(|| handle.load().is_some()).await; + let counts = indexer.track_resident_blocks().unwrap(); + indexer + .apply_event(crate::test_utils::make_store_event_with_dp_rank( + 1, + &[1, 2], + 0, + )) + .await; + let worker = WorkerWithDpRank::new(1, 0); + wait_until(|| counts.load().get(worker) == Some(2)).await; match &indexer { Indexer::Single { primary, .. } => primary.shutdown(), Indexer::Concurrent { primary, .. } => primary.shutdown(), Indexer::Remote { .. } | Indexer::None => unreachable!(), } - // A dead primary reads as None, not as a frozen or all-zero snapshot. - wait_until(|| handle.load().is_none()).await; + // A stopped primary's ranks read None, not counts nothing updates anymore. + wait_until(|| counts.load().get(worker).is_none()).await; } } - #[tokio::test(start_paused = true)] - async fn resident_block_counts_withdraw_while_a_read_misses_the_deadline() { - let worker = WorkerWithDpRank::new(1, 0); - let published = ResidentBlockCountsHandle::default(); - let mut reads = 0; - let read = move || { - reads += 1; - let blocks = reads; - async move { - // The second read waits behind a backlog, past the 300 ms deadline. - if blocks == 2 { - tokio::time::sleep(Duration::from_secs(1)).await; - } - Ok(Some([(worker, blocks)].into_iter().collect())) - } - }; - tokio::spawn(publish_resident_block_counts( - read, - published.clone(), - Duration::from_millis(100), - CancellationToken::new(), - )); - let held = || published.load().as_ref().map(|counts| counts.get(worker)); - - tokio::time::sleep(Duration::from_millis(50)).await; - assert_eq!(held(), Some(1)); - // The second read starts at 100 ms, so its deadline passes at 400 ms. - tokio::time::sleep(Duration::from_millis(300)).await; - assert_eq!(held(), Some(1), "within the deadline"); - tokio::time::sleep(Duration::from_millis(100)).await; - assert_eq!(held(), None, "past the deadline"); - // The late read ends at 1.1 s and is dropped; the next read answers in time. - tokio::time::sleep(Duration::from_millis(700)).await; - assert_eq!(held(), Some(3)); - } - - #[tokio::test] - async fn resident_block_counts_need_a_local_primary() { - let indexer = Indexer::None; - assert!(indexer.resident_block_counts().await.unwrap().is_none()); - let poller = indexer - .spawn_resident_block_counts_poller(Duration::from_millis(5), CancellationToken::new()); - assert!(poller.is_none()); + #[test] + fn resident_block_counts_need_a_local_primary() { + assert!(Indexer::None.track_resident_blocks().is_none()); } } diff --git a/lib/kv-router/src/services/selection/core/tests.rs b/lib/kv-router/src/services/selection/core/tests.rs index 1f3a4d9c4fab..0bcfe0183953 100644 --- a/lib/kv-router/src/services/selection/core/tests.rs +++ b/lib/kv-router/src/services/selection/core/tests.rs @@ -782,7 +782,7 @@ fn resident_blocks_core( } #[tokio::test] -async fn declaring_policy_reads_polled_resident_blocks() { +async fn declaring_policy_reads_live_resident_blocks() { let (core, observed) = resident_blocks_core( crate::scheduling::selector::WorkerInputs::CACHE | crate::scheduling::selector::WorkerInputs::RESIDENT_BLOCKS, @@ -804,7 +804,7 @@ async fn declaring_policy_reads_polled_resident_blocks() { .await .unwrap(); - // The partition's poller publishes in the background, so select until the count arrives. + // The store applies on the indexer's own thread, so select until the count arrives. tokio::time::timeout(Duration::from_secs(5), async { loop { core.select(select_request()).await.expect("select"); @@ -815,19 +815,17 @@ async fn declaring_policy_reads_polled_resident_blocks() { } }) .await - .expect("a declaring policy should read the polled count"); + .expect("a declaring policy should read the live count"); } #[tokio::test] -async fn only_a_declaring_policy_polls_resident_blocks() { +async fn only_a_declaring_policy_tracks_resident_blocks() { use crate::scheduling::selector::WorkerInputs; let indexer = crate::services::indexer::backend::create_indexer(16, 1); let key = default_key(); - let poll = |inputs| { - super::workers::poll_resident_block_counts(&key, inputs, &indexer, CancellationToken::new()) - }; - assert!(poll(WorkerInputs::CACHE | WorkerInputs::LOAD).is_none()); - assert!(poll(WorkerInputs::CACHE | WorkerInputs::RESIDENT_BLOCKS).is_some()); + let track = |inputs| super::workers::track_resident_blocks(&key, inputs, &indexer); + assert!(track(WorkerInputs::CACHE | WorkerInputs::LOAD).is_none()); + assert!(track(WorkerInputs::CACHE | WorkerInputs::RESIDENT_BLOCKS).is_some()); } #[tokio::test] diff --git a/lib/kv-router/src/services/selection/core/workers.rs b/lib/kv-router/src/services/selection/core/workers.rs index fedf518f700b..13a507ebaf52 100644 --- a/lib/kv-router/src/services/selection/core/workers.rs +++ b/lib/kv-router/src/services/selection/core/workers.rs @@ -8,7 +8,7 @@ use std::collections::HashSet; use super::reservations::{ReservationIndexObserver, spawn_reservation_index_sweep}; use super::*; -use crate::indexer::{RESIDENT_BLOCK_COUNTS_POLL_INTERVAL, ResidentBlockCountsHandle}; +use crate::indexer::ResidentBlockCountsHandle; use crate::plugins::worker_selection::WorkerInputs; use crate::sequences::topology::MAX_DATA_PARALLEL_RANKS_PER_WORKER; @@ -346,12 +346,7 @@ impl SelectionCore { block_size, )) }); - if let Some(counts) = poll_resident_block_counts( - &key, - inputs, - &indexer, - self.cancel_token.child_token(), - ) { + if let Some(counts) = track_resident_blocks(&key, inputs, &indexer) { selector = selector.with_resident_block_counts(counts); } let scheduler = LocalScheduler::new( @@ -438,19 +433,17 @@ impl SelectionCore { } } -/// Polls `indexer` for a policy that reads resident block counts. Warns when the partition has +/// Starts live resident block counts for a policy that reads them. Warns when the partition has /// no local indexer, since those components can only ever read `None`. -pub(super) fn poll_resident_block_counts( +pub(super) fn track_resident_blocks( key: &RoutingPartitionId, inputs: WorkerInputs, indexer: &Indexer, - cancel: CancellationToken, ) -> Option { if !inputs.contains(WorkerInputs::RESIDENT_BLOCKS) { return None; } - let counts = - indexer.spawn_resident_block_counts_poller(RESIDENT_BLOCK_COUNTS_POLL_INTERVAL, cancel); + let counts = indexer.track_resident_blocks(); if counts.is_none() { tracing::warn!( partition = %key, From 96b381481e88d8c36af42d3570ac0997d4307629 Mon Sep 17 00:00:00 2001 From: Akira Lonske Date: Wed, 14 Oct 2026 19:25:02 -0400 Subject: [PATCH 07/12] fix(kv-router): read resident block counts from the CRTC lane lookup #16134 replaced the CRTC event thread's per-rank lookup map with `LaneLookup`. Counter updates now read `LaneLookup::block_count`, which no longer needs to be test-only, and the backfill reads `block_counts()`. Signed-off-by: Akira Lonske --- .../concurrent_radix_tree_compressed/lane_lookup.rs | 1 - .../concurrent_radix_tree_compressed/sync_impl.rs | 12 ++++-------- 2 files changed, 4 insertions(+), 9 deletions(-) diff --git a/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/lane_lookup.rs b/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/lane_lookup.rs index 806fe5e4bade..a0120caf33c0 100644 --- a/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/lane_lookup.rs +++ b/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/lane_lookup.rs @@ -318,7 +318,6 @@ impl LaneLookup { }); } - #[cfg(test)] pub(super) fn block_count(&self, worker: WorkerWithDpRank) -> Option { self.workers.get(&worker).map(BlockLookup::len) } diff --git a/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/sync_impl.rs b/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/sync_impl.rs index be1e34b46dc9..cf245732af7b 100644 --- a/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/sync_impl.rs +++ b/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/sync_impl.rs @@ -11,11 +11,11 @@ use crate::indexer::{ /// Publishes `worker`'s block count once a policy reads resident blocks. fn track( resident_blocks: &mut Option, - lookup: &FxHashMap, + lookup: &LaneLookup, worker: WorkerWithDpRank, ) { if let Some(tracker) = resident_blocks { - tracker.update(worker, lookup.get(&worker).map(WorkerLookup::len)); + tracker.update(worker, lookup.block_count(worker)); } } @@ -191,12 +191,8 @@ impl SyncIndexer for ConcurrentRadixTreeCompressed { let _ = sender.send(()); } WorkerTask::TrackResidentBlocks(counts) => { - resident_blocks = Some(ResidentBlockTracker::new( - counts, - lookup - .iter() - .map(|(worker, worker_lookup)| (*worker, worker_lookup.len())), - )); + resident_blocks = + Some(ResidentBlockTracker::new(counts, lookup.block_counts())); } WorkerTask::Terminate => { break; From 2f901dcc09634639dff0f194810a35bb719c7306 Mon Sep 17 00:00:00 2001 From: Akira Lonske Date: Thu, 15 Oct 2026 15:21:25 -0400 Subject: [PATCH 08/12] fix(kv-router): make resident block counts agree across indexer variants A live clear on the single-threaded indexer went through the rank reset, which withdraws the rank's counter, so the rank read None where the thread-pool indexer reads 0. It now applies the clear the way the thread-pool indexer does. With KV events that indexer has no prune manager or LRU, so the tree changes the same way. A rank now starts counting at its first block, so the empty block map a remove for an unknown rank can leave reads None on both indexers. The resident block test now clears through the frontend's event path and removes a whole worker, so dropping either tracking call fails it. Two tests flush instead of polling, resident_blocks_write measures tracking on single-block events, and the docs say when a rank reads None or 0. Signed-off-by: Akira Lonske --- .../router/custom-worker-selection.mdx | 2 +- .../router/custom-policy-example/README.md | 2 +- lib/kv-router/benches/worker_selection.rs | 65 ++++++++++--------- lib/kv-router/src/indexer/radix_tree.rs | 3 +- lib/kv-router/src/indexer/resident_blocks.rs | 19 ++++-- .../src/plugins/worker_selection/inputs.rs | 4 +- lib/kv-router/src/services/indexer/backend.rs | 18 +++-- .../src/services/selection/core/tests.rs | 15 +---- 8 files changed, 67 insertions(+), 61 deletions(-) diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx index 85f90712df50..7b78b48e1d27 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx @@ -454,7 +454,7 @@ A filter or scorer reads groups through `WorkerCandidate::cache()` or `load()`. `WorkerCandidate::preferred_taint_multiplier()` and `ScoredWorkerCandidate::preferred_taint_multiplier()` return the optional cost multiplier from preferred routing constraints. A filter, scorer, or picker must request `WorkerInputs::PREFERRED_TAINT` before reading the multiplier. Without that declaration, the component receives `None`, even when another component requested the multiplier. Exact hard-pinned requests also do not materialize it. Required routing constraints remain host-enforced eligibility rules. -`WorkerCandidate::resident_blocks()` and `WorkerInputView::resident_blocks(row)` return how many blocks the router's indexer holds for the worker rank, including cached blocks no running request uses; use `LOAD` for work in flight. The indexer thread that owns a rank updates its count after every event it applies, so a count is current as of the last applied event, the same freshness as the `CACHE` overlap scores, and can trail events still queued for the indexer. Reads are live, so a count can change between two reads in one selection. With KV events, the count is the device-tier blocks the indexer currently tracks from engine events. With approximate routing, it is the blocks routed to the rank within the TTL, or the distinct full blocks the LRU cache policy retains, which leaves out its partial tail blocks and duplicate copies. The count is unweighted and not clamped to `total_kv_blocks`. It measures occupancy, not reusable cache: it includes blocks that prefix matching can no longer reach, such as descendants of an evicted block, until their own remove events arrive, and lost remove events inflate it until the rank is removed. A rank the indexer has seen but holds no blocks for, such as after a clear, reads `0`. A component that did not request `WorkerInputs::RESIDENT_BLOCKS` receives `None`, and so does every component for a rank the indexer has not seen or has removed, with a remote indexer, or after an indexer thread stops; a worker whose KV events never arrive therefore reads `None` rather than looking empty. A policy that requests `RESIDENT_BLOCKS` must also request `CACHE`, because the counts come from the indexer that `CACHE` starts. The router checks this whenever it builds a policy for a routing partition: the frontend fails to start, and the standalone selection service marks the partition's workers unschedulable with the reason. With KV events every router replica sees the same counts. With approximate routing each replica counts only the requests it routed, except a replica that serves its indexer (`serve_indexer`), whose counts also include the routing decisions other replicas send it. +`WorkerCandidate::resident_blocks()` and `WorkerInputView::resident_blocks(row)` return how many blocks the router's indexer holds for the worker rank, including cached blocks no running request uses; use `LOAD` for work in flight. The indexer thread that owns a rank updates its count after every event it applies, so a count is current as of the last applied event, the same freshness as the `CACHE` overlap scores, and can trail events still queued for the indexer. Reads are live, so a count can change between two reads in one selection. With KV events, the count is the device-tier blocks the indexer currently tracks from engine events. With approximate routing, it is the blocks routed to the rank within the TTL, or the distinct full blocks the LRU cache policy retains, which leaves out its partial tail blocks and duplicate copies. The count is unweighted and not clamped to `total_kv_blocks`. It measures occupancy, not reusable cache: it includes blocks that prefix matching can no longer reach, such as descendants of an evicted block, until their own remove events arrive, and lost remove events inflate it until the rank is removed. A rank reads `None` until the indexer holds a block for it while counting. After that, a clear sets its count to `0`, and removing the rank or its worker returns it to `None`. Counting starts when each indexer thread reaches the policy's request, so ranks on a thread with queued events read `None` until it catches up. A component that did not request `WorkerInputs::RESIDENT_BLOCKS` receives `None`, and so does every component with a remote indexer or after an indexer thread stops; a worker whose KV events never arrive therefore reads `None` rather than looking empty. A policy that requests `RESIDENT_BLOCKS` must also request `CACHE`, because the counts come from the indexer that `CACHE` starts. The router checks this whenever it builds a policy for a routing partition: the frontend fails to start, and the standalone selection service marks the partition's workers unschedulable with the reason. With KV events every router replica sees the same counts. With approximate routing each replica counts only the requests it routed, except a replica that serves its indexer (`serve_indexer`), whose counts also include the routing decisions other replicas send it. ## Link the Policy Into Dynamo diff --git a/examples/router/custom-policy-example/README.md b/examples/router/custom-policy-example/README.md index 48040725d943..0c8d9d7dd5ba 100644 --- a/examples/router/custom-policy-example/README.md +++ b/examples/router/custom-policy-example/README.md @@ -24,7 +24,7 @@ Declare `WorkerInputs::CACHE` to read cache data through `candidate.cache()` (fi Preferred routing taints are optional candidate metadata. A filter, scorer, or picker must request `WorkerInputs::PREFERRED_TAINT` before reading `preferred_taint_multiplier()` from a candidate; otherwise, that component receives `None`, even if another component requested it. Exact hard-pinned requests also do not materialize it. Required routing taints remain Dynamo eligibility rules. -Declare `WorkerInputs::RESIDENT_BLOCKS` to read `resident_blocks()` from a candidate, or `input.resident_blocks(row)` in a picker. It returns how many blocks the router's indexer holds for that worker rank, including cached blocks no running request uses, as of the last event the indexer applied. Without the declaration, with a remote indexer, or for a rank the indexer has not seen or has removed, it returns `None`. A policy that requests `RESIDENT_BLOCKS` must also request `CACHE`. +Declare `WorkerInputs::RESIDENT_BLOCKS` to read `resident_blocks()` from a candidate, or `input.resident_blocks(row)` in a picker. It returns how many blocks the router's indexer holds for that worker rank, including cached blocks no running request uses, as of the last event the indexer applied. Without the declaration, with a remote indexer, or for a rank the indexer has held no blocks for or has removed, it returns `None`; a cleared rank returns `0`. A policy that requests `RESIDENT_BLOCKS` must also request `CACHE`. ## Pick a Starting Point diff --git a/lib/kv-router/benches/worker_selection.rs b/lib/kv-router/benches/worker_selection.rs index 3f1459621e98..45616fe57822 100644 --- a/lib/kv-router/benches/worker_selection.rs +++ b/lib/kv-router/benches/worker_selection.rs @@ -15,7 +15,7 @@ use dynamo_kv_router::indexer::{ concurrent_radix_tree_compressed::ConcurrentRadixTreeCompressed, }; use dynamo_kv_router::protocols::{ - ExternalSequenceBlockHash, KvCacheEvent, KvCacheEventData, KvCacheStoreData, + ExternalSequenceBlockHash, KvCacheEvent, KvCacheEventData, KvCacheRemoveData, KvCacheStoreData, KvCacheStoredBlockData, LocalBlockHash, RouterEvent, RoutingConstraints, WorkerConfigLike, WorkerId, WorkerWithDpRank, compute_seq_hash_for_block, }; @@ -498,43 +498,54 @@ fn kv_event(worker: WorkerId, event_id: u64, data: KvCacheEventData) -> RouterEv ) } -/// Event-thread cost of keeping resident block counts: stores then clears for 128 workers on -/// 8 event threads, untracked and tracked. A concurrent reader belongs in the pinned Mooncake -/// replay (`--resident-blocks-reader-cpu`): unpinned, any busy thread skews this bench. +/// Event-thread cost of keeping resident block counts, untracked and tracked: 128 workers on 8 +/// event threads each store a 16-block chain one block per event, then remove it the same way. +/// Single-block events are the cheapest to apply, so the per-event tracking cost is the largest +/// share of them. Nothing reads the counters here; unpinned, a busy reader thread skews this +/// bench. The pinned Mooncake replay's `--resident-blocks-reader-cpu` covers a reader, at that +/// trace's larger events. fn resident_blocks_write(c: &mut Criterion) { const WORKERS: u64 = 128; - const BLOCKS: u64 = 64; + const BLOCKS: u64 = 16; let runtime = tokio::runtime::Builder::new_current_thread() .enable_all() .build() .unwrap(); - let stores: Vec<_> = (0..WORKERS) - .map(|worker| { + let events: Vec<_> = (0..WORKERS) + .flat_map(|worker| { let local: Vec<_> = (0..BLOCKS) .map(|i| LocalBlockHash(worker * BLOCKS + i)) .collect(); - let blocks = local - .iter() - .zip(compute_seq_hash_for_block(&local)) - .map(|(&tokens_hash, hash)| KvCacheStoredBlockData { - block_hash: ExternalSequenceBlockHash(hash), - tokens_hash, - mm_extra_info: None, - }) + let hashes: Vec<_> = compute_seq_hash_for_block(&local) + .into_iter() + .map(ExternalSequenceBlockHash) .collect(); - let store = KvCacheStoreData { - parent_hash: None, - start_position: None, - blocks, - }; - kv_event(worker, 0, KvCacheEventData::Stored(store)) + let stores = (0..hashes.len()).map(|i| { + let store = KvCacheStoreData { + parent_hash: i.checked_sub(1).map(|parent| hashes[parent]), + start_position: None, + blocks: vec![KvCacheStoredBlockData { + block_hash: hashes[i], + tokens_hash: local[i], + mm_extra_info: None, + }], + }; + kv_event(worker, i as u64, KvCacheEventData::Stored(store)) + }); + let removes = (0..hashes.len()).rev().map(|i| { + let remove = KvCacheRemoveData { + block_hashes: vec![hashes[i]], + }; + kv_event(worker, BLOCKS + i as u64, KvCacheEventData::Removed(remove)) + }); + stores.chain(removes).collect::>() }) .collect(); let mut group = c.benchmark_group("resident_blocks_write"); group.sample_size(20); group.warm_up_time(Duration::from_secs(1)); group.measurement_time(Duration::from_secs(3)); - group.throughput(Throughput::Elements(WORKERS * 2)); + group.throughput(Throughput::Elements(events.len() as u64)); for (name, tracked) in [("untracked", false), ("tracked", true)] { let indexer = ThreadPoolIndexer::new(ConcurrentRadixTreeCompressed::new(), 8, 32); if tracked { @@ -543,14 +554,8 @@ fn resident_blocks_write(c: &mut Criterion) { group.bench_function(name, |b| { b.iter(|| { runtime.block_on(async { - for store in &stores { - indexer.apply_event(store.clone()).await; - } - indexer.flush().await; - for worker in 0..WORKERS { - indexer - .apply_event(kv_event(worker, 1, KvCacheEventData::Cleared)) - .await; + for event in &events { + indexer.apply_event(event.clone()).await; } indexer.flush().await; }) diff --git a/lib/kv-router/src/indexer/radix_tree.rs b/lib/kv-router/src/indexer/radix_tree.rs index 5c9edf3ccd89..8a907006b05f 100644 --- a/lib/kv-router/src/indexer/radix_tree.rs +++ b/lib/kv-router/src/indexer/radix_tree.rs @@ -773,7 +773,8 @@ impl RadixTree { } } - /// Drops `worker`'s blocks without forgetting the rank, as a clear does. + /// Drops `worker`'s blocks and lookup entry. Callers decide the rank's resident block counter: + /// a clear leaves it at 0, a removal withdraws it. fn drop_rank(&mut self, worker: WorkerWithDpRank) { let Some(blocks) = self.lookup.remove(&worker) else { return; diff --git a/lib/kv-router/src/indexer/resident_blocks.rs b/lib/kv-router/src/indexer/resident_blocks.rs index f60de91d9bb0..ae83242bcae5 100644 --- a/lib/kv-router/src/indexer/resident_blocks.rs +++ b/lib/kv-router/src/indexer/resident_blocks.rs @@ -31,8 +31,8 @@ pub(crate) struct ResidentBlockCounts { } impl ResidentBlockCounts { - /// Blocks held for `worker` as of the last event applied to it; `None` for a rank the indexer - /// has not seen or has removed. + /// Blocks held for `worker` as of the last event applied to it; `None` for a rank that has held + /// no blocks while tracked, or that the indexer removed. pub(crate) fn get(&self, worker: WorkerWithDpRank) -> Option { self.counters .get(&worker) @@ -121,13 +121,14 @@ pub(crate) struct ResidentBlockTracker { } impl ResidentBlockTracker { - /// Starts tracking, with counters for the ranks the thread already holds. + /// Starts tracking, with counters for the ranks the thread already holds blocks for. pub(crate) fn new( counts: ResidentBlockCountsHandle, ranks: impl IntoIterator, ) -> Self { let owned: FxHashMap<_, _> = ranks .into_iter() + .filter(|&(_, blocks)| blocks > 0) .map(|(worker, blocks)| (worker, counter(blocks))) .collect(); let added: Vec<_> = owned @@ -138,12 +139,13 @@ impl ResidentBlockTracker { Self { counts, owned } } - /// Records `worker`'s block count after an event. `None` means the thread holds no blocks - /// for it, as after a clear; a rank already counted then reads 0. + /// Records `worker`'s block count after an event. A rank is counted from its first block on; + /// `None` means the thread holds no blocks for it, as after a clear, and a counted rank then + /// reads 0. pub(crate) fn update(&mut self, worker: WorkerWithDpRank, blocks: Option) { if let Some(counter) = self.owned.get(&worker) { counter.store(blocks.unwrap_or(0) as u64, Ordering::Relaxed); - } else if let Some(blocks) = blocks { + } else if let Some(blocks) = blocks.filter(|&blocks| blocks > 0) { let counter = counter(blocks); self.owned.insert(worker, counter.clone()); self.counts.publish(&[(worker, counter)]); @@ -186,10 +188,13 @@ mod tests { WorkerWithDpRank::new(2, 0), WorkerWithDpRank::new(3, 0), ); - let mut tracker = ResidentBlockTracker::new(counts.clone(), [(held, 4), (emptied, 2)]); + // Empty block maps, as a remove for an unknown rank can leave, start no count. + let mut tracker = + ResidentBlockTracker::new(counts.clone(), [(held, 4), (emptied, 2), (unseen, 0)]); tracker.update(held, Some(5)); tracker.update(emptied, None); tracker.update(removed, Some(7)); + tracker.update(unseen, Some(0)); let read = |worker| counts.load().get(worker); assert_eq!( [read(held), read(emptied), read(removed), read(unseen)], diff --git a/lib/kv-router/src/plugins/worker_selection/inputs.rs b/lib/kv-router/src/plugins/worker_selection/inputs.rs index d81d6077c22d..ab4680fb0dab 100644 --- a/lib/kv-router/src/plugins/worker_selection/inputs.rs +++ b/lib/kv-router/src/plugins/worker_selection/inputs.rs @@ -72,8 +72,8 @@ impl<'a> WorkerCandidate<'a> { /// Blocks the router's indexer holds for this worker rank as of the last event it applied, only /// when this component declared RESIDENT_BLOCKS: device blocks from KV events, or the router's /// prediction in approximate modes. Unweighted and not clamped to capacity. None means access - /// was not requested, the host has no local indexer, or the indexer has not seen the rank or - /// has removed it. + /// was not requested, the host has no local indexer, or the indexer has held no blocks for the + /// rank while counting or has removed it. pub fn resident_blocks(self) -> Option { if !self.inputs.contains(WorkerInputs::RESIDENT_BLOCKS) { return None; diff --git a/lib/kv-router/src/services/indexer/backend.rs b/lib/kv-router/src/services/indexer/backend.rs index 06dc5aa6abdb..0357d96d2169 100644 --- a/lib/kv-router/src/services/indexer/backend.rs +++ b/lib/kv-router/src/services/indexer/backend.rs @@ -468,9 +468,7 @@ impl Indexer { } => { if is_clear { if targets_primary { - primary - .reset_worker_dp_rank_and_wait(event.worker_id, event.event.dp_rank) - .await?; + primary.apply_event_and_wait(event.clone()).await?; } for indexer in lower_tier.all() { indexer.apply_event_and_wait(event.clone()).await?; @@ -1645,10 +1643,12 @@ mod tests { "{num_threads} indexer thread(s)" ); - // A cleared rank stays known and empty; a removed rank is gone. + // A cleared rank stays known and empty; a removed rank is gone. The clear takes the + // path the frontend admits live engine events through. indexer - .apply_event(make_clear_event_with_dp_rank(2, 0)) - .await; + .try_apply_event(make_clear_event_with_dp_rank(2, 0)) + .await + .unwrap(); indexer.remove_worker_dp_rank(1, 1).await; flush(&indexer).await; assert_eq!( @@ -1666,6 +1666,9 @@ mod tests { .get(WorkerWithDpRank::new(1, 0)), Some(2) ); + indexer.remove_worker(1).await; + flush(&indexer).await; + assert_eq!(held(1, 0), None, "{num_threads} indexer thread(s)"); } } @@ -1682,7 +1685,8 @@ mod tests { )) .await; let worker = WorkerWithDpRank::new(1, 0); - wait_until(|| counts.load().get(worker) == Some(2)).await; + flush(&indexer).await; + assert_eq!(counts.load().get(worker), Some(2)); match &indexer { Indexer::Single { primary, .. } => primary.shutdown(), Indexer::Concurrent { primary, .. } => primary.shutdown(), diff --git a/lib/kv-router/src/services/selection/core/tests.rs b/lib/kv-router/src/services/selection/core/tests.rs index 08399b6cc75c..062c01910cc1 100644 --- a/lib/kv-router/src/services/selection/core/tests.rs +++ b/lib/kv-router/src/services/selection/core/tests.rs @@ -801,19 +801,10 @@ async fn declaring_policy_reads_live_resident_blocks() { )) .await .unwrap(); + entry.indexer.flush().await; - // The store applies on the indexer's own thread, so select until the count arrives. - tokio::time::timeout(Duration::from_secs(5), async { - loop { - core.select(select_request()).await.expect("select"); - if observed.lock().last() == Some(&Some(3)) { - break; - } - tokio::time::sleep(Duration::from_millis(20)).await; - } - }) - .await - .expect("a declaring policy should read the live count"); + core.select(select_request()).await.expect("select"); + assert_eq!(observed.lock().last(), Some(&Some(3))); } #[tokio::test] From ec56b4828e29a669611c696be07d54ade84929fb Mon Sep 17 00:00:00 2001 From: Akira Lonske Date: Tue, 20 Oct 2026 15:18:58 -0400 Subject: [PATCH 09/12] refactor(kv-router): skip no-op counter map swaps and trim the resident blocks docs publish and unpublish return early on an empty slice, so a tracker started on a thread with no counted ranks, or dropped with nothing owned, no longer clones the published map for a no-op rcu. remove's own emptiness check moves there. The resident_blocks() paragraph in custom-worker-selection.mdx now carries the contract (what it counts, freshness, None vs 0, the CACHE requirement) in 229 words instead of 370, with the per-mode and per-replica details in a list below it. The test that only restated the Indexer::None arm of track_resident_blocks is dropped. Signed-off-by: Akira Lonske --- .../router/custom-worker-selection.mdx | 7 ++++++- lib/kv-router/src/indexer/resident_blocks.rs | 10 +++++++--- lib/kv-router/src/services/indexer/backend.rs | 5 ----- 3 files changed, 13 insertions(+), 9 deletions(-) diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx index 7b78b48e1d27..b5e82a0b8d1b 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx @@ -454,7 +454,12 @@ A filter or scorer reads groups through `WorkerCandidate::cache()` or `load()`. `WorkerCandidate::preferred_taint_multiplier()` and `ScoredWorkerCandidate::preferred_taint_multiplier()` return the optional cost multiplier from preferred routing constraints. A filter, scorer, or picker must request `WorkerInputs::PREFERRED_TAINT` before reading the multiplier. Without that declaration, the component receives `None`, even when another component requested the multiplier. Exact hard-pinned requests also do not materialize it. Required routing constraints remain host-enforced eligibility rules. -`WorkerCandidate::resident_blocks()` and `WorkerInputView::resident_blocks(row)` return how many blocks the router's indexer holds for the worker rank, including cached blocks no running request uses; use `LOAD` for work in flight. The indexer thread that owns a rank updates its count after every event it applies, so a count is current as of the last applied event, the same freshness as the `CACHE` overlap scores, and can trail events still queued for the indexer. Reads are live, so a count can change between two reads in one selection. With KV events, the count is the device-tier blocks the indexer currently tracks from engine events. With approximate routing, it is the blocks routed to the rank within the TTL, or the distinct full blocks the LRU cache policy retains, which leaves out its partial tail blocks and duplicate copies. The count is unweighted and not clamped to `total_kv_blocks`. It measures occupancy, not reusable cache: it includes blocks that prefix matching can no longer reach, such as descendants of an evicted block, until their own remove events arrive, and lost remove events inflate it until the rank is removed. A rank reads `None` until the indexer holds a block for it while counting. After that, a clear sets its count to `0`, and removing the rank or its worker returns it to `None`. Counting starts when each indexer thread reaches the policy's request, so ranks on a thread with queued events read `None` until it catches up. A component that did not request `WorkerInputs::RESIDENT_BLOCKS` receives `None`, and so does every component with a remote indexer or after an indexer thread stops; a worker whose KV events never arrive therefore reads `None` rather than looking empty. A policy that requests `RESIDENT_BLOCKS` must also request `CACHE`, because the counts come from the indexer that `CACHE` starts. The router checks this whenever it builds a policy for a routing partition: the frontend fails to start, and the standalone selection service marks the partition's workers unschedulable with the reason. With KV events every router replica sees the same counts. With approximate routing each replica counts only the requests it routed, except a replica that serves its indexer (`serve_indexer`), whose counts also include the routing decisions other replicas send it. +`WorkerCandidate::resident_blocks()` and `WorkerInputView::resident_blocks(row)` return how many blocks the router's indexer holds for the worker rank, including cached blocks no running request uses; use `LOAD` for work in flight. The count is unweighted, not clamped to `total_kv_blocks`, and measures occupancy rather than reusable cache: blocks that prefix matching can no longer reach, such as descendants of an evicted block, count until their own remove events arrive, and a lost remove event inflates the count until the rank is removed. A count is current as of the last event its indexer thread applied, the same freshness as the `CACHE` overlap scores, so it can trail queued events and change between two reads in one selection. A rank reads `None` until its indexer thread has reached the policy's request and holds a block for it, `0` after a clear, and `None` again once the rank or its worker is removed. A component also reads `None` when it did not request `WorkerInputs::RESIDENT_BLOCKS`, with a remote indexer, or after an indexer thread stops, so a worker whose KV events never arrive never looks empty. A policy that requests `RESIDENT_BLOCKS` must also request `CACHE`, because the counts come from the indexer that `CACHE` starts. The router checks this for every policy it builds for a routing partition: the frontend fails to start, and the standalone selection service marks the partition's workers unschedulable with the reason. + +What the count contains, and whether router replicas agree on it, depends on the routing mode: + +- **KV events:** the device-tier blocks the indexer currently tracks from engine events. Every router replica sees the same counts. +- **Approximate routing:** the blocks routed to the rank within the TTL, or the distinct full blocks the LRU cache policy retains, which leaves out its partial tail blocks and duplicate copies. Each replica counts only the requests it routed, except a replica that serves its indexer (`serve_indexer`), whose counts also include the routing decisions other replicas send it. ## Link the Policy Into Dynamo diff --git a/lib/kv-router/src/indexer/resident_blocks.rs b/lib/kv-router/src/indexer/resident_blocks.rs index ae83242bcae5..194438f9e5c8 100644 --- a/lib/kv-router/src/indexer/resident_blocks.rs +++ b/lib/kv-router/src/indexer/resident_blocks.rs @@ -85,6 +85,9 @@ impl ResidentBlockCountsHandle { } fn publish(&self, added: &[(WorkerWithDpRank, Arc)]) { + if added.is_empty() { + return; + } self.counts.rcu(|current| { let mut next = ResidentBlockCounts::clone(current); next.counters.extend( @@ -97,6 +100,9 @@ impl ResidentBlockCountsHandle { } fn unpublish(&self, removed: &[(WorkerWithDpRank, Arc)]) { + if removed.is_empty() { + return; + } self.counts.rcu(|current| { let mut next = ResidentBlockCounts::clone(current); for (worker, counter) in removed { @@ -162,9 +168,7 @@ impl ResidentBlockTracker { } keep }); - if !gone.is_empty() { - self.counts.unpublish(&gone); - } + self.counts.unpublish(&gone); } } diff --git a/lib/kv-router/src/services/indexer/backend.rs b/lib/kv-router/src/services/indexer/backend.rs index 0357d96d2169..6e73aa0debb0 100644 --- a/lib/kv-router/src/services/indexer/backend.rs +++ b/lib/kv-router/src/services/indexer/backend.rs @@ -1696,11 +1696,6 @@ mod tests { wait_until(|| counts.load().get(worker).is_none()).await; } } - - #[test] - fn resident_block_counts_need_a_local_primary() { - assert!(Indexer::None.track_resident_blocks().is_none()); - } } pub fn create_indexer(block_size: u32, num_threads: usize) -> Indexer { From c8625efd789f5967973e5215c0ce2469eff67d6e Mon Sep 17 00:00:00 2001 From: Akira Lonske Date: Sat, 24 Oct 2026 18:45:52 +0000 Subject: [PATCH 10/12] fix(kv-router): pin the resident block reader before the bench trial starts The reader thread now reports its CPU pin to the coordinator, which aborts before the trial when the pin fails, instead of running a trial without a reader and discarding it. The bench guide shows how to free a CPU for the reader on the documented 128-core binding. Signed-off-by: Akira Lonske --- lib/bench/kv_router/INDEXER_BENCH.md | 4 +++- lib/bench/kv_router/mooncake_bench.rs | 27 ++++++++++++++++++++------- 2 files changed, 23 insertions(+), 8 deletions(-) diff --git a/lib/bench/kv_router/INDEXER_BENCH.md b/lib/bench/kv_router/INDEXER_BENCH.md index 59c819d6e40c..4c5e98c88584 100644 --- a/lib/bench/kv_router/INDEXER_BENCH.md +++ b/lib/bench/kv_router/INDEXER_BENCH.md @@ -166,7 +166,9 @@ and compare `achieved_block_ops_per_sec` against an untracked run. Add `--resident-blocks-reader-cpu ` to also read every counter in a loop from a thread pinned to that CPU, as selections do, so the result includes the cache-line traffic between event threads and readers. Keep that CPU outside the backend and -issuer masks. +issuer masks. The 128-core binding above uses every CPU, so take one out of the +backend mask first, for example `--backend-cpus 4-62,69-127 +--resident-blocks-reader-cpu 63`. ### Sweep diff --git a/lib/bench/kv_router/mooncake_bench.rs b/lib/bench/kv_router/mooncake_bench.rs index a52adbd1de33..982d40718d20 100644 --- a/lib/bench/kv_router/mooncake_bench.rs +++ b/lib/bench/kv_router/mooncake_bench.rs @@ -310,20 +310,33 @@ async fn run_open_loop_for_config( .zip(args.resident_blocks_reader_cpu) .map(|(counts, cpu)| { let stop = stop.clone(); - std::thread::spawn(move || { - pin_current_thread_to_cpus(&[cpu])?; + let (pinned_tx, pinned_rx) = std::sync::mpsc::sync_channel(1); + let handle = std::thread::spawn(move || { + let pinned = pin_current_thread_to_cpus(&[cpu]); + let failed = pinned.is_err(); + let _ = pinned_tx.send(pinned); + if failed { + return; + } while !stop.load(Ordering::Relaxed) { std::hint::black_box(counts.bench_sum()); } - std::io::Result::Ok(()) - }) + }); + (handle, pinned_rx) }); + // Don't start the trial until the reader is pinned, so a failed pin + // aborts here instead of producing a run with no reader. + if let Some((_, pinned_rx)) = &reader { + pinned_rx.recv().map_err(|_| { + anyhow::anyhow!("resident block reader exited before pinning") + })??; + } let result = run_backend(config.short_name(), indexer, trial, open_config).await; stop.store(true, Ordering::Relaxed); - if let Some(reader) = reader { - reader + if let Some((handle, _)) = reader { + handle .join() - .map_err(|_| anyhow::anyhow!("resident block reader panicked"))??; + .map_err(|_| anyhow::anyhow!("resident block reader panicked"))?; } result } From 99d4fd6ac8c94780dce11a9788351e9e593cfa8f Mon Sep 17 00:00:00 2001 From: Akira Lonske Date: Tue, 3 Nov 2026 09:56:59 -0500 Subject: [PATCH 11/12] refactor(kv-router): move resident block counts under the experimental plugin namespace Export RESIDENT_BLOCKS and the two resident_blocks() accessors from plugins::worker_selection::experimental::resident_blocks, as a sealed extension trait, so the stable WorkerInputs surface stays unchanged. The counters and host plumbing do not move. Tests, benches, the frontend prepare() gate, and the docs follow the new path. Signed-off-by: Akira Lonske --- .../router/custom-worker-selection.mdx | 50 ++++++++-- .../router/custom-policy-example/README.md | 2 +- lib/kv-router/benches/worker_selection.rs | 5 +- lib/kv-router/src/plugins/worker_selection.rs | 1 + .../plugins/worker_selection/experimental.rs | 9 ++ .../experimental/resident_blocks.rs | 95 +++++++++++++++++++ .../src/plugins/worker_selection/inputs.rs | 21 ++-- .../src/scheduling/selector/policy.rs | 3 + .../src/services/selection/core/tests.rs | 1 + lib/llm/src/kv_router.rs | 18 ++-- 10 files changed, 172 insertions(+), 33 deletions(-) create mode 100644 lib/kv-router/src/plugins/worker_selection/experimental.rs create mode 100644 lib/kv-router/src/plugins/worker_selection/experimental/resident_blocks.rs diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx index c8ef440d1cd8..b34df6efe641 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx @@ -53,6 +53,7 @@ Use `dynamo_kv_router::plugins` for plugin registration and these modules for pl |---|---| | `plugins::worker_selection` | `WorkerFilter`, `WorkerScorer`, `WorkerPicker`, `WorkerSelectionPolicy`, request context, worker signal accessors, configuration, and factories | | `plugins` | `RouterPluginRegistry` and the resolved `RouterPlugins` bundle | +| `plugins::worker_selection::experimental` | Unstable extensions, such as the [resident block counts](#experimental-resident-block-counts). No compatibility guarantee | Dynamo owns scheduler queues, eligibility checks, reservations, lifecycle synchronization, and the default selector implementation. Low-level host interfaces such as `WorkerSelector` are used to connect Dynamo components; external plugins use `WorkerSelectionPolicy`. @@ -370,7 +371,7 @@ Each worker policy is constructed once per routing partition. An empty bundle re Worker identity and ordinary request metadata are always available. Dynamo calculates optional per-worker signals only for groups that a filter, scorer, or picker requests. -Routing cache resources start only if a component requests `CACHE`. `serve_indexer`, `enable_session_prefix_index`, and `RESIDENT_BLOCKS` also require `CACHE`. +Routing cache resources start only if a component requests `CACHE`. `serve_indexer`, `enable_session_prefix_index`, and the experimental `RESIDENT_BLOCKS` input also require `CACHE`. ### Request Context @@ -446,7 +447,6 @@ If a component needs no optional worker data, return `WorkerInputs::NONE`. Combi | `LOAD` | `decode_cost_blocks()` | Projected active decode footprint, including this request's additional active blocks | | `LOAD` | `active_requests()` | Requests currently active on the worker | | `PREFERRED_TAINT` | `preferred_taint_multiplier()` | Optional cost multiplier from matching preferred routing taints | -| `RESIDENT_BLOCKS` | `resident_blocks()` | Blocks the router's indexer tracks for the worker rank, including cached blocks no running request uses; current as of the last event the indexer applied and not clamped to capacity | The GPU, CPU, and disk overlap accessors return raw tier counts. A missing tier or worker entry is zero; Dynamo does not substitute its weighted effective-overlap estimate. Each custom scorer chooses how to combine the raw counts. @@ -454,13 +454,6 @@ A filter or scorer reads groups through `WorkerCandidate::cache()` or `load()`. `WorkerCandidate::preferred_taint_multiplier()` and `ScoredWorkerCandidate::preferred_taint_multiplier()` return the optional cost multiplier from preferred routing constraints. A filter, scorer, or picker must request `WorkerInputs::PREFERRED_TAINT` before reading the multiplier. Without that declaration, the component receives `None`, even when another component requested the multiplier. Exact hard-pinned requests also do not materialize it. Required routing constraints remain host-enforced eligibility rules. -`WorkerCandidate::resident_blocks()` and `WorkerInputView::resident_blocks(row)` return how many blocks the router's indexer holds for the worker rank, including cached blocks no running request uses; use `LOAD` for work in flight. The count is unweighted, not clamped to `total_kv_blocks`, and measures occupancy rather than reusable cache: blocks that prefix matching can no longer reach, such as descendants of an evicted block, count until their own remove events arrive, and a lost remove event inflates the count until the rank is removed. A count is current as of the last event its indexer thread applied, the same freshness as the `CACHE` overlap scores, so it can trail queued events and change between two reads in one selection. A rank reads `None` until its indexer thread has reached the policy's request and holds a block for it, `0` after a clear, and `None` again once the rank or its worker is removed. A component also reads `None` when it did not request `WorkerInputs::RESIDENT_BLOCKS`, with a remote indexer, or after an indexer thread stops, so a worker whose KV events never arrive never looks empty. A policy that requests `RESIDENT_BLOCKS` must also request `CACHE`, because the counts come from the indexer that `CACHE` starts. The router checks this for every policy it builds for a routing partition: the frontend fails to start, and the standalone selection service marks the partition's workers unschedulable with the reason. - -What the count contains, and whether router replicas agree on it, depends on the routing mode: - -- **KV events:** the device-tier blocks the indexer currently tracks from engine events. Every router replica sees the same counts. -- **Approximate routing:** the blocks routed to the rank within the TTL, or the distinct full blocks the LRU cache policy retains, which leaves out its partial tail blocks and duplicate copies. Each replica counts only the requests it routed, except a replica that serves its indexer (`serve_indexer`), whose counts also include the routing decisions other replicas send it. - ## Link the Policy Into Dynamo Both paths use the same policy crate, catalog, and YAML file. Choose the process that owns worker selection. @@ -584,6 +577,45 @@ impl RequestClassifier for InspectPrefixes { } ``` +## Experimental Resident Block Counts + +A policy can read how many blocks the router's indexer holds for each worker rank, including cached blocks no running request uses, and charge a placement for the cache it would evict; use `LOAD` for work in flight. The API lives in `plugins::worker_selection::experimental`, whose items have no compatibility guarantee and may change or be removed in any release. Declare `resident_blocks::RESIDENT_BLOCKS` together with `WorkerInputs::CACHE`, then import the `ResidentBlocks` trait to read `candidate.resident_blocks()` in a filter or scorer, or `ResidentBlocksByRow` to read `input.resident_blocks(row)` in a picker: + +```rust +use dynamo_kv_router::plugins::worker_selection::experimental::resident_blocks::{ + RESIDENT_BLOCKS, ResidentBlocks, +}; + +/// Charges one unit per block a placement could evict. +struct EvictionPressure; + +impl WorkerScorer for EvictionPressure { + fn required_worker_inputs(&self) -> WorkerInputs { + WorkerInputs::CACHE | RESIDENT_BLOCKS + } + + fn score( + &mut self, + _context: &WorkerSelectionContext<'_>, + candidates: WorkerCandidates<'_>, + costs: &mut [f64], + ) -> Result<(), WorkerSelectionPolicyError> { + for (candidate, cost) in candidates.iter().zip(costs) { + // None means the host has no count for this rank; treat it as neutral. + *cost = candidate.resident_blocks().unwrap_or(0) as f64; + } + Ok(()) + } +} +``` + +The count is unweighted, not clamped to `total_kv_blocks`, and measures occupancy rather than reusable cache: blocks that prefix matching can no longer reach, such as descendants of an evicted block, count until their own remove events arrive, and a lost remove event inflates the count until the rank is removed. A count is current as of the last event its indexer thread applied, the same freshness as the `CACHE` overlap scores, so it can trail queued events and change between two reads in one selection. A rank reads `None` until its indexer thread has reached the policy's request and holds a block for it, `0` after a clear, and `None` again once the rank or its worker is removed. A component also reads `None` when it did not declare `RESIDENT_BLOCKS`, with a remote indexer, or after an indexer thread stops, so a worker whose KV events never arrive never looks empty. A policy that declares `RESIDENT_BLOCKS` must also declare `CACHE`, because the counts come from the indexer that `CACHE` starts. The router checks this for every policy it builds for a routing partition: the frontend fails to start, and the standalone selection service marks the partition's workers unschedulable with the reason. + +What the count contains, and whether router replicas agree on it, depends on the routing mode: + +- **KV events:** the device-tier blocks the indexer currently tracks from engine events. Every router replica sees the same counts. +- **Approximate routing:** the blocks routed to the rank within the TTL, or the distinct full blocks the LRU cache policy retains, which leaves out its partial tail blocks and duplicate copies. Each replica counts only the requests it routed, except a replica that serves its indexer (`serve_indexer`), whose counts also include the routing decisions other replicas send it. + ## Policy Contract - Return `true` from a filter to keep a worker and `false` to reject it. diff --git a/examples/router/custom-policy-example/README.md b/examples/router/custom-policy-example/README.md index 6730f36d95cf..821c3fb41644 100644 --- a/examples/router/custom-policy-example/README.md +++ b/examples/router/custom-policy-example/README.md @@ -24,7 +24,7 @@ Declare `WorkerInputs::CACHE` to read cache data through `candidate.cache()` (fi Preferred routing taints are optional candidate metadata. A filter, scorer, or picker must request `WorkerInputs::PREFERRED_TAINT` before reading `preferred_taint_multiplier()` from a candidate; otherwise, that component receives `None`, even if another component requested it. Exact hard-pinned requests also do not materialize it. Required routing taints remain Dynamo eligibility rules. -Declare `WorkerInputs::RESIDENT_BLOCKS` to read `resident_blocks()` from a candidate, or `input.resident_blocks(row)` in a picker. It returns how many blocks the router's indexer holds for that worker rank, including cached blocks no running request uses, as of the last event the indexer applied. Without the declaration, with a remote indexer, or for a rank the indexer has held no blocks for or has removed, it returns `None`; a cleared rank returns `0`. A policy that requests `RESIDENT_BLOCKS` must also request `CACHE`. +Declare `plugins::worker_selection::experimental::resident_blocks::RESIDENT_BLOCKS` together with `WorkerInputs::CACHE`, and import that module's `ResidentBlocks` trait to read `resident_blocks()` from a candidate, or `ResidentBlocksByRow` to read `input.resident_blocks(row)` in a picker. It returns how many blocks the router's indexer holds for that worker rank, including cached blocks no running request uses, as of the last event the indexer applied. Without the declaration, with a remote indexer, or for a rank the indexer has held no blocks for or has removed, it returns `None`; a cleared rank returns `0`. The `experimental` module has no compatibility guarantee. ## Pick a Starting Point diff --git a/lib/kv-router/benches/worker_selection.rs b/lib/kv-router/benches/worker_selection.rs index 45616fe57822..ea6a6b76aeb7 100644 --- a/lib/kv-router/benches/worker_selection.rs +++ b/lib/kv-router/benches/worker_selection.rs @@ -14,6 +14,9 @@ use dynamo_kv_router::indexer::{ KvIndexerInterface, ThreadPoolIndexer, concurrent_radix_tree_compressed::ConcurrentRadixTreeCompressed, }; +use dynamo_kv_router::plugins::worker_selection::experimental::resident_blocks::{ + RESIDENT_BLOCKS, ResidentBlocks, +}; use dynamo_kv_router::protocols::{ ExternalSequenceBlockHash, KvCacheEvent, KvCacheEventData, KvCacheRemoveData, KvCacheStoreData, KvCacheStoredBlockData, LocalBlockHash, RouterEvent, RoutingConstraints, WorkerConfigLike, @@ -137,7 +140,7 @@ struct ResidentBlocksScorer; impl WorkerScorer for ResidentBlocksScorer { fn required_worker_inputs(&self) -> WorkerInputs { - WorkerInputs::CACHE | WorkerInputs::LOAD | WorkerInputs::RESIDENT_BLOCKS + WorkerInputs::CACHE | WorkerInputs::LOAD | RESIDENT_BLOCKS } fn score( diff --git a/lib/kv-router/src/plugins/worker_selection.rs b/lib/kv-router/src/plugins/worker_selection.rs index e62afbd9d134..97cbd18bec3b 100644 --- a/lib/kv-router/src/plugins/worker_selection.rs +++ b/lib/kv-router/src/plugins/worker_selection.rs @@ -5,6 +5,7 @@ mod config; mod context; +pub mod experimental; mod inputs; mod runtime; diff --git a/lib/kv-router/src/plugins/worker_selection/experimental.rs b/lib/kv-router/src/plugins/worker_selection/experimental.rs new file mode 100644 index 000000000000..ac069bbc5851 --- /dev/null +++ b/lib/kv-router/src/plugins/worker_selection/experimental.rs @@ -0,0 +1,9 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Unstable worker-selection extensions. +//! +//! Items in this module have no compatibility guarantee. They may change or be removed in any +//! release. Each extension lives in its own submodule. + +pub mod resident_blocks; diff --git a/lib/kv-router/src/plugins/worker_selection/experimental/resident_blocks.rs b/lib/kv-router/src/plugins/worker_selection/experimental/resident_blocks.rs new file mode 100644 index 000000000000..d8aa0d5bf5c5 --- /dev/null +++ b/lib/kv-router/src/plugins/worker_selection/experimental/resident_blocks.rs @@ -0,0 +1,95 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Per-rank resident block counts for worker-selection policies. +//! +//! A component that declares [`RESIDENT_BLOCKS`] next to [`WorkerInputs::CACHE`] can read how +//! many blocks the router's indexer holds for each worker rank, including cached blocks that no +//! running request uses, and charge a placement for the cache it would evict. Use +//! [`WorkerInputs::LOAD`] for work in flight. +//! +//! The indexer's event threads keep the counts live, so a count is current as of the last event +//! its thread applied, the same freshness as the `CACHE` overlap scores. It is unweighted and not +//! clamped to the rank's capacity. In KV-event mode it is the device-tier blocks the indexer +//! tracks, and every router replica sees the same counts. In approximate modes it is the blocks +//! this replica routed to the rank within the TTL, or the distinct full blocks its LRU policy +//! retains. +//! +//! A read is `None` when the component did not declare [`RESIDENT_BLOCKS`], when the host has no +//! local indexer, when the indexer has held no blocks for the rank while counting or has removed +//! it, and after an indexer thread stops. A rank cleared of blocks reads `Some(0)`. +//! +//! [`RESIDENT_BLOCKS`] requires [`WorkerInputs::CACHE`], because the counts come from the indexer +//! that `CACHE` starts. The host checks this for every policy it builds for a routing partition. +//! +//! ``` +//! use dynamo_kv_router::plugins::worker_selection::experimental::resident_blocks::{ +//! RESIDENT_BLOCKS, ResidentBlocks, +//! }; +//! use dynamo_kv_router::plugins::worker_selection::{ +//! WorkerCandidates, WorkerInputs, WorkerScorer, WorkerSelectionContext, +//! WorkerSelectionPolicyError, +//! }; +//! +//! /// Charges one unit per block a placement could evict. +//! struct EvictionPressure; +//! +//! impl WorkerScorer for EvictionPressure { +//! fn required_worker_inputs(&self) -> WorkerInputs { +//! WorkerInputs::CACHE | RESIDENT_BLOCKS +//! } +//! +//! fn score( +//! &mut self, +//! _context: &WorkerSelectionContext<'_>, +//! candidates: WorkerCandidates<'_>, +//! costs: &mut [f64], +//! ) -> Result<(), WorkerSelectionPolicyError> { +//! for (candidate, cost) in candidates.iter().zip(costs) { +//! // None means the host has no count for this rank; treat it as neutral. +//! *cost = candidate.resident_blocks().unwrap_or(0) as f64; +//! } +//! Ok(()) +//! } +//! } +//! ``` + +use crate::plugins::worker_selection::{WorkerCandidate, WorkerInputView, WorkerInputs}; + +/// Request per-rank counts of the blocks the router's indexer tracks. Requires +/// [`WorkerInputs::CACHE`]. +pub const RESIDENT_BLOCKS: WorkerInputs = WorkerInputs::RESIDENT_BLOCKS; + +mod sealed { + pub trait Sealed {} +} + +/// Reads a candidate's resident block count in a filter or scorer. +pub trait ResidentBlocks: sealed::Sealed { + /// Blocks the router's indexer holds for this worker rank, or `None` when this component did + /// not declare [`RESIDENT_BLOCKS`] or the host has no count for the rank (see the module docs). + fn resident_blocks(self) -> Option; +} + +impl sealed::Sealed for WorkerCandidate<'_> {} + +impl ResidentBlocks for WorkerCandidate<'_> { + fn resident_blocks(self) -> Option { + self.resident_block_count() + } +} + +/// Reads resident block counts by candidate row in a picker. +pub trait ResidentBlocksByRow: sealed::Sealed { + /// [`ResidentBlocks::resident_blocks`] for the candidate at `row`, or `None` for an + /// out-of-range row. + fn resident_blocks(self, row: usize) -> Option; +} + +impl sealed::Sealed for WorkerInputView<'_> {} + +impl ResidentBlocksByRow for WorkerInputView<'_> { + fn resident_blocks(self, row: usize) -> Option { + self.resident_block_count(row) + } +} diff --git a/lib/kv-router/src/plugins/worker_selection/inputs.rs b/lib/kv-router/src/plugins/worker_selection/inputs.rs index ab4680fb0dab..b9d2f5ae36df 100644 --- a/lib/kv-router/src/plugins/worker_selection/inputs.rs +++ b/lib/kv-router/src/plugins/worker_selection/inputs.rs @@ -69,12 +69,9 @@ impl<'a> WorkerCandidate<'a> { } } - /// Blocks the router's indexer holds for this worker rank as of the last event it applied, only - /// when this component declared RESIDENT_BLOCKS: device blocks from KV events, or the router's - /// prediction in approximate modes. Unweighted and not clamped to capacity. None means access - /// was not requested, the host has no local indexer, or the indexer has held no blocks for the - /// rank while counting or has removed it. - pub fn resident_blocks(self) -> Option { + /// Backs `experimental::resident_blocks::ResidentBlocks`, which documents the contract: the + /// count only when this component declared RESIDENT_BLOCKS and the host has one for the rank. + pub(crate) fn resident_block_count(self) -> Option { if !self.inputs.contains(WorkerInputs::RESIDENT_BLOCKS) { return None; } @@ -155,9 +152,9 @@ impl WorkerInputs { pub const LOAD: Self = Self(1 << 1); /// Request preferred-taint routing metadata. pub const PREFERRED_TAINT: Self = Self(1 << 2); - /// Request per-rank counts of blocks the router's indexer tracks (routing predictions in - /// approximate modes). Requires CACHE. - pub const RESIDENT_BLOCKS: Self = Self(1 << 3); + /// Request per-rank resident block counts. Public only as + /// `experimental::resident_blocks::RESIDENT_BLOCKS`; requires CACHE. + pub(crate) const RESIDENT_BLOCKS: Self = Self(1 << 3); /// Request host-owned active-request counts. pub const OCCUPANCY: Self = Self(1 << 5); #[cfg(any(test, feature = "bench"))] @@ -391,9 +388,9 @@ impl<'a> WorkerInputView<'a> { self.load } - /// Return [`WorkerCandidate::resident_blocks`] for candidate `row` when the picker requested - /// [`WorkerInputs::RESIDENT_BLOCKS`], or None for an out-of-range row. - pub fn resident_blocks(self, row: usize) -> Option { + /// Backs `experimental::resident_blocks::ResidentBlocksByRow`: the count for candidate `row` + /// when the picker declared RESIDENT_BLOCKS, or None for an out-of-range row. + pub(crate) fn resident_block_count(self, row: usize) -> Option { let counts = self.resident_blocks?; self.candidates .get(row) diff --git a/lib/kv-router/src/scheduling/selector/policy.rs b/lib/kv-router/src/scheduling/selector/policy.rs index cea2135cc2a5..328893683612 100644 --- a/lib/kv-router/src/scheduling/selector/policy.rs +++ b/lib/kv-router/src/scheduling/selector/policy.rs @@ -509,6 +509,9 @@ impl ResidentBlockCountsWriter { #[cfg(test)] mod tests { use crate::plugins::worker_selection::WorkerInputView; + use crate::plugins::worker_selection::experimental::resident_blocks::{ + ResidentBlocks, ResidentBlocksByRow, + }; use crate::protocols::WorkerWithDpRank; use crate::scheduling::SessionContext; use std::{ diff --git a/lib/kv-router/src/services/selection/core/tests.rs b/lib/kv-router/src/services/selection/core/tests.rs index 781e58fffe4c..4825e7ce6e11 100644 --- a/lib/kv-router/src/services/selection/core/tests.rs +++ b/lib/kv-router/src/services/selection/core/tests.rs @@ -8,6 +8,7 @@ use super::reservations::{ Reservation, ReservationClaim, ReservationIndexObserver, sweep_reservation_index, }; use super::*; +use crate::plugins::worker_selection::experimental::resident_blocks::ResidentBlocksByRow; use crate::protocols::{ActiveSequenceEvent, ActiveSequenceEventData}; use crate::protocols::{RoutingConstraints, StorageTier}; use crate::services::common::replica_sync::HostReplicaChannels; diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index 0f35d34669ad..b92c0d82f598 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -11,6 +11,7 @@ use std::{ use anyhow::Result; #[cfg(test)] use dynamo_kv_router::WorkerSelectionPolicy; +use dynamo_kv_router::plugins::worker_selection::experimental::resident_blocks::RESIDENT_BLOCKS; use dynamo_kv_router::{ DEFAULT_ROUTING_GROUP, KvSchedulerError, PrefillLoadEstimator, RoutingPartitionRef, SessionPrefixIndexer, SharedKvCache, TrackingHashAlgorithm, TrackingHashContext, @@ -156,9 +157,10 @@ impl SelectionPolicySource { !config.enable_session_prefix_index, "enable_session_prefix_index requires the worker-selection policy to declare WorkerInputs::CACHE" ); - // Resident block counts come from the indexer that CACHE starts. + // Resident block counts come from the indexer that CACHE starts. The input is + // experimental; this check is the frontend's only use of that namespace. anyhow::ensure!( - !prepared.inputs().contains(WorkerInputs::RESIDENT_BLOCKS), + !prepared.inputs().contains(RESIDENT_BLOCKS), "RESIDENT_BLOCKS requires the worker-selection policy to declare WorkerInputs::CACHE" ); } @@ -2479,19 +2481,15 @@ mod tests { #[tokio::test] async fn resident_blocks_require_declared_cache_input() { let config = KvRouterConfig::default(); - let error = picker_policy(|| Box::new(DeclaredInputsPicker(WorkerInputs::RESIDENT_BLOCKS))) + let error = picker_policy(|| Box::new(DeclaredInputsPicker(RESIDENT_BLOCKS))) .prepare(&config, WorkerType::Aggregated, "decode", Some("model")) .err() .expect("without CACHE there is no indexer to count blocks"); assert!(error.to_string().contains("RESIDENT_BLOCKS")); assert!(error.to_string().contains("WorkerInputs::CACHE")); - picker_policy(|| { - Box::new(DeclaredInputsPicker( - WorkerInputs::RESIDENT_BLOCKS | WorkerInputs::CACHE, - )) - }) - .prepare(&config, WorkerType::Aggregated, "decode", Some("model")) - .expect("CACHE starts the indexer the counts come from"); + picker_policy(|| Box::new(DeclaredInputsPicker(RESIDENT_BLOCKS | WorkerInputs::CACHE))) + .prepare(&config, WorkerType::Aggregated, "decode", Some("model")) + .expect("CACHE starts the indexer the counts come from"); } #[tokio::test] From c96a8a1bfe47674a6b32bab17dacaec6589b15f4 Mon Sep 17 00:00:00 2001 From: Akira Lonske Date: Tue, 3 Nov 2026 15:37:56 -0500 Subject: [PATCH 12/12] test(kv-router): pin resident block counts on the approximate-LRU event path The LRU-emitted stores in the compressed tree's lane loop updated the counters with no test behind them: deleting that call failed nothing. The new test tracks before any block lands, materializes two blocks through a lease, and expects Some(2). Also state in the contract that only filters, scorers, and pickers receive counts; a request classifier that declares the input reads None. Signed-off-by: Akira Lonske --- .../router/custom-worker-selection.mdx | 2 +- lib/kv-router/src/indexer/approximate_lru.rs | 54 +++++++++++++++++++ .../experimental/resident_blocks.rs | 4 +- 3 files changed, 58 insertions(+), 2 deletions(-) diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx index 483e9b906622..507fc9376676 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx @@ -658,7 +658,7 @@ impl WorkerScorer for EvictionPressure { } ``` -The count is unweighted, not clamped to `total_kv_blocks`, and measures occupancy rather than reusable cache: blocks that prefix matching can no longer reach, such as descendants of an evicted block, count until their own remove events arrive, and a lost remove event inflates the count until the rank is removed. A count is current as of the last event its indexer thread applied, the same freshness as the `CACHE` overlap scores, so it can trail queued events and change between two reads in one selection. A rank reads `None` until its indexer thread has reached the policy's request and holds a block for it, `0` after a clear, and `None` again once the rank or its worker is removed. A component also reads `None` when it did not declare `RESIDENT_BLOCKS`, with a remote indexer, or after an indexer thread stops, so a worker whose KV events never arrive never looks empty. A policy that declares `RESIDENT_BLOCKS` must also declare `CACHE`, because the counts come from the indexer that `CACHE` starts. The router checks this for every policy it builds for a routing partition: the frontend fails to start, and the standalone selection service marks the partition's workers unschedulable with the reason. +The count is unweighted, not clamped to `total_kv_blocks`, and measures occupancy rather than reusable cache: blocks that prefix matching can no longer reach, such as descendants of an evicted block, count until their own remove events arrive, and a lost remove event inflates the count until the rank is removed. A count is current as of the last event its indexer thread applied, the same freshness as the `CACHE` overlap scores, so it can trail queued events and change between two reads in one selection. A rank reads `None` until its indexer thread has reached the policy's request and holds a block for it, `0` after a clear, and `None` again once the rank or its worker is removed. A component also reads `None` when it did not declare `RESIDENT_BLOCKS`, with a remote indexer, or after an indexer thread stops, so a worker whose KV events never arrive never looks empty. Only filters, scorers, and pickers receive counts: a request classifier that declares `RESIDENT_BLOCKS` reads `None`. A policy that declares `RESIDENT_BLOCKS` must also declare `CACHE`, because the counts come from the indexer that `CACHE` starts. The router checks this for every policy it builds for a routing partition: the frontend fails to start, and the standalone selection service marks the partition's workers unschedulable with the reason. What the count contains, and whether router replicas agree on it, depends on the routing mode: diff --git a/lib/kv-router/src/indexer/approximate_lru.rs b/lib/kv-router/src/indexer/approximate_lru.rs index 37668e25e0a8..0fc9cf5d39f2 100644 --- a/lib/kv-router/src/indexer/approximate_lru.rs +++ b/lib/kv-router/src/indexer/approximate_lru.rs @@ -1852,4 +1852,58 @@ mod tests { assert_eq!(stats.leases, 0); assert_eq!(stats.inactive_blocks, 2); } + + #[tokio::test] + async fn lru_output_blocks_update_resident_block_counts() { + let indexer = ThreadPoolIndexer::new_with_metrics_and_approximate_retention( + ConcurrentRadixTreeCompressed::new(), + 2, + 4, + None, + Some(ApproximateRetentionConfig::Lru { + fallback_ttl: PruneConfig { + ttl: Duration::from_secs(120), + }, + }), + ); + let worker = worker(); + // Track before any block lands, so the count comes from the LRU-emitted + // events rather than the backfill. + let counts = indexer.track_resident_blocks(); + assert_eq!(counts.load().get(worker), None); + + let local_hashes = vec![LocalBlockHash(11), LocalBlockHash(22)]; + let sequence_hashes = compute_seq_hash_for_block(&local_hashes); + indexer + .set_approximate_lru_capacity(worker, 1, Some(4)) + .await + .unwrap(); + let lease = indexer + .begin_approximate_lru_request(worker, 1, attempt(1)) + .unwrap(); + lease + .acquire( + vec![ApproximateLruBlock { + local_hash: local_hashes[0], + sequence_hash: sequence_hashes[0], + }], + 0, + ) + .await + .unwrap(); + lease + .materialize( + Some(sequence_hashes[0]), + vec![ApproximateLruBlock { + local_hash: local_hashes[1], + sequence_hash: sequence_hashes[1], + }], + 1, + 0, + ) + .unwrap(); + lease.finish().await.unwrap(); + + assert_eq!(counts.load().get(worker), Some(2)); + } } diff --git a/lib/kv-router/src/plugins/worker_selection/experimental/resident_blocks.rs b/lib/kv-router/src/plugins/worker_selection/experimental/resident_blocks.rs index d8aa0d5bf5c5..920f5cfaa19f 100644 --- a/lib/kv-router/src/plugins/worker_selection/experimental/resident_blocks.rs +++ b/lib/kv-router/src/plugins/worker_selection/experimental/resident_blocks.rs @@ -17,7 +17,9 @@ //! //! A read is `None` when the component did not declare [`RESIDENT_BLOCKS`], when the host has no //! local indexer, when the indexer has held no blocks for the rank while counting or has removed -//! it, and after an indexer thread stops. A rank cleared of blocks reads `Some(0)`. +//! it, and after an indexer thread stops. A rank cleared of blocks reads `Some(0)`. Only filters, +//! scorers, and pickers receive counts: a request classifier that declares [`RESIDENT_BLOCKS`] +//! reads `None`. //! //! [`RESIDENT_BLOCKS`] requires [`WorkerInputs::CACHE`], because the counts come from the indexer //! that `CACHE` starts. The host checks this for every policy it builds for a routing partition.