From 414adef060f41977e3a4888cdb35cb0b1f4b8858 Mon Sep 17 00:00:00 2001 From: Mick Date: Sun, 20 Sep 2026 16:33:23 +0800 Subject: [PATCH] [CI] skip srt rust extension builds for diffusion-only PRs (#40293) Co-authored-by: Mick Qian --- .github/workflows/_pr-test-check-changes.yml | 1 + .github/workflows/pr-test-multimodal-gen.yml | 32 ----- .github/workflows/pr-test.yml | 17 ++- scripts/ci/cuda/ci_install_dependency.sh | 37 ++++-- scripts/ci/test_diffusion_rust_gate.py | 121 +++++++++++++++++++ 5 files changed, 163 insertions(+), 45 deletions(-) create mode 100644 scripts/ci/test_diffusion_rust_gate.py diff --git a/.github/workflows/_pr-test-check-changes.yml b/.github/workflows/_pr-test-check-changes.yml index 6355e6e98..4070b5ca2 100644 --- a/.github/workflows/_pr-test-check-changes.yml +++ b/.github/workflows/_pr-test-check-changes.yml @@ -89,6 +89,7 @@ jobs: - ".github/workflows/pr-gate.yml" - ".github/actions/**" - "python/pyproject.toml" + - "python/setup.py" - "python/sglang/!(multimodal_gen|kernels)/**/!(*.md)" - "python/sglang/kernels/!(*.md)" - "python/sglang/kernels/!(aot)/**/!(*.md)" diff --git a/.github/workflows/pr-test-multimodal-gen.yml b/.github/workflows/pr-test-multimodal-gen.yml index fb650d87e..1baa2c41d 100644 --- a/.github/workflows/pr-test-multimodal-gen.yml +++ b/.github/workflows/pr-test-multimodal-gen.yml @@ -9,10 +9,6 @@ on: sgl_kernel: required: true type: string - rust_ext_artifact: - description: 'Artifact of prebuilt Rust extension modules, from rust-ext-build. Empty, or a download that fails, falls back to the cache; a miss there compiles during install.' - type: string - default: '' runner_config: required: true type: string @@ -107,10 +103,6 @@ jobs: merge-multiple: true pattern: wheel-python3.10-cuda* - - uses: ./.github/actions/download-rust-ext - with: - artifact_name: ${{ inputs.rust_ext_artifact }} - - name: Install dependencies timeout-minutes: 20 run: | @@ -178,10 +170,6 @@ jobs: merge-multiple: true pattern: wheel-python3.10-cuda* - - uses: ./.github/actions/download-rust-ext - with: - artifact_name: ${{ inputs.rust_ext_artifact }} - - name: Install dependencies timeout-minutes: 20 run: | @@ -247,10 +235,6 @@ jobs: merge-multiple: true pattern: wheel-python3.10-cuda* - - uses: ./.github/actions/download-rust-ext - with: - artifact_name: ${{ inputs.rust_ext_artifact }} - - name: Install dependencies timeout-minutes: 20 run: | @@ -313,10 +297,6 @@ jobs: merge-multiple: true pattern: wheel-python3.10-cuda* - - uses: ./.github/actions/download-rust-ext - with: - artifact_name: ${{ inputs.rust_ext_artifact }} - - name: Install dependencies timeout-minutes: 20 run: | @@ -387,10 +367,6 @@ jobs: merge-multiple: true pattern: wheel-python3.10-cuda* - - uses: ./.github/actions/download-rust-ext - with: - artifact_name: ${{ inputs.rust_ext_artifact }} - - name: Install dependencies timeout-minutes: 20 run: | @@ -436,10 +412,6 @@ jobs: merge-multiple: true pattern: wheel-python3.10-cuda* - - uses: ./.github/actions/download-rust-ext - with: - artifact_name: ${{ inputs.rust_ext_artifact }} - - name: Install dependencies timeout-minutes: 20 env: @@ -505,10 +477,6 @@ jobs: merge-multiple: true pattern: wheel-python3.10-cuda* - - uses: ./.github/actions/download-rust-ext - with: - artifact_name: ${{ inputs.rust_ext_artifact }} - - name: Install dependencies timeout-minutes: 20 run: | diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index 5c6d1bd07..1b8ba5f38 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -210,8 +210,8 @@ jobs: # =============================================== Rust extensions ==================================================== - # Unlike sgl-kernel, not gated on source changes: the cache key is the source - # hash, so the job always runs but only compiles on a miss. + # diffusion-only PRs do not consume SRT Rust extensions; other runs seed or + # reuse the source-hash cache rust-ext-build: needs: [check-changes, call-gate] # Ungated on the event, unlike the jobs above: a cache entry is visible to @@ -223,6 +223,11 @@ jobs: if: | !cancelled() && needs.check-changes.result == 'success' && + (github.event_name != 'pull_request' || + needs.check-changes.outputs.multimodal_gen != 'true' || + needs.check-changes.outputs.main_package == 'true' || + needs.check-changes.outputs.jit_kernel == 'true' || + needs.check-changes.outputs.sgl_kernel == 'true') && (needs.call-gate.result == 'success' || needs.call-gate.result == 'skipped') uses: ./.github/workflows/_pr-test-rust-ext-build.yml with: @@ -240,6 +245,11 @@ jobs: if: | !cancelled() && needs.check-changes.result == 'success' && + (github.event_name != 'pull_request' || + needs.check-changes.outputs.multimodal_gen != 'true' || + needs.check-changes.outputs.main_package == 'true' || + needs.check-changes.outputs.jit_kernel == 'true' || + needs.check-changes.outputs.sgl_kernel == 'true') && (needs.call-gate.result == 'success' || needs.call-gate.result == 'skipped') uses: ./.github/workflows/_pr-test-rust-ext-build.yml with: @@ -387,7 +397,7 @@ jobs: secrets: inherit call-multimodal-gen-tests: - needs: [check-changes, call-gate, sgl-kernel-build-wheels, rust-ext-build] + needs: [check-changes, call-gate, sgl-kernel-build-wheels] if: | always() && !cancelled() && @@ -404,7 +414,6 @@ jobs: test_parallel_dispatch: ${{ inputs.test_parallel_dispatch == true && 'true' || 'false' }} caller_needs_failure: ${{ (needs.call-gate.result == 'failure' || needs.sgl-kernel-build-wheels.result == 'failure' || needs.check-changes.result == 'failure') && 'true' || 'false' }} skip_pr_test_health_check: ${{ (inputs.skip_pr_test_health_check == true || inputs.test_parallel_dispatch == true || inputs.run_all_tests == true) && 'true' || 'false' }} - rust_ext_artifact: ${{ needs.rust-ext-build.outputs.artifact_name }} secrets: inherit base-c-test-4-gpu-h100: diff --git a/scripts/ci/cuda/ci_install_dependency.sh b/scripts/ci/cuda/ci_install_dependency.sh index d69f3a253..b87569412 100755 --- a/scripts/ci/cuda/ci_install_dependency.sh +++ b/scripts/ci/cuda/ci_install_dependency.sh @@ -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]}" diff --git a/scripts/ci/test_diffusion_rust_gate.py b/scripts/ci/test_diffusion_rust_gate.py new file mode 100644 index 000000000..455844ac4 --- /dev/null +++ b/scripts/ci/test_diffusion_rust_gate.py @@ -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()