Use the updated ci connect

Author: nitins17

.github/workflows/bazel_cpu_rbe.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:
   run_bazel_rbe_cpu_tests:
@@ -31,6 +41,8 @@ jobs:
       - uses: actions/checkout@v3
       # Halt for testing
       - name: Wait For Connection
-        uses: ./actions/ci_connection/
+        uses: google-ml-infra/actions/ci_connection@main
+        with:
+          halt-dispatch-input: ${{ inputs.halt-for-connection }}
       - name: Run Bazel CPU Tests
         run: ./ci/run_bazel_test_cpu_rbe.sh

.github/workflows/bazel_gpu_non_rbe.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:
@@ -25,7 +35,9 @@ jobs:
       - uses: actions/checkout@v3
       # Halt for testing
       - name: Wait For Connection
-        uses: ./actions/ci_connection/
+        uses: google-ml-infra/actions/ci_connection@main
+        with:
+          halt-dispatch-input: ${{ inputs.halt-for-connection }}
       - name: Build jaxlib
         run: ./ci/build_artifacts.sh "ci/envs/build_artifacts/jaxlib.env"
       - name: Build jax-cuda-plugin

.github/workflows/bazel_gpu_rbe.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:
@@ -25,6 +35,8 @@ jobs:
       - uses: actions/checkout@v3
       # Halt for testing
       - name: Wait For Connection
-        uses: ./actions/ci_connection/
+        uses: google-ml-infra/actions/ci_connection@main
+        with:
+          halt-dispatch-input: ${{ inputs.halt-for-connection }}
       - name: Run Bazel GPU tests using RBE
         run: ./ci/run_bazel_test_gpu_rbe.sh

.github/workflows/build_artifacts.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'
   workflow_call:
 
 jobs:

.github/workflows/pytest_cpu.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:
@@ -32,7 +42,9 @@ jobs:
       - uses: actions/checkout@v3
       # Halt for testing
       - name: Wait For Connection
-        uses: ./actions/ci_connection/
+        uses: google-ml-infra/actions/ci_connection@main
+        with:
+          halt-dispatch-input: ${{ inputs.halt-for-connection }}
       - name: Build jaxlib
         run: ./ci/build_artifacts.sh "ci/envs/build_artifacts/jaxlib.env"
       - name: Install pytest

.github/workflows/pytest_gpu.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:
@@ -25,7 +35,9 @@ jobs:
       - uses: actions/checkout@v3
       # Halt for testing
       - name: Wait For Connection
-        uses: ./actions/ci_connection/
+        uses: google-ml-infra/actions/ci_connection@main
+        with:
+          halt-dispatch-input: ${{ inputs.halt-for-connection }}
       - name: Build jaxlib
         run: ./ci/build_artifacts.sh "ci/envs/build_artifacts/jaxlib.env"
       - name: Build jax-cuda-plugin

.github/workflows/pytest_tpu.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:
   run_tests: