diff --git a/model_gateway/src/policies/round_robin.rs b/model_gateway/src/policies/round_robin.rs index 6ca4ba0cd..eacd162fe 100644 --- a/model_gateway/src/policies/round_robin.rs +++ b/model_gateway/src/policies/round_robin.rs @@ -165,4 +165,51 @@ mod tests { policy.reset(); assert_eq!(policy.select_worker(&workers, &info), Some(0)); } + + #[test] + fn test_independent_policies_cover_all_workers_across_two_pools() { + fn make_workers(prefix: &str, n: usize) -> Vec> { + (0..n) + .map(|i| { + Arc::new( + BasicWorkerBuilder::new(format!("http://{prefix}{i}:8000")) + .worker_type(WorkerType::Regular) + .health_config(no_health_check()) + .build(), + ) as Arc + }) + .collect() + } + + let prefill_workers = make_workers("p", 4); + let decode_workers = make_workers("d", 4); + let info = SelectWorkerInfo::default(); + + let shared = RoundRobinPolicy::new(); + let mut shared_prefill = [0usize; 4]; + let mut shared_decode = [0usize; 4]; + for _ in 0..40 { + let p = shared.select_worker(&prefill_workers, &info).unwrap(); + let d = shared.select_worker(&decode_workers, &info).unwrap(); + shared_prefill[p] += 1; + shared_decode[d] += 1; + } + assert_eq!(shared_prefill, [20, 0, 20, 0]); + assert_eq!(shared_decode, [0, 20, 0, 20]); + + let prefill_policy = RoundRobinPolicy::new(); + let decode_policy = RoundRobinPolicy::new(); + let mut indep_prefill = [0usize; 4]; + let mut indep_decode = [0usize; 4]; + for _ in 0..40 { + let p = prefill_policy + .select_worker(&prefill_workers, &info) + .unwrap(); + let d = decode_policy.select_worker(&decode_workers, &info).unwrap(); + indep_prefill[p] += 1; + indep_decode[d] += 1; + } + assert_eq!(indep_prefill, [10, 10, 10, 10]); + assert_eq!(indep_decode, [10, 10, 10, 10]); + } } diff --git a/model_gateway/src/routers/grpc/common/stages/worker_selection.rs b/model_gateway/src/routers/grpc/common/stages/worker_selection.rs index 8e55e603d..72fa6a677 100644 --- a/model_gateway/src/routers/grpc/common/stages/worker_selection.rs +++ b/model_gateway/src/routers/grpc/common/stages/worker_selection.rs @@ -370,8 +370,9 @@ impl WorkerSelectionStage { return None; } - // Select using policies - let policy = self.policy_registry.get_policy_or_default(model_id); + // Independent P/D policies so stateful ones (e.g. round_robin) don't share a counter. + let prefill_policy = self.policy_registry.get_prefill_policy(); + let decode_policy = self.policy_registry.get_decode_policy(); // Get cached hash ring for consistent hashing (O(log n) lookup) let hash_ring = self.worker_registry.get_hash_ring(model_id); @@ -385,29 +386,28 @@ impl WorkerSelectionStage { hash_ring, leg: WorkerLeg::Prefill, }; - let prefill_idx = self - .policy_registry - .select_worker(&policy, &available_prefill, &info)?; + let prefill_idx = + self.policy_registry + .select_worker(&prefill_policy, &available_prefill, &info)?; info.leg = WorkerLeg::Decode; - let decode_idx = self - .policy_registry - .select_worker(&policy, &available_decode, &info)?; + let decode_idx = + self.policy_registry + .select_worker(&decode_policy, &available_decode, &info)?; let model = model_id; - let policy_name = policy.name(); // Record worker selection metrics for both prefill and decode Metrics::record_worker_selection( metrics_labels::WORKER_PREFILL, metrics_labels::CONNECTION_GRPC, model, - policy_name, + prefill_policy.name(), ); Metrics::record_worker_selection( metrics_labels::WORKER_DECODE, metrics_labels::CONNECTION_GRPC, model, - policy_name, + decode_policy.name(), ); Some(( @@ -660,3 +660,119 @@ fn hex_encode(bytes: &[u8]) -> String { } out } + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use openai_protocol::worker::HealthCheckConfig; + + use super::*; + use crate::{ + config::types::PolicyConfig, + policies::PolicyFactory, + worker::{BasicWorkerBuilder, ConnectionMode, ModelCard}, + }; + + fn no_health_check() -> HealthCheckConfig { + HealthCheckConfig { + disable_health_check: true, + ..Default::default() + } + } + + fn register_pd_workers( + registry: &WorkerRegistry, + model_id: &str, + n: usize, + ) -> (Vec, Vec) { + let mut prefill_urls = Vec::with_capacity(n); + let mut decode_urls = Vec::with_capacity(n); + + for i in 0..n { + let url = format!("grpc://127.0.0.1:{}", 8000 + i); + prefill_urls.push(url.clone()); + registry + .register(Arc::new( + BasicWorkerBuilder::new(url) + .model(ModelCard::new(model_id)) + .worker_type(WorkerType::Prefill) + .connection_mode(ConnectionMode::Grpc) + .health_config(no_health_check()) + .build(), + )) + .unwrap(); + } + + for i in 0..n { + let url = format!("grpc://127.0.0.1:{}", 8100 + i); + decode_urls.push(url.clone()); + registry + .register(Arc::new( + BasicWorkerBuilder::new(url) + .model(ModelCard::new(model_id)) + .worker_type(WorkerType::Decode) + .connection_mode(ConnectionMode::Grpc) + .health_config(no_health_check()) + .build(), + )) + .unwrap(); + } + + (prefill_urls, decode_urls) + } + + #[test] + fn select_pd_pair_round_robin_covers_all_workers_with_independent_policies() { + let model_id = "test-model"; + let worker_registry = Arc::new(WorkerRegistry::new()); + let (prefill_urls, decode_urls) = register_pd_workers(&worker_registry, model_id, 4); + + let policy_registry = Arc::new(PolicyRegistry::new(PolicyConfig::RoundRobin)); + // Mirror production PD startup: create two independent RoundRobin instances. + policy_registry + .set_prefill_policy(PolicyFactory::create_from_config(&PolicyConfig::RoundRobin)); + policy_registry + .set_decode_policy(PolicyFactory::create_from_config(&PolicyConfig::RoundRobin)); + + let prefill_policy = policy_registry.get_prefill_policy(); + let decode_policy = policy_registry.get_decode_policy(); + assert_eq!(prefill_policy.name(), "round_robin"); + assert_eq!(decode_policy.name(), "round_robin"); + assert!( + !Arc::ptr_eq(&prefill_policy, &decode_policy), + "prefill/decode must not share one RoundRobinPolicy counter" + ); + + let stage = WorkerSelectionStage::new( + worker_registry, + policy_registry, + WorkerSelectionMode::PrefillDecode, + ); + + let mut prefill_hits: HashMap = HashMap::new(); + let mut decode_hits: HashMap = HashMap::new(); + for _ in 0..40 { + let (prefill, decode, _) = stage + .select_pd_pair(model_id, None, None, None) + .expect("select_pd_pair should return a pair"); + *prefill_hits.entry(prefill.url().to_string()).or_default() += 1; + *decode_hits.entry(decode.url().to_string()).or_default() += 1; + } + + for url in &prefill_urls { + assert_eq!( + prefill_hits.get(url).copied().unwrap_or(0), + 10, + "prefill worker {url} should receive 10 of 40 selections" + ); + } + for url in &decode_urls { + assert_eq!( + decode_hits.get(url).copied().unwrap_or(0), + 10, + "decode worker {url} should receive 10 of 40 selections" + ); + } + } +}