[CI] Show test partition assignments after checkout (#20085)

Co-authored-by: Alison Shao <alisonshao@mac.lan>

.github/workflows/pr-test.yml

@@ -96,6 +96,10 @@ jobs:
         with:
           ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
 
+      - name: Show test partition assignments
+        continue-on-error: true
+        run: python3 test/show_partitions.py
+
       - name: Determine run mode
         id: run-mode
         run: |

python/sglang/test/ci/ci_register.py

@@ -8,6 +8,7 @@ __all__ = [
     "HWBackend",
     "CIRegistry",
     "collect_tests",
+    "auto_partition",
     "register_cpu_ci",
     "register_cuda_ci",
     "register_amd_ci",
@@ -196,6 +197,32 @@ def ut_parse_one_file(filename: str) -> List[CIRegistry]:
     return visitor.registries
 
 
+def auto_partition(files: List[CIRegistry], rank: int, size: int) -> 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.
+    """
+    if not files or size <= 0:
+        return []
+
+    # 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))
+
+    partitions: List[List[CIRegistry]] = [[] for _ in range(size)]
+    partition_sums = [0.0] * size
+
+    # Greedily assign each file to the partition with the smallest current total time
+    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
+
+    if rank < size:
+        return partitions[rank]
+    return []
+
+
 def collect_tests(files: list[str], sanity_check: bool = True) -> List[CIRegistry]:
     ci_tests = []
     for file in files:

test/run_suite.py

@@ -5,7 +5,12 @@ from typing import List
 
 import tabulate
 
-from sglang.test.ci.ci_register import CIRegistry, HWBackend, collect_tests
+from sglang.test.ci.ci_register import (
+    CIRegistry,
+    HWBackend,
+    auto_partition,
+    collect_tests,
+)
 from sglang.test.ci.ci_utils import run_unittest_files
 
 HW_MAPPING = {
@@ -117,33 +122,6 @@ def filter_tests(
     return enabled_tests, skipped_tests
 
 
-def auto_partition(files: List[CIRegistry], rank, size):
-    """
-    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.
-    """
-    if not files or size <= 0:
-        return []
-
-    # Sort files by estimated_time in descending order (LPT heuristic).
-    # Use filename as tie-breaker to ensure deterministic partitioning
-    # regardless of glob ordering.
-    sorted_files = sorted(files, key=lambda f: (-f.est_time, f.filename))
-
-    partitions = [[] for _ in range(size)]
-    partition_sums = [0.0] * size
-
-    # Greedily assign each file to the partition with the smallest current total time
-    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
-
-    if rank < size:
-        return partitions[rank]
-    return []
-
-
 def pretty_print_tests(
     args, ci_tests: List[CIRegistry], skipped_tests: List[CIRegistry]
 ):

test/show_partitions.py

@@ -0,0 +1,146 @@
+#!/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)
+
+CIRegistry = _ci_register.CIRegistry
+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. IS_BLACKWELL=1)
+    r"python3\s+run_suite\.py\b"
+    r"(?=.*\s--hw\s+(?P<hw>\S+))"
+    r"(?=.*\s--suite\s+(?P<suite>\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} partition(s), "
+            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()