Skip to content

Create a workflow to run benchmarks #2

Create a workflow to run benchmarks

Create a workflow to run benchmarks #2

Workflow file for this run

name: Benchmarks
on:
pull_request:
branches:
- main
jobs:
build:
strategy:
matrix:
runner: ["linux-x86-g2-48-l4-4gpu"]
runs-on: ${{ matrix.runner }}
container:
image: "gcr.io/tensorflow-testing/nosla-cuda12.3-cudnn9.1-ubuntu20.04-manylinux2014-multipython:latest"
env:
JAXCI_HERMETIC_PYTHON_VERSION: 3.11
steps:
- uses: actions/checkout@v3
- name: Install dependencies
run: |
python -m pip install --upgrade pip
pip install pytest
pip install absl-py
# pip install -U jax
# pip install -U "jax[cuda12]"
pip install google-benchmark
- name: Run Multiprocess GPU Test
run: |
python -m pytest tests/multiprocess_gpu_test.py