feat(testing): cut over daily test suite to Kueue scheduling

.pre-commit-config.yaml

@@ -74,15 +74,6 @@ repos:
     files: 'tools/cloud-build/daily-tests/builds/.*\.yaml'
     pass_filenames: false
     require_serial: true
-  - id: validate-daily-test-builds
-    name: validate-daily-test-builds
-    entry: tools/cloud-build/daily-tests/validate_daily_test_builds.py
-    language: python
-    language_version: python3
-    additional_dependencies: ['pyyaml']
-    files: 'tools/cloud-build/daily-tests/builds/.*\.(yaml|yml)'
-    pass_filenames: true
-    require_serial: true
   - id: pytest-check
     name: pytest-check
     entry: pytest

tools/cloud-build/check_retriable_error.sh

@@ -0,0 +1,38 @@
+#!/bin/bash
+# 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.
+
+# Returns 0 if a retriable error (e.g. Terraform capacity error) is found in the log, 1 otherwise.
+# This script is called by Kueue pipelines to determine if the job should be requeued.
+LOG_FILE=$1
+if [ -z "$LOG_FILE" ]; then
+	echo "Usage: $0 <log_file>" >&2
+	exit 2
+fi
+
+if [ ! -f "$LOG_FILE" ]; then
+	echo "Error: Log file '$LOG_FILE' does not exist or is not a regular file." >&2
+	exit 2
+fi
+
+# Define all retriable errors here.
+# Note: "Couldn't find a zone to deploy" and "ERROR: ZONE not found" are not included here
+# because find_available_zone.sh now internally loops and waits for zone capacity to maintain Kueue locks.
+RETRIABLE_ERRORS="ZONE_RESOURCE_POOL_EXHAUSTED|does not have enough resources available|not enough resources available|stockout"
+
+if grep -q -i -E "$RETRIABLE_ERRORS" "$LOG_FILE"; then
+	exit 0
+else
+	exit 1
+fi

tools/cloud-build/daily-tests/builds/slurm-gcp-v6-rocky8.yaml

@@ -26,42 +26,238 @@ tags:
 substitutions:
   _TEST_PREFIX: "" # Default to no prefix
 
-timeout: 14400s  # 4hr
+timeout: 86400s
+queueTtl: 86400s  # 24hr
+
+options:
+  defaultLogsBucketBehavior: REGIONAL_USER_OWNED_BUCKET
+
 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
+# 1. Build and push workspace container containing the current code layer
+- id: build-workspace-image
+  name: gcr.io/cloud-builders/docker
+  entrypoint: /bin/bash
+  args:
+  - -c
+  - |
+    docker build -t us-central1-docker.pkg.dev/$PROJECT_ID/hpc-toolkit-repo/test-runner:$BUILD_ID -f- . <<EOF
+    FROM us-central1-docker.pkg.dev/$PROJECT_ID/hpc-toolkit-repo/test-runner:latest
+    COPY . /workspace
+    WORKDIR /workspace
+    EOF
+    docker push us-central1-docker.pkg.dev/$PROJECT_ID/hpc-toolkit-repo/test-runner:$BUILD_ID
+
+# 2. Retrieve GCS Bucket name from Secret Manager and generate Job YAML
+- id: generate-job-manifest
   name: gcr.io/cloud-builders/gcloud
