[CI] Dispatch base-a-test-cpu through its own reusable stage workflow (#33461)

.github/workflows/_pr-test-check-changes.yml

@@ -79,6 +79,7 @@ jobs:
           filters: |
             main_package:
               - ".github/workflows/pr-test.yml"
+              - ".github/workflows/_pr-test-*.yml"
               - ".github/workflows/pr-gate.yml"
               - ".github/actions/**"
               - "python/pyproject.toml"

.github/workflows/_pr-test-stage-cpu.yml

@@ -0,0 +1,149 @@
+name: PR Test Stage (CPU)
+# Reusable workflow for the hosted-runner CPU test stage. Mirrors
+# _pr-test-stage.yml's input contract -- `check_changes` / `caller_inputs`
+# bundles, with `partitions` forwarded separately so matrix expressions stay
+# single-fromJson -- but keeps its own uv pip / protoc / rust-cache install,
+# which shares no step with the self-hosted GPU stages.
+
+on:
+  workflow_call:
+    inputs:
+      self_name:
+        description: 'Caller job key; used for partitions[suite] lookup and the suite name.'
+        type: string
+        required: true
+      check_changes:
+        description: 'toJson(needs.check-changes.outputs). Read via fromJson(...).main_package / continue_on_error / partition_model_sha.'
+        type: string
+        required: true
+      caller_inputs:
+        description: 'toJson(inputs) from pr-test.yml. Read via fromJson(...).git_ref / skip_pr_test_health_check / test_parallel_dispatch / run_all_tests.'
+        type: string
+        required: true
+      partitions:
+        description: 'check-changes.outputs.partitions raw -- kept separate to avoid double-fromJson in matrix expressions.'
+        type: string
+        required: true
+      run_timeout_minutes:
+        description: 'timeout-minutes for the Run test step. Required so compute_partitions.py can read it from pr-test.yml without a duplicated default constant.'
+        type: string
+        required: true
+      rust_ext_artifact:
+        description: 'Artifact of prebuilt Rust extension modules, from rust-ext-build. Empty, or a download that fails, means this stage compiles them during install.'
+        type: string
+        default: ''
+
+# Reusable workflows do NOT inherit the caller's workflow-level env across the
+# workflow_call boundary, so pr-test.yml's env is redeclared here -- minus the
+# CUDA-only entries it also sets, which have no consumer on a hosted CPU runner.
+env:
+  SGLANG_IS_IN_CI: true
+  SKIP_PR_TEST_HEALTH_CHECK: ${{ (fromJson(inputs.caller_inputs).skip_pr_test_health_check || fromJson(inputs.caller_inputs).test_parallel_dispatch || fromJson(inputs.caller_inputs).run_all_tests) && 'true' || 'false' }}
+  PR_TEST_BYPASS_MAINTENANCE_ON_MAIN: ${{ github.ref == 'refs/heads/main' && 'true' || 'false' }}
+  USE_VENV: false
+
+jobs:
+  run:
+    name: ${{ inputs.self_name }} (${{ matrix.partition }})
+    # Gated here rather than on the caller so a main_package-less run still
+    # emits a skipped entry for wait-for-jobs to match. The schedule /
+    # parallel-dispatch / !failure() half stays on the caller, where the
+    # `needs` context it guards actually lives.
+    if: fromJson(inputs.check_changes).main_package == 'true'
+    runs-on: ubuntu-latest
+    timeout-minutes: 240
+    env:
+      HF_HOME: ${{ github.workspace }}/.hf-cache
+    strategy:
+      fail-fast: false
+      max-parallel: ${{ fromJson(inputs.partitions)[inputs.self_name].max_parallel }}
+      matrix:
+        partition: ${{ fromJson(inputs.partitions)[inputs.self_name].arr }}
+    steps:
+      - name: Free disk space
+        run: |
+          sudo rm -rf /usr/share/dotnet /usr/local/lib/android /opt/ghc
+          df -h
+
+      - name: Checkout code
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ fromJson(inputs.caller_inputs).git_ref || github.sha }}
+
+      - uses: ./.github/actions/check-pr-test-health
+
+      - uses: ./.github/actions/check-maintenance
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.10'
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v5
+
+      # This stage compiled the workspace too - 7+ minutes per partition on
+      # billable hosted minutes. rust-ext-build's modules need an older glibc than
+      # this runner has, which is the safe direction, and both pin Python 3.10.
+      - name: Download prebuilt Rust extensions
+        id: rust_ext
+        if: ${{ inputs.rust_ext_artifact != '' }}
+        continue-on-error: true
+        uses: actions/download-artifact@v4
+        with:
+          name: ${{ inputs.rust_ext_artifact }}
+          path: python/sglang/srt/
+
+      # Both only serve the fallback where this stage compiles the extensions
+      # itself, so they follow the download's outcome, not the job output: an
+      # expired artifact still needs cargo and a warm target dir here. Otherwise
+      # rust-cache restores ~1 GB per partition for nothing, on an over-quota cache.
+      - name: Install protoc + Rust toolchain
+        if: ${{ steps.rust_ext.outcome != 'success' }}
+        timeout-minutes: 10
+        run: bash scripts/ci/utils/install_rust_protoc.sh
+
+      - name: Rust cache (rust/ workspace)
+        if: ${{ steps.rust_ext.outcome != 'success' }}
+        uses: Swatinem/rust-cache@v2
+        with:
+          workspaces: rust
+          shared-key: "sglang-grpc-cpu"
+          save-if: ${{ matrix.partition == 0 }}
+
+      # uv pip targets a venv by default; setup-python has no venv — install into that interpreter (see UV_SYSTEM_PYTHON in https://docs.astral.sh/uv/guides/integration/github/)
+      - name: Install dependencies
+        timeout-minutes: 20
+        env:
+          UV_SYSTEM_PYTHON: "1"
+          SGLANG_BUILD_RUST_EXTS: ${{ steps.rust_ext.outcome == 'success' && 'none' || '' }}
+        run: |
+          uv pip install -e "python[dev]" --index-strategy unsafe-best-match --prerelease allow
+
+      # Hosted runners are ephemeral, so models are re-fetched every run and the
+      # Hub occasionally returns 429s. Persist the HF cache in GitHub's cache
+      # storage (rolling key + restore-keys) so warm runs never hit the network.
+      - name: Cache HF hub
+        uses: actions/cache@v4
+        with:
+          path: ${{ github.workspace }}/.hf-cache
+          key: hf-cpu-${{ matrix.partition }}-${{ github.run_id }}
+          restore-keys: hf-cpu-${{ matrix.partition }}-
+
+      # Pinned SHA so every shard splits against the same snapshot; without
+      # the file run_suite falls back to the (drifting) in-source est_time.
+      - name: Fetch live partition model
+        if: fromJson(inputs.check_changes).partition_model_sha != ''
+        run: |
+          rm -f /tmp/partition-model.json
+          URL="https://raw.githubusercontent.com/sgl-project/sglang-ci-stats/${{ fromJson(inputs.check_changes).partition_model_sha }}/model.json"
+          curl --fail --silent --show-error --max-time 15 --retry 3 --retry-delay 2 \
+            "$URL" -o /tmp/partition-model.json
+
+      - name: Run test
+        timeout-minutes: ${{ fromJson(inputs.run_timeout_minutes) }}
+        env:
+          CONTINUE_ON_ERROR_FLAG: ${{ fromJson(inputs.check_changes).continue_on_error == 'true' && '--continue-on-error' || '' }}
+        run: |
+          cd test/
+          python3 run_suite.py --hw cpu --suite ${{ inputs.self_name }} --auto-partition-id ${{ matrix.partition }} --auto-partition-size ${{ fromJson(inputs.partitions)[inputs.self_name].size }} --partition-model-file /tmp/partition-model.json $CONTINUE_ON_ERROR_FLAG

