From db2242f82f1f6b726788b080aef8a20820f17db7 Mon Sep 17 00:00:00 2001 From: Manjunath Gaonkar Date: Wed, 15 Jul 2026 13:32:34 -0500 Subject: [PATCH 1/6] Enable additional single-GPU bazel tests in ROCm CI (#816) Drop 43 of the 45 test-target exclusions in build/rocm/ci_test_targets.txt, leaving only two: - buffer_callback_test_gpu, whose py_import deps are incompatible with pre-built plugin wheels (build_jaxlib=false). - pallas:ops_test_mgpu_gpu, a multi-accelerator test that is excluded from single_gpu runs anyway. Most of the dropped exclusions traced back to a complex-GEMM gap that no longer reproduces: C64 and C128 matmul return correct results on ROCm 7.14. The three teardown-hang exclusions (cholesky_update_test, pallas:pallas_vmap_test and pallas:gpu_paged_attention_test) each passed five consecutive runs with no hang. This brings the suite to 114 single-GPU GPU targets. Validated on MI300X (gfx942); the bazel CI pool is gfx950, so CI is the real gate. --- build/rocm/ci_test_targets.txt | 46 +++------------------------------- 1 file changed, 3 insertions(+), 43 deletions(-) diff --git a/build/rocm/ci_test_targets.txt b/build/rocm/ci_test_targets.txt index 55b28e7fc425..9174ea492bd7 100644 --- a/build/rocm/ci_test_targets.txt +++ b/build/rocm/ci_test_targets.txt @@ -7,50 +7,10 @@ //tests/pallas:gpu_tests //tests/pallas:backend_independent_tests //jaxlib/tools:check_gpu_wheel_sources_test --//tests/pallas:pallas_test_gpu --//tests/pallas:ops_test_gpu + +# Multi-accelerator test (not single-GPU; excluded from single_gpu runs anyway). -//tests/pallas:ops_test_mgpu_gpu --//tests/pallas:pallas_shape_poly_test_gpu --//tests/pallas:pallas_vmap_test_gpu --//tests/pallas:triton_pallas_test_gpu --//tests/pallas:gpu_paged_attention_test_gpu --//tests:export_harnesses_multi_platform_test_gpu --//tests:jet_test_gpu --//tests:lax_autodiff_test_gpu --//tests:lax_numpy_setops_test_gpu --//tests:lax_numpy_test_gpu --//tests:lax_test_gpu --//tests:linalg_test_gpu --//tests:logging_test_gpu --//tests:random_lax_test_gpu --//tests:scipy_signal_test_gpu --//tests:stax_test_gpu --//tests:ode_test_gpu --//tests:lobpcg_test_gpu --//tests:scipy_stats_test_gpu --//tests:nn_test_gpu --//tests:lax_scipy_sparse_test_gpu --//tests:lax_scipy_spectral_dac_test_gpu --//tests:lax_scipy_special_functions_test_gpu --//tests:cholesky_update_test_gpu --//tests:api_test_gpu --//tests:ann_test_gpu --//tests:experimental_rnn_test_gpu --//tests:lax_vmap_test_gpu --//tests:qdwh_test_gpu --//tests:scaled_dot_test_gpu --//tests:scipy_spatial_test_gpu --//tests:shape_poly_test_gpu --//tests:sparsify_test_gpu --//tests:lax_numpy_reducers_test_gpu --//tests:scipy_optimize_test_gpu --//tests:lax_numpy_einsum_test_gpu --//tests:lax_scipy_test_gpu --//tests:sparse_test_gpu --//tests:sparse_bcoo_bcsr_test_gpu --//tests:svd_test_gpu --//tests:eigh_test_gpu --//tests:image_test_gpu + # buffer_callback_test_gpu uses py_import deps that are incompatible with # pre-built plugin wheels (build_jaxlib=false); exclude it unconditionally. -//tests:buffer_callback_test_gpu From 1983294b68149a595502d888b1009b0f79e19c22 Mon Sep 17 00:00:00 2001 From: Oluwapelumi <88652093+pelumi1163@users.noreply.github.com> Date: Thu, 6 Aug 2026 20:54:33 +0100 Subject: [PATCH 2/6] Move Bazel test to the jax-base image (#831) --- .github/workflows/bazel_rocm.yml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/bazel_rocm.yml b/.github/workflows/bazel_rocm.yml index 63497e73d034..04ffffd23ee2 100644 --- a/.github/workflows/bazel_rocm.yml +++ b/.github/workflows/bazel_rocm.yml @@ -115,7 +115,7 @@ jobs: shell: bash runs-on: ${{ inputs.runner }} container: - image: ghcr.io/rocm/jax-dev-ubu24.${{ inputs.rocm-tag }}:latest # zizmor: ignore[unpinned-images] + image: ghcr.io/rocm/jax-base-ubu24.${{ inputs.rocm-tag }}:latest # zizmor: ignore[unpinned-images] volumes: - /data:/data options: >- @@ -185,8 +185,8 @@ jobs: env: ROCM_TAG: ${{ inputs.rocm-tag }} run: | - IMAGE="ghcr.io/rocm/jax-dev-ubu24.${ROCM_TAG}" - REPO="rocm/jax-dev-ubu24.${ROCM_TAG}" + IMAGE="ghcr.io/rocm/jax-base-ubu24.${ROCM_TAG}" + REPO="rocm/jax-base-ubu24.${ROCM_TAG}" TOKEN=$(curl -sf "https://ghcr.io/token?service=ghcr.io&scope=repository:${REPO}:pull" | \ python3 -c "import sys,json; print(json.load(sys.stdin)['token'])") echo "::add-mask::${TOKEN}" From f682b257a40fee469c4308f2f5845771b5ad2deb Mon Sep 17 00:00:00 2001 From: Marco Minutoli Date: Tue, 4 Aug 2026 18:10:41 +0000 Subject: [PATCH 3/6] [ROCm] nightly: build jax/jaxlib on fork runners Upstream CI has stopped publishing jax/jaxlib nightly wheels to GCS (gs://jax-nightly-artifacts/latest), leaving the ROCm nightly test legs running against missing/stale base wheels. On the ROCm fork, build the base jax and jaxlib wheels on the same amd-do-linux.jax.cpu runners as the ROCm plugin/pjrt wheels and upload them to the same S3 prefix, so the ROCm test legs consume freshly built wheels instead of the broken GCS bucket. Upstream (jax-ml) behavior is unchanged. - ci/build_rocm_artifacts.sh: allow "jax"/"jaxlib" artifacts; skip auditwheel for the pure-python jax wheel. - wheel_tests_nightly_release.yml: new fork-only build-rocm-jax-jaxlib job; wire it into the ROCm test legs and select S3 as the jax/jaxlib source on the fork. - pytest_rocm.yml / bazel_rocm.yml: new optional jax-jaxlib-source input (default gcs), forwarded to the download action. - download-jax-rocm-wheels: when source is s3, fetch jax/jaxlib from the CloudFront listing alongside the ROCm wheels instead of GCS. (cherry picked from commit 8654db624806ff686405864af739aae8d888a431) --- .../download-jax-rocm-wheels/action.yml | 39 +++++++++++++-- .github/workflows/bazel_rocm.yml | 5 ++ .github/workflows/pytest_rocm.yml | 9 ++++ .../workflows/wheel_tests_nightly_release.yml | 50 ++++++++++++++++++- ci/build_rocm_artifacts.sh | 15 ++++-- 5 files changed, 110 insertions(+), 8 deletions(-) diff --git a/.github/actions/download-jax-rocm-wheels/action.yml b/.github/actions/download-jax-rocm-wheels/action.yml index 1b9c18c70302..b1f91eceed5f 100644 --- a/.github/actions/download-jax-rocm-wheels/action.yml +++ b/.github/actions/download-jax-rocm-wheels/action.yml @@ -27,6 +27,9 @@ inputs: s3_download_uri: description: "S3 URI for ROCm plugin/PJRT wheels (use 'latest' to resolve via LATEST pointer)" default: 'latest' + jax-jaxlib-source: + description: "Where to fetch the base jax/jaxlib wheels from: 'gcs' (upstream nightly bucket) or 's3' (built alongside the ROCm wheels on the fork's runners)" + default: 'gcs' permissions: {} runs: using: "composite" @@ -74,7 +77,11 @@ runs: mkdir -p $(pwd)/dist # Disable Authentication on the ROCm runners. gcloud config set auth/disable_credentials True - if [[ "${INPUTS_DOWNLOAD_JAX_FROM_GCS}" == "1" ]]; then + if [[ "${INPUTS_JAX_JAXLIB_SOURCE}" == "s3" ]]; then + # jax is fetched from S3/CloudFront in the jaxlib "head" block below, + # alongside jaxlib and the ROCm plugin/pjrt wheels. + echo "jax wheel will be downloaded from S3/CloudFront alongside the ROCm wheels." + elif [[ "${INPUTS_DOWNLOAD_JAX_FROM_GCS}" == "1" ]]; then gcloud storage cp -r "${INPUTS_GCS_DOWNLOAD_URI}"/jax*py3*none*any.whl $(pwd)/dist/ else echo "JAX wheel won't be downloaded, only jaxlib pre-built wheel is tested." @@ -86,7 +93,9 @@ runs: echo "JAX only release. Only downloading the jax wheel from the release bucket." else if [[ ${INPUTS_JAXLIB_VERSION} == "head" ]]; then - gcloud storage cp -r "${INPUTS_GCS_DOWNLOAD_URI}/jaxlib*${PYTHON_MAJOR_MINOR}*${OS}*${ARCH}*.whl" $(pwd)/dist/ + if [[ "${INPUTS_JAX_JAXLIB_SOURCE}" != "s3" ]]; then + gcloud storage cp -r "${INPUTS_GCS_DOWNLOAD_URI}/jaxlib*${PYTHON_MAJOR_MINOR}*${OS}*${ARCH}*.whl" $(pwd)/dist/ + fi # Resolve ROCm plugin/PJRT wheels path via CloudFront. ROCM_WHEELS_BASE_URL="https://d22q5eopkfeftw.cloudfront.net" @@ -112,8 +121,31 @@ runs: exit 1 fi + DOWNLOAD_URLS=("${WHEELS_URL}/${PJRT_FILE}" "${WHEELS_URL}/${PLUGIN_FILE}") + + # On the fork, the base jax/jaxlib wheels are built alongside the + # ROCm wheels and uploaded to the same S3 prefix, so fetch them from + # the same CloudFront listing instead of the (broken) GCS bucket. + if [[ "${INPUTS_JAX_JAXLIB_SOURCE}" == "s3" ]]; then + JAXLIB_FILE=$(echo "$FILES" | grep "jaxlib-" | grep "${PYTHON_MAJOR_MINOR}" | head -n1 || true) + if [[ -z "${JAXLIB_FILE}" ]]; then + echo "Could not find jaxlib wheel in ${WHEELS_URL}" + exit 1 + fi + DOWNLOAD_URLS+=("${WHEELS_URL}/${JAXLIB_FILE}") + + if [[ "${INPUTS_DOWNLOAD_JAX_FROM_GCS}" == "1" ]]; then + JAX_FILE=$(echo "$FILES" | grep -E "jax-[0-9].*-py3-none-any\.whl" | head -n1 || true) + if [[ -z "${JAX_FILE}" ]]; then + echo "Could not find jax wheel in ${WHEELS_URL}" + exit 1 + fi + DOWNLOAD_URLS+=("${WHEELS_URL}/${JAX_FILE}") + fi + fi + PYTHON_BIN="python${INPUTS_PYTHON}" - "${PYTHON_BIN}" -m pip download --no-deps --dest "$(pwd)/dist" "${WHEELS_URL}/${PJRT_FILE}" "${WHEELS_URL}/${PLUGIN_FILE}" + "${PYTHON_BIN}" -m pip download --no-deps --dest "$(pwd)/dist" "${DOWNLOAD_URLS[@]}" elif [[ ${INPUTS_JAXLIB_VERSION} == "pypi_latest" ]]; then PYTHON=python${INPUTS_PYTHON} $PYTHON -m pip download jaxlib jax-rocm${JAXCI_ROCM_VERSION}-pjrt jax-rocm${JAXCI_ROCM_VERSION}-plugin --dest $(pwd)/dist/ --no-deps @@ -134,6 +166,7 @@ runs: INPUTS_PYTHON: ${{ inputs.python }} INPUTS_S3_DOWNLOAD_URI: ${{ inputs.s3_download_uri }} INPUTS_ROCM_FULL_VERSION: ${{ inputs.rocm-release-version || inputs.rocm-version }} + INPUTS_JAX_JAXLIB_SOURCE: ${{ inputs.jax-jaxlib-source }} - name: Skip the test run if the wheel artifacts were not downloaded successfully shell: bash if: steps.download-wheel-artifacts.outcome == 'failure' diff --git a/.github/workflows/bazel_rocm.yml b/.github/workflows/bazel_rocm.yml index 04ffffd23ee2..75b6d8dbad89 100644 --- a/.github/workflows/bazel_rocm.yml +++ b/.github/workflows/bazel_rocm.yml @@ -87,6 +87,10 @@ on: description: "S3 URI for ROCm plugin/PJRT wheels (use 'latest' to resolve via LATEST pointer)" default: 'latest' type: string + jax-jaxlib-source: + description: "Where to fetch base jax/jaxlib wheels from: 'gcs' or 's3'" + default: 'gcs' + type: string gate-targets-file: description: "Path to a file listing Bazel targets to run (one per line). If empty, runs the full default test suite." type: string @@ -175,6 +179,7 @@ jobs: jaxlib-version: ${{ inputs.jaxlib-version }} gcs_download_uri: ${{ inputs.gcs_download_uri }} s3_download_uri: ${{ inputs.s3_download_uri }} + jax-jaxlib-source: ${{ inputs.jax-jaxlib-source }} # Halt for testing - name: Wait For Connection uses: google-ml-infra/actions/ci_connection@7f5ca0c263a81ed09ea276524c1b9192f1304e3c diff --git a/.github/workflows/pytest_rocm.yml b/.github/workflows/pytest_rocm.yml index 9737e56ca0cd..8dd0df3e8a6a 100644 --- a/.github/workflows/pytest_rocm.yml +++ b/.github/workflows/pytest_rocm.yml @@ -66,6 +66,10 @@ on: description: "S3 URI for ROCm plugin/PJRT wheels (use 'latest' to resolve via LATEST pointer)" type: string default: 'latest' + jax-jaxlib-source: + description: "Where to fetch base jax/jaxlib wheels from: 'gcs' or 's3'" + type: string + default: 'gcs' halt-for-connection: description: 'Should this workflow run wait for a remote connection?' type: string @@ -108,6 +112,10 @@ on: description: "S3 URI for ROCm plugin/PJRT wheels (use 'latest' to resolve via LATEST pointer)" default: 'latest' type: string + jax-jaxlib-source: + description: "Where to fetch base jax/jaxlib wheels from: 'gcs' or 's3'" + default: 'gcs' + type: string halt-for-connection: description: 'Should this workflow run wait for a remote connection?' type: string @@ -162,6 +170,7 @@ jobs: skip-download-jaxlib-and-plugins-from-gcs: ${{ inputs.skip-download-jaxlib-and-plugins-from-gcs }} gcs_download_uri: ${{ inputs.gcs_download_uri }} s3_download_uri: ${{ inputs.s3_download_uri }} + jax-jaxlib-source: ${{ inputs.jax-jaxlib-source }} - name: Install Python dependencies run: | $JAXCI_PYTHON -m pip install uv~=0.12.3 diff --git a/.github/workflows/wheel_tests_nightly_release.yml b/.github/workflows/wheel_tests_nightly_release.yml index d31bb147bbd6..4de09534277f 100644 --- a/.github/workflows/wheel_tests_nightly_release.yml +++ b/.github/workflows/wheel_tests_nightly_release.yml @@ -201,9 +201,49 @@ jobs: upload_artifacts_to_s3: true s3_upload_uri: 's3://jax-ci-amd/rocm-wheels/${{ github.repository }}/${{ github.ref_name }}/${{ matrix.rocm.release-version }}/wheel-tests-nightly/${{ github.run_number }}/${{ github.run_attempt }}' + # Upstream CI frequently stops publishing jax/jaxlib nightly wheels to GCS. + # On the ROCm fork, build the base jax and jaxlib wheels ourselves on the same + # runners as the ROCm plugin/pjrt wheels and upload them to the same S3 prefix, + # so the ROCm test legs can consume them instead of the (stale) GCS bucket. + build-rocm-jax-jaxlib: + if: github.repository_owner == 'ROCm' + uses: ./.github/workflows/build_rocm_artifacts.yml + permissions: + id-token: write + contents: read + actions: read + strategy: + fail-fast: false + matrix: + runner: ["amd-do-linux.jax.cpu"] + artifact: ["jax", "jaxlib"] + python: ["3.12", "3.13", "3.14"] + rocm: [ + {label: "TheRock 7.14.0", wheel-version: "7", release-version: "7.14.0", manylinux-image: "ghcr.io/rocm/jax-manylinux_2_28-therock-7.14:latest" }, + {label: "TheRock latest", wheel-version: "10", release-version: "therock-latest", manylinux-image: "ghcr.io/rocm/jax-manylinux_2_28-therock-latest:latest"}, + ] + exclude: + # jax is a pure-python (py3-none-any) wheel; build it once per ROCm + # variant rather than once per Python version. + - artifact: "jax" + python: "3.13" + - artifact: "jax" + python: "3.14" + name: "Build ROCm jax/jaxlib (${{ matrix.rocm.label }})" + with: + runner: ${{ matrix.runner }} + artifact: ${{ matrix.artifact }} + python: ${{ matrix.python }} + rocm-version: ${{ matrix.rocm.wheel-version }} + rocm-release-version: ${{ matrix.rocm.release-version }} + manylinux-image: ${{ matrix.rocm.manylinux-image }} + clone_main_xla: 0 + upload_artifacts_to_s3: true + s3_upload_uri: 's3://jax-ci-amd/rocm-wheels/${{ github.repository }}/${{ github.ref_name }}/${{ matrix.rocm.release-version }}/wheel-tests-nightly/${{ github.run_number }}/${{ github.run_attempt }}' + run-pytest-rocm: if: ${{ !cancelled() && (github.repository_owner == 'jax-ml' || github.repository_owner == 'ROCm') }} - needs: [build-rocm-artifacts] + needs: [build-rocm-artifacts, build-rocm-jax-jaxlib] permissions: id-token: write contents: read @@ -229,10 +269,13 @@ jobs: skip-download-jaxlib-and-plugins-from-gcs: ${{inputs.skip-download-jaxlib-and-plugins-from-gcs}} gcs_download_uri: ${{inputs.gcs_download_uri}} s3_download_uri: 's3://jax-ci-amd/rocm-wheels/${{ github.repository }}/${{ github.ref_name }}/${{ matrix.rocm.release-version }}/wheel-tests-nightly/${{ github.run_number }}/${{ github.run_attempt }}' + # On the ROCm fork, use the jax/jaxlib wheels built above (S3) instead of + # the upstream GCS nightly bucket, which is frequently stale/broken. + jax-jaxlib-source: ${{ github.repository_owner == 'ROCm' && 's3' || 'gcs' }} run-bazel-test-rocm: if: ${{ !cancelled() && (github.repository_owner == 'jax-ml' || github.repository_owner == 'ROCm') }} - needs: [build-rocm-artifacts] + needs: [build-rocm-artifacts, build-rocm-jax-jaxlib] permissions: id-token: write contents: read @@ -258,6 +301,9 @@ jobs: gcs_download_uri: ${{inputs.gcs_download_uri}} halt-for-connection: ${{inputs.halt-for-connection}} s3_download_uri: 's3://jax-ci-amd/rocm-wheels/${{ github.repository }}/${{ github.ref_name }}/${{ matrix.rocm.release-version }}/wheel-tests-nightly/${{ github.run_number }}/${{ github.run_attempt }}' + # On the ROCm fork, use the jax/jaxlib wheels built above (S3) instead of + # the upstream GCS nightly bucket, which is frequently stale/broken. + jax-jaxlib-source: ${{ github.repository_owner == 'ROCm' && 's3' || 'gcs' }} build_jaxlib: ${{ 'false' }} build_jax: ${{ 'false' }} jaxlib-version: "head" diff --git a/ci/build_rocm_artifacts.sh b/ci/build_rocm_artifacts.sh index 4bfacf04616a..77097036ce9e 100755 --- a/ci/build_rocm_artifacts.sh +++ b/ci/build_rocm_artifacts.sh @@ -37,7 +37,13 @@ fi # Set up the build environment. source "ci/utilities/setup_build_environment.sh" -allowed_artifacts=("jax-rocm-plugin" "jax-rocm-pjrt") +# The base "jax" and "jaxlib" wheels are also built here on the ROCm fork's +# runners so the ROCm test legs no longer depend on the (frequently broken) +# upstream nightly wheels published to GCS. jax is a pure-python wheel and +# jaxlib is the backend-neutral CPU base (ROCm support lives in the separate +# plugin/pjrt wheels), so both build correctly with the ROCm bazel config used +# below. +allowed_artifacts=("jax-rocm-plugin" "jax-rocm-pjrt" "jax" "jaxlib") if [[ ! " ${allowed_artifacts[*]} " =~ " ${artifact} " ]]; then echo "Error: Invalid artifact: $artifact. Allowed values are: ${allowed_artifacts[*]}" @@ -101,5 +107,8 @@ python build/build.py build --wheels="$artifact" \ $wheel_version_suffix_flag \ "${rocm_path_flags[@]}" -# Verify manylinux compliance. -./ci/utilities/run_auditwheel.sh +# Verify manylinux compliance. The pure-python "jax" wheel has no platform +# wheels to inspect (run_auditwheel.sh errors when it finds none), so skip it. +if [[ "$artifact" != "jax" ]]; then + ./ci/utilities/run_auditwheel.sh +fi From 5b73f8f2e73c6bb3669dd6474232f2d566c4b359 Mon Sep 17 00:00:00 2001 From: magaonka Date: Thu, 13 Aug 2026 14:09:34 -0500 Subject: [PATCH 4/6] [ROCm] Tag ROCm CI job names with N / N+1 release tracks - Label the nightly and continuous ROCm matrices `N (7.14.0)` and `N+1 (therock-latest)` so the release track is identifiable from the job name. - Report the ROCm release version rather than the wheel packaging major in the build, pytest and Bazel job names, and report the GPU count as `gfx950 -GPU`. - Remove the duplicate continuous Bazel ROCm leg: `gfx950.8` ran the same `--config=single_gpu` suite as `gfx950.4` under an identical job name. - Apply the same labels to the fork-only `build-rocm-jax-jaxlib` nightly job. The Bazel job name sits inside a Presubmit Naming Check block, so the presubmit job names change from `ROCM=7` to `ROCM=7.14.0`. The internal github-rocm-presubmits and github-rocm-blocking-presubmits checks need to be updated to match. wheel_tests_rocm_release.yml keeps its `TheRock ` labels: it validates one specific ROCm release, so it is single-track. --- .github/workflows/bazel_rocm.yml | 2 +- .github/workflows/build_rocm_artifacts.yml | 2 +- .github/workflows/pytest_rocm.yml | 6 ++--- .github/workflows/wheel_tests_continuous.yml | 14 +++++------ .../workflows/wheel_tests_nightly_release.yml | 24 +++++++++---------- 5 files changed, 24 insertions(+), 24 deletions(-) diff --git a/.github/workflows/bazel_rocm.yml b/.github/workflows/bazel_rocm.yml index 75b6d8dbad89..ced268a5272f 100644 --- a/.github/workflows/bazel_rocm.yml +++ b/.github/workflows/bazel_rocm.yml @@ -141,7 +141,7 @@ jobs: HSA_NO_SCRATCH_RECLAIM: "1" # Begin Presubmit Naming Check - name modification requires internal check to be updated - name: "linux x86, py${{ inputs.python }}, ROCM=${{ inputs.rocm-version }}, jaxlib=${{ inputs.jaxlib-version }}, x64=${{ inputs.enable-x64 }}, build_jax=${{ inputs.build_jax }}, build_jaxlib=${{ inputs.build_jaxlib }}" + name: "linux x86, py${{ inputs.python }}, ROCM=${{ inputs.rocm-release-version || inputs.rocm-version }}, jaxlib=${{ inputs.jaxlib-version }}, x64=${{ inputs.enable-x64 }}, build_jax=${{ inputs.build_jax }}, build_jaxlib=${{ inputs.build_jaxlib }}" # End Presubmit Naming Check github-rocm-presubmits steps: - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 diff --git a/.github/workflows/build_rocm_artifacts.yml b/.github/workflows/build_rocm_artifacts.yml index b8765259731e..b0214835864b 100644 --- a/.github/workflows/build_rocm_artifacts.yml +++ b/.github/workflows/build_rocm_artifacts.yml @@ -135,7 +135,7 @@ jobs: JAXCI_CLONE_MAIN_XLA: "${{ inputs.clone_main_xla }}" JAXCI_OUTPUT_DIR: "${{ github.workspace }}/dist" - name: "${{ inputs.artifact }}, py ${{ inputs.python }}, ROCm ${{ inputs.rocm-version }}" + name: "${{ inputs.artifact }}, py ${{ inputs.python }}, ROCm ${{ inputs.rocm-release-version || inputs.rocm-version }}" steps: - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 diff --git a/.github/workflows/pytest_rocm.yml b/.github/workflows/pytest_rocm.yml index 8dd0df3e8a6a..8312641b0f03 100644 --- a/.github/workflows/pytest_rocm.yml +++ b/.github/workflows/pytest_rocm.yml @@ -142,9 +142,9 @@ jobs: username: ${{ github.actor }} password: ${{ secrets.GITHUB_TOKEN }} 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, 'gfx950.1') && 'gfx950.1') || - (contains(inputs.runner, 'gfx950.4') && 'gfx950.4') || - (contains(inputs.runner, 'gfx950.8') && 'gfx950.8') }}, ROCm ${{ inputs.rocm-version }}, py${{ inputs.python }}" + name: "${{ (contains(inputs.runner, 'gfx950.1') && 'gfx950 1-GPU') || + (contains(inputs.runner, 'gfx950.4') && 'gfx950 4-GPU') || + (contains(inputs.runner, 'gfx950.8') && 'gfx950 8-GPU') }}, ROCm ${{ inputs.rocm-release-version || inputs.rocm-version }}, py${{ inputs.python }}" env: JAXCI_HERMETIC_PYTHON_VERSION: "${{ inputs.python }}" diff --git a/.github/workflows/wheel_tests_continuous.yml b/.github/workflows/wheel_tests_continuous.yml index 6ffabed8b27f..0201c219c05d 100644 --- a/.github/workflows/wheel_tests_continuous.yml +++ b/.github/workflows/wheel_tests_continuous.yml @@ -504,14 +504,14 @@ jobs: artifact: ["jax-rocm-plugin", "jax-rocm-pjrt"] python: ["3.12", "3.13", "3.14"] rocm: [ - {label: "TheRock 7.14.0", wheel-version: "7", release-version: "7.14.0", manylinux-image: "ghcr.io/rocm/jax-manylinux_2_28-therock-7.14:latest" }, + {label: "N (7.14.0)", wheel-version: "7", release-version: "7.14.0", manylinux-image: "ghcr.io/rocm/jax-manylinux_2_28-therock-7.14:latest" }, ] exclude: - artifact: "jax-rocm-pjrt" python: "3.13" - artifact: "jax-rocm-pjrt" python: "3.14" - name: "Build ROCm artifacts (${{ matrix.rocm.label }})" + name: "Build ROCm artifacts ${{ matrix.rocm.label }}" with: runner: ${{ matrix.runner }} artifact: ${{ matrix.artifact }} @@ -545,9 +545,9 @@ jobs: 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: "TheRock 7.14.0", wheel-version: "7", release-version: "7.14.0", tag: "therock-7.14"}, + {label: "N (7.14.0)", wheel-version: "7", release-version: "7.14.0", tag: "therock-7.14"}, ] - name: "Pytest ROCm (${{ matrix.rocm.label }})" + name: "Pytest ROCm ${{ matrix.rocm.label }}" with: runner: ${{ matrix.runner }} python: ${{ matrix.python }} @@ -577,13 +577,13 @@ jobs: strategy: fail-fast: false matrix: - runner: ["amd-do-linux.jax.gpu.gfx950.4", "amd-do-linux.jax.gpu.gfx950.8"] + runner: ["amd-do-linux.jax.gpu.gfx950.4"] python: ["3.12"] rocm: [ - {label: "TheRock 7.14.0", wheel-version: "7", release-version: "7.14.0", tag: "therock-7.14"}, + {label: "N (7.14.0)", wheel-version: "7", release-version: "7.14.0", tag: "therock-7.14"}, ] enable-x64: [0] - name: "Bazel ROCm (${{ matrix.rocm.label }})" + name: "Bazel ROCm ${{ matrix.rocm.label }}" with: runner: ${{ matrix.runner }} python: ${{ matrix.python }} diff --git a/.github/workflows/wheel_tests_nightly_release.yml b/.github/workflows/wheel_tests_nightly_release.yml index 4de09534277f..eb15f52de335 100644 --- a/.github/workflows/wheel_tests_nightly_release.yml +++ b/.github/workflows/wheel_tests_nightly_release.yml @@ -181,15 +181,15 @@ jobs: artifact: ["jax-rocm-plugin", "jax-rocm-pjrt"] python: ["3.12", "3.13", "3.14"] rocm: [ - {label: "TheRock 7.14.0", wheel-version: "7", release-version: "7.14.0", manylinux-image: "ghcr.io/rocm/jax-manylinux_2_28-therock-7.14:latest" }, - {label: "TheRock latest", wheel-version: "10", release-version: "therock-latest", manylinux-image: "ghcr.io/rocm/jax-manylinux_2_28-therock-latest:latest"}, + {label: "N (7.14.0)", wheel-version: "7", release-version: "7.14.0", manylinux-image: "ghcr.io/rocm/jax-manylinux_2_28-therock-7.14:latest" }, + {label: "N+1 (therock-latest)", wheel-version: "10", release-version: "therock-latest", manylinux-image: "ghcr.io/rocm/jax-manylinux_2_28-therock-latest:latest"}, ] exclude: - artifact: "jax-rocm-pjrt" python: "3.13" - artifact: "jax-rocm-pjrt" python: "3.14" - name: "Build ROCm artifacts (${{ matrix.rocm.label }})" + name: "Build ROCm artifacts ${{ matrix.rocm.label }}" with: runner: ${{ matrix.runner }} artifact: ${{ matrix.artifact }} @@ -219,8 +219,8 @@ jobs: artifact: ["jax", "jaxlib"] python: ["3.12", "3.13", "3.14"] rocm: [ - {label: "TheRock 7.14.0", wheel-version: "7", release-version: "7.14.0", manylinux-image: "ghcr.io/rocm/jax-manylinux_2_28-therock-7.14:latest" }, - {label: "TheRock latest", wheel-version: "10", release-version: "therock-latest", manylinux-image: "ghcr.io/rocm/jax-manylinux_2_28-therock-latest:latest"}, + {label: "N (7.14.0)", wheel-version: "7", release-version: "7.14.0", manylinux-image: "ghcr.io/rocm/jax-manylinux_2_28-therock-7.14:latest" }, + {label: "N+1 (therock-latest)", wheel-version: "10", release-version: "therock-latest", manylinux-image: "ghcr.io/rocm/jax-manylinux_2_28-therock-latest:latest"}, ] exclude: # jax is a pure-python (py3-none-any) wheel; build it once per ROCm @@ -229,7 +229,7 @@ jobs: python: "3.13" - artifact: "jax" python: "3.14" - name: "Build ROCm jax/jaxlib (${{ matrix.rocm.label }})" + name: "Build ROCm jax/jaxlib ${{ matrix.rocm.label }}" with: runner: ${{ matrix.runner }} artifact: ${{ matrix.artifact }} @@ -255,10 +255,10 @@ jobs: 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: "TheRock 7.14.0", wheel-version: "7", release-version: "7.14.0", tag: "therock-7.14"}, - {label: "TheRock latest", wheel-version: "10", release-version: "therock-latest", tag: "therock-latest"}, + {label: "N (7.14.0)", wheel-version: "7", release-version: "7.14.0", tag: "therock-7.14"}, + {label: "N+1 (therock-latest)", wheel-version: "10", release-version: "therock-latest", tag: "therock-latest"}, ] - name: "Pytest ROCm (${{ matrix.rocm.label }})" + name: "Pytest ROCm ${{ matrix.rocm.label }}" with: runner: ${{ matrix.runner }} python: ${{ matrix.python }} @@ -286,11 +286,11 @@ jobs: runner: ["amd-do-linux.jax.gpu.gfx950.1"] python: ["3.12", "3.13", "3.14"] rocm: [ - {label: "TheRock 7.14.0", wheel-version: "7", release-version: "7.14.0", tag: "therock-7.14"}, - {label: "TheRock latest", wheel-version: "10", release-version: "therock-latest", tag: "therock-latest"}, + {label: "N (7.14.0)", wheel-version: "7", release-version: "7.14.0", tag: "therock-7.14"}, + {label: "N+1 (therock-latest)", wheel-version: "10", release-version: "therock-latest", tag: "therock-latest"}, ] enable-x64: [0] - name: "Bazel ROCm (${{ matrix.rocm.label }})" + name: "Bazel ROCm ${{ matrix.rocm.label }}" with: runner: ${{ matrix.runner }} python: ${{ matrix.python }} From 262203f4df5c78d1fd768b4a6d270fdbd2479077 Mon Sep 17 00:00:00 2001 From: Kunreddy Date: Fri, 14 Aug 2026 12:54:53 -0500 Subject: [PATCH 5/6] Mount /home/runner/_work in ROCm workflow containers. Bind-mount the runner workspace into GitHub container: jobs so checkout and artifacts work on ARC runners that use hostPath _work volumes. Co-authored-by: Cursor --- .github/workflows/bazel_rocm.yml | 1 + .github/workflows/pytest_rocm.yml | 2 ++ 2 files changed, 3 insertions(+) diff --git a/.github/workflows/bazel_rocm.yml b/.github/workflows/bazel_rocm.yml index ced268a5272f..7bbee7570489 100644 --- a/.github/workflows/bazel_rocm.yml +++ b/.github/workflows/bazel_rocm.yml @@ -122,6 +122,7 @@ jobs: image: ghcr.io/rocm/jax-base-ubu24.${{ inputs.rocm-tag }}:latest # zizmor: ignore[unpinned-images] volumes: - /data:/data + - /home/runner/_work:/home/runner/_work options: >- --device=/dev/kfd --device=/dev/dri diff --git a/.github/workflows/pytest_rocm.yml b/.github/workflows/pytest_rocm.yml index 8312641b0f03..f3ad083587c6 100644 --- a/.github/workflows/pytest_rocm.yml +++ b/.github/workflows/pytest_rocm.yml @@ -141,6 +141,8 @@ jobs: credentials: username: ${{ github.actor }} password: ${{ secrets.GITHUB_TOKEN }} + volumes: + - /home/runner/_work:/home/runner/_work 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, 'gfx950.1') && 'gfx950 1-GPU') || (contains(inputs.runner, 'gfx950.4') && 'gfx950 4-GPU') || From d90f745462e4522c2f4e4d9e7a78f9925675c2e3 Mon Sep 17 00:00:00 2001 From: Kunreddy Date: Fri, 14 Aug 2026 17:05:54 -0500 Subject: [PATCH 6/6] pytest_rocm: replace JS actions with shell for no-dind container: support Co-authored-by: Cursor --- .github/workflows/pytest_rocm.yml | 30 ++++++++++++++++++++---------- 1 file changed, 20 insertions(+), 10 deletions(-) diff --git a/.github/workflows/pytest_rocm.yml b/.github/workflows/pytest_rocm.yml index f3ad083587c6..c4cedf5cec61 100644 --- a/.github/workflows/pytest_rocm.yml +++ b/.github/workflows/pytest_rocm.yml @@ -159,9 +159,11 @@ jobs: HSA_NO_SCRATCH_RECLAIM: "1" steps: - - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - with: - persist-credentials: false + - name: Checkout + run: | + git config --global --add safe.directory "$GITHUB_WORKSPACE" + git fetch --depth=1 origin ${{ github.sha }} + git checkout ${{ github.sha }} - name: Download JAX ROCm wheels uses: ./.github/actions/download-jax-rocm-wheels with: @@ -180,9 +182,10 @@ jobs: $JAXCI_PYTHON -m pip install pytest-json-report # Halt for testing - name: Wait For Connection - uses: google-ml-infra/actions/ci_connection@7f5ca0c263a81ed09ea276524c1b9192f1304e3c - with: - halt-dispatch-input: ${{ inputs.halt-for-connection }} + run: | + if [[ "${{ inputs.halt-for-connection }}" == "yes" ]]; then + echo "halt-for-connection not supported without dind" && exit 1 + fi - name: Run Pytest ROCm tests timeout-minutes: 120 run: ./ci/run_pytest_rocm.sh @@ -201,10 +204,17 @@ jobs: tar -czf logs.tar.gz logs - name: Configure AWS Credentials if: ${{ !cancelled() }} - uses: aws-actions/configure-aws-credentials@e6de054238d6b7531b4efff3b6587d9aade6a06c # v6.2.3 - with: - role-to-assume: arn:aws:iam::661452401056:role/jax-ci-amd-s3-oidc - aws-region: us-east-1 + run: | + OIDC_TOKEN=$(curl -sH "Authorization: bearer $ACTIONS_ID_TOKEN_REQUEST_TOKEN" \ + "${ACTIONS_ID_TOKEN_REQUEST_URL}&audience=sts.amazonaws.com" | jq -r .value) + CREDS=$(aws sts assume-role-with-web-identity \ + --role-arn arn:aws:iam::661452401056:role/jax-ci-amd-s3-oidc \ + --role-session-name gha \ + --web-identity-token "$OIDC_TOKEN" \ + --region us-east-1) + echo "AWS_ACCESS_KEY_ID=$(echo $CREDS | jq -r .Credentials.AccessKeyId)" >> $GITHUB_ENV + echo "AWS_SECRET_ACCESS_KEY=$(echo $CREDS | jq -r .Credentials.SecretAccessKey)" >> $GITHUB_ENV + echo "AWS_SESSION_TOKEN=$(echo $CREDS | jq -r .Credentials.SessionToken)" >> $GITHUB_ENV - name: Upload test-artifacts to AMD S3 if: ${{ !cancelled() }} env: