From 90a9455ff304076a1ffa3a34115a29fa90180632 Mon Sep 17 00:00:00 2001 From: Marco Minutoli Date: Thu, 21 May 2026 22:51:53 +0000 Subject: [PATCH 1/5] [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 159613552d09b01c4341544051a69ded5641abb9 Mon Sep 17 00:00:00 2001 From: gulsumgudukbay Date: Fri, 5 Jun 2026 21:53:19 +0000 Subject: [PATCH 2/5] [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 ea0ac83fabb4f921a82f952d599b16681f503da9 Mon Sep 17 00:00:00 2001 From: gulsumgudukbay Date: Mon, 8 Jun 2026 20:07:42 +0000 Subject: [PATCH 3/5] [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 aa92dbfeed1b..665aabfd0f50 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 cccdf5d8b324..e4df65fa08d1 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 88e42dccb269cb9f0aaf15b97a5f93880dc8c4f2 Mon Sep 17 00:00:00 2001 From: Pakize Sanal Date: Thu, 18 Jun 2026 15:43:15 -0500 Subject: [PATCH 4/5] 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 665aabfd0f50..aa92dbfeed1b 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 e4df65fa08d1..cccdf5d8b324 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 b96c0a91e525d64b09780e194534e711c58a31f7 Mon Sep 17 00:00:00 2001 From: Pakize Sanal Date: Thu, 18 Jun 2026 13:52:23 -0500 Subject: [PATCH 5/5] 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 +