Skip to content

fix parameters pass through for trtllm #509

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
slin1237 merged 2 commits into from
Feb 23, 2026
Merged

fix parameters pass through for trtllm #509

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
2ebedd3
fix parameters pass through for trtllm
gongwei-130 Feb 21, 2026
b866556
fix lint
gongwei-130 Feb 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 5 additions & 8 deletions bindings/python/src/smg/serve.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -363,13 +363,7 @@ def _add_trtllm_stub_args(parser: argparse.ArgumentParser) -> None:
"""
group = parser.add_argument_group("TensorRT-LLM Options")
group.add_argument("--model", type=str, help="Model path (HuggingFace ID or local path)")
group.add_argument("--tp-size", type=int, help="Tensor parallel size (overrides config file)")
group.add_argument(
"--config",
type=str,
required=False,
help="Config file path (YAML, optional - must contain tensor_parallel_size if provided)",
)
group.add_argument("--tp_size", type=int, help="Tensor parallel size (overrides config file)")

@chatgpt-codex-connector chatgpt-codex-connector Bot Feb 22, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Restore --tp-size flag spelling for TensorRT-LLM

Changing the stub flag from --tp-size to --tp_size breaks existing CLI usage and silently drops TP size in the orchestrator path: users passing --tp-size now leave args.tp_size unset, so TrtllmWorkerLauncher._get_tp_size() falls back to defaults while the raw token is only forwarded in backend_args. In multi-worker runs this mis-sizes CUDA_VISIBLE_DEVICES and can cause incorrect GPU allocation or worker startup failures.

Useful? React with 👍 / 👎.



BACKEND_ARG_ADDERS = {
Expand Down Expand Up @@ -482,7 +476,10 @@ def parse_serve_args(
_import_backend_args(backend, parser)
RouterArgs.add_cli_args(parser, use_router_prefix=True, exclude_host_port=True)

args = parser.parse_args(argv)
if backend == "trtllm":
args, _ = parser.parse_known_args(argv)
Comment on lines +479 to +480

@chatgpt-codex-connector chatgpt-codex-connector Bot Feb 22, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Keep --config in parsed args before TP sizing

Using parse_known_args for trtllm while no longer defining --config means config is discarded from args, but _get_tp_size() depends on args.config to read tensor parallel size from YAML. As a result, launches that specify TP only via --config now default to TP=1 for GPU env assignment, which can overlap GPUs across data-parallel workers and break startup.

Useful? React with 👍 / 👎.

@chatgpt-codex-connector chatgpt-codex-connector Bot Feb 23, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Reject unknown non-backend flags in trtllm parsing

Using parse_known_args here and dropping the unknown list makes typos in SMG-owned CLI flags silently bypass validation for trtllm. In that case the mistyped option falls back to defaults in args (for example, a misspelled worker port flag keeps worker_base_port=31000) and is forwarded as backend passthrough, which can break worker startup if TensorRT-LLM does not recognize it.

Useful? React with 👍 / 👎.

Comment on lines +479 to +480

@chatgpt-codex-connector chatgpt-codex-connector Bot Feb 23, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Preserve --config in parsed trtllm args

Using parse_known_args for the trtllm pass-2 parse here means options that are no longer declared (notably --config after this change) are silently dropped from args. TrtllmWorkerLauncher._get_tp_size() relies on args.config to read tensor-parallel size from YAML before computing CUDA_VISIBLE_DEVICES, so runs that set TP only via config now fall back to tp=1, which can assign overlapping GPUs and fail multi-worker startup.

Useful? React with 👍 / 👎.

else:
args = parser.parse_args(argv)
return backend, args, backend_args
Comment on lines +479 to 483

@coderabbitai coderabbitai Bot Feb 22, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick | 🔵 Trivial

Log the discarded tokens from parse_known_args to aid diagnosis

Unknown args in pass 2 are silently dropped into _. A user who misspells a serve-level flag (e.g., --tp-size instead of --tp_size) will get no feedback: pass 2 silently ignores it, args.tp_size stays None, and GPU assignment defaults to 1. Adding a debug log of the discarded tokens makes this easier to catch.

♻️ Proposed change 
     if backend == "trtllm":
-        args, _ = parser.parse_known_args(argv)
+        args, unknown = parser.parse_known_args(argv)
+        if unknown:
+            logger.debug(
+                "trtllm: ignoring unrecognized args in pass-2 parse (will be forwarded via backend_args): %s",
+                unknown,
+            )
     else:
         args = parser.parse_args(argv)
🤖 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 479 - 483, When backend ==
"trtllm" the call to parser.parse_known_args(argv) drops unknown tokens into the
discard variable `_` without feedback; update the branch handling
parse_known_args to detect if `_` is non-empty and emit a debug/warn log listing
those discarded tokens (include context like the backend and argv), so callers
see misspelled flags; keep existing behavior of returning backend, args,
backend_args and do not change the parse logic otherwise (refer to
parser.parse_known_args, the `_` discard variable, args, and the backend ==
"trtllm" branch).



Expand Down
9 changes: 6 additions & 3 deletions bindings/python/tests/test_serve.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -158,9 +158,12 @@ class TestImportBackendArgs:
def test_trtllm_adds_model_arg(self):
parser = argparse.ArgumentParser()
_import_backend_args("trtllm", parser)
args = parser.parse_args(["--model", "/path/to/model", "--config", "/path/to/config.yml"])
args, backend_args = parser.parse_known_args(
["--model", "/path/to/model", "--config", "/path/to/config.yml"]
)
assert args.model == "/path/to/model"
assert args.config == "/path/to/config.yml"
assert "--config" in backend_args
assert "/path/to/config.yml" in backend_args
Comment on lines 158 to +166

@coderabbitai coderabbitai Bot Feb 23, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick | 🔵 Trivial

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
# Find the serve.py file and examine the implementations

echo "=== Looking for serve.py ==="
fd -t f "serve.py" bindings/python

echo ""
echo "=== Finding _add_trtllm_stub_args ==="
rg -A 20 "def _add_trtllm_stub_args" bindings/python

echo ""
echo "=== Finding _get_tp_size ==="
rg -A 10 "def _get_tp_size" bindings/python

Repository: lightseekorg/smg

Length of output: 5081

🏁 Script executed:

#!/bin/bash
# Get the full TrtllmWorkerLauncher._get_tp_size implementation
echo "=== Full TrtllmWorkerLauncher._get_tp_size ==="
rg -A 30 "class TrtllmWorkerLauncher" bindings/python/src/smg/serve.py | grep -A 30 "_get_tp_size"

echo ""
echo "=== Check if args.config is referenced in TrtllmWorkerLauncher._get_tp_size ==="
sed -n '/class TrtllmWorkerLauncher/,/^class /p' bindings/python/src/smg/serve.py | grep -A 25 "def _get_tp_size" | head -30

Repository: lightseekorg/smg

Length of output: 2147

🏁 Script executed:

#!/bin/bash
# Search for tp_size coverage in tests
echo "=== Searching for tp_size in test file ==="
rg "tp_size" bindings/python/tests/test_serve.py

echo ""
echo "=== Full TestAddTrtllmStubArgs test ==="
sed -n '/class TestAddTrtllmStubArgs/,/^class /p' bindings/python/tests/test_serve.py | head -30

Repository: lightseekorg/smg

Length of output: 2537

Correctly reflects --config removal from stub; --tp_size stub argument lacks direct test coverage.

The switch to parse_known_args and the backend_args assertions accurately exercise the new behavior where --config is no longer a stub argument. However, _add_trtllm_stub_args does add --tp_size as a new argument, yet TestAddTrtllmStubArgs (lines 183-196) lacks a test for it—only testing --model. The _get_tp_size() method is thoroughly tested elsewhere, but the stub argument parser itself should verify that --tp_size can be parsed correctly, following the existing pattern of test_adds_model_arg.

Note: TrtllmWorkerLauncher._get_tp_size safely accesses args.config via getattr(args, "config", None), so the removal of --config from stub args does not cause breakage; it simply falls back to the default.

🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@bindings/python/tests/test_serve.py` around lines 158 - 166, Add a test that
verifies the trtllm stub parser accepts and returns the --tp_size flag: update
or add a test in TestAddTrtllmStubArgs that calls _add_trtllm_stub_args (or uses
_import_backend_args for "trtllm") to build an argparse.Parser, then call
parser.parse_known_args with ["--model","/path/to/model","--tp_size","4"] (or
similar) and assert args.model and that the parsed value or backend_args reflect
the tp_size as expected; reference _add_trtllm_stub_args, TestAddTrtllmStubArgs,
TrtllmWorkerLauncher._get_tp_size and parser.parse_known_args to locate the
relevant code to modify.


def test_sglang_import_error(self):
"""sglang is not installed in test env, so parser.error should be called."""
Expand Down Expand Up @@ -304,7 +307,7 @@ def test_two_pass_extracts_backend_first(self):
def test_unknown_arg_rejected_in_pass2(self):
"""Unknown args should be rejected by the full parser in pass 2."""
with pytest.raises(SystemExit):
parse_serve_args(["--backend", "trtllm", "--totally-unknown-flag"])
parse_serve_args(["--backend", "sglang", "--totally-unknown-flag"])
Comment on lines 307 to +310

@coderabbitai coderabbitai Bot Feb 23, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

test_unknown_arg_rejected_in_pass2 exits for the wrong reason and duplicates test_sglang_explicit_exits.

Using "sglang" as the backend causes _import_backend_args("sglang", parser) to call parser.error(...) (import failure) and raise SystemExit(2) before pass-2 parsing runs at all. The --totally-unknown-flag token is never evaluated, so the test does not verify what its docstring claims. It is also functionally identical to test_sglang_explicit_exits (lines 267-270), adding no incremental coverage.

The previous "trtllm" backend was also wrong post-PR: trtllm now uses parse_known_args, so the unknown flag would be silently swallowed, not rejected.

To genuinely test pass-2 unknown-arg rejection, mock _import_backend_args so a non-trtllm backend does not error on import:

🛠️ Suggested fix 
 def test_unknown_arg_rejected_in_pass2(self):
     """Unknown args should be rejected by the full parser in pass 2."""
-    with pytest.raises(SystemExit):
-        parse_serve_args(["--backend", "sglang", "--totally-unknown-flag"])
+    # Patch _import_backend_args so the backend loads successfully, letting
+    # pass-2 parse_args be the one to reject the unknown flag.
+    with patch("smg.serve._import_backend_args"):
+        with pytest.raises(SystemExit):
+            parse_serve_args(["--backend", "sglang", "--totally-unknown-flag"])
🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@bindings/python/tests/test_serve.py` around lines 307 - 310, The test
currently exits during backend import (via _import_backend_args) instead of
during pass-2 parsing; fix by stubbing/mocking _import_backend_args in
test_unknown_arg_rejected_in_pass2 so it does nothing (e.g., monkeypatch
_import_backend_args to a no-op or return None), call
parse_serve_args(["--backend", "sglang", "--totally-unknown-flag"]), and assert
it raises SystemExit from the full parser; keep the test name but ensure it no
longer duplicates test_sglang_explicit_exits by verifying the unknown-flag
rejection rather than import failure.



# ---------------------------------------------------------------------------
Expand Down