[Testing do not merge] Add ROCm wheel build and test pipeline to continuous CI

.bazelrc

@@ -347,6 +347,10 @@ common:ci_windows_amd64 --extra_toolchains="@xla//tools/toolchains/win2022/20241
 common:ci_windows_amd64 --host_linkopt=/FORCE:MULTIPLE --linkopt=/FORCE:MULTIPLE
 common:ci_windows_amd64 --color=yes
 
+# Used for rules_python bootstrap 1.8.0+
+build --@rules_python//python/config_settings:bootstrap_impl=script --repo_env=RULES_PYTHON_ENABLE_PIPSTAR=0
+test --@rules_python//python/config_settings:bootstrap_impl=script --repo_env=RULES_PYTHON_ENABLE_PIPSTAR=0
+
 # #############################################################################
 # RBE config options below. These inherit the CI configs above and set the
 # remote execution backend and authentication options required to run builds

.github/workflows/bazel_rocm.yml

@@ -113,11 +113,12 @@ jobs:
       JAXCI_WRITE_TO_BAZEL_REMOTE_CACHE: ${{ inputs.write_to_bazel_remote_cache }}
       JAXCI_BUILD_JAX: ${{ inputs.build_jax }}
       JAXCI_BUILD_JAXLIB: ${{ inputs.build_jaxlib }}
+      JAXCI_CLONE_MAIN_XLA: ${{ inputs.clone_main_xla }}
 # Begin Presubmit Naming Check - name modification requires internal check to be updated
     name: "linux x86, jaxlib=${{ inputs.jaxlib-version }}, ROCM=${{ inputs.rocm-version }}, Python=${{ inputs.python }}, 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@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # v6.0.0
+      - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
         with:
           persist-credentials: false
       - name: ROCm Info

.github/workflows/build_rocm_artifacts.yml

