Files
keras-team--keras/.github/workflows/tpu_tests.yml
T
2026-07-13 12:20:15 +08:00

88 lines
2.6 KiB
YAML

name: Keras TPU Tests
on:
push:
branches: [master]
pull_request:
types: [unlabeled]
release:
types: [created]
permissions:
contents: read
concurrency:
group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
cancel-in-progress: true
jobs:
test-in-container:
# Only run on pushes to master, releases or "kokoro:force-run" unlabel
if: |
github.event_name == 'push' ||
github.event_name == 'release' ||
(github.event.action == 'unlabeled' && github.event.label.name == 'kokoro:force-run')
strategy:
fail-fast: false
matrix:
include:
- backend: jax
multi_device: false
runner: linux-x86-ct6e-44-1tpu
name: Run tests on TPU
- backend: jax
multi_device: true
runner: linux-x86-ct5lp-112-4tpu
name: Run tests on multi-TPU
container:
image: python:3.11-slim
options: --privileged --network host
name: ${{ matrix.name }}
runs-on: ${{ matrix.runner }}
timeout-minutes: 120
env:
KERAS_HOME: .github/workflows/config/${{ matrix.backend }}
steps:
- name: Checkout ${{ github.ref }}
uses: actions/checkout@v6
- name: Install binary dependencies
# curl, git, gpg are for codecov/codecov-action
run: |
apt-get update
apt-get -y install curl
apt-get -y install git
apt-get -y install gnupg
- name: Install Python dependencies
run: pip install --no-cache-dir -r requirements-${{ matrix.backend }}-tpu.txt
- name: Verify JAX Installation
run: python3 -c "import jax; print('JAX devices:', jax.devices(), jax.devices()[0].device_kind); assert jax.default_backend() == 'tpu'"
- name: Run Tests
if: ${{ !matrix.multi_device }}
run: pytest keras --ignore keras/src/applications --cov=keras --cov-config=pyproject.toml
- name: Run Multi-device Tests
if: ${{ matrix.multi_device }}
run: pytest keras -m multi_device --cov=keras --cov-config=pyproject.toml --cov-append
- name: Convert Coverage
run: coverage xml --omit='keras/src/applications/*,keras/api' -o core-coverage.xml
- name: Upload Coverage
uses: codecov/codecov-action@e79a6962e0d4c0c17b229090214935d2e33f8354 # v5
with:
env_vars: KERAS_HOME
flags: keras,keras-tpu,keras-${{ matrix.backend }}${{ matrix.multi_device == 'true' && '-multi' || '' }}
files: core-coverage.xml
token: ${{ secrets.CODECOV_TOKEN }}
fail_ci_if_error: false