TheRock/rocm-systems/.github/workflows/therock-rccl-test-jax-collective.yml

name: TheRock Test JAX Collective for rccl

on:
  workflow_call:
    inputs:
      amdgpu_families:
        type: string
      artifact_group:
        type: string
      test_runs_on:
        type: string
      artifact_run_id:
        type: string
      notify_email:
        type: string
        default: ''
      notify_webhook:
        type: string
        default: ''
  workflow_dispatch:
    inputs:
      amdgpu_families:
        type: string
      artifact_group:
        type: string
      test_runs_on:
        type: string
      artifact_run_id:
        type: string
      notify_email:
        description: 'Email address to send test summary report'
        type: string
        default: ''
      notify_webhook:
        description: 'Teams webhook URL for test summary notifications'
        type: string
        default: ''

permissions:
  contents: read

jobs:
  test_jax_collective:
    name: 'Test JAX collective smoke'
    runs-on: ${{ inputs.test_runs_on }}
    container:
      image: ghcr.io/rocm/no_rocm_image_ubuntu24_04@sha256:4150afe4759d14822f0e3f8930e1124f26e11f68b5c7b91ec9a02b20b1ebbb98
      options: --ipc host
        --group-add video
        --device /dev/kfd
        --device /dev/dri
        --group-add 110
        --group-add 993
        --group-add 992
        --ulimit memlock=-1:-1
        --security-opt seccomp=unconfined
        --env-file /etc/podinfo/gha-gpu-isolation-settings
        --user 0:0
    defaults:
      run:
        shell: bash
    env:
      VENV_DIR: ${{ github.workspace }}/.venv
      ARTIFACT_RUN_ID: "${{ inputs.artifact_run_id }}"
      OUTPUT_ARTIFACTS_DIR: "./build"
      JAX_SRC: ${{ github.workspace }}/jax-src
    steps:
      - name: Checkout Repository
        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
        with:
          repository: "ROCm/TheRock"
          ref: d8e27dff1d165462ffa5bd50057a075b943c97c5 # 2026-07-13

      - name: Checkout rocm-systems scripts
        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
        with:
          sparse-checkout: .github/scripts
          path: rocm-systems

      - name: Run setup test environment workflow
        uses: './.github/actions/setup_test_environment'
        with:
          ARTIFACT_RUN_ID: ${{ env.ARTIFACT_RUN_ID }}
          AMDGPU_FAMILIES: ${{ inputs.amdgpu_families }}
          ARTIFACT_GROUP: ${{ inputs.artifact_group }}
          OUTPUT_ARTIFACTS_DIR: ${{ env.OUTPUT_ARTIFACTS_DIR }}
          VENV_DIR: ${{ env.VENV_DIR }}
          FETCH_ARTIFACT_ARGS: "--rccl --tests"
          IS_PR_FROM_FORK: ${{ github.event.pull_request.head.repo.fork }}

      - name: Install JAX and test dependencies
        timeout-minutes: 15
        # Intentionally installs latest nightly wheels against pinned test sources.
        # The 10 smoke tests use stable pmap/shard_map APIs that are compatible
        # across JAX versions. This lets us catch RCCL regressions with current JAX.
        #
        # xla_rocm_plugin.so has RUNPATH entries that resolve math libraries
        # (hipBLASLt, hipFFT, MIOpen, rocBLAS, etc.) from the _rocm_sdk_libraries
        # site-package. The arch-specific v2 index provides the real binary
        # wheels; the whl-multi-arch index only has stub metapackages.
        # The test script overrides librccl.so with the CI-built version.
        run: |
          AMDGPU_LOWER=$(echo "${{ inputs.amdgpu_families }}" | tr '[:upper:]' '[:lower:]')
          uv pip install --prerelease=allow \
            rocm-sdk-core "rocm-sdk-libraries-${AMDGPU_LOWER}" \
            --index-url "https://rocm.nightlies.amd.com/v2/${{ inputs.amdgpu_families }}/"
          uv pip install --prerelease=allow \
            jax-rocm7-plugin jax-rocm7-pjrt jax \
            --index-url https://rocm.nightlies.amd.com/whl-multi-arch/ \
            --extra-index-url https://pypi.org/simple/
          uv pip install \
            pytest \
            pytest-timeout \
            hypothesis \
            numpy \
            absl-py

      - name: Run JAX collective smoke tests
        timeout-minutes: 10
        env:
          NCCL_DEBUG: INFO
          NOTIFY_EMAIL: ${{ inputs.notify_email }}
          NOTIFY_WEBHOOK: ${{ inputs.notify_webhook }}
        run: |
          NOTIFY_ARGS=()
          [ -n "$NOTIFY_EMAIL" ] && NOTIFY_ARGS+=(--notify-email "$NOTIFY_EMAIL")
          [ -n "$NOTIFY_WEBHOOK" ] && NOTIFY_ARGS+=(--notify-webhook "$NOTIFY_WEBHOOK")
          python rocm-systems/.github/scripts/test_jax_collective.py \
            --artifact-dir ${{ env.OUTPUT_ARTIFACTS_DIR }} \
            --jax-src ${{ env.JAX_SRC }} \
            --results-log ${{ env.JAX_SRC }}/jax_collective_results.log \
            "${NOTIFY_ARGS[@]}"

      - name: Upload test results
        if: always()
        uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
        with:
          name: jax-collective-results-${{ inputs.amdgpu_families }}
          path: |
            ${{ env.JAX_SRC }}/jax_collective_results.log
            ${{ env.JAX_SRC }}/jax_collective_results.xml
            ${{ env.JAX_SRC }}/jax_collective_summary.txt
          retention-days: 14