ci: build the stages ci:run-wsl needs on narrowed builds

build_tools/github_actions/amdgpu_family_matrix.py

@@ -918,6 +918,13 @@ def get_cpu_test_runner(platform: str) -> str:
 # Opt-in label for also testing a family's Linux build on its test-runs-on-wsl runner.
 WSL_RUN_LABEL = "ci:run-wsl"
 
+# Artifacts that CI control labels need beyond the components' own artifacts.
+# Tests on WSL also install the WSL libhsakmt (see install_rocm_from_artifacts),
+# so a narrowed build adds the stage that produces it.
+CI_LABEL_REQUIRED_ARTIFACTS: dict[str, list[str]] = {
+    WSL_RUN_LABEL: ["wsl-rocdxg"],
+}
+
 
 def get_wsl_test_runner(
     platform_info: dict, *, labels: list[str], test_runs_on: str, build_variant: str

build_tools/github_actions/configure_multi_arch_ci.py

@@ -58,7 +58,7 @@
 from _therock_utils.build_topology import get_topology
 
 from amdgpu_family_matrix import (
-    WSL_RUN_LABEL,
+    CI_LABEL_REQUIRED_ARTIFACTS,
     all_build_variants,
     get_all_families_for_trigger_types,
     get_cpu_test_runner,
@@ -153,6 +153,27 @@ def _resolve_skipped_stages(build_stages: list[str]) -> list[str]:
     return [stage for stage in all_stages if stage not in build_stages]
 
 
+def _resolve_effective_build_stages(
+    requested_stages: list[str], labels: list[str]
+) -> list[str]:
+    """Add to the requested build stages the stages the requested tests need.
+
+    Labels can require artifacts beyond the components' own (for example a
+    test platform's runtime pieces). The topology resolves those artifacts to
+    their producing stages and the stages those depend on. A full build
+    (``requested_stages`` empty) already includes every stage.
+    """
+    if not requested_stages:
+        return []
+    required_artifacts = sorted(
+        {a for label in labels for a in CI_LABEL_REQUIRED_ARTIFACTS.get(label, [])}
+    )
+    if not required_artifacts:
+        return list(requested_stages)
+    required_stages = get_topology().get_stages_for_artifacts(required_artifacts)
+    return list(requested_stages) + sorted(required_stages - set(requested_stages))
+
+
 # Maps build stages to the test labels that can run with artifacts from those
 # stages. Used to auto-filter tests when build_stages narrows the build graph.
 # Test labels not listed here require a full build.
@@ -203,10 +224,6 @@ def _get_allowed_test_labels_for_stages(build_stages: list[str]) -> list[str] |
     return sorted(allowed) if allowed else []
 
 
-# ci:run-wsl tests also need the wsl-rocdxg stage's artifact (the WSL libhsakmt).
-WSL_REQUIRED_STAGE = "wsl-rocdxg"
-
-
 def _apply_stage_test_labels(
     platform: str,
     labels: list[str],
@@ -278,8 +295,9 @@ class CIInputs:
     prebuilt_stages: str = ""
     baseline_run_id: str = ""
 
-    # Allowlist of build stages to run. Empty means every stage, which is the
-    # default for all existing callers.
+    # Build stages the caller requested. Empty means every stage, which is the
+    # default for all existing callers. decide_jobs() adds the stages that
+    # produce artifacts the requested tests need.
     build_stages: list[str] = field(default_factory=list)
     # Repository to query for baseline runs (for cross-repo artifact reuse)
     baseline_repository: str = ""
@@ -299,9 +317,8 @@ def log(self) -> None:
     def validate(self) -> None:
         """Validate inputs for consistency (fail-fast behavior).
 
-        Raises ValueError for conflicting build selections, test labels requiring
-        stages not in build_stages, or ci:run-wsl on a partial build without the
-        wsl-rocdxg stage.
+        Raises ValueError for conflicting build selections or test labels
+        requiring stages not in build_stages.
         """
         for job in ("native-linux", "python-packages", "pytorch", "jax"):
             build_label = f"ci:build-{job}"
@@ -318,15 +335,6 @@ def validate(self) -> None:
                 _apply_stage_test_labels(
                     platform, labels, self.build_stages, allowed_labels
                 )
-            if (
-                WSL_RUN_LABEL in self.linux_test_labels
-                and WSL_REQUIRED_STAGE not in self.build_stages
-            ):
-                raise ValueError(
-                    f"{WSL_RUN_LABEL} needs the {WSL_REQUIRED_STAGE} stage (the WSL "
-                    f"libhsakmt), which is not in build_stages {self.build_stages}. "
-                    f"Add {WSL_REQUIRED_STAGE} to build_stages."
-                )
 
     @property
     def is_pull_request(self) -> bool:
@@ -444,7 +452,7 @@ def from_environ() -> "CIInputs":
             linux_test_labels=linux_test_labels,
             windows_test_labels=windows_test_labels,
             prebuilt_stages=os.environ.get("PREBUILT_STAGES", ""),
-            build_stages=_parse_comma_list(os.environ.get("BUILD_STAGES", "")),
+            build_stages=build_stages,
             baseline_run_id=os.environ.get("BASELINE_RUN_ID", ""),
             # Which repo to query for baseline_run_id. Defaults to THEROCK_REPOSITORY.
             baseline_repository=os.environ.get("BASELINE_REPOSITORY")
@@ -1413,14 +1421,21 @@ def decide_jobs(
         if auto_stage_reuse.applied_reuse_stages and auto_stage_reuse.baseline_run_id:
             baseline_run_id = auto_stage_reuse.baseline_run_id
 
-    # A build_stages allowlist narrows the build graph: everything outside it
+    # The requested build_stages narrow the build graph to those stages plus
+    # the ones the requested tests need (the effective stages). Everything else
     # is skipped outright, with no artifacts built or copied. This takes
     # precedence over prebuilt/reuse decisions, which only matter for stages
     # that are part of the build in the first place.
-    skipped_stages = _resolve_skipped_stages(ci_inputs.build_stages)
+    effective_stages = _resolve_effective_build_stages(
+        ci_inputs.build_stages, ci_inputs.pr_labels
+    )
+    skipped_stages = _resolve_skipped_stages(effective_stages)
     if skipped_stages:
-        print(f"  build_stages allowlist: {ci_inputs.build_stages}")
-        print(f"  Skipping stages outside the allowlist: {skipped_stages}")
+        print(f"  Requested build stages: {ci_inputs.build_stages}")
+        added_stages = [s for s in effective_stages if s not in ci_inputs.build_stages]
+        if added_stages:
+            print(f"  Added for the requested tests: {added_stages}")
+        print(f"  Skipping stages outside the effective stages: {skipped_stages}")
         for stage in skipped_stages:
             stage_decisions[stage] = JobAction.SKIP
 

build_tools/github_actions/tests/configure_multi_arch_ci_test.py

@@ -218,29 +218,34 @@ def test_build_stages_with_only_control_labels_auto_selects_tests(self):
             ["ci:run-wsl", "test:hip-tests", "test:kfdtest", "test:rocrtst"],
         )
 
-    def test_run_wsl_partial_build_requires_wsl_rocdxg_stage(self):
-        """WSL tests need the wsl-rocdxg artifact (the WSL libhsakmt)."""
-        with self.assertRaisesRegex(ValueError, "wsl-rocdxg"):
-            _run_from_environ(
-                event_name="workflow_dispatch",
-                event_payload={},
-                extra_env={
-                    "LINUX_AMDGPU_FAMILIES": "gfx110X",
-                    "PR_LABELS": "ci:run-wsl,test:hip-tests",
-                    "BUILD_STAGES": "compiler-runtime,runtime-tests",
-                },
-            )
-
-    def test_run_wsl_full_build_needs_no_stage_check(self):
-        inputs = cm.CIInputs(
-            run_id="12345",
-            event_name="pull_request",
-            commit_ref="feature",
-            base_ref="HEAD^",
-            build_variant="release",
-            linux_test_labels=["test:hip-tests", "ci:run-wsl"],
-        )
-        inputs.validate()
+    def test_run_wsl_adds_wsl_rocdxg_to_partial_build(self):
+        """A narrowed build (e.g. an external repo PR) still builds wsl-rocdxg,
+        the WSL libhsakmt that ci:run-wsl tests need, on top of the requested
+        stages, which stay unchanged."""
+        for labels, expect_wsl in (("ci:run-wsl,ci:gfx110X-all", True), ("", False)):
+            with self.subTest(labels=labels):
+                inputs = _run_from_environ(
+                    event_name="pull_request",
+                    event_payload={
+                        "pull_request": {
+                            "labels": [
+                                {"name": name} for name in labels.split(",") if name
+                            ],
+                            "base": {"ref": "develop"},
+                        }
+                    },
+                    extra_env={"BUILD_STAGES": "compiler-runtime,dctools-core"},
+                )
+                self.assertEqual(
+                    inputs.build_stages, ["compiler-runtime", "dctools-core"]
+                )
+                skipped = cm.decide_jobs(
+                    inputs, cm.GitContext(), cm.TargetSelection()
+                ).build_rocm.skipped_stages
+                self.assertEqual("wsl-rocdxg" not in skipped, expect_wsl)
+                self.assertNotIn("compiler-runtime", skipped)
+                self.assertNotIn("dctools-core", skipped)
+                self.assertIn("math-libs", skipped)
 
 
 class TestCIInputsFromEnviron(unittest.TestCase):