ci(jax): take a subset by running the script that holds it

ROCm/jax split the multi-accelerator tests into ci/run_pytest_rocm_multi.sh
rather than adding a variable that picks a subset, so --test-subset now says
which script to run. A ref without that script cannot run the subset on its
own, and saying so beats falling back to the whole suite on a runner that was
taken for the subset.

The environment the two scripts share moved to ci/utilities/rocm_test_env.sh,
which is sourced when it is there. A ref from before the split is still read
from the marked section of its one script, since the release branches take
this one at a time.

Author: gulsumgudukbay

.github/workflows/multi_arch_build_linux_jax_wheels.yml

@@ -23,9 +23,9 @@ on:
         default: ""
       test_scope:
         description: >-
-          "short" runs the single-accelerator tests on the 1-GPU runner,
-          "full" adds the multi-accelerator tests on the multi-GPU runner, and
-          "auto" picks from release_type.
+          "short" runs the suite on the 1-GPU runner, which is the
+          single-accelerator tests, "full" adds the multi-accelerator tests on
+          the multi-GPU runner, and "auto" picks from release_type.
         type: string
         default: "auto"
       python_version:

.github/workflows/test_multi_arch_linux_jax_wheels.yml

@@ -64,8 +64,8 @@ on:
         type: string
       test_subset:
         description: >-
-          Accelerator subset to run: "all", "single" or "multi". "multi" needs
-          test_runs_on to be a multi-GPU runner.
+          Accelerator subset to run: "all", or "multi" for the multi-accelerator
+          tests alone, which needs test_runs_on to be a multi-GPU runner.
         required: false
         type: string
         default: "all"
@@ -131,8 +131,8 @@ on:
         default: "linux-gfx942-1gpu-ccs-csp-ossci-rocm"
       test_subset:
         description: >-
-          Accelerator subset to run: "all", "single" or "multi". "multi" needs
-          test_runs_on to be a multi-GPU runner.
+          Accelerator subset to run: "all", or "multi" for the multi-accelerator
+          tests alone, which needs test_runs_on to be a multi-GPU runner.
         type: string
         default: "all"
 permissions:

build_tools/github_actions/configure_jax_test_matrix.py

@@ -6,12 +6,13 @@
 
 The suite splits in two by the "multiaccelerator" pytest marker. Only that
 subset needs more than one GPU, and it is a small part of the run, so the two go
-to different runners: the single-accelerator tests to the family's 1-GPU runner
-and the multi-accelerator tests to its multi-GPU runner.
+to different runners: the family's 1-GPU runner takes the suite, which on one
+GPU is the single-accelerator tests, and its multi-GPU runner takes the
+multi-accelerator script.
 
 Multi-GPU runners are scarce, so the second job is only worth its queue slot
 when full testing is asked for. Short testing, which is what a pull request
-gets, runs the single-accelerator subset alone.
+gets, runs the suite on the 1-GPU runner alone.
 """
 
 import argparse
@@ -35,9 +36,8 @@
 TEST_SCOPES = [SCOPE_AUTO, SCOPE_SHORT, SCOPE_FULL]
 FULL_SCOPE_RELEASE_TYPES = ["nightly", "prerelease"]
 
-# Mirrors JAXCI_ROCM_TEST_SUBSET in ROCm/jax ci/envs/default.env, which is what
-# run_jax_tests.py passes these to.
-SUBSET_SINGLE = "single"
+# --test-subset of run_jax_tests.py, which is which ROCm/jax suite script runs.
+SUBSET_ALL = "all"
 SUBSET_MULTI = "multi"
 
 
@@ -79,7 +79,7 @@ def build_test_matrix(
 
     single_runner = entry.get("test-runs-on")
     if single_runner:
-        include.append({"test_subset": SUBSET_SINGLE, "test_runs_on": single_runner})
+        include.append({"test_subset": SUBSET_ALL, "test_runs_on": single_runner})
     else:
         print(f"No {platform} test runner for {target}, so no tests will run")
 

build_tools/github_actions/tests/configure_jax_test_matrix_test.py

@@ -58,32 +58,33 @@ def matrix(self, target="gfx94X-dcgpu", platform="linux", scope="short") -> dict
             target=target, platform=platform, scope=scope
         )
 
-    def test_short_scope_runs_only_the_single_accelerator_subset(self):
+    def test_short_scope_runs_only_the_single_gpu_job(self):
+        # On one GPU the suite is the single-accelerator tests.
         matrix = self.matrix(scope="short")
 
-        self.assertEqual(subsets(matrix), ["single"])
-        self.assertEqual(runner_for(matrix, "single"), GFX94X_SINGLE_GPU)
+        self.assertEqual(subsets(matrix), ["all"])
+        self.assertEqual(runner_for(matrix, "all"), GFX94X_SINGLE_GPU)
 
     def test_full_scope_adds_the_multi_accelerator_subset(self):
         matrix = self.matrix(scope="full")
 
-        self.assertEqual(subsets(matrix), ["single", "multi"])
-        self.assertEqual(runner_for(matrix, "single"), GFX94X_SINGLE_GPU)
+        self.assertEqual(subsets(matrix), ["all", "multi"])
+        self.assertEqual(runner_for(matrix, "all"), GFX94X_SINGLE_GPU)
         self.assertEqual(runner_for(matrix, "multi"), GFX94X_MULTI_GPU)
 
     def test_the_outer_family_key_resolves_too(self):
         matrix = self.matrix(target="gfx94x", scope="full")
 
-        self.assertEqual(runner_for(matrix, "single"), GFX94X_SINGLE_GPU)
+        self.assertEqual(runner_for(matrix, "all"), GFX94X_SINGLE_GPU)
         self.assertEqual(runner_for(matrix, "multi"), GFX94X_MULTI_GPU)
 
     def test_a_family_without_a_multi_gpu_runner_skips_that_subset(self):
         # Those tests need several GPUs, and a 1-GPU runner would skip every one
         # of them while reporting a pass.
         matrix = self.matrix(target="gfx1151", platform="windows", scope="full")
 
-        self.assertEqual(subsets(matrix), ["single"])
-        self.assertEqual(runner_for(matrix, "single"), "windows-gfx1151-gpu-rocm")
+        self.assertEqual(subsets(matrix), ["all"])
+        self.assertEqual(runner_for(matrix, "all"), "windows-gfx1151-gpu-rocm")
 
     def test_an_unknown_family_is_an_error(self):
         with self.assertRaises(ValueError):
@@ -117,7 +118,7 @@ def test_the_matrix_a_workflow_consumes(self):
             json.loads(outputs["matrix"]),
             {
                 "include": [
-                    {"test_subset": "single", "test_runs_on": GFX94X_SINGLE_GPU},
+                    {"test_subset": "all", "test_runs_on": GFX94X_SINGLE_GPU},
                     {"test_subset": "multi", "test_runs_on": GFX94X_MULTI_GPU},
                 ]
             },
@@ -128,7 +129,7 @@ def test_a_pull_request_gets_one_job_on_the_single_gpu_runner(self):
 
         self.assertEqual(
             json.loads(outputs["matrix"]),
-            {"include": [{"test_subset": "single", "test_runs_on": GFX94X_SINGLE_GPU}]},
+            {"include": [{"test_subset": "all", "test_runs_on": GFX94X_SINGLE_GPU}]},
         )
 
     def test_an_unknown_release_type_is_rejected(self):