wire the rust server modules into lib, runtime, and tokenizer manager (#32877)

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.

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;
 

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<Py<PyBytes>>,
+    /// The raw-data plane today just all requests' raw little-endian int64
+    /// ids, concatenated; sliced per request via `lengths`.
+    data: Py<PyBytes>,
+    /// Per-request token count (0 for control requests).
+    lengths: Vec<u32>,
+}
+
+/// 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<String>,
+        ingress_ring_cap: usize,
+        egress_ring_cap: usize,
+        channel_cap: usize,
+        cores: Option<Vec<usize>>,
+        server_args_json: &str,
+    ) -> PyResult<Self> {
+        // 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::<pyo3::exceptions::PyValueError, _>(format!(
+                    "bad server_args_json: {e}"
+                ))
+            })?;
+        server_args.validate_mandatory().map_err(|e| {
+            PyErr::new::<pyo3::exceptions::PyValueError, _>(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::<pyo3::exceptions::PyValueError, _>(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::<pyo3::exceptions::PyValueError, _>(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<IngressBatch> {
+        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<PyBackedBytes>) -> 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<tracing_appender::non_blocking::WorkerGuard> =
+    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::<Server>()?;
+    m.add_class::<IngressBatch>()?;
     Ok(())
 }

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};
 

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<Vec<JoinHandle<()>>>,
+    /// The single shutdown sender.
+    shutdown_tx: Mutex<Option<flume::Sender<()>>>,
+}
+
+/// 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<JoinHandle<()>> = 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<Runtime, String> {
+    // 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::<TmEvent>(cfg.rust_server_args.channel_cap);
+    let (tok_tx, tok_rx) =
+        flume::bounded::<crate::message::Request>(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::<DetokMsg>(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::<crate::tokenizer_manager::AbortSource>();
+    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<dyn tokenizer::TextTokenizer> =
+            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}");
+    }
+}

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};