Merge branch 'develop' into shubham/a4x-debug

Author: shubpal07

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

@@ -0,0 +1,69 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+---
+tags:
+- m.vpc
+- m.service-account
+- m.gke-cluster
+- m.gke-node-pool
+- m.kubectl-apply
+- gke
+
+timeout: 7200s
+
+steps:
+# While using static network names we are guarding against more than 1 instance running at a time (for multi-group tests)
+- id: check_for_running_build
+  name: gcr.io/cloud-builders/gcloud
+  script: "tools/cloud-build/check_running_build.sh tools/cloud-build/daily-tests/builds/gke-tpu-v6e.yaml"
+
+- id: gke-tpu-v6e
+  name: us-central1-docker.pkg.dev/$PROJECT_ID/hpc-toolkit-repo/test-runner
+  entrypoint: /bin/bash
+  env:
+  - "ANSIBLE_HOST_KEY_CHECKING=false"
+  - "ANSIBLE_CONFIG=/workspace/tools/cloud-build/ansible.cfg"
+  args:
+  - -c
+  - |
+    set -x -e
+    cd /workspace && make
+    BUILD_ID_FULL=$BUILD_ID
+    BUILD_ID_SHORT=$${BUILD_ID_FULL:0:6}
+    EXAMPLE_BP=examples/gke-tpu-v6/gke-tpu-v6.yaml
+    # adding vm to act as remote node
+    echo '  - id: remote-node'                                                            >> $${EXAMPLE_BP}
+    echo '    source: modules/compute/vm-instance'                                        >> $${EXAMPLE_BP}
+    echo '    use: [gke-tpu-v6-net-0]'                                                    >> $${EXAMPLE_BP}
+    echo '    settings:'                                                                  >> $${EXAMPLE_BP}
+    echo '      machine_type: n2-standard-2'                                              >> $${EXAMPLE_BP}
+    echo '      name_prefix: remote-node'                                                 >> $${EXAMPLE_BP}
+    echo '      add_deployment_name_before_prefix: true'                                  >> $${EXAMPLE_BP}
+    echo ''
+    echo '  - id: job_template_hostname'                                                   >> $${EXAMPLE_BP}
+    echo '    source: modules/compute/gke-job-template'                                    >> $${EXAMPLE_BP}
+    echo '    use: [gke-tpu-v6-pool]'                                                      >> $${EXAMPLE_BP}
+    echo '    settings:'                                                                   >> $${EXAMPLE_BP}
+    echo '      image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest'         >> $${EXAMPLE_BP}
+    echo '      command:'                                                                  >> $${EXAMPLE_BP}
+    echo '      - /bin/bash'                                                               >> $${EXAMPLE_BP}
+    echo '      - -c'                                                                      >> $${EXAMPLE_BP}
+    echo '      - python -c '\''import jax; print(jax.device_count(), "TPU cores")'\'''    >> $${EXAMPLE_BP}
+    echo '      node_count: 1'                                                             >> $${EXAMPLE_BP}
+    echo '    outputs: [instructions]'                                                     >> $${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="@tools/cloud-build/daily-tests/tests/gke-tpu-v6e.yml"

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

@@ -0,0 +1,41 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+---
+test_name: gke-tpu-v6e
+deployment_name: gke-tpu-v6e-{{ build }}
+workspace: /workspace
+blueprint_yaml: "{{ workspace }}/examples/gke-tpu-v6/gke-tpu-v6.yaml"
+network: "{{ deployment_name }}-net-0"
+region: us-central2
+zone: us-central2-b
+remote_node: "{{ deployment_name }}-remote-node-0"
+machine_type: ct6e-standard-4t
+extended_reservation: cloudtpu-20251024020500-953082093
+num_slices: 1
+tpu_topology: 2x2
+static_node_count: 1
+cli_deployment_vars:
+  region: "{{ region }}"
+  zone: "{{ zone }}"
+  num_slices: "{{ num_slices }}"
+  machine_type: "{{ machine_type }}"
+  tpu_topology: "{{ tpu_topology }}"
+  static_node_count: "{{ static_node_count }}"
+  authorized_cidr: "{{ build_ip.stdout }}/32"
+  reservation: "{{ extended_reservation }}"
+custom_vars:
+  project: "{{ project }}"
+post_deploy_tests:
+- test-validation/test-gke-job.yml