Merge branch 'rocm-main' into ci-upstream-sync-110_1

charleshofer · charleshofer · commit 31c1f25425cd · 2025-02-10T18:02:25.000Z
diff --git a/.github/workflows/ci-build.yaml b/.github/workflows/ci-build.yaml
@@ -1,4 +1,4 @@
-name: CI
+name: ROCm CPU CI
 
 # We test all supported Python versions as follows:
 # - 3.10 : Documentation build
@@ -11,10 +11,10 @@ on:
   # but only for the main branch
   push:
     branches:
-      - main
+      - rocm-main
   pull_request:
     branches:
-      - main
+      - rocm-main
 
 permissions:
   contents: read  # to fetch code
@@ -42,12 +42,8 @@ jobs:
       - run: pre-commit run --show-diff-on-failure --color=always --all-files
 
   build:
-    # Don't execute in fork due to runner type
-    if: github.repository == 'jax-ml/jax'
     name: "build ${{ matrix.name-prefix }} (py ${{ matrix.python-version }} on ubuntu-20.04, x64=${{ matrix.enable-x64}})"
-    runs-on: linux-x86-n2-32
-    container:
-      image: index.docker.io/library/ubuntu@sha256:6d8d9799fe6ab3221965efac00b4c34a2bcc102c086a58dff9e19a08b913c7ef # ratchet:ubuntu:20.04
+    runs-on: ROCM-Ubuntu
     timeout-minutes: 60
     strategy:
       matrix:
@@ -65,10 +61,6 @@ jobs:
             num_generated_cases: 1
     steps:
     - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
-    - name: Image Setup
-      run: |
-        apt update
-        apt install -y libssl-dev
     - name: Set up Python ${{ matrix.python-version }}
       uses: actions/setup-python@42375524e23c412d93fb67b49958b491fce71c38  # v5.4.0
       with:
@@ -95,12 +87,12 @@ jobs:
         echo "JAX_THREEFRY_PARTITIONABLE=$JAX_THREEFRY_PARTITIONABLE"
         echo "JAX_ENABLE_CHECKS=$JAX_ENABLE_CHECKS"
         echo "JAX_SKIP_SLOW_TESTS=$JAX_SKIP_SLOW_TESTS"
-        pytest -n auto --tb=short --maxfail=20 tests examples
+        pytest -n 4 --tb=short --maxfail=20 tests examples
 
 
   documentation:
     name: Documentation - test code snippets
-    runs-on: ubuntu-latest
+    runs-on: ROCM-Ubuntu
     timeout-minutes: 10
     strategy:
       matrix:
@@ -128,19 +120,13 @@ jobs:
 
   documentation_render:
     name: Documentation - render documentation
-    runs-on: linux-x86-n2-16
-    container:
-      image: index.docker.io/library/ubuntu@sha256:6d8d9799fe6ab3221965efac00b4c34a2bcc102c086a58dff9e19a08b913c7ef # ratchet:ubuntu:20.04
-    timeout-minutes: 10
+    runs-on: ubuntu-latest
+    timeout-minutes: 20
     strategy:
       matrix:
         python-version: ['3.10']
     steps:
     - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
-    - name: Image Setup
-      run: |
-        apt update
-        apt install -y libssl-dev libsqlite3-dev
     - name: Set up Python ${{ matrix.python-version }}
       uses: actions/setup-python@42375524e23c412d93fb67b49958b491fce71c38  # v5.4.0
       with:
@@ -193,9 +179,7 @@ jobs:
 
   ffi:
     name: FFI example
-    runs-on: linux-x86-g2-16-l4-1gpu
-    container:
-      image: index.docker.io/tensorflow/build:latest-python3.12@sha256:48e99608fe9434ada5b14e19fdfd8e64f4cfc83aacd328b9c2101b210e984295 # ratchet:index.docker.io/tensorflow/build:latest-python3.12
+    runs-on: ROCM-Ubuntu
     timeout-minutes: 30
     steps:
     - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
@@ -206,7 +190,7 @@ jobs:
     - name: Install JAX
       run: |
         pip install uv
