sglang rust server environ fsm error id gen (#32240)
This commit is contained in:
@@ -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<T: Copy + std::fmt::Debug>(name: &str, default: T, parse: impl Fn(&str) -> Option<T>) -> 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);
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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<T> = std::result::Result<T, Error>;
|
||||||
@@ -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(())
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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_<uuid4 hex>`, 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);
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -11,6 +11,11 @@
|
|||||||
//! All are non-blocking, so the GIL is never held across a wait.
|
//! All are non-blocking, so the GIL is never held across a wait.
|
||||||
#![allow(dead_code)] // TODO: remove when the consumer PR lands
|
#![allow(dead_code)] // TODO: remove when the consumer PR lands
|
||||||
|
|
||||||
|
mod environ;
|
||||||
|
mod error;
|
||||||
|
mod fsm;
|
||||||
|
mod ids;
|
||||||
|
|
||||||
use pyo3::prelude::*;
|
use pyo3::prelude::*;
|
||||||
|
|
||||||
#[pymodule]
|
#[pymodule]
|
||||||
|
|||||||
Reference in New Issue
Block a user