[gRPC] Extract gRPC servicer into standalone package (#20478)
Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
This commit is contained in:
4
3rdparty/amd/wheel/sglang/pyproject.toml
vendored
4
3rdparty/amd/wheel/sglang/pyproject.toml
vendored
@@ -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/)
|
||||
|
||||
@@ -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]]
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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 = [
|
||||
|
||||
@@ -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]
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -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.
|
||||
@@ -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
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user