Adding necessary test files and changes to test using spot VMs

Author: simrankaurb

community/modules/compute/schedmd-slurm-gcp-v6-nodeset/main.tf

@@ -14,7 +14,7 @@
 
 locals {
   # This label allows for billing report tracking based on module.
-  labels = merge(var.labels, { ghpc_module = "schedmd-slurm-gcp-v6-nodeset", ghpc_role = "compute" })
+  labels = merge(var.labels, { ghpc_module = "schedmd-slurm-gcp-v6-nodeset", ghpc_role = "compute", spot = var.enable_spot_vm })
 }
 
 module "instance_validation" {

tools/cloud-build/daily-tests/ansible_playbooks/slurm-integration-test.yml

@@ -191,6 +191,7 @@
   ignore_unreachable: true  # ensure always block will run even if SSH fails
   vars:
     key_type: "munge"
+    integration_tests_failed: false  # Initialize failure flag
   tasks:
   - name: Slurm Test Block
     vars:
@@ -209,7 +210,7 @@
     - name: Wait until Slurm key exists
       ansible.builtin.wait_for:
         path: /etc/slurm/slurm.key
-        timeout: 600 # Waits for up to 10 minutes for the file to appear
+        timeout: 600  # Waits for up to 10 minutes for the file to appear
       when: key_type == 'slurm'
 
     - name: Count Slurm nodes
@@ -223,17 +224,43 @@
       retries: 60
       delay: 15
 
-    - name: Run Integration tests for Cluster Toolkit
-      ansible.builtin.include_tasks: "{{ test }}"
-      vars:
-        login_node: "{{ login_node }}"
-        custom_vars: "{{ custom_vars }}"
-      loop: "{{ post_deploy_tests }}"
-      loop_control:
-        loop_var: test
+    - name: Block for running main integration tests
+      block:
+      - name: Run Integration tests for Cluster Toolkit
+        ansible.builtin.include_tasks: "{{ test }}"
+        vars:
+          login_node: "{{ login_node }}"
+          custom_vars: "{{ custom_vars }}"
+        loop: "{{ post_deploy_tests }}"
+        loop_control:
+          loop_var: test
+      rescue:
+      - name: Set flag indicating test failure
+        ansible.builtin.set_fact:
+          integration_tests_failed: true
+      - name: Display test failure message
+        ansible.builtin.debug:
+          msg: "A task within the integration tests failed. Preemption check might run."
 
     ## Always cleanup, even on failure
     always:
+    - name: Check if tests failed
+      ansible.builtin.debug:
+        var: integration_tests_failed
+
+    - name: Check if spot is enable
+      ansible.builtin.debug:
+        var: custom_vars.enable_spot | default(false) | bool
+
+    # Conditional Preemption Check
+    - name: Check for recent preemptions if tests failed and Spot is enabled
+      ansible.builtin.include_tasks: test-validation/test-a3m-preemption.yml
+      when:
+      - integration_tests_failed | bool
+      - custom_vars.enable_spot | default(false) | bool
+      vars:
+        custom_vars: "{{ custom_vars }}"
+
     - name: Ensure all nodes are powered down
       vars:
         wait_for_compute_nodes_to_go_down: true

tools/cloud-build/daily-tests/ansible_playbooks/test-validation/test-a3m-preemption.yml

@@ -0,0 +1,90 @@
+# Copyright 2023 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: Set up variables for preemption query
+  ansible.builtin.set_fact:
+    zone: "{{ custom_vars.zone }}"
+    slurm_cluster_name: "{{ custom_vars.slurm_cluster_name }}"
+    node_pattern: "{{ custom_vars.slurm_cluster_name }}-a3mega-*"
+    gcp_project: "{{ custom_vars.project_id }}"
+    query_window: "24h" # How far back to look for preemptions
+
+- name: Fail if gcp_project is not set
+  ansible.builtin.fail:
+    msg: "gcp_project variable is not defined. Please ensure custom_vars.project_id is passed."
+  when: gcp_project is not defined or gcp_project == ""
+
+- name: Fail if zone is not set
+  ansible.builtin.fail:
+    msg: "zone variable is not defined. Please ensure custom_vars.zone is passed."
+  when: zone is not defined or zone == ""
+
+- name: Fail if slurm_cluster_name is not set
+  ansible.builtin.fail:
+    msg: "slurm_cluster_name variable is not defined. Please ensure custom_vars.slurm_cluster_name is passed."
+  when: slurm_cluster_name is not defined or slurm_cluster_name == ""
+
+- name: Get current A3 Mega node instance details
+  ansible.builtin.shell: |
+    gcloud compute instances list --project={{ gcp_project }} --filter="name ~ '^{{ node_pattern }}' AND zone={{ zone }}" --format="json(name,id)"
+  register: a3_nodes_raw
+  changed_when: false
+  environment:
+    CLOUDSDK_CORE_PROJECT: "{{ gcp_project }}"
+  delegate_to: localhost
+
+- name: Parse instance details for query
+  ansible.builtin.set_fact:
+    a3_nodes: "{{ a3_nodes_raw.stdout | from_json }}"
+
+- name: Build instance ID filter string
+  ansible.builtin.set_fact:
+    instance_id_filter: "{{ a3_nodes | map(attribute='id') | map('regex_replace', '^(.*)$', 'resource.labels.instance_id=\"\\1\"') | join(' OR ') }}"
+  when: a3_nodes is defined and a3_nodes | length > 0
+
+- name: Debug instance ID filter
+  ansible.builtin.debug:
+    var: instance_id_filter
+  when: instance_id_filter is defined and instance_id_filter != ""
+
+- name: Query Cloud Audit Logs for any preemption events for these nodes
+  ansible.builtin.shell: |
+    gcloud logging read '
+      resource.type="gce_instance"
+      AND protoPayload.methodName="compute.instances.preempted"
+      AND log_id("cloudaudit.googleapis.com/system_event")
+      AND ({{ instance_id_filter }})
+    ' --project={{ gcp_project }} --freshness={{ query_window }} --format="json"
+  register: audit_logs_raw
+  changed_when: false
+  environment:
+    CLOUDSDK_CORE_PROJECT: "{{ gcp_project }}"
+  delegate_to: localhost
+  when: instance_id_filter is defined and instance_id_filter != ""
+
+- name: Parse preemption logs
+  ansible.builtin.set_fact:
+    preemption_logs: "{{ audit_logs_raw.stdout | from_json }}"
+  when: audit_logs_raw.stdout is defined and audit_logs_raw.stdout != "" and audit_logs_raw.stdout != "[]"
+
+- name: Report preempted nodes
+  ansible.builtin.debug:
+    msg: "Preemption Events Found in the last {{ query_window }}: {{ preemption_logs }}"
+  when: preemption_logs is defined and preemption_logs | length > 0
+
+- name: Report no preemptions found
+  ansible.builtin.debug:
+    msg: "No preemption events found for nodes matching '{{ node_pattern }}' in zone {{ zone }} in the last {{ query_window }}."
+  when: preemption_logs is not defined or preemption_logs | length == 0

tools/cloud-build/daily-tests/builds/ml-a3-megagpu-onspot-slurm-ubuntu.yaml

@@ -0,0 +1,60 @@
+# 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.custom-image
+- m.startup-script
+- slurm6
+- m.filestore
+- m.schedmd-slurm-gcp-v6-controller
+- m.schedmd-slurm-gcp-v6-login
+- m.schedmd-slurm-gcp-v6-nodeset
+- m.schedmd-slurm-gcp-v6-partition
+- m.vpc
+- m.cloud-storage-bucket
+- m.multivpc
+- m.private-service-access
+
+# cloudbuild.yaml
+timeout: 14400s  # 4hr
+steps:
+# While using static network names we are gaurding 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: "/workspace/tools/cloud-build/check_running_build.sh /workspace/tools/cloud-build/daily-tests/builds/ml-a3-mega-slurm-ubuntu.yaml"
+
+- id: ml-a3-megagpu-slurm
+  name: us-central1-docker.pkg.dev/$PROJECT_ID/hpc-toolkit-repo/test-runner
+  entrypoint: /bin/bash
+  env:
+  # General Ansible configuration
+  - "ANSIBLE_HOST_KEY_CHECKING=false"
+  - "ANSIBLE_CONFIG=/workspace/tools/cloud-build/ansible.cfg"
+
+  # Configuration Variables for the A3-MegaGPU Test (These will change for other tests)
+  - "MACHINE_TYPE=a3-megagpu-8g"
+  - "IMAGE_FAMILY=slurm-a3mega"
+  - "NUM_NODES=3"
+  - "BLUEPRINT_PATH=/workspace/examples/machine-learning/a3-megagpu-8g/a3mega-slurm-blueprint.yaml"
+  - "TEST_VARS_FILE=@tools/cloud-build/daily-tests/tests/ml-a3-megagpu-onspot-slurm-ubuntu.yml"
+  - "INSTANCE_PREFIX=a3" # Use a generic prefix that the script will complete
+  - "PROJECT_ID=$PROJECT_ID"
+  - "BUILD_ID=$BUILD_ID"
+  args:
+  # Execute the script, which now relies on the exported environment variables.
+  - -c
+  - |
+    set -x
+    /workspace/tools/cloud-build/instance_provisioning.sh

tools/cloud-build/daily-tests/tests/ml-a3-megagpu-onspot-slurm-ubuntu.yml

@@ -0,0 +1,59 @@
+# Copyright 2024 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.
+
+---
+
+# region, zone must be defined in build file with --extra-vars flag!
+test_name: a3m-slurm
+deployment_name: a3m-slurm-{{ build }}
+slurm_cluster_name: "a3m{{ build[0:4] }}"
+workspace: /workspace
+blueprint_yaml: "{{ workspace }}/examples/machine-learning/a3-megagpu-8g/a3mega-slurm-blueprint.yaml"
+login_node: "{{ slurm_cluster_name }}-login-*"
+controller_node: "{{ slurm_cluster_name }}-controller"
+network: "{{ deployment_name }}-net-0"
+sub_network: "{{ deployment_name }}-sub-net-0"
+post_deploy_tests:
+- test-validation/test-mounts.yml
+- test-validation/test-partitions.yml
+- test-validation/test-enroot.yml
+- test-validation/test-gpus-slurm.yml
+post_destroy_tasks:
+- post-destroy-tasks/delete-image.yml
+custom_vars:
+  gpu_count: 8
+  gpu_partition: a3mega
+  test_persistenced: true
+  partitions:
+  - a3mega
+  - debug
+  mounts:
+  - /home
+  project_id: "{{ project }}"
+  zone: "{{ zone }}"
+  slurm_cluster_name: "{{ slurm_cluster_name }}"
+  enable_spot: true
+cli_deployment_vars:
+  network_name_system: "{{ network }}"
+  subnetwork_name_system: "{{ sub_network}}"
+  region: "{{ region }}"
+  zone: "{{ zone }}"
+  slurm_cluster_name: "{{ slurm_cluster_name }}"
+  disk_size_gb: 200
+  a3mega_cluster_size: 2
+  enable_ops_agent: "true"
+  enable_nvidia_dcgm: "true"
+  enable_nvidia_persistenced: true
+  final_image_family: "{{ deployment_name}}-u22"
+  a3mega_enable_spot_vm: true