From c21375b5e5436f4a1b812fc6ce1b2bb71910183a Mon Sep 17 00:00:00 2001 From: Simo Lin Date: Sat, 28 Feb 2026 17:58:18 -0800 Subject: [PATCH 1/6] feat(policy): add event-driven routing tier to CacheAwarePolicy Add a new highest-priority routing tier to CacheAwarePolicy that uses PositionalIndexer overlap scoring from KvEventMonitor when available. This enables precise KV cache hit routing based on real backend cache state rather than approximate prefix matching. Three-tier dispatch order (gRPC path): 1. Event-driven: PositionalIndexer overlap scoring (when monitor has data for the model) 2. Approximate token tree: existing TokenTree prefix matching (fallback) 3. Min-load: when workers are imbalanced (unchanged) HTTP path remains unchanged (string tree). What changed: - policies/cache_aware.rs: add kv_monitor field, set_kv_event_monitor setter, try_event_driven_routing and select_worker_with_events methods, update select_worker dispatch to 3-tier, add 9 new tests - policies/mod.rs: add block_size to CacheAwareConfig (default 16) - core/kv_event_monitor.rs: add manual Debug impl - config/types.rs: add block_size to PolicyConfig::CacheAware with serde default for backward compat - config/builder.rs, validation.rs, factory.rs, main.rs, bindings/python/src/lib.rs, test files: propagate block_size Refs: Phase 1, Task 6 of KV cache event-driven routing Signed-off-by: Simo Lin --- bindings/python/src/lib.rs | 1 + model_gateway/src/config/builder.rs | 1 + model_gateway/src/config/types.rs | 13 + model_gateway/src/config/validation.rs | 5 + model_gateway/src/core/kv_event_monitor.rs | 11 +- model_gateway/src/main.rs | 1 + model_gateway/src/policies/cache_aware.rs | 430 +++++++++++++++++- model_gateway/src/policies/factory.rs | 3 + model_gateway/src/policies/mod.rs | 5 + .../cache_aware_backward_compat_test.rs | 2 + .../tests/routing/test_pd_routing.rs | 1 + 11 files changed, 466 insertions(+), 7 deletions(-) diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs index a1038ca46b..8800642044 100644 --- a/bindings/python/src/lib.rs +++ b/bindings/python/src/lib.rs @@ -477,6 +477,7 @@ impl Router { balance_rel_threshold: self.balance_rel_threshold, eviction_interval_secs: self.eviction_interval_secs, max_tree_size: self.max_tree_size, + block_size: 16, }, PolicyType::PowerOfTwo => ConfigPolicyConfig::PowerOfTwo { load_check_interval_secs: 5, diff --git a/model_gateway/src/config/builder.rs b/model_gateway/src/config/builder.rs index 72e9675991..0b3ffb901f 100644 --- a/model_gateway/src/config/builder.rs +++ b/model_gateway/src/config/builder.rs @@ -122,6 +122,7 @@ impl RouterConfigBuilder { balance_rel_threshold, eviction_interval_secs, max_tree_size, + block_size: 16, }; self } diff --git a/model_gateway/src/config/types.rs b/model_gateway/src/config/types.rs index 4643c050d8..6735b05a7d 100644 --- a/model_gateway/src/config/types.rs +++ b/model_gateway/src/config/types.rs @@ -250,6 +250,8 @@ pub enum PolicyConfig { balance_rel_threshold: f32, eviction_interval_secs: u64, max_tree_size: usize, + #[serde(default = "default_block_size")] + block_size: usize, }, #[serde(rename = "power_of_two")] @@ -307,6 +309,10 @@ pub enum PolicyConfig { }, } +fn default_block_size() -> usize { + 16 +} + fn default_prefix_token_count() -> usize { 256 } @@ -773,6 +779,7 @@ mod tests { balance_rel_threshold: 1.5, eviction_interval_secs: 300, max_tree_size: 1000, + block_size: 16, }; assert_eq!(cache_aware.name(), "cache_aware"); @@ -794,6 +801,7 @@ mod tests { balance_rel_threshold: 1.5, eviction_interval_secs: 300, max_tree_size: 1000, + block_size: 16, }; let json = serde_json::to_string(&cache_aware).unwrap(); assert!(json.contains("\"type\":\"cache_aware\"")); @@ -816,6 +824,7 @@ mod tests { balance_rel_threshold: 2.0, eviction_interval_secs: 600, max_tree_size: 5000, + block_size: 16, }; match cache_aware { @@ -825,6 +834,7 @@ mod tests { balance_rel_threshold, eviction_interval_secs, max_tree_size, + .. } => { assert!((cache_threshold - 0.75).abs() < 0.0001); assert_eq!(balance_abs_threshold, 20); @@ -1218,6 +1228,7 @@ mod tests { balance_rel_threshold: 1.1, eviction_interval_secs: 60, max_tree_size: 1000, + block_size: 16, }), decode_policy: Some(PolicyConfig::PowerOfTwo { load_check_interval_secs: 60, @@ -1248,6 +1259,7 @@ mod tests { balance_rel_threshold: 1.1, eviction_interval_secs: 60, max_tree_size: 1000, + block_size: 16, }), decode_policy: None, }; @@ -1304,6 +1316,7 @@ mod tests { balance_rel_threshold: 1.5, eviction_interval_secs: 300, max_tree_size: 2000, + block_size: 16, }; match pd.get_prefill_policy(&main_policy) { diff --git a/model_gateway/src/config/validation.rs b/model_gateway/src/config/validation.rs index aa226423b6..f8373aa79d 100644 --- a/model_gateway/src/config/validation.rs +++ b/model_gateway/src/config/validation.rs @@ -173,6 +173,7 @@ impl ConfigValidator { balance_rel_threshold, eviction_interval_secs, max_tree_size, + .. } => { if !(0.0..=1.0).contains(cache_threshold) { return Err(ConfigError::InvalidValue { @@ -742,6 +743,7 @@ mod tests { balance_rel_threshold: 1.1, eviction_interval_secs: 60, max_tree_size: 1000, + block_size: 16, }, ); @@ -761,6 +763,7 @@ mod tests { balance_rel_threshold: 1.1, eviction_interval_secs: 60, max_tree_size: 1000, + block_size: 16, }, ); @@ -815,6 +818,7 @@ mod tests { balance_rel_threshold: 1.1, eviction_interval_secs: 60, max_tree_size: 1000, + block_size: 16, }, ); @@ -859,6 +863,7 @@ mod tests { balance_rel_threshold: 1.1, eviction_interval_secs: 60, max_tree_size: 1000, + block_size: 16, }), decode_policy: Some(PolicyConfig::PowerOfTwo { load_check_interval_secs: 60, diff --git a/model_gateway/src/core/kv_event_monitor.rs b/model_gateway/src/core/kv_event_monitor.rs index a81262b775..5057154122 100644 --- a/model_gateway/src/core/kv_event_monitor.rs +++ b/model_gateway/src/core/kv_event_monitor.rs @@ -10,7 +10,7 @@ //! - `on_worker_removed` — aborts task, removes worker from indexer //! - `stop` — aborts all tasks, clears state -use std::{collections::HashMap, sync::Arc, time::Duration}; +use std::{collections::HashMap, fmt, sync::Arc, time::Duration}; use dashmap::DashMap; use kv_index::{compute_content_hash, ApplyError, PositionalIndexer, SequenceHash, StoredBlock}; @@ -408,6 +408,15 @@ impl Drop for KvEventMonitor { } } +impl fmt::Debug for KvEventMonitor { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("KvEventMonitor") + .field("models", &self.indexers.len()) + .field("jump_size", &self.jump_size) + .finish() + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/model_gateway/src/main.rs b/model_gateway/src/main.rs index 5660424869..550a4d1e25 100644 --- a/model_gateway/src/main.rs +++ b/model_gateway/src/main.rs @@ -754,6 +754,7 @@ impl CliArgs { balance_rel_threshold: self.balance_rel_threshold, eviction_interval_secs: self.eviction_interval, max_tree_size: self.max_tree_size, + block_size: 16, }, "power_of_two" => PolicyConfig::PowerOfTwo { load_check_interval_secs: 5, diff --git a/model_gateway/src/policies/cache_aware.rs b/model_gateway/src/policies/cache_aware.rs index fdbbf3b15f..ef844670cd 100644 --- a/model_gateway/src/policies/cache_aware.rs +++ b/model_gateway/src/policies/cache_aware.rs @@ -62,7 +62,7 @@ use std::sync::Arc; use dashmap::DashMap; -use kv_index::{TokenTree, Tree}; +use kv_index::{compute_request_content_hashes, PositionalIndexer, TokenTree, Tree}; use rand::Rng; use smg_mesh::{OptionalMeshSyncManager, TreeInsertOp, TreeOperation}; use tracing::{debug, warn}; @@ -71,7 +71,7 @@ use super::{ get_healthy_worker_indices, normalize_model_key, utils::PeriodicTask, CacheAwareConfig, LoadBalancingPolicy, SelectWorkerInfo, }; -use crate::core::{Worker, UNKNOWN_MODEL_ID}; +use crate::core::{KvEventMonitor, Worker, UNKNOWN_MODEL_ID}; /// Cache-aware routing policy /// @@ -93,6 +93,10 @@ pub struct CacheAwarePolicy { token_trees: Arc>>, mesh_sync: OptionalMeshSyncManager, _eviction_task: Option, + /// Event-driven KV cache monitor for overlap scoring (gRPC workers only). + /// Set via `set_kv_event_monitor`. When present and the indexer has data for + /// a model, event-driven routing takes priority over approximate trees. + kv_monitor: Option>, } impl CacheAwarePolicy { @@ -148,6 +152,7 @@ impl CacheAwarePolicy { token_trees, mesh_sync: None, _eviction_task: eviction_task, + kv_monitor: None, } } @@ -159,6 +164,11 @@ impl CacheAwarePolicy { } } + /// Set event-driven KV cache monitor (can be called after construction) + pub fn set_kv_event_monitor(&mut self, monitor: Option>) { + self.kv_monitor = monitor; + } + /// Initialize the trees with worker URLs (used only during initial setup) /// Initializes both string trees (HTTP) and token trees (gRPC) for each model. pub fn init_workers(&self, workers: &[Arc]) { @@ -451,13 +461,17 @@ impl LoadBalancingPolicy for CacheAwarePolicy { return self.select_worker_min_load(workers, info, &healthy_indices, model_id); } - // Use cache-aware routing when balanced - // gRPC requests use token tree, HTTP requests use string tree + // Use cache-aware routing when balanced (three tiers) if let Some(tokens) = request_tokens { - // gRPC path: use token-based tree + // gRPC path: try event-driven first, then approximate token tree + if let Some(idx) = + self.try_event_driven_routing(workers, tokens, &healthy_indices, model_id) + { + return Some(idx); + } self.select_worker_with_tokens(workers, tokens, &healthy_indices, model_id) } else { - // HTTP path: use string-based tree + // HTTP path: approximate string tree (unchanged) let text = request_text.unwrap_or(""); self.select_worker_with_text(workers, text, &healthy_indices, model_id) } @@ -490,6 +504,78 @@ impl LoadBalancingPolicy for CacheAwarePolicy { // Private helper methods for select_worker impl CacheAwarePolicy { + /// Try event-driven routing via KvEventMonitor overlap scoring. + /// + /// Returns `Some(idx)` if the monitor has an indexer with data for this model + /// and a worker was selected. Returns `None` to fall back to approximate trees. + fn try_event_driven_routing( + &self, + workers: &[Arc], + tokens: &[u32], + healthy_indices: &[usize], + model_id: &str, + ) -> Option { + let monitor = self.kv_monitor.as_ref()?; + let indexer = monitor.get_indexer(model_id)?; + if indexer.current_size() == 0 { + return None; + } + self.select_worker_with_events(workers, tokens, healthy_indices, &indexer) + } + + /// Select worker using event-driven overlap scoring (gRPC path, Tier 1). + /// + /// Computes content hashes from request tokens, queries the PositionalIndexer + /// for overlap scores, and selects the worker with the best match. Tie-breaks + /// by load (lower wins) then tree size (smaller wins). + /// + /// Returns `None` if the request is too short for any full block or no workers + /// have cached data matching the request — caller should fall back to approximate trees. + fn select_worker_with_events( + &self, + workers: &[Arc], + tokens: &[u32], + healthy_indices: &[usize], + indexer: &PositionalIndexer, + ) -> Option { + let content_hashes = compute_request_content_hashes(tokens, self.config.block_size); + if content_hashes.is_empty() { + return None; + } + + let overlap = indexer.find_matches(&content_hashes); + if overlap.scores.is_empty() { + return None; + } + + let best_idx = healthy_indices + .iter() + .max_by(|&&a, &&b| { + let score_a = overlap.scores.get(workers[a].url()).copied().unwrap_or(0); + let score_b = overlap.scores.get(workers[b].url()).copied().unwrap_or(0); + score_a + .cmp(&score_b) + .then_with(|| workers[b].load().cmp(&workers[a].load())) + .then_with(|| { + let size_a = overlap + .tree_sizes + .get(workers[a].url()) + .copied() + .unwrap_or(0); + let size_b = overlap + .tree_sizes + .get(workers[b].url()) + .copied() + .unwrap_or(0); + size_b.cmp(&size_a) + }) + }) + .copied()?; + + workers[best_idx].increment_processed(); + Some(best_idx) + } + /// Select worker using token-based tree (gRPC path) fn select_worker_with_tokens( &self, @@ -620,6 +706,8 @@ impl Default for CacheAwarePolicy { #[cfg(test)] mod tests { + use kv_index::{compute_content_hash, SequenceHash, StoredBlock}; + use super::*; use crate::core::{BasicWorkerBuilder, WorkerType}; @@ -693,6 +781,7 @@ mod tests { balance_rel_threshold: 2.0, eviction_interval_secs: 0, // Disable eviction thread max_tree_size: 10000, + block_size: 16, }); let worker1 = BasicWorkerBuilder::new("http://w1:8000") @@ -973,4 +1062,333 @@ mod tests { .unwrap(); assert_eq!(idx, 0); } + + // ----------------------------------------------------------------------- + // Event-driven routing tests + // ----------------------------------------------------------------------- + + /// Helper: create a PositionalIndexer and store blocks for a worker. + /// `token_chunks` is a list of token-id slices — each becomes one block. + fn setup_indexer_with_blocks( + worker_url: &str, + token_chunks: &[&[u32]], + jump_size: usize, + ) -> Arc { + let indexer = Arc::new(PositionalIndexer::new(jump_size)); + let blocks: Vec = token_chunks + .iter() + .enumerate() + .map(|(i, tokens)| StoredBlock { + seq_hash: SequenceHash(i as u64 + 1), + content_hash: compute_content_hash(tokens), + }) + .collect(); + indexer.apply_stored(worker_url, &blocks, None).unwrap(); + indexer + } + + fn test_config() -> CacheAwareConfig { + CacheAwareConfig { + eviction_interval_secs: 0, + block_size: 4, // small block size for easy test setup + ..Default::default() + } + } + + #[test] + fn test_event_driven_selects_best_overlap() { + let policy = CacheAwarePolicy::with_config(test_config()); + let workers: Vec> = vec![ + Arc::new( + BasicWorkerBuilder::new("http://w1:8000") + .worker_type(WorkerType::Regular) + .build(), + ), + Arc::new( + BasicWorkerBuilder::new("http://w2:8000") + .worker_type(WorkerType::Regular) + .build(), + ), + ]; + policy.init_workers(&workers); + + // Store 4 blocks for w1: tokens [1..16] in blocks of 4 + let indexer = setup_indexer_with_blocks( + "http://w1:8000", + &[ + &[1, 2, 3, 4], + &[5, 6, 7, 8], + &[9, 10, 11, 12], + &[13, 14, 15, 16], + ], + 4, + ); + + // Query with matching tokens — should select w1 + let result = policy.select_worker_with_events( + &workers, + &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16], + &[0, 1], + &indexer, + ); + assert_eq!(result, Some(0)); // w1 + } + + #[test] + fn test_event_driven_no_overlap_returns_none() { + let policy = CacheAwarePolicy::with_config(test_config()); + let workers: Vec> = vec![Arc::new( + BasicWorkerBuilder::new("http://w1:8000") + .worker_type(WorkerType::Regular) + .build(), + )]; + policy.init_workers(&workers); + + // Store blocks for tokens [1..8] + let indexer = + setup_indexer_with_blocks("http://w1:8000", &[&[1, 2, 3, 4], &[5, 6, 7, 8]], 4); + + // Query with completely different tokens — no overlap + let result = policy.select_worker_with_events( + &workers, + &[100, 200, 300, 400, 500, 600, 700, 800], + &[0], + &indexer, + ); + assert_eq!(result, None); // no match → falls back + } + + #[test] + fn test_event_driven_load_tiebreak() { + let policy = CacheAwarePolicy::with_config(test_config()); + + let w1 = BasicWorkerBuilder::new("http://w1:8000") + .worker_type(WorkerType::Regular) + .build(); + let w2 = BasicWorkerBuilder::new("http://w2:8000") + .worker_type(WorkerType::Regular) + .build(); + + // Give w1 higher load + for _ in 0..10 { + w1.increment_load(); + } + + let workers: Vec> = vec![Arc::new(w1), Arc::new(w2)]; + policy.init_workers(&workers); + + // Store same blocks for both workers (equal overlap) + let indexer = Arc::new(PositionalIndexer::new(4)); + let blocks = vec![StoredBlock { + seq_hash: SequenceHash(1), + content_hash: compute_content_hash(&[1, 2, 3, 4]), + }]; + indexer + .apply_stored("http://w1:8000", &blocks, None) + .unwrap(); + let blocks2 = vec![StoredBlock { + seq_hash: SequenceHash(1), + content_hash: compute_content_hash(&[1, 2, 3, 4]), + }]; + indexer + .apply_stored("http://w2:8000", &blocks2, None) + .unwrap(); + + // Equal overlap → tie-break by load → w2 wins (lower load) + let result = policy.select_worker_with_events(&workers, &[1, 2, 3, 4], &[0, 1], &indexer); + assert_eq!(result, Some(1)); // w2 (lower load) + } + + #[test] + fn test_event_driven_tree_size_tiebreak() { + let policy = CacheAwarePolicy::with_config(test_config()); + let workers: Vec> = vec![ + Arc::new( + BasicWorkerBuilder::new("http://w1:8000") + .worker_type(WorkerType::Regular) + .build(), + ), + Arc::new( + BasicWorkerBuilder::new("http://w2:8000") + .worker_type(WorkerType::Regular) + .build(), + ), + ]; + policy.init_workers(&workers); + + let indexer = Arc::new(PositionalIndexer::new(4)); + + // Both workers have block [1,2,3,4] (equal overlap, equal load) + let block = vec![StoredBlock { + seq_hash: SequenceHash(1), + content_hash: compute_content_hash(&[1, 2, 3, 4]), + }]; + indexer + .apply_stored("http://w1:8000", &block, None) + .unwrap(); + + // w2 has the same block plus extra blocks → larger tree + let block2 = vec![StoredBlock { + seq_hash: SequenceHash(1), + content_hash: compute_content_hash(&[1, 2, 3, 4]), + }]; + indexer + .apply_stored("http://w2:8000", &block2, None) + .unwrap(); + let extra = vec![StoredBlock { + seq_hash: SequenceHash(2), + content_hash: compute_content_hash(&[5, 6, 7, 8]), + }]; + indexer + .apply_stored("http://w2:8000", &extra, Some(SequenceHash(1))) + .unwrap(); + + // Equal overlap, equal load → tie-break by tree size → w1 wins (smaller) + let result = policy.select_worker_with_events(&workers, &[1, 2, 3, 4], &[0, 1], &indexer); + assert_eq!(result, Some(0)); // w1 (smaller tree) + } + + #[test] + fn test_event_driven_short_request_returns_none() { + let policy = CacheAwarePolicy::with_config(test_config()); // block_size=4 + let workers: Vec> = vec![Arc::new( + BasicWorkerBuilder::new("http://w1:8000") + .worker_type(WorkerType::Regular) + .build(), + )]; + + let indexer = setup_indexer_with_blocks("http://w1:8000", &[&[1, 2, 3, 4]], 4); + + // Request shorter than block_size → no full blocks → None + let result = policy.select_worker_with_events(&workers, &[1, 2, 3], &[0], &indexer); + assert_eq!(result, None); + } + + #[test] + fn test_three_tier_dispatch_prefers_events() { + let policy = CacheAwarePolicy::with_config(test_config()); + let workers: Vec> = vec![ + Arc::new( + BasicWorkerBuilder::new("http://w1:8000") + .worker_type(WorkerType::Regular) + .build(), + ), + Arc::new( + BasicWorkerBuilder::new("http://w2:8000") + .worker_type(WorkerType::Regular) + .build(), + ), + ]; + policy.init_workers(&workers); + + // Event indexer: only w1 has cached blocks + let indexer = + setup_indexer_with_blocks("http://w1:8000", &[&[1, 2, 3, 4], &[5, 6, 7, 8]], 4); + + // Event-driven path selects w1 (it has the cached blocks) + let result = policy.select_worker_with_events( + &workers, + &[1, 2, 3, 4, 5, 6, 7, 8], + &[0, 1], + &indexer, + ); + assert_eq!(result, Some(0)); // w1 via event-driven + + // Without event data (no overlap), returns None → would fall back + let result = + policy.select_worker_with_events(&workers, &[100, 200, 300, 400], &[0, 1], &indexer); + assert_eq!(result, None); // no match → caller falls back to approximate tree + } + + #[test] + fn test_three_tier_empty_indexer_falls_back() { + let policy = CacheAwarePolicy::with_config(test_config()); + let workers: Vec> = vec![ + Arc::new( + BasicWorkerBuilder::new("http://w1:8000") + .worker_type(WorkerType::Regular) + .build(), + ), + Arc::new( + BasicWorkerBuilder::new("http://w2:8000") + .worker_type(WorkerType::Regular) + .build(), + ), + ]; + policy.init_workers(&workers); + + // Empty indexer — no data + let indexer = Arc::new(PositionalIndexer::new(4)); + + // Should return None → caller falls back to approximate tree + let result = policy.select_worker_with_events( + &workers, + &[1, 2, 3, 4, 5, 6, 7, 8], + &[0, 1], + &indexer, + ); + assert_eq!(result, None); + } + + #[test] + fn test_set_kv_event_monitor() { + let mut policy = CacheAwarePolicy::with_config(test_config()); + + // Initially no monitor + assert!(policy.kv_monitor.is_none()); + + // Set monitor + let monitor = Arc::new(KvEventMonitor::new(Some(4))); + policy.set_kv_event_monitor(Some(Arc::clone(&monitor))); + assert!(policy.kv_monitor.is_some()); + + // get_indexer returns None for unknown model + assert!(monitor.get_indexer("nonexistent").is_none()); + + // Clear monitor + policy.set_kv_event_monitor(None); + assert!(policy.kv_monitor.is_none()); + } + + #[test] + fn test_event_driven_imbalanced_skips_events() { + let mut policy = CacheAwarePolicy::with_config(CacheAwareConfig { + balance_abs_threshold: 5, + balance_rel_threshold: 2.0, + eviction_interval_secs: 0, + block_size: 4, + ..Default::default() + }); + + let w1 = BasicWorkerBuilder::new("http://w1:8000") + .worker_type(WorkerType::Regular) + .build(); + let w2 = BasicWorkerBuilder::new("http://w2:8000") + .worker_type(WorkerType::Regular) + .build(); + + // Create heavy imbalance: w1 has 20 load, w2 has 0 + for _ in 0..20 { + w1.increment_load(); + } + + let workers: Vec> = vec![Arc::new(w1), Arc::new(w2)]; + policy.init_workers(&workers); + + // Even though we set up event monitor, imbalance check fires first + let monitor = Arc::new(KvEventMonitor::new(Some(4))); + policy.set_kv_event_monitor(Some(monitor)); + + // With imbalance, select_worker should pick min-load (w2), not event-driven + let idx = policy + .select_worker( + &workers, + &SelectWorkerInfo { + tokens: Some(&[1, 2, 3, 4]), + ..Default::default() + }, + ) + .unwrap(); + assert_eq!(idx, 1); // w2 (min load), regardless of event data + } } diff --git a/model_gateway/src/policies/factory.rs b/model_gateway/src/policies/factory.rs index ae3ad845e5..cc09768b09 100644 --- a/model_gateway/src/policies/factory.rs +++ b/model_gateway/src/policies/factory.rs @@ -25,6 +25,7 @@ impl PolicyFactory { balance_rel_threshold, eviction_interval_secs, max_tree_size, + block_size, } => { let config = CacheAwareConfig { cache_threshold: *cache_threshold, @@ -32,6 +33,7 @@ impl PolicyFactory { balance_rel_threshold: *balance_rel_threshold, eviction_interval_secs: *eviction_interval_secs, max_tree_size: *max_tree_size, + block_size: *block_size, }; Arc::new(CacheAwarePolicy::with_config(config)) } @@ -114,6 +116,7 @@ mod tests { balance_rel_threshold: 1.5, eviction_interval_secs: 30, max_tree_size: 1000, + block_size: 16, }); assert_eq!(policy.name(), "cache_aware"); diff --git a/model_gateway/src/policies/mod.rs b/model_gateway/src/policies/mod.rs index 07d68c6159..21663c6462 100644 --- a/model_gateway/src/policies/mod.rs +++ b/model_gateway/src/policies/mod.rs @@ -97,6 +97,10 @@ pub struct CacheAwareConfig { pub balance_rel_threshold: f32, pub eviction_interval_secs: u64, pub max_tree_size: usize, + /// Backend KV cache block size (tokens per block) for event-driven routing. + /// Used by `compute_request_content_hashes` to chunk request tokens into blocks. + /// Must match the backend's block size. Default: 16 (SGLang page size). + pub block_size: usize, } impl Default for CacheAwareConfig { @@ -107,6 +111,7 @@ impl Default for CacheAwareConfig { balance_rel_threshold: 1.1, eviction_interval_secs: 30, max_tree_size: 10000, + block_size: 16, } } } diff --git a/model_gateway/tests/routing/cache_aware_backward_compat_test.rs b/model_gateway/tests/routing/cache_aware_backward_compat_test.rs index c9f746cea3..f4c299fced 100644 --- a/model_gateway/tests/routing/cache_aware_backward_compat_test.rs +++ b/model_gateway/tests/routing/cache_aware_backward_compat_test.rs @@ -13,6 +13,7 @@ fn test_backward_compatibility_with_empty_model_id() { balance_rel_threshold: 1.5, eviction_interval_secs: 0, // Disable background eviction for testing max_tree_size: 100, + block_size: 16, }; let policy = CacheAwarePolicy::with_config(config); @@ -62,6 +63,7 @@ fn test_mixed_model_ids() { balance_rel_threshold: 1.5, eviction_interval_secs: 0, max_tree_size: 100, + block_size: 16, }; let policy = CacheAwarePolicy::with_config(config); diff --git a/model_gateway/tests/routing/test_pd_routing.rs b/model_gateway/tests/routing/test_pd_routing.rs index 4489fd5f03..c8a8adfd80 100644 --- a/model_gateway/tests/routing/test_pd_routing.rs +++ b/model_gateway/tests/routing/test_pd_routing.rs @@ -165,6 +165,7 @@ mod pd_routing_unit_tests { balance_rel_threshold: 1.2, eviction_interval_secs: 60, max_tree_size: 1000000, + block_size: 16, }, ), ( From 3270d7677ec917daef36f793b073235e946efe33 Mon Sep 17 00:00:00 2001 From: Simo Lin Date: Sun, 1 Mar 2026 19:46:44 -0800 Subject: [PATCH 2/6] refactor(policy): restructure event-driven routing as mutually exclusive type MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Restructure CacheAwarePolicy from cascading tiers (event → token tree fallback) to three mutually exclusive routing types selected by worker connection mode and KV event availability: 1. Event-Driven (gRPC + KV events): PositionalIndexer overlap scoring 2. Approximate Token Tree (gRPC, no events): TokenTree prefix matching 3. Approximate String Tree (HTTP): Tree character prefix matching What changed: - model_gateway/src/policies/cache_aware.rs: - Replace try_event_driven_routing + select_worker_with_events with has_event_indexer, select_worker_event_driven, and score_overlap - Dispatch logic uses if/else (mutually exclusive) instead of fallthrough - Event-driven path is self-contained: overlap match or min-load fallback, never falls through to approximate token tree - Rewrite file header to document all three types - Add 12 new tests: 6 score_overlap unit tests (best match, no match, load tiebreak, tree size tiebreak, short request, partial match) and 6 select_worker_event_driven integration tests (overlap selects cached, no overlap uses min-load, short request min-load, no monitor token tree, set/clear monitor, imbalanced skips events) - model_gateway/src/core/kv_event_monitor.rs: - Change indexers field to pub(crate) for test access from cache_aware tests Why: The cascading tier design (try event-driven, fall back to token tree) was incorrect — when event data exists for a model, approximate tree routing would produce stale/conflicting decisions. The three types are mutually exclusive by design: if we have real cache state from KV events, we should not also consult an approximate tree. When event-driven has no overlap (cold start, novel tokens), min-load is the correct fallback. Signed-off-by: Simon Lin Signed-off-by: Simo Lin --- model_gateway/src/core/kv_event_monitor.rs | 2 +- model_gateway/src/policies/cache_aware.rs | 370 +++++++++++++++------ 2 files changed, 261 insertions(+), 111 deletions(-) diff --git a/model_gateway/src/core/kv_event_monitor.rs b/model_gateway/src/core/kv_event_monitor.rs index 5057154122..43ee727f69 100644 --- a/model_gateway/src/core/kv_event_monitor.rs +++ b/model_gateway/src/core/kv_event_monitor.rs @@ -38,7 +38,7 @@ const MAX_RECONNECT_DELAY_MS: u64 = 30_000; /// (one per `model_id`). Workers serving the same model share the same indexer. pub struct KvEventMonitor { /// Per-model positional indexers: model_id → shared indexer. - indexers: DashMap>, + pub(crate) indexers: DashMap>, /// Per-worker subscription handles: worker_url → subscription info. /// Mutex matches LoadMonitor pattern for atomic abort + remove. worker_handles: Mutex>, diff --git a/model_gateway/src/policies/cache_aware.rs b/model_gateway/src/policies/cache_aware.rs index ef844670cd..7e3a518e88 100644 --- a/model_gateway/src/policies/cache_aware.rs +++ b/model_gateway/src/policies/cache_aware.rs @@ -1,62 +1,44 @@ /* Cache-Aware Load Balancing Router - This router combines two strategies to optimize both cache utilization and request distribution: + When load is balanced, uses cache-aware routing. When imbalanced, uses + shortest-queue. A system is imbalanced when both: + (max - min) > abs_threshold AND max > rel_threshold * min - 1. Cache-Aware Routing (Approximate Tree) - 2. Load Balancing (Shortest Queue with Balance Thresholds) + Three types of cache-aware routing (mutually exclusive, selected by + worker connection mode and KV event availability): - The router dynamically switches between these strategies based on load conditions: - - Uses load balancing when the system is imbalanced - - Uses cache-aware routing when the system is balanced - - A system is considered imbalanced if both conditions are met: - 1. (max - min) > abs_threshold - 2. max > rel_threshold * min + 1. Event-Driven (gRPC + KV events) + ------------------------------------------- + Uses PositionalIndexer overlap scoring from KvEventMonitor. Routes based + on actual backend KV cache state. Selects the worker with the highest + overlap count; tie-breaks by load (lower) then tree size (smaller). + Falls back to min-load when no cache overlap exists. - Strategy Details: + 2. Approximate Token Tree (gRPC, no KV events) + ------------------------------------------- + Maintains a TokenTree per model tracking which token prefixes were routed + where. If match_rate > cache_threshold, routes to the best-matching worker. + Otherwise routes to the worker with the smallest tree (most cache capacity). - 1. Cache-Aware Routing (Approximate Tree) + 3. Approximate String Tree (HTTP) ------------------------------------------- - This strategy maintains an approximate radix tree for each worker based on request history, - eliminating the need for direct cache state queries. The tree stores raw text characters - instead of token IDs to avoid tokenization overhead. - - Process: - a. For each request, find the worker with the highest prefix match - b. If match rate > cache_threshold: - Route to the worker with highest match (likely has relevant data cached) - c. If match rate ≤ cache_threshold: - Route to the worker with smallest tree size (most available cache capacity) - d. Background maintenance: - Periodically evict least recently used leaf nodes to prevent memory overflow - - 2. Load Balancing (Shortest Queue) + Same algorithm as (2) but operates on raw text characters instead of + token IDs, avoiding tokenization overhead. + + Load Balancing (Shortest Queue) ------------------------------------------- - This strategy tracks pending request counts per worker and routes new requests - to the least busy worker when the system is detected to be imbalanced. + When the system is imbalanced, routes to the least busy worker regardless + of cache affinity. Configuration Parameters: ------------------------ - 1. cache_threshold: (float, 0.0 to 1.0) - Minimum prefix match ratio to use highest-match routing. - Below this threshold, routes to worker with most available cache space. - - 2. balance_abs_threshold: (integer) - Absolute difference threshold for load imbalance detection. - System is potentially imbalanced if (max_load - min_load) > abs_threshold - - 3. balance_rel_threshold: (float) - Relative ratio threshold for load imbalance detection. - System is potentially imbalanced if max_load > min_load * rel_threshold - Used in conjunction with abs_threshold to determine final imbalance state. - - 4. eviction_interval_secs: (integer) - Interval between LRU eviction cycles for the approximate trees. - - 5. max_tree_size: (integer) - Maximum nodes per tree. When exceeded, LRU leaf nodes are evicted - during the next eviction cycle. + cache_threshold: Min prefix match ratio for highest-match routing (0.0-1.0) + balance_abs_threshold: Absolute load diff threshold for imbalance detection + balance_rel_threshold: Relative load ratio threshold for imbalance detection + eviction_interval_secs: Interval between LRU eviction cycles + max_tree_size: Max nodes per approximate tree before eviction + block_size: Backend KV cache block size for event-driven routing */ use std::sync::Arc; @@ -461,17 +443,17 @@ impl LoadBalancingPolicy for CacheAwarePolicy { return self.select_worker_min_load(workers, info, &healthy_indices, model_id); } - // Use cache-aware routing when balanced (three tiers) + // Cache-aware routing when balanced — three types (mutually exclusive): + // 1. Event-driven: PositionalIndexer overlap scoring (gRPC + KV events) + // 2. Approximate token tree: TokenTree prefix matching (gRPC, no events) + // 3. Approximate string tree: Tree prefix matching (HTTP) if let Some(tokens) = request_tokens { - // gRPC path: try event-driven first, then approximate token tree - if let Some(idx) = - self.try_event_driven_routing(workers, tokens, &healthy_indices, model_id) - { - return Some(idx); + if self.has_event_indexer(model_id) { + self.select_worker_event_driven(workers, tokens, &healthy_indices, model_id) + } else { + self.select_worker_with_tokens(workers, tokens, &healthy_indices, model_id) } - self.select_worker_with_tokens(workers, tokens, &healthy_indices, model_id) } else { - // HTTP path: approximate string tree (unchanged) let text = request_text.unwrap_or(""); self.select_worker_with_text(workers, text, &healthy_indices, model_id) } @@ -504,34 +486,52 @@ impl LoadBalancingPolicy for CacheAwarePolicy { // Private helper methods for select_worker impl CacheAwarePolicy { - /// Try event-driven routing via KvEventMonitor overlap scoring. + /// Check if an event-driven indexer exists for this model. + /// Used to choose between event-driven and approximate tree routing. + fn has_event_indexer(&self, model_id: &str) -> bool { + self.kv_monitor + .as_ref() + .and_then(|m| m.get_indexer(model_id)) + .is_some() + } + + /// Event-driven routing: PositionalIndexer overlap scoring (Type 1). /// - /// Returns `Some(idx)` if the monitor has an indexer with data for this model - /// and a worker was selected. Returns `None` to fall back to approximate trees. - fn try_event_driven_routing( + /// Self-contained — when overlap is found, selects the worker with the best + /// cache match. When no overlap (cold start, novel tokens, short request), + /// falls back to min-load. Does NOT fall back to approximate token tree. + fn select_worker_event_driven( &self, workers: &[Arc], tokens: &[u32], healthy_indices: &[usize], model_id: &str, ) -> Option { - let monitor = self.kv_monitor.as_ref()?; - let indexer = monitor.get_indexer(model_id)?; - if indexer.current_size() == 0 { - return None; + let indexer = self.kv_monitor.as_ref()?.get_indexer(model_id)?; + + if let Some(idx) = self.score_overlap(workers, tokens, healthy_indices, &indexer) { + return Some(idx); } - self.select_worker_with_events(workers, tokens, healthy_indices, &indexer) + + // No cache overlap — min-load fallback (no token tree involved) + let min_idx = healthy_indices + .iter() + .min_by_key(|&&idx| workers[idx].load()) + .copied()?; + debug!( + worker = workers[min_idx].url(), + model_id, "Event-driven routing: no overlap, min-load fallback" + ); + workers[min_idx].increment_processed(); + Some(min_idx) } - /// Select worker using event-driven overlap scoring (gRPC path, Tier 1). - /// - /// Computes content hashes from request tokens, queries the PositionalIndexer - /// for overlap scores, and selects the worker with the best match. Tie-breaks - /// by load (lower wins) then tree size (smaller wins). + /// Score healthy workers by PositionalIndexer overlap and select the best. /// - /// Returns `None` if the request is too short for any full block or no workers - /// have cached data matching the request — caller should fall back to approximate trees. - fn select_worker_with_events( + /// Returns `Some(idx)` if at least one worker has cached blocks matching the + /// request. Returns `None` if the request is too short for a full block or + /// no workers have matching data. + fn score_overlap( &self, workers: &[Arc], tokens: &[u32], @@ -548,6 +548,7 @@ impl CacheAwarePolicy { return None; } + // Select worker with best overlap; tie-break by load (lower) then tree size (smaller) let best_idx = healthy_indices .iter() .max_by(|&&a, &&b| { @@ -572,6 +573,15 @@ impl CacheAwarePolicy { }) .copied()?; + debug!( + worker = workers[best_idx].url(), + score = overlap + .scores + .get(workers[best_idx].url()) + .copied() + .unwrap_or(0), + "Event-driven routing: overlap match" + ); workers[best_idx].increment_processed(); Some(best_idx) } @@ -1064,7 +1074,7 @@ mod tests { } // ----------------------------------------------------------------------- - // Event-driven routing tests + // Event-driven routing tests (Type 1: PositionalIndexer overlap scoring) // ----------------------------------------------------------------------- /// Helper: create a PositionalIndexer and store blocks for a worker. @@ -1095,8 +1105,10 @@ mod tests { } } + // -- score_overlap unit tests (scoring helper) -- + #[test] - fn test_event_driven_selects_best_overlap() { + fn test_score_overlap_selects_best_match() { let policy = CacheAwarePolicy::with_config(test_config()); let workers: Vec> = vec![ Arc::new( @@ -1125,7 +1137,7 @@ mod tests { ); // Query with matching tokens — should select w1 - let result = policy.select_worker_with_events( + let result = policy.score_overlap( &workers, &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16], &[0, 1], @@ -1135,7 +1147,7 @@ mod tests { } #[test] - fn test_event_driven_no_overlap_returns_none() { + fn test_score_overlap_no_match_returns_none() { let policy = CacheAwarePolicy::with_config(test_config()); let workers: Vec> = vec![Arc::new( BasicWorkerBuilder::new("http://w1:8000") @@ -1144,22 +1156,21 @@ mod tests { )]; policy.init_workers(&workers); - // Store blocks for tokens [1..8] let indexer = setup_indexer_with_blocks("http://w1:8000", &[&[1, 2, 3, 4], &[5, 6, 7, 8]], 4); - // Query with completely different tokens — no overlap - let result = policy.select_worker_with_events( + // Completely different tokens — no overlap → None + let result = policy.score_overlap( &workers, &[100, 200, 300, 400, 500, 600, 700, 800], &[0], &indexer, ); - assert_eq!(result, None); // no match → falls back + assert_eq!(result, None); } #[test] - fn test_event_driven_load_tiebreak() { + fn test_score_overlap_load_tiebreak() { let policy = CacheAwarePolicy::with_config(test_config()); let w1 = BasicWorkerBuilder::new("http://w1:8000") @@ -1195,12 +1206,12 @@ mod tests { .unwrap(); // Equal overlap → tie-break by load → w2 wins (lower load) - let result = policy.select_worker_with_events(&workers, &[1, 2, 3, 4], &[0, 1], &indexer); + let result = policy.score_overlap(&workers, &[1, 2, 3, 4], &[0, 1], &indexer); assert_eq!(result, Some(1)); // w2 (lower load) } #[test] - fn test_event_driven_tree_size_tiebreak() { + fn test_score_overlap_tree_size_tiebreak() { let policy = CacheAwarePolicy::with_config(test_config()); let workers: Vec> = vec![ Arc::new( @@ -1244,12 +1255,12 @@ mod tests { .unwrap(); // Equal overlap, equal load → tie-break by tree size → w1 wins (smaller) - let result = policy.select_worker_with_events(&workers, &[1, 2, 3, 4], &[0, 1], &indexer); + let result = policy.score_overlap(&workers, &[1, 2, 3, 4], &[0, 1], &indexer); assert_eq!(result, Some(0)); // w1 (smaller tree) } #[test] - fn test_event_driven_short_request_returns_none() { + fn test_score_overlap_short_request_returns_none() { let policy = CacheAwarePolicy::with_config(test_config()); // block_size=4 let workers: Vec> = vec![Arc::new( BasicWorkerBuilder::new("http://w1:8000") @@ -1260,12 +1271,12 @@ mod tests { let indexer = setup_indexer_with_blocks("http://w1:8000", &[&[1, 2, 3, 4]], 4); // Request shorter than block_size → no full blocks → None - let result = policy.select_worker_with_events(&workers, &[1, 2, 3], &[0], &indexer); + let result = policy.score_overlap(&workers, &[1, 2, 3], &[0], &indexer); assert_eq!(result, None); } #[test] - fn test_three_tier_dispatch_prefers_events() { + fn test_score_overlap_partial_match() { let policy = CacheAwarePolicy::with_config(test_config()); let workers: Vec> = vec![ Arc::new( @@ -1281,27 +1292,163 @@ mod tests { ]; policy.init_workers(&workers); - // Event indexer: only w1 has cached blocks - let indexer = - setup_indexer_with_blocks("http://w1:8000", &[&[1, 2, 3, 4], &[5, 6, 7, 8]], 4); + let indexer = Arc::new(PositionalIndexer::new(4)); + + // w1 has 4 blocks cached + let blocks_w1: Vec = (0..4) + .map(|i| StoredBlock { + seq_hash: SequenceHash(i as u64 + 1), + content_hash: compute_content_hash(&[ + (i * 4 + 1) as u32, + (i * 4 + 2) as u32, + (i * 4 + 3) as u32, + (i * 4 + 4) as u32, + ]), + }) + .collect(); + indexer + .apply_stored("http://w1:8000", &blocks_w1, None) + .unwrap(); - // Event-driven path selects w1 (it has the cached blocks) - let result = policy.select_worker_with_events( + // w2 has only the first 2 blocks (partial overlap with same request) + let blocks_w2: Vec = (0..2) + .map(|i| StoredBlock { + seq_hash: SequenceHash(i as u64 + 1), + content_hash: compute_content_hash(&[ + (i * 4 + 1) as u32, + (i * 4 + 2) as u32, + (i * 4 + 3) as u32, + (i * 4 + 4) as u32, + ]), + }) + .collect(); + indexer + .apply_stored("http://w2:8000", &blocks_w2, None) + .unwrap(); + + // Query with all 4 blocks worth of tokens → w1 wins (higher overlap: 4 vs 2) + let result = policy.score_overlap( &workers, - &[1, 2, 3, 4, 5, 6, 7, 8], + &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16], &[0, 1], &indexer, ); - assert_eq!(result, Some(0)); // w1 via event-driven + assert_eq!(result, Some(0)); // w1 (higher overlap) + } + + // -- select_worker_event_driven integration tests -- + + #[test] + fn test_event_driven_overlap_selects_cached_worker() { + let mut policy = CacheAwarePolicy::with_config(test_config()); + let workers: Vec> = vec![ + Arc::new( + BasicWorkerBuilder::new("http://w1:8000") + .worker_type(WorkerType::Regular) + .build(), + ), + Arc::new( + BasicWorkerBuilder::new("http://w2:8000") + .worker_type(WorkerType::Regular) + .build(), + ), + ]; + policy.init_workers(&workers); + + // Set up monitor with indexer data for "unknown" model + let monitor = Arc::new(KvEventMonitor::new(Some(4))); + let indexer = + setup_indexer_with_blocks("http://w1:8000", &[&[1, 2, 3, 4], &[5, 6, 7, 8]], 4); + monitor.indexers.insert("unknown".to_string(), indexer); + policy.set_kv_event_monitor(Some(monitor)); + + // Full dispatch: should use event-driven and select w1 + let idx = policy + .select_worker( + &workers, + &SelectWorkerInfo { + tokens: Some(&[1, 2, 3, 4, 5, 6, 7, 8]), + ..Default::default() + }, + ) + .unwrap(); + assert_eq!(idx, 0); // w1 (has cached blocks) + } + + #[test] + fn test_event_driven_no_overlap_uses_min_load() { + let mut policy = CacheAwarePolicy::with_config(test_config()); + + let w1 = BasicWorkerBuilder::new("http://w1:8000") + .worker_type(WorkerType::Regular) + .build(); + let w2 = BasicWorkerBuilder::new("http://w2:8000") + .worker_type(WorkerType::Regular) + .build(); + // Give w1 higher load so min-load picks w2 + for _ in 0..3 { + w1.increment_load(); + } + + let workers: Vec> = vec![Arc::new(w1), Arc::new(w2)]; + policy.init_workers(&workers); + + // Monitor has indexer with data, but tokens don't match + let monitor = Arc::new(KvEventMonitor::new(Some(4))); + let indexer = setup_indexer_with_blocks("http://w1:8000", &[&[1, 2, 3, 4]], 4); + monitor.indexers.insert("unknown".to_string(), indexer); + policy.set_kv_event_monitor(Some(monitor)); - // Without event data (no overlap), returns None → would fall back - let result = - policy.select_worker_with_events(&workers, &[100, 200, 300, 400], &[0, 1], &indexer); - assert_eq!(result, None); // no match → caller falls back to approximate tree + // No overlap → event-driven falls back to min-load (not token tree) + let idx = policy + .select_worker( + &workers, + &SelectWorkerInfo { + tokens: Some(&[100, 200, 300, 400]), + ..Default::default() + }, + ) + .unwrap(); + assert_eq!(idx, 1); // w2 (min load), NOT token tree result } #[test] - fn test_three_tier_empty_indexer_falls_back() { + fn test_event_driven_short_request_uses_min_load() { + let mut policy = CacheAwarePolicy::with_config(test_config()); // block_size=4 + + let w1 = BasicWorkerBuilder::new("http://w1:8000") + .worker_type(WorkerType::Regular) + .build(); + let w2 = BasicWorkerBuilder::new("http://w2:8000") + .worker_type(WorkerType::Regular) + .build(); + for _ in 0..3 { + w1.increment_load(); + } + + let workers: Vec> = vec![Arc::new(w1), Arc::new(w2)]; + policy.init_workers(&workers); + + let monitor = Arc::new(KvEventMonitor::new(Some(4))); + let indexer = setup_indexer_with_blocks("http://w1:8000", &[&[1, 2, 3, 4]], 4); + monitor.indexers.insert("unknown".to_string(), indexer); + policy.set_kv_event_monitor(Some(monitor)); + + // Request shorter than block_size → no full blocks → min-load fallback + let idx = policy + .select_worker( + &workers, + &SelectWorkerInfo { + tokens: Some(&[1, 2, 3]), + ..Default::default() + }, + ) + .unwrap(); + assert_eq!(idx, 1); // w2 (min load) + } + + #[test] + fn test_no_monitor_uses_token_tree() { let policy = CacheAwarePolicy::with_config(test_config()); let workers: Vec> = vec![ Arc::new( @@ -1317,17 +1464,20 @@ mod tests { ]; policy.init_workers(&workers); - // Empty indexer — no data - let indexer = Arc::new(PositionalIndexer::new(4)); + // No kv_monitor → has_event_indexer returns false → uses token tree + assert!(!policy.has_event_indexer("unknown")); - // Should return None → caller falls back to approximate tree - let result = policy.select_worker_with_events( - &workers, - &[1, 2, 3, 4, 5, 6, 7, 8], - &[0, 1], - &indexer, - ); - assert_eq!(result, None); + // Should still route (via token tree, not event-driven) + let idx = policy + .select_worker( + &workers, + &SelectWorkerInfo { + tokens: Some(&[1, 2, 3, 4]), + ..Default::default() + }, + ) + .unwrap(); + assert!(idx < 2); // valid worker selected } #[test] @@ -1351,7 +1501,7 @@ mod tests { } #[test] - fn test_event_driven_imbalanced_skips_events() { + fn test_imbalanced_skips_event_driven() { let mut policy = CacheAwarePolicy::with_config(CacheAwareConfig { balance_abs_threshold: 5, balance_rel_threshold: 2.0, From 105332109cb6575bdbd2cfda82be2eaf2dcfb097 Mon Sep 17 00:00:00 2001 From: Simo Lin Date: Mon, 2 Mar 2026 07:35:21 -0800 Subject: [PATCH 3/6] feat(config): add per-model kv_block_size to WorkerSpec and KvEventMonitor MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add kv_block_size as a configurable per-worker field in WorkerSpec, discoverable during worker registration. This replaces the hardcoded block_size=16 default with a per-model value that users can control via the workers API. What changed: - protocols/src/worker.rs: add kv_block_size: Option to WorkerSpec with serde support (default None, skip_serializing_if None) - model_gateway/src/core/kv_event_monitor.rs: add block_sizes DashMap per-model storage, block_size()/set_block_size() public API, seed from WorkerSpec during on_worker_added, cleanup on worker removal and stop - model_gateway/src/policies/cache_aware.rs: resolve per-model block_size from monitor (falling back to config default), pass explicit block_size to score_overlap instead of reading from config Why: Block size is a static property of the backend engine (e.g. 16 for SGLang page size). Different models or backends may use different block sizes. Hardcoding 16 in the config prevents multi-model deployments with heterogeneous block sizes. Following Dynamo's pattern where block_size is per-model via ModelDeploymentCard, discovered during registration. How: Discovery flow: WorkerSpec.kv_block_size (user sets via API) → on_worker_added seeds KvEventMonitor.block_sizes → CacheAwarePolicy reads monitor.block_size(model_id) at query time, falls back to CacheAwareConfig.block_size (default 16) if not set. Tests: 3 new tests (set_block_size, stop_clears_block_sizes, event_driven_uses_monitor_block_size), 6 existing tests updated for new score_overlap signature. Signed-off-by: Simo Lin --- model_gateway/src/core/kv_event_monitor.rs | 55 +++++++++++++ model_gateway/src/policies/cache_aware.rs | 90 +++++++++++++++++++--- protocols/src/worker.rs | 7 ++ 3 files changed, 141 insertions(+), 11 deletions(-) diff --git a/model_gateway/src/core/kv_event_monitor.rs b/model_gateway/src/core/kv_event_monitor.rs index 43ee727f69..27e4b1dc28 100644 --- a/model_gateway/src/core/kv_event_monitor.rs +++ b/model_gateway/src/core/kv_event_monitor.rs @@ -39,6 +39,9 @@ const MAX_RECONNECT_DELAY_MS: u64 = 30_000; pub struct KvEventMonitor { /// Per-model positional indexers: model_id → shared indexer. pub(crate) indexers: DashMap>, + /// Per-model block sizes learned from KV events or set via WorkerSpec. + /// Used by CacheAwarePolicy to chunk request tokens at query time. + block_sizes: DashMap, /// Per-worker subscription handles: worker_url → subscription info. /// Mutex matches LoadMonitor pattern for atomic abort + remove. worker_handles: Mutex>, @@ -71,6 +74,7 @@ impl KvEventMonitor { let jump_size = jump_size.unwrap_or(DEFAULT_JUMP_SIZE).max(1); Self { indexers: DashMap::new(), + block_sizes: DashMap::new(), worker_handles: Mutex::new(HashMap::new()), jump_size, } @@ -102,6 +106,11 @@ impl KvEventMonitor { .or_insert_with(|| Arc::new(PositionalIndexer::new(self.jump_size))) .clone(); + // Seed block_size from WorkerSpec if set and not already known. + if let Some(bs) = worker.metadata().spec.kv_block_size { + self.block_sizes.entry(model_id.clone()).or_insert(bs); + } + let worker = Arc::clone(worker); let worker_url = url.clone(); @@ -151,6 +160,7 @@ impl KvEventMonitor { if should_remove_indexer { self.indexers.remove(&sub.model_id); + self.block_sizes.remove(&sub.model_id); } } @@ -176,6 +186,7 @@ impl KvEventMonitor { } self.indexers.clear(); + self.block_sizes.clear(); } /// Get the indexer for a model (used by `CacheAwarePolicy` for queries). @@ -183,6 +194,19 @@ impl KvEventMonitor { self.indexers.get(model_id).map(|r| Arc::clone(&r)) } + /// Get the block size for a model (learned from events or set via `set_block_size`). + pub fn block_size(&self, model_id: &str) -> Option { + self.block_sizes.get(model_id).map(|v| *v) + } + + /// Set the block size for a model (e.g. from WorkerSpec during registration). + /// Does not overwrite a value already learned from events. + pub fn set_block_size(&self, model_id: &str, block_size: usize) { + self.block_sizes + .entry(model_id.to_string()) + .or_insert(block_size); + } + /// Check if any subscription is running. pub async fn is_running(&self) -> bool { !self.worker_handles.lock().await.is_empty() @@ -412,6 +436,7 @@ impl fmt::Debug for KvEventMonitor { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("KvEventMonitor") .field("models", &self.indexers.len()) + .field("block_sizes", &self.block_sizes.len()) .field("jump_size", &self.jump_size) .finish() } @@ -734,4 +759,34 @@ mod tests { let monitor = KvEventMonitor::new(None); monitor.on_worker_removed("http://nonexistent:8000").await; } + + // ----------------------------------------------------------------------- + // block_size learning + // ----------------------------------------------------------------------- + + #[test] + fn test_set_block_size() { + let monitor = KvEventMonitor::new(None); + + // Initially no block_size + assert!(monitor.block_size("llama").is_none()); + + // Set it + monitor.set_block_size("llama", 32); + assert_eq!(monitor.block_size("llama"), Some(32)); + + // set_block_size doesn't overwrite existing value + monitor.set_block_size("llama", 64); + assert_eq!(monitor.block_size("llama"), Some(32)); + } + + #[tokio::test] + async fn test_stop_clears_block_sizes() { + let monitor = KvEventMonitor::new(None); + monitor.set_block_size("llama", 16); + assert_eq!(monitor.block_size("llama"), Some(16)); + + monitor.stop().await; + assert!(monitor.block_size("llama").is_none()); + } } diff --git a/model_gateway/src/policies/cache_aware.rs b/model_gateway/src/policies/cache_aware.rs index 7e3a518e88..6e131a2a4f 100644 --- a/model_gateway/src/policies/cache_aware.rs +++ b/model_gateway/src/policies/cache_aware.rs @@ -507,9 +507,17 @@ impl CacheAwarePolicy { healthy_indices: &[usize], model_id: &str, ) -> Option { - let indexer = self.kv_monitor.as_ref()?.get_indexer(model_id)?; + let monitor = self.kv_monitor.as_ref()?; + let indexer = monitor.get_indexer(model_id)?; - if let Some(idx) = self.score_overlap(workers, tokens, healthy_indices, &indexer) { + // Per-model block_size: learned from events > config default + let block_size = monitor + .block_size(model_id) + .unwrap_or(self.config.block_size); + + if let Some(idx) = + Self::score_overlap(workers, tokens, healthy_indices, &indexer, block_size) + { return Some(idx); } @@ -532,13 +540,13 @@ impl CacheAwarePolicy { /// request. Returns `None` if the request is too short for a full block or /// no workers have matching data. fn score_overlap( - &self, workers: &[Arc], tokens: &[u32], healthy_indices: &[usize], indexer: &PositionalIndexer, + block_size: usize, ) -> Option { - let content_hashes = compute_request_content_hashes(tokens, self.config.block_size); + let content_hashes = compute_request_content_hashes(tokens, block_size); if content_hashes.is_empty() { return None; } @@ -1137,11 +1145,12 @@ mod tests { ); // Query with matching tokens — should select w1 - let result = policy.score_overlap( + let result = CacheAwarePolicy::score_overlap( &workers, &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16], &[0, 1], &indexer, + 4, ); assert_eq!(result, Some(0)); // w1 } @@ -1160,11 +1169,12 @@ mod tests { setup_indexer_with_blocks("http://w1:8000", &[&[1, 2, 3, 4], &[5, 6, 7, 8]], 4); // Completely different tokens — no overlap → None - let result = policy.score_overlap( + let result = CacheAwarePolicy::score_overlap( &workers, &[100, 200, 300, 400, 500, 600, 700, 800], &[0], &indexer, + 4, ); assert_eq!(result, None); } @@ -1206,7 +1216,7 @@ mod tests { .unwrap(); // Equal overlap → tie-break by load → w2 wins (lower load) - let result = policy.score_overlap(&workers, &[1, 2, 3, 4], &[0, 1], &indexer); + let result = CacheAwarePolicy::score_overlap(&workers, &[1, 2, 3, 4], &[0, 1], &indexer, 4); assert_eq!(result, Some(1)); // w2 (lower load) } @@ -1255,13 +1265,12 @@ mod tests { .unwrap(); // Equal overlap, equal load → tie-break by tree size → w1 wins (smaller) - let result = policy.score_overlap(&workers, &[1, 2, 3, 4], &[0, 1], &indexer); + let result = CacheAwarePolicy::score_overlap(&workers, &[1, 2, 3, 4], &[0, 1], &indexer, 4); assert_eq!(result, Some(0)); // w1 (smaller tree) } #[test] fn test_score_overlap_short_request_returns_none() { - let policy = CacheAwarePolicy::with_config(test_config()); // block_size=4 let workers: Vec> = vec![Arc::new( BasicWorkerBuilder::new("http://w1:8000") .worker_type(WorkerType::Regular) @@ -1271,7 +1280,7 @@ mod tests { let indexer = setup_indexer_with_blocks("http://w1:8000", &[&[1, 2, 3, 4]], 4); // Request shorter than block_size → no full blocks → None - let result = policy.score_overlap(&workers, &[1, 2, 3], &[0], &indexer); + let result = CacheAwarePolicy::score_overlap(&workers, &[1, 2, 3], &[0], &indexer, 4); assert_eq!(result, None); } @@ -1327,11 +1336,12 @@ mod tests { .unwrap(); // Query with all 4 blocks worth of tokens → w1 wins (higher overlap: 4 vs 2) - let result = policy.score_overlap( + let result = CacheAwarePolicy::score_overlap( &workers, &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16], &[0, 1], &indexer, + 4, ); assert_eq!(result, Some(0)); // w1 (higher overlap) } @@ -1500,6 +1510,64 @@ mod tests { assert!(policy.kv_monitor.is_none()); } + #[test] + fn test_event_driven_uses_monitor_block_size() { + // Test that event-driven routing uses monitor's learned block_size + // instead of config default when available. + let mut policy = CacheAwarePolicy::with_config(CacheAwareConfig { + block_size: 4, // config default + eviction_interval_secs: 0, + ..Default::default() + }); + + let workers: Vec> = vec![ + Arc::new( + BasicWorkerBuilder::new("http://w1:8000") + .worker_type(WorkerType::Regular) + .build(), + ), + Arc::new( + BasicWorkerBuilder::new("http://w2:8000") + .worker_type(WorkerType::Regular) + .build(), + ), + ]; + policy.init_workers(&workers); + + let monitor = Arc::new(KvEventMonitor::new(Some(4))); + + // Store blocks using block_size=8 (tokens chunked in groups of 8) + let indexer = Arc::new(PositionalIndexer::new(4)); + let block = vec![StoredBlock { + seq_hash: SequenceHash(1), + content_hash: compute_content_hash(&[1, 2, 3, 4, 5, 6, 7, 8]), + }]; + indexer + .apply_stored("http://w1:8000", &block, None) + .unwrap(); + monitor + .indexers + .insert("unknown".to_string(), indexer.clone()); + + // Set block_size=8 in monitor (simulating learned from events) + monitor.set_block_size("unknown", 8); + + policy.set_kv_event_monitor(Some(monitor)); + + // Query with 8 tokens — with block_size=8, this is one full block + // With config block_size=4, this would be two blocks and wouldn't match + let idx = policy + .select_worker( + &workers, + &SelectWorkerInfo { + tokens: Some(&[1, 2, 3, 4, 5, 6, 7, 8]), + ..Default::default() + }, + ) + .unwrap(); + assert_eq!(idx, 0); // w1 has the cached block + } + #[test] fn test_imbalanced_skips_event_driven() { let mut policy = CacheAwarePolicy::with_config(CacheAwareConfig { diff --git a/protocols/src/worker.rs b/protocols/src/worker.rs index 7b62889bbb..52108d127d 100644 --- a/protocols/src/worker.rs +++ b/protocols/src/worker.rs @@ -481,6 +481,12 @@ pub struct WorkerSpec { #[serde(default, skip_serializing_if = "Option::is_none")] pub kv_role: Option, + /// KV cache block size (tokens per block) for event-driven routing. + /// When set, overrides the router-level default for this worker's model. + /// Typically matches the backend engine's page size (e.g. 16 for SGLang). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub kv_block_size: Option, + /// Per-worker health check overrides (partial — only `Some` fields override router defaults). #[serde(default, skip_serializing_if = "HealthCheckUpdate::is_empty")] pub health: HealthCheckUpdate, @@ -517,6 +523,7 @@ impl WorkerSpec { dp_size: None, kv_connector: None, kv_role: None, + kv_block_size: None, health: HealthCheckUpdate::default(), max_connection_attempts: default_max_connection_attempts(), load_monitor_interval_secs: None, From 112a2436741237329ca7ddaf5adacea3635da6da Mon Sep 17 00:00:00 2001 From: Simo Lin Date: Mon, 2 Mar 2026 08:29:09 -0800 Subject: [PATCH 4/6] fix: address PR review comments on cache-aware event routing - Add block_size > 0 validation in config validator to prevent panic from chunks_exact(0) in compute_request_content_hashes - Filter zero-overlap workers in score_overlap to avoid selecting a worker with no cache match as if it were an overlap hit - Simplify score_overlap ranking with max_by_key + tuple instead of manual max_by with then_with chains Signed-off-by: Simon Lin Signed-off-by: Simo Lin --- model_gateway/src/config/validation.rs | 10 ++++++- model_gateway/src/policies/cache_aware.rs | 33 ++++++++--------------- 2 files changed, 20 insertions(+), 23 deletions(-) diff --git a/model_gateway/src/config/validation.rs b/model_gateway/src/config/validation.rs index f8373aa79d..d29d47f8b2 100644 --- a/model_gateway/src/config/validation.rs +++ b/model_gateway/src/config/validation.rs @@ -173,8 +173,16 @@ impl ConfigValidator { balance_rel_threshold, eviction_interval_secs, max_tree_size, - .. + block_size, } => { + if *block_size == 0 { + return Err(ConfigError::InvalidValue { + field: "block_size".to_string(), + value: block_size.to_string(), + reason: "Must be > 0".to_string(), + }); + } + if !(0.0..=1.0).contains(cache_threshold) { return Err(ConfigError::InvalidValue { field: "cache_threshold".to_string(), diff --git a/model_gateway/src/policies/cache_aware.rs b/model_gateway/src/policies/cache_aware.rs index 6e131a2a4f..c935c488d8 100644 --- a/model_gateway/src/policies/cache_aware.rs +++ b/model_gateway/src/policies/cache_aware.rs @@ -556,30 +556,19 @@ impl CacheAwarePolicy { return None; } - // Select worker with best overlap; tie-break by load (lower) then tree size (smaller) + // Select worker with best overlap among those that actually match. + // Tie-break: lower load, then smaller tree size. let best_idx = healthy_indices .iter() - .max_by(|&&a, &&b| { - let score_a = overlap.scores.get(workers[a].url()).copied().unwrap_or(0); - let score_b = overlap.scores.get(workers[b].url()).copied().unwrap_or(0); - score_a - .cmp(&score_b) - .then_with(|| workers[b].load().cmp(&workers[a].load())) - .then_with(|| { - let size_a = overlap - .tree_sizes - .get(workers[a].url()) - .copied() - .unwrap_or(0); - let size_b = overlap - .tree_sizes - .get(workers[b].url()) - .copied() - .unwrap_or(0); - size_b.cmp(&size_a) - }) - }) - .copied()?; + .copied() + .filter(|&idx| overlap.scores.get(workers[idx].url()).copied().unwrap_or(0) > 0) + .max_by_key(|&idx| { + let url = workers[idx].url(); + let score = overlap.scores.get(url).copied().unwrap_or(0); + let load = workers[idx].load(); + let tree_size = overlap.tree_sizes.get(url).copied().unwrap_or(0); + (score, std::cmp::Reverse(load), std::cmp::Reverse(tree_size)) + })?; debug!( worker = workers[best_idx].url(), From de02057cfe290e5609ffa03b81fc219e8e007423 Mon Sep 17 00:00:00 2001 From: Simo Lin Date: Mon, 2 Mar 2026 09:44:55 -0800 Subject: [PATCH 5/6] feat: wire KvEventMonitor into workflow and expose block_size CLI arg Wire the KvEventMonitor into the production workflow so event-driven cache-aware routing activates automatically for eligible workers. AppContext: - Add kv_event_monitor field (Option>) - Create monitor in with_kv_event_monitor() when policy is cache_aware and connection mode is gRPC - Inject monitor into PolicyRegistry at startup PolicyRegistry: - Add kv_event_monitor field with RwLock for thread-safe late init - Add set_kv_event_monitor() method (follows set_mesh_sync pattern) - Inject monitor into new CacheAwarePolicy instances during create_policy_from_type("cache_aware") UpdatePoliciesStep: - Call monitor.on_worker_added(worker) for gRPC workers with cache_aware policy, starting KV event subscription RemoveFromPolicyRegistryStep: - Call monitor.on_worker_removed(worker_url) before policy cleanup, stopping subscription and removing worker from indexer block_size CLI parameter (replaces hardcoded 16): - router_args.py: add block_size field + --block-size argument - lib.rs (Python bindings): add block_size to Router struct, signature, constructor, use self.block_size in CacheAware config - main.rs: add --block-size CLI arg, use self.block_size in config Bug fix: - kv_event_monitor.rs: stop() now always clears indexers and block_sizes, not only when subscriptions exist Signed-off-by: Simo Lin --- bindings/python/src/lib.rs | 6 ++- bindings/python/src/smg/router_args.py | 7 ++++ model_gateway/src/app_context.rs | 41 ++++++++++++++++++- model_gateway/src/core/kv_event_monitor.rs | 22 +++++----- .../local/remove_from_policy_registry.rs | 5 +++ .../steps/worker/shared/update_policies.rs | 13 +++++- model_gateway/src/main.rs | 6 ++- model_gateway/src/policies/registry.rs | 22 +++++++++- model_gateway/src/service_discovery.rs | 1 + 9 files changed, 106 insertions(+), 17 deletions(-) diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs index 8800642044..2b9f559282 100644 --- a/bindings/python/src/lib.rs +++ b/bindings/python/src/lib.rs @@ -375,6 +375,7 @@ struct Router { balance_rel_threshold: f32, eviction_interval_secs: u64, max_tree_size: usize, + block_size: usize, max_idle_secs: u64, assignment_mode: String, max_payload_size: usize, @@ -477,7 +478,7 @@ impl Router { balance_rel_threshold: self.balance_rel_threshold, eviction_interval_secs: self.eviction_interval_secs, max_tree_size: self.max_tree_size, - block_size: 16, + block_size: self.block_size, }, PolicyType::PowerOfTwo => ConfigPolicyConfig::PowerOfTwo { load_check_interval_secs: 5, @@ -700,6 +701,7 @@ impl Router { balance_rel_threshold = 1.5, eviction_interval_secs = 120, max_tree_size = 2usize.pow(26), + block_size = 16, max_idle_secs = 14400, assignment_mode = String::from("random"), max_payload_size = 512 * 1024 * 1024, @@ -793,6 +795,7 @@ impl Router { balance_rel_threshold: f32, eviction_interval_secs: u64, max_tree_size: usize, + block_size: usize, max_idle_secs: u64, assignment_mode: String, max_payload_size: usize, @@ -895,6 +898,7 @@ impl Router { balance_rel_threshold, eviction_interval_secs, max_tree_size, + block_size, max_idle_secs, assignment_mode, max_payload_size, diff --git a/bindings/python/src/smg/router_args.py b/bindings/python/src/smg/router_args.py index 7555bf5221..1fc594369f 100644 --- a/bindings/python/src/smg/router_args.py +++ b/bindings/python/src/smg/router_args.py @@ -36,6 +36,7 @@ class RouterArgs: balance_rel_threshold: float = 1.5 eviction_interval_secs: int = 60 max_tree_size: int = 2**26 + block_size: int = 16 max_idle_secs: int = 4 * 3600 assignment_mode: str = "random" # Mode for manual policy new routing key assignment max_payload_size: int = 512 * 1024 * 1024 # 512MB default for large batches @@ -332,6 +333,12 @@ def add_cli_args( default=RouterArgs.max_tree_size, help="Maximum size of the approximation tree for cache-aware routing", ) + routing_group.add_argument( + f"--{prefix}block-size", + type=int, + default=RouterArgs.block_size, + help="KV cache block size for event-driven cache-aware routing (default: 16)", + ) routing_group.add_argument( f"--{prefix}max-idle-secs", type=int, diff --git a/model_gateway/src/app_context.rs b/model_gateway/src/app_context.rs index e345ba57bb..62e8ba1e0c 100644 --- a/model_gateway/src/app_context.rs +++ b/model_gateway/src/app_context.rs @@ -16,7 +16,10 @@ use tracing::debug; use crate::{ config::RouterConfig, - core::{steps::WorkflowEngines, JobQueue, LoadMonitor, WorkerRegistry, WorkerService}, + core::{ + steps::WorkflowEngines, JobQueue, KvEventMonitor, LoadMonitor, WorkerRegistry, + WorkerService, + }, middleware::TokenBucket, observability::inflight_tracker::InFlightRequestTracker, policies::PolicyRegistry, @@ -59,6 +62,7 @@ pub struct AppContext { pub wasm_manager: Option>, pub worker_service: Arc, pub inflight_tracker: Arc, + pub kv_event_monitor: Option>, } impl std::fmt::Debug for AppContext { @@ -87,6 +91,7 @@ pub struct AppContextBuilder { workflow_engines: Option>>, mcp_orchestrator: Option>>>, wasm_manager: Option>, + kv_event_monitor: Option>, } impl AppContext { @@ -132,6 +137,7 @@ impl AppContextBuilder { workflow_engines: None, mcp_orchestrator: None, wasm_manager: None, + kv_event_monitor: None, } } @@ -232,6 +238,11 @@ impl AppContextBuilder { self } + pub fn kv_event_monitor(mut self, kv_event_monitor: Option>) -> Self { + self.kv_event_monitor = kv_event_monitor; + self + } + pub fn build(self) -> Result { let router_config = self .router_config @@ -289,6 +300,7 @@ impl AppContextBuilder { wasm_manager: self.wasm_manager, worker_service, inflight_tracker: InFlightRequestTracker::new(), + kv_event_monitor: self.kv_event_monitor, }) } @@ -314,6 +326,7 @@ impl AppContextBuilder { .with_mcp_orchestrator(&router_config) .await? .with_wasm_manager(&router_config) + .with_kv_event_monitor(&router_config) .router_config(router_config)) } @@ -544,6 +557,32 @@ impl AppContextBuilder { Ok(self) } + /// Create KV event monitor for event-driven cache-aware routing. + /// + /// The monitor is created when the default policy is cache_aware and + /// connection mode is gRPC. It stays dormant until workers are added + /// via `on_worker_added` in the UpdatePoliciesStep. + fn with_kv_event_monitor(mut self, config: &RouterConfig) -> Self { + use crate::config::types::PolicyConfig; + + let is_cache_aware = matches!(config.policy, PolicyConfig::CacheAware { .. }); + let is_grpc = matches!(config.connection_mode, crate::core::ConnectionMode::Grpc); + + if is_cache_aware && is_grpc { + let monitor = Arc::new(KvEventMonitor::new(None)); + debug!("Created KV event monitor for event-driven cache-aware routing"); + + // Inject monitor into PolicyRegistry so new CacheAwarePolicy instances get it + if let Some(ref registry) = self.policy_registry { + registry.set_kv_event_monitor(Some(Arc::clone(&monitor))); + } + + self.kv_event_monitor = Some(monitor); + } + + self + } + /// Create wasm manager if enabled in config fn with_wasm_manager(mut self, config: &RouterConfig) -> Self { self.wasm_manager = if config.enable_wasm { diff --git a/model_gateway/src/core/kv_event_monitor.rs b/model_gateway/src/core/kv_event_monitor.rs index 27e4b1dc28..7b494eb1f4 100644 --- a/model_gateway/src/core/kv_event_monitor.rs +++ b/model_gateway/src/core/kv_event_monitor.rs @@ -171,18 +171,16 @@ impl KvEventMonitor { std::mem::take(&mut *handles) }; - if subscriptions.is_empty() { - return; - } - - info!( - count = subscriptions.len(), - "Stopping all KV event subscriptions" - ); - for (url, sub) in subscriptions { - debug!(worker_url = %url, "Aborting KV event subscription"); - sub.handle.abort(); - let _ = sub.handle.await; + if !subscriptions.is_empty() { + info!( + count = subscriptions.len(), + "Stopping all KV event subscriptions" + ); + for (url, sub) in subscriptions { + debug!(worker_url = %url, "Aborting KV event subscription"); + sub.handle.abort(); + let _ = sub.handle.await; + } } self.indexers.clear(); diff --git a/model_gateway/src/core/steps/worker/local/remove_from_policy_registry.rs b/model_gateway/src/core/steps/worker/local/remove_from_policy_registry.rs index 58a06b63fb..d4422e974d 100644 --- a/model_gateway/src/core/steps/worker/local/remove_from_policy_registry.rs +++ b/model_gateway/src/core/steps/worker/local/remove_from_policy_registry.rs @@ -38,6 +38,11 @@ impl StepExecutor for RemoveFromPolicyRegistryStep { let model_id = worker.model_id().to_string(); let worker_url = worker.url(); + // Remove from KV event monitor (stops subscription, removes from indexer) + if let Some(ref monitor) = app_context.kv_event_monitor { + monitor.on_worker_removed(worker_url).await; + } + // Remove from cache-aware policy app_context .policy_registry diff --git a/model_gateway/src/core/steps/worker/shared/update_policies.rs b/model_gateway/src/core/steps/worker/shared/update_policies.rs index 0985f52588..36b570e2b1 100644 --- a/model_gateway/src/core/steps/worker/shared/update_policies.rs +++ b/model_gateway/src/core/steps/worker/shared/update_policies.rs @@ -8,7 +8,7 @@ use wfaas::{ StepExecutor, StepResult, WorkflowContext, WorkflowData, WorkflowError, WorkflowResult, }; -use crate::core::{steps::workflow_data::WorkerRegistrationData, Worker}; +use crate::core::{steps::workflow_data::WorkerRegistrationData, ConnectionMode, Worker}; /// Unified step to update policy registry for registered workers. /// @@ -127,6 +127,17 @@ impl StepExecutor for UpdatePolicie } } + // Start KV event subscription for gRPC workers with cache_aware policy + if let Some(ref monitor) = app_context.kv_event_monitor { + if let Some(policy) = app_context.policy_registry.get_policy(&model_id) { + if policy.name() == "cache_aware" + && *worker.connection_mode() == ConnectionMode::Grpc + { + monitor.on_worker_added(worker).await; + } + } + } + if !updated_models.contains(&model_id) { updated_models.push(model_id); } diff --git a/model_gateway/src/main.rs b/model_gateway/src/main.rs index 550a4d1e25..21e6ab9241 100644 --- a/model_gateway/src/main.rs +++ b/model_gateway/src/main.rs @@ -169,6 +169,10 @@ struct CliArgs { #[arg(long, default_value_t = 67108864, help_heading = "Routing Policy")] max_tree_size: usize, + /// KV cache block size for event-driven cache-aware routing + #[arg(long, default_value_t = 16, help_heading = "Routing Policy")] + block_size: usize, + /// Maximum idle time in seconds before eviction (for manual policy) #[arg(long, default_value_t = 14400, help_heading = "Routing Policy")] max_idle_secs: u64, @@ -754,7 +758,7 @@ impl CliArgs { balance_rel_threshold: self.balance_rel_threshold, eviction_interval_secs: self.eviction_interval, max_tree_size: self.max_tree_size, - block_size: 16, + block_size: self.block_size, }, "power_of_two" => PolicyConfig::PowerOfTwo { load_check_interval_secs: 5, diff --git a/model_gateway/src/policies/registry.rs b/model_gateway/src/policies/registry.rs index 36f1fe710e..90bd445905 100644 --- a/model_gateway/src/policies/registry.rs +++ b/model_gateway/src/policies/registry.rs @@ -13,7 +13,10 @@ use tracing::{debug, info, warn}; /// All subsequent workers of the same model use the established policy. /// When the last worker of a model is removed, the policy mapping is cleaned up. use super::{BucketPolicy, CacheAwarePolicy, LoadBalancingPolicy, PolicyFactory}; -use crate::{config::types::PolicyConfig, core::Worker}; +use crate::{ + config::types::PolicyConfig, + core::{KvEventMonitor, Worker}, +}; /// Registry for managing model-to-policy mappings #[derive(Clone)] @@ -37,6 +40,10 @@ pub struct PolicyRegistry { /// When None, the registry works independently without mesh synchronization /// Uses RwLock for thread-safe access when setting mesh_sync after initialization mesh_sync: Arc>, + + /// Optional KV event monitor for event-driven cache-aware routing. + /// When set, new CacheAwarePolicy instances are injected with this monitor. + kv_event_monitor: Arc>>>, } impl PolicyRegistry { @@ -51,6 +58,7 @@ impl PolicyRegistry { prefill_policy: Arc::new(OnceLock::new()), decode_policy: Arc::new(OnceLock::new()), mesh_sync: Arc::new(RwLock::new(None)), + kv_event_monitor: Arc::new(RwLock::new(None)), } } @@ -60,6 +68,12 @@ impl PolicyRegistry { *guard = mesh_sync; } + /// Set KV event monitor (thread-safe, can be called after initialization) + pub fn set_kv_event_monitor(&self, monitor: Option>) { + let mut guard = self.kv_event_monitor.write(); + *guard = monitor; + } + /// Called when a worker is added /// Returns the policy that should be used for this worker's model pub fn on_worker_added( @@ -202,6 +216,12 @@ impl PolicyRegistry { let guard = self.mesh_sync.read(); cache_aware.set_mesh_sync(guard.clone()); } + { + let guard = self.kv_event_monitor.read(); + if let Some(ref monitor) = *guard { + cache_aware.set_kv_event_monitor(Some(Arc::clone(monitor))); + } + } Arc::new(cache_aware) } else { PolicyFactory::create_by_name(policy_type).unwrap_or_else(|| { diff --git a/model_gateway/src/service_discovery.rs b/model_gateway/src/service_discovery.rs index e51bd296b4..194884e813 100644 --- a/model_gateway/src/service_discovery.rs +++ b/model_gateway/src/service_discovery.rs @@ -858,6 +858,7 @@ mod tests { router_config, )), inflight_tracker: InFlightRequestTracker::new(), + kv_event_monitor: None, }) } From 59b18a9c730799bdb76a1cf14ca4c9a59b575c03 Mon Sep 17 00:00:00 2001 From: Simo Lin Date: Mon, 2 Mar 2026 10:58:43 -0800 Subject: [PATCH 6/6] fix: address P1/P2 review issues in event-driven cache-aware routing Interior mutability for kv_monitor (P1): - Change CacheAwarePolicy.kv_monitor from Option> to RwLock>> so set_kv_event_monitor works via &self on policies already behind Arc. - PolicyRegistry.set_kv_event_monitor now propagates to all existing cache-aware policies (default, prefill, decode, per-model), fixing the issue where default_policy created before the monitor was available would never receive it. Empty indexer fallback (P1): - has_event_indexer now checks current_size() > 0, so an empty indexer (startup, reconnect) falls through to the approximate token tree instead of taking the event-driven path and landing on min-load. - Add test_empty_indexer_falls_through_to_token_tree test. Monitor creation scope (P1): - Remove gRPC connection_mode check from with_kv_event_monitor in AppContextBuilder. The monitor is cheap (empty DashMaps) and UpdatePoliciesStep already gates subscriptions on cache_aware && gRPC. This allows dynamic/per-model gRPC workers to use event-driven routing even when the global connection_mode is not gRPC. Model key normalization (P2): - Add normalize_model_id helper to KvEventMonitor that maps empty model_id to UNKNOWN_MODEL_ID, matching routing's normalize_model_key. Block size validation (P2): - Skip kv_block_size=0 from WorkerSpec instead of storing it, which would poison overlap scoring with zero-length chunks. TOCTOU fix in on_worker_removed (P2): - Re-acquire worker_handles lock after aborting the subscription task to check if the model's last worker was truly removed. Prevents race where concurrent on_worker_added for the same model could repopulate handles between lock release and indexer removal. Signed-off-by: Simo Lin --- model_gateway/src/app_context.rs | 13 +-- model_gateway/src/core/kv_event_monitor.rs | 43 ++++++--- model_gateway/src/policies/cache_aware.rs | 104 +++++++++++++++++---- model_gateway/src/policies/registry.rs | 33 ++++++- 4 files changed, 152 insertions(+), 41 deletions(-) diff --git a/model_gateway/src/app_context.rs b/model_gateway/src/app_context.rs index 62e8ba1e0c..19125b1648 100644 --- a/model_gateway/src/app_context.rs +++ b/model_gateway/src/app_context.rs @@ -559,20 +559,21 @@ impl AppContextBuilder { /// Create KV event monitor for event-driven cache-aware routing. /// - /// The monitor is created when the default policy is cache_aware and - /// connection mode is gRPC. It stays dormant until workers are added - /// via `on_worker_added` in the UpdatePoliciesStep. + /// The monitor is created when the default policy is cache_aware, regardless + /// of connection mode. The monitor itself is cheap (empty DashMaps) and stays + /// dormant until workers are added. The UpdatePoliciesStep gates subscriptions + /// on `cache_aware && gRPC`, so HTTP workers are never subscribed. fn with_kv_event_monitor(mut self, config: &RouterConfig) -> Self { use crate::config::types::PolicyConfig; let is_cache_aware = matches!(config.policy, PolicyConfig::CacheAware { .. }); - let is_grpc = matches!(config.connection_mode, crate::core::ConnectionMode::Grpc); - if is_cache_aware && is_grpc { + if is_cache_aware { let monitor = Arc::new(KvEventMonitor::new(None)); debug!("Created KV event monitor for event-driven cache-aware routing"); - // Inject monitor into PolicyRegistry so new CacheAwarePolicy instances get it + // Inject monitor into PolicyRegistry — propagates to default_policy + // and any other existing cache-aware policies. if let Some(ref registry) = self.policy_registry { registry.set_kv_event_monitor(Some(Arc::clone(&monitor))); } diff --git a/model_gateway/src/core/kv_event_monitor.rs b/model_gateway/src/core/kv_event_monitor.rs index 7b494eb1f4..d75a5b721e 100644 --- a/model_gateway/src/core/kv_event_monitor.rs +++ b/model_gateway/src/core/kv_event_monitor.rs @@ -20,7 +20,7 @@ use smg_grpc_client::common_proto::{ use tokio::{sync::Mutex, task::JoinHandle}; use tracing::{debug, info, warn}; -use crate::core::{ConnectionMode, Worker}; +use crate::core::{ConnectionMode, Worker, UNKNOWN_MODEL_ID}; /// Default jump size for new `PositionalIndexer` instances. const DEFAULT_JUMP_SIZE: usize = 64; @@ -87,7 +87,8 @@ impl KvEventMonitor { /// Duplicate calls for the same worker URL are no-ops. pub async fn on_worker_added(&self, worker: &Arc) { let url = worker.url().to_string(); - let model_id = worker.model_id().to_string(); + // Normalize model_id to match routing's normalize_model_key — empty → "unknown". + let model_id = Self::normalize_model_id(worker.model_id()); if *worker.connection_mode() == ConnectionMode::Http { debug!(worker_url = %url, "HTTP worker, skipping KV event subscription"); @@ -106,9 +107,13 @@ impl KvEventMonitor { .or_insert_with(|| Arc::new(PositionalIndexer::new(self.jump_size))) .clone(); - // Seed block_size from WorkerSpec if set and not already known. + // Seed block_size from WorkerSpec if set, valid, and not already known. if let Some(bs) = worker.metadata().spec.kv_block_size { - self.block_sizes.entry(model_id.clone()).or_insert(bs); + if bs > 0 { + self.block_sizes.entry(model_id.clone()).or_insert(bs); + } else { + warn!(worker_url = %url, "Worker reports kv_block_size=0, ignoring"); + } } let worker = Arc::clone(worker); @@ -134,15 +139,10 @@ impl KvEventMonitor { /// Stop the KV event subscription for a worker and remove it from the indexer. pub async fn on_worker_removed(&self, worker_url: &str) { - // Single lock acquisition: remove subscription and check remaining atomically - // to avoid TOCTOU race with concurrent on_worker_added for the same model. - let (subscription, should_remove_indexer) = { + // Remove subscription from handles under lock. + let subscription = { let mut handles = self.worker_handles.lock().await; - let sub = handles.remove(worker_url); - let should_remove = sub - .as_ref() - .is_some_and(|s| !handles.values().any(|other| other.model_id == s.model_id)); - (sub, should_remove) + handles.remove(worker_url) }; let Some(sub) = subscription else { @@ -158,6 +158,15 @@ impl KvEventMonitor { indexer.remove_worker(worker_url); } + // Re-check under lock whether this was the last worker for the model. + // Must re-acquire lock after abort to avoid TOCTOU with concurrent + // on_worker_added that may have added a new worker for the same model + // between our first lock release and this point. + let should_remove_indexer = { + let handles = self.worker_handles.lock().await; + !handles.values().any(|other| other.model_id == sub.model_id) + }; + if should_remove_indexer { self.indexers.remove(&sub.model_id); self.block_sizes.remove(&sub.model_id); @@ -210,6 +219,16 @@ impl KvEventMonitor { !self.worker_handles.lock().await.is_empty() } + /// Normalize model_id to match routing's `normalize_model_key`. + /// Empty model IDs map to UNKNOWN_MODEL_ID for consistent keying. + fn normalize_model_id(model_id: &str) -> String { + if model_id.is_empty() { + UNKNOWN_MODEL_ID.to_string() + } else { + model_id.to_string() + } + } + // ----------------------------------------------------------------------- // Subscription loop // ----------------------------------------------------------------------- diff --git a/model_gateway/src/policies/cache_aware.rs b/model_gateway/src/policies/cache_aware.rs index c935c488d8..b695b4f3ea 100644 --- a/model_gateway/src/policies/cache_aware.rs +++ b/model_gateway/src/policies/cache_aware.rs @@ -45,6 +45,7 @@ use std::sync::Arc; use dashmap::DashMap; use kv_index::{compute_request_content_hashes, PositionalIndexer, TokenTree, Tree}; +use parking_lot::RwLock; use rand::Rng; use smg_mesh::{OptionalMeshSyncManager, TreeInsertOp, TreeOperation}; use tracing::{debug, warn}; @@ -78,7 +79,9 @@ pub struct CacheAwarePolicy { /// Event-driven KV cache monitor for overlap scoring (gRPC workers only). /// Set via `set_kv_event_monitor`. When present and the indexer has data for /// a model, event-driven routing takes priority over approximate trees. - kv_monitor: Option>, + /// Uses `RwLock` for interior mutability so the monitor can be injected into + /// policies already behind `Arc`. + kv_monitor: RwLock>>, } impl CacheAwarePolicy { @@ -134,7 +137,7 @@ impl CacheAwarePolicy { token_trees, mesh_sync: None, _eviction_task: eviction_task, - kv_monitor: None, + kv_monitor: RwLock::new(None), } } @@ -146,9 +149,10 @@ impl CacheAwarePolicy { } } - /// Set event-driven KV cache monitor (can be called after construction) - pub fn set_kv_event_monitor(&mut self, monitor: Option>) { - self.kv_monitor = monitor; + /// Set event-driven KV cache monitor (thread-safe, can be called after construction). + /// Uses interior mutability so this works on policies behind `Arc`. + pub fn set_kv_event_monitor(&self, monitor: Option>) { + *self.kv_monitor.write() = monitor; } /// Initialize the trees with worker URLs (used only during initial setup) @@ -486,13 +490,16 @@ impl LoadBalancingPolicy for CacheAwarePolicy { // Private helper methods for select_worker impl CacheAwarePolicy { - /// Check if an event-driven indexer exists for this model. - /// Used to choose between event-driven and approximate tree routing. + /// Check if an event-driven indexer exists with data for this model. + /// Returns false when the indexer is empty (startup, reconnect) so + /// routing falls through to the approximate token tree instead of + /// taking the event-driven path with no data and landing on min-load. fn has_event_indexer(&self, model_id: &str) -> bool { - self.kv_monitor + let guard = self.kv_monitor.read(); + guard .as_ref() .and_then(|m| m.get_indexer(model_id)) - .is_some() + .is_some_and(|indexer| indexer.current_size() > 0) } /// Event-driven routing: PositionalIndexer overlap scoring (Type 1). @@ -507,7 +514,8 @@ impl CacheAwarePolicy { healthy_indices: &[usize], model_id: &str, ) -> Option { - let monitor = self.kv_monitor.as_ref()?; + let guard = self.kv_monitor.read(); + let monitor = guard.as_ref()?; let indexer = monitor.get_indexer(model_id)?; // Per-model block_size: learned from events > config default @@ -1339,7 +1347,7 @@ mod tests { #[test] fn test_event_driven_overlap_selects_cached_worker() { - let mut policy = CacheAwarePolicy::with_config(test_config()); + let policy = CacheAwarePolicy::with_config(test_config()); let workers: Vec> = vec![ Arc::new( BasicWorkerBuilder::new("http://w1:8000") @@ -1376,7 +1384,7 @@ mod tests { #[test] fn test_event_driven_no_overlap_uses_min_load() { - let mut policy = CacheAwarePolicy::with_config(test_config()); + let policy = CacheAwarePolicy::with_config(test_config()); let w1 = BasicWorkerBuilder::new("http://w1:8000") .worker_type(WorkerType::Regular) @@ -1413,7 +1421,7 @@ mod tests { #[test] fn test_event_driven_short_request_uses_min_load() { - let mut policy = CacheAwarePolicy::with_config(test_config()); // block_size=4 + let policy = CacheAwarePolicy::with_config(test_config()); // block_size=4 let w1 = BasicWorkerBuilder::new("http://w1:8000") .worker_type(WorkerType::Regular) @@ -1481,29 +1489,29 @@ mod tests { #[test] fn test_set_kv_event_monitor() { - let mut policy = CacheAwarePolicy::with_config(test_config()); + let policy = CacheAwarePolicy::with_config(test_config()); // Initially no monitor - assert!(policy.kv_monitor.is_none()); + assert!(policy.kv_monitor.read().is_none()); - // Set monitor + // Set monitor (works via &self thanks to interior mutability) let monitor = Arc::new(KvEventMonitor::new(Some(4))); policy.set_kv_event_monitor(Some(Arc::clone(&monitor))); - assert!(policy.kv_monitor.is_some()); + assert!(policy.kv_monitor.read().is_some()); // get_indexer returns None for unknown model assert!(monitor.get_indexer("nonexistent").is_none()); // Clear monitor policy.set_kv_event_monitor(None); - assert!(policy.kv_monitor.is_none()); + assert!(policy.kv_monitor.read().is_none()); } #[test] fn test_event_driven_uses_monitor_block_size() { // Test that event-driven routing uses monitor's learned block_size // instead of config default when available. - let mut policy = CacheAwarePolicy::with_config(CacheAwareConfig { + let policy = CacheAwarePolicy::with_config(CacheAwareConfig { block_size: 4, // config default eviction_interval_secs: 0, ..Default::default() @@ -1559,7 +1567,7 @@ mod tests { #[test] fn test_imbalanced_skips_event_driven() { - let mut policy = CacheAwarePolicy::with_config(CacheAwareConfig { + let policy = CacheAwarePolicy::with_config(CacheAwareConfig { balance_abs_threshold: 5, balance_rel_threshold: 2.0, eviction_interval_secs: 0, @@ -1598,4 +1606,60 @@ mod tests { .unwrap(); assert_eq!(idx, 1); // w2 (min load), regardless of event data } + + #[test] + fn test_empty_indexer_falls_through_to_token_tree() { + // When the monitor has an indexer for a model but the indexer is empty + // (startup, reconnect), routing should fall through to the token tree + // instead of taking the event-driven path and landing on min-load. + let policy = CacheAwarePolicy::with_config(test_config()); + let workers: Vec> = vec![ + Arc::new( + BasicWorkerBuilder::new("http://w1:8000") + .worker_type(WorkerType::Regular) + .build(), + ), + Arc::new( + BasicWorkerBuilder::new("http://w2:8000") + .worker_type(WorkerType::Regular) + .build(), + ), + ]; + policy.init_workers(&workers); + + // Set up monitor with an empty indexer + let monitor = Arc::new(KvEventMonitor::new(Some(4))); + let empty_indexer = Arc::new(PositionalIndexer::new(4)); + monitor + .indexers + .insert("unknown".to_string(), empty_indexer); + policy.set_kv_event_monitor(Some(monitor)); + + // Empty indexer → has_event_indexer returns false → falls through to token tree + assert!(!policy.has_event_indexer("unknown")); + + // Route a request — should use token tree, not event-driven min-load + let idx = policy + .select_worker( + &workers, + &SelectWorkerInfo { + tokens: Some(&[1, 2, 3, 4]), + ..Default::default() + }, + ) + .unwrap(); + assert!(idx < 2); // valid worker via token tree + + // Route the same tokens again — token tree should route to same worker (cache hit) + let idx2 = policy + .select_worker( + &workers, + &SelectWorkerInfo { + tokens: Some(&[1, 2, 3, 4]), + ..Default::default() + }, + ) + .unwrap(); + assert_eq!(idx, idx2); // token tree cache affinity preserved + } } diff --git a/model_gateway/src/policies/registry.rs b/model_gateway/src/policies/registry.rs index 90bd445905..1eecc5f2db 100644 --- a/model_gateway/src/policies/registry.rs +++ b/model_gateway/src/policies/registry.rs @@ -68,10 +68,37 @@ impl PolicyRegistry { *guard = mesh_sync; } - /// Set KV event monitor (thread-safe, can be called after initialization) + /// Set KV event monitor (thread-safe, can be called after initialization). + /// Propagates to all existing cache-aware policies (including default, prefill, decode). pub fn set_kv_event_monitor(&self, monitor: Option>) { - let mut guard = self.kv_event_monitor.write(); - *guard = monitor; + { + let mut guard = self.kv_event_monitor.write(); + guard.clone_from(&monitor); + } + + // Propagate to existing cache-aware policies so they don't miss the monitor. + // This covers the default_policy (created before the monitor was available) + // and any model/PD policies that were already set up. + Self::maybe_inject_monitor(&self.default_policy, monitor.as_ref()); + if let Some(p) = self.prefill_policy.get() { + Self::maybe_inject_monitor(p, monitor.as_ref()); + } + if let Some(p) = self.decode_policy.get() { + Self::maybe_inject_monitor(p, monitor.as_ref()); + } + for entry in self.model_policies.iter() { + Self::maybe_inject_monitor(entry.value(), monitor.as_ref()); + } + } + + /// Inject KV event monitor into a policy if it's cache-aware. + fn maybe_inject_monitor( + policy: &Arc, + monitor: Option<&Arc>, + ) { + if let Some(cache_aware) = policy.as_any().downcast_ref::() { + cache_aware.set_kv_event_monitor(monitor.cloned()); + } } /// Called when a worker is added