feat: add gb200 and sglang commands to slurm scripts #1804
Conversation
ishandhanani
requested review from
biswapanda,
tedzhouhk,
rmccorm4,
alec-flowers and
grahamking
as code owners
ishandhanani
changed the title
WalkthroughThe script for SLURM job worker setup was updated to support an alternative SGLang-based execution mode, toggled by a new Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant WorkerSetupScript
participant SGLangServer
participant DynamoWorker
participant NatsEtcd
participant SGLangLB
User->>WorkerSetupScript: Run with/without --use-sglang-commands
alt use_sglang_commands = true
WorkerSetupScript->>SGLangLB: (decode node only) Launch mini load balancer
WorkerSetupScript->>SGLangServer: Launch SGLang server with arguments
else use_sglang_commands = false
WorkerSetupScript->>NatsEtcd: (prefill node, rank 0) Start nats-server and etcd
WorkerSetupScript->>DynamoWorker: Launch Dynamo worker with arguments
end
Possibly related PRs
Poem
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
Documentation and Community
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
examples/sglang/slurm_jobs/scripts/worker_setup.py(9 hunks)
🧰 Additional context used
🧠 Learnings (2)
📓 Common learnings
Learnt from: fsaady
PR: ai-dynamo/dynamo#1730
File: examples/sglang/slurm_jobs/scripts/worker_setup.py:230-244
Timestamp: 2025-07-03T10:14:30.538Z
Learning: In examples/sglang/slurm_jobs/scripts/worker_setup.py, background processes (like nats-server, etcd) are intentionally left running even if later processes fail. This design choice allows users to manually connect to nodes and debug issues without having to restart the entire SLURM job from scratch, providing operational flexibility for troubleshooting in cluster environments.
Learnt from: fsaady
PR: ai-dynamo/dynamo#1730
File: examples/sglang/slurm_jobs/job_script_template.j2:59-59
Timestamp: 2025-07-02T13:20:28.800Z
Learning: In the SLURM job script template at examples/sglang/slurm_jobs/job_script_template.j2, the `--total_nodes` parameter represents the total nodes per worker type (prefill or decode), not the total nodes in the entire cluster. Each worker type needs to know its own group size for distributed coordination.
Learnt from: fsaady
PR: ai-dynamo/dynamo#1730
File: examples/sglang/slurm_jobs/scripts/worker_setup.py:113-116
Timestamp: 2025-07-03T09:44:41.140Z
Learning: In examples/sglang/slurm_jobs/scripts/worker_setup.py, logging the full configuration file content is acceptable because the config file is public, contains only placeholder replacements (no sensitive data), and provides debugging benefits for users who may want to create configurations based on the logged output.
examples/sglang/slurm_jobs/scripts/worker_setup.py (4)
Learnt from: fsaady
PR: ai-dynamo/dynamo#1730
File: examples/sglang/slurm_jobs/scripts/worker_setup.py:230-244
Timestamp: 2025-07-03T10:14:30.538Z
Learning: In examples/sglang/slurm_jobs/scripts/worker_setup.py, background processes (like nats-server, etcd) are intentionally left running even if later processes fail. This design choice allows users to manually connect to nodes and debug issues without having to restart the entire SLURM job from scratch, providing operational flexibility for troubleshooting in cluster environments.
Learnt from: fsaady
PR: ai-dynamo/dynamo#1730
File: examples/sglang/slurm_jobs/job_script_template.j2:59-59
Timestamp: 2025-07-02T13:20:28.800Z
Learning: In the SLURM job script template at examples/sglang/slurm_jobs/job_script_template.j2, the `--total_nodes` parameter represents the total nodes per worker type (prefill or decode), not the total nodes in the entire cluster. Each worker type needs to know its own group size for distributed coordination.
Learnt from: fsaady
PR: ai-dynamo/dynamo#1730
File: examples/sglang/slurm_jobs/scripts/worker_setup.py:113-116
Timestamp: 2025-07-03T09:44:41.140Z
Learning: In examples/sglang/slurm_jobs/scripts/worker_setup.py, logging the full configuration file content is acceptable because the config file is public, contains only placeholder replacements (no sensitive data), and provides debugging benefits for users who may want to create configurations based on the logged output.
Learnt from: nnshah1
PR: ai-dynamo/dynamo#1444
File: tests/fault_tolerance/configs/agg_tp_1_dp_8.yaml:31-38
Timestamp: 2025-07-01T15:33:53.262Z
Learning: In fault tolerance test configurations, the `resources` section under `ServiceArgs` specifies resources per individual worker, not total resources for all workers. So `workers: 8` with `gpu: '1'` means 8 workers × 1 GPU each = 8 GPUs total.
🪛 GitHub Actions: Pre Merge Validation of (ai-dynamo/dynamo/refs/pull/1804/merge) by ishandhanani.
examples/sglang/slurm_jobs/scripts/worker_setup.py
[error] 199-400: Black formatting check failed. The file was reformatted by the black hook. Run 'black --write examples/sglang/slurm_jobs/scripts/worker_setup.py' to fix code style issues.
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Build and Test - vllm
🔇 Additional comments (5)
examples/sglang/slurm_jobs/scripts/worker_setup.py (5)
180-184: The command line argument implementation looks good.The new
--use-sglang-commandsflag is properly defined with clear help text and follows standard argparse conventions.
290-305: Background process handling follows established patterns.The SGLang mini load balancer is correctly started in the background using the existing
run_commandfunction, which is consistent with how other background processes are managed in this codebase.
244-274: Command construction logic is well-structured.The dynamic command construction based on the
use_sglang_commandsflag is clean and maintains consistency between the two execution paths. The reuse of command arguments is efficient and reduces code duplication.Also applies to: 306-335
202-203: Function signature updates are correct.The addition of the
use_sglang_commandsparameter to both setup functions follows the existing patterns and maintains backward compatibility through the main function's argument parsing.Also applies to: 283-284
362-362: Parameter passing and logging are properly implemented.The new flag is correctly passed to both setup functions and appropriately logged for debugging purposes.
Also applies to: 371-371, 380-380
ishandhanani
changed the title
ishandhanani
requested review from
piotrm-nvidia,
ryanolson,
paulhendricks,
tmonty12,
GuanLuo,
kkranen and
jthomson04
as code owners
This PR adds in the h100 and gb200 large scale disagg + wide ep SGLang commands detailed in sgl-project/sglang#6017 and sgl-project/sglang#7227. The goal here is to be able to use the same slurm script to collect benchmarking numbers on sglang and dynamo
Summary by CodeRabbit
--use-sglang-commands) to enable the SGLang-based mode.