diff --git a/.github/workflows/ci-build.yaml b/.github/workflows/ci-build.yaml index 2b3362516165..83edccbf0754 100644 --- 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@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0 with: @@ -109,7 +101,7 @@ jobs: documentation: name: Documentation - test code snippets - runs-on: ubuntu-latest + runs-on: ROCM-Ubuntu timeout-minutes: 10 strategy: matrix: @@ -146,19 +138,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@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0 with: @@ -229,9 +215,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 @@ -250,7 +234,7 @@ jobs: path: ${{ steps.pip-cache.outputs.dir }} key: ${{ runner.os }}-pip-ffi-examples-${{ hashFiles('**/setup.py', '**/requirements.txt', '**/test-requirements.txt', 'examples/**/pyproject.toml') }} - name: Install JAX - run: pip install .[cuda12] + run: pip install . - name: Build and install example project run: python -m pip install -v ./examples/ffi[test] env: @@ -259,7 +243,7 @@ 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: diff --git a/.github/workflows/rocm-nightly-upstream-sync.yml b/.github/workflows/rocm-nightly-upstream-sync.yml new file mode 100644 index 000000000000..2f169088147c --- /dev/null +++ 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 new file mode 100644 index 000000000000..bd14fa050577 --- /dev/null +++ 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 index e0dc14468871..1613fca23568 100644 --- a/.github/workflows/upstream-nightly.yml +++ b/.github/workflows/upstream-nightly.yml @@ -22,7 +22,7 @@ on: jobs: upstream-dev: - runs-on: ubuntu-20.04-16core + runs-on: ROCM-Ubuntu permissions: contents: read checks: write # for upload-artifact diff --git a/build/rocm/README.md b/build/rocm/README.md index 58427826f73f..450736547052 100644 --- 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/upload_wheels.sh b/build/rocm/upload_wheels.sh new file mode 100644 index 000000000000..129c87006af6 --- /dev/null +++ b/build/rocm/upload_wheels.sh @@ -0,0 +1,53 @@ +#!/bin/bash + +# Check for user-supplied arguments. +if [[ $# -lt 2 ]]; then + echo "Usage: $0 " + 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/tests/lax_scipy_test.py b/tests/lax_scipy_test.py index c30e8da651be..55ab4f634dfc 100644 --- a/tests/lax_scipy_test.py +++ b/tests/lax_scipy_test.py @@ -445,6 +445,7 @@ def testSphHarmOrderOneDegreeOne(self): self.assertAllClose(actual, expected, rtol=1e-8, atol=6e-8) + @unittest.skip(reason="https://github.com/jax-ml/jax/pull/25675") @jtu.sample_product( [dict(l_max=l_max, num_z=num_z) for l_max, num_z in zip([1, 3, 8, 10], [2, 6, 7, 8])