diff --git a/.github/workflows/_pr-test-check-changes.yml b/.github/workflows/_pr-test-check-changes.yml index 3da5c824a..df04e35ff 100644 --- a/.github/workflows/_pr-test-check-changes.yml +++ b/.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 diff --git a/.github/workflows/_pr-test-stage.yml b/.github/workflows/_pr-test-stage.yml index 38102dddf..b7f65f136 100644 --- a/.github/workflows/_pr-test-stage.yml +++ b/.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 diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index ce6ddf768..84c074139 100644 --- a/.github/workflows/pr-test.yml +++ b/.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 diff --git a/python/sglang/test/ci/ci_register.py b/python/sglang/test/ci/ci_register.py index 7c6234940..edf830493 100644 --- a/python/sglang/test/ci/ci_register.py +++ b/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] diff --git a/scripts/ci/utils/compute_partitions.py b/scripts/ci/utils/compute_partitions.py index 5cb48faa7..1c419384f 100644 --- a/scripts/ci/utils/compute_partitions.py +++ b/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)['']`. @@ -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") diff --git a/test/registered/dsv4/test_deepseek_v4_flash_fp4_megamoe_b200.py b/test/registered/dsv4/test_deepseek_v4_flash_fp4_megamoe_b200.py index 5c5c480f6..1c4df29c4 100644 --- a/test/registered/dsv4/test_deepseek_v4_flash_fp4_megamoe_b200.py +++ b/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 diff --git a/test/run_suite.py b/test/run_suite.py index dd9b1a774..eb3ebc938 100644 --- a/test/run_suite.py +++ b/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