diff --git a/rust/sglang-server/src/api_server/native_api.rs b/rust/sglang-server/src/api_server/native_api.rs new file mode 100644 index 000000000..4b1d0bee1 --- /dev/null +++ b/rust/sglang-server/src/api_server/native_api.rs @@ -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 { + 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 { + let timeout = + std::time::Duration::from_secs(crate::environ::env_u64("SGLANG_HEALTH_CHECK_TIMEOUT", 20)); + let probe = get(move |state: State| 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, 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_` 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, + body: Result, 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, + 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, + 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, +) -> (usize, mpsc::Receiver, Vec) { + 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)>, + mut guard: AbortGuard, + incremental: bool, + with_index: bool, +) -> impl futures::Stream { + async_stream::stream! { + use futures::StreamExt; + + let n = receivers.len(); + let rid_strs: Vec = receivers.iter().map(|(rid, _)| rid.clone()).collect(); + let mut accs: Vec = + (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" + ); + } + } +} diff --git a/rust/sglang-server/src/runtime/threads.rs b/rust/sglang-server/src/runtime/threads.rs new file mode 100644 index 000000000..5e56ae4d4 --- /dev/null +++ b/rust/sglang-server/src/runtime/threads.rs @@ -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, + pub(super) tok: Vec, + pub(super) detok: Vec, + pub(super) tm: Vec, +} + +pub(super) fn plan_cores(cfg: &RuntimeConfig) -> Option { + // `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 = 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 = 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 = 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) { + 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>, i: usize) -> Option { + 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, + stage: impl Runnable, + threads: &mut Vec>, +) { + 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( + name: &str, + cores: Option>, + count: usize, + threads: &mut Vec>, + 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); + } +}