From be7cc173075c1c906f9a6589d04ba20eb03064d5 Mon Sep 17 00:00:00 2001 From: Rain Jiang <96632942+rainj-me@users.noreply.github.com> Date: Fri, 24 Jul 2026 13:55:10 -0700 Subject: [PATCH] sglang rust server environ fsm error id gen (#32240) --- rust/sglang-server/src/environ.rs | 101 ++++++++++++++++++++++ rust/sglang-server/src/error.rs | 53 ++++++++++++ rust/sglang-server/src/fsm.rs | 135 ++++++++++++++++++++++++++++++ rust/sglang-server/src/ids.rs | 92 ++++++++++++++++++++ rust/sglang-server/src/lib.rs | 5 ++ 5 files changed, 386 insertions(+) create mode 100644 rust/sglang-server/src/environ.rs create mode 100644 rust/sglang-server/src/error.rs create mode 100644 rust/sglang-server/src/fsm.rs create mode 100644 rust/sglang-server/src/ids.rs diff --git a/rust/sglang-server/src/environ.rs b/rust/sglang-server/src/environ.rs new file mode 100644 index 000000000..e1a691e5f --- /dev/null +++ b/rust/sglang-server/src/environ.rs @@ -0,0 +1,101 @@ +//! Env-var parsing with the semantics of Python `sglang.srt.environ.EnvField`: +//! 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. +pub fn env_bool(name: &str, default: bool) -> bool { + read(name, default, |raw| match raw.to_lowercase().as_str() { + "true" | "1" | "yes" | "y" => Some(true), + "false" | "0" | "no" | "n" => Some(false), + _ => None, + }) +} + +/// Deliberately restricted unsigned parser — NOT Python `int()` semantics. +/// Accepts only what `u64::from_str` does: ASCII digits with an optional +/// leading `+`, up to `u64::MAX`. Inputs Python's `EnvInt` would accept — +/// surrounding whitespace (`" 45 "`), digit-group underscores (`"4_5"`), +/// negatives, values above `u64::MAX`, non-ASCII digits — warn and fall back +/// to the default, like any other invalid value. Callers are counts/sizes, so +/// strictness over parity is intentional here. +pub fn env_u64(name: &str, default: u64) -> u64 { + read(name, default, |raw| raw.parse().ok()) +} + +/// Shared read-or-default: unset → default; a set-but-unparsable value warns +/// and falls back to the default (mirrors `EnvField.get`'s `warnings.warn`). +fn read(name: &str, default: T, parse: impl Fn(&str) -> Option) -> T { + let Ok(raw) = std::env::var(name) else { + return default; + }; + match parse(&raw) { + Some(v) => v, + None => { + tracing::warn!(name, value = %raw, ?default, "invalid env value; using default"); + default + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// The accepted literal sets are copied from Python `EnvBool.parse` + /// (`true/1/yes/y` / `false/0/no/n`, case-insensitive; invalid → default) — + /// parity pins, not this crate's invention. + #[test] + fn env_bool_matches_python_envbool_parse() { + // Unique var name per case: tests in this binary run concurrently and + // share the process environment. + for (raw, want) in [ + ("true", true), + ("1", true), + ("YES", true), + ("y", true), + ("false", false), + ("0", false), + ("No", false), + ("n", false), + // Invalid → default (here: true), matching the warn-and-default path. + ("off", true), + ("2", true), + ] { + let name = format!("SGLANG_TEST_ENV_BOOL_{raw}"); + unsafe { std::env::set_var(&name, raw) }; + assert_eq!(env_bool(&name, true), want, "value {raw:?}"); + } + assert!(env_bool("SGLANG_TEST_ENV_BOOL_UNSET", true)); + assert!(!env_bool("SGLANG_TEST_ENV_BOOL_UNSET", false)); + } + + /// `env_u64`: strict `u64::from_str` grammar; everything else — including + /// int()-valid inputs the doc calls out as deliberately rejected — → default. + #[test] + fn env_u64_parses_or_defaults() { + for (i, (raw, want)) in [ + ("45", 45), + ("+45", 45), // u64::from_str allows a leading `+` + // Invalid → default. + ("20s", 20), + ("", 20), + // int()-valid but deliberately rejected → default. + (" 45 ", 20), + ("4_5", 20), + ("-1", 20), + ("18446744073709551616", 20), // u64::MAX + 1 + ("١٢", 20), // non-ASCII digits + ] + .into_iter() + .enumerate() + { + let name = format!("SGLANG_TEST_ENV_U64_{i}"); + unsafe { std::env::set_var(&name, raw) }; + assert_eq!(env_u64(&name, 20), want, "value {raw:?}"); + } + assert_eq!(env_u64("SGLANG_TEST_ENV_U64_UNSET", 20), 20); + } +} diff --git a/rust/sglang-server/src/error.rs b/rust/sglang-server/src/error.rs new file mode 100644 index 000000000..a1eff4082 --- /dev/null +++ b/rust/sglang-server/src/error.rs @@ -0,0 +1,53 @@ +//! 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; + +// Some variants are emitted only once their stage matures (real validation, +// the deferred Encoder, HF detok). They are part of the stable error surface. +#[allow(dead_code)] +#[derive(Debug, Clone, Error)] +pub enum Error { + #[error("validation failed: {0}")] + Validation(String), + + #[error("tokenize failed: {0}")] + Tokenize(String), + + #[error("encode failed: {0}")] + Encode(String), + + #[error("detokenize failed: {0}")] + Detokenize(String), + + /// Ingress ring full / scheduler not draining. Surfaced as backpressure. + #[error("ingress queue full")] + QueueFull, + + /// Client went away mid-stream. Drives `Aborted`, not `Failed`. + #[error("client disconnected")] + Disconnected, + + #[error("serialization error: {0}")] + Codec(String), + + #[error("internal error: {0}")] + Internal(String), +} + +impl Error { + /// HTTP status to surface for the non-streaming error path. Mirrors the + /// codes used in the Python `_create_error_response`. + pub fn http_status(&self) -> u16 { + match self { + Error::Validation(_) => 400, + Error::Disconnected => 499, // nginx-style client closed request + Error::QueueFull => 503, + _ => 500, + } + } +} + +#[allow(dead_code)] +pub type Result = std::result::Result; diff --git a/rust/sglang-server/src/fsm.rs b/rust/sglang-server/src/fsm.rs new file mode 100644 index 000000000..9206fa708 --- /dev/null +++ b/rust/sglang-server/src/fsm.rs @@ -0,0 +1,135 @@ +//! Request lifecycle FSM. +//! +//! The state lives *inside* the owned request struct (see [`crate::message`]), +//! so transitions are in-place mutations on a single owner — no shared state, +//! no locks. Each pipeline stage drives the transition for its own phase and +//! then moves the request to the next stage's channel. +//! +//! Port of the design enum: +//! ```text +//! Received, Validating, Normalizing, Encoding, Tokenizing, Queued, +//! Streaming { chunks_sent }, Finalizing, Completed, Failed(Error), Aborted +//! ``` +#![allow(dead_code)] // TODO: remove when the consumer PR lands + +use crate::error::Error; + +// `Failed(Error)` carries the cause for observability even where it isn't read +// back yet; `EncodeDone` belongs to the deferred Encoder edge. +#[allow(dead_code)] +#[derive(Debug, Clone)] +pub enum RequestState { + Received, + Validating, + /// Generate-only: sampling params normalized + verified before routing. + Normalizing, + Encoding, + Tokenizing, + Queued, + Streaming { + chunks_sent: u64, + }, + Finalizing, + Completed, + Failed(Error), + Aborted, +} + +/// Outcome of validation, selecting the ingress branch. +#[derive(Debug, Clone, Copy)] +pub enum ValidationOutcome { + /// Has multimodal inputs → Encoding. Deferred: no encoder yet. + #[allow(dead_code)] + HasMultimodal, + /// Plain text → Tokenizing. + NeedsTokenize, + /// Caller already supplied token ids → straight to Queued. + AlreadyTokenized, +} + +/// Events that drive transitions. Each variant maps 1:1 to an edge in the +/// design's transition table. +#[allow(dead_code)] // EncodeDone is the deferred Encoder edge. +#[derive(Debug)] +pub enum Event { + // --- ingress --- + Validated(ValidationOutcome), + NeedsNormalize, + EncodeDone, + TokenizeDone, + SchedulerPicked, + // --- egress --- + Chunk { finish: bool }, + FinalFrameSent, + // --- terminal (valid from any state) --- + Error(Error), + Disconnect, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum TransitionError { + /// The (state, event) pair has no defined edge. + Illegal, +} + +impl RequestState { + /// Whether this is a terminal state (no further transitions expected). + pub fn is_terminal(&self) -> bool { + matches!( + self, + RequestState::Completed | RequestState::Failed(_) | RequestState::Aborted + ) + } + + /// Apply `event`, mutating in place. Returns `Err(Illegal)` for undefined + /// edges so the caller can decide whether to log-and-drop or fail the req. + /// + /// Terminal events (`Error`/`Disconnect`) are accepted from *any* non-terminal + /// state, matching `(*, Error | Disconnect) -> Failed | Aborted`. + pub fn apply(&mut self, event: Event) -> Result<(), TransitionError> { + use Event::*; + use RequestState::*; + use ValidationOutcome::*; + + // Wildcard terminal edges first. + match &event { + Error(e) => { + if !self.is_terminal() { + *self = Failed(e.clone()); + } + return Ok(()); + } + Disconnect => { + if !self.is_terminal() { + *self = Aborted; + } + return Ok(()); + } + _ => {} + } + + let next = match (&*self, &event) { + // ingress + (Received, Validated(_)) => Validating, + // Generate requests pass through Normalizing (sampling-param + // normalize/verify); control requests skip straight to Queued. + (Validating, NeedsNormalize) => Normalizing, + (Validating, Validated(AlreadyTokenized)) => Queued, + (Normalizing, Validated(HasMultimodal)) => Encoding, + (Normalizing, Validated(NeedsTokenize)) => Tokenizing, + (Normalizing, Validated(AlreadyTokenized)) => Queued, + (Encoding, EncodeDone) => Tokenizing, + (Tokenizing, TokenizeDone) => Queued, + (Queued, SchedulerPicked) => Streaming { chunks_sent: 0 }, + // egress + (Streaming { chunks_sent }, Chunk { finish: false }) => Streaming { + chunks_sent: chunks_sent + 1, + }, + (Streaming { .. }, Chunk { finish: true }) => Finalizing, + (Finalizing, FinalFrameSent) => Completed, + _ => return Err(TransitionError::Illegal), + }; + *self = next; + Ok(()) + } +} diff --git a/rust/sglang-server/src/ids.rs b/rust/sglang-server/src/ids.rs new file mode 100644 index 000000000..5f3e9f0cc --- /dev/null +++ b/rust/sglang-server/src/ids.rs @@ -0,0 +1,92 @@ +//! Lightweight identifiers used across pipeline stages. +#![allow(dead_code)] // TODO: remove when the consumer PR lands + +use std::{ + collections::hash_map::DefaultHasher, + fmt::{Debug, Display, Formatter, Result}, + hash::{Hash, Hasher}, +}; + +use uuid::Uuid; + +/// Process-local request key, derived by hashing the client-visible rid string +/// (`from_rid`). Cheap to copy, used as the routing key on the egress side +/// (detok shard selection) and to correlate scheduler output chunks back to the +/// originating connection. NOT the identity the client sees — that is the rid +/// string; this is its stable 64-bit digest. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)] +pub struct RidHash(pub u64); + +impl RidHash { + /// Derive the routing key from a rid string. `DefaultHasher::new()` uses + /// fixed keys, so every stage (ingress push, egress decode, abort) computes + /// the same id from the same rid — no shared map. + pub fn from_rid(rid: &str) -> Self { + let mut h = DefaultHasher::new(); + rid.hash(&mut h); + RidHash(h.finish()) + } + + /// Shard index for `n` detokenizer shards. Pure function of the id so the + /// ingress and egress sides agree without any shared map. + #[inline] + pub fn shard(self, n: usize) -> usize { + debug_assert!(n > 0); + (self.0 as usize) % n + } +} + +impl Display for RidHash { + fn fmt(&self, f: &mut Formatter<'_>) -> Result { + write!(f, "req-{}", self.0) + } +} + +/// Mint a fresh client-visible request id — uuid4 hex, matching the Python +/// server's `uuid.uuid4().hex`. +pub fn new_rid() -> String { + Uuid::new_v4().simple().to_string() +} + +/// Health-probe rid prefix — MUST match the Python server's +/// `sglang.srt.constants.HEALTH_CHECK_RID_PREFIX`, so scheduler logs / crash +/// dumps and any prefix-gated logic recognize probes from either server. +pub const HEALTH_CHECK_RID_PREFIX: &str = "HEALTH_CHECK"; + +/// Mint a health-probe rid: `HEALTH_CHECK_`, the Python server's +/// `f"{HEALTH_CHECK_RID_PREFIX}_{uuid.uuid4().hex}"` format. +pub fn new_health_check_rid() -> String { + format!("{HEALTH_CHECK_RID_PREFIX}_{}", new_rid()) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Cross-language format guard: Python rids are `uuid.uuid4().hex` — 32 + /// lowercase hex chars, no hyphens. `.simple()` is the matching uuid-crate + /// encoding; swapping it for the default `to_string()` (36 chars, + /// hyphenated) would silently break the parity. + #[test] + fn rid_matches_python_uuid4_hex_format() { + let rid = new_rid(); + assert_eq!(rid.len(), 32); + assert!( + rid.chars() + .all(|c| c.is_ascii_hexdigit() && !c.is_ascii_uppercase()), + "rid must be lowercase hex: {rid}" + ); + } + + /// Cross-language literal guard: the prefix is dictated by Python's + /// `constants.HEALTH_CHECK_RID_PREFIX` ("HEALTH_CHECK"); drifting silently + /// would break prefix-gated handling (e.g. the disagg encode server). + #[test] + fn health_rid_matches_python_convention() { + assert_eq!(HEALTH_CHECK_RID_PREFIX, "HEALTH_CHECK"); + let rid = new_health_check_rid(); + // "HEALTH_CHECK_" + 32 hex chars + assert!(rid.starts_with("HEALTH_CHECK_")); + assert_eq!(rid.len(), "HEALTH_CHECK_".len() + 32); + } +} diff --git a/rust/sglang-server/src/lib.rs b/rust/sglang-server/src/lib.rs index ab23f963c..f4e369418 100644 --- a/rust/sglang-server/src/lib.rs +++ b/rust/sglang-server/src/lib.rs @@ -11,6 +11,11 @@ //! All are non-blocking, so the GIL is never held across a wait. #![allow(dead_code)] // TODO: remove when the consumer PR lands +mod environ; +mod error; +mod fsm; +mod ids; + use pyo3::prelude::*; #[pymodule]