Backport: Mount /home/runner/_work in ROCm workflow containers

.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'

.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
@@ -118,6 +122,7 @@ jobs:
       image: ghcr.io/rocm/jax-dev-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
@@ -175,6 +180,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

.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
@@ -133,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') ||
         (contains(inputs.runner, 'gfx950.4') && 'gfx950.4') ||
@@ -162,6 +172,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.11.2

.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"

build/rocm/rocm.bazelrc

@@ -10,7 +10,7 @@ common:rocm --action_env=TF_ROCM_CLANG="1"
 common:rocm --action_env=HIPCC_COMPILE_FLAGS_APPEND=--offload-compress
 common:rocm --copt=-Wno-gnu-offsetof-extensions
 common:rocm --copt=-Qunused-arguments
-common:rocm --repo_env=TF_ROCM_AMDGPU_TARGETS="gfx908,gfx90a,gfx942,gfx950,gfx1030,gfx1100,gfx1101,gfx1200,gfx1201"
+common:rocm --repo_env=TF_ROCM_AMDGPU_TARGETS="gfx900,gfx906,gfx908,gfx90a,gfx942,gfx950,gfx1010,gfx1011,gfx1012,gfx1030,gfx1031,gfx1032,gfx1033,gfx1034,gfx1035,gfx1036,gfx1100,gfx1101,gfx1102,gfx1103,gfx1150,gfx1151,gfx1152,gfx1153,gfx1200,gfx1201"
 # Used for @xla//build_tools/rocm:parallel_gpu_execute
 common:rocm --legacy_external_runfiles=true
 

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

third_party/xla/revision.bzl

@@ -17,9 +17,9 @@
 # To update XLA to a new revision,
 # a) update XLA_COMMIT to the new git commit hash
 # b) get the sha256 hash of the commit by running:
-#    curl -L https://github.com/openxla/xla/archive/{git_hash}.tar.gz | sha256sum
+#    curl -L https://github.com/ROCm/xla/archive/{git_hash}.tar.gz | sha256sum
 #    and update XLA_SHA256 with the result.
 
 # buildifier: disable=module-docstring
-XLA_COMMIT = "c54647791964d0eb2062ea2f293bd199c48045b2"
-XLA_SHA256 = "7d08c2cb2267e7be25314682404662867ae474d8c1498f3eea44314bf83e497f"
+XLA_COMMIT = "9df7d1a5b395aa557be11a673861fd7a08f75d44"
+XLA_SHA256 = "c291f639021bf4d304f6d4f39d0ba103bd848034f5fe444c8c604f4468ce51b2"

third_party/xla/workspace.bzl

@@ -21,7 +21,7 @@ def repo():
         name = "xla",
         sha256 = XLA_SHA256,
         strip_prefix = "xla-{commit}".format(commit = XLA_COMMIT),
-        urls = tf_mirror_urls("https://github.com/openxla/xla/archive/{commit}.tar.gz".format(commit = XLA_COMMIT)),
+        urls = tf_mirror_urls("https://github.com/ROCm/xla/archive/{commit}.tar.gz".format(commit = XLA_COMMIT)),
     )
 
     # For development, one often wants to make changes to the TF repository as well