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
@@ -115,9 +119,10 @@ 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
+        - /home/runner/_work:/home/runner/_work
       options: >-
         --device=/dev/kfd
         --device=/dev/dri
@@ -137,7 +142,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
@@ -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
@@ -185,8 +191,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}"

.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

.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,10 +141,12 @@ 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') ||
-        (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 }}"
@@ -149,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:
@@ -162,16 +174,18 @@ 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
           $JAXCI_PYTHON -m uv pip install -r build/test-requirements.txt
           $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
@@ -190,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:

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