From d86f6cafe66bb2acaeffce3ca195642fbde47235 Mon Sep 17 00:00:00 2001 From: erman-gurses <99776114+erman-gurses@users.noreply.github.com> Date: Fri, 26 Jun 2026 10:44:25 -0700 Subject: [PATCH] feat: Add JAX release support v0.10.0 (#6054) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary ISSUE ID: https://github.com/ROCm/TheRock/issues/3878 Related PR: https://github.com/ROCm/rockrel/pull/65 Extend the JAX multi-arch release workflows to support both `rocm-jaxlib-v0.9.1` and `rocm-jaxlib-v0.10.0` release lines. ## Changes ### Release workflow - Add matrix support for: - `rocm-jaxlib-v0.9.1` - `rocm-jaxlib-v0.10.0` - Add release-line specific metadata: - `jax_repository` - `build_mode` - `gfx_arch` ### Build workflow #### rocm-jaxlib-v0.9.1 - Continue using the existing native build flow. - Build from `ROCm/rocm-jax`. - Continue using `build/ci_build`. #### rocm-jaxlib-v0.10.0 - Build from `ROCm/jax`. - Build a manylinux image using the release branch Dockerfile from `ROCm/rocm-jax`. - Build wheels inside the manylinux container using `build/build.py`. ### Testing On rockrel: https://github.com/ROCm/rockrel/actions/runs/28217952215/job/83613621212 https://github.com/ROCm/rockrel/actions/runs/28224689859 - Preserve existing GPU-family based runner selection through `configure_target_run.py`. - Add a JAX device visibility sanity check before running tests: ```bash python -c "import jax; print(jax.local_devices())" - Verified matrix expansion across Python 3.11–3.14. - Verified native build path for rocm-jaxlib-v0.9.1. - Verified manylinux build path for rocm-jaxlib-v0.10.0. - Verified release artifact publication. - Verified JAX device visibility check executes before tests. TODO: Run JAX builds inside a top-level GitHub Actions job container, similar to existing PyTorch workflows. --- .../multi_arch_build_linux_jax_wheels.yml | 134 ++++++++++++--- .../workflows/multi_arch_release_linux.yml | 3 +- .../multi_arch_release_linux_jax_wheels.yml | 25 ++- .../test_linux_jax_wheels_partial.yml | 75 ++++++--- external-builds/jax/README.md | 156 +++++++++++------- 5 files changed, 278 insertions(+), 115 deletions(-) diff --git a/.github/workflows/multi_arch_build_linux_jax_wheels.yml b/.github/workflows/multi_arch_build_linux_jax_wheels.yml index f43068d75bc..9aa160c09fe 100644 --- a/.github/workflows/multi_arch_build_linux_jax_wheels.yml +++ b/.github/workflows/multi_arch_build_linux_jax_wheels.yml @@ -25,18 +25,31 @@ on: description: "Python version to build wheels for." type: string required: true + jax_repository: + description: "Repository containing the JAX release branch." + type: string + required: true + build_mode: + description: "Build mode: native or manylinux." + type: string + required: true + gfx_arch: + description: "ROCm device package selector used for the manylinux build (e.g. device-gfx942, device-gfx950, or device-all)." + type: string + required: false + default: "" jax_ref: description: "JAX / rocm-jax release ref." type: string required: true - build_jaxlib: - description: "Whether this matrix cell should build jaxlib from JAX source." - type: boolean - required: true rocm_version: description: "ROCm package version to install and build against." type: string required: true + rocm_package_find_links_url: + description: "URL for pip --find-links to install ROCm packages for the JAX manylinux build." + type: string + default: "" tar_url: description: "URL to the TheRock tarball used for the build." type: string @@ -75,18 +88,34 @@ on: description: "Python version to build wheels for." type: string required: true + jax_repository: + description: "GitHub repository in owner/repo format (e.g. ROCm/jax or ROCm/rocm-jax)." + type: string + default: "ROCm/jax" + build_mode: + description: "Build mode." + type: choice + options: + - native + - manylinux + required: true + gfx_arch: + description: "ROCm device package selector used for the manylinux build (e.g. device-gfx942, device-gfx950, or device-all)." + type: string + required: false + default: "" jax_ref: description: "JAX / rocm-jax release ref." type: string required: true - build_jaxlib: - description: "Whether this matrix cell should build jaxlib from JAX source." - type: boolean - required: true rocm_version: description: "ROCm package version to install and build against." type: string required: true + rocm_package_find_links_url: + description: "URL for pip --find-links to install ROCm packages for the JAX manylinux build." + type: string + default: "" tar_url: description: "URL to the TheRock tarball used for the build." type: string @@ -103,7 +132,6 @@ on: description: "Branch, tag, or SHA to checkout. Defaults to the triggering ref." type: string default: "" - run-name: Build Multi-Arch Linux JAX Wheels (${{ inputs.release_type }}, ${{ inputs.rocm_version }}, Multiarch) permissions: @@ -121,7 +149,7 @@ jobs: jax_plugin_version: ${{ steps.write_jax_versions.outputs.jax_plugin_version }} jax_pjrt_version: ${{ steps.write_jax_versions.outputs.jax_pjrt_version }} env: - PACKAGE_DIST_DIR: ${{ github.workspace }}/rocm-jax/jax_rocm_plugin/wheelhouse + MANYLINUX_IMAGE_TAG: therock-jax-manylinux:${{ inputs.jax_ref }}-${{ inputs.python_version }} steps: - name: Checkout TheRock @@ -138,13 +166,14 @@ jobs: ref: ${{ inputs.jax_ref }} - name: Checkout JAX - if: ${{ inputs.build_jaxlib }} + if: ${{ inputs.build_mode == 'manylinux' }} uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3 with: - path: jax - repository: ROCm/jax + path: jax-source + repository: ${{ inputs.jax_repository }} ref: ${{ inputs.jax_ref }} + - name: Configure Git Identity run: | git config --global user.name "therockbot" @@ -155,21 +184,20 @@ jobs: with: python-version: ${{ inputs.python_version }} - - name: Install boto3 + - name: Install python deps for CI run: | - python -m pip install boto3 + pip install -r external-builds/jax/requirements-jax.txt - - name: Determine build arguments - env: - BUILD_JAXLIB: ${{ inputs.build_jaxlib }} + - name: Set package dist dir run: | - if [ "$BUILD_JAXLIB" = "true" ]; then - echo "SOURCE_ARG=--jax-source-dir=$GITHUB_WORKSPACE/jax" >> "$GITHUB_ENV" + if [ "${{ inputs.build_mode }}" = "manylinux" ]; then + echo "PACKAGE_DIST_DIR=${GITHUB_WORKSPACE}/jax-source/dist" >> "$GITHUB_ENV" else - echo "SOURCE_ARG=" >> "$GITHUB_ENV" + echo "PACKAGE_DIST_DIR=${GITHUB_WORKSPACE}/rocm-jax/jax_rocm_plugin/wheelhouse" >> "$GITHUB_ENV" fi - - name: Build JAX Wheels + - name: Build JAX Wheels (native, v0.9.1) + if: ${{ inputs.build_mode == 'native' }} working-directory: rocm-jax env: ROCM_VERSION: ${{ inputs.rocm_version }} @@ -181,8 +209,64 @@ jobs: --python-versions="${PYTHON_VERSION}" \ --rocm-version="${ROCM_VERSION}" \ --therock-path="${TAR_URL}" \ - ${SOURCE_ARG} \ dist_wheels + # TODO: THEROCK_INDEX_URL need to be renamed in the future to + # THEROCK_FIND_LINKS_URL to avoid confusion with pip --index-url. + - name: Build manylinux image + if: ${{ inputs.build_mode == 'manylinux' }} + working-directory: jax-source + env: + THEROCK_VERSION: "" + GFX_ARCH: ${{ inputs.gfx_arch }} + ROCM_PACKAGE_FIND_LINKS_URL: ${{ inputs.rocm_package_find_links_url }} + run: | + cp ../rocm-jax/docker/manylinux/Dockerfile.jax-manylinux_2_28-therock \ + Dockerfile.jax-manylinux_2_28-therock + + docker build \ + -t "${MANYLINUX_IMAGE_TAG}" \ + --file=Dockerfile.jax-manylinux_2_28-therock \ + --build-arg=THEROCK_INDEX_URL="${ROCM_PACKAGE_FIND_LINKS_URL}" \ + --build-arg=THEROCK_VERSION="${THEROCK_VERSION}" \ + --build-arg=GFX_ARCH="${GFX_ARCH}" \ + --progress=plain \ + . + + - name: Determine wheel version suffix + run: | + python build_tools/github_actions/determine_version.py \ + --rocm-version "${{ inputs.rocm_version }}" + + - name: Build JAX Wheels in manylinux container + if: ${{ inputs.build_mode == 'manylinux' }} + working-directory: jax-source + env: + ROCM_VERSION: ${{ inputs.rocm_version }} + PYTHON_VERSION: ${{ inputs.python_version }} + run: | + docker run --rm \ + --user root \ + --env ROCM_VERSION="${ROCM_VERSION}" \ + --env PYTHON_VERSION="${PYTHON_VERSION}" \ + --env ML_WHEEL_VERSION_SUFFIX="${version_suffix}" \ + --env PACKAGE_DIST_DIR="/workspace/jax-source/dist" \ + --volume "${GITHUB_WORKSPACE}:/workspace" \ + --workdir /workspace/jax-source \ + "${MANYLINUX_IMAGE_TAG}" \ + bash -lc ' + python build/build.py build \ + --wheels=jax-rocm-plugin,jax-rocm-pjrt \ + --python_version="${PYTHON_VERSION}" \ + --bazel_startup_options=--bazelrc=build/rocm/rocm.bazelrc \ + --bazel_options=--config=rocm_release_wheel \ + --bazel_options=--repo_env=ROCM_PATH=$(rocm-sdk path --root) \ + --bazel_options=--repo_env=ML_WHEEL_TYPE=release \ + --bazel_options=--repo_env=ML_WHEEL_VERSION_SUFFIX="${ML_WHEEL_VERSION_SUFFIX}" \ + --bazel_options=--//jaxlib/tools:jaxlib_git_hash=$(git rev-parse HEAD) \ + --verbose \ + --detailed_timestamped_log \ + --output_path=$(pwd)/dist + ' - name: Extract JAX versions from built wheels id: write_jax_versions @@ -202,8 +286,6 @@ jobs: --source-dir="${{ env.PACKAGE_DIST_DIR }}" \ --release-type="${{ inputs.release_type }}" - - generate_target_to_run: name: Generate target_to_run if: ${{ inputs.test_amdgpu_family != '' }} @@ -245,7 +327,7 @@ jobs: test_amdgpu_family: ${{ inputs.test_amdgpu_family }} package_index_url: ${{ needs.build_jax_wheels.outputs.package_index_url }} rocm_version: ${{ inputs.rocm_version }} - tar_url: ${{ inputs.tar_url }} + rocm_package_find_links_url: ${{ inputs.rocm_package_find_links_url }} release_type: ${{ inputs.release_type }} python_version: ${{ inputs.python_version }} repository: ${{ inputs.repository || github.repository }} diff --git a/.github/workflows/multi_arch_release_linux.yml b/.github/workflows/multi_arch_release_linux.yml index de8d1b0e31c..fdfec0d8e63 100644 --- a/.github/workflows/multi_arch_release_linux.yml +++ b/.github/workflows/multi_arch_release_linux.yml @@ -251,7 +251,7 @@ jobs: } trigger_release_jax_wheels: - needs: [build_artifacts, build_tarballs] + needs: [build_artifacts, build_tarballs, build_python_packages] name: Trigger Release JAX Wheels runs-on: ubuntu-24.04 permissions: @@ -266,6 +266,7 @@ jobs: "test_amdgpu_family": "${{ contains(fromJSON(inputs.build_config).dist_amdgpu_families, 'gfx94X-dcgpu') && 'gfx94X-dcgpu' || '' }}", "release_type": "${{ inputs.release_type }}", "rocm_version": "${{ inputs.rocm_package_version }}", + "rocm_package_find_links_url": "${{ needs.build_python_packages.outputs.package_find_links_url }}", "tar_url": "${{ fromJSON(needs.build_tarballs.outputs.tarball_urls).multiarch }}", "repository": "${{ inputs.repository || github.repository }}", "ref": "${{ inputs.ref || '' }}" diff --git a/.github/workflows/multi_arch_release_linux_jax_wheels.yml b/.github/workflows/multi_arch_release_linux_jax_wheels.yml index dfaf69e8b81..0c3c62eba15 100644 --- a/.github/workflows/multi_arch_release_linux_jax_wheels.yml +++ b/.github/workflows/multi_arch_release_linux_jax_wheels.yml @@ -3,7 +3,6 @@ # Multi-arch release entry point for Linux JAX wheels. # -# Currently this workflow is uses only rocm-jaxlib-v0.9.1 # This workflow owns the release matrix and delegates each matrix cell to the # reusable single-build workflow. @@ -24,6 +23,10 @@ on: description: "ROCm package version to build against." type: string required: true + rocm_package_find_links_url: + description: "URL for pip --find-links to install ROCm packages for the JAX manylinux build." + type: string + default: "" tar_url: description: "URL to the TheRock tarball used for the build." type: string @@ -52,6 +55,10 @@ on: description: "ROCm package version to build against." type: string required: true + rocm_package_find_links_url: + description: "URL for pip --find-links to install ROCm packages for the JAX manylinux build." + type: string + default: "" tar_url: description: "URL to the TheRock tarball used for the build." type: string @@ -69,7 +76,6 @@ run-name: Release Linux JAX Wheels (${{ inputs.release_type }}, ${{ inputs.rocm_ permissions: contents: read -#TODO: Add rocm-jaxlib-v0.10.0 once it is released and tested. jobs: build_jax_wheels: name: Build | py ${{ matrix.python_version }} | jax ${{ matrix.jax_ref }} @@ -79,9 +85,17 @@ jobs: python_version: ["3.11", "3.12", "3.13", "3.14"] jax_ref: - "rocm-jaxlib-v0.9.1" + - "rocm-jaxlib-v0.10.0" include: - jax_ref: "rocm-jaxlib-v0.9.1" - build_jaxlib: false + jax_repository: "ROCm/rocm-jax" + build_mode: "native" + gfx_arch: "" + + - jax_ref: "rocm-jaxlib-v0.10.0" + jax_repository: "ROCm/jax" + build_mode: "manylinux" + gfx_arch: "device-all" permissions: id-token: write @@ -91,8 +105,11 @@ jobs: test_amdgpu_family: ${{ inputs.test_amdgpu_family }} python_version: ${{ matrix.python_version }} jax_ref: ${{ matrix.jax_ref }} - build_jaxlib: ${{ matrix.build_jaxlib }} + jax_repository: ${{ matrix.jax_repository }} + build_mode: ${{ matrix.build_mode }} + gfx_arch: ${{ matrix.gfx_arch }} rocm_version: ${{ inputs.rocm_version }} + rocm_package_find_links_url: ${{ inputs.rocm_package_find_links_url }} tar_url: ${{ inputs.tar_url }} release_type: ${{ inputs.release_type }} repository: ${{ inputs.repository || github.repository }} diff --git a/.github/workflows/test_linux_jax_wheels_partial.yml b/.github/workflows/test_linux_jax_wheels_partial.yml index 0dc4cc3eedb..2166d65daf9 100644 --- a/.github/workflows/test_linux_jax_wheels_partial.yml +++ b/.github/workflows/test_linux_jax_wheels_partial.yml @@ -18,10 +18,10 @@ on: description: ROCm version required: true type: string - tar_url: - description: URL to TheRock tarball - required: true + rocm_package_find_links_url: + description: "URL for pip --find-links to install ROCm packages for the JAX manylinux build." type: string + default: "" release_type: description: "Release type" type: string @@ -76,8 +76,8 @@ on: description: ROCm version type: string default: "7.14.0.dev0" - tar_url: - description: URL to TheRock tarball + rocm_package_find_links_url: + description: "URL for pip --find-links to install ROCm packages for the JAX manylinux build." type: string default: "" release_type: @@ -134,6 +134,7 @@ jobs: container: image: ghcr.io/rocm/no_rocm_image_ubuntu24_04@sha256:405945a40deaff9db90b9839c0f41d4cba4a383c1a7459b28627047bf6302a26 options: >- + --shm-size=128g --device /dev/kfd --device /dev/dri --group-add render @@ -149,9 +150,9 @@ jobs: env: VIRTUAL_ENV: ${{ github.workspace }}/.venv - TAR_URL: ${{ inputs.tar_url }} PYTHON_VERSION: ${{ inputs.python_version }} WHEEL_INDEX_URL: ${{ inputs.package_index_url }} + ROCM_PACKAGE_FIND_LINKS_URL: ${{ inputs.rocm_package_find_links_url }} steps: - name: Checkout @@ -192,19 +193,6 @@ jobs: run: | pip install -r external-builds/jax/requirements-jax.txt - - name: Configure ROCm from TheRock tarball - run: | - DEST="/opt/rocm-${{ inputs.rocm_version }}" - TARBALL="${RUNNER_TEMP}/therock.tar.gz" - - rm -rf "${DEST}" - mkdir -p "${DEST}" - - wget -O "${TARBALL}" "${TAR_URL}" - tar -xzf "${TARBALL}" -C "${DEST}" - - ln -sfn "${DEST}" /opt/rocm - - name: Determine JAX package versions run: | if [ -n "${{ inputs.jax_plugin_version }}" ] && [ -n "${{ inputs.jax_version }}" ]; then @@ -219,6 +207,27 @@ jobs: --rocm-version "${{ inputs.rocm_version }}" \ --jax-requirements rocm-jax/build/requirements.txt fi + - name: Compute ROCm device extras + id: expand + run: | + python build_tools/github_actions/expand_amdgpu_families.py \ + --amdgpu-families "${{ inputs.test_amdgpu_family }}" \ + --output-mode=device-extras + + - name: Install ROCm packages + run: | + python -m pip install --upgrade pip + + pkg="rocm[libraries]" + if [ -n "${{ steps.expand.outputs.device_extras }}" ]; then + pkg="${pkg%]},${{ steps.expand.outputs.device_extras }}]" + fi + + python -m pip install --pre \ + --find-links "${ROCM_PACKAGE_FIND_LINKS_URL}" \ + "${pkg}==${{ inputs.rocm_version }}" + + echo "ROCM_ROOT=$(rocm-sdk path --root)" >> "$GITHUB_ENV" - name: Install JAX wheels from package index run: | @@ -244,11 +253,31 @@ jobs: - name: Run JAX tests env: - JAX_PLATFORMS: rocm - HIP_DEVICE_LIB_PATH: /opt/rocm/lib/llvm/amdgcn/bitcode - JAX_ROCM_PLUGIN_INTERNAL_BITCODE_PATH: /opt/rocm/lib/llvm/amdgcn/bitcode - JAX_ROCM_PLUGIN_INTERNAL_LLD_PATH: /opt/rocm/lib/llvm/bin + # Avoid large upfront GPU memory reservation. + XLA_PYTHON_CLIENT_PREALLOCATE: "false" + + # Set XLA allocator. + XLA_PYTHON_CLIENT_ALLOCATOR: platform + + # Reduce compilation/autotune memory pressure. + XLA_FLAGS: >- + --xla_gpu_force_compilation_parallelism=1 + --xla_gpu_enable_nccl_comm_splitting=false + --xla_gpu_enable_command_buffer= + run: | + set -euxo pipefail + + unset ROCM_ROOT + unset HIP_DEVICE_LIB_PATH + unset JAX_ROCM_PLUGIN_INTERNAL_BITCODE_PATH + unset JAX_ROCM_PLUGIN_INTERNAL_LLD_PATH + unset HIP_VISIBLE_DEVICES + unset ROCR_VISIBLE_DEVICES + + python -c "import jax; print(jax.local_devices())" + + echo "=== Starting full test suite ===" pytest jax/tests/multi_device_test.py -q --log-cli-level=INFO pytest jax/tests/core_test.py -q --log-cli-level=INFO pytest jax/tests/util_test.py -q --log-cli-level=INFO diff --git a/external-builds/jax/README.md b/external-builds/jax/README.md index d2cd1ce6cfa..0cf58a5f4a2 100644 --- a/external-builds/jax/README.md +++ b/external-builds/jax/README.md @@ -30,15 +30,15 @@ For upstream JAX development references, see: ### Supported JAX versions -Support for JAX is provided via stable release branches of -[ROCm/rocm-jax](https://github.com/ROCm/rocm-jax). +Support for JAX is provided via stable release branches. + +JAX 0.9.1 uses release branch from [ROCm/rocm-jax](https://github.com/ROCm/rocm-jax). +Starting with JAX 0.10.0, build support uses the [ROCm/jax](https://github.com/ROCm/jax) repository. | JAX version | Linux | Windows | | ----------- | --------------------------------------------------------------------------------------------------------------- | ---------------- | -| 0.9.2 | ✅ Supported via [ROCm/rocm-jax `rocm-jaxlib-v0.9.2`](https://github.com/ROCm/rocm-jax/tree/rocm-jaxlib-v0.9.2) | ❌ Not supported | +| 0.10.0 | ✅ Supported via [ROCm/jax `rocm-jaxlib-v0.10.0`](https://github.com/ROCm/jax/tree/rocm-jaxlib-v0.10.0) | ❌ Not supported | | 0.9.1 | ✅ Supported via [ROCm/rocm-jax `rocm-jaxlib-v0.9.1`](https://github.com/ROCm/rocm-jax/tree/rocm-jaxlib-v0.9.1) | ❌ Not supported | -| 0.9.0 | ✅ Supported via [ROCm/rocm-jax `rocm-jaxlib-v0.9.0`](https://github.com/ROCm/rocm-jax/tree/rocm-jaxlib-v0.9.0) | ❌ Not supported | -| 0.8.2 | ✅ Supported via [ROCm/rocm-jax `rocm-jaxlib-v0.8.2`](https://github.com/ROCm/rocm-jax/tree/rocm-jaxlib-v0.8.2) | ❌ Not supported | See also: @@ -61,60 +61,79 @@ This repository builds the following ROCm-enabled JAX artifacts: ### How building with TheRock differs from upstream The upstream [rocm-jax build instructions](https://github.com/ROCm/rocm-jax/blob/master/BUILDING.md) -assume that a stable ROCm version is already installed on the system. When -building through TheRock, we instead build ROCm from the latest sources and -provide it via **tarballs** with arbitrary install locations. +assume that a stable ROCm version is already installed on the system. + +TheRock currently supports two build paths depending on the JAX release branch: -> [!IMPORTANT] -> JAX currently requires a **tarball-based** ROCm installation. TheRock Python -> wheel packages (`rocm[libraries,devel]`) are **not yet supported** for -> building JAX. This is because the JAX build system expects a traditional -> filesystem-based ROCm installation layout (e.g. `/opt/rocm-/`) that -> tarballs provide but Python wheel packages do not. +- **JAX 0.9.1** uses the legacy tarball-based build flow via + `build/ci_build --therock-path`. +- **JAX 0.10.0** builds against ROCm Python packages installed from the + TheRock multi-arch Python package index. ### Prerequisites - **OS**: Linux (supported distributions with ROCm) - **Python**: 3.12 recommended -- **Compiler**: Clang (provided via the TheRock tarball) -- **ROCm**: Provided via a TheRock tarball (see - [RELEASES.md](/RELEASES.md#installing-from-tarballs) for tarball download - locations and installation instructions) +- **Compiler**: + - JAX 0.9.1: Clang provided by the TheRock tarball + - JAX 0.10.0: Clang provided by the manylinux build environment +- **ROCm**: + - JAX 0.9.1: TheRock tarball + - JAX 0.10.0: ROCm Python packages from the TheRock multi-arch package index ### Steps -1. Checkout rocm-jax and jax: +1. Checkout the source repository for your JAX version. + + **JAX 0.9.1** ```bash git clone https://github.com/ROCm/rocm-jax.git git clone https://github.com/ROCm/jax.git + pushd rocm-jax - git checkout rocm-jaxlib-v + git checkout rocm-jaxlib-v0.9.1 popd + pushd jax - git checkout rocm-jaxlib-v + git checkout rocm-jaxlib-v0.9.1 + popd + ``` + + **JAX 0.10.0** + + ```bash + git clone https://github.com/ROCm/jax.git + + pushd jax + git checkout rocm-jaxlib-v0.10.0 popd ``` 1. Choose your configuration: - - **JAX version**: e.g. `0.9.1`, `0.9.0`, or `0.8.2` + - **JAX version**: e.g. `0.9.1` or `0.10.0` - **Python version**: e.g. `3.12` - - **TheRock tarball**: A tarball URL, a local tarball file path, or a - directory containing a ROCm installation. Nightly tarballs are available - at . -1. Build all wheels: + For **JAX 0.9.1**: + + - TheRock tarball URL, local tarball, or extracted ROCm installation. + + For **JAX 0.10.0**: + + - TheRock multi-arch Python package index. + +1. Build JAX 0.9.1 (legacy tarball flow) ```bash pushd rocm-jax - python3 build/ci_build \ - --compiler=clang \ - --python-versions="3.12" \ - --rocm-version="" \ - --therock-path="" \ - --jax-source-dir="" \ - dist_wheels + PYTHON_VERSION= + ROCM_VERSION= + + python3 build/ci_build --therock-path "" + --python-versions="$PYTHON_VERSION" + --rocm-version="$ROCM_VERSION" + dist_wheels popd ``` @@ -124,17 +143,40 @@ provide it via **tarballs** with arbitrary install locations. > For JAX >= 0.9.1, jaxlib is installed from upstream PyPI, so this flag > can be omitted. +1. Build JAX 0.10.0 (multi-arch package flow) + + JAX 0.10.0 builds are performed by the GitHub Actions workflow: + + - `.github/workflows/multi_arch_build_linux_jax_wheels.yml` + + The workflow installs ROCm Python packages from the configured TheRock + multi-arch package index before building `jax_rocm7_plugin` and + `jax_rocm7_pjrt`. + 1. Locate built wheels: + **JAX 0.9.1** + After a successful build, wheels will be available in: ```text rocm-jax/jax_rocm_plugin/wheelhouse/ ``` -For more detailed build options (including building `jax_rocm7_plugin` and -`jax_rocm7_pjrt` wheels individually), see -[ROCm/rocm-jax BUILDING.md](https://github.com/ROCm/rocm-jax/blob/master/BUILDING.md#building). + **JAX 0.10.0** + + After a successful build, wheels will be available in: + + ```text + jax/dist/ + ``` + +For more detailed build options, see the build instructions for the JAX +release branch you are using. + +- JAX 0.9.1: [ROCm/rocm-jax BUILDING.md](https://github.com/ROCm/rocm-jax/blob/master/BUILDING.md#building) +- JAX 0.10.0: See the `ROCm/jax` repository and the + `.github/workflows/multi_arch_build_linux_jax_wheels.yml` workflow in TheRock. ## Test instructions @@ -165,27 +207,26 @@ For more detailed build options (including building `jax_rocm7_plugin` and 1. Install requirements: ```bash - pip install -r external-builds/jax/requirements-jax.txt + cd jax + pip install -r build/test-requirements.txt + pip install pytest-html pytest-csv uv pytest-json-report ``` -1. Install ROCm from TheRock tarball: +1. Install ROCm Python packages: ```bash - python build_tools/install_rocm_from_artifacts.py \ - --release "" \ - --artifact-group "" \ - --output-dir "/opt/rocm-" + pip install \ + --index-url \ + "rocm[libraries,device-]==" ``` - For detailed instructions and example usage, see the - [TheRock RELEASES.md](https://github.com/ROCm/TheRock/blob/main/RELEASES.md#automated-tarball-extraction). - 1. Install JAX wheels from the package index: ```bash - # Install jaxlib, jax_rocm7_plugin, and jax_rocm7_pjrt from the GPU-family index - pip install --index-url "https://rocm.nightlies.amd.com/v2//" \ - jaxlib jax_rocm7_plugin jax_rocm7_pjrt + pip install \ + --index-url \ + jax_rocm7_plugin \ + jax_rocm7_pjrt # Install jax from PyPI to match the version pip install jax== @@ -200,21 +241,14 @@ For more detailed build options (including building `jax_rocm7_plugin` and pytest jax_tests/tests/scipy_stats_test.py -q --log-cli-level=INFO ``` -We are planning to expand our test coverage and update the testing workflow. -Upcoming changes will include running smoke tests, unit tests, and multi-GPU -tests using the `pip install` packaging method for improved reliability and -consistency. Tracking issue: -[ROCm/TheRock#2592](https://github.com/ROCm/TheRock/issues/2592) - ## Nightly releases ### Gating releases with JAX tests -With passing builds we upload `jaxlib`, `jax_rocm7_plugin`, and -`jax_rocm7_pjrt` wheels to subfolders of the `v2-staging` directory in the -nightly release S3 bucket with a public URL at -. +Successful builds publish JAX wheels to the nightly multi-arch Python package +index: + + -Only with passing JAX tests we promote passed wheels to the `v2` directory in -the nightly release S3 bucket with a public URL at -. +The published wheels are validated by the JAX test workflow as part of the +nightly release process before being made available for use.