diff --git a/python/sglang/test/ci/ci_utils.py b/python/sglang/test/ci/ci_utils.py index 447a14a54..18c3994c2 100644 --- a/python/sglang/test/ci/ci_utils.py +++ b/python/sglang/test/ci/ci_utils.py @@ -3,9 +3,10 @@ import subprocess import threading import time from dataclasses import dataclass -from typing import Callable, List, Optional +from typing import Callable, List, Optional, Union from sglang.srt.utils.common import kill_process_tree +from sglang.test.ci.ci_register import CIRegistry @dataclass @@ -39,7 +40,9 @@ def run_with_timeout( def run_unittest_files( - files: List[TestFile], timeout_per_file: float, continue_on_error: bool = False + files: Union[List[TestFile], List[CIRegistry]], + timeout_per_file: float, + continue_on_error: bool = False, ): """ Run a list of test files. @@ -56,7 +59,12 @@ def run_unittest_files( failed_tests = [] for i, file in enumerate(files): - filename, estimated_time = file.name, file.estimated_time + if isinstance(file, CIRegistry): + filename, estimated_time = file.filename, file.est_time + else: + # FIXME: remove this branch after migrating all tests to use CIRegistry + filename, estimated_time = file.name, file.estimated_time + process = None def run_one_file(filename): diff --git a/test/run_suite.py b/test/run_suite.py index 49ef52ee9..32a5fd744 100644 --- a/test/run_suite.py +++ b/test/run_suite.py @@ -3,8 +3,10 @@ import glob import sys from typing import List +import tabulate + from sglang.test.ci.ci_register import CIRegistry, HWBackend, collect_tests -from sglang.test.ci.ci_utils import TestFile, run_unittest_files +from sglang.test.ci.ci_utils import run_unittest_files HW_MAPPING = { "cpu": HWBackend.CPU, @@ -53,7 +55,6 @@ def filter_tests( if t.backend == hw and t.suite == suite and t.nightly == nightly ] - ret = [] valid_suites = ( NIGHTLY_SUITES.get(hw, []) if nightly else PER_COMMIT_SUITES.get(hw, []) ) @@ -63,17 +64,13 @@ def filter_tests( f"Warning: Unknown suite {suite} for backend {hw.name}, nightly={nightly}" ) - for t in ci_tests: - if t.disabled is None: - ret.append(t) - print(f"Including test {t.filename}") - else: - print(f"Skipping disabled test {t.filename} due to: {t.disabled}") + enabled_tests = [t for t in ci_tests if t.disabled is None] + skipped_tests = [t for t in ci_tests if t.disabled is not None] - return ret + return enabled_tests, skipped_tests -def auto_partition(files, rank, size): +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. @@ -82,7 +79,7 @@ def auto_partition(files, rank, size): return [] # Sort files by estimated_time in descending order (LPT heuristic) - sorted_files = sorted(files, key=lambda f: f.estimated_time, reverse=True) + sorted_files = sorted(files, key=lambda f: f.est_time, reverse=True) partitions = [[] for _ in range(size)] partition_sums = [0.0] * size @@ -91,13 +88,51 @@ def auto_partition(files, rank, size): 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.estimated_time + 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] +): + hw = HW_MAPPING[args.hw] + suite = args.suite + nightly = args.nightly + if args.auto_partition_size: + partition_info = ( + f"{args.auto_partition_id + 1}/{args.auto_partition_size} " + f"(0-based id={args.auto_partition_id})" + ) + else: + partition_info = "full" + + headers = ["Hardware", "Suite", "Nightly", "Partition"] + rows = [[hw.name, suite, str(nightly), partition_info]] + msg = tabulate.tabulate(rows, headers=headers, tablefmt="psql") + print(msg + "\n") + + if skipped_tests: + print(f"⚠️ Skipped {len(skipped_tests)} test(s):") + for t in skipped_tests: + reason = t.disabled or "disabled" + print(f" - {t.filename} (reason: {reason})") + + print() + + if len(ci_tests) == 0: + print(f"No tests found for hw={hw.name}, suite={suite}, nightly={nightly}") + print("This is expected during incremental migration. Skipping.") + return + + total_est_time = sum(t.est_time for t in ci_tests) + print(f"✅ Enabled {len(ci_tests)} test(s) (est total {total_est_time:.1f}s):") + for t in ci_tests: + print(f" - {t.filename} (est_time={t.est_time})") + + def run_a_suite(args): hw = HW_MAPPING[args.hw] suite = args.suite @@ -111,29 +146,19 @@ def run_a_suite(args): sanity_check = False # Allow files without registration during migration else: files = glob.glob("registered/**/*.py", recursive=True) - sanity_check = ( - True # Strict: all registered files must have proper registration - ) + # Strict: all registered files must have proper registration + sanity_check = True - ci_tests = filter_tests( - collect_tests(files, sanity_check=sanity_check), hw, suite, nightly - ) - test_files = [TestFile(t.filename, t.est_time) for t in ci_tests] - - if not test_files: - print(f"No tests found for hw={hw.name}, suite={suite}, nightly={nightly}") - print("This is expected during incremental migration. Skipping.") - return 0 + all_tests = collect_tests(files, sanity_check=sanity_check) + ci_tests, skipped_tests = filter_tests(all_tests, hw, suite, nightly) if auto_partition_size: - test_files = auto_partition(test_files, auto_partition_id, auto_partition_size) + ci_tests = auto_partition(ci_tests, auto_partition_id, auto_partition_size) - print( - f"Running {len(test_files)} test(s) for hw={hw.name}, suite={suite}, nightly={nightly}" - ) + pretty_print_tests(args, ci_tests, skipped_tests) return run_unittest_files( - test_files, + ci_tests, timeout_per_file=args.timeout_per_file, continue_on_error=args.continue_on_error, )