bazel_cuda_h100_b200.yml

name: CI - Bazel H100 and B200 CUDA tests
# The workflow runs if any of the following conditions are met:
# - H100 and B200 on Workflow dispatch.
# - H100 and B200 on scheduled every two hours.
# - H100 and B200 on PR to main that has the 'CI Optional GPU Presubmit' label.
on:
  # Runs on PR if label "CI Optional GPU Presubmit" is present.
  workflow_dispatch:
    inputs:
      halt-for-connection:
        description: 'Should this workflow run wait for a remote connection?'
        type: choice
        required: true
        default: 'no'
        options:
        - 'yes'
        - 'no'
      python_version:
        description: 'Hermetic Python version for tests (for example, 3.14).'
        type: string
        required: true
        default: '3.14'
      xla_track:
        description: 'XLA revision used by tests.'
        type: choice
        required: true
        default: 'head'
        options:
        - 'pinned'
        - 'head'
        - 'commit'
      xla_commit:
        description: 'XLA commit SHA used by tests (only respected if xla_track is "commit").'
        type: string
        required: false
      jobs_to_run:
        description: 'Which jobs to run?'
        type: choice
        required: true
        default: 'all'
        options:
        - 'all'
        - 'single_b200'
        - 'multi_h100'
      bazel_targets:
        description: 'Custom Bazel targets to run (space-separated, overrides defaults)'
        type: string
        required: false
  pull_request:
    branches:
      - main
    types: [ labeled, synchronize, opened, reopened ]
  schedule:
    - cron: "0 */2 * * *" # Run every 2 hours
permissions: {}

# - For PRs, group by PR number and cancel stale in-progress runs.
# - For workflow_dispatch events, group by run_id so multiple dispatches can run in parallel.
# - Never cancel scheduled runs.
concurrency:
  group: >-
    ${{ github.workflow }}-
    ${{
        (github.event_name == 'pull_request'
         && format('pr-{0}', github.event.pull_request.number))

        || (github.event_name == 'workflow_dispatch'
            && format('dispatch-{0}', github.run_id))

        || format('ref-{0}', github.ref)
    }}
  cancel-in-progress: >-
    ${{
        github.ref != 'refs/heads/main' &&
        github.event_name == 'pull_request'
    }}

defaults:
  run:
    shell: bash

