[https://nvbugs/6541343][fix] Add slurm_wait_all_ranks() — a job+step-keyed marker barrier on the shared… - #17367
Open
trtllm-agent wants to merge 1 commit into
Open
Conversation
The only install lock lives under $resourcePathNode (/tmp) in slurm_install.sh, which is node-local, so its wait loop fences just the $SLURM_LOCALID peers on the same node -- a node can never observe another node's lock. Nothing then stops slurm_run.sh from reaching `eval $pytestCommand` on one node while another is still installing, and the per-rank work in between skews the ranks further: non-zero ranks cover rank 0's coverage-config write with a blind `sleep 30`, and slurm_setup_runtime_env shells out to pip3. Pytest's first action is `import tensorrt_llm`, whose module-scope MPI collective must be entered by every rank. Under --mpi=pmix, which is added exactly when nodeCount > 1, that collective has a 300s fence timeout, so a node whose pip3 install stalls (up to the 2700s retry budget) makes the collective abort every rank rather than merely run late. The ranks die between pytest setup and teardown, which leaves the nodeid in unfinished_test.txt and makes generate_timeout_xml.py synthesize the "Test terminated unexpectedly" this bug reports -- there is no traceback for the test body. Note that PMIX_MCA_gds=hash does not mitigate this: a fence that times out never exchanges the modex regardless of GDS mode, and the pml_ucx errors seen alongside it are downstream of the same missing exchange. Add a marker barrier on the shared $jobWorkspace immediately before `eval $pytestCommand`. It counts SLURM_NTASKS rank markers rather than nodes, so the fenced set is exactly the set that enters the aborting collective, and placing it after the block that wipes SLURM_* keeps it a no-op for single-node and disaggregated benchmark/server runs, which reach it with SLURM_NTASKS unset. The marker directory is keyed per job and per step because $jobWorkspace outlives a step, so a later step must not be released by an earlier one's markers. The wait is bounded above the 2700s pip3 budget so a genuinely dead rank fails the stage with a clear message instead of hanging until the partition walltime. Verified by asserting elapsed seconds, since a no-op barrier also returns 0: single rank, unset workspace and unset SLURM_NTASKS return at 0s; a peer 25s late makes rank 0 block 30s; a dead rank fails bounded with rc=1; markers from another job and from an earlier step of the same job both correctly time out; all ranks concurrent release at 0s; and only rank 0 logs progress. A live 2-node 8-rank probe with rank 7 delayed 20s converged every rank on both nodes at 20-30s. The target test then passed end to end on a 2-node 8-rank pmix launch (1 passed, 399.82s) with no PMIx fence errors, reaching all three accuracy phases: GSM8K 94.845 against a 92.217 threshold and CnnDailymail rouge1 30.240 against 26.716. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
Contributor
WalkthroughThe Slurm runner adds a shared barrier before pytest. Each rank writes a readiness marker in a job- and step-specific workspace directory. Rank 0 reports progress. The barrier bypasses single-rank or workspace-less runs and times out after 3600 seconds. ChangesSlurm barrier
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant SlurmRank
participant JobWorkspace
participant Pytest
SlurmRank->>JobWorkspace: Create rank readiness marker
SlurmRank->>JobWorkspace: Wait for all configured rank markers
JobWorkspace-->>SlurmRank: Return readiness progress
SlurmRank->>Pytest: Start pytest after barrier
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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: 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 `@jenkins/scripts/slurm_run.sh`:
- Around line 75-85: Add the standard NVIDIA copyright header to the top of
slurm_run.sh, using 2026 as the latest modification year. Preserve the existing
script content and behavior unchanged.
🪄 Autofix
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: CHILL
Plan: Enterprise
Run ID: 87e29a98-56b2-458d-aa3f-d764e78df970
📒 Files selected for processing (1)
jenkins/scripts/slurm_run.sh
Comment on lines
+75
to
+85
| # The install lock in slurm_install.sh lives under $resourcePathNode (/tmp), so it | ||
| # is node-local: its wait loop only fences $SLURM_LOCALID peers on the same node, | ||
| # and a node can never observe another node's lock. Nothing else stops one node | ||
| # from reaching `eval $pytestCommand` below while another is still installing, and | ||
| # the per-rank work above skews the ranks further (non-zero ranks cover the | ||
| # coverage-config write with a blind `sleep 30`, and slurm_setup_runtime_env shells | ||
| # out to pip3). Pytest's first action is `import tensorrt_llm`, whose module-scope | ||
| # MPI collective must be entered by every rank; under --mpi=pmix -- added exactly | ||
| # when nodeCount > 1 -- that collective has a 300s fence timeout, so the skew | ||
| # aborts every rank instead of merely running late. Fence every rank on the shared | ||
| # $jobWorkspace so they enter pytest together. |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '1,25p' jenkins/scripts/slurm_run.shRepository: NVIDIA/TensorRT-LLM
Length of output: 1013
Add the NVIDIA copyright header
jenkins/scripts/slurm_run.sh has no NVIDIA copyright header. Add the header with 2026 as the latest modification year.
🤖 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 `@jenkins/scripts/slurm_run.sh` around lines 75 - 85, Add the standard NVIDIA
copyright header to the top of slurm_run.sh, using 2026 as the latest
modification year. Preserve the existing script content and behavior unchanged.
Source: Coding guidelines
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.
Summary
$resourcePathNode=/tmp), so nothing fences ranks across nodes before pytest; the 300s PMIx fence timeout onimport tensorrt_llm's module-scope MPI collective then aborts every rank when one node installs late.slurm_wait_all_ranks()— a job+step-keyed marker barrier on the shared$jobWorkspacecountingSLURM_NTASKSrank markers — immediately beforeeval $pytestCommand, placed after theSLURM_*wipe so single-node/disagg paths no-op.pytest tests/integration/defs/accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_nvfp4_multi_gpus[throughput] -vTest plan
Links
Dev Engineer Review
slurm_wait_all_ranks()injenkins/scripts/slurm_run.sh.$jobWorkspace.SLURM_NTASKSranks before pytest starts.SLURM_*environment reset and does not affect single-node or disaggregated paths.QA Engineer Review
No test changes.