refactor(zmq): clients finalize their own requests (no more stage-level transport probes) - #2068
Merged
Merged
Conversation
📝 WalkthroughSummary by CodeRabbit
WalkthroughThe change centralizes backend-specific stop finalization, adds router-owned stop obligations to response state, applies a ZMQ EOS backstop, and restricts ZMQ dispatch to supported runtimes. Harmony and regular request paths now use the centralized finalization flow. ChangesStop finalization and runtime routing
Estimated code review effort: 3 (Moderate) | ~25 minutes Possibly related PRs
Suggested labels: 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 |
github-actions
Bot
added
grpc
|
👋 The PR description doesn't fully follow PULL_REQUEST_TEMPLATE.md:
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}
|
|
||
| // gRPC vLLM matches stops server-side: nothing left for the router. | ||
| let mut req = vllm_request(vec!["."], vec![]); | ||
| assert!(resolve_string_stops(&mut req, Some(&mock_tokenizer()), false).is_empty()); |
There was a problem hiding this comment.
🟡 Nit: This test covers SGLang (always returns obligations) and gRPC vLLM (token_only_wire=false → empty), but doesn't assert the return value for the two ZMQ arms that also return obligations: vLLM with token_only_wire=true and TokenSpeed with token_only_wire=true.
The existing vllm_zmq_resolves_string_stops_and_injects_eos test above verifies side-effects (stop cleared, token ids set) but discards the return value. Worth adding a couple of assertions here to close the gap, e.g.:
// ZMQ vLLM: engine can't match strings, so all come back as obligations.
let mut req = vllm_request(vec![".", "Hello world"], vec![]);
let obligations = resolve_string_stops(&mut req, Some(&mock_tokenizer()), true);
assert_eq!(obligations, vec![".".to_string(), "Hello world".to_string()]);
There was a problem hiding this comment.
Clean refactor — the single-owner finalization model is a clear improvement over the scattered is_zmq() probes. The behavioral change for Harmony + SGLang gRPC (now correctly returning stop obligations) is an improvement, not a regression. One minor test coverage nit posted inline.
0 🔴 Important · 1 🟡 Nit · 0 🟣 Pre-existing
slin1237
force-pushed
the
zmq/phase0-arch-fixes
branch
from
August 7, 2026 05:46
eb76224 to
e08e75e
Compare
slin1237
force-pushed
the
zmq/phase1-client-finalization
branch
from
August 7, 2026 05:56
e82264a to
6f6ef14
Compare
…stop duty The four request-building stages read builder_client.is_zmq() and threaded the flag into helpers::resolve_string_stops — a client property enforced two layers above the client. Move the decision to the client: - BackendClient::finalize_generate_request resolves string stops its engine cannot match (token-only wires, SGLang skip_tokenizer_init) and returns the router's residual obligation: the stop strings the engine will never see, which response processing must trim from output text. - resolve_string_stops returns those stripped strings instead of discarding them; its transport flag is now supplied by the client, and no stage carries is_zmq knowledge anymore. - Obligation wiring into response processing lands with the follow-up convergence commit; stages currently discard the value, preserving today's behavior exactly. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
The EOS story was split across two layers: connect-time model-dir resolution (EosTokenIds, zmq_client) and a request-time tokenizer fold buried inside the shared stop-resolution helper — vLLM-ZMQ policy living in a stage helper every backend shares. Move the fold into zmq_client::fold_tokenizer_eos_backstop, invoked by BackendClient::finalize_generate_request for vLLM-over-ZMQ only (the tokenizer-less EngineCore stops at EOS only when the frontend supplies ids; TokenSpeed's scheduler stops itself). resolve_string_stops is now purely about stop strings, and every EOS mechanism — connect-time ids, translate-time merge, request-time backstop — lives in one file. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
…_zmq Harmony's response-processing stage inspected the selected client (is_zmq on the single/decode leg) to decide whether the router must scan channel text for string stops — the last response-side transport probe. Converge it onto the request-building contract instead: - HarmonyRequestBuildingStage finalizes its built request through BackendClient::finalize_generate_request like the regular stages, and stores the returned obligation on ResponseState. - router_stop_strings() now just reads that stored obligation; empty for engines that match stops server-side. No transport inspection. - The regular request-building stages store the same obligation instead of discarding it (parity for Chat: builders map request.stop into the proto, so the obligation equals what the old probe read; Responses builders don't carry stops yet, matching the probe's empty result). The regular pipeline's response side intentionally still builds its decoder from request params: MLX's proto has no string-stop field, so the built request cannot report those strings as obligations yet, and the ungated decoder remains the mechanism that trims them. Folding that path onto obligations is future work once MLX stops ride the proto. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
Six dispatch sites treated "any ZMQ runtime that isn't TokenSpeed" as vLLM via catch-all arms — an Unspecified (or future) runtime would silently build vLLM-shaped requests instead of failing. connect() already admits only vLLM/TokenSpeed over ZMQ, so make that contract visible at every dispatch point: - backend_client.rs: the four build_*_request ZMQ matches name RuntimeType::Vllm explicitly and return a hard error for anything else. - harmony request building: the Zmq catch-all becomes a Vllm-guarded arm plus an explicit unsupported-runtime error. - multimodal assemble: the ZMQ arm drops its Unspecified fallback; the existing not-supported error now catches it. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
The repo lints deny paths of 4+ segments (absolute-paths-max-segments = 3); bring helpers and fold_tokenizer_eos_backstop into scope in backend_client, and MockTokenizer in the zmq_client tests. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
slin1237
force-pushed
the
zmq/phase1-client-finalization
branch
from
August 7, 2026 07:32
6f6ef14 to
45453dd
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
model_gateway/src/routers/grpc/backend_client.rs (1)
68-83: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy lift🟡 Nit Add combined finalization tests.
The supplied tests cover
resolve_string_stopsandfold_tokenizer_eos_backstopseparately. Add tests forfinalize_generate_requestthat verify vLLM ZMQ receives stop IDs plus EOS IDs, TokenSpeed ZMQ receives no EOS IDs, and gRPC vLLM keeps server-side string stops.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/backend_client.rs` around lines 68 - 83, Add combined tests for finalize_generate_request covering vLLM ZMQ requests retaining both resolved stop IDs and folded EOS IDs, TokenSpeed ZMQ requests receiving no EOS IDs, and vLLM gRPC requests preserving server-side string stops. Reuse the existing test helpers and fixtures for resolve_string_stops, fold_tokenizer_eos_backstop, runtime selection, and request construction, then run the pr-test-analyzer agent to verify coverage.Source: Coding guidelines
🤖 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.
Nitpick comments:
In `@model_gateway/src/routers/grpc/backend_client.rs`:
- Around line 68-83: Add combined tests for finalize_generate_request covering
vLLM ZMQ requests retaining both resolved stop IDs and folded EOS IDs,
TokenSpeed ZMQ requests receiving no EOS IDs, and vLLM gRPC requests preserving
server-side string stops. Reuse the existing test helpers and fixtures for
resolve_string_stops, fold_tokenizer_eos_backstop, runtime selection, and
request construction, then run the pr-test-analyzer agent to verify coverage.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 444ff063-e6bc-413d-9d73-f60a1118eaf8
📒 Files selected for processing (11)
model_gateway/src/routers/grpc/backend_client.rsmodel_gateway/src/routers/grpc/common/stages/helpers.rsmodel_gateway/src/routers/grpc/context.rsmodel_gateway/src/routers/grpc/harmony/stages/request_building.rsmodel_gateway/src/routers/grpc/harmony/stages/response_processing.rsmodel_gateway/src/routers/grpc/multimodal/assemble.rsmodel_gateway/src/routers/grpc/regular/stages/chat/request_building.rsmodel_gateway/src/routers/grpc/regular/stages/completion/request_building.rsmodel_gateway/src/routers/grpc/regular/stages/generate/request_building.rsmodel_gateway/src/routers/grpc/regular/stages/messages/request_building.rsmodel_gateway/src/routers/grpc/zmq_client.rs
2 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
Stacked on #2065. Phase 1 of the ZMQ architecture cleanup: the pipeline had one structural fact to model — vLLM EngineCore has no tokenizer — and it was being expressed as
is_zmq()probes and catch-all runtime arms scattered across stages. This PR moves those decisions into the layer that already is the abstraction:BackendClient. No new types, tables, or layers.Modifications (one commit each, reviewable in order)
BackendClient::finalize_generate_requestresolves stringstops its engine can't match (token-only wires, SGLangskip_tokenizer_init) and returns the router's residual obligation — the stop strings the engine will never see. All four regular request-building stages drop theiris_zmq()probes.zmq_client::fold_tokenizer_eos_backstop, next to the connect-timeEosTokenIds.resolve_string_stopsis now purely about stop strings.is_zmqon the single/decode leg) — the last response-side transport probe is gone. Regular-pipeline response side intentionally unchanged: MLX's proto carries no string-stop field, so its trimming still rides the ungated decoder (noted in the commit message as future work)._ => vLLMcatch-alls (4× backend_client build matches, harmony'sZmq(_)arm, assemble'sVllm | Unspecified) become explicitRuntimeType::Vllmarms with hard errors for anything else —connect()already admits only vLLM/TokenSpeed, so this makes that contract visible at every dispatch point.Checklist
cargo fmtcleancargo test -p smg --lib: 1435 passed (the single failure,middleware::metrics::…interner, fails identically on clean main under the full parallel suite — pre-existing flake, passes in isolation)