From 654fc02cf1bcc6db19dfdb860ede49f8b2a5afb3 Mon Sep 17 00:00:00 2001 From: Simo Lin Date: Fri, 13 Mar 2026 09:13:29 -0700 Subject: [PATCH] [gRPC] Extract gRPC servicer into standalone package (#20478) Signed-off-by: Simo Lin --- 3rdparty/amd/wheel/sglang/pyproject.toml | 4 +- python/pyproject.toml | 5 +- python/pyproject_cpu.toml | 4 +- python/pyproject_npu.toml | 4 +- python/pyproject_other.toml | 4 +- python/pyproject_xpu.toml | 4 +- python/sglang/srt/entrypoints/grpc_server.py | 1548 +---------------- .../sglang/srt/grpc/grpc_request_manager.py | 1022 ----------- python/sglang/srt/grpc/health_servicer.py | 189 -- python/sglang/srt/grpc/scheduler_launcher.py | 198 --- python/sglang/srt/grpc/utils.py | 21 - 11 files changed, 15 insertions(+), 2988 deletions(-) delete mode 100644 python/sglang/srt/grpc/grpc_request_manager.py delete mode 100644 python/sglang/srt/grpc/health_servicer.py delete mode 100644 python/sglang/srt/grpc/scheduler_launcher.py delete mode 100644 python/sglang/srt/grpc/utils.py diff --git a/3rdparty/amd/wheel/sglang/pyproject.toml b/3rdparty/amd/wheel/sglang/pyproject.toml index 0d9f93154..6f0cbad6a 100644 --- a/3rdparty/amd/wheel/sglang/pyproject.toml +++ b/3rdparty/amd/wheel/sglang/pyproject.toml @@ -63,9 +63,7 @@ runtime_common = [ "uvicorn", "uvloop", "xgrammar==0.1.27", - "smg-grpc-proto>=0.3.3", - "grpcio>=1.78.0", - "grpcio-reflection>=1.78.0", + "smg-grpc-servicer>=0.5.0", ] # ROCm specific packages (https://repo.radeon.com/rocm/manylinux/) diff --git a/python/pyproject.toml b/python/pyproject.toml index 39703571d..cde1cab99 100755 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -78,10 +78,7 @@ dependencies = [ "watchfiles", "xgrammar==0.1.27", - "smg-grpc-proto>=0.4.1", - "grpcio>=1.78.0", - "grpcio-reflection>=1.78.0", - "grpcio-health-checking>=1.78.0", + "smg-grpc-servicer>=0.5.0", ] [[tool.uv.index]] diff --git a/python/pyproject_cpu.toml b/python/pyproject_cpu.toml index 0087f823a..ae22b1127 100644 --- a/python/pyproject_cpu.toml +++ b/python/pyproject_cpu.toml @@ -67,9 +67,7 @@ dependencies = [ "uvicorn", "uvloop", "xgrammar==0.1.27", - "smg-grpc-proto>=0.4.1", - "grpcio>=1.78.0", - "grpcio-reflection>=1.78.0", + "smg-grpc-servicer>=0.5.0", ] [project.optional-dependencies] diff --git a/python/pyproject_npu.toml b/python/pyproject_npu.toml index 9f7caac12..94417f6d9 100644 --- a/python/pyproject_npu.toml +++ b/python/pyproject_npu.toml @@ -61,9 +61,7 @@ dependencies = [ "uvicorn", "uvloop", "xgrammar==0.1.27", - "smg-grpc-proto>=0.4.1", - "grpcio>=1.78.0", - "grpcio-reflection>=1.78.0", + "smg-grpc-servicer>=0.5.0", ] [project.optional-dependencies] diff --git a/python/pyproject_other.toml b/python/pyproject_other.toml index 9008341e6..c7a68e913 100755 --- a/python/pyproject_other.toml +++ b/python/pyproject_other.toml @@ -63,9 +63,7 @@ runtime_common = [ "uvicorn", "uvloop", "xgrammar==0.1.27", - "smg-grpc-proto>=0.4.1", - "grpcio>=1.78.0", - "grpcio-reflection>=1.78.0", + "smg-grpc-servicer>=0.5.0", ] tracing = [ diff --git a/python/pyproject_xpu.toml b/python/pyproject_xpu.toml index e43514f79..113bf3eda 100644 --- a/python/pyproject_xpu.toml +++ b/python/pyproject_xpu.toml @@ -66,9 +66,7 @@ dependencies = [ "uvicorn", "uvloop", # "xgrammar==0.1.24", , xgrammar depends on CUDA PyTorch and Triton only - "smg-grpc-proto>=0.4.1", - "grpcio>=1.78.0", - "grpcio-reflection>=1.78.0", + "smg-grpc-servicer>=0.5.0", ] [project.optional-dependencies] diff --git a/python/sglang/srt/entrypoints/grpc_server.py b/python/sglang/srt/entrypoints/grpc_server.py index c85aeace2..674431bf4 100644 --- a/python/sglang/srt/entrypoints/grpc_server.py +++ b/python/sglang/srt/entrypoints/grpc_server.py @@ -1,1545 +1,15 @@ """ -Standalone gRPC Server for SGLang - Fully separated from HTTP server. -Uses GrpcRequestManager for orchestration without tokenization. +Thin gRPC server wrapper — delegates to smg-grpc-servicer package. """ -import asyncio -import dataclasses -import json -import logging -import os -import signal -import threading -import time -from concurrent import futures -from datetime import datetime, timezone -from typing import AsyncIterator, Dict, Optional -import grpc -import msgspec -import numpy as np -import torch -import zmq -import zmq.asyncio -from google.protobuf.json_format import MessageToDict -from google.protobuf.struct_pb2 import Struct -from google.protobuf.timestamp_pb2 import Timestamp -from grpc_health.v1 import health_pb2_grpc -from grpc_reflection.v1alpha import reflection -from smg_grpc_proto import sglang_scheduler_pb2, sglang_scheduler_pb2_grpc -from smg_grpc_proto.generated import common_pb2 - -import sglang -from sglang.srt.configs.model_config import ModelConfig -from sglang.srt.disaggregation.kv_events import ( - AllBlocksCleared, - BlockRemoved, - BlockStored, - KVEventBatch, - KVEventsConfig, - ZmqEventPublisher, -) -from sglang.srt.disaggregation.utils import FAKE_BOOTSTRAP_HOST, DisaggregationMode -from sglang.srt.grpc.grpc_request_manager import GrpcRequestManager -from sglang.srt.grpc.health_servicer import SGLangHealthServicer -from sglang.srt.grpc.scheduler_launcher import launch_scheduler_process_only -from sglang.srt.grpc.utils import abort_code_from_output -from sglang.srt.managers.disagg_service import start_disagg_service -from sglang.srt.managers.io_struct import ( - GetLoadsReqOutput, - TokenizedEmbeddingReqInput, - TokenizedGenerateReqInput, -) -from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem -from sglang.srt.sampling.sampling_params import SamplingParams as SGLSamplingParams -from sglang.srt.server_args import ServerArgs -from sglang.srt.utils import kill_process_tree -from sglang.utils import get_exception_traceback - -logger = logging.getLogger(__name__) -HEALTH_CHECK_TIMEOUT = int(os.getenv("SGLANG_HEALTH_CHECK_TIMEOUT", 20)) - - -def _convert_loads_to_protobuf( - result: GetLoadsReqOutput, -) -> sglang_scheduler_pb2.SchedulerLoad: - """Convert GetLoadsReqOutput dataclass to protobuf SchedulerLoad message.""" - scheduler_load = sglang_scheduler_pb2.SchedulerLoad( - dp_rank=result.dp_rank, - num_running_reqs=result.num_running_reqs, - num_waiting_reqs=result.num_waiting_reqs, - num_total_reqs=result.num_running_reqs + result.num_waiting_reqs, - num_used_tokens=result.num_used_tokens, - max_total_num_tokens=result.max_total_num_tokens, - token_usage=result.token_usage, - gen_throughput=result.gen_throughput, - cache_hit_rate=result.cache_hit_rate, - utilization=result.utilization, - max_running_requests=result.max_running_requests, - ) - - # Add optional sections using CopyFrom for proper protobuf assignment - if result.memory: - scheduler_load.memory.CopyFrom( - sglang_scheduler_pb2.MemoryMetrics( - weight_gb=result.memory.weight_gb, - kv_cache_gb=result.memory.kv_cache_gb, - graph_gb=result.memory.graph_gb, - token_capacity=result.memory.token_capacity, - ) - ) - - if result.speculative: - scheduler_load.speculative.CopyFrom( - sglang_scheduler_pb2.SpeculativeMetrics( - accept_length=result.speculative.accept_length, - accept_rate=result.speculative.accept_rate, - ) - ) - - if result.lora: - scheduler_load.lora.CopyFrom( - sglang_scheduler_pb2.LoRAMetrics( - slots_used=result.lora.slots_used, - slots_total=result.lora.slots_total, - utilization=result.lora.utilization, - ) - ) - - if result.disaggregation: - scheduler_load.disaggregation.CopyFrom( - sglang_scheduler_pb2.DisaggregationMetrics( - mode=result.disaggregation.mode, - prefill_prealloc_queue_reqs=result.disaggregation.prefill_prealloc_queue_reqs, - prefill_inflight_queue_reqs=result.disaggregation.prefill_inflight_queue_reqs, - decode_prealloc_queue_reqs=result.disaggregation.decode_prealloc_queue_reqs, - decode_transfer_queue_reqs=result.disaggregation.decode_transfer_queue_reqs, - decode_retracted_queue_reqs=result.disaggregation.decode_retracted_queue_reqs, - kv_transfer_speed_gb_s=result.disaggregation.kv_transfer_speed_gb_s, - kv_transfer_latency_ms=result.disaggregation.kv_transfer_latency_ms, - ) - ) - - if result.queues: - scheduler_load.queues.CopyFrom( - sglang_scheduler_pb2.QueueMetrics( - waiting=result.queues.waiting, - grammar=result.queues.grammar, - paused=result.queues.paused, - retracted=result.queues.retracted, - ) - ) - - return scheduler_load - - -def _compute_aggregate_protobuf( - loads: list, -) -> sglang_scheduler_pb2.AggregateMetrics: - """Compute aggregate metrics from list of SchedulerLoad protobuf messages.""" - if not loads: - return sglang_scheduler_pb2.AggregateMetrics() - - n = len(loads) - total_running = sum(load.num_running_reqs for load in loads) - total_waiting = sum(load.num_waiting_reqs for load in loads) - - return sglang_scheduler_pb2.AggregateMetrics( - total_running_reqs=total_running, - total_waiting_reqs=total_waiting, - total_reqs=total_running + total_waiting, - avg_token_usage=round(sum(load.token_usage for load in loads) / n, 4), - avg_throughput=round(sum(load.gen_throughput for load in loads) / n, 2), - avg_utilization=round(sum(load.utilization for load in loads) / n, 4), - ) - - -class SGLangSchedulerServicer(sglang_scheduler_pb2_grpc.SglangSchedulerServicer): - """ - Standalone gRPC service implementation using GrpcRequestManager. - Fully separated from HTTP server with its own process and no shared globals. - """ - - def __init__( - self, - request_manager: GrpcRequestManager, - server_args: ServerArgs, - model_info: Dict, - scheduler_info: Dict, - health_servicer: Optional[SGLangHealthServicer] = None, - ): - """Initialize the standalone gRPC service.""" - self.request_manager = request_manager - self.server_args = server_args - self.model_info = model_info - self.scheduler_info = scheduler_info - self.start_time = time.time() - self.health_servicer = health_servicer - self.mm_receiver = None - if ( - self.server_args.language_only - and self.server_args.encoder_transfer_backend == "zmq_to_scheduler" - ): - from sglang.srt.disaggregation import encode_receiver as mm_receiver - - self.mm_receiver = mm_receiver.create_mm_receiver(self.server_args) - - # Parse KV events config for SubscribeKvEvents support - self._kv_events_config: Optional[KVEventsConfig] = None - self._kv_event_id_counter = 0 - if server_args.kv_events_config: - try: - self._kv_events_config = KVEventsConfig.from_cli( - server_args.kv_events_config - ) - if self._kv_events_config.publisher != "zmq": - logger.info( - "KV events publisher is '%s', SubscribeKvEvents disabled", - self._kv_events_config.publisher, - ) - self._kv_events_config = None - else: - logger.info( - "KV events enabled: endpoint=%s", - self._kv_events_config.endpoint, - ) - except Exception as e: - logger.warning("Failed to parse kv_events_config: %s", e) - - # Start the request manager's event loop using auto_create_handle_loop - self.request_manager.auto_create_handle_loop() - - logger.info("gRPC scheduler servicer initialized") - - async def Generate( - self, - request: sglang_scheduler_pb2.GenerateRequest, - context: grpc.aio.ServicerContext, - ) -> AsyncIterator[sglang_scheduler_pb2.GenerateResponse]: - """Handle generation requests with streaming responses.""" - logger.info(f"Receive generation request: {request.request_id}") - - try: - # Convert gRPC request to internal format - tokenized_req = self._convert_generate_request(request) - self._handle_epd_disaggregation_encode_request(request, tokenized_req) - - # Submit to request manager (automatically handles n>1) - response_generator = self.request_manager.generate_request( - obj=tokenized_req, - request_id=request.request_id, - grpc_context=context, - ) - - async for output in response_generator: - # Handle batch responses (for n>1 non-streaming) - if isinstance(output, list): - for batch_output in output: - if "error" in batch_output: - await context.abort( - abort_code_from_output(batch_output), - batch_output["error"], - ) - else: - # All non-error batch outputs are final responses - yield self._create_completion_response( - request.request_id, batch_output - ) - else: - # Handle single response (for streaming or n=1 non-streaming) - if "error" in output: - await context.abort( - abort_code_from_output(output), - output["error"], - ) - elif request.stream: - yield self._create_chunk_response(request.request_id, output) - if output.get("finished", False): - yield self._create_completion_response( - request.request_id, output - ) - else: - # Non-streaming n=1: single completion response - yield self._create_completion_response( - request.request_id, output - ) - - except grpc.aio.AbortError: - raise - except ValueError as e: - logger.warning(f"Generate invalid request {request.request_id}: {e}") - await context.abort(grpc.StatusCode.INVALID_ARGUMENT, str(e)) - except Exception as e: - logger.error( - f"Generate failed for request {request.request_id}: {e}\n" - f"{get_exception_traceback()}" - ) - await context.abort(grpc.StatusCode.INTERNAL, str(e)) - - async def Embed( - self, - request: sglang_scheduler_pb2.EmbedRequest, - context: grpc.aio.ServicerContext, - ) -> sglang_scheduler_pb2.EmbedResponse: - """Handle embedding requests.""" - logger.info(f"Receive embedding request: {request.request_id}") - - try: - tokenized_req = self._convert_embed_request(request) - - future = await self.request_manager.embedding_request( - obj=tokenized_req, - request_id=request.request_id, - ) - - result = await future - - return sglang_scheduler_pb2.EmbedResponse( - request_id=request.request_id, - complete=sglang_scheduler_pb2.EmbedComplete( - embedding=result["embedding"], - prompt_tokens=result.get("prompt_tokens", 0), - cached_tokens=0, - embedding_dim=len(result["embedding"]), - ), - ) - - except grpc.aio.AbortError: - raise - except ValueError as e: - logger.warning(f"Embed invalid request {request.request_id}: {e}") - await context.abort(grpc.StatusCode.INVALID_ARGUMENT, str(e)) - except Exception as e: - logger.error( - f"Embed failed for request {request.request_id}: {e}\n" - f"{get_exception_traceback()}" - ) - await context.abort(grpc.StatusCode.INTERNAL, str(e)) - - async def HealthCheck( - self, - request: sglang_scheduler_pb2.HealthCheckRequest, - context: grpc.aio.ServicerContext, - ) -> sglang_scheduler_pb2.HealthCheckResponse: - """ - Check the health of the inference server by sending a special request to generate one token. - Similar to HTTP server's /health endpoint. - """ - rid = f"HEALTH_CHECK_{time.time()}" - logger.info(f"Receive health check request: {rid}") - - if self.request_manager.gracefully_exit: - logger.info( - "Health check request received during shutdown. Returning unhealthy." - ) - return sglang_scheduler_pb2.HealthCheckResponse( - healthy=False, message="Server is shutting down" - ) - - # Create a special health check request - sampling_params = SGLSamplingParams(max_new_tokens=1, temperature=0.0) - sampling_params.normalize(tokenizer=None) - - # Create health check request - is_generation = self.scheduler_info.get("is_generation") - if is_generation is None: - is_generation = not self.server_args.is_embedding - - if is_generation: - health_req = TokenizedGenerateReqInput( - rid=rid, - input_text="", - input_ids=[0], - sampling_params=sampling_params, - return_logprob=False, - logprob_start_len=-1, - top_logprobs_num=0, - stream=False, - mm_inputs=None, - token_ids_logprob=None, - ) - # Set disaggregation params if needed - if self.server_args.disaggregation_mode != DisaggregationMode.NULL.value: - health_req.bootstrap_host = FAKE_BOOTSTRAP_HOST - health_req.bootstrap_room = 0 - else: - sampling_params.max_new_tokens = 0 - health_req = TokenizedEmbeddingReqInput( - rid=rid, - input_text="", - input_ids=[0], - image_inputs={"mm_items": []}, - token_type_ids=[0], - sampling_params=sampling_params, - ) - - # Submit health check request - 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" - ) - - async def Abort( - self, - request: sglang_scheduler_pb2.AbortRequest, - _context: grpc.aio.ServicerContext, - ) -> sglang_scheduler_pb2.AbortResponse: - """Abort an ongoing request.""" - logger.info(f"Receive abort request: {request.request_id}") - - try: - success = await self.request_manager.abort_request(request.request_id) - - return sglang_scheduler_pb2.AbortResponse( - success=success, - message=f"Request {request.request_id} {'aborted' if success else 'not found'}", - ) - except Exception as e: - logger.error( - f"Abort failed for request {request.request_id}: {e}\n" - f"{get_exception_traceback()}" - ) - return sglang_scheduler_pb2.AbortResponse( - success=False, - message=str(e), - ) - - async def GetModelInfo( - self, - _request: sglang_scheduler_pb2.GetModelInfoRequest, - _context: grpc.aio.ServicerContext, - ) -> sglang_scheduler_pb2.GetModelInfoResponse: - """Get model information.""" - logger.debug("Receive model info request") - - is_generation = self.scheduler_info.get("is_generation") - if is_generation is None: - is_generation = not self.server_args.is_embedding - - return sglang_scheduler_pb2.GetModelInfoResponse( - model_path=self.server_args.model_path, - tokenizer_path=self.server_args.tokenizer_path or "", - is_generation=is_generation, - preferred_sampling_params=( - self.server_args.preferred_sampling_params or "" - ), - weight_version=self.server_args.weight_version or "", - served_model_name=self.server_args.served_model_name, - max_context_length=self.model_info["max_context_length"], - vocab_size=self.model_info["vocab_size"], - supports_vision=self.model_info["supports_vision"], - model_type=self.model_info.get("model_type") or "", - architectures=self.model_info.get("architectures") or [], - eos_token_ids=self.model_info["eos_token_ids"], - pad_token_id=self.model_info["pad_token_id"], - bos_token_id=self.model_info["bos_token_id"], - max_req_input_len=self.model_info["max_req_input_len"], - # Classification model support - id2label_json=self.model_info.get("id2label_json") or "", - num_labels=self.model_info.get("num_labels") or 0, - ) - - async def GetServerInfo( - self, - _request: sglang_scheduler_pb2.GetServerInfoRequest, - _context: grpc.aio.ServicerContext, - ) -> sglang_scheduler_pb2.GetServerInfoResponse: - """Get server information.""" - logger.debug("Receive server info request") - - server_args_dict = dataclasses.asdict(self.server_args) - server_args_struct = Struct() - - def make_serializable(obj): - if obj is None: - return None - elif isinstance(obj, (str, int, float, bool)): - return obj - elif isinstance(obj, (list, tuple, set)): - return [make_serializable(item) for item in obj] - elif isinstance(obj, dict): - return {k: make_serializable(v) for k, v in obj.items()} - else: - return str(obj) - - serializable_args = make_serializable(server_args_dict) - server_args_struct.update(serializable_args) - - # Convert scheduler_info to Struct - scheduler_info_struct = Struct() - scheduler_info_struct.update(self.scheduler_info) - - # Get runtime state from request manager - manager_state = self.request_manager.get_server_info() - - # Calculate uptime - uptime = time.time() - self.start_time - - # Create timestamp - start_timestamp = Timestamp() - start_timestamp.FromSeconds(int(self.start_time)) - - return sglang_scheduler_pb2.GetServerInfoResponse( - server_args=server_args_struct, - scheduler_info=scheduler_info_struct, - active_requests=manager_state["active_requests"], - is_paused=manager_state["paused"], - last_receive_timestamp=manager_state["last_receive_time"], - uptime_seconds=uptime, - sglang_version=sglang.__version__, - server_type="grpc", - start_time=start_timestamp, - ) - - async def GetLoads( - self, - request: sglang_scheduler_pb2.GetLoadsRequest, - context: grpc.aio.ServicerContext, - ) -> sglang_scheduler_pb2.GetLoadsResponse: - """ - Get comprehensive load metrics for all DP ranks. - - Uses the communicator pattern to fetch real-time metrics, - providing full parity with the HTTP /v1/loads endpoint. - """ - logger.debug("Receive get loads request") - - include = list(request.include) if request.include else ["all"] - dp_rank = request.dp_rank if request.HasField("dp_rank") else None - - try: - results = await self.request_manager.get_loads( - include=include, dp_rank=dp_rank - ) - except ValueError as e: - # Validation error (e.g., invalid include sections) - context.set_code(grpc.StatusCode.INVALID_ARGUMENT) - context.set_details(str(e)) - return sglang_scheduler_pb2.GetLoadsResponse() - except asyncio.TimeoutError: - context.set_code(grpc.StatusCode.DEADLINE_EXCEEDED) - context.set_details("Timeout waiting for scheduler response") - return sglang_scheduler_pb2.GetLoadsResponse() - except Exception as e: - logger.error(f"GetLoads failed: {e}\n{get_exception_traceback()}") - context.set_code(grpc.StatusCode.INTERNAL) - context.set_details(f"Failed to get load metrics: {e}") - return sglang_scheduler_pb2.GetLoadsResponse() - - loads = [_convert_loads_to_protobuf(r) for r in results] - - return sglang_scheduler_pb2.GetLoadsResponse( - timestamp=datetime.now(timezone.utc).isoformat(), - version=sglang.__version__, - dp_rank_count=len(loads), - loads=loads, - aggregate=_compute_aggregate_protobuf(loads), - ) - - async def SubscribeKvEvents( - self, - request: common_pb2.SubscribeKvEventsRequest, - context: grpc.aio.ServicerContext, - ) -> AsyncIterator[common_pb2.KvEventBatch]: - """Bridge internal ZMQ KV cache events to gRPC server-streaming. - - Uses the ZMQ publisher's native sequence numbers as gRPC sequence - numbers directly. - """ - if self._kv_events_config is None: - await context.abort( - grpc.StatusCode.UNIMPLEMENTED, - "KV cache events not enabled. Start SGLang with " - '--kv-events-config \'{"publisher": "zmq"}\'', - ) - return - - config = self._kv_events_config - - # Resolve the PUB endpoint to a connectable address. - # The publisher binds to e.g. "tcp://*:5557"; we connect to localhost. - pub_endpoint = config.endpoint.replace("*", "127.0.0.1") - - # For DP attention, each rank publishes on port + rank with - # independent sequence counters. Subscribing to multiple ranks - # on one socket interleaves independent counters, breaking gap - # detection. For now, subscribe to rank 0 only. - # TODO(phase2): per-rank virtual workers or merged renumbering. - pub_endpoint = ZmqEventPublisher.offset_endpoint_port(pub_endpoint, 0) - - zmq_ctx = zmq.asyncio.Context.instance() - sub_socket = zmq_ctx.socket(zmq.SUB) - sub_socket.subscribe(config.topic.encode("utf-8")) - sub_socket.connect(pub_endpoint) - - logger.info("SubscribeKvEvents: connected to ZMQ endpoint %s", pub_endpoint) - - # Send response headers immediately so the tonic client's - # subscribe_kv_events().await resolves without waiting for the first - # yielded event (grpc.aio defers headers until first yield otherwise). - await context.send_initial_metadata(()) - - decoder = msgspec.msgpack.Decoder(KVEventBatch) - - # Stream live events using the ZMQ publisher's native seq numbers. - try: - while not context.cancelled(): - try: - frames = await asyncio.wait_for( - sub_socket.recv_multipart(), timeout=1.0 - ) - except asyncio.TimeoutError: - continue - - # ZMQ multipart: [topic, seq_bytes, payload] - if len(frames) < 3: - continue - - zmq_seq = int.from_bytes(frames[1], "big") - payload = frames[2] - - try: - raw_batch = decoder.decode(payload) - except Exception as e: - logger.warning("Failed to decode KV event batch: %s", e) - continue - - yield self._convert_kv_event_batch(raw_batch, zmq_seq) - except asyncio.CancelledError: - pass - finally: - sub_socket.close(linger=0) - logger.info("SubscribeKvEvents: stream closed") - - def _convert_kv_event_batch( - self, raw_batch: KVEventBatch, seq_num: int - ) -> common_pb2.KvEventBatch: - """Convert a ZMQ KVEventBatch to proto KvEventBatch.""" - proto_batch = common_pb2.KvEventBatch( - sequence_number=seq_num, - timestamp=raw_batch.ts, - ) - if raw_batch.attn_dp_rank is not None: - proto_batch.dp_rank = raw_batch.attn_dp_rank - - for event in raw_batch.events: - proto_event = self._convert_kv_event(event) - if proto_event is not None: - proto_batch.events.append(proto_event) - - return proto_batch - - def _convert_kv_event(self, event) -> Optional[common_pb2.KvCacheEvent]: - """Convert a single raw KV event to proto KvCacheEvent.""" - self._kv_event_id_counter += 1 - event_id = self._kv_event_id_counter - - if isinstance(event, BlockStored): - # SGLang emits one BlockStored per page with block_hashes=[single_hash] - # and token_ids containing only that page's tokens. - blocks = [] - for i, bh in enumerate(event.block_hashes): - start = i * event.block_size - end = start + event.block_size - block = common_pb2.KvBlock( - block_hash=bh, - token_ids=event.token_ids[start:end], - block_size=event.block_size, - ) - if event.lora_id is not None: - block.lora_id = event.lora_id - blocks.append(block) - - stored = common_pb2.KvBlocksStored(blocks=blocks) - if event.parent_block_hash is not None: - stored.parent_block_hash = event.parent_block_hash - - return common_pb2.KvCacheEvent(event_id=event_id, stored=stored) - - elif isinstance(event, BlockRemoved): - return common_pb2.KvCacheEvent( - event_id=event_id, - removed=common_pb2.KvBlocksRemoved(block_hashes=event.block_hashes), - ) - - elif isinstance(event, AllBlocksCleared): - return common_pb2.KvCacheEvent( - event_id=event_id, cleared=common_pb2.KvCacheCleared() - ) - - return None - - def _handle_epd_disaggregation_encode_request( - self, - grpc_req: sglang_scheduler_pb2.GenerateRequest, - tokenized_req: TokenizedGenerateReqInput, - ) -> None: - if not self.mm_receiver: - return - - image_urls = list(grpc_req.mm_inputs.image_urls) - if not image_urls: - return - - encode_req = self.mm_receiver.build_and_send_encode_request( - image_urls=image_urls, - rid=grpc_req.request_id, - ) - tokenized_req.need_wait_for_image = bool(encode_req.need_wait_for_image) - tokenized_req.num_items_assigned = encode_req.num_items_assigned - - # Helper methods for request/response conversion - - def _convert_generate_request( - self, grpc_req: sglang_scheduler_pb2.GenerateRequest - ) -> TokenizedGenerateReqInput: - """Convert gRPC GenerateRequest to internal format.""" - - # Extract tokenized input - if not grpc_req.HasField("tokenized"): - raise ValueError("Tokenized input must be provided") - - input_text = grpc_req.tokenized.original_text - input_ids = list(grpc_req.tokenized.input_ids) - - # Convert sampling params - sampling_params = self._convert_sampling_params(grpc_req.sampling_params) - sampling_params.normalize(tokenizer=None) - - # Extract disaggregated params if present - bootstrap_host = None - bootstrap_port = None - bootstrap_room = None - if grpc_req.HasField("disaggregated_params"): - # Don't use 'or None' as it treats 0 as falsy - bootstrap_host = ( - grpc_req.disaggregated_params.bootstrap_host - if grpc_req.disaggregated_params.bootstrap_host - else None - ) - bootstrap_port = ( - grpc_req.disaggregated_params.bootstrap_port - if grpc_req.disaggregated_params.bootstrap_port - else None - ) - bootstrap_room = ( - grpc_req.disaggregated_params.bootstrap_room - ) # Can be 0, don't use 'or None' - - # Parse multimodal inputs if present - mm_inputs = None - if grpc_req.HasField("mm_inputs") and grpc_req.mm_inputs.HasField( - "pixel_values" - ): - mm_inputs = self._parse_mm_inputs(grpc_req.mm_inputs) - - # Create request - return TokenizedGenerateReqInput( - rid=grpc_req.request_id, - input_text=input_text, - input_ids=input_ids, - mm_inputs=mm_inputs, - sampling_params=sampling_params, - return_logprob=grpc_req.return_logprob, - logprob_start_len=( - grpc_req.logprob_start_len - if grpc_req.logprob_start_len is not None - else -1 - ), - top_logprobs_num=grpc_req.top_logprobs_num or 0, - stream=grpc_req.stream or False, - lora_id=grpc_req.lora_id if grpc_req.lora_id else None, - token_ids_logprob=( - list(grpc_req.token_ids_logprob) if grpc_req.token_ids_logprob else None - ), - bootstrap_host=bootstrap_host, - bootstrap_port=bootstrap_port, - bootstrap_room=bootstrap_room, - ) - - @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) - - def _parse_mm_inputs(self, mm_proto) -> dict: - """Parse proto MultimodalInputs into the mm_inputs dict expected by scheduler.""" - # Decode pixel_values from typed TensorData field - pixel_values = self._decode_tensor_data(mm_proto.pixel_values) - - # Decode model-specific tensors - model_specific_data = {} - for key, tensor_data in mm_proto.model_specific_tensors.items(): - model_specific_data[key] = self._decode_tensor_data(tensor_data) - - # Convert placeholder ranges to offsets: list of (start, end_inclusive) - offsets = [ - (p.offset, p.offset + p.length - 1) for p in mm_proto.mm_placeholders - ] - if not offsets: - logger.warning( - "No mm_placeholders from Rust gateway — token expansion may have " - "failed to find the placeholder token in input_ids. " - "Check that placeholder_token_id matches the tokenized image token." - ) - offsets = None - - mm_item = MultimodalDataItem( - modality=Modality.IMAGE, - feature=pixel_values, - model_specific_data=model_specific_data, - offsets=offsets, - ) - - result = {"mm_items": [mm_item]} - - if mm_proto.HasField("im_token_id"): - result["im_token_id"] = mm_proto.im_token_id - - return result - - def _convert_embed_request( - self, grpc_req: sglang_scheduler_pb2.EmbedRequest - ) -> TokenizedEmbeddingReqInput: - """Convert gRPC EmbedRequest to internal format.""" - - # Extract tokenized input - if not grpc_req.HasField("tokenized"): - raise ValueError("Tokenized input must be provided") - - input_text = grpc_req.tokenized.original_text - input_ids = list(grpc_req.tokenized.input_ids) - - # Convert sampling params - sampling_params = self._convert_sampling_params(grpc_req.sampling_params) - - # For embedding requests, max_new_tokens should be 0. - # The scheduler logic expects an integer, not None. - sampling_params.max_new_tokens = 0 - - sampling_params.normalize(tokenizer=None) - - return TokenizedEmbeddingReqInput( - rid=grpc_req.request_id, - input_text=input_text, - input_ids=input_ids, - image_inputs={"mm_items": []}, - token_type_ids=list(grpc_req.token_type_ids), - sampling_params=sampling_params, - ) - - def _convert_sampling_params( - self, grpc_params: sglang_scheduler_pb2.SamplingParams - ) -> SGLSamplingParams: - """Convert gRPC SamplingParams to internal format.""" - - # Handle constraint types - regex = None - json_schema = None - ebnf_grammar = None - structural_tag = None - - if grpc_params.HasField("regex"): - regex = grpc_params.regex - elif grpc_params.HasField("json_schema"): - json_schema = grpc_params.json_schema - elif grpc_params.HasField("ebnf_grammar"): - ebnf_grammar = grpc_params.ebnf_grammar - elif grpc_params.HasField("structural_tag"): - structural_tag = grpc_params.structural_tag - - # Handle optional parameters conversion - custom_params = ( - MessageToDict(grpc_params.custom_params) - if grpc_params.HasField("custom_params") - else None - ) - max_new_tokens = ( - grpc_params.max_new_tokens - if grpc_params.HasField("max_new_tokens") - else None - ) - stream_interval = ( - grpc_params.stream_interval - if grpc_params.HasField("stream_interval") - else None - ) - logit_bias = dict(grpc_params.logit_bias) if grpc_params.logit_bias else None - stop = list(grpc_params.stop) if grpc_params.stop else None - stop_token_ids = ( - list(grpc_params.stop_token_ids) if grpc_params.stop_token_ids else None - ) - - return SGLSamplingParams( - temperature=grpc_params.temperature, - top_p=grpc_params.top_p, - top_k=grpc_params.top_k, - min_p=grpc_params.min_p, - frequency_penalty=grpc_params.frequency_penalty, - presence_penalty=grpc_params.presence_penalty, - repetition_penalty=grpc_params.repetition_penalty, - max_new_tokens=max_new_tokens, - min_new_tokens=grpc_params.min_new_tokens, - stop=stop, - stop_token_ids=stop_token_ids, - skip_special_tokens=grpc_params.skip_special_tokens, - spaces_between_special_tokens=grpc_params.spaces_between_special_tokens, - no_stop_trim=grpc_params.no_stop_trim, - regex=regex, - json_schema=json_schema, - ebnf=ebnf_grammar, - structural_tag=structural_tag, - n=grpc_params.n, - ignore_eos=grpc_params.ignore_eos, - stream_interval=stream_interval, - logit_bias=logit_bias, - custom_params=custom_params, - ) - - def _convert_output_logprobs_to_proto( - self, logprobs_data: Dict - ) -> Optional[sglang_scheduler_pb2.OutputLogProbs]: - """Convert output logprobs dict to proto (no None values, plain floats).""" - if not logprobs_data: - return None - - token_logprobs_val = logprobs_data.get("token_logprobs_val", []) - token_logprobs_idx = logprobs_data.get("token_logprobs_idx", []) - top_logprobs_val = logprobs_data.get("top_logprobs_val", []) - top_logprobs_idx = logprobs_data.get("top_logprobs_idx", []) - - # Build TopLogProbs entries - top_logprobs_proto = [] - if top_logprobs_val and top_logprobs_idx: - for val_list, idx_list in zip(top_logprobs_val, top_logprobs_idx): - top_logprobs_proto.append( - sglang_scheduler_pb2.TopLogProbs( - values=val_list, - token_ids=idx_list, - ) - ) - - return sglang_scheduler_pb2.OutputLogProbs( - token_logprobs=token_logprobs_val, # Plain float array - token_ids=token_logprobs_idx, - top_logprobs=top_logprobs_proto, - ) - - def _convert_input_logprobs_to_proto( - self, logprobs_data: Dict - ) -> Optional[sglang_scheduler_pb2.InputLogProbs]: - """Convert input logprobs dict to proto (first token is None, wrapped in InputTokenLogProb).""" - if not logprobs_data: - return None - - token_logprobs_val = logprobs_data.get("token_logprobs_val", []) - token_logprobs_idx = logprobs_data.get("token_logprobs_idx", []) - top_logprobs_val = logprobs_data.get("top_logprobs_val", []) - top_logprobs_idx = logprobs_data.get("top_logprobs_idx", []) - - # Wrap values in InputTokenLogProb (None for first token, value for others) - token_logprobs_wrapped = [ - ( - sglang_scheduler_pb2.InputTokenLogProb() - if x is None - else sglang_scheduler_pb2.InputTokenLogProb(value=x) - ) - for x in token_logprobs_val - ] - - # Build TopLogProbs entries - top_logprobs_proto = [] - if top_logprobs_val and top_logprobs_idx: - for val_list, idx_list in zip(top_logprobs_val, top_logprobs_idx): - top_logprobs_proto.append( - sglang_scheduler_pb2.TopLogProbs( - values=val_list, - token_ids=idx_list, - ) - ) - - return sglang_scheduler_pb2.InputLogProbs( - token_logprobs=token_logprobs_wrapped, - token_ids=token_logprobs_idx, - top_logprobs=top_logprobs_proto, - ) - - def _create_chunk_response( - self, request_id: str, output: Dict - ) -> sglang_scheduler_pb2.GenerateResponse: - """Create a streaming chunk response.""" - meta_info = output.get("meta_info", {}) - - # Convert output logprobs if present - output_logprobs_proto = self._convert_output_logprobs_to_proto( - output.get("output_logprobs") - ) - - # Convert input logprobs if present (only in first chunk) - input_logprobs_proto = self._convert_input_logprobs_to_proto( - output.get("input_logprobs") - ) - - return sglang_scheduler_pb2.GenerateResponse( - request_id=request_id, - chunk=sglang_scheduler_pb2.GenerateStreamChunk( - token_ids=output.get("token_ids", []), - prompt_tokens=meta_info.get("prompt_tokens", 0), - completion_tokens=meta_info.get("completion_tokens", 0), - cached_tokens=meta_info.get("cached_tokens", 0), - output_logprobs=output_logprobs_proto, - input_logprobs=input_logprobs_proto, - index=output.get("index", 0), - ), - ) - - def _create_completion_response( - self, request_id: str, output: Dict - ) -> sglang_scheduler_pb2.GenerateResponse: - """Create a completion response.""" - - # Extract meta info and finish reason details - meta_info = output.get("meta_info", {}) - finish_reason_data = meta_info.get("finish_reason") - - # Determine finish reason, default is stop - finish_reason = "stop" - if finish_reason_data: - if isinstance(finish_reason_data, dict): - finish_reason_type = finish_reason_data.get("type") - else: - # Handle legacy string format - finish_reason_type = finish_reason_data - - if finish_reason_type == "length": - finish_reason = "length" - elif finish_reason_type == "abort": - finish_reason = "abort" - - # Extract matched_stop information - matched_stop_kwargs = {} - if isinstance(finish_reason_data, dict) and "matched" in finish_reason_data: - matched = finish_reason_data["matched"] - if isinstance(matched, int): - matched_stop_kwargs["matched_token_id"] = matched - elif isinstance(matched, str): - matched_stop_kwargs["matched_stop_str"] = matched - - # Convert output logprobs if present - output_logprobs_proto = self._convert_output_logprobs_to_proto( - output.get("output_logprobs") - ) - - # Convert input logprobs if present - input_logprobs_proto = self._convert_input_logprobs_to_proto( - output.get("input_logprobs") - ) - - return sglang_scheduler_pb2.GenerateResponse( - request_id=request_id, - complete=sglang_scheduler_pb2.GenerateComplete( - output_ids=output.get("token_ids", []), - finish_reason=finish_reason, - prompt_tokens=meta_info.get("prompt_tokens", 0), - completion_tokens=meta_info.get( - "completion_tokens", len(output.get("token_ids", [])) - ), - cached_tokens=meta_info.get("cached_tokens", 0), - output_logprobs=output_logprobs_proto, - input_logprobs=input_logprobs_proto, - index=output.get("index", 0), - **matched_stop_kwargs, - ), - ) - - async def shutdown(self): - """Shutdown the service.""" - logger.info("Shutting down gRPC service") - - # Mark health service as NOT_SERVING before shutdown - if self.health_servicer: - self.health_servicer.set_not_serving() - - # Shutdown request manager (handles its own tasks) - await self.request_manager.shutdown() - - -async def serve_grpc( - server_args: ServerArgs, - model_info: Optional[Dict] = None, -): +async def serve_grpc(server_args, model_info=None): """Start the standalone gRPC server with integrated scheduler.""" - - # Start bootstrap server BEFORE launching scheduler processes (only in PREFILL mode) - # This ensures the bootstrap server is ready when prefill schedulers try to register - bootstrap_server = None - if server_args.disaggregation_mode == "prefill": - bootstrap_server = start_disagg_service(server_args) - if bootstrap_server: - logger.info( - f"Bootstrap server started for disaggregation mode on {server_args.host}:{server_args.disaggregation_bootstrap_port}" - ) - - # Launch only the scheduler process(es) (no tokenizer/detokenizer needed for gRPC) - logger.info("Launching scheduler process(es)...") - scheduler_info, port_args, scheduler_procs = launch_scheduler_process_only( - server_args=server_args, - ) - - # Load model config to get HF config info (same as TokenizerManager does) - model_config = ModelConfig.from_server_args(server_args) - - # Update model info from scheduler info and model config - if model_info is None: - # Extract classification labels from HuggingFace config (if available) - # Match logic in serving_classify.py::_get_id2label_mapping - hf_config = model_config.hf_config - id2label = getattr(hf_config, "id2label", None) - num_labels = getattr(hf_config, "num_labels", 0) or 0 - - # If no id2label but num_labels exists, create default mapping - if not id2label and num_labels: - id2label = {i: f"LABEL_{i}" for i in range(num_labels)} - elif id2label and not num_labels: - num_labels = len(id2label) - - # Convert to JSON string for proto transport - # id2label is a dict like {0: "negative", 1: "positive"} - id2label_json = json.dumps(id2label) if id2label else "" - - model_info = { - "model_name": server_args.model_path, - "max_context_length": scheduler_info.get( - "max_total_num_tokens", server_args.context_length or 8192 - ), - "vocab_size": scheduler_info.get("vocab_size", 128256), - "supports_vision": scheduler_info.get("supports_vision", False), - "model_type": getattr(hf_config, "model_type", None), - "architectures": getattr(hf_config, "architectures", None), - "max_req_input_len": scheduler_info.get("max_req_input_len", 8192), - "eos_token_ids": scheduler_info.get("eos_token_ids", []), - "pad_token_id": scheduler_info.get("pad_token_id", 0), - "bos_token_id": scheduler_info.get("bos_token_id", 1), - # Classification model support - "id2label_json": id2label_json, - "num_labels": num_labels or 0, - } - - # Create request manager with the correct port args - # Note: We pass None for bootstrap_server since it's already started above - request_manager = GrpcRequestManager( - server_args=server_args, - port_args=port_args, - bootstrap_server=bootstrap_server, - ) - - # Create gRPC server - server = grpc.aio.server( - futures.ThreadPoolExecutor(max_workers=10), - options=[ - ("grpc.max_send_message_length", 1024 * 1024 * 256), - ("grpc.max_receive_message_length", 1024 * 1024 * 256), - # Allow client HTTP/2 keepalive pings every 10s+. - # Without this, the gRPC C-core default (300s minimum) causes - # GOAWAY when clients send pings more frequently during long - # requests (e.g. prefill) where no DATA frames flow. - ("grpc.http2.min_recv_ping_interval_without_data_ms", 10000), - ("grpc.keepalive_permit_without_calls", True), - ], - ) - - # Create standard health service (for Kubernetes probes) - health_servicer = SGLangHealthServicer( - request_manager=request_manager, - scheduler_info=scheduler_info, - ) - health_pb2_grpc.add_HealthServicer_to_server(health_servicer, server) - - # Add SGLang service - servicer = SGLangSchedulerServicer( - request_manager=request_manager, - server_args=server_args, - model_info=model_info, - scheduler_info=scheduler_info, - health_servicer=health_servicer, - ) - sglang_scheduler_pb2_grpc.add_SglangSchedulerServicer_to_server(servicer, server) - - # Enable reflection - SERVICE_NAMES = ( - sglang_scheduler_pb2.DESCRIPTOR.services_by_name["SglangScheduler"].full_name, - "grpc.health.v1.Health", - reflection.SERVICE_NAME, - ) - reflection.enable_server_reflection(SERVICE_NAMES, server) - - # Start server - listen_addr = f"{server_args.host}:{server_args.port}" - if server_args.ssl_certfile and server_args.ssl_keyfile: - if server_args.ssl_keyfile_password: - raise ValueError( - "gRPC mode does not support encrypted SSL key files " - "(--ssl-keyfile-password). Please provide an unencrypted key " - "file when using --grpc-mode." - ) - - def _read_ssl_file(filepath: str, description: str) -> bytes: - try: - with open(filepath, "rb") as f: - return f.read() - except OSError as e: - raise ValueError( - f"Failed to read {description} '{filepath}': {e}" - ) from e - - private_key = _read_ssl_file(server_args.ssl_keyfile, "SSL key file") - certificate_chain = _read_ssl_file( - server_args.ssl_certfile, "SSL certificate file" - ) - root_certificates = None - if server_args.ssl_ca_certs: - root_certificates = _read_ssl_file( - server_args.ssl_ca_certs, "SSL CA certificates file" - ) - - if server_args.enable_ssl_refresh: - # Use dynamic credentials so gRPC re-reads certs on each - # new connection via the fetcher callback. - _cert_mtime = os.path.getmtime(server_args.ssl_certfile) - _key_mtime = os.path.getmtime(server_args.ssl_keyfile) - _ca_mtime = ( - os.path.getmtime(server_args.ssl_ca_certs) - if server_args.ssl_ca_certs - else None - ) - - def _cert_config_fetcher(): - nonlocal _cert_mtime, _key_mtime, _ca_mtime - try: - new_cert_mt = os.path.getmtime(server_args.ssl_certfile) - new_key_mt = os.path.getmtime(server_args.ssl_keyfile) - new_ca_mt = ( - os.path.getmtime(server_args.ssl_ca_certs) - if server_args.ssl_ca_certs - else None - ) - - if ( - new_cert_mt == _cert_mtime - and new_key_mt == _key_mtime - and new_ca_mt == _ca_mtime - ): - return None # No change - - new_key = _read_ssl_file(server_args.ssl_keyfile, "SSL key file") - new_cert = _read_ssl_file( - server_args.ssl_certfile, "SSL certificate file" - ) - new_root = None - if server_args.ssl_ca_certs: - new_root = _read_ssl_file( - server_args.ssl_ca_certs, - "SSL CA certificates file", - ) - - logger.info("gRPC SSL certificate change detected, reloading.") - config = grpc.ssl_server_certificate_configuration( - [(new_key, new_cert)], - root_certificates=new_root, - ) - - # Update mtimes only after successful reload - _cert_mtime = new_cert_mt - _key_mtime = new_key_mt - _ca_mtime = new_ca_mt - - return config - except Exception: - logger.exception( - "Failed to reload gRPC SSL certificates — " - "continuing with previous certificates." - ) - return None - - try: - initial_config = grpc.ssl_server_certificate_configuration( - [(private_key, certificate_chain)], - root_certificates=root_certificates, - ) - credentials = grpc.dynamic_ssl_server_credentials( - initial_config, - _cert_config_fetcher, - ) - except Exception as e: - raise ValueError( - f"Failed to create gRPC dynamic SSL credentials. " - f"Verify that --ssl-keyfile and --ssl-certfile contain " - f"valid, matching PEM data. Underlying error: {e}" - ) from e - logger.info("gRPC SSL certificate auto-refresh enabled.") - else: - try: - credentials = grpc.ssl_server_credentials( - [(private_key, certificate_chain)], # pairs: (key, cert) - root_certificates=root_certificates, - ) - except Exception as e: - raise ValueError( - f"Failed to create gRPC SSL credentials. Verify that " - f"--ssl-keyfile and --ssl-certfile contain valid, matching " - f"PEM data. Underlying error: {e}" - ) from e - bound_port = server.add_secure_port(listen_addr, credentials) - if bound_port == 0: - raise RuntimeError( - f"Failed to bind gRPC TLS server to {listen_addr}. " - f"Check that the port is available and SSL credentials are valid." - ) - logger.info(f"gRPC server (TLS) listening on {listen_addr}") - else: - server.add_insecure_port(listen_addr) - logger.info(f"gRPC server listening on {listen_addr}") - - await server.start() - - # Start warmup in a separate thread - warmup_thread = threading.Thread( - target=_wait_and_warmup_grpc, - args=(server_args, health_servicer), - ) - warmup_thread.start() - - # Handle shutdown signals - loop = asyncio.get_running_loop() - stop_event = asyncio.Event() - - def signal_handler(): - logger.info("Received shutdown signal") - stop_event.set() - - for sig in (signal.SIGTERM, signal.SIGINT): - loop.add_signal_handler(sig, signal_handler) - try: - await stop_event.wait() - finally: - logger.info("Shutting down gRPC server") - - # Shutdown request manager first - this closes ZMQ sockets and stops background tasks - await servicer.shutdown() - - # Stop the gRPC server - await server.stop(5.0) - - # Wait for warmup thread to finish - if warmup_thread.is_alive(): - logger.info("Waiting for warmup thread to finish...") - warmup_thread.join(timeout=5.0) - - # Terminate scheduler processes before exiting to avoid atexit hang - # The scheduler processes have SIGINT ignored, so they won't get KeyboardInterrupt - for i, proc in enumerate(scheduler_procs): - if proc.is_alive(): - logger.info(f"Terminating scheduler process {i}...") - proc.terminate() - proc.join(timeout=2.0) - if proc.is_alive(): - logger.warning( - f"Scheduler process {i} did not terminate, killing..." - ) - proc.kill() - proc.join(timeout=1.0) - - logger.info("All scheduler processes terminated") - - -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), - ], - ) - stub = sglang_scheduler_pb2_grpc.SglangSchedulerStub(channel) - - # Wait until the server is launched (poll GetModelInfo) - success = False - last_error = None - for _ in range(120): - time.sleep(1) - try: - request = sglang_scheduler_pb2.GetModelInfoRequest() - response = stub.GetModelInfo(request, timeout=5) - success = True - break - except Exception as e: - last_error = str(e) - pass - - if not success: - error_msg = f"gRPC server warmup failed: Could not connect to server after 120 seconds. Last error: {last_error}" - logger.error(error_msg) - channel.close() - kill_process_tree(os.getpid()) - return False - - # Get model info to determine if it's generation or embedding - is_generation = response.is_generation - - # Send a warmup request - logger.info("Sending warmup request to gRPC server...") - max_new_tokens = 8 if is_generation else 1 - - if is_generation: - warmup_request_kwargs = { - "request_id": f"WARMUP_{time.time()}", - "tokenized": sglang_scheduler_pb2.TokenizedInput( - input_ids=[ - 123, - 456, - 789, - 234, - 567, - 890, - 345, - ], # Random-looking but safe token IDs - original_text="warmup request", - ), - "sampling_params": sglang_scheduler_pb2.SamplingParams( - temperature=0.0, - max_new_tokens=max_new_tokens, - ), - "stream": False, - } - - # Set disaggregation params if needed - if server_args.disaggregation_mode != DisaggregationMode.NULL.value: - warmup_request_kwargs["disaggregated_params"] = ( - sglang_scheduler_pb2.DisaggregatedParams( - bootstrap_host=FAKE_BOOTSTRAP_HOST, - bootstrap_room=0, - ) - ) - - warmup_request = sglang_scheduler_pb2.GenerateRequest( - **warmup_request_kwargs - ) - - # Send the warmup request - try: - responses = list(stub.Generate(warmup_request, timeout=600)) - # Check if we got a valid response - if responses and not responses[-1].HasField("error"): - logger.info("gRPC warmup request completed successfully") - success = True - else: - error_msg = ( - responses[-1].error.message if responses else "No response" - ) - logger.warning(f"gRPC warmup request returned error: {error_msg}") - success = False - except Exception as e: - error_msg = f"gRPC warmup request failed: {e}" - logger.error(error_msg) - channel.close() - kill_process_tree(os.getpid()) - return False - else: - # For embedding models - warmup_request = sglang_scheduler_pb2.EmbedRequest( - request_id=f"WARMUP_{time.time()}", - tokenized=sglang_scheduler_pb2.TokenizedInput( - input_ids=[10, 11, 12], - original_text="test embedding", - ), - ) - - try: - response = stub.Embed(warmup_request, timeout=600) - if not response.HasField("error"): - logger.info("gRPC warmup request completed successfully") - success = True - else: - logger.warning( - f"gRPC warmup request returned error: {response.error.message}" - ) - success = False - except Exception as e: - error_msg = f"gRPC warmup request failed: {e}" - logger.error(error_msg) - channel.close() - kill_process_tree(os.getpid()) - return False - - channel.close() - return success - - except Exception as e: - error_msg = ( - f"gRPC warmup failed with exception: {e}\n{get_exception_traceback()}" - ) - logger.error(error_msg) - try: - channel.close() - except Exception: - pass - kill_process_tree(os.getpid()) - return False - - -def _wait_and_warmup_grpc( - server_args: ServerArgs, - health_servicer: Optional[SGLangHealthServicer] = None, -): - """Wait for gRPC server to be ready and execute warmup.""" - if not server_args.skip_server_warmup: - if not _execute_grpc_server_warmup(server_args): - return - else: - logger.info("Skipping gRPC server warmup (skip_server_warmup=True)") - - # Mark health service as SERVING after warmup completes - if health_servicer: - health_servicer.set_serving() - - logger.info("The server is fired up and ready to roll!") + from smg_grpc_servicer.sglang.server import serve_grpc as _serve_grpc + except ImportError: + raise ImportError( + "gRPC mode requires the smg-grpc-servicer package. " + "Install it with: pip install smg-grpc-servicer[sglang]" + ) from None + await _serve_grpc(server_args, model_info) diff --git a/python/sglang/srt/grpc/grpc_request_manager.py b/python/sglang/srt/grpc/grpc_request_manager.py deleted file mode 100644 index 959f43e49..000000000 --- a/python/sglang/srt/grpc/grpc_request_manager.py +++ /dev/null @@ -1,1022 +0,0 @@ -""" -gRPC Request Manager - Orchestrates request lifecycle without tokenization. -Mimics TokenizerManager's state management and ZMQ communication patterns. -""" - -import asyncio -import copy -import dataclasses -import logging -import os -import signal -import sys -import threading -import uuid -from typing import Any, AsyncGenerator, Dict, List, Optional, Union - -import grpc -import zmq -import zmq.asyncio - -from sglang.srt.disaggregation.utils import DisaggregationMode -from sglang.srt.managers.io_struct import ( - AbortReq, - BatchEmbeddingOutput, - BatchTokenIDOutput, - GetLoadsReqInput, - GetLoadsReqOutput, - HealthCheckOutput, - TokenizedEmbeddingReqInput, - TokenizedGenerateReqInput, -) -from sglang.srt.observability.req_time_stats import ( - APIServerReqTimeStats, - calibrate_time_diff, - real_time, -) -from sglang.srt.server_args import PortArgs, ServerArgs -from sglang.srt.utils import get_or_create_event_loop, get_zmq_socket, kill_process_tree -from sglang.utils import get_exception_traceback - -logger = logging.getLogger(__name__) - - -class _GrpcCommunicator: - """ - Communicator for request/response patterns with scheduler. - - Thread-safe and handles the async request/response cycle with proper - timeout handling to prevent hangs if the scheduler becomes unresponsive. - """ - - DEFAULT_TIMEOUT = 30.0 # seconds - - def __init__(self, sender: zmq.Socket, fan_out: int = 1): - self._sender = sender - self._fan_out = fan_out - self._result_event: Optional[asyncio.Event] = None - self._result_values: Optional[List[Any]] = None - self._lock = asyncio.Lock() - - async def __call__(self, obj, timeout: float = DEFAULT_TIMEOUT) -> List[Any]: - """ - Send request and wait for response(s). - - Args: - obj: Request object to send to scheduler - timeout: Maximum time to wait for response (seconds) - - Returns: - List of response objects from scheduler(s) - - Raises: - asyncio.TimeoutError: If no response within timeout - """ - async with self._lock: - # Initialize state BEFORE sending to avoid race condition - self._result_event = asyncio.Event() - self._result_values = [] - - # Send request to scheduler - if obj: - self._sender.send_pyobj(obj) - - try: - # Wait for response(s) with timeout - await asyncio.wait_for(self._result_event.wait(), timeout=timeout) - return self._result_values - finally: - # Always clean up state - self._result_event = None - self._result_values = None - - def handle_recv(self, recv_obj: Any): - """ - Handle received response from scheduler. - - Called by handle_loop when a matching response type is received. - Safe to call even if no request is pending (will be ignored). - """ - if self._result_values is not None and self._result_event is not None: - self._result_values.append(recv_obj) - if len(self._result_values) >= self._fan_out: - self._result_event.set() - - -class GrpcSignalHandler: - """Minimal signal handler for gRPC server - delegates real crash handling to scheduler.""" - - def __init__(self, grpc_manager): - self.grpc_manager = grpc_manager - - def sigterm_handler(self, signum=None, frame=None): - """Handle SIGTERM by gracefully shutting down gRPC server.""" - logger.warning( - f"SIGTERM received. {signum=} {frame=}. Shutting down gRPC server..." - ) - self.grpc_manager.gracefully_exit = True - - def running_phase_sigquit_handler(self, signum=None, frame=None): - """Handle SIGQUIT from failed scheduler process.""" - logger.error( - "Received SIGQUIT from scheduler process. Scheduler failed, shutting down gRPC server." - ) - logger.info( - "Note: Crash dumps are handled by the scheduler process, not the gRPC server." - ) - # Just exit cleanly - the scheduler handles crash dumps - kill_process_tree(os.getpid(), include_parent=True) - - -@dataclasses.dataclass -class GrpcReqState: - """State tracking for a gRPC request.""" - - # Request identification - request_id: str - grpc_context: Optional[grpc.aio.ServicerContext] - - # Communication - out_queue: asyncio.Queue - finished: bool - event: asyncio.Event - obj: Union[TokenizedGenerateReqInput, TokenizedEmbeddingReqInput] - - # Metrics (same as TokenizerManager's ReqState) - time_stats: APIServerReqTimeStats - last_completion_tokens: int = 1 - - # Streaming state - stream_finished: bool = False - input_logprobs_sent: bool = False # Track if input logprobs were sent in streaming - - # Token accumulation (for non-streaming) - output_ids: List[int] = dataclasses.field(default_factory=list) - input_token_logprobs_val: List[float] = dataclasses.field(default_factory=list) - input_token_logprobs_idx: List[int] = dataclasses.field(default_factory=list) - output_token_logprobs_val: List[float] = dataclasses.field(default_factory=list) - output_token_logprobs_idx: List[int] = dataclasses.field(default_factory=list) - input_top_logprobs_val: List[List[float]] = dataclasses.field(default_factory=list) - input_top_logprobs_idx: List[List[int]] = dataclasses.field(default_factory=list) - output_top_logprobs_val: List[List[float]] = dataclasses.field(default_factory=list) - output_top_logprobs_idx: List[List[int]] = dataclasses.field(default_factory=list) - - # Session state - session_id: Optional[str] = None - is_session_request: bool = False - - -class GrpcRequestManager: - """ - Manages gRPC request lifecycle, mimicking TokenizerManager's orchestration - behaviors without tokenization. - """ - - def __init__( - self, - server_args: ServerArgs, - port_args: PortArgs, - bootstrap_server=None, - ): - """Initialize the gRPC request manager.""" - self.server_args = server_args - self.port_args = port_args - - # ZMQ Communication Setup (same pattern as TokenizerManager) - self.context = zmq.asyncio.Context(2) - - # Socket for receiving outputs from scheduler - self.recv_from_scheduler = get_zmq_socket( - self.context, zmq.PULL, port_args.detokenizer_ipc_name, bind=True - ) - - # Socket for sending requests to scheduler - self.send_to_scheduler = get_zmq_socket( - self.context, zmq.PUSH, port_args.scheduler_input_ipc_name, bind=True - ) - - # State Management (from TokenizerManager) - self.rid_to_state: Dict[str, GrpcReqState] = {} - self.asyncio_tasks: set = set() - self.gracefully_exit = False - self.no_create_loop = False - self.event_loop = None - - # Pause/Resume Control - self.is_pause = False - self.is_pause_cond = asyncio.Condition() - - # Metrics - self.last_receive_tstamp = real_time() - - # Crash dump for debugging - self.crash_dump_request_list = [] - self.crash_dump_performed = False - - # disaggregation mode - self.disaggregation_mode = DisaggregationMode(server_args.disaggregation_mode) - - # Bootstrap server (passed from serve_grpc, not started here) - self.bootstrap_server = bootstrap_server - - # Communicators for request/response patterns with scheduler - # Note: These must be initialized after send_to_scheduler socket is created - self.get_loads_communicator = _GrpcCommunicator( - self.send_to_scheduler, fan_out=server_args.dp_size - ) - - logger.info( - f"GrpcRequestManager initialized with ZMQ IPC: " - f"recv={port_args.detokenizer_ipc_name}, " - f"send={port_args.scheduler_input_ipc_name}" - ) - if self.bootstrap_server: - logger.info( - f"Bootstrap server initialized for disaggregation mode: " - f"{server_args.disaggregation_mode}" - ) - - async def generate_request( - self, - obj: TokenizedGenerateReqInput, - request_id: Optional[str] = None, - grpc_context: Optional[grpc.aio.ServicerContext] = None, - ) -> AsyncGenerator[Union[Dict, List[Dict]], None]: - """ - Submit a generation request to the scheduler with n>1 parallel sampling support. - - This method implements the same two-phase approach as tokenizer_manager.py: - 1. Phase 1: Send prefix caching request (max_new_tokens=0) - 2. Phase 2: Send n generation requests that reuse the cached prefix - - Yields individual responses for streaming, or aggregated responses for non-streaming. - """ - n = getattr(obj.sampling_params, "n", 1) - - if n <= 1: - async for response in self._handle_single_request( - obj, request_id, grpc_context - ): - yield response - return - - # N>1 handling - two-phase approach - logger.debug(f"Multiple sampling request (n={n}), using two-phase approach") - - # Generate base request ID if not provided - if request_id is None: - base_request_id = f"grpc-{uuid.uuid4().hex}" - else: - base_request_id = request_id - - # Phase 1: Cache the common prefix - logger.debug(f"Phase 1: Caching prefix for request {base_request_id}") - prefix_obj = copy.copy(obj) - prefix_obj.sampling_params = copy.copy(obj.sampling_params) - prefix_obj.sampling_params.max_new_tokens = 0 # Prefill-only - prefix_obj.sampling_params.n = 1 # Don't replicate prefix request - - # Send prefix caching request and consume response - async for _ in self._handle_single_request( - prefix_obj, f"{base_request_id}-prefix", grpc_context - ): - # Consume prefix response (usually just one chunk with finish_reason) - pass - - logger.debug(f"Phase 1 completed: Prefix cached for {base_request_id}") - - # Phase 2: Generate n parallel requests - logger.debug(f"Phase 2: Generating {n} parallel requests") - generators = [] - request_ids = [] - - for i in range(n): - # Create individual generation request - gen_obj = copy.copy(obj) - gen_obj.sampling_params = copy.copy(obj.sampling_params) - gen_obj.sampling_params.n = 1 # Each request generates 1 response - - gen_request_id = f"{base_request_id}-{i}" - request_ids.append(gen_request_id) - - # Start generation request - generators.append( - self._handle_single_request(gen_obj, gen_request_id, grpc_context) - ) - - # Handle response aggregation - is_stream = getattr(obj, "stream", False) - - if not is_stream: - # Non-streaming: collect all responses and return as batch - logger.debug(f"Non-streaming mode: collecting {n} responses") - responses = [] - for generator in generators: - async for response in generator: - responses.append(response) - yield responses # Return all responses as a batch - else: - # Streaming mode: multiplex responses with index for ordering - logger.debug(f"Streaming mode: multiplexing {n} streams") - rid_to_index = {rid: i for i, rid in enumerate(request_ids)} - - # Create async tasks for all generators - task_map = {} - for generator in generators: - task = asyncio.create_task(generator.__anext__()) - task_map[task] = generator - - # Process responses as they arrive - while task_map: - done, _ = await asyncio.wait( - task_map.keys(), return_when=asyncio.FIRST_COMPLETED - ) - - for task in done: - generator = task_map.pop(task) - try: - response = await task - - # Add index for client-side ordering - if isinstance(response, dict): - response_rid = response.get("request_id", "") - if response_rid in rid_to_index: - response["index"] = rid_to_index[response_rid] - - yield response - - # Create next task for this generator - next_task = asyncio.create_task(generator.__anext__()) - task_map[next_task] = generator - - except StopAsyncIteration: - # This generator is finished - pass - - async def _handle_single_request( - self, - obj: TokenizedGenerateReqInput, - request_id: Optional[str] = None, - grpc_context: Optional[grpc.aio.ServicerContext] = None, - ): - """Handle a single request - core implementation without n>1 logic.""" - # Generate request ID if not provided - if request_id is None: - request_id = f"grpc-{uuid.uuid4().hex}" - - obj.rid = request_id - - self._req_stats_init(obj, grpc_context) - state = self.rid_to_state[request_id] - self.record_request_for_crash_dump(obj) - - try: - # Send to scheduler - let exceptions bubble up to grpc_server.py - state.time_stats.set_api_server_dispatch_time() - await self._send_to_scheduler(obj) - state.time_stats.set_api_server_dispatch_finish_time() - - is_stream = getattr(obj, "stream", False) - - while True: - try: - response = await state.out_queue.get() - - if is_stream: - yield response - - # Non-streaming: yield final response with accumulated tokens from state - if isinstance(response, dict) and response.get("finished", False): - if not is_stream: - final_response = response.copy() - final_response["token_ids"] = state.output_ids - yield final_response - break - - except asyncio.CancelledError: - # Task was cancelled by gRPC framework when client disconnected - logger.info(f"Request {request_id} cancelled by client") - await self.abort_request(request_id) - raise # Re-raise to let gRPC server handle cleanup - - finally: - # Always clean up request state when exiting - self._cleanup_request_state(request_id) - - def _cleanup_request_state(self, request_id: str): - """Clean up local request state (does not notify scheduler).""" - if request_id in self.rid_to_state: - del self.rid_to_state[request_id] - - async def embedding_request( - self, - obj: TokenizedEmbeddingReqInput, - request_id: Optional[str] = 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 - - async def abort_request(self, request_id: str) -> bool: - """Abort a running request. - - Sends abort request to scheduler and marks local state as finished - to stop processing any further outputs from the scheduler. - """ - # Skip aborting health check requests (they clean themselves up) - if request_id.startswith("HEALTH_CHECK"): - return False - - # Mark state as finished immediately to stop processing scheduler outputs - state = self.rid_to_state.get(request_id) - if state: - state.finished = True - state.stream_finished = True - logger.debug(f"Marked request {request_id} as aborted locally") - - # Send abort to scheduler - the scheduler will send AbortReq back - # which will be handled by _handle_abort_req - abort_req = AbortReq(rid=request_id) - try: - await self._send_to_scheduler(abort_req) - logger.debug(f"Sent abort to scheduler for request {request_id}") - except Exception as e: - logger.error(f"Failed to send abort request to scheduler: {e}") - return False - - return True - - async def handle_loop(self): - """ - Main event loop - processes outputs from scheduler. - Mimics TokenizerManager's handle_loop. - """ - while not self.gracefully_exit: - try: - # Receive from scheduler - recv_obj = await self.recv_from_scheduler.recv_pyobj() - self.last_receive_tstamp = real_time() - - # Check for pause (optimized: check flag before acquiring lock) - if self.is_pause: - async with self.is_pause_cond: - while self.is_pause: - await self.is_pause_cond.wait() - - # Handle different output types - if isinstance(recv_obj, BatchTokenIDOutput): - await self._handle_batch_output(recv_obj) - elif isinstance(recv_obj, BatchEmbeddingOutput): - await self._handle_embedding_output(recv_obj) - elif isinstance(recv_obj, HealthCheckOutput): - await self._handle_health_check_output(recv_obj) - elif isinstance(recv_obj, AbortReq): - await self._handle_abort_req(recv_obj) - elif isinstance(recv_obj, GetLoadsReqOutput): - # Route to communicator for request/response pattern - self.get_loads_communicator.handle_recv(recv_obj) - else: - logger.warning(f"Unknown output type: {type(recv_obj)}") - - except zmq.error.Again: - # Timeout, check if we should exit - if self.gracefully_exit: - break - continue - except zmq.error.ZMQError as e: - # Socket closed or other ZMQ error - exit cleanly if shutting down - if self.gracefully_exit: - logger.debug(f"ZMQ recv interrupted during shutdown: {e}") - break - logger.error( - f"ZMQ error in handle loop: {e}\n{get_exception_traceback()}" - ) - break - except Exception as e: - logger.error(f"Handle loop error: {e}\n{get_exception_traceback()}") - if self.gracefully_exit: - break - - def _convert_logprob_style( - self, - state: GrpcReqState, - batch_out: BatchTokenIDOutput, - batch_index: int, - ): - """ - Convert and accumulate logprobs from batch output to state. - Follows the same logic as tokenizer_manager.convert_logprob_style. - """ - # Early exit if no input logprobs at all - if batch_out.input_token_logprobs_val is None: - return - - # Accumulate input token logprobs (only if list is non-empty) - if len(batch_out.input_token_logprobs_val) > 0: - state.input_token_logprobs_val.extend( - batch_out.input_token_logprobs_val[batch_index] - ) - state.input_token_logprobs_idx.extend( - batch_out.input_token_logprobs_idx[batch_index] - ) - - # Always accumulate output token logprobs - state.output_token_logprobs_val.extend( - batch_out.output_token_logprobs_val[batch_index] - ) - state.output_token_logprobs_idx.extend( - batch_out.output_token_logprobs_idx[batch_index] - ) - - # Handle top logprobs if requested - if state.obj.top_logprobs_num > 0: - # Accumulate input top logprobs (only if list is non-empty) - if len(batch_out.input_top_logprobs_val) > 0: - state.input_top_logprobs_val.extend( - batch_out.input_top_logprobs_val[batch_index] - ) - state.input_top_logprobs_idx.extend( - batch_out.input_top_logprobs_idx[batch_index] - ) - - # Always accumulate output top logprobs - state.output_top_logprobs_val.extend( - batch_out.output_top_logprobs_val[batch_index] - ) - state.output_top_logprobs_idx.extend( - batch_out.output_top_logprobs_idx[batch_index] - ) - - async def _handle_batch_output(self, batch_out: BatchTokenIDOutput): - """Handle batch generation output from scheduler.""" - # Collect all queue.put() tasks for parallel execution - put_tasks = [] - cleanup_tasks = [] - - # Process each request in the batch - for i, rid in enumerate(batch_out.rids): - if rid not in self.rid_to_state: - continue - - state = self.rid_to_state[rid] - - # Skip if already aborted/finished locally (client cancelled) - if state.finished: - logger.debug(f"Skipping output for aborted request {rid}") - continue - - # Update metrics - if state.time_stats.first_token_time == 0.0: - state.time_stats.set_first_token_time() - else: - state.time_stats.set_last_time() - - # Extract output for this request - output_data = { - "request_id": rid, - "token_ids": batch_out.output_ids[i] if batch_out.output_ids else [], - "finished": batch_out.finished_reasons[i] is not None, - "meta_info": { - "prompt_tokens": ( - batch_out.prompt_tokens[i] if batch_out.prompt_tokens else 0 - ), - "completion_tokens": ( - batch_out.completion_tokens[i] - if batch_out.completion_tokens - else 0 - ), - "cached_tokens": ( - batch_out.cached_tokens[i] if batch_out.cached_tokens else 0 - ), - "finish_reason": ( - batch_out.finished_reasons[i] - if batch_out.finished_reasons[i] - else None - ), - }, - } - - # Accumulate logprobs (following tokenizer_manager pattern) - # Use getattr for safe access - not all request types have return_logprob - # (e.g., TokenizedEmbeddingReqInput) - if getattr(state.obj, "return_logprob", False): - self._convert_logprob_style(state, batch_out, i) - - # Send input logprobs based if available - if ( - getattr(state.obj, "return_logprob", False) - and state.obj.logprob_start_len >= 0 - and state.input_token_logprobs_val - ): - if state.obj.stream and not state.input_logprobs_sent: - # Streaming: send input logprobs once in first chunk that has them - output_data["input_logprobs"] = { - "token_logprobs_val": state.input_token_logprobs_val, - "token_logprobs_idx": state.input_token_logprobs_idx, - "top_logprobs_val": state.input_top_logprobs_val, - "top_logprobs_idx": state.input_top_logprobs_idx, - } - state.input_logprobs_sent = True - elif not state.obj.stream and output_data["finished"]: - # Non-streaming: send input logprobs in final chunk - output_data["input_logprobs"] = { - "token_logprobs_val": state.input_token_logprobs_val, - "token_logprobs_idx": state.input_token_logprobs_idx, - "top_logprobs_val": state.input_top_logprobs_val, - "top_logprobs_idx": state.input_top_logprobs_idx, - } - - # Send output logprobs if available - if ( - getattr(state.obj, "return_logprob", False) - and batch_out.output_token_logprobs_val - and i < len(batch_out.output_token_logprobs_val) - ): - if state.obj.stream: - # For streaming: send incremental logprobs (only new tokens in this chunk) - # NOTE: this is different than TokenizerManager, which always accumulates - def get_part(attr_name): - source_list = getattr(batch_out, attr_name, None) - return ( - source_list[i] - if source_list and i < len(source_list) - else [] - ) - - output_data["output_logprobs"] = { - "token_logprobs_val": batch_out.output_token_logprobs_val[i], - "token_logprobs_idx": get_part("output_token_logprobs_idx"), - "top_logprobs_val": get_part("output_top_logprobs_val"), - "top_logprobs_idx": get_part("output_top_logprobs_idx"), - } - elif output_data["finished"]: - # Non-streaming: send cumulative output logprobs in final chunk - output_data["output_logprobs"] = { - "token_logprobs_val": state.output_token_logprobs_val, - "token_logprobs_idx": state.output_token_logprobs_idx, - "top_logprobs_val": state.output_top_logprobs_val, - "top_logprobs_idx": state.output_top_logprobs_idx, - } - - # Update state for accumulation - if output_data["token_ids"]: - state.output_ids.extend(output_data["token_ids"]) - - # Add queue.put() to parallel task list - put_tasks.append(state.out_queue.put(output_data)) - - # Handle completion - if output_data["finished"]: - state.finished = True - state.time_stats.set_finished_time() - state.stream_finished = True - state.event.set() - - # 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))) - - # Execute all queue.put() operations in parallel - if put_tasks: - await asyncio.gather(*put_tasks, return_exceptions=True) - - async def _handle_embedding_output(self, batch_out: BatchEmbeddingOutput): - """Handle batch embedding output from scheduler.""" - for i, rid in enumerate(batch_out.rids): - if rid not in self.rid_to_state: - continue - - state = self.rid_to_state[rid] - - # Create result - result = { - "request_id": rid, - "embedding": batch_out.embeddings[i], - "prompt_tokens": ( - batch_out.prompt_tokens[i] if batch_out.prompt_tokens else 0 - ), - "finish_reason": ( - batch_out.finished_reasons[i] - if batch_out.finished_reasons - else None - ), - } - - # Send result - await state.out_queue.put(result) - - # Mark as finished - state.finished = True - state.time_stats.set_finished_time() - state.event.set() - - async def _handle_health_check_output(self, health_out: HealthCheckOutput): - """Handle health check output from scheduler.""" - rid = health_out.rid - - if rid not in self.rid_to_state: - logger.warning(f"Health check output for unknown request: {rid}") - return - - state = self.rid_to_state[rid] - - # Create health check result - result = { - "request_id": rid, - "healthy": True, # If we got a response, scheduler is healthy - "output_text": ( - health_out.output_str if hasattr(health_out, "output_str") else "" - ), - "finish_reason": ( - health_out.finish_reason - if hasattr(health_out, "finish_reason") - else "stop" - ), - } - - # Send result - await state.out_queue.put(result) - - # Mark as finished - state.finished = True - state.time_stats.set_finished_time() - state.event.set() - - async def _handle_abort_req(self, recv_obj: AbortReq): - """Handle abort request from scheduler. - - The scheduler sends AbortReq back to notify us that a request was aborted, - either due to explicit abort_request() call or scheduler-initiated abort - (priority preemption, queue full, KV cache pressure, etc). - """ - # Skip health check requests - if recv_obj.rid.startswith("HEALTH_CHECK"): - return - - # Check if request still exists - if recv_obj.rid not in self.rid_to_state: - logger.debug( - f"Abort request for {recv_obj.rid} not in local state (may have already finished or not started yet)" - ) - return - - state = self.rid_to_state[recv_obj.rid] - - # Mark as finished - state.finished = True - state.stream_finished = True - - # Create abort response - if recv_obj.finished_reason: - # Scheduler provided a specific finish reason (e.g., priority preemption, queue full) - abort_response = { - "request_id": recv_obj.rid, - "error": recv_obj.finished_reason.get("message", "Request aborted"), - "finished": True, - "meta_info": { - "id": recv_obj.rid, - "finish_reason": recv_obj.finished_reason, - }, - } - else: - # Generic abort (e.g., explicit abort_request call) - abort_response = { - "request_id": recv_obj.rid, - "error": "Request aborted", - "finished": True, - "meta_info": { - "id": recv_obj.rid, - "finish_reason": { - "type": "abort", - "message": "Abort before prefill", - }, - "prompt_tokens": 0, - "completion_tokens": 0, - }, - } - - # Send abort notification to output queue - await state.out_queue.put(abort_response) - - # Wake up any waiting coroutines - state.event.set() - - logger.debug(f"Handled abort request for {recv_obj.rid}") - - async def _send_to_scheduler(self, obj): - """Send an object to the scheduler via ZMQ.""" - try: - self.send_to_scheduler.send_pyobj(obj) - except Exception as e: - logger.error(f"Failed to send to scheduler: {e}") - raise - - def record_request_for_crash_dump(self, obj): - """Record request for potential crash dump.""" - if len(self.crash_dump_request_list) < 100: - self.crash_dump_request_list.append( - { - "time": real_time(), - "request_id": getattr(obj, "rid", "unknown"), - "type": type(obj).__name__, - } - ) - - 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) - - # Shutdown bootstrap server if running - if self.bootstrap_server: - logger.info("Shutting down bootstrap server") - try: - if hasattr(self.bootstrap_server, "shutdown"): - if asyncio.iscoroutinefunction(self.bootstrap_server.shutdown): - await self.bootstrap_server.shutdown() - else: - self.bootstrap_server.shutdown() - except Exception as e: - logger.warning(f"Error shutting down bootstrap server: {e}") - - # Close ZMQ sockets - self.recv_from_scheduler.close() - self.send_to_scheduler.close() - - # Terminate the ZMQ context - this is critical for asyncio loop to exit cleanly - self.context.term() - - logger.info("GrpcRequestManager shutdown complete") - - def get_server_info(self) -> Dict[str, Any]: - """Get server information for health checks.""" - return { - "active_requests": len(self.rid_to_state), - "paused": self.is_pause, - "last_receive_time": self.last_receive_tstamp, - } - - async def get_loads( - self, include: List[str], dp_rank: Optional[int] = None - ) -> List[GetLoadsReqOutput]: - """ - Get comprehensive load metrics from the scheduler. - - This method uses the communicator pattern to send GetLoadsReqInput to the - scheduler and wait for GetLoadsReqOutput responses. - - Args: - include: List of metric sections to include (core, memory, spec, lora, disagg, queues, all) - dp_rank: Optional DP rank filter (None for all ranks) - - Returns: - List of GetLoadsReqOutput objects, one per scheduler/DP rank - """ - req = GetLoadsReqInput(include=include, dp_rank=dp_rank) - results = await self.get_loads_communicator(req) - - # Filter by dp_rank if specified - if dp_rank is not None: - results = [r for r in results if r.dp_rank == dp_rank] - - return results - - def auto_create_handle_loop(self): - """Automatically create and start the handle_loop task, matching TokenizerManager pattern.""" - if self.no_create_loop: - return - - self.no_create_loop = True - loop = get_or_create_event_loop() - self.asyncio_tasks.add( - loop.create_task(print_exception_wrapper(self.handle_loop)) - ) - - self.event_loop = loop - - # We only add signal handler when the tokenizer manager is in the main thread - # due to the CPython limitation. - if threading.current_thread() is threading.main_thread(): - signal_handler = GrpcSignalHandler(self) - loop.add_signal_handler(signal.SIGTERM, signal_handler.sigterm_handler) - # Update the signal handler for the process. It overrides the sigquit handler in the launch phase. - loop.add_signal_handler( - signal.SIGQUIT, signal_handler.running_phase_sigquit_handler - ) - - self.asyncio_tasks.add( - loop.create_task(print_exception_wrapper(self.sigterm_watchdog)) - ) - - async def sigterm_watchdog(self): - """Watchdog to handle SIGTERM gracefully, matching TokenizerManager pattern.""" - while not self.gracefully_exit: - await asyncio.sleep(1.0) - - def _req_stats_init( - self, - obj: Union[TokenizedGenerateReqInput, TokenizedEmbeddingReqInput], - grpc_context: Optional[grpc.ServicerContext] = None, - ): - calibrate_time_diff() - # Create and register request state - # TODO: support log_request - # TODO: support request tracing - time_stats = APIServerReqTimeStats(disagg_mode=self.disaggregation_mode) - state = GrpcReqState( - request_id=obj.rid, - grpc_context=grpc_context, - out_queue=asyncio.Queue(), - finished=False, - event=asyncio.Event(), - obj=obj, - time_stats=time_stats, - ) - - # Track session if needed - if hasattr(obj, "session_params") and obj.session_params: - state.session_id = obj.session_params.session_id - state.is_session_request = True - - self.rid_to_state[obj.rid] = state - time_stats.set_created_time() - - -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) - sys.exit(1) diff --git a/python/sglang/srt/grpc/health_servicer.py b/python/sglang/srt/grpc/health_servicer.py deleted file mode 100644 index db3db2cc0..000000000 --- a/python/sglang/srt/grpc/health_servicer.py +++ /dev/null @@ -1,189 +0,0 @@ -""" -Standard gRPC health check service implementation for Kubernetes probes. - -This module implements the grpc.health.v1.Health service protocol, enabling -native Kubernetes gRPC health probes for liveness and readiness checks. -""" - -import logging -import time -from typing import AsyncIterator - -import grpc -from grpc_health.v1 import health_pb2, health_pb2_grpc - -logger = logging.getLogger(__name__) - - -class SGLangHealthServicer(health_pb2_grpc.HealthServicer): - """ - Standard gRPC health check service implementation for Kubernetes probes. - Implements grpc.health.v1.Health protocol. - - Supports two service levels: - 1. Overall server health (service="") - for liveness probes - 2. SGLang service health (service="sglang.grpc.scheduler.SglangScheduler") - for readiness probes - - Health status lifecycle: - - NOT_SERVING: Initial state, model loading, or shutting down - - SERVING: Model loaded and ready to serve requests - """ - - # Service names we support - OVERALL_SERVER = "" # Empty string for overall server health - SGLANG_SERVICE = "sglang.grpc.scheduler.SglangScheduler" - - def __init__(self, request_manager, scheduler_info: dict): - """ - Initialize health servicer. - - Args: - request_manager: GrpcRequestManager instance for checking server state - scheduler_info: Dict containing scheduler metadata - """ - self.request_manager = request_manager - self.scheduler_info = scheduler_info - self._serving_status = {} - - # Initially set to NOT_SERVING until model is loaded - self._serving_status[self.OVERALL_SERVER] = ( - health_pb2.HealthCheckResponse.NOT_SERVING - ) - self._serving_status[self.SGLANG_SERVICE] = ( - health_pb2.HealthCheckResponse.NOT_SERVING - ) - - logger.info("Standard gRPC health service initialized") - - def set_serving(self): - """Mark services as SERVING - call this after model is loaded.""" - self._serving_status[self.OVERALL_SERVER] = ( - health_pb2.HealthCheckResponse.SERVING - ) - self._serving_status[self.SGLANG_SERVICE] = ( - health_pb2.HealthCheckResponse.SERVING - ) - logger.info("Health service status set to SERVING") - - def set_not_serving(self): - """Mark services as NOT_SERVING - call this during shutdown.""" - self._serving_status[self.OVERALL_SERVER] = ( - health_pb2.HealthCheckResponse.NOT_SERVING - ) - self._serving_status[self.SGLANG_SERVICE] = ( - health_pb2.HealthCheckResponse.NOT_SERVING - ) - logger.info("Health service status set to NOT_SERVING") - - async def Check( - self, - request: health_pb2.HealthCheckRequest, - context: grpc.aio.ServicerContext, - ) -> health_pb2.HealthCheckResponse: - """ - Standard health check for Kubernetes probes. - - Args: - request: Contains service name ("" for overall, or specific service) - context: gRPC context - - Returns: - HealthCheckResponse with SERVING/NOT_SERVING/SERVICE_UNKNOWN status - """ - service_name = request.service - logger.debug(f"Health check request for service: '{service_name}'") - - # Check if shutting down - if self.request_manager.gracefully_exit: - logger.debug("Health check: Server is shutting down") - return health_pb2.HealthCheckResponse( - status=health_pb2.HealthCheckResponse.NOT_SERVING - ) - - # Overall server health - just check if process is alive - if service_name == self.OVERALL_SERVER: - status = self._serving_status.get( - self.OVERALL_SERVER, health_pb2.HealthCheckResponse.NOT_SERVING - ) - logger.debug( - f"Overall health check: {health_pb2.HealthCheckResponse.ServingStatus.Name(status)}" - ) - return health_pb2.HealthCheckResponse(status=status) - - # Specific service health - check if ready to serve - elif service_name == self.SGLANG_SERVICE: - # Additional checks for service readiness - - # Check base status first - base_status = self._serving_status.get( - self.SGLANG_SERVICE, health_pb2.HealthCheckResponse.NOT_SERVING - ) - - if base_status != health_pb2.HealthCheckResponse.SERVING: - logger.debug("Service health check: NOT_SERVING (base status)") - return health_pb2.HealthCheckResponse(status=base_status) - - # Check if scheduler is responsive (received data recently) - time_since_last_receive = ( - time.time() - self.request_manager.last_receive_tstamp - ) - - # If no recent activity and we have active requests, might be stuck - # NOTE: 30s timeout is hardcoded. This is more conservative than - # HEALTH_CHECK_TIMEOUT (20s) used for custom HealthCheck RPC. - # Consider making this configurable via environment variable in the future - # if different workloads need different responsiveness thresholds. - if ( - time_since_last_receive > 30 - and len(self.request_manager.rid_to_state) > 0 - ): - logger.warning( - f"Service health check: Scheduler not responsive " - f"({time_since_last_receive:.1f}s since last receive, " - f"{len(self.request_manager.rid_to_state)} pending requests)" - ) - return health_pb2.HealthCheckResponse( - status=health_pb2.HealthCheckResponse.NOT_SERVING - ) - - logger.debug("Service health check: SERVING") - return health_pb2.HealthCheckResponse( - status=health_pb2.HealthCheckResponse.SERVING - ) - - # Unknown service - else: - logger.debug(f"Health check for unknown service: '{service_name}'") - context.set_code(grpc.StatusCode.NOT_FOUND) - context.set_details(f"Unknown service: {service_name}") - return health_pb2.HealthCheckResponse( - status=health_pb2.HealthCheckResponse.SERVICE_UNKNOWN - ) - - async def Watch( - self, - request: health_pb2.HealthCheckRequest, - context: grpc.aio.ServicerContext, - ) -> AsyncIterator[health_pb2.HealthCheckResponse]: - """ - Streaming health check - sends updates when status changes. - - For now, just send 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 when status changes - """ - service_name = request.service - logger.debug(f"Health watch request for service: '{service_name}'") - - # Send current status - response = await self.Check(request, context) - yield response - - # Note: Full Watch implementation would monitor status changes - # and stream updates. For K8s probes, Check is sufficient. diff --git a/python/sglang/srt/grpc/scheduler_launcher.py b/python/sglang/srt/grpc/scheduler_launcher.py deleted file mode 100644 index a35dd0471..000000000 --- a/python/sglang/srt/grpc/scheduler_launcher.py +++ /dev/null @@ -1,198 +0,0 @@ -""" -Scheduler process management for gRPC server. - -This module handles launching and managing scheduler processes for the gRPC server, -including tensor parallelism, pipeline parallelism, and data parallelism configurations. -""" - -import logging -import multiprocessing as mp -import signal -from typing import Dict, List, Optional, Tuple - -from sglang.srt.managers.data_parallel_controller import ( - run_data_parallel_controller_process, -) -from sglang.srt.managers.scheduler import run_scheduler_process -from sglang.srt.server_args import PortArgs, ServerArgs -from sglang.srt.utils import configure_logger, numa_utils -from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter - -logger = logging.getLogger(__name__) - - -def run_scheduler_with_signal_handling(*args, **kwargs): - """ - Wrapper for run_scheduler_process that ignores SIGINT. - - The scheduler process should not handle Ctrl+C - it should only terminate - when the parent gRPC server exits (via kill_itself_when_parent_died). - - Args: - *args: Positional arguments for run_scheduler_process - **kwargs: Keyword arguments for run_scheduler_process - """ - # Ignore SIGINT in this subprocess - let the parent handle it - signal.signal(signal.SIGINT, signal.SIG_IGN) - - # Now run the actual scheduler process - run_scheduler_process(*args, **kwargs) - - -def launch_scheduler_process_only( - server_args: ServerArgs, - port_args: Optional[PortArgs] = None, -) -> Tuple[Dict, PortArgs, List[mp.Process]]: - """ - Launch only the scheduler process(es) without tokenizer/detokenizer. - - This function handles all scheduler startup logic including: - - Tensor parallelism (tp_size) - - Pipeline parallelism (pp_size) - - Data parallelism (dp_size) - - Multi-node distributed setup - - Args: - server_args: Server configuration - port_args: Port configuration (created if None) - - Returns: - Tuple of (scheduler_info, port_args, scheduler_processes): - - scheduler_info: Dict with model metadata and configuration - - port_args: Port configuration used for IPC - - scheduler_processes: List of launched scheduler Process objects - - Raises: - RuntimeError: If any scheduler process fails to initialize - """ - # Configure global environment - configure_logger(server_args) - server_args.check_server_args() - - # Fix CUDA multiprocessing issues - must be called before any CUDA operations - mp.set_start_method("spawn", force=True) - - # Allocate ports for inter-process communications - if port_args is None: - port_args = PortArgs.init_new(server_args) - logger.info(f"{server_args=}") - - scheduler_procs = [] - - if server_args.dp_size == 1: - # Single data parallel group - launch TP/PP schedulers - memory_saver_adapter = TorchMemorySaverAdapter.create( - enable=server_args.enable_memory_saver - ) - scheduler_pipe_readers = [] - - # Calculate TP/PP distribution across nodes - nnodes_per_tp_group = max(server_args.nnodes // server_args.pp_size, 1) - tp_size_per_node = server_args.tp_size // nnodes_per_tp_group - tp_rank_range = range( - tp_size_per_node * (server_args.node_rank % nnodes_per_tp_group), - tp_size_per_node * (server_args.node_rank % nnodes_per_tp_group + 1), - ) - - pp_size_per_node = max(server_args.pp_size // server_args.nnodes, 1) - pp_rank_range = range( - pp_size_per_node * (server_args.node_rank // nnodes_per_tp_group), - pp_size_per_node * (server_args.node_rank // nnodes_per_tp_group + 1), - ) - - # Launch scheduler for each TP/PP rank combination - for pp_rank in pp_rank_range: - for tp_rank in tp_rank_range: - reader, writer = mp.Pipe(duplex=False) - - # Calculate GPU ID for this rank - gpu_id = ( - server_args.base_gpu_id - + ((pp_rank % pp_size_per_node) * tp_size_per_node) - + (tp_rank % tp_size_per_node) * server_args.gpu_id_step - ) - - # Calculate parallelism ranks (matching engine.py logic) - attn_dp_size = ( - server_args.dp_size if server_args.enable_dp_attention else 1 - ) - attn_tp_size = ( - server_args.tp_size // attn_dp_size // server_args.attn_cp_size - ) - attn_cp_rank = (tp_rank // attn_tp_size) % server_args.attn_cp_size - moe_dp_rank = tp_rank // ( - server_args.tp_size // server_args.moe_dp_size - ) - moe_ep_rank = ( - tp_rank - % (server_args.tp_size // server_args.moe_dp_size) - // ( - server_args.tp_size - // server_args.moe_dp_size - // server_args.ep_size - ) - ) - - # Create scheduler process - proc = mp.Process( - target=run_scheduler_with_signal_handling, - args=( - server_args, - port_args, - gpu_id, - tp_rank, - attn_cp_rank, - moe_dp_rank, - moe_ep_rank, - pp_rank, - None, # dp_rank - writer, - ), - ) - - with memory_saver_adapter.configure_subprocess(), numa_utils.configure_subprocess( - server_args, gpu_id - ): - proc.start() - - scheduler_procs.append(proc) - scheduler_pipe_readers.append(reader) - else: - # Data parallelism - launch data parallel controller - reader, writer = mp.Pipe(duplex=False) - scheduler_pipe_readers = [reader] - - proc = mp.Process( - target=run_data_parallel_controller_process, - args=(server_args, port_args, writer), - ) - proc.start() - scheduler_procs.append(proc) - - # TODO(CatherineSue): handle cases for multi-node - - # Wait for all scheduler processes to be ready - scheduler_infos = [] - for i, reader in enumerate(scheduler_pipe_readers): - try: - data = reader.recv() - except EOFError: - logger.error( - f"Rank {i} scheduler is dead. Please check if there are relevant logs." - ) - scheduler_procs[i].join() - logger.error(f"Exit code: {scheduler_procs[i].exitcode}") - raise RuntimeError(f"Failed to initialize scheduler rank {i}") - - if data.get("status") != "ready": - raise RuntimeError( - f"Scheduler rank {i} initialization failed: {data.get('error', 'Unknown error')}" - ) - scheduler_infos.append(data) - - logger.info( - f"All {len(scheduler_procs)} scheduler process(es) initialized successfully" - ) - - # Return the first scheduler's info (they should all be the same) - return scheduler_infos[0], port_args, scheduler_procs diff --git a/python/sglang/srt/grpc/utils.py b/python/sglang/srt/grpc/utils.py deleted file mode 100644 index 2258a725f..000000000 --- a/python/sglang/srt/grpc/utils.py +++ /dev/null @@ -1,21 +0,0 @@ -"""gRPC utility functions.""" - -from http import HTTPStatus - -import grpc - -_HTTP_TO_GRPC_CODE = { - HTTPStatus.BAD_REQUEST: grpc.StatusCode.INVALID_ARGUMENT, - HTTPStatus.SERVICE_UNAVAILABLE: grpc.StatusCode.UNAVAILABLE, - HTTPStatus.INTERNAL_SERVER_ERROR: grpc.StatusCode.INTERNAL, -} - - -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