maybe?

.github/workflows/tpu_unit_tests.yaml

@@ -31,6 +31,5 @@ jobs:
         run: |
           export TPU_NAME=ci-run-${{ github.run_id }}
           eval "$(ssh-agent -s)"
-          ssh-add ~/.ssh/google_compute_engine
-          infra/babysit-tpu-vm.sh $TPU_NAME -z ${{ matrix.tpu-zone }} -t v4-8 --preemptible -s infra/helpers/setup-tpu-vm-tests.sh -b ${{ github.sha }} -- \
+          infra/babysit-tpu-vm.sh $TPU_NAME -z ${{ matrix.tpu-zone }} -t v4-8 --preemptible -s infra/helpers/setup-tpu-vm-tests.sh -b ${{ github.sha }} --retries 3 -- \
             bash levanter/infra/run.sh pytest tests

infra/babysit-tpu-vm.sh

@@ -59,6 +59,8 @@ CMD_ARGS_STR=$(printf ' %s' "${CMD_ARGS[@]}")
 CMD_ARGS_STR=${CMD_ARGS_STR:1}
 CMD_ARGS_STR="RUN_ID=${RUN_ID} ${CMD_ARGS_STR}"
 
+TRIES=0
+
 # check if the VM is running
 # if not, spin it up
 # if it is, just run the command
@@ -82,6 +84,13 @@ while true; do
         break
       else
         echo "Command failed"
+        TRIES=$((TRIES+1))
+        if [ $RETRIES -ge 0 ]; then
+          if [ $TRIES -ge $RETRIES ]; then
+            echo "Command failed $TRIES times, exiting"
+            break
+          fi
+        fi
       fi
     fi
   else

infra/helpers/parse-tpu-creation-args.sh

@@ -23,6 +23,7 @@ AUTODELETE=true
 SETUP_SCRIPT="$SCRIPT_DIR/helpers/setup-tpu-vm.sh"
 SUBNETWORK="default"
 USE_ALPHA=false
+RETRIES=-1  # how many times babysit-tpu-vm.sh should retry before giving up. -1 means infinite
 
 if [ -z "$GIT_BRANCH" ]; then
     GIT_BRANCH=$(git rev-parse --abbrev-ref HEAD)
@@ -86,6 +87,11 @@ while [[ $# -gt 0 ]]; do
       USE_ALPHA="true"
       shift # past argument
       ;;
+    --retries)
+      RETRIES="$2"
+      shift # past argument
+      shift # past value
+      ;;
     *)    # unknown option, assume it's the vm name if it doesn't start with a dash
       if [[ $1 == -* ]]; then
         echo "Error: unknown option $1" >&2