Skip to content

Add a workflow to run benchmarks #3

Add a workflow to run benchmarks

Add a workflow to run benchmarks #3

Workflow file for this run

name: Benchmarks
on:
push:
branches:
- main
pull_request:
branches:
- main
jobs:
multiprocess_gpu_test:
runs-on: ubuntu-latest # For self-hosted runners with GPUs, update accordingly.
strategy:
matrix:
gpu-count: [2, 4, 8] # Adjust based on the runner's available GPUs
steps:
- uses: actions/checkout@v3
- name: Install dependencies
run: |
python -m pip install --upgrade pip
pip install -U jax
pip install -U "jax[cuda12]"
- name: Run Multiprocess GPU Test
run: |
export CUDA_VISIBLE_DEVICES="0-$((${matrix.gpu-count} - 1))"
python -m pytest -n ${matrix.gpu-count} tests/multiprocess_gpu_test.py