Implement Mosaic GPU detection and Auto-Skips

.github/workflows/pytest_rocm_abort.yml

@@ -0,0 +1,162 @@
+# CI - Pytest ROCm (Abort Support)
+#
+# This workflow runs the ROCm tests with Pytest in ROCm GHCR containers,
+# using the ROCm `pytest-abort` retry wrapper to detect/retry aborts/crashes.
+#
+# It can be triggered manually via workflow_dispatch or called by other workflows
+# via workflow_call.
+#
+# It consists of the following job:
+# run-tests:
+#    - Runs in ROCm container (ghcr.io/rocm/jax-base-ubu24-rocm*:latest)
+#    - Downloads the JAX and jaxlib wheels from GCS, and ROCm plugins from latest release.
+#    - Executes the `run_pytest_rocm_abort.sh` script, which installs wheel artifacts and
+#      runs the ROCm tests with Pytest under `pytest-abort-retry`.
+name: CI - Pytest ROCm (Abort Support)
+
+on:
+  workflow_dispatch:
+    inputs:
+      runner:
+        description: "Which runner should the workflow run on?"
+        type: choice
+        default: "linux-x86-64-4gpu-amd"
+        options:
+          - "linux-x86-64-1gpu-amd"
+          - "linux-x86-64-4gpu-amd"
+          - "linux-x86-64-8gpu-amd"
+      python:
+        description: "Which Python version to use?"
+        type: choice
+        default: "3.11"
+        options:
+          - "3.11"
+          - "3.12"
+      rocm-version:
+        description: "Which ROCm version to test?"
+        type: choice
+        default: "7.2.0"
+        options:
+          - "7.2.0"
+      rocm-tag:
+        description: "ROCm tag for container image (e.g., rocm720)"
+        type: string
+        default: "rocm720"
+      jaxlib-version:
+        description: "Which jaxlib version to use? (head/pypi_latest)"
+        type: choice
+        default: "head"
+        options:
+          - "head"
+          - "pypi_latest"
+      skip-download-jaxlib-and-plugins-from-gcs:
+        description: "Whether to skip downloading the jaxlib and plugins from GCS (e.g for testing a jax only release)"
+        type: choice
+        default: '0'
+        options:
+          - '0'
+          - '1'
+      gcs_download_uri:
+        description: "GCS location prefix from where the artifacts should be downloaded"
+        type: string
+        default: 'gs://jax-nightly-artifacts/latest'
+      halt-for-connection:
+        description: 'Should this workflow run wait for a remote connection?'
+        type: string
+        default: 'no'
+  workflow_call:
+    inputs:
+      runner:
+        description: "Which runner should the workflow run on?"
+        type: string
+        default: "linux-x86-64-4gpu-amd"
+      python:
+        description: "Which Python version to use?"
+        type: string
+        default: "3.11"
+      rocm-version:
+        description: "Which ROCm version to test?"
+        type: string
+        default: "7.2.0"
+      rocm-tag:
+        description: "ROCm tag for container image (e.g., rocm720)"
+        type: string
+        default: "rocm720"
+      jaxlib-version:
+        description: "Which jaxlib version to use? (head/pypi_latest)"
+        type: string
+        default: "head"
+      skip-download-jaxlib-and-plugins-from-gcs:
+        description: "Whether to skip downloading the jaxlib and plugins from GCS (e.g for testing a jax only release)"
+        default: '0'
+        type: string
+      gcs_download_uri:
+        description: "GCS location prefix from where the artifacts should be downloaded"
+        default: 'gs://jax-nightly-artifacts/latest'
+        type: string
+
+permissions: {}
+
+env:
+  UV_DEFAULT_INDEX: "https://us-python.pkg.dev/ml-oss-artifacts-published/pypi-mirror/simple"
+
+jobs:
+  run-tests:
+    defaults:
+      run:
+        # Set the shell to bash as GitHub actions run with /bin/sh by default
+        shell: bash
+    runs-on: ${{ inputs.runner }}
+    continue-on-error: true
+    # Run in ROCm GHCR container with GPU access
+    container:
+      image: ghcr.io/rocm/jax-base-ubu24.${{ inputs.rocm-tag }}:latest
+      credentials:
+        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, '1gpu') && '1gpu') ||
+        (contains(inputs.runner, '4gpu') && '4gpu') ||
+        (contains(inputs.runner, '8gpu') && '8gpu') }}, ROCm ${{ inputs.rocm-version }}, py${{ inputs.python }}"
+
+    env:
+      JAXCI_HERMETIC_PYTHON_VERSION: "${{ inputs.python }}"
+      JAXCI_PYTHON: "python${{ inputs.python }}"
+      JAXCI_ENABLE_X64: "0"
+
+    steps:
+      - uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # v6.0.0
+        with:
+          persist-credentials: false
+      - name: Download JAX ROCm wheels
+        uses: ./.github/actions/download-jax-rocm-wheels
+        with:
+          python: ${{ inputs.python }}
+          rocm-version: ${{ inputs.rocm-version }}
+          jaxlib-version: ${{ inputs.jaxlib-version }}
+          skip-download-jaxlib-and-plugins-from-gcs: ${{ inputs.skip-download-jaxlib-and-plugins-from-gcs }}
+          gcs_download_uri: ${{ inputs.gcs_download_uri }}
+        env:
+          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+      - name: Install Python dependencies
+        run: |
+          $JAXCI_PYTHON -m pip install uv~=0.5.30
+          $JAXCI_PYTHON -m uv pip install -r build/test-requirements.txt
+      # Halt for testing
+      - name: Wait For Connection
+        uses: google-ml-infra/actions/ci_connection@7f5ca0c263a81ed09ea276524c1b9192f1304e3c
+        with:
+          halt-dispatch-input: ${{ inputs.halt-for-connection }}
+      - name: Run Pytest ROCm tests (abort support)
+        timeout-minutes: 180
+        run: ./ci/run_pytest_rocm_abort.sh
+      - name: Upload pytest results to artifact
+        if: always()
+        uses: actions/upload-artifact@v4
+        with:
+          name: logs_abort
+          path: |
+            logs_abort/
+          if-no-files-found: warn
+          retention-days: 2
+          overwrite: true

