Allow profiling ranks in string format with comma as separator

[juntaowww]

Summary

Allow profiling ranks in string format with comma as separator. The list format in cmd_args would make it DSE target.

Test Plan

manual run

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro

Run ID: ec3b3454-94e5-4c30-9696-3901f9e94ae7

📥 Commits

Reviewing files that changed from the base of the PR and between 89287e8 and 8054958.

📒 Files selected for processing (1)

• tests/workloads/megatron_bridge/test_command_gen_strategy_slurm.py

---

📝 Walkthrough

Walkthrough

Updated the MegatronBridgeCmdArgs.profiling_ranks field type to accept str in addition to int and List[int], and added a Field description documenting comma-separated rank IDs (e.g. "0,4,8"). Added a unit test verifying string-format input is propagated into generated Slurm wrapper content.

Changes

Cohort / File(s)	Summary
Megatron Bridge Configuration src/cloudai/workloads/megatron_bridge/megatron_bridge.py	Expanded profiling_ranks field type from Optional[Union[int, List[int]]] to Optional[Union[int, str, List[int]]] and updated the Field(...) description to document comma-separated string input support.
Unit Tests — Megatron Bridge tests/workloads/megatron_bridge/test_command_gen_strategy_slurm.py	Added test_profiling_ranks_string_format to assert that a string "0,1,2,3" for profiling_ranks with enable_nsys produces a Slurm wrapper containing --profiling_ranks 0,1,2,3 and that the TestRun is not a DSE job before generation.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Poem

🐰 I nibble on commas, hop in delight,

"0,1,2,3" shines in the wrapper's light,
No more strict ints to bind my paws,
Flexibility hops in, without a pause.

🚥 Pre-merge checks | ✅ 2

✅ Passed checks (2 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title accurately describes the main change: allowing profiling ranks to be specified as a comma-separated string format in addition to existing formats.
Description check	✅ Passed	The description is related to the changeset, explaining the motivation (avoiding DSE target designation) and test approach, though it lacks technical depth.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 2

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@src/cloudai/workloads/megatron_bridge/megatron_bridge.py`:
- Line 116: The type annotation for the profiling_ranks Field should use modern
PEP 604 syntax: replace Optional[Union[int, str, List[int]]] with int | str |
list[int] | None for clarity and consistency; remove any now-unused typing
imports (Optional, Union, List) or update imports accordingly and ensure the
annotation is applied on the profiling_ranks Field definition in
megatron_bridge.py.
- Around line 116-119: Add a Pydantic field validator for profiling_ranks to
validate and normalize comma-separated string input: when profiling_ranks is a
str, ensure it matches only digits separated by commas (e.g. "0,4,8"), split and
convert to a list[int] (or raise a ValidationError on invalid format); keep
existing support for int and list[int]. Attach the validator to the
profiling_ranks field (use the model's validator/@field_validator for
"profiling_ranks") so downstream command generation always receives a validated
list[int] or a single int.

🪄 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: Pro

Run ID: 21a82a8b-87a4-4cbf-bc66-1a1d6a2de461

📥 Commits

Reviewing files that changed from the base of the PR and between f0ffef9 and 89287e8.

📒 Files selected for processing (1)

• src/cloudai/workloads/megatron_bridge/megatron_bridge.py

[podkidyshev]

@juntaowww let me know when to merge

[juntaowww]

good to merge now, thanks!