ci: dynamic partition + LPT from live sglang-ci-stats model (#25263)

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

@@ -34,6 +34,8 @@ on:
         value: ${{ jobs.run.outputs.multimodal_gen }}
       partitions:
         value: ${{ jobs.run.outputs.partitions }}
+      partition_model_sha:
+        value: ${{ jobs.run.outputs.partition_model_sha }}
       b200_runner:
         value: ${{ jobs.run.outputs.b200_runner }}
       enable_retry:
@@ -60,6 +62,7 @@ jobs:
       jit_kernel: ${{ steps.filter-api.outputs.jit_kernel || steps.filter.outputs.jit_kernel || steps.run-mode.outputs.run_all_tests }}
       multimodal_gen: ${{ steps.filter-api.outputs.multimodal_gen || steps.filter.outputs.multimodal_gen || steps.run-mode.outputs.run_all_tests }}
       partitions: ${{ steps.partitions.outputs.partitions }}
+      partition_model_sha: ${{ steps.partition-model-sha.outputs.sha }}
       b200_runner: ${{ steps.set-runner.outputs.b200_runner }}
       enable_retry: ${{ steps.set-retry.outputs.enable_retry }}
       continue_on_error: ${{ steps.set-continue-on-error.outputs.continue_on_error }}
@@ -233,6 +236,33 @@ jobs:
           fi
           echo "full=$FULL" >> "$GITHUB_OUTPUT"
 
+      - name: Resolve sglang-ci-stats SHA
+        id: partition-model-sha
+        env:
+          GH_TOKEN: ${{ github.token }}
+        run: |
+          # Pin all shards to one immutable commit so dispatch and every
+          # runtime LPT use the same model snapshot. Soft fail -> static.
+          SHA=$(gh api repos/sgl-project/sglang-ci-stats/commits/main --jq '.sha' 2>/dev/null || true)
+          if [[ -n "$SHA" ]]; then
+            echo "Pinned sglang-ci-stats@$SHA"
+            echo "sha=$SHA" >> "$GITHUB_OUTPUT"
+          else
+            echo "::warning::Could not resolve sglang-ci-stats SHA; using in-source est_time"
+            echo "sha=" >> "$GITHUB_OUTPUT"
+          fi
+
+      - name: Fetch live partition model
+        if: steps.partition-model-sha.outputs.sha != ''
+        run: |
+          # SHA resolved -> require fetch (curl --retry). Soft fallback
+          # would risk cross-shard LPT divergence on transient curl flake.
+          rm -f /tmp/partition-model.json
+          URL="https://raw.githubusercontent.com/sgl-project/sglang-ci-stats/${{ steps.partition-model-sha.outputs.sha }}/model.json"
+          curl --fail --silent --show-error --max-time 15 --retry 3 --retry-delay 2 \
+            "$URL" -o /tmp/partition-model.json
+          echo "Fetched partition-model.json ($(wc -c < /tmp/partition-model.json) bytes)"
+
       - name: Compute partitions
         id: partitions
         run: |
@@ -243,6 +273,7 @@ jobs:
           # See scripts/ci/utils/compute_partitions.py.
           python3 scripts/ci/utils/compute_partitions.py \
             --full-parallel ${{ steps.parallel-mode.outputs.full }} \
+            --partition-model-file /tmp/partition-model.json \
             >> "$GITHUB_OUTPUT"
 
       - name: Set B200 runner tag

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

@@ -35,9 +35,9 @@ on:
         type: string
         required: true
       run_timeout_minutes:
-        description: 'timeout-minutes for the Run test step.'
+        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
-        default: '30'
+        required: true
       timeout_per_file:
         description: 'run_suite.py --timeout-per-file value (empty = unset).'
         type: string
@@ -147,6 +147,16 @@ jobs:
           [ -f "${SGLANG_CI_VENV_PATH:-/dev/null}/env.sh" ] && source "${SGLANG_CI_VENV_PATH}/env.sh"
           python3 scripts/ci/cuda/warmup_server.py ${{ inputs.warmup_server_models }}
 
+      - name: Fetch live partition model
+        if: fromJson(inputs.check_changes).partition_model_sha != ''
+        run: |
+          # SHA resolved by check-changes -> require fetch (curl --retry)
+          # so all shards stay on the same snapshot.
+          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:
@@ -156,6 +166,7 @@ jobs:
           python3 run_suite.py --hw cuda --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 \
             ${{ inputs.timeout_per_file && format('--timeout-per-file {0}', inputs.timeout_per_file) || '' }} \
             $CONTINUE_ON_ERROR_FLAG
 

.github/workflows/pr-test.yml

@@ -372,6 +372,7 @@ jobs:
       check_changes: ${{ toJson(needs.check-changes.outputs) }}
       caller_inputs: ${{ toJson(inputs) }}
       partitions: ${{ needs.check-changes.outputs.partitions }}
+      run_timeout_minutes: '30'
     secrets: inherit
 
   # Runs on H100 (80GB, SM90) - tests that don't pass on 5090 (FA3, FP8, high VRAM, etc.)
@@ -386,6 +387,7 @@ jobs:
       check_changes: ${{ toJson(needs.check-changes.outputs) }}
       caller_inputs: ${{ toJson(inputs) }}
       partitions: ${{ needs.check-changes.outputs.partitions }}
+      run_timeout_minutes: '30'
       timeout_per_file: '1800'
     secrets: inherit
 
@@ -400,6 +402,7 @@ jobs:
       check_changes: ${{ toJson(needs.check-changes.outputs) }}
       caller_inputs: ${{ toJson(inputs) }}
       partitions: ${{ needs.check-changes.outputs.partitions }}
+      run_timeout_minutes: '30'
     secrets: inherit
 
   stage-b-test-4-gpu-b200:
@@ -461,6 +464,7 @@ jobs:
       check_changes: ${{ toJson(needs.check-changes.outputs) }}
       caller_inputs: ${{ toJson(inputs) }}
       partitions: ${{ needs.check-changes.outputs.partitions }}
+      run_timeout_minutes: '30'
     secrets: inherit
 
   stage-c-test-8-gpu-h200:
@@ -474,6 +478,7 @@ jobs:
       check_changes: ${{ toJson(needs.check-changes.outputs) }}
       caller_inputs: ${{ toJson(inputs) }}
       partitions: ${{ needs.check-changes.outputs.partitions }}
+      run_timeout_minutes: '30'
       # Per-model TP must match the test's launch in test/registered/ -- see
       # FALLBACK_ARGS in scripts/ci/cuda/warmup_deep_gemm.py for extra dp/ep
       # flags. Only models that actually invoke DeepGEMM kernels at runtime
@@ -494,6 +499,7 @@ jobs:
       check_changes: ${{ toJson(needs.check-changes.outputs) }}
       caller_inputs: ${{ toJson(inputs) }}
       partitions: ${{ needs.check-changes.outputs.partitions }}
+      run_timeout_minutes: '30'
     secrets: inherit
 
   stage-c-test-deepep-4-gpu-h100:
@@ -507,6 +513,7 @@ jobs:
       check_changes: ${{ toJson(needs.check-changes.outputs) }}
       caller_inputs: ${{ toJson(inputs) }}
       partitions: ${{ needs.check-changes.outputs.partitions }}
+      run_timeout_minutes: '30'
       warmup_deep_gemm_models: 'lmsys/sglang-ci-dsv3-test:4'
       warmup_server_models: 'lmsys/sglang-ci-dsv3-test:4'
     secrets: inherit
@@ -538,6 +545,7 @@ jobs:
       check_changes: ${{ toJson(needs.check-changes.outputs) }}
       caller_inputs: ${{ toJson(inputs) }}
       partitions: ${{ needs.check-changes.outputs.partitions }}
+      run_timeout_minutes: '30'
       timeout_per_file: '1800'
     secrets: inherit
 
@@ -552,6 +560,7 @@ jobs:
       check_changes: ${{ toJson(needs.check-changes.outputs) }}
       caller_inputs: ${{ toJson(inputs) }}
       partitions: ${{ needs.check-changes.outputs.partitions }}
+      run_timeout_minutes: '45'
       timeout_per_file: '1800'
     secrets: inherit
 
@@ -566,6 +575,7 @@ jobs:
       check_changes: ${{ toJson(needs.check-changes.outputs) }}
       caller_inputs: ${{ toJson(inputs) }}
       partitions: ${{ needs.check-changes.outputs.partitions }}
+      run_timeout_minutes: '30'
       timeout_per_file: '1800'
     secrets: inherit
 

python/sglang/test/ci/ci_register.py

@@ -287,17 +287,30 @@ def ut_parse_one_file(filename: str) -> Tuple[List[CIRegistry], bool]:
     return visitor.registries, visitor.has_main_entry
 
 
-def auto_partition(files: List[CIRegistry], rank: int, size: int) -> List[CIRegistry]:
+def auto_partition(
+    files: List[CIRegistry],
+    rank: int,
+    size: int,
+    live_est: Optional[dict] = None,
+) -> List[CIRegistry]:
     """Partition files into `size` sublists with approximately equal sums of
     estimated times using a greedy algorithm (LPT heuristic), and return the
     partition for the specified rank.
+
+    `live_est`: optional `filename -> est seconds` overrides; missing
+    files fall back to in-source `est_time`.
     """
     if not files or size <= 0:
         return []
 
+    def est_of(f: CIRegistry) -> float:
+        if live_est is not None and f.filename in live_est:
+            return live_est[f.filename]
+        return f.est_time
+
     # Sort by estimated_time descending; filename as tie-breaker for
     # deterministic partitioning regardless of glob ordering.
-    sorted_files = sorted(files, key=lambda f: (-f.est_time, f.filename))
+    sorted_files = sorted(files, key=lambda f: (-est_of(f), f.filename))
 
     partitions: List[List[CIRegistry]] = [[] for _ in range(size)]
     partition_sums = [0.0] * size
@@ -306,7 +319,7 @@ def auto_partition(files: List[CIRegistry], rank: int, size: int) -> List[CIRegi
     for file in sorted_files:
         min_sum_idx = min(range(size), key=partition_sums.__getitem__)
         partitions[min_sum_idx].append(file)
-        partition_sums[min_sum_idx] += file.est_time
+        partition_sums[min_sum_idx] += est_of(file)
 
     if rank < size:
         return partitions[rank]

scripts/ci/utils/compute_partitions.py

@@ -1,4 +1,4 @@
-"""Sum CIRegistry est_time per per-commit suite and emit one $GITHUB_OUTPUT line
+"""Sum est_time per per-commit suite and emit one $GITHUB_OUTPUT line
 keyed by suite name. Consumed by pr-test.yml stage jobs as
 `fromJson(needs.check-changes.outputs.partitions)['<suite>']`.
 
@@ -13,6 +13,8 @@ import math
 import os
 from collections import defaultdict
 
+import yaml  # PyYAML; preinstalled on ubuntu-latest GHA runners.
+
 REPO_ROOT = os.path.dirname(
     os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 )
@@ -39,16 +41,36 @@ _STAGE_A_OVERRIDES = {
     "stage-a-test-1-gpu-small": 1,
 }
 
-# Per-partition wall-clock target. ~20 min avg naive; worst-case LPT 4/3
-# imbalance is ~27 min, still below the 30-min job-level timeout that acts
-# as the real safety net. No LPT slop applied — we lean on the runtime
-# timeout + the explicit MAX_PARTITION_SECONDS sanity check rather than
-# padding partition count.
-TARGET_SECONDS = 20 * 60
+_REUSABLE_STAGE_USES = "./.github/workflows/_pr-test-stage.yml"
 
-# Hard ceiling. Exceeded → raise, forcing the maintainer to split a slow file
-# or bump TARGET_SECONDS deliberately.
-MAX_PARTITION_SECONDS = 30 * 60
+
+def load_run_timeouts(pr_test_yml_path: str) -> dict:
+    """Map `self_name -> run_timeout_minutes` from pr-test.yml. The input
+    is required in `_pr-test-stage.yml` -- KeyError surfaces missing.
+    Inline stage-a-test-cpu is skipped (uses `_STAGE_A_OVERRIDES`)."""
+    with open(pr_test_yml_path) as f:
+        wf = yaml.safe_load(f)
+    timeouts = {}
+    for job_id, job in (wf.get("jobs") or {}).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"])
+    if not timeouts:
+        raise RuntimeError(
+            f"load_run_timeouts: no jobs matched uses={_REUSABLE_STAGE_USES!r} "
+            f"in {pr_test_yml_path}. The reusable workflow path likely "
+            "changed -- update _REUSABLE_STAGE_USES."
+        )
+    return timeouts
+
+
+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."""
+    return 0.75 * run_timeouts[suite] * 60
 
 
 def discover_files(repo_root: str) -> list[str]:
@@ -70,43 +92,85 @@ def discover_files(repo_root: str) -> list[str]:
     return files
 
 
+def load_partition_model(path):
+    """Read sglang-ci-stats' model.json; None on missing/unparsable.
+    Cross-repo schema -- guard against non-dict top-level."""
+    if not path or not os.path.exists(path):
+        return None
+    try:
+        with open(path) as f:
+            data = json.load(f)
+    except (OSError, json.JSONDecodeError):
+        return None
+    return data if isinstance(data, dict) else None
+
+
 def compute_max_parallel(size: int) -> int:
-    return max(size // 4, 1)
+    return max(size // 3, 1)
 
 
-def compute_partitions(tests, full_parallel=False):
+def compute_partitions(
+    tests, repo_root, run_timeouts, partition_model=None, full_parallel=False
+):
     """Group per-commit tests by suite and emit partition metadata.
 
-    `full_parallel=True` (scheduled cron or `high priority` PR) sets
-    max_parallel = size, lifting the matrix-fanout throttle.
+    `run_timeouts`: `suite -> minutes` from `load_run_timeouts`.
+    `partition_model`: optional sglang-ci-stats `model.json`; per-file
+    `est` and per-suite `(coeff, bias)` each fall back independently to
+    in-source `est_time` / `(1.0, 0.0)`.
+    `full_parallel=True` lifts the matrix-fanout throttle.
     """
+    # Allowlist: stages pr-test.yml dispatches. Stress / weekly /
+    # nightly-* live in test/registered/ but pr-test doesn't run them.
+    dispatched_suites = set(run_timeouts) | set(_STAGE_A_OVERRIDES)
     suite_tests = defaultdict(list)
     for t in tests:
         if t.backend not in _TARGET_BACKENDS:
             continue
         if t.nightly or t.disabled is not None:
             continue
+        if t.effective_suite not in dispatched_suites:
+            continue
         suite_tests[t.effective_suite].append(t)
 
+    est_table = (partition_model or {}).get("est", {})
+    fit_table = (partition_model or {}).get("fit", {})
+
     result = {}
     for suite, group in suite_tests.items():
-        total = sum(t.est_time for t in group)
+        live_est = est_table.get(suite, {})
+        total = 0.0
+        for t in group:
+            relpath = os.path.relpath(t.filename, repo_root)
+            total += live_est.get(relpath, t.est_time)
+
+        fit = fit_table.get(suite) or {}
+        coeff = fit.get("coeff", 1.0)
+        bias = fit.get("bias", 0.0)
+
+        # Each shard pays `bias` once, so size >= coeff*total / (target-bias).
         if suite in _STAGE_A_OVERRIDES:
             size = _STAGE_A_OVERRIDES[suite]
             max_parallel = size
         else:
-            size = max(1, math.ceil(total / TARGET_SECONDS))
+            target = per_shard_target_seconds(suite, run_timeouts)
+            budget = target - bias
+            if budget <= 0:
+                raise RuntimeError(
+                    f"Suite {suite!r}: fit bias={bias}s >= target={target}s. "
+                    "Investigate the fit or raise the stage's run_timeout_minutes."
+                )
+            ideal_size = math.ceil(coeff * total / budget)
+            # ideal_size > len(group) -> slowest single file alone exceeds
+            # the per-shard budget; surface via raise instead of empty shard.
+            if ideal_size > len(group):
+                raise RuntimeError(
+                    f"Suite {suite!r}: needs {ideal_size} shards but has only "
+                    f"{len(group)} test file(s). target={target:.0f}s, "
+                    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)
-        # Check naive average (total/size). LPT can be ~4/3 of that in
-        # worst case; the 30-min job timeout enforces the real ceiling at
-        # runtime. This build-time check fails fast on egregious misconfigs.
-        if total / size > MAX_PARTITION_SECONDS:
-            raise RuntimeError(
-                f"Suite {suite!r}: total est_time {total:.0f}s / size {size} "
-                f"= {total / size:.0f}s exceeds MAX_PARTITION_SECONDS "
-                f"({MAX_PARTITION_SECONDS}s). Split a slow file or raise "
-                f"TARGET_SECONDS deliberately."
-            )
         result[suite] = {
             "size": size,
             "arr": list(range(size)),
@@ -130,14 +194,32 @@ def main():
         default="false",
         help="Lift the max_parallel throttle (set by schedule / `high priority`)",
     )
+    parser.add_argument(
+        "--partition-model-file",
+        default=None,
+        help="Path to sglang-ci-stats model.json (omit/missing -> static fallback)",
+    )
+    parser.add_argument(
+        "--pr-test-yml",
+        default=os.path.join(REPO_ROOT, ".github", "workflows", "pr-test.yml"),
+        help="Path to pr-test.yml; per-stage `run_timeout_minutes` is read from here.",
+    )
     args = parser.parse_args()
 
     files = discover_files(args.repo_root)
     # Warn-not-fail on unregistered files: run_suite.py catches this at
     # test-execution time with sanity_check=True; dispatch should keep going.
     all_tests = collect_tests(files, sanity_check=False)
+    partition_model = load_partition_model(args.partition_model_file)
+    run_timeouts = load_run_timeouts(args.pr_test_yml)
 
-    result = compute_partitions(all_tests, full_parallel=(args.full_parallel == "true"))
+    result = compute_partitions(
+        all_tests,
+        repo_root=args.repo_root,
+        run_timeouts=run_timeouts,
+        partition_model=partition_model,
+        full_parallel=(args.full_parallel == "true"),
+    )
     payload = json.dumps(result, separators=(",", ":"), sort_keys=True)
     if args.output_format == "gha":
         print(f"partitions={payload}")
@@ -148,9 +230,16 @@ def main():
     if summary_path:
         with open(summary_path, "a") as f:
             f.write("## Partitions\n\n")
+            if partition_model is None:
+                src_note = "no live model -- static est_time + (coeff=1, bias=0)"
+            else:
+                src_note = (
+                    f"live model `data_as_of={partition_model.get('data_as_of')}`, "
+                    f"`n_runs={partition_model.get('n_runs')}`"
+                )
             f.write(
                 f"`full_parallel={args.full_parallel}` "
-                f"(`size//4` throttle is lifted when true)\n\n"
+                f"(`size//3` throttle is lifted when true); {src_note}\n\n"
             )
             f.write("| Suite | size | max_parallel |\n")
             f.write("|---|---:|---:|\n")

test/registered/dsv4/test_deepseek_v4_flash_fp4_megamoe_b200.py

@@ -21,7 +21,7 @@ from sglang.test.test_utils import (
     try_cached_model,
 )
 
-register_cuda_ci(est_time=1800, suite="stage-c-test-dsv4-4-gpu-b200")
+register_cuda_ci(est_time=900, suite="stage-c-test-dsv4-4-gpu-b200")
 
 MODEL = "deepseek-ai/DeepSeek-V4-Flash"
 SERVER_LAUNCH_TIMEOUT = 3600

test/run_suite.py

@@ -1,8 +1,9 @@
 import argparse
 import glob
+import json
 import os
 import sys
-from typing import List
+from typing import Dict, List, Optional
 
 import tabulate
 
@@ -222,6 +223,29 @@ def pretty_print_tests(
     print(msg, flush=True)
 
 
+def load_live_est(
+    partition_model_file: Optional[str], suite: str, repo_root: str
+) -> Optional[Dict[str, float]]:
+    """`CIRegistry.filename -> est seconds` from `model.json est[suite]`;
+    None on any miss (caller falls back to in-source `est_time`)."""
+    if not partition_model_file or not os.path.exists(partition_model_file):
+        return None
+    try:
+        with open(partition_model_file) as f:
+            partition_model = json.load(f)
+    except (OSError, json.JSONDecodeError):
+        return None
+    if not isinstance(partition_model, dict):
+        return None
+    suite_est = partition_model.get("est", {}).get(suite)
+    if not isinstance(suite_est, dict) or not suite_est:
+        return None
+    return {
+        os.path.join(repo_root, relpath): float(elapsed)
+        for relpath, elapsed in suite_est.items()
+    }
+
+
 def run_a_suite(args):
     hw = HW_MAPPING[args.hw]
     suite = args.suite
@@ -261,7 +285,20 @@ def run_a_suite(args):
     ci_tests, skipped_tests = filter_tests(all_tests, hw, suite, nightly)
 
     if auto_partition_size:
-        ci_tests = auto_partition(ci_tests, auto_partition_id, auto_partition_size)
+        live_est = load_live_est(args.partition_model_file, suite, repo_root)
+        if live_est is not None:
+            print(
+                f"LPT: {len(live_est)} live est entries from {args.partition_model_file}",
+                flush=True,
+            )
+        else:
+            print(
+                f"LPT: no live est ({args.partition_model_file!r}); using in-source est_time",
+                flush=True,
+            )
+        ci_tests = auto_partition(
+            ci_tests, auto_partition_id, auto_partition_size, live_est=live_est
+        )
 
     pretty_print_tests(args, ci_tests, skipped_tests)
 
@@ -343,6 +380,12 @@ def main():
         default=600,
         help="Additional timeout in seconds when retry is enabled (default: 600)",
     )
+    parser.add_argument(
+        "--partition-model-file",
+        type=str,
+        default=None,
+        help="Path to sglang-ci-stats model.json for live LPT est; missing/malformed -> in-source est_time fallback.",
+    )
     args = parser.parse_args()
 
     # Validate auto-partition arguments