.github/workflows/pr-test.yml

@@ -295,108 +295,17 @@ jobs:
     if: |
       always() &&
       needs.check-changes.result == 'success' &&
-      ((github.event_name == 'schedule' || inputs.test_parallel_dispatch == true) || (!failure() && !cancelled())) &&
+      ((github.event_name == 'schedule' || inputs.test_parallel_dispatch == true) || (!failure() && !cancelled()))
-      (needs.check-changes.outputs.main_package == 'true')
+    uses: ./.github/workflows/_pr-test-stage-cpu.yml
-    runs-on: ubuntu-latest
+    with:
-    timeout-minutes: 240
+      self_name: base-a-test-cpu
-    env:
+      check_changes: ${{ toJson(needs.check-changes.outputs) }}
-      HF_HOME: ${{ github.workspace }}/.hf-cache
+      caller_inputs: ${{ toJson(inputs) }}
-    strategy:
+      partitions: ${{ needs.check-changes.outputs.partitions }}
-      fail-fast: false
+      run_timeout_minutes: '15'
-      max-parallel: ${{ fromJson(needs.check-changes.outputs.partitions)['base-a-test-cpu'].max_parallel }}
+      rust_ext_artifact: ${{ needs.rust-ext-build.outputs.artifact_name }}
-      matrix:
+    # No `secrets: inherit`: this stage has no secret consumer, unlike the GPU
-        partition: ${{ fromJson(needs.check-changes.outputs.partitions)['base-a-test-cpu'].arr }}
+    # stages' coredump upload. GITHUB_TOKEN and permissions inherit regardless.
-    steps:
-      - name: Free disk space
-        run: |
-          sudo rm -rf /usr/share/dotnet /usr/local/lib/android /opt/ghc
-          df -h
-
-      - name: Checkout code
-        uses: actions/checkout@v4
-        with:
-          ref: ${{ inputs.git_ref || github.sha }}
-
-      - uses: ./.github/actions/check-pr-test-health
-
-      - uses: ./.github/actions/check-maintenance
-
-      - name: Set up Python
-        uses: actions/setup-python@v5
-        with:
-          python-version: '3.10'
-
-      - name: Install uv
-        uses: astral-sh/setup-uv@v5
-
-      # This stage compiled the workspace too - 7+ minutes per partition on
-      # billable hosted minutes. rust-ext-build's modules need an older glibc than
-      # this runner has, which is the safe direction, and both pin Python 3.10.
-      - name: Download prebuilt Rust extensions
-        id: rust_ext
-        if: ${{ needs.rust-ext-build.outputs.artifact_name != '' }}
-        continue-on-error: true
-        uses: actions/download-artifact@v4
-        with:
-          name: ${{ needs.rust-ext-build.outputs.artifact_name }}
-          path: python/sglang/srt/
-
-      # Both only serve the fallback where this stage compiles the extensions
-      # itself, so they follow the download's outcome, not the job output: an
-      # expired artifact still needs cargo and a warm target dir here. Otherwise
-      # rust-cache restores ~1 GB per partition for nothing, on an over-quota cache.
-      - name: Install protoc + Rust toolchain
-        if: ${{ steps.rust_ext.outcome != 'success' }}
-        timeout-minutes: 10
-        run: bash scripts/ci/utils/install_rust_protoc.sh
-
-      - name: Rust cache (rust/ workspace)
-        if: ${{ steps.rust_ext.outcome != 'success' }}
-        uses: Swatinem/rust-cache@v2
-        with:
-          workspaces: rust
-          shared-key: "sglang-grpc-cpu"
-          save-if: ${{ matrix.partition == 0 }}
-
-      # uv pip targets a venv by default; setup-python has no venv — install into that interpreter (see UV_SYSTEM_PYTHON in https://docs.astral.sh/uv/guides/integration/github/)
-      - name: Install dependencies
-        timeout-minutes: 20
-        env:
-          UV_SYSTEM_PYTHON: "1"
-          SGLANG_BUILD_RUST_EXTS: ${{ steps.rust_ext.outcome == 'success' && 'none' || '' }}
-        run: |
-          uv pip install -e "python[dev]" --index-strategy unsafe-best-match --prerelease allow
-
-      # Hosted runners are ephemeral, so models are re-fetched every run and the
-      # Hub occasionally returns 429s. Persist the HF cache in GitHub's cache
-      # storage (rolling key + restore-keys) so warm runs never hit the network.
-      - name: Cache HF hub
-        uses: actions/cache@v4
-        with:
-          path: ${{ github.workspace }}/.hf-cache
-          key: hf-cpu-${{ matrix.partition }}-${{ github.run_id }}
-          restore-keys: hf-cpu-${{ matrix.partition }}-
-
-      # Pinned SHA so every shard splits against the same snapshot; without
-      # the file run_suite falls back to the (drifting) in-source est_time.
-      - name: Fetch live partition model
-        if: needs.check-changes.outputs.partition_model_sha != ''
-        run: |
-          rm -f /tmp/partition-model.json
-          URL="https://raw.githubusercontent.com/sgl-project/sglang-ci-stats/${{ needs.check-changes.outputs.partition_model_sha }}/model.json"
-          curl --fail --silent --show-error --max-time 15 --retry 3 --retry-delay 2 \
-            "$URL" -o /tmp/partition-model.json
-
-      # compute_partitions reads this back as the per-shard budget, so it
-      # drives the fanout rather than capping it -- shrinking it buys more
-      # shards, not shorter ones.
-      - name: Run test
-        timeout-minutes: 15
-        env:
-          CONTINUE_ON_ERROR_FLAG: ${{ needs.check-changes.outputs.continue_on_error == 'true' && '--continue-on-error' || '' }}
-        run: |
-          cd test/
-          python3 run_suite.py --hw cpu --suite base-a-test-cpu --auto-partition-id ${{ matrix.partition }} --auto-partition-size ${{ fromJson(needs.check-changes.outputs.partitions)['base-a-test-cpu'].size }} --partition-model-file /tmp/partition-model.json $CONTINUE_ON_ERROR_FLAG
 
   # Runs on 5090 (32GB, SM120)
   base-b-test-1-gpu-small:

