Skip to content

Commit ceb5305

Browse files
yaoyu-33yaoyu-33
committed
chore: merge current main into Muse Glimmer support
Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
2 parents eb4dd2f + 2b8cb21 commit ceb5305

70 files changed

Lines changed: 2907 additions & 209 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎.dev.commit‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1 +1 @@
1-
cae55c1f0991abd3e391d0029b15bcb35723d3f6
1+
12eed0a5ac3f1ced6a3a4f76727db530bdbf7391

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

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,14 @@ inputs:
5959
runner:
6060
description: "Runner to use for test"
6161
required: true
62+
shard-id:
63+
description: "Zero-based pytest shard index"
64+
required: false
65+
default: ""
66+
num-shards:
67+
description: "Total number of pytest shards"
68+
required: false
69+
default: ""
6270

6371
runs:
6472
using: "composite"
@@ -114,9 +122,14 @@ runs:
114122
echo "::group::Start test container"
115123
if [[ "$MOUNT_FS" == "true" ]]; then
116124
VOLUME_ARGS="--volume ${{ inputs.test-data-path }}:/home/TestData --env HF_HUB_OFFLINE=1"
125+
HF_HOME=/home/TestData/HF_HOME
126+
NEMO_HOME=/home/TestData/nemo_home
117127
else
118128
VOLUME_ARGS=""
129+
HF_HOME=/home/ubuntu/.cache/huggingface
130+
NEMO_HOME=/home/ubuntu/.cache/nemo
119131
fi
132+
HF_MODULES_CACHE=/home/ubuntu/.cache/huggingface/modules
120133
121134
cmd=$(cat <<RUN_TEST_EOF
122135
#!/bin/bash
@@ -131,10 +144,13 @@ runs:
131144
--cpus=40 \
132145
--env GHA_RUNNER=${{ inputs.runner }} \
133146
--env HYDRA_FULL_ERROR=1 \
134-
--env HF_HOME=/home/TestData/HF_HOME \
135-
--env NEMO_HOME=/home/TestData/nemo_home \
147+
--env HF_HOME=$HF_HOME \
148+
--env HF_MODULES_CACHE=$HF_MODULES_CACHE \
149+
--env NEMO_HOME=$NEMO_HOME \
136150
--env RUN_ID=${{ github.run_id }} \
137151
--env UV_NO_SYNC=1 \
152+
--env SHARD_ID=${{ inputs.shard-id }} \
153+
--env NUM_SHARDS=${{ inputs.num-shards }} \
138154
--workdir /opt/Megatron-Bridge \
139155
$VOLUME_ARGS \
140156
${{ inputs.container-image }} \

‎.github/scripts/test_docker_dependency_layers.sh‎

Lines changed: 48 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ set -euo pipefail
33

44
dockerfile="${1:-docker/Dockerfile.ci}"
55
workflow="${2:-.github/workflows/cicd-main.yml}"
6+
fw_final_dockerfile="${3:-docker/Dockerfile.fw_final}"
67

78
baseline_arg_line=$(grep -n '^ARG BASELINE_MCORE_REF$' "$dockerfile" | cut -d: -f1)
89
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
2425
((baseline_line < dispatched_copy_line))
2526
((dispatched_copy_line < delta_line))
2627

28+
fw_final_from_line=$(grep -n '^FROM ${NEMO_FW_FINAL_BASE_IMAGE} AS nemo_fw_final$' "$fw_final_dockerfile" | cut -d: -f1)
29+
fw_final_root_line=$(grep -n '^USER root$' "$fw_final_dockerfile" | cut -d: -f1)
30+
fw_final_clone_line=$(grep -n '^RUN git clone --depth 1 https://github.com/NVIDIA-NeMo/NeMo.git' "$fw_final_dockerfile" | cut -d: -f1)
31+
[[ -n "$fw_final_from_line" ]]
32+
[[ -n "$fw_final_root_line" ]]
33+
[[ -n "$fw_final_clone_line" ]]
34+
((fw_final_from_line < fw_final_root_line))
35+
((fw_final_root_line < fw_final_clone_line))
36+
2737
grep -q -- '--mount=type=cache,target=/root/.cache/uv' "$dockerfile"
2838
if grep -q -- '--mount=type=secret,id=GH_TOKEN' "$dockerfile"; then
2939
echo "Baseline MCore clone must not require a GitHub token" >&2
3040
exit 1
3141
fi
3242
mcore_reinstall_line=$(grep -n 'uv pip install --no-deps --reinstall -e 3rdparty/Megatron-LM' "$dockerfile" | cut -d: -f1)
3343
helper_assertion_line=$(grep -n "find 3rdparty/Megatron-LM/megatron/core/datasets -maxdepth 1 -name 'helpers_cpp\*\.so'" "$dockerfile" | cut -d: -f1)
34-
final_copy_line=$(grep -n '^COPY --chmod=644 \. /opt/Megatron-Bridge$' "$dockerfile" | cut -d: -f1)
44+
final_copy_line=$(grep -nE '^COPY (--chmod=644|--chown=1000:1000) \. /opt/Megatron-Bridge$' "$dockerfile" | cut -d: -f1)
3545
[[ -n "$mcore_reinstall_line" ]]
3646
[[ -n "$helper_assertion_line" ]]
3747
[[ -n "$final_copy_line" ]]
@@ -232,17 +242,54 @@ if "$temporary_dir/revision-validator" "$revision_repo" not-a-full-sha; then
232242
exit 1
233243
fi
234244

245+
if ! grep -A1 '^te = \[$' pyproject.toml | grep -Fxq ' "megatron-core[te]",'; then
246+
echo "Bridge's TE extra must enable the selected MCore ref's TE extra" >&2
247+
exit 1
248+
fi
249+
grep -Fq '{ name = "megatron-core", extra = ["te"] }' uv.lock
235250
if grep -q 'transformer-engine @ git+https://github.com/NVIDIA/TransformerEngine.git@' pyproject.toml || \
236251
grep -q '^name = "transformer-engine"$' pyproject.toml; then
237252
echo "Bridge must inherit the TransformerEngine source and metadata from the selected MCore ref" >&2
238253
exit 1
239254
fi
240255

256+
lock_package_field() {
257+
local lock_file="$1"
258+
local package="$2"
259+
local field="$3"
260+
261+
awk -v package="$package" -v field="$field" '
262+
$0 == "[[package]]" { in_package = 1; package_matches = 0; next }
263+
in_package && $0 ~ /^\[\[/ { in_package = 0; package_matches = 0 }
264+
in_package && $0 == "name = \"" package "\"" { package_matches = 1; next }
265+
package_matches && index($0, field " = ") == 1 { print; exit }
266+
' "$lock_file"
267+
}
268+
269+
mlm_te_source=$(lock_package_field 3rdparty/Megatron-LM/uv.lock transformer-engine source)
270+
bridge_te_source=$(lock_package_field uv.lock transformer-engine source)
271+
if [[ ! "$mlm_te_source" =~ ^source\ =\ \{\ git\ =\ \"https://github.com/NVIDIA/TransformerEngine\.git\?rev=[0-9a-f]{40}#[0-9a-f]{40}\"\ \}$ ]]; then
272+
echo "Selected MCore lock must pin TransformerEngine to a full source revision" >&2
273+
exit 1
274+
fi
275+
if [[ "$bridge_te_source" != "$mlm_te_source" ]]; then
276+
echo "Bridge must lock the TransformerEngine source selected by MCore" >&2
277+
exit 1
278+
fi
279+
241280
composite_action=".github/actions/test-template/action.yml"
242281
if grep -qE 'mcore_(commit|ref)|MCORE_COMMIT|uv sync --all-extras --all-groups' "$composite_action"; then
243282
echo "Test template must use the already-validated MCore in the built image" >&2
244283
exit 1
245284
fi
285+
grep -Fq 'VOLUME_ARGS="--volume ${{ inputs.test-data-path }}:/home/TestData --env HF_HUB_OFFLINE=1"' "$composite_action"
286+
grep -q 'MOUNT_FS: ${{ inputs.is_unit_test == '\''false'\'' }}' "$composite_action"
287+
grep -q 'HF_HOME=/home/TestData/HF_HOME' "$composite_action"
288+
grep -q 'NEMO_HOME=/home/TestData/nemo_home' "$composite_action"
289+
grep -q 'HF_HOME=/home/ubuntu/.cache/huggingface' "$composite_action"
290+
grep -q 'NEMO_HOME=/home/ubuntu/.cache/nemo' "$composite_action"
291+
grep -q 'HF_MODULES_CACHE=/home/ubuntu/.cache/huggingface/modules' "$composite_action"
292+
grep -q -- '--env HF_MODULES_CACHE=\$HF_MODULES_CACHE' "$composite_action"
246293

247294
# The baseline dependency layer must be structurally independent of the mutable
248295
# dispatched checkout. CI validates the ordering statically so this regression

‎.github/scripts/test_queue_bypass.sh‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
#!/usr/bin/env bash
2+
set -euo pipefail
3+
4+
workflow="${1:-.github/workflows/cicd-main.yml}"
5+
6+
awk '
7+
function check_job() {
8+
if (!in_queue_chain || !has_success) {
9+
return
10+
}
11+
12+
checked++
13+
if (!allows_queue_bypass) {
14+
print job " uses success() but does not allow the intentional queue skip" > "/dev/stderr"
15+
failed = 1
16+
}
17+
}
18+
19+
/^ [[:alnum:]_-]+:$/ {
20+
check_job()
21+
job = $1
22+
sub(/:$/, "", job)
23+
in_queue_chain = (in_queue_chain || job == "cicd-compute-build-matrix")
24+
if (job == "Nemo_CICD_Test") {
25+
in_queue_chain = 0
26+
}
27+
has_success = 0
28+
allows_queue_bypass = 0
29+
next
30+
}
31+
32+
in_queue_chain && /success\(\)/ {
33+
has_success = 1
34+
}
35+
36+
in_queue_chain && /needs\.configure\.outputs\.skip_queue == '\''true'\'' && !failure\(\)/ {
37+
allows_queue_bypass = 1
38+
}
39+
40+
END {
41+
check_job()
42+
if (checked == 0) {
43+
print "No queue-dependent success() conditions found" > "/dev/stderr"
44+
exit 1
45+
}
46+
exit failed
47+
}
48+
' "$workflow"

‎.github/workflows/cicd-main.yml‎

Lines changed: 34 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,7 @@ jobs:
9090
expect_l1: ${{ steps.configure.outputs.expect_l1 }}
9191
expect_l2: ${{ steps.configure.outputs.expect_l2 }}
9292
perf_scripts_only: ${{ steps.configure.outputs.perf_scripts_only }}
93+
skip_queue: ${{ steps.configure.outputs.skip_queue }}
9394
merge_sha: ${{ steps.merge-sha.outputs.merge_sha }}
9495
steps:
9596
- name: Get PR info
@@ -122,6 +123,7 @@ jobs:
122123
EVENT_NAME: ${{ github.event_name }}
123124
REF: ${{ github.ref }}
124125
TEST_SUITE: ${{ github.event.inputs.test_suite }}
126+
PR_HEAD_LABEL: ${{ fromJSON(steps.get-pr-info.outputs.pr-info || '{}').head.label || '' }}
125127
run: |
126128
PR_NUMBER=${{ fromJSON(steps.get-pr-info.outputs.pr-info || '{}').number }}
127129
@@ -131,6 +133,9 @@ jobs:
131133
NEEDS_MORE_TESTS=$(echo "$LABELS" | jq 'any(. == "needs-more-tests")')
132134
FULL_TEST_SUITE=$(echo "$LABELS" | jq 'any(. == "full-test-suite")')
133135
136+
SKIP_QUEUE=false
137+
[[ "$PR_HEAD_LABEL" == svcnemo-autobot:chore/bump-mcore-* ]] && SKIP_QUEUE=true
138+
134139
# Detect if every changed file lives under scripts/performance/
135140
PERF_SCRIPTS_ONLY=false
136141
if [[ -n "$PR_NUMBER" ]]; then
@@ -173,6 +178,7 @@ jobs:
173178
echo "docs_only=$DOCS_ONLY" | tee -a "$GITHUB_OUTPUT"
174179
echo "needs_more_tests=$NEEDS_MORE_TESTS" | tee -a "$GITHUB_OUTPUT"
175180
echo "full_test_suite=$FULL_TEST_SUITE" | tee -a "$GITHUB_OUTPUT"
181+
echo "skip_queue=$SKIP_QUEUE" | tee -a "$GITHUB_OUTPUT"
176182
echo "expect_l0=$EXPECT_L0" | tee -a "$GITHUB_OUTPUT"
177183
echo "expect_l1=$EXPECT_L1" | tee -a "$GITHUB_OUTPUT"
178184
echo "expect_l2=$EXPECT_L2" | tee -a "$GITHUB_OUTPUT"
@@ -281,6 +287,7 @@ jobs:
281287
run: |
282288
.github/scripts/test_cache_keys.sh
283289
.github/scripts/test_docker_dependency_layers.sh
290+
.github/scripts/test_queue_bypass.sh
284291
285292
- name: Check lint
286293
run: |
@@ -296,7 +303,8 @@ jobs:
296303
!(needs.pre-flight.outputs.is_ci_workload == 'true'
297304
|| needs.pre-flight.outputs.is_deployment_workflow == 'true'
298305
|| needs.configure.outputs.docs_only == 'true'
299-
|| needs.configure.outputs.perf_scripts_only == 'true')
306+
|| needs.configure.outputs.perf_scripts_only == 'true'
307+
|| needs.configure.outputs.skip_queue == 'true')
300308
&& github.event_name != 'merge_group'
301309
steps:
302310
- name: Running CI tests
@@ -310,6 +318,7 @@ jobs:
310318
(
311319
success()
312320
|| needs.pre-flight.outputs.is_ci_workload == 'true'
321+
|| (needs.configure.outputs.skip_queue == 'true' && !failure())
313322
|| needs.pre-flight.outputs.force_run_all == 'true'
314323
|| github.event_name == 'merge_group'
315324
)
@@ -340,7 +349,7 @@ jobs:
340349
echo "matrix=$MATRIX" | tee -a "$GITHUB_OUTPUT"
341350
342351
cicd-container-build:
343-
needs: [pre-flight, configure, cicd-wait-in-queue, cicd-compute-build-matrix]
352+
needs: [pre-flight, configure, cicd-compute-build-matrix]
344353
strategy:
345354
fail-fast: false
346355
matrix: ${{ fromJson(needs.cicd-compute-build-matrix.outputs.matrix) }}
@@ -350,6 +359,7 @@ jobs:
350359
(
351360
success()
352361
|| needs.pre-flight.outputs.is_ci_workload == 'true'
362+
|| (needs.configure.outputs.skip_queue == 'true' && !failure())
353363
|| needs.pre-flight.outputs.force_run_all == 'true'
354364
|| github.event_name == 'merge_group'
355365
)
@@ -529,12 +539,13 @@ jobs:
529539
(
530540
success()
531541
|| needs.pre-flight.outputs.is_ci_workload == 'true'
542+
|| (needs.configure.outputs.skip_queue == 'true' && !failure())
532543
|| needs.pre-flight.outputs.force_run_all == 'true'
533544
|| github.event_name == 'merge_group'
534545
)
535546
&& !cancelled()
536547
&& needs.configure.outputs.perf_scripts_only != 'true'
537-
needs: [pre-flight, configure, cicd-wait-in-queue, cicd-container-build]
548+
needs: [pre-flight, configure, cicd-container-build]
538549
runs-on: ${{ needs.pre-flight.outputs.runner_prefix }}
539550
name: Launch_Import_Check
540551
steps:
@@ -556,18 +567,23 @@ jobs:
556567
--skip-file /opt/import-check/import_check_skip.txt
557568
558569
cicd-unit-tests-core:
570+
strategy:
571+
fail-fast: false
572+
matrix:
573+
shard-id: [0, 1, 2]
559574
if: |
560575
(
561576
success()
562577
|| needs.pre-flight.outputs.is_ci_workload == 'true'
578+
|| (needs.configure.outputs.skip_queue == 'true' && !failure())
563579
|| needs.pre-flight.outputs.force_run_all == 'true'
564580
|| github.event_name == 'merge_group'
565581
)
566582
&& !cancelled()
567583
&& (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))
568-
needs: [pre-flight, configure, cicd-wait-in-queue, cicd-container-build]
584+
needs: [pre-flight, configure, cicd-container-build]
569585
runs-on: ${{ needs.pre-flight.outputs.runner_prefix }}
570-
name: Launch_Unit_Tests_Core
586+
name: Launch_Unit_Tests_Core_${{ matrix.shard-id }}
571587
env:
572588
TRANSFORMERS_OFFLINE: "1"
573589
HF_HUB_OFFLINE: "1"
@@ -589,8 +605,10 @@ jobs:
589605
uses: ./.github/actions/test-template
590606
with:
591607
script: Launch_Unit_Tests_Core
592-
timeout: 18
608+
timeout: 30
593609
is_unit_test: "true"
610+
shard-id: ${{ matrix.shard-id }}
611+
num-shards: 3
594612
github-token: ${{ steps.app-token.outputs.token }}
595613
container-image: ${{ needs.pre-flight.outputs.registry }}/megatron-bridge:${{ needs.configure.outputs.merge_sha }}
596614
test-data-path: ${{ needs.pre-flight.outputs.test_data_path }}
@@ -601,12 +619,13 @@ jobs:
601619
(
602620
success()
603621
|| needs.pre-flight.outputs.is_ci_workload == 'true'
622+
|| (needs.configure.outputs.skip_queue == 'true' && !failure())
604623
|| needs.pre-flight.outputs.force_run_all == 'true'
605624
|| github.event_name == 'merge_group'
606625
)
607626
&& !cancelled()
608627
&& (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))
609-
needs: [pre-flight, configure, cicd-wait-in-queue, cicd-container-build]
628+
needs: [pre-flight, configure, cicd-container-build]
610629
runs-on: ${{ needs.pre-flight.outputs.runner_prefix }}
611630
name: Launch_Unit_Tests_Diffusion
612631
env:
@@ -649,6 +668,7 @@ jobs:
649668
(
650669
success()
651670
|| needs.pre-flight.outputs.is_ci_workload == 'true'
671+
|| (needs.configure.outputs.skip_queue == 'true' && !failure())
652672
|| needs.pre-flight.outputs.force_run_all == 'true'
653673
|| github.event_name == 'merge_group'
654674
)
@@ -711,6 +731,7 @@ jobs:
711731
(
712732
success()
713733
|| needs.pre-flight.outputs.is_ci_workload == 'true'
734+
|| (needs.configure.outputs.skip_queue == 'true' && !failure())
714735
|| needs.pre-flight.outputs.force_run_all == 'true'
715736
|| github.event_name == 'merge_group'
716737
)
@@ -758,6 +779,7 @@ jobs:
758779
(
759780
success()
760781
|| needs.pre-flight.outputs.is_ci_workload == 'true'
782+
|| (needs.configure.outputs.skip_queue == 'true' && !failure())
761783
|| needs.pre-flight.outputs.force_run_all == 'true'
762784
|| github.event_name == 'merge_group'
763785
)
@@ -806,6 +828,7 @@ jobs:
806828
(
807829
success()
808830
|| needs.pre-flight.outputs.is_ci_workload == 'true'
831+
|| (needs.configure.outputs.skip_queue == 'true' && !failure())
809832
|| needs.pre-flight.outputs.force_run_all == 'true'
810833
|| github.event_name == 'merge_group'
811834
)
@@ -893,6 +916,7 @@ jobs:
893916
(
894917
success()
895918
|| needs.pre-flight.outputs.is_ci_workload == 'true'
919+
|| (needs.configure.outputs.skip_queue == 'true' && !failure())
896920
|| needs.pre-flight.outputs.force_run_all == 'true'
897921
|| github.event_name == 'merge_group'
898922
)
@@ -955,6 +979,7 @@ jobs:
955979
(
956980
success()
957981
|| needs.pre-flight.outputs.is_ci_workload == 'true'
982+
|| (needs.configure.outputs.skip_queue == 'true' && !failure())
958983
|| needs.pre-flight.outputs.force_run_all == 'true'
959984
|| github.event_name == 'merge_group'
960985
)
@@ -1010,6 +1035,7 @@ jobs:
10101035
(
10111036
success()
10121037
|| needs.pre-flight.outputs.is_ci_workload == 'true'
1038+
|| (needs.configure.outputs.skip_queue == 'true' && !failure())
10131039
|| needs.pre-flight.outputs.force_run_all == 'true'
10141040
|| github.event_name == 'merge_group'
10151041
)
@@ -1066,6 +1092,7 @@ jobs:
10661092
(
10671093
success()
10681094
|| needs.pre-flight.outputs.is_ci_workload == 'true'
1095+
|| (needs.configure.outputs.skip_queue == 'true' && !failure())
10691096
|| needs.pre-flight.outputs.force_run_all == 'true'
10701097
|| github.event_name == 'merge_group'
10711098
)

0 commit comments

Comments
 (0)