From 6a0c55fd6c48f79ff48008cbcd849a54b6ccc0da Mon Sep 17 00:00:00 2001 From: Alex Nails Date: Sat, 5 Sep 2026 15:30:32 -0700 Subject: [PATCH] [CI] Pin the Rust TreeCore build to the resolved libtorch instead of interpreter discovery (#37696) Co-authored-by: Claude Opus 5 (1M context) --- .../sglang/srt/rust_extensions/torch_build.py | 6 ++- scripts/ci/cuda/ci_install_dependency.sh | 15 ------ test/registered/rust/test_rust_extension.py | 7 ++- .../tools/test_cuda_ci_install_dependency.py | 46 ------------------- 4 files changed, 11 insertions(+), 63 deletions(-) delete mode 100644 test/registered/unit/tools/test_cuda_ci_install_dependency.py diff --git a/python/sglang/srt/rust_extensions/torch_build.py b/python/sglang/srt/rust_extensions/torch_build.py index f4fd40adb..65629b90c 100644 --- a/python/sglang/srt/rust_extensions/torch_build.py +++ b/python/sglang/srt/rust_extensions/torch_build.py @@ -74,7 +74,11 @@ def torch_build_configuration( cxx11_abi = bool(torch_module._C._GLIBCXX_USE_CXX11_ABI) environment = dict(os.environ if base_environment is None else base_environment) - environment["LIBTORCH_USE_PYTORCH"] = "1" + environment.pop("LIBTORCH_USE_PYTORCH", None) + environment["LIBTORCH"] = os.fspath(torch_root) + environment["LIBTORCH_INCLUDE"] = os.fspath(torch_root) + environment["LIBTORCH_LIB"] = os.fspath(torch_root) + environment["LIBTORCH_CXX11_ABI"] = "1" if cxx11_abi else "0" # tch 0.24 targets Torch 2.11. The compatibility header below covers the # API removals in the supported 2.12/2.13 builds, after this explicit gate. environment["LIBTORCH_BYPASS_VERSION_CHECK"] = "1" diff --git a/scripts/ci/cuda/ci_install_dependency.sh b/scripts/ci/cuda/ci_install_dependency.sh index eb57a8a97..067c64c0a 100755 --- a/scripts/ci/cuda/ci_install_dependency.sh +++ b/scripts/ci/cuda/ci_install_dependency.sh @@ -324,20 +324,6 @@ setup_cargo_cache() { mark_step_done "${FUNCNAME[0]}" } -invalidate_torch_rust_cache() { - if [ "${SGLANG_BUILD_RUST_EXTS:-}" = "none" ]; then - mark_step_done "${FUNCNAME[0]}" - return - fi - - # uv's editable build uses a temporary torch path. Rebuild these units - # under the lock so Cargo does not reuse that path in a later job. - cargo clean --release --manifest-path "${REPO_ROOT}/rust/sglang-radix-tree/Cargo.toml" \ - -p torch-sys -p sglang-radix-tree - - mark_step_done "${FUNCNAME[0]}" -} - release_cargo_cache_lock() { if [ "${CARGO_TARGET_LOCK_HELD:-0}" = "1" ]; then flock --unlock 9 @@ -918,7 +904,6 @@ main() { install_pytorch_stack install_cuda12_deepep_wheel setup_cargo_cache - invalidate_torch_rust_cache install_sglang release_cargo_cache_lock # Diffusion B200 CI imports torch inside install_sglang_kernel after removing diff --git a/test/registered/rust/test_rust_extension.py b/test/registered/rust/test_rust_extension.py index 60ec85345..14765771a 100644 --- a/test/registered/rust/test_rust_extension.py +++ b/test/registered/rust/test_rust_extension.py @@ -441,10 +441,15 @@ crate-type = ["cdylib"] "PATH": "/usr/bin", "CXXFLAGS": "-O2", "RUSTFLAGS": "-Ctarget-cpu=x86-64", + "LIBTORCH_USE_PYTORCH": "1", }, ) - self.assertEqual(build.environment["LIBTORCH_USE_PYTORCH"], "1") + self.assertNotIn("LIBTORCH_USE_PYTORCH", build.environment) + self.assertEqual(build.environment["LIBTORCH"], str(torch_root)) + self.assertEqual(build.environment["LIBTORCH_INCLUDE"], str(torch_root)) + self.assertEqual(build.environment["LIBTORCH_LIB"], str(torch_root)) + self.assertEqual(build.environment["LIBTORCH_CXX11_ABI"], "1") self.assertEqual(build.environment["LIBTORCH_BYPASS_VERSION_CHECK"], "1") self.assertIn(str(compat_header), build.environment["CXXFLAGS"]) self.assertIn( diff --git a/test/registered/unit/tools/test_cuda_ci_install_dependency.py b/test/registered/unit/tools/test_cuda_ci_install_dependency.py deleted file mode 100644 index d6679c6b5..000000000 --- a/test/registered/unit/tools/test_cuda_ci_install_dependency.py +++ /dev/null @@ -1,46 +0,0 @@ -import importlib.util -import re -import unittest -from pathlib import Path - -REPO_ROOT = Path(__file__).resolve().parents[4] -CI_REGISTER_PATH = REPO_ROOT / "python" / "sglang" / "test" / "ci" / "ci_register.py" -INSTALL_SCRIPT = REPO_ROOT / "scripts" / "ci" / "cuda" / "ci_install_dependency.sh" - - -def _load_module(name, path): - spec = importlib.util.spec_from_file_location(name, path) - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) - return module - - -register_cpu_ci = _load_module("ci_register", CI_REGISTER_PATH).register_cpu_ci -register_cpu_ci(est_time=0, suite="base-a-test-cpu") - - -class TestCudaCiInstallDependencyTorchCache(unittest.TestCase): - def test_rebuilds_torch_extensions_before_editable_install(self): - script = INSTALL_SCRIPT.read_text() - - self.assertRegex( - script, - re.compile( - r"setup_cargo_cache\s*\n" - r"\s*invalidate_torch_rust_cache\s*\n" - r"\s*install_sglang" - ), - ) - self.assertIn('"${SGLANG_BUILD_RUST_EXTS:-}" = "none"', script) - self.assertRegex( - script, - re.compile( - r"cargo clean --release --manifest-path " - r"\"\$\{REPO_ROOT\}/rust/sglang-radix-tree/Cargo\.toml\"\s*\\\n" - r"\s*-p torch-sys -p sglang-radix-tree" - ), - ) - - -if __name__ == "__main__": - unittest.main()