testing-nccl-test-dymanic-branch

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

@@ -13,15 +13,77 @@
 # limitations under the License.
 
 ---
+# Detect PR number from Cloud Build
+- name: Get PR number
+  ansible.builtin.set_fact:
+    pr_number: "{{ lookup('env','_PR_NUMBER') | default('', true) }}"
+
+# Detect branch name (manual run or Cloud Build trigger)
+- name: Get branch name
+  ansible.builtin.set_fact:
+    branch_name: "{{ lookup('env','BRANCH_NAME') | default('', true) }}"
+
+# Determine which branch should run
+- name: Decide branch to run
+  ansible.builtin.set_fact:
+    target_branch: >-
+      {% if pr_number != '' %}
+      PR
+      {% elif branch_name != '' %}
+      {{ branch_name }}
+      {% else %}
+      develop
+      {% endif %}
+
+# Show trigger information
+- name: Show build info
+  ansible.builtin.debug:
+    msg:
+    - "PR Number      : {{ pr_number }}"
+    - "Branch Name    : {{ branch_name }}"
+    - "Target Branch  : {{ target_branch }}"
+
+# Remove old repo if exists
+- name: Remove existing cluster-toolkit repo
+  ansible.builtin.file:
+    path: "{{ ansible_user_dir }}/cluster-toolkit"
+    state: absent
+
+# Clone repo (default develop)
+- name: Clone cluster-toolkit repository
+  ansible.builtin.git:
+    repo: https://github.com/GoogleCloudPlatform/cluster-toolkit.git
+    dest: "{{ ansible_user_dir }}/cluster-toolkit"
+    version: develop
+    force: yes
+
+# Fetch and checkout PR branch
+- name: Checkout PR branch safely
+  ansible.builtin.shell: |
+    git fetch origin pull/{{ pr_number }}/head:pr-{{ pr_number }}
+    git checkout pr-{{ pr_number }}
+  args:
+    chdir: "{{ ansible_user_dir }}/cluster-toolkit"
+  when: pr_number != ''
+
+# Checkout normal branch
+- name: Checkout specified branch
+  ansible.builtin.shell: |
+    git checkout {{ branch_name }}
+  args:
+    chdir: "{{ ansible_user_dir }}/cluster-toolkit"
+  when:
+  - pr_number == ''
+  - branch_name != ''
+
+# Run NCCL tests
 - name: Run prerequisite NCCL scripts
-  shell: |
+  ansible.builtin.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 }}
+    cd {{ ansible_user_dir }}/cluster-toolkit/{{ nccl_test_path }}
     bash import_pytorch_container.sh
     sbatch --wait build-nccl-tests.sh
+
   args:
     chdir: "{{ ansible_user_dir }}"
     executable: /bin/bash

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

@@ -29,6 +29,7 @@ tags:
 
 substitutions:
   _TEST_PREFIX: "" # Default to no prefix
+  _PR_NUMBER: ""
 
 timeout: 14400s  # 4hr
 steps:
@@ -48,6 +49,8 @@ steps:
   - "INSTANCE_PREFIX=a3msp"
   - "PROJECT_ID=$PROJECT_ID"
   - "BUILD_ID=$BUILD_ID"
+  - "PR_NUMBER=$_PR_NUMBER"
+  - "BRANCH_NAME=$BRANCH_NAME"
   - "OPTIONS_GCS_PATH=gs://hpc-ctk1357/a3moptions.txt"
   - "ENABLE_SPOT_FALLBACK=true"
   args: