[diffusion] CI: fix auto-partition (#23076)

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)")

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(

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:<filename>".
     """
-    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")