Create a workflow to run benchmarks #4

	name: Benchmarks

	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:
	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: Build jaxlib
	env:
	JAXCI_CLONE_MAIN_XLA: 1
	run: ./ci/build_artifacts.sh "jaxlib"
	# - name: Build jax-cuda-plugin
	# env:
	# JAXCI_CLONE_MAIN_XLA: 1
	# run: ./ci/build_artifacts.sh "jax-cuda-plugin"
	# - name: Build jax-cuda-pjrt
	# env:
	# JAXCI_CLONE_MAIN_XLA: 1
	# run: ./ci/build_artifacts.sh "jax-cuda-pjrt"
	- 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

Provide feedback