jobs:
  run_tests:
    if: >-
      ${{
          github.repository_owner == 'jax-ml'
          && github.event.repository.fork == false
          && (github.event_name == 'schedule'
              || (github.event_name == 'workflow_dispatch'
                  && (github.event.inputs.jobs_to_run == 'all'
                      || github.event.inputs.jobs_to_run == 'single_b200'))
              || contains(github.event.pull_request.labels.*.name,
                          'CI Optional GPU Presubmit'))
      }}
    runs-on: linux-x86-a4-224-b200-1gpu
    container: 'us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest'  # zizmor: ignore[unpinned-images]
    name: "Bazel single B200 CUDA tests"
    steps:
      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1  # v7.0.1
        with:
          persist-credentials: false
      - name: Wait For Connection
        uses: google-ml-infra/actions/ci_connection@7f5ca0c263a81ed09ea276524c1b9192f1304e3c
        with:
          halt-dispatch-input: ${{ inputs.halt-for-connection }}
      - name: Run Bazel single B200 CUDA Tests
        env:
          JAXCI_HERMETIC_PYTHON_VERSION: ${{ github.event.inputs.python_version || '3.14' }}
          JAXCI_XLA_TRACK: ${{ github.event.inputs.xla_track || 'head' }}
          JAXCI_XLA_COMMIT: ${{ github.event.inputs.xla_commit }}
          JAXCI_BAZEL_TARGETS: |
            //tests:cudnn_fusion_test_gpu
            //tests:scaled_matmul_stablehlo_test_gpu
            //tests:fused_attention_stablehlo_test_gpu
            //tests:nn_test_gpu
            //tests/pallas:gpu_tests
            //tests/mosaic:gpu_tests
          JAXCI_TEST_TAG_FILTERS: '-multiaccelerator'
          JAXCI_USE_PARALLEL_ACCELERATOR_RUNNER: '1'
          JAXCI_LOCAL_TEST_JOBS: '8'
          JAXCI_EXCLUDE_TEST_TARGETS: 'PmapTest.testSizeOverflow|.*InterpretTest.*'
          JAXCI_TEST_TIMEOUT: '420'
          CUSTOM_BAZEL_TARGETS: ${{ github.event.inputs.bazel_targets }}
        run: |
          if [[ -n "${CUSTOM_BAZEL_TARGETS}" ]]; then
            export JAXCI_BAZEL_TARGETS="${CUSTOM_BAZEL_TARGETS// /$'\n'}"
          fi
          bash ./ci/run_bazel_cuda_targeted_tests.sh
      - name: Upload test artifacts
        if: always()
        continue-on-error: true
        uses: ./.github/actions/upload-test-artifacts
        with:
          artifact-name: bazel-cuda-h100-single
          inputs_json: ${{ toJSON(inputs) }}
          matrix_json: ${{ toJSON(matrix) }}
  run_multiaccelerator_tests:
    if: >-
      ${{
          github.repository_owner == 'jax-ml'
          && github.event.repository.fork == false
          && (
              github.event_name == 'schedule'
              || (github.event_name == 'workflow_dispatch'
                  && (github.event.inputs.jobs_to_run == 'all'
                      || github.event.inputs.jobs_to_run == 'multi_h100'))
              || contains(github.event.pull_request.labels.*.name,
                          'CI Optional GPU Presubmit')
          )
      }}
    runs-on: linux-x86-a3-8g-h100-8gpu
    container: 'us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest'  # zizmor: ignore[unpinned-images]
    name: "Bazel multiple H100 CUDA tests"
    steps:
      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1  # v7.0.1
        with:
          persist-credentials: false
      - name: Wait For Connection
        uses: google-ml-infra/actions/ci_connection@9f8d7d8d9954430d0fef4a8cce2cb4e3844c200d
        with:
          halt-dispatch-input: ${{ inputs.halt-for-connection }}
      - name: Run Bazel multiple H100 CUDA Tests
        env:
          JAXCI_HERMETIC_PYTHON_VERSION: ${{ github.event.inputs.python_version || '3.14' }}
          JAXCI_XLA_TRACK: ${{ github.event.inputs.xla_track || 'head' }}
          JAXCI_XLA_COMMIT: ${{ github.event.inputs.xla_commit }}
          JAXCI_BAZEL_TARGETS: |
            //tests/mosaic:gpu_tests
            //tests/pallas:gpu_tests
            //tests:array_interoperability_test_gpu
            //tests:cudnn_fusion_test_gpu
            //tests:fused_attention_stablehlo_test_gpu
            //tests:gpu_tests
            //tests:python_callback_test_gpu
            //tests:ragged_collective_test_gpu
            //tests/multiprocess:gpu_tests
            //jax/experimental/jax2tf/tests/multiprocess:gpu_tests
          JAXCI_TEST_TAG_FILTERS: 'multiaccelerator'
          # Limit to 4 concurrent jobs to reduce multi-GPU resource contention (OOMs and
          # timeouts), since every test runs against all 8 H100s.
          JAXCI_LOCAL_TEST_JOBS: '4'
          JAXCI_EXCLUDE_TEST_TARGETS: 'PmapTest.testSizeOverflow|.*InterpretTest.*'
          CUSTOM_BAZEL_TARGETS: ${{ github.event.inputs.bazel_targets }}
        run: |
          if [[ -n "${CUSTOM_BAZEL_TARGETS}" ]]; then
            export JAXCI_BAZEL_TARGETS="${CUSTOM_BAZEL_TARGETS// /$'\n'}"
          fi
          bash ./ci/run_bazel_cuda_targeted_tests.sh
      - name: Upload test artifacts
        if: always()
        continue-on-error: true
        uses: ./.github/actions/upload-test-artifacts
        with:
          artifact-name: bazel-cuda-h100-multi
          inputs_json: ${{ toJSON(inputs) }}
          matrix_json: ${{ toJSON(matrix) }}