@@ -0,0 +1,180 @@
+# CI - Build ROCm Artifacts
+# This workflow builds ROCm wheels (jax-rocm-plugin, jax-rocm-pjrt) in a ROCm container
+# and uploads them to an S3 bucket. It can be triggered manually via workflow_dispatch or
+# called by other workflows via workflow_call.
+name: CI - Build ROCm Artifacts
+
+on:
+  workflow_dispatch:
+    inputs:
+      runner:
+        description: "Which runner should the workflow run on?"
+        type: choice
+        default: "linux-x86-64-1gpu-amd"
+        options:
+        - "linux-x86-64-1gpu-amd"
+      artifact:
+        description: "Which ROCm artifact to build?"
+        type: choice
+        default: "jax-rocm-plugin"
+        options:
+        - "jax-rocm-plugin"
+        - "jax-rocm-pjrt"
+      python:
+        description: "Which python version should the artifact be built for?"
+        type: choice
+        default: "3.12"
+        options:
+        - "3.11"
+        - "3.12"
+        - "3.13"
+        - "3.14"
+      rocm-version:
+        description: "Which ROCm version to build for?"
+        type: string
+        default: "7"
+      clone_main_xla:
+        description: "Should latest XLA be used?"
+        type: choice
+        default: "1"
+        options:
+        - "1"
+        - "0"
+      halt-for-connection:
+        description: 'Should this workflow run wait for a remote connection?'
+        type: choice
+        default: 'no'
+        options:
+        - 'yes'
+        - 'no'
+  workflow_call:
+    inputs:
+      runner:
+        description: "Which runner should the workflow run on?"
+        type: string
+        default: "linux-x86-64-1gpu-amd"
+      artifact:
+        description: "Which ROCm artifact to build?"
+        type: string
+        default: "jax-rocm-plugin"
+      python:
+        description: "Which python version should the artifact be built for?"
+        type: string
+        default: "3.12"
+      rocm-version:
+        description: "Which ROCm version to build for?"
+        type: string
+        default: "7"
+      clone_main_xla:
+        description: "Should latest XLA be used?"
+        type: string
+        default: "1"
+      upload_artifacts_to_s3:
+        description: "Should the artifacts be uploaded to S3?"
+        default: true
+        type: boolean
+      s3_upload_uri:
+        description: "S3 location prefix to where the artifacts should be uploaded"
+        default: 's3://jax-ci-amd/jax-rocm-plugin-wheels'
+        type: string
+      halt-for-connection:
+        description: 'Should this workflow run wait for a remote connection?'
+        type: string
+        default: 'no'
+    secrets:
+      AWS_ACCESS_KEY_ID:
+        required: true
+      AWS_SECRET_ACCESS_KEY:
+        required: true
+      S3_BUCKET_NAME:
+        required: true
+    outputs:
+      s3_upload_uri:
+        description: "S3 location prefix to where the artifacts were uploaded"
+        value: ${{ jobs.build-artifacts.outputs.s3_upload_uri }}
+
+permissions:
+  id-token: write
+  contents: read
+  actions: read
+
+jobs:
+  build-artifacts:
+    defaults:
+      run:
+        shell: bash
+    runs-on: ${{ inputs.runner }}
+    container:
+      image: "ghcr.io/rocm/jax-manylinux_2_28-rocm-7.2.0:latest"
+      volumes:
+        - /data:/data
+      options: >-
+        --device=/dev/kfd
+        --device=/dev/dri
+        --security-opt seccomp=unconfined
+        --group-add video
+        --shm-size 64G
+
+    env:
+      JAXCI_HERMETIC_PYTHON_VERSION: "${{ inputs.python }}"
+      JAXCI_OUTPUT_DIR: "${{ github.workspace }}/dist"
+
+    name: "${{ inputs.artifact }}, py ${{ inputs.python }}, ROCm ${{ inputs.rocm-version }}"
+
+    outputs:
+      s3_upload_uri: ${{ steps.store-s3-upload-uri.outputs.s3_upload_uri }}
+
+    steps:
+      - uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # v6.0.0
+        with:
+          persist-credentials: false
+      - name: Install Bazelisk
+        run: |
+          curl -fSsL -o /usr/local/bin/bazel https://github.com/bazelbuild/bazelisk/releases/latest/download/bazelisk-linux-amd64
+          chmod +x /usr/local/bin/bazel
+      - name: Install AWS CLI
+        run: |
+          curl -fSsL "https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip" -o /tmp/awscliv2.zip
+          unzip -q /tmp/awscliv2.zip -d /tmp
+          /tmp/aws/install
+          rm -rf /tmp/awscliv2.zip /tmp/aws
+      - name: ROCm Info
+        run: rocminfo
+      # Halt for testing
+      - name: Wait For Connection
+        uses: google-ml-infra/actions/ci_connection@7f5ca0c263a81ed09ea276524c1b9192f1304e3c
+        with:
+          halt-dispatch-input: ${{ inputs.halt-for-connection }}
+      - name: Create dist dir
+        run: mkdir -p $JAXCI_OUTPUT_DIR/
+      - name: Build ${{ inputs.artifact }}
+        timeout-minutes: 120
+        run: |
+          bazel --bazelrc=build/rocm/rocm.bazelrc run \
+                --config=rocm_release_wheel \
+                --config=rocm_rbe \
+                --repo_env=HERMETIC_PYTHON_VERSION="${{ inputs.python }}" \
+                $DEPLOY_TARGET -- $JAXCI_OUTPUT_DIR/
+        env:
+          DEPLOY_TARGET: ${{ inputs.artifact == 'jax-rocm-plugin' && '//jaxlib/tools:deploy_rocm_plugin_wheel' || '//jaxlib/tools:deploy_rocm_pjrt_wheel' }}
+      - name: Configure AWS Credentials
+        if: ${{ inputs.upload_artifacts_to_s3 }}
+        uses: aws-actions/configure-aws-credentials@v4
+        with:
+          role-to-assume: arn:aws:iam::661452401056:role/jax-ci-rocm-jax-s3-oidc
+          aws-region: us-east-1
+      - name: Upload artifacts to S3
+        if: ${{ inputs.upload_artifacts_to_s3 }}
+        run: |
+          echo "Uploading wheels to S3..."
+          ls -lh $JAXCI_OUTPUT_DIR/*.whl
+          aws s3 cp --only-show-errors --recursive $JAXCI_OUTPUT_DIR/ "${INPUTS_S3_UPLOAD_URI}"/
+          echo "Upload complete."
+        env:
+          INPUTS_S3_UPLOAD_URI: ${{ inputs.s3_upload_uri }}
+      - name: Store the S3 upload URI as an output
+        if: ${{ inputs.upload_artifacts_to_s3 }}
+        id: store-s3-upload-uri
+        run: echo "s3_upload_uri=${INPUTS_S3_UPLOAD_URI}" >> "$GITHUB_OUTPUT"
+        env:
+          INPUTS_S3_UPLOAD_URI: ${{ inputs.s3_upload_uri }}

.github/workflows/pytest_tpu.yml

@@ -93,7 +93,11 @@ jobs:
           gcs_download_uri: ${{ inputs.gcs_download_uri }}
       - name: Install Python dependencies
         run: |
-          $JAXCI_PYTHON -m uv pip install -r build/test-requirements.txt -r build/collect-profile-requirements.txt
+          $JAXCI_PYTHON -m uv pip install -r build/test-requirements.txt
+          # xprof depends on cffi, which doesn't support Python 3.13 free-threaded.
+          if [[ "${JAXCI_HERMETIC_PYTHON_VERSION}" != "3.13-nogil" ]]; then
+            $JAXCI_PYTHON -m uv pip install -r build/collect-profile-requirements.txt
+          fi
       - name: Set up libtpu wheels
         run: |
           if [[ "${INPUTS_LIBTPU_VERSION_TYPE}" == "nightly" ]]; then
@@ -124,10 +128,6 @@ jobs:
       - name: Run Pytest TPU tests
         timeout-minutes: ${{ github.event_name == 'pull_request' && 30 || 210 }}
         run: |
-          if [[ ${INPUTS_PYTHON} == "3.13-nogil" ]]; then
-            echo "Uninstalling xprof as it is not compatible with python 3.13t."
-            $JAXCI_PYTHON -m uv pip uninstall xprof
-          fi
           ./ci/run_pytest_tpu.sh
         env:
           INPUTS_PYTHON: ${{ inputs.python }}

.github/workflows/rocm-ci.yml

@@ -46,7 +46,7 @@ jobs:
             dist_docker \
             --image-tag $TEST_IMAGE
       - name: Archive jax wheels
-        uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0
+        uses: actions/upload-artifact@bbbca2ddaa5d8feaa63e36b76fdaad77386f024f # v7.0.0
         with:
           name: rocm_jax_r${{ env.ROCM_VERSION }}_py${{ env.PYTHON_VERSION  }}_id${{ github.run_id }}
           path: ${{ env.WORKSPACE_DIR }}/dist/*.whl

.github/workflows/wheel_tests_continuous.yml

@@ -1,6 +1,6 @@
 # CI - Wheel Tests (Continuous)
 #
-# This workflow builds JAX artifacts and runs CPU/TPU/CUDA tests.
+# This workflow builds JAX artifacts and runs CPU/TPU/CUDA/ROCm tests.
 #
 # It orchestrates the following:
 # 1. build-jaxlib-artifact: Calls the `build_artifacts.yml` workflow to build jaxlib and
@@ -22,15 +22,26 @@
 #                           that was built in the previous step and runs TPU tests.
 # 9. run-bazel-test-tpu:    Calls the `bazel_test_tpu.yml` workflow which
 #                           runs Bazel TPU tests with py_import.
+# 10. build-rocm-artifacts: Calls the `build_rocm_artifacts.yml` workflow to build ROCm plugin/pjrt
+#                           wheels and uploads them to an S3 bucket.
+# 11. run-pytest-rocm:      Calls the `pytest_rocm.yml` workflow which downloads the jaxlib and
+#                           ROCm artifacts and runs the ROCm tests.
+# 12. run-bazel-test-rocm:  Calls the `bazel_rocm.yml` workflow which runs the ROCm Bazel tests.
 
 name: CI - Wheel Tests (Continuous)
 permissions:
-      contents: read
+  id-token: write
+  contents: read
+  actions: read
 
 on:
   schedule:
     - cron: "0 */3 * * *" # Run once every 3 hours
   workflow_dispatch: # allows triggering the workflow run manually
