Skip to content

Commit 92f1d92

Browse files
vikramvs-ggvikramvs-gg
authored
Merge pull request GoogleCloudPlatform#4828 from vikramvs-gg/a4x_intgr_test
Adding integration test for GKE A4X
2 parents c17a052 + 24dbf41 commit 92f1d92

5 files changed

Lines changed: 568 additions & 8 deletions

File tree

‎tools/cloud-build/daily-tests/ansible_playbooks/test-validation/test-gke-a4x-kueue-config.yml‎

Lines changed: 243 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,243 @@
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: Kueue Configuration and Verification Tasks
16+
block:
17+
- name: Assert variables are defined
18+
ansible.builtin.assert:
19+
that:
20+
- region is defined
21+
- custom_vars.project is defined
22+
23+
- name: Get cluster credentials for kubectl
24+
delegate_to: localhost
25+
ansible.builtin.command: gcloud container clusters get-credentials {{ deployment_name }} --region {{ region }} --project {{ custom_vars.project }} --verbosity=debug
26+
27+
- name: Create the host topology kueue job
28+
delegate_to: localhost
29+
ansible.builtin.shell: |
30+
TEMPLATE_PATH="{{ workspace }}/tools/cloud-build/daily-tests/blueprints/kueue-config-files/host-topology-tas-small-job.yaml"
31+
OUTPUT_PATH="/tmp/host-topology-job.yaml"
32+
sed -e 's|image:.*|image: busybox|' \
33+
-e 's|args:.*|command: ["sleep", "10s"]|' \
34+
-e "s/tas-user-queue/{{ instance_type }}/g" \
35+
"${TEMPLATE_PATH}" > "${OUTPUT_PATH}"
36+
kubectl create -f "${OUTPUT_PATH}" -v=9
37+
args:
38+
executable: /bin/bash
39+
changed_when: False
40+
41+
- name: Verify Kueue Topology a4x-default exists with correct levels
42+
delegate_to: localhost
43+
ansible.builtin.shell: |
44+
# Get the YAML for the Topology and check for specific lines
45+
kubectl get topology a4x-default -o yaml | grep -q -E "nodeLabel: cloud.google.com/gce-topology-block" &&
46+
kubectl get topology a4x-default -o yaml | grep -q -E "nodeLabel: cloud.google.com/gce-topology-subblock" &&
47+
kubectl get topology a4x-default -o yaml | grep -q -E "nodeLabel: cloud.google.com/gke-nodepool" &&
48+
kubectl get topology a4x-default -o yaml | grep -q -E "nodeLabel: cloud.google.com/gce-topology-host" &&
49+
kubectl get topology a4x-default -o yaml | grep -q -E "nodeLabel: kubernetes.io/hostname"
50+
register: topology_validation
51+
until: topology_validation.rc == 0
52+
retries: 5
53+
delay: 5
54+
- name: Debug Topology Validation Output
55+
ansible.builtin.debug:
56+
msg: "Topology Validation Failed. Output: {{ topology_validation.stdout }}"
57+
when: topology_validation.rc != 0
58+
59+
- name: Verify Kueue ResourceFlavor exists with correct properties
60+
delegate_to: localhost
61+
ansible.builtin.shell: |
62+
# Get the YAML for the ResourceFlavor and check for key properties
63+
kubectl get resourceflavor {{instance_type}} -o yaml | grep -q -E "topologyName: a4x-default" &&
64+
kubectl get resourceflavor {{instance_type}} -o yaml | grep -q -E "cloud.google.com/gke-accelerator: {{ accelerator_type }}" &&
65+
kubectl get resourceflavor {{instance_type}} -o yaml | grep -q -E "key: nvidia.com/gpu"
66+
register: resourceflavor_validation
67+
until: resourceflavor_validation.rc == 0
68+
retries: 5
69+
delay: 5
70+
- name: Debug ResourceFlavor Validation Output
71+
ansible.builtin.debug:
72+
msg: "ResourceFlavor Validation Failed. Output: {{ resourceflavor_validation.stdout }}"
73+
when: resourceflavor_validation.rc != 0
74+
75+
- name: Verify Kueue ClusterQueue exists with correct quota and covered resources
76+
delegate_to: localhost
77+
ansible.builtin.shell: |
78+
# Get the YAML for the ClusterQueue and check for key properties
79+
kubectl get clusterqueue {{instance_type}} -o yaml | grep -q -E "nominalQuota: \"{{ num_gpus }}\""
80+
register: clusterqueue_validation
81+
until: clusterqueue_validation.rc == 0
82+
retries: 5
83+
delay: 5
84+
- name: Debug ClusterQueue Validation Output
85+
ansible.builtin.debug:
86+
msg: "ClusterQueue Validation Failed. Output: {{ clusterqueue_validation.stdout }}"
87+
when: clusterqueue_validation.rc != 0
88+
89+
- name: Verify Kueue LocalQueue {{instance_type}} exists and links to ClusterQueue
90+
delegate_to: localhost
91+
ansible.builtin.shell: |
92+
# Get the YAML for the LocalQueue and check its clusterQueue link
93+
kubectl get localqueue {{instance_type}} -n default -o yaml | grep -q -E "clusterQueue: {{instance_type}}"
94+
register: localqueue_validation
95+
until: localqueue_validation.rc == 0
96+
retries: 5
97+
delay: 5
98+
- name: Debug LocalQueue Validation Output
99+
ansible.builtin.debug:
100+
msg: "LocalQueue Validation Failed. Output: {{ localqueue_validation.stdout }}"
101+
when: localqueue_validation.rc != 0
102+
103+
- name: Check for host topologyAssignment in workloads
104+
delegate_to: localhost
105+
ansible.builtin.shell: |
106+
kubectl get workloads -o yaml | grep "topologyAssignment" -A 5
107+
register: tas_check_host
108+
until: tas_check_host.stdout | length > 0
109+
retries: 20
110+
delay: 5
111+
changed_when: false
112+
113+
- name: Print the topologyAssignment details
114+
ansible.builtin.debug:
115+
msg: "Found topologyAssignment in workload status:\n{{ tas_check_host.stdout }}"
116+
117+
- name: Ensure all pods are on the same host
118+
delegate_to: localhost
119+
ansible.builtin.shell: |
120+
kubectl get pods -v=9 \
121+
-o custom-columns="Name:.metadata.name,Host:.spec.nodeSelector.cloud\.google\.com/gce-topology-host" | \
122+
sort -k2 | uniq -f 1 | wc -l
123+
register: unique_host_count
124+
until: unique_host_count.stdout | int == 2
125+
retries: 10
126+
delay: 10
127+
128+
- name: Delete the host topology kueue job
129+
delegate_to: localhost
130+
ansible.builtin.shell: |
131+
kubectl delete --all jobs -v=9
132+
args:
133+
executable: /bin/bash
134+
changed_when: False
135+
136+
- name: Create the rack topology kueue job
137+
delegate_to: localhost
138+
ansible.builtin.shell: |
139+
TEMPLATE_PATH="{{ workspace }}/tools/cloud-build/daily-tests/blueprints/kueue-config-files/rack-topology-tas-small-job.yaml"
140+
OUTPUT_PATH="/tmp/rack-topology-job.yaml"
141+
sed -e 's|image:.*|image: busybox|' \
142+
-e 's|args:.*|command: ["sleep", "10s"]|' \
143+
-e "s/tas-user-queue/{{ instance_type }}/g" \
144+
"${TEMPLATE_PATH}" > "${OUTPUT_PATH}"
145+
kubectl create -f "${OUTPUT_PATH}" -v=9
146+
args:
147+
executable: /bin/bash
148+
changed_when: False
149+
150+
- name: Check for rack topologyAssignment in workloads
151+
delegate_to: localhost
152+
ansible.builtin.shell: |
153+
kubectl get workloads -o yaml | grep "topologyAssignment" -A 5
154+
register: tas_check_rack
155+
until: tas_check_rack.stdout | length > 0
156+
retries: 20
157+
delay: 5
158+
changed_when: false
159+
160+
- name: Print the topologyAssignment details
161+
ansible.builtin.debug:
162+
msg: "Found topologyAssignment in workload status:\n{{ tas_check_rack.stdout }}"
163+
164+
- name: Ensure all pods are on the same rack
165+
delegate_to: localhost
166+
ansible.builtin.shell: |
167+
kubectl get pods -v=9 \
168+
-o custom-columns="Name:.metadata.name,Host:.spec.nodeSelector.cloud\.google\.com/gce-topology-subblock" | \
169+
sort -k2 | uniq -f 1 | wc -l
170+
register: unique_host_count
171+
until: unique_host_count.stdout | int == 2
172+
retries: 10
173+
delay: 10
174+
175+
- name: Delete the rack topology kueue job
176+
delegate_to: localhost
177+
ansible.builtin.shell: |
178+
kubectl delete --all jobs -v=9
179+
args:
180+
executable: /bin/bash
181+
changed_when: False
182+
183+
- name: Create the block topology kueue job
184+
delegate_to: localhost
185+
ansible.builtin.shell: |
186+
TEMPLATE_PATH="{{ workspace }}/tools/cloud-build/daily-tests/blueprints/kueue-config-files/block-topology-tas-small-job.yaml"
187+
OUTPUT_PATH="/tmp/block-topology-job.yaml"
188+
sed -e 's|image:.*|image: busybox|' \
189+
-e 's|args:.*|command: ["sleep", "10s"]|' \
190+
-e "s/tas-user-queue/{{ instance_type }}/g" \
191+
"${TEMPLATE_PATH}" > "${OUTPUT_PATH}"
192+
kubectl create -f "${OUTPUT_PATH}" -v=9
193+
args:
194+
executable: /bin/bash
195+
changed_when: False
196+
197+
- name: Check for block topologyAssignment in workloads
198+
delegate_to: localhost
199+
ansible.builtin.shell: |
200+
kubectl get workloads -o yaml | grep "topologyAssignment" -A 5
201+
register: tas_check_block
202+
until: tas_check_block.stdout | length > 0
203+
retries: 20
204+
delay: 5
205+
changed_when: false
206+
207+
- name: Print the topologyAssignment details
208+
ansible.builtin.debug:
209+
msg: "Found topologyAssignment in workload status:\n{{ tas_check_block.stdout }}"
210+
211+
- name: Ensure all pods are on the same block
212+
delegate_to: localhost
213+
ansible.builtin.shell: |
214+
kubectl get pods -v=9 \
215+
-o custom-columns="Name:.metadata.name,Host:.spec.nodeSelector.cloud\.google\.com/gce-topology-block" | \
216+
sort -k2 | uniq -f 1 | wc -l
217+
register: unique_host_count
218+
until: unique_host_count.stdout | int == 2
219+
retries: 10
220+
delay: 10
221+
222+
- name: Wait for job to complete
223+
delegate_to: localhost
224+
ansible.builtin.command: |
225+
kubectl get job --field-selector status.successful=2
226+
register: job_completion
227+
until: job_completion.stdout_lines | length > 1
228+
retries: 10
229+
delay: 5
230+
231+
- name: Print job_completion debug output
232+
ansible.builtin.debug:
233+
var: job_completion.stdout_lines
234+
always:
235+
- name: Clean up temporary job files (always run)
236+
delegate_to: localhost
237+
ansible.builtin.file:
238+
path: "{{ item }}"
239+
state: absent
240+
with_items:
241+
- "/tmp/host-topology-job.yaml"
242+
- "/tmp/rack-topology-job.yaml"
243+
- "/tmp/block-topology-job.yaml"

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

Lines changed: 137 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,137 @@
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: Assert variables are defined
16+
ansible.builtin.assert:
17+
that:
18+
- region is defined
19+
- custom_vars.project is defined
20+
21+
- name: Get cluster credentials for kubectl
22+
delegate_to: localhost
23+
ansible.builtin.shell: |
24+
gcloud container clusters get-credentials {{ deployment_name }} --region {{ region }} --project {{ custom_vars.project }} --verbosity=debug
25+
26+
- name: Get cluster nodes to verify connection and see node status
27+
delegate_to: localhost
28+
ansible.builtin.shell: |
29+
kubectl get nodes -o wide
30+
register: nodes_status
31+
32+
- name: Print node status
33+
ansible.builtin.debug:
34+
msg: "{{ nodes_status.stdout_lines }}"
35+
36+
- name: Run the NCCL test
37+
delegate_to: localhost
38+
ansible.builtin.shell: |
39+
kubectl create -f {{ workspace }}/examples/gke-a4x/nccl-jobset-example.yaml -v=9
40+
args:
41+
executable: /bin/bash
42+
43+
- name: Get JobSet status after creation
44+
delegate_to: localhost
45+
ansible.builtin.shell: |
46+
sleep 10 # Wait a moment for the jobset to be fully recognized
47+
kubectl get jobset --show-labels
48+
register: jobset_status
49+
50+
- name: Print JobSet status
51+
ansible.builtin.debug:
52+
msg: "{{ jobset_status.stdout_lines }}"
53+
54+
- name: Block to wait for job completion and capture failures
55+
block:
56+
- name: Wait for Job to hit 2/2 completions
57+
delegate_to: localhost
58+
ansible.builtin.shell: |
59+
kubectl get job --field-selector status.successful=2 -v=9
60+
register: job_completion
61+
until: job_completion.stdout_lines | length > 1
62+
retries: 20
63+
delay: 60
64+
rescue:
65+
- name: "FAILURE: Job did not complete in time. Capturing debug info."
66+
ansible.builtin.debug:
67+
msg: "The JobSet failed to complete within the expected time. See debug output below."
68+
69+
- name: Get all pods status on failure
70+
delegate_to: localhost
71+
ansible.builtin.shell: |
72+
kubectl get pods -o wide
73+
register: pods_on_failure
74+
75+
- name: Print all pods status on failure
76+
ansible.builtin.debug:
77+
msg: "{{ pods_on_failure.stdout_lines }}"
78+
79+
- name: Describe JobSet on failure
80+
delegate_to: localhost
81+
ansible.builtin.shell: |
82+
kubectl describe jobset nccl-all
83+
register: jobset_describe_on_failure
84+
85+
- name: Print JobSet description on failure
86+
ansible.builtin.debug:
87+
msg: "{{ jobset_describe_on_failure.stdout_lines }}"
88+
89+
- name: Describe pods on failure
90+
delegate_to: localhost
91+
ansible.builtin.shell: |
92+
kubectl describe pods --selector=jobset.x-k8s.io/jobset-name=nccl-all
93+
register: pods_describe_on_failure
94+
95+
- name: Print pod description on failure
96+
ansible.builtin.debug:
97+
msg: "{{ pods_describe_on_failure.stdout_lines }}"
98+
99+
- name: Fail the playbook
100+
ansible.builtin.fail:
101+
msg: "NCCL JobSet did not complete successfully."
102+
103+
- name: "SUCCESS: Job completed. Fetching logs for validation."
104+
ansible.builtin.debug:
105+
msg: "JobSet completed successfully. Now fetching logs to validate bandwidth."
106+
107+
- name: Fetch logs from the 0-0 pod and save to pod_logs.txt
108+
delegate_to: localhost
109+
ansible.builtin.shell: |
110+
pod_0_0_name="$(kubectl get pods --no-headers -o custom-columns="Name:.metadata.name" | awk '{if ($1 ~ "0-0") print $0}')"
111+
kubectl logs ${pod_0_0_name} > pod_logs.txt
112+
cat pod_logs.txt
113+
register: nccl_test_logs
114+
115+
- name: Print the NCCL test logs from 0-0 pod
116+
debug:
117+
msg: "{{nccl_test_logs.stdout}}"
118+
119+
- name: Ensure average bus bandwidth is >= 100 GB/s
120+
delegate_to: localhost
121+
ansible.builtin.shell: |
122+
grep -o '# Avg bus bandwidth : [0-9\.]*' pod_logs.txt | cut -d ':' -f 2 | awk '{$1=$1;print}'
123+
register: avg_bus_bandwidth
124+
failed_when: avg_bus_bandwidth.stdout | float < 100
125+
126+
- name: "SUCCESS: Bandwidth validation passed. Cleaning up resources."
127+
ansible.builtin.debug:
128+
msg: "Average bus bandwidth is >= 100 GB/s. Test passed. Cleaning up JobSet."
129+
130+
- name: Clean up
131+
delegate_to: localhost
132+
ansible.builtin.shell: |
133+
kubectl delete jobset --all -v=9
134+
135+
- name: "Cleanup complete."
136+
ansible.builtin.debug:
137+
msg: "JobSet has been deleted."

0 commit comments

Comments
 (0)