Add pre-test hook support for MegatronBridge workload#957
Merged
podkidyshev merged 8 commits intoJul 9, 2026
Conversation
saivishal1999
requested review from
jeffnvidia,
podkidyshev and
srivatsankrishnan
as code owners
July 2, 2026 07:09
Contributor
|
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:
📝 WalkthroughWalkthroughThis PR adds optional pre-test sbatch handling for Megatron-Bridge Slurm command generation, refactors shared SBATCH resource directive assembly, and introduces a new NCCL ChangesMegatron-Bridge pre-hook sbatch flow
Shared SBATCH resource directive assembly
NCCL bus bandwidth threshold
Estimated code review effort: 4 (Complex) | ~45 minutes 🚥 Pre-merge checks | ✅ 4✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 5
🤖 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_bridge/slurm_command_gen_strategy.py`:
- Around line 133-136: The pre-hook sbatch generation in build_sbatch_lines is
dropping excluded nodes when get_cached_nodes_spec() only provides num_nodes.
Update the sbatch line assembly so self.test_run.exclude_nodes is preserved even
in the --nodes path, not just when node_list is present, and keep the existing
behavior around num_nodes in SlurmCommandGenStrategy while ensuring excluded
nodes are still passed through.
- Around line 117-118: The pre-test node assignment in
slurm_command_gen_strategy.py is using the read-only TestRun.nnodes property,
which will fail before sbatch generation. Update the loop in the strategy method
that iterates over self.test_run.pre_test.test_runs to assign the underlying
mutable field num_nodes instead of nnodes, keeping the existing logic that
copies the parent test run’s node count into each pre-test run.
- Line 139: The pre-hook submission in gen_pre_test / SlurmCommandGenStrategy
currently only propagates PRE_TEST_SUCCESS, so a failed pre-test can still let
the sbatch wrapper exit successfully and release the main job. Update the
command assembly around sbatch_lines.extend and pre_hook_cmds to explicitly exit
non-zero when the pre-test check fails, using the gen_pre_test-generated success
flag as the gate so the pre-hook sbatch cannot complete as success on failure.
- Around line 344-350: Make the include_slurm_params parameter keyword-only in
_build_launcher_parts so it cannot be passed positionally and the FBT001/FBT002
warning is eliminated. Update the _build_launcher_parts function signature to
force include_slurm_params as a keyword-only argument, and ensure any internal
call sites use the named argument form; the existing keyword-based caller should
continue to work unchanged.
- Around line 268-270: The shell wrapper in the Slurm command generation logic
is inserting base Slurm args directly into the ADDITIONAL_SLURM_PARAMS
assignment, which can allow shell expansion and produce a leading separator when
the base value is empty. Update the Megatron bridge’s slurm_command_gen_strategy
around the ADDITIONAL_SLURM_PARAMS construction to safely quote/escape
base_slurm_params before embedding it in the shell snippet, and adjust the
PRE_HOOK_JOB_ID append logic so it only adds the dependency separator when
ADDITIONAL_SLURM_PARAMS already has content. Use the existing command assembly
in the strategy method that emits these lines to keep the fix localized.
🪄 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: dce884bc-1a4a-49c7-bf8d-0aa679ff96f6
📒 Files selected for processing (1)
src/cloudai/workloads/megatron_bridge/slurm_command_gen_strategy.py
Contributor
|
@saivishal1999 please resolve coderabbit comments |
podkidyshev
requested changes
Jul 3, 2026
Contributor
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_bridge/slurm_command_gen_strategy.py`:
- Around line 299-301: The `capture_nodelist` parameter in
`SlurmCommandGenStrategy` should be keyword-only to match the existing
`include_slurm_params` pattern and clear the FBT001/FBT002 warning. Update the
affected method signature in `slurm_command_gen_strategy.py` to make
`capture_nodelist` keyword-only, and ensure any call sites use the keyword form
while keeping the existing named invocation compatible.
🪄 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: b2b19501-1bf3-4c77-9c80-043760100c93
📒 Files selected for processing (3)
src/cloudai/workloads/megatron_bridge/slurm_command_gen_strategy.pysrc/cloudai/workloads/nccl_test/nccl.pysrc/cloudai/workloads/nccl_test/slurm_command_gen_strategy.py
Contributor
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/nccl_test/slurm_command_gen_strategy.py`:
- Around line 34-43: The `_mpi` suffix check is duplicated in both
`gen_srun_prefix` and `generate_test_command`, so extract it into a shared
helper or property on this strategy class such as an `is_single_process`
accessor. Update both call sites to use that single source of truth in
`NCCLTestDefinition`-based logic so the naming convention is maintained in one
place and the two methods cannot drift.
🪄 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: 49d63a27-0267-4e3e-9ec4-b187c1038374
📒 Files selected for processing (1)
src/cloudai/workloads/nccl_test/slurm_command_gen_strategy.py
Contributor
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/cloudai/workloads/megatron_bridge/slurm_command_gen_strategy.py (1)
315-330: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winUnbounded
while truecan hang the submit process forever.This loop only breaks on
_state = "RUNNING". Two reachable paths never satisfy that:
- The pre-hook fails/gets cancelled while
PENDING, so it leaves the queue andsqueuereturns empty.- A short pre-hook transitions
PENDING → RUNNING → COMPLETEDbetween two 10s polls, soRUNNINGis never observed.In both cases
_statebecomes empty and the loop spins indefinitely (sleep 10), blocking the wrapper on the submit node. Additionally, Line 328 appendsnodelist=unconditionally, so if capture failsPRE_HOOK_NODESis empty and--additional_slurm_paramsgets a danglingnodelist=.Break out on a terminal/empty state and only append the nodelist when non-empty.
🐛 Proposed fix
"while true; do", ' _state=$(squeue -j "$PRE_HOOK_JOB_ID" -h -o "%T" 2>/dev/null || true)', ' if [ "$_state" = "RUNNING" ]; then', ' PRE_HOOK_NODES=$(squeue -j "$PRE_HOOK_JOB_ID" -h -o "%N" 2>/dev/null | head -1 || true)', " break", " fi", + ' if [ -z "$_state" ]; then', + ' echo "Pre-hook $PRE_HOOK_JOB_ID left the queue before RUNNING was observed." >&2', + " break", + " fi", " sleep 10", "done", - 'ADDITIONAL_SLURM_PARAMS="${ADDITIONAL_SLURM_PARAMS};nodelist=${PRE_HOOK_NODES}"', + 'if [ -n "$PRE_HOOK_NODES" ]; then', + ' ADDITIONAL_SLURM_PARAMS="${ADDITIONAL_SLURM_PARAMS};nodelist=${PRE_HOOK_NODES}"', + "fi", "",🤖 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_bridge/slurm_command_gen_strategy.py` around lines 315 - 330, The nodelist capture logic in slurm_command_gen_strategy.py can hang forever because the `while true` loop in the `capture_nodelist` block only exits on `RUNNING` and never handles empty or terminal states. Update the `capture_nodelist` branch in `slurm_command_gen_strategy.py` to break on failed/cancelled/completed or empty `_state`, and only add `nodelist=${PRE_HOOK_NODES}` to `ADDITIONAL_SLURM_PARAMS` when `PRE_HOOK_NODES` is non-empty. Use the existing `capture_nodelist` setup, `PRE_HOOK_JOB_ID`, and `PRE_HOOK_NODES` symbols to locate the fix.
🤖 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.
Outside diff comments:
In `@src/cloudai/workloads/megatron_bridge/slurm_command_gen_strategy.py`:
- Around line 315-330: The nodelist capture logic in
slurm_command_gen_strategy.py can hang forever because the `while true` loop in
the `capture_nodelist` block only exits on `RUNNING` and never handles empty or
terminal states. Update the `capture_nodelist` branch in
`slurm_command_gen_strategy.py` to break on failed/cancelled/completed or empty
`_state`, and only add `nodelist=${PRE_HOOK_NODES}` to `ADDITIONAL_SLURM_PARAMS`
when `PRE_HOOK_NODES` is non-empty. Use the existing `capture_nodelist` setup,
`PRE_HOOK_JOB_ID`, and `PRE_HOOK_NODES` symbols to locate the fix.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Enterprise
Run ID: 6c358501-1776-4c18-8ca3-00d4ad44528b
📒 Files selected for processing (2)
src/cloudai/systems/slurm/slurm_command_gen_strategy.pysrc/cloudai/workloads/megatron_bridge/slurm_command_gen_strategy.py
MegatronBridge's gen_exec_command uses setup_experiment.py to generate its own sbatch, bypassing CloudAI's normal pre_test injection. This adds support by submitting the pre-hook as a separate sbatch first, then passing dependency=afterok:<job_id> to setup_experiment.py via --additional_slurm_params so Slurm gates training on pre-hook success. - _gen_pre_hook_sbatch: generates standalone sbatch for pre-hook tests, injecting main test's nnodes so srun allocates all training nodes - _collect_additional_slurm_params: extracted from _build_launcher_parts so the base params can be built in Python and dependency appended at bash runtime after the pre-hook job ID is known - _wrap_launcher_for_job_id_and_quiet_output: submits pre-hook sbatch, captures job ID, builds --additional_slurm_params with dependency Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
tr.nnodes is a read-only property backed by num_nodes; set num_nodes directly so srun gets -N<main_job_nodes> instead of -N1. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…Bridge - _gen_pre_hook_sbatch: loop-based independent per-node alltoalls (32 x srun --nodelist=$_node), busbw threshold check via awk (min_busbw from hook TOML), exit 1 on failure to block training, forwards extra_srun_args (e.g. --segment) as #SBATCH directives so pre-hook gets same node placement constraints as training - _wrap_launcher_for_job_id_and_quiet_output: submit pre-hook as separate sbatch, pass dependency=afterok:<id> to setup_experiment.py; poll indefinitely until pre-hook RUNNING, capture nodelist, pass to training via --additional_slurm_params - NCCLCmdArgs: add optional min_busbw field for per-hook threshold configuration - NcclTestSlurmCommandGenStrategy: exclude min_busbw from CLI args forwarded to nccl binary Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Switch alltoall_perf from MPI multi-process to single-process multi-GPU mode: replace --mpi=pmix with --ntasks=1 and use -g <ngpus> instead of --ngpus so one process drives all GPUs directly.
…thod Extract resource-related SBATCH directives (reservation, distribution, nodelist/exclude, gpus-per-node/gres, ntasks-per-node, time limit, extra_sbatch_args) into a new base class method _append_resource_directives. Refactor _append_sbatch_directives to call it, and update MegatronBridge _gen_pre_hook_sbatch to use it instead of reimplementing a subset of directives — fixing missing: distribution, exclude_nodes, ntasks-per-node, gres, and extra_sbatch_args in the pre-hook.
saivishal1999
force-pushed
the
spothula/nccl-alltoall-pre-hook
branch
from
July 8, 2026 23:52
1d26948 to
aa0aba2
Compare
Contributor
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/systems/slurm/slurm_command_gen_strategy.py`:
- Line 390: The assignment in SlurmCommandGenStrategy._add_reservation usage is
redundant because the method mutates the existing content list and returns the
same object. Update the surrounding code in slurm_command_gen_strategy.py to
call _add_reservation(content) without reassigning content, keeping the in-place
mutation pattern consistent and removing the unnecessary no-op.
In `@src/cloudai/workloads/megatron_bridge/slurm_command_gen_strategy.py`:
- Around line 315-330: The nodelist capture loop in the pre-hook wrapper can
wait forever if the job never reaches RUNNING or exits between polls; update the
logic in the nodelist_lines block to handle terminal states and add a maximum
wait timeout. Use the existing PRE_HOOK_JOB_ID and squeue polling to detect
COMPLETED/FAILED/CANCELLED (or empty output after the job leaves the queue),
then break with an error instead of looping forever. Keep the current
capture_nodelist behavior, but make the shell snippet fail fast when the
pre-hook is no longer runnable.
🪄 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: e979cbb7-15d5-4a9e-a7cf-5f8a813df648
📒 Files selected for processing (4)
src/cloudai/systems/slurm/slurm_command_gen_strategy.pysrc/cloudai/workloads/megatron_bridge/slurm_command_gen_strategy.pysrc/cloudai/workloads/nccl_test/nccl.pysrc/cloudai/workloads/nccl_test/slurm_command_gen_strategy.py
- Fix infinite nodelist poll loop: detect terminal Slurm states
(FAILED/CANCELLED/COMPLETED/TIMEOUT/gone) and exit with error instead
of spinning forever
- Remove redundant reassignment of _add_reservation return value; the
method mutates content in-place
- Extract _is_single_process property in NcclTestSlurmCommandGenStrategy
to deduplicate the endswith("_mpi") check across gen_srun_prefix and
generate_test_command
- Fix docstring and formatting issues flagged by ruff
- Assert pre_test is not None before accessing .test_runs in _gen_pre_hook_sbatch; pyright can't infer the caller already guards it - Handle list branch of subtest_name Union type in _is_single_process; subtest_name is Union[Literal[...], list[Literal[...]]] so .endswith() requires narrowing to str first
Terminal state detection was already added; this adds a 3600s (1h) hard ceiling so the loop cannot spin indefinitely when the pre-hook job stays PENDING longer than expected. Exits with an error after timeout.
podkidyshev
reviewed
Jul 9, 2026
podkidyshev
reviewed
Jul 9, 2026
podkidyshev
reviewed
Jul 9, 2026
podkidyshev
reviewed
Jul 9, 2026
podkidyshev
approved these changes
Jul 9, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
MegatronBridge's gen_exec_command uses setup_experiment.py to generate its own sbatch, bypassing CloudAI's normal pre_test injection. This adds support by submitting the pre-hook as a separate sbatch first, then passing dependency=afterok:<job_id> to setup_experiment.py via --additional_slurm_params so Slurm gates training on pre-hook success.