From dfb1a19595f7b10061338f73112906a4864f6651 Mon Sep 17 00:00:00 2001 From: Marco Minutoli Date: Thu, 21 May 2026 22:51:53 +0000 Subject: [PATCH 1/9] [ROCm] Add TheRock 7.13 and TheRock latest to CI test matrices Extends the wheel test continuous and nightly release workflows to run against TheRock 7.13.0 and TheRock latest builds in addition to ROCm 7.2.0. --- .github/workflows/wheel_tests_continuous.yml | 12 +++++++++--- .github/workflows/wheel_tests_nightly_release.yml | 12 +++++++++--- 2 files changed, 18 insertions(+), 6 deletions(-) diff --git a/.github/workflows/wheel_tests_continuous.yml b/.github/workflows/wheel_tests_continuous.yml index 73e3bdfa385c..575d2cd68c60 100644 --- a/.github/workflows/wheel_tests_continuous.yml +++ b/.github/workflows/wheel_tests_continuous.yml @@ -358,7 +358,9 @@ jobs: 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" @@ -391,7 +393,9 @@ jobs: runner: ["linux-x86-64-1gpu-amd", "linux-x86-64-4gpu-amd", "linux-x86-64-8gpu-amd"] 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: @@ -418,7 +422,9 @@ jobs: runner: ["linux-x86-64-4gpu-amd", "linux-x86-64-8gpu-amd"] 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 }})" diff --git a/.github/workflows/wheel_tests_nightly_release.yml b/.github/workflows/wheel_tests_nightly_release.yml index 9c5dfb9ce04c..f2b43d140d7b 100644 --- a/.github/workflows/wheel_tests_nightly_release.yml +++ b/.github/workflows/wheel_tests_nightly_release.yml @@ -187,7 +187,9 @@ jobs: 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" @@ -220,7 +222,9 @@ jobs: runner: ["linux-x86-64-1gpu-amd", "linux-x86-64-4gpu-amd", "linux-x86-64-8gpu-amd"] 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"}, + {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: @@ -247,7 +251,9 @@ jobs: runner: ["linux-x86-64-1gpu-amd"] 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"}, + {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 }})" From d430e3bbedbdaf8a90dde8857e51cbe111c3c00e Mon Sep 17 00:00:00 2001 From: gulsumgudukbay Date: Fri, 5 Jun 2026 21:53:19 +0000 Subject: [PATCH 2/9] [ROCm] Embed TheRock-aware wheel RUNPATH at link time Consolidate the ROCm wheel RUNPATH strategy in jaxlib/rocm/rocm_rpath.bzl and embed it at link time via _WHEEL_RPATHS instead of a post-build patchelf step (applied under rocm_path_type=link_only, which the release wheel config sets). A single wheel resolves ROCm libs across TheRock pip wheels (1- and 2-deep $ORIGIN layouts), TheRock tarballs under /opt/rocm-, and legacy /opt/rocm; the loader silently skips $ORIGIN entries whose directories don't exist. Also pin the hermetic build ROCm to the image's TheRock version via ROCM_DISTRO_URL so the wheel's ROCm SONAMEs match the runtime image. --- jaxlib/rocm/rocm_rpath.bzl | 1 - 1 file changed, 1 deletion(-) diff --git a/jaxlib/rocm/rocm_rpath.bzl b/jaxlib/rocm/rocm_rpath.bzl index c8c0c1bf4fa3..0030a41e0b5f 100644 --- a/jaxlib/rocm/rocm_rpath.bzl +++ b/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 From b0ab89091dedccfb7329a3f861dbfeb5fe5bfc88 Mon Sep 17 00:00:00 2001 From: gulsumgudukbay Date: Mon, 8 Jun 2026 20:07:42 +0000 Subject: [PATCH 3/9] [ROCm] Put TheRock SDK bin on PATH via so rocminfo/rocm-smi resolve in /opt/rocm-less images From cb097f7a03d212724e8ccbc7e7231d612d221106 Mon Sep 17 00:00:00 2001 From: Marco Minutoli Date: Fri, 19 Jun 2026 11:37:40 -0700 Subject: [PATCH 4/9] [ROCm] Update CI runner labels to amd-do-linux.jax.gpu.gfx950.{1,4,8} (#807) --- .github/workflows/bazel_rocm.yml | 2 +- .github/workflows/bazel_rocm_presubmit.yml | 2 +- .github/workflows/build_rocm_artifacts.yml | 6 +++--- .github/workflows/pytest_rocm.yml | 10 +++++----- .github/workflows/rocm-ci.yml | 2 +- .github/workflows/wheel_tests_continuous.yml | 6 +++--- .github/workflows/wheel_tests_nightly_release.yml | 6 +++--- 7 files changed, 17 insertions(+), 17 deletions(-) diff --git a/.github/workflows/bazel_rocm.yml b/.github/workflows/bazel_rocm.yml index aa92dbfeed1b..b9cace44acd7 100644 --- a/.github/workflows/bazel_rocm.yml +++ b/.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 diff --git a/.github/workflows/bazel_rocm_presubmit.yml b/.github/workflows/bazel_rocm_presubmit.yml index 6a3a4bbbf5ec..7749795e08fa 100644 --- a/.github/workflows/bazel_rocm_presubmit.yml +++ b/.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"] diff --git a/.github/workflows/build_rocm_artifacts.yml b/.github/workflows/build_rocm_artifacts.yml index cccdf5d8b324..64a39f961b30 100644 --- a/.github/workflows/build_rocm_artifacts.yml +++ b/.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 diff --git a/.github/workflows/pytest_rocm.yml b/.github/workflows/pytest_rocm.yml index 8d695851ba4c..d1f1d9aac486 100644 --- a/.github/workflows/pytest_rocm.yml +++ b/.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 diff --git a/.github/workflows/rocm-ci.yml b/.github/workflows/rocm-ci.yml index 144d0521172b..f184cdf2a25c 100644 --- a/.github/workflows/rocm-ci.yml +++ b/.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 }} diff --git a/.github/workflows/wheel_tests_continuous.yml b/.github/workflows/wheel_tests_continuous.yml index 575d2cd68c60..8251ab7e7d3e 100644 --- a/.github/workflows/wheel_tests_continuous.yml +++ b/.github/workflows/wheel_tests_continuous.yml @@ -354,7 +354,7 @@ 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: [ @@ -390,7 +390,7 @@ 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"}, @@ -419,7 +419,7 @@ 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"}, diff --git a/.github/workflows/wheel_tests_nightly_release.yml b/.github/workflows/wheel_tests_nightly_release.yml index f2b43d140d7b..b547491913f9 100644 --- a/.github/workflows/wheel_tests_nightly_release.yml +++ b/.github/workflows/wheel_tests_nightly_release.yml @@ -183,7 +183,7 @@ 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: [ @@ -219,7 +219,7 @@ 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"}, @@ -248,7 +248,7 @@ 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"}, From 4da859133d635b003162b92c41763e3e3ac203bc Mon Sep 17 00:00:00 2001 From: gulsumgudukbay Date: Mon, 8 Jun 2026 20:07:42 +0000 Subject: [PATCH 5/9] [ROCm] Put TheRock SDK bin on PATH via so rocminfo/rocm-smi resolve in /opt/rocm-less images --- .github/workflows/bazel_rocm.yml | 4 ++++ .github/workflows/build_rocm_artifacts.yml | 4 ++++ 2 files changed, 8 insertions(+) diff --git a/.github/workflows/bazel_rocm.yml b/.github/workflows/bazel_rocm.yml index b9cace44acd7..75eff691487d 100644 --- a/.github/workflows/bazel_rocm.yml +++ b/.github/workflows/bazel_rocm.yml @@ -136,10 +136,14 @@ jobs: # SDK bin dir, so put it on PATH via rocm-sdk. apt-ROCm images lack # rocm-sdk and already have rocminfo on PATH. if command -v rocm-sdk >/dev/null 2>&1; then +<<<<<<< HEAD sdk_bin="$(rocm-sdk path --bin)" if [[ -n "$sdk_bin" ]]; then export PATH="$sdk_bin:$PATH" fi +======= + export PATH="$(rocm-sdk path --bin):$PATH" +>>>>>>> b8c52b15be ([ROCm] Put TheRock SDK bin on PATH via so rocminfo/rocm-smi resolve in /opt/rocm-less images) fi rocminfo - name: Configure AWS Credentials diff --git a/.github/workflows/build_rocm_artifacts.yml b/.github/workflows/build_rocm_artifacts.yml index 64a39f961b30..43a7eb3d9b72 100644 --- a/.github/workflows/build_rocm_artifacts.yml +++ b/.github/workflows/build_rocm_artifacts.yml @@ -154,10 +154,14 @@ jobs: # SDK bin dir, so put it on PATH via rocm-sdk. apt-ROCm images lack # rocm-sdk and already have rocminfo on PATH. if command -v rocm-sdk >/dev/null 2>&1; then +<<<<<<< HEAD sdk_bin="$(rocm-sdk path --bin)" if [[ -n "$sdk_bin" ]]; then export PATH="$sdk_bin:$PATH" fi +======= + export PATH="$(rocm-sdk path --bin):$PATH" +>>>>>>> b8c52b15be ([ROCm] Put TheRock SDK bin on PATH via so rocminfo/rocm-smi resolve in /opt/rocm-less images) fi rocminfo # Halt for testing From f4b6c808199b09a5a5b94a854e644e7d3fc8813d Mon Sep 17 00:00:00 2001 From: Pakize Sanal Date: Thu, 18 Jun 2026 15:43:15 -0500 Subject: [PATCH 6/9] Resolve ROCm workflow conflict markers --- .github/workflows/bazel_rocm.yml | 4 ---- .github/workflows/build_rocm_artifacts.yml | 4 ---- 2 files changed, 8 deletions(-) diff --git a/.github/workflows/bazel_rocm.yml b/.github/workflows/bazel_rocm.yml index 75eff691487d..b9cace44acd7 100644 --- a/.github/workflows/bazel_rocm.yml +++ b/.github/workflows/bazel_rocm.yml @@ -136,14 +136,10 @@ jobs: # SDK bin dir, so put it on PATH via rocm-sdk. apt-ROCm images lack # rocm-sdk and already have rocminfo on PATH. if command -v rocm-sdk >/dev/null 2>&1; then -<<<<<<< HEAD sdk_bin="$(rocm-sdk path --bin)" if [[ -n "$sdk_bin" ]]; then export PATH="$sdk_bin:$PATH" fi -======= - export PATH="$(rocm-sdk path --bin):$PATH" ->>>>>>> b8c52b15be ([ROCm] Put TheRock SDK bin on PATH via so rocminfo/rocm-smi resolve in /opt/rocm-less images) fi rocminfo - name: Configure AWS Credentials diff --git a/.github/workflows/build_rocm_artifacts.yml b/.github/workflows/build_rocm_artifacts.yml index 43a7eb3d9b72..64a39f961b30 100644 --- a/.github/workflows/build_rocm_artifacts.yml +++ b/.github/workflows/build_rocm_artifacts.yml @@ -154,14 +154,10 @@ jobs: # SDK bin dir, so put it on PATH via rocm-sdk. apt-ROCm images lack # rocm-sdk and already have rocminfo on PATH. if command -v rocm-sdk >/dev/null 2>&1; then -<<<<<<< HEAD sdk_bin="$(rocm-sdk path --bin)" if [[ -n "$sdk_bin" ]]; then export PATH="$sdk_bin:$PATH" fi -======= - export PATH="$(rocm-sdk path --bin):$PATH" ->>>>>>> b8c52b15be ([ROCm] Put TheRock SDK bin on PATH via so rocminfo/rocm-smi resolve in /opt/rocm-less images) fi rocminfo # Halt for testing From b7935f2fbb5dd97df94ea96acc438cdc53adf631 Mon Sep 17 00:00:00 2001 From: Pakize Sanal Date: Thu, 18 Jun 2026 13:52:23 -0500 Subject: [PATCH 7/9] Isolate MIOpen cache per xdist worker to avoid lock contention --- conftest.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/conftest.py b/conftest.py index 3abbeef26136..1b57b5d00ef2 100644 --- a/conftest.py +++ b/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 + From 2327338d112f0367fe46c4fd7a47f639ae77ea6c Mon Sep 17 00:00:00 2001 From: Pakize Sanal Date: Thu, 18 Jun 2026 18:20:28 -0500 Subject: [PATCH 8/9] Cap xdist workers per GPU to avoid oversubscription --- ci/run_pytest_rocm.sh | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/ci/run_pytest_rocm.sh b/ci/run_pytest_rocm.sh index 92043dba49fa..19162a896285 100755 --- a/ci/run_pytest_rocm.sh +++ b/ci/run_pytest_rocm.sh @@ -102,6 +102,14 @@ 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 + +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" From 63140c7c41ef3a46c541252a656c657564f78c1b Mon Sep 17 00:00:00 2001 From: Pakize Sanal Date: Fri, 19 Jun 2026 07:46:50 -0500 Subject: [PATCH 9/9] Temp: Sweep per-GPU worker count to find the best value --- .github/workflows/pytest_rocm.yml | 8 +++++++- .github/workflows/wheel_tests_nightly_release.yml | 6 ------ ci/run_pytest_rocm.sh | 5 +++-- 3 files changed, 10 insertions(+), 9 deletions(-) diff --git a/.github/workflows/pytest_rocm.yml b/.github/workflows/pytest_rocm.yml index d1f1d9aac486..061edd4b2b93 100644 --- a/.github/workflows/pytest_rocm.yml +++ b/.github/workflows/pytest_rocm.yml @@ -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 diff --git a/.github/workflows/wheel_tests_nightly_release.yml b/.github/workflows/wheel_tests_nightly_release.yml index b547491913f9..2750d6954284 100644 --- a/.github/workflows/wheel_tests_nightly_release.yml +++ b/.github/workflows/wheel_tests_nightly_release.yml @@ -188,8 +188,6 @@ jobs: python: ["3.12", "3.13", "3.14"] rocm: [ {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" @@ -223,8 +221,6 @@ jobs: python: ["3.12", "3.13", "3.14"] rocm: [ {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: @@ -252,8 +248,6 @@ jobs: python: ["3.12", "3.13", "3.14"] rocm: [ {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 }})" diff --git a/ci/run_pytest_rocm.sh b/ci/run_pytest_rocm.sh index 19162a896285..4228f6d1f96b 100755 --- a/ci/run_pytest_rocm.sh +++ b/ci/run_pytest_rocm.sh @@ -103,8 +103,9 @@ if [[ $host_memory_limit -lt $num_processes ]]; then fi # Cap workers per GPU to avoid oversubscribing a single device. -max_workers_per_gpu=8 - +# 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:-}')" 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"