[CI] skip srt rust extension builds for diffusion-only PRs (#40293)

Co-authored-by: Mick Qian <mickqian@users.noreply.github.com>
This commit is contained in:
Mick
2026-09-20 16:33:23 +08:00
committed by GitHub
co-authored by Mick Qian
parent dc002c85fc
commit 414adef060
5 changed files with 163 additions and 45 deletions
+28 -9
View File
@@ -42,6 +42,14 @@ configure_environment() {
esac
OPTIONAL_DEPS="${1:-}"
if [ "$OPTIONAL_DEPS" = "diffusion" ]; then
export SGLANG_BUILD_RUST_EXTS=none
export SGLANG_RUST_BUILD_MODE=never
if [ -n "${GITHUB_ENV:-}" ]; then
echo "SGLANG_BUILD_RUST_EXTS=none" >> "$GITHUB_ENV"
echo "SGLANG_RUST_BUILD_MODE=never" >> "$GITHUB_ENV"
fi
fi
# Whether to create a uv venv (set USE_VENV=1). Default: 0.
USE_VENV="${USE_VENV:-0}"
@@ -275,6 +283,12 @@ clean_site_packages() {
rm -rf "$SITE_PACKAGES/sglang"
fi
# diffusion does not use the SRT Rust extensions
if [ "$OPTIONAL_DEPS" = "diffusion" ]; then
mark_step_done "${FUNCNAME[0]}"
return
fi
# Install protoc + Rust toolchain (needed by setuptools-rust, e.g. the native gRPC extension)
bash "${SCRIPT_DIR}/../utils/install_rust_protoc.sh"
export PATH="${CARGO_HOME:-$HOME/.cargo}/bin:${PATH}"
@@ -295,7 +309,7 @@ clean_site_packages() {
setup_cargo_cache() {
if [ "${SGLANG_BUILD_RUST_EXTS:-}" = "none" ]; then
echo "Using prebuilt Rust extensions; skipping Cargo target setup"
echo "Rust extension compilation disabled; skipping Cargo target setup"
mark_step_done "${FUNCNAME[0]}"
return
fi
@@ -462,6 +476,10 @@ install_pytorch_stack() {
}
require_prebuilt_rust_exts() {
if [ "$OPTIONAL_DEPS" = "diffusion" ]; then
mark_step_done "${FUNCNAME[0]}"
return
fi
# Stages whose download succeeded set this to none. Runs before
# setup_pip_toolchain uninstalls sglang, so clearing it here still reaches
# install_sglang below - setup.py reads it from the environment at build time.
@@ -787,7 +805,7 @@ verify_imports() {
# One process; torch/cutlass do not import sglang, so the find_spec check
# still runs ahead of any sglang import.
SGLANG_EXPECTED_INIT="${REPO_ROOT}/python/sglang/__init__.py" python3 -c '
SGLANG_CI_OPTIONAL_DEPS="$OPTIONAL_DEPS" SGLANG_EXPECTED_INIT="${REPO_ROOT}/python/sglang/__init__.py" python3 -c '
import ctypes
import importlib.metadata
import os
@@ -830,13 +848,14 @@ print(f"sglang resolves to {spec.origin}")
# Import, not find_spec: the finders locate an extension without dlopening it,
# so a .so that cannot load passes find_spec and only fails inside some suite.
import importlib
for mod in ("server", "grpc", "multimodal"):
name = f"sglang.srt.rust_extensions._{mod}"
try:
importlib.import_module(name)
except Exception as exc:
raise SystemExit(f"{name} is present but does not load: {exc!r}")
print(f"{name} loads")
if os.environ["SGLANG_CI_OPTIONAL_DEPS"] != "diffusion":
for mod in ("server", "grpc", "multimodal"):
name = f"sglang.srt.rust_extensions._{mod}"
try:
importlib.import_module(name)
except Exception as exc:
raise SystemExit(f"{name} is present but does not load: {exc!r}")
print(f"{name} loads")
'
mark_step_done "${FUNCNAME[0]}"
+121
View File
@@ -0,0 +1,121 @@
import os
import re
import subprocess
import tempfile
import unittest
from pathlib import Path
import yaml
ROOT = Path(__file__).resolve().parents[2]
class TestDiffusionRustGate(unittest.TestCase):
def test_build_gate(self):
jobs = yaml.safe_load((ROOT / ".github/workflows/pr-test.yml").read_text())[
"jobs"
]
for name in ("rust-ext-build", "rust-ext-build-aarch64"):
for event in ("pull_request", "schedule", "workflow_dispatch"):
for main, jit, kernel, diffusion in (
(False, False, False, True),
(True, False, False, True),
(False, True, False, True),
(False, False, True, True),
(False, False, False, False),
):
with self.subTest(
job=name, event=event, flags=(main, jit, kernel, diffusion)
):
expression = jobs[name]["if"]
values = {
"github.event_name": event,
"needs.check-changes.result": "success",
"needs.call-gate.result": "success",
"needs.check-changes.outputs.main_package": str(
main
).lower(),
"needs.check-changes.outputs.jit_kernel": str(jit).lower(),
"needs.check-changes.outputs.sgl_kernel": str(
kernel
).lower(),
"needs.check-changes.outputs.multimodal_gen": str(
diffusion
).lower(),
}
for key, value in values.items():
expression = expression.replace(key, repr(value))
expression = expression.replace("!cancelled()", "True")
expression = expression.replace("&&", " and ").replace(
"||", " or "
)
actual = eval(
" ".join(expression.split()), {"__builtins__": {}}
)
expected = (
event != "pull_request"
or not diffusion
or main
or jit
or kernel
)
self.assertEqual(actual, expected)
caller = jobs["call-multimodal-gen-tests"]
self.assertNotIn("rust-ext-build", caller["needs"])
self.assertNotIn("rust_ext_artifact", caller["with"])
callee = (ROOT / ".github/workflows/pr-test-multimodal-gen.yml").read_text()
self.assertNotIn("download-rust-ext", callee)
self.assertNotIn("rust_ext_artifact", callee)
def test_installer_preserves_srt_fallback(self):
source = (ROOT / "scripts/ci/cuda/ci_install_dependency.sh").read_text()
# exercise the actual installer functions, replacing only package-manager I/O
functions = "\n".join(
re.search(rf"^{name}\(\) \{{\n.*?^\}}", source, re.M | re.S)[0]
for name in (
"configure_environment",
"require_prebuilt_rust_exts",
"setup_cargo_cache",
)
)
for extra, expected in (("diffusion", "none:never"), ("", ":auto")):
with self.subTest(extra=extra), tempfile.TemporaryDirectory() as tmp:
script = (
functions
+ """
mark_step_done() { :; }
python3() { if [ "$1" = "-c" ]; then echo .test.so; fi; }
uv() { :; }
pip() { :; }
configure_environment "$1"
require_prebuilt_rust_exts
if [ "$1" = diffusion ]; then setup_cargo_cache; fi
printf 'RESULT=%s:%s\n' "$SGLANG_BUILD_RUST_EXTS" "$SGLANG_RUST_BUILD_MODE"
"""
)
env = dict(
os.environ,
GITHUB_ENV=f"{tmp}/env",
USE_VENV="0",
SGLANG_BUILD_RUST_EXTS="none",
SGLANG_RUST_BUILD_MODE="never",
)
result = subprocess.run(
["bash", "-eu", "-c", script, "test", extra],
cwd=tmp,
env=env,
capture_output=True,
text=True,
check=True,
)
self.assertIn(f"RESULT={expected}", result.stdout)
if extra == "diffusion":
self.assertIn(
"SGLANG_BUILD_RUST_EXTS=none",
Path(env["GITHUB_ENV"]).read_text(),
)
if __name__ == "__main__":
unittest.main()