V0.9.0 merger

.github/workflows/build-base-docker.yml

@@ -10,14 +10,22 @@ on:
       - 'build/ci_build'
       - '.github/workflows/build-base-docker.yml'
       - 'docker/Dockerfile*'
+      - 'docker/manylinux/*'
+  push:
+    branches:
+      - 'master'
+    paths:
+      - 'build/ci_build'
+      - '.github/workflows/build-base-docker.yml'
+      - 'docker/Dockerfile*'
+      - 'docker/manylinux/*'
 
 jobs:
   build-base-images:
     runs-on: ${{ matrix.runner-label }}
     strategy:
       fail-fast: false
       matrix:
-        ubuntu-version: ["22", "24"]
         rocm-version: ["7.1.1", "7.2.0"]
         install-llvm: [true, false]
         include:
@@ -64,7 +72,7 @@ jobs:
             --rocm-version="${{ matrix.rocm-version }}" \
             $BUILD_ARGS \
             build_base_dockers \
-            --filter="ubu${{ matrix.ubuntu-version }}" \
+            --filter="ubu24" \
             ${{ matrix.install-llvm && '--install-llvm --llvm-version 18' || '' }}
       - name: Authenticate to GitHub Container Registry
         run: |
@@ -73,16 +81,15 @@ jobs:
       - name: Push docker images
         env:
           ROCM_VERSION: ${{ matrix.rocm-version }}
-          UBUNTU_VERSION: ${{ matrix.ubuntu-version }}
           INSTALL_LLVM: ${{ matrix.install-llvm }}
         run: |
           # Construct image tag based on matrix values
           # ROCm version tag removes dots (7.1.1 -> 711, 7.2.0 -> 720)
           rocm_tag="rocm${ROCM_VERSION//.}"
           if [ "$INSTALL_LLVM" = "true" ]; then
-            image_tag="jax-dev-ubu${UBUNTU_VERSION}.${rocm_tag}"
+            image_tag="jax-dev-ubu24.${rocm_tag}"
           else
-            image_tag="jax-base-ubu${UBUNTU_VERSION}.${rocm_tag}"
+            image_tag="jax-base-ubu24.${rocm_tag}"
           fi
 
           # Push with commit SHA tag
@@ -98,3 +105,68 @@ jobs:
             docker tag "${image_tag}" "${ghcr_image_latest}"
             docker push "${ghcr_image_latest}"
           fi
+
+  build-manylinux-builder-images:
+    runs-on: ${{ matrix.runner-label }}
+    strategy:
+      fail-fast: false
+      matrix:
+        rocm-version: ["7.1.1", "7.2.0"]
+        include:
+          - rocm-version: "7.1.1"
+            runner-label: "linux-x86-64-1gpu-amd"
+          - rocm-version: "7.2.0"
+            runner-label: "linux-x86-64-1gpu-amd"
+    steps:
+      - name: Clean up old runs
+        run: |
+          ls -lah
+          # Make sure that we own all of the files so that we have permissions to delete them
+          docker run --rm -v "./:/rocm-jax" ubuntu \
+            /bin/bash -c "shopt -s dotglob; chown -R $UID /rocm-jax/* || true"
+          # Remove any old work directories from this machine
+          rm -rf * || true
+          ls -lah
+          # Clean up any docker stuff that's more than a week old
+          docker system prune -a --filter "until=168h"
+          # Stop any containers running for more than 12 hours. No CI job should take this long.
+          docker ps --format="{{.RunningFor}} {{.Names}}" | grep hours \
+            | awk -F: '{if($1>12)print$1}' | awk ' {print $4} ' | xargs docker stop || true
+      - uses: actions/checkout@v4
+      - name: Build docker images
+        run: |
+          python3 build/ci_build \
+            --rocm-version="${{ matrix.rocm-version }}" \
+            --rocm-build-job="${{ matrix.rocm-build-job }}" \
+            --rocm-build-num="${{ matrix.rocm-build-num }}" \
+            build_manylinux_dockers
+      - name: Login to GitHub Container Registry
+        uses: docker/login-action@v3
+        with:
+          registry: ghcr.io
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+      - name: Push docker images
+        env:
+          ROCM_VERSION: ${{ matrix.rocm-version }}
+          ROCM_BUILD_JOB: >-
+            ${{ matrix.rocm-build-job && format('-{0}', inputs.rocm-build-job) || '' }}
+          ROCM_BUILD_NUM: >-
+            ${{ matrix.rocm-build-num && format('-{0}', inputs.rocm-build-num) || '' }}
+        run: |
+          image_tag="ghcr.io/rocm/jax-manylinux_2_28-rocm-${ROCM_VERSION}${ROCM_BUILD_JOB}${ROCM_BUILD_NUM}"
+
+          # Push with commit SHA tag
+          sha_image_tag="${image_tag}:${GITHUB_SHA}"
+          echo "Image name (SHA): ${sha_image_tag}"
+          docker tag "${image_tag}" "${sha_image_tag}"
+          docker push "${sha_image_tag}"
+
+          # Push with latest tag (only for schedule and workflow_dispatch, not PRs)
+          if [ "${{ github.event_name }}" != "pull_request" ]; then
+            latest_image_tag="${image_tag}:latest"
+            echo "Image Name (latest): ${latest_image_tag}"
+            docker tag "${image_tag}" "${latest_image_tag}"
+            docker push "${latest_image_tag}"
+          fi
+

