feat(grpc): serve /v1/responses in PD mode - #1956
Conversation
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
github-actions
Bot
added
grpc
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Run ID: 📒 Files selected for processing (3)
📝 WalkthroughWalkthroughChangesPD Responses API support
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant APIClient
participant GrpcRouter
participant MCPOrchestrator
APIClient->>GrpcRouter: create Responses request
GrpcRouter->>MCPOrchestrator: process response context
MCPOrchestrator-->>GrpcRouter: response or stream events
GrpcRouter-->>APIClient: response or events
Possibly related issues
Possibly related PRs
Suggested labels: Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
e48154c to
57e714c
Compare
![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 `@e2e_test/router/test_pd_responses.py`:
- Around line 59-63: Update the assertions in the response event test to require
a terminal response.completed event rather than accepting response.in_progress,
and additionally assert that the streamed events contain at least one
output-text delta. Keep the existing response.created validation unchanged.
In `@model_gateway/src/routers/grpc/router.rs`:
- Around line 378-379: Make GrpcRouter::new runtime-neutral by changing the
Harmony initialization through RequestPipeline::build so its encoding load does
not rely on block_in_place. Use async-compatible offloading for the blocking
work, preserving the existing unsupported (endpoint, mode) handling while
ensuring initialization works on current-thread Tokio runtimes.
🪄 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: e373df7b-bd3f-47a8-b05a-872a75cc6cf1
📒 Files selected for processing (3)
e2e_test/router/test_pd_responses.pymodel_gateway/src/routers/factory.rsmodel_gateway/src/routers/grpc/router.rs
| // `None` when the (endpoint, mode) combo is unsupported; those endpoints 501. | ||
| let harmony_pipeline = RequestPipeline::build(Endpoint::Harmony, mode, &configured_deps); |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
rg -n -C2 'tokio::main|new_current_thread|flavor\s*=\s*"current_thread"|GrpcRouter::new|RouterFactory::create_router' model_gatewayRepository: lightseekorg/smg
Length of output: 16936
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the PD router construction and any block_in_place usage around the cited area.
sed -n '340,430p' model_gateway/src/routers/grpc/router.rs
echo '---'
rg -n -C3 'block_in_place|Harmony|PrefillDecode|RequestPipeline::build' model_gateway/src/routers/grpc/router.rs model_gateway/src/routers/factory.rs model_gateway/src/health.rs model_gateway/tests -g '!**/target/**'Repository: lightseekorg/smg
Length of output: 30219
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Look for any block_in_place usage along the router construction path.
rg -n -C3 'block_in_place|spawn_blocking|RequestPipeline::build|Endpoint::Harmony|GrpcRouter::new' model_gateway/src/routers model_gateway/testsRepository: lightseekorg/smg
Length of output: 16548
model_gateway/src/routers/grpc/router.rs:379 — Make PD router init runtime-neutral. GrpcRouter::new eagerly builds Harmony here, and Harmony’s encoding load still uses block_in_place, which will panic on a current-thread Tokio runtime. Use async offloading instead of relying on multi-thread-only tests.
🤖 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 378 - 379, Make
GrpcRouter::new runtime-neutral by changing the Harmony initialization through
RequestPipeline::build so its encoding load does not rely on block_in_place. Use
async-compatible offloading for the blocking work, preserving the existing
unsupported (endpoint, mode) handling while ensuring initialization works on
current-thread Tokio runtimes.
Build the responses contexts for every mode except EPD, so PD serves the full responses surface (regular + Harmony, MCP tool loop, streaming, storage) through the existing PD pipeline. Add gRPC PD responses e2e coverage. Closes #1946 Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
57e714c to
2006b80
Compare
![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: 2006b80e9b
ℹ️ 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".
| let (responses_context, harmony_responses_context) = if mode == Mode::EncodePrefillDecode { | ||
| (None, None) | ||
| } else { |
There was a problem hiding this comment.
When --storage-context-headers or a storage hook is configured, PD /v1/responses now enters this branch during router construction, so current_request_context() below is evaluated at startup rather than under the request middleware's task-local. The resulting responses_context is then reused by the non-Harmony responses path, so persisted PD responses lose tenant/user header context and hooks cannot apply per-request scoping. Please build the ResponsesContext (or refresh its request context) inside route_responses_impl for each request, similar to the Harmony branch.
Useful? React with 👍 / 👎.
1 participant
Description
Problem
/v1/responsesreturns501 Not Implementedon the gRPC PD router (#1946). Since #1923 the gap is purely the router gate: PD shares the Mode-parameterized pipeline that already does prefill/decode pair selection, bootstrap injection, and dual dispatch for chat, and the responses stack touches workers only through that pipeline — but the responses contexts were built only forMode::Regular.Solution
Build the responses contexts in every mode except EPD. With the contexts present, the existing endpoint guards pass and the full responses surface — regular + Harmony stacks, MCP tool loop, streaming, conversation history, storage-backed retrieval/cancel — serves over PD unchanged. PD router construction now requires the MCP orchestrator, which production always initializes (
AppContext::from_config).EPD keeps its 501: the responses stack requires the Harmony pipeline, which EPD legitimately doesn't build.
Behavior note: each MCP tool-loop iteration under PD performs a fresh pair selection and bootstrap room, identical to a client-driven chat loop.
Closes #1946.
Changes
grpc/router.rs— flip the responses-context gate frommode == Regulartomode == EPD ? (None, None) : build both; mode-aware harmony-pipeline construction error; correct stale comments (some predated this PR: the struct doc claimedharmony_pipelineis Regular-only while PD has built and served it for Harmony chat since refactor(grpc): unify regular/PD/EPD into one Mode-parameterized router + EncodeStage #1923).grpc/router.rstests —pd_retry_tests→pd_tests: the fixture now initializes a realMcpOrchestrator(multi-thread runtime for the Harmony encodingblock_in_placeload, matchingfactory.rs); new tests pin PD behavior (404 for unknown model / missing response id, previously 501) and EPD's kept 501s.factory.rs— stale test-helper comment.e2e_test/router/test_pd_responses.py— newTestPDResponsesGrpcoverpd_grpc(sglang + vllm,gpu(2), Llama-3.1-8B,--history-backend memory, openai + smg clients): basic create, streaming events,previous_response_idchaining,store=false. Placed underrouter/soe2e-2gpu-pdruns it on all three legs (sglang, vllm-nixl, vllm-mooncake). First e2e coverage of responses over PD.Test Plan
cargo test -p smg: lib 1225 passed; all integration binaries green (api_tests 106, routing_tests 93, spec_test 96, security_tests 50, reliability_tests 26, mcp_test 23, wasm_test 17; 0 failures).pd_router_serves_responses_and_cancel,epd_router_501s_responses_and_cancel.cargo clippy --all-targets -- -D warnings,cargo +nightly fmt --check,ruff check/ruff format --checkon the e2e file.pytest e2e_test/router/test_pd_responses.py(2 GPUs; sglang orE2E_RUNTIME=vllm).Checklist
cargo +nightly fmtpassescargo clippy --all-targets --all-features -- -D warningspassesSummary by CodeRabbit
New Features
Bug Fixes
/v1/responsesno longer returns “not implemented”.store=falseresponses aren’t retrievable after creation.Tests