Repository navigation
fix(pd): abort both PD requests when one side hits a transport error - #844
Conversation
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request addresses a critical issue where PD disaggregation requests could hang indefinitely if one of the concurrent prefill or decode operations failed due to a transport error. By switching to Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
|
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 (1)
📝 WalkthroughWalkthroughReplaced Changes
Sequence Diagram(s)sequenceDiagram
participant Client as Client
participant Router as PD Router
participant Prefill as Worker Prefill
participant Decode as Worker Decode
Client->>Router: /generate request
Router->>Prefill: prefill_request (concurrent)
Router->>Decode: decode_request (concurrent)
alt transport error from one worker
Prefill--xRouter: transport error
Router->>Router: log transport error\nskip worker record_outcome
Router->>Client: 502 Bad Gateway (transport error)
else both requests complete
Prefill-->>Router: prefill_resp
Decode-->>Router: decode_resp
Router->>Client: aggregated success response
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
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
The pull request successfully addresses the issue of PD disaggregation requests hanging by replacing tokio::join! with tokio::try_join!. This change correctly ensures that if one side of the concurrent requests encounters a transport error, the other is immediately cancelled, preventing indefinite waits and improving the robustness of the system. The error handling and outcome recording for both workers in case of a transport error are well implemented.
![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: 6e9faa09d2
ℹ️ 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".
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@model_gateway/src/routers/http/pd_router.rs`:
- Around line 599-602: The transport error path currently calls
prefill.record_outcome(false) and decode.record_outcome(false) and then returns
error::bad_gateway, which causes the caller to double-count failures (the caller
also records outcomes after this function returns); remove the two outcome
recordings from this error branch so only the caller records failures—i.e.,
delete the calls to prefill.record_outcome(false) and
decode.record_outcome(false) in the block that logs "PD request transport error,
both sides aborted" and keep the error::bad_gateway return.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: b5bba8ff-4690-44d4-8244-ba9828cd169c
📒 Files selected for processing (1)
model_gateway/src/routers/http/pd_router.rs
6e9faa0 to
67bba3e
Compare
There was a problem hiding this comment.
♻️ Duplicate comments (1)
model_gateway/src/routers/http/pd_router.rs (1)
596-599:⚠️ Potential issue | 🟠 MajorRemove duplicate failure recording on transport-error path.
Line 597 and Line 598 already mark both workers failed, and the caller records failures again at Line 360 and Line 361 after this function returns
502. This double-counts a single transport failure.Proposed fix
Err(e) => { error!("PD request transport error, both sides aborted: {e}"); - prefill.record_outcome(false); - decode.record_outcome(false); return error::bad_gateway( "PD disaggregation request failed", format!("Transport error: {e}"), ); }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/routers/http/pd_router.rs` around lines 596 - 599, The transport-error branch in pd_router.rs currently calls prefill.record_outcome(false) and decode.record_outcome(false) before returning error::bad_gateway, which duplicates failure recording already done by the caller; remove those two record_outcome(false) calls (the ones immediately before the return of error::bad_gateway) so that failures are recorded only once by the caller and not double-counted.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Duplicate comments:
In `@model_gateway/src/routers/http/pd_router.rs`:
- Around line 596-599: The transport-error branch in pd_router.rs currently
calls prefill.record_outcome(false) and decode.record_outcome(false) before
returning error::bad_gateway, which duplicates failure recording already done by
the caller; remove those two record_outcome(false) calls (the ones immediately
before the return of error::bad_gateway) so that failures are recorded only once
by the caller and not double-counted.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 247f4b01-7e5b-4ce0-a2a1-7c23e1409359
📒 Files selected for processing (1)
model_gateway/src/routers/http/pd_router.rs
When prefill and decode requests are sent concurrently with tokio::join!, if one side fails with a transport error (e.g., connection reset), the other side hangs forever waiting for a PD bootstrap that will never come. Switch from tokio::join! to tokio::try_join! so that when either request returns Err, the other future is dropped (cancelled). Both sides are recorded as failed and a bad_gateway response is returned immediately. The health check path (health_generate) still uses tokio::join! since health checks are quick and don't involve bootstrap. Closes #831 Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
67bba3e to
7411c52
Compare
1 participant
Summary
Fixes PD disaggregation requests hanging forever when one side hits a transport error.
Closes #831
Problem
tokio::join!waits for both prefill and decode requests to complete. If prefill fails with a connection reset, the decode request hangs waiting for a PD bootstrap from prefill that will never come — until the bootstrap timeout (default 60s).Fix
Replace
tokio::join!withtokio::try_join!. When either request returnsErr(transport error),try_join!drops the other future, cancelling the in-flight request immediately. Both workers are recorded as failed (circuit breaker), and a502 Bad Gatewayis returned.What changed
execute_pd_request_onceusestokio::try_join!instead oftokio::join!for the prefill+decode concurrent send. On transport error, returnsbad_gatewayimmediately.health_generate) still usestokio::join!— health checks don't involve bootstrap.Test plan
cargo check -p smgpassescargo clippy -p smg -- -D warningspassesSummary by CodeRabbit