feat(gateway): Filter /workers by model - #1672
Conversation
📝 WalkthroughWalkthroughThe ChangesWorker listing with model filtering
Sequence DiagramsequenceDiagram
participant Client
participant REST as /workers endpoint
participant WorkerService
Client->>REST: GET /workers?model=llama-3
REST->>REST: extract ListWorkersQuery
REST->>WorkerService: list_workers(Some("llama-3"))
WorkerService->>WorkerService: iterate registered workers
WorkerService->>WorkerService: skip workers without model support
WorkerService->>WorkerService: compute counts from filtered set
WorkerService-->>REST: ListWorkersResult
REST-->>Client: filtered worker list
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 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 |
![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 introduces model-based filtering for the worker listing endpoint. It updates the REST handler to accept an optional model query parameter and modifies the list_workers service method to filter workers and compute statistics dynamically based on the filtered results. Unit tests have also been added to verify this behavior. The review feedback suggests two key improvements: first, handling empty or whitespace-only query parameters as None to avoid filtering by empty strings, and second, optimizing the worker retrieval loop to avoid unnecessary heap allocations and cloning of worker IDs on large registries.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| async fn list_workers_rest( | ||
| State(state): State<Arc<AppState>>, | ||
| Query(query): Query<ListWorkersQuery>, | ||
| ) -> Response { | ||
| state | ||
| .context | ||
| .worker_service | ||
| .list_workers(query.model.as_deref()) | ||
| .into_response() | ||
| } |
There was a problem hiding this comment.
Empty or whitespace-only model query parameters (e.g., ?model=) should be treated as None rather than attempting to filter by an empty string, which would return no workers. Filtering them out ensures robust handling of empty query parameters.
| async fn list_workers_rest( | |
| State(state): State<Arc<AppState>>, | |
| Query(query): Query<ListWorkersQuery>, | |
| ) -> Response { | |
| state | |
| .context | |
| .worker_service | |
| .list_workers(query.model.as_deref()) | |
| .into_response() | |
| } | |
| async fn list_workers_rest( | |
| State(state): State<Arc<AppState>>, | |
| Query(query): Query<ListWorkersQuery>, | |
| ) -> Response { | |
| let model = query.model.as_deref().filter(|s| !s.trim().is_empty()); | |
| state | |
| .context | |
| .worker_service | |
| .list_workers(model) | |
| .into_response() | |
| } |
| for (worker_id, worker) in self.worker_registry.get_all_with_ids() { | ||
| if let Some(model) = model { | ||
| if !worker.supports_model(model) { | ||
| continue; | ||
| } | ||
| } | ||
| match worker.worker_type() { | ||
| WorkerType::Prefill => prefill_count += 1, | ||
| WorkerType::Decode => decode_count += 1, | ||
| WorkerType::Regular => regular_count += 1, | ||
| } | ||
| let mut info = worker_to_info(&worker); | ||
| info.id = worker_id.as_str().to_string(); | ||
| worker_infos.push(info); | ||
| } |
There was a problem hiding this comment.
Using get_all_with_ids() clones the WorkerId (which contains a String) and the Arc<dyn Worker> for every single worker in the registry. In a gateway serving 10,000 workers, this results in 10,000 heap allocations even if only a few workers match the filter.
By using get_all() instead, we only clone the Arcs (which are cheap atomic increments with no heap allocations). We can then filter the workers first, and only look up and clone the WorkerIds for the matching workers using get_id_by_url(). This drastically reduces heap allocations and CPU overhead.
Additionally, for types like WorkerId that are frequently cloned on hot paths and represent a small, repeated set of values, we should use an interned string type like Arc<str> to make clones cheap (atomic reference count increments) and eliminate heap allocations entirely.
for worker in self.worker_registry.get_all() {
if let Some(model) = model {
if !worker.supports_model(model) {
continue;
}
}
match worker.worker_type() {
WorkerType::Prefill => prefill_count += 1,
WorkerType::Decode => decode_count += 1,
WorkerType::Regular => regular_count += 1,
}
let mut info = worker_to_info(&worker);
if let Some(worker_id) = self.worker_registry.get_id_by_url(worker.url()) {
info.id = worker_id.as_str().to_string();
}
worker_infos.push(info);
}References
- For types that are frequently cloned on hot paths and represent a small, repeated set of values (e.g., worker IDs or tenant IDs), use an interned string type like
Arc<str>to improve performance by making clones cheap (atomic reference count increments).
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
![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 current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@model_gateway/src/worker/service.rs`:
- Around line 346-375: The current list_workers(&self, model: Option<&str>) ->
ListWorkersResult implementation only returns total and per-type counts and thus
no longer matches the published protocol schema (which expects fields like
stats.total_workers, healthy_workers, total_models, total_load and nested
by_type). Update this function to return the shared response type used by the
REST endpoint (or extend ListWorkersResult and its serializer) so filtered
responses remain wire-compatible with crates::protocols::worker (lines
~800-824): compute and populate stats.total_workers, healthy_workers,
total_models, total_load and the by_type breakdown in the returned object (reuse
worker_to_info outputs for workers list), and add/adjust unit tests for the
serializer/REST handler (e.g., list_workers_rest) to verify the filtered
response matches the protocol schema.
- Around line 353-357: The optional model filter (the local variable named
`model`) must be normalized so blank or whitespace-only strings are treated as
None before the loop that checks `worker.supports_model(model)`; update the code
around the `model` binding in `service.rs` to trim and convert
`Some("")`/whitespace to `None` (so the later `if let Some(model) = model { if
!worker.supports_model(model) { continue; } }` logic behaves correctly and does
not erroneously exclude concrete-model workers when the query was empty).
🪄 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: fef8bf1e-2e78-425e-a9db-81c6b1b7bf31
📒 Files selected for processing (2)
model_gateway/src/server.rsmodel_gateway/src/worker/service.rs
| pub fn list_workers(&self, model: Option<&str>) -> ListWorkersResult { | ||
| let mut worker_infos = Vec::new(); | ||
| let mut prefill_count = 0; | ||
| let mut decode_count = 0; | ||
| let mut regular_count = 0; | ||
|
|
||
| for (worker_id, worker) in self.worker_registry.get_all_with_ids() { | ||
| if let Some(model) = model { | ||
| if !worker.supports_model(model) { | ||
| continue; | ||
| } | ||
| } | ||
| match worker.worker_type() { | ||
| WorkerType::Prefill => prefill_count += 1, | ||
| WorkerType::Decode => decode_count += 1, | ||
| WorkerType::Regular => regular_count += 1, | ||
| } | ||
| let mut info = worker_to_info(&worker); | ||
| info.id = worker_id.as_str().to_string(); | ||
| worker_infos.push(info); | ||
| } | ||
|
|
||
| ListWorkersResult { | ||
| total: worker_infos.len(), | ||
| workers: worker_infos, | ||
| total: stats.total_workers, | ||
| prefill_count: stats.prefill_workers, | ||
| decode_count: stats.decode_workers, | ||
| regular_count: stats.regular_workers, | ||
| prefill_count, | ||
| decode_count, | ||
| regular_count, | ||
| } | ||
| } |
There was a problem hiding this comment.
Keep filtered /workers responses aligned with the shared protocol schema.
Line 346 now recomputes only total and per-type counters, but the published /workers contract in crates/protocols/src/worker.rs:800-824 includes stats.total_workers, healthy_workers, total_models, total_load, and nested by_type. Since list_workers_rest returns this service result directly, filtered responses still cannot preserve that contract or recompute the full filtered stats set. Please switch this path to the shared response type, or extend ListWorkersResult/its serializer and tests so the filtered endpoint stays wire-compatible.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@model_gateway/src/worker/service.rs` around lines 346 - 375, The current
list_workers(&self, model: Option<&str>) -> ListWorkersResult implementation
only returns total and per-type counts and thus no longer matches the published
protocol schema (which expects fields like stats.total_workers, healthy_workers,
total_models, total_load and nested by_type). Update this function to return the
shared response type used by the REST endpoint (or extend ListWorkersResult and
its serializer) so filtered responses remain wire-compatible with
crates::protocols::worker (lines ~800-824): compute and populate
stats.total_workers, healthy_workers, total_models, total_load and the by_type
breakdown in the returned object (reuse worker_to_info outputs for workers
list), and add/adjust unit tests for the serializer/REST handler (e.g.,
list_workers_rest) to verify the filtered response matches the protocol schema.
| if let Some(model) = model { | ||
| if !worker.supports_model(model) { | ||
| continue; | ||
| } | ||
| } |
There was a problem hiding this comment.
Normalize blank model filters before matching.
Line 353 treats Some("") and whitespace-only values as real model IDs. With GET /workers?model=, every concrete-model worker is filtered out while wildcard workers still pass supports_model, so the optional filter returns a misleading partial inventory. Normalize blank input to None before entering the loop.
Suggested fix
pub fn list_workers(&self, model: Option<&str>) -> ListWorkersResult {
+ let model = model.map(str::trim).filter(|model| !model.is_empty());
let mut worker_infos = Vec::new();
let mut prefill_count = 0;
let mut decode_count = 0;
let mut regular_count = 0;
for (worker_id, worker) in self.worker_registry.get_all_with_ids() {
if let Some(model) = model {
if !worker.supports_model(model) {
continue;
}
}🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@model_gateway/src/worker/service.rs` around lines 353 - 357, The optional
model filter (the local variable named `model`) must be normalized so blank or
whitespace-only strings are treated as None before the loop that checks
`worker.supports_model(model)`; update the code around the `model` binding in
`service.rs` to trim and convert `Some("")`/whitespace to `None` (so the later
`if let Some(model) = model { if !worker.supports_model(model) { continue; } }`
logic behaves correctly and does not erroneously exclude concrete-model workers
when the query was empty).
This allows you to filter the `GET /workers` by model ID, improving performance / reducing CPU overhead in gateways that serve many workers (~10 thousand with DP-aware) spread across ~100 models. Signed-off-by: Eric Zhang <ekzhang1@gmail.com>
43e7a46 to
4a7fee8
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
model_gateway/src/worker/service.rs (1)
346-375:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winNormalize blank
modelfilters before matching.The optional
modelparameter is not normalized, so empty or whitespace-only query strings (?model=or?model=%20) arrive asSome("")orSome(" ")rather thanNone. Line 353'sif let Some(model) = modelthen callsworker.supports_model(""), which may yield unexpected filtering behavior (e.g., filtering out concrete-model workers while wildcard workers pass). Normalize blank input toNonebefore the loop.🛡️ Proposed fix
pub fn list_workers(&self, model: Option<&str>) -> ListWorkersResult { + let model = model.map(str::trim).filter(|m| !m.is_empty()); let mut worker_infos = Vec::new(); let mut prefill_count = 0; let mut decode_count = 0;🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@model_gateway/src/worker/service.rs` around lines 346 - 375, In list_workers, normalize the optional model filter before the loop by trimming whitespace and turning empty strings into None (e.g., let normalized_model = model.map(|m| m.trim()).and_then(|s| if s.is_empty() { None } else { Some(s) });), then use normalized_model instead of model when checking and calling worker.supports_model; this ensures "" or " " are treated as None and avoids unexpected filtering in worker.supports_model.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@model_gateway/src/worker/service.rs`:
- Around line 464-540: Add a new unit test in the tests module that verifies
empty and whitespace-only model filters are treated like None: use make_service
and register_worker to create the same two workers as in
test_list_workers_unfiltered_returns_all, then call
service.list_workers(Some("")) and service.list_workers(Some(" ")) and assert
each returns the same workers.len(), total, and regular_count as
service.list_workers(None); reference WorkerService::list_workers, make_service,
and register_worker to locate where to add the test.
---
Duplicate comments:
In `@model_gateway/src/worker/service.rs`:
- Around line 346-375: In list_workers, normalize the optional model filter
before the loop by trimming whitespace and turning empty strings into None
(e.g., let normalized_model = model.map(|m| m.trim()).and_then(|s| if
s.is_empty() { None } else { Some(s) });), then use normalized_model instead of
model when checking and calling worker.supports_model; this ensures "" or " "
are treated as None and avoids unexpected filtering in worker.supports_model.
🪄 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: f00cec13-c274-49ae-a3d0-bbf92efc8124
📒 Files selected for processing (2)
model_gateway/src/server.rsmodel_gateway/src/worker/service.rs
2 participants
Description
Problem
This allows you to filter the
GET /workersby model ID, improving performance / reducing CPU overhead in gateways that serve many workers (~10 thousand with DP-aware) spread across ~100 models.Solution
/workersnow takes a?model=query parameter that filters by workers. The output response schema is exactly the same as before, just filtered now. No performance impact as it's just a cheap check while iterating the list to construct theVec<WorkerInfo>(already done)Changes
/workersendpointTest Plan
/workersendpointChecklist
cargo +nightly fmtpassescargo clippy --all-targets --all-features -- -D warningspassesSummary by CodeRabbit
New Features
Tests