fix(pd): route DP logical workers via base endpoint - #1522
Conversation
Signed-off-by: Aurick Qiao <aurick@thinkingmachines.ai>
|
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)
📝 WalkthroughWalkthroughThe PD proxy's dual-dispatch logic was refactored to compute prefill and decode DP-rank metadata independently and inject them selectively into JSON payloads. Request construction for downstream legs now uses the selected ChangesDP-rank metadata and worker-based request routing
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Suggested labels
Suggested reviewers
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 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 |
![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
This pull request refactors the PDRouter to improve the handling of data parallel (DP) ranks and updates the request construction logic to utilize the Worker trait. Specifically, it introduces explicit tracking of prefill and decode ranks and replaces the api_path utility with worker.endpoint_url. Feedback highlights a potential regression in URL path joining that could lead to double slashes and suggests verifying that the injected JSON keys for DP ranks align with external protocol specifications.
| if let Some(p_rank) = prefill_rank { | ||
| Self::inject_dp_rank_to_json( | ||
| &mut prefill_json_request, | ||
| p_rank, | ||
| "routed_dp_rank", | ||
| ); | ||
| Self::inject_dp_rank_to_json( | ||
| &mut decode_json_request, | ||
| p_rank, | ||
| "disagg_prefill_dp_rank", | ||
| ); | ||
| } | ||
| if let Some(d_rank) = decode_rank { | ||
| Self::inject_dp_rank_to_json( | ||
| &mut decode_json_request, | ||
| d_rank, | ||
| "routed_dp_rank", | ||
| ); | ||
| } |
There was a problem hiding this comment.
The PD router is injecting routed_dp_rank and disagg_prefill_dp_rank. While data_parallel_rank is the standard internal key, please verify if these names are required by the external PD disaggregation protocol. Per repository rules, alignment with external specifications takes priority over internal consistency for protocol data structures.
References
- For protocol data structures that mirror an external API, prioritize alignment with the external specification over internal consistency.
| let endpoint_url = worker.endpoint_url(route); | ||
| let mut request = client.post(endpoint_url).json(json_request); |
There was a problem hiding this comment.
The previous implementation used api_path(url, route), which safely handles leading/trailing slashes to avoid double slashes in the resulting URL. worker.endpoint_url(route) currently performs a simple concatenation: format!("{}{}", self.base_url(), route). Consider using api_path(worker.base_url(), route) to maintain the previous safety.
| let endpoint_url = worker.endpoint_url(route); | |
| let mut request = client.post(endpoint_url).json(json_request); | |
| let endpoint_url = super::pd_types::api_path(worker.base_url(), route); | |
| let mut request = client.post(endpoint_url).json(json_request); |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Ports smg-project/smg#1522 to this fork's dispatch path and current sglang field names (routed_dp_rank / disagg_prefill_dp_rank). Three related bugs in the DP-aware PD-disaggregation path, found by forcing --router-dp-aware and exercising real per-rank worker registration end-to-end: - pd_router.rs dispatched requests using Worker::url(), which is a worker's logical "@rank" identity for DP-aware workers, not a dispatchable address. Dispatch now goes through endpoint_url()/ base_url(), and decode requests are stamped with disagg_prefill_dp_rank so decode knows which prefill rank's KV cache to pull. - DPAwareWorkerBuilder derived bootstrap_host by URL-parsing the "@rank"-suffixed logical url, which parses the rank as the hostname (e.g. "http://host:port@1" resolves host as numeric "1" -> 0.0.0.1), breaking the decode->prefill bootstrap handshake. bootstrap_host is now derived from base_url instead. - WorkerInfo (the /workers API) had no way to report a DP-aware worker's real address at all; added base_url and dp_rank fields.
2 participants
Description
Problem
HTTP PD disaggregation can fail when routing to DP-aware logical workers. In that configuration, SMG represents each logical DP worker with a URL suffix such as
http://host:port@rank. The HTTP PD router used that logical URL as the network request target, so URL parsing treated the value after@as the host. This can surface as 502 responses with errors like:PD disaggregation request failedPD request transport errorerror sending request for url (http://0.0.0.0/generate)This affects HTTP PD disaggregation configurations where DP-aware worker discovery or explicit DP worker registration creates logical workers with
@rankURLs.Solution
Use the worker endpoint URL for actual HTTP dispatch. For DP-aware logical workers,
endpoint_url()uses the worker'sdp_base_url, while preserving the logical worker URL for identity and selection.The PD router now also propagates selected DP ranks into the prefill and decode request bodies so the backend receives the selected logical ranks after dispatch goes to the base endpoint.
Changes
worker.endpoint_url(route)instead of the logical worker URL.routed_dp_rankfor selected prefill and decode workers when DP rank metadata is present.disagg_prefill_dp_rankinto decode requests when the selected prefill worker has DP rank metadata.http://127.0.0.1:30000@2dispatches tohttp://127.0.0.1:30000/generate.Test Plan
cargo +nightly fmt --all --checkcargo test -p smg test_build_post_uses_dp_base_url_for_logical_workerhttp://0.0.0.0/generate; after this change, requests completed successfully and gateway logs showed selected DP ranks being propagated.Checklist
cargo +nightly fmtpassescargo clippy --all-targets --all-features -- -D warningspassesSummary by CodeRabbit
Improvements
Refactor
Tests