Skip to content

Commit 636dcf9

Browse files
simrankaurbsimrankaurb
authored
feat: Add TPU support to zone finding script (GoogleCloudPlatform#5006)
2 parents f274872 + ce97100 commit 636dcf9

3 files changed

Lines changed: 136 additions & 52 deletions

File tree

‎tools/cloud-build/daily-tests/builds/slurm-gcp-v6-tpu.yaml‎

Lines changed: 23 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -35,14 +35,33 @@ steps:
3535
env:
3636
- "ANSIBLE_HOST_KEY_CHECKING=false"
3737
- "ANSIBLE_CONFIG=/workspace/tools/cloud-build/ansible.cfg"
38+
- "MACHINE_TYPE=tpu"
39+
- "TPU_RUNTIME_VERSION=tpu-vm-tf-2.14.0"
40+
- "ACCELERATOR_TYPE=v3-8"
41+
- "PROJECT_ID=$PROJECT_ID"
42+
- "NUM_NODES=2"
43+
- "INSTANCE_PREFIX=v3sp"
44+
- "BUILD_ID=$BUILD_ID"
45+
- "OPTIONS_GCS_PATH=gs://hpc-ctk1357/tpuv3options.txt"
3846
args:
3947
- -c
4048
- |
41-
set -x -e
49+
set -e -u -o pipefail
50+
echo "Sourcing find_available_zone.sh to determine zone."
51+
source /workspace/tools/cloud-build/find_available_zone.sh
52+
if [ -z "$${ZONE:-}" ]; then
53+
echo "ERROR: ZONE not found" >&2
54+
exit 1
55+
fi
56+
set -x
4257
cd /workspace && make
43-
BUILD_ID_FULL=$BUILD_ID
44-
BUILD_ID_SHORT=$${BUILD_ID_FULL:0:6}
58+
REGION="$${ZONE%-*}"
59+
BUILD_ID_SHORT="$${BUILD_ID:0:6}"
60+
BLUEPRINT="/workspace/community/examples/hpc-slurm6-tpu.yaml"
61+
sed -i 's/^[ ]*preemptible: false/ preemptible: true/' $${BLUEPRINT}
4562
4663
ansible-playbook tools/cloud-build/daily-tests/ansible_playbooks/slurm-integration-test.yml \
47-
--user=sa_106486320838376751393 --extra-vars="project=${PROJECT_ID} build=$${BUILD_ID_SHORT}" \
64+
--user=sa_106486320838376751393 \
65+
--extra-vars="project=${PROJECT_ID} build=$${BUILD_ID_SHORT}" \
66+
--extra-vars="region=$${REGION} zone=$${ZONE}" \
4867
--extra-vars="@tools/cloud-build/daily-tests/tests/slurm-v6-tpu.yml"

‎tools/cloud-build/daily-tests/tests/slurm-v6-tpu.yml‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,10 +21,9 @@ deployment_name: "v6-tpu-{{ build }}"
2121
slurm_cluster_name: "v6tpu{{ build[0:5] }}"
2222

2323
cli_deployment_vars:
24-
region: us-central1
25-
zone: us-central1-a
24+
region: "{{ region }}"
25+
zone: "{{ zone }}"
2626

27-
zone: us-central1-a
2827
workspace: /workspace
2928
blueprint_yaml: "{{ workspace }}/community/examples/hpc-slurm6-tpu.yaml"
3029
network: "{{ deployment_name }}-net"

‎tools/cloud-build/find_available_zone.sh‎

Lines changed: 111 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ generate_instance_names() {
2727
done
2828
}
2929

30-
cleanup_instances() {
30+
cleanup_vm_instances() {
3131
local project=$1
3232
local zone=$2
3333
local prefix=$3
@@ -49,6 +49,24 @@ cleanup_instances() {
4949
fi
5050
}
5151

52+
cleanup_tpu_nodes() {
53+
local project=$1
54+
local zone=$2
55+
# Shift out project and zone
56+
shift 2
57+
local tpu_names=("$@")
58+
59+
if [[ "${#tpu_names[@]}" -eq 0 ]]; then
60+
return
61+
fi
62+
63+
for tpu_name in "${tpu_names[@]}"; do
64+
if ! DELETE_OUTPUT=$(gcloud compute tpus tpu-vm delete "${tpu_name}" --project="${project}" --zone="${zone}" --quiet 2>&1); then
65+
echo "ERROR IN DELETING TPU RESOURCE ${tpu_name}: ${DELETE_OUTPUT}" >&2
66+
fi
67+
done
68+
}
69+
5270
if ! GCS_CONTENT=$(gcloud storage cat "${OPTIONS_GCS_PATH}"); then
5371
echo "ERROR: Failed to read ${OPTIONS_GCS_PATH}." >&2
5472
exit 1
@@ -60,7 +78,6 @@ while IFS= read -r line; do
6078
ZONES_ARRAY+=("${line}")
6179
fi
6280
done <<<"${GCS_CONTENT}"
63-
6481
if [[ "${#ZONES_ARRAY[@]}" -eq 0 ]]; then
6582
echo "ERROR: No valid zones found in ${OPTIONS_GCS_PATH}" >&2
6683
exit 1
@@ -71,55 +88,104 @@ SUCCESS=false
7188

7289
# Loop through all zones to find capacity
7390
for ZONE in "${ZONES_ARRAY[@]}"; do
74-
readarray -t INSTANCE_NAMES_ARRAY < <(generate_instance_names "${FULL_INSTANCE_PREFIX}" "${NUM_NODES}")
75-
# Join instance names with commas for the --predefined-names flag
76-
instance_names_str=$(
77-
IFS=,
78-
echo "${INSTANCE_NAMES_ARRAY[*]}"
79-
)
80-
81-
declare -a GCLOUD_CMD
82-
GCLOUD_CMD=(
83-
gcloud compute instances bulk create
84-
--predefined-names="${instance_names_str}"
85-
--project="${PROJECT_ID}"
86-
--zone="${ZONE}"
87-
--machine-type="${MACHINE_TYPE}"
88-
--provisioning-model="${PROVISIONING_MODEL}"
89-
--instance-termination-action="${TERMINATION_ACTION}"
90-
--no-address
91-
--quiet
92-
--min-count="${MIN_NODES}"
93-
)
94-
if CREATE_OUTPUT=$("${GCLOUD_CMD[@]}" 2>&1); then
95-
instance_list_output=$(gcloud compute instances list --project="${PROJECT_ID}" --zones="${ZONE}" \
96-
--filter="name ~ ^${FULL_INSTANCE_PREFIX}" --format='value(name)')
97-
98-
if [[ $? -eq 0 && -n "${instance_list_output}" ]]; then
99-
readarray -t created_instances < <(echo "${instance_list_output}")
100-
NUM_CREATED=$((${#created_instances[@]}))
101-
echo "${ZONE} ${NUM_CREATED}"
102-
cleanup_instances "${PROJECT_ID}" "${ZONE}" "${FULL_INSTANCE_PREFIX}"
103-
104-
if [[ "${NUM_CREATED}" -ge "${MIN_NODES}" ]]; then
105-
SELECTED_ZONE="${ZONE}"
106-
SUCCESS=true
107-
break
91+
if [[ "${MACHINE_TYPE}" == "tpu" ]]; then
92+
# --- TPU Capacity Check ---
93+
for i in $(seq 1 "${NUM_NODES}"); do
94+
TPU_NAME="${FULL_INSTANCE_PREFIX}$(printf "%02d" "$i")"
95+
declare -a GCLOUD_TPU_CMD
96+
GCLOUD_TPU_CMD=(
97+
gcloud compute tpus tpu-vm create "${TPU_NAME}"
98+
--project="${PROJECT_ID}"
99+
--zone="${ZONE}"
100+
--accelerator-type="${ACCELERATOR_TYPE}"
101+
--version="${TPU_RUNTIME_VERSION}"
102+
--spot
103+
--quiet
104+
)
105+
106+
if ! TPU_CREATE_OUTPUT=$("${GCLOUD_TPU_CMD[@]}" 2>&1); then
107+
echo "ERROR: Unexpected error during TPU create in ${ZONE}: ${TPU_CREATE_OUTPUT}" >&2
108+
fi
109+
done
110+
111+
declare -a CREATED_TPU_NAMES=()
112+
TPU_CREATED_COUNT=0
113+
if tpu_list_output=$(gcloud compute tpus tpu-vm list --project="${PROJECT_ID}" --zone="${ZONE}" \
114+
--filter="name ~ ${FULL_INSTANCE_PREFIX}" --format='value(name)' 2>/dev/null); then
115+
if [[ -n "${tpu_list_output}" ]]; then
116+
readarray -t CREATED_TPU_NAMES <<<"${tpu_list_output}"
117+
TPU_CREATED_COUNT=${#CREATED_TPU_NAMES[@]}
118+
echo "INFO: Found ${TPU_CREATED_COUNT} TPU nodes: ${CREATED_TPU_NAMES[*]}"
108119
else
109-
echo "ERROR: Bulk create succeeded but only ${NUM_CREATED} instances found, less than min_count ${MIN_NODES} in ${ZONE}." >&2
120+
echo "INFO: No matching TPU nodes found in ${ZONE} via list."
110121
fi
111122
else
112-
echo "ERROR: Bulk create command succeeded in ${ZONE}, but failed to list the created instances or none found." >&2
113-
cleanup_instances "${PROJECT_ID}" "${ZONE}" "${FULL_INSTANCE_PREFIX}"
123+
echo "ERROR: Failed to list TPU nodes in ${ZONE}: ${tpu_list_output}" >&2
114124
fi
125+
126+
if [[ "${TPU_CREATED_COUNT}" -ge 1 ]]; then
127+
SELECTED_ZONE="${ZONE}"
128+
SUCCESS=true
129+
fi
130+
131+
cleanup_tpu_nodes "${PROJECT_ID}" "${ZONE}" "${CREATED_TPU_NAMES[@]}"
115132
else
116-
# Command failed
117-
if [[ "${CREATE_OUTPUT}" != *"INSUFFICIENT_CAPACITY"* &&
118-
"${CREATE_OUTPUT}" != *"ZONE_RESOURCE_POOL_EXHAUSTED"* ]]; then
119-
echo "ERROR: Unexpected error during bulk create in ${ZONE}: ${CREATE_OUTPUT}" >&2
133+
readarray -t INSTANCE_NAMES_ARRAY < <(generate_instance_names "${FULL_INSTANCE_PREFIX}" "${NUM_NODES}")
134+
instance_names_str=$(
135+
IFS=,
136+
echo "${INSTANCE_NAMES_ARRAY[*]}"
137+
)
138+
139+
declare -a GCLOUD_CMD
140+
GCLOUD_CMD=(
141+
gcloud compute instances bulk create
142+
--predefined-names="${instance_names_str}"
143+
--project="${PROJECT_ID}"
144+
--zone="${ZONE}"
145+
--machine-type="${MACHINE_TYPE}"
146+
--provisioning-model="${PROVISIONING_MODEL}"
147+
--instance-termination-action="${TERMINATION_ACTION}"
148+
--no-address
149+
--quiet
150+
--min-count="${MIN_NODES}"
151+
)
152+
echo "INFO: Attempting to bulk create ${NUM_NODES} VMs in ${ZONE}..."
153+
if CREATE_OUTPUT=$("${GCLOUD_CMD[@]}" 2>&1); then
154+
if instance_list_output=$(gcloud compute instances list --project="${PROJECT_ID}" --zones="${ZONE}" \
155+
--filter="name ~ ^${FULL_INSTANCE_PREFIX}" --format='value(name)'); then
156+
if [[ -n "${instance_list_output}" ]]; then
157+
readarray -t created_instances < <(echo "${instance_list_output}")
158+
NUM_CREATED=$((${#created_instances[@]}))
159+
cleanup_vm_instances "${PROJECT_ID}" "${ZONE}" "${FULL_INSTANCE_PREFIX}"
160+
161+
if [[ "${NUM_CREATED}" -ge "${MIN_NODES}" ]]; then
162+
SELECTED_ZONE="${ZONE}"
163+
SUCCESS=true
164+
echo "INFO: Found sufficient VM capacity in ${ZONE}."
165+
else
166+
echo "ERROR: Bulk create & list succeeded in ${ZONE}, but only ${NUM_CREATED} instances found, less than min_count ${MIN_NODES}." >&2
167+
fi
168+
else
169+
echo "ERROR: Bulk create command apparently succeeded in ${ZONE}, but LIST command found no instances with the prefix." >&2
170+
cleanup_vm_instances "${PROJECT_ID}" "${ZONE}" "${FULL_INSTANCE_PREFIX}"
171+
fi
172+
else
173+
echo "ERROR: Bulk create command succeeded in ${ZONE}, but the command to list instances failed." >&2
174+
cleanup_vm_instances "${PROJECT_ID}" "${ZONE}" "${FULL_INSTANCE_PREFIX}"
175+
fi
176+
else
177+
if [[ "${CREATE_OUTPUT}" != *"INSUFFICIENT_CAPACITY"* &&
178+
"${CREATE_OUTPUT}" != *"ZONE_RESOURCE_POOL_EXHAUSTED"* ]]; then
179+
echo "ERROR: Unexpected error during bulk create in ${ZONE}: ${CREATE_OUTPUT}" >&2
180+
else
181+
echo "INFO: Insufficient VM capacity for bulk create in ${ZONE}."
182+
fi
183+
cleanup_vm_instances "${PROJECT_ID}" "${ZONE}" "${FULL_INSTANCE_PREFIX}"
120184
fi
121-
# bulk create with --min-count should handle rollback on failure, but cleanup just in case.
122-
cleanup_instances "${PROJECT_ID}" "${ZONE}" "${FULL_INSTANCE_PREFIX}"
185+
fi
186+
187+
if [[ "${SUCCESS}" == "true" ]]; then
188+
break
123189
fi
124190
done
125191

0 commit comments

Comments
 (0)