""" Test runner for multimodal_gen that manages test suites and parallel execution. For diffusion 1-gpu/2-gpu suites, cases are partitioned by estimated runtime using LPT so each CI shard has a similar total runtime. """ import argparse import copy import json import os import random import subprocess import sys import time from dataclasses import dataclass from pathlib import Path import tabulate from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.test.partitioning import ( PartitionItem, partition_items_by_lpt, ) from sglang.multimodal_gen.test.runner.pytest_runner import ( partition_items_by_index, run_pytest, ) from sglang.multimodal_gen.test.server.testcase_configs import ( BASELINE_CONFIG, DiffusionTestCase, ) # TODO: remove duplicated code if current_platform.is_npu(): from sglang.multimodal_gen.test.server.ascend.testcase_configs_npu import ( _UPDATE_WEIGHTS_FROM_DISK_TEST_FILE, COMPONENT_ACCURACY_SUITES, DEFAULT_EST_TIME_SECONDS, DEFAULT_STANDALONE_EST_TIME_SECONDS, FILE_SUITES, PARAMETRIZED_CASE_GROUPS, STANDALONE_FILES, STARTUP_OVERHEAD_SECONDS, SUITES, ) else: from sglang.multimodal_gen.test.server.gpu_cases import ( # noqa: F401 It is used by ci scripts _UPDATE_WEIGHTS_FROM_DISK_TEST_FILE, _UPDATE_WEIGHTS_MODEL_PAIR_ENV, _UPDATE_WEIGHTS_MODEL_PAIR_IDS, COMPONENT_ACCURACY_FILE_NUM_GPUS, COMPONENT_ACCURACY_SUITES, DEFAULT_EST_TIME_SECONDS, DEFAULT_STANDALONE_EST_TIME_SECONDS, FILE_SUITES, ONE_GPU_CASES, PARAMETRIZED_CASE_GROUPS, STANDALONE_FILE_EST_TIMES, STANDALONE_FILES, STARTUP_OVERHEAD_SECONDS, STRICT_SUITES, SUITES, TWO_GPU_CASES, ) logger = init_logger(__name__) @dataclass(frozen=True) class PartitionAssignment: case_ids: list[str] standalone_files: list[str] estimated_time: float | None = None missing_standalone_estimates: list[str] | None = None def get_case_est_time(case_id: str) -> float: scenario = BASELINE_CONFIG.scenarios.get(case_id) if scenario is None: return DEFAULT_EST_TIME_SECONDS if scenario.estimated_full_test_time_s is not None: return scenario.estimated_full_test_time_s return scenario.expected_e2e_ms / 1000.0 + STARTUP_OVERHEAD_SECONDS def get_standalone_file_est_time( suite: str, standalone_file: str ) -> tuple[float, bool]: suite_est_times = STANDALONE_FILE_EST_TIMES.get(suite, {}) if standalone_file not in suite_est_times: return DEFAULT_STANDALONE_EST_TIME_SECONDS, True return suite_est_times[standalone_file], False def get_all_standalone_file_est_times() -> dict[str, dict[str, float]]: return copy.deepcopy(STANDALONE_FILE_EST_TIMES) def validate_standalone_file_est_times() -> dict[str, list[str]]: missing_by_suite: dict[str, list[str]] = {} for suite, standalone_files in STANDALONE_FILES.items(): suite_est_times = STANDALONE_FILE_EST_TIMES.get(suite, {}) missing = [ standalone_file for standalone_file in standalone_files if standalone_file not in suite_est_times ] if missing: missing_by_suite[suite] = missing return missing_by_suite def auto_partition( cases: list[DiffusionTestCase], rank: int, size: int ) -> list[DiffusionTestCase]: if not cases or size <= 0: return [] case_by_id = {case.id: case for case in cases} items = [ PartitionItem(kind="case", item_id=case.id, est_time=get_case_est_time(case.id)) for case in cases ] partitions = partition_items_by_lpt(items, size) if rank >= len(partitions): return [] return [case_by_id[item.item_id] for item in partitions[rank]] def get_suite_files_rel(suite: str, parametrized_only: bool = False) -> list[str]: if parametrized_only and suite in PARAMETRIZED_CASE_GROUPS: return [filename for filename, _ in PARAMETRIZED_CASE_GROUPS[suite]] return SUITES[suite] def _normalize_standalone_key(standalone_file: str) -> str: return f"standalone:{standalone_file}" def parse_partition_plan( suite: str, partition_id: int, total_partitions: int, plan_json: str, ) -> PartitionAssignment: plan = json.loads(plan_json) if plan.get("suite") != suite: raise ValueError( f"Partition plan suite mismatch: expected {suite!r}, " f"got {plan.get('suite')!r}" ) partition_count = plan.get("partition_count") if partition_count != total_partitions: raise ValueError( f"Partition count mismatch for suite {suite!r}: " f"plan={partition_count}, matrix={total_partitions}" ) partitions = plan.get("partitions", []) selected_partition = None for partition in partitions: if partition.get("part") == partition_id: selected_partition = partition break if selected_partition is None: raise ValueError( f"Partition {partition_id} not found in plan for suite {suite!r}" ) return PartitionAssignment( case_ids=list(selected_partition.get("case_ids", [])), standalone_files=list(selected_partition.get("standalone_files", [])), estimated_time=selected_partition.get("estimated_time"), missing_standalone_estimates=list( selected_partition.get("missing_standalone_estimates", []) ), ) def _merge_execution_results( executed_cases: list[str], case_results: dict[str, str], new_executed_cases: list[str], new_case_results: dict[str, str], ) -> None: executed_cases.extend( case_id for case_id in new_executed_cases if case_id not in executed_cases ) case_results.update(new_case_results) def _format_standalone_estimate_snippet( suite: str, standalone_file: str, measured_full_test_time_s: float ) -> str: return ( f'"{suite}": {{\n' f' "{standalone_file}": {measured_full_test_time_s:.1f},\n' f"}}" ) def _print_missing_standalone_estimate_message( suite: str, standalone_file: str, measured_full_test_time_s: float, ) -> None: snippet = _format_standalone_estimate_snippet( suite, standalone_file, measured_full_test_time_s ) logger.error( f'\n{"=" * 60}\n' f'Add standalone estimate for suite "{suite}" and file "{standalone_file}":\n\n' f"File: python/sglang/multimodal_gen/test/run_suite.py\n\n" f"Current partition used fallback estimate: " f"{DEFAULT_STANDALONE_EST_TIME_SECONDS:.1f}s\n\n" f"{snippet}\n" f'{"=" * 60}\n' ) def _run_standalone_file( suite: str, standalone_rel: str, target_dir: Path, extra_filter: str | None = None, ) -> tuple[int, list[str], dict[str, str], dict]: if standalone_rel == _UPDATE_WEIGHTS_FROM_DISK_TEST_FILE: _maybe_pin_update_weights_model_pair([standalone_rel]) est_time, used_fallback_estimate = get_standalone_file_est_time( suite, standalone_rel ) standalone_file = _resolve_suite_files(target_dir, [standalone_rel], strict=True)[0] junit_xml_path = str( target_dir / f"junit_results_{suite}_{Path(standalone_rel).stem}.xml" ) start_time = time.perf_counter() exit_code, _, _ = run_pytest( [standalone_file], filter_expr=extra_filter, junit_xml_path=junit_xml_path, ) measured_full_test_time_s = round(time.perf_counter() - start_time, 1) standalone_key = _normalize_standalone_key(standalone_rel) measurement = { "suite": suite, "standalone_file": standalone_rel, "measured_full_test_time_s": measured_full_test_time_s, "used_fallback_estimate": used_fallback_estimate, "fallback_estimate_s": DEFAULT_STANDALONE_EST_TIME_SECONDS, "had_configured_estimate": not used_fallback_estimate, "configured_or_fallback_estimate_s": est_time, } if used_fallback_estimate: _print_missing_standalone_estimate_message( suite, standalone_rel, measured_full_test_time_s ) return ( exit_code, [standalone_key], {standalone_key: "pass" if exit_code == 0 else "fail"}, measurement, ) def parse_args(): suite_choices = sorted(set(FILE_SUITES) | set(PARAMETRIZED_CASE_GROUPS)) parser = argparse.ArgumentParser(description="Run multimodal_gen test suite") parser.add_argument( "--suite", type=str, required=True, choices=suite_choices, help="The test suite to run.", ) parser.add_argument( "--partition-id", type=int, default=0, help="Index of the current partition (for parallel execution)", ) parser.add_argument( "--total-partitions", type=int, default=1, help="Total number of partitions", ) parser.add_argument( "--base-dir", type=str, default="server", help="Base directory for tests relative to this script's parent", ) parser.add_argument( "-k", "--filter", type=str, default=None, help="Pytest filter expression (passed to pytest -k)", ) parser.add_argument( "--continue-on-error", action="store_true", default=False, help="Continue running remaining tests even if one fails.", ) parser.add_argument( "--partition-plan-json", type=str, default=None, help="Full partition plan JSON for the current suite.", ) return parser.parse_args() def write_execution_report( suite: str, partition_id: int, total_partitions: int, executed_cases: list[str], is_standalone: bool = False, standalone_file: str | None = None, case_results: dict[str, str] | None = None, missing_standalone_estimates: list[str] | None = None, standalone_measurements: list[dict] | None = None, ) -> str: report = { "suite": suite, "partition_id": partition_id, "total_partitions": total_partitions, "is_standalone": is_standalone, "standalone_file": standalone_file, "executed_cases": executed_cases, "case_results": case_results or {}, "missing_standalone_estimates": missing_standalone_estimates or [], "standalone_measurements": standalone_measurements or [], } report_filename = f"execution_report_{suite}_{partition_id}.json" report_path = Path(__file__).parent / report_filename with open(report_path, "w", encoding="utf-8") as f: json.dump(report, f, indent=2) logger.info("Execution report written to: %s", report_path) return str(report_path) def run_component_accuracy_files(files, filter_expr=None, continue_on_error=False): exit_code = 0 for file_path in files: file_name = Path(file_path).name num_gpus = COMPONENT_ACCURACY_FILE_NUM_GPUS.get(file_name, 1) if num_gpus > 1: cmd = [ sys.executable, "-m", "torch.distributed.run", f"--nproc_per_node={num_gpus}", "-m", "pytest", "-s", "-v", ] else: cmd = [sys.executable, "-m", "pytest", "-s", "-v"] if filter_expr: cmd.extend(["-k", filter_expr]) cmd.append(file_path) print(f"Running command: {' '.join(cmd)}") file_exit_code = subprocess.call(cmd) if file_exit_code == 5: print( "No tests collected (exit code 5). This is expected when filters " "deselect all tests in a file. Treating as success." ) file_exit_code = 0 if file_exit_code != 0 and exit_code == 0: exit_code = file_exit_code if file_exit_code != 0 and not continue_on_error: return file_exit_code return exit_code def _is_in_ci() -> bool: return os.environ.get("SGLANG_IS_IN_CI", "").lower() in ("1", "true", "yes", "on") def _maybe_pin_update_weights_model_pair(suite_files_rel: list[str]) -> None: if not _is_in_ci(): return if _UPDATE_WEIGHTS_FROM_DISK_TEST_FILE not in suite_files_rel: return if os.environ.get(_UPDATE_WEIGHTS_MODEL_PAIR_ENV): print( f"Using preset {_UPDATE_WEIGHTS_MODEL_PAIR_ENV}=" f"{os.environ[_UPDATE_WEIGHTS_MODEL_PAIR_ENV]}" ) return selected_pair = random.choice(_UPDATE_WEIGHTS_MODEL_PAIR_IDS) os.environ[_UPDATE_WEIGHTS_MODEL_PAIR_ENV] = selected_pair print(f"Selected {_UPDATE_WEIGHTS_MODEL_PAIR_ENV}={selected_pair} for this CI run") def _resolve_suite_files( target_dir: Path, suite_files_rel: list[str], strict: bool ) -> list[str]: suite_files_abs = [] for f_rel in suite_files_rel: f_abs = target_dir / f_rel if not f_abs.exists(): msg = f"Test file {f_rel} not found in {target_dir}." if strict: print(f"Error: {msg}") sys.exit(1) print(f"Warning: {msg} Skipping.") continue suite_files_abs.append(str(f_abs)) return suite_files_abs def _run_file_suite(args, target_dir: Path) -> int: suite_files_rel = FILE_SUITES[args.suite] _maybe_pin_update_weights_model_pair(suite_files_rel) suite_files_abs = _resolve_suite_files( target_dir, suite_files_rel, args.suite in STRICT_SUITES ) if not suite_files_abs: print(f"No valid test files found for suite '{args.suite}'.") return 1 if args.suite in STRICT_SUITES else 0 exit_code, _, _ = run_pytest( suite_files_abs, filter_expr=args.filter, junit_xml_path=None, ) return exit_code def _get_dynamic_suite_cases(suite: str) -> list[DiffusionTestCase]: cases = [] for _, case_group in PARAMETRIZED_CASE_GROUPS[suite]: cases.extend(case_group) return cases def _get_parametrized_files_for_case_ids( suite: str, case_ids: set[str], target_dir: Path ) -> list[str]: files = [] for filename, case_group in PARAMETRIZED_CASE_GROUPS[suite]: if any(case.id in case_ids for case in case_group): file_path = target_dir / filename if file_path.exists(): files.append(str(file_path)) else: logger.warning("Test file %s not found in %s", filename, target_dir) return files def _get_standalone_file(target_dir: Path, suite: str, index: int) -> str | None: standalone_files = STANDALONE_FILES.get(suite, []) if index < 0 or index >= len(standalone_files): return None file_path = target_dir / standalone_files[index] if file_path.exists(): return str(file_path) logger.warning( "Standalone test file %s not found in %s", standalone_files[index], target_dir ) return None def _run_dynamic_suite(args, target_dir: Path) -> int: if args.partition_plan_json: assignment = parse_partition_plan( suite=args.suite, partition_id=args.partition_id, total_partitions=args.total_partitions, plan_json=args.partition_plan_json, ) rows = [[args.suite, f"{args.partition_id + 1}/{args.total_partitions}"]] print(tabulate.tabulate(rows, headers=["Suite", "Partition"], tablefmt="psql")) total_est_time = 0.0 executed_cases: list[str] = [] case_results: dict[str, str] = {} missing_standalone_estimates: list[str] = [] standalone_measurements: list[dict] = [] overall_exit_code = 0 if assignment.case_ids: case_id_set = set(assignment.case_ids) total_est_time += sum( get_case_est_time(case_id) for case_id in assignment.case_ids ) suite_files = _get_parametrized_files_for_case_ids( args.suite, case_id_set, target_dir ) if not suite_files: print( f"No valid parametrized test files found for suite '{args.suite}'." ) return 0 partition_filter = " or ".join( f"[{case_id}]" for case_id in assignment.case_ids ) filter_expr = ( f"({partition_filter}) and ({args.filter})" if args.filter else partition_filter ) print( f"Running {len(assignment.case_ids)} parametrized cases with estimated total " f"{sum(get_case_est_time(case_id) for case_id in assignment.case_ids):.1f}s:" ) for case_id in assignment.case_ids: print(f" - case: {case_id} ({get_case_est_time(case_id):.1f}s)") print(f"Test files: {[Path(f).name for f in suite_files]}") print(f"Filter expression: {filter_expr}") junit_xml_path = str( target_dir / f"junit_results_{args.suite}_{args.partition_id}.xml" ) exit_code, new_executed_cases, new_case_results = run_pytest( suite_files, filter_expr=filter_expr, junit_xml_path=junit_xml_path, ) _merge_execution_results( executed_cases, case_results, new_executed_cases, new_case_results ) if exit_code != 0 and overall_exit_code == 0: overall_exit_code = exit_code if exit_code != 0 and not args.continue_on_error: write_execution_report( suite=args.suite, partition_id=args.partition_id, total_partitions=args.total_partitions, executed_cases=executed_cases, is_standalone=False, standalone_file=None, case_results=case_results, missing_standalone_estimates=missing_standalone_estimates, standalone_measurements=standalone_measurements, ) return overall_exit_code if assignment.standalone_files: standalone_estimate = sum( get_standalone_file_est_time(args.suite, standalone_file)[0] for standalone_file in assignment.standalone_files ) total_est_time += standalone_estimate print( f"Running {len(assignment.standalone_files)} standalone file(s) with estimated total " f"{standalone_estimate:.1f}s:" ) for standalone_file in assignment.standalone_files: est_time, used_fallback_estimate = get_standalone_file_est_time( args.suite, standalone_file ) fallback_suffix = ( f", fallback estimate {DEFAULT_STANDALONE_EST_TIME_SECONDS:.1f}s" if used_fallback_estimate else "" ) print( f" - standalone: {standalone_file} " f"({est_time:.1f}s{fallback_suffix})" ) for standalone_file in assignment.standalone_files: exit_code, new_executed_cases, new_case_results, measurement = ( _run_standalone_file( args.suite, standalone_file, target_dir, extra_filter=args.filter, ) ) if measurement["used_fallback_estimate"]: missing_standalone_estimates.append(standalone_file) standalone_measurements.append(measurement) _merge_execution_results( executed_cases, case_results, new_executed_cases, new_case_results, ) if exit_code != 0 and overall_exit_code == 0: overall_exit_code = exit_code if exit_code != 0 and not args.continue_on_error: break print(f"Partition estimated total time: {total_est_time:.1f}s") write_execution_report( suite=args.suite, partition_id=args.partition_id, total_partitions=args.total_partitions, executed_cases=executed_cases, is_standalone=False, standalone_file=None, case_results=case_results, missing_standalone_estimates=missing_standalone_estimates, standalone_measurements=standalone_measurements, ) return overall_exit_code all_cases = _get_dynamic_suite_cases(args.suite) standalone_files = STANDALONE_FILES.get(args.suite, []) parametrized_partitions = args.total_partitions - len(standalone_files) if parametrized_partitions < 0: print( f"Error: total_partitions ({args.total_partitions}) must be >= " f"standalone files ({len(standalone_files)})" ) return 1 if args.partition_id < parametrized_partitions: if not all_cases: print(f"No cases found for suite '{args.suite}'.") return 0 my_cases = auto_partition(all_cases, args.partition_id, parametrized_partitions) if not my_cases: print( f"No cases assigned to partition {args.partition_id}. Exiting success." ) write_execution_report( suite=args.suite, partition_id=args.partition_id, total_partitions=args.total_partitions, executed_cases=[], is_standalone=False, missing_standalone_estimates=[], standalone_measurements=[], ) return 0 case_ids = [case.id for case in my_cases] case_id_set = set(case_ids) total_est_time = sum(get_case_est_time(case.id) for case in my_cases) suite_files = _get_parametrized_files_for_case_ids( args.suite, case_id_set, target_dir ) if not suite_files: print(f"No valid parametrized test files found for suite '{args.suite}'.") return 0 partition_filter = " or ".join(f"[{case_id}]" for case_id in case_ids) filter_expr = ( f"({partition_filter}) and ({args.filter})" if args.filter else partition_filter ) rows = [[args.suite, f"{args.partition_id + 1}/{args.total_partitions}"]] print(tabulate.tabulate(rows, headers=["Suite", "Partition"], tablefmt="psql")) print( f"Running {len(my_cases)} cases with estimated total " f"{total_est_time:.1f}s:" ) for case in my_cases: print(f" - {case.id} ({get_case_est_time(case.id):.1f}s)") print(f"Test files: {[Path(f).name for f in suite_files]}") print(f"Filter expression: {filter_expr}") junit_xml_path = str( target_dir / f"junit_results_{args.suite}_{args.partition_id}.xml" ) exit_code, executed_cases, case_results = run_pytest( suite_files, filter_expr=filter_expr, junit_xml_path=junit_xml_path, ) write_execution_report( suite=args.suite, partition_id=args.partition_id, total_partitions=args.total_partitions, executed_cases=executed_cases, is_standalone=False, case_results=case_results, missing_standalone_estimates=[], standalone_measurements=[], ) return exit_code standalone_idx = args.partition_id - parametrized_partitions if standalone_idx >= len(standalone_files): print( f"ERROR: Standalone partition index {standalone_idx} exceeds available " f"standalone files ({len(standalone_files)}) for suite '{args.suite}'." ) return 1 standalone_rel = standalone_files[standalone_idx] print( f"Suite: {args.suite} | Partition: {args.partition_id + 1}/{args.total_partitions} (standalone)" ) print(f"Running standalone test file: {Path(standalone_rel).name}") exit_code, executed_cases, case_results, measurement = _run_standalone_file( args.suite, standalone_rel, target_dir, extra_filter=args.filter, ) write_execution_report( suite=args.suite, partition_id=args.partition_id, total_partitions=args.total_partitions, executed_cases=executed_cases, is_standalone=True, standalone_file=standalone_rel, case_results=case_results, missing_standalone_estimates=( [standalone_rel] if measurement["used_fallback_estimate"] else [] ), standalone_measurements=[measurement], ) return exit_code def main(): args = parse_args() validate_standalone_file_est_times() test_root_dir = Path(__file__).resolve().parent target_dir = test_root_dir / args.base_dir if not target_dir.exists(): print(f"Error: Target directory {target_dir} does not exist.") sys.exit(1) if args.suite in COMPONENT_ACCURACY_SUITES: suite_files_rel = FILE_SUITES[args.suite] suite_files_abs = _resolve_suite_files( target_dir, suite_files_rel, args.suite in STRICT_SUITES ) if not suite_files_abs: print(f"No valid test files found for suite '{args.suite}'.") sys.exit(1 if args.suite in STRICT_SUITES else 0) my_files = partition_items_by_index( suite_files_abs, args.partition_id, args.total_partitions ) partition_info = ( f"{args.partition_id + 1}/{args.total_partitions} " f"(0-based id={args.partition_id})" ) headers = ["Suite", "Partition"] rows = [[args.suite, partition_info]] msg = tabulate.tabulate(rows, headers=headers, tablefmt="psql") + "\n" msg += f"Enabled {len(my_files)} file(s):\n" for file_path in my_files: msg += f" - {file_path}\n" print(msg, flush=True) print( f"Suite: {args.suite} | Partition: {args.partition_id}/{args.total_partitions}" ) print(f"Selected {len(suite_files_abs)} files:") for f in suite_files_abs: print(f" - {os.path.basename(f)}") if not my_files: print("No files assigned to this partition. Exiting success.") sys.exit(0) print(f"Running {len(my_files)} files in this shard: {', '.join(my_files)}") exit_code = run_component_accuracy_files( my_files, filter_expr=args.filter, continue_on_error=args.continue_on_error, ) msg = "\n" + tabulate.tabulate(rows, headers=headers, tablefmt="psql") + "\n" msg += f"Executed {len(my_files)} file(s):\n" for file_path in my_files: msg += f" - {file_path}\n" print(msg, flush=True) elif args.suite in PARAMETRIZED_CASE_GROUPS: exit_code = _run_dynamic_suite(args, target_dir) else: exit_code = _run_file_suite(args, target_dir) sys.exit(exit_code) if __name__ == "__main__": main()