minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Rust] Bound multimodal media ingress (#37967)

Browse Source
This commit is contained in:
Lianmin Zheng
2026-09-04 05:22:44 -07:00
committed by GitHub
parent e4adf63275
commit 978cc228ca
7 changed files with 274 additions and 74 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rust/sglang-server/src/api_server/native_api.rs
+4 -1
View File
@@ -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 {
rust/sglang-server/src/api_server/prefetch.rs
+131 -18
View File
@@ -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<String, usize>,
) -> Result<(), String> {
// The item budget rejects before a single byte is fetched.
let plan = |mm: &Option<Box<MmData>>| -> Result<Vec<String>, 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::<usize>();
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<String>, total_bytes: u64) -> Result<Vec<Bytes>, 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<String>, total_bytes: u64) -> Result<Vec<Byt
// an API worker. Those threads are pinned round-robin over the api
// core set (see `on_thread_start` in `runtime::start`) — off the
// CPU-bound stages, and mostly I/O-parked, so sharing is fine.
tokio::task::spawn_blocking(move || fetch_bytes_budgeted(&src, &budget))
.await
.map_err(|e| format!("media prefetch: {e}"))?
.map(Bytes::from)
tokio::task::spawn_blocking(move || {
if src.starts_with('/') || src.starts_with("file://") {
fetch_local_file_budgeted(&src, &budget)
} else {
fetch_bytes_budgeted(&src, &budget)
}
})
.await
.map_err(|e| format!("media prefetch: {e}"))?
.map(Bytes::from)
}
}))
.await
@@ -121,6 +159,36 @@ mod tests {
}
}
#[tokio::test]
async fn mixed_modalities_preserve_image_video_audio_order() {
let base = std::env::temp_dir().join(format!("sglang-prefetch-mm-{}", std::process::id()));
std::fs::create_dir_all(&base).unwrap();
let paths = [base.join("image"), base.join("video"), base.join("audio")];
for (path, body) in
paths
.iter()
.zip([b"image".as_ref(), b"video".as_ref(), b"audio".as_ref()])
{
std::fs::write(path, body).unwrap();
}
let mut requests = vec![GenerateRequest {
mm: Some(Box::new(MmData {
image_data: Some(Value::from(paths[0].display().to_string())),
video_data: Some(Value::from(paths[1].display().to_string())),
audio_data: Some(Value::from(paths[2].display().to_string())),
..Default::default()
})),
..Default::default()
}];
prefetch_all(&mut requests, &BTreeMap::new()).await.unwrap();
std::fs::remove_dir_all(base).ok();
let fetched = &requests[0].mm.as_ref().unwrap().prefetched;
assert_eq!(
fetched.iter().map(Bytes::as_ref).collect::<Vec<_>>(),
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::<usize>(), 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}");
}
}
rust/sglang-server/src/message/config.rs
+7
View File
@@ -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<PreferredSamplingParams>,
/// Per-modality media-count limits from `--limit-mm-data-per-request`.
pub limit_mm_data_per_request: BTreeMap<String, usize>,
/// 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<PreferredSamplingParams>,
limit_mm_data_per_request: BTreeMap<String, usize>,
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,