From 5de89ea9420af45c19a5ded8818ba5d38311213a Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 17 Apr 2026 22:37:24 +0800 Subject: [PATCH] [diffusion] CI: fix auto-partition (#23076) --- .../test/server/perf_baselines.json | 31 ++++---- .../diffusion/compute_diffusion_partitions.py | 37 ++++++++-- .../utils/diffusion/diffusion_case_parser.py | 71 +++++++++++++++---- .../diffusion/verify_diffusion_coverage.py | 29 ++++++-- 4 files changed, 128 insertions(+), 40 deletions(-) diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines.json b/python/sglang/multimodal_gen/test/server/perf_baselines.json index 509ab0af9..d9223a5e4 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines.json @@ -1378,26 +1378,27 @@ }, "fast_hunyuan_video": { "stages_ms": { - "DecodingStage": 4586.75, "InputValidationStage": 0.06, - "DenoisingStage": 5828.43, - "TextEncodingStage": 310.35, - "LatentPreparationStage": 0.14, - "TimestepPreparationStage": 32.2 + "TextEncodingStage": 321.95, + "TimestepPreparationStage": 28.98, + "LatentPreparationStage": 0.13, + "DenoisingStage": 5898.72, + "DecodingStage": 4736.19, + "per_frame_generation": null }, "denoise_step_ms": { - "0": 296.26, - "1": 853.35, - "2": 853.4, - "3": 852.1, - "4": 853.82, - "5": 854.99 + "0": 286.89, + "1": 1115.43, + "2": 1118.06, + "3": 1124.53, + "4": 1130.91, + "5": 1119.61 }, - "expected_e2e_ms": 11636.47, - "expected_avg_denoise_ms": 970.56, - "expected_median_denoise_ms": 1104.44, - "estimated_full_test_time_s": 131.6 + "expected_e2e_ms": 11816.27, + "expected_avg_denoise_ms": 982.57, + "expected_median_denoise_ms": 1118.83 }, + "wan2_2_i2v_a14b_2gpu": { "stages_ms": { "InputValidationStage": 18.1, diff --git a/scripts/ci/utils/diffusion/compute_diffusion_partitions.py b/scripts/ci/utils/diffusion/compute_diffusion_partitions.py index a1af41027..50569dfe9 100755 --- a/scripts/ci/utils/diffusion/compute_diffusion_partitions.py +++ b/scripts/ci/utils/diffusion/compute_diffusion_partitions.py @@ -17,9 +17,9 @@ from pathlib import Path from diffusion_case_parser import ( BASELINE_REL_PATH, RUN_SUITE_REL_PATH, - TESTCASE_CONFIG_REL_PATH, DiffusionSuiteInfo, collect_diffusion_suites, + resolve_case_config_path, ) SUITE_OUTPUT_NAMES = { @@ -37,6 +37,29 @@ class PartitionItem: used_fallback_estimate: bool = False +def validate_suite_case_coverage(suites: dict[str, DiffusionSuiteInfo]) -> None: + """ + Guardrail: dynamic diffusion suites must contain parametrized cases. + """ + suites_with_no_cases = [] + for suite_name in SUITE_OUTPUT_NAMES: + suite_info = suites.get(suite_name) + if suite_info is None: + print(f"Error: Required suite '{suite_name}' not found in parsed suites.") + sys.exit(1) + if len(suite_info.cases) == 0: + suites_with_no_cases.append(suite_name) + + if suites_with_no_cases: + joined = ", ".join(suites_with_no_cases) + print( + "Error: Parsed zero parametrized cases for diffusion suites: " + f"{joined}. This usually means run_suite case imports changed but " + "diffusion parser logic was not updated." + ) + sys.exit(1) + + def compute_partition_count( total_time_seconds: float, min_time_seconds: float, @@ -229,22 +252,24 @@ def main(): script_dir = Path(__file__).resolve().parent repo_root = script_dir.parent.parent.parent.parent - testcase_config_path = repo_root / TESTCASE_CONFIG_REL_PATH baseline_path = repo_root / BASELINE_REL_PATH run_suite_path = repo_root / RUN_SUITE_REL_PATH - if not testcase_config_path.exists(): - print(f"Error: Testcase config not found: {testcase_config_path}") - sys.exit(1) if not run_suite_path.exists(): print(f"Error: Run suite not found: {run_suite_path}") sys.exit(1) + try: + case_config_path = resolve_case_config_path(repo_root, run_suite_path) + except (RuntimeError, FileNotFoundError) as exc: + print(f"Error: {exc}") + sys.exit(1) suites = collect_diffusion_suites( - testcase_config_path, + case_config_path, run_suite_path, baseline_path, ) + validate_suite_case_coverage(suites) print("=== Diffusion Partition Computation ===") print(f"Min partition time: {args.min_time}s ({args.min_time/60:.1f} min)") diff --git a/scripts/ci/utils/diffusion/diffusion_case_parser.py b/scripts/ci/utils/diffusion/diffusion_case_parser.py index d5839714e..651f47c02 100755 --- a/scripts/ci/utils/diffusion/diffusion_case_parser.py +++ b/scripts/ci/utils/diffusion/diffusion_case_parser.py @@ -2,14 +2,16 @@ """ AST-based parser for diffusion test cases. -This module parses testcase_configs.py and run_suite.py using AST -to extract test case information without requiring sglang dependencies. -Designed to run on lightweight CI runners (ubuntu-latest). +This module parses the diffusion case source and run_suite.py using AST to +extract test case information without requiring sglang dependencies. The case +source file is discovered from ONE_GPU_CASES/TWO_GPU_CASES imports in +run_suite.py so CI keeps a single source of truth. Usage: # From sibling scripts in this directory: - from diffusion_case_parser import collect_diffusion_suites - suites = collect_diffusion_suites(testcase_config_path, run_suite_path, baseline_path) + from diffusion_case_parser import collect_diffusion_suites, resolve_case_config_path + case_config_path = resolve_case_config_path(repo_root, run_suite_path) + suites = collect_diffusion_suites(case_config_path, run_suite_path, baseline_path) """ import ast @@ -34,9 +36,6 @@ DEFAULT_EST_TIME_SECONDS = 300.0 STARTUP_OVERHEAD_SECONDS = 120.0 # Paths relative to repository root -TESTCASE_CONFIG_REL_PATH = ( - "python/sglang/multimodal_gen/test/server/testcase_configs.py" -) BASELINE_REL_PATH = "python/sglang/multimodal_gen/test/server/perf_baselines.json" RUN_SUITE_REL_PATH = "python/sglang/multimodal_gen/test/run_suite.py" @@ -65,7 +64,7 @@ class DiffusionSuiteInfo: class DiffusionTestCaseVisitor(ast.NodeVisitor): """ - AST visitor to extract DiffusionTestCase definitions from testcase_configs.py. + AST visitor to extract DiffusionTestCase definitions from the case config. Parses assignments like: ONE_GPU_CASES_A: list[DiffusionTestCase] = [ @@ -121,6 +120,50 @@ class DiffusionTestCaseVisitor(ast.NodeVisitor): return None +def resolve_case_config_path(repo_root: Path, run_suite_path: Path) -> Path: + """ + Resolve the diffusion case config path from run_suite imports. + + run_suite.py must import BOTH ONE_GPU_CASES and TWO_GPU_CASES from the same + module. That imported module is treated as the single source of truth. + """ + with open(run_suite_path, "r", encoding="utf-8") as f: + content = f.read() + + tree = ast.parse(content, filename=str(run_suite_path)) + one_gpu_module: Optional[str] = None + two_gpu_module: Optional[str] = None + + for node in ast.walk(tree): + if not isinstance(node, ast.ImportFrom) or not node.module: + continue + imported_names = {alias.name for alias in node.names} + if "ONE_GPU_CASES" in imported_names: + one_gpu_module = node.module + if "TWO_GPU_CASES" in imported_names: + two_gpu_module = node.module + + if one_gpu_module is None or two_gpu_module is None: + raise RuntimeError( + "run_suite.py must import BOTH ONE_GPU_CASES and TWO_GPU_CASES." + ) + if one_gpu_module != two_gpu_module: + raise RuntimeError( + "run_suite.py imports ONE_GPU_CASES and TWO_GPU_CASES from different " + f"modules: {one_gpu_module} vs {two_gpu_module}" + ) + + rel_path = Path(*one_gpu_module.split(".")).with_suffix(".py") + candidates = [repo_root / rel_path, repo_root / "python" / rel_path] + case_config_path = next((path for path in candidates if path.exists()), None) + if case_config_path is None: + raise FileNotFoundError( + "Resolved case config from run_suite does not exist. Checked: " + + ", ".join(str(path) for path in candidates) + ) + return case_config_path + + class RunSuiteVisitor(ast.NodeVisitor): """ AST visitor to extract standalone metadata from run_suite.py. @@ -217,7 +260,7 @@ def get_case_est_time(case_id: str, baselines: Dict[str, float]) -> float: def parse_testcase_configs(config_path: Path) -> Dict[str, List[str]]: """ - Parse testcase_configs.py to extract case IDs. + Parse a diffusion case config file to extract case IDs. Returns: Dictionary mapping list name to case IDs. @@ -272,7 +315,7 @@ def validate_standalone_est_times( def collect_diffusion_suites( - testcase_config_path: Path, + case_config_path: Path, run_suite_path: Path, baseline_path: Path, ) -> Dict[str, DiffusionSuiteInfo]: @@ -280,15 +323,15 @@ def collect_diffusion_suites( Collect all diffusion test suite information using AST parsing. Args: - testcase_config_path: Path to testcase_configs.py + case_config_path: Path to case config (resolved from run_suite.py) run_suite_path: Path to run_suite.py baseline_path: Path to perf_baselines.json Returns: Dictionary mapping suite name to DiffusionSuiteInfo. """ - # Parse case IDs from testcase_configs.py - case_lists = parse_testcase_configs(testcase_config_path) + # Parse case IDs from the single source case config. + case_lists = parse_testcase_configs(case_config_path) # Parse standalone files from run_suite.py standalone_files, standalone_est_times = parse_run_suite_standalone_data( diff --git a/scripts/ci/utils/diffusion/verify_diffusion_coverage.py b/scripts/ci/utils/diffusion/verify_diffusion_coverage.py index 5659fa828..b327e38bb 100755 --- a/scripts/ci/utils/diffusion/verify_diffusion_coverage.py +++ b/scripts/ci/utils/diffusion/verify_diffusion_coverage.py @@ -21,8 +21,8 @@ from pathlib import Path from diffusion_case_parser import ( BASELINE_REL_PATH, RUN_SUITE_REL_PATH, - TESTCASE_CONFIG_REL_PATH, collect_diffusion_suites, + resolve_case_config_path, ) DYNAMIC_SUITES = {"1-gpu", "2-gpu"} @@ -39,18 +39,18 @@ def load_execution_reports(reports_dir: Path) -> list[dict]: def get_expected_cases(repo_root: Path) -> dict[str, set[str]]: """ - Get all expected cases from testcase_configs.py and run_suite.py. + Get all expected cases from case config and run_suite.py. Returns: Dictionary mapping suite name to set of expected case IDs. Standalone files are represented as "standalone:". """ - testcase_config_path = repo_root / TESTCASE_CONFIG_REL_PATH baseline_path = repo_root / BASELINE_REL_PATH run_suite_path = repo_root / RUN_SUITE_REL_PATH + case_config_path = resolve_case_config_path(repo_root, run_suite_path) suites = collect_diffusion_suites( - testcase_config_path, + case_config_path, run_suite_path, baseline_path, ) @@ -65,6 +65,21 @@ def get_expected_cases(repo_root: Path) -> dict[str, set[str]]: case_ids.add(f"standalone:{standalone_file}") expected[suite_name] = case_ids + empty_dynamic_suites = [ + suite_name + for suite_name in DYNAMIC_SUITES + if suite_name in expected + and not any( + not case_id.startswith("standalone:") for case_id in expected[suite_name] + ) + ] + if empty_dynamic_suites: + raise RuntimeError( + "Parsed zero parametrized cases for diffusion suites: " + + ", ".join(sorted(empty_dynamic_suites)) + + ". Refuse to pass coverage verification." + ) + return expected @@ -268,7 +283,11 @@ def main(): sys.exit(1) # Get expected cases - expected = get_expected_cases(repo_root) + try: + expected = get_expected_cases(repo_root) + except (RuntimeError, FileNotFoundError) as exc: + print(f"\nERROR: {exc}") + sys.exit(1) print("\nExpected cases by suite:") for suite, cases in expected.items(): print(f" {suite}: {len(cases)} cases")