feat(grpc): support batched completion prompts in the pipeline - #1957
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 (12)
💤 Files with no reviewable changes (1)
📝 WalkthroughWalkthroughgRPC-backed ChangesgRPC completion batching
Estimated code review effort: 4 (Complex) | ~60 minutes Possibly related PRs
Suggested reviewers: Poem
Sequence Diagram(s)sequenceDiagram
participant Client
participant CompletionPreparationStage
participant RequestExecutionStage
participant BackendWorkers
participant CompletionResponseProcessor
Client->>CompletionPreparationStage: submit prompt array
CompletionPreparationStage->>RequestExecutionStage: create per-prompt batch plan
RequestExecutionStage->>BackendWorkers: concurrently dispatch prompt requests
BackendWorkers-->>RequestExecutionStage: return per-prompt results
RequestExecutionStage-->>CompletionResponseProcessor: provide ExecutionResult::Batch
CompletionResponseProcessor-->>Client: return indexed choices and aggregated usage
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
![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: 1409671c39
ℹ️ 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 dispatches = requests.into_iter().map(|request| { | ||
| let mut clients = clients.clone(); |
There was a problem hiding this comment.
When prompt is an array, this branch starts one backend Generate RPC per prompt, but RequestExecutionStage::execute still creates only one LoadGuards::new(...) before entering the match. In load-aware deployments (least_load, power_of_two, prefix-hash load checks), a large prompt batch therefore looks like a single active request while it is actually occupying the selected worker with many concurrent backend streams, so subsequent traffic can keep selecting an already overloaded worker. Consider acquiring/releasing a guard per sub-request or otherwise scaling the in-flight load by requests.len() for this batch path.
Useful? React with 👍 / 👎.
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| }); | ||
|
|
||
| let mut results = Vec::new(); | ||
| for result in join_all(dispatches).await { |
There was a problem hiding this comment.
🟡 Nit: join_all waits for all dispatches to complete before any error is checked, so this isn't truly fail-fast — if the first dispatch errors in 1 ms but the second takes 30 s, we still wait the full 30 s. try_join_all (already used in preparation.rs) would drop remaining futures on first Err, actually cancelling in-flight dispatches and matching the doc comment's intent:
| for result in join_all(dispatches).await { | |
| let results: Vec<_> = try_join_all(dispatches).await?; |
(with use futures::future::try_join_all; replacing the join_all import)
The practical impact is small since dispatches are typically fast (establishing connections, not consuming streams), but it's a cleaner fit for the stated semantics.
There was a problem hiding this comment.
Thorough, well-structured implementation. The batch state is mode-agnostic as described, scalar requests keep their exact plan shape, and the streaming coordinator/per-stream-unit split is clean. Tests cover the key contracts (index mapping, usage aggregation, echo, n>1).
One minor nit on join_all vs try_join_all in execute_batch_dispatch — see inline comment.
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: 5
🤖 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/completions/test_basic.py`:
- Around line 275-294: The streaming test helper and its callers must detect
duplicate usage chunks and exercise multi-choice global indexing. Update
_collect_stream_by_index to count usage-bearing chunks and assert exactly one
was emitted, change the relevant streaming requests from n=1 to n=2, and require
choice indices 0 through 3 to validate prompt_index * n offsets while preserving
the existing text and finish-reason assertions.
- Around line 320-333: Strengthen test_batch_non_streaming_with_n to enable echo
and verify prompt-major ordering directly: require response.choices indices in
response order to be [0, 1, 2, 3], and assert choices 0–1 contain the first
prompt while choices 2–3 contain the second. Replace the sorted-index-only
assertion while preserving the existing count and n=2 setup.
In `@model_gateway/src/routers/grpc/common/stages/request_execution.rs`:
- Around line 339-368: Update execute_batch_dispatch and the surrounding execute
load-guard flow so each fanned-out sub-request contributes one unit of worker
load, rather than the entire batch counting as one request. Reuse the existing
guard mechanism by creating a guard per dispatched request or weighting it by
requests.len(), while preserving the current batch result and error propagation
behavior.
In `@model_gateway/src/routers/grpc/regular/processor.rs`:
- Around line 844-859: Update the finish_reason handling in the processor’s
completion branch to mirror process_completion_streaming_chunks: preserve empty,
"stop", "length", and supported "content_filter" values, but coerce unknown raw
reasons and unknown JSON type values to "stop" instead of passing them through.
Keep the existing JSON parsing behavior while ensuring non-streaming responses
match streaming responses.
In `@model_gateway/tests/routing/grpc_completion_batch_test.rs`:
- Around line 121-128: Update the test helper completion_request to serialize
its single prompt as a scalar string rather than an array, while preserving the
existing model and max_tokens fields. Ensure the one-item request exercises
scalar completion behavior, and apply the same correction to the additional
request construction around the referenced batch test.
🪄 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: 8bcc36b8-f984-4dfd-9c42-bacfb2b89f55
📒 Files selected for processing (12)
e2e_test/completions/test_basic.pymodel_gateway/src/routers/grpc/common/response_collection.rsmodel_gateway/src/routers/grpc/common/stages/request_execution.rsmodel_gateway/src/routers/grpc/context.rsmodel_gateway/src/routers/grpc/harmony/streaming.rsmodel_gateway/src/routers/grpc/regular/processor.rsmodel_gateway/src/routers/grpc/regular/stages/completion/preparation.rsmodel_gateway/src/routers/grpc/regular/stages/completion/request_building.rsmodel_gateway/src/routers/grpc/regular/stages/completion/response_processing.rsmodel_gateway/src/routers/grpc/regular/streaming.rsmodel_gateway/tests/routing/grpc_completion_batch_test.rsmodel_gateway/tests/routing/mod.rs
💤 Files with no reviewable changes (1)
- model_gateway/src/routers/grpc/regular/stages/completion/response_processing.rs
| let finish_reason = { | ||
| let reason = complete.finish_reason(); | ||
| if reason.is_empty() { | ||
| None | ||
| } else if reason == "stop" || reason == "length" { | ||
| Some(reason.to_string()) | ||
| } else if let Ok(json) = serde_json::from_str::<serde_json::Value>(reason) { | ||
| json.get("type").and_then(|v| v.as_str()).map(|s| match s { | ||
| "length" => "length".to_string(), | ||
| "stop" => "stop".to_string(), | ||
| other => other.to_string(), | ||
| }) | ||
| } else { | ||
| Some(reason.to_string()) | ||
| } | ||
| }; |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Unmapped finish_reason values pass through here but are coerced to "stop" in the streaming path.
For unknown finish-reason strings (or unknown JSON type values) this non-streaming branch emits other.to_string() / reason.to_string() verbatim, whereas process_completion_streaming_chunks maps anything outside stop/length/content_filter to "stop". This produces non-OpenAI-spec finish_reason values and diverges from the streaming response for the same backend output. Consider mirroring the streaming coercion for consistency.
🤖 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/regular/processor.rs` around lines 844 - 859,
Update the finish_reason handling in the processor’s completion branch to mirror
process_completion_streaming_chunks: preserve empty, "stop", "length", and
supported "content_filter" values, but coerce unknown raw reasons and unknown
JSON type values to "stop" instead of passing them through. Keep the existing
JSON parsing behavior while ensuring non-streaming responses match streaming
responses.
1409671 to
22cdf08
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 22cdf086e4
ℹ️ 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".
| for result in join_all(dispatches).await { | ||
| results.push(result?); |
There was a problem hiding this comment.
For batched prompts where one backend Generate start returns an error while another sub-dispatch is still queued or slow to establish, join_all(dispatches).await waits for every dispatch future to finish before this loop observes the first Err. That means the client can wait until the slow sibling times out, and any streams that already started are not dropped/aborted promptly, despite the intended fail-fast behavior. Consider using a fail-fast combinator such as try_join_all (or an equivalent ordered collection that drops pending futures on the first error).
Useful? React with 👍 / 👎.
22cdf08 to
cb2d86c
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: cb2d86ce86
ℹ️ 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".
| StringOrArray::Array(texts) => texts.iter().map(String::as_str).collect(), | ||
| }; | ||
|
|
||
| let outcomes = join_all(units.into_iter().enumerate().map(|(prompt_index, unit)| { |
There was a problem hiding this comment.
For streaming prompt arrays, this join_all does not inspect any per-prompt Err until every sibling stream has finished. When one backend stream fails early while another prompt is still generating or stalls, the client will keep waiting on the remaining stream and will not receive the error SSE until that sibling ends, instead of failing the request promptly as the single-prompt streaming path does. Consider using a fail-fast collection/select that sends the error and drops the remaining stream units.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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/regular/processor.rs`:
- Around line 780-798: Update process_non_streaming_completion_response to
collect all per-prompt responses concurrently rather than awaiting
collect_responses sequentially inside the prompt loop. Mirror
process_completion_streaming_response by creating one collection future per
prompt, awaiting them together with join_all (or the existing concurrency
helper), then process the collected results in prompt order while preserving
current response construction and error handling.
🪄 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: 7663582b-b739-4d9e-a7c0-d563cb06b27e
📒 Files selected for processing (12)
e2e_test/completions/test_basic.pymodel_gateway/src/routers/grpc/common/response_collection.rsmodel_gateway/src/routers/grpc/common/stages/request_execution.rsmodel_gateway/src/routers/grpc/context.rsmodel_gateway/src/routers/grpc/harmony/streaming.rsmodel_gateway/src/routers/grpc/regular/processor.rsmodel_gateway/src/routers/grpc/regular/stages/completion/preparation.rsmodel_gateway/src/routers/grpc/regular/stages/completion/request_building.rsmodel_gateway/src/routers/grpc/regular/stages/completion/response_processing.rsmodel_gateway/src/routers/grpc/regular/streaming.rsmodel_gateway/tests/routing/grpc_completion_batch_test.rsmodel_gateway/tests/routing/mod.rs
💤 Files with no reviewable changes (1)
- model_gateway/src/routers/grpc/regular/stages/completion/response_processing.rs
Accept prompt arrays for /v1/completions natively in the gRPC pipeline: preparation tokenizes each prompt, request building emits one backend request per prompt under a shared client-visible id, execution dispatches them concurrently in every mode, and response processing merges typed results with prompt-major global choice indices, summed usage, one aggregated streaming usage chunk, and a single [DONE]. Closes #1903 Co-authored-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
cb2d86c to
c4e7e9e
Compare
…pagation Replays the codex/fix-glm-tool-parser work as a single commit on top of main. - glm4_moe: prefix (not substring) model resolution with basename matching; reject lone '</' at EOF; gate streaming tool-name validation on !tools.is_empty(); exclude argument strings from orphan marker detection; accept short literal marker prefixes in the complete path; register glm45/glm47 aliases. - grpc/regular: propagate tool-parse errors in the Messages API and preserve streaming logprobs; use the Anthropic error format for Messages streaming tool-parse failures; fix tool-call boundary leak and redundant model lookup; extract shared test fakes into a test_fakes module. - sglang servicer: prevent gRPC requests from failing before scheduling and ensure upstream output failures can never be reported as success. factory.rs merges cleanly with main's Sarashina parser registration; processor.rs and streaming.rs merge with main's batched completion streaming (smg-project#1957), preserving both the batched fan-out and the branch's error propagation (bad_gateway / upstream_output_parse_failed). Co-Authored-By: Claude <noreply@anthropic.com>
1 participant
Description
Problem
CompletionRequestacceptsprompt: string[], but the gRPC completion pipeline rejected every array inCompletionPreparationStagewithbatch_prompts_not_supported. This blocks OpenAI-compatible batched completions for gRPC-backed engines (#1903).Solution
Support batch natively inside the Mode-parameterized completion pipeline, rather than fanning out at the router boundary as #1904 did (that approach re-parsed the router's own serialized SSE/JSON at the edge — with SSE chunk-fragmentation hazards — multiplied retry/metrics per sub-request, and predates the #1923 router unification):
PreparationOutput::Completion; multi-prompt routing text joins the prompts, mirroring the HTTP router'sextract_text_for_routing.GenerateRequestper prompt under one shared client-visiblecmpl_id; backend sub-ids get a-p{i}suffix so abort-on-drop and PD correlation stay unique, and PD bootstrap rooms are minted per sub-request.ExecutionPlan::Batch/ExecutionResult::Batch: sub-requests dispatch concurrently with the mode's shape (single worker, or PD dual dispatch — SGLang/TokenSpeed parallel, vLLM sequential KV relay). Fail-fast: the first failed dispatch fails the batch and dropped streams abort backend-side.CompletionResponsewith prompt-major global indices (prompt_index * n + choice_index) and usage summed across prompts (per-prompt max prompt tokens); streaming drives all per-prompt chunk loops into one SSE channel with offset indices and emits a single aggregated usage chunk, one metrics record, and one[DONE]. Per-promptechofalls out naturally.Scalar requests keep their exact plan shape and wire behavior (a one-element array now behaves identically to a scalar prompt), and the whole thing works in Regular, PD, and EPD because the batch state is mode-agnostic.
Within a prompt, non-streaming choice indices keep the pre-existing arrival-order semantics (SGLang non-streaming Complete frames carry index 0 for every choice, so backend-reported indices are not usable there); the batch offset is applied on top.
Supersedes #1904; keeps its global-index/usage contract and test intent. Co-authored with @lucifer1004.
Closes #1903.
Changes
grpc/context.rs—CompletionItem, list-shapedPreparationOutput::Completion(+ batch-awarerouting_text()/token_ids()),ExecutionPlan::Batch,ExecutionResult::Batch,ClientSelection: Clone.grpc/regular/stages/completion/preparation.rs— accept arrays; concurrent per-prompt tokenization (rejection removed).grpc/regular/stages/completion/request_building.rs— per-prompt proto build + PD injection under a shared id.grpc/common/stages/request_execution.rs— concurrentBatchdispatch arm.grpc/regular/processor.rs,grpc/regular/streaming.rs— typed batch merge; streaming usage/metrics/[DONE]centralized in a coordinator (scalar event order unchanged); the stale arraydebug_assertin streaming echo is gone.grpc/common/response_collection.rs,grpc/harmony/streaming.rs, chat/generate/messages streaming entries — defensiveBatcharms.tests/routing/grpc_completion_batch_test.rs(real tokenizer; arrays pass preparation and hit the same worker-selection wall as scalars in Regular/PD/EPD), context unit tests for the new plan/prep contracts, and e2eTestCompletionBatch(e2e_test/completions/test_basic.py): non-streaming indices/usage,n>1prompt-major indices, per-prompt echo mapping, streaming per-index deltas + aggregated usage — across sglang/vllm/tokenspeed.Test Plan
cargo test -p smg: lib 1225 passed; all integration binaries green (api_tests 106, routing_tests 94, spec_test 96, security_tests 50, reliability_tests 26, mcp_test 23, wasm_test 17; 0 failures).cargo clippy --all-targets -- -D warnings,cargo +nightly fmt --check,ruff check/ruff format --checkon the e2e file.pytest e2e_test/completions/test_basic.py -k TestCompletionBatch(1 GPU; sglang, orE2E_RUNTIME=vllm).Checklist
cargo +nightly fmtpassescargo clippy --all-targets --all-features -- -D warningspassesSummary by CodeRabbit