Skip to content

Commit 68ae6e0

Browse files
agrawalkhushi18agrawalkhushi18
committed
Adding a post-deploy test to assert tpu device count
1 parent 4c3246b commit 68ae6e0

4 files changed

Lines changed: 151 additions & 3 deletions

File tree

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

Lines changed: 146 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,146 @@
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+
- custom_vars.expected_tpu_count is defined
21+
22+
- name: Get cluster credentials for kubectl
23+
delegate_to: localhost
24+
ansible.builtin.command: gcloud container clusters get-credentials {{ deployment_name }} --region {{ region }} --project {{ custom_vars.project }} --verbosity=debug
25+
26+
- name: Identify Job File
27+
delegate_to: localhost
28+
ansible.builtin.shell: ls {{ workspace }}/{{ deployment_name }}/primary/my-job*.yaml | head -n 1
29+
register: job_file_result
30+
changed_when: False
31+
32+
- name: Register Job Variables
33+
ansible.builtin.set_fact:
34+
job_file_path: "{{ job_file_result.stdout }}"
35+
job_name: "{{ job_file_result.stdout | basename | splitext | first }}"
36+
37+
- name: Execute the job
38+
delegate_to: localhost
39+
ansible.builtin.command: kubectl create -f {{ job_file_path }} -v=9
40+
changed_when: False
41+
42+
- name: Block to wait for job completion and capture failures
43+
block:
44+
- name: Wait for job to complete
45+
delegate_to: localhost
46+
ansible.builtin.command: kubectl wait --for=condition=complete "job/{{ job_name }}" --timeout=10m
47+
register: job_wait_result
48+
49+
- name: Get job pods
50+
delegate_to: localhost
51+
ansible.builtin.shell: |
52+
kubectl get pods --selector=job-name={{ job_name }} -o jsonpath='{.items[0].metadata.name}'
53+
register: job_pod_name
54+
changed_when: false
55+
56+
- name: Get job logs
57+
delegate_to: localhost
58+
ansible.builtin.shell: |
59+
kubectl logs {{ job_pod_name.stdout }}
60+
register: job_logs
61+
changed_when: false
62+
63+
- name: Print job output
64+
ansible.builtin.debug:
65+
msg: "Job Output: {{ job_logs.stdout }}"
66+
67+
- name: Verify TPU count
68+
ansible.builtin.assert:
69+
that:
70+
- (custom_vars.expected_tpu_count | string ~ ' TPU cores') in job_logs.stdout
71+
fail_msg: "Expected '{{ custom_vars.expected_tpu_count }} TPU cores' in logs, but got: {{ job_logs.stdout }}"
72+
73+
rescue:
74+
- name: "FAILURE: Job did not complete or verification failed. Capturing debug info."
75+
ansible.builtin.debug:
76+
msg: "The Job failed to complete within the expected time or failed verification. See debug output below."
77+
78+
- name: Get all pods status on failure
79+
delegate_to: localhost
80+
ansible.builtin.shell: |
81+
kubectl get pods -o wide
82+
register: pods_on_failure
83+
changed_when: false
84+
85+
- name: Print all pods status on failure
86+
ansible.builtin.debug:
87+
msg: "{{ pods_on_failure.stdout_lines }}"
88+
89+
- name: Describe the job on failure
90+
delegate_to: localhost
91+
ansible.builtin.shell: |
92+
kubectl describe job {{ job_name }}
93+
register: job_describe_on_failure
94+
changed_when: false
95+
96+
- name: Print job description on failure
97+
ansible.builtin.debug:
98+
msg: "{{ job_describe_on_failure.stdout_lines }}"
99+
100+
- name: Get pod name on failure
101+
delegate_to: localhost
102+
ansible.builtin.shell: |
103+
kubectl get pods --selector=job-name={{ job_name }} --no-headers -o custom-columns=":metadata.name"
104+
register: pod_name_on_failure
105+
changed_when: false
106+
107+
- name: Describe the pod on failure
108+
delegate_to: localhost
109+
ansible.builtin.shell: |
110+
kubectl describe pod {{ item }}
111+
loop: "{{ pod_name_on_failure.stdout_lines }}"
112+
register: pod_describe_on_failure
113+
changed_when: false
114+
when: pod_name_on_failure.stdout_lines | length > 0
115+
116+
- name: Print pod description on failure
117+
ansible.builtin.debug:
118+
msg: "{{ pod_describe_on_failure.results }}"
119+
when: pod_name_on_failure.stdout_lines | length > 0
120+
121+
- name: Get logs from the pod on failure
122+
delegate_to: localhost
123+
ansible.builtin.shell: |
124+
kubectl logs {{ item }}
125+
loop: "{{ pod_name_on_failure.stdout_lines }}"
126+
register: pod_logs_on_failure
127+
changed_when: false
128+
when: pod_name_on_failure.stdout_lines | length > 0
129+
130+
- name: Print pod logs on failure
131+
ansible.builtin.debug:
132+
msg: "{{ pod_logs_on_failure.results }}"
133+
when: pod_name_on_failure.stdout_lines | length > 0
134+
135+
- name: Fail the playbook
136+
ansible.builtin.fail:
137+
msg: "GKE Job failed."
138+
139+
- name: Print job_completion debug output
140+
ansible.builtin.debug:
141+
var: job_wait_result.stdout_lines
142+
143+
- name: Clean up
144+
delegate_to: localhost
145+
ansible.builtin.shell: |
146+
kubectl delete job {{ job_name }} -v=9

‎tools/cloud-build/daily-tests/builds/gke-tpu-v6e.yaml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,7 @@ steps:
6060
echo ' command:' >> $${EXAMPLE_BP}
6161
echo ' - /bin/bash' >> $${EXAMPLE_BP}
6262
echo ' - -c' >> $${EXAMPLE_BP}
63-
echo ' - python -c '\''import jax; print(jax.device_count(), "TPU cores")'\''' >> $${EXAMPLE_BP}
63+
echo ' - python -c '\''import warnings; warnings.filterwarnings("ignore"); import jax; print(jax.device_count(), "TPU cores")'\''' >> $${EXAMPLE_BP}
6464
echo ' node_count: 1' >> $${EXAMPLE_BP}
6565
echo ' outputs: [instructions]' >> $${EXAMPLE_BP}
6666

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,5 +37,6 @@ cli_deployment_vars:
3737
reservation: "{{ extended_reservation }}"
3838
custom_vars:
3939
project: "{{ project }}"
40+
expected_tpu_count: 8
4041
post_deploy_tests:
41-
- test-validation/test-gke-job.yml
42+
- test-validation/test-gke-tpu.yml

‎tools/cloud-build/daily-tests/tests/gke-tpu-v6e.yml‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,5 +37,6 @@ cli_deployment_vars:
3737
reservation: "{{ extended_reservation }}"
3838
custom_vars:
3939
project: "{{ project }}"
40+
expected_tpu_count: 4
4041
post_deploy_tests:
41-
- test-validation/test-gke-job.yml
42+
- test-validation/test-gke-tpu.yml

0 commit comments

Comments
 (0)