Merge branch 'main' into khazic/feat/dspark-core

Author: kamran-nvidia

.dev.commit

@@ -1 +1 @@
-af91b6b5cc5aec01458b8e6d24ccb0f4717e49e3
+12eed0a5ac3f1ced6a3a4f76727db530bdbf7391

.github/CODEOWNERS

@@ -5,8 +5,11 @@ pyproject.toml @nvidia-nemo/automation
 ruff.toml @nvidia-nemo/automation
 .pre-commit-config.yaml
 
-scripts/performance/argument_parser.py @nvidia-nemo/automation @malay-nagda @erhoo82
-scripts/performance/setup_experiment.py @nvidia-nemo/automation @malay-nagda @erhoo82
-scripts/performance/bootstrap.py @nvidia-nemo/automation @malay-nagda @erhoo82
-scripts/performance/run_script.py @nvidia-nemo/automation @malay-nagda @erhoo82
-scripts/performance/run_recipe.py @nvidia-nemo/automation @malay-nagda @erhoo82
+scripts/performance @malay-nagda @dingqingy-nv
+src/megatron/bridge/perf_recipes @malay-nagda @dingqingy-nv
+
+scripts/performance/argument_parser.py @nvidia-nemo/automation
+scripts/performance/setup_experiment.py @nvidia-nemo/automation
+scripts/performance/bootstrap.py @nvidia-nemo/automation
+scripts/performance/run_script.py @nvidia-nemo/automation
+scripts/performance/run_recipe.py @nvidia-nemo/automation

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

@@ -43,8 +43,8 @@ inputs:
     description: "Has Azure credentials"
     required: false
     default: "false"
-  PAT:
-    description: "GitHub Personal Access Token"
+  github-token:
+    description: "GitHub token"
     required: true
   container-image:
     description: "Container image to use for test"
@@ -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 }} \
@@ -159,7 +175,7 @@ runs:
 
         docker exec -t nemo_container_${{ github.run_id }} bash -c '
           set -e
-          export GH_TOKEN=${{ inputs.PAT }}
+          export GH_TOKEN=${{ inputs.github-token }}
 
           timeout $(( ${{ inputs.timeout }} * 60 ))s bash tests/${{ inputs.is_unit_test == 'true' && 'unit_tests' || format('functional_tests/launch_scripts/{0}', inputs.script_dir) }}/${{ inputs.script }}.sh || EXIT_CODE=$?
 

.github/scripts/test_docker_dependency_layers.sh

@@ -3,6 +3,7 @@ set -euo pipefail
 
 dockerfile="${1:-docker/Dockerfile.ci}"
 workflow="${2:-.github/workflows/cicd-main.yml}"
+fw_final_dockerfile="${3:-docker/Dockerfile.fw_final}"
 
 baseline_arg_line=$(grep -n '^ARG BASELINE_MCORE_REF$' "$dockerfile" | cut -d: -f1)
 baseline_clone_line=$(grep -n 'git clone --filter=blob:none --no-checkout' "$dockerfile" | cut -d: -f1)
@@ -24,14 +25,23 @@ trap 'rm -rf "$temporary_dir"' EXIT
 ((baseline_line < dispatched_copy_line))
 ((dispatched_copy_line < delta_line))
 
+fw_final_from_line=$(grep -n '^FROM ${NEMO_FW_FINAL_BASE_IMAGE} AS nemo_fw_final$' "$fw_final_dockerfile" | cut -d: -f1)
+fw_final_root_line=$(grep -n '^USER root$' "$fw_final_dockerfile" | cut -d: -f1)
+fw_final_clone_line=$(grep -n '^RUN git clone --depth 1 https://github.com/NVIDIA-NeMo/NeMo.git' "$fw_final_dockerfile" | cut -d: -f1)
+[[ -n "$fw_final_from_line" ]]
+[[ -n "$fw_final_root_line" ]]
+[[ -n "$fw_final_clone_line" ]]
+((fw_final_from_line < fw_final_root_line))
+((fw_final_root_line < fw_final_clone_line))
+
 grep -q -- '--mount=type=cache,target=/root/.cache/uv' "$dockerfile"
 if grep -q -- '--mount=type=secret,id=GH_TOKEN' "$dockerfile"; then
   echo "Baseline MCore clone must not require a GitHub token" >&2
   exit 1
 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" ]]
@@ -52,7 +62,7 @@ if grep -q 'MCORE_REF="${{ github.event.inputs.mcore_ref }}"' "$workflow"; then
   echo "Dispatch inputs must enter shell steps through env" >&2
   exit 1
 fi
-test "$(grep -c 'validate_mcore_revision.sh "$MCORE_REPO" "$MCORE_REF"' "$workflow")" = 2
+test "$(grep -c 'validate_mcore_revision.sh "$MCORE_REPO" "$MCORE_REF" "$TRIGGERED_BY"' "$workflow")" = 2
 test "$(grep -c 'git fetch "$MCORE_REPO" "$MCORE_REF"' "$workflow")" = 2
 test "$(grep -c 'git checkout "$MCORE_REF"' "$workflow")" = 2
 test "$(grep -c 'test "$(git rev-parse HEAD)" = "$MCORE_REF"' "$workflow")" = 1
