[None][infra] Select modeling_v2 CI tests by the node a change lands in

jenkins/scripts/cbts/cbts/command/main.py

@@ -101,7 +101,7 @@ def build_rules(
         TestListRule(yaml_index, stages, repo_root=repo_root),
         VisualGenRule(yaml_index, stages),
         SpecDecRule(yaml_index, stages),
-        ModelingV2Rule(yaml_index, stages),
+        ModelingV2Rule(yaml_index, stages, repo_root=repo_root),
         AgentFlowRule(yaml_index, stages),
         OpenEngineRule(yaml_index, stages),
         DocsRule(yaml_index, stages),

jenkins/scripts/cbts/cbts/rules/modeling_v2.py

@@ -18,17 +18,41 @@
 (checkpoint, GPU arch, parallel topology), assembled from a catalog of
 op wrappers.
 
-Block selection — entry-pattern based only:
-It has no `condition.terms.backend` of its own; its entries sit in
-`backend: pytorch` blocks beside everything else. A block belongs to
-modeling_v2 iff one of its `tests:` entries matches a marker in
-`_MV2_ENTRY_PATTERNS`. Those markers are exact by construction rather
-than by luck: every unit test file in the subtree is named
-`test_modeling_v2_*` precisely so it cannot collide with the upstream
+Block selection: a block belongs to modeling_v2 iff one of its `tests:`
+entries matches a marker in `_MV2_ENTRY_PATTERNS`. Those markers are exact
+by construction rather than by luck: every unit test file in the subtree is
+named `test_modeling_v2_*` precisely so it cannot collide with the upstream
 test of the same op, and the accuracy gates live in their own
-`accuracy/modeling_v2/` directory. So there is no substring that could
-claim an unrelated entry, and no `mtp_nextn=0`-style carve-out is
-needed.
+`accuracy/modeling_v2/` directory.
+
+Entry selection: each changed file maps to the tests it can affect.
+
+- `models/` is a tree of nodes. A directory is a node when it, or a
+  directory below it, holds a target's `modeling.py`; a family directory
+  (`models/<family>/`) is always one. Other directories -- `ops/` and the
+  like -- serve the node above them. A changed file belongs to the nearest
+  node at or above it, and selects that node's unit tests
+  (`unittest/_torch/modeling_v2/models/<node>`) plus the accuracy gate of
+  every target under it. A target's gate is `Test<TargetClass>` in
+  `accuracy/modeling_v2/test_<family>.py`, the class name read from the
+  import line in the family's `routing.py` that routes to it.
+- A target's `cells.py` is input configurations for catalog ops and
+  nothing else: it changes no forward. It selects only the unit tests of
+  the catalog ops it imports, at their mirrored paths
+  (`catalog/<area>/<op>.py` -> `unittest/_torch/modeling_v2/<area>/test_modeling_v2_<op>.py`).
+- Anything else -- the catalog, the resolver, `models/__init__.py` -- can
+  reach every target and selects every modeling_v2 entry.
+
+The tests directly under `unittest/_torch/modeling_v2/` check the tree as a
+whole (claims, routing, contracts), so any change keeps them.
+
+A unit test prefix that no entry of a block sits at or under keeps the
+longest entry of that block that covers it instead: a family is one
+directory entry, and a change to one of its targets runs the family's
+unit tests. Whenever the tree on disk cannot answer -- a node directory
+deleted with its `routing.py`, a `cells.py` that cannot be parsed or names
+an op whose test is not at the mirrored path -- the file falls back to the
+wider selection.
 
 Outward fallback: not needed, and that is a property of the design
 rather than an accident. Nothing beyond the resolver module is imported
@@ -60,9 +84,12 @@
 
 from __future__ import annotations
 
+import ast
+import re
+from pathlib import Path
 from typing import Optional
 
-from cbts.blocks import Stage, YAMLIndex, _entry_target
+from cbts.blocks import Stage, YAMLIndex, _entry_target, _target_in_filter_subtree
 
 from ._helpers import resolve_affected_stages, stages_by_yaml_stem
 from .base import PRInputs, Rule, RuleResult
@@ -71,20 +98,27 @@
 # Source-path prefixes the rule may claim. Tests under tests/** are left
 # to TestsDefRule; the two scopes combine via _TESTSONLY_FAMILY.
 _MV2_SRC_PREFIXES: tuple[str, ...] = ("tensorrt_llm/_torch/_experimental/modeling_v2/",)
+_MV2_PACKAGE = "tensorrt_llm._torch._experimental.modeling_v2"
+
+_UNIT_ROOT = "unittest/_torch/modeling_v2"
+_ACCURACY_ROOT = "accuracy/modeling_v2"
 
 # Substrings that mark a test entry as modeling_v2. All are unambiguous:
-#   - "unittest/_torch/modeling_v2/" → the op-level catalog matrix, taken
-#     as whole-directory entries (one on l0_gb300, one on
-#     l0_gb300_multi_gpus for the 4-rank collectives)
+#   - "unittest/_torch/modeling_v2/" → the op and target unit tests
 #   - "accuracy/modeling_v2/" → the per-model accuracy gates
 #   - "test_modeling_v2_" → any future unit file named by the subtree's
 #     own convention
 _MV2_ENTRY_PATTERNS: tuple[str, ...] = (
-    "unittest/_torch/modeling_v2/",
-    "accuracy/modeling_v2/",
+    f"{_UNIT_ROOT}/",
+    f"{_ACCURACY_ROOT}/",
     "test_modeling_v2_",
 )
 
+# A family's `routing.py` imports each target it routes to by this line, from
+# the target directory directly under the family. The claims test reads the
+# same line with the same pattern.
+_TARGET_IMPORT = re.compile(r"^\s*from \.(\w+)\.modeling import (\w+)$", re.MULTILINE)
+
 
 def _is_mv2_claim(path: str) -> bool:
     """Decide whether ModelingV2Rule claims `path`.
@@ -110,31 +144,130 @@ def _mv2_entries(block) -> list[str]:
     return [t for t in block.tests if _entry_is_mv2(t)]
 
 
+def _is_whole_tree_test(target: str) -> bool:
+    """A test directly under the unit root, which checks the tree as a whole."""
+    rest = target.removeprefix(f"{_UNIT_ROOT}/").split("::")[0]
+    return target.startswith(f"{_UNIT_ROOT}/") and "/" not in rest and rest.endswith(".py")
+
+
 def _is_perf_sanity_stem(stem: str) -> bool:
     """True for perf-sanity yaml stems (`l0_*_perf_sanity*`)."""
     return "perf_sanity" in stem
 
 
+def _kept_entries(entries: list[str], prefixes: set[str]) -> list[str]:
+    """The entries of one block that `prefixes` select.
+
+    An entry is kept when it sits at or under a prefix. A prefix with no such
+    entry in the block keeps the longest entry that covers it instead.
+    """
+    targets = {entry: _entry_target(entry) for entry in entries}
+    kept = {e for e, t in targets.items() if _is_whole_tree_test(t)}
+    for prefix in prefixes:
+        under = {e for e, t in targets.items() if _target_in_filter_subtree(t, prefix)}
+        if not under:
+            covering = [e for e, t in targets.items() if _target_in_filter_subtree(prefix, t)]
+            under = {max(covering, key=lambda e: len(targets[e]))} if covering else set()
+        kept |= under
+    return [e for e in entries if e in kept]
+
+
 class ModelingV2Rule(Rule):
     name = "modelingv2"
     needs_diff_for: tuple[str, ...] = ()
 
-    def __init__(self, yaml_index: YAMLIndex, stages: dict[str, Stage]) -> None:
+    def __init__(self, yaml_index: YAMLIndex, stages: dict[str, Stage], repo_root: Path) -> None:
         self.yaml_index = yaml_index
         self._stages_by_yaml = stages_by_yaml_stem(stages)
+        self._src_root = Path(repo_root) / _MV2_SRC_PREFIXES[0]
+        self._tests_root = Path(repo_root) / "tests"
+
+    def _is_node(self, rel_dir: Path) -> bool:
+        """Whether `models/<rel_dir>` is a node: a family, or holds a target at or below it."""
+        path = self._src_root / "models" / rel_dir
+        return len(rel_dir.parts) == 1 or (path.is_dir() and any(path.rglob("modeling.py")))
+
+    def _gates_under(self, node: Path) -> Optional[set[str]]:
+        """Accuracy-gate prefixes of every target at or under `node`; None if unreadable."""
+        family = node.parts[0]
+        routing = self._src_root / "models" / family / "routing.py"
+        if not routing.is_file():
+            return None
+        gates = set()
+        for target, target_cls in _TARGET_IMPORT.findall(routing.read_text(encoding="utf-8")):
+            leaf = Path(family, target)
+            if leaf == node or node in leaf.parents:
+                gates.add(f"{_ACCURACY_ROOT}/test_{family}.py::Test{target_cls}")
+        return gates
+
+    def _cells_tests(self, rel: Path) -> Optional[set[str]]:
+        """Unit tests of the catalog ops a `cells.py` imports; None if not resolvable."""
+        try:
+            tree = ast.parse((self._src_root / rel).read_text(encoding="utf-8"))
+        except (OSError, SyntaxError, UnicodeDecodeError):
+            return None
+        package = rel.parent.parts
+        tests = set()
+        for node in ast.walk(tree):
+            if not isinstance(node, ast.ImportFrom):
+                continue
+            if node.level:
+                if node.level > len(package) + 1:
+                    return None
+                base = package[: len(package) - node.level + 1]
+                module = (*base, *(node.module or "").split("."))
+            elif (node.module or "").startswith(_MV2_PACKAGE + "."):
+                module = tuple(node.module[len(_MV2_PACKAGE) + 1 :].split("."))
+            else:
+                continue
+            if module[:1] != ("catalog",) or len(module) != 3 or module[1].startswith("_"):
+                if module == ("catalog", "_op"):
+                    continue
+                return None
+            test = f"{_UNIT_ROOT}/{module[1]}/test_modeling_v2_{module[2]}.py"
+            if not (self._tests_root / test).is_file():
+                return None
+            tests.add(test)
+        return tests
+
+    def _prefixes_for(self, path: str) -> Optional[set[str]]:
+        """Filter prefixes a changed file selects; None selects every modeling_v2 entry."""
+        rel = Path(path[len(_MV2_SRC_PREFIXES[0]) :])
+        if rel.parts[0] != "models" or len(rel.parts) < 3:
+            return None
+        if rel.name == "cells.py":
+            tests = self._cells_tests(rel)
+            if tests is not None:
+                return tests
+        node = Path(*rel.parts[1:-1])
+        while not self._is_node(node):
+            node = node.parent
+        gates = self._gates_under(node)
+        if gates is None:
+            return None
+        return {f"{_UNIT_ROOT}/models/{node.as_posix()}"} | gates
 
     def apply(self, pr: PRInputs) -> Optional[RuleResult]:
         claimed = {f for f in pr.changed_files if _is_mv2_claim(f)}
         if not claimed:
             return None
 
+        prefixes: Optional[set[str]] = set()
+        for path in sorted(claimed):
+            selected = self._prefixes_for(path)
+            if selected is None:
+                prefixes = None
+                break
+            prefixes |= selected
+
         block_filters: dict[tuple[str, int], dict[str, set[str]]] = {}
         for block in self.yaml_index.blocks:
             entries = _mv2_entries(block)
+            if prefixes is not None:
+                entries = _kept_entries(entries, prefixes)
             if not entries:
                 continue
-            key = (block.yaml_stem, block.block_index)
-            prefix_dict = block_filters.setdefault(key, {})
+            prefix_dict = block_filters.setdefault((block.yaml_stem, block.block_index), {})
             for entry in entries:
                 target = _entry_target(entry)
                 if target:
@@ -168,6 +301,7 @@ def apply(self, pr: PRInputs) -> Optional[RuleResult]:
             perfsanity_relevant=perfsanity_relevant,
             reason=(
                 f"modelingv2: {len(claimed)} modeling_v2 source file(s) → "
-                f"{len(block_filters)} modeling_v2 block(s), {len(affected)} stage(s)"
+                f"{'every modeling_v2 entry' if prefixes is None else f'{len(prefixes)} test prefix(es)'}, "
+                f"{len(block_filters)} block(s), {len(affected)} stage(s)"
             ),
         )

tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py

@@ -27,7 +27,6 @@
 """
 
 from ._router_index import (
-    MODELING_V2_ROUTERS,
     MODELING_V2_VALIDATE_ENV,
     ModelingV2Mode,
     NoTarget,
@@ -36,7 +35,6 @@
 )
 
 __all__ = [
-    "MODELING_V2_ROUTERS",
     "MODELING_V2_VALIDATE_ENV",
     "ModelingV2Mode",
     "NoTarget",