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

Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
2026-03-13 09:13:29 -07:00
parent be7a0311a0
commit 654fc02cf1
11 changed files with 15 additions and 2988 deletions
@@ -63,9 +63,7 @@ runtime_common = [
  "uvicorn",
  "uvloop",
  "xgrammar==0.1.27",
-  "smg-grpc-proto>=0.3.3",
+  "smg-grpc-servicer>=0.5.0",
  "grpcio>=1.78.0",
  "grpcio-reflection>=1.78.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",
+  "smg-grpc-servicer>=0.5.0",
  "grpcio>=1.78.0",
  "grpcio-reflection>=1.78.0",
  "grpcio-health-checking>=1.78.0",
 ]
 [[tool.uv.index]]
@@ -67,9 +67,7 @@ dependencies = [
  "uvicorn",
  "uvloop",
  "xgrammar==0.1.27",
-  "smg-grpc-proto>=0.4.1",
+  "smg-grpc-servicer>=0.5.0",
  "grpcio>=1.78.0",
  "grpcio-reflection>=1.78.0",
 ]
 [project.optional-dependencies]
@@ -61,9 +61,7 @@ dependencies = [
  "uvicorn",
  "uvloop",
  "xgrammar==0.1.27",
-  "smg-grpc-proto>=0.4.1",
+  "smg-grpc-servicer>=0.5.0",
  "grpcio>=1.78.0",
  "grpcio-reflection>=1.78.0",
 ]
 [project.optional-dependencies]
@@ -63,9 +63,7 @@ runtime_common = [
  "uvicorn",
  "uvloop",
  "xgrammar==0.1.27",
-  "smg-grpc-proto>=0.4.1",
+  "smg-grpc-servicer>=0.5.0",
  "grpcio>=1.78.0",
  "grpcio-reflection>=1.78.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",
+  "smg-grpc-servicer>=0.5.0",
  "grpcio>=1.78.0",
  "grpcio-reflection>=1.78.0",
 ]
 [project.optional-dependencies]
@@ -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