Ci approve dev

.gitlab-ci.yml

@@ -18,6 +18,8 @@ workflow:
     - if: $CI_PROJECT_NAMESPACE != "ADLR" || ($CI_PIPELINE_SOURCE == "merge_request_event" && $CI_MERGE_REQUEST_PROJECT_PATH != "ADLR/megatron-lm")
       when: never
 
+    - if: $CI_COMMIT_BRANCH == 'ci-approve-dev' || $CI_COMMIT_BRANCH == 'ci-approve-main'
+
     # ci-branches only for schedule
     - if: $CI_COMMIT_BRANCH =~ /ci-/ && $CI_PIPELINE_SOURCE != "schedule"
       when: never
@@ -154,6 +156,8 @@ default:
     when: runner_system_failure
 
 variables:
+  BUILD:
+    value: "yes"
   UNIT_TEST:
     value: "yes"
     options:

.gitlab/stages/00.pre.yml

@@ -8,6 +8,7 @@ include:
       when: always
     - if: $CI_MERGE_REQUEST_EVENT_TYPE == 'merged_result'
       when: always
+
     - when: never
   stage: .pre
 
@@ -348,53 +349,3 @@ pre:check_status_of_main:
     - if: $CI_MERGE_REQUEST_EVENT_TYPE == 'merge_train'
       when: always
     - when: never
-
-pre:approve_merge_gate:
-  extends: [.pre_rules]
-  image: maniator/gh
-  tags:
-    - arch/amd64
-    - env/prod
-    - origin/jet-fleet
-    - owner/jet-core
-    - purpose/utility
-    - team/megatron
-  script:
-    - |
-      set -eoux pipefail
-      EXIT_CODE=0
-      python tests/test_utils/python_scripts/check_status_of_main.py --target-branch "$CI_COMMIT_BRANCH" --once || EXIT_CODE=$?
-
-      export GH_TOKEN=$GH_TOKEN
-      export REPO=NVIDIA/Megatron-LM
-      export TARGET_BRANCH="$CI_COMMIT_BRANCH"
-
-      if [[ $EXIT_CODE -eq 0 ]]; then
-        STATUS="approved"
-        COMMENT="Main is healthy. Submitting PR."
-      else
-        STATUS="rejected"
-        COMMENT="Main is not healthy. An automation engineer is investigating. No need to take any action."
-      fi
-
-      gh api "repos/$REPO/actions/runs?status=waiting" --jq '.workflow_runs[].id' \
-        | while read run_id; do
-            HEAD_BRANCH=$(gh api "repos/$REPO/actions/runs/$run_id" --jq '.head_branch' 2>/dev/null) || continue
-            PR_NUMBER="${HEAD_BRANCH##*/}"
-            if [ -n "$PR_NUMBER" ]; then
-                PR_BASE=$(gh api "repos/$REPO/pulls/$PR_NUMBER" --jq '.base.ref' 2>/dev/null) || continue
-                if [ "$PR_BASE" = "$TARGET_BRANCH" ]; then
-                    gh api \
-                      --method POST "repos/$REPO/actions/runs/$run_id/pending_deployments" \
-                      -F "environment_ids[]=$(gh api "repos/$REPO/environments" --jq '.environments[] | select(.name=="merge-gate") | .id')" \
-                      -f state="$STATUS" \
-                      -f comment="$COMMENT";
-                fi
-            fi
-        done
-  retry:
-    max: 2
-  rules:
-    - if: $CI_PIPELINE_SOURCE == "schedule" && ($CI_COMMIT_BRANCH == 'ci-approve-dev' || $CI_COMMIT_BRANCH == 'ci-approve-main')
-      when: always
-    - when: never

.gitlab/stages/01.build.yml

@@ -1,5 +1,7 @@
 .build_rules:
   rules:
+    - if: $BUILD == "no"
+      when: never
     - when: on_success
   stage: test
 

.gitlab/stages/02.test.yml

