diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs index 13ce48ae0..7e4e58a64 100644 --- a/bindings/python/src/lib.rs +++ b/bindings/python/src/lib.rs @@ -484,6 +484,8 @@ struct Router { epd_disaggregation: bool, encode_urls: Option)>>, encode_policy: Option, + multimodal_tensor_transport: Option, + multimodal_shm_min_bytes: Option, } impl Router { @@ -527,6 +529,23 @@ impl Router { DiscoveryConfig, MetricsConfig, PolicyConfig as ConfigPolicyConfig, RoutingMode, }; + // Validate the transport mode up front. The CLI (value_parser) and the + // argparse path (choices) already reject bad values; this covers direct + // programmatic `RouterArgs` use, matching the CLI/Rust parsing contract. + let multimodal_tensor_transport = self + .multimodal_tensor_transport + .as_deref() + .map(|value| { + config::TransportMode::parse(value).ok_or_else(|| { + config::ConfigError::InvalidValue { + field: "multimodal_tensor_transport".to_string(), + value: value.to_string(), + reason: "expected 'inline', 'shm', or 'auto'".to_string(), + } + }) + }) + .transpose()?; + let convert_policy = |policy: &PolicyType| -> config::ConfigResult { Ok(match policy { PolicyType::Random => ConfigPolicyConfig::Random, @@ -785,6 +804,8 @@ impl Router { .maybe_storage_hook_wasm_path(self.storage_hook_wasm_path.as_deref()) .enable_wasm(self.enable_wasm) .dp_aware(self.dp_aware) + .multimodal_tensor_transport(multimodal_tensor_transport) + .multimodal_shm_min_bytes(self.multimodal_shm_min_bytes) .routing_key_override(config::RoutingKeyOverrideConfig { enabled: self.routing_key_override, eviction_interval_secs: self.eviction_interval_secs, @@ -930,6 +951,8 @@ impl Router { epd_disaggregation = false, encode_urls = None, encode_policy = None, + multimodal_tensor_transport = None, + multimodal_shm_min_bytes = None, ))] #[expect(clippy::too_many_arguments)] #[expect( @@ -1054,6 +1077,8 @@ impl Router { epd_disaggregation: bool, encode_urls: Option)>>, encode_policy: Option, + multimodal_tensor_transport: Option, + multimodal_shm_min_bytes: Option, ) -> PyResult { let mut all_urls = worker_urls.clone(); @@ -1192,6 +1217,8 @@ impl Router { epd_disaggregation, encode_urls, encode_policy, + multimodal_tensor_transport, + multimodal_shm_min_bytes, }) } diff --git a/bindings/python/src/smg/router_args.py b/bindings/python/src/smg/router_args.py index 24659a3e2..213266ce8 100644 --- a/bindings/python/src/smg/router_args.py +++ b/bindings/python/src/smg/router_args.py @@ -72,6 +72,8 @@ class RouterArgs: max_payload_size: int = 512 * 1024 * 1024 # 512MB default for large batches bucket_adjust_interval_secs: int = 5 dp_aware: bool = False + multimodal_tensor_transport: str | None = None + multimodal_shm_min_bytes: int | None = None routing_key_override: bool = False dp_minimum_tokens_scheduler: bool = False enable_igw: bool = False # Enable IGW (Inter-Gateway) mode for multi-model support @@ -565,6 +567,21 @@ def add_cli_args( help="Interval in seconds between load monitor checks for PowerOfTwo routing (default: 10)", ) + # Multimodal tensor transport + parser.add_argument( + f"--{prefix}multimodal-tensor-transport", + type=str, + choices=["inline", "shm", "auto"], + default=RouterArgs.multimodal_tensor_transport, + help="Multimodal tensor transport mode: inline (default), shm, or auto", + ) + parser.add_argument( + f"--{prefix}multimodal-shm-min-bytes", + type=int, + default=RouterArgs.multimodal_shm_min_bytes, + help="Minimum multimodal tensor size (bytes) before the SHM transport is used", + ) + # Logging configuration logging_group.add_argument( f"--{prefix}log-dir", diff --git a/crates/protocols/src/worker.rs b/crates/protocols/src/worker.rs index 7715a0c8a..d13d42126 100644 --- a/crates/protocols/src/worker.rs +++ b/crates/protocols/src/worker.rs @@ -670,6 +670,17 @@ pub struct WorkerSpec { /// Falls back to the global `load_monitor_interval_secs` from router config. #[serde(default, skip_serializing_if = "Option::is_none")] pub load_monitor_interval_secs: Option, + + /// Per-worker multimodal tensor transport override (`inline` | `shm` | `auto`). + /// Overrides the router-level `multimodal_tensor_transport` for this worker + /// (e.g. force `shm` for a co-located worker, `inline` for a remote one). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub multimodal_tensor_transport: Option, + + /// Per-worker minimum multimodal tensor size (bytes) before the SHM transport + /// is used. Overrides the router-level `multimodal_shm_min_bytes`. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub multimodal_shm_min_bytes: Option, } impl WorkerSpec { @@ -700,10 +711,65 @@ impl WorkerSpec { resilience: ResilienceUpdate::default(), max_connection_attempts: default_max_connection_attempts(), load_monitor_interval_secs: None, + multimodal_tensor_transport: None, + multimodal_shm_min_bytes: None, + } + } +} + +/// Multimodal tensor transport mode for large payloads. +/// +/// - `Inline`: always carry tensor bytes in the gRPC message. +/// - `Shm`: use same-host `/dev/shm` when SMG can write it. +/// - `Auto`: use `/dev/shm` only when the receiving worker is verified to share +/// SMG's `/dev/shm`; otherwise fall back to inline. +#[derive( + Debug, Clone, Copy, PartialEq, Eq, Hash, Default, Serialize, Deserialize, schemars::JsonSchema, +)] +#[serde(rename_all = "lowercase")] +pub enum TransportMode { + #[default] + Inline, + Shm, + Auto, +} + +impl TransportMode { + /// Parse from a case-insensitive string (`inline` | `shm` | `auto`). + pub fn parse(value: &str) -> Option { + match value.trim().to_ascii_lowercase().as_str() { + "inline" => Some(Self::Inline), + "shm" => Some(Self::Shm), + "auto" => Some(Self::Auto), + _ => None, + } + } + + /// Canonical lowercase name. + pub fn as_str(self) -> &'static str { + match self { + Self::Inline => "inline", + Self::Shm => "shm", + Self::Auto => "auto", } } } +impl std::fmt::Display for TransportMode { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(self.as_str()) + } +} + +impl std::str::FromStr for TransportMode { + type Err = String; + + fn from_str(s: &str) -> Result { + Self::parse(s) + .ok_or_else(|| format!("invalid transport mode '{s}'; expected inline|shm|auto")) + } +} + // ── API types ─────────────────────────────────────────────────────── /// Worker information for API responses. diff --git a/docs/reference/configuration.md b/docs/reference/configuration.md index 67b9caef2..bd5648c0c 100644 --- a/docs/reference/configuration.md +++ b/docs/reference/configuration.md @@ -937,16 +937,28 @@ smg \ | `JWT_JWKS_URI` | `--jwt-jwks-uri` | JWKS URI | | `CONTROL_PLANE_API_KEYS` | `--control-plane-api-keys` | Control plane API keys | -### TokenSpeed Multimodal Tensor Transport +### Multimodal Tensor Transport -These env-only variables tune how the router ships preprocessed multimodal -tensors (image/video encoder inputs) to a TokenSpeed worker. They do not affect +Controls how the router ships preprocessed multimodal tensors (image/video +encoder inputs and model-specific tensors) to the worker. This does not affect accuracy — the inline and shared-memory paths produce byte-identical tensors. +Shared memory is currently used by the TokenSpeed backend. + +Resolution precedence (highest first): per-worker `WorkerSpec` override → +router config / CLI flag → `SMG_MM_*` environment variable → built-in default. + +| CLI Flag | Config / WorkerSpec field | Default | Description | +|----------|---------------------------|---------|-------------| +| `--multimodal-tensor-transport` | `multimodal_tensor_transport` | `inline` | Transport for large MM tensors: `inline` (gRPC bytes), `shm` (use `/dev/shm` whenever the router can write it — the operator asserts co-location), or `auto` (use `/dev/shm` only when the worker is *verified* to share it). In `auto`, the router compares the worker's advertised `/dev/shm` namespace token (`GetServerInfo`) to its own and uses SHM only on a match; otherwise it falls back to inline. | +| `--multimodal-shm-min-bytes` | `multimodal_shm_min_bytes` | `65536` | Minimum tensor size (bytes) before the SHM path is used; smaller tensors stay inline. | + +Both settings can be overridden per worker via the matching `WorkerSpec` fields +(e.g. force `shm` for a co-located worker and `inline` for a remote one). | Environment Variable | Default | Description | |---------------------|---------|-------------| -| `SMG_TOKENSPEED_MM_TENSOR_TRANSPORT` | `inline` | Transport for large MM tensors: `inline` (gRPC bytes), `shm` (always use `/dev/shm`), or `auto` (use `/dev/shm` only when the worker is *verified* to share it). In `auto`, the router compares the worker's advertised `/dev/shm` namespace token (`GetServerInfo`) to its own and uses SHM only on a match; otherwise it falls back to inline. No locality configuration is needed. | -| `SMG_TOKENSPEED_MM_SHM_MIN_BYTES` | `65536` | Minimum tensor size (bytes) before the SHM path is used; smaller tensors stay inline. | +| `SMG_MM_TENSOR_TRANSPORT` | `inline` | Env fallback for `--multimodal-tensor-transport`. Legacy alias: `SMG_TOKENSPEED_MM_TENSOR_TRANSPORT`. | +| `SMG_MM_SHM_MIN_BYTES` | `65536` | Env fallback for `--multimodal-shm-min-bytes`. Legacy alias: `SMG_TOKENSPEED_MM_SHM_MIN_BYTES`. | | `SMG_LOG_MM_TIMING` | `false` | Log per-stage multimodal preprocessing/assembly timing at `INFO`. Accepts `1`/`true`/`yes`. | The TokenSpeed gRPC servicer (worker side) reads two companion variables: diff --git a/model_gateway/src/config/builder.rs b/model_gateway/src/config/builder.rs index e019b54bb..6877961cc 100644 --- a/model_gateway/src/config/builder.rs +++ b/model_gateway/src/config/builder.rs @@ -1,5 +1,6 @@ use std::collections::HashMap; +use openai_protocol::worker::TransportMode; use smg_mcp::McpConfig; use super::{ @@ -213,6 +214,18 @@ impl RouterConfigBuilder { self } + /// Global multimodal tensor transport mode (per-worker specs can override). + pub fn multimodal_tensor_transport(mut self, mode: Option) -> Self { + self.config.multimodal_tensor_transport = mode; + self + } + + /// Global minimum multimodal tensor size (bytes) before SHM transport is used. + pub fn multimodal_shm_min_bytes(mut self, bytes: Option) -> Self { + self.config.multimodal_shm_min_bytes = bytes; + self + } + // ==================== Rate Limiting ==================== pub fn max_concurrent_requests(mut self, max: i32) -> Self { diff --git a/model_gateway/src/config/types.rs b/model_gateway/src/config/types.rs index a1979f9ff..44afa8f53 100644 --- a/model_gateway/src/config/types.rs +++ b/model_gateway/src/config/types.rs @@ -1,6 +1,7 @@ use std::collections::HashMap; use openai_protocol::worker::HealthCheckConfig as ProtocolHealthCheckConfig; +pub use openai_protocol::worker::TransportMode; use serde::{Deserialize, Serialize}; // Re-export storage config types from data_connector pub use smg_data_connector::{ @@ -43,6 +44,16 @@ pub struct RouterConfig { /// observability from routing. #[serde(default)] pub engine_metrics: bool, + /// Global multimodal tensor transport mode (`inline` | `shm` | `auto`). + /// Per-worker `WorkerSpec.multimodal_tensor_transport` overrides this; when + /// unset, falls back to `SMG_MM_TENSOR_TRANSPORT`, then `inline`. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub multimodal_tensor_transport: Option, + /// Global minimum multimodal tensor size (bytes) before SHM transport is used. + /// Per-worker `WorkerSpec.multimodal_shm_min_bytes` overrides this; falls back + /// to `SMG_MM_SHM_MIN_BYTES`, then 64 KiB. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub multimodal_shm_min_bytes: Option, pub dp_aware: bool, #[serde(default)] pub dp_minimum_tokens_scheduler: bool, @@ -743,6 +754,8 @@ impl Default for RouterConfig { worker_startup_check_interval_secs: 30, load_monitor_interval_secs: 10, engine_metrics: false, + multimodal_tensor_transport: None, + multimodal_shm_min_bytes: None, dp_aware: false, dp_minimum_tokens_scheduler: false, api_key: None, diff --git a/model_gateway/src/main.rs b/model_gateway/src/main.rs index cc4dbb934..2b7da4887 100644 --- a/model_gateway/src/main.rs +++ b/model_gateway/src/main.rs @@ -1,6 +1,7 @@ use std::collections::HashMap; use clap::{ArgAction, Parser, Subcommand, ValueEnum}; +use openai_protocol::worker::TransportMode; use rand::{distr::Alphanumeric, RngExt}; use smg::{ config::{ @@ -146,6 +147,12 @@ enum Commands { }, } +/// Parse the `--multimodal-tensor-transport` value into a `TransportMode`. +fn parse_transport_mode(value: &str) -> Result { + TransportMode::parse(value) + .ok_or_else(|| format!("invalid value '{value}'; expected inline, shm, or auto")) +} + #[derive(Parser, Debug)] struct CliArgs { // ==================== Worker Configuration ==================== @@ -308,6 +315,17 @@ struct CliArgs { #[arg(long, default_value_t = false, help_heading = "Load Monitoring")] engine_metrics: bool, + /// Multimodal tensor transport mode: `inline` (default), `shm` (same-host + /// /dev/shm), or `auto` (shm only when the worker shares /dev/shm). A + /// per-worker `WorkerSpec.multimodal_tensor_transport` overrides this. + #[arg(long, value_parser = parse_transport_mode, help_heading = "Multimodal")] + multimodal_tensor_transport: Option, + + /// Minimum multimodal tensor size (bytes) before the SHM transport is used. + /// Overridable per worker via `WorkerSpec.multimodal_shm_min_bytes`. + #[arg(long, help_heading = "Multimodal")] + multimodal_shm_min_bytes: Option, + // ==================== Service Discovery (Kubernetes) ==================== /// Enable Kubernetes service discovery #[arg( @@ -1338,6 +1356,8 @@ impl CliArgs { .worker_startup_check_interval_secs(self.worker_startup_check_interval) .load_monitor_interval_secs(self.load_monitor_interval) .engine_metrics(self.engine_metrics) + .multimodal_tensor_transport(self.multimodal_tensor_transport) + .multimodal_shm_min_bytes(self.multimodal_shm_min_bytes) .max_concurrent_requests(self.max_concurrent_requests) .queue_size(self.queue_size) .queue_timeout_secs(self.queue_timeout_secs) @@ -1708,6 +1728,37 @@ mod tests { ); } + /// The multimodal transport flags must reach both `RouterConfig` and the + /// wrapped `ServerConfig.router_config`. Two-path config-plumbing guard. + #[test] + fn multimodal_transport_flows_into_both_configs() { + let cli = cli_args_from(&[ + "--multimodal-tensor-transport", + "shm", + "--multimodal-shm-min-bytes", + "1024", + ]); + + let router_config = cli.to_router_config(vec![], vec![]).unwrap(); + assert_eq!( + router_config.multimodal_tensor_transport, + Some(TransportMode::Shm), + "transport mode must reach RouterConfig via to_router_config" + ); + assert_eq!(router_config.multimodal_shm_min_bytes, Some(1024)); + + let server_config = cli.to_server_config(router_config).unwrap(); + assert_eq!( + server_config.router_config.multimodal_tensor_transport, + Some(TransportMode::Shm), + "transport mode must survive into ServerConfig via to_server_config" + ); + assert_eq!( + server_config.router_config.multimodal_shm_min_bytes, + Some(1024) + ); + } + /// Default is off: the flag stays false through both conversions so /// existing deployments keep the routing-gated polling behavior. #[test] diff --git a/model_gateway/src/routers/grpc/epd_encode.rs b/model_gateway/src/routers/grpc/epd_encode.rs index 428d9ffe5..2578b8235 100644 --- a/model_gateway/src/routers/grpc/epd_encode.rs +++ b/model_gateway/src/routers/grpc/epd_encode.rs @@ -77,15 +77,17 @@ pub(crate) enum PreparedEncodeItem { TokenSpeed { item: Option, shm_enabled: bool, + shm_min_bytes: usize, cleanup_on_drop: bool, }, } impl PreparedEncodeItem { - fn tokenspeed(item: TokenSpeedMultimodalItem, shm_enabled: bool) -> Self { + fn tokenspeed(item: TokenSpeedMultimodalItem, shm_enabled: bool, shm_min_bytes: usize) -> Self { Self::TokenSpeed { item: Some(item), shm_enabled, + shm_min_bytes, cleanup_on_drop: true, } } @@ -99,6 +101,7 @@ impl PreparedEncodeItem { Self::TokenSpeed { item, shm_enabled, + shm_min_bytes, cleanup_on_drop, } => { let item = item @@ -111,6 +114,7 @@ impl PreparedEncodeItem { TokenSpeedMultimodalData { items: vec![item], shm_enabled: *shm_enabled, + shm_min_bytes: *shm_min_bytes, } .into_proto(), ), @@ -254,10 +258,11 @@ fn prepare_tokenspeed_items( ) -> Result> { let tokenspeed_mm = assemble_tokenspeed(precomputed, workers, false)?; let shm_enabled = tokenspeed_mm.shm_enabled; + let shm_min_bytes = tokenspeed_mm.shm_min_bytes; Ok(tokenspeed_mm .items .into_iter() - .map(|item| PreparedEncodeItem::tokenspeed(item, shm_enabled)) + .map(|item| PreparedEncodeItem::tokenspeed(item, shm_enabled, shm_min_bytes)) .collect()) } diff --git a/model_gateway/src/routers/grpc/multimodal/assemble.rs b/model_gateway/src/routers/grpc/multimodal/assemble.rs index db6926b7b..107108ebe 100644 --- a/model_gateway/src/routers/grpc/multimodal/assemble.rs +++ b/model_gateway/src/routers/grpc/multimodal/assemble.rs @@ -19,7 +19,7 @@ use super::{ model_specific_to_tensor_bytes, serialize_array_as_tokenspeed_tensor, serialize_encoder_input, serialize_model_specific, slice_array_axis0, }, - transport::{resolve_tokenspeed_shm_enabled, tokenspeed_encoder_input_dtype}, + transport::{mm_encoder_input_dtype, resolve_mm_shm_enabled, resolve_mm_shm_min_bytes}, MultimodalIntermediate, PrecomputedMultimodalIntermediate, }; use crate::routers::grpc::{ @@ -213,6 +213,7 @@ pub(crate) fn assemble_tokenspeed( struct TokenSpeedAssemblyOptions { shm_enabled: bool, + shm_min_bytes: usize, encoder_input_dtype: String, skip_pixel_values: bool, } @@ -223,8 +224,9 @@ fn tokenspeed_assembly_options( skip_pixel_values: bool, ) -> TokenSpeedAssemblyOptions { TokenSpeedAssemblyOptions { - shm_enabled: resolve_tokenspeed_shm_enabled(workers, skip_pixel_values), - encoder_input_dtype: tokenspeed_encoder_input_dtype(modality, workers), + shm_enabled: resolve_mm_shm_enabled(workers, skip_pixel_values), + shm_min_bytes: resolve_mm_shm_min_bytes(workers), + encoder_input_dtype: mm_encoder_input_dtype(modality, workers), skip_pixel_values, } } @@ -237,6 +239,7 @@ fn assemble_tokenspeed_with_options( let total_started = Instant::now(); let TokenSpeedAssemblyOptions { shm_enabled, + shm_min_bytes, encoder_input_dtype, skip_pixel_values, } = options; @@ -284,6 +287,7 @@ fn assemble_tokenspeed_with_options( &item_encoder_input, &encoder_input_dtype, shm_enabled, + shm_min_bytes, ) }; let encoder_input_serialize_ms = encoder_input_started.elapsed().as_secs_f64() * 1000.0; @@ -339,7 +343,11 @@ fn assemble_tokenspeed_with_options( ); } - Ok(TokenSpeedMultimodalData { items, shm_enabled }) + Ok(TokenSpeedMultimodalData { + items, + shm_enabled, + shm_min_bytes, + }) } fn precomputed_multimodal_item_count( @@ -537,6 +545,7 @@ mod tests { content_hash: vec![], }], shm_enabled: true, + shm_min_bytes: 0, }; (PendingTokenSpeedAssembly::new(data), path) } diff --git a/model_gateway/src/routers/grpc/multimodal/mod.rs b/model_gateway/src/routers/grpc/multimodal/mod.rs index ebd798498..e3f69c5fb 100644 --- a/model_gateway/src/routers/grpc/multimodal/mod.rs +++ b/model_gateway/src/routers/grpc/multimodal/mod.rs @@ -41,6 +41,7 @@ pub(crate) use detect::{chat_modalities, has_multimodal_content_messages}; pub(crate) use process::{ process_multimodal, process_multimodal_messages, resolve_placeholder_token, }; +pub(crate) use transport::init_mm_transport_defaults; /// Whether verbose multimodal timing logs are enabled via `SMG_LOG_MM_TIMING`. /// Read from the environment once and cached; the flag is not expected to change diff --git a/model_gateway/src/routers/grpc/multimodal/serialize.rs b/model_gateway/src/routers/grpc/multimodal/serialize.rs index cbcd37d17..2c6c744e7 100644 --- a/model_gateway/src/routers/grpc/multimodal/serialize.rs +++ b/model_gateway/src/routers/grpc/multimodal/serialize.rs @@ -11,7 +11,7 @@ use tracing::{info, warn}; use super::log_mm_timing_enabled; use crate::routers::grpc::proto_wrapper::{ - tokenspeed_mm_shm_min_bytes, write_tokenspeed_shm_with, TensorBytes, TokenSpeedTensor, + write_tokenspeed_shm_with, TensorBytes, TokenSpeedTensor, }; /// Serialize the primary encoder input ndarray to raw little-endian f32 bytes + shape. @@ -55,6 +55,7 @@ pub(super) fn serialize_array_as_tokenspeed_tensor( encoder_input: &ArrayViewD<'_, f32>, dtype: &str, shm_enabled: bool, + shm_min_bytes: usize, ) -> TokenSpeedTensor { let dtype = match canonical_float_dtype(dtype).as_deref() { Some("float32") => "float32".to_string(), @@ -76,7 +77,7 @@ pub(super) fn serialize_array_as_tokenspeed_tensor( }; let nbytes = encoder_input.len() * element_size; - if shm_enabled && nbytes >= tokenspeed_mm_shm_min_bytes() { + if shm_enabled && nbytes >= shm_min_bytes { let started = Instant::now(); match write_tokenspeed_shm_with(nbytes, |output| { fill_array_as_dtype(output, encoder_input, &dtype) @@ -461,4 +462,33 @@ mod tests { assert_eq!(direct, expected); assert_eq!(serialize_array(&fortran_item), (expected, vec![1, 2])); } + + #[test] + #[cfg(target_os = "linux")] + fn shm_min_bytes_threshold_gates_inline_vs_shm() { + use std::path::Path; + + use crate::routers::grpc::proto_wrapper::TokenSpeedTensorStorage; + + // float32 = 4 bytes/elem; threshold of 16 bytes falls at 4 elements. + let below = ArrayD::from_shape_vec(IxDyn(&[3]), vec![1.0_f32, 2.0, 3.0]).unwrap(); + let at = ArrayD::from_shape_vec(IxDyn(&[4]), vec![1.0_f32, 2.0, 3.0, 4.0]).unwrap(); + + // Below the threshold: stays inline even with SHM enabled. + let tensor = serialize_array_as_tokenspeed_tensor(&below.view(), "float32", true, 16); + assert!(matches!(tensor.storage, TokenSpeedTensorStorage::Inline(_))); + + // At/above the threshold: uses SHM (/dev/shm is writable on Linux CI). + let tensor = serialize_array_as_tokenspeed_tensor(&at.view(), "float32", true, 16); + match tensor.storage { + TokenSpeedTensorStorage::Shm(handle) => { + let path = Path::new("/dev/shm").join(&handle.name); + assert!(path.exists()); + let _ = std::fs::remove_file(&path); + } + TokenSpeedTensorStorage::Inline(_) => { + panic!("expected SHM at/above the threshold when /dev/shm is writable") + } + } + } } diff --git a/model_gateway/src/routers/grpc/multimodal/transport.rs b/model_gateway/src/routers/grpc/multimodal/transport.rs index c0b0c6448..fdf923aa9 100644 --- a/model_gateway/src/routers/grpc/multimodal/transport.rs +++ b/model_gateway/src/routers/grpc/multimodal/transport.rs @@ -1,31 +1,194 @@ //! Multimodal tensor transport resolution. //! //! Decides how large multimodal tensors travel from the gateway to the worker: -//! the SHM-vs-inline transport mode, the encoder-input wire dtype, and the -//! `/dev/shm` namespace verification that makes the SHM path safe. All values -//! resolve from environment + worker labels today. +//! the SHM-vs-inline transport mode, its size threshold, the encoder-input wire +//! dtype, and the `/dev/shm` namespace verification that makes the SHM path safe. +//! +//! Resolution precedence for the transport mode and SHM threshold: +//! per-worker `WorkerSpec` override → router config (seeded once at startup via +//! [`init_mm_transport_defaults`]) → `SMG_MM_*` env (with the legacy +//! `SMG_TOKENSPEED_MM_*` names as a fallback) → built-in default (`inline`, +//! 64 KiB). use std::sync::{Arc, OnceLock}; use llm_multimodal::Modality; +use openai_protocol::worker::TransportMode; use tracing::{info, warn}; -use crate::routers::grpc::{ - context::WorkerSelection, - proto_wrapper::{ - tokenspeed_mm_shm_min_bytes, tokenspeed_mm_tensor_transport_mode, - tokenspeed_shm_dev_writable, - }, -}; +use crate::routers::grpc::{context::WorkerSelection, proto_wrapper::mm_shm_dev_writable}; + +const DEFAULT_SHM_MIN_BYTES: usize = 64 * 1024; + +/// Router-level transport defaults, resolved once at startup from `RouterConfig` +/// (falling back to env, then built-in defaults). Per-worker `WorkerSpec` +/// overrides take precedence over these at request time. +#[derive(Debug, Clone, Copy)] +struct MmTransportDefaults { + mode: TransportMode, + shm_min_bytes: usize, +} + +static DEFAULTS: OnceLock = OnceLock::new(); + +/// Seed the process-wide transport defaults from router config. Config values +/// win; unset values fall back to env, then the built-in defaults. Call once at +/// startup before serving; idempotent (first call wins). +pub(crate) fn init_mm_transport_defaults( + mode: Option, + shm_min_bytes: Option, +) { + let resolved = MmTransportDefaults { + mode: mode + .or_else(mm_tensor_transport_mode_from_env) + .unwrap_or_default(), + shm_min_bytes: shm_min_bytes + .or_else(mm_shm_min_bytes_from_env) + .unwrap_or(DEFAULT_SHM_MIN_BYTES), + }; + let _ = DEFAULTS.set(resolved); + log_transport_config_once(resolved); +} + +/// The resolved router-level defaults. If [`init_mm_transport_defaults`] was +/// never called (e.g. in tests), resolve lazily from env + built-in defaults. +fn mm_transport_defaults() -> MmTransportDefaults { + if let Some(defaults) = DEFAULTS.get() { + return *defaults; + } + MmTransportDefaults { + mode: mm_tensor_transport_mode_from_env().unwrap_or_default(), + shm_min_bytes: mm_shm_min_bytes_from_env().unwrap_or(DEFAULT_SHM_MIN_BYTES), + } +} + +fn mm_tensor_transport_mode_from_env() -> Option { + static LEGACY_WARNED: OnceLock<()> = OnceLock::new(); + let raw = env_with_deprecated_alias( + "SMG_MM_TENSOR_TRANSPORT", + "SMG_TOKENSPEED_MM_TENSOR_TRANSPORT", + &LEGACY_WARNED, + )?; + match TransportMode::parse(&raw) { + Some(mode) => Some(mode), + None => { + log_unknown_transport_once(&raw); + None + } + } +} + +fn mm_shm_min_bytes_from_env() -> Option { + static LEGACY_WARNED: OnceLock<()> = OnceLock::new(); + let raw = env_with_deprecated_alias( + "SMG_MM_SHM_MIN_BYTES", + "SMG_TOKENSPEED_MM_SHM_MIN_BYTES", + &LEGACY_WARNED, + )?; + match raw.parse::() { + Ok(value) => Some(value), + Err(_) => { + log_invalid_shm_min_bytes_once(&raw); + None + } + } +} + +/// Read the canonical env var, falling back to the deprecated alias. When the +/// value comes from the alias, log a one-time migration warning (guarded by +/// `warned`, one warning per variable). +fn env_with_deprecated_alias( + canonical: &str, + deprecated: &str, + warned: &OnceLock<()>, +) -> Option { + if let Some(value) = read_env_nonempty(canonical) { + return Some(value); + } + let value = read_env_nonempty(deprecated)?; + warned.get_or_init(|| { + warn!( + deprecated, + canonical, + "Deprecated multimodal transport env var is set; migrate to the canonical name" + ); + }); + Some(value) +} + +fn read_env_nonempty(name: &str) -> Option { + std::env::var(name) + .ok() + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()) +} + +/// Resolve whether large multimodal tensors should use the SHM transport for +/// this request: per-worker override → router default. `shm` forces SHM whenever +/// SMG can write `/dev/shm` (the operator asserts co-location); `auto` also +/// requires the receiving worker leg to be verified as sharing SMG's `/dev/shm`; +/// `inline` (the default) keeps the gRPC path. +pub(super) fn resolve_mm_shm_enabled( + workers: Option<&WorkerSelection>, + skip_pixel_values: bool, +) -> bool { + let mode = + worker_transport_mode_override(workers).unwrap_or_else(|| mm_transport_defaults().mode); + match mode { + TransportMode::Shm => mm_shm_dev_writable(), + TransportMode::Auto => { + worker_shares_dev_shm(workers, skip_pixel_values) && mm_shm_dev_writable() + } + TransportMode::Inline => false, + } +} -pub(super) fn tokenspeed_encoder_input_dtype( +/// Resolve the SHM size threshold (bytes) for this request: per-worker override +/// → router default. +pub(super) fn resolve_mm_shm_min_bytes(workers: Option<&WorkerSelection>) -> usize { + worker_shm_min_bytes_override(workers).unwrap_or_else(|| mm_transport_defaults().shm_min_bytes) +} + +fn worker_transport_mode_override(workers: Option<&WorkerSelection>) -> Option { + primary_worker(workers)? + .metadata() + .spec + .multimodal_tensor_transport +} + +fn worker_shm_min_bytes_override(workers: Option<&WorkerSelection>) -> Option { + primary_worker(workers)? + .metadata() + .spec + .multimodal_shm_min_bytes +} + +/// The worker whose per-worker overrides apply. Multimodal tensors are sent to +/// wherever the vision encoder runs: the encode worker in EPD (so its spec wins), +/// otherwise the single/prefill worker that does the encoding itself. +fn primary_worker(workers: Option<&WorkerSelection>) -> Option<&Arc> { + match workers? { + WorkerSelection::Single { worker } => Some(worker), + WorkerSelection::Disaggregated { + encode_assignments, + prefill, + .. + } => encode_assignments + .as_ref() + .and_then(|assignments| assignments.first()) + .map(|assignment| &assignment.worker) + .or(Some(prefill)), + } +} + +pub(super) fn mm_encoder_input_dtype( modality: Modality, workers: Option<&WorkerSelection>, ) -> String { - if let Some(dtype) = tokenspeed_encoder_input_dtype_from_env(modality) { + if let Some(dtype) = mm_encoder_input_dtype_from_env(modality) { return dtype; } - if let Some(dtype) = tokenspeed_encoder_input_dtype_from_worker(workers) { + if let Some(dtype) = mm_encoder_input_dtype_from_worker(workers) { return dtype; } // Default to bf16 on the wire: the engine casts encoder_input to the model @@ -35,7 +198,7 @@ pub(super) fn tokenspeed_encoder_input_dtype( "bfloat16".to_string() } -fn tokenspeed_encoder_input_dtype_from_env(modality: Modality) -> Option { +fn mm_encoder_input_dtype_from_env(modality: Modality) -> Option { static IMAGE_DTYPE: OnceLock> = OnceLock::new(); static VIDEO_DTYPE: OnceLock> = OnceLock::new(); static AUDIO_DTYPE: OnceLock> = OnceLock::new(); @@ -61,12 +224,8 @@ fn cached_env_dtype(cell: &'static OnceLock>, name: &str) -> Opti .clone() } -fn tokenspeed_encoder_input_dtype_from_worker(workers: Option<&WorkerSelection>) -> Option { - let worker = match workers? { - WorkerSelection::Single { worker } => worker, - WorkerSelection::Disaggregated { prefill, .. } => prefill, - }; - worker +fn mm_encoder_input_dtype_from_worker(workers: Option<&WorkerSelection>) -> Option { + primary_worker(workers)? .metadata() .spec .labels @@ -75,47 +234,34 @@ fn tokenspeed_encoder_input_dtype_from_worker(workers: Option<&WorkerSelection>) .cloned() } -/// Resolve whether large multimodal tensors should use the SHM transport for -/// this request. `shm` and `auto` require the receiving worker leg to share -/// SMG's `/dev/shm`; anything else (including unset or `inline`) keeps the -/// inline gRPC path. -pub(super) fn resolve_tokenspeed_shm_enabled( - workers: Option<&WorkerSelection>, - skip_pixel_values: bool, -) -> bool { - let mode = tokenspeed_mm_tensor_transport_mode(); - log_tokenspeed_transport_config_once(&mode); - match mode.as_str() { - // SHM only ever happens when SMG can actually write /dev/shm. - "shm" | "auto" => { - worker_shares_dev_shm(workers, skip_pixel_values) && tokenspeed_shm_dev_writable() - } - "" | "inline" => false, - other => { - log_unknown_tokenspeed_transport_once(other); - false - } - } -} - -fn log_tokenspeed_transport_config_once(mode: &str) { +fn log_transport_config_once(defaults: MmTransportDefaults) { static LOGGED: OnceLock<()> = OnceLock::new(); LOGGED.get_or_init(|| { info!( - mode, - shm_min_bytes = tokenspeed_mm_shm_min_bytes(), - dev_writable = tokenspeed_shm_dev_writable(), - "TokenSpeed multimodal tensor transport configured" + mode = %defaults.mode, + shm_min_bytes = defaults.shm_min_bytes, + dev_writable = mm_shm_dev_writable(), + "Multimodal tensor transport configured" + ); + }); +} + +fn log_unknown_transport_once(value: &str) { + static WARNED: OnceLock<()> = OnceLock::new(); + WARNED.get_or_init(|| { + warn!( + value, + "Unknown multimodal tensor transport value; expected inline|shm|auto, using inline" ); }); } -fn log_unknown_tokenspeed_transport_once(value: &str) { +fn log_invalid_shm_min_bytes_once(value: &str) { static WARNED: OnceLock<()> = OnceLock::new(); WARNED.get_or_init(|| { warn!( value, - "Unknown SMG_TOKENSPEED_MM_TENSOR_TRANSPORT value; expected inline|shm|auto, using inline" + "Invalid multimodal SHM min-bytes value; expected a non-negative integer, using default" ); }); } diff --git a/model_gateway/src/routers/grpc/proto_wrapper.rs b/model_gateway/src/routers/grpc/proto_wrapper.rs index a9c4201d0..e4a1e251d 100644 --- a/model_gateway/src/routers/grpc/proto_wrapper.rs +++ b/model_gateway/src/routers/grpc/proto_wrapper.rs @@ -109,6 +109,10 @@ pub struct TokenSpeedMultimodalData { /// transport? Computed upstream from the transport mode and (for `auto`) /// worker locality, so `into_proto` does not re-read the environment. pub shm_enabled: bool, + /// Resolved per-request SHM size threshold (bytes): tensors smaller than this + /// stay inline even when `shm_enabled`. Computed upstream (worker override → + /// router config → env → default) so `into_proto` does not re-read the env. + pub shm_min_bytes: usize, } #[derive(Debug)] @@ -276,17 +280,18 @@ impl TokenSpeedMultimodalData { /// each item's encoder_input afterward via `clear_mm_pixel_values`. pub fn into_proto(self) -> tokenspeed::MultimodalInputs { let shm_enabled = self.shm_enabled; + let shm_min_bytes = self.shm_min_bytes; let items = self .items .into_iter() - .map(|item| item.into_proto(shm_enabled)) + .map(|item| item.into_proto(shm_enabled, shm_min_bytes)) .collect(); tokenspeed::MultimodalInputs { items } } } impl TokenSpeedMultimodalItem { - fn into_proto(self, shm_enabled: bool) -> tokenspeed::MultimodalItem { + fn into_proto(self, shm_enabled: bool, shm_min_bytes: usize) -> tokenspeed::MultimodalItem { let placeholders = self .mm_placeholders .into_iter() @@ -296,10 +301,14 @@ impl TokenSpeedMultimodalItem { let model_specific_tensors = self .model_specific_tensors .into_iter() - .map(|(k, v)| (k, tensor_bytes_to_tokenspeed(v, shm_enabled))) + .map(|(k, v)| (k, tensor_bytes_to_tokenspeed(v, shm_enabled, shm_min_bytes))) .collect::>(); - let encoder_input = Some(tokenspeed_tensor_to_proto(self.encoder_input, shm_enabled)); + let encoder_input = Some(tokenspeed_tensor_to_proto( + self.encoder_input, + shm_enabled, + shm_min_bytes, + )); tokenspeed::MultimodalItem { modality: match self.modality { @@ -319,6 +328,7 @@ impl TokenSpeedMultimodalItem { fn tokenspeed_tensor_to_proto( value: TokenSpeedTensor, shm_enabled: bool, + shm_min_bytes: usize, ) -> tokenspeed::TensorData { use crate::observability::metrics::Metrics; let TokenSpeedTensor { @@ -328,7 +338,9 @@ fn tokenspeed_tensor_to_proto( } = value; let payload = match storage { // Inline storage is metered inside tokenspeed_tensor_payload. - TokenSpeedTensorStorage::Inline(data) => tokenspeed_tensor_payload(data, shm_enabled), + TokenSpeedTensorStorage::Inline(data) => { + tokenspeed_tensor_payload(data, shm_enabled, shm_min_bytes) + } // Encoder input already written directly to SHM upstream — meter it here. TokenSpeedTensorStorage::Shm(handle) => { Metrics::record_mm_tensor("tokenspeed", "shm", handle.nbytes as usize); @@ -343,17 +355,25 @@ fn tokenspeed_tensor_to_proto( } } -fn tensor_bytes_to_tokenspeed(value: TensorBytes, shm_enabled: bool) -> tokenspeed::TensorData { +fn tensor_bytes_to_tokenspeed( + value: TensorBytes, + shm_enabled: bool, + shm_min_bytes: usize, +) -> tokenspeed::TensorData { let TensorBytes { data, shape, dtype } = value; tokenspeed::TensorData { shape, dtype, - payload: Some(tokenspeed_tensor_payload(data, shm_enabled)), + payload: Some(tokenspeed_tensor_payload(data, shm_enabled, shm_min_bytes)), } } -fn tokenspeed_tensor_payload(data: Vec, shm_enabled: bool) -> tokenspeed::tensor_data::Payload { +fn tokenspeed_tensor_payload( + data: Vec, + shm_enabled: bool, + min_bytes: usize, +) -> tokenspeed::tensor_data::Payload { use crate::observability::metrics::Metrics; let log_timing = log_tokenspeed_mm_timing_enabled(); let nbytes = data.len(); @@ -365,7 +385,6 @@ fn tokenspeed_tensor_payload(data: Vec, shm_enabled: bool) -> tokenspeed::te return tokenspeed::tensor_data::Payload::Inline(data); } - let min_bytes = tokenspeed_mm_shm_min_bytes(); if nbytes < min_bytes { if log_timing { tracing::info!( @@ -410,27 +429,6 @@ fn log_tokenspeed_mm_timing_enabled() -> bool { .unwrap_or(false) } -/// Multimodal tensor transport mode for the TokenSpeed backend. -/// -/// This only governs multimodal tensor payloads (encoder inputs and -/// model-specific tensors); prompt `input_ids` are always sent inline. Set via -/// `SMG_TOKENSPEED_MM_TENSOR_TRANSPORT`. -pub fn tokenspeed_mm_tensor_transport_mode() -> String { - std::env::var("SMG_TOKENSPEED_MM_TENSOR_TRANSPORT") - .unwrap_or_default() - .trim() - .to_ascii_lowercase() -} - -/// Minimum multimodal tensor size (bytes) before the SHM transport is used. -/// Set via `SMG_TOKENSPEED_MM_SHM_MIN_BYTES`. Defaults to 64 KiB. -pub fn tokenspeed_mm_shm_min_bytes() -> usize { - std::env::var("SMG_TOKENSPEED_MM_SHM_MIN_BYTES") - .ok() - .and_then(|value| value.parse::().ok()) - .unwrap_or(64 * 1024) -} - static TOKENSPEED_SHM_COUNTER: AtomicU64 = AtomicU64::new(0); fn write_tokenspeed_shm(data: &[u8]) -> std::io::Result { @@ -442,7 +440,7 @@ fn write_tokenspeed_shm(data: &[u8]) -> std::io::Result { /// Whether SMG can actually create+write files under `/dev/shm`. Probed once; /// when false the SHM transport cannot work, so `auto`/`shm` must stay inline. -pub fn tokenspeed_shm_dev_writable() -> bool { +pub fn mm_shm_dev_writable() -> bool { static WRITABLE: OnceLock = OnceLock::new(); *WRITABLE.get_or_init(|| { let name = format!("smg-tokenspeed-probe-{}", process::id()); @@ -1863,6 +1861,7 @@ mod tests { content_hash: vec![7; 32], }], shm_enabled: false, + shm_min_bytes: 0, } .into_proto(); @@ -1905,6 +1904,7 @@ mod tests { content_hash: vec![7; 32], }], shm_enabled: false, + shm_min_bytes: 0, } .into_proto(); @@ -1930,6 +1930,7 @@ mod tests { dtype: "uint32".to_string(), }, false, + 0, ); assert_eq!( @@ -1963,6 +1964,7 @@ mod tests { content_hash: vec![7; 32], }], shm_enabled: true, + shm_min_bytes: 0, } .into_proto(); diff --git a/model_gateway/src/server.rs b/model_gateway/src/server.rs index 94f38d369..863fd111e 100644 --- a/model_gateway/src/server.rs +++ b/model_gateway/src/server.rs @@ -1002,6 +1002,14 @@ pub async fn startup(config: ServerConfig) -> Result<(), Box