support rust sglang server (#29799)
This commit is contained in:
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user