From 667bbfd53249a2d6fee95324f84c19e30ac7f626 Mon Sep 17 00:00:00 2001
From: Oliver Koenig <okoenig@nvidia.com>
Date: Sat, 31 Aug 2024 10:13:24 -0700
Subject: [PATCH] ADLR/megatron-lm!2019 - ci: Create CI branches

---
 .gitlab/stages/00.pre.yml              | 28 +++++++++++++++++++++++---
 .gitlab/stages/01.tests.yml            |  7 ++++++-
 .gitlab/stages/02.functional-tests.yml |  2 +-
 3 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/.gitlab/stages/00.pre.yml b/.gitlab/stages/00.pre.yml
index 3afdaf5d9c..1c7b120b75 100644
--- a/.gitlab/stages/00.pre.yml
+++ b/.gitlab/stages/00.pre.yml
@@ -3,7 +3,7 @@ include:
 
 mirror_to_github:
   rules:
-    - if: '$CI_COMMIT_REF_PROTECTED == "true"'
+    - if: '$CI_COMMIT_BRANCH == $CI_DEFAULT_BRANCH'
     - when: never
   tags: [mcore-docker-node-small]
   stage: .pre
@@ -14,7 +14,29 @@ mirror_to_github:
     - git checkout $CI_COMMIT_BRANCH
     - git remote add github https://ko3n1g:$GH_TOKEN@github.com/NVIDIA/Megatron-LM.git || true
     - git push -u github $CI_COMMIT_BRANCH
-  
+
+create_ci_branches:
+  rules:
+    - if: '$CI_COMMIT_BRANCH == $CI_DEFAULT_BRANCH'
+  parallel:
+    matrix:
+      - branch: ci-unit-test-extended
+      - branch: ci-rebuild-mcore-nemo-image
+      - branch: ci-mr-a100
+      - branch: ci-nightly-a100
+      - branch: ci-weekly-a100
+      - branch: ci-weekly-h100
+      - branch: ci-pre-release
+  tags: [mcore-docker-node-small]
+  stage: .pre
+  image: python:3.10
+  variables:
+    GIT_STRATEGY: "clone"
+  script:
+    - git remote set-url origin "https://gitlab-ci-token:${PROJECT_ACCESS_TOKEN_MCORE}@${GITLAB_ENDPOINT}/adlr/megatron-lm.git"
+    - git switch --force-create $branch;
+    - git push --force -u origin $branch
+
 label_merge_request:
   rules: 
     - if: $CI_PIPELINE_SOURCE == "merge_request_event"
@@ -44,7 +66,7 @@ clean_docker_node:
   tags: [mcore-docker-node-small]
   script:
     - export DOCKER_HOST='unix:///var/run/docker.sock'
-    - docker system prune -a --filter "until=48h" -f
+    - docker system prune -a --filter "until=48h" -f || true
 
 check_milestone:
   rules: 
diff --git a/.gitlab/stages/01.tests.yml b/.gitlab/stages/01.tests.yml
index 04f7a6ab7f..44ded54afd 100644
--- a/.gitlab/stages/01.tests.yml
+++ b/.gitlab/stages/01.tests.yml
@@ -1,6 +1,6 @@
 .tests_common:
   rules:
-    - if: $CI_MERGE_REQUEST_TARGET_BRANCH_PROTECTED != "true"
+    - if: $CI_PIPELINE_SOURCE == 'merge_request_event' && $CI_MERGE_REQUEST_TARGET_BRANCH_PROTECTED != "true"
       allow_failure: true
       when: always
     - when: always
@@ -66,6 +66,11 @@ build_image:
       docker push ${IMAGE}:${CI_PIPELINE_ID}
       docker push ${IMAGE}:${CI_MERGE_REQUEST_IID:-noop}
 
+      if [[ "$CI_COMMIT_BRANCH" == "ci-nightly-a100" ]]; then
+        docker tag ${IMAGE}:${CI_PIPELINE_ID} ${IMAGE}:nightly
+        docker push ${IMAGE}:nightly
+      fi
+
       if [[ "$CI_COMMIT_BRANCH" == "$CI_DEFAULT_BRANCH" ]]; then
         docker tag ${IMAGE}:${CI_PIPELINE_ID} ${IMAGE}:buildcache
         docker push ${IMAGE}:buildcache
diff --git a/.gitlab/stages/02.functional-tests.yml b/.gitlab/stages/02.functional-tests.yml
index 19f98e2730..a79259bf4c 100644
--- a/.gitlab/stages/02.functional-tests.yml
+++ b/.gitlab/stages/02.functional-tests.yml
@@ -1,7 +1,7 @@
 .jet_common:
   stage: functional_tests
   rules:
-    - if: $FUNCTIONAL_TEST == "yes" && $CI_MERGE_REQUEST_TARGET_BRANCH_PROTECTED != "true"
+    - if: $FUNCTIONAL_TEST == "yes" && ($CI_PIPELINE_SOURCE == 'merge_request_event' && $CI_MERGE_REQUEST_TARGET_BRANCH_PROTECTED != "true")
       allow_failure: true
     - if: $FUNCTIONAL_TEST == "yes"
     - when: never