sglang rust server tokenizer manager, ring and runtime (#32358)
This commit is contained in:
@@ -7,10 +7,10 @@
|
|||||||
//!
|
//!
|
||||||
//! Port of the design enum:
|
//! Port of the design enum:
|
||||||
//! ```text
|
//! ```text
|
||||||
//! Received, Validating, Normalizing, Encoding, Tokenizing, Queued,
|
//! Received, Validating, Normalizing, Encoding, Tokenizing, PreSendValidating,
|
||||||
//! Streaming { chunks_sent }, Finalizing, Completed, Failed(Error), Aborted
|
//! Queued, Streaming { chunks_sent }, Finalizing, Completed, Failed(Error),
|
||||||
|
//! Aborted
|
||||||
//! ```
|
//! ```
|
||||||
#![allow(dead_code)] // TODO: remove when the consumer PR lands
|
|
||||||
|
|
||||||
use crate::error::Error;
|
use crate::error::Error;
|
||||||
|
|
||||||
@@ -25,6 +25,10 @@ pub enum RequestState {
|
|||||||
Normalizing,
|
Normalizing,
|
||||||
Encoding,
|
Encoding,
|
||||||
Tokenizing,
|
Tokenizing,
|
||||||
|
/// Every branch converges here with its final `input_ids`, for the checks
|
||||||
|
/// that need the tokenized length (the input + `max_new_tokens` ceiling).
|
||||||
|
/// The last state before the request leaves Rust.
|
||||||
|
PreSendValidating,
|
||||||
Queued,
|
Queued,
|
||||||
Streaming {
|
Streaming {
|
||||||
chunks_sent: u64,
|
chunks_sent: u64,
|
||||||
@@ -43,7 +47,7 @@ pub enum ValidationOutcome {
|
|||||||
HasMultimodal,
|
HasMultimodal,
|
||||||
/// Plain text → Tokenizing.
|
/// Plain text → Tokenizing.
|
||||||
NeedsTokenize,
|
NeedsTokenize,
|
||||||
/// Caller already supplied token ids → straight to Queued.
|
/// Caller already supplied token ids → straight to the pre-send checks.
|
||||||
AlreadyTokenized,
|
AlreadyTokenized,
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -57,9 +61,13 @@ pub enum Event {
|
|||||||
NeedsNormalize,
|
NeedsNormalize,
|
||||||
EncodeDone,
|
EncodeDone,
|
||||||
TokenizeDone,
|
TokenizeDone,
|
||||||
|
/// The pre-send checks passed; the request may be pushed to the ring.
|
||||||
|
PreSendValidated,
|
||||||
SchedulerPicked,
|
SchedulerPicked,
|
||||||
// --- egress ---
|
// --- egress ---
|
||||||
Chunk { finish: bool },
|
Chunk {
|
||||||
|
finish: bool,
|
||||||
|
},
|
||||||
FinalFrameSent,
|
FinalFrameSent,
|
||||||
// --- terminal (valid from any state) ---
|
// --- terminal (valid from any state) ---
|
||||||
Error(Error),
|
Error(Error),
|
||||||
@@ -112,14 +120,17 @@ impl RequestState {
|
|||||||
// ingress
|
// ingress
|
||||||
(Received, Validated(_)) => Validating,
|
(Received, Validated(_)) => Validating,
|
||||||
// Generate requests pass through Normalizing (sampling-param
|
// Generate requests pass through Normalizing (sampling-param
|
||||||
// normalize/verify); control requests skip straight to Queued.
|
// normalize/verify); control requests skip it, having none.
|
||||||
(Validating, NeedsNormalize) => Normalizing,
|
(Validating, NeedsNormalize) => Normalizing,
|
||||||
(Validating, Validated(AlreadyTokenized)) => Queued,
|
(Validating, Validated(AlreadyTokenized)) => PreSendValidating,
|
||||||
(Normalizing, Validated(HasMultimodal)) => Encoding,
|
(Normalizing, Validated(HasMultimodal)) => Encoding,
|
||||||
(Normalizing, Validated(NeedsTokenize)) => Tokenizing,
|
(Normalizing, Validated(NeedsTokenize)) => Tokenizing,
|
||||||
(Normalizing, Validated(AlreadyTokenized)) => Queued,
|
(Normalizing, Validated(AlreadyTokenized)) => PreSendValidating,
|
||||||
(Encoding, EncodeDone) => Tokenizing,
|
(Encoding, EncodeDone) => Tokenizing,
|
||||||
(Tokenizing, TokenizeDone) => Queued,
|
// Every ingress branch funnels through the pre-send checks, so they
|
||||||
|
// run exactly once per request no matter how it got its ids.
|
||||||
|
(Tokenizing, TokenizeDone) => PreSendValidating,
|
||||||
|
(PreSendValidating, PreSendValidated) => Queued,
|
||||||
(Queued, SchedulerPicked) => Streaming { chunks_sent: 0 },
|
(Queued, SchedulerPicked) => Streaming { chunks_sent: 0 },
|
||||||
// egress
|
// egress
|
||||||
(Streaming { chunks_sent }, Chunk { finish: false }) => Streaming {
|
(Streaming { chunks_sent }, Chunk { finish: false }) => Streaming {
|
||||||
@@ -133,3 +144,59 @@ impl RequestState {
|
|||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
fn after(mut state: RequestState, event: Event) -> RequestState {
|
||||||
|
state.apply(event).expect("edge must exist");
|
||||||
|
state
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Every ingress branch — control, client-supplied ids, and text through the
|
||||||
|
/// tokenizer pool — must land in `PreSendValidating`, because that is where
|
||||||
|
/// the checks needing the final `input_ids` run. A branch that reached
|
||||||
|
/// `Queued` directly would skip them silently.
|
||||||
|
#[test]
|
||||||
|
fn every_branch_reaches_the_ring_through_pre_send_validating() {
|
||||||
|
for from in [
|
||||||
|
after(
|
||||||
|
RequestState::Validating,
|
||||||
|
Event::Validated(ValidationOutcome::AlreadyTokenized),
|
||||||
|
),
|
||||||
|
after(
|
||||||
|
RequestState::Normalizing,
|
||||||
|
Event::Validated(ValidationOutcome::AlreadyTokenized),
|
||||||
|
),
|
||||||
|
after(RequestState::Tokenizing, Event::TokenizeDone),
|
||||||
|
] {
|
||||||
|
assert!(
|
||||||
|
matches!(from, RequestState::PreSendValidating),
|
||||||
|
"branch bypassed the pre-send checks: {from:?}"
|
||||||
|
);
|
||||||
|
assert!(matches!(
|
||||||
|
after(from, Event::PreSendValidated),
|
||||||
|
RequestState::Queued
|
||||||
|
));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The converse: `Queued` has no other in-edge, so the checks can't be skipped
|
||||||
|
/// by emitting the wrong event, and can't run twice.
|
||||||
|
#[test]
|
||||||
|
fn queued_has_no_other_in_edge() {
|
||||||
|
for mut state in [
|
||||||
|
RequestState::Validating,
|
||||||
|
RequestState::Normalizing,
|
||||||
|
RequestState::Tokenizing,
|
||||||
|
RequestState::Queued,
|
||||||
|
] {
|
||||||
|
assert_eq!(
|
||||||
|
state.apply(Event::PreSendValidated),
|
||||||
|
Err(TransitionError::Illegal),
|
||||||
|
"only PreSendValidating may enter Queued"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -16,6 +16,9 @@ mod error;
|
|||||||
mod fsm;
|
mod fsm;
|
||||||
mod ids;
|
mod ids;
|
||||||
mod message;
|
mod message;
|
||||||
|
mod ring;
|
||||||
|
mod runtime;
|
||||||
|
mod tokenizer_manager;
|
||||||
mod utils;
|
mod utils;
|
||||||
|
|
||||||
use pyo3::prelude::*;
|
use pyo3::prelude::*;
|
||||||
|
|||||||
@@ -15,6 +15,65 @@ mod request;
|
|||||||
mod sampling;
|
mod sampling;
|
||||||
mod types;
|
mod types;
|
||||||
|
|
||||||
pub(crate) use request::GenerateRequest;
|
pub use egress::{ChunkEvent, EgressSink};
|
||||||
|
pub use request::{GenerateRequest, RequestKind};
|
||||||
pub(crate) use sampling::{SamplingParams, SamplingParamsInput};
|
pub(crate) use sampling::{SamplingParams, SamplingParamsInput};
|
||||||
pub(crate) use types::{OneOrMany, OneOrManyItem, TokenIds};
|
pub(crate) use types::{OneOrMany, OneOrManyItem, TokenIds};
|
||||||
|
|
||||||
|
use bytes::Bytes;
|
||||||
|
|
||||||
|
use crate::fsm::RequestState;
|
||||||
|
use crate::ids::Rid;
|
||||||
|
|
||||||
|
/// The owned request as it travels ingress stages (single owner, so `state` is
|
||||||
|
/// mutated lock-free). Common fields here; variant data in [`RequestKind`].
|
||||||
|
#[derive(Debug)]
|
||||||
|
pub struct Request {
|
||||||
|
/// Client-visible request id (uuid hex) — what the scheduler wire and
|
||||||
|
/// `meta_info.id` carry.
|
||||||
|
pub rid: Rid,
|
||||||
|
pub state: RequestState,
|
||||||
|
/// Back-channel to the client connection for egress frames.
|
||||||
|
pub sink: EgressSink,
|
||||||
|
/// Discriminant + variant body (generate vs control).
|
||||||
|
pub kind: RequestKind,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// One ingress-ring entry, split columnar: the scalar `header` (msgpack, `input_ids`
|
||||||
|
/// omitted) + the raw int64 `ids` cell, so the big tensor never goes through msgpack.
|
||||||
|
#[derive(Debug)]
|
||||||
|
pub struct IngressMsg {
|
||||||
|
pub header: Bytes,
|
||||||
|
pub ids: Bytes,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Messages to a Detokenizer shard. `Register` carries the per-request sink for
|
||||||
|
/// the shard's local `rid -> sink` map. The rid STRING is the identity: `Rid::hash`
|
||||||
|
/// picks the shard (collisions there merely co-locate, which is harmless), but two
|
||||||
|
/// distinct rids that hash alike must not be the same map entry — that evicted one
|
||||||
|
/// client's sink and delivered their tokens to the other's connection. Equal rids
|
||||||
|
/// cannot reach here from different requests: `Rid::from_client` uniquifies every
|
||||||
|
/// client-supplied one.
|
||||||
|
pub enum DetokMsg {
|
||||||
|
Register {
|
||||||
|
/// Client-visible rid string — kept in `DetokState` so the shard can
|
||||||
|
/// emit `TmEvent::Abort(rid)` (the wire needs the string, not the hash).
|
||||||
|
rid: Rid,
|
||||||
|
sink: EgressSink,
|
||||||
|
/// Decode logprob token ids to text here (CPU-bound) not on the api threads.
|
||||||
|
decode_logprob_text: bool,
|
||||||
|
/// `SamplingParams.no_stop_trim`: keep the matched stop; default trims it.
|
||||||
|
no_stop_trim: bool,
|
||||||
|
},
|
||||||
|
/// One decode step's chunks for *this shard*. Batched because `tm-egress` blocks
|
||||||
|
/// per send, so one message per request cost ~1.3 µs × batch (5.1x at 4096).
|
||||||
|
Chunks(Vec<ChunkEvent>),
|
||||||
|
/// Control result: one already-serialized payload delivered to the sink verbatim.
|
||||||
|
Result { rid: Rid, payload: bytes::Bytes },
|
||||||
|
/// Terminal per-request failure → an `Error` to the sink (a 400, not a crash).
|
||||||
|
Fail { rid: Rid, message: String },
|
||||||
|
/// Drop the `rid -> sink` entry for a request rejected before the scheduler
|
||||||
|
/// (the rejecting stage already answered the client); else `Register` leaks one
|
||||||
|
/// entry.
|
||||||
|
Deregister { rid: Rid },
|
||||||
|
}
|
||||||
|
|||||||
@@ -486,6 +486,10 @@ fn fan_out<T: OneOrManyItem + Clone + HeapBytes>(
|
|||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
|
|
||||||
|
/// Vocab size for tests that aren't about the vocab bound (see
|
||||||
|
/// `sampling::tests::TEST_VOCAB`).
|
||||||
|
const TEST_VOCAB: u64 = 1000;
|
||||||
|
|
||||||
fn requests(body: &str) -> Result<(Vec<GenerateRequest>, bool), Error> {
|
fn requests(body: &str) -> Result<(Vec<GenerateRequest>, bool), Error> {
|
||||||
serde_json::from_str::<GenerateBody>(body)
|
serde_json::from_str::<GenerateBody>(body)
|
||||||
.unwrap()
|
.unwrap()
|
||||||
@@ -558,7 +562,7 @@ mod tests {
|
|||||||
// Parallel sampling is rejected where Python reads it — in the params,
|
// Parallel sampling is rejected where Python reads it — in the params,
|
||||||
// at normalization (the ingress step), not here.
|
// at normalization (the ingress step), not here.
|
||||||
let (mut ps, _) = requests(r#"{"text": "a", "sampling_params": {"n": 2}}"#).unwrap();
|
let (mut ps, _) = requests(r#"{"text": "a", "sampling_params": {"n": 2}}"#).unwrap();
|
||||||
assert!(ps[0].sampling_params.normalize(false, None).is_err());
|
assert!(ps[0].sampling_params.normalize(false, TEST_VOCAB).is_err());
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Unported `GenerateReqInput` fields are IGNORED, not rejected.
|
/// Unported `GenerateReqInput` fields are IGNORED, not rejected.
|
||||||
|
|||||||
@@ -0,0 +1,251 @@
|
|||||||
|
//! The two Rust↔Python boundary queues.
|
||||||
|
//!
|
||||||
|
//! In embedded mode the Rust frontend threads and the Python scheduler loop
|
||||||
|
//! share one process, so these are in-process `flume` channels — literal
|
||||||
|
//! `mpsc`/`mpmc`, no shared memory, no serialization beyond the msgpack bytes
|
||||||
|
//! the payload already is.
|
||||||
|
//!
|
||||||
|
//! GIL note: the Python side only ever calls the *non-blocking* `drain` /
|
||||||
|
//! `try_push` methods while holding the GIL, and the Rust worker threads only
|
||||||
|
//! ever push/drain raw `Bytes` — neither side touches a `PyObject` off-thread,
|
||||||
|
//! so the producer threads never need the GIL.
|
||||||
|
|
||||||
|
use std::sync::Mutex;
|
||||||
|
use std::time::Duration;
|
||||||
|
|
||||||
|
use bytes::Bytes;
|
||||||
|
|
||||||
|
use crate::message::IngressMsg;
|
||||||
|
|
||||||
|
/// Ingress: TokenizerManager → scheduler `recv_requests`.
|
||||||
|
/// Producers are Rust TM workers; the single consumer is the Python thread.
|
||||||
|
/// Carries [`IngressMsg`] (columnar: scalar header + raw int64 ids cell), not a
|
||||||
|
/// single msgpack blob, so the large `input_ids` tensor bypasses msgpack.
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct IngressProducer {
|
||||||
|
tx: flume::Sender<IngressMsg>,
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct IngressConsumer {
|
||||||
|
rx: flume::Receiver<IngressMsg>,
|
||||||
|
/// One-slot buffer holding a message consumed by a blocking [`wait`] so the
|
||||||
|
/// scheduler can park on idle without losing it — the next [`drain`] returns
|
||||||
|
/// it first. Only ever touched by the single consumer (the Python thread),
|
||||||
|
/// so contention is nil; the `Mutex` is just for interior mutability across
|
||||||
|
/// the `&self` methods.
|
||||||
|
///
|
||||||
|
/// [`wait`]: IngressConsumer::wait
|
||||||
|
/// [`drain`]: IngressConsumer::drain
|
||||||
|
stash: Mutex<Option<IngressMsg>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A drained ingress batch in **columnar** (struct-of-arrays) form. The `ids`
|
||||||
|
/// cells are kept *un-concatenated* so the pyo3 boundary can copy them straight
|
||||||
|
/// into one `PyBytes` (no intermediate buffer); `ids_total` is their summed
|
||||||
|
/// length, precomputed for that single allocation.
|
||||||
|
#[derive(Default)]
|
||||||
|
pub struct IngressColumns {
|
||||||
|
/// Per-request scalar msgpack header (`input_ids` omitted).
|
||||||
|
pub headers: Vec<Bytes>,
|
||||||
|
/// Per-request raw little-endian int64 ids cell (empty for control reqs).
|
||||||
|
pub ids: Vec<Bytes>,
|
||||||
|
/// Per-request token count (`ids` cell length / 8).
|
||||||
|
pub lengths: Vec<u32>,
|
||||||
|
/// Sum of all `ids` cell byte lengths.
|
||||||
|
pub ids_total: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl IngressProducer {
|
||||||
|
/// Non-blocking push. Returns `false` on a full ring (backpressure) so the
|
||||||
|
/// caller can fail the request rather than block a worker thread.
|
||||||
|
#[inline]
|
||||||
|
pub fn try_push(&self, msg: IngressMsg) -> bool {
|
||||||
|
self.tx.try_send(msg).is_ok()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl IngressConsumer {
|
||||||
|
/// Drain up to `max` messages into a columnar [`IngressColumns`], returning
|
||||||
|
/// immediately when the ring runs dry — mirrors the scheduler's existing
|
||||||
|
/// `zmq.NOBLOCK` loop in `request_receiver._pull_raw_reqs`. Splitting headers
|
||||||
|
/// from ids here (off the GIL) leaves `recv_requests` a thin marshaling shim.
|
||||||
|
///
|
||||||
|
/// Non-blocking by construction: `try_recv` returns `Err(TryRecvError::Empty)`
|
||||||
|
/// instantly when the ring is empty, and `Err(_) => break` exits the loop
|
||||||
|
/// right away.
|
||||||
|
pub fn drain(&self, max: usize) -> IngressColumns {
|
||||||
|
let mut batch = IngressColumns::default();
|
||||||
|
// A message parked by a prior blocking `wait` is delivered first.
|
||||||
|
if let Some(m) = self.stash.lock().unwrap().take() {
|
||||||
|
push_msg(&mut batch, m);
|
||||||
|
}
|
||||||
|
while batch.headers.len() < max {
|
||||||
|
match self.rx.try_recv() {
|
||||||
|
Ok(m) => push_msg(&mut batch, m),
|
||||||
|
Err(_) => break, // Empty or Disconnected -> stop now
|
||||||
|
}
|
||||||
|
}
|
||||||
|
batch
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Park up to `timeout` for at least one incoming message, so the idle
|
||||||
|
/// scheduler loop sleeps instead of spinning at 100% CPU. The message is
|
||||||
|
/// **stashed, not returned** — the next [`drain`](Self::drain) yields it —
|
||||||
|
/// so this composes with the existing non-blocking drain flow. Returns
|
||||||
|
/// whether a message is now available. `flume` wakes the parked thread the
|
||||||
|
/// instant a producer pushes, so this adds no latency to real requests.
|
||||||
|
pub fn wait(&self, timeout: Duration) -> bool {
|
||||||
|
if self.stash.lock().unwrap().is_some() {
|
||||||
|
return true;
|
||||||
|
}
|
||||||
|
match self.rx.recv_timeout(timeout) {
|
||||||
|
Ok(m) => {
|
||||||
|
*self.stash.lock().unwrap() = Some(m);
|
||||||
|
true
|
||||||
|
}
|
||||||
|
Err(_) => false, // Timeout or Disconnected
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Append one drained message's columnar cells to the batch.
|
||||||
|
#[inline]
|
||||||
|
fn push_msg(batch: &mut IngressColumns, m: IngressMsg) {
|
||||||
|
batch.ids_total += m.ids.len();
|
||||||
|
batch.lengths.push((m.ids.len() / 8) as u32); // int64 cell → tokens
|
||||||
|
batch.headers.push(m.header);
|
||||||
|
batch.ids.push(m.ids);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Egress: scheduler output (`push_chunk`) → Rust egress dispatcher.
|
||||||
|
/// The single producer is the Python thread; the consumer is the dispatcher.
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct EgressProducer {
|
||||||
|
tx: flume::Sender<Bytes>,
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct EgressConsumer {
|
||||||
|
rx: flume::Receiver<Bytes>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl EgressProducer {
|
||||||
|
/// Blocking push: parks until the ring has space, so a full ring applies
|
||||||
|
/// backpressure to the scheduler instead of dropping output the scheduler has
|
||||||
|
/// already committed (advanced `send_token_offset` for). The GIL is released
|
||||||
|
/// around the call, so parking here doesn't stall other Python threads.
|
||||||
|
/// `false` only when the consumer is gone (runtime shutdown), where the frame
|
||||||
|
/// is unavoidably lost.
|
||||||
|
pub fn push(&self, msg: Bytes) -> bool {
|
||||||
|
self.tx.send(msg).is_ok()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Non-blocking push, so the pyo3 boundary can try to hand the frame over
|
||||||
|
/// while still holding the GIL and detach only when it would actually park.
|
||||||
|
/// Releasing the GIL is not free: reacquiring it waits out the interpreter's
|
||||||
|
/// switch interval (5 ms by default), which dwarfs the sub-microsecond push
|
||||||
|
/// it was protecting.
|
||||||
|
///
|
||||||
|
/// Hands the frame BACK on a full ring (`Err(Some(msg))`) so the caller can
|
||||||
|
/// retry it under [`push`](Self::push) without rebuilding it. `Err(None)` is
|
||||||
|
/// the consumer being gone (shutdown), where the frame is unavoidably lost.
|
||||||
|
#[inline]
|
||||||
|
pub fn try_push(&self, msg: Bytes) -> Result<(), Option<Bytes>> {
|
||||||
|
match self.tx.try_send(msg) {
|
||||||
|
Ok(()) => Ok(()),
|
||||||
|
Err(flume::TrySendError::Full(msg)) => Err(Some(msg)),
|
||||||
|
Err(flume::TrySendError::Disconnected(_)) => Err(None),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl EgressConsumer {
|
||||||
|
/// The underlying receiver, so the dispatcher can drain it via
|
||||||
|
/// [`tokenizer_manager::recv`](crate::tokenizer_manager::recv) (data + shutdown select).
|
||||||
|
pub fn receiver(&self) -> &flume::Receiver<Bytes> {
|
||||||
|
&self.rx
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Build both halves of a bounded ring.
|
||||||
|
pub fn ingress_ring(cap: usize) -> (IngressProducer, IngressConsumer) {
|
||||||
|
let (tx, rx) = flume::bounded(cap);
|
||||||
|
(
|
||||||
|
IngressProducer { tx },
|
||||||
|
IngressConsumer {
|
||||||
|
rx,
|
||||||
|
stash: Mutex::new(None),
|
||||||
|
},
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn egress_ring(cap: usize) -> (EgressProducer, EgressConsumer) {
|
||||||
|
let (tx, rx) = flume::bounded(cap);
|
||||||
|
(EgressProducer { tx }, EgressConsumer { rx })
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
fn msg(h: &'static [u8]) -> IngressMsg {
|
||||||
|
IngressMsg {
|
||||||
|
header: Bytes::from_static(h),
|
||||||
|
ids: Bytes::new(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// `wait` parks when empty (times out), stashes a pushed message
|
||||||
|
/// non-destructively, and the next `drain` returns it.
|
||||||
|
#[test]
|
||||||
|
fn wait_stashes_then_drain_returns_it() {
|
||||||
|
let (tx, rx) = ingress_ring(8);
|
||||||
|
// Empty ring → times out, nothing stashed.
|
||||||
|
assert!(!rx.wait(Duration::from_millis(1)));
|
||||||
|
// Push one, then wait stashes it (returns true).
|
||||||
|
assert!(tx.try_push(msg(b"a")));
|
||||||
|
assert!(rx.wait(Duration::from_millis(200)));
|
||||||
|
// Idempotent: already stashed, returns true without touching the ring.
|
||||||
|
assert!(rx.wait(Duration::from_millis(1)));
|
||||||
|
// Drain yields the stashed message, then the ring is empty.
|
||||||
|
assert_eq!(rx.drain(16).headers.len(), 1);
|
||||||
|
assert!(rx.drain(16).headers.is_empty());
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A blocked `wait` is woken the instant a producer pushes (no polling).
|
||||||
|
#[test]
|
||||||
|
fn wait_wakes_on_push() {
|
||||||
|
let (tx, rx) = ingress_ring(8);
|
||||||
|
std::thread::spawn(move || {
|
||||||
|
std::thread::sleep(Duration::from_millis(20));
|
||||||
|
let _ = tx.try_push(msg(b"a"));
|
||||||
|
});
|
||||||
|
// Generous timeout, but it should return well before it as soon as the
|
||||||
|
// push lands.
|
||||||
|
assert!(rx.wait(Duration::from_secs(5)));
|
||||||
|
assert_eq!(rx.drain(16).headers.len(), 1);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A full egress ring parks the producer until the consumer drains — the
|
||||||
|
/// committed frame is delivered in order, never dropped.
|
||||||
|
#[test]
|
||||||
|
fn egress_push_blocks_until_drained() {
|
||||||
|
let (tx, rx) = egress_ring(1);
|
||||||
|
assert!(tx.push(Bytes::from_static(b"a"))); // fits; ring now full
|
||||||
|
let t = std::thread::spawn(move || tx.push(Bytes::from_static(b"b")));
|
||||||
|
// The parked push can't have completed while the ring is full.
|
||||||
|
std::thread::sleep(Duration::from_millis(20));
|
||||||
|
// Drain one → frees a slot → the parked push lands.
|
||||||
|
assert_eq!(rx.receiver().recv().unwrap(), Bytes::from_static(b"a"));
|
||||||
|
assert!(t.join().unwrap(), "push should succeed once space frees");
|
||||||
|
assert_eq!(rx.receiver().recv().unwrap(), Bytes::from_static(b"b"));
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A closed ring (consumer gone → shutdown) returns `false` instead of
|
||||||
|
/// parking forever, so a scheduler blocked in `push` unblocks on teardown.
|
||||||
|
#[test]
|
||||||
|
fn egress_push_returns_false_when_closed() {
|
||||||
|
let (tx, rx) = egress_ring(1);
|
||||||
|
drop(rx);
|
||||||
|
assert!(!tx.push(Bytes::from_static(b"x")));
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,18 @@
|
|||||||
|
//! Runtime bootstrap: wires channels, pins CPU-bound pools, starts the tokio
|
||||||
|
//! API server, and returns a handle the Python boundary uses for
|
||||||
|
//! `recv_requests` (ingress drain) and `push_batch` (egress push).
|
||||||
|
//!
|
||||||
|
//! Thread layout:
|
||||||
|
//! * API server — tokio multi-thread runtime (I/O bound), pinned core set A
|
||||||
|
//! * Tokenizer — N pinned OS threads (CPU bound), core set B
|
||||||
|
//! * Detokenizer — M pinned OS threads / shards (CPU bound), core set C
|
||||||
|
//! * TM ingress — 1 thread driving the ingress FSM
|
||||||
|
//! * TM egress — 1 thread draining the egress ring → detok shards
|
||||||
|
//!
|
||||||
|
//! Keeping CPU-bound tokenize/detokenize off the async executor avoids stalling
|
||||||
|
//! axum's worker threads.
|
||||||
|
|
||||||
|
mod config;
|
||||||
|
mod runnable;
|
||||||
|
|
||||||
|
// Re-export so stages keep importing `crate::runtime::Runnable`.
|
||||||
@@ -0,0 +1,203 @@
|
|||||||
|
//! Runtime configuration: the rust-server boot knobs
|
||||||
|
//! ([`RustServerServerArgs`]), the typed view of the scheduler's `server_args`
|
||||||
|
//! dump ([`ServerArgs`] / [`ModelConfig`]), and the [`RuntimeConfig`] pairing
|
||||||
|
//! them for `runtime::start`.
|
||||||
|
|
||||||
|
use std::net::SocketAddr;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
/// Boot knobs specific to the embedded rust server — none of these exist in
|
||||||
|
/// the Python `server_args` dump (see [`ServerArgs`]); they arrive as explicit
|
||||||
|
/// `Server::start` parameters.
|
||||||
|
#[derive(Clone, Debug)]
|
||||||
|
pub struct RustServerServerArgs {
|
||||||
|
pub http_addr: SocketAddr,
|
||||||
|
pub api_worker_num: usize,
|
||||||
|
pub ingress_ring_cap: usize,
|
||||||
|
pub egress_ring_cap: usize,
|
||||||
|
pub channel_cap: usize,
|
||||||
|
/// CPU core ids the pools pin to (e.g. this rank's NUMA-local cores minus
|
||||||
|
/// the scheduler's reserved launch cores). `None` → run unpinned.
|
||||||
|
pub cores: Option<Vec<usize>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for RustServerServerArgs {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self {
|
||||||
|
http_addr: "127.0.0.1:30000".parse().unwrap(),
|
||||||
|
api_worker_num: 2,
|
||||||
|
ingress_ring_cap: 8192,
|
||||||
|
egress_ring_cap: 8192,
|
||||||
|
channel_cap: 8192,
|
||||||
|
cores: None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Clone, Debug)]
|
||||||
|
pub struct RuntimeConfig {
|
||||||
|
/// Rust-server-only boot knobs (listen address, pool/ring sizes, pinning).
|
||||||
|
pub rust_server_args: RustServerServerArgs,
|
||||||
|
/// The scheduler's `server_args` dump (worker counts, tokenizer source,
|
||||||
|
/// config-endpoint metadata). `Arc` so cloning the config (and, downstream,
|
||||||
|
/// each `AppState`) is cheap; immutable after construction.
|
||||||
|
pub server_args: Arc<ServerArgs>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for RuntimeConfig {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self {
|
||||||
|
rust_server_args: RustServerServerArgs::default(),
|
||||||
|
server_args: Arc::new(
|
||||||
|
ServerArgs::from_json("{}").expect("empty server_args blob parses"),
|
||||||
|
),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The scheduler's startup blob (`RustServer._build_server_args`) parsed once into
|
||||||
|
/// typed fields: values are post-`__post_init__`, unknown keys (e.g. `api_key`) are dropped.
|
||||||
|
#[derive(Debug, serde::Deserialize)]
|
||||||
|
pub struct ServerArgs {
|
||||||
|
/// HF repo id / local dir of the model, reported by `/get_model_info`.
|
||||||
|
#[serde(default)]
|
||||||
|
pub model_path: String,
|
||||||
|
/// Model name reported by `/v1/models` and `/server_info`.
|
||||||
|
#[serde(default)]
|
||||||
|
pub served_model_name: String,
|
||||||
|
/// Tokenizer source (model dir / `tokenizer.json` / HF repo id). Empty only
|
||||||
|
/// in minimal standalone blobs — then boot requires `skip_tokenizer_init`.
|
||||||
|
#[serde(default)]
|
||||||
|
pub tokenizer_path: String,
|
||||||
|
/// HF revision, used only when `tokenizer_path` is a repo id. `None` → main.
|
||||||
|
#[serde(default)]
|
||||||
|
pub revision: Option<String>,
|
||||||
|
/// HTTP bind address (see [`Self::bind`]).
|
||||||
|
#[serde(default = "default_host")]
|
||||||
|
pub host: String,
|
||||||
|
#[serde(default = "default_port")]
|
||||||
|
pub port: u16,
|
||||||
|
/// Log levels driving the access log — uvicorn runs at
|
||||||
|
/// `log_level_http or log_level` (see [`Self::http_access_log_enabled`]).
|
||||||
|
#[serde(default = "default_log_level")]
|
||||||
|
pub log_level: String,
|
||||||
|
#[serde(default)]
|
||||||
|
pub log_level_http: Option<String>,
|
||||||
|
/// Pinned tokenizer threads / detok shards (Python asserts both ≥ 1).
|
||||||
|
#[serde(default = "default_worker_num")]
|
||||||
|
pub tokenizer_worker_num: usize,
|
||||||
|
#[serde(default = "default_worker_num")]
|
||||||
|
pub detokenizer_worker_num: usize,
|
||||||
|
/// Token-ids-in / token-ids-out mode: no tokenizer load, raw `output_ids`
|
||||||
|
/// frames (drives the `Skip` detok backend and the ingress branch).
|
||||||
|
#[serde(default)]
|
||||||
|
pub skip_tokenizer_init: bool,
|
||||||
|
/// Streamed `/generate` frames carry per-step deltas instead of cumulative
|
||||||
|
/// text. Matches the Python `TokenizerManager`.
|
||||||
|
#[serde(default)]
|
||||||
|
pub incremental_streaming_output: bool,
|
||||||
|
/// The resolved Python `ModelConfig`, attached to the blob at dump time.
|
||||||
|
#[serde(default)]
|
||||||
|
pub model_config: ModelConfig,
|
||||||
|
/// Default sampling params advertised by `/get_model_info`, verbatim from
|
||||||
|
/// `server_args.preferred_sampling_params` (a JSON object or null).
|
||||||
|
#[serde(default)]
|
||||||
|
pub preferred_sampling_params: Option<serde_json::Value>,
|
||||||
|
/// Over-long inputs are truncated to fit the context instead of 400ing, and
|
||||||
|
/// `max_new_tokens` is clamped rather than rejected (Python
|
||||||
|
/// `TokenizerManager._validate_one_request`).
|
||||||
|
#[serde(default)]
|
||||||
|
pub allow_auto_truncate: bool,
|
||||||
|
/// `return_hidden_states` is refused unless the server was launched with it:
|
||||||
|
/// the scheduler simply won't produce them, so the request would 200 with the
|
||||||
|
/// field silently missing.
|
||||||
|
#[serde(default)]
|
||||||
|
pub enable_return_hidden_states: bool,
|
||||||
|
/// Output slots reserved per request on top of its input (eagle stores draft
|
||||||
|
/// tokens there). Not a `server_args` field — `TokenizerManager` derives it and
|
||||||
|
/// `RustServer._build_server_args` stamps it in, so both sides count alike.
|
||||||
|
#[serde(default)]
|
||||||
|
pub num_reserved_tokens: u64,
|
||||||
|
/// Launch-time stamps (not `server_args` fields): sglang package version
|
||||||
|
/// and the scheduler-derived KV token capacity, reported by `/server_info`.
|
||||||
|
#[serde(default)]
|
||||||
|
pub version: Option<String>,
|
||||||
|
#[serde(default)]
|
||||||
|
pub max_total_num_tokens: Option<u64>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The slice of the resolved Python `ModelConfig` the rust server reads.
|
||||||
|
#[derive(Debug, Default, serde::Deserialize)]
|
||||||
|
pub struct ModelConfig {
|
||||||
|
/// Resolved context length (`max_model_len` in `/v1/models`); mandatory at
|
||||||
|
/// boot ([`ServerArgs::validate_mandatory`]).
|
||||||
|
#[serde(default)]
|
||||||
|
pub context_len: Option<u64>,
|
||||||
|
/// Bounds client-supplied token ids — ingress 400s out-of-vocab ids before
|
||||||
|
/// they crash the scheduler's embedding lookup; mandatory at
|
||||||
|
/// boot ([`ServerArgs::validate_mandatory`]).
|
||||||
|
#[serde(default)]
|
||||||
|
pub vocab_size: Option<u64>,
|
||||||
|
}
|
||||||
|
|
||||||
|
fn default_host() -> String {
|
||||||
|
"127.0.0.1".into()
|
||||||
|
}
|
||||||
|
fn default_port() -> u16 {
|
||||||
|
30000
|
||||||
|
}
|
||||||
|
fn default_log_level() -> String {
|
||||||
|
"info".into()
|
||||||
|
}
|
||||||
|
fn default_worker_num() -> usize {
|
||||||
|
1
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ServerArgs {
|
||||||
|
/// Parse the blob; errors on malformed JSON or a wrongly-typed field.
|
||||||
|
pub fn from_json(s: &str) -> Result<Self, String> {
|
||||||
|
serde_json::from_str(s).map_err(|e| e.to_string())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Fail fast at startup if a field an endpoint depends on is missing.
|
||||||
|
pub fn validate_mandatory(&self) -> Result<(), String> {
|
||||||
|
if self.served_model_name.is_empty() {
|
||||||
|
return Err("no 'served_model_name' in server_args".into());
|
||||||
|
}
|
||||||
|
if self.model_config.context_len.is_none() {
|
||||||
|
return Err("no resolvable context length (model_config.context_len)".into());
|
||||||
|
}
|
||||||
|
if self.model_config.vocab_size.is_none() {
|
||||||
|
return Err("no resolvable vocab size (model_config.vocab_size)".into());
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Bind address `host:port`. `host` is expected to be an IP — the result is
|
||||||
|
/// parsed as a `SocketAddr`.
|
||||||
|
pub fn bind(&self) -> String {
|
||||||
|
format!("{}:{}", self.host, self.port)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Whether the HTTP access log is emitted, mirroring the Python server:
|
||||||
|
/// uvicorn runs at `log_level_http or log_level` and prints access lines
|
||||||
|
/// only at info/debug. `--log-level-http warning` turns them off.
|
||||||
|
pub fn http_access_log_enabled(&self) -> bool {
|
||||||
|
let level = self
|
||||||
|
.log_level_http
|
||||||
|
.as_deref()
|
||||||
|
.filter(|s| !s.is_empty())
|
||||||
|
.unwrap_or(&self.log_level);
|
||||||
|
matches!(
|
||||||
|
level.to_ascii_lowercase().as_str(),
|
||||||
|
"trace" | "debug" | "info"
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Pinned API threads for the embedded HTTP api-server. Python `server_args`
|
||||||
|
/// has no such field — this is derived: enough to cover the widest pool.
|
||||||
|
pub fn api_worker_num(&self) -> usize {
|
||||||
|
4.max(self.tokenizer_worker_num)
|
||||||
|
.max(self.detokenizer_worker_num)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,10 @@
|
|||||||
|
//! The [`Runnable`] stage trait — the one contract every pipeline stage
|
||||||
|
//! (CPU-bound worker or TM router) implements to be spawned by the runtime.
|
||||||
|
|
||||||
|
/// A pipeline stage that owns its channel handles + config and runs a blocking
|
||||||
|
/// loop until its inbox closes. Lets the runtime spawn stages uniformly via
|
||||||
|
/// `threads::spawn_stage` / `threads::spawn_pool` instead of free `run_*` functions with
|
||||||
|
/// positional handles. Implemented by every CPU-bound worker and TM router.
|
||||||
|
pub trait Runnable: Send + 'static {
|
||||||
|
fn run(self);
|
||||||
|
}
|
||||||
@@ -0,0 +1,87 @@
|
|||||||
|
//! TokenizerManager — owns the request lifecycle across two isolated threads:
|
||||||
|
//!
|
||||||
|
//! * [`ingress`] — drives the ingress FSM (Received → Validating →
|
||||||
|
//! Normalizing → {Tokenizing | PreSendValidating}) and pushes tokenized requests to the
|
||||||
|
//! scheduler ring.
|
||||||
|
//! * [`egress`] — drains the scheduler-output ring and routes each chunk to
|
||||||
|
//! the owning detokenizer shard.
|
||||||
|
//!
|
||||||
|
//! The two run on separate pinned threads with no shared state, connected to
|
||||||
|
//! the rest of the pipeline only through `flume` channels: [`TmEvent`] into the
|
||||||
|
//! ingress loop, [`Senders`] fanning out to the pools.
|
||||||
|
|
||||||
|
use crate::ids::Rid;
|
||||||
|
use crate::message::{DetokMsg, Request};
|
||||||
|
|
||||||
|
/// Blocking receive that also wakes on shutdown: returns `None` when `rx` closes
|
||||||
|
/// *or* the `shutdown` sender is dropped.
|
||||||
|
pub fn recv<T>(rx: &flume::Receiver<T>, shutdown: &flume::Receiver<()>) -> Option<T> {
|
||||||
|
flume::Selector::new()
|
||||||
|
.recv(rx, |r| r.ok())
|
||||||
|
.recv(shutdown, |_| None)
|
||||||
|
.wait()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Events into the TokenizerManager ingress loop. API server + tokenizer pool
|
||||||
|
/// share this one inbox, keeping the loop a single consumer (no `select`).
|
||||||
|
pub enum TmEvent {
|
||||||
|
/// A freshly received request from the API server.
|
||||||
|
Ingress(Request),
|
||||||
|
/// A request back from the tokenizer pool: `PreSendValidating` (ids filled) on success,
|
||||||
|
/// or `Failed` on a tokenize error. `drive` handles both.
|
||||||
|
Tokenized(Request),
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Producer-side handles, cloned into every stage that needs to emit.
|
||||||
|
/// Who asked for an abort. Both variants do the same work in
|
||||||
|
/// [`Ingress::on_abort`](crate::tokenizer_manager::ingress::Ingress) — deregister
|
||||||
|
/// the detok entry, tell the scheduler to stop — and the source is kept for
|
||||||
|
/// diagnostics.
|
||||||
|
///
|
||||||
|
/// There is no in-flight rid registry to keep consistent, and so no release
|
||||||
|
/// ordering to get wrong: [`Rid::from_client`] makes every client-supplied rid
|
||||||
|
/// internally unique, so a resubmit of the "same" rid is a different `Rid` and
|
||||||
|
/// cannot be tangled up with an abort still in flight for the original.
|
||||||
|
#[derive(Clone, Debug)]
|
||||||
|
pub enum AbortSource {
|
||||||
|
/// From an `AbortGuard` drop. Owns the release.
|
||||||
|
Guard(Rid),
|
||||||
|
/// From a detokenizer terminal path. Aborts the scheduler work.
|
||||||
|
Detok(Rid),
|
||||||
|
}
|
||||||
|
|
||||||
|
impl AbortSource {
|
||||||
|
pub fn rid(&self) -> &Rid {
|
||||||
|
match self {
|
||||||
|
Self::Guard(rid) | Self::Detok(rid) => rid,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct Senders {
|
||||||
|
/// → TokenizerManager ingress loop.
|
||||||
|
pub tm: flume::Sender<TmEvent>,
|
||||||
|
/// → the same loop, but UNBOUNDED and abort-only.
|
||||||
|
///
|
||||||
|
/// Aborts cannot share the bounded inbox. `try_send` there drops them exactly
|
||||||
|
/// when they matter most — under overload — leaving the scheduler generating
|
||||||
|
/// for a dead connection; and the caller then faces a false choice between
|
||||||
|
/// releasing the rid (a live entry can be overwritten by a resubmit) and
|
||||||
|
/// holding it (a permanent leak). An unbounded lane removes the dilemma: an
|
||||||
|
/// abort is a small `String` and is always accepted, so releases can be
|
||||||
|
/// unconditional again. It cannot grow without bound in practice — one entry
|
||||||
|
/// per in-flight request, each already bounded by the inbox that admitted it.
|
||||||
|
pub abort: flume::Sender<AbortSource>,
|
||||||
|
/// → Tokenizer pool (CPU-bound, pinned threads).
|
||||||
|
pub tok: flume::Sender<Request>,
|
||||||
|
/// → Detokenizer shards, indexed by `Rid::shard(detok.len())`.
|
||||||
|
pub detok: Vec<flume::Sender<DetokMsg>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Senders {
|
||||||
|
#[inline]
|
||||||
|
pub fn detok_for(&self, rid: &Rid) -> &flume::Sender<DetokMsg> {
|
||||||
|
&self.detok[rid.shard(self.detok.len())]
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user