diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs index 142f10ee0..61e0f9157 100644 --- a/bindings/python/src/lib.rs +++ b/bindings/python/src/lib.rs @@ -753,6 +753,7 @@ impl Router { match self.backend { BackendType::Vllm => Some(worker::RuntimeType::Vllm), BackendType::Tokenspeed => Some(worker::RuntimeType::TokenSpeed), + BackendType::Sglang => Some(worker::RuntimeType::Sglang), _ => None, } } else { diff --git a/bindings/python/src/smg/_sglang_zmq_launcher.py b/bindings/python/src/smg/_sglang_zmq_launcher.py new file mode 100644 index 000000000..f61c0602e --- /dev/null +++ b/bindings/python/src/smg/_sglang_zmq_launcher.py @@ -0,0 +1,121 @@ +"""Headless SGLang scheduler launcher for the SMG direct-ZMQ backend. + +On the direct-ZMQ path SMG plays the SGLang TokenizerManager role itself: it +binds the PUSH input and PULL output ``ipc://`` sockets, sends tokenized +generate requests, and receives batched token outputs straight off the wire. +This module therefore launches *only* the scheduler process(es) — no +TokenizerManager, no DetokenizerManager — wired to SMG's two sockets through a +hand-built ``PortArgs``, with ``skip_tokenizer_init`` so the scheduler runs +tokenizer-free and routes its outputs back over the tokenizer socket (which SMG +owns). + +Unmodified SGLang exposes no CLI flag or env var to point a bare scheduler at +externally-chosen sockets, so SMG drives SGLang's own ``PortArgs`` / +``Engine._launch_scheduler_processes`` primitives directly. Every argument other +than the two SMG-owned endpoints (``--smg-input-ipc`` / ``--smg-output-ipc``) is +a native SGLang server argument forwarded verbatim. + +The scheduler defaults to pickle over ZMQ; SMG only decodes msgpack, so +``SGLANG_USE_PICKLE_IPC`` is forced off before SGLang is imported (the flag is +snapshotted at import time). + +Invoked as a subprocess by the ``sglang`` launcher in ``serve.py``. +""" + +import os + +# Must precede any sglang import: io_struct snapshots this at module load to +# choose msgpack vs pickle for the ZMQ wire, and SMG only speaks msgpack. Hard +# override (not setdefault): an inherited "1" would silently pickle the wire. +os.environ["SGLANG_USE_PICKLE_IPC"] = "0" + +import logging # noqa: E402 +import sys # noqa: E402 + +logger = logging.getLogger("smg.sglang_zmq_launcher") + +_INPUT_FLAG = "--smg-input-ipc" +_OUTPUT_FLAG = "--smg-output-ipc" + + +def _extract_flag(argv: list[str], flag: str) -> tuple[str, list[str]]: + """Pop ``flag VALUE`` or ``flag=VALUE`` from argv; return (value, rest). + + These are SMG-owned endpoints, not SGLang args, so they are stripped before + the remainder is handed to SGLang's own parser. + """ + for i, token in enumerate(argv): + if token == flag: + if i + 1 >= len(argv): + raise SystemExit(f"{flag} requires a value") + return argv[i + 1], argv[:i] + argv[i + 2 :] + if token.startswith(flag + "="): + return token.split("=", 1)[1], argv[:i] + argv[i + 1 :] + raise SystemExit(f"{flag} is required") + + +def main(argv: list[str] | None = None) -> None: + argv = list(sys.argv[1:] if argv is None else argv) + # scheduler_input_ipc_name: SMG PUSHes requests here, the scheduler PULLs. + input_ipc, argv = _extract_flag(argv, _INPUT_FLAG) + # tokenizer_ipc_name: under skip_tokenizer_init the scheduler PUSHes outputs + # here, and SMG PULLs them. + output_ipc, argv = _extract_flag(argv, _OUTPUT_FLAG) + + from sglang.srt.entrypoints.engine import ( + Engine, + _set_envs_and_config, + _set_gc, + ) + from sglang.srt.managers.scheduler import run_scheduler_process + from sglang.srt.plugins import load_plugins + from sglang.srt.server_args import PortArgs, prepare_server_args + from sglang.srt.utils import configure_logger + + # SMG tokenizes upstream and drives the tokenizer<->scheduler ZMQ wire + # directly, so the scheduler must run without its own tokenizer. Pass this as + # a CLI flag: SGLang resolves and freezes server_args in prepare_server_args, + # so it can no longer be set afterward. + if "--skip-tokenizer-init" not in argv: + argv.append("--skip-tokenizer-init") + server_args = prepare_server_args(argv) + + # Mirror the head of Engine._launch_subprocesses so the scheduler subprocess + # inherits the same logging, env/config, plugins, and mp start method it + # would under a normal SGLang launch. + configure_logger(server_args) + _set_envs_and_config(server_args) + load_plugins() + server_args.check_server_args() + _set_gc(server_args) + + # Reuse SGLang's own port derivation (nccl port, instance id, the unused + # detokenizer/rpc/metrics ipc names), then repoint the two SMG-owned sockets + # at the endpoints SMG binds. + port_args = PortArgs.init_new(server_args) + port_args.scheduler_input_ipc_name = input_ipc + port_args.tokenizer_ipc_name = output_ipc + + result, procs = Engine._launch_scheduler_processes( + server_args, port_args, run_scheduler_process + ) + # Block until every scheduler rank has finished loading; the SMG router + # gates request admission on its own ZMQ readiness, so once the schedulers + # are up this process just keeps them alive. + result.wait_for_ready() + logger.info( + "SGLang scheduler(s) ready on input=%s output=%s; SMG owns the tokenizer ZMQ wire", + input_ipc, + output_ipc, + ) + result.block_until_scheduler_exits() + + # A scheduler exiting means the engine is gone; propagate a non-zero status + # (rather than a silent exit 0) so the parent launcher sees the failure + # instead of leaving the router pushing to a dead socket. + if any(proc.exitcode for proc in procs): + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/bindings/python/src/smg/serve.py b/bindings/python/src/smg/serve.py index e202de8f5..7c879a3a6 100644 --- a/bindings/python/src/smg/serve.py +++ b/bindings/python/src/smg/serve.py @@ -37,6 +37,17 @@ def _zmq_ipc_url(port: int) -> str: return f"ipc://{_ZMQ_SOCKET_DIR}/engine-{port}" +def _zmq_data_addresses(ipc_url: str) -> tuple[str, str]: + """The ipc:// data-plane sockets SMG binds for a worker: (input, output). + + Mirrors `zmq_socket_addresses` in model_gateway/src/worker/worker.rs: SMG + binds `-in.sock` for requests and `-out.sock` for outputs. An + engine that connects directly (no TCP handshake, e.g. SGLang) must dial these + exact endpoints. + """ + return f"{ipc_url}-in.sock", f"{ipc_url}-out.sock" + + def _zmq_handshake_port(ipc_url: str) -> int: """The tcp handshake port SMG derives from an ipc:// URL. @@ -189,11 +200,23 @@ class SglangWorkerLauncher(WorkerLauncher): """Launcher for sglang inference workers.""" def _get_tp_size(self, args: argparse.Namespace) -> int: - return getattr(args, "tensor_parallel_size", 1) + # sglang's ServerArgs registers --tp-size with dest `tp_size` (the alias + # --tensor-parallel-size shares that dest), so read tp_size first, then + # fall back to the alias and finally a single-GPU default. The default + # applies only when neither attribute is set; an explicit value (even a + # non-positive one) is preserved so gpu_env can reject it. + if hasattr(args, "tp_size"): + return int(args.tp_size) + if hasattr(args, "tensor_parallel_size"): + return int(args.tensor_parallel_size) + return 1 def build_command( self, args: argparse.Namespace, backend_args: list[str], host: str, port: int ) -> list[str]: + if getattr(args, "connection_mode", "grpc") == "zmq": + return self._build_zmq_command(args, backend_args, port) + cmd = [ sys.executable, "-m", @@ -216,6 +239,37 @@ def build_command( cmd.extend(self._filter_backend_args(backend_args, ["--model-path", "--host", "--port"])) return cmd + def _build_zmq_command( + self, args: argparse.Namespace, backend_args: list[str], port: int + ) -> list[str]: + """Launch a headless SGLang scheduler wired to SMG's ZMQ sockets. + + Unmodified SGLang exposes no CLI to point a bare scheduler at + externally-chosen sockets, so this dispatches the `_sglang_zmq_launcher` + module, which drives SGLang's own scheduler-launch primitives on the two + ipc endpoints SMG binds. Each worker is one standalone scheduler; running + several is dense data parallelism as N independent ZMQ workers. + """ + input_ipc, output_ipc = _zmq_data_addresses(_zmq_ipc_url(port)) + cmd = [ + sys.executable, + "-m", + "smg._sglang_zmq_launcher", + "--smg-input-ipc", + input_ipc, + "--smg-output-ipc", + output_ipc, + "--model-path", + getattr(args, "model_path", ""), + ] + cmd.extend( + self._filter_backend_args( + backend_args, + ["--smg-input-ipc", "--smg-output-ipc", "--model-path", "--host", "--port"], + ) + ) + return cmd + class VllmWorkerLauncher(WorkerLauncher): """Launcher for vLLM inference workers.""" @@ -654,7 +708,7 @@ def add_serve_args(parser: argparse.ArgumentParser) -> None: help=( "Connection mode for workers (default: grpc). Note: trtllm only " "supports grpc, tokenspeed only supports zmq, and zmq is otherwise " - "only supported for the vllm backend" + "only supported for the vllm and sglang backends" ), ) # Router host/port - may be overridden by backend (e.g. sglang) @@ -735,10 +789,14 @@ def parse_serve_args( # ZMQ direct-backend is a same-host engine connection; only vLLM EngineCore # and TokenSpeed speak a supported ZMQ wire protocol. - if serve_router_args.connection_mode == "zmq" and backend not in ("vllm", "tokenspeed"): + if serve_router_args.connection_mode == "zmq" and backend not in ( + "vllm", + "tokenspeed", + "sglang", + ): pre_parser.error( - "connection-mode zmq is only supported for the vllm and tokenspeed " - f"backends, not {backend}" + "connection-mode zmq is only supported for the vllm, tokenspeed, and " + f"sglang backends, not {backend}" ) # Pass 2: full parser with backend-specific args; resolve so backend can override @@ -760,6 +818,12 @@ def parse_serve_args( else: args = parser.parse_args(argv) + # Some backends declare data parallelism under their own dest (sglang uses + # dp_size), which resolves over the serve-level --data-parallel-size. Give + # the orchestrator one canonical field regardless of backend. + if not hasattr(args, "data_parallel_size"): + args.data_parallel_size = getattr(args, "dp_size", 1) + return backend, args, backend_args diff --git a/crates/engine_zmq_client/examples/live_probe.rs b/crates/engine_zmq_client/examples/live_probe.rs index 9dcbca45f..768ef552b 100644 --- a/crates/engine_zmq_client/examples/live_probe.rs +++ b/crates/engine_zmq_client/examples/live_probe.rs @@ -65,13 +65,17 @@ async fn main() { .expect("handshake with engine"); let engine = &transport.engines[0]; + let ready = engine + .ready_response + .as_ref() + .expect("handshake engine registers a ready response"); eprintln!( "[probe] engine registered: vllm_version={} max_model_len={} num_gpu_blocks={} block_size={} dtype={} engine_index={:?}", - engine.ready_response.vllm_version, - engine.ready_response.max_model_len, - engine.ready_response.num_gpu_blocks, - engine.ready_response.block_size, - engine.ready_response.dtype.as_str(), + ready.vllm_version, + ready.max_model_len, + ready.num_gpu_blocks, + ready.block_size, + ready.dtype.as_str(), engine.engine_id.engine_index(), ); diff --git a/crates/engine_zmq_client/src/connector.rs b/crates/engine_zmq_client/src/connector.rs index 876ffc6cf..4735459ef 100644 --- a/crates/engine_zmq_client/src/connector.rs +++ b/crates/engine_zmq_client/src/connector.rs @@ -17,15 +17,16 @@ use futures::Stream; use parking_lot::Mutex; use tokio::{sync::mpsc, task::JoinHandle}; use tracing::{trace, warn}; -use zeromq::RouterSendHalf; use crate::{ error::{Error, Result}, protocol::{ - tokenspeed::TokenSpeedProtocol, vllm::VllmProtocol, EngineBatch, EngineLoad, EngineOutput, - EngineProtocol, + sglang::SglangProtocol, tokenspeed::TokenSpeedProtocol, vllm::VllmProtocol, EngineBatch, + EngineLoad, EngineOutput, EngineProtocol, + }, + transport::{ + run_output_loop, send_message, ConnectedEngine, ConnectedTransport, EngineId, InputSocket, }, - transport::{run_output_loop, send_message, ConnectedEngine, ConnectedTransport, EngineId}, }; /// The vLLM EngineCore connector (the original engine surface). @@ -36,6 +37,10 @@ pub type EngineCoreStream = RequestStream; pub type TokenSpeedClient = Client; /// The per-request output stream for the TokenSpeed connector. pub type TokenSpeedStream = RequestStream; +/// The SGLang connector (no-handshake PUSH/PULL, tag-dispatched). +pub type SglangClient = Client; +/// The per-request output stream for the SGLang connector. +pub type SglangStream = RequestStream; type OutputSender = mpsc::UnboundedSender>; type OutputReceiver = mpsc::UnboundedReceiver>; @@ -103,8 +108,8 @@ impl RequestRegistry { } struct ClientInner { - /// Shared input ROUTER send half (serialized across concurrent submits). - input_send: tokio::sync::Mutex, + /// Shared input socket, ROUTER or PUSH (serialized across concurrent submits). + input_send: tokio::sync::Mutex, engines: Vec, registry: Mutex>, /// Latest per-rank load, keyed by engine index. SMG's DP load signal. @@ -144,7 +149,7 @@ impl ClientInner

{ async fn send_to_engine( &self, engine_id: &EngineId, - request_type: Bytes, + request_type: Option, payload: Vec, aux_frames: Vec, ) -> Result<()> { @@ -711,4 +716,113 @@ mod tests { // Terminal output ends the stream. assert!(stream.next().await.is_none()); } + + /// The generic connector drives the SGLang protocol over the no-handshake + /// PUSH/PULL transport: no identity or request-type frame, a tag-dispatched + /// `TokenizedGenerateReqInput` in, and `BatchTokenIDOutput` batches back. The + /// engine replies with the pinned Python wire vectors, exercising the real + /// bytes end to end (the output struct is decode-only, so it cannot be + /// re-encoded in-process). + #[tokio::test] + async fn sglang_client_submits_and_streams() { + use crate::{ + mock_engine::MockPushPullEngine, + protocol::sglang::{ + request::{AbortReq, TokenizedGenerateReqInput}, + sampling::SamplingParams, + token_ids::TokenIdArray, + }, + transport::connect_push_pull, + }; + + // Pinned Python `BatchTokenIDOutput` vectors for rid "req-00000001", + // output_ids [[15496]] — still-generating (finish nil) then finished + // (`{"type":"stop"}`). Captured in `protocol::sglang::output` tests. + const OUTPUT_STILL_GENERATING: &str = + "dc0032b24261746368546f6b656e49444f757470757491ac7265712d3030303030303031c091\ + c091a09192a171c408883c00000000000091009192a171c408883c00000000000091c391c391\ + c29103910091019100c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0\ + c0c0c0c0c0c0"; + const OUTPUT_FINISHED: &str = + "dc0032b24261746368546f6b656e49444f757470757491ac7265712d3030303030303031c091\ + 82a474797065a473746f70a76d6174636865640291a09192a171c408883c0000000000009100\ + 9192a171c408883c00000000000091c391c391c29103910091019100c0c0c0c0c0c0c0c0c0c0\ + c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0"; + + fn from_hex(hex: &str) -> Bytes { + let hex: String = hex.chars().filter(|c| !c.is_whitespace()).collect(); + Bytes::from( + (0..hex.len()) + .step_by(2) + .map(|i| u8::from_str_radix(&hex[i..i + 2], 16).unwrap()) + .collect::>(), + ) + } + + let ns = IpcNamespace::new().unwrap(); + let (input, output) = (ns.input_endpoint(), ns.output_endpoint()); + std::mem::forget(ns); + + // Frontend binds PUSH input + PULL output; the mock scheduler connects in. + let (transport, engine) = tokio::join!( + connect_push_pull("127.0.0.1", Some(&input), Some(&output), TIMEOUT), + MockPushPullEngine::connect(&input, &output), + ); + let client = SglangClient::new(transport.unwrap()); + let mut engine = engine.unwrap(); + + let request = TokenizedGenerateReqInput { + rid: "req-00000001".to_string(), + input_ids: TokenIdArray(vec![9906, 11, 1917]), + sampling_params: SamplingParams { + max_new_tokens: Some(4), + ..SamplingParams::default() + }, + stream: true, + ..TokenizedGenerateReqInput::default() + }; + let mut stream = client.submit(request).await.unwrap(); + + // The scheduler sees a single-frame payload: no identity, no type frame. + let frames = engine.recv_request().await.unwrap(); + assert_eq!(frames.len(), 1); + let received: TokenizedGenerateReqInput = decode_msgpack(frames[0].as_ref()).unwrap(); + assert_eq!(received.rid, "req-00000001"); + assert_eq!(received.input_ids, TokenIdArray(vec![9906, 11, 1917])); + + engine + .send_output(vec![from_hex(OUTPUT_STILL_GENERATING)]) + .await + .unwrap(); + engine + .send_output(vec![from_hex(OUTPUT_FINISHED)]) + .await + .unwrap(); + + let first = stream.next().await.unwrap().unwrap(); + assert_eq!(first.output_ids, vec![15496]); + assert!(!first.finished()); + let second = stream.next().await.unwrap().unwrap(); + assert_eq!(second.finish_reason.as_deref(), Some("stop")); + assert!(second.finished()); + // Terminal output ends the stream. + assert!(stream.next().await.is_none()); + + // A fresh request dropped before finishing auto-aborts with a tagged + // AbortReq (still no identity/type frame on the PUSH input). + let request = TokenizedGenerateReqInput { + rid: "req-00000002".to_string(), + input_ids: TokenIdArray(vec![1]), + stream: true, + ..TokenizedGenerateReqInput::default() + }; + let stream = client.submit(request).await.unwrap(); + let add = engine.recv_request().await.unwrap(); + assert_eq!(add.len(), 1); + drop(stream); + let abort = engine.recv_request().await.unwrap(); + assert_eq!(abort.len(), 1); + let decoded: AbortReq = decode_msgpack(abort[0].as_ref()).unwrap(); + assert_eq!(decoded, AbortReq::new("req-00000002")); + } } diff --git a/crates/engine_zmq_client/src/lib.rs b/crates/engine_zmq_client/src/lib.rs index c8b9ec3fb..2dc1c77b0 100644 --- a/crates/engine_zmq_client/src/lib.rs +++ b/crates/engine_zmq_client/src/lib.rs @@ -36,11 +36,12 @@ pub mod transport; pub mod mock_engine; pub use connector::{ - Client, EngineCoreClient, EngineCoreStream, RequestStream, TokenSpeedClient, TokenSpeedStream, + Client, EngineCoreClient, EngineCoreStream, RequestStream, SglangClient, SglangStream, + TokenSpeedClient, TokenSpeedStream, }; pub use error::{Error, Result}; pub use protocol::{EngineBatch, EngineOutput, EngineProtocol}; pub use transport::{ - connect_handshake, run_output_loop, send_message, ConnectedEngine, ConnectedTransport, - EngineId, ENGINE_CORE_DEAD_SENTINEL, + connect_handshake, connect_push_pull, run_output_loop, send_message, ConnectedEngine, + ConnectedTransport, EngineId, InputSocket, ENGINE_CORE_DEAD_SENTINEL, }; diff --git a/crates/engine_zmq_client/src/mock_engine.rs b/crates/engine_zmq_client/src/mock_engine.rs index 1ea597a3d..2378ab283 100644 --- a/crates/engine_zmq_client/src/mock_engine.rs +++ b/crates/engine_zmq_client/src/mock_engine.rs @@ -15,7 +15,7 @@ use tokio::time::{sleep, timeout}; use zeromq::{ prelude::{Socket, SocketRecv, SocketSend}, util::PeerIdentity, - DealerSocket, PushSocket, SocketOptions, ZmqMessage, + DealerSocket, PullSocket, PushSocket, SocketOptions, ZmqMessage, }; use crate::{ @@ -190,6 +190,53 @@ impl MockEngine { } } +/// A mock engine over the no-handshake PUSH/PULL topology (the SGLang scheduler +/// role): it CONNECTS a PULL to the frontend's bound PUSH input to receive +/// requests, and a PUSH to the frontend's bound PULL output to send outputs. +/// There is no handshake and no identity frame. +pub struct MockPushPullEngine { + input: PullSocket, + output: PushSocket, +} + +impl MockPushPullEngine { + /// Connect to a frontend that has already bound its PUSH input and PULL + /// output sockets. + pub async fn connect(input_address: &str, output_address: &str) -> Result { + wait_for_endpoint(input_address).await?; + let mut input = PullSocket::new(); + input.connect(input_address).await?; + + wait_for_endpoint(output_address).await?; + let mut output = PushSocket::new(); + output.connect(output_address).await?; + + Ok(Self { input, output }) + } + + /// Receive one request's raw frames (`[payload, aux..]`; no identity or type + /// frame on this topology). + pub async fn recv_request(&mut self) -> Result> { + Ok(self.input.recv().await?.into_vec()) + } + + /// Push one raw multi-frame output message back to the frontend. + pub async fn send_output(&mut self, frames: Vec) -> Result<()> { + let mut iter = frames.into_iter(); + let Some(first) = iter.next() else { + return Err(Error::UnexpectedHandshakeMessage { + message: "mock engine output needs at least one frame".to_string(), + }); + }; + let mut message = ZmqMessage::from(first); + for frame in iter { + message.push_back(frame); + } + self.output.send(message).await?; + Ok(()) + } +} + fn ready_message(status: &str) -> ReadyMessage { ReadyMessage { status: Some(status.to_string()), diff --git a/crates/engine_zmq_client/src/protocol/mod.rs b/crates/engine_zmq_client/src/protocol/mod.rs index ecf3f6640..3216bd7fb 100644 --- a/crates/engine_zmq_client/src/protocol/mod.rs +++ b/crates/engine_zmq_client/src/protocol/mod.rs @@ -11,13 +11,59 @@ //! provides one implementation. use bytes::Bytes; +use serde::de::{Deserialize, Error as _, IgnoredAny, SeqAccess}; use crate::Result; pub mod handshake; +pub mod sglang; pub mod tokenspeed; pub mod vllm; +/// Read the next positional element of a msgspec `array_like` struct, failing +/// loudly when the array is shorter than the modeled prefix (every modeled field +/// is required on decode). +pub(crate) fn next_field<'de, A, T>( + seq: &mut A, + name: &'static str, +) -> std::result::Result +where + A: SeqAccess<'de>, + T: Deserialize<'de>, +{ + seq.next_element::()? + .ok_or_else(|| A::Error::custom(format!("missing positional field `{name}`"))) +} + +/// Validate the msgspec tag string at element 0. A wrong tag means the payload +/// is a different message type — fail loudly instead of misreading fields. +pub(crate) fn expect_tag<'de, A>( + seq: &mut A, + expected: &'static str, +) -> std::result::Result<(), A::Error> +where + A: SeqAccess<'de>, +{ + let tag: String = next_field(seq, "_tag")?; + if tag != expected { + return Err(A::Error::custom(format!( + "wrong msgspec tag: expected `{expected}`, got `{tag}`" + ))); + } + Ok(()) +} + +/// Drain positional elements beyond the modeled prefix. Engines append new +/// fields at the end of their structs over time, so unknown trailing elements +/// are skipped rather than treated as a decode error. +pub(crate) fn drain_trailing<'de, A>(seq: &mut A) -> std::result::Result<(), A::Error> +where + A: SeqAccess<'de>, +{ + while seq.next_element::()?.is_some() {} + Ok(()) +} + /// Engine-neutral per-rank load signal. Each protocol maps its native scheduler /// stats into this shape (vLLM maps `SchedulerStats`; TokenSpeed carries none /// yet), so the connector and gateway consume one load type regardless of @@ -79,10 +125,12 @@ pub trait EngineProtocol: Send + Sync + 'static { /// The typed per-request output streamed back. type Output: EngineOutput + Send + 'static; - /// The single-byte request-type frame prepended to an add-request. - fn add_frame() -> Bytes; - /// The single-byte request-type frame prepended to an abort. - fn abort_frame() -> Bytes; + /// The single-byte request-type frame prepended to an add-request, or + /// `None` for protocols that dispatch by payload tag alone (no type frame). + fn add_frame() -> Option; + /// The single-byte request-type frame prepended to an abort, or `None` for + /// protocols with no type frame. + fn abort_frame() -> Option; /// The request's id (the registry routing key). fn request_id(request: &Self::Request) -> &str; diff --git a/crates/engine_zmq_client/src/protocol/sglang/mod.rs b/crates/engine_zmq_client/src/protocol/sglang/mod.rs new file mode 100644 index 000000000..3525c230c --- /dev/null +++ b/crates/engine_zmq_client/src/protocol/sglang/mod.rs @@ -0,0 +1,178 @@ +//! SGLang wire protocol. +//! +//! Clean-room port of SGLang's native inter-process messages +//! (`sglang/srt/managers/io_struct.py`, `sglang/srt/sampling/sampling_params.py`). +//! Every data-plane message is a Python `msgspec.Struct(tag=True, kw_only=True, +//! array_like=True)`: it rides the wire as a positional msgpack array whose +//! element 0 is the class-name tag string, followed by the fields in declaration +//! order. Field order is the wire contract — append only, never reorder. Decoders +//! here validate the tag and skip trailing fields they do not model (SGLang +//! appends fields over time); the encoder emits the shortest valid prefix, since +//! `msgspec` fills missing trailing fields from their defaults. +//! +//! Three structural differences from the TokenSpeed protocol drive this module: +//! +//! - **No request-type frame.** SGLang dispatches purely by the msgspec tag +//! (element 0), so [`SglangProtocol::add_frame`] / [`SglangProtocol::abort_frame`] +//! return `None` and the abort is a tagged [`AbortReq`] struct, not a bare +//! id list. +//! - **`SamplingParams` is an untagged positional array.** It is +//! `array_like=True` with no tag, so it rides the wire as a bare msgpack array +//! in field order (not keyed by name and no leading tag element); SMG emits +//! every field, carrying SGLang defaults for the ones it does not set, and the +//! scheduler normalizes on receipt. +//! - **Token ids are typed arrays.** `input_ids` / `output_ids` use SGLang's +//! `array.array('q', ...)` form — the 2-tuple `["q", ]` — not +//! a plain msgpack integer list (see [`token_ids`]). +//! +//! Only the text-generation (skip-tokenizer) path is typed. + +pub mod output; +pub mod request; +pub mod sampling; +pub mod token_ids; + +use bytes::Bytes; + +use crate::{ + codec::{decode_msgpack, encode_msgpack}, + error::Result, + protocol::{ + sglang::{ + output::{BatchTokenIDOutput, SglangOutput}, + request::{AbortReq, TokenizedGenerateReqInput}, + }, + EngineBatch, EngineProtocol, + }, +}; + +/// The SGLang engine protocol: drives [`TokenizedGenerateReqInput`] over the +/// tag-dispatched ZMQ transport and decodes [`BatchTokenIDOutput`] back. +pub struct SglangProtocol; + +impl EngineProtocol for SglangProtocol { + type Request = TokenizedGenerateReqInput; + type Output = SglangOutput; + + fn add_frame() -> Option { + // SGLang dispatches by the msgspec tag alone — no request-type frame. + None + } + + fn abort_frame() -> Option { + None + } + + fn request_id(request: &Self::Request) -> &str { + &request.rid + } + + fn data_parallel_rank(_request: &Self::Request) -> Option { + // The modeled request prefix carries no DP-rank field, so requests route + // to the sole engine (single-engine ZMQ). DP fan-out is future work. + None + } + + fn validate(_request: &Self::Request) -> Result<()> { + // The tokenized text path has no fields this client cannot represent. + Ok(()) + } + + fn encode_add(request: &Self::Request) -> Result<(Vec, Vec)> { + // Text path carries no aux tensor frames. + Ok((encode_msgpack(request)?, Vec::new())) + } + + fn encode_abort(request_id: &str) -> Result> { + // The abort is a tagged AbortReq struct (not a bare id list): the + // scheduler matches `rid` against in-flight requests. + encode_msgpack(&AbortReq::new(request_id)) + } + + fn decode_batch(frames: &[Bytes]) -> Result> { + // Output messages are `[payload, aux...]`. The token-id batch carries no + // tensor fields today, so aux frames are unexpected but not fatal. + if frames.len() > 1 { + tracing::debug!( + aux_frames = frames.len() - 1, + "ignoring aux frames on an SGLang output (BatchTokenIDOutput has \ + no tensor fields on the text path)" + ); + } + let payload = frames.first().map(AsRef::as_ref).unwrap_or_default(); + let batch: BatchTokenIDOutput = decode_msgpack(payload)?; + let outputs = batch.into_outputs()?; + let finished_request_ids = outputs + .iter() + .filter(|output| output.finish_reason.is_some()) + .map(|output| output.request_id.clone()) + .collect(); + Ok(EngineBatch { + // Single-engine ZMQ: SGLang batches carry no engine index and no + // piggybacked scheduler load on this path. + engine_index: 0, + outputs, + finished_request_ids, + load: None, + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::codec::decode_msgpack; + + /// The pinned finished-variant output vector (see `output.rs`): one request + /// "req-00000001" with `finished_reasons = [{"type":"stop","matched":2}]`. + const PYTHON_OUTPUT_FINISHED_VECTOR: &str = + "dc0032b24261746368546f6b656e49444f757470757491ac7265712d3030303030303031c091\ + 82a474797065a473746f70a76d6174636865640291a09192a171c408883c0000000000009100\ + 9192a171c408883c00000000000091c391c391c29103910091019100c0c0c0c0c0c0c0c0c0c0\ + c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0"; + + fn from_hex(hex: &str) -> Vec { + let hex: String = hex.chars().filter(|c| !c.is_whitespace()).collect(); + (0..hex.len()) + .step_by(2) + .map(|i| u8::from_str_radix(&hex[i..i + 2], 16).unwrap()) + .collect() + } + + #[test] + fn no_request_type_frames() { + assert!(SglangProtocol::add_frame().is_none()); + assert!(SglangProtocol::abort_frame().is_none()); + } + + #[test] + fn encode_abort_emits_a_tagged_abort_req() { + let payload = SglangProtocol::encode_abort("req-1").unwrap(); + let decoded: AbortReq = decode_msgpack(&payload).unwrap(); + assert_eq!(decoded, AbortReq::new("req-1")); + } + + #[test] + fn decode_batch_maps_outputs_and_finished_ids() { + let frames = vec![Bytes::from(from_hex(PYTHON_OUTPUT_FINISHED_VECTOR))]; + let decoded = SglangProtocol::decode_batch(&frames).unwrap(); + assert_eq!(decoded.outputs.len(), 1); + assert_eq!(decoded.outputs[0].request_id, "req-00000001"); + assert_eq!(decoded.outputs[0].output_ids, vec![15496]); + assert_eq!( + decoded.finished_request_ids, + vec!["req-00000001".to_string()] + ); + assert!(decoded.load.is_none()); + } + + #[test] + fn decode_batch_tolerates_aux_frames() { + let frames = vec![ + Bytes::from(from_hex(PYTHON_OUTPUT_FINISHED_VECTOR)), + Bytes::from_static(b"opaque-aux-frame"), + ]; + let decoded = SglangProtocol::decode_batch(&frames).unwrap(); + assert_eq!(decoded.outputs.len(), 1); + } +} diff --git a/crates/engine_zmq_client/src/protocol/sglang/output.rs b/crates/engine_zmq_client/src/protocol/sglang/output.rs new file mode 100644 index 000000000..54175495f --- /dev/null +++ b/crates/engine_zmq_client/src/protocol/sglang/output.rs @@ -0,0 +1,368 @@ +// SGLang per-step batched output — the native `BatchTokenIDOutput` from +// `io_struct.py`, a tagged `msgspec.Struct(array_like=True)`: on the wire it is +// a positional msgpack array with the class-name tag string as element 0. +// **Field order is the wire contract** — do not reorder. SGLang carries far +// more columns than SMG consumes; the decoder reads the modeled prefix (through +// the output logprob columns), consuming the intervening columns positionally, +// and skips everything past it. + +use serde::{ + de::{IgnoredAny, SeqAccess, Visitor}, + Deserialize, Deserializer, +}; + +use crate::{ + error::{Error, Result}, + protocol::{expect_tag, next_field, sglang::token_ids::TokenIdArray, EngineOutput}, +}; + +/// The msgspec tag for [`BatchTokenIDOutput`] (element 0 on the wire). +pub const BATCH_TOKEN_ID_OUTPUT_TAG: &str = "BatchTokenIDOutput"; + +/// A finish reason entry — SGLang encodes each as a dict such as +/// `{"type": "stop", "matched": 2}` (or `nil` while the request is still +/// generating). Only the `type` discriminant matters to SMG; the extra keys +/// (`matched` / `length` / ...) are ignored on decode. +#[derive(Debug, Clone, PartialEq, Deserialize)] +pub struct FinishReason { + /// The finish-reason kind (`"stop"`, `"length"`, `"abort"`). + #[serde(rename = "type")] + pub kind: String, +} + +/// A batch of per-request token outputs from one scheduler step. Every field is +/// a column indexed in parallel by request: `rids[i]` owns `output_ids[i]`, +/// `finished_reasons[i]`, and the token-count columns. +/// +/// The logprob columns are the SGLang nesting `List[Optional[List[Optional[T]]]]` +/// — outer per request, inner per token — and are absent (`None`) unless the +/// request asked for logprobs. +#[derive(Debug, Clone, PartialEq, Default)] +pub struct BatchTokenIDOutput { + /// Request ids, one per column. + pub rids: Vec, + /// Finish reason per request; `None` while still generating. + pub finished_reasons: Vec>, + /// Newly generated token ids per request this step. + pub output_ids: Vec, + /// Prompt token count per request. + pub prompt_tokens: Vec, + /// Reasoning-phase token count per request. + pub reasoning_tokens: Vec, + /// Completion token count so far, per request. + pub completion_tokens: Vec, + /// Prefix-cache-hit token count per request. + pub cached_tokens: Vec, + /// Sampled-token logprob values per request (per-token inner list). + pub output_token_logprobs_val: Option>>>>, + /// Token id each logprob belongs to, per request (per-token inner list). + pub output_token_logprobs_idx: Option>>>>, +} + +impl<'de> Deserialize<'de> for BatchTokenIDOutput { + fn deserialize>(deserializer: D) -> std::result::Result { + struct BatchVisitor; + + impl<'de> Visitor<'de> for BatchVisitor { + type Value = BatchTokenIDOutput; + + fn expecting(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "a tagged BatchTokenIDOutput positional array") + } + + fn visit_seq>( + self, + mut seq: A, + ) -> std::result::Result { + expect_tag(&mut seq, BATCH_TOKEN_ID_OUTPUT_TAG)?; + // rids / http_worker_ipcs are the BaseBatchReq prefix. + let rids = + next_field::<_, Option>>(&mut seq, "rids")?.unwrap_or_default(); + next_field::<_, IgnoredAny>(&mut seq, "http_worker_ipcs")?; + let finished_reasons = next_field(&mut seq, "finished_reasons")?; + next_field::<_, IgnoredAny>(&mut seq, "decoded_texts")?; + next_field::<_, IgnoredAny>(&mut seq, "decode_ids")?; + next_field::<_, IgnoredAny>(&mut seq, "read_offsets")?; + let output_ids = + next_field::<_, Option>>(&mut seq, "output_ids")? + .unwrap_or_default(); + next_field::<_, IgnoredAny>(&mut seq, "skip_special_tokens")?; + next_field::<_, IgnoredAny>(&mut seq, "spaces_between_special_tokens")?; + next_field::<_, IgnoredAny>(&mut seq, "no_stop_trim")?; + let prompt_tokens = next_field(&mut seq, "prompt_tokens")?; + let reasoning_tokens = next_field(&mut seq, "reasoning_tokens")?; + let completion_tokens = next_field(&mut seq, "completion_tokens")?; + let cached_tokens = next_field(&mut seq, "cached_tokens")?; + next_field::<_, IgnoredAny>(&mut seq, "input_token_logprobs_val")?; + next_field::<_, IgnoredAny>(&mut seq, "input_token_logprobs_idx")?; + let output_token_logprobs_val = next_field(&mut seq, "output_token_logprobs_val")?; + let output_token_logprobs_idx = next_field(&mut seq, "output_token_logprobs_idx")?; + // SGLang appends many more columns; skip everything past here. + while seq.next_element::()?.is_some() {} + Ok(BatchTokenIDOutput { + rids, + finished_reasons, + output_ids, + prompt_tokens, + reasoning_tokens, + completion_tokens, + cached_tokens, + output_token_logprobs_val, + output_token_logprobs_idx, + }) + } + } + + deserializer.deserialize_seq(BatchVisitor) + } +} + +/// One request's slice of a [`BatchTokenIDOutput`], in the engine-neutral shape +/// the connector routes to per-request streams. +#[derive(Debug, Clone, PartialEq, Default)] +pub struct SglangOutput { + /// The request id this output belongs to. + pub request_id: String, + /// Newly generated token ids this step. + pub output_ids: Vec, + /// Finish reason; `None` while the request is still generating. + pub finish_reason: Option, + /// Prompt token count. + pub prompt_tokens: u32, + /// Reasoning-phase token count. + pub reasoning_tokens: u32, + /// Completion token count so far. + pub completion_tokens: u32, + /// Prefix-cache-hit token count. + pub cached_tokens: u32, + /// Sampled-token logprob value per decoded token this step. Empty when + /// logprobs were not requested. + pub output_logprobs_val: Vec, + /// Token id each logprob in `output_logprobs_val` belongs to. Empty when + /// logprobs were not requested. + pub output_logprobs_idx: Vec, +} + +impl EngineOutput for SglangOutput { + fn request_id(&self) -> &str { + &self.request_id + } + + fn finished(&self) -> bool { + self.finish_reason.is_some() + } +} + +impl BatchTokenIDOutput { + /// Split the parallel columns into one [`SglangOutput`] per request. Errors + /// if the required columns are ragged (a length mismatch is a protocol bug). + pub fn into_outputs(self) -> Result> { + let n = self.rids.len(); + + // `output_ids` may arrive as a whole-field `nil` (no new tokens this + // step); treat that as an empty column so it stays aligned. + let output_ids = if self.output_ids.is_empty() { + vec![TokenIdArray::default(); n] + } else { + self.output_ids + }; + + let ragged = self.finished_reasons.len() != n + || output_ids.len() != n + || self.prompt_tokens.len() != n + || self.reasoning_tokens.len() != n + || self.completion_tokens.len() != n + || self.cached_tokens.len() != n; + if ragged { + return Err(Error::Decode { + target_type: "BatchTokenIDOutput", + message: format!( + "ragged columns: rids={n}, finished_reasons={}, output_ids={}, \ + prompt_tokens={}, reasoning_tokens={}, completion_tokens={}, \ + cached_tokens={}", + self.finished_reasons.len(), + output_ids.len(), + self.prompt_tokens.len(), + self.reasoning_tokens.len(), + self.completion_tokens.len(), + self.cached_tokens.len(), + ), + }); + } + + let BatchTokenIDOutput { + rids, + finished_reasons, + prompt_tokens, + reasoning_tokens, + completion_tokens, + cached_tokens, + output_token_logprobs_val, + output_token_logprobs_idx, + .. + } = self; + + Ok(rids + .into_iter() + .zip(finished_reasons) + .zip(output_ids) + .zip(prompt_tokens) + .zip(reasoning_tokens) + .zip(completion_tokens) + .zip(cached_tokens) + .enumerate() + .map( + |(index, ((((((rid, reason), ids), prompt), reasoning), completion), cached))| { + SglangOutput { + request_id: rid, + output_ids: ids.0, + finish_reason: reason.map(|reason| reason.kind), + prompt_tokens: prompt, + reasoning_tokens: reasoning, + completion_tokens: completion, + cached_tokens: cached, + output_logprobs_val: per_request_logprobs( + output_token_logprobs_val.as_ref(), + index, + ), + output_logprobs_idx: per_request_logprobs( + output_token_logprobs_idx.as_ref(), + index, + ), + } + }, + ) + .collect()) + } +} + +/// Flatten request `index`'s logprob column into a dense vec, dropping the +/// per-token `Option` nesting (`None` inner values are omitted). Returns an +/// empty vec when the column is absent or the request had no logprobs. +fn per_request_logprobs( + column: Option<&Vec>>>>, + index: usize, +) -> Vec { + column + .and_then(|rows| rows.get(index)) + .and_then(Option::as_ref) + .map(|values| values.iter().filter_map(|value| *value).collect()) + .unwrap_or_default() +} + +#[cfg(test)] +mod tests { + use rmpv::Value; + + use super::*; + use crate::codec::{decode_msgpack, decode_value}; + + fn from_hex(hex: &str) -> Vec { + let hex: String = hex.chars().filter(|c| !c.is_whitespace()).collect(); + (0..hex.len()) + .step_by(2) + .map(|i| u8::from_str_radix(&hex[i..i + 2], 16).unwrap()) + .collect() + } + + /// The pinned output vector (still-generating variant): one request + /// "req-00000001", output_ids [[15496]], finish nil, prompt 3 / completion 1. + const PYTHON_OUTPUT_VECTOR: &str = + "dc0032b24261746368546f6b656e49444f757470757491ac7265712d3030303030303031c091\ + c091a09192a171c408883c00000000000091009192a171c408883c00000000000091c391c391\ + c29103910091019100c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0\ + c0c0c0c0c0c0"; + + /// The finished variant: finished_reasons carries `{"type":"stop","matched":2}`. + const PYTHON_OUTPUT_FINISHED_VECTOR: &str = + "dc0032b24261746368546f6b656e49444f757470757491ac7265712d3030303030303031c091\ + 82a474797065a473746f70a76d6174636865640291a09192a171c408883c0000000000009100\ + 9192a171c408883c00000000000091c391c391c29103910091019100c0c0c0c0c0c0c0c0c0c0\ + c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0"; + + #[test] + fn python_output_vector_decodes_still_generating() { + let decoded: BatchTokenIDOutput = decode_msgpack(&from_hex(PYTHON_OUTPUT_VECTOR)).unwrap(); + assert_eq!(decoded.rids, vec!["req-00000001".to_string()]); + assert_eq!(decoded.finished_reasons, vec![None]); + assert_eq!(decoded.output_ids, vec![TokenIdArray(vec![15496])]); + assert_eq!(decoded.prompt_tokens, vec![3]); + assert_eq!(decoded.completion_tokens, vec![1]); + + let outputs = decoded.into_outputs().unwrap(); + assert_eq!(outputs.len(), 1); + assert_eq!(outputs[0].request_id, "req-00000001"); + assert_eq!(outputs[0].output_ids, vec![15496]); + assert_eq!(outputs[0].finish_reason, None); + assert!(!outputs[0].finished()); + assert_eq!(outputs[0].prompt_tokens, 3); + assert_eq!(outputs[0].completion_tokens, 1); + } + + #[test] + fn python_output_vector_decodes_finished() { + let decoded: BatchTokenIDOutput = + decode_msgpack(&from_hex(PYTHON_OUTPUT_FINISHED_VECTOR)).unwrap(); + assert_eq!( + decoded.finished_reasons, + vec![Some(FinishReason { + kind: "stop".to_string() + })] + ); + + let outputs = decoded.into_outputs().unwrap(); + assert_eq!(outputs[0].finish_reason.as_deref(), Some("stop")); + assert!(outputs[0].finished()); + } + + #[test] + fn decode_rejects_wrong_tag() { + let mut bytes = from_hex(PYTHON_OUTPUT_VECTOR); + // Corrupt a char inside the tag: "BatchTokenIDOutput" -> "BXtch...". + bytes[5] = b'X'; + let error = decode_msgpack::(&bytes).unwrap_err(); + assert!(error.to_string().contains("wrong msgspec tag"), "{error}"); + } + + #[test] + fn into_outputs_rejects_ragged_columns() { + let batch = BatchTokenIDOutput { + rids: vec!["a".into(), "b".into()], + finished_reasons: vec![None, None], + output_ids: vec![TokenIdArray(vec![10])], + prompt_tokens: vec![3, 4], + reasoning_tokens: vec![0, 0], + completion_tokens: vec![1, 1], + cached_tokens: vec![0, 0], + output_token_logprobs_val: None, + output_token_logprobs_idx: None, + }; + assert!(batch.into_outputs().is_err()); + } + + #[test] + fn into_outputs_flattens_optional_logprobs() { + let batch = BatchTokenIDOutput { + rids: vec!["a".into()], + finished_reasons: vec![None], + output_ids: vec![TokenIdArray(vec![10, 11])], + prompt_tokens: vec![3], + reasoning_tokens: vec![0], + completion_tokens: vec![2], + cached_tokens: vec![0], + output_token_logprobs_val: Some(vec![Some(vec![Some(-0.5), Some(-0.25)])]), + output_token_logprobs_idx: Some(vec![Some(vec![Some(10), Some(11)])]), + }; + let outputs = batch.into_outputs().unwrap(); + assert_eq!(outputs[0].output_logprobs_val, vec![-0.5, -0.25]); + assert_eq!(outputs[0].output_logprobs_idx, vec![10, 11]); + } + + #[test] + fn output_vector_is_a_positional_array() { + let Value::Array(array) = decode_value(&from_hex(PYTHON_OUTPUT_VECTOR)).unwrap() else { + panic!("expected positional array"); + }; + assert_eq!(array[0], Value::from(BATCH_TOKEN_ID_OUTPUT_TAG)); + } +} diff --git a/crates/engine_zmq_client/src/protocol/sglang/request.rs b/crates/engine_zmq_client/src/protocol/sglang/request.rs new file mode 100644 index 000000000..3ea033fc1 --- /dev/null +++ b/crates/engine_zmq_client/src/protocol/sglang/request.rs @@ -0,0 +1,312 @@ +// SGLang tokenized generate request and abort — the native +// `TokenizedGenerateReqInput` and `AbortReq` from `io_struct.py`, tagged +// `msgspec.Struct(array_like=True)`: each rides the wire as a positional +// msgpack array whose element 0 is the class-name tag string, followed by the +// fields in declaration order. **Field order is the wire contract** — do not +// reorder. Unlike TokenSpeed, SGLang dispatches purely by this tag (there is no +// single-byte request-type frame), so the abort is a tagged struct rather than +// a bare list of ids. + +use serde::{ + de::{IgnoredAny, SeqAccess, Visitor}, + ser::SerializeTuple, + Deserialize, Deserializer, Serialize, Serializer, +}; + +use crate::protocol::{ + expect_tag, next_field, + sglang::{sampling::SamplingParams, token_ids::TokenIdArray}, +}; + +/// The msgspec tag for [`TokenizedGenerateReqInput`] (element 0 on the wire). +pub const TOKENIZED_GENERATE_REQ_INPUT_TAG: &str = "TokenizedGenerateReqInput"; +/// The msgspec tag for [`AbortReq`] (element 0 on the wire). +pub const ABORT_REQ_TAG: &str = "AbortReq"; + +/// SGLang tokenized generate request sent from frontend to scheduler. +/// +/// Models the leading prefix of the Python class, through `stream` — the fields +/// SMG sets on the token-id (skip-tokenizer) path. The encoder emits exactly +/// this 14-element array (tag + 13 fields); the scheduler's decoder fills every +/// later field from its defaults, since msgspec tolerates missing trailing +/// fields. The decoder here accepts full-length arrays and skips the unmodeled +/// trailing fields. +/// +/// The three fields between `input_ids` and `sampling_params` — `input_embeds`, +/// `mm_inputs`, `token_type_ids` — are unused on the text path; they are emitted +/// as `nil` and skipped on decode, but must be present to keep the positional +/// layout aligned. +#[derive(Debug, Clone, PartialEq)] +pub struct TokenizedGenerateReqInput { + /// Request id (the routing/registry key). + pub rid: String, + /// In-process HTTP-worker return address; unused on this transport. + pub http_worker_ipc: Option, + /// Original prompt text. `None` on the token-id path (SMG detokenizes + /// downstream of the engine, so only ids are sent). + pub input_text: Option, + /// Pre-tokenized prompt token ids (SMG tokenizes upstream), encoded as + /// SGLang's `array.array('q', ...)` wire form. + pub input_ids: TokenIdArray, + /// Sampling parameters (nested positional array). + pub sampling_params: SamplingParams, + /// Whether to return the sampled token's logprob for this request. + pub return_logprob: bool, + /// Prompt-logprob start offset. Neutral `-1`: prompt logprobs are not + /// supported on this wire. + pub logprob_start_len: i32, + /// Output top-k logprob count. Neutral `0`: only the sampled token's + /// logprob is materialized. + pub top_logprobs_num: u32, + /// Token ids to report logprobs for. Neutral `None`: not supported. + pub token_ids_logprob: Option>, + /// Whether to stream outputs incrementally. + pub stream: bool, +} + +impl Default for TokenizedGenerateReqInput { + fn default() -> Self { + Self { + rid: String::new(), + http_worker_ipc: None, + input_text: None, + input_ids: TokenIdArray::default(), + sampling_params: SamplingParams::default(), + return_logprob: false, + logprob_start_len: -1, + top_logprobs_num: 0, + token_ids_logprob: None, + stream: false, + } + } +} + +impl Serialize for TokenizedGenerateReqInput { + fn serialize(&self, serializer: S) -> Result { + let mut tuple = serializer.serialize_tuple(14)?; + tuple.serialize_element(TOKENIZED_GENERATE_REQ_INPUT_TAG)?; + tuple.serialize_element(&self.rid)?; + tuple.serialize_element(&self.http_worker_ipc)?; + tuple.serialize_element(&self.input_text)?; + tuple.serialize_element(&self.input_ids)?; + // input_embeds / mm_inputs / token_type_ids: unused on the text path. + tuple.serialize_element(&None::<()>)?; + tuple.serialize_element(&None::<()>)?; + tuple.serialize_element(&None::<()>)?; + tuple.serialize_element(&self.sampling_params)?; + tuple.serialize_element(&self.return_logprob)?; + tuple.serialize_element(&self.logprob_start_len)?; + tuple.serialize_element(&self.top_logprobs_num)?; + tuple.serialize_element(&self.token_ids_logprob)?; + tuple.serialize_element(&self.stream)?; + tuple.end() + } +} + +impl<'de> Deserialize<'de> for TokenizedGenerateReqInput { + fn deserialize>(deserializer: D) -> Result { + struct ReqVisitor; + + impl<'de> Visitor<'de> for ReqVisitor { + type Value = TokenizedGenerateReqInput; + + fn expecting(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "a tagged TokenizedGenerateReqInput positional array") + } + + fn visit_seq>(self, mut seq: A) -> Result { + expect_tag(&mut seq, TOKENIZED_GENERATE_REQ_INPUT_TAG)?; + let rid = next_field(&mut seq, "rid")?; + let http_worker_ipc = next_field(&mut seq, "http_worker_ipc")?; + let input_text = next_field(&mut seq, "input_text")?; + let input_ids = next_field(&mut seq, "input_ids")?; + // Consume the unused middle fields to keep positions aligned. + next_field::<_, IgnoredAny>(&mut seq, "input_embeds")?; + next_field::<_, IgnoredAny>(&mut seq, "mm_inputs")?; + next_field::<_, IgnoredAny>(&mut seq, "token_type_ids")?; + let request = TokenizedGenerateReqInput { + rid, + http_worker_ipc, + input_text, + input_ids, + sampling_params: next_field(&mut seq, "sampling_params")?, + return_logprob: next_field(&mut seq, "return_logprob")?, + logprob_start_len: next_field(&mut seq, "logprob_start_len")?, + top_logprobs_num: next_field(&mut seq, "top_logprobs_num")?, + token_ids_logprob: next_field(&mut seq, "token_ids_logprob")?, + stream: next_field(&mut seq, "stream")?, + }; + // SGLang appends fields over time; skip everything past `stream`. + while seq.next_element::()?.is_some() {} + Ok(request) + } + } + + deserializer.deserialize_seq(ReqVisitor) + } +} + +/// SGLang abort request. On the tag-dispatched wire an abort is a tagged +/// [`AbortReq`] struct (not a bare id list): the scheduler matches `rid` against +/// in-flight requests. `abort_all` and the two message fields are unused by SMG. +#[derive(Debug, Clone, PartialEq)] +pub struct AbortReq { + /// The request id to abort. + pub rid: String, + /// Whether to abort every in-flight request (SMG always aborts one rid). + pub abort_all: bool, +} + +impl AbortReq { + /// An abort for a single request id. + pub fn new(rid: impl Into) -> Self { + Self { + rid: rid.into(), + abort_all: false, + } + } +} + +impl Serialize for AbortReq { + fn serialize(&self, serializer: S) -> Result { + let mut tuple = serializer.serialize_tuple(6)?; + tuple.serialize_element(ABORT_REQ_TAG)?; + tuple.serialize_element(&self.rid)?; + // http_worker_ipc / finished_reason / abort_message: unused by SMG. + tuple.serialize_element(&None::<()>)?; + tuple.serialize_element(&self.abort_all)?; + tuple.serialize_element(&None::<()>)?; + tuple.serialize_element(&None::<()>)?; + tuple.end() + } +} + +impl<'de> Deserialize<'de> for AbortReq { + fn deserialize>(deserializer: D) -> Result { + struct AbortVisitor; + + impl<'de> Visitor<'de> for AbortVisitor { + type Value = AbortReq; + + fn expecting(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "a tagged AbortReq positional array") + } + + fn visit_seq>(self, mut seq: A) -> Result { + expect_tag(&mut seq, ABORT_REQ_TAG)?; + let rid = next_field(&mut seq, "rid")?; + next_field::<_, IgnoredAny>(&mut seq, "http_worker_ipc")?; + let abort_all = next_field(&mut seq, "abort_all")?; + while seq.next_element::()?.is_some() {} + Ok(AbortReq { rid, abort_all }) + } + } + + deserializer.deserialize_seq(AbortVisitor) + } +} + +#[cfg(test)] +mod tests { + use rmpv::Value; + + use super::*; + use crate::{ + codec::{decode_msgpack, decode_value, encode_msgpack}, + protocol::sglang::sampling::SamplingParams, + }; + + /// The pinned request vector captured from the Python encoder: rid + /// "req-00000001", input_text "Hello, world", input_ids [9906, 11, 1917], + /// sampling {max_new_tokens 64, temperature 0.7, top_p 0.9, top_k 50}, + /// stream true. The Python side emits all 45 elements (sampling_params is a + /// nested 30-element positional array at index 8). + const PYTHON_REQUEST_VECTOR: &str = + "dc002db9546f6b656e697a656447656e6572617465526571496e707574ac7265712d30303030\ + 30303031c0ac48656c6c6f2c20776f726c6492a171c418b2260000000000000b000000000000\ + 007d07000000000000c0c0c0dc001e40c0c0c0cb3fe6666666666666cb3feccccccccccccd32c\ + b0000000000000000cb0000000000000000cb0000000000000000cb3ff0000000000000000\ + 1c0c0c0c0c2c3c3c2c0c0c0c0c0c00000c2c2ff00c0c3c2c2c2c200c2c0c0c0c0c0c0c0c0c0c0\ + c0c0c0c2c0c0c2c2c2c0c0c0c0c0c0"; + + fn from_hex(hex: &str) -> Vec { + let hex: String = hex.chars().filter(|c| !c.is_whitespace()).collect(); + (0..hex.len()) + .step_by(2) + .map(|i| u8::from_str_radix(&hex[i..i + 2], 16).unwrap()) + .collect() + } + + fn vector_request() -> TokenizedGenerateReqInput { + TokenizedGenerateReqInput { + rid: "req-00000001".to_string(), + input_text: Some("Hello, world".to_string()), + input_ids: TokenIdArray(vec![9906, 11, 1917]), + sampling_params: SamplingParams { + max_new_tokens: Some(64), + temperature: 0.7, + top_p: 0.9, + top_k: 50, + ..SamplingParams::default() + }, + stream: true, + ..TokenizedGenerateReqInput::default() + } + } + + #[test] + fn python_request_vector_decodes() { + let decoded: TokenizedGenerateReqInput = + decode_msgpack(&from_hex(PYTHON_REQUEST_VECTOR)).unwrap(); + assert_eq!(decoded, vector_request()); + assert_eq!(decoded.input_ids, TokenIdArray(vec![9906, 11, 1917])); + assert_eq!(decoded.logprob_start_len, -1); + assert!(decoded.stream); + assert!(!decoded.return_logprob); + } + + #[test] + fn encoder_emits_tagged_prefix_through_stream() { + let encoded = encode_msgpack(&vector_request()).unwrap(); + let Value::Array(array) = decode_value(&encoded).unwrap() else { + panic!("expected positional array"); + }; + // tag + 13 fields; the three unused middle fields are nil. + assert_eq!(array.len(), 14); + assert_eq!(array[0], Value::from(TOKENIZED_GENERATE_REQ_INPUT_TAG)); + assert_eq!(array[1], Value::from("req-00000001")); + assert_eq!(array[5], Value::Nil); // input_embeds + assert_eq!(array[6], Value::Nil); // mm_inputs + assert_eq!(array[7], Value::Nil); // token_type_ids + assert_eq!(array[13], Value::from(true)); // stream + + // Round-trip: the prefix encoding decodes to the same request as the + // full-length Python vector. + let roundtripped: TokenizedGenerateReqInput = decode_msgpack(&encoded).unwrap(); + assert_eq!(roundtripped, vector_request()); + let from_vector: TokenizedGenerateReqInput = + decode_msgpack(&from_hex(PYTHON_REQUEST_VECTOR)).unwrap(); + assert_eq!(roundtripped, from_vector); + } + + #[test] + fn request_decode_rejects_wrong_tag() { + let mut bytes = from_hex(PYTHON_REQUEST_VECTOR); + // Corrupt one tag byte inside "TokenizedGenerateReqInput". + bytes[4] = b'X'; + let error = decode_msgpack::(&bytes).unwrap_err(); + assert!(error.to_string().contains("wrong msgspec tag"), "{error}"); + } + + #[test] + fn abort_req_matches_python_bytes() { + // The pinned abort vector: [tag, "req-00000001", nil, false, nil, nil]. + const PYTHON_ABORT_VECTOR: &str = "96a841626f7274526571ac7265712d3030303030303031c0c2c0c0"; + let encoded = encode_msgpack(&AbortReq::new("req-00000001")).unwrap(); + assert_eq!(encoded, from_hex(PYTHON_ABORT_VECTOR)); + + let decoded: AbortReq = decode_msgpack(&encoded).unwrap(); + assert_eq!(decoded, AbortReq::new("req-00000001")); + assert!(!decoded.abort_all); + } +} diff --git a/crates/engine_zmq_client/src/protocol/sglang/sampling.rs b/crates/engine_zmq_client/src/protocol/sglang/sampling.rs new file mode 100644 index 000000000..40036a43b --- /dev/null +++ b/crates/engine_zmq_client/src/protocol/sglang/sampling.rs @@ -0,0 +1,341 @@ +// SGLang `SamplingParams` — a Python `msgspec.Struct(kw_only=True, +// array_like=True)` (sglang/srt/sampling/sampling_params.py), so it rides the +// wire as a positional msgpack array (no tag element), fields in declaration +// order. **Field order is the wire contract** — do not reorder. +// +// This is a faithful codec: [`SamplingParams::default`] mirrors SGLang's own +// class defaults so the positional array is always well-formed, and a producer +// sets only the fields it means to drive. Whether the internal (tokenizer- +// derived) fields are pre-resolved is the producer's policy, not the codec's — +// on the direct path SMG owns the tokenizer and sends them already resolved. + +use std::collections::HashMap; + +use serde::{ + de::{SeqAccess, Visitor}, + ser::SerializeTuple, + Deserialize, Deserializer, Serialize, Serializer, +}; + +use crate::codec::OpaqueValue; + +/// SGLang's "consider the whole vocabulary" `top_k` sentinel (`1 << 30`). SMG +/// forwards this when no explicit cutoff is requested, matching the SGLang +/// default so the scheduler samples over the full distribution. +pub const TOP_K_ALL: i32 = 1 << 30; + +/// The number of positional fields SGLang's `SamplingParams` encodes (the full +/// `array_like` struct is emitted, no `omit_defaults`). +const FIELD_COUNT: usize = 30; + +/// Engine-facing sampling parameters for SGLang text generation. +/// +/// Field order and count mirror the Python `SamplingParams` exactly: the struct +/// rides the wire as a 30-element positional array. Genuinely-optional SGLang +/// fields are `Option` (encoded as `nil` when unset); the rest carry concrete +/// values defaulting to the SGLang default. +#[derive(Debug, Clone, PartialEq)] +pub struct SamplingParams { + /// Maximum number of tokens to generate (SGLang default `128`, but + /// `Optional`; `None` leaves it unbounded by this field). + pub max_new_tokens: Option, + /// Stop strings. SMG rejects stop strings upstream on this wire, so unset. + pub stop: Option>, + /// Token ids that stop generation (a Python set; encoded as an array). + pub stop_token_ids: Option>, + /// Stop regexes. SMG rejects constraints upstream, so unset. + pub stop_regex: Option>, + /// Controls randomness (SGLang default `1.0`). + pub temperature: f64, + /// Cumulative probability threshold for nucleus sampling (default `1.0`). + pub top_p: f64, + /// Maximum number of top tokens to consider. Defaults to [`TOP_K_ALL`]. + pub top_k: i32, + /// Minimum probability threshold for token sampling (default `0.0`). + pub min_p: f64, + /// Frequency penalty applied by the sampler (default `0.0`). + pub frequency_penalty: f64, + /// Presence penalty applied by the sampler (default `0.0`). + pub presence_penalty: f64, + /// Repetition penalty applied by the sampler (default `1.0`). + pub repetition_penalty: f64, + /// Minimum number of tokens to generate before EOS / stop handling + /// (default `0`). + pub min_new_tokens: u32, + /// OpenAI-compat fanout count. SMG fans out `n > 1` itself, so this is + /// always `1` on the wire (default `1`). + pub n: u32, + /// Structured-output JSON schema. SMG rejects constraints upstream. + pub json_schema: Option, + /// Structured-output regex. SMG rejects constraints upstream. + pub regex: Option, + /// Structured-output EBNF grammar. SMG rejects constraints upstream. + pub ebnf: Option, + /// Structured-output structural tag. SMG rejects constraints upstream. + pub structural_tag: Option, + /// Ignore the EOS token and keep generating until another stop condition + /// (default `false`). + pub ignore_eos: bool, + /// Whether detokenization skips special tokens (default `true`). + pub skip_special_tokens: bool, + /// Whether detokenization inserts spaces between special tokens + /// (default `true`). + pub spaces_between_special_tokens: bool, + /// Whether stop sequences are kept in the output text (default `false`). + pub no_stop_trim: bool, + /// Streaming flush interval override. Not set by SMG. + pub stream_interval: Option, + /// Per-token logit bias, keyed by stringified token id. SMG rejects + /// logit_bias upstream (no support on this backend). + pub logit_bias: Option>, + /// Random seed. `None` lets the engine derive one so all ranks agree. + pub sampling_seed: Option, + /// Free-form engine extension parameters. Not set by SMG. + pub custom_params: Option, + /// Normalized stop strings, resolved by the engine. Not set by SMG. + pub stop_strs: Option>, + /// Normalized stop regexes, resolved by the engine. Not set by SMG. + pub stop_regex_strs: Option>, + /// Longest stop string in tokens, resolved by the engine (default `0`). + pub stop_str_max_len: u32, + /// Longest stop regex in tokens, resolved by the engine (default `0`). + pub stop_regex_max_len: u32, + /// True once the tokenizer-derived fields (`stop_strs`, `stop_str_max_len`, + /// …) are resolved. Defaults to `false` to match SGLang's class default; + /// producers that resolve those fields themselves set it `true`. + pub is_normalized: bool, +} + +impl Default for SamplingParams { + fn default() -> Self { + // Mirror the SGLang class defaults so an untouched instance encodes to + // the same positional array SGLang itself would emit. + Self { + max_new_tokens: Some(128), + stop: None, + stop_token_ids: None, + stop_regex: None, + temperature: 1.0, + top_p: 1.0, + top_k: TOP_K_ALL, + min_p: 0.0, + frequency_penalty: 0.0, + presence_penalty: 0.0, + repetition_penalty: 1.0, + min_new_tokens: 0, + n: 1, + json_schema: None, + regex: None, + ebnf: None, + structural_tag: None, + ignore_eos: false, + skip_special_tokens: true, + spaces_between_special_tokens: true, + no_stop_trim: false, + stream_interval: None, + logit_bias: None, + sampling_seed: None, + custom_params: None, + stop_strs: None, + stop_regex_strs: None, + stop_str_max_len: 0, + stop_regex_max_len: 0, + is_normalized: false, + } + } +} + +impl Serialize for SamplingParams { + fn serialize(&self, serializer: S) -> Result { + let mut tuple = serializer.serialize_tuple(FIELD_COUNT)?; + tuple.serialize_element(&self.max_new_tokens)?; + tuple.serialize_element(&self.stop)?; + tuple.serialize_element(&self.stop_token_ids)?; + tuple.serialize_element(&self.stop_regex)?; + tuple.serialize_element(&self.temperature)?; + tuple.serialize_element(&self.top_p)?; + tuple.serialize_element(&self.top_k)?; + tuple.serialize_element(&self.min_p)?; + tuple.serialize_element(&self.frequency_penalty)?; + tuple.serialize_element(&self.presence_penalty)?; + tuple.serialize_element(&self.repetition_penalty)?; + tuple.serialize_element(&self.min_new_tokens)?; + tuple.serialize_element(&self.n)?; + tuple.serialize_element(&self.json_schema)?; + tuple.serialize_element(&self.regex)?; + tuple.serialize_element(&self.ebnf)?; + tuple.serialize_element(&self.structural_tag)?; + tuple.serialize_element(&self.ignore_eos)?; + tuple.serialize_element(&self.skip_special_tokens)?; + tuple.serialize_element(&self.spaces_between_special_tokens)?; + tuple.serialize_element(&self.no_stop_trim)?; + tuple.serialize_element(&self.stream_interval)?; + tuple.serialize_element(&self.logit_bias)?; + tuple.serialize_element(&self.sampling_seed)?; + tuple.serialize_element(&self.custom_params)?; + tuple.serialize_element(&self.stop_strs)?; + tuple.serialize_element(&self.stop_regex_strs)?; + tuple.serialize_element(&self.stop_str_max_len)?; + tuple.serialize_element(&self.stop_regex_max_len)?; + tuple.serialize_element(&self.is_normalized)?; + tuple.end() + } +} + +impl<'de> Deserialize<'de> for SamplingParams { + fn deserialize>(deserializer: D) -> Result { + struct ParamsVisitor; + + impl<'de> Visitor<'de> for ParamsVisitor { + type Value = SamplingParams; + + fn expecting(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "a SamplingParams positional array") + } + + fn visit_seq>(self, mut seq: A) -> Result { + let default = SamplingParams::default(); + // Each position falls back to the SGLang default when the array + // is shorter than the full field list (msgspec omits nothing + // today, but tolerate short arrays for forward compatibility). + macro_rules! field { + ($name:ident) => { + seq.next_element()?.unwrap_or(default.$name) + }; + } + let params = SamplingParams { + max_new_tokens: seq.next_element()?.unwrap_or(default.max_new_tokens), + stop: seq.next_element()?.unwrap_or(default.stop), + stop_token_ids: seq.next_element()?.unwrap_or(default.stop_token_ids), + stop_regex: seq.next_element()?.unwrap_or(default.stop_regex), + temperature: field!(temperature), + top_p: field!(top_p), + top_k: field!(top_k), + min_p: field!(min_p), + frequency_penalty: field!(frequency_penalty), + presence_penalty: field!(presence_penalty), + repetition_penalty: field!(repetition_penalty), + min_new_tokens: field!(min_new_tokens), + n: field!(n), + json_schema: seq.next_element()?.unwrap_or(default.json_schema), + regex: seq.next_element()?.unwrap_or(default.regex), + ebnf: seq.next_element()?.unwrap_or(default.ebnf), + structural_tag: seq.next_element()?.unwrap_or(default.structural_tag), + ignore_eos: field!(ignore_eos), + skip_special_tokens: field!(skip_special_tokens), + spaces_between_special_tokens: field!(spaces_between_special_tokens), + no_stop_trim: field!(no_stop_trim), + stream_interval: seq.next_element()?.unwrap_or(default.stream_interval), + logit_bias: seq.next_element()?.unwrap_or(default.logit_bias), + sampling_seed: seq.next_element()?.unwrap_or(default.sampling_seed), + custom_params: seq.next_element()?.unwrap_or(default.custom_params), + stop_strs: seq.next_element()?.unwrap_or(default.stop_strs), + stop_regex_strs: seq.next_element()?.unwrap_or(default.stop_regex_strs), + stop_str_max_len: field!(stop_str_max_len), + stop_regex_max_len: field!(stop_regex_max_len), + is_normalized: field!(is_normalized), + }; + // SGLang appends fields over time; skip everything past the + // modeled prefix. + while seq.next_element::()?.is_some() {} + Ok(params) + } + } + + deserializer.deserialize_seq(ParamsVisitor) + } +} + +#[cfg(test)] +mod tests { + use rmpv::Value; + + use super::*; + use crate::codec::{decode_msgpack, decode_value, encode_msgpack}; + + /// The pinned SamplingParams array captured from the installed SGLang + /// encoder (`msgspec.msgpack`, `SGLANG_USE_PICKLE_IPC=0`) for + /// `SamplingParams(max_new_tokens=64, temperature=0.7, top_p=0.9, top_k=50)` + /// — a 30-element positional array. + const PYTHON_SAMPLING_VECTOR: &str = + "dc001e40c0c0c0cb3fe6666666666666cb3feccccccccccccd32cb0000000000000000cb00000000\ + 00000000cb0000000000000000cb3ff00000000000000001c0c0c0c0c2c3c3c2c0c0c0c0c0c00000c2"; + + fn python_sampling_bytes() -> Vec { + let hex: String = PYTHON_SAMPLING_VECTOR + .chars() + .filter(|c| !c.is_whitespace()) + .collect(); + (0..hex.len()) + .step_by(2) + .map(|i| u8::from_str_radix(&hex[i..i + 2], 16).unwrap()) + .collect() + } + + fn vector_sampling() -> SamplingParams { + SamplingParams { + max_new_tokens: Some(64), + temperature: 0.7, + top_p: 0.9, + top_k: 50, + ..SamplingParams::default() + } + } + + #[test] + fn python_sampling_vector_decodes() { + let decoded: SamplingParams = decode_msgpack(&python_sampling_bytes()).unwrap(); + assert_eq!(decoded, vector_sampling()); + // Fields SMG never set carry the SGLang defaults. + assert_eq!(decoded.min_p, 0.0); + assert_eq!(decoded.n, 1); + assert!(!decoded.is_normalized); + } + + #[test] + fn encoder_matches_python_bytes() { + let encoded = encode_msgpack(&vector_sampling()).unwrap(); + assert_eq!(encoded, python_sampling_bytes()); + } + + #[test] + fn encodes_as_a_positional_array() { + let encoded = encode_msgpack(&vector_sampling()).unwrap(); + let Value::Array(elements) = decode_value(&encoded).unwrap() else { + panic!("expected a msgpack array (array_like=True struct)"); + }; + assert_eq!(elements.len(), FIELD_COUNT); + assert_eq!(elements[0], Value::from(64u32)); // max_new_tokens + assert_eq!(elements[6], Value::from(50)); // top_k + assert_eq!(elements[12], Value::from(1)); // n + assert_eq!(elements[29], Value::from(false)); // is_normalized + } + + #[test] + fn roundtrips_through_the_array_wire_form() { + let params = SamplingParams { + max_new_tokens: Some(128), + stop_token_ids: Some(vec![2, 3]), + temperature: 0.5, + top_p: 0.95, + top_k: 40, + min_p: 0.05, + frequency_penalty: 0.1, + sampling_seed: Some(42), + ignore_eos: true, + ..SamplingParams::default() + }; + let encoded = encode_msgpack(¶ms).unwrap(); + assert_eq!(decode_msgpack::(&encoded).unwrap(), params); + } + + #[test] + fn default_encodes_all_fields() { + let encoded = encode_msgpack(&SamplingParams::default()).unwrap(); + let Value::Array(elements) = decode_value(&encoded).unwrap() else { + panic!("expected a msgpack array"); + }; + assert_eq!(elements.len(), FIELD_COUNT); + assert_eq!(elements[6], Value::from(TOP_K_ALL)); // default top_k sentinel + } +} diff --git a/crates/engine_zmq_client/src/protocol/sglang/token_ids.rs b/crates/engine_zmq_client/src/protocol/sglang/token_ids.rs new file mode 100644 index 000000000..408a8cc8e --- /dev/null +++ b/crates/engine_zmq_client/src/protocol/sglang/token_ids.rs @@ -0,0 +1,159 @@ +// SGLang encodes token-id sequences as Python `array.array('q', ...)`, which +// msgspec serializes as the 2-tuple `[typecode, raw_bytes]` — the typecode +// string `"q"` (signed 64-bit) followed by a msgpack `bin` of little-endian +// int64s, not a plain msgpack integer list. This module carries that wire form +// as a newtype over the engine-neutral `Vec` token ids SMG uses. + +use rmpv::Value; +use serde::{ + de::Error as _, ser::SerializeTuple, Deserialize, Deserializer, Serialize, Serializer, +}; + +/// The msgspec typecode string for a Python `array.array('q', ...)`: a signed +/// 64-bit little-endian integer array. Token ids ride the wire under this code. +const TYPECODE_INT64: &str = "q"; + +/// Number of bytes per `int64` element in the raw-bytes half of the tuple. +const INT64_BYTES: usize = 8; + +/// A token-id sequence in SGLang's `array.array('q', ...)` wire form. +/// +/// On the wire it is the 2-element msgpack array `["q", ]`. +/// In memory it is the engine-neutral `Vec` token ids: SMG tokenizes +/// upstream and the engine speaks token ids on the skip-tokenizer path. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct TokenIdArray(pub Vec); + +/// Serialize `&[u8]` as a msgpack `bin` (serde's default treats it as a `u8` +/// sequence, which would emit an integer array instead of the raw buffer). +struct RawBytes<'a>(&'a [u8]); + +impl Serialize for RawBytes<'_> { + fn serialize(&self, serializer: S) -> Result { + serializer.serialize_bytes(self.0) + } +} + +impl Serialize for TokenIdArray { + fn serialize(&self, serializer: S) -> Result { + let mut raw = Vec::with_capacity(self.0.len() * INT64_BYTES); + for &id in &self.0 { + raw.extend_from_slice(&i64::from(id).to_le_bytes()); + } + let mut tuple = serializer.serialize_tuple(2)?; + tuple.serialize_element(TYPECODE_INT64)?; + tuple.serialize_element(&RawBytes(&raw))?; + tuple.end() + } +} + +impl<'de> Deserialize<'de> for TokenIdArray { + fn deserialize>(deserializer: D) -> Result { + // Decode through a dynamic value: the second element is a msgpack `bin`, + // which serde cannot map onto a typed field without `serde_bytes`. + let value = Value::deserialize(deserializer)?; + let Value::Array(items) = value else { + return Err(D::Error::custom(format!( + "expected a 2-element [typecode, bytes] token-id array, got {value:?}" + ))); + }; + let [code, bytes] = items.as_slice() else { + return Err(D::Error::custom(format!( + "expected exactly 2 elements in a token-id array, got {}", + items.len() + ))); + }; + let code = code.as_str().ok_or_else(|| { + D::Error::custom(format!("token-id array typecode is not a string: {code:?}")) + })?; + if code != TYPECODE_INT64 { + return Err(D::Error::custom(format!( + "unsupported token-id array typecode `{code}`, expected `{TYPECODE_INT64}`" + ))); + } + let Value::Binary(raw) = bytes else { + return Err(D::Error::custom(format!( + "token-id array payload is not a msgpack bin: {bytes:?}" + ))); + }; + if raw.len() % INT64_BYTES != 0 { + return Err(D::Error::custom(format!( + "token-id array byte length {} is not a multiple of {INT64_BYTES}", + raw.len() + ))); + } + raw.chunks_exact(INT64_BYTES) + .map(|chunk| { + let mut buf = [0u8; INT64_BYTES]; + buf.copy_from_slice(chunk); + let id = i64::from_le_bytes(buf); + u32::try_from(id) + .map_err(|_| D::Error::custom(format!("token id {id} is out of range for u32"))) + }) + .collect::, D::Error>>() + .map(TokenIdArray) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::codec::{decode_msgpack, decode_value, encode_msgpack}; + + #[test] + fn token_id_array_encodes_as_typed_bin_tuple() { + // The golden wire bytes for [9906, 11, 1917] captured from the Python + // msgspec encoder: 2-array ["q", bin24 of three little-endian int64s]. + let encoded = encode_msgpack(&TokenIdArray(vec![9906, 11, 1917])).unwrap(); + let expected = "92a171c418b2260000000000000b000000000000007d07000000000000"; + assert_eq!(hex(&encoded), expected); + + let Value::Array(items) = decode_value(&encoded).unwrap() else { + panic!("expected a 2-element array"); + }; + assert_eq!(items.len(), 2); + assert_eq!(items[0], Value::from("q")); + } + + #[test] + fn token_id_array_roundtrips() { + for ids in [ + vec![], + vec![0], + vec![1, 2, 3], + vec![9906, 11, 1917, u32::MAX], + ] { + let encoded = encode_msgpack(&TokenIdArray(ids.clone())).unwrap(); + assert_eq!( + decode_msgpack::(&encoded).unwrap(), + TokenIdArray(ids) + ); + } + } + + #[test] + fn decode_rejects_wrong_typecode() { + // A float32 array ("f") is not a valid token-id sequence. + let value = Value::Array(vec![Value::from("f"), Value::Binary(vec![0; 8])]); + let mut bytes = Vec::new(); + rmpv::encode::write_value(&mut bytes, &value).unwrap(); + let error = decode_msgpack::(&bytes).unwrap_err(); + assert!(error + .to_string() + .contains("unsupported token-id array typecode")); + } + + #[test] + fn decode_rejects_ragged_byte_length() { + // 7 bytes is not a whole number of int64 elements. + let value = Value::Array(vec![Value::from("q"), Value::Binary(vec![0; 7])]); + let mut bytes = Vec::new(); + rmpv::encode::write_value(&mut bytes, &value).unwrap(); + let error = decode_msgpack::(&bytes).unwrap_err(); + assert!(error.to_string().contains("not a multiple of")); + } + + fn hex(bytes: &[u8]) -> String { + bytes.iter().map(|b| format!("{b:02x}")).collect() + } +} diff --git a/crates/engine_zmq_client/src/protocol/tokenspeed/mod.rs b/crates/engine_zmq_client/src/protocol/tokenspeed/mod.rs index 0ef2eed60..3c526a5cd 100644 --- a/crates/engine_zmq_client/src/protocol/tokenspeed/mod.rs +++ b/crates/engine_zmq_client/src/protocol/tokenspeed/mod.rs @@ -23,8 +23,10 @@ pub mod request; pub mod sampling; use bytes::Bytes; -use serde::de::{Deserialize, Error as _, IgnoredAny, SeqAccess}; +// The msgspec `array_like` decode helpers are shared across engine protocols; +// re-exported here so this module's submodules keep importing them unchanged. +pub(crate) use crate::protocol::{drain_trailing, expect_tag, next_field}; use crate::{ codec::{decode_msgpack, encode_msgpack}, error::Result, @@ -37,49 +39,6 @@ use crate::{ }, }; -/// Read the next positional element, failing loudly when the array is shorter -/// than the modeled prefix (every modeled field is required on decode). -pub(crate) fn next_field<'de, A, T>( - seq: &mut A, - name: &'static str, -) -> std::result::Result -where - A: SeqAccess<'de>, - T: Deserialize<'de>, -{ - seq.next_element::()? - .ok_or_else(|| A::Error::custom(format!("missing positional field `{name}`"))) -} - -/// Validate the msgspec tag string at element 0. A wrong tag means the payload -/// is a different message type — fail loudly instead of misreading fields. -pub(crate) fn expect_tag<'de, A>( - seq: &mut A, - expected: &'static str, -) -> std::result::Result<(), A::Error> -where - A: SeqAccess<'de>, -{ - let tag: String = next_field(seq, "_tag")?; - if tag != expected { - return Err(A::Error::custom(format!( - "wrong msgspec tag: expected `{expected}`, got `{tag}`" - ))); - } - Ok(()) -} - -/// Drain positional elements beyond the modeled prefix. TokenSpeed appends new -/// fields at the end of its structs, so unknown trailing elements are skipped -/// rather than treated as a decode error. -pub(crate) fn drain_trailing<'de, A>(seq: &mut A) -> std::result::Result<(), A::Error> -where - A: SeqAccess<'de>, -{ - while seq.next_element::()?.is_some() {} - Ok(()) -} - /// The TokenSpeed engine protocol: drives [`TokenizedGenerateReqInput`] over /// the shared ZMQ transport and decodes [`BatchTokenIDOutSlim`] back. pub struct TokenSpeedProtocol; @@ -88,12 +47,12 @@ impl EngineProtocol for TokenSpeedProtocol { type Request = TokenizedGenerateReqInput; type Output = TokenSpeedOutput; - fn add_frame() -> Bytes { - TokenSpeedRequestType::Add.to_frame() + fn add_frame() -> Option { + Some(TokenSpeedRequestType::Add.to_frame()) } - fn abort_frame() -> Bytes { - TokenSpeedRequestType::Abort.to_frame() + fn abort_frame() -> Option { + Some(TokenSpeedRequestType::Abort.to_frame()) } fn request_id(request: &Self::Request) -> &str { @@ -192,8 +151,14 @@ mod tests { #[test] fn request_type_frames_match_wire_contract() { - assert_eq!(TokenSpeedProtocol::add_frame().as_ref(), b"\x00"); - assert_eq!(TokenSpeedProtocol::abort_frame().as_ref(), b"\x01"); + assert_eq!( + TokenSpeedProtocol::add_frame().as_deref(), + Some(&b"\x00"[..]) + ); + assert_eq!( + TokenSpeedProtocol::abort_frame().as_deref(), + Some(&b"\x01"[..]) + ); } #[test] diff --git a/crates/engine_zmq_client/src/protocol/vllm/mod.rs b/crates/engine_zmq_client/src/protocol/vllm/mod.rs index 44d96abdd..2aeecbeac 100644 --- a/crates/engine_zmq_client/src/protocol/vllm/mod.rs +++ b/crates/engine_zmq_client/src/protocol/vllm/mod.rs @@ -66,12 +66,12 @@ impl EngineProtocol for VllmProtocol { type Request = EngineCoreRequest; type Output = EngineCoreOutput; - fn add_frame() -> Bytes { - EngineCoreRequestType::Add.to_frame() + fn add_frame() -> Option { + Some(EngineCoreRequestType::Add.to_frame()) } - fn abort_frame() -> Bytes { - EngineCoreRequestType::Abort.to_frame() + fn abort_frame() -> Option { + Some(EngineCoreRequestType::Abort.to_frame()) } fn request_id(request: &Self::Request) -> &str { diff --git a/crates/engine_zmq_client/src/transport.rs b/crates/engine_zmq_client/src/transport.rs index 8b2599969..5608a1293 100644 --- a/crates/engine_zmq_client/src/transport.rs +++ b/crates/engine_zmq_client/src/transport.rs @@ -16,12 +16,13 @@ use std::{ }; use bytes::Bytes; +use futures::{channel::mpsc as futures_mpsc, StreamExt}; use tokio::{sync::mpsc, time::timeout}; use tracing::{debug, error, info, trace, warn}; use zeromq::{ prelude::{Socket, SocketRecv, SocketSend}, util::PeerIdentity, - PullSocket, RouterSendHalf, RouterSocket, ZmqError, ZmqMessage, + PullSocket, PushSocket, RouterSendHalf, RouterSocket, SocketEvent, ZmqError, ZmqMessage, }; use crate::{ @@ -107,27 +108,57 @@ impl TryFrom for PeerIdentity { } } -/// Per-engine handshake result collected while bootstrapping one shared -/// transport. +/// Per-engine result collected while bootstrapping one shared transport. #[derive(Clone, Debug)] pub struct ConnectedEngine { /// The identity of the connected engine. pub engine_id: EngineId, - /// Post-init configuration received on the input socket registration. - pub ready_response: EngineCoreReadyResponse, + /// Post-init configuration received on the input socket registration, when + /// the topology carries one. Handshake engines (vLLM/TokenSpeed) register a + /// [`EngineCoreReadyResponse`]; a no-handshake PUSH/PULL engine (SGLang) + /// signals readiness out-of-band, so it has none. + pub ready_response: Option, } -/// The connected shared transport plus all registered engines after a -/// successful startup handshake. +/// The frontend's request-sending socket. +/// +/// Handshake engines (vLLM/TokenSpeed) use a ROUTER, where every message leads +/// with the engine identity frame that selects the DP rank. A no-handshake +/// PUSH/PULL engine (SGLang) uses a PUSH into the scheduler's PULL, with no +/// identity frame — the sole engine is implied. +pub enum InputSocket { + /// ROUTER send half: messages lead with the engine identity frame. + Router(RouterSendHalf), + /// PUSH socket: messages carry no routing identity frame. + Push(PushSocket), +} + +impl InputSocket { + /// Whether this input leads each message with an engine identity frame. + fn is_router(&self) -> bool { + matches!(self, Self::Router(_)) + } + + /// Send one already-framed message over the underlying socket. + async fn send(&mut self, message: ZmqMessage) -> Result<()> { + match self { + Self::Router(half) => half.send(message).await?, + Self::Push(socket) => socket.send(message).await?, + } + Ok(()) + } +} + +/// The connected shared transport plus all registered engines. pub struct ConnectedTransport { /// Local address of the shared input socket engines connect to for requests. pub input_address: String, /// Local address of the shared output socket engines connect to for outputs. pub output_address: String, - /// All engines connected through the startup handshake. + /// All engines connected on this transport. pub engines: Vec, - /// The sending half of the shared input ROUTER socket. - pub input_send: RouterSendHalf, + /// The sending half of the shared input socket (ROUTER or PUSH). + pub input_send: InputSocket, /// The shared output PULL socket for receiving all engines' responses. pub output_socket: PullSocket, } @@ -278,11 +309,88 @@ pub async fn connect_handshake( input_address, output_address, engines, - input_send, + input_send: InputSocket::Router(input_send), + output_socket, + }) +} + +/// Connect to a single same-host engine over a no-handshake PUSH/PULL topology. +/// +/// The frontend BINDS a PUSH input socket (the engine's scheduler PULLs +/// requests) and a PULL output socket (the engine PUSHes outputs back). There is +/// no startup handshake and no per-engine registration, so a single engine is +/// synthesized with no post-init config. `ipc://` is supported via explicit +/// addresses; otherwise an ephemeral `tcp://host:0` is used. +/// +/// A bound PUSH with no attached peer fails sends immediately instead of +/// queueing, so this blocks (up to `peer_timeout`) until the engine has dialed +/// into the input socket before returning a usable transport. +pub async fn connect_push_pull( + local_host: &str, + local_input_address: Option<&str>, + local_output_address: Option<&str>, + peer_timeout: Duration, +) -> Result { + let mut input_socket = PushSocket::new(); + // Register the monitor before binding so the peer-attach event cannot be + // missed between bind and the wait below. + let mut input_monitor = input_socket.monitor(); + let input_bind = local_input_address + .map(str::to_owned) + .unwrap_or_else(|| format!("tcp://{local_host}:0")); + let input_address = input_socket.bind(&input_bind).await?.to_string(); + + let mut output_socket = PullSocket::new(); + let output_bind = local_output_address + .map(str::to_owned) + .unwrap_or_else(|| format!("tcp://{local_host}:0")); + let output_address = output_socket.bind(&output_bind).await?.to_string(); + + info!(%input_address, %output_address, "bound local PUSH/PULL transport sockets"); + + wait_for_input_peer(&mut input_monitor, peer_timeout).await?; + info!(%input_address, "engine attached to input socket"); + + Ok(ConnectedTransport { + input_address, + output_address, + // No handshake: synthesize the sole engine with no registration payload. + engines: vec![ConnectedEngine { + engine_id: EngineId::from_engine_index(0), + ready_response: None, + }], + input_send: InputSocket::Push(input_socket), output_socket, }) } +/// Block until a peer attaches to the bound input socket, bounded by +/// `peer_timeout`. The socket monitor emits `Accepted` when a peer connects; +/// other lifecycle events (e.g. `Listening`) precede it and are skipped. +async fn wait_for_input_peer( + monitor: &mut futures_mpsc::Receiver, + peer_timeout: Duration, +) -> Result<()> { + loop { + let event = + timeout(peer_timeout, monitor.next()) + .await + .map_err(|_| Error::HandshakeTimeout { + stage: "input peer", + timeout: peer_timeout, + })?; + match event { + Some(SocketEvent::Accepted(..)) => return Ok(()), + Some(_) => continue, + None => { + return Err(Error::Transport(ZmqError::Socket( + "input socket monitor closed before a peer attached", + ))) + } + } + } +} + /// Bind the shared input (ROUTER) and output (PULL) sockets. `ipc://` is /// supported via explicit addresses; otherwise an ephemeral `tcp://host:0` is used. async fn bind_local_sockets( @@ -396,34 +504,48 @@ async fn wait_for_input_registrations( })?; Ok(ConnectedEngine { engine_id, - ready_response, + ready_response: Some(ready_response), }) }) .collect() } -/// Send an encoded request to one engine over the shared input ROUTER socket. -/// Frames: `[engine_id, request_type, payload, aux..]`. +/// Send an encoded request to one engine over the shared input socket. +/// +/// On a ROUTER input the frames are `[engine_id, request_type?, payload, aux..]`: +/// the leading identity frame selects the engine. On a PUSH input there is no +/// routing identity, so the frames are `[request_type?, payload, aux..]` for the +/// sole connected engine. `request_type` is omitted for protocols that dispatch +/// by payload tag alone (no type frame). pub async fn send_message( - input_send: &mut RouterSendHalf, + input: &mut InputSocket, engine_id: &EngineId, - request_type: Bytes, + request_type: Option, payload: Bytes, aux_frames: Vec, ) -> Result<()> { - let mut message = ZmqMessage::from(engine_id.to_frame()); - message.push_back(request_type); - message.push_back(payload); - for frame in aux_frames { + let mut frames: Vec = Vec::with_capacity(aux_frames.len() + 3); + if input.is_router() { + frames.push(engine_id.to_frame()); + } + if let Some(request_type) = request_type { + frames.push(request_type); + } + frames.push(payload); + frames.extend(aux_frames); + + // `payload` is always pushed, so `frames` is never empty. + let frame_count = frames.len(); + let mut iter = frames.into_iter(); + let Some(first) = iter.next() else { + return Err(unexpected_handshake("cannot send an empty ZMQ message")); + }; + let mut message = ZmqMessage::from(first); + for frame in iter { message.push_back(frame); } - trace!( - ?engine_id, - frame_count = message.len(), - "sending ZMQ message" - ); - input_send.send(message).await?; - Ok(()) + trace!(?engine_id, frame_count, "sending ZMQ message"); + input.send(message).await } /// Receive engine outputs from the shared PULL socket, decode them with the @@ -538,7 +660,11 @@ mod tests { EngineId::from_engine_index(0) ); assert_eq!( - transport.engines[0].ready_response.vllm_version, + transport.engines[0] + .ready_response + .as_ref() + .expect("handshake engine registers a ready response") + .vllm_version, "test-vllm-version" ); @@ -560,7 +686,7 @@ mod tests { send_message( &mut input_send, &engine_id, - EngineCoreRequestType::Add.to_frame(), + Some(EngineCoreRequestType::Add.to_frame()), Bytes::from(payload), Vec::new(), ) diff --git a/crates/mock_worker/src/zmq.rs b/crates/mock_worker/src/zmq.rs index 8c3c81c63..c97b4fcee 100644 --- a/crates/mock_worker/src/zmq.rs +++ b/crates/mock_worker/src/zmq.rs @@ -321,7 +321,14 @@ mod tests { .await .expect("frontend handshake"); let client = EngineCoreClient::new(transport); - assert_eq!(client.engines()[0].ready_response.data_parallel_rank, 0); + assert_eq!( + client.engines()[0] + .ready_response + .as_ref() + .expect("handshake engine registers a ready response") + .data_parallel_rank, + 0 + ); let mut stream = client .submit(EngineCoreRequest { diff --git a/model_gateway/src/main.rs b/model_gateway/src/main.rs index 5df56c4c2..8bd4142ae 100644 --- a/model_gateway/src/main.rs +++ b/model_gateway/src/main.rs @@ -1410,6 +1410,7 @@ impl CliArgs { match self.backend { Some(Backend::Vllm) => Some(RuntimeType::Vllm), Some(Backend::Tokenspeed) => Some(RuntimeType::TokenSpeed), + Some(Backend::Sglang) => Some(RuntimeType::Sglang), _ => None, } } else { diff --git a/model_gateway/src/routers/grpc/zmq_client.rs b/model_gateway/src/routers/grpc/zmq_client.rs index e3fc0b692..60d65308e 100644 --- a/model_gateway/src/routers/grpc/zmq_client.rs +++ b/model_gateway/src/routers/grpc/zmq_client.rs @@ -2,8 +2,8 @@ // // ZMQ backend adapter (gateway glue): presents the vLLM engine surface (the same // proto request/response types as `VllmEngineClient`) but speaks ZMQ directly to -// a same-host engine (vLLM EngineCore or TokenSpeed) via `engine-zmq-client`, -// bypassing the gRPC Python servicer. +// a same-host engine (vLLM EngineCore, TokenSpeed, or SGLang) via +// `engine-zmq-client`, bypassing the gRPC Python servicer. // // This bridges the gateway's proto request-execution pipeline to the raw ZMQ // transport, so it lives with the router (which owns `GrpcClient`/`ProtoStream`), @@ -19,9 +19,18 @@ use std::{ }; use engine_zmq_client::{ - connect_handshake, - connector::{EngineCoreClient, EngineCoreStream, TokenSpeedClient, TokenSpeedStream}, + connect_handshake, connect_push_pull, + connector::{ + EngineCoreClient, EngineCoreStream, SglangClient, SglangStream, TokenSpeedClient, + TokenSpeedStream, + }, protocol::{ + sglang::{ + output::SglangOutput, + request::TokenizedGenerateReqInput as SglangGenerateReqInput, + sampling::{SamplingParams as SglangSamplingParams, TOP_K_ALL as SGLANG_TOP_K_ALL}, + token_ids::TokenIdArray, + }, tokenspeed::{ output::TokenSpeedOutput, request::TokenizedGenerateReqInput, sampling::SamplingParams as TokenSpeedSamplingParams, @@ -45,15 +54,18 @@ use crate::worker::RuntimeType; /// binds). Shared with the worker-side socket derivation. pub(crate) const ZMQ_LOOPBACK_HOST: &str = "127.0.0.1"; -/// The engine protocol a ZMQ backend speaks. Both share the transport and -/// handshake; only the request/output struct shapes and the translation to/from -/// SMG proto differ. +/// The engine protocol a ZMQ backend speaks. vLLM and TokenSpeed share the +/// handshake transport; SGLang uses the no-handshake PUSH/PULL transport. Only +/// the request/output struct shapes and the translation to/from SMG proto +/// differ. #[derive(Clone)] enum ZmqBackend { /// vLLM EngineCore. Vllm(Arc), /// TokenSpeed. TokenSpeed(Arc), + /// SGLang (no handshake; tag-dispatched PUSH/PULL). + Sglang(Arc), } /// Direct ZMQ connection to a same-host engine (vLLM EngineCore or TokenSpeed), @@ -67,13 +79,16 @@ pub struct ZmqEngineClient { } impl ZmqEngineClient { - /// Bind the frontend sockets and complete the handshake with the engine(s), - /// which must already be running and dialing `handshake_address`. + /// Bind the frontend sockets and connect to the engine(s). For the handshake + /// wires (vLLM EngineCore, TokenSpeed) the engines must already be running and + /// dialing `handshake_address`; for SGLang there is no handshake — SMG binds + /// the PUSH/PULL data plane and the scheduler connects in (`handshake_address` + /// is unused). /// /// `input_address`/`output_address` are the `ipc://` data-plane endpoints the /// engines connect to (chosen by SMG). `engine_count` is the number of DP - /// ranks to await. `runtime` selects the wire protocol spoken over the shared - /// transport (vLLM EngineCore vs TokenSpeed). + /// ranks to await. `runtime` selects the wire protocol (vLLM EngineCore, + /// TokenSpeed, or SGLang). pub async fn connect( handshake_address: &str, input_address: &str, @@ -83,31 +98,49 @@ impl ZmqEngineClient { runtime: RuntimeType, timeout: Duration, ) -> Result> { - // Single-engine scope for TokenSpeed: its wire carries no DP-rank routing - // yet (`data_parallel_rank` is always `None`), so more than one engine - // would silently send all traffic to engine 0. Reject it loudly until - // DP>1 lands. The engine count is known here (the handshake awaits it). - if matches!(runtime, RuntimeType::TokenSpeed) && engine_count > 1 { + // Single-engine scope for TokenSpeed and SGLang: their wires carry no + // DP-rank routing yet (`data_parallel_rank` is always `None`), so more + // than one engine would silently send all traffic to engine 0. Reject it + // loudly until DP>1 lands. + if matches!(runtime, RuntimeType::TokenSpeed | RuntimeType::Sglang) && engine_count > 1 { return Err(format!( - "TokenSpeed ZMQ backend supports a single engine only (got \ + "{runtime} ZMQ backend supports a single engine only (got \ engine_count={engine_count}); DP>1 is not yet supported" ) .into()); } // No silent fallback: any other runtime has no ZMQ engine adapter. - // Reject before the handshake — no such engine ever dials in, so the - // handshake would just block for the full timeout. + // Reject before connecting — for the handshake wires no such engine ever + // dials in, so the handshake would just block for the full timeout. if !matches!( runtime, - RuntimeType::Vllm | RuntimeType::TokenSpeed | RuntimeType::Unspecified + RuntimeType::Vllm + | RuntimeType::TokenSpeed + | RuntimeType::Sglang + | RuntimeType::Unspecified ) { return Err(format!( "ZMQ direct backend has no engine implementation for runtime \ - {runtime}; only vllm and tokenspeed are supported" + {runtime}; only vllm, tokenspeed, and sglang are supported" ) .into()); } + // SGLang speaks the no-handshake PUSH/PULL transport: SMG binds the input + // and output sockets and the scheduler connects in. There is no TCP + // handshake, so `handshake_address` is unused for this runtime. + if matches!(runtime, RuntimeType::Sglang) { + let transport = connect_push_pull( + ZMQ_LOOPBACK_HOST, + Some(input_address), + Some(output_address), + timeout, + ) + .await?; + let backend = ZmqBackend::Sglang(Arc::new(SglangClient::new(transport))); + return Ok(Self { backend, model_id }); + } + let transport = connect_handshake( handshake_address, engine_count, @@ -135,15 +168,18 @@ impl ZmqEngineClient { match &self.backend { ZmqBackend::Vllm(_) => RuntimeType::Vllm, ZmqBackend::TokenSpeed(_) => RuntimeType::TokenSpeed, + ZmqBackend::Sglang(_) => RuntimeType::Sglang, } } - /// The engines connected on the shared transport (same handshake for both - /// protocols). + /// The engines connected on the transport. The handshake wires register a + /// per-rank `ready_response`; SGLang's no-handshake wire synthesizes a single + /// engine with none. fn engines(&self) -> &[ConnectedEngine] { match &self.backend { ZmqBackend::Vllm(client) => client.engines(), ZmqBackend::TokenSpeed(client) => client.engines(), + ZmqBackend::Sglang(client) => client.engines(), } } @@ -184,6 +220,16 @@ impl ZmqEngineClient { } Ok(ZmqGenerateStream::TokenSpeed(streams)) } + ZmqBackend::Sglang(client) => { + let mut streams = SelectAll::new(); + for (index, sub) in subs.into_iter().enumerate() { + let request = + translate_request_sglang(sub).map_err(tonic::Status::invalid_argument)?; + let stream = client.submit(request).await.map_err(zmq_status)?; + streams.push(SglangGenerateStream::new(stream, index as u32)); + } + Ok(ZmqGenerateStream::Sglang(streams)) + } } } @@ -193,6 +239,7 @@ impl ZmqEngineClient { match &self.backend { ZmqBackend::Vllm(client) => client.is_alive(), ZmqBackend::TokenSpeed(client) => client.is_alive(), + ZmqBackend::Sglang(client) => client.is_alive(), } } @@ -210,17 +257,20 @@ impl ZmqEngineClient { } /// Latest per-rank load for one engine index, if the backend carries it. - /// vLLM piggybacks it on every batch; TokenSpeed does not (always `None`). + /// vLLM piggybacks it on every batch; TokenSpeed and SGLang do not (always + /// `None`). fn engine_load(&self, engine_index: u32) -> Option { match &self.backend { ZmqBackend::Vllm(client) => client.engine_load(engine_index), ZmqBackend::TokenSpeed(client) => client.engine_load(engine_index), + ZmqBackend::Sglang(client) => client.engine_load(engine_index), } } /// Per-rank load from the piggybacked `scheduler_stats` (SMG's DP routing - /// signal), in the same shape as the gRPC `GetLoads` response. TokenSpeed - /// carries no piggybacked load, so its response has no per-rank entries. + /// signal), in the same shape as the gRPC `GetLoads` response. TokenSpeed and + /// SGLang carry no piggybacked load, so their responses have no per-rank + /// entries. pub fn get_loads(&self) -> WorkerLoadResponse { let loads: Vec = self .engines() @@ -246,12 +296,14 @@ impl ZmqEngineClient { /// Model info derived from the handshake `EngineCoreReadyResponse` plus the /// configured model id (the engine does not report tokenizer/vocab metadata, - /// so those come from worker config). + /// so those come from worker config). SGLang's no-handshake wire has no ready + /// response, so the context length falls back to `0` (unknown). pub fn get_model_info(&self) -> vllm::GetModelInfoResponse { let max_context_length = self .engines() .first() - .map(|e| e.ready_response.max_model_len) + .and_then(|e| e.ready_response.as_ref()) + .map(|ready| ready.max_model_len) .unwrap_or(0); vllm::GetModelInfoResponse { model_path: self.model_id.clone(), @@ -263,16 +315,19 @@ impl ZmqEngineClient { } } - /// Server info derived from the handshake response. + /// Server info derived from the handshake response. SGLang's no-handshake + /// wire has no ready response, so the DP size falls back to `1`. pub fn get_server_info(&self) -> vllm::GetServerInfoResponse { let data_parallel_size = self .engines() .first() - .map(|e| e.ready_response.data_parallel_size) + .and_then(|e| e.ready_response.as_ref()) + .map(|ready| ready.data_parallel_size) .unwrap_or(1); let server_type = match &self.backend { ZmqBackend::Vllm(_) => "vllm", ZmqBackend::TokenSpeed(_) => "tokenspeed", + ZmqBackend::Sglang(_) => "sglang", }; vllm::GetServerInfoResponse { data_parallel_size: i32::try_from(data_parallel_size).unwrap_or(i32::MAX), @@ -295,6 +350,8 @@ pub enum ZmqGenerateStream { Vllm(SelectAll), /// TokenSpeed outputs. TokenSpeed(SelectAll), + /// SGLang outputs. + Sglang(SelectAll), } impl ZmqGenerateStream { @@ -303,6 +360,7 @@ impl ZmqGenerateStream { match self { Self::Vllm(streams) => streams.next().await, Self::TokenSpeed(streams) => streams.next().await, + Self::Sglang(streams) => streams.next().await, } } @@ -476,11 +534,59 @@ impl Stream for VllmGenerateStream { } } -/// Streaming generate output for one TokenSpeed sub-request, mapping each -/// `TokenSpeedOutput` to a vLLM-proto `GenerateResponse`, tagged with this -/// sub's choice `index`. -pub struct TokenSpeedGenerateStream { - inner: TokenSpeedStream, +/// One tick of a "slim batch" backend's per-request output, reduced to the +/// fields the shared stream mapper consumes. TokenSpeed and SGLang both report +/// cumulative per-request token counts and per-tick sampled-token logprobs in +/// this same shape, so one mapper serves both. +struct SlimTick { + prompt_tokens: u32, + cached_tokens: u32, + completion_tokens: u32, + output_ids: Vec, + finish_reason: Option, + output_logprobs_val: Vec, + output_logprobs_idx: Vec, +} + +/// A "slim batch" per-request output (TokenSpeed or SGLang), reducible to the +/// shared [`SlimTick`] the stream mapper consumes. +trait SlimOutput { + fn into_tick(self) -> SlimTick; +} + +impl SlimOutput for TokenSpeedOutput { + fn into_tick(self) -> SlimTick { + SlimTick { + prompt_tokens: self.prompt_tokens, + cached_tokens: self.cached_tokens, + completion_tokens: self.completion_tokens, + output_ids: self.output_ids, + finish_reason: self.finish_reason, + output_logprobs_val: self.output_logprobs_val, + output_logprobs_idx: self.output_logprobs_idx, + } + } +} + +impl SlimOutput for SglangOutput { + fn into_tick(self) -> SlimTick { + SlimTick { + prompt_tokens: self.prompt_tokens, + cached_tokens: self.cached_tokens, + completion_tokens: self.completion_tokens, + output_ids: self.output_ids, + finish_reason: self.finish_reason, + output_logprobs_val: self.output_logprobs_val, + output_logprobs_idx: self.output_logprobs_idx, + } + } +} + +/// Streaming generate output for one sub-request of a "slim batch" backend +/// (TokenSpeed or SGLang), mapping each output tick to a vLLM-proto +/// `GenerateResponse`, tagged with this sub's choice `index`. +pub struct SlimGenerateStream { + inner: S, state: StreamState, /// Choice index stamped on every chunk/complete (0 for n=1; the fan-out /// position for n>1) — the proto field the pipeline demuxes choices by. @@ -491,8 +597,8 @@ pub struct TokenSpeedGenerateStream { pending: Option, } -impl TokenSpeedGenerateStream { - fn new(inner: TokenSpeedStream, index: u32) -> Self { +impl SlimGenerateStream { + fn new(inner: S, index: u32) -> Self { Self { inner, state: StreamState::default(), @@ -500,91 +606,104 @@ impl TokenSpeedGenerateStream { pending: None, } } +} - fn map_output(&mut self, output: TokenSpeedOutput) -> vllm::GenerateResponse { - let state = &mut self.state; - // TokenSpeed reports per-request token counts directly (cumulative for - // completions), rather than vLLM's per-output prefill-stats deltas. - if output.prompt_tokens > 0 { - state.prompt_tokens = output.prompt_tokens; - } - if output.cached_tokens > 0 { - state.cached_tokens = output.cached_tokens; - } - state.completion_tokens = output.completion_tokens; - state.output_ids.extend(output.output_ids.iter().copied()); - - // Sampled-token logprobs, if requested. The proto column is `float`, so - // downcast the wire's `f64` values. Chunks carry this tick's increment; - // the terminal `Complete` carries the cumulative set, so accumulate - // into `state` and drain it on finish. - let chunk_logprobs = - (!output.output_logprobs_val.is_empty()).then(|| vllm::OutputLogProbs { - token_logprobs: output - .output_logprobs_val - .iter() - .map(|&lp| lp as f32) - .collect(), - token_ids: output.output_logprobs_idx.clone(), - ..Default::default() - }); - state - .output_logprobs_val - .extend(output.output_logprobs_val.iter().map(|&lp| lp as f32)); - state - .output_logprobs_idx - .extend(output.output_logprobs_idx.iter().copied()); +/// TokenSpeed's per-sub output stream. +pub type TokenSpeedGenerateStream = SlimGenerateStream; +/// SGLang's per-sub output stream. +pub type SglangGenerateStream = SlimGenerateStream; - let response = match output.finish_reason { - Some(reason) => { - let complete = vllm::GenerateResponse { - response: Some(vllm::generate_response::Response::Complete( - vllm::GenerateComplete { - output_ids: std::mem::take(&mut state.output_ids), - finish_reason: normalize_finish_reason(&reason).to_string(), - prompt_tokens: state.prompt_tokens, - completion_tokens: state.completion_tokens, - cached_tokens: state.cached_tokens, - output_logprobs: state.take_complete_logprobs(), - index: self.index, - ..Default::default() - }, - )), - }; - if output.output_ids.is_empty() { - return complete; - } - // The finish tick carried new tokens: emit them as a `Chunk` - // first and hold the `Complete` for the next poll. - let chunk = vllm::generate_response::Response::Chunk(vllm::GenerateStreamChunk { - token_ids: output.output_ids, - prompt_tokens: state.prompt_tokens, - completion_tokens: state.completion_tokens, - cached_tokens: state.cached_tokens, - output_logprobs: chunk_logprobs, - index: self.index, - ..Default::default() - }); - self.pending = Some(complete); - chunk +/// Map one "slim batch" output tick to a vLLM-proto response. These backends +/// report per-request token counts directly (cumulative for completions), rather +/// than vLLM's per-output prefill-stats deltas. +fn map_slim_tick( + state: &mut StreamState, + index: u32, + pending: &mut Option, + tick: SlimTick, +) -> vllm::GenerateResponse { + if tick.prompt_tokens > 0 { + state.prompt_tokens = tick.prompt_tokens; + } + if tick.cached_tokens > 0 { + state.cached_tokens = tick.cached_tokens; + } + state.completion_tokens = tick.completion_tokens; + state.output_ids.extend(tick.output_ids.iter().copied()); + + // Sampled-token logprobs, if requested. The proto column is `float`, so + // downcast the wire's `f64` values. Chunks carry this tick's increment; the + // terminal `Complete` carries the cumulative set, so accumulate into `state` + // and drain it on finish. + let chunk_logprobs = (!tick.output_logprobs_val.is_empty()).then(|| vllm::OutputLogProbs { + token_logprobs: tick + .output_logprobs_val + .iter() + .map(|&lp| lp as f32) + .collect(), + token_ids: tick.output_logprobs_idx.clone(), + ..Default::default() + }); + state + .output_logprobs_val + .extend(tick.output_logprobs_val.iter().map(|&lp| lp as f32)); + state + .output_logprobs_idx + .extend(tick.output_logprobs_idx.iter().copied()); + + let response = match tick.finish_reason { + Some(reason) => { + let complete = vllm::GenerateResponse { + response: Some(vllm::generate_response::Response::Complete( + vllm::GenerateComplete { + output_ids: std::mem::take(&mut state.output_ids), + finish_reason: normalize_finish_reason(&reason).to_string(), + prompt_tokens: state.prompt_tokens, + completion_tokens: state.completion_tokens, + cached_tokens: state.cached_tokens, + output_logprobs: state.take_complete_logprobs(), + index, + ..Default::default() + }, + )), + }; + if tick.output_ids.is_empty() { + return complete; } - None => vllm::generate_response::Response::Chunk(vllm::GenerateStreamChunk { - token_ids: output.output_ids, + // The finish tick carried new tokens: emit them as a `Chunk` first + // and hold the `Complete` for the next poll. + let chunk = vllm::generate_response::Response::Chunk(vllm::GenerateStreamChunk { + token_ids: tick.output_ids, prompt_tokens: state.prompt_tokens, completion_tokens: state.completion_tokens, cached_tokens: state.cached_tokens, output_logprobs: chunk_logprobs, - index: self.index, + index, ..Default::default() - }), - }; - vllm::GenerateResponse { - response: Some(response), + }); + *pending = Some(complete); + chunk } + None => vllm::generate_response::Response::Chunk(vllm::GenerateStreamChunk { + token_ids: tick.output_ids, + prompt_tokens: state.prompt_tokens, + completion_tokens: state.completion_tokens, + cached_tokens: state.cached_tokens, + output_logprobs: chunk_logprobs, + index, + ..Default::default() + }), + }; + vllm::GenerateResponse { + response: Some(response), } } -impl Stream for TokenSpeedGenerateStream { +impl Stream for SlimGenerateStream +where + S: Stream> + Unpin, + O: SlimOutput, +{ type Item = Result; fn poll_next( @@ -597,7 +716,12 @@ impl Stream for TokenSpeedGenerateStream { return Poll::Ready(Some(Ok(pending))); } match std::pin::Pin::new(&mut this.inner).poll_next(cx) { - Poll::Ready(Some(Ok(output))) => Poll::Ready(Some(Ok(this.map_output(output)))), + Poll::Ready(Some(Ok(output))) => Poll::Ready(Some(Ok(map_slim_tick( + &mut this.state, + this.index, + &mut this.pending, + output.into_tick(), + )))), Poll::Ready(Some(Err(error))) => Poll::Ready(Some(Err(zmq_status(error)))), Poll::Ready(None) => Poll::Ready(None), Poll::Pending => Poll::Pending, @@ -769,6 +893,156 @@ fn translate_sampling_tokenspeed(sp: vllm::SamplingParams) -> TokenSpeedSampling params } +/// Translate a vLLM-proto generate request into an SGLang +/// `TokenizedGenerateReqInput`. ZMQ mode requires pre-tokenized input (SMG +/// tokenizes upstream). +fn translate_request_sglang(req: vllm::GenerateRequest) -> Result { + let input_ids = match req.input { + Some(vllm::generate_request::Input::Tokenized(tokenized)) => tokenized.input_ids, + Some(vllm::generate_request::Input::Text(_)) => { + return Err("ZMQ mode requires pre-tokenized input (TokenizedInput)".to_string()); + } + None => { + return Err("ZMQ mode requires pre-tokenized input; no input provided".to_string()); + } + }; + let stream = req.stream; + // Single-engine SGLang: a pinned DP rank other than 0 cannot be honored (the + // tag-dispatched wire carries no DP-rank routing). + if req.data_parallel_rank.is_some_and(|rank| rank != 0) { + return Err(format!( + "invalid data_parallel_rank {:?}: the SGLang ZMQ backend is single-engine", + req.data_parallel_rank + )); + } + if let Some(sp) = req.sampling_params.as_ref() { + // The output wire (`BatchTokenIDOutput`) materializes only the sampled + // token's logprob, so a top-k request (count above 1, or `-1` = "all") + // cannot be honored — reject it rather than silently return fewer. Counts + // of 0 or 1 are the plain sampled-token logprob and are wired end-to-end. + if sp.logprobs.is_some_and(|n| !(0..=1).contains(&n)) { + return Err("top_logprobs are not supported over the SGLang ZMQ backend".to_string()); + } + // Prompt logprobs are not decoded from the output wire. + if sp.prompt_logprobs.is_some() { + return Err( + "prompt logprobs are not supported over the SGLang ZMQ backend".to_string(), + ); + } + // The response_format / forced-tool-choice constraint oneof is not + // translated onto the SGLang structured-output fields yet; dropping it + // would return unconstrained text. + if sp.constraint.is_some() { + return Err( + "structured output constraints are not supported over the ZMQ backend yet" + .to_string(), + ); + } + // Stop strings are not forwarded: skip-tokenizer-init mode runs the + // scheduler without a tokenizer, so it cannot match stop strings, and the + // gateway's stop decoder does not enforce them — they would be ignored and + // the request would run to max_tokens. + if !sp.stop.is_empty() { + return Err( + "stop strings are not supported over the SGLang ZMQ backend yet; \ + use stop_token_ids" + .to_string(), + ); + } + // logit_bias is not translated onto the SGLang wire. + if !sp.logit_bias.is_empty() { + return Err("logit_bias is not supported over the SGLang ZMQ backend".to_string()); + } + } + let return_logprob = req + .sampling_params + .as_ref() + .is_some_and(|sp| sp.logprobs.is_some()); + Ok(SglangGenerateReqInput { + rid: req.request_id, + input_ids: TokenIdArray(input_ids), + sampling_params: req + .sampling_params + .map(translate_sampling_sglang) + .unwrap_or_else(sglang_default_sampling), + return_logprob, + stream, + // The remaining fields keep their neutral defaults; SMG owns the + // tokenizer on this path, so the sampling params it sends are already + // normalized (see `translate_sampling_sglang`). + ..SglangGenerateReqInput::default() + }) +} + +/// SGLang treats a temperature this close to zero as a request for greedy +/// (argmax) decoding. Mirrors `_SAMPLING_EPS` in SGLang's sampling params. +const SGLANG_SAMPLING_EPS: f64 = 1e-6; + +/// SGLang's own defaults, presented already-normalized (see +/// [`translate_sampling_sglang`] for why SMG normalizes). Used for requests that +/// carry no sampling params at all. +fn sglang_default_sampling() -> SglangSamplingParams { + SglangSamplingParams { + // The direct path skips the scheduler's tokenizer-manager normalization, + // so present the tokenizer-derived stop fields as an empty (resolved) + // list rather than the unresolved `None`, and mark the struct normalized. + stop_strs: Some(Vec::new()), + stop_regex_strs: Some(Vec::new()), + is_normalized: true, + ..SglangSamplingParams::default() + } +} + +/// Map vLLM-proto sampling params onto SGLang's native `SamplingParams`. The +/// engine struct rides the wire as a full positional array, so every field +/// carries a value; the ones SMG does not drive keep their SGLang defaults (via +/// `..default()`). The gateway's vLLM builder already resolves the proto defaults +/// (temperature/top_p → 1.0), so the resolved values are forwarded. `n > 1` is +/// fanned out before translation, so `n` is left at its default of 1 +/// (single-sample on the wire). +/// +/// SMG owns the tokenizer on this direct path, so the scheduler never runs its +/// tokenizer-manager normalization (`SamplingParams.normalize`). SMG therefore +/// presents params the scheduler can consume as-is: the tokenizer-derived stop +/// lists are emitted resolved (empty — SMG applies stop strings upstream), +/// `is_normalized` is set, and greedy decoding is resolved here rather than in +/// the engine's skipped `__post_init__`. +fn translate_sampling_sglang(sp: vllm::SamplingParams) -> SglangSamplingParams { + let mut temperature = f64::from(sp.temperature.unwrap_or(1.0)); + // vLLM proto uses `0` for "all tokens"; forward SGLang's own "all tokens" + // sentinel in that case, otherwise the explicit positive cutoff. + let mut top_k = if sp.top_k == 0 { + SGLANG_TOP_K_ALL + } else { + i32::try_from(sp.top_k).unwrap_or(i32::MAX) + }; + // Greedy: SGLang collapses a ~zero temperature to argmax (temperature 1.0, + // top_k 1). Normally done in `__post_init__`, which the normalized struct + // skips, so resolve it here. + if (0.0..SGLANG_SAMPLING_EPS).contains(&temperature) { + temperature = 1.0; + top_k = 1; + } + SglangSamplingParams { + max_new_tokens: sp.max_tokens, + stop_token_ids: (!sp.stop_token_ids.is_empty()).then_some(sp.stop_token_ids), + temperature, + top_p: f64::from(sp.top_p), + top_k, + min_p: f64::from(sp.min_p), + frequency_penalty: f64::from(sp.frequency_penalty), + presence_penalty: f64::from(sp.presence_penalty), + repetition_penalty: f64::from(sp.repetition_penalty), + // Proto `0` is "no floor" — which is also the SGLang default. + min_new_tokens: sp.min_tokens, + ignore_eos: sp.ignore_eos, + // A negative seed is a "no seed" sentinel; drop it (the engine then + // derives one so all ranks agree). + sampling_seed: sp.seed.and_then(|seed| u64::try_from(seed).ok()), + ..sglang_default_sampling() + } +} + /// Translate a vLLM-proto generate request into an `EngineCoreRequest`. ZMQ mode /// requires pre-tokenized input (SMG tokenizes upstream). fn translate_request(req: vllm::GenerateRequest) -> Result { @@ -868,11 +1142,12 @@ fn finish_reason_str(reason: EngineCoreFinishReason) -> &'static str { } } -/// Normalize a TokenSpeed wire finish-reason string into the canonical set the -/// gateway's response layer exact-matches (`stop`, `length`, `abort`, `error`) — -/// the same set the vLLM path emits via [`finish_reason_str`]. TokenSpeed emits -/// `stop`/`length`/`abort`; an unknown value falls back to `stop` with a warning -/// so a non-canonical string never mis-renders downstream. +/// Normalize a "slim batch" (TokenSpeed / SGLang) wire finish-reason string into +/// the canonical set the gateway's response layer exact-matches (`stop`, +/// `length`, `abort`, `error`) — the same set the vLLM path emits via +/// [`finish_reason_str`]. Both engines emit `stop`/`length`/`abort`; an unknown +/// value falls back to `stop` with a warning so a non-canonical string never +/// mis-renders downstream. fn normalize_finish_reason(reason: &str) -> &'static str { match reason { "stop" => "stop", @@ -882,7 +1157,7 @@ fn normalize_finish_reason(reason: &str) -> &'static str { other => { tracing::warn!( finish_reason = other, - "unknown TokenSpeed finish_reason; defaulting to \"stop\"" + "unknown ZMQ engine finish_reason; defaulting to \"stop\"" ); "stop" } @@ -1219,6 +1494,142 @@ mod tests { engine_task.await.unwrap(); } + /// End-to-end over ipc:// for an SGLang backend on the no-handshake + /// PUSH/PULL topology: the adapter frames a tag-dispatched + /// `TokenizedGenerateReqInput` (no identity or type frame), and maps + /// `BatchTokenIDOutput` batches back to vLLM-proto responses. The mock + /// scheduler replies with the pinned Python wire vectors, exercising the real + /// bytes end to end (the output struct is decode-only, so it cannot be + /// re-encoded in-process). + #[tokio::test] + async fn generate_e2e_translates_and_streams_sglang() { + use engine_zmq_client::{ + codec::decode_msgpack, mock_engine::MockPushPullEngine, + protocol::sglang::request::TokenizedGenerateReqInput as WireReq, + }; + + // Pinned Python `BatchTokenIDOutput` vectors for rid "req-00000001", + // output_ids [[15496]] — still-generating (finish nil) then finished + // (`{"type":"stop"}`). Captured in `protocol::sglang::output` tests. + const OUTPUT_STILL_GENERATING: &str = + "dc0032b24261746368546f6b656e49444f757470757491ac7265712d3030303030303031c091\ + c091a09192a171c408883c00000000000091009192a171c408883c00000000000091c391c391\ + c29103910091019100c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0\ + c0c0c0c0c0c0"; + const OUTPUT_FINISHED: &str = + "dc0032b24261746368546f6b656e49444f757470757491ac7265712d3030303030303031c091\ + 82a474797065a473746f70a76d6174636865640291a09192a171c408883c0000000000009100\ + 9192a171c408883c00000000000091c391c391c29103910091019100c0c0c0c0c0c0c0c0c0c0\ + c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0"; + + fn from_hex(hex: &str) -> bytes::Bytes { + let hex: String = hex.chars().filter(|c| !c.is_whitespace()).collect(); + bytes::Bytes::from( + (0..hex.len()) + .step_by(2) + .map(|i| u8::from_str_radix(&hex[i..i + 2], 16).unwrap()) + .collect::>(), + ) + } + + let dir = tempfile::tempdir().unwrap(); + let ep = |name: &str| format!("ipc://{}", dir.path().join(name).display()); + // SGLang binds no handshake socket; the adapter ignores the handshake + // argument on this runtime and binds only the PUSH input + PULL output. + let (input, output) = (ep("in.sock"), ep("out.sock")); + + let (client, engine) = tokio::join!( + ZmqEngineClient::connect( + "ipc:///unused", + &input, + &output, + 1, + "m".to_string(), + RuntimeType::Sglang, + Duration::from_secs(10) + ), + MockPushPullEngine::connect(&input, &output), + ); + let client = client.expect("adapter connect"); + let mut engine = engine.expect("mock engine"); + + #[expect( + clippy::disallowed_methods, + reason = "engine task ends after responding" + )] + let engine_task = tokio::spawn(async move { + // The scheduler sees a single-frame payload: no identity, no type frame. + let frames = engine.recv_request().await.unwrap(); + assert_eq!(frames.len(), 1); + let request: WireReq = decode_msgpack(frames[0].as_ref()).unwrap(); + assert_eq!(request.rid, "req-00000001"); + assert_eq!(request.input_ids.0, vec![1, 2, 3]); + assert_eq!(request.sampling_params.max_new_tokens, Some(2)); + // No logprobs requested -> the flag stays false. + assert!(!request.return_logprob); + + engine + .send_output(vec![from_hex(OUTPUT_STILL_GENERATING)]) + .await + .unwrap(); + engine + .send_output(vec![from_hex(OUTPUT_FINISHED)]) + .await + .unwrap(); + }); + + let req = vllm::GenerateRequest { + request_id: "req-00000001".to_string(), + input: Some(vllm::generate_request::Input::Tokenized( + vllm::TokenizedInput { + original_text: String::new(), + input_ids: vec![1, 2, 3], + }, + )), + sampling_params: Some(vllm::SamplingParams { + max_tokens: Some(2), + ..Default::default() + }), + stream: true, + ..Default::default() + }; + let mut stream = client.generate(req).await.expect("generate"); + + let first = stream.next().await.expect("chunk item").expect("chunk ok"); + match first.response { + Some(vllm::generate_response::Response::Chunk(chunk)) => { + assert_eq!(chunk.token_ids, vec![15496]); + assert_eq!(chunk.prompt_tokens, 3); + } + other => panic!("expected chunk, got {other:?}"), + } + // The finish tick carried a new token, so its delta is emitted as a + // chunk before the (cumulative) terminal complete. + let second = stream.next().await.expect("chunk item").expect("chunk ok"); + match second.response { + Some(vllm::generate_response::Response::Chunk(chunk)) => { + assert_eq!(chunk.token_ids, vec![15496]); + } + other => panic!("expected chunk, got {other:?}"), + } + let third = stream + .next() + .await + .expect("complete item") + .expect("complete ok"); + match third.response { + Some(vllm::generate_response::Response::Complete(complete)) => { + assert_eq!(complete.output_ids, vec![15496, 15496]); + assert_eq!(complete.finish_reason, "stop"); + assert_eq!(complete.prompt_tokens, 3); + } + other => panic!("expected complete, got {other:?}"), + } + assert!(stream.next().await.is_none()); + + engine_task.await.unwrap(); + } + #[test] fn finish_reasons_map_to_vllm_strings() { assert_eq!(finish_reason_str(EngineCoreFinishReason::Length), "length"); @@ -1370,6 +1781,159 @@ mod tests { assert!(translate_request_tokenspeed(req).is_ok()); } + #[test] + fn sglang_sampling_maps_defaults_and_seed() { + // Proto top_k=0 ("all tokens") maps to the engine's own "all" sentinel; + // a negative seed is dropped. + let mapped = translate_sampling_sglang(vllm::SamplingParams { + top_k: 0, + seed: Some(-1), + max_tokens: Some(8), + ..Default::default() + }); + assert_eq!(mapped.top_k, SGLANG_TOP_K_ALL); + assert_eq!(mapped.sampling_seed, None); + assert_eq!(mapped.max_new_tokens, Some(8)); + // Unset temperature resolves to the proto default 1.0 (forwarded). + assert_eq!(mapped.temperature, 1.0); + // SMG owns the tokenizer, so it presents already-normalized params: the + // tokenizer-derived stop lists are resolved (empty) and the struct is + // flagged normalized so the scheduler consumes it as-is. + assert!(mapped.is_normalized); + assert_eq!(mapped.stop_strs, Some(Vec::new())); + assert_eq!(mapped.stop_regex_strs, Some(Vec::new())); + + // An explicit positive cutoff and non-negative seed ride through. + let mapped = translate_sampling_sglang(vllm::SamplingParams { + top_k: 40, + seed: Some(7), + ..Default::default() + }); + assert_eq!(mapped.top_k, 40); + assert_eq!(mapped.sampling_seed, Some(7)); + + // Empty stop_token_ids ride as None; min_tokens=0 is the default floor. + let mapped = translate_sampling_sglang(vllm::SamplingParams::default()); + assert_eq!(mapped.stop_token_ids, None); + assert_eq!(mapped.min_new_tokens, 0); + + // A real min_tokens floor is forwarded. + let mapped = translate_sampling_sglang(vllm::SamplingParams { + min_tokens: 3, + ..Default::default() + }); + assert_eq!(mapped.min_new_tokens, 3); + } + + #[test] + fn sglang_sampling_resolves_greedy_decoding() { + // A ~zero temperature is greedy: SGLang expects temperature 1.0 + top_k 1 + // (argmax). The normalized struct skips the engine's `__post_init__`, so + // the translation resolves it here. + let mapped = translate_sampling_sglang(vllm::SamplingParams { + temperature: Some(0.0), + top_k: 50, + ..Default::default() + }); + assert_eq!(mapped.temperature, 1.0); + assert_eq!(mapped.top_k, 1); + + // A normal temperature leaves the explicit cutoff untouched (the proto + // field is f32, so compare against the widened value). + let mapped = translate_sampling_sglang(vllm::SamplingParams { + temperature: Some(0.7), + top_k: 50, + ..Default::default() + }); + assert_eq!(mapped.temperature, f64::from(0.7_f32)); + assert_eq!(mapped.top_k, 50); + } + + #[test] + fn sglang_plain_logprobs_set_return_logprob() { + // Counts 0 and 1 are the plain sampled-token case; both accepted. + for count in [0, 1] { + let req = translate_request_sglang(tokenized_req(vllm::SamplingParams { + logprobs: Some(count), + ..Default::default() + })) + .expect("plain logprobs accepted"); + assert!(req.return_logprob); + } + + // No logprobs -> the flag stays false. + let req = translate_request_sglang(tokenized_req(vllm::SamplingParams::default())) + .expect("no logprobs accepted"); + assert!(!req.return_logprob); + } + + #[test] + fn sglang_rejects_top_logprobs_and_prompt_logprobs() { + // Top-k logprobs (count > 1) cannot be materialized from the output wire. + assert!( + translate_request_sglang(tokenized_req(vllm::SamplingParams { + logprobs: Some(5), + ..Default::default() + })) + .is_err() + ); + // "all" (count -1) cannot be honored. + assert!( + translate_request_sglang(tokenized_req(vllm::SamplingParams { + logprobs: Some(-1), + ..Default::default() + })) + .is_err() + ); + // Prompt logprobs are not decoded from the output wire. + assert!( + translate_request_sglang(tokenized_req(vllm::SamplingParams { + prompt_logprobs: Some(1), + ..Default::default() + })) + .is_err() + ); + } + + #[test] + fn sglang_rejects_unsupported_sampling_features() { + // Structured-output constraints are not translated onto the wire. + let err = translate_request_sglang(tokenized_req(vllm::SamplingParams { + constraint: Some(vllm::sampling_params::Constraint::JsonObject(true)), + ..Default::default() + })) + .expect_err("constraint rejected"); + assert!(err.contains("structured output"), "{err}"); + + // Stop strings cannot be matched without a tokenizer on this backend. + let err = translate_request_sglang(tokenized_req(vllm::SamplingParams { + stop: vec!["".to_string()], + ..Default::default() + })) + .expect_err("stop strings rejected"); + assert!(err.contains("stop_token_ids"), "{err}"); + + // logit_bias has no wire slot. + let err = translate_request_sglang(tokenized_req(vllm::SamplingParams { + logit_bias: HashMap::from([(7, 1.0)]), + ..Default::default() + })) + .expect_err("logit_bias rejected"); + assert!(err.contains("logit_bias"), "{err}"); + } + + #[test] + fn sglang_rejects_nonzero_dp_rank() { + // Single-engine backend: only rank 0 (or none) is valid. + let mut req = tokenized_req(vllm::SamplingParams::default()); + req.data_parallel_rank = Some(1); + assert!(translate_request_sglang(req).is_err()); + + let mut req = tokenized_req(vllm::SamplingParams::default()); + req.data_parallel_rank = Some(0); + assert!(translate_request_sglang(req).is_ok()); + } + #[test] fn vllm_rejects_unsupported_sampling_features() { // Structured-output constraints are not translated onto the wire. diff --git a/model_gateway/src/worker/worker.rs b/model_gateway/src/worker/worker.rs index a117fc545..c43ac9271 100644 --- a/model_gateway/src/worker/worker.rs +++ b/model_gateway/src/worker/worker.rs @@ -187,6 +187,42 @@ async fn ensure_ipc_socket_dir(base_url: &str) -> WorkerResult<()> { Ok(()) } +/// Remove an `ipc://` socket file stranded by a crashed predecessor. libzmq +/// refuses to `bind()` over an existing ipc endpoint file (it fails with +/// `EADDRINUSE`), so a socket left behind by an earlier process that did not +/// unbind cleanly would block every future bind at the same path. Only an actual +/// socket file is removed: a regular file or directory at the path is left in +/// place and surfaced as an error, so an unexpected collision is never silently +/// clobbered. A missing path is a no-op. +/// +/// This targets the crashed-predecessor case. In the one-worker-per-rank model a +/// given ipc path has a single binder, so this does not race a live peer. +async fn remove_stale_ipc_socket(ipc_url: &str) -> WorkerResult<()> { + let path = ipc_url.strip_prefix("ipc://").unwrap_or(ipc_url); + let fail = |reason: String| WorkerError::ConnectionFailed { + url: ipc_url.to_string(), + reason, + }; + let meta = match tokio::fs::symlink_metadata(path).await { + Ok(meta) => meta, + Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(()), + Err(e) => return Err(fail(format!("failed to stat ipc socket {path}: {e}"))), + }; + #[cfg(unix)] + { + use std::os::unix::fs::FileTypeExt; + if !meta.file_type().is_socket() { + return Err(fail(format!( + "ipc socket path {path} exists but is not a socket; refusing to remove it" + ))); + } + } + let _ = &meta; + tokio::fs::remove_file(path) + .await + .map_err(|e| fail(format!("failed to remove stale ipc socket {path}: {e}"))) +} + /// Bind the SMG-side ZMQ sockets and complete the handshake with the engine. /// Shared by the lazy client accessor and the background handshake driver so /// both go through the exact same connect path. `base_url` is the `ipc://` URL; @@ -200,6 +236,11 @@ async fn connect_zmq_backend( let (handshake, input, output) = zmq_socket_addresses(&base_url, handshake_override.as_deref())?; ensure_ipc_socket_dir(&base_url).await?; + // Clear the data-plane sockets a crashed predecessor may have stranded; + // libzmq will not bind over an existing ipc endpoint file. The handshake is + // tcp://, so it needs no such cleanup. + remove_stale_ipc_socket(&input).await?; + remove_stale_ipc_socket(&output).await?; tracing::info!("Binding ZMQ client for worker {base_url} (handshake={handshake})"); match ZmqEngineClient::connect( &handshake, @@ -2804,6 +2845,45 @@ mod tests { assert!(ensure_ipc_socket_dir(&url).await.is_err()); } + #[tokio::test] + async fn remove_stale_ipc_socket_is_a_noop_when_absent() { + let base = tempfile::tempdir().unwrap(); + let url = format!("ipc://{}/gone.sock", base.path().display()); + // A path that never existed is fine — the first launch has nothing to + // clean up. + remove_stale_ipc_socket(&url).await.unwrap(); + } + + #[cfg(unix)] + #[tokio::test] + async fn remove_stale_ipc_socket_removes_a_stranded_socket() { + use std::os::unix::net::UnixListener; + + let base = tempfile::tempdir().unwrap(); + let path = base.path().join("engine-in.sock"); + // Binding a unix listener creates a real socket file, standing in for + // one a crashed predecessor left behind. + let listener = UnixListener::bind(&path).unwrap(); + drop(listener); // closing the listener does not unlink the file + assert!(path.exists(), "socket file should linger after close"); + + let url = format!("ipc://{}", path.display()); + remove_stale_ipc_socket(&url).await.unwrap(); + assert!(!path.exists(), "stale socket must be removed before bind"); + } + + #[cfg(unix)] + #[tokio::test] + async fn remove_stale_ipc_socket_refuses_a_non_socket() { + let base = tempfile::tempdir().unwrap(); + let path = base.path().join("engine-in.sock"); + // A regular file at the socket path is unexpected — never clobber it. + std::fs::write(&path, b"not a socket").unwrap(); + let url = format!("ipc://{}", path.display()); + assert!(remove_stale_ipc_socket(&url).await.is_err()); + assert!(path.exists(), "a non-socket path must be left untouched"); + } + /// A ZMQ client whose engine dies must be evicted by the health probe (the /// connection can't reconnect in place — liveness is latched), and the /// handshake guard reset so a later probe rebinds the sockets for a diff --git a/model_gateway/src/workflow/steps/local/create_worker.rs b/model_gateway/src/workflow/steps/local/create_worker.rs index b5ffc63ec..0abffc36e 100644 --- a/model_gateway/src/workflow/steps/local/create_worker.rs +++ b/model_gateway/src/workflow/steps/local/create_worker.rs @@ -117,17 +117,20 @@ impl StepExecutor for CreateLocalWorkerStep { } })?; - // Only vLLM EngineCore and TokenSpeed speak the ZMQ direct-backend wire. - // Fail registration here rather than letting the connect-time rejection - // strand the worker in Pending. + // vLLM EngineCore, TokenSpeed, and SGLang speak the ZMQ direct-backend + // wire. Fail registration here rather than letting the connect-time + // rejection strand the worker in Pending. if *connection_mode == ConnectionMode::Zmq - && !matches!(runtime_type, RuntimeType::Vllm | RuntimeType::TokenSpeed) + && !matches!( + runtime_type, + RuntimeType::Vllm | RuntimeType::TokenSpeed | RuntimeType::Sglang + ) { return Err(WorkflowError::StepFailed { step_id: StepId::new("create_worker"), message: format!( - "ZMQ worker {} has unsupported runtime {}: only vllm and tokenspeed \ - are supported over the ZMQ direct backend", + "ZMQ worker {} has unsupported runtime {}: only vllm, tokenspeed, \ + and sglang are supported over the ZMQ direct backend", config.url, runtime_type ), });