From d7a4c830e55b2f85a262f2475dff4a7465667455 Mon Sep 17 00:00:00 2001 From: Rain Jiang <96632942+rainj-me@users.noreply.github.com> Date: Wed, 29 Jul 2026 18:42:12 -0700 Subject: [PATCH] sglang rust server tokenizer manager, ring and runtime (#32358) --- rust/sglang-server/src/fsm.rs | 85 ++++++- rust/sglang-server/src/lib.rs | 3 + rust/sglang-server/src/message.rs | 61 ++++- rust/sglang-server/src/message/request.rs | 6 +- rust/sglang-server/src/ring.rs | 251 ++++++++++++++++++++ rust/sglang-server/src/runtime.rs | 18 ++ rust/sglang-server/src/runtime/config.rs | 203 ++++++++++++++++ rust/sglang-server/src/runtime/runnable.rs | 10 + rust/sglang-server/src/tokenizer_manager.rs | 87 +++++++ 9 files changed, 713 insertions(+), 11 deletions(-) create mode 100644 rust/sglang-server/src/ring.rs create mode 100644 rust/sglang-server/src/runtime.rs create mode 100644 rust/sglang-server/src/runtime/config.rs create mode 100644 rust/sglang-server/src/runtime/runnable.rs create mode 100644 rust/sglang-server/src/tokenizer_manager.rs diff --git a/rust/sglang-server/src/fsm.rs b/rust/sglang-server/src/fsm.rs index 9206fa708..afb1f58a7 100644 --- a/rust/sglang-server/src/fsm.rs +++ b/rust/sglang-server/src/fsm.rs @@ -7,10 +7,10 @@ //! //! Port of the design enum: //! ```text -//! Received, Validating, Normalizing, Encoding, Tokenizing, Queued, -//! Streaming { chunks_sent }, Finalizing, Completed, Failed(Error), Aborted +//! Received, Validating, Normalizing, Encoding, Tokenizing, PreSendValidating, +//! Queued, Streaming { chunks_sent }, Finalizing, Completed, Failed(Error), +//! Aborted //! ``` -#![allow(dead_code)] // TODO: remove when the consumer PR lands use crate::error::Error; @@ -25,6 +25,10 @@ pub enum RequestState { Normalizing, Encoding, Tokenizing, + /// Every branch converges here with its final `input_ids`, for the checks + /// that need the tokenized length (the input + `max_new_tokens` ceiling). + /// The last state before the request leaves Rust. + PreSendValidating, Queued, Streaming { chunks_sent: u64, @@ -43,7 +47,7 @@ pub enum ValidationOutcome { HasMultimodal, /// Plain text → Tokenizing. NeedsTokenize, - /// Caller already supplied token ids → straight to Queued. + /// Caller already supplied token ids → straight to the pre-send checks. AlreadyTokenized, } @@ -57,9 +61,13 @@ pub enum Event { NeedsNormalize, EncodeDone, TokenizeDone, + /// The pre-send checks passed; the request may be pushed to the ring. + PreSendValidated, SchedulerPicked, // --- egress --- - Chunk { finish: bool }, + Chunk { + finish: bool, + }, FinalFrameSent, // --- terminal (valid from any state) --- Error(Error), @@ -112,14 +120,17 @@ impl RequestState { // ingress (Received, Validated(_)) => Validating, // Generate requests pass through Normalizing (sampling-param - // normalize/verify); control requests skip straight to Queued. + // normalize/verify); control requests skip it, having none. (Validating, NeedsNormalize) => Normalizing, - (Validating, Validated(AlreadyTokenized)) => Queued, + (Validating, Validated(AlreadyTokenized)) => PreSendValidating, (Normalizing, Validated(HasMultimodal)) => Encoding, (Normalizing, Validated(NeedsTokenize)) => Tokenizing, - (Normalizing, Validated(AlreadyTokenized)) => Queued, + (Normalizing, Validated(AlreadyTokenized)) => PreSendValidating, (Encoding, EncodeDone) => Tokenizing, - (Tokenizing, TokenizeDone) => Queued, + // Every ingress branch funnels through the pre-send checks, so they + // run exactly once per request no matter how it got its ids. + (Tokenizing, TokenizeDone) => PreSendValidating, + (PreSendValidating, PreSendValidated) => Queued, (Queued, SchedulerPicked) => Streaming { chunks_sent: 0 }, // egress (Streaming { chunks_sent }, Chunk { finish: false }) => Streaming { @@ -133,3 +144,59 @@ impl RequestState { Ok(()) } } + +#[cfg(test)] +mod tests { + use super::*; + + fn after(mut state: RequestState, event: Event) -> RequestState { + state.apply(event).expect("edge must exist"); + state + } + + /// Every ingress branch — control, client-supplied ids, and text through the + /// tokenizer pool — must land in `PreSendValidating`, because that is where + /// the checks needing the final `input_ids` run. A branch that reached + /// `Queued` directly would skip them silently. + #[test] + fn every_branch_reaches_the_ring_through_pre_send_validating() { + for from in [ + after( + RequestState::Validating, + Event::Validated(ValidationOutcome::AlreadyTokenized), + ), + after( + RequestState::Normalizing, + Event::Validated(ValidationOutcome::AlreadyTokenized), + ), + after(RequestState::Tokenizing, Event::TokenizeDone), + ] { + assert!( + matches!(from, RequestState::PreSendValidating), + "branch bypassed the pre-send checks: {from:?}" + ); + assert!(matches!( + after(from, Event::PreSendValidated), + RequestState::Queued + )); + } + } + + /// The converse: `Queued` has no other in-edge, so the checks can't be skipped + /// by emitting the wrong event, and can't run twice. + #[test] + fn queued_has_no_other_in_edge() { + for mut state in [ + RequestState::Validating, + RequestState::Normalizing, + RequestState::Tokenizing, + RequestState::Queued, + ] { + assert_eq!( + state.apply(Event::PreSendValidated), + Err(TransitionError::Illegal), + "only PreSendValidating may enter Queued" + ); + } + } +} diff --git a/rust/sglang-server/src/lib.rs b/rust/sglang-server/src/lib.rs index f733bb9dd..1d4e6e8fb 100644 --- a/rust/sglang-server/src/lib.rs +++ b/rust/sglang-server/src/lib.rs @@ -16,6 +16,9 @@ mod error; mod fsm; mod ids; mod message; +mod ring; +mod runtime; +mod tokenizer_manager; mod utils; use pyo3::prelude::*; diff --git a/rust/sglang-server/src/message.rs b/rust/sglang-server/src/message.rs index 2eb79682c..81d91e37f 100644 --- a/rust/sglang-server/src/message.rs +++ b/rust/sglang-server/src/message.rs @@ -15,6 +15,65 @@ mod request; mod sampling; mod types; -pub(crate) use request::GenerateRequest; +pub use egress::{ChunkEvent, EgressSink}; +pub use request::{GenerateRequest, RequestKind}; pub(crate) use sampling::{SamplingParams, SamplingParamsInput}; pub(crate) use types::{OneOrMany, OneOrManyItem, TokenIds}; + +use bytes::Bytes; + +use crate::fsm::RequestState; +use crate::ids::Rid; + +/// The owned request as it travels ingress stages (single owner, so `state` is +/// mutated lock-free). Common fields here; variant data in [`RequestKind`]. +#[derive(Debug)] +pub struct Request { + /// Client-visible request id (uuid hex) — what the scheduler wire and + /// `meta_info.id` carry. + pub rid: Rid, + pub state: RequestState, + /// Back-channel to the client connection for egress frames. + pub sink: EgressSink, + /// Discriminant + variant body (generate vs control). + pub kind: RequestKind, +} + +/// One ingress-ring entry, split columnar: the scalar `header` (msgpack, `input_ids` +/// omitted) + the raw int64 `ids` cell, so the big tensor never goes through msgpack. +#[derive(Debug)] +pub struct IngressMsg { + pub header: Bytes, + pub ids: Bytes, +} + +/// Messages to a Detokenizer shard. `Register` carries the per-request sink for +/// the shard's local `rid -> sink` map. The rid STRING is the identity: `Rid::hash` +/// picks the shard (collisions there merely co-locate, which is harmless), but two +/// distinct rids that hash alike must not be the same map entry — that evicted one +/// client's sink and delivered their tokens to the other's connection. Equal rids +/// cannot reach here from different requests: `Rid::from_client` uniquifies every +/// client-supplied one. +pub enum DetokMsg { + Register { + /// Client-visible rid string — kept in `DetokState` so the shard can + /// emit `TmEvent::Abort(rid)` (the wire needs the string, not the hash). + rid: Rid, + sink: EgressSink, + /// Decode logprob token ids to text here (CPU-bound) not on the api threads. + decode_logprob_text: bool, + /// `SamplingParams.no_stop_trim`: keep the matched stop; default trims it. + no_stop_trim: bool, + }, + /// One decode step's chunks for *this shard*. Batched because `tm-egress` blocks + /// per send, so one message per request cost ~1.3 µs × batch (5.1x at 4096). + Chunks(Vec), + /// Control result: one already-serialized payload delivered to the sink verbatim. + Result { rid: Rid, payload: bytes::Bytes }, + /// Terminal per-request failure → an `Error` to the sink (a 400, not a crash). + Fail { rid: Rid, message: String }, + /// Drop the `rid -> sink` entry for a request rejected before the scheduler + /// (the rejecting stage already answered the client); else `Register` leaks one + /// entry. + Deregister { rid: Rid }, +} diff --git a/rust/sglang-server/src/message/request.rs b/rust/sglang-server/src/message/request.rs index 9b2c5e3d9..00277a543 100644 --- a/rust/sglang-server/src/message/request.rs +++ b/rust/sglang-server/src/message/request.rs @@ -486,6 +486,10 @@ fn fan_out( mod tests { use super::*; + /// Vocab size for tests that aren't about the vocab bound (see + /// `sampling::tests::TEST_VOCAB`). + const TEST_VOCAB: u64 = 1000; + fn requests(body: &str) -> Result<(Vec, bool), Error> { serde_json::from_str::(body) .unwrap() @@ -558,7 +562,7 @@ mod tests { // Parallel sampling is rejected where Python reads it — in the params, // at normalization (the ingress step), not here. let (mut ps, _) = requests(r#"{"text": "a", "sampling_params": {"n": 2}}"#).unwrap(); - assert!(ps[0].sampling_params.normalize(false, None).is_err()); + assert!(ps[0].sampling_params.normalize(false, TEST_VOCAB).is_err()); } /// Unported `GenerateReqInput` fields are IGNORED, not rejected. diff --git a/rust/sglang-server/src/ring.rs b/rust/sglang-server/src/ring.rs new file mode 100644 index 000000000..f8d0fa1f2 --- /dev/null +++ b/rust/sglang-server/src/ring.rs @@ -0,0 +1,251 @@ +//! The two Rust↔Python boundary queues. +//! +//! In embedded mode the Rust frontend threads and the Python scheduler loop +//! share one process, so these are in-process `flume` channels — literal +//! `mpsc`/`mpmc`, no shared memory, no serialization beyond the msgpack bytes +//! the payload already is. +//! +//! GIL note: the Python side only ever calls the *non-blocking* `drain` / +//! `try_push` methods while holding the GIL, and the Rust worker threads only +//! ever push/drain raw `Bytes` — neither side touches a `PyObject` off-thread, +//! so the producer threads never need the GIL. + +use std::sync::Mutex; +use std::time::Duration; + +use bytes::Bytes; + +use crate::message::IngressMsg; + +/// Ingress: TokenizerManager → scheduler `recv_requests`. +/// Producers are Rust TM workers; the single consumer is the Python thread. +/// Carries [`IngressMsg`] (columnar: scalar header + raw int64 ids cell), not a +/// single msgpack blob, so the large `input_ids` tensor bypasses msgpack. +#[derive(Clone)] +pub struct IngressProducer { + tx: flume::Sender, +} + +pub struct IngressConsumer { + rx: flume::Receiver, + /// One-slot buffer holding a message consumed by a blocking [`wait`] so the + /// scheduler can park on idle without losing it — the next [`drain`] returns + /// it first. Only ever touched by the single consumer (the Python thread), + /// so contention is nil; the `Mutex` is just for interior mutability across + /// the `&self` methods. + /// + /// [`wait`]: IngressConsumer::wait + /// [`drain`]: IngressConsumer::drain + stash: Mutex>, +} + +/// A drained ingress batch in **columnar** (struct-of-arrays) form. The `ids` +/// cells are kept *un-concatenated* so the pyo3 boundary can copy them straight +/// into one `PyBytes` (no intermediate buffer); `ids_total` is their summed +/// length, precomputed for that single allocation. +#[derive(Default)] +pub struct IngressColumns { + /// Per-request scalar msgpack header (`input_ids` omitted). + pub headers: Vec, + /// Per-request raw little-endian int64 ids cell (empty for control reqs). + pub ids: Vec, + /// Per-request token count (`ids` cell length / 8). + pub lengths: Vec, + /// Sum of all `ids` cell byte lengths. + pub ids_total: usize, +} + +impl IngressProducer { + /// Non-blocking push. Returns `false` on a full ring (backpressure) so the + /// caller can fail the request rather than block a worker thread. + #[inline] + pub fn try_push(&self, msg: IngressMsg) -> bool { + self.tx.try_send(msg).is_ok() + } +} + +impl IngressConsumer { + /// Drain up to `max` messages into a columnar [`IngressColumns`], returning + /// immediately when the ring runs dry — mirrors the scheduler's existing + /// `zmq.NOBLOCK` loop in `request_receiver._pull_raw_reqs`. Splitting headers + /// from ids here (off the GIL) leaves `recv_requests` a thin marshaling shim. + /// + /// Non-blocking by construction: `try_recv` returns `Err(TryRecvError::Empty)` + /// instantly when the ring is empty, and `Err(_) => break` exits the loop + /// right away. + pub fn drain(&self, max: usize) -> IngressColumns { + let mut batch = IngressColumns::default(); + // A message parked by a prior blocking `wait` is delivered first. + if let Some(m) = self.stash.lock().unwrap().take() { + push_msg(&mut batch, m); + } + while batch.headers.len() < max { + match self.rx.try_recv() { + Ok(m) => push_msg(&mut batch, m), + Err(_) => break, // Empty or Disconnected -> stop now + } + } + batch + } + + /// Park up to `timeout` for at least one incoming message, so the idle + /// scheduler loop sleeps instead of spinning at 100% CPU. The message is + /// **stashed, not returned** — the next [`drain`](Self::drain) yields it — + /// so this composes with the existing non-blocking drain flow. Returns + /// whether a message is now available. `flume` wakes the parked thread the + /// instant a producer pushes, so this adds no latency to real requests. + pub fn wait(&self, timeout: Duration) -> bool { + if self.stash.lock().unwrap().is_some() { + return true; + } + match self.rx.recv_timeout(timeout) { + Ok(m) => { + *self.stash.lock().unwrap() = Some(m); + true + } + Err(_) => false, // Timeout or Disconnected + } + } +} + +/// Append one drained message's columnar cells to the batch. +#[inline] +fn push_msg(batch: &mut IngressColumns, m: IngressMsg) { + batch.ids_total += m.ids.len(); + batch.lengths.push((m.ids.len() / 8) as u32); // int64 cell → tokens + batch.headers.push(m.header); + batch.ids.push(m.ids); +} + +/// Egress: scheduler output (`push_chunk`) → Rust egress dispatcher. +/// The single producer is the Python thread; the consumer is the dispatcher. +#[derive(Clone)] +pub struct EgressProducer { + tx: flume::Sender, +} + +pub struct EgressConsumer { + rx: flume::Receiver, +} + +impl EgressProducer { + /// Blocking push: parks until the ring has space, so a full ring applies + /// backpressure to the scheduler instead of dropping output the scheduler has + /// already committed (advanced `send_token_offset` for). The GIL is released + /// around the call, so parking here doesn't stall other Python threads. + /// `false` only when the consumer is gone (runtime shutdown), where the frame + /// is unavoidably lost. + pub fn push(&self, msg: Bytes) -> bool { + self.tx.send(msg).is_ok() + } + + /// Non-blocking push, so the pyo3 boundary can try to hand the frame over + /// while still holding the GIL and detach only when it would actually park. + /// Releasing the GIL is not free: reacquiring it waits out the interpreter's + /// switch interval (5 ms by default), which dwarfs the sub-microsecond push + /// it was protecting. + /// + /// Hands the frame BACK on a full ring (`Err(Some(msg))`) so the caller can + /// retry it under [`push`](Self::push) without rebuilding it. `Err(None)` is + /// the consumer being gone (shutdown), where the frame is unavoidably lost. + #[inline] + pub fn try_push(&self, msg: Bytes) -> Result<(), Option> { + match self.tx.try_send(msg) { + Ok(()) => Ok(()), + Err(flume::TrySendError::Full(msg)) => Err(Some(msg)), + Err(flume::TrySendError::Disconnected(_)) => Err(None), + } + } +} + +impl EgressConsumer { + /// The underlying receiver, so the dispatcher can drain it via + /// [`tokenizer_manager::recv`](crate::tokenizer_manager::recv) (data + shutdown select). + pub fn receiver(&self) -> &flume::Receiver { + &self.rx + } +} + +/// Build both halves of a bounded ring. +pub fn ingress_ring(cap: usize) -> (IngressProducer, IngressConsumer) { + let (tx, rx) = flume::bounded(cap); + ( + IngressProducer { tx }, + IngressConsumer { + rx, + stash: Mutex::new(None), + }, + ) +} + +pub fn egress_ring(cap: usize) -> (EgressProducer, EgressConsumer) { + let (tx, rx) = flume::bounded(cap); + (EgressProducer { tx }, EgressConsumer { rx }) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn msg(h: &'static [u8]) -> IngressMsg { + IngressMsg { + header: Bytes::from_static(h), + ids: Bytes::new(), + } + } + + /// `wait` parks when empty (times out), stashes a pushed message + /// non-destructively, and the next `drain` returns it. + #[test] + fn wait_stashes_then_drain_returns_it() { + let (tx, rx) = ingress_ring(8); + // Empty ring → times out, nothing stashed. + assert!(!rx.wait(Duration::from_millis(1))); + // Push one, then wait stashes it (returns true). + assert!(tx.try_push(msg(b"a"))); + assert!(rx.wait(Duration::from_millis(200))); + // Idempotent: already stashed, returns true without touching the ring. + assert!(rx.wait(Duration::from_millis(1))); + // Drain yields the stashed message, then the ring is empty. + assert_eq!(rx.drain(16).headers.len(), 1); + assert!(rx.drain(16).headers.is_empty()); + } + + /// A blocked `wait` is woken the instant a producer pushes (no polling). + #[test] + fn wait_wakes_on_push() { + let (tx, rx) = ingress_ring(8); + std::thread::spawn(move || { + std::thread::sleep(Duration::from_millis(20)); + let _ = tx.try_push(msg(b"a")); + }); + // Generous timeout, but it should return well before it as soon as the + // push lands. + assert!(rx.wait(Duration::from_secs(5))); + assert_eq!(rx.drain(16).headers.len(), 1); + } + + /// A full egress ring parks the producer until the consumer drains — the + /// committed frame is delivered in order, never dropped. + #[test] + fn egress_push_blocks_until_drained() { + let (tx, rx) = egress_ring(1); + assert!(tx.push(Bytes::from_static(b"a"))); // fits; ring now full + let t = std::thread::spawn(move || tx.push(Bytes::from_static(b"b"))); + // The parked push can't have completed while the ring is full. + std::thread::sleep(Duration::from_millis(20)); + // Drain one → frees a slot → the parked push lands. + assert_eq!(rx.receiver().recv().unwrap(), Bytes::from_static(b"a")); + assert!(t.join().unwrap(), "push should succeed once space frees"); + assert_eq!(rx.receiver().recv().unwrap(), Bytes::from_static(b"b")); + } + + /// A closed ring (consumer gone → shutdown) returns `false` instead of + /// parking forever, so a scheduler blocked in `push` unblocks on teardown. + #[test] + fn egress_push_returns_false_when_closed() { + let (tx, rx) = egress_ring(1); + drop(rx); + assert!(!tx.push(Bytes::from_static(b"x"))); + } +} diff --git a/rust/sglang-server/src/runtime.rs b/rust/sglang-server/src/runtime.rs new file mode 100644 index 000000000..cd9751fed --- /dev/null +++ b/rust/sglang-server/src/runtime.rs @@ -0,0 +1,18 @@ +//! Runtime bootstrap: wires channels, pins CPU-bound pools, starts the tokio +//! API server, and returns a handle the Python boundary uses for +//! `recv_requests` (ingress drain) and `push_batch` (egress push). +//! +//! Thread layout: +//! * API server — tokio multi-thread runtime (I/O bound), pinned core set A +//! * Tokenizer — N pinned OS threads (CPU bound), core set B +//! * Detokenizer — M pinned OS threads / shards (CPU bound), core set C +//! * TM ingress — 1 thread driving the ingress FSM +//! * TM egress — 1 thread draining the egress ring → detok shards +//! +//! Keeping CPU-bound tokenize/detokenize off the async executor avoids stalling +//! axum's worker threads. + +mod config; +mod runnable; + +// Re-export so stages keep importing `crate::runtime::Runnable`. diff --git a/rust/sglang-server/src/runtime/config.rs b/rust/sglang-server/src/runtime/config.rs new file mode 100644 index 000000000..ca5a557a9 --- /dev/null +++ b/rust/sglang-server/src/runtime/config.rs @@ -0,0 +1,203 @@ +//! Runtime configuration: the rust-server boot knobs +//! ([`RustServerServerArgs`]), the typed view of the scheduler's `server_args` +//! dump ([`ServerArgs`] / [`ModelConfig`]), and the [`RuntimeConfig`] pairing +//! them for `runtime::start`. + +use std::net::SocketAddr; +use std::sync::Arc; + +/// Boot knobs specific to the embedded rust server — none of these exist in +/// the Python `server_args` dump (see [`ServerArgs`]); they arrive as explicit +/// `Server::start` parameters. +#[derive(Clone, Debug)] +pub struct RustServerServerArgs { + pub http_addr: SocketAddr, + pub api_worker_num: usize, + pub ingress_ring_cap: usize, + pub egress_ring_cap: usize, + pub channel_cap: usize, + /// CPU core ids the pools pin to (e.g. this rank's NUMA-local cores minus + /// the scheduler's reserved launch cores). `None` → run unpinned. + pub cores: Option>, +} + +impl Default for RustServerServerArgs { + fn default() -> Self { + Self { + http_addr: "127.0.0.1:30000".parse().unwrap(), + api_worker_num: 2, + ingress_ring_cap: 8192, + egress_ring_cap: 8192, + channel_cap: 8192, + cores: None, + } + } +} + +#[derive(Clone, Debug)] +pub struct RuntimeConfig { + /// Rust-server-only boot knobs (listen address, pool/ring sizes, pinning). + pub rust_server_args: RustServerServerArgs, + /// The scheduler's `server_args` dump (worker counts, tokenizer source, + /// config-endpoint metadata). `Arc` so cloning the config (and, downstream, + /// each `AppState`) is cheap; immutable after construction. + pub server_args: Arc, +} + +impl Default for RuntimeConfig { + fn default() -> Self { + Self { + rust_server_args: RustServerServerArgs::default(), + server_args: Arc::new( + ServerArgs::from_json("{}").expect("empty server_args blob parses"), + ), + } + } +} + +/// The scheduler's startup blob (`RustServer._build_server_args`) parsed once into +/// typed fields: values are post-`__post_init__`, unknown keys (e.g. `api_key`) are dropped. +#[derive(Debug, serde::Deserialize)] +pub struct ServerArgs { + /// HF repo id / local dir of the model, reported by `/get_model_info`. + #[serde(default)] + pub model_path: String, + /// Model name reported by `/v1/models` and `/server_info`. + #[serde(default)] + pub served_model_name: String, + /// Tokenizer source (model dir / `tokenizer.json` / HF repo id). Empty only + /// in minimal standalone blobs — then boot requires `skip_tokenizer_init`. + #[serde(default)] + pub tokenizer_path: String, + /// HF revision, used only when `tokenizer_path` is a repo id. `None` → main. + #[serde(default)] + pub revision: Option, + /// HTTP bind address (see [`Self::bind`]). + #[serde(default = "default_host")] + pub host: String, + #[serde(default = "default_port")] + pub port: u16, + /// Log levels driving the access log — uvicorn runs at + /// `log_level_http or log_level` (see [`Self::http_access_log_enabled`]). + #[serde(default = "default_log_level")] + pub log_level: String, + #[serde(default)] + pub log_level_http: Option, + /// Pinned tokenizer threads / detok shards (Python asserts both ≥ 1). + #[serde(default = "default_worker_num")] + pub tokenizer_worker_num: usize, + #[serde(default = "default_worker_num")] + pub detokenizer_worker_num: usize, + /// Token-ids-in / token-ids-out mode: no tokenizer load, raw `output_ids` + /// frames (drives the `Skip` detok backend and the ingress branch). + #[serde(default)] + pub skip_tokenizer_init: bool, + /// Streamed `/generate` frames carry per-step deltas instead of cumulative + /// text. Matches the Python `TokenizerManager`. + #[serde(default)] + pub incremental_streaming_output: bool, + /// The resolved Python `ModelConfig`, attached to the blob at dump time. + #[serde(default)] + pub model_config: ModelConfig, + /// Default sampling params advertised by `/get_model_info`, verbatim from + /// `server_args.preferred_sampling_params` (a JSON object or null). + #[serde(default)] + pub preferred_sampling_params: Option, + /// Over-long inputs are truncated to fit the context instead of 400ing, and + /// `max_new_tokens` is clamped rather than rejected (Python + /// `TokenizerManager._validate_one_request`). + #[serde(default)] + pub allow_auto_truncate: bool, + /// `return_hidden_states` is refused unless the server was launched with it: + /// the scheduler simply won't produce them, so the request would 200 with the + /// field silently missing. + #[serde(default)] + pub enable_return_hidden_states: bool, + /// Output slots reserved per request on top of its input (eagle stores draft + /// tokens there). Not a `server_args` field — `TokenizerManager` derives it and + /// `RustServer._build_server_args` stamps it in, so both sides count alike. + #[serde(default)] + pub num_reserved_tokens: u64, + /// Launch-time stamps (not `server_args` fields): sglang package version + /// and the scheduler-derived KV token capacity, reported by `/server_info`. + #[serde(default)] + pub version: Option, + #[serde(default)] + pub max_total_num_tokens: Option, +} + +/// The slice of the resolved Python `ModelConfig` the rust server reads. +#[derive(Debug, Default, serde::Deserialize)] +pub struct ModelConfig { + /// Resolved context length (`max_model_len` in `/v1/models`); mandatory at + /// boot ([`ServerArgs::validate_mandatory`]). + #[serde(default)] + pub context_len: Option, + /// Bounds client-supplied token ids — ingress 400s out-of-vocab ids before + /// they crash the scheduler's embedding lookup; mandatory at + /// boot ([`ServerArgs::validate_mandatory`]). + #[serde(default)] + pub vocab_size: Option, +} + +fn default_host() -> String { + "127.0.0.1".into() +} +fn default_port() -> u16 { + 30000 +} +fn default_log_level() -> String { + "info".into() +} +fn default_worker_num() -> usize { + 1 +} + +impl ServerArgs { + /// Parse the blob; errors on malformed JSON or a wrongly-typed field. + pub fn from_json(s: &str) -> Result { + serde_json::from_str(s).map_err(|e| e.to_string()) + } + + /// Fail fast at startup if a field an endpoint depends on is missing. + pub fn validate_mandatory(&self) -> Result<(), String> { + if self.served_model_name.is_empty() { + return Err("no 'served_model_name' in server_args".into()); + } + if self.model_config.context_len.is_none() { + return Err("no resolvable context length (model_config.context_len)".into()); + } + if self.model_config.vocab_size.is_none() { + return Err("no resolvable vocab size (model_config.vocab_size)".into()); + } + Ok(()) + } + + /// Bind address `host:port`. `host` is expected to be an IP — the result is + /// parsed as a `SocketAddr`. + pub fn bind(&self) -> String { + format!("{}:{}", self.host, self.port) + } + + /// Whether the HTTP access log is emitted, mirroring the Python server: + /// uvicorn runs at `log_level_http or log_level` and prints access lines + /// only at info/debug. `--log-level-http warning` turns them off. + pub fn http_access_log_enabled(&self) -> bool { + let level = self + .log_level_http + .as_deref() + .filter(|s| !s.is_empty()) + .unwrap_or(&self.log_level); + matches!( + level.to_ascii_lowercase().as_str(), + "trace" | "debug" | "info" + ) + } + + /// Pinned API threads for the embedded HTTP api-server. Python `server_args` + /// has no such field — this is derived: enough to cover the widest pool. + pub fn api_worker_num(&self) -> usize { + 4.max(self.tokenizer_worker_num) + .max(self.detokenizer_worker_num) + } +} diff --git a/rust/sglang-server/src/runtime/runnable.rs b/rust/sglang-server/src/runtime/runnable.rs new file mode 100644 index 000000000..66b7535d7 --- /dev/null +++ b/rust/sglang-server/src/runtime/runnable.rs @@ -0,0 +1,10 @@ +//! The [`Runnable`] stage trait — the one contract every pipeline stage +//! (CPU-bound worker or TM router) implements to be spawned by the runtime. + +/// A pipeline stage that owns its channel handles + config and runs a blocking +/// loop until its inbox closes. Lets the runtime spawn stages uniformly via +/// `threads::spawn_stage` / `threads::spawn_pool` instead of free `run_*` functions with +/// positional handles. Implemented by every CPU-bound worker and TM router. +pub trait Runnable: Send + 'static { + fn run(self); +} diff --git a/rust/sglang-server/src/tokenizer_manager.rs b/rust/sglang-server/src/tokenizer_manager.rs new file mode 100644 index 000000000..a830a7b4a --- /dev/null +++ b/rust/sglang-server/src/tokenizer_manager.rs @@ -0,0 +1,87 @@ +//! TokenizerManager — owns the request lifecycle across two isolated threads: +//! +//! * [`ingress`] — drives the ingress FSM (Received → Validating → +//! Normalizing → {Tokenizing | PreSendValidating}) and pushes tokenized requests to the +//! scheduler ring. +//! * [`egress`] — drains the scheduler-output ring and routes each chunk to +//! the owning detokenizer shard. +//! +//! The two run on separate pinned threads with no shared state, connected to +//! the rest of the pipeline only through `flume` channels: [`TmEvent`] into the +//! ingress loop, [`Senders`] fanning out to the pools. + +use crate::ids::Rid; +use crate::message::{DetokMsg, Request}; + +/// Blocking receive that also wakes on shutdown: returns `None` when `rx` closes +/// *or* the `shutdown` sender is dropped. +pub fn recv(rx: &flume::Receiver, shutdown: &flume::Receiver<()>) -> Option { + flume::Selector::new() + .recv(rx, |r| r.ok()) + .recv(shutdown, |_| None) + .wait() +} + +/// Events into the TokenizerManager ingress loop. API server + tokenizer pool +/// share this one inbox, keeping the loop a single consumer (no `select`). +pub enum TmEvent { + /// A freshly received request from the API server. + Ingress(Request), + /// A request back from the tokenizer pool: `PreSendValidating` (ids filled) on success, + /// or `Failed` on a tokenize error. `drive` handles both. + Tokenized(Request), +} + +/// Producer-side handles, cloned into every stage that needs to emit. +/// Who asked for an abort. Both variants do the same work in +/// [`Ingress::on_abort`](crate::tokenizer_manager::ingress::Ingress) — deregister +/// the detok entry, tell the scheduler to stop — and the source is kept for +/// diagnostics. +/// +/// There is no in-flight rid registry to keep consistent, and so no release +/// ordering to get wrong: [`Rid::from_client`] makes every client-supplied rid +/// internally unique, so a resubmit of the "same" rid is a different `Rid` and +/// cannot be tangled up with an abort still in flight for the original. +#[derive(Clone, Debug)] +pub enum AbortSource { + /// From an `AbortGuard` drop. Owns the release. + Guard(Rid), + /// From a detokenizer terminal path. Aborts the scheduler work. + Detok(Rid), +} + +impl AbortSource { + pub fn rid(&self) -> &Rid { + match self { + Self::Guard(rid) | Self::Detok(rid) => rid, + } + } +} + +#[derive(Clone)] +pub struct Senders { + /// → TokenizerManager ingress loop. + pub tm: flume::Sender, + /// → the same loop, but UNBOUNDED and abort-only. + /// + /// Aborts cannot share the bounded inbox. `try_send` there drops them exactly + /// when they matter most — under overload — leaving the scheduler generating + /// for a dead connection; and the caller then faces a false choice between + /// releasing the rid (a live entry can be overwritten by a resubmit) and + /// holding it (a permanent leak). An unbounded lane removes the dilemma: an + /// abort is a small `String` and is always accepted, so releases can be + /// unconditional again. It cannot grow without bound in practice — one entry + /// per in-flight request, each already bounded by the inbox that admitted it. + pub abort: flume::Sender, + /// → Tokenizer pool (CPU-bound, pinned threads). + pub tok: flume::Sender, + /// → Detokenizer shards, indexed by `Rid::shard(detok.len())`. + pub detok: Vec>, +} + +impl Senders { + #[inline] + pub fn detok_for(&self, rid: &Rid) -> &flume::Sender { + &self.detok[rid.shard(self.detok.len())] + } +}