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

feat(drain): graceful shutdown and admin drain

Browse Source 
feat: add admin drain to sglang
This commit is contained in:
laoyao0822
2026-03-25 21:11:50 +08:00
committed by wxiwnd
parent 150d4db41d
commit 9b5ecc8c1f
5 changed files with 499 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

64
python/sglang/srt/entrypoints/http_server.py
View File

@@ -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": "<admin_key>"}.
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 #####

15
python/sglang/srt/managers/scheduler.py
View File

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

345
python/sglang/srt/managers/scheduler_graceful_shutdown_mixin.py Normal file
View File

@@ -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}")

73
python/sglang/srt/managers/tokenizer_manager.py
View File

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

7
python/sglang/srt/server_args.py
View File

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