From 3312645a307453893a00778592f105581e3d1c3d Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Thu, 30 Jul 2026 12:46:19 -0700 Subject: [PATCH] wire the rust server modules into lib, runtime, and tokenizer manager (#32877) --- rust/sglang-server/src/environ.rs | 1 - rust/sglang-server/src/error.rs | 1 - rust/sglang-server/src/lib.rs | 225 ++++++++- rust/sglang-server/src/message.rs | 10 +- rust/sglang-server/src/runtime.rs | 492 ++++++++++++++++++++ rust/sglang-server/src/tokenizer_manager.rs | 6 + 6 files changed, 729 insertions(+), 6 deletions(-) diff --git a/rust/sglang-server/src/environ.rs b/rust/sglang-server/src/environ.rs index e1a691e5f..afa70ff25 100644 --- a/rust/sglang-server/src/environ.rs +++ b/rust/sglang-server/src/environ.rs @@ -2,7 +2,6 @@ //! unset → default, invalid → warn + default (never an error). One shared //! parser per type — call sites pass their variable name + default instead of //! each hand-rolling a reader. -#![allow(dead_code)] // TODO: remove when the consumer PR lands /// Python `EnvBool.parse`: true = `true/1/yes/y`, false = `false/0/no/n` /// (case-insensitive); anything else is invalid. diff --git a/rust/sglang-server/src/error.rs b/rust/sglang-server/src/error.rs index a1eff4082..6b5feb8c1 100644 --- a/rust/sglang-server/src/error.rs +++ b/rust/sglang-server/src/error.rs @@ -1,6 +1,5 @@ //! Error type shared by all stages. Kept `Clone` so a single failure can be //! reported to the client stream and logged without moving ownership around. -#![allow(dead_code)] // TODO: remove when the consumer PR lands use thiserror::Error; diff --git a/rust/sglang-server/src/lib.rs b/rust/sglang-server/src/lib.rs index 1d4e6e8fb..e82bf57a0 100644 --- a/rust/sglang-server/src/lib.rs +++ b/rust/sglang-server/src/lib.rs @@ -9,8 +9,9 @@ //! * `push_result` — Python scheduler thread pushes one control result. //! //! All are non-blocking, so the GIL is never held across a wait. -#![allow(dead_code)] // TODO: remove when the consumer PR lands +mod api_server; +mod detokenizer; mod environ; mod error; mod fsm; @@ -18,12 +19,232 @@ mod ids; mod message; mod ring; mod runtime; +mod tokenizer; mod tokenizer_manager; mod utils; +use std::net::SocketAddr; + use pyo3::prelude::*; +use pyo3::pybacked::PyBackedBytes; +use pyo3::types::PyBytes; + +use crate::runtime::{Runtime, RuntimeConfig}; + +/// Columnar ingress batch handed to Python by [`Server::recv_requests`]. +/// `frozen`: immutable snapshot, so field access never contends on a borrow. +#[pyclass(frozen, get_all)] +struct IngressBatch { + /// One msgpack scalar header per request (`input_ids` omitted). + headers: Vec>, + /// The raw-data plane today just all requests' raw little-endian int64 + /// ids, concatenated; sliced per request via `lengths`. + data: Py, + /// Per-request token count (0 for control requests). + lengths: Vec, +} + +/// Handle owned by the Python scheduler process. Construct once via +/// [`Server::start`], then poll it from the scheduler event loop. +#[pyclass] +struct Server { + rt: Runtime, +} + +#[pymethods] +impl Server { + /// Boot the frontend (spawns all threads) and return immediately. + #[new] + #[pyo3(signature = ( + http_addr = None, + ingress_ring_cap = 8192, + egress_ring_cap = 8192, + channel_cap = 8192, + cores = None, + + server_args_json = "{}", + ))] + // pyo3 `#[new]` constructor: the wide arg list is the Python-facing boot + // surface (all optional overrides), not a call-site ergonomics problem. + #[allow(clippy::too_many_arguments)] + fn start( + http_addr: Option, + ingress_ring_cap: usize, + egress_ring_cap: usize, + channel_cap: usize, + cores: Option>, + server_args_json: &str, + ) -> PyResult { + // Static server metadata (server_args + model_config) dumped by the + // scheduler; parse and validate mandatory fields now so a bad/missing + // field is a boot error, not a request-time 500. + let server_args: runtime::ServerArgs = runtime::ServerArgs::from_json(server_args_json) + .map_err(|e| { + PyErr::new::(format!( + "bad server_args_json: {e}" + )) + })?; + server_args.validate_mandatory().map_err(|e| { + PyErr::new::(format!("server_args: {e}")) + })?; + // The HTTP listen address, tokenizer source/threads/shards all live in the + // `server_args` blob; resolve them from there so the scheduler doesn't + // re-pass them. The explicit params stay as optional overrides for + // standalone callers (tests) that construct a `Server` without a full + // `server_args`. + let http_addr: SocketAddr = http_addr + .unwrap_or_else(|| server_args.bind()) + .parse() + .map_err(|e| { + PyErr::new::(format!("bad http_addr: {e}")) + })?; + + let cfg = RuntimeConfig { + rust_server_args: runtime::RustServerServerArgs { + http_addr, + api_worker_num: server_args.api_worker_num(), + ingress_ring_cap, + egress_ring_cap, + channel_cap, + cores, + }, + server_args: std::sync::Arc::new(server_args), + }; + let rt = runtime::start(cfg).map_err(|e| { + PyErr::new::(format!("runtime start failed: {e}")) + })?; + Ok(Server { rt }) + } + + /// Non-blocking drain of the ingress ring, returned **columnar** as an + /// [`IngressBatch`] so the large `input_ids` tensor never goes through + /// msgpack (see the field docs for the layout). The `ids` cells are copied + /// **directly into the result `bytes`** (one copy, no intermediate buffer). + /// + /// Runs entirely GIL-held, deliberately. `drain` is a `try_recv` loop plus an + /// uncontended stash lock (the Python thread is the only consumer), so it + /// cannot block — there is nothing for a detach to overlap with. And detaching + /// is far from free: reacquiring the GIL waits out the interpreter's switch + /// interval, so a `py.detach` here cost up to 5 ms whenever another Python + /// thread was runnable, to cover ~0.2 µs of work. Held, the whole call is a + /// fraction of a microsecond on an empty ring. + #[pyo3(signature = (max = 256))] + fn recv_requests(&self, py: Python<'_>, max: usize) -> PyResult { + let cols = self.rt.ingress.drain(max); + let headers = cols + .headers + .iter() + .map(|h| PyBytes::new(py, h).unbind()) + .collect(); + // Single pass: copy each raw ids cell straight into the output `bytes`. + let data = PyBytes::new_with(py, cols.ids_total, |buf| { + let mut pos = 0; + for cell in &cols.ids { + let end = pos + cell.len(); + buf[pos..end].copy_from_slice(cell); + pos = end; + } + Ok(()) + })? + .unbind(); + Ok(IngressBatch { + headers, + data, + lengths: cols.lengths, + }) + } + + /// Park up to `timeout_ms` for an incoming request so the idle scheduler loop + /// sleeps instead of spinning at 100% CPU. Returns `True` when a request is + /// ready (the next `recv_requests` includes it). The GIL is released while + /// parked, and `flume` wakes the moment a request is pushed, so this adds no + /// latency to real requests — only the idle wait is bounded by `timeout_ms`. + #[pyo3(signature = (timeout_ms = 1000))] + fn wait_ingress(&self, py: Python<'_>, timeout_ms: u64) -> bool { + py.detach(|| { + self.rt + .ingress + .wait(std::time::Duration::from_millis(timeout_ms)) + }) + } + + /// Push a whole decode batch as ONE frame: a columnar msgpack `header` plus + /// the raw `data_cols` (per-column `bytes`), concatenated here. Blocks for + /// backpressure; `False` only on shutdown. + /// + /// Framed and pushed with the GIL HELD, detaching only if the ring is full. + /// This runs on the scheduler's CUDA-launch thread every decode step, where the + /// unconditional detach was the single worst boundary cost: framing is + /// ~0.1–0.2 µs, but reacquiring the GIL waits out the interpreter's switch + /// interval (5 ms by default) whenever another Python thread is runnable — + /// 17–50% of a 10–30 ms decode step, landing nondeterministically. Held, the + /// whole boundary is ~1.3 µs per step. + /// + /// The slow path keeps its detach because a full ring genuinely parks: the + /// scheduler must feel backpressure rather than drop output it has already + /// committed to. It essentially never fires — measured headroom is ~100×. + fn push_batch(&self, py: Python<'_>, header: &[u8], data_cols: Vec) -> bool { + let cols: Vec<&[u8]> = data_cols.iter().map(|d| d.as_ref()).collect(); + self.push_frame(py, crate::message::frame_egress_batch_cols(header, &cols)) + } + + /// Push a control-request result. Blocks for backpressure; `False` only on + /// shutdown. + fn push_result(&self, py: Python<'_>, rid: &str, payload: &[u8]) -> bool { + self.push_frame(py, crate::message::frame_egress_result(rid, payload)) + } + + /// Route a terminal failure back to request `rid`. Blocks for backpressure; + /// `False` only on shutdown. + fn push_error(&self, py: Python<'_>, rid: &str, message: &str) -> bool { + self.push_frame(py, crate::message::frame_egress_error(rid, message)) + } + + /// Signal all threads to stop (best effort). + fn shutdown(&self) { + self.rt.request_shutdown(); + } +} + +impl Server { + /// Hand one already-framed egress message to the ring: GIL-held when it fits, + /// detaching only to park on a full ring. Shared by every push path — they + /// differ solely in how the frame is built. `false` only on shutdown. + #[inline] + fn push_frame(&self, py: Python<'_>, frame: bytes::Bytes) -> bool { + match self.rt.egress.try_push(frame) { + Ok(()) => true, + // Consumer gone (shutdown): the frame is unavoidably lost. + Err(None) => false, + // Full: the scheduler must block here so backpressure reaches it, and + // blocking is exactly when releasing the GIL pays for itself. + Err(Some(frame)) => py.detach(|| self.rt.egress.push(frame)), + } + } +} + +/// Keeps the non-blocking log writer's background thread alive for the process +/// lifetime (dropping the guard would stop log delivery). +static LOG_GUARD: std::sync::OnceLock = + std::sync::OnceLock::new(); #[pymodule] -fn _core(_m: &Bound<'_, PyModule>) -> PyResult<()> { +fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> { + // Initialize tracing once; ignore if already set by the host process. + // Non-blocking writer: emitting threads (axum workers, egress, detok) only + // enqueue; a dedicated thread does the stdout formatting-flush + syscall. + // The queue is bounded and lossy — under extreme pressure log lines are + // dropped instead of stalling request threads. + let (writer, guard) = tracing_appender::non_blocking(std::io::stdout()); + let _ = LOG_GUARD.set(guard); + let _ = tracing_subscriber::fmt() + .with_env_filter( + tracing_subscriber::EnvFilter::try_from_default_env() + .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")), + ) + .with_writer(writer) + .try_init(); + m.add_class::()?; + m.add_class::()?; Ok(()) } diff --git a/rust/sglang-server/src/message.rs b/rust/sglang-server/src/message.rs index 81d91e37f..91fdacf34 100644 --- a/rust/sglang-server/src/message.rs +++ b/rust/sglang-server/src/message.rs @@ -15,8 +15,14 @@ mod request; mod sampling; mod types; -pub use egress::{ChunkEvent, EgressSink}; -pub use request::{GenerateRequest, RequestKind}; +pub use egress::{ + ChunkEvent, ChunkExtras, EGRESS_TAG_BATCH, EGRESS_TAG_ERROR, EGRESS_TAG_RESULT, EgressItem, + EgressSink, SinkError, for_each_chunk, frame_egress_batch_cols, frame_egress_error, + frame_egress_result, +}; +pub use finish_reason::Matched; +pub(crate) use io_struct::{AbortReq, ControlRequest, GetInternalStateReq}; +pub use request::{GenerateBody, GenerateRequest, RequestKind}; pub(crate) use sampling::{SamplingParams, SamplingParamsInput}; pub(crate) use types::{OneOrMany, OneOrManyItem, TokenIds}; diff --git a/rust/sglang-server/src/runtime.rs b/rust/sglang-server/src/runtime.rs index cd9751fed..7119b0a86 100644 --- a/rust/sglang-server/src/runtime.rs +++ b/rust/sglang-server/src/runtime.rs @@ -12,7 +12,499 @@ //! Keeping CPU-bound tokenize/detokenize off the async executor avoids stalling //! axum's worker threads. +use std::sync::{Arc, Mutex}; +use std::thread::JoinHandle; + mod config; mod runnable; +mod threads; + +pub use config::{RuntimeConfig, RustServerServerArgs, ServerArgs}; + +use crate::message::DetokMsg; +use crate::ring::{ + EgressConsumer, EgressProducer, IngressConsumer, IngressProducer, egress_ring, ingress_ring, +}; +use crate::runtime::threads::{plan_cores, spawn_pool}; +use crate::tokenizer_manager::{Senders, TmEvent}; +use crate::{api_server, detokenizer, tokenizer, tokenizer_manager}; // Re-export so stages keep importing `crate::runtime::Runnable`. +pub use runnable::Runnable; + +/// Live runtime. Held by the pyo3 bridge; the Python boundary reads `ingress` +/// and `egress`. `request_shutdown` (also run on `Drop`) stops every stage. +pub struct Runtime { + pub ingress: IngressConsumer, + pub egress: EgressProducer, + /// Worker join handles, joined by `request_shutdown` / `Drop`. + threads: Mutex>>, + /// The single shutdown sender. + shutdown_tx: Mutex>>, +} + +/// Deadline for joining worker threads on shutdown. Past it we abandon the join +/// so process teardown can't deadlock on a worker that somehow failed to exit. +const SHUTDOWN_JOIN_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5); + +impl Runtime { + /// Stop the runtime and join every worker thread (with a bounded wait). + /// + /// Dropping `shutdown_tx` wakes the tm-ingress/tm-egress selectors (which + /// otherwise never see their inbox close — one self-holds a `tm` sender, the + /// other's inbox is the Python-fed egress ring). Those exit and drop their + /// `Senders` clones; the api thread's `serve` returns non-gracefully, so its + /// `block_on` unwinds and the api tokio runtime is dropped — cancelling + /// in-flight handlers, whose `AbortGuard`s release the remaining clones. With + /// every clone gone the tok/detok channels close and those workers exit. + /// + /// In-flight requests are **aborted**, not drained — this is the hard-stop + /// path (also run on `Drop`). Clients of aborted requests retry. + pub fn request_shutdown(&self) { + drop(self.shutdown_tx.lock().unwrap().take()); + let handles: Vec> = self.threads.lock().unwrap().drain(..).collect(); + if handles.is_empty() { + return; // Idempotent: a `Drop` after an explicit shutdown has nothing to join. + } + // Join off-thread and wait with a deadline: a stuck worker can't wedge exit. + let (done_tx, done_rx) = flume::bounded::<()>(1); + std::thread::spawn(move || { + for h in handles { + let _ = h.join(); + } + let _ = done_tx.send(()); + }); + if done_rx.recv_timeout(SHUTDOWN_JOIN_TIMEOUT).is_err() { + tracing::warn!( + "shutdown: workers did not exit within {SHUTDOWN_JOIN_TIMEOUT:?}; abandoning join" + ); + } + } +} + +impl Drop for Runtime { + fn drop(&mut self) { + self.request_shutdown(); + } +} + +/// Boot the whole frontend. Returns once threads are spawned (non-blocking), +/// so the Python caller regains control of the GIL immediately. `Err` on a +/// startup misconfiguration (e.g. no tokenizer for a non-skip server). +pub fn start(cfg: RuntimeConfig) -> Result { + // Bind the API server port before spawning any thread, so an unavailable + // port (EADDRINUSE) is a hard startup error. + let listener = std::net::TcpListener::bind(cfg.rust_server_args.http_addr) + .map_err(|e| format!("bind {} failed: {e}", cfg.rust_server_args.http_addr))?; + listener + .set_nonblocking(true) + .map_err(|e| format!("listener set_nonblocking failed: {e}"))?; + + let (shutdown_tx, shutdown_rx) = flume::unbounded::<()>(); + let mut threads = Vec::new(); + let plan = plan_cores(&cfg); + + // --- rings (Rust ↔ Python) --- + let (ingress_tx, ingress_rx): (IngressProducer, IngressConsumer) = + ingress_ring(cfg.rust_server_args.ingress_ring_cap); + let (egress_tx, egress_rx): (EgressProducer, EgressConsumer) = + egress_ring(cfg.rust_server_args.egress_ring_cap); + + // --- inter-stage channels --- + let (tm_tx, tm_rx) = flume::bounded::(cfg.rust_server_args.channel_cap); + let (tok_tx, tok_rx) = + flume::bounded::(cfg.rust_server_args.channel_cap); + let detokenizer_worker_num = cfg.server_args.detokenizer_worker_num; + let mut detok_tx = Vec::with_capacity(detokenizer_worker_num); + let mut detok_rx = Vec::with_capacity(detokenizer_worker_num); + for _ in 0..detokenizer_worker_num { + let (tx, rx) = flume::bounded::(cfg.rust_server_args.channel_cap); + detok_tx.push(tx); + detok_rx.push(rx); + } + + // Aborts get their own UNBOUNDED lane: on the bounded inbox they are dropped + // exactly under the overload that makes them necessary (see `Senders::abort`). + let (abort_tx, abort_rx) = flume::unbounded::(); + let senders = Senders { + tm: tm_tx.clone(), + abort: abort_tx.clone(), + tok: tok_tx, + detok: detok_tx, + }; + + // `skip_tokenizer_init`: clients send token ids and receive token ids — no + // tokenizer is loaded, and the egress emits raw `output_ids` (no decode). + let skip_tokenizer_init = cfg.server_args.skip_tokenizer_init; + + // The same instance is shared by the tokenizer pool (encode) and the detok + // shards (decode); `None` only under `skip_tokenizer_init`. + let dyn_tokenizer = tokenizer::load_tokenizer( + // Empty only in minimal standalone blobs (the Python dump always + // resolves it); empty → no tokenizer, allowed only under + // `skip_tokenizer_init`. + (!cfg.server_args.tokenizer_path.is_empty()).then_some(&*cfg.server_args.tokenizer_path), + cfg.server_args.revision.as_deref(), + skip_tokenizer_init, + )?; + + // --- Detokenizer shards (pinned, CPU bound) --- + { + // Default: a real tokenizer decodes to text. `None` (→ `Skip`, raw + // `output_ids`) only happens under `skip_tokenizer_init` — + // `load_tokenizer` rejects a non-skip server with no tokenizer. + let backend = match &dyn_tokenizer { + Some(t) => detokenizer::DetokenizerBackend::Dynamo(t.clone()), + None => detokenizer::DetokenizerBackend::Skip, + }; + let detok_cores = plan.as_ref().map(|p| p.detok.clone()); + // Each shard owns its receiver outright (one consumer per shard), so the + // owned `detok_rx` Vec is moved out element-by-element via the iterator. + let count = detok_rx.len(); + let mut rxs = detok_rx.into_iter(); + spawn_pool("detokenizer", detok_cores, count, &mut threads, |i| { + detokenizer::DetokenizerWorker::new( + i, + rxs.next().unwrap(), + backend.clone(), + abort_tx.clone(), + ) + }); + } + + // --- Tokenizer pool (pinned, CPU bound) --- + // Only spawned when a real tokenizer is loaded; under `skip_tokenizer_init` + // there is none and ingress never routes to the pool, so we skip it. + if let Some(t) = &dyn_tokenizer { + // Reuse the single loaded tokenizer (shared with the detok shards). + let tokenizer: Arc = + Arc::new(tokenizer::DynamoTokenizer::new(t.clone())); + let tok_cores = plan.as_ref().map(|p| p.tok.clone()); + // Workers share the MPMC inbox (`tok_rx`) and the read-only backend, so + // each gets a cheap clone of both. + spawn_pool( + "tokenizer", + tok_cores, + cfg.server_args.tokenizer_worker_num, + &mut threads, + |_i| tokenizer::TokenizerWorker::new(tok_rx.clone(), tm_tx.clone(), tokenizer.clone()), + ); + } + + // Egress heartbeat: bumped per drained frame, watched by `/health_generate`. + let egress_activity: tokenizer_manager::ActivityCounter = + Arc::new(std::sync::atomic::AtomicU64::new(0)); + + // --- Egress dispatcher: drains egress ring → routes chunks to shards --- + { + // First TM core; egress is the hotter router (every output token). One + // worker today via `spawn_pool`, so sharding by `Rid::shard` later (see + // `TM_CORES`) is just a larger count + per-shard receivers. + let cores = plan + .as_ref() + .and_then(|p| p.tm.first().copied()) + .map(|c| vec![c]); + let mut egress_rx = Some(egress_rx); // moved into the single worker + let activity = egress_activity.clone(); + let shutdown_rx = shutdown_rx.clone(); + spawn_pool("tm-egress", cores, 1, &mut threads, |_| { + tokenizer_manager::Egress::new( + egress_rx.take().unwrap(), + senders.clone(), + activity.clone(), + shutdown_rx.clone(), + ) + }); + } + + // --- TokenizerManager ingress loop --- + { + // Second TM core when present, else share the first (1-core / API-set + // fallback) — still off the CPU-bound pool cores either way. + let cores = plan + .as_ref() + .and_then(|p| p.tm.get(1).or_else(|| p.tm.first()).copied()) + .map(|c| vec![c]); + let limits = tokenizer_manager::Limits::try_from(&*cfg.server_args) + .map_err(|e| format!("ingress limits: {e}"))?; + let mut parts = Some((tm_rx, ingress_tx)); // moved into the single worker + let shutdown_rx = shutdown_rx.clone(); + spawn_pool("tm-ingress", cores, 1, &mut threads, |_| { + let (tm_rx, ingress_tx) = parts.take().unwrap(); + tokenizer_manager::Ingress::new( + tm_rx, + abort_rx.clone(), + senders.clone(), + ingress_tx, + limits.clone(), + shutdown_rx.clone(), + ) + }); + } + + // --- API server (tokio, I/O bound) --- + { + let cfg = cfg.clone(); + let api_cores = plan.as_ref().map(|p| p.api.clone()); + let senders = senders.clone(); + let api_activity = egress_activity.clone(); + let shutdown_rx = shutdown_rx.clone(); + let handle = std::thread::Builder::new() + .name("api-runtime".into()) + .spawn(move || { + let mut builder = tokio::runtime::Builder::new_multi_thread(); + builder + .worker_threads(cfg.rust_server_args.api_worker_num) + .enable_all(); + if let Some(cores) = api_cores { + let next = std::sync::atomic::AtomicUsize::new(0); + builder.on_thread_start(move || { + let idx = next.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + if let Some(c) = cores.get(idx % cores.len()) { + core_affinity::set_for_current(*c); + } + }); + } + let rt = builder.build().expect("build api runtime"); + rt.block_on(api_server::serve( + listener, + senders, + cfg.rust_server_args.channel_cap, + cfg.server_args.clone(), + // Egress heartbeat watched by `/health_generate`. + api_activity, + shutdown_rx, + )) + }) + .expect("spawn api runtime"); + threads.push(handle); + } + + Ok(Runtime { + ingress: ingress_rx, + egress: egress_tx, + threads: Mutex::new(threads), + shutdown_tx: Mutex::new(Some(shutdown_tx)), + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Minimal boot args. `skip_tokenizer_init` avoids loading a tokenizer/detok + /// model; `model_config` carries the two fields `Limits::from_server_args` + /// requires. They are mandatory at boot, so a fixture without them panics the + /// runtime instead of exercising what these tests are about — `start` does not + /// run `ServerArgs::validate_mandatory` itself, `Server::start` does. + const TEST_SERVER_ARGS: &str = r#"{ + "skip_tokenizer_init": true, + "model_config": {"context_len": 2048, "vocab_size": 1000} + }"#; + + /// Regression: `request_shutdown` must actually stop the API server — it joins + /// the api thread once the listener closes, so the port stops accepting. + /// (Previously it set an unread flag and the port kept accepting.) + #[test] + fn request_shutdown_closes_listener() { + // Pick a free port: bind :0, read the assigned addr, release it. + let probe = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let addr = probe.local_addr().unwrap(); + drop(probe); + + // `skip_tokenizer_init` → no tokenizer/detok model load; minimal boot. + let server_args = ServerArgs::from_json(TEST_SERVER_ARGS).unwrap(); + let cfg = RuntimeConfig { + rust_server_args: RustServerServerArgs { + http_addr: addr, + api_worker_num: 1, + ..Default::default() + }, + server_args: Arc::new(server_args), + }; + // Bind is synchronous in `start`, so the port is already accepting. + let rt = start(cfg).expect("start runtime"); + assert!( + std::net::TcpStream::connect(addr).is_ok(), + "server not listening on {addr} after start returned", + ); + + // Joins the api thread; the listener is closed by the time it returns. + rt.request_shutdown(); + + assert!( + std::net::TcpStream::connect(addr).is_err(), + "port still accepting connections after shutdown", + ); + } + + /// Regression: shutdown must return promptly even with an in-flight `/generate`. + /// No scheduler drains the ingress ring or feeds the egress ring here, so the + /// handler parks on its egress channel forever. Graceful shutdown would wait + /// for it (deadlock → only the 5s bounded-join fallback returns); the + /// non-graceful path cancels the handler via the api runtime drop, whose + /// `AbortGuard` releases the last `Senders` clone so the workers exit. + #[test] + fn shutdown_returns_with_in_flight_request() { + use std::io::Write; + use std::time::{Duration, Instant}; + + let probe = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let addr = probe.local_addr().unwrap(); + drop(probe); + + let server_args = ServerArgs::from_json(TEST_SERVER_ARGS).unwrap(); + let cfg = RuntimeConfig { + rust_server_args: RustServerServerArgs { + http_addr: addr, + api_worker_num: 1, + ..Default::default() + }, + server_args: Arc::new(server_args), + }; + let rt = start(cfg).expect("start runtime"); + + // Fire a request that will block (already-tokenized → valid → pushed to the + // ring, then the handler awaits egress frames that never arrive). + let mut conn = std::net::TcpStream::connect(addr).expect("connect"); + let body = r#"{"input_ids":[1,2,3],"stream":false,"sampling_params":{"max_new_tokens":8}}"#; + let req = format!( + "POST /generate HTTP/1.1\r\nHost: t\r\nContent-Type: application/json\r\n\ + Content-Length: {}\r\nConnection: close\r\n\r\n{}", + body.len(), + body + ); + conn.write_all(req.as_bytes()).unwrap(); + conn.flush().unwrap(); + std::thread::sleep(Duration::from_millis(300)); // reach the blocked state + + let t = Instant::now(); + rt.request_shutdown(); + let elapsed = t.elapsed(); + assert!( + elapsed < Duration::from_secs(3), + "shutdown took {elapsed:?} with an in-flight request (deadlock?)", + ); + drop(conn); + } + + /// Regression: a >2MB body must reach the JSON layer and fail on its + /// *content* (unknown field → 4xx), never on size (413). + #[test] + fn accepts_multi_megabyte_generate_body() { + use std::io::{Read, Write}; + + let probe = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let addr = probe.local_addr().unwrap(); + drop(probe); + + let server_args = ServerArgs::from_json(TEST_SERVER_ARGS).unwrap(); + let cfg = RuntimeConfig { + rust_server_args: RustServerServerArgs { + http_addr: addr, + api_worker_num: 1, + ..Default::default() + }, + server_args: Arc::new(server_args), + }; + let rt = start(cfg).expect("start runtime"); + + // ~3MB of input_ids plus a `text`, which is mutually exclusive with them: + // the body parses in full and is then rejected by `into_requests` with a + // 400, proving it got past any size limit (a 413 would fire before + // parsing). The rejection must come from OUR validation, not from serde — + // an unknown field used to serve here, but unknown fields are now ignored + // to match Python, so such a body would be accepted, dispatched to a ring + // nobody drains in this test, and hang the connection. + let ids = "1,".repeat(1_500_000); + let body = format!( + r#"{{"input_ids":[{}1],"text":"x","sampling_params":{{"max_new_tokens":1}}}}"#, + ids + ); + assert!(body.len() > 2 * 1024 * 1024, "test body must exceed 2MB"); + + let mut conn = std::net::TcpStream::connect(addr).expect("connect"); + let req = format!( + "POST /generate HTTP/1.1\r\nHost: t\r\nContent-Type: application/json\r\n\ + Content-Length: {}\r\nConnection: close\r\n\r\n{}", + body.len(), + body + ); + conn.write_all(req.as_bytes()).unwrap(); + conn.flush().unwrap(); + + let mut response = String::new(); + conn.read_to_string(&mut response).unwrap(); + let status_line = response.lines().next().unwrap_or(""); + let code: u16 = status_line + .split_whitespace() + .nth(1) + .and_then(|c| c.parse().ok()) + .unwrap_or(0); + // A 400 from the mutually-exclusive-inputs check proves the body was read + // and parsed in full; 413 would mean it was rejected on size beforehand. + assert!( + (400..500).contains(&code) && code != 413, + "expected a JSON-layer 4xx (not 413), got: {status_line}" + ); + + rt.request_shutdown(); + } + + /// Regression: a port conflict must fail `start` (so the scheduler doesn't + /// advertise ready), not return an `Ok` runtime whose listener never binds. + #[test] + fn start_fails_on_port_conflict() { + // Hold the port so the runtime's bind conflicts (EADDRINUSE). + let hog = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let addr = hog.local_addr().unwrap(); + + let server_args = ServerArgs::from_json(TEST_SERVER_ARGS).unwrap(); + let cfg = RuntimeConfig { + rust_server_args: RustServerServerArgs { + http_addr: addr, + api_worker_num: 1, + ..Default::default() + }, + server_args: Arc::new(server_args), + }; + let err = match start(cfg) { + Ok(_) => panic!("bind conflict must fail startup, got Ok"), + Err(e) => e, + }; + assert!(err.contains("bind"), "error should mention bind: {err}"); + } + + /// `server_args` missing a mandatory `model_config` field must be a startup + /// ERROR, not a panic. + /// + /// `Limits::try_from` is fallible and the ingress loop is built inside a + /// `spawn_pool` closure, so resolving it there would put the failure on a + /// freshly spawned worker thread — a thread `start` never inspects. The boot + /// would report success and the server would accept connections with no + /// ingress loop behind them, hanging every request instead of refusing to + /// start. Only `Server::start` runs `validate_mandatory`, so `start` cannot + /// assume these fields are present. + #[test] + fn start_fails_when_model_config_is_incomplete() { + let probe = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let addr = probe.local_addr().unwrap(); + drop(probe); + + // Boots fine in every other respect — only `model_config` is absent. + let server_args = ServerArgs::from_json(r#"{"skip_tokenizer_init": true}"#).unwrap(); + let cfg = RuntimeConfig { + rust_server_args: RustServerServerArgs { + http_addr: addr, + api_worker_num: 1, + ..Default::default() + }, + server_args: Arc::new(server_args), + }; + let err = match start(cfg) { + Ok(_) => panic!("an incomplete model_config must not boot, got Ok"), + Err(e) => e, + }; + assert!(err.contains("ingress limits"), "{err}"); + } +} diff --git a/rust/sglang-server/src/tokenizer_manager.rs b/rust/sglang-server/src/tokenizer_manager.rs index a830a7b4a..29bf73e91 100644 --- a/rust/sglang-server/src/tokenizer_manager.rs +++ b/rust/sglang-server/src/tokenizer_manager.rs @@ -10,6 +10,12 @@ //! the rest of the pipeline only through `flume` channels: [`TmEvent`] into the //! ingress loop, [`Senders`] fanning out to the pools. +mod egress; +mod ingress; + +pub use egress::{ActivityCounter, Egress}; +pub use ingress::{Ingress, Limits}; + use crate::ids::Rid; use crate::message::{DetokMsg, Request};