diff --git a/model_gateway/src/routers/grpc/client.rs b/model_gateway/src/routers/grpc/client.rs index 990a291d1..e92422e60 100644 --- a/model_gateway/src/routers/grpc/client.rs +++ b/model_gateway/src/routers/grpc/client.rs @@ -184,6 +184,19 @@ impl GrpcClient { matches!(self, Self::TokenSpeed(_)) } + /// Runtime type backing this client. Lets shared logic (e.g. the multimodal + /// capability matrix) key on the backend without matching every variant. + pub fn runtime_type(&self) -> crate::worker::RuntimeType { + use crate::worker::RuntimeType; + match self { + Self::Sglang(_) => RuntimeType::Sglang, + Self::Vllm(_) => RuntimeType::Vllm, + Self::Trtllm(_) => RuntimeType::Trtllm, + Self::Mlx(_) => RuntimeType::Mlx, + Self::TokenSpeed(_) => RuntimeType::TokenSpeed, + } + } + pub async fn connect( url: &str, runtime_type: &str, diff --git a/model_gateway/src/routers/grpc/common/stages/worker_selection.rs b/model_gateway/src/routers/grpc/common/stages/worker_selection.rs index f32a7f446..a8d06bfb7 100644 --- a/model_gateway/src/routers/grpc/common/stages/worker_selection.rs +++ b/model_gateway/src/routers/grpc/common/stages/worker_selection.rs @@ -170,6 +170,23 @@ impl PipelineStage for WorkerSelectionStage { } }; + // Reject an unsupported (backend, modality) combination now that the + // runtime is known, before request building fetches/preprocesses media + // only to fail deep in assembly. The prefill leg builds the request in + // disaggregated mode, so its runtime is the one that must support the + // request's modalities. + if let Some(intermediate) = multimodal_intermediate(prep) { + if let Err(err) = multimodal::ensure_backend_supports_modalities( + selection_runtime(&workers), + intermediate, + ) { + return Err(error::bad_request( + "multimodal_not_supported", + format!("{err}"), + )); + } + } + ctx.state.workers = Some(workers); Ok(None) } @@ -179,6 +196,30 @@ impl PipelineStage for WorkerSelectionStage { } } +/// Runtime of the leg that builds the generate request: the sole worker in +/// regular mode, the prefill worker in disaggregated (PD/EPD) mode. +fn selection_runtime(workers: &WorkerSelection) -> RuntimeType { + match workers { + WorkerSelection::Single { worker } => worker.metadata().spec.runtime_type, + WorkerSelection::Disaggregated { runtime_type, .. } => *runtime_type, + } +} + +/// Borrow the request's multimodal intermediate, if any. +fn multimodal_intermediate( + prep: &PreparationOutput, +) -> Option<&multimodal::MultimodalIntermediate> { + match prep { + PreparationOutput::Chat { + processed_messages, .. + } + | PreparationOutput::Messages { + processed_messages, .. + } => processed_messages.multimodal_intermediate.as_ref(), + _ => None, + } +} + impl WorkerSelectionStage { fn select_single_worker( &self, @@ -558,16 +599,7 @@ impl WorkerSelectionStage { } fn encode_item_hashes(prep: &PreparationOutput) -> anyhow::Result>> { - let intermediate = match prep { - PreparationOutput::Chat { - processed_messages, .. - } - | PreparationOutput::Messages { - processed_messages, .. - } => processed_messages.multimodal_intermediate.as_ref(), - _ => None, - }; - let Some(intermediate) = intermediate else { + let Some(intermediate) = multimodal_intermediate(prep) else { return Ok(Vec::new()); }; multimodal::encode_routing_hashes(intermediate) diff --git a/model_gateway/src/routers/grpc/multimodal/assemble.rs b/model_gateway/src/routers/grpc/multimodal/assemble.rs index a18028392..f016b6643 100644 --- a/model_gateway/src/routers/grpc/multimodal/assemble.rs +++ b/model_gateway/src/routers/grpc/multimodal/assemble.rs @@ -19,6 +19,7 @@ use smg_grpc_client::common_proto as common; use tracing::{info, warn}; use super::{ + capability::ensure_backend_supports_modalities, log_mm_timing_enabled, serialize::{ model_specific_to_tensor_bytes, serialize_array_as_tokenspeed_tensor, @@ -66,17 +67,21 @@ async fn assemble_multimodal_data_impl( omit_prefill_pixels: bool, ) -> Result { validate_intermediate(&intermediate)?; + // Defense in depth: worker selection already rejected unsupported (backend, + // modality) combinations. Re-assert here against the same capability matrix + // so assembly can never silently mishandle a modality the backend lacks. + ensure_client_supports_intermediate(client, &intermediate)?; match client { GrpcClient::Sglang(_) => { - let batch = into_single_image_batch(intermediate, "SGLang")?; + let batch = into_single_batch(intermediate, "SGLang")?; Ok(MultimodalData::Sglang(assemble_sglang(batch)?)) } GrpcClient::Vllm(_) => { - let batch = into_single_vision_batch(intermediate, "vLLM")?; + let batch = into_single_batch(intermediate, "vLLM")?; Ok(MultimodalData::Vllm(assemble_vllm(batch, workers)?)) } GrpcClient::Trtllm(_) => { - let batch = into_single_image_batch(intermediate, "TRT-LLM")?; + let batch = into_single_batch(intermediate, "TRT-LLM")?; Ok(MultimodalData::Trtllm(assemble_trtllm(batch)?)) } GrpcClient::TokenSpeed(_) => { @@ -132,46 +137,34 @@ impl Drop for PendingTokenSpeedAssembly { } } -fn into_single_image_batch( +/// Backends other than TokenSpeed take a single preprocessed batch (one +/// modality). The per-modality capability is enforced by +/// [`ensure_client_supports_intermediate`]; this only enforces the structural +/// single-batch constraint (multiple modalities in one request are TokenSpeed- +/// only). +fn into_single_batch( intermediate: MultimodalIntermediate, backend: &str, ) -> Result { anyhow::ensure!( intermediate.batches().len() == 1, - "{backend} multimodal path requires exactly one image batch; got {} batches", + "{backend} multimodal path requires exactly one batch; got {} batches", intermediate.batches().len() ); - let mut batches = intermediate.into_batches(); - let batch = batches + intermediate + .into_batches() .pop() - .context("multimodal intermediate is missing its sole batch")?; - anyhow::ensure!( - matches!(&batch.media, MediaBatch::Images(_)), - "{backend} multimodal path currently supports image inputs only; got {}", - batch.media.modality() - ); - Ok(batch) + .context("multimodal intermediate is missing its sole batch") } -fn into_single_vision_batch( - intermediate: MultimodalIntermediate, - backend: &str, -) -> Result { - anyhow::ensure!( - intermediate.batches().len() == 1, - "{backend} multimodal path requires exactly one vision batch; got {} batches", - intermediate.batches().len() - ); - let mut batches = intermediate.into_batches(); - let batch = batches - .pop() - .context("multimodal intermediate is missing its sole batch")?; - anyhow::ensure!( - matches!(&batch.media, MediaBatch::Images(_) | MediaBatch::Videos(_)), - "{backend} multimodal path currently supports image and video inputs only; got {}", - batch.media.modality() - ); - Ok(batch) +/// Defense-in-depth capability assertion mirroring the early worker-selection +/// check, keyed on the concrete backend client. See +/// [`crate::routers::grpc::multimodal::capability::ensure_backend_supports_modalities`]. +fn ensure_client_supports_intermediate( + client: &GrpcClient, + intermediate: &MultimodalIntermediate, +) -> Result<()> { + ensure_backend_supports_modalities(client.runtime_type(), intermediate) } fn assemble_sglang( diff --git a/model_gateway/src/routers/grpc/multimodal/capability.rs b/model_gateway/src/routers/grpc/multimodal/capability.rs new file mode 100644 index 000000000..8c7517506 --- /dev/null +++ b/model_gateway/src/routers/grpc/multimodal/capability.rs @@ -0,0 +1,170 @@ +//! Backend x modality capability: the single source of truth for which gRPC +//! engine supports which input modality. +//! +//! Previously this was implicit and duplicated across assembly (per-backend +//! `into_single_image_batch` / `into_single_vision_batch` / ad-hoc `bail!`s with +//! divergent messages). Centralizing it here lets the pipeline reject an +//! unsupported (engine, modality) request early -- at worker selection, before +//! any media is fetched or preprocessed -- with one consistent message, and lets +//! assembly assert against the same matrix as defense in depth. + +use anyhow::Result; +use llm_multimodal::Modality; + +use super::MultimodalIntermediate; +use crate::worker::RuntimeType; + +/// Whether `runtime` accepts multimodal inputs of `modality`. +/// +/// This is the authoritative capability matrix. Truth, transcribed from the +/// per-backend assembly arms: +/// - Image: SGLang, vLLM, TRT-LLM, TokenSpeed +/// - ImageEmbeds: TokenSpeed only (pre-computed embeddings are an EPD feature) +/// - Video: vLLM, TokenSpeed +/// - Audio: TokenSpeed +/// - MLX: none +/// +/// `Modality::ImageEmbeds` cannot actually reach the early check today because a +/// [`crate::routers::grpc::multimodal::MediaBatch`] only ever yields +/// Image/Video/Audio; it is kept in the matrix for correctness/defense in depth. +pub(crate) fn runtime_supports_modality(runtime: RuntimeType, modality: Modality) -> bool { + match modality { + Modality::Image => matches!( + runtime, + RuntimeType::Sglang | RuntimeType::Vllm | RuntimeType::Trtllm | RuntimeType::TokenSpeed + ), + Modality::ImageEmbeds => matches!(runtime, RuntimeType::TokenSpeed), + Modality::Video => matches!(runtime, RuntimeType::Vllm | RuntimeType::TokenSpeed), + Modality::Audio => matches!(runtime, RuntimeType::TokenSpeed), + } +} + +/// Reject early if the selected backend does not support every modality present +/// in the request. Runs at worker selection, once the runtime is known but +/// before media is fetched/preprocessed, so an unsupported combination fails +/// fast with one clear message instead of dying deep in assembly. +pub(crate) fn ensure_backend_supports_modalities( + runtime: RuntimeType, + intermediate: &MultimodalIntermediate, +) -> Result<()> { + for batch in intermediate.batches() { + let modality = batch.media.modality(); + anyhow::ensure!( + runtime_supports_modality(runtime, modality), + "backend {runtime} does not support {modality} inputs" + ); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// The full backend x modality matrix, mirroring the pre-refactor assembly + /// dispatch so a behavior change surfaces here. + #[test] + fn capability_matrix_matches_backend_support() { + use Modality::{Audio, Image, ImageEmbeds, Video}; + use RuntimeType::{Mlx, Sglang, TokenSpeed, Trtllm, Vllm}; + + // (runtime, modality, expected_supported) + let cases = [ + (Sglang, Image, true), + (Sglang, Video, false), + (Sglang, Audio, false), + (Vllm, Image, true), + (Vllm, Video, true), + (Vllm, Audio, false), + (Trtllm, Image, true), + (Trtllm, Video, false), + (Trtllm, Audio, false), + (TokenSpeed, Image, true), + (TokenSpeed, Video, true), + (TokenSpeed, Audio, true), + (Mlx, Image, false), + (Mlx, Video, false), + (Mlx, Audio, false), + // ImageEmbeds is a TokenSpeed-only EPD feature. + (Sglang, ImageEmbeds, false), + (Vllm, ImageEmbeds, false), + (Trtllm, ImageEmbeds, false), + (TokenSpeed, ImageEmbeds, true), + (Mlx, ImageEmbeds, false), + ]; + + for (runtime, modality, expected) in cases { + assert_eq!( + runtime_supports_modality(runtime, modality), + expected, + "runtime={runtime} modality={modality} expected supported={expected}" + ); + } + } + + /// Non-gRPC runtimes are never routed to the multimodal gRPC path; they + /// support nothing here. + #[test] + fn non_grpc_runtimes_support_no_modality() { + for runtime in [RuntimeType::Unspecified, RuntimeType::External] { + for modality in [Modality::Image, Modality::Video, Modality::Audio] { + assert!(!runtime_supports_modality(runtime, modality)); + } + } + } + + fn single_image_intermediate() -> MultimodalIntermediate { + use std::{collections::HashMap, sync::Arc}; + + use llm_multimodal::{ + EncoderFieldLayouts, ImageDetail, ImageFrame, ImageSource, PlaceholderRange, + PreprocessedEncoderInputs, + }; + use ndarray::{ArrayD, IxDyn}; + + use super::super::{MediaBatch, PrecomputedMultimodalIntermediate, PromptBinding}; + + MultimodalIntermediate::try_new(vec![PrecomputedMultimodalIntermediate { + preprocessed: PreprocessedEncoderInputs { + encoder_input: ArrayD::from_shape_vec(IxDyn(&[1, 1]), vec![1.0]).unwrap(), + feature_token_counts: vec![1], + item_sizes: vec![(1, 1)], + model_specific: HashMap::new(), + }, + media: MediaBatch::Images(vec![Arc::new(ImageFrame::new( + image::DynamicImage::new_rgb8(1, 1), + bytes::Bytes::from_static(b"image"), + ImageDetail::Auto, + ImageSource::InlineBytes, + "image-hash".to_string(), + ))]), + bindings: vec![PromptBinding { + item_index: 0, + prompt_ordinal: 0, + structural: PlaceholderRange { + offset: 0, + length: 1, + }, + patches: vec![], + }], + placeholder_token_id: Some(10), + field_layouts: EncoderFieldLayouts::default(), + keep_on_cpu_keys: vec![], + }]) + .unwrap() + } + + /// The early check accepts a supported (backend, modality) pair and rejects + /// an unsupported one with the single consistent message. + #[test] + fn early_check_gates_backend_modality() { + let intermediate = single_image_intermediate(); + + // Image on SGLang is supported. + assert!(ensure_backend_supports_modalities(RuntimeType::Sglang, &intermediate).is_ok()); + + // Image on MLX is not. + let err = ensure_backend_supports_modalities(RuntimeType::Mlx, &intermediate).unwrap_err(); + assert_eq!(err.to_string(), "backend mlx does not support image inputs"); + } +} diff --git a/model_gateway/src/routers/grpc/multimodal/mod.rs b/model_gateway/src/routers/grpc/multimodal/mod.rs index 38e591240..5aa2d8035 100644 --- a/model_gateway/src/routers/grpc/multimodal/mod.rs +++ b/model_gateway/src/routers/grpc/multimodal/mod.rs @@ -24,6 +24,7 @@ use llm_multimodal::{ }; mod assemble; +mod capability; mod config; mod detect; mod pixel_cache; @@ -36,6 +37,7 @@ pub(crate) use assemble::{ assemble_multimodal_data, assemble_multimodal_data_after_encode, assemble_tokenspeed_for_encode, encode_routing_hashes, }; +pub(crate) use capability::ensure_backend_supports_modalities; pub(crate) use config::{ load_preprocessor_config_file, load_video_preprocessor_config, MultimodalComponents, MultimodalConfigRegistry, MultimodalModelConfig,