[None][fix] Pass CUDA_VISIBLE_DEVICES as script arg instead of srun --export#12370
Merged
Merged
Conversation
…-export srun --export cannot reliably pass comma-separated values (like CUDA_VISIBLE_DEVICES=0,1) inside shared pyxis containers. The shared container retains the environment from its initial creation, ignoring subsequent --export overrides. This caused all workers to see all GPUs instead of their assigned subset, leading to GPU OOM during model load. Fix: pass CUDA_VISIBLE_DEVICES as a command-line argument to start_worker.sh and export it inside the script, bypassing srun --export entirely for this variable. Signed-off-by: Xianjie <[email protected]>
Collaborator
Author
|
/bot skip --comment "slurm scripts are not protected by CI" |
Contributor
📝 WalkthroughWalkthroughThe changes relocate GPU device visibility configuration from the Python submission script to the shell worker startup script, with Changes
Estimated code review effort🎯 2 (Simple) | ⏱️ ~10 minutes 🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
📝 Coding Plan
Comment Tip CodeRabbit can enforce grammar and style rules using `languagetool`.Configure the |
![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 the current code and only fix it if needed.
Inline comments:
In `@examples/disaggregated/slurm/benchmark/submit.py`:
- Line 520: The current code builds a single cuda_devices string from gpu_ids of
the first node (variable gpu_ids -> cuda_devices) which can misconfigure
multi-node, heterogeneous allocations; change the logic to compute
CUDA_VISIBLE_DEVICES per node using each node's local GPU id list (iterate over
the allocation's nodes or use the existing structure that contains per-node GPU
ids), and ensure wherever cuda_devices is injected into the job environment or
srun/ssh launch it uses the node-specific string (e.g., construct cuda_devices
inside the per-node launch loop or map keyed by node hostname) instead of
reusing the first-node value.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 998bd6fd-4329-4462-b881-069677676d08
📒 Files selected for processing (2)
examples/disaggregated/slurm/benchmark/start_worker.shexamples/disaggregated/slurm/benchmark/submit.py
dc3671
approved these changes
Mar 20, 2026
Collaborator
|
PR_Github #39667 [ skip ] triggered by Bot. Commit: |
Shixiaowei02
approved these changes
Mar 20, 2026
Collaborator
|
PR_Github #39667 [ skip ] completed with state |
longcheng-nv
pushed a commit
to longcheng-nv/TensorRT-LLM
that referenced
this pull request
Mar 31, 2026
…-export (NVIDIA#12370) Signed-off-by: Xianjie <[email protected]>
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
4 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
--exportcannot reliably pass comma-separated values (likeCUDA_VISIBLE_DEVICES=0,1) inside shared pyxis containers. The shared container retains the environment from its initial creation, ignoring subsequent--exportoverrides.0,1,2,3) instead of their assigned subset, leading to GPU OOM during model load when multiple servers share the same node.CUDA_VISIBLE_DEVICESas a command-line argument tostart_worker.shandexportit inside the script, bypassing srun--exportentirely for this variable.Test plan
🤖 Generated with Claude Code
Summary by CodeRabbit
Release Notes