+  pull_request:
+    paths:
+      - '.github/workflows/wheel_tests_continuous.yml'
+      - '.github/workflows/build_rocm_artifacts.yml'
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
@@ -285,4 +296,73 @@ jobs:
       gcs_download_uri: ${{ needs.build-jaxlib-artifact.outputs.gcs_upload_uri }}
       build_jaxlib: "wheel"
       build_jax: "wheel"
-      clone_main_xla: 1
+      clone_main_xla: 1
+
+  build-rocm-artifacts:
+    uses: ./.github/workflows/build_rocm_artifacts.yml
+    secrets: inherit
+    permissions:
+      id-token: write
+      contents: read
+      actions: read
+    strategy:
+        fail-fast: false
+        matrix:
+          runner: ["linux-x86-64-1gpu-amd"]
+          artifact: ["jax-rocm-plugin", "jax-rocm-pjrt"]
+          python: ["3.11"]
+          rocm-version: ["7"]
+    name: "Build ${{ format('{0}', 'ROCm') }} artifacts"
+    with:
+      runner: ${{ matrix.runner }}
+      artifact: ${{ matrix.artifact }}
+      python: ${{ matrix.python }}
+      rocm-version: ${{ matrix.rocm-version }}
+      clone_main_xla: 1
+      upload_artifacts_to_s3: true
+      s3_upload_uri: 's3://jax-ci-amd/rocm-wheels/wheel-tests-continuous/${{ github.run_number }}/${{ github.run_attempt }}'
+
+  run-pytest-rocm:
+    if: ${{ !cancelled() }}
+    needs: [build-jax-artifact, build-jaxlib-artifact, build-rocm-artifacts]
+    uses: ./.github/workflows/pytest_rocm.yml
+    strategy:
+        fail-fast: false
+        matrix:
+          runner: ["linux-x86-64-1gpu-amd", "linux-x86-64-4gpu-amd", "linux-x86-64-8gpu-amd"]
+          python: ["3.11"]
+          rocm: [
+            {version: "7.2.0", tag: "rocm720"},
+          ]
+    name: "Pytest ROCm (JAX artifacts version = ${{ format('{0}', 'head') }})"
+    with:
+      runner: ${{ matrix.runner }}
+      python: ${{ matrix.python }}
+      rocm-version: ${{ matrix.rocm.version }}
+      rocm-tag: ${{ matrix.rocm.tag }}
+      jaxlib-version: "head"
+      gcs_download_uri: ${{ needs.build-jaxlib-artifact.outputs.gcs_upload_uri }}
+
+  run-bazel-test-rocm:
+    if: ${{ !cancelled() }}
+    needs: [build-jax-artifact, build-jaxlib-artifact, build-rocm-artifacts]
+    uses: ./.github/workflows/bazel_rocm.yml
+    strategy:
+        fail-fast: false
+        matrix:
+          runner: ["linux-x86-64-4gpu-amd", "linux-x86-64-8gpu-amd"]
+          python: ["3.11"]
+          rocm-version: ["7"]
+          enable-x64: [0]
+    name: "Bazel ROCm tests (JAX artifacts version = ${{ format('{0}', 'head') }})"
+    with:
+      runner: ${{ matrix.runner }}
+      python: ${{ matrix.python }}
+      rocm-version: ${{ matrix.rocm-version }}
+      enable-x64: ${{ matrix.enable-x64 }}
+      gcs_download_uri: ${{ needs.build-jaxlib-artifact.outputs.gcs_upload_uri }}
+      build_jaxlib: "false"
+      build_jax: "false"
+      jaxlib-version: "head"
+      run_multiaccelerator_tests: "false"
+      clone_main_xla: 1