Merge remote-tracking branch 'upstream/main' into update-lazy-trt-compile

Signed-off-by: Boris Fomitchev <bfomitchev@nvidia.com>

Author: borisfom

.github/scripts/notify.py

@@ -0,0 +1,69 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.
+#
+# 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.
+import os
+
+import requests
+from github import Github
+
+
+def send_slack_notification():
+    # Get environment variables
+    gh_token = os.environ.get('GH_TOKEN')
+    webhook_url = os.environ.get('SLACK_WEBHOOK')
+    repository = os.environ.get('REPOSITORY')
+    run_id = os.environ.get('RUN_ID')
+    server_url = os.environ.get('SERVER_URL', 'https://github.com')
+    pr_number = int(os.environ.get('PR_NUMBER'))
+
+    # Get failure info from GitHub API
+    gh = Github(gh_token)
+    repo = gh.get_repo(repository)
+    pr = repo.get_pull(pr_number)
+
+    # Get failed jobs
+    failed_jobs = [job.name for job in repo.get_workflow_run(int(run_id)).jobs() if job.conclusion == 'failure']
+
+    # Build message blocks
+    blocks = [
+        {
+            "type": "section",
+            "text": {
+                "type": "mrkdwn",
+                "text": (
+                    f"*<{server_url}/{repository}/pull/{pr_number}|PR#{pr_number}: {pr.title.replace('`', '')}>*\n"
+                    f"• Author: <{server_url}/{pr.user.login}|{pr.user.login}>\n"
+                    f"• Branch: <{server_url}/{pr.head.repo.full_name}/tree/{pr.head.ref}|{pr.head.ref}>\n"
+                    f"• Pipeline: <{server_url}/{repository}/actions/runs/{run_id}|View Run>\n"
+                    f"• Failed Jobs:\n"
+                    + "\n".join(
+                        [
+                            f"    • <{server_url}/{repository}/actions/runs/{run_id}|{job.split('/')[-1]}>"
+                            for job in failed_jobs
+                            if job.split('/')[-1] != 'Nemo_CICD_Test'
+                        ]
+                    )
+                ),
+            },
+        }
+    ]
+
+    print({"blocks": blocks})
+
+    # Send to Slack
+    response = requests.post(webhook_url, json={"blocks": blocks})
+    response.raise_for_status()
+
+
+if __name__ == "__main__":
+    send_slack_notification()

.github/workflows/_bump_mcore_tag.yml

@@ -0,0 +1,54 @@
+name: Bump Megatron Tag Template
+on:
+  workflow_call:
+    inputs:
+      nemo-target-branch:
+        required: true
+        type: string
+        description: "The target branch to bump"
+      mcore-target-branch:
+        required: true
+        type: string
+        description: "The target branch to bump"
+    secrets:
+      PAT:
+        required: true
+
+jobs:
+  update-branch:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+        with:
+          ref: main
+
+      - name: Set Git config
+        run: |
+          git config --local user.email "actions@github.com"
+          git config --local user.name "Github Actions"
+      - name: Merge ${{ inputs.nemo-target-branch }} back to base branch
+        run: |
+          git fetch --unshallow
+          # Try to checkout and pull existing branch, create new one if it doesn't exist
+          if git ls-remote --exit-code origin ${{ inputs.nemo-target-branch }}; then
+            git checkout ${{ inputs.nemo-target-branch }}
+            git pull origin ${{ inputs.nemo-target-branch }}
+            git merge --no-ff ${{ inputs.nemo-target-branch }} -m "chore: Auto-merge main back to ${{ inputs.nemo-target-branch }}"
+          else
+            git checkout -b ${{ inputs.nemo-target-branch }} main
+          fi
+          git push -u origin ${{ inputs.nemo-target-branch }}
+
+  mcore:
+    uses: NVIDIA/NeMo-FW-CI-templates/.github/workflows/_bump_yamlfile.yml@v0.27.1
+    needs: [update-branch]
+    with:
+      source-repository: NVIDIA/Megatron-LM
+      source-ref: ${{ inputs.mcore-target-branch }}
+      yaml-path: '."vcs-dependencies"."megatron-lm".ref'
+      file: requirements/manifest.json
+      base-branch: weekly-bump-${{ inputs.nemo-target-branch }}
+      cicd-labels: Run CICD,no-fail-fast
+      pr-reviewers: ${{ inputs.pr-reviewers }}
+    secrets:
+      PAT: ${{ secrets.PAT }}

.github/workflows/_test_template.yml

@@ -31,6 +31,11 @@ on:
         description: Upload coverage as unit test
         required: false
         default: false
+      TESTS_TO_RUN:
+        type: string
+        description: Tests to run
+        required: false
+        default: '["all"]'
     outputs:
       conclusion:
         description: Conclusion of main test step
@@ -44,6 +49,8 @@ on:
 jobs:
   main:
     runs-on: ${{ inputs.RUNNER }}
+    name: ${{ inputs.SCRIPT }}
+    if: contains(fromJSON(inputs.TESTS_TO_RUN), inputs.SCRIPT) || contains(fromJSON(inputs.TESTS_TO_RUN), 'all')
     outputs:
       conclusion: ${{ steps.check.conclusion }}
       log: ${{ steps.check.outputs.log }}
@@ -62,26 +69,56 @@ jobs:
         run: |
           docker pull nemoci.azurecr.io/nemo_container:${{ github.run_id }}
 
+      - name: Clean repos
+        run: |
+
+      - name: Install jq
+        run: |
+          curl -sS https://webi.sh/jq | sh
+
+      - name: Create UUID
+        id: uuid
+        run: |
+          echo "id=$(uuidgen)" >> "$GITHUB_OUTPUT"
+
+      - name: Checkout NeMo
+        uses: actions/checkout@v2
+        with:
+          repository: NVIDIA/NeMo
+          path: ${{ github.run_id }}/${{steps.uuid.outputs.id }}/NeMo
+
       - name: Start container
         run: |
           mkdir -p $DIR
 
+          # Map of runner names to GPU device configurations
+          declare -A GPU_CONFIGS=(
+            ["myVm-01"]="0,1"
+            ["myVm-02"]="2,3" 
+            ["myVm-03"]="4,5"
+            ["myVm-04"]="6,7"
+          )
+
           ARG=("")
-          if [[ "${{ inputs.RUNNER }}" != *cpu* ]]; then
+          if [[ -n "${GPU_CONFIGS[${{ runner.name }}]}" ]]; then
+            ARG=("--runtime=nvidia --cpus="40" --memory="400g" --gpus '\"device=${GPU_CONFIGS[${{ runner.name }}]}\"'")
+          elif [[ "${{ inputs.RUNNER }}" != *cpu* ]]; then
             ARG=("--runtime=nvidia --gpus all")
           fi
 
           cmd=$(cat <<RUN_TEST_EOF
           #!/bin/bash
-          docker container rm -f nemo_container_${{ github.run_id }} || true
+          docker container rm -f nemo_container_${{ github.run_id }}_${{ runner.name }} || true
           docker run \
             --rm \
             -d \
-            --name nemo_container_${{ github.run_id }} ${ARG[@]} \
+            --name nemo_container_${{ github.run_id }}_${{ runner.name }} ${ARG[@]} \
             --shm-size=64g \
             --env TRANSFORMERS_OFFLINE=0 \
             --env HYDRA_FULL_ERROR=1 \
             --env HF_HOME=/home/TestData/HF_HOME \
+            --env RUN_ID=${{ github.run_id }} \
+            --volume $(pwd)/${{ github.run_id }}/${{steps.uuid.outputs.id }}/NeMo:/workspace \
             --volume /mnt/datadrive/TestData:/home/TestData nemoci.azurecr.io/nemo_container:${{ github.run_id }} \
             bash -c "sleep $(( ${{ inputs.TIMEOUT }} * 60 + 60 ))"
           RUN_TEST_EOF
@@ -107,7 +144,10 @@ jobs:
           (
             set -e
 
-            docker exec -t nemo_container_${{ github.run_id }} bash -c 'RUN_ID=${{ github.run_id }} bash tests/functional_tests/$SCRIPT.sh && echo "Finished successfully." || echo "Did not finish."'
+            docker exec -t nemo_container_${{ github.run_id }}_${{ runner.name }} bash -c '\
+              cp -r /opt/Megatron-LM/ /workspace/ && \
+              bash tests/functional_tests/$SCRIPT.sh && \
+              echo "Finished successfully." || echo "Did not finish."'
           ) 2>&1 | tee $DIR/err.log
 
           RUN_TEST_EOF
@@ -137,10 +177,10 @@ jobs:
           potential_infra_failure=$(cat $DIR/err.log | grep -Eqiw "device" && echo true || echo false)
           echo "potential_infra_failure=$potential_infra_failure" >> "$GITHUB_OUTPUT"
 
-          docker exec nemo_container_${{ github.run_id }} coverage combine
-          docker exec nemo_container_${{ github.run_id }} coverage xml
-          docker cp nemo_container_${{ github.run_id }}:/workspace/.coverage $DIR/.coverage 
-          docker cp nemo_container_${{ github.run_id }}:/workspace/coverage.xml $DIR/coverage.xml
+          docker exec nemo_container_${{ github.run_id }}_${{ runner.name }} coverage combine
+          docker exec nemo_container_${{ github.run_id }}_${{ runner.name }} coverage xml
+          docker cp nemo_container_${{ github.run_id }}_${{ runner.name }}:/workspace/.coverage $DIR/.coverage 
+          docker cp nemo_container_${{ github.run_id }}_${{ runner.name }}:/workspace/coverage.xml $DIR/coverage.xml
 
           coverage_report=coverage-${{ steps.create.outputs.coverage-prefix }}-${{ github.run_id }}-$(uuidgen)
           echo "coverage_report=$coverage_report" >> "$GITHUB_OUTPUT"
@@ -162,7 +202,7 @@ jobs:
       - name: Test coverage
         shell: bash -x -e -u -o pipefail {0}
         run: |
-          docker exec -t nemo_container_${{ github.run_id }} coverage report -i
+          docker exec -t nemo_container_${{ github.run_id }}_${{ runner.name }} coverage report -i
 
       - name: Upload artifacts
         uses: actions/upload-artifact@v4
@@ -174,14 +214,9 @@ jobs:
             ${{ github.run_id }}/.coverage
           include-hidden-files: true
 
-      - uses: "NVIDIA/NeMo/.github/actions/cancel-workflow@main"
-        if: failure() && inputs.IS_OPTIONAL == false && github.event_name == 'pull_request' && !contains(github.event.pull_request.labels.*.name, 'no-fail-fast')
-      - name: after_script
-        if: always() && inputs.AFTER_SCRIPT != ':'
-        run: |
-          docker exec nemo_container_${{ github.run_id }} bash -c '${{ inputs.AFTER_SCRIPT }}'
-
       - name: Container shutdown
         if: always()
         run: |
-          docker container rm -f nemo_container_${{ github.run_id }} || true
+          docker exec nemo_container_${{ github.run_id }}_${{ runner.name }} bash -c "chown -R $(id -u):$(id -g) /workspace"
+          rm -rf $(pwd)/${{ github.run_id }}/${{steps.uuid.outputs.id }} || true
+          docker container rm -f nemo_container_${{ github.run_id }}_${{ runner.name }} || true