Skip to content

MBridge: time limit managed by test run #849

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
podkidyshev wants to merge 1 commit into
base: main
Choose a base branch
from
Open

MBridge: time limit managed by test run #849

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 0 additions & 1 deletion src/cloudai/workloads/megatron_bridge/megatron_bridge.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,6 @@ class MegatronBridgeCmdArgs(CmdArgs):
# Slurm/launcher-level
gpu_type: str = Field(default="gb200")
log_dir: str = Field(default="")
time_limit: str = Field(default="00:05:00")
container_image: str = Field(default="")
num_gpus: int = Field(default=8)
gpus_per_node: int = Field(default=8)
Expand Down
2 changes: 1 addition & 1 deletion src/cloudai/workloads/megatron_bridge/slurm_command_gen_strategy.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -325,7 +325,7 @@ def add_field(field: str, flag: str, value: Any) -> None:
add("-p", self.system.default_partition)
add_field("gpu_type", "-g", args.gpu_type)
add_field("log_dir", "-l", args.log_dir)
add_field("time_limit", "-t", args.time_limit)
add("-t", self.test_run.time_limit)
if container_path:
add_field("container_image", "-i", container_path)
add_field("compute_dtype", "-c", args.compute_dtype)
Expand Down
2 changes: 1 addition & 1 deletion tests/ref_data/megatron-bridge.sbatch
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@ if [ "${WANDB_INSTALL_RC}" -ne 0 ]; then
fi

LAUNCH_RC=0
NEMORUN_HOME="__OUTPUT_DIR__/output" __INSTALL_DIR__/Run__main-venv/bin/python __INSTALL_DIR__/Megatron-Bridge__main/scripts/performance/setup_experiment.py -p main -i __OUTPUT_DIR__/output/megatron_bridge_image.sqsh -hf dummy_token -ng 8 -gn 4 --golden_values_path cloudai_megatron_bridge_golden_values.json -cm __INSTALL_DIR__/Megatron-Bridge__main:/opt/Megatron-Bridge -cb 'export CUDA_VISIBLE_DEVICES=0,1,2,3' -cb 'export NCCL_DEBUG=INFO' -m qwen3 -mr 30b_a3b --detach false --additional_slurm_params 'gpus-per-node=4;gres=gpu:4' >>"$LOG" 2>&1 || LAUNCH_RC=$?
NEMORUN_HOME="__OUTPUT_DIR__/output" __INSTALL_DIR__/Run__main-venv/bin/python __INSTALL_DIR__/Megatron-Bridge__main/scripts/performance/setup_experiment.py -p main -t 00:20:00 -i __OUTPUT_DIR__/output/megatron_bridge_image.sqsh -hf dummy_token -ng 8 -gn 4 --golden_values_path cloudai_megatron_bridge_golden_values.json -cm __INSTALL_DIR__/Megatron-Bridge__main:/opt/Megatron-Bridge -cb 'export CUDA_VISIBLE_DEVICES=0,1,2,3' -cb 'export NCCL_DEBUG=INFO' -m qwen3 -mr 30b_a3b --detach false --additional_slurm_params 'gpus-per-node=4;gres=gpu:4' >>"$LOG" 2>&1 || LAUNCH_RC=$?


JOB_ID=""
Expand Down
5 changes: 3 additions & 2 deletions tests/test_acceptance.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -173,8 +173,8 @@ def partial_tr(slurm_system: SlurmSystem) -> partial[TestRun]:
return partial(TestRun, num_nodes=1, nodes=[], output_path=slurm_system.output_path)


def create_test_run(partial_tr: partial[TestRun], name: str, test_definition: TestDefinition) -> TestRun:
tr = partial_tr(name=name, test=test_definition)
def create_test_run(partial_tr: partial[TestRun], name: str, test_definition: TestDefinition, **kwargs) -> TestRun:
tr = partial_tr(name=name, test=test_definition, **kwargs)
return tr
Comment on lines +176 to 178
Copy link
Contributor

@coderabbitai coderabbitai bot Mar 25, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 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.



Expand Down Expand Up @@ -576,6 +576,7 @@ def test_req(request, slurm_system: SlurmSystem, partial_tr: partial[TestRun]) -
)
],
),
time_limit="00:20:00",
),
"vllm": lambda: create_test_run(
partial_tr,
Expand Down
Loading