.github/workflows/build-docker.yml

@@ -73,26 +73,18 @@ jobs:
           EXTRA_CR_TAG: ${{ inputs.extra-cr-tag }}
           ROCM_VERSION: ${{ inputs.rocm-version }}
         run: |
-          ubu22_img="ghcr.io/rocm/jax-ubu22.rocm${ROCM_VERSION//.}:${GITHUB_SHA}"
           ubu24_img="ghcr.io/rocm/jax-ubu24.rocm${ROCM_VERSION//.}:${GITHUB_SHA}"
-          echo "Ubuntu 22 image name: ${ubu22_img}"
           echo "Ubuntu 24 image name: ${ubu24_img}"
-          docker tag "jax-ubu22.rocm${ROCM_VERSION//.}" "${ubu22_img}"
           docker tag "jax-ubu24.rocm${ROCM_VERSION//.}" "${ubu24_img}"
-          docker push "${ubu22_img}"
           docker push "${ubu24_img}"
       - name: Push extra tags
         if: ${{ inputs.extra-cr-tag }}
         env:
           EXTRA_CR_TAG: ${{ inputs.extra-cr-tag }}
           ROCM_VERSION: ${{ inputs.rocm-version }}
         run: |
-          ubu22_img="ghcr.io/rocm/jax-ubu22.rocm${ROCM_VERSION//.}:${EXTRA_CR_TAG}"
           ubu24_img="ghcr.io/rocm/jax-ubu24.rocm${ROCM_VERSION//.}:${EXTRA_CR_TAG}"
-          echo "Ubuntu 22 image name: ${ubu22_img}"
           echo "Ubuntu 24 image name: ${ubu24_img}"
-          docker tag "jax-ubu22.rocm${ROCM_VERSION//.}" "${ubu22_img}"
           docker tag "jax-ubu24.rocm${ROCM_VERSION//.}" "${ubu24_img}"
-          docker push "${ubu22_img}"
           docker push "${ubu24_img}"
 

.github/workflows/build-wheels.yml

@@ -32,6 +32,10 @@ on:
         required: false
         type: string
         default: 'rocm-jaxlib-v0.9.0'
+      builder-image:
+        required: false
+        type: string
+        default: ''
     secrets:
       rbe_ci_key:
         required: true
@@ -71,6 +75,10 @@ jobs:
           repository: ${{ inputs.jax-repo }}
           ref: ${{ inputs.jax-ref }}
           path: jax
+      - name: Authenticate to GitHub Container Registry
+        run: |
+          echo "${{ secrets.GITHUB_TOKEN }}" \
+            | docker login ghcr.io -u ${{ github.actor }} --password-stdin
       - name: Get RBE cluster keys
         env:
           RBE_CI_CERT: ${{ secrets.rbe_ci_cert }}
@@ -88,7 +96,8 @@ jobs:
             --rocm-build-num="${{ inputs.rocm-build-num }}" \
             --jax-source-dir="./jax" \
             dist_wheels \
-            --rbe
+            --rbe \
+            --builder-image="${{ inputs.builder-image }}"
       - name: Archive plugin wheels
         uses: actions/upload-artifact@v4
         with:

.github/workflows/ci-ut.yml

