support rust sglang server (#29799)

This commit is contained in:
Rain Jiang
2026-07-31 11:56:31 -07:00
committed by GitHub
parent 77c77a3da8
commit 4af8ddb576
23 changed files with 1314 additions and 84 deletions
+79 -3
View File
@@ -2,6 +2,8 @@
python3 -m unittest test_srt_endpoint.TestSRTEndpoint.test_simple_decode
python3 -m unittest test_srt_endpoint.TestSRTEndpoint.test_logprob_with_chunked_prefill
python3 -m unittest test_srt_endpoint.TestTokenizeDetokenize
python3 -m unittest test_srt_endpoint.TestRustServerEndpoint
python3 -m unittest test_srt_endpoint.TestRustServerLogprob
"""
import json
@@ -24,17 +26,22 @@ from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
is_rust_server_built,
popen_launch_server,
run_logprob_check,
)
register_cuda_ci(est_time=160, stage="base-b", runner_config="1-gpu-small")
register_amd_ci(est_time=160, suite="stage-b-test-1-gpu-small-amd")
register_cuda_ci(est_time=260, stage="base-b", runner_config="1-gpu-small")
register_amd_ci(est_time=260, suite="stage-b-test-1-gpu-small-amd")
SERVER_ENV = {"SGLANG_USE_PICKLE_IPC": "0"}
class TestSRTEndpoint(CustomTestCase):
# Extra server-launch env; subclasses override to run the same suite
# against a different server flavor (e.g. SGLANG_RUST_SERVER=1).
env = {}
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
@@ -46,7 +53,7 @@ class TestSRTEndpoint(CustomTestCase):
# The tiny logprob chunk size routes this file's logprob tests
# through the multi-chunk stitching path (requests at or below 64
# rows still cover the non-chunked path).
env={**SERVER_ENV, "SGLANG_LOGPROB_CHUNK_SIZE": "64"},
env={**cls.env, **SERVER_ENV, "SGLANG_LOGPROB_CHUNK_SIZE": "64"},
other_args=(
"--enable-custom-logit-processor",
"--mem-fraction-static",
@@ -844,5 +851,74 @@ class TestTokenizeDetokenize(CustomTestCase):
self.assertEqual(r2.status_code, 500)
# ---------------------------------------------------------------------------
# Embedded Rust server (SGLANG_RUST_SERVER=1): rerun the whole endpoint suite
# against the rust api-server/tokenizer/detokenizer stack — the logprob tests
# exercise the columnar egress wire (`push_generation` extras -> Rust
# `BatchHeader`/`for_each_chunk` -> detok reshape) end to end. Suite
# surface the rust server does not implement yet is skipped explicitly below.
# ---------------------------------------------------------------------------
@unittest.skipUnless(
is_rust_server_built(),
"embedded rust server extension not built (e.g. AMD suite)",
)
class TestRustServerEndpoint(TestSRTEndpoint):
env = {"SGLANG_RUST_SERVER": "1"}
_RUST_TODO = "not implemented by the embedded Rust server yet"
@unittest.skip(f"custom_logit_processor request field {_RUST_TODO}")
def test_custom_logit_processor(self):
pass
@unittest.skip(f"custom_logit_processor request field {_RUST_TODO}")
def test_custom_logit_processor_batch_mixed(self):
pass
@unittest.skip(f"custom_logit_processor request field {_RUST_TODO}")
def test_stateful_custom_logit_processor(self):
pass
@unittest.skip(f"custom_logit_processor request field {_RUST_TODO}")
def test_stateful_custom_logit_processor_batch_mixed(self):
pass
@unittest.skip(f"/flush_cache endpoint + cached_tokens meta {_RUST_TODO}")
def test_cache_tokens(self):
pass
def test_greedy_token_equals_top1(self):
"""Cross-column alignment guard for the columnar logprob wire: at
temperature 0 the chosen token must BE the top-1 entry of its own
position. A column shifted across requests or positions (the failure
mode a truncation-tolerant reader would mask) breaks this instantly."""
response = requests.post(
self.base_url + "/generate",
json={
"text": ["The capital of France is", "I have a very good idea on"],
"sampling_params": {"temperature": 0, "max_new_tokens": 8},
"return_logprob": True,
"top_logprobs_num": 5,
"logprob_start_len": 0,
},
)
self.assertEqual(response.status_code, 200, response.text)
for res in response.json():
meta = res["meta_info"]
out_lp = meta["output_token_logprobs"]
top = meta["output_top_logprobs"]
self.assertEqual(len(out_lp), meta["completion_tokens"])
self.assertEqual(len(top), len(out_lp))
# First prompt token's logprob is the None sentinel; it must
# survive the NaN wire encoding and come back as null.
self.assertIsNone(meta["input_token_logprobs"][0][0])
for (lp, tid, _), pos_top in zip(out_lp, top):
self.assertEqual(len(pos_top), 5)
self.assertEqual(pos_top[0][1], tid)
self.assertAlmostEqual(pos_top[0][0], lp, places=4)
vals = [t[0] for t in pos_top]
self.assertEqual(vals, sorted(vals, reverse=True))
if __name__ == "__main__":
unittest.main()
+34 -3
View File
@@ -9,13 +9,21 @@ from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
is_rust_server_built,
popen_launch_server,
)
register_cuda_ci(est_time=62, stage="base-b", runner_config="1-gpu-large")
# Two classes run from this file: the default server plus the Rust-frontend
# variant (when the embedded extension is built), each launches a server + eval.
register_cuda_ci(est_time=124, stage="base-b", runner_config="1-gpu-large")
class TestModeloptFP8(CustomTestCase):
# Extra server env; the Rust-frontend subclass sets SGLANG_RUST_SERVER here.
env = None
# Eval endpoint. The Rust server exposes only the native `/generate`, so its
# subclass overrides this to "generate".
api = "completion"
@classmethod
def setUpClass(cls):
@@ -25,7 +33,15 @@ class TestModeloptFP8(CustomTestCase):
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=["--quantization", "modelopt_fp8"],
other_args=[
"--quantization",
"modelopt_fp8",
"--tokenizer-worker-num",
"2",
"--detokenizer-worker-num",
"2",
],
env=cls.env,
)
@classmethod
@@ -38,7 +54,7 @@ class TestModeloptFP8(CustomTestCase):
base_url=self.base_url,
model=self.model,
eval_name="gsm8k",
api="completion",
api=self.api,
max_tokens=512,
num_examples=200,
num_threads=200,
@@ -48,5 +64,20 @@ class TestModeloptFP8(CustomTestCase):
self.assertGreater(metrics["score"], 0.70)
@unittest.skipUnless(
is_rust_server_built(),
"embedded rust server extension not built",
)
class TestModeloptFP8WithRustServer(TestModeloptFP8):
"""Same model + eval, but served through the embedded Rust frontend
(`SGLANG_RUST_SERVER`). Guards the Rust tokenizer/detokenizer/completions path
against accuracy regressions: a bug there drops gsm8k score below the same
0.70 bar the default frontend must clear. Uses the native `/generate` endpoint
(the only API the Rust server exposes)."""
env = {"SGLANG_RUST_SERVER": "1"}
api = "generate"
if __name__ == "__main__":
unittest.main()
@@ -0,0 +1,76 @@
"""Run the `rust/` Cargo workspace's unit tests from the CPU CI suite.
The `rust/` workspace (sglang-grpc, sglang-mm, sglang-server) is compiled into
the wheel by setuptools-rust, but until now nothing ran `cargo test` in CI --
`.github/workflows/pr-test-rust.yml` and `pr-benchmark-rust.yml` are both
path-scoped to `sgl-model-gateway/**`, a different workspace. `lint.yml` covers
rustfmt/clippy via the pre-commit hooks, so this file only adds the test run.
The debug profile is deliberate: these are pure-logic tests (no timing or
codegen assertions), and the release profile costs a full LTO build for the
same coverage.
"""
import shutil
import subprocess
import unittest
from pathlib import Path
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
# base-c-test-cpu is where this was asked for, and it matches the repo's
# base-a + base-c dual-registration convention -- but base-c-test-cpu currently
# has no runner job in any workflow (it was carved out of base-b in #28623 to
# *reduce* CPU CI scope), so base-a-test-cpu is what actually executes.
register_cpu_ci(est_time=300, suite="base-a-test-cpu")
# repo root: test/registered/rust/<this file>
RUST_WORKSPACE = Path(__file__).resolve().parents[3] / "rust"
# Not `est_time`: that is a scheduling hint for partition balancing (a rough
# average), this is a hard ceiling for the worst case. The 136 tests run in ~1s;
# what varies is the build. Cache-warm the workspace crates recompile in ~15s,
# but a Swatinem/rust-cache miss rebuilds all ~370 dependencies -- measured at
# 48s on 4 fast cores, so several minutes on a hosted runner.
#
# Capped below the 600s `timeout-minutes` on the suite's "Run test" step so a
# hang fails here, with output, instead of being killed as an opaque job
# timeout. The harness `--timeout-per-file` (1200s) is looser still.
BUILD_AND_RUN_TIMEOUT_S = 300
class TestCargoWorkspace(CustomTestCase):
def test_cargo_test_workspace(self):
# Not skipUnless: cargo is a hard dependency of the editable install
# (setuptools-rust builds sglang-grpc), so a missing toolchain is a
# broken environment, and a silently-skipped CI test is worthless.
self.assertIsNotNone(
shutil.which("cargo"),
"cargo not found on PATH; install a Rust toolchain "
"(scripts/ci/utils/install_rust_protoc.sh)",
)
self.assertTrue(
(RUST_WORKSPACE / "Cargo.toml").is_file(),
f"rust workspace manifest not found at {RUST_WORKSPACE}",
)
proc = subprocess.run(
["cargo", "test", "--workspace"],
cwd=RUST_WORKSPACE,
capture_output=True,
text=True,
timeout=BUILD_AND_RUN_TIMEOUT_S,
)
# Print unconditionally so a green run still shows which tests ran.
print(proc.stdout)
self.assertEqual(
proc.returncode,
0,
f"`cargo test --workspace` failed in {RUST_WORKSPACE}\n"
f"--- stdout ---\n{proc.stdout}\n--- stderr ---\n{proc.stderr}",
)
if __name__ == "__main__":
unittest.main()