[CI] Size the CPU stage from the live partition model (#33329)

.github/workflows/pr-test.yml

@@ -334,13 +334,26 @@ jobs:
           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: 10
+        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 }} $CONTINUE_ON_ERROR_FLAG
+          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

@@ -35,29 +35,48 @@ HWBackend = _ci_register.HWBackend
 # pr-test-amd.yml / pr-test-npu.yml have their own dispatch.
 _TARGET_BACKENDS = {HWBackend.CUDA, HWBackend.CPU}
 
-# base-a is the critical-path entry gate; pin its fanout to sanity-coverage
+# Single-shard sanity gate on the critical path; pinned rather than sized
-# defaults instead of est_time. max_parallel = size (no throttle).
+# from est_time. max_parallel = size (no throttle).
 _BASE_A_OVERRIDES = {
-    "base-a-test-cpu": 8,
     "base-a-test-1-gpu-small": 1,
 }
 
 _REUSABLE_STAGE_USES = "./.github/workflows/_pr-test-stage.yml"
 
+# Inlined in pr-test.yml rather than dispatched through the reusable stage,
+# so there is no `run_timeout_minutes` input to read the budget from.
+_INLINE_SUITE_JOBS = {"base-a-test-cpu"}
+
 
 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.
-    Inline base-a-test-cpu is skipped (uses `_BASE_A_OVERRIDES`)."""
+    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 (wf.get("jobs") or {}).items():
+    for job_id, job in jobs.items():
         if not isinstance(job, dict) or job.get("uses") != _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:
+        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} "
@@ -176,7 +195,13 @@ def compute_partitions(
                     f"coeff={coeff}, bias={bias}s, total_est={total:.0f}s."
                 )
             size = max(1, ideal_size)
-            max_parallel = size if full_parallel else compute_max_parallel(size)
+            # The throttle rations scarce self-hosted GPU runners; hosted
+            # ubuntu-latest is elastic, so capping CPU shards only serializes
+            # them.
+            unthrottled = full_parallel or all(
+                t.backend == HWBackend.CPU for t in group
+            )
+            max_parallel = size if unthrottled else compute_max_parallel(size)
         result[suite] = {
             "size": size,
             "arr": list(range(size)),