[Rust] Bound multimodal media ingress (#37967)
This commit is contained in:
@@ -107,10 +107,10 @@ Adding one = a `MmFamilyProcessor` impl in `src/<model>/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
|
||||
|
||||
|
||||
@@ -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<Vec<u8>, 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<Vec<u8>, 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<Vec<u8>, 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<Vec<u8>, 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<u8>, budget: &ByteBudget) -> Result<Vec<u8>, 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<Vec<u8>, 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<Vec<u8>, 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.
|
||||
|
||||
@@ -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<std::borrow::Cow<'_, [u8]>, String> {
|
||||
}
|
||||
}
|
||||
|
||||
fn add_media_bytes(total: u64, next: usize) -> Result<u64, String> {
|
||||
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}");
|
||||
}
|
||||
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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}");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user