sglang 0.5.3rc2__py3-none-any.whl → 0.5.4__py3-none-any.whl

sglang/srt/{entrypoints → grpc}/grpc_request_manager.py

@@ -263,8 +263,8 @@ class GrpcRequestManager:
                         response = await task
 
                         # Add index for client-side ordering
-                        if isinstance(response, dict) and "meta_info" in response:
-                            response_rid = response["meta_info"].get("id", "")
+                        if isinstance(response, dict):
+                            response_rid = response.get("request_id", "")
                             if response_rid in rid_to_index:
                                 response["index"] = rid_to_index[response_rid]
 
@@ -318,13 +318,8 @@ class GrpcRequestManager:
             is_stream = getattr(obj, "stream", False)
 
             while True:
-                # Client cancelled - notify scheduler and exit
-                if grpc_context and grpc_context.cancelled():
-                    await self.abort_request(request_id)
-                    return
-
                 try:
-                    response = await asyncio.wait_for(state.out_queue.get(), timeout=4)
+                    response = await state.out_queue.get()
 
                     if is_stream:
                         yield response
@@ -337,13 +332,11 @@ class GrpcRequestManager:
                             yield final_response
                         break
 
-                except asyncio.TimeoutError:
-                    # Timeout waiting for response - abort and cleanup
-                    logger.warning(
-                        f"Timeout waiting for response for request {request_id}"
-                    )
+                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)
-                    return
+                    raise  # Re-raise to let gRPC server handle cleanup
 
         finally:
             # Always clean up request state when exiting
@@ -397,9 +390,7 @@ class GrpcRequestManager:
         # Wait for result in background
         async def wait_for_result():
             try:
-                # Wait for completion
                 await state.event.wait()
-                # Get result from queue
                 result = await state.out_queue.get()
                 future.set_result(result)
             except Exception as e:
@@ -413,43 +404,34 @@ class GrpcRequestManager:
         return future
 
     async def abort_request(self, request_id: str) -> bool:
-        """Abort a running request."""
-        if request_id not in self.rid_to_state:
-            return False
+        """Abort a running request.
 
-        # Send abort to scheduler
-        abort_req = AbortReq(rid=request_id)
-        try:
-            await self._send_to_scheduler(abort_req)
-        except Exception as e:
-            logger.error(f"Failed to send abort request: {e}")
+        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 as finished
+        # 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
-            state.event.set()
+            logger.debug(f"Marked request {request_id} as aborted locally")
 
-            # Send abort notification to output queue
-            await state.out_queue.put({"error": "Request aborted", "abort": True})
+        # 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 pause_generation(self):
-        """Pause generation processing."""
-        async with self.is_pause_cond:
-            self.is_pause = True
-            logger.info("Generation paused")
-
-    async def resume_generation(self):
-        """Resume generation processing."""
-        async with self.is_pause_cond:
-            self.is_pause = False
-            self.is_pause_cond.notify_all()
-            logger.info("Generation resumed")
-
     async def handle_loop(self):
         """
         Main event loop - processes outputs from scheduler.
@@ -461,10 +443,11 @@ class GrpcRequestManager:
                 recv_obj = await self.recv_from_scheduler.recv_pyobj()
                 self.last_receive_tstamp = time.time()
 
-                # Check for pause
-                async with self.is_pause_cond:
-                    while self.is_pause:
-                        await self.is_pause_cond.wait()
+                # 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):
@@ -473,6 +456,8 @@ class GrpcRequestManager:
                     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)
                 else:
                     logger.warning(f"Unknown output type: {type(recv_obj)}")
 
@@ -547,6 +532,11 @@ class GrpcRequestManager:
 
     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 = []
+        now = time.time()
+
         # Process each request in the batch
         for i, rid in enumerate(batch_out.rids):
             if rid not in self.rid_to_state:
@@ -554,8 +544,12 @@ class GrpcRequestManager:
 
             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
-            now = time.time()
             if state.first_token_time == 0.0:
                 state.first_token_time = now
             state.last_time = now
@@ -649,7 +643,8 @@ class GrpcRequestManager:
             if output_data["token_ids"]:
                 state.output_ids.extend(output_data["token_ids"])
 
-            await state.out_queue.put(output_data)
+            # Add queue.put() to parallel task list
+            put_tasks.append(state.out_queue.put(output_data))
 
             # Handle completion
             if output_data["finished"]:
@@ -659,12 +654,16 @@ class GrpcRequestManager:
                 state.event.set()
 
                 # Remove from tracking after a delay
-                async def cleanup():
+                async def cleanup(request_id):
                     await asyncio.sleep(5.0)
-                    if rid in self.rid_to_state:
-                        del self.rid_to_state[rid]
+                    if request_id in self.rid_to_state:
+                        del self.rid_to_state[request_id]
 
-                asyncio.create_task(cleanup())
+                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."""
@@ -726,6 +725,67 @@ class GrpcRequestManager:
         state.finished_time = time.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:

sglang/srt/grpc/health_servicer.py

@@ -0,0 +1,189 @@
+"""
+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.

sglang/srt/grpc/scheduler_launcher.py

@@ -0,0 +1,181 @@
+"""
+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, prepare_model_and_tokenizer
+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=}")
+
+    # Prepare model and tokenizer paths
+    server_args.model_path, server_args.tokenizer_path = prepare_model_and_tokenizer(
+        server_args.model_path, server_args.tokenizer_path
+    )
+
+    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 MoE expert parallel rank
+                moe_ep_rank = tp_rank // (server_args.tp_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,
+                        moe_ep_rank,
+                        pp_rank,
+                        None,  # dp_rank
+                        writer,
+                    ),
+                )
+
+                with memory_saver_adapter.configure_subprocess():
+                    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