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 ddad76e78468..bb41fba913d8 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` or a scorer requests `KV_TRANSFER`. `serve_indexer` and `enable_session_prefix_index` require `CACHE`. +Routing cache resources start only if a component requests `CACHE` or a scorer requests `KV_TRANSFER`. `serve_indexer`, `enable_session_prefix_index`, and the experimental `RESIDENT_BLOCKS` input require `CACHE`. ### Request Context @@ -628,6 +629,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( + &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. 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: + +- **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 e0ca90675092..e9af13e9b808 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 `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 | Crate | Use it for | diff --git a/lib/bench/kv_router/INDEXER_BENCH.md b/lib/bench/kv_router/INDEXER_BENCH.md index bf095e7c12ec..4f6e7084a8d1 100644 --- a/lib/bench/kv_router/INDEXER_BENCH.md +++ b/lib/bench/kv_router/INDEXER_BENCH.md @@ -158,6 +158,18 @@ 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. 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 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 55420b2c57cf..a4c775ba762f 100644 --- a/lib/bench/kv_router/mooncake_bench.rs +++ b/lib/bench/kv_router/mooncake_bench.rs @@ -28,6 +28,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; @@ -163,6 +164,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)] @@ -316,6 +327,11 @@ async fn run_open_loop_for_config( let metrics = 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 => { @@ -343,7 +359,42 @@ async fn run_open_loop_for_config( args.lower_tier.kvcr_blocks() > 0, ) }); - run_backend(config.short_name(), indexer, lower_tier, 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(); + 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()); + } + }); + (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, lower_tier, trial, open_config).await; + stop.store(true, Ordering::Relaxed); + if let Some((handle, _)) = reader { + handle + .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 003f78fc89bb..33252dfa0485 100644 --- a/lib/kv-router/benches/worker_selection.rs +++ b/lib/kv-router/benches/worker_selection.rs @@ -6,15 +6,25 @@ //! 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::plugins::worker_selection::experimental::resident_blocks::{ + RESIDENT_BLOCKS, ResidentBlocks, +}; use dynamo_kv_router::plugins::worker_selection::{ WorkerCandidate, WorkerCandidates, WorkerFilter, WorkerInputView, WorkerInputs, WorkerPicker, WorkerScorer, WorkerSelectionContext, WorkerSelectionPolicy, }; use dynamo_kv_router::protocols::{ - RoutingConstraints, WorkerConfigLike, WorkerId, WorkerWithDpRank, + ExternalSequenceBlockHash, KvCacheEvent, KvCacheEventData, KvCacheRemoveData, KvCacheStoreData, + KvCacheStoredBlockData, LocalBlockHash, RouterEvent, RoutingConstraints, WorkerConfigLike, + WorkerId, WorkerWithDpRank, compute_seq_hash_for_block, }; use dynamo_kv_router::scheduling::{OverlapSignals, ScheduleMode}; use dynamo_kv_router::{ @@ -142,6 +152,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 | RESIDENT_BLOCKS + } + + fn score( + &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 { @@ -409,6 +455,158 @@ 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, writer) = 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, 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, 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 = 16; + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let events: Vec<_> = (0..WORKERS) + .flat_map(|worker| { + let local: Vec<_> = (0..BLOCKS) + .map(|i| LocalBlockHash(worker * BLOCKS + i)) + .collect(); + let hashes: Vec<_> = compute_seq_hash_for_block(&local) + .into_iter() + .map(ExternalSequenceBlockHash) + .collect(); + 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(events.len() as u64)); + 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 event in &events { + indexer.apply_event(event.clone()).await; + } + indexer.flush().await; + }) + }) + }); + indexer.shutdown(); + } + group.finish(); +} + fn default_policy_wrapper(c: &mut Criterion) { let config = KvRouterConfig { router_temperature: 0.0, @@ -507,6 +705,8 @@ criterion_group!( worker_selection, 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/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/indexer/concurrent_radix_tree_compressed/lane_lookup.rs b/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/lane_lookup.rs index 0427890132f6..c5157dfca6f4 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 @@ -343,7 +343,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 89efd9ac4fe5..591fab83af14 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: &LaneLookup, + worker: WorkerWithDpRank, +) { + if let Some(tracker) = resident_blocks { + tracker.update(worker, lookup.block_count(worker)); + } +} // ============================================================================ // SyncIndexer implementation for ConcurrentRadixTreeCompressed @@ -20,6 +33,7 @@ impl SyncIndexer for ConcurrentRadixTreeCompressed { let mut lookup = LaneLookup::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::DumpEvents(_sender) => { let _ = _sender.send(Ok(Vec::new())); @@ -159,6 +189,10 @@ impl SyncIndexer for ConcurrentRadixTreeCompressed { self.drain_pending_unlinks(&mut lookup, None); let _ = sender.send(()); } + WorkerTask::TrackResidentBlocks(counts) => { + resident_blocks = + Some(ResidentBlockTracker::new(counts, lookup.block_counts())); + } WorkerTask::Terminate => { break; } diff --git a/lib/kv-router/src/indexer/kv_indexer.rs b/lib/kv-router/src/indexer/kv_indexer.rs index ecf97d1f25fa..5d76d3aaf688 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, ResidentBlockCountsHandle, + RoutingDecisionRequest, panic_payload_message, }; use crate::indexer::pruning::{BlockEntry, PruneConfig, WorkerPruneManager}; use crate::protocols::*; @@ -336,6 +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, + /// Starts live per-rank block counts in the actor. + #[cfg_attr(not(feature = "standalone-indexer"), allow(dead_code))] + 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. @@ -430,6 +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 (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); @@ -460,6 +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 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); @@ -498,6 +508,10 @@ impl KvIndexer { } } + Some(counts) = track_resident_blocks_rx.recv() => { + trie.track_resident_blocks(counts); + } + Some(mutation) = mutation_rx.recv() => { apply_mutation(&mut trie, mutation, &counters, &prune_manager); } @@ -656,6 +670,8 @@ impl KvIndexer { remove_worker_tx, remove_worker_dp_rank_tx, get_workers_tx, + track_resident_blocks_tx, + resident_blocks: Arc::default(), dump_tx, flush_tx, routing_tx, @@ -774,6 +790,19 @@ impl KvIndexer { pub fn get_workers_sender(&self) -> mpsc::Sender { self.get_workers_tx.clone() } + + /// 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) 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() + } } #[async_trait] diff --git a/lib/kv-router/src/indexer/lower_tier.rs b/lib/kv-router/src/indexer/lower_tier.rs index 29b067a02da5..43d48efc0187 100644 --- a/lib/kv-router/src/indexer/lower_tier.rs +++ b/lib/kv-router/src/indexer/lower_tier.rs @@ -1430,6 +1430,8 @@ impl SyncIndexer for LowerTierIndexer { 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/mod.rs b/lib/kv-router/src/indexer/mod.rs index 5a9e98205605..c34b10b4afc7 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,8 @@ pub use lower_tier_indexers::*; pub use metrics::*; #[cfg(feature = "bench")] pub use observation::*; +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 fd360fc86047..9d64218b8a7c 100644 --- a/lib/kv-router/src/indexer/positional.rs +++ b/lib/kv-router/src/indexer/positional.rs @@ -322,6 +322,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 7cefaeefd086..8a907006b05f 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}; +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( @@ -750,10 +760,22 @@ 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 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; }; @@ -774,6 +796,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 { @@ -874,6 +903,16 @@ impl RadixTree { self.lookup.values().map(FxHashMap::len).sum() } + /// 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)] 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..194438f9e5c8 --- /dev/null +++ b/lib/kv-router/src/indexer/resident_blocks.rs @@ -0,0 +1,213 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Live per-rank resident block counts. +//! +//! 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::sync::atomic::{AtomicU64, Ordering}; + +use arc_swap::{ArcSwap, Guard}; +use crossbeam_utils::CachePadded; +use rustc_hash::FxHashMap; + +use crate::protocols::WorkerWithDpRank; + +/// Padded so counters written by different event threads never share a cache line. +type Counter = CachePadded; + +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 { + counters: FxHashMap>, +} + +impl ResidentBlockCounts { + /// 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) + .map(|counter| counter.load(Ordering::Relaxed)) + } +} + +/// 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 })), + } + } + + /// 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)]) { + if added.is_empty() { + return; + } + 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)]) { + if removed.is_empty() { + return; + } + 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 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 ResidentBlockTracker { + /// 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 + .iter() + .map(|(worker, counter)| (*worker, counter.clone())) + .collect(); + counts.publish(&added); + Self { counts, owned } + } + + /// 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.filter(|&blocks| blocks > 0) { + 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 + }); + 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), + ); + // 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)], + [Some(5), Some(0), Some(7), 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 dfe0dd0db89c..5d6c4876a1c4 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, ResidentBlockCountsHandle, + ShardSizeSnapshot, SyncIndexer, WorkerLookupStats, WorkerTask, panic_payload_message, }; #[cfg(feature = "bench")] use super::{ @@ -83,6 +83,12 @@ 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, + + /// 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, @@ -304,17 +310,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(|| { @@ -366,6 +376,8 @@ impl ThreadPoolIndexer { num_workers, 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, @@ -628,27 +640,55 @@ 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) } + /// [`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) 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 { self.worker_lookup_stats() .await @@ -1604,6 +1644,51 @@ mod tests { assert_score(&indexer, &[20], rank1, 1).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_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(); + // 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_eq!(read.expect("stats read hung").worker_count(), 0); + } + #[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 c4ee982a0eed..b7376ed5ee17 100644 --- a/lib/kv-router/src/indexer/types.rs +++ b/lib/kv-router/src/indexer/types.rs @@ -654,6 +654,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.rs b/lib/kv-router/src/plugins/worker_selection.rs index 0fb894db97c6..d24731e177dd 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; pub(crate) 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..08edb3409206 --- /dev/null +++ b/lib/kv-router/src/plugins/worker_selection/experimental/resident_blocks.rs @@ -0,0 +1,97 @@ +// 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)`. 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. +//! +//! ``` +//! 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( +//! &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 b22e9926ee1e..ad6ddf133cb6 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, WorkerConfigLike, WorkerId, WorkerWithDpRank}; use crate::sequences::WorkerLoadProjection; use rustc_hash::FxHashMap; @@ -103,6 +104,17 @@ impl<'a> WorkerCandidate<'a> { None } } + + /// 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; + } + self.cache_snapshot + .resident_blocks + .and_then(|counts| counts.get(self.data.worker)) + } } /// Borrowed candidate batch restricted to one scorer's declared inputs. @@ -190,6 +202,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 resident block counts. Public only as + /// `experimental::resident_blocks::RESIDENT_BLOCKS`; requires CACHE. + pub(crate) const RESIDENT_BLOCKS: Self = Self(1 << 3); /// Request each row's active request count, read through /// [`WorkerInputView::active_requests`]. Picker only: filters and scorers that declare it /// receive nothing. @@ -231,6 +246,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>, } /// Rows and request-wide cache facts owned by a reader outside worker selection. @@ -281,6 +298,7 @@ impl WorkerInputRows { let cache_snapshot = CacheSnapshot { shared_hits: self.shared_hits.as_ref(), has_tier_matches: self.has_tier_matches, + resident_blocks: None, }; let info = |row: usize| WorkerInfo::new(self.rows[row].worker, self.capacities[row]); read(WorkerCandidates::new( @@ -482,6 +500,8 @@ pub struct WorkerInputView<'a> { rows: PickerRows<'a>, cache: Option>, load: Option<&'a [WorkerLoadInput]>, + // Set only on the KV path, for a picker that declared RESIDENT_BLOCKS. + resident_blocks: Option<&'a ResidentBlockCounts>, } /// Where a picker's rows come from. @@ -697,6 +717,7 @@ impl<'a> WorkerInputView<'a> { load: Option<&'a [WorkerLoadInput]>, active_requests: Option<&'a dyn Fn(usize) -> u64>, info: WorkerInfoLookup<'a>, + resident_blocks: Option<&'a ResidentBlockCounts>, ) -> Self { Self { rows: PickerRows::Materialized { @@ -706,6 +727,7 @@ impl<'a> WorkerInputView<'a> { }, cache, load, + resident_blocks, } } @@ -724,6 +746,7 @@ impl<'a> WorkerInputView<'a> { }, cache: None, load: None, + resident_blocks: None, } } @@ -816,4 +839,13 @@ impl<'a> WorkerInputView<'a> { } } } + + /// 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) + .and_then(|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 5bccaf8d681b..d38d44068cdc 100644 --- a/lib/kv-router/src/scheduling/selector/mod.rs +++ b/lib/kv-router/src/scheduling/selector/mod.rs @@ -47,6 +47,7 @@ use super::filter::{RoutingEligibility, WorkerEligibilityError}; use super::types::{ KvSchedulerError, SchedulingRequest, SchedulingResponse, WorkerSelectionPolicyError, }; +use crate::indexer::ResidentBlockCounts; use crate::protocols::{ SelectionLogFields, SelectionTrace, SelectionTraceRow, SelectionTraceSpec, WorkerConfigLike, WorkerId, WorkerSelectionResult, WorkerWithDpRank, @@ -344,6 +345,7 @@ fn cache_snapshot(request: &SchedulingRequest) -> CacheSnapshot<'_> { 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, } } @@ -626,6 +628,7 @@ fn select_worker_with_policy( worker_type: &'static str, state: WorkerSelectionPolicyStateRef<'_>, input: WorkerSelectionInput<'_, C>, + resident_blocks: Option<&ResidentBlockCounts>, trace: Option, prepared: Option, ) -> Result<(WorkerSelectionResult, Option), KvSchedulerError> { @@ -646,6 +649,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) => { @@ -714,6 +718,10 @@ fn select_worker_with_policy( .contains(WorkerInputs::ACTIVE_REQUESTS) .then_some(&active_requests), &info, + 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 2c8955246a9d..bb68253d45cc 100644 --- a/lib/kv-router/src/scheduling/selector/policy.rs +++ b/lib/kv-router/src/scheduling/selector/policy.rs @@ -14,6 +14,7 @@ use super::{ select_worker_with_policy, worker_info, }; +use crate::indexer::ResidentBlockCountsHandle; use crate::kv_hints::{KvTransferSourceSummary, KvTransferTarget}; use crate::protocols::{ SelectionTrace, SelectionTraceSpec, WorkerConfigLike, WorkerId, WorkerSelectionResult, @@ -79,6 +80,7 @@ pub struct WorkerSelectionPolicy { worker_label: &'static str, state: WorkerSelectionPolicyState, exclusive_affinity: bool, + resident_blocks: Option, runtime: Option>, } @@ -147,6 +149,7 @@ impl WorkerSelectionPolicy { Self { worker_label, exclusive_affinity: false, + resident_blocks: None, runtime: None, state: WorkerSelectionPolicyState::Composed(RefCell::new(ComposedPolicyState { filters, @@ -174,6 +177,28 @@ 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 + } + + /// 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, ResidentBlockCountsWriter) { + let counts = ResidentBlockCountsHandle::fixed(counts); + ( + self.with_resident_block_counts(counts.clone()), + ResidentBlockCountsWriter(counts), + ) + } + /// Opt in to runtime parameter updates with a pure replacement builder. /// The initial parameters define the patchable keys and reset values. Replacements must /// require no new worker-input groups and preserve affinity behavior; queue and tracking state are retained. @@ -221,6 +246,8 @@ impl WorkerSelectionPolicy { } let response = state.config(); replacement.runtime = Some(Box::new(state)); + // The host attached the counts to this instance; the replacement keeps reading them. + replacement.resident_blocks = self.resident_blocks.take(); *self = replacement; Ok(response) } @@ -298,6 +325,7 @@ impl WorkerSelectionPolicy { Self { worker_label, exclusive_affinity: false, + resident_blocks: None, runtime: None, state: WorkerSelectionPolicyState::Reference(Box::new(kv_router_config), picker), } @@ -788,13 +816,41 @@ impl WorkerSelectionPolicy { WorkerSelectionPolicyStateRef::Composed(state) } }; - select_worker_with_policy(self.worker_label, state, input, trace, prepared) + // Held here so the per-selection snapshot borrows it and stays plain data. + let resident_blocks = self + .resident_blocks + .as_ref() + .map(ResidentBlockCountsHandle::load); + select_worker_with_policy( + self.worker_label, + state, + input, + resident_blocks.as_deref().map(|counts| &**counts), + trace, + prepared, + ) + } +} + +/// 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; + use crate::plugins::worker_selection::experimental::resident_blocks::{ + ResidentBlocks, ResidentBlocksByRow, + }; use crate::protocols::WorkerWithDpRank; use crate::scheduling::SessionContext; use std::{ @@ -1587,6 +1643,233 @@ mod tests { assert_eq!(workers[&0].taint_reads.get(), 1); } + #[test] + fn resident_blocks_reach_only_declaring_components() { + use crate::indexer::ResidentBlockCountsHandle; + + // 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 } + } + + 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( + &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( + &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(()) + } + } + + /// Picks the fewest resident blocks; an undeclared picker reads None and takes row 0. + struct ResidencyPicker(WorkerInputs); + + impl WorkerPicker for ResidencyPicker { + fn required_worker_inputs(&self) -> WorkerInputs { + self.0 + } + + fn pick( + &self, + _context: &WorkerSelectionContext<'_>, + input: WorkerInputView<'_>, + ) -> 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() { + 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_or(0)) + .unwrap()) + } + } + + let workers = HashMap::from([ + (3, TaintedWorkerConfig::default()), + (8, TaintedWorkerConfig::default()), + ]); + let request = base_request(16); + let counts = ResidentBlockCountsHandle::fixed([ + (WorkerWithDpRank::new(3, 0), 12), + (WorkerWithDpRank::new(8, 0), 0), + ]); + let probes = || -> Vec> { + vec![ + Box::new(ResidencyProbe(WorkerInputs::RESIDENT_BLOCKS)), + Box::new(ResidencyProbe(WorkerInputs::NONE)), + ] + }; + 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(); + if picker_inputs.contains(WorkerInputs::RESIDENT_BLOCKS) { + 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( + &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 with no tracked ranks 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 runtime_replacement_keeps_resident_block_counts() { + use std::sync::Arc; + + use crate::indexer::ResidentBlockCountsHandle; + use crate::plugins::worker_selection::RuntimePolicyUpdate; + + struct ExpectCount(u64); + + impl WorkerPicker for ExpectCount { + fn required_worker_inputs(&self) -> WorkerInputs { + WorkerInputs::RESIDENT_BLOCKS + } + + fn pick( + &self, + _context: &WorkerSelectionContext<'_>, + input: WorkerInputView<'_>, + ) -> Result { + assert_eq!(input.resident_blocks(0), Some(self.0)); + Ok(0) + } + } + + fn build() -> WorkerSelectionPolicy { + WorkerSelectionPolicy::new( + KvRouterConfig::default(), + "test", + Vec::new(), + Box::new(ExpectCount(5)), + ) + } + + let workers = HashMap::from([(0, TaintedWorkerConfig::default())]); + let request = base_request(16); + let counts = ResidentBlockCountsHandle::fixed([(WorkerWithDpRank::new(0, 0), 5)]); + let mut policy = build() + .with_resident_block_counts(counts) + .with_runtime_config(serde_json::Map::new(), Arc::new(|_| Ok(build()))); + // The replacement is built without the host's counts; the update carries them over. + policy + .update_runtime_config(RuntimePolicyUpdate::reset(0)) + .unwrap(); + 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 492a1cc5ce58..68b321401f17 100644 --- a/lib/kv-router/src/scheduling/selector/reference.rs +++ b/lib/kv-router/src/scheduling/selector/reference.rs @@ -669,6 +669,7 @@ impl WorkerSelector for DefaultWorkerSelector { input, None, None, + None, ) .map(|(result, _)| result) } diff --git a/lib/kv-router/src/services/indexer/backend.rs b/lib/kv-router/src/services/indexer/backend.rs index 8aabdc9cf821..232789ccf559 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, ResidentBlockCountsHandle, RoutingDecisionHashes, SyncIndexer, + TieredMatchDetails, TieredMatchProvider, record_unsupported_residency_event, }; use crate::protocols::{ DpRank, ExternalSequenceBlockHash, KvCacheEventData, LocalBlockHash, OverlapScores, @@ -466,9 +466,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?; @@ -868,6 +866,17 @@ impl Indexer { Self::Remote { .. } | Self::None => Ok(ApproximateLruStats::default()), } } + + /// 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))] + pub(crate) fn track_resident_blocks(&self) -> Option { + match self { + Self::Single { primary, .. } => Some(primary.track_resident_blocks()), + Self::Concurrent { primary, .. } => Some(primary.track_resident_blocks()), + Self::Remote { .. } | Self::None => None, + } + } } #[async_trait] @@ -1296,6 +1305,13 @@ mod tests { ); let flat = indexer.find_matches(sequence).await.unwrap(); assert_eq!(flat.scores.get(&worker).copied(), Some(2)); + let counts = indexer.track_resident_blocks().unwrap(); + flush(&indexer).await; + assert_eq!( + counts.load().get(worker), + Some(2), + "{num_threads} thread(s): routed prefix" + ); // The recorded prefix is dumped like any other primary state, so a // recovering peer inherits it. @@ -1613,6 +1629,105 @@ 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_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 [ + // 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; + assert_eq!( + [held(1, 0), held(1, 1), held(2, 0), held(3, 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. The clear takes the + // path the frontend admits live engine events through. + indexer + .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!( + [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) + ); + indexer.remove_worker(1).await; + flush(&indexer).await; + assert_eq!(held(1, 0), None, "{num_threads} indexer thread(s)"); + } + } + + #[tokio::test] + async fn resident_block_counts_withdraw_when_the_primary_stops() { + for num_threads in [1, 4] { + let indexer = create_indexer(16, num_threads); + 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); + flush(&indexer).await; + assert_eq!(counts.load().get(worker), Some(2)); + match &indexer { + Indexer::Single { primary, .. } => primary.shutdown(), + Indexer::Concurrent { primary, .. } => primary.shutdown(), + Indexer::Remote { .. } | Indexer::None => unreachable!(), + } + // A stopped primary's ranks read None, not counts nothing updates anymore. + wait_until(|| counts.load().get(worker).is_none()).await; + } + } } 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 0295b6ba2f23..24cae5dde231 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; @@ -992,6 +993,114 @@ async fn wire_explicit_pin_does_not_report_affinity_overlap_loss() { ); } +/// Picker that declares `inputs`, records the resident block count of row 0, and takes it. +struct ResidentBlocksRecorder { + inputs: crate::plugins::worker_selection::WorkerInputs, + observed: Arc>>>, +} + +impl crate::plugins::worker_selection::WorkerPicker for ResidentBlocksRecorder { + fn required_worker_inputs(&self) -> crate::plugins::worker_selection::WorkerInputs { + self.inputs + } + + fn pick( + &self, + _context: &crate::plugins::worker_selection::WorkerSelectionContext<'_>, + input: crate::plugins::worker_selection::WorkerInputView<'_>, + ) -> Result { + self.observed.lock().push(input.resident_blocks(0)); + Ok(0) + } +} + +fn resident_blocks_core( + inputs: crate::plugins::worker_selection::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| { + WorkerSelectionPolicy::new( + config.clone(), + worker_type.as_str(), + Vec::new(), + Box::new(ResidentBlocksRecorder { + inputs, + observed: Arc::clone(&factory_observed), + }), + ) + }); + let core = core_with( + test_config(true), + SelectionHost::default(), + Some(factory), + WorkerType::Aggregated, + None, + ); + (core, observed) +} + +#[tokio::test] +async fn declaring_policy_reads_live_resident_blocks() { + let (core, observed) = resident_blocks_core( + crate::plugins::worker_selection::WorkerInputs::CACHE + | crate::plugins::worker_selection::WorkerInputs::RESIDENT_BLOCKS, + ); + 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(); + entry.indexer.flush().await; + + core.select(select_request()).await.expect("select"); + assert_eq!(observed.lock().last(), Some(&Some(3))); +} + +#[tokio::test] +async fn only_a_declaring_policy_tracks_resident_blocks() { + use crate::plugins::worker_selection::WorkerInputs; + let indexer = crate::services::indexer::backend::create_indexer(16, 1); + let key = default_key(); + 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] +async fn partition_rejects_resident_blocks_without_cache() { + let (core, _) = + resident_blocks_core(crate::plugins::worker_selection::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 + ); + assert!( + core.indexer_registry.get_indexer(&default_key()).is_none(), + "a rejected partition opened an indexer" + ); +} + // Keep clock- and poll-order tests independent of the native indexer thread. struct ReadyIndexer; #[async_trait::async_trait] diff --git a/lib/kv-router/src/services/selection/core/workers.rs b/lib/kv-router/src/services/selection/core/workers.rs index 55b1e28f4714..cdc96f34b63a 100644 --- a/lib/kv-router/src/services/selection/core/workers.rs +++ b/lib/kv-router/src/services/selection/core/workers.rs @@ -9,6 +9,7 @@ use std::collections::HashSet; use super::hint::hint_capable_partition; use super::reservations::{ReservationIndexObserver, spawn_reservation_index_sweep}; use super::*; +use crate::indexer::ResidentBlockCountsHandle; use crate::sequences::topology::MAX_DATA_PARALLEL_RANKS_PER_WORKER; impl SelectionCore { @@ -281,6 +282,27 @@ 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 = + WorkerSelector::::required_worker_inputs(&selector); + 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 @@ -338,18 +360,10 @@ impl SelectionCore { block_size, )) }); - let selector = (self.worker_selection_policy_factory)( - &self.kv_router_config, - self.worker_type, - key.as_ref(), - ); - let scores_kv_transfer = - WorkerSelector::::required_worker_inputs(&selector) - .contains(WorkerInputs::KV_TRANSFER); - let profile = self - .kv_router_config - .policy_profile(Some(&key.model_name)) - .map_err(|error| SelectionError::BadRequest(error.to_string()))?; + let scores_kv_transfer = inputs.contains(WorkerInputs::KV_TRANSFER); + if let Some(counts) = track_resident_blocks(&key, inputs, &indexer) { + selector = selector.with_resident_block_counts(counts); + } let scheduler = LocalScheduler::new( slots, workers_rx, @@ -440,3 +454,23 @@ impl SelectionCore { } } } + +/// 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 track_resident_blocks( + key: &RoutingPartitionId, + inputs: WorkerInputs, + indexer: &Indexer, +) -> Option { + if !inputs.contains(WorkerInputs::RESIDENT_BLOCKS) { + return None; + } + let counts = indexer.track_resident_blocks(); + 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 +} diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index 41c7362d7839..f8928c501ee6 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::plugins::worker_selection::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,6 +157,12 @@ 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. The input is + // experimental; this check is the frontend's only use of that namespace. + anyhow::ensure!( + !prepared.inputs().contains(RESIDENT_BLOCKS), + "RESIDENT_BLOCKS requires the worker-selection policy to declare WorkerInputs::CACHE" + ); } Ok(prepared) } @@ -2275,6 +2282,22 @@ mod tests { } } + struct DeclaredInputsPicker(WorkerInputs); + + impl WorkerPicker for DeclaredInputsPicker { + fn required_worker_inputs(&self) -> WorkerInputs { + self.0 + } + + fn pick( + &self, + _context: &WorkerSelectionContext<'_>, + _input: WorkerInputView<'_>, + ) -> Result { + unreachable!("capability construction test does not select a worker") + } + } + fn fixed_policy( expected_shared_blocks: Option, worker: WorkerWithDpRank, @@ -2509,6 +2532,20 @@ mod tests { } } + #[tokio::test] + async fn resident_blocks_require_declared_cache_input() { + let config = KvRouterConfig::default(); + 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(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(