@@ -99,6 +109,13 @@ ancestor_sha=$(git -C "$revision_worktree" rev-parse HEAD)
 printf 'main\n' >>"$revision_worktree/file"
 git -C "$revision_worktree" commit -qam main
 main_sha=$(git -C "$revision_worktree" rev-parse HEAD)
+git -C "$revision_worktree" switch -q -c dev
+printf 'dev\n' >>"$revision_worktree/file"
+git -C "$revision_worktree" commit -qam dev
+dev_ancestor_sha=$(git -C "$revision_worktree" rev-parse HEAD)
+printf 'dev tip\n' >>"$revision_worktree/file"
+git -C "$revision_worktree" commit -qam dev-tip
+dev_sha=$(git -C "$revision_worktree" rev-parse HEAD)
 git -C "$revision_worktree" branch contributor "$ancestor_sha"
 git -C "$revision_worktree" switch -q contributor
 printf 'approved\n' >>"$revision_worktree/file"
@@ -122,6 +139,7 @@ git -C "$revision_worktree" add unapproved
 git -C "$revision_worktree" commit -q -m unapproved
 unapproved_sha=$(git -C "$revision_worktree" rev-parse HEAD)
 git -C "$revision_worktree" push -q "$revision_repo" "$main_sha:refs/heads/main"
+git -C "$revision_worktree" push -q "$revision_repo" "$dev_sha:refs/heads/dev"
 git -C "$revision_worktree" push -q "$revision_repo" "$mirror_sha:refs/heads/pull-request/123"
 git -C "$revision_worktree" push -q "$revision_repo" "$merge_sha:refs/pull/123/merge"
 git -C "$revision_worktree" push -q \
@@ -140,9 +158,69 @@ sed "s#\.github/scripts/validate_mcore_repo\.sh#$temporary_dir/revision-bin/repo
   "$revision_validator" >"$temporary_dir/revision-validator"
 chmod +x "$temporary_dir/revision-validator"
 "$temporary_dir/revision-validator" "$revision_repo" "$ancestor_sha"
+"$temporary_dir/revision-validator" "$revision_repo" "$dev_ancestor_sha"
+"$temporary_dir/revision-validator" "$revision_repo" "$dev_sha"
 "$temporary_dir/revision-validator" "$revision_repo" "$mirror_sha"
 "$temporary_dir/revision-validator" "$revision_repo" "$merge_sha"
 "$temporary_dir/revision-validator" "$revision_repo" "$queue_sha"
+git -C "$revision_worktree" push -q \
+  "$revision_repo" \
+  ":refs/heads/gh-readonly-queue/main/pr-123-$main_sha"
+if "$temporary_dir/revision-validator" "$revision_repo" "$queue_sha"; then
+  echo "MCore revision validation accepted a deleted queue ref without source-run evidence" >&2
+  exit 1
+fi
+cat >"$temporary_dir/revision-bin/gh" <<'EOF'
+#!/usr/bin/env bash
+[[ "$1" == "api" && "$2" == "repos/NVIDIA/Megatron-LM/actions/runs/$TEST_RUN_ID" ]] || exit 1
+printf '%s\t%s\t%s\t%s\n' \
+  "${RUN_REPO:-NVIDIA/Megatron-LM}" \
+  "${RUN_EVENT:-merge_group}" \
+  "${RUN_SHA:-}" \
+  "${RUN_BRANCH:-gh-readonly-queue/main/pr-123-0000000000000000000000000000000000000000}"
+EOF
+chmod +x "$temporary_dir/revision-bin/gh"
+TEST_RUN_ID=123456
+trigger_url="https://github.com/NVIDIA/Megatron-LM/actions/runs/$TEST_RUN_ID"
+PATH="$temporary_dir/revision-bin:$PATH" \
+  TEST_RUN_ID="$TEST_RUN_ID" \
+  RUN_SHA="$queue_sha" \
+  RUN_BRANCH="gh-readonly-queue/main/pr-123-$main_sha" \
+  "$temporary_dir/revision-validator" "$revision_repo" "$queue_sha" "$trigger_url"
+if PATH="$temporary_dir/revision-bin:$PATH" \
+  TEST_RUN_ID="$TEST_RUN_ID" \
+  RUN_SHA="$unapproved_sha" \
+  RUN_BRANCH="gh-readonly-queue/main/pr-123-$main_sha" \
+  "$temporary_dir/revision-validator" "$revision_repo" "$queue_sha" "$trigger_url"; then
+  echo "MCore revision validation accepted source-run metadata for a different SHA" >&2
+  exit 1
+fi
+if PATH="$temporary_dir/revision-bin:$PATH" \
+  TEST_RUN_ID="$TEST_RUN_ID" \
+  RUN_SHA="$queue_sha" \
+  RUN_BRANCH="pull-request/123" \
+  "$temporary_dir/revision-validator" "$revision_repo" "$queue_sha" "$trigger_url"; then
+  echo "MCore revision validation accepted a non-queue source run" >&2
+  exit 1
+fi
+if PATH="$temporary_dir/revision-bin:$PATH" \
+  TEST_RUN_ID="$TEST_RUN_ID" \
+  RUN_REPO="example/Megatron-LM" \
+  RUN_SHA="$queue_sha" \
+  RUN_BRANCH="gh-readonly-queue/main/pr-123-$main_sha" \
+  "$temporary_dir/revision-validator" "$revision_repo" "$queue_sha" "$trigger_url"; then
+  echo "MCore revision validation accepted an untrusted source-run repository" >&2
+  exit 1
+fi
+if PATH="$temporary_dir/revision-bin:$PATH" \
+  TEST_RUN_ID="$TEST_RUN_ID" \
+  RUN_EVENT="push" \
+  RUN_SHA="$queue_sha" \
+  RUN_BRANCH="gh-readonly-queue/main/pr-123-$main_sha" \
+  "$temporary_dir/revision-validator" "$revision_repo" "$queue_sha" "$trigger_url"; then
+  echo "MCore revision validation accepted a non-merge-group source run" >&2
+  exit 1
+fi
 if "$temporary_dir/revision-validator" "$revision_repo" "$queue_sha"x; then
   echo "MCore revision validation accepted a malformed queue SHA" >&2
   exit 1
