[Radix Cache] Add Rust TreeCore backend with shared parity tests (#32710)
Co-authored-by: alphabetc1 <2508695655@qq.com> Co-authored-by: ispobock <ispobaoke@gmail.com>
This commit is contained in:
Jialin Ouyang
co-authored by
alphabetc1
ispobock
alphabetc1
ispobock
parent
52e1c24744
commit
9cf157c252
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user