[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:
@@ -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]}"
|
||||
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user