diff --git a/Cargo.lock b/Cargo.lock index e5b4ca52839e..28fc6c4250dd 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2767,6 +2767,7 @@ dependencies = [ "reqwest 0.12.28", "rmp-serde", "rstest 0.18.2", + "rustc-hash 2.1.2", "serde", "serde_json", "serial_test", diff --git a/lib/bindings/kvbm/Cargo.lock b/lib/bindings/kvbm/Cargo.lock index b3dff3b0a03f..1e72282f9020 100644 --- a/lib/bindings/kvbm/Cargo.lock +++ b/lib/bindings/kvbm/Cargo.lock @@ -2217,6 +2217,7 @@ dependencies = [ "rayon", "reqwest 0.12.28", "rmp-serde", + "rustc-hash 2.1.2", "serde", "serde_json", "sha2 0.10.9", diff --git a/lib/bindings/python/Cargo.lock b/lib/bindings/python/Cargo.lock index eacb4c06dce9..165e9f51500b 100644 --- a/lib/bindings/python/Cargo.lock +++ b/lib/bindings/python/Cargo.lock @@ -2268,6 +2268,7 @@ dependencies = [ "rayon", "reqwest 0.12.28", "rmp-serde", + "rustc-hash 2.1.2", "serde", "serde_json", "sha2 0.10.9", diff --git a/lib/bindings/python/src/dynamo/prometheus_names.py b/lib/bindings/python/src/dynamo/prometheus_names.py index bf602a0ae205..ce59b4a32c28 100644 --- a/lib/bindings/python/src/dynamo/prometheus_names.py +++ b/lib/bindings/python/src/dynamo/prometheus_names.py @@ -127,6 +127,26 @@ class frontend_service: NON_MAX_OVERLAP_SELECTIONS_TOTAL = "non_max_overlap_selections_total" # Effective KV overlap blocks lost by non-max-overlap selections OVERLAP_BLOCKS_LOST = "overlap_blocks_lost" + # Input tokens for cache-reuse observations started by the router + CACHE_LOSS_OBSERVATION_INPUT_TOKENS_TOTAL = ( + "cache_loss_observation_input_tokens_total" + ) + # Raw token observations at each cache-reuse funnel stage + CACHE_LOSS_FUNNEL_TOKENS_TOTAL = "cache_loss_funnel_tokens_total" + # Prompt-prefix tokens previously seen in this router's bounded history + KV_HISTORY_CACHED_PREFIX_TOKENS_TOTAL = "kv_history_cached_prefix_tokens_total" + # Cache-reuse observations by completion status + CACHE_LOSS_OBSERVATIONS_TOTAL = "cache_loss_observations_total" + # Distinct canonical block hashes retained by cache history + CACHE_LOSS_HISTORY_UNIQUE_HASHES = "cache_loss_history_unique_hashes" + # Tokens represented by distinct cache-history entries + CACHE_LOSS_HISTORY_REPRESENTED_TOKENS = "cache_loss_history_represented_tokens" + # Estimated bytes used by retained cache-history entries + CACHE_LOSS_HISTORY_ESTIMATED_BYTES = "cache_loss_history_estimated_bytes" + # Configured distinct-block capacity of cache history + CACHE_LOSS_HISTORY_CAPACITY_BLOCKS = "cache_loss_history_capacity_blocks" + # Configured byte budget of cache history + CACHE_LOSS_HISTORY_CAPACITY_BYTES = "cache_loss_history_capacity_bytes" # Number of cached tokens (prefix cache hits) per request CACHED_TOKENS = "cached_tokens" # Tokenizer latency in milliseconds diff --git a/lib/llm/Cargo.toml b/lib/llm/Cargo.toml index 084ae6e9dadf..f6722e96fd2c 100644 --- a/lib/llm/Cargo.toml +++ b/lib/llm/Cargo.toml @@ -109,6 +109,7 @@ oneshot = { workspace = true } parking_lot = { workspace = true } prometheus = { workspace = true } rand = { workspace = true } +rustc-hash = "2.1.1" serde = { workspace = true } serde_json = { workspace = true } sha2 = "0.10" diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index ebdbd624ab9f..24727236573f 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -57,6 +57,7 @@ use futures::stream; pub use dynamo_kv_router::protocols; pub use dynamo_kv_router::scheduling; +mod cache_history; pub(crate) mod embedded; pub mod encoder_router; pub mod indexer; diff --git a/lib/llm/src/kv_router/cache_history.rs b/lib/llm/src/kv_router/cache_history.rs new file mode 100644 index 000000000000..c41b7a0a1b3e --- /dev/null +++ b/lib/llm/src/kv_router/cache_history.rs @@ -0,0 +1,320 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +mod sync; + +use std::{ + collections::VecDeque, + sync::{ + Arc, OnceLock, + atomic::{AtomicUsize, Ordering}, + }, +}; + +use parking_lot::{Mutex, RwLock}; +use rustc_hash::FxHashSet; + +pub(crate) const CACHE_REUSE_HISTORY_ENABLED_ENV: &str = "DYN_ROUTER_CACHE_REUSE_HISTORY"; +pub(crate) const HISTORY_BYTES_ENV: &str = "DYN_ROUTER_CACHE_REUSE_HISTORY_BYTES"; +pub(crate) const DEFAULT_HISTORY_BYTES: usize = 256 * 1024 * 1024; +pub(crate) const ESTIMATED_BYTES_PER_ENTRY: usize = 32; +const HISTORY_SHARDS: usize = 64; + +pub(crate) fn enabled() -> bool { + dynamo_runtime::config::env_is_truthy(CACHE_REUSE_HISTORY_ENABLED_ENV) +} + +/// Per-`RoutingHost` FIFO membership history for canonical sequence hashes. +/// +/// Reads touch only one shard at a time. The FIFO lock serializes insertions and +/// evictions so a recomputation cannot race an eviction and leave membership +/// inconsistent. Duplicate hashes neither consume capacity nor refresh age. Each +/// host applies the configured byte budget independently. +/// With router replica sync enabled, peer completions also enter the same FIFO; +/// otherwise the history remains local. Peer updates are best-effort, without replay. +pub(crate) struct CacheHistory { + block_tokens: u64, + capacity_entries: usize, + capacity_bytes: usize, + fifo: Mutex>, + shards: Box<[RwLock>]>, + retained_entries: AtomicUsize, + sync: OnceLock, +} + +impl CacheHistory { + pub(crate) fn from_env(block_tokens: u32) -> Arc { + let requested_bytes = match std::env::var(HISTORY_BYTES_ENV) { + Ok(value) => match value.parse::() { + Ok(bytes) if bytes > 0 => bytes, + _ => { + tracing::warn!( + value, + effective_bytes = DEFAULT_HISTORY_BYTES, + "Invalid DYN_ROUTER_CACHE_REUSE_HISTORY_BYTES; using default" + ); + DEFAULT_HISTORY_BYTES + } + }, + Err(_) => DEFAULT_HISTORY_BYTES, + }; + Arc::new(Self::with_byte_budget(block_tokens, requested_bytes)) + } + + fn with_byte_budget(block_tokens: u32, requested_bytes: usize) -> Self { + assert!( + block_tokens > 0, + "cache history block size must be positive" + ); + let capacity_bytes = requested_bytes.max(ESTIMATED_BYTES_PER_ENTRY); + let capacity_entries = capacity_bytes / ESTIMATED_BYTES_PER_ENTRY; + let shards = (0..HISTORY_SHARDS) + .map(|_| RwLock::new(FxHashSet::default())) + .collect::>() + .into_boxed_slice(); + Self { + block_tokens: u64::from(block_tokens), + capacity_entries, + capacity_bytes, + fifo: Mutex::new(VecDeque::new()), + shards, + retained_entries: AtomicUsize::new(0), + sync: OnceLock::new(), + } + } + + #[cfg(test)] + pub(crate) fn with_capacity(capacity_entries: usize, block_tokens: u32) -> Self { + Self::with_byte_budget( + block_tokens, + capacity_entries.saturating_mul(ESTIMATED_BYTES_PER_ENTRY), + ) + } + + pub(crate) fn previously_computed_tokens(&self, sequence_hashes: &[u64]) -> u64 { + let blocks = sequence_hashes + .iter() + .take_while(|&&hash| { + let shard = self.shard(hash); + self.shards[shard].read().contains(&hash) + }) + .count() as u64; + blocks.saturating_mul(self.block_tokens) + } + + pub(crate) fn record_completed(&self, sequence_hashes: I) -> Option + where + I: Iterator + Clone, + { + let stats = self.insert(sequence_hashes.clone()); + if let Some(sync) = self.sync.get() { + sync.publish(sequence_hashes); + } + stats + } + + fn insert(&self, sequence_hashes: I) -> Option + where + I: Iterator + Clone, + { + if sequence_hashes.clone().all(|hash| { + let shard = self.shard(hash); + self.shards[shard].read().contains(&hash) + }) { + return None; + } + + let mut fifo = self.fifo.lock(); + let initial_len = fifo.len(); + for hash in sequence_hashes { + let shard = self.shard(hash); + let mut membership = self.shards[shard].write(); + if membership.contains(&hash) { + continue; + } + + if fifo.len() == self.capacity_entries { + let evicted = fifo.pop_front().expect("history capacity is positive"); + let evicted_shard = self.shard(evicted); + let removed = if evicted_shard == shard { + membership.remove(&evicted) + } else { + self.shards[evicted_shard].write().remove(&evicted) + }; + debug_assert!(removed, "FIFO and membership shards diverged"); + } + + let inserted = membership.insert(hash); + debug_assert!(inserted, "FIFO writer must serialize membership updates"); + fifo.push_back(hash); + } + self.retained_entries.store(fifo.len(), Ordering::Relaxed); + (fifo.len() != initial_len).then(|| self.stats_for_entries(fifo.len())) + } + + pub(crate) fn stats(&self) -> CacheHistoryStats { + let retained_entries = self.retained_entries.load(Ordering::Relaxed); + self.stats_for_entries(retained_entries) + } + + fn stats_for_entries(&self, retained_entries: usize) -> CacheHistoryStats { + CacheHistoryStats { + retained_entries, + represented_tokens: (retained_entries as u64).saturating_mul(self.block_tokens), + estimated_retained_bytes: retained_entries.saturating_mul(ESTIMATED_BYTES_PER_ENTRY), + capacity_entries: self.capacity_entries, + capacity_bytes: self.capacity_bytes, + } + } + + fn shard(&self, hash: u64) -> usize { + let mut mixed = hash; + mixed ^= mixed >> 33; + mixed = mixed.wrapping_mul(0xff51_afd7_ed55_8ccd); + mixed ^= mixed >> 33; + (mixed as usize) & (HISTORY_SHARDS - 1) + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) struct CacheHistoryStats { + pub(crate) retained_entries: usize, + pub(crate) represented_tokens: u64, + pub(crate) estimated_retained_bytes: usize, + pub(crate) capacity_entries: usize, + pub(crate) capacity_bytes: usize, +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::Barrier; + + #[derive(Clone, Default)] + struct CapturedLogs(Arc>>); + + impl std::io::Write for CapturedLogs { + fn write(&mut self, bytes: &[u8]) -> std::io::Result { + self.0.lock().unwrap().extend_from_slice(bytes); + Ok(bytes.len()) + } + + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } + } + + fn history_with_captured_logs(value: Option<&str>) -> (CacheHistoryStats, String) { + let logs = CapturedLogs::default(); + let writer = logs.clone(); + let _guard = tracing::subscriber::set_default( + tracing_subscriber::fmt() + .without_time() + .with_ansi(false) + .with_max_level(tracing::Level::WARN) + .with_writer(move || writer.clone()) + .finish(), + ); + let stats = temp_env::with_var(HISTORY_BYTES_ENV, value, || { + CacheHistory::from_env(16).stats() + }); + let output = String::from_utf8(logs.0.lock().unwrap().clone()).unwrap(); + (stats, output) + } + + #[test] + fn longest_prefix_stops_at_first_unseen_hash() { + let history = CacheHistory::with_capacity(4, 16); + assert!(history.record_completed([10, 20, 30].into_iter()).is_some()); + + assert_eq!(history.previously_computed_tokens(&[10, 20, 99, 30]), 32); + } + + #[test] + fn duplicate_does_not_consume_capacity_or_refresh_fifo_age() { + let history = CacheHistory::with_capacity(2, 8); + assert!(history.record_completed([10, 20].into_iter()).is_some()); + history.record_completed([10, 30].into_iter()); + + assert_eq!(history.previously_computed_tokens(&[10]), 0); + assert_eq!(history.previously_computed_tokens(&[20, 30]), 16); + assert_eq!(history.stats().retained_entries, 2); + } + + #[test] + fn duplicate_only_completion_reports_no_change() { + let history = CacheHistory::with_capacity(2, 8); + assert!(history.record_completed([10, 20].into_iter()).is_some()); + assert!(history.record_completed([10, 20, 10].into_iter()).is_none()); + } + + #[test] + fn concurrent_recomputations_and_evictions_keep_fifo_consistent() { + let history = Arc::new(CacheHistory::with_capacity(8, 8)); + let barrier = Arc::new(Barrier::new(8)); + std::thread::scope(|scope| { + for thread in 0..8_u64 { + let history = history.clone(); + let barrier = barrier.clone(); + scope.spawn(move || { + barrier.wait(); + for offset in 0..64_u64 { + history + .record_completed([offset % 4, thread * 64 + offset + 10].into_iter()); + } + }); + } + }); + + let fifo = history.fifo.lock(); + assert_eq!(fifo.len(), history.stats().retained_entries); + let unique = fifo.iter().copied().collect::>(); + assert_eq!(unique.len(), fifo.len()); + for &hash in fifo.iter() { + assert!(history.shards[history.shard(hash)].read().contains(&hash)); + } + } + + #[test] + fn environment_variable_names_are_stable() { + assert_eq!( + CACHE_REUSE_HISTORY_ENABLED_ENV, + "DYN_ROUTER_CACHE_REUSE_HISTORY" + ); + assert_eq!(HISTORY_BYTES_ENV, "DYN_ROUTER_CACHE_REUSE_HISTORY_BYTES"); + } + + #[test] + fn disabled_by_default() { + temp_env::with_var_unset(CACHE_REUSE_HISTORY_ENABLED_ENV, || assert!(!enabled())); + } + + #[test] + fn explicit_flag_enables_history() { + temp_env::with_var(CACHE_REUSE_HISTORY_ENABLED_ENV, Some("true"), || { + assert!(enabled()) + }); + } + + #[test] + fn invalid_byte_budget_warns_and_uses_default() { + for value in ["0", "64m", " 268435456 "] { + let (stats, output) = history_with_captured_logs(Some(value)); + assert_eq!(stats.capacity_bytes, DEFAULT_HISTORY_BYTES); + assert!(output.contains("Invalid DYN_ROUTER_CACHE_REUSE_HISTORY_BYTES; using default")); + assert!(output.contains(value), "captured log: {output}"); + assert!(output.contains(&format!("effective_bytes={DEFAULT_HISTORY_BYTES}"))); + } + } + + #[test] + fn valid_or_unset_byte_budget_does_not_warn() { + let (stats, output) = history_with_captured_logs(None); + assert_eq!(stats.capacity_bytes, DEFAULT_HISTORY_BYTES); + assert!(output.is_empty()); + + let (stats, output) = history_with_captured_logs(Some("67108864")); + assert_eq!(stats.capacity_bytes, 67_108_864); + assert!(output.is_empty()); + } +} diff --git a/lib/llm/src/kv_router/cache_history/sync.rs b/lib/llm/src/kv_router/cache_history/sync.rs new file mode 100644 index 000000000000..cb07b8e286b9 --- /dev/null +++ b/lib/llm/src/kv_router/cache_history/sync.rs @@ -0,0 +1,386 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::{ + future::Future, + sync::{Arc, Weak}, + time::Duration, +}; + +use dynamo_runtime::{ + component::Endpoint, + traits::DistributedRuntimeProvider, + transports::event_plane::{EventPublisher, EventSubscriber}, +}; +use serde::{Deserialize, Serialize}; +use tokio::sync::mpsc; +use tokio_util::sync::CancellationToken; + +use super::CacheHistory; +use crate::kv_router::{KvRouter, metrics::RouterRequestMetrics}; + +const SUBJECT: &str = "cache-history-v1"; +const MAX_BATCH_HASHES: usize = 4096; +const QUEUE_BATCHES: usize = 256; +const INITIAL_RETRY_DELAY: Duration = Duration::from_millis(100); +const MAX_RETRY_DELAY: Duration = Duration::from_secs(5); + +#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)] +struct Domain { + model: String, + block_size: u32, + is_eagle: bool, +} + +#[derive(Serialize, Deserialize)] +struct CompletedHashes { + source: String, + domain: Domain, + hashes: Vec, +} + +fn apply_peer( + history: &CacheHistory, + source: &str, + domain: &Domain, + event: CompletedHashes, +) -> Option { + if event.source == source || event.domain != *domain || event.hashes.len() > MAX_BATCH_HASHES { + return None; + } + history.insert(event.hashes.into_iter()) +} + +pub(super) struct HistorySync { + sender: mpsc::Sender>, + cancel: CancellationToken, +} + +impl Drop for HistorySync { + fn drop(&mut self) { + self.cancel.cancel(); + } +} + +impl HistorySync { + pub(super) fn publish(&self, mut hashes: impl Iterator) { + loop { + let batch: Vec<_> = hashes.by_ref().take(MAX_BATCH_HASHES).collect(); + if batch.is_empty() { + break; + } + if self.sender.try_send(batch).is_err() { + tracing::trace!("Cache history replica queue unavailable; dropping peer update"); + break; + } + } + } +} + +impl CacheHistory { + pub(crate) fn start_replica_sync( + self: &Arc, + router: &KvRouter, + metrics: Arc, + ) { + if !router.kv_router_config().router_replica_sync { + return; + } + let endpoint = router.client().endpoint.clone(); + let domain = Domain { + model: router.tracking_model_name.clone(), + block_size: router.block_size(), + is_eagle: router.is_eagle(), + }; + self.start_sync(endpoint, domain, metrics); + } + + fn start_sync( + self: &Arc, + endpoint: Endpoint, + domain: Domain, + metrics: Arc, + ) { + let cancel = CancellationToken::new(); + let (sender, receiver) = mpsc::channel(QUEUE_BATCHES); + if self + .sync + .set(HistorySync { + sender, + cancel: cancel.clone(), + }) + .is_err() + { + return; + } + let history = Arc::downgrade(self); + tokio::spawn(async move { + let result = tokio::select! { + _ = cancel.cancelled() => return, + result = run(endpoint, domain, history, receiver, metrics, cancel.clone()) => result, + }; + if let Err(error) = result { + tracing::warn!(%error, "Cache history replica sync stopped"); + } + }); + } +} + +async fn run( + endpoint: Endpoint, + domain: Domain, + history: Weak, + mut receiver: mpsc::Receiver>, + metrics: Arc, + cancel: CancellationToken, +) -> anyhow::Result<()> { + let transport = endpoint.drt().default_event_transport_kind(); + let source = uuid::Uuid::new_v4().to_string(); + loop { + let connected = connect_with_retry( + || async { + let subscriber = + EventSubscriber::for_endpoint_with_transport(&endpoint, SUBJECT, transport) + .await? + .typed::(); + let publisher = + EventPublisher::for_endpoint_with_transport(&endpoint, SUBJECT, transport) + .await?; + Ok((publisher, subscriber)) + }, + &cancel, + ) + .await; + let Some((publisher, mut subscriber)) = connected else { + return Ok(()); + }; + loop { + tokio::select! { + _ = cancel.cancelled() => return Ok(()), + hashes = receiver.recv() => { + let Some(hashes) = hashes else { return Ok(()) }; + let event = CompletedHashes { source: source.clone(), domain: domain.clone(), hashes }; + if let Err(error) = publisher.publish(&event).await { + tracing::warn!(%error, "Failed to publish completed cache history hashes; reconnecting"); + break; + } + } + event = subscriber.next() => { + match event { + Some(Ok((_, event))) => { + let Some(history) = history.upgrade() else { return Ok(()) }; + if let Some(stats) = apply_peer(&history, &source, &domain, event) { + metrics.set_cache_history_retained(stats); + } + } + Some(Err(error)) => tracing::warn!(%error, "Invalid cache history replica event"), + None => { + tracing::warn!("Cache history replica stream ended; reconnecting"); + break; + } + } + } + } + } + tokio::select! { + _ = cancel.cancelled() => return Ok(()), + _ = tokio::time::sleep(INITIAL_RETRY_DELAY) => {} + } + } +} + +async fn connect_with_retry(mut connect: F, cancel: &CancellationToken) -> Option +where + F: FnMut() -> Fut, + Fut: Future>, +{ + let mut delay = INITIAL_RETRY_DELAY; + loop { + let result = tokio::select! { + _ = cancel.cancelled() => return None, + result = connect() => result, + }; + match result { + Ok(connection) => return Some(connection), + Err(error) => { + tracing::warn!(%error, "Cache history replica transport unavailable; retrying") + } + } + tokio::select! { + _ = cancel.cancelled() => return None, + _ = tokio::time::sleep(delay) => {} + } + delay = (delay * 2).min(MAX_RETRY_DELAY); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test(start_paused = true)] + async fn transient_connect_failure_retries_with_backoff() { + let cancel = CancellationToken::new(); + let mut attempts = 0; + let start = tokio::time::Instant::now(); + let result = connect_with_retry( + || { + attempts += 1; + std::future::ready(if attempts < 3 { + Err(anyhow::anyhow!("offline")) + } else { + Ok(42) + }) + }, + &cancel, + ) + .await; + assert_eq!(result, Some(42)); + assert_eq!(attempts, 3); + assert_eq!(start.elapsed(), Duration::from_millis(300)); + } + + #[tokio::test(start_paused = true)] + async fn cancellation_stops_retrying_an_unavailable_transport() { + let cancel = CancellationToken::new(); + let stop = cancel.clone(); + tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(50)).await; + stop.cancel(); + }); + let mut attempts = 0; + let result = connect_with_retry( + || { + attempts += 1; + std::future::ready(Err::<(), _>(anyhow::anyhow!("offline"))) + }, + &cancel, + ) + .await; + assert!(result.is_none()); + assert_eq!(attempts, 1); + } + + fn domain() -> Domain { + Domain { + model: "model-a".into(), + block_size: 16, + is_eagle: false, + } + } + + fn event(hashes: &[u64]) -> CompletedHashes { + CompletedHashes { + source: "peer".into(), + domain: domain(), + hashes: hashes.to_vec(), + } + } + + #[test] + fn peer_and_local_insertions_have_identical_fifo_behavior_without_rebroadcast() { + let local = CacheHistory::with_capacity(3, 16); + let peer = CacheHistory::with_capacity(3, 16); + let (sender, mut receiver) = mpsc::channel(8); + assert!( + peer.sync + .set(HistorySync { + sender, + cancel: CancellationToken::new() + }) + .is_ok() + ); + for batch in [&[10, 20, 30][..], &[10, 40], &[20, 50], &[20, 50]] { + local.record_completed(batch.iter().copied()); + apply_peer(&peer, "local", &domain(), event(batch)); + assert_eq!(*local.fifo.lock(), *peer.fifo.lock()); + assert_eq!(local.stats(), peer.stats()); + } + assert!(receiver.try_recv().is_err()); + peer.record_completed([20, 50].into_iter()); + assert_eq!(receiver.try_recv().unwrap(), vec![20, 50]); + assert_eq!(*local.fifo.lock(), *peer.fifo.lock()); + } + + #[test] + fn rejects_self_and_incompatible_hash_domains() { + let history = CacheHistory::with_capacity(3, 16); + apply_peer(&history, "peer", &domain(), event(&[10])); + for incompatible in [ + Domain { + model: "model-b".into(), + ..domain() + }, + Domain { + block_size: 8, + ..domain() + }, + Domain { + is_eagle: true, + ..domain() + }, + ] { + apply_peer(&history, "local", &incompatible, event(&[10])); + } + assert_eq!(history.stats().retained_entries, 0); + apply_peer(&history, "local", &domain(), event(&[10])); + assert_eq!(history.previously_computed_tokens(&[10]), 16); + } + + #[test] + fn publication_is_chunked_and_does_not_block_when_full() { + let (sender, mut receiver) = mpsc::channel(2); + let sync = HistorySync { + sender, + cancel: CancellationToken::new(), + }; + sync.publish(0..MAX_BATCH_HASHES as u64 + 3); + assert_eq!(receiver.try_recv().unwrap().len(), MAX_BATCH_HASHES); + assert_eq!(receiver.try_recv().unwrap().len(), 3); + sync.publish(0..3 * MAX_BATCH_HASHES as u64); + assert_eq!(receiver.len(), 2); + } + + #[tokio::test] + #[serial_test::serial] + async fn two_frontends_learn_completed_hashes_over_the_event_plane() -> anyhow::Result<()> { + let runtime = dynamo_runtime::Runtime::from_current()?; + let distributed = dynamo_runtime::DistributedRuntime::new( + runtime, + dynamo_runtime::distributed::DistributedConfig::process_local(), + ) + .await?; + let component = distributed + .namespace(format!("history-sync-{}", uuid::Uuid::new_v4()))? + .component("workers")?; + let endpoint = component.endpoint("generate"); + let metrics = RouterRequestMetrics::from_component(&component); + let first = Arc::new(CacheHistory::with_capacity(3, 16)); + let second = Arc::new(CacheHistory::with_capacity(3, 16)); + first.start_sync(endpoint.clone(), domain(), metrics.clone()); + second.start_sync(endpoint, domain(), metrics); + tokio::time::timeout(std::time::Duration::from_secs(5), async { + loop { + first.record_completed([10, 20].into_iter()); + if second.previously_computed_tokens(&[10, 20]) == 32 { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + } + }) + .await?; + tokio::time::timeout(std::time::Duration::from_secs(5), async { + loop { + second.record_completed([10, 20, 30].into_iter()); + if first.previously_computed_tokens(&[10, 20, 30]) == 48 { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + } + }) + .await?; + drop((first, second)); + distributed.shutdown(); + Ok(()) + } +} diff --git a/lib/llm/src/kv_router/metrics.rs b/lib/llm/src/kv_router/metrics.rs index 1ab197f68b23..f9990fb42f82 100644 --- a/lib/llm/src/kv_router/metrics.rs +++ b/lib/llm/src/kv_router/metrics.rs @@ -63,6 +63,7 @@ use prometheus::{ }; use crate::http::service::metrics::generate_log_buckets; +use crate::kv_router::cache_history::{self, CacheHistoryStats}; use crate::protocols::common::timing::WORKER_TYPE_PREFILL; use dynamo_kv_router::indexer::ApproximateLruStats; @@ -1037,6 +1038,21 @@ pub struct RouterRequestMetrics { pub shared_cache_beyond_blocks: prometheus::Histogram, pub non_max_overlap_selections_total: IntCounterVec, pub overlap_blocks_lost: HistogramVec, + pub(crate) cache_history: Option, +} + +#[derive(Clone)] +pub(crate) struct CacheHistoryMetrics { + observation_input_tokens_total: IntCounter, + f0_tokens_total: IntCounter, + f1_tokens_total: IntCounter, + complete_observations_total: IntCounter, + incomplete_observations_total: IntCounter, + retained_entries: IntGauge, + represented_tokens: IntGauge, + estimated_retained_bytes: IntGauge, + capacity_entries: IntGauge, + capacity_bytes: IntGauge, } static ROUTER_REQUEST_METRICS: OnceLock> = OnceLock::new(); @@ -1158,6 +1174,77 @@ impl RouterRequestMetrics { .expect("failed to create router_overlap_blocks_lost"); non_max_overlap_selections_total.with_label_values(&[WORKER_TYPE_PREFILL]); overlap_blocks_lost.with_label_values(&[WORKER_TYPE_PREFILL]); + let cache_history = cache_history::enabled().then(|| { + let observation_input_tokens_total = metrics + .create_intcounter( + &router_metric(frontend_service::CACHE_LOSS_OBSERVATION_INPUT_TOKENS_TOTAL), + "Input tokens for cache-reuse observations started by the router", + extra_labels, + ) + .expect( + "failed to create router_cache_loss_observation_input_tokens_total", + ); + let funnel_tokens_total = metrics + .create_intcountervec( + &router_metric(frontend_service::CACHE_LOSS_FUNNEL_TOKENS_TOTAL), + "Raw token observations at each cache-reuse funnel stage", + &["stage"], + extra_labels, + ) + .expect("failed to create router_cache_loss_funnel_tokens_total"); + let f0_tokens_total = funnel_tokens_total.with_label_values(&["f0"]); + let f1_tokens_total = metrics + .create_intcounter( + &router_metric(frontend_service::KV_HISTORY_CACHED_PREFIX_TOKENS_TOTAL), + "Prompt-prefix tokens previously seen in this router's bounded history", + extra_labels, + ) + .expect("failed to create router_kv_history_cached_prefix_tokens_total"); + let observations_total = metrics + .create_intcountervec( + &router_metric(frontend_service::CACHE_LOSS_OBSERVATIONS_TOTAL), + "Cache-reuse observations by completion status", + &["result"], + extra_labels, + ) + .expect("failed to create router_cache_loss_observations_total"); + let complete_observations_total = + observations_total.with_label_values(&["complete"]); + let incomplete_observations_total = + observations_total.with_label_values(&["incomplete"]); + let gauge = |suffix, help| { + metrics + .create_intgauge(&router_metric(suffix), help, extra_labels) + .unwrap_or_else(|error| panic!("failed to create {suffix}: {error}")) + }; + CacheHistoryMetrics { + observation_input_tokens_total, + f0_tokens_total, + f1_tokens_total, + complete_observations_total, + incomplete_observations_total, + retained_entries: gauge( + frontend_service::CACHE_LOSS_HISTORY_UNIQUE_HASHES, + "Distinct canonical block hashes retained by the last cache-history-enabled RoutingHost to update this process-global gauge", + ), + represented_tokens: gauge( + frontend_service::CACHE_LOSS_HISTORY_REPRESENTED_TOKENS, + "Tokens represented by the last cache-history-enabled RoutingHost to update this process-global gauge", + ), + estimated_retained_bytes: gauge( + frontend_service::CACHE_LOSS_HISTORY_ESTIMATED_BYTES, + "Estimated retained bytes for the last cache-history-enabled RoutingHost to update this process-global gauge", + ), + capacity_entries: gauge( + frontend_service::CACHE_LOSS_HISTORY_CAPACITY_BLOCKS, + "Per-RoutingHost distinct-block capacity; process-global gauge reports the last initialized cache-history-enabled host", + ), + capacity_bytes: gauge( + frontend_service::CACHE_LOSS_HISTORY_CAPACITY_BYTES, + "Per-RoutingHost byte budget; process-global gauge reports the last initialized cache-history-enabled host", + ), + } + }); Arc::new(Self { requests_started_total, requests_total, @@ -1171,6 +1258,7 @@ impl RouterRequestMetrics { shared_cache_beyond_blocks, non_max_overlap_selections_total, overlap_blocks_lost, + cache_history, }) }) .clone() @@ -1197,6 +1285,55 @@ impl RouterRequestMetrics { .with_label_values(&[worker_type]) .observe(overlap_blocks_lost); } + + pub(crate) fn observe_cache_history_input(&self, prompt_tokens: u64) { + let Some(metrics) = &self.cache_history else { + return; + }; + metrics.observation_input_tokens_total.inc_by(prompt_tokens); + } + + pub(crate) fn observe_cache_history_complete( + &self, + prompt_tokens: u64, + previously_seen_tokens: u64, + ) { + let Some(metrics) = &self.cache_history else { + return; + }; + metrics.f0_tokens_total.inc_by(prompt_tokens); + metrics.f1_tokens_total.inc_by(previously_seen_tokens); + metrics.complete_observations_total.inc(); + } + + pub(crate) fn observe_cache_history_incomplete(&self) { + if let Some(metrics) = &self.cache_history { + metrics.incomplete_observations_total.inc(); + } + } + + pub(crate) fn set_cache_history_capacity(&self, stats: CacheHistoryStats) { + let Some(metrics) = &self.cache_history else { + return; + }; + let gauge = |value: usize| i64::try_from(value).unwrap_or(i64::MAX); + metrics.capacity_entries.set(gauge(stats.capacity_entries)); + metrics.capacity_bytes.set(gauge(stats.capacity_bytes)); + } + + pub(crate) fn set_cache_history_retained(&self, stats: CacheHistoryStats) { + let Some(metrics) = &self.cache_history else { + return; + }; + let gauge = |value: usize| i64::try_from(value).unwrap_or(i64::MAX); + metrics.retained_entries.set(gauge(stats.retained_entries)); + metrics + .represented_tokens + .set(i64::try_from(stats.represented_tokens).unwrap_or(i64::MAX)); + metrics + .estimated_retained_bytes + .set(gauge(stats.estimated_retained_bytes)); + } } /// Process-local observability for the experimental approximate LRU. diff --git a/lib/llm/src/kv_router/routing_host.rs b/lib/llm/src/kv_router/routing_host.rs index f7648ee73829..32f49f6cb606 100644 --- a/lib/llm/src/kv_router/routing_host.rs +++ b/lib/llm/src/kv_router/routing_host.rs @@ -29,7 +29,12 @@ use futures::stream::{self, StreamExt}; use tracing::Instrument; use crate::{ - kv_router::{KvRouter, metrics::RouterRequestMetrics, to_worker_selection_session_context}, + kv_router::{ + KvRouter, + cache_history::{self, CacheHistory}, + metrics::RouterRequestMetrics, + to_worker_selection_session_context, + }, lora::{LoadEstimator, LoraFilter}, preprocessor::PreprocessedRequest, protocols::common::{ @@ -57,7 +62,7 @@ use cancellation::{CleanupBudget, DispatchCancellation, StagedKv, await_with_cle use kv_selection::{RoutingRequestParts, SelectionOptions, WorkerSelection}; use occupancy::HostedOccupancy; pub(crate) use request_guard::prompt_private_blocks; -use request_guard::{KvRequestCleanup, LoraLoadGuard, RequestGuard}; +use request_guard::{CacheHistoryTracking, KvRequestCleanup, LoraLoadGuard, RequestGuard}; const OUTPUT_REPLAY_ID_ANNOTATION_KEY: &str = "output_replay_id"; const OUTPUT_REPLAY_CONSUMER_RUNTIME_KEY: &str = "output_replay_consumer"; @@ -296,6 +301,7 @@ pub struct RoutingHost { inner: PushRouter>, policy: RoutingPolicy, request_metrics: Arc, + cache_history: Option>, affinity: Option, session_affinity_mode: SessionAffinityMode, hosted_occupancy: Option, @@ -435,6 +441,12 @@ impl RoutingHost { // and the standalone router create RoutingHost, so this covers both. let request_metrics = RouterRequestMetrics::from_component(kv_router.client().endpoint.component()); + let cache_history = + cache_history::enabled().then(|| CacheHistory::from_env(kv_router.block_size())); + if let Some(history) = &cache_history { + request_metrics.set_cache_history_capacity(history.stats()); + history.start_replica_sync(&kv_router, request_metrics.clone()); + } RoutingHost { inner, @@ -444,6 +456,7 @@ impl RoutingHost { .as_ref() .map(AffinityCoordinator::mode) .unwrap_or_default(), + cache_history, affinity, hosted_occupancy: None, lora: None, @@ -521,6 +534,7 @@ impl RoutingHost { .as_ref() .map(AffinityCoordinator::mode) .unwrap_or_default(), + cache_history: None, affinity, hosted_occupancy, lora: lora diff --git a/lib/llm/src/kv_router/routing_host/kv.rs b/lib/llm/src/kv_router/routing_host/kv.rs index 88498369b9fe..eff004a1cdaf 100644 --- a/lib/llm/src/kv_router/routing_host/kv.rs +++ b/lib/llm/src/kv_router/routing_host/kv.rs @@ -410,6 +410,19 @@ impl RoutingHost { let chooser = self.kv_router(); let block_size = chooser.block_size() as usize; let selected_worker = selection.worker; + let history_prompt_hashes = self.cache_history.as_ref().and_then(|_| { + if chooser.indexer().records_routing_decisions() { + selection + .routing_hashes + .as_ref() + .map(|hashes| hashes.sequence_hashes.clone()) + } else { + selection + .routing_hashes + .take() + .map(|hashes| hashes.sequence_hashes) + } + }); let mut guard = match cleanup { Some(cleanup) => RequestGuard::new_kv_with_cleanup( self.request_metrics.clone(), @@ -427,6 +440,20 @@ impl RoutingHost { selection.request_lifecycle.take(), ), }; + if let (Some(history), Some(prompt_hashes)) = + (self.cache_history.as_ref(), history_prompt_hashes) + { + guard.track_cache_history( + CacheHistoryTracking::new( + Arc::clone(history), + prompt_hashes, + routing_parts.token_ids.len() as u64, + ), + request, + chooser.block_size(), + chooser.is_eagle(), + ); + } let record_result: Result<(), Error> = async { if !is_query_only && chooser.indexer().records_routing_decisions() { diff --git a/lib/llm/src/kv_router/routing_host/kv_selection.rs b/lib/llm/src/kv_router/routing_host/kv_selection.rs index 00f16608e8fa..7ec14a8c78ea 100644 --- a/lib/llm/src/kv_router/routing_host/kv_selection.rs +++ b/lib/llm/src/kv_router/routing_host/kv_selection.rs @@ -203,8 +203,9 @@ impl RoutingHost { return Err(anyhow::anyhow!(error)); } } - let return_routing_hashes = - !is_query_only && self.kv_router().indexer().records_routing_decisions(); + let return_routing_hashes = !is_query_only + && (self.kv_router().indexer().records_routing_decisions() + || self.cache_history.is_some()); let SelectionOptions { pinned_target, affinity_target, diff --git a/lib/llm/src/kv_router/routing_host/request_guard.rs b/lib/llm/src/kv_router/routing_host/request_guard.rs index 7622637cbf6f..07e597015d7d 100644 --- a/lib/llm/src/kv_router/routing_host/request_guard.rs +++ b/lib/llm/src/kv_router/routing_host/request_guard.rs @@ -5,8 +5,9 @@ use std::{collections::HashMap, sync::Arc}; use crate::{ kv_router::{ - KvRouter, indexer::ApproximateRequestLease, metrics::RouterRequestMetrics, - prefill_router::BYPASS_REMOTE_PREFILL_ANNOTATION, request_lease::RequestAttemptLease, + KvRouter, cache_history::CacheHistory, indexer::ApproximateRequestLease, + metrics::RouterRequestMetrics, prefill_router::BYPASS_REMOTE_PREFILL_ANNOTATION, + request_lease::RequestAttemptLease, }, lora::LoadEstimator, preprocessor::PreprocessedRequest, @@ -68,6 +69,60 @@ struct MaterializedOutputBlocks { private_blocks: usize, } +pub(super) struct CacheHistoryTracking { + history: Arc, + prompt_hashes: Vec, + output_hashes: Vec, + prompt_tokens: u64, + previously_seen_tokens: u64, +} + +impl CacheHistoryTracking { + pub(super) fn new( + history: Arc, + prompt_hashes: Vec, + prompt_tokens: u64, + ) -> Self { + let previously_seen_tokens = history.previously_computed_tokens(&prompt_hashes); + Self { + history, + prompt_hashes, + output_hashes: Vec::new(), + prompt_tokens, + previously_seen_tokens, + } + } +} + +struct CacheHistoryFinalization { + prompt_tokens: u64, + previously_seen_tokens: u64, + retained: Option, +} + +fn finalize_cache_history( + tracking: &mut Option, + record_completed: bool, +) -> Option { + let tracking = tracking.take()?; + let retained = record_completed + .then(|| { + tracking.history.record_completed( + tracking + .prompt_hashes + .iter() + .copied() + .chain(tracking.output_hashes.iter().copied()), + ) + }) + .flatten(); + Some(CacheHistoryFinalization { + prompt_tokens: tracking.prompt_tokens, + previously_seen_tokens: tracking.previously_seen_tokens, + retained, + }) +} + pub(crate) fn prompt_private_blocks( token_count: usize, complete_blocks: usize, @@ -570,10 +625,12 @@ pub(super) struct RequestGuard { observability: RequestObservability, output_blocks: OutputBlockTracker, approximate_lru: Option, + approximate_lru_active: bool, output_hashes: Option, record_itl_at_completion: bool, prefill_marked: bool, migration_state: Option, + cache_history: Option, _lora_load: Option, } @@ -644,10 +701,12 @@ impl RequestGuard { expected_output_tokens, ), approximate_lru, + approximate_lru_active: false, output_hashes, record_itl_at_completion: false, prefill_marked: false, migration_state: request.migration_state.clone(), + cache_history: None, _lora_load: None, } } @@ -673,10 +732,12 @@ impl RequestGuard { // when the request completes rather than observing every streamed token. output_blocks: OutputBlockTracker::new(false, request.token_ids.len(), 1, None), approximate_lru: None, + approximate_lru_active: false, output_hashes: None, record_itl_at_completion: true, prefill_marked: false, migration_state: request.migration_state.clone(), + cache_history: None, _lora_load: lora_load, } } @@ -717,6 +778,24 @@ impl RequestGuard { self.approximate_lru.is_some() } + pub(super) fn track_cache_history( + &mut self, + tracking: CacheHistoryTracking, + request: &PreprocessedRequest, + block_size: u32, + is_eagle: bool, + ) { + self.observability + .request_metrics() + .observe_cache_history_input(tracking.prompt_tokens); + let parent_hash = tracking.prompt_hashes.last().copied(); + let output_hashes = self + .output_hashes + .get_or_insert_with(|| CanonicalOutputTracker::new(request, block_size, is_eagle)); + output_hashes.set_prompt_parent(parent_hash); + self.cache_history = Some(tracking); + } + pub(super) async fn acquire_approximate_lru( &mut self, hashes: RoutingDecisionHashes, @@ -730,8 +809,11 @@ impl RequestGuard { return Ok(()); }; let mode = lease.acquire(hashes, private_blocks).await?; + self.approximate_lru_active = mode == ApproximateAcquireMode::Lru; if mode != ApproximateAcquireMode::Lru { - self.output_hashes = None; + if self.cache_history.is_none() { + self.output_hashes = None; + } return Ok(()); } if let Some(output_hashes) = self.output_hashes.as_mut() { @@ -785,29 +867,41 @@ impl RequestGuard { } } - if self + let track_approximate = self .cleanup .lifecycle() .is_some_and(RequestAttemptLease::is_active) - && let (Some(data), Some(output_hashes), Some(lease)) = ( - item.data.as_ref(), - self.output_hashes.as_mut(), - self.approximate_lru.as_ref(), - ) - && let Some(materialized) = - output_hashes.observe(data.index.unwrap_or(0), &data.token_ids) - && let Err(error) = lease.materialize( - materialized.parent_hash, - materialized.blocks, - materialized.start_position, - materialized.private_blocks, - ) - { - tracing::warn!( - request_id = self.cleanup.context_id().unwrap_or("stateless"), - %error, - "Failed to materialize approximate LRU output blocks" - ); + && self.approximate_lru_active + && self.approximate_lru.is_some(); + let materialized = (track_approximate || self.cache_history.is_some()) + .then(|| { + let data = item.data.as_ref()?; + self.output_hashes + .as_mut()? + .observe(data.index.unwrap_or(0), &data.token_ids) + }) + .flatten(); + if let Some(materialized) = materialized { + if let Some(history) = self.cache_history.as_mut() { + history + .output_hashes + .extend(materialized.blocks.iter().map(|block| block.sequence_hash)); + } + if track_approximate + && let Some(lease) = self.approximate_lru.as_ref() + && let Err(error) = lease.materialize( + materialized.parent_hash, + materialized.blocks, + materialized.start_position, + materialized.private_blocks, + ) + { + tracing::warn!( + request_id = self.cleanup.context_id().unwrap_or("stateless"), + %error, + "Failed to materialize approximate LRU output blocks" + ); + } } self.observability.observe_tokens(new_tokens); let cumulative_osl = self.observability.cumulative_osl(); @@ -843,6 +937,7 @@ impl RequestGuard { lifecycle.complete(); } // Metrics must observe the completed request before cleanup releases its state. + self.finish_cache_history(true); self.observability .record_metrics(self.record_itl_at_completion); self.cleanup.finish().await; @@ -869,18 +964,75 @@ impl RequestGuard { if let Some(lifecycle) = self.cleanup.request_lifecycle_mut() { lifecycle.abort(error.map(crate::protocols::common::preprocessor::owned_abort_error)); } + self.finish_cache_history(false); self.cleanup.finish().await; } + + fn finish_cache_history(&mut self, record_completed: bool) { + let Some(finalization) = finalize_cache_history(&mut self.cache_history, record_completed) + else { + return; + }; + if record_completed { + self.observability + .request_metrics() + .observe_cache_history_complete( + finalization.prompt_tokens, + finalization.previously_seen_tokens, + ); + } else { + self.observability + .request_metrics() + .observe_cache_history_incomplete(); + } + if let Some(stats) = finalization.retained { + self.observability + .request_metrics() + .set_cache_history_retained(stats); + } + } } impl Drop for RequestGuard { fn drop(&mut self) { + self.finish_cache_history(false); self.observability .record_metrics(self.record_itl_at_completion); // RequestCleanup drops immediately afterward and performs resource cleanup. } } +#[cfg(test)] +mod cache_history_tests { + use super::*; + + #[test] + fn completed_request_observes_once_and_records_membership() { + let history = Arc::new(CacheHistory::with_capacity(4, 8)); + let mut tracking = Some(CacheHistoryTracking::new(history.clone(), vec![10], 8)); + + let first = finalize_cache_history(&mut tracking, true).unwrap(); + assert_eq!(first.prompt_tokens, 8); + assert_eq!(first.previously_seen_tokens, 0); + assert!(first.retained.is_some()); + assert_eq!(history.previously_computed_tokens(&[10]), 8); + assert!(finalize_cache_history(&mut tracking, true).is_none()); + } + + #[test] + fn aborted_request_observes_once_without_recording_membership() { + let history = Arc::new(CacheHistory::with_capacity(4, 8)); + let mut tracking = Some(CacheHistoryTracking::new(history.clone(), vec![10], 8)); + + let first = finalize_cache_history(&mut tracking, false).unwrap(); + assert_eq!(first.prompt_tokens, 8); + assert_eq!(first.previously_seen_tokens, 0); + assert!(first.retained.is_none()); + assert_eq!(history.previously_computed_tokens(&[10]), 0); + assert!(finalize_cache_history(&mut tracking, false).is_none()); + } +} + #[cfg(test)] mod output_block_tests { use super::OutputBlockTracker; @@ -1117,6 +1269,7 @@ mod prefill_start_tests { ) .unwrap(), overlap_blocks_lost: hist_vec("overlap_blocks_lost"), + cache_history: None, }) } diff --git a/lib/llm/src/kv_router/routing_host/tests.rs b/lib/llm/src/kv_router/routing_host/tests.rs index 83d10187f890..e96e03118eaa 100644 --- a/lib/llm/src/kv_router/routing_host/tests.rs +++ b/lib/llm/src/kv_router/routing_host/tests.rs @@ -58,6 +58,100 @@ fn request() -> PreprocessedRequest { .unwrap() } +#[tokio::test] +#[serial_test::serial] +async fn cache_history_replica_sync_requires_both_flags() { + use dynamo_runtime::discovery::{DiscoveryQuery, EventChannelQuery}; + + for history_enabled in [true, false] { + for replica_sync in [true, false] { + temp_env::async_with_vars( + [( + cache_history::CACHE_REUSE_HISTORY_ENABLED_ENV, + Some(if history_enabled { "true" } else { "false" }), + )], + async { + let runtime = Runtime::from_current().unwrap(); + let distributed = + DistributedRuntime::new(runtime, DistributedConfig::process_local()) + .await + .unwrap(); + let endpoint = distributed + .namespace(format!("history-gate-{}", uuid::Uuid::new_v4())) + .unwrap() + .component("workers") + .unwrap() + .endpoint("generate"); + let query = DiscoveryQuery::EventChannels(EventChannelQuery::endpoint_topic( + endpoint.id(), + "cache-history-v1", + )); + let client = endpoint.client().await.unwrap(); + let (_workers_tx, workers) = watch::channel(HashMap::new()); + let config = KvRouterConfig { + skip_initial_worker_wait: true, + use_kv_events: false, + router_track_active_blocks: false, + router_replica_sync: replica_sync, + ..Default::default() + }; + let chooser = KvRouter::new( + endpoint, + client.clone(), + workers, + None, + 16, + SelectionPolicySource::Registry, + Some(config), + None, + "decode", + None, + false, + None, + None, + ) + .await + .unwrap(); + let inner = PushRouter::from_client(client, RouterMode::KV) + .await + .unwrap(); + let host = RoutingHost::new(inner, Arc::new(chooser), None).unwrap(); + assert_eq!(host.cache_history.is_some(), history_enabled); + let expected = usize::from(history_enabled && replica_sync); + tokio::time::timeout(Duration::from_secs(5), async { + let start = Instant::now(); + loop { + let count = distributed + .discovery() + .list(query.clone()) + .await + .unwrap() + .len(); + if expected == 0 { + assert_eq!( + count, 0, + "history={history_enabled}, sync={replica_sync}" + ); + if start.elapsed() >= Duration::from_millis(200) { + break; + } + } else if count == expected { + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("cache history publisher did not follow the enable flags"); + drop(host); + distributed.shutdown(); + }, + ) + .await; + } + } +} + async fn test_load_context(client: &Client) -> Arc { RoutingLoadContext::start( client.clone(), diff --git a/lib/runtime/src/metrics/prometheus_names.rs b/lib/runtime/src/metrics/prometheus_names.rs index 41d125f1c5eb..37733480b234 100644 --- a/lib/runtime/src/metrics/prometheus_names.rs +++ b/lib/runtime/src/metrics/prometheus_names.rs @@ -220,6 +220,34 @@ pub mod frontend_service { /// Effective KV overlap blocks lost by non-max-overlap selections pub const OVERLAP_BLOCKS_LOST: &str = "overlap_blocks_lost"; + /// Input tokens for cache-reuse observations started by the router + pub const CACHE_LOSS_OBSERVATION_INPUT_TOKENS_TOTAL: &str = + "cache_loss_observation_input_tokens_total"; + + /// Raw token observations at each cache-reuse funnel stage + pub const CACHE_LOSS_FUNNEL_TOKENS_TOTAL: &str = "cache_loss_funnel_tokens_total"; + + /// Prompt-prefix tokens previously seen in this router's bounded history + pub const KV_HISTORY_CACHED_PREFIX_TOKENS_TOTAL: &str = "kv_history_cached_prefix_tokens_total"; + + /// Cache-reuse observations by completion status + pub const CACHE_LOSS_OBSERVATIONS_TOTAL: &str = "cache_loss_observations_total"; + + /// Distinct canonical block hashes retained by cache history + pub const CACHE_LOSS_HISTORY_UNIQUE_HASHES: &str = "cache_loss_history_unique_hashes"; + + /// Tokens represented by distinct cache-history entries + pub const CACHE_LOSS_HISTORY_REPRESENTED_TOKENS: &str = "cache_loss_history_represented_tokens"; + + /// Estimated bytes used by retained cache-history entries + pub const CACHE_LOSS_HISTORY_ESTIMATED_BYTES: &str = "cache_loss_history_estimated_bytes"; + + /// Configured distinct-block capacity of cache history + pub const CACHE_LOSS_HISTORY_CAPACITY_BLOCKS: &str = "cache_loss_history_capacity_blocks"; + + /// Configured byte budget of cache history + pub const CACHE_LOSS_HISTORY_CAPACITY_BYTES: &str = "cache_loss_history_capacity_bytes"; + /// Number of cached tokens (prefix cache hits) per request pub const CACHED_TOKENS: &str = "cached_tokens"; @@ -1001,6 +1029,18 @@ pub fn clamp_u64_to_i64(value: u64) -> i64 { mod tests { use super::*; + #[test] + fn test_cache_history_prefix_metric_name() { + assert_eq!( + build_component_metric_name(&format!( + "{}{}", + router_request::METRIC_PREFIX, + frontend_service::KV_HISTORY_CACHED_PREFIX_TOKENS_TOTAL + )), + "dynamo_component_router_kv_history_cached_prefix_tokens_total" + ); + } + #[test] fn test_sanitize_frontend_prometheus_prefix() { // Test that valid prefixes remain unchanged