Disable GPU-specific IRFFT lowering test on non-GPU devices. #42334
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| 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) }} |