-        uv pip install --system .[cuda12]
+        uv pip install --system .
     - name: Build and install example project
       run: uv pip install --system ./examples/ffi[test]
       env:
@@ -215,10 +199,11 @@ jobs:
         # a different toolchain. GCC is the default compiler on the
         # 'ubuntu-latest' runner, but we still set this explicitly just to be
         # clear.
-        CMAKE_ARGS: -DCMAKE_CXX_COMPILER=g++ -DJAX_FFI_EXAMPLE_ENABLE_CUDA=ON
+        CMAKE_ARGS: -DCMAKE_CXX_COMPILER=g++ #-DJAX_FFI_EXAMPLE_ENABLE_CUDA=ON
     - name: Run CPU tests
       run: python -m pytest examples/ffi/tests
       env:
         JAX_PLATFORM_NAME: cpu
     - name: Run GPU tests
       run: python -m pytest examples/ffi/tests
+
diff --git a/.github/workflows/rocm-ci.yml b/.github/workflows/rocm-ci.yml
@@ -1,23 +1,25 @@
-name: ROCm GPU Post-Merge Check
+name: ROCm GPU CI
 
 on:
-  # Trigger the workflow after a push into the main branch
+  # Trigger the workflow on push or pull request,
+  # but only for the rocm-main branch
   push:
     branches:
-      - main
-
-permissions:
-  contents: read
+      - rocm-main
+  pull_request:
+    branches:
+      - rocm-main
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
+  cancel-in-progress: true
 
 jobs:
-  build-jax-in-docker:
-    runs-on: linux-x86_64-cirrascale-64-8gpu-amd-mi250
+  build-jax-in-docker: # strategy and matrix come here
+    runs-on:  mi-250
     env:
       BASE_IMAGE: "ubuntu:22.04"
-      TEST_IMAGE: ubuntu-jax-upstream-${{ github.run_id }}_${{ github.run_number }}_${{ github.run_attempt }}
+      TEST_IMAGE: ubuntu-jax-${{ github.run_id }}_${{ github.run_number }}_${{ github.run_attempt }}
       PYTHON_VERSION: "3.10"
       ROCM_VERSION: "6.2.4"
       WORKSPACE_DIR: workdir_${{ github.run_id }}_${{ github.run_number }}_${{ github.run_attempt }}
@@ -32,6 +34,9 @@ jobs:
           ls
       - name: Print system info
         run: |
+          whoami
+          printenv
+          df -h
           rocm-smi
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
         with:
@@ -50,10 +55,9 @@ jobs:
         uses: actions/upload-artifact@v4
         with:
           name: rocm_jax_r${{ env.ROCM_VERSION }}_py${{ env.PYTHON_VERSION  }}_id${{ github.run_id }}
