[diffusion] optimize: optimize lingbot performance (#27023)

scripts/ci/utils/diffusion/diffusion_case_parser.py

@@ -80,6 +80,21 @@ class DiffusionTestCaseVisitor(ast.NodeVisitor):
 
     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
+
+        for stmt in node.body:
+            if isinstance(stmt, ast.Expr):
+                self._process_expr(stmt.value)
+
+        self.generic_visit(node)
 
     def visit_Assign(self, node: ast.Assign):
         self._process_assignment(node.targets, node.value)
@@ -121,6 +136,26 @@ class DiffusionTestCaseVisitor(ast.NodeVisitor):
         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):
@@ -167,9 +202,20 @@ class DiffusionTestCaseVisitor(ast.NodeVisitor):
         }:
             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:
     """