TheRock/.github/workflows/test_linux_jax_wheels.yml

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

name: Test Linux JAX Wheels

on:
  workflow_call:
    inputs:
      amdgpu_family:
        required: true
        type: string
      release_type:
        required: true
        type: string
      package_index_url:
        description: Base CloudFront URL for the Python package index
        required: true
        type: string
      rocm_version:
        description: ROCm version (optional, informational)
        required: false
        type: string
      tar_url:
        description: URL to TheRock tarball to configure ROCm
        required: true
        type: string
      python_version:
        description: Python version(s) to test (e.g., "3.12")
        required: true
        type: string
      repository:
        description: "Repository to checkout. Otherwise, defaults to `github.repository`."
        type: string
      jax_ref:
        description: rocm-jax repository ref/branch to check out
        required: false
        type: string
        default: "rocm-jaxlib-v0.11.0"
      jax_version:
        description: "Base JAX version (e.g. 0.8.2) for installing jax from PyPI. Extracted from built wheels by write_jax_versions.py."
        required: false
        type: string
      jaxlib_version:
        description: "jaxlib wheel version (e.g. 0.9.0+rocm7 or 0.8.2+rocm7.12.0). Extracted from built wheels by write_jax_versions.py."
        required: false
        type: string
      jax_plugin_version:
        description: "jax_rocm7_plugin wheel version. Extracted from built wheels by write_jax_versions.py."
        required: false
        type: string
      jax_pjrt_version:
        description: "jax_rocm7_pjrt wheel version. Extracted from built wheels by write_jax_versions.py."
        required: false
        type: string
      ref:
        description: "Branch, tag or SHA to checkout. Defaults to the reference or SHA that triggered the workflow."
        type: string
      test_runs_on:
        required: true
        type: string

  workflow_dispatch:
    inputs:
      amdgpu_family:
        type: choice
        options:
          - gfx101X-dgpu
          - gfx103X-all
          - gfx110X-all
          - gfx1150
          - gfx1151
          - gfx1152
          - gfx1153
          - gfx120X-all
          - gfx900
          - gfx90c
          - gfx906
          - gfx908
          - gfx90a
          - gfx94X-dcgpu
          - gfx950-dcgpu
          - gfx125X-dcgpu
        default: gfx94X-dcgpu
      release_type:
        description: 'Release type; developer-triggered jobs should use "dev"'
        type: choice
        options:
          - dev
        default: dev
      package_index_url:
        description: Base CloudFront URL for the Python package index
        required: true
        type: string
        default: https://rocm.nightlies.amd.com/v2-staging/
      rocm_version:
        description: ROCm version
        required: false
        type: string
      tar_url:
        description: URL to TheRock tarball to configure ROCm
        required: true
        type: string
      python_version:
        description: Python version(s) to test (e.g., "3.12")
        required: true
        type: string
        default: "3.12"
      jax_ref:
        description: rocm-jax repository ref/branch to check out
        required: false
        type: string
        default: "rocm-jaxlib-v0.11.0"
      jax_version:
        description: "Base JAX version (e.g. 0.8.2). Leave empty to auto-detect from rocm-jax requirements."
        required: false
        type: string
      jaxlib_version:
        description: "jaxlib wheel version (e.g. 0.9.0+rocm7 or 0.8.2+rocm7.12.0). Leave empty to auto-compute from rocm_version."
        required: false
        type: string
      jax_plugin_version:
        description: "jax_rocm7_plugin wheel version. Leave empty to use jaxlib_version."
        required: false
        type: string
      jax_pjrt_version:
        description: "jax_rocm7_pjrt wheel version. Leave empty to use jaxlib_version."
        required: false
        type: string
      test_runs_on:
        description: Runner label to use. The selected runner should have a GPU supported by amdgpu_family
        required: true
        type: string
        default: "linux-gfx942-1gpu-ccs-csp-ossci-rocm"
      ref:
        description: "Branch, tag or SHA to checkout. Defaults to the reference or SHA that triggered the workflow."
        type: string

permissions:
  contents: read
  packages: read

