refactor(zmq): dedup engine scaffolding and move connect mechanics to the client layer - #2073
Conversation
github-actions
Bot
added
grpc
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughSummary by CodeRabbit
WalkthroughHarmony request construction now uses shared backend dispatch and stop-token handling. ZMQ connection setup moved into the client module. vLLM and TokenSpeed share stream and fan-out handling. Worker connection logic delegates to the client module. ChangesHarmony and ZMQ backend integration
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Worker
participant ZmqClient
participant Backend
participant StreamMapper
Worker->>ZmqClient: connect_for_worker(...)
ZmqClient->>Backend: establish ZMQ connection
Backend-->>ZmqClient: connected client
Worker->>Backend: submit generation request
Backend-->>StreamMapper: stream output ticks
StreamMapper-->>Worker: chunks and completion
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ 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 |
|
👋 The PR description doesn't fully follow
Please update the PR description so reviewers have the context they need. |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| req.include_stop_token_in_output = true; | ||
| } | ||
| debug!( | ||
| stop_token_count = harmony_stop_ids.len(), | ||
| "Injected Harmony stop tokens" | ||
| ); | ||
| } | ||
|
|
||
| // The client resolves string `stop`s its engine can't match and |
There was a problem hiding this comment.
🟡 Nit: The catch-all _ arm here returns a String error that the caller wraps as error::bad_request (HTTP 400). The old code used error::internal_error (HTTP 500) for the equivalent BackendClient::Zmq(zmq_client) arm with an unsupported runtime, which is semantically more accurate — this is a wiring bug, not a malformed client request.
Practically unreachable since connect() only admits vLLM/TokenSpeed runtimes, but if it's ever hit, a 400 would mislead the caller into thinking their request was wrong rather than pointing at an internal misconfiguration.
There was a problem hiding this comment.
Clean refactoring. The four changes (shared fan-out scaffolding, shared emit_tick, consolidated Harmony request builder, ZMQ connect mechanics moved to client layer) are all behavior-preserving and well-tested. One nit filed about the catch-all error type change (internal_error → bad_request) in the Harmony builder — practically unreachable but semantically off.
0 🔴 Important · 1 🟡 Nit · 0 🟣 Pre-existing
![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
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
model_gateway/src/routers/grpc/proto_wrapper.rs (1)
1095-1128: 📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win🔴 Important — The new method was inserted between the
#[expect]attribute andas_sglang.Lines 1095-1098 hold
#[expect(clippy::panic, ...)]and line 1094 holds the doc comment/// Get SGLang variant (panics if not SGLang). Both now attach toextend_stop_token_ids. Two consequences follow:
extend_stop_token_idscontains nopanic!, so the expectation is unfulfilled. Clippy reportsunfulfilled_lint_expectation.as_sglangat line 1128 still callspanic!but has lost its allowance and its doc comment. Ifclippy::panicis denied for this crate, the lint gate fails there.Place
extend_stop_token_idsafteras_sglangso the attribute and doc comment return to the accessor.🔧 Proposed fix: restore the accessor attribute and move the new method
impl ProtoGenerateRequest { - /// Get SGLang variant (panics if not SGLang) - #[expect( - clippy::panic, - reason = "typed accessor: caller guarantees variant via is_sglang() check" - )] /// Append stop token ids to the request's sampling params (TRT-LLM keeps /// them on the request itself). Requests without sampling params are left /// unchanged, matching the per-engine injection this replaces. pub fn extend_stop_token_ids(&mut self, ids: &[u32]) { @@ Self::Trtllm(req) => req.stop_token_ids.extend_from_slice(ids), } } + /// Get SGLang variant (panics if not SGLang) + #[expect( + clippy::panic, + reason = "typed accessor: caller guarantees variant via is_sglang() check" + )] pub fn as_sglang(&self) -> &sglang::GenerateRequest {🤖 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/routers/grpc/proto_wrapper.rs` around lines 1095 - 1128, Move extend_stop_token_ids so it appears after as_sglang and its full implementation, restoring the existing doc comment and #[expect(clippy::panic, ...)] attribute directly above as_sglang. Ensure the attribute applies to the panic-containing accessor and the new method no longer inherits either annotation.
🧹 Nitpick comments (5)
model_gateway/src/routers/grpc/harmony/stages/request_building.rs (2)
130-141: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value🟡 Nit — Replace the wildcard arm with explicit variants to keep exhaustiveness.
The request-id match at lines 88-116 already returns an error for
Generate,Completion,Embedding,Classify, andMessages. Both theEmbeddingarm and the_arm here are unreachable. The_arm also disables exhaustiveness checking. If a newRequestTypevariant is added later, this match silently accepts it and returns a generic error instead of failing to compile.List the remaining variants in one reject arm.
♻️ Proposed refactor
RequestType::Responses(request) => HarmonyBody::Responses(request.as_ref()), - RequestType::Embedding(_) => { - return Err(error::bad_request( - "harmony_embedding_not_supported", - "Embedding requests are not supported with Harmony models".to_string(), - )); - } - _ => { + RequestType::Generate(_) + | RequestType::Completion(_) + | RequestType::Embedding(_) + | RequestType::Classify(_) + | RequestType::Messages(_) => { return Err(error::bad_request( "unsupported_request_type", "Unsupported request type for Harmony models".to_string(), )); }🤖 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/routers/grpc/harmony/stages/request_building.rs` around lines 130 - 141, Update the RequestType match in the request-building flow to replace the wildcard arm with one explicit reject arm listing every remaining RequestType variant, while retaining the existing Embedding-specific error and behavior for already handled variants. Preserve exhaustiveness checking so adding a new variant requires updating this match.
152-158: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value🟡 Nit — The single error path maps wiring bugs to HTTP 400.
build_harmony_protoreturns two distinct classes of error. Builder failures come from client input. The fallback arm at lines 342-346 returns"unsupported backend runtime ...", which the comment describes as a wiring bug. Both now becomebad_request("invalid_request_parameters", ...). A server misconfiguration is then reported to the caller as a 400 and counted as a client error.Distinguish the two classes. One option is a small error enum from
build_harmony_proto; a simpler option is to keep the unsupported-runtime case as aninternal_error.🤖 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/routers/grpc/harmony/stages/request_building.rs` around lines 152 - 158, Update the error handling around HarmonyRequestBuildingStage::execute and build_harmony_proto to distinguish client validation failures from unsupported backend runtime failures. Preserve bad_request("invalid_request_parameters", ...) for invalid input, but map the unsupported-runtime fallback to internal_error so wiring or server configuration issues are reported as server errors.model_gateway/src/routers/grpc/zmq_client.rs (3)
242-242: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value🟡 Nit —
failis an identity closure.
let fail = |reason: String| reason;returns its argument unchanged. Every call site wraps aformat!that is already the final error value. Remove the closure and pass the formatted string directly.🤖 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/routers/grpc/zmq_client.rs` at line 242, Remove the identity closure fail and update each of its call sites to pass the existing format! result directly as the error value, preserving the current messages and behavior.
1409-1416: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win🟡 Nit — Add a test for the symlinked-parent rejection.
The doc comment at Line 243 states that
symlink_metadatais used so "a symlinked parent must not redirect the checks (or the sockets) into a directory we did not verify." That property is the security reason for choosingsymlink_metadataovermetadata, and no test pins it. A future change back tometadatawould pass the whole suite.💚 Proposed test
#[cfg(unix)] #[tokio::test] async fn ensure_ipc_socket_dir_rejects_a_symlinked_parent() { let base = tempfile::tempdir().unwrap(); let real = base.path().join("real"); std::fs::create_dir(&real).unwrap(); let link = base.path().join("link"); std::os::unix::fs::symlink(&real, &link).unwrap(); let url = format!("ipc://{}/x.ipc", link.display()); assert!(ensure_ipc_socket_dir(&url).await.is_err()); }As per coding guidelines: "Run the pr-test-analyzer agent to verify that tests adequately cover new or changed functionality."
🤖 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/routers/grpc/zmq_client.rs` around lines 1409 - 1416, Add a Unix-only async test alongside ensure_ipc_socket_dir_rejects_a_non_directory_parent that creates a real directory, symlinks a parent path to it, and verifies ensure_ipc_socket_dir returns an error for an IPC URL using the symlinked parent. Keep the test focused on rejecting symlinked parents and use the existing tempfile and tokio test patterns.Source: Coding guidelines
304-314: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low value🟡 Nit — Blocking filesystem calls run on the runtime thread here.
model_dir.is_dir()andEosTokenIds::from_model_dirperform synchronous filesystem I/O inside anasync fn. The adjacentensure_ipc_socket_dirwas deliberately made async for exactly this reason (see its doc comment at Line 228). The cost is small — one stat plus two small JSON reads per connect — but the two paths now follow different policies. Consider wrapping the EOS resolution intokio::task::spawn_blockingor converting it totokio::fs.🤖 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/routers/grpc/zmq_client.rs` around lines 304 - 314, Move the synchronous EOS resolution in the connect flow into a blocking-safe async path: wrap both model_dir.is_dir() and EosTokenIds::from_model_dir in tokio::task::spawn_blocking, or replace their filesystem access with tokio::fs while preserving the existing warning and default behavior for non-directory model IDs. Handle the async task/result errors consistently with the surrounding connection logic.
🤖 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/routers/grpc/harmony/stages/request_building.rs`:
- Around line 285-347: Add unit tests for build_harmony_proto covering ZMQ vLLM
and TokenSpeed clients with both Chat and Responses requests, verifying each
dispatches to the correct ProtoGenerateRequest variant. Also test the
unsupported-runtime fallback returns the expected error, using existing
request/client test helpers and keeping the builder behavior unchanged.
In `@model_gateway/src/routers/grpc/zmq_client.rs`:
- Around line 922-933: Update TokenSpeedGenerateStream::map_output to detect the
wire-level "error" finish reason before normalizing it or calling
state.emit_tick, and return the corresponding Result::Err instead of emitting a
normal completion. Match the existing vLLM error mapping behavior.
---
Outside diff comments:
In `@model_gateway/src/routers/grpc/proto_wrapper.rs`:
- Around line 1095-1128: Move extend_stop_token_ids so it appears after
as_sglang and its full implementation, restoring the existing doc comment and
#[expect(clippy::panic, ...)] attribute directly above as_sglang. Ensure the
attribute applies to the panic-containing accessor and the new method no longer
inherits either annotation.
---
Nitpick comments:
In `@model_gateway/src/routers/grpc/harmony/stages/request_building.rs`:
- Around line 130-141: Update the RequestType match in the request-building flow
to replace the wildcard arm with one explicit reject arm listing every remaining
RequestType variant, while retaining the existing Embedding-specific error and
behavior for already handled variants. Preserve exhaustiveness checking so
adding a new variant requires updating this match.
- Around line 152-158: Update the error handling around
HarmonyRequestBuildingStage::execute and build_harmony_proto to distinguish
client validation failures from unsupported backend runtime failures. Preserve
bad_request("invalid_request_parameters", ...) for invalid input, but map the
unsupported-runtime fallback to internal_error so wiring or server configuration
issues are reported as server errors.
In `@model_gateway/src/routers/grpc/zmq_client.rs`:
- Line 242: Remove the identity closure fail and update each of its call sites
to pass the existing format! result directly as the error value, preserving the
current messages and behavior.
- Around line 1409-1416: Add a Unix-only async test alongside
ensure_ipc_socket_dir_rejects_a_non_directory_parent that creates a real
directory, symlinks a parent path to it, and verifies ensure_ipc_socket_dir
returns an error for an IPC URL using the symlinked parent. Keep the test
focused on rejecting symlinked parents and use the existing tempfile and tokio
test patterns.
- Around line 304-314: Move the synchronous EOS resolution in the connect flow
into a blocking-safe async path: wrap both model_dir.is_dir() and
EosTokenIds::from_model_dir in tokio::task::spawn_blocking, or replace their
filesystem access with tokio::fs while preserving the existing warning and
default behavior for non-directory model IDs. Handle the async task/result
errors consistently with the surrounding connection logic.
🪄 Autofix
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: CHILL
Plan: Pro Plus
Run ID: 7405db94-b302-41ca-8581-b2a82df27bef
📒 Files selected for processing (4)
model_gateway/src/routers/grpc/harmony/stages/request_building.rsmodel_gateway/src/routers/grpc/proto_wrapper.rsmodel_gateway/src/routers/grpc/zmq_client.rsmodel_gateway/src/worker/worker.rs
fan_out_requests and fan_out_tokenspeed_requests were structurally identical (clone n subs, suffix rids, collapse n to 1) differing only in vLLM's per-sub seed derivation. Extract fan_out_n; each engine keeps just its sampling tweak closure. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
…streams VllmGenerateStream and TokenSpeedGenerateStream carried near-identical ~50-line emission tails: build the Complete, and when the finish tick also carried tokens, emit a Chunk first and park the Complete for the next poll. Move that dance into StreamState::emit_tick; each mapper now contributes only its engine-specific inputs (finish-reason string, matched_stop, token ids). vLLM's engine-error surfacing stays in its mapper; TokenSpeed documents its absent matched_stop. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
…(backend, kind) The builder was a ~330-line 7-arm backend match, each arm repeating the same 4-way request-type match, per-arm error mapping, and reject arms — and the two ZMQ arms re-duplicated the gRPC vLLM/TokenSpeed arms verbatim. - Resolve the request kind once (HarmonyBody: Chat with the modified body applied, or Responses), with the rejects stated a single time. - build_harmony_proto is one compact (backend x kind) match where every arm is just the engine's builder call, sharing a single error path at the call site. vLLM and TokenSpeed build through static translators, so one arm each now covers both gRPC and direct-ZMQ — the transport duplication is gone, and an impossible runtime remains a hard error. - The 5-arm Harmony stop-id injection block becomes ProtoGenerateRequest::extend_stop_token_ids, with TRT-LLM's include_stop_token_in_output kept as the one engine-specific line. Net -164 lines with no behavior change. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
worker/worker.rs owned the entire ZMQ connect path — handshake-port derivation (FNV-1a, pinned against serve.py), socket-address layout, ipc-dir preparation, EOS resolution, and the engine handshake — while reverse-importing ZmqEngineClient from routers::grpc to do it. The worker layer was reaching into the router layer for its transport. Move it all behind zmq_client::connect_for_worker: one connect path, plain-reason errors, owned by the module that owns every other piece of the ZMQ client (EOS, translate, streams). The worker keeps a thin wrapper that maps into its client slot and WorkerError, and its reverse import shrinks to the BackendClient/GrpcClient types it stores. The pinned FNV vectors and address/ipc-dir tests move with the code. Also corrects the stale EOS warning: connect-time ids being unavailable falls back to the tokenizer EOS set at request time, it does not mean generation runs to max_tokens. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
c1f35d4 to
52e319b
Compare
| let Some(parent) = Path::new(path).parent() else { | ||
| return Ok(()); | ||
| }; | ||
| let fail = |reason: String| reason; |
There was a problem hiding this comment.
🟡 Nit: fail is now an identity closure (|reason: String| reason) — a leftover from the move out of worker.rs where it wrapped into WorkerError::ConnectionFailed. Every call site is just fail(format!(...)) which is format!(...). You can drop the binding and pass the format!() strings directly to Err / .map_err().
| let fail = |reason: String| reason; |
|
Review comments addressed in f7c1ec8:
|
… error classes - extend_stop_token_ids was inserted between as_sglang's doc + #[expect(clippy::panic)] and the fn itself, hijacking both (and leaving as_sglang's panic unshielded — the CI lint failure). Restore the attachment and give the new method its own placement, plus a test covering the sampling-params/TRT-LLM/missing-params variants. - TokenSpeed streams now surface an engine 'error' finish reason as a tonic error before emit_tick, mirroring the vLLM stream's guard, instead of emitting a normal Complete for a failed request. - Harmony build failures are typed: builder rejections stay 400, the impossible backend/runtime pairing is a wiring bug and returns 500 (restoring the pre-refactor semantics the collapse had lost). - Drop the identity 'fail' closure left over from the worker.rs move. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
f7c1ec8 to
baea129
Compare
1 participant
Motivation
Phase 2 of the ZMQ architecture cleanup (follows #2065/#2068): remove the duplication blocks and the one layering inversion the audit identified. No behavior changes.
Modifications (one commit each)
fan_out_requests/fan_out_tokenspeed_requestswere ~90% identical;fan_out_nowns the clone/suffix/collapse dance, each engine keeps only its sampling-tweak closure.StreamState::emit_tick; mappers contribute only finish reason, matched stop, and token ids.HarmonyBodyresolution +build_harmony_proto, with one shared error path. vLLM/TokenSpeed build via static translators, so one arm each now covers both gRPC and direct-ZMQ. The 5-arm stop-id injection becomesProtoGenerateRequest::extend_stop_token_ids. Net −164 lines.zmq_client— handshake-port derivation (FNV pinned against serve.py, vectors moved with it), socket layout, ipc-dir prep, EOS resolution, and the handshake now live inzmq_client::connect_for_worker;worker.rskeeps a thin wrapper and stops reaching into the router layer for its transport. Also fixes the misleading EOS warning (connect-time miss falls back to tokenizer EOS at request time — it never meant 'runs to max_tokens').Checklist
cargo fmtclean; clippy--all-targetsclean on changed filescargo test -p smg --lib: 1435 passed (single failure is the pre-existingmiddleware::metricsinterner flake, identical on clean main)