[mm] rust-server: native multimodal processing for Qwen VL (integrate sglang-mm, e2e) (#32365)
Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Kan Wu
co-authored by
Claude Fable 5
Cursor
Claude Fable 5
Cursor
parent
dea07b348b
commit
32e5d788bd
@@ -43,14 +43,25 @@ dynamo-parsers = "7.0.1"
|
||||
dynamo-protocols = "5.1.0"
|
||||
dynamo-renderer = "5.0.0"
|
||||
flume = "0.12.0"
|
||||
# Safe POD slice casts (feature buffers viewed as bytes for the shm copy).
|
||||
bytemuck = "1"
|
||||
itertools = "0.14"
|
||||
hf-hub = { version = "0.4", default-features = false }
|
||||
# POSIX shm for the MM feature fan-out (`mm::ShmSegment`).
|
||||
libc = "0.2"
|
||||
# Same major as the workspace pyo3: the zero-copy MM drain (`take_mm`) moves
|
||||
# Rust vectors into numpy arrays.
|
||||
numpy = "0.29.0"
|
||||
rmp-serde = "1"
|
||||
rmpv = { version = "1", features = ["with-serde"] }
|
||||
# Pinned EXACTLY: this crate's accepted grammar defines the
|
||||
# "anything Rust admits, Python can compile" invariant in `message::sampling`.
|
||||
# A minor bump can widen it and silently reopen a scheduler-killing hole.
|
||||
regex-syntax = "=0.8.11"
|
||||
# The pure-Rust core of the MM pipeline. `default-features = false` drops the
|
||||
# pyo3 bindings so it links as a plain rlib; renamed so `use sglang_mm::…` reads
|
||||
# naturally while the crate keeps its own artifact name in the shared target/.
|
||||
sglang_mm = { package = "sglang-mm", path = "../sglang-mm", default-features = false }
|
||||
|
||||
[dev-dependencies]
|
||||
# `Router::oneshot` for handler-level router tests.
|
||||
|
||||
@@ -10,6 +10,7 @@ mod log;
|
||||
mod native_api;
|
||||
mod openai;
|
||||
mod pd_bootstrap;
|
||||
mod prefetch;
|
||||
mod submit;
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
@@ -154,7 +154,7 @@ async fn generate(
|
||||
let stream = body.stream;
|
||||
// Fan `text`/`input_ids`/`sampling_params` (scalar or list) into per-request
|
||||
// payloads. `is_batch` = list form → the response is a JSON array.
|
||||
let (payloads, is_batch) = match body.into_requests() {
|
||||
let (mut payloads, is_batch) = match body.into_requests() {
|
||||
Ok(v) => v,
|
||||
// The error carries its own status (a bad batch is `Validation` → 400).
|
||||
Err(e) => {
|
||||
@@ -162,6 +162,11 @@ async fn generate(
|
||||
return pre_submit_error(code, &e.to_string(), stream);
|
||||
}
|
||||
};
|
||||
// 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 {
|
||||
return pre_submit_error(StatusCode::BAD_REQUEST, &e, stream);
|
||||
}
|
||||
if !is_batch {
|
||||
// `into_requests` guarantees exactly one payload for a non-batch body.
|
||||
let payload = payloads
|
||||
|
||||
@@ -0,0 +1,199 @@
|
||||
//! Resolve I/O-backed media sources on the API runtime, before MM dispatch.
|
||||
//!
|
||||
//! 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
|
||||
//! [`crate::message::MmData::prefetched`], which
|
||||
//! [`crate::message::mm_payload::to_mm_input`] swaps back in.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use bytes::Bytes;
|
||||
use sglang_mm::common::fetch::{ByteBudget, fetch_bytes_budgeted};
|
||||
use sglang_mm::driver::{MAX_ITEMS_PER_REQUEST, MAX_REQUEST_BYTES};
|
||||
use tokio::sync::Semaphore;
|
||||
|
||||
use crate::message::mm_payload::{io_sources, item_count};
|
||||
use crate::message::{GenerateRequest, MmData};
|
||||
|
||||
/// Global bound on concurrent media fetches across all in-flight requests;
|
||||
/// excess acquisitions queue on the semaphore without holding a thread.
|
||||
static PERMITS: Semaphore = Semaphore::const_new(32);
|
||||
|
||||
/// Fill [`MmData::prefetched`] for every request, all fetches across the batch
|
||||
/// concurrent. Any failure rejects the call (a 400, as on the Python path).
|
||||
///
|
||||
/// The driver's budgets ([`MAX_ITEMS_PER_REQUEST`], [`MAX_REQUEST_BYTES`]) are
|
||||
/// 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> {
|
||||
// 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 {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
if item_count(image_data) > MAX_ITEMS_PER_REQUEST {
|
||||
return Err(format!(
|
||||
"multimodal request exceeds {MAX_ITEMS_PER_REQUEST} media items"
|
||||
));
|
||||
}
|
||||
Ok(io_sources(image_data))
|
||||
};
|
||||
let plans = requests
|
||||
.iter()
|
||||
.map(|r| plan(&r.mm))
|
||||
.collect::<Result<Vec<_>, String>>()?;
|
||||
let fetches = plans
|
||||
.into_iter()
|
||||
.map(|sources| fetch_ordered(sources, MAX_REQUEST_BYTES));
|
||||
let fetched = futures::future::try_join_all(fetches).await?;
|
||||
for (req, bytes) in requests.iter_mut().zip(fetched) {
|
||||
if !bytes.is_empty() {
|
||||
req.mm.as_mut().expect("sources came from mm").prefetched = bytes;
|
||||
}
|
||||
}
|
||||
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.
|
||||
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| {
|
||||
let budget = Arc::clone(&budget);
|
||||
async move {
|
||||
let _permit = PERMITS.acquire().await.expect("semaphore never closed");
|
||||
// Blocking I/O: parks a lazily-spawned blocking-pool thread, never
|
||||
// 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)
|
||||
}
|
||||
}))
|
||||
.await
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rmpv::Value;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn serve(bodies: Vec<Vec<u8>>) -> std::net::SocketAddr {
|
||||
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
std::thread::spawn(move || {
|
||||
for body in bodies {
|
||||
use std::io::{BufRead, Write};
|
||||
let (stream, _) = listener.accept().unwrap();
|
||||
let mut reader = std::io::BufReader::new(stream);
|
||||
let mut line = String::new();
|
||||
while reader.read_line(&mut line).unwrap() > 2 {
|
||||
line.clear(); // headers until the blank line
|
||||
}
|
||||
let mut stream = reader.into_inner();
|
||||
write!(
|
||||
stream,
|
||||
"HTTP/1.1 200 OK\r\nContent-Length: {}\r\n\r\n",
|
||||
body.len()
|
||||
)
|
||||
.unwrap();
|
||||
stream.write_all(&body).unwrap();
|
||||
}
|
||||
});
|
||||
addr
|
||||
}
|
||||
|
||||
fn mm_request(image_data: Value) -> GenerateRequest {
|
||||
GenerateRequest {
|
||||
mm: Some(Box::new(MmData {
|
||||
image_data: Some(image_data),
|
||||
..Default::default()
|
||||
})),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
/// URLs and file paths resolve concurrently into `prefetched` in source
|
||||
/// order; CPU-only sources and mm-free requests are untouched.
|
||||
#[tokio::test]
|
||||
async fn resolves_io_sources() {
|
||||
let addr = serve(vec![b"one".to_vec(), b"two".to_vec()]);
|
||||
let path = std::env::temp_dir().join(format!("sglang-prefetch-{}", std::process::id()));
|
||||
std::fs::write(&path, b"zzz").unwrap();
|
||||
let mut requests = vec![
|
||||
mm_request(Value::Array(vec![
|
||||
Value::from(format!("http://{addr}/a.png")),
|
||||
Value::from("data:image/png;base64,x"),
|
||||
Value::from(format!("http://{addr}/b.png")),
|
||||
Value::from(path.display().to_string()),
|
||||
])),
|
||||
GenerateRequest::default(),
|
||||
];
|
||||
prefetch_all(&mut requests).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
|
||||
// between the two URLs; all three bodies must arrive.
|
||||
let mut got: Vec<&[u8]> = fetched.iter().map(|b| b.as_ref()).collect();
|
||||
got.sort();
|
||||
assert_eq!(got, vec![b"one".as_ref(), b"two".as_ref(), b"zzz".as_ref()]);
|
||||
assert!(requests[1].mm.is_none());
|
||||
}
|
||||
|
||||
#[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();
|
||||
assert!(err.contains("media fetch"), "{err}");
|
||||
}
|
||||
|
||||
/// The item budget rejects before any source is touched: all of these would
|
||||
/// fail to fetch, so a fetch error would prove fetching started.
|
||||
#[tokio::test]
|
||||
async fn item_budget_rejects_before_fetching() {
|
||||
let sources: Vec<Value> = (0..=MAX_ITEMS_PER_REQUEST)
|
||||
.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();
|
||||
assert_eq!(
|
||||
err,
|
||||
format!("multimodal request exceeds {MAX_ITEMS_PER_REQUEST} media items")
|
||||
);
|
||||
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]
|
||||
async fn byte_budget_is_shared_across_sources() {
|
||||
let addr = serve(vec![vec![b'a'; 4096], vec![b'b'; 4096]]);
|
||||
let sources = vec![
|
||||
format!("http://{addr}/a.png"),
|
||||
format!("http://{addr}/b.png"),
|
||||
];
|
||||
// Room for one body, not both.
|
||||
let err = fetch_ordered(sources, 6144).await.err().unwrap();
|
||||
assert!(err.contains("request media byte budget"), "{err}");
|
||||
}
|
||||
|
||||
/// ...and a fitting set still fetches: the budget never over-rejects.
|
||||
#[tokio::test]
|
||||
async fn byte_budget_admits_a_fitting_request() {
|
||||
let addr = serve(vec![vec![b'a'; 4096], vec![b'b'; 4096]]);
|
||||
let sources = vec![
|
||||
format!("http://{addr}/a.png"),
|
||||
format!("http://{addr}/b.png"),
|
||||
];
|
||||
let fetched = fetch_ordered(sources, MAX_REQUEST_BYTES).await.unwrap();
|
||||
assert_eq!(fetched.iter().map(|b| b.len()).sum::<usize>(), 8192);
|
||||
}
|
||||
}
|
||||
@@ -14,8 +14,6 @@
|
||||
|
||||
use crate::error::Error;
|
||||
|
||||
// `Failed(Error)` carries the cause for observability even where it isn't read
|
||||
// back yet; `EncodeDone` belongs to the deferred Encoder edge.
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum RequestState {
|
||||
Received,
|
||||
@@ -41,8 +39,8 @@ pub enum RequestState {
|
||||
/// Outcome of validation, selecting the ingress branch.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub enum ValidationOutcome {
|
||||
/// Has multimodal inputs → Encoding. Deferred: no encoder yet.
|
||||
#[allow(dead_code)]
|
||||
/// Has multimodal inputs → Encoding, where an MM worker runs the native
|
||||
/// pipeline and returns the final expanded `input_ids`.
|
||||
HasMultimodal,
|
||||
/// Plain text → Tokenizing.
|
||||
NeedsTokenize,
|
||||
@@ -52,7 +50,6 @@ pub enum ValidationOutcome {
|
||||
|
||||
/// Events that drive transitions. Each variant maps 1:1 to an edge in the
|
||||
/// design's transition table.
|
||||
#[allow(dead_code)] // EncodeDone is the deferred Encoder edge.
|
||||
#[derive(Debug)]
|
||||
pub enum Event {
|
||||
// --- ingress ---
|
||||
@@ -125,7 +122,11 @@ impl RequestState {
|
||||
(Normalizing, Validated(HasMultimodal)) => Encoding,
|
||||
(Normalizing, Validated(NeedsTokenize)) => Tokenizing,
|
||||
(Normalizing, Validated(AlreadyTokenized)) => PreSendValidating,
|
||||
(Encoding, EncodeDone) => Tokenizing,
|
||||
// The MM worker returns the *final* placeholder-expanded input_ids,
|
||||
// so an encoded request skips the tokenizer pool — but not the
|
||||
// pre-send checks: expanded image tokens count against the same
|
||||
// input + max_new_tokens ceiling as tokenized text.
|
||||
(Encoding, EncodeDone) => PreSendValidating,
|
||||
// Every ingress branch funnels through the pre-send checks, so they
|
||||
// run exactly once per request no matter how it got its ids.
|
||||
(Tokenizing, TokenizeDone) => PreSendValidating,
|
||||
|
||||
@@ -17,6 +17,7 @@ mod error;
|
||||
mod fsm;
|
||||
mod ids;
|
||||
mod message;
|
||||
mod mm;
|
||||
mod ring;
|
||||
mod runtime;
|
||||
mod tokenizer;
|
||||
@@ -31,6 +32,21 @@ use pyo3::types::PyBytes;
|
||||
|
||||
use crate::runtime::{Runtime, RuntimeConfig};
|
||||
|
||||
/// One drained MM result (see [`Server::take_mm`]). Exactly one of
|
||||
/// `features`/`shm_names` is `Some`: inline features for single-rank serving
|
||||
/// (zero-copy into numpy), or one POSIX segment name per item when the scheduler
|
||||
/// broadcasts across TP ranks and Python wraps each in a `ShmPointerMMData`.
|
||||
#[pyclass(frozen, get_all)]
|
||||
struct MmHandoff {
|
||||
features: Option<Py<numpy::PyArray1<f32>>>,
|
||||
shm_names: Option<Vec<String>>,
|
||||
grids: Vec<(u32, u32, u32)>,
|
||||
hashes: Vec<u64>,
|
||||
offsets: Vec<(u32, u32)>,
|
||||
mrope: Py<numpy::PyArray1<i64>>,
|
||||
mrope_delta: i64,
|
||||
}
|
||||
|
||||
/// Columnar ingress batch handed to Python by [`Server::recv_requests`].
|
||||
/// `frozen`: immutable snapshot, so field access never contends on a borrow.
|
||||
#[pyclass(frozen, get_all)]
|
||||
@@ -199,6 +215,54 @@ impl Server {
|
||||
self.push_frame(py, crate::message::frame_egress_error(rid, message))
|
||||
}
|
||||
|
||||
/// Spawn the MM worker pool for the pipeline in `spec_json` (built from the
|
||||
/// resolved processor config; see `NativeMmHost.resolve_native_spec`).
|
||||
/// Image-only requests are processed entirely in Rust and parked for
|
||||
/// [`Server::take_mm`]; anything the pipeline cannot serve is rejected back to
|
||||
/// the client — there is no Python fallback.
|
||||
fn start_mm_workers(&self, spec_json: &str, workers: usize) -> PyResult<()> {
|
||||
let ctx = mm::Context::new(
|
||||
spec_json,
|
||||
self.rt.tokenizer.clone(),
|
||||
self.rt.mm_sidecar.clone(),
|
||||
)
|
||||
.map_err(PyErr::new::<pyo3::exceptions::PyValueError, _>)?;
|
||||
self.rt.spawn_mm_pool(workers, std::sync::Arc::new(ctx));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Pop the MM result for `rid` — parked strictly before the request reached
|
||||
/// the ingress ring — or `None` if there is none. The numeric buffers become
|
||||
/// 1-D numpy arrays that take **ownership** of the Rust vectors, no copy.
|
||||
///
|
||||
/// Runs on the scheduler loop (`RustServer.drain`, under the GIL) between
|
||||
/// decode steps, so any per-byte work here — memcpy or hashing, tens of MB
|
||||
/// per image-heavy request — would stall every running request's ITL. Hence
|
||||
/// the worker-precomputed `hashes`.
|
||||
fn take_mm(&self, py: Python<'_>, rid: &str) -> Option<MmHandoff> {
|
||||
use numpy::IntoPyArray;
|
||||
|
||||
let res = self.rt.mm_sidecar.take(rid)?;
|
||||
let (features, shm_names) = match res.features {
|
||||
mm::FeatureStore::Inline(v) => (Some(v.into_pyarray(py).unbind()), None),
|
||||
// The segments — and the duty to unlink — move to Python here;
|
||||
// `materialize()` unlinks after the post-broadcast clone on each rank.
|
||||
mm::FeatureStore::Shm(segments) => (
|
||||
None,
|
||||
Some(segments.into_iter().map(|s| s.into_name()).collect()),
|
||||
),
|
||||
};
|
||||
Some(MmHandoff {
|
||||
features,
|
||||
shm_names,
|
||||
grids: res.grids.iter().map(|g| (g[0], g[1], g[2])).collect(),
|
||||
hashes: res.hashes,
|
||||
offsets: res.offsets,
|
||||
mrope: res.mrope.into_pyarray(py).unbind(),
|
||||
mrope_delta: res.mrope_delta,
|
||||
})
|
||||
}
|
||||
|
||||
/// Signal all threads to stop (best effort).
|
||||
fn shutdown(&self) {
|
||||
self.rt.request_shutdown();
|
||||
@@ -245,5 +309,6 @@ fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
.try_init();
|
||||
m.add_class::<Server>()?;
|
||||
m.add_class::<IngressBatch>()?;
|
||||
m.add_class::<MmHandoff>()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -11,6 +11,7 @@
|
||||
mod egress;
|
||||
mod finish_reason;
|
||||
mod io_struct;
|
||||
pub mod mm_payload;
|
||||
mod request;
|
||||
mod sampling;
|
||||
mod types;
|
||||
@@ -22,7 +23,10 @@ pub use egress::{
|
||||
};
|
||||
pub use finish_reason::Matched;
|
||||
pub(crate) use io_struct::{AbortReq, ControlRequest, GetInternalStateReq};
|
||||
pub use request::{GenerateBody, GenerateRequest, RequestKind};
|
||||
pub use request::{GenerateBody, GenerateRequest, MmRequest, MmWorkItem, RequestKind};
|
||||
// Constructed directly only by tests: `api_server::prefetch` fills its
|
||||
// `prefetched` field, everything else gets it packed inside a `GenerateRequest`.
|
||||
pub use request::MmData;
|
||||
pub(crate) use sampling::{SamplingParams, SamplingParamsInput};
|
||||
pub(crate) use types::{OneOrMany, OneOrManyItem, TokenIds};
|
||||
|
||||
|
||||
@@ -0,0 +1,222 @@
|
||||
//! Convert a parked request's [`MmWorkItem`] into the typed [`MmInput`] the
|
||||
//! `sglang-mm` driver consumes — an in-process handoff, nothing serialized.
|
||||
//!
|
||||
//! Every `Err` rejects the request back to the client; the message says whether
|
||||
//! the input is malformed or merely outside the pipeline's scope (video/audio,
|
||||
//! precomputed features, …).
|
||||
|
||||
use bytes::Bytes;
|
||||
use rmpv::Value;
|
||||
use sglang_mm::driver::{ImageSource, MmInput};
|
||||
|
||||
use super::request::MmWorkItem;
|
||||
|
||||
/// True for sources the API layer must resolve before MM dispatch: I/O — network
|
||||
/// *or* disk, since a network mount can hang past any HTTP timeout — never runs
|
||||
/// on the fixed MM worker pool (see `api_server::prefetch`). `data:` and bare
|
||||
/// base64 are pure CPU and stay on the worker. Lives next to `collect_images` so
|
||||
/// the prefetch walk and the parse walk cannot drift.
|
||||
pub fn is_io_source(src: &str) -> bool {
|
||||
src.starts_with("http://")
|
||||
|| src.starts_with("https://")
|
||||
|| src.starts_with("file://")
|
||||
|| src.starts_with('/')
|
||||
}
|
||||
|
||||
/// The I/O-backed sources of an `image_data` value, in `collect_images` order.
|
||||
pub fn io_sources(value: &Value) -> Vec<String> {
|
||||
let mut out = Vec::new();
|
||||
let mut walk = |value: &Value| {
|
||||
if let Some(src) = value.as_str().filter(|s| is_io_source(s)) {
|
||||
out.push(src.to_owned());
|
||||
}
|
||||
};
|
||||
if let Value::Array(values) = value {
|
||||
values.iter().for_each(&mut walk);
|
||||
} else {
|
||||
walk(value);
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// How many media items an `image_data` value contributes, walked the way
|
||||
/// [`collect_images`] walks it, so the item budget can reject before fetching.
|
||||
pub fn item_count(value: &Value) -> usize {
|
||||
match value {
|
||||
Value::Nil => 0,
|
||||
Value::Array(values) => values.iter().map(item_count).sum(),
|
||||
_ => 1,
|
||||
}
|
||||
}
|
||||
|
||||
/// I/O-backed sources are swapped for their `work.prefetched` bytes (in
|
||||
/// [`io_sources`] order); one left without an entry is an internal error here,
|
||||
/// never a fetch.
|
||||
pub fn to_mm_input(work: MmWorkItem) -> Result<MmInput, String> {
|
||||
let present = |v: &Option<Value>| v.as_ref().is_some_and(value_present);
|
||||
if present(&work.video_data) || present(&work.audio_data) {
|
||||
return Err("unsupported modality: video/audio input".into());
|
||||
}
|
||||
let mut images = Vec::new();
|
||||
if let Some(image_data) = &work.image_data {
|
||||
collect_images(image_data, &mut work.prefetched.iter(), &mut images)?;
|
||||
}
|
||||
if images.is_empty() {
|
||||
return Err("no raw image sources in mm input".into());
|
||||
}
|
||||
Ok(MmInput {
|
||||
text: work.text,
|
||||
input_ids: work.input_ids,
|
||||
images,
|
||||
})
|
||||
}
|
||||
|
||||
fn collect_images(
|
||||
value: &Value,
|
||||
prefetched: &mut std::slice::Iter<Bytes>,
|
||||
out: &mut Vec<ImageSource>,
|
||||
) -> Result<(), String> {
|
||||
match value {
|
||||
Value::Nil => Ok(()),
|
||||
Value::String(value) => {
|
||||
let value = value
|
||||
.as_str()
|
||||
.ok_or_else(|| "non-utf8 image source".to_string())?;
|
||||
if is_io_source(value) {
|
||||
let bytes = prefetched
|
||||
.next()
|
||||
.ok_or_else(|| "I/O-backed image source was not prefetched".to_string())?;
|
||||
out.push(ImageSource::Bytes(bytes.to_vec()));
|
||||
} else {
|
||||
out.push(ImageSource::String(value.to_owned()));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
Value::Binary(value) => {
|
||||
out.push(ImageSource::Bytes(value.clone()));
|
||||
Ok(())
|
||||
}
|
||||
Value::Array(values) => {
|
||||
for value in values {
|
||||
match value {
|
||||
Value::String(_) | Value::Binary(_) | Value::Nil => {
|
||||
collect_images(value, prefetched, out)?
|
||||
}
|
||||
_ => {
|
||||
return Err("unsupported image_data shape: nested/typed item".into());
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
_ => Err("unsupported image_data shape".into()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Rust mirror of Python `has_valid_data`: `nil` and (recursively) empty or
|
||||
/// all-nil lists don't count as multimodal input. Shared with the ingress
|
||||
/// `has_multimodal` check so routing and parsing cannot drift.
|
||||
pub fn value_present(value: &Value) -> bool {
|
||||
match value {
|
||||
Value::Nil => false,
|
||||
Value::Array(values) => values.iter().any(value_present),
|
||||
_ => true,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn image_work(image: Value) -> MmWorkItem {
|
||||
MmWorkItem {
|
||||
text: Some("prompt".into()),
|
||||
image_data: Some(image),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn converts_string_and_list_images() {
|
||||
let one = to_mm_input(image_work(Value::from("data:image/png;base64,x"))).unwrap();
|
||||
assert_eq!(one.images.len(), 1);
|
||||
let many = to_mm_input(image_work(Value::Array(vec![
|
||||
Value::from("a"),
|
||||
Value::from("b"),
|
||||
])))
|
||||
.unwrap();
|
||||
assert_eq!(many.images.len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unsupported_modalities_and_shapes_rejected() {
|
||||
let video = MmWorkItem {
|
||||
video_data: Some(Value::from("video.mp4")),
|
||||
..Default::default()
|
||||
};
|
||||
assert!(to_mm_input(video).err().unwrap().contains("video/audio"));
|
||||
|
||||
let dict = Value::Map(vec![(Value::from("format"), Value::from("x"))]);
|
||||
assert!(
|
||||
to_mm_input(image_work(Value::Array(vec![dict])))
|
||||
.err()
|
||||
.unwrap()
|
||||
.contains("image_data shape")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_video_audio_lists_are_not_modalities() {
|
||||
// Mirrors Python `has_valid_data`: nil / empty lists don't count.
|
||||
let work = MmWorkItem {
|
||||
input_ids: Some(vec![1]),
|
||||
image_data: Some(Value::from("a")),
|
||||
video_data: Some(Value::Array(vec![])),
|
||||
audio_data: Some(Value::Array(vec![Value::Array(vec![])])),
|
||||
..Default::default()
|
||||
};
|
||||
assert_eq!(to_mm_input(work).unwrap().images.len(), 1);
|
||||
}
|
||||
|
||||
/// I/O-backed sources (URLs, file paths) take their prefetched bytes in walk
|
||||
/// order; one left unfetched errors, so no I/O can reach an MM worker.
|
||||
#[test]
|
||||
fn io_sources_use_prefetched_bytes() {
|
||||
let image = Value::Array(vec![
|
||||
Value::from("http://a/x.png"),
|
||||
Value::from("data:image/png;base64,x"),
|
||||
Value::from("/mnt/nfs/y.png"),
|
||||
]);
|
||||
assert_eq!(io_sources(&image), vec!["http://a/x.png", "/mnt/nfs/y.png"]);
|
||||
|
||||
let mut work = image_work(image.clone());
|
||||
work.prefetched = vec![Bytes::from_static(b"aa"), Bytes::from_static(b"bb")];
|
||||
let input = to_mm_input(work).unwrap();
|
||||
let as_bytes = |i: usize| match &input.images[i] {
|
||||
ImageSource::Bytes(b) => b.as_slice(),
|
||||
other => panic!("expected bytes, got {other:?}"),
|
||||
};
|
||||
assert_eq!(as_bytes(0), b"aa");
|
||||
assert_eq!(as_bytes(2), b"bb");
|
||||
assert!(matches!(&input.images[1], ImageSource::String(_)));
|
||||
|
||||
let err = to_mm_input(image_work(image)).err().unwrap();
|
||||
assert!(err.contains("not prefetched"), "{err}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn image_free_work_rejected() {
|
||||
assert!(
|
||||
to_mm_input(image_work(Value::Nil))
|
||||
.err()
|
||||
.unwrap()
|
||||
.contains("no raw image sources")
|
||||
);
|
||||
assert!(
|
||||
to_mm_input(MmWorkItem::default())
|
||||
.err()
|
||||
.unwrap()
|
||||
.contains("no raw image sources")
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -106,6 +106,23 @@ pub struct GenerateBody {
|
||||
pub routed_dp_rank: Option<i64>,
|
||||
#[serde(default)]
|
||||
pub disagg_prefill_dp_rank: Option<i64>,
|
||||
// Multimodal inputs, permissive `Value` so any shape Python's
|
||||
// `GenerateReqInput` accepts (URL / base64 / list / list-of-lists) parses.
|
||||
// `into_requests` fans them out per the Python
|
||||
// `_normalize_{image,video,audio}_data` batch rules.
|
||||
#[serde(default)]
|
||||
pub image_data: Option<rmpv::Value>,
|
||||
/// Caller-supplied per-item content hashes (hex) overriding the computed
|
||||
/// ones, so an external router's keys align with the prefix cache. Single
|
||||
/// requests only: Python declares the batched shapes but `__getitem__` never
|
||||
/// forwards them, so a batch is rejected here rather than answered with
|
||||
/// hashes it did not ask for.
|
||||
#[serde(default)]
|
||||
pub mm_hashes: Option<rmpv::Value>,
|
||||
#[serde(default)]
|
||||
pub video_data: Option<rmpv::Value>,
|
||||
#[serde(default)]
|
||||
pub audio_data: Option<rmpv::Value>,
|
||||
}
|
||||
|
||||
impl GenerateBody {
|
||||
@@ -135,6 +152,10 @@ impl GenerateBody {
|
||||
decode_tp_size,
|
||||
routed_dp_rank,
|
||||
disagg_prefill_dp_rank,
|
||||
image_data,
|
||||
video_data,
|
||||
audio_data,
|
||||
mm_hashes,
|
||||
// Unported `GenerateReqInput` fields land here and are dropped, as they
|
||||
// are on the Python path.
|
||||
..
|
||||
@@ -320,10 +341,26 @@ impl GenerateBody {
|
||||
let bootstrap_pair_keys =
|
||||
flatten_column(fan_out(bootstrap_pair_key, n, "bootstrap_pair_key")?);
|
||||
let decode_tp_sizes = flatten_column(fan_out(decode_tp_size, n, "decode_tp_size")?);
|
||||
// `mm_hashes` has no batch form: honoring it only here would give the two
|
||||
// servers different prefix-cache keys for the same body. Reject instead of
|
||||
// dropping it silently as Python does — the field exists to align a
|
||||
// caller's keys, so ignoring it returns subtly wrong ones.
|
||||
if is_batch && mm_value_present(&mm_hashes) {
|
||||
return Err(Error::Validation(
|
||||
"mm_hashes is not supported for batch requests; send one request per prompt".into(),
|
||||
));
|
||||
}
|
||||
// Multimodal columns; see `split_mm_column` for the Python parity rules.
|
||||
let images = split_mm_column(image_data, n, is_batch, MmBroadcast::WrapInList)
|
||||
.map_err(|e| Error::Validation(format!("image_data: {e}")))?;
|
||||
let videos = split_mm_column(video_data, n, is_batch, MmBroadcast::AsIs)
|
||||
.map_err(|e| Error::Validation(format!("video_data: {e}")))?;
|
||||
let audios = split_mm_column(audio_data, n, is_batch, MmBroadcast::AsIs)
|
||||
.map_err(|e| Error::Validation(format!("audio_data: {e}")))?;
|
||||
|
||||
// Every column above is exactly `n` long, so zip them by value: each
|
||||
// request takes ownership of its cell, with no indexing or bounds checks.
|
||||
let requests = izip!(
|
||||
let mut requests: Vec<GenerateRequest> = izip!(
|
||||
rids,
|
||||
texts,
|
||||
id_lists,
|
||||
@@ -338,6 +375,9 @@ impl GenerateBody {
|
||||
bootstrap_rooms,
|
||||
bootstrap_pair_keys,
|
||||
decode_tp_sizes,
|
||||
images,
|
||||
videos,
|
||||
audios,
|
||||
)
|
||||
.map(
|
||||
|(
|
||||
@@ -355,6 +395,9 @@ impl GenerateBody {
|
||||
bootstrap_room,
|
||||
bootstrap_pair_key,
|
||||
decode_tp_size,
|
||||
image_data,
|
||||
video_data,
|
||||
audio_data,
|
||||
)| GenerateRequest {
|
||||
rid,
|
||||
text,
|
||||
@@ -381,13 +424,126 @@ impl GenerateBody {
|
||||
decode_tp_size,
|
||||
routed_dp_rank,
|
||||
disagg_prefill_dp_rank,
|
||||
mm: pack_mm(image_data, video_data, audio_data),
|
||||
},
|
||||
)
|
||||
.collect();
|
||||
// Single requests only (batches rejected above). Malformed entries are
|
||||
// dropped here and warned about in `mm::apply_caller_hashes`, never a 400.
|
||||
if !is_batch
|
||||
&& let (Some(rmpv::Value::Array(vals)), Some(req)) = (mm_hashes, requests.first_mut())
|
||||
&& let Some(mm) = req.mm.as_deref_mut()
|
||||
{
|
||||
mm.mm_hashes = vals
|
||||
.iter()
|
||||
.filter_map(|v| v.as_str().map(str::to_owned))
|
||||
.collect();
|
||||
}
|
||||
Ok((requests, is_batch))
|
||||
}
|
||||
}
|
||||
|
||||
/// Box the per-item mm values, `None` when the item has none — the common
|
||||
/// text-only case keeps `GenerateRequest` slim.
|
||||
fn pack_mm(
|
||||
image_data: Option<rmpv::Value>,
|
||||
video_data: Option<rmpv::Value>,
|
||||
audio_data: Option<rmpv::Value>,
|
||||
) -> Option<Box<MmData>> {
|
||||
if image_data.is_none() && video_data.is_none() && audio_data.is_none() {
|
||||
return None;
|
||||
}
|
||||
Some(Box::new(MmData {
|
||||
image_data,
|
||||
video_data,
|
||||
audio_data,
|
||||
..Default::default()
|
||||
}))
|
||||
}
|
||||
|
||||
/// How a scalar mm value broadcasts across a batch: images become a one-image
|
||||
/// list per item (`[[img]] * num` in Python `_normalize_image_data`),
|
||||
/// video/audio broadcast bare (`[v] * num` in `_normalize_video_data`).
|
||||
#[derive(Clone, Copy)]
|
||||
enum MmBroadcast {
|
||||
WrapInList,
|
||||
AsIs,
|
||||
}
|
||||
|
||||
/// Fan one mm field into per-item values, mirroring Python's
|
||||
/// `_normalize_{image,video,audio}_data`:
|
||||
/// * `None` / empty list → `None` for every item;
|
||||
/// * single request → the raw value passes through (the processor wraps a
|
||||
/// non-list into a one-element list);
|
||||
/// * batch + scalar → broadcast to every item, per `MmBroadcast`;
|
||||
/// * batch + list → per item, length must equal the batch size.
|
||||
fn split_mm_column(
|
||||
v: Option<rmpv::Value>,
|
||||
n: usize,
|
||||
is_batch: bool,
|
||||
broadcast: MmBroadcast,
|
||||
) -> Result<Vec<Option<rmpv::Value>>, String> {
|
||||
let Some(v) = v else {
|
||||
return Ok(vec![None; n]);
|
||||
};
|
||||
if v.is_nil() {
|
||||
return Ok(vec![None; n]);
|
||||
}
|
||||
if !is_batch {
|
||||
return Ok(vec![Some(v)]);
|
||||
}
|
||||
match v {
|
||||
rmpv::Value::Array(items) if items.is_empty() => Ok(vec![None; n]),
|
||||
rmpv::Value::Array(items) => {
|
||||
if items.len() != n {
|
||||
return Err(format!(
|
||||
"list length {} does not match batch size {n}",
|
||||
items.len()
|
||||
));
|
||||
}
|
||||
Ok(items.into_iter().map(Some).collect())
|
||||
}
|
||||
scalar => {
|
||||
// A broadcast deep-clones once per prompt — same blow-up as
|
||||
// sampling_params above, so bound the product before any clone.
|
||||
check_broadcast_budget(scalar.heap_bytes(), n, "value").map_err(|e| e.to_string())?;
|
||||
Ok(match broadcast {
|
||||
MmBroadcast::WrapInList => vec![Some(rmpv::Value::Array(vec![scalar])); n],
|
||||
MmBroadcast::AsIs => vec![Some(scalar); n],
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// One request handed to the MM worker pool: the rid to correlate the result,
|
||||
/// plus the owned inputs from [`GenerateRequest::take_mm_work`].
|
||||
#[derive(Debug)]
|
||||
pub struct MmRequest {
|
||||
pub rid: crate::ids::Rid,
|
||||
pub work: MmWorkItem,
|
||||
}
|
||||
|
||||
/// The parked request's fields the MM worker owns; converted to the driver input
|
||||
/// by [`super::mm_payload::to_mm_input`].
|
||||
#[derive(Debug, Default)]
|
||||
pub struct MmWorkItem {
|
||||
pub text: Option<String>,
|
||||
pub input_ids: Option<Vec<i32>>,
|
||||
pub image_data: Option<rmpv::Value>,
|
||||
pub video_data: Option<rmpv::Value>,
|
||||
pub audio_data: Option<rmpv::Value>,
|
||||
/// See [`MmData::prefetched`].
|
||||
pub prefetched: Vec<Bytes>,
|
||||
/// See [`GenerateBody::mm_hashes`].
|
||||
pub mm_hashes: Vec<String>,
|
||||
}
|
||||
|
||||
/// Whether an optional mm field counts as multimodal input, via the same
|
||||
/// `value_present` the MM worker's payload parser uses.
|
||||
fn mm_value_present(v: &Option<rmpv::Value>) -> bool {
|
||||
v.as_ref().is_some_and(super::mm_payload::value_present)
|
||||
}
|
||||
|
||||
/// Request variant — selects the ingress branch, scheduler wire message, and
|
||||
/// egress shape. Each owns its body, so generate/control fields stay type-separate.
|
||||
#[derive(Debug)]
|
||||
@@ -475,6 +631,26 @@ pub struct GenerateRequest {
|
||||
/// so these are pure passthrough for the scheduler/LB protocol.
|
||||
pub routed_dp_rank: Option<i64>,
|
||||
pub disagg_prefill_dp_rank: Option<i64>,
|
||||
/// Multimodal inputs, carried opaquely. Consumed by the Encoding stage,
|
||||
/// which ships them to the MM worker pool; never read by the tokenizer or
|
||||
/// serialized onto the scheduler header. Boxed so the common text-only
|
||||
/// request doesn't grow every `Request` moved between stages.
|
||||
pub mm: Option<Box<MmData>>,
|
||||
}
|
||||
|
||||
/// The opaque multimodal fields of one request (see [`GenerateRequest::mm`]).
|
||||
#[derive(Debug, Default)]
|
||||
pub struct MmData {
|
||||
pub image_data: Option<rmpv::Value>,
|
||||
pub video_data: Option<rmpv::Value>,
|
||||
pub audio_data: Option<rmpv::Value>,
|
||||
/// Bytes of `image_data`'s I/O-backed sources, resolved by
|
||||
/// `api_server::prefetch` in `mm_payload::io_sources` order so MM workers
|
||||
/// never block on I/O. Out-of-band: the values above stay as the client
|
||||
/// sent them.
|
||||
pub prefetched: Vec<bytes::Bytes>,
|
||||
/// See [`GenerateBody::mm_hashes`]; applied by the MM worker.
|
||||
pub mm_hashes: Vec<String>,
|
||||
}
|
||||
|
||||
impl GenerateRequest {
|
||||
@@ -483,11 +659,33 @@ impl GenerateRequest {
|
||||
self.input_ids.as_ref().is_some_and(|v| !v.is_empty())
|
||||
}
|
||||
|
||||
/// Multimodal detection hook. Deferred (Encoder stubbed): always false until mm
|
||||
/// fields are wired in.
|
||||
#[allow(dead_code)]
|
||||
/// True when the request carries a usable multimodal payload — the mirror of
|
||||
/// Python `GenerateReqInput.contains_mm_input()`.
|
||||
pub fn has_multimodal(&self) -> bool {
|
||||
false
|
||||
self.mm.as_ref().is_some_and(|mm| {
|
||||
mm_value_present(&mm.image_data)
|
||||
|| mm_value_present(&mm.video_data)
|
||||
|| mm_value_present(&mm.audio_data)
|
||||
})
|
||||
}
|
||||
|
||||
/// Carve out the MM worker's inputs: `text` is cloned (the scheduler header
|
||||
/// still needs it), `input_ids` is taken (the expanded ids replace it), and
|
||||
/// the mm values move wholesale.
|
||||
pub fn take_mm_work(&mut self) -> MmWorkItem {
|
||||
let mut work = MmWorkItem {
|
||||
text: self.text.clone(),
|
||||
input_ids: self.input_ids.take(),
|
||||
..Default::default()
|
||||
};
|
||||
if let Some(m) = self.mm.as_deref_mut() {
|
||||
work.image_data = m.image_data.take();
|
||||
work.video_data = m.video_data.take();
|
||||
work.audio_data = m.audio_data.take();
|
||||
work.prefetched = std::mem::take(&mut m.prefetched);
|
||||
work.mm_hashes = std::mem::take(&mut m.mm_hashes);
|
||||
}
|
||||
work
|
||||
}
|
||||
|
||||
pub fn encode_header(&self) -> Result<Bytes, Error> {
|
||||
@@ -539,6 +737,23 @@ impl<T: HeapBytes> HeapBytes for Option<T> {
|
||||
self.as_ref().map_or(0, HeapBytes::heap_bytes)
|
||||
}
|
||||
}
|
||||
impl HeapBytes for rmpv::Value {
|
||||
fn heap_bytes(&self) -> usize {
|
||||
use rmpv::Value;
|
||||
const NODE: usize = std::mem::size_of::<rmpv::Value>();
|
||||
match self {
|
||||
Value::String(s) => s.as_bytes().len(),
|
||||
Value::Binary(b) => b.len(),
|
||||
Value::Ext(_, b) => b.len(),
|
||||
Value::Array(items) => items.iter().map(|v| NODE + v.heap_bytes()).sum(),
|
||||
Value::Map(entries) => entries
|
||||
.iter()
|
||||
.map(|(k, v)| 2 * NODE + k.heap_bytes() + v.heap_bytes())
|
||||
.sum(),
|
||||
_ => 0,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Collapse `fan_out`'s nullable-element output: outer `None` (field absent /
|
||||
/// scalar broadcast of nothing) and inner `None` (an explicit `null` list
|
||||
@@ -743,8 +958,9 @@ mod tests {
|
||||
}
|
||||
|
||||
/// The native `bench_serving` payload (a `GenerateReqInput` superset) parses:
|
||||
/// its `lora_path`/`return_routed_experts`/`image_data` are accepted-but-ignored,
|
||||
/// so `split` succeeds and drops them while the real fields survive.
|
||||
/// its `lora_path`/`return_routed_experts` are accepted-but-ignored and a
|
||||
/// `null` `image_data` means "no multimodal input", so `split` succeeds
|
||||
/// while the real fields survive.
|
||||
#[test]
|
||||
fn accepts_bench_serving_payload() {
|
||||
let (ps, is_batch) = requests(
|
||||
@@ -758,6 +974,111 @@ mod tests {
|
||||
assert_eq!(ps.len(), 1);
|
||||
assert_eq!(ps[0].text.as_deref(), Some("hi"));
|
||||
assert!(ps[0].stream);
|
||||
assert!(!ps[0].has_multimodal());
|
||||
}
|
||||
|
||||
/// Mm columns fan out per Python `_normalize_{image,video}_data`: a single
|
||||
/// request passes the raw value through; a batch broadcasts a scalar image as
|
||||
/// `[img]` per item, maps a list per item with matching lengths, and treats
|
||||
/// `null`/`[]` as absent.
|
||||
#[test]
|
||||
fn split_mm_fanout_matches_python_normalize() {
|
||||
let image_of = |p: &GenerateRequest| p.mm.as_ref().unwrap().image_data.clone().unwrap();
|
||||
|
||||
// Single request: raw value passes through untouched.
|
||||
let (ps, _) = requests(r#"{"text": "a", "image_data": "http://x/i.jpg"}"#).unwrap();
|
||||
assert_eq!(image_of(&ps[0]).as_str(), Some("http://x/i.jpg"));
|
||||
assert!(ps[0].has_multimodal());
|
||||
|
||||
// Batch + scalar image: broadcast, wrapped as a one-image list per item.
|
||||
let (ps, _) = requests(r#"{"text": ["a", "b"], "image_data": "u"}"#).unwrap();
|
||||
for p in &ps {
|
||||
assert_eq!(image_of(p).as_array().unwrap().len(), 1);
|
||||
assert!(p.has_multimodal());
|
||||
}
|
||||
|
||||
// Batch + per-item list: element i goes to item i.
|
||||
let (ps, _) = requests(r#"{"text": ["a", "b"], "image_data": ["u1", "u2"]}"#).unwrap();
|
||||
assert_eq!(image_of(&ps[0]).as_str(), Some("u1"));
|
||||
assert_eq!(image_of(&ps[1]).as_str(), Some("u2"));
|
||||
|
||||
// Batch + wrong-length list is a 400.
|
||||
assert!(requests(r#"{"text": ["a", "b"], "image_data": ["u1"]}"#).is_err());
|
||||
|
||||
// null / [] mean "no multimodal input".
|
||||
let (ps, _) = requests(r#"{"text": "a", "image_data": null}"#).unwrap();
|
||||
assert!(!ps[0].has_multimodal());
|
||||
let (ps, _) = requests(r#"{"text": "a", "image_data": []}"#).unwrap();
|
||||
assert!(!ps[0].has_multimodal());
|
||||
|
||||
// Batch + scalar video: broadcast bare (not wrapped), per Python
|
||||
// `_normalize_video_data`.
|
||||
let (ps, _) = requests(r#"{"text": ["a", "b"], "video_data": "v"}"#).unwrap();
|
||||
let video = ps[1].mm.as_ref().unwrap().video_data.clone().unwrap();
|
||||
assert_eq!(video.as_str(), Some("v"));
|
||||
assert!(ps[1].has_multimodal());
|
||||
}
|
||||
|
||||
/// A scalar broadcast is budget-checked before the deep clones (16 MiB ×
|
||||
/// 4096 prompts would be 64 GiB and an abort); per-item lists clone nothing
|
||||
/// and are never charged.
|
||||
#[test]
|
||||
fn oversized_mm_broadcast_rejected() {
|
||||
let big = rmpv::Value::from("x".repeat(MAX_BROADCAST_CLONE_BYTES / 2 + 1));
|
||||
let err = split_mm_column(Some(big.clone()), 2, true, MmBroadcast::WrapInList)
|
||||
.err()
|
||||
.unwrap();
|
||||
assert!(err.contains("broadcast"), "{err}");
|
||||
// A per-item list of the same total size moves, not clones: accepted.
|
||||
let list = rmpv::Value::Array(vec![big, rmpv::Value::from("y")]);
|
||||
assert!(split_mm_column(Some(list), 2, true, MmBroadcast::WrapInList).is_ok());
|
||||
// Small scalars broadcast fine.
|
||||
let small = rmpv::Value::from("u1");
|
||||
assert!(split_mm_column(Some(small), 2, true, MmBroadcast::AsIs).is_ok());
|
||||
}
|
||||
|
||||
/// `mm_hashes` rides only on single requests (Python `__getitem__`
|
||||
/// parity: batches drop it) and moves into the work item.
|
||||
#[test]
|
||||
fn mm_hashes_single_only() {
|
||||
let (mut ps, _) =
|
||||
requests(r#"{"text": "a", "image_data": "u", "mm_hashes": ["a1b2", "0xff"]}"#).unwrap();
|
||||
assert_eq!(ps[0].mm.as_ref().unwrap().mm_hashes, vec!["a1b2", "0xff"]);
|
||||
assert_eq!(ps[0].take_mm_work().mm_hashes, vec!["a1b2", "0xff"]);
|
||||
assert!(ps[0].mm.as_ref().unwrap().mm_hashes.is_empty());
|
||||
|
||||
// A batch cannot carry hashes (Python drops them), so it is rejected...
|
||||
for body in [
|
||||
r#"{"text": ["a", "b"], "image_data": ["u", "v"], "mm_hashes": [["x"], ["y"]]}"#,
|
||||
r#"{"text": ["a", "b"], "image_data": ["u", "v"], "mm_hashes": ["x", "y"]}"#,
|
||||
] {
|
||||
let err = requests(body).err().unwrap();
|
||||
assert!(matches!(err, Error::Validation(_)), "{body}: {err:?}");
|
||||
}
|
||||
// ...while an absent or empty field is not a payload and must still pass.
|
||||
for body in [
|
||||
r#"{"text": ["a", "b"], "image_data": ["u", "v"], "mm_hashes": null}"#,
|
||||
r#"{"text": ["a", "b"], "image_data": ["u", "v"], "mm_hashes": []}"#,
|
||||
] {
|
||||
assert!(requests(body).is_ok(), "{body}");
|
||||
}
|
||||
}
|
||||
|
||||
/// `take_mm_work` clones `text` (the scheduler header still needs it) and
|
||||
/// moves everything the worker owns out of the request.
|
||||
#[test]
|
||||
fn mm_work_item_takes_owned_fields() {
|
||||
let (mut ps, _) =
|
||||
requests(r#"{"text": "hi", "image_data": ["u1", "u2"], "audio_data": "a"}"#).unwrap();
|
||||
let work = ps[0].take_mm_work();
|
||||
assert_eq!(work.text.as_deref(), Some("hi"));
|
||||
assert!(work.input_ids.is_none());
|
||||
assert_eq!(work.image_data.unwrap().as_array().unwrap().len(), 2);
|
||||
assert!(work.video_data.is_none());
|
||||
assert_eq!(work.audio_data.unwrap().as_str(), Some("a"));
|
||||
// Moved out, not cloned; `text` survives for the header.
|
||||
assert!(ps[0].mm.as_ref().unwrap().image_data.is_none());
|
||||
assert_eq!(ps[0].text.as_deref(), Some("hi"));
|
||||
}
|
||||
|
||||
/// The body limit is disabled, so an unbounded batch turns a small body into an
|
||||
|
||||
@@ -0,0 +1,413 @@
|
||||
//! Multimodal worker pool.
|
||||
//!
|
||||
//! Rust threads drain requests parked in `Encoding` and run the `sglang-mm`
|
||||
//! pipeline registered by `Server.start_mm_workers` (decode → preprocess →
|
||||
//! placeholder expansion → M-RoPE, GIL-free). Each worker parks the result
|
||||
//! buffers in the rid-keyed [`Sidecar`] and returns only the expanded ids;
|
||||
//! Python attaches the buffers at drain time (`Server.take_mm`). Inputs the
|
||||
//! pipeline cannot serve are rejected to the client — no Python fallback.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use crate::message::MmRequest;
|
||||
use crate::runtime::Runnable;
|
||||
use crate::tokenizer::TextTokenizer;
|
||||
use crate::tokenizer_manager::TmEvent;
|
||||
|
||||
/// A named POSIX shared-memory segment owning its name: dropped → unlinked.
|
||||
///
|
||||
/// Written by an MM worker so the TP broadcast carries a ~100-byte
|
||||
/// `ShmPointerMMData` stub instead of the ~20 MB feature tensor, and every
|
||||
/// rank maps it in parallel. Python's `materialize()` unlinks after cloning;
|
||||
/// this `Drop` covers the paths where the buffers never reach Python (aborted
|
||||
/// while parked, late result purged).
|
||||
pub struct ShmSegment {
|
||||
name: String,
|
||||
}
|
||||
|
||||
impl ShmSegment {
|
||||
/// Create `/dev/shm/{name}` holding exactly `bytes`. No leading slash —
|
||||
/// the name must suit Python's `SharedMemory(name=…)` (shm_open adds one).
|
||||
pub fn create(name: String, bytes: &[u8]) -> Result<Self, String> {
|
||||
let c_name = std::ffi::CString::new(format!("/{name}"))
|
||||
.map_err(|_| "shm name contains NUL".to_string())?;
|
||||
// SAFETY: plain POSIX calls on a name we own; every handle created
|
||||
// below is closed/unmapped on all paths.
|
||||
unsafe {
|
||||
let fd = libc::shm_open(
|
||||
c_name.as_ptr(),
|
||||
libc::O_CREAT | libc::O_EXCL | libc::O_RDWR,
|
||||
0o600,
|
||||
);
|
||||
if fd < 0 {
|
||||
return Err(format!(
|
||||
"shm_open({name}): {}",
|
||||
std::io::Error::last_os_error()
|
||||
));
|
||||
}
|
||||
let segment = Self { name }; // unlink from here on any failure
|
||||
if libc::ftruncate(fd, bytes.len() as libc::off_t) != 0 {
|
||||
let e = std::io::Error::last_os_error();
|
||||
libc::close(fd);
|
||||
return Err(format!("ftruncate({}): {e}", segment.name));
|
||||
}
|
||||
let ptr = libc::mmap(
|
||||
std::ptr::null_mut(),
|
||||
bytes.len(),
|
||||
libc::PROT_WRITE,
|
||||
libc::MAP_SHARED,
|
||||
fd,
|
||||
0,
|
||||
);
|
||||
libc::close(fd);
|
||||
if ptr == libc::MAP_FAILED {
|
||||
return Err(format!(
|
||||
"mmap({}): {}",
|
||||
segment.name,
|
||||
std::io::Error::last_os_error()
|
||||
));
|
||||
}
|
||||
std::ptr::copy_nonoverlapping(bytes.as_ptr(), ptr.cast::<u8>(), bytes.len());
|
||||
libc::munmap(ptr, bytes.len());
|
||||
Ok(segment)
|
||||
}
|
||||
}
|
||||
|
||||
/// Hand the segment — and the duty to unlink — to the caller (Python, at
|
||||
/// drain time).
|
||||
pub fn into_name(self) -> String {
|
||||
std::mem::take(&mut std::mem::ManuallyDrop::new(self).name)
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for ShmSegment {
|
||||
fn drop(&mut self) {
|
||||
if let Ok(c_name) = std::ffi::CString::new(format!("/{}", self.name)) {
|
||||
// SAFETY: unlinking a name we created; ENOENT (already unlinked
|
||||
// by Python's materialize) is fine to ignore.
|
||||
unsafe { libc::shm_unlink(c_name.as_ptr()) };
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Unique segment names: the pid separates server restarts (a crash can leak
|
||||
/// segments under the old pid), the counter separates results within one.
|
||||
fn shm_name(item: usize) -> String {
|
||||
static COUNTER: AtomicU64 = AtomicU64::new(0);
|
||||
let n = COUNTER.fetch_add(1, Ordering::Relaxed);
|
||||
format!("sglmm-{}-{n}-{item}", std::process::id())
|
||||
}
|
||||
|
||||
/// Python parity: caller hashes override the computed ones so an external
|
||||
/// router's keys align with the prefix cache. A length mismatch or malformed
|
||||
/// entry warns and keeps the computed hash — never blocks the request.
|
||||
fn apply_caller_hashes(hashes: &mut [u64], caller: &[String]) {
|
||||
if caller.is_empty() {
|
||||
return;
|
||||
}
|
||||
if caller.len() != hashes.len() {
|
||||
tracing::warn!(
|
||||
caller = caller.len(),
|
||||
items = hashes.len(),
|
||||
"mm_hashes length != mm item count; ignoring caller hashes"
|
||||
);
|
||||
return;
|
||||
}
|
||||
for (hash, entry) in hashes.iter_mut().zip(caller) {
|
||||
match parse_caller_hash(entry) {
|
||||
Some(v) => *hash = v,
|
||||
None => tracing::warn!(%entry, "malformed mm_hashes entry; keeping computed hash"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Hex of any width, as Python's `int(hex_hash, 16)` takes it (a full SHA-256
|
||||
/// being the common case), keeping the low 64 bits — only the low 30 are
|
||||
/// observable, through `_compute_pad_value`.
|
||||
fn parse_caller_hash(entry: &str) -> Option<u64> {
|
||||
let hex = entry.strip_prefix("0x").unwrap_or(entry);
|
||||
if hex.is_empty() || !hex.bytes().all(|b| b.is_ascii_hexdigit()) {
|
||||
return None;
|
||||
}
|
||||
u64::from_str_radix(&hex[hex.len().saturating_sub(16)..], 16).ok()
|
||||
}
|
||||
|
||||
/// One parked result: the buffers the drain-time Python adapter needs (the
|
||||
/// expanded `input_ids` travel separately, via `TmEvent::MmEncoded`). The qwen
|
||||
/// drain shape (`sglang_mm::qwen_vl::pack_drain`); generalizes to a
|
||||
/// named-tensor handoff once a family needs a different one.
|
||||
pub struct MmSidecarEntry {
|
||||
pub features: FeatureStore,
|
||||
pub grids: Vec<[u32; 3]>,
|
||||
pub hashes: Vec<u64>,
|
||||
pub offsets: Vec<(u32, u32)>,
|
||||
pub mrope: Vec<i64>,
|
||||
pub mrope_delta: i64,
|
||||
}
|
||||
|
||||
/// Where a result's feature buffers live between worker and drain.
|
||||
pub enum FeatureStore {
|
||||
/// In-process; the drain wraps them zero-copy. Single-rank serving, or the
|
||||
/// shm fallback. Under TP the whole buffer would ride `broadcast_pyobj`.
|
||||
Inline(Vec<f32>),
|
||||
/// One POSIX segment per item, written by the worker; only the names cross
|
||||
/// ranks. See [`ShmSegment`].
|
||||
Shm(Vec<ShmSegment>),
|
||||
}
|
||||
|
||||
/// Results parked between a worker's `MmEncoded` and the scheduler drain, keyed
|
||||
/// by rid. Owns the lifecycle so entries never leak: [`park`](Self::park)
|
||||
/// strictly before `MmEncoded`, [`take`](Self::take) at the drain,
|
||||
/// [`purge`](Self::purge) for requests that die while parked.
|
||||
#[derive(Clone, Default)]
|
||||
pub struct Sidecar(Arc<Mutex<HashMap<String, MmSidecarEntry>>>);
|
||||
|
||||
impl Sidecar {
|
||||
pub fn park(&self, rid: String, entry: MmSidecarEntry) {
|
||||
self.0.lock().unwrap().insert(rid, entry);
|
||||
}
|
||||
pub fn take(&self, rid: &str) -> Option<MmSidecarEntry> {
|
||||
self.0.lock().unwrap().remove(rid)
|
||||
}
|
||||
pub fn purge(&self, rid: &str) {
|
||||
self.0.lock().unwrap().remove(rid);
|
||||
}
|
||||
}
|
||||
|
||||
/// Shared state of the mm path, built once at `start_mm_workers`.
|
||||
pub struct Context {
|
||||
pub family: Box<dyn sglang_mm::pipeline::MmFamilyProcessor>,
|
||||
/// `None` under `skip_tokenizer_init` (requests must carry `input_ids`).
|
||||
pub tokenizer: Option<Arc<dyn TextTokenizer>>,
|
||||
pub sidecar: Sidecar,
|
||||
/// Park feature buffers in POSIX shm. Set by the Python launcher
|
||||
/// (`NativeMmHost._use_feature_shm`) exactly when the scheduler broadcasts
|
||||
/// across TP ranks and will unwrap `ShmPointerMMData`.
|
||||
pub feature_shm: bool,
|
||||
}
|
||||
|
||||
impl Context {
|
||||
pub fn new(
|
||||
spec_json: &str,
|
||||
tokenizer: Option<Arc<dyn TextTokenizer>>,
|
||||
sidecar: Sidecar,
|
||||
) -> Result<Self, String> {
|
||||
let feature_shm = serde_json::from_str::<serde_json::Value>(spec_json)
|
||||
.ok()
|
||||
.and_then(|v| v.get("feature_shm").and_then(|b| b.as_bool()))
|
||||
.unwrap_or(false);
|
||||
Ok(Self {
|
||||
family: sglang_mm::registry::pipeline_from_spec(spec_json)?,
|
||||
tokenizer,
|
||||
sidecar,
|
||||
feature_shm,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Run the pipeline for one request. `Ok` returns the final expanded ids, the
|
||||
/// buffers already parked; `Err` rejects the request back to the client.
|
||||
fn process(
|
||||
ctx: &Context,
|
||||
rid: &crate::ids::Rid,
|
||||
mut work: crate::message::MmWorkItem,
|
||||
) -> Result<Vec<i32>, String> {
|
||||
let caller_hashes = std::mem::take(&mut work.mm_hashes);
|
||||
let input = crate::message::mm_payload::to_mm_input(work)?;
|
||||
let output = sglang_mm::driver::process(ctx.family.as_ref(), input, |text| {
|
||||
let tokenizer = ctx.tokenizer.as_ref().ok_or_else(|| {
|
||||
"skip_tokenizer_init is set: multimodal text prompts require input_ids".to_string()
|
||||
})?;
|
||||
tokenizer.encode(text).map_err(|error| error.to_string())
|
||||
})?;
|
||||
let mut drain = sglang_mm::qwen_vl::pack_drain(output)?;
|
||||
apply_caller_hashes(&mut drain.hashes, &caller_hashes);
|
||||
let features = if ctx.feature_shm {
|
||||
park_features_in_shm(&drain.features, &drain.grids)
|
||||
} else {
|
||||
FeatureStore::Inline(drain.features)
|
||||
};
|
||||
ctx.sidecar.park(
|
||||
rid.as_str().to_owned(),
|
||||
MmSidecarEntry {
|
||||
features,
|
||||
grids: drain.grids,
|
||||
hashes: drain.hashes,
|
||||
offsets: drain.offsets,
|
||||
mrope: drain.mrope,
|
||||
mrope_delta: drain.mrope_delta,
|
||||
},
|
||||
);
|
||||
Ok(drain.input_ids)
|
||||
}
|
||||
|
||||
/// Split the flat feature buffer per item (`t*h*w` rows per grid) and park each
|
||||
/// slice in its own segment. Any shm failure (`/dev/shm` full, odd shape) falls
|
||||
/// back to inline, as Python's `_wrap_shm_or_inline` does: degrade to the slow
|
||||
/// path, never fail the request.
|
||||
fn park_features_in_shm(features: &[f32], grids: &[[u32; 3]]) -> FeatureStore {
|
||||
let total_rows: usize = grids
|
||||
.iter()
|
||||
.map(|g| g[0] as usize * g[1] as usize * g[2] as usize)
|
||||
.sum();
|
||||
if total_rows == 0 || !features.len().is_multiple_of(total_rows) {
|
||||
return FeatureStore::Inline(features.to_vec());
|
||||
}
|
||||
let dim = features.len() / total_rows;
|
||||
let mut segments = Vec::with_capacity(grids.len());
|
||||
let mut row = 0usize;
|
||||
for (item, grid) in grids.iter().enumerate() {
|
||||
let rows = grid[0] as usize * grid[1] as usize * grid[2] as usize;
|
||||
let slice = &features[row * dim..(row + rows) * dim];
|
||||
row += rows;
|
||||
match ShmSegment::create(shm_name(item), bytemuck::cast_slice(slice)) {
|
||||
Ok(segment) => segments.push(segment),
|
||||
Err(error) => {
|
||||
tracing::warn!(%error, "mm: shm feature transport failed; falling back to inline");
|
||||
return FeatureStore::Inline(features.to_vec());
|
||||
}
|
||||
}
|
||||
}
|
||||
FeatureStore::Shm(segments)
|
||||
}
|
||||
|
||||
/// One MM worker, spawned via `Runtime::spawn_mm_pool` (which owns the
|
||||
/// pinning policy for this pool — see its docs).
|
||||
pub struct MmWorker {
|
||||
rx: flume::Receiver<MmRequest>,
|
||||
tm: flume::Sender<TmEvent>,
|
||||
ctx: Arc<Context>,
|
||||
}
|
||||
|
||||
impl MmWorker {
|
||||
pub fn new(
|
||||
rx: flume::Receiver<MmRequest>,
|
||||
tm: flume::Sender<TmEvent>,
|
||||
ctx: Arc<Context>,
|
||||
) -> Self {
|
||||
Self { rx, tm, ctx }
|
||||
}
|
||||
}
|
||||
|
||||
impl Runnable for MmWorker {
|
||||
/// Drain until the mm channel closes (tm-ingress drops its sender on
|
||||
/// shutdown). One request at a time, so the pool size bounds MM
|
||||
/// concurrency; an error rejects the request back to the client.
|
||||
fn run(self) {
|
||||
while let Ok(req) = self.rx.recv() {
|
||||
let rid = req.rid;
|
||||
let event = match process(&self.ctx, &rid, req.work) {
|
||||
Ok(input_ids) => {
|
||||
tracing::debug!(%rid, tokens = input_ids.len(), "mm: processed");
|
||||
TmEvent::MmEncoded { rid, input_ids }
|
||||
}
|
||||
Err(message) => {
|
||||
tracing::warn!(%rid, %message, "mm processing rejected");
|
||||
TmEvent::MmFailed { rid, message }
|
||||
}
|
||||
};
|
||||
if self.tm.send(event).is_err() {
|
||||
return; // tm-ingress gone: shutdown
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// Caller hashes override computed ones; mismatched lengths and malformed
|
||||
/// entries fall back per item, never reject (Python parity).
|
||||
#[test]
|
||||
fn caller_hashes_override_with_fallback() {
|
||||
let mut hashes = vec![1, 2, 3];
|
||||
apply_caller_hashes(&mut hashes, &[]);
|
||||
assert_eq!(hashes, [1, 2, 3]);
|
||||
|
||||
apply_caller_hashes(&mut hashes, &["ff".into()]); // length mismatch
|
||||
assert_eq!(hashes, [1, 2, 3]);
|
||||
|
||||
apply_caller_hashes(&mut hashes, &["ff".into(), "not-hex".into(), "0x10".into()]);
|
||||
assert_eq!(hashes, [0xff, 2, 0x10]);
|
||||
}
|
||||
|
||||
/// A full SHA-256 (what routers send) keeps its low 64 bits rather than
|
||||
/// falling back, so the pad value matches Python's wide `int`.
|
||||
#[test]
|
||||
fn caller_hashes_accept_arbitrary_width() {
|
||||
let sha256 = "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855";
|
||||
let mut hashes = vec![1];
|
||||
apply_caller_hashes(&mut hashes, &[sha256.into()]);
|
||||
assert_eq!(hashes, [0xa495991b7852b855]);
|
||||
assert_eq!(hashes[0] % (1 << 30), 944_945_237); // int(sha256, 16) % (1 << 30)
|
||||
|
||||
// Width alone is never malformed; a non-hex digit still is.
|
||||
assert_eq!(parse_caller_hash(&"f".repeat(64)), Some(u64::MAX));
|
||||
assert_eq!(parse_caller_hash("0x"), None);
|
||||
assert_eq!(parse_caller_hash(""), None);
|
||||
}
|
||||
|
||||
fn shm_path(name: &str) -> std::path::PathBuf {
|
||||
std::path::Path::new("/dev/shm").join(name)
|
||||
}
|
||||
|
||||
/// The segment holds exactly the written bytes and dropping it unlinks —
|
||||
/// the leak guard for results purged before Python takes them.
|
||||
#[test]
|
||||
fn segment_roundtrip_and_drop_unlinks() {
|
||||
let name = shm_name(0);
|
||||
let payload: Vec<u8> = (0..255u8).collect();
|
||||
let segment = ShmSegment::create(name.clone(), &payload).unwrap();
|
||||
assert_eq!(std::fs::read(shm_path(&name)).unwrap(), payload);
|
||||
drop(segment);
|
||||
assert!(!shm_path(&name).exists(), "drop must unlink");
|
||||
}
|
||||
|
||||
/// `into_name` transfers the unlink duty to the caller (Python's
|
||||
/// `materialize()`), so the segment must survive the handoff.
|
||||
#[test]
|
||||
fn into_name_disarms_the_unlink() {
|
||||
let segment = ShmSegment::create(shm_name(0), &[1, 2, 3]).unwrap();
|
||||
let name = segment.into_name();
|
||||
assert!(shm_path(&name).exists(), "handoff must not unlink");
|
||||
// manual cleanup for the test
|
||||
let c = std::ffi::CString::new(format!("/{name}")).unwrap();
|
||||
unsafe { libc::shm_unlink(c.as_ptr()) };
|
||||
}
|
||||
|
||||
/// Per-item slicing follows the grid row counts, so Python's
|
||||
/// `(rows, feature_dim)` reshape of a segment sees only its own item.
|
||||
#[test]
|
||||
fn park_splits_features_by_grid() {
|
||||
// Two items: grids (1,2,2)=4 rows and (1,1,2)=2 rows, dim=3.
|
||||
let features: Vec<f32> = (0..18).map(|i| i as f32).collect();
|
||||
let grids = [[1, 2, 2], [1, 1, 2]];
|
||||
let FeatureStore::Shm(segments) = park_features_in_shm(&features, &grids) else {
|
||||
panic!("expected shm store");
|
||||
};
|
||||
assert_eq!(segments.len(), 2);
|
||||
let read = |seg: &ShmSegment| -> Vec<u8> { std::fs::read(shm_path(&seg.name)).unwrap() };
|
||||
assert_eq!(
|
||||
read(&segments[0]),
|
||||
bytemuck::cast_slice::<f32, u8>(&features[..12])
|
||||
);
|
||||
assert_eq!(
|
||||
read(&segments[1]),
|
||||
bytemuck::cast_slice::<f32, u8>(&features[12..])
|
||||
);
|
||||
}
|
||||
|
||||
/// A degenerate shape must degrade to inline, never a shm-side panic.
|
||||
#[test]
|
||||
fn shape_surprise_falls_back_inline() {
|
||||
let features = vec![0.0f32; 7]; // not divisible by 2 rows
|
||||
let grids = [[1, 1, 2]];
|
||||
assert!(matches!(
|
||||
park_features_in_shm(&features, &grids),
|
||||
FeatureStore::Inline(_)
|
||||
));
|
||||
}
|
||||
}
|
||||
@@ -8,6 +8,8 @@
|
||||
//! * Detokenizer — M pinned OS threads / shards (CPU bound), core set C
|
||||
//! * TM ingress — 1 thread driving the ingress FSM
|
||||
//! * TM egress — 1 thread draining the egress ring → detok shards
|
||||
//! * MM workers — K unpinned OS threads, spawned late via
|
||||
//! [`Runtime::spawn_mm_pool`] (multimodal models only)
|
||||
//!
|
||||
//! Keeping CPU-bound tokenize/detokenize off the async executor avoids stalling
|
||||
//! axum's worker threads.
|
||||
@@ -38,6 +40,18 @@ pub use runnable::Runnable;
|
||||
pub struct Runtime {
|
||||
pub ingress: IngressConsumer,
|
||||
pub egress: EgressProducer,
|
||||
/// Requests parked in `Encoding`, drained by the MM worker pool
|
||||
/// (`Server.start_mm_workers`). Stays empty for non-multimodal models —
|
||||
/// ingress never routes to it.
|
||||
pub mm: flume::Receiver<crate::message::MmRequest>,
|
||||
/// Back-channel for the MM workers' `MmEncoded` / `MmFailed` into tm-ingress.
|
||||
pub tm: flume::Sender<TmEvent>,
|
||||
/// The loaded tokenizer, shared with the MM worker path (`None` under
|
||||
/// `skip_tokenizer_init`).
|
||||
pub tokenizer: Option<Arc<dyn tokenizer::TextTokenizer>>,
|
||||
/// MM results parked between a worker's `MmEncoded` and the scheduler drain
|
||||
/// (`Server.take_mm`).
|
||||
pub mm_sidecar: crate::mm::Sidecar,
|
||||
/// Worker join handles, joined by `request_shutdown` / `Drop`.
|
||||
threads: Mutex<Vec<JoinHandle<()>>>,
|
||||
/// The single shutdown sender.
|
||||
@@ -49,6 +63,21 @@ pub struct Runtime {
|
||||
const SHUTDOWN_JOIN_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5);
|
||||
|
||||
impl Runtime {
|
||||
/// Spawn `workers` `mm-worker-{i}` threads into the shutdown join set —
|
||||
/// late, once Python has built the mm spec (`Server::start_mm_workers`).
|
||||
///
|
||||
/// Deliberately unpinned: the threads inherit the launch thread's affinity,
|
||||
/// already narrowed by `RustServer.launch` to the server cores, so bursty
|
||||
/// MM preprocessing floats over that whole set (rather than owning cores
|
||||
/// that idle between bursts) and never preempts the scheduler's reserved
|
||||
/// cores.
|
||||
pub fn spawn_mm_pool(&self, workers: usize, ctx: Arc<crate::mm::Context>) {
|
||||
let mut threads = self.threads.lock().unwrap();
|
||||
spawn_pool("mm-worker", None, workers.max(1), &mut threads, |_| {
|
||||
crate::mm::MmWorker::new(self.mm.clone(), self.tm.clone(), ctx.clone())
|
||||
});
|
||||
}
|
||||
|
||||
/// Stop the runtime and join every worker thread (with a bounded wait).
|
||||
///
|
||||
/// Dropping `shutdown_tx` wakes the tm-ingress/tm-egress selectors (which
|
||||
@@ -107,6 +136,10 @@ pub fn start(cfg: RuntimeConfig) -> Result<Runtime, String> {
|
||||
let (tm_tx, tm_rx) = flume::bounded::<TmEvent>(cfg.rust_server_args.channel_cap);
|
||||
let (tok_tx, tok_rx) =
|
||||
flume::bounded::<crate::message::Request>(cfg.rust_server_args.channel_cap);
|
||||
// Encoding → MM worker pool. Bounded like the other stage edges so a slow
|
||||
// pool back-pressures instead of buffering unboundedly.
|
||||
let (mm_tx, mm_rx) =
|
||||
flume::bounded::<crate::message::MmRequest>(cfg.rust_server_args.channel_cap);
|
||||
let detokenizer_worker_num = cfg.server_args.detokenizer_worker_num;
|
||||
let mut detok_tx = Vec::with_capacity(detokenizer_worker_num);
|
||||
let mut detok_rx = Vec::with_capacity(detokenizer_worker_num);
|
||||
@@ -140,6 +173,14 @@ pub fn start(cfg: RuntimeConfig) -> Result<Runtime, String> {
|
||||
cfg.server_args.revision.as_deref(),
|
||||
skip_tokenizer_init,
|
||||
)?;
|
||||
// The `TextTokenizer` view of it, shared by the tokenizer pool and the MM
|
||||
// worker path (which encodes the placeholder-expanded prompt itself).
|
||||
let text_tokenizer: Option<Arc<dyn tokenizer::TextTokenizer>> = dyn_tokenizer
|
||||
.as_ref()
|
||||
.map(|t| Arc::new(tokenizer::DynamoTokenizer::new(t.clone())) as _);
|
||||
|
||||
// Shared: MM workers park, the Python drain pops, tm-ingress purges.
|
||||
let mm_sidecar: crate::mm::Sidecar = Default::default();
|
||||
|
||||
// --- Detokenizer shards (pinned, CPU bound) ---
|
||||
{
|
||||
@@ -168,10 +209,9 @@ pub fn start(cfg: RuntimeConfig) -> Result<Runtime, String> {
|
||||
// --- Tokenizer pool (pinned, CPU bound) ---
|
||||
// Only spawned when a real tokenizer is loaded; under `skip_tokenizer_init`
|
||||
// there is none and ingress never routes to the pool, so we skip it.
|
||||
if let Some(t) = &dyn_tokenizer {
|
||||
if let Some(tokenizer) = &text_tokenizer {
|
||||
// Reuse the single loaded tokenizer (shared with the detok shards).
|
||||
let tokenizer: Arc<dyn tokenizer::TextTokenizer> =
|
||||
Arc::new(tokenizer::DynamoTokenizer::new(t.clone()));
|
||||
let tokenizer = tokenizer.clone();
|
||||
let tok_cores = plan.as_ref().map(|p| p.tok.clone());
|
||||
// Workers share the MPMC inbox (`tok_rx`) and the read-only backend, so
|
||||
// each gets a cheap clone of both.
|
||||
@@ -220,6 +260,11 @@ pub fn start(cfg: RuntimeConfig) -> Result<Runtime, String> {
|
||||
.map(|c| vec![c]);
|
||||
let limits = tokenizer_manager::Limits::try_from(&*cfg.server_args)
|
||||
.map_err(|e| format!("ingress limits: {e}"))?;
|
||||
let mm = tokenizer_manager::Mm {
|
||||
enabled: cfg.server_args.model_is_multimodal(),
|
||||
tx: mm_tx,
|
||||
sidecar: mm_sidecar.clone(),
|
||||
};
|
||||
let mut parts = Some((tm_rx, ingress_tx)); // moved into the single worker
|
||||
let shutdown_rx = shutdown_rx.clone();
|
||||
spawn_pool("tm-ingress", cores, 1, &mut threads, |_| {
|
||||
@@ -230,6 +275,7 @@ pub fn start(cfg: RuntimeConfig) -> Result<Runtime, String> {
|
||||
senders.clone(),
|
||||
ingress_tx,
|
||||
limits.clone(),
|
||||
mm.clone(),
|
||||
shutdown_rx.clone(),
|
||||
)
|
||||
});
|
||||
@@ -282,6 +328,10 @@ pub fn start(cfg: RuntimeConfig) -> Result<Runtime, String> {
|
||||
Ok(Runtime {
|
||||
ingress: ingress_rx,
|
||||
egress: egress_tx,
|
||||
mm: mm_rx,
|
||||
tm: tm_tx,
|
||||
tokenizer: text_tokenizer,
|
||||
mm_sidecar,
|
||||
threads: Mutex::new(threads),
|
||||
shutdown_tx: Mutex::new(Some(shutdown_tx)),
|
||||
})
|
||||
|
||||
@@ -158,6 +158,11 @@ pub struct ModelConfig {
|
||||
/// boot ([`ServerArgs::validate_mandatory`]).
|
||||
#[serde(default)]
|
||||
pub vocab_size: Option<u64>,
|
||||
/// Whether the model accepts multimodal inputs. Gates the MM Encoding branch
|
||||
/// in tm-ingress; `false` silently ignores mm fields, as the Python
|
||||
/// `TokenizerManager` does with `mm_processor is None`.
|
||||
#[serde(default)]
|
||||
pub is_multimodal: bool,
|
||||
/// Resolved default sampling parameters, stamped by
|
||||
/// `RustServer._build_server_args` from Python's
|
||||
/// `ModelConfig.get_default_sampling_params()`. Already gated on
|
||||
@@ -259,6 +264,12 @@ impl ServerArgs {
|
||||
self.disaggregation_mode == "prefill"
|
||||
}
|
||||
|
||||
/// Whether the served model is multimodal, from the scheduler's dump. See
|
||||
/// [`ModelConfig::is_multimodal`].
|
||||
pub fn model_is_multimodal(&self) -> bool {
|
||||
self.model_config.is_multimodal
|
||||
}
|
||||
|
||||
/// Bind address `host:port`. `host` is expected to be an IP — the result is
|
||||
/// parsed as a `SocketAddr`, so a bare IPv6 host gets bracketed.
|
||||
pub fn bind(&self) -> String {
|
||||
|
||||
@@ -14,7 +14,7 @@ mod egress;
|
||||
mod ingress;
|
||||
|
||||
pub use egress::{ActivityCounter, Egress};
|
||||
pub use ingress::{Ingress, Limits};
|
||||
pub use ingress::{Ingress, Limits, Mm};
|
||||
|
||||
use crate::ids::Rid;
|
||||
use crate::message::{DetokMsg, Request};
|
||||
@@ -36,6 +36,13 @@ pub enum TmEvent {
|
||||
/// A request back from the tokenizer pool: `PreSendValidating` (ids filled) on success,
|
||||
/// or `Failed` on a tokenize error. `drive` handles both.
|
||||
Tokenized(Request),
|
||||
/// An MM worker finished a request parked in `Encoding`: `input_ids` are the
|
||||
/// final placeholder-expanded prompt ids. The buffers ride the rid-keyed
|
||||
/// sidecar (`Server.take_mm`), not this event.
|
||||
MmEncoded { rid: Rid, input_ids: Vec<i32> },
|
||||
/// An MM worker rejected a request parked in `Encoding` (bad media URL,
|
||||
/// unsupported modality, preprocess error, …).
|
||||
MmFailed { rid: Rid, message: String },
|
||||
}
|
||||
|
||||
/// Producer-side handles, cloned into every stage that needs to emit.
|
||||
|
||||
@@ -17,14 +17,17 @@
|
||||
//! The egress edges (Streaming/Finalizing/Completed) are driven on the egress
|
||||
//! side (see `egress` + `detokenizer`).
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use bytes::Bytes;
|
||||
|
||||
use crate::error::Error;
|
||||
use crate::fsm::{Event, RequestState, ValidationOutcome};
|
||||
use crate::ids::Rid;
|
||||
|
||||
use crate::message::{
|
||||
AbortReq, ControlRequest, DetokMsg, EgressItem, GenerateRequest, IngressMsg, Request,
|
||||
RequestKind,
|
||||
AbortReq, ControlRequest, DetokMsg, EgressItem, GenerateRequest, IngressMsg, MmRequest,
|
||||
Request, RequestKind,
|
||||
};
|
||||
use crate::ring::IngressProducer;
|
||||
use crate::runtime::{Runnable, ServerArgs};
|
||||
@@ -41,9 +44,29 @@ pub struct Ingress {
|
||||
senders: Senders,
|
||||
ingress: IngressProducer,
|
||||
limits: Limits,
|
||||
mm: Mm,
|
||||
/// Requests parked in `Encoding` while an MM worker processes their media;
|
||||
/// resumed by `MmEncoded` / `MmFailed`. Only this thread touches it, so no
|
||||
/// lock.
|
||||
pending_mm: HashMap<Rid, Request>,
|
||||
shutdown: flume::Receiver<()>,
|
||||
}
|
||||
|
||||
/// The ingress side of the MM path.
|
||||
#[derive(Clone)]
|
||||
pub struct Mm {
|
||||
/// Whether the model is multimodal. When false, mm fields are silently
|
||||
/// ignored, as the Python `TokenizerManager` does with `mm_processor is
|
||||
/// None`.
|
||||
pub enabled: bool,
|
||||
/// → MM worker pool (spawned via `Server.start_mm_workers`).
|
||||
pub tx: flume::Sender<MmRequest>,
|
||||
/// Results sidecar. Purged here when a late result arrives for a request
|
||||
/// that is no longer parked; otherwise it would leak, since only the
|
||||
/// scheduler drain pops entries.
|
||||
pub sidecar: crate::mm::Sidecar,
|
||||
}
|
||||
|
||||
/// Longest client-supplied rid accepted. It keys the detok table and travels on
|
||||
/// every chunk, so its length is a recurring cost; Python mints 32-byte uuid hex.
|
||||
const MAX_RID_LEN: usize = 128;
|
||||
@@ -104,6 +127,7 @@ impl Ingress {
|
||||
senders: Senders,
|
||||
ingress: IngressProducer,
|
||||
limits: Limits,
|
||||
mm: Mm,
|
||||
shutdown: flume::Receiver<()>,
|
||||
) -> Self {
|
||||
Self {
|
||||
@@ -112,6 +136,8 @@ impl Ingress {
|
||||
senders,
|
||||
ingress,
|
||||
limits,
|
||||
mm,
|
||||
pending_mm: HashMap::new(),
|
||||
shutdown,
|
||||
}
|
||||
}
|
||||
@@ -124,7 +150,7 @@ enum Lane {
|
||||
}
|
||||
|
||||
impl Runnable for Ingress {
|
||||
fn run(self) {
|
||||
fn run(mut self) {
|
||||
loop {
|
||||
// Select, not a drain-then-block: an abort arriving while the inbox is
|
||||
// idle must still be handled at once.
|
||||
@@ -139,6 +165,12 @@ impl Runnable for Ingress {
|
||||
Some(Lane::Event(TmEvent::Ingress(req) | TmEvent::Tokenized(req))) => {
|
||||
self.drive(req)
|
||||
}
|
||||
Some(Lane::Event(TmEvent::MmEncoded { rid, input_ids })) => {
|
||||
self.on_mm_encoded(rid, input_ids)
|
||||
}
|
||||
Some(Lane::Event(TmEvent::MmFailed { rid, message })) => {
|
||||
self.on_mm_failed(rid, message)
|
||||
}
|
||||
None => {
|
||||
// Shutdown, or the inbox closed. Drain whatever is still queued
|
||||
// on the abort lane first: those requests are in flight on the
|
||||
@@ -168,6 +200,9 @@ impl Ingress {
|
||||
if err.http_status() == 500 {
|
||||
tracing::error!(rid = %req.rid, error = %err, "ingress rejected request");
|
||||
}
|
||||
// A rejected request never reaches the scheduler drain, so purge any
|
||||
// parked MM result (no-op for the common non-mm request).
|
||||
self.mm.sidecar.purge(req.rid.as_str());
|
||||
let _ = req.state.apply(Event::Error(err.clone()));
|
||||
let _ = req.sink.try_send(EgressItem::Error(err)); // client may be gone
|
||||
if registered {
|
||||
@@ -178,11 +213,12 @@ impl Ingress {
|
||||
}
|
||||
|
||||
/// Drive a request through its ingress states until it terminates (failed or
|
||||
/// pushed to the ring) or is handed to the tokenizer pool (re-entering as a
|
||||
/// `Tokenized` event). Each arm acts and advances the FSM; the loop
|
||||
/// re-dispatches. The arms are the design table's states, `Failed` the single
|
||||
/// reject path.
|
||||
fn drive(&self, mut req: Request) {
|
||||
/// pushed to the ring), is handed to the tokenizer pool (re-entering as a
|
||||
/// `Tokenized` event), or is parked in `pending_mm` awaiting an MM worker
|
||||
/// (re-entering via `MmEncoded` / `MmFailed`). Each arm acts and advances
|
||||
/// the FSM; the loop re-dispatches. The arms are the design table's states,
|
||||
/// `Failed` the single reject path.
|
||||
fn drive(&mut self, mut req: Request) {
|
||||
// Flipped once `register_detok` succeeds; `fail` must not deregister before
|
||||
// that (see `fail`). A pool return re-enters `drive` already registered.
|
||||
let mut registered = !matches!(req.state, RequestState::Received);
|
||||
@@ -235,6 +271,13 @@ impl Ingress {
|
||||
.normalize(self.limits.skip_tokenizer_init, self.limits.vocab_size)
|
||||
{
|
||||
Err(e) => Err(e),
|
||||
// The native pipeline produces the final input_ids,
|
||||
// so it wins even over a pre-tokenized prompt (which
|
||||
// still needs placeholder expansion) — the same
|
||||
// precedence as the Python TokenizerManager.
|
||||
Ok(()) if self.mm.enabled && g.has_multimodal() => {
|
||||
Ok(ValidationOutcome::HasMultimodal)
|
||||
}
|
||||
// Client ids skip the pool; text goes to the tokenizer.
|
||||
Ok(()) if g.already_tokenized() => {
|
||||
Ok(ValidationOutcome::AlreadyTokenized)
|
||||
@@ -252,6 +295,40 @@ impl Ingress {
|
||||
}
|
||||
}
|
||||
}
|
||||
// Hand off to the MM worker pool and park the request; it
|
||||
// re-enters via `MmEncoded` (→ PreSendValidating) or `MmFailed`
|
||||
// (→ reject). Doesn't loop.
|
||||
RequestState::Encoding => {
|
||||
let work = {
|
||||
let RequestKind::Generate(g) = &mut req.kind else {
|
||||
self.fail(
|
||||
&mut req,
|
||||
Error::Internal("non-generate request in Encoding".into()),
|
||||
registered,
|
||||
);
|
||||
return;
|
||||
};
|
||||
g.take_mm_work()
|
||||
};
|
||||
let msg = MmRequest {
|
||||
rid: req.rid.clone(),
|
||||
work,
|
||||
};
|
||||
// Full = the pool can't keep up, so back-pressure like a full
|
||||
// ingress ring. Disconnected = pool gone.
|
||||
if let Err(e) = self.mm.tx.try_send(msg) {
|
||||
let err = match e {
|
||||
flume::TrySendError::Full(_) => Error::QueueFull,
|
||||
flume::TrySendError::Disconnected(_) => {
|
||||
Error::Internal("mm worker pool gone".into())
|
||||
}
|
||||
};
|
||||
self.fail(&mut req, err, registered);
|
||||
return;
|
||||
}
|
||||
self.pending_mm.insert(req.rid.clone(), req);
|
||||
return;
|
||||
}
|
||||
// Hand off to the tokenizer pool; it returns the request as a
|
||||
// `Tokenized` event (PreSendValidating, or Failed on error).
|
||||
// Doesn't loop.
|
||||
@@ -392,6 +469,34 @@ impl Ingress {
|
||||
}
|
||||
}
|
||||
|
||||
/// An MM worker finished a parked request: fill in the final expanded
|
||||
/// `input_ids`, advance `Encoding → PreSendValidating`, and resume driving
|
||||
/// (pre-send checks → ring). No pending entry means the request was already
|
||||
/// rejected or aborted, so the result is dropped.
|
||||
fn on_mm_encoded(&mut self, rid: Rid, input_ids: Vec<i32>) {
|
||||
let Some(mut req) = self.pending_mm.remove(&rid) else {
|
||||
tracing::debug!(rid = %rid, "mm result for unknown/finished request; dropped");
|
||||
// It will never reach the scheduler drain, so purge or leak.
|
||||
self.mm.sidecar.purge(rid.as_str());
|
||||
return;
|
||||
};
|
||||
if let RequestKind::Generate(g) = &mut req.kind {
|
||||
g.input_ids = Some(input_ids);
|
||||
}
|
||||
let _ = req.state.apply(Event::EncodeDone); // Encoding → PreSendValidating
|
||||
self.drive(req);
|
||||
}
|
||||
|
||||
/// An MM worker failed a parked request (bad URL, processor error): reject it
|
||||
/// back to the client, as Python turns a per-request exception into a 400.
|
||||
fn on_mm_failed(&mut self, rid: Rid, message: String) {
|
||||
let Some(mut req) = self.pending_mm.remove(&rid) else {
|
||||
tracing::debug!(rid = %rid, "mm failure for unknown/finished request; dropped");
|
||||
return;
|
||||
};
|
||||
self.fail(&mut req, Error::Encode(message), true); // parked ⇒ registered
|
||||
}
|
||||
|
||||
/// Client disconnected (or a detok terminal): deregister the sink, then push an
|
||||
/// `AbortReq(rid)` so the scheduler stops generating for it.
|
||||
///
|
||||
@@ -400,8 +505,15 @@ impl Ingress {
|
||||
/// That wastes GPU work until the request finishes on its own, but it cannot be
|
||||
/// misdelivered — the rid is unique to this request for the process's lifetime
|
||||
/// ([`Rid::from_client`]), so no later request can ever answer to it.
|
||||
fn on_abort(&self, source: AbortSource) {
|
||||
///
|
||||
/// A request parked in `pending_mm` is cancelled here, so the worker's late
|
||||
/// result lands in `on_mm_encoded`'s no-entry branch and purges the sidecar —
|
||||
/// no generation runs for output nobody will read.
|
||||
fn on_abort(&mut self, source: AbortSource) {
|
||||
let rid = source.rid().clone();
|
||||
if self.pending_mm.remove(&rid).is_some() {
|
||||
tracing::debug!(rid = %rid, "abort cancelled request parked for MM");
|
||||
}
|
||||
let _ = self
|
||||
.senders
|
||||
.detok_for(&rid)
|
||||
@@ -616,12 +728,14 @@ mod tests {
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
/// An `Ingress` plus its detok-shard receiver, ring consumer (keep alive —
|
||||
/// dropping it closes the ring → false QueueFull), and tm inbox sender.
|
||||
/// dropping it closes the ring → false QueueFull), tm inbox sender, and the
|
||||
/// mm-pool receiver (keep alive — dropping it makes mm submits fail).
|
||||
fn make_ingress() -> (
|
||||
Ingress,
|
||||
flume::Receiver<DetokMsg>,
|
||||
IngressConsumer,
|
||||
flume::Sender<TmEvent>,
|
||||
flume::Receiver<MmRequest>,
|
||||
) {
|
||||
make_ingress_with(test_limits())
|
||||
}
|
||||
@@ -633,6 +747,7 @@ mod tests {
|
||||
flume::Receiver<DetokMsg>,
|
||||
IngressConsumer,
|
||||
flume::Sender<TmEvent>,
|
||||
flume::Receiver<MmRequest>,
|
||||
) {
|
||||
make_ingress_inner(test_limits(), abort_rx)
|
||||
}
|
||||
@@ -644,6 +759,7 @@ mod tests {
|
||||
flume::Receiver<DetokMsg>,
|
||||
IngressConsumer,
|
||||
flume::Sender<TmEvent>,
|
||||
flume::Receiver<MmRequest>,
|
||||
) {
|
||||
let (abort_tx, abort_rx) = flume::unbounded::<AbortSource>();
|
||||
std::mem::forget(abort_tx); // keep the lane open; tests end by dropping tm_tx
|
||||
@@ -658,6 +774,7 @@ mod tests {
|
||||
flume::Receiver<DetokMsg>,
|
||||
IngressConsumer,
|
||||
flume::Sender<TmEvent>,
|
||||
flume::Receiver<MmRequest>,
|
||||
) {
|
||||
let (tok_tx, _tok_rx) = flume::unbounded();
|
||||
let (detok_tx, detok_rx) = flume::unbounded();
|
||||
@@ -669,12 +786,30 @@ mod tests {
|
||||
};
|
||||
let (ingress_producer, consumer) = ingress_ring(16);
|
||||
let (tm_tx, tm_rx) = flume::unbounded();
|
||||
let (mm_tx, mm_rx) = flume::unbounded();
|
||||
// Keep the shutdown sender alive (leak) so its branch never fires — tests
|
||||
// end `run` by dropping `tm_tx`, not by shutdown.
|
||||
let (sd_tx, sd_rx) = flume::unbounded::<()>();
|
||||
std::mem::forget(sd_tx);
|
||||
let ingress = Ingress::new(tm_rx, abort_rx, senders, ingress_producer, limits, sd_rx);
|
||||
(ingress, detok_rx, consumer, tm_tx)
|
||||
let ingress = Ingress::new(
|
||||
tm_rx,
|
||||
abort_rx,
|
||||
senders,
|
||||
ingress_producer,
|
||||
limits,
|
||||
test_mm(mm_tx, true),
|
||||
sd_rx,
|
||||
);
|
||||
(ingress, detok_rx, consumer, tm_tx, mm_rx)
|
||||
}
|
||||
|
||||
/// An [`Mm`] over `tx` with a fresh sidecar.
|
||||
fn test_mm(tx: flume::Sender<MmRequest>, enabled: bool) -> Mm {
|
||||
Mm {
|
||||
enabled,
|
||||
tx,
|
||||
sidecar: Default::default(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Both abort sources do the same two things: drop the detok entry so no
|
||||
@@ -696,7 +831,7 @@ mod tests {
|
||||
let (ingress_producer, consumer) = ingress_ring(16);
|
||||
let (sd_tx, sd_rx) = flume::unbounded::<()>();
|
||||
std::mem::forget(sd_tx);
|
||||
let ingress = Ingress::new(
|
||||
let mut ingress = Ingress::new(
|
||||
flume::unbounded().1,
|
||||
flume::unbounded().1,
|
||||
Senders {
|
||||
@@ -707,6 +842,7 @@ mod tests {
|
||||
},
|
||||
ingress_producer,
|
||||
test_limits(),
|
||||
test_mm(flume::unbounded().0, true),
|
||||
sd_rx,
|
||||
);
|
||||
|
||||
@@ -960,7 +1096,7 @@ mod tests {
|
||||
/// to the ring, after registration — so it must be deregistered, not leaked.
|
||||
#[test]
|
||||
fn over_context_request_deregisters_and_never_reaches_the_ring() {
|
||||
let (ingress, detok_rx, consumer, _tm_tx) = make_ingress_with(Limits {
|
||||
let (mut ingress, detok_rx, consumer, _tm_tx, _mm_rx) = make_ingress_with(Limits {
|
||||
context_len: 4,
|
||||
..test_limits()
|
||||
});
|
||||
@@ -993,7 +1129,7 @@ mod tests {
|
||||
/// pins. Nothing may reach the scheduler ring.
|
||||
#[test]
|
||||
fn detokenize_flows_register_then_decode_and_skips_the_ring() {
|
||||
let (ingress, detok_rx, consumer, _tm_tx) = make_ingress();
|
||||
let (mut ingress, detok_rx, consumer, _tm_tx, _mm_rx) = make_ingress();
|
||||
let (tx, mut rx) = mpsc::channel(8);
|
||||
ingress.drive(Request {
|
||||
rid: "41".into(),
|
||||
@@ -1028,7 +1164,7 @@ mod tests {
|
||||
/// leak and no decode job to drop).
|
||||
#[test]
|
||||
fn detokenize_negative_ids_reject_before_registration() {
|
||||
let (ingress, detok_rx, consumer, _tm_tx) = make_ingress();
|
||||
let (mut ingress, detok_rx, consumer, _tm_tx, _mm_rx) = make_ingress();
|
||||
let (tx, mut rx) = mpsc::channel(8);
|
||||
ingress.drive(Request {
|
||||
rid: "43".into(),
|
||||
@@ -1071,7 +1207,15 @@ mod tests {
|
||||
let (_tm_tx, tm_rx) = flume::unbounded();
|
||||
let (sd_tx, sd_rx) = flume::unbounded::<()>();
|
||||
std::mem::forget(sd_tx);
|
||||
let ingress = Ingress::new(tm_rx, abort_rx, senders, producer, test_limits(), sd_rx);
|
||||
let mut ingress = Ingress::new(
|
||||
tm_rx,
|
||||
abort_rx,
|
||||
senders,
|
||||
producer,
|
||||
test_limits(),
|
||||
test_mm(flume::unbounded().0, true),
|
||||
sd_rx,
|
||||
);
|
||||
|
||||
ingress.on_abort(AbortSource::Guard("pushed".into()));
|
||||
ingress.on_abort(AbortSource::Guard("dropped".into()));
|
||||
@@ -1108,7 +1252,7 @@ mod tests {
|
||||
#[test]
|
||||
fn pre_registration_failure_does_not_deregister() {
|
||||
// Rejected inside `validate` (out-of-vocab id), which runs before registration.
|
||||
let (ingress, detok_rx, _consumer, _tm_tx) = make_ingress();
|
||||
let (mut ingress, detok_rx, _consumer, _tm_tx, _mm_rx) = make_ingress();
|
||||
let mut req = generate_req(41, SamplingParams::default());
|
||||
if let RequestKind::Generate(g) = &mut req.kind {
|
||||
g.input_ids = Some(vec![2_000_000_000]);
|
||||
@@ -1121,7 +1265,7 @@ mod tests {
|
||||
);
|
||||
|
||||
// A post-registration reject still deregisters (the leak fix stays fixed).
|
||||
let (ingress, detok_rx, _consumer, _tm_tx) = make_ingress();
|
||||
let (mut ingress, detok_rx, _consumer, _tm_tx, _mm_rx) = make_ingress();
|
||||
ingress.drive(generate_req(
|
||||
42,
|
||||
SamplingParams {
|
||||
@@ -1140,7 +1284,7 @@ mod tests {
|
||||
/// sees `Register` then `Deregister`. Regression for RSS growth on bad input.
|
||||
#[test]
|
||||
fn rejected_request_deregisters_from_shard() {
|
||||
let (ingress, detok_rx, _consumer, _tm_tx) = make_ingress();
|
||||
let (mut ingress, detok_rx, _consumer, _tm_tx, _mm_rx) = make_ingress();
|
||||
// top_p = 2.0 is outside (0, 1], so `SamplingParams::normalize` rejects it.
|
||||
let bad = SamplingParams {
|
||||
top_p: 2.0,
|
||||
@@ -1167,7 +1311,7 @@ mod tests {
|
||||
/// and kills the scheduler process (`make_ingress` bounds vocab at 1000).
|
||||
#[test]
|
||||
fn out_of_vocab_input_ids_rejected() {
|
||||
let (ingress, detok_rx, _consumer, _tm_tx) = make_ingress();
|
||||
let (mut ingress, detok_rx, _consumer, _tm_tx, _mm_rx) = make_ingress();
|
||||
let mut req = generate_req(21, SamplingParams::default());
|
||||
if let RequestKind::Generate(g) = &mut req.kind {
|
||||
g.input_ids = Some(vec![1, 2_000_000_000]);
|
||||
@@ -1185,7 +1329,7 @@ mod tests {
|
||||
/// Same guard for negative ids and for `token_ids_logprob` entries.
|
||||
#[test]
|
||||
fn negative_and_logprob_token_ids_rejected() {
|
||||
let (ingress, detok_rx, _consumer, _tm_tx) = make_ingress();
|
||||
let (mut ingress, detok_rx, _consumer, _tm_tx, _mm_rx) = make_ingress();
|
||||
let mut req = generate_req(22, SamplingParams::default());
|
||||
if let RequestKind::Generate(g) = &mut req.kind {
|
||||
g.input_ids = Some(vec![-1]);
|
||||
@@ -1196,7 +1340,7 @@ mod tests {
|
||||
Ok(_) => panic!("negative token id must not be admitted"),
|
||||
}
|
||||
|
||||
let (ingress, detok_rx, _consumer, _tm_tx) = make_ingress();
|
||||
let (mut ingress, detok_rx, _consumer, _tm_tx, _mm_rx) = make_ingress();
|
||||
let mut req = generate_req(23, SamplingParams::default());
|
||||
if let RequestKind::Generate(g) = &mut req.kind {
|
||||
g.token_ids_logprob = Some(vec![999_999]);
|
||||
@@ -1211,7 +1355,7 @@ mod tests {
|
||||
/// A valid request is registered and handed onward — never deregistered.
|
||||
#[test]
|
||||
fn admitted_request_keeps_registration() {
|
||||
let (ingress, detok_rx, _consumer, _tm_tx) = make_ingress();
|
||||
let (mut ingress, detok_rx, _consumer, _tm_tx, _mm_rx) = make_ingress();
|
||||
// Empty map → all sampling defaults, passes normalization.
|
||||
ingress.drive(generate_req(9, SamplingParams::default()));
|
||||
|
||||
@@ -1229,7 +1373,7 @@ mod tests {
|
||||
/// path and deregistered, not leaked.
|
||||
#[test]
|
||||
fn tokenize_failure_deregisters_via_ingress() {
|
||||
let (ingress, detok_rx, _consumer, tm_tx) = make_ingress();
|
||||
let (ingress, detok_rx, _consumer, tm_tx, _mm_rx) = make_ingress();
|
||||
// The pool marks a failed encode as `Failed(err)` before returning it.
|
||||
let mut req = generate_req(11, SamplingParams::default());
|
||||
let _ = req
|
||||
@@ -1253,7 +1397,7 @@ mod tests {
|
||||
fn abort_deregisters_from_shard() {
|
||||
// Aborts arrive on their own unbounded lane now, not the request inbox.
|
||||
let (abort_tx, abort_rx) = flume::unbounded::<AbortSource>();
|
||||
let (ingress, detok_rx, _consumer, tm_tx) = make_ingress_with_abort(abort_rx);
|
||||
let (ingress, detok_rx, _consumer, tm_tx, _mm_rx) = make_ingress_with_abort(abort_rx);
|
||||
abort_tx.send(AbortSource::Guard("rid-13".into())).unwrap();
|
||||
drop(abort_tx);
|
||||
drop(tm_tx);
|
||||
@@ -1270,7 +1414,7 @@ mod tests {
|
||||
/// rejected; its registration is untouched.
|
||||
#[test]
|
||||
fn tokenized_return_pushes_without_deregister() {
|
||||
let (ingress, detok_rx, _consumer, tm_tx) = make_ingress();
|
||||
let (ingress, detok_rx, _consumer, tm_tx, _mm_rx) = make_ingress();
|
||||
let mut req = generate_req(15, SamplingParams::default());
|
||||
// Simulate a successful pool return: ids filled, PreSendValidating.
|
||||
if let RequestKind::Generate(g) = &mut req.kind {
|
||||
@@ -1293,7 +1437,7 @@ mod tests {
|
||||
#[test]
|
||||
fn tokenize_pool_gone_deregisters() {
|
||||
// `make_ingress` drops the tok receiver, so `tok.send` fails.
|
||||
let (ingress, detok_rx, _consumer, _tm_tx) = make_ingress();
|
||||
let (mut ingress, detok_rx, _consumer, _tm_tx, _mm_rx) = make_ingress();
|
||||
// No ids → NeedsTokenize → Tokenizing branch.
|
||||
let mut req = generate_req(21, SamplingParams::default());
|
||||
if let RequestKind::Generate(g) = &mut req.kind {
|
||||
@@ -1311,4 +1455,157 @@ mod tests {
|
||||
);
|
||||
assert!(detok_rx.try_recv().is_err(), "no further shard messages");
|
||||
}
|
||||
|
||||
/// Build a generate request carrying an image. The parked entry and the
|
||||
/// `MmEncoded` resume path agree on identity via the rid string.
|
||||
fn mm_generate_req(rid: &str) -> Request {
|
||||
let (tx, _rx) = mpsc::channel(8);
|
||||
Request {
|
||||
rid: rid.to_string().into(),
|
||||
state: RequestState::Received,
|
||||
sink: EgressSink::Local(tx),
|
||||
kind: RequestKind::Generate(Box::new(GenerateRequest {
|
||||
rid: rid.to_string().into(),
|
||||
text: Some("<image> hi".into()),
|
||||
mm: Some(Box::new(crate::message::MmData {
|
||||
image_data: Some(rmpv::Value::from("data:image/jpeg;base64,xxxx")),
|
||||
..Default::default()
|
||||
})),
|
||||
..Default::default()
|
||||
})),
|
||||
}
|
||||
}
|
||||
|
||||
/// An abort while the request is parked for MM cancels it: the pending
|
||||
/// entry is removed, the worker's late result is dropped, and its parked
|
||||
/// sidecar entry is purged — no scheduler work runs for a dead client.
|
||||
#[test]
|
||||
fn abort_cancels_parked_mm_request() {
|
||||
let (mut ingress, _detok_rx, consumer, _tm_tx, mm_rx) = make_ingress();
|
||||
ingress.drive(mm_generate_req("mm-gone"));
|
||||
mm_rx.try_recv().expect("parked to mm pool");
|
||||
|
||||
// The worker parks its result, as it always does before MmEncoded.
|
||||
ingress.mm.sidecar.park(
|
||||
"mm-gone".into(),
|
||||
crate::mm::MmSidecarEntry {
|
||||
features: crate::mm::FeatureStore::Inline(vec![]),
|
||||
grids: vec![],
|
||||
hashes: vec![],
|
||||
offsets: vec![],
|
||||
mrope: vec![],
|
||||
mrope_delta: 0,
|
||||
},
|
||||
);
|
||||
ingress.on_abort(AbortSource::Guard("mm-gone".to_string().into()));
|
||||
assert_eq!(consumer.drain(16).headers.len(), 1, "only the AbortReq");
|
||||
|
||||
// The late result must be dropped, not queued, and the sidecar purged.
|
||||
ingress.on_mm_encoded("mm-gone".to_string().into(), vec![5, 6]);
|
||||
assert!(
|
||||
consumer.drain(16).headers.is_empty(),
|
||||
"cancelled, not queued"
|
||||
);
|
||||
assert!(ingress.mm.sidecar.take("mm-gone").is_none(), "entry purged");
|
||||
}
|
||||
|
||||
/// A multimodal request parks in `Encoding` (submitted to the mm worker
|
||||
/// pool, not the tokenizer pool, not the ring) until `MmEncoded` resumes
|
||||
/// it → ring.
|
||||
#[test]
|
||||
fn mm_request_parks_then_mm_encoded_pushes_to_ring() {
|
||||
let (mut ingress, _detok_rx, consumer, _tm_tx, mm_rx) = make_ingress();
|
||||
ingress.drive(mm_generate_req("mm-1"));
|
||||
|
||||
// Submitted to the mm pool with the typed work item; nothing on the ring yet.
|
||||
let sub = mm_rx.try_recv().expect("mm pool must receive the request");
|
||||
assert_eq!(sub.rid.as_str(), "mm-1");
|
||||
assert_eq!(sub.work.text.as_deref(), Some("<image> hi"));
|
||||
assert!(sub.work.input_ids.is_none(), "no client input_ids");
|
||||
assert_eq!(
|
||||
sub.work.image_data.as_ref().and_then(|v| v.as_str()),
|
||||
Some("data:image/jpeg;base64,xxxx")
|
||||
);
|
||||
assert!(consumer.drain(16).headers.is_empty(), "parked, not queued");
|
||||
|
||||
// The worker returns the final expanded ids → pushed to the ring.
|
||||
ingress.on_mm_encoded("mm-1".to_string().into(), vec![5, 6, 7, 8]);
|
||||
let batch = consumer.drain(16);
|
||||
assert_eq!(batch.headers.len(), 1);
|
||||
assert_eq!(
|
||||
batch.lengths,
|
||||
vec![4],
|
||||
"expanded ids ride the columnar cell"
|
||||
);
|
||||
}
|
||||
|
||||
/// A worker failure rejects the parked request (deregister, no ring push).
|
||||
#[test]
|
||||
fn mm_failure_rejects_parked_request() {
|
||||
let (mut ingress, detok_rx, consumer, _tm_tx, _mm_rx) = make_ingress();
|
||||
ingress.drive(mm_generate_req("mm-2"));
|
||||
assert!(
|
||||
matches!(detok_rx.try_recv(), Ok(DetokMsg::Register { .. })),
|
||||
"registered before parking",
|
||||
);
|
||||
|
||||
ingress.on_mm_failed("mm-2".to_string().into(), "bad image".into());
|
||||
assert!(
|
||||
matches!(detok_rx.try_recv(), Ok(DetokMsg::Deregister { rid })
|
||||
if rid.as_str() == "mm-2"),
|
||||
"mm failure must deregister",
|
||||
);
|
||||
assert!(consumer.drain(16).headers.is_empty(), "nothing queued");
|
||||
}
|
||||
|
||||
/// On a non-multimodal model (`Mm::enabled == false`), image_data is silently
|
||||
/// ignored and the request tokenizes as plain text — the Python
|
||||
/// TokenizerManager behavior when `mm_processor is None`.
|
||||
#[test]
|
||||
fn mm_fields_ignored_when_disabled() {
|
||||
let (tok_tx, tok_rx) = flume::unbounded();
|
||||
let (detok_tx, _detok_rx) = flume::unbounded();
|
||||
let senders = Senders {
|
||||
tm: flume::unbounded().0,
|
||||
abort: flume::unbounded().0,
|
||||
tok: tok_tx,
|
||||
detok: vec![detok_tx],
|
||||
};
|
||||
let (ingress_producer, _consumer) = ingress_ring(16);
|
||||
let (_tm_tx, tm_rx) = flume::unbounded();
|
||||
let (mm_tx, mm_rx) = flume::unbounded();
|
||||
let (abort_tx, abort_rx) = flume::unbounded::<AbortSource>();
|
||||
std::mem::forget(abort_tx);
|
||||
let (sd_tx, sd_rx) = flume::unbounded::<()>();
|
||||
std::mem::forget(sd_tx);
|
||||
let mut ingress = Ingress::new(
|
||||
tm_rx,
|
||||
abort_rx,
|
||||
senders,
|
||||
ingress_producer,
|
||||
test_limits(),
|
||||
test_mm(mm_tx, false),
|
||||
sd_rx,
|
||||
);
|
||||
|
||||
ingress.drive(mm_generate_req("mm-3"));
|
||||
assert!(
|
||||
mm_rx.try_recv().is_err(),
|
||||
"mm disabled: nothing submitted to the mm channel",
|
||||
);
|
||||
assert!(
|
||||
tok_rx.try_recv().is_ok(),
|
||||
"request must fall through to plain tokenization",
|
||||
);
|
||||
}
|
||||
|
||||
/// A late mm result for a rid that is no longer parked is dropped without
|
||||
/// panicking (e.g. hash-collision overwrite) — regression guard.
|
||||
#[test]
|
||||
fn late_mm_result_is_dropped() {
|
||||
let (mut ingress, _detok_rx, consumer, _tm_tx, _mm_rx) = make_ingress();
|
||||
ingress.on_mm_encoded("ghost".to_string().into(), vec![1]);
|
||||
ingress.on_mm_failed("ghost".to_string().into(), "boom".into());
|
||||
assert!(consumer.drain(16).headers.is_empty());
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user