@@ -0,0 +1,96 @@
+name: CI Unit Tests
+
+on:
+  push:
+    branches:
+      - master
+      - 'rocm-jaxlib-v*'
+  pull_request:
+    branches:
+      - master
+      - 'rocm-jaxlib-v*'
+  workflow_dispatch:
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+permissions:
+  contents: read
+
+jobs:
+  build-and-test:
+    name: build-and-test (${{ matrix.mode.name }})
+    runs-on: linux-x86-64-4gpu-amd-gfx942
+    strategy:
+      fail-fast: false
+      matrix:
+        mode:
+          - {name: "py3.11", python_version: "3.11", config: ""}
+          - {name: "py3.12", python_version: "3.12", config: ""}
+          - {name: "py3.13", python_version: "3.13", config: ""}
+          - {name: "py3.14", python_version: "3.14", config: ""}
+          - {name: "asan", python_version: "3.11", config: "--config=asan"}
+    container:
+      # note this image shall match the one defined in platform/linux:tf_linux_gpu
+      image: rocm/tensorflow-build@sha256:7fcfbd36b7ac8f6b0805b37c4248e929e31cf5ee3af766c8409dd70d5ab65faa
+      options: >-
+        -w ${{ github.workspace }}/jax_rocm_plugin
+        --device=/dev/kfd
+        --device=/dev/dri
+        --group-add video
+        --cap-add=SYS_PTRACE
+        --security-opt seccomp=unconfined
+        --shm-size 16G
+    defaults:
+      run:
+        working-directory: jax_rocm_plugin
+    steps:
+      - name: Checkout plugin repo
+        uses: actions/checkout@v4
+
+      - name: Get RBE cluster keys
+        env:
+          RBE_CI_CERT: ${{ secrets.RBE_CI_CERT }}
+          RBE_CI_KEY: ${{ secrets.RBE_CI_KEY }}
+        run: |
+          echo "$RBE_CI_CERT" >> ci-cert.crt
+          echo "$RBE_CI_KEY" >> ci-cert.key
+
+      - name: Run single-GPU unit tests
+        if: always()
+        run: |
+          bash build/rocm/ci_run_jax_ut.sh \
+            --config=rocm_sgpu \
+            --config=rocm_rbe \
+            --repo_env=HERMETIC_PYTHON_VERSION=${{ matrix.mode.python_version }} \
+            ${{ matrix.mode.config }} \
+            --curses=no \
+            --color=yes \
+            -- \
+            @jax//tests:gpu_tests \
+            @jax//tests:backend_independent_tests \
+            $(build/rocm/targets_to_ignore.sh)
+
+      - name: Run multi-GPU unit tests
+        if: always()
+        run: |
+          bash build/rocm/ci_run_jax_ut.sh \
+            --config=rocm_mgpu \
+            --config=rocm_rbe \
+            --repo_env=HERMETIC_PYTHON_VERSION=${{ matrix.mode.python_version }} \
+            ${{ matrix.mode.config }} \
+            --curses=no \
+            --color=yes \
+            --strategy=TestRunner=local \
+            -- \
+            @jax//tests:gpu_tests \
+            @jax//tests:backend_independent_tests \
+            $(build/rocm/targets_to_ignore.sh)
+
+      - name: Upload logs to artifact
+        if: always()
+        uses: actions/upload-artifact@v4
+        with:
+          name: logs-rbe-py${{ matrix.mode.name }}
+          path: jax_rocm_plugin/logs/

.github/workflows/ci.yml

@@ -22,11 +22,10 @@ jobs:
         rocm-version: ["7.1.1"]
     uses: ./.github/workflows/build-wheels.yml
     with:
-      # TODO: Add back Python 3.13 when we're ready to move to a more recent version of XLA. 3.13
-      #       fails with a complaint abou the pipes module.
-      python-versions: "3.11,3.12"
+      python-versions: "3.11,3.12,3.13,3.14"
       rocm-version: ${{ matrix.rocm-version }}
       runner-label: '["linux-x86-64-1gpu-amd"]'
+      builder-image: "search"
     secrets:
       rbe_ci_cert: ${{ secrets.RBE_CI_CERT }}
       rbe_ci_key: ${{ secrets.RBE_CI_KEY }}
@@ -86,7 +85,7 @@ jobs:
           ROCM_VERSION: ${{ matrix.rocm-version }}
         run: |
           docker run --rm --device=/dev/kfd --device=/dev/dri --group-add video \
-            "ghcr.io/rocm/jax-ubu22.rocm${ROCM_VERSION//.}:${GITHUB_SHA}" \
+            "ghcr.io/rocm/jax-ubu24.rocm${ROCM_VERSION//.}:${GITHUB_SHA}" \
             rocm-smi -a || true
       - name: Download wheel artifacts
         uses: actions/download-artifact@v4
@@ -105,6 +104,6 @@ jobs:
         # TODO: Add the tests/linalg_test.py test back once we fix the XLAClient thing.
         run: |
           python3 build/ci_build test \
-            "ghcr.io/rocm/jax-ubu22.rocm${ROCM_VERSION//.}:${GITHUB_SHA}" \
+            "ghcr.io/rocm/jax-ubu24.rocm${ROCM_VERSION//.}:${GITHUB_SHA}" \
             --test-cmd "bash ci/jax_rbe/pr_setup.sh && ci/jax_rbe/pr_test.sh 0.9.0 3.12"