diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml index 6b46d92aacb..9999aacc224 100644 --- a/.gitlab-ci.yml +++ b/.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: diff --git a/.gitlab/stages/00.pre.yml b/.gitlab/stages/00.pre.yml index 5c74073ff14..dca3a7b47ae 100644 --- a/.gitlab/stages/00.pre.yml +++ b/.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 diff --git a/.gitlab/stages/01.build.yml b/.gitlab/stages/01.build.yml index 2fd9e1f32e6..0658daaa9ec 100644 --- a/.gitlab/stages/01.build.yml +++ b/.gitlab/stages/01.build.yml @@ -1,5 +1,7 @@ .build_rules: rules: + - if: $BUILD == "no" + when: never - when: on_success stage: test diff --git a/.gitlab/stages/02.test.yml b/.gitlab/stages/02.test.yml index 699bef68181..c8542ff7409 100644 --- a/.gitlab/stages/02.test.yml +++ b/.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' diff --git a/.gitlab/stages/03.integration-tests.yml b/.gitlab/stages/03.integration-tests.yml index df4d84234bb..824721b9fb1 100644 --- a/.gitlab/stages/03.integration-tests.yml +++ b/.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 diff --git a/.gitlab/stages/04.functional-tests.yml b/.gitlab/stages/04.functional-tests.yml index 7fe8aad0771..4c50e3158e3 100644 --- a/.gitlab/stages/04.functional-tests.yml +++ b/.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 diff --git a/.gitlab/stages/05.publish.yml b/.gitlab/stages/05.publish.yml index 695479179c5..c925a92d1c3 100644 --- a/.gitlab/stages/05.publish.yml +++ b/.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 diff --git a/tests/test_utils/python_scripts/approve_merge_gate.py b/tests/test_utils/python_scripts/approve_merge_gate.py new file mode 100644 index 00000000000..dbd4ef99b44 --- /dev/null +++ b/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()