Test generated workflow matrix schemas

build_tools/github_actions/configure_jax_release_matrix.py

@@ -53,12 +53,18 @@ def generate_jax_matrix(
     matrix: list[dict[str, object]] = []
     for py in versions:
         for ref_cfg in JAX_REFS:
+            # These row keys are the contract with workflow files which use them
+            # via matrix.<key> expressions. Empty values are allowed when the
+            # workflow handles them explicitly, but undefined keys are not.
             matrix.append(
                 {
                     "python_version": py,
                     "jax_ref": ref_cfg["jax_ref"],
                     "jax_repository": ref_cfg["jax_repository"],
                     "build_mode": ref_cfg["build_mode"],
+                    # gfx_arch is intentionally empty for native JAX builds and
+                    # non-empty for manylinux builds. This direct lookup raises
+                    # KeyError if JAX_REFS omits the key.
                     "gfx_arch": ref_cfg["gfx_arch"],
                 }
             )

build_tools/github_actions/configure_pytorch_release_matrix.py

@@ -162,6 +162,9 @@ def generate_pytorch_matrix_for_release_type(
             families = _filter_families(amdgpu_families, exclude)
             if not families:
                 continue
+            # These row keys are the contract with workflow files which use them
+            # via matrix.<key> expressions. Empty values are allowed when the
+            # workflow handles them explicitly, but undefined keys are not.
             row: dict[str, str] = {
                 "python_version": py,
                 "pytorch_git_ref": ref,

build_tools/github_actions/configure_rocm_python_test_matrix.py

@@ -111,12 +111,18 @@ def build_rocm_python_test_matrix(
         # ]
         for environment in test_environments:
             test_runs_on = _select_test_runs_on(family_info)
+            # These row keys are the contract with workflow files which use them
+            # via matrix.<key> expressions. Empty values are allowed when the
+            # workflow handles them explicitly, but undefined keys are not.
             matrix.append(
                 {
                     "amdgpu_family": amdgpu_family,
                     "test_runs_on": test_runs_on,
                     "python_version": environment.python_version,
                     "container_image_name": environment.container_image_name,
+                    # container_image_url is intentionally empty for native
+                    # Windows test environments and non-empty for Linux
+                    # container test environments.
                     "container_image_url": environment.container_image_url,
                 }
             )

build_tools/github_actions/tests/configure_jax_release_matrix_test.py

@@ -9,6 +9,12 @@
 sys.path.insert(0, os.fspath(Path(__file__).parent.parent))
 
 import configure_jax_release_matrix as m
+from workflow_utils import (
+    WORKFLOWS_DIR,
+    get_matrix_references,
+    get_workflow_job,
+    load_workflow,
+)
 
 
 class ConfigureJaxReleaseMatrixTest(unittest.TestCase):
@@ -35,6 +41,39 @@ def test_explicit_python_version_narrows_matrix(self):
             {"3.12"},
         )
 
+    def test_generated_rows_cover_workflow_matrix_inputs(self):
+        # workflow file like:
+        #
+        #   matrix:
+        #     include: ${{ fromJSON(needs.setup_matrix.outputs.jax_matrix) }}
+        #
+        # Then it passes values to the build workflow using expressions like:
+        #   with:
+        #     test_amdgpu_family: ${{ inputs.test_amdgpu_family }}
+        #     python_version: ${{ matrix.python_version }}
+        #     jax_ref: ${{ matrix.jax_ref }}
+        #
+        # This test checks that all `${{ matrix. }}` values are produced for
+        # every row in the generated matrix. It intentionally does not check
+        # that every generated key is consumed by each workflow; if we want to
+        # enforce exact schemas, do that with generator-local tests.
+
+        workflow = load_workflow(
+            WORKFLOWS_DIR / "multi_arch_release_linux_jax_wheels.yml"
+        )
+        job = get_workflow_job(workflow, "build_jax_wheels")
+        matrix_references = get_matrix_references(job["with"])
+
+        matrix = m.generate_jax_matrix(["3.12"])
+
+        self.assertGreater(len(matrix), 0)
+        for row in matrix:
+            # This checks the row schema, not whether values are truthy. Empty
+            # values are allowed, such as gfx_arch="" for native JAX builds.
+            # Undefined values are not: if the workflow reads `matrix.unknown`,
+            # this test fails until the generator emits that key for every row.
+            self.assertEqual(matrix_references - set(row), set())
+
 
 if __name__ == "__main__":
     unittest.main()

build_tools/github_actions/tests/configure_pytorch_release_matrix_test.py

@@ -9,6 +9,12 @@
 sys.path.insert(0, os.fspath(Path(__file__).parent.parent))
 
 import configure_pytorch_release_matrix as m
+from workflow_utils import (
+    WORKFLOWS_DIR,
+    get_matrix_references,
+    get_workflow_job,
+    load_workflow,
+)
 
 
 class ConfigurePytorchReleaseMatrixTest(unittest.TestCase):
@@ -119,6 +125,63 @@ def test_unknown_explicit_ref_keeps_families(self):
             ],
         )
 
+    def test_generated_rows_cover_workflow_matrix_inputs(self):
+        # The generate_pytorch_matrix_for_release_type script produces matrix
+        # JSON for use in workflow files like:
+        #
+        #   matrix:
+        #     include: ${{ fromJSON(needs.setup_matrix.outputs.pytorch_matrix) }}
+        #
+        # Then it passes values to the build workflow using expressions like:
+        #   with:
+        #     amdgpu_families: ${{ matrix.amdgpu_families }}
+        #     python_version: ${{ matrix.python_version }}
+        #     pytorch_git_ref: ${{ matrix.pytorch_git_ref }}
+        #
+        # This test checks that all `${{ matrix. }}` values are produced for
+        # every row in the generated matrix. It intentionally does not check
+        # that every generated key is consumed by each workflow; if we want to
+        # enforce exact schemas, do that with generator-local tests.
+        test_cases = [
+            (
+                "multi_arch_release_linux_pytorch_wheels.yml",
+                "build_pytorch_wheels",
+                "linux",
+                "dev",
+            ),
+            (
+                "multi_arch_release_windows_pytorch_wheels.yml",
+                "build_pytorch_wheels",
+                "windows",
+                "dev",
+            ),
+            ("multi_arch_ci_linux.yml", "build_pytorch_wheel_fat", "linux", "ci"),
+            ("multi_arch_ci_windows.yml", "build_pytorch_wheel_fat", "windows", "ci"),
+        ]
+
+        for workflow_filename, job_name, platform, release_type in test_cases:
+            with self.subTest(workflow_filename=workflow_filename):
+                workflow = load_workflow(WORKFLOWS_DIR / workflow_filename)
+                job = get_workflow_job(workflow, job_name)
+                matrix_references = get_matrix_references(job["with"])
+
+                matrix = m.generate_pytorch_matrix_for_release_type(
+                    release_type=release_type,
+                    python_versions=["3.12"],
+                    pytorch_git_refs=["release/2.10"],
+                    amdgpu_families="gfx94X-dcgpu",
+                    platform=platform,
+                )
+
+                self.assertGreater(len(matrix), 0)
+                for row in matrix:
+                    # This checks the row schema, not whether values are
+                    # truthy. Empty values are allowed when a workflow handles
+                    # them explicitly. Undefined values are not: if the
+                    # workflow reads `matrix.unknown`, this test fails until
+                    # the generator emits that key for every row.
+                    self.assertEqual(matrix_references - set(row), set())
+
 
 if __name__ == "__main__":
     unittest.main()

build_tools/github_actions/tests/configure_rocm_python_test_matrix_test.py

@@ -10,6 +10,12 @@
 sys.path.insert(0, os.fspath(Path(__file__).parent.parent))
 
 import configure_rocm_python_test_matrix as m
+from workflow_utils import (
+    WORKFLOWS_DIR,
+    get_matrix_references,
+    get_workflow_job,
+    load_workflow,
+)
 
 
 class ConfigureRocmPythonTestMatrixTest(unittest.TestCase):
@@ -103,6 +109,54 @@ def test_unknown_platform_errors(self):
                 platform="not-a-platform",
             )
 
