Repository navigation
Conversation
github-actions
Bot
added
python-bindings
📝 WalkthroughWalkthroughCentralized tensor-parallel (TP) size resolution: a shared resolver normalizes TP CLI/YAML inputs, stores a canonical TP on parsed args, and WorkerLauncher.gpu_env() uses this value across SGLang, vLLM, and TRT-LLM backends. Changes
Sequence Diagram(s)sequenceDiagram
participant CLI
participant Parser
participant Resolver
participant Launcher
participant Backend
CLI->>Parser: provide args (aliases: --tp_size, --tp-size, --tensor-parallel-size)
Parser->>Resolver: request canonical TP (check namespace & config)
Resolver->>Parser: store `_smg_worker_tp_size` on namespace
Launcher->>Parser: read `_smg_worker_tp_size` (preferred)
alt `_smg_worker_tp_size` missing
Launcher->>Backend: call backend._get_tp_size() (may read config/args)
Backend->>Resolver: (backend-specific) resolver returns TP
Resolver->>Launcher: TP value
end
Launcher->>Launcher: compute CUDA_VISIBLE_DEVICES slice using TP and dp_rank
Launcher->>Backend: launch worker with computed GPU env
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~35 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
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. Comment |
|
Hi @smfirmin, the DCO sign-off check has failed. All commits must include a To fix existing commits: # Sign off the last N commits (replace N with the number of unsigned commits)
git rebase HEAD~N --signoff
git push --force-with-leaseTo sign off future commits automatically:
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request centralizes Tensor Parallel (TP) size resolution across SGLang, vLLM, and TRT-LLM backends by introducing normalized resolution helpers and updating worker launchers. Feedback highlights the omission of the --config argument in the TRT-LLM parser and suggests implementing explicit encoding and type validation when parsing YAML configuration files.
3246ba2 to
5447bd4
Compare
|
@gongwei-130 can you take a look at this PR |
1808a50 to
6609400
Compare
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 `@bindings/python/src/smg/serve.py`:
- Around line 73-74: Open the YAML config file with an explicit encoding (e.g.,
use open(config_path, encoding='utf-8')) to avoid platform-dependent behavior,
and after yaml.safe_load(f) validate that the returned config is a mapping/dict
(the variable config) — if it is None or not a dict, raise a clear error
indicating the config is empty or malformed; update the code paths that use
config to rely on this validated dict (references: the open(config_path) call
and the config variable resulting from yaml.safe_load).
🪄 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: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 1cfc4316-e5d5-43e1-b3a8-49daa75827cf
📒 Files selected for processing (2)
bindings/python/src/smg/serve.pybindings/python/tests/test_serve.py
Signed-off-by: Sydney Firmin <sydney.firmin@oracle.com>
582d79a to
f6ec618
Compare
|
This pull request has been automatically marked as stale because it has not had any activity within 14 days. It will be automatically closed if no further activity occurs within 16 days. Leave a comment if you feel this pull request should remain open. Thank you! |
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 `@bindings/python/src/smg/serve.py`:
- Around line 483-490: parse_serve_args() is caching _smg_worker_tp_size before
args.config is available because the stub parser never parses --config; update
the flow so the TRT-LLM config is parsed into the same namespace before caching
TP (either by declaring --config on the stub parser or by explicitly
loading/parsing the config YAML into args.config prior to computing
_smg_worker_tp_size), ensuring _resolve_trtllm_tp_size() sees args.config and
gpu_env() gets the correct tensor_parallel_size (tensor_parallel_size,
_smg_worker_tp_size, backend_args, args.config, _resolve_trtllm_tp_size,
gpu_env()).
🪄 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: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: b152ff33-ec74-4a61-988c-99bca8e4ae75
📒 Files selected for processing (2)
bindings/python/src/smg/serve.pybindings/python/tests/test_serve.py
| group.add_argument( | ||
| "--tp_size", | ||
| "--tp-size", | ||
| "--tensor-parallel-size", | ||
| dest="tensor_parallel_size", | ||
| type=int, | ||
| help="Tensor parallel size (overrides config file)", | ||
| ) |
There was a problem hiding this comment.
Parse --config into the TRT-LLM namespace before caching TP.
parse_serve_args() now computes _smg_worker_tp_size from args, but --config is still left in backend_args because this stub parser never declares it. On a config-only TRT-LLM launch, _resolve_trtllm_tp_size() never sees args.config, caches 1, and gpu_env() slices GPUs for TP=1 even though the worker can still start with tensor_parallel_size > 1 from YAML. That recreates the same invalid-device-ordinal failure this PR is trying to eliminate.
🐛 Minimal fix
def _add_trtllm_stub_args(parser: argparse.ArgumentParser) -> None:
"""Add TensorRT-LLM specific arguments.
@@
group = parser.add_argument_group("TensorRT-LLM Options")
+ group.add_argument(
+ "--config",
+ type=str,
+ help="Path to TRT-LLM YAML config",
+ )
group.add_argument(
"--model",
"--model-path",
dest="model_path",
type=str,🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@bindings/python/src/smg/serve.py` around lines 483 - 490, parse_serve_args()
is caching _smg_worker_tp_size before args.config is available because the stub
parser never parses --config; update the flow so the TRT-LLM config is parsed
into the same namespace before caching TP (either by declaring --config on the
stub parser or by explicitly loading/parsing the config YAML into args.config
prior to computing _smg_worker_tp_size), ensuring _resolve_trtllm_tp_size() sees
args.config and gpu_env() gets the correct tensor_parallel_size
(tensor_parallel_size, _smg_worker_tp_size, backend_args, args.config,
_resolve_trtllm_tp_size, gpu_env()).
|
This pull request has been automatically marked as stale because it has not had any activity within 14 days. It will be automatically closed if no further activity occurs within 16 days. Leave a comment if you feel this pull request should remain open. Thank you! |
|
@gongwei-130 have you had a chance to review this? |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 2f114cca4a
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| "--tp_size", | ||
| "--tp-size", | ||
| "--tensor-parallel-size", | ||
| dest="tensor_parallel_size", |
There was a problem hiding this comment.
For TRT-LLM, NVIDIA's trtllm-serve help documents the CLI spelling as --tensor_parallel_size, --tp_size <tensor_parallel_size>. If a user passes that native underscore form through smg serve --backend trtllm --tensor_parallel_size 4, parse_known_args leaves it in backend_args so the worker still launches with TP=4, but this stub never parses it and _resolve_worker_tp_size falls back to 1 for CUDA_VISIBLE_DEVICES; with DP workers this recreates the invalid GPU slicing this change is trying to prevent. Add --tensor_parallel_size to this alias list as well.
Useful? React with 👍 / 👎.
|
This pull request has been automatically marked as stale because it has not had any activity within 14 days. It will be automatically closed if no further activity occurs within 16 days. Leave a comment if you feel this pull request should remain open. Thank you! |
|
This pull request has been automatically closed due to inactivity. Please feel free to reopen if you intend to continue working on it. Thank you! |
3 participants
Description
Problem
smg servecomputes workerCUDA_VISIBLE_DEVICESbefore launching backend workers. For SGLang and vLLM, that computation depended on backend-specific parsed argument names.In practice, SGLang accepts
--tensor-parallel-size/--tp-size, but its TP value may be normalized differently from the raw CLI spelling. This causes SMG to compute worker GPU visibility as iftp=1even when the worker itself starts withtp=4, leading toCUDA error: invalid device ordinalduring worker startup.There was a related mismatch risk for TRT-LLM as well: SMG could normalize hyphenated TP flags for the launched worker command, while still failing to use the same TP value for its own GPU assignment.
Solution
Normalize tensor parallel size once during
parse_serve_args()into a single SMG-internal value used for worker GPU assignment.This PR:
smg serveworker launchgpu_env()use that canonical value instead of inferring TP from backend parser internalsChanges
bindings/python/src/smg/serve.pyparse_serve_args()gpu_env()to use the canonical normalized TP value0so existing validation still fails fast--tp_size,--tp-size, and--tensor-parallel-sizefor SMG-side GPU assignmentTest Plan
Repro before fix:
smg servewith SGLang and--tensor-parallel-size 4or--tp-size 4.CUDA_VISIBLE_DEVICES=0.tp_size=4and then fail withCUDA error: invalid device ordinal.Expected behavior after fix:
smg servewith SGLang and--tensor-parallel-size 4.4before worker launch.gpu_env()should assign:dp_rank=0 -> CUDA_VISIBLE_DEVICES=0,1,2,3dp_rank=1 -> CUDA_VISIBLE_DEVICES=4,5,6,7Targeted test run used for this change:
Result:
30 passed, 105 deselectedNote:
smg_rsextension mismatch during import, so the targetedservetest subset was run with a temporary import stub forsmg.smg_rs.Checklist
cargo +nightly fmtpassescargo clippy --all-targets --all-features -- -D warningspassesSummary by CodeRabbit
Improvements
Tests