Enable tpu7x daily tests to run on Spot VMs (#5072)

Author: simrankaurb

tools/cloud-build/daily-tests/builds/gke-tpu-7x.yaml

@@ -36,13 +36,26 @@ steps:
   env:
   - "ANSIBLE_HOST_KEY_CHECKING=false"
   - "ANSIBLE_CONFIG=/workspace/tools/cloud-build/ansible.cfg"
+  - "MACHINE_TYPE=tpu7x-standard-4t"
+  - "PROJECT_ID=$PROJECT_ID"
+  - "NUM_NODES=2"
+  - "INSTANCE_PREFIX=tpu7xsp"
+  - "BUILD_ID=$BUILD_ID"
+  - "OPTIONS_GCS_PATH=gs://hpc-ctk1357/tpu7xoptions.txt"
   args:
   - -c
   - |
-    set -x -e
+    set -e -u -o pipefail
+    echo "Sourcing find_available_zone.sh to determine zone."
+    source /workspace/tools/cloud-build/find_available_zone.sh
+    if [ -z "$${ZONE:-}" ]; then
+      echo "ERROR: ZONE not found" >&2
+      exit 1
+    fi
+    set -x
     cd /workspace && make
-    BUILD_ID_FULL=$BUILD_ID
-    BUILD_ID_SHORT=$${BUILD_ID_FULL:0:6}
+    REGION="$${ZONE%-*}"
+    BUILD_ID_SHORT="$${BUILD_ID:0:6}"
     EXAMPLE_BP=examples/gke-tpu-7x/gke-tpu-7x.yaml
     # adding vm to act as remote node
     echo '  - id: remote-node'                                                            >> $${EXAMPLE_BP}
@@ -64,6 +77,9 @@ steps:
     echo '      - python -c '\''import jax; print(jax.device_count(), "TPU cores")'\'''    >> $${EXAMPLE_BP}
     echo '      node_count: 1'                                                             >> $${EXAMPLE_BP}
     echo '    outputs: [instructions]'                                                     >> $${EXAMPLE_BP}
+    sed -i -e '/reservation_affinity:/,+3c\      spot: true' $${EXAMPLE_BP}
+    sed -i '/reservation/d' $${EXAMPLE_BP}
     ansible-playbook tools/cloud-build/daily-tests/ansible_playbooks/base-integration-test.yml \
         --user=sa_106486320838376751393 --extra-vars="project=${PROJECT_ID} build=$${BUILD_ID_SHORT}" \
+        --extra-vars="region=$${REGION} zone=$${ZONE}" \
         --extra-vars="@tools/cloud-build/daily-tests/tests/gke-tpu-7x.yml"

tools/cloud-build/daily-tests/tests/gke-tpu-7x.yml

@@ -18,11 +18,8 @@ deployment_name: gke-tpu-7x-{{ build }}
 workspace: /workspace
 blueprint_yaml: "{{ workspace }}/examples/gke-tpu-7x/gke-tpu-7x.yaml"
 network: "{{ deployment_name }}-net"
-region: us-central1
-zone: us-central1-c
 remote_node: "{{ deployment_name }}-remote-node-0"
 machine_type: tpu7x-standard-4t
-extended_reservation: cloudtpu-20251001223000-1287580847
 num_slices: 1
 tpu_topology: 2x2x1
 static_node_count: 1
@@ -34,9 +31,11 @@ cli_deployment_vars:
   tpu_topology: "{{ tpu_topology }}"
   static_node_count: "{{ static_node_count }}"
   authorized_cidr: "{{ build_ip.stdout }}/32"
-  reservation: "{{ extended_reservation }}"
 custom_vars:
   project: "{{ project }}"
   expected_tpu_count: 8
+  instance_labels:
+    tpu7x_onspot: true
+  enable_spot: true
 post_deploy_tests:
 - test-validation/test-gke-tpu.yml