chore: import upstream snapshot with attribution

scripts/ci/utils/diffusion/diffusion_case_parser.py

@@ -0,0 +1,517 @@
+#!/usr/bin/env python3
+"""
+AST-based parser for diffusion test cases.
+
+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, 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
+import json
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Dict, List, Optional
+
+# Mapping from list variable names to suite names
+CASE_LIST_TO_SUITE = {
+    "ONE_GPU_CASES": "1-gpu",
+    "ONE_GPU_CASES_A": "1-gpu",
+    "ONE_GPU_CASES_B": "1-gpu",
+    "ONE_GPU_CASES_C": "1-gpu-b200",
+    "ONE_GPU_MODELOPT_FP8_CASES": "1-gpu",
+    "ONE_GPU_MODELOPT_CASES": "1-gpu-b200",
+    "ONE_GPU_B200_CASES": "1-gpu-b200",
+    "TWO_GPU_CASES": "2-gpu",
+    "TWO_GPU_CASES_A": "2-gpu",
+    "TWO_GPU_CASES_B": "2-gpu",
+}
+
+# Default estimated time for cases without baseline (5 minutes)
+DEFAULT_EST_TIME_SECONDS = 300.0
+
+# Fixed overhead for server startup when estimated_full_test_time_s is not set
+STARTUP_OVERHEAD_SECONDS = 120.0
+
+# Paths relative to repository root
+BASELINE_REL_PATH = "python/sglang/multimodal_gen/test/server/perf_baselines"
+BASELINE_PLATFORM_ORDER = ("h100", "b200", "5090")
+RUN_SUITE_REL_PATH = "python/sglang/multimodal_gen/test/run_suite.py"
+
+USE_NPU_CONFIGS = os.getenv("USE_NPU_CONFIGS", "0").lower() in ("1", "true")
+
+if USE_NPU_CONFIGS:
+    BASELINE_REL_PATH = (
+        "python/sglang/multimodal_gen/test/server/perf_baselines_npu.json"
+    )
+    CASE_LIST_TO_SUITE = {
+        "ONE_NPU_CASES": "1-npu",
+        "TWO_NPU_CASES": "2-npu",
+    }
+
+
+@dataclass
+class DiffusionCaseInfo:
+    """Information about a single diffusion test case."""
+
+    case_id: str  # e.g., "qwen_image_t2i"
+    suite: str  # "1-gpu" or "2-gpu"
+    est_time: float  # estimated time in seconds
+
+
+@dataclass
+class DiffusionSuiteInfo:
+    """Complete information for a test suite."""
+
+    suite: str  # "1-gpu" or "2-gpu"
+    cases: List[DiffusionCaseInfo]  # parametrized test cases
+    standalone_files: List[str]  # standalone test files
+    standalone_est_times: Dict[str, float]  # standalone file -> estimated seconds
+    missing_standalone_estimates: List[
+        str
+    ]  # standalone files without configured estimate
+
+
+class DiffusionTestCaseVisitor(ast.NodeVisitor):
+    """
+    AST visitor to extract DiffusionTestCase definitions from the case config.
+
+    Parses assignments like:
+        ONE_GPU_CASES_A: list[DiffusionTestCase] = [
+            DiffusionTestCase("case_id", ...),
+            ...
+        ]
+    """
+
+    def __init__(self):
+        self.cases: Dict[str, List[str]] = {}  # list_name -> [case_id, ...]
+        self.factory_case_ids: Dict[str, str] = {}
+
+    def visit_Module(self, node: ast.Module):
+        for stmt in node.body:
+            if not isinstance(stmt, ast.FunctionDef):
+                continue
+            case_id = self._extract_factory_case_id(stmt)
+            if case_id:
+                self.factory_case_ids[stmt.name] = case_id
+
+        self.generic_visit(node)
+
+    def visit_Expr(self, node: ast.Expr):
+        """Handle ``LIST.append(...)`` mutations at any nesting level.
+
+        Previously only module-top-level ``ast.Expr`` statements were scanned for
+        ``.append()`` calls, so cases registered under a platform guard such as
+        ``if not current_platform.is_hip(): ONE_GPU_CASES.append(...)`` (used by
+        ``hunyuan3d_shape_gen`` and ``turbo_wan2_1_t2v_1.3b``) were invisible to
+        the partition planner and therefore never scheduled in CI. Visiting every
+        ``Expr`` lets ``generic_visit`` reach appends inside ``if``/``else`` blocks.
+        """
+        self._process_expr(node.value)
+        self.generic_visit(node)
+
+    def visit_Assign(self, node: ast.Assign):
+        self._process_assignment(node.targets, node.value)
+        self.generic_visit(node)
+
+    def visit_AnnAssign(self, node: ast.AnnAssign):
+        if node.target and node.value:
+            self._process_assignment([node.target], node.value)
+        self.generic_visit(node)
+
+    def visit_AugAssign(self, node: ast.AugAssign):
+        self._process_aug_assignment(node.target, node.op, node.value)
+        self.generic_visit(node)
+
+    def _process_assignment(self, targets: List[ast.AST], value: ast.AST):
+        """Process an assignment to extract case IDs."""
+        for target in targets:
+            if isinstance(target, ast.Name):
+                list_name = target.id
+                case_ids = self._extract_case_ids(value)
+                if case_ids is not None:
+                    self.cases[list_name] = case_ids
+
+    def _process_aug_assignment(self, target: ast.AST, op: ast.AST, value: ast.AST):
+        """Process `+=` style assignment to merge case lists."""
+        if not isinstance(target, ast.Name) or not isinstance(op, ast.Add):
+            return
+
+        if isinstance(value, ast.Name):
+            target_suite = CASE_LIST_TO_SUITE.get(target.id)
+            value_suite = CASE_LIST_TO_SUITE.get(value.id)
+            if target_suite and value_suite and target_suite != value_suite:
+                return
+
+        rhs_case_ids = self._extract_case_ids(value)
+        if rhs_case_ids is None:
+            return
+
+        lhs_case_ids = self.cases.get(target.id, [])
+        self.cases[target.id] = [*lhs_case_ids, *rhs_case_ids]
+
+    def _process_expr(self, node: ast.AST):
+        """Process list mutation calls such as `ONE_GPU_CASES.append(...)`."""
+        if not isinstance(node, ast.Call):
+            return
+        if not isinstance(node.func, ast.Attribute):
+            return
+        if node.func.attr != "append":
+            return
+        if not isinstance(node.func.value, ast.Name):
+            return
+        list_name = node.func.value.id
+        if list_name not in CASE_LIST_TO_SUITE:
+            return
+        if len(node.args) != 1:
+            return
+
+        case_id = self._extract_case_id_from_call(node.args[0])
+        if case_id:
+            self.cases.setdefault(list_name, []).append(case_id)
+
+    def _extract_case_ids(self, node: ast.AST) -> Optional[List[str]]:
+        """Extract case IDs from a supported expression."""
+        if isinstance(node, ast.List):
+            return self._extract_case_ids_from_list(node)
+
+        if isinstance(node, ast.Name):
+            # Reference to a previously parsed list variable.
+            if node.id not in self.cases:
+                return None
+            return list(self.cases[node.id])
+
+        if isinstance(node, ast.BinOp) and isinstance(node.op, ast.Add):
+            left_ids = self._extract_case_ids(node.left)
+            right_ids = self._extract_case_ids(node.right)
+            if left_ids is None or right_ids is None:
+                return None
+            return [*left_ids, *right_ids]
+
+        return None
+
+    def _extract_case_ids_from_list(self, node: ast.List) -> List[str]:
+        """Extract case IDs from a literal list of DiffusionTestCase calls."""
+        case_ids = []
+        for elt in node.elts:
+            if isinstance(elt, ast.Starred):
+                starred_case_ids = self._extract_case_ids(elt.value)
+                if starred_case_ids:
+                    case_ids.extend(starred_case_ids)
+                continue
+            case_id = self._extract_case_id_from_call(elt)
+            if case_id:
+                case_ids.append(case_id)
+        return case_ids
+
+    def _extract_case_id_from_call(self, node: ast.AST) -> Optional[str]:
+        """Extract case_id from DiffusionTestCase(...) call."""
+        if not isinstance(node, ast.Call):
+            return None
+
+        # First positional argument is the case_id.
+        if isinstance(node.func, ast.Name) and node.func.id in {
+            "DiffusionTestCase",
+            "_make_modelopt_ci_case",
+        }:
+            if node.args and isinstance(node.args[0], ast.Constant):
+                return node.args[0].value
+        if isinstance(node.func, ast.Name) and not node.args:
+            return self.factory_case_ids.get(node.func.id)
+
+        return None
+
+    def _extract_factory_case_id(self, node: ast.FunctionDef) -> Optional[str]:
+        for child in ast.walk(node):
+            if not isinstance(child, ast.Return) or child.value is None:
+                continue
+            case_id = self._extract_case_id_from_call(child.value)
+            if case_id:
+                return case_id
+        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.
+
+    Parses:
+        STANDALONE_FILES = {
+            "1-gpu": ["test_lora_format_adapter.py"],
+            "2-gpu": [],
+        }
+    """
+
+    def __init__(self):
+        self.standalone_files: Dict[str, List[str]] = {}
+        self.standalone_est_times: Dict[str, Dict[str, float]] = {}
+
+    def visit_Assign(self, node: ast.Assign):
+        for target in node.targets:
+            if isinstance(target, ast.Name) and target.id == "STANDALONE_FILES":
+                self.standalone_files = self._extract_file_dict(node.value)
+            if (
+                isinstance(target, ast.Name)
+                and target.id == "STANDALONE_FILE_EST_TIMES"
+            ):
+                self.standalone_est_times = self._extract_est_time_dict(node.value)
+        self.generic_visit(node)
+
+    def _extract_file_dict(self, node: ast.AST) -> Dict[str, List[str]]:
+        """Extract dictionary of suite -> file list."""
+        result = {}
+        if isinstance(node, ast.Dict):
+            for key, value in zip(node.keys, node.values):
+                if isinstance(key, ast.Constant) and isinstance(value, ast.List):
+                    suite = key.value
+                    files = [
+                        elt.value for elt in value.elts if isinstance(elt, ast.Constant)
+                    ]
+                    result[suite] = files
+        return result
+
+    def _extract_est_time_dict(self, node: ast.AST) -> Dict[str, Dict[str, float]]:
+        """Extract dictionary of suite -> standalone file -> estimated seconds."""
+        result = {}
+        if not isinstance(node, ast.Dict):
+            return result
+
+        for key, value in zip(node.keys, node.values):
+            if not isinstance(key, ast.Constant) or not isinstance(value, ast.Dict):
+                continue
+
+            suite = key.value
+            suite_est_times = {}
+            for inner_key, inner_value in zip(value.keys, value.values):
+                if not (
+                    isinstance(inner_key, ast.Constant)
+                    and isinstance(inner_value, ast.Constant)
+                ):
+                    continue
+                suite_est_times[inner_key.value] = float(inner_value.value)
+            result[suite] = suite_est_times
+
+        return result
+
+
+def _iter_baseline_paths(baseline_path: Path) -> List[Path]:
+    if baseline_path.is_file():
+        return [baseline_path]
+    if not baseline_path.is_dir():
+        return []
+
+    ordered_paths = [
+        baseline_path / f"{platform}.json" for platform in BASELINE_PLATFORM_ORDER
+    ]
+    ordered_paths.extend(
+        path
+        for path in sorted(baseline_path.glob("*.json"))
+        if path not in ordered_paths
+    )
+    return [path for path in ordered_paths if path.exists()]
+
+
+def load_baselines(baseline_path: Path) -> Dict[str, float]:
+    """
+    Load performance baselines from a JSON file or platform baseline directory.
+
+    Returns:
+        Dictionary mapping case_id to estimated time in seconds.
+    """
+    baselines = {}
+    for path in _iter_baseline_paths(baseline_path):
+        with open(path, "r", encoding="utf-8") as f:
+            data = json.load(f)
+
+        scenarios = data.get("scenarios", {})
+        for case_id, scenario in scenarios.items():
+            if scenario.get("estimated_full_test_time_s") is not None:
+                est_time = scenario["estimated_full_test_time_s"]
+            else:
+                expected_e2e_ms = scenario.get("expected_e2e_ms", 0)
+                est_time = expected_e2e_ms / 1000.0 + STARTUP_OVERHEAD_SECONDS
+            baselines.setdefault(case_id, est_time)
+
+    return baselines
+
+
+def get_case_est_time(case_id: str, baselines: Dict[str, float]) -> float:
+    """Get estimated time for a case, with fallback to default."""
+    return baselines.get(case_id, DEFAULT_EST_TIME_SECONDS)
+
+
+def parse_testcase_configs(config_path: Path) -> Dict[str, List[str]]:
+    """
+    Parse a diffusion case config file to extract case IDs.
+
+    Returns:
+        Dictionary mapping list name to case IDs.
+        e.g., {"ONE_GPU_CASES_A": ["qwen_image_t2i", ...], ...}
+    """
+    with open(config_path, "r", encoding="utf-8") as f:
+        content = f.read()
+
+    tree = ast.parse(content, filename=str(config_path))
+    visitor = DiffusionTestCaseVisitor()
+    visitor.visit(tree)
+
+    return visitor.cases
+
+
+def parse_run_suite_standalone_data(
+    run_suite_path: Path,
+) -> tuple[Dict[str, List[str]], Dict[str, Dict[str, float]]]:
+    """
+    Parse run_suite.py to extract standalone file metadata.
+
+    Returns:
+        Tuple of:
+          - suite -> standalone file list
+          - suite -> standalone file -> estimated seconds
+    """
+    with open(run_suite_path, "r", encoding="utf-8") as f:
+        content = f.read()
+
+    tree = ast.parse(content, filename=str(run_suite_path))
+    visitor = RunSuiteVisitor()
+    visitor.visit(tree)
+
+    return visitor.standalone_files, visitor.standalone_est_times
+
+
+def validate_standalone_est_times(
+    standalone_files: Dict[str, List[str]],
+    standalone_est_times: Dict[str, Dict[str, float]],
+) -> Dict[str, List[str]]:
+    missing_by_suite = {}
+    for suite, files in standalone_files.items():
+        suite_est_times = standalone_est_times.get(suite, {})
+        missing = [
+            standalone_file
+            for standalone_file in files
+            if standalone_file not in suite_est_times
+        ]
+        if missing:
+            missing_by_suite[suite] = missing
+    return missing_by_suite
+
+
+def collect_diffusion_suites(
+    case_config_path: Path,
+    run_suite_path: Path,
+    baseline_path: Path,
+) -> Dict[str, DiffusionSuiteInfo]:
+    """
+    Collect all diffusion test suite information using AST parsing.
+
+    Args:
+        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/ or a single baseline JSON file
+
+    Returns:
+        Dictionary mapping suite name to DiffusionSuiteInfo.
+    """
+    # 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(
+        run_suite_path
+    )
+    missing_standalone_estimates = validate_standalone_est_times(
+        standalone_files, standalone_est_times
+    )
+
+    # Load baselines for time estimation
+    baselines = load_baselines(baseline_path)
+
+    # Build suite info
+    suites = {}
+    for list_name, suite in CASE_LIST_TO_SUITE.items():
+        case_ids = case_lists.get(list_name, [])
+        cases = [
+            DiffusionCaseInfo(
+                case_id=cid,
+                suite=suite,
+                est_time=get_case_est_time(cid, baselines),
+            )
+            for cid in case_ids
+        ]
+
+        if suite not in suites:
+            suites[suite] = DiffusionSuiteInfo(
+                suite=suite,
+                cases=[],
+                standalone_files=standalone_files.get(suite, []),
+                standalone_est_times=dict(standalone_est_times.get(suite, {})),
+                missing_standalone_estimates=list(
+                    missing_standalone_estimates.get(suite, [])
+                ),
+            )
+        suites[suite].cases.extend(cases)
+
+    # Dedupe duplicated case IDs while preserving first-seen order.
+    for suite_info in suites.values():
+        seen_case_ids = set()
+        deduped_cases = []
+        for case in suite_info.cases:
+            if case.case_id in seen_case_ids:
+                continue
+            seen_case_ids.add(case.case_id)
+            deduped_cases.append(case)
+        suite_info.cases = deduped_cases
+
+    return suites