jobs:
  test_jax_wheels:
    name: Test JAX Wheels | ${{ inputs.amdgpu_family }}
    runs-on: ${{ inputs.test_runs_on }}
    container:
      image: ghcr.io/rocm/no_rocm_image_ubuntu24_04@sha256:405945a40deaff9db90b9839c0f41d4cba4a383c1a7459b28627047bf6302a26
      # DNS workaround: These --dns entries are added to resolve CloudFront/runner DNS issues
      # They should NOT be required for testing JAX and can be removed once DNS issues are resolved
      # TODO: Remove these lines when issue is fixed - https://github.com/ROCm/TheRock/issues/3014
      options:  >-
          --device /dev/kfd
          --device /dev/dri
          --group-add render
          --group-add video
          --user root
          --env-file /etc/podinfo/gha-gpu-isolation-settings
          --dns 8.8.8.8
          --dns 8.8.4.4
    defaults:
      run:
        shell: bash
    env:
      VIRTUAL_ENV: ${{ github.workspace }}/.venv
      AMDGPU_FAMILY: ${{ inputs.amdgpu_family }}
      THEROCK_TAR_URL: ${{ inputs.tar_url }}
      PYTHON_VERSION: ${{ inputs.python_version }}
      WHEEL_INDEX_URL: ${{ inputs.package_index_url }}/${{ inputs.amdgpu_family }}

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

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

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

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

      - name: System deps, venv configure
        run: |
          python3 -m venv "${VIRTUAL_ENV}"
          echo "PATH=${VIRTUAL_ENV}/bin:${PATH}" >> "$GITHUB_ENV"
          python3 build_tools/setup_venv.py "${VIRTUAL_ENV}" --activate-in-future-github-actions-steps

      - name: Install base JAX test requirements
        run: |
          # This script sets up the venv and activates it across steps; keep it consistent
          pip install -r external-builds/jax/requirements-jax.txt

      - name: Configure ROCm from TheRock tarball
        env:
          ROCM_VERSION: ${{ inputs.rocm_version }}
          AMDGPU_FAMILY: ${{ inputs.amdgpu_family }}
        run: |
          DEST="/opt/rocm-${{ inputs.rocm_version }}"
          # Install directly from TheRock release buckets (nightly/dev) using the provided version
          python build_tools/install_rocm_from_artifacts.py \
            --release "${{ inputs.rocm_version }}" \
            --artifact-group "${{ inputs.amdgpu_family }}" \
            --output-dir "${DEST}"
          # Temporary CI workaround.
          # TODO(#4201): Remove this /opt/rocm symlink once JAX wheel
          # packaging/RUNPATH fixes are validated and CI no longer depends on
          # /opt/rocm -> /opt/rocm-<version> path translation.
          ln -sfn "${DEST}" /opt/rocm

      - name: Determine JAX package versions
        run: |
          # The plugin/pjrt wheels embed the ROCm major version in their package
          # name (jax_rocm7_*, jax_rocm10_*), tracking the ROCm release they were
          # built against.
          rocm_major="${{ inputs.rocm_version }}"
          rocm_major="${rocm_major%%.*}"
          if [ -z "${rocm_major}" ]; then
            echo "::error::rocm_version is required to determine the JAX plugin package name"
            exit 1
          fi
          echo "JAX_PLUGIN_PACKAGE=jax_rocm${rocm_major}_plugin" >> "$GITHUB_ENV"
          echo "JAX_PJRT_PACKAGE=jax_rocm${rocm_major}_pjrt" >> "$GITHUB_ENV"

          # Use versions passed from the build workflow when available.
          # For manual workflow_dispatch, compute the version suffix using
          # compute_jax_package_version.py (matches rocm-jax ci_build format).
          if [ -n "${{ inputs.jax_plugin_version }}" ] && [ -n "${{ inputs.jax_version }}" ]; then
            echo "Using versions passed from build workflow"
            echo "JAXLIB_VERSION=${{ inputs.jaxlib_version }}" >> "$GITHUB_ENV"
            echo "JAX_VERSION=${{ inputs.jax_version }}" >> "$GITHUB_ENV"
            echo "JAX_PLUGIN_VERSION=${{ inputs.jax_plugin_version }}" >> "$GITHUB_ENV"
            echo "JAX_PJRT_VERSION=${{ inputs.jax_pjrt_version }}" >> "$GITHUB_ENV"
          else
            echo "Computing versions from rocm_version and rocm-jax requirements"
            pip install packaging
            python build_tools/github_actions/compute_jax_package_version.py \
              --rocm-version "${{ inputs.rocm_version }}" \
              --jax-requirements rocm-jax/build/requirements.txt
          fi

      - name: Install JAX wheels from package index
        run: |
          echo "Installing JAX packages:"
          echo "  JAXLIB_VERSION=${JAXLIB_VERSION}"
          echo "  JAX_PLUGIN_VERSION=${JAX_PLUGIN_VERSION}"
          echo "  JAX_PJRT_VERSION=${JAX_PJRT_VERSION}"
          echo "  JAX_VERSION=${JAX_VERSION}"
          echo "  JAX_PLUGIN_PACKAGE=${JAX_PLUGIN_PACKAGE}"
          echo "  JAX_PJRT_PACKAGE=${JAX_PJRT_PACKAGE}"
          # Install plugin/pjrt from the GPU-family index
          pip install --index-url "${{ env.WHEEL_INDEX_URL }}" \
            "${JAX_PLUGIN_PACKAGE}==${JAX_PLUGIN_VERSION}" \
            "${JAX_PJRT_PACKAGE}==${JAX_PJRT_VERSION}"
          # For JAX <= 0.9.0, jaxlib is built and installed from GPU index.
          # For JAX >= 0.9.1, jaxlib is not built - install jax and jaxlib from upstream PyPI.
          if [ -n "${JAXLIB_VERSION}" ]; then
            pip install --index-url "${{ env.WHEEL_INDEX_URL }}" "jaxlib==${JAXLIB_VERSION}"
            pip install jax==${JAX_VERSION}
          else
            pip install jax==${JAX_VERSION} jaxlib==${JAX_VERSION}
          fi

      - 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
        run: |
          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
          pytest jax/tests/scipy_stats_test.py -q --log-cli-level=INFO