diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index 6c2feef7a..cdaa5d6a9 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -95,10 +95,6 @@ jobs: with: github-token: ${{ github.token }} - - name: Show test partition assignments - continue-on-error: true - run: python3 test/show_partitions.py - - name: Determine run mode id: run-mode run: | diff --git a/test/show_partitions.py b/test/show_partitions.py deleted file mode 100755 index cb77791da..000000000 --- a/test/show_partitions.py +++ /dev/null @@ -1,145 +0,0 @@ -#!/usr/bin/env python3 -"""Print test partition assignments for all CI stages. - -Parses stage configs (hw, suite, partition_size) directly from pr-test.yml -so they never go out of sync. Uses the same deterministic LPT algorithm as -run_suite.py to show which test files land in which partition. - -Runs on ubuntu-latest without third-party deps (stdlib + ci_register). -""" - -import glob -import importlib.util -import os -import re - -# Import ci_register directly by file path to avoid pulling in the full -# sglang package (which has heavy deps like numpy, torch, etc.). -_ci_register_path = os.path.join( - os.path.dirname(__file__), "..", "python", "sglang", "test", "ci", "ci_register.py" -) -_spec = importlib.util.spec_from_file_location("ci_register", _ci_register_path) -_ci_register = importlib.util.module_from_spec(_spec) -_spec.loader.exec_module(_ci_register) - -HWBackend = _ci_register.HWBackend -auto_partition = _ci_register.auto_partition -collect_tests = _ci_register.collect_tests - -HW_MAPPING = { - "cpu": HWBackend.CPU, - "cuda": HWBackend.CUDA, - "amd": HWBackend.AMD, - "npu": HWBackend.NPU, -} - -# Regex to detect run_suite.py invocations from pr-test.yml. -# Uses lookaheads so --hw and --suite can appear in any order. -_RUN_SUITE_RE = re.compile( - r"^\s*(?:\w+=\S+\s+)?" # optional env prefix (e.g. VAR=value) - r"python3\s+run_suite\.py\b" - r"(?=.*\s--hw\s+(?P\S+))" - r"(?=.*\s--suite\s+(?P\S+))" -) -_PARTITION_SIZE_RE = re.compile(r"--auto-partition-size\s+(\d+)") - - -def parse_stage_configs(workflow_path): - """Parse (hw, suite, partition_count) tuples from pr-test.yml.""" - configs = [] - seen = set() - with open(workflow_path) as f: - for line in f: - # Skip commented-out lines - stripped = line.lstrip() - if stripped.startswith("#"): - continue - m = _RUN_SUITE_RE.search(line) - if not m: - continue - hw = m.group("hw") - suite = m.group("suite") - size_match = _PARTITION_SIZE_RE.search(line) - size = int(size_match.group(1)) if size_match else 1 - key = (hw, suite) - if key not in seen: - seen.add(key) - configs.append((hw, suite, size)) - return configs - - -def main(): - repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) - test_dir = os.path.join(repo_root, "test") - workflow_path = os.path.join(repo_root, ".github", "workflows", "pr-test.yml") - - os.chdir(test_dir) - - stage_configs = parse_stage_configs(workflow_path) - - files = [ - f - for f in glob.glob("registered/**/*.py", recursive=True) - if not f.endswith("/conftest.py") and not f.endswith("/__init__.py") - ] - all_tests = collect_tests(files, sanity_check=True) - - summary_file = os.environ.get("GITHUB_STEP_SUMMARY") - md_lines = [] - - def out(line=""): - print(line) - md_lines.append(line) - - out("## Test Partition Assignments") - out() - - for hw_str, suite, partition_count in stage_configs: - hw = HW_MAPPING.get(hw_str) - if hw is None: - continue - - suite_tests = [ - t - for t in all_tests - if t.backend == hw and t.suite == suite and not t.nightly - ] - enabled = [t for t in suite_tests if t.disabled is None] - disabled = [t for t in suite_tests if t.disabled is not None] - - total_time = sum(t.est_time for t in enabled) - out( - f"### {suite} (hw={hw_str}, {partition_count} partitions, " - f"{len(enabled)} enabled, {len(disabled)} disabled, est {total_time:.0f}s total)" - ) - out() - - if disabled: - out(f"**Disabled ({len(disabled)}):**") - for t in disabled: - out(f"- ~~{t.filename}~~ (reason: {t.disabled})") - out() - - if not enabled: - out("_(no enabled tests)_") - out() - continue - - out("| Partition | Tests | Est Time | Files |") - out("|-----------|-------|----------|-------|") - for rank in range(partition_count): - part_tests = auto_partition(enabled, rank, partition_count) - part_time = sum(t.est_time for t in part_tests) - file_list = ", ".join( - f"`{t.filename}` ({t.est_time:.0f}s)" for t in part_tests - ) - out(f"| {rank} | {len(part_tests)} | {part_time:.0f}s | {file_list} |") - out() - - if summary_file: - with open(summary_file, "a") as f: - f.write("\n".join(md_lines) + "\n") - - -if __name__ == "__main__": - main()