From 704ec6f6fb02c976a5f2004f18821162993060a4 Mon Sep 17 00:00:00 2001 From: Jun Liu Date: Sat, 25 Jul 2026 22:20:07 +0900 Subject: [PATCH] fix(router): support model aliases end to end MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A worker registers under a canonical model ID and may declare aliases next to it. Routing only ever matched the canonical ID, so a request that named an alias found no workers and got a 404. The registry now keeps an alias index next to the model index. It is separate so aliases stay out of `/v1/models` and out of the statistics, and `get_by_model` falls back to it. Every other per-model map — hash rings, retry overrides, load balancing policies — stays keyed by the canonical ID, so each request entry point resolves the alias exactly once and passes the canonical ID from there on: - the gRPC pipeline resolves in `RequestContext::new` and rewrites both `input.model_id` and the request's own `model` field, so worker selection, tokenizer lookup, parser selection and tool call ID format all read the canonical ID without resolving anything themselves; - both HTTP routers resolve at the top of their request path and rewrite the `model` field of the body they forward. The backend was registered under the canonical ID and has never heard of the alias. One visible consequence: the response reports the canonical model ID, not the alias the client sent, matching how the OpenAI API answers with the model it actually ran. `/v1/responses` gates on the client-supplied name before the pipeline canonicalizes it, so it gets `WorkerRegistry::contains_model`, which accepts both spellings. The `unknown` wildcard is not a registered name and stays rejected. Alias conflicts are resolved so that a name is either a canonical model ID or an alias, never both: - an alias naming a registered model is refused, and an alias already recorded is dropped when a model later claims that name — otherwise the name would start resolving to a different model once the real one's last worker left; - when two models declare the same alias the first registration wins, and losing that worker hands the alias to a remaining model that declares it rather than stranding it; - the alias entry is held across that whole decision, so a concurrent registration cannot be erased between the check and the delete. Per-model retry overrides are keyed by the canonical ID and are read before the pipeline canonicalizes, so `GrpcRouter::resolve_retry_config` resolves the alias itself. Without it a request naming an alias silently took the router default instead of the worker's override, on chat, generate, messages and completion alike. Tests cover the registry rules, both HTTP request paths end to end against the body the worker actually received, the multipart form, the rerank response, the gRPC pipeline, the Responses canonicalization, the retry override under an alias, and the Responses gate. Signed-off-by: Jun Liu --- .../routers/grpc/common/responses/utils.rs | 55 +- .../grpc/common/stages/dispatch_metadata.rs | 12 +- model_gateway/src/routers/grpc/context.rs | 147 +++--- model_gateway/src/routers/grpc/pipeline.rs | 168 ++++-- model_gateway/src/routers/grpc/router.rs | 149 +++++- model_gateway/src/routers/http/mod.rs | 39 ++ model_gateway/src/routers/http/pd_router.rs | 43 +- model_gateway/src/routers/http/router.rs | 81 ++- model_gateway/src/routers/router_manager.rs | 28 + model_gateway/src/worker/registry.rs | 490 +++++++++++++++++- model_gateway/tests/common/mock_worker.rs | 138 ++++- model_gateway/tests/routing/mod.rs | 1 + .../tests/routing/model_alias_test.rs | 234 +++++++++ .../tests/routing/pd_routing_test.rs | 99 +++- 14 files changed, 1550 insertions(+), 134 deletions(-) create mode 100644 model_gateway/tests/routing/model_alias_test.rs diff --git a/model_gateway/src/routers/grpc/common/responses/utils.rs b/model_gateway/src/routers/grpc/common/responses/utils.rs index 05ce82844..70213a071 100644 --- a/model_gateway/src/routers/grpc/common/responses/utils.rs +++ b/model_gateway/src/routers/grpc/common/responses/utils.rs @@ -98,14 +98,17 @@ pub(crate) async fn ensure_mcp_connection( Ok((false, Vec::new())) } -/// Validate that workers are available for the requested model +/// Validate that workers are available for the requested model. +/// +/// Runs on the client-supplied name, before the pipeline canonicalizes it, so +/// it has to accept aliases as well as canonical model IDs. `contains_model` +/// covers both; listing `get_models()` and testing membership would reject +/// every alias here. pub(crate) fn validate_worker_availability( worker_registry: &Arc, model: &str, ) -> Option { - let available_models = worker_registry.get_models(); - - if !available_models.contains(&model.to_string()) { + if !worker_registry.contains_model(model) { return Some(error::model_not_found(model)); } @@ -175,3 +178,47 @@ pub(crate) async fn persist_response_if_needed( } } } + +#[cfg(test)] +mod tests { + use openai_protocol::{model_card::ModelCard, worker::HealthCheckConfig}; + + use super::*; + use crate::worker::{BasicWorkerBuilder, UNKNOWN_MODEL_ID}; + + fn registry_with_aliased_worker() -> Arc { + let registry = Arc::new(WorkerRegistry::new()); + let worker = BasicWorkerBuilder::new("http://worker:8080") + .model(ModelCard::new("canonical-model").with_alias("model-alias")) + .health_config(HealthCheckConfig { + disable_health_check: true, + ..Default::default() + }) + .build(); + registry.register_or_replace(Arc::new(worker)); + registry + } + + #[test] + fn worker_availability_accepts_alias_and_preserves_unknown_rejection() { + let registry = registry_with_aliased_worker(); + + assert!(validate_worker_availability(®istry, "canonical-model").is_none()); + assert!(validate_worker_availability(®istry, "model-alias").is_none()); + + let response = validate_worker_availability(®istry, UNKNOWN_MODEL_ID) + .expect("unknown model should remain rejected for Responses"); + assert_eq!(response.status(), http::StatusCode::NOT_FOUND); + } + + #[test] + fn worker_availability_rejects_alias_once_its_worker_is_gone() { + let registry = registry_with_aliased_worker(); + let worker_id = registry.get_id_by_url("http://worker:8080").unwrap(); + assert!(registry.remove(&worker_id).is_some()); + + let response = validate_worker_availability(®istry, "model-alias") + .expect("alias must stop resolving with no workers behind it"); + assert_eq!(response.status(), http::StatusCode::NOT_FOUND); + } +} diff --git a/model_gateway/src/routers/grpc/common/stages/dispatch_metadata.rs b/model_gateway/src/routers/grpc/common/stages/dispatch_metadata.rs index f5b27e732..73215c019 100644 --- a/model_gateway/src/routers/grpc/common/stages/dispatch_metadata.rs +++ b/model_gateway/src/routers/grpc/common/stages/dispatch_metadata.rs @@ -27,14 +27,16 @@ impl PipelineStage for DispatchMetadataStage { })?; let request_id = execution_plan.request_id().to_string(); + // The model the response reports. `RequestContext::new` already + // rewrote every one of these to the canonical model ID, so a request + // that arrived under an alias is answered under the canonical name. let model = match &ctx.input.request_type { RequestType::Chat(req) => req.model.clone(), RequestType::Completion(req) => req.model.clone(), - RequestType::Generate(_req) => { - // Generate requests don't have a model field - // Use model_id from input - ctx.input.model_id.clone() - } + // `GenerateRequest` carries a model field too, but callers of the + // native `/generate` route may leave it empty, so prefer the + // model the router resolved. + RequestType::Generate(_req) => ctx.input.model_id.clone(), RequestType::Responses(req) => req.model.clone(), RequestType::Embedding(req) => req.model.clone(), RequestType::Classify(req) => req.model.clone(), diff --git a/model_gateway/src/routers/grpc/context.rs b/model_gateway/src/routers/grpc/context.rs index 8ff18fd72..d2ba6862a 100644 --- a/model_gateway/src/routers/grpc/context.rs +++ b/model_gateway/src/routers/grpc/context.rs @@ -32,7 +32,7 @@ use super::{ }; use crate::{ middleware::TenantRequestMeta, - worker::{RuntimeType, Worker, WorkerLoadGuard}, + worker::{RuntimeType, Worker, WorkerLoadGuard, WorkerRegistry}, }; /// Main request processing context @@ -50,6 +50,7 @@ pub(crate) struct RequestContext { pub(crate) struct RequestInput { pub request_type: RequestType, pub headers: Option, + /// Canonical model ID used after aliases are resolved at request entry. pub model_id: String, pub tenant_request_meta: Option, } @@ -67,6 +68,29 @@ pub(crate) enum RequestType { } impl RequestType { + /// Overwrite the request's own `model` field. + /// + /// Callers hold the request behind an `Arc` that the retry loop also + /// holds, so `Arc::make_mut` copies the request here. That cost is paid + /// only on the alias path — [`RequestContext::new`] skips this call + /// entirely when the client already used the canonical model ID. + fn set_model(&mut self, model_id: &str) { + fn replace(model: &mut String, model_id: &str) { + model.clear(); + model.push_str(model_id); + } + + match self { + Self::Chat(request) => replace(&mut Arc::make_mut(request).model, model_id), + Self::Generate(request) => replace(&mut Arc::make_mut(request).model, model_id), + Self::Completion(request) => replace(&mut Arc::make_mut(request).model, model_id), + Self::Responses(request) => replace(&mut Arc::make_mut(request).model, model_id), + Self::Embedding(request) => replace(&mut Arc::make_mut(request).model, model_id), + Self::Classify(request) => replace(&mut Arc::make_mut(request).model, model_id), + Self::Messages(request) => replace(&mut Arc::make_mut(request).model, model_id), + } + } + /// Client-supplied backend request id (`rid`), where the protocol carries /// one. Responses ids are storage-owned (`resp_*`) and never client-set. pub fn rid(&self) -> Option<&str> { @@ -112,6 +136,7 @@ impl std::fmt::Display for FinalResponse { /// Shared components (injected once at creation) pub(crate) struct SharedComponents { pub tokenizer_registry: Arc, + pub worker_registry: Arc, pub tool_parser_factory: ToolParserFactory, pub reasoning_parser_factory: ReasoningParserFactory, /// Configured tool parser name (from CLI `--tool-call-parser`) @@ -448,16 +473,32 @@ pub(crate) struct ResponseState { } impl RequestContext { - /// Create context for chat completion request - pub fn for_chat( - request: Arc, + /// Build a context, resolving a model alias to its canonical model ID. + /// + /// This is the single place the gRPC pipeline canonicalizes. Both + /// `input.model_id` and the request's own `model` field are rewritten, so + /// every stage below — worker selection, tokenizer lookup, parser + /// selection, tool call ID format — reads the canonical ID without + /// resolving anything itself. + /// + /// One visible consequence: the response reports the canonical model ID, + /// not the alias the client sent. That matches how the OpenAI API answers + /// with the model it actually ran. + fn new( + mut request_type: RequestType, headers: Option, - model_id: String, + mut model_id: String, components: Arc, ) -> Self { + if let Some(canonical_model_id) = components.worker_registry.resolve_model_alias(&model_id) + { + model_id.clear(); + model_id.push_str(&canonical_model_id); + request_type.set_model(&model_id); + } Self { input: RequestInput { - request_type: RequestType::Chat(request), + request_type, headers, model_id, tenant_request_meta: None, @@ -467,6 +508,16 @@ impl RequestContext { } } + /// Create context for chat completion request + pub fn for_chat( + request: Arc, + headers: Option, + model_id: String, + components: Arc, + ) -> Self { + Self::new(RequestType::Chat(request), headers, model_id, components) + } + /// Create context for generate request pub fn for_generate( request: Arc, @@ -474,16 +525,12 @@ impl RequestContext { model_id: String, components: Arc, ) -> Self { - Self { - input: RequestInput { - request_type: RequestType::Generate(request), - headers, - model_id, - tenant_request_meta: None, - }, + Self::new( + RequestType::Generate(request), + headers, + model_id, components, - state: ProcessingState::default(), - } + ) } /// Create context for completion request @@ -493,16 +540,12 @@ impl RequestContext { model_id: String, components: Arc, ) -> Self { - Self { - input: RequestInput { - request_type: RequestType::Completion(request), - headers, - model_id, - tenant_request_meta: None, - }, + Self::new( + RequestType::Completion(request), + headers, + model_id, components, - state: ProcessingState::default(), - } + ) } /// Create context for Responses API request @@ -512,16 +555,12 @@ impl RequestContext { model_id: String, components: Arc, ) -> Self { - Self { - input: RequestInput { - request_type: RequestType::Responses(request), - headers, - model_id, - tenant_request_meta: None, - }, + Self::new( + RequestType::Responses(request), + headers, + model_id, components, - state: ProcessingState::default(), - } + ) } /// Create context for embedding request @@ -531,16 +570,12 @@ impl RequestContext { model_id: String, components: Arc, ) -> Self { - Self { - input: RequestInput { - request_type: RequestType::Embedding(request), - headers, - model_id, - tenant_request_meta: None, - }, + Self::new( + RequestType::Embedding(request), + headers, + model_id, components, - state: ProcessingState::default(), - } + ) } /// Create context for classify request @@ -550,16 +585,12 @@ impl RequestContext { model_id: String, components: Arc, ) -> Self { - Self { - input: RequestInput { - request_type: RequestType::Classify(request), - headers, - model_id, - tenant_request_meta: None, - }, + Self::new( + RequestType::Classify(request), + headers, + model_id, components, - state: ProcessingState::default(), - } + ) } /// Create context for messages request @@ -569,16 +600,12 @@ impl RequestContext { model_id: String, components: Arc, ) -> Self { - Self { - input: RequestInput { - request_type: RequestType::Messages(request), - headers, - model_id, - tenant_request_meta: None, - }, + Self::new( + RequestType::Messages(request), + headers, + model_id, components, - state: ProcessingState::default(), - } + ) } /// Get chat request (panics if not chat) diff --git a/model_gateway/src/routers/grpc/pipeline.rs b/model_gateway/src/routers/grpc/pipeline.rs index 83306318e..e71715904 100644 --- a/model_gateway/src/routers/grpc/pipeline.rs +++ b/model_gateway/src/routers/grpc/pipeline.rs @@ -418,23 +418,21 @@ impl RequestPipeline { tenant_request_meta: Option, ) -> Response { let start = Instant::now(); - // Clone Arc for metrics (cheap atomic increment) to avoid borrow issues - let request_for_metrics = Arc::clone(&request); let streaming = request.stream; + let mut ctx = RequestContext::for_chat(request, headers, model_id, components); + ctx.input.tenant_request_meta = tenant_request_meta; + let model = ctx.input.model_id.clone(); // Record request start Metrics::record_router_request( metrics_labels::ROUTER_GRPC, self.backend_type, metrics_labels::CONNECTION_GRPC, - &request_for_metrics.model, + &model, metrics_labels::ENDPOINT_CHAT, bool_to_static_str(streaming), ); - let mut ctx = RequestContext::for_chat(request, headers, model_id, components); - ctx.input.tenant_request_meta = tenant_request_meta; - for stage in self.stages.iter() { match stage.execute(&mut ctx).await { Ok(Some(response)) => { @@ -443,7 +441,7 @@ impl RequestPipeline { metrics_labels::ROUTER_GRPC, self.backend_type, metrics_labels::CONNECTION_GRPC, - &request_for_metrics.model, + &model, metrics_labels::ENDPOINT_CHAT, start.elapsed(), ); @@ -455,7 +453,7 @@ impl RequestPipeline { metrics_labels::ROUTER_GRPC, self.backend_type, metrics_labels::CONNECTION_GRPC, - &request_for_metrics.model, + &model, metrics_labels::ENDPOINT_CHAT, error_type_from_status(response.status()), ); @@ -475,7 +473,7 @@ impl RequestPipeline { metrics_labels::ROUTER_GRPC, self.backend_type, metrics_labels::CONNECTION_GRPC, - &request_for_metrics.model, + &model, metrics_labels::ENDPOINT_CHAT, start.elapsed(), ); @@ -491,14 +489,12 @@ impl RequestPipeline { "execute_chat", "Chat", &response_type, - &request_for_metrics.model, - metrics_labels::ENDPOINT_CHAT, - ), - None => self.no_response_produced( - "execute_chat", - &request_for_metrics.model, + &model, metrics_labels::ENDPOINT_CHAT, ), + None => { + self.no_response_produced("execute_chat", &model, metrics_labels::ENDPOINT_CHAT) + } } } @@ -513,6 +509,9 @@ impl RequestPipeline { ) -> Response { let start = Instant::now(); let streaming = request.stream; + let mut ctx = RequestContext::for_generate(request, headers, model_id, components); + ctx.input.tenant_request_meta = tenant_request_meta; + let model_id = ctx.input.model_id.clone(); // Record request start Metrics::record_router_request( @@ -524,9 +523,6 @@ impl RequestPipeline { bool_to_static_str(streaming), ); - let mut ctx = RequestContext::for_generate(request, headers, model_id.clone(), components); - ctx.input.tenant_request_meta = tenant_request_meta; - for stage in self.stages.iter() { match stage.execute(&mut ctx).await { Ok(Some(response)) => { @@ -603,8 +599,10 @@ impl RequestPipeline { tenant_request_meta: Option, ) -> Response { let start = Instant::now(); - let model = request.model.clone(); let streaming = request.stream; + let mut ctx = RequestContext::for_completion(request, headers, model_id, components); + ctx.input.tenant_request_meta = tenant_request_meta; + let model = ctx.input.model_id.clone(); Metrics::record_router_request( metrics_labels::ROUTER_GRPC, @@ -615,9 +613,6 @@ impl RequestPipeline { bool_to_static_str(streaming), ); - let mut ctx = RequestContext::for_completion(request, headers, model_id, components); - ctx.input.tenant_request_meta = tenant_request_meta; - for stage in self.stages.iter() { match stage.execute(&mut ctx).await { Ok(Some(response)) => { @@ -693,6 +688,9 @@ impl RequestPipeline { components: Arc, tenant_request_meta: Option, ) -> Response { + let mut ctx = RequestContext::for_embedding(request, headers, model_id, components); + ctx.input.tenant_request_meta = tenant_request_meta; + let model_id = ctx.input.model_id.clone(); debug!( "execute_embeddings: Starting execution for model: {}", &model_id @@ -709,9 +707,6 @@ impl RequestPipeline { bool_to_static_str(false), ); - let mut ctx = RequestContext::for_embedding(request, headers, model_id.clone(), components); - ctx.input.tenant_request_meta = tenant_request_meta; - for stage in self.stages.iter() { debug!("execute_embeddings: Executing stage: {}", stage.name()); match stage.execute(&mut ctx).await { @@ -795,6 +790,9 @@ impl RequestPipeline { components: Arc, tenant_request_meta: Option, ) -> Response { + let mut ctx = RequestContext::for_classify(request, headers, model_id, components); + ctx.input.tenant_request_meta = tenant_request_meta; + let model_id = ctx.input.model_id.clone(); debug!( "execute_classify: Starting execution for model: {}", &model_id @@ -811,9 +809,6 @@ impl RequestPipeline { bool_to_static_str(false), // Classify is never streaming ); - let mut ctx = RequestContext::for_classify(request, headers, model_id.clone(), components); - ctx.input.tenant_request_meta = tenant_request_meta; - for stage in self.stages.iter() { debug!("execute_classify: Executing stage: {}", stage.name()); match stage.execute(&mut ctx).await { @@ -899,20 +894,20 @@ impl RequestPipeline { ) -> Response { let start = Instant::now(); let streaming = request.stream.unwrap_or(false); + let mut ctx = RequestContext::for_messages(request, headers, model_id, components); + ctx.input.tenant_request_meta = tenant_request_meta; + let model = ctx.input.model_id.clone(); // Record request start Metrics::record_router_request( metrics_labels::ROUTER_GRPC, self.backend_type, metrics_labels::CONNECTION_GRPC, - &request.model, + &model, metrics_labels::ENDPOINT_MESSAGES, bool_to_static_str(streaming), ); - let mut ctx = RequestContext::for_messages(request.clone(), headers, model_id, components); - ctx.input.tenant_request_meta = tenant_request_meta; - for stage in self.stages.iter() { match stage.execute(&mut ctx).await { Ok(Some(response)) => { @@ -921,7 +916,7 @@ impl RequestPipeline { metrics_labels::ROUTER_GRPC, self.backend_type, metrics_labels::CONNECTION_GRPC, - &request.model, + &model, metrics_labels::ENDPOINT_MESSAGES, start.elapsed(), ); @@ -933,7 +928,7 @@ impl RequestPipeline { metrics_labels::ROUTER_GRPC, self.backend_type, metrics_labels::CONNECTION_GRPC, - &request.model, + &model, metrics_labels::ENDPOINT_MESSAGES, error_type_from_status(response.status()), ); @@ -953,7 +948,7 @@ impl RequestPipeline { metrics_labels::ROUTER_GRPC, self.backend_type, metrics_labels::CONNECTION_GRPC, - &request.model, + &model, metrics_labels::ENDPOINT_MESSAGES, start.elapsed(), ); @@ -969,12 +964,12 @@ impl RequestPipeline { "execute_messages", "Messages", &response_type, - &request.model, + &model, metrics_labels::ENDPOINT_MESSAGES, ), None => self.no_response_produced( "execute_messages", - &request.model, + &model, metrics_labels::ENDPOINT_MESSAGES, ), } @@ -1434,3 +1429,102 @@ mod build_parity_tests { } } } + +#[cfg(test)] +mod alias_pipeline_tests { + use llm_tokenizer::{traits::Tokenizer, MockTokenizer, TokenizerRegistry}; + use openai_protocol::{ + generate::GenerateRequest, model_card::ModelCard, worker::HealthCheckConfig, + }; + use serde_json::json; + + use super::*; + use crate::{ + config::types::PolicyConfig, + worker::{BasicWorkerBuilder, ConnectionMode, RuntimeType, WorkerType}, + }; + + const CANONICAL_MODEL: &str = "canonical-model"; + const MODEL_ALIAS: &str = "model-alias"; + + fn register_pd_worker(registry: &WorkerRegistry, url: &str, worker_type: WorkerType) { + let worker = BasicWorkerBuilder::new(url) + .worker_type(worker_type) + .connection_mode(ConnectionMode::Grpc) + .runtime_type(RuntimeType::Sglang) + .model(ModelCard::new(CANONICAL_MODEL).with_alias(MODEL_ALIAS)) + .health_config(HealthCheckConfig { + disable_health_check: true, + ..Default::default() + }) + .build(); + registry.register(Arc::new(worker)).unwrap(); + } + + #[tokio::test] + async fn pd_generate_alias_is_canonical_before_preparation() { + let worker_registry = Arc::new(WorkerRegistry::new()); + register_pd_worker( + &worker_registry, + "grpc://prefill:30000", + WorkerType::Prefill, + ); + register_pd_worker(&worker_registry, "grpc://decode:30000", WorkerType::Decode); + + let tokenizer_registry = Arc::new(TokenizerRegistry::new()); + let tokenizer = Arc::new(MockTokenizer::new()) as Arc; + tokenizer_registry + .load("tokenizer-id", CANONICAL_MODEL, "test", || async move { + Ok(tokenizer) + }) + .await + .unwrap(); + + let policy_registry = Arc::new(PolicyRegistry::new(PolicyConfig::RoundRobin)); + let deps = PipelineDeps::pair(worker_registry.clone(), policy_registry); + let pipeline = RequestPipeline::build(Endpoint::Chat, Mode::PrefillDecode, &deps).unwrap(); + let components = Arc::new(SharedComponents { + tokenizer_registry, + worker_registry, + tool_parser_factory: ToolParserFactory::default(), + reasoning_parser_factory: ReasoningParserFactory::default(), + configured_tool_parser: None, + configured_reasoning_parser: None, + multimodal: None, + }); + let request: GenerateRequest = serde_json::from_value(json!({ + "model": MODEL_ALIAS, + "text": "Hello" + })) + .unwrap(); + let mut ctx = RequestContext::for_generate( + Arc::new(request), + None, + MODEL_ALIAS.to_string(), + components, + ); + + assert_eq!(ctx.input.model_id, CANONICAL_MODEL); + assert_eq!(ctx.generate_request().model, CANONICAL_MODEL); + + for stage in pipeline.stages.iter() { + assert!(stage.execute(&mut ctx).await.unwrap().is_none()); + if ctx.state.workers.is_some() { + break; + } + } + + assert_eq!(ctx.input.model_id, CANONICAL_MODEL); + assert_eq!(ctx.generate_request().model, CANONICAL_MODEL); + assert!(ctx.state.tokenizer.is_some()); + match ctx.state.workers.as_ref().unwrap() { + WorkerSelection::Disaggregated { + prefill, decode, .. + } => { + assert_eq!(prefill.url(), "grpc://prefill:30000"); + assert_eq!(decode.url(), "grpc://decode:30000"); + } + WorkerSelection::Single { .. } => panic!("expected PD worker selection"), + } + } +} diff --git a/model_gateway/src/routers/grpc/router.rs b/model_gateway/src/routers/grpc/router.rs index b9e572685..9b84c4d4b 100644 --- a/model_gateway/src/routers/grpc/router.rs +++ b/model_gateway/src/routers/grpc/router.rs @@ -1,4 +1,4 @@ -use std::sync::Arc; +use std::{borrow::Cow, sync::Arc}; use async_trait::async_trait; use axum::{ @@ -348,6 +348,7 @@ impl GrpcRouter { // Create shared components for pipeline let shared_components = Arc::new(SharedComponents { tokenizer_registry: tokenizer_registry.clone(), + worker_registry: worker_registry.clone(), tool_parser_factory: tool_parser_factory.clone(), reasoning_parser_factory: reasoning_parser_factory.clone(), configured_tool_parser: ctx.configured_tool_parser.clone(), @@ -437,7 +438,14 @@ impl GrpcRouter { /// The per-model retry override registered by a worker, else the router /// default. Applied at every retrying endpoint /// (chat/generate/messages/completion) in every mode. + /// + /// Resolves the alias itself. Retry overrides are keyed by canonical model + /// ID, and this runs before the pipeline canonicalizes in + /// `RequestContext::new`, so an alias would miss the override and fall + /// back to the router default without saying so. fn resolve_retry_config(&self, model_id: &str) -> RetryConfig { + let canonical_model = self.worker_registry.resolve_model_alias(model_id); + let model_id = canonical_model.as_deref().unwrap_or(model_id); self.worker_registry .get_retry_config(model_id) .unwrap_or_else(|| self.retry_config.clone()) @@ -612,6 +620,10 @@ impl GrpcRouter { return error_response; } + let (body, canonical_model_id) = + canonicalize_responses_request(&self.worker_registry, body, model_id); + let model_id = canonical_model_id.as_ref(); + // Choose implementation based on Harmony model detection (checks worker metadata) let is_harmony = HarmonyDetector::is_harmony_model_in_registry(&self.worker_registry, &body.model); @@ -634,10 +646,11 @@ impl GrpcRouter { ); if body.stream.unwrap_or(false) { - serve_harmony_responses_stream(&harmony_ctx, body.clone(), tenant_meta.clone()) + serve_harmony_responses_stream(&harmony_ctx, body.into_owned(), tenant_meta.clone()) .await } else { - match serve_harmony_responses(&harmony_ctx, body.clone(), tenant_meta.clone()).await + match serve_harmony_responses(&harmony_ctx, body.into_owned(), tenant_meta.clone()) + .await { Ok(response) => axum::Json(response).into_response(), Err(error_response) => error_response, @@ -646,7 +659,7 @@ impl GrpcRouter { } else { responses::route_responses( responses_context, - Arc::new(body.clone()), + Arc::new(body.into_owned()), headers.cloned(), tenant_meta.clone(), model_id.to_string(), @@ -1007,6 +1020,34 @@ impl RouterTrait for GrpcRouter { } } +/// Resolve a Responses request's model alias, for both the routing decisions +/// and the request the Responses layer will read. +/// +/// The Responses layer builds its SSE events and its tool call responses from +/// the request rather than from the pipeline context, so the alias has to be +/// gone before dispatch. Otherwise a plain answer reports the canonical ID +/// while a tool call answer reports the alias. +/// +/// Returns both halves together so a caller cannot take the canonical model ID +/// and forget the body, or the reverse. The common path borrows both untouched +/// and allocates nothing. +fn canonicalize_responses_request<'a>( + worker_registry: &WorkerRegistry, + body: &'a ResponsesRequest, + model_id: &'a str, +) -> (Cow<'a, ResponsesRequest>, Cow<'a, str>) { + let Some(canonical_model) = worker_registry.resolve_model_alias(model_id) else { + return (Cow::Borrowed(body), Cow::Borrowed(model_id)); + }; + let mut canonical_body = body.clone(); + canonical_body.model.clear(); + canonical_body.model.push_str(&canonical_model); + ( + Cow::Owned(canonical_body), + Cow::Owned(canonical_model.to_string()), + ) +} + #[cfg(test)] mod tests { use axum::http::StatusCode; @@ -1198,6 +1239,7 @@ mod pd_tests { use std::sync::{Arc, OnceLock}; use llm_tokenizer::registry::TokenizerRegistry; + use openai_protocol::{model_card::ModelCard, worker::HealthCheckConfig}; use reasoning_parser::ParserFactory as ReasoningParserFactory; use smg_data_connector::{ MemoryConversationItemStorage, MemoryConversationStorage, MemoryResponseStorage, @@ -1210,7 +1252,7 @@ mod pd_tests { config::{PolicyConfig, RouterConfig, RoutingMode}, policies::PolicyRegistry, tenant::TenantKey, - worker::WorkerRegistry, + worker::{BasicWorkerBuilder, WorkerRegistry}, }; fn pd_routing_mode() -> RoutingMode { @@ -1344,4 +1386,101 @@ mod pd_tests { let cancel = router.cancel_response(None, "resp_missing").await; assert_eq!(cancel.status(), StatusCode::NOT_IMPLEMENTED); } + + /// A Responses request addressed to an alias must reach the Responses + /// layer under the canonical model ID. That layer builds its SSE events + /// and tool-call responses from this request rather than from the pipeline + /// context, so leaving the alias here is what made the streaming and + /// tool-call paths disagree with the plain-answer path about which model + /// ran. + #[tokio::test] + async fn responses_request_is_canonical_before_the_responses_layer_sees_it() { + let ctx = grpc_ctx(pd_routing_mode()).await; + for (url, worker_type) in [ + ("grpc://prefill:30000", WorkerType::Prefill), + ("grpc://decode:30000", WorkerType::Decode), + ] { + let worker = BasicWorkerBuilder::new(url) + .worker_type(worker_type) + .connection_mode(ConnectionMode::Grpc) + .model(ModelCard::new("canonical-model").with_alias("model-alias")) + .health_config(HealthCheckConfig { + disable_health_check: true, + ..Default::default() + }) + .build(); + ctx.worker_registry + .register(Arc::new(worker)) + .expect("register worker"); + } + + let request = responses_request("model-alias"); + let (body, model_id) = + canonicalize_responses_request(&ctx.worker_registry, &request, "model-alias"); + assert_eq!( + body.model, "canonical-model", + "the Responses layer must see the canonical model ID" + ); + assert_eq!( + model_id, "canonical-model", + "worker selection must use the canonical model ID" + ); + + // The canonical ID is left alone, and costs no copy. + let already_canonical = responses_request("canonical-model"); + let (body, model_id) = canonicalize_responses_request( + &ctx.worker_registry, + &already_canonical, + "canonical-model", + ); + assert!(matches!(body, Cow::Borrowed(_))); + assert!(matches!(model_id, Cow::Borrowed(_))); + } + + /// A per-model retry override must survive an alias. The override is keyed + /// by canonical model ID and is read before the pipeline canonicalizes, so + /// without its own resolution an alias silently falls back to the router + /// default. + #[tokio::test] + async fn retry_override_survives_a_model_alias() { + let ctx = grpc_ctx(pd_routing_mode()).await; + let worker = BasicWorkerBuilder::new("grpc://decode:30000") + .worker_type(WorkerType::Decode) + .connection_mode(ConnectionMode::Grpc) + .model(ModelCard::new("canonical-model").with_alias("model-alias")) + .health_config(HealthCheckConfig { + disable_health_check: true, + ..Default::default() + }) + .build(); + ctx.worker_registry + .register(Arc::new(worker)) + .expect("register worker"); + + let router = GrpcRouter::new(&ctx, Mode::PrefillDecode).expect("pd router"); + let default_retries = router.retry_config.max_retries; + let override_retries = default_retries + 7; + ctx.worker_registry.set_model_retry_config( + "canonical-model", + RetryConfig { + max_retries: override_retries, + ..RetryConfig::default() + }, + true, + ); + + assert_eq!( + router.resolve_retry_config("canonical-model").max_retries, + override_retries + ); + assert_eq!( + router.resolve_retry_config("model-alias").max_retries, + override_retries, + "an aliased request must get the same retry override as the canonical ID" + ); + assert_eq!( + router.resolve_retry_config("unrelated-model").max_retries, + default_retries + ); + } } diff --git a/model_gateway/src/routers/http/mod.rs b/model_gateway/src/routers/http/mod.rs index 3f31b6f86..beeef2386 100644 --- a/model_gateway/src/routers/http/mod.rs +++ b/model_gateway/src/routers/http/mod.rs @@ -3,3 +3,42 @@ pub mod pd_router; pub mod pd_types; pub mod router; + +use serde_json::Value; + +/// Rewrite the `model` field of an outbound request body. +/// +/// A worker is registered under its canonical model ID, so a request that +/// arrived under an alias must reach the backend under the canonical name — +/// the backend has never heard of the alias. Both HTTP routers resolve the +/// alias for their own routing decisions; this puts the same answer into the +/// body they forward. +/// +/// Leaves a body without a `model` field alone. Inserting the key would change +/// the request the client wrote. +pub(crate) fn set_request_model(json: &mut Value, canonical_model_id: &str) { + if let Some(model) = json.get_mut("model") { + *model = Value::String(canonical_model_id.to_owned()); + } +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + + #[test] + fn set_request_model_replaces_an_existing_model_field() { + let mut body = json!({"model": "GLM-5.2-Coding", "stream": false}); + set_request_model(&mut body, "GLM-5.2"); + assert_eq!(body, json!({"model": "GLM-5.2", "stream": false})); + } + + #[test] + fn set_request_model_leaves_a_body_without_the_field_untouched() { + let mut body = json!({"text": "Hello"}); + set_request_model(&mut body, "GLM-5.2"); + assert_eq!(body, json!({"text": "Hello"})); + } +} diff --git a/model_gateway/src/routers/http/pd_router.rs b/model_gateway/src/routers/http/pd_router.rs index 6fc9b44a3..0f4637ce5 100644 --- a/model_gateway/src/routers/http/pd_router.rs +++ b/model_gateway/src/routers/http/pd_router.rs @@ -286,12 +286,17 @@ impl PDRouter { &self, headers: Option<&HeaderMap>, original_request: &T, - context: PDRequestContext<'_>, + mut context: PDRequestContext<'_>, ) -> Response { let start_time = Instant::now(); let route = context.route; - let model = context.model_id; + // Resolve once, here, so every registry, policy and metrics lookup + // below is keyed by the canonical model ID. Only `get_by_model` + // understands aliases; retry configs, hash rings and policies do not. + let canonical_model = self.worker_registry.resolve_model_alias(context.model_id); + let model = canonical_model.as_deref().unwrap_or(context.model_id); + context.model_id = model; let endpoint = route_to_endpoint(route); // Record request start (Layer 2) @@ -346,6 +351,10 @@ impl PDRouter { Ok(v) => v, Err(e) => return Self::handle_serialization_error(e), }; + // The prefill and decode workers only know the + // canonical name, so forward that, not the alias the + // client sent. + super::set_request_model(&mut json_request, model); json_request = match Self::inject_bootstrap_into_value( json_request, @@ -1532,6 +1541,8 @@ impl RouterTrait for PDRouter { #[cfg(test)] mod tests { + use openai_protocol::model_card::ModelCard; + use super::*; use crate::{ config::PolicyConfig, @@ -1696,6 +1707,34 @@ mod tests { assert!(prefill.is_healthy()); } + #[tokio::test] + async fn test_select_pd_pair_accepts_model_alias() { + let router = create_test_pd_router(); + for (url, worker_type) in [ + ("http://prefill", WorkerType::Prefill), + ("http://decode", WorkerType::Decode), + ] { + let worker = BasicWorkerBuilder::new(url) + .worker_type(worker_type) + .model(ModelCard::new("GLM-5.2").with_alias("GLM-5.2-Coding")) + .build(); + worker.set_status(openai_protocol::worker::WorkerStatus::Ready); + router.worker_registry.register(Arc::new(worker)).unwrap(); + } + + let (prefill, decode) = router + .select_pd_pair(None, "GLM-5.2-Coding", None) + .await + .expect("alias should select a PD pair"); + assert_eq!(prefill.url(), "http://prefill"); + assert_eq!(decode.url(), "http://decode"); + + assert!(router + .select_pd_pair(None, "GLM-5.2-Unknown", None) + .await + .is_err()); + } + #[tokio::test] async fn test_empty_worker_lists() { let router = create_test_pd_router(); diff --git a/model_gateway/src/routers/http/router.rs b/model_gateway/src/routers/http/router.rs index 3e588f334..38188a169 100644 --- a/model_gateway/src/routers/http/router.rs +++ b/model_gateway/src/routers/http/router.rs @@ -247,6 +247,12 @@ impl Router { let start = Instant::now(); let is_stream = typed_req.is_stream(); let text = typed_req.extract_text_for_routing(); + // Resolve once, here, so every registry, policy and metrics lookup + // below is keyed by the canonical model ID. Only `get_by_model` + // understands aliases; retry configs, hash rings and policies do not, + // and an alias would silently fall back to router defaults. + let canonical_model = self.worker_registry.resolve_model_alias(model_id); + let model_id = canonical_model.as_deref().unwrap_or(model_id); let model = model_id; let endpoint = route_to_endpoint(route); @@ -271,7 +277,15 @@ impl Router { // operation per attempt |_: u32| async { let res = self - .route_typed_request_once(headers, typed_req, route, model_id, is_stream, &text) + .route_typed_request_once( + headers, + typed_req, + route, + model_id, + canonical_model.as_deref(), + is_stream, + &text, + ) .await; // Need to be outside `route_typed_request_once` because that function has multiple return paths @@ -322,12 +336,17 @@ impl Router { response } + #[expect( + clippy::too_many_arguments, + reason = "per-attempt state threaded from route_typed_request; a struct would only move the arity" + )] async fn route_typed_request_once( &self, headers: Option<&HeaderMap>, typed_req: &T, route: &'static str, model_id: &str, + canonical_model: Option<&str>, is_stream: bool, text: &str, ) -> Response { @@ -375,6 +394,7 @@ impl Router { headers, typed_req, route, + canonical_model, worker.as_ref(), is_stream, load_guard, @@ -524,6 +544,13 @@ impl Router { let start = Instant::now(); let is_stream = body.is_stream(); let text = body.extract_text_for_routing(); + // Resolve once, here, for the same reason as `route_typed_request`: + // only `get_by_model` understands aliases, so the policy and hash ring + // lookups below would silently fall back to router defaults on an + // alias. This path cannot reuse that resolution because multipart + // never goes through `route_typed_request`. + let canonical_model = self.worker_registry.resolve_model_alias(model_id); + let model_id = canonical_model.as_deref().unwrap_or(model_id); let endpoint = route_to_endpoint(route); Metrics::record_router_request( @@ -648,7 +675,7 @@ impl Router { events::RequestSentEvent { url: worker.url() }.emit(); - let form = match build_transcription_form(body, audio) { + let form = match build_transcription_form(body, audio, canonical_model.as_deref()) { Ok(f) => f, Err(e) => { let resp = error::bad_request("multipart_build_failed", e); @@ -877,12 +904,21 @@ impl Router { response } - // Send typed request directly without conversion + // Send typed request directly without conversion. + // + // `canonical_model` is set only when the client addressed the model by an + // alias. The worker was registered under the canonical ID and has never + // heard of the alias, so the body it receives carries the canonical name. + #[expect( + clippy::too_many_arguments, + reason = "per-request state threaded from route_typed_request_once; a struct would only move the arity" + )] async fn send_typed_request( &self, headers: Option<&HeaderMap>, typed_req: &T, route: &'static str, + canonical_model: Option<&str>, worker: &dyn Worker, is_stream: bool, load_guard: Option, @@ -890,7 +926,7 @@ impl Router { let api_key = worker.api_key().cloned(); let endpoint_url = worker.endpoint_url(route); - let json_val = match serde_json::to_value(typed_req) { + let mut json_val = match serde_json::to_value(typed_req) { Ok(j) => j, Err(e) => { return error::bad_request( @@ -900,6 +936,10 @@ impl Router { } }; + if let Some(canonical_model) = canonical_model { + super::set_request_model(&mut json_val, canonical_model); + } + let mut json_val = match worker.prepare_request(json_val) { Ok(prepared) => prepared, Err(e) => { @@ -1012,15 +1052,23 @@ impl Router { } } + /// Build the public rerank response. + /// + /// Rerank is the one HTTP route whose response the gateway constructs + /// itself instead of passing the worker's through, so the model it reports + /// has to be canonicalized here. `canonical_model` is set only when the + /// client addressed the model by an alias; reporting the alias would make + /// this route disagree with every other one about which model ran. async fn build_rerank_response( req: &RerankRequest, + canonical_model: Option<&str>, response: Response, ) -> anyhow::Result { let (_, response_body) = response.into_parts(); let body_bytes = to_bytes(response_body, usize::MAX).await?; let rerank_results = serde_json::from_slice::>(&body_bytes)?; - let mut rerank_response = - RerankResponse::new(rerank_results, req.model.clone(), req.rid.clone()); + let model = canonical_model.map_or_else(|| req.model.clone(), ToOwned::to_owned); + let mut rerank_response = RerankResponse::new(rerank_results, model, req.rid.clone()); // Sorting is handled by Python worker (serving_rerank.py) if let Some(top_k) = req.top_k { rerank_response.apply_top_k(top_k); @@ -1032,7 +1080,16 @@ impl Router { } } -fn build_transcription_form(body: &TranscriptionRequest, audio: AudioFile) -> Result { +/// Build the multipart body forwarded to the worker. +/// +/// `canonical_model` is set only when the client addressed the model by an +/// alias. The worker was registered under the canonical ID and has never heard +/// of the alias, so that is the name the form carries. +fn build_transcription_form( + body: &TranscriptionRequest, + audio: AudioFile, + canonical_model: Option<&str>, +) -> Result { let AudioFile { bytes, file_name, @@ -1050,9 +1107,10 @@ fn build_transcription_form(body: &TranscriptionRequest, audio: AudioFile) -> Re .map_err(|e| format!("Invalid audio content-type '{ct}': {e}"))?; } - let mut form = Form::new() - .part("file", file_part) - .text("model", body.model.clone()); + let mut form = Form::new().part("file", file_part).text( + "model", + canonical_model.map_or_else(|| body.model.clone(), ToOwned::to_owned), + ); if let Some(ref language) = body.language { form = form.text("language", language.clone()); @@ -1257,11 +1315,12 @@ impl RouterTrait for Router { body: &RerankRequest, model_id: &str, ) -> Response { + let canonical_model = self.worker_registry.resolve_model_alias(model_id); let response = self .route_typed_request(headers, body, "/v1/rerank", model_id) .await; if response.status().is_success() { - match Self::build_rerank_response(body, response).await { + match Self::build_rerank_response(body, canonical_model.as_deref(), response).await { Ok(rerank_response) => rerank_response, Err(e) => { error!("Failed to build rerank response: {}", e); diff --git a/model_gateway/src/routers/router_manager.rs b/model_gateway/src/routers/router_manager.rs index abaf682e1..924d8ba90 100644 --- a/model_gateway/src/routers/router_manager.rs +++ b/model_gateway/src/routers/router_manager.rs @@ -1200,4 +1200,32 @@ mod tests { assert_eq!(router.router_type(), "epd"); } } + + #[test] + fn weighted_routing_accepts_model_alias() { + let registry = Arc::new(WorkerRegistry::new()); + for (url, worker_type) in [ + ("http://prefill:8080", WorkerType::Prefill), + ("http://decode:8080", WorkerType::Decode), + ] { + let worker = BasicWorkerBuilder::new(url) + .worker_type(worker_type) + .connection_mode(ConnectionMode::Http) + .model(ModelCard::new("canonical-model").with_alias("model-alias")) + .circuit_breaker_config(CircuitBreakerConfig::default()) + .build(); + registry.register(Arc::new(worker)).unwrap(); + } + + let mut manager = RouterManager::new(registry, reqwest::Client::new()); + manager.enable_igw = true; + let manager = Arc::new(manager); + manager.register_router(router_ids::HTTP_REGULAR, Arc::new(StubRouter)); + manager.register_router(router_ids::HTTP_PD, Arc::new(PdStubRouter)); + + let router = manager + .select_router_for_request(Some("model-alias")) + .expect("alias should select the PD router"); + assert_eq!(router.router_type(), "pd"); + } } diff --git a/model_gateway/src/worker/registry.rs b/model_gateway/src/worker/registry.rs index ab9b05abe..e871c3ffb 100644 --- a/model_gateway/src/worker/registry.rs +++ b/model_gateway/src/worker/registry.rs @@ -30,7 +30,7 @@ use crate::{ event::WorkerEvent, hash_ring::HashRing, worker::{RuntimeType, WorkerType}, - ConnectionMode, Worker, DEFAULT_SAMPLING_PARAMS_LABEL, + ConnectionMode, Worker, DEFAULT_SAMPLING_PARAMS_LABEL, UNKNOWN_MODEL_ID, }, }; @@ -84,6 +84,9 @@ pub struct WorkerDescriptor { /// Updates create new snapshots (copy-on-write semantics). type ModelIndex = Arc]>>>; +/// Model alias to canonical model ID. +type ModelAliasIndex = Arc>>; + /// Worker registry with model-based indexing #[derive(Debug)] pub struct WorkerRegistry { @@ -94,6 +97,15 @@ pub struct WorkerRegistry { /// Uses Arc<[T]> instead of Arc>> for lock-free reads. model_index: ModelIndex, + /// Alias index kept separate from `model_index` so aliases do not appear as + /// models in discovery or statistics. + /// + /// Lock order: code that holds an entry of this map may then take a + /// `model_index` lock, never the reverse. Every alias write below reads + /// `model_index` while holding its own entry, so taking them in the other + /// order somewhere else would deadlock. + model_alias_index: ModelAliasIndex, + /// Consistent hash rings per model for O(log n) routing. /// Rebuilt on worker add/remove (copy-on-write). hash_rings: Arc>>, @@ -139,6 +151,7 @@ impl WorkerRegistry { Self { workers: Arc::new(DashMap::new()), model_index: Arc::new(DashMap::new()), + model_alias_index: Arc::new(DashMap::new()), hash_rings: Arc::new(DashMap::new()), type_workers: Arc::new(DashMap::new()), connection_workers: Arc::new(DashMap::new()), @@ -268,6 +281,11 @@ impl WorkerRegistry { /// Returns `Some(ring)` if any workers are registered for this model, /// `None` otherwise. The ring is pre-built and updated on worker add /// or remove, so reads are allocation-free apart from the Arc clone. + /// + /// Keyed by canonical model ID only, like every other per-model map on + /// this registry except [`Self::get_by_model`]. A caller holding a + /// client-supplied name resolves it with [`Self::resolve_model_alias`] + /// once at request entry and passes the canonical ID from there on. pub fn get_hash_ring(&self, model_id: &str) -> Option> { self.hash_rings.get(model_id).map(|r| Arc::clone(&r)) } @@ -279,19 +297,61 @@ impl WorkerRegistry { /// Empty worker slice constant returned when a lookup has no matches. const EMPTY_WORKERS: &'static [Arc] = &[]; - /// Return all workers serving a model as an immutable shared slice. + /// Return all workers serving a canonical model or alias. /// /// This is the fastest possible read path: the model index already /// stores the slice as an `Arc<[_]>`, so the return value is just an /// atomic refcount bump with zero contention. Returns an empty shared /// slice when the model is unknown. pub fn get_by_model(&self, model_id: &str) -> Arc<[Arc]> { - self.model_index + if let Some(workers) = self.model_index.get(model_id) { + return Arc::clone(&workers); + } + self.model_alias_index .get(model_id) - .map(|workers| Arc::clone(&workers)) + .and_then(|canonical_id| { + self.model_index + .get(canonical_id.as_ref()) + .map(|workers| Arc::clone(&workers)) + }) .unwrap_or_else(|| Arc::from(Self::EMPTY_WORKERS)) } + /// Resolve an alias to its canonical model ID without copying the string. + /// + /// Returns `None` for canonical model IDs, the `unknown` wildcard, and + /// unknown names. Callers can use the original input when this returns + /// `None`. + /// + /// # Rewriting an outbound request with the result + /// + /// Safe only where the alias is a second name for the same model, which is + /// what a self-hosted worker declares. It is NOT safe for external + /// providers: `workflow::steps::external::discover_models` groups a + /// provider's date-stamped variants under the shortest name, so `gpt-4o` + /// becomes canonical and `gpt-4o-2024-08-06` becomes its alias. Those name + /// different models — one floats to the provider's current release, the + /// other is pinned — and substituting one for the other would silently run + /// a model the client did not ask for. + /// + /// External workers reach only the OpenAI, Anthropic and Gemini routers + /// ([`RouterManager::select_router_for_workers`] gives them priority, and + /// single-router mode picks by routing mode), none of which rewrite the + /// outbound model. Keep it that way: canonicalize registry lookups there + /// if needed, never the request body. + /// + /// [`RouterManager::select_router_for_workers`]: crate::routers::RouterManager + pub fn resolve_model_alias(&self, model_id: &str) -> Option> { + if model_id == UNKNOWN_MODEL_ID || self.model_index.contains_key(model_id) { + return None; + } + + let canonical_id = self.model_alias_index.get(model_id)?; + self.model_index + .contains_key(canonical_id.as_ref()) + .then(|| Arc::clone(&canonical_id)) + } + /// Return all workers of a given type as an immutable shared slice. /// /// Unified with [`Self::get_by_model`] on `Arc<[_]>` so callers can @@ -474,6 +534,28 @@ impl WorkerRegistry { .collect() } + /// Whether at least one worker serves this name, as a canonical model ID + /// or as an alias. + /// + /// [`Self::get_models`] lists canonical IDs only, so a membership test + /// written against it rejects every alias. Endpoints that gate on "is + /// this model servable" use this instead. The `unknown` wildcard is not a + /// registered name and stays rejected. + pub fn contains_model(&self, model_id: &str) -> bool { + if self.model_has_workers(model_id) { + return true; + } + self.model_alias_index + .get(model_id) + .is_some_and(|canonical_id| self.model_has_workers(canonical_id.as_ref())) + } + + fn model_has_workers(&self, canonical_id: &str) -> bool { + self.model_index + .get(canonical_id) + .is_some_and(|workers| !workers.is_empty()) + } + /// Return the number of registered workers. pub fn len(&self) -> usize { self.workers.len() @@ -747,6 +829,7 @@ impl WorkerRegistry { for added_model in new_models.difference(&old_models) { self.add_worker_to_model_index(added_model, new_worker.clone()); self.rebuild_hash_ring(added_model); + self.drop_alias_shadowed_by_model(added_model); } // For models that stayed the same, update the worker reference in the index for kept_model in old_models.intersection(&new_models) { @@ -754,6 +837,19 @@ impl WorkerRegistry { self.rebuild_hash_ring(kept_model); } + // Update aliases after canonical indexes so alias removal can inspect + // the final worker set for each canonical model. + for model in old_worker.models() { + for alias in model.aliases { + self.remove_model_alias_if_unused(&alias, &model.id); + } + } + for model in new_worker.models() { + for alias in model.aliases { + self.add_model_alias(&alias, &model.id); + } + } + self.warn_on_sampling_defaults_divergence_for_worker(&new_worker); if old_worker.worker_type() != new_worker.worker_type() { @@ -983,6 +1079,11 @@ impl WorkerRegistry { self.model_retry_configs.remove(&model_id); } } + for model in worker.models() { + for alias in model.aliases { + self.remove_model_alias_if_unused(&alias, &model.id); + } + } if let Some(mut type_workers) = self.type_workers.get_mut(worker.worker_type()) { type_workers.retain(|id| id != worker_id); } @@ -1203,6 +1304,15 @@ impl WorkerRegistry { for model_id in Self::worker_model_ids(&worker) { self.add_worker_to_model_index(&model_id, worker.clone()); self.rebuild_hash_ring(&model_id); + // Run after the model index so `add_model_alias` below sees the + // model IDs this worker just contributed, and never while an index + // entry is held (see the lock order on `model_alias_index`). + self.drop_alias_shadowed_by_model(&model_id); + } + for model in worker.models() { + for alias in model.aliases { + self.add_model_alias(&alias, &model.id); + } } self.warn_on_sampling_defaults_divergence_for_worker(&worker); @@ -1285,6 +1395,144 @@ impl WorkerRegistry { self.rebuild_hash_ring(model_id); } + fn add_model_alias(&self, alias: &str, canonical_id: &str) { + if alias == canonical_id { + return; + } + if alias == UNKNOWN_MODEL_ID { + tracing::warn!( + alias, + canonical_id, + "Ignoring model alias reserved for wildcard routing" + ); + return; + } + // A name is either a canonical model ID or an alias, never both. + // Registered models win, so the alias is dropped rather than kept as a + // shadow that would take over once the model's last worker leaves. + // + // The collision check runs while the alias entry is held (alias + // before model, per the lock order on `model_alias_index`). Checked + // before taking the entry, a concurrent registration of a model + // named `alias` could slip between the check and the insert, run its + // `drop_alias_shadowed_by_model` against a still-absent entry, and + // leave the alias inserted here in place. Holding the entry makes + // that cleanup wait until the insert is visible. + let entry = self.model_alias_index.entry(alias.to_string()); + if self.model_index.contains_key(alias) { + tracing::warn!( + alias, + canonical_id, + "Ignoring model alias that collides with a registered model ID" + ); + // An occupied entry here is stale: the invariant above says the + // name cannot be both. Remove it instead of leaving it to take + // over once the model's last worker leaves. + if let Entry::Occupied(entry) = entry { + entry.remove(); + } + return; + } + + match entry { + Entry::Occupied(entry) if entry.get().as_ref() != canonical_id => tracing::warn!( + alias, + existing_canonical_id = entry.get().as_ref(), + ignored_canonical_id = canonical_id, + "Model alias already maps to a different canonical model" + ), + Entry::Occupied(_) => {} + Entry::Vacant(entry) => { + entry.insert(Arc::from(canonical_id)); + } + } + } + + /// Drop an alias entry that a newly registered model ID now shadows. + /// + /// Registration order decides which of the two arrives first, so the + /// collision has to be resolved from both sides: [`Self::add_model_alias`] + /// refuses an alias that names an existing model, and this refuses to keep + /// an existing alias that a new model now names. + fn drop_alias_shadowed_by_model(&self, model_id: &str) { + if let Some((alias, shadowed_canonical_id)) = self.model_alias_index.remove(model_id) { + tracing::warn!( + alias, + shadowed_canonical_id = shadowed_canonical_id.as_ref(), + "Dropping model alias that collides with a registered model ID" + ); + } + } + + /// Release `canonical_id`'s claim on `alias` after its declaring workers + /// are gone. + /// + /// Two models may declare the same alias; [`Self::add_model_alias`] keeps + /// whichever registered first and warns about the rest. Deleting the alias + /// outright when that winner leaves would strand the losers, which still + /// advertise it, so the alias is handed to a remaining model that declares + /// it and only deleted when none does. + /// + /// Holds the alias entry across the whole decision. Checking outside the + /// entry would let a concurrent registration insert the same alias between + /// the check and the delete, and the delete would erase it. + fn remove_model_alias_if_unused(&self, alias: &str, canonical_id: &str) { + let Entry::Occupied(mut entry) = self.model_alias_index.entry(alias.to_string()) else { + return; + }; + if entry.get().as_ref() != canonical_id { + return; + } + if self.model_declares_alias(canonical_id, alias) { + return; + } + + match self.find_model_declaring_alias(alias, canonical_id) { + Some(next_canonical_id) => { + tracing::debug!( + alias, + previous_canonical_id = canonical_id, + next_canonical_id = %next_canonical_id, + "Handing model alias to another model that declares it" + ); + entry.insert(Arc::from(next_canonical_id.as_str())); + } + None => { + entry.remove(); + } + } + } + + /// Whether any worker still serving `canonical_id` declares `alias`. + fn model_declares_alias(&self, canonical_id: &str, alias: &str) -> bool { + self.model_index.get(canonical_id).is_some_and(|workers| { + workers + .iter() + .any(|worker| Self::worker_declares_alias(worker, canonical_id, alias)) + }) + } + + /// Find a registered model other than `exclude` that declares `alias`. + fn find_model_declaring_alias(&self, alias: &str, exclude: &str) -> Option { + self.model_index.iter().find_map(|entry| { + let canonical_id = entry.key(); + if canonical_id == exclude { + return None; + } + entry + .value() + .iter() + .any(|worker| Self::worker_declares_alias(worker, canonical_id, alias)) + .then(|| canonical_id.clone()) + }) + } + + fn worker_declares_alias(worker: &Arc, canonical_id: &str, alias: &str) -> bool { + worker.models().into_iter().any(|model| { + model.id == canonical_id && model.aliases.iter().any(|candidate| candidate == alias) + }) + } + /// Shared backend for [`Self::transition_status`] and /// [`Self::transition_status_if_revision`]. Holds the per-worker /// mutation lock for the full read-modify-emit sequence. @@ -1526,6 +1774,21 @@ mod tests { Arc::new(builder.build()) } + fn worker_with_model_aliases( + url: &str, + canonical_id: &str, + aliases: &[&str], + worker_type: WorkerType, + ) -> Arc { + Arc::new( + BasicWorkerBuilder::new(url) + .model(ModelCard::new(canonical_id).with_aliases(aliases.iter().copied())) + .worker_type(worker_type) + .health_config(no_health_check()) + .build(), + ) + } + fn assert_sampling_defaults_group_values( registry: &WorkerRegistry, model_id: &str, @@ -2229,6 +2492,225 @@ mod tests { assert_eq!(stats.total_models, 2); } + #[test] + fn test_model_alias_resolves_to_canonical_model() { + let registry = WorkerRegistry::new(); + let aliased = worker_with_model_aliases( + "http://aliased:8080", + "GLM-5.2", + &["GLM-5.2-Coding"], + WorkerType::Prefill, + ); + let canonical_only = + worker_with_model_aliases("http://canonical:8080", "GLM-5.2", &[], WorkerType::Decode); + + registry.register(aliased).unwrap(); + registry.register(canonical_only).unwrap(); + + let canonical_workers = registry.get_by_model("GLM-5.2"); + assert_eq!(canonical_workers.len(), 2); + let alias_workers = registry.get_by_model("GLM-5.2-Coding"); + assert_eq!(alias_workers.len(), 2); + assert_eq!( + registry + .get_workers_filtered( + Some("GLM-5.2-Coding"), + Some(WorkerType::Prefill), + None, + None, + false, + ) + .len(), + 1 + ); + assert_eq!( + registry + .get_workers_filtered( + Some("GLM-5.2-Coding"), + Some(WorkerType::Decode), + None, + None, + false, + ) + .len(), + 1 + ); + + assert_eq!( + registry.resolve_model_alias("GLM-5.2-Coding").as_deref(), + Some("GLM-5.2") + ); + assert!(registry.resolve_model_alias("GLM-5.2").is_none()); + assert!(registry.get_by_model("glm-5.2-coding").is_empty()); + assert_eq!(registry.get_models(), vec!["GLM-5.2"]); + assert_eq!(registry.stats().total_models, 1); + } + + #[test] + fn test_model_alias_tracks_shared_workers_through_removal() { + let registry = WorkerRegistry::new(); + let first = worker_with_model_aliases( + "http://first:8080", + "GLM-5.2", + &["GLM-5.2-Coding"], + WorkerType::Prefill, + ); + let second = worker_with_model_aliases( + "http://second:8080", + "GLM-5.2", + &["GLM-5.2-Coding"], + WorkerType::Decode, + ); + let first_id = registry.register(first).unwrap(); + let second_id = registry.register(second).unwrap(); + + assert_eq!(registry.get_by_model("GLM-5.2-Coding").len(), 2); + registry.remove(&first_id).unwrap(); + let remaining = registry.get_by_model("GLM-5.2-Coding"); + assert_eq!(remaining.len(), 1); + assert_eq!(remaining[0].url(), "http://second:8080"); + + registry.remove(&second_id).unwrap(); + assert!(registry.get_by_model("GLM-5.2-Coding").is_empty()); + assert!(registry.resolve_model_alias("GLM-5.2-Coding").is_none()); + } + + #[test] + fn test_model_alias_replace_refreshes_kept_and_changed_aliases() { + let registry = WorkerRegistry::new(); + let original: Arc = Arc::new( + BasicWorkerBuilder::new("http://worker:8080") + .model(ModelCard::new("GLM-5.2").with_alias("old")) + .build(), + ); + let worker_id = registry.register(original).unwrap(); + let replacement: Arc = Arc::new( + BasicWorkerBuilder::new("http://worker:8080") + .model(ModelCard::new("GLM-5.2").with_alias("new")) + .build(), + ); + + assert!(registry.replace(&worker_id, replacement)); + + assert!(registry.get_by_model("old").is_empty()); + assert_eq!(registry.get_by_model("new").len(), 1); + } + + #[test] + fn test_contains_model_accepts_canonical_id_and_alias() { + let registry = WorkerRegistry::new(); + let worker = worker_with_model_aliases( + "http://worker:8080", + "GLM-5.2", + &["GLM-5.2-Coding"], + WorkerType::Regular, + ); + let worker_id = registry.register(worker).unwrap(); + + assert!(registry.contains_model("GLM-5.2")); + assert!(registry.contains_model("GLM-5.2-Coding")); + assert!(!registry.contains_model("GLM-5.2-Unknown")); + // The wildcard is a routing sentinel, never a registered name. + assert!(!registry.contains_model(UNKNOWN_MODEL_ID)); + + assert!(registry.remove(&worker_id).is_some()); + assert!(!registry.contains_model("GLM-5.2")); + assert!(!registry.contains_model("GLM-5.2-Coding")); + } + + #[test] + fn test_model_alias_is_handed_over_when_its_owner_leaves() { + // Two models declare the same alias. `add_model_alias` keeps the + // first, so removing the first must not strand the second, which + // still advertises the alias. + let registry = WorkerRegistry::new(); + let owner = worker_with_model_aliases( + "http://owner:8080", + "GLM-5.2", + &["shared-alias"], + WorkerType::Regular, + ); + let loser = worker_with_model_aliases( + "http://loser:8080", + "Qwen3", + &["shared-alias"], + WorkerType::Regular, + ); + let owner_id = registry.register(owner).unwrap(); + registry.register(loser).unwrap(); + + assert_eq!( + registry.resolve_model_alias("shared-alias").as_deref(), + Some("GLM-5.2") + ); + + assert!(registry.remove(&owner_id).is_some()); + + assert_eq!( + registry.resolve_model_alias("shared-alias").as_deref(), + Some("Qwen3"), + "alias must move to the remaining model that declares it" + ); + let workers = registry.get_by_model("shared-alias"); + assert_eq!(workers.len(), 1); + assert_eq!(workers[0].url(), "http://loser:8080"); + } + + #[test] + fn test_model_alias_never_shadows_a_registered_model_id() { + // Registered later than the model it collides with. + let registry = WorkerRegistry::new(); + let canonical = + worker_with_model_aliases("http://canonical:8080", "GLM-5.2", &[], WorkerType::Regular); + let colliding = worker_with_model_aliases( + "http://colliding:8080", + "Qwen3", + &["GLM-5.2"], + WorkerType::Regular, + ); + let canonical_id = registry.register(canonical).unwrap(); + registry.register(colliding).unwrap(); + + assert!(registry.resolve_model_alias("GLM-5.2").is_none()); + // The real model has gone, so the name must resolve to nothing at all + // rather than start pointing at Qwen3. + assert!(registry.remove(&canonical_id).is_some()); + assert!(registry.resolve_model_alias("GLM-5.2").is_none()); + assert!(registry.get_by_model("GLM-5.2").is_empty()); + } + + #[test] + fn test_registered_model_id_evicts_a_colliding_alias() { + // Registered earlier than the model it collides with: same invariant, + // resolved from the other side. + let registry = WorkerRegistry::new(); + let colliding = worker_with_model_aliases( + "http://colliding:8080", + "Qwen3", + &["GLM-5.2"], + WorkerType::Regular, + ); + registry.register(colliding).unwrap(); + assert_eq!( + registry.resolve_model_alias("GLM-5.2").as_deref(), + Some("Qwen3") + ); + + let canonical = + worker_with_model_aliases("http://canonical:8080", "GLM-5.2", &[], WorkerType::Regular); + let canonical_id = registry.register(canonical).unwrap(); + + assert!(registry.resolve_model_alias("GLM-5.2").is_none()); + assert_eq!(registry.get_by_model("GLM-5.2").len(), 1); + assert_eq!( + registry.get_by_model("GLM-5.2")[0].url(), + "http://canonical:8080" + ); + + assert!(registry.remove(&canonical_id).is_some()); + assert!(registry.get_by_model("GLM-5.2").is_empty()); + } + #[test] fn test_replace_same_url_refreshes_all_model_indexes() { let registry = WorkerRegistry::new(); diff --git a/model_gateway/tests/common/mock_worker.rs b/model_gateway/tests/common/mock_worker.rs index 045c523c3..092d062fe 100755 --- a/model_gateway/tests/common/mock_worker.rs +++ b/model_gateway/tests/common/mock_worker.rs @@ -12,7 +12,7 @@ use std::{ }; use axum::{ - extract::{Json, Path, State}, + extract::{Json, Multipart, Path, State}, http::StatusCode, response::{ sse::{Event, KeepAlive}, @@ -155,6 +155,96 @@ fn clear_scheduler_controls(port: u16) { } } +/// Records the JSON body of every chat request a mock worker receives. +/// +/// A test can only read the mock's canned answer, which says nothing about +/// what the gateway actually forwarded. This exposes the forwarded body, so a +/// test can assert on rewrites the gateway performs on the way out — for +/// instance that a request addressed to a model alias reaches the worker under +/// the canonical model ID. +/// +/// Kept in a port-keyed table rather than in [`MockWorkerConfig`], for the +/// same reason as [`SchedulerControls`]: adding a knob must not force every +/// existing `MockWorkerConfig { .. }` literal in the suite to gain a field. +#[derive(Default)] +pub struct RequestRecorder { + bodies: Mutex>, +} + +impl RequestRecorder { + pub fn new() -> Arc { + Arc::new(Self::default()) + } + + /// Every body received so far, oldest first. + #[expect( + clippy::expect_used, + reason = "test helper - panicking on failure is intentional" + )] + pub fn bodies(&self) -> Vec { + self.bodies + .lock() + .expect("request recorder mutex poisoned") + .clone() + } + + /// The single body received, panicking unless exactly one arrived. + #[expect( + clippy::expect_used, + reason = "test helper - panicking on failure is intentional" + )] + pub fn only_body(&self) -> serde_json::Value { + let bodies = self.bodies(); + assert_eq!( + bodies.len(), + 1, + "expected exactly one recorded request, got {}", + bodies.len() + ); + bodies.into_iter().next().expect("length checked above") + } +} + +static REQUEST_RECORDERS: OnceLock>>> = OnceLock::new(); + +fn request_recorders_table() -> &'static Mutex>> { + REQUEST_RECORDERS.get_or_init(|| Mutex::new(HashMap::new())) +} + +/// Attach a recorder to the mock worker that will bind `port`. Call before the +/// worker starts, like [`set_scheduler_controls`]. +#[expect( + clippy::expect_used, + reason = "test helper - panicking on failure is intentional" +)] +pub fn set_request_recorder(port: u16, recorder: Arc) { + request_recorders_table() + .lock() + .expect("request recorders mutex poisoned") + .insert(port, recorder); +} + +fn record_request(port: u16, body: &serde_json::Value) { + let recorder = request_recorders_table() + .lock() + .ok() + .and_then(|table| table.get(&port).cloned()); + if let Some(recorder) = recorder { + if let Ok(mut bodies) = recorder.bodies.lock() { + bodies.push(body.clone()); + } + } +} + +/// Remove a port's recorder on teardown so a later worker reusing the port +/// does not append to it. Tolerant of a poisoned mutex since it runs from +/// `Drop`. +fn clear_request_recorder(port: u16) { + if let Ok(mut table) = request_recorders_table().lock() { + table.remove(&port); + } +} + /// Configuration for mock worker behavior #[derive(Clone)] pub struct MockWorkerConfig { @@ -233,6 +323,10 @@ impl MockWorker { .route("/v1/messages", post(messages_handler)) .route("/v1/completions", post(completions_handler)) .route("/v1/rerank", post(rerank_handler)) + .route( + "/v1/audio/transcriptions", + post(audio_transcriptions_handler), + ) .route("/v1/responses", post(responses_handler)) .route("/v1/responses/{response_id}", get(responses_get_handler)) .route( @@ -293,10 +387,11 @@ impl Drop for MockWorker { if let Some(shutdown_tx) = self.shutdown_tx.take() { let _ = shutdown_tx.send(()); } - // Prune our scheduler controls so a later worker reusing this port - // doesn't inherit stale state. + // Prune our scheduler controls and recorder so a later worker reusing + // this port doesn't inherit stale state. if let Some(port) = self.bound_port { clear_scheduler_controls(port); + clear_request_recorder(port); } } } @@ -585,6 +680,9 @@ async fn chat_completions_handler( Json(payload): Json, ) -> Response { let config = config.read().await; + // Before any early return, so a test still sees what arrived even when the + // mock is configured to fail the request. + record_request(config.port, &payload); if should_fail(&config) { return ( @@ -1530,6 +1628,39 @@ fn response_exists_for_port(port: u16, response_id: &str) -> bool { .unwrap_or(false) } +/// Minimal `/v1/audio/transcriptions` handler. +/// +/// Collects the multipart text fields into a JSON object and records it, so a +/// test can assert on the model name the gateway put in the form. The audio +/// part itself is drained and discarded. +async fn audio_transcriptions_handler( + State(config): State>>, + mut multipart: Multipart, +) -> Response { + let config = config.read().await; + + let mut fields = serde_json::Map::new(); + while let Ok(Some(field)) = multipart.next_field().await { + let Some(name) = field.name().map(ToOwned::to_owned) else { + continue; + }; + if name == "file" { + let _ = field.bytes().await; + continue; + } + if let Ok(value) = field.text().await { + fields.insert(name, serde_json::Value::String(value)); + } + } + record_request(config.port, &serde_json::Value::Object(fields)); + + if should_fail(&config) { + return (StatusCode::INTERNAL_SERVER_ERROR, "Simulated failure").into_response(); + } + + Json(json!({"text": "mock transcription"})).into_response() +} + // Minimal rerank handler returning mock results; router shapes final response #[expect( clippy::unwrap_used, @@ -1540,6 +1671,7 @@ async fn rerank_handler( Json(payload): Json, ) -> impl IntoResponse { let config = config.read().await; + record_request(config.port, &payload); // Simulate response delay if config.response_delay_ms > 0 { diff --git a/model_gateway/tests/routing/mod.rs b/model_gateway/tests/routing/mod.rs index f5bfbb5ba..3ad99c245 100644 --- a/model_gateway/tests/routing/mod.rs +++ b/model_gateway/tests/routing/mod.rs @@ -5,6 +5,7 @@ pub mod grpc_completion_batch_test; pub mod header_forwarding_test; pub mod load_balancing_test; pub mod manual_routing_test; +pub mod model_alias_test; pub mod payload_size_test; pub mod pd_routing_test; pub mod policy_registry_integration; diff --git a/model_gateway/tests/routing/model_alias_test.rs b/model_gateway/tests/routing/model_alias_test.rs new file mode 100644 index 000000000..b1418a634 --- /dev/null +++ b/model_gateway/tests/routing/model_alias_test.rs @@ -0,0 +1,234 @@ +//! Model alias integration tests for the regular (non-PD) HTTP router. +//! +//! A worker is registered under its canonical model ID and declares aliases +//! alongside it. A client may address the model by either name, but the worker +//! itself only knows the canonical one, so the gateway resolves the alias for +//! its own routing decisions and forwards the canonical name downstream. +//! +//! The PD equivalent lives in `pd_routing_test::test_pd_model_alias`. + +use std::sync::Arc; + +use axum::{ + body::{to_bytes, Body}, + extract::Request, + http::{header::CONTENT_TYPE, StatusCode}, +}; +use openai_protocol::model_card::ModelCard; +use serde_json::json; +use smg::worker::BasicWorkerBuilder; +use tower::ServiceExt; + +use crate::common::{ + mock_worker::{set_request_recorder, RequestRecorder}, + AppTestContext, TestWorkerConfig, +}; + +#[cfg(test)] +mod model_alias_tests { + use super::*; + + const WORKER_PORT: u16 = 19860; + const CANONICAL_MODEL: &str = "GLM-5.2"; + const MODEL_ALIAS: &str = "GLM-5.2-Coding"; + + /// Re-register the started worker under a canonical model ID plus an + /// alias. Mock workers advertise a generic model card, so the alias has to + /// be installed after startup. + fn declare_alias(ctx: &AppTestContext, url: &str) { + let registry = &ctx.app_context.worker_registry; + let worker_id = registry.get_id_by_url(url).unwrap(); + let worker = registry.get(&worker_id).unwrap(); + let mut spec = worker.metadata().spec.as_ref().clone(); + spec.models = vec![ModelCard::new(CANONICAL_MODEL).with_alias(MODEL_ALIAS)].into(); + let replacement = BasicWorkerBuilder::from_spec(spec) + .health_config(worker.metadata().health_config.clone()) + .health_endpoint(&worker.metadata().health_endpoint) + .build(); + assert!(registry.replace(&worker_id, Arc::new(replacement))); + } + + fn chat_request(model: &str) -> Request { + Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header(CONTENT_TYPE, "application/json") + .body(Body::from( + json!({ + "model": model, + "messages": [{"role": "user", "content": "Hello"}], + "stream": false + }) + .to_string(), + )) + .unwrap() + } + + /// The worker must receive the canonical model ID, not the alias the + /// client sent. Forwarding the alias would hand the backend a model it + /// cannot serve. + #[tokio::test] + async fn test_regular_routing_forwards_canonical_model_for_alias() { + let recorder = RequestRecorder::new(); + set_request_recorder(WORKER_PORT, Arc::clone(&recorder)); + + let ctx = AppTestContext::new(vec![TestWorkerConfig::healthy(WORKER_PORT)]).await; + let app = ctx.create_app(); + declare_alias(&ctx, &format!("http://127.0.0.1:{WORKER_PORT}")); + + let response = app + .clone() + .oneshot(chat_request(MODEL_ALIAS)) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + + let forwarded = recorder.only_body(); + assert_eq!( + forwarded["model"], CANONICAL_MODEL, + "worker must receive the canonical model ID, got {}", + forwarded["model"] + ); + + ctx.shutdown().await; + } + + /// The canonical ID keeps working and is forwarded unchanged. + #[tokio::test] + async fn test_regular_routing_leaves_canonical_model_untouched() { + let recorder = RequestRecorder::new(); + set_request_recorder(WORKER_PORT + 1, Arc::clone(&recorder)); + + let ctx = AppTestContext::new(vec![TestWorkerConfig::healthy(WORKER_PORT + 1)]).await; + let app = ctx.create_app(); + declare_alias(&ctx, &format!("http://127.0.0.1:{}", WORKER_PORT + 1)); + + let response = app + .clone() + .oneshot(chat_request(CANONICAL_MODEL)) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(recorder.only_body()["model"], CANONICAL_MODEL); + + ctx.shutdown().await; + } + + /// Rerank is the one HTTP route whose response the gateway builds itself + /// rather than passing the worker's through, so it needs its own check + /// that the reported model is the one that actually ran. + /// + /// Uses `/rerank`, not `/v1/rerank`: the latter takes only `query` and + /// `documents` and hardcodes the model to the `unknown` wildcard, so no + /// alias can reach the router through it. + #[tokio::test] + async fn test_rerank_forwards_and_reports_the_canonical_model() { + let recorder = RequestRecorder::new(); + set_request_recorder(WORKER_PORT + 3, Arc::clone(&recorder)); + + let ctx = AppTestContext::new(vec![TestWorkerConfig::healthy(WORKER_PORT + 3)]).await; + let app = ctx.create_app(); + declare_alias(&ctx, &format!("http://127.0.0.1:{}", WORKER_PORT + 3)); + + let request = Request::builder() + .method("POST") + .uri("/rerank") + .header(CONTENT_TYPE, "application/json") + .body(Body::from( + json!({ + "model": MODEL_ALIAS, + "query": "what is rust", + "documents": ["a systems language", "a kind of oxide"] + }) + .to_string(), + )) + .unwrap(); + let response = app.clone().oneshot(request).await.unwrap(); + assert_eq!(response.status(), StatusCode::OK); + + assert_eq!( + recorder.only_body()["model"], + CANONICAL_MODEL, + "worker must receive the canonical model ID" + ); + + let body = to_bytes(response.into_body(), usize::MAX).await.unwrap(); + let json: serde_json::Value = serde_json::from_slice(&body).unwrap(); + assert_eq!( + json["model"], CANONICAL_MODEL, + "rerank must report the model that ran, got {}", + json["model"] + ); + + ctx.shutdown().await; + } + + /// `/v1/audio/transcriptions` builds a multipart form instead of a JSON + /// body and never passes through `route_typed_request`, so it canonicalizes + /// on its own path. + #[tokio::test] + async fn test_transcription_form_carries_the_canonical_model() { + let recorder = RequestRecorder::new(); + set_request_recorder(WORKER_PORT + 4, Arc::clone(&recorder)); + + let ctx = AppTestContext::new(vec![TestWorkerConfig::healthy(WORKER_PORT + 4)]).await; + let app = ctx.create_app(); + declare_alias(&ctx, &format!("http://127.0.0.1:{}", WORKER_PORT + 4)); + + let boundary = "alias-test-boundary"; + let form = format!( + "--{boundary}\r\n\ + Content-Disposition: form-data; name=\"file\"; filename=\"a.wav\"\r\n\ + Content-Type: audio/wav\r\n\r\n\ + RIFFmock\r\n\ + --{boundary}\r\n\ + Content-Disposition: form-data; name=\"model\"\r\n\r\n\ + {MODEL_ALIAS}\r\n\ + --{boundary}--\r\n" + ); + let request = Request::builder() + .method("POST") + .uri("/v1/audio/transcriptions") + .header( + CONTENT_TYPE, + format!("multipart/form-data; boundary={boundary}"), + ) + .body(Body::from(form)) + .unwrap(); + + let response = app.clone().oneshot(request).await.unwrap(); + assert_eq!(response.status(), StatusCode::OK); + + assert_eq!( + recorder.only_body()["model"], + CANONICAL_MODEL, + "the multipart form must carry the canonical model ID" + ); + + ctx.shutdown().await; + } + + /// A name that is neither a canonical model ID nor an alias is still + /// rejected, and nothing reaches the worker. + #[tokio::test] + async fn test_regular_routing_rejects_an_unknown_model() { + let recorder = RequestRecorder::new(); + set_request_recorder(WORKER_PORT + 2, Arc::clone(&recorder)); + + let ctx = AppTestContext::new(vec![TestWorkerConfig::healthy(WORKER_PORT + 2)]).await; + let app = ctx.create_app(); + declare_alias(&ctx, &format!("http://127.0.0.1:{}", WORKER_PORT + 2)); + + let response = app + .clone() + .oneshot(chat_request("GLM-5.2-Nonexistent")) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::NOT_FOUND); + let body = to_bytes(response.into_body(), usize::MAX).await.unwrap(); + assert!(!body.is_empty()); + assert!(recorder.bodies().is_empty()); + + ctx.shutdown().await; + } +} diff --git a/model_gateway/tests/routing/pd_routing_test.rs b/model_gateway/tests/routing/pd_routing_test.rs index 510bb1fed..50dae8830 100644 --- a/model_gateway/tests/routing/pd_routing_test.rs +++ b/model_gateway/tests/routing/pd_routing_test.rs @@ -2,17 +2,22 @@ //! //! Tests for prefill-decode disaggregation routing mode. +use std::sync::Arc; + use axum::{ - body::Body, + body::{to_bytes, Body}, extract::Request, http::{header::CONTENT_TYPE, StatusCode}, }; +use openai_protocol::model_card::ModelCard; use serde_json::json; -use smg::config::RouterConfig; +use smg::{config::RouterConfig, worker::BasicWorkerBuilder}; use tower::ServiceExt; use crate::common::{ - mock_worker::{HealthStatus, MockWorkerConfig, WorkerType}, + mock_worker::{ + set_request_recorder, HealthStatus, MockWorkerConfig, RequestRecorder, WorkerType, + }, AppTestContext, TestWorkerConfig, }; @@ -20,6 +25,9 @@ use crate::common::{ mod pd_routing_tests { use super::*; + const CANONICAL_MODEL: &str = "GLM-5.2"; + const MODEL_ALIAS: &str = "GLM-5.2-Coding"; + /// Test basic PD mode routing with prefill and decode workers #[tokio::test] async fn test_pd_mode_basic_routing() { @@ -151,6 +159,91 @@ mod pd_routing_tests { ctx.shutdown().await; } + /// A request addressed to a model alias must reach both PD workers under + /// the canonical model ID. The workers were registered under the canonical + /// ID and have never heard of the alias, so forwarding the alias would + /// hand them a model they cannot serve. + #[tokio::test] + async fn test_pd_model_alias() { + let prefill_url = "http://127.0.0.1:19840".to_string(); + let decode_url = "http://127.0.0.1:19841".to_string(); + + let prefill_recorder = RequestRecorder::new(); + let decode_recorder = RequestRecorder::new(); + set_request_recorder(19840, Arc::clone(&prefill_recorder)); + set_request_recorder(19841, Arc::clone(&decode_recorder)); + + let mut config = RouterConfig::builder() + .prefill_decode_mode(vec![(prefill_url.clone(), None)], vec![decode_url.clone()]) + .round_robin_policy() + .host("127.0.0.1") + .port(3804) + .max_payload_size(256 * 1024 * 1024) + .request_timeout_secs(600) + .worker_startup_timeout_secs(5) + .worker_startup_check_interval_secs(1) + .max_concurrent_requests(64) + .queue_timeout_secs(60) + .build_unchecked(); + config.health_check.disable_health_check = true; + + let ctx = AppTestContext::new_with_config( + config, + vec![ + TestWorkerConfig::prefill(19840), + TestWorkerConfig::decode(19841), + ], + ) + .await; + let app = ctx.create_app(); + + let registry = &ctx.app_context.worker_registry; + for url in [&prefill_url, &decode_url] { + let worker_id = registry.get_id_by_url(url).unwrap(); + let worker = registry.get(&worker_id).unwrap(); + let mut spec = worker.metadata().spec.as_ref().clone(); + spec.models = vec![ModelCard::new(CANONICAL_MODEL).with_alias(MODEL_ALIAS)].into(); + let replacement = BasicWorkerBuilder::from_spec(spec) + .health_config(worker.metadata().health_config.clone()) + .health_endpoint(&worker.metadata().health_endpoint) + .build(); + assert!(registry.replace(&worker_id, Arc::new(replacement))); + } + + let alias_request = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header(CONTENT_TYPE, "application/json") + .body(Body::from( + json!({ + "model": MODEL_ALIAS, + "messages": [{"role": "user", "content": "Hello"}], + "stream": false + }) + .to_string(), + )) + .unwrap(); + let alias_response = app.clone().oneshot(alias_request).await.unwrap(); + assert_eq!(alias_response.status(), StatusCode::OK); + let alias_body = to_bytes(alias_response.into_body(), usize::MAX) + .await + .unwrap(); + let alias_json: serde_json::Value = serde_json::from_slice(&alias_body).unwrap(); + assert_eq!(alias_json["model"], "mock-model"); + + // The part the response body cannot show: what the gateway forwarded. + for (leg, recorder) in [("prefill", &prefill_recorder), ("decode", &decode_recorder)] { + let forwarded = recorder.only_body(); + assert_eq!( + forwarded["model"], CANONICAL_MODEL, + "{leg} worker must receive the canonical model ID, got {}", + forwarded["model"] + ); + } + + ctx.shutdown().await; + } + /// A non-streaming PD request must emit the SMG-only PD metrics, including /// the honest `smg_pd_ttft_seconds`. Runs on a current-thread runtime so the /// thread-local Prometheus recorder captures emissions from the request path.