diff --git a/sgl-model-gateway/src/core/steps/worker/local/remove_from_policy_registry.rs b/sgl-model-gateway/src/core/steps/worker/local/remove_from_policy_registry.rs index c720e0c47..c103037a2 100644 --- a/sgl-model-gateway/src/core/steps/worker/local/remove_from_policy_registry.rs +++ b/sgl-model-gateway/src/core/steps/worker/local/remove_from_policy_registry.rs @@ -38,11 +38,17 @@ impl StepExecutor for RemoveFromPolicyRegistryStep { let model_id = worker.model_id().to_string(); let worker_url = worker.url(); - // Remove from cache-aware policy + // Remove from cache-aware policy (model_policies; covers the regular Router) app_context .policy_registry .remove_worker_from_cache_aware(&model_id, worker_url); + // PD mode keeps prefill/decode cache-aware policies separate from + // model_policies, so also drop the worker from the matching pool's policy. + app_context + .policy_registry + .remove_pd_worker_from_cache_aware(worker.as_ref()); + // Notify policy registry app_context.policy_registry.on_worker_removed(&model_id); } diff --git a/sgl-model-gateway/src/core/steps/worker/shared/update_policies.rs b/sgl-model-gateway/src/core/steps/worker/shared/update_policies.rs index bc586d206..cefb2d8af 100644 --- a/sgl-model-gateway/src/core/steps/worker/shared/update_policies.rs +++ b/sgl-model-gateway/src/core/steps/worker/shared/update_policies.rs @@ -141,6 +141,23 @@ impl StepExecutor for UpdatePolicie } } + // Initialize cache-aware policies for PD mode (prefill_policy / decode_policy + // are separate instances from model_policies and are not touched by the + // per-model loop above). `init_workers` is idempotent so re-running on each + // registration is safe. + let decode_workers = app_context.worker_registry.get_decode_workers(); + let prefill_is_cache_aware = + app_context.policy_registry.get_prefill_policy().name() == "cache_aware"; + let decode_is_cache_aware = + app_context.policy_registry.get_decode_policy().name() == "cache_aware"; + if (prefill_is_cache_aware && !prefill_workers.is_empty()) + || (decode_is_cache_aware && !decode_workers.is_empty()) + { + app_context + .policy_registry + .init_pd_cache_aware_policies(&prefill_workers, &decode_workers); + } + debug!( "Updated policies for {} workers across {} models", workers.len(), diff --git a/sgl-model-gateway/src/policies/cache_aware.rs b/sgl-model-gateway/src/policies/cache_aware.rs index dfc3ef646..4a67d9dcc 100644 --- a/sgl-model-gateway/src/policies/cache_aware.rs +++ b/sgl-model-gateway/src/policies/cache_aware.rs @@ -71,13 +71,40 @@ use super::{ get_healthy_worker_indices, normalize_model_key, tree::Tree, utils::PeriodicTask, CacheAwareConfig, LoadBalancingPolicy, SelectWorkerInfo, }; -use crate::core::{Worker, UNKNOWN_MODEL_ID}; +use crate::core::{Worker, WorkerType, UNKNOWN_MODEL_ID}; + +/// Tag used to isolate prefill/decode/regular worker pools in the cache_aware tree key. +/// +/// Trees are keyed by `pool::model` so that an alternating prefill→decode call sequence +/// for the same model cannot evict each other's tenants. Without this isolation, the +/// `tree.insert(text, url)` at the end of every `select_worker` call would overwrite +/// the previous pool's tenant for the same prompt and collapse cache_aware into a +/// flip-flop between pools. +fn pool_tag(worker_type: &WorkerType) -> &'static str { + match worker_type { + WorkerType::Regular => "regular", + WorkerType::Prefill { .. } => "prefill", + WorkerType::Decode => "decode", + } +} + +fn make_tree_key(pool: &str, model: &str) -> String { + format!("{}::{}", pool, model) +} + +fn tree_key_for_worker(worker: &dyn Worker) -> String { + make_tree_key( + pool_tag(worker.worker_type()), + normalize_model_key(worker.model_id()), + ) +} /// Cache-aware routing policy /// /// Routes requests based on cache affinity when load is balanced, /// switches to shortest-queue routing when load is imbalanced. -/// Maintains separate trees per model for multi-model support. +/// Maintains separate trees per `(pool, model)` so that prefill, decode, and +/// regular worker pools cannot evict each other's tenants. /// Supports mesh synchronization of tree operations across cluster nodes. /// When mesh is not enabled, the policy works independently without synchronization. #[derive(Debug)] @@ -106,13 +133,13 @@ impl CacheAwarePolicy { "Eviction", move || { for tree_ref in trees_clone.iter() { - let model_id = tree_ref.key(); + let tree_key = tree_ref.key(); let tree = tree_ref.value(); tree.evict_tenant_by_size(max_tree_size); debug!( - "Cache eviction completed for model {}, max_size: {}", - model_id, max_tree_size + "Cache eviction completed for {}, max_size: {}", + tree_key, max_tree_size ); } }, @@ -139,24 +166,22 @@ impl CacheAwarePolicy { /// Initialize the tree with worker URLs (used only during initial setup) pub fn init_workers(&self, workers: &[Arc]) { - // Group workers by model - let mut model_workers: std::collections::HashMap>> = + // Group workers by (pool, model) so each pool gets its own isolated tree. + let mut grouped: std::collections::HashMap>> = std::collections::HashMap::new(); for worker in workers { - let tree_key = normalize_model_key(worker.model_id()); - model_workers - .entry(tree_key.to_string()) + grouped + .entry(tree_key_for_worker(worker.as_ref())) .or_default() .push(worker); } - // Initialize tree for each model - for (tree_key, model_workers) in model_workers { + for (tree_key, pool_workers) in grouped { let tree = self .trees .entry(tree_key) .or_insert_with(|| Arc::new(Tree::new())); - for worker in model_workers { + for worker in pool_workers { tree.insert("", worker.url()); } } @@ -164,27 +189,18 @@ impl CacheAwarePolicy { /// Add a single worker to the tree (incremental update) pub fn add_worker(&self, worker: &dyn Worker) { - let tree_key = normalize_model_key(worker.model_id()); + let tree_key = tree_key_for_worker(worker); let tree = self .trees - .entry(tree_key.to_string()) + .entry(tree_key) .or_insert_with(|| Arc::new(Tree::new())); tree.insert("", worker.url()); } - /// Add a worker by URL and model (for backward compatibility) - pub fn add_worker_by_url(&self, url: &str, model_id: &str) { - let tree = self - .trees - .entry(model_id.to_string()) - .or_insert_with(|| Arc::new(Tree::new())); - tree.insert("", url); - } - /// Remove a worker from the tree pub fn remove_worker(&self, worker: &dyn Worker) { - let tree_key = normalize_model_key(worker.model_id()); - if let Some(tree) = self.trees.get(tree_key) { + let tree_key = tree_key_for_worker(worker); + if let Some(tree) = self.trees.get(&tree_key) { tree.remove_tenant(worker.url()); } } @@ -207,11 +223,11 @@ impl CacheAwarePolicy { // In a full implementation, we might want to query mesh for all tree states for tree_ref in self.trees.iter() { - let model_id = tree_ref.key(); - if let Some(tree_state) = mesh_sync.get_tree_state(model_id) { + let tree_key = tree_ref.key(); + if let Some(tree_state) = mesh_sync.get_tree_state(tree_key) { debug!( - "Restoring tree state for model {} with {} operations", - model_id, + "Restoring tree state for {} with {} operations", + tree_key, tree_state.operations.len() ); @@ -232,20 +248,30 @@ impl CacheAwarePolicy { } } - /// Normalize model_id for mesh synchronization - /// Converts empty model_id to UNKNOWN_MODEL_ID for consistency - fn normalize_mesh_model_id(model_id: &str) -> &str { - if model_id.is_empty() { + /// Normalize a tree key for mesh synchronization, converting an accidentally + /// empty key to `UNKNOWN_MODEL_ID` for consistency. In current code the + /// composite `pool::model` key is never empty, so this is defensive. + fn normalize_mesh_model_id(tree_key: &str) -> &str { + if tree_key.is_empty() { UNKNOWN_MODEL_ID } else { - model_id + tree_key } } - /// Apply remote tree operation from mesh - /// This is called when receiving tree state updates from other nodes - pub fn apply_remote_tree_operation(&self, model_id: &str, operation: &TreeOperation) { - let tree_key = Self::normalize_mesh_model_id(model_id); + /// Apply remote tree operation from mesh. + /// + /// `mesh_key` is the opaque key the operation was originally synced under; + /// `select_worker` / `select_worker_min_load` send tree operations to mesh + /// keyed by the composite `pool::model`, and any future receive path is + /// expected to forward that same string back here unchanged. The argument + /// is kept as `&str` so the mesh layer can stay key-agnostic. + /// + /// Note: `PolicyRegistry::apply_remote_tree_operation` (the only forwarder) + /// currently has no in-process callers; the receive path is not yet wired, + /// so this method is reachable only via tests today. + pub fn apply_remote_tree_operation(&self, mesh_key: &str, operation: &TreeOperation) { + let tree_key = Self::normalize_mesh_model_id(mesh_key); let tree = self .trees @@ -256,15 +282,15 @@ impl CacheAwarePolicy { TreeOperation::Insert(insert_op) => { tree.insert(&insert_op.text, &insert_op.tenant); debug!( - "Applied remote tree insert: model={}, text={}, tenant={}", - model_id, insert_op.text, insert_op.tenant + "Applied remote tree insert: key={}, text={}, tenant={}", + mesh_key, insert_op.text, insert_op.tenant ); } TreeOperation::Remove(remove_op) => { tree.remove_tenant(&remove_op.tenant); debug!( - "Applied remote tree remove: model={}, tenant={}", - model_id, remove_op.tenant + "Applied remote tree remove: key={}, tenant={}", + mesh_key, remove_op.tenant ); } } @@ -273,13 +299,10 @@ impl CacheAwarePolicy { /// Run cache eviction to prevent unbounded growth pub fn evict_cache(&self, max_size: usize) { for tree_ref in self.trees.iter() { - let model_id = tree_ref.key(); + let tree_key = tree_ref.key(); let tree = tree_ref.value(); tree.evict_tenant_by_size(max_size); - debug!( - "Cache eviction for model {}, max_size: {}", - model_id, max_size - ); + debug!("Cache eviction for {}, max_size: {}", tree_key, max_size); } } @@ -288,7 +311,7 @@ impl CacheAwarePolicy { workers: &[Arc], request_text: &Option<&str>, healthy_indices: &[usize], - model_id: &str, + tree_key: &str, max_load: usize, min_load: usize, ) -> Option { @@ -312,7 +335,7 @@ impl CacheAwarePolicy { if let Some(text) = request_text { // Get the tree reference without locking the entire HashMap // DashMap only locks the specific shard containing this key - let tree = self.trees.get(model_id).map(|entry| entry.value().clone()); + let tree = self.trees.get(tree_key).map(|entry| entry.value().clone()); if let Some(tree) = tree { let worker_url = workers[min_load_idx].url(); @@ -326,15 +349,17 @@ impl CacheAwarePolicy { text: text.to_string(), tenant: worker_url.to_string(), }); - let mesh_model_id = Self::normalize_mesh_model_id(model_id); - if let Err(e) = mesh_sync.sync_tree_operation(mesh_model_id.to_string(), op) { + let mesh_key = Self::normalize_mesh_model_id(tree_key); + if let Err(e) = mesh_sync.sync_tree_operation(mesh_key.to_string(), op) { warn!("Failed to sync tree insert operation to mesh: {}", e); } } } else { - debug!( - "Warning: No tree found for model '{}', skipping cache update", - model_id + warn!( + "cache_aware: no tree found for key '{}', skipping cache update — \ + pool tree was not seeded (init_pd_cache_aware_policies missed or \ + a race during worker registration)", + tree_key ); } } @@ -360,9 +385,10 @@ impl LoadBalancingPolicy for CacheAwarePolicy { return None; } - // Determine the model for this set of workers (router pre-filters by model) - // All workers should be from the same model - let model_id = normalize_model_key(workers[healthy_indices[0]].model_id()); + // Determine the (pool, model) key for this set of workers — the router pre-filters + // so every healthy worker here belongs to the same pool and same model. + let pivot = workers[healthy_indices[0]].as_ref(); + let tree_key = tree_key_for_worker(pivot); // Get current load statistics - compute min/max in single pass without allocation let (min_load, max_load) = workers.iter().fold((usize::MAX, 0usize), |(min, max), w| { @@ -380,7 +406,7 @@ impl LoadBalancingPolicy for CacheAwarePolicy { workers, &request_text, &healthy_indices, - model_id, + &tree_key, max_load, min_load, ); @@ -391,7 +417,7 @@ impl LoadBalancingPolicy for CacheAwarePolicy { // Get the tree reference without locking the entire HashMap // DashMap only locks the specific shard containing this key - let tree = self.trees.get(model_id).map(|entry| entry.value().clone()); + let tree = self.trees.get(&tree_key).map(|entry| entry.value().clone()); if let Some(tree) = tree { // Now we work with the tree without holding the HashMap lock @@ -430,8 +456,8 @@ impl LoadBalancingPolicy for CacheAwarePolicy { text: text.to_string(), tenant: workers[idx].url().to_string(), }); - let mesh_model_id = Self::normalize_mesh_model_id(model_id); - if let Err(e) = mesh_sync.sync_tree_operation(mesh_model_id.to_string(), op) { + let mesh_key = Self::normalize_mesh_model_id(&tree_key); + if let Err(e) = mesh_sync.sync_tree_operation(mesh_key.to_string(), op) { warn!("Failed to sync tree insert operation to mesh: {}", e); } } @@ -454,8 +480,8 @@ impl LoadBalancingPolicy for CacheAwarePolicy { let op = TreeOperation::Remove(TreeRemoveOp { tenant: tenant_url.to_string(), }); - let mesh_model_id = Self::normalize_mesh_model_id(model_id); - if let Err(e) = mesh_sync.sync_tree_operation(mesh_model_id.to_string(), op) { + let mesh_key = Self::normalize_mesh_model_id(&tree_key); + if let Err(e) = mesh_sync.sync_tree_operation(mesh_key.to_string(), op) { warn!("Failed to sync tree remove operation to mesh: {}", e); } } @@ -464,12 +490,14 @@ impl LoadBalancingPolicy for CacheAwarePolicy { // Fallback to first healthy worker healthy_indices.first().copied() } else { - // No tree for this model, log warning and use random selection - debug!( - "Warning: No tree found for model '{}', using random worker selection", - model_id + warn!( + "cache_aware: no tree found for key '{}', falling back to random \ + worker selection — pool tree was not seeded \ + (init_pd_cache_aware_policies missed or a race during worker \ + registration); cache affinity is effectively disabled until this \ + clears", + tree_key ); - // Return a random healthy worker let mut rng = rand::rng(); let random_idx = rng.random_range(0..healthy_indices.len()); Some(healthy_indices[random_idx]) @@ -711,8 +739,11 @@ mod tests { .await .unwrap(); - // Verify tree operation was synced to mesh (under UNKNOWN_MODEL_ID since no model was specified) - let tree_state = mesh_sync.get_tree_state(UNKNOWN_MODEL_ID); + // Verify tree operation was synced to mesh under the composite `pool::model` + // key — workers here are Regular and no model was specified, so the key is + // `regular::UNKNOWN_MODEL_ID`. + let expected_key = format!("regular::{}", UNKNOWN_MODEL_ID); + let tree_state = mesh_sync.get_tree_state(&expected_key); assert!(tree_state.is_some()); let tree = tree_state.unwrap(); assert!(!tree.operations.is_empty()); @@ -884,4 +915,597 @@ mod tests { .unwrap(); assert_eq!(idx, 0); } + + fn make_prefill(url: &str) -> Arc { + Arc::new( + BasicWorkerBuilder::new(url) + .worker_type(WorkerType::Prefill { + bootstrap_port: Some(9000), + }) + .build(), + ) + } + + fn make_decode(url: &str) -> Arc { + Arc::new( + BasicWorkerBuilder::new(url) + .worker_type(WorkerType::Decode) + .build(), + ) + } + + /// PD setup with two separate `CacheAwarePolicy` instances — the production + /// wiring. Each pool's tree is seeded only with its own workers. Across a + /// 4-turn growing prompt, each pool must stick to one worker. + #[tokio::test] + async fn test_pd_pool_isolation_two_policies() { + let config = CacheAwareConfig { + eviction_interval_secs: 0, + ..Default::default() + }; + let prefill_policy = CacheAwarePolicy::with_config(config.clone()); + let decode_policy = CacheAwarePolicy::with_config(config); + + let prefill_workers: Vec> = vec![ + make_prefill("http://prefill0:8000"), + make_prefill("http://prefill1:8000"), + ]; + let decode_workers: Vec> = vec![ + make_decode("http://decode0:8000"), + make_decode("http://decode1:8000"), + ]; + prefill_policy.init_workers(&prefill_workers); + decode_policy.init_workers(&decode_workers); + + let turns = [ + "turn1", + "turn1 turn2", + "turn1 turn2 turn3", + "turn1 turn2 turn3 turn4", + ]; + + let mut prefill_idx: Option = None; + let mut decode_idx: Option = None; + for prompt in turns { + let info = SelectWorkerInfo { + request_text: Some(prompt), + ..Default::default() + }; + let p = prefill_policy + .select_worker(&prefill_workers, &info) + .await + .expect("prefill pool returns a worker"); + let d = decode_policy + .select_worker(&decode_workers, &info) + .await + .expect("decode pool returns a worker"); + match prefill_idx { + None => prefill_idx = Some(p), + Some(pinned) => assert_eq!( + p, pinned, + "prefill should stay pinned across turns (prompt={prompt:?})" + ), + } + match decode_idx { + None => decode_idx = Some(d), + Some(pinned) => assert_eq!( + d, pinned, + "decode should stay pinned across turns (prompt={prompt:?})" + ), + } + } + } + + /// Regression: even if a single `CacheAwarePolicy` instance is incorrectly + /// wired to both pools, pool-aware tree keys must keep their state disjoint. + /// The pre-fix code shared one trie keyed by `model_id`, so alternating + /// prefill/decode `tree.insert` calls overwrote each other and the policy + /// degenerated into worker-flipping random selection. + #[tokio::test] + async fn test_pd_pool_isolation_shared_policy_regression() { + let config = CacheAwareConfig { + eviction_interval_secs: 0, + ..Default::default() + }; + let policy = CacheAwarePolicy::with_config(config); + + let prefill_workers: Vec> = vec![ + make_prefill("http://prefill0:8000"), + make_prefill("http://prefill1:8000"), + ]; + let decode_workers: Vec> = vec![ + make_decode("http://decode0:8000"), + make_decode("http://decode1:8000"), + ]; + + // One instance, mixed init — pool-aware keys split the trees internally. + let mut combined: Vec> = Vec::new(); + combined.extend(prefill_workers.iter().cloned()); + combined.extend(decode_workers.iter().cloned()); + policy.init_workers(&combined); + + let turns = [ + "turn1", + "turn1 turn2", + "turn1 turn2 turn3", + "turn1 turn2 turn3 turn4", + ]; + + let mut prefill_idx: Option = None; + let mut decode_idx: Option = None; + for prompt in turns { + let info = SelectWorkerInfo { + request_text: Some(prompt), + ..Default::default() + }; + let p = policy + .select_worker(&prefill_workers, &info) + .await + .expect("prefill pool returns a worker"); + let d = policy + .select_worker(&decode_workers, &info) + .await + .expect("decode pool returns a worker"); + + assert!( + prefill_workers[p].url().starts_with("http://prefill"), + "prefill call must return a prefill index, got {} (prompt={prompt:?})", + prefill_workers[p].url() + ); + assert!( + decode_workers[d].url().starts_with("http://decode"), + "decode call must return a decode index, got {} (prompt={prompt:?})", + decode_workers[d].url() + ); + + match prefill_idx { + None => prefill_idx = Some(p), + Some(pinned) => assert_eq!( + p, pinned, + "prefill should stay pinned across turns (prompt={prompt:?})" + ), + } + match decode_idx { + None => decode_idx = Some(d), + Some(pinned) => assert_eq!( + d, pinned, + "decode should stay pinned across turns (prompt={prompt:?})" + ), + } + } + } + + /// Removing a PD worker via the composite-key `remove_worker(&dyn Worker)` path + /// must drop it from its own pool's tree without touching the other pool. This + /// covers `PolicyRegistry::remove_pd_worker_from_cache_aware`, which routes the + /// removal here based on `worker.worker_type()`. + #[tokio::test] + async fn test_pd_pool_isolation_remove_worker() { + let config = CacheAwareConfig { + eviction_interval_secs: 0, + ..Default::default() + }; + let policy = CacheAwarePolicy::with_config(config); + + let prefill0 = make_prefill("http://prefill0:8000"); + let prefill1 = make_prefill("http://prefill1:8000"); + let decode0 = make_decode("http://decode0:8000"); + let decode1 = make_decode("http://decode1:8000"); + + let prefill_workers: Vec> = vec![prefill0.clone(), prefill1.clone()]; + let decode_workers: Vec> = vec![decode0.clone(), decode1.clone()]; + let combined: Vec> = vec![ + prefill0.clone(), + prefill1.clone(), + decode0.clone(), + decode1.clone(), + ]; + policy.init_workers(&combined); + + // Seed both trees with affinity for one prompt. + let prompt = "shared prefix to seed the cache_aware trees"; + let info = SelectWorkerInfo { + request_text: Some(prompt), + ..Default::default() + }; + policy + .select_worker(&prefill_workers, &info) + .await + .expect("seed prefill"); + policy + .select_worker(&decode_workers, &info) + .await + .expect("seed decode"); + + let prefill_key = format!("prefill::{}", UNKNOWN_MODEL_ID); + let decode_key = format!("decode::{}", UNKNOWN_MODEL_ID); + + let prefill_before = policy + .trees + .get(&prefill_key) + .map(|t| t.value().prefix_match_with_counts(prompt).tenant) + .expect("prefill tree seeded"); + let decode_before = policy + .trees + .get(&decode_key) + .map(|t| t.value().prefix_match_with_counts(prompt).tenant) + .expect("decode tree seeded"); + assert!( + prefill_before.starts_with("http://prefill"), + "prefill tree should hold a prefill tenant before removal, got {prefill_before}" + ); + assert!( + decode_before.starts_with("http://decode"), + "decode tree should hold a decode tenant before removal, got {decode_before}" + ); + + // Drop prefill0 via the composite-key removal path. + policy.remove_worker(prefill0.as_ref()); + + // The prefill tree must no longer point at prefill0 for the seeded prompt. + let prefill_after = policy + .trees + .get(&prefill_key) + .map(|t| t.value().prefix_match_with_counts(prompt).tenant) + .expect("prefill tree still exists"); + assert_ne!( + &*prefill_after, + prefill0.url(), + "prefill0 should be gone from the prefill tree" + ); + + // The decode tree must be byte-for-byte unchanged. + let decode_after = policy + .trees + .get(&decode_key) + .map(|t| t.value().prefix_match_with_counts(prompt).tenant) + .expect("decode tree still exists"); + assert_eq!( + decode_after, decode_before, + "removing a prefill worker must not touch the decode tree" + ); + } + + /// Shared setup for `PolicyRegistry::remove_pd_worker_from_cache_aware` tests: + /// build a registry whose prefill and decode policies are separate + /// `CacheAwarePolicy` instances seeded with the matching pool's workers, then + /// return the registry, the per-pool policy handles (for tree inspection), and + /// representative workers from each pool. + #[allow(clippy::type_complexity)] + fn pd_registry_with_cache_aware_pools() -> ( + Arc, + Arc, + Arc, + Arc, + Arc, + Arc, + Arc, + ) { + let registry = Arc::new(crate::policies::PolicyRegistry::new( + crate::config::types::PolicyConfig::RoundRobin, + )); + let no_eviction = CacheAwareConfig { + eviction_interval_secs: 0, + ..Default::default() + }; + let prefill_ca = Arc::new(CacheAwarePolicy::with_config(no_eviction.clone())); + let decode_ca = Arc::new(CacheAwarePolicy::with_config(no_eviction)); + registry.set_prefill_policy(prefill_ca.clone() as Arc); + registry.set_decode_policy(decode_ca.clone() as Arc); + + let prefill0 = make_prefill("http://prefill0:8000"); + let prefill1 = make_prefill("http://prefill1:8000"); + let decode0 = make_decode("http://decode0:8000"); + let decode1 = make_decode("http://decode1:8000"); + + let prefill_workers: Vec> = vec![prefill0.clone(), prefill1.clone()]; + let decode_workers: Vec> = vec![decode0.clone(), decode1.clone()]; + registry.init_pd_cache_aware_policies(&prefill_workers, &decode_workers); + + ( + registry, prefill_ca, decode_ca, prefill0, prefill1, decode0, decode1, + ) + } + + /// Seed both pool trees so each has a known tenant for `prompt`, then return + /// the (prefill_tenant, decode_tenant) snapshot to compare against after a + /// dispatched removal. + async fn seed_pd_pools( + prefill_ca: &CacheAwarePolicy, + decode_ca: &CacheAwarePolicy, + prefill_workers: &[Arc], + decode_workers: &[Arc], + prompt: &str, + ) -> (Arc, Arc) { + let info = SelectWorkerInfo { + request_text: Some(prompt), + ..Default::default() + }; + prefill_ca + .select_worker(prefill_workers, &info) + .await + .expect("prefill seed"); + decode_ca + .select_worker(decode_workers, &info) + .await + .expect("decode seed"); + + let prefill_key = format!("prefill::{}", UNKNOWN_MODEL_ID); + let decode_key = format!("decode::{}", UNKNOWN_MODEL_ID); + let prefill_tenant = prefill_ca + .trees + .get(&prefill_key) + .map(|t| t.value().prefix_match_with_counts(prompt).tenant) + .expect("prefill tree seeded"); + let decode_tenant = decode_ca + .trees + .get(&decode_key) + .map(|t| t.value().prefix_match_with_counts(prompt).tenant) + .expect("decode tree seeded"); + (prefill_tenant, decode_tenant) + } + + /// A `Prefill` worker passed to `remove_pd_worker_from_cache_aware` must hit + /// the registry's `prefill_policy` and leave `decode_policy` untouched. + /// Catches a dispatch swap like `Prefill => self.decode_policy.get()`. + #[tokio::test] + async fn test_registry_remove_pd_worker_prefill_dispatches_to_prefill_policy() { + let (registry, prefill_ca, decode_ca, prefill0, prefill1, decode0, decode1) = + pd_registry_with_cache_aware_pools(); + let prefill_workers: Vec> = vec![prefill0.clone(), prefill1.clone()]; + let decode_workers: Vec> = vec![decode0.clone(), decode1.clone()]; + + let prompt = "prefix used to seed both pool trees"; + let (_prefill_before, decode_before) = seed_pd_pools( + &prefill_ca, + &decode_ca, + &prefill_workers, + &decode_workers, + prompt, + ) + .await; + + registry.remove_pd_worker_from_cache_aware(prefill0.as_ref()); + + let prefill_key = format!("prefill::{}", UNKNOWN_MODEL_ID); + let decode_key = format!("decode::{}", UNKNOWN_MODEL_ID); + let prefill_after = prefill_ca + .trees + .get(&prefill_key) + .map(|t| t.value().prefix_match_with_counts(prompt).tenant) + .expect("prefill tree still exists"); + assert_ne!( + &*prefill_after, + prefill0.url(), + "registry dispatch must drop prefill0 from the prefill pool's tree" + ); + let decode_after = decode_ca + .trees + .get(&decode_key) + .map(|t| t.value().prefix_match_with_counts(prompt).tenant) + .expect("decode tree still exists"); + assert_eq!( + decode_after, decode_before, + "removing a prefill worker must not touch the decode pool's tree" + ); + } + + /// Mirror of the prefill dispatch test for `Decode`. Catches a dispatch swap + /// in the other direction (`Decode => self.prefill_policy.get()`). + #[tokio::test] + async fn test_registry_remove_pd_worker_decode_dispatches_to_decode_policy() { + let (registry, prefill_ca, decode_ca, prefill0, prefill1, decode0, decode1) = + pd_registry_with_cache_aware_pools(); + let prefill_workers: Vec> = vec![prefill0.clone(), prefill1.clone()]; + let decode_workers: Vec> = vec![decode0.clone(), decode1.clone()]; + + let prompt = "prefix used to seed both pool trees"; + let (prefill_before, _decode_before) = seed_pd_pools( + &prefill_ca, + &decode_ca, + &prefill_workers, + &decode_workers, + prompt, + ) + .await; + + registry.remove_pd_worker_from_cache_aware(decode0.as_ref()); + + let prefill_key = format!("prefill::{}", UNKNOWN_MODEL_ID); + let decode_key = format!("decode::{}", UNKNOWN_MODEL_ID); + let decode_after = decode_ca + .trees + .get(&decode_key) + .map(|t| t.value().prefix_match_with_counts(prompt).tenant) + .expect("decode tree still exists"); + assert_ne!( + &*decode_after, + decode0.url(), + "registry dispatch must drop decode0 from the decode pool's tree" + ); + let prefill_after = prefill_ca + .trees + .get(&prefill_key) + .map(|t| t.value().prefix_match_with_counts(prompt).tenant) + .expect("prefill tree still exists"); + assert_eq!( + prefill_after, prefill_before, + "removing a decode worker must not touch the prefill pool's tree" + ); + } + + /// `remove_pd_worker_from_cache_aware` must short-circuit on `Regular` + /// workers and silently ignore non-cache_aware policies (`name() != "cache_aware"`). + /// Both branches are no-ops: neither pool tree changes, and no downcast panic. + #[tokio::test] + async fn test_registry_remove_pd_worker_regular_and_non_cache_aware_noop() { + // (a) Regular worker: should early-return regardless of policy state. + let (registry, prefill_ca, decode_ca, prefill0, prefill1, decode0, decode1) = + pd_registry_with_cache_aware_pools(); + let prefill_workers: Vec> = vec![prefill0.clone(), prefill1.clone()]; + let decode_workers: Vec> = vec![decode0.clone(), decode1.clone()]; + + let prompt = "regular-noop seed prompt"; + let (prefill_before, decode_before) = seed_pd_pools( + &prefill_ca, + &decode_ca, + &prefill_workers, + &decode_workers, + prompt, + ) + .await; + + let regular: Arc = Arc::new( + BasicWorkerBuilder::new("http://regular0:8000") + .worker_type(WorkerType::Regular) + .build(), + ); + registry.remove_pd_worker_from_cache_aware(regular.as_ref()); + + let prefill_key = format!("prefill::{}", UNKNOWN_MODEL_ID); + let decode_key = format!("decode::{}", UNKNOWN_MODEL_ID); + let prefill_after = prefill_ca + .trees + .get(&prefill_key) + .map(|t| t.value().prefix_match_with_counts(prompt).tenant) + .expect("prefill tree still exists"); + let decode_after = decode_ca + .trees + .get(&decode_key) + .map(|t| t.value().prefix_match_with_counts(prompt).tenant) + .expect("decode tree still exists"); + assert_eq!( + prefill_after, prefill_before, + "Regular worker dispatch must not touch the prefill tree" + ); + assert_eq!( + decode_after, decode_before, + "Regular worker dispatch must not touch the decode tree" + ); + + // (b) Non-cache_aware policy: PD pool is round_robin. The downcast must + // be skipped (no panic) and the call must be a no-op. + let registry = + crate::policies::PolicyRegistry::new(crate::config::types::PolicyConfig::RoundRobin); + let rr_prefill: Arc = + Arc::new(crate::policies::RoundRobinPolicy::new()); + let rr_decode: Arc = + Arc::new(crate::policies::RoundRobinPolicy::new()); + registry.set_prefill_policy(rr_prefill); + registry.set_decode_policy(rr_decode); + // No panic, no downcast — this would fault if the guard + // `policy.name() == "cache_aware"` were dropped. + registry.remove_pd_worker_from_cache_aware(prefill0.as_ref()); + registry.remove_pd_worker_from_cache_aware(decode0.as_ref()); + } + + /// `init_pd_cache_aware_policies` must seed only the pool whose policy is + /// cache_aware AND whose worker list is non-empty. Covers all four corners: + /// both seeded, only-prefill-cache_aware, empty-worker short-circuit, and the + /// non-cache_aware side staying a no-op. + #[tokio::test] + async fn test_registry_init_pd_cache_aware_policies_gating() { + let no_eviction = CacheAwareConfig { + eviction_interval_secs: 0, + ..Default::default() + }; + let prefill_key = format!("prefill::{}", UNKNOWN_MODEL_ID); + let decode_key = format!("decode::{}", UNKNOWN_MODEL_ID); + + let prefill0 = make_prefill("http://prefill0:8000"); + let decode0 = make_decode("http://decode0:8000"); + let prefill_workers: Vec> = vec![prefill0.clone()]; + let decode_workers: Vec> = vec![decode0.clone()]; + + // (a) Both pools are cache_aware with workers → both trees seeded under + // the correct composite key. + { + let registry = crate::policies::PolicyRegistry::new( + crate::config::types::PolicyConfig::RoundRobin, + ); + let prefill_ca = Arc::new(CacheAwarePolicy::with_config(no_eviction.clone())); + let decode_ca = Arc::new(CacheAwarePolicy::with_config(no_eviction.clone())); + registry.set_prefill_policy(prefill_ca.clone() as Arc); + registry.set_decode_policy(decode_ca.clone() as Arc); + + registry.init_pd_cache_aware_policies(&prefill_workers, &decode_workers); + + assert!( + prefill_ca.trees.contains_key(&prefill_key), + "prefill cache_aware policy must be seeded under '{prefill_key}'" + ); + assert!( + decode_ca.trees.contains_key(&decode_key), + "decode cache_aware policy must be seeded under '{decode_key}'" + ); + assert!( + !prefill_ca.trees.contains_key(&decode_key), + "prefill_workers must not seed the decode tree key" + ); + assert!( + !decode_ca.trees.contains_key(&prefill_key), + "decode_workers must not seed the prefill tree key" + ); + } + + // (b) Only prefill is cache_aware (decode is round_robin) → prefill seeded, + // decode side skipped silently (no downcast, no panic). + { + let registry = crate::policies::PolicyRegistry::new( + crate::config::types::PolicyConfig::RoundRobin, + ); + let prefill_ca = Arc::new(CacheAwarePolicy::with_config(no_eviction.clone())); + let decode_rr: Arc = + Arc::new(crate::policies::RoundRobinPolicy::new()); + registry.set_prefill_policy(prefill_ca.clone() as Arc); + registry.set_decode_policy(decode_rr); + + registry.init_pd_cache_aware_policies(&prefill_workers, &decode_workers); + + assert!( + prefill_ca.trees.contains_key(&prefill_key), + "prefill cache_aware side must seed even when decode side is non-cache_aware" + ); + } + + // (c) Both cache_aware but prefill worker list is empty → prefill tree + // NOT seeded (the inner `!is_empty()` guard short-circuits); decode side + // is still seeded. + { + let registry = crate::policies::PolicyRegistry::new( + crate::config::types::PolicyConfig::RoundRobin, + ); + let prefill_ca = Arc::new(CacheAwarePolicy::with_config(no_eviction.clone())); + let decode_ca = Arc::new(CacheAwarePolicy::with_config(no_eviction.clone())); + registry.set_prefill_policy(prefill_ca.clone() as Arc); + registry.set_decode_policy(decode_ca.clone() as Arc); + + registry.init_pd_cache_aware_policies(&[], &decode_workers); + + assert!( + prefill_ca.trees.is_empty(), + "empty prefill worker list must not seed the prefill tree" + ); + assert!( + decode_ca.trees.contains_key(&decode_key), + "decode side must still seed when only the prefill list is empty" + ); + } + + // (d) Both worker lists empty → neither pool seeded (init is a full no-op). + { + let registry = crate::policies::PolicyRegistry::new( + crate::config::types::PolicyConfig::RoundRobin, + ); + let prefill_ca = Arc::new(CacheAwarePolicy::with_config(no_eviction.clone())); + let decode_ca = Arc::new(CacheAwarePolicy::with_config(no_eviction)); + registry.set_prefill_policy(prefill_ca.clone() as Arc); + registry.set_decode_policy(decode_ca.clone() as Arc); + + registry.init_pd_cache_aware_policies(&[], &[]); + + assert!(prefill_ca.trees.is_empty()); + assert!(decode_ca.trees.is_empty()); + } + } } diff --git a/sgl-model-gateway/src/policies/registry.rs b/sgl-model-gateway/src/policies/registry.rs index 2e0e4b9eb..e34156cd9 100644 --- a/sgl-model-gateway/src/policies/registry.rs +++ b/sgl-model-gateway/src/policies/registry.rs @@ -369,6 +369,31 @@ impl PolicyRegistry { } } + /// Remove a PD worker from the matching pool's cache-aware policy (lock-free). + /// + /// `prefill_policy` and `decode_policy` are separate `CacheAwarePolicy` instances + /// from `model_policies`, so the regular `remove_worker_from_cache_aware` does not + /// touch them. This dispatches to the right pool based on `worker.worker_type()`. + pub fn remove_pd_worker_from_cache_aware(&self, worker: &dyn Worker) { + let policy = match worker.worker_type() { + crate::core::WorkerType::Prefill { .. } => self.prefill_policy.get(), + crate::core::WorkerType::Decode => self.decode_policy.get(), + crate::core::WorkerType::Regular => return, + }; + if let Some(policy) = policy { + if policy.name() == "cache_aware" { + if let Some(cache_aware) = policy.as_any().downcast_ref::() { + cache_aware.remove_worker(worker); + debug!( + "Removed PD worker {} ({}) from cache-aware policy", + worker.url(), + worker.worker_type() + ); + } + } + } + } + /// Initialize bucket policies for PD mode - lock-free pub fn init_pd_bucket_policies(&self, prefill_workers: &[Arc]) { // Initialize prefill policy if it's bucket (lock-free via OnceLock::get)