From 7d7ab4b5c6314fe5ba6bb55a8ffc5538f9ed238b Mon Sep 17 00:00:00 2001 From: Rain Jiang <96632942+rainj-me@users.noreply.github.com> Date: Fri, 21 Aug 2026 16:37:02 -0700 Subject: [PATCH] Rainj me/rust server refactor2 (#35239) --- .github/workflows/_pr-test-check-changes.yml | 13 + .github/workflows/_pr-test-stage-cpu.yml | 5 +- .github/workflows/lint.yml | 17 - python/sglang/srt/environ.py | 3 + python/sglang/srt/managers/rust_server.py | 152 ++++- .../scheduler_components/idle_sleeper.py | 4 +- python/sglang/srt/managers/utils.py | 7 +- rust/sglang-mm/src/registry.rs | 39 +- rust/sglang-server/src/api_server.rs | 96 +-- rust/sglang-server/src/api_server/app.rs | 104 +++ rust/sglang-server/src/api_server/common.rs | 58 +- .../api_server/disaggregation/bootstrap.rs | 41 +- rust/sglang-server/src/api_server/frame.rs | 4 +- rust/sglang-server/src/api_server/guard.rs | 16 +- rust/sglang-server/src/api_server/log.rs | 2 +- .../src/api_server/native_api.rs | 110 +-- rust/sglang-server/src/api_server/openai.rs | 35 +- .../src/api_server/openai/chat.rs | 44 +- .../src/api_server/openai/completions.rs | 45 +- .../src/api_server/openai/models.rs | 7 +- .../src/api_server/openai/template.rs | 2 +- .../src/api_server/openai/test_utils.rs | 66 +- .../src/api_server/openai/tools.rs | 6 +- rust/sglang-server/src/api_server/prefetch.rs | 8 +- rust/sglang-server/src/api_server/submit.rs | 33 +- rust/sglang-server/src/lib.rs | 286 ++++---- rust/sglang-server/src/message.rs | 105 +-- rust/sglang-server/src/message/config.rs | 638 ++++++++++++++++++ rust/sglang-server/src/message/detok.rs | 42 ++ .../src/message/finish_reason.rs | 6 +- rust/sglang-server/src/{ => message}/ids.rs | 14 +- rust/sglang-server/src/message/io_struct.rs | 6 +- rust/sglang-server/src/message/request.rs | 62 +- .../src/message/{egress.rs => response.rs} | 112 ++- rust/sglang-server/src/message/sampling.rs | 30 +- rust/sglang-server/src/mm.rs | 413 ------------ rust/sglang-server/src/multi_modality.rs | 6 + .../payload.rs} | 5 +- rust/sglang-server/src/multi_modality/shm.rs | 122 ++++ .../src/multi_modality/sidecar.rs | 121 ++++ .../src/multi_modality/worker.rs | 187 +++++ rust/sglang-server/src/runtime/config.rs | 312 --------- rust/sglang-server/src/runtime/runnable.rs | 10 - rust/sglang-server/src/tokenizer_manager.rs | 106 +-- .../{ring.rs => tokenizer_manager/channel.rs} | 120 ++-- .../{ => tokenizer_manager}/detokenizer.rs | 89 ++- .../{egress.rs => from_scheduler.rs} | 77 +-- .../{ingress.rs => to_scheduler.rs} | 365 +++++----- .../src/{ => tokenizer_manager}/tokenizer.rs | 18 +- .../src/tokenizer_manager/wiring.rs | 77 +++ rust/sglang-server/src/utils.rs | 6 + rust/sglang-server/src/{ => utils}/environ.rs | 0 rust/sglang-server/src/{ => utils}/error.rs | 4 +- rust/sglang-server/src/{ => utils}/fsm.rs | 16 +- rust/sglang-server/src/utils/logging.rs | 23 + rust/sglang-server/src/utils/regex.rs | 33 +- rust/sglang-server/src/{ => utils}/runtime.rs | 294 ++++---- .../src/{runtime => utils}/threads.rs | 40 +- test/registered/rust/test_run_rust_tests.py | 58 ++ .../multimodal/rust/qwen/test_e2e_parity.py | 2 +- .../rust/shared/test_build_native_mm.py | 2 +- 61 files changed, 2531 insertions(+), 2193 deletions(-) create mode 100644 rust/sglang-server/src/api_server/app.rs create mode 100644 rust/sglang-server/src/message/config.rs create mode 100644 rust/sglang-server/src/message/detok.rs rename rust/sglang-server/src/{ => message}/ids.rs (92%) rename rust/sglang-server/src/message/{egress.rs => response.rs} (93%) delete mode 100644 rust/sglang-server/src/mm.rs create mode 100644 rust/sglang-server/src/multi_modality.rs rename rust/sglang-server/src/{message/mm_payload.rs => multi_modality/payload.rs} (97%) create mode 100644 rust/sglang-server/src/multi_modality/shm.rs create mode 100644 rust/sglang-server/src/multi_modality/sidecar.rs create mode 100644 rust/sglang-server/src/multi_modality/worker.rs delete mode 100644 rust/sglang-server/src/runtime/config.rs delete mode 100644 rust/sglang-server/src/runtime/runnable.rs rename rust/sglang-server/src/{ring.rs => tokenizer_manager/channel.rs} (67%) rename rust/sglang-server/src/{ => tokenizer_manager}/detokenizer.rs (90%) rename rust/sglang-server/src/tokenizer_manager/{egress.rs => from_scheduler.rs} (75%) rename rust/sglang-server/src/tokenizer_manager/{ingress.rs => to_scheduler.rs} (85%) rename rust/sglang-server/src/{ => tokenizer_manager}/tokenizer.rs (96%) create mode 100644 rust/sglang-server/src/tokenizer_manager/wiring.rs rename rust/sglang-server/src/{ => utils}/environ.rs (100%) rename rust/sglang-server/src/{ => utils}/error.rs (92%) rename rust/sglang-server/src/{ => utils}/fsm.rs (95%) create mode 100644 rust/sglang-server/src/utils/logging.rs rename rust/sglang-server/src/{ => utils}/runtime.rs (63%) rename rust/sglang-server/src/{runtime => utils}/threads.rs (81%) create mode 100644 test/registered/rust/test_run_rust_tests.py diff --git a/.github/workflows/_pr-test-check-changes.yml b/.github/workflows/_pr-test-check-changes.yml index 7dd63e1ca..2dfae1948 100644 --- a/.github/workflows/_pr-test-check-changes.yml +++ b/.github/workflows/_pr-test-check-changes.yml @@ -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 }} |" diff --git a/.github/workflows/_pr-test-stage-cpu.yml b/.github/workflows/_pr-test-stage-cpu.yml index cdb5242b3..651a0aa8d 100644 --- a/.github/workflows/_pr-test-stage-cpu.yml +++ b/.github/workflows/_pr-test-stage-cpu.yml @@ -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 diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index c4bd0feda..0b2487619 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -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: diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index af10cf771..86b95e63f 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -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) diff --git a/python/sglang/srt/managers/rust_server.py b/python/sglang/srt/managers/rust_server.py index 4c596bed9..5922be596 100644 --- a/python/sglang/srt/managers/rust_server.py +++ b/python/sglang/srt/managers/rust_server.py @@ -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( diff --git a/python/sglang/srt/managers/scheduler_components/idle_sleeper.py b/python/sglang/srt/managers/scheduler_components/idle_sleeper.py index cef2e3419..5f85391ac 100644 --- a/python/sglang/srt/managers/scheduler_components/idle_sleeper.py +++ b/python/sglang/srt/managers/scheduler_components/idle_sleeper.py @@ -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 diff --git a/python/sglang/srt/managers/utils.py b/python/sglang/srt/managers/utils.py index 69d00c9af..42db8b4e0 100644 --- a/python/sglang/srt/managers/utils.py +++ b/python/sglang/srt/managers/utils.py @@ -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) diff --git a/rust/sglang-mm/src/registry.rs b/rust/sglang-mm/src/registry.rs index 061dbc9f3..d74c441a9 100644 --- a/rust/sglang-mm/src/registry.rs +++ b/rust/sglang-mm/src/registry.rs @@ -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, 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, 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) } diff --git a/rust/sglang-server/src/api_server.rs b/rust/sglang-server/src/api_server.rs index 204582e73..e45b13747 100644 --- a/rust/sglang-server/src/api_server.rs +++ b/rust/sglang-server/src/api_server.rs @@ -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, - chat_formatter: Option, - /// 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, - 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 state, so it cannot merge into the - // Router 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::(), - ); - 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"); - } - } -} diff --git a/rust/sglang-server/src/api_server/app.rs b/rust/sglang-server/src/api_server/app.rs new file mode 100644 index 000000000..13ec61143 --- /dev/null +++ b/rust/sglang-server/src/api_server/app.rs @@ -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` — 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, + pub(super) chat_formatter: Option, + /// 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, + 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 state, so it cannot merge into the + // Router> 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::(), + ); + 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"); + } + } +} diff --git a/rust/sglang-server/src/api_server/common.rs b/rust/sglang-server/src/api_server/common.rs index e7f3025b5..5e6012c79 100644 --- a/rust/sglang-server/src/api_server/common.rs +++ b/rust/sglang-server/src/api_server/common.rs @@ -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 { +pub(super) fn routes() -> Router> { 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 { .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) -> Response { +async fn model_info(State(state): State>) -> 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) -> 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) -> Response { +async fn server_info(State(state): State>) -> Response { let bytes = match await_control_result( &state, - ControlRequest::GetInternalStateReq(GetInternalStateReq::new( - crate::ids::Rid::new().to_string(), - )), + 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. diff --git a/rust/sglang-server/src/api_server/disaggregation/bootstrap.rs b/rust/sglang-server/src/api_server/disaggregation/bootstrap.rs index 2b0a8a3a5..7e3cb49e6 100644 --- a/rust/sglang-server/src/api_server/disaggregation/bootstrap.rs +++ b/rust/sglang-server/src/api_server/disaggregation/bootstrap.rs @@ -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) -> Router { /// Drop room entries async fn cleanup_sweeper(state: Arc) { - 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 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); diff --git a/rust/sglang-server/src/api_server/frame.rs b/rust/sglang-server/src/api_server/frame.rs index f58d42427..64e5f7a42 100644 --- a/rust/sglang-server/src/api_server/frame.rs +++ b/rust/sglang-server/src/api_server/frame.rs @@ -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; } diff --git a/rust/sglang-server/src/api_server/guard.rs b/rust/sglang-server/src/api_server/guard.rs index 728d03ac5..8b53260f5 100644 --- a/rust/sglang-server/src/api_server/guard.rs +++ b/rust/sglang-server/src/api_server/guard.rs @@ -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) -> 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(); diff --git a/rust/sglang-server/src/api_server/log.rs b/rust/sglang-server/src/api_server/log.rs index 0cf9ab596..1c26fa2d9 100644 --- a/rust/sglang-server/src/api_server/log.rs +++ b/rust/sglang-server/src/api_server/log.rs @@ -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). diff --git a/rust/sglang-server/src/api_server/native_api.rs b/rust/sglang-server/src/api_server/native_api.rs index 835c3fb81..2db165407 100644 --- a/rust/sglang-server/src/api_server/native_api.rs +++ b/rust/sglang-server/src/api_server/native_api.rs @@ -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 { +pub(super) fn routes() -> Router> { 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 { +fn health_routes() -> Router> { let timeout = - std::time::Duration::from_secs(crate::environ::env_u64("SGLANG_HEALTH_CHECK_TIMEOUT", 20)); - let probe = get(move |state: State| 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>| 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 { 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, timeout: std::time::Duration) -> Response { +/// Python's `last_receive_tstamp`). +async fn health_generate( + State(state): State>, + 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, 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, 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, + State(state): State>, body: Result, 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, + rx: &mut mpsc::Receiver, 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, -) -> (usize, mpsc::Receiver, Vec) { + mut rx: mpsc::Receiver, +) -> (usize, mpsc::Receiver, Vec) { 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, RequestTiming)>, + receivers: Vec<(Rid, mpsc::Receiver, 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, - ) -> (Rid, mpsc::Receiver, RequestTiming) { + rx: mpsc::Receiver, + ) -> (Rid, mpsc::Receiver, 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); diff --git a/rust/sglang-server/src/api_server/openai.rs b/rust/sglang-server/src/api_server/openai.rs index 554050371..9bd8b493e 100644 --- a/rust/sglang-server/src/api_server/openai.rs +++ b/rust/sglang-server/src/api_server/openai.rs @@ -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 { +pub(super) fn routes() -> Router> { Router::new() .merge(models::routes()) .merge(completions::routes()) @@ -47,7 +50,7 @@ pub(super) fn load_chat_support(server_args: &ServerArgs) -> Option, 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, + mut rx: mpsc::Receiver, guard: &mut AbortGuard, rid: &Rid, ) -> Result { 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, Response> { +) -> Result, 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, -) -> futures::stream::BoxStream<'static, (usize, Option)> { + rx: mpsc::Receiver, +) -> futures::stream::BoxStream<'static, (usize, Option)> { 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))), diff --git a/rust/sglang-server/src/api_server/openai/chat.rs b/rust/sglang-server/src/api_server/openai/chat.rs index df3ed1725..6483e35aa 100644 --- a/rust/sglang-server/src/api_server/openai/chat.rs +++ b/rust/sglang-server/src/api_server/openai/chat.rs @@ -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 { +pub(super) fn routes() -> Router> { Router::new().route("/v1/chat/completions", post(chat_completions)) } async fn chat_completions( - State(state): State, + State(state): State>, body: Result, JsonRejection>, ) -> Response { let request = match body { @@ -283,7 +288,7 @@ pub(super) fn chat_sampling( tool_choice: &DynamoToolChoice, tools: &[ToolDefinition], parallel_tool_calls: Option, - server_args: &crate::runtime::ServerArgs, + server_args: &ServerArgs, ) -> Result { 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)>, + submitted: Vec<(usize, Rid, mpsc::Receiver)>, 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)>, + submitted: Vec<(usize, Rid, mpsc::Receiver)>, 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() ])) diff --git a/rust/sglang-server/src/api_server/openai/completions.rs b/rust/sglang-server/src/api_server/openai/completions.rs index 7147c685f..5fd45b842 100644 --- a/rust/sglang-server/src/api_server/openai/completions.rs +++ b/rust/sglang-server/src/api_server/openai/completions.rs @@ -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 { +pub(super) fn routes() -> Router> { 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, + pub(super) rx: mpsc::Receiver, } #[derive(Debug, Default)] pub(super) struct ChoiceExtensions { @@ -58,7 +61,7 @@ pub(super) struct ChoiceExtensions { } async fn completions( - State(state): State, + State(state): State>, body: Result, 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 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, diff --git a/rust/sglang-server/src/api_server/openai/models.rs b/rust/sglang-server/src/api_server/openai/models.rs index f8a56b7d1..124366c07 100644 --- a/rust/sglang-server/src/api_server/openai/models.rs +++ b/rust/sglang-server/src/api_server/openai/models.rs @@ -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 { +pub(super) fn routes() -> Router> { Router::new() .route("/v1/models", get(available_models)) .route("/v1/models/{model}", get(retrieve_model)) @@ -18,12 +19,12 @@ pub(super) fn routes() -> Router { /// `GET /v1/models` — OpenAI-compatible model list. Served from `server_args`; /// no scheduler round-trip. -async fn available_models(State(state): State) -> Response { +async fn available_models(State(state): State>) -> Response { let base = model_card(&state); Json(serde_json::json!({ "object": "list", "data": [base] })).into_response() } -async fn retrieve_model(State(state): State, Path(model): Path) -> Response { +async fn retrieve_model(State(state): State>, Path(model): Path) -> Response { if model != state.server_args.served_model_name { return openai_error( StatusCode::NOT_FOUND, diff --git a/rust/sglang-server/src/api_server/openai/template.rs b/rust/sglang-server/src/api_server/openai/template.rs index 3f832c7ef..f6131f686 100644 --- a/rust/sglang-server/src/api_server/openai/template.rs +++ b/rust/sglang-server/src/api_server/openai/template.rs @@ -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", diff --git a/rust/sglang-server/src/api_server/openai/test_utils.rs b/rust/sglang-server/src/api_server/openai/test_utils.rs index 68eeea61f..d93303e04 100644 --- a/rust/sglang-server/src/api_server/openai/test_utils.rs +++ b/rust/sglang-server/src/api_server/openai/test_utils.rs @@ -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, + tokio::sync::mpsc::Sender, ) { 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), - tokio::sync::mpsc::Sender, + (usize, Rid, tokio::sync::mpsc::Receiver), + tokio::sync::mpsc::Sender, ) { 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 { - 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 { + 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![], } } diff --git a/rust/sglang-server/src/api_server/openai/tools.rs b/rust/sglang-server/src/api_server/openai/tools.rs index 84c04a8de..c6d61963f 100644 --- a/rust/sglang-server/src/api_server/openai/tools.rs +++ b/rust/sglang-server/src/api_server/openai/tools.rs @@ -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; diff --git a/rust/sglang-server/src/api_server/prefetch.rs b/rust/sglang-server/src/api_server/prefetch.rs index 4d5af2436..18b2a2e1e 100644 --- a/rust/sglang-server/src/api_server/prefetch.rs +++ b/rust/sglang-server/src/api_server/prefetch.rs @@ -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. diff --git a/rust/sglang-server/src/api_server/submit.rs b/rust/sglang-server/src/api_server/submit.rs index fa2f23985..afe79018c 100644 --- a/rust/sglang-server/src/api_server/submit.rs +++ b/rust/sglang-server/src/api_server/submit.rs @@ -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_` 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), Response> { +) -> Result<(Rid, mpsc::Receiver), 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::(state.egress_buf); + let (tx, rx) = mpsc::channel::(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(_) => { diff --git a/rust/sglang-server/src/lib.rs b/rust/sglang-server/src/lib.rs index 2743ed648..e2b9d63c5 100644 --- a/rust/sglang-server/src/lib.rs +++ b/rust/sglang-server/src/lib.rs @@ -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>>, + /// *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>, - 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, + /// *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>, + /// *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>, /// 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, - ingress_ring_cap: usize, - egress_ring_cap: usize, + to_scheduler_cap: usize, + from_scheduler_cap: usize, channel_cap: usize, cores: Option>, - server_args_json: &str, ) -> PyResult { - // 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::(format!( - "bad server_args_json: {e}" - )) - })?; - server_args.validate_mandatory().map_err(|e| { - PyErr::new::(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::(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::(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 { - let cols = self.rt.ingress.drain(max); + fn recv_requests(&self, py: Python<'_>, max: usize) -> PyResult { + 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) -> bool { + fn push_decode_result_batch( + &self, + py: Python<'_>, + header: &[u8], + data_cols: Vec, + ) -> 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::)?; + .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 { + /// 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 { 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 = - 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::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; m.add_class::()?; - m.add_class::()?; - m.add_class::()?; + m.add_class::()?; + m.add_class::()?; Ok(()) } diff --git a/rust/sglang-server/src/message.rs b/rust/sglang-server/src/message.rs index 5747e926c..ff5904462 100644 --- a/rust/sglang-server/src/message.rs +++ b/rust/sglang-server/src/message.rs @@ -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), - /// 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 }, - /// 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; diff --git a/rust/sglang-server/src/message/config.rs b/rust/sglang-server/src/message/config.rs new file mode 100644 index 000000000..61a35d819 --- /dev/null +++ b/rust/sglang-server/src/message/config.rs @@ -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>, +} + +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, +} + +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, + /// 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, + /// 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, + /// 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, + /// 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, + /// Parser selected by `--tool-call-parser`. + pub tool_call_parser: Option, + /// 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, + /// 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, + /// 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, + load_format: Option, + weight_version: Option, + host: String, + port: u16, + log_level: String, + log_level_http: Option, + chat_template: Option, + tool_call_parser: Option, + reasoning_parser: Option, + 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, + 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 { + let text = obj.extract::()?; + 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, + pub top_p: Option, + pub top_k: Option, + pub min_p: Option, + pub repetition_penalty: Option, +} + +#[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, + top_p: Option, + top_k: Option, + min_p: Option, + repetition_penalty: Option, + ) -> 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 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()); + } +} diff --git a/rust/sglang-server/src/message/detok.rs b/rust/sglang-server/src/message/detok.rs new file mode 100644 index 000000000..af18c12d6 --- /dev/null +++ b/rust/sglang-server/src/message/detok.rs @@ -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), + /// 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 }, + /// 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 }, +} diff --git a/rust/sglang-server/src/message/finish_reason.rs b/rust/sglang-server/src/message/finish_reason.rs index d394a5581..a72011ee7 100644 --- a/rust/sglang-server/src/message/finish_reason.rs +++ b/rust/sglang-server/src/message/finish_reason.rs @@ -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}; diff --git a/rust/sglang-server/src/ids.rs b/rust/sglang-server/src/message/ids.rs similarity index 92% rename from rust/sglang-server/src/ids.rs rename to rust/sglang-server/src/message/ids.rs index 3b3ba9640..7d26846c1 100644 --- a/rust/sglang-server/src/ids.rs +++ b/rust/sglang-server/src/message/ids.rs @@ -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 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 = OnceLock::new(); let hash = SEED.get_or_init(RandomState::new).hash_one(&id); Rid { id, hash } diff --git a/rust/sglang-server/src/message/io_struct.rs b/rust/sglang-server/src/message/io_struct.rs index edc9b7f2e..cd8c8fe34 100644 --- a/rust/sglang-server/src/message/io_struct.rs +++ b/rust/sglang-server/src/message/io_struct.rs @@ -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 diff --git a/rust/sglang-server/src/message/request.rs b/rust/sglang-server/src/message/request.rs index dc35dd94b..11bbbca85 100644 --- a/rust/sglang-server/src/message/request.rs +++ b/rust/sglang-server/src/message/request.rs @@ -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, @@ -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) -> 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), /// 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), /// 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, pub video_data: Option, pub audio_data: Option, /// 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, @@ -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()); } diff --git a/rust/sglang-server/src/message/egress.rs b/rust/sglang-server/src/message/response.rs similarity index 93% rename from rust/sglang-server/src/message/egress.rs rename to rust/sglang-server/src/message/response.rs index db448da9e..8e7ce35bb 100644 --- a/rust/sglang-server/src/message/egress.rs +++ b/rust/sglang-server/src/message/response.rs @@ -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), +pub enum ResponseSink { + Local(mpsc::Sender), } -/// 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; +pub type ResponseSource = mpsc::Receiver; -/// 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> { /// 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, } -/// 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 = 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 = [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 = 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 = [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 = [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::(); @@ -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!( diff --git a/rust/sglang-server/src/message/sampling.rs b/rust/sglang-server/src/message/sampling.rs index 99a051eaa..7fb633222 100644 --- a/rust/sglang-server/src/message/sampling.rs +++ b/rust/sglang-server/src/message/sampling.rs @@ -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. diff --git a/rust/sglang-server/src/mm.rs b/rust/sglang-server/src/mm.rs deleted file mode 100644 index d08f6b9c2..000000000 --- a/rust/sglang-server/src/mm.rs +++ /dev/null @@ -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 { - 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::(), 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 { - 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, - pub offsets: Vec<(u32, u32)>, - pub mrope: Vec, - 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), - /// One POSIX segment per item, written by the worker; only the names cross - /// ranks. See [`ShmSegment`]. - Shm(Vec), -} - -/// 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>>); - -impl Sidecar { - pub fn park(&self, rid: String, entry: MmSidecarEntry) { - self.0.lock().unwrap().insert(rid, entry); - } - pub fn take(&self, rid: &str) -> Option { - 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, - /// `None` under `skip_tokenizer_init` (requests must carry `input_ids`). - pub tokenizer: Option>, - 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>, - sidecar: Sidecar, - ) -> Result { - let feature_shm = serde_json::from_str::(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, 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, - tm: flume::Sender, - ctx: Arc, -} - -impl MmWorker { - pub fn new( - rx: flume::Receiver, - tm: flume::Sender, - ctx: Arc, - ) -> 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 = (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 = (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 { std::fs::read(shm_path(&seg.name)).unwrap() }; - assert_eq!( - read(&segments[0]), - bytemuck::cast_slice::(&features[..12]) - ); - assert_eq!( - read(&segments[1]), - bytemuck::cast_slice::(&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(_) - )); - } -} diff --git a/rust/sglang-server/src/multi_modality.rs b/rust/sglang-server/src/multi_modality.rs new file mode 100644 index 000000000..3729b574c --- /dev/null +++ b/rust/sglang-server/src/multi_modality.rs @@ -0,0 +1,6 @@ +//! Multimodal worker pool. + +pub mod payload; +mod shm; +pub mod sidecar; +pub mod worker; diff --git a/rust/sglang-server/src/message/mm_payload.rs b/rust/sglang-server/src/multi_modality/payload.rs similarity index 97% rename from rust/sglang-server/src/message/mm_payload.rs rename to rust/sglang-server/src/multi_modality/payload.rs index 397f88034..f047c3bf7 100644 --- a/rust/sglang-server/src/message/mm_payload.rs +++ b/rust/sglang-server/src/multi_modality/payload.rs @@ -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, diff --git a/rust/sglang-server/src/multi_modality/shm.rs b/rust/sglang-server/src/multi_modality/shm.rs new file mode 100644 index 000000000..e16c0de47 --- /dev/null +++ b/rust/sglang-server/src/multi_modality/shm.rs @@ -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 { + 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::(), 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 = (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()) }; + } +} diff --git a/rust/sglang-server/src/multi_modality/sidecar.rs b/rust/sglang-server/src/multi_modality/sidecar.rs new file mode 100644 index 000000000..fb1bd55a0 --- /dev/null +++ b/rust/sglang-server/src/multi_modality/sidecar.rs @@ -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, + pub offsets: Vec<(u32, u32)>, + pub mrope: Vec, + 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), + /// One POSIX segment per item, written by the worker; only the names cross + /// ranks. See [`ShmSegment`]. + Shm(Vec), +} + +/// 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>>); + +impl Sidecar { + pub fn park(&self, rid: String, entry: MmSidecarEntry) { + self.0.lock().unwrap().insert(rid, entry); + } + pub fn take(&self, rid: &str) -> Option { + 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 = (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 { std::fs::read(shm_path(&seg.name)).unwrap() }; + assert_eq!( + read(&segments[0]), + bytemuck::cast_slice::(&features[..12]) + ); + assert_eq!( + read(&segments[1]), + bytemuck::cast_slice::(&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(_) + )); + } +} diff --git a/rust/sglang-server/src/multi_modality/worker.rs b/rust/sglang-server/src/multi_modality/worker.rs new file mode 100644 index 000000000..391265aeb --- /dev/null +++ b/rust/sglang-server/src/multi_modality/worker.rs @@ -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 { + 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, + /// `None` under `skip_tokenizer_init` (requests must carry `input_ids`). + pub tokenizer: Option>, + 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>, + sidecar: Sidecar, + ) -> Result { + 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, 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, + tm: flume::Sender, + ctx: Arc, +} + +impl MmWorker { + pub fn new( + rx: flume::Receiver, + tm: flume::Sender, + ctx: Arc, + ) -> 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); + } +} diff --git a/rust/sglang-server/src/runtime/config.rs b/rust/sglang-server/src/runtime/config.rs deleted file mode 100644 index b12a1294f..000000000 --- a/rust/sglang-server/src/runtime/config.rs +++ /dev/null @@ -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>, -} - -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, -} - -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, - /// 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, - /// 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, - /// 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, - /// 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, - /// Parser selected by `--tool-call-parser`. - #[serde(default)] - pub tool_call_parser: Option, - /// 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, - /// 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, - /// 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, - #[serde(default)] - pub max_total_num_tokens: Option, -} - -/// 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, - /// 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, - /// 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, - #[serde(default)] - pub top_p: Option, - #[serde(default)] - pub top_k: Option, - #[serde(default)] - pub min_p: Option, - #[serde(default)] - pub repetition_penalty: Option, -} - -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 { - 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) - } -} diff --git a/rust/sglang-server/src/runtime/runnable.rs b/rust/sglang-server/src/runtime/runnable.rs deleted file mode 100644 index 66b7535d7..000000000 --- a/rust/sglang-server/src/runtime/runnable.rs +++ /dev/null @@ -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); -} diff --git a/rust/sglang-server/src/tokenizer_manager.rs b/rust/sglang-server/src/tokenizer_manager.rs index e5b25cc4f..915084539 100644 --- a/rust/sglang-server/src/tokenizer_manager.rs +++ b/rust/sglang-server/src/tokenizer_manager.rs @@ -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(rx: &flume::Receiver, shutdown: &flume::Receiver<()>) -> Option { - 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 }, - /// 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, - /// → 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, - /// → Tokenizer pool (CPU-bound, pinned threads). - pub tok: flume::Sender, - /// → Detokenizer shards, indexed by `Rid::shard(detok.len())`. - pub detok: Vec>, -} - -impl Senders { - #[inline] - pub fn detok_for(&self, rid: &Rid) -> &flume::Sender { - &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; diff --git a/rust/sglang-server/src/ring.rs b/rust/sglang-server/src/tokenizer_manager/channel.rs similarity index 67% rename from rust/sglang-server/src/ring.rs rename to rust/sglang-server/src/tokenizer_manager/channel.rs index f8d0fa1f2..67e6e7540 100644 --- a/rust/sglang-server/src/ring.rs +++ b/rust/sglang-server/src/tokenizer_manager/channel.rs @@ -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, +pub struct ToSchedulerTx { + tx: flume::Sender, } -pub struct IngressConsumer { - rx: flume::Receiver, +pub struct ToSchedulerRx { + rx: flume::Receiver, /// 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>, + /// [`wait`]: ToSchedulerRx::wait + /// [`drain`]: ToSchedulerRx::drain + stash: Mutex>, } -/// 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, /// 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, } -pub struct EgressConsumer { +pub struct FromSchedulerRx { rx: flume::Receiver, } -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> { 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 { &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"))); } diff --git a/rust/sglang-server/src/detokenizer.rs b/rust/sglang-server/src/tokenizer_manager/detokenizer.rs similarity index 90% rename from rust/sglang-server/src/detokenizer.rs rename to rust/sglang-server/src/tokenizer_manager/detokenizer.rs index 71503f780..98a6b7f3c 100644 --- a/rust/sglang-server/src/detokenizer.rs +++ b/rust/sglang-server/src/tokenizer_manager/detokenizer.rs @@ -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>, - /// 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, @@ -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: &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::(1); - tx.try_send(EgressItem::Frame(ChunkEvent::default())) + let (tx, _rx) = mpsc::channel::(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::(4); + let (tx, mut rx) = mpsc::channel::(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::(4); - let (tx_b, mut rx_b) = mpsc::channel::(4); + let (tx_a, mut rx_a) = mpsc::channel::(4); + let (tx_b, mut rx_b) = mpsc::channel::(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| match rx.try_recv() { - Ok(EgressItem::Frame(ev)) => ev.token_ids, + let ids = |rx: &mut mpsc::Receiver| 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, ) -> ChunkEvent { - let (tx, mut rx) = mpsc::channel::(4); + let (tx, mut rx) = mpsc::channel::(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:?}"), } } diff --git a/rust/sglang-server/src/tokenizer_manager/egress.rs b/rust/sglang-server/src/tokenizer_manager/from_scheduler.rs similarity index 75% rename from rust/sglang-server/src/tokenizer_manager/egress.rs rename to rust/sglang-server/src/tokenizer_manager/from_scheduler.rs index 1ad6df0ec..0b5dca72d 100644 --- a/rust/sglang-server/src/tokenizer_manager/egress.rs +++ b/rust/sglang-server/src/tokenizer_manager/from_scheduler.rs @@ -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; -/// 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> = (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); diff --git a/rust/sglang-server/src/tokenizer_manager/ingress.rs b/rust/sglang-server/src/tokenizer_manager/to_scheduler.rs similarity index 85% rename from rust/sglang-server/src/tokenizer_manager/ingress.rs rename to rust/sglang-server/src/tokenizer_manager/to_scheduler.rs index aef1abe4e..37a8cf222 100644 --- a/rust/sglang-server/src/tokenizer_manager/ingress.rs +++ b/rust/sglang-server/src/tokenizer_manager/to_scheduler.rs @@ -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, +pub struct Intake { + tok_manager_rx: flume::Receiver, /// Unbounded abort lane (see [`Senders::abort`]). Selected against `rx` so an /// abort is handled promptly even while the bounded inbox is saturated. abort_rx: flume::Receiver, senders: Senders, - ingress: IngressProducer, + 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 { - 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, + tok_manager_rx: flume::Receiver, abort_rx: flume::Receiver, 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, - IngressConsumer, + ToSchedulerRx, flume::Sender, flume::Receiver, ) { - make_ingress_with(test_limits()) + make_intake_with(test_limits()) } - fn make_ingress_with_abort( + fn make_intake_with_abort( abort_rx: flume::Receiver, ) -> ( - Ingress, + Intake, flume::Receiver, - IngressConsumer, + ToSchedulerRx, flume::Sender, flume::Receiver, ) { - 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, - IngressConsumer, + ToSchedulerRx, flume::Sender, flume::Receiver, ) { let (abort_tx, abort_rx) = flume::unbounded::(); 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, ) -> ( - Ingress, + Intake, flume::Receiver, - IngressConsumer, + ToSchedulerRx, flume::Sender, flume::Receiver, ) { 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::(); - 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::(); 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::(); - 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(" 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::(); 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()); } } diff --git a/rust/sglang-server/src/tokenizer.rs b/rust/sglang-server/src/tokenizer_manager/tokenizer.rs similarity index 96% rename from rust/sglang-server/src/tokenizer.rs rename to rust/sglang-server/src/tokenizer_manager/tokenizer.rs index b0664df56..36d0622f2 100644 --- a/rust/sglang-server/src/tokenizer.rs +++ b/rust/sglang-server/src/tokenizer_manager/tokenizer.rs @@ -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()), diff --git a/rust/sglang-server/src/tokenizer_manager/wiring.rs b/rust/sglang-server/src/tokenizer_manager/wiring.rs new file mode 100644 index 000000000..f50adfa6b --- /dev/null +++ b/rust/sglang-server/src/tokenizer_manager/wiring.rs @@ -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(rx: &flume::Receiver, shutdown: &flume::Receiver<()>) -> Option { + 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 }, + /// 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, + /// → the same loop, but UNBOUNDED and abort-only. + pub abort_tx: flume::Sender, + /// → Tokenizer pool (CPU-bound, pinned threads). + pub tokenizer_tx: flume::Sender, + /// → Detokenizer shards, indexed by `Rid::shard(detok.len())`. + pub detokenizer_tx: Vec>, +} + +impl Senders { + #[inline] + pub fn detok_for(&self, rid: &Rid) -> &flume::Sender { + &self.detokenizer_tx[rid.shard(self.detokenizer_tx.len())] + } +} diff --git a/rust/sglang-server/src/utils.rs b/rust/sglang-server/src/utils.rs index b434cf3a7..2e0e777ac 100644 --- a/rust/sglang-server/src/utils.rs +++ b/rust/sglang-server/src/utils.rs @@ -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; diff --git a/rust/sglang-server/src/environ.rs b/rust/sglang-server/src/utils/environ.rs similarity index 100% rename from rust/sglang-server/src/environ.rs rename to rust/sglang-server/src/utils/environ.rs diff --git a/rust/sglang-server/src/error.rs b/rust/sglang-server/src/utils/error.rs similarity index 92% rename from rust/sglang-server/src/error.rs rename to rust/sglang-server/src/utils/error.rs index 6b5feb8c1..44667b19c 100644 --- a/rust/sglang-server/src/error.rs +++ b/rust/sglang-server/src/utils/error.rs @@ -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`. diff --git a/rust/sglang-server/src/fsm.rs b/rust/sglang-server/src/utils/fsm.rs similarity index 95% rename from rust/sglang-server/src/fsm.rs rename to rust/sglang-server/src/utils/fsm.rs index 5777891dd..c28b6de2f 100644 --- a/rust/sglang-server/src/fsm.rs +++ b/rust/sglang-server/src/utils/fsm.rs @@ -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. diff --git a/rust/sglang-server/src/utils/logging.rs b/rust/sglang-server/src/utils/logging.rs new file mode 100644 index 000000000..b85b65350 --- /dev/null +++ b/rust/sglang-server/src/utils/logging.rs @@ -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 = 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(); +} diff --git a/rust/sglang-server/src/utils/regex.rs b/rust/sglang-server/src/utils/regex.rs index 92106b0a7..e66e132ad 100644 --- a/rust/sglang-server/src/utils/regex.rs +++ b/rust/sglang-server/src/utils/regex.rs @@ -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, 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 { 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, diff --git a/rust/sglang-server/src/runtime.rs b/rust/sglang-server/src/utils/runtime.rs similarity index 63% rename from rust/sglang-server/src/runtime.rs rename to rust/sglang-server/src/utils/runtime.rs index 0cfb9d920..ba1db76cb 100644 --- a/rust/sglang-server/src/runtime.rs +++ b/rust/sglang-server/src/utils/runtime.rs @@ -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, - /// Back-channel for the MM workers' `MmEncoded` / `MmFailed` into tm-ingress. - pub tm: flume::Sender, + /// request never routes to it. + pub to_mm_worker_rx: flume::Receiver, + /// Back-channel for the MM workers' `MmEncoded` / `MmFailed` into to_scheduler. + pub from_mm_worker_tx: flume::Sender, /// The loaded tokenizer, shared with the MM worker path (`None` under /// `skip_tokenizer_init`). pub tokenizer: Option>, /// 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>>, /// 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) { + pub fn spawn_mm_pool(&self, workers: usize, ctx: Arc) { 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> = 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 { 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::(cfg.rust_server_args.channel_cap); - let (tok_tx, tok_rx) = - flume::bounded::(cfg.rust_server_args.channel_cap); + let (tok_manager_tx, tok_manager_rx) = + flume::bounded::(cfg.rust_server_args.channel_cap); + let (tokenizer_tx, tokenizer_rx) = + flume::bounded::(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::(cfg.rust_server_args.channel_cap); + let (mm_worker_tx, mm_worker_rx) = + flume::bounded::(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::(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::(); + let (abort_tx, abort_rx) = flume::unbounded::(); 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 { .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 { 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 { // --- 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 { 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 { }); } - // --- 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 { .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 { 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 { .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 { }); } 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 { } 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 { #[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}"); - } } diff --git a/rust/sglang-server/src/runtime/threads.rs b/rust/sglang-server/src/utils/threads.rs similarity index 81% rename from rust/sglang-server/src/runtime/threads.rs rename to rust/sglang-server/src/utils/threads.rs index 5e56ae4d4..758e098bb 100644 --- a/rust/sglang-server/src/runtime/threads.rs +++ b/rust/sglang-server/src/utils/threads.rs @@ -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 { _ => 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 { let mut it = cores.into_iter(); let api: Vec = 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( 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>, 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() +} diff --git a/test/registered/rust/test_run_rust_tests.py b/test/registered/rust/test_run_rust_tests.py new file mode 100644 index 000000000..085dcf1b7 --- /dev/null +++ b/test/registered/rust/test_run_rust_tests.py @@ -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() diff --git a/test/registered/unit/multimodal/rust/qwen/test_e2e_parity.py b/test/registered/unit/multimodal/rust/qwen/test_e2e_parity.py index 75cec9d64..62f0e235d 100644 --- a/test/registered/unit/multimodal/rust/qwen/test_e2e_parity.py +++ b/test/registered/unit/multimodal/rust/qwen/test_e2e_parity.py @@ -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, diff --git a/test/registered/unit/multimodal/rust/shared/test_build_native_mm.py b/test/registered/unit/multimodal/rust/shared/test_build_native_mm.py index c29ad76cf..5a5b7298f 100644 --- a/test/registered/unit/multimodal/rust/shared/test_build_native_mm.py +++ b/test/registered/unit/multimodal/rust/shared/test_build_native_mm.py @@ -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,