mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-17 12:26:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

Enable New bazel presubmits for pull requests.

Browse Source 
At this time these workflows will currently be non-blocking for submission.

PiperOrigin-RevId: 704825599
This commit is contained in:
Michael Hudgins 2024-12-10 13:25:18 -08:00 committed by jax authors
parent e6d6c4ef8a
commit ac92aaaf3a
2 changed files with 16 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
.github/workflows/bazel_cpu_rbe.yml vendored
View File

@ -11,6 +11,9 @@ on:
options:
- 'yes'
- 'no'
pull_request:
branches:
- main
concurrency:
group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
@ -21,15 +24,18 @@ jobs:
if: github.event.repository.fork == false
strategy:
matrix:
runner: ["linux-x86-n2-16", "linux-arm64-t2a-16"]
runner: ["linux-x86-n2-16", "linux-arm64-c4a-16"]
enable-x_64: [1, 0]
runs-on: ${{ matrix.runner }}
# TODO(b/369382309): Replace Linux Arm64 container with the ml-build container once it is available
container: ${{ (contains(matrix.runner, 'linux-x86') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build:latest') ||
(contains(matrix.runner, 'linux-arm64') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/linux-arm64-arc-container:latest') }}
(contains(matrix.runner, 'linux-arm64') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build-arm64:latest') }}
env:
JAXCI_HERMETIC_PYTHON_VERSION: "3.12"
JAXCI_ENABLE_X64: ${{ matrix.enable-x_64 }}
name: "Bazel CPU tests (${{ matrix.runner }}, Python 3.12, x64=${{ matrix.enable-x_64 }})"
steps:
- uses: actions/checkout@v3

7
.github/workflows/bazel_gpu_rbe.yml vendored
View File

@ -11,6 +11,9 @@ on:
options:
- 'yes'
- 'no'
pull_request:
branches:
- main
concurrency:
group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
@ -22,12 +25,16 @@ jobs:
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 GPU tests (${{ matrix.runner }}, Python 3.12, x64=${{ matrix.enable-x_64 }})"
steps:
- uses: actions/checkout@v3