From 978cc228caabaaaab19fa5ca88f50d6170cb4326 Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Fri, 4 Sep 2026 05:22:44 -0700 Subject: [PATCH] [Rust] Bound multimodal media ingress (#37967) --- python/sglang/srt/rust_server/config.py | 2 + rust/sglang-mm/README.md | 8 +- rust/sglang-mm/src/common/fetch.rs | 140 ++++++++++++---- rust/sglang-mm/src/driver.rs | 37 ++--- .../src/api_server/native_api.rs | 5 +- rust/sglang-server/src/api_server/prefetch.rs | 149 +++++++++++++++--- rust/sglang-server/src/message/config.rs | 7 + 7 files changed, 274 insertions(+), 74 deletions(-) diff --git a/python/sglang/srt/rust_server/config.py b/python/sglang/srt/rust_server/config.py index 8ed8fd42d..06f622d7e 100644 --- a/python/sglang/srt/rust_server/config.py +++ b/python/sglang/srt/rust_server/config.py @@ -11,6 +11,7 @@ from sglang.srt.arg_groups.overrides import resolving_view from sglang.srt.managers.utils import compute_num_reserved_tokens from sglang.srt.runtime_context import ( get_disagg, + get_mm, get_model, get_observability, get_serving, @@ -81,6 +82,7 @@ def _build_server_args(scheduler: Scheduler) -> ServerArgs: if get_serving().preferred_sampling_params is not None else None ), + limit_mm_data_per_request=get_mm().limit_mm_data_per_request or {}, allow_auto_truncate=get_serving().allow_auto_truncate, enable_return_hidden_states=sa.enable_return_hidden_states, # Not a `server_args` field: `TokenizerManager` derives it, and the diff --git a/rust/sglang-mm/README.md b/rust/sglang-mm/README.md index a5bd958ec..6cd884627 100644 --- a/rust/sglang-mm/README.md +++ b/rust/sglang-mm/README.md @@ -107,10 +107,10 @@ Adding one = a `MmFamilyProcessor` impl in `src//mod.rs` plus a `common::fetch` matches the Python `get_image_bytes` semantics (`REQUEST_TIMEOUT` env, `HTTP(S)_PROXY` / `ALL_PROXY` / `NO_PROXY` including -IPv4-CIDR and `host:port` entries) with two deliberate differences: every -source form is capped at 64 MiB — plus 64 items / 256 MiB per request in the -driver — and `file://` URLs actually work (the Python helper passes the -un-stripped URL to `open()`). +IPv4-CIDR and `host:port` entries) with deliberate safety bounds: every media +source contributes to a shared 64-item / 1.25 GiB request budget, and each +remote I/O stream has an additional 64 MiB cap. `file://` URLs also work (the +Python helper passes the un-stripped URL to `open()`). ## Python API diff --git a/rust/sglang-mm/src/common/fetch.rs b/rust/sglang-mm/src/common/fetch.rs index a5d8ceb3b..a30564e7f 100644 --- a/rust/sglang-mm/src/common/fetch.rs +++ b/rust/sglang-mm/src/common/fetch.rs @@ -10,9 +10,8 @@ use std::sync::OnceLock; use base64::Engine; -/// Cap on any single resolved payload — HTTP, file, or base64 — so no source -/// form can exhaust memory (the Python path has no such cap; oversized -/// payloads reject the request here). +/// Cap on one remotely fetched payload. Inline base64 and trusted local files +/// use their caller's whole-request budget instead. pub const MAX_FETCH_BYTES: u64 = 64 << 20; /// Charge granularity of a streaming read: the most an in-flight source can @@ -39,6 +38,17 @@ impl ByteBudget { .map_err(|_| ()) } + fn remaining(&self) -> u64 { + self.0.load(std::sync::atomic::Ordering::Acquire) + } + + /// Charge bytes which were already materialized by an earlier pipeline + /// stage. This lets the consumer apply one whole-request bound across + /// prefetched I/O and inline payloads without reading the source twice. + pub fn charge_existing(&self, n: usize, what: &str) -> Result<(), String> { + self.claim(n as u64).map_err(|()| over_budget(what)) + } + /// Give back bytes claimed for a chunk but not filled by the read. fn release(&self, n: u64) { self.0.fetch_add(n, std::sync::atomic::Ordering::AcqRel); @@ -52,8 +62,40 @@ pub fn fetch_bytes(src: &str) -> Result, String> { fetch_bytes_budgeted(src, &ByteBudget::new(MAX_FETCH_BYTES)) } +/// Read a trusted local media path without applying the per-source remote cap. +/// +/// Python's media security limit is specifically a URL-download limit. Local +/// video fixtures and mounted production assets are commonly larger than 64 +/// MiB, so applying [`MAX_FETCH_BYTES`] to them breaks requests which the Python +/// frontend accepts. They still consume the caller's whole-request budget. +/// Reject non-regular files and charge their size before reading so a request +/// cannot turn devices or a huge sparse file into an unbounded allocation. +pub fn fetch_local_file_budgeted(src: &str, budget: &ByteBudget) -> Result, String> { + let path = src.strip_prefix("file://").unwrap_or(src); + let file = std::fs::File::open(path).map_err(|e| format!("media fetch: {path}: {e}"))?; + let metadata = file + .metadata() + .map_err(|e| format!("media fetch: stat {path}: {e}"))?; + if !metadata.is_file() { + return Err(format!("media fetch: {path}: not a regular file")); + } + let expected = metadata.len(); + budget.claim(expected).map_err(|()| over_budget(path))?; + let mut buf = Vec::with_capacity(usize::try_from(expected).unwrap_or(usize::MAX)); + let read = file + .take(expected.saturating_add(1)) + .read_to_end(&mut buf) + .map_err(|e| format!("media fetch: read {path}: {e}"))? as u64; + if read > expected { + return Err(format!("media fetch: {path}: changed size while reading")); + } + budget.release(expected - read); + Ok(buf) +} + /// [`fetch_bytes`] against a caller-owned allowance, for resolving several -/// sources under one whole-request bound. [`MAX_FETCH_BYTES`] still caps each. +/// sources under one whole-request bound. [`MAX_FETCH_BYTES`] still caps I/O +/// streams; already-resident base64 is bounded by `budget`. pub fn fetch_bytes_budgeted(src: &str, budget: &ByteBudget) -> Result, String> { if src.starts_with("http://") || src.starts_with("https://") { return http_get(src, budget); @@ -69,19 +111,47 @@ pub fn fetch_bytes_budgeted(src: &str, budget: &ByteBudget) -> Result, S .split_once(',') .ok_or_else(|| "media fetch: malformed data: URL".to_string())? .1; - return charge_decoded(b64(encoded)?, budget); + return decode_base64_budgeted(encoded, budget); } // Python treats any other string as bare base64. - charge_decoded(b64(src)?, budget) + decode_base64_budgeted(src, budget) } -/// Base64 payloads are already resident in the request body — they cannot -/// amplify the way a download can, so they charge once decoded, not per chunk. -fn charge_decoded(decoded: Vec, budget: &ByteBudget) -> Result, String> { - budget - .claim(decoded.len() as u64) - .map_err(|()| over_budget("base64 payload"))?; - Ok(decoded) +/// Reserve the maximum decoded size before allocating. The reservation is +/// reconciled with the exact size afterwards because trailing padding can +/// reduce the result by up to two bytes. +fn decode_base64_budgeted(encoded: &str, budget: &ByteBudget) -> Result, String> { + let encoded = encoded.trim(); + let padding = encoded + .as_bytes() + .iter() + .rev() + .take_while(|&&byte| byte == b'=') + .take(2) + .count() as u64; + let estimate = (encoded.len() as u64) + .checked_add(3) + .and_then(|n| n.checked_div(4)) + .and_then(|n| n.checked_mul(3)) + .and_then(|n| n.checked_sub(padding)) + .ok_or_else(|| over_budget("base64 payload"))?; + let remaining = budget.remaining(); + budget.claim(estimate).map_err(|()| { + format!( + "{} (decoded size {estimate} bytes, {remaining} bytes remaining)", + over_budget("base64 payload") + ) + })?; + match base64::engine::general_purpose::STANDARD.decode(encoded.as_bytes()) { + Ok(decoded) => { + budget.release(estimate - decoded.len() as u64); + Ok(decoded) + } + Err(error) => { + budget.release(estimate); + Err(format!("media fetch: base64 decode: {error}")) + } + } } fn over_budget(what: &str) -> String { @@ -122,21 +192,6 @@ fn read_capped(mut reader: impl Read, what: &str, budget: &ByteBudget) -> Result } } -fn b64(encoded: &str) -> Result, String> { - // Slightly laxer than Python's `pybase64.b64decode(validate=True)`: - // surrounding whitespace (e.g. a trailing newline) is trimmed here. - let encoded = encoded.trim(); - // Reject by encoded length before allocating the decode buffer. - if encoded.len() as u64 / 4 * 3 > MAX_FETCH_BYTES { - return Err(format!( - "media fetch: base64 payload exceeds {MAX_FETCH_BYTES} bytes" - )); - } - base64::engine::general_purpose::STANDARD - .decode(encoded.as_bytes()) - .map_err(|e| format!("media fetch: base64 decode: {e}")) -} - /// Shared pooled agent honoring `HTTP_PROXY`/`HTTPS_PROXY`/`ALL_PROXY`, as the /// Python `requests` session does. fn http_agent() -> &'static ureq::Agent { @@ -268,6 +323,14 @@ mod tests { assert!(fetch_bytes("/definitely/not/here.jpg").is_err()); } + #[test] + fn trusted_local_reader_rejects_non_regular_files() { + let err = fetch_local_file_budgeted("/dev/zero", &ByteBudget::new(1024)) + .err() + .unwrap(); + assert!(err.contains("not a regular file"), "{err}"); + } + /// A non-regular file must hit the byte cap, not exhaust memory. #[test] fn unbounded_file_capped() { @@ -275,12 +338,29 @@ mod tests { assert!(err.contains("exceeds"), "{err}"); } - /// Oversized base64 is rejected from its encoded length, before decoding. + /// The convenience API keeps its 64 MiB budget, while a server request may + /// supply a larger bounded allowance for already-resident inline media. #[test] - fn oversized_base64_rejected() { + fn inline_base64_uses_the_supplied_request_budget() { let encoded = "A".repeat((MAX_FETCH_BYTES / 3 * 4 + 8) as usize); let err = fetch_bytes(&encoded).err().unwrap(); assert!(err.contains("exceeds"), "{err}"); + let decoded = fetch_bytes_budgeted(&encoded, &ByteBudget::new(MAX_FETCH_BYTES + 16)) + .expect("larger request budget admits inline media over the remote-fetch cap"); + assert!(decoded.len() as u64 > MAX_FETCH_BYTES); + } + + #[test] + fn base64_uses_the_callers_exact_budget() { + let encoded = base64::engine::general_purpose::STANDARD.encode(b"a"); + assert_eq!( + fetch_bytes_budgeted(&encoded, &ByteBudget::new(1)).unwrap(), + b"a" + ); + let err = fetch_bytes_budgeted(&encoded, &ByteBudget::new(0)) + .err() + .unwrap(); + assert!(err.contains("request media byte budget"), "{err}"); } /// One budget spans sources: each fits alone, the set does not. diff --git a/rust/sglang-mm/src/driver.rs b/rust/sglang-mm/src/driver.rs index 9b70a52ea..80fdc4ee4 100644 --- a/rust/sglang-mm/src/driver.rs +++ b/rust/sglang-mm/src/driver.rs @@ -8,10 +8,10 @@ use crate::common::{self, fetch, par, token_layout}; use crate::pipeline::{DecodedMedia, MmFamilyProcessor, PositionOutput, ProcessedItem}; -/// Per-request bounds: together with [`fetch::MAX_FETCH_BYTES`] they cap what -/// one request can make the pipeline buffer. +/// Per-request bounds. Every source consumes the aggregate budget; +/// [`fetch::MAX_FETCH_BYTES`] additionally caps each remote I/O stream. pub const MAX_ITEMS_PER_REQUEST: usize = 64; -pub const MAX_REQUEST_BYTES: u64 = 256 << 20; +pub const MAX_REQUEST_BYTES: u64 = 1280 << 20; /// One raw image source from the request. #[derive(Debug)] @@ -59,6 +59,16 @@ fn resolve(source: &ImageSource) -> Result, String> { } } +fn add_media_bytes(total: u64, next: usize) -> Result { + let total = total.saturating_add(next as u64); + if total > MAX_REQUEST_BYTES { + return Err(format!( + "multimodal request exceeds {MAX_REQUEST_BYTES} total media bytes" + )); + } + Ok(total) +} + /// Run one request through the pipeline. Any `Err` rejects the request back /// to the client — including inputs merely outside the pipeline's scope /// (video/audio, precomputed features, undecodable images), since there is @@ -85,12 +95,7 @@ pub fn process( let mut total: u64 = 0; for source in &input.images { let bytes = resolve(source)?; - total += bytes.len() as u64; - if total > MAX_REQUEST_BYTES { - return Err(format!( - "multimodal request exceeds {MAX_REQUEST_BYTES} total media bytes" - )); - } + total = add_media_bytes(total, bytes.len())?; fetched.push(bytes); } let processed: Vec<(ProcessedItem, u64)> = @@ -240,18 +245,8 @@ mod tests { .unwrap(); assert!(err.contains("media items"), "{err}"); - let chunk = (MAX_REQUEST_BYTES / 2 + 1) as usize; - let too_big = MmInput { - text: None, - input_ids: Some(vec![1, 1]), - images: vec![ - ImageSource::Bytes(vec![0; chunk]), - ImageSource::Bytes(vec![0; chunk]), - ], - }; - let err = process(family.as_ref(), too_big, |_| unreachable!()) - .err() - .unwrap(); + let err = add_media_bytes(MAX_REQUEST_BYTES / 2, (MAX_REQUEST_BYTES / 2 + 1) as usize) + .unwrap_err(); assert!(err.contains("total media bytes"), "{err}"); } diff --git a/rust/sglang-server/src/api_server/native_api.rs b/rust/sglang-server/src/api_server/native_api.rs index d65a76eed..71b791f55 100644 --- a/rust/sglang-server/src/api_server/native_api.rs +++ b/rust/sglang-server/src/api_server/native_api.rs @@ -232,7 +232,10 @@ async fn generate( let timing = RequestTiming::new(); // Media I/O (URL downloads, file reads) happens here, on the API runtime // — never on the MM worker pool (see `prefetch`). - if let Err(e) = super::prefetch::prefetch_all(&mut payloads).await { + if let Err(e) = + super::prefetch::prefetch_all(&mut payloads, &state.server_args.limit_mm_data_per_request) + .await + { return native_error(StatusCode::BAD_REQUEST, &e, stream); } if !is_batch { diff --git a/rust/sglang-server/src/api_server/prefetch.rs b/rust/sglang-server/src/api_server/prefetch.rs index 18b2a2e1e..b011a9538 100644 --- a/rust/sglang-server/src/api_server/prefetch.rs +++ b/rust/sglang-server/src/api_server/prefetch.rs @@ -2,16 +2,18 @@ //! //! The MM worker pool is fixed, core-pinned CPU capacity: a slow image host — or //! a file on a hanging network mount — must never occupy it, and a request's -//! images must download concurrently, not in `n * REQUEST_TIMEOUT`. URLs and -//! file paths resolve here through `sglang-mm`'s `fetch_bytes_budgeted` (one -//! owner for proxy/timeout/cap semantics) and ride out-of-band as +//! images must download concurrently, not in `n * REQUEST_TIMEOUT`. Remote and +//! inline sources resolve through `sglang-mm`'s `fetch_bytes_budgeted` (one +//! owner for proxy/timeout/cap semantics); trusted local files skip the remote +//! per-source cap but share the whole-request budget. Resolved bytes ride out-of-band as //! [`crate::message::request::MmData::prefetched`], which //! [`crate::multi_modality::payload::to_mm_input`] swaps back in. +use std::collections::BTreeMap; use std::sync::Arc; use bytes::Bytes; -use sglang_mm::common::fetch::{ByteBudget, fetch_bytes_budgeted}; +use sglang_mm::common::fetch::{ByteBudget, fetch_bytes_budgeted, fetch_local_file_budgeted}; use sglang_mm::driver::{MAX_ITEMS_PER_REQUEST, MAX_REQUEST_BYTES}; use tokio::sync::Semaphore; @@ -29,18 +31,46 @@ static PERMITS: Semaphore = Semaphore::const_new(32); /// enforced *here* rather than in `sglang_mm::driver::process`, where 64 sources /// of 64 MiB would already be resident. The driver keeps its own checks as the /// backstop for callers without a prefetch layer. -pub async fn prefetch_all(requests: &mut [GenerateRequest]) -> Result<(), String> { +pub async fn prefetch_all( + requests: &mut [GenerateRequest], + modality_limits: &BTreeMap, +) -> Result<(), String> { // The item budget rejects before a single byte is fetched. let plan = |mm: &Option>| -> Result, String> { - let Some(image_data) = mm.as_deref().and_then(|m| m.image_data.as_ref()) else { + let Some(mm) = mm.as_deref() else { return Ok(Vec::new()); }; - if item_count(image_data) > MAX_ITEMS_PER_REQUEST { + let values = [ + ("image", mm.image_data.as_ref()), + ("video", mm.video_data.as_ref()), + ("audio", mm.audio_data.as_ref()), + ]; + let items = values + .iter() + .filter_map(|(_, value)| *value) + .map(item_count) + .sum::(); + if items > MAX_ITEMS_PER_REQUEST { return Err(format!( "multimodal request exceeds {MAX_ITEMS_PER_REQUEST} media items" )); } - Ok(io_sources(image_data)) + for (modality, value) in values { + let count = value.map(item_count).unwrap_or_default(); + if let Some(limit) = modality_limits.get(modality) + && count > *limit + { + let display = modality[..1].to_uppercase() + &modality[1..]; + return Err(format!( + "{display} count {count} exceeds limit {limit} per request." + )); + } + } + Ok(values + .iter() + .filter_map(|(_, value)| *value) + .flat_map(io_sources) + .collect()) }; let plans = requests .iter() @@ -58,9 +88,11 @@ pub async fn prefetch_all(requests: &mut [GenerateRequest]) -> Result<(), String Ok(()) } -/// Resolve one request's sources concurrently (globally bounded), in order, -/// against one shared `total_bytes` allowance. Overflow rejects mid-download and -/// `try_join_all` drops the rest, so queued sources never start. +/// Resolve one request's sources concurrently (globally bounded), in order. +/// All inputs share `total_bytes`; trusted local files skip only the remote +/// per-source cap, matching Python's URL-only security limit. Overflow rejects +/// before or during I/O and `try_join_all` drops the rest, so queued sources +/// never start. async fn fetch_ordered(sources: Vec, total_bytes: u64) -> Result, String> { let budget = Arc::new(ByteBudget::new(total_bytes)); futures::future::try_join_all(sources.into_iter().map(|src| { @@ -71,10 +103,16 @@ async fn fetch_ordered(sources: Vec, total_bytes: u64) -> Result>(), + vec![b"image".as_ref(), b"video".as_ref(), b"audio".as_ref()] + ); + } + /// URLs and file paths resolve concurrently into `prefetched` in source /// order; CPU-only sources and mm-free requests are untouched. #[tokio::test] @@ -137,7 +205,7 @@ mod tests { ])), GenerateRequest::default(), ]; - prefetch_all(&mut requests).await.unwrap(); + prefetch_all(&mut requests, &BTreeMap::new()).await.unwrap(); std::fs::remove_file(&path).ok(); let fetched = &requests[0].mm.as_ref().unwrap().prefetched; // The one-shot server answers in accept order, so contents may swap @@ -151,7 +219,10 @@ mod tests { #[tokio::test] async fn failed_download_rejects() { let mut requests = vec![mm_request(Value::from("http://127.0.0.1:1/nope.png"))]; - let err = prefetch_all(&mut requests).await.err().unwrap(); + let err = prefetch_all(&mut requests, &BTreeMap::new()) + .await + .err() + .unwrap(); assert!(err.contains("media fetch"), "{err}"); } @@ -163,7 +234,10 @@ mod tests { .map(|i| Value::from(format!("/definitely/not/here-{i}.png"))) .collect(); let mut requests = vec![mm_request(Value::Array(sources))]; - let err = prefetch_all(&mut requests).await.err().unwrap(); + let err = prefetch_all(&mut requests, &BTreeMap::new()) + .await + .err() + .unwrap(); assert_eq!( err, format!("multimodal request exceeds {MAX_ITEMS_PER_REQUEST} media items") @@ -171,6 +245,25 @@ mod tests { assert!(requests[0].mm.as_ref().unwrap().prefetched.is_empty()); } + #[tokio::test] + async fn per_modality_budget_rejects_before_fetching() { + let mut requests = vec![GenerateRequest { + mm: Some(Box::new(MmData { + image_data: Some(Value::Array(vec![ + Value::from("/definitely/not/here-0.png"), + Value::from("/definitely/not/here-1.png"), + ])), + video_data: Some(Value::Array(vec![Value::from("/definitely/not/here.mp4")])), + ..Default::default() + })), + ..Default::default() + }]; + let limits = BTreeMap::from([("image".to_owned(), 1), ("video".to_owned(), 1)]); + let err = prefetch_all(&mut requests, &limits).await.err().unwrap(); + assert_eq!(err, "Image count 2 exceeds limit 1 per request."); + assert!(requests[0].mm.as_ref().unwrap().prefetched.is_empty()); + } + /// Sources legal alone but collectively over the limit are rejected while /// downloading, not once every body is resident. #[tokio::test] @@ -196,4 +289,24 @@ mod tests { let fetched = fetch_ordered(sources, MAX_REQUEST_BYTES).await.unwrap(); assert_eq!(fetched.iter().map(|b| b.len()).sum::(), 8192); } + + #[tokio::test] + async fn local_files_share_the_request_budget() { + let base = std::env::temp_dir().join(format!( + "sglang-prefetch-local-budget-{}", + std::process::id() + )); + std::fs::create_dir_all(&base).unwrap(); + let first = base.join("first.mp4"); + let second = base.join("second.mp4"); + std::fs::write(&first, b"first").unwrap(); + std::fs::write(&second, b"second").unwrap(); + + let sources = vec![first.display().to_string(), second.display().to_string()]; + let fetched = fetch_ordered(sources.clone(), 11).await.unwrap(); + let error = fetch_ordered(sources, 10).await.err().unwrap(); + std::fs::remove_dir_all(base).ok(); + assert_eq!(fetched.len(), 2); + assert!(error.contains("request media byte budget"), "{error}"); + } } diff --git a/rust/sglang-server/src/message/config.rs b/rust/sglang-server/src/message/config.rs index 9680d4f7a..b22ef9218 100644 --- a/rust/sglang-server/src/message/config.rs +++ b/rust/sglang-server/src/message/config.rs @@ -15,6 +15,7 @@ //! [`PreferredSamplingParams`] — are the only Python-facing code in this file; //! the rest is pure Rust. +use std::collections::BTreeMap; use std::net::SocketAddr; use std::sync::Arc; @@ -129,6 +130,8 @@ pub struct ServerArgs { /// Launch-time sampling defaults merged beneath per-request values and /// advertised by `/get_model_info`. pub preferred_sampling_params: Option, + /// Per-modality media-count limits from `--limit-mm-data-per-request`. + pub limit_mm_data_per_request: BTreeMap, /// Over-long inputs are truncated to fit the context instead of 400ing, and /// `max_new_tokens` is clamped rather than rejected (Python /// `TokenizerManager._validate_one_request`). @@ -172,6 +175,7 @@ impl ServerArgs { disaggregation_mode, model_config, preferred_sampling_params, + limit_mm_data_per_request, allow_auto_truncate, enable_return_hidden_states, num_reserved_tokens, @@ -202,6 +206,7 @@ impl ServerArgs { disaggregation_mode: DisaggregationMode, model_config: ModelConfig, preferred_sampling_params: Option, + limit_mm_data_per_request: BTreeMap, allow_auto_truncate: bool, enable_return_hidden_states: bool, num_reserved_tokens: u64, @@ -230,6 +235,7 @@ impl ServerArgs { disaggregation_mode, model_config, preferred_sampling_params, + limit_mm_data_per_request, allow_auto_truncate, enable_return_hidden_states, num_reserved_tokens, @@ -266,6 +272,7 @@ impl Default for ServerArgs { disaggregation_mode: DisaggregationMode::Null, model_config: ModelConfig::default(), preferred_sampling_params: None, + limit_mm_data_per_request: BTreeMap::new(), allow_auto_truncate: false, enable_return_hidden_states: false, num_reserved_tokens: 0,