[diffusion] CI: minor refactor CI (#23576)

python/sglang/multimodal_gen/test/run_suite.py

@@ -65,7 +65,7 @@ FILE_SUITES = {
         "test_component_accuracy_2_gpu.py",
     ],
     "1-gpu-b200": [
-        "test_server_c.py",
+        "test_server_b200.py",
     ],
 }
 

python/sglang/multimodal_gen/test/server/gpu_cases.py

@@ -46,7 +46,7 @@ from sglang.multimodal_gen.test.test_utils import (
 
 # All test cases with clean default values
 # To test different models, simply add more DiffusionCase entries
-ONE_GPU_CASES_A: list[DiffusionTestCase] = [
+ONE_GPU_CASES: list[DiffusionTestCase] = [
     # === Text to Image (T2I) ===
     DiffusionTestCase(
         "qwen_image_t2i",
@@ -167,9 +167,6 @@ ONE_GPU_CASES_A: list[DiffusionTestCase] = [
             extras={"enable_upscaling": True, "upscaling_scale": 4},
         ),
     ),
-]
-
-ONE_GPU_CASES_B: list[DiffusionTestCase] = [
     # === Text to Video (T2V) ===
     DiffusionTestCase(
         "wan2_1_t2v_1.3b",
@@ -339,7 +336,7 @@ ONE_GPU_CASES_B: list[DiffusionTestCase] = [
 
 # Skip hunyuan3d on AMD: marching_cubes surface extraction produces invalid SDF on ROCm.
 if not current_platform.is_hip():
-    ONE_GPU_CASES_B.append(
+    ONE_GPU_CASES.append(
         DiffusionTestCase(
             "hunyuan3d_shape_gen",
             DiffusionServerArgs(
@@ -352,7 +349,7 @@ if not current_platform.is_hip():
     )
 # Skip turbowan on AMD: Triton requires 81920 shared memory, but AMD only has 65536.
 if not current_platform.is_hip():
-    ONE_GPU_CASES_B.append(
+    ONE_GPU_CASES.append(
         DiffusionTestCase(
             "turbo_wan2_1_t2v_1.3b",
             DiffusionServerArgs(
@@ -364,9 +361,9 @@ if not current_platform.is_hip():
 # Skip all ModelOpt tests on AMD: FP8 requires torch._scaled_mm (HIPBLAS_STATUS_NOT_SUPPORTED
 # on ROCm), NVFP4 requires flashinfer or sgl_kernel FP4 kernels (CUDA-only)
 if current_platform.is_hip():
-    ONE_GPU_CASES_C = []
+    ONE_GPU_MODELOPT_CASES = []
 else:
-    ONE_GPU_CASES_C = [
+    ONE_GPU_MODELOPT_CASES = [
         _make_modelopt_ci_case(
             "flux1_modelopt_fp8_t2i",
             model_path=DEFAULT_FLUX_1_DEV_MODEL_NAME_FOR_TEST,
@@ -414,7 +411,7 @@ else:
         ),
     ]
 
-TWO_GPU_CASES_A = [
+TWO_GPU_CASES = [
     DiffusionTestCase(
         "wan2_2_i2v_a14b_2gpu",
         DiffusionServerArgs(
@@ -528,9 +525,6 @@ TWO_GPU_CASES_A = [
         ),
         TI2V_sampling_params,
     ),
-]
-
-TWO_GPU_CASES_B = [
     DiffusionTestCase(
         "wan2_1_i2v_14b_480P_2gpu",
         DiffusionServerArgs(
@@ -632,7 +626,7 @@ TWO_GPU_CASES_B = [
 
 if not current_platform.is_hip():
     # Flux2 multi-image edit with cache-dit, regression test
-    ONE_GPU_CASES_B.append(
+    ONE_GPU_CASES.append(
         DiffusionTestCase(
             "flux_2_ti2i_multi_image_cache_dit",
             DiffusionServerArgs(
@@ -643,7 +637,5 @@ if not current_platform.is_hip():
         )
     )
 
-ONE_GPU_CASES = [*ONE_GPU_CASES_A, *ONE_GPU_CASES_B, *ONE_GPU_CASES_C]
+ONE_GPU_CASES += ONE_GPU_MODELOPT_CASES
-TWO_GPU_CASES_A = _with_default_num_gpus(TWO_GPU_CASES_A, 2)
+TWO_GPU_CASES = _with_default_num_gpus(TWO_GPU_CASES, 2)
-TWO_GPU_CASES_B = _with_default_num_gpus(TWO_GPU_CASES_B, 2)
-TWO_GPU_CASES = [*TWO_GPU_CASES_A, *TWO_GPU_CASES_B]

python/sglang/multimodal_gen/test/server/test_server_b200.py

@@ -7,7 +7,7 @@ from __future__ import annotations
 import pytest
 
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
-from sglang.multimodal_gen.test.server.gpu_cases import ONE_GPU_CASES_C
+from sglang.multimodal_gen.test.server.gpu_cases import ONE_GPU_MODELOPT_CASES
 from sglang.multimodal_gen.test.server.test_server_common import (  # noqa: F401
     DiffusionServerBase,
     diffusion_server,
@@ -20,7 +20,7 @@ logger = init_logger(__name__)
 class TestDiffusionServerOneGpuB200(DiffusionServerBase):
     """B200-targeted CI tests for 1-GPU ModelOpt diffusion cases."""
 
-    @pytest.fixture(params=ONE_GPU_CASES_C, ids=lambda c: c.id)
+    @pytest.fixture(params=ONE_GPU_MODELOPT_CASES, ids=lambda c: c.id)
     def case(self, request) -> DiffusionTestCase:
         """Provide a DiffusionTestCase for each 1-GPU B200 test."""
         return request.param

python/sglang/multimodal_gen/test/server/testcase_configs.py

@@ -5,12 +5,12 @@ Usage:
 
 pytest python/sglang/multimodal_gen/test/server/test_server_1_gpu.py
 # for a single testcase, look for the name of the testcase in ONE_GPU_CASES,
-# ONE_GPU_CASES_C, or TWO_GPU_CASES
+# ONE_GPU_MODELOPT_CASES, or TWO_GPU_CASES
 pytest python/sglang/multimodal_gen/test/server/test_server_1_gpu.py -k qwen_image_t2i
 
 
 To add a new testcase:
-1. add your testcase with case-id: `my_new_test_case_id` to the appropriate `*_CASES_*` list
+1. add your testcase with case-id: `my_new_test_case_id` to `ONE_GPU_CASES`, `ONE_GPU_MODELOPT_CASES`, or `TWO_GPU_CASES`
 2. run `SGLANG_GEN_BASELINE=1 pytest -s python/sglang/multimodal_gen/test/server/ -k my_new_test_case_id`
 3. insert or override the corresponding scenario in `scenarios` section of perf_baselines.json with the output baseline of step-2
 

scripts/ci/utils/diffusion/diffusion_case_parser.py

@@ -22,9 +22,12 @@ 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_CASES": "1-gpu-b200",
+    "TWO_GPU_CASES": "2-gpu",
     "TWO_GPU_CASES_A": "2-gpu",
     "TWO_GPU_CASES_B": "2-gpu",
 }
@@ -85,22 +88,60 @@ class DiffusionTestCaseVisitor(ast.NodeVisitor):
             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 if it's a known list."""
+        """Process an assignment to extract case IDs."""
         for target in targets:
-            if isinstance(target, ast.Name) and target.id in CASE_LIST_TO_SUITE:
+            if isinstance(target, ast.Name):
                 list_name = target.id
-                case_ids = self._extract_case_ids_from_list(value)
+                case_ids = self._extract_case_ids(value)
                 if case_ids is not None:
                     self.cases[list_name] = case_ids
 
-    def _extract_case_ids_from_list(self, node: ast.AST) -> Optional[List[str]]:
+    def _process_aug_assignment(self, target: ast.AST, op: ast.AST, value: ast.AST):
-        """Extract case IDs from a literal list of DiffusionTestCase calls."""
+        """Process `+=` style assignment to merge case lists."""
-        if not isinstance(node, ast.List):
+        if not isinstance(target, ast.Name) or not isinstance(op, ast.Add):
-            return None
+            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 _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)
@@ -369,4 +410,15 @@ def collect_diffusion_suites(
             )
         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