Cap xdist workers per GPU to avoid oversubscription

.github/workflows/bazel_rocm.yml

@@ -28,7 +28,7 @@ on:
       runner:
         description: "Which runner should the workflow run on?"
         type: string
-        default: "linux-x86-64-1gpu-amd"
+        default: "amd-do-linux.jax.gpu.gfx950.1"
       python:
         description: "Which python version to test?"
         type: string

.github/workflows/bazel_rocm_presubmit.yml

@@ -33,7 +33,7 @@ jobs:
         fail-fast: false # don't cancel all jobs on failure
         matrix:
           python: ["3.14"]
-          runner: ["linux-x86-64-1gpu-amd"]
+          runner: ["amd-do-linux.jax.gpu.gfx950.1"]
           jaxlib-version: ["head", "pypi_latest"]
           enable-x64: [1]
           rocm-version: ["7"]

.github/workflows/build_rocm_artifacts.yml

@@ -10,9 +10,9 @@ on:
       runner:
         description: "Which runner should the workflow run on?"
         type: choice
-        default: "linux-x86-64-1gpu-amd"
+        default: "amd-do-linux.jax.gpu.gfx950.1"
         options:
-        - "linux-x86-64-1gpu-amd"
+        - "amd-do-linux.jax.gpu.gfx950.1"
       artifact:
         description: "Which ROCm artifact to build?"
         type: choice
@@ -63,7 +63,7 @@ on:
       runner:
         description: "Which runner should the workflow run on?"
         type: string
-        default: "linux-x86-64-1gpu-amd"
+        default: "amd-do-linux.jax.gpu.gfx950.1"
       artifact:
         description: "Which ROCm artifact to build?"
         type: string

.github/workflows/pytest_rocm.yml

@@ -19,11 +19,11 @@ on:
       runner:
         description: "Which runner should the workflow run on?"
         type: choice
-        default: "linux-x86-64-4gpu-amd"
+        default: "amd-do-linux.jax.gpu.gfx950.4"
         options:
-          - "linux-x86-64-1gpu-amd"
-          - "linux-x86-64-4gpu-amd"
-          - "linux-x86-64-8gpu-amd"
+          - "amd-do-linux.jax.gpu.gfx950.1"
+          - "amd-do-linux.jax.gpu.gfx950.4"
+          - "amd-do-linux.jax.gpu.gfx950.8"
       python:
         description: "Which Python version to use?"
         type: choice
@@ -75,7 +75,7 @@ on:
       runner:
         description: "Which runner should the workflow run on?"
         type: string
-        default: "linux-x86-64-4gpu-amd"
+        default: "amd-do-linux.jax.gpu.gfx950.4"
       python:
         description: "Which Python version to use?"
         type: string
@@ -117,6 +117,10 @@ env:
 
 jobs:
   run-tests:
+    strategy:
+      fail-fast: false
+      matrix:
+        workers_per_gpu: [16, 8, 4, 2, 1]
     defaults:
       run:
         # Set the shell to bash as GitHub actions run with /bin/sh by default
@@ -132,7 +136,7 @@ jobs:
       options: --device=/dev/kfd --device=/dev/dri --security-opt seccomp=unconfined --group-add video --shm-size 64G --env-file /etc/podinfo/gha-gpu-isolation-settings
     name: "${{ (contains(inputs.runner, '1gpu') && '1gpu') ||
         (contains(inputs.runner, '4gpu') && '4gpu') ||
-        (contains(inputs.runner, '8gpu') && '8gpu') }}, ROCm ${{ inputs.rocm-version }}, py${{ inputs.python }}"
+        (contains(inputs.runner, '8gpu') && '8gpu') }}, n=${{ matrix.workers_per_gpu }}, ROCm ${{ inputs.rocm-version }}, py${{ inputs.python }}"
 
     env:
       JAXCI_HERMETIC_PYTHON_VERSION: "${{ inputs.python }}"
@@ -143,6 +147,8 @@ jobs:
       # warning when this is unset; the env var must be in place before HIP
       # initialization, which happens at JAX import time.
       HSA_NO_SCRATCH_RECLAIM: "1"
+      # Temp: Sweep the xdist worker cap across matrix values
+      JAXCI_MAX_WORKERS_PER_GPU: ${{ matrix.workers_per_gpu }}
 
     steps:
       - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3

.github/workflows/rocm-ci.yml