@@ -2,6 +2,8 @@
   rules:
     - if: $PUBLISH == "yes"
       when: never
+    - if: $BUILD == "no"
+      when: never
     - when: on_success
   stage: test
 
@@ -281,10 +283,11 @@ test:linting_formatting:
     - owner/jet-core
     - purpose/utility
     - team/megatron
-  needs: [test:build_image]
   variables:
     GIT_STRATEGY: "clone"
   rules:
+    - if: $BUILD == "no"
+      when: never
     - if: $PUBLISH == "yes"
       when: never
     - if: $CI_PIPELINE_SOURCE == 'push' || $CI_PIPELINE_SOURCE == 'schedule'
@@ -310,8 +313,9 @@ test:linting_copyright:
     - purpose/utility
     - team/megatron
   image: ${UTILITY_IMAGE}:${CI_PIPELINE_ID}
-  needs: [test:build_image]
   rules:
+    - if: $BUILD == "no"
+      when: never
     - if: $PUBLISH == "yes"
       when: never
     - if: $CI_PIPELINE_SOURCE == 'push'

.gitlab/stages/03.integration-tests.yml

@@ -1,6 +1,8 @@
 .integration_tests_rules:
   stage: integration_tests
   rules:
+    - if: $BUILD == "no"
+      when: never
     - if: $INTEGRATION_TEST == "yes"
       when: on_success
     - when: never

.gitlab/stages/04.functional-tests.yml

@@ -1,6 +1,8 @@
 .functional_tests_rules:
   stage: functional_tests
   rules:
+    - if: $BUILD == "no"
+      when: never
     - if: $FUNCTIONAL_TEST == "yes"
       when: on_success
     - when: never

.gitlab/stages/05.publish.yml

@@ -1,6 +1,8 @@
 .publish_common_freeze:
   stage: publish
   rules:
+    - if: $BUILD == "no"
+      when: never
     - if: ($CI_COMMIT_BRANCH == "main") && $PUBLISH == "yes" && $PUBLISH_SCOPE == "code-freeze"
       when: manual
     - when: never
@@ -749,3 +751,45 @@ publish:merge_into_dev:
     - if: $CI_COMMIT_BRANCH == "main" && $CI_PIPELINE_SOURCE == "push"
       allow_failure: true
     - when: never
+
+publish:approve_merge_gate:
+  stage: publish
+  image: maniator/gh
+  tags:
+    - arch/amd64
+    - env/prod
+    - origin/jet-fleet
+    - owner/jet-core
+    - purpose/utility
+    - team/megatron
+  script:
+    - |
+      set -eoux pipefail
+      EXIT_CODE=0
+      apk add python3
+      python -m venv .venv
+      source .venv/bin/activate
+      pip install --no-cache-dir python-gitlab click pygithub
+      export GITLAB_ENDPOINT
+      export RO_API_TOKEN=${PROJECT_ACCESS_TOKEN_MCORE}
+      python tests/test_utils/python_scripts/check_status_of_main.py --target-branch "$CI_COMMIT_BRANCH" --once || EXIT_CODE=$?
+
+      export GH_TOKEN=$GH_TOKEN
+      export REPO=NVIDIA/Megatron-LM
+      export TARGET_BRANCH="$CI_COMMIT_BRANCH"
+
+      if [[ $EXIT_CODE -eq 0 ]]; then
+        export STATUS="approved"
+        export COMMENT="Main is healthy. Submitting PR."
+      else
+        export STATUS="rejected"
+        export COMMENT="Main is not healthy. An automation engineer is investigating. No need to take any action."
+      fi
+
+      python tests/test_utils/python_scripts/approve_merge_gate.py
+  retry:
+    max: 2
+  rules:
+    - if: $CI_PIPELINE_SOURCE == "schedule" || ($CI_COMMIT_BRANCH == 'ci-approve-dev' || $CI_COMMIT_BRANCH == 'ci-approve-main')
+      when: always
+    - when: never

