From f65fa04748af4ac678f683c3872895810474eeff Mon Sep 17 00:00:00 2001 From: Arthur Cheng Date: Wed, 24 Dec 2025 00:27:58 -0800 Subject: [PATCH] [model-gateway]Enable IGW mode with gRPC router and auto enable IGW when service discovery is turned on (#15459) --- sgl-model-gateway/src/app_context.rs | 26 +++-- sgl-model-gateway/src/main.rs | 22 ++--- .../src/routers/router_manager.rs | 95 +++++++++++++++---- 3 files changed, 99 insertions(+), 44 deletions(-) diff --git a/sgl-model-gateway/src/app_context.rs b/sgl-model-gateway/src/app_context.rs index 79086efd9..02c56f8ab 100644 --- a/sgl-model-gateway/src/app_context.rs +++ b/sgl-model-gateway/src/app_context.rs @@ -8,9 +8,7 @@ use tracing::{debug, info}; use crate::{ config::RouterConfig, - core::{ - ConnectionMode, JobQueue, LoadMonitor, WorkerRegistry, WorkerService, UNKNOWN_MODEL_ID, - }, + core::{JobQueue, LoadMonitor, WorkerRegistry, WorkerService, UNKNOWN_MODEL_ID}, data_connector::{ create_storage, ConversationItemStorage, ConversationStorage, ResponseStorage, }, @@ -290,8 +288,8 @@ impl AppContextBuilder { .with_client(&router_config, request_timeout_secs)? .maybe_rate_limiter(&router_config) .with_tokenizer_registry(&router_config)? - .maybe_reasoning_parser_factory(&router_config) - .maybe_tool_parser_factory(&router_config) + .with_reasoning_parser_factory() + .with_tool_parser_factory() .with_worker_registry() .with_policy_registry(&router_config) .with_storage(&router_config)? @@ -435,19 +433,17 @@ impl AppContextBuilder { Ok(Some(tokenizer)) } - /// Create reasoning parser factory for gRPC mode - fn maybe_reasoning_parser_factory(mut self, config: &RouterConfig) -> Self { - if matches!(config.connection_mode, ConnectionMode::Grpc { .. }) { - self.reasoning_parser_factory = Some(ReasoningParserFactory::new()); - } + /// Create reasoning parser factory for gRPC mode or IGW mode + fn with_reasoning_parser_factory(mut self) -> Self { + // Initialize reasoning parser factory + self.reasoning_parser_factory = Some(ReasoningParserFactory::new()); self } - /// Create tool parser factory for gRPC mode - fn maybe_tool_parser_factory(mut self, config: &RouterConfig) -> Self { - if matches!(config.connection_mode, ConnectionMode::Grpc { .. }) { - self.tool_parser_factory = Some(ToolParserFactory::new()); - } + /// Create tool parser factory for gRPC mode or IGW mode + fn with_tool_parser_factory(mut self) -> Self { + // Initialize tool parser factory + self.tool_parser_factory = Some(ToolParserFactory::new()); self } diff --git a/sgl-model-gateway/src/main.rs b/sgl-model-gateway/src/main.rs index 9ad931f9b..468d870a9 100644 --- a/sgl-model-gateway/src/main.rs +++ b/sgl-model-gateway/src/main.rs @@ -516,26 +516,20 @@ impl CliArgs { &self, prefill_urls: Vec<(String, Option)>, ) -> ConfigResult { - let mode = if self.enable_igw { - RoutingMode::Regular { - worker_urls: vec![], - } - } else if matches!(self.backend, Backend::Openai) { + // Determine routing mode based on backend type and PD disaggregation flag + // IGW mode doesn't change routing mode, only affects router initialization + let mode = if matches!(self.backend, Backend::Openai) { RoutingMode::OpenAI { worker_urls: self.worker_urls.clone(), } } else if self.pd_disaggregation { - let decode_urls = self.decode.clone(); - - // Allow empty URLs to support dynamic worker addition RoutingMode::PrefillDecode { prefill_urls, - decode_urls, + decode_urls: self.decode.clone(), prefill_policy: self.prefill_policy.as_ref().map(|p| self.parse_policy(p)), decode_policy: self.decode_policy.as_ref().map(|p| self.parse_policy(p)), } } else { - // Allow empty URLs to support dynamic worker addition RoutingMode::Regular { worker_urls: self.worker_urls.clone(), } @@ -762,11 +756,17 @@ fn main() -> Result<(), Box> { let cli = Cli::parse_from(filtered_args); // Handle subcommands or use direct args - let cli_args = match cli.command { + let mut cli_args = match cli.command { Some(Commands::Launch { args }) => args, None => cli.router_args, }; + // Automatically enable IGW mode when service discovery is turned on + if cli_args.service_discovery && !cli_args.enable_igw { + println!("INFO: IGW mode automatically enabled because service discovery is turned on"); + cli_args.enable_igw = true; + } + println!("SGLang Router starting..."); println!("Host: {}:{}", cli_args.host, cli_args.port); let mode_str = if cli_args.enable_igw { diff --git a/sgl-model-gateway/src/routers/router_manager.rs b/sgl-model-gateway/src/routers/router_manager.rs index a088d8bde..7b0588935 100644 --- a/sgl-model-gateway/src/routers/router_manager.rs +++ b/sgl-model-gateway/src/routers/router_manager.rs @@ -93,6 +93,23 @@ impl RouterManager { } } + // Always create gRPC Regular router in IGW mode + match RouterFactory::create_grpc_router(app_context).await { + Ok(grpc_regular) => { + info!("Created gRPC Regular router"); + manager.register_router( + RouterId::new("grpc-regular".to_string()), + Arc::from(grpc_regular), + ); + } + Err(e) => { + warn!("Failed to create gRPC Regular router: {e}"); + } + } + + info!("PD disaggregation auto-enabled for IGW mode, creating PD routers"); + + // Create HTTP PD router match RouterFactory::create_pd_router( None, None, @@ -111,11 +128,28 @@ impl RouterManager { } } - // TODO: Add gRPC routers once we have dynamic tokenizer loading + // Create gRPC PD router + match RouterFactory::create_grpc_pd_router( + None, + None, + &config.router_config.policy, + app_context, + ) + .await + { + Ok(grpc_pd) => { + info!("Created gRPC PD router"); + manager + .register_router(RouterId::new("grpc-pd".to_string()), Arc::from(grpc_pd)); + } + Err(e) => { + warn!("Failed to create gRPC PD router: {e}"); + } + } info!( "RouterManager initialized with {} routers for multi-router mode", - manager.router_count() + manager.router_count(), ); } else { info!("Initializing RouterManager in single-router mode"); @@ -236,25 +270,40 @@ impl RouterManager { pub fn get_router_for_model(&self, model_id: &str) -> Option> { let workers = self.worker_registry.get_by_model(model_id); - if !workers.is_empty() { - let has_pd_workers = workers.iter().any(|w| { - matches!( + // Find the best worker type and derive router ID from it + // Priority: grpc-pd (3) > http-pd (2) > grpc-regular (1) > http-regular (0) + let best_score = 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 router_id = if has_pd_workers { - RouterId::new("http-pd".to_string()) - } else { - RouterId::new("http-regular".to_string()) + match (is_grpc, is_pd) { + (true, true) => 3, // grpc-pd (best) + (false, true) => 2, // http-pd + (true, false) => 1, // grpc-regular + (false, false) => 0, // http-regular + } + }) + .max(); + + if let Some(score) = best_score { + let router_id = match score { + 3 => "grpc-pd", + 2 => "http-pd", + 1 => "grpc-regular", + _ => "http-regular", }; - if let Some(router) = self.routers.get(&router_id) { + if let Some(router) = self.routers.get(&RouterId::new(router_id.to_string())) { return Some(router.clone()); } } + // Fallback to default router let default_router = self.default_router.read().unwrap(); if let Some(ref default_id) = *default_router { self.routers.get(default_id).map(|r| r.clone()) @@ -340,12 +389,22 @@ impl RouterTrait for RouterManager { } async fn health_generate(&self, _req: Request) -> Response { - // TODO: Should check if any router has healthy workers - ( - StatusCode::SERVICE_UNAVAILABLE, - "No routers with healthy workers available", - ) - .into_response() + // IGW readiness: return 200 if at least one router has healthy workers + let has_healthy_workers = self + .worker_registry + .get_all() + .iter() + .any(|w| w.is_healthy()); + + if has_healthy_workers { + (StatusCode::OK, "At least one router has healthy workers").into_response() + } else { + ( + StatusCode::SERVICE_UNAVAILABLE, + "No routers with healthy workers available", + ) + .into_response() + } } async fn get_server_info(&self, _req: Request) -> Response {