add the rust server native api handlers and runtime threads (#32876)
This commit is contained in:
Rain Jiang
@@ -0,0 +1,602 @@
|
||||
//! The native SGLang data-plane endpoints: `/generate` (submit a request, then
|
||||
//! either fold egress frames to one unary JSON response or relay them as SSE
|
||||
//! `data: {json}` … `[DONE]`, byte-compatible with Python
|
||||
//! `http_server.generate_request`) and `/health` + `/health_generate` (which
|
||||
//! round-trip a 1-token generate probe). Frame shaping (`meta_info`, logprob
|
||||
//! tuples, cumulative vs incremental streams) lives here, as does
|
||||
//! generate-request submission (`submit`); the shared `AppState` lives in the
|
||||
//! parent `api_server` module.
|
||||
|
||||
use std::convert::Infallible;
|
||||
|
||||
use axum::{
|
||||
Json, Router,
|
||||
extract::State,
|
||||
extract::rejection::JsonRejection,
|
||||
http::StatusCode,
|
||||
response::{
|
||||
IntoResponse, Response,
|
||||
sse::{Event, Sse},
|
||||
},
|
||||
routing::{get, post},
|
||||
};
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
use super::AppState;
|
||||
use super::frame::{
|
||||
OutputAccumulator, cumulative_frame_string, error_value, frame_value, stream_frame_string,
|
||||
tag_value,
|
||||
};
|
||||
use super::guard::AbortGuard;
|
||||
use super::submit::{pre_submit_error, submit};
|
||||
use crate::environ::env_bool;
|
||||
use crate::ids::Rid;
|
||||
use crate::message::{EgressItem, GenerateBody, GenerateRequest, RequestKind, SamplingParams};
|
||||
|
||||
/// The routes this module owns, mounted by `api_server::serve`.
|
||||
pub(super) fn routes() -> Router<AppState> {
|
||||
Router::new()
|
||||
.route("/generate", post(generate))
|
||||
.merge(health_routes())
|
||||
}
|
||||
|
||||
/// `/health` + `/health_generate`. Both env knobs are resolved ONCE here, at
|
||||
/// router build (server startup) — changing them on a live process needs a
|
||||
/// restart. The deep-probe handler is built once with
|
||||
/// `SGLANG_HEALTH_CHECK_TIMEOUT` frozen in and serves `/health_generate`
|
||||
/// always; `SGLANG_ENABLE_HEALTH_ENDPOINT_GENERATION` (default true, mirroring
|
||||
/// Python) decides whether `/health` shares it or is a plain 200 (routing the
|
||||
/// request already proves the frontend is up).
|
||||
fn health_routes() -> Router<AppState> {
|
||||
let timeout =
|
||||
std::time::Duration::from_secs(crate::environ::env_u64("SGLANG_HEALTH_CHECK_TIMEOUT", 20));
|
||||
let probe = get(move |state: State<AppState>| health_generate(state, timeout));
|
||||
let health = if env_bool("SGLANG_ENABLE_HEALTH_ENDPOINT_GENERATION", true) {
|
||||
probe.clone()
|
||||
} else {
|
||||
get(|| async { StatusCode::OK.into_response() })
|
||||
};
|
||||
Router::new()
|
||||
.route("/health", health)
|
||||
.route("/health_generate", probe)
|
||||
}
|
||||
|
||||
/// `GET /health_generate` — deep health: confirm the scheduler → detok path is
|
||||
/// producing output. 200 iff the egress heartbeat advances within `timeout`
|
||||
/// (from `SGLANG_HEALTH_CHECK_TIMEOUT`, frozen at router build), else 503.
|
||||
/// (`/health` uses the same handler when its env gate is on.)
|
||||
///
|
||||
/// Fires a pre-tokenized 1-token probe (`input_ids = [0]`, skips the tokenizer) so
|
||||
/// an idle pipeline produces a frame, then watches the *global*
|
||||
/// [`AppState::egress_activity`] counter (not the probe's own rid) — so a busy
|
||||
/// server passes immediately and a backlog never false-503s (the analogue of
|
||||
/// Python's `last_receive_tstamp`). The `HEALTH_CHECK` skip + `http_worker_ipc`
|
||||
/// ack are irrelevant here: this single-process server owns the egress ring.
|
||||
async fn health_generate(State(state): State<AppState>, timeout: std::time::Duration) -> Response {
|
||||
let baseline = state
|
||||
.egress_activity
|
||||
.load(std::sync::atomic::Ordering::Relaxed);
|
||||
|
||||
// Fire the probe (the heartbeat is the signal, not its own response). A busy
|
||||
// scheduler skips it with no terminal frame, so its detok registration is
|
||||
// cleaned up only by the `AbortGuard` below.
|
||||
let probe = GenerateRequest {
|
||||
// The `HEALTH_CHECK_<uuid>` rid form
|
||||
rid: Rid::new_health_check(),
|
||||
input_ids: Some(vec![0]),
|
||||
// One greedy token: the cheapest round-trip that still produces a frame.
|
||||
sampling_params: SamplingParams {
|
||||
max_new_tokens: Some(1),
|
||||
temperature: 0.0,
|
||||
..Default::default()
|
||||
},
|
||||
stream: false,
|
||||
..Default::default()
|
||||
};
|
||||
let (rid, _keepalive) =
|
||||
match submit(&state, RequestKind::Generate(Box::new(probe)), false).await {
|
||||
// Hold the receiver so the probe's sink stays open until it completes.
|
||||
Ok(v) => v,
|
||||
Err(resp) => return resp,
|
||||
};
|
||||
// Deregister on drop (never disarmed): a busy-skipped probe has no terminal
|
||||
// frame, so without this abort it leaks one detok entry per call.
|
||||
let _abort_guard = AbortGuard::new(state.senders.clone(), rid);
|
||||
|
||||
// Watch the heartbeat advance (timeout frozen at router build, default 20s).
|
||||
let deadline = tokio::time::Instant::now() + timeout;
|
||||
loop {
|
||||
if state
|
||||
.egress_activity
|
||||
.load(std::sync::atomic::Ordering::Relaxed)
|
||||
!= baseline
|
||||
{
|
||||
return StatusCode::OK.into_response();
|
||||
}
|
||||
if tokio::time::Instant::now() >= deadline {
|
||||
return StatusCode::SERVICE_UNAVAILABLE.into_response();
|
||||
}
|
||||
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
|
||||
}
|
||||
}
|
||||
|
||||
/// `POST /generate` — the native generation endpoint. Splits the body
|
||||
/// into per-request payloads (a scalar body → one, a list body → a batch) and
|
||||
/// dispatches to the single or batch path; a malformed body is a 400 before
|
||||
/// anything reaches the scheduler.
|
||||
///
|
||||
/// The body is extracted as a `Result` so a deserialization failure is answered
|
||||
/// with **400** (Python's status for a bad request) carrying serde's field-level
|
||||
/// message, instead of axum's default 422.
|
||||
async fn generate(
|
||||
State(state): State<AppState>,
|
||||
body: Result<Json<GenerateBody>, JsonRejection>,
|
||||
) -> Response {
|
||||
let body = match body {
|
||||
Ok(Json(body)) => body,
|
||||
// A body that fails to parse has no readable `stream` flag, so this one
|
||||
// can only answer unary — as Python's does (FastAPI rejects before its
|
||||
// handler runs).
|
||||
Err(rejection) => {
|
||||
return pre_submit_error(StatusCode::BAD_REQUEST, &rejection.body_text(), false);
|
||||
}
|
||||
};
|
||||
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() {
|
||||
Ok(v) => v,
|
||||
// The error carries its own status (a bad batch is `Validation` → 400).
|
||||
Err(e) => {
|
||||
let code = StatusCode::from_u16(e.http_status()).unwrap_or(StatusCode::BAD_REQUEST);
|
||||
return pre_submit_error(code, &e.to_string(), stream);
|
||||
}
|
||||
};
|
||||
if !is_batch {
|
||||
// `into_requests` guarantees exactly one payload for a non-batch body.
|
||||
let payload = payloads
|
||||
.into_iter()
|
||||
.next()
|
||||
.expect("into_requests yields >=1 payload");
|
||||
generate_single(&state, payload, stream).await
|
||||
} else {
|
||||
generate_batch(&state, payloads, stream).await
|
||||
}
|
||||
}
|
||||
|
||||
/// Answer an error raised *before* anything was submitted, in the shape the client
|
||||
/// asked for.
|
||||
///
|
||||
/// A single (non-batched) `/generate`: submit one request, then either stream its
|
||||
/// SSE frames or fold to one unary response.
|
||||
async fn generate_single(state: &AppState, req: GenerateRequest, stream: bool) -> Response {
|
||||
// `return_text_in_logprobs` is decoded on the detok shard into `*_txt`, so
|
||||
// `frame_value` just reads them — no tokenizer needed here.
|
||||
let (rid_str, mut rx) = match submit(state, RequestKind::Generate(Box::new(req)), stream).await
|
||||
{
|
||||
Ok(v) => v,
|
||||
Err(resp) => return resp,
|
||||
};
|
||||
// Abort on client disconnect: the guard fires when dropped before the request
|
||||
// finishes (axum drops the handler/SSE stream). Disarmed on a natural terminal.
|
||||
// `rid_str` is the response `meta_info.id`, reused for every frame.
|
||||
let mut guard = AbortGuard::new(state.senders.clone(), rid_str.clone());
|
||||
// Cumulative frames (SGLang default) vs per-step deltas.
|
||||
let incremental = state.server_args.incremental_streaming_output;
|
||||
|
||||
if stream {
|
||||
// A single request is a 1-element batch without the `index` field — reuse
|
||||
// the same stream so the frame/abort/truncation logic lives in one place.
|
||||
use futures::StreamExt;
|
||||
let s = generation_event_stream(vec![(rid_str, rx)], guard, incremental, false)
|
||||
.map(|data| Ok::<_, Infallible>(Event::default().data(data)));
|
||||
Sse::new(s).into_response()
|
||||
} else {
|
||||
// Unary: fold to the terminal, respond once. Disarm only on a real terminal
|
||||
// (a truncation leaves the guard armed so the scheduler work is aborted).
|
||||
let (status, value, terminal) = drain_unary(&mut rx, rid_str.client_facing()).await;
|
||||
if terminal {
|
||||
guard.disarm(&rid_str);
|
||||
}
|
||||
(status, Json(value)).into_response()
|
||||
}
|
||||
}
|
||||
|
||||
/// Fold a unary request to its terminal → (HTTP status, result/`error` JSON, saw-terminal);
|
||||
/// `false` = truncation, caller keeps the abort guard armed. Shared by single + batch.
|
||||
async fn drain_unary(
|
||||
rx: &mut mpsc::Receiver<EgressItem>,
|
||||
rid_str: &str,
|
||||
) -> (StatusCode, serde_json::Value, bool) {
|
||||
let mut acc = OutputAccumulator::default();
|
||||
while let Some(item) = rx.recv().await {
|
||||
match item {
|
||||
EgressItem::Frame(out) => acc.fold(&out),
|
||||
EgressItem::Done(out) => {
|
||||
acc.fold(&out);
|
||||
let final_out = acc.into_output();
|
||||
// A validation abort carries its own HTTP status + diagnostic.
|
||||
if let Some((code, message)) = final_out
|
||||
.finish_reason
|
||||
.as_ref()
|
||||
.and_then(|f| f.abort_status())
|
||||
{
|
||||
let status =
|
||||
StatusCode::from_u16(code).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
|
||||
return (status, error_value(code, message), true);
|
||||
}
|
||||
return (StatusCode::OK, frame_value(&final_out, rid_str), true);
|
||||
}
|
||||
EgressItem::Error(e) => {
|
||||
let code = e.http_status();
|
||||
let status =
|
||||
StatusCode::from_u16(code).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
|
||||
return (status, error_value(code, &e.to_string()), true);
|
||||
}
|
||||
EgressItem::Control(_) => continue, // never on `/generate`
|
||||
}
|
||||
}
|
||||
// Sender dropped without a terminal item: the shard dropped this request (a
|
||||
// truncation — a client disconnect would have dropped the handler future).
|
||||
(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
error_value(500, "response truncated before completion"),
|
||||
false,
|
||||
)
|
||||
}
|
||||
|
||||
/// Batch `/generate`: submit all sub-requests first (scheduler runs them together),
|
||||
/// then either (unary) drain each in order into a JSON array, or (streaming)
|
||||
/// multiplex their streams into one SSE response, each frame carrying its `index`.
|
||||
/// One [`AbortGuard`] covers the batch. A failed unary item is its own
|
||||
/// `{ "error": … }` entry; the batch response is 200.
|
||||
async fn generate_batch(
|
||||
state: &AppState,
|
||||
requests: Vec<GenerateRequest>,
|
||||
stream: bool,
|
||||
) -> Response {
|
||||
// No cross-item rid collision to worry about: `into_requests` rejected duplicate
|
||||
// rids within this batch, and `Rid::from_client` made each one unique against
|
||||
// every other in-flight request.
|
||||
let mut guard = AbortGuard::new_empty(state.senders.clone());
|
||||
let mut receivers = Vec::with_capacity(requests.len());
|
||||
for req in requests {
|
||||
match submit(state, RequestKind::Generate(Box::new(req)), stream).await {
|
||||
Ok((rid, rx)) => {
|
||||
guard.arm(rid.clone());
|
||||
receivers.push((rid, rx));
|
||||
}
|
||||
Err(resp) => return resp,
|
||||
}
|
||||
}
|
||||
|
||||
if stream {
|
||||
// Multiplex the N streams (mirrors the Python `_handle_batch_request` path);
|
||||
// `guard` moves into the stream so a disconnect aborts what's unfinished.
|
||||
use futures::StreamExt;
|
||||
let incremental = state.server_args.incremental_streaming_output;
|
||||
let s = generation_event_stream(receivers, guard, incremental, true)
|
||||
.map(|data| Ok::<_, Infallible>(Event::default().data(data)));
|
||||
Sse::new(s).into_response()
|
||||
} else {
|
||||
// Unary: drain each in order (already all submitted, so they run together).
|
||||
let mut results = Vec::with_capacity(receivers.len());
|
||||
for (rid_str, mut rx) in receivers {
|
||||
let (_status, value, terminal) = drain_unary(&mut rx, rid_str.client_facing()).await;
|
||||
if terminal {
|
||||
guard.disarm(&rid_str);
|
||||
}
|
||||
results.push(value);
|
||||
}
|
||||
(StatusCode::OK, Json(serde_json::Value::Array(results))).into_response()
|
||||
}
|
||||
}
|
||||
|
||||
/// Await the next item from `rx`, then drain whatever queued behind it (so the caller
|
||||
/// can coalesce a backlog, as Python's `state.out_list` does), handing the receiver
|
||||
/// back for `FuturesUnordered` to re-poll. Empty result = channel closed.
|
||||
async fn recv_indexed(
|
||||
index: usize,
|
||||
mut rx: mpsc::Receiver<EgressItem>,
|
||||
) -> (usize, mpsc::Receiver<EgressItem>, Vec<EgressItem>) {
|
||||
let mut items = Vec::new();
|
||||
match rx.recv().await {
|
||||
Some(item) => items.push(item),
|
||||
None => return (index, rx, items), // closed
|
||||
}
|
||||
while let Ok(item) = rx.try_recv() {
|
||||
items.push(item);
|
||||
}
|
||||
(index, rx, items)
|
||||
}
|
||||
|
||||
/// Multiplex `receivers` (one per request) into SSE `data` strings + a final `[DONE]`;
|
||||
/// `with_index` tags each frame (batch only), `incremental` = delta vs cumulative,
|
||||
/// `guard` aborts unfinished on drop.
|
||||
fn generation_event_stream(
|
||||
receivers: Vec<(Rid, mpsc::Receiver<EgressItem>)>,
|
||||
mut guard: AbortGuard,
|
||||
incremental: bool,
|
||||
with_index: bool,
|
||||
) -> impl futures::Stream<Item = String> {
|
||||
async_stream::stream! {
|
||||
use futures::StreamExt;
|
||||
|
||||
let n = receivers.len();
|
||||
let rid_strs: Vec<Rid> = receivers.iter().map(|(rid, _)| rid.clone()).collect();
|
||||
let mut accs: Vec<OutputAccumulator> =
|
||||
(0..n).map(|_| OutputAccumulator::default()).collect();
|
||||
|
||||
// Batch position, tagged onto every frame (a single request omits it).
|
||||
let idx = |i: usize| with_index.then_some(i);
|
||||
|
||||
// Poll all receivers concurrently; re-arm a receiver's future after each
|
||||
// non-terminal frame so its stream keeps flowing.
|
||||
let mut futs = futures::stream::FuturesUnordered::new();
|
||||
for (i, (_, rx)) in receivers.into_iter().enumerate() {
|
||||
futs.push(recv_indexed(i, rx));
|
||||
}
|
||||
|
||||
while let Some((i, rx, items)) = futs.next().await {
|
||||
if items.is_empty() {
|
||||
// Channel closed with no terminal → truncation for this item;
|
||||
// leave its rid armed so the scheduler work is aborted.
|
||||
yield tag_value(error_value(500, "response truncated before completion"), idx(i));
|
||||
continue;
|
||||
}
|
||||
|
||||
// Cumulative frames supersede one another, so a drained backlog collapses
|
||||
// to its last (Python's `out_list[-1]`); deltas can't be dropped.
|
||||
let mut coalesced = false; // a cumulative frame is pending
|
||||
let mut terminal = None; // (finish_reason) of a `Done` in this batch
|
||||
let mut failed = None; // an `Error` in this batch
|
||||
|
||||
for item in items {
|
||||
match item {
|
||||
EgressItem::Frame(out) => {
|
||||
accs[i].fold(&out);
|
||||
if incremental {
|
||||
yield stream_frame_string(out, &accs[i], true, rid_strs[i].client_facing(), idx(i));
|
||||
} else {
|
||||
coalesced = true;
|
||||
}
|
||||
}
|
||||
EgressItem::Done(out) => {
|
||||
accs[i].fold(&out);
|
||||
terminal = Some(out);
|
||||
}
|
||||
EgressItem::Error(e) => failed = Some(e),
|
||||
EgressItem::Control(_) => {} // never on /generate
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(e) = failed {
|
||||
yield tag_value(error_value(e.http_status(), &e.to_string()), idx(i));
|
||||
guard.disarm(&rid_strs[i]);
|
||||
} else if let Some(out) = terminal {
|
||||
// A validation abort → an error object, not a frame. The final frame
|
||||
// carries the full cumulative state, so any coalesced ones are moot.
|
||||
yield match out.finish_reason.as_ref().and_then(|f| f.abort_status()) {
|
||||
Some((code, message)) => tag_value(error_value(code, message), idx(i)),
|
||||
None => stream_frame_string(out, &accs[i], incremental, rid_strs[i].client_facing(), idx(i)),
|
||||
};
|
||||
guard.disarm(&rid_strs[i]); // terminal → not re-pushed
|
||||
} else {
|
||||
if coalesced {
|
||||
yield cumulative_frame_string(&accs[i], rid_strs[i].client_facing(), idx(i));
|
||||
}
|
||||
futs.push(recv_indexed(i, rx)); // keep this item flowing
|
||||
}
|
||||
}
|
||||
yield "[DONE]".to_string();
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::message::ChunkEvent;
|
||||
use crate::tokenizer_manager::Senders;
|
||||
use futures::StreamExt;
|
||||
fn senders() -> Senders {
|
||||
Senders {
|
||||
tm: flume::unbounded().0,
|
||||
abort: flume::unbounded().0,
|
||||
tok: flume::unbounded().0,
|
||||
detok: vec![],
|
||||
}
|
||||
}
|
||||
|
||||
fn frame(rid: u64, text: &str) -> EgressItem {
|
||||
EgressItem::Frame(ChunkEvent {
|
||||
rid: Rid::from(rid.to_string()),
|
||||
text: text.into(),
|
||||
completion_tokens: 1,
|
||||
..Default::default()
|
||||
})
|
||||
}
|
||||
fn done(rid: u64, text: &str) -> EgressItem {
|
||||
EgressItem::Done(ChunkEvent {
|
||||
rid: Rid::from(rid.to_string()),
|
||||
text: text.into(),
|
||||
completion_tokens: 1,
|
||||
// Parsed from the wire map Python emits, not a hand-built enum.
|
||||
finish_reason: Some(
|
||||
serde_json::from_value(serde_json::json!({"type": "length", "length": 1}))
|
||||
.expect("finish reason must parse"),
|
||||
),
|
||||
..Default::default()
|
||||
})
|
||||
}
|
||||
fn parse(s: &str) -> serde_json::Value {
|
||||
serde_json::from_str(s).expect("frame is JSON")
|
||||
}
|
||||
|
||||
/// Two sub-requests' frames interleave into one stream, each tagged with its
|
||||
/// batch `index`; text accumulates per item; `[DONE]` comes only after both
|
||||
/// terminate, then the stream ends.
|
||||
#[tokio::test]
|
||||
async fn interleaves_indexes_and_accumulates() {
|
||||
let (tx0, rx0) = mpsc::channel(8);
|
||||
let (tx1, rx1) = mpsc::channel(8);
|
||||
let receivers = vec![("10".into(), rx0), ("11".into(), rx1)];
|
||||
let stream =
|
||||
generation_event_stream(receivers, AbortGuard::new_empty(senders()), false, true);
|
||||
futures::pin_mut!(stream);
|
||||
|
||||
// Drive deterministically: exactly one channel has data before each poll.
|
||||
tx0.send(frame(10, "a")).await.unwrap();
|
||||
let v = parse(&stream.next().await.unwrap());
|
||||
assert_eq!(v["index"], 0);
|
||||
assert_eq!(v["text"], "a");
|
||||
|
||||
tx1.send(frame(11, "b")).await.unwrap();
|
||||
let v = parse(&stream.next().await.unwrap());
|
||||
assert_eq!(v["index"], 1);
|
||||
assert_eq!(v["text"], "b");
|
||||
|
||||
tx0.send(done(10, "!")).await.unwrap();
|
||||
let v = parse(&stream.next().await.unwrap());
|
||||
assert_eq!(v["index"], 0);
|
||||
assert_eq!(v["text"], "a!", "cumulative per item");
|
||||
assert_eq!(v["meta_info"]["finish_reason"]["type"], "length");
|
||||
|
||||
tx1.send(done(11, "?")).await.unwrap();
|
||||
let v = parse(&stream.next().await.unwrap());
|
||||
assert_eq!(v["index"], 1);
|
||||
assert_eq!(v["text"], "b?");
|
||||
|
||||
assert_eq!(stream.next().await.unwrap(), "[DONE]");
|
||||
assert!(stream.next().await.is_none());
|
||||
}
|
||||
|
||||
/// A per-item error is surfaced with its `index` and doesn't end the batch;
|
||||
/// `[DONE]` still waits for the other item.
|
||||
#[tokio::test]
|
||||
async fn per_item_error_carries_index() {
|
||||
let (tx0, rx0) = mpsc::channel(8);
|
||||
let (tx1, rx1) = mpsc::channel(8);
|
||||
let receivers = vec![("10".into(), rx0), ("11".into(), rx1)];
|
||||
let stream =
|
||||
generation_event_stream(receivers, AbortGuard::new_empty(senders()), false, true);
|
||||
futures::pin_mut!(stream);
|
||||
|
||||
tx0.send(EgressItem::Error(crate::error::Error::Validation(
|
||||
"bad".into(),
|
||||
)))
|
||||
.await
|
||||
.unwrap();
|
||||
let v = parse(&stream.next().await.unwrap());
|
||||
assert_eq!(v["index"], 0);
|
||||
assert_eq!(v["error"]["code"], 400);
|
||||
|
||||
tx1.send(done(11, "ok")).await.unwrap();
|
||||
let v = parse(&stream.next().await.unwrap());
|
||||
assert_eq!(v["index"], 1);
|
||||
|
||||
assert_eq!(stream.next().await.unwrap(), "[DONE]");
|
||||
}
|
||||
|
||||
/// `incremental=true`: each frame carries this step's **delta** text/output_ids,
|
||||
/// but `meta_info.completion_tokens` stays cumulative (matching Python).
|
||||
#[tokio::test]
|
||||
async fn incremental_emits_deltas_with_cumulative_count() {
|
||||
let (tx, rx) = mpsc::channel(8);
|
||||
let receivers = vec![("10".into(), rx)];
|
||||
let stream =
|
||||
generation_event_stream(receivers, AbortGuard::new_empty(senders()), true, true);
|
||||
futures::pin_mut!(stream);
|
||||
|
||||
tx.send(frame(10, "Hello")).await.unwrap();
|
||||
let v = parse(&stream.next().await.unwrap());
|
||||
assert_eq!(v["text"], "Hello");
|
||||
assert_eq!(v["meta_info"]["completion_tokens"], 1);
|
||||
|
||||
tx.send(frame(10, " world")).await.unwrap();
|
||||
let v = parse(&stream.next().await.unwrap());
|
||||
assert_eq!(v["text"], " world", "delta, not cumulative 'Hello world'");
|
||||
assert_eq!(
|
||||
v["meta_info"]["completion_tokens"], 2,
|
||||
"count stays cumulative"
|
||||
);
|
||||
|
||||
tx.send(done(10, "!")).await.unwrap();
|
||||
let v = parse(&stream.next().await.unwrap());
|
||||
assert_eq!(v["text"], "!");
|
||||
assert_eq!(v["meta_info"]["completion_tokens"], 3);
|
||||
assert_eq!(v["meta_info"]["finish_reason"]["type"], "length");
|
||||
|
||||
assert_eq!(stream.next().await.unwrap(), "[DONE]");
|
||||
}
|
||||
|
||||
/// The single-request shape (`with_index=false`, one receiver) omits the
|
||||
/// `index` field entirely, and still terminates with `[DONE]`.
|
||||
#[tokio::test]
|
||||
async fn single_shape_omits_index() {
|
||||
let (tx, rx) = mpsc::channel(8);
|
||||
let receivers = vec![("10".into(), rx)];
|
||||
let stream =
|
||||
generation_event_stream(receivers, AbortGuard::new_empty(senders()), false, false);
|
||||
futures::pin_mut!(stream);
|
||||
|
||||
tx.send(done(10, "hi")).await.unwrap();
|
||||
let v = parse(&stream.next().await.unwrap());
|
||||
assert_eq!(v["text"], "hi");
|
||||
assert!(v.get("index").is_none(), "single response has no index");
|
||||
|
||||
assert_eq!(stream.next().await.unwrap(), "[DONE]");
|
||||
}
|
||||
|
||||
/// A backlog of cumulative chunks collapses to a single frame carrying the latest
|
||||
/// state — each cumulative frame supersedes the last, so emitting the intermediate
|
||||
/// ones ships the full O(T) payload again for nothing. Mirrors the Python waiter's
|
||||
/// `out = out_list[-1]`. This is the whole point of draining in `recv_indexed`.
|
||||
#[tokio::test]
|
||||
async fn cumulative_backlog_coalesces_to_latest() {
|
||||
let (tx, rx) = mpsc::channel(8);
|
||||
let receivers = vec![("10".into(), rx)];
|
||||
let stream =
|
||||
generation_event_stream(receivers, AbortGuard::new_empty(senders()), false, false);
|
||||
futures::pin_mut!(stream);
|
||||
|
||||
// Three chunks queued before the stream is ever polled (a client falling behind).
|
||||
tx.send(frame(10, "a")).await.unwrap();
|
||||
tx.send(frame(10, "b")).await.unwrap();
|
||||
tx.send(frame(10, "c")).await.unwrap();
|
||||
|
||||
let v = parse(&stream.next().await.unwrap());
|
||||
assert_eq!(v["text"], "abc", "one frame, full cumulative text");
|
||||
assert_eq!(v["meta_info"]["completion_tokens"], 3, "no tokens lost");
|
||||
|
||||
// The terminal frame still carries everything, and only then does [DONE] land.
|
||||
tx.send(done(10, "!")).await.unwrap();
|
||||
let v = parse(&stream.next().await.unwrap());
|
||||
assert_eq!(v["text"], "abc!");
|
||||
assert_eq!(v["meta_info"]["finish_reason"]["type"], "length");
|
||||
assert_eq!(stream.next().await.unwrap(), "[DONE]");
|
||||
}
|
||||
|
||||
/// Incremental frames are *deltas*, so a backlog must emit every one — dropping
|
||||
/// any would silently lose tokens. Only the cumulative protocol may coalesce.
|
||||
#[tokio::test]
|
||||
async fn incremental_backlog_emits_every_delta() {
|
||||
let (tx, rx) = mpsc::channel(8);
|
||||
let receivers = vec![("10".into(), rx)];
|
||||
let stream =
|
||||
generation_event_stream(receivers, AbortGuard::new_empty(senders()), true, false);
|
||||
futures::pin_mut!(stream);
|
||||
|
||||
tx.send(frame(10, "a")).await.unwrap();
|
||||
tx.send(frame(10, "b")).await.unwrap();
|
||||
tx.send(frame(10, "c")).await.unwrap();
|
||||
|
||||
for (n, expect) in [(1, "a"), (2, "b"), (3, "c")] {
|
||||
let v = parse(&stream.next().await.unwrap());
|
||||
assert_eq!(v["text"], expect, "delta {n} must not be dropped");
|
||||
assert_eq!(
|
||||
v["meta_info"]["completion_tokens"], n,
|
||||
"count stays cumulative"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,142 @@
|
||||
//! Thread-group machinery: CPU-core partitioning and the pinned-thread
|
||||
//! spawners (for [`Runnable`] stages) used by `runtime::start`.
|
||||
//!
|
||||
//! Adding a new thread group (encoder, weight loader, KV-cache offloader, …) is
|
||||
//! three small steps and no spawn boilerplate:
|
||||
//! 1. a struct implementing [`Runnable`];
|
||||
//! 2. a core set for it (a field on [`CorePlan`] + a slice in [`plan_cores`]);
|
||||
//! 3. one [`spawn_pool`] (N pinned workers) or [`spawn_stage`] (singleton) call.
|
||||
|
||||
use std::thread::JoinHandle;
|
||||
|
||||
use core_affinity::CoreId;
|
||||
|
||||
use super::{Runnable, RuntimeConfig};
|
||||
|
||||
/// Cores reserved for the two TokenizerManager router threads (`tm-ingress`,
|
||||
/// `tm-egress`) — light, latency-sensitive channel routers, so one core each.
|
||||
///
|
||||
/// TODO(tm-scaling): both TM threads are single-consumer serialization points,
|
||||
/// each with its own ceiling. `tm-ingress` runs validate + `normalize_sampling_params`
|
||||
/// for *every* request before fanning out to the (pooled) tokenizer workers, so a
|
||||
/// high request-arrival / short-request workload is bounded by that one thread's
|
||||
/// per-request cost (kept O(fields), see `sampling::normalize_sampling_params`).
|
||||
/// Sharding ingress by rid — like the tokenizer/detok pools — lifts that ceiling.
|
||||
///
|
||||
/// `tm-egress` is a head-of-line ceiling of a different kind — it
|
||||
/// does a *blocking* send per chunk to the owning detok shard, so one slow shard
|
||||
/// stalls the dispatcher and thus every shard (see `Egress::route`). Sharding the
|
||||
/// dispatcher alone doesn't fix it: each egress-ring frame is a whole batch fanned
|
||||
/// to *all* shards, so any dispatcher still blocks on the slow one. The real fix
|
||||
/// is a per-shard egress ring (the scheduler pushing each request's output to its
|
||||
/// shard's ring), each drained by its own dispatcher — at which point this needs
|
||||
/// one core per ingress/egress shard rather than a fixed 2.
|
||||
const TM_CORES: usize = 2;
|
||||
|
||||
/// Partition the machine's cores into four disjoint sets: the I/O-bound API
|
||||
/// pool, the CPU-bound tokenizer and detokenizer pools, and the two TM router
|
||||
/// threads. Falls back to no pinning if affinity isn't available or there aren't
|
||||
/// enough cores for the (CPU-bound) pools. A new thread group adds a field here
|
||||
/// and a slice in [`plan_cores`].
|
||||
pub(super) struct CorePlan {
|
||||
pub(super) api: Vec<CoreId>,
|
||||
pub(super) tok: Vec<CoreId>,
|
||||
pub(super) detok: Vec<CoreId>,
|
||||
pub(super) tm: Vec<CoreId>,
|
||||
}
|
||||
|
||||
pub(super) fn plan_cores(cfg: &RuntimeConfig) -> Option<CorePlan> {
|
||||
// `cores` carries the pinning decision: `None`/empty → run unpinned. The
|
||||
// caller (Python `_partition_cores`) passes this rank's NUMA-local cores
|
||||
// minus the scheduler's reserved launch cores.
|
||||
let cores: Vec<CoreId> = match &cfg.rust_server_args.cores {
|
||||
Some(ids) if !ids.is_empty() => ids.iter().map(|&id| CoreId { id }).collect(),
|
||||
_ => return None,
|
||||
};
|
||||
if cores.len()
|
||||
< cfg.rust_server_args.api_worker_num
|
||||
+ cfg.server_args.tokenizer_worker_num
|
||||
+ cfg.server_args.detokenizer_worker_num
|
||||
{
|
||||
tracing::warn!(
|
||||
available = cores.len(),
|
||||
"not enough cores to pin all pools; running unpinned"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
let mut it = cores.into_iter();
|
||||
let api: Vec<CoreId> = it
|
||||
.by_ref()
|
||||
.take(cfg.rust_server_args.api_worker_num)
|
||||
.collect();
|
||||
let tok = it
|
||||
.by_ref()
|
||||
.take(cfg.server_args.tokenizer_worker_num)
|
||||
.collect();
|
||||
let detok = it
|
||||
.by_ref()
|
||||
.take(cfg.server_args.detokenizer_worker_num)
|
||||
.collect();
|
||||
// The two TM router threads get up to `TM_CORES` leftover cores; when none
|
||||
// are spare they fall back to the API set so they never float onto the
|
||||
// CPU-bound tokenizer/detok cores.
|
||||
let mut tm: Vec<CoreId> = it.by_ref().take(TM_CORES).collect();
|
||||
if tm.is_empty() {
|
||||
tm = api.clone();
|
||||
}
|
||||
Some(CorePlan {
|
||||
api,
|
||||
tok,
|
||||
detok,
|
||||
tm,
|
||||
})
|
||||
}
|
||||
|
||||
/// Pin the calling thread to `core` if one was assigned (no-op otherwise).
|
||||
fn pin_current(core: Option<CoreId>) {
|
||||
if let Some(c) = core {
|
||||
core_affinity::set_for_current(c);
|
||||
}
|
||||
}
|
||||
|
||||
/// Pick the pinned core for worker `i` from an optional pool core set.
|
||||
pub(super) fn pool_core(cores: &Option<Vec<CoreId>>, i: usize) -> Option<CoreId> {
|
||||
cores.as_ref().and_then(|c| c.get(i).copied())
|
||||
}
|
||||
|
||||
/// Spawn a single [`Runnable`] stage on a named thread, optionally pinned.
|
||||
/// Used by [`spawn_pool`]; every group goes through the pool spawner now.
|
||||
fn spawn_stage(
|
||||
name: &str,
|
||||
core: Option<CoreId>,
|
||||
stage: impl Runnable,
|
||||
threads: &mut Vec<JoinHandle<()>>,
|
||||
) {
|
||||
let handle = std::thread::Builder::new()
|
||||
.name(name.to_string())
|
||||
.spawn(move || {
|
||||
pin_current(core);
|
||||
stage.run();
|
||||
})
|
||||
.expect("spawn stage");
|
||||
threads.push(handle);
|
||||
}
|
||||
|
||||
/// Spawn a pool of `count` [`Runnable`] workers, each pinned to `cores[i]` (when
|
||||
/// available) and named `{name}-{i}`. `build(i)` constructs worker `i` — cloning
|
||||
/// shared handles, or moving a per-worker resource out of a captured iterator.
|
||||
pub(super) fn spawn_pool<R, F>(
|
||||
name: &str,
|
||||
cores: Option<Vec<CoreId>>,
|
||||
count: usize,
|
||||
threads: &mut Vec<JoinHandle<()>>,
|
||||
mut build: F,
|
||||
) where
|
||||
R: Runnable,
|
||||
F: FnMut(usize) -> R,
|
||||
{
|
||||
for i in 0..count {
|
||||
let core = pool_core(&cores, i);
|
||||
spawn_stage(&format!("{name}-{i}"), core, build(i), threads);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user