@@ -13,7 +13,7 @@ concurrency:
 jobs:
   build-jax-in-docker:
     if: github.repository_owner == 'jax-ml' || github.repository_owner == 'ROCm'
-    runs-on: linux-x86_64-cirrascale-64-8gpu-amd-mi250
+    runs-on: amd-do-linux.jax.gpu.gfx950.8
     env:
       BASE_IMAGE: "ubuntu:22.04"
       TEST_IMAGE: ubuntu-jax-upstream-${{ github.run_id }}_${{ github.run_number }}_${{ github.run_attempt }}

.github/workflows/wheel_tests_continuous.yml

@@ -354,11 +354,13 @@ jobs:
     strategy:
         fail-fast: false
         matrix:
-          runner: ["linux-x86-64-1gpu-amd"]
+          runner: ["amd-do-linux.jax.gpu.gfx950.1"]
           artifact: ["jax-rocm-plugin", "jax-rocm-pjrt"]
           python: ["3.12", "3.13", "3.14"]
           rocm: [
-            {label: "ROCm 7.2.0", wheel-version: "7", release-version: "7.2.0", manylinux-image: ""},
+            {label: "ROCm 7.2.0",     wheel-version: "7", release-version: "7.2.0",        manylinux-image: ""},
+            {label: "TheRock 7.13",   wheel-version: "7", release-version: "therock-7.13.0",  manylinux-image: "ghcr.io/rocm/jax-manylinux_2_28-therock-7.13.0:latest"},
+            {label: "TheRock latest", wheel-version: "7", release-version: "therock-latest", manylinux-image: "ghcr.io/rocm/jax-manylinux_2_28-therock-latest:latest"},
           ]
           exclude:
             - artifact: "jax-rocm-pjrt"
@@ -388,10 +390,12 @@ jobs:
     strategy:
         fail-fast: false
         matrix:
-          runner: ["linux-x86-64-1gpu-amd", "linux-x86-64-4gpu-amd", "linux-x86-64-8gpu-amd"]
+          runner: ["amd-do-linux.jax.gpu.gfx950.1", "amd-do-linux.jax.gpu.gfx950.4", "amd-do-linux.jax.gpu.gfx950.8"]
           python: ["3.12"]
           rocm: [
-            {label: "ROCm 7.2.0", wheel-version: "7", release-version: "7.2.0", tag: "rocm720"},
+            {label: "ROCm 7.2.0",     wheel-version: "7", release-version: "7.2.0",        tag: "rocm720"},
+            {label: "TheRock 7.13",   wheel-version: "7", release-version: "therock-7.13.0",  tag: "therock-7.13.0"},
+            {label: "TheRock latest", wheel-version: "7", release-version: "therock-latest", tag: "therock-latest"},
           ]
     name: "Pytest ROCm (${{ matrix.rocm.label }})"
     with:
@@ -415,10 +419,12 @@ jobs:
     strategy:
         fail-fast: false
         matrix:
-          runner: ["linux-x86-64-4gpu-amd", "linux-x86-64-8gpu-amd"]
+          runner: ["amd-do-linux.jax.gpu.gfx950.4", "amd-do-linux.jax.gpu.gfx950.8"]
           python: ["3.12"]
           rocm: [
-            {label: "ROCm 7.2.0", wheel-version: "7", release-version: "7.2.0", tag: "rocm720"},
+            {label: "ROCm 7.2.0",     wheel-version: "7", release-version: "7.2.0",        tag: "rocm720"},
+            {label: "TheRock 7.13",   wheel-version: "7", release-version: "therock-7.13.0",  tag: "therock-7.13.0"},
+            {label: "TheRock latest", wheel-version: "7", release-version: "therock-latest", tag: "therock-latest"},
           ]
           enable-x64: [0]
     name: "Bazel ROCm (${{ matrix.rocm.label }})"

.github/workflows/wheel_tests_nightly_release.yml

@@ -183,11 +183,11 @@ jobs:
     strategy:
         fail-fast: false
         matrix:
-          runner: ["linux-x86-64-1gpu-amd"]
+          runner: ["amd-do-linux.jax.gpu.gfx950.1"]
           artifact: ["jax-rocm-plugin", "jax-rocm-pjrt"]
           python: ["3.12", "3.13", "3.14"]
           rocm: [
-            {label: "ROCm 7.2.0", wheel-version: "7", release-version: "7.2.0", manylinux-image: ""},
+            {label: "ROCm 7.2.0",     wheel-version: "7", release-version: "7.2.0",        manylinux-image: ""},
           ]
           exclude:
             - artifact: "jax-rocm-pjrt"
