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..061edd4b2b93 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 @@ -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/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 73e3bdfa385c..8251ab7e7d3e 100644 --- a/.github/workflows/wheel_tests_continuous.yml +++ b/.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 }})" diff --git a/.github/workflows/wheel_tests_nightly_release.yml b/.github/workflows/wheel_tests_nightly_release.yml index 9c5dfb9ce04c..2750d6954284 100644 --- a/.github/workflows/wheel_tests_nightly_release.yml +++ b/.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 }})" diff --git a/ci/run_pytest_rocm.sh b/ci/run_pytest_rocm.sh index 92043dba49fa..4228f6d1f96b 100755 --- a/ci/run_pytest_rocm.sh +++ b/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:-}')" +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" 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