Update benchmarks.yml

Author: juliagmt-google

.github/workflows/benchmarks.yml

@@ -4,6 +4,16 @@ 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:
@@ -20,6 +30,18 @@ jobs:
 
     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
@@ -31,4 +53,3 @@ jobs:
       - name: Run Multiprocess GPU Test
         run: |
           python -m pytest tests/multiprocess_gpu_test.py
-