From 6e48c13497115e4e836fe719463120ec476600c4 Mon Sep 17 00:00:00 2001 From: Rain Jiang <96632942+rainj-me@users.noreply.github.com> Date: Wed, 29 Jul 2026 17:16:06 -0700 Subject: [PATCH] sglang rust server egress message (#32342) --- rust/sglang-server/src/message.rs | 2 + rust/sglang-server/src/message/egress.rs | 1190 +++++++++++++++++ .../src/message/finish_reason.rs | 217 +++ 3 files changed, 1409 insertions(+) create mode 100644 rust/sglang-server/src/message/egress.rs create mode 100644 rust/sglang-server/src/message/finish_reason.rs diff --git a/rust/sglang-server/src/message.rs b/rust/sglang-server/src/message.rs index 9b6e2754f..2eb79682c 100644 --- a/rust/sglang-server/src/message.rs +++ b/rust/sglang-server/src/message.rs @@ -8,6 +8,8 @@ //! `SamplingParams` port), [`io_struct`] (the scheduler wire structs), [`types`] //! (the shared wire-shape adapters both directions use). +mod egress; +mod finish_reason; mod io_struct; mod request; mod sampling; diff --git a/rust/sglang-server/src/message/egress.rs b/rust/sglang-server/src/message/egress.rs new file mode 100644 index 000000000..24d5b6d0e --- /dev/null +++ b/rust/sglang-server/src/message/egress.rs @@ -0,0 +1,1190 @@ +//! The egress (response) direction: the per-request back-channel the API +//! handler drains ([`EgressSink`] / [`EgressItem`]), the egress-ring frame +//! encodings (batch / control result / error), and the columnar batch decode +//! into per-request [`ChunkEvent`]s. + +use bytes::Bytes; +use serde::{Deserialize, Serialize}; +use tokio::sync::mpsc; + +use super::TokenIds; +use super::finish_reason::FinishReason; +use crate::error::Error; +use crate::ids::Rid; + +/// Per-request back-channel the detok shard writes egress frames to and the API +/// handler drains for SSE; bounded, and receiver-drop (disconnect) = stream end. +#[derive(Clone, Debug)] +pub enum EgressSink { + Local(mpsc::Sender), +} + +/// Why an [`EgressSink::try_send`] failed: `Full` = client backpressure, `Closed` +/// = client gone. Both terminal for a stream; the caller distinguishes for logging. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum SinkError { + Full, + Closed, +} + +impl EgressSink { + /// Non-blocking send. `Err(Full)` = backpressure, `Err(Closed)` = client gone. + pub fn try_send(&self, item: EgressItem) -> Result<(), SinkError> { + match self { + EgressSink::Local(tx) => tx.try_send(item).map_err(|e| match e { + mpsc::error::TrySendError::Full(_) => SinkError::Full, + mpsc::error::TrySendError::Closed(_) => SinkError::Closed, + }), + } + } +} + +#[allow(dead_code)] // the receiver half is created inline in api_server::submit. +pub type EgressSource = mpsc::Receiver; + +/// What the connection handler receives on the egress stream: a detok-decoded +/// [`ChunkEvent`] (handler formats it), a verbatim control payload, or an error. +#[derive(Debug)] +pub enum EgressItem { + /// An intermediate streamed generation step (only sent for streaming reqs). + Frame(ChunkEvent), + /// The final generation step. + Done(ChunkEvent), + /// A control-request result: one verbatim payload (e.g. `/server_info`), + /// delivered as-is with no per-protocol formatting. + Control(Bytes), + /// Terminal failure: handler emits an error frame (stream) or status (unary). + Error(Error), +} + +/// Egress-ring frame tag (first byte, prepended Rust-side; Python wire unchanged): +/// a single control-request result payload. +pub const EGRESS_TAG_RESULT: u8 = 1; +/// A whole decode batch: msgpack columnar header + one concatenated raw buffer; +/// tm-egress decodes it into per-request [`ChunkEvent`]s (no per-request FFI). +pub const EGRESS_TAG_BATCH: u8 = 2; +/// A per-request failure `[rid, message]`: the Python drain couldn't decode a +/// header, so it routes a 400 back to that request instead of crashing the loop. +pub const EGRESS_TAG_ERROR: u8 = 3; + +/// Read `n` little-endian f32s from `data` at `*off`, advancing `*off`. `None` when +/// the range runs past the buffer (a malformed / positional-ABI-drifted frame): the +/// caller rejects the whole frame. Bounds-checked via `data.get` — clamping only the +/// end is unsafe because a prior bad length can push `*off` past `len`, making the +/// range reversed (`start > end`) and the slice panic. +fn take_f32(data: &[u8], off: &mut usize, n: usize) -> Option> { + let start = *off; + let end = start.checked_add(n.checked_mul(4)?)?; + let out = data + .get(start..end)? + .chunks_exact(4) + .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])) + .collect(); + *off = end; + Some(out) +} + +/// Read `n` little-endian i32s from `data` at `*off`, advancing `*off`. `None` past +/// the buffer end (see [`take_f32`]). +fn take_i32(data: &[u8], off: &mut usize, n: usize) -> Option> { + let start = *off; + let end = start.checked_add(n.checked_mul(4)?)?; + let out = data + .get(start..end)? + .chunks_exact(4) + .map(|c| i32::from_le_bytes([c[0], c[1], c[2], c[3]])) + .collect(); + *off = end; + Some(out) +} + +/// Frame a decode batch: `[BATCH tag][u32 header len][header][data cols…]`. The +/// caller's `data_cols` are concatenated straight into the frame (one copy, no +/// `b"".join`); `header` is the msgpack [`BatchHeader`]. Runs off the GIL. +pub fn frame_egress_batch_cols(header: &[u8], data_cols: &[&[u8]]) -> Bytes { + let data_len: usize = data_cols.iter().map(|c| c.len()).sum(); + let mut buf = Vec::with_capacity(1 + 4 + header.len() + data_len); + buf.push(EGRESS_TAG_BATCH); + buf.extend_from_slice(&(header.len() as u32).to_le_bytes()); + buf.extend_from_slice(header); + for col in data_cols { + buf.extend_from_slice(col); + } + Bytes::from(buf) +} + +/// Columnar scalar header for a whole decode batch. All numeric fields are +/// `#[serde(default)]`; the hot path (no extras) emits just the first four. +#[derive(Debug, Default, Serialize, Deserialize)] +pub struct BatchHeader { + /// Request ids, as the same strings Python holds (`Req.rid`, uuid hex) — + /// hashed back to the internal routing key in `decode_one` + /// (`Rid::shard`), mirroring the control path. The wire has no + /// rid-shape coupling; any string is a valid rid. + pub rids: Vec, + pub finish_reasons: Vec>, + pub prompt_tokens: Vec, + pub tok_lens: Vec, + #[serde(default)] + pub out_lp_lens: Vec, + #[serde(default)] + pub in_lp_lens: Vec, + #[serde(default)] + pub out_top_reqlens: Vec, + #[serde(default)] + pub out_top_poslens: Vec, + #[serde(default)] + pub in_top_reqlens: Vec, + #[serde(default)] + pub in_top_poslens: Vec, + #[serde(default)] + pub out_tid_reqlens: Vec, + #[serde(default)] + pub out_tid_poslens: Vec, + #[serde(default)] + pub in_tid_reqlens: Vec, + #[serde(default)] + pub in_tid_poslens: Vec, + #[serde(default)] + pub hidden_reqlens: Vec, + #[serde(default)] + pub hidden_poslens: Vec, +} + +/// Read a request's flat logprob column (`l` val/idx pairs) from `data` at cursors +/// `cv`/`ci`, advancing them. Only this request is copied — no whole-column buffer. +/// `None` if either read runs past the buffer (see [`take_f32`]). +fn take_flat( + data: &[u8], + cv: &mut usize, + ci: &mut usize, + l: usize, +) -> Option<(Vec, Vec)> { + Some((take_f32(data, cv, l)?, take_i32(data, ci, l)?)) +} + +/// Read `np` per-position lengths from `poslens` at `*pcur`, advancing it. `None` +/// when the range runs past the column — the header's `reqlens` promised more +/// positions than `poslens` carries, so this request's lengths are unknowable. +/// Clamping instead would return a short `lens`, and since `lens` also drives how +/// far the val/idx cursors advance, every later request in that column would read +/// from the wrong offset: silently wrong logprobs rather than a rejected frame. +fn take_poslens(poslens: &[u32], pcur: &mut usize, np: usize) -> Option> { + let start = *pcur; + let end = start.checked_add(np)?; + let lens = poslens.get(start..end)?.to_vec(); + *pcur = end; + Some(lens) +} + +/// Read a request's ragged logprob column (`np` positions): its per-position `lens` +/// from `poslens` (advancing `pcur`), then that many val/idx from `cv`/`ci`. `None` +/// if either read runs past its buffer (see [`take_poslens`], [`take_f32`]). +fn take_ragged( + data: &[u8], + cv: &mut usize, + ci: &mut usize, + poslens: &[u32], + pcur: &mut usize, + np: usize, +) -> Option<(Vec, Vec, Vec)> { + let lens = take_poslens(poslens, pcur, np)?; + let nv: usize = lens.iter().map(|&x| x as usize).sum(); + Some((take_f32(data, cv, nv)?, take_i32(data, ci, nv)?, lens)) +} + +/// Like [`take_ragged`] but for hidden states — a val column + row `poslens`, no +/// idx column. `None` if either read runs past its buffer (see [`take_poslens`], +/// [`take_f32`]). +fn take_hidden( + data: &[u8], + cv: &mut usize, + poslens: &[u32], + pcur: &mut usize, + nr: usize, +) -> Option<(Vec, Vec)> { + let lens = take_poslens(poslens, pcur, nr)?; + let nv: usize = lens.iter().map(|&x| x as usize).sum(); + Some((take_f32(data, cv, nv)?, lens)) +} + +/// Decode a batch egress frame (tag stripped), calling `route` with each request's +/// [`ChunkEvent`] as it's decoded — one pass, no intermediate `Vec`, peak memory +/// one request. Column order matches `push_generation`. +/// +/// `ok == false` means the frame was rejected. The caller discards everything it +/// routed and fails the frame's requests instead of forwarding a partial fan-out +/// (see `tokenizer_manager::egress`), so a rejected frame delivers nothing — +/// `rids` exists precisely so those requests can be failed rather than left +/// waiting for a `Done` that no longer exists. +pub fn for_each_chunk(body: &[u8], mut route: impl FnMut(ChunkEvent)) -> Decoded { + let mut decoded = Decoded::default(); + if body.len() < 4 { + return decoded; + } + let hlen = u32::from_le_bytes([body[0], body[1], body[2], body[3]]) as usize; + let Some(header) = body.get(4..4 + hlen) else { + return decoded; + }; + // Every rejection past the header slice names its requests by re-reading the + // rid column, rather than the whole frame paying a clone for the error path. + macro_rules! reject { + () => {{ + decoded.rids = recover_rids(header); + return decoded; + }}; + } + + let data = &body[4 + hlen..]; + let mut h = match rmp_serde::from_slice::(header) { + Ok(h) => h, + Err(_) => reject!(), + }; + let n = h.rids.len(); + // Every column must agree with `rids`. A SHORT column contributes nothing to + // `base`, so the `base != data.len()` check below cannot see it: a short + // `tok_lens` delivered a 200 with an empty completion, and a short + // `finish_reasons` dropped the terminal marker so the request never completed + // and its unary drain pended forever. The producer already asserts this for the + // extras columns; the four core ones were unchecked. + if h.finish_reasons.len() != n || h.prompt_tokens.len() != n || h.tok_lens.len() != n { + reject!() + } + // The per-request extras columns are either absent (no request asked) or one + // entry per request — never partial. + let per_req_ok = |c: &[u32]| c.is_empty() || c.len() == n; + if !per_req_ok(&h.out_lp_lens) + || !per_req_ok(&h.in_lp_lens) + || !per_req_ok(&h.out_top_reqlens) + || !per_req_ok(&h.in_top_reqlens) + || !per_req_ok(&h.out_tid_reqlens) + || !per_req_ok(&h.in_tid_reqlens) + || !per_req_ok(&h.hidden_reqlens) + { + reject!() + } + let sum = |v: &[u32]| v.iter().map(|&x| x as usize).sum::(); + // Each ragged family's per-request counts must consume its position column + // EXACTLY. Only the deficit was caught (`take_poslens` runs off the end); a + // surplus left positions unread and delivered a truncated row with a 200. + if sum(&h.out_top_reqlens) != h.out_top_poslens.len() + || sum(&h.in_top_reqlens) != h.in_top_poslens.len() + || sum(&h.out_tid_reqlens) != h.out_tid_poslens.len() + || sum(&h.in_tid_reqlens) != h.in_tid_poslens.len() + || sum(&h.hidden_reqlens) != h.hidden_poslens.len() + { + reject!() + } + + // Per-column byte cursors, advanced per request — no whole-column read. Columns + // are concatenated in exactly this order, every element 4 bytes. + let mut base = 0usize; + let mut col = |count: usize| -> usize { + let start = base; + base += count * 4; + start + }; + // Each val/idx column pair shares one element count — sum it once. + let n_ids = sum(&h.tok_lens); + let n_olp = sum(&h.out_lp_lens); + let n_ilp = sum(&h.in_lp_lens); + let n_ot = sum(&h.out_top_poslens); + let n_it = sum(&h.in_top_poslens); + let n_od = sum(&h.out_tid_poslens); + let n_id = sum(&h.in_tid_poslens); + let n_h = sum(&h.hidden_poslens); + let mut c_ids = col(n_ids); + let mut c_olp_v = col(n_olp); + let mut c_olp_i = col(n_olp); + let mut c_ilp_v = col(n_ilp); + let mut c_ilp_i = col(n_ilp); + let mut c_ot_v = col(n_ot); + let mut c_ot_i = col(n_ot); + let mut c_it_v = col(n_it); + let mut c_it_i = col(n_it); + let mut c_od_v = col(n_od); + let mut c_od_i = col(n_od); + let mut c_id_v = col(n_id); + let mut c_id_i = col(n_id); + let mut c_h_v = col(n_h); + + // `col` summed every column's span into `base`, so a truncated frame is caught + // here — the one rejection that is genuinely whole-frame, since it precedes the + // routing loop. Past this point a failure can only be partial. + // + // The check is EQUALITY, not `>`: every column length is header-determined, so + // a data buffer longer than `base` means the header and the buffer disagree. + // Accepting the surplus let a producer-side val/idx mismatch through with a + // 200, every later column reading off by the difference. + if base != data.len() { + reject!() + } + + // Mirror of Python's `has_extra` guard: checking once per frame lets the + // per-request loop skip the extras machinery entirely on a plain decode frame. + let has_extras = !(h.out_lp_lens.is_empty() + && h.in_lp_lens.is_empty() + && h.out_top_reqlens.is_empty() + && h.in_top_reqlens.is_empty() + && h.out_tid_reqlens.is_empty() + && h.in_tid_reqlens.is_empty() + && h.hidden_reqlens.is_empty()); + + // Position cursors into the header's per-request `poslens` (ragged + hidden). + let (mut p_ot, mut p_it, mut p_od, mut p_id, mut p_h) = + (0usize, 0usize, 0usize, 0usize, 0usize); + let lens_i = |v: &[u32], i: usize| v.get(i).copied().unwrap_or(0) as usize; + + // Decode one request's slice of every column, advancing the cursors. `None` if a + // read overruns `data` (belt-and-suspenders past the upfront check, which already + // covers every column's span — so this is unreachable today). It stops the loop + // before slicing out of bounds, but requests decoded earlier in the frame are + // already routed; making that abort atomic would mean buffering the whole frame, + // which is exactly what the streaming decode avoids. + let mut decode_one = |i: usize| -> Option { + let token_ids = take_i32(data, &mut c_ids, lens_i(&h.tok_lens, i))?; + + // Plain decode frame (no request in the batch asked for logprobs/hidden): + // the extras columns are all zero-width, so skip reading them entirely. + let extras = if !has_extras { + None + } else { + let (out_lp_val, out_lp_idx) = + take_flat(data, &mut c_olp_v, &mut c_olp_i, lens_i(&h.out_lp_lens, i))?; + let (in_lp_val, in_lp_idx) = + take_flat(data, &mut c_ilp_v, &mut c_ilp_i, lens_i(&h.in_lp_lens, i))?; + let (out_top_val, out_top_idx, out_top_lens) = take_ragged( + data, + &mut c_ot_v, + &mut c_ot_i, + &h.out_top_poslens, + &mut p_ot, + lens_i(&h.out_top_reqlens, i), + )?; + let (in_top_val, in_top_idx, in_top_lens) = take_ragged( + data, + &mut c_it_v, + &mut c_it_i, + &h.in_top_poslens, + &mut p_it, + lens_i(&h.in_top_reqlens, i), + )?; + let (out_tid_val, out_tid_idx, out_tid_lens) = take_ragged( + data, + &mut c_od_v, + &mut c_od_i, + &h.out_tid_poslens, + &mut p_od, + lens_i(&h.out_tid_reqlens, i), + )?; + let (in_tid_val, in_tid_idx, in_tid_lens) = take_ragged( + data, + &mut c_id_v, + &mut c_id_i, + &h.in_tid_poslens, + &mut p_id, + lens_i(&h.in_tid_reqlens, i), + )?; + let (hidden_val, hidden_lens) = take_hidden( + data, + &mut c_h_v, + &h.hidden_poslens, + &mut p_h, + lens_i(&h.hidden_reqlens, i), + )?; + + // Even in an extras batch, most requests carry none — box only if this + // one actually does, so its `ChunkEvent` stays the small common frame. + let ex = ChunkExtras { + out_lp_val, + out_lp_idx, + in_lp_val, + in_lp_idx, + out_top_val, + out_top_idx, + out_top_lens, + in_top_val, + in_top_idx, + in_top_lens, + out_tid_val, + out_tid_idx, + out_tid_lens, + in_tid_val, + in_tid_idx, + in_tid_lens, + hidden_val, + hidden_lens, + // Explicit, NOT `..Default::default()` — same reason as `ChunkEvent` + // below: a new column must fail to compile here until it is decoded. + out_lp_txt: Vec::new(), + in_lp_txt: Vec::new(), + out_top_txt: Vec::new(), + in_top_txt: Vec::new(), + out_tid_txt: Vec::new(), + in_tid_txt: Vec::new(), + }; + (!ex.is_empty()).then(|| Box::new(ex)) + }; + + Some(ChunkEvent { + // Any string is a valid rid; hash to the routing key. An unknown + // rid routes to a shard whose table has no entry → dropped there. + rid: std::mem::take(&mut h.rids[i]).into(), + token_ids, + finish_reason: h.finish_reasons.get(i).cloned().flatten(), + prompt_tokens: h.prompt_tokens.get(i).copied().unwrap_or(0), + extras, + // Listed explicitly, NOT `..Default::default()`: a new column added to + // `ChunkEvent` and wired into the response must fail to compile here + // until it is actually decoded. With the struct-update syntax it + // compiled clean and silently shipped zeros. + text: String::new(), + completion_tokens: 0, + }) + }; + + for i in 0..n { + let Some(ev) = decode_one(i) else { reject!() }; + route(ev); + } + decoded.ok = true; + decoded +} + +/// Recover the rid column straight from the header bytes. +/// +/// Read through `rmpv` rather than a serde tuple: `(Vec,)` decodes ONLY a +/// 1-element array, while real headers carry 4 or 16 columns, so it returned +/// `Err(LengthMismatch(1))` every time and named nobody. Arity independence is the +/// whole point — the trigger for the typed-decode-failure path is Python appending +/// a column an older Rust build does not know. +/// +/// Called only when a frame is rejected, which is why the accepted path no longer +/// clones the column up front. That clone was discarded unread on every good frame +/// and cost ~27% of the whole decode at batch 4096, plus a third of the crate's +/// steady-state allocations. It also has to be a re-read rather than a snapshot: +/// `decode_one` moves each rid out of the header as it goes, so by the time the +/// routing loop can fail, the decoded header no longer holds the earlier ones. +fn recover_rids(header: &[u8]) -> Vec { + rmpv::decode::read_value(&mut &header[..]) + .ok() + .and_then(|v| match v { + rmpv::Value::Array(cols) => cols.into_iter().next(), + _ => None, + }) + .and_then(|c| match c { + rmpv::Value::Array(rids) => Some( + rids.iter() + .filter_map(|r| r.as_str().map(Rid::from)) + .collect(), + ), + _ => None, + }) + .unwrap_or_default() +} + +/// Outcome of [`for_each_chunk`]: whether the frame was accepted, plus the rids it +/// named. `rids` is populated as soon as the header parses, so a caller can fail +/// every request in a rejected frame — including ones whose chunk never decoded, +/// which would otherwise wait forever for a `Done` that no longer exists. +#[derive(Debug, Default)] +pub struct Decoded { + pub ok: bool, + pub rids: Vec, +} + +/// Frame a control result `[rid, payload]` for the egress ring (tag prepended). +pub fn frame_egress_result(rid: &str, payload: &[u8]) -> Bytes { + use rmpv::Value; + let arr = Value::Array(vec![Value::from(rid), Value::Binary(payload.to_vec())]); + let mut buf = Vec::with_capacity(1 + payload.len() + rid.len() + 8); + buf.push(EGRESS_TAG_RESULT); + let _ = rmpv::encode::write_value(&mut buf, &arr); + Bytes::from(buf) +} + +/// Frame a per-request failure `[rid, message]` for the egress ring — routes a +/// terminal error back to the owning request (→ HTTP 400) instead of crashing. +pub fn frame_egress_error(rid: &str, message: &str) -> Bytes { + use rmpv::Value; + let arr = Value::Array(vec![Value::from(rid), Value::from(message)]); + let mut buf = Vec::with_capacity(1 + rid.len() + message.len() + 8); + buf.push(EGRESS_TAG_ERROR); + let _ = rmpv::encode::write_value(&mut buf, &arr); + Bytes::from(buf) +} + +/// One scheduler output increment — the common, always-present frame. `token_ids` +/// / `prompt_tokens` / `finish_reason` arrive from Python (pre-decode); the detok +/// shard fills `text` in place. Deltas — fold with `OutputAccumulator` for +/// cumulative. Logprobs + hidden states (rare, and large) live behind the boxed +/// [`ChunkExtras`] (`None` unless requested) so this frame stays small even when +/// the decoder builds an inline array at up to batch 4096 per step. +/// +/// Not a wire type — built by `for_each_chunk` from the columnar [`BatchHeader`] +/// frame and moved between stages in-process (never serialized), so no serde. +#[derive(Debug, Clone, Default)] +pub struct ChunkEvent { + /// Client-visible rid — the request's IDENTITY. Moved out of the frame header + /// (which owns it and drops it), so carrying it costs no allocation. The shard + /// is still chosen by `Rid::shard`, but a hash collision there now only + /// co-locates two requests instead of merging them. + pub rid: Rid, + /// New token ids for this step. Empty allowed (e.g. metadata-only frames). + pub token_ids: TokenIds, + /// `None` while streaming, the [`FinishReason`] on the final chunk. + pub finish_reason: Option, + /// Prompt token count for this request (constant across its chunks). + pub prompt_tokens: u32, + /// Decoded text **delta** for this chunk (empty in skip mode / on partial UTF-8), + /// filled by the detok shard. `token_ids` doubles as `output_ids`; + /// `completion_tokens` is this chunk's count. + pub text: String, + pub completion_tokens: u64, + /// Logprob + hidden-state columns — `None` unless the request asked for them. + /// Boxed to keep the common token/text/finish frame small at large decode + /// batches (the decoder allocates it only when a column is non-empty). + pub extras: Option>, +} + +/// Logprob + hidden-state columns for a [`ChunkEvent`], allocated only when the +/// request enabled logprobs / hidden states. Columnar `val`/`idx` (+ ragged `lens`) +/// buffers arrive pre-decode; the detok shard fills the parallel `*_txt` columns +/// when `return_text_in_logprobs` is set. In-process only — no serde (see +/// [`ChunkEvent`]). +#[derive(Debug, Clone, Default)] +pub struct ChunkExtras { + /// Output-token logprobs (parallel `val`/`idx`, one entry per new output token). + pub out_lp_val: Vec, + pub out_lp_idx: Vec, + /// Input (prefill) token logprobs, sent once on the first chunk. + pub in_lp_val: Vec, + pub in_lp_idx: Vec, + /// Top-k logprobs (2-level ragged): flat `val`/`idx` + per-position `lens` (0 = + /// null). Output = per-step delta, input = once on the first chunk. + pub out_top_val: Vec, + pub out_top_idx: Vec, + pub out_top_lens: Vec, + pub in_top_val: Vec, + pub in_top_idx: Vec, + pub in_top_lens: Vec, + /// Token-ids logprobs (same ragged layout); set only when `token_ids_logprob` was. + pub out_tid_val: Vec, + pub out_tid_idx: Vec, + pub out_tid_lens: Vec, + pub in_tid_val: Vec, + pub in_tid_idx: Vec, + pub in_tid_lens: Vec, + /// Hidden states (dense f32): flat buffer + per-row lengths. Last-writer-wins + /// across chunks (the final message has the full set). + pub hidden_val: Vec, + pub hidden_lens: Vec, + /// Decoded logprob token text (`return_text_in_logprobs`), parallel to the + /// `*_idx` buffers; empty when not requested (the tuple's text slot stays null). + pub out_lp_txt: Vec, + pub in_lp_txt: Vec, + pub out_top_txt: Vec, + pub in_top_txt: Vec, + pub out_tid_txt: Vec, + pub in_tid_txt: Vec, +} + +impl ChunkExtras { + /// True when no logprob / hidden column carries data — lets the decoder skip the + /// box allocation for the common (extras-free) frame. + fn is_empty(&self) -> bool { + self.out_lp_val.is_empty() + && self.in_lp_val.is_empty() + && self.out_top_lens.is_empty() + && self.in_top_lens.is_empty() + && self.out_tid_lens.is_empty() + && self.in_tid_lens.is_empty() + && self.hidden_lens.is_empty() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::message::finish_reason::{FinishKind, Matched}; + + #[test] + fn batch_cols_match_single_joined_buffer() { + let header = [1u8, 2, 3]; + let a = [10u8, 11]; + let b = [12u8, 13, 14]; + let multi = frame_egress_batch_cols(&header, &[&a[..], &b[..]]); + let joined: Vec = a.iter().chain(&b).copied().collect(); + let single = frame_egress_batch_cols(&header, &[joined.as_slice()]); + assert_eq!(multi, single); + assert_eq!(multi[0], EGRESS_TAG_BATCH); + assert_eq!( + u32::from_le_bytes([multi[1], multi[2], multi[3], multi[4]]), + 3 + ); + assert_eq!(&multi[5..8], &header); // header + assert_eq!(&multi[8..], &[10, 11, 12, 13, 14]); // columns end-to-end + } + + /// A batch frame (the fast path) decodes into per-request ChunkEvents, with + /// token ids sliced from the single concatenated buffer by `tok_lens`. The + /// header is a msgspec-style positional array (what Python emits). + #[test] + fn decodes_batch_frame() { + use rmpv::Value; + // 3 requests: rids "1","2","3"; finish [nil, {type:stop,matched:5}, nil]; + // prompt_tokens [4,5,6]; tok_lens [2,0,1] -> ids [10,11 | (none) | 12]. + let stop = Value::Map(vec![ + (Value::from("type"), Value::from("stop")), + (Value::from("matched"), Value::from(5)), + ]); + let header_arr = Value::Array(vec![ + Value::Array(vec![Value::from("1"), Value::from("2"), Value::from("3")]), + Value::Array(vec![Value::Nil, stop, Value::Nil]), + Value::Array(vec![ + Value::from(4u32), + Value::from(5u32), + Value::from(6u32), + ]), + Value::Array(vec![ + Value::from(2u32), + Value::from(0u32), + Value::from(1u32), + ]), + ]); + let mut header = Vec::new(); + rmpv::encode::write_value(&mut header, &header_arr).unwrap(); + let data: Vec = [10i32, 11, 12] + .iter() + .flat_map(|x| x.to_le_bytes()) + .collect(); + + let framed = frame_egress_batch_cols(&header, &[&data]); + assert_eq!(framed[0], EGRESS_TAG_BATCH); + let mut events = Vec::new(); + assert!(for_each_chunk(&framed[1..], |ev| events.push(ev)).ok); + assert_eq!(events.len(), 3); + assert_eq!(events[0].rid, Rid::from("1")); + assert_eq!(events[0].token_ids, vec![10, 11]); + assert_eq!(events[0].prompt_tokens, 4); + assert!(events[0].finish_reason.is_none()); + assert_eq!(events[1].rid, Rid::from("2")); + assert!(events[1].token_ids.is_empty()); + // The whole reason survives msgpack (type + matched), not just the type. + assert_eq!( + events[1].finish_reason, + Some( + FinishKind::Stop { + matched: Some(Matched::Token(5)) + } + .into() + ) + ); + assert_eq!(events[2].rid, Rid::from("3")); + assert_eq!(events[2].token_ids, vec![12]); + assert_eq!(events[2].prompt_tokens, 6); + // A plain decode frame carries no extras columns at all, so the per-frame + // `has_extras` guard must skip the extras machinery entirely for every + // request (this is the tm-egress hot path — see `for_each_chunk`). + assert!(events.iter().all(|e| e.extras.is_none())); + } + + /// A header whose column lengths exceed the data buffer (a Python/Rust + /// positional-ABI drift, or a truncated frame) is rejected: `for_each_chunk` + /// returns false and routes nothing — it must NOT panic the sole egress thread + /// on an out-of-bounds slice. Built the way Python emits (positional msgpack + /// header + concatenated data columns). + #[test] + fn rejects_frame_with_lengths_past_data() { + use rmpv::Value; + // 1 request: tok_lens[0]=10 claims 40 bytes and out_lp_lens[0]=1 puts the + // logprob column's base past the 4-byte data buffer. The old clamp-only-`end` + // code advanced the cursor past `len`, then sliced `data[40..4]` (start > end) + // and panicked. + let header_arr = Value::Array(vec![ + Value::Array(vec![Value::from("1")]), // rids + Value::Array(vec![Value::Nil]), // finish_reasons + Value::Array(vec![Value::from(0u32)]), // prompt_tokens + Value::Array(vec![Value::from(10u32)]), // tok_lens (claims 40 bytes) + Value::Array(vec![Value::from(1u32)]), // out_lp_lens (base now past data) + ]); + let mut header = Vec::new(); + rmpv::encode::write_value(&mut header, &header_arr).unwrap(); + let data: Vec = [0i32].iter().flat_map(|x| x.to_le_bytes()).collect(); // 4 bytes + + let framed = frame_egress_batch_cols(&header, &[&data]); + let mut routed = 0usize; + let decoded = for_each_chunk(&framed[1..], |_| routed += 1); + assert!(!decoded.ok, "malformed frame must be rejected, not decoded"); + assert_eq!(routed, 0, "no request may be routed from a rejected frame"); + // The rid is still reported, so the caller can fail the request that was + // waiting on this frame instead of letting it hang. + assert_eq!(decoded.rids, vec![Rid::from("1")]); + } + + /// A header whose `reqlens` claim more positions than `poslens` carries is + /// rejected, not truncated. This drift passes the upfront `base > data.len()` + /// check — the data columns are exactly as long as `poslens` says — so only the + /// per-column bound catches it. Clamping (the old behavior) handed req0 a short + /// `lens`, which also under-advanced the val/idx cursors, so req1 read from the + /// wrong offset: a frame that decodes "successfully" into wrong logprobs. + #[test] + fn rejects_poslens_shorter_than_reqlens_claims() { + use rmpv::Value; + let f = |xs: &[f32]| -> Vec { xs.iter().flat_map(|x| x.to_le_bytes()).collect() }; + let i = |xs: &[i32]| -> Vec { xs.iter().flat_map(|x| x.to_le_bytes()).collect() }; + let arr_u = |xs: &[u32]| Value::Array(xs.iter().map(|&x| Value::from(x)).collect()); + let rids = Value::Array(vec![Value::from("1"), Value::from("2")]); + let finish = Value::Array(vec![Value::Nil, Value::Nil]); + + // Ragged column: reqs claim 2 + 1 = 3 positions, `out_top_poslens` has 2. + let header_arr = Value::Array(vec![ + rids.clone(), + finish.clone(), + arr_u(&[0, 0]), // prompt + arr_u(&[1, 1]), // tok_lens + arr_u(&[0, 0]), // out_lp_lens + arr_u(&[0, 0]), // in_lp_lens + arr_u(&[2, 1]), // out_top_reqlens — 3 positions claimed + arr_u(&[2, 2]), // out_top_poslens — only 2 supplied + ]); + let mut header = Vec::new(); + rmpv::encode::write_value(&mut header, &header_arr).unwrap(); + let mut data = Vec::new(); + data.extend(i(&[10, 20])); // token_ids + data.extend(f(&[-0.1, -0.2, -0.3, -0.4])); // out_top_val (sum of poslens = 4) + data.extend(i(&[1, 2, 3, 4])); // out_top_idx + let framed = frame_egress_batch_cols(&header, &[&data]); + let mut routed = 0usize; + assert!( + !for_each_chunk(&framed[1..], |_| routed += 1).ok, + "ragged poslens drift must be rejected, not truncated" + ); + + // Same drift in the hidden column, which has no idx pair: 2 rows claimed, 1 supplied. + let header_arr = Value::Array(vec![ + rids, + finish, + arr_u(&[0, 0]), // prompt + arr_u(&[1, 1]), // tok_lens + arr_u(&[0, 0]), // out_lp_lens + arr_u(&[0, 0]), // in_lp_lens + arr_u(&[0, 0]), // out_top_reqlens + arr_u(&[]), // out_top_poslens + arr_u(&[0, 0]), // in_top_reqlens + arr_u(&[]), // in_top_poslens + arr_u(&[0, 0]), // out_tid_reqlens + arr_u(&[]), // out_tid_poslens + arr_u(&[0, 0]), // in_tid_reqlens + arr_u(&[]), // in_tid_poslens + arr_u(&[2, 0]), // hidden_reqlens — 2 rows claimed + arr_u(&[3]), // hidden_poslens — only 1 supplied + ]); + let mut header = Vec::new(); + rmpv::encode::write_value(&mut header, &header_arr).unwrap(); + let mut data = Vec::new(); + data.extend(i(&[10, 20])); // token_ids + data.extend(f(&[0.1, 0.2, 0.3])); // hidden_val (sum of poslens = 3) + let framed = frame_egress_batch_cols(&header, &[&data]); + let mut routed = 0usize; + assert!( + !for_each_chunk(&framed[1..], |_| routed += 1).ok, + "hidden poslens drift must be rejected, not truncated" + ); + } + + /// A frame that fails at request 0 buckets nothing, so bucket-driven cleanup + /// would leave every request in it hanging. The rids come from the header, + /// which is fully parsed before the decode loop. + #[test] + fn rejected_frame_still_reports_all_its_rids() { + use rmpv::Value; + let header_arr = Value::Array(vec![ + Value::Array(vec![Value::from("r0"), Value::from("r1")]), // rids + Value::Array(vec![Value::Nil, Value::Nil]), + Value::Array(vec![Value::from(0u32), Value::from(0u32)]), + Value::Array(vec![Value::from(9u32), Value::from(9u32)]), // tok_lens past data + ]); + let mut header = Vec::new(); + rmpv::encode::write_value(&mut header, &header_arr).unwrap(); + let framed = frame_egress_batch_cols(&header, &[&[0u8; 4][..]]); + let mut routed = 0usize; + let decoded = for_each_chunk(&framed[1..], |_| routed += 1); + assert!(!decoded.ok); + assert_eq!(routed, 0, "fails at request 0 — nothing bucketed"); + assert_eq!( + decoded.rids, + vec![Rid::from("r0"), Rid::from("r1")], + "both requests must be nameable so the caller can fail them" + ); + } + + /// A data buffer LONGER than the header's columns means the two disagree — a + /// producer-side val/idx mismatch. Accepting the surplus delivered another + /// column's bytes as logprobs with a 200. + #[test] + fn frame_longer_than_its_columns_is_rejected() { + use rmpv::Value; + let header_arr = Value::Array(vec![ + Value::Array(vec![Value::from("1")]), + Value::Array(vec![Value::Nil]), + Value::Array(vec![Value::from(0u32)]), + Value::Array(vec![Value::from(1u32)]), // tok_lens: 1 id = 4 bytes + ]); + let mut header = Vec::new(); + rmpv::encode::write_value(&mut header, &header_arr).unwrap(); + let data: Vec = vec![0u8; 8]; // 4 bytes too many + let framed = frame_egress_batch_cols(&header, &[&data]); + let decoded = for_each_chunk(&framed[1..], |_| {}); + assert!(!decoded.ok, "header and data must agree exactly"); + } + + /// Ingress/egress rid agreement: the rid decoded from the frame must be the + /// one Python sent, AND both sides must derive the same shard from it. The + /// partition key is memoized inside `Rid`, so a per-conversion hasher seed + /// would send a request's chunks to a shard that never registered it. + #[test] + fn uuid_rid_decodes_to_the_same_rid_and_shard() { + use rmpv::Value; + let rid = Rid::from("9f86d081884c7d659a2feaa0c55ad015"); + let header_arr = Value::Array(vec![ + Value::Array(vec![Value::from(rid.to_string())]), // rids + Value::Array(vec![Value::Nil]), // finish_reasons + Value::Array(vec![Value::from(1u32)]), // prompt_tokens + Value::Array(vec![Value::from(1u32)]), // tok_lens + ]); + let mut header = Vec::new(); + rmpv::encode::write_value(&mut header, &header_arr).unwrap(); + let data: Vec = [0i32].iter().flat_map(|x| x.to_le_bytes()).collect(); + + let framed = frame_egress_batch_cols(&header, &[&data]); + let mut events = Vec::new(); + assert!(for_each_chunk(&framed[1..], |ev| events.push(ev)).ok); + assert_eq!(events.len(), 1); + // The equality below reuses `from_rid`, so on its own it would hold for + // any decoder that hashed *something*. These pin the two failure modes it + // is meant to catch: a decoder that parsed the rid as an integer (a uuid + // is not numeric, so it would fall back to 0), and one that keyed off the + // request's position in the batch rather than its rid. + assert_eq!( + events[0].rid, rid, + "the rid IS the identity — carried, not hashed" + ); + assert_ne!(events[0].rid, Rid::from("0")); + // Same string, independently constructed → same shard, on any shard count. + for shards in [1usize, 3, 8, 64] { + assert_eq!( + events[0].rid.shard(shards), + Rid::from("9f86d081884c7d659a2feaa0c55ad015").shard(shards), + "the partition key must not depend on WHERE the Rid was built" + ); + } + assert_ne!( + events[0].rid, + Rid::from("another-rid"), + "distinct rids stay distinct — identity is the string, not a digest" + ); + } + + /// A batch frame carrying the numeric columns (extras path): 2 requests, + /// req0 with output logprobs + top-k + hidden, req1 empty. Verifies the + /// column-major data split by the header's reqlens/poslens. + #[test] + fn decodes_batch_frame_with_extras() { + use rmpv::Value; + let f = |xs: &[f32]| -> Vec { xs.iter().flat_map(|x| x.to_le_bytes()).collect() }; + let i = |xs: &[i32]| -> Vec { xs.iter().flat_map(|x| x.to_le_bytes()).collect() }; + let arr_u = |xs: &[u32]| Value::Array(xs.iter().map(|&x| Value::from(x)).collect()); + // header: rids, finish, prompt, tok_lens, out_lp_lens, in_lp_lens, + // out_top_reqlens, out_top_poslens, in_top_*, out_tid_*, in_tid_*, + // hidden_reqlens, hidden_poslens + let header_arr = Value::Array(vec![ + Value::Array(vec![Value::from("1"), Value::from("2")]), // rids + Value::Array(vec![Value::Nil, Value::Nil]), // finish + arr_u(&[3, 4]), // prompt + arr_u(&[1, 1]), // tok_lens + arr_u(&[2, 0]), // out_lp_lens + arr_u(&[0, 0]), // in_lp_lens + arr_u(&[1, 0]), // out_top_reqlens (req0: 1 pos) + arr_u(&[2]), // out_top_poslens (that pos: k=2) + arr_u(&[0, 0]), // in_top_reqlens + arr_u(&[]), // in_top_poslens + arr_u(&[0, 0]), // out_tid_reqlens + arr_u(&[]), // out_tid_poslens + arr_u(&[0, 0]), // in_tid_reqlens + arr_u(&[]), // in_tid_poslens + arr_u(&[1, 0]), // hidden_reqlens (req0: 1 row) + arr_u(&[3]), // hidden_poslens (dim 3) + ]); + let mut header = Vec::new(); + rmpv::encode::write_value(&mut header, &header_arr).unwrap(); + let mut data = Vec::new(); + data.extend(i(&[10, 20])); // token_ids: req0=[10], req1=[20] + data.extend(f(&[-0.5, -0.6])); // out_lp_val (req0, 2) + data.extend(i(&[10, 99])); // out_lp_idx + data.extend(f(&[-0.1, -0.2])); // out_top_val (1 pos, k=2) + data.extend(i(&[10, 11])); // out_top_idx + data.extend(f(&[0.1, 0.2, 0.3])); // hidden_val (1 row, dim 3) + + let framed = frame_egress_batch_cols(&header, &[&data]); + let mut events = Vec::new(); + assert!(for_each_chunk(&framed[1..], |ev| events.push(ev)).ok); + assert_eq!(events.len(), 2); + assert_eq!(events[0].token_ids, vec![10]); + let ex0 = events[0] + .extras + .as_deref() + .expect("req0 has logprob/hidden extras"); + assert_eq!(ex0.out_lp_val, vec![-0.5, -0.6]); + assert_eq!(ex0.out_lp_idx, vec![10, 99]); + assert_eq!(ex0.out_top_val, vec![-0.1, -0.2]); + assert_eq!(ex0.out_top_lens, vec![2]); + assert_eq!(ex0.hidden_val, vec![0.1, 0.2, 0.3]); + assert_eq!(ex0.hidden_lens, vec![3]); + // req1 has a token id but no numeric columns → no extras box allocated. + assert_eq!(events[1].token_ids, vec![20]); + assert!(events[1].extras.is_none()); + } + + /// An inactive extras family may arrive as EMPTY columns or as `batch_size` + /// zeros, and both must decode to exactly the same events. + /// + /// This is the contract the Python producer's per-family gating relies on. It + /// used to run all seven families whenever any one of them was active, so an + /// inactive family shipped 4096 zeros per column; it now skips `accept`, which + /// leaves those columns empty. `per_req_ok` admits either and `lens_i` reads 0 + /// past the end — but nothing pinned that, so a later tightening of the column + /// validation (say, requiring every column to be `batch_size` long) would + /// silently start rejecting whole frames, hanging every request in them. + #[test] + fn empty_and_zero_filled_inactive_families_decode_alike() { + use rmpv::Value; + let f = |xs: &[f32]| -> Vec { xs.iter().flat_map(|x| x.to_le_bytes()).collect() }; + let i = |xs: &[i32]| -> Vec { xs.iter().flat_map(|x| x.to_le_bytes()).collect() }; + let arr_u = |xs: &[u32]| Value::Array(xs.iter().map(|&x| Value::from(x)).collect()); + + // Only `out_lp` is active; the other six families are inactive. `zeros` + // picks how they are spelled: the old producer's per-request zeros, or the + // gated producer's empty column. + let build = |zeros: bool| { + // Inactive: the old producer ran `accept` anyway and appended a 0 per + // request; the gated producer skips it and leaves the column empty. + let reqlens = if zeros { arr_u(&[0, 0]) } else { arr_u(&[]) }; + let header_arr = Value::Array(vec![ + Value::Array(vec![Value::from("1"), Value::from("2")]), // rids + Value::Array(vec![Value::Nil, Value::Nil]), // finish + arr_u(&[3, 4]), // prompt + arr_u(&[1, 1]), // tok_lens + arr_u(&[2, 0]), // out_lp_lens (ACTIVE) + reqlens.clone(), // in_lp_lens + reqlens.clone(), // out_top_reqlens + arr_u(&[]), // out_top_poslens + reqlens.clone(), // in_top_reqlens + arr_u(&[]), // in_top_poslens + reqlens.clone(), // out_tid_reqlens + arr_u(&[]), // out_tid_poslens + reqlens.clone(), // in_tid_reqlens + arr_u(&[]), // in_tid_poslens + reqlens, // hidden_reqlens + arr_u(&[]), // hidden_poslens + ]); + let mut header = Vec::new(); + rmpv::encode::write_value(&mut header, &header_arr).unwrap(); + let mut data = Vec::new(); + data.extend(i(&[10, 20])); // token_ids + data.extend(f(&[-0.5, -0.6])); // out_lp_val (req0) + data.extend(i(&[10, 99])); // out_lp_idx + let framed = frame_egress_batch_cols(&header, &[&data]); + let mut events = Vec::new(); + assert!( + for_each_chunk(&framed[1..], |ev| events.push(ev)).ok, + "frame rejected (zeros={zeros})" + ); + events + }; + + let old = build(true); + let new = build(false); + assert_eq!(old.len(), 2); + assert_eq!(old.len(), new.len()); + for (o, n) in old.iter().zip(&new) { + assert_eq!(o.rid, n.rid); + assert_eq!(o.token_ids, n.token_ids); + match (o.extras.as_deref(), n.extras.as_deref()) { + (Some(a), Some(b)) => { + assert_eq!(a.out_lp_val, b.out_lp_val); + assert_eq!(a.out_lp_idx, b.out_lp_idx); + // Every inactive family stays empty in BOTH spellings. + assert!(a.in_lp_val.is_empty() && b.in_lp_val.is_empty()); + assert!(a.out_top_lens.is_empty() && b.out_top_lens.is_empty()); + assert!(a.in_top_lens.is_empty() && b.in_top_lens.is_empty()); + assert!(a.out_tid_lens.is_empty() && b.out_tid_lens.is_empty()); + assert!(a.in_tid_lens.is_empty() && b.in_tid_lens.is_empty()); + assert!(a.hidden_lens.is_empty() && b.hidden_lens.is_empty()); + } + (None, None) => {} + _ => panic!("extras presence differs between the two spellings"), + } + } + // req0 really did carry its active family through both spellings. + assert_eq!( + new[0] + .extras + .as_deref() + .expect("req0 has out_lp") + .out_lp_val, + vec![-0.5, -0.6] + ); + } + + /// All SEVEN extras families in one frame, each with a distinct length AND + /// distinct values. The existing extras test exercises only `out_lp` / + /// `out_top` / `hidden`, so transposing a header pair — `in_top_*` with + /// `out_tid_*`, say — leaves every assertion passing while the client receives + /// another request's logprobs under the wrong key. Lengths differ per family + /// (2/1/2/1/2/1/3 elements) so a swap misaligns the cursors too. + #[test] + fn decodes_all_extras_families_without_transposition() { + use rmpv::Value; + let f = |xs: &[f32]| -> Vec { xs.iter().flat_map(|x| x.to_le_bytes()).collect() }; + let i = |xs: &[i32]| -> Vec { xs.iter().flat_map(|x| x.to_le_bytes()).collect() }; + let arr_u = |xs: &[u32]| Value::Array(xs.iter().map(|&x| Value::from(x)).collect()); + + let header_arr = Value::Array(vec![ + Value::Array(vec![Value::from("1")]), // rids + Value::Array(vec![Value::Nil]), // finish + arr_u(&[9]), // prompt + arr_u(&[1]), // tok_lens + arr_u(&[2]), // out_lp_lens (2 flat) + arr_u(&[1]), // in_lp_lens (1 flat) + arr_u(&[1]), // out_top_reqlens (1 position… + arr_u(&[2]), // out_top_poslens …k=2) + arr_u(&[1]), // in_top_reqlens (1 position… + arr_u(&[1]), // in_top_poslens …k=1) + arr_u(&[1]), // out_tid_reqlens (1 position… + arr_u(&[2]), // out_tid_poslens …2 ids) + arr_u(&[1]), // in_tid_reqlens (1 position… + arr_u(&[1]), // in_tid_poslens …1 id) + arr_u(&[1]), // hidden_reqlens (1 row… + arr_u(&[3]), // hidden_poslens …dim 3) + ]); + let mut header = Vec::new(); + rmpv::encode::write_value(&mut header, &header_arr).unwrap(); + + // Concatenated in `for_each_chunk`'s column order. + let mut data = Vec::new(); + data.extend(i(&[100])); // token_ids + data.extend(f(&[-1.1, -1.2])); // out_lp_val + data.extend(i(&[11, 12])); // out_lp_idx + data.extend(f(&[-2.1])); // in_lp_val + data.extend(i(&[21])); // in_lp_idx + data.extend(f(&[-3.1, -3.2])); // out_top_val + data.extend(i(&[31, 32])); // out_top_idx + data.extend(f(&[-4.1])); // in_top_val + data.extend(i(&[41])); // in_top_idx + data.extend(f(&[-5.1, -5.2])); // out_tid_val + data.extend(i(&[51, 52])); // out_tid_idx + data.extend(f(&[-6.1])); // in_tid_val + data.extend(i(&[61])); // in_tid_idx + data.extend(f(&[7.1, 7.2, 7.3])); // hidden_val + + let framed = frame_egress_batch_cols(&header, &[&data]); + let mut events = Vec::new(); + assert!(for_each_chunk(&framed[1..], |ev| events.push(ev)).ok); + assert_eq!(events.len(), 1); + let ex = events[0].extras.as_deref().expect("extras present"); + + assert_eq!(ex.out_lp_val, vec![-1.1, -1.2]); + assert_eq!(ex.out_lp_idx, vec![11, 12]); + assert_eq!(ex.in_lp_val, vec![-2.1]); + assert_eq!(ex.in_lp_idx, vec![21]); + assert_eq!(ex.out_top_val, vec![-3.1, -3.2]); + assert_eq!(ex.out_top_idx, vec![31, 32]); + assert_eq!(ex.out_top_lens, vec![2]); + assert_eq!(ex.in_top_val, vec![-4.1]); + assert_eq!(ex.in_top_idx, vec![41]); + assert_eq!(ex.in_top_lens, vec![1]); + assert_eq!(ex.out_tid_val, vec![-5.1, -5.2]); + assert_eq!(ex.out_tid_idx, vec![51, 52]); + assert_eq!(ex.out_tid_lens, vec![2]); + assert_eq!(ex.in_tid_val, vec![-6.1]); + assert_eq!(ex.in_tid_idx, vec![61]); + assert_eq!(ex.in_tid_lens, vec![1]); + assert_eq!(ex.hidden_val, vec![7.1, 7.2, 7.3]); + assert_eq!(ex.hidden_lens, vec![3]); + // Every byte of the data buffer was consumed by exactly one family. + assert_eq!(events[0].token_ids, vec![100]); + } + + /// Two DISTINCT rids that hash to the same shard must stay separate requests. + /// Identity is the rid string now, so a shard-hash collision can only co-locate + /// them — it can no longer merge their `DetokState`, which used to evict one + /// client's sink and deliver their tokens to the other's connection. + #[test] + fn colliding_rids_stay_distinct_requests() { + use rmpv::Value; + let header_arr = Value::Array(vec![ + Value::Array(vec![Value::from("alice"), Value::from("bob")]), + Value::Array(vec![Value::Nil, Value::Nil]), + Value::Array(vec![Value::from(0u32), Value::from(0u32)]), + Value::Array(vec![Value::from(1u32), Value::from(1u32)]), + ]); + let mut header = Vec::new(); + rmpv::encode::write_value(&mut header, &header_arr).unwrap(); + let data: Vec = [7i32, 8].iter().flat_map(|x| x.to_le_bytes()).collect(); + let framed = frame_egress_batch_cols(&header, &[&data]); + let mut events = Vec::new(); + assert!(for_each_chunk(&framed[1..], |ev| events.push(ev)).ok); + // Each chunk carries its OWN rid — the value a shard keys its table on. + assert_eq!(events[0].rid, Rid::from("alice")); + assert_eq!(events[1].rid, Rid::from("bob")); + assert_ne!(events[0].rid, events[1].rid); + } + + /// The common frame must stay small: logprob/hidden columns are boxed behind + /// `ChunkExtras`, so the inline decode array is a few KiB — not MiB — even at + /// batch 4096. A regression that inlines a rare column would blow this up. + /// + /// The budget moved 128 → 144 when the rid gained its memoized partition key + /// (`String` + `u64`). That costs 8 bytes × batch per decode step and buys not + /// re-hashing the rid on every chunk in the egress bucketing loop — a + /// deliberate trade, not drift. + #[test] + fn chunk_event_frame_stays_small() { + let sz = std::mem::size_of::(); + assert!( + sz <= 144, + "ChunkEvent grew to {sz} bytes; keep rare columns behind ChunkExtras" + ); + } +} + +#[cfg(test)] +mod rid_recovery_tests { + use super::*; + + /// A header this build cannot type-decode (Python appended a column) must + /// still name its requests, or every one of them hangs. The previous + /// `(Vec,)` tuple decoded ONLY a 1-element array, so it failed on every + /// real header — arity independence is the entire point of this path. + #[test] + fn rid_recovery_works_at_every_header_arity() { + use rmpv::Value; + for extra_cols in [0usize, 3, 15, 16] { + let mut cols = vec![Value::Array(vec![Value::from("a"), Value::from("b")])]; + // Columns of a type this build would reject (strings where u32 is + // expected) — the "Python widened a column" case. + cols.extend((0..extra_cols).map(|_| Value::from("unexpected"))); + let mut header = Vec::new(); + rmpv::encode::write_value(&mut header, &Value::Array(cols)).unwrap(); + let framed = frame_egress_batch_cols(&header, &[]); + let decoded = for_each_chunk(&framed[1..], |_| {}); + assert!(!decoded.ok, "arity {extra_cols}: must reject"); + assert_eq!( + decoded.rids, + vec![Rid::from("a"), Rid::from("b")], + "arity {extra_cols}: rids must survive so the caller can fail them" + ); + } + } +} diff --git a/rust/sglang-server/src/message/finish_reason.rs b/rust/sglang-server/src/message/finish_reason.rs new file mode 100644 index 000000000..6f9b58503 --- /dev/null +++ b/rust/sglang-server/src/message/finish_reason.rs @@ -0,0 +1,217 @@ +//! The terminal finish reason: Python's `FinishReasonDict` — what +//! `BaseFinishReason.to_json()` (schedule_batch.py) puts on the egress wire, and +//! what the API echoes back as `meta_info.finish_reason`. Ingress has no +//! counterpart; it rides in the [`BatchHeader`](super::egress::BatchHeader) and on +//! each terminal [`ChunkEvent`](super::ChunkEvent). + +use serde::{Deserialize, Serialize}; + +/// The stop that ended a request — Python's `matched` key, typed +/// `Union[str, int, List[int]]`: a stop token id, a stop string +/// (`FINISH_MATCHED_STR` / `FINISHED_MATCHED_REGEX`), or a multi-token stop +/// sequence. Untagged because the three wire shapes are disjoint, so the shape +/// alone picks the arm. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(untagged)] +pub enum Matched { + Token(i64), + Str(String), + Tokens(Vec), +} + +/// The finish reasons this build knows, keyed by the `type` tag Python writes. +/// Every field is optional so a reason missing one still classifies (and keeps its +/// tag) instead of falling through to [`FinishReason::Unknown`]. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "lowercase")] +pub enum FinishKind { + /// `FINISH_MATCHED_TOKEN` / `FINISH_MATCHED_STR` / `FINISHED_MATCHED_REGEX` — + /// all three report `type: "stop"`, to match the OpenAI API's value. + Stop { + #[serde(default)] + matched: Option, + }, + /// `FINISH_LENGTH` — hit `max_new_tokens` (or the context limit). + Length { + #[serde(default)] + length: Option, + }, + /// `FINISH_ABORT` — a scheduler-side termination. Boxed because it is by far + /// the widest variant and the rarest: unboxed it would set the size of every + /// [`ChunkEvent`], including the plain stop/length ones (see + /// `chunk_event_frame_stays_small`). + Abort(Box), +} + +/// The `FINISH_ABORT` payload. `status_code`/`err_type` are `None` for a plain +/// abort and set for a request error (e.g. over-context → 400); Python emits all +/// three keys, nulls included, so none of them is skipped on the way out. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct AbortReason { + #[serde(default)] + pub message: Option, + #[serde(default)] + pub status_code: Option, + #[serde(default)] + pub err_type: Option, +} + +/// A terminal finish reason exactly as `BaseFinishReason.to_json()` +/// (schedule_batch.py) puts it on the wire — Python's `FinishReasonDict`. It +/// round-trips: the API echoes it verbatim as `meta_info.finish_reason`, so the +/// serialized form must stay key-for-key what Python would have sent. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(untagged)] +pub enum FinishReason { + Known(FinishKind), + /// A `type` this build doesn't know, kept as its raw map and echoed unchanged. + /// This arm is why the outer enum is untagged: a finish reason added Python-side + /// must not fail the header decode, which rejects the whole frame — every + /// request in the batch, not just the one that carried it. + Unknown(serde_json::Map), +} + +impl From for FinishReason { + fn from(kind: FinishKind) -> Self { + FinishReason::Known(kind) + } +} + +impl FinishReason { + /// The stop this request matched, if it stopped on one. `None` for + /// length/abort and for an unknown type. + pub fn matched(&self) -> Option<&Matched> { + match self { + FinishReason::Known(FinishKind::Stop { matched }) => matched.as_ref(), + _ => None, + } + } + + /// `Some((status, message))` when this is an abort carrying a `status_code` — + /// a scheduler-side request error the API surfaces as that HTTP status instead + /// of as a normal completion. A plain abort (no code) reads as `None`. + pub fn abort_status(&self) -> Option<(u16, &str)> { + match self { + FinishReason::Known(FinishKind::Abort(a)) => Some(( + a.status_code?, + a.message.as_deref().unwrap_or("request aborted"), + )), + _ => None, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Parse a finish reason from the map Python puts on the wire, so these tests + /// exercise the real deserialization and not just the classifier. + fn fr(v: serde_json::Value) -> Option { + Some(serde_json::from_value(v).expect("finish reason must parse")) + } + + /// Classify through an `Option`, the shape both API paths actually hold. + /// Owned message so a caller can classify a temporary inline. + fn abort_status(f: &Option) -> Option<(u16, String)> { + let (code, message) = f.as_ref()?.abort_status()?; + Some((code, message.to_string())) + } + + /// The classifier both API paths use: a validation abort yields its + /// `(code, message)` (the streaming path turns this into an SSE error event + /// instead of a normal `Done` frame); anything else yields `None`. + #[test] + fn abort_status_extracts_code_and_message() { + let (code, msg) = abort_status(&fr(serde_json::json!({ + "type": "abort", "message": "over the limit", "status_code": 400 + }))) + .expect("validation abort → (code, message)"); + assert_eq!(code, 400); + assert_eq!(msg, "over the limit"); + // Normal finish, bare abort (no status), and no finish → not an error. + assert!(abort_status(&fr(serde_json::json!({"type": "stop"}))).is_none()); + assert!(abort_status(&fr(serde_json::json!({"type": "abort"}))).is_none()); + assert!(abort_status(&None).is_none()); + } + + /// A normal finish, a bare abort (no status), and no finish are not errors + /// (the unary path returns them as a 200 result frame). + #[test] + fn non_error_finishes_stay_ok() { + assert!(abort_status(&fr(serde_json::json!({"type": "stop", "matched": 5}))).is_none()); + assert!(abort_status(&fr(serde_json::json!({"type": "length", "length": 8}))).is_none()); + assert!( + abort_status(&fr( + serde_json::json!({"type": "abort", "message": "Aborted"}) + )) + .is_none() + ); + assert!(abort_status(&None).is_none()); + } + + /// The `matched` accessor reads only a stop, and reads every shape of one. + #[test] + fn matched_reads_stops_only() { + let m = |v| fr(v).unwrap().matched().cloned(); + assert_eq!( + m(serde_json::json!({"type": "stop", "matched": 9})), + Some(Matched::Token(9)) + ); + assert_eq!( + m(serde_json::json!({"type": "stop", "matched": ""})), + Some(Matched::Str("".into())) + ); + assert_eq!( + m(serde_json::json!({"type": "stop", "matched": [9, 10]})), + Some(Matched::Tokens(vec![9, 10])) + ); + // A stop with no `matched`, and the non-stop reasons. + assert_eq!(m(serde_json::json!({"type": "stop"})), None); + assert_eq!(m(serde_json::json!({"type": "length", "length": 8})), None); + assert_eq!(m(serde_json::json!({"type": "abort"})), None); + } + + /// Python's abort dict always carries `status_code`/`err_type`, null included, + /// and the reason is echoed verbatim into `meta_info.finish_reason` — so the + /// typed form must serialize back key-for-key, nulls and all. A dropped null + /// key is a silently changed response body. + #[test] + fn finish_reason_round_trips_python_shapes() { + for wire in [ + serde_json::json!({"type": "stop", "matched": 9}), + serde_json::json!({"type": "stop", "matched": ""}), + serde_json::json!({"type": "stop", "matched": [9, 10]}), + serde_json::json!({"type": "length", "length": 8}), + serde_json::json!({ + "type": "abort", "message": "Aborted", + "status_code": null, "err_type": null + }), + serde_json::json!({ + "type": "abort", "message": "over the limit", + "status_code": 400, "err_type": "BadRequestError" + }), + ] { + let parsed: FinishReason = serde_json::from_value(wire.clone()).unwrap(); + assert!( + matches!(parsed, FinishReason::Known(_)), + "must classify, not fall back to Unknown: {wire}" + ); + assert_eq!(serde_json::to_value(&parsed).unwrap(), wire); + } + } + + /// A `type` this build doesn't know must not fail the decode — the header holds + /// the whole batch, so a rejected frame drops every request in it, not just the + /// one carrying the new reason. It is kept verbatim and echoed unchanged. + #[test] + fn unknown_finish_type_is_preserved_not_rejected() { + let wire = serde_json::json!({"type": "tool_calls", "name": "search"}); + let parsed: FinishReason = serde_json::from_value(wire.clone()).unwrap(); + assert!(matches!(parsed, FinishReason::Unknown(_))); + assert_eq!(serde_json::to_value(&parsed).unwrap(), wire); + // Unknown reasons classify as neither an abort nor a matched stop. + assert!(parsed.abort_status().is_none()); + assert!(parsed.matched().is_none()); + } +}