+    def test_generated_rows_cover_workflow_matrix_inputs(self):
+        # The build_rocm_python_test_matrix script produces matrix JSON for use
+        # in workflow files like:
+        #
+        #   matrix:
+        #     include: ${{ fromJSON(inputs.build_config).test_python_packages_matrix }}
+        #
+        # Then it passes values to the test workflow using expressions like:
+        #   with:
+        #     amdgpu_family: ${{ matrix.amdgpu_family }}
+        #     test_runs_on: ${{ matrix.test_runs_on }}
+        #     python_version: ${{ matrix.python_version }}
+        #
+        # This test checks that all `${{ matrix. }}` values are produced for
+        # every row in the generated matrix. It intentionally does not check
+        # that every generated key is consumed by each workflow; if we want to
+        # enforce exact schemas, do that with generator-local tests.
+        test_cases = [
+            ("multi_arch_ci_linux.yml", "linux"),
+            ("multi_arch_ci_windows.yml", "windows"),
+        ]
+
+        for workflow_filename, platform in test_cases:
+            with self.subTest(workflow_filename=workflow_filename):
+                workflow = load_workflow(WORKFLOWS_DIR / workflow_filename)
+                job = get_workflow_job(workflow, "test_python_packages_per_family")
+                matrix_references = get_matrix_references(job["with"])
+
+                matrix = m.build_rocm_python_test_matrix(
+                    per_family_info=[
+                        {
+                            "amdgpu_family": "gfxMOCKTEST",
+                            "test-runs-on": "mock-runner",
+                        }
+                    ],
+                    platform=platform,
+                )
+
+                self.assertGreater(len(matrix), 0)
+                for row in matrix:
+                    # This checks the row schema, not whether values are
+                    # truthy. Empty values are allowed, such as
+                    # container_image_url="" for native Windows tests.
+                    # Undefined values are not: if the workflow reads
+                    # `matrix.unknown`, this test fails until the generator
+                    # emits that key for every row.
+                    self.assertEqual(matrix_references - set(row), set())
+
 
 if __name__ == "__main__":
     unittest.main()

