feat(grpc): add VllmHealthServicer for standard gRPC health checking (grpc.health.v1) - #885
Conversation
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
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 gap in vLLM's gRPC capabilities by introducing a standard health check servicer. This enhancement is crucial for operational environments, particularly those leveraging Kubernetes, as it provides a standardized and robust mechanism for monitoring the health and readiness of vLLM gRPC services, thereby improving reliability and manageability. 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 (2)
📝 WalkthroughWalkthroughAdded a new vLLM gRPC health servicer Changes
Sequence DiagramsequenceDiagram
autonumber
actor Client
participant VllmHealthServicer as "VllmHealthServicer\n(rpc)"
participant AsyncLLM as "async_llm"
participant GRPC_CTX as "gRPC Context"
Client->>VllmHealthServicer: Check(HealthCheckRequest(service))
alt _shutting_down == True
VllmHealthServicer-->>Client: HealthCheckResponse(NOT_SERVING)
else Supported service ("" or "vllm.grpc.engine.VllmEngine")
VllmHealthServicer->>AsyncLLM: await check_health()
alt check_health succeeds
AsyncLLM-->>VllmHealthServicer: success
VllmHealthServicer-->>Client: HealthCheckResponse(SERVING)
else check_health raises
AsyncLLM-->>VllmHealthServicer: exception
VllmHealthServicer-->>Client: HealthCheckResponse(NOT_SERVING)
end
else Unknown service
VllmHealthServicer->>GRPC_CTX: set_code(NOT_FOUND), set_details("Unknown service: " + request.service)
VllmHealthServicer-->>Client: HealthCheckResponse(SERVICE_UNKNOWN)
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related issues
Possibly related PRs
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)
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 introduces a new VllmHealthServicer to implement the standard gRPC health check protocol for vLLM, enabling Kubernetes probes to monitor the health of the AsyncLLM engine. The servicer supports both overall server health and specific VllmEngine service health checks. Feedback includes adding a type hint for the async_llm parameter in the constructor for improved type safety and clarity. Additionally, a critical issue was identified in the Watch method, where its current implementation violates the gRPC Health Checking Protocol by incorrectly setting the RPC status for unknown services; a direct re-implementation of the health checking logic within Watch is suggested to ensure protocol compliance.
![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/health_servicer.py`:
- Around line 88-110: The Watch method must not call Check (which sets context
code to NOT_FOUND) and must be converted from a single-yield to a persistent
stream: compute the initial status from your internal status store (e.g. lookup
in self._status_map or equivalent) and if missing send a
HealthCheckResponse(status=health_pb2.HealthCheckResponse.SERVICE_UNKNOWN)
without calling context.set_code(NOT_FOUND); yield that initial
HealthCheckResponse, then enter a loop that polls for status changes (or awaits
notifications) and yields new HealthCheckResponse messages whenever the status
changes, sleeping briefly between polls (e.g. asyncio.sleep(poll_interval)) and
terminating only when the gRPC context is cancelled; ensure you reference Watch,
Check, HealthCheckResponse and avoid invoking Check so the stream stays open for
unknown services.
- Around line 45-48: The VllmHealthServicer.set_not_serving() is never wired
into the server lifecycle so shutdown probes stay SERVING; mirror the sglang
server wiring: create and register a VllmHealthServicer instance with the gRPC
server, pass that instance into VllmEngineServicer on construction/registration,
and ensure the global graceful-shutdown handler calls
VllmHealthServicer.set_not_serving() before stopping the server (i.e., call
set_not_serving() in the same shutdown sequence that currently calls
server.stop/graceful drain as in the sglang example); look for
VllmHealthServicer, VllmEngineServicer and the server startup/shutdown block to
add the instantiation, registration, and shutdown call.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 468496dc-afcd-4e82-90db-684ff3ea1f4d
📒 Files selected for processing (2)
grpc_servicer/smg_grpc_servicer/vllm/__init__.pygrpc_servicer/smg_grpc_servicer/vllm/health_servicer.py
| async def Watch( | ||
| self, | ||
| request: health_pb2.HealthCheckRequest, | ||
| context: grpc.aio.ServicerContext, | ||
| ) -> AsyncIterator[health_pb2.HealthCheckResponse]: | ||
| """ | ||
| Streaming health check - sends current status once. | ||
|
|
||
| For now, sends current status once (Kubernetes doesn't use Watch). | ||
| A full implementation would monitor status changes and stream updates. | ||
|
|
||
| Args: | ||
| request: Contains service name | ||
| context: gRPC context | ||
|
|
||
| Yields: | ||
| HealthCheckResponse messages | ||
| """ | ||
| service_name = request.service | ||
| logger.debug(f"Health watch request for service: '{service_name}'") | ||
|
|
||
| response = await self.Check(request, context) | ||
| yield response |
There was a problem hiding this comment.
🧩 Analysis chain
🌐 Web query:
In grpc.health.v1, what is the expected behavior of Health/Watch for unknown services, and should Watch be a continuous stream of status updates instead of a single response?
💡 Result:
For unknown services in grpc.health.v1, the Health/Watch RPC has a specific expected behavior distinct from the Check RPC. When Watch receives a HealthCheckRequest for an unknown service, the server immediately sends a single HealthCheckResponse with status SERVICE_UNKNOWN (value 3) but does NOT terminate the call. The stream remains open, and the server MUST send subsequent messages if the service becomes known later (e.g., changes to SERVING or NOT_SERVING). If the stream terminates with UNIMPLEMENTED, clients should not retry; other terminations trigger retries with backoff. Yes, Watch is designed as a continuous bidirectional streaming RPC (server-streaming from the proto definition: rpc Watch(HealthCheckRequest) returns (stream HealthCheckResponse)). The server sends an initial status immediately upon request, then additional updates whenever the service's status changes over time. It is explicitly not a single unary response but a persistent stream for ongoing monitoring.
Citations:
- 1: https://grpc.io/grpc/core/md_doc_health-checking.html
- 2: https://pkg.go.dev/google.golang.org/grpc/health/grpc_health_v1
- 3: https://grpc.github.io/grpc/python/grpc_health_checking.html
- 4: https://github.com/grpc/grpc-proto/blob/master/grpc/health/v1/health.proto
- 5: https://grpc.io/docs/guides/health-checking/
- 6: https://github.com/grpc/grpc/blob/master/doc/health-checking.md
🏁 Script executed:
cat -n grpc_servicer/smg_grpc_servicer/vllm/health_servicer.py | head -120Repository: lightseekorg/smg
Length of output: 4741
Watch implementation violates gRPC health.v1 streaming semantics in two ways.
-
For unknown services: Line 109 calls
Check(), which sets the gRPC context code toNOT_FOUND(line 84). This terminates the stream with aNOT_FOUNDerror. Per the spec,Watchshould send aSERVICE_UNKNOWNresponse and keep the stream open without terminating. -
Single response instead of stream: Line 110 yields once then closes.
Watchmust be a persistent stream that sends initial status immediately and continues to emit updates when service status changes, not a single unary response.
The suggested fix addresses both issues by bypassing Check() (avoiding the problematic NOT_FOUND context code) and implementing a continuous stream with periodic health checks.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@grpc_servicer/smg_grpc_servicer/vllm/health_servicer.py` around lines 88 -
110, The Watch method must not call Check (which sets context code to NOT_FOUND)
and must be converted from a single-yield to a persistent stream: compute the
initial status from your internal status store (e.g. lookup in self._status_map
or equivalent) and if missing send a
HealthCheckResponse(status=health_pb2.HealthCheckResponse.SERVICE_UNKNOWN)
without calling context.set_code(NOT_FOUND); yield that initial
HealthCheckResponse, then enter a loop that polls for status changes (or awaits
notifications) and yields new HealthCheckResponse messages whenever the status
changes, sleeping briefly between polls (e.g. asyncio.sleep(poll_interval)) and
terminating only when the gRPC context is cancelled; ensure you reference Watch,
Check, HealthCheckResponse and avoid invoking Check so the stream stays open for
unknown services.
There was a problem hiding this comment.
Acknowledged. The current single-yield Watch matches SGLangHealthServicer in the same repo (sglang/health_servicer.py:165). Converting Watch to a persistent stream is a larger change that should be done for both SGLang and vLLM servicers together as a follow-up PR. The context.set_code() side effect was already fixed in the previous commit by inlining the status computation.
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Check the current state of the vLLM health servicer Watch method
cat -n grpc_servicer/smg_grpc_servicer/vllm/health_servicer.py
echo "---SGLang health servicer Watch---"
# Check the SGLang health servicer Watch method for comparison
grep -n "Watch\|yield\|set_code\|Check" grpc_servicer/smg_grpc_servicer/sglang/health_servicer.py
Length of output: 8352
@V2arK, confirmed — the current branch already addresses both points:
-
context.set_code()side effect — fixed. Lines 113–130 of the currenthealth_servicer.pyinline the status computation inWatchwith an explicit comment explaining the rationale, soCheck()is no longer called fromWatch. This is actually an improvement overSGLangHealthServicer, which still delegates toself.Check()in itsWatch(line 164 ofsglang/health_servicer.py). -
Single-yield
Watch— intentional. The pattern is consistent withSGLangHealthServicer, and the docstring documents the limitation and rationale. The follow-up refactor covering both servicers together makes sense.
Noting the design decision for future reviews.
✏️ Learnings added
Learnt from: V2arK
Repo: lightseekorg/smg PR: 885
File: grpc_servicer/smg_grpc_servicer/vllm/health_servicer.py:0-0
Timestamp: 2026-03-24T18:53:57.311Z
Learning: In repo lightseekorg/smg, the single-yield Watch() pattern in grpc health servicers is intentional and consistent across both `grpc_servicer/smg_grpc_servicer/sglang/health_servicer.py` and `grpc_servicer/smg_grpc_servicer/vllm/health_servicer.py`. Kubernetes only uses the Check RPC, so Watch yields one status response then closes. A persistent-stream Watch is a known future improvement to be implemented for both servicers together in a follow-up PR. Do not flag the single-yield Watch as a protocol violation in these files.
Note: Learnings are effective only in the context of similar code segments. To apply general coding guidelines, please configure review instructions. You can manage existing learnings in the app.
🧠 Learnings used
Learnt from: zhaowenzi
Repo: lightseekorg/smg PR: 807
File: model_gateway/src/routers/openai/responses/streaming.rs:821-855
Timestamp: 2026-03-18T21:57:03.433Z
Learning: In repo lightseekorg/smg, file model_gateway/src/routers/openai/responses/streaming.rs: The early `return` statements inside `handle_streaming_with_tool_interception` (on `tx` send failures in `forward_streaming_event`, `send_mcp_list_tools_events`, and the `is_in_progress`/`mcp_list_tools_sent` branches) are pre-existing behavior that predates PR `#807`. They cause the persistence phase (final response / conversation-backed storage writes at the end of the tool loop) to be skipped when the client disconnects mid-stream with `store=true` or a conversation-backed request. This is a known pre-existing gap in the MCP streaming path, not a regression introduced by the storage context header changes in PR `#807`.
Learnt from: CatherineSue
Repo: lightseekorg/smg PR: 638
File: grpc_servicer/smg_grpc_servicer/vllm/servicer.py:321-323
Timestamp: 2026-03-05T04:48:40.754Z
Learning: In `grpc_servicer/smg_grpc_servicer/vllm/servicer.py`, the `last_receive_timestamp` field in `GetServerInfoResponse` is intentionally set to `time.time()` at read time (with an inline `# TODO looks wrong?` comment). This is a pre-existing TODO carried over from vLLM and is not to be fixed within the current PR scope.
Learnt from: hyeongyun0916
Repo: lightseekorg/smg PR: 784
File: grpc_servicer/smg_grpc_servicer/vllm/render_servicer.py:65-70
Timestamp: 2026-03-19T10:00:53.513Z
Learning: In repo lightseekorg/smg, the convention across all gRPC servicer implementations (e.g., grpc_servicer/smg_grpc_servicer/vllm/servicer.py, grpc_servicer/smg_grpc_servicer/vllm/render_servicer.py) is to NOT add an explicit `return` after `await context.abort(...)` calls, because `context.abort()` raises `grpc.aio.AbortError` internally making any subsequent return unreachable dead code. This is an intentional consistency decision across the codebase.
Learnt from: CatherineSue
Repo: lightseekorg/smg PR: 690
File: model_gateway/src/core/worker_registry.rs:667-670
Timestamp: 2026-03-10T05:04:49.809Z
Learning: In repo lightseekorg/smg, file model_gateway/src/core/worker_registry.rs: `any_external_worker_supports_model` intentionally uses `healthy_only = true` for two reasons: (1) The 503 ("service unavailable") path in `select_worker_for_model` is specifically for the circuit-breaker case — healthy workers whose circuit breaker is open — while workers failing health checks fall through to 404. (2) Unhealthy workers have stale model lists (models registered at startup may no longer be accurate) and should not be trusted for model existence checks. This design follows the upstream sglang pattern from sgl-project/sglang#15611. Do not flag `healthy_only = true` in `any_external_worker_supports_model` as a bug.
Learnt from: hyeongyun0916
Repo: lightseekorg/smg PR: 784
File: grpc_servicer/smg_grpc_servicer/vllm/render_servicer.py:98-103
Timestamp: 2026-03-19T10:01:20.396Z
Learning: In repo lightseekorg/smg, the convention across all gRPC servicer implementations (e.g., vllm/servicer.py, vllm/render_servicer.py) is to NOT add an explicit `return` after `await context.abort(...)` calls. The codebase relies on `grpc.aio.AbortError` being raised by `context.abort()` to stop execution, and omitting the explicit `return` is an intentional, consistent design decision.
Learnt from: CatherineSue
Repo: lightseekorg/smg PR: 690
File: model_gateway/src/core/worker_registry.rs:667-670
Timestamp: 2026-03-10T04:59:38.803Z
Learning: In repo lightseekorg/smg, file model_gateway/src/core/worker_registry.rs: `any_external_worker_supports_model` intentionally uses `healthy_only = true`. The 503 ("service unavailable") path in `select_worker_for_model` is specifically for the circuit-breaker case — healthy workers whose circuit breaker is open. Workers that are genuinely unhealthy (failing health checks) are intentionally excluded: the model falls through to 404 for those. This design follows the upstream sglang pattern established in sgl-project/sglang#15611 ("[model-gateway] return 503 when all workers are circuit-broken"). Do not flag the `healthy_only = true` argument in this method as a bug.
Learnt from: pallasathena92
Repo: lightseekorg/smg PR: 687
File: model_gateway/src/routers/openai/realtime/webrtc.rs:238-289
Timestamp: 2026-03-11T01:29:56.655Z
Learning: In repo lightseekorg/smg, file model_gateway/src/routers/openai/realtime/webrtc.rs: `Metrics::record_router_request` is already emitted in router.rs (around line 1152) before `handle_realtime_webrtc` is called, and `Metrics::record_router_error` is emitted inside `handle_realtime_webrtc` for the no-workers case. The missing instrumentation is success/duration recording after `setup_and_spawn_bridge` returns — this is a metrics improvement deferred to a follow-up PR, not a correctness gap. Do not flag missing success/duration metrics as a blocking issue for PR `#687`.
Learnt from: key4ng
Repo: lightseekorg/smg PR: 867
File: tui/src/app.rs:571-596
Timestamp: 2026-03-23T20:26:39.459Z
Learning: In lightseekorg/smg, tui/src/app.rs, the "toggle-health" command uses `worker.is_healthy` (not a `disable_health_check` field) to compute the new `disable_health_check` value because the worker list API does not expose `disable_health_check` in its response. The intended behavior is: healthy worker → disable health checks (maintenance mode), unhealthy worker → re-enable health checks (recovery). Do not flag this as a bug; it is an intentional heuristic given the API limitation.
Learnt from: hyeongyun0916
Repo: lightseekorg/smg PR: 784
File: grpc_servicer/smg_grpc_servicer/vllm/render_servicer.py:94-96
Timestamp: 2026-03-17T09:41:10.752Z
Learning: Enforce the repository-wide convention for gRPC INTERNAL errors in all servicer implementations: log the exception on the server (using logger.exception or logger.error) and then call context.abort(grpc.StatusCode.INTERNAL, str(e)) without replacing the message. Do not substitute a generic message. This applies to all servicer files (e.g., vllm/servicer.py, sglang/servicer.py, vllm/render_servicer.py) and should be treated as an intentional, consistent design decision across the codebase.
If you found this review helpful, would you consider giving us a shout-out on X?
Thank you for using CodeRabbit!
There was a problem hiding this comment.
A follow up PR sounds good. @V2arK lmk if you want to take on this task. Appreciate the contribution
There was a problem hiding this comment.
Seems like the humans are having a chat. I'll hop back into my burrow for now. If you need me again, just tag @coderabbitai in a new comment, and I'll come hopping out!
There was a problem hiding this comment.
@CatherineSue For the Watch continuous stream follow-up, I'm considering two approaches:
- Inline: each servicer implements its own Watch loop independently (duplicated but self-contained)
- Mixin: extract the Watch loop into a shared HealthWatchMixin, servicers only provide_compute_watch_status() and _is_shutting_down()
Given there are only two servicers today, I'm leaning toward (1) for simplicity. Any preference?
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/health_servicer.py`:
- Around line 34-42: Update the __init__ signature of the Health servicer to
type-hint the async_llm parameter as "AsyncLLM" (use a string forward reference
to avoid circular imports) so it matches VllmEngineServicer; if you prefer
explicit typing, add a TYPE_CHECKING block to import AsyncLLM (from vllm import
AsyncLLM) guarded by typing.TYPE_CHECKING and annotate self.async_llm
accordingly in the __init__ of the health servicer class.
- Around line 119-120: The Watch method currently swallows all exceptions;
mirror the Check method's behavior by logging the exception instead of silently
ignoring it: inside the except Exception block in HealthServicer.Watch, call
logger.debug (or logger.exception if you prefer full error visibility) with a
brief message and exc_info=True so the traceback is captured when debugging is
enabled; this keeps behavior consistent with Check (which uses logger.exception)
while avoiding noisy logs at higher levels.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 0a60ca28-8fd2-4d2c-a97b-ab9cba214558
📒 Files selected for processing (1)
grpc_servicer/smg_grpc_servicer/vllm/health_servicer.py
|
Note: the 3 CI failures ( |
57d0352 to
fef5e88
Compare
|
WASM should be fixed. Rebased the commits. |
CatherineSue
left a comment
CatherineSue
left a comment
There was a problem hiding this comment.
Thanks for contributing this critical piece. LGTM
…protocol Signed-off-by: Honglin Cao <Caohonglin317@hotmail.com>
- Add forward reference type hint for async_llm parameter - Inline health check logic in Watch() to avoid Check()'s context.set_code() side effect on streaming responses Signed-off-by: Honglin Cao <Caohonglin317@hotmail.com>
- Add TYPE_CHECKING guard for AsyncLLM type hint on __init__ - Add logger.debug(exc_info=True) in Watch except block for consistency with Check's logger.exception() Signed-off-by: Honglin Cao <Caohonglin317@hotmail.com>
fef5e88 to
b365f6f
Compare
My pleasure! |
…(grpc.health.v1) (smg-project#885) Signed-off-by: Honglin Cao <Caohonglin317@hotmail.com>
3 participants
Description
Problem
smg-grpc-servicerprovidesSGLangHealthServicerfor SGLang's standard gRPC health check (grpc.health.v1.Health), but the vLLM side has no equivalent. vLLM gRPC deployments cannot use Kubernetes native gRPC health probes (livenessProbe.grpc/readinessProbe.grpc, GA since K8s 1.27).Solution
Add
VllmHealthServicerinsmg_grpc_servicer/vllm/health_servicer.py, mirroringSGLangHealthServicerin structure and placement. Health status delegates toAsyncLLM.check_health()-- the sameEngineClientprotocol method used by vLLM's HTTP/healthendpoint.Companion vLLM PR: vllm-project/vllm#38016
Changes
smg_grpc_servicer/vllm/health_servicer.pyVllmHealthServicer(health_pb2_grpc.HealthServicer)withTYPE_CHECKINGtype hint forAsyncLLMCheck(): callsawait async_llm.check_health(), returns SERVING / NOT_SERVING / SERVICE_UNKNOWN. Logs exceptions vialogger.exception().Watch(): inlines status computation (avoidsCheck()'scontext.set_code()side effect on streaming responses). Logs exceptions vialogger.debug(exc_info=True). Single-yield matchingSGLangHealthServicerbehavior; persistent streaming deferred to a follow-up PR for both servicers.set_not_serving(): sets_shutting_downflag for graceful shutdown. Called by vLLM'sserve_grpc()in itsfinallyblock (companion PR).""(liveness) and"vllm.grpc.engine.VllmEngine"(readiness)smg_grpc_servicer/vllm/__init__.pyto exportVllmHealthServicerTest Plan
Tests are in the companion vLLM PR (test file). All verified on NVIDIA H200 MIG (1g.18gb) with
facebook/opt-125m.Unit tests: 13/13 passed
test_check_serving_overalltest_check_serving_vllm_servicetest_check_not_serving_engine_erroredtest_check_not_serving_shutting_downtest_check_unknown_service_statustest_check_unknown_service_grpc_codetest_check_shutting_down_overrides_healthytest_check_logs_exception_on_errortest_watch_yields_servingtest_watch_yields_not_servingtest_watch_yields_exactly_oncetest_watch_unknown_servicetest_set_not_serving_sets_flagE2E tests: 10/10 scenarios passed
grpc.health.v1.Health{"healthy":true}Checklist
cargo +nightly fmtpassescargo clippy --all-targets --all-features -- -D warningspassesSummary by CodeRabbit