88 lines
2.6 KiB
YAML
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
|