cicd-approve-test-queue.yml

# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
#
# 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: Approve Test Queue

on:
  schedule:
    - cron: "*/5 * * * *" # Runs every 5 minutes
  # Scheduled events can be delayed, so refill released queue slots immediately.
  workflow_run:
    workflows: ["CICD NeMo"]
    types: [completed]
  workflow_dispatch: # Allows manual triggering

concurrency:
  group: approve-test-queue
  cancel-in-progress: false

jobs:
  approve-queue:
    runs-on: ubuntu-latest
    environment: main
    if: github.repository == 'NVIDIA-NeMo/Megatron-Bridge'
    strategy:
      matrix:
        branch: [main, others, workflow_dispatch]
        contributor_type: [internal, external]
    steps:
      - name: Checkout repository
        uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6

      - name: Set up Python
        uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6
        with:
          python-version: "3.12"

      - name: Install dependencies
        run: |
          python -m pip install --upgrade pip
          pip install requests

      - name: Download SSO users list
        run: |
          gh release download v0.1.0 \
            --repo NVIDIA-GitHub-Management/github-audits \
            --pattern users_sso.json \
            --output users_sso.json || echo '{}' > users_sso.json
        env:
          GH_TOKEN: ${{ secrets.NVIDIA_MANAGEMENT_ORG_PAT }}

      - name: Approve waiting deployments
        env:
          GITHUB_TOKEN: ${{ secrets.PAT }}
          MAX_CONCURRENCY: ${{ vars.MAX_CONCURRENCY || 1 }}
          MAX_CONCURRENCY_EXTERNAL: ${{ vars.MAX_CONCURRENCY_EXTERNAL || 3 }}
          MAX_CONCURRENCY_WORKFLOW_DISPATCH: ${{ vars.MAX_CONCURRENCY || 1 }}
          CONTRIBUTOR_TYPE: ${{ matrix.contributor_type }}
          MATRIX_BRANCH: ${{ matrix.branch }}
          SSO_USERS_FILE: users_sso.json
          PYTHONUNBUFFERED: 1
        shell: python
        run: |
          import os
          import json
          import requests
          import re
          import time

          # GitHub API configuration
          GITHUB_TOKEN = os.environ["GITHUB_TOKEN"]
          REPO = os.environ["GITHUB_REPOSITORY"]
          CONTRIBUTOR_TYPE = os.environ["CONTRIBUTOR_TYPE"]
          MATRIX_BRANCH = os.environ["MATRIX_BRANCH"]
          if MATRIX_BRANCH == "workflow_dispatch":
              MAX_CONCURRENCY = int(os.environ["MAX_CONCURRENCY_WORKFLOW_DISPATCH"])
              API_BASE = f"https://api.github.com/repos/{REPO}"
              WORKFLOW_NAME = "CICD NeMo"
          else:
              if CONTRIBUTOR_TYPE == "external":
                  MAX_CONCURRENCY = int(os.environ["MAX_CONCURRENCY_EXTERNAL"])
              else:
                  MAX_CONCURRENCY = int(os.environ["MAX_CONCURRENCY"])
              API_BASE = "https://api.github.com/repos/NVIDIA-NeMo/Megatron-Bridge"
              WORKFLOW_NAME = "CICD NeMo"

          # Load SSO users for internal/external classification
          with open(os.environ["SSO_USERS_FILE"]) as f:
              sso_users = json.load(f)

          # Headers for GitHub API
          headers = {
              "Authorization": f"token {GITHUB_TOKEN}",
              "Accept": "application/vnd.github.v3+json",
              "X-GitHub-Api-Version": "2022-11-28",
          }

          def make_request(endpoint, method="GET", data=None, max_retries=5):
              """Make a request to the GitHub API with retry on transient errors."""
              url = f"{API_BASE}/{endpoint}"
              for attempt in range(max_retries):
                  try:
                      if method == "GET":
                          response = requests.get(url, headers=headers, timeout=30)
                      else:
                          response = requests.post(url, headers=headers, json=data, timeout=30)
                      if response.status_code == 429:
                          retry_after = int(response.headers.get("Retry-After", 2 ** attempt))
                          print(f"Rate limited on {endpoint}, retrying in {retry_after}s (attempt {attempt + 1}/{max_retries})")
                          time.sleep(retry_after)
                          continue
                      if response.status_code >= 500:
                          delay = 2 ** attempt
                          print(f"Server error {response.status_code} on {endpoint}, retrying in {delay}s (attempt {attempt + 1}/{max_retries})")
                          time.sleep(delay)
                          continue
                      response.raise_for_status()
                      return response.json()
                  except (requests.exceptions.ConnectionError, requests.exceptions.Timeout) as e:
                      delay = 2 ** attempt
                      print(f"Transient error on {endpoint}: {e}, retrying in {delay}s (attempt {attempt + 1}/{max_retries})")
                      time.sleep(delay)
                  except requests.exceptions.RequestException as e:
                      print(f"Error making request to {endpoint}: {str(e)}")
                      if hasattr(e, 'response') and e.response is not None:
                          print(f"Response: {e.response.text}")
                      return None
              print(f"Max retries ({max_retries}) exceeded for {endpoint}")
              return None

          def is_internal_contributor(pr_info):
              """Return True if the PR author is a member of NVIDIA or NVIDIA-NeMo org (is_org_member)."""
              login = pr_info.get("user", {}).get("login", "")
              org_roles = sso_users.get(login, {}).get("org_roles", [])
              return any(role in ("NVIDIA:Member", "NVIDIA-NeMo:Member") for role in org_roles)

          def get_pr_base_branch(workflow_run):
              """
              Return the base branch of the PR associated with a workflow run, or None.
              Extracts PR number from head branch like 'pull-request/1913' and fetches PR info.
              Returns (base_branch, pr_info) tuple, or (None, None) if not a PR run.
              """
              print(workflow_run.get("head_branch", ""))
              head_branch = workflow_run.get("head_branch", "")
              match = re.match(r"pull-request/(\d+)", head_branch)
              if not match:
                  return None, None  # Not a PR branch pattern

              pr_number = int(match.group(1))

              # Fetch PR info from GitHub API
              pr_info = make_request(f"pulls/{pr_number}")
              if not pr_info:
                  print(f"Failed to fetch PR #{pr_number}")
                  return None, None

              base_branch = pr_info.get("base", {}).get("ref")
              return base_branch, pr_info

          def is_internal_actor(workflow_run):
              """Return True if the actor who triggered the workflow run is an NVIDIA/NVIDIA-NeMo member."""
              login = (workflow_run.get("triggering_actor") or workflow_run.get("actor") or {}).get("login", "")
              org_roles = sso_users.get(login, {}).get("org_roles", [])
              return any(role in ("NVIDIA:Member", "NVIDIA-NeMo:Member") for role in org_roles)

          def is_pr_run(workflow_run):
              """Return True if this run was triggered by a PR (head_branch matches pull-request/<number>)."""
              return bool(re.match(r"pull-request/\d+", workflow_run.get("head_branch", "")))

          def is_workflow_dispatch_run(workflow_run):
              """Return True if this run was manually triggered (head_branch starts with mcore-testing-)."""
              return workflow_run.get("head_branch", "").startswith("mcore-testing-")

          def matches_queue(workflow_run, target_branch, contributor_type):
              """
              Return True if the workflow run belongs to this queue cell:
              matching target branch AND matching contributor type (internal/external).

              workflow_dispatch runs (head_branch: mcore-testing-*) are routed to the 'workflow_dispatch' queue only.
              PR runs (head_branch: pull-request/<n>) are routed to 'main' or 'others' queues only.
              """
              if target_branch == "workflow_dispatch":
                  if not is_workflow_dispatch_run(workflow_run):
                      return False
                  internal = is_internal_actor(workflow_run)
                  contributor_match = (contributor_type == "internal") == internal
                  if contributor_match:
                      actor = (workflow_run.get("triggering_actor") or workflow_run.get("actor") or {}).get("login", "unknown")
                      print(f"workflow_dispatch run by {actor}, contributor_type={contributor_type} (internal={internal})")
                  return contributor_match

              # PR queue: skip non-PR runs
              if not is_pr_run(workflow_run):
                  return False

              base_branch, pr_info = get_pr_base_branch(workflow_run)
              if base_branch is None:
                  return False

              branch_match = (
                  (base_branch == target_branch) or
                  (base_branch != "main" and base_branch != "dev" and target_branch == "others")
              )
              if not branch_match:
                  return False

              pr_number = re.match(r"pull-request/(\d+)", workflow_run.get("head_branch", "")).group(1)
              internal = is_internal_contributor(pr_info)
              contributor_match = (contributor_type == "internal") == internal
              if branch_match and contributor_match:
                  print(f"PR #{pr_number} targets {target_branch}, contributor_type={contributor_type} (internal={internal})")
              return branch_match and contributor_match

          # Get current running and queued workflows
          print(f"\n=== Queue cell: branch=${{ matrix.branch }}, contributor_type={CONTRIBUTOR_TYPE} ===")
          print("Fetching workflow runs...")
          queued_resp = make_request("actions/runs?status=queued")
          if queued_resp is None:
              print("Failed to fetch queued workflow runs after retries, exiting")
              exit(1)
          queued_workflow_runs = queued_resp.get("workflow_runs", [])
          in_progress_resp = make_request("actions/runs?status=in_progress")
          if in_progress_resp is None:
              print("Failed to fetch in-progress workflow runs after retries, exiting")
              exit(1)
          in_progress_workflow_runs = in_progress_resp.get("workflow_runs", [])

          def log_and_filter(runs, label):
              cicd_runs = [r for r in runs if r["name"] == WORKFLOW_NAME]
              print(f"{label}: {len(runs)} total, {len(cicd_runs)} {WORKFLOW_NAME}")
              for r in cicd_runs:
                  actor = (r.get("triggering_actor") or r.get("actor") or {}).get("login", "unknown")
                  matched = matches_queue(r, "${{ matrix.branch }}", CONTRIBUTOR_TYPE)
                  print(f"  run={r['id']} head_branch={r.get('head_branch')} event={r.get('event')} actor={actor} -> matched={matched}")
              return [r for r in cicd_runs if matches_queue(r, "${{ matrix.branch }}", CONTRIBUTOR_TYPE)]

          queued_workflow_runs = log_and_filter(queued_workflow_runs, "queued")
          in_progress_workflow_runs = log_and_filter(in_progress_workflow_runs, "in_progress")

          # Count running and queued workflows
          queued_workflows = len(queued_workflow_runs)
          in_progress_workflows = len(in_progress_workflow_runs)

          total_workflows = queued_workflows + in_progress_workflows
          print(f"Current queued workflows (PRs targeting ${{ matrix.branch }}, {CONTRIBUTOR_TYPE}): {queued_workflows}")
          print(f"Current running workflows (PRs targeting ${{ matrix.branch }}, {CONTRIBUTOR_TYPE}): {in_progress_workflows}")
          print(f"Total workflows: {total_workflows}")
          print(f"Max concurrency: {MAX_CONCURRENCY}")

          if total_workflows >= MAX_CONCURRENCY:
              print("Maximum concurrency reached, no new approvals will be made")
              exit(0)

          # Get waiting CI workflows for test environment
          print("Fetching waiting deployments...")
          waiting_resp = make_request("actions/runs?status=waiting")
          if waiting_resp is None:
              print("Failed to fetch waiting workflow runs after retries, exiting")
              exit(1)
          pending_workflows = log_and_filter(waiting_resp.get("workflow_runs", []), "waiting")

          # Sort deployments by creation date (oldest first)
          print("Sorting workflows...")
          pending_workflows = sorted(pending_workflows, key=lambda x: x["created_at"])

          # Process each deployment
          print(f"Processing {len(pending_workflows)} pending workflows...")
          for workflow in pending_workflows:
              if total_workflows >= MAX_CONCURRENCY:
                  print("Maximum concurrency reached, stopping approvals")
                  break

              workflow_id = workflow["id"]
              workflow_name = workflow["display_title"]
              print(f"Approving workflow {workflow_name} with Run Id: {workflow_id}")

              deployment_url = f"actions/runs/{workflow_id}/pending_deployments"
              deployments = make_request(deployment_url)
              if not deployments:
                  print(f"Failed to fetch pending deployments for run {workflow_id}")
                  exit(1)
              deployment = deployments[0]
              environment_id = deployment["environment"]["id"]

              # Approve the deployment
              status_data = {
                  "environment_ids": [environment_id],
                  "state": "approved",
                  "comment": "Automatically approved by queue manager"
              }
              result = make_request(deployment_url, method="POST", data=status_data)

              if result:
                  total_workflows += 1
              else:
                  print(f"Failed to approve deployment {deployment['id']}")
                  exit(1)
  notify:
    if: failure()
    runs-on: ubuntu-latest
    needs: [approve-queue]
    steps:
      - name: Notify
        env:
          SLACK_WEBHOOK: ${{ secrets.SLACK_CI_CHANNEL_WEBHOOK }}
          SLACK_WEBHOOK_ADMIN: <!subteam^${{ secrets.SLACK_TEAM_GROUP_ID }}>
          GITHUB_RUN_ID: ${{ github.run_id }}
          GITHUB_REPOSITORY: ${{ github.repository }}
        run: |
          curl -X POST \
            -H 'Content-type: application/json' \
            --data "{\"text\":\":robot_joy: <https://github.com/${GITHUB_REPOSITORY}/actions/runs/${GITHUB_RUN_ID}|Test-queue-approval-bot workflow> failed. Please review manually.\n\ncc ${SLACK_WEBHOOK_ADMIN}\"}" \
            $SLACK_WEBHOOK