mirror of
https://github.com/ROCm/jax.git
synced 2025-04-15 19:36:06 +00:00
46 lines
1.3 KiB
YAML
46 lines
1.3 KiB
YAML
name: CI - Bazel CUDA tests (RBE)
|
|
|
|
on:
|
|
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'
|
|
pull_request:
|
|
branches:
|
|
- main
|
|
|
|
concurrency:
|
|
group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
|
|
cancel-in-progress: true
|
|
|
|
jobs:
|
|
run_tests:
|
|
if: github.event.repository.fork == false
|
|
strategy:
|
|
matrix:
|
|
runner: ["linux-x86-n2-16"]
|
|
enable-x_64: [1, 0]
|
|
|
|
runs-on: ${{ matrix.runner }}
|
|
container: 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build:latest'
|
|
|
|
env:
|
|
JAXCI_HERMETIC_PYTHON_VERSION: "3.12"
|
|
JAXCI_ENABLE_X64: ${{ matrix.enable-x_64 }}
|
|
|
|
name: "Bazel single accelerator CUDA tests (${{ matrix.runner }}, Python 3.12, x64=${{ matrix.enable-x_64 }})"
|
|
|
|
steps:
|
|
- uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
|
|
- name: Wait For Connection
|
|
uses: google-ml-infra/actions/ci_connection@main
|
|
with:
|
|
halt-dispatch-input: ${{ inputs.halt-for-connection }}
|
|
- name: Run Bazel CUDA Tests with RBE
|
|
run: ./ci/run_bazel_test_cuda_rbe.sh |