MBridge: time limit managed by test run

[podkidyshev]

Summary

Manage time_limit consistently with other workloads via TestRun.time_limit

Test Plan

• Automated CI

Additional Notes

• Closes CLOUDAI-33

[coderabbitai]

📝 Walkthrough

Walkthrough

The changes relocate the time_limit parameter from command-line arguments to the test run configuration layer. The parameter is removed from MegatronBridgeCmdArgs, the command generation strategy now uses test_run.time_limit directly, and test utilities are updated to support passing time limits at the test run level rather than through command arguments.

Changes

Cohort / File(s)	Summary
Megatron Bridge Time-Limit Refactoring src/cloudai/workloads/megatron_bridge/megatron_bridge.py, src/cloudai/workloads/megatron_bridge/slurm_command_gen_strategy.py	Removed time_limit field from MegatronBridgeCmdArgs and updated command generation to use self.test_run.time_limit instead of args.time_limit for the -t launcher argument.
Test Infrastructure Updates tests/test_acceptance.py, tests/ref_data/megatron-bridge.sbatch	Extended create_test_run() to accept **kwargs for flexible test run parameter injection; added time_limit="00:20:00" to megatron-bridge test instantiation; updated sbatch reference with explicit -t 00:20:00 argument.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~12 minutes

Poem

🐰 The time limit hopped along its way,
From args to test_run, hip hooray!
No defaults to carry with a frown,
Just test_run's wisdom in the town! ⏰

🚥 Pre-merge checks | ✅ 2

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description check	✅ Passed	The description clearly relates to the changeset, explaining that time_limit management is being standardized via TestRun.time_limit.
Title check	✅ Passed	The title accurately reflects the main change: removing time_limit from MegatronBridgeCmdArgs and managing it via TestRun.time_limit instead, which aligns with the PR objective to manage time_limit consistently with other workloads.

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

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch ipod/mbrdige-time

---

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

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

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

Inline comments:
In `@tests/test_acceptance.py`:
- Around line 176-178: Add a type annotation for the variadic kwargs in
create_test_run: change the signature of create_test_run to annotate **kwargs as
typing.Any (or Mapping[str, Any]) and import Any from typing at the top of the
file so static analysis recognizes the type for forwarded TestRun fields like
time_limit; keep the rest of the function unchanged and ensure the import is
added to the existing typing imports.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro

Run ID: 0656b233-29d3-4fa4-be8f-8a0664ea8ea9

📥 Commits

Reviewing files that changed from the base of the PR and between 96f185e and bd5acee.

📒 Files selected for processing (4)

• src/cloudai/workloads/megatron_bridge/megatron_bridge.py
• src/cloudai/workloads/megatron_bridge/slurm_command_gen_strategy.py
• tests/ref_data/megatron-bridge.sbatch
• tests/test_acceptance.py

💤 Files with no reviewable changes (1)

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

[coderabbitai]

🧹 Nitpick | 🔵 Trivial

LGTM! Consider adding type annotation for **kwargs.

The **kwargs passthrough pattern correctly allows forwarding additional TestRun fields (like time_limit) to the partial constructor. This is a clean approach to support the new time_limit handling without modifying every test case.

The static analysis flags a missing type annotation for **kwargs. This is a minor improvement that could enhance clarity:

🔧 Optional: Add type annotation

-def create_test_run(partial_tr: partial[TestRun], name: str, test_definition: TestDefinition, **kwargs) -> TestRun:
+def create_test_run(partial_tr: partial[TestRun], name: str, test_definition: TestDefinition, **kwargs: Any) -> TestRun:

This would require adding Any to the imports from typing.

🧰 Tools

🪛 Ruff (0.15.6)

[warning] 176-176: Missing type annotation for **kwargs

(ANN003)

---

[warning] 178-178: Unnecessary assignment to tr before return statement

Remove unnecessary assignment

(RET504)

🤖 Prompt for AI Agents

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

In `@tests/test_acceptance.py` around lines 176 - 178, Add a type annotation for
the variadic kwargs in create_test_run: change the signature of create_test_run
to annotate **kwargs as typing.Any (or Mapping[str, Any]) and import Any from
typing at the top of the file so static analysis recognizes the type for
forwarded TestRun fields like time_limit; keep the rest of the function
unchanged and ensure the import is added to the existing typing imports.