minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Rainj me/rust server refactor2 (#35239)

Browse Source
This commit is contained in:
Rain Jiang
2026-08-21 16:37:02 -07:00
committed by GitHub
parent fe8f9d7457
commit 7d7ab4b5c6
61 changed files with 2531 additions and 2193 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/_pr-test-check-changes.yml
+13
View File
@@ -25,6 +25,8 @@ on:
value: ${{ jobs.run.outputs.jit_kernel }}
multimodal_gen:
value: ${{ jobs.run.outputs.multimodal_gen }}
rust_workspace:
value: ${{ jobs.run.outputs.rust_workspace }}
partitions:
value: ${{ jobs.run.outputs.partitions }}
partition_model_sha:
@@ -45,6 +47,7 @@ jobs:
sgl_kernel: ${{ steps.filter.outputs.sgl_kernel }}
jit_kernel: ${{ steps.filter.outputs.jit_kernel || steps.run-mode.outputs.run_all_tests }}
multimodal_gen: ${{ steps.filter.outputs.multimodal_gen || steps.run-mode.outputs.run_all_tests }}
rust_workspace: ${{ steps.filter.outputs.rust_workspace || steps.run-mode.outputs.run_all_tests }}
partitions: ${{ steps.partitions.outputs.partitions }}
partition_model_sha: ${{ steps.partition-model-sha.outputs.sha }}
runs_on_map: ${{ steps.runner-map.outputs.runs_on_map }}
@@ -89,6 +92,8 @@ jobs:
- "scripts/ci/cuda/*"
- "scripts/ci/utils/*"
- "test/**/!(*.md)"
- "rust/**"
- "proto/sglang/runtime/v1/sglang.proto"
multimodal_gen:
- ".github/workflows/pr-test.yml"
- ".github/workflows/pr-test-multimodal-gen.yml"
@@ -113,6 +118,13 @@ jobs:
# Intentionally excludes ".github/workflows/pr-test-sgl-kernel.yml" —
# see API-side detector below for rationale.
- "python/sglang/kernels/aot/**/!(*.md|THIRDPARTYNOTICES.txt|LICENSE)"
rust_workspace:
# Gates the cargo test in test/registered/rust/; exported to the
# CPU stage as its negation, SGLANG_SKIP_RUST_TESTS.
- "rust/**"
- "proto/sglang/runtime/v1/sglang.proto"
- "test/registered/rust/**"
- ".github/workflows/_pr-test-*.yml"
- name: Determine full-parallel mode
id: parallel-mode
@@ -223,6 +235,7 @@ jobs:
echo "| sgl_kernel | ${{ steps.filter.outputs.sgl_kernel }} |"
echo "| jit_kernel | ${{ steps.filter.outputs.jit_kernel || steps.run-mode.outputs.run_all_tests }} |"
echo "| multimodal_gen | ${{ steps.filter.outputs.multimodal_gen || steps.run-mode.outputs.run_all_tests }} |"
echo "| rust_workspace | ${{ steps.filter.outputs.rust_workspace || steps.run-mode.outputs.run_all_tests }} |"
echo "| b200_runner | ${{ steps.set-runner.outputs.b200_runner }} |"
echo "| enable_retry | ${{ steps.set-retry.outputs.enable_retry }} |"
echo "| continue_on_error | ${{ steps.set-continue-on-error.outputs.continue_on_error }} |"
.github/workflows/_pr-test-stage-cpu.yml
+4 -1
View File
@@ -13,7 +13,7 @@ on:
type: string
required: true
check_changes:
description: 'toJson(needs.check-changes.outputs). Read via fromJson(...).main_package / continue_on_error / partition_model_sha.'
description: 'toJson(needs.check-changes.outputs). Read via fromJson(...).main_package / rust_workspace / continue_on_error / partition_model_sha.'
type: string
required: true
caller_inputs:
@@ -153,6 +153,9 @@ jobs:
timeout-minutes: ${{ fromJson(inputs.run_timeout_minutes) }}
env:
CONTINUE_ON_ERROR_FLAG: ${{ fromJson(inputs.check_changes).continue_on_error == 'true' && '--continue-on-error' || '' }}
# Normalized to a literal true/false: EnvBool rejects an empty string,
# which is what a missing check-changes output would expand to.
SGLANG_SKIP_RUST_TESTS: ${{ fromJson(inputs.check_changes).rust_workspace == 'true' && 'false' || 'true' }}
run: |
cd test/
python3 run_suite.py --hw cpu --suite ${{ inputs.self_name }} --auto-partition-id ${{ matrix.partition }} --auto-partition-size ${{ fromJson(inputs.partitions)[inputs.self_name].size }} --partition-model-file /tmp/partition-model.json $CONTINUE_ON_ERROR_FLAG
.github/workflows/lint.yml
-17
View File
@@ -13,16 +13,6 @@ jobs:
with:
fetch-depth: 0
- name: Detect Rust workspace changes
id: paths
uses: dorny/paths-filter@v3
with:
filters: |
rust_workspace:
- 'rust/**'
- 'proto/sglang/runtime/v1/sglang.proto'
- '.github/workflows/lint.yml'
# Fail-fast gate: docs_new/ was renamed to docs/ (#32123). git's rename
# detection silently re-adds NEW files under docs_new/ on merge without a
# conflict, which would resurrect the directory. We check the checked-out
@@ -60,13 +50,6 @@ jobs:
- name: Run pre-commit checks
run: SKIP=no-commit-to-branch pre-commit run --all-files --show-diff-on-failure
# Not in the rust-ext build job: its cache key covers the built .so files,
# and a test script is not a build input. Tests take ~1s; the timeout is
# for a cold cache, which codegens the dependency graph first.
- name: Run rust/ workspace tests
if: steps.paths.outputs.rust_workspace == 'true'
run: cd rust && timeout 900 cargo test --workspace
- name: Run lychee docs checks (offline references)
uses: lycheeverse/lychee-action@8646ba30535128ac92d33dfc9133794bfdd9b411 # v2
with:
python/sglang/srt/environ.py
+3
View File
@@ -351,6 +351,9 @@ class Envs:
# ===================================================================
SGLANG_IS_IN_CI = EnvBool(False)
SGLANG_IS_IN_CI_AMD = EnvBool(False)
# Set to true by the check-changes CI job when a PR touches nothing under
# rust/; default false so local and scheduled runs never skip the cargo tests.
SGLANG_SKIP_RUST_TESTS = EnvBool(False)
SGLANG_TEST_MAX_RETRY = EnvInt(None)
# Expand jit_kernel test grids to their full parameter ranges (nightly).
SGLANG_JIT_KERNEL_RUN_FULL_TESTS = EnvBool(False)
python/sglang/srt/managers/rust_server.py
+116 -36
View File
@@ -4,13 +4,14 @@ 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
typed `server_args` handoff, and control-response routing — out of `scheduler.py`. The
scheduler holds an `Optional[RustServer]` and delegates to it.
"""
from __future__ import annotations
import importlib
import json
import logging
import os
from array import array
@@ -37,7 +38,7 @@ if TYPE_CHECKING:
from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.managers.io_struct import BatchTokenIDOutput
from sglang.srt.managers.scheduler import Scheduler
from sglang.srt.rust_extensions._server import Server
from sglang.srt.rust_extensions._server import MmSpec, Server, ServerArgs
from sglang.srt.server_args import ServerArgs
logger = logging.getLogger(__name__)
@@ -45,8 +46,10 @@ logger = logging.getLogger(__name__)
class NativeMmSpec(msgspec.Struct, frozen=True, kw_only=True):
"""Resolved parameters of the native Rust MM pipeline for one model,
consumed by the Rust worker pool (:meth:`rust_json`) and the drain
adapter (:meth:`NativeMmHost.build_native_mm`)."""
consumed by the Rust worker pool (as the typed extension ``MmSpec``, see
:meth:`RustServer._build_mm_spec`), the ``_multimodal`` parity API
(:meth:`rust_json`) and the drain adapter
(:meth:`NativeMmHost.build_native_mm`)."""
family: str
feature_shm: bool
@@ -73,7 +76,9 @@ class NativeMmSpec(msgspec.Struct, frozen=True, kw_only=True):
return 3 * self.temporal_patch_size * self.patch_size * self.patch_size
def rust_json(self) -> str:
"""The subset `sglang_mm::registry::pipeline_from_spec` parses."""
"""The subset `sglang_mm::registry::pipeline_from_spec` parses — the
JSON form the ``_multimodal`` parity API takes; the server itself is
handed the typed ``MmSpec`` instead."""
fields = (f for f in self.__struct_fields__ if f not in self.DRAIN_ONLY)
return msgspec.json.encode({f: getattr(self, f) for f in fields}).decode()
@@ -272,7 +277,7 @@ class NativeMmHost:
@staticmethod
def build_native_mm(spec: NativeMmSpec, entry):
"""Drain-time adapter: wrap the Rust-produced buffers of one ``MmHandoff``
"""Drain-time adapter: wrap the Rust-produced buffers of one ``MmEncodeResult``
into the scheduler's ``MultimodalProcessorOutput``. Wrapping only — load,
resize, patchify, token expansion and M-RoPE all ran in Rust.
@@ -405,10 +410,10 @@ class RustServer:
)
server = Server(
cls._build_server_args(scheduler),
# None -> run unpinned; the list carries the pinning decision.
cores=server_cores,
http_addr=http_addr,
server_args_json=cls._build_server_args(scheduler),
)
# Multimodal models must have a native Rust pipeline — there is no Python
@@ -443,7 +448,7 @@ class RustServer:
f"(supported: {', '.join(supported)}; "
"images only). Unset SGLANG_RUST_SERVER to serve this model."
)
server.start_mm_workers(mm_spec.rust_json(), mm_host.mm_workers)
server.start_mm_workers(cls._build_mm_spec(mm_spec), mm_host.mm_workers)
# Narrow the scheduler thread only after the server threads are launched.
if launch_cores is not None:
@@ -466,11 +471,11 @@ class RustServer:
return cls(server, mm_spec=mm_spec)
def wait_ingress(self, timeout_ms: int) -> None:
def wait_request(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)
self.server.wait_request(timeout_ms)
def drain(self, max_recv: int) -> List[Any]:
"""Ingress: non-blocking drain of the in-process ring → list of decoded
@@ -484,8 +489,10 @@ class RustServer:
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`.
never waits: the ring drain is `try_recv` (returns the instant the ring
is dry, capped at `max_recv`) and the rest is one memcpy per header
plus one for the concatenated ids — same contract as `zmq.NOBLOCK`.
Parking for work is :meth:`wait_request`, which does release the GIL.
"""
limit = max_recv if max_recv > 0 else self._max_per_poll
batch = self.server.recv_requests(limit)
@@ -553,7 +560,7 @@ class RustServer:
# rendering happens in Rust.
encoded = msgspec.msgpack.encode(payload, enc_hook=str)
self.server.push_result(recv_req.rid, encoded)
self.server.push_control_result(recv_req.rid, encoded)
def push_generation(self, payload: BatchTokenIDOutput) -> None:
"""Egress redirect for generation output (replaces the zmq detokenizer).
@@ -698,38 +705,111 @@ class RustServer:
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):
if not self.server.push_decode_result_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``)."""
def _build_mm_spec(spec: NativeMmSpec) -> MmSpec:
"""The typed MM handoff for ``Server.start_mm_workers``: the
:class:`NativeMmSpec` fields the Rust pipeline consumes, as the Rust
extension's own ``MmSpec`` class (same required-keyword contract as
:meth:`_build_server_args`; ``family`` / ``resample`` become the
extension's ``MmFamily`` / ``MmResample`` enums)."""
from sglang.srt.rust_extensions import load_rust_extension
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
# Resolved default sampling params (generation_config.json when
# `--sampling-defaults model`, {} otherwise). The rust server consumes
# these for omitted temperature/top_p in chat conversions instead of
# hard-coding the OpenAI terminal defaults.
model_config["default_sampling_params"] = (
scheduler.model_config.get_default_sampling_params()
ext = load_rust_extension("sglang.srt.rust_extensions._server")
family = {"qwen_vl": ext.MmFamily.QwenVl}[spec.family]
resample = {"aten_u8": ext.MmResample.AtenU8, "pil": ext.MmResample.Pil}[
spec.resample
]
return ext.MmSpec(
family=family,
feature_shm=spec.feature_shm,
image_token_id=spec.image_token_id,
patch_size=spec.patch_size,
merge_size=spec.merge_size,
temporal_patch_size=spec.temporal_patch_size,
min_pixels=spec.min_pixels,
max_pixels=spec.max_pixels,
image_mean=spec.image_mean,
image_std=spec.image_std,
resample=resample,
)
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()
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 _build_server_args(scheduler: Scheduler) -> ServerArgs:
"""The typed launch handoff for the scheduler's embedded Rust server:
the ``server_args`` fields it reads, the already-resolved
``model_config``, and launch-time facts — as the Rust extension's own
``ServerArgs`` class. Its constructor takes every field as a required
keyword (see ``rust/sglang-server/src/message/config.rs``), so a
missing, extra or mistyped field fails here at boot rather than
running on a silently-defaulted knob."""
from sglang.srt.rust_extensions import load_rust_extension
ext = load_rust_extension("sglang.srt.rust_extensions._server")
sa = scheduler.server_args
mc = scheduler.model_config
disaggregation_mode = {
"null": ext.DisaggregationMode.Null,
"prefill": ext.DisaggregationMode.Prefill,
"decode": ext.DisaggregationMode.Decode,
}[sa.disaggregation_mode]
return ext.ServerArgs(
model_path=sa.model_path,
served_model_name=sa.served_model_name,
tokenizer_path=sa.tokenizer_path,
revision=sa.revision,
load_format=sa.load_format,
weight_version=sa.weight_version,
host=sa.host,
port=sa.port,
log_level=sa.log_level,
log_level_http=sa.log_level_http,
chat_template=sa.chat_template,
tool_call_parser=sa.tool_call_parser,
reasoning_parser=sa.reasoning_parser,
stream_response_default_include_usage=sa.stream_response_default_include_usage,
tokenizer_worker_num=sa.tokenizer_worker_num,
detokenizer_worker_num=sa.detokenizer_worker_num,
skip_tokenizer_init=sa.skip_tokenizer_init,
incremental_streaming_output=sa.incremental_streaming_output,
disaggregation_mode=disaggregation_mode,
model_config=ext.ModelConfig(
context_len=mc.context_len,
vocab_size=mc.vocab_size,
is_multimodal=mc.is_multimodal,
# Resolved default sampling params (generation_config.json when
# `--sampling-defaults model`, {} otherwise). The rust server
# consumes these for omitted temperature/top_p in chat
# conversions instead of hard-coding the OpenAI terminal
# defaults.
default_sampling_params=ext.DefaultSamplingParams(
**mc.get_default_sampling_params()
),
),
# `preferred_sampling_params` is deliberately absent: `launch`
# refuses to start when it is set, so the Rust server never needs it.
preferred_sampling_params=(
json.dumps(sa.preferred_sampling_params)
if sa.preferred_sampling_params is not None
else None
),
allow_auto_truncate=sa.allow_auto_truncate,
enable_return_hidden_states=sa.enable_return_hidden_states,
# Not a `server_args` field: `TokenizerManager` derives it, and the
# rust ingress needs the same number for its total-token check.
num_reserved_tokens=compute_num_reserved_tokens(),
# 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).
version=__version__,
max_total_num_tokens=scheduler.max_total_num_tokens,
)
@staticmethod
def _partition_cores(
python/sglang/srt/managers/scheduler_components/idle_sleeper.py
+2 -2
View File
@@ -46,7 +46,7 @@ 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
Instead park directly on the ring: ``wait_request`` 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.
@@ -59,7 +59,7 @@ class RustServerIdleSleeper:
self.empty_cache_interval = envs.SGLANG_EMPTY_CACHE_INTERVAL.get()
def maybe_sleep(self):
self.rust_server.wait_ingress(self.timeout_ms)
self.rust_server.wait_request(self.timeout_ms)
if (
self.empty_cache_interval > 0
and real_time() - self.last_empty_time > self.empty_cache_interval
python/sglang/srt/managers/utils.py
+2 -5
View File
@@ -383,11 +383,8 @@ def compute_num_reserved_tokens() -> int:
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. Both stamp the number once at
launch, so it has to cover every step an adaptive-spec run may switch to:
it reads the bags for the candidate-table ceiling and the current
`topk * steps`, not the untouched startup record.
`server_args` handoff (`RustServer._build_server_args`), which needs the same
number to run the total-token check in Rust.
"""
spec = get_spec()
algorithm = SpeculativeAlgorithm.from_string(spec.speculative_algorithm)
rust/sglang-mm/src/registry.rs
+25 -14
View File
@@ -67,21 +67,32 @@ pub fn default_registry() -> ProcessorRegistry {
}
// --- Server (pure-Rust) request pipeline ---
/// Build a family processor from the Python-side spec JSON. `Err` on an
/// unknown family or malformed spec — the caller treats that as "no Rust
/// pipeline".
/// The resolved parameters of one family pipeline — the typed form of the
/// Python-side spec, one variant per family arm. `sglang-server` builds it
/// directly from its `MmSpec` pyclass; the JSON parity API reaches it through
/// [`pipeline_from_spec`], where the `family` key selects the variant.
#[derive(Clone, Debug, serde::Deserialize)]
#[serde(tag = "family", rename_all = "snake_case")]
pub enum PipelineSpec {
QwenVl(crate::qwen_vl::QwenVlSpec),
}
/// Build a family processor from a typed spec. `Err` when the family
/// rejects its parameters (e.g. a zero patch size).
pub fn build_pipeline(
spec: PipelineSpec,
) -> Result<Box<dyn crate::pipeline::MmFamilyProcessor>, String> {
match spec {
PipelineSpec::QwenVl(spec) => Ok(Box::new(crate::qwen_vl::QwenVlProcessor::new(spec)?)),
}
}
/// Build a family processor from the Python-side spec JSON
/// (`{"family": ..., resolved processor params}`). `Err` on an unknown family
/// or malformed spec — the caller treats that as "no Rust pipeline".
pub fn pipeline_from_spec(
json: &str,
) -> Result<Box<dyn crate::pipeline::MmFamilyProcessor>, String> {
#[derive(serde::Deserialize)]
struct Header {
family: String,
}
let header: Header = serde_json::from_str(json).map_err(|e| format!("mm spec: {e}"))?;
match header.family.as_str() {
"qwen_vl" => Ok(Box::new(crate::qwen_vl::QwenVlProcessor::from_spec_json(
json,
)?)),
other => Err(format!("unknown mm family: {other}")),
}
let spec: PipelineSpec = serde_json::from_str(json).map_err(|e| format!("mm spec: {e}"))?;
build_pipeline(spec)
}
rust/sglang-server/src/api_server.rs
+1 -95
View File
@@ -3,6 +3,7 @@
//! `/generate` submits a `Request` then awaits one `Done` (unary) or relays SSE
//! frames (`data: {json}` … `[DONE]`), byte-compatible with Python
//! `http_server.generate_request`; `/server_info` reuses it for one control result.
pub mod app;
mod common;
mod disaggregation;
mod frame;
@@ -12,98 +13,3 @@ mod native_api;
mod openai;
mod prefetch;
mod submit;
use std::sync::Arc;
use axum::Router;
use crate::runtime::ServerArgs;
use crate::tokenizer_manager::ActivityCounter;
use crate::tokenizer_manager::Senders;
use disaggregation::bootstrap as pd_bootstrap;
/// Shared handler state: submission handles, immutable server configuration,
/// and the API-owned chat formatter.
#[derive(Clone)]
struct AppState {
senders: Senders,
egress_buf: usize,
server_args: Arc<ServerArgs>,
chat_formatter: Option<openai::ChatFormatter>,
/// Egress heartbeat (bumped per drained ring frame).
egress_activity: ActivityCounter,
}
pub async fn serve(
listener: std::net::TcpListener,
senders: Senders,
egress_buf: usize,
server_args: Arc<ServerArgs>,
egress_activity: ActivityCounter,
// The SAME set ingress releases from — see `Ingress::on_abort`. Constructing a
// local one here would leave the api server admitting rids that nothing ever
// releases.
shutdown: flume::Receiver<()>,
) {
let chat_formatter = openai::load_chat_support(&server_args);
let state = AppState {
senders,
egress_buf,
server_args: server_args.clone(),
chat_formatter,
egress_activity,
};
// Each endpoint module registers its own routes and merges here.
let router = Router::new()
.merge(common::routes())
.merge(native_api::routes())
.merge(openai::routes());
// TODO(auth): no API-key boundary yet. Python gates every route (except
// /health*, /metrics*, OPTIONS) via `add_api_key_middleware`; until ported,
// a configured `api_key` does NOT protect these routes.
//
// No body limit, matching the Python server.
let mut app = router
.layer(axum::extract::DefaultBodyLimit::disable())
.with_state(state);
// Prefill-only KV bootstrap registry. Merged AFTER `with_state` — its
// router carries its own Arc<Registry> state, so it cannot merge into the
// Router<AppState> above — and before `log::apply`, so bootstrap traffic
// shows in the access log.
if server_args.enable_pd_bootstrap() {
let (routes, sweeper) = pd_bootstrap::router_and_sweeper();
tokio::spawn(sweeper); // cancelled with the runtime on shutdown
app = app.merge(routes);
tracing::info!("PD KV bootstrap registry mounted on the api listener");
}
// Apply logging and access log middleware.
let app = log::apply(app, &server_args);
// The listener was already bound synchronously in `runtime::start` (so a port
// conflict fails startup); adopt it into the tokio reactor here.
let listener = match tokio::net::TcpListener::from_std(listener) {
Ok(l) => l,
Err(e) => {
tracing::error!(error = %e, "failed to adopt pre-bound listener");
return;
}
};
// `with_connect_info` exposes the peer address to the access-log middleware.
let serve = axum::serve(
listener,
app.into_make_service_with_connect_info::<std::net::SocketAddr>(),
);
tokio::select! {
r = serve => {
if let Err(e) = r {
tracing::error!(error = %e, "axum serve exited");
}
}
_ = shutdown.recv_async() => {
tracing::info!("shutdown: stopping accepts, aborting in-flight handlers");
}
}
}
rust/sglang-server/src/api_server/app.rs
+104
View File
@@ -0,0 +1,104 @@
//! Router assembly and the shared handler state: every endpoint module
//! registers its routes here, and [`serve`] runs the assembled app on the
//! pre-bound listener until shutdown.
use std::sync::Arc;
use axum::Router;
use super::disaggregation::bootstrap as pd_bootstrap;
use super::{common, log, native_api, openai};
use crate::message::config::ServerArgs;
use crate::tokenizer_manager::from_scheduler::ActivityCounter;
use crate::tokenizer_manager::wiring::Senders;
/// Shared handler state: submission handles, immutable server configuration,
/// and the API-owned chat formatter.
///
/// axum clones the router state into **every** request, so it is mounted as
/// `Arc<AppState>` — one refcount bump per request instead of cloning each
/// `flume::Sender` and the chat formatter. Deliberately not `Clone`, so it
/// can only be shared through that `Arc`.
pub(super) struct AppState {
pub(super) senders: Senders,
pub(super) response_buf: usize,
pub(super) server_args: Arc<ServerArgs>,
pub(super) chat_formatter: Option<openai::ChatFormatter>,
/// Response heartbeat (bumped per drained ring frame).
pub(super) response_activity: ActivityCounter,
}
pub async fn serve(
listener: std::net::TcpListener,
senders: Senders,
response_buf: usize,
server_args: Arc<ServerArgs>,
response_activity: ActivityCounter,
// The runtime's shutdown signal, shared with every worker stage: it fires
// (disconnects) when `Runtime::request_shutdown` drops the sender, at
// which point `serve` stops accepting and its in-flight handlers are
// aborted with the api runtime.
shutdown: flume::Receiver<()>,
) {
let chat_formatter = openai::load_chat_support(&server_args);
let state = Arc::new(AppState {
senders,
response_buf,
server_args: server_args.clone(),
chat_formatter,
response_activity,
});
// Each endpoint module registers its own routes and merges here.
let router = Router::new()
.merge(common::routes())
.merge(native_api::routes())
.merge(openai::routes());
// TODO(auth): no API-key boundary yet. Python gates every route (except
// /health*, /metrics*, OPTIONS) via `add_api_key_middleware`; until ported,
// a configured `api_key` does NOT protect these routes.
//
// No body limit, matching the Python server.
let mut app = router
.layer(axum::extract::DefaultBodyLimit::disable())
.with_state(state);
// Prefill-only KV bootstrap registry. Merged AFTER `with_state` — its
// router carries its own Arc<Registry> state, so it cannot merge into the
// Router<Arc<AppState>> above — and before `log::apply`, so bootstrap traffic
// shows in the access log.
if server_args.enable_pd_bootstrap() {
let (routes, sweeper) = pd_bootstrap::router_and_sweeper();
tokio::spawn(sweeper); // cancelled with the runtime on shutdown
app = app.merge(routes);
tracing::info!("PD KV bootstrap registry mounted on the api listener");
}
// Apply logging and access log middleware.
let app = log::apply(app, &server_args);
// The listener was already bound synchronously in `runtime::start` (so a port
// conflict fails startup); adopt it into the tokio reactor here.
let listener = match tokio::net::TcpListener::from_std(listener) {
Ok(l) => l,
Err(e) => {
tracing::error!(error = %e, "failed to adopt pre-bound listener");
return;
}
};
// `with_connect_info` exposes the peer address to the access-log middleware.
let serve = axum::serve(
listener,
app.into_make_service_with_connect_info::<std::net::SocketAddr>(),
);
tokio::select! {
r = serve => {
if let Err(e) = r {
tracing::error!(error = %e, "axum serve exited");
}
}
_ = shutdown.recv_async() => {
tracing::info!("shutdown: stopping accepts, aborting in-flight handlers");
}
}
}
rust/sglang-server/src/api_server/common.rs
+29 -29
View File
@@ -12,17 +12,21 @@ use axum::{
response::{IntoResponse, Response},
routing::get,
};
use std::sync::Arc;
use super::AppState;
use super::app::AppState;
use super::guard::AbortGuard;
use super::submit::submit;
use crate::message::{ControlRequest, EgressItem, GetInternalStateReq, RequestKind};
use crate::runtime::ServerArgs;
use crate::message::config::ServerArgs;
use crate::message::ids::Rid;
use crate::message::io_struct::{ControlRequest, GetInternalStateReq};
use crate::message::request::RequestKind;
use crate::message::response::ResponseItem;
/// The routes this module owns, mounted by `api_server::serve`.
pub(super) fn routes() -> Router<AppState> {
pub(super) fn routes() -> Router<Arc<AppState>> {
Router::new()
// Control-plane: reuses the ingress FSM (no tokenization), returns one
// Control-plane: reuses the request FSM (no tokenization), returns one
// non-streamed JSON result. Adding one = a route line + its struct tag.
.route("/server_info", get(server_info))
// Static config, no scheduler round-trip. `/get_model_info` (+ `/model_info`
@@ -31,7 +35,7 @@ pub(super) fn routes() -> Router<AppState> {
.route("/model_info", get(model_info))
}
/// Submit a control request through the ingress FSM (no tokenization) and await the
/// Submit a control request through the request FSM (no tokenization) and await the
/// scheduler's single msgpack result (a `structs.asdict` named map). Returns the
/// raw bytes, or an error `Response` to return as-is.
async fn await_control_result(
@@ -50,34 +54,27 @@ async fn await_control_result(
guard.disarm(&rid); // completed normally — nothing to abort
}
match received {
Some(EgressItem::Control(bytes)) => Ok(bytes),
Some(EgressItem::Error(e)) => {
Some(ResponseItem::Control(bytes)) => Ok(bytes),
Some(ResponseItem::Error(e)) => {
let code =
StatusCode::from_u16(e.http_status()).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
Err((code, e.to_string()).into_response())
}
// A control request never receives generation frames or service-call data.
Some(EgressItem::Frame(_)) | Some(EgressItem::Done(_)) | Some(EgressItem::Data(_)) => {
Err((
StatusCode::INTERNAL_SERVER_ERROR,
"unexpected generation output for control request",
)
.into_response())
}
Some(ResponseItem::Frame(_))
| Some(ResponseItem::Done(_))
| Some(ResponseItem::Data(_)) => Err((
StatusCode::INTERNAL_SERVER_ERROR,
"unexpected generation output for control request",
)
.into_response()),
None => Err((StatusCode::from_u16(499).unwrap(), "request aborted").into_response()),
}
}
/// `GET /get_model_info` (+ `/model_info` alias) — static model metadata from
/// `server_args` (no scheduler round-trip); `is_generation` always true.
///
/// Under `SGLANG_RUST_SERVER=1` this is the only `/model_info` a client can
/// reach — `launch_server` never mounts the Python app — so it answers the same
/// keys. It answers them from the launch blob, which is the whole of this
/// server's config knowledge: `server_args` is parsed once at boot and held
/// behind an `Arc`, and no route mounted here changes weights or parsers, so
/// the launch values are also the current ones.
async fn model_info(State(state): State<AppState>) -> Response {
async fn model_info(State(state): State<Arc<AppState>>) -> Response {
let sa = &state.server_args;
let body = serde_json::json!({
"model_path": sa.model_path,
@@ -111,12 +108,10 @@ async fn model_info(State(state): State<AppState>) -> Response {
/// `api_key`/`admin_api_key`; see [`shape_server_info`]).
///
/// TODO(server_info): Python also includes `kv_events`; add once plumbed.
async fn server_info(State(state): State<AppState>) -> Response {
async fn server_info(State(state): State<Arc<AppState>>) -> Response {
let bytes = match await_control_result(
&state,
ControlRequest::GetInternalStateReq(GetInternalStateReq::new(
crate::ids::Rid::new().to_string(),
)),
ControlRequest::GetInternalStateReq(GetInternalStateReq::new(Rid::new().to_string())),
)
.await
{
@@ -216,8 +211,13 @@ mod tests {
let mut msgpack = Vec::new();
rmpv::encode::write_value(&mut msgpack, &outer).unwrap();
let sa =
ServerArgs::from_json(r#"{"model_path": "/m", "api_key": "secret-token"}"#).unwrap();
// `api_key` is deliberately NOT a `ServerArgs` field — the typed schema
// cannot carry it — so the only place it could leak from is the raw
// scheduler dump shaped above.
let sa = ServerArgs {
model_path: "/m".into(),
..Default::default()
};
let out = shape_server_info(&msgpack, &sa).unwrap();
let text = String::from_utf8(out.clone()).unwrap();
// No secret leaks anywhere in the serialized response.
rust/sglang-server/src/api_server/disaggregation/bootstrap.rs
+21 -20
View File
@@ -15,6 +15,7 @@ use axum::routing::{post, put};
use axum::{Json, Router};
use serde::{Deserialize, Serialize};
use crate::utils::environ;
use crate::utils::response::json_error;
use crate::utils::serialize::{parse_int, parse_int_opt, parse_int_vec};
@@ -337,7 +338,7 @@ fn router(state: Arc<Registry>) -> Router {
/// Drop room entries
async fn cleanup_sweeper(state: Arc<Registry>) {
let cleanup_interval = Duration::from_secs(crate::environ::env_u64(
let cleanup_interval = Duration::from_secs(environ::env_u64(
ENTRY_CLEANUP_INTERVAL_ENV,
ENTRY_CLEANUP_INTERVAL_DEFAULT_SECS,
));
@@ -356,7 +357,10 @@ pub(crate) fn router_and_sweeper() -> (Router, impl std::future::Future<Output =
#[cfg(test)]
mod tests {
use super::*;
use crate::runtime::{Runtime, RuntimeConfig, RustServerServerArgs, ServerArgs};
use crate::message::config::{
DisaggregationMode, RuntimeConfig, RustServerServerArgs, ServerArgs,
};
use crate::runtime::Runtime;
use std::io::{Read, Write};
use std::net::SocketAddr;
@@ -413,35 +417,36 @@ mod tests {
const SENTINEL: &str =
"/route?prefill_dp_rank=-1&prefill_cp_rank=-1&target_tp_rank=-1&target_pp_rank=-1";
/// Minimal prefill boot blob (same shape as the `runtime` tests): no
/// tokenizer load, the two mandatory `model_config` fields, and the
/// prefill role that mounts the registry.
const TEST_SERVER_ARGS: &str = r#"{
"skip_tokenizer_init": true,
"disaggregation_mode": "prefill",
"model_config": {"context_len": 2048, "vocab_size": 1000}
}"#;
/// Minimal boot config (same shape as the `runtime` tests): no tokenizer
/// load, a complete default `model_config`, and the given PD role.
fn test_server_args(disaggregation_mode: DisaggregationMode) -> ServerArgs {
ServerArgs {
skip_tokenizer_init: true,
disaggregation_mode,
..Default::default()
}
}
/// Pick a free port (probe-bind pattern, as in the `runtime` tests) and
/// boot the full runtime there with the bootstrap registry mounted — the
/// registry serves on the api listener, so these tests also pin the merge
/// wiring (including the `enable_pd_bootstrap()` derivation from the
/// blob), not just the handlers.
/// role), not just the handlers.
fn start_on_free_port() -> (Runtime, SocketAddr) {
start_runtime(TEST_SERVER_ARGS)
start_runtime(test_server_args(DisaggregationMode::Prefill))
}
fn start_runtime(server_args_json: &str) -> (Runtime, SocketAddr) {
fn start_runtime(server_args: ServerArgs) -> (Runtime, SocketAddr) {
let probe = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
let addr = probe.local_addr().unwrap();
drop(probe);
let cfg = RuntimeConfig {
rust_server_args: RustServerServerArgs {
http_addr: addr,
api_worker_num: 1,
http_api_worker_num: 1,
..Default::default()
},
server_args: Arc::new(ServerArgs::from_json(server_args_json).unwrap()),
server_args: Arc::new(server_args),
};
(crate::runtime::start(cfg).expect("start runtime"), addr)
}
@@ -593,11 +598,7 @@ mod tests {
/// hiding a misdirected decode/router behind its retry loop.
#[test]
fn routes_absent_off_prefill() {
let non_prefill = r#"{
"skip_tokenizer_init": true,
"model_config": {"context_len": 2048, "vocab_size": 1000}
}"#;
let (_rt, addr) = start_runtime(non_prefill);
let (_rt, addr) = start_runtime(test_server_args(DisaggregationMode::Null));
let (status, _) = request(addr, "GET", SENTINEL, None);
assert_eq!(status, 404);
rust/sglang-server/src/api_server/frame.rs
+2 -2
View File
@@ -4,7 +4,7 @@
//! abort frames). No HTTP here — the sibling `native_api` module owns the handlers
//! and streams; it calls these per frame.
use crate::message::{ChunkEvent, ChunkExtras};
use crate::message::response::{ChunkEvent, ChunkExtras};
/// The text slot of a `[logprob, token_id, text]` tuple: the decoded token when
/// `return_text_in_logprobs` supplied a text buffer, else `null`.
@@ -152,7 +152,7 @@ fn hidden_states_rows(vals: &[f32], lens: &[u32]) -> serde_json::Value {
// `get`, not a clamped index: clamping only the END leaves `off` past
// `vals.len()` after one over-long row, making the next range reversed
// (`start > end`) — which panics on the api thread rather than yielding
// an empty row. Same reasoning as the egress decoder's `take_f32`.
// an empty row. Same reasoning as the decoder's `take_f32`.
rows.push(serde_json::json!(vals.get(off..off + l).unwrap_or(&[])));
off += l;
}
rust/sglang-server/src/api_server/guard.rs
+8 -8
View File
@@ -5,8 +5,8 @@
use std::collections::HashSet;
use crate::ids::Rid;
use crate::tokenizer_manager::{AbortSource, Senders};
use crate::message::ids::Rid;
use crate::tokenizer_manager::wiring::{AbortSource, Senders};
/// Aborts still-in-flight rids on drop. Each rid is disarmed on natural finish;
/// whatever remains at drop is aborted.
@@ -64,7 +64,7 @@ impl Drop for AbortGuard {
// The lane is unbounded, so this send only fails at shutdown, when the loop
// is gone and nothing is generating anyway.
for rid in self.rids.drain() {
let _ = self.senders.abort.send(AbortSource::Guard(rid));
let _ = self.senders.abort_tx.send(AbortSource::Guard(rid));
}
}
}
@@ -75,10 +75,10 @@ mod tests {
fn senders_with_abort(abort: flume::Sender<AbortSource>) -> Senders {
Senders {
tm: flume::unbounded().0,
abort,
tok: flume::unbounded().0,
detok: vec![],
tok_manager_tx: flume::unbounded().0,
abort_tx: abort,
tokenizer_tx: flume::unbounded().0,
detokenizer_tx: vec![],
}
}
@@ -105,7 +105,7 @@ mod tests {
/// An armed guard aborts its rid on drop — exactly the cleanup a busy-skipped
/// `/health_generate` probe relies on. It never sees a terminal frame here, so
/// dropping the guard is the only path that deregisters its detok sink (via the
/// ingress `on_abort`). Regression for the detok-entry leak per health probe.
/// request `on_abort`). Regression for the detok-entry leak per health probe.
#[test]
fn armed_guard_aborts_on_drop() {
let (tm_tx, tm_rx) = flume::unbounded();
rust/sglang-server/src/api_server/log.rs
+1 -1
View File
@@ -6,7 +6,7 @@
use axum::{Router, response::Response};
use crate::runtime::ServerArgs;
use crate::message::config::ServerArgs;
/// Install the access-log middleware when `server_args` enables it; identity
/// otherwise (the layer is never installed, so disabled stays zero-cost).
rust/sglang-server/src/api_server/native_api.rs
+56 -54
View File
@@ -1,5 +1,5 @@
//! The native SGLang data-plane endpoints: `/generate` (submit a request, then
//! either fold egress frames to one unary JSON response or relay them as SSE
//! either fold decode frames to one unary JSON response or relay them as SSE
//! `data: {json}` … `[DONE]`, byte-compatible with Python
//! `http_server.generate_request`) and `/health` + `/health_generate` (which
//! round-trip a 1-token generate probe). Frame shaping (`meta_info`, logprob
@@ -7,10 +7,9 @@
//! generate-request submission (`submit`); the shared `AppState` lives in the
//! parent `api_server` module.
use std::{
convert::Infallible,
time::{Duration, Instant},
};
use std::convert::Infallible;
use std::sync::Arc;
use std::time::{Duration, Instant};
use axum::{
Json, Router,
@@ -25,18 +24,20 @@ use axum::{
};
use tokio::sync::mpsc;
use super::AppState;
use super::app::AppState;
use super::frame::{
OutputAccumulator, cumulative_frame_string, frame_value, stream_frame_string, tag_value,
};
use super::guard::AbortGuard;
use super::submit::submit;
use crate::environ::env_bool;
use crate::ids::Rid;
use crate::message::{
ChunkEvent, EgressItem, GenerateBody, GenerateRequest, RequestKind, SamplingParams,
use crate::message::ids::Rid;
use crate::message::request::{GenerateBody, GenerateRequest, RequestKind};
use crate::message::response::{ChunkEvent, ResponseItem};
use crate::message::sampling::SamplingParams;
use crate::utils::{
environ,
response::{error_response, error_value},
};
use crate::utils::response::{error_response, error_value};
/// API-local timing for one request.
///
@@ -78,7 +79,7 @@ impl RequestTiming {
}
/// The routes this module owns, mounted by `api_server::serve`.
pub(super) fn routes() -> Router<AppState> {
pub(super) fn routes() -> Router<Arc<AppState>> {
Router::new()
.route("/generate", post(generate))
.merge(health_routes())
@@ -97,11 +98,11 @@ pub(super) fn native_error(code: StatusCode, message: &str, stream: bool) -> Res
/// always; `SGLANG_ENABLE_HEALTH_ENDPOINT_GENERATION` (default true, mirroring
/// Python) decides whether `/health` shares it or is a plain 200 (routing the
/// request already proves the frontend is up).
fn health_routes() -> Router<AppState> {
fn health_routes() -> Router<Arc<AppState>> {
let timeout =
std::time::Duration::from_secs(crate::environ::env_u64("SGLANG_HEALTH_CHECK_TIMEOUT", 20));
let probe = get(move |state: State<AppState>| health_generate(state, timeout));
let health = if env_bool("SGLANG_ENABLE_HEALTH_ENDPOINT_GENERATION", true) {
std::time::Duration::from_secs(environ::env_u64("SGLANG_HEALTH_CHECK_TIMEOUT", 20));
let probe = get(move |state: State<Arc<AppState>>| health_generate(state, timeout));
let health = if environ::env_bool("SGLANG_ENABLE_HEALTH_ENDPOINT_GENERATION", true) {
probe.clone()
} else {
get(|| async { StatusCode::OK.into_response() })
@@ -116,19 +117,21 @@ fn health_routes() -> Router<AppState> {
const FAKE_BOOTSTRAP_HOST: &str = "2.2.2.2";
/// `GET /health_generate` — deep health: confirm the scheduler → detok path is
/// producing output. 200 iff the egress heartbeat advances within `timeout`
/// producing output. 200 if the response heartbeat advances within `timeout`
/// (from `SGLANG_HEALTH_CHECK_TIMEOUT`, frozen at router build), else 503.
/// (`/health` uses the same handler when its env gate is on.)
///
/// Fires a pre-tokenized 1-token probe (`input_ids = [0]`, skips the tokenizer) so
/// an idle pipeline produces a frame, then watches the *global*
/// [`AppState::egress_activity`] counter (not the probe's own rid) — so a busy
/// [`AppState::response_activity`] counter (not the probe's own rid) — so a busy
/// server passes immediately and a backlog never false-503s (the analogue of
/// Python's `last_receive_tstamp`). The `HEALTH_CHECK` skip + `http_worker_ipc`
/// ack are irrelevant here: this single-process server owns the egress ring.
async fn health_generate(State(state): State<AppState>, timeout: std::time::Duration) -> Response {
/// Python's `last_receive_tstamp`).
async fn health_generate(
State(state): State<Arc<AppState>>,
timeout: std::time::Duration,
) -> Response {
let baseline = state
.egress_activity
.response_activity
.load(std::sync::atomic::Ordering::Relaxed);
// Fire the probe (the heartbeat is the signal, not its own response). A busy
@@ -167,7 +170,7 @@ async fn health_generate(State(state): State<AppState>, timeout: std::time::Dura
let deadline = tokio::time::Instant::now() + timeout;
loop {
if state
.egress_activity
.response_activity
.load(std::sync::atomic::Ordering::Relaxed)
!= baseline
{
@@ -189,7 +192,7 @@ async fn health_generate(State(state): State<AppState>, timeout: std::time::Dura
/// with **400** (Python's status for a bad request) carrying serde's field-level
/// message, instead of axum's default 422.
async fn generate(
State(state): State<AppState>,
State(state): State<Arc<AppState>>,
body: Result<Json<GenerateBody>, JsonRejection>,
) -> Response {
let body = match body {
@@ -280,18 +283,18 @@ async fn generate_single(
/// Fold a unary request to its terminal → (HTTP status, result/`error` JSON, saw-terminal);
/// `false` = truncation, caller keeps the abort guard armed. Shared by single + batch.
async fn drain_unary(
rx: &mut mpsc::Receiver<EgressItem>,
rx: &mut mpsc::Receiver<ResponseItem>,
rid_str: &str,
mut timing: RequestTiming,
) -> (StatusCode, serde_json::Value, bool) {
let mut acc = OutputAccumulator::default();
while let Some(item) = rx.recv().await {
match item {
EgressItem::Frame(out) => {
ResponseItem::Frame(out) => {
timing.observe_first_output();
acc.fold(&out);
}
EgressItem::Done(out) => {
ResponseItem::Done(out) => {
timing.observe_first_output();
timing.finish();
acc.fold(&out);
@@ -310,14 +313,14 @@ async fn drain_unary(
add_e2e_latency(&mut value, &timing);
return (StatusCode::OK, value, true);
}
EgressItem::Error(e) => {
ResponseItem::Error(e) => {
timing.finish();
let code = e.http_status();
let status =
StatusCode::from_u16(code).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
return (status, error_value(code, &e.to_string()), true);
}
EgressItem::Control(_) | EgressItem::Data(_) => continue, // never on `/generate`
ResponseItem::Control(_) | ResponseItem::Data(_) => continue, // never on `/generate`
}
}
// Sender dropped without a terminal item: the shard dropped this request (a
@@ -393,8 +396,8 @@ async fn generate_batch(
/// back for `FuturesUnordered` to re-poll. Empty result = channel closed.
async fn recv_indexed(
index: usize,
mut rx: mpsc::Receiver<EgressItem>,
) -> (usize, mpsc::Receiver<EgressItem>, Vec<EgressItem>) {
mut rx: mpsc::Receiver<ResponseItem>,
) -> (usize, mpsc::Receiver<ResponseItem>, Vec<ResponseItem>) {
let mut items = Vec::new();
match rx.recv().await {
Some(item) => items.push(item),
@@ -410,7 +413,7 @@ async fn recv_indexed(
/// `with_index` tags each frame (batch only), `incremental` = delta vs cumulative,
/// `guard` aborts unfinished on drop.
fn generation_event_stream(
receivers: Vec<(Rid, mpsc::Receiver<EgressItem>, RequestTiming)>,
receivers: Vec<(Rid, mpsc::Receiver<ResponseItem>, RequestTiming)>,
mut guard: AbortGuard,
incremental: bool,
with_index: bool,
@@ -456,7 +459,7 @@ fn generation_event_stream(
for item in items {
match item {
EgressItem::Frame(out) => {
ResponseItem::Frame(out) => {
timings[i].observe_first_output();
accs[i].fold(&out);
if incremental {
@@ -465,17 +468,17 @@ fn generation_event_stream(
coalesced = true;
}
}
EgressItem::Done(out) => {
ResponseItem::Done(out) => {
timings[i].observe_first_output();
timings[i].finish();
accs[i].fold(&out);
terminal = Some(out);
}
EgressItem::Error(e) => {
ResponseItem::Error(e) => {
timings[i].finish();
failed = Some(e);
}
EgressItem::Control(_) | EgressItem::Data(_) => {} // never on /generate
ResponseItem::Control(_) | ResponseItem::Data(_) => {} // never on /generate
}
}
@@ -539,29 +542,30 @@ fn terminal_stream_frame_string(
#[cfg(test)]
mod tests {
use super::*;
use crate::message::ChunkEvent;
use crate::tokenizer_manager::Senders;
use crate::message::response::ChunkEvent;
use crate::tokenizer_manager::wiring::Senders;
use crate::utils::error::Error;
use futures::StreamExt;
use std::time::Duration;
fn senders() -> Senders {
Senders {
tm: flume::unbounded().0,
abort: flume::unbounded().0,
tok: flume::unbounded().0,
detok: vec![],
tok_manager_tx: flume::unbounded().0,
abort_tx: flume::unbounded().0,
tokenizer_tx: flume::unbounded().0,
detokenizer_tx: vec![],
}
}
fn frame(rid: u64, text: &str) -> EgressItem {
EgressItem::Frame(ChunkEvent {
fn frame(rid: u64, text: &str) -> ResponseItem {
ResponseItem::Frame(ChunkEvent {
rid: Rid::from(rid.to_string()),
text: text.into(),
completion_tokens: 1,
..Default::default()
})
}
fn done(rid: u64, text: &str) -> EgressItem {
EgressItem::Done(ChunkEvent {
fn done(rid: u64, text: &str) -> ResponseItem {
ResponseItem::Done(ChunkEvent {
rid: Rid::from(rid.to_string()),
text: text.into(),
completion_tokens: 1,
@@ -579,8 +583,8 @@ mod tests {
fn timed_receiver(
rid: u64,
rx: mpsc::Receiver<EgressItem>,
) -> (Rid, mpsc::Receiver<EgressItem>, RequestTiming) {
rx: mpsc::Receiver<ResponseItem>,
) -> (Rid, mpsc::Receiver<ResponseItem>, RequestTiming) {
(
Rid::from(rid.to_string()),
rx,
@@ -630,7 +634,7 @@ mod tests {
#[tokio::test]
async fn unary_terminal_meta_info_matches_python_semantics() {
let (tx, mut rx) = mpsc::channel(2);
tx.send(EgressItem::Done(ChunkEvent {
tx.send(ResponseItem::Done(ChunkEvent {
rid: "internal-rid".into(),
text: "ok".into(),
token_ids: vec![7, 8],
@@ -723,11 +727,9 @@ mod tests {
generation_event_stream(receivers, AbortGuard::new_empty(senders()), false, true);
futures::pin_mut!(stream);
tx0.send(EgressItem::Error(crate::error::Error::Validation(
"bad".into(),
)))
.await
.unwrap();
tx0.send(ResponseItem::Error(Error::Validation("bad".into())))
.await
.unwrap();
let v = parse(&stream.next().await.unwrap());
assert_eq!(v["index"], 0);
assert_eq!(v["error"]["code"], 400);
rust/sglang-server/src/api_server/openai.rs
+19 -16
View File
@@ -6,6 +6,7 @@
use axum::{Router, http::StatusCode, response::Response};
use futures::StreamExt;
use std::sync::Arc;
use tokio::sync::mpsc;
mod chat;
@@ -17,19 +18,21 @@ mod tools;
pub(super) use template::ChatFormatter;
use super::AppState;
use super::app::AppState;
use super::frame::OutputAccumulator;
use super::guard::AbortGuard;
use super::submit::submit;
use crate::ids::Rid;
use crate::message::{ChunkEvent, EgressItem, GenerateRequest, RequestKind};
use crate::runtime::ServerArgs;
use crate::message::config::ServerArgs;
use crate::message::ids::Rid;
use crate::message::request::{GenerateRequest, RequestKind};
use crate::message::response::{ChunkEvent, ResponseItem};
use crate::tokenizer_manager::tokenizer;
use crate::utils::response::error_response;
const MAX_OPENAI_CHOICES: usize = 4096;
/// The routes this module owns, mounted by `api_server::serve`.
pub(super) fn routes() -> Router<AppState> {
pub(super) fn routes() -> Router<Arc<AppState>> {
Router::new()
.merge(models::routes())
.merge(completions::routes())
@@ -47,7 +50,7 @@ pub(super) fn load_chat_support(server_args: &ServerArgs) -> Option<ChatFormatte
if server_args.skip_tokenizer_init || server_args.tokenizer_path.is_empty() {
return None;
}
let config_file = crate::tokenizer::resolve_model_file(
let config_file = tokenizer::resolve_model_file(
&server_args.tokenizer_path,
server_args.revision.as_deref(),
"tokenizer_config.json",
@@ -114,25 +117,25 @@ pub(super) fn openai_error(code: StatusCode, message: impl Into<String>, stream:
/// `guard` on a natural terminal, and map errors / validation aborts /
/// truncation to `(status, message)` for the OpenAI error shape.
async fn collect_output(
mut rx: mpsc::Receiver<EgressItem>,
mut rx: mpsc::Receiver<ResponseItem>,
guard: &mut AbortGuard,
rid: &Rid,
) -> Result<ChunkEvent, (StatusCode, String)> {
let mut accumulator = OutputAccumulator::default();
let output = loop {
match rx.recv().await {
Some(EgressItem::Frame(output)) => accumulator.fold(&output),
Some(EgressItem::Done(output)) => {
Some(ResponseItem::Frame(output)) => accumulator.fold(&output),
Some(ResponseItem::Done(output)) => {
accumulator.fold(&output);
break accumulator.into_output();
}
Some(EgressItem::Error(error)) => {
Some(ResponseItem::Error(error)) => {
guard.disarm(rid);
let status = StatusCode::from_u16(error.http_status())
.unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
return Err((status, error.to_string()));
}
Some(EgressItem::Control(_)) | Some(EgressItem::Data(_)) => {}
Some(ResponseItem::Control(_)) | Some(ResponseItem::Data(_)) => {}
None => {
return Err((
StatusCode::INTERNAL_SERVER_ERROR,
@@ -160,7 +163,7 @@ async fn submit_generation(
request: GenerateRequest,
stream: bool,
guard: &mut AbortGuard,
) -> Result<mpsc::Receiver<EgressItem>, Response> {
) -> Result<mpsc::Receiver<ResponseItem>, Response> {
match submit(state, RequestKind::Generate(Box::new(request)), stream).await {
Ok((rid, rx)) => {
guard.arm(rid);
@@ -177,17 +180,17 @@ async fn submit_generation(
}
}
fn indexed_egress_stream(
fn indexed_decode_stream(
index: usize,
rx: mpsc::Receiver<EgressItem>,
) -> futures::stream::BoxStream<'static, (usize, Option<EgressItem>)> {
rx: mpsc::Receiver<ResponseItem>,
) -> futures::stream::BoxStream<'static, (usize, Option<ResponseItem>)> {
futures::stream::unfold((rx, false), move |(mut rx, finished)| async move {
if finished {
return None;
}
match rx.recv().await {
Some(item) => {
let finished = matches!(item, EgressItem::Done(_) | EgressItem::Error(_));
let finished = matches!(item, ResponseItem::Done(_) | ResponseItem::Error(_));
Some(((index, Some(item)), (rx, finished)))
}
None => Some(((index, None), (rx, true))),
rust/sglang-server/src/api_server/openai/chat.rs
+23 -21
View File
@@ -2,6 +2,7 @@
use std::collections::BTreeMap;
use std::convert::Infallible;
use std::sync::Arc;
use axum::{
Json, Router,
@@ -33,18 +34,22 @@ use super::tools::{
parse_chat_tool_calls,
};
use super::{
AppState, ChatFormatter, collect_output, contains_media, error_payload, indexed_egress_stream,
AppState, ChatFormatter, collect_output, contains_media, error_payload, indexed_decode_stream,
openai_error, submit_generation, unix_seconds_u32,
};
use crate::ids::Rid;
use crate::message::{ChunkExtras, EgressItem, GenerateRequest, OneOrMany, SamplingParams};
use crate::message::config::{DefaultSamplingParams, ServerArgs};
use crate::message::ids::Rid;
use crate::message::request::GenerateRequest;
use crate::message::response::{ChunkExtras, ResponseItem};
use crate::message::sampling::SamplingParams;
use crate::message::types::OneOrMany;
pub(super) fn routes() -> Router<AppState> {
pub(super) fn routes() -> Router<Arc<AppState>> {
Router::new().route("/v1/chat/completions", post(chat_completions))
}
async fn chat_completions(
State(state): State<AppState>,
State(state): State<Arc<AppState>>,
body: Result<Json<CreateChatCompletionRequest>, JsonRejection>,
) -> Response {
let request = match body {
@@ -283,7 +288,7 @@ pub(super) fn chat_sampling(
tool_choice: &DynamoToolChoice,
tools: &[ToolDefinition],
parallel_tool_calls: Option<bool>,
server_args: &crate::runtime::ServerArgs,
server_args: &ServerArgs,
) -> Result<SamplingParams, String> {
let mut sampling = chat_sampling_params(
request,
@@ -299,7 +304,7 @@ pub(super) fn chat_sampling(
sampling
.normalize(
server_args.skip_tokenizer_init,
server_args.model_config.vocab_size.unwrap_or(u64::MAX),
server_args.model_config.vocab_size,
)
.map_err(|error| error.to_string())?;
Ok(sampling)
@@ -352,10 +357,7 @@ impl SamplingDefaults {
};
/// The resolved model defaults (empty in `--sampling-defaults openai`
/// mode), which slot between the user's values and the OpenAI terminals.
pub(super) fn with_model_defaults(
mut self,
model: &crate::runtime::DefaultSamplingParams,
) -> SamplingDefaults {
pub(super) fn with_model_defaults(mut self, model: &DefaultSamplingParams) -> SamplingDefaults {
self.temperature = model.temperature;
self.top_p = model.top_p;
self
@@ -421,7 +423,7 @@ pub(super) fn chat_sampling_params(
#[allow(clippy::too_many_arguments)]
pub(super) async fn unary_chat(
submitted: Vec<(usize, Rid, mpsc::Receiver<EgressItem>)>,
submitted: Vec<(usize, Rid, mpsc::Receiver<ResponseItem>)>,
mut guard: AbortGuard,
response_id: String,
model: String,
@@ -505,7 +507,7 @@ pub(super) async fn unary_chat(
#[allow(clippy::too_many_arguments)]
pub(super) fn chat_event_stream(
submitted: Vec<(usize, Rid, mpsc::Receiver<EgressItem>)>,
submitted: Vec<(usize, Rid, mpsc::Receiver<ResponseItem>)>,
mut guard: AbortGuard,
response_id: String,
model: String,
@@ -541,7 +543,7 @@ pub(super) fn chat_event_stream(
for (index, rid, rx) in submitted {
rids.push(rid);
streams.push(indexed_egress_stream(index, rx));
streams.push(indexed_decode_stream(index, rx));
yield Annotated {
data: Some(CreateChatCompletionStreamResponse {
id: response_id.clone(),
@@ -578,12 +580,12 @@ pub(super) fn chat_event_stream(
continue;
};
let output = match item {
EgressItem::Frame(output) => output,
EgressItem::Done(output) => {
ResponseItem::Frame(output) => output,
ResponseItem::Done(output) => {
guard.disarm(&rids[index]);
output
}
EgressItem::Error(error) => {
ResponseItem::Error(error) => {
guard.disarm(&rids[index]);
yield Annotated {
data: None,
@@ -594,7 +596,7 @@ pub(super) fn chat_event_stream(
};
continue;
}
EgressItem::Control(_) | EgressItem::Data(_) => continue,
ResponseItem::Control(_) | ResponseItem::Data(_) => continue,
};
if let Some((code, message)) = output
.finish_reason
@@ -854,8 +856,8 @@ mod tests {
merge_template_stops, unary_chat,
};
use crate::api_server::guard::AbortGuard;
use crate::message::ChunkExtras;
use crate::runtime::DefaultSamplingParams;
use crate::message::config::DefaultSamplingParams;
use crate::message::response::ChunkExtras;
use axum::http::StatusCode;
use dynamo_protocols::types::{CreateChatCompletionRequest, Stop};
use futures::StreamExt;
@@ -923,7 +925,7 @@ mod tests {
));
assert_eq!(
formatter.stop_strs(),
Some(crate::message::OneOrMany::Many(vec![
Some(crate::message::types::OneOrMany::Many(vec![
"<|endoftext|>".into(),
"<|im_end|>".into()
]))
rust/sglang-server/src/api_server/openai/completions.rs
+22 -23
View File
@@ -2,6 +2,7 @@
use std::collections::BTreeMap;
use std::convert::Infallible;
use std::sync::Arc;
use axum::{
Json, Router,
@@ -23,16 +24,18 @@ use tokio::sync::mpsc;
use super::super::guard::AbortGuard;
use super::super::submit::submit;
use super::{
AppState, MAX_OPENAI_CHOICES, collect_output, error_payload, indexed_egress_stream,
AppState, MAX_OPENAI_CHOICES, collect_output, error_payload, indexed_decode_stream,
openai_error, submit_generation, unix_seconds_u32,
};
use crate::ids::Rid;
use crate::message::{
ChunkEvent, ChunkExtras, EgressItem, GenerateRequest, Matched, OneOrMany, RequestKind,
SamplingParams, TokenIds,
};
use crate::message::finish_reason::Matched;
use crate::message::ids::Rid;
use crate::message::request::{GenerateRequest, RequestKind};
use crate::message::response::{ChunkEvent, ChunkExtras, ResponseItem};
use crate::message::sampling::SamplingParams;
use crate::message::types::{OneOrMany, TokenIds};
use crate::utils::error::Error;
pub(super) fn routes() -> Router<AppState> {
pub(super) fn routes() -> Router<Arc<AppState>> {
Router::new().route("/v1/completions", post(completions))
}
@@ -47,7 +50,7 @@ pub(super) struct SubmittedChoice {
pub(super) prompt_index: usize,
pub(super) rid: Rid,
pub(super) echo: String,
pub(super) rx: mpsc::Receiver<EgressItem>,
pub(super) rx: mpsc::Receiver<ResponseItem>,
}
#[derive(Debug, Default)]
pub(super) struct ChoiceExtensions {
@@ -58,7 +61,7 @@ pub(super) struct ChoiceExtensions {
}
async fn completions(
State(state): State<AppState>,
State(state): State<Arc<AppState>>,
body: Result<Json<CreateCompletionRequest>, JsonRejection>,
) -> Response {
let request = match body {
@@ -123,11 +126,7 @@ async fn completions(
};
if let Err(error) = sampling.normalize(
state.server_args.skip_tokenizer_init,
state
.server_args
.model_config
.vocab_size
.unwrap_or(u64::MAX),
state.server_args.model_config.vocab_size,
) {
return openai_error(StatusCode::BAD_REQUEST, error.to_string(), false);
}
@@ -252,17 +251,17 @@ async fn decode_prompt_echo(state: &AppState, token_ids: TokenIds) -> Result<Str
));
};
match rx.recv().await {
Some(EgressItem::Data(payload)) => String::from_utf8(payload.to_vec()).map_err(|_| {
Some(ResponseItem::Data(payload)) => String::from_utf8(payload.to_vec()).map_err(|_| {
openai_error(
StatusCode::INTERNAL_SERVER_ERROR,
"detokenized prompt is not valid UTF-8",
false,
)
}),
Some(EgressItem::Error(crate::error::Error::Validation(message))) => {
Some(ResponseItem::Error(Error::Validation(message))) => {
Err(openai_error(StatusCode::BAD_REQUEST, &message, false))
}
Some(EgressItem::Error(error)) => {
Some(ResponseItem::Error(error)) => {
let status = StatusCode::from_u16(error.http_status())
.unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
Err(openai_error(
@@ -538,7 +537,7 @@ pub(super) fn completion_event_stream(
rids.push(choice.rid);
prompt_indexes.push(choice.prompt_index);
echoes.push(choice.echo);
streams.push(indexed_egress_stream(index, choice.rx));
streams.push(indexed_decode_stream(index, choice.rx));
}
let mut events = futures::stream::select_all(streams);
@@ -548,17 +547,17 @@ pub(super) fn completion_event_stream(
continue;
};
let output = match item {
EgressItem::Frame(output) => output,
EgressItem::Done(output) => {
ResponseItem::Frame(output) => output,
ResponseItem::Done(output) => {
guard.disarm(&rids[index]);
output
}
EgressItem::Error(error) => {
ResponseItem::Error(error) => {
guard.disarm(&rids[index]);
yield error_payload(StatusCode::from_u16(error.http_status()).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR), error.to_string()).to_string();
continue;
}
EgressItem::Control(_) | EgressItem::Data(_) => continue,
ResponseItem::Control(_) | ResponseItem::Data(_) => continue,
};
if let Some((code, message)) = output
@@ -739,7 +738,7 @@ mod tests {
completion_prompt_specs, completion_response_value, unary_completion,
};
use crate::api_server::guard::AbortGuard;
use crate::message::ChunkExtras;
use crate::message::response::ChunkExtras;
use axum::http::StatusCode;
use dynamo_protocols::types::{
Choice, CreateCompletionRequest, CreateCompletionResponse, Prompt,
rust/sglang-server/src/api_server/openai/models.rs
+4 -3
View File
@@ -7,10 +7,11 @@ use axum::{
response::{IntoResponse, Response},
routing::get,
};
use std::sync::Arc;
use super::{AppState, openai_error, unix_seconds_u32};
pub(super) fn routes() -> Router<AppState> {
pub(super) fn routes() -> Router<Arc<AppState>> {
Router::new()
.route("/v1/models", get(available_models))
.route("/v1/models/{model}", get(retrieve_model))
@@ -18,12 +19,12 @@ pub(super) fn routes() -> Router<AppState> {
/// `GET /v1/models` — OpenAI-compatible model list. Served from `server_args`;
/// no scheduler round-trip.
async fn available_models(State(state): State<AppState>) -> Response {
async fn available_models(State(state): State<Arc<AppState>>) -> Response {
let base = model_card(&state);
Json(serde_json::json!({ "object": "list", "data": [base] })).into_response()
}
async fn retrieve_model(State(state): State<AppState>, Path(model): Path<String>) -> Response {
async fn retrieve_model(State(state): State<Arc<AppState>>, Path(model): Path<String>) -> Response {
if model != state.server_args.served_model_name {
return openai_error(
StatusCode::NOT_FOUND,
rust/sglang-server/src/api_server/openai/template.rs
+1 -1
View File
@@ -18,7 +18,7 @@ use dynamo_renderer::{ChatTemplate, ContextMixins, PromptContextMixin, PromptFor
use serde_json::Value;
use thiserror::Error;
use crate::message::OneOrMany;
use crate::message::types::OneOrMany;
const SUPPORTED_STYLES: &[&str] = &[
"ADD_COLON_SINGLE",
rust/sglang-server/src/api_server/openai/test_utils.rs
+29 -37
View File
@@ -19,21 +19,21 @@ use serde_json::json;
use tower::util::ServiceExt;
use super::{openai_error, routes};
use crate::ids::Rid;
use crate::message::{ChunkEvent, EgressItem};
use crate::runtime::ServerArgs;
use crate::tokenizer_manager::Senders;
use crate::message::config::ServerArgs;
use crate::message::ids::Rid;
use crate::message::response::{ChunkEvent, ResponseItem};
use crate::tokenizer_manager::wiring::Senders;
pub(super) fn senders() -> Senders {
Senders {
tm: flume::unbounded().0,
abort: flume::unbounded().0,
tok: flume::unbounded().0,
detok: vec![],
tok_manager_tx: flume::unbounded().0,
abort_tx: flume::unbounded().0,
tokenizer_tx: flume::unbounded().0,
detokenizer_tx: vec![],
}
}
pub(super) fn chunk(rid: &str, text: &str, done: bool) -> EgressItem {
pub(super) fn chunk(rid: &str, text: &str, done: bool) -> ResponseItem {
let output = ChunkEvent {
rid: rid.into(),
text: text.into(),
@@ -50,20 +50,20 @@ pub(super) fn chunk(rid: &str, text: &str, done: bool) -> EgressItem {
..Default::default()
};
if done {
EgressItem::Done(output)
ResponseItem::Done(output)
} else {
EgressItem::Frame(output)
ResponseItem::Frame(output)
}
}
/// A submitted legacy completion choice with its egress channel.
/// A submitted legacy completion choice.
pub(super) fn submitted(
index: usize,
prompt_index: usize,
rid: &str,
) -> (
super::completions::SubmittedChoice,
tokio::sync::mpsc::Sender<EgressItem>,
tokio::sync::mpsc::Sender<ResponseItem>,
) {
let (tx, rx) = tokio::sync::mpsc::channel(8);
(
@@ -78,41 +78,33 @@ pub(super) fn submitted(
)
}
/// A submitted chat choice (the tuple `chat_event_stream` consumes) with its
/// egress channel.
/// A submitted chat choice (the tuple `chat_event_stream` consumes).
pub(super) fn chat_submitted(
index: usize,
rid: &str,
) -> (
(usize, Rid, tokio::sync::mpsc::Receiver<EgressItem>),
tokio::sync::mpsc::Sender<EgressItem>,
(usize, Rid, tokio::sync::mpsc::Receiver<ResponseItem>),
tokio::sync::mpsc::Sender<ResponseItem>,
) {
let (tx, rx) = tokio::sync::mpsc::channel(8);
((index, rid.into(), rx), tx)
}
// ---------------------------------------------------------------------
// Handler-level tests: full router, real extractors, no scheduler. A
// request that reaches `submit` with an OPEN tm lane would wait on the
// egress receiver forever, so submission-reaching cases use `senders_closed`
// (503) and everything else fails validation before submit.
// ---------------------------------------------------------------------
pub(super) fn server_args() -> Arc<ServerArgs> {
Arc::new(
serde_json::from_value(serde_json::json!({ "served_model_name": "model" }))
.expect("ServerArgs must deserialize"),
)
Arc::new(ServerArgs {
served_model_name: "model".into(),
..Default::default()
})
}
pub(super) fn app_state(senders: Senders) -> super::AppState {
super::AppState {
pub(super) fn app_state(senders: Senders) -> Arc<super::AppState> {
Arc::new(super::AppState {
senders,
egress_buf: 8,
response_buf: 8,
server_args: server_args(),
chat_formatter: None,
egress_activity: Default::default(),
}
response_activity: Default::default(),
})
}
pub(super) fn senders_closed() -> Senders {
@@ -126,10 +118,10 @@ pub(super) fn senders_closed() -> Senders {
let (tok_tx, tok_rx) = flume::unbounded();
drop(tok_rx);
Senders {
tm: tm_tx,
abort: abort_tx,
tok: tok_tx,
detok: vec![],
tok_manager_tx: tm_tx,
abort_tx,
tokenizer_tx: tok_tx,
detokenizer_tx: vec![],
}
}
rust/sglang-server/src/api_server/openai/tools.rs
+4 -2
View File
@@ -36,7 +36,8 @@ use dynamo_protocols::types::{
Role,
};
use crate::message::{ChunkEvent, SamplingParams};
use crate::message::response::ChunkEvent;
use crate::message::sampling::SamplingParams;
/// Canonicalize a tool-call parser name onto the dynamo-parsers registry keys.
///
@@ -263,7 +264,8 @@ mod tests {
apply_tool_constraint, chat_delta, chat_finish_reason, dynamo_parser_name,
dynamo_tool_choice, parse_chat_tool_calls,
};
use crate::message::{ChunkEvent, SamplingParams};
use crate::message::response::ChunkEvent;
use crate::message::sampling::SamplingParams;
use dynamo_parsers::tool_calling::jail::{Annotated, apply_tool_calling_jail};
use dynamo_parsers::{ToolChoice as DynamoToolChoice, ToolDefinition};
use dynamo_protocols::types::CreateChatCompletionStreamResponse as StreamResponse;
rust/sglang-server/src/api_server/prefetch.rs
+4 -4
View File
@@ -5,8 +5,8 @@
//! images must download concurrently, not in `n * REQUEST_TIMEOUT`. URLs and
//! file paths resolve here through `sglang-mm`'s `fetch_bytes_budgeted` (one
//! owner for proxy/timeout/cap semantics) and ride out-of-band as
//! [`crate::message::MmData::prefetched`], which
//! [`crate::message::mm_payload::to_mm_input`] swaps back in.
//! [`crate::message::request::MmData::prefetched`], which
//! [`crate::multi_modality::payload::to_mm_input`] swaps back in.
use std::sync::Arc;
@@ -15,8 +15,8 @@ use sglang_mm::common::fetch::{ByteBudget, fetch_bytes_budgeted};
use sglang_mm::driver::{MAX_ITEMS_PER_REQUEST, MAX_REQUEST_BYTES};
use tokio::sync::Semaphore;
use crate::message::mm_payload::{io_sources, item_count};
use crate::message::{GenerateRequest, MmData};
use crate::message::request::{GenerateRequest, MmData};
use crate::multi_modality::payload::{io_sources, item_count};
/// Global bound on concurrent media fetches across all in-flight requests;
/// excess acquisitions queue on the semaphore without holding a thread.
rust/sglang-server/src/api_server/submit.rs
+20 -13
View File
@@ -1,18 +1,20 @@
//! Request submission into the ingress pipeline, shared by every endpoint
//! Request submission into the to_scheduler pipeline, shared by every endpoint
//! module: mint the client-visible rid (uuid hex, Python-parity), build the
//! `Request`, and hand it to the TM with an egress receiver for the response.
//! `Request`, and hand it to the TM with a receiver for the response.
use axum::{http::StatusCode, response::Response};
use tokio::sync::mpsc;
use super::{AppState, native_api::native_error};
use crate::fsm::RequestState;
use crate::ids::Rid;
use crate::message::{EgressItem, EgressSink, Request, RequestKind};
use crate::tokenizer_manager::TmEvent;
use super::app::AppState;
use super::native_api::native_error;
use crate::message::ids::Rid;
use crate::message::request::{Request, RequestKind};
use crate::message::response::{ResponseItem, ResponseSink};
use crate::tokenizer_manager::wiring::TmEvent;
use crate::utils::fsm::RequestState;
/// Submit one request; returns the rid, its hashed routing key, and the egress
/// receiver. Every request arrives with its final rid — a generate request from
/// Submit one request; returns its rid and the response receiver. Every
/// request arrives with its final rid — a generate request from
/// `into_requests` (or the `HEALTH_CHECK_<uuid>` the health probe sets), a
/// control request from its constructor — so this only echoes it back.
pub(super) async fn submit(
@@ -21,7 +23,7 @@ pub(super) async fn submit(
// `stream`: the client is reading an SSE stream, so it expects 200 plus an
// error frame rather than a 4xx — `utils::response::error_response`'s rule.
stream: bool,
) -> Result<(Rid, mpsc::Receiver<EgressItem>), Response> {
) -> Result<(Rid, mpsc::Receiver<ResponseItem>), Response> {
let rid = match &kind {
// Generate rids are already final: `GenerateBody::into_requests` normalized the
// client's, or minted one. Control requests have no client-facing rid.
@@ -37,14 +39,19 @@ pub(super) async fn submit(
// sent for `meta_info.id`.
// Async-aware send so a full TM inbox yields (backpressure) instead of parking
// a thread; Err only when the inbox is closed (shutdown).
let (tx, rx) = mpsc::channel::<EgressItem>(state.egress_buf);
let (tx, rx) = mpsc::channel::<ResponseItem>(state.response_buf);
let request = Request {
rid: rid.clone(),
state: RequestState::Received,
sink: EgressSink::Local(tx),
sink: ResponseSink::Local(tx),
kind,
};
match state.senders.tm.send_async(TmEvent::Ingress(request)).await {
match state
.senders
.tok_manager_tx
.send_async(TmEvent::Intake(request))
.await
{
Ok(()) => Ok((rid, rx)),
// `SendError` has a single meaning — the channel is disconnected.
Err(_) => {
rust/sglang-server/src/lib.rs
+130 -156
View File
@@ -1,26 +1,19 @@
//! sglang-server: a multi-threaded Rust frontend (API server → TokenizerManager
//! sglang-server: a multi-threaded Rust frontend (HTTP server → TokenizerManager
//! → Tokenizer/Detokenizer) embedded in the Python scheduler process.
//!
//! Pipeline stages 1–5 are pure Rust and never touch a `PyObject`, so they run
//! concurrently with the Python scheduler without contending for the GIL. The
//! only GIL crossings are the boundary methods on [`Server`]:
//! * `recv_requests` — Python scheduler thread drains the ingress ring.
//! * `push_batch` — Python scheduler thread pushes one output batch.
//! * `push_result` — Python scheduler thread pushes one control result.
//!
//! All are non-blocking, so the GIL is never held across a wait.
//! This file is the Python↔Rust boundary: it registers the pyo3 module
//! (`_server`) and the classes exposed to the scheduler — the boot config
//! ([`ServerArgs`] and its parts, constructed by keyword from Python; their
//! `#[pyclass]`es and constructors live in `message::config`), [`Server`]
//! (boot, `recv_requests`/`wait_request`, `push_*`, MM handoff, shutdown),
//! [`RequestBatch`] and [`MmEncodeResult`]. Everything behind that boundary —
//! receiving requests, encoding multimodal inputs, tokenizing, detokenizing,
//! SSE streaming, and so on — is implemented purely in Rust and never touches
//! a `PyObject`.
mod api_server;
mod detokenizer;
mod environ;
mod error;
mod fsm;
mod ids;
mod message;
mod mm;
mod ring;
mod runtime;
mod tokenizer;
mod multi_modality;
mod tokenizer_manager;
mod utils;
@@ -30,27 +23,49 @@ use pyo3::prelude::*;
use pyo3::pybacked::PyBackedBytes;
use pyo3::types::PyBytes;
use crate::runtime::{Runtime, RuntimeConfig};
use crate::message::config::{
DefaultSamplingParams, DisaggregationMode, MmFamily, MmResample, MmSpec, ModelConfig,
RuntimeConfig, RustServerServerArgs, ServerArgs,
};
use crate::utils::{logging, runtime};
/// One drained MM result (see [`Server::take_mm`]). Exactly one of
/// `features`/`shm_names` is `Some`: inline features for single-rank serving
/// (zero-copy into numpy), or one POSIX segment name per item when the scheduler
/// broadcasts across TP ranks and Python wraps each in a `ShmPointerMMData`.
/// A `ValueError` for a boot-time failure, as `"{context}: {err}"`.
fn value_error(context: &str, err: impl std::fmt::Display) -> PyErr {
pyo3::exceptions::PyValueError::new_err(format!("{context}: {err}"))
}
/// One drained MM result (see [`Server::take_mm`]), consumed by
/// `RustServer.build_native_mm` to build the scheduler's
/// `MultimodalProcessorOutput`.
#[pyclass(frozen, get_all)]
struct MmHandoff {
struct MmEncodeResult {
/// *Generic.* All items' `pixel_values` concatenated, flat `f32` of logical
/// shape `[sum(t*h*w), feature_dim]`; `Some` on the inline (single-rank) path.
features: Option<Py<numpy::PyArray1<f32>>>,
/// *Generic.* Per-item POSIX shm segment name holding that item's features
/// (`[t*h*w, feature_dim]` f32); `Some` on the TP-broadcast path.
shm_names: Option<Vec<String>>,
grids: Vec<(u32, u32, u32)>,
/// *Generic.* Per-item content hash of the raw source bytes (or the caller's
/// `mm_hashes` override), precomputed so the drain never re-hashes.
hashes: Vec<u64>,
/// *Generic.* Per-item inclusive `(start, end)` placeholder-token span in the
/// expanded `input_ids`.
offsets: Vec<(u32, u32)>,
/// *Qwen-VL specific.* Per-item `image_grid_thw` `(t, h, w)` in patch units;
/// `t*h*w` is also the item's row count in `features`.
grids: Vec<(u32, u32, u32)>,
/// *Qwen-VL specific.* M-RoPE position ids, flat `i64` of row-major shape
/// `[3, seq_len]` (temporal, height, width rows).
mrope: Py<numpy::PyArray1<i64>>,
/// *Qwen-VL specific.* M-RoPE delta, `max(mrope) + 1 - seq_len`, that decode
/// adds to the plain sequence position.
mrope_delta: i64,
}
/// Columnar ingress batch handed to Python by [`Server::recv_requests`].
/// Columnar request batch handed to Python by [`Server::recv_requests`].
/// `frozen`: immutable snapshot, so field access never contends on a borrow.
#[pyclass(frozen, get_all)]
struct IngressBatch {
struct RequestBatch {
/// One msgpack scalar header per request (`input_ids` omitted).
headers: Vec<Py<PyBytes>>,
/// The raw-data plane today just all requests' raw little-endian int64
@@ -64,121 +79,92 @@ struct IngressBatch {
/// [`Server::start`], then poll it from the scheduler event loop.
#[pyclass]
struct Server {
rt: Runtime,
rt: runtime::Runtime,
}
#[pymethods]
impl Server {
/// Boot the frontend (spawns all threads) and return immediately.
/// `server_args` is the scheduler's [`ServerArgs`]; the rest are
/// rust-server-only overrides.
#[new]
#[pyo3(signature = (
server_args,
http_addr = None,
ingress_ring_cap = 8192,
egress_ring_cap = 8192,
to_scheduler_cap = 8192,
from_scheduler_cap = 8192,
channel_cap = 8192,
cores = None,
server_args_json = "{}",
))]
// pyo3 `#[new]` constructor: the wide arg list is the Python-facing boot
// surface (all optional overrides), not a call-site ergonomics problem.
#[allow(clippy::too_many_arguments)]
fn start(
server_args: ServerArgs,
http_addr: Option<String>,
ingress_ring_cap: usize,
egress_ring_cap: usize,
to_scheduler_cap: usize,
from_scheduler_cap: usize,
channel_cap: usize,
cores: Option<Vec<usize>>,
server_args_json: &str,
) -> PyResult<Self> {
// Static server metadata (server_args + model_config) dumped by the
// scheduler; parse and validate mandatory fields now so a bad/missing
// field is a boot error, not a request-time 500.
let server_args: runtime::ServerArgs = runtime::ServerArgs::from_json(server_args_json)
.map_err(|e| {
PyErr::new::<pyo3::exceptions::PyValueError, _>(format!(
"bad server_args_json: {e}"
))
})?;
server_args.validate_mandatory().map_err(|e| {
PyErr::new::<pyo3::exceptions::PyValueError, _>(format!("server_args: {e}"))
})?;
// The HTTP listen address, tokenizer source/threads/shards all live in the
// `server_args` blob; resolve them from there so the scheduler doesn't
// re-pass them. The explicit params stay as optional overrides for
// standalone callers (tests) that construct a `Server` without a full
// `server_args`.
// `server_args` already arrived typed (pyo3 rejected any missing/extra/
// mistyped field when Python constructed it); only value checks remain.
server_args
.validate()
.map_err(|e| value_error("server_args", e))?;
// The HTTP listen address, tokenizer source/threads/shards all live in
// `server_args`; resolve them from there so the scheduler doesn't re-pass
// them. The explicit params stay as optional overrides (per-DP-rank port,
// pinning) and for standalone callers.
let http_addr: SocketAddr = http_addr
.unwrap_or_else(|| server_args.bind())
.parse()
.map_err(|e| {
PyErr::new::<pyo3::exceptions::PyValueError, _>(format!("bad http_addr: {e}"))
})?;
.map_err(|e| value_error("bad http_addr", e))?;
let cfg = RuntimeConfig {
rust_server_args: runtime::RustServerServerArgs {
rust_server_args: RustServerServerArgs {
http_addr,
api_worker_num: server_args.api_worker_num(),
ingress_ring_cap,
egress_ring_cap,
http_api_worker_num: server_args.http_api_worker_num(),
to_scheduler_cap,
from_scheduler_cap,
channel_cap,
cores,
},
server_args: std::sync::Arc::new(server_args),
};
let rt = runtime::start(cfg).map_err(|e| {
PyErr::new::<pyo3::exceptions::PyValueError, _>(format!("runtime start failed: {e}"))
})?;
let rt = runtime::start(cfg).map_err(|e| value_error("runtime start failed", e))?;
Ok(Server { rt })
}
/// Non-blocking drain of the ingress ring, returned **columnar** as an
/// [`IngressBatch`] so the large `input_ids` tensor never goes through
/// msgpack (see the field docs for the layout). The `ids` cells are copied
/// **directly into the result `bytes`** (one copy, no intermediate buffer).
///
/// Runs entirely GIL-held, deliberately. `drain` is a `try_recv` loop plus an
/// uncontended stash lock (the Python thread is the only consumer), so it
/// cannot block — there is nothing for a detach to overlap with. And detaching
/// is far from free: reacquiring the GIL waits out the interpreter's switch
/// interval, so a `py.detach` here cost up to 5 ms whenever another Python
/// thread was runnable, to cover ~0.2 µs of work. Held, the whole call is a
/// fraction of a microsecond on an empty ring.
/// Non-blocking drain of the to_scheduler channel, returned **columnar** as an
/// [`RequestBatch`] so the large `input_ids` tensor never goes through
/// msgpack (see the field docs for the layout).
#[pyo3(signature = (max = 256))]
fn recv_requests(&self, py: Python<'_>, max: usize) -> PyResult<IngressBatch> {
let cols = self.rt.ingress.drain(max);
fn recv_requests(&self, py: Python<'_>, max: usize) -> PyResult<RequestBatch> {
let cols = self.rt.to_scheduler_rx.drain(max);
let headers = cols
.headers
.iter()
.map(|h| PyBytes::new(py, h).unbind())
.collect();
// Single pass: copy each raw ids cell straight into the output `bytes`.
let data = PyBytes::new_with(py, cols.ids_total, |buf| {
let mut pos = 0;
for cell in &cols.ids {
let end = pos + cell.len();
buf[pos..end].copy_from_slice(cell);
pos = end;
}
cols.copy_ids_into(buf);
Ok(())
})?
.unbind();
Ok(IngressBatch {
})?;
Ok(RequestBatch {
headers,
data,
data: data.unbind(),
lengths: cols.lengths,
})
}
/// Park up to `timeout_ms` for an incoming request so the idle scheduler loop
/// sleeps instead of spinning at 100% CPU. Returns `True` when a request is
/// ready (the next `recv_requests` includes it). The GIL is released while
/// parked, and `flume` wakes the moment a request is pushed, so this adds no
/// latency to real requests — only the idle wait is bounded by `timeout_ms`.
/// sleeps instead of spinning at 100% CPU.
#[pyo3(signature = (timeout_ms = 1000))]
fn wait_ingress(&self, py: Python<'_>, timeout_ms: u64) -> bool {
fn wait_request(&self, py: Python<'_>, timeout_ms: u64) -> bool {
py.detach(|| {
self.rt
.ingress
.to_scheduler_rx
.wait(std::time::Duration::from_millis(timeout_ms))
})
}
@@ -186,73 +172,74 @@ impl Server {
/// Push a whole decode batch as ONE frame: a columnar msgpack `header` plus
/// the raw `data_cols` (per-column `bytes`), concatenated here. Blocks for
/// backpressure; `False` only on shutdown.
///
/// Framed and pushed with the GIL HELD, detaching only if the ring is full.
/// This runs on the scheduler's CUDA-launch thread every decode step, where the
/// unconditional detach was the single worst boundary cost: framing is
/// ~0.1–0.2 µs, but reacquiring the GIL waits out the interpreter's switch
/// interval (5 ms by default) whenever another Python thread is runnable —
/// 17–50% of a 10–30 ms decode step, landing nondeterministically. Held, the
/// whole boundary is ~1.3 µs per step.
///
/// The slow path keeps its detach because a full ring genuinely parks: the
/// scheduler must feel backpressure rather than drop output it has already
/// committed to. It essentially never fires — measured headroom is ~100×.
fn push_batch(&self, py: Python<'_>, header: &[u8], data_cols: Vec<PyBackedBytes>) -> bool {
fn push_decode_result_batch(
&self,
py: Python<'_>,
header: &[u8],
data_cols: Vec<PyBackedBytes>,
) -> bool {
let cols: Vec<&[u8]> = data_cols.iter().map(|d| d.as_ref()).collect();
self.push_frame(py, crate::message::frame_egress_batch_cols(header, &cols))
self.push_frame(
py,
crate::message::response::frame_decode_batch_cols(header, &cols),
)
}
/// Push a control-request result. Blocks for backpressure; `False` only on
/// shutdown.
fn push_result(&self, py: Python<'_>, rid: &str, payload: &[u8]) -> bool {
self.push_frame(py, crate::message::frame_egress_result(rid, payload))
fn push_control_result(&self, py: Python<'_>, rid: &str, payload: &[u8]) -> bool {
self.push_frame(
py,
crate::message::response::frame_control_result(rid, payload),
)
}
/// Route a terminal failure back to request `rid`. Blocks for backpressure;
/// `False` only on shutdown.
fn push_error(&self, py: Python<'_>, rid: &str, message: &str) -> bool {
self.push_frame(py, crate::message::frame_egress_error(rid, message))
self.push_frame(py, crate::message::response::frame_error(rid, message))
}
/// Spawn the MM worker pool for the pipeline in `spec_json` (built from the
/// resolved processor config; see `NativeMmHost.resolve_native_spec`).
/// Image-only requests are processed entirely in Rust and parked for
/// [`Server::take_mm`]; anything the pipeline cannot serve is rejected back to
/// the client — there is no Python fallback.
fn start_mm_workers(&self, spec_json: &str, workers: usize) -> PyResult<()> {
let ctx = mm::Context::new(
spec_json,
/// Spawn the MM worker pool for the pipeline in `spec` (built from the
/// resolved processor config; see `NativeMmHost.resolve_native_spec` and
/// `RustServer._build_mm_spec`). Image-only requests are processed entirely
/// in Rust and parked for [`Server::take_mm`]; anything the pipeline cannot
/// serve is rejected back to the client — there is no Python fallback.
fn start_mm_workers(&self, spec: MmSpec, workers: usize) -> PyResult<()> {
let ctx = multi_modality::worker::Context::new(
spec,
self.rt.tokenizer.clone(),
self.rt.mm_sidecar.clone(),
)
.map_err(PyErr::new::<pyo3::exceptions::PyValueError, _>)?;
.map_err(|e| value_error("mm spec", e))?;
self.rt.spawn_mm_pool(workers, std::sync::Arc::new(ctx));
Ok(())
}
/// Pop the MM result for `rid` — parked strictly before the request reached
/// the ingress ring — or `None` if there is none. The numeric buffers become
/// 1-D numpy arrays that take **ownership** of the Rust vectors, no copy.
/// the to_scheduler channel — or `None` if there is none. The numeric
/// buffers become 1-D numpy arrays that take **ownership** of the Rust
/// vectors, no copy.
///
/// Runs on the scheduler loop (`RustServer.drain`, under the GIL) between
/// decode steps, so any per-byte work here — memcpy or hashing, tens of MB
/// per image-heavy request — would stall every running request's ITL. Hence
/// the worker-precomputed `hashes`.
fn take_mm(&self, py: Python<'_>, rid: &str) -> Option<MmHandoff> {
/// Runs on the scheduler loop between decode steps, so any per-byte work
/// here — memcpy or hashing, tens of MB per image-heavy request — would
/// stall every running request's ITL. Hence the worker-precomputed `hashes`.
fn take_mm(&self, py: Python<'_>, rid: &str) -> Option<MmEncodeResult> {
use numpy::IntoPyArray;
let res = self.rt.mm_sidecar.take(rid)?;
let (features, shm_names) = match res.features {
mm::FeatureStore::Inline(v) => (Some(v.into_pyarray(py).unbind()), None),
multi_modality::sidecar::FeatureStore::Inline(v) => {
(Some(v.into_pyarray(py).unbind()), None)
}
// The segments — and the duty to unlink — move to Python here;
// `materialize()` unlinks after the post-broadcast clone on each rank.
mm::FeatureStore::Shm(segments) => (
multi_modality::sidecar::FeatureStore::Shm(segments) => (
None,
Some(segments.into_iter().map(|s| s.into_name()).collect()),
),
};
Some(MmHandoff {
Some(MmEncodeResult {
features,
shm_names,
grids: res.grids.iter().map(|g| (g[0], g[1], g[2])).collect(),
@@ -270,45 +257,32 @@ impl Server {
}
impl Server {
/// Hand one already-framed egress message to the ring: GIL-held when it fits,
/// detaching only to park on a full ring. Shared by every push path — they
/// differ solely in how the frame is built. `false` only on shutdown.
/// Hand one already-framed message to the ring. Shared by every push path —
/// they differ solely in how the frame is built. `false` only on shutdown.
#[inline]
fn push_frame(&self, py: Python<'_>, frame: bytes::Bytes) -> bool {
match self.rt.egress.try_push(frame) {
match self.rt.from_scheduler_tx.try_push(frame) {
Ok(()) => true,
// Consumer gone (shutdown): the frame is unavoidably lost.
Err(None) => false,
// Full: the scheduler must block here so backpressure reaches it, and
// blocking is exactly when releasing the GIL pays for itself.
Err(Some(frame)) => py.detach(|| self.rt.egress.push(frame)),
// Full: the scheduler must block here so backpressure reaches it.
Err(Some(frame)) => py.detach(|| self.rt.from_scheduler_tx.push(frame)),
}
}
}
/// Keeps the non-blocking log writer's background thread alive for the process
/// lifetime (dropping the guard would stop log delivery).
static LOG_GUARD: std::sync::OnceLock<tracing_appender::non_blocking::WorkerGuard> =
std::sync::OnceLock::new();
#[pymodule]
fn _server(m: &Bound<'_, PyModule>) -> PyResult<()> {
// Initialize tracing once; ignore if already set by the host process.
// Non-blocking writer: emitting threads (axum workers, egress, detok) only
// enqueue; a dedicated thread does the stdout formatting-flush + syscall.
// The queue is bounded and lossy — under extreme pressure log lines are
// dropped instead of stalling request threads.
let (writer, guard) = tracing_appender::non_blocking(std::io::stdout());
let _ = LOG_GUARD.set(guard);
let _ = tracing_subscriber::fmt()
.with_env_filter(
tracing_subscriber::EnvFilter::try_from_default_env()
.unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
)
.with_writer(writer)
.try_init();
logging::init_tracing();
m.add_class::<DisaggregationMode>()?;
m.add_class::<DefaultSamplingParams>()?;
m.add_class::<ModelConfig>()?;
m.add_class::<ServerArgs>()?;
m.add_class::<MmFamily>()?;
m.add_class::<MmResample>()?;
m.add_class::<MmSpec>()?;
m.add_class::<Server>()?;
m.add_class::<IngressBatch>()?;
m.add_class::<MmHandoff>()?;
m.add_class::<RequestBatch>()?;
m.add_class::<MmEncodeResult>()?;
Ok(())
}
rust/sglang-server/src/message.rs
+11 -94
View File
@@ -1,96 +1,13 @@
//! Messages moved between stages via `flume` (zero-copy moves); variable-length
//! buffers are `bytes::Bytes`, so egress fan-out to detok shards is a refcount bump.
//! Grouped by flow direction: [`request`] (the `/generate` body fan-out, the
//! in-flight request bodies + scheduler ingress wire), [`egress`]
//! (the response back-channel + egress-ring frames and decoded chunk events),
//! [`finish_reason`] (the terminal reason a request ended, Python's
//! `FinishReasonDict`), [`sampling`] (sampling-params normalization, the Python
//! `SamplingParams` port), [`io_struct`] (the scheduler wire structs), [`types`]
//! (the shared wire-shape adapters both directions use).
//! buffers are `bytes::Bytes`, so fanning one out to several detok shards is a
//! refcount bump, not a copy.
mod egress;
mod finish_reason;
mod io_struct;
pub mod mm_payload;
mod request;
mod sampling;
mod types;
pub use egress::{
ChunkEvent, ChunkExtras, EGRESS_TAG_BATCH, EGRESS_TAG_ERROR, EGRESS_TAG_RESULT, EgressItem,
EgressSink, SinkError, for_each_chunk, frame_egress_batch_cols, frame_egress_error,
frame_egress_result,
};
pub use finish_reason::Matched;
pub(crate) use io_struct::{AbortReq, ControlRequest, GetInternalStateReq};
pub use request::{GenerateBody, GenerateRequest, MmRequest, MmWorkItem, RequestKind};
// Constructed directly only by tests: `api_server::prefetch` fills its
// `prefetched` field, everything else gets it packed inside a `GenerateRequest`.
pub use request::MmData;
pub(crate) use sampling::{SamplingParams, SamplingParamsInput};
pub(crate) use types::{OneOrMany, OneOrManyItem, TokenIds};
use bytes::Bytes;
use crate::fsm::RequestState;
use crate::ids::Rid;
/// The owned request as it travels ingress stages (single owner, so `state` is
/// mutated lock-free). Common fields here; variant data in [`RequestKind`].
#[derive(Debug)]
pub struct Request {
/// Client-visible request id (uuid hex) — what the scheduler wire and
/// `meta_info.id` carry.
pub rid: Rid,
pub state: RequestState,
/// Back-channel to the client connection for egress frames.
pub sink: EgressSink,
/// Discriminant + variant body (generate vs control).
pub kind: RequestKind,
}
/// One ingress-ring entry, split columnar: the scalar `header` (msgpack, `input_ids`
/// omitted) + the raw int64 `ids` cell, so the big tensor never goes through msgpack.
#[derive(Debug)]
pub struct IngressMsg {
pub header: Bytes,
pub ids: Bytes,
}
/// Messages to a Detokenizer shard. `Register` carries the per-request sink for
/// the shard's local `rid -> sink` map. The rid STRING is the identity: `Rid::hash`
/// picks the shard (collisions there merely co-locate, which is harmless), but two
/// distinct rids that hash alike must not be the same map entry — that evicted one
/// client's sink and delivered their tokens to the other's connection. Equal rids
/// cannot reach here from different requests: `Rid::from_client` uniquifies every
/// client-supplied one.
pub enum DetokMsg {
Register {
/// Client-visible rid string — kept in `DetokState` so the shard can
/// emit `TmEvent::Abort(rid)` (the wire needs the string, not the hash).
rid: Rid,
sink: EgressSink,
/// Decode logprob token ids to text here (CPU-bound) not on the api threads.
decode_logprob_text: bool,
/// `SamplingParams.no_stop_trim`: keep the matched stop; default trims it.
no_stop_trim: bool,
},
/// One decode step's chunks for *this shard*. Batched because `tm-egress` blocks
/// per send, so one message per request cost ~1.3 µs × batch (5.1x at 4096).
Chunks(Vec<ChunkEvent>),
/// Decode a complete token-id sequence — the backend of
/// [`RequestKind::Detokenize`], the one request kind the detok stage itself
/// answers (it never reaches the scheduler ring). Sent by tm-ingress right
/// after the same rid's `Register` on the same channel (FIFO), so the shard
/// delivers the text through the registered sink like a control `Result`
/// and drops the entry.
Decode { rid: Rid, token_ids: Vec<u32> },
/// Control result: one already-serialized payload delivered to the sink verbatim.
Result { rid: Rid, payload: bytes::Bytes },
/// Terminal per-request failure → an `Error` to the sink (a 400, not a crash).
Fail { rid: Rid, message: String },
/// Drop the `rid -> sink` entry for a request rejected before the scheduler
/// (the rejecting stage already answered the client); else `Register` leaks one
/// entry.
Deregister { rid: Rid },
}
pub mod config;
pub mod detok;
pub mod finish_reason;
pub mod ids;
pub mod io_struct;
pub mod request;
pub mod response;
pub mod sampling;
pub mod types;
rust/sglang-server/src/message/config.rs
+638
View File
@@ -0,0 +1,638 @@
//! Runtime configuration: the rust-server boot knobs
//! ([`RustServerServerArgs`]), the scheduler's typed `server_args` handoff
//! ([`ServerArgs`] / [`ModelConfig`]), the [`RuntimeConfig`] pairing them for
//! `runtime::start`, and the native MM pipeline handoff ([`MmSpec`]).
//!
//! [`ServerArgs`] / [`ModelConfig`] / [`DefaultSamplingParams`] /
//! [`DisaggregationMode`] / [`MmSpec`] / [`MmFamily`] / [`MmResample`] are
//! also `#[pyclass]`es: the Python scheduler (`RustServer._build_server_args`
//! / `_build_mm_spec`) constructs them directly by keyword and hands them to
//! `Server`. There is one schema — this file — and
//! pyo3 enforces it at construction: every field is a required, typed
//! constructor argument, so a drifted caller fails at boot rather than running
//! on a silently-defaulted knob. The `#[pyo3::pymethods]` constructors below
//! each struct — plus the one hand-written extraction,
//! [`PreferredSamplingParams`] — are the only Python-facing code in this file;
//! the rest is pure Rust.
use std::net::SocketAddr;
use std::sync::Arc;
use serde::Serialize;
/// Boot knobs specific to the embedded rust server — none of these exist in
/// the Python-built [`ServerArgs`]; they arrive as explicit
/// `Server::start` parameters.
#[derive(Clone, Debug)]
pub struct RustServerServerArgs {
pub http_addr: SocketAddr,
pub http_api_worker_num: usize,
pub to_scheduler_cap: usize,
pub from_scheduler_cap: usize,
pub channel_cap: usize,
/// CPU core ids the pools pin to (e.g. this rank's NUMA-local cores minus
/// the scheduler's reserved launch cores). `None` → run unpinned.
pub cores: Option<Vec<usize>>,
}
impl Default for RustServerServerArgs {
fn default() -> Self {
Self {
http_addr: "127.0.0.1:30000".parse().unwrap(),
http_api_worker_num: 2,
to_scheduler_cap: 8192,
from_scheduler_cap: 8192,
channel_cap: 8192,
cores: None,
}
}
}
#[derive(Clone, Debug)]
pub struct RuntimeConfig {
/// Rust-server-only boot knobs (listen address, pool/ring sizes, pinning).
pub rust_server_args: RustServerServerArgs,
/// The scheduler's [`ServerArgs`] (worker counts, tokenizer source,
/// config-endpoint metadata). `Arc` so cloning the config (and, downstream,
/// each `AppState`) is cheap; immutable after construction.
pub server_args: Arc<ServerArgs>,
}
impl Default for RuntimeConfig {
fn default() -> Self {
Self {
rust_server_args: RustServerServerArgs::default(),
server_args: Arc::new(ServerArgs::default()),
}
}
}
/// The scheduler's launch-time handoff (`RustServer._build_server_args`):
/// the `server_args` fields the rust server reads, the resolved
/// [`ModelConfig`], and launch-time stamps. Values are post-`__post_init__`
/// (all paths and names resolved). Constructed from Python via the `#[new]` in
/// `lib.rs`, whose keyword parameters are exactly these fields.
#[pyo3::pyclass(frozen, from_py_object, module = "sglang.srt.rust_extensions._server")]
#[derive(Clone, Debug)]
pub struct ServerArgs {
/// HF repo id / local dir of the model, reported by `/get_model_info`.
pub model_path: String,
/// Model name reported by `/v1/models` and `/server_info`.
pub served_model_name: String,
/// Tokenizer source (model dir / `tokenizer.json` / HF repo id). Empty only
/// in standalone (test) configs — then boot requires `skip_tokenizer_init`.
pub tokenizer_path: String,
/// HF revision, used only when `tokenizer_path` is a repo id. `None` → main.
pub revision: Option<String>,
/// Weight format selected by `--load-format`, reported by `/get_model_info`.
/// The blob carries the post-`__post_init__` value (`auto` is already
/// narrowed to `gguf` / `mistral` / `runai_streamer` / `remote` where the
/// checkpoint demands it). Not consumed for loading -- the scheduler owns
/// that; `None` only when the blob omits the key.
pub load_format: Option<String>,
/// Operator-supplied weight version, reported by `/model_info`. Defaults to
/// `"default"` on the Python side, so it is present in every blob; `None`
/// only when the blob omits the key.
pub weight_version: Option<String>,
/// HTTP bind address (see [`Self::bind`]).
pub host: String,
pub port: u16,
/// Log levels driving the access log — uvicorn runs at
/// `log_level_http or log_level` (see [`Self::http_access_log_enabled`]).
pub log_level: String,
pub log_level_http: Option<String>,
/// Optional built-in chat-template name or path to a Jinja/legacy JSON
/// template file. Without an override, uses the tokenizer config template.
pub chat_template: Option<String>,
/// Parser selected by `--tool-call-parser`.
pub tool_call_parser: Option<String>,
/// Reasoning splitter selected by `--reasoning-parser` (e.g. deepseek-r1).
/// When set, chat completions strip the model's reasoning markers out of
/// `content` into `reasoning_content` — both unary and streaming.
pub reasoning_parser: Option<String>,
/// Python's global default for whether an SSE stream ends with a usage chunk.
pub stream_response_default_include_usage: bool,
/// Pinned tokenizer threads / detok shards (Python asserts both ≥ 1).
pub tokenizer_worker_num: usize,
pub detokenizer_worker_num: usize,
/// Token-ids-in / token-ids-out mode: no tokenizer load, raw `output_ids`
/// frames.
pub skip_tokenizer_init: bool,
/// Streamed `/generate` frames carry per-step deltas instead of cumulative
/// text. Matches the Python `TokenizerManager`.
pub incremental_streaming_output: bool,
/// PD-disaggregation role. (On prefill, the KV bootstrap registry is mounted
/// on the api router — see [`Self::enable_pd_bootstrap`].)
pub disaggregation_mode: DisaggregationMode,
/// The resolved Python `ModelConfig`, attached at handoff time.
pub model_config: ModelConfig,
/// Default sampling params advertised by `/get_model_info`, verbatim from
/// `server_args.preferred_sampling_params` (a JSON object or null).
pub preferred_sampling_params: Option<PreferredSamplingParams>,
/// Over-long inputs are truncated to fit the context instead of 400ing, and
/// `max_new_tokens` is clamped rather than rejected (Python
/// `TokenizerManager._validate_one_request`).
pub allow_auto_truncate: bool,
/// `return_hidden_states` is refused unless the server was launched with it:
/// the scheduler simply won't produce them, so the request would 200 with the
/// field silently missing.
pub enable_return_hidden_states: bool,
/// Output slots reserved per request on top of its input (eagle stores draft
/// tokens there). Not a `server_args` field — `TokenizerManager` derives it and
/// `RustServer._build_server_args` stamps it in, so both sides count alike.
pub num_reserved_tokens: u64,
/// Launch-time stamps (not `server_args` fields): sglang package version
/// and the scheduler-derived KV token capacity, reported by `/server_info`.
pub version: String,
pub max_total_num_tokens: u64,
}
#[pyo3::pymethods]
impl ServerArgs {
#[new]
#[pyo3(signature = (*,
model_path,
served_model_name,
tokenizer_path,
revision,
load_format,
weight_version,
host,
port,
log_level,
log_level_http,
chat_template,
tool_call_parser,
reasoning_parser,
stream_response_default_include_usage,
tokenizer_worker_num,
detokenizer_worker_num,
skip_tokenizer_init,
incremental_streaming_output,
disaggregation_mode,
model_config,
preferred_sampling_params,
allow_auto_truncate,
enable_return_hidden_states,
num_reserved_tokens,
version,
max_total_num_tokens,
))]
// The parameter list IS the schema; one keyword per field, all required.
#[allow(clippy::too_many_arguments)]
fn py_new(
model_path: String,
served_model_name: String,
tokenizer_path: String,
revision: Option<String>,
load_format: Option<String>,
weight_version: Option<String>,
host: String,
port: u16,
log_level: String,
log_level_http: Option<String>,
chat_template: Option<String>,
tool_call_parser: Option<String>,
reasoning_parser: Option<String>,
stream_response_default_include_usage: bool,
tokenizer_worker_num: usize,
detokenizer_worker_num: usize,
skip_tokenizer_init: bool,
incremental_streaming_output: bool,
disaggregation_mode: DisaggregationMode,
model_config: ModelConfig,
preferred_sampling_params: Option<PreferredSamplingParams>,
allow_auto_truncate: bool,
enable_return_hidden_states: bool,
num_reserved_tokens: u64,
version: String,
max_total_num_tokens: u64,
) -> Self {
Self {
model_path,
served_model_name,
tokenizer_path,
revision,
load_format,
weight_version,
host,
port,
log_level,
log_level_http,
chat_template,
tool_call_parser,
reasoning_parser,
stream_response_default_include_usage,
tokenizer_worker_num,
detokenizer_worker_num,
skip_tokenizer_init,
incremental_streaming_output,
disaggregation_mode,
model_config,
preferred_sampling_params,
allow_auto_truncate,
enable_return_hidden_states,
num_reserved_tokens,
version,
max_total_num_tokens,
}
}
}
impl Default for ServerArgs {
/// A standalone (test) config: no model, no tokenizer, unified role, but a
/// complete `model_config` so the runtime boots. Real launches never use
/// this — Python supplies every field.
fn default() -> Self {
Self {
model_path: String::new(),
served_model_name: String::new(),
tokenizer_path: String::new(),
revision: None,
load_format: None,
weight_version: None,
host: "127.0.0.1".into(),
port: 30000,
log_level: "info".into(),
log_level_http: None,
chat_template: None,
tool_call_parser: None,
reasoning_parser: None,
stream_response_default_include_usage: false,
tokenizer_worker_num: 1,
detokenizer_worker_num: 1,
skip_tokenizer_init: false,
incremental_streaming_output: false,
disaggregation_mode: DisaggregationMode::Null,
model_config: ModelConfig::default(),
preferred_sampling_params: None,
allow_auto_truncate: false,
enable_return_hidden_states: false,
num_reserved_tokens: 0,
version: String::new(),
max_total_num_tokens: 0,
}
}
}
/// `--preferred-sampling-params`, carried verbatim: `/get_model_info` echoes
/// whatever Python advertises, and the keys are whatever `SamplingParams`
/// accepts, so there is no fixed field list to model as a `#[pyclass]`.
#[derive(Clone, Debug, Serialize)]
#[serde(transparent)]
pub struct PreferredSamplingParams(pub serde_json::Value);
impl<'py> pyo3::FromPyObject<'_, 'py> for PreferredSamplingParams {
type Error = pyo3::PyErr;
fn extract(obj: pyo3::Borrowed<'_, 'py, pyo3::PyAny>) -> pyo3::PyResult<Self> {
let text = obj.extract::<String>()?;
serde_json::from_str(&text).map(Self).map_err(|e| {
pyo3::exceptions::PyValueError::new_err(format!(
"preferred_sampling_params is not valid JSON: {e}"
))
})
}
}
/// PD-disaggregation role, the values of `--disaggregation-mode`. Exposed to
/// Python as an enum (`DisaggregationMode.Null` / `.Prefill` / `.Decode`).
#[pyo3::pyclass(
eq,
frozen,
from_py_object,
module = "sglang.srt.rust_extensions._server"
)]
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum DisaggregationMode {
/// Unified prefill + decode.
Null,
Prefill,
Decode,
}
/// The slice of the resolved Python `ModelConfig` the rust server reads.
#[pyo3::pyclass(frozen, from_py_object, module = "sglang.srt.rust_extensions._server")]
#[derive(Clone, Debug)]
pub struct ModelConfig {
/// Resolved context length (`max_model_len` in `/v1/models`); the ceiling
/// for input + `max_new_tokens`.
pub context_len: u64,
/// Bounds client-supplied token ids — return 400s out-of-vocab ids before
/// they crash the scheduler's embedding lookup.
pub vocab_size: u64,
/// Whether the model accepts multimodal inputs. Gates the MM Encoding branch
/// in to-scheduler; `false` silently ignores mm fields, as the Python
/// `TokenizerManager` does with `mm_processor is None`.
pub is_multimodal: bool,
/// Resolved default sampling parameters, from Python's
/// `ModelConfig.get_default_sampling_params()`. Already gated on
/// `--sampling-defaults`: holds the model's generation_config.json values
/// in "model" mode, and is all-`None` in "openai" mode. Consumed when a chat
/// request omits `temperature`/`top_p` — the conversion must not skip
/// straight to the OpenAI terminal defaults.
pub default_sampling_params: DefaultSamplingParams,
}
#[pyo3::pymethods]
impl ModelConfig {
#[new]
#[pyo3(signature = (*, context_len, vocab_size, is_multimodal, default_sampling_params))]
fn py_new(
context_len: u64,
vocab_size: u64,
is_multimodal: bool,
default_sampling_params: DefaultSamplingParams,
) -> Self {
Self {
context_len,
vocab_size,
is_multimodal,
default_sampling_params,
}
}
}
impl Default for ModelConfig {
/// Test-only: a small but complete model so the runtime boots.
fn default() -> Self {
Self {
context_len: 2048,
vocab_size: 1000,
is_multimodal: false,
default_sampling_params: DefaultSamplingParams::default(),
}
}
}
/// One `SamplingParams` field per key `get_default_sampling_params()` may emit
/// (`repetition_penalty`, `temperature`, `top_k`, `top_p`, `min_p`); `None`
/// where the generation config does not set it.
///
/// `top_k` / `min_p` / `repetition_penalty` are carried for parity with the
/// Python dict but not yet consumed: the Dynamo chat request type only carries
/// `temperature` and `top_p`, so the conversion resolves just those two.
#[pyo3::pyclass(frozen, from_py_object, module = "sglang.srt.rust_extensions._server")]
#[derive(Clone, Debug, Default)]
#[allow(dead_code)]
pub struct DefaultSamplingParams {
pub temperature: Option<f64>,
pub top_p: Option<f64>,
pub top_k: Option<i64>,
pub min_p: Option<f64>,
pub repetition_penalty: Option<f64>,
}
#[pyo3::pymethods]
impl DefaultSamplingParams {
#[new]
#[pyo3(signature = (*, temperature = None, top_p = None, top_k = None, min_p = None, repetition_penalty = None))]
fn py_new(
temperature: Option<f64>,
top_p: Option<f64>,
top_k: Option<i64>,
min_p: Option<f64>,
repetition_penalty: Option<f64>,
) -> Self {
Self {
temperature,
top_p,
top_k,
min_p,
repetition_penalty,
}
}
}
/// The native MM pipeline handoff, built by `RustServer._build_mm_spec` from
/// the resolved `NativeMmSpec` and passed to `Server.start_mm_workers`. Same
/// contract as [`ServerArgs`]: every field is a required, typed constructor
/// keyword, so a drifted Python caller fails at boot.
#[pyo3::pyclass(frozen, from_py_object, module = "sglang.srt.rust_extensions._server")]
#[derive(Clone, Debug)]
pub struct MmSpec {
/// Park feature buffers in POSIX shm rather than inline. Set by the Python
/// launcher (`NativeMmHost._use_feature_shm`) exactly when the scheduler
/// broadcasts across TP ranks and will unwrap `ShmPointerMMData`.
pub feature_shm: bool,
/// The family pipeline and its resolved processor parameters.
pub pipeline: sglang_mm::registry::PipelineSpec,
}
#[pyo3::pymethods]
impl MmSpec {
/// The parameter list is flat because every family so far shares the
/// Qwen-VL processor geometry; a family with different knobs adds its own
/// keywords and match arm here.
#[new]
#[pyo3(signature = (*,
family,
feature_shm,
image_token_id,
patch_size,
merge_size,
temporal_patch_size,
min_pixels,
max_pixels,
image_mean,
image_std,
resample,
))]
#[allow(clippy::too_many_arguments)]
fn py_new(
family: MmFamily,
feature_shm: bool,
image_token_id: i32,
patch_size: usize,
merge_size: usize,
temporal_patch_size: usize,
min_pixels: usize,
max_pixels: usize,
image_mean: [f32; 3],
image_std: [f32; 3],
resample: MmResample,
) -> Self {
use sglang_mm::registry::PipelineSpec;
let pipeline = match family {
MmFamily::QwenVl => PipelineSpec::QwenVl(sglang_mm::qwen_vl::QwenVlSpec {
image_token_id,
patch_size,
merge_size,
temporal_patch_size,
min_pixels,
max_pixels,
image_mean,
image_std,
resample: resample.into(),
}),
};
Self {
feature_shm,
pipeline,
}
}
}
/// Which `sglang_mm` family pipeline serves the model — one variant per
/// [`sglang_mm::registry::PipelineSpec`] arm. Exposed to Python as an enum
/// (`MmFamily.QwenVl`); `NativeMmFamily.name` maps onto it at handoff.
#[pyo3::pyclass(
eq,
frozen,
from_py_object,
module = "sglang.srt.rust_extensions._server"
)]
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum MmFamily {
QwenVl,
}
/// The HF image processor the native resize must reproduce bit-exactly (see
/// [`sglang_mm::qwen_vl::Resampler`]). Exposed to Python as an enum
/// (`MmResample.AtenU8` / `.Pil`); `NativeMmFamily.image_processors` maps each
/// processor class onto it.
#[pyo3::pyclass(
eq,
frozen,
from_py_object,
module = "sglang.srt.rust_extensions._server"
)]
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum MmResample {
/// `Qwen2VLImageProcessor` / `…Fast` — torchvision on a uint8 tensor.
AtenU8,
/// `Qwen2VLImageProcessorPil`, behind `--disable-fast-image-processor`.
Pil,
}
impl From<MmResample> for sglang_mm::qwen_vl::Resampler {
fn from(r: MmResample) -> Self {
match r {
MmResample::AtenU8 => Self::AtenU8,
MmResample::Pil => Self::Pil,
}
}
}
fn join_host_port(host: &str, port: u16) -> String {
if host.contains(':') && !host.starts_with('[') {
format!("[{host}]:{port}") // bare IPv6 (`::`) needs brackets to bind
} else {
format!("{host}:{port}")
}
}
impl ServerArgs {
/// Fail fast at startup on values the types cannot express.
pub fn validate(&self) -> Result<(), String> {
if self.served_model_name.is_empty() {
return Err("empty 'served_model_name' in server_args".into());
}
Ok(())
}
/// True on a prefill or decode node — requests need bootstrap routing.
pub fn is_disaggregation(&self) -> bool {
self.disaggregation_mode != DisaggregationMode::Null
}
/// Serve the PD KV bootstrap registry on the api listener: every prefill
/// rust server hosts it, unconditionally — no extra topology gating. KV
/// managers and decode nodes reach the registry at the resolved
/// `disaggregation_bootstrap_port`, which rust-server mode aliases to the
/// api port, so whichever prefill server that port names is the one that
/// receives the registrations.
pub fn enable_pd_bootstrap(&self) -> bool {
self.disaggregation_mode == DisaggregationMode::Prefill
}
/// Whether the served model is multimodal, from the scheduler's config. See
/// [`ModelConfig::is_multimodal`].
pub fn model_is_multimodal(&self) -> bool {
self.model_config.is_multimodal
}
/// Bind address `host:port`. `host` is expected to be an IP — the result is
/// parsed as a `SocketAddr`, so a bare IPv6 host gets bracketed.
pub fn bind(&self) -> String {
join_host_port(&self.host, self.port)
}
/// Whether the HTTP access log is emitted, mirroring the Python server:
/// uvicorn runs at `log_level_http or log_level` and prints access lines
/// only at info/debug. `--log-level-http warning` turns them off.
pub fn http_access_log_enabled(&self) -> bool {
let level = self
.log_level_http
.as_deref()
.filter(|s| !s.is_empty())
.unwrap_or(&self.log_level);
matches!(
level.to_ascii_lowercase().as_str(),
"trace" | "debug" | "info"
)
}
/// Pinned API threads for the embedded HTTP api-server. Python `server_args`
/// has no such field — this is derived: enough to cover the widest pool.
pub fn http_api_worker_num(&self) -> usize {
4.max(self.tokenizer_worker_num)
.max(self.detokenizer_worker_num)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn bind_brackets_bare_ipv6() {
let sa = ServerArgs {
host: "::".into(),
port: 30001,
..Default::default()
};
assert_eq!(sa.bind(), "[::]:30001");
assert_eq!(ServerArgs::default().bind(), "127.0.0.1:30000");
}
#[test]
fn pd_role_derivations() {
let prefill = ServerArgs {
disaggregation_mode: DisaggregationMode::Prefill,
..Default::default()
};
assert!(prefill.is_disaggregation());
assert!(prefill.enable_pd_bootstrap());
let decode = ServerArgs {
disaggregation_mode: DisaggregationMode::Decode,
..Default::default()
};
assert!(decode.is_disaggregation());
assert!(!decode.enable_pd_bootstrap());
assert!(!ServerArgs::default().is_disaggregation());
}
#[test]
fn validate_requires_served_model_name() {
assert!(ServerArgs::default().validate().is_err());
let sa = ServerArgs {
served_model_name: "m".into(),
..Default::default()
};
assert!(sa.validate().is_ok());
}
/// `--log-level-http` overrides `--log-level` for the access log; unset or
/// empty falls through.
#[test]
fn access_log_follows_http_level_then_global() {
let mut sa = ServerArgs::default();
assert!(sa.http_access_log_enabled());
sa.log_level_http = Some("warning".into());
assert!(!sa.http_access_log_enabled());
sa.log_level_http = Some(String::new());
sa.log_level = "error".into();
assert!(!sa.http_access_log_enabled());
}
}
rust/sglang-server/src/message/detok.rs
+42
View File
@@ -0,0 +1,42 @@
//! Messages to a Detokenizer shard.
use super::ids::Rid;
use super::response::{ChunkEvent, ResponseSink};
/// Messages to a Detokenizer shard. `Register` carries the per-request sink for
/// the shard's local `rid -> sink` map. The rid STRING is the identity: `Rid::hash`
/// picks the shard (collisions there merely co-locate, which is harmless), but two
/// distinct rids that hash alike must not be the same map entry — that evicted one
/// client's sink and delivered their tokens to the other's connection. Equal rids
/// cannot reach here from different requests: `Rid::from_client` uniquifies every
/// client-supplied one.
pub enum DetokMsg {
Register {
/// Client-visible rid string — kept in `DetokState` so the shard can
/// emit `TmEvent::Abort(rid)` (the wire needs the string, not the hash).
rid: Rid,
sink: ResponseSink,
/// Decode logprob token ids to text here (CPU-bound) not on the api threads.
decode_logprob_text: bool,
/// `SamplingParams.no_stop_trim`: keep the matched stop; default trims it.
no_stop_trim: bool,
},
/// One decode step's chunks for *this shard*. Batched because `from-scheduler` blocks
/// per send.
Chunks(Vec<ChunkEvent>),
/// Decode a complete token-id sequence — the backend of
/// [`RequestKind::Detokenize`](super::RequestKind::Detokenize), the one
/// request kind the detok stage itself answers (it never reaches the
/// scheduler ring). Sent by to-scheduler right after the same rid's `Register`
/// on the same channel (FIFO), so the shard delivers the text through the
/// registered sink like a control `Result` and drops the entry.
Decode { rid: Rid, token_ids: Vec<u32> },
/// Control result: one already-serialized payload delivered to the sink verbatim.
Result { rid: Rid, payload: bytes::Bytes },
/// Terminal per-request failure → an `Error` to the sink (a 400, not a crash).
Fail { rid: Rid, message: String },
/// Drop the `rid -> sink` entry for a request rejected before the scheduler
/// (the rejecting stage already answered the client); else `Register` leaks one
/// entry.
Deregister { rid: Rid },
}
rust/sglang-server/src/message/finish_reason.rs
+2 -4
View File
@@ -1,8 +1,6 @@
//! The terminal finish reason: Python's `FinishReasonDict` — what
//! `BaseFinishReason.to_json()` (schedule_batch.py) puts on the egress wire, and
//! what the API echoes back as `meta_info.finish_reason`. Ingress has no
//! counterpart; it rides in the [`BatchHeader`](super::egress::BatchHeader) and on
//! each terminal [`ChunkEvent`](super::ChunkEvent).
//! `BaseFinishReason.to_json()` (schedule_batch.py) puts on the response, and
//! what the API echoes back as `meta_info.finish_reason`.
use serde::{Deserialize, Serialize};
rust/sglang-server/src/ids.rs → rust/sglang-server/src/message/ids.rs
+3 -11
View File
@@ -115,8 +115,7 @@ impl Rid {
}
}
/// Shard index for `n` detokenizer shards. Pure function of the id so the
/// ingress and egress sides agree without any shared map.
/// Shard index for `n` detokenizer shards.
#[inline]
pub fn shard(&self, n: usize) -> usize {
debug_assert!(n > 0);
@@ -126,16 +125,9 @@ impl Rid {
impl From<String> for Rid {
fn from(id: String) -> Self {
// ONE seed per process, not one per conversion. Ingress and egress each
// ONE seed per process, not one per conversion. To-scheduler and from-scheduler each
// build a `Rid` from the same string and must agree on the shard without a
// shared map — a fresh `RandomState` here would hash the same rid two
// different ways, so chunks would arrive at a shard that never registered
// the request and be dropped.
//
// The seed is random rather than fixed because rids are client-supplied:
// with public keys, colliding rids are an offline ~2^32 search. Collisions
// are only a shard co-location now (identity is the string), but a keyed
// hash also stops an attacker from stacking every request onto one shard.
// shared map.
static SEED: OnceLock<RandomState> = OnceLock::new();
let hash = SEED.get_or_init(RandomState::new).hash_one(&id);
Rid { id, hash }
rust/sglang-server/src/message/io_struct.rs
+4 -2
View File
@@ -6,9 +6,11 @@
use bytes::Bytes;
use serde::Serialize;
use super::request::GenerateRequest;
use super::sampling::SamplingParams;
use super::types::TokenIds;
use super::types::{Tagged, control_messages, wire_struct};
use super::{GenerateRequest, SamplingParams, TokenIds};
use crate::error::Error;
use crate::utils::error::Error;
wire_struct! {
/// The scheduler's `TokenizedGenerateReqInput`. Keep in lockstep with the
rust/sglang-server/src/message/request.rs
+44 -18
View File
@@ -1,7 +1,5 @@
//! The `/generate` request path: the HTTP body and its per-request fan-out
//! ([`GenerateBody`] → [`GenerateRequest`]s), the variant bodies, and the
//! scheduler ingress encodings (`TokenizedGenerateReqInput` header,
//! control/abort, `IngressMsg`).
//! ([`GenerateBody`] → [`GenerateRequest`]s).
use std::collections::HashSet;
use std::sync::LazyLock;
@@ -11,10 +9,12 @@ use itertools::izip;
use serde::Deserialize;
use super::io_struct::{ControlRequest, TokenizedGenerateReqInput};
use super::{OneOrMany, OneOrManyItem, SamplingParams, SamplingParamsInput, TokenIds};
use crate::environ::env_u64;
use crate::error::Error;
use crate::ids::Rid;
use super::response::ResponseSink;
use super::sampling::{SamplingParams, SamplingParamsInput};
use super::types::{OneOrMany, OneOrManyItem, TokenIds};
use crate::message::ids::Rid;
use crate::utils::fsm::RequestState;
use crate::utils::{environ::env_u64, error::Error};
/// Hard cap on how many scheduler requests one `/generate` HTTP call may expand
/// into. Every column below is allocated per item before anything is dispatched,
@@ -519,12 +519,12 @@ fn split_mm_column(
/// plus the owned inputs from [`GenerateRequest::take_mm_work`].
#[derive(Debug)]
pub struct MmRequest {
pub rid: crate::ids::Rid,
pub rid: Rid,
pub work: MmWorkItem,
}
/// The parked request's fields the MM worker owns; converted to the driver input
/// by [`super::mm_payload::to_mm_input`].
/// by [`crate::multi_modality::payload::to_mm_input`].
#[derive(Debug, Default)]
pub struct MmWorkItem {
pub text: Option<String>,
@@ -541,17 +541,40 @@ pub struct MmWorkItem {
/// Whether an optional mm field counts as multimodal input, via the same
/// `value_present` the MM worker's payload parser uses.
fn mm_value_present(v: &Option<rmpv::Value>) -> bool {
v.as_ref().is_some_and(super::mm_payload::value_present)
v.as_ref()
.is_some_and(crate::multi_modality::payload::value_present)
}
/// Request variant — selects the ingress branch, scheduler wire message, and
/// egress shape. Each owns its body, so generate/control fields stay type-separate.
/// The owned request as it travels request stages (single owner, so `state` is
/// mutated lock-free). Common fields here; variant data in [`RequestKind`].
#[derive(Debug)]
pub struct Request {
/// Client-visible request id (uuid hex) — what the scheduler wire and
/// `meta_info.id` carry.
pub rid: Rid,
pub state: RequestState,
/// Back-channel to the client connection for response frames.
pub sink: ResponseSink,
/// Discriminant + variant body (generate vs control).
pub kind: RequestKind,
}
/// One to_scheduler channel entry, split columnar: the scalar `header` (msgpack, `input_ids`
/// omitted) + the raw int64 `ids` cell, so the big tensor never goes through msgpack.
#[derive(Debug)]
pub struct SchedulerRequest {
pub header: Bytes,
pub ids: Bytes,
}
/// Request variant — selects the request branch, scheduler wire message, and
/// response shape. Each owns its body, so generate/control fields stay type-separate.
#[derive(Debug)]
pub enum RequestKind {
/// `/generate`: tokenize (if needed) then push a `TokenizedGenerateReqInput`.
Generate(Box<GenerateRequest>),
/// A control endpoint (e.g. `/server_info`, `/health`): no tokenization, and
/// the egress is a single non-streamed JSON result.
/// the response is a single non-streamed JSON result.
Control(Box<ControlRequest>),
/// Internal service call: decode a complete token-id sequence to text. Walks
/// the same FSM as every request (validate → register → Queued), but the
@@ -595,7 +618,7 @@ pub struct GenerateRequest {
/// by the pool before the header is built; never reaches the scheduler wire.
pub skip_special_tokens: bool,
/// Sampling params (defaults when the client sent none, as in Python);
/// normalized + verified at ingress, then serialized into the header.
/// normalized + verified, then serialized into the header.
pub sampling_params: SamplingParams,
/// Whether the client asked for SSE streaming.
pub stream: bool,
@@ -639,13 +662,16 @@ pub struct GenerateRequest {
}
/// The opaque multimodal fields of one request (see [`GenerateRequest::mm`]).
///
/// Constructed directly only by tests: `api_server::prefetch` fills its
/// `prefetched` field, everything else gets it packed inside a `GenerateRequest`.
#[derive(Debug, Default)]
pub struct MmData {
pub image_data: Option<rmpv::Value>,
pub video_data: Option<rmpv::Value>,
pub audio_data: Option<rmpv::Value>,
/// Bytes of `image_data`'s I/O-backed sources, resolved by
/// `api_server::prefetch` in `mm_payload::io_sources` order so MM workers
/// `api_server::prefetch` in `payload::io_sources` order so MM workers
/// never block on I/O. Out-of-band: the values above stay as the client
/// sent them.
pub prefetched: Vec<bytes::Bytes>,
@@ -693,8 +719,8 @@ impl GenerateRequest {
}
/// `input_ids` widened to raw little-endian int64 bytes (the scheduler's
/// `array("q")` columnar cell — rides the ingress ring outside msgpack). Empty
/// when not tokenized.
/// `array("q")` columnar cell — rides the to-scheduler channel outside
/// msgpack). Empty when not tokenized.
pub fn encode_data_buf(&self) -> Bytes {
let ids = self.input_ids.as_deref().unwrap_or(&[]);
let mut buf = Vec::with_capacity(ids.len() * 8);
@@ -882,7 +908,7 @@ mod tests {
assert!(requests(r#"{"text": "a", "input_ids": [1]}"#).is_err());
assert!(requests(r#"{"stream": true}"#).is_err());
// Parallel sampling is rejected where Python reads it — in the params,
// at normalization (the ingress step), not here.
// at normalization, not here.
let (mut ps, _) = requests(r#"{"text": "a", "sampling_params": {"n": 2}}"#).unwrap();
assert!(ps[0].sampling_params.normalize(false, TEST_VOCAB).is_err());
}
rust/sglang-server/src/message/egress.rs → rust/sglang-server/src/message/response.rs
+53 -59
View File
@@ -1,25 +1,25 @@
//! The egress (response) direction: the per-request back-channel the API
//! handler drains ([`EgressSink`] / [`EgressItem`]), the egress-ring frame
//! encodings (batch / control result / error), and the columnar batch decode
//! into per-request [`ChunkEvent`]s.
//! The response direction: the per-request back-channel the API handler
//! drains ([`ResponseSink`] / [`ResponseItem`]), the response frame encodings
//! (batch / control result / error), and the columnar batch decode into
//! per-request [`ChunkEvent`]s.
use bytes::Bytes;
use serde::{Deserialize, Serialize};
use tokio::sync::mpsc;
use super::TokenIds;
use super::finish_reason::FinishReason;
use crate::error::Error;
use crate::ids::Rid;
use super::types::TokenIds;
use crate::message::ids::Rid;
use crate::utils::error::Error;
/// Per-request back-channel the detok shard writes egress frames to and the API
/// Per-request back-channel the detok shard writes decode frames to and the API
/// handler drains for SSE; bounded, and receiver-drop (disconnect) = stream end.
#[derive(Clone, Debug)]
pub enum EgressSink {
Local(mpsc::Sender<EgressItem>),
pub enum ResponseSink {
Local(mpsc::Sender<ResponseItem>),
}
/// Why an [`EgressSink::try_send`] failed: `Full` = client backpressure, `Closed`
/// Why an [`ResponseSink::try_send`] failed: `Full` = client backpressure, `Closed`
/// = client gone. Both terminal for a stream; the caller distinguishes for logging.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SinkError {
@@ -27,11 +27,11 @@ pub enum SinkError {
Closed,
}
impl EgressSink {
impl ResponseSink {
/// Non-blocking send. `Err(Full)` = backpressure, `Err(Closed)` = client gone.
pub fn try_send(&self, item: EgressItem) -> Result<(), SinkError> {
pub fn try_send(&self, item: ResponseItem) -> Result<(), SinkError> {
match self {
EgressSink::Local(tx) => tx.try_send(item).map_err(|e| match e {
ResponseSink::Local(tx) => tx.try_send(item).map_err(|e| match e {
mpsc::error::TrySendError::Full(_) => SinkError::Full,
mpsc::error::TrySendError::Closed(_) => SinkError::Closed,
}),
@@ -40,12 +40,12 @@ impl EgressSink {
}
#[allow(dead_code)] // the receiver half is created inline in api_server::submit.
pub type EgressSource = mpsc::Receiver<EgressItem>;
pub type ResponseSource = mpsc::Receiver<ResponseItem>;
/// What the connection handler receives on the egress stream: a detok-decoded
/// What the connection handler receives on the decode stream: a detok-decoded
/// [`ChunkEvent`] (handler formats it), a verbatim control payload, or an error.
#[derive(Debug)]
pub enum EgressItem {
pub enum ResponseItem {
/// An intermediate streamed generation step (only sent for streaming reqs).
Frame(ChunkEvent),
/// The final generation step.
@@ -62,15 +62,15 @@ pub enum EgressItem {
Error(Error),
}
/// Egress-ring frame tag (first byte, prepended Rust-side; Python wire unchanged):
/// Response frame tag (first byte, prepended Rust-side; Python wire unchanged):
/// a single control-request result payload.
pub const EGRESS_TAG_RESULT: u8 = 1;
pub const DISPATCH_TAG_RESULT: u8 = 1;
/// A whole decode batch: msgpack columnar header + one concatenated raw buffer;
/// tm-egress decodes it into per-request [`ChunkEvent`]s (no per-request FFI).
pub const EGRESS_TAG_BATCH: u8 = 2;
/// from-scheduler decodes it into per-request [`ChunkEvent`]s (no per-request FFI).
pub const DISPATCH_TAG_BATCH: u8 = 2;
/// A per-request failure `[rid, message]`: the Python drain couldn't decode a
/// header, so it routes a 400 back to that request instead of crashing the loop.
pub const EGRESS_TAG_ERROR: u8 = 3;
pub const DISPATCH_TAG_ERROR: u8 = 3;
/// Read `n` little-endian f32s from `data` at `*off`, advancing `*off`. `None` when
/// the range runs past the buffer (a malformed / positional-ABI-drifted frame): the
@@ -105,11 +105,11 @@ fn take_i32(data: &[u8], off: &mut usize, n: usize) -> Option<Vec<i32>> {
/// Frame a decode batch: `[BATCH tag][u32 header len][header][data cols…]`. The
/// caller's `data_cols` are concatenated straight into the frame (one copy, no
/// `b"".join`); `header` is the msgpack [`BatchHeader`]. Runs off the GIL.
pub fn frame_egress_batch_cols(header: &[u8], data_cols: &[&[u8]]) -> Bytes {
/// `b"".join`); `header` is the msgpack [`BatchHeader`].
pub fn frame_decode_batch_cols(header: &[u8], data_cols: &[&[u8]]) -> Bytes {
let data_len: usize = data_cols.iter().map(|c| c.len()).sum();
let mut buf = Vec::with_capacity(1 + 4 + header.len() + data_len);
buf.push(EGRESS_TAG_BATCH);
buf.push(DISPATCH_TAG_BATCH);
buf.extend_from_slice(&(header.len() as u32).to_le_bytes());
buf.extend_from_slice(header);
for col in data_cols {
@@ -213,13 +213,13 @@ fn take_hidden(
Some((take_f32(data, cv, nv)?, lens))
}
/// Decode a batch egress frame (tag stripped), calling `route` with each request's
/// Decode a batch frame (tag stripped), calling `route` with each request's
/// [`ChunkEvent`] as it's decoded — one pass, no intermediate `Vec`, peak memory
/// one request. Column order matches `push_generation`.
///
/// `ok == false` means the frame was rejected. The caller discards everything it
/// routed and fails the frame's requests instead of forwarding a partial fan-out
/// (see `tokenizer_manager::egress`), so a rejected frame delivers nothing —
/// (see `tokenizer_manager::from_scheduler`), so a rejected frame delivers nothing —
/// `rids` exists precisely so those requests can be failed rather than left
/// waiting for a `Done` that no longer exists.
pub fn for_each_chunk(body: &[u8], mut route: impl FnMut(ChunkEvent)) -> Decoded {
@@ -498,23 +498,23 @@ pub struct Decoded {
pub rids: Vec<Rid>,
}
/// Frame a control result `[rid, payload]` for the egress ring (tag prepended).
pub fn frame_egress_result(rid: &str, payload: &[u8]) -> Bytes {
/// Frame a control result `[rid, payload]` for the response ring (tag prepended).
pub fn frame_control_result(rid: &str, payload: &[u8]) -> Bytes {
use rmpv::Value;
let arr = Value::Array(vec![Value::from(rid), Value::Binary(payload.to_vec())]);
let mut buf = Vec::with_capacity(1 + payload.len() + rid.len() + 8);
buf.push(EGRESS_TAG_RESULT);
buf.push(DISPATCH_TAG_RESULT);
let _ = rmpv::encode::write_value(&mut buf, &arr);
Bytes::from(buf)
}
/// Frame a per-request failure `[rid, message]` for the egress ring — routes a
/// Frame a per-request failure `[rid, message]` for the response — routes a
/// terminal error back to the owning request (→ HTTP 400) instead of crashing.
pub fn frame_egress_error(rid: &str, message: &str) -> Bytes {
pub fn frame_error(rid: &str, message: &str) -> Bytes {
use rmpv::Value;
let arr = Value::Array(vec![Value::from(rid), Value::from(message)]);
let mut buf = Vec::with_capacity(1 + rid.len() + message.len() + 8);
buf.push(EGRESS_TAG_ERROR);
buf.push(DISPATCH_TAG_ERROR);
let _ = rmpv::encode::write_value(&mut buf, &arr);
Bytes::from(buf)
}
@@ -618,11 +618,11 @@ mod tests {
let header = [1u8, 2, 3];
let a = [10u8, 11];
let b = [12u8, 13, 14];
let multi = frame_egress_batch_cols(&header, &[&a[..], &b[..]]);
let multi = frame_decode_batch_cols(&header, &[&a[..], &b[..]]);
let joined: Vec<u8> = a.iter().chain(&b).copied().collect();
let single = frame_egress_batch_cols(&header, &[joined.as_slice()]);
let single = frame_decode_batch_cols(&header, &[joined.as_slice()]);
assert_eq!(multi, single);
assert_eq!(multi[0], EGRESS_TAG_BATCH);
assert_eq!(multi[0], DISPATCH_TAG_BATCH);
assert_eq!(
u32::from_le_bytes([multi[1], multi[2], multi[3], multi[4]]),
3
@@ -664,8 +664,8 @@ mod tests {
.flat_map(|x| x.to_le_bytes())
.collect();
let framed = frame_egress_batch_cols(&header, &[&data]);
assert_eq!(framed[0], EGRESS_TAG_BATCH);
let framed = frame_decode_batch_cols(&header, &[&data]);
assert_eq!(framed[0], DISPATCH_TAG_BATCH);
let mut events = Vec::new();
assert!(for_each_chunk(&framed[1..], |ev| events.push(ev)).ok);
assert_eq!(events.len(), 3);
@@ -690,13 +690,13 @@ mod tests {
assert_eq!(events[2].prompt_tokens, 6);
// A plain decode frame carries no extras columns at all, so the per-frame
// `has_extras` guard must skip the extras machinery entirely for every
// request (this is the tm-egress hot path — see `for_each_chunk`).
// request (this is the from-scheduler hot path — see `for_each_chunk`).
assert!(events.iter().all(|e| e.extras.is_none()));
}
/// A header whose column lengths exceed the data buffer (a Python/Rust
/// positional-ABI drift, or a truncated frame) is rejected: `for_each_chunk`
/// returns false and routes nothing — it must NOT panic the sole egress thread
/// returns false and routes nothing — it must NOT panic the sole from_scheduler thread
/// on an out-of-bounds slice. Built the way Python emits (positional msgpack
/// header + concatenated data columns).
#[test]
@@ -717,7 +717,7 @@ mod tests {
rmpv::encode::write_value(&mut header, &header_arr).unwrap();
let data: Vec<u8> = [0i32].iter().flat_map(|x| x.to_le_bytes()).collect(); // 4 bytes
let framed = frame_egress_batch_cols(&header, &[&data]);
let framed = frame_decode_batch_cols(&header, &[&data]);
let mut routed = 0usize;
let decoded = for_each_chunk(&framed[1..], |_| routed += 1);
assert!(!decoded.ok, "malformed frame must be rejected, not decoded");
@@ -759,7 +759,7 @@ mod tests {
data.extend(i(&[10, 20])); // token_ids
data.extend(f(&[-0.1, -0.2, -0.3, -0.4])); // out_top_val (sum of poslens = 4)
data.extend(i(&[1, 2, 3, 4])); // out_top_idx
let framed = frame_egress_batch_cols(&header, &[&data]);
let framed = frame_decode_batch_cols(&header, &[&data]);
let mut routed = 0usize;
assert!(
!for_each_chunk(&framed[1..], |_| routed += 1).ok,
@@ -790,7 +790,7 @@ mod tests {
let mut data = Vec::new();
data.extend(i(&[10, 20])); // token_ids
data.extend(f(&[0.1, 0.2, 0.3])); // hidden_val (sum of poslens = 3)
let framed = frame_egress_batch_cols(&header, &[&data]);
let framed = frame_decode_batch_cols(&header, &[&data]);
let mut routed = 0usize;
assert!(
!for_each_chunk(&framed[1..], |_| routed += 1).ok,
@@ -812,7 +812,7 @@ mod tests {
]);
let mut header = Vec::new();
rmpv::encode::write_value(&mut header, &header_arr).unwrap();
let framed = frame_egress_batch_cols(&header, &[&[0u8; 4][..]]);
let framed = frame_decode_batch_cols(&header, &[&[0u8; 4][..]]);
let mut routed = 0usize;
let decoded = for_each_chunk(&framed[1..], |_| routed += 1);
assert!(!decoded.ok);
@@ -839,12 +839,12 @@ mod tests {
let mut header = Vec::new();
rmpv::encode::write_value(&mut header, &header_arr).unwrap();
let data: Vec<u8> = vec![0u8; 8]; // 4 bytes too many
let framed = frame_egress_batch_cols(&header, &[&data]);
let framed = frame_decode_batch_cols(&header, &[&data]);
let decoded = for_each_chunk(&framed[1..], |_| {});
assert!(!decoded.ok, "header and data must agree exactly");
}
/// Ingress/egress rid agreement: the rid decoded from the frame must be the
/// Request/response rid agreement: the rid decoded from the frame must be the
/// one Python sent, AND both sides must derive the same shard from it. The
/// partition key is memoized inside `Rid`, so a per-conversion hasher seed
/// would send a request's chunks to a shard that never registered it.
@@ -862,7 +862,7 @@ mod tests {
rmpv::encode::write_value(&mut header, &header_arr).unwrap();
let data: Vec<u8> = [0i32].iter().flat_map(|x| x.to_le_bytes()).collect();
let framed = frame_egress_batch_cols(&header, &[&data]);
let framed = frame_decode_batch_cols(&header, &[&data]);
let mut events = Vec::new();
assert!(for_each_chunk(&framed[1..], |ev| events.push(ev)).ok);
assert_eq!(events.len(), 1);
@@ -931,7 +931,7 @@ mod tests {
data.extend(i(&[10, 11])); // out_top_idx
data.extend(f(&[0.1, 0.2, 0.3])); // hidden_val (1 row, dim 3)
let framed = frame_egress_batch_cols(&header, &[&data]);
let framed = frame_decode_batch_cols(&header, &[&data]);
let mut events = Vec::new();
assert!(for_each_chunk(&framed[1..], |ev| events.push(ev)).ok);
assert_eq!(events.len(), 2);
@@ -999,7 +999,7 @@ mod tests {
data.extend(i(&[10, 20])); // token_ids
data.extend(f(&[-0.5, -0.6])); // out_lp_val (req0)
data.extend(i(&[10, 99])); // out_lp_idx
let framed = frame_egress_batch_cols(&header, &[&data]);
let framed = frame_decode_batch_cols(&header, &[&data]);
let mut events = Vec::new();
assert!(
for_each_chunk(&framed[1..], |ev| events.push(ev)).ok,
@@ -1093,7 +1093,7 @@ mod tests {
data.extend(i(&[61])); // in_tid_idx
data.extend(f(&[7.1, 7.2, 7.3])); // hidden_val
let framed = frame_egress_batch_cols(&header, &[&data]);
let framed = frame_decode_batch_cols(&header, &[&data]);
let mut events = Vec::new();
assert!(for_each_chunk(&framed[1..], |ev| events.push(ev)).ok);
assert_eq!(events.len(), 1);
@@ -1137,7 +1137,7 @@ mod tests {
let mut header = Vec::new();
rmpv::encode::write_value(&mut header, &header_arr).unwrap();
let data: Vec<u8> = [7i32, 8].iter().flat_map(|x| x.to_le_bytes()).collect();
let framed = frame_egress_batch_cols(&header, &[&data]);
let framed = frame_decode_batch_cols(&header, &[&data]);
let mut events = Vec::new();
assert!(for_each_chunk(&framed[1..], |ev| events.push(ev)).ok);
// Each chunk carries its OWN rid — the value a shard keys its table on.
@@ -1147,13 +1147,7 @@ mod tests {
}
/// The common frame must stay small: logprob/hidden columns are boxed behind
/// `ChunkExtras`, so the inline decode array is a few KiB — not MiB — even at
/// batch 4096. A regression that inlines a rare column would blow this up.
///
/// The budget moved 128 → 144 when the rid gained its memoized partition key
/// (`String` + `u64`). That costs 8 bytes × batch per decode step and buys not
/// re-hashing the rid on every chunk in the egress bucketing loop — a
/// deliberate trade, not drift.
/// `ChunkExtras`.
#[test]
fn chunk_event_frame_stays_small() {
let sz = std::mem::size_of::<ChunkEvent>();
@@ -1182,7 +1176,7 @@ mod rid_recovery_tests {
cols.extend((0..extra_cols).map(|_| Value::from("unexpected")));
let mut header = Vec::new();
rmpv::encode::write_value(&mut header, &Value::Array(cols)).unwrap();
let framed = frame_egress_batch_cols(&header, &[]);
let framed = frame_decode_batch_cols(&header, &[]);
let decoded = for_each_chunk(&framed[1..], |_| {});
assert!(!decoded.ok, "arity {extra_cols}: must reject");
assert_eq!(
rust/sglang-server/src/message/sampling.rs
+3 -27
View File
@@ -2,29 +2,6 @@
//! (python/sglang/srt/sampling/sampling_params.py): every field, plus its
//! `__post_init__` → `normalize` → `verify` pipeline (run in that order, as
//! `TokenizerManager._create_tokenized_object` does).
//!
//! The embedded Rust server replaces the Python `TokenizerManager`, which is the
//! only place those three run on the normal (zmq) path. Running them here, in the
//! ingress `Normalizing` FSM step, keeps the per-request CPU (notably the
//! stop-string work) off the scheduler's latency-critical loop. We set
//! `is_normalized=true` on the wire so the scheduler's `__post_init__` and
//! `normalize` early-return; its `verify` is likewise skipped (we did it here).
//!
//! KEEP IN SYNC with `sampling_params.py`: the field list, defaults and ranges
//! below mirror that file, and the struct is serialized by field name into the
//! `TokenizedGenerateReqInput` header, so a renamed/added Python field must be
//! mirrored here (an unknown key would be silently dropped by msgspec).
//!
//! Two deliberate deviations, both safe over-estimates or stricter:
//! * `stop_str_max_len` is the stop string's **UTF-8 byte length** — a provably
//! safe over-estimate of its token length (a token spans ≥ 1 byte, so
//! `bytes ≥ tokens`; `chars` is *not* a bound — one char can be several
//! tokens, e.g. `𓀀` → 3). The scheduler uses it only as a match-window
//! *size* (capped at the output length), so an over-estimate matches the same
//! stops — only an under-estimate misses. Python encodes each stop with the
//! tokenizer for the exact count; the byte bound avoids needing it here.
//! * `n > 1` (parallel sampling) is rejected — the rust egress maps one rid to
//! one response, so every sample past the first would be dropped.
use std::collections::BTreeMap;
use std::fmt;
@@ -33,9 +10,8 @@ use serde::de::value::{MapAccessDeserializer, SeqAccessDeserializer};
use serde::de::{MapAccess, SeqAccess, Visitor};
use serde::{Deserialize, Deserializer, Serialize};
use super::OneOrMany;
use crate::error::Error;
use crate::utils::regex::RegexPattern;
use super::types::OneOrMany;
use crate::utils::{error::Error, regex::RegexPattern};
/// `_SAMPLING_EPS` — temperatures in `[0, eps)` mean greedy decoding.
const SAMPLING_EPS: f64 = 1e-6;
@@ -495,7 +471,7 @@ impl SamplingParams {
"Only one of regex, json_schema, or ebnf can be set".into()
));
}
// Not a Python restriction: the rust egress maps one rid to one response,
// Not a Python restriction: the rust from_scheduler maps one rid to one response,
// so parallel sampling would drop all but the first sample. This is the
// only place it is rejected — `n` lives in `sampling_params`, where
// Python reads it, and the `/generate` body has no `n` of its own.
rust/sglang-server/src/mm.rs
-413
View File
@@ -1,413 +0,0 @@
//! Multimodal worker pool.
//!
//! Rust threads drain requests parked in `Encoding` and run the `sglang-mm`
//! pipeline registered by `Server.start_mm_workers` (decode → preprocess →
//! placeholder expansion → M-RoPE, GIL-free). Each worker parks the result
//! buffers in the rid-keyed [`Sidecar`] and returns only the expanded ids;
//! Python attaches the buffers at drain time (`Server.take_mm`). Inputs the
//! pipeline cannot serve are rejected to the client — no Python fallback.
use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use crate::message::MmRequest;
use crate::runtime::Runnable;
use crate::tokenizer::TextTokenizer;
use crate::tokenizer_manager::TmEvent;
/// A named POSIX shared-memory segment owning its name: dropped → unlinked.
///
/// Written by an MM worker so the TP broadcast carries a ~100-byte
/// `ShmPointerMMData` stub instead of the ~20 MB feature tensor, and every
/// rank maps it in parallel. Python's `materialize()` unlinks after cloning;
/// this `Drop` covers the paths where the buffers never reach Python (aborted
/// while parked, late result purged).
pub struct ShmSegment {
name: String,
}
impl ShmSegment {
/// Create `/dev/shm/{name}` holding exactly `bytes`. No leading slash —
/// the name must suit Python's `SharedMemory(name=…)` (shm_open adds one).
pub fn create(name: String, bytes: &[u8]) -> Result<Self, String> {
let c_name = std::ffi::CString::new(format!("/{name}"))
.map_err(|_| "shm name contains NUL".to_string())?;
// SAFETY: plain POSIX calls on a name we own; every handle created
// below is closed/unmapped on all paths.
unsafe {
let fd = libc::shm_open(
c_name.as_ptr(),
libc::O_CREAT | libc::O_EXCL | libc::O_RDWR,
0o600,
);
if fd < 0 {
return Err(format!(
"shm_open({name}): {}",
std::io::Error::last_os_error()
));
}
let segment = Self { name }; // unlink from here on any failure
if libc::ftruncate(fd, bytes.len() as libc::off_t) != 0 {
let e = std::io::Error::last_os_error();
libc::close(fd);
return Err(format!("ftruncate({}): {e}", segment.name));
}
let ptr = libc::mmap(
std::ptr::null_mut(),
bytes.len(),
libc::PROT_WRITE,
libc::MAP_SHARED,
fd,
0,
);
libc::close(fd);
if ptr == libc::MAP_FAILED {
return Err(format!(
"mmap({}): {}",
segment.name,
std::io::Error::last_os_error()
));
}
std::ptr::copy_nonoverlapping(bytes.as_ptr(), ptr.cast::<u8>(), bytes.len());
libc::munmap(ptr, bytes.len());
Ok(segment)
}
}
/// Hand the segment — and the duty to unlink — to the caller (Python, at
/// drain time).
pub fn into_name(self) -> String {
std::mem::take(&mut std::mem::ManuallyDrop::new(self).name)
}
}
impl Drop for ShmSegment {
fn drop(&mut self) {
if let Ok(c_name) = std::ffi::CString::new(format!("/{}", self.name)) {
// SAFETY: unlinking a name we created; ENOENT (already unlinked
// by Python's materialize) is fine to ignore.
unsafe { libc::shm_unlink(c_name.as_ptr()) };
}
}
}
/// Unique segment names: the pid separates server restarts (a crash can leak
/// segments under the old pid), the counter separates results within one.
fn shm_name(item: usize) -> String {
static COUNTER: AtomicU64 = AtomicU64::new(0);
let n = COUNTER.fetch_add(1, Ordering::Relaxed);
format!("sglmm-{}-{n}-{item}", std::process::id())
}
/// Python parity: caller hashes override the computed ones so an external
/// router's keys align with the prefix cache. A length mismatch or malformed
/// entry warns and keeps the computed hash — never blocks the request.
fn apply_caller_hashes(hashes: &mut [u64], caller: &[String]) {
if caller.is_empty() {
return;
}
if caller.len() != hashes.len() {
tracing::warn!(
caller = caller.len(),
items = hashes.len(),
"mm_hashes length != mm item count; ignoring caller hashes"
);
return;
}
for (hash, entry) in hashes.iter_mut().zip(caller) {
match parse_caller_hash(entry) {
Some(v) => *hash = v,
None => tracing::warn!(%entry, "malformed mm_hashes entry; keeping computed hash"),
}
}
}
/// Hex of any width, as Python's `int(hex_hash, 16)` takes it (a full SHA-256
/// being the common case), keeping the low 64 bits — only the low 30 are
/// observable, through `_compute_pad_value`.
fn parse_caller_hash(entry: &str) -> Option<u64> {
let hex = entry.strip_prefix("0x").unwrap_or(entry);
if hex.is_empty() || !hex.bytes().all(|b| b.is_ascii_hexdigit()) {
return None;
}
u64::from_str_radix(&hex[hex.len().saturating_sub(16)..], 16).ok()
}
/// One parked result: the buffers the drain-time Python adapter needs (the
/// expanded `input_ids` travel separately, via `TmEvent::MmEncoded`). The qwen
/// drain shape (`sglang_mm::qwen_vl::pack_drain`); generalizes to a
/// named-tensor handoff once a family needs a different one.
pub struct MmSidecarEntry {
pub features: FeatureStore,
pub grids: Vec<[u32; 3]>,
pub hashes: Vec<u64>,
pub offsets: Vec<(u32, u32)>,
pub mrope: Vec<i64>,
pub mrope_delta: i64,
}
/// Where a result's feature buffers live between worker and drain.
pub enum FeatureStore {
/// In-process; the drain wraps them zero-copy. Single-rank serving, or the
/// shm fallback. Under TP the whole buffer would ride `broadcast_pyobj`.
Inline(Vec<f32>),
/// One POSIX segment per item, written by the worker; only the names cross
/// ranks. See [`ShmSegment`].
Shm(Vec<ShmSegment>),
}
/// Results parked between a worker's `MmEncoded` and the scheduler drain, keyed
/// by rid. Owns the lifecycle so entries never leak: [`park`](Self::park)
/// strictly before `MmEncoded`, [`take`](Self::take) at the drain,
/// [`purge`](Self::purge) for requests that die while parked.
#[derive(Clone, Default)]
pub struct Sidecar(Arc<Mutex<HashMap<String, MmSidecarEntry>>>);
impl Sidecar {
pub fn park(&self, rid: String, entry: MmSidecarEntry) {
self.0.lock().unwrap().insert(rid, entry);
}
pub fn take(&self, rid: &str) -> Option<MmSidecarEntry> {
self.0.lock().unwrap().remove(rid)
}
pub fn purge(&self, rid: &str) {
self.0.lock().unwrap().remove(rid);
}
}
/// Shared state of the mm path, built once at `start_mm_workers`.
pub struct Context {
pub family: Box<dyn sglang_mm::pipeline::MmFamilyProcessor>,
/// `None` under `skip_tokenizer_init` (requests must carry `input_ids`).
pub tokenizer: Option<Arc<dyn TextTokenizer>>,
pub sidecar: Sidecar,
/// Park feature buffers in POSIX shm. Set by the Python launcher
/// (`NativeMmHost._use_feature_shm`) exactly when the scheduler broadcasts
/// across TP ranks and will unwrap `ShmPointerMMData`.
pub feature_shm: bool,
}
impl Context {
pub fn new(
spec_json: &str,
tokenizer: Option<Arc<dyn TextTokenizer>>,
sidecar: Sidecar,
) -> Result<Self, String> {
let feature_shm = serde_json::from_str::<serde_json::Value>(spec_json)
.ok()
.and_then(|v| v.get("feature_shm").and_then(|b| b.as_bool()))
.unwrap_or(false);
Ok(Self {
family: sglang_mm::registry::pipeline_from_spec(spec_json)?,
tokenizer,
sidecar,
feature_shm,
})
}
}
/// Run the pipeline for one request. `Ok` returns the final expanded ids, the
/// buffers already parked; `Err` rejects the request back to the client.
fn process(
ctx: &Context,
rid: &crate::ids::Rid,
mut work: crate::message::MmWorkItem,
) -> Result<Vec<i32>, String> {
let caller_hashes = std::mem::take(&mut work.mm_hashes);
let input = crate::message::mm_payload::to_mm_input(work)?;
let output = sglang_mm::driver::process(ctx.family.as_ref(), input, |text| {
let tokenizer = ctx.tokenizer.as_ref().ok_or_else(|| {
"skip_tokenizer_init is set: multimodal text prompts require input_ids".to_string()
})?;
tokenizer.encode(text).map_err(|error| error.to_string())
})?;
let mut drain = sglang_mm::qwen_vl::pack_drain(output)?;
apply_caller_hashes(&mut drain.hashes, &caller_hashes);
let features = if ctx.feature_shm {
park_features_in_shm(&drain.features, &drain.grids)
} else {
FeatureStore::Inline(drain.features)
};
ctx.sidecar.park(
rid.as_str().to_owned(),
MmSidecarEntry {
features,
grids: drain.grids,
hashes: drain.hashes,
offsets: drain.offsets,
mrope: drain.mrope,
mrope_delta: drain.mrope_delta,
},
);
Ok(drain.input_ids)
}
/// Split the flat feature buffer per item (`t*h*w` rows per grid) and park each
/// slice in its own segment. Any shm failure (`/dev/shm` full, odd shape) falls
/// back to inline, as Python's `_wrap_shm_or_inline` does: degrade to the slow
/// path, never fail the request.
fn park_features_in_shm(features: &[f32], grids: &[[u32; 3]]) -> FeatureStore {
let total_rows: usize = grids
.iter()
.map(|g| g[0] as usize * g[1] as usize * g[2] as usize)
.sum();
if total_rows == 0 || !features.len().is_multiple_of(total_rows) {
return FeatureStore::Inline(features.to_vec());
}
let dim = features.len() / total_rows;
let mut segments = Vec::with_capacity(grids.len());
let mut row = 0usize;
for (item, grid) in grids.iter().enumerate() {
let rows = grid[0] as usize * grid[1] as usize * grid[2] as usize;
let slice = &features[row * dim..(row + rows) * dim];
row += rows;
match ShmSegment::create(shm_name(item), bytemuck::cast_slice(slice)) {
Ok(segment) => segments.push(segment),
Err(error) => {
tracing::warn!(%error, "mm: shm feature transport failed; falling back to inline");
return FeatureStore::Inline(features.to_vec());
}
}
}
FeatureStore::Shm(segments)
}
/// One MM worker, spawned via `Runtime::spawn_mm_pool` (which owns the
/// pinning policy for this pool — see its docs).
pub struct MmWorker {
rx: flume::Receiver<MmRequest>,
tm: flume::Sender<TmEvent>,
ctx: Arc<Context>,
}
impl MmWorker {
pub fn new(
rx: flume::Receiver<MmRequest>,
tm: flume::Sender<TmEvent>,
ctx: Arc<Context>,
) -> Self {
Self { rx, tm, ctx }
}
}
impl Runnable for MmWorker {
/// Drain until the mm channel closes (tm-ingress drops its sender on
/// shutdown). One request at a time, so the pool size bounds MM
/// concurrency; an error rejects the request back to the client.
fn run(self) {
while let Ok(req) = self.rx.recv() {
let rid = req.rid;
let event = match process(&self.ctx, &rid, req.work) {
Ok(input_ids) => {
tracing::debug!(%rid, tokens = input_ids.len(), "mm: processed");
TmEvent::MmEncoded { rid, input_ids }
}
Err(message) => {
tracing::warn!(%rid, %message, "mm processing rejected");
TmEvent::MmFailed { rid, message }
}
};
if self.tm.send(event).is_err() {
return; // tm-ingress gone: shutdown
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
/// Caller hashes override computed ones; mismatched lengths and malformed
/// entries fall back per item, never reject (Python parity).
#[test]
fn caller_hashes_override_with_fallback() {
let mut hashes = vec![1, 2, 3];
apply_caller_hashes(&mut hashes, &[]);
assert_eq!(hashes, [1, 2, 3]);
apply_caller_hashes(&mut hashes, &["ff".into()]); // length mismatch
assert_eq!(hashes, [1, 2, 3]);
apply_caller_hashes(&mut hashes, &["ff".into(), "not-hex".into(), "0x10".into()]);
assert_eq!(hashes, [0xff, 2, 0x10]);
}
/// A full SHA-256 (what routers send) keeps its low 64 bits rather than
/// falling back, so the pad value matches Python's wide `int`.
#[test]
fn caller_hashes_accept_arbitrary_width() {
let sha256 = "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855";
let mut hashes = vec![1];
apply_caller_hashes(&mut hashes, &[sha256.into()]);
assert_eq!(hashes, [0xa495991b7852b855]);
assert_eq!(hashes[0] % (1 << 30), 944_945_237); // int(sha256, 16) % (1 << 30)
// Width alone is never malformed; a non-hex digit still is.
assert_eq!(parse_caller_hash(&"f".repeat(64)), Some(u64::MAX));
assert_eq!(parse_caller_hash("0x"), None);
assert_eq!(parse_caller_hash(""), None);
}
fn shm_path(name: &str) -> std::path::PathBuf {
std::path::Path::new("/dev/shm").join(name)
}
/// The segment holds exactly the written bytes and dropping it unlinks —
/// the leak guard for results purged before Python takes them.
#[test]
fn segment_roundtrip_and_drop_unlinks() {
let name = shm_name(0);
let payload: Vec<u8> = (0..255u8).collect();
let segment = ShmSegment::create(name.clone(), &payload).unwrap();
assert_eq!(std::fs::read(shm_path(&name)).unwrap(), payload);
drop(segment);
assert!(!shm_path(&name).exists(), "drop must unlink");
}
/// `into_name` transfers the unlink duty to the caller (Python's
/// `materialize()`), so the segment must survive the handoff.
#[test]
fn into_name_disarms_the_unlink() {
let segment = ShmSegment::create(shm_name(0), &[1, 2, 3]).unwrap();
let name = segment.into_name();
assert!(shm_path(&name).exists(), "handoff must not unlink");
// manual cleanup for the test
let c = std::ffi::CString::new(format!("/{name}")).unwrap();
unsafe { libc::shm_unlink(c.as_ptr()) };
}
/// Per-item slicing follows the grid row counts, so Python's
/// `(rows, feature_dim)` reshape of a segment sees only its own item.
#[test]
fn park_splits_features_by_grid() {
// Two items: grids (1,2,2)=4 rows and (1,1,2)=2 rows, dim=3.
let features: Vec<f32> = (0..18).map(|i| i as f32).collect();
let grids = [[1, 2, 2], [1, 1, 2]];
let FeatureStore::Shm(segments) = park_features_in_shm(&features, &grids) else {
panic!("expected shm store");
};
assert_eq!(segments.len(), 2);
let read = |seg: &ShmSegment| -> Vec<u8> { std::fs::read(shm_path(&seg.name)).unwrap() };
assert_eq!(
read(&segments[0]),
bytemuck::cast_slice::<f32, u8>(&features[..12])
);
assert_eq!(
read(&segments[1]),
bytemuck::cast_slice::<f32, u8>(&features[12..])
);
}
/// A degenerate shape must degrade to inline, never a shm-side panic.
#[test]
fn shape_surprise_falls_back_inline() {
let features = vec![0.0f32; 7]; // not divisible by 2 rows
let grids = [[1, 1, 2]];
assert!(matches!(
park_features_in_shm(&features, &grids),
FeatureStore::Inline(_)
));
}
}
rust/sglang-server/src/multi_modality.rs
+6
View File
@@ -0,0 +1,6 @@
//! Multimodal worker pool.
pub mod payload;
mod shm;
pub mod sidecar;
pub mod worker;
rust/sglang-server/src/message/mm_payload.rs → rust/sglang-server/src/multi_modality/payload.rs
+2 -3
View File
@@ -9,7 +9,7 @@ use bytes::Bytes;
use rmpv::Value;
use sglang_mm::driver::{ImageSource, MmInput};
use super::request::MmWorkItem;
use crate::message::request::MmWorkItem;
/// True for sources the API layer must resolve before MM dispatch: I/O — network
/// *or* disk, since a network mount can hang past any HTTP timeout — never runs
@@ -114,8 +114,7 @@ fn collect_images(
}
/// Rust mirror of Python `has_valid_data`: `nil` and (recursively) empty or
/// all-nil lists don't count as multimodal input. Shared with the ingress
/// `has_multimodal` check so routing and parsing cannot drift.
/// all-nil lists don't count as multimodal input.
pub fn value_present(value: &Value) -> bool {
match value {
Value::Nil => false,
rust/sglang-server/src/multi_modality/shm.rs
+122
View File
@@ -0,0 +1,122 @@
//! POSIX shared-memory transport for feature tensors.
use std::sync::atomic::{AtomicU64, Ordering};
/// A named POSIX shared-memory segment owning its name: dropped → unlinked.
///
/// Written by an MM worker so the TP broadcast carries a ~100-byte
/// `ShmPointerMMData` stub instead of the ~20 MB feature tensor, and every
/// rank maps it in parallel. Python's `materialize()` unlinks after cloning;
/// this `Drop` covers the paths where the buffers never reach Python (aborted
/// while parked, late result purged).
pub struct ShmSegment {
pub(super) name: String,
}
impl ShmSegment {
/// Create `/dev/shm/{name}` holding exactly `bytes`. No leading slash —
/// the name must suit Python's `SharedMemory(name=…)` (shm_open adds one).
pub fn create(name: String, bytes: &[u8]) -> Result<Self, String> {
let c_name = std::ffi::CString::new(format!("/{name}"))
.map_err(|_| "shm name contains NUL".to_string())?;
// SAFETY: plain POSIX calls on a name we own; every handle created
// below is closed/unmapped on all paths.
unsafe {
let fd = libc::shm_open(
c_name.as_ptr(),
libc::O_CREAT | libc::O_EXCL | libc::O_RDWR,
0o600,
);
if fd < 0 {
return Err(format!(
"shm_open({name}): {}",
std::io::Error::last_os_error()
));
}
let segment = Self { name }; // unlink from here on any failure
if libc::ftruncate(fd, bytes.len() as libc::off_t) != 0 {
let e = std::io::Error::last_os_error();
libc::close(fd);
return Err(format!("ftruncate({}): {e}", segment.name));
}
let ptr = libc::mmap(
std::ptr::null_mut(),
bytes.len(),
libc::PROT_WRITE,
libc::MAP_SHARED,
fd,
0,
);
libc::close(fd);
if ptr == libc::MAP_FAILED {
return Err(format!(
"mmap({}): {}",
segment.name,
std::io::Error::last_os_error()
));
}
std::ptr::copy_nonoverlapping(bytes.as_ptr(), ptr.cast::<u8>(), bytes.len());
libc::munmap(ptr, bytes.len());
Ok(segment)
}
}
/// Hand the segment — and the duty to unlink — to the caller (Python, at
/// drain time).
pub fn into_name(self) -> String {
std::mem::take(&mut std::mem::ManuallyDrop::new(self).name)
}
}
impl Drop for ShmSegment {
fn drop(&mut self) {
if let Ok(c_name) = std::ffi::CString::new(format!("/{}", self.name)) {
// SAFETY: unlinking a name we created; ENOENT (already unlinked
// by Python's materialize) is fine to ignore.
unsafe { libc::shm_unlink(c_name.as_ptr()) };
}
}
}
/// Unique segment names: the pid separates server restarts (a crash can leak
/// segments under the old pid), the counter separates results within one.
pub(super) fn shm_name(item: usize) -> String {
static COUNTER: AtomicU64 = AtomicU64::new(0);
let n = COUNTER.fetch_add(1, Ordering::Relaxed);
format!("sglmm-{}-{n}-{item}", std::process::id())
}
/// Test helper shared with the sidecar's parking tests.
#[cfg(test)]
pub(super) fn shm_path(name: &str) -> std::path::PathBuf {
std::path::Path::new("/dev/shm").join(name)
}
#[cfg(test)]
mod tests {
use super::*;
/// The segment holds exactly the written bytes and dropping it unlinks —
/// the leak guard for results purged before Python takes them.
#[test]
fn segment_roundtrip_and_drop_unlinks() {
let name = shm_name(0);
let payload: Vec<u8> = (0..255u8).collect();
let segment = ShmSegment::create(name.clone(), &payload).unwrap();
assert_eq!(std::fs::read(shm_path(&name)).unwrap(), payload);
drop(segment);
assert!(!shm_path(&name).exists(), "drop must unlink");
}
/// `into_name` transfers the unlink duty to the caller (Python's
/// `materialize()`), so the segment must survive the handoff.
#[test]
fn into_name_disarms_the_unlink() {
let segment = ShmSegment::create(shm_name(0), &[1, 2, 3]).unwrap();
let name = segment.into_name();
assert!(shm_path(&name).exists(), "handoff must not unlink");
// manual cleanup for the test
let c = std::ffi::CString::new(format!("/{name}")).unwrap();
unsafe { libc::shm_unlink(c.as_ptr()) };
}
}
rust/sglang-server/src/multi_modality/sidecar.rs
+121
View File
@@ -0,0 +1,121 @@
//! Rid-keyed parking of finished results between an MM worker and the
//! scheduler drain.
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use super::shm::{ShmSegment, shm_name};
/// One parked result: the buffers the drain-time Python adapter needs (the
/// expanded `input_ids` travel separately, via `TmEvent::MmEncoded`). The qwen
/// drain shape (`sglang_mm::qwen_vl::pack_drain`); generalizes to a
/// named-tensor handoff once a family needs a different one.
///
/// Constructed from outside the module only by tests; the worker parks every
/// real entry itself.
pub struct MmSidecarEntry {
pub features: FeatureStore,
pub grids: Vec<[u32; 3]>,
pub hashes: Vec<u64>,
pub offsets: Vec<(u32, u32)>,
pub mrope: Vec<i64>,
pub mrope_delta: i64,
}
/// Where a result's feature buffers live between worker and drain.
pub enum FeatureStore {
/// In-process; the drain wraps them zero-copy. Single-rank serving, or the
/// shm fallback. Under TP the whole buffer would ride `broadcast_pyobj`.
Inline(Vec<f32>),
/// One POSIX segment per item, written by the worker; only the names cross
/// ranks. See [`ShmSegment`].
Shm(Vec<ShmSegment>),
}
/// Results parked between a worker's `MmEncoded` and the scheduler drain, keyed
/// by rid. Owns the lifecycle so entries never leak: [`park`](Self::park)
/// strictly before `MmEncoded`, [`take`](Self::take) at the drain,
/// [`purge`](Self::purge) for requests that die while parked.
#[derive(Clone, Default)]
pub struct Sidecar(Arc<Mutex<HashMap<String, MmSidecarEntry>>>);
impl Sidecar {
pub fn park(&self, rid: String, entry: MmSidecarEntry) {
self.0.lock().unwrap().insert(rid, entry);
}
pub fn take(&self, rid: &str) -> Option<MmSidecarEntry> {
self.0.lock().unwrap().remove(rid)
}
pub fn purge(&self, rid: &str) {
self.0.lock().unwrap().remove(rid);
}
}
/// Split the flat feature buffer per item (`t*h*w` rows per grid) and park each
/// slice in its own segment. Any shm failure (`/dev/shm` full, odd shape) falls
/// back to inline, as Python's `_wrap_shm_or_inline` does: degrade to the slow
/// path, never fail the request.
pub(super) fn park_features_in_shm(features: &[f32], grids: &[[u32; 3]]) -> FeatureStore {
let total_rows: usize = grids
.iter()
.map(|g| g[0] as usize * g[1] as usize * g[2] as usize)
.sum();
if total_rows == 0 || !features.len().is_multiple_of(total_rows) {
return FeatureStore::Inline(features.to_vec());
}
let dim = features.len() / total_rows;
let mut segments = Vec::with_capacity(grids.len());
let mut row = 0usize;
for (item, grid) in grids.iter().enumerate() {
let rows = grid[0] as usize * grid[1] as usize * grid[2] as usize;
let slice = &features[row * dim..(row + rows) * dim];
row += rows;
match ShmSegment::create(shm_name(item), bytemuck::cast_slice(slice)) {
Ok(segment) => segments.push(segment),
Err(error) => {
tracing::warn!(%error, "mm: shm feature transport failed; falling back to inline");
return FeatureStore::Inline(features.to_vec());
}
}
}
FeatureStore::Shm(segments)
}
#[cfg(test)]
mod tests {
use super::super::shm::shm_path;
use super::*;
/// Per-item slicing follows the grid row counts, so Python's
/// `(rows, feature_dim)` reshape of a segment sees only its own item.
#[test]
fn park_splits_features_by_grid() {
// Two items: grids (1,2,2)=4 rows and (1,1,2)=2 rows, dim=3.
let features: Vec<f32> = (0..18).map(|i| i as f32).collect();
let grids = [[1, 2, 2], [1, 1, 2]];
let FeatureStore::Shm(segments) = park_features_in_shm(&features, &grids) else {
panic!("expected shm store");
};
assert_eq!(segments.len(), 2);
let read = |seg: &ShmSegment| -> Vec<u8> { std::fs::read(shm_path(&seg.name)).unwrap() };
assert_eq!(
read(&segments[0]),
bytemuck::cast_slice::<f32, u8>(&features[..12])
);
assert_eq!(
read(&segments[1]),
bytemuck::cast_slice::<f32, u8>(&features[12..])
);
}
/// A degenerate shape must degrade to inline, never a shm-side panic.
#[test]
fn shape_surprise_falls_back_inline() {
let features = vec![0.0f32; 7]; // not divisible by 2 rows
let grids = [[1, 1, 2]];
assert!(matches!(
park_features_in_shm(&features, &grids),
FeatureStore::Inline(_)
));
}
}
rust/sglang-server/src/multi_modality/worker.rs
+187
View File
@@ -0,0 +1,187 @@
//! The worker pool: drain MM requests, run the `sglang-mm` pipeline, park
//! the result buffers.
use std::sync::Arc;
use super::sidecar::{FeatureStore, MmSidecarEntry, Sidecar, park_features_in_shm};
use crate::message::config::MmSpec;
use crate::message::ids::Rid;
use crate::message::request::MmRequest;
use crate::tokenizer_manager::tokenizer::TextTokenizer;
use crate::tokenizer_manager::wiring::TmEvent;
use crate::utils::runtime::Runnable;
/// Python parity: caller hashes override the computed ones so an external
/// router's keys align with the prefix cache. A length mismatch or malformed
/// entry warns and keeps the computed hash — never blocks the request.
fn apply_caller_hashes(hashes: &mut [u64], caller: &[String]) {
if caller.is_empty() {
return;
}
if caller.len() != hashes.len() {
tracing::warn!(
caller = caller.len(),
items = hashes.len(),
"mm_hashes length != mm item count; ignoring caller hashes"
);
return;
}
for (hash, entry) in hashes.iter_mut().zip(caller) {
match parse_caller_hash(entry) {
Some(v) => *hash = v,
None => tracing::warn!(%entry, "malformed mm_hashes entry; keeping computed hash"),
}
}
}
/// Hex of any width, as Python's `int(hex_hash, 16)` takes it (a full SHA-256
/// being the common case), keeping the low 64 bits — only the low 30 are
/// observable, through `_compute_pad_value`.
fn parse_caller_hash(entry: &str) -> Option<u64> {
let hex = entry.strip_prefix("0x").unwrap_or(entry);
if hex.is_empty() || !hex.bytes().all(|b| b.is_ascii_hexdigit()) {
return None;
}
u64::from_str_radix(&hex[hex.len().saturating_sub(16)..], 16).ok()
}
/// Shared state of the mm path, built once at `start_mm_workers`.
pub struct Context {
pub family: Box<dyn sglang_mm::pipeline::MmFamilyProcessor>,
/// `None` under `skip_tokenizer_init` (requests must carry `input_ids`).
pub tokenizer: Option<Arc<dyn TextTokenizer>>,
pub sidecar: Sidecar,
/// Park feature buffers in POSIX shm. Set by the Python launcher
/// (`NativeMmHost._use_feature_shm`) exactly when the scheduler broadcasts
/// across TP ranks and will unwrap `ShmPointerMMData`.
pub feature_shm: bool,
}
impl Context {
pub fn new(
spec: MmSpec,
tokenizer: Option<Arc<dyn TextTokenizer>>,
sidecar: Sidecar,
) -> Result<Self, String> {
Ok(Self {
family: sglang_mm::registry::build_pipeline(spec.pipeline)?,
tokenizer,
sidecar,
feature_shm: spec.feature_shm,
})
}
}
/// Run the pipeline for one request. `Ok` returns the final expanded ids, the
/// buffers already parked; `Err` rejects the request back to the client.
fn process(
ctx: &Context,
rid: &Rid,
mut work: crate::message::request::MmWorkItem,
) -> Result<Vec<i32>, String> {
let caller_hashes = std::mem::take(&mut work.mm_hashes);
let input = super::payload::to_mm_input(work)?;
let output = sglang_mm::driver::process(ctx.family.as_ref(), input, |text| {
let tokenizer = ctx.tokenizer.as_ref().ok_or_else(|| {
"skip_tokenizer_init is set: multimodal text prompts require input_ids".to_string()
})?;
tokenizer.encode(text).map_err(|error| error.to_string())
})?;
let mut drain = sglang_mm::qwen_vl::pack_drain(output)?;
apply_caller_hashes(&mut drain.hashes, &caller_hashes);
let features = if ctx.feature_shm {
park_features_in_shm(&drain.features, &drain.grids)
} else {
FeatureStore::Inline(drain.features)
};
ctx.sidecar.park(
rid.as_str().to_owned(),
MmSidecarEntry {
features,
grids: drain.grids,
hashes: drain.hashes,
offsets: drain.offsets,
mrope: drain.mrope,
mrope_delta: drain.mrope_delta,
},
);
Ok(drain.input_ids)
}
/// One MM worker, spawned via `Runtime::spawn_mm_pool` (which owns the
/// pinning policy for this pool — see its docs).
pub struct MmWorker {
rx: flume::Receiver<MmRequest>,
tm: flume::Sender<TmEvent>,
ctx: Arc<Context>,
}
impl MmWorker {
pub fn new(
rx: flume::Receiver<MmRequest>,
tm: flume::Sender<TmEvent>,
ctx: Arc<Context>,
) -> Self {
Self { rx, tm, ctx }
}
}
impl Runnable for MmWorker {
/// Drain until the mm channel closes (to-scheduler drops its sender on
/// shutdown). One request at a time, so the pool size bounds MM
/// concurrency; an error rejects the request back to the client.
fn run(self) {
while let Ok(req) = self.rx.recv() {
let rid = req.rid;
let event = match process(&self.ctx, &rid, req.work) {
Ok(input_ids) => {
tracing::debug!(%rid, tokens = input_ids.len(), "mm: processed");
TmEvent::MmEncoded { rid, input_ids }
}
Err(message) => {
tracing::warn!(%rid, %message, "mm processing rejected");
TmEvent::MmFailed { rid, message }
}
};
if self.tm.send(event).is_err() {
return; // to-scheduler gone: shutdown
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
/// Caller hashes override computed ones; mismatched lengths and malformed
/// entries fall back per item, never reject (Python parity).
#[test]
fn caller_hashes_override_with_fallback() {
let mut hashes = vec![1, 2, 3];
apply_caller_hashes(&mut hashes, &[]);
assert_eq!(hashes, [1, 2, 3]);
apply_caller_hashes(&mut hashes, &["ff".into()]); // length mismatch
assert_eq!(hashes, [1, 2, 3]);
apply_caller_hashes(&mut hashes, &["ff".into(), "not-hex".into(), "0x10".into()]);
assert_eq!(hashes, [0xff, 2, 0x10]);
}
/// A full SHA-256 (what routers send) keeps its low 64 bits rather than
/// falling back, so the pad value matches Python's wide `int`.
#[test]
fn caller_hashes_accept_arbitrary_width() {
let sha256 = "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855";
let mut hashes = vec![1];
apply_caller_hashes(&mut hashes, &[sha256.into()]);
assert_eq!(hashes, [0xa495991b7852b855]);
assert_eq!(hashes[0] % (1 << 30), 944_945_237); // int(sha256, 16) % (1 << 30)
// Width alone is never malformed; a non-hex digit still is.
assert_eq!(parse_caller_hash(&"f".repeat(64)), Some(u64::MAX));
assert_eq!(parse_caller_hash("0x"), None);
assert_eq!(parse_caller_hash(""), None);
}
}
rust/sglang-server/src/runtime/config.rs
-312
View File
@@ -1,312 +0,0 @@
//! Runtime configuration: the rust-server boot knobs
//! ([`RustServerServerArgs`]), the typed view of the scheduler's `server_args`
//! dump ([`ServerArgs`] / [`ModelConfig`]), and the [`RuntimeConfig`] pairing
//! them for `runtime::start`.
use std::net::SocketAddr;
use std::sync::Arc;
/// Boot knobs specific to the embedded rust server — none of these exist in
/// the Python `server_args` dump (see [`ServerArgs`]); they arrive as explicit
/// `Server::start` parameters.
#[derive(Clone, Debug)]
pub struct RustServerServerArgs {
pub http_addr: SocketAddr,
pub api_worker_num: usize,
pub ingress_ring_cap: usize,
pub egress_ring_cap: usize,
pub channel_cap: usize,
/// CPU core ids the pools pin to (e.g. this rank's NUMA-local cores minus
/// the scheduler's reserved launch cores). `None` → run unpinned.
pub cores: Option<Vec<usize>>,
}
impl Default for RustServerServerArgs {
fn default() -> Self {
Self {
http_addr: "127.0.0.1:30000".parse().unwrap(),
api_worker_num: 2,
ingress_ring_cap: 8192,
egress_ring_cap: 8192,
channel_cap: 8192,
cores: None,
}
}
}
#[derive(Clone, Debug)]
pub struct RuntimeConfig {
/// Rust-server-only boot knobs (listen address, pool/ring sizes, pinning).
pub rust_server_args: RustServerServerArgs,
/// The scheduler's `server_args` dump (worker counts, tokenizer source,
/// config-endpoint metadata). `Arc` so cloning the config (and, downstream,
/// each `AppState`) is cheap; immutable after construction.
pub server_args: Arc<ServerArgs>,
}
impl Default for RuntimeConfig {
fn default() -> Self {
Self {
rust_server_args: RustServerServerArgs::default(),
server_args: Arc::new(
ServerArgs::from_json("{}").expect("empty server_args blob parses"),
),
}
}
}
/// The scheduler's startup blob (`RustServer._build_server_args`) parsed once into
/// typed fields: values are post-`__post_init__`; unrelated unknown keys are dropped.
#[derive(Debug, serde::Deserialize)]
pub struct ServerArgs {
/// HF repo id / local dir of the model, reported by `/get_model_info`.
#[serde(default)]
pub model_path: String,
/// Model name reported by `/v1/models` and `/server_info`.
#[serde(default)]
pub served_model_name: String,
/// Tokenizer source (model dir / `tokenizer.json` / HF repo id). Empty only
/// in minimal standalone blobs — then boot requires `skip_tokenizer_init`.
#[serde(default)]
pub tokenizer_path: String,
/// HF revision, used only when `tokenizer_path` is a repo id. `None` → main.
#[serde(default)]
pub revision: Option<String>,
/// Weight format selected by `--load-format`, reported by `/get_model_info`.
/// The blob carries the post-`__post_init__` value (`auto` is already
/// narrowed to `gguf` / `mistral` / `runai_streamer` / `remote` where the
/// checkpoint demands it). Not consumed for loading -- the scheduler owns
/// that; `None` only when the blob omits the key.
#[serde(default)]
pub load_format: Option<String>,
/// Operator-supplied weight version, reported by `/model_info`. Defaults to
/// `"default"` on the Python side, so it is present in every blob; `None`
/// only when the blob omits the key.
#[serde(default)]
pub weight_version: Option<String>,
/// HTTP bind address (see [`Self::bind`]).
#[serde(default = "default_host")]
pub host: String,
#[serde(default = "default_port")]
pub port: u16,
/// Log levels driving the access log — uvicorn runs at
/// `log_level_http or log_level` (see [`Self::http_access_log_enabled`]).
#[serde(default = "default_log_level")]
pub log_level: String,
#[serde(default)]
pub log_level_http: Option<String>,
/// Optional built-in chat-template name or path to a Jinja/legacy JSON
/// template file. Without an override, uses the tokenizer config template.
#[serde(default)]
pub chat_template: Option<String>,
/// Parser selected by `--tool-call-parser`.
#[serde(default)]
pub tool_call_parser: Option<String>,
/// Reasoning splitter selected by `--reasoning-parser` (e.g. deepseek-r1).
/// When set, chat completions strip the model's reasoning markers out of
/// `content` into `reasoning_content` — both unary and streaming.
#[serde(default)]
pub reasoning_parser: Option<String>,
/// Python's global default for whether an SSE stream ends with a usage chunk.
#[serde(default)]
pub stream_response_default_include_usage: bool,
/// Pinned tokenizer threads / detok shards (Python asserts both ≥ 1).
#[serde(default = "default_worker_num")]
pub tokenizer_worker_num: usize,
#[serde(default = "default_worker_num")]
pub detokenizer_worker_num: usize,
/// Token-ids-in / token-ids-out mode: no tokenizer load, raw `output_ids`
/// frames (drives the `Skip` detok backend and the ingress branch).
#[serde(default)]
pub skip_tokenizer_init: bool,
/// Streamed `/generate` frames carry per-step deltas instead of cumulative
/// text. Matches the Python `TokenizerManager`.
#[serde(default)]
pub incremental_streaming_output: bool,
/// PD-disaggregation role: `"null"` (unified), `"prefill"`, or `"decode"`.
/// (On prefill, the KV bootstrap registry is mounted on the api router —
/// see [`Self::enable_pd_bootstrap`].)
#[serde(default = "default_disaggregation_mode")]
pub disaggregation_mode: String,
/// The resolved Python `ModelConfig`, attached to the blob at dump time.
#[serde(default)]
pub model_config: ModelConfig,
/// Default sampling params advertised by `/get_model_info`, verbatim from
/// `server_args.preferred_sampling_params` (a JSON object or null).
#[serde(default)]
pub preferred_sampling_params: Option<serde_json::Value>,
/// Over-long inputs are truncated to fit the context instead of 400ing, and
/// `max_new_tokens` is clamped rather than rejected (Python
/// `TokenizerManager._validate_one_request`).
#[serde(default)]
pub allow_auto_truncate: bool,
/// `return_hidden_states` is refused unless the server was launched with it:
/// the scheduler simply won't produce them, so the request would 200 with the
/// field silently missing.
#[serde(default)]
pub enable_return_hidden_states: bool,
/// Output slots reserved per request on top of its input (eagle stores draft
/// tokens there). Not a `server_args` field — `TokenizerManager` derives it and
/// `RustServer._build_server_args` stamps it in, so both sides count alike.
#[serde(default)]
pub num_reserved_tokens: u64,
/// Launch-time stamps (not `server_args` fields): sglang package version
/// and the scheduler-derived KV token capacity, reported by `/server_info`.
#[serde(default)]
pub version: Option<String>,
#[serde(default)]
pub max_total_num_tokens: Option<u64>,
}
/// The slice of the resolved Python `ModelConfig` the rust server reads.
#[derive(Debug, Default, serde::Deserialize)]
pub struct ModelConfig {
/// Resolved context length (`max_model_len` in `/v1/models`); mandatory at
/// boot ([`ServerArgs::validate_mandatory`]).
#[serde(default)]
pub context_len: Option<u64>,
/// Bounds client-supplied token ids — ingress 400s out-of-vocab ids before
/// they crash the scheduler's embedding lookup; mandatory at
/// boot ([`ServerArgs::validate_mandatory`]).
#[serde(default)]
pub vocab_size: Option<u64>,
/// Whether the model accepts multimodal inputs. Gates the MM Encoding branch
/// in tm-ingress; `false` silently ignores mm fields, as the Python
/// `TokenizerManager` does with `mm_processor is None`.
#[serde(default)]
pub is_multimodal: bool,
/// Resolved default sampling parameters, stamped by
/// `RustServer._build_server_args` from Python's
/// `ModelConfig.get_default_sampling_params()`. Already gated on
/// `--sampling-defaults`: holds the model's generation_config.json values
/// in "model" mode, and is empty in "openai" mode. Consumed when a chat
/// request omits `temperature`/`top_p` — the conversion must not skip
/// straight to the OpenAI terminal defaults.
#[serde(default)]
pub default_sampling_params: DefaultSamplingParams,
}
/// One `SamplingParams` field per key `get_default_sampling_params()` may emit
/// (`repetition_penalty`, `temperature`, `top_k`, `top_p`, `min_p`), filtered
/// to values the generation config actually sets — hence all `Option`.
///
/// `top_k` / `min_p` / `repetition_penalty` are parsed for parity with the
/// Python dict but not yet consumed: the Dynamo chat request type only carries
/// `temperature` and `top_p`, so the conversion resolves just those two.
#[derive(Debug, Default, serde::Deserialize)]
#[allow(dead_code)]
pub struct DefaultSamplingParams {
#[serde(default)]
pub temperature: Option<f64>,
#[serde(default)]
pub top_p: Option<f64>,
#[serde(default)]
pub top_k: Option<i64>,
#[serde(default)]
pub min_p: Option<f64>,
#[serde(default)]
pub repetition_penalty: Option<f64>,
}
fn join_host_port(host: &str, port: u16) -> String {
if host.contains(':') && !host.starts_with('[') {
format!("[{host}]:{port}") // bare IPv6 (`::`) needs brackets to bind
} else {
format!("{host}:{port}")
}
}
fn default_host() -> String {
"127.0.0.1".into()
}
fn default_port() -> u16 {
30000
}
fn default_log_level() -> String {
"info".into()
}
fn default_disaggregation_mode() -> String {
"null".into()
}
fn default_worker_num() -> usize {
1
}
impl ServerArgs {
/// Parse the blob; errors on malformed JSON or a wrongly-typed field.
pub fn from_json(s: &str) -> Result<Self, String> {
serde_json::from_str(s).map_err(|e| e.to_string())
}
/// Fail fast at startup if a field an endpoint depends on is missing.
pub fn validate_mandatory(&self) -> Result<(), String> {
if self.served_model_name.is_empty() {
return Err("no 'served_model_name' in server_args".into());
}
if self.model_config.context_len.is_none() {
return Err("no resolvable context length (model_config.context_len)".into());
}
if self.model_config.vocab_size.is_none() {
return Err("no resolvable vocab size (model_config.vocab_size)".into());
}
if !matches!(
self.disaggregation_mode.as_str(),
"null" | "prefill" | "decode"
) {
return Err(format!(
"unknown disaggregation_mode '{}' in server_args",
self.disaggregation_mode
));
}
Ok(())
}
/// True on a prefill or decode node — requests need bootstrap routing.
pub fn is_disaggregation(&self) -> bool {
self.disaggregation_mode != "null"
}
/// Serve the PD KV bootstrap registry on the api listener: every prefill
/// rust server hosts it, unconditionally — no extra topology gating. KV
/// managers and decode nodes reach the registry at the resolved
/// `disaggregation_bootstrap_port`, which rust-server mode aliases to the
/// api port, so whichever prefill server that port names is the one that
/// receives the registrations.
pub fn enable_pd_bootstrap(&self) -> bool {
self.disaggregation_mode == "prefill"
}
/// Whether the served model is multimodal, from the scheduler's dump. See
/// [`ModelConfig::is_multimodal`].
pub fn model_is_multimodal(&self) -> bool {
self.model_config.is_multimodal
}
/// Bind address `host:port`. `host` is expected to be an IP — the result is
/// parsed as a `SocketAddr`, so a bare IPv6 host gets bracketed.
pub fn bind(&self) -> String {
join_host_port(&self.host, self.port)
}
/// Whether the HTTP access log is emitted, mirroring the Python server:
/// uvicorn runs at `log_level_http or log_level` and prints access lines
/// only at info/debug. `--log-level-http warning` turns them off.
pub fn http_access_log_enabled(&self) -> bool {
let level = self
.log_level_http
.as_deref()
.filter(|s| !s.is_empty())
.unwrap_or(&self.log_level);
matches!(
level.to_ascii_lowercase().as_str(),
"trace" | "debug" | "info"
)
}
/// Pinned API threads for the embedded HTTP api-server. Python `server_args`
/// has no such field — this is derived: enough to cover the widest pool.
pub fn api_worker_num(&self) -> usize {
4.max(self.tokenizer_worker_num)
.max(self.detokenizer_worker_num)
}
}
rust/sglang-server/src/runtime/runnable.rs
-10
View File
@@ -1,10 +0,0 @@
//! The [`Runnable`] stage trait — the one contract every pipeline stage
//! (CPU-bound worker or TM router) implements to be spawned by the runtime.
/// A pipeline stage that owns its channel handles + config and runs a blocking
/// loop until its inbox closes. Lets the runtime spawn stages uniformly via
/// `threads::spawn_stage` / `threads::spawn_pool` instead of free `run_*` functions with
/// positional handles. Implemented by every CPU-bound worker and TM router.
pub trait Runnable: Send + 'static {
fn run(self);
}
rust/sglang-server/src/tokenizer_manager.rs
+7 -99
View File
@@ -1,100 +1,8 @@
//! TokenizerManager — owns the request lifecycle across two isolated threads:
//!
//! * [`ingress`] — drives the ingress FSM (Received → Validating →
//! Normalizing → {Tokenizing | PreSendValidating}) and pushes tokenized requests to the
//! scheduler ring.
//! * [`egress`] — drains the scheduler-output ring and routes each chunk to
//! the owning detokenizer shard.
//!
//! The two run on separate pinned threads with no shared state, connected to
//! the rest of the pipeline only through `flume` channels: [`TmEvent`] into the
//! ingress loop, [`Senders`] fanning out to the pools.
//! TokenizerManager
mod egress;
mod ingress;
pub use egress::{ActivityCounter, Egress};
pub use ingress::{Ingress, Limits, Mm};
use crate::ids::Rid;
use crate::message::{DetokMsg, Request};
/// Blocking receive that also wakes on shutdown: returns `None` when `rx` closes
/// *or* the `shutdown` sender is dropped.
pub fn recv<T>(rx: &flume::Receiver<T>, shutdown: &flume::Receiver<()>) -> Option<T> {
flume::Selector::new()
.recv(rx, |r| r.ok())
.recv(shutdown, |_| None)
.wait()
}
/// Events into the TokenizerManager ingress loop. API server + tokenizer pool
/// share this one inbox, keeping the loop a single consumer (no `select`).
pub enum TmEvent {
/// A freshly received request from the API server.
Ingress(Request),
/// A request back from the tokenizer pool: `PreSendValidating` (ids filled) on success,
/// or `Failed` on a tokenize error. `drive` handles both.
Tokenized(Request),
/// An MM worker finished a request parked in `Encoding`: `input_ids` are the
/// final placeholder-expanded prompt ids. The buffers ride the rid-keyed
/// sidecar (`Server.take_mm`), not this event.
MmEncoded { rid: Rid, input_ids: Vec<i32> },
/// An MM worker rejected a request parked in `Encoding` (bad media URL,
/// unsupported modality, preprocess error, …).
MmFailed { rid: Rid, message: String },
}
/// Producer-side handles, cloned into every stage that needs to emit.
/// Who asked for an abort. Both variants do the same work in
/// [`Ingress::on_abort`](crate::tokenizer_manager::ingress::Ingress) — deregister
/// the detok entry, tell the scheduler to stop — and the source is kept for
/// diagnostics.
///
/// There is no in-flight rid registry to keep consistent, and so no release
/// ordering to get wrong: [`Rid::from_client`] makes every client-supplied rid
/// internally unique, so a resubmit of the "same" rid is a different `Rid` and
/// cannot be tangled up with an abort still in flight for the original.
#[derive(Clone, Debug)]
pub enum AbortSource {
/// From an `AbortGuard` drop. Owns the release.
Guard(Rid),
/// From a detokenizer terminal path. Aborts the scheduler work.
Detok(Rid),
}
impl AbortSource {
pub fn rid(&self) -> &Rid {
match self {
Self::Guard(rid) | Self::Detok(rid) => rid,
}
}
}
#[derive(Clone)]
pub struct Senders {
/// → TokenizerManager ingress loop.
pub tm: flume::Sender<TmEvent>,
/// → the same loop, but UNBOUNDED and abort-only.
///
/// Aborts cannot share the bounded inbox. `try_send` there drops them exactly
/// when they matter most — under overload — leaving the scheduler generating
/// for a dead connection; and the caller then faces a false choice between
/// releasing the rid (a live entry can be overwritten by a resubmit) and
/// holding it (a permanent leak). An unbounded lane removes the dilemma: an
/// abort is a small `String` and is always accepted, so releases can be
/// unconditional again. It cannot grow without bound in practice — one entry
/// per in-flight request, each already bounded by the inbox that admitted it.
pub abort: flume::Sender<AbortSource>,
/// → Tokenizer pool (CPU-bound, pinned threads).
pub tok: flume::Sender<Request>,
/// → Detokenizer shards, indexed by `Rid::shard(detok.len())`.
pub detok: Vec<flume::Sender<DetokMsg>>,
}
impl Senders {
#[inline]
pub fn detok_for(&self, rid: &Rid) -> &flume::Sender<DetokMsg> {
&self.detok[rid.shard(self.detok.len())]
}
}
pub mod channel;
pub mod detokenizer;
pub mod from_scheduler;
pub mod to_scheduler;
pub mod tokenizer;
pub mod wiring;
rust/sglang-server/src/ring.rs → rust/sglang-server/src/tokenizer_manager/channel.rs
+58 -62
View File
@@ -4,47 +4,42 @@
//! share one process, so these are in-process `flume` channels — literal
//! `mpsc`/`mpmc`, no shared memory, no serialization beyond the msgpack bytes
//! the payload already is.
//!
//! GIL note: the Python side only ever calls the *non-blocking* `drain` /
//! `try_push` methods while holding the GIL, and the Rust worker threads only
//! ever push/drain raw `Bytes` — neither side touches a `PyObject` off-thread,
//! so the producer threads never need the GIL.
use std::sync::Mutex;
use std::time::Duration;
use bytes::Bytes;
use crate::message::IngressMsg;
use crate::message::request::SchedulerRequest;
/// Ingress: TokenizerManager → scheduler `recv_requests`.
/// ToSchedulerTx: TokenizerManager → scheduler `recv_requests`.
/// Producers are Rust TM workers; the single consumer is the Python thread.
/// Carries [`IngressMsg`] (columnar: scalar header + raw int64 ids cell), not a
/// Carries [`SchedulerRequest`] (columnar: scalar header + raw int64 ids cell), not a
/// single msgpack blob, so the large `input_ids` tensor bypasses msgpack.
#[derive(Clone)]
pub struct IngressProducer {
tx: flume::Sender<IngressMsg>,
pub struct ToSchedulerTx {
tx: flume::Sender<SchedulerRequest>,
}
pub struct IngressConsumer {
rx: flume::Receiver<IngressMsg>,
pub struct ToSchedulerRx {
rx: flume::Receiver<SchedulerRequest>,
/// One-slot buffer holding a message consumed by a blocking [`wait`] so the
/// scheduler can park on idle without losing it — the next [`drain`] returns
/// it first. Only ever touched by the single consumer (the Python thread),
/// so contention is nil; the `Mutex` is just for interior mutability across
/// the `&self` methods.
///
/// [`wait`]: IngressConsumer::wait
/// [`drain`]: IngressConsumer::drain
stash: Mutex<Option<IngressMsg>>,
/// [`wait`]: ToSchedulerRx::wait
/// [`drain`]: ToSchedulerRx::drain
stash: Mutex<Option<SchedulerRequest>>,
}
/// A drained ingress batch in **columnar** (struct-of-arrays) form. The `ids`
/// A drained request batch in **columnar** (struct-of-arrays) form. The `ids`
/// cells are kept *un-concatenated* so the pyo3 boundary can copy them straight
/// into one `PyBytes` (no intermediate buffer); `ids_total` is their summed
/// length, precomputed for that single allocation.
#[derive(Default)]
pub struct IngressColumns {
pub struct RequestColumns {
/// Per-request scalar msgpack header (`input_ids` omitted).
pub headers: Vec<Bytes>,
/// Per-request raw little-endian int64 ids cell (empty for control reqs).
@@ -55,26 +50,39 @@ pub struct IngressColumns {
pub ids_total: usize,
}
impl IngressProducer {
impl RequestColumns {
/// Concatenate the `ids` cells into `buf`, which must be exactly
/// `ids_total` bytes — the pyo3 boundary hands in the freshly allocated
/// `PyBytes` so the ids are copied once, straight to their destination.
pub fn copy_ids_into(&self, mut buf: &mut [u8]) {
debug_assert_eq!(buf.len(), self.ids_total);
for cell in &self.ids {
let (dst, rest) = buf.split_at_mut(cell.len());
dst.copy_from_slice(cell);
buf = rest;
}
}
}
impl ToSchedulerTx {
/// Non-blocking push. Returns `false` on a full ring (backpressure) so the
/// caller can fail the request rather than block a worker thread.
#[inline]
pub fn try_push(&self, msg: IngressMsg) -> bool {
pub fn try_push(&self, msg: SchedulerRequest) -> bool {
self.tx.try_send(msg).is_ok()
}
}
impl IngressConsumer {
/// Drain up to `max` messages into a columnar [`IngressColumns`], returning
impl ToSchedulerRx {
/// Drain up to `max` messages into a columnar [`RequestColumns`], returning
/// immediately when the ring runs dry — mirrors the scheduler's existing
/// `zmq.NOBLOCK` loop in `request_receiver._pull_raw_reqs`. Splitting headers
/// from ids here (off the GIL) leaves `recv_requests` a thin marshaling shim.
/// `zmq.NOBLOCK` loop in `request_receiver._pull_raw_reqs`.
///
/// Non-blocking by construction: `try_recv` returns `Err(TryRecvError::Empty)`
/// instantly when the ring is empty, and `Err(_) => break` exits the loop
/// right away.
pub fn drain(&self, max: usize) -> IngressColumns {
let mut batch = IngressColumns::default();
pub fn drain(&self, max: usize) -> RequestColumns {
let mut batch = RequestColumns::default();
// A message parked by a prior blocking `wait` is delivered first.
if let Some(m) = self.stash.lock().unwrap().take() {
push_msg(&mut batch, m);
@@ -110,44 +118,32 @@ impl IngressConsumer {
/// Append one drained message's columnar cells to the batch.
#[inline]
fn push_msg(batch: &mut IngressColumns, m: IngressMsg) {
fn push_msg(batch: &mut RequestColumns, m: SchedulerRequest) {
batch.ids_total += m.ids.len();
batch.lengths.push((m.ids.len() / 8) as u32); // int64 cell → tokens
batch.headers.push(m.header);
batch.ids.push(m.ids);
}
/// Egress: scheduler output (`push_chunk`) → Rust egress dispatcher.
/// The single producer is the Python thread; the consumer is the dispatcher.
/// Scheduler output (`Server.push_decode_result_batch` / `push_control_result`
/// / `push_error`) → Rust response dispatcher. The single producer is the
/// Python thread; the consumer is the dispatcher.
#[derive(Clone)]
pub struct EgressProducer {
pub struct FromSchedulerTx {
tx: flume::Sender<Bytes>,
}
pub struct EgressConsumer {
pub struct FromSchedulerRx {
rx: flume::Receiver<Bytes>,
}
impl EgressProducer {
/// Blocking push: parks until the ring has space, so a full ring applies
/// backpressure to the scheduler instead of dropping output the scheduler has
/// already committed (advanced `send_token_offset` for). The GIL is released
/// around the call, so parking here doesn't stall other Python threads.
/// `false` only when the consumer is gone (runtime shutdown), where the frame
/// is unavoidably lost.
impl FromSchedulerTx {
/// Blocking push.
pub fn push(&self, msg: Bytes) -> bool {
self.tx.send(msg).is_ok()
}
/// Non-blocking push, so the pyo3 boundary can try to hand the frame over
/// while still holding the GIL and detach only when it would actually park.
/// Releasing the GIL is not free: reacquiring it waits out the interpreter's
/// switch interval (5 ms by default), which dwarfs the sub-microsecond push
/// it was protecting.
///
/// Hands the frame BACK on a full ring (`Err(Some(msg))`) so the caller can
/// retry it under [`push`](Self::push) without rebuilding it. `Err(None)` is
/// the consumer being gone (shutdown), where the frame is unavoidably lost.
/// Non-blocking push.
#[inline]
pub fn try_push(&self, msg: Bytes) -> Result<(), Option<Bytes>> {
match self.tx.try_send(msg) {
@@ -158,37 +154,37 @@ impl EgressProducer {
}
}
impl EgressConsumer {
impl FromSchedulerRx {
/// The underlying receiver, so the dispatcher can drain it via
/// [`tokenizer_manager::recv`](crate::tokenizer_manager::recv) (data + shutdown select).
/// [`wiring::recv`](crate::tokenizer_manager::wiring::recv) (data + shutdown select).
pub fn receiver(&self) -> &flume::Receiver<Bytes> {
&self.rx
}
}
/// Build both halves of a bounded ring.
pub fn ingress_ring(cap: usize) -> (IngressProducer, IngressConsumer) {
pub fn to_scheduler(cap: usize) -> (ToSchedulerTx, ToSchedulerRx) {
let (tx, rx) = flume::bounded(cap);
(
IngressProducer { tx },
IngressConsumer {
ToSchedulerTx { tx },
ToSchedulerRx {
rx,
stash: Mutex::new(None),
},
)
}
pub fn egress_ring(cap: usize) -> (EgressProducer, EgressConsumer) {
pub fn from_scheduler(cap: usize) -> (FromSchedulerTx, FromSchedulerRx) {
let (tx, rx) = flume::bounded(cap);
(EgressProducer { tx }, EgressConsumer { rx })
(FromSchedulerTx { tx }, FromSchedulerRx { rx })
}
#[cfg(test)]
mod tests {
use super::*;
fn msg(h: &'static [u8]) -> IngressMsg {
IngressMsg {
fn msg(h: &'static [u8]) -> SchedulerRequest {
SchedulerRequest {
header: Bytes::from_static(h),
ids: Bytes::new(),
}
@@ -198,7 +194,7 @@ mod tests {
/// non-destructively, and the next `drain` returns it.
#[test]
fn wait_stashes_then_drain_returns_it() {
let (tx, rx) = ingress_ring(8);
let (tx, rx) = to_scheduler(8);
// Empty ring → times out, nothing stashed.
assert!(!rx.wait(Duration::from_millis(1)));
// Push one, then wait stashes it (returns true).
@@ -214,7 +210,7 @@ mod tests {
/// A blocked `wait` is woken the instant a producer pushes (no polling).
#[test]
fn wait_wakes_on_push() {
let (tx, rx) = ingress_ring(8);
let (tx, rx) = to_scheduler(8);
std::thread::spawn(move || {
std::thread::sleep(Duration::from_millis(20));
let _ = tx.try_push(msg(b"a"));
@@ -225,11 +221,11 @@ mod tests {
assert_eq!(rx.drain(16).headers.len(), 1);
}
/// A full egress ring parks the producer until the consumer drains — the
/// A full from_scheduler channel parks the producer until the consumer drains — the
/// committed frame is delivered in order, never dropped.
#[test]
fn egress_push_blocks_until_drained() {
let (tx, rx) = egress_ring(1);
fn response_push_blocks_until_drained() {
let (tx, rx) = from_scheduler(1);
assert!(tx.push(Bytes::from_static(b"a"))); // fits; ring now full
let t = std::thread::spawn(move || tx.push(Bytes::from_static(b"b")));
// The parked push can't have completed while the ring is full.
@@ -243,8 +239,8 @@ mod tests {
/// A closed ring (consumer gone → shutdown) returns `false` instead of
/// parking forever, so a scheduler blocked in `push` unblocks on teardown.
#[test]
fn egress_push_returns_false_when_closed() {
let (tx, rx) = egress_ring(1);
fn response_push_returns_false_when_closed() {
let (tx, rx) = from_scheduler(1);
drop(rx);
assert!(!tx.push(Bytes::from_static(b"x")));
}
rust/sglang-server/src/detokenizer.rs → rust/sglang-server/src/tokenizer_manager/detokenizer.rs
+44 -45
View File
@@ -17,22 +17,26 @@
//! `skip_tokenizer_init` is set) the backend is `Skip`: no decoding, the raw
//! `output_ids` are emitted instead of text.
//!
//! Per-chunk egress flow (no FSM state change inside Streaming):
//! Per-chunk response flow (no FSM state change inside Streaming):
//! ChunkEvent{finish:None} -> step ids -> delta -> Server frame
//! ChunkEvent{finish:Some} -> step ids -> delta -> final frame
use std::collections::HashMap;
use crate::error::Error;
use crate::fsm::{Event, RequestState};
use crate::ids::Rid;
use crate::message::DetokMsg;
use crate::message::{ChunkEvent, EgressItem, EgressSink, Matched, SinkError, TokenIds};
use crate::runtime::Runnable;
use crate::tokenizer_manager::AbortSource;
use crate::message::detok::DetokMsg;
use crate::message::finish_reason::Matched;
use crate::message::ids::Rid;
use crate::message::response::{ChunkEvent, ResponseItem, ResponseSink, SinkError};
use crate::message::types::TokenIds;
use crate::tokenizer_manager::wiring::AbortSource;
use crate::utils::runtime::Runnable;
use crate::utils::{
error::Error,
fsm::{Event, RequestState},
};
/// Default for `skip_special_tokens` (SGLang's SamplingParams default). The
/// per-request value isn't available on the egress side yet; see the note in
/// per-request value isn't available on the response yet; see the note in
/// `DetokenizerBackend::new_decoder`.
const SKIP_SPECIAL_TOKENS: bool = true;
@@ -126,7 +130,7 @@ impl DetokenizerBackend {
}
struct DetokState {
sink: EgressSink,
sink: ResponseSink,
/// `return_text_in_logprobs`: whether to decode this request's logprob token
/// ids to text (in this shard) for the `[logprob, token_id, text]` tuples.
decode_logprob_text: bool,
@@ -140,15 +144,14 @@ struct DetokState {
/// cumulative view where a consumer needs it (every unary response and the
/// cumulative SGLang `/generate` stream); OpenAI streaming forwards deltas.
decoder: Option<Box<dyn StreamDecoder>>,
/// Egress half of the lifecycle FSM. Lives here because the ingress
/// `Request` (and its FSM) was handed to the scheduler when queued; the
/// shard is the sole owner of the request's egress state, so no lock.
/// Response half of the lifecycle FSM. Lives here because the `Request` (and
/// its FSM) was handed to the scheduler when queued; the shard is the sole
/// owner of the response state, so no lock.
fsm: RequestState,
}
/// One detokenizer shard: owns a *local* `rid -> DetokState` map (single accessor,
/// no lock) and the egress backend. Spawned (pinned) per shard as a [`Runnable`];
/// a given rid is routed to exactly one shard.
/// no lock) and the detokenizer backend.
pub struct DetokenizerWorker {
shard: usize,
rx: flume::Receiver<DetokMsg>,
@@ -182,7 +185,7 @@ impl Runnable for DetokenizerWorker {
// Plain `recv`: exits when the `DetokMsg` channel closes (every `Senders`
// clone gone). On shutdown that happens once the API runtime drop cancels
// in-flight handlers (their `AbortGuard`s release the last clones) and
// tm-ingress/tm-egress exit — no shutdown signal needed here.
// to-scheduler/from-scheduler exit — no shutdown signal needed here.
while let Ok(msg) = self.rx.recv() {
match msg {
DetokMsg::Register {
@@ -203,7 +206,7 @@ impl Runnable for DetokenizerWorker {
},
);
}
// One decode step's chunks for this shard, batched by tm-egress.
// One decode step's chunks for this shard, batched by from-scheduler.
DetokMsg::Chunks(evs) => {
for ev in evs {
handle_chunk(&mut table, ev, &self.backend, &self.abort);
@@ -224,7 +227,7 @@ impl Runnable for DetokenizerWorker {
}
}
/// The `RequestKind::Detokenize` backend stage: tm-ingress queued this rid's
/// The `RequestKind::Detokenize` backend stage: to-scheduler queued this rid's
/// `Register` just before on this same channel, so the entry exists — deliver
/// the decoded text (or the error) through the registered sink and drop it,
/// like a one-result control request. No scheduler abort on failure: this kind
@@ -237,8 +240,8 @@ fn handle_decode(
) {
if let Some(mut st) = table.remove(rid) {
let item = match backend.decode_once(token_ids) {
Ok(text) => EgressItem::Data(text.into()),
Err(e) => EgressItem::Error(e),
Ok(text) => ResponseItem::Data(text.into()),
Err(e) => ResponseItem::Error(e),
};
let _ = st.sink.try_send(item);
st.fsm = RequestState::Completed;
@@ -249,8 +252,8 @@ fn handle_decode(
/// single `Done` frame — no detokenization, no streaming.
fn handle_result(table: &mut HashMap<Rid, DetokState>, rid: &Rid, payload: bytes::Bytes) {
if let Some(mut st) = table.remove(rid) {
let _ = st.sink.try_send(EgressItem::Control(payload));
// Egress FSM: a control request goes straight to Completed (no Streaming
let _ = st.sink.try_send(ResponseItem::Control(payload));
// Response FSM: a control request goes straight to Completed (no Streaming
// / Finalizing states — single response, never streamed).
st.fsm = RequestState::Completed;
}
@@ -274,7 +277,7 @@ fn handle_fail(
let _ = abort.send(AbortSource::Detok(rid.clone()));
let _ = st
.sink
.try_send(EgressItem::Error(Error::Internal(message)));
.try_send(ResponseItem::Error(Error::Internal(message)));
st.fsm = RequestState::Completed;
}
}
@@ -328,7 +331,7 @@ fn handle_chunk(
// — the other two terminal paths (disconnect, fail) both abort.
let _ = st.fsm.apply(Event::Error(e.clone()));
let _ = abort.send(AbortSource::Detok(rid.clone()));
let _ = st.sink.try_send(EgressItem::Error(e));
let _ = st.sink.try_send(ResponseItem::Error(e));
table.remove(&rid);
return;
}
@@ -365,7 +368,7 @@ fn handle_chunk(
if finished {
// The Done frame *is* the final frame: Finalizing → Completed.
let sent = st.sink.try_send(EgressItem::Done(ev)).is_ok();
let sent = st.sink.try_send(ResponseItem::Done(ev)).is_ok();
let _ = st.fsm.apply(if sent {
Event::FinalFrameSent
} else {
@@ -379,7 +382,7 @@ fn handle_chunk(
// silently dropping the frame would truncate the response and still look
// like success at EOS. So treat both as terminal: drop the request AND
// abort scheduler work for it.
if let Err(e) = st.sink.try_send(EgressItem::Frame(ev)) {
if let Err(e) = st.sink.try_send(ResponseItem::Frame(ev)) {
match e {
SinkError::Full => {
tracing::warn!(
@@ -441,15 +444,15 @@ mod tests {
#[test]
fn full_sink_drops_request_and_aborts_scheduler() {
// Capacity-1 sink, pre-filled so the next send hits `Full`.
let (tx, _rx) = mpsc::channel::<EgressItem>(1);
tx.try_send(EgressItem::Frame(ChunkEvent::default()))
let (tx, _rx) = mpsc::channel::<ResponseItem>(1);
tx.try_send(ResponseItem::Frame(ChunkEvent::default()))
.unwrap();
let mut table = HashMap::new();
table.insert(
Rid::from("1"),
DetokState {
sink: EgressSink::Local(tx),
sink: ResponseSink::Local(tx),
decode_logprob_text: false,
no_stop_trim: false,
decoder: None,
@@ -510,19 +513,15 @@ mod tests {
}
/// A `Decode` job answers through the REGISTERED sink and consumes the
/// entry — the `RequestKind::Detokenize` egress contract. Uses the `Skip`
/// backend, whose decode error must arrive as an `Error` item (not vanish):
/// dropping it leaves the submitter awaiting a reply forever. (Unlike
/// `handle_fail` there is deliberately no abort lane in the signature —
/// this kind never reached the ring, so there is no scheduler work to stop.)
/// entry.
#[test]
fn decode_answers_via_registered_sink_and_consumes_the_entry() {
let (tx, mut rx) = mpsc::channel::<EgressItem>(4);
let (tx, mut rx) = mpsc::channel::<ResponseItem>(4);
let mut table = HashMap::new();
table.insert(
Rid::from("d1"),
DetokState {
sink: EgressSink::Local(tx),
sink: ResponseSink::Local(tx),
decode_logprob_text: false,
no_stop_trim: false,
decoder: None,
@@ -537,7 +536,7 @@ mod tests {
&DetokenizerBackend::Skip,
);
let Ok(EgressItem::Error(err)) = rx.try_recv() else {
let Ok(ResponseItem::Error(err)) = rx.try_recv() else {
panic!("the decode error must reach the sink, not vanish");
};
assert!(matches!(err, Error::Validation(_)));
@@ -562,11 +561,11 @@ mod tests {
/// deterministically, without needing to find a real 64-bit collision.
#[test]
fn co_located_requests_keep_their_own_sinks() {
let (tx_a, mut rx_a) = mpsc::channel::<EgressItem>(4);
let (tx_b, mut rx_b) = mpsc::channel::<EgressItem>(4);
let (tx_a, mut rx_a) = mpsc::channel::<ResponseItem>(4);
let (tx_b, mut rx_b) = mpsc::channel::<ResponseItem>(4);
let mut table = HashMap::new();
let state = |tx| DetokState {
sink: EgressSink::Local(tx),
sink: ResponseSink::Local(tx),
decode_logprob_text: false,
no_stop_trim: false,
decoder: None,
@@ -594,8 +593,8 @@ mod tests {
&tm_tx,
);
let ids = |rx: &mut mpsc::Receiver<EgressItem>| match rx.try_recv() {
Ok(EgressItem::Frame(ev)) => ev.token_ids,
let ids = |rx: &mut mpsc::Receiver<ResponseItem>| match rx.try_recv() {
Ok(ResponseItem::Frame(ev)) => ev.token_ids,
other => panic!("expected a frame, got {other:?}"),
};
assert_eq!(
@@ -618,12 +617,12 @@ mod tests {
finish_reason: serde_json::Value,
ids: Vec<i32>,
) -> ChunkEvent {
let (tx, mut rx) = mpsc::channel::<EgressItem>(4);
let (tx, mut rx) = mpsc::channel::<ResponseItem>(4);
let mut table = HashMap::new();
table.insert(
Rid::from("1"),
DetokState {
sink: EgressSink::Local(tx),
sink: ResponseSink::Local(tx),
decode_logprob_text: false,
no_stop_trim,
decoder: None, // skip mode → output_ids passthrough
@@ -643,7 +642,7 @@ mod tests {
};
handle_chunk(&mut table, ev, &DetokenizerBackend::Skip, &tm_tx);
match rx.try_recv() {
Ok(EgressItem::Done(out)) => out,
Ok(ResponseItem::Done(out)) => out,
other => panic!("expected Done, got {other:?}"),
}
}
rust/sglang-server/src/tokenizer_manager/egress.rs → rust/sglang-server/src/tokenizer_manager/from_scheduler.rs
+35 -42
View File
@@ -1,51 +1,44 @@
//! TokenizerManager egress thread — drains the egress ring (scheduler output
//! pushed from Python) and routes each message to the detok shard that owns its
//! `Rid::shard`. Routing is a pure function of the rid, so it matches the shard
//! the request registered with on ingress — no shared map, no lock.
//!
//! The ring carries a 1-byte frame tag: `BATCH` (a whole decode batch, fanned
//! out here into per-request chunks), `RESULT` (a single control-request JSON
//! payload, e.g. `/server_info`), or `ERROR` (a terminal per-request failure the
//! scheduler ingress couldn't decode, routed back as a 400).
//! TokenizerManager dispatcher thread — drains the from_scheduler channel and
//! routes each message to the detok shard that owns its `Rid::shard`.
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use bytes::Bytes;
use crate::ids::Rid;
use crate::message::DetokMsg;
use crate::message::{
ChunkEvent, EGRESS_TAG_BATCH, EGRESS_TAG_ERROR, EGRESS_TAG_RESULT, for_each_chunk,
use crate::message::detok::DetokMsg;
use crate::message::ids::Rid;
use crate::message::response::{
ChunkEvent, DISPATCH_TAG_BATCH, DISPATCH_TAG_ERROR, DISPATCH_TAG_RESULT, for_each_chunk,
};
use crate::ring::EgressConsumer;
use crate::runtime::Runnable;
use crate::tokenizer_manager::{Senders, recv};
use crate::tokenizer_manager::channel::FromSchedulerRx;
use crate::tokenizer_manager::wiring::{Senders, recv};
/// A monotonic counter bumped once per egress-ring frame the dispatcher drains.
/// A monotonic counter bumped once per from_scheduler frame the dispatcher drains.
/// It's the rust-native equivalent of the Python `TokenizerManager`'s
/// `last_receive_tstamp`: `/health_generate` watches it advance to confirm the
/// scheduler → detok path is alive (the value itself is meaningless).
pub type ActivityCounter = Arc<AtomicU64>;
/// Egress dispatcher stage. Owns the egress-ring consumer + the detok-shard
/// Dispatcher dispatcher stage. Owns the from_scheduler consumer + the detok-shard
/// senders, so the runtime spawns it as a [`Runnable`].
pub struct Egress {
egress: EgressConsumer,
pub struct Dispatcher {
from_scheduler_rx: FromSchedulerRx,
senders: Senders,
activity: ActivityCounter,
shutdown: flume::Receiver<()>,
}
impl Egress {
impl Dispatcher {
pub fn new(
egress: EgressConsumer,
from_scheduler_rx: FromSchedulerRx,
senders: Senders,
activity: ActivityCounter,
shutdown: flume::Receiver<()>,
) -> Self {
Self {
egress,
from_scheduler_rx,
senders,
activity,
shutdown,
@@ -53,20 +46,20 @@ impl Egress {
}
}
impl Runnable for Egress {
impl Runnable for Dispatcher {
fn run(self) {
// Reused across frames (`clear` keeps capacity) — steady state allocates nothing.
let shards = self.senders.detok.len();
let shards = self.senders.detokenizer_tx.len();
let mut buckets: Vec<Vec<ChunkEvent>> = (0..shards).map(|_| Vec::new()).collect();
while let Some(bytes) = recv(self.egress.receiver(), &self.shutdown) {
while let Some(bytes) = recv(self.from_scheduler_rx.receiver(), &self.shutdown) {
let Some((&tag, body)) = bytes.split_first() else {
continue;
};
match tag {
// A whole decode batch: bucket each request by the shard owning its
// rid, then hand each shard its chunks in one send.
EGRESS_TAG_BATCH => {
DISPATCH_TAG_BATCH => {
for b in buckets.iter_mut() {
b.clear();
}
@@ -97,13 +90,13 @@ impl Runnable for Egress {
// log the same line as the recoverable one.
if decoded.rids.is_empty() {
tracing::error!(
"egress: bad batch frame named NO rids; any request in \
"from_scheduler: bad batch frame named NO rids; any request in \
it will hang (header undecodable, or empty rid column)"
);
} else {
tracing::warn!(
rids = decoded.rids.len(),
"egress: bad batch frame; failing its requests"
"from_scheduler: bad batch frame; failing its requests"
);
}
for b in buckets.iter_mut() {
@@ -113,7 +106,7 @@ impl Runnable for Egress {
// 500, not 400: the client's request was fine — the
// scheduler's own output frame was not.
let shard = rid.shard(shards);
let _ = self.senders.detok[shard].send(DetokMsg::Fail {
let _ = self.senders.detokenizer_tx[shard].send(DetokMsg::Fail {
rid,
message: "internal error: malformed scheduler output frame".into(),
});
@@ -125,36 +118,36 @@ impl Runnable for Egress {
continue;
}
let chunks = DetokMsg::Chunks(std::mem::take(b));
if self.senders.detok[i].send(chunks).is_err() {
tracing::error!("egress: detok shard closed");
if self.senders.detokenizer_tx[i].send(chunks).is_err() {
tracing::error!("from_scheduler: detok shard closed");
}
}
// Any frame off the ring = the scheduler produced output → alive.
self.activity.fetch_add(1, Ordering::Relaxed);
}
EGRESS_TAG_RESULT => {
DISPATCH_TAG_RESULT => {
if let Some((rid, msg)) = decode_result(body) {
self.route(&rid, msg);
}
}
EGRESS_TAG_ERROR => {
DISPATCH_TAG_ERROR => {
if let Some((rid, msg)) = decode_error(body) {
self.route(&rid, msg);
}
}
other => tracing::warn!(tag = other, "egress: unknown frame tag"),
other => tracing::warn!(tag = other, "from_scheduler: unknown frame tag"),
}
}
}
}
impl Egress {
impl Dispatcher {
/// Route one message to the shard owning `rid`. HOL ceiling: a slow shard stalls
/// this thread; the fix is a per-shard egress ring (see `threads::TM_CORES`).
/// this thread; the fix is a per-shard from_scheduler channel.
#[inline]
fn route(&self, rid: &Rid, msg: DetokMsg) {
if self.senders.detok_for(rid).send(msg).is_err() {
tracing::error!("egress: detok shard closed");
tracing::error!("from_scheduler: detok shard closed");
}
}
}
@@ -194,15 +187,15 @@ fn decode_error(body: &[u8]) -> Option<(Rid, DetokMsg)> {
#[cfg(test)]
mod tests {
use super::*;
use crate::message::DetokMsg;
use crate::message::frame_egress_error;
use crate::message::detok::DetokMsg;
use crate::message::response::frame_error;
/// A framed error round-trips: `frame_egress_error` → tag stripped →
/// A framed error round-trips: `frame_error` → tag stripped →
/// `decode_error` yields the rid + a `Fail` carrying the message.
#[test]
fn error_frame_roundtrips_to_fail() {
let framed = frame_egress_error("42", "invalid request: bad field");
assert_eq!(framed[0], EGRESS_TAG_ERROR);
let framed = frame_error("42", "invalid request: bad field");
assert_eq!(framed[0], DISPATCH_TAG_ERROR);
let (rid, msg) = decode_error(&framed[1..]).expect("decodes");
let want = Rid::from("42");
assert_eq!(rid, want);
rust/sglang-server/src/tokenizer_manager/ingress.rs → rust/sglang-server/src/tokenizer_manager/to_scheduler.rs
+169 -196
View File
@@ -1,48 +1,33 @@
//! TokenizerManager — ingress side.
//!
//! [`Ingress`] is a single-consumer stage draining one inbox fed by both the API
//! server (fresh requests) and the Tokenizer pool (returned requests). It owns
//! the request while driving the ingress FSM and hands it off by *moving* it to
//! the next stage; nothing here is shared, so no locks.
//!
//! Edges driven here (from the design table):
//! Received → Validating
//! Validating → Normalizing (generate: sampling-param normalize/verify)
//! Validating → PreSendValidating (control: no tokenize, no sampling params)
//! Normalizing → {Encoding | Tokenizing | PreSendValidating} (by ValidationOutcome)
//! Tokenizing → PreSendValidating (on TokenizeDone, when the request returns)
//! PreSendValidating → Queued (checks needing the tokenized length)
//! Queued → ring (handed to the scheduler)
//!
//! The egress edges (Streaming/Finalizing/Completed) are driven on the egress
//! side (see `egress` + `detokenizer`).
//! TokenizerManager — to_scheduler side.
use std::collections::HashMap;
use bytes::Bytes;
use crate::error::Error;
use crate::fsm::{Event, RequestState, ValidationOutcome};
use crate::ids::Rid;
use crate::message::{
AbortReq, ControlRequest, DetokMsg, EgressItem, GenerateRequest, IngressMsg, MmRequest,
Request, RequestKind,
use crate::message::config::ServerArgs;
use crate::message::detok::DetokMsg;
use crate::message::ids::Rid;
use crate::message::io_struct::{AbortReq, ControlRequest};
use crate::message::request::{GenerateRequest, MmRequest, Request, RequestKind, SchedulerRequest};
use crate::message::response::ResponseItem;
use crate::runtime::Runnable;
use crate::tokenizer_manager::channel::ToSchedulerTx;
use crate::tokenizer_manager::wiring::{AbortSource, Senders, TmEvent};
use crate::utils::{
error::Error,
fsm::{Event, RequestState, ValidationOutcome},
};
use crate::ring::IngressProducer;
use crate::runtime::{Runnable, ServerArgs};
use crate::tokenizer_manager::{AbortSource, Senders, TmEvent};
/// Ingress FSM dispatcher stage. Owns its inbox + downstream handles, so the
/// Intake FSM dispatcher stage. Owns its inbox + downstream handles, so the
/// runtime spawns it as a [`Runnable`] rather than calling a free `run_*` fn
/// with positional arguments.
pub struct Ingress {
rx: flume::Receiver<TmEvent>,
pub struct Intake {
tok_manager_rx: flume::Receiver<TmEvent>,
/// Unbounded abort lane (see [`Senders::abort`]). Selected against `rx` so an
/// abort is handled promptly even while the bounded inbox is saturated.
abort_rx: flume::Receiver<AbortSource>,
senders: Senders,
ingress: IngressProducer,
to_scheduler_tx: ToSchedulerTx,
limits: Limits,
mm: Mm,
/// Requests parked in `Encoding` while an MM worker processes their media;
@@ -52,7 +37,7 @@ pub struct Ingress {
shutdown: flume::Receiver<()>,
}
/// The ingress side of the MM path.
/// The intake side of the MM path.
#[derive(Clone)]
pub struct Mm {
/// Whether the model is multimodal. When false, mm fields are silently
@@ -64,32 +49,22 @@ pub struct Mm {
/// Results sidecar. Purged here when a late result arrives for a request
/// that is no longer parked; otherwise it would leak, since only the
/// scheduler drain pops entries.
pub sidecar: crate::mm::Sidecar,
pub sidecar: crate::multi_modality::sidecar::Sidecar,
}
/// Longest client-supplied rid accepted. It keys the detok table and travels on
/// every chunk, so its length is a recurring cost; Python mints 32-byte uuid hex.
const MAX_RID_LEN: usize = 128;
/// What ingress admits, resolved once at boot from the scheduler's `server_args`.
/// A struct rather than more positional `new` arguments — these grew from two to
/// six, and every one of them is a `u64`/`bool` that would be trivial to swap at
/// a call site.
///
/// NOT `Default`-able on purpose. `vocab_size` and `context_len` are mandatory,
/// and their zero value is the most restrictive setting there is — a derived
/// `Default` would silently build limits that reject every request rather than
/// failing loudly. Tests construct these explicitly (see `test_limits`).
/// Resolved once at boot from the scheduler's `server_args`.
#[derive(Clone, Debug)]
pub struct Limits {
/// Token-ids-in mode: a generate request must arrive already tokenized.
pub skip_tokenizer_init: bool,
/// `model_config.vocab_size`; bounds client-supplied token ids. Mandatory —
/// [`ServerArgs::validate_mandatory`](crate::runtime::ServerArgs) rejects a
/// boot without it, so ingress can check unconditionally.
/// `model_config.vocab_size`; bounds client-supplied token ids. A required
/// field of the `ServerArgs` schema, so intake can check unconditionally.
pub vocab_size: u64,
/// `model_config.context_len`, the ceiling for input + `max_new_tokens`.
/// Mandatory, as above.
pub context_len: u64,
/// Output slots reserved on top of the input (eagle draft tokens).
pub num_reserved_tokens: u64,
@@ -99,42 +74,34 @@ pub struct Limits {
pub enable_return_hidden_states: bool,
}
impl TryFrom<&ServerArgs> for Limits {
type Error = Error;
fn try_from(sa: &ServerArgs) -> Result<Self, Self::Error> {
Ok(Self {
impl From<&ServerArgs> for Limits {
fn from(sa: &ServerArgs) -> Self {
Self {
skip_tokenizer_init: sa.skip_tokenizer_init,
vocab_size: sa
.model_config
.vocab_size
.ok_or_else(|| Error::Validation("vocab_size missing".into()))?,
context_len: sa
.model_config
.context_len
.ok_or_else(|| Error::Validation("context_len missing".into()))?,
vocab_size: sa.model_config.vocab_size,
context_len: sa.model_config.context_len,
num_reserved_tokens: sa.num_reserved_tokens,
allow_auto_truncate: sa.allow_auto_truncate,
enable_return_hidden_states: sa.enable_return_hidden_states,
})
}
}
}
impl Ingress {
impl Intake {
pub fn new(
rx: flume::Receiver<TmEvent>,
tok_manager_rx: flume::Receiver<TmEvent>,
abort_rx: flume::Receiver<AbortSource>,
senders: Senders,
ingress: IngressProducer,
to_scheduler_tx: ToSchedulerTx,
limits: Limits,
mm: Mm,
shutdown: flume::Receiver<()>,
) -> Self {
Self {
rx,
tok_manager_rx,
abort_rx,
senders,
ingress,
to_scheduler_tx,
limits,
mm,
pending_mm: HashMap::new(),
@@ -149,20 +116,20 @@ enum Lane {
Event(TmEvent),
}
impl Runnable for Ingress {
impl Runnable for Intake {
fn run(mut self) {
loop {
// Select, not a drain-then-block: an abort arriving while the inbox is
// idle must still be handled at once.
let next = flume::Selector::new()
.recv(&self.abort_rx, |r| r.ok().map(Lane::Abort))
.recv(&self.rx, |r| r.ok().map(Lane::Event))
.recv(&self.tok_manager_rx, |r| r.ok().map(Lane::Event))
.recv(&self.shutdown, |_| None)
.wait();
match next {
Some(Lane::Abort(rid)) => self.on_abort(rid),
// A fresh request and one returning from the tokenizer pool.
Some(Lane::Event(TmEvent::Ingress(req) | TmEvent::Tokenized(req))) => {
Some(Lane::Event(TmEvent::Intake(req) | TmEvent::Tokenized(req))) => {
self.drive(req)
}
Some(Lane::Event(TmEvent::MmEncoded { rid, input_ids })) => {
@@ -186,7 +153,7 @@ impl Runnable for Ingress {
}
}
impl Ingress {
impl Intake {
/// Reject a request: → `Failed`, notify the client, deregister (unconditional
/// — a no-op when nothing was registered).
/// `registered` says whether this request ever reached `register_detok`. It
@@ -198,13 +165,13 @@ impl Ingress {
fn fail(&self, req: &mut Request, err: Error, registered: bool) {
// Log only server faults (500); 4xx/499/503 are expected and would spam.
if err.http_status() == 500 {
tracing::error!(rid = %req.rid, error = %err, "ingress rejected request");
tracing::error!(rid = %req.rid, error = %err, "intake rejected request");
}
// A rejected request never reaches the scheduler drain, so purge any
// parked MM result (no-op for the common non-mm request).
self.mm.sidecar.purge(req.rid.as_str());
let _ = req.state.apply(Event::Error(err.clone()));
let _ = req.sink.try_send(EgressItem::Error(err)); // client may be gone
let _ = req.sink.try_send(ResponseItem::Error(err)); // client may be gone
if registered {
let _ = self.senders.detok_for(&req.rid).send(DetokMsg::Deregister {
rid: req.rid.clone(),
@@ -212,7 +179,7 @@ impl Ingress {
}
}
/// Drive a request through its ingress states until it terminates (failed or
/// Drive a request through its intake states until it terminates (failed or
/// pushed to the ring), is handed to the tokenizer pool (re-entering as a
/// `Tokenized` event), or is parked in `pending_mm` awaiting an MM worker
/// (re-entering via `MmEncoded` / `MmFailed`). Each arm acts and advances
@@ -315,7 +282,7 @@ impl Ingress {
work,
};
// Full = the pool can't keep up, so back-pressure like a full
// ingress ring. Disconnected = pool gone.
// to_scheduler channel. Disconnected = pool gone.
if let Err(e) = self.mm.tx.try_send(msg) {
let err = match e {
flume::TrySendError::Full(_) => Error::QueueFull,
@@ -333,7 +300,7 @@ impl Ingress {
// `Tokenized` event (PreSendValidating, or Failed on error).
// Doesn't loop.
RequestState::Tokenizing => {
if let Err(err) = self.senders.tok.send(req) {
if let Err(err) = self.senders.tokenizer_tx.send(req) {
// Pool gone (workers exited); flume hands the request back.
let mut req = err.into_inner();
// Past `Received`, so registration happened.
@@ -377,12 +344,12 @@ impl Ingress {
self.fail(&mut req, e, registered);
return;
}
// Unreachable (egress states never reach here). Reject via `fail`/
// Unreachable (request states never reach here). Reject via `fail`/
// return (not apply + continue, which would spin on a terminal state).
other => {
self.fail(
&mut req,
Error::Internal(format!("unexpected ingress state: {other:?}")),
Error::Internal(format!("unexpected state: {other:?}")),
registered,
);
return;
@@ -391,7 +358,7 @@ impl Ingress {
}
}
/// Register the egress sink with the owning detok shard (by id) so the response
/// Register the response sink with the owning detok shard (by id) so the response
/// has a home. Carries the per-request detok flags — `return_text_in_logprobs`
/// (decode logprob text on this shard) and `no_stop_trim` (keep the matched
/// stop in the output) — so the shard needs no back-reference to the request.
@@ -443,9 +410,9 @@ impl Ingress {
}
}
/// Push a bare control request (`[tag, rid, nil]`) onto the ingress ring. The
/// Push a bare control request (`[tag, rid, nil]`) onto the to_scheduler channel. The
/// scheduler dispatches it (e.g. `GetInternalStateReq`) and replies via the
/// egress ring as a single `Result`.
/// from_scheduler channel as a single `Result`.
fn push_control_to_ring(&self, mut req: Request) {
let encode = match &req.kind {
RequestKind::Control(control) => control.encode(),
@@ -461,7 +428,7 @@ impl Ingress {
}
};
// Control requests carry no tensor cell — empty `ids`.
if !self.ingress.try_push(IngressMsg {
if !self.to_scheduler_tx.try_push(SchedulerRequest {
header,
ids: Bytes::new(),
}) {
@@ -523,13 +490,13 @@ impl Ingress {
// for, so report the miss rather than assuming the scheduler was told.
match ControlRequest::AbortReq(AbortReq::new(rid.as_str().to_string(), false)).encode() {
Ok(header) => {
if !self.ingress.try_push(IngressMsg {
if !self.to_scheduler_tx.try_push(SchedulerRequest {
header,
ids: Bytes::new(),
}) {
tracing::error!(
rid = %rid,
"abort dropped: ingress ring full; the scheduler keeps generating \
"abort dropped: to_scheduler channel is full; the scheduler keeps generating \
for this request until it finishes on its own"
);
}
@@ -539,7 +506,7 @@ impl Ingress {
}
/// Serialize the tokenized request to its `TokenizedGenerateReqInput` wire and
/// push it onto the ingress ring for the scheduler. On backpressure, fail it.
/// push it onto the to_scheduler channel for the scheduler. On backpressure, fail it.
fn push_to_ring(&self, mut req: Request) {
// Only generate requests reach here (control uses `push_control_to_ring`).
// Validate + serialize while borrowing `g` immutably; the resulting `Bytes`
@@ -561,10 +528,13 @@ impl Ingress {
}
};
if !self.ingress.try_push(IngressMsg { header, ids }) {
if !self
.to_scheduler_tx
.try_push(SchedulerRequest { header, ids })
{
self.fail(&mut req, Error::QueueFull, true); // registered
}
// On success the scheduler owns the request (egress arrives by rid); we
// On success the scheduler owns the request (response arrives by rid); we
// drop our `Request` here — the detok shard holds the sink.
}
}
@@ -722,85 +692,87 @@ fn check_total_tokens(g: &mut GenerateRequest, limits: &Limits) -> Result<(), Er
#[cfg(test)]
mod tests {
use super::*;
use crate::fsm::RequestState;
use crate::message::{EgressSink, GenerateRequest, SamplingParams};
use crate::ring::{IngressConsumer, ingress_ring};
use crate::message::request::GenerateRequest;
use crate::message::response::ResponseSink;
use crate::message::sampling::SamplingParams;
use crate::tokenizer_manager::channel::{ToSchedulerRx, to_scheduler};
use crate::utils::fsm::RequestState;
use tokio::sync::mpsc;
/// An `Ingress` plus its detok-shard receiver, ring consumer (keep alive —
/// dropping it closes the ring → false QueueFull), tm inbox sender, and the
/// An `Intake` plus its detok-shard receiver, to_scheduler channel consumer (keep alive —
/// dropping it closes the channel → false QueueFull), tm inbox sender, and the
/// mm-pool receiver (keep alive — dropping it makes mm submits fail).
fn make_ingress() -> (
Ingress,
fn make_intake() -> (
Intake,
flume::Receiver<DetokMsg>,
IngressConsumer,
ToSchedulerRx,
flume::Sender<TmEvent>,
flume::Receiver<MmRequest>,
) {
make_ingress_with(test_limits())
make_intake_with(test_limits())
}
fn make_ingress_with_abort(
fn make_intake_with_abort(
abort_rx: flume::Receiver<AbortSource>,
) -> (
Ingress,
Intake,
flume::Receiver<DetokMsg>,
IngressConsumer,
ToSchedulerRx,
flume::Sender<TmEvent>,
flume::Receiver<MmRequest>,
) {
make_ingress_inner(test_limits(), abort_rx)
make_intake_inner(test_limits(), abort_rx)
}
fn make_ingress_with(
fn make_intake_with(
limits: Limits,
) -> (
Ingress,
Intake,
flume::Receiver<DetokMsg>,
IngressConsumer,
ToSchedulerRx,
flume::Sender<TmEvent>,
flume::Receiver<MmRequest>,
) {
let (abort_tx, abort_rx) = flume::unbounded::<AbortSource>();
std::mem::forget(abort_tx); // keep the lane open; tests end by dropping tm_tx
make_ingress_inner(limits, abort_rx)
make_intake_inner(limits, abort_rx)
}
fn make_ingress_inner(
fn make_intake_inner(
limits: Limits,
abort_rx: flume::Receiver<AbortSource>,
) -> (
Ingress,
Intake,
flume::Receiver<DetokMsg>,
IngressConsumer,
ToSchedulerRx,
flume::Sender<TmEvent>,
flume::Receiver<MmRequest>,
) {
let (tok_tx, _tok_rx) = flume::unbounded();
let (detok_tx, detok_rx) = flume::unbounded();
let senders = Senders {
tm: flume::unbounded().0,
abort: flume::unbounded().0,
tok: tok_tx,
detok: vec![detok_tx],
tok_manager_tx: flume::unbounded().0,
abort_tx: flume::unbounded().0,
tokenizer_tx: tok_tx,
detokenizer_tx: vec![detok_tx],
};
let (ingress_producer, consumer) = ingress_ring(16);
let (to_scheduler_tx, consumer) = to_scheduler(16);
let (tm_tx, tm_rx) = flume::unbounded();
let (mm_tx, mm_rx) = flume::unbounded();
// Keep the shutdown sender alive (leak) so its branch never fires — tests
// end `run` by dropping `tm_tx`, not by shutdown.
let (sd_tx, sd_rx) = flume::unbounded::<()>();
std::mem::forget(sd_tx);
let ingress = Ingress::new(
let intake = Intake::new(
tm_rx,
abort_rx,
senders,
ingress_producer,
to_scheduler_tx,
limits,
test_mm(mm_tx, true),
sd_rx,
);
(ingress, detok_rx, consumer, tm_tx, mm_rx)
(intake, detok_rx, consumer, tm_tx, mm_rx)
}
/// An [`Mm`] over `tx` with a fresh sidecar.
@@ -828,25 +800,25 @@ mod tests {
AbortSource::Detok("x".into()),
] {
let (detok_tx, detok_rx) = flume::unbounded::<DetokMsg>();
let (ingress_producer, consumer) = ingress_ring(16);
let (to_scheduler_tx, consumer) = to_scheduler(16);
let (sd_tx, sd_rx) = flume::unbounded::<()>();
std::mem::forget(sd_tx);
let mut ingress = Ingress::new(
let mut intake = Intake::new(
flume::unbounded().1,
flume::unbounded().1,
Senders {
tm: flume::unbounded().0,
abort: flume::unbounded().0,
tok: flume::unbounded().0,
detok: vec![detok_tx],
tok_manager_tx: flume::unbounded().0,
abort_tx: flume::unbounded().0,
tokenizer_tx: flume::unbounded().0,
detokenizer_tx: vec![detok_tx],
},
ingress_producer,
to_scheduler_tx,
test_limits(),
test_mm(flume::unbounded().0, true),
sd_rx,
);
ingress.on_abort(source.clone());
intake.on_abort(source.clone());
assert!(
matches!(detok_rx.try_recv(), Ok(DetokMsg::Deregister { rid }) if rid.as_str() == "x"),
@@ -887,7 +859,7 @@ mod tests {
Request {
rid: id.to_string().into(),
state: RequestState::Received,
sink: EgressSink::Local(tx),
sink: ResponseSink::Local(tx),
kind: RequestKind::Generate(Box::new(GenerateRequest {
rid: id.to_string().into(),
input_ids: Some(vec![1, 2, 3]),
@@ -985,7 +957,7 @@ mod tests {
/// `max_new_tokens: null` means "no cap", NOT "skip the checks" — the input
/// alone must still fit. Gating the whole function on `max_new_tokens` let an
/// over-long prompt through to the scheduler with no ingress error at all.
/// over-long prompt through to the scheduler with no error at all.
/// Python compares with `>=`: a prompt that exactly fills the window leaves no
/// room to generate.
#[test]
@@ -1096,11 +1068,11 @@ mod tests {
/// to the ring, after registration — so it must be deregistered, not leaked.
#[test]
fn over_context_request_deregisters_and_never_reaches_the_ring() {
let (mut ingress, detok_rx, consumer, _tm_tx, _mm_rx) = make_ingress_with(Limits {
let (mut intake, detok_rx, consumer, _tm_tx, _mm_rx) = make_intake_with(Limits {
context_len: 4,
..test_limits()
});
ingress.drive(generate_req(
intake.drive(generate_req(
33,
SamplingParams {
max_new_tokens: Some(64),
@@ -1129,12 +1101,12 @@ mod tests {
/// pins. Nothing may reach the scheduler ring.
#[test]
fn detokenize_flows_register_then_decode_and_skips_the_ring() {
let (mut ingress, detok_rx, consumer, _tm_tx, _mm_rx) = make_ingress();
let (mut intake, detok_rx, consumer, _tm_tx, _mm_rx) = make_intake();
let (tx, mut rx) = mpsc::channel(8);
ingress.drive(Request {
intake.drive(Request {
rid: "41".into(),
state: RequestState::Received,
sink: EgressSink::Local(tx),
sink: ResponseSink::Local(tx),
kind: RequestKind::Detokenize {
token_ids: vec![7, 8, 9],
},
@@ -1155,7 +1127,10 @@ mod tests {
consumer.drain(16).headers.is_empty(),
"must never reach the scheduler"
);
assert!(rx.try_recv().is_err(), "no egress until the shard answers");
assert!(
rx.try_recv().is_err(),
"no response until the shard answers"
);
}
/// Negative ids cannot decode (the shard's domain is `&[u32]`): rejected by
@@ -1164,17 +1139,17 @@ mod tests {
/// leak and no decode job to drop).
#[test]
fn detokenize_negative_ids_reject_before_registration() {
let (mut ingress, detok_rx, consumer, _tm_tx, _mm_rx) = make_ingress();
let (mut intake, detok_rx, consumer, _tm_tx, _mm_rx) = make_intake();
let (tx, mut rx) = mpsc::channel(8);
ingress.drive(Request {
intake.drive(Request {
rid: "43".into(),
state: RequestState::Received,
sink: EgressSink::Local(tx),
sink: ResponseSink::Local(tx),
kind: RequestKind::Detokenize {
token_ids: vec![1, -1],
},
});
let Ok(EgressItem::Error(err)) = rx.try_recv() else {
let Ok(ResponseItem::Error(err)) = rx.try_recv() else {
panic!("sink must receive the validation error");
};
assert_eq!(err.http_status(), 400);
@@ -1198,16 +1173,16 @@ mod tests {
let (detok_tx, detok_rx) = flume::unbounded();
let (abort_tx, abort_rx) = flume::unbounded::<AbortSource>();
let senders = Senders {
tm: flume::unbounded().0,
abort: abort_tx,
tok: tok_tx,
detok: vec![detok_tx],
tok_manager_tx: flume::unbounded().0,
abort_tx,
tokenizer_tx: tok_tx,
detokenizer_tx: vec![detok_tx],
};
let (producer, _consumer) = ingress_ring(1);
let (producer, _consumer) = to_scheduler(1);
let (_tm_tx, tm_rx) = flume::unbounded();
let (sd_tx, sd_rx) = flume::unbounded::<()>();
std::mem::forget(sd_tx);
let mut ingress = Ingress::new(
let mut intake = Intake::new(
tm_rx,
abort_rx,
senders,
@@ -1217,8 +1192,8 @@ mod tests {
sd_rx,
);
ingress.on_abort(AbortSource::Guard("pushed".into()));
ingress.on_abort(AbortSource::Guard("dropped".into()));
intake.on_abort(AbortSource::Guard("pushed".into()));
intake.on_abort(AbortSource::Guard("dropped".into()));
// Both deregisters land regardless of whether the ring accepted the push.
for expected in ["pushed", "dropped"] {
@@ -1252,12 +1227,12 @@ mod tests {
#[test]
fn pre_registration_failure_does_not_deregister() {
// Rejected inside `validate` (out-of-vocab id), which runs before registration.
let (mut ingress, detok_rx, _consumer, _tm_tx, _mm_rx) = make_ingress();
let (mut intake, detok_rx, _consumer, _tm_tx, _mm_rx) = make_intake();
let mut req = generate_req(41, SamplingParams::default());
if let RequestKind::Generate(g) = &mut req.kind {
g.input_ids = Some(vec![2_000_000_000]);
}
ingress.drive(req);
intake.drive(req);
assert!(
detok_rx.try_recv().is_err(),
"a pre-registration reject must send NOTHING to the shard — a Deregister \
@@ -1265,8 +1240,8 @@ mod tests {
);
// A post-registration reject still deregisters (the leak fix stays fixed).
let (mut ingress, detok_rx, _consumer, _tm_tx, _mm_rx) = make_ingress();
ingress.drive(generate_req(
let (mut intake, detok_rx, _consumer, _tm_tx, _mm_rx) = make_intake();
intake.drive(generate_req(
42,
SamplingParams {
top_p: 2.0, // rejected by `normalize`, after registration
@@ -1284,13 +1259,13 @@ mod tests {
/// sees `Register` then `Deregister`. Regression for RSS growth on bad input.
#[test]
fn rejected_request_deregisters_from_shard() {
let (mut ingress, detok_rx, _consumer, _tm_tx, _mm_rx) = make_ingress();
let (mut intake, detok_rx, _consumer, _tm_tx, _mm_rx) = make_intake();
// top_p = 2.0 is outside (0, 1], so `SamplingParams::normalize` rejects it.
let bad = SamplingParams {
top_p: 2.0,
..Default::default()
};
ingress.drive(generate_req(7, bad));
intake.drive(generate_req(7, bad));
assert!(
matches!(detok_rx.try_recv(), Ok(DetokMsg::Register { rid, .. }) if rid.as_str() == "7"),
@@ -1307,16 +1282,16 @@ mod tests {
}
/// Regression: an out-of-vocabulary client token id must be rejected at
/// ingress with a 400 — passed through, it reaches the embedding lookup
/// and kills the scheduler process (`make_ingress` bounds vocab at 1000).
/// with a 400 — passed through, it reaches the embedding lookup
/// and kills the scheduler process (`make_intake` bounds vocab at 1000).
#[test]
fn out_of_vocab_input_ids_rejected() {
let (mut ingress, detok_rx, _consumer, _tm_tx, _mm_rx) = make_ingress();
let (mut intake, detok_rx, _consumer, _tm_tx, _mm_rx) = make_intake();
let mut req = generate_req(21, SamplingParams::default());
if let RequestKind::Generate(g) = &mut req.kind {
g.input_ids = Some(vec![1, 2_000_000_000]);
}
ingress.drive(req);
intake.drive(req);
// Rejected before registration: the only shard message is nothing at
// all, or a Deregister if registration happened first — never a push.
match detok_rx.try_recv() {
@@ -1329,23 +1304,23 @@ mod tests {
/// Same guard for negative ids and for `token_ids_logprob` entries.
#[test]
fn negative_and_logprob_token_ids_rejected() {
let (mut ingress, detok_rx, _consumer, _tm_tx, _mm_rx) = make_ingress();
let (mut intake, detok_rx, _consumer, _tm_tx, _mm_rx) = make_intake();
let mut req = generate_req(22, SamplingParams::default());
if let RequestKind::Generate(g) = &mut req.kind {
g.input_ids = Some(vec![-1]);
}
ingress.drive(req);
intake.drive(req);
match detok_rx.try_recv() {
Err(_) | Ok(DetokMsg::Deregister { .. }) => {}
Ok(_) => panic!("negative token id must not be admitted"),
}
let (mut ingress, detok_rx, _consumer, _tm_tx, _mm_rx) = make_ingress();
let (mut intake, detok_rx, _consumer, _tm_tx, _mm_rx) = make_intake();
let mut req = generate_req(23, SamplingParams::default());
if let RequestKind::Generate(g) = &mut req.kind {
g.token_ids_logprob = Some(vec![999_999]);
}
ingress.drive(req);
intake.drive(req);
match detok_rx.try_recv() {
Err(_) | Ok(DetokMsg::Deregister { .. }) => {}
Ok(_) => panic!("out-of-vocab token_ids_logprob must not be admitted"),
@@ -1355,9 +1330,9 @@ mod tests {
/// A valid request is registered and handed onward — never deregistered.
#[test]
fn admitted_request_keeps_registration() {
let (mut ingress, detok_rx, _consumer, _tm_tx, _mm_rx) = make_ingress();
let (mut intake, detok_rx, _consumer, _tm_tx, _mm_rx) = make_intake();
// Empty map → all sampling defaults, passes normalization.
ingress.drive(generate_req(9, SamplingParams::default()));
intake.drive(generate_req(9, SamplingParams::default()));
assert!(
matches!(detok_rx.try_recv(), Ok(DetokMsg::Register { rid, .. }) if rid.as_str() == "9"),
@@ -1372,8 +1347,8 @@ mod tests {
/// A pool return in `Failed` state (failed encode) is rejected via the same
/// path and deregistered, not leaked.
#[test]
fn tokenize_failure_deregisters_via_ingress() {
let (ingress, detok_rx, _consumer, tm_tx, _mm_rx) = make_ingress();
fn tokenize_failure_deregisters_via_intake() {
let (intake, detok_rx, _consumer, tm_tx, _mm_rx) = make_intake();
// The pool marks a failed encode as `Failed(err)` before returning it.
let mut req = generate_req(11, SamplingParams::default());
let _ = req
@@ -1382,7 +1357,7 @@ mod tests {
tm_tx.send(TmEvent::Tokenized(req)).unwrap();
// Close the inbox so the run loop returns after draining the one event.
drop(tm_tx);
ingress.run();
intake.run();
assert!(
matches!(detok_rx.try_recv(), Ok(DetokMsg::Deregister { rid }) if rid.as_str() == "11"),
@@ -1397,11 +1372,11 @@ mod tests {
fn abort_deregisters_from_shard() {
// Aborts arrive on their own unbounded lane now, not the request inbox.
let (abort_tx, abort_rx) = flume::unbounded::<AbortSource>();
let (ingress, detok_rx, _consumer, tm_tx, _mm_rx) = make_ingress_with_abort(abort_rx);
let (intake, detok_rx, _consumer, tm_tx, _mm_rx) = make_intake_with_abort(abort_rx);
abort_tx.send(AbortSource::Guard("rid-13".into())).unwrap();
drop(abort_tx);
drop(tm_tx);
ingress.run();
intake.run();
assert!(
matches!(detok_rx.try_recv(), Ok(DetokMsg::Deregister { rid }) if rid.as_str() == "rid-13"),
@@ -1414,7 +1389,7 @@ mod tests {
/// rejected; its registration is untouched.
#[test]
fn tokenized_return_pushes_without_deregister() {
let (ingress, detok_rx, _consumer, tm_tx, _mm_rx) = make_ingress();
let (intake, detok_rx, _consumer, tm_tx, _mm_rx) = make_intake();
let mut req = generate_req(15, SamplingParams::default());
// Simulate a successful pool return: ids filled, PreSendValidating.
if let RequestKind::Generate(g) = &mut req.kind {
@@ -1423,7 +1398,7 @@ mod tests {
req.state = RequestState::PreSendValidating;
tm_tx.send(TmEvent::Tokenized(req)).unwrap();
drop(tm_tx);
ingress.run();
intake.run();
// Pushed to the ring; the shard sees nothing.
assert!(
@@ -1436,14 +1411,12 @@ mod tests {
/// deregistered, not silently dropped.
#[test]
fn tokenize_pool_gone_deregisters() {
// `make_ingress` drops the tok receiver, so `tok.send` fails.
let (mut ingress, detok_rx, _consumer, _tm_tx, _mm_rx) = make_ingress();
// No ids → NeedsTokenize → Tokenizing branch.
let (mut intake, detok_rx, _consumer, _tm_tx, _mm_rx) = make_intake();
let mut req = generate_req(21, SamplingParams::default());
if let RequestKind::Generate(g) = &mut req.kind {
g.input_ids = None;
}
ingress.drive(req);
intake.drive(req);
assert!(
matches!(detok_rx.try_recv(), Ok(DetokMsg::Register { rid, .. }) if rid.as_str() == "21"),
@@ -1463,11 +1436,11 @@ mod tests {
Request {
rid: rid.to_string().into(),
state: RequestState::Received,
sink: EgressSink::Local(tx),
sink: ResponseSink::Local(tx),
kind: RequestKind::Generate(Box::new(GenerateRequest {
rid: rid.to_string().into(),
text: Some("<image> hi".into()),
mm: Some(Box::new(crate::message::MmData {
mm: Some(Box::new(crate::message::request::MmData {
image_data: Some(rmpv::Value::from("data:image/jpeg;base64,xxxx")),
..Default::default()
})),
@@ -1481,15 +1454,15 @@ mod tests {
/// sidecar entry is purged — no scheduler work runs for a dead client.
#[test]
fn abort_cancels_parked_mm_request() {
let (mut ingress, _detok_rx, consumer, _tm_tx, mm_rx) = make_ingress();
ingress.drive(mm_generate_req("mm-gone"));
let (mut intake, _detok_rx, consumer, _tm_tx, mm_rx) = make_intake();
intake.drive(mm_generate_req("mm-gone"));
mm_rx.try_recv().expect("parked to mm pool");
// The worker parks its result, as it always does before MmEncoded.
ingress.mm.sidecar.park(
intake.mm.sidecar.park(
"mm-gone".into(),
crate::mm::MmSidecarEntry {
features: crate::mm::FeatureStore::Inline(vec![]),
crate::multi_modality::sidecar::MmSidecarEntry {
features: crate::multi_modality::sidecar::FeatureStore::Inline(vec![]),
grids: vec![],
hashes: vec![],
offsets: vec![],
@@ -1497,16 +1470,16 @@ mod tests {
mrope_delta: 0,
},
);
ingress.on_abort(AbortSource::Guard("mm-gone".to_string().into()));
intake.on_abort(AbortSource::Guard("mm-gone".to_string().into()));
assert_eq!(consumer.drain(16).headers.len(), 1, "only the AbortReq");
// The late result must be dropped, not queued, and the sidecar purged.
ingress.on_mm_encoded("mm-gone".to_string().into(), vec![5, 6]);
intake.on_mm_encoded("mm-gone".to_string().into(), vec![5, 6]);
assert!(
consumer.drain(16).headers.is_empty(),
"cancelled, not queued"
);
assert!(ingress.mm.sidecar.take("mm-gone").is_none(), "entry purged");
assert!(intake.mm.sidecar.take("mm-gone").is_none(), "entry purged");
}
/// A multimodal request parks in `Encoding` (submitted to the mm worker
@@ -1514,8 +1487,8 @@ mod tests {
/// it → ring.
#[test]
fn mm_request_parks_then_mm_encoded_pushes_to_ring() {
let (mut ingress, _detok_rx, consumer, _tm_tx, mm_rx) = make_ingress();
ingress.drive(mm_generate_req("mm-1"));
let (mut intake, _detok_rx, consumer, _tm_tx, mm_rx) = make_intake();
intake.drive(mm_generate_req("mm-1"));
// Submitted to the mm pool with the typed work item; nothing on the ring yet.
let sub = mm_rx.try_recv().expect("mm pool must receive the request");
@@ -1529,7 +1502,7 @@ mod tests {
assert!(consumer.drain(16).headers.is_empty(), "parked, not queued");
// The worker returns the final expanded ids → pushed to the ring.
ingress.on_mm_encoded("mm-1".to_string().into(), vec![5, 6, 7, 8]);
intake.on_mm_encoded("mm-1".to_string().into(), vec![5, 6, 7, 8]);
let batch = consumer.drain(16);
assert_eq!(batch.headers.len(), 1);
assert_eq!(
@@ -1542,14 +1515,14 @@ mod tests {
/// A worker failure rejects the parked request (deregister, no ring push).
#[test]
fn mm_failure_rejects_parked_request() {
let (mut ingress, detok_rx, consumer, _tm_tx, _mm_rx) = make_ingress();
ingress.drive(mm_generate_req("mm-2"));
let (mut intake, detok_rx, consumer, _tm_tx, _mm_rx) = make_intake();
intake.drive(mm_generate_req("mm-2"));
assert!(
matches!(detok_rx.try_recv(), Ok(DetokMsg::Register { .. })),
"registered before parking",
);
ingress.on_mm_failed("mm-2".to_string().into(), "bad image".into());
intake.on_mm_failed("mm-2".to_string().into(), "bad image".into());
assert!(
matches!(detok_rx.try_recv(), Ok(DetokMsg::Deregister { rid })
if rid.as_str() == "mm-2"),
@@ -1566,29 +1539,29 @@ mod tests {
let (tok_tx, tok_rx) = flume::unbounded();
let (detok_tx, _detok_rx) = flume::unbounded();
let senders = Senders {
tm: flume::unbounded().0,
abort: flume::unbounded().0,
tok: tok_tx,
detok: vec![detok_tx],
tok_manager_tx: flume::unbounded().0,
abort_tx: flume::unbounded().0,
tokenizer_tx: tok_tx,
detokenizer_tx: vec![detok_tx],
};
let (ingress_producer, _consumer) = ingress_ring(16);
let (to_scheduler_tx, _consumer) = to_scheduler(16);
let (_tm_tx, tm_rx) = flume::unbounded();
let (mm_tx, mm_rx) = flume::unbounded();
let (abort_tx, abort_rx) = flume::unbounded::<AbortSource>();
std::mem::forget(abort_tx);
let (sd_tx, sd_rx) = flume::unbounded::<()>();
std::mem::forget(sd_tx);
let mut ingress = Ingress::new(
let mut intake = Intake::new(
tm_rx,
abort_rx,
senders,
ingress_producer,
to_scheduler_tx,
test_limits(),
test_mm(mm_tx, false),
sd_rx,
);
ingress.drive(mm_generate_req("mm-3"));
intake.drive(mm_generate_req("mm-3"));
assert!(
mm_rx.try_recv().is_err(),
"mm disabled: nothing submitted to the mm channel",
@@ -1603,9 +1576,9 @@ mod tests {
/// panicking (e.g. hash-collision overwrite) — regression guard.
#[test]
fn late_mm_result_is_dropped() {
let (mut ingress, _detok_rx, consumer, _tm_tx, _mm_rx) = make_ingress();
ingress.on_mm_encoded("ghost".to_string().into(), vec![1]);
ingress.on_mm_failed("ghost".to_string().into(), "boom".into());
let (mut intake, _detok_rx, consumer, _tm_tx, _mm_rx) = make_intake();
intake.on_mm_encoded("ghost".to_string().into(), vec![1]);
intake.on_mm_failed("ghost".to_string().into(), "boom".into());
assert!(consumer.drain(16).headers.is_empty());
}
}
rust/sglang-server/src/tokenizer.rs → rust/sglang-server/src/tokenizer_manager/tokenizer.rs
+10 -8
View File
@@ -14,11 +14,11 @@
use std::path::Path;
use std::sync::Arc;
use crate::error::Error;
use crate::fsm::Event;
use crate::message::{Request, RequestKind, TokenIds};
use crate::message::request::{Request, RequestKind};
use crate::message::types::TokenIds;
use crate::runtime::Runnable;
use crate::tokenizer_manager::TmEvent;
use crate::tokenizer_manager::wiring::TmEvent;
use crate::utils::{error::Error, fsm::Event};
/// Pluggable text→token-ids backend. `Send + Sync` so one instance is shared
/// (read-only) across all pinned workers.
@@ -230,8 +230,10 @@ impl Runnable for TokenizerWorker {
#[cfg(test)]
mod tests {
use super::*;
use crate::fsm::RequestState;
use crate::message::{EgressSink, GenerateRequest, RequestKind, SamplingParams};
use crate::message::request::{GenerateRequest, RequestKind};
use crate::message::response::ResponseSink;
use crate::message::sampling::SamplingParams;
use crate::utils::fsm::RequestState;
use tokio::sync::mpsc;
/// One token per whitespace-separated word, so a stop's token count differs
@@ -266,7 +268,7 @@ mod tests {
.send(Request {
rid: "1".into(),
state: RequestState::Tokenizing,
sink: EgressSink::Local(sink_tx),
sink: ResponseSink::Local(sink_tx),
kind: RequestKind::Generate(Box::new(GenerateRequest {
rid: "1".into(),
text: Some("hello world".into()),
@@ -326,7 +328,7 @@ mod tests {
.send(Request {
rid: "1".into(),
state: RequestState::Tokenizing,
sink: EgressSink::Local(tokio::sync::mpsc::channel(4).0),
sink: ResponseSink::Local(tokio::sync::mpsc::channel(4).0),
kind: RequestKind::Generate(Box::new(GenerateRequest {
rid: "1".into(),
text: Some("hi".into()),
rust/sglang-server/src/tokenizer_manager/wiring.rs
+77
View File
@@ -0,0 +1,77 @@
//! The flume fabric between stages: the request-loop inbox ([`TmEvent`]), the
//! abort lane ([`AbortSource`]), the producer-side handles ([`Senders`]), and
//! the shutdown-aware [`recv`].
use crate::message::detok::DetokMsg;
use crate::message::ids::Rid;
use crate::message::request::Request;
/// Blocking receive that also wakes on shutdown: returns `None` when `rx` closes
/// *or* the `shutdown` sender is dropped.
pub fn recv<T>(rx: &flume::Receiver<T>, shutdown: &flume::Receiver<()>) -> Option<T> {
flume::Selector::new()
.recv(rx, |r| r.ok())
.recv(shutdown, |_| None)
.wait()
}
/// Events into the TokenizerManager request loop. API server + tokenizer pool
/// share this one inbox, keeping the loop a single consumer (no `select`).
pub enum TmEvent {
/// A freshly received request from the API server.
Intake(Request),
/// A request back from the tokenizer pool: `PreSendValidating` (ids filled)
/// on success, or `Failed` on a tokenize error. `drive` handles both.
Tokenized(Request),
/// An MM worker finished a request parked in `Encoding`: `input_ids` are the
/// final placeholder-expanded prompt ids. The buffers ride the rid-keyed
/// sidecar (`Server.take_mm`), not this event.
MmEncoded { rid: Rid, input_ids: Vec<i32> },
/// An MM worker rejected a request parked in `Encoding` (bad media URL,
/// unsupported modality, preprocess error, …).
MmFailed { rid: Rid, message: String },
}
/// The source of the abort request. Both variants do the same work in
/// [`Intake::on_abort`] — deregister the detok entry, tell the scheduler to
/// stop — and the source is kept for diagnostics.
///
/// There is no in-flight rid registry to keep consistent, and so no release
/// ordering to get wrong: [`Rid::from_client`] makes every client-supplied rid
/// internally unique, so a resubmit of the "same" rid is a different `Rid` and
/// cannot be tangled up with an abort still in flight for the original.
#[derive(Clone, Debug)]
pub enum AbortSource {
/// From an `AbortGuard` drop. Owns the release.
Guard(Rid),
/// From a detokenizer terminal path. Aborts the scheduler work.
Detok(Rid),
}
impl AbortSource {
pub fn rid(&self) -> &Rid {
match self {
Self::Guard(rid) | Self::Detok(rid) => rid,
}
}
}
/// Producer-side handles, cloned into every stage that needs to emit.
#[derive(Clone)]
pub struct Senders {
/// → TokenizerManager loop.
pub tok_manager_tx: flume::Sender<TmEvent>,
/// → the same loop, but UNBOUNDED and abort-only.
pub abort_tx: flume::Sender<AbortSource>,
/// → Tokenizer pool (CPU-bound, pinned threads).
pub tokenizer_tx: flume::Sender<Request>,
/// → Detokenizer shards, indexed by `Rid::shard(detok.len())`.
pub detokenizer_tx: Vec<flume::Sender<DetokMsg>>,
}
impl Senders {
#[inline]
pub fn detok_for(&self, rid: &Rid) -> &flume::Sender<DetokMsg> {
&self.detokenizer_tx[rid.shard(self.detokenizer_tx.len())]
}
}
rust/sglang-server/src/utils.rs
+6
View File
@@ -1,6 +1,12 @@
//! Shared helpers with no home in a pipeline stage.
pub mod environ;
pub mod error;
pub mod fsm;
pub mod logging;
pub mod regex;
pub mod response;
pub mod runtime;
pub mod serialize;
pub mod sock;
pub mod threads;
rust/sglang-server/src/environ.rs → rust/sglang-server/src/utils/environ.rs
View File
rust/sglang-server/src/error.rs → rust/sglang-server/src/utils/error.rs
+2 -2
View File
@@ -20,8 +20,8 @@ pub enum Error {
#[error("detokenize failed: {0}")]
Detokenize(String),
/// Ingress ring full / scheduler not draining. Surfaced as backpressure.
#[error("ingress queue full")]
/// To-scheduler channel full. Surfaced as backpressure.
#[error("to_scheduler channel full")]
QueueFull,
/// Client went away mid-stream. Drives `Aborted`, not `Failed`.
rust/sglang-server/src/fsm.rs → rust/sglang-server/src/utils/fsm.rs
+8 -8
View File
@@ -12,7 +12,7 @@
//! Aborted
//! ```
use crate::error::Error;
use super::error::Error;
#[derive(Debug, Clone)]
pub enum RequestState {
@@ -36,7 +36,7 @@ pub enum RequestState {
Aborted,
}
/// Outcome of validation, selecting the ingress branch.
/// Outcome of validation.
#[derive(Debug, Clone, Copy)]
pub enum ValidationOutcome {
/// Has multimodal inputs → Encoding, where an MM worker runs the native
@@ -52,7 +52,7 @@ pub enum ValidationOutcome {
/// design's transition table.
#[derive(Debug)]
pub enum Event {
// --- ingress ---
// --- request ---
Validated(ValidationOutcome),
NeedsNormalize,
EncodeDone,
@@ -60,7 +60,7 @@ pub enum Event {
/// The pre-send checks passed; the request may be pushed to the ring.
PreSendValidated,
SchedulerPicked,
// --- egress ---
// --- response ---
Chunk {
finish: bool,
},
@@ -113,7 +113,7 @@ impl RequestState {
}
let next = match (&*self, &event) {
// ingress
// request
(Received, Validated(_)) => Validating,
// Generate requests pass through Normalizing (sampling-param
// normalize/verify); control requests skip it, having none.
@@ -127,12 +127,12 @@ impl RequestState {
// pre-send checks: expanded image tokens count against the same
// input + max_new_tokens ceiling as tokenized text.
(Encoding, EncodeDone) => PreSendValidating,
// Every ingress branch funnels through the pre-send checks, so they
// Every to-scheduler branch funnels through the pre-send checks, so they
// run exactly once per request no matter how it got its ids.
(Tokenizing, TokenizeDone) => PreSendValidating,
(PreSendValidating, PreSendValidated) => Queued,
(Queued, SchedulerPicked) => Streaming { chunks_sent: 0 },
// egress
// response
(Streaming { chunks_sent }, Chunk { finish: false }) => Streaming {
chunks_sent: chunks_sent + 1,
},
@@ -154,7 +154,7 @@ mod tests {
state
}
/// Every ingress branch — control, client-supplied ids, and text through the
/// Every to-scheduler branch — control, client-supplied ids, and text through the
/// tokenizer pool — must land in `PreSendValidating`, because that is where
/// the checks needing the final `input_ids` run. A branch that reached
/// `Queued` directly would skip them silently.
rust/sglang-server/src/utils/logging.rs
+23
View File
@@ -0,0 +1,23 @@
//! Process-wide `tracing` setup for the embedded server.
use std::sync::OnceLock;
use tracing_appender::non_blocking::WorkerGuard;
use tracing_subscriber::EnvFilter;
/// Keeps the non-blocking log writer's background thread alive for the process
/// lifetime (dropping the guard would stop log delivery).
static LOG_GUARD: OnceLock<WorkerGuard> = OnceLock::new();
/// Install the global `tracing` subscriber once; a no-op if the host process
/// (or an earlier call) already set one.
pub fn init_tracing() {
let (writer, guard) = tracing_appender::non_blocking(std::io::stdout());
let _ = LOG_GUARD.set(guard);
let _ = tracing_subscriber::fmt()
.with_env_filter(
EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info")),
)
.with_writer(writer)
.try_init();
}
rust/sglang-server/src/utils/regex.rs
+4 -29
View File
@@ -8,7 +8,7 @@
use std::collections::HashMap;
use std::sync::{LazyLock, Mutex};
use crate::error::Error;
use super::error::Error;
/// `MAX_LEN` from Python's `get_max_seq_length`: the bound for an *unbounded* stop
/// regex (`\d+`, `.*`, …) or one we can't statically size — the scheduler then
@@ -221,7 +221,7 @@ const ADMISSION_CACHE_CAP: usize = 512;
/// it is HIR translation, which expands `\w`/`\W` into large Unicode class unions.
/// A 256-byte `\W`-heavy pattern (exactly [`MAX_STOP_REGEX_LEN`]) measures 574 µs,
/// and a request may carry [`MAX_STOP_REGEX_COUNT`] of them — 18 ms of admission on
/// the single ingress thread, re-derived from scratch on every request. It
/// the single to-scheduler thread, re-derived from scratch on every request. It
/// multiplies through a batch, because one `sampling_params` object broadcasts to
/// every item: a 13.6 KB body measured **1.01 s**, during which that thread serves
/// no other request, no abort and no health probe.
@@ -234,7 +234,7 @@ const ADMISSION_CACHE_CAP: usize = 512;
/// Cleared wholesale when full rather than evicted one at a time: that is what
/// CPython's `re` does, and it keeps the hot path one lookup with no LRU
/// bookkeeping. The lock is held across a hash lookup and nothing else, and is
/// taken almost exclusively by the one ingress thread.
/// taken almost exclusively by the one to-scheduler thread.
static ADMISSION_CACHE: LazyLock<Mutex<HashMap<Box<str>, usize>>> =
LazyLock::new(|| Mutex::new(HashMap::new()));
@@ -322,24 +322,6 @@ impl<'a> RegexPattern<'a> {
/// Validate a `stop_regex` before it can reach the scheduler, returning the parsed
/// AST so the caller can derive its bound without parsing again.
///
/// Two independent classes of rejection, for two different reasons:
///
/// * **Dialect.** CPython's `re` is the engine that actually runs this pattern, and
/// `regex-syntax` is neither a superset nor a subset of it. The rows where
/// `regex-syntax` is *wider* are the dangerous ones — a pattern admitted here but
/// uncompilable there reaches `re.search` on the decode hot path, where the raised
/// error is uncaught and takes the scheduler down. Hence the invariant is
/// one-directional: **anything admitted here must compile in Python**, while
/// rejecting a pattern Python would have accepted costs one client a 400. The
/// asymmetry is deliberate, and it is why the checks below only ever add
/// rejections.
///
/// * **Cost.** The scheduler re-matches this against the output tail on *every*
/// decode step, inside `re`'s C loop with the GIL held, where no timeout or signal
/// can interrupt it. Repetition counts, nested unbounded repetitions, quantified
/// assertions and ambiguous alternations are refused on those grounds alone —
/// they are all valid Python.
fn validate(pattern: &str) -> Result<regex_syntax::ast::Ast, Error> {
reject_python_incompatible(pattern)?;
let ast = regex_syntax::ast::parse::ParserBuilder::new()
@@ -815,14 +797,7 @@ mod tests {
// ---- AMBIGUITY (rounds 6-8). Every one compiles cleanly on both sides
// and raises nothing, so the `except (re.error, RecursionError)` seatbelt
// in `_check_str_based_finish` is irrelevant: the match simply never
// returns, inside GIL-holding CPython C that no watchdog can preempt.
//
// Two distinct kill modes, both represented:
// * unbounded bound -> `_stop_match_tail_len` hands `re.search` the
// WHOLE accumulated output, so cost grows every decode step;
// * finite bound -> a fixed but ruinous cost paid EVERY step forever,
// and `MAX_STOP_REGEX_COUNT` allows 64 patterns per request.
// Timings on a matching subject; see the module docs for the method.
// returns.
case(
"(?:.|.)*Z",
Policy::MustReject,
rust/sglang-server/src/runtime.rs → rust/sglang-server/src/utils/runtime.rs
+123 -171
View File
@@ -1,14 +1,14 @@
//! Runtime bootstrap: wires channels, pins CPU-bound pools, starts the tokio
//! API server, and returns a handle the Python boundary uses for
//! `recv_requests` (ingress drain) and `push_batch` (egress push).
//! `recv_requests` and `push_decode_result_batch`.
//!
//! Thread layout:
//! * API server — tokio multi-thread runtime (I/O bound), pinned core set A
//! * Tokenizer — N pinned OS threads (CPU bound), core set B
//! * Detokenizer — M pinned OS threads / shards (CPU bound), core set C
//! * TM ingress — 1 thread driving the ingress FSM
//! * TM egress — 1 thread draining the egress ring → detok shards
//! * MM workers — K unpinned OS threads, spawned late via
//! * API server — tokio multi-thread runtime (I/O bound), pinned core set A
//! * Tokenizer — N pinned OS threads (CPU bound), core set B
//! * Detokenizer — M pinned OS threads (CPU bound), core set C
//! * To_scheduler — 1 thread driving the FSM
//! * From_scheduler — 1 thread draining the scheduler → detok shards
//! * MM workers — K unpinned OS threads, spawned late via
//! [`Runtime::spawn_mm_pool`] (multimodal models only)
//!
//! Keeping CPU-bound tokenize/detokenize off the async executor avoids stalling
@@ -17,41 +17,42 @@
use std::sync::{Arc, Mutex};
use std::thread::JoinHandle;
mod config;
mod runnable;
mod threads;
use crate::message::config::RuntimeConfig;
use crate::message::detok::DetokMsg;
pub use config::{DefaultSamplingParams, RuntimeConfig, RustServerServerArgs, ServerArgs};
use crate::message::DetokMsg;
use crate::ring::{
EgressConsumer, EgressProducer, IngressConsumer, IngressProducer, egress_ring, ingress_ring,
use super::threads::{join_all_with_timeout, plan_cores, spawn_pool};
use crate::tokenizer_manager::channel::{
FromSchedulerRx, FromSchedulerTx, ToSchedulerRx, ToSchedulerTx, from_scheduler, to_scheduler,
};
use crate::runtime::threads::{plan_cores, spawn_pool};
use crate::tokenizer_manager::{Senders, TmEvent};
use crate::tokenizer_manager::wiring::{Senders, TmEvent};
use crate::utils::sock::bind_tcp_listener;
use crate::{api_server, detokenizer, tokenizer, tokenizer_manager};
use crate::{
api_server, tokenizer_manager, tokenizer_manager::detokenizer, tokenizer_manager::tokenizer,
};
// Re-export so stages keep importing `crate::runtime::Runnable`.
pub use runnable::Runnable;
/// A pipeline stage that owns its channel handles + config and runs a blocking
/// loop until its inbox closes.
pub trait Runnable: Send + 'static {
fn run(self);
}
/// Live runtime. Held by the pyo3 bridge; the Python boundary reads `ingress`
/// and `egress`. `request_shutdown` (also run on `Drop`) stops every stage.
/// Live runtime. Held by the pyo3 bridge; the Python boundary reads the `to_scheduler_rx` channel,
/// and write to `from_scheduler_tx` channel. `request_shutdown` (also run on `Drop`) stops every stage.
pub struct Runtime {
pub ingress: IngressConsumer,
pub egress: EgressProducer,
pub to_scheduler_rx: ToSchedulerRx,
pub from_scheduler_tx: FromSchedulerTx,
/// Requests parked in `Encoding`, drained by the MM worker pool
/// (`Server.start_mm_workers`). Stays empty for non-multimodal models —
/// ingress never routes to it.
pub mm: flume::Receiver<crate::message::MmRequest>,
/// Back-channel for the MM workers' `MmEncoded` / `MmFailed` into tm-ingress.
pub tm: flume::Sender<TmEvent>,
/// request never routes to it.
pub to_mm_worker_rx: flume::Receiver<crate::message::request::MmRequest>,
/// Back-channel for the MM workers' `MmEncoded` / `MmFailed` into to_scheduler.
pub from_mm_worker_tx: flume::Sender<TmEvent>,
/// The loaded tokenizer, shared with the MM worker path (`None` under
/// `skip_tokenizer_init`).
pub tokenizer: Option<Arc<dyn tokenizer::TextTokenizer>>,
/// MM results parked between a worker's `MmEncoded` and the scheduler drain
/// (`Server.take_mm`).
pub mm_sidecar: crate::mm::Sidecar,
pub mm_sidecar: crate::multi_modality::sidecar::Sidecar,
/// Worker join handles, joined by `request_shutdown` / `Drop`.
threads: Mutex<Vec<JoinHandle<()>>>,
/// The single shutdown sender.
@@ -71,40 +72,23 @@ impl Runtime {
/// MM preprocessing floats over that whole set (rather than owning cores
/// that idle between bursts) and never preempts the scheduler's reserved
/// cores.
pub fn spawn_mm_pool(&self, workers: usize, ctx: Arc<crate::mm::Context>) {
pub fn spawn_mm_pool(&self, workers: usize, ctx: Arc<crate::multi_modality::worker::Context>) {
let mut threads = self.threads.lock().unwrap();
spawn_pool("mm-worker", None, workers.max(1), &mut threads, |_| {
crate::mm::MmWorker::new(self.mm.clone(), self.tm.clone(), ctx.clone())
crate::multi_modality::worker::MmWorker::new(
self.to_mm_worker_rx.clone(),
self.from_mm_worker_tx.clone(),
ctx.clone(),
)
});
}
/// Stop the runtime and join every worker thread (with a bounded wait).
///
/// Dropping `shutdown_tx` wakes the tm-ingress/tm-egress selectors (which
/// otherwise never see their inbox close — one self-holds a `tm` sender, the
/// other's inbox is the Python-fed egress ring). Those exit and drop their
/// `Senders` clones; the api thread's `serve` returns non-gracefully, so its
/// `block_on` unwinds and the api tokio runtime is dropped — cancelling
/// in-flight handlers, whose `AbortGuard`s release the remaining clones. With
/// every clone gone the tok/detok channels close and those workers exit.
///
/// In-flight requests are **aborted**, not drained — this is the hard-stop
/// path (also run on `Drop`). Clients of aborted requests retry.
pub fn request_shutdown(&self) {
drop(self.shutdown_tx.lock().unwrap().take());
let handles: Vec<JoinHandle<()>> = self.threads.lock().unwrap().drain(..).collect();
if handles.is_empty() {
return; // Idempotent: a `Drop` after an explicit shutdown has nothing to join.
}
// Join off-thread and wait with a deadline: a stuck worker can't wedge exit.
let (done_tx, done_rx) = flume::bounded::<()>(1);
std::thread::spawn(move || {
for h in handles {
let _ = h.join();
}
let _ = done_tx.send(());
});
if done_rx.recv_timeout(SHUTDOWN_JOIN_TIMEOUT).is_err() {
// Idempotent: a `Drop` after an explicit shutdown finds nothing to join.
let handles = std::mem::take(&mut *self.threads.lock().unwrap());
if !join_all_with_timeout(handles, SHUTDOWN_JOIN_TIMEOUT) {
tracing::warn!(
"shutdown: workers did not exit within {SHUTDOWN_JOIN_TIMEOUT:?}; abandoning join"
);
@@ -118,55 +102,55 @@ impl Drop for Runtime {
}
}
/// Boot the whole frontend. Returns once threads are spawned (non-blocking),
/// so the Python caller regains control of the GIL immediately. `Err` on a
/// startup misconfiguration (e.g. no tokenizer for a non-skip server).
/// Boot the whole frontend. Returns once threads are spawned (non-blocking).
/// `Err` on a startup misconfiguration (e.g. no tokenizer for a non-skip server).
pub fn start(cfg: RuntimeConfig) -> Result<Runtime, String> {
let (shutdown_tx, shutdown_rx) = flume::unbounded::<()>();
let mut threads = Vec::new();
let plan = plan_cores(&cfg);
// --- rings (Rust ↔ Python) ---
let (ingress_tx, ingress_rx): (IngressProducer, IngressConsumer) =
ingress_ring(cfg.rust_server_args.ingress_ring_cap);
let (egress_tx, egress_rx): (EgressProducer, EgressConsumer) =
egress_ring(cfg.rust_server_args.egress_ring_cap);
let (to_scheduler_tx, to_scheduler_rx): (ToSchedulerTx, ToSchedulerRx) =
to_scheduler(cfg.rust_server_args.to_scheduler_cap);
let (from_scheduler_tx, from_scheduler_rx): (FromSchedulerTx, FromSchedulerRx) =
from_scheduler(cfg.rust_server_args.from_scheduler_cap);
// --- inter-stage channels ---
let (tm_tx, tm_rx) = flume::bounded::<TmEvent>(cfg.rust_server_args.channel_cap);
let (tok_tx, tok_rx) =
flume::bounded::<crate::message::Request>(cfg.rust_server_args.channel_cap);
let (tok_manager_tx, tok_manager_rx) =
flume::bounded::<TmEvent>(cfg.rust_server_args.channel_cap);
let (tokenizer_tx, tokenizer_rx) =
flume::bounded::<crate::message::request::Request>(cfg.rust_server_args.channel_cap);
// Encoding → MM worker pool. Bounded like the other stage edges so a slow
// pool back-pressures instead of buffering unboundedly.
let (mm_tx, mm_rx) =
flume::bounded::<crate::message::MmRequest>(cfg.rust_server_args.channel_cap);
let (mm_worker_tx, mm_worker_rx) =
flume::bounded::<crate::message::request::MmRequest>(cfg.rust_server_args.channel_cap);
let detokenizer_worker_num = cfg.server_args.detokenizer_worker_num;
let mut detok_tx = Vec::with_capacity(detokenizer_worker_num);
let mut detok_rx = Vec::with_capacity(detokenizer_worker_num);
let mut detokenizer_tx = Vec::with_capacity(detokenizer_worker_num);
let mut detokenizer_rx = Vec::with_capacity(detokenizer_worker_num);
for _ in 0..detokenizer_worker_num {
let (tx, rx) = flume::bounded::<DetokMsg>(cfg.rust_server_args.channel_cap);
detok_tx.push(tx);
detok_rx.push(rx);
detokenizer_tx.push(tx);
detokenizer_rx.push(rx);
}
// Aborts get their own UNBOUNDED lane: on the bounded inbox they are dropped
// exactly under the overload that makes them necessary (see `Senders::abort`).
let (abort_tx, abort_rx) = flume::unbounded::<crate::tokenizer_manager::AbortSource>();
let (abort_tx, abort_rx) = flume::unbounded::<crate::tokenizer_manager::wiring::AbortSource>();
let senders = Senders {
tm: tm_tx.clone(),
abort: abort_tx.clone(),
tok: tok_tx,
detok: detok_tx,
tok_manager_tx: tok_manager_tx.clone(),
abort_tx: abort_tx.clone(),
tokenizer_tx,
detokenizer_tx,
};
// `skip_tokenizer_init`: clients send token ids and receive token ids — no
// tokenizer is loaded, and the egress emits raw `output_ids` (no decode).
// tokenizer is loaded, and the server emits raw `output_ids` (no decode).
let skip_tokenizer_init = cfg.server_args.skip_tokenizer_init;
// The same instance is shared by the tokenizer pool (encode) and the detok
// shards (decode); `None` only under `skip_tokenizer_init`.
let dyn_tokenizer = tokenizer::load_tokenizer(
// Empty only in minimal standalone blobs (the Python dump always
// Empty only in standalone (test) configs (the Python handoff always
// resolves it); empty → no tokenizer, allowed only under
// `skip_tokenizer_init`.
(!cfg.server_args.tokenizer_path.is_empty()).then_some(&*cfg.server_args.tokenizer_path),
@@ -179,8 +163,8 @@ pub fn start(cfg: RuntimeConfig) -> Result<Runtime, String> {
.as_ref()
.map(|t| Arc::new(tokenizer::DynamoTokenizer::new(t.clone())) as _);
// Shared: MM workers park, the Python drain pops, tm-ingress purges.
let mm_sidecar: crate::mm::Sidecar = Default::default();
// Shared: MM workers park, the Python drain pops.
let mm_sidecar: crate::multi_modality::sidecar::Sidecar = Default::default();
// --- Detokenizer shards (pinned, CPU bound) ---
{
@@ -194,12 +178,12 @@ pub fn start(cfg: RuntimeConfig) -> Result<Runtime, String> {
let detok_cores = plan.as_ref().map(|p| p.detok.clone());
// Each shard owns its receiver outright (one consumer per shard), so the
// owned `detok_rx` Vec is moved out element-by-element via the iterator.
let count = detok_rx.len();
let mut rxs = detok_rx.into_iter();
let count = detokenizer_rx.len();
let mut detokenizer_rxs = detokenizer_rx.into_iter();
spawn_pool("detokenizer", detok_cores, count, &mut threads, |i| {
detokenizer::DetokenizerWorker::new(
i,
rxs.next().unwrap(),
detokenizer_rxs.next().unwrap(),
backend.clone(),
abort_tx.clone(),
)
@@ -208,7 +192,7 @@ pub fn start(cfg: RuntimeConfig) -> Result<Runtime, String> {
// --- Tokenizer pool (pinned, CPU bound) ---
// Only spawned when a real tokenizer is loaded; under `skip_tokenizer_init`
// there is none and ingress never routes to the pool, so we skip it.
// there is none and request never routes to the pool, so we skip it.
if let Some(tokenizer) = &text_tokenizer {
// Reuse the single loaded tokenizer (shared with the detok shards).
let tokenizer = tokenizer.clone();
@@ -220,29 +204,35 @@ pub fn start(cfg: RuntimeConfig) -> Result<Runtime, String> {
tok_cores,
cfg.server_args.tokenizer_worker_num,
&mut threads,
|_i| tokenizer::TokenizerWorker::new(tok_rx.clone(), tm_tx.clone(), tokenizer.clone()),
|_i| {
tokenizer::TokenizerWorker::new(
tokenizer_rx.clone(),
tok_manager_tx.clone(),
tokenizer.clone(),
)
},
);
}
// Egress heartbeat: bumped per drained frame, watched by `/health_generate`.
let egress_activity: tokenizer_manager::ActivityCounter =
// Response heartbeat: bumped per drained frame, watched by `/health_generate`.
let response_activity: tokenizer_manager::from_scheduler::ActivityCounter =
Arc::new(std::sync::atomic::AtomicU64::new(0));
// --- Egress dispatcher: drains egress ring → routes chunks to shards ---
// --- Response dispatcher: drains from_scheduler channel → routes chunks to shards ---
{
// First TM core; egress is the hotter router (every output token). One
// First TM core; from_scheduler is the hotter router (every output token). One
// worker today via `spawn_pool`, so sharding by `Rid::shard` later (see
// `TM_CORES`) is just a larger count + per-shard receivers.
let cores = plan
.as_ref()
.and_then(|p| p.tm.first().copied())
.map(|c| vec![c]);
let mut egress_rx = Some(egress_rx); // moved into the single worker
let activity = egress_activity.clone();
let mut from_scheduler_rx = Some(from_scheduler_rx); // moved into the single worker
let activity = response_activity.clone();
let shutdown_rx = shutdown_rx.clone();
spawn_pool("tm-egress", cores, 1, &mut threads, |_| {
tokenizer_manager::Egress::new(
egress_rx.take().unwrap(),
spawn_pool("from-scheduler", cores, 1, &mut threads, |_| {
tokenizer_manager::from_scheduler::Dispatcher::new(
from_scheduler_rx.take().unwrap(),
senders.clone(),
activity.clone(),
shutdown_rx.clone(),
@@ -250,7 +240,7 @@ pub fn start(cfg: RuntimeConfig) -> Result<Runtime, String> {
});
}
// --- TokenizerManager ingress loop ---
// --- TokenizerManager to_scheduler loop ---
{
// Second TM core when present, else share the first (1-core / API-set
// fallback) — still off the CPU-bound pool cores either way.
@@ -258,22 +248,21 @@ pub fn start(cfg: RuntimeConfig) -> Result<Runtime, String> {
.as_ref()
.and_then(|p| p.tm.get(1).or_else(|| p.tm.first()).copied())
.map(|c| vec![c]);
let limits = tokenizer_manager::Limits::try_from(&*cfg.server_args)
.map_err(|e| format!("ingress limits: {e}"))?;
let mm = tokenizer_manager::Mm {
let limits = tokenizer_manager::to_scheduler::Limits::from(&*cfg.server_args);
let mm = tokenizer_manager::to_scheduler::Mm {
enabled: cfg.server_args.model_is_multimodal(),
tx: mm_tx,
tx: mm_worker_tx,
sidecar: mm_sidecar.clone(),
};
let mut parts = Some((tm_rx, ingress_tx)); // moved into the single worker
let mut parts = Some((tok_manager_rx, to_scheduler_tx)); // moved into the single worker
let shutdown_rx = shutdown_rx.clone();
spawn_pool("tm-ingress", cores, 1, &mut threads, |_| {
let (tm_rx, ingress_tx) = parts.take().unwrap();
tokenizer_manager::Ingress::new(
tm_rx,
spawn_pool("to-scheduler", cores, 1, &mut threads, |_| {
let (tok_manager_rx, to_scheduler_tx) = parts.take().unwrap();
tokenizer_manager::to_scheduler::Intake::new(
tok_manager_rx,
abort_rx.clone(),
senders.clone(),
ingress_tx,
to_scheduler_tx,
limits.clone(),
mm.clone(),
shutdown_rx.clone(),
@@ -286,7 +275,7 @@ pub fn start(cfg: RuntimeConfig) -> Result<Runtime, String> {
let cfg = cfg.clone();
let api_cores = plan.as_ref().map(|p| p.api.clone());
let senders = senders.clone();
let api_activity = egress_activity.clone();
let response_activity = response_activity.clone();
let shutdown_rx = shutdown_rx.clone();
// Bind synchronously so an unavailable port (EADDRINUSE) is a hard
// startup error. The `?` drops `shutdown_tx`/`senders`, which stops the
@@ -299,7 +288,7 @@ pub fn start(cfg: RuntimeConfig) -> Result<Runtime, String> {
.spawn(move || {
let mut builder = tokio::runtime::Builder::new_multi_thread();
builder
.worker_threads(cfg.rust_server_args.api_worker_num)
.worker_threads(cfg.rust_server_args.http_api_worker_num)
.enable_all();
if let Some(cores) = api_cores {
let next = std::sync::atomic::AtomicUsize::new(0);
@@ -311,13 +300,13 @@ pub fn start(cfg: RuntimeConfig) -> Result<Runtime, String> {
});
}
let rt = builder.build().expect("build api runtime");
rt.block_on(api_server::serve(
rt.block_on(api_server::app::serve(
listener,
senders,
cfg.rust_server_args.channel_cap,
cfg.server_args.clone(),
// Egress heartbeat watched by `/health_generate`.
api_activity,
// Response heartbeat watched by `/health_generate`.
response_activity,
shutdown_rx,
))
})
@@ -326,10 +315,10 @@ pub fn start(cfg: RuntimeConfig) -> Result<Runtime, String> {
}
Ok(Runtime {
ingress: ingress_rx,
egress: egress_tx,
mm: mm_rx,
tm: tm_tx,
to_scheduler_rx,
from_scheduler_tx,
to_mm_worker_rx: mm_worker_rx,
from_mm_worker_tx: tok_manager_tx,
tokenizer: text_tokenizer,
mm_sidecar,
threads: Mutex::new(threads),
@@ -340,16 +329,16 @@ pub fn start(cfg: RuntimeConfig) -> Result<Runtime, String> {
#[cfg(test)]
mod tests {
use super::*;
use crate::message::config::{RuntimeConfig, RustServerServerArgs, ServerArgs};
/// Minimal boot args. `skip_tokenizer_init` avoids loading a tokenizer/detok
/// model; `model_config` carries the two fields `Limits::from_server_args`
/// requires. They are mandatory at boot, so a fixture without them panics the
/// runtime instead of exercising what these tests are about — `start` does not
/// run `ServerArgs::validate_mandatory` itself, `Server::start` does.
const TEST_SERVER_ARGS: &str = r#"{
"skip_tokenizer_init": true,
"model_config": {"context_len": 2048, "vocab_size": 1000}
}"#;
/// Minimal boot config: no tokenizer load, complete `model_config` (from
/// `Default`), unified role.
fn test_server_args() -> ServerArgs {
ServerArgs {
skip_tokenizer_init: true,
..Default::default()
}
}
/// Regression: `request_shutdown` must actually stop the API server — it joins
/// the api thread once the listener closes, so the port stops accepting.
@@ -362,11 +351,11 @@ mod tests {
drop(probe);
// `skip_tokenizer_init` → no tokenizer/detok model load; minimal boot.
let server_args = ServerArgs::from_json(TEST_SERVER_ARGS).unwrap();
let server_args = test_server_args();
let cfg = RuntimeConfig {
rust_server_args: RustServerServerArgs {
http_addr: addr,
api_worker_num: 1,
http_api_worker_num: 1,
..Default::default()
},
server_args: Arc::new(server_args),
@@ -387,12 +376,8 @@ mod tests {
);
}
/// Regression: shutdown must return promptly even with an in-flight `/generate`.
/// No scheduler drains the ingress ring or feeds the egress ring here, so the
/// handler parks on its egress channel forever. Graceful shutdown would wait
/// for it (deadlock → only the 5s bounded-join fallback returns); the
/// non-graceful path cancels the handler via the api runtime drop, whose
/// `AbortGuard` releases the last `Senders` clone so the workers exit.
/// Regression: shutdown must return promptly even with an in-flight
/// `/generate`.
#[test]
fn shutdown_returns_with_in_flight_request() {
use std::io::Write;
@@ -402,11 +387,11 @@ mod tests {
let addr = probe.local_addr().unwrap();
drop(probe);
let server_args = ServerArgs::from_json(TEST_SERVER_ARGS).unwrap();
let server_args = test_server_args();
let cfg = RuntimeConfig {
rust_server_args: RustServerServerArgs {
http_addr: addr,
api_worker_num: 1,
http_api_worker_num: 1,
..Default::default()
},
server_args: Arc::new(server_args),
@@ -414,7 +399,7 @@ mod tests {
let rt = start(cfg).expect("start runtime");
// Fire a request that will block (already-tokenized → valid → pushed to the
// ring, then the handler awaits egress frames that never arrive).
// ring, then the handler awaits decode frames that never arrive).
let mut conn = std::net::TcpStream::connect(addr).expect("connect");
let body = r#"{"input_ids":[1,2,3],"stream":false,"sampling_params":{"max_new_tokens":8}}"#;
let req = format!(
@@ -447,11 +432,11 @@ mod tests {
let addr = probe.local_addr().unwrap();
drop(probe);
let server_args = ServerArgs::from_json(TEST_SERVER_ARGS).unwrap();
let server_args = test_server_args();
let cfg = RuntimeConfig {
rust_server_args: RustServerServerArgs {
http_addr: addr,
api_worker_num: 1,
http_api_worker_num: 1,
..Default::default()
},
server_args: Arc::new(server_args),
@@ -508,11 +493,11 @@ mod tests {
let hog = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
let addr = hog.local_addr().unwrap();
let server_args = ServerArgs::from_json(TEST_SERVER_ARGS).unwrap();
let server_args = test_server_args();
let cfg = RuntimeConfig {
rust_server_args: RustServerServerArgs {
http_addr: addr,
api_worker_num: 1,
http_api_worker_num: 1,
..Default::default()
},
server_args: Arc::new(server_args),
@@ -523,37 +508,4 @@ mod tests {
};
assert!(err.contains("bind"), "error should mention bind: {err}");
}
/// `server_args` missing a mandatory `model_config` field must be a startup
/// ERROR, not a panic.
///
/// `Limits::try_from` is fallible and the ingress loop is built inside a
/// `spawn_pool` closure, so resolving it there would put the failure on a
/// freshly spawned worker thread — a thread `start` never inspects. The boot
/// would report success and the server would accept connections with no
/// ingress loop behind them, hanging every request instead of refusing to
/// start. Only `Server::start` runs `validate_mandatory`, so `start` cannot
/// assume these fields are present.
#[test]
fn start_fails_when_model_config_is_incomplete() {
let probe = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
let addr = probe.local_addr().unwrap();
drop(probe);
// Boots fine in every other respect — only `model_config` is absent.
let server_args = ServerArgs::from_json(r#"{"skip_tokenizer_init": true}"#).unwrap();
let cfg = RuntimeConfig {
rust_server_args: RustServerServerArgs {
http_addr: addr,
api_worker_num: 1,
..Default::default()
},
server_args: Arc::new(server_args),
};
let err = match start(cfg) {
Ok(_) => panic!("an incomplete model_config must not boot, got Ok"),
Err(e) => e,
};
assert!(err.contains("ingress limits"), "{err}");
}
}
rust/sglang-server/src/runtime/threads.rs → rust/sglang-server/src/utils/threads.rs
+24 -16
View File
@@ -8,29 +8,22 @@
//! 3. one [`spawn_pool`] (N pinned workers) or [`spawn_stage`] (singleton) call.
use std::thread::JoinHandle;
use std::time::Duration;
use core_affinity::CoreId;
use super::{Runnable, RuntimeConfig};
use super::runtime::Runnable;
use crate::message::config::RuntimeConfig;
/// Cores reserved for the two TokenizerManager router threads (`tm-ingress`,
/// `tm-egress`) — light, latency-sensitive channel routers, so one core each.
/// Cores reserved for the two TokenizerManager router threads (`to-scheduler`,
/// `from-scheduler`) — light, latency-sensitive channel routers, so one core each.
///
/// TODO(tm-scaling): both TM threads are single-consumer serialization points,
/// each with its own ceiling. `tm-ingress` runs validate + `normalize_sampling_params`
/// each with its own ceiling. `to-scheduler` runs validate + `normalize_sampling_params`
/// for *every* request before fanning out to the (pooled) tokenizer workers, so a
/// high request-arrival / short-request workload is bounded by that one thread's
/// per-request cost (kept O(fields), see `sampling::normalize_sampling_params`).
/// Sharding ingress by rid — like the tokenizer/detok pools — lifts that ceiling.
///
/// `tm-egress` is a head-of-line ceiling of a different kind — it
/// does a *blocking* send per chunk to the owning detok shard, so one slow shard
/// stalls the dispatcher and thus every shard (see `Egress::route`). Sharding the
/// dispatcher alone doesn't fix it: each egress-ring frame is a whole batch fanned
/// to *all* shards, so any dispatcher still blocks on the slow one. The real fix
/// is a per-shard egress ring (the scheduler pushing each request's output to its
/// shard's ring), each drained by its own dispatcher — at which point this needs
/// one core per ingress/egress shard rather than a fixed 2.
/// Sharding to-scheduler by rid — like the tokenizer/detok pools — lifts that ceiling.
const TM_CORES: usize = 2;
/// Partition the machine's cores into four disjoint sets: the I/O-bound API
@@ -54,7 +47,7 @@ pub(super) fn plan_cores(cfg: &RuntimeConfig) -> Option<CorePlan> {
_ => return None,
};
if cores.len()
< cfg.rust_server_args.api_worker_num
< cfg.rust_server_args.http_api_worker_num
+ cfg.server_args.tokenizer_worker_num
+ cfg.server_args.detokenizer_worker_num
{
@@ -67,7 +60,7 @@ pub(super) fn plan_cores(cfg: &RuntimeConfig) -> Option<CorePlan> {
let mut it = cores.into_iter();
let api: Vec<CoreId> = it
.by_ref()
.take(cfg.rust_server_args.api_worker_num)
.take(cfg.rust_server_args.http_api_worker_num)
.collect();
let tok = it
.by_ref()
@@ -140,3 +133,18 @@ pub(super) fn spawn_pool<R, F>(
spawn_stage(&format!("{name}-{i}"), core, build(i), threads);
}
}
/// Join every handle, giving up after `timeout`.
pub(super) fn join_all_with_timeout(handles: Vec<JoinHandle<()>>, timeout: Duration) -> bool {
if handles.is_empty() {
return true;
}
let (done_tx, done_rx) = std::sync::mpsc::channel::<()>();
std::thread::spawn(move || {
for h in handles {
let _ = h.join();
}
let _ = done_tx.send(());
});
done_rx.recv_timeout(timeout).is_ok()
}
test/registered/rust/test_run_rust_tests.py
+58
View File
@@ -0,0 +1,58 @@
"""Run the `rust/` Cargo workspace's unit tests from the CPU CI suite."""
import shutil
import subprocess
import unittest
from pathlib import Path
from sglang.srt.environ import envs
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
BUILD_AND_RUN_TIMEOUT_S = 900
RUST_WORKSPACE = Path(__file__).resolve().parents[3] / "rust"
register_cpu_ci(est_time=900, suite="base-a-test-cpu")
# Exported by _pr-test-stage-cpu.yml as the negation of the check-changes
# rust_workspace paths filter; it defaults to false, so only a CI run that
# positively detected no rust/ changes skips the cargo build.
@unittest.skipIf(
envs.SGLANG_SKIP_RUST_TESTS.get(),
"SGLANG_SKIP_RUST_TESTS is set (no rust/ workspace changes per CI check-changes)",
)
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()
test/registered/unit/multimodal/rust/qwen/test_e2e_parity.py
+1 -1
View File
@@ -82,7 +82,7 @@ class TestQwenE2eParity(CustomTestCase):
ids, features, grids, hashes, offsets, mrope, delta = DRIVER(
PROMPT_PER_IMAGE * len(sources), sources, spec.rust_json()
)
# The shape of Rust's MmHandoff, inline transport (test_build_native_mm
# The shape of Rust's MmEncodeResult, inline transport (test_build_native_mm
# pins the shm shape).
handoff = SimpleNamespace(
features=features,
test/registered/unit/multimodal/rust/shared/test_build_native_mm.py
+1 -1
View File
@@ -52,7 +52,7 @@ class TestBuildNativeMm(CustomTestCase):
features = np.arange(30, dtype=np.float32)
output = NativeMmHost.build_native_mm(
self.spec,
SimpleNamespace( # the shape of Rust's MmHandoff
SimpleNamespace( # the shape of Rust's MmEncodeResult
grids=self.GRIDS,
hashes=self.HASHES,
offsets=self.OFFSETS,