diff --git a/python/sglang/benchmark/serving.py b/python/sglang/benchmark/serving.py index 4bd5e5f01..73a7ce619 100644 --- a/python/sglang/benchmark/serving.py +++ b/python/sglang/benchmark/serving.py @@ -35,6 +35,7 @@ from typing import Any, AsyncGenerator, Callable, Dict, List, Optional, Tuple, U import aiohttp import numpy as np +import orjson import requests from tqdm.asyncio import tqdm from transformers import AutoTokenizer, PreTrainedTokenizerBase @@ -710,12 +711,18 @@ async def async_request_sglang_generate( if not chunk_bytes: continue - chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data: ") + # Cumulative chunks make parsing O(n^2) per request on this + # single asyncio thread; orjson on raw bytes is ~2.2x cheaper. + sse_data = ( + chunk_bytes[6:] + if chunk_bytes.startswith(b"data: ") + else chunk_bytes + ) latency = time.perf_counter() - st - if chunk == "[DONE]": + if sse_data == b"[DONE]": pass else: - data = json.loads(chunk) + data = orjson.loads(sse_data) _meta_info = data.get("meta_info") or {} if _meta_info.get("spec_accept_length") is not None: diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py index 033da16d0..aa7933fa8 100644 --- a/python/sglang/srt/entrypoints/engine.py +++ b/python/sglang/srt/entrypoints/engine.py @@ -55,6 +55,7 @@ from sglang.srt.entrypoints.engine_info_bootstrap_server import ( ) from sglang.srt.entrypoints.engine_score_mixin import EngineScoreMixin from sglang.srt.entrypoints.EngineBase import EngineBase +from sglang.srt.environ import envs from sglang.srt.managers.data_parallel_controller import ( SCHEDULER_PIDS_ARG, run_data_parallel_controller_process, @@ -135,7 +136,7 @@ class SchedulerInitResult: scheduler_infos: List[Dict[str, Any]] all_child_pids: List[int] = dataclasses.field(default_factory=list) wait_for_ready: Callable[[], None] = lambda: None - wait_for_completion: Callable[[], None] = lambda: None + block_until_scheduler_exits: Callable[[], None] = lambda: None engine_info_bootstrap_server: Optional[Any] = None @@ -231,6 +232,14 @@ class Engine(EngineScoreMixin, EngineBase): self.server_args = server_args logger.info(f"{server_args=}") + # Rust Server is not supported with the offline Engine API + if envs.SGLANG_RUST_SERVER.get(): + raise ValueError( + "SGLANG_RUST_SERVER is not supported with the offline Engine " + "API; it only replaces the HTTP server path (`sglang serve`). " + "Unset SGLANG_RUST_SERVER to use sgl.Engine." + ) + # Pre-initialize tokenizer_manager so the atexit handler in # shutdown() won't hit AttributeError. self.tokenizer_manager = None @@ -900,7 +909,7 @@ class Engine(EngineScoreMixin, EngineBase): if SCHEDULER_PIDS_ARG in info: all_child_pids.extend(info[SCHEDULER_PIDS_ARG]) - def wait_for_completion(): + def block_until_scheduler_exits(): for proc in scheduler_procs: proc.join() logger.error( @@ -913,7 +922,7 @@ class Engine(EngineScoreMixin, EngineBase): scheduler_infos=scheduler_infos, all_child_pids=all_child_pids, wait_for_ready=wait_for_ready, - wait_for_completion=wait_for_completion, + block_until_scheduler_exits=block_until_scheduler_exits, ), scheduler_procs, ) @@ -1074,7 +1083,7 @@ class Engine(EngineScoreMixin, EngineBase): server_args.host, server_args.port, server_args.enable_metrics ) - scheduler_init_result.wait_for_completion() + scheduler_init_result.block_until_scheduler_exits() return ( None, None, @@ -1084,6 +1093,29 @@ class Engine(EngineScoreMixin, EngineBase): weight_cache_daemon_procs, ) + # The embedded Rust server (started inside the rank-0 scheduler) owns + # the API server, tokenization, and detokenization. In that mode we do + # not start the Python detokenizer subprocess(es) or tokenizer manager. + # Do not use RayEngine with the Rust server, as it is not supported. + if envs.SGLANG_RUST_SERVER.get(): + scheduler_init_result.wait_for_ready() + # Set up subprocess liveness watchdog to detect crashes + processes = list(scheduler_procs or []) + names = [f"scheduler_{i}" for i in range(len(processes))] + subprocess_watchdog = SubprocessWatchdog( + processes=processes, process_names=names + ) + subprocess_watchdog.start() + + return ( + None, + None, + port_args, + scheduler_init_result, + subprocess_watchdog, + None, + ) + # Launch detokenizer process(es) — optionally fronted by a router when # detokenizer_worker_num > 1. detoken_procs, detoken_names = cls._launch_detokenizer_subprocesses( diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index 2a843f156..88b51f204 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -2184,10 +2184,20 @@ def _execute_server_warmup(server_args: ServerArgs): and server_args.disaggregation_mode == "null" and model_info["is_generation"] ): + served_model_name = "" + if not envs.SGLANG_RUST_SERVER.get(): + served_model_name = _global_state.tokenizer_manager.served_model_name + else: + # _global_state.tokenizer_manager is not initialized in the rust server, + # so we need to get the model name from the model_info + served_model_name = model_info.get( + "model_path", server_args.served_model_name + ) + served_model_name = served_model_name or server_args.model_path # TODO: ChatCompletionRequest does not have bootstrap info required by disaggregation mode, disable image-warmup for now # Only use chat completions format for generation models, not embedding models json_data = { - "model": _global_state.tokenizer_manager.served_model_name, + "model": served_model_name, "messages": [ { "role": "user", @@ -2236,9 +2246,15 @@ def _execute_server_warmup(server_args: ServerArgs): verify=ssl_verify, ) assert res.status_code == 200, f"{res.text}" - _global_state.tokenizer_manager.server_status = ServerStatus.Up + # Skip server_status update for Rust server + if not envs.SGLANG_RUST_SERVER.get(): + _global_state.tokenizer_manager.server_status = ServerStatus.Up else: + # TODO: @rainj-me fix this when Rust server supports disaggregation + assert ( + not envs.SGLANG_RUST_SERVER.get() + ), "Rust server is not supported for disaggregation warmup for now" logger.info(f"Start of pd disaggregation warmup ...") status_codes = asyncio.run( _send_disaggregation_warmup_requests( @@ -2691,13 +2707,27 @@ def launch_server( run_detokenizer_process_func=run_detokenizer_process_func, ) - _setup_and_run_http_server( - server_args, - tokenizer_manager, - template_manager, - port_args, - scheduler_init_result.scheduler_infos, - subprocess_watchdog, - execute_warmup_func=execute_warmup_func, - launch_callback=launch_callback, - ) + if envs.SGLANG_RUST_SERVER.get(): + # The Rust server serves api-server, tokenizer, and detokenizer, so the + # main process has no Python HTTP server / tokenizer manager to run. + # Run a warmup /generate before advertising readiness: the Rust /health + # and /get_model_info endpoints are static (200 as soon as the server + # binds, before any forward pass), so without this the first real request + # pays the cold-start cost (observed as a >60s first generation). + if not server_args.skip_server_warmup: + _execute_server_warmup(server_args) + logger.info("The server is fired up and ready to roll!") + if launch_callback is not None: + launch_callback() + scheduler_init_result.block_until_scheduler_exits() + else: + _setup_and_run_http_server( + server_args, + tokenizer_manager, + template_manager, + port_args, + scheduler_init_result.scheduler_infos, + subprocess_watchdog, + execute_warmup_func=execute_warmup_func, + launch_callback=launch_callback, + ) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 0787565a5..b46a8063e 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -1258,6 +1258,13 @@ class Envs: SGLANG_KV_CANARY_SWA_DIVERGENCE_STATS_INTERVAL = EnvInt(0) SGLANG_KV_CANARY_ENABLE_MHA_V = EnvBool(False) + # =================================================================== + # Rust Server specific envs. + # =================================================================== + SGLANG_RUST_SERVER = EnvBool(False) + # Most batched requests one /generate HTTP call may expand into. + SGLANG_MAX_BATCH_REQS_PER_HTTP_REQ = EnvInt(4096) + envs = Envs() EnvField._allow_set_name = False diff --git a/python/sglang/srt/managers/rust_server.py b/python/sglang/srt/managers/rust_server.py new file mode 100644 index 000000000..655db276c --- /dev/null +++ b/python/sglang/srt/managers/rust_server.py @@ -0,0 +1,392 @@ +"""Embedded Rust server lifecycle for the scheduler. + +The Rust server replaces the Python api-server + `TokenizerManager` + +`DetokenizerManager` stack (hence this module sits beside them in `managers/`), +running them as Rust threads inside the scheduler process. This wrapper keeps +all `SGLANG_RUST_SERVER` plumbing — startup, CPU-core partitioning, the +`server_args` blob, and control-response routing — out of `scheduler.py`. The +scheduler holds an `Optional[RustServer]` and delegates to it. +""" + +from __future__ import annotations + +import logging +import os +from array import array +from itertools import chain +from typing import TYPE_CHECKING, Any, List, Optional, Tuple + +import msgspec + +from sglang.srt.managers.utils import ( + MsgpackDecodeError, + compute_num_reserved_tokens, + msgpack_decode_explained, +) +from sglang.srt.utils.flatten import ( + FlatPairColumns, + NestedRowColumns, + RaggedPairColumns, +) +from sglang.version import __version__ + +if TYPE_CHECKING: + from sglang.srt.managers.io_struct import BatchTokenIDOutput + from sglang.srt.managers.scheduler import Scheduler + from sglang.srt.server._core import Server + +logger = logging.getLogger(__name__) + + +class RustServer: + """Owns the embedded multi-threaded Rust server (``sglang_server.Server``). + + The server owns the api-server, tokenizermanager, tokenizer, and detokenizer + all implemented as Rust threads in scheduler process. + """ + + def __init__(self, server: Server, max_per_poll: int = 256): + self.server = server + self._max_per_poll = max_per_poll + + @classmethod + def launch(cls, scheduler: Scheduler) -> RustServer: + """Start the embedded Rust server threads and bind the listen port. + + The caller gates this (``SGLANG_RUST_SERVER`` + rank 0); this always + creates. + """ + from sglang.srt.server._core import Server + + # Force turn off HF tokenizers rayon's unpinned global thread pool. + os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") + + server_args = scheduler.server_args + # `TokenizerManager` merges these under each request's own sampling params + # (`{**preferred, **obj.sampling_params}`), and this server replaces that + # manager wholesale — so honouring the flag is not implemented here yet. + # Refuse rather than run: silently dropping it means generating with + # sampling the operator did not configure, and `/get_model_info` would go on + # advertising values no request ever receives. + if server_args.preferred_sampling_params: + raise ValueError( + "SGLANG_RUST_SERVER does not yet apply --preferred-sampling-params " + "(the Python TokenizerManager merges it into every request; the rust " + "ingress has no equivalent). Launch without SGLANG_RUST_SERVER, or " + "drop --preferred-sampling-params and send those values per request." + ) + http_addr = f"{server_args.host}:{server_args.port}" + launch_cores, server_cores = cls._partition_cores() + + server = Server( + # None -> run unpinned; the list carries the pinning decision. + cores=server_cores, + http_addr=http_addr, + server_args_json=cls._build_server_args(scheduler), + ) + + # Narrow the scheduler thread only after the server threads are launched. + if launch_cores is not None: + try: + # pid 0 == this thread (the scheduler event-loop / launch thread). + os.sched_setaffinity(0, set(launch_cores)) + except OSError as e: + logger.warning("rust server: cannot pin scheduler launch thread: %s", e) + + logger.info( + "SGLANG_RUST_SERVER enabled, Rust server listen on %s", + http_addr, + ) + + return cls(server) + + def wait_ingress(self, timeout_ms: int) -> None: + """Block until a request is pushed into the in-process ring or the timeout + elapses. + """ + self.server.wait_ingress(timeout_ms) + + def drain(self, max_recv: int) -> List[Any]: + """Ingress: non-blocking drain of the in-process ring → list of decoded + request objects. The scheduler's request receiver calls this instead of + polling the zmq socket when `rust_server_mode` is set. + + The transfer is **columnar**: `recv_requests` returns an `IngressBatch` + of scalar msgpack `headers` (with `input_ids` omitted) plus one + concatenated raw int64 `data` buffer and per-request `lengths`, so the + large `input_ids` lists never go through msgpack. Each header is `msgpack_decode`d (yielding + the same `TokenizedGenerateReqInput` / control objects the zmq path + produces, so the IPC schema is tracked automatically) and its `input_ids` + slice is wrapped as the `array("q")` the scheduler expects. `recv_requests` + releases the GIL for the drain + concat, so this never holds the GIL + across a wait — same contract as `zmq.NOBLOCK`. + """ + limit = max_recv if max_recv > 0 else self._max_per_poll + batch = self.server.recv_requests(limit) + # Bind once: each attribute access converts the rust vec to a fresh list. + headers, data, lengths = batch.headers, batch.data, batch.lengths + if not headers: + return [] + + ids_view = memoryview(data) + out = [] + pos = 0 # byte offset into ids_buf + for header, n in zip(headers, lengths): + nbytes = n * 8 + try: + obj = msgpack_decode_explained(header) + except MsgpackDecodeError as e: + # Return 400 for malformed request field (e.g. token_ids_logprob=[[0]]. + logger.warning( + "rust ingress: dropping undecodable request %s: %s", e.rid, e.reason + ) + if e.rid is not None: + self.server.push_error(e.rid, f"invalid request: {e.reason}") + pos += nbytes + continue + if n: # generate request: attach its int64 ids slice as array("q") + ids = array("q") + ids.frombytes(ids_view[pos : pos + nbytes]) + obj.input_ids = ids + pos += nbytes + out.append(obj) + return out + + def push_control_output(self, recv_req, output) -> None: + """Push a control-request response through the egress ring to the waiting + request (routed by rid), encoded as **msgpack** (the ring's native + format). + + A msgspec struct is converted to a *named map* (``structs.asdict``, since + the IPC structs are ``array_like`` and would otherwise lose field names) + so the Rust api_server can shape it per-endpoint (e.g. /server_info) + before rendering JSON to the client — keeping JSON formatting off the + scheduler's GIL. + """ + + # Invariant: control requests always carry a rust-minted rid; without + # one the response is unroutable, so fail loudly rather than drop it. + assert ( + recv_req.rid is not None + ), f"control response without rid: {type(output).__name__}" + # No local try/except: a failed push propagates to run_scheduler_process's + # outer handler, which logs the full traceback (scheduler-fatal either way). + payload = ( + msgspec.structs.asdict(output) + if isinstance(output, msgspec.Struct) + else output + ) + # enc_hook stringifies non-native types (paths, enums); JSON + # rendering happens in Rust. + encoded = msgspec.msgpack.encode(payload, enc_hook=str) + + self.server.push_result(recv_req.rid, encoded) + + def push_generation(self, payload: BatchTokenIDOutput) -> None: + """Egress redirect for generation output (replaces the zmq detokenizer). + + Push the WHOLE batch into the Rust egress ring as one frame (-> detokenizer + shards -> client streams), mirroring the ingress ``input_ids`` split so the + bulk numeric columns never go through msgpack: + + - ``header``: msgpack ``BatchHeader`` positional array — the per-request + scalar columns (``rids, finish_reasons, prompt_tokens, tok_lens``) plus + the shape metadata for the optional families (``*_lens`` element counts + for the flat logprob columns, ``*_reqlens``/``*_poslens`` for the ragged + and hidden ones). + - ``data``: the raw little-endian numeric buffer — every column is a + 4-byte element (``f32`` values, ``i32`` indices), concatenated in the + order the Rust ``for_each_chunk`` reads them. + + Logprobs are columnar: output families are per-step deltas, input + (prefill) families ride once on the first chunk. Ragged families (top-k, + token-ids) flatten a per-position ``list[list]`` into flat ``val``/``idx`` + buffers plus a per-position ``lens`` vector (0 = null position). Hidden + states flatten to rows of floats (one row per output position). + """ + output_ids = payload.output_ids or [] + prompt_tokens = payload.prompt_tokens or [] + + # Hot-path guard: almost no decode step wants logprobs / hidden states, + # so only then pay the per-request flatten + buffer packing below. + has_extra = bool( + payload.output_token_logprobs_val + or payload.input_token_logprobs_val + or payload.output_top_logprobs_val + or payload.input_top_logprobs_val + or payload.output_token_ids_logprobs_val + or payload.input_token_ids_logprobs_val + or payload.output_hidden_states + ) + + # Runs on the scheduler's CUDA-launch thread every decode step, so each + # Python-level pass over the batch costs inter-token latency: `rids` are + # the plain rid strings (hashed to a routing key on the Rust side with a + # per-process seed, off the GIL — not parsed; a rid is any string), + # `finished_reasons` already `dict | None`, and `output_ids` entries are + # always `array("i")` (never None) so `map(len)` and a bare + # `chain.from_iterable` stay in C. + rids = payload.rids + finish_reasons = payload.finished_reasons + tok_lens = list(map(len, output_ids)) + flat_ids = array("i", chain.from_iterable(output_ids)) + + # Column order here MUST match BatchHeader (header_cols) and + # for_each_chunk's read order (data_cols); the extras contribution + # is ordered by the `extras` tuple below. + header_cols = [rids, finish_reasons, prompt_tokens, tok_lens] + data_cols = [flat_ids.tobytes()] + + if has_extra: + # The `extras` tuple is the SINGLE source of the extras column + # order — it must match the Rust ``BatchHeader`` fields and + # ``for_each_chunk``'s read order. + # + # TODO(perf): the per-request flatten assumes the logprob/hidden + # columns are ragged, non-contiguous nested Python lists — which is + # only an assumption. The scheduler moves these off the GPU with + # `tensor.tolist()`, so revisit whether the upstream values are + # still contiguous tensors; if so, ship raw bytes + a shape + # descriptor and skip the flatten entirely. + batch_size = len(rids) + extras = ( + FlatPairColumns( + "output_token_logprobs", + payload.output_token_logprobs_val or [], + payload.output_token_logprobs_idx or [], + ), + FlatPairColumns( + "input_token_logprobs", + payload.input_token_logprobs_val or [], + payload.input_token_logprobs_idx or [], + first_none_to_nan=True, + ), + RaggedPairColumns( + "output_top_logprobs", + payload.output_top_logprobs_val or [], + payload.output_top_logprobs_idx or [], + ), + RaggedPairColumns( + "input_top_logprobs", + payload.input_top_logprobs_val or [], + payload.input_top_logprobs_idx or [], + ), + RaggedPairColumns( + "output_token_ids_logprobs", + payload.output_token_ids_logprobs_val or [], + payload.output_token_ids_logprobs_idx or [], + ), + RaggedPairColumns( + "input_token_ids_logprobs", + payload.input_token_ids_logprobs_val or [], + payload.input_token_ids_logprobs_idx or [], + ), + NestedRowColumns( + "output_hidden_states", payload.output_hidden_states or [] + ), + ) + + # Every column is all-or-nothing per payload — which is also what makes + # a family's emptiness a reliable "nobody asked for this" signal. + active = [] + for extra in extras: + populated = False + for name, col in extra.columns(): + assert len(col) in ( + 0, + batch_size, + ), f"extras column {name}: {len(col)} entries for a batch of {batch_size}" + populated |= len(col) > 0 + if populated: + active.append(extra) + + # Flatten only the families someone asked for. `has_extra` above is a + # per-FRAME guard, so one client enabling logprobs used to drag all + # seven families through the per-request loop: at B=4096 that is 28,672 + # bound-method calls per decode step, materializing 12 columns of 4096 + # zeros nobody reads. Measured 0.37 ms -> 7.90 ms GIL-held per step, + # i.e. 25-75% of a decode step added to the scheduler's critical path. + # + # Skipping `accept` leaves a family's buffers empty, which is exactly + # the wire form the Rust decoder already treats as absent (`per_req_ok` + # admits an empty column, `lens_i` reads 0 for every request). The + # `header_cols`/`data_cols` loops below still walk all seven, so column + # ORDER and arity are unchanged — an inactive family contributes empty + # columns in place rather than disappearing. + for extra in active: + accept = extra.accept # hoisted: this is the hottest loop here + for i in range(batch_size): + accept(i) + + for extra in extras: + header_cols += extra.header_cols() + data_cols += extra.data_cols() + + header = msgspec.msgpack.encode(header_cols) + # Pass the raw column list; the Rust side concatenates it into the frame + # with the GIL released. + if not self.server.push_batch(header, data_cols): + logger.warning( + "Rust egress closed; dropped batch of %d requests during shutdown", + len(rids), + ) + + @staticmethod + def _build_server_args(scheduler: Scheduler) -> str: + """JSON blob of the scheduler's ``server_args`` for its embedded Rust + server (carries the already-resolved ``model_config``).""" + + server_args = dict(vars(scheduler.server_args)) + model_config = dict(vars(scheduler.model_config)) + model_config["hf_config"] = None # HF config is not JSON-serializable + server_args["model_config"] = model_config + # Launch-time facts Python's /server_info reports from scheduler_info / + # the package — stamped here so the rust endpoint can serve them + # statically (no scheduler round-trip). + server_args["version"] = __version__ + # Not a `server_args` field: `TokenizerManager` derives it, and the rust + # ingress needs the same number for its total-token check. + server_args["num_reserved_tokens"] = compute_num_reserved_tokens( + scheduler.server_args + ) + server_args["max_total_num_tokens"] = scheduler.max_total_num_tokens + + return msgspec.json.encode(server_args, enc_hook=str).decode("utf-8") + + @staticmethod + def _partition_cores() -> Tuple[Optional[List[int]], Optional[List[int]]]: + """Split this rank's allowed cores into ``(launch_cores, server_cores)``. + + Pure computation — no affinity is changed here. Both sets are a subset + of this rank's NUMA-local cores (when affinity/NUMA bind is on), so the + partition stays NUMA-local. Returns ``(None, None)`` (server runs + unpinned, confined only by the process affinity) when the platform has + no affinity API or too few cores to split. + """ + if not hasattr(os, "sched_getaffinity"): + return None, None + try: + allowed = sorted(os.sched_getaffinity(0)) + except OSError as e: + logger.warning("rust server: cannot read cpu affinity: %s", e) + return None, None + + # Need enough cores to reserve launch cores and still pin the pools. + if len(allowed) < 4: + logger.info( + "rust server: only %d cores allowed; running pools unpinned", + len(allowed), + ) + return None, None + + # Keep a small slice for the launch loop; cap at 2 (the event loop is + # effectively serial) and never take more than a quarter of the cores. + reserve = min(2, len(allowed) // 4) + launch_cores = allowed[:reserve] + server_cores = allowed[reserve:] + logger.info( + "rust server cores=%s, scheduler launch cores=%s", + server_cores, + launch_cores, + ) + return launch_cores, server_cores diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 4b3548de3..98b4341cb 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -1450,7 +1450,28 @@ class Req(ReqDllmMixin): # Check stop regex if len(self.sampling_params.stop_regex_strs) > 0: for stop_regex_str in self.sampling_params.stop_regex_strs: - if re.search(stop_regex_str, tail_str): + # Seatbelt, not validation: patterns are checked at ingress + # (Python's `normalize`, or the rust server's stricter + # `stop_regex_bound`). This runs per decode step on the hot + # path, so an `re.error` escaping here would take the whole + # scheduler down over one malformed request. Fail that request + # instead. + try: + matched = re.search(stop_regex_str, tail_str) + except (re.error, RecursionError) as e: + logger.warning( + "req %s: invalid stop_regex %r (%s); aborting the request", + self.rid, + stop_regex_str, + e, + ) + self.finished_reason = FINISH_ABORT( + f"invalid stop_regex {stop_regex_str!r}: {e}", + HTTPStatus.BAD_REQUEST, + "BadRequestError", + ) + break + if matched: self.finished_reason = FINISHED_MATCHED_REGEX( matched=stop_regex_str ) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 861c13fb4..b26525b21 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -167,6 +167,7 @@ from sglang.srt.managers.prefill_delayer import ( PrefillDelayer, PrefillDelayerSinglePassExecutor, ) +from sglang.srt.managers.rust_server import RustServer from sglang.srt.managers.schedule_batch import ( FINISH_ABORT, MultimodalInputs, @@ -185,7 +186,10 @@ from sglang.srt.managers.scheduler_components.batch_result_processor import ( ) from sglang.srt.managers.scheduler_components.dp_attn import SchedulerDPAttnAdapter from sglang.srt.managers.scheduler_components.flush_wrapper import SchedulerFlushWrapper -from sglang.srt.managers.scheduler_components.idle_sleeper import IdleSleeper +from sglang.srt.managers.scheduler_components.idle_sleeper import ( + IdleSleeper, + RustServerIdleSleeper, +) from sglang.srt.managers.scheduler_components.invariant_checker import ( SchedulerInvariantChecker, create_scheduler_watchdog, @@ -597,6 +601,10 @@ class Scheduler( self.maybe_init_scripted_scheduler_hook() + # Start the embedded Rust frontend (rank 0) before the request receiver, + # which reads self.rust_ring_recv to pick its ingress transport. + self.maybe_init_rust_server() + self.init_request_receiver() self.init_dp_attn_adapter() @@ -672,9 +680,12 @@ class Scheduler( ) self.load_snapshot_writer = None + self.recv_from_tokenizer = None + if not is_rank_zero: return + self.recv_from_tokenizer = self.ipc_channels.recv_from_tokenizer dp_rank = self.ps.dp_rank if self.ps.dp_rank is not None else 0 try: self.load_snapshot_writer = create_load_snapshot_writer( @@ -1748,11 +1759,16 @@ class Scheduler( output = self._request_dispatcher(recv_req) if output is not None: - if not isinstance(output, RpcReqOutput): - self.ipc_channels.send_to_tokenizer.send_output(output, recv_req) - else: + if self.rust_server is not None: + # Embedded Rust server: every control-request response goes + # back through the egress ring (the zmq tokenizer socket is + # not consumed); the Rust api_server shapes it per-endpoint. + self.rust_server.push_control_output(recv_req, output) + elif isinstance(output, RpcReqOutput): if self.ipc_channels.recv_from_rpc is not None: sock_send(self.ipc_channels.recv_from_rpc, output) + else: + self.ipc_channels.send_to_tokenizer.send_output(output, recv_req) self.flush_wrapper.check_pending() if self.external_corpus_manager is not None: @@ -1808,9 +1824,33 @@ class Scheduler( else: self.scripted_scheduler_hook = None + def maybe_init_rust_server(self) -> None: + """Start the embedded Rust server (rank 0) if ``SGLANG_RUST_SERVER`` is + set, and point the ingress receiver at it. All the plumbing lives in + ``RustServer`` (scheduler_components/rust_scheduler.py).""" + + is_rank_zero = ( + self.ps.pp_rank == 0 + and self.ps.attn_tp_rank == 0 + and self.ps.attn_cp_rank == 0 + ) + if not (envs.SGLANG_RUST_SERVER.get() and is_rank_zero): + # Always define the attribute: init_output_streamer and the + # process_input_requests hook read self.rust_server unconditionally. + self.rust_server = None + return + + rust_server = RustServer.launch(self) + self.rust_server = rust_server + # The rust server *is* the ingress source: SchedulerRequestReceiver + # drains its request ring (rust_server_mode) instead of a zmq socket. + self.recv_from_tokenizer = rust_server + # Park the idle loop on the request ring within the rank-0 rust-server + self.idle_sleeper = RustServerIdleSleeper(rust_server) + def init_request_receiver(self) -> None: self.request_receiver = SchedulerRequestReceiver( - recv_from_tokenizer=self.ipc_channels.recv_from_tokenizer, + recv_from_tokenizer=self.recv_from_tokenizer, recv_from_rpc=self.ipc_channels.recv_from_rpc, recv_skipper=self.recv_skipper, input_blocker=self.input_blocker, @@ -1936,6 +1976,7 @@ class Scheduler( spec_algorithm=self.spec_algorithm, disaggregation_mode=self.disaggregation_mode, enable_hicache_storage=lambda: self.enable_hicache_storage, + rust_server=self.rust_server, ) def init_batch_result_processor(self) -> None: diff --git a/python/sglang/srt/managers/scheduler_components/idle_sleeper.py b/python/sglang/srt/managers/scheduler_components/idle_sleeper.py index 38c0c7a2e..cef2e3419 100644 --- a/python/sglang/srt/managers/scheduler_components/idle_sleeper.py +++ b/python/sglang/srt/managers/scheduler_components/idle_sleeper.py @@ -1,9 +1,16 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + import zmq from sglang.srt.environ import envs from sglang.srt.observability.req_time_stats import real_time from sglang.srt.platforms import current_platform +if TYPE_CHECKING: + from sglang.srt.managers.rust_server import RustServer + class IdleSleeper: """ @@ -33,3 +40,29 @@ class IdleSleeper: ): self.last_empty_time = real_time() current_platform.empty_cache() + + +class RustServerIdleSleeper: + """Idle sleeper for the embedded Rust server. + + The Rust ingress is an in-process request ring, not a zmq socket. + Instead park directly on the ring: ``wait_ingress`` blocks until + a request is pushed — the request ring wakes the parked thread + the instant a producer pushes, so there's no added latency for real + requests — or the timeout elapses. + """ + + def __init__(self, rust_server: RustServer, timeout_ms: int = 1000): + self.rust_server = rust_server + self.timeout_ms = timeout_ms + self.last_empty_time = real_time() + self.empty_cache_interval = envs.SGLANG_EMPTY_CACHE_INTERVAL.get() + + def maybe_sleep(self): + self.rust_server.wait_ingress(self.timeout_ms) + if ( + self.empty_cache_interval > 0 + and real_time() - self.last_empty_time > self.empty_cache_interval + ): + self.last_empty_time = real_time() + current_platform.empty_cache() diff --git a/python/sglang/srt/managers/scheduler_components/output_streamer.py b/python/sglang/srt/managers/scheduler_components/output_streamer.py index 278ccf428..b06c66e33 100644 --- a/python/sglang/srt/managers/scheduler_components/output_streamer.py +++ b/python/sglang/srt/managers/scheduler_components/output_streamer.py @@ -3,6 +3,7 @@ from __future__ import annotations import logging from dataclasses import dataclass, field from typing import ( + TYPE_CHECKING, Any, Callable, List, @@ -29,6 +30,10 @@ from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache from sglang.srt.server_args import ServerArgs from sglang.srt.speculative.spec_info import SpeculativeAlgorithm +if TYPE_CHECKING: + from sglang.srt.managers.rust_server import RustServer + + logger = logging.getLogger(__name__) @@ -45,6 +50,10 @@ class SchedulerOutputStreamer: spec_algorithm: SpeculativeAlgorithm disaggregation_mode: DisaggregationMode enable_hicache_storage: Callable[[], bool] + # When SGLANG_RUST_SERVER is on, generation output is pushed to the embedded + # Rust egress ring via `rust_server.push_generation` instead of the zmq + # detokenizer. None otherwise. (Rust-specific state lives in RustServer.) + rust_server: Optional[RustServer] = None _test_stream_output_count: int = 0 def _get_storage_backend_type(self) -> str: @@ -147,6 +156,7 @@ class SchedulerOutputStreamer: default_stream_interval=self.server_args.stream_interval, default_force_stream_interval=DEFAULT_FORCE_STREAM_INTERVAL, get_cached_tokens_details=self.get_cached_tokens_details, + rust_server_mode=self.rust_server is not None, ) for req in reqs: if req is skip_req: @@ -165,7 +175,10 @@ class SchedulerOutputStreamer: is_idle_batch=is_idle_batch, ) if payload is not None: - self.send_to_detokenizer.send_output(payload) + if self.rust_server is not None: + self.rust_server.push_generation(payload) + else: + self.send_to_detokenizer.send_output(payload) def _maybe_log_time_stats(self, *, req: Req) -> None: if ( @@ -260,7 +273,6 @@ class _GenerationStreamAccumulator: default_stream_interval: int default_force_stream_interval: int get_cached_tokens_details: Callable[[Req], Optional[CachedTokensDetails]] - rids: list = field(default_factory=list) http_worker_ipcs: list = field(default_factory=list) finished_reasons: list = field(default_factory=list) @@ -307,6 +319,10 @@ class _GenerationStreamAccumulator: output_token_ids_logprobs_idx: Optional[list] = None output_token_sampling_mask: Optional[list] = None output_token_sampling_logprobs: Optional[list] = None + # Rust server mode: the Rust detokenizer reconstructs text/ids from the raw + # output tokens itself and never consumes the scheduler's incremental-detok + # offsets (decode_ids / read_offset), so that per-step bookkeeping is skipped. + rust_server_mode: bool = False def __post_init__(self) -> None: if self.return_hidden_states: @@ -367,34 +383,40 @@ class _GenerationStreamAccumulator: send_token_offset = req.send_token_offset send_output_token_logprobs_offset = req.send_output_token_logprobs_offset self.rids.append(req.rid) - self.http_worker_ipcs.append(req.http_worker_ipc) self.finished_reasons.append( req.finished_reason.to_json() if req.finished_reason else None ) - self.decoded_texts.append(req.decoded_text) - decode_ids, read_offset = req.init_incremental_detokenize() - - self.decode_ids_list.append(decode_ids[req.send_decode_id_offset :]) - # Exclude the tokens after stop condition output_ids_ = req.output_ids_through_stop - - req.send_decode_id_offset = len(decode_ids) - self.read_offsets.append(read_offset) self.output_ids.append(output_ids_[send_token_offset:]) req.send_token_offset = len(output_ids_) - self.skip_special_tokens.append(req.sampling_params.skip_special_tokens) - self.spaces_between_special_tokens.append( - req.sampling_params.spaces_between_special_tokens - ) - self.no_stop_trim.append(req.sampling_params.no_stop_trim) self.prompt_tokens.append(len(req.origin_input_ids)) - self.reasoning_tokens.append(req.reasoning_tokens) - self.completion_tokens.append(len(output_ids_)) - self.cached_tokens.append(req.cached_tokens) - # Collect detailed cache breakdown if available - self.cached_tokens_details.append(self.get_cached_tokens_details(req)) + if not self.rust_server_mode: + # Everything below feeds the Python DetokenizerManager / + # TokenizerManager (incremental detok, meta_info, per-request metrics) + # or gets pickled into the payload (time_stats). The Rust server + # replaces those stages and builds its own metadata from the + # ChunkEvent, so `push_generation` never reads these — skip the whole + # block. The parallel lists stay empty; the payload goes straight to + # `push_generation`, which only indexes the fields appended above. + self.http_worker_ipcs.append(req.http_worker_ipc) + self.decoded_texts.append(req.decoded_text) + decode_ids, read_offset = req.init_incremental_detokenize() + self.decode_ids_list.append(decode_ids[req.send_decode_id_offset :]) + req.send_decode_id_offset = len(decode_ids) + self.read_offsets.append(read_offset) + self.skip_special_tokens.append(req.sampling_params.skip_special_tokens) + self.spaces_between_special_tokens.append( + req.sampling_params.spaces_between_special_tokens + ) + self.no_stop_trim.append(req.sampling_params.no_stop_trim) + self.reasoning_tokens.append(req.reasoning_tokens) + self.completion_tokens.append(len(output_ids_)) + self.cached_tokens.append(req.cached_tokens) + + # Collect detailed cache breakdown if available + self.cached_tokens_details.append(self.get_cached_tokens_details(req)) # Multimodal prompt token counts. In disagg decode mode the prefill node # already computed these and transferred them via the metadata buffer diff --git a/python/sglang/srt/managers/scheduler_components/request_receiver.py b/python/sglang/srt/managers/scheduler_components/request_receiver.py index 11df155a3..d5a93fc1f 100644 --- a/python/sglang/srt/managers/scheduler_components/request_receiver.py +++ b/python/sglang/srt/managers/scheduler_components/request_receiver.py @@ -15,6 +15,7 @@ import zmq from torch.distributed import barrier from sglang.srt.disaggregation.utils import prepare_abort +from sglang.srt.environ import envs from sglang.srt.managers.io_struct import ( BatchTokenizedEmbeddingReqInput, BatchTokenizedGenerateReqInput, @@ -35,6 +36,7 @@ from sglang.srt.utils.nvtx_utils import scheduler_nvtx_method if TYPE_CHECKING: from sglang.srt.configs.model_config import ModelConfig from sglang.srt.distributed.parallel_state_wrapper import ParallelState + from sglang.srt.managers.rust_server import RustServer from sglang.srt.server_args import ServerArgs from sglang.test.scripted_runtime.scheduler_hook import ScriptedSchedulerHook from sglang.test.scripted_runtime.tokenizer_recv_proxy import ( @@ -44,7 +46,7 @@ if TYPE_CHECKING: @dataclass(kw_only=True, slots=True, frozen=True) class SchedulerRequestReceiver: - recv_from_tokenizer: Union[zmq.Socket, ScriptedTokenizerRecvProxy] + recv_from_tokenizer: Union[zmq.Socket, ScriptedTokenizerRecvProxy, RustServer] recv_from_rpc: Optional[zmq.Socket] recv_skipper: Any input_blocker: Any @@ -103,6 +105,15 @@ class SchedulerRequestReceiver: if self.ps.attn_tp_rank == 0 and self.ps.attn_cp_rank == 0: recv_reqs = [] + # Rust ringbuffer backend: drain the in-process ring fed by the + # embedded Rust TokenizerManager instead of a zmq socket. Same + # non-blocking, msgpack-decoded contract as the zmq path below. + if envs.SGLANG_RUST_SERVER.get(): + recv_reqs.extend( + self.recv_from_tokenizer.drain(self.max_recv_per_poll) + ) + return recv_reqs + while True: try: if self.recv_limit_reached(len(recv_reqs)): diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index c8934b27e..0d7237db8 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -94,7 +94,10 @@ from sglang.srt.managers.schedule_batch import MultimodalDataItem from sglang.srt.managers.scheduler_input_blocker import input_blocker_guard_region from sglang.srt.managers.tokenizer_control_mixin import TokenizerControlMixin from sglang.srt.managers.tokenizer_manager_score_mixin import TokenizerManagerScoreMixin -from sglang.srt.managers.utils import is_health_check_generate_req +from sglang.srt.managers.utils import ( + compute_num_reserved_tokens, + is_health_check_generate_req, +) from sglang.srt.observability.cpu_monitor import start_cpu_monitor_thread from sglang.srt.observability.metrics_collector import ( STAT_LOGGER_ROLE_TOKENIZER, @@ -117,7 +120,6 @@ from sglang.srt.server_args import ( ServerArgs, set_global_server_args_for_tokenizer, ) -from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.utils import ( configure_gc_warning, freeze_gc, @@ -400,18 +402,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): self.max_req_input_len = None # Will be set later in engine.py self.enable_priority_scheduling = server_args.enable_priority_scheduling self.default_priority_value = server_args.default_priority_value - speculative_algorithm = SpeculativeAlgorithm.from_string( - server_args.speculative_algorithm - ) - if speculative_algorithm.is_eagle(): - # In the current eagle implementation, we store the draft tokens in the output token slots, - # so we need to reserve the space for the draft tokens. - self.num_reserved_tokens = max( - server_args.speculative_eagle_topk * server_args.speculative_num_steps, - server_args.max_speculative_num_draft_tokens, - ) - else: - self.num_reserved_tokens = 0 + self.num_reserved_tokens = compute_num_reserved_tokens(server_args) self.validate_total_tokens = True def init_tokenizer_and_processor(self): diff --git a/python/sglang/srt/managers/utils.py b/python/sglang/srt/managers/utils.py index 94ae2a307..fe883c264 100644 --- a/python/sglang/srt/managers/utils.py +++ b/python/sglang/srt/managers/utils.py @@ -2,20 +2,25 @@ from __future__ import annotations import dataclasses import logging +import re from dataclasses import dataclass from typing import TYPE_CHECKING, Any, List, Optional, Union +import msgspec import torch from sglang.srt.constants import HEALTH_CHECK_RID_PREFIX from sglang.srt.eplb.expert_distribution import ExpertDistributionMetrics from sglang.srt.layers.logits_processor import LogitsProcessorOutput +from sglang.srt.managers import io_struct from sglang.srt.managers.schedule_batch import Req from sglang.srt.model_executor.forward_batch_info import PPProxyTensors +from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.state_capturer.base import TopkCaptureOutput if TYPE_CHECKING: from sglang.srt.managers.scheduler import GenerationBatchResult + from sglang.srt.server_args import ServerArgs from sglang.srt.speculative.eagle_info import EagleDraftInput @@ -318,3 +323,73 @@ class EmbeddingBatchResult: def is_health_check_generate_req(recv_req): rid = getattr(recv_req, "rid", None) return rid is not None and rid.startswith(HEALTH_CHECK_RID_PREFIX) + + +class MsgpackDecodeError(ValueError): + """A msgpack frame the typed decoder rejected, with the failure explained: + ``rid`` (when recoverable from the raw tagged array) and a human-readable + ``reason`` whose leading ``$[]`` array index is resolved to the struct + field name. + """ + + def __init__(self, rid: Optional[str], reason: str): + super().__init__(reason) + self.rid = rid + self.reason = reason + + +def msgpack_decode_explained(data: bytes) -> Any: + """`io_struct.msgpack_decode`, but a rejected frame raises + `MsgpackDecodeError` carrying the rid (recovered via an untyped re-decode of + the tagged array) and a reason with the failing field named — for callers + that must report the failure back to a client (e.g. the rust ingress) + instead of just crashing.""" + # TODO: the hook_custom_types() currently only apply for unit tests, once it + # esclate to the main code, we can provide a function to access the _all_types + + try: + return io_struct.msgpack_decode(data) + except Exception as e: + msg = str(e) + try: + arr = msgspec.msgpack.decode(data) + except Exception: + arr = None + if not (isinstance(arr, (list, tuple)) and arr): + raise MsgpackDecodeError(None, msg) from e + # Tagged array_like layout is [tag, *fields]; rid is the first field of + # every BaseReq struct. + rid = str(arr[1]) if len(arr) > 1 and arr[1] is not None else None + tag_to_fields = { + cls.__struct_config__.tag: cls.__struct_fields__ + for cls in io_struct._all_types + if isinstance(cls, type) and issubclass(cls, msgspec.Struct) + } + fields = tag_to_fields.get(arr[0]) + if fields is not None: + # Leading ``$[]`` in a msgspec ValidationError path, e.g. + # ``$[12][0]``. + m = re.search(r"\$\[(\d+)\]", msg) + if m is not None: + idx = int(m.group(1)) + if 1 <= idx <= len(fields): + msg = f"{msg[:m.start()]}$.{fields[idx - 1]}{msg[m.end():]}" + raise MsgpackDecodeError(rid, msg) from e + + +def compute_num_reserved_tokens(server_args: ServerArgs) -> int: + """Output token slots reserved per request, on top of its input. + + The current eagle implementation stores draft tokens in the output token + slots, so the context budget has to account for them; every other algorithm + reserves nothing. Shared by `TokenizerManager` and the rust server's + `server_args` blob (`RustServer._build_server_args`), which needs the same + number to run the total-token check in Rust. + """ + algorithm = SpeculativeAlgorithm.from_string(server_args.speculative_algorithm) + if not algorithm.is_eagle(): + return 0 + return max( + server_args.speculative_eagle_topk * server_args.speculative_num_steps, + server_args.max_speculative_num_draft_tokens, + ) diff --git a/python/sglang/srt/ray/engine.py b/python/sglang/srt/ray/engine.py index 57c1d16b8..efbbf09e7 100644 --- a/python/sglang/srt/ray/engine.py +++ b/python/sglang/srt/ray/engine.py @@ -401,7 +401,7 @@ class RayEngine(Engine): actor.run_event_loop.remote() for actor in scheduler_actors ] - def wait_for_completion(): + def block_until_scheduler_exits(): try: ray.get(event_loop_refs) except Exception as e: @@ -410,7 +410,7 @@ class RayEngine(Engine): return ( RaySchedulerInitResult( scheduler_infos=scheduler_infos, - wait_for_completion=wait_for_completion, + block_until_scheduler_exits=block_until_scheduler_exits, scheduler_actors=scheduler_actors, ), None, @@ -489,7 +489,7 @@ class RayEngine(Engine): event_loop_refs = controller.event_loop_refs - def wait_for_completion(): + def block_until_scheduler_exits(): try: ray.get(event_loop_refs) except Exception as e: @@ -497,6 +497,6 @@ class RayEngine(Engine): return RaySchedulerInitResult( scheduler_infos=scheduler_infos, - wait_for_completion=wait_for_completion, + block_until_scheduler_exits=block_until_scheduler_exits, scheduler_actors=controller.scheduler_actors, ) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 9905744fd..e23e34624 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -7377,7 +7377,7 @@ class ServerArgs: "Please choose one tokenizer batching approach." ) - if self.skip_tokenizer_init: + if self.skip_tokenizer_init and not envs.SGLANG_RUST_SERVER.get(): # Tokenizer workers still serve HTTP / state / output work, so # their fanout is preserved; detokenizer workers only decode. if self.detokenizer_worker_num != 1: @@ -9190,14 +9190,13 @@ class PortArgs: dist_init_host = na.host dist_init_port = na.port - # We need 5 consecutive ports from port_base for: - # port_base, detokenizer, rpc, metrics, scheduler. - # In multi-node, all nodes derive ports independently from - # dist_init_port, so the derivation must be deterministic - # (no availability-based search). If incrementing would - # overflow the valid TCP range, decrement instead. - NUM_DERIVED_PORTS = 5 - if server_args.is_ep_joiner: + # Reserve port_base+0..NUM_DERIVED_PORTS-1 (6 fixed ports + dp_size + # rust-path slots); derive from server_args only (never dp_rank) so + # every init_new call agrees, decrementing below dist_init_port on + # overflow. + is_rust_server = envs.SGLANG_RUST_SERVER.get() + NUM_DERIVED_PORTS = 6 if not is_rust_server else 6 + server_args.dp_size + if server_args.is_ep_scale_joiner: port_base = server_args.port + ZMQ_TCP_PORT_DELTA if port_base + NUM_DERIVED_PORTS > 65535: port_base = server_args.port - ZMQ_TCP_PORT_DELTA diff --git a/python/sglang/srt/utils/flatten.py b/python/sglang/srt/utils/flatten.py new file mode 100644 index 000000000..0ff6b4e12 --- /dev/null +++ b/python/sglang/srt/utils/flatten.py @@ -0,0 +1,186 @@ +"""Flatten ragged (nested, variable-length) structures into flat value buffers +plus per-position length vectors — the columnar wire layout used by the embedded +Rust server's egress path (see ``managers/rust_server.py``). + +The ``*Columns`` classes accumulate one batch column-family each: feed them one +request cell at a time (``accept``), then read the header contribution (length +vectors) and data contribution (raw ``array`` buffers). +""" + +from __future__ import annotations + +from array import array +from typing import List + + +def flatten_ragged(per_pos_val, per_pos_idx): + """Flatten per-position ``list[Optional[list]]`` val/idx pairs into flat + buffers + a shared ``lens`` vector (falsy position -> len 0 -> Rust ``null``). + idx must mirror val exactly (asserted): the wire pairs both buffers by the + one ``lens`` vector, so divergence shifts or drops token ids downstream.""" + + flat_val: List[float] = [] + flat_idx: List[int] = [] + lens: List[int] = [] + if not per_pos_val: + assert not per_pos_idx, ( + f"ragged idx column has {len(per_pos_idx)} positions but the val " + "column is empty" + ) + return flat_val, flat_idx, lens + assert per_pos_idx is not None and len(per_pos_idx) == len(per_pos_val), ( + f"ragged idx column has {len(per_pos_idx) if per_pos_idx else 0} " + f"positions, val column has {len(per_pos_val)}" + ) + for p, pv in enumerate(per_pos_val): + pi = per_pos_idx[p] + if pv: + # A truthy position holds only real logprobs; a `None`/empty position + # is the falsy branch below (len 0), so no per-value None check here. + assert pi is not None and len(pi) == len(pv), ( + f"position {p}: idx len " + f"{len(pi) if pi is not None else None} != val len {len(pv)}" + ) + flat_val.extend(pv) + flat_idx.extend(pi) + lens.append(len(pv)) + else: + assert not pi, f"position {p}: idx has {len(pi)} entries but val is empty" + lens.append(0) + return flat_val, flat_idx, lens + + +def flatten_hidden(hs): + """Flatten one request's hidden states into a flat ``val`` buffer plus a + per-row ``lens`` vector (one row per output position). Each top-level element + becomes a single row; the Rust side reshapes back to ``list[list[float]]``, + matching ``meta_info["hidden_states"]``'s common per-position-vector shape. + """ + vals: List[float] = [] + lens: List[int] = [] + if not hs: + return vals, lens + for row in hs: + flat = _flatten_floats(row) + vals.extend(flat) + lens.append(len(flat)) + return vals, lens + + +def _flatten_floats(x): + """Recursively flatten a (possibly nested) float structure into a flat list + of floats — handles the ``float | list[float]`` union inside a hidden-state + chunk.""" + if isinstance(x, (int, float)): + return [float(x)] + out: List[float] = [] + for e in x: + out.extend(_flatten_floats(e)) + return out + + +class FlatPairColumns: + """A flat val/idx column pair (e.g. per-token logprob values + token ids): + per-request element counts in the header, concatenated f32 + i32 buffers in + the data. ``first_none_to_nan`` maps a leading ``None`` cell element to NaN + (the input-logprob first-prompt-token sentinel).""" + + def __init__(self, name, vals, idxs, first_none_to_nan=False): + self.name = name + self.vals = vals + self.idxs = idxs + self.first_none_to_nan = first_none_to_nan + self.v = array("f") + self.i = array("i") + self.lens = [] + + def columns(self): + return ((f"{self.name}_val", self.vals), (f"{self.name}_idx", self.idxs)) + + def accept(self, j): + vv = (self.vals[j] if self.vals else None) or [] + ii = (self.idxs[j] if self.idxs else None) or [] + # Parity assert, the flat twin of `flatten_ragged`'s: only `len(vv)` is + # recorded in `lens`, but both buffers are extended, so a longer idx column + # silently pushes every LATER column's offset out by the difference. The + # decoder cannot catch it — the data buffer only grows, so the receiver's + # bounds check still passes and it hands the client another column's bytes + # reinterpreted as logprobs, with a 200. + assert len(ii) == len( + vv + ), f"{self.name}: request {j} has {len(ii)} idx entries but {len(vv)} vals" + if self.first_none_to_nan and vv and vv[0] is None: + self.v.append(float("nan")) + self.v.extend(vv[1:]) + else: + self.v.extend(vv) + self.i.extend(ii) + self.lens.append(len(vv)) + + def header_cols(self): + return [self.lens] + + def data_cols(self): + return [self.v.tobytes(), self.i.tobytes()] + + +class RaggedPairColumns: + """A per-position ragged val/idx column pair (e.g. top-k / token-ids + logprobs): per-request position counts + a flat per-position length stream + in the header, concatenated f32/i32 buffers in the data.""" + + def __init__(self, name, vals, idxs): + self.name = name + self.vals = vals + self.idxs = idxs + self.v = array("f") + self.i = array("i") + self.pos = [] + self.req = [] + + def columns(self): + return ((f"{self.name}_val", self.vals), (f"{self.name}_idx", self.idxs)) + + def accept(self, j): + fv, fi, lens = flatten_ragged( + self.vals[j] if self.vals else None, + self.idxs[j] if self.idxs else None, + ) + self.v.extend(fv) + self.i.extend(fi) + self.pos.extend(lens) + self.req.append(len(lens)) + + def header_cols(self): + return [self.req, self.pos] + + def data_cols(self): + return [self.v.tobytes(), self.i.tobytes()] + + +class NestedRowColumns: + """A nested-rows float column (e.g. hidden states): per-request row counts + + per-row length stream in the header, one concatenated f32 buffer in the + data.""" + + def __init__(self, name, rows): + self.name = name + self.rows = rows + self.v = array("f") + self.pos = [] + self.req = [] + + def columns(self): + return ((self.name, self.rows),) + + def accept(self, j): + hv, hlens = flatten_hidden(self.rows[j] if self.rows else None) + self.v.extend(hv) + self.pos.extend(hlens) + self.req.append(len(hlens)) + + def header_cols(self): + return [self.req, self.pos] + + def data_cols(self): + return [self.v.tobytes()] diff --git a/python/sglang/test/run_eval.py b/python/sglang/test/run_eval.py index 3ea3db287..4b8cc109f 100644 --- a/python/sglang/test/run_eval.py +++ b/python/sglang/test/run_eval.py @@ -16,6 +16,7 @@ from sglang.test.simple_eval_common import ( ChatCompletionSampler, CompletionSampler, Eval, + GenerateSampler, make_report, set_ulimit, ) @@ -81,6 +82,13 @@ def run_eval_once(args, base_url: str, eval_obj: Eval) -> dict: **common_kwargs, stop=stop, ) + elif api_mode == "generate": + # SGLang-native `/generate` (raw text + sampling_params), same stop defaults. + stop = getattr(args, "stop", ["Question", "Assistant:", "<|separator|>"]) + sampler = GenerateSampler( + **common_kwargs, + stop=stop, + ) else: sampler = ChatCompletionSampler( **common_kwargs, @@ -454,8 +462,8 @@ if __name__ == "__main__": "--api", type=str, default="chat", - choices=["chat", "completion"], - help="API mode: 'chat' for /v1/chat/completions, 'completion' for /v1/completions", + choices=["chat", "completion", "generate"], + help="API mode: 'chat' for /v1/chat/completions, 'completion' for /v1/completions, 'generate' for SGLang-native /generate", ) parser.add_argument("--num-examples", type=int) parser.add_argument("--num-threads", type=int, default=512) diff --git a/python/sglang/test/simple_eval_common.py b/python/sglang/test/simple_eval_common.py index 235916554..6970643c1 100644 --- a/python/sglang/test/simple_eval_common.py +++ b/python/sglang/test/simple_eval_common.py @@ -254,6 +254,100 @@ class CompletionSampler(SamplerBase): return "" +class GenerateSampler(SamplerBase): + """ + Sample from SGLang's native ``/generate`` endpoint (not the OpenAI-compatible + API). Sends raw text prompts with `sampling_params`, so it exercises the same + path as `bench_serving` rather than the `/v1/completions` wrapper. + + `base_url` is the OpenAI-style URL the eval harness builds (``.../v1``); the + trailing ``/v1`` is stripped to reach the server root's ``/generate``. + """ + + def __init__( + self, + base_url: str = None, + model: Optional[str] = None, + temperature: float = 0.0, + top_p: float = 1.0, + max_tokens: int = 2048, + stop: Optional[List[str]] = None, + ): + self.client = LargerHttpxClient() + + # The harness passes the OpenAI base (`.../v1`); `/generate` lives at root. + root = (base_url or "http://127.0.0.1:30000/v1").rstrip("/") + if root.endswith("/v1"): + root = root[: -len("/v1")] + self.generate_url = f"{root}/generate" + + # `/generate` serves the loaded model and ignores a model field, so `model` + # is informational only; fill it from `/get_model_info` when unset. + if model is None: + try: + info = self.client.get(f"{root}/get_model_info").json() + model = info.get("model_path") + except Exception: + model = None + + self.model = model + self.temperature = temperature + self.top_p = top_p + self.max_tokens = max_tokens + self.stop = stop + self._completion_tokens: list[int] = [] + print( + f"GenerateSampler initialized with {self.generate_url=} {self.model=} " + f"{self.temperature=} {self.max_tokens=} {self.stop=}" + ) + + def _pack_message(self, role: str, content: Any): + return {"role": str(role), "content": content} + + def __call__(self, message_list: MessageList) -> str: + # Extract raw text from message list (eval objects pack prompt as a single user message) + prompt = "\n".join( + msg["content"] + for msg in message_list + if isinstance(msg.get("content"), str) + ) + payload = { + "text": prompt, + "sampling_params": { + "temperature": self.temperature, + "top_p": self.top_p, + "max_new_tokens": self.max_tokens, + "stop": self.stop, + }, + "stream": False, + } + trial = 0 + while trial < 6: + try: + response = self.client.post(self.generate_url, json=payload) + # A 400 is a malformed request, not a transient failure — don't retry. + if response.status_code == 400: + print("Bad Request Error", response.text) + return "" + response.raise_for_status() + data = response.json() + meta_info = data.get("meta_info") or {} + completion_tokens = meta_info.get("completion_tokens") + if completion_tokens is not None: + self._completion_tokens.append(completion_tokens) + return data.get("text") or "" + except Exception as e: + exception_backoff = 2**trial + print( + f"Rate limit exception so wait and retry {trial} after {exception_backoff} sec", + e, + ) + time.sleep(exception_backoff) + trial += 1 + print(f"All retry attempts exhausted for request. Returning empty response.") + return "" + + QUERY_TEMPLATE_MULTICHOICE = """ Answer the following multiple choice question. The last line of your response should be of the following format: 'Answer: $LETTER' (without quotes) where LETTER is one of ABCD. Think step by step before answering. diff --git a/python/sglang/test/test_utils.py b/python/sglang/test/test_utils.py index 6140b032f..7a6b1f641 100644 --- a/python/sglang/test/test_utils.py +++ b/python/sglang/test/test_utils.py @@ -4,6 +4,7 @@ import argparse import asyncio import copy import doctest +import importlib.util import inspect import json import logging @@ -202,6 +203,23 @@ def is_h200_system(): return envs.IS_H200.get() +def is_rust_server_built(): + """Return whether the embedded Rust server extension (``SGLANG_RUST_SERVER``) + is importable. + + ``sglang/srt/server/`` is not in the source tree — it is produced by + ``setup.py build_rust --inplace``, so on a build without it ``find_spec`` + raises ``ModuleNotFoundError`` for the missing *parent* package rather than + returning ``None`` for the missing leaf. Suites gate a rust-server subclass on + this at class-definition time, so letting that escape would fail the whole + module import instead of skipping the one class. + """ + try: + return importlib.util.find_spec("sglang.srt.server._core") is not None + except ModuleNotFoundError: + return False + + def _use_cached_default_models(model_repo: str): cache_dir = os.getenv("DEFAULT_MODEL_CACHE_DIR") if cache_dir and model_repo: diff --git a/rust/sglang-server/Cargo.toml b/rust/sglang-server/Cargo.toml index 5429ce03d..bb022754a 100644 --- a/rust/sglang-server/Cargo.toml +++ b/rust/sglang-server/Cargo.toml @@ -31,8 +31,9 @@ uuid = { workspace = true } axum = { version = "0.8.9", features = ["json", "tokio"] } core_affinity = "0.8" -# the dynamo-tokenizers is deps on hf-hub, should bump version together -dynamo-tokenizers = "1.5.3" +# pin the dynamo-tokenizers for now, 1.7.0 enables`serde_json/preserve_order` +# transitively, which reorders every JSON object. +dynamo-tokenizers = "=1.5.3" flume = "0.12.0" itertools = "0.14" hf-hub = { version = "0.4", default-features = false } diff --git a/rust/sglang-server/src/tokenizer.rs b/rust/sglang-server/src/tokenizer.rs index ee90d605a..2cc43bb26 100644 --- a/rust/sglang-server/src/tokenizer.rs +++ b/rust/sglang-server/src/tokenizer.rs @@ -144,6 +144,20 @@ impl Runnable for TokenizerWorker { tracing::error!("tokenizer pool received a non-generate request"); continue; }; + // Size the scheduler's stop-match window in TOKENS, as Python's + // `normalize(tokenizer)` does. + let stop_tokens = g + .sampling_params + .stop_strs + .iter() + // A stop that won't encode falls back to its byte length rather + // than failing the request: still an over-estimate, never an + // under-estimate, so the scheduler cannot miss that stop. + .map(|s| self.tokenizer.encode(s).map_or(s.len(), |ids| ids.len())) + .max(); + if let Some(n) = stop_tokens { + g.sampling_params.stop_str_max_len = n; + } match self.tokenizer.encode(g.text.as_deref().unwrap_or("")) { Ok(ids) => { g.input_ids = Some(ids); @@ -160,3 +174,68 @@ impl Runnable for TokenizerWorker { } } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::fsm::RequestState; + use crate::message::{EgressSink, GenerateRequest, RequestKind, SamplingParams}; + use tokio::sync::mpsc; + + /// One token per whitespace-separated word, so a stop's token count differs + /// from its byte count and the two units cannot be confused. + struct WordTokenizer; + impl TextTokenizer for WordTokenizer { + fn encode(&self, text: &str) -> Result { + Ok(text.split_whitespace().map(|_| 1i32).collect()) + } + } + + /// The scheduler's stop-match window must reach the wire as a TOKEN count, as + /// Python's `normalize(tokenizer)` produces. + /// + /// `Normalizing` leaves a UTF-8 BYTE count there — a safe over-estimate, but it + /// makes the scheduler decode a longer tail on EVERY decode step of EVERY + /// request (14 tokens vs 6 for a typical stop set). This stage owns the + /// tokenizer, so it is where the exact count is resolved. + #[test] + fn tokenizing_replaces_the_byte_window_with_a_token_count() { + let (req_tx, req_rx) = flume::unbounded::(); + let (tm_tx, tm_rx) = flume::unbounded::(); + + // 8 bytes vs 3 "tokens" under WordTokenizer — units are distinguishable. + let sp = SamplingParams { + stop_strs: vec!["a bb ccc".to_string(), "dd".to_string()], + stop_str_max_len: 8, // what `normalize_stops` left: max BYTE length + ..Default::default() + }; + let (sink_tx, _sink_rx) = mpsc::channel(4); + req_tx + .send(Request { + rid: "1".into(), + state: RequestState::Tokenizing, + sink: EgressSink::Local(sink_tx), + kind: RequestKind::Generate(Box::new(GenerateRequest { + rid: "1".into(), + text: Some("hello world".into()), + sampling_params: sp, + ..Default::default() + })), + }) + .expect("send"); + drop(req_tx); // closes the loop after one request + + TokenizerWorker::new(req_rx, tm_tx, Arc::new(WordTokenizer)).run(); + + let TmEvent::Tokenized(req) = tm_rx.try_recv().expect("returned") else { + panic!("expected Tokenized"); + }; + let RequestKind::Generate(g) = &req.kind else { + panic!("expected generate"); + }; + assert_eq!( + g.sampling_params.stop_str_max_len, 3, + "must be the max TOKEN count (3), not the byte count (8)" + ); + } +} diff --git a/test/registered/core/test_srt_endpoint.py b/test/registered/core/test_srt_endpoint.py index 81f3fa7ba..af5cc9a19 100644 --- a/test/registered/core/test_srt_endpoint.py +++ b/test/registered/core/test_srt_endpoint.py @@ -2,6 +2,8 @@ python3 -m unittest test_srt_endpoint.TestSRTEndpoint.test_simple_decode python3 -m unittest test_srt_endpoint.TestSRTEndpoint.test_logprob_with_chunked_prefill python3 -m unittest test_srt_endpoint.TestTokenizeDetokenize +python3 -m unittest test_srt_endpoint.TestRustServerEndpoint +python3 -m unittest test_srt_endpoint.TestRustServerLogprob """ import json @@ -24,17 +26,22 @@ from sglang.test.test_utils import ( DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, DEFAULT_URL_FOR_TEST, CustomTestCase, + is_rust_server_built, popen_launch_server, run_logprob_check, ) -register_cuda_ci(est_time=160, stage="base-b", runner_config="1-gpu-small") -register_amd_ci(est_time=160, suite="stage-b-test-1-gpu-small-amd") +register_cuda_ci(est_time=260, stage="base-b", runner_config="1-gpu-small") +register_amd_ci(est_time=260, suite="stage-b-test-1-gpu-small-amd") SERVER_ENV = {"SGLANG_USE_PICKLE_IPC": "0"} class TestSRTEndpoint(CustomTestCase): + # Extra server-launch env; subclasses override to run the same suite + # against a different server flavor (e.g. SGLANG_RUST_SERVER=1). + env = {} + @classmethod def setUpClass(cls): cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST @@ -46,7 +53,7 @@ class TestSRTEndpoint(CustomTestCase): # The tiny logprob chunk size routes this file's logprob tests # through the multi-chunk stitching path (requests at or below 64 # rows still cover the non-chunked path). - env={**SERVER_ENV, "SGLANG_LOGPROB_CHUNK_SIZE": "64"}, + env={**cls.env, **SERVER_ENV, "SGLANG_LOGPROB_CHUNK_SIZE": "64"}, other_args=( "--enable-custom-logit-processor", "--mem-fraction-static", @@ -844,5 +851,74 @@ class TestTokenizeDetokenize(CustomTestCase): self.assertEqual(r2.status_code, 500) +# --------------------------------------------------------------------------- +# Embedded Rust server (SGLANG_RUST_SERVER=1): rerun the whole endpoint suite +# against the rust api-server/tokenizer/detokenizer stack — the logprob tests +# exercise the columnar egress wire (`push_generation` extras -> Rust +# `BatchHeader`/`for_each_chunk` -> detok reshape) end to end. Suite +# surface the rust server does not implement yet is skipped explicitly below. +# --------------------------------------------------------------------------- +@unittest.skipUnless( + is_rust_server_built(), + "embedded rust server extension not built (e.g. AMD suite)", +) +class TestRustServerEndpoint(TestSRTEndpoint): + env = {"SGLANG_RUST_SERVER": "1"} + + _RUST_TODO = "not implemented by the embedded Rust server yet" + + @unittest.skip(f"custom_logit_processor request field {_RUST_TODO}") + def test_custom_logit_processor(self): + pass + + @unittest.skip(f"custom_logit_processor request field {_RUST_TODO}") + def test_custom_logit_processor_batch_mixed(self): + pass + + @unittest.skip(f"custom_logit_processor request field {_RUST_TODO}") + def test_stateful_custom_logit_processor(self): + pass + + @unittest.skip(f"custom_logit_processor request field {_RUST_TODO}") + def test_stateful_custom_logit_processor_batch_mixed(self): + pass + + @unittest.skip(f"/flush_cache endpoint + cached_tokens meta {_RUST_TODO}") + def test_cache_tokens(self): + pass + + def test_greedy_token_equals_top1(self): + """Cross-column alignment guard for the columnar logprob wire: at + temperature 0 the chosen token must BE the top-1 entry of its own + position. A column shifted across requests or positions (the failure + mode a truncation-tolerant reader would mask) breaks this instantly.""" + response = requests.post( + self.base_url + "/generate", + json={ + "text": ["The capital of France is", "I have a very good idea on"], + "sampling_params": {"temperature": 0, "max_new_tokens": 8}, + "return_logprob": True, + "top_logprobs_num": 5, + "logprob_start_len": 0, + }, + ) + self.assertEqual(response.status_code, 200, response.text) + for res in response.json(): + meta = res["meta_info"] + out_lp = meta["output_token_logprobs"] + top = meta["output_top_logprobs"] + self.assertEqual(len(out_lp), meta["completion_tokens"]) + self.assertEqual(len(top), len(out_lp)) + # First prompt token's logprob is the None sentinel; it must + # survive the NaN wire encoding and come back as null. + self.assertIsNone(meta["input_token_logprobs"][0][0]) + for (lp, tid, _), pos_top in zip(out_lp, top): + self.assertEqual(len(pos_top), 5) + self.assertEqual(pos_top[0][1], tid) + self.assertAlmostEqual(pos_top[0][0], lp, places=4) + vals = [t[0] for t in pos_top] + self.assertEqual(vals, sorted(vals, reverse=True)) + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/quant/test_modelopt_fp8.py b/test/registered/quant/test_modelopt_fp8.py index d753b559f..0e5b76fe0 100644 --- a/test/registered/quant/test_modelopt_fp8.py +++ b/test/registered/quant/test_modelopt_fp8.py @@ -9,13 +9,21 @@ from sglang.test.test_utils import ( DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, DEFAULT_URL_FOR_TEST, CustomTestCase, + is_rust_server_built, popen_launch_server, ) -register_cuda_ci(est_time=62, stage="base-b", runner_config="1-gpu-large") +# Two classes run from this file: the default server plus the Rust-frontend +# variant (when the embedded extension is built), each launches a server + eval. +register_cuda_ci(est_time=124, stage="base-b", runner_config="1-gpu-large") class TestModeloptFP8(CustomTestCase): + # Extra server env; the Rust-frontend subclass sets SGLANG_RUST_SERVER here. + env = None + # Eval endpoint. The Rust server exposes only the native `/generate`, so its + # subclass overrides this to "generate". + api = "completion" @classmethod def setUpClass(cls): @@ -25,7 +33,15 @@ class TestModeloptFP8(CustomTestCase): cls.model, cls.base_url, timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - other_args=["--quantization", "modelopt_fp8"], + other_args=[ + "--quantization", + "modelopt_fp8", + "--tokenizer-worker-num", + "2", + "--detokenizer-worker-num", + "2", + ], + env=cls.env, ) @classmethod @@ -38,7 +54,7 @@ class TestModeloptFP8(CustomTestCase): base_url=self.base_url, model=self.model, eval_name="gsm8k", - api="completion", + api=self.api, max_tokens=512, num_examples=200, num_threads=200, @@ -48,5 +64,20 @@ class TestModeloptFP8(CustomTestCase): self.assertGreater(metrics["score"], 0.70) +@unittest.skipUnless( + is_rust_server_built(), + "embedded rust server extension not built", +) +class TestModeloptFP8WithRustServer(TestModeloptFP8): + """Same model + eval, but served through the embedded Rust frontend + (`SGLANG_RUST_SERVER`). Guards the Rust tokenizer/detokenizer/completions path + against accuracy regressions: a bug there drops gsm8k score below the same + 0.70 bar the default frontend must clear. Uses the native `/generate` endpoint + (the only API the Rust server exposes).""" + + env = {"SGLANG_RUST_SERVER": "1"} + api = "generate" + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/rust/test_cargo_workspace.py b/test/registered/rust/test_cargo_workspace.py new file mode 100644 index 000000000..e287aa16d --- /dev/null +++ b/test/registered/rust/test_cargo_workspace.py @@ -0,0 +1,76 @@ +"""Run the `rust/` Cargo workspace's unit tests from the CPU CI suite. + +The `rust/` workspace (sglang-grpc, sglang-mm, sglang-server) is compiled into +the wheel by setuptools-rust, but until now nothing ran `cargo test` in CI -- +`.github/workflows/pr-test-rust.yml` and `pr-benchmark-rust.yml` are both +path-scoped to `sgl-model-gateway/**`, a different workspace. `lint.yml` covers +rustfmt/clippy via the pre-commit hooks, so this file only adds the test run. + +The debug profile is deliberate: these are pure-logic tests (no timing or +codegen assertions), and the release profile costs a full LTO build for the +same coverage. +""" + +import shutil +import subprocess +import unittest +from pathlib import Path + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +# base-c-test-cpu is where this was asked for, and it matches the repo's +# base-a + base-c dual-registration convention -- but base-c-test-cpu currently +# has no runner job in any workflow (it was carved out of base-b in #28623 to +# *reduce* CPU CI scope), so base-a-test-cpu is what actually executes. +register_cpu_ci(est_time=300, suite="base-a-test-cpu") + +# repo root: test/registered/rust/ +RUST_WORKSPACE = Path(__file__).resolve().parents[3] / "rust" + +# Not `est_time`: that is a scheduling hint for partition balancing (a rough +# average), this is a hard ceiling for the worst case. The 136 tests run in ~1s; +# what varies is the build. Cache-warm the workspace crates recompile in ~15s, +# but a Swatinem/rust-cache miss rebuilds all ~370 dependencies -- measured at +# 48s on 4 fast cores, so several minutes on a hosted runner. +# +# Capped below the 600s `timeout-minutes` on the suite's "Run test" step so a +# hang fails here, with output, instead of being killed as an opaque job +# timeout. The harness `--timeout-per-file` (1200s) is looser still. +BUILD_AND_RUN_TIMEOUT_S = 300 + + +class TestCargoWorkspace(CustomTestCase): + def test_cargo_test_workspace(self): + # Not skipUnless: cargo is a hard dependency of the editable install + # (setuptools-rust builds sglang-grpc), so a missing toolchain is a + # broken environment, and a silently-skipped CI test is worthless. + self.assertIsNotNone( + shutil.which("cargo"), + "cargo not found on PATH; install a Rust toolchain " + "(scripts/ci/utils/install_rust_protoc.sh)", + ) + self.assertTrue( + (RUST_WORKSPACE / "Cargo.toml").is_file(), + f"rust workspace manifest not found at {RUST_WORKSPACE}", + ) + + proc = subprocess.run( + ["cargo", "test", "--workspace"], + cwd=RUST_WORKSPACE, + capture_output=True, + text=True, + timeout=BUILD_AND_RUN_TIMEOUT_S, + ) + # Print unconditionally so a green run still shows which tests ran. + print(proc.stdout) + self.assertEqual( + proc.returncode, + 0, + f"`cargo test --workspace` failed in {RUST_WORKSPACE}\n" + f"--- stdout ---\n{proc.stdout}\n--- stderr ---\n{proc.stderr}", + ) + + +if __name__ == "__main__": + unittest.main()