.github/workflows/wheel_tests_nightly_release_abort.yml

@@ -0,0 +1,53 @@
+# CI - Wheel Tests (Nightly/Release) (ROCm abort only)
+#
+# This workflow runs only the ROCm wheel tests using the abort/retry wrapper workflow.
+name: CI - Wheel Tests (Nightly/Release) (ROCm abort only)
+
+on:
+  workflow_dispatch:
+    inputs:
+      gcs_download_uri:
+        description: "GCS location URI from where the artifacts should be downloaded"
+        required: true
+        default: 'gs://jax-nightly-artifacts/latest'
+        type: string
+      skip-download-jaxlib-and-plugins-from-gcs:
+        description: "Whether to download only the jax wheel from GCS (e.g for testing a jax only release)"
+        required: true
+        default: '0'
+        type: string
+      halt-for-connection:
+        description: 'Should this workflow run wait for a remote connection? (yes/no)'
+        required: false
+        default: 'no'
+        type: string
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
+  cancel-in-progress: true
+permissions: {}
+
+env:
+  UV_DEFAULT_INDEX: "https://us-python.pkg.dev/ml-oss-artifacts-published/pypi-mirror/simple"
+
+jobs:
+  run-pytest-rocm:
+    uses: ./.github/workflows/pytest_rocm_abort.yml
+    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"]
+          python: ["3.11", "3.12", "3.13", "3.14"]
+          rocm: [
+            {version: "7.2.0", tag: "rocm720"},
+          ]
+    name: "Pytest ROCm abort (JAX artifacts version = ${{ startsWith(github.ref_name, 'release/') && 'latest release' || 'nightly' }})"
+    with:
+      runner: ${{ matrix.runner }}
+      python: ${{ matrix.python }}
+      rocm-version: ${{ matrix.rocm.version }}
+      rocm-tag: ${{ matrix.rocm.tag }}
+      jaxlib-version: "head"
+      skip-download-jaxlib-and-plugins-from-gcs: ${{inputs.skip-download-jaxlib-and-plugins-from-gcs}}
+      gcs_download_uri: ${{inputs.gcs_download_uri}}
+      halt-for-connection: ${{inputs.halt-for-connection}}

.gitignore

@@ -31,3 +31,7 @@ jax.iml
 /include/
 /lib/
 /share/
+
+/compile_commands.json
+/strace.txt
+/external

build/build.py

@@ -638,6 +638,10 @@ async def main():
       )
 
   if "rocm" in args.wheels:
+    if not args.configure_only:
+      print("ERROR: This repo is not used for building the ROCm JAX plugins. Please use the new plugin repo: https://github.com/ROCm/rocm-jax")
+      exit(1)
+
     wheel_build_command_base.append("--config=rocm_base")
     wheel_build_command_base.append("--config=rocm")
     if clang_local:

build/rocm-test-requirements.txt

@@ -0,0 +1,24 @@
+absl-py
+build
+cloudpickle
+colorama>=0.4.4
+filelock
+flatbuffers
+hypothesis
+mpmath>=1.3
+pillow>=10.4.0
+# TODO(kanglan): Remove once psutil from portpicker supports python 3.13t
+portpicker; python_version<"3.13"
+pytest-xdist
+pytest-json-report
+pytest-html
+pytest-csv
+pytest-rerunfailures
+pytest-html-merger
+pytest-reportlog
+wheel
+rich
+setuptools
+matplotlib
+opt-einsum
+auditwheel