refactor(grpc): optimize vLLM servicer and remove duplicate server.py - #662
Conversation
github-actions
Bot
added
the
documentation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughRemoved the vLLM gRPC server module and its export, updated README to use the public gRPC entrypoint, and refactored servicer internals to use torch dtypes, add caching and vectorized masking, and tighten kv_transfer_params validation. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
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 refactors the vLLM gRPC servicer to significantly improve its efficiency and streamline its architecture. It addresses several performance bottlenecks by optimizing tensor deserialization, logprobs construction, and placeholder mask generation, leading to a more performant system. Furthermore, it eliminates a redundant server entrypoint by adopting the upstream vLLM implementation and centralizes validation logic for Highlights
Changelog
Activity
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. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![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 a solid set of optimizations and refactorings for the vLLM gRPC servicer, including removing the numpy dependency for tensor deserialization, vectorizing mask building, caching tensor conversions, removing the duplicate server.py, and consolidating validation logic. However, a potential Server-Side Request Forgery (SSRF) vulnerability was identified in the handling of kv_transfer_params where the remote_host is not properly validated against an allow-list or restricted to authorized internal addresses. Additionally, there is one minor suggestion for a micro-optimization to further improve the code.
![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 the current code and only fix it if needed.
Inline comments:
In `@grpc_servicer/smg_grpc_servicer/vllm/__init__.py`:
- Line 5: The change removed serve_grpc from the public exports causing a
compatibility break; restore a deprecated shim named serve_grpc in
grpc_servicer/smg_grpc_servicer/vllm/__init__.py that imports and calls the
current implementation (or forwards to VllmEngineServicer where appropriate),
add "serve_grpc" back into __all__ alongside "VllmEngineServicer", and mark the
shim with a deprecation warning (e.g., using warnings.warn) so consumers get
notified while existing imports continue to work for one release.
In `@grpc_servicer/smg_grpc_servicer/vllm/servicer.py`:
- Around line 455-459: The INFO-level logging of KV transfer parameters in the
request hot path (the logger.info call that prints
"kv_transfer_params={remote_host=%s, remote_port=%d}") should be removed or
demoted to debug and redact sensitive topology; update the logging inside the
Generate request path (in servicer.py where the kv_transfer_params logger.info
is used) to either delete the log statement or change it to logger.debug and
avoid emitting raw remote_host/remote_port (e.g., redact or replace with masked
values or a boolean flag like "kv_transfer_params_present"). Ensure the change
touches the exact logger.info invocation so the hot path no longer emits
INFO-level host/port details.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 6c5f1c15-6b48-47e7-90bc-d8f26f904ebd
📒 Files selected for processing (4)
grpc_servicer/README.mdgrpc_servicer/smg_grpc_servicer/vllm/__init__.pygrpc_servicer/smg_grpc_servicer/vllm/server.pygrpc_servicer/smg_grpc_servicer/vllm/servicer.py
💤 Files with no reviewable changes (1)
- grpc_servicer/smg_grpc_servicer/vllm/server.py
There was a problem hiding this comment.
♻️ Duplicate comments (1)
grpc_servicer/smg_grpc_servicer/vllm/servicer.py (1)
449-464:⚠️ Potential issue | 🟠 MajorKeep the validation, but drop the hot-path host/port INFO log.
The validation move is fine, but Line 455 still emits internal topology for every KV-transfer request on the Generate path. This should be debug-level at most and ideally redacted.
Suggested change
- logger.info( - "kv_transfer_params={remote_host=%s, remote_port=%d}", - remote_host, - remote_port, - ) + logger.debug("kv_transfer_params received")🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@grpc_servicer/smg_grpc_servicer/vllm/servicer.py` around lines 449 - 464, The validation for kv_transfer_params (checking kv_transfer_params.remote_host and 1 <= kv_transfer_params.remote_port <= 65535) should remain, but replace the hot-path logger.info call in servicer.py with a lower-verbosity, redacted debug log: change logger.info(...) to logger.debug(...) and avoid printing raw internal topology values (use placeholders like "<redacted>" or log only that kv_transfer_params was provided without including remote_host/remote_port), leaving extra_args construction (kv_transfer_params mapping) intact so behavior is unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Duplicate comments:
In `@grpc_servicer/smg_grpc_servicer/vllm/servicer.py`:
- Around line 449-464: The validation for kv_transfer_params (checking
kv_transfer_params.remote_host and 1 <= kv_transfer_params.remote_port <= 65535)
should remain, but replace the hot-path logger.info call in servicer.py with a
lower-verbosity, redacted debug log: change logger.info(...) to
logger.debug(...) and avoid printing raw internal topology values (use
placeholders like "<redacted>" or log only that kv_transfer_params was provided
without including remote_host/remote_port), leaving extra_args construction
(kv_transfer_params mapping) intact so behavior is unchanged.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 392b95c9-6452-4ef3-a9a6-595335705921
📒 Files selected for processing (1)
grpc_servicer/smg_grpc_servicer/vllm/servicer.py
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 `@grpc_servicer/smg_grpc_servicer/vllm/servicer.py`:
- Around line 582-585: The condition checking whether to append top_logprobs is
functionally correct but stylistically inconsistent with the check in
_build_output_logprobs; change the conditional around proto.top_logprobs (the
block appending VllmEngineServicer._build_top_logprobs) to match the simpler
style used in _build_output_logprobs (e.g. use "if num_top_logprobs:"), so both
places use the same falsy-check pattern for num_top_logprobs.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 133b7d33-bd23-46c9-94eb-fe538cb86a73
📒 Files selected for processing (1)
grpc_servicer/smg_grpc_servicer/vllm/servicer.py
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 `@grpc_servicer/smg_grpc_servicer/vllm/servicer.py`:
- Around line 42-47: The _tensor_from_proto function currently lets
torch.frombuffer()/reshape raise RuntimeError for malformed payloads, causing
server-side INTERNAL errors; wrap the deserialization/reshape logic in a
try/except that catches RuntimeError (and any torch-specific buffer/shape
errors) and re-raise them as ValueError with a clear message so clients get
INVALID_ARGUMENT; reference _tensor_from_proto, vllm_engine_pb2.TensorData and
_PROTO_DTYPE_MAP when locating the change.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 35fca93e-feb0-41b2-9393-b32c341b1625
📒 Files selected for processing (1)
grpc_servicer/smg_grpc_servicer/vllm/servicer.py
Guard _build_top_logprobs in _build_input_logprobs to avoid allocating empty proto objects for every prompt token. Vectorize placeholder mask building with torch ops. Eliminate numpy intermediate in tensor deserialization. Cache flat sizes tensor conversion. Extract _build_top_logprobs to module level. Signed-off-by: Chang Su <chang.s.su@oracle.com>
…points.grpc_server Signed-off-by: Chang Su <chang.s.su@oracle.com>
…ams_from_proto Signed-off-by: Chang Su <chang.s.su@oracle.com>
Signed-off-by: Chang Su <chang.s.su@oracle.com>
Signed-off-by: Chang Su <chang.s.su@oracle.com>
Signed-off-by: Chang Su <chang.s.su@oracle.com>
Signed-off-by: Chang Su <chang.s.su@oracle.com>
torch.topk() already returns entries sorted descending by logprob value, and Python dict insertion order preserves this (3.7+). Signed-off-by: Chang Su <chang.s.su@oracle.com>
![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: 54c10a81d7
ℹ️ 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".
1 participant
Description
Problem
The vLLM gRPC servicer had several efficiency issues (unnecessary allocations in logprobs building, unvectorized placeholder mask ops, numpy intermediate in tensor deserialization) and maintained a duplicate
server.pythat is now upstream in vLLM.Solution
Optimize hot paths in the servicer, remove the duplicate server entrypoint, and consolidate kv_transfer_params validation.
Changes
Optimize tensor ops and logprobs building (
servicer.py):_build_top_logprobsin_build_input_logprobsto avoid allocating empty proto objects for every prompt token when top logprobs aren't requested_tensor_from_proto— usetorch.frombuffer(bytearray(...))directly (1 copy instead of 2, removed numpy dependency).flatten().to(torch.int64)callsRemove duplicate
server.py: The canonical entrypoint is nowvllm.entrypoints.grpc_serverupstream. Updated__init__.pyandREADME.mdreferences.Move kv_transfer_params validation into
_sampling_params_from_proto: Consolidates validation and logging into one place instead of inline inGenerate(). RaisesValueErroron invalid params, caught by existing error handler.Test Plan
Checklist
cargo +nightly fmtpassescargo clippy --all-targets --all-features -- -D warningspassesSummary by CodeRabbit
Documentation
Refactor
Bug Fixes
Chores