From d5015ba0710cf2881e7e20619fac5c8f134b20ac Mon Sep 17 00:00:00 2001 From: Eric Zhang Date: Mon, 4 May 2026 16:56:44 -0700 Subject: [PATCH] feat(gateway): Smooth transition from regular->PD This adds support for smoothly transitioning from regular routing to PD routing and back, as well as between HTTP and gRPC worker engines. Before if there was at least 1 P or D worker, all requests would be routed to PD mode, even if there were other regular workers. After this change, we weighted route requests between PD and regular workers based on the total worker count, and only send to prefill-decode disaggregated workers if min(P, D) >= 1. Motivation here is to gradually transition a production deployment from regular to PD disaggregated mode without overloading a single worker with a storm of traffic as soon as it starts up. Instead, traffic will gradually be routed to the PD workers as they start up. Also cleaned up the `x-prefer-pd` header which is only used when running in IGW mode without a model ID passed in, deprecated usage. Signed-off-by: Eric Zhang --- model_gateway/src/routers/router_manager.rs | 301 ++++++++++++-------- 1 file changed, 188 insertions(+), 113 deletions(-) diff --git a/model_gateway/src/routers/router_manager.rs b/model_gateway/src/routers/router_manager.rs index 8615ede3b0..4218e7e4aa 100644 --- a/model_gateway/src/routers/router_manager.rs +++ b/model_gateway/src/routers/router_manager.rs @@ -6,7 +6,6 @@ use std::{collections::HashSet, sync::Arc}; -use arc_swap::ArcSwap; use async_trait::async_trait; use axum::{ body::Body, @@ -55,7 +54,7 @@ use crate::{ RouterFactory, RouterTrait, }, server::ServerConfig, - worker::{ConnectionMode, ProviderType, RuntimeType, WorkerRegistry, WorkerType}, + worker::{ConnectionMode, ProviderType, RuntimeType, Worker, WorkerRegistry, WorkerType}, }; pub struct RouterManager { @@ -63,7 +62,6 @@ pub struct RouterManager { client: reqwest::Client, gateway_api_key: Option, routers: Arc>>, - routers_snapshot: ArcSwap>>, default_router: Arc>>, skill_service: Option>, enable_igw: bool, @@ -76,10 +74,9 @@ impl RouterManager { client, gateway_api_key: None, routers: Arc::new(DashMap::new()), - routers_snapshot: ArcSwap::from_pointee(Vec::new()), default_router: Arc::new(std::sync::RwLock::new(None)), skill_service: None, - enable_igw: false, // Will be set properly in from_config + enable_igw: false, } } @@ -172,10 +169,6 @@ impl RouterManager { pub fn register_router(&self, id: RouterId, router: Arc) { self.routers.insert(id.clone(), router); - // Update the lock-free snapshot for fast per-request iteration - let new_snapshot: Vec<_> = self.routers.iter().map(|e| e.value().clone()).collect(); - self.routers_snapshot.store(Arc::new(new_snapshot)); - let mut default_router = self .default_router .write() @@ -198,54 +191,93 @@ impl RouterManager { self.routers.len() } - pub fn get_router_for_model(&self, model_id: &str) -> Option> { - let workers = self.worker_registry.get_by_model(model_id); + /// Selects a router by weighting all four router types (grpc-pd, http-pd, grpc-regular, + /// http-regular) by their worker counts. PD routers only receive weight when both prefill + /// and decode workers are present on the same protocol; an incomplete pair contributes 0. + /// + /// Weighting the router selection lets operators gradually migrate traffic between + /// HTTP / gRPC and regular / prefill-decode disaggregation workers. + fn pick_router_by_weights( + &self, + grpc_pd: usize, + http_pd: usize, + grpc_regular: usize, + http_regular: usize, + ) -> Option> { + let options: [(usize, &RouterId); 4] = [ + (grpc_pd, &router_ids::GRPC_PD), + (http_pd, &router_ids::HTTP_PD), + (grpc_regular, &router_ids::GRPC_REGULAR), + (http_regular, &router_ids::HTTP_REGULAR), + ]; + + let total: usize = options.iter().map(|(w, _)| w).sum(); + if total == 0 { + return None; + } - // Find the best router ID based on worker capabilities - // Priority: external (provider-specific) > grpc-pd > http-pd > grpc-regular > http-regular - let best_router_id = workers - .iter() - .map(|w| { - let is_pd = matches!(w.worker_type(), WorkerType::Prefill | WorkerType::Decode); - let is_grpc = matches!(w.connection_mode(), ConnectionMode::Grpc); - let is_external = matches!(w.metadata().spec.runtime_type, RuntimeType::External); - - if is_external { - // Route external workers to the correct provider-specific router - let router_id = match w.provider_for_model(model_id) { + let pick = ((rand::random::() * total as f64) as usize).min(total - 1); + let mut cum = 0usize; + for (weight, router_id) in &options { + cum += weight; + if pick < cum { + return self.routers.get(*router_id).map(|r| r.clone()); + } + } + None + } + + fn select_router_for_workers( + &self, + workers: &[Arc], + model_id: Option<&str>, + ) -> Option> { + // External workers take highest priority when a model is known. + if let Some(model) = model_id { + for w in workers { + if matches!(w.metadata().spec.runtime_type, RuntimeType::External) { + let router_id = match w.provider_for_model(model) { Some(ProviderType::Gemini) => &router_ids::HTTP_GEMINI, Some(ProviderType::Anthropic) => &router_ids::HTTP_ANTHROPIC, _ => &router_ids::HTTP_OPENAI, }; - return (4, router_id); - } - - match (is_grpc, is_pd) { - (true, true) => (3, &router_ids::GRPC_PD), - (false, true) => (2, &router_ids::HTTP_PD), - (true, false) => (1, &router_ids::GRPC_REGULAR), - (false, false) => (0, &router_ids::HTTP_REGULAR), + return self.routers.get(router_id).map(|r| r.clone()); } - }) - .max_by_key(|(score, _)| *score) - .map(|(_, id)| id); + } + } - if let Some(router_id) = best_router_id { - if let Some(router) = self.routers.get(router_id) { - return Some(router.clone()); + let mut grpc_prefill = 0; + let mut http_prefill = 0; + let mut grpc_decode = 0; + let mut http_decode = 0; + let mut grpc_regular = 0; + let mut http_regular = 0; + + for w in workers { + match (w.worker_type(), w.connection_mode()) { + (WorkerType::Prefill, ConnectionMode::Grpc) => grpc_prefill += 1, + (WorkerType::Prefill, ConnectionMode::Http) => http_prefill += 1, + (WorkerType::Decode, ConnectionMode::Grpc) => grpc_decode += 1, + (WorkerType::Decode, ConnectionMode::Http) => http_decode += 1, + (WorkerType::Regular, ConnectionMode::Grpc) => grpc_regular += 1, + (WorkerType::Regular, ConnectionMode::Http) => http_regular += 1, } } - // Fallback to default router - let default_router = self - .default_router - .read() - .unwrap_or_else(|e| e.into_inner()); - if let Some(ref default_id) = *default_router { - self.routers.get(default_id).map(|r| r.clone()) + // We need at least one prefill and one decode worker to handle requests + // in PD disaggregation mode. + let grpc_pd = if grpc_prefill > 0 && grpc_decode > 0 { + grpc_prefill + grpc_decode } else { - None - } + 0 + }; + let http_pd = if http_prefill > 0 && http_decode > 0 { + http_prefill + http_decode + } else { + 0 + }; + + self.pick_router_by_weights(grpc_pd, http_pd, grpc_regular, http_regular) } fn requires_explicit_generate_model(&self, model_id: &str) -> bool { @@ -254,7 +286,6 @@ impl RouterManager { pub fn select_router_for_request( &self, - headers: Option<&HeaderMap>, model_id: Option<&str>, ) -> Option> { // In single-router mode (enable_igw=false), always use the default router @@ -273,55 +304,22 @@ impl RouterManager { } } - let prefer_pd = headers - .and_then(|h| { - h.get("x-prefer-pd") - .and_then(|v| v.to_str().ok()) - .map(|s| s == "true" || s == "1") - }) - .unwrap_or(false); - - let (num_regular_workers, num_pd_workers) = self.worker_registry.get_worker_distribution(); - let mut best_router = None; - let mut best_score = -1.0; - - // Extract router validity check into a closure to reduce redundancy - let is_router_valid = - |is_pd: bool| (is_pd && num_pd_workers > 0) || (!is_pd && num_regular_workers > 0); - - if let Some(model) = model_id { - // Efficient Single Lookup for Specific Model - if let Some(router) = self.get_router_for_model(model) { - if is_router_valid(router.is_pd_mode()) { - return Some(router); - } - } + let workers = if let Some(model) = model_id { + self.worker_registry.get_by_model(model).to_vec() } else { - // ZERO-ALLOCATION Snapshot Iteration (Hot Path Optimization) - // Atomic load avoids heap allocations and DashMap shard locks per-request - let routers_snapshot = self.routers_snapshot.load(); - for router in routers_snapshot.iter() { - let mut score = 1.0; - - let is_pd = router.is_pd_mode(); - if prefer_pd && is_pd { - score += 2.0; - } else if !prefer_pd && !is_pd { - score += 1.0; - } - // TODO: Once routers expose worker stats, we can evaluate: - // - Average worker priority vs priority_threshold - // - Average worker cost vs max_cost - // - Current load and health status - - if score > best_score && is_router_valid(is_pd) { - best_score = score; - best_router = Some(Arc::clone(router)); - } - } - } + self.worker_registry.get_all() + }; - best_router + self.select_router_for_workers(&workers, model_id) + .or_else(|| { + let default = self + .default_router + .read() + .unwrap_or_else(|e| e.into_inner()); + default + .as_ref() + .and_then(|id| self.routers.get(id).map(|r| r.clone())) + }) } /// Build a response from self-hosted registry models (excludes external workers). @@ -436,7 +434,7 @@ impl RouterTrait for RouterManager { } async fn health_generate(&self, _req: Request) -> Response { - let router = self.select_router_for_request(None, None); + let router = self.select_router_for_request(None); if let Some(router) = router { router.health_generate(_req).await } else { @@ -449,7 +447,7 @@ impl RouterTrait for RouterManager { } async fn get_server_info(&self, req: Request) -> Response { - let router = self.select_router_for_request(None, None); + let router = self.select_router_for_request(None); if let Some(router) = router { router.get_server_info(req).await } else { @@ -536,7 +534,7 @@ impl RouterTrait for RouterManager { ); } - let router = self.select_router_for_request(headers, Some(model_id)); + let router = self.select_router_for_request(Some(model_id)); if let Some(router) = router { router @@ -558,7 +556,7 @@ impl RouterTrait for RouterManager { body: &ChatCompletionRequest, model_id: &str, ) -> Response { - let router = self.select_router_for_request(headers, Some(model_id)); + let router = self.select_router_for_request(Some(model_id)); if let Some(router) = router { router @@ -580,7 +578,7 @@ impl RouterTrait for RouterManager { body: &CompletionRequest, model_id: &str, ) -> Response { - let router = self.select_router_for_request(headers, Some(model_id)); + let router = self.select_router_for_request(Some(model_id)); if let Some(router) = router { router @@ -606,7 +604,7 @@ impl RouterTrait for RouterManager { return route_error::reserved_skill_tool_name(error); } - let router = self.select_router_for_request(headers, Some(model_id)); + let router = self.select_router_for_request(Some(model_id)); if let Some(router) = router { let skill_manifest = match resolve_messages_skill_manifest( self.skill_service.as_deref(), @@ -646,7 +644,7 @@ impl RouterTrait for RouterManager { return route_error::reserved_skill_tool_name(error); } - let router = self.select_router_for_request(headers, Some(model_id)); + let router = self.select_router_for_request(Some(model_id)); if let Some(router) = router { let skill_manifest = match resolve_responses_skill_manifest( self.skill_service.as_deref(), @@ -683,7 +681,7 @@ impl RouterTrait for RouterManager { model_id: Option<&str>, ) -> Response { let selected_model = model_id.or(body.model.as_deref()).or(body.agent.as_deref()); - let router = self.select_router_for_request(headers, selected_model); + let router = self.select_router_for_request(selected_model); if let Some(router) = router { router @@ -699,7 +697,7 @@ impl RouterTrait for RouterManager { } async fn cancel_response(&self, headers: Option<&HeaderMap>, response_id: &str) -> Response { - let router = self.select_router_for_request(headers, None); + let router = self.select_router_for_request(None); if let Some(router) = router { router.cancel_response(headers, response_id).await } else { @@ -718,7 +716,7 @@ impl RouterTrait for RouterManager { body: &EmbeddingRequest, model_id: &str, ) -> Response { - let router = self.select_router_for_request(headers, Some(model_id)); + let router = self.select_router_for_request(Some(model_id)); if let Some(router) = router { router @@ -740,7 +738,7 @@ impl RouterTrait for RouterManager { body: &ClassifyRequest, model_id: &str, ) -> Response { - let router = self.select_router_for_request(headers, Some(model_id)); + let router = self.select_router_for_request(Some(model_id)); if let Some(router) = router { router @@ -763,7 +761,7 @@ impl RouterTrait for RouterManager { audio: AudioFile, model_id: &str, ) -> Response { - let router = self.select_router_for_request(headers, Some(model_id)); + let router = self.select_router_for_request(Some(model_id)); if let Some(router) = router { router @@ -785,7 +783,7 @@ impl RouterTrait for RouterManager { body: &RerankRequest, model_id: &str, ) -> Response { - let router = self.select_router_for_request(headers, Some(model_id)); + let router = self.select_router_for_request(Some(model_id)); if let Some(router) = router { router @@ -806,7 +804,7 @@ impl RouterTrait for RouterManager { body: &RealtimeSessionCreateRequest, ) -> Response { let model = body.model.as_deref(); - let router = self.select_router_for_request(headers, model); + let router = self.select_router_for_request(model); if let Some(router) = router { router.route_realtime_session(headers, body).await } else { @@ -824,7 +822,7 @@ impl RouterTrait for RouterManager { body: &RealtimeClientSecretCreateRequest, ) -> Response { let model = body.session.model.as_deref(); - let router = self.select_router_for_request(headers, model); + let router = self.select_router_for_request(model); if let Some(router) = router { router.route_realtime_client_secret(headers, body).await } else { @@ -842,7 +840,7 @@ impl RouterTrait for RouterManager { body: &RealtimeTranscriptionSessionCreateRequest, ) -> Response { let model = body.model.as_deref(); - let router = self.select_router_for_request(headers, model); + let router = self.select_router_for_request(model); if let Some(router) = router { router .route_realtime_transcription_session(headers, body) @@ -857,7 +855,7 @@ impl RouterTrait for RouterManager { } async fn route_realtime_ws(&self, req: Request, model: &str) -> Response { - let router = self.select_router_for_request(None, Some(model)); + let router = self.select_router_for_request(Some(model)); if let Some(router) = router { router.route_realtime_ws(req, model).await } else { @@ -870,7 +868,7 @@ impl RouterTrait for RouterManager { } async fn route_realtime_webrtc(&self, req: Request, model: &str) -> Response { - let router = self.select_router_for_request(None, Some(model)); + let router = self.select_router_for_request(Some(model)); if let Some(router) = router { router.route_realtime_webrtc(req, model).await } else { @@ -903,13 +901,15 @@ impl std::fmt::Debug for RouterManager { #[cfg(test)] mod tests { + use std::collections::HashMap; + use async_trait::async_trait; use super::*; use crate::{ middleware::{RouteRequestMeta, TenantKey}, routers::factory::router_ids, - worker::WorkerRegistry, + worker::{BasicWorkerBuilder, CircuitBreakerConfig, WorkerRegistry}, }; #[derive(Debug)] @@ -936,6 +936,30 @@ mod tests { } } + #[derive(Debug)] + struct PdStubRouter; + + #[async_trait] + impl RouterTrait for PdStubRouter { + fn as_any(&self) -> &dyn std::any::Any { + self + } + + async fn route_generate( + &self, + _headers: Option<&HeaderMap>, + _tenant_meta: &TenantRequestMeta, + _body: &GenerateRequest, + _model_id: &str, + ) -> Response { + (StatusCode::OK, "pd-routed").into_response() + } + + fn router_type(&self) -> &'static str { + "pd" + } + } + fn test_manager(enable_igw: bool) -> Arc { let mut manager = RouterManager::new(Arc::new(WorkerRegistry::new()), reqwest::Client::new()); @@ -984,4 +1008,55 @@ mod tests { assert_eq!(response.status(), StatusCode::OK); } + + #[test] + fn weighted_routing_splits_40_pd_60_regular() { + let registry = Arc::new(WorkerRegistry::new()); + + let mut url_idx = 0; + let mut add_workers = |wtype: WorkerType, count: usize| { + for _ in 0..count { + let mut labels = HashMap::new(); + labels.insert("model_id".to_string(), "model-x".to_string()); + let worker = BasicWorkerBuilder::new(format!("http://w{url_idx}:8080")) + .worker_type(wtype) + .connection_mode(ConnectionMode::Http) + .labels(labels) + .circuit_breaker_config(CircuitBreakerConfig::default()) + .build(); + registry.register(Arc::new(worker)).unwrap(); + url_idx += 1; + } + }; + + // Try adding 2 prefill, 2 decode workers, and 6 regular workers. + // We should send 40% of traffic to PD and 60% to regular. + add_workers(WorkerType::Prefill, 2); + add_workers(WorkerType::Decode, 2); + add_workers(WorkerType::Regular, 6); + + let mut manager = RouterManager::new(registry, reqwest::Client::new()); + manager.enable_igw = true; + let manager = Arc::new(manager); + manager.register_router(router_ids::HTTP_PD, Arc::new(PdStubRouter)); + manager.register_router(router_ids::HTTP_REGULAR, Arc::new(StubRouter)); + + let n = 10_000; + let pd_count = (0..n) + .filter(|_| { + manager + .select_router_for_request(Some("model-x")) + .map(|r| r.router_type() == "pd") + .unwrap_or(false) + }) + .count(); + + let pd_ratio = pd_count as f64 / n as f64; + let expected = 0.4; // 4 PD workers / 10 total + let tolerance = 0.05; + assert!( + (pd_ratio - expected).abs() < tolerance, + "PD ratio {pd_ratio:.3} was outside expected {expected} ± {tolerance}", + ); + } }