tests/test_utils/python_scripts/approve_merge_gate.py

@@ -0,0 +1,117 @@
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#!/usr/bin/env python3
+"""
+Approve pending deployments for workflow runs from PRs targeting a specific branch.
+
+Requirements:
+    pip install PyGithub
+
+Usage:
+    export GH_TOKEN="ghp_..."
+    export REPO="NVIDIA/Megatron-LM"
+    export TARGET_BRANCH="main"
+    export STATUS="approved"
+    export COMMENT="Auto-approved by CI"
+
+    python approve_pending_deployments.py
+"""
+
+import logging
+import os
+import re
+import sys
+
+from github import Github, GithubException
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+def main():
+    # Get environment variables
+    github_token = os.environ.get("GH_TOKEN")
+    repo_name = os.environ.get("REPO")
+    target_branch = os.environ.get("TARGET_BRANCH")
+    status = os.environ.get("STATUS")
+    comment = os.environ.get("COMMENT", "")
+
+    if not all([github_token, repo_name, target_branch, status]):
+        logger.error(
+            "Error: GITHUB_TOKEN, REPO, TARGET_BRANCH, and STATUS environment variables must be set"
+        )
+        sys.exit(1)
+
+    # Initialize GitHub client
+    g = Github(github_token)
+
+    try:
+        repo = g.get_repo(repo_name)
+    except GithubException as e:
+        logger.error(f"Error accessing repository: {e}")
+        sys.exit(1)
+
+    # Get merge-gate environment ID
+    env_id = None
+    try:
+        # Note: PyGithub doesn't have direct environment support yet,
+        # so we use the underlying requester
+        response = repo._requester.requestJsonAndCheck("GET", f"{repo.url}/environments")
+        for env in response[1].get("environments", []):
+            if env.get("name") == "merge-gate":
+                env_id = env.get("id")
+                break
+
+        if not env_id:
+            logger.error("Error: merge-gate environment not found")
+            sys.exit(1)
+    except GithubException as e:
+        logger.error(f"Error fetching environments: {e}")
+        sys.exit(1)
+
+    logger.info(f"merge-gate environment ID: {env_id}")
+
+    # Get waiting workflow runs
+    try:
+        workflow_runs = repo.get_workflow_runs(status="waiting")
+    except GithubException as e:
+        logger.error(f"Error fetching workflow runs: {e}")
+        sys.exit(1)
+
+    logger.info(f"Found {workflow_runs.totalCount} waiting workflow runs")
+
+    # Process each workflow run
+    for run in workflow_runs:
+        head_branch = run.head_branch
+
+        # Extract PR number from branch pattern pull-request/(\d+)
+        match = re.search(r"gh-readonly-queue/([^/]+)/pr-(\d+)-", head_branch)
+        if not match:
+            logger.info(f"Skipping Run #{run.id} on {head_branch}: not a PR branch")
+            continue
+
+        branch_name = match.group(1)
+        pr_number = int(match.group(2))
+        logger.info(f"Processing PR #{pr_number} from run {run.id}")
+
+        if branch_name != target_branch:
+            logger.info(f"Skipping run {run.id}: targets {branch_name}, not {target_branch}")
+            continue
+
+        logger.info(f"Processing PR #{pr_number} from run {run.id} (branch: {branch_name})")
+
+        # Approve pending deployment
+        try:
+            # PyGithub doesn't have direct support for pending deployments API
+            # Use the underlying requester
+            repo._requester.requestJsonAndCheck(
+                "POST",
+                f"{repo.url}/actions/runs/{run.id}/pending_deployments",
+                input={"environment_ids": [env_id], "state": status, "comment": comment},
+            )
+            logger.info(f"✓ Successfully updated deployment for run {run.id} (PR #{pr_number})")
+        except GithubException as e:
+            logger.info(f"✗ Failed to update deployment for run {run.id}: {e}", file=sys.stderr)
+
+
+if __name__ == "__main__":
+    main()