[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
@@ -277,11 +277,24 @@ clean_site_packages() {
|
||||
}
|
||||
|
||||
setup_cargo_cache() {
|
||||
if [ "${SGLANG_BUILD_RUST_EXTS:-}" = "none" ]; then
|
||||
echo "Using prebuilt Rust extensions; skipping Cargo target setup"
|
||||
mark_step_done "${FUNCNAME[0]}"
|
||||
return
|
||||
fi
|
||||
|
||||
# actions/checkout's `git clean -ffdx` deletes the gitignored in-repo
|
||||
# rust/target, so every job recompiles the whole dependency graph. Move the
|
||||
# target dir out of the tree: setuptools-rust has no target-dir option of its
|
||||
# own and defers to CARGO_TARGET_DIR, which uv passes to the build backend.
|
||||
export CARGO_TARGET_DIR="${HOME}/.cache/sglang-cargo-target"
|
||||
local cargo_target_lock="${HOME}/.cache/sglang-cargo-target.lock"
|
||||
mkdir -p "${HOME}/.cache"
|
||||
exec 9>"${cargo_target_lock}"
|
||||
echo "Waiting for exclusive cargo target lock: ${cargo_target_lock}"
|
||||
flock --exclusive 9
|
||||
CARGO_TARGET_LOCK_HELD=1
|
||||
echo "Acquired cargo target lock"
|
||||
mkdir -p "${CARGO_TARGET_DIR}"
|
||||
|
||||
# Same disk-pressure guard as the uv cache in ci_cleanup_venv.sh (which
|
||||
@@ -298,6 +311,15 @@ setup_cargo_cache() {
|
||||
mark_step_done "${FUNCNAME[0]}"
|
||||
}
|
||||
|
||||
release_cargo_cache_lock() {
|
||||
if [ "${CARGO_TARGET_LOCK_HELD:-0}" = "1" ]; then
|
||||
flock --unlock 9
|
||||
exec 9>&-
|
||||
CARGO_TARGET_LOCK_HELD=0
|
||||
echo "Released cargo target lock"
|
||||
fi
|
||||
}
|
||||
|
||||
setup_pip_toolchain() {
|
||||
if [ "$USE_VENV" = "1" ]; then
|
||||
# The bootstrap upgrade hit system pip; this upgrades the venv's own.
|
||||
@@ -473,10 +495,19 @@ require_prebuilt_rust_exts() {
|
||||
for module in server grpc multimodal; do
|
||||
[ -f "python/sglang/srt/rust_extensions/_${module}${suffix}" ] || missing+=("${module}")
|
||||
done
|
||||
[ -f "python/sglang/srt/mem_cache/rust_tree_core/mem_cache${suffix}" ] \
|
||||
|| missing+=("mem_cache")
|
||||
[ -f "python/sglang/srt/mem_cache/rust_tree_core/mem_cache_inspection${suffix}" ] \
|
||||
|| missing+=("mem_cache_inspection")
|
||||
if [ ${#missing[@]} -gt 0 ]; then
|
||||
echo "::warning::no prebuilt Rust extension ${suffix} for: ${missing[*]}; building from source"
|
||||
ls -l python/sglang/srt/rust_extensions/_*.so 2>/dev/null || echo "(no extension modules at all)"
|
||||
ls -l python/sglang/srt/mem_cache/rust_tree_core/mem_cache*.so 2>/dev/null || true
|
||||
export SGLANG_BUILD_RUST_EXTS=
|
||||
export SGLANG_RUST_BUILD_MODE=auto
|
||||
if [ -n "${GITHUB_ENV:-}" ]; then
|
||||
echo "SGLANG_RUST_BUILD_MODE=auto" >> "${GITHUB_ENV}"
|
||||
fi
|
||||
mark_step_done "${FUNCNAME[0]}"
|
||||
return
|
||||
fi
|
||||
@@ -846,14 +877,15 @@ main() {
|
||||
install_apt_packages
|
||||
install_gdrcopy
|
||||
clean_site_packages
|
||||
setup_cargo_cache
|
||||
require_prebuilt_rust_exts
|
||||
setup_pip_toolchain
|
||||
remove_stale_cuda12_nvidia_wheels
|
||||
uninstall_stale_flashinfer
|
||||
install_pytorch_stack
|
||||
install_cuda12_deepep_wheel
|
||||
setup_cargo_cache
|
||||
install_sglang
|
||||
release_cargo_cache_lock
|
||||
install_nccl
|
||||
# Diffusion B200 CI imports torch inside install_sglang_kernel after removing
|
||||
# stale CUDA 12 NVIDIA wheels, so opt into one early LD_LIBRARY_PATH refresh.
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
#!/bin/bash
|
||||
# Copy the built PyO3 extension modules into rust-ext-staging/rust_extensions/ for
|
||||
# upload-artifact. Shared by both jobs of _pr-test-rust-ext-build.yml, so the
|
||||
# archive layout and the module-count check cannot drift between them.
|
||||
# Copy the built PyO3 extension modules into their package-relative paths under
|
||||
# rust-ext-staging/. Shared by both jobs of _pr-test-rust-ext-build.yml, so the
|
||||
# archive layout and module-count checks cannot drift between them.
|
||||
#
|
||||
# MAX_GLIBC (optional): also reject a module requiring a newer GLIBC symbol
|
||||
# version than the test runners have. Only set where the modules were just
|
||||
@@ -32,6 +32,23 @@ for module in server grpc multimodal; do
|
||||
cp "${found[@]}" rust-ext-staging/rust_extensions/
|
||||
built+=("${found[@]}")
|
||||
done
|
||||
|
||||
mkdir -p rust-ext-staging/mem_cache/rust_tree_core
|
||||
for module in mem_cache mem_cache_inspection; do
|
||||
tree_core=(python/sglang/srt/mem_cache/rust_tree_core/"${module}".*.so)
|
||||
if [ ${#tree_core[@]} -eq 0 ]; then
|
||||
echo "::error::no Rust TreeCore ${module} extension module found"
|
||||
exit 1
|
||||
fi
|
||||
tree_core_suffixes=$(printf '%s\n' "${tree_core[@]##*/${module}}" | sort)
|
||||
if [ "${tree_core_suffixes}" != "${expected_suffixes}" ]; then
|
||||
echo "::error::Rust TreeCore ${module} extension does not match the interpreter set"
|
||||
printf 'have:\n%s\nwant:\n%s\n' "${tree_core_suffixes}" "${expected_suffixes}"
|
||||
exit 1
|
||||
fi
|
||||
cp "${tree_core[@]}" rust-ext-staging/mem_cache/rust_tree_core/
|
||||
built+=("${tree_core[@]}")
|
||||
done
|
||||
max_allowed="${MAX_GLIBC:-}"
|
||||
[ -n "${max_allowed}" ] || exit 0
|
||||
|
||||
|
||||
Executable
+203
@@ -0,0 +1,203 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Repair an SGLang wheel and smoke-test its production Rust TreeCore."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import textwrap
|
||||
import zipfile
|
||||
from email.parser import BytesParser
|
||||
from pathlib import Path, PurePosixPath
|
||||
|
||||
_LIBTORCH_EXCLUDES = (
|
||||
"libc10.so",
|
||||
"libc10_cuda.so",
|
||||
"libtorch.so",
|
||||
"libtorch_cpu.so",
|
||||
"libtorch_cuda.so",
|
||||
"libtorch_python.so",
|
||||
)
|
||||
_TREE_CORE_DIR = PurePosixPath("sglang/srt/mem_cache/rust_tree_core")
|
||||
_BINDING_CLASSES = (
|
||||
"RustUnifiedTreeCoreBinding",
|
||||
"RustBigramUnifiedTreeCoreBinding",
|
||||
"TreeCoreInitParamsBinding",
|
||||
)
|
||||
|
||||
|
||||
def _single_wheel(directory: Path) -> Path:
|
||||
wheels = sorted(directory.glob("*.whl"))
|
||||
if len(wheels) != 1:
|
||||
raise RuntimeError(f"expected one wheel in {directory}, found {wheels}")
|
||||
return wheels[0]
|
||||
|
||||
|
||||
def _metadata(wheel: Path) -> tuple[str, str]:
|
||||
with zipfile.ZipFile(wheel) as archive:
|
||||
metadata_files = [
|
||||
name for name in archive.namelist() if name.endswith(".dist-info/METADATA")
|
||||
]
|
||||
if len(metadata_files) != 1:
|
||||
raise RuntimeError(
|
||||
f"expected one METADATA file in {wheel}, found {metadata_files}"
|
||||
)
|
||||
metadata = BytesParser().parsebytes(archive.read(metadata_files[0]))
|
||||
return str(metadata["Name"]), str(metadata["Version"])
|
||||
|
||||
|
||||
def _smoke_test_tree_core(wheel: Path) -> None:
|
||||
with tempfile.TemporaryDirectory(prefix="sglang-wheel-smoke-") as temp_dir:
|
||||
root = Path(temp_dir)
|
||||
with zipfile.ZipFile(wheel) as archive:
|
||||
names = archive.namelist()
|
||||
inspection_modules = [
|
||||
name
|
||||
for name in names
|
||||
if PurePosixPath(name).parent == _TREE_CORE_DIR
|
||||
and PurePosixPath(name).name.startswith("mem_cache_inspection")
|
||||
and name.endswith(".so")
|
||||
]
|
||||
if inspection_modules:
|
||||
raise RuntimeError(
|
||||
f"production wheel contains inspection modules: {inspection_modules}"
|
||||
)
|
||||
production_modules = [
|
||||
name
|
||||
for name in names
|
||||
if PurePosixPath(name).parent == _TREE_CORE_DIR
|
||||
and PurePosixPath(name).name.startswith("mem_cache.")
|
||||
and name.endswith(".so")
|
||||
]
|
||||
if len(production_modules) != 1:
|
||||
raise RuntimeError(
|
||||
"expected one production Rust TreeCore module, found "
|
||||
f"{production_modules}"
|
||||
)
|
||||
|
||||
install_dir = root / "installed"
|
||||
subprocess.run(
|
||||
[
|
||||
sys.executable,
|
||||
"-m",
|
||||
"pip",
|
||||
"install",
|
||||
"--disable-pip-version-check",
|
||||
"--no-compile",
|
||||
"--no-deps",
|
||||
"--no-index",
|
||||
"--target",
|
||||
os.fspath(install_dir),
|
||||
os.fspath(wheel),
|
||||
],
|
||||
check=True,
|
||||
)
|
||||
|
||||
smoke_script = textwrap.dedent(f"""
|
||||
import sys
|
||||
import types
|
||||
from pathlib import Path
|
||||
|
||||
site_packages = Path({os.fspath(install_dir)!r}).resolve()
|
||||
sys.path.insert(0, str(site_packages))
|
||||
package = types.ModuleType("sglang")
|
||||
package.__package__ = "sglang"
|
||||
package.__path__ = [str(site_packages / "sglang")]
|
||||
sys.modules["sglang"] = package
|
||||
|
||||
from sglang.srt.mem_cache.rust_tree_core.extension import bindings
|
||||
|
||||
module_path = Path(bindings.__file__).resolve()
|
||||
if site_packages not in module_path.parents:
|
||||
raise RuntimeError(
|
||||
f"loaded TreeCore outside installed wheel: {{module_path}}"
|
||||
)
|
||||
if bindings.__name__ != "sglang.srt.mem_cache.rust_tree_core.mem_cache":
|
||||
raise RuntimeError(
|
||||
f"loaded unexpected TreeCore module: {{bindings.__name__}}"
|
||||
)
|
||||
for class_name in {_BINDING_CLASSES!r}:
|
||||
binding = getattr(bindings, class_name, None)
|
||||
if binding is None:
|
||||
raise RuntimeError(
|
||||
f"production TreeCore is missing {{class_name}}"
|
||||
)
|
||||
inspection_methods = [
|
||||
name for name in dir(binding) if name.startswith("inspect_")
|
||||
]
|
||||
if inspection_methods:
|
||||
raise RuntimeError(
|
||||
f"production {{class_name}} exposes inspection methods: "
|
||||
f"{{inspection_methods}}"
|
||||
)
|
||||
|
||||
from array import array
|
||||
|
||||
hashes = bindings.get_hash_str(array("q", [1, 2]), None, 1)
|
||||
if len(hashes) != 2 or any(len(value) != 64 for value in hashes):
|
||||
raise RuntimeError(f"unexpected TreeCore hash result: {{hashes}}")
|
||||
""")
|
||||
environment = os.environ.copy()
|
||||
environment["SGLANG_RUST_BUILD_MODE"] = "never"
|
||||
environment.pop("PYTHONPATH", None)
|
||||
subprocess.run(
|
||||
[sys.executable, "-I", "-c", smoke_script],
|
||||
cwd=root,
|
||||
env=environment,
|
||||
check=True,
|
||||
)
|
||||
|
||||
|
||||
def _write_github_outputs(path: Path, *, wheel: Path, version: str) -> None:
|
||||
with path.open("a", encoding="utf-8") as output:
|
||||
output.write(f"wheel_filename={wheel.name}\n")
|
||||
output.write(f"wheel_version={version}\n")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("wheel_dir", type=Path)
|
||||
parser.add_argument("--github-output", type=Path)
|
||||
args = parser.parse_args()
|
||||
|
||||
wheel_dir = args.wheel_dir.resolve()
|
||||
source_wheel = _single_wheel(wheel_dir)
|
||||
with tempfile.TemporaryDirectory(
|
||||
prefix="sglang-wheel-repair-", dir=wheel_dir.parent
|
||||
) as repair_dir:
|
||||
command = [
|
||||
sys.executable,
|
||||
"-m",
|
||||
"auditwheel",
|
||||
"repair",
|
||||
os.fspath(source_wheel),
|
||||
"--wheel-dir",
|
||||
repair_dir,
|
||||
]
|
||||
for library in _LIBTORCH_EXCLUDES:
|
||||
command.extend(("--exclude", library))
|
||||
subprocess.run(command, check=True)
|
||||
|
||||
repaired_wheel = _single_wheel(Path(repair_dir))
|
||||
name, version = _metadata(repaired_wheel)
|
||||
if name.casefold() != "sglang":
|
||||
raise RuntimeError(f"expected sglang wheel, found {name!r}")
|
||||
_smoke_test_tree_core(repaired_wheel)
|
||||
|
||||
destination = wheel_dir / repaired_wheel.name
|
||||
source_wheel.unlink()
|
||||
shutil.move(repaired_wheel, destination)
|
||||
|
||||
if args.github_output is not None:
|
||||
_write_github_outputs(
|
||||
args.github_output.resolve(), wheel=destination, version=version
|
||||
)
|
||||
print(f"Prepared {destination.name} (sglang {version})")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user