ci(jax): run JAX tests on the multi-GPU runner by gulsumgudukbay · Pull Request #6853 · ROCm/TheRock

gulsumgudukbay · 2026-07-24T21:22:54Z

Motivation

The JAX suite splits in two by the multiaccelerator pytest marker, and only that subset needs more than one GPU: roughly 855 tests against 29k. Sending the whole run to a multi-GPU runner spends a scarce 8-GPU slot on tests that fit on one.
Multi-GPU capacity is the binding constraint, so the question is not only which runner each subset gets but how often the multi-GPU one is taken at all.
This replaces the original version of this PR, which preferred a multi-GPU runner whenever the family had one.

Technical Details

A matrix instead of a runner lookup. configure_target_run.py becomes configure_jax_test_matrix.py, named for what it produces and modelled on configure_pytorch_test_matrix.py. It pairs the single-accelerator tests with the family's test-runs-on and the multi-accelerator ones with its test-runs-on-multi-gpu.
The script decides the cadence too. A prerelease takes the multi-GPU runner every time; a nightly takes it one day a week (Sunday, UTC); CI and dev builds never do. --test-scope short|full overrides either, so a workflow or a person can ask for the pair on any day. The rule is unit-tested against fixed dates rather than the clock.
The day comes from the clock, as it does for the PyTorch pipeline. A full re-run onto another day therefore gives that day's matrix; the configure job annotates that case and names the test_scope: full dispatch that runs the multi-accelerator tests anyway.
What that costs the pool. The nightly JAX matrix is 15 build cells (4 refs × 4 Python versions, less 3.11 on 0.11.0), so weekly rather than nightly is about 15 multi-GPU jobs a week instead of 105. A pull request never asks for one.
A subset is a script, not a flag. run_jax_tests.py --test-subset all|multi runs ci/run_pytest_rocm.sh or ci/run_pytest_rocm_multi.sh (ROCm/jax PRs [ROCm] Split the multi-accelerator tests into their own script jax#832, [ROCm] Split the multi-accelerator tests into their own script jax#833, [ROCm] Split the multi-accelerator tests into their own script jax#834, [ROCm] Split the multi-accelerator tests into their own script jax#835). Asking for multi against a checkout that has no such script is an error, on a dry run too, rather than a job that quietly runs everything on an 8-GPU machine.
The environment comes from the checkout. The runner reads it from ci/utilities/rocm_test_env.sh, falling back to the marked section of the older single script, so both layouts work while the release branches move.
A family with no multi-GPU runner drops that job rather than sending it to a 1-GPU machine, where every test in it would skip and the job would still pass.

How a run reaches its jobs

A scope decides which runners a run takes, not how many tests each one runs: short and full both run the whole single-accelerator suite.

flowchart TD
    A["release_type + test_scope"] --> B{"test_scope"}
    B -->|short| S["scope = short"]
    B -->|full| F["scope = full"]
    B -->|auto| C{"release_type"}
    C -->|prerelease| F
    C -->|nightly| D{"Sunday, UTC?"}
    C -->|"ci, dev"| S
    D -->|yes| F
    D -->|no| S
    S --> J1["test_subset=all<br/>test-runs-on, 1 GPU<br/>ci/run_pytest_rocm.sh"]
    F --> J1
    F --> J2["test_subset=multi<br/>test-runs-on-multi-gpu<br/>ci/run_pytest_rocm_multi.sh"]

A family with no test-runs-on-multi-gpu drops the second job, and one with neither label runs nothing; both are ::warning:: annotations on the configure job rather than silent.

Test Plan

pre-commit, including actionlint on the workflows.
Unit tests across build_tools/github_actions/tests.
Run the matrix script for each release type, on and off the weekly day.
Dry-run each subset against a real ROCm/jax checkout.

Test Result

pre-commit clean. 739 unit tests pass; the four failures (stage_reuse_decision_test x2, configure_multi_arch_ci_test, post_stage_upload_test) fail identically on main.
The matrix holds: ci gives one job on linux-gfx942-1gpu-ccs-csp-ossci-rocm; nightly gives the same on a Tuesday and adds linux-gfx942-8gpu-ossci-rocm on a Sunday; prerelease gives both whatever day it is;
--test-subset all execs ci/run_pytest_rocm.sh and --test-subset multi execs ci/run_pytest_rocm_multi.sh.
GPU test execution on a real CI runner: TODO before merge.

#6753 builds on this to add test sizes, so the single-GPU job can run a PR-sized selection.

Submission Checklist

Look over the contributing guidelines at https://github.com/ROCm/TheRock/blob/main/CONTRIBUTING.md.

therock-pr-bot · 2026-07-24T21:24:11Z

✅ All Checks Passed — Ready for Review

Check	Status	Details
📝 PR Description	✅ Pass	—
⛔ Forbidden Files	✅ Pass	—
🧪 Unit Test	✅ Pass	—
🔎 pre-commit	✅ Pass	—
🚫 Draft PR	🔜 To Be Enabled	—
🚩 Feature Flag	🔜 To Be Enabled	—
📊 Code Coverage	🔜 To Be Enabled	—
🤖 therock-pr-bot	✅ Pass	—

🎉 All checks passed! This PR is ready for review.

📖 Need help? See the Policy FAQ for details on every check and how to fix failures.

🙋 Wish to Override Policy?

therock-pr-bot · 2026-07-24T21:24:29Z

🎉 All checks passed! This PR is ready for review.

erman-gurses

Would be good if you share the test details - how did you test it and what is the output.

gulsumgudukbay · 2026-07-27T15:35:00Z

Would be good if you share the test details - how did you test it and what is the output.

please see the PR description @erman-gurses

ScottTodd · 2026-07-27T15:58:52Z

Requested review from @geomin12 and @amd-shiraz for capacity planning. 8 GPU runners are generally more expensive and in shorter supply than 1 GPU runners.

amd-shiraz · 2026-07-27T20:43:48Z

+        # JAX tests use the multi-GPU runner for the family.
        runner_label = configure_target_run.get_runner_label("gfx94X-dcgpu", "linux")
-        self.assertEqual(runner_label, "linux-gfx942-1gpu-ccs-csp-ossci-rocm")
+        self.assertEqual(runner_label, "linux-gfx942-8gpu-ossci-rocm")


@gulsumgudukbay how many builds is JAX tests going to add on a daily basis ? including nightlies, PRs etc. this is a multi GPU runner and we only have 10 nodes and the capacity on these are already maxed out.

For PR CI 1 JAX version (0.11.0) will be present and for nightly we have 4 versions (0.10.0, 0.10.1, 0.10.2 and 0.11.0).

ScottTodd · 2026-07-27T22:09:56Z

More context... for PyTorch we run standard tests on 1-GPU runners then "full" tests on a mix of 1-GPU and multi-GPU runners at lower frequency. See

https://github.com/ROCm/TheRock/blob/main/.github/workflows/multi_arch_build_portable_linux_pytorch_wheels.yml
https://github.com/ROCm/TheRock/blob/main/.github/workflows/test_pytorch_wheels_full.yml
https://github.com/ROCm/TheRock/blob/main/external-builds/pytorch/generate_test_sharding_matrix.py (note that only "distributed" tests run on the multi-GPU runners, other tests run on 1-GPU runners in separate shards)

We should probably do something similar for JAX. Would need to see the data for how much runner time is needed for each subset of the tests.

gulsumgudukbay · 2026-07-27T23:38:44Z

Working with @amd-shiraz offline for capacity numbers.
cc: @ScottTodd

gulsumgudukbay · 2026-07-27T23:43:09Z

Also @ScottTodd Agreed on following the PyTorch model and the JAX equivalent of that split already exists upstream, in ci/run_pytest_rocm.sh in ROCm/jax (the script #6558 switches us to). It splits the suite by the multiaccelerator pytest marker into two separate invocations:

# single-accelerator, parallel via xdist
pytest -n $num_processes -m "not multiaccelerator" \
  --json-report --json-report-file=logs/pytest_results_single.json \
  --junitxml=test-artifacts/junit-single.xml tests
# multi-accelerator, only when more than one GPU is present
if [[ $gpu_count -gt 1 ]]; then
  pytest -m "multiaccelerator" \
    --json-report --json-report-file=logs/pytest_results_multi.json \
    --junitxml=test-artifacts/junit-multi.xml tests
fi

So the "distributed vs everything else" classification is already done for us by the marker, and per-test durations are already written to the JSON/JUnit reports on every run, they're just not uploaded as artifacts today.

Two things worth noting for the capacity discussion (cc @amd-shiraz):

On a 1-GPU runner the multi-accelerator subset is skipped entirely (Skipping multi-accelerator tests (only 1 GPU detected)). So the multi-GPU runner isn't about speed, without it we get zero multi-device coverage.
num_processes is capped at 16 regardless of GPU count, but JAX_ENABLE_ROCM_XDIST=$gpu_count spreads those workers across the available GPUs. On 1 GPU that's 16 workers contending on one device; on 8 GPUs it's ~2 per device. So the single-accelerator subset may also be affected, but that's measurable rather than something I want to assert.
Plan to get you the data, without merging this PR: since test_runs_on is already a workflow_dispatch input on test_linux_jax_wheels_partial.yml, I'll dispatch Charlie's #6558 branch twice, once on linux-gfx942-1gpu-ccs-csp-ossci-rocm and once on linux-gfx942-8gpu-ossci-rocm, and add an artifact upload for pytest_results_{single,multi}.json so we get per-test timings. That gives us runner time for each subset on each runner type, and I'll post the numbers here.

On this PR specifically: you and @amd-shiraz are right that it's too blunt as written. It was a fast PR to get other PRs unblocked and routes the whole JAX test job to the 8-GPU runner, including tests that only need 1 GPU. I'll rework it to match the PyTorch shape (default 1-GPU runner plus a separate multi-GPU runner input, with only the multiaccelerator subset routed there) so our 8-GPU footprint is limited to what genuinely requires it.

amd-shiraz · 2026-07-27T23:51:21Z

thanks @gulsumgudukbay sharing some capacity backlogs we are seeing already hence we are being caution before taking more.

gulsumgudukbay · 2026-07-28T01:09:09Z

Also @ScottTodd Agreed on following the PyTorch model and the JAX equivalent of that split already exists upstream, in ci/run_pytest_rocm.sh in ROCm/jax (the script #6558 switches us to). It splits the suite by the multiaccelerator pytest marker into two separate invocations:
# single-accelerator, parallel via xdist
pytest -n $num_processes -m "not multiaccelerator" \
  --json-report --json-report-file=logs/pytest_results_single.json \
  --junitxml=test-artifacts/junit-single.xml tests
# multi-accelerator, only when more than one GPU is present
if [[ $gpu_count -gt 1 ]]; then
  pytest -m "multiaccelerator" \
    --json-report --json-report-file=logs/pytest_results_multi.json \
    --junitxml=test-artifacts/junit-multi.xml tests
fi
So the "distributed vs everything else" classification is already done for us by the marker, and per-test durations are already written to the JSON/JUnit reports on every run, they're just not uploaded as artifacts today.

Two things worth noting for the capacity discussion (cc @amd-shiraz):

On a 1-GPU runner the multi-accelerator subset is skipped entirely (Skipping multi-accelerator tests (only 1 GPU detected)). So the multi-GPU runner isn't about speed, without it we get zero multi-device coverage. num_processes is capped at 16 regardless of GPU count, but JAX_ENABLE_ROCM_XDIST=$gpu_count spreads those workers across the available GPUs. On 1 GPU that's 16 workers contending on one device; on 8 GPUs it's ~2 per device. So the single-accelerator subset may also be affected, but that's measurable rather than something I want to assert. Plan to get you the data, without merging this PR: since test_runs_on is already a workflow_dispatch input on test_linux_jax_wheels_partial.yml, I'll dispatch Charlie's #6558 branch twice, once on linux-gfx942-1gpu-ccs-csp-ossci-rocm and once on linux-gfx942-8gpu-ossci-rocm, and add an artifact upload for pytest_results_{single,multi}.json so we get per-test timings. That gives us runner time for each subset on each runner type, and I'll post the numbers here.

On this PR specifically: you and @amd-shiraz are right that it's too blunt as written. It was a fast PR to get other PRs unblocked and routes the whole JAX test job to the 8-GPU runner, including tests that only need 1 GPU. I'll rework it to match the PyTorch shape (default 1-GPU runner plus a separate multi-GPU runner input, with only the multiaccelerator subset routed there) so our 8-GPU footprint is limited to what genuinely requires it.

@amd-shiraz I started two jobs (1 and 8-GPU runner)

1-GPU: https://github.com/ROCm/TheRock/actions/runs/30325982453
8-GPU: https://github.com/ROCm/TheRock/actions/runs/30325989065

These runs are ran using the following commands:

# ---- 1-GPU baseline (single-accelerator subset only; multi is auto-skipped) ----
gh workflow run test_linux_jax_wheels_partial.yml --ref cph-expand-jax-tests \
  -f test_amdgpu_family=gfx94X-dcgpu \
  -f jax_ref=rocm-jaxlib-v0.11.0 \
  -f python_version=3.13 \
  -f rocm_version=7.15.0a20260723 \
  -f package_index_url=https://rocm.devreleases.amd.com/whl-multi-arch/ \
  -f rocm_package_find_links_url=https://therock-nightly-artifacts.s3.amazonaws.com/29968645587-linux/python/index.html \
  -f jax_version=0.11.0 \
  -f jax_plugin_version=0.11.0+rocm7.15.0a20260723 \
  -f jax_pjrt_version=0.11.0+rocm7.15.0a20260723 \
  -f test_runs_on=linux-gfx942-1gpu-ccs-csp-ossci-rocm \
  -f visible_devices=0
# ---- 8-GPU (gets BOTH subsets) ----
gh workflow run test_linux_jax_wheels_partial.yml --ref cph-expand-jax-tests \
  -f test_amdgpu_family=gfx94X-dcgpu \
  -f jax_ref=rocm-jaxlib-v0.11.0 \
  -f python_version=3.13 \
  -f rocm_version=7.15.0a20260723 \
  -f package_index_url=https://rocm.devreleases.amd.com/whl-multi-arch/ \
  -f rocm_package_find_links_url=https://therock-nightly-artifacts.s3.amazonaws.com/29968645587-linux/python/index.html \
  -f jax_version=0.11.0 \
  -f jax_plugin_version=0.11.0+rocm7.15.0a20260723 \
  -f jax_pjrt_version=0.11.0+rocm7.15.0a20260723 \
  -f test_runs_on=linux-gfx942-8gpu-ossci-rocm \
  -f visible_devices=8

on branch users/gulsumgudukbay/jax-timing-probe

geomin12

as noted by Shiraz, let's calculate velocity so we don't run into queue issues with this machine. please work with @amd-shiraz about this

gulsumgudukbay · 2026-07-28T15:03:47Z

Also posting my results here @geomin12 @amd-shiraz @erman-gurses
JAX full-suite timings, one runner type vs the other. Same inputs, same commit; the only difference is the runner and an explicit device mask on the 1-GPU side so it actually gets one GPU.

	1 GPU: run 1	1 GPU: run 2	8 GPUs: run 1	8 GPUs: run 2
Single-accelerator subset	46:45	44:57	14:49	14:33
Multi-accelerator subset	not run	not run	3:13	3:07
Queue wait	8 min	1 min	89 min	129 min
Setup before tests	2 min	6 min	32 min	13 min
End to end	58 min	52 min	141 min	161 min

Details worth noting:

The multi-accelerator subset is 855 tests. Everything else 64,698 collected, 29,217 executed, is single-accelerator.
- Eight GPUs give a 3.2x speedup on the single-accelerator subset (46:45 to 14:49).
The single-accelerator subset produced identical results on both (4 failed, 29,217 passed, 35,477 skipped), so only speed differs.

Runs: 1-GPU, 8-GPU.

The entire JAX suite finishes in about 18 minutes of 8-GPU node time, 14:49 for the single-accelerator tests plus 3:13 for the multi-accelerator ones. So this is a bounded, sub-half-hour nightly ask rather than a standing allocation.

## Summary JAX test jobs were discarding the GPU isolation supplied by the runner, letting a job on a 1-GPU runner use every GPU on the host. This removes the two `unset` calls responsible. ## Problem `test_linux_jax_wheels_partial.yml` unset `HIP_VISIBLE_DEVICES` and `ROCR_VISIBLE_DEVICES` before running tests. Those are exactly the variables the runner provides through `--env-file /etc/podinfo/gha-gpu-isolation-settings`, and since the container is given `/dev/kfd` and `/dev/dri` outright, that env file is the only isolation mechanism in play. So a job scheduled on `linux-gfx942-1gpu-ccs-csp-ossci-rocm` enumerates all 8 GPUs and allocates across them. ## Impact In [run 30032481350](https://github.com/ROCm/TheRock/actions/runs/30032481350) all 15 test jobs saw 8 devices, and three failed with `RESOURCE_EXHAUSTED` on GPUs they were never assigned, all in `multi_device_test.py::MultiDeviceTest::test_lax_full_like_efficient`: | Job | Allocation | Allocator | |---|---|---| | [py3.14 / v0.10.2](https://github.com/ROCm/TheRock/actions/runs/30032481350/job/89307813608) | 12 GiB | `GPU_5_bfc` | | [py3.11 / v0.10.2](https://github.com/ROCm/TheRock/actions/runs/30032481350/job/89308971448) | 12 GiB | `GPU_7_bfc` | | [py3.13 / v0.10.1](https://github.com/ROCm/TheRock/actions/runs/30032481350/job/89308509210) | 16 GiB | `XLA_backend_5_bfc` | A 12 GiB allocation is trivial for a 192 GiB MI300X, so those GPUs were occupied by other tenants. This produces flaky failures that look like JAX or ROCm bugs, and our jobs can equally OOM other workloads on the same host. ## Why this fix Five workflows mount the isolation env file `test_pytorch_wheels.yml`, `test_pytorch_wheels_full.yml`, `test_rocm_wheels.yml`, `test_linux_jax_wheels.yml`, and this one. Only this one unsets the masks; removing them restores parity. The remaining unsets (`ROCM_ROOT`, `HIP_DEVICE_LIB_PATH`, `JAX_ROCM_PLUGIN_INTERNAL_*`) are unrelated to GPU visibility, they exist because the image is deliberately `/opt/rocm`-less and are left in place. No test logic, runner selection, or matrix configuration is touched. ## Note on multi-device coverage With isolation respected, jobs on 1-GPU runners see one GPU, so `run_pytest_rocm.sh` will skip the `multiaccelerator` subset (it gates on `gpu_count > 1`). Those tests need a properly labeled multi-GPU runner, tracked separately in #6853. Measured cost of that subset is 3:13 for 855 tests. ## Tests **Mechanism confirmed.** To show the masks were the cause rather than a correlation, I re-applied them explicitly after the unsets on a 1-GPU runner ([run 30325982453](https://github.com/ROCm/TheRock/actions/runs/30325982453)): ``` Number of GPUs detected: 1 [RocmDevice(id=0)] ``` `rocminfo` and JAX both reported one device, versus 8 in every unmasked run on the same runner label. **Verification of this branch.** [Run 30375175749](https://github.com/ROCm/TheRock/actions/runs/30375175749) on `linux-gfx942-1gpu-ccs-csp-ossci-rocm`, covering `multi_device_test.py` (the test that OOM'd above), `core_test.py`, `util_test.py` and `scipy_stats_test.py`. | Check | Expected | Result | |---|---|---| | Devices visible to JAX | `[RocmDevice(id=0)]` | `[RocmDevice(id=0)]` - 1 device | | Cross-tenant OOM | none | none | | Overall conclusion | success | success | Per-file results: | Test file | Result | |---|---| | `multi_device_test.py` | 2 passed, 18 skipped (3.87s) | | `core_test.py` | 480 passed (40.04s) | | `util_test.py` | 19 passed, 6 subtests passed (0.87s) | | `scipy_stats_test.py` | 982 passed (282.69s) | The `multi_device_test.py` numbers are the expected consequence of restoring isolation. On the same file **before** the fix, [job 89307813608](https://github.com/ROCm/TheRock/actions/runs/30032481350/job/89307813608) ran 19 passed / 1 failed, the failure being the cross-tenant OOM because the job could see 8 GPUs. With one GPU visible, 18 of those 20 tests correctly skip as requiring multiple devices, and `test_lax_full_like_efficient` no longer OOMs. That skipped coverage is real and is what #6853 exists to restore, on a runner where multiple GPUs are actually allocated to us rather than borrowed from neighbours. ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/TheRock/blob/main/CONTRIBUTING.md.

ScottTodd

Okay, it sounds like we don't have enough capacity to always require multi-GPU for all JAX tests, and there are related PRs that are moving JAX tests up from nightly builds to also CI builds.

I think we should have a script like https://github.com/ROCm/TheRock/blob/main/build_tools/github_actions/configure_pytorch_test_matrix.py for JAX that decides which jobs to run with which runners, choosing 1 GPU runners for "not multiaccelerator" tests when short testing is requested (e.g. on a pull request) and choosing multi-GPU runners for "multiaccelerator" tests when full testing is requested (e.g. once a week for the "nightly" trigger" or more often, always for the "prerelease" release type, etc.).

I don't think we should always use multi-GPU runners whenever they are available, that's going to put heavy pressure on runner queues without giving sufficient control over when such jobs are run.

ScottTodd · 2026-08-03T18:12:59Z

 # SPDX-License-Identifier: MIT

-"""This file looks up the test-runs-on runner label for a given GPU family.
+"""This file looks up the runner label to run JAX tests on for a GPU family.


If this script is only used for JAX, it should be renamed. See what we have for PyTorch:

configure_pytorch_release_matrix.py

configure_pytorch_test_matrix.py

and JAX:

configure_jax_release_matrix.py

this file --> configure_jax_test_matrix.py (ideally, following the same structure) or configure_jax_target_run.py (minimal change without refactoring the workflows)

Both done. configure_target_run.py is now configure_jax_test_matrix.py, modelled on configure_pytorch_test_matrix.py: it pairs the not multiaccelerator tests with the family's 1-GPU runner and the multiaccelerator ones with its multi-GPU runner, and it decides how often (prerelease every time, nightly one day a week), --test-scope to override.
A family with no multi-GPU runner drops that job instead of sending it to a 1-GPU machine, where every test in it would skip and still report a pass.

gulsumgudukbay · 2026-08-11T22:26:00Z

@amd-shiraz the nightly JAX matrix is 15 cells (4 refs x 4 Python versions, minus 3.11 on 0.11.0);
PR CI is 1.
As written before, every nightly cell would have taken an 8-GPU runner: ~105 jobs a week. now with the recent change, only the multi-accelerator subset goes there, and only on Sundays for nightly plus prerelease runs: ~15 a week plus the prerelease. PR CI never touches that pool and it runs on the 1-GPU runner.
Could we add the multi GPU runner to JAX CI with these changes?

cc: @ScottTodd, @jayfurmanek

ScottTodd

I don't expect I'll have time to serve as primary reviewer on these PRs. Will defer to Erman who can tag me in as needed.

gulsumgudukbay · 2026-08-12T04:12:07Z

https://github.com/ROCm/TheRock/actions/runs/31568065672 is the test for this PR

erman-gurses

Please see my comments for initial review:

In general, I think it worth to add small flow/design diagram to the PR description that shows how release_type / test_scope map to the generated matrix, runners, and test_subset. I believe this will make much easier to understand and review.

This PR also needs to be reviewed by @amd-shiraz and/or @geomin12 for the final review.

erman-gurses · 2026-08-12T21:04:38Z

+SCOPE_SHORT = "short"
+SCOPE_FULL = "full"
+SCOPE_AUTO = "auto"


We can add TEST prefix in here to give more information for the future as well as specific comment for like what do SHORT / FULL / AUTO mean?

Done: they are TEST_SCOPE_SHORT / TEST_SCOPE_FULL / TEST_SCOPE_AUTO now, commented with what each selects. "short" runs the whole single-accelerator suite on the 1-GPU runner, "full" adds the multi-accelerator tests on the multi-GPU runner, and "auto" reads the release type. The comment also says the part that was easy to misread: a scope decides which runners a run takes, not which tests each one runs.

erman-gurses · 2026-08-12T21:08:42Z

+SUBSET_ALL = "all"
+SUBSET_MULTI = "multi"


Same here to add meaningful prefix. What do all vs multi mean?

Done: TEST_SUBSET_ALL / TEST_SUBSET_MULTI, commented with the script behind each. "all" is ci/run_pytest_rocm.sh, which is the whole suite if there are multiple GPUs on a host, but is only the single-GPU test suite if there is only 1 GPU on the host, and "multi" is ci/run_pytest_rocm_multi.sh, the multi-accelerator tests alone.
run_jax_tests.py uses the same two names, so the workflow input, the matrix and the runner all read alike.

erman-gurses · 2026-08-12T21:20:22Z

+        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")
+


What if we missed single_runner in here? Both short and full appear to require the single-GPU all job, but currently short would disable testing entirely and full could run only the multi subset if a multi-GPU runner exists. This could make a missing/changed runner configuration silently reduce the test coverage. Shouldn't we raise an error in here instead of printing?

An empty test-runs-on is a real configuration rather than a mistake. gfx90a on Windows has a build and no test hardware. So raising would fail the build job for every such family.
It is a ::warning:: annotation now instead of a bare print, so it lands in the run summary, and a new test covers the case: the matrix comes back empty and annotated rather than quietly multi-only.

erman-gurses · 2026-08-12T21:27:21Z

@@ -122,13 +129,19 @@ on:
        description: Runner label to use. The selected runner should have a GPU supported by test_amdgpu_family
        type: string
        default: "linux-gfx942-1gpu-ccs-csp-ossci-rocm"


Should we guard against test_subset=multi being used with a single-GPU test_runs_on? Since inputs can be changed independently in workflow_dispatch, so selecting multi while leaving the default 1-GPU runner might result in the multi-accelerator tests being skipped while the job still passes.

It is enforced already in ci/run_pytest_rocm_multi.sh exits non-zero when gpu_count <= 1 ("Multi-accelerator tests need more than one GPU"), so dispatching multi at a 1-GPU runner fails the job instead of skipping every test and passing.
On this side the input is now a choice of all / multi and its description says it fails unless test_runs_on is a multi-GPU runner.

erman-gurses · 2026-08-12T21:32:25Z

+def today_utc() -> date:
+    """The day the run is happening, which the weekly rule below reads."""
+    return datetime.now(timezone.utc).date()


From agent:

The PR defines Sunday UTC as the weekly multi-GPU day. But today is evaluated when configure_jax_test_matrix.py actually runs. That means the same nightly workflow can produce a different matrix when rerun on a different day: Original Sunday nightly: auto + nightly + Sunday → full → all + multi Rerun Monday: auto + nightly + Monday → short → all only The unit tests verify Sunday versus Tuesday, but they don't address reruns preserving the original scheduled cadence.

Fixed in 854ae9a: the configure job reads the run's created_at from the API and passes it as --run-date, so the weekly rule reads the day the run was created instead of the day it happens to be evaluated. A re-run keeps created_at at the first attempt while run_started_at moves to the new one, so re-running a Sunday nightly on Monday now repeats the matrix it had, which is the case that matters when a job failed due to for example infrastructure issues. It needs actions: read, granted in this workflow and in the calling release workflow, and falls back to today if the token cannot read the run; a unit test covers a nightly with --run-date on a Sunday producing both jobs while the clock says Tuesday.

This seems fairly complicated, see my other comments (#6753 (review)). Simpler would be to match pytorch:

TheRock/.github/workflows/multi_arch_build_portable_linux_pytorch_wheels.yml

Lines 561 to 570 in 35ed46c

steps:

- name: Check full test cadence

id: cadence

# Release branches test daily; the nightly branch tests weekly (Sunday).

run: |

if [ "${PYTORCH_GIT_REF}" != "nightly" ] || [ "$(date -u +%u)" = "7" ]; then

echo "dispatch=true" >> "$GITHUB_OUTPUT"

else

echo "dispatch=false" >> "$GITHUB_OUTPUT"

fi

erman-gurses · 2026-08-12T21:41:06Z

+      test_scope:
+        description: >-
+          "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"


Minor naming question - short still runs the full single-GPU suite and only skips the multi-GPU subset, right? Would something like single / multi (or similar) make the distinction clearer? I initially read short as a reduced test selection.

yes, short runs the whole single-accelerator suite and only drops the multi-GPU job, so the input description and the constants now say that a scope picks runners rather than narrowing the suite. The stacked #6753 removes the ambiguity entirely by replacing test_scope with test_size (small / medium / large), where only small is a reduced selection.

gulsumgudukbay · 2026-08-14T02:59:26Z

@erman-gurses thanks! all seven comments are addressed in 4aeddb9 and 854ae9a, and the description now carries a diagram of how release_type and test_scope map to the scope, the runners and test_subset, with the "no runner configured" branches called out. The Test Plan and Test Result sections list what was run: pre-commit with actionlint, the build_tools/github_actions unit tests, the matrix script per release type on and off the weekly day, and a dry run of each subset against a real ROCm/jax checkout; a GPU run on a real CI runner is still outstanding before merge (https://github.com/ROCm/TheRock/actions/runs/31568065672 did pass but I will restart a new instance of it with the updated branch).

Summary of the changes: the constants are TEST_SCOPE_* and TEST_SUBSET_* with comments saying what each selects, a family with no test runner is now a ::warning:: annotation covered by a test rather than a bare print, the weekly day is read from the run's created_at so a rerun repeats the matrix of the attempt it repeats, and test_subset is a choice in workflow_dispatch. cc @amd-shiraz @geomin12 for the capacity side: weekly rather than nightly is about 15 multi-GPU jobs a week instead of 105, and a pull request never asks for one.

gulsumgudukbay · 2026-08-14T04:37:57Z

@erman-gurses test run is below:
Run 31770261531, test_scope=full. Exercises the renamed constants, the configure job's new created_at lookup with actions: read, and the two-job matrix, so both the 1-GPU all job and the 8-GPU multi job should appear.

gulsumgudukbay · 2026-08-14T16:09:38Z

The previous job failed due to setup and network speed timing out the whole job. https://github.com/ROCm/TheRock/actions/runs/31817753726 this is the new job

erman-gurses

LGTM - Thanks addressing all the comments

gulsumgudukbay · 2026-08-17T02:05:37Z

newest build workflow: https://github.com/ROCm/TheRock/actions/runs/31985867454

## What breaks today Every JAX nightly test job fails before it runs a single test: Error: ci/run_pytest_rocm.sh has no section between '# Set up the generic test environment variables' and '# Run tests'. `run_jax_tests.py` evaluates the suite's environment so that the serial retry pass runs under the same settings as the run it is checking. It read that environment out of a marked section of `ci/run_pytest_rocm.sh`. ROCm/jax has since moved it into `ci/utilities/rocm_test_env.sh`, and every release branch from `rocm-jaxlib-v0.10.0` through `rocm-jaxlib-v0.11.0` now ships the file with no section left behind, so the lookup fails on all of them. ## The change Source `ci/utilities/rocm_test_env.sh` when the checkout has it, and fall back to the section when it does not, since an older `--jax-ref` is still valid. The file only exports, echoes and queries devices, so sourcing it runs nothing. Split out of #6853 so the nightly is unblocked without waiting on that stack. Solves #7320 ## Test plan - [x] `pytest build_tools/github_actions/tests/run_jax_tests_test.py` (77 passed) - [x] New coverage for the split layout, and the error now names both places to fix - [x] Nightly JAX test job gets past the environment step

The suite script runs the single-accelerator tests and then the multiaccelerator ones, so a job cannot take just the subset its runner has the GPUs for. Pass JAXCI_ROCM_TEST_SUBSET through as --test-subset. A jax ref from before that variable ignores it and runs everything, which would look like the subset had been covered, so asking for a subset such a ref does not read stops the run.

The suite splits in two by the "multiaccelerator" marker, and only that subset needs more than one GPU. Picking a multi-GPU runner whenever one exists therefore spent a scarce 8-GPU slot to run ~29k tests that fit on one GPU. Replace the runner lookup with a matrix that pairs each subset with the runner it needs: the single-accelerator tests on the family's 1-GPU runner, and, when full testing is asked for, the multi-accelerator tests on its multi-GPU runner. Release types that run per change get the first job alone, so a pull request no longer queues for 8 GPUs, and --test-scope lets a workflow put the full pair on its own cadence. A family with no multi-GPU runner drops that job rather than sending it to a 1-GPU runner, which would skip every test in it and still pass. The script is named for what it produces now, rather than after the generic configure_target_run.py it started from.

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.

The nightly JAX matrix is 15 build cells, so a multi-GPU job per cell is 105 jobs a week against a pool of 10 nodes whose queues are already full. The multi-accelerator subset does not change fast enough to be worth that. Let the matrix script decide how often, not only what: a prerelease takes that runner every time, a nightly takes it one day a week, and --test-scope still overrides both for anything that wants the pair on another cadence.

…mpty Review follow-up: prefix the constants, spell out that a scope picks runners rather than narrowing the suite, annotate a family with no test runner instead of printing past it, and offer the subset as a choice.

A rerun of a Sunday nightly on Monday configured itself for Monday and dropped the multi-accelerator job, so a run that failed on infrastructure could not be repeated as it was. The API keeps created_at at the first attempt while run_started_at moves, so the configure job reads that.

Installing the ROCm packages has taken over an hour on its own, so a 120 minute job budget cut the single-accelerator suite off mid-run. Give the job 240 and hold each subset to its own step budget instead.

gulsumgudukbay · 2026-08-19T02:19:40Z

Rebased onto main (which now carries #7387) and re-validated on the nightly packages the HUD runs against (10.1.0a20260818): https://github.com/ROCm/TheRock/actions/runs/32202408475

Where this stack overlapped #7387, the merged version of the environment lookup wins: suite_env_program returns the file it read alongside the program, and the diagnostics name that file. What is left here is the second suite script and the subset-to-runner selection.

Note that a dispatched run is release_type: dev, so the scope resolves to short and this run takes the 1-GPU runner only, the same as a nightly on any day but Sunday. The multi-GPU leg needs a run with test_scope: full.

gulsumgudukbay · 2026-08-19T02:23:59Z

Full-scope run on the same nightly packages, for the leg a dev/nightly dispatch does not take: https://github.com/ROCm/TheRock/actions/runs/32208281220 (test_scope: full, jax rocm-jaxlib-v0.11.0, py3.12)

Its configure job selected both subsets, all on linux-gfx942-1gpu-ccs-csp-ossci-rocm and multi on linux-gfx942-8gpu-ossci-rocm, which is what this PR is for.

ScottTodd

No major concerns from me on this.

ScottTodd · 2026-08-19T21:22:32Z

+        run: |
+          # created_at stays at the first attempt while run_started_at moves, so
+          # a rerun the next day still gets the day it was scheduled for. Falling
+          # back to today keeps the run going if the token cannot read it.
+          created_at=$(gh api "repos/${GITHUB_REPOSITORY}/actions/runs/${GITHUB_RUN_ID}" --jq .created_at || date -u +%FT%TZ)
+          python ./build_tools/github_actions/configure_jax_test_matrix.py \
+            --target "${{ inputs.test_amdgpu_family }}" \
+            --platform linux \
+            --test-scope "${{ inputs.test_scope }}" \
+            --release-type "${{ inputs.release_type }}" \
+            --run-date "${created_at%%T*}"


Another pattern we could use here would be #1236 - have a single "schedule release" job run on schedule and freeze all dynamic, time-related inputs in that job, then have all other jobs spawned from that one use fixed values. Here we could even have multi_arch_release_linux.yml do the freeze, without needing a new "schedule release" job.

I'd actually prefer to follow what PyTorch does for consistency/simplicity. The retriggering behavior isn't going to come up too often and we could solve for pytorch too:

TheRock/.github/workflows/multi_arch_build_portable_linux_pytorch_wheels.yml

Lines 561 to 570 in 35ed46c

steps:

- name: Check full test cadence

id: cadence

# Release branches test daily; the nightly branch tests weekly (Sunday).

run: |

if [ "${PYTORCH_GIT_REF}" != "nightly" ] || [ "$(date -u +%u)" = "7" ]; then

echo "dispatch=true" >> "$GITHUB_OUTPUT"

else

echo "dispatch=false" >> "$GITHUB_OUTPUT"

fi

Done in 94b87db. i drropped the created_at lookup, the token, --run-date and the actions: read grant on this workflow and the release caller; the weekly day is the wall clock now, same as the PyTorch pipeline.
the rule and its Sunday/Tuesday tests stay in wants_multi_gpu(), so the workflow still only reads the matrix.

this does mean a nightly re-run onto another day loses that week's multi-accelerator run, so the configure job prints a ::notice:: naming the day and the test_size: large dispatch that runs those tests on demand, rather than skipping them quietly.

Review follow-up: reading the run creation date from the API cost a token, an actions: read grant on this workflow and its caller, and a fallback, to keep one case honest: re-running a whole Sunday nightly on a later day. Drop it and read the clock, which is how the PyTorch pipeline picks its weekly day. That case now loses the week of multi-accelerator coverage, so the configure job annotates a nightly that is off the day and names the dispatch that runs those tests anyway.

amd-shiraz

discussed offline on capacity side. lg to proceed for single + multi gpu gfx942 wise.

gulsumgudukbay · 2026-08-19T22:14:54Z

Validation on the nightly packages the HUD runs against (10.1.0a20260819): https://github.com/ROCm/TheRock/actions/runs/32307755867

It runs the configure job without the API call, the token or the actions: read grant, since the weekly day now comes from the clock as it does for PyTorch. release_type: dev resolves to the short scope, so this is the leg a dev or nightly dispatch takes off the weekly day.

github-project-automation Bot added this to TheRock Triage Jul 24, 2026

github-project-automation Bot moved this to TODO in TheRock Triage Jul 24, 2026

gulsumgudukbay mentioned this pull request Jul 24, 2026

ci(jax): release v0.10.0-v0.11.0, build single ref in PR CI, drop v0.9.1/native #6708

Merged

1 task

gulsumgudukbay requested a review from ScottTodd July 24, 2026 21:55

erman-gurses reviewed Jul 25, 2026

View reviewed changes

gulsumgudukbay requested a review from erman-gurses July 27, 2026 15:39

ScottTodd requested review from amd-shiraz and geomin12 July 27, 2026 15:58

gulsumgudukbay requested review from geomin12 and removed request for geomin12 July 27, 2026 16:13

gulsumgudukbay force-pushed the users/gulsumgudukbay/jax-test-multi-gpu-runner branch 2 times, most recently from 20a9d67 to 6373b5e Compare July 27, 2026 20:40

amd-shiraz requested changes Jul 27, 2026

View reviewed changes

geomin12 reviewed Jul 28, 2026

View reviewed changes

gulsumgudukbay mentioned this pull request Jul 28, 2026

ci(jax): respect GPU isolation on JAX test runners #6924

Merged

1 task

gulsumgudukbay requested a review from amd-shiraz July 29, 2026 15:07

ScottTodd reviewed Aug 3, 2026

View reviewed changes

gulsumgudukbay force-pushed the users/gulsumgudukbay/jax-test-multi-gpu-runner branch 2 times, most recently from 9c05238 to 687bd9e Compare August 10, 2026 06:19

gulsumgudukbay mentioned this pull request Aug 11, 2026

ci(jax): run the JAX suite at three sizes #6753

Open

1 task

gulsumgudukbay force-pushed the users/gulsumgudukbay/jax-test-multi-gpu-runner branch from ed88f03 to b87760f Compare August 11, 2026 22:16

gulsumgudukbay requested a review from ScottTodd August 11, 2026 22:39

ScottTodd reviewed Aug 11, 2026

View reviewed changes

erman-gurses reviewed Aug 12, 2026

View reviewed changes

erman-gurses approved these changes Aug 14, 2026

View reviewed changes

gulsumgudukbay mentioned this pull request Aug 14, 2026

ci(jax): read the suite environment from where it lives now #7387

Merged

3 tasks

gulsumgudukbay force-pushed the users/gulsumgudukbay/jax-test-multi-gpu-runner branch from b153cf6 to 00860be Compare August 17, 2026 01:43

gulsumgudukbay and others added 8 commits August 19, 2026 00:27

Update JAX Test runner to multi-GPU runner

af07b54

ci(jax): leave the test job room for its setup

56dee38

Installing the ROCm packages has taken over an hour on its own, so a 120 minute job budget cut the single-accelerator suite off mid-run. Give the job 240 and hold each subset to its own step budget instead.

gulsumgudukbay force-pushed the users/gulsumgudukbay/jax-test-multi-gpu-runner branch from 00860be to 56dee38 Compare August 19, 2026 00:40

ScottTodd approved these changes Aug 19, 2026

View reviewed changes

amd-shiraz approved these changes Aug 19, 2026

View reviewed changes

	steps:
	- name: Check full test cadence
	id: cadence
	# Release branches test daily; the nightly branch tests weekly (Sunday).
	run: \|
	if [ "${PYTORCH_GIT_REF}" != "nightly" ] \|\| [ "$(date -u +%u)" = "7" ]; then
	echo "dispatch=true" >> "$GITHUB_OUTPUT"
	else
	echo "dispatch=false" >> "$GITHUB_OUTPUT"
	fi

		SUBSET_ALL = "all"
		SUBSET_MULTI = "multi"

Uh oh!

Conversation

gulsumgudukbay commented Jul 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Technical Details

How a run reaches its jobs

Test Plan

Test Result

Submission Checklist

Uh oh!

therock-pr-bot Bot commented Jul 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

✅ All Checks Passed — Ready for Review

Uh oh!

therock-pr-bot Bot commented Jul 24, 2026

Uh oh!

erman-gurses left a comment

Choose a reason for hiding this comment

Uh oh!

gulsumgudukbay commented Jul 27, 2026

Uh oh!

ScottTodd commented Jul 27, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ScottTodd commented Jul 27, 2026

Uh oh!

gulsumgudukbay commented Jul 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

gulsumgudukbay commented Jul 27, 2026

Uh oh!

amd-shiraz commented Jul 27, 2026

Uh oh!

gulsumgudukbay commented Jul 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

geomin12 left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gulsumgudukbay commented Jul 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ScottTodd left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gulsumgudukbay commented Aug 11, 2026

Uh oh!

ScottTodd left a comment

Choose a reason for hiding this comment

Uh oh!

gulsumgudukbay commented Aug 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

erman-gurses left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gulsumgudukbay Aug 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gulsumgudukbay Aug 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gulsumgudukbay commented Jul 24, 2026 •

edited

Loading

therock-pr-bot Bot commented Jul 24, 2026 •

edited

Loading

gulsumgudukbay commented Jul 27, 2026 •

edited

Loading

gulsumgudukbay commented Jul 28, 2026 •

edited

Loading

geomin12 left a comment •

edited

Loading

gulsumgudukbay commented Jul 28, 2026 •

edited

Loading

gulsumgudukbay commented Aug 12, 2026 •

edited

Loading

erman-gurses left a comment •

edited

Loading

gulsumgudukbay Aug 14, 2026 •

edited

Loading

gulsumgudukbay Aug 14, 2026 •

edited

Loading

gulsumgudukbay Aug 14, 2026 •

edited

Loading

gulsumgudukbay Aug 14, 2026 •

edited

Loading

gulsumgudukbay Aug 14, 2026 •

edited

Loading

ScottTodd Aug 19, 2026 •

edited

Loading

erman-gurses Aug 12, 2026 •

edited

Loading

gulsumgudukbay Aug 14, 2026 •

edited

Loading

gulsumgudukbay commented Aug 14, 2026 •

edited

Loading