Create benchmarks.yml

Author: juliagmt-google

.github/workflows/benchmarks.yml

@@ -0,0 +1,63 @@
+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
+      # 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
+        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: Run Bazel GPU tests locally
+        run: ./ci/run_bazel_test_gpu_non_rbe.sh 
+      - 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
+