Merge pull request GoogleCloudPlatform#4796 from saara-tyagi27/a3m-nccl-test

Add NCCL integration test for A3 Mega GPUs

Author: saara-tyagi27

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

@@ -0,0 +1,97 @@
+# 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.
+
+---
+- name: Run prerequisite NCCL scripts
+  shell: |
+    set -x -e
+    cd {{ ansible_user_dir }}
+    rm -rf cluster-toolkit
+    git clone https://github.com/GoogleCloudPlatform/cluster-toolkit.git
+    cd cluster-toolkit/{{ nccl_test_path }}
+    bash import_pytorch_container.sh
+    sbatch --wait build-nccl-tests.sh
+  args:
+    chdir: "{{ ansible_user_dir }}"
+    executable: /bin/bash
+  register: build_result
+
+- name: Wait for aperture devices on compute nodes
+  shell: |
+    srun -N 2 --partition=a3mega --pty bash -c '
+      max_attempts=60 # Wait for up to 5 minute
+      attempt=0
+      while [ -z "$(ls -A /dev/aperture_devices 2>/dev/null)" ]; do
+        echo "Node $(hostname): Waiting for aperture devices... (Attempt $(($attempt + 1)))"
+        if [ $attempt -ge $max_attempts ]; then
+          echo "ERROR: Node $(hostname): Aperture devices not found after $max_attempts attempts." >&2
+          exit 1
+        fi
+        attempt=$(($attempt + 1))
+        sleep 5
+      done
+      echo "Node $(hostname): Aperture devices are ready."
+    '
+  args:
+    chdir: "{{ ansible_user_dir }}/cluster-toolkit/{{ nccl_test_path }}"
+    executable: /bin/bash
+
+- name: Submit NCCL test job and get job ID
+  shell: "sbatch run-nccl-tests.sh | awk '{print $4}'"
+  args:
+    chdir: "{{ ansible_user_dir }}/cluster-toolkit/{{ nccl_test_path }}"
+    executable: /bin/bash
+  register: sbatch_output
+  changed_when: false
+
+- name: Set job_id fact
+  set_fact:
+    job_id: "{{ sbatch_output.stdout | trim }}"
+
+- name: Wait for NCCL test job to complete
+  command: "sacct -j {{ job_id }} -n -o State"
+  register: job_status
+  until: "'COMPLETED' in job_status.stdout or 'FAILED' in job_status.stdout or 'TIMEOUT' in job_status.stdout"
+  retries: 60 # wait for 30 minutes
+  delay: 30
+  changed_when: false
+
+- name: Check final job status
+  fail:
+    msg: "NCCL test job {{ job_id }} failed with status {{ job_status.stdout }}. Check slurm-{{ job_id }}.out on the login node."
+  when: "'COMPLETED' not in job_status.stdout"
+
+- name: Verify NCCL test output
+  shell: "cat slurm-{{ job_id }}.out | grep '# Avg bus bandwidth'"
+  args:
+    chdir: "{{ ansible_user_dir }}/cluster-toolkit/{{ nccl_test_path }}"
+    executable: /bin/bash
+  register: nccl_output
+  changed_when: false
+
+- name: Display NCCL Test Result
+  debug:
+    msg: "NCCL Test Result: {{ nccl_output.stdout }}"
+  when: nccl_output.stdout is defined and nccl_output.stdout != ""
+
+- name: Extract average bus bandwidth
+  shell: "echo '{{ nccl_output.stdout }}' | awk -F':' '{print $NF}'"
+  register: avg_bus_bandwidth
+  changed_when: false
+  when: nccl_output.stdout is defined and nccl_output.stdout != ""
+
+- name: Ensure average bus bandwidth is sufficient
+  fail:
+    msg: "Average bus bandwidth is {{ avg_bus_bandwidth.stdout | trim }} GB/s, which is below the threshold of 30 GB/s."
+  when: (avg_bus_bandwidth.stdout | default('0') | float) < 30

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

@@ -25,12 +25,14 @@ controller_node: "{{ slurm_cluster_name }}-controller"
 region: us-west4
 zone: us-west4-a
 network: "{{ deployment_name }}-net-0"
+nccl_test_path: "examples/machine-learning/a3-megagpu-8g/nccl-tests"
 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
+- test-validation/test-nccl.yml
 post_destroy_tasks:
 - post-destroy-tasks/delete-image.yml
 custom_vars: