Resolve file-run labels from the PR snapshot

tests/ci/ci_register.py

@@ -158,8 +158,9 @@ def _extract_list_constant(node: ast.AST, *, context: str = "value") -> list:
 
 
 class RegistryVisitor(ast.NodeVisitor):
-    def __init__(self, filename: str):
+    def __init__(self, filename: str, *, known_labels=KNOWN_LABELS):
         self.filename = filename
+        self.known_labels = known_labels
         self.registries: list[CIRegistry] = []
 
     def _parse_call_args(self, func_call: ast.Call, func_name: str) -> CIRegistry:
@@ -234,9 +235,9 @@ def _parse_call_args(self, func_call: ast.Call, func_name: str) -> CIRegistry:
                 f"CPU has no GPU generation and ROCm is a separate backend"
             )
 
-        unknown = [label for label in labels if label not in KNOWN_LABELS]
+        unknown = [label for label in labels if label not in self.known_labels]
         if unknown:
-            valid_list = ", ".join(sorted(KNOWN_LABELS))
+            valid_list = ", ".join(sorted(self.known_labels))
             raise ValueError(
                 f"{self.filename}: unknown labels {unknown} in {func_name}(); "
                 f"valid labels: [{valid_list}]. "
@@ -302,11 +303,11 @@ def visit_Module(self, node):
                 self.registries.append(cr)
 
 
-def ut_parse_one_file(filename: str) -> list[CIRegistry]:
+def ut_parse_one_file(filename: str, *, known_labels=KNOWN_LABELS) -> list[CIRegistry]:
     with open(filename) as f:
         file_content = f.read()
     tree = ast.parse(file_content, filename=filename)
-    visitor = RegistryVisitor(filename=filename)
+    visitor = RegistryVisitor(filename=filename, known_labels=known_labels)
     visitor.visit(tree)
     return visitor.registries
 
@@ -363,10 +364,10 @@ def _make_implicit_cpu_registry(filename: str) -> CIRegistry:
     )
 
 
-def collect_tests(files: list[str], sanity_check: bool = True) -> list[CIRegistry]:
+def collect_tests(files: list[str], sanity_check: bool = True, *, known_labels=KNOWN_LABELS) -> list[CIRegistry]:
     ci_tests: list[CIRegistry] = []
     for file in files:
-        registries = ut_parse_one_file(file)
+        registries = ut_parse_one_file(file, known_labels=known_labels)
         if _is_implicit_fast_cpu_path(file):
             # tests/fast/ is CPU-only by location;
             for r in registries:

tests/ci/file_run.py

@@ -5,6 +5,7 @@
 may import only the stdlib and the dependency-free registry modules.
 """
 
+import ast
 import json
 import os
 import re
@@ -104,6 +105,31 @@ def plan_file_run(all_tests, test_file: str, image_tag: str) -> dict[str, str]:
     }
 
 
+def _read_snapshot_labels() -> dict[str, str]:
+    path = Path("tests/ci/labels.py")
+    if path.is_symlink():
+        raise FileRunError(f"CI label registry must not be a symlink: {path}")
+    try:
+        tree = ast.parse(path.read_text(), filename=str(path))
+        definitions = [
+            node.value
+            for node in tree.body
+            if isinstance(node, ast.AnnAssign)
+            and isinstance(node.target, ast.Name)
+            and node.target.id == "KNOWN_LABELS"
+        ]
+        if len(definitions) != 1:
+            raise FileRunError(f"{path} must define KNOWN_LABELS once as a literal dictionary")
+        labels = ast.literal_eval(definitions[0])
+    except (OSError, SyntaxError, ValueError, TypeError) as error:
+        raise FileRunError(f"cannot read CI label registry {path}: {error}") from error
+    if not isinstance(labels, dict) or not all(
+        isinstance(key, str) and isinstance(value, str) for key, value in labels.items()
+    ):
+        raise FileRunError(f"{path}: KNOWN_LABELS must be a literal string-to-string dictionary")
+    return labels
+
+
 def collect_snapshot_tests(source_root: str | Path):
     """Parse a source snapshot as data, without importing its Python modules."""
     try:
@@ -116,7 +142,8 @@ def collect_snapshot_tests(source_root: str | Path):
     previous_directory = Path.cwd()
     try:
         os.chdir(root)
-        return collect_tests(_discover_regular_ci_files(), sanity_check=True)
+        files = _discover_regular_ci_files()
+        return collect_tests(files, sanity_check=True, known_labels=_read_snapshot_labels())
     finally:
         os.chdir(previous_directory)