Repository navigation
fix: unify --model/--model-path and revert in-gateway served-model-name - #537
Conversation
github-actions
Bot
added
python-bindings
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughRemoved router-level served_model_name and its propagation across Python bindings and Rust config; simplified worker identity to a single canonical model_id, consolidated tokenizer registration to one per worker, and added CLI aliasing/normalization for model vs model-path. Changes
Sequence Diagram(s)sequenceDiagram
participant CLI as CLI (args parser)
participant Router as Router
participant Registry as Registry
participant Worker as Worker
participant Tokenizer as Tokenizer
CLI->>Router: start (model_path / aliases parsed)
Router->>Registry: register worker (canonical model_id)
Registry->>Worker: create worker entry (with model_id)
Worker->>Tokenizer: submit AddTokenizer(job for model_id)
Registry->>Router: get_models() (for validation)
Router->>CLI: accept requests if model_id ∈ get_models()
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related issues
Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 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)
Comment |
Summary of ChangesHello @slin1237, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request focuses on unifying command-line argument handling for model paths across different backends, particularly addressing inconsistencies between trtllm and sglang. It also reverts a previously introduced feature related to served model names, choosing to implement it directly in the upstream TensorRT-LLM project. The changes ensure consistent behavior and simplify configuration across the model gateway. Highlights
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
|
Hi @slin1237, 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}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@model_gateway/src/core/worker_registry.rs`:
- Around line 332-344: The code stores an empty slice back into self.model_index
when removing the last worker, leaving a tombstone key that causes
rebuild_hash_ring(&model_id) to see an empty ring; change the removal logic in
the block that builds new_workers (using worker.model_id() and worker_url) to
check if new_workers.is_empty() and if so call
self.model_index.remove(&model_id) instead of writing an empty slice, otherwise
replace the entry as before, so rebuild_hash_ring sees no entry for that model
when its last worker is removed.
- Around line 264-273: The current model_index update always appends the worker
to the existing snapshot, which can create duplicate entries for the same worker
URL; change the and_modify closure used on model_index.entry(model_id.clone())
to rebuild new_workers by collecting existing workers but excluding any whose
URL equals worker.url() (or the appropriate unique address getter on the Worker
trait) before pushing the new worker, then replace existing with the deduped Arc
slice; keep the or_insert_with path unchanged.
ℹ️ Review info
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
📒 Files selected for processing (11)
bindings/python/src/lib.rsbindings/python/src/smg/router_args.pybindings/python/src/smg/serve.pymodel_gateway/src/config/builder.rsmodel_gateway/src/config/types.rsmodel_gateway/src/core/steps/worker/local/create_worker.rsmodel_gateway/src/core/steps/worker/local/submit_tokenizer_job.rsmodel_gateway/src/core/worker.rsmodel_gateway/src/core/worker_registry.rsmodel_gateway/src/main.rsmodel_gateway/src/routers/grpc/common/responses/utils.rs
💤 Files with no reviewable changes (4)
- model_gateway/src/config/types.rs
- bindings/python/src/lib.rs
- model_gateway/src/config/builder.rs
- model_gateway/src/core/worker.rs
![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 removes the served_model_name option from the router configuration and related code. The served_model_name allowed overriding the model name exposed to clients, but this functionality is being removed. The changes affect the Python bindings, the model gateway's configuration, worker creation, tokenizer job submission, worker registry, command-line arguments, and gRPC response handling.
![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: 14c0efae65
ℹ️ 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".
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 225-230: The normalization currently replaces hyphens in the
entire token which turns "--model-path" into "__model_path" and breaks matching;
update the logic where `normalized` is computed (in serve.py near
`backend_args`, `cmd.extend`, and `_filter_backend_args` usage) to only replace
hyphens in the flag name portion after the leading "--" (e.g., for items
starting with "--" split into prefix and name, transform name =
name.replace("-", "_"), then rejoin as prefix + name) while leaving non-flag
args unchanged so `_filter_backend_args` correctly matches ["--model",
"--model_path", ...].
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e4277f3c45
ℹ️ 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".
c0c37d1 to
19743db
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 19743dba2f
ℹ️ 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".
There was a problem hiding this comment.
♻️ Duplicate comments (3)
bindings/python/src/smg/serve.py (1)
237-240:⚠️ Potential issue | 🔴 CriticalFix flag normalization: it currently breaks
--prefixes.Line 237 transforms
--model-pathinto__model_path, so Line 239 filtering misses expected keys and invalid flags leak into TRT-LLM arguments.🐛 Proposed fix
- normalized = [a.replace("-", "_") if a.startswith("--") else a for a in backend_args] + normalized = ["--" + a[2:].replace("-", "_") if a.startswith("--") else a for a in backend_args]#!/bin/bash python - <<'PY' args = ["--model-path", "/tmp/model", "--tensor-parallel-size=2"] bad = [a.replace("-", "_") if a.startswith("--") else a for a in args] good = ["--" + a[2:].replace("-", "_") if a.startswith("--") else a for a in args] print("bad :", bad) print("good:", good) PY🤖 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 237 - 240, The normalization step for backend_args incorrectly replaces the leading "--" (producing "__model_path") so _filter_backend_args misses flags; change the list comprehension that produces normalized (used before cmd.extend and _filter_backend_args) to preserve the "--" prefix and only replace internal hyphens with underscores (i.e., transform "--model-path" -> "--model_path"), then pass that normalized list into self._filter_backend_args(["--model", "--model_path", "--host", "--port"]) so the expected keys match and no invalid flags leak into TRT-LLM arguments.model_gateway/src/core/worker_registry.rs (2)
264-273:⚠️ Potential issue | 🟠 MajorDeduplicate by worker URL when updating
model_indexsnapshots.Line 269 currently clones all existing entries and Line 270 blindly pushes the worker again, so re-registering the same URL creates duplicates in the model bucket.
🔧 Proposed fix
self.model_index .entry(model_id.clone()) .and_modify(|existing| { - // Create new snapshot with the additional worker - let mut new_workers: Vec<Arc<dyn Worker>> = existing.iter().cloned().collect(); + // Rebuild snapshot without stale copy for this URL, then append latest worker + let mut new_workers: Vec<Arc<dyn Worker>> = existing + .iter() + .filter(|w| w.url() != worker.url()) + .cloned() + .collect(); new_workers.push(worker.clone()); *existing = Arc::from(new_workers.into_boxed_slice()); }) .or_insert_with(|| Arc::from(vec![worker.clone()].into_boxed_slice()));🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/core/worker_registry.rs` around lines 264 - 273, When updating self.model_index for a model_id, avoid pushing a worker with the same URL into the snapshot: instead of blindly cloning existing and pushing worker.clone() in the and_modify closure, deduplicate by comparing each existing worker's URL (use the Worker URL accessor/method on the Arc<dyn Worker>) and only include existing entries whose URL != worker.url(), then append the new worker; update the and_modify block that references existing and worker.clone() to build new_workers by filtering out same-URL entries before pushing the worker, so registering the same URL won't create duplicates.
332-343:⚠️ Potential issue | 🟠 MajorRemove model key when its last worker is deleted.
Line 339 stores an empty slice instead of deleting the key, leaving tombstone model entries and stale ring bookkeeping.
🔧 Proposed fix
let model_id = worker.model_id().to_string(); if let Some(mut entry) = self.model_index.get_mut(&model_id) { let new_workers: Vec<Arc<dyn Worker>> = entry .iter() .filter(|w| w.url() != worker_url) .cloned() .collect(); - *entry = Arc::from(new_workers.into_boxed_slice()); + if new_workers.is_empty() { + drop(entry); + self.model_index.remove(&model_id); + } else { + *entry = Arc::from(new_workers.into_boxed_slice()); + } }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/core/worker_registry.rs` around lines 332 - 343, The code in the block that updates self.model_index for model_id currently stores an empty slice when the last worker is removed, leaving a tombstone entry; change the logic in the removal path inside the model_index update so that if the filtered new_workers Vec is empty you call self.model_index.remove(&model_id) to delete the key instead of assigning an empty Arc slice, otherwise assign the new Arc as before; keep the subsequent call to self.rebuild_hash_ring(&model_id) as-is so the ring is rebuilt for that model_id.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Duplicate comments:
In `@bindings/python/src/smg/serve.py`:
- Around line 237-240: The normalization step for backend_args incorrectly
replaces the leading "--" (producing "__model_path") so _filter_backend_args
misses flags; change the list comprehension that produces normalized (used
before cmd.extend and _filter_backend_args) to preserve the "--" prefix and only
replace internal hyphens with underscores (i.e., transform "--model-path" ->
"--model_path"), then pass that normalized list into
self._filter_backend_args(["--model", "--model_path", "--host", "--port"]) so
the expected keys match and no invalid flags leak into TRT-LLM arguments.
In `@model_gateway/src/core/worker_registry.rs`:
- Around line 264-273: When updating self.model_index for a model_id, avoid
pushing a worker with the same URL into the snapshot: instead of blindly cloning
existing and pushing worker.clone() in the and_modify closure, deduplicate by
comparing each existing worker's URL (use the Worker URL accessor/method on the
Arc<dyn Worker>) and only include existing entries whose URL != worker.url(),
then append the new worker; update the and_modify block that references existing
and worker.clone() to build new_workers by filtering out same-URL entries before
pushing the worker, so registering the same URL won't create duplicates.
- Around line 332-343: The code in the block that updates self.model_index for
model_id currently stores an empty slice when the last worker is removed,
leaving a tombstone entry; change the logic in the removal path inside the
model_index update so that if the filtered new_workers Vec is empty you call
self.model_index.remove(&model_id) to delete the key instead of assigning an
empty Arc slice, otherwise assign the new Arc as before; keep the subsequent
call to self.rebuild_hash_ring(&model_id) as-is so the ring is rebuilt for that
model_id.
ℹ️ Review info
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
📒 Files selected for processing (11)
bindings/python/src/lib.rsbindings/python/src/smg/router_args.pybindings/python/src/smg/serve.pymodel_gateway/src/config/builder.rsmodel_gateway/src/config/types.rsmodel_gateway/src/core/steps/worker/local/create_worker.rsmodel_gateway/src/core/steps/worker/local/submit_tokenizer_job.rsmodel_gateway/src/core/worker.rsmodel_gateway/src/core/worker_registry.rsmodel_gateway/src/main.rsmodel_gateway/src/routers/grpc/common/responses/utils.rs
💤 Files with no reviewable changes (4)
- model_gateway/src/core/worker.rs
- bindings/python/src/lib.rs
- model_gateway/src/config/builder.rs
- model_gateway/src/config/types.rs
19743db to
b48de3c
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b48de3c60f
ℹ️ 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".
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
bindings/python/tests/test_serve.py (1)
536-543:⚠️ Potential issue | 🟡 MinorThis vLLM test no longer validates model propagation.
VllmWorkerLauncher.build_command()readsargs.model, but this test now sets onlymodel_path, so it can pass even if--modelis empty.🔧 Suggested fix
- args = argparse.Namespace(model_path="/tmp/model", connection_mode="http") + args = argparse.Namespace(model="/tmp/model", connection_mode="http") backend_args = ["--trust-remote-code"] cmd = launcher.build_command(args, backend_args, "127.0.0.1", 31000) assert "vllm.entrypoints.openai.api_server" in cmd + assert "--model" in cmd + assert "/tmp/model" in cmd🤖 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 536 - 543, The test sets argparse.Namespace(model_path=...) but VllmWorkerLauncher.build_command() reads args.model, so the test can pass without verifying model propagation; update the test to set args.model (e.g., args.model="/tmp/model") or ensure the Namespace includes both model and model_path, then add an assertion that the produced cmd contains the "--model" flag and the expected model path to validate model propagation (reference VllmWorkerLauncher.build_command and the test in test_serve.py).
♻️ Duplicate comments (2)
model_gateway/src/core/worker_registry.rs (2)
332-344:⚠️ Potential issue | 🟠 MajorRemove
model_indexkey when the last worker is removed.Writing back an empty slice leaves a tombstone model bucket and an empty hash-ring entry for that model.
🔧 Suggested fix
let model_id = worker.model_id().to_string(); if let Some(mut entry) = self.model_index.get_mut(&model_id) { let new_workers: Vec<Arc<dyn Worker>> = entry .iter() .filter(|w| w.url() != worker_url) .cloned() .collect(); - *entry = Arc::from(new_workers.into_boxed_slice()); + if new_workers.is_empty() { + drop(entry); + self.model_index.remove(&model_id); + } else { + *entry = Arc::from(new_workers.into_boxed_slice()); + } }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/core/worker_registry.rs` around lines 332 - 344, When removing a worker from self.model_index in the code around model_index.get_mut, detect if filtering out worker_url results in zero workers and in that case remove the model_id key from self.model_index instead of writing back an empty slice; otherwise update the entry as you do now. After that, still call self.rebuild_hash_ring(&model_id) so the hash ring reflects the removal. This change affects the block that builds new_workers from entry.iter().filter(...).cloned().collect() and the subsequent assignment to *entry.
264-273:⚠️ Potential issue | 🟠 MajorDeduplicate model snapshot entries by worker URL during register updates.
Re-registering the same worker URL currently appends another copy into the model bucket, which can skew routing/load behavior.
🔧 Suggested fix
self.model_index .entry(model_id.clone()) .and_modify(|existing| { - // Create new snapshot with the additional worker - let mut new_workers: Vec<Arc<dyn Worker>> = existing.iter().cloned().collect(); + // Rebuild without stale copy for this URL, then append latest worker + let mut new_workers: Vec<Arc<dyn Worker>> = existing + .iter() + .filter(|w| w.url() != worker.url()) + .cloned() + .collect(); new_workers.push(worker.clone()); *existing = Arc::from(new_workers.into_boxed_slice()); }) .or_insert_with(|| Arc::from(vec![worker.clone()].into_boxed_slice()));🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/core/worker_registry.rs` around lines 264 - 273, When updating model_index for model_id, avoid appending duplicate workers by URL: in the and_modify closure (where existing is converted to new_workers), check for an existing worker whose url() (or the appropriate identifier method on dyn Worker) equals worker.url(); if found, either skip pushing or replace that entry with worker.clone(), otherwise push; ensure the or_insert_with path still inserts a single worker. Modify the closure around new_workers.push(...) to filter/compare by the worker URL and only add when not present.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Outside diff comments:
In `@bindings/python/tests/test_serve.py`:
- Around line 536-543: The test sets argparse.Namespace(model_path=...) but
VllmWorkerLauncher.build_command() reads args.model, so the test can pass
without verifying model propagation; update the test to set args.model (e.g.,
args.model="/tmp/model") or ensure the Namespace includes both model and
model_path, then add an assertion that the produced cmd contains the "--model"
flag and the expected model path to validate model propagation (reference
VllmWorkerLauncher.build_command and the test in test_serve.py).
---
Duplicate comments:
In `@model_gateway/src/core/worker_registry.rs`:
- Around line 332-344: When removing a worker from self.model_index in the code
around model_index.get_mut, detect if filtering out worker_url results in zero
workers and in that case remove the model_id key from self.model_index instead
of writing back an empty slice; otherwise update the entry as you do now. After
that, still call self.rebuild_hash_ring(&model_id) so the hash ring reflects the
removal. This change affects the block that builds new_workers from
entry.iter().filter(...).cloned().collect() and the subsequent assignment to
*entry.
- Around line 264-273: When updating model_index for model_id, avoid appending
duplicate workers by URL: in the and_modify closure (where existing is converted
to new_workers), check for an existing worker whose url() (or the appropriate
identifier method on dyn Worker) equals worker.url(); if found, either skip
pushing or replace that entry with worker.clone(), otherwise push; ensure the
or_insert_with path still inserts a single worker. Modify the closure around
new_workers.push(...) to filter/compare by the worker URL and only add when not
present.
ℹ️ Review info
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
📒 Files selected for processing (12)
bindings/python/src/lib.rsbindings/python/src/smg/router_args.pybindings/python/src/smg/serve.pybindings/python/tests/test_serve.pymodel_gateway/src/config/builder.rsmodel_gateway/src/config/types.rsmodel_gateway/src/core/steps/worker/local/create_worker.rsmodel_gateway/src/core/steps/worker/local/submit_tokenizer_job.rsmodel_gateway/src/core/worker.rsmodel_gateway/src/core/worker_registry.rsmodel_gateway/src/main.rsmodel_gateway/src/routers/grpc/common/responses/utils.rs
💤 Files with no reviewable changes (4)
- model_gateway/src/config/builder.rs
- bindings/python/src/lib.rs
- model_gateway/src/config/types.rs
- model_gateway/src/core/worker.rs
This reverts commit ae1f82e. The --served-model-name feature should be implemented upstream in TensorRT-LLM rather than worked around in the gateway. A proper upstream implementation allows the served name to flow through both gRPC and HTTP server paths natively. Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
…istration Three related fixes for TensorRT-LLM backend support: 1. Unify --model and --model-path: trtllm uses --model while sglang uses --model-path, but both refer to the same thing. Map trtllm's --model to dest=model_path so it populates the same namespace key as sglang, allowing RouterArgs.from_cli_args to find it via the standard fallback. 2. Fix from_cli_args prefix fallback: when use_router_prefix=True, the prefixed key (e.g. router_model_path) always exists in the namespace with default=None, which blocked the fallback to the unprefixed key (model_path) even when the user didn't pass --router-model-path. Now only prefer the prefixed version when it is explicitly set. 3. Add --model as alias for --model-path in the Rust CLI (smg launch) and Python RouterArgs for consistency across backends. Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
TRT-LLM Click options use underscores (e.g. --tensor_parallel_size) while SGLang/vLLM use hyphens. Normalize backend args passed to TRT-LLM so users can use either form consistently across backends. Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
b48de3c to
74606ee
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (2)
model_gateway/src/core/worker_registry.rs (2)
332-343:⚠️ Potential issue | 🟠 MajorRemove model-index key when its last worker is removed.
Line 333-Line 340 stores an empty slice in
model_indexfor the model, leaving a tombstone key and allowing an empty hash ring to persist.Proposed fix
let model_id = worker.model_id().to_string(); if let Some(mut entry) = self.model_index.get_mut(&model_id) { let new_workers: Vec<Arc<dyn Worker>> = entry .iter() .filter(|w| w.url() != worker_url) .cloned() .collect(); - *entry = Arc::from(new_workers.into_boxed_slice()); + if new_workers.is_empty() { + drop(entry); + self.model_index.remove(&model_id); + } else { + *entry = Arc::from(new_workers.into_boxed_slice()); + } } // Rebuild hash ring for this model self.rebuild_hash_ring(&model_id);🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/core/worker_registry.rs` around lines 332 - 343, The current removal logic for a worker leaves an empty slice in self.model_index for a model, creating a tombstone; change the branch that handles the filtered new_workers in the block that mutates self.model_index (where model_id is derived from worker.model_id() and worker_url is used in the filter) so that if new_workers.is_empty() you call self.model_index.remove(&model_id) to delete the key entirely, otherwise assign *entry = Arc::from(new_workers.into_boxed_slice()) as before; keep the subsequent call to self.rebuild_hash_ring(&model_id) so the hash ring is rebuilt for removed keys as well.
264-273:⚠️ Potential issue | 🟠 MajorPrevent duplicate entries when re-registering an existing worker URL.
Line 264-Line 271 rebuilds
new_workersby cloning all existing entries and then pushing the new worker, so the same URL can appear multiple times after re-registration.Proposed fix
let model_id = worker.model_id().to_string(); self.model_index .entry(model_id.clone()) .and_modify(|existing| { - // Create new snapshot with the additional worker - let mut new_workers: Vec<Arc<dyn Worker>> = existing.iter().cloned().collect(); + // Rebuild without stale copy for this URL, then append latest worker + let mut new_workers: Vec<Arc<dyn Worker>> = existing + .iter() + .filter(|w| w.url() != worker.url()) + .cloned() + .collect(); new_workers.push(worker.clone()); *existing = Arc::from(new_workers.into_boxed_slice()); }) .or_insert_with(|| Arc::from(vec![worker.clone()].into_boxed_slice()));🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/core/worker_registry.rs` around lines 264 - 273, The current model_index update always appends the worker clone which can create duplicates; inside the entry(...).and_modify closure (the block that builds new_workers from existing and pushes worker.clone()) first check whether an equivalent worker already exists by comparing the worker URL/identifier (e.g. call the worker URL accessor such as worker.url() or worker.address() against each entry in existing) and only push the new worker if no match is found; update the logic in the and_modify closure to dedupe before creating the new Arc slice so re-registering the same worker URL will not add a duplicate entry.
🤖 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/tests/test_serve.py`:
- Line 536: The test instantiates args with the wrong key: it sets model_path
but VllmWorkerLauncher.build_command() reads args.model, so update the test to
pass the expected attribute (e.g., args = argparse.Namespace(model="/tmp/model",
connection_mode="http")) and ensure the assertion that validates model
propagation into VllmWorkerLauncher.build_command() still runs; look for the
test function around the args assignment and the
VllmWorkerLauncher.build_command() call to change the Namespace key and preserve
the existing propagation assertion.
---
Duplicate comments:
In `@model_gateway/src/core/worker_registry.rs`:
- Around line 332-343: The current removal logic for a worker leaves an empty
slice in self.model_index for a model, creating a tombstone; change the branch
that handles the filtered new_workers in the block that mutates self.model_index
(where model_id is derived from worker.model_id() and worker_url is used in the
filter) so that if new_workers.is_empty() you call
self.model_index.remove(&model_id) to delete the key entirely, otherwise assign
*entry = Arc::from(new_workers.into_boxed_slice()) as before; keep the
subsequent call to self.rebuild_hash_ring(&model_id) so the hash ring is rebuilt
for removed keys as well.
- Around line 264-273: The current model_index update always appends the worker
clone which can create duplicates; inside the entry(...).and_modify closure (the
block that builds new_workers from existing and pushes worker.clone()) first
check whether an equivalent worker already exists by comparing the worker
URL/identifier (e.g. call the worker URL accessor such as worker.url() or
worker.address() against each entry in existing) and only push the new worker if
no match is found; update the logic in the and_modify closure to dedupe before
creating the new Arc slice so re-registering the same worker URL will not add a
duplicate entry.
ℹ️ Review info
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
📒 Files selected for processing (12)
bindings/python/src/lib.rsbindings/python/src/smg/router_args.pybindings/python/src/smg/serve.pybindings/python/tests/test_serve.pymodel_gateway/src/config/builder.rsmodel_gateway/src/config/types.rsmodel_gateway/src/core/steps/worker/local/create_worker.rsmodel_gateway/src/core/steps/worker/local/submit_tokenizer_job.rsmodel_gateway/src/core/worker.rsmodel_gateway/src/core/worker_registry.rsmodel_gateway/src/main.rsmodel_gateway/src/routers/grpc/common/responses/utils.rs
💤 Files with no reviewable changes (4)
- model_gateway/src/core/worker.rs
- model_gateway/src/config/builder.rs
- bindings/python/src/lib.rs
- model_gateway/src/config/types.rs
| """When connection_mode is http, --grpc-mode should not be present.""" | ||
| launcher = VllmWorkerLauncher() | ||
| args = argparse.Namespace(model="/tmp/model", connection_mode="http") | ||
| args = argparse.Namespace(model_path="/tmp/model", connection_mode="http") |
There was a problem hiding this comment.
Use model in this vLLM launcher test to avoid coverage drift.
Line 536 passes model_path, but VllmWorkerLauncher.build_command() reads args.model. This test now misses model propagation validation for vLLM.
Proposed fix
- args = argparse.Namespace(model_path="/tmp/model", connection_mode="http")
+ args = argparse.Namespace(model="/tmp/model", connection_mode="http")📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| args = argparse.Namespace(model_path="/tmp/model", connection_mode="http") | |
| args = argparse.Namespace(model="/tmp/model", connection_mode="http") |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@bindings/python/tests/test_serve.py` at line 536, The test instantiates args
with the wrong key: it sets model_path but VllmWorkerLauncher.build_command()
reads args.model, so update the test to pass the expected attribute (e.g., args
= argparse.Namespace(model="/tmp/model", connection_mode="http")) and ensure the
assertion that validates model propagation into
VllmWorkerLauncher.build_command() still runs; look for the test function around
the args assignment and the VllmWorkerLauncher.build_command() call to change
the Namespace key and preserve the existing propagation assertion.
2 participants
Summary
Revert feat: add --served-model-name support for model aliasing #521 (
--served-model-namein gateway): This feature should be implemented upstream in TensorRT-LLM (like SGLang and vLLM do), not worked around in the gateway. A corresponding TRT-LLM change adds--served-model-namenatively to TRT-LLM's serve command.Fix tokenizer registration for trtllm backend:
smg serve --backend trtllm --model <path>failed to register the tokenizer because trtllm uses--modelwhile sglang uses--model-path. The gateway'sRouterArgs.from_cli_argsonly looked formodel_pathin the namespace, missing trtllm'smodel. Fixed by mapping trtllm's--modeltodest=model_pathso it uses the same namespace key as sglang.Fix
from_cli_argsprefix fallback: Whenuse_router_prefix=True, the prefixed key (e.g.router_model_path) always existed in the namespace withdefault=None, which blocked the fallback to the unprefixedmodel_patheven when the user didn't pass--router-model-path. Now only prefers the prefixed version when explicitly set.Add
--modelas alias for--model-pathin both Rust CLI and Python RouterArgs for consistency across backends.Test plan
smg serve --backend trtllm --model <path>registers tokenizer correctlysmg serve --backend trtllm --model <path> --served-model-name abcuses model path for tokenizer, abc as model namesmg serve --backend sglang --model-path <path>continues to work as beforesmg launch --model <path> --worker-urls ...works (--model alias)Summary by CodeRabbit
New Features
Refactor
Tests