feat(grpc): support batched completion prompts - #1904
Conversation
Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>
📝 WalkthroughWalkthroughAdds gRPC support for batched completion prompts by executing each prompt independently and merging non-streaming or streaming responses. Regular and PD routers attempt batch handling before using the existing scalar completion path. CI now validates FFmpeg-dependent vLLM imports. ChangesgRPC completion batching
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant GrpcRouter
participant BatchExecutor
participant CompletionPipeline
participant SSEStream
Client->>GrpcRouter: CompletionRequest with prompt array
GrpcRouter->>BatchExecutor: execute_batch(request)
BatchExecutor->>CompletionPipeline: execute scalar request per prompt
CompletionPipeline-->>BatchExecutor: completion responses or streams
BatchExecutor->>SSEStream: merge indexed events and usage
SSEStream-->>Client: merged completion response and single [DONE]
Possibly related PRs
Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ 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 |
github-actions
Bot
added
grpc
|
@coderabbitai review |
✅ Action performedReview finished.
|
![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 introduces batching support for completion requests by splitting requests with multiple prompts into individual single-prompt requests, executing them concurrently, and merging their responses for both streaming and non-streaming modes. A critical correctness issue was identified in the streaming merge logic, where SSE stream chunk fragmentation is not handled. This can cause partial chunks to fail JSON parsing and bypass the re-indexing and metadata merging logic. The reviewer provided a detailed suggestion to buffer incoming bytes and split them into complete lines before processing.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
vLLM 0.25 eagerly imports torchcodec, which dlopens FFmpeg shared libraries at import; the e2e runner image ships none, so every vLLM worker died importing vllm. Install distro FFmpeg in the vLLM setup and add an import canary so a missing-lib regression fails at install rather than mid-test. Signed-off-by: lightseek-bot <243258330+lightseek-bot@users.noreply.github.com>
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 `@scripts/ci_install_vllm.sh`:
- Around line 58-63: Update the dependency installation step in
scripts/ci_install_vllm.sh to explicitly install torchcodec before the import
canary executes. Keep the existing python3 import check for torch, torchcodec,
and vllm unchanged so it validates all required imports.
🪄 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: cc5c8d51-1d06-48f2-a48b-86cb7dd02161
📒 Files selected for processing (1)
scripts/ci_install_vllm.sh
![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: c03c00e67f
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| # Import canary: fail here (not mid-e2e) if vLLM's eager torchcodec import | ||
| # can't find the FFmpeg shared libs installed above (torch first so its | ||
| # bundled CUDA libraries are loaded) | ||
| python3 -c "import torch, torchcodec, vllm" |
There was a problem hiding this comment.
When the resolver selects an allowed vLLM version before 0.25 from the still-wide vllm>=0.22.1 install spec, torchcodec is not guaranteed to be installed, so this unconditional canary fails setup even though those vLLM versions do not need the FFmpeg/torchcodec path. Either raise the vLLM lower bound to the first version that requires torchcodec or make this import conditional on torchcodec being present.
Useful? React with 👍 / 👎.
Local review summaryI do not think this is ready to merge yet. The happy-path merge behavior is coherent and the focused tests pass, but there is one resource-safety issue that I consider a merge blocker. Merge blocker
Other correctness and resilience concerns
Validation
|
|
Hi @lucifer1004, this PR has merge conflicts that must be resolved before it can be merged. Please rebase your branch: git fetch origin main
git rebase origin/main
# resolve any conflicts, then:
git push --force-with-lease |
mergify
Bot
added
the
needs-rebase
|
This pull request has been automatically marked as stale because it has not had any activity within 14 days. It will be automatically closed if no further activity occurs within 16 days. Leave a comment if you feel this pull request should remain open. Thank you! |
3 participants
Description
Problem
CompletionRequestacceptsprompt: string[], but the gRPC completion pipeline rejects every prompt array. This blocks OpenAI-compatible batched completions for gRPC-backed engines such as TokenSpeed.Solution
Fan out prompt arrays at the gRPC router boundary because each backend
GenerateRequestcarries one prompt token sequence. Merge the scalar results back into one OpenAI-compatible response for both streaming and non-streaming requests.The scalar completion path and the HTTP router remain unchanged. The shared helper is used by regular, PD, and EPD gRPC routers.
Closes #1903
Changes
n,echo, stops, suffixes, and streaming;[DONE];Test Plan
cargo test -p smg completion_batchn=2produce global non-streaming indices[0, 1, 2, 3]and summed usage;[0, 2], summed usage, one response ID, and one[DONE]event.cargo testcargo +1.95 clippy --workspace --all-targets -- -D warningsSKIP=clippy prek run --all-filesLocal
cargo clippy --all-targets --all-features -- -D warningsis pending because this machine does not have system OpenCV. On stable 1.97, currentmainalso trips the pre-existingunneeded_wildcard_patternlint inworker/monitor.rs:744; neither condition is introduced by this PR.Checklist
cargo +nightly fmtpassescargo clippy --all-targets --all-features -- -D warningspassesSummary by CodeRabbit