Add integration test for gke-tpu-v5e

tools/cloud-build/daily-tests/ansible_playbooks/test-validation/test-gcluster-job-pathways.yml

@@ -0,0 +1,264 @@
+# Copyright 2026 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.
+
+---
+
+- name: Assert variables are defined
+  ansible.builtin.assert:
+    that:
+    - region is defined
+    - custom_vars.project is defined
+    - custom_vars.gcs_bucket is defined
+    - custom_vars.tpu_type is defined
+    - custom_vars.num_slices is defined
+    - custom_vars.expected_tpu_count is defined
+    - custom_vars.workload_image is defined
+    - custom_vars.workload_command is defined
+    - custom_vars.expected_log_output is defined
+
+- name: Get cluster credentials for GCluster operations
+  delegate_to: localhost
+  ansible.builtin.command: gcloud container clusters get-credentials {{ deployment_name }} --region {{ region }} --project {{ custom_vars.project }} --verbosity=debug
+
+- name: Set test local facts
+  ansible.builtin.set_fact:
+    job_name: "pw-{{ build | default('test') }}"
+    dry_run_file: "/tmp/{{ deployment_name }}-pw-dry-run.yaml"
+    templated_command: "{{ custom_vars.workload_command | replace('GCS_BUCKET', custom_vars.gcs_bucket) }}"
+
+# ==========================================
+# WORKFLOW 1: Dry-Run Validation & Env Check
+# ==========================================
+
+- name: Configure Docker authentication for GCR registry checks
+  delegate_to: localhost
+  ansible.builtin.command: gcloud auth configure-docker gcr.io --quiet
+  changed_when: false
+
+- name: Configure Docker authentication for regional Artifact Registry checks
+  delegate_to: localhost
+  ansible.builtin.command: gcloud auth configure-docker {{ region }}-docker.pkg.dev --quiet
+  changed_when: false
+
+- name: Execute GCluster Job Submit Dry-Run
+  delegate_to: localhost
+  ansible.builtin.command:
+    argv:
+    - "{{ workspace }}/gcluster"
+    - "job"
+    - "submit"
+    - "--pathways"
+    - "--name"
+    - "{{ job_name }}"
+    - "--num-slices"
+    - "{{ custom_vars.num_slices }}"
+    - "--compute-type"
+    - "{{ custom_vars.tpu_type }}"
+    - "--project"
+    - "{{ custom_vars.project }}"
+    - "--location"
+    - "{{ region }}"
+    - "--cluster"
+    - "{{ deployment_name }}"
+    - "--image"
+    - "{{ custom_vars.workload_image }}"
+    - "--pathways-gcs-location"
+    - "gs://{{ custom_vars.gcs_bucket }}/pathways-tests-{{ deployment_name }}"
+    - "--command"
+    - "{{ templated_command }}"
+    - "--pathways-proxy-env"
+    - "TEST_PROXY_VAR=proxy_val"
+    - "--pathways-server-env"
+    - "TEST_SERVER_VAR=server_val"
+    - "--pathways-worker-env"
+    - "TEST_WORKER_VAR=worker_val"
+    - "--dry-run-out"
+    - "{{ dry_run_file }}"
+  changed_when: false
+
+- name: Read generated Dry-Run manifest
+  delegate_to: localhost
+  ansible.builtin.slurp:
+    src: "{{ dry_run_file }}"
+  register: dry_run_content
+
+- name: Parse and Assert Dry-Run manifest content
+  ansible.builtin.set_fact:
+    dry_run_yaml: "{{ dry_run_content.content | b64decode | from_yaml }}"
+
+- name: Validate manifest is a JobSet
+  ansible.builtin.assert:
+    that:
+    - dry_run_yaml.kind == "JobSet"
+    fail_msg: "Generated manifest is not a JobSet, kind found: {{ dry_run_yaml.kind }}"
+
+# Helper queries to assert correct env variable injection in proxy, RM server, and worker specs
+- name: Validate Env Variables in dry-run manifest
+  ansible.builtin.assert:
+    that:
+    # 1. Assert proxy container env exists
+    - >
+      dry_run_yaml.spec.replicatedJobs
+      | selectattr('name', 'equalto', 'pathways-head') | map(attribute='template.spec.template.spec.initContainers') | flatten
+      | selectattr('name', 'equalto', 'pathways-proxy') | map(attribute='env') | flatten
+      | selectattr('name', 'equalto', 'TEST_PROXY_VAR') | selectattr('value', 'equalto', 'proxy_val') | list | length > 0
+    # 2. Assert RM server container env exists
+    - >
+      dry_run_yaml.spec.replicatedJobs
+      | selectattr('name', 'equalto', 'pathways-head') | map(attribute='template.spec.template.spec.initContainers') | flatten
+      | selectattr('name', 'equalto', 'pathways-rm') | map(attribute='env') | flatten
+      | selectattr('name', 'equalto', 'TEST_SERVER_VAR') | selectattr('value', 'equalto', 'server_val') | list | length > 0
+    # 3. Assert worker container env exists
+    - >
+      dry_run_yaml.spec.replicatedJobs
+      | selectattr('name', 'equalto', 'worker') | map(attribute='template.spec.template.spec.containers') | flatten
+      | selectattr('name', 'equalto', 'pathways-worker') | map(attribute='env') | flatten
+      | selectattr('name', 'equalto', 'TEST_WORKER_VAR') | selectattr('value', 'equalto', 'worker_val') | list | length > 0
+    fail_msg: "Environment variables failed to inject properly in dry-run output"
+
+# ==========================================
+# WORKFLOW 2: Submit & Environment Verification
+# ==========================================
+
+- name: Submit the actual Pathways job to GKE
+  delegate_to: localhost
+  ansible.builtin.command:
+    argv:
+    - "{{ workspace }}/gcluster"
+    - "job"
+    - "submit"
+    - "--pathways"
+    - "--name"
+    - "{{ job_name }}"
+    - "--num-slices"
+    - "{{ custom_vars.num_slices }}"
+    - "--compute-type"
+    - "{{ custom_vars.tpu_type }}"
+    - "--project"
+    - "{{ custom_vars.project }}"
+    - "--location"
+    - "{{ region }}"
+    - "--cluster"
+    - "{{ deployment_name }}"
+    - "--image"
+    - "{{ custom_vars.workload_image }}"
+    - "--pathways-gcs-location"
+    - "gs://{{ custom_vars.gcs_bucket }}/pathways-tests-{{ deployment_name }}"
+    - "--command"
+    - "{{ templated_command }}"
+  register: submission_output
+  changed_when: false
+
+- name: Block to wait for job completion, assert logs and clean up on failure
+  block:
+  - name: Wait for Pathways JobSet resource creation
+    delegate_to: localhost
+    ansible.builtin.command: kubectl get jobset "{{ job_name }}"
+    register: jobset_check
+    until: jobset_check.rc == 0
+    retries: 12
+    delay: 5
+    changed_when: false
+
+  - name: Wait for Pathways Proxy and RM sidecars to initialize
+    delegate_to: localhost
+    ansible.builtin.command: kubectl wait --for=condition=Initialized pod -l jobset.sigs.k8s.io/jobset-name={{ job_name }} --timeout=600s
+    changed_when: false
+
+  - name: Wait for Pathways Workload container to start running
+    delegate_to: localhost
+    ansible.builtin.shell: >
+      kubectl get pod -l jobset.sigs.k8s.io/jobset-name={{ job_name }}
+      -o jsonpath='{.items[*].status.containerStatuses[?(@.name=="workload-container")].state}'
+      | grep -E "running|terminated"
+    register: workload_state
+    until: workload_state.rc == 0
+    retries: 60
+    delay: 5
+    changed_when: false
+
+  - name: Stream Pathways Workload logs for JAX TPU execution result
+    delegate_to: localhost
+    ansible.builtin.shell: >
+      timeout 600 bash -c '
+      {{ workspace }}/gcluster job logs "{{ job_name }}" -f
+      --cluster {{ deployment_name }}
+      --location {{ region }}
+      --project {{ custom_vars.project }}
+      --main-only
+      | grep -m 1 "JAX TPU Result"
+      '
+    register: streaming_logs
+
+  - name: Print streaming logs result
+    ansible.builtin.debug:
+      msg: "{{ streaming_logs.stdout }}"
+
+  - name: Assert JAX TPU computation and device count verification
+    ansible.builtin.assert:
+      that:
+      - streaming_logs.stdout is search("JAX TPU Result")
+      - streaming_logs.stdout is search(custom_vars.expected_log_output)
+      fail_msg: "Failed to find expected JAX output in streaming logs."
+
+  rescue:
+  - name: "FAILURE: Job failed execution. Gathering GCluster diagnostic inspection."
+    ansible.builtin.debug:
+      msg: "Workload failed or did not print the expected output. Commencing diagnostics."
+
+  - name: Run GCluster diagnostic inspection on failure
+    delegate_to: localhost
+    ansible.builtin.command: >
+      {{ workspace }}/gcluster job inspect
+      --name "{{ job_name }}"
+      --show
+      --cluster {{ deployment_name }}
+      --location {{ region }}
+      --project {{ custom_vars.project }}
+    register: inspect_report
+    changed_when: false
+    ignore_errors: true
+
+  - name: Print GCluster diagnostic inspection report
+    ansible.builtin.debug:
+      msg: "{{ inspect_report.stdout_lines | default('No diagnostic report generated') }}"
+
+  - name: Get all job container logs on failure (GCluster)
+    delegate_to: localhost
+    ansible.builtin.command: >
+      {{ workspace }}/gcluster job logs "{{ job_name }}"
+      --cluster {{ deployment_name }}
+      --location {{ region }}
+      --project {{ custom_vars.project }}
+    register: all_logs_failure
+    changed_when: false
+    ignore_errors: true
+
+  - name: Print all job logs on failure
+    ansible.builtin.debug:
+      msg: "{{ all_logs_failure.stdout | default('No job logs found') }}"
+
+  - name: Force failure of Ansible run
+    ansible.builtin.fail:
+      msg: "Integration test failed to verify GCluster pathways job execution."
+
+  always:
+  - name: Clean up workload (cancel Pathways job)
+    delegate_to: localhost
+    ansible.builtin.command: >
+      {{ workspace }}/gcluster job cancel "{{ job_name }}"
+      --cluster {{ deployment_name }}
+      --location {{ region }}
+      --project {{ custom_vars.project }}
+    ignore_errors: true

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

