diff --git a/tools/cloud-build/daily-tests/ansible_playbooks/test-validation/test-gcluster-job-pathways.yml b/tools/cloud-build/daily-tests/ansible_playbooks/test-validation/test-gcluster-job-pathways.yml new file mode 100644 index 0000000000..9b5bc33097 --- /dev/null +++ b/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 diff --git a/tools/cloud-build/daily-tests/builds/gke-tpu-v5e.yaml b/tools/cloud-build/daily-tests/builds/gke-tpu-v5e.yaml new file mode 100644 index 0000000000..aa55346f2e --- /dev/null +++ b/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' diff --git a/tools/cloud-build/daily-tests/tests/gke-tpu-v5e.yml b/tools/cloud-build/daily-tests/tests/gke-tpu-v5e.yml new file mode 100644 index 0000000000..1c41de706d --- /dev/null +++ b/tools/cloud-build/daily-tests/tests/gke-tpu-v5e.yml @@ -0,0 +1,52 @@ +# 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. + +--- +test_name: gke-tpu-v5e +deployment_name: gke-tpu-v5e-{{ build }} +workspace: /workspace +blueprint_yaml: "{{ workspace }}/examples/gke-tpu-v5e/gke-tpu-v5e.yaml" +network: "{{ deployment_name }}-net" +remote_node: "{{ deployment_name }}-remote-node-0" +machine_type: ct5lp-hightpu-4t +num_slices: 1 +tpu_topology: 2x2 +cli_deployment_vars: + region: "{{ region }}" + zone: "{{ zone }}" + num_slices: "{{ num_slices }}" + machine_type: "{{ machine_type }}" + tpu_topology: "{{ tpu_topology }}" + authorized_cidr: "{{ build_ip.stdout }}/32" +custom_vars: + project: "{{ project }}" + expected_tpu_count: 4 + gcs_bucket: "ml-diagnostics-test" + image_tag: "us-central1-docker.pkg.dev/hpc-toolkit-dev/ml-diagnostics-test/sample-diagnostics-job:1.0.0" + tpu_accelerator: tpu-v5-lite-podslice + tpu_topology: 2x2 + k8s_service_account: workload-identity-k8s-sa + jobset_name: sample-tpu-jobset + tpu_type: v5litepod-4 + num_slices: 1 + workload_image: "us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest" + workload_command: >- + pip install pathwaysutils && python -c 'import pathwaysutils; pathwaysutils.initialize(); import jax; dc = jax.device_count(); print("JAX Device count:", dc); assert dc == 4, f"Expected 4 devices, got {dc}"; a = jax.random.normal(jax.random.PRNGKey(0), (2, 2)); print("JAX TPU Result:", a)' + expected_log_output: "JAX TPU Result:" +instance_type: "tpuv5e" +scheduler: "gke" +post_deploy_tests: +- test-validation/test-gke-tpu.yml +- test-validation/test-gke-ml-diagnostics.yml +- test-validation/test-gcluster-job-pathways.yml