MegatronRun: support "true"/"false" string convention for bare CLI flags#967
MegatronRun: support "true"/"false" string convention for bare CLI flags#967saivishal1999 wants to merge 6 commits into
Conversation
…boolean flags Extend cmd_args property to treat "true" the same as "" (append bare flag) and skip "false" values entirely. Enables DSE to sweep boolean flags via ["true", "false"] parameter lists in TOML configs.
saivishal1999
requested review from
jeffnvidia,
podkidyshev and
srivatsankrishnan
as code owners
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughChangesMegatron CLI flags
Estimated code review effort: 2 (Simple) | ~10 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 4✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
saivishal1999
changed the title
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/cloudai/workloads/megatron_run/megatron_run.py`:
- Around line 87-92: Restrict the "true"/"false"/empty-string sentinel
conversion in the argument-dumping logic to recognized boolean flags only. For
non-boolean fields such as tensorboard_dir and arbitrary extra arguments,
preserve the literal value unchanged; retain bare-flag behavior for boolean
"true" or empty values and omission for boolean "false".
- Around line 84-92: Add parameterized regression coverage for the
argument-conversion logic that builds flags from args, covering string sentinels
"true" (emits a presence-only flag) and "false" (omits the flag), while also
verifying ordinary values remain unchanged and the existing empty-string
behavior is preserved.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Enterprise
Run ID: 6d23fce5-4452-4e52-935f-30b4d22881d7
📒 Files selected for processing (1)
src/cloudai/workloads/megatron_run/megatron_run.py
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/cloudai/workloads/megatron_run/megatron_run.py`:
- Around line 86-87: Broaden the result annotation in the model serialization
method containing model_dump to reflect preserved native values such as
integers, Path instances, and extra-field values, rather than restricting it to
str or list[str]. Keep the existing model_dump behavior unchanged and ensure the
declared return type covers every value placed into result.
In `@tests/workloads/megatron_run/test_megatron_run.py`:
- Line 28: Update the BASE_ARGS fixture definition to avoid a mutable
class-level dictionary warning: use a dictionary literal and annotate it as
ClassVar, or move the constant to module scope. Preserve the existing
docker_image_url and run_script values.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Enterprise
Run ID: bb07c6d9-b7f2-4e43-b360-036b6a9c3f14
📒 Files selected for processing (2)
src/cloudai/workloads/megatron_run/megatron_run.pytests/workloads/megatron_run/test_megatron_run.py
|
|
||
|
|
||
| class TestMegatronRunCmdArgs: | ||
| BASE_ARGS = dict(docker_image_url="http://url", run_script=Path(__file__)) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
Avoid a mutable class-level fixture dictionary.
Ruff flags BASE_ARGS with RUF012, and dict(...) also triggers C408. Use a dictionary literal and annotate it as ClassVar, or move it to module scope.
🧰 Tools
🪛 GitHub Actions: CI / 3_Linting.txt
[error] 28-28: ruff check (pre-commit hook id: ruff-check) failed: RUF012 Mutable class attributes should be annotated with typing.ClassVar.
🪛 GitHub Actions: CI / Linting
[error] 28-28: ruff check failed (RUF012): Mutable class attributes should be annotated with typing.ClassVar. Offending code: BASE_ARGS = dict(...).
🪛 Ruff (0.15.21)
[warning] 28-28: Unnecessary dict() call (rewrite as a literal)
Rewrite as a literal
(C408)
[warning] 28-28: Mutable default value for class attribute
(RUF012)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/workloads/megatron_run/test_megatron_run.py` at line 28, Update the
BASE_ARGS fixture definition to avoid a mutable class-level dictionary warning:
use a dictionary literal and annotate it as ClassVar, or move the constant to
module scope. Preserve the existing docker_image_url and run_script values.
Source: Linters/SAST tools
saivishal1999
force-pushed
the
spothula/megatron-run-bool-flag
branch
from
e1bea1c to
ed881a8
Compare
There was a problem hiding this comment.
♻️ Duplicate comments (2)
src/cloudai/workloads/megatron_run/megatron_run.py (2)
87-87:⚠️ Potential issue | 🟡 MinorKeep the return annotation consistent with preserved native values.
Line [87] restricts
resultto strings and string lists, but the loop preservesmodel_dump()values unchanged. Fields such asglobal_batch_sizeandlog_intervalare integers, while path fields returnPathobjects; extra fields are unconstrained. Broaden the annotation (for example, todict[str, Any]) or normalize values before returning. This is the same unresolved issue noted in the prior review.#!/bin/bash python - <<'PY' from pathlib import Path from pydantic import BaseModel class Sample(BaseModel): count: int = 1 path: Path = Path("/tmp") dumped = Sample().model_dump() assert isinstance(dumped["count"], int) assert isinstance(dumped["path"], Path) PY🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/cloudai/workloads/megatron_run/megatron_run.py` at line 87, Broaden the return/result annotation in the megatron run flow around the result dictionary to reflect that model_dump() values are preserved unchanged, including integers, Path objects, and unconstrained extra fields; use dict[str, Any] or normalize every value before returning, while keeping the existing native-value behavior.
90-96: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRemove the unnecessary
elifaftercontinue.Ruff
RET507flags Line [93]. Since Line [92] exits the iteration, use a standaloneif:Proposed change
if k in bool_flags or k in extra_keys: if v == "false": continue - elif v in ("true", ""): + if v in ("true", ""): result[flag] = "" else: result[flag] = v🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/cloudai/workloads/megatron_run/megatron_run.py` around lines 90 - 96, Update the flag-processing logic in the visible loop so the condition handling "false" continues immediately, then evaluate the "true" or empty-value case with a standalone if instead of an elif. Preserve the existing result[flag] assignments and behavior for all other values.Source: Linters/SAST tools
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Duplicate comments:
In `@src/cloudai/workloads/megatron_run/megatron_run.py`:
- Line 87: Broaden the return/result annotation in the megatron run flow around
the result dictionary to reflect that model_dump() values are preserved
unchanged, including integers, Path objects, and unconstrained extra fields; use
dict[str, Any] or normalize every value before returning, while keeping the
existing native-value behavior.
- Around line 90-96: Update the flag-processing logic in the visible loop so the
condition handling "false" continues immediately, then evaluate the "true" or
empty-value case with a standalone if instead of an elif. Preserve the existing
result[flag] assignments and behavior for all other values.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Enterprise
Run ID: 9c3f6795-1d27-497a-9120-ed9ae3eb98f0
📒 Files selected for processing (1)
src/cloudai/workloads/megatron_run/megatron_run.py
…ared return type model_dump() returns native Python types (int, Path) for typed fields; converting to str at insertion keeps the dict[str, Union[str, list[str]]] annotation accurate without cascading changes to the base class.
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/cloudai/workloads/megatron_run/megatron_run.py`:
- Around line 96-98: Update the boolean-style flag detection near bool_flags in
the Megatron argument-building logic to explicitly include
recompute_activations, even though its default is None. Ensure the existing
result construction emits a bare --recompute-activations for true and omits the
flag for false, while preserving normal string-value handling for other flags.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Enterprise
Run ID: d435f5d8-da7c-4764-8f90-3fabdc2b8d52
📒 Files selected for processing (1)
src/cloudai/workloads/megatron_run/megatron_run.py
model_dump() preserves native Python types (int, Path); the prior Union[str, list[str]] annotation was narrower than what was actually stored. Existing test_tokenizer_model confirms Path values must be preserved.
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/cloudai/workloads/megatron_run/megatron_run.py`:
- Around line 82-85: Update MegatronRunCmdArgs.cmd_args boolean-style flag
detection to explicitly include recompute_activations alongside fields
identified by an empty-string default and model_extra keys. Ensure true emits
the bare --recompute-activations flag and false omits the flag, preserving
existing handling for other arguments.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Enterprise
Run ID: 2fda69c7-8cb6-4423-aa28-be1bce01f4a0
📒 Files selected for processing (1)
src/cloudai/workloads/megatron_run/megatron_run.py
default=None keeps it omittable; "true" now emits bare --recompute-activations and "false" suppresses it, consistent with the bool flag convention.
1 participant
Currently the only way to emit a bare CLI flag in MegatronRun is to set
the parameter to
"". This makes it impossible for a DSE search spaceto ever disable a boolean flag — there is no value that means "skip this flag."
This PR adds a string convention to
MegatronRunCmdArgs.cmd_args:"true"→ emits the bare flag (e.g.--mock-data)"false"→ skips the flag entirely""→ emits the bare flag (backward compatible, existing TOMLs unchanged)This allows DSE configs to use
["true", "false"]as a search space totoggle boolean flags on/off across trials.
Existing TOMLs using
= ""continue to work without modification.