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

[Radix Cache] Add Rust TreeCore backend with shared parity tests (#32710)

Browse Source 
Co-authored-by: alphabetc1 <2508695655@qq.com>
Co-authored-by: ispobock <ispobaoke@gmail.com>
This commit is contained in:
Jialin Ouyang
2026-09-01 00:26:20 +08:00
committed by GitHub
co-authored by alphabetc1 ispobock
parent 52e1c24744
commit 9cf157c252
72 changed files with 39973 additions and 396 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/rust/test_run_mem_cache_rust_tests.py
+64
View File
@@ -0,0 +1,64 @@
"""Run the standalone mem-cache crate's native Rust unit tests."""
import shutil
import subprocess
import unittest
from pathlib import Path
from sglang.srt.environ import envs
from sglang.srt.rust_extensions.torch_build import torch_build_configuration
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"
MEM_CACHE_MANIFEST = RUST_WORKSPACE / "mem-cache" / "Cargo.toml"
register_cpu_ci(est_time=900, suite="base-a-test-cpu")
@unittest.skipIf(
envs.SGLANG_SKIP_RUST_TESTS.get(),
"SGLANG_SKIP_RUST_TESTS is set (no rust/ workspace changes per CI check-changes)",
)
class TestMemCacheCargo(CustomTestCase):
def test_mem_cache_native_tests(self):
self.assertIsNotNone(
shutil.which("cargo"),
"cargo not found on PATH; install a Rust toolchain "
"(scripts/ci/utils/install_rust_protoc.sh)",
)
self.assertTrue(
MEM_CACHE_MANIFEST.is_file(),
f"mem-cache manifest not found at {MEM_CACHE_MANIFEST}",
)
build = torch_build_configuration(
compat_header=MEM_CACHE_MANIFEST.parent / "torch_2_13_compat.h",
python_module="sglang.srt.mem_cache.rust_tree_core.mem_cache",
)
proc = subprocess.run(
[
"cargo",
"test",
"--manifest-path",
str(MEM_CACHE_MANIFEST),
"--locked",
"--no-default-features",
],
cwd=RUST_WORKSPACE,
env=build.environment,
capture_output=True,
text=True,
timeout=BUILD_AND_RUN_TIMEOUT_S,
)
print(proc.stdout)
self.assertEqual(
proc.returncode,
0,
f"mem-cache native tests failed\n"
f"--- stdout ---\n{proc.stdout}\n--- stderr ---\n{proc.stderr}",
)
if __name__ == "__main__":
unittest.main()
test/registered/rust/test_run_rust_tests.py
+20 -17
View File
@@ -1,4 +1,4 @@
"""Run the `rust/` Cargo workspace's unit tests from the CPU CI suite."""
"""Run the repository's native Rust unit tests from the CPU CI suite."""
import shutil
import subprocess
@@ -11,7 +11,6 @@ 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")
@@ -23,6 +22,24 @@ register_cpu_ci(est_time=900, suite="base-a-test-cpu")
"SGLANG_SKIP_RUST_TESTS is set (no rust/ workspace changes per CI check-changes)",
)
class TestCargoWorkspace(CustomTestCase):
def _run_cargo(self, args: list[str], *, cwd: Path, env: dict | None = None):
proc = subprocess.run(
["cargo", *args],
cwd=cwd,
env=env,
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 {' '.join(args)}` failed in {cwd}\n"
f"--- stdout ---\n{proc.stdout}\n--- stderr ---\n{proc.stderr}",
)
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
@@ -37,21 +54,7 @@ class TestCargoWorkspace(CustomTestCase):
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}",
)
self._run_cargo(["test", "--workspace"], cwd=RUST_WORKSPACE)
if __name__ == "__main__":
test/registered/rust/test_rust_extension.py
+212 -3
View File
@@ -9,11 +9,12 @@ import time
import unittest
from pathlib import Path
from tempfile import TemporaryDirectory
from types import ModuleType
from types import ModuleType, SimpleNamespace
from unittest import mock
from sglang.srt.rust_extensions import load_rust_extension
from sglang.srt.rust_extensions import loader as rust_extension
from sglang.srt.rust_extensions.torch_build import torch_build_configuration
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
@@ -84,6 +85,74 @@ crate-type = ["cdylib"]
fingerprint.assert_not_called()
cargo_build.assert_not_called()
def test_bundled_named_variant_never_touches_source_or_cargo(self):
bundled = ModuleType("demo._inspection")
with (
mock.patch.object(
rust_extension.importlib, "import_module", return_value=bundled
) as import_module,
mock.patch.object(rust_extension, "_discover_crate") as discover,
mock.patch.object(rust_extension, "_build_context") as fingerprint,
mock.patch.object(rust_extension, "_cargo_build") as cargo_build,
):
self.assertIs(
load_rust_extension(
"demo._core",
mode="never",
workspace=Path("/workspace/not-present"),
additional_features=("inspection",),
extension_module="demo._inspection",
),
bundled,
)
import_module.assert_called_once_with("demo._inspection")
discover.assert_not_called()
fingerprint.assert_not_called()
cargo_build.assert_not_called()
def test_auto_ignores_a_stale_bundled_extension_in_a_source_tree(self):
with TemporaryDirectory() as directory:
root = Path(directory)
workspace = self._workspace(root)
(workspace / "demo/lib.rs").write_text(
"fn source_changed() {}\n", encoding="utf-8"
)
stale = ModuleType("demo._core")
built = ModuleType("demo._core")
artifact = root / "libdemo_extension.so"
artifact.write_bytes(b"fresh extension")
context = rust_extension._BuildContext(
"changed-source", "fingerprint", "target"
)
with (
mock.patch.object(
rust_extension, "_import_bundled_extension", return_value=stale
) as bundled_import,
mock.patch.object(
rust_extension, "_build_context", return_value=context
),
mock.patch.object(
rust_extension, "_source_digest", return_value="changed-source"
),
mock.patch.object(
rust_extension, "_cargo_build", return_value=artifact
) as cargo_build,
mock.patch.object(
rust_extension, "_load_extension_from_path", return_value=built
),
):
self.assertIs(
load_rust_extension(
"demo._core",
mode="auto",
workspace=workspace,
cache_dir=root / "cache",
),
built,
)
bundled_import.assert_not_called()
cargo_build.assert_called_once()
def test_discovery_reads_crate_manifest_metadata(self):
with TemporaryDirectory() as directory:
workspace = self._workspace(Path(directory))
@@ -126,6 +195,18 @@ crate-type = ["cdylib"]
changed_flags.target_fingerprint,
)
inspection = rust_extension._build_context(
crate,
features=(*crate.features, "inspection"),
extension_module="demo._inspection",
build_fingerprint={"torch": "2.13"},
)
self.assertNotEqual(changed_source.fingerprint, inspection.fingerprint)
self.assertNotEqual(
changed_source.target_fingerprint,
inspection.target_fingerprint,
)
def test_auto_builds_once_then_uses_cache(self):
with TemporaryDirectory() as directory:
root = Path(directory)
@@ -155,13 +236,19 @@ crate-type = ["cdylib"]
):
self.assertIs(
rust_extension.load_rust_extension(
"demo._core", workspace=workspace, cache_dir=root / "cache"
"demo._core",
mode="auto",
workspace=workspace,
cache_dir=root / "cache",
),
loaded,
)
self.assertIs(
rust_extension.load_rust_extension(
"demo._core", workspace=workspace, cache_dir=root / "cache"
"demo._core",
mode="auto",
workspace=workspace,
cache_dir=root / "cache",
),
loaded,
)
@@ -270,6 +357,122 @@ crate-type = ["cdylib"]
],
)
def test_variant_uses_its_own_module_name_features_and_environment(self):
with TemporaryDirectory() as directory:
root = Path(directory)
workspace = self._workspace(root)
artifact = root / "libdemo_extension.so"
artifact.write_bytes(b"extension")
context = rust_extension._BuildContext("source", "fingerprint", "target")
loaded = ModuleType("demo._inspection")
environment = {"CUSTOM_BUILD_INPUT": "value"}
with (
mock.patch.object(
rust_extension, "_import_bundled_extension", return_value=None
) as bundled_import,
mock.patch.object(
rust_extension, "_build_context", return_value=context
) as build_context,
mock.patch.object(
rust_extension, "_source_digest", return_value="source"
),
mock.patch.object(
rust_extension, "_cargo_build", return_value=artifact
) as cargo_build,
mock.patch.object(
rust_extension,
"_load_extension_from_path",
return_value=loaded,
) as load_from_path,
):
self.assertIs(
load_rust_extension(
"demo._core",
mode="auto",
workspace=workspace,
cache_dir=root / "cache",
additional_features=("inspection",),
extension_module="demo._inspection",
build_environment=environment,
build_fingerprint={"native": "abi"},
),
loaded,
)
bundled_import.assert_not_called()
self.assertEqual(
build_context.call_args.kwargs,
{
"features": ("python", "inspection"),
"build_fingerprint": {"native": "abi"},
"extension_module": "demo._inspection",
},
)
self.assertEqual(
cargo_build.call_args.kwargs,
{
"features": ("python", "inspection"),
"build_environment": environment,
},
)
self.assertEqual(load_from_path.call_args.args[0], "demo._inspection")
def test_torch_build_configuration_is_versioned_and_relocatable(self):
with TemporaryDirectory() as directory:
root = Path(directory)
torch_root = root / "torch"
(torch_root / "lib").mkdir(parents=True)
torch_init = torch_root / "__init__.py"
torch_init.write_text("", encoding="utf-8")
compat_header = root / "compat.h"
compat_header.write_text("// compatibility\n", encoding="utf-8")
fake_torch = SimpleNamespace(
__version__="2.13.0+cu130",
__file__=str(torch_init),
compiled_with_cxx11_abi=lambda: True,
version=SimpleNamespace(cuda="13.0", hip=None),
)
build = torch_build_configuration(
compat_header=compat_header,
python_module="sglang.srt.mem_cache.rust_tree_core.mem_cache",
torch_module=fake_torch,
base_environment={
"PATH": "/usr/bin",
"CXXFLAGS": "-O2",
"RUSTFLAGS": "-Ctarget-cpu=x86-64",
},
)
self.assertEqual(build.environment["LIBTORCH_USE_PYTORCH"], "1")
self.assertEqual(build.environment["LIBTORCH_BYPASS_VERSION_CHECK"], "1")
self.assertIn(str(compat_header), build.environment["CXXFLAGS"])
self.assertIn(
"$ORIGIN/../../../../torch/lib", build.environment["RUSTFLAGS"]
)
self.assertIn(str(torch_root / "lib"), build.environment["RUSTFLAGS"])
self.assertEqual(build.fingerprint["torch_version"], "2.13.0+cu130")
self.assertTrue(build.fingerprint["torch_cxx11_abi"])
wheel_build = torch_build_configuration(
compat_header=compat_header,
python_module="sglang.srt.mem_cache.rust_tree_core.mem_cache",
torch_module=fake_torch,
base_environment={},
include_absolute_rpath=False,
)
self.assertNotIn(
str(torch_root / "lib"), wheel_build.environment["RUSTFLAGS"]
)
self.assertFalse(wheel_build.fingerprint["include_absolute_rpath"])
fake_torch.__version__ = "2.14.0"
with self.assertRaisesRegex(RuntimeError, "PyTorch 2.11 through 2.13"):
torch_build_configuration(
compat_header=compat_header,
python_module="sglang.srt.mem_cache.rust_tree_core.mem_cache",
torch_module=fake_torch,
)
def test_filesystem_lock_serializes_processes(self):
with TemporaryDirectory() as directory:
lock_path = Path(directory) / "build.lock"
@@ -337,6 +540,12 @@ crate-type = ["cdylib"]
"sglang_mm_core",
("python", "parallel"),
),
(
"sglang.srt.mem_cache.rust_tree_core.mem_cache",
"mem_cache",
"mem_cache",
("python-extension",),
),
):
crate = rust_extension._discover_crate(
rust_extension._RUST_WORKSPACE, python_module