-          path: ${{ env.WORKSPACE_DIR }}/dist/*.whl
-          retention-days: 2
+          path: ./dist/*.whl
       - name: Run tests
         run: |
           cd $WORKSPACE_DIR
-          python3 build/rocm/ci_build test $TEST_IMAGE
+          python3 build/rocm/ci_build test $TEST_IMAGE --test-cmd "pytest tests/core_test.py"
 
diff --git a/.github/workflows/rocm-nightly-upstream-sync.yml b/.github/workflows/rocm-nightly-upstream-sync.yml
@@ -0,0 +1,52 @@
+# Pulls the latest changes from upstream into main and opens a PR to merge
+# them into rocm-main branch.
+
+name: ROCm Nightly Upstream Sync
+on:
+  workflow_dispatch:
+  schedule:
+    - cron: '0 6 * * 1-5'
+permissions:
+  contents: write
+  pull-requests: write
+env:
+  SYNC_BRANCH_NAME: ci-upstream-sync-${{ github.run_number }}_${{ github.run_attempt }}
+jobs:
+  sync-main:
+    runs-on: ubuntu-latest
+    steps:
+      - run: |
+          gh auth status
+          gh repo sync rocm/jax -b main
+        env:
+          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+  create-sync-branch:
+    needs: sync-main
+    runs-on: ubuntu-latest
+    env:
+      GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
+      - name: Create branch
+        run: |
+          git fetch
+          git checkout origin/main 
+          git checkout -b $SYNC_BRANCH_NAME
+          # Try and merge rocm-main into this new branch so that we don't run upstream's CI code
+          git config --global user.email "github-actions@github.com"
+          git config --global user.name "GitHub Actions"
+          git merge origin/rocm-main || true
+          # If the merge creates conflicts, we want to abort and push to origin anyways so that a dev can resolve the conflicts
+          git merge --abort || true
+          git push origin HEAD
+  open-sync-pr:
+    needs: create-sync-branch
+    runs-on: ubuntu-latest
+    steps:
+      - run: |
+          gh pr create --repo $GITHUB_REPOSITORY --head $SYNC_BRANCH_NAME --base rocm-main --title "CI: $(date +%x) upstream sync" --body "Daily sync with upstream"
+          gh pr merge --repo $GITHUB_REPOSITORY --merge --auto $SYNC_BRANCH_NAME
+        env:
+          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+
diff --git a/.github/workflows/rocm-open-upstream-pr.yml b/.github/workflows/rocm-open-upstream-pr.yml
@@ -0,0 +1,41 @@
+name: ROCm Open Upstream PR
+on:
+  pull_request:
+    types: [ labeled ]
+    branches: [ rocm-main ]
+jobs:
+  open-upstream:
+    if: ${{ github.event.label.name == 'open-upstream' }}
+    permissions:
+      contents: write
+      pull-requests: write
+    runs-on: ubuntu-latest
+    env:
+      NEW_BRANCH_NAME: "${{ github.head_ref }}-upstream"
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
+      - name: Rebase code to main
+        run: |
+          git config --global user.email "github-actions@github.com"
+          git config --global user.name "Github Actions"
+          git fetch
+          git checkout -b $NEW_BRANCH_NAME origin/${{ github.head_ref }}
+          git rebase --onto origin/main origin/rocm-main
+          # Force push here so that we don't run into conflicts with the origin branch
+          git push origin HEAD --force
+      - name: Leave link to create PR
+        env:
+          GH_TOKEN: ${{ github.token }}
+        run: |
+          # Bash is not friendly with newline characters, so make our own
+          NL=$'\n'
+          # Encode the PR title and body for passing as URL get parameters
+          TITLE_ENC=$(jq -rn --arg x "[ROCm] ${{ github.event.pull_request.title }}" '$x|@uri')
+          BODY_ENC=$(jq -rn --arg x $"${{ github.event.pull_request.body }}${NL}${NL}Created from: rocm/jax#${{ github.event.pull_request.number }}" '$x|@uri')
+          # Create a link to the that will open up a new PR form to upstream and autofill the fields
+          CREATE_PR_LINK="https://github.com/jax-ml/jax/compare/main...ROCm:jax:$NEW_BRANCH_NAME?expand=1&title=$TITLE_ENC&body=$BODY_ENC"
+          # Add a comment with the link to the PR
+          COMMENT_BODY="Feature branch from main is ready. [Create a new PR][1] destined for upstream?${NL}${NL}[1]: $CREATE_PR_LINK"
+          gh pr comment ${{ github.event.pull_request.number }} --repo rocm/jax --body "$COMMENT_BODY"
+
diff --git a/.github/workflows/upstream-nightly.yml b/.github/workflows/upstream-nightly.yml
@@ -22,7 +22,7 @@ on:
 
 jobs:
   upstream-dev:
-    runs-on: ubuntu-latest
+    runs-on: ROCM-Ubuntu
     permissions:
       contents: read
       issues: write  # for failed-build-issue
diff --git a/build/rocm/README.md b/build/rocm/README.md
@@ -207,3 +207,4 @@ This will generate three wheels in the `dist/` directory:
 ### Simplified Build Script
 
 For a streamlined process, consider using the `jax/build/rocm/dev_build_rocm.py` script.
+
diff --git a/build/rocm/ci_build b/build/rocm/ci_build
@@ -307,6 +307,7 @@ def parse_args():
         "--test-cmd",
         default="./build/rocm/build_rocm.sh && ./build/rocm/run_single_gpu.py -c && ./build/rocm/run_multi_gpu.sh",
     )
+    testp.add_argument("--test-cmd", default="./build/rocm/build_rocm.sh && ./build/rocm/run_single_gpu.py -c && ./build/rocm/run_multi_gpu.sh")
 
     ddp = subp.add_parser("dist_docker")
     ddp.add_argument("--dockerfile", default="build/rocm/Dockerfile.ms")
diff --git a/build/rocm/upload_wheels.sh b/build/rocm/upload_wheels.sh
@@ -0,0 +1,53 @@
+#!/bin/bash
+
+# Check for user-supplied arguments.
+if [[ $# -lt 2 ]]; then
+    echo "Usage: $0 <jax_home_directory> <version>"
+    exit 1
+fi
+
+# Set JAX_HOME and RELEASE_VERSION from user arguments.
+JAX_HOME=$1
+RELEASE_VERSION=$2
+WHEELHOUSE="$JAX_HOME/wheelhouse"
+
+# Projects to upload separately to PyPI.
+PROJECTS=("jax_rocm60_pjrt" "jax_rocm60_plugin")
+
+# PyPI API Token.
+PYPI_API_TOKEN=${PYPI_API_TOKEN:-"pypi-replace_with_token"}
+
+# Ensure the specified JAX_HOME and wheelhouse directories exists.
+if [[ ! -d "$JAX_HOME" ]]; then
+    echo "Error: The specified JAX_HOME directory does not exist: $JAX_HOME"
+    exit 1
+fi
+if [[ ! -d "$WHEELHOUSE" ]]; then
+    echo "Error: The wheelhouse directory does not exist: $WHEELHOUSE"
+    exit 1
+fi
+
+upload_and_release_project() {
+    local project=$1
+
+    echo "Searching for wheels matching project: $project version: $RELEASE_VERSION..."
+    wheels=($(ls $WHEELHOUSE | grep "^${project}-${RELEASE_VERSION}[.-].*\.whl"))
+    if [[ ${#wheels[@]} -eq 0 ]]; then
+        echo "No wheels found for project: $project version: $RELEASE_VERSION. Skipping..."
+        return
+    fi
+    echo "Found wheels for $project: ${wheels[*]}"
+
+    echo "Uploading wheels for $project version $RELEASE_VERSION to PyPI..."
+    for wheel in "${wheels[@]}"; do
+        twine upload --verbose --repository pypi --non-interactive --username "__token__" --password "$PYPI_API_TOKEN" "$WHEELHOUSE/$wheel"
+    done
+}
+
+# Install twine if not already installed.
+python -m pip install --upgrade twine
+
+# Iterate over each project and upload its wheels.
+for project in "${PROJECTS[@]}"; do
+    upload_and_release_project $project
+done
diff --git a/rocm-downstream-dev-guide.md b/rocm-downstream-dev-guide.md
diff --git a/tests/lax_scipy_test.py b/tests/lax_scipy_test.py
diff --git a/tests/pallas/tpu_pallas_distributed_test.py b/tests/pallas/tpu_pallas_distributed_test.py
diff --git a/third_party/xla/workspace.bzl b/third_party/xla/workspace.bzl

Original file line number	Diff line number	Diff line change
@@ -207,3 +207,4 @@ This will generate three wheels in the `dist/` directory:
`207`	`207`	`### Simplified Build Script`
`208`	`208`
`209`	`209`	For a streamlined process, consider using the `jax/build/rocm/dev_build_rocm.py` script.
	`210`	`+`
Original file line number	Diff line number	Diff line change
`@@ -307,6 +307,7 @@ def parse_args():`
`307`	`307`	`"--test-cmd",`
`308`	`308`	`default="./build/rocm/build_rocm.sh && ./build/rocm/run_single_gpu.py -c && ./build/rocm/run_multi_gpu.sh",`
`309`	`309`	`)`
	`310`	`+ testp.add_argument("--test-cmd", default="./build/rocm/build_rocm.sh && ./build/rocm/run_single_gpu.py -c && ./build/rocm/run_multi_gpu.sh")`
`310`	`311`
`311`	`312`	`ddp = subp.add_parser("dist_docker")`
`312`	`313`	`ddp.add_argument("--dockerfile", default="build/rocm/Dockerfile.ms")`