TheRock/.github/workflows/multi_arch_build_linux_jax_wheels.yml

# Copyright Advanced Micro Devices, Inc.
# SPDX-License-Identifier: MIT

# Reusable single-build workflow for the Linux JAX wheel flow.
#
# This workflow is called once per release matrix cell from
# test_linux_jax_wheels_partial.yml
#
# Built shape:
#   - One build/ci_build invocation for one python_version and jax_ref
#     matrix cell.
#   - Uploads the resulting wheels to the multi-arch release bucket.
#   - Runs the JAX test subset against the published multi-arch index.

name: Multi-Arch Build Linux JAX Wheels

on:
  workflow_call:
    inputs:
      test_amdgpu_family:
        description: "GPU family to use for JAX tests; empty means skip tests."
        type: string
        default: ""
      python_version:
        description: "Python version to build wheels for."
        type: string
        required: true
      jax_repository:
        description: "Repository containing the JAX release branch."
        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
      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: ""
      release_type:
        description: "Release type."
        type: string
        required: true
      repository:
        description: "Repository to checkout. Defaults to github.repository."
        type: string
        default: ""
      ref:
        description: "Branch, tag, or SHA to checkout. Defaults to the triggering ref."
        type: string
        default: ""
    outputs:
      package_index_url:
        description: "Package index URL for the multi-arch release bucket."
        value: ${{ jobs.build_jax_wheels.outputs.package_index_url }}
      jax_version:
        value: ${{ jobs.build_jax_wheels.outputs.jax_version }}
      jaxlib_version:
        value: ${{ jobs.build_jax_wheels.outputs.jaxlib_version }}
      jax_plugin_version:
        value: ${{ jobs.build_jax_wheels.outputs.jax_plugin_version }}
      jax_pjrt_version:
        value: ${{ jobs.build_jax_wheels.outputs.jax_pjrt_version }}
  workflow_dispatch:
    inputs:
      test_amdgpu_family:
        description: "Representative GPU family to use for JAX tests."
        type: string
        default: "gfx94X-dcgpu"
      python_version:
        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"
      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
      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: ""
      release_type:
        description: "Release type."
        type: string
        required: true
      repository:
        description: "Repository to checkout. Defaults to github.repository."
        type: string
        default: ""
      ref:
        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:
  contents: read
  id-token: write

jobs:
  build_jax_wheels:
    name: Build | py ${{ inputs.python_version }} | jax ${{ inputs.jax_ref }}
    runs-on: ${{ github.repository_owner == 'ROCm' && 'azure-linux-scale-rocm' || 'ubuntu-24.04' }}
    outputs:
      package_index_url: ${{ steps.publish_jax_wheels.outputs.package_index_url }}
      jax_version: ${{ steps.write_jax_versions.outputs.jax_version }}
      jaxlib_version: ${{ steps.write_jax_versions.outputs.jaxlib_version }}
      jax_plugin_version: ${{ steps.write_jax_versions.outputs.jax_plugin_version }}
      jax_pjrt_version: ${{ steps.write_jax_versions.outputs.jax_pjrt_version }}
    env:
      MANYLINUX_IMAGE_TAG: therock-jax-manylinux:${{ inputs.jax_ref }}-${{ inputs.python_version }}
      PACKAGE_DIST_DIR: ${{ github.workspace }}/jax-source/dist

    steps:
      - name: Checkout TheRock
        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
        with:
          repository: ${{ inputs.repository || github.repository }}
          ref: ${{ inputs.ref || github.ref_name }}

      - name: Checkout rocm-jax
        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
        with:
          path: rocm-jax
          repository: ROCm/rocm-jax
          ref: ${{ inputs.jax_ref }}

      - name: Checkout JAX
        id: checkout_jax
        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
        with:
          path: jax-source
          repository: ${{ inputs.jax_repository }}
          ref: ${{ inputs.jax_ref }}


      - name: Configure Git Identity
        run: |
          git config --global user.name "therockbot"
          git config --global user.email "therockbot@amd.com"

      - name: Set up Python
        uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0
        with:
          python-version: ${{ inputs.python_version }}

      - name: Install python deps for CI
        run: |
          pip install -r external-builds/jax/requirements-jax.txt

      # 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
        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
        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="${{ steps.checkout_jax.outputs.commit }}" \
                --verbose \
                --detailed_timestamped_log \
                --output_path=$(pwd)/dist
            '

      - name: Extract JAX versions from built wheels
        id: write_jax_versions
        run: |
          python3 ./build_tools/github_actions/write_jax_versions.py \
            --dist-dir "${{ env.PACKAGE_DIST_DIR }}"

      - name: Configure AWS Credentials
        uses: ./.github/actions/configure_aws_artifacts_credentials
        with:
          release_type: ${{ inputs.release_type }}

      - name: Upload JAX wheels to release bucket
        id: publish_jax_wheels
        run: |
          python build_tools/github_actions/publish_jax_to_release_bucket.py \
            --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 != '' }}
    runs-on: ubuntu-24.04
    outputs:
      test_runs_on: ${{ steps.configure.outputs.test-runs-on }}
    steps:
      - name: Checking out repository
        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
        with:
          repository: ${{ inputs.repository || github.repository }}
          ref: ${{ inputs.ref || github.ref_name }}

      - name: Checkout CI config
        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
        with:
          repository: ROCm/therock-ci-config
          ref: main
          path: ci-config
        continue-on-error: true

      - name: Generating target to run
        id: configure
        env:
          TARGET: ${{ inputs.test_amdgpu_family }}
          PLATFORM: linux
          CI_CONFIG_PATH: ci-config
        run: python ./build_tools/github_actions/configure_target_run.py

  test_jax_wheels:
    if: ${{ inputs.test_amdgpu_family != '' }}
    needs: [build_jax_wheels, generate_target_to_run]
    name: Test JAX wheels | py ${{ inputs.python_version }} | jax ${{ inputs.jax_ref }}
    permissions:
      contents: read
      id-token: write
    uses: ./.github/workflows/test_linux_jax_wheels_partial.yml
    with:
      test_amdgpu_family: ${{ inputs.test_amdgpu_family }}
      package_index_url: ${{ needs.build_jax_wheels.outputs.package_index_url }}
      rocm_version: ${{ inputs.rocm_version }}
      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 }}
      ref: ${{ inputs.ref }}
      jax_ref: ${{ inputs.jax_ref }}
      jax_version: ${{ needs.build_jax_wheels.outputs.jax_version }}
      jaxlib_version: ${{ needs.build_jax_wheels.outputs.jaxlib_version }}
      jax_plugin_version: ${{ needs.build_jax_wheels.outputs.jax_plugin_version }}
      jax_pjrt_version: ${{ needs.build_jax_wheels.outputs.jax_pjrt_version }}
      test_runs_on: ${{ needs.generate_target_to_run.outputs.test_runs_on }}