scripts/ci/utils/compute_partitions.py

@@ -41,46 +41,31 @@ _BASE_A_OVERRIDES = {
     "base-a-test-1-gpu-small": 1,
 }
 
-_REUSABLE_STAGE_USES = "./.github/workflows/_pr-test-stage.yml"
+# Every stage dispatches through one of these; CPU has its own because its
-
+# hosted-runner install shares no step with the self-hosted GPU stages.
-# Inlined in pr-test.yml rather than dispatched through the reusable stage,
+_REUSABLE_STAGE_USES = {
-# so there is no `run_timeout_minutes` input to read the budget from.
+    "./.github/workflows/_pr-test-stage.yml",
-_INLINE_SUITE_JOBS = {"base-a-test-cpu"}
+    "./.github/workflows/_pr-test-stage-cpu.yml",
+}
 
 
 def load_run_timeouts(pr_test_yml_path: str) -> dict:
     """Map `self_name -> run_timeout_minutes` from one pr-test*.yml. The input
-    is required in `_pr-test-stage.yml` -- KeyError surfaces missing.
+    is required in both reusable stage workflows -- KeyError surfaces missing."""
-    Inline suites (`_INLINE_SUITE_JOBS`) contribute their `Run test` step
-    timeout so they size off the same budget as dispatched stages."""
     with open(pr_test_yml_path) as f:
         wf = yaml.safe_load(f)
     jobs = wf.get("jobs") or {}
     timeouts = {}
     for job_id, job in jobs.items():
-        if not isinstance(job, dict) or job.get("uses") != _REUSABLE_STAGE_USES:
+        if not isinstance(job, dict) or job.get("uses") not in _REUSABLE_STAGE_USES:
             continue
         with_ = job.get("with") or {}
         suite = with_.get("self_name", job_id)
         timeouts[suite] = int(with_["run_timeout_minutes"])
-    for suite in _INLINE_SUITE_JOBS.intersection(jobs):
-        budgets = [
-            s["timeout-minutes"]
-            for s in ((jobs.get(suite) or {}).get("steps") or [])
-            if isinstance(s, dict)
-            and s.get("name") == "Run test"
-            and "timeout-minutes" in s
-        ]
-        if len(budgets) != 1:
-            raise RuntimeError(
-                f"load_run_timeouts: inline suite {suite!r} needs exactly one "
-                f"`Run test` step with `timeout-minutes` in {pr_test_yml_path}."
-            )
-        timeouts[suite] = int(budgets[0])
     if not timeouts:
         raise RuntimeError(
-            f"load_run_timeouts: no jobs matched uses={_REUSABLE_STAGE_USES!r} "
+            f"load_run_timeouts: no jobs matched uses in {_REUSABLE_STAGE_USES!r} "
-            f"in {pr_test_yml_path}. The reusable workflow path likely "
+            f"in {pr_test_yml_path}. A reusable workflow path likely "
             "changed -- update _REUSABLE_STAGE_USES."
         )
     return timeouts
@@ -89,7 +74,10 @@ def load_run_timeouts(pr_test_yml_path: str) -> dict:
 def per_shard_target_seconds(suite: str, run_timeouts: dict) -> float:
     """Per-shard wall budget = 0.75 * stage timeout. 0.75 is the inverse
     of LPT's 4/3 worst-case approximation ratio, so the most imbalanced
-    LPT shard fills exactly the timeout."""
+    LPT shard fills exactly the timeout.
+
+    A stage's `run_timeout_minutes` therefore drives its fanout rather than
+    capping it: shrinking it buys more shards, not shorter ones."""
     return 0.75 * run_timeouts[suite] * 60