-  script: "tools/cloud-build/check_running_build.sh tools/cloud-build/daily-tests/builds/slurm-gcp-v6-rocky8.yaml"
+  entrypoint: /bin/bash
+  secretEnv: ['GCLUSTER_GCS_PATH', 'SA_EMAIL', 'TRIAGE_GCS_BUCKET', 'TRIAGE_PROJECT_NUMBER', 'TRIAGE_INVOKER_SA', 'TRIAGE_CLOUD_RUN_URL']
+  args:
+  - -c
+  - |
+    BUILD_ID_SHORT=$$(echo "$BUILD_ID" | cut -c1-6)
+
+    # Determine the priority class based on the test prefix.
+    if [ "${_TEST_PREFIX}" == "hotfix-" ]; then
+      PRIORITY_LINE="priorityClassName: emergency-hotfix-priority"
+    else
+      PRIORITY_LINE=""
+    fi
+
+    SA_UNIQUE_ID=$(gcloud iam service-accounts describe "$$SA_EMAIL" --project="$PROJECT_ID" --format="value(uniqueId)")
+    OSLOGIN_USER="sa_$${SA_UNIQUE_ID}"
+
+    cat <<EOF > job.yaml
+    apiVersion: batch/v1
+    kind: Job
+    metadata:
+      name: slurm-gcp-v6-rocky8-$$BUILD_ID_SHORT
+      namespace: default
+      labels:
+        kueue.x-k8s.io/queue-name: local-queue-test-locks
+        build-id: "$BUILD_ID"
+    spec:
+      suspend: true # Required for Kueue: the job must be created in suspended state for Kueue to manage it
+      backoffLimit: 0
+      template:
+        spec:
+          $$PRIORITY_LINE
+          terminationGracePeriodSeconds: 1800 # Grants 30 minutes for terraform destroy to finish gracefully if cancelled
+          serviceAccountName: test-kueue-cluster-runner-ksa
+          restartPolicy: Never
+          containers:
+          - name: runner
+            image: us-central1-docker.pkg.dev/$PROJECT_ID/hpc-toolkit-repo/test-runner:$BUILD_ID
+            env:
+            - name: ANSIBLE_HOST_KEY_CHECKING
+              value: "false"
+            - name: ANSIBLE_CONFIG
+              value: "/workspace/tools/cloud-build/ansible.cfg"
+            - name: _TEST_PREFIX
+              value: "${_TEST_PREFIX}"
+            - name: TRIAGE_GCS_BUCKET
+              value: "$$TRIAGE_GCS_BUCKET"
+            - name: TRIAGE_PROJECT_NUMBER
+              value: "$$TRIAGE_PROJECT_NUMBER"
+            - name: TRIAGE_INVOKER_SA
+              value: "$$TRIAGE_INVOKER_SA"
+            - name: TRIAGE_CLOUD_RUN_URL
+              value: "$$TRIAGE_CLOUD_RUN_URL"
+            command: ["/bin/bash", "-c"]
+            args:
+            - |
+              set -x -e
+              cd /workspace
+
+              DEPLOYMENT_NAME="rock-8-$$BUILD_ID_SHORT"
+
+              RUN_CLEANUP=true
+              # Trap function that runs the rescue playbook (terraform destroy)
+              # if the pod is terminated by Kueue, fails, or is cancelled.
+              cleanup_pod() {
+                local exit_code=\$$?
+                trap - EXIT SIGTERM SIGINT ERR
+                set +e
+
+                if [ "\$$RUN_CLEANUP" = "false" ]; then
+                  exit 0
+                fi
+                if [ \$$exit_code -eq 0 ]; then exit_code=1; fi
+
+                echo ""
+                echo "=========================================================================="
+                echo "CAUGHT SIGTERM OR SCRIPT ERROR!"
+                echo "Halting primary Ansible execution..."
+                echo "=========================================================================="
+                if [ -n "\$${ANSIBLE_PID:-}" ]; then
+                    kill -TERM \$$ANSIBLE_PID 2>/dev/null || true
+                    wait \$$ANSIBLE_PID 2>/dev/null || true
+                fi
+
+                echo ""
+                echo "INITIATING RESCUE PLAYBOOK: Destroying leaked infrastructure for \$$DEPLOYMENT_NAME..."
+                echo "- hosts: localhost" > /workspace/cleanup-playbook.yml
+                echo "  tasks:" >> /workspace/cleanup-playbook.yml
+                echo "  - ansible.builtin.include_tasks:" >> /workspace/cleanup-playbook.yml
+                echo "      file: tools/cloud-build/daily-tests/ansible_playbooks/tasks/rescue_gcluster_failure.yml" >> /workspace/cleanup-playbook.yml
+
+                ansible-playbook /workspace/cleanup-playbook.yml -e deployment_name="\$$DEPLOYMENT_NAME" -e workspace="/workspace" || true
 
-- id: slurm-gcp-v6-rocky8
-  name: us-central1-docker.pkg.dev/$PROJECT_ID/hpc-toolkit-repo/test-runner
+                echo "Graceful cleanup finished."
+                exit \$$exit_code
+              }
+              trap cleanup_pod EXIT SIGTERM SIGINT
+
+              if [ "${_TEST_PREFIX}" == "daily-" ]; then
+                  gsutil cp gs://$$GCLUSTER_GCS_PATH/latest/gcluster-bundle.zip .
+                  unzip -o gcluster-bundle.zip
+                  chmod +x gcluster
+              else
+                  make
+              fi
+              BLUEPRINT="examples/hpc-slurm.yaml"
+              bash tools/add_ttl_label.sh \$$BLUEPRINT
+              python3 tools/fix_vpc_name.py \$$BLUEPRINT "${_TEST_PREFIX}"
+
+              ansible-playbook tools/cloud-build/daily-tests/ansible_playbooks/slurm-integration-test.yml \
+                --extra-vars="test_prefix=${_TEST_PREFIX} use_fixed_vpc=true" \
+                --user="$$OSLOGIN_USER" --extra-vars="project=$PROJECT_ID build=$$BUILD_ID_SHORT" \
+                --extra-vars="@tools/cloud-build/daily-tests/tests/slurm-v6-rocky8.yml" \
+                --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" &
+              ANSIBLE_PID=\$$!
+
+              wait \$$ANSIBLE_PID
+              RUN_CLEANUP=false
+            resources:
+              requests:
+                cpu: 200m
+                memory: "2Gi"
+                test-locks/slurm: 1
+                test-locks/slurm-gcp-v6-rocky8: 1
+              limits:
+                cpu: 1
+                memory: "2Gi"
+                test-locks/slurm: 1
+                test-locks/slurm-gcp-v6-rocky8: 1
+    EOF
+
+# 3. Submit and Monitor GKE Kueue Job with Retry
+- id: submit-and-monitor-gke-job
+  name: gcr.io/cloud-builders/gcloud
+  entrypoint: /bin/bash
+  args:
+  - -c
+  - |
+    gcloud container clusters get-credentials test-kueue-cluster --region=us-central1
+    BUILD_ID_SHORT=$$(echo "$BUILD_ID" | cut -c1-6)
+    JOB_NAME="slurm-gcp-v6-rocky8-$$BUILD_ID_SHORT"
+
+    # Cloud Build trap: If the Cloud Build step itself receives a cancellation signal,
+    # delete the GKE job so the pod receives a SIGTERM and performs infrastructure cleanup.
+    cleanup_cb() {
+      echo ""
+      echo "=========================================================================="
+      echo "PIPELINE CANCELLED: The Cloud Build step received a termination signal."
+      echo "Deleting GKE Kueue Job ($$JOB_NAME) to force the Pod to clean itself up!"
+      echo "=========================================================================="
+      kubectl delete job "$$JOB_NAME" -n default || true
+      exit 1
+    }
+    trap cleanup_cb SIGTERM SIGINT
+
+    MAX_RETRIES=280
+    RETRY_DELAY=300
+    ATTEMPT=1
+
+    while true; do
+      echo "=== ATTEMPT $$ATTEMPT: Submitting Kueue Job ==="
+      kubectl apply -f /workspace/job.yaml
+
+      set +e
+      ( bash tools/cloud-build/monitor_kueue_job.sh \
+          test-kueue-cluster \
+          us-central1 \
+          "$$JOB_NAME" \
+          default | tee /workspace/job_logs.txt; exit $${PIPESTATUS[0]} ) &
+      MONITOR_PID=$$!
+      wait $$MONITOR_PID
+      EXIT_CODE=$$?
+      set -e
+
+      if [ $$EXIT_CODE -eq 0 ]; then
+        echo "Job succeeded!"
+        break
+      fi
+
+      # Check if the failure was specifically due to a lack of GCP zone capacity.
+      # If so, retry. If it's a real error (like a terraform syntax error), fail immediately.
+      if bash tools/cloud-build/check_retriable_error.sh /workspace/job_logs.txt; then
+        echo "WARNING: Retriable error detected. Kueue Job has already been deleted. Retrying in $$RETRY_DELAY seconds..."
+      else
+        echo "ERROR: Test failed due to an actual error (not zone capacity). Failing pipeline." >&2
+        exit 1
+      fi
+
+      if [ $$ATTEMPT -ge $$MAX_RETRIES ]; then
+        echo "ERROR: Job failed to find zone capacity after $$MAX_RETRIES attempts." >&2
+        exit 1
+      fi
+
+      sleep $$RETRY_DELAY
+      ATTEMPT=$$((ATTEMPT + 1))
+    done
+
+# 4. Cleanup the built image from Registry to save costs
+- id: cleanup-image
+  name: gcr.io/cloud-builders/gcloud
   entrypoint: /bin/bash
