Add argparse_utils and enhance typer_utils for megatron runner#769
Add argparse_utils and enhance typer_utils for megatron runner#769
Conversation
Add dumper CLI arguments (--dumper-enable, --dumper-dir, per-phase config), dumper_utils.py for SGLang/Megatron dumper integration, model.py hooks for forward-only and forward-backward phases, rollout env var plumbing, source patcher wiring in training actors, and basic e2e test.
Add _maybe_apply_dumper_overrides to disable heartbeats, force single rollout, and disable eval/save when --dumper-enable is set.
Add conftest_dumper.py with shared source patcher YAML configs and comparator helpers. Expand test_dumper.py with full MoE parallelism coverage, field verification, and cross-framework (SGLang vs Megatron) activation comparison. Update dumper_utils.py to nest engine dumps under engines/ subdirectory.
Add DataclassArgparseBridge for dataclass-to-argparse conversion, refactor typer_utils with dataclass_cli decorator, and add parallel_utils for running parallel CLI commands. Include tests.
fzyzcjy
requested review from
guapisolo,
maocheng23 and
yueming-yuan
as code owners
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly improves the handling of command-line arguments by introducing new utilities that bridge Python dataclasses with popular CLI parsing libraries like Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a helpful DataclassArgparseBridge and enhances typer_utils with more features and robust testing. The changes are well-structured and the new utilities are accompanied by comprehensive tests. I've found a couple of areas for improvement in argparse_utils.py to enhance compatibility and correctness, particularly around type checking for optional types and handling of boolean flags with True defaults.
| def _is_optional(tp: type) -> tuple[bool, type | None]: | ||
| if not isinstance(tp, types.UnionType): | ||
| return False, None | ||
|
|
||
| args: tuple[type, ...] = tp.__args__ | ||
| non_none: list[type] = [a for a in args if a is not type(None)] | ||
| if type(None) not in args or len(non_none) != 1: | ||
| return False, None | ||
|
|
||
| return True, non_none[0] |
There was a problem hiding this comment.
The current implementation of _is_optional using isinstance(tp, types.UnionType) is only compatible with Python 3.10+. To support typing.Union from older Python versions, it's better to use typing.get_origin and typing.get_args. This will make the utility more robust across different Python versions.
| def _is_optional(tp: type) -> tuple[bool, type | None]: | |
| if not isinstance(tp, types.UnionType): | |
| return False, None | |
| args: tuple[type, ...] = tp.__args__ | |
| non_none: list[type] = [a for a in args if a is not type(None)] | |
| if type(None) not in args or len(non_none) != 1: | |
| return False, None | |
| return True, non_none[0] | |
| def _is_optional(tp: type) -> tuple[bool, type | None]: | |
| origin = get_origin(tp) | |
| if origin is Union or (hasattr(types, "UnionType") and origin is types.UnionType): | |
| args = get_args(tp) | |
| non_none: list[type] = [a for a in args if a is not type(None)] | |
| if type(None) in args and len(non_none) == 1: | |
| return True, non_none[0] | |
| return False, None |
| if _is_bool(tp): | ||
| group.add_argument(flag, dest=dest, action="store_true", default=False) | ||
| continue |
There was a problem hiding this comment.
The current handling of boolean fields assumes they default to False. If a dataclass field is defined with my_flag: bool = True, this implementation will incorrectly set the argument's default to False. This means if the flag isn't provided on the command line, the value will be False, ignoring the dataclass default.
A more robust approach would be to respect the dataclass default, for example by using argparse.BooleanOptionalAction (Python 3.9+) or manually implementing both --flag and --no-flag arguments. This would likely require updating to_cli_args as well to be consistent.
1 participant
No description provided.