Skip to content

[data] fix: Support Energon 7 metadata fields#4089

Closed
yaoyu-33 wants to merge 3 commits into
mainfrom
yuya/mcore-main-autofix-20260530-pr4087
Closed

[data] fix: Support Energon 7 metadata fields#4089
yaoyu-33 wants to merge 3 commits into
mainfrom
yuya/mcore-main-autofix-20260530-pr4087

Conversation

@yaoyu-33

@yaoyu-33 yaoyu-33 commented May 30, 2026

Copy link
Copy Markdown
Contributor

Summary

  • Adds Energon metadata helpers that construct sample and batch kwargs from the installed Energon base dataclasses.
  • Updates Bridge Energon task encoders and focused tests to work with the Energon 7 metadata contract while retaining Energon 6 compatibility.

Original Bump

  • Bump PR: chore(beep boop 🤖): Bump uv.lock (main, mcore-main) (2026-05-30) #4087
  • Target: main
  • Classification: Fix must ship with the bump
  • Root cause: MCore commit 52d1d681d6000e5da867b9b100c14648e4f85de9 updates the MCore Energon requirement from ~=6.0 to ~=7.0, and the bump lockfile moves megatron-energon from 6.0.1 to 7.3.2. Energon 7 removes the singular sample __subflavor__ field and requires batch __key__ / __restore_key__ metadata, which broke Bridge unit tests in PR chore(beep boop 🤖): Bump uv.lock (main, mcore-main) (2026-05-30) #4087.
  • Guards: Added local Energon 6 compatibility guards in src/megatron/bridge/data/energon/metadata.py with TODO removal condition: remove when Bridge no longer supports megatron-energon 6.x. Removed stale guards: none.

Validation

  • 2026-05-30, CW interactive partition: uv run pre-commit run --all-files - passed.
  • 2026-05-30, CW interactive partition: uv run python -m pytest tests/unit_tests/data/energon/test_hf_encoder_task_encoder.py tests/unit_tests/data/energon/test_nemotron_omni_task_encoder.py tests/unit_tests/recipes/qwen_vl/data/energon/test_task_encoder.py tests/unit_tests/diffusion/data/common/test_diffusion_sample.py tests/unit_tests/diffusion/data/common/test_diffusion_task_encoder.py tests/unit_tests/diffusion/data/flux/test_flux_taskencoder.py tests/unit_tests/diffusion/data/wan/test_wan_taskencoder.py tests/unit_tests/models/nemotron_omni/test_nemotron_omni_conversion.py -q - passed, 59 passed, 32 warnings in 1.24s.

dimapihtar and others added 2 commits May 30, 2026 06:10
@dimapihtar @github-actions
chore(beep boop 🤖): Bump uv.lock (main, mcore-main) (2026-05-30)
732aedb 
Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
@yaoyu-33
[data] fix: Support Energon 7 metadata fields
6f742fc 
Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
@copy-pr-bot

copy-pr-bot Bot commented May 30, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@yaoyu-33

yaoyu-33 commented May 30, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 6f742fc

@yaoyu-33 yaoyu-33 mentioned this pull request May 30, 2026
Closed
claude[bot]
claude Bot reviewed May 30, 2026
View reviewed changes
Comment thread src/megatron/bridge/data/energon/hf_encoder_task_encoder.py Outdated
@claude

claude Bot commented May 30, 2026

Copy link
Copy Markdown
Contributor

test

@claude

claude Bot commented May 30, 2026

Copy link
Copy Markdown
Contributor

Light Review

Clean compatibility layer — centralizing the Energon 6/7 metadata logic into metadata.py is the right call.

Findings

  1. Minor bug — duplicated list comprehension (hf_encoder_task_encoder.py:348): keys rebuilds [s.key for s in samples] instead of reusing the keys variable from line 345. The other two batch sites (nemotron_omni_task_encoder.py, qwen_vl/task_encoder.py) correctly use keys=keys. See inline comment.

  2. Missing unit tests for metadata.py itself: The new helpers have meaningful branching logic (subflavor presence, Batch being a non-dataclass, restore_keys default generation). They are exercised indirectly through the encoder tests, but a small focused test of sample_metadata_kwargs and batch_metadata_kwargs would make the compat contract explicit.

  3. Functional test not updated: tests/functional_tests/test_groups/data/energon/test_hf_encoder_task_encoder.py has its own inline compat logic (lines 40-57) that duplicates what sample_metadata_kwargs now provides. Not blocking, but worth a follow-up.

Overall the change looks correct and well-scoped.

Suggested test cases

No perf tests impacted.

@yaoyu-33 yaoyu-33 added area:data Dataset builders, preprocessing, and samplers bug Something isn't working full-test-suite needs-review PR is ready for code review and waiting on a reviewer labels May 30, 2026
@yaoyu-33 @claude
Apply suggestions from code review
c0215a2 
Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com>
Signed-off-by: Yu Yao <54727607+yaoyu-33@users.noreply.github.com>
@yaoyu-33

yaoyu-33 commented May 30, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test c0215a2

@yaoyu-33 yaoyu-33 mentioned this pull request May 30, 2026
Merged
@yaoyu-33

yaoyu-33 commented May 30, 2026

Copy link
Copy Markdown
Contributor Author

Closing as superseded by #4090. PR #4090 now carries the Energon 7 metadata fix from this PR plus the diffusion @stateless cooker fix, so only one bump PR needs to merge.

@yaoyu-33 yaoyu-33 closed this May 30, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@claude claude[bot] claude[bot] left review comments

Assignees

No one assigned

Labels

area:data Dataset builders, preprocessing, and samplers bug Something isn't working full-test-suite needs-review PR is ready for code review and waiting on a reviewer

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@yaoyu-33 @dimapihtar