@@ -217,10 +217,10 @@ jobs:
     strategy:
         fail-fast: false # don't cancel all jobs on failure
         matrix:
-          runner: ["linux-x86-64-1gpu-amd", "linux-x86-64-4gpu-amd", "linux-x86-64-8gpu-amd"]
+          runner: ["amd-do-linux.jax.gpu.gfx950.1", "amd-do-linux.jax.gpu.gfx950.4", "amd-do-linux.jax.gpu.gfx950.8"]
           python: ["3.12", "3.13", "3.14"]
           rocm: [
-            {label: "ROCm 7.2.0", wheel-version: "7", release-version: "7.2.0", tag: "rocm720"},
+            {label: "ROCm 7.2.0",     wheel-version: "7", release-version: "7.2.0",        tag: "rocm720"},
           ]
     name: "Pytest ROCm (${{ matrix.rocm.label }})"
     with:
@@ -244,10 +244,10 @@ jobs:
     strategy:
         fail-fast: false # don't cancel all jobs on failure
         matrix:
-          runner: ["linux-x86-64-1gpu-amd"]
+          runner: ["amd-do-linux.jax.gpu.gfx950.1"]
           python: ["3.12", "3.13", "3.14"]
           rocm: [
-            {label: "ROCm 7.2.0", wheel-version: "7", release-version: "7.2.0", tag: "rocm720"},
+            {label: "ROCm 7.2.0",     wheel-version: "7", release-version: "7.2.0",        tag: "rocm720"},
           ]
           enable-x64: [0]
     name: "Bazel ROCm (${{ matrix.rocm.label }})"

ci/run_pytest_rocm.sh

@@ -102,6 +102,15 @@ if [[ $host_memory_limit -lt $num_processes ]]; then
   echo "Adjusting num_processes to match host memory limit: $num_processes"
 fi
 
+# Cap workers per GPU to avoid oversubscribing a single device.
+# max_workers_per_gpu=8
+max_workers_per_gpu="${JAXCI_MAX_WORKERS_PER_GPU:-8}"
+echo "max_workers_per_gpu=$max_workers_per_gpu (JAXCI_MAX_WORKERS_PER_GPU='${JAXCI_MAX_WORKERS_PER_GPU:-<unset>}')"
+if [[ $((gpu_count * max_workers_per_gpu)) -lt $num_processes ]]; then
+  num_processes=$((gpu_count * max_workers_per_gpu))
+  echo "Capping num_processes to $max_workers_per_gpu/GPU: $num_processes"
+fi
+
 if [[ 16 -lt $num_processes ]]; then
   num_processes=16
   echo "Reducing num_processes to $num_processes"

conftest.py

@@ -95,3 +95,12 @@ def pytest_collection() -> None:
     # HIP_VISIBLE_DEVICES to all GPUs; override to "0" so HIP doesn't try to
     # enable agents that ROCr just hid.
     os.environ["HIP_VISIBLE_DEVICES"] = "0"
+
+    # MIOpen serializes kernel lookup through a shared user-DB lock under
+    # /tmp/miopen-lockfiles. Use a per-worker cache directory to avoid lock
+    # contention when xdist workers share the same GPU.
+    miopen_cache_dir = f"/tmp/miopen-{xdist_worker_name}"
+    os.makedirs(miopen_cache_dir, exist_ok=True)
+    os.environ["MIOPEN_USER_DB_PATH"] = miopen_cache_dir
+    os.environ["MIOPEN_CUSTOM_CACHE_DIR"] = miopen_cache_dir
+

jaxlib/rocm/rocm_rpath.bzl

@@ -65,7 +65,6 @@ def _rocm_wheel_rpaths():
     where TheRock is in the container's system Python). The legacy `/opt/rocm`
     fallback is appended by the caller so it stays last.
     """
-
     # `_rocm_sdk_core/lib` holds the non-arch runtime libs; `_rocm_sdk_libraries
     # [...]/lib` holds the math libs + per-arch kernel data. site_libs covers the
     # multi-arch + per-family packages; core_libs is the cross-Python/absolute