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

Co-authored-by: Mick Qian <mickqian@users.noreply.github.com>

.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)"

.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: |

.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
+  # diffusion-only PRs do not consume SRT Rust extensions; other runs seed or
-  # hash, so the job always runs but only compiles on a miss.
+  # 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:

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"):
+if os.environ["SGLANG_CI_OPTIONAL_DEPS"] != "diffusion":
-    name = f"sglang.srt.rust_extensions._{mod}"
+    for mod in ("server", "grpc", "multimodal"):
-    try:
+        name = f"sglang.srt.rust_extensions._{mod}"
-        importlib.import_module(name)
+        try:
-    except Exception as exc:
+            importlib.import_module(name)
-        raise SystemExit(f"{name} is present but does not load: {exc!r}")
+        except Exception as exc:
-    print(f"{name} loads")
+            raise SystemExit(f"{name} is present but does not load: {exc!r}")
+        print(f"{name} loads")
 '
 
     mark_step_done "${FUNCNAME[0]}"

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()