@@ -164,11 +242,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/scripts/validate_mcore_revision.sh

@@ -3,6 +3,7 @@ set -euo pipefail
 
 repo="${1:-}"
 revision="${2:-}"
+triggered_by="${3:-}"
 
 if ! .github/scripts/validate_mcore_repo.sh "$repo"; then
   exit 1
@@ -14,27 +15,32 @@ fi
 refs=$(git ls-remote \
   "$repo" \
   "refs/heads/main" \
+  "refs/heads/dev" \
   "refs/heads/pull-request/*" \
-  "refs/heads/gh-readonly-queue/main/pr-*" \
+  "refs/heads/gh-readonly-queue/*/pr-*" \
   "refs/pull/*/merge")
 object_store=$(mktemp -d)
 trap 'rm -rf "$object_store"' EXIT
 git -C "$object_store" init --quiet
 
 main_sha=$(awk '$2 == "refs/heads/main" {print $1}' <<<"$refs")
-if [[ -n "$main_sha" ]] && \
-  git -C "$object_store" fetch --quiet --filter=blob:none --no-tags "$repo" "$main_sha" "$revision"; then
-  if git -C "$object_store" merge-base --is-ancestor "$revision" "$main_sha"; then
-    exit 0
+dev_sha=$(awk '$2 == "refs/heads/dev" {print $1}' <<<"$refs")
+for protected_sha in "$main_sha" "$dev_sha"; do
+  if [[ -n "$protected_sha" ]] && \
+    git -C "$object_store" fetch --quiet --filter=blob:none --no-tags "$repo" "$protected_sha" "$revision"; then
+    if git -C "$object_store" merge-base --is-ancestor "$revision" "$protected_sha"; then
+      exit 0
+    fi
   fi
-fi
+done
 
 while IFS=$'\t' read -r sha ref; do
   if [[ "$sha" != "$revision" ]]; then
     continue
   fi
-  if [[ "$ref" =~ ^refs/heads/pull-request/[0-9]+$ ]] || \
-    [[ "$ref" =~ ^refs/heads/gh-readonly-queue/main/pr-[0-9]+-[0-9a-f]{40}$ ]]; then
+  if [[ "$ref" == "refs/heads/dev" ]] || \
+    [[ "$ref" =~ ^refs/heads/pull-request/[0-9]+$ ]] || \
+    [[ "$ref" =~ ^refs/heads/gh-readonly-queue/(main|dev)/pr-[0-9]+-[0-9a-f]{40}$ ]]; then
     exit 0
   fi
 done <<<"$refs"
@@ -53,4 +59,22 @@ while IFS=$'\t' read -r merge_sha merge_ref; do
   fi
 done <<<"$refs"
 
+# GitHub deletes merge-queue refs when a group leaves the queue, which can
+# happen before MCore's remotely triggered MBridge run starts. In that case,
+# authenticate the immutable source run and require it to describe this exact
+# official merge-group SHA. Arbitrary workflow-dispatch inputs and fork SHAs
+# remain rejected.
+if [[ "$triggered_by" =~ ^https://github.com/NVIDIA/Megatron-LM/actions/runs/([1-9][0-9]*)$ ]]; then
+  run_id="${BASH_REMATCH[1]}"
+  run_metadata=$(gh api "repos/NVIDIA/Megatron-LM/actions/runs/$run_id" --jq \
+    '[.repository.full_name, .event, .head_sha, .head_branch] | @tsv' 2>/dev/null) || run_metadata=""
+  IFS=$'\t' read -r run_repo run_event run_sha run_branch <<<"$run_metadata"
+  if [[ "$run_repo" == "NVIDIA/Megatron-LM" && \
+    "$run_event" == "merge_group" && \
+    "$run_sha" == "$revision" && \
+    "$run_branch" =~ ^gh-readonly-queue/(main|dev)/pr-[0-9]+-[0-9a-f]{40}$ ]]; then
+    exit 0
+  fi
+fi
+
 exit 1