chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[gRPC] Extract gRPC servicer into standalone package (#20478)

Browse Source 
Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
This commit is contained in:
Simo Lin
2026-03-13 09:13:29 -07:00
committed by GitHub
parent be7a0311a0
commit 654fc02cf1
11 changed files with 15 additions and 2988 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
python/pyproject.toml
View File

@@ -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]]

4
python/pyproject_cpu.toml
View File

@@ -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]

4
python/pyproject_npu.toml
View File

@@ -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]

4
python/pyproject_other.toml
View File

@@ -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 = [

4
python/pyproject_xpu.toml
View File

@@ -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]

1548
python/sglang/srt/entrypoints/grpc_server.py
View File

File diff suppressed because it is too large Load Diff

1022
python/sglang/srt/grpc/grpc_request_manager.py
View File

File diff suppressed because it is too large Load Diff

189
python/sglang/srt/grpc/health_servicer.py
View File

@@ -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.

198
python/sglang/srt/grpc/scheduler_launcher.py
View File

@@ -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

21
python/sglang/srt/grpc/utils.py
View File

@@ -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