fix(zmq): correct transport-identity leaks in metrics, admin, and PD paths - #2065
Conversation
…paths Phase-0 fixes from the ZMQ architecture audit: places where ZMQ workers were miscounted, mislabeled, or could panic because code assumed every gRPC-pipeline worker speaks gRPC. - PD/EPD worker-selection metrics and PD execution error records now label each leg with the worker's actual connection mode instead of hardcoding grpc. - flush_cache_all skips ZMQ workers (no cache-flush RPC) instead of reporting every one as failed, and says so in the result message. - GrpcRouter's PD debug output counts all gRPC-pipeline workers (gRPC and ZMQ), not just ConnectionMode::Grpc. - inject_sglang_bootstrap_metadata matches the proto variant defensively instead of using the panicking as_sglang_mut accessor. - EPD encode rejection for TokenSpeed-over-ZMQ now names the transport limitation rather than claiming the engine is unsupported. - Document the ZMQ response-side contract at the ProtoStream seam: all ZMQ engines emit vllm-shaped responses, so engine logic must key on the worker's runtime_type, never the response variant. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
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 (8)
🚧 Files skipped from review as they are similar to previous changes (8)
📝 WalkthroughSummary by CodeRabbit
WalkthroughThe changes add ZMQ-specific validation, protect SGLang metadata injection, record actual worker connection modes in routing metrics, count gRPC-pipeline workers, and skip unsupported ZMQ workers during cache flushing. ChangesZMQ worker handling
Estimated code review effort: 3 (Moderate) | ~20 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 |
|
👋 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}
| .into_iter() | ||
| .partition(|w| !matches!(w.connection_mode(), ConnectionMode::Zmq)); | ||
| let zmq_skipped = zmq_workers.len(); | ||
| let grpc_workers = total_workers - http_workers - zmq_skipped; |
There was a problem hiding this comment.
🟡 Nit: total_workers still counts ZMQ workers, but grpc_workers now excludes them. This breaks the prior invariant total_workers == http_workers + grpc_workers that the JSON API response in FlushCacheResult::into_response exposes (total_workers, total_http_workers, total_grpc_workers).
API consumers doing total - http - grpc to sanity-check will get a non-zero remainder with no structured field to explain it. The message string mentions ZMQ but the structured fields don't. Consider either adding a zmq_skipped field to the struct+response, or redefining total_workers to exclude ZMQ so the arithmetic stays clean.
There was a problem hiding this comment.
Clean mechanical PR that fixes transport-identity leaks across metrics, admin, and PD paths. All changes are well-scoped and correct — the dynamic connection-mode labeling, defensive let-else in helpers, ZMQ partition in flush_cache, and the EPD encode guard all look good. One minor nit on the FlushCacheResult API invariant posted inline.
Summary: 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: 3
🧹 Nitpick comments (3)
model_gateway/src/routers/grpc/common/stages/helpers.rs (1)
455-461: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win🟡 Nit: Add a regression test for the non-SGLang request path.
Pass a non-SGLang
ProtoGenerateRequestvariant while the selected worker runtime isRuntimeType::Sglang. Assert that the helper returns without a panic and leaves the request unchanged.As per coding guidelines, run the
pr-test-analyzeragent to verify coverage for 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/common/stages/helpers.rs` around lines 455 - 461, The defensive non-SGLang branch in the PD bootstrap metadata injection helper lacks regression coverage. Add a test that selects RuntimeType::Sglang, passes a non-Sglang ProtoGenerateRequest, verifies the helper completes without panicking, and confirms the request remains unchanged; also run the pr-test-analyzer agent to validate coverage.Source: Coding guidelines
model_gateway/src/routers/grpc/router.rs (1)
827-840: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win🟡 Nit: Add a regression test for ZMQ worker counts.
Register one ZMQ prefill worker and one ZMQ decode worker in
Mode::PrefillDecode. Format theGrpcRouterwithDebug. Assert that both reported counts include the ZMQ workers.As per coding guidelines, run the
pr-test-analyzeragent to verify coverage for 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/router.rs` around lines 827 - 840, Add a regression test covering the GrpcRouter debug output with one ZMQ prefill worker and one ZMQ decode worker registered in Mode::PrefillDecode. Format GrpcRouter with Debug and assert both prefill_workers_count and decode_workers_count include the ZMQ workers; then run the pr-test-analyzer agent to verify coverage.Source: Coding guidelines
model_gateway/src/routers/grpc/common/stages/request_execution.rs (1)
28-49: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win🟡 Nit: Add mixed-transport tests for PD error metrics.
Exercise parallel PD with a ZMQ prefill worker and a gRPC decode worker. Exercise sequential PD with the opposite transport assignment. Force prefill startup, prefill stream, and decode startup failures. Assert that each failure uses the selected leg's
CONNECTION_ZMQorCONNECTION_GRPClabel.As per coding guidelines, run the
pr-test-analyzeragent to verify coverage for changed functionality.Also applies to: 510-518, 528-532, 671-681, 700-704, 789-795
🤖 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/common/stages/request_execution.rs` around lines 28 - 49, Extend the request-execution metric tests around pd_leg_labels to cover mixed transports: parallel PD with ZMQ prefill and gRPC decode, and sequential PD with the reverse assignment. Force prefill startup, prefill stream, and decode startup failures, asserting each metric uses the failed leg’s CONNECTION_ZMQ or CONNECTION_GRPC label. 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.
Inline comments:
In `@model_gateway/src/routers/grpc/common/stages/encode.rs`:
- Around line 323-331: Update select_encode_prefill_decode_workers and/or
assign_encode_workers to reject every direct-ZMQ TokenSpeed encode worker before
EPD planning, not only a direct-ZMQ prefill client. Exclude TokenSpeed encode
workers whose transport lacks the required encode dispatch path, or validate
each assigned encode worker so multimodal items cannot reach dispatch() with an
unsupported worker; retain the existing gRPC TokenSpeed prefill behavior.
In `@model_gateway/src/routers/grpc/proto_wrapper.rs`:
- Around line 1974-1979: Update the ZMQ response-consumption paths in
harmony/streaming.rs, regular/streaming.rs, and common/response_formatting.rs to
dispatch using the worker’s runtime_type() or an equivalent ZMQ-backed flag, not
ProtoGenerateResponse::is_vllm(). Ensure ZMQ-backed TokenSpeed responses avoid
vLLM accumulation and formatting logic while preserving existing behavior for
other runtimes.
In `@model_gateway/src/worker/manager.rs`:
- Around line 867-879: Update the empty-capable-worker handling in the
cache-flush flow around the workers.is_empty() early return to construct the
final result message before returning, including the count of skipped ZMQ
workers. Ensure both the returned message and corresponding log identify nonzero
skipped ZMQ workers, including when all registered workers are ZMQ, while
preserving the existing behavior for flush-capable workers.
---
Nitpick comments:
In `@model_gateway/src/routers/grpc/common/stages/helpers.rs`:
- Around line 455-461: The defensive non-SGLang branch in the PD bootstrap
metadata injection helper lacks regression coverage. Add a test that selects
RuntimeType::Sglang, passes a non-Sglang ProtoGenerateRequest, verifies the
helper completes without panicking, and confirms the request remains unchanged;
also run the pr-test-analyzer agent to validate coverage.
In `@model_gateway/src/routers/grpc/common/stages/request_execution.rs`:
- Around line 28-49: Extend the request-execution metric tests around
pd_leg_labels to cover mixed transports: parallel PD with ZMQ prefill and gRPC
decode, and sequential PD with the reverse assignment. Force prefill startup,
prefill stream, and decode startup failures, asserting each metric uses the
failed leg’s CONNECTION_ZMQ or CONNECTION_GRPC label. Run the pr-test-analyzer
agent to verify coverage.
In `@model_gateway/src/routers/grpc/router.rs`:
- Around line 827-840: Add a regression test covering the GrpcRouter debug
output with one ZMQ prefill worker and one ZMQ decode worker registered in
Mode::PrefillDecode. Format GrpcRouter with Debug and assert both
prefill_workers_count and decode_workers_count include the ZMQ workers; then run
the pr-test-analyzer agent to verify coverage.
🪄 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: fee4b990-1b1f-424c-bbf9-95650fea71ae
📒 Files selected for processing (7)
model_gateway/src/routers/grpc/common/stages/encode.rsmodel_gateway/src/routers/grpc/common/stages/helpers.rsmodel_gateway/src/routers/grpc/common/stages/request_execution.rsmodel_gateway/src/routers/grpc/common/stages/worker_selection.rsmodel_gateway/src/routers/grpc/proto_wrapper.rsmodel_gateway/src/routers/grpc/router.rsmodel_gateway/src/worker/manager.rs
| // Every ZMQ engine (including TokenSpeed) emits vllm-shaped | ||
| // responses: the adapter translates wire output into | ||
| // `vllm::GenerateResponse`, so variant checks like | ||
| // `is_tokenspeed()` on a response are unreliable for ZMQ-backed | ||
| // streams. Key response-side engine logic on the worker's | ||
| // `runtime_type()`, never on the response variant. |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/usr/bin/env bash
set -euo pipefail
rg -n -C 6 \
'is_tokenspeed|ProtoResponseVariant::|into_response\(\)' \
model_gateway/src || trueRepository: smg-project/smg
Length of output: 50372
🏁 Script executed:
#!/usr/bin/env bash
set -euo pipefail
echo "Changed files:"
git diff --name-only || true
echo
echo "proto_wrapper proto_variant definitions/context:"
rg -n -C 8 'enum ProtoGenerateResponse|is_tokenspeed|ProtoResponseVariant|into_response|self::ProtoGenerateResponse::|ProtoGenerateResponse::' model_gateway/src/routers/grpc/proto_wrapper.rs || true
echo
echo "Search response variant engine dispatches across model_gateway:"
rg -n 'is_tokenspeed\(\)|is_vllm\(\)|ProtoGenerateResponse::|runtime_type\(\)' model_gateway/src \
--glob '*.rs' \
|| true
echo
echo "ProtoStream Zmq implementation context:"
rg -n -C 10 'enum ProtoStream|ProtoStream::Zmq|Zmq|ProtoGenerateResponse::Vllm|generate_stream|into_stream' model_gateway/src/routers/grpc/proto_wrapper.rs || trueRepository: smg-project/smg
Length of output: 17467
🏁 Script executed:
#!/usr/bin/env bash
set -euo pipefail
echo "Response variant dispatch sites with engine branching:"
for f in \
model_gateway/src/routers/grpc/harmony/streaming.rs \
model_gateway/src/routers/grpc/regular/streaming.rs \
model_gateway/src/routers/grpc/common/response_formatting.rs
do
if [ -f "$f" ]; then
echo "===== $f ====="
rg -n -C 12 'is_vllm\(\)|is_tokenspeed\(\)|ProtoGenerateStreamChunk|ProtoGenerateComplete|runtime_type' "$f" \
|| true
fi
done
echo
echo "Callers of ProtoStream next and runtime_type context:"
rg -n -C 8 'ProtoStream::|proto_stream|\.next\(\)|runtime_type\(\)|disaggregated_runtime_type' model_gateway/src/routers/grpc \
--glob '*.rs' \
| sed -n '1,260p'Repository: smg-project/smg
Length of output: 30556
🔴 Important: ZMQ response consumers still branch by response variant.
ProtoStream::Zmq wraps TokenSpeed as ProtoGenerateResponse::Vllm, but current ZMQ response formatting paths in harmony/streaming.rs, regular/streaming.rs, and common/response_formatting.rs still treat chunks by is_vllm(), which routes ZMQ TokenSpeed chunks through vLLM accumulation logic. Use runtime_type() or a ZMQ-backed flag for response-side engine dispatch instead of the response variant.
🤖 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 1974 - 1979,
Update the ZMQ response-consumption paths in harmony/streaming.rs,
regular/streaming.rs, and common/response_formatting.rs to dispatch using the
worker’s runtime_type() or an equivalent ZMQ-backed flag, not
ProtoGenerateResponse::is_vllm(). Ensure ZMQ-backed TokenSpeed responses avoid
vLLM accumulation and formatting logic while preserving existing behavior for
other runtimes.
Source: Coding guidelines
|
Addressed all review comments in eb76224:
|
…g, tests - Exclude direct-ZMQ workers from the EPD encode pool: encode dispatch is a gRPC encoder RPC, so a ZMQ TokenSpeed encode worker would be assigned items only to fail at dispatch against its ipc:// URL. - FlushCacheResult gains a zmq_workers field (serde-defaulted, so the wire format stays backward compatible) restoring the exact total = http + grpc + zmq invariant in the JSON response, and the all-ZMQ early return now reports the skipped workers in message + log. - Regression tests: non-SGLang PD bootstrap injection no longer panics, GrpcRouter debug counts include ZMQ PD workers, and pd_leg_labels yields each PD leg's own transport label. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
eb76224 to
e08e75e
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
1 participant
Motivation
An architecture audit of the direct-ZMQ stack (post #2054–#2059) found a set of spots where code assumes every gRPC-pipeline worker speaks gRPC — so ZMQ workers get miscounted, mislabeled in metrics, or hit panicking accessors. This is the small mechanical batch (Phase 0); the capability-model refactor that removes the underlying pattern follows after the stack lands.
Modifications
as_metric_label()) instead of hardcodedgrpc.flush_cache_allpartitions ZMQ workers out of the fan-out (they have no cache-flush RPC) instead of reporting them all as failed; counts and message updated.GrpcRouter's PD/EPD debug output counts all gRPC-pipeline workers (gRPC + ZMQ) rather than onlyConnectionMode::Grpc.inject_sglang_bootstrap_metadatamatches the proto variant with alet-else+ warn instead of the panickingas_sglang_mutaccessor.ProtoStream::Zmqseam now documents that every ZMQ engine emits vllm-shaped responses, so engine-specific response logic must key onruntime_type(), never the response variant.Checklist
cargo fmtcargo test -p smg --lib(1431 passed; the one failure,middleware::metrics::distinct_ids_on_matched_route_do_not_grow_interner, fails identically on clean main under the full parallel suite — pre-existing flake, passes in isolation)