build_tools/github_actions/tests/workflow_utils.py

@@ -3,11 +3,13 @@
 
 """Shared helpers for workflow YAML tests."""
 
+import re
 from pathlib import Path
 
 import yaml
 
 WORKFLOWS_DIR = Path(__file__).resolve().parents[3] / ".github" / "workflows"
+_MATRIX_REFERENCE_RE = re.compile(r"\bmatrix\.([A-Za-z_][A-Za-z0-9_]*)")
 
 
 def load_workflow(path: Path) -> dict:
@@ -16,6 +18,66 @@ def load_workflow(path: Path) -> dict:
         return yaml.safe_load(f)
 
 
+def get_workflow_job(workflow: dict, job_name: str) -> dict:
+    """Returns a workflow job definition.
+
+    For a workflow with:
+        jobs:
+          build_wheels:
+            uses: ./.github/workflows/build_wheels.yml
+            with:
+              python_version: ${{ matrix.python_version }}
+
+    get_workflow_job(workflow, "build_wheels") returns the dictionary
+    containing the uses/with blocks for that job.
+    """
+    jobs = workflow.get("jobs")
+    if not isinstance(jobs, dict):
+        raise KeyError("workflow has no jobs block")
+
+    job = jobs[job_name]
+    if not isinstance(job, dict):
+        raise KeyError(f"workflow job {job_name!r} is not a mapping")
+    return job
+
+
+def get_matrix_references(value: object) -> set[str]:
+    """Extracts top-level matrix keys referenced by a workflow YAML value.
+
+    For a workflow value with:
+        with:
+          python_version: ${{ matrix.python_version }}
+          package_url: >-
+            ${{
+              format('{0}/{1}/index.html',
+                  needs.build.outputs.package_find_links_url,
+                  matrix.amdgpu_family)
+            }}
+
+    get_matrix_references(value) returns:
+        {"python_version", "amdgpu_family"}
+
+    Nested matrix objects like matrix.family_info.amdgpu_family return only the
+    top-level matrix key, {"family_info"}.
+    """
+    if isinstance(value, str):
+        return set(_MATRIX_REFERENCE_RE.findall(value))
+
+    if isinstance(value, dict):
+        references = set()
+        for child_value in value.values():
+            references.update(get_matrix_references(child_value))
+        return references
+
+    if isinstance(value, list):
+        references = set()
+        for child_value in value:
+            references.update(get_matrix_references(child_value))
+        return references
+
+    return set()
+
+
 def _get_workflow_dispatch_block(workflow: dict) -> dict | None:
     """Returns the workflow_dispatch block, or None."""
     # PyYAML parses the unquoted YAML key `on:` as boolean True.