Demo PR to test connecting to a GitHub actions runner

.github/workflows/build_jaxlib.yml

@@ -0,0 +1,30 @@
+name: Build Jaxlib
+
+# dummy; a small change so that "wait for connection" detects the github labels applied on the PR
+on:
+  pull_request:
+    branches:
+      - main
+
+jobs:
+  build:
+    strategy:
+      matrix:
+        runner: ["arc-linux-x86-n2-64"]
+
+    runs-on: ${{ matrix.runner }}
+    container:
+      image: "index.docker.io/tensorflow/build@sha256:e042f39c18a01012349ad42ae69f6c9fb6e8cc912ddf362d47f6caf1f97f54c3"
+
+    env:
+      ENV_FILE: ci/envs/build_artifacts/jaxlib
+      JAXCI_USE_DOCKER: 0
+      JAXCI_USE_RBE: 1
+
+    steps:
+      - uses: actions/checkout@v3
+      # Halt for testing
+      - name: Wait For Connection
+        uses: ./actions/ci_connection/
+      - name: Run build script
+        run: ./ci/build_artifacts.sh

actions/ci_connection/notify_connection.py

@@ -14,6 +14,7 @@
 
 import time
 import threading
+import os
 import subprocess
 from multiprocessing.connection import Client
 
@@ -42,8 +43,22 @@ def timer(conn):
     timer_thread.start()
 
     print("Entering interactive bash session")
+
+    # Hard-coded for now for demo purposes.
+    next_command = "bash ci/build_artifacts.sh"
+    # Print the "next" commands to be run
+    # TODO: actually get this data from workflow files
+    print(f"The next command that would have run is:\n\n{next_command}")
+
+    # Set the hardcoded envs for testing purposes
+    # TODO: sync env vars
+    sub_env = os.environ.copy()
+    sub_env["ENV_FILE"] = "ci/envs/build_artifacts/jaxlib"
+    sub_env["JAXCI_USE_DOCKER"] = "0"
+    sub_env["JAXCI_USE_RBE"] = "1"
+
     # Enter interactive bash session
-    subprocess.run(["/bin/bash", "-i"])
+    subprocess.run(["/bin/bash", "-i"], env=sub_env)
 
     print("Exiting interactive bash session")
     with lock:

actions/ci_connection/wait_for_connection.py

@@ -20,9 +20,9 @@
 import sys
 
 last_time = time.time()
-timeout = 600  # 10 minutes for initial connection
+timeout = 1800  # 30 minutes for initial connection
 keep_alive_timeout = (
-  900  # 30 minutes for keep-alive if no closed message (allow for reconnects)
+  3600  # 30 minutes for keep-alive if no closed message (allow for reconnects)
 )
 
 
@@ -102,7 +102,7 @@ def timer():
 
   print("Googler connection only\nSee go/<insert final golink> for details")
   print(
-    f"Connection string: ml-actions-connect  --runner={host} --ns={ns} --loc={location} --cluster={cluster} --halt_directory={actions_path}"
+    f"Connection string: ml-actions-connect  --runner={host} --ns={ns} --loc={location} --cluster={cluster} --halt_directory={actions_path} --project=ml-velocity-actions-testing"
   )
 
   # Thread is running as a daemon so it will quit when the