Skip to content

Commit 591a00a

Browse files
simrankaurbsimrankaurb
authored
Enable Spot VM testing for a3-ultragpu instances (#4862)
2 parents 39ff641 + 0676a82 commit 591a00a

6 files changed

Lines changed: 381 additions & 1 deletion

File tree

‎tools/cloud-build/daily-tests/ansible_playbooks/slurm-integration-test.yml‎

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -209,7 +209,7 @@
209209
- name: Wait until Slurm key exists
210210
ansible.builtin.wait_for:
211211
path: /etc/slurm/slurm.key
212-
timeout: 600 # Waits for up to 10 minutes for the file to appear
212+
timeout: 600 # Waits for up to 10 minutes for the file to appear
213213
when: key_type == 'slurm'
214214

215215
- name: Count Slurm nodes
@@ -223,6 +223,10 @@
223223
retries: 60
224224
delay: 15
225225

226+
- name: Add instance labels for spot instances
227+
ansible.builtin.include_tasks: tasks/add-instance-labels.yml
228+
when: custom_vars is defined and custom_vars.enable_spot | default(false) | bool
229+
226230
- name: Run Integration tests for Cluster Toolkit
227231
ansible.builtin.include_tasks: "{{ test }}"
228232
vars:
@@ -232,6 +236,19 @@
232236
loop_control:
233237
loop_var: test
234238

239+
rescue:
240+
- name: Display test failure message
241+
ansible.builtin.debug:
242+
msg: "A task within the integration tests failed. Conditional preemption check will run for spot instances."
243+
244+
- name: Check for recent preemptions for spot instances
245+
ansible.builtin.include_tasks: test-validation/test-preemption.yml
246+
when: custom_vars is defined and custom_vars.enable_spot | default(false) | bool
247+
248+
- name: Propagate original failure after rescue tasks
249+
ansible.builtin.fail:
250+
msg: "Integration tests failed. Rescue tasks were executed."
251+
235252
## Always cleanup, even on failure
236253
always:
237254
- name: Ensure all nodes are powered down

‎tools/cloud-build/daily-tests/ansible_playbooks/tasks/add-instance-labels.yml‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
---
15+
- name: Check if instance_labels are provided and are a dictionary
16+
ansible.builtin.assert:
17+
that:
18+
- custom_vars.instance_labels is defined
19+
- custom_vars.instance_labels is mapping
20+
- custom_vars.instance_labels | length > 0
21+
fail_msg: "custom_vars.instance_labels must be a non-empty dictionary."
22+
23+
# Creates label string of the form "key1=value1,key2=value2" to add to the instances created during the test.
24+
# Also, transforms key and value to lowercase to be compliant with the gcloud command.
25+
- name: Prepare label string
26+
ansible.builtin.set_fact:
27+
labels_string: "{{ keys | zip(values) | map('join', '=') | join(',') }}"
28+
vars:
29+
items: "{{ custom_vars.instance_labels.items() }}"
30+
keys: "{{ items | map(attribute=0) | map('string') | map('lower') }}"
31+
values: "{{ items | map(attribute=1) | map('string') | map('lower') }}"
32+
33+
- name: Gather instance information
34+
ansible.builtin.import_tasks: tasks/get_instance_ids.yml
35+
36+
- name: Parse instance list
37+
ansible.builtin.set_fact:
38+
gce_instances_list: "{{ instances.stdout | default('[]') | from_json }}"
39+
40+
- name: Report if no instances found
41+
ansible.builtin.debug:
42+
msg: "No instances found with label ghpc_deployment={{ deployment_name }} in project {{ project }}."
43+
when: gce_instances_list | length == 0
44+
45+
- name: Add spot label to each node
46+
ansible.builtin.command: >-
47+
gcloud compute instances update {{ item.name }}
48+
--project={{ project }}
49+
--zone={{ item.zone | basename }}
50+
--update-labels={{ labels_string }}
51+
loop: "{{ gce_instances_list }}"
52+
loop_control:
53+
label: "{{ item.name }}"
54+
when: gce_instances_list | length > 0

‎tools/cloud-build/daily-tests/ansible_playbooks/test-validation/test-preemption.yml‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
---
16+
- name: Populate instance list if empty or undefined
17+
when: gce_instances_list | default([]) | length == 0
18+
block:
19+
- name: Gather instance information
20+
ansible.builtin.import_tasks: tasks/get_instance_ids.yml
21+
22+
- name: Parse instance list
23+
ansible.builtin.set_fact:
24+
gce_instances_list: "{{ instances.stdout | default('[]') | from_json }}"
25+
26+
- name: Check for Preemption Events
27+
when: gce_instances_list | length > 0
28+
block:
29+
# Creating a filter string of the form "resource.labels.instance_id=id1 OR resource.labels.instance_id=id2"
30+
- name: Build instance ID filter string
31+
ansible.builtin.set_fact:
32+
instance_id_filter: "{{ gce_instances_list | map(attribute='id') | map('regex_replace', '^(.*)$', 'resource.labels.instance_id=\"\\1\"') | join(' OR ') }}"
33+
34+
- name: Query Cloud Audit Logs for any preemption events for these nodes
35+
ansible.builtin.shell: |
36+
gcloud logging read '
37+
resource.type="gce_instance"
38+
AND protoPayload.methodName="compute.instances.preempted"
39+
AND log_id("cloudaudit.googleapis.com/system_event")
40+
AND ({{ instance_id_filter }})
41+
' --project={{ project }} --freshness=4h --format="json"
42+
register: preemption_logs
43+
44+
- name: Display Instance Names being checked
45+
ansible.builtin.debug:
46+
msg: "Instance Names being checked for preemption in {{ deployment_name }}: {{ gce_instances_list | map(attribute='name') | list }}"
47+
48+
# Displays preempted nodes if any, else, gives message that no preemption events found
49+
- name: Report preemption status
50+
vars:
51+
preemption_logs_parsed: "{{ preemption_logs.stdout | from_json }}"
52+
ansible.builtin.debug:
53+
msg: >-
54+
{% if preemption_logs_parsed | length > 0 %}
55+
Preemption Events Found in the last 4 hours: {{ preemption_logs_parsed }}
56+
{% else %}
57+
No preemption events found for the checked instances in {{ deployment_name }} in the last 4 hours.
58+
{% endif %}

‎tools/cloud-build/daily-tests/builds/ml-a3-ultragpu-onspot-slurm.yaml‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
---
16+
tags:
17+
- m.custom-image
18+
- m.cloud-storage-bucket
19+
- m.pre-existing-network-storage
20+
- m.filestore
21+
- m.gpu-rdma-vpc
22+
- m.schedmd-slurm-gcp-v6-controller
23+
- m.schedmd-slurm-gcp-v6-login
24+
- m.schedmd-slurm-gcp-v6-nodeset
25+
- m.schedmd-slurm-gcp-v6-partition
26+
- m.startup-script
27+
- m.vpc
28+
- slurm6
29+
30+
timeout: 14400s # 4hr
31+
steps:
32+
# While using static network names we are gaurding against more than 1 instance running at a time (for multi-group tests)
33+
- id: check_for_running_build
34+
name: gcr.io/cloud-builders/gcloud
35+
script: "tools/cloud-build/check_running_build.sh tools/cloud-build/daily-tests/builds/ml-a3-ultragpu-onspot-slurm.yaml"
36+
37+
- id: ml-a3-ultragpu-onspot-slurm
38+
name: us-central1-docker.pkg.dev/$PROJECT_ID/hpc-toolkit-repo/test-runner
39+
entrypoint: /bin/bash
40+
env:
41+
- "ANSIBLE_HOST_KEY_CHECKING=false"
42+
- "ANSIBLE_CONFIG=/workspace/tools/cloud-build/ansible.cfg"
43+
- "PROJECT_ID=$PROJECT_ID"
44+
- "NUM_NODES=4"
45+
- "IMAGE_PROJECT=ubuntu-os-accelerator-images"
46+
- "IMAGE_NAME=ubuntu-accelerator-2204-amd64-with-nvidia-570-v20250722"
47+
- "MACHINE_TYPE=a3-ultragpu-8g"
48+
- "INSTANCE_PREFIX=a3usp"
49+
- "BUILD_ID=$BUILD_ID"
50+
- "OPTIONS_GCS_PATH=gs://hpc-ctk1357/a3uoptions.txt"
51+
args:
52+
- -c
53+
- |
54+
set -e -u -o pipefail
55+
echo "Sourcing find_available_zone.sh to determine zone."
56+
source /workspace/tools/cloud-build/find_available_zone.sh
57+
if [ -z "$${ZONE:-}" ]; then
58+
echo "ERROR: ZONE not found" >&2
59+
exit 1
60+
fi
61+
set -x
62+
cd /workspace && make
63+
REGION="$${ZONE%-*}"
64+
BUILD_ID_SHORT="$${BUILD_ID:0:6}"
65+
BLUEPRINT="/workspace/examples/machine-learning/a3-ultragpu-8g/a3ultra-slurm-blueprint.yaml"
66+
sed -i -e '/deletion_protection:/{n;s/enabled: true/enabled: false/}' $${BLUEPRINT}
67+
sed -i -e '/reason:/d' $${BLUEPRINT}
68+
ansible-playbook tools/cloud-build/daily-tests/ansible_playbooks/slurm-integration-test.yml \
69+
--user=sa_106486320838376751393 \
70+
--extra-vars="project=$${PROJECT_ID} build=$${BUILD_ID_SHORT}" \
71+
--extra-vars="region=$${REGION} zone=$${ZONE}" \
72+
--extra-vars="@tools/cloud-build/daily-tests/tests/ml-a3-ultragpu-onspot-slurm.yml"

‎tools/cloud-build/daily-tests/tests/ml-a3-ultragpu-onspot-slurm.yml‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
---
16+
17+
# region, zone must be defined in build file with --extra-vars flag!
18+
test_name: a3u-onspot-slurm
19+
deployment_name: a3u-onspot-slurm-{{ build }}
20+
slurm_cluster_name: "a3usp{{ build[0:5] }}"
21+
workspace: /workspace
22+
blueprint_yaml: "{{ workspace }}/examples/machine-learning/a3-ultragpu-8g/a3ultra-slurm-blueprint.yaml"
23+
login_node: "{{ slurm_cluster_name }}-slurm-login-*"
24+
controller_node: "{{ slurm_cluster_name }}-controller"
25+
network: "{{ test_name }}-net-0"
26+
post_deploy_tests:
27+
- test-validation/test-mounts.yml
28+
- test-validation/test-partitions.yml
29+
- test-validation/test-default-partition.yml
30+
- test-validation/test-enroot.yml
31+
- test-validation/test-gpus-slurm.yml
32+
post_destroy_tasks:
33+
- post-destroy-tasks/delete-image.yml
34+
custom_vars:
35+
gpu_count: 8
36+
gpu_partition: a3ultra
37+
test_persistenced: true
38+
partitions:
39+
- a3ultra
40+
mounts:
41+
- /home
42+
- /gcs
43+
instance_labels:
44+
a3ultra_onspot: true
45+
enable_spot: true
46+
cli_deployment_vars:
47+
region: "{{ region }}"
48+
zone: "{{ zone }}"
49+
slurm_cluster_name: "{{ slurm_cluster_name }}"
50+
disk_size_gb: 200
51+
a3u_cluster_size: 2
52+
base_network_name: "{{ test_name }}"
53+
a3u_enable_spot_vm: true

0 commit comments

Comments
 (0)