feat(grpc_servicer): add sglang support with multi-backend extras - #745
Conversation
github-actions
Bot
added
documentation
📝 WalkthroughWalkthroughThis PR adds comprehensive SGLang backend support to the gRPC servicer package alongside existing vLLM support. The changes include version bump to 0.5.0, optional dependency configuration, documentation updates reflecting support for both backends, and new SGLang service modules implementing request management, process supervision, gRPC server integration, and health checking. Changes
Sequence DiagramsequenceDiagram
participant Client
participant SGLangServicer
participant GrpcRequestManager
participant ZMQ Scheduler
participant SchedulerProcess
Client->>SGLangServicer: Generate Request
SGLangServicer->>GrpcRequestManager: generate_request() async
GrpcRequestManager->>ZMQ Scheduler: PUSH request via socket
ZMQ Scheduler->>SchedulerProcess: Process tokens
SchedulerProcess->>ZMQ Scheduler: BatchTokenIDOutput
ZMQ Scheduler->>GrpcRequestManager: PULL output
GrpcRequestManager->>GrpcRequestManager: _handle_batch_output()
GrpcRequestManager->>SGLangServicer: Yield response chunk
SGLangServicer->>Client: Stream GenerateResponse
Client->>SGLangServicer: HealthCheck Request
SGLangServicer->>GrpcRequestManager: Health-check generate request
GrpcRequestManager->>ZMQ Scheduler: PUSH health request
ZMQ Scheduler->>SchedulerProcess: Perform check
SchedulerProcess->>ZMQ Scheduler: HealthCheckOutput
ZMQ Scheduler->>GrpcRequestManager: PULL health response
GrpcRequestManager->>SGLangServicer: Health status
SGLangServicer->>Client: HealthCheckResponse
Estimated code review effort🎯 4 (Complex) | ⏱️ ~70 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)
📝 Coding Plan
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 significantly enhances the Highlights
Changelog
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
|
|
Hi @slin1237, the DCO sign-off check has failed. All commits must include a To fix existing commits: # Sign off the last N commits (replace N with the number of unsigned commits)
git rebase HEAD~N --signoff
git push --force-with-leaseTo sign off future commits automatically:
|
![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 is a significant architectural improvement, refactoring smg-grpc-servicer to support multiple backends like vLLM and SGLang through optional extras. The introduction of the SGLang gRPC servicer is well-executed, incorporating a dedicated request manager, health checks, and a server entrypoint that addresses production considerations such as TLS and graceful shutdown. The implementation is robust and aligns with existing SGLang patterns. My feedback includes a couple of suggestions to enhance the clarity and reliability of the health check mechanisms, specifically addressing the use of magic numbers and ensuring robust, individual timeouts for health checks in line with established repository guidelines.
![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: a4f756ccb7
ℹ️ 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".
Move vllm from hard dependency to optional extra and add sglang extra. Base package now only depends on smg-grpc-proto and grpcio packages. - Restructure pyproject.toml: vllm and sglang as optional extras - Add smg_grpc_servicer/sglang/ package (servicer, server, request manager, health servicer, scheduler launcher, utils) - Update README and DEVELOPMENT docs for multi-backend usage - Add grpc_servicer install to ci_install_sglang.sh Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
Apply ruff and ruff-format fixes, fix remaining E501 long lines. Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 11
🤖 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/sglang/request_manager.py`:
- Around line 673-679: The delayed cleanup coroutine created in the cleanup()
closure (cleanup(rid)) is only appended to cleanup_tasks and not registered on
the component-wide task tracker, so it won't be cancelled on shutdown; update
the task creation so the created task is also added to self.asyncio_tasks (or
the existing task-tracking collection) after asyncio.create_task(cleanup(rid))
so shutdown logic can cancel it (reference the cleanup coroutine,
cleanup_tasks.append(...), and self.asyncio_tasks to locate the change).
- Around line 401-446: The background task started in embedding_request (the
inner coroutine wait_for_result) is not tracked, so create the task as a
variable (e.g., task = asyncio.create_task(wait_for_result())) and add it to the
servicer's task collection (self.asyncio_tasks) immediately; also ensure the
task is removed from self.asyncio_tasks when it completes (inside
wait_for_result's finally block or via task.add_done_callback) so it can be
cancelled/cleaned up during shutdown. Target symbols: embedding_request,
wait_for_result, self.asyncio_tasks.
- Around line 819-842: The shutdown method currently awaits
asyncio.gather(*list(self.asyncio_tasks), return_exceptions=True) twice; remove
the second duplicate gather call at the end of GrpcRequestManager.shutdown so
tasks are awaited only once after cancellation (keep the cancellation loop and
the first await to let cancelled tasks process), and retain the logic that
iterates rid_to_state to put shutdown errors into state.out_queue and set
state.finished/state.event; ensure references are to shutdown,
self.asyncio_tasks, and rid_to_state/state.out_queue/state.event so the correct
code block is edited.
- Around line 953-965: The exception handler calls
func.__self__.dump_requests_before_crash() but GrpcRequestManager lacks that
method, causing an AttributeError; add a dump_requests_before_crash(self) method
on GrpcRequestManager that safely reads the existing crash_dump_request_list
(use the same locking/structure used by record_request_for_crash_dump()),
serializes or logs the stored requests (e.g., to logger.error or a crash dump
file) with clear context, and ensure it is no-op safe if the list is empty;
alternatively, if you prefer not to add a method, modify print_exception_wrapper
to check for hasattr(..., "dump_requests_before_crash") before calling it to
avoid the AttributeError.
In `@grpc_servicer/smg_grpc_servicer/sglang/scheduler_launcher.py`:
- Around line 71-72: Replace the unconditional mp.set_start_method("spawn",
force=True) with a safe check: call mp.get_start_method(allow_none=True) and
only call mp.set_start_method("spawn") if the current method is None or
different, and when you must override, log a warning (using the module logger or
processLogger) indicating the previous method to avoid silently masking
configuration; reference mp.set_start_method and mp.get_start_method in
scheduler_launcher.py.
In `@grpc_servicer/smg_grpc_servicer/sglang/server.py`:
- Around line 319-326: The warmup currently always uses grpc.insecure_channel
which fails when the server is TLS-enabled; update _execute_grpc_server_warmup
to detect server_args.ssl_certfile (and ssl_keyfile) and, when present, load the
server certificate (server_args.ssl_certfile), create
grpc.ssl_channel_credentials using that cert as the root certificate, and use
grpc.secure_channel with those credentials and the same channel options;
otherwise fallback to the existing grpc.insecure_channel path. Ensure you
reference _execute_grpc_server_warmup, server_args.ssl_certfile, and
server_args.ssl_keyfile when making the change.
In `@grpc_servicer/smg_grpc_servicer/sglang/servicer.py`:
- Around line 637-641: The _kv_event_id_counter is incremented unsafely in
_convert_kv_event which can produce duplicate IDs across concurrent
SubscribeKvEvents streams; either make the counter thread-safe by replacing the
shared integer with a global iterator (e.g., itertools.count()) or protect
increments with an asyncio.Lock around accesses to _kv_event_id_counter in
_convert_kv_event, or alternatively move the counter to be a per-stream local
variable inside SubscribeKvEvents if IDs only need to be unique per
stream—update the implementation of _convert_kv_event and SubscribeKvEvents
accordingly to use the chosen approach and ensure the unique-symbols
_kv_event_id_counter, _convert_kv_event, and SubscribeKvEvents are adjusted
together.
- Around line 358-392: Replace the manual polling loop with awaiting the
run_health_check task via asyncio.wait_for to ensure we observe the specific
health-check result rather than unrelated scheduler activity: create the task
from run_health_check (function run_health_check uses
self.request_manager.generate_request), then use try/except around await
asyncio.wait_for(task, timeout=HEALTH_CHECK_TIMEOUT) to return a healthy
response on True, handle asyncio.TimeoutError to cancel the task and return the
timeout unhealthy response, and handle other exceptions to log and return
unhealthy; always call self.request_manager._cleanup_request_state(rid) in a
finally block so request state is cleaned up whether the check succeeds, times
out, or errors, and avoid cancelling the task when it has already completed.
- Around line 759-766: The _decode_tensor_data function currently maps only
"float32" and "int64" (dtype_map) and silently defaults unknown
tensor_data.dtype to float32, which can corrupt data; update _decode_tensor_data
to: expand dtype_map to include other expected proto dtypes (e.g., "float16",
"bfloat16"/"bf16", "int32", "uint8", "bool", etc.) mapped to the correct numpy
dtypes, use that mapping to select np_dtype from tensor_data.dtype, and if
tensor_data.dtype is not in the map raise a clear exception (ValueError) naming
tensor_data.dtype so callers fail fast rather than silently decoding as float32;
keep the existing np.frombuffer(...).reshape(...) and torch.from_numpy(arr)
flow.
In `@grpc_servicer/smg_grpc_servicer/sglang/utils.py`:
- Around line 14-21: The function abort_code_from_output assumes `output` is a
dict and will raise if `None`; add a defensive guard at the start of
abort_code_from_output to handle non-dict/None inputs (e.g., if output is None
or not isinstance(output, dict): return grpc.StatusCode.INTERNAL), then proceed
to access meta_info and use _HTTP_TO_GRPC_CODE as before so the function never
calls .get on a None value.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: daba203a-a142-4281-8660-4c1065ea9399
📒 Files selected for processing (11)
grpc_servicer/DEVELOPMENT.mdgrpc_servicer/README.mdgrpc_servicer/pyproject.tomlgrpc_servicer/smg_grpc_servicer/sglang/__init__.pygrpc_servicer/smg_grpc_servicer/sglang/health_servicer.pygrpc_servicer/smg_grpc_servicer/sglang/request_manager.pygrpc_servicer/smg_grpc_servicer/sglang/scheduler_launcher.pygrpc_servicer/smg_grpc_servicer/sglang/server.pygrpc_servicer/smg_grpc_servicer/sglang/servicer.pygrpc_servicer/smg_grpc_servicer/sglang/utils.pyscripts/ci_install_sglang.sh
| async def embedding_request( | ||
| self, | ||
| obj: TokenizedEmbeddingReqInput, | ||
| request_id: str | None = None, | ||
| ) -> asyncio.Future: | ||
| """ | ||
| Submit an embedding request to the scheduler. | ||
| Returns a future that will contain the embedding result. | ||
| """ | ||
| # Generate request ID if not provided | ||
| if request_id is None: | ||
| request_id = f"grpc-embed-{uuid.uuid4().hex}" | ||
|
|
||
| obj.rid = request_id | ||
|
|
||
| self._req_stats_init(obj) | ||
| state = self.rid_to_state[request_id] | ||
|
|
||
| # Create future for result | ||
| future = asyncio.Future() | ||
|
|
||
| # Send to scheduler | ||
| try: | ||
| state.time_stats.set_api_server_dispatch_time() | ||
| await self._send_to_scheduler(obj) | ||
| state.time_stats.set_api_server_dispatch_finish_time() | ||
| except Exception as e: | ||
| del self.rid_to_state[request_id] | ||
| future.set_exception(e) | ||
| return future | ||
|
|
||
| # Wait for result in background | ||
| async def wait_for_result(): | ||
| try: | ||
| await state.event.wait() | ||
| result = await state.out_queue.get() | ||
| future.set_result(result) | ||
| except Exception as e: | ||
| future.set_exception(e) | ||
| finally: | ||
| # Clean up | ||
| if request_id in self.rid_to_state: | ||
| del self.rid_to_state[request_id] | ||
|
|
||
| asyncio.create_task(wait_for_result()) | ||
| return future |
There was a problem hiding this comment.
Embedding request creates untracked background task.
The asyncio.create_task(wait_for_result()) at Line 445 creates a task that isn't added to self.asyncio_tasks, which means it won't be properly cancelled during shutdown. This could lead to orphaned tasks.
🛠️ Proposed fix to track the background task
# Wait for result in background
async def wait_for_result():
try:
await state.event.wait()
result = await state.out_queue.get()
future.set_result(result)
except Exception as e:
future.set_exception(e)
finally:
# Clean up
if request_id in self.rid_to_state:
del self.rid_to_state[request_id]
+ self.asyncio_tasks.discard(task)
- asyncio.create_task(wait_for_result())
+ task = asyncio.create_task(wait_for_result())
+ self.asyncio_tasks.add(task)
return future🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@grpc_servicer/smg_grpc_servicer/sglang/request_manager.py` around lines 401 -
446, The background task started in embedding_request (the inner coroutine
wait_for_result) is not tracked, so create the task as a variable (e.g., task =
asyncio.create_task(wait_for_result())) and add it to the servicer's task
collection (self.asyncio_tasks) immediately; also ensure the task is removed
from self.asyncio_tasks when it completes (inside wait_for_result's finally
block or via task.add_done_callback) so it can be cancelled/cleaned up during
shutdown. Target symbols: embedding_request, wait_for_result,
self.asyncio_tasks.
| # Remove from tracking after a delay | ||
| async def cleanup(request_id): | ||
| await asyncio.sleep(5.0) | ||
| if request_id in self.rid_to_state: | ||
| del self.rid_to_state[request_id] | ||
|
|
||
| cleanup_tasks.append(asyncio.create_task(cleanup(rid))) |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Cleanup tasks are not tracked for shutdown.
The delayed cleanup tasks created at Lines 674-679 are not added to self.asyncio_tasks, meaning they won't be cancelled during shutdown. While this is a 5-second delay and likely harmless, tracking them would be more robust.
♻️ Suggested tracking for cleanup tasks
# Remove from tracking after a delay
async def cleanup(request_id):
await asyncio.sleep(5.0)
if request_id in self.rid_to_state:
del self.rid_to_state[request_id]
+ self.asyncio_tasks.discard(asyncio.current_task())
- cleanup_tasks.append(asyncio.create_task(cleanup(rid)))
+ task = asyncio.create_task(cleanup(rid))
+ self.asyncio_tasks.add(task)
+ cleanup_tasks.append(task)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@grpc_servicer/smg_grpc_servicer/sglang/request_manager.py` around lines 673 -
679, The delayed cleanup coroutine created in the cleanup() closure
(cleanup(rid)) is only appended to cleanup_tasks and not registered on the
component-wide task tracker, so it won't be cancelled on shutdown; update the
task creation so the created task is also added to self.asyncio_tasks (or the
existing task-tracking collection) after asyncio.create_task(cleanup(rid)) so
shutdown logic can cancel it (reference the cleanup coroutine,
cleanup_tasks.append(...), and self.asyncio_tasks to locate the change).
| async def shutdown(self): | ||
| """Gracefully shutdown the request manager.""" | ||
| logger.info("Shutting down GrpcRequestManager") | ||
| self.gracefully_exit = True | ||
|
|
||
| # Cancel all asyncio tasks FIRST - this will interrupt blocked recv() calls | ||
| for task in list(self.asyncio_tasks): | ||
| if not task.done(): | ||
| task.cancel() | ||
|
|
||
| # Give tasks a moment to process cancellation | ||
| if self.asyncio_tasks: | ||
| await asyncio.gather(*list(self.asyncio_tasks), return_exceptions=True) | ||
|
|
||
| # Cancel all pending requests | ||
| for rid, state in list(self.rid_to_state.items()): | ||
| if not state.finished: | ||
| await state.out_queue.put({"error": "Server shutting down", "shutdown": True}) | ||
| state.finished = True | ||
| state.event.set() | ||
|
|
||
| # Wait for tasks to complete | ||
| if self.asyncio_tasks: | ||
| await asyncio.gather(*list(self.asyncio_tasks), return_exceptions=True) |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Redundant asyncio.gather calls during shutdown.
The shutdown method calls asyncio.gather(*list(self.asyncio_tasks), ...) twice (Lines 831 and 842). The first gather after task cancellation should be sufficient, and the second is redundant since tasks should already be done.
♻️ Proposed simplification
# Cancel all asyncio tasks FIRST - this will interrupt blocked recv() calls
for task in list(self.asyncio_tasks):
if not task.done():
task.cancel()
- # Give tasks a moment to process cancellation
- if self.asyncio_tasks:
- await asyncio.gather(*list(self.asyncio_tasks), return_exceptions=True)
-
# Cancel all pending requests
for rid, state in list(self.rid_to_state.items()):
if not state.finished:
await state.out_queue.put({"error": "Server shutting down", "shutdown": True})
state.finished = True
state.event.set()
- # Wait for tasks to complete
+ # Wait for cancelled tasks to complete
if self.asyncio_tasks:
await asyncio.gather(*list(self.asyncio_tasks), return_exceptions=True)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@grpc_servicer/smg_grpc_servicer/sglang/request_manager.py` around lines 819 -
842, The shutdown method currently awaits
asyncio.gather(*list(self.asyncio_tasks), return_exceptions=True) twice; remove
the second duplicate gather call at the end of GrpcRequestManager.shutdown so
tasks are awaited only once after cancellation (keep the cancellation loop and
the first await to let cancelled tasks process), and retain the logic that
iterates rid_to_state to put shutdown errors into state.out_queue and set
state.finished/state.event; ensure references are to shutdown,
self.asyncio_tasks, and rid_to_state/state.out_queue/state.event so the correct
code block is edited.
| async def print_exception_wrapper(func): | ||
| """ | ||
| Sometimes an asyncio function does not print exception. | ||
| We do another wrapper to handle the exception. | ||
| """ | ||
| try: | ||
| await func() | ||
| except Exception: | ||
| traceback = get_exception_traceback() | ||
| logger.error(f"GrpcRequestManager hit an exception: {traceback}") | ||
| if hasattr(func, "__self__") and isinstance(func.__self__, GrpcRequestManager): | ||
| func.__self__.dump_requests_before_crash() | ||
| kill_process_tree(os.getpid(), include_parent=True) |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Search for dump_requests_before_crash method definition in the codebase
# Check if the method exists anywhere in the sglang package
rg -n "def dump_requests_before_crash" --type py
# Check what crash dump related methods exist in the request manager
rg -n "crash_dump|dump_request" grpc_servicer/Repository: lightseekorg/smg
Length of output: 831
🏁 Script executed:
# Search for any other references to dump_requests_before_crash
rg "dump_requests_before_crash" grpc_servicer/
# Get all public methods in GrpcRequestManager class
rg -A 2 "^\s*def " grpc_servicer/smg_grpc_servicer/sglang/request_manager.py | grep -E "^\s*def " | head -30Repository: lightseekorg/smg
Length of output: 753
🏁 Script executed:
# Get the full print_exception_wrapper function with more context
sed -n '950,970p' grpc_servicer/smg_grpc_servicer/sglang/request_manager.pyRepository: lightseekorg/smg
Length of output: 662
🏁 Script executed:
# Look at the record_request_for_crash_dump method to understand the pattern
sed -n '808,825p' grpc_servicer/smg_grpc_servicer/sglang/request_manager.py
# Search for any methods that might write/dump crash information
rg -n "def.*dump|def.*crash|def.*write.*request" grpc_servicer/smg_grpc_servicer/sglang/request_manager.py | head -20
# Check if crash_dump_request_list is used anywhere for dumping/writing
rg -n "crash_dump_request_list" grpc_servicer/smg_grpc_servicer/sglang/request_manager.pyRepository: lightseekorg/smg
Length of output: 988
🏁 Script executed:
# Check the class definition and inheritance of GrpcRequestManager
rg -A 5 "^class GrpcRequestManager" grpc_servicer/smg_grpc_servicer/sglang/request_manager.pyRepository: lightseekorg/smg
Length of output: 217
Missing dump_requests_before_crash method will cause AttributeError in exception handler.
The print_exception_wrapper function calls func.__self__.dump_requests_before_crash() at line 964, but this method is not defined in GrpcRequestManager. The class records requests via record_request_for_crash_dump() and maintains crash_dump_request_list, but has no method to actually dump or serialize this data before crashing. When an exception occurs and this handler is triggered, an AttributeError will be raised instead of performing the intended crash diagnostics.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@grpc_servicer/smg_grpc_servicer/sglang/request_manager.py` around lines 953 -
965, The exception handler calls func.__self__.dump_requests_before_crash() but
GrpcRequestManager lacks that method, causing an AttributeError; add a
dump_requests_before_crash(self) method on GrpcRequestManager that safely reads
the existing crash_dump_request_list (use the same locking/structure used by
record_request_for_crash_dump()), serializes or logs the stored requests (e.g.,
to logger.error or a crash dump file) with clear context, and ensure it is no-op
safe if the list is empty; alternatively, if you prefer not to add a method,
modify print_exception_wrapper to check for hasattr(...,
"dump_requests_before_crash") before calling it to avoid the AttributeError.
| # Fix CUDA multiprocessing issues - must be called before any CUDA operations | ||
| mp.set_start_method("spawn", force=True) |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
set_start_method("spawn", force=True) may mask configuration issues.
Using force=True will override any previously set start method without warning. If another part of the codebase has already set a different method, this silently overrides it. Consider logging when forcing or checking the current method first.
♻️ Suggested improvement for visibility
# Fix CUDA multiprocessing issues - must be called before any CUDA operations
+ current_method = mp.get_start_method(allow_none=True)
+ if current_method and current_method != "spawn":
+ logger.warning(
+ f"Overriding multiprocessing start method from '{current_method}' to 'spawn' for CUDA compatibility"
+ )
mp.set_start_method("spawn", force=True)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@grpc_servicer/smg_grpc_servicer/sglang/scheduler_launcher.py` around lines 71
- 72, Replace the unconditional mp.set_start_method("spawn", force=True) with a
safe check: call mp.get_start_method(allow_none=True) and only call
mp.set_start_method("spawn") if the current method is None or different, and
when you must override, log a warning (using the module logger or processLogger)
indicating the previous method to avoid silently masking configuration;
reference mp.set_start_method and mp.get_start_method in scheduler_launcher.py.
| grpc_url = f"{server_args.host}:{server_args.port}" | ||
| channel = grpc.insecure_channel( | ||
| grpc_url, | ||
| options=[ | ||
| ("grpc.max_send_message_length", 1024 * 1024 * 256), | ||
| ("grpc.max_receive_message_length", 1024 * 1024 * 256), | ||
| ], | ||
| ) |
There was a problem hiding this comment.
Warmup uses insecure channel but server may be TLS-enabled.
_execute_grpc_server_warmup always creates an insecure_channel, but if ssl_certfile and ssl_keyfile are configured, the server runs with TLS. This will cause warmup to fail with connection errors.
🐛 Proposed fix to handle TLS in warmup
def _execute_grpc_server_warmup(server_args: ServerArgs):
"""Execute warmup for gRPC server by checking health and sending test request."""
try:
# Connect to the gRPC server
grpc_url = f"{server_args.host}:{server_args.port}"
- channel = grpc.insecure_channel(
- grpc_url,
- options=[
- ("grpc.max_send_message_length", 1024 * 1024 * 256),
- ("grpc.max_receive_message_length", 1024 * 1024 * 256),
- ],
- )
+ channel_options = [
+ ("grpc.max_send_message_length", 1024 * 1024 * 256),
+ ("grpc.max_receive_message_length", 1024 * 1024 * 256),
+ ]
+ if server_args.ssl_certfile and server_args.ssl_keyfile:
+ # Use secure channel for TLS-enabled server
+ # For self-signed certs in warmup, we can use the CA cert if provided
+ if server_args.ssl_ca_certs:
+ with open(server_args.ssl_ca_certs, "rb") as f:
+ root_certs = f.read()
+ credentials = grpc.ssl_channel_credentials(root_certificates=root_certs)
+ else:
+ # Trust the server's certificate (for self-signed without CA)
+ credentials = grpc.ssl_channel_credentials()
+ channel = grpc.secure_channel(grpc_url, credentials, options=channel_options)
+ else:
+ channel = grpc.insecure_channel(grpc_url, options=channel_options)
stub = sglang_scheduler_pb2_grpc.SglangSchedulerStub(channel)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@grpc_servicer/smg_grpc_servicer/sglang/server.py` around lines 319 - 326, The
warmup currently always uses grpc.insecure_channel which fails when the server
is TLS-enabled; update _execute_grpc_server_warmup to detect
server_args.ssl_certfile (and ssl_keyfile) and, when present, load the server
certificate (server_args.ssl_certfile), create grpc.ssl_channel_credentials
using that cert as the root certificate, and use grpc.secure_channel with those
credentials and the same channel options; otherwise fallback to the existing
grpc.insecure_channel path. Ensure you reference _execute_grpc_server_warmup,
server_args.ssl_certfile, and server_args.ssl_keyfile when making the change.
| async def run_health_check(): | ||
| try: | ||
| async for _ in self.request_manager.generate_request( | ||
| obj=health_req, | ||
| request_id=rid, | ||
| ): | ||
| # Got at least one response, server is healthy | ||
| return True | ||
| except Exception as e: | ||
| logger.warning(f"Health check failed: {e}") | ||
| return False | ||
| return False | ||
|
|
||
| task = asyncio.create_task(run_health_check()) | ||
|
|
||
| # Wait for response with timeout | ||
| tic = time.time() | ||
| while time.time() < tic + HEALTH_CHECK_TIMEOUT: | ||
| await asyncio.sleep(1) | ||
| # Check if we got a response from scheduler | ||
| if self.request_manager.last_receive_tstamp > tic: | ||
| task.cancel() | ||
| # Clean up health check state | ||
| self.request_manager._cleanup_request_state(rid) | ||
| return sglang_scheduler_pb2.HealthCheckResponse( | ||
| healthy=True, message="Health check passed" | ||
| ) | ||
|
|
||
| # Timeout - server not responding | ||
| task.cancel() | ||
| self.request_manager._cleanup_request_state(rid) | ||
| logger.warning(f"Health check timeout after {HEALTH_CHECK_TIMEOUT}s") | ||
| return sglang_scheduler_pb2.HealthCheckResponse( | ||
| healthy=False, message=f"Health check timeout after {HEALTH_CHECK_TIMEOUT}s" | ||
| ) |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Health check polling loop may miss actual completion.
The health check relies on last_receive_tstamp > tic to detect scheduler response, but this could match any scheduler activity, not specifically the health check response. Additionally, the run_health_check task is cancelled even when the health check succeeds (Line 379), which may log unnecessary cancellation warnings.
Consider awaiting the task result directly with timeout instead of polling last_receive_tstamp.
♻️ Suggested improvement using task result
task = asyncio.create_task(run_health_check())
- # Wait for response with timeout
- tic = time.time()
- while time.time() < tic + HEALTH_CHECK_TIMEOUT:
- await asyncio.sleep(1)
- # Check if we got a response from scheduler
- if self.request_manager.last_receive_tstamp > tic:
- task.cancel()
- # Clean up health check state
- self.request_manager._cleanup_request_state(rid)
- return sglang_scheduler_pb2.HealthCheckResponse(
- healthy=True, message="Health check passed"
- )
-
- # Timeout - server not responding
- task.cancel()
- self.request_manager._cleanup_request_state(rid)
- logger.warning(f"Health check timeout after {HEALTH_CHECK_TIMEOUT}s")
- return sglang_scheduler_pb2.HealthCheckResponse(
- healthy=False, message=f"Health check timeout after {HEALTH_CHECK_TIMEOUT}s"
- )
+ try:
+ healthy = await asyncio.wait_for(task, timeout=HEALTH_CHECK_TIMEOUT)
+ return sglang_scheduler_pb2.HealthCheckResponse(
+ healthy=healthy,
+ message="Health check passed" if healthy else "Health check failed",
+ )
+ except TimeoutError:
+ task.cancel()
+ logger.warning(f"Health check timeout after {HEALTH_CHECK_TIMEOUT}s")
+ return sglang_scheduler_pb2.HealthCheckResponse(
+ healthy=False, message=f"Health check timeout after {HEALTH_CHECK_TIMEOUT}s"
+ )
+ finally:
+ self.request_manager._cleanup_request_state(rid)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@grpc_servicer/smg_grpc_servicer/sglang/servicer.py` around lines 358 - 392,
Replace the manual polling loop with awaiting the run_health_check task via
asyncio.wait_for to ensure we observe the specific health-check result rather
than unrelated scheduler activity: create the task from run_health_check
(function run_health_check uses self.request_manager.generate_request), then use
try/except around await asyncio.wait_for(task, timeout=HEALTH_CHECK_TIMEOUT) to
return a healthy response on True, handle asyncio.TimeoutError to cancel the
task and return the timeout unhealthy response, and handle other exceptions to
log and return unhealthy; always call
self.request_manager._cleanup_request_state(rid) in a finally block so request
state is cleaned up whether the check succeeds, times out, or errors, and avoid
cancelling the task when it has already completed.
| def _convert_kv_event(self, event) -> common_pb2.KvCacheEvent | None: | ||
| """Convert a single raw KV event to proto KvCacheEvent.""" | ||
| self._kv_event_id_counter += 1 | ||
| event_id = self._kv_event_id_counter | ||
|
|
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Event ID counter is not thread-safe for concurrent subscribers.
The _kv_event_id_counter is incremented without synchronization. If multiple clients call SubscribeKvEvents concurrently, this could lead to duplicate event IDs across streams. Consider using itertools.count() or asyncio.Lock if unique IDs are required across all streams, or document that IDs are per-stream.
♻️ Option 1: Per-stream counter (document current behavior)
If event IDs only need to be unique within a single stream, move the counter into the SubscribeKvEvents method as a local variable.
async def SubscribeKvEvents(
self,
request: common_pb2.SubscribeKvEventsRequest,
context: grpc.aio.ServicerContext,
) -> AsyncIterator[common_pb2.KvEventBatch]:
...
+ event_id_counter = 0
...
- yield self._convert_kv_event_batch(raw_batch, zmq_seq)
+ event_id_counter, proto_batch = self._convert_kv_event_batch(
+ raw_batch, zmq_seq, event_id_counter
+ )
+ yield proto_batch🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@grpc_servicer/smg_grpc_servicer/sglang/servicer.py` around lines 637 - 641,
The _kv_event_id_counter is incremented unsafely in _convert_kv_event which can
produce duplicate IDs across concurrent SubscribeKvEvents streams; either make
the counter thread-safe by replacing the shared integer with a global iterator
(e.g., itertools.count()) or protect increments with an asyncio.Lock around
accesses to _kv_event_id_counter in _convert_kv_event, or alternatively move the
counter to be a per-stream local variable inside SubscribeKvEvents if IDs only
need to be unique per stream—update the implementation of _convert_kv_event and
SubscribeKvEvents accordingly to use the chosen approach and ensure the
unique-symbols _kv_event_id_counter, _convert_kv_event, and SubscribeKvEvents
are adjusted together.
| @staticmethod | ||
| def _decode_tensor_data(tensor_data): | ||
| """Decode a proto TensorData message into a torch.Tensor.""" | ||
| dtype_map = {"float32": np.float32, "int64": np.int64} | ||
| np_dtype = dtype_map.get(tensor_data.dtype, np.float32) | ||
| shape = list(tensor_data.shape) | ||
| arr = np.frombuffer(tensor_data.data, dtype=np_dtype).reshape(shape) | ||
| return torch.from_numpy(arr) |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Limited dtype support in tensor decoding.
The _decode_tensor_data method only supports float32 and int64 dtypes, defaulting to float32 for unknown types. If the proto can contain other dtypes (e.g., float16, bfloat16 for optimized models), this silent fallback could cause data corruption.
🛡️ Proposed fix to handle unknown dtypes explicitly
`@staticmethod`
def _decode_tensor_data(tensor_data):
"""Decode a proto TensorData message into a torch.Tensor."""
- dtype_map = {"float32": np.float32, "int64": np.int64}
- np_dtype = dtype_map.get(tensor_data.dtype, np.float32)
+ dtype_map = {
+ "float32": np.float32,
+ "float64": np.float64,
+ "int32": np.int32,
+ "int64": np.int64,
+ "float16": np.float16,
+ }
+ np_dtype = dtype_map.get(tensor_data.dtype)
+ if np_dtype is None:
+ raise ValueError(f"Unsupported tensor dtype: {tensor_data.dtype}")
shape = list(tensor_data.shape)
arr = np.frombuffer(tensor_data.data, dtype=np_dtype).reshape(shape)
return torch.from_numpy(arr)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@grpc_servicer/smg_grpc_servicer/sglang/servicer.py` around lines 759 - 766,
The _decode_tensor_data function currently maps only "float32" and "int64"
(dtype_map) and silently defaults unknown tensor_data.dtype to float32, which
can corrupt data; update _decode_tensor_data to: expand dtype_map to include
other expected proto dtypes (e.g., "float16", "bfloat16"/"bf16", "int32",
"uint8", "bool", etc.) mapped to the correct numpy dtypes, use that mapping to
select np_dtype from tensor_data.dtype, and if tensor_data.dtype is not in the
map raise a clear exception (ValueError) naming tensor_data.dtype so callers
fail fast rather than silently decoding as float32; keep the existing
np.frombuffer(...).reshape(...) and torch.from_numpy(arr) flow.
| def abort_code_from_output(output: dict) -> grpc.StatusCode: | ||
| """Map a scheduler error output to the appropriate gRPC status code.""" | ||
| finish_reason = output.get("meta_info", {}).get("finish_reason") | ||
| if isinstance(finish_reason, dict): | ||
| status_code = finish_reason.get("status_code") | ||
| if status_code is not None: | ||
| return _HTTP_TO_GRPC_CODE.get(status_code, grpc.StatusCode.INTERNAL) | ||
| return grpc.StatusCode.INTERNAL |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Consider handling None input defensively.
If output is None, output.get("meta_info", {}) will raise AttributeError. Consider adding a guard:
♻️ Suggested defensive handling
def abort_code_from_output(output: dict) -> grpc.StatusCode:
"""Map a scheduler error output to the appropriate gRPC status code."""
+ if not output:
+ return grpc.StatusCode.INTERNAL
finish_reason = output.get("meta_info", {}).get("finish_reason")📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def abort_code_from_output(output: dict) -> grpc.StatusCode: | |
| """Map a scheduler error output to the appropriate gRPC status code.""" | |
| finish_reason = output.get("meta_info", {}).get("finish_reason") | |
| if isinstance(finish_reason, dict): | |
| status_code = finish_reason.get("status_code") | |
| if status_code is not None: | |
| return _HTTP_TO_GRPC_CODE.get(status_code, grpc.StatusCode.INTERNAL) | |
| return grpc.StatusCode.INTERNAL | |
| def abort_code_from_output(output: dict) -> grpc.StatusCode: | |
| """Map a scheduler error output to the appropriate gRPC status code.""" | |
| if not output: | |
| return grpc.StatusCode.INTERNAL | |
| finish_reason = output.get("meta_info", {}).get("finish_reason") | |
| if isinstance(finish_reason, dict): | |
| status_code = finish_reason.get("status_code") | |
| if status_code is not None: | |
| return _HTTP_TO_GRPC_CODE.get(status_code, grpc.StatusCode.INTERNAL) | |
| return grpc.StatusCode.INTERNAL |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@grpc_servicer/smg_grpc_servicer/sglang/utils.py` around lines 14 - 21, The
function abort_code_from_output assumes `output` is a dict and will raise if
`None`; add a defensive guard at the start of abort_code_from_output to handle
non-dict/None inputs (e.g., if output is None or not isinstance(output, dict):
return grpc.StatusCode.INTERNAL), then proceed to access meta_info and use
_HTTP_TO_GRPC_CODE as before so the function never calls .get on a None value.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 8b2d833725
ℹ️ 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".
| logger.warning("Failed to decode KV event batch: %s", e) | ||
| continue | ||
|
|
||
| yield self._convert_kv_event_batch(raw_batch, zmq_seq) |
There was a problem hiding this comment.
SubscribeKvEvents ignores request.start_sequence_number and always forwards the live ZMQ sequence (zmq_seq) immediately, so reconnects cannot replay from the caller’s last acknowledged sequence. In the gateway monitor path (subscribe_kv_events(last_seq)), a transient disconnect after sequence N will reconnect at a much larger live sequence and repeatedly trigger gap detection/reconnect, preventing KV cache state from converging.
Useful? React with 👍 / 👎.
| dtype_map = {"float32": np.float32, "int64": np.int64} | ||
| np_dtype = dtype_map.get(tensor_data.dtype, np.float32) |
There was a problem hiding this comment.
_decode_tensor_data only maps float32 and int64, then silently falls back to float32 for unknown dtypes. The protobuf TensorData contract explicitly includes uint32, so any multimodal tensor sent as uint32 is reinterpreted as float32 bytes, corrupting tensor values and breaking multimodal request handling for affected models.
Useful? React with 👍 / 👎.
| async for _ in self.request_manager.generate_request( | ||
| obj=health_req, | ||
| request_id=rid, |
There was a problem hiding this comment.
In the embedding-model branch, HealthCheck builds a TokenizedEmbeddingReqInput but sends it through generate_request; that path waits for a finished flag that embedding outputs never set, so the health task cannot complete normally and is only canceled by the timestamp polling loop. This makes probe success depend on generic scheduler traffic timing rather than a completed embedding health request.
Useful? React with 👍 / 👎.
1 participant
Summary
smg-grpc-servicerto use optional extras instead of hard dependencies:[vllm]for vLLM,[sglang]for SGLangsmg_grpc_servicer/sglang/package — extracts sglang's gRPC server code (servicer, request manager, health servicer, scheduler launcher, server entrypoint) into smgBreaking change
smg-grpc-servicerno longer has a hard dependency onvllm. Consumers that relied on the implicit vllm dep must update tosmg-grpc-servicer[vllm].Test plan
pip install -e grpc_servicer/succeeds without vllm or sglang installedpip install smg-grpc-servicer[vllm]pulls in vllmpip install smg-grpc-servicer[sglang]pulls in sglang--grpc-modeSummary by CodeRabbit
Release Notes
New Features
Documentation
Chores
smg-grpc-servicer[vllm]orsmg-grpc-servicer[sglang]as needed