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.