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