Skip to content

Add new CI scripts as part of JAX CI Rework #187

Add new CI scripts as part of JAX CI Rework

Add new CI scripts as part of JAX CI Rework #187

Workflow file for this run

name: Run Pytest CPU tests
on:
pull_request:
branches:
- main
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'
jobs:
build:
continue-on-error: true
defaults:
run:
# Explicitly set the shell to bash to override the default Windows environment, i.e, cmd.
shell: bash
strategy:
matrix:
runner: ["windows-x86-n2-64", "linux-x86-n2-64", "linux-arm64-t2a-48"]
python: ["3.10"]
runs-on: ${{ matrix.runner }}
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, 'windows-x86') && null) }}
env:
JAXCI_CLONE_MAIN_XLA: 1
JAXCI_HERMETIC_PYTHON_VERSION: ${{ matrix.python }}
steps:
- uses: actions/checkout@v3
# Halt for testing
- name: Wait For Connection
uses: google-ml-infra/actions/ci_connection@main
with:
halt-dispatch-input: ${{ inputs.halt-for-connection }}
- name: Build jaxlib
run: ./ci/build_artifacts.sh "ci/envs/build_artifacts/jaxlib.env"
- name: Install pytest
env:
JAXCI_PYTHON: python${{ matrix.python }}
run: $JAXCI_PYTHON -m pip install pytest
- name: Install dependencies
env:
JAXCI_PYTHON: python${{ matrix.python }}
run: $JAXCI_PYTHON -m pip install --upgrade numpy=="2.0.0" scipy=="1.13.1" wheel build pytest-xdist absl-py opt-einsum colorama portpicker matplotlib 'importlib_metadata>=4.6' hypothesis flatbuffers filelock
- name: Run Pytest CPU tests
run: ./ci/run_pytest_cpu.sh