-  env:
-  - "BUILD_ID=$BUILD_ID"
-  - "ANSIBLE_HOST_KEY_CHECKING=false"
-  - "ANSIBLE_CONFIG=/workspace/tools/cloud-build/ansible.cfg"
   args:
   - -c
   - |
-    set -x -e
-    cd /workspace
-    bash tools/get_binary.sh "${_TEST_PREFIX}"
-    BUILD_ID_FULL=$BUILD_ID
-    BUILD_ID_SHORT=$${BUILD_ID_FULL:0:6}
-    BLUEPRINT="examples/hpc-slurm.yaml"
-    bash tools/add_ttl_label.sh $${BLUEPRINT}
-    python3 tools/fix_vpc_name.py $${BLUEPRINT} "${_TEST_PREFIX}"
-
-    ansible-playbook tools/cloud-build/daily-tests/ansible_playbooks/slurm-integration-test.yml \
-      --extra-vars="test_prefix=${_TEST_PREFIX}" \
-      --extra-vars="use_fixed_vpc=true" \
-      --user=sa_106486320838376751393 --extra-vars="project=${PROJECT_ID} build=$${BUILD_ID_SHORT} full_build_id=$BUILD_ID" \
-      --extra-vars="@tools/cloud-build/daily-tests/tests/slurm-v6-rocky8.yml" \
-        --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', 'TRIAGE_GCS_BUCKET', 'TRIAGE_PROJECT_NUMBER', 'TRIAGE_INVOKER_SA', 'TRIAGE_CLOUD_RUN_URL']
+    gcloud artifacts docker images delete us-central1-docker.pkg.dev/$PROJECT_ID/hpc-toolkit-repo/test-runner:$BUILD_ID --quiet
 availableSecrets:
   secretManager:
   - versionName: projects/${PROJECT_ID}/secrets/gcluster-develop-release-bucket/versions/latest
@@ -74,3 +270,5 @@ availableSecrets:
     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/sa-email/versions/latest
+    env: 'SA_EMAIL'