diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index acf60ccfe..6fff9e6b0 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -401,6 +401,23 @@ from sglang.srt.entrypoints.v1_loads import router as v1_loads_router app.include_router(v1_loads_router) +@app.middleware("http") +async def check_graceful_shutdown(request: Request, call_next): + """Reject all new requests with 503 when the server is shutting down gracefully. + + This middleware runs before any route handler, so the HTTP response has not + started yet and we can safely return a 503 status code — even for streaming + endpoints that would otherwise have already sent 200 headers. + """ + if _global_state is not None and _global_state.tokenizer_manager.gracefully_exit: + if not request.url.path.startswith("/admin/"): + return ORJSONResponse( + content={"error": {"message": "Server is shutting down. No new requests are accepted."}}, + status_code=503, + ) + return await call_next(request) + + @app.exception_handler(HTTPException) async def validation_exception_handler(request: Request, exc: HTTPException): """Enrich HTTP exception with status code and other details. @@ -561,6 +578,16 @@ async def health_generate(request: Request) -> Response: ) _global_state.tokenizer_manager.rid_to_state.pop(rid, None) _global_state.tokenizer_manager.server_status = ServerStatus.UnHealthy + + # Trigger graceful shutdown: stop accepting new requests and drain + exit + if os.getenv("SGLANG_GRACEFUL_EXIT_ON_HEALTH_FAILURE", "1").lower() in ("1", "true"): + logger.error( + "Triggering graceful shutdown due to health check failure. " + "The server will stop accepting new requests and exit after draining. " + "Set SGLANG_GRACEFUL_EXIT_ON_HEALTH_FAILURE=0 to disable this behavior." + ) + _global_state.tokenizer_manager.gracefully_exit = True + return Response(status_code=503) @@ -1390,6 +1417,43 @@ async def continue_generation(obj: ContinueGenerationReqInput, request: Request) ) +##### Admin API endpoints ##### + + +@app.post("/admin/drain") +async def admin_drain(request: Request): + """Initiate graceful drain and shutdown. + + Requires a JSON body with {"password": ""}. + Once triggered, the server rejects all new requests with 503, + waits for in-flight requests to complete, then auto-shuts down. + """ + admin_key = _global_state.tokenizer_manager.server_args.admin_key + if not admin_key: + raise HTTPException(status_code=403, detail="Admin key not configured on this server") + + try: + body = await request.json() + except Exception: + raise HTTPException(status_code=400, detail="Invalid JSON body") + + if body.get("password") != admin_key: + raise HTTPException(status_code=403, detail="Invalid password") + + _global_state.tokenizer_manager.gracefully_exit = True + remaining = len(_global_state.tokenizer_manager.rid_to_state) + logger.info( + f"Admin drain initiated. Remaining in-flight requests: {remaining}. " + f"Server will shut down after all requests complete." + ) + return ORJSONResponse( + { + "message": "Drain initiated. Server will shut down after all in-flight requests complete.", + "remaining_requests": remaining, + } + ) + + ##### OpenAI-compatible API endpoints ##### diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index f31303e8b..f4534b028 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -154,6 +154,9 @@ from sglang.srt.managers.schedule_batch import ( Req, ScheduleBatch, ) +from sglang.srt.managers.scheduler_graceful_shutdown_mixin import ( + SchedulerGracefulShutdownMixin, +) from sglang.srt.managers.schedule_policy import ( AddReqResult, PrefillAdder, @@ -276,6 +279,7 @@ class Scheduler( SchedulerRuntimeCheckerMixin, SchedulerPPMixin, SchedulerDPAttnMixin, + SchedulerGracefulShutdownMixin, SchedulerDllmMixin, ): """A scheduler that manages a tensor parallel GPU worker.""" @@ -3471,6 +3475,7 @@ def run_scheduler_process( trace_set_thread_info(thread_label, tp_rank, dp_rank) # Create a scheduler and run the event loop + scheduler = None try: scheduler = Scheduler( server_args, @@ -3490,7 +3495,13 @@ def run_scheduler_process( # Run the event loop (blocks until shutdown) scheduler.run_event_loop() - except Exception: + except Exception as e: traceback = get_exception_traceback() - logger.error(f"Scheduler hit an exception: {traceback}") + if scheduler is not None and scheduler.is_unrecoverable_error(e): + logger.critical( + f"Scheduler hit an unrecoverable error: {traceback}" + ) + scheduler.handle_graceful_shutdown(e) + else: + logger.error(f"Scheduler hit an exception: {traceback}") parent_process.send_signal(signal.SIGQUIT) diff --git a/python/sglang/srt/managers/scheduler_graceful_shutdown_mixin.py b/python/sglang/srt/managers/scheduler_graceful_shutdown_mixin.py new file mode 100644 index 000000000..327570ff5 --- /dev/null +++ b/python/sglang/srt/managers/scheduler_graceful_shutdown_mixin.py @@ -0,0 +1,345 @@ +"""Mixin for graceful shutdown of the scheduler on unrecoverable errors (CUDA, DeepEP, NCCL, etc.).""" + +from __future__ import annotations + +import logging +from http import HTTPStatus +from typing import TYPE_CHECKING, List + +import torch + +from sglang.srt.managers.io_struct import AbortReq +from sglang.srt.managers.schedule_batch import FINISH_ABORT, Req + +if TYPE_CHECKING: + from sglang.srt.managers.scheduler import Scheduler + +logger = logging.getLogger(__name__) + +# Keywords that indicate an unrecoverable GPU/communication error +_UNRECOVERABLE_ERROR_KEYWORDS = [ + "cuda error", + "cudaerror", + "cuda_error", + "deepep", + "nccl", + "cudaerrorlaunchfailure", + "unspecified launch failure", + "device-side assert", + "out of memory", +] + + +class SchedulerGracefulShutdownMixin: + """Mixin that adds graceful shutdown capability to the Scheduler. + + When an unrecoverable error (e.g., CUDA error, DeepEP timeout, NCCL error) is detected: + 1. All in-flight requests are failed gracefully: + - Requests that the user has already received data for (send_token_offset > 0): + receive finish_reason + usage + [DONE] via stream_output + - Requests that the user hasn't received any data for: + receive a 503 error via AbortReq + 2. All queued requests are aborted with 503 + 3. The scheduler exits the event loop cleanly + """ + + @staticmethod + def is_unrecoverable_error(error: Exception) -> bool: + """Check if an exception represents an unrecoverable GPU/communication error.""" + # Check for torch CUDA errors + if hasattr(torch, "AcceleratorError") and isinstance( + error, torch.AcceleratorError + ): + return True + if hasattr(torch.cuda, "CudaError") and isinstance( + error, torch.cuda.CudaError + ): + return True + + # Check error message for known patterns + if isinstance(error, (RuntimeError, OSError)): + msg = str(error).lower() + for keyword in _UNRECOVERABLE_ERROR_KEYWORDS: + if keyword in msg: + return True + + return False + + def handle_graceful_shutdown(self: Scheduler, error: Exception) -> None: + """Handle an unrecoverable error by gracefully failing all requests and exiting. + + This method does NOT perform any CUDA operations since the GPU state may be corrupted. + It only sends zmq messages to notify upstream components about the failure. + """ + error_msg = ( + f"Unrecoverable error, scheduler shutting down: " + f"{type(error).__name__}: {error}" + ) + logger.critical(error_msg) + + try: + self._fail_all_inflight_requests(error_msg) + except Exception as cleanup_error: + logger.error( + f"Error during graceful shutdown (inflight requests): {cleanup_error}" + ) + + try: + self._fail_all_queued_requests(error_msg) + except Exception as cleanup_error: + logger.error( + f"Error during graceful shutdown (queued requests): {cleanup_error}" + ) + + try: + self._fail_overlap_buffered_requests(error_msg) + except Exception as cleanup_error: + logger.error( + f"Error during graceful shutdown (overlap buffer): {cleanup_error}" + ) + + try: + self._fail_disagg_requests(error_msg) + except Exception as cleanup_error: + logger.error( + f"Error during graceful shutdown (disagg queues): {cleanup_error}" + ) + + logger.critical( + "Scheduler graceful shutdown complete. All requests have been notified." + ) + + def _fail_all_inflight_requests(self: Scheduler, error_msg: str) -> None: + """Fail all requests currently in the running batch and current batch. + + Requests are split into two groups based on whether the user has already + received data (send_token_offset > 0): + - User received data: finish via stream_output with finish_reason + [DONE] + - User not received data: abort with 503 via AbortReq + """ + # Collect all unique in-flight requests + seen_rids = set() + reqs_user_received: List[Req] = [] + reqs_user_not_received: List[Req] = [] + + for batch in [self.running_batch, self.cur_batch, self.last_batch]: + if batch is None: + continue + for req in batch.reqs: + if req.rid not in seen_rids and not req.finished(): + seen_rids.add(req.rid) + if req.send_token_offset > 0: + reqs_user_received.append(req) + else: + reqs_user_not_received.append(req) + + # For requests the user has already received data from: + # Send a final finish via stream_output (finish_reason + usage + [DONE]) + # Use status_code=None so the API layer does NOT return an HTTP error + if reqs_user_received: + for req in reqs_user_received: + req.finished_reason = FINISH_ABORT( + message=error_msg, + status_code=None, + err_type=None, + ) + if req.finished_len is None: + req.finished_len = len(req.output_ids) + + try: + self.stream_output(reqs_user_received, False) + except Exception as e: + logger.error( + f"Failed to stream output for inflight requests: {e}" + ) + # Fallback: send individual AbortReq for each request + for req in reqs_user_received: + try: + self._send_abort_for_req(req, error_msg) + except Exception: + pass + + # For requests the user hasn't received data from: send 503 abort + for req in reqs_user_not_received: + try: + self._send_abort_for_req(req, error_msg) + except Exception as e: + logger.error( + f"Failed to send abort for request {req.rid}: {e}" + ) + + total = len(reqs_user_received) + len(reqs_user_not_received) + if total > 0: + logger.warning( + f"Failed {total} in-flight requests " + f"({len(reqs_user_received)} with data already sent to user, " + f"{len(reqs_user_not_received)} without)" + ) + + def _fail_all_queued_requests(self: Scheduler, error_msg: str) -> None: + """Fail all requests in the waiting queue and grammar queue.""" + count = 0 + + # Fail waiting queue + for req in self.waiting_queue: + self._send_abort_for_req(req, error_msg) + count += 1 + self.waiting_queue.clear() + + # Fail grammar queue + for req in self.grammar_queue: + if hasattr(req, "grammar") and req.grammar is not None: + try: + req.grammar.cancel() + except Exception: + pass + self._send_abort_for_req(req, error_msg) + count += 1 + self.grammar_queue.clear() + + # Fail chunked request if any + if self.chunked_req is not None: + if self.chunked_req.send_token_offset > 0: + # User already received data for this chunked request + self.chunked_req.finished_reason = FINISH_ABORT( + message=error_msg, + status_code=None, + err_type=None, + ) + if self.chunked_req.finished_len is None: + self.chunked_req.finished_len = len(self.chunked_req.output_ids) + try: + self.stream_output([self.chunked_req], False) + except Exception: + self._send_abort_for_req(self.chunked_req, error_msg) + else: + self._send_abort_for_req(self.chunked_req, error_msg) + self.chunked_req = None + count += 1 + + if count > 0: + logger.warning(f"Aborted {count} queued requests") + + def _fail_overlap_buffered_requests(self: Scheduler, error_msg: str) -> None: + """Fail requests buffered in overlap scheduling's result_queue.""" + if not self.enable_overlap or not hasattr(self, "result_queue"): + return + + count = 0 + while self.result_queue: + try: + tmp_batch, _ = self.result_queue.popleft() + + reqs_user_received = [] + reqs_user_not_received = [] + + for req in tmp_batch.reqs: + if req.finished(): + continue + if req.send_token_offset > 0: + reqs_user_received.append(req) + else: + reqs_user_not_received.append(req) + + # Finish requests the user has already received data from + for req in reqs_user_received: + req.finished_reason = FINISH_ABORT( + message=error_msg, + status_code=None, + err_type=None, + ) + if req.finished_len is None: + req.finished_len = len(req.output_ids) + count += 1 + + if reqs_user_received: + try: + self.stream_output(reqs_user_received, False) + except Exception: + for req in reqs_user_received: + self._send_abort_for_req(req, error_msg) + + # Abort requests the user hasn't received data from + for req in reqs_user_not_received: + self._send_abort_for_req(req, error_msg) + count += 1 + + except Exception as e: + logger.error(f"Error processing overlap buffer: {e}") + break + + if count > 0: + logger.warning(f"Failed {count} overlap-buffered requests") + + def _fail_disagg_requests(self: Scheduler, error_msg: str) -> None: + """Fail all requests in disaggregation queues.""" + from sglang.srt.disaggregation.utils import DisaggregationMode + + count = 0 + + if self.disaggregation_mode == DisaggregationMode.DECODE: + # Fail prealloc queue + if hasattr(self, "disagg_decode_prealloc_queue"): + queue = self.disagg_decode_prealloc_queue + + if hasattr(queue, "queue"): + for decode_req in queue.queue: + self._send_abort_for_req(decode_req.req, error_msg) + count += 1 + queue.queue.clear() + + if hasattr(queue, "retracted_queue"): + for req in queue.retracted_queue: + self._send_abort_for_req(req, error_msg) + count += 1 + queue.retracted_queue.clear() + + if hasattr(queue, "pending_route_queue"): + for req in queue.pending_route_queue: + self._send_abort_for_req(req, error_msg) + count += 1 + queue.pending_route_queue.clear() + + # Fail transfer queue + if hasattr(self, "disagg_decode_transfer_queue"): + if hasattr(self.disagg_decode_transfer_queue, "queue"): + for decode_req in self.disagg_decode_transfer_queue.queue: + self._send_abort_for_req(decode_req.req, error_msg) + count += 1 + self.disagg_decode_transfer_queue.queue.clear() + + elif self.disaggregation_mode == DisaggregationMode.PREFILL: + # Fail bootstrap queue + if hasattr(self, "disagg_prefill_bootstrap_queue"): + if hasattr(self.disagg_prefill_bootstrap_queue, "queue"): + for req in self.disagg_prefill_bootstrap_queue.queue: + self._send_abort_for_req(req, error_msg) + count += 1 + self.disagg_prefill_bootstrap_queue.queue.clear() + + # Fail inflight queue + if hasattr(self, "disagg_prefill_inflight_queue"): + for req in self.disagg_prefill_inflight_queue: + self._send_abort_for_req(req, error_msg) + count += 1 + self.disagg_prefill_inflight_queue.clear() + + if count > 0: + logger.warning(f"Aborted {count} disaggregation-queued requests") + + def _send_abort_for_req(self: Scheduler, req: Req, error_msg: str) -> None: + """Send an AbortReq to the tokenizer for a single request (503 error).""" + try: + abort_req = AbortReq( + rid=req.rid, + finished_reason={ + "type": "abort", + "status_code": int(HTTPStatus.SERVICE_UNAVAILABLE), + "message": error_msg, + }, + ) + # Propagate http_worker_ipc for multi-http-worker routing + abort_req.http_worker_ipc = getattr(req, "http_worker_ipc", None) + self.send_to_tokenizer.send_output(abort_req) + except Exception as e: + logger.error(f"Failed to send abort for request {req.rid}: {e}") diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index d7a324afd..4550b0381 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -30,7 +30,7 @@ from datetime import datetime from enum import Enum from http import HTTPStatus from typing import Any, Awaitable, Dict, List, Optional, Tuple, Union - +import time import fastapi import uvloop import zmq @@ -483,6 +483,7 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi obj: Union[GenerateReqInput, EmbeddingReqInput], request: Optional[fastapi.Request] = None, ): + created_time = obj.received_time if obj.received_time else time.time() self.auto_create_handle_loop() # Normalize the request @@ -2125,10 +2126,21 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi remaining_rids = list(self.rid_to_state.keys()) if self.server_status == ServerStatus.UnHealthy: - # if health check failed, we should exit immediately + # Detokenizer is hung, in-flight requests cannot complete either. + # Before force exiting, properly finish/abort all pending requests + # so that clients get a proper response instead of a connection drop. logger.error( - "Signal SIGTERM received while health check failed. Force exiting." + f"Health check failed (detokenizer hang). " + f"{remain_num_req} requests still pending. " + f"Finishing pending requests before force exiting. " + f"{remaining_rids=}" ) + + self._finish_all_pending_requests_on_shutdown() + + # Give the event loop time to wake up coroutines and send HTTP responses + await asyncio.sleep(2) + self.dump_requests_before_crash() self.force_exit_handler() break @@ -2157,6 +2169,61 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi """Put some custom force exit logic here.""" pass + def _finish_all_pending_requests_on_shutdown(self): + """Finish all pending requests when the server is shutting down due to detokenizer hang. + + For each pending request in rid_to_state: + - Streaming requests that already sent data to the client: finish normally with + finish_reason="stop" so the client receives [DONE] cleanly. + - Non-streaming requests or streaming requests that haven't sent data yet: abort + with 503 so the client gets a proper error response instead of a connection drop. + """ + error_msg = "Server is shutting down due to detokenizer hang." + finished_count = 0 + aborted_count = 0 + + for rid, state in list(self.rid_to_state.items()): + if state.finished: + continue + + is_stream = getattr(state.obj, "stream", False) + already_sent_data = state.response_sent_to_client_ts > 0 + + if is_stream and already_sent_data: + # Streaming request that already sent partial data to client. + # The HTTP 200 headers are already sent, so we can't change the status code. + finish_reason = {"type": "stop", "matched": None} + finished_count += 1 + else: + # Non-streaming request, or streaming request before the first chunk. + # Return a 503 abort so the client gets a proper error response. + finish_reason = { + "type": "abort", + "status_code": int(HTTPStatus.SERVICE_UNAVAILABLE), + "message": error_msg, + } + aborted_count += 1 + + state.finished = True + out = { + "text": state.text, + "output_ids": state.output_ids, + "meta_info": { + "id": rid, + "finish_reason": finish_reason, + "prompt_tokens": 0, + "completion_tokens": len(state.output_ids), + "cached_tokens": 0, + }, + } + state.out_list.append(out) + state.event.set() + + logger.info( + f"Shutdown: finished {finished_count} streaming requests normally, " + f"aborted {aborted_count} requests with 503." + ) + def _handle_abort_req(self, recv_obj: AbortReq): if is_health_check_generate_req(recv_obj): return diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index f94da569a..d9f1b5614 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -423,6 +423,7 @@ class ServerArgs: # API related api_key: Optional[str] = None admin_api_key: Optional[str] = None + admin_key: Optional[str] = None served_model_name: Optional[str] = None weight_version: str = "default" chat_template: Optional[str] = None @@ -4329,6 +4330,12 @@ class ServerArgs: "When set, admin endpoints require this key and do NOT accept --api-key." ), ) + parser.add_argument( + "--admin-key", + type=str, + default=ServerArgs.admin_key, + help="Set admin key for privileged endpoints like /admin/drain.", + ) parser.add_argument( "--served-model-name", type=str,