-
Notifications
You must be signed in to change notification settings - Fork 2.9k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add GitHub action workflow for Bazel CUDA continuous tests
PiperOrigin-RevId: 705238265
- Loading branch information
1 parent
9fb2976
commit dc7a9b3
Showing
4 changed files
with
109 additions
and
3 deletions.
There are no files selected for viewing
This file contains 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
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,81 @@ | ||
# CI - Bazel CUDA tests (Non-RBE) | ||
# | ||
# This workflow runs the CUDA tests with Bazel. It can only be triggered by other workflows via | ||
# `workflow_call`. It is used by the `CI - Wheel Tests` workflows to run the Bazel CUDA tests. | ||
# | ||
# It consists of the following job: | ||
# run-tests: | ||
# - Downloads the jaxlib and CUDA artifacts from a GCS bucket. | ||
# - Executes the `run_bazel_test_cuda_non_rbe.sh` script, which performs the following actions: | ||
# - Installs the downloaded wheel artifacts. | ||
# - Runs the CUDA tests with Bazel. | ||
name: CI - Bazel CUDA tests (Non-RBE) | ||
|
||
on: | ||
workflow_call: | ||
inputs: | ||
runner: | ||
description: "Which runner should the workflow run on?" | ||
type: string | ||
required: true | ||
default: "linux-x86-n2-16" | ||
python: | ||
description: "Which python version to test?" | ||
type: string | ||
required: true | ||
default: "3.12" | ||
enable-x64: | ||
description: "Should x64 mode be enabled?" | ||
type: string | ||
required: true | ||
default: "0" | ||
gcs_download_uri: | ||
description: "GCS location URI from where the artifacts should be downloaded" | ||
required: true | ||
default: 'gs://general-ml-ci-transient/jax-github-actions/jax-fork/${{ github.workflow }}/${{ github.run_number }}/${{ github.run_attempt }}' | ||
type: string | ||
halt-for-connection: | ||
description: 'Should this workflow run wait for a remote connection?' | ||
type: boolean | ||
required: false | ||
default: false | ||
|
||
jobs: | ||
run-tests: | ||
runs-on: ${{ inputs.runner }} | ||
container: "us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/nosla-cuda12.3-cudnn9.1-ubuntu20.04-manylinux2014-multipython:latest" | ||
|
||
env: | ||
JAXCI_HERMETIC_PYTHON_VERSION: ${{ inputs.python }} | ||
JAXCI_ENABLE_X64: ${{ inputs.enable-x64 }} | ||
|
||
name: "Bazel single accelerator and multi-accelerator CUDA tests (${{ inputs.runner }}, Python ${{ inputs.python }}, x64=${{ inputs.enable-x64 }})" | ||
|
||
steps: | ||
- uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 | ||
- name: Set env vars for use in artifact download URL | ||
run: | | ||
os=$(uname -s | awk '{print tolower($0)}') | ||
arch=$(uname -m) | ||
# Get the major and minor version of Python. | ||
# E.g if JAXCI_HERMETIC_PYTHON_VERSION=3.10, then python_major_minor=310 | ||
python_major_minor=$(echo "$JAXCI_HERMETIC_PYTHON_VERSION" | tr -d '.') | ||
echo "OS=${os}" >> $GITHUB_ENV | ||
echo "ARCH=${arch}" >> $GITHUB_ENV | ||
echo "PYTHON_MAJOR_MINOR=${python_major_minor}" >> $GITHUB_ENV | ||
- name: Download the wheel artifacts from GCS | ||
run: >- | ||
mkdir -p $(pwd)/dist && | ||
gsutil -m cp -r "${{ inputs.gcs_download_uri }}/jaxlib*${PYTHON_MAJOR_MINOR}*${OS}*${ARCH}*.whl" $(pwd)/dist/ && | ||
gsutil -m cp -r "${{ inputs.gcs_download_uri }}/jax*cuda*plugin*${PYTHON_MAJOR_MINOR}*${OS}*${ARCH}*.whl" $(pwd)/dist/ && | ||
gsutil -m cp -r "${{ inputs.gcs_download_uri }}/jax*cuda*pjrt*${OS}*${ARCH}*.whl" $(pwd)/dist/ | ||
# Halt for testing | ||
- 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 (Non-RBE) | ||
timeout-minutes: 60 | ||
run: ./ci/run_bazel_test_cuda_non_rbe.sh |
This file contains 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
This file contains 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
This file contains 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