@@ -0,0 +1,125 @@
+# Copyright 2026 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
+
+substitutions:
+  _TEST_PREFIX: "" # Default to no prefix
+
+timeout: 10800s
+
+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-v5e.yaml"
+
+- id: gke-tpu-v5e
+  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"
+  - "MACHINE_TYPE=tpu"
+  - "TPU_RUNTIME_VERSION=v2-alpha-tpuv5-lite"
+  - "ACCELERATOR_TYPE=v5litepod-4"
+  - "PROJECT_ID=$PROJECT_ID"
+  - "NUM_NODES=1"
+  - "INSTANCE_PREFIX=v5esp"
+  - "BUILD_ID=$BUILD_ID"
+  - "OPTIONS_GCS_PATH=gs://hpc-ctk1357/tpuv5eoptions.txt"
+  - "ENABLE_SPOT_FALLBACK=true"
+  args:
+  - -c
+  - |
+    set -e -u -o pipefail
+    echo "Sourcing find_available_zone.sh to determine zone."
+    source /workspace/tools/cloud-build/find_available_zone.sh
+    if [ -z "$${ZONE:-}" ]; then
+      echo "ERROR: ZONE not found" >&2
+      exit 1
+    fi
+    set -x
+    cd /workspace
+    bash tools/get_binary.sh "${_TEST_PREFIX}"
+    REGION="$${ZONE%-*}"
+    BUILD_ID_SHORT="$${BUILD_ID:0:6}"
+
+    EXAMPLE_BP=examples/gke-tpu-v5e/gke-tpu-v5e.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-v5e-net]'                                                     >> $${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-v5e-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 warnings; warnings.filterwarnings("ignore"); import jax; print(jax.device_count(), "TPU cores")'\'''    >> $${EXAMPLE_BP}
+    echo '      node_count: 1'                                                             >> $${EXAMPLE_BP}
+    echo '    outputs: [instructions]'                                                     >> $${EXAMPLE_BP}
+
+    ENABLE_SPOT="true"
+    GKE_VARS_FILE="tools/cloud-build/daily-tests/tests/gke-tpu-v5e.yml"
+    if [[ "$$PROVISIONING_MODEL" == "STANDARD" ]]; then
+      ENABLE_SPOT="false"
+      sed -i '/instance_labels:/,+1d' "$${GKE_VARS_FILE}"
+      sed -i '/enable_spot: true/d' "$${GKE_VARS_FILE}"
+    else
+      echo "INFO: Using $${GKE_VARS_FILE} as it is for SPOT provisioning."
+    fi
+
+    sed -i -e '/reservation_affinity:/,+3c\      spot: '"$$ENABLE_SPOT"'' $${EXAMPLE_BP}
+    sed -i '/reservation/d' $${EXAMPLE_BP}
+    bash tools/add_ttl_label.sh "$${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} full_build_id=$BUILD_ID chs_repo=$${CHS_REPO}" \
+        --extra-vars="region=$${REGION} zone=$${ZONE}" \
+        --extra-vars="enable_spot=$${ENABLE_SPOT}" \
+        --extra-vars="@$${GKE_VARS_FILE}" \
+        --extra-vars="triage_gcs_bucket_override=$$TRIAGE_GCS_BUCKET" \
+        --extra-vars="triage_project_number_override=$$TRIAGE_PROJECT_NUMBER" \
+        --extra-vars="triage_invoker_sa_override=$$TRIAGE_INVOKER_SA" \
+        --extra-vars="triage_cloud_run_url_override=$$TRIAGE_CLOUD_RUN_URL"
+  secretEnv: ['GCLUSTER_GCS_PATH', 'CHS_REPO', 'TRIAGE_GCS_BUCKET', 'TRIAGE_PROJECT_NUMBER', 'TRIAGE_INVOKER_SA', 'TRIAGE_CLOUD_RUN_URL']
+availableSecrets:
+  secretManager:
+  - versionName: projects/${PROJECT_ID}/secrets/gcluster-develop-release-bucket/versions/latest
+    env: 'GCLUSTER_GCS_PATH'
+  - versionName: projects/${PROJECT_ID}/secrets/triage-gcs-bucket/versions/latest
+    env: 'TRIAGE_GCS_BUCKET'
+  - versionName: projects/${PROJECT_ID}/secrets/triage-project-number/versions/latest
+    env: 'TRIAGE_PROJECT_NUMBER'
+  - versionName: projects/${PROJECT_ID}/secrets/triage-invoker-sa/versions/latest
+    env: 'TRIAGE_INVOKER_SA'
+  - versionName: projects/${PROJECT_ID}/secrets/triage-cloud-run-url/versions/latest
+    env: 'TRIAGE_CLOUD_RUN_URL'
+  - versionName: projects/${PROJECT_ID}/secrets/cluster-health-scanner/versions/latest
+    env: 'CHS_REPO'