fix(api): return 429 when server is busy

lightllm/server/api_anthropic.py

@@ -840,6 +840,22 @@ async def _openai_sse_to_anthropic_events(
                 logger.debug("Skipping non-JSON SSE payload: %r", payload)
                 continue
 
+            if "error" in chunk and "choices" not in chunk:
+                error = chunk["error"]
+                error_type = error.get("type")
+                if error.get("code") == 429 or error_type in ("RateLimitError", "rate_limit_error"):
+                    error_type = "rate_limit_error"
+                elif error_type != "invalid_request_error":
+                    error_type = "api_error"
+                yield _sse_event(
+                    "error",
+                    {
+                        "type": "error",
+                        "error": {"type": error_type, "message": error.get("message", "generation failed")},
+                    },
+                )
+                return
+
             # final_output_tokens is sourced exclusively from the trailing usage
             # chunk emitted by chat_completions_impl; we intentionally do not
             # estimate it per delta because that would diverge from the
@@ -1096,7 +1112,7 @@ def _rewrap_openai_error_as_anthropic(resp: JSONResponse) -> JSONResponse:
 async def anthropic_messages_impl(raw_request: Request) -> Response:
     # Lazy imports to avoid pulling in heavy server deps at module import time.
     from .api_models import ChatCompletionRequest, ChatCompletionResponse
-    from .api_openai import chat_completions_impl
+    from .api_openai import chat_completions_impl, prime_pd_master_streaming_response
 
     try:
         raw_body = await raw_request.json()
@@ -1135,6 +1151,8 @@ async def anthropic_messages_impl(raw_request: Request) -> Response:
                 return _rewrap_openai_error_as_anthropic(downstream)
             return downstream
 
+        downstream = await prime_pd_master_streaming_response(downstream)
+
         message_id = f"msg_{uuid.uuid4().hex[:24]}"
         anthropic_stream = _openai_sse_to_anthropic_events(
             downstream.body_iterator, requested_model=requested_model, message_id=message_id

lightllm/server/api_http.py

@@ -52,7 +52,7 @@
 from lightllm.utils.shm_port_args import get_shm_port_args
 from dataclasses import asdict, dataclass, is_dataclass
 
-from .api_openai import chat_completions_impl, completions_impl
+from .api_openai import chat_completions_impl, completions_impl, prime_pd_master_streaming_response
 from .api_models import (
     ChatCompletionRequest,
     ChatCompletionResponse,
@@ -161,6 +161,8 @@ def create_error_response(
             err_type = "InternalServerError"
         elif status_code == HTTPStatus.NOT_FOUND:
             err_type = "NotFoundError"
+        elif status_code == HTTPStatus.TOO_MANY_REQUESTS:
+            err_type = "RateLimitError"
         else:
             err_type = "BadRequestError"
 
@@ -171,6 +173,17 @@ def create_error_response(
     )
 
 
+def create_server_busy_response(exc: ServerBusyError) -> JSONResponse:
+    status = HTTPStatus(exc.status_code)
+    return create_error_response(status, str(exc), err_type="RateLimitError")
+
+
+@app.exception_handler(ServerBusyError)
+async def server_busy_exception_handler(request: Request, exc: ServerBusyError) -> JSONResponse:
+    logger.warning(str(exc))
+    return create_server_busy_response(exc)
+
+
 @app.get("/liveness")
 @app.post("/liveness")
 def liveness():
@@ -294,8 +307,8 @@ async def generate(request: Request) -> Response:
     try:
         return await g_objs.g_generate_func(request, g_objs.httpserver_manager)
     except ServerBusyError as e:
-        logger.error("%s", str(e), exc_info=True)
-        return create_error_response(HTTPStatus.SERVICE_UNAVAILABLE, str(e))
+        logger.warning(str(e))
+        return create_server_busy_response(e)
     except ValueError as e:
         return create_error_response(HTTPStatus.BAD_REQUEST, str(e))
     except ClientDisconnected as e:
@@ -314,10 +327,11 @@ async def generate_stream(request: Request) -> Response:
         )
 
     try:
-        return await g_objs.g_generate_stream_func(request, g_objs.httpserver_manager)
+        response = await g_objs.g_generate_stream_func(request, g_objs.httpserver_manager)
+        return await prime_pd_master_streaming_response(response)
     except ServerBusyError as e:
-        logger.error("%s", str(e), exc_info=True)
-        return create_error_response(HTTPStatus.SERVICE_UNAVAILABLE, str(e))
+        logger.warning(str(e))
+        return create_server_busy_response(e)
     except ValueError as e:
         return create_error_response(HTTPStatus.BAD_REQUEST, str(e))
     except ClientDisconnected as e:
@@ -337,6 +351,9 @@ async def get_score(request: Request) -> Response:
 
     try:
         return await lightllm_get_score(request, g_objs.httpserver_manager)
+    except ServerBusyError as e:
+        logger.warning(str(e))
+        return create_server_busy_response(e)
     except ClientDisconnected as e:
         logger.warning(str(e))
         return Response(status_code=499)
@@ -368,11 +385,12 @@ async def chat_completions(request: ChatCompletionRequest, raw_request: Request)
 
     try:
         resp = await chat_completions_impl(request, raw_request)
+        resp = await prime_pd_master_streaming_response(resp)
     except ValueError as e:
         return create_error_response(HTTPStatus.BAD_REQUEST, str(e))
     except ServerBusyError as e:
         logger.warning(str(e))
-        return create_error_response(HTTPStatus.SERVICE_UNAVAILABLE, str(e))
+        return create_server_busy_response(e)
     except ClientDisconnected as e:
         logger.warning(str(e))
         return Response(status_code=499)
@@ -388,11 +406,12 @@ async def completions(request: CompletionRequest, raw_request: Request) -> Respo
 
     try:
         resp = await completions_impl(request, raw_request)
+        resp = await prime_pd_master_streaming_response(resp)
     except ValueError as e:
         return create_error_response(HTTPStatus.BAD_REQUEST, str(e))
     except ServerBusyError as e:
         logger.warning(str(e))
-        return create_error_response(HTTPStatus.SERVICE_UNAVAILABLE, str(e))
+        return create_server_busy_response(e)
     except ClientDisconnected as e:
         logger.warning(str(e))
         return Response(status_code=499)
@@ -405,10 +424,15 @@ async def anthropic_messages(raw_request: Request) -> Response:
         return create_error_response(
             HTTPStatus.EXPECTATION_FAILED, "service in pd mode dont recv reqs from http interface"
         )
-    from .api_anthropic import anthropic_messages_impl
+    from .api_anthropic import _anthropic_error_response, anthropic_messages_impl
 
     try:
-        return await anthropic_messages_impl(raw_request)
+        response = await anthropic_messages_impl(raw_request)
+        return await prime_pd_master_streaming_response(response)
+    except ServerBusyError as e:
+        logger.warning(str(e))
+        g_objs.metric_client.counter_inc("lightllm_request_failure")
+        return _anthropic_error_response(HTTPStatus(e.status_code), str(e))
     except ClientDisconnected as e:
         logger.warning(str(e))
         return Response(status_code=499)
@@ -423,10 +447,11 @@ async def openai_responses(raw_request: Request) -> Response:
     from .api_responses import responses_impl
 
     try:
-        return await responses_impl(raw_request)
+        response = await responses_impl(raw_request)
+        return await prime_pd_master_streaming_response(response)
     except ServerBusyError as e:
-        logger.error("%s", str(e), exc_info=True)
-        return create_error_response(HTTPStatus.SERVICE_UNAVAILABLE, str(e))
+        logger.warning(str(e))
+        return create_server_busy_response(e)
     except ValueError as e:
         return create_error_response(HTTPStatus.BAD_REQUEST, str(e))
     except ClientDisconnected as e:

lightllm/server/api_openai.py

@@ -30,7 +30,7 @@
 from .httpserver_for_pd_master.manager import HttpServerManagerForPDMaster
 from .api_lightllm import lightllm_get_score
 from lightllm.utils.envs_utils import get_env_start_args, get_lightllm_websocket_max_message_size
-from lightllm.utils.error_utils import ClientDisconnected
+from lightllm.utils.error_utils import ClientDisconnected, ServerBusyError
 
 from lightllm.utils.log_utils import init_logger
 from lightllm.server.metrics.manager import MetricClient
@@ -60,15 +60,53 @@
 logger = init_logger(__name__)
 
 
+async def prime_pd_master_streaming_response(response: Response) -> Response:
+    """In PD-master mode, read the first stream item before sending the HTTP status."""
+    if not isinstance(response, StreamingResponse):
+        return response
+    if get_env_start_args().run_mode != "pd_master":
+        return response
+
+    body_iterator = response.body_iterator
+    try:
+        first_item = await body_iterator.__anext__()
+    except StopAsyncIteration:
+        return response
+
+    async def replay_stream():
+        try:
+            yield first_item
+            async for item in body_iterator:
+                yield item
+        finally:
+            close = getattr(body_iterator, "aclose", None)
+            if close is not None:
+                await close()
+
+    response.body_iterator = replay_stream()
+    return response
+
+
 async def _safe_stream_wrapper(stream_generator):
     """Wrap a streaming generator to catch ValueError (e.g. input too long) and yield an SSE error
     event instead of letting the exception propagate to Starlette which prints a long traceback."""
+    stream_started = False
     try:
         async for item in stream_generator:
             yield item
+            stream_started = True
     except ValueError as e:
         error_data = json.dumps({"error": {"message": str(e), "type": "invalid_request_error"}}, ensure_ascii=False)
         yield f"data: {error_data}\n\n"
+    except ServerBusyError as e:
+        if not stream_started:
+            raise
+        logger.error("Generation interrupted after the stream started: %s", e.message)
+        error_data = json.dumps(
+            {"error": {"message": e.message, "type": "server_error", "code": "stream_error"}},
+            ensure_ascii=False,
+        )
+        yield f"data: {error_data}\n\n"
     except ClientDisconnected as e:
         logger.warning(str(e))
         # Client is gone — there's no point yielding more SSE chunks. Stop quietly.

lightllm/server/api_responses.py

@@ -528,7 +528,12 @@ def open_item(kind: str, item: Dict[str, Any]):
 
     if failed_error is not None:
         response["status"] = "failed"
-        response["error"] = {"code": "server_error", "message": failed_error.get("message", "generation failed")}
+        error_code = failed_error.get("code")
+        if error_code == 429 or failed_error.get("type") in ("RateLimitError", "rate_limit_error"):
+            error_code = "rate_limit_error"
+        else:
+            error_code = "server_error"
+        response["error"] = {"code": error_code, "message": failed_error.get("message", "generation failed")}
         yield event("response.failed", {"response": response})
         return
 
@@ -547,7 +552,7 @@ def open_item(kind: str, item: Dict[str, Any]):
 
 async def responses_impl(raw_request: Request) -> Response:
     from .api_models import ChatCompletionRequest, ChatCompletionResponse
-    from .api_openai import chat_completions_impl, create_error_response
+    from .api_openai import chat_completions_impl, create_error_response, prime_pd_master_streaming_response
 
     try:
         body = await raw_request.json()
@@ -582,6 +587,7 @@ async def responses_impl(raw_request: Request) -> Response:
     if chat_request.stream:
         if not isinstance(downstream, StreamingResponse):
             return downstream
+        downstream = await prime_pd_master_streaming_response(downstream)
         return StreamingResponse(
             _openai_sse_to_responses_events(downstream.body_iterator, body),
             media_type="text/event-stream",

lightllm/utils/error_utils.py

@@ -7,17 +7,17 @@
 class ServerBusyError(Exception):
     """Custom exception for server busy/overload situations"""
 
-    def __init__(self, message="Server is busy, please try again later", status_code=503):
+    def __init__(self, message="Server is busy, please try again later", status_code=429):
         """
         Initialize the ServerBusyError
 
         Args:
             message (str): Error message to display
-            status_code (int): HTTP status code (default 503 Service Unavailable)
+            status_code (int): HTTP status code (default 429 Too Many Requests)
         """
         super().__init__(message)
         self.message = message
-        self.status_code = status_code  # HTTP 503 Service Unavailable
+        self.status_code = status_code  # HTTP 429 Too Many Requests
 
     def __str__(self):
         """String representation of the error"""