From e2c65af2299c17de11d66af09a8e1f5ba45f04a5 Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Thu, 30 Jul 2026 12:46:17 -0700 Subject: [PATCH] add the rust server ingress request validation and api server common types (#32873) --- rust/sglang-server/src/api_server/common.rs | 226 +++++++ rust/sglang-server/src/api_server/log.rs | 41 ++ rust/sglang-server/src/api_server/openai.rs | 55 ++ .../src/tokenizer_manager/ingress.rs | 566 ++++++++++++++++++ 4 files changed, 888 insertions(+) create mode 100644 rust/sglang-server/src/api_server/common.rs create mode 100644 rust/sglang-server/src/api_server/log.rs create mode 100644 rust/sglang-server/src/api_server/openai.rs create mode 100644 rust/sglang-server/src/tokenizer_manager/ingress.rs diff --git a/rust/sglang-server/src/api_server/common.rs b/rust/sglang-server/src/api_server/common.rs new file mode 100644 index 000000000..70a7b7323 --- /dev/null +++ b/rust/sglang-server/src/api_server/common.rs @@ -0,0 +1,226 @@ +//! Common control-plane endpoints — `/server_info`, `/get_model_info` +//! (+ `/model_info` alias), plus the control-request submission path +//! (`await_control_result`, on the shared `submit`). Data-plane endpoints (incl. `/health*`, +//! which round-trips a generate probe) live in the sibling `native_api` and +//! `openai` modules; the shared `AppState` lives in the parent +//! `api_server` module. + +use axum::{ + Router, + extract::State, + http::StatusCode, + response::{IntoResponse, Response}, + routing::get, +}; + +use super::AppState; +use super::guard::AbortGuard; +use super::submit::submit; +use crate::message::{ControlRequest, EgressItem, GetInternalStateReq, RequestKind}; +use crate::runtime::ServerArgs; + +/// The routes this module owns, mounted by `api_server::serve`. +pub(super) fn routes() -> Router { + Router::new() + // Control-plane: reuses the ingress FSM (no tokenization), returns one + // non-streamed JSON result. Adding one = a route line + its struct tag. + .route("/server_info", get(server_info)) + // Static config, no scheduler round-trip. `/get_model_info` (+ `/model_info` + // alias). + .route("/get_model_info", get(model_info)) + .route("/model_info", get(model_info)) +} + +/// Submit a control request through the ingress FSM (no tokenization) and await the +/// scheduler's single msgpack result (a `structs.asdict` named map). Returns the +/// raw bytes, or an error `Response` to return as-is. +async fn await_control_result( + state: &AppState, + control: ControlRequest, +) -> Result { + let (rid, mut rx) = submit(state, RequestKind::Control(Box::new(control)), false).await?; + // Control requests register a detok entry like any other, and only + // `handle_result` removes it — so a request that never produces one (a stalled + // scheduler, a client that hangs up mid-await) leaves the entry behind. A + // monitor polling `/server_info` then leaks one `DetokState` per poll, forever. + // The guard deregisters on drop; it is disarmed below when the result lands. + let mut guard = AbortGuard::new(state.senders.clone(), rid.clone()); + let received = rx.recv().await; + if received.is_some() { + guard.disarm(&rid); // completed normally — nothing to abort + } + match received { + Some(EgressItem::Control(bytes)) => Ok(bytes), + Some(EgressItem::Error(e)) => { + let code = + StatusCode::from_u16(e.http_status()).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR); + Err((code, e.to_string()).into_response()) + } + // A control request never receives generation frames. + Some(EgressItem::Frame(_)) | Some(EgressItem::Done(_)) => Err(( + StatusCode::INTERNAL_SERVER_ERROR, + "unexpected generation output for control request", + ) + .into_response()), + None => Err((StatusCode::from_u16(499).unwrap(), "request aborted").into_response()), + } +} + +/// `GET /get_model_info` (+ `/model_info` alias) — static model metadata from +/// `server_args` (no scheduler round-trip); `is_generation` always true. +async fn model_info(State(state): State) -> Response { + let sa = &state.server_args; + let body = serde_json::json!({ + "model_path": sa.model_path, + "tokenizer_path": sa.tokenizer_path, + "is_generation": true, + // Python's `TokenizerManager` merges this into every request + // (`{**preferred, **client}`); this server has no equivalent yet, so + // `RustServer.launch` REFUSES to start when it is set. It can therefore + // only be null here — echoing it keeps the field's shape. + "preferred_sampling_params": sa.preferred_sampling_params, + "weight_version": serde_json::Value::Null, + }); + ( + StatusCode::OK, + [("content-type", "application/json")], + serde_json::to_vec(&body).unwrap_or_default(), + ) + .into_response() +} + +/// `GET /server_info` — surface only an allowlist ([`INTERNAL_STATE_ALLOWLIST`] + +/// curated [`ServerArgs`] accessors), never the raw server-args dump (embeds +/// `api_key`/`admin_api_key`; see [`shape_server_info`]). +/// +/// TODO(server_info): Python also includes `kv_events`; add once plumbed. +async fn server_info(State(state): State) -> Response { + let bytes = match await_control_result( + &state, + ControlRequest::GetInternalStateReq(GetInternalStateReq::new( + crate::ids::Rid::new().to_string(), + )), + ) + .await + { + Ok(b) => b, + Err(resp) => return resp, + }; + match shape_server_info(&bytes, &state.server_args) { + Ok(json) => (StatusCode::OK, [("content-type", "application/json")], json).into_response(), + Err(e) => { + tracing::error!(error = %e, "server_info: shaping failed"); + ( + StatusCode::INTERNAL_SERVER_ERROR, + "bad server_info response", + ) + .into_response() + } + } +} + +/// Runtime-metric keys `get_internal_state` adds atop the server-args dump. We copy +/// ONLY these out of `internal_state` (an allowlist), so the co-mingled +/// `api_key`/`admin_api_key` can never reach the response. +const INTERNAL_STATE_ALLOWLIST: &[&str] = &[ + "last_gen_throughput", + "memory_usage", + "effective_max_running_requests_per_dp", + "avg_spec_accept_length", + "step_time_dict", +]; + +fn shape_server_info(msgpack: &[u8], server_args: &ServerArgs) -> Result, String> { + // GetInternalStateReqOutput asdict → `{ "internal_state": { server-args dump + + // metrics }, ... }`. Pull that inner map out (it is NOT safe to expose whole). + let mut obj: serde_json::Map = + rmp_serde::from_slice(msgpack).map_err(|e| e.to_string())?; + let internal = match obj.remove("internal_state") { + Some(serde_json::Value::Object(m)) => m, + _ => serde_json::Map::new(), + }; + + // Copy only the allowlisted runtime metrics — never the raw server-args dump. + let mut state_out = serde_json::Map::new(); + for &k in INTERNAL_STATE_ALLOWLIST { + match internal.get(k) { + Some(v) if !v.is_null() => { + state_out.insert(k.to_string(), v.clone()); + } + _ => {} + } + } + + // Top-level non-secret config from typed accessors (structurally can't surface + // a key field, unlike the raw dump). + let response = serde_json::json!({ + "model_path": server_args.model_path, + "served_model_name": server_args.served_model_name, + "tokenizer_path": server_args.tokenizer_path, + "max_context_length": server_args.model_config.context_len, + "max_total_num_tokens": server_args.max_total_num_tokens, + "version": server_args.version, + "internal_states": [serde_json::Value::Object(state_out)], + }); + serde_json::to_vec(&response).map_err(|e| e.to_string()) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// The scheduler's `internal_state` embeds the full server-args dump (incl. + /// `api_key`/`admin_api_key`). `/server_info` must surface only the allowlisted + /// runtime metrics + curated config — never the secrets — and must not re-nest + /// the dump under `internal_states[].internal_state`. + #[test] + fn shape_server_info_excludes_secrets_and_dump() { + // GetInternalStateReqOutput.asdict → { "internal_state": { …dump+metrics… } }. + let internal = rmpv::Value::Map(vec![ + ( + rmpv::Value::from("api_key"), + rmpv::Value::from("secret-token"), + ), + ( + rmpv::Value::from("admin_api_key"), + rmpv::Value::from("admin-token"), + ), + (rmpv::Value::from("model_path"), rmpv::Value::from("/m")), + ( + rmpv::Value::from("last_gen_throughput"), + rmpv::Value::from(1.5), + ), + ( + rmpv::Value::from("effective_max_running_requests_per_dp"), + rmpv::Value::from(32), + ), + ]); + let outer = rmpv::Value::Map(vec![(rmpv::Value::from("internal_state"), internal)]); + let mut msgpack = Vec::new(); + rmpv::encode::write_value(&mut msgpack, &outer).unwrap(); + + let sa = + ServerArgs::from_json(r#"{"model_path": "/m", "api_key": "secret-token"}"#).unwrap(); + let out = shape_server_info(&msgpack, &sa).unwrap(); + let text = String::from_utf8(out.clone()).unwrap(); + // No secret leaks anywhere in the serialized response. + assert!(!text.contains("secret-token"), "api_key leaked: {text}"); + assert!( + !text.contains("admin-token"), + "admin_api_key leaked: {text}" + ); + + let v: serde_json::Value = serde_json::from_slice(&out).unwrap(); + // Allowlisted metric surfaced; the whole dump did not. + let state0 = &v["internal_states"][0]; + assert_eq!(state0["last_gen_throughput"], 1.5); + assert_eq!(state0["effective_max_running_requests_per_dp"], 32); + assert!( + state0.get("internal_state").is_none(), + "must not re-nest the dump under internal_state" + ); + assert!(state0.get("api_key").is_none()); + // Curated top-level config comes from typed accessors, not the dump. + assert_eq!(v["model_path"], "/m"); + } +} diff --git a/rust/sglang-server/src/api_server/log.rs b/rust/sglang-server/src/api_server/log.rs new file mode 100644 index 000000000..0cf9ab596 --- /dev/null +++ b/rust/sglang-server/src/api_server/log.rs @@ -0,0 +1,41 @@ +//! HTTP access logging — one INFO line per request, content-matching the +//! Python server's uvicorn access log. Gated exactly like uvicorn's +//! (`--log-level-http warning` turns it off, see +//! `ServerArgs::http_access_log_enabled`); when disabled the middleware is not +//! installed at all — zero cost. + +use axum::{Router, response::Response}; + +use crate::runtime::ServerArgs; + +/// Install the access-log middleware when `server_args` enables it; identity +/// otherwise (the layer is never installed, so disabled stays zero-cost). +pub(super) fn apply(app: Router, server_args: &ServerArgs) -> Router { + if server_args.http_access_log_enabled() { + app.layer(axum::middleware::from_fn(access_log)) + } else { + app + } +} + +/// Access log — one INFO line per request, content-matching the Python server's +/// uvicorn access log (`127.0.0.1:54232 - "GET /model_info HTTP/1.1" 200 OK`). +/// Logged when the response head is ready; for SSE that's stream start, same as +/// uvicorn. +async fn access_log( + axum::extract::ConnectInfo(peer): axum::extract::ConnectInfo, + req: axum::extract::Request, + next: axum::middleware::Next, +) -> Response { + let method = req.method().clone(); + let uri = req.uri().clone(); + let version = req.version(); + let res = next.run(req).await; + let status = res.status(); + tracing::info!( + "{peer} - \"{method} {uri} {version:?}\" {} {}", + status.as_u16(), + status.canonical_reason().unwrap_or("") + ); + res +} diff --git a/rust/sglang-server/src/api_server/openai.rs b/rust/sglang-server/src/api_server/openai.rs new file mode 100644 index 000000000..22678f6d3 --- /dev/null +++ b/rust/sglang-server/src/api_server/openai.rs @@ -0,0 +1,55 @@ +//! OpenAI-compatible endpoints: `/v1/completions`, `/v1/chat/completions`, and +//! `/v1/models`. Each runs the same tokenize→generate→detok pipeline as +//! `/generate` and shapes the neutral [`ChunkEvent`] delta into OpenAI types +//! (`dynamo-protocols`), with chat-template rendering (`dynamo-renderer`) and +//! reasoning / tool-call parsing (`dynamo-parsers`). +//! +//! Mounted on the shared [`AppState`](super::AppState) by the parent +//! `api_server` module; the submit machinery and control plane live there. + +use axum::{ + extract::State, + http::StatusCode, + response::{IntoResponse, Response}, +}; + +use axum::{Router, routing::get}; + +use super::AppState; + +/// The routes this module owns, mounted by `api_server::serve`. +pub(super) fn routes() -> Router { + // `/v1/models` is OpenAI-compatible; completions/chat land here too. + Router::new().route("/v1/models", get(available_models)) +} + +/// `GET /v1/models` — OpenAI-compatible model list. Served from `server_args`; +/// no scheduler round-trip. Mirrors `http_server.available_models`. +/// +/// TODO(v1/models): when `--enable-lora`, append a `ModelCard` per loaded LoRA +/// adapter (`id=lora_name, root=lora_path, parent=served_model_name, +/// max_model_len=None`). Adapters load/unload at runtime, so that part needs a +/// control-request query to the scheduler's LoRA registry. +async fn available_models(State(state): State) -> Response { + let created = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_secs()) + .unwrap_or(0); + let name = &state.server_args.served_model_name; + let base = serde_json::json!({ + "id": name, + "object": "model", + "created": created, + "owned_by": "sglang", + "root": name, + "parent": serde_json::Value::Null, + "max_model_len": state.server_args.model_config.context_len, + }); + let list = serde_json::json!({ "object": "list", "data": [base] }); + ( + StatusCode::OK, + [("content-type", "application/json")], + serde_json::to_vec(&list).unwrap_or_default(), + ) + .into_response() +} diff --git a/rust/sglang-server/src/tokenizer_manager/ingress.rs b/rust/sglang-server/src/tokenizer_manager/ingress.rs new file mode 100644 index 000000000..d1514d6df --- /dev/null +++ b/rust/sglang-server/src/tokenizer_manager/ingress.rs @@ -0,0 +1,566 @@ +//! TokenizerManager — ingress side. +//! +//! [`Ingress`] is a single-consumer stage draining one inbox fed by both the API +//! server (fresh requests) and the Tokenizer pool (returned requests). It owns +//! the request while driving the ingress FSM and hands it off by *moving* it to +//! the next stage; nothing here is shared, so no locks. +//! +//! Edges driven here (from the design table): +//! Received → Validating +//! Validating → Normalizing (generate: sampling-param normalize/verify) +//! Validating → PreSendValidating (control: no tokenize, no sampling params) +//! Normalizing → {Encoding | Tokenizing | PreSendValidating} (by ValidationOutcome) +//! Tokenizing → PreSendValidating (on TokenizeDone, when the request returns) +//! PreSendValidating → Queued (checks needing the tokenized length) +//! Queued → ring (handed to the scheduler) +//! +//! The egress edges (Streaming/Finalizing/Completed) are driven on the egress +//! side (see `egress` + `detokenizer`). + +use bytes::Bytes; + +use crate::error::Error; +use crate::fsm::{Event, RequestState, ValidationOutcome}; + +use crate::message::{ + AbortReq, ControlRequest, DetokMsg, EgressItem, GenerateRequest, IngressMsg, Request, + RequestKind, +}; +use crate::ring::IngressProducer; +use crate::runtime::{Runnable, ServerArgs}; +use crate::tokenizer_manager::{AbortSource, Senders, TmEvent}; + +/// Ingress FSM dispatcher stage. Owns its inbox + downstream handles, so the +/// runtime spawns it as a [`Runnable`] rather than calling a free `run_*` fn +/// with positional arguments. +pub struct Ingress { + rx: flume::Receiver, + /// Unbounded abort lane (see [`Senders::abort`]). Selected against `rx` so an + /// abort is handled promptly even while the bounded inbox is saturated. + abort_rx: flume::Receiver, + senders: Senders, + ingress: IngressProducer, + limits: Limits, + shutdown: flume::Receiver<()>, +} + +/// Longest client-supplied rid accepted. It keys the detok table and travels on +/// every chunk, so its length is a recurring cost; Python mints 32-byte uuid hex. +const MAX_RID_LEN: usize = 128; + +/// What ingress admits, resolved once at boot from the scheduler's `server_args`. +/// A struct rather than more positional `new` arguments — these grew from two to +/// six, and every one of them is a `u64`/`bool` that would be trivial to swap at +/// a call site. +/// +/// NOT `Default`-able on purpose. `vocab_size` and `context_len` are mandatory, +/// and their zero value is the most restrictive setting there is — a derived +/// `Default` would silently build limits that reject every request rather than +/// failing loudly. Tests construct these explicitly (see `test_limits`). +#[derive(Clone, Debug)] +pub struct Limits { + /// Token-ids-in mode: a generate request must arrive already tokenized. + pub skip_tokenizer_init: bool, + /// `model_config.vocab_size`; bounds client-supplied token ids. Mandatory — + /// [`ServerArgs::validate_mandatory`](crate::runtime::ServerArgs) rejects a + /// boot without it, so ingress can check unconditionally. + pub vocab_size: u64, + /// `model_config.context_len`, the ceiling for input + `max_new_tokens`. + /// Mandatory, as above. + pub context_len: u64, + /// Output slots reserved on top of the input (eagle draft tokens). + pub num_reserved_tokens: u64, + /// Clamp `max_new_tokens` to what fits instead of rejecting the request. + pub allow_auto_truncate: bool, + /// Whether the server can produce hidden states at all. + pub enable_return_hidden_states: bool, +} + +impl TryFrom<&ServerArgs> for Limits { + type Error = Error; + + fn try_from(sa: &ServerArgs) -> Result { + Ok(Self { + skip_tokenizer_init: sa.skip_tokenizer_init, + vocab_size: sa + .model_config + .vocab_size + .ok_or_else(|| Error::Validation("vocab_size missing".into()))?, + context_len: sa + .model_config + .context_len + .ok_or_else(|| Error::Validation("context_len missing".into()))?, + num_reserved_tokens: sa.num_reserved_tokens, + allow_auto_truncate: sa.allow_auto_truncate, + enable_return_hidden_states: sa.enable_return_hidden_states, + }) + } +} + +impl Ingress { + pub fn new( + rx: flume::Receiver, + abort_rx: flume::Receiver, + senders: Senders, + ingress: IngressProducer, + limits: Limits, + shutdown: flume::Receiver<()>, + ) -> Self { + Self { + rx, + abort_rx, + senders, + ingress, + limits, + shutdown, + } + } +} + +/// Which lane produced the next item. +enum Lane { + Abort(AbortSource), + Event(TmEvent), +} + +impl Runnable for Ingress { + fn run(self) { + loop { + // Select, not a drain-then-block: an abort arriving while the inbox is + // idle must still be handled at once. + let next = flume::Selector::new() + .recv(&self.abort_rx, |r| r.ok().map(Lane::Abort)) + .recv(&self.rx, |r| r.ok().map(Lane::Event)) + .recv(&self.shutdown, |_| None) + .wait(); + match next { + Some(Lane::Abort(rid)) => self.on_abort(rid), + // A fresh request and one returning from the tokenizer pool. + Some(Lane::Event(TmEvent::Ingress(req) | TmEvent::Tokenized(req))) => { + self.drive(req) + } + None => { + // Shutdown, or the inbox closed. Drain whatever is still queued + // on the abort lane first: those requests are in flight on the + // scheduler, and the selector may report the closed inbox before + // it ever looks at a pending abort. + while let Ok(source) = self.abort_rx.try_recv() { + self.on_abort(source); + } + return; + } + } + } + } +} + +impl Ingress { + /// Reject a request: → `Failed`, notify the client, deregister (unconditional + /// — a no-op when nothing was registered). + /// `registered` says whether this request ever reached `register_detok`. It + /// must: `Deregister`'s handler is a bare `table.remove(&rid)`, so a + /// request rejected BEFORE registering would evict whatever entry currently + /// holds that key — a concurrent request's sink — leaving that client with no + /// terminal frame and a hung connection. Python cannot hit this because it + /// validates before `rid_to_state[obj.rid] = state`. + fn fail(&self, req: &mut Request, err: Error, registered: bool) { + // Log only server faults (500); 4xx/499/503 are expected and would spam. + if err.http_status() == 500 { + tracing::error!(rid = %req.rid, error = %err, "ingress rejected request"); + } + let _ = req.state.apply(Event::Error(err.clone())); + let _ = req.sink.try_send(EgressItem::Error(err)); // client may be gone + if registered { + let _ = self.senders.detok_for(&req.rid).send(DetokMsg::Deregister { + rid: req.rid.clone(), + }); + } + } + + /// Drive a request through its ingress states until it terminates (failed or + /// pushed to the ring) or is handed to the tokenizer pool (re-entering as a + /// `Tokenized` event). Each arm acts and advances the FSM; the loop + /// re-dispatches. The arms are the design table's states, `Failed` the single + /// reject path. + fn drive(&self, mut req: Request) { + // Flipped once `register_detok` succeeds; `fail` must not deregister before + // that (see `fail`). A pool return re-enters `drive` already registered. + let mut registered = !matches!(req.state, RequestState::Received); + loop { + match req.state.clone() { + // Validate, then register the sink before the request leaves Rust. + // Failures move to `Failed` and fall through to the reject arm. + RequestState::Received => { + if let Err(e) = validate(&mut req, &self.limits) { + let _ = req.state.apply(Event::Error(e)); // → Failed + continue; + } + if !self.register_detok(&req) { + let _ = req + .state + .apply(Event::Error(Error::Internal("detok shard gone".into()))); + continue; + } + registered = true; + // `validate` advanced Received → Validating; keep driving. + } + // Control skips normalization (no sampling params) straight to the + // pre-send checks; generate goes to Normalizing. + RequestState::Validating => match &req.kind { + RequestKind::Control(_) => { + let _ = req + .state + .apply(Event::Validated(ValidationOutcome::AlreadyTokenized)); + } + RequestKind::Generate(_) => { + let _ = req.state.apply(Event::NeedsNormalize); + } + }, + // Normalize + verify sampling params (off the scheduler loop), then + // pick the branch; a bad param becomes `Failed`. + RequestState::Normalizing => { + let outcome = { + let RequestKind::Generate(g) = &mut req.kind else { + // Unreachable (control never reaches here); reject so a + // bug can't leak/hang a registered request. + self.fail( + &mut req, + Error::Internal("non-generate request in Normalizing".into()), + registered, + ); + return; + }; + match g + .sampling_params + .normalize(self.limits.skip_tokenizer_init, self.limits.vocab_size) + { + Err(e) => Err(e), + // Client ids skip the pool; text goes to the tokenizer. + Ok(()) if g.already_tokenized() => { + Ok(ValidationOutcome::AlreadyTokenized) + } + Ok(()) => Ok(ValidationOutcome::NeedsTokenize), + } + }; + match outcome { + Err(e) => { + let _ = req.state.apply(Event::Error(e)); // → Failed + } + Ok(o) => { + // AlreadyTokenized → Queued, NeedsTokenize → Tokenizing. + let _ = req.state.apply(Event::Validated(o)); + } + } + } + // Hand off to the tokenizer pool; it returns the request as a + // `Tokenized` event (PreSendValidating, or Failed on error). + // Doesn't loop. + RequestState::Tokenizing => { + if let Err(err) = self.senders.tok.send(req) { + // Pool gone (workers exited); flume hands the request back. + let mut req = err.into_inner(); + // Past `Received`, so registration happened. + self.fail( + &mut req, + Error::Internal("tokenizer pool gone".into()), + true, + ); + } + return; + } + // The checks that need the final `input_ids`: every branch + // converges here (client ids arrive directly, text arrives from + // the tokenizer pool), so they run once per request regardless of + // how it was tokenized. `validate` runs too early — at `Received` + // a text request has no ids yet. + RequestState::PreSendValidating => { + if let RequestKind::Generate(g) = &mut req.kind + && let Err(e) = check_total_tokens(g, &self.limits) + { + let _ = req.state.apply(Event::Error(e)); // → Failed + continue; + } + let _ = req.state.apply(Event::PreSendValidated); // → Queued + } + // Push the wire message (control frame or generate payload) to the ring. + RequestState::Queued => { + // `matches!` reads the discriminant without holding a borrow, + // so `req` can be moved into the push below. + if matches!(req.kind, RequestKind::Generate(_)) { + self.push_to_ring(req); + } else { + self.push_control_to_ring(req); + } + return; + } + // The single reject path for every post-register failure. + RequestState::Failed(e) => { + self.fail(&mut req, e, registered); + return; + } + // Unreachable (egress states never reach here). Reject via `fail`/ + // return (not apply + continue, which would spin on a terminal state). + other => { + self.fail( + &mut req, + Error::Internal(format!("unexpected ingress state: {other:?}")), + registered, + ); + return; + } + } + } + } + + /// Register the egress sink with the owning detok shard (by id) so the response + /// has a home. Carries the per-request detok flags — `return_text_in_logprobs` + /// (decode logprob text on this shard) and `no_stop_trim` (keep the matched + /// stop in the output) — so the shard needs no back-reference to the request. + /// Returns `false` if the shard is gone. + fn register_detok(&self, req: &Request) -> bool { + let (decode_logprob_text, no_stop_trim) = match &req.kind { + RequestKind::Generate(g) => ( + g.return_text_in_logprobs.unwrap_or(false), + g.sampling_params.no_stop_trim, + ), + RequestKind::Control(_) => (false, false), + }; + self.senders + .detok_for(&req.rid) + .send(DetokMsg::Register { + rid: req.rid.clone(), + sink: req.sink.clone(), + decode_logprob_text, + no_stop_trim, + }) + .is_ok() + } + + /// Push a bare control request (`[tag, rid, nil]`) onto the ingress ring. The + /// scheduler dispatches it (e.g. `GetInternalStateReq`) and replies via the + /// egress ring as a single `Result`. + fn push_control_to_ring(&self, mut req: Request) { + let encode = match &req.kind { + RequestKind::Control(control) => control.encode(), + _ => Err(Error::Internal( + "non-control request reached push_control_to_ring".into(), + )), + }; + let header = match encode { + Ok(b) => b, + Err(e) => { + self.fail(&mut req, e, true); // on the push path: registered + return; + } + }; + // Control requests carry no tensor cell — empty `ids`. + if !self.ingress.try_push(IngressMsg { + header, + ids: Bytes::new(), + }) { + self.fail(&mut req, Error::QueueFull, true); // registered + } + } + + /// Client disconnected (or a detok terminal): deregister the sink, then push an + /// `AbortReq(rid)` so the scheduler stops generating for it. + /// + /// A failed push is logged, not retried: the scheduler keeps generating and the + /// chunks arrive for a rid no longer in the detok table, where they are dropped. + /// That wastes GPU work until the request finishes on its own, but it cannot be + /// misdelivered — the rid is unique to this request for the process's lifetime + /// ([`Rid::from_client`]), so no later request can ever answer to it. + fn on_abort(&self, source: AbortSource) { + let rid = source.rid().clone(); + let _ = self + .senders + .detok_for(&rid) + .send(DetokMsg::Deregister { rid: rid.clone() }); + + // The ring is BOUNDED and drops pushes under exactly the load this matters + // for, so report the miss rather than assuming the scheduler was told. + match ControlRequest::AbortReq(AbortReq::new(rid.as_str().to_string(), false)).encode() { + Ok(header) => { + if !self.ingress.try_push(IngressMsg { + header, + ids: Bytes::new(), + }) { + tracing::error!( + rid = %rid, + "abort dropped: ingress ring full; the scheduler keeps generating \ + for this request until it finishes on its own" + ); + } + } + Err(e) => tracing::error!(rid = %rid, error = %e, "abort encode failed"), + } + } + + /// Serialize the tokenized request to its `TokenizedGenerateReqInput` wire and + /// push it onto the ingress ring for the scheduler. On backpressure, fail it. + fn push_to_ring(&self, mut req: Request) { + // Only generate requests reach here (control uses `push_control_to_ring`). + // Validate + serialize while borrowing `g` immutably; the resulting `Bytes` + // own their data, so the borrow ends before any `fail(&mut req)`. + let serialized = match &req.kind { + RequestKind::Generate(g) if g.already_tokenized() => g + .encode_header() + .map(|header| (header, g.encode_data_buf())), + RequestKind::Generate(_) => Err(Error::Tokenize("empty input_ids".into())), + _ => Err(Error::Internal( + "non-generate request reached push_to_ring".into(), + )), + }; + let (header, ids) = match serialized { + Ok(v) => v, + Err(e) => { + self.fail(&mut req, e, true); // on the push path: registered + return; + } + }; + + if !self.ingress.try_push(IngressMsg { header, ids }) { + self.fail(&mut req, Error::QueueFull, true); // registered + } + // On success the scheduler owns the request (egress arrives by rid); we + // drop our `Request` here — the detok shard holds the sink. + } +} + +/// `Received → Validating` + admissibility check. Under `skip_tokenizer_init` a +/// generate request must already carry token ids (no tokenizer to byte-encode +/// text); control requests carry none and are exempt. +fn validate(req: &mut Request, limits: &Limits) -> Result<(), Error> { + let (skip_tokenizer_init, vocab_size) = (limits.skip_tokenizer_init, limits.vocab_size); + let _ = req + .state + .apply(Event::Validated(ValidationOutcome::NeedsTokenize)); + + // The rid is the request's identity everywhere downstream: it keys the detok + // table, and it rides on EVERY chunk of EVERY decode step. An unbounded + // client-supplied rid is therefore a per-step cost, not a one-off. Python's is + // a 32-byte uuid hex, so this is generous. + // Measured on the CLIENT-facing form: the uniquifier `Rid::from_client` appends + // is this server's own overhead, and charging the client for bytes it did not + // send would reject a rid exactly at the documented limit. + let client_rid_len = req.rid.client_facing().len(); + if client_rid_len > MAX_RID_LEN { + return Err(Error::Validation(format!( + "rid is {client_rid_len} bytes, over the {MAX_RID_LEN}-byte limit" + ))); + } + if skip_tokenizer_init + && matches!(&req.kind, RequestKind::Generate(g) if !g.already_tokenized()) + { + // `Validation` (400), not `Tokenize` (500): the client sent a request this + // server cannot serve, which is their error to fix — Python 400s it too. + return Err(Error::Validation( + "skip_tokenizer_init is set: request must provide input_ids".into(), + )); + } + + // Client-supplied token ids must be in-vocabulary: an out-of-range id + // reaches the embedding lookup and kills the scheduler process, so 400 + // here instead — mirroring the Python `TokenizerManager` validation. + if let RequestKind::Generate(g) = &req.kind { + if let Some(ids) = &g.input_ids { + for &id in ids { + if id < 0 || id as u64 >= vocab_size { + return Err(Error::Validation(format!( + "input_ids contains out-of-vocabulary token id {id}; \ + valid range is [0, {vocab_size})" + ))); + } + } + } + if let Some(ids) = &g.token_ids_logprob { + for &id in ids { + if id < 0 || id as u64 >= vocab_size { + return Err(Error::Validation(format!( + "token_ids_logprob contains out-of-vocabulary token id \ + {id}; valid range is [0, {vocab_size})" + ))); + } + } + } + } + + // The scheduler only computes hidden states when launched for it, so without + // this the request would 200 with `meta_info.hidden_states` silently absent + // (Python `TokenizerManager._validate_one_request`). + if !limits.enable_return_hidden_states + && matches!(&req.kind, RequestKind::Generate(g) if g.return_hidden_states) + { + return Err(Error::Validation( + "The server is not configured to return the hidden states. \ + Please set `--enable-return-hidden-states` to enable this feature." + .into(), + )); + } + + Ok(()) +} + +/// The context-window checks that need the tokenized length, mirroring Python +/// `TokenizerManager._validate_one_request`: the input alone must fit, and then +/// input + `max_new_tokens` must fit. Without them the scheduler silently clamps +/// and the client gets a 200 with a truncated completion instead of an actionable +/// 400. +/// +/// Under `allow_auto_truncate` both clamp instead of rejecting — the launch flag +/// opted into that. +fn check_total_tokens(g: &mut GenerateRequest, limits: &Limits) -> Result<(), Error> { + let max_req_len = limits.context_len; + // Python counts the reserved slots as part of the input, so a request can be + // rejected for them even when the prompt alone fits. + let input_len = + g.input_ids.as_ref().map_or(0, |ids| ids.len()) as u64 + limits.num_reserved_tokens; + + // Input length first, and unconditionally: `max_new_tokens: null` means "no + // cap", which must not disable this. Python's comparison is `>=` — a prompt + // that exactly fills the window leaves no room to generate. + if input_len >= max_req_len { + if !limits.allow_auto_truncate { + return Err(Error::Validation(format!( + "The input ({input_len} tokens) is longer than the model's context \ + length ({max_req_len} tokens)." + ))); + } + if let Some(ids) = &mut g.input_ids { + ids.truncate(max_req_len as usize); + } + } + let input_len = + g.input_ids.as_ref().map_or(0, |ids| ids.len()) as u64 + limits.num_reserved_tokens; + + let Some(max_new_tokens) = g.sampling_params.max_new_tokens else { + return Ok(()); // no cap requested → nothing to add to the input length + }; + let total = input_len.saturating_add(max_new_tokens.max(0) as u64); + if total <= max_req_len { + return Ok(()); + } + if !limits.allow_auto_truncate { + return Err(Error::Validation(format!( + "Requested token count exceeds the model's maximum context length of \ + {max_req_len} tokens. You requested a total of {total} tokens: {input_len} \ + tokens from the input messages and {max_new_tokens} tokens for the \ + completion. Please reduce the number of tokens in the input messages or \ + the completion to fit within the limit." + ))); + } + let clamped = max_req_len.saturating_sub(input_len) as i64; + // Re-check what the clamp can break. `verify` already ran (in Normalizing), so + // lowering `max_new_tokens` here can leave `min_new_tokens > max_new_tokens` — + // and `is_normalized: true` stops the scheduler from re-verifying, so nothing + // downstream would catch it. Python validates before it verifies; we can't + // reorder the FSM, so we re-assert the one invariant the clamp can violate. + if g.sampling_params.min_new_tokens > clamped { + return Err(Error::Validation(format!( + "min_new_tokens must be in [0, max_new_tokens({clamped})], got {}", + g.sampling_params.min_new_tokens + ))); + } + g.sampling_params.max_new_tokens = Some(clamped); + Ok(()) +}