From 44001752012b576f35ff15367e385353a7435b43 Mon Sep 17 00:00:00 2001 From: yangxurui Date: Mon, 3 Aug 2026 16:09:06 +0800 Subject: [PATCH 1/4] fix(router): use independent P/D policies in gRPC PD selection Sharing one round_robin counter across prefill and decode made each role only see every other worker index. Signed-off-by: yangxurui Co-authored-by: longcat-ia-team <312376714+longcat-ia-team@users.noreply.github.com> Co-authored-by: Cursor --- model_gateway/src/policies/round_robin.rs | 49 +++++++++++++++++++ .../grpc/common/stages/worker_selection.rs | 26 +++++----- 2 files changed, 64 insertions(+), 11 deletions(-) diff --git a/model_gateway/src/policies/round_robin.rs b/model_gateway/src/policies/round_robin.rs index 6ca4ba0cd..0eaa6a100 100644 --- a/model_gateway/src/policies/round_robin.rs +++ b/model_gateway/src/policies/round_robin.rs @@ -165,4 +165,53 @@ 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 d0eca45a5..e6891e931 100644 --- a/model_gateway/src/routers/grpc/common/stages/worker_selection.rs +++ b/model_gateway/src/routers/grpc/common/stages/worker_selection.rs @@ -365,8 +365,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); @@ -380,29 +381,32 @@ 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(( From fa63664f7077eebe70b26f6c04a48c3b09c638a1 Mon Sep 17 00:00:00 2001 From: yangxurui Date: Mon, 3 Aug 2026 20:51:53 +0800 Subject: [PATCH 2/4] test(router): cover select_pd_pair round_robin via PolicyRegistry Add a WorkerSelectionStage regression that selects 4P+4D with independent round_robin policies through get_prefill_policy / get_decode_policy, locking the PD wiring (not just RoundRobinPolicy). Signed-off-by: yangxurui Co-authored-by: longcat-ia-team <312376714+longcat-ia-team@users.noreply.github.com> Co-authored-by: Cursor --- .../grpc/common/stages/worker_selection.rs | 117 ++++++++++++++++++ 1 file changed, 117 insertions(+) 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 e6891e931..4b7246af5 100644 --- a/model_gateway/src/routers/grpc/common/stages/worker_selection.rs +++ b/model_gateway/src/routers/grpc/common/stages/worker_selection.rs @@ -658,3 +658,120 @@ 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, 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" + ); + } + } +} From 9e2b38c6b0a2813ab948fbd7dc210f1492b0f338 Mon Sep 17 00:00:00 2001 From: yangxurui Date: Tue, 4 Aug 2026 10:29:11 +0800 Subject: [PATCH 3/4] fix(test): import ConnectionMode in select_pd_pair regression The test module cannot see the parent's private use import, which broke CI clippy/libtest compilation. Signed-off-by: yangxurui Co-authored-by: longcat-ia-team <312376714+longcat-ia-team@users.noreply.github.com> Co-authored-by: Cursor --- .../src/routers/grpc/common/stages/worker_selection.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 9e0e946d4..33bebeea2 100644 --- a/model_gateway/src/routers/grpc/common/stages/worker_selection.rs +++ b/model_gateway/src/routers/grpc/common/stages/worker_selection.rs @@ -675,7 +675,7 @@ mod tests { use crate::{ config::types::PolicyConfig, policies::PolicyFactory, - worker::{BasicWorkerBuilder, ModelCard}, + worker::{BasicWorkerBuilder, ConnectionMode, ModelCard}, }; fn no_health_check() -> HealthCheckConfig { From dda4843257f53be6a0031548d45685482933d8f7 Mon Sep 17 00:00:00 2001 From: yangxurui Date: Tue, 4 Aug 2026 20:07:12 +0800 Subject: [PATCH 4/4] style: rustfmt PD round-robin selection and tests Align formatting with nightly rustfmt so CI fmt --check passes. Signed-off-by: yangxurui Co-authored-by: longcat-ia-team <312376714+longcat-ia-team@users.noreply.github.com> Co-authored-by: Cursor --- model_gateway/src/policies/round_robin.rs | 4 +--- .../grpc/common/stages/worker_selection.rs | 21 +++++++------------ 2 files changed, 9 insertions(+), 16 deletions(-) diff --git a/model_gateway/src/policies/round_robin.rs b/model_gateway/src/policies/round_robin.rs index 0eaa6a100..eacd162fe 100644 --- a/model_gateway/src/policies/round_robin.rs +++ b/model_gateway/src/policies/round_robin.rs @@ -205,9 +205,7 @@ mod tests { let p = prefill_policy .select_worker(&prefill_workers, &info) .unwrap(); - let d = decode_policy - .select_worker(&decode_workers, &info) - .unwrap(); + let d = decode_policy.select_worker(&decode_workers, &info).unwrap(); indep_prefill[p] += 1; indep_decode[d] += 1; } 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 33bebeea2..72fa6a677 100644 --- a/model_gateway/src/routers/grpc/common/stages/worker_selection.rs +++ b/model_gateway/src/routers/grpc/common/stages/worker_selection.rs @@ -386,17 +386,13 @@ impl WorkerSelectionStage { hash_ring, leg: WorkerLeg::Prefill, }; - let prefill_idx = self.policy_registry.select_worker( - &prefill_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( - &decode_policy, - &available_decode, - &info, - )?; + let decode_idx = + self.policy_registry + .select_worker(&decode_policy, &available_decode, &info)?; let model = model_id; @@ -734,9 +730,8 @@ mod tests { 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_prefill_policy(PolicyFactory::create_from_config(&PolicyConfig::RoundRobin)); policy_registry .set_decode_policy(PolicyFactory::create_from_config(&PolicyConfig::RoundRobin));