merge: sync chat template kwargs with main

Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>

Author: yaoyu-33

.dev.commit

@@ -1 +1 @@
-cae55c1f0991abd3e391d0029b15bcb35723d3f6
+12eed0a5ac3f1ced6a3a4f76727db530bdbf7391

.github/actions/test-template/action.yml

@@ -59,6 +59,14 @@ inputs:
   runner:
     description: "Runner to use for test"
     required: true
+  shard-id:
+    description: "Zero-based pytest shard index"
+    required: false
+    default: ""
+  num-shards:
+    description: "Total number of pytest shards"
+    required: false
+    default: ""
 
 runs:
   using: "composite"
@@ -114,9 +122,14 @@ runs:
         echo "::group::Start test container"
         if [[ "$MOUNT_FS" == "true" ]]; then
           VOLUME_ARGS="--volume ${{ inputs.test-data-path }}:/home/TestData --env HF_HUB_OFFLINE=1"
+          HF_HOME=/home/TestData/HF_HOME
+          NEMO_HOME=/home/TestData/nemo_home
         else
           VOLUME_ARGS=""
+          HF_HOME=/home/ubuntu/.cache/huggingface
+          NEMO_HOME=/home/ubuntu/.cache/nemo
         fi
+        HF_MODULES_CACHE=/home/ubuntu/.cache/huggingface/modules
 
         cmd=$(cat <<RUN_TEST_EOF
         #!/bin/bash
@@ -131,10 +144,13 @@ runs:
           --cpus=40 \
           --env GHA_RUNNER=${{ inputs.runner }} \
           --env HYDRA_FULL_ERROR=1 \
-          --env HF_HOME=/home/TestData/HF_HOME \
-          --env NEMO_HOME=/home/TestData/nemo_home \
+          --env HF_HOME=$HF_HOME \
+          --env HF_MODULES_CACHE=$HF_MODULES_CACHE \
+          --env NEMO_HOME=$NEMO_HOME \
           --env RUN_ID=${{ github.run_id }} \
           --env UV_NO_SYNC=1 \
+          --env SHARD_ID=${{ inputs.shard-id }} \
+          --env NUM_SHARDS=${{ inputs.num-shards }} \
           --workdir /opt/Megatron-Bridge \
           $VOLUME_ARGS \
           ${{ inputs.container-image }} \

.github/scripts/test_docker_dependency_layers.sh

@@ -31,7 +31,7 @@ if grep -q -- '--mount=type=secret,id=GH_TOKEN' "$dockerfile"; then
 fi
 mcore_reinstall_line=$(grep -n 'uv pip install --no-deps --reinstall -e 3rdparty/Megatron-LM' "$dockerfile" | cut -d: -f1)
 helper_assertion_line=$(grep -n "find 3rdparty/Megatron-LM/megatron/core/datasets -maxdepth 1 -name 'helpers_cpp\*\.so'" "$dockerfile" | cut -d: -f1)
-final_copy_line=$(grep -n '^COPY --chmod=644 \. /opt/Megatron-Bridge$' "$dockerfile" | cut -d: -f1)
+final_copy_line=$(grep -nE '^COPY (--chmod=644|--chown=1000:1000) \. /opt/Megatron-Bridge$' "$dockerfile" | cut -d: -f1)
 [[ -n "$mcore_reinstall_line" ]]
 [[ -n "$helper_assertion_line" ]]
 [[ -n "$final_copy_line" ]]
@@ -232,17 +232,54 @@ if "$temporary_dir/revision-validator" "$revision_repo" not-a-full-sha; then
   exit 1
 fi
 
+if ! grep -A1 '^te = \[$' pyproject.toml | grep -Fxq '    "megatron-core[te]",'; then
+  echo "Bridge's TE extra must enable the selected MCore ref's TE extra" >&2
+  exit 1
+fi
+grep -Fq '{ name = "megatron-core", extra = ["te"] }' uv.lock
 if grep -q 'transformer-engine @ git+https://github.com/NVIDIA/TransformerEngine.git@' pyproject.toml || \
   grep -q '^name = "transformer-engine"$' pyproject.toml; then
   echo "Bridge must inherit the TransformerEngine source and metadata from the selected MCore ref" >&2
   exit 1
 fi
 
+lock_package_field() {
+  local lock_file="$1"
+  local package="$2"
+  local field="$3"
+
+  awk -v package="$package" -v field="$field" '
+    $0 == "[[package]]" { in_package = 1; package_matches = 0; next }
+    in_package && $0 ~ /^\[\[/ { in_package = 0; package_matches = 0 }
+    in_package && $0 == "name = \"" package "\"" { package_matches = 1; next }
+    package_matches && index($0, field " = ") == 1 { print; exit }
+  ' "$lock_file"
+}
+
+mlm_te_source=$(lock_package_field 3rdparty/Megatron-LM/uv.lock transformer-engine source)
+bridge_te_source=$(lock_package_field uv.lock transformer-engine source)
+if [[ ! "$mlm_te_source" =~ ^source\ =\ \{\ git\ =\ \"https://github.com/NVIDIA/TransformerEngine\.git\?rev=[0-9a-f]{40}#[0-9a-f]{40}\"\ \}$ ]]; then
+  echo "Selected MCore lock must pin TransformerEngine to a full source revision" >&2
+  exit 1
+fi
+if [[ "$bridge_te_source" != "$mlm_te_source" ]]; then
+  echo "Bridge must lock the TransformerEngine source selected by MCore" >&2
+  exit 1
+fi
+
 composite_action=".github/actions/test-template/action.yml"
 if grep -qE 'mcore_(commit|ref)|MCORE_COMMIT|uv sync --all-extras --all-groups' "$composite_action"; then
   echo "Test template must use the already-validated MCore in the built image" >&2
   exit 1
 fi
+grep -Fq 'VOLUME_ARGS="--volume ${{ inputs.test-data-path }}:/home/TestData --env HF_HUB_OFFLINE=1"' "$composite_action"
+grep -q 'MOUNT_FS: ${{ inputs.is_unit_test == '\''false'\'' }}' "$composite_action"
+grep -q 'HF_HOME=/home/TestData/HF_HOME' "$composite_action"
+grep -q 'NEMO_HOME=/home/TestData/nemo_home' "$composite_action"
+grep -q 'HF_HOME=/home/ubuntu/.cache/huggingface' "$composite_action"
+grep -q 'NEMO_HOME=/home/ubuntu/.cache/nemo' "$composite_action"
+grep -q 'HF_MODULES_CACHE=/home/ubuntu/.cache/huggingface/modules' "$composite_action"
+grep -q -- '--env HF_MODULES_CACHE=\$HF_MODULES_CACHE' "$composite_action"
 
 # The baseline dependency layer must be structurally independent of the mutable
 # dispatched checkout. CI validates the ordering statically so this regression

.github/scripts/test_queue_bypass.sh

@@ -0,0 +1,48 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+workflow="${1:-.github/workflows/cicd-main.yml}"
+
+awk '
+  function check_job() {
+    if (!in_queue_chain || !has_success) {
+      return
+    }
+
+    checked++
+    if (!allows_queue_bypass) {
+      print job " uses success() but does not allow the intentional queue skip" > "/dev/stderr"
+      failed = 1
+    }
+  }
+
+  /^  [[:alnum:]_-]+:$/ {
+    check_job()
+    job = $1
+    sub(/:$/, "", job)
+    in_queue_chain = (in_queue_chain || job == "cicd-compute-build-matrix")
+    if (job == "Nemo_CICD_Test") {
+      in_queue_chain = 0
+    }
+    has_success = 0
+    allows_queue_bypass = 0
+    next
+  }
+
+  in_queue_chain && /success\(\)/ {
+    has_success = 1
+  }
+
+  in_queue_chain && /needs\.configure\.outputs\.skip_queue == '\''true'\'' && !failure\(\)/ {
+    allows_queue_bypass = 1
+  }
+
+  END {
+    check_job()
+    if (checked == 0) {
+      print "No queue-dependent success() conditions found" > "/dev/stderr"
+      exit 1
+    }
+    exit failed
+  }
+' "$workflow"

.github/workflows/cicd-main.yml

@@ -90,6 +90,7 @@ jobs:
       expect_l1:         ${{ steps.configure.outputs.expect_l1 }}
       expect_l2:         ${{ steps.configure.outputs.expect_l2 }}
       perf_scripts_only: ${{ steps.configure.outputs.perf_scripts_only }}
+      skip_queue:        ${{ steps.configure.outputs.skip_queue }}
       merge_sha:         ${{ steps.merge-sha.outputs.merge_sha }}
     steps:
       - name: Get PR info
@@ -122,6 +123,7 @@ jobs:
           EVENT_NAME:    ${{ github.event_name }}
           REF:           ${{ github.ref }}
           TEST_SUITE:    ${{ github.event.inputs.test_suite }}
+          PR_HEAD_LABEL: ${{ fromJSON(steps.get-pr-info.outputs.pr-info || '{}').head.label || '' }}
         run: |
           PR_NUMBER=${{ fromJSON(steps.get-pr-info.outputs.pr-info || '{}').number }}
 
@@ -131,6 +133,9 @@ jobs:
           NEEDS_MORE_TESTS=$(echo "$LABELS" | jq 'any(. == "needs-more-tests")')
           FULL_TEST_SUITE=$(echo "$LABELS"  | jq 'any(. == "full-test-suite")')
 
+          SKIP_QUEUE=false
+          [[ "$PR_HEAD_LABEL" == svcnemo-autobot:chore/bump-mcore-* ]] && SKIP_QUEUE=true
+
           # Detect if every changed file lives under scripts/performance/
           PERF_SCRIPTS_ONLY=false
           if [[ -n "$PR_NUMBER" ]]; then
@@ -173,6 +178,7 @@ jobs:
           echo "docs_only=$DOCS_ONLY"               | tee -a "$GITHUB_OUTPUT"
           echo "needs_more_tests=$NEEDS_MORE_TESTS"   | tee -a "$GITHUB_OUTPUT"
           echo "full_test_suite=$FULL_TEST_SUITE"     | tee -a "$GITHUB_OUTPUT"
+          echo "skip_queue=$SKIP_QUEUE"               | tee -a "$GITHUB_OUTPUT"
           echo "expect_l0=$EXPECT_L0"                 | tee -a "$GITHUB_OUTPUT"
           echo "expect_l1=$EXPECT_L1"                 | tee -a "$GITHUB_OUTPUT"
           echo "expect_l2=$EXPECT_L2"                 | tee -a "$GITHUB_OUTPUT"
@@ -281,6 +287,7 @@ jobs:
         run: |
           .github/scripts/test_cache_keys.sh
           .github/scripts/test_docker_dependency_layers.sh
+          .github/scripts/test_queue_bypass.sh
 
       - name: Check lint
         run: |
@@ -296,7 +303,8 @@ jobs:
       !(needs.pre-flight.outputs.is_ci_workload == 'true'
       || needs.pre-flight.outputs.is_deployment_workflow == 'true'
       || needs.configure.outputs.docs_only == 'true'
-      || needs.configure.outputs.perf_scripts_only == 'true')
+      || needs.configure.outputs.perf_scripts_only == 'true'
+      || needs.configure.outputs.skip_queue == 'true')
       && github.event_name != 'merge_group'
     steps:
       - name: Running CI tests
@@ -310,6 +318,7 @@ jobs:
       (
         success()
         || needs.pre-flight.outputs.is_ci_workload == 'true'
+        || (needs.configure.outputs.skip_queue == 'true' && !failure())
         || needs.pre-flight.outputs.force_run_all == 'true'
         || github.event_name == 'merge_group'
       )
@@ -340,7 +349,7 @@ jobs:
           echo "matrix=$MATRIX" | tee -a "$GITHUB_OUTPUT"
 
   cicd-container-build:
-    needs: [pre-flight, configure, cicd-wait-in-queue, cicd-compute-build-matrix]
+    needs: [pre-flight, configure, cicd-compute-build-matrix]
     strategy:
       fail-fast: false
       matrix: ${{ fromJson(needs.cicd-compute-build-matrix.outputs.matrix) }}
@@ -350,6 +359,7 @@ jobs:
       (
         success()
         || needs.pre-flight.outputs.is_ci_workload == 'true'
+        || (needs.configure.outputs.skip_queue == 'true' && !failure())
         || needs.pre-flight.outputs.force_run_all == 'true'
         || github.event_name == 'merge_group'
       )
@@ -529,12 +539,13 @@ jobs:
       (
         success()
         || needs.pre-flight.outputs.is_ci_workload == 'true'
+        || (needs.configure.outputs.skip_queue == 'true' && !failure())
         || needs.pre-flight.outputs.force_run_all == 'true'
         || github.event_name == 'merge_group'
       )
       && !cancelled()
       && needs.configure.outputs.perf_scripts_only != 'true'
-    needs: [pre-flight, configure, cicd-wait-in-queue, cicd-container-build]
+    needs: [pre-flight, configure, cicd-container-build]
     runs-on: ${{ needs.pre-flight.outputs.runner_prefix }}
     name: Launch_Import_Check
     steps:
@@ -556,18 +567,23 @@ jobs:
               --skip-file /opt/import-check/import_check_skip.txt
 
   cicd-unit-tests-core:
+    strategy:
+      fail-fast: false
+      matrix:
+        shard-id: [0, 1, 2]
     if: |
       (
         success()
         || needs.pre-flight.outputs.is_ci_workload == 'true'
+        || (needs.configure.outputs.skip_queue == 'true' && !failure())
         || needs.pre-flight.outputs.force_run_all == 'true'
         || github.event_name == 'merge_group'
       )
       && !cancelled()
       && (github.event.inputs.test_suite == '' || github.event.inputs.test_suite == 'all' || github.event.inputs.test_suite == 'unit-only' || contains('L0 L1 L2', github.event.inputs.test_suite))
-    needs: [pre-flight, configure, cicd-wait-in-queue, cicd-container-build]
+    needs: [pre-flight, configure, cicd-container-build]
     runs-on: ${{ needs.pre-flight.outputs.runner_prefix }}
-    name: Launch_Unit_Tests_Core
+    name: Launch_Unit_Tests_Core_${{ matrix.shard-id }}
     env:
       TRANSFORMERS_OFFLINE: "1"
       HF_HUB_OFFLINE: "1"
@@ -589,8 +605,10 @@ jobs:
         uses: ./.github/actions/test-template
         with:
           script: Launch_Unit_Tests_Core
-          timeout: 18
+          timeout: 30
           is_unit_test: "true"
+          shard-id: ${{ matrix.shard-id }}
+          num-shards: 3
           github-token: ${{ steps.app-token.outputs.token }}
           container-image: ${{ needs.pre-flight.outputs.registry }}/megatron-bridge:${{ needs.configure.outputs.merge_sha }}
           test-data-path: ${{ needs.pre-flight.outputs.test_data_path }}
@@ -601,12 +619,13 @@ jobs:
       (
         success()
         || needs.pre-flight.outputs.is_ci_workload == 'true'
+        || (needs.configure.outputs.skip_queue == 'true' && !failure())
         || needs.pre-flight.outputs.force_run_all == 'true'
         || github.event_name == 'merge_group'
       )
       && !cancelled()
       && (github.event.inputs.test_suite == '' || github.event.inputs.test_suite == 'all' || github.event.inputs.test_suite == 'unit-only' || contains('L0 L1 L2', github.event.inputs.test_suite))
-    needs: [pre-flight, configure, cicd-wait-in-queue, cicd-container-build]
+    needs: [pre-flight, configure, cicd-container-build]
     runs-on: ${{ needs.pre-flight.outputs.runner_prefix }}
     name: Launch_Unit_Tests_Diffusion
     env:
@@ -649,6 +668,7 @@ jobs:
       (
         success()
         || needs.pre-flight.outputs.is_ci_workload == 'true'
+        || (needs.configure.outputs.skip_queue == 'true' && !failure())
         || needs.pre-flight.outputs.force_run_all == 'true'
         || github.event_name == 'merge_group'
       )
@@ -711,6 +731,7 @@ jobs:
       (
         success()
         || needs.pre-flight.outputs.is_ci_workload == 'true'
+        || (needs.configure.outputs.skip_queue == 'true' && !failure())
         || needs.pre-flight.outputs.force_run_all == 'true'
         || github.event_name == 'merge_group'
       )
@@ -758,6 +779,7 @@ jobs:
       (
         success()
         || needs.pre-flight.outputs.is_ci_workload == 'true'
+        || (needs.configure.outputs.skip_queue == 'true' && !failure())
         || needs.pre-flight.outputs.force_run_all == 'true'
         || github.event_name == 'merge_group'
       )
@@ -806,6 +828,7 @@ jobs:
       (
         success()
         || needs.pre-flight.outputs.is_ci_workload == 'true'
+        || (needs.configure.outputs.skip_queue == 'true' && !failure())
         || needs.pre-flight.outputs.force_run_all == 'true'
         || github.event_name == 'merge_group'
       )
@@ -893,6 +916,7 @@ jobs:
       (
         success()
         || needs.pre-flight.outputs.is_ci_workload == 'true'
+        || (needs.configure.outputs.skip_queue == 'true' && !failure())
         || needs.pre-flight.outputs.force_run_all == 'true'
         || github.event_name == 'merge_group'
       )
@@ -955,6 +979,7 @@ jobs:
       (
         success()
         || needs.pre-flight.outputs.is_ci_workload == 'true'
+        || (needs.configure.outputs.skip_queue == 'true' && !failure())
         || needs.pre-flight.outputs.force_run_all == 'true'
         || github.event_name == 'merge_group'
       )
@@ -1010,6 +1035,7 @@ jobs:
       (
         success()
         || needs.pre-flight.outputs.is_ci_workload == 'true'
+        || (needs.configure.outputs.skip_queue == 'true' && !failure())
         || needs.pre-flight.outputs.force_run_all == 'true'
         || github.event_name == 'merge_group'
       )
@@ -1066,6 +1092,7 @@ jobs:
       (
         success()
         || needs.pre-flight.outputs.is_ci_workload == 'true'
+        || (needs.configure.outputs.skip_queue == 'true' && !failure())
         || needs.pre-flight.outputs.force_run_all == 'true'
         || github.event_name == 'merge_group'
       )

.github/workflows/claude-review.yml

@@ -24,13 +24,15 @@ jobs:
   auto-review:
     name: Claude Review (auto)
     needs: [pre-flight]
-    # Auto-review non-draft PRs from NVIDIA authors on open / reopen / ready-for-review.
+    # Auto-review same-repository, non-draft PRs from NVIDIA authors on open / reopen / ready-for-review.
+    # Fork pull_request workflows cannot access the secrets or OIDC token required by the review action.
     # Subsequent commits do not re-trigger; re-request a review by commenting
     # "/claude review" (handled by the manual-review job below).
     # Skip bot-generated PRs: dependency bumps (`bump-ci-container-*`) and
     # auto cherry-picks (`cherry-pick-*`).
     if: |
       github.event_name == 'pull_request' &&
+      github.event.pull_request.head.repo.full_name == github.repository &&
       needs.pre-flight.outputs.is_member == 'true' &&
       github.event.pull_request.draft == false &&
       !startsWith(github.event.pull_request.head.ref, 'bump-ci-container-') &&

.main.commit

@@ -1 +1 @@
-24bad8e677d22625d86ef2a54c9506b6e4992c93
+14346b65a2d0790e451919858f7771078105c5f0