diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index abc7a0970..6f0f5f73f 100644 --- a/.github/workflows/pr-test.yml +++ b/.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: | diff --git a/python/sglang/test/ci/ci_register.py b/python/sglang/test/ci/ci_register.py index fdd72de44..5c564e416 100644 --- a/python/sglang/test/ci/ci_register.py +++ b/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: diff --git a/test/run_suite.py b/test/run_suite.py index 6805160cc..aba22bafd 100644 --- a/test/run_suite.py +++ b/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] ): diff --git a/test/show_partitions.py b/test/show_partitions.py new file mode 100755 index 000000000..49426b457 --- /dev/null +++ b/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\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} 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()