diff --git a/lightllm/server/api_anthropic.py b/lightllm/server/api_anthropic.py index 720072476..7ee311568 100644 --- a/lightllm/server/api_anthropic.py +++ b/lightllm/server/api_anthropic.py @@ -8,6 +8,7 @@ chat_completions_impl and re-emits it as the Anthropic event sequence (message_start, content_block_*, message_delta, message_stop). """ + from __future__ import annotations import asyncio @@ -33,6 +34,8 @@ from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.log_utils import init_logger +from .api_errors import is_rate_limit_error + logger = init_logger(__name__) _cached_adapter: Any = None @@ -780,6 +783,15 @@ def _sse_event(event_type: str, data_obj: Dict[str, Any]) -> bytes: return f"event: {event_type}\ndata: {json.dumps(data_obj)}\n\n".encode("utf-8") +def _anthropic_error_type(error: Dict[str, Any]) -> str: + error_type = error.get("type") + if is_rate_limit_error(error): + return "rate_limit_error" + if error_type == "invalid_request_error": + return error_type + return "api_error" + + async def _openai_sse_to_anthropic_events( openai_body_iterator, requested_model: str, @@ -842,16 +854,14 @@ async def _openai_sse_to_anthropic_events( 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")}, + "error": { + "type": _anthropic_error_type(error), + "message": error.get("message", "generation failed"), + }, }, ) return @@ -1112,7 +1122,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, prime_pd_master_streaming_response + from .api_openai import chat_completions_impl try: raw_body = await raw_request.json() @@ -1143,21 +1153,19 @@ async def anthropic_messages_impl(raw_request: Request) -> Response: downstream = await chat_completions_impl(chat_request, raw_request) if is_stream: - from fastapi.responses import StreamingResponse + from .api_stream_obj import CustomStreamingResponse - if not isinstance(downstream, StreamingResponse): + if not isinstance(downstream, CustomStreamingResponse): # chat_completions_impl returned an OpenAI-format error — rewrap it. if isinstance(downstream, JSONResponse): 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 ) - return StreamingResponse(anthropic_stream, media_type="text/event-stream") + return CustomStreamingResponse(anthropic_stream, media_type="text/event-stream") if not isinstance(downstream, ChatCompletionResponse): if isinstance(downstream, JSONResponse): diff --git a/lightllm/server/api_errors.py b/lightllm/server/api_errors.py new file mode 100644 index 000000000..59cfbe8fe --- /dev/null +++ b/lightllm/server/api_errors.py @@ -0,0 +1,40 @@ +"""Shared helpers for OpenAI-compatible API error responses.""" + +from http import HTTPStatus +from typing import Any, Mapping + +from fastapi.responses import JSONResponse + +from lightllm.utils.error_utils import ServerBusyError + +_RATE_LIMIT_ERROR_TYPES = {"RateLimitError", "rate_limit_error"} + + +def is_rate_limit_error(error: Mapping[str, Any]) -> bool: + return error.get("code") == HTTPStatus.TOO_MANY_REQUESTS or error.get("type") in _RATE_LIMIT_ERROR_TYPES + + +def create_error_response( + status_code: HTTPStatus, message: str, err_type: str = None, param: str = None +) -> JSONResponse: + if err_type is None: + if status_code.value >= 500: + 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" + + from .api_http import g_objs + + g_objs.metric_client.counter_inc("lightllm_request_failure") + return JSONResponse( + {"error": {"message": message, "type": err_type, "param": param, "code": status_code.value}}, + status_code=status_code.value, + ) + + +def create_server_busy_response(exc: ServerBusyError) -> JSONResponse: + return create_error_response(HTTPStatus(exc.status_code), str(exc)) diff --git a/lightllm/server/api_http.py b/lightllm/server/api_http.py index 1cda03b96..d985b5e66 100755 --- a/lightllm/server/api_http.py +++ b/lightllm/server/api_http.py @@ -37,7 +37,7 @@ from typing import Callable from lightllm.server import TokenLoad from fastapi import BackgroundTasks, FastAPI, Request -from fastapi.responses import Response, StreamingResponse, JSONResponse +from fastapi.responses import Response, JSONResponse from lightllm.server.core.objs.sampling_params import SamplingParams from lightllm.server.core.objs import StartArgs from .multimodal_params import MultimodalParams @@ -52,7 +52,8 @@ 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, prime_pd_master_streaming_response +from .api_errors import create_error_response, create_server_busy_response +from .api_openai import chat_completions_impl, completions_impl from .api_models import ( ChatCompletionRequest, ChatCompletionResponse, @@ -153,34 +154,19 @@ async def send_wrapper(message): app.add_middleware(_AccessLogMiddleware) -def create_error_response( - status_code: HTTPStatus, message: str, err_type: str = None, param: str = None -) -> JSONResponse: - if err_type is None: - if status_code.value >= 500: - 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" - - g_objs.metric_client.counter_inc("lightllm_request_failure") - return JSONResponse( - {"error": {"message": message, "type": err_type, "param": param, "code": status_code.value}}, - status_code=status_code.value, - ) - - -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)) + + # Streaming responses can raise during their first body iteration, after + # the route handler has already returned. Preserve the Anthropic error + # envelope for that deferred failure path as well. + if request.url.path == "/v1/messages": + from .api_anthropic import _anthropic_error_response + + g_objs.metric_client.counter_inc("lightllm_request_failure") + return _anthropic_error_response(HTTPStatus(exc.status_code), str(exc)) + return create_server_busy_response(exc) @@ -327,8 +313,7 @@ async def generate_stream(request: Request) -> Response: ) try: - response = await g_objs.g_generate_stream_func(request, g_objs.httpserver_manager) - return await prime_pd_master_streaming_response(response) + return await g_objs.g_generate_stream_func(request, g_objs.httpserver_manager) except ServerBusyError as e: logger.warning(str(e)) return create_server_busy_response(e) @@ -385,7 +370,6 @@ 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: @@ -406,7 +390,6 @@ 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: @@ -427,8 +410,7 @@ async def anthropic_messages(raw_request: Request) -> Response: from .api_anthropic import _anthropic_error_response, anthropic_messages_impl try: - response = await anthropic_messages_impl(raw_request) - return await prime_pd_master_streaming_response(response) + return await anthropic_messages_impl(raw_request) except ServerBusyError as e: logger.warning(str(e)) g_objs.metric_client.counter_inc("lightllm_request_failure") @@ -447,8 +429,7 @@ async def openai_responses(raw_request: Request) -> Response: from .api_responses import responses_impl try: - response = await responses_impl(raw_request) - return await prime_pd_master_streaming_response(response) + return await responses_impl(raw_request) except ServerBusyError as e: logger.warning(str(e)) return create_server_busy_response(e) diff --git a/lightllm/server/api_http_rl.py b/lightllm/server/api_http_rl.py index f4ead5031..b73d76170 100644 --- a/lightllm/server/api_http_rl.py +++ b/lightllm/server/api_http_rl.py @@ -7,7 +7,7 @@ → (多数) RlOpReq → Router → Model RlBackendOps 路由在模块级 ``router`` 上注册，由 ``api_http`` ``include_router`` 挂载。 -``g_objs`` / ``create_error_response`` 在 handler 内懒导入，避免与 api_http 循环依赖。 +``g_objs`` 在 handler 内懒导入，避免与 api_http 循环依赖。 """ from http import HTTPStatus @@ -29,14 +29,14 @@ ) from lightllm.utils.log_utils import init_logger +from .api_errors import create_error_response + logger = init_logger(__name__) router = APIRouter() async def handle_request_common(request_obj, handler): - from .api_http import create_error_response - try: ret: RlOpRsp = await handler(request_obj) if ret.success: @@ -51,7 +51,7 @@ async def handle_request_common(request_obj, handler): @router.post("/abort_request") async def abort_request(request: AbortReq, raw_request: Request): """Abort a request.""" - from .api_http import create_error_response, g_objs + from .api_http import g_objs try: success, msg = await g_objs.httpserver_manager.abort_request(request) diff --git a/lightllm/server/api_lightllm.py b/lightllm/server/api_lightllm.py index 6a0abe81b..fe6547d63 100644 --- a/lightllm/server/api_lightllm.py +++ b/lightllm/server/api_lightllm.py @@ -1,10 +1,11 @@ import collections from typing import AsyncGenerator from fastapi import BackgroundTasks, Request -from fastapi.responses import Response, StreamingResponse +from fastapi.responses import Response from lightllm.server.core.objs.sampling_params import SamplingParams from .multimodal_params import MultimodalParams from .httpserver.manager import HttpServerManager +from .api_stream_obj import CustomStreamingResponse import ujson as json @@ -164,6 +165,6 @@ async def stream_results() -> AsyncGenerator[bytes, None]: from .api_openai import _safe_stream_wrapper background_tasks = BackgroundTasks() - return StreamingResponse( + return CustomStreamingResponse( _safe_stream_wrapper(stream_results()), media_type="text/event-stream", background=background_tasks ) diff --git a/lightllm/server/api_openai.py b/lightllm/server/api_openai.py index 413d05c99..e8570369a 100644 --- a/lightllm/server/api_openai.py +++ b/lightllm/server/api_openai.py @@ -23,7 +23,7 @@ from typing import Callable from lightllm.server import TokenLoad from fastapi import BackgroundTasks, FastAPI, Request, WebSocket, WebSocketDisconnect -from fastapi.responses import Response, StreamingResponse, JSONResponse +from fastapi.responses import Response, JSONResponse from lightllm.server.core.objs.sampling_params import SamplingParams from .multimodal_params import MultimodalParams from .httpserver.manager import HttpServerManager @@ -37,6 +37,8 @@ from lightllm.utils.envs_utils import get_unique_server_name from dataclasses import dataclass +from .api_errors import create_error_response +from .api_stream_obj import CustomStreamingResponse from .api_models import ( ChatCompletionRequest, CompletionRequest, @@ -60,46 +62,18 @@ 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 + """Convert generation errors to SSE events after the response has started.""" + first_chunk_sent = False try: async for item in stream_generator: + first_chunk_sent = True 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: + if not first_chunk_sent: raise logger.error("Generation interrupted after the stream started: %s", e.message) error_data = json.dumps( @@ -125,26 +99,6 @@ def _serialize_sse_chunk(chunk, choice_nulls=(), response_nulls=()): return json.dumps(d, ensure_ascii=False) -def create_error_response( - status_code: HTTPStatus, message: str, err_type: str = None, param: str = None -) -> JSONResponse: - from .api_http import g_objs - - if err_type is None: - if status_code.value >= 500: - err_type = "InternalServerError" - elif status_code == HTTPStatus.NOT_FOUND: - err_type = "NotFoundError" - else: - err_type = "BadRequestError" - - g_objs.metric_client.counter_inc("lightllm_request_failure") - return JSONResponse( - {"error": {"message": message, "type": err_type, "param": param, "code": status_code.value}}, - status_code=status_code.value, - ) - - def _process_tool_call_id( tool_call_parser, call_item: ToolCallItem, @@ -412,18 +366,18 @@ async def chat_completions_impl(request: ChatCompletionRequest, raw_request: Req prompt_cache_len_dict[sub_req_id] = metadata.get("prompt_cache_len", 0) choices = [] sub_ids = list(final_output_dict.keys())[: request.n] + prompt_tokens = prompt_tokens_dict[sub_ids[0]] + completion_tokens = sum(count_output_tokens_dict[sub_req_id] for sub_req_id in sub_ids) + cached_tokens = prompt_cache_len_dict.get(sub_ids[0], 0) + usage = UsageInfo( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + prompt_tokens_details=PromptTokensDetails(cached_tokens=cached_tokens), + ) + for i in range(request.n): sub_req_id = sub_ids[i] - prompt_tokens = prompt_tokens_dict[sub_req_id] - completion_tokens = count_output_tokens_dict[sub_req_id] - cached_tokens = prompt_cache_len_dict.get(sub_req_id, 0) - usage = UsageInfo( - prompt_tokens=prompt_tokens, - completion_tokens=completion_tokens, - total_tokens=prompt_tokens + completion_tokens, - prompt_tokens_details=PromptTokensDetails(cached_tokens=cached_tokens), - ) - finish_reason = finish_reason_dict[sub_req_id] text = "".join(final_output_dict[sub_req_id]) @@ -802,7 +756,7 @@ async def stream_results() -> AsyncGenerator[bytes, None]: yield "data: [DONE]\n\n".encode("utf-8") background_tasks = BackgroundTasks() - return StreamingResponse( + return CustomStreamingResponse( _safe_stream_wrapper(stream_results()), media_type="text/event-stream", background=background_tasks ) @@ -913,7 +867,7 @@ async def _process_prompts_completion( prompts[0], sampling_params, multimodal_params, raw_request, request, created_time ) - async def process_single_prompt(prompt: Union[str, List[int]], prompt_index: int): + async def process_single_prompt(prompt: Union[str, List[int]]): if len(prompts) > 1: individual_sampling_params = SamplingParams() individual_sampling_params.init(tokenizer=g_objs.httpserver_manager.tokenizer, **sampling_params_dict) @@ -930,14 +884,12 @@ async def process_single_prompt(prompt: Union[str, List[int]], prompt_index: int prompt, individual_sampling_params, multimodal_params, request=raw_request ) - return await _collect_generation_results( - generator, request, prompt_str, prompt_index, individual_sampling_params - ) + return await _collect_generation_results(generator, request, prompt_str, individual_sampling_params) - tasks = [asyncio.create_task(process_single_prompt(prompt, i)) for i, prompt in enumerate(prompts)] + tasks = [asyncio.create_task(process_single_prompt(prompt)) for prompt in prompts] - results = await asyncio.gather(*tasks) - return _build_completion_response(results, request, created_time, len(prompts) > 1) + results_by_prompt = await asyncio.gather(*tasks) + return _build_completion_response(results_by_prompt, request, created_time, len(prompts) > 1) async def _handle_streaming_completion( @@ -1010,72 +962,82 @@ async def stream_results() -> AsyncGenerator[bytes, None]: yield "data: [DONE]\n\n" background_tasks = BackgroundTasks() - return StreamingResponse( + return CustomStreamingResponse( _safe_stream_wrapper(stream_results()), media_type="text/event-stream", background=background_tasks ) async def _collect_generation_results( - generator, request: CompletionRequest, prompt: str, prompt_index: int, sampling_params: SamplingParams + generator, request: CompletionRequest, prompt: str, sampling_params: SamplingParams ): - final_output = [] - count_output_tokens = 0 - finish_reason = None - prompt_tokens = 0 - prompt_cache_len = 0 - token_infos = [] if request.logprobs is not None else None - prompt_logprobs = None - prompt_token_ids = None - is_first_metadata = True + final_outputs = collections.defaultdict(list) + output_token_counts = collections.defaultdict(int) + finish_reasons = {} + prompt_tokens = {} + prompt_cache_lens = {} + token_infos = collections.defaultdict(list) + prompt_logprobs = {} + prompt_token_ids = {} async for sub_req_id, request_output, metadata, finish_status in generator: - if is_first_metadata: - prompt_logprobs = metadata.get("prompt_logprobs", None) - prompt_token_ids = metadata.get("prompt_token_ids", None) - is_first_metadata = False - - count_output_tokens += 1 - final_output.append(request_output) - - if request.logprobs is not None and token_infos is not None: - token_info = { - "text": request_output, - "logprob": metadata.get("logprob", None), - "id": metadata.get("id", None), - } - token_infos.append(token_info) + if sub_req_id not in prompt_token_ids: + prompt_logprobs[sub_req_id] = metadata.get("prompt_logprobs") + prompt_token_ids[sub_req_id] = metadata.get("prompt_token_ids") + + output_token_counts[sub_req_id] += 1 + final_outputs[sub_req_id].append(request_output) + + if request.logprobs is not None: + token_infos[sub_req_id].append( + { + "text": request_output, + "logprob": metadata.get("logprob"), + "id": metadata.get("id"), + } + ) if finish_status.is_finished(): - finish_reason = finish_status.get_finish_reason() - prompt_tokens = metadata["prompt_tokens"] - prompt_cache_len = metadata.get("prompt_cache_len", 0) - - # 处理停止序列剔除 - final_text = "".join(final_output) - if finish_reason == "stop" and sampling_params.stop_sequences.size > 0: - valid_stop_strings = sampling_params.stop_sequences.to_strings() - for stop_str in valid_stop_strings: - stop_index = final_text.rfind(stop_str, max(0, len(final_text) - len(stop_str) - 20), len(final_text)) - if stop_index != -1: - logger.debug(f"removed stop sequence in tail: '{final_text[stop_index:]}'") - final_text = final_text[:stop_index] - break + finish_reasons[sub_req_id] = finish_status.get_finish_reason() + prompt_tokens[sub_req_id] = metadata["prompt_tokens"] + prompt_cache_lens[sub_req_id] = metadata.get("prompt_cache_len", 0) + + results = [] + for sub_req_id in sorted(final_outputs)[: request.n]: + final_text = "".join(final_outputs[sub_req_id]) + finish_reason = finish_reasons.get(sub_req_id) + + if finish_reason == "stop" and sampling_params.stop_sequences.size > 0: + for stop_str in sampling_params.stop_sequences.to_strings(): + stop_index = final_text.rfind( + stop_str, + max(0, len(final_text) - len(stop_str) - 20), + len(final_text), + ) + if stop_index != -1: + logger.debug("removed stop sequence in tail: '%s'", final_text[stop_index:]) + final_text = final_text[:stop_index] + break + + results.append( + { + "text": final_text, + "finish_reason": finish_reason, + "prompt_tokens": prompt_tokens.get(sub_req_id, 0), + "prompt_cache_len": prompt_cache_lens.get(sub_req_id, 0), + "completion_tokens": output_token_counts[sub_req_id], + "token_infos": (token_infos[sub_req_id] if request.logprobs is not None else None), + "prompt_logprobs": prompt_logprobs[sub_req_id], + "prompt_token_ids": prompt_token_ids[sub_req_id], + "prompt_text": prompt, + } + ) - return { - "index": prompt_index, - "text": final_text, - "finish_reason": finish_reason, - "prompt_tokens": prompt_tokens, - "prompt_cache_len": prompt_cache_len, - "completion_tokens": count_output_tokens, - "token_infos": token_infos, - "prompt_logprobs": prompt_logprobs, - "prompt_token_ids": prompt_token_ids, - "prompt_text": prompt, - } + return results -def _build_completion_response(results: List[Dict], request: CompletionRequest, created_time: int, is_batch: bool): +def _build_completion_response( + results_by_prompt: List[List[Dict]], request: CompletionRequest, created_time: int, is_batch: bool +): from .api_http import g_objs choices = [] @@ -1083,24 +1045,28 @@ def _build_completion_response(results: List[Dict], request: CompletionRequest, total_completion_tokens = 0 total_cached_tokens = 0 - for result in results: - text = result["text"] - if request.echo: - text = result["prompt_text"] + text - - logprobs_data = _build_logprobs_data(result, request, g_objs.httpserver_manager.tokenizer) - - choice = CompletionChoice( - index=result["index"], - text=text, - finish_reason=result["finish_reason"], - logprobs=CompletionLogprobs(**logprobs_data) if logprobs_data else None, - ) - choices.append(choice) - - total_prompt_tokens += result["prompt_tokens"] - total_completion_tokens += result["completion_tokens"] - total_cached_tokens += result.get("prompt_cache_len", 0) + for prompt_results in results_by_prompt: + if not prompt_results: + continue + + total_prompt_tokens += prompt_results[0]["prompt_tokens"] + total_cached_tokens += prompt_results[0].get("prompt_cache_len", 0) + + for result in prompt_results: + text = result["text"] + if request.echo: + text = result["prompt_text"] + text + + logprobs_data = _build_logprobs_data(result, request, g_objs.httpserver_manager.tokenizer) + choices.append( + CompletionChoice( + index=len(choices), + text=text, + finish_reason=result["finish_reason"], + logprobs=CompletionLogprobs(**logprobs_data) if logprobs_data else None, + ) + ) + total_completion_tokens += result["completion_tokens"] usage = UsageInfo( prompt_tokens=total_prompt_tokens, diff --git a/lightllm/server/api_responses.py b/lightllm/server/api_responses.py index 0cad4e63d..cad3f4fe3 100644 --- a/lightllm/server/api_responses.py +++ b/lightllm/server/api_responses.py @@ -7,10 +7,13 @@ from typing import Any, AsyncGenerator, Dict, List, Optional from fastapi import Request -from fastapi.responses import JSONResponse, Response, StreamingResponse +from fastapi.responses import JSONResponse, Response from lightllm.utils.log_utils import init_logger +from .api_errors import create_error_response, is_rate_limit_error +from .api_stream_obj import CustomStreamingResponse + logger = init_logger(__name__) @@ -528,11 +531,7 @@ def open_item(kind: str, item: Dict[str, Any]): if failed_error is not None: response["status"] = "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" + error_code = "rate_limit_error" if is_rate_limit_error(failed_error) else "server_error" response["error"] = {"code": error_code, "message": failed_error.get("message", "generation failed")} yield event("response.failed", {"response": response}) return @@ -552,7 +551,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, prime_pd_master_streaming_response + from .api_openai import chat_completions_impl try: body = await raw_request.json() @@ -585,10 +584,9 @@ async def responses_impl(raw_request: Request) -> Response: downstream = await chat_completions_impl(chat_request, raw_request) if chat_request.stream: - if not isinstance(downstream, StreamingResponse): + if not isinstance(downstream, CustomStreamingResponse): return downstream - downstream = await prime_pd_master_streaming_response(downstream) - return StreamingResponse( + return CustomStreamingResponse( _openai_sse_to_responses_events(downstream.body_iterator, body), media_type="text/event-stream", ) diff --git a/lightllm/server/api_stream_obj.py b/lightllm/server/api_stream_obj.py new file mode 100644 index 000000000..f4765a53f --- /dev/null +++ b/lightllm/server/api_stream_obj.py @@ -0,0 +1,83 @@ +"""Custom streaming response behavior for PD-master request admission. + +Starlette's ``StreamingResponse`` sends ``http.response.start`` before it +starts iterating over the response body. This is normally desirable because +the client receives the HTTP status and headers immediately. However, the +generation body used by LightLLM is an async generator, so its code does not +run when the ``StreamingResponse`` object is created. In PD-master mode, node +selection and admission waiting happen only when that generator is iterated. +Consequently, a ``ServerBusyError`` may be raised during the first iteration, +after Starlette has already sent HTTP 200. + +An HTTP status cannot be changed after ``http.response.start`` has been sent. +This module therefore delays that event in PD-master mode until the body has +produced its first chunk. If admission fails before then, no response has +started and FastAPI's exception handler can still return HTTP 429. Once the +first chunk exists, the request has passed this initial admission point and +the response starts normally. + +Only PD-master mode uses the delayed behavior. Other modes retain Starlette's +original implementation so clients and proxies receive headers immediately +and keep their existing response-header timeout semantics. +""" + +from fastapi.responses import StreamingResponse +from starlette.types import Send + +from lightllm.utils.envs_utils import get_env_start_args + + +class CustomStreamingResponse(StreamingResponse): + """Send the PD-master HTTP status only after the first body chunk is ready. + + The first chunk is sent directly after the response headers; it is not + discarded or replayed through another iterator. Empty streams still send + the configured status followed by an empty final body. + + This mechanism only covers failures raised before the first chunk. After + the response has started, HTTP no longer allows changing its status code; + later generation failures must be reported in the stream body (for + example, as an SSE error event) or by closing the connection. + """ + + async def stream_response(self, send: Send) -> None: + # Preserve Starlette's immediate response-start behavior outside + # PD-master mode. Delaying every streaming response would make normal + # first-token latency count against response-header timeouts. + if get_env_start_args().run_mode != "pd_master": + await super().stream_response(send) + return + + async def send_chunk(chunk): + if not isinstance(chunk, (bytes, memoryview)): + chunk = chunk.encode(self.charset) + await send({"type": "http.response.body", "body": chunk, "more_body": True}) + + async def send_response_start(): + # Read status and headers at send time. The first body iteration + # may update them while performing admission or request setup. + await send( + { + "type": "http.response.start", + "status": self.status_code, + "headers": self.raw_headers, + } + ) + + # Iterating here starts the otherwise-lazy generation pipeline. A + # ServerBusyError raised before the first yield escapes without an + # http.response.start event, allowing FastAPI to produce HTTP 429. + async for first_chunk in self.body_iterator: + await send_response_start() + await send_chunk(first_chunk) + break + else: + # An empty iterator is still a valid response and needs headers. + await send_response_start() + + # The first chunk has already been sent; continue from the same + # iterator without restarting or duplicating it. + async for chunk in self.body_iterator: + await send_chunk(chunk) + + await send({"type": "http.response.body", "body": b"", "more_body": False}) diff --git a/lightllm/server/api_tgi.py b/lightllm/server/api_tgi.py index 2c6962323..e7118e5e2 100755 --- a/lightllm/server/api_tgi.py +++ b/lightllm/server/api_tgi.py @@ -2,11 +2,12 @@ import collections from typing import AsyncGenerator from fastapi import BackgroundTasks, Request -from fastapi.responses import Response, StreamingResponse, JSONResponse +from fastapi.responses import Response, JSONResponse from fastapi.encoders import jsonable_encoder from lightllm.server.core.objs.sampling_params import SamplingParams from .multimodal_params import MultimodalParams from .httpserver.manager import HttpServerManager +from .api_stream_obj import CustomStreamingResponse import json RETURN_LIST = os.getenv("RETURN_LIST", "FALSE").upper() in ["ON", "TRUE", "1"] @@ -191,6 +192,6 @@ async def stream_results() -> AsyncGenerator[bytes, None]: from .api_openai import _safe_stream_wrapper background_tasks = BackgroundTasks() - return StreamingResponse( + return CustomStreamingResponse( _safe_stream_wrapper(stream_results()), media_type="text/event-stream", background=background_tasks ) diff --git a/test/test_api/test_completion_multiple_choices.py b/test/test_api/test_completion_multiple_choices.py new file mode 100644 index 000000000..98b897506 --- /dev/null +++ b/test/test_api/test_completion_multiple_choices.py @@ -0,0 +1,130 @@ +import asyncio +from types import SimpleNamespace + +from lightllm.server import api_openai +from lightllm.server.api_models import ChatCompletionRequest, CompletionRequest +from lightllm.server.api_openai import ( + _build_completion_response, + _collect_generation_results, +) + + +class _FinishStatus: + def __init__(self, finished=False, reason=None): + self.finished = finished + self.reason = reason + + def is_finished(self): + return self.finished + + def get_finish_reason(self): + return self.reason + + +def test_collect_generation_results_keeps_choices_separate(monkeypatch): + async def generate_results(): + metadata = { + "prompt_tokens": 4, + "prompt_cache_len": 1, + "prompt_token_ids": [1, 2, 3, 4], + } + yield 82, "C", {**metadata, "logprob": -0.5, "id": 14}, _FinishStatus() + yield 81, "B", {**metadata, "logprob": -0.2, "id": 11}, _FinishStatus() + yield 80, "A", {**metadata, "logprob": -0.1, "id": 10}, _FinishStatus() + yield 82, "3", {**metadata, "logprob": -0.6, "id": 15}, _FinishStatus(True, "length") + yield 81, "2", {**metadata, "logprob": -0.4, "id": 13}, _FinishStatus(True, "length") + yield 80, "1", {**metadata, "logprob": -0.3, "id": 12}, _FinishStatus(True, "length") + + request = CompletionRequest( + model="test-model", + prompt="Prompt", + n=3, + best_of=3, + max_tokens=2, + logprobs=1, + ) + sampling_params = SimpleNamespace(stop_sequences=SimpleNamespace(size=0)) + + results = asyncio.run(_collect_generation_results(generate_results(), request, "Prompt", sampling_params)) + + assert [result["text"] for result in results] == ["A1", "B2", "C3"] + assert [result["completion_tokens"] for result in results] == [2, 2, 2] + assert [[token["id"] for token in result["token_infos"]] for result in results] == [ + [10, 12], + [11, 13], + [14, 15], + ] + + from lightllm.server.api_http import g_objs + + monkeypatch.setattr(g_objs, "httpserver_manager", SimpleNamespace(tokenizer=None), raising=False) + response = _build_completion_response([results], request, created_time=123, is_batch=False) + + assert [choice.index for choice in response.choices] == [0, 1, 2] + assert [choice.text for choice in response.choices] == ["A1", "B2", "C3"] + assert response.usage.prompt_tokens == 4 + assert response.usage.completion_tokens == 6 + assert response.usage.total_tokens == 10 + assert response.usage.prompt_tokens_details.cached_tokens == 1 + + second_prompt_results = [ + { + **result, + "prompt_tokens": 5, + "prompt_cache_len": 2, + "prompt_text": "Another prompt", + } + for result in results + ] + batch_response = _build_completion_response( + [results, second_prompt_results], request, created_time=123, is_batch=True + ) + + assert [choice.index for choice in batch_response.choices] == list(range(6)) + assert batch_response.usage.prompt_tokens == 9 + assert batch_response.usage.completion_tokens == 12 + assert batch_response.usage.total_tokens == 21 + assert batch_response.usage.prompt_tokens_details.cached_tokens == 3 + + +def test_non_streaming_chat_usage_sums_all_choices(monkeypatch): + async def generate_results(): + metadata = {"prompt_tokens": 4, "prompt_cache_len": 1} + for sub_req_id, text in [(80, "A"), (81, "B"), (82, "C")]: + yield sub_req_id, text, metadata, _FinishStatus() + for sub_req_id, text in [(80, "1"), (81, "2"), (82, "3")]: + yield sub_req_id, text, metadata, _FinishStatus(True, "length") + + class _SamplingParams: + def init(self, **_kwargs): + pass + + def verify(self): + pass + + async def build_prompt(_request, _tools): + return "Prompt" + + manager = SimpleNamespace(tokenizer=None, generate=lambda *_args, **_kwargs: generate_results()) + monkeypatch.setattr(api_openai, "SamplingParams", _SamplingParams) + monkeypatch.setattr(api_openai, "build_prompt", build_prompt) + monkeypatch.setattr(api_openai, "get_env_start_args", lambda: SimpleNamespace(reasoning_parser=None)) + + from lightllm.server.api_http import g_objs + + monkeypatch.setattr(g_objs, "httpserver_manager", manager, raising=False) + request = ChatCompletionRequest( + model="test-model", + messages=[{"role": "user", "content": "Hello"}], + n=3, + max_tokens=2, + ) + + response = asyncio.run(api_openai.chat_completions_impl(request, SimpleNamespace())) + + assert [choice.index for choice in response.choices] == [0, 1, 2] + assert [choice.message.content for choice in response.choices] == ["A1", "B2", "C3"] + assert response.usage.prompt_tokens == 4 + assert response.usage.completion_tokens == 6 + assert response.usage.total_tokens == 10 + assert response.usage.prompt_tokens_details.cached_tokens == 1 diff --git a/test/test_api/test_server_busy_handling.py b/test/test_api/test_server_busy_handling.py new file mode 100644 index 000000000..4650919a0 --- /dev/null +++ b/test/test_api/test_server_busy_handling.py @@ -0,0 +1,194 @@ +import asyncio +import json +from types import SimpleNamespace + +import httpx +import pytest +from fastapi import FastAPI, Request +from fastapi.responses import JSONResponse + +from lightllm.server import api_anthropic, api_errors, api_http, api_openai, api_stream_obj +from lightllm.server.api_stream_obj import CustomStreamingResponse +from lightllm.utils.error_utils import ServerBusyError + + +class _MetricClient: + def __init__(self): + self.counters = [] + + def counter_inc(self, name): + self.counters.append(name) + + +def test_api_modules_share_error_response_factory(): + assert api_http.create_error_response is api_errors.create_error_response + assert api_openai.create_error_response is api_errors.create_error_response + + +def test_server_busy_response_is_rate_limited(monkeypatch): + metric_client = _MetricClient() + monkeypatch.setattr(api_http.g_objs, "metric_client", metric_client) + + response = api_errors.create_server_busy_response(ServerBusyError()) + body = json.loads(response.body) + + assert response.status_code == 429 + assert body["error"]["type"] == "RateLimitError" + assert body["error"]["code"] == 429 + assert metric_client.counters == ["lightllm_request_failure"] + + +@pytest.mark.parametrize( + ("error", "expected"), + [ + ({"code": 429, "type": "server_error"}, "rate_limit_error"), + ({"type": "RateLimitError"}, "rate_limit_error"), + ({"type": "invalid_request_error"}, "invalid_request_error"), + ({"type": "server_error"}, "api_error"), + ], +) +def test_anthropic_error_type(error, expected): + assert api_anthropic._anthropic_error_type(error) == expected + + +def test_pd_master_stream_starts_response_after_first_chunk(monkeypatch): + monkeypatch.setattr(api_stream_obj, "get_env_start_args", lambda: SimpleNamespace(run_mode="pd_master")) + + async def run(): + response = None + + async def generate(): + response.status_code = 201 + yield "first" + yield "second" + + messages = [] + + async def send(message): + messages.append(message) + + response = CustomStreamingResponse(generate()) + await response.stream_response(send) + return messages + + messages = asyncio.run(run()) + assert [message["type"] for message in messages] == [ + "http.response.start", + "http.response.body", + "http.response.body", + "http.response.body", + ] + assert [message.get("body") for message in messages[1:]] == [b"first", b"second", b""] + assert messages[0]["status"] == 201 + + +def test_pd_master_stream_propagates_busy_error_before_response_start(monkeypatch): + monkeypatch.setattr(api_stream_obj, "get_env_start_args", lambda: SimpleNamespace(run_mode="pd_master")) + + async def run(): + async def generate(): + if False: + yield + raise ServerBusyError() + + messages = [] + + async def send(message): + messages.append(message) + + with pytest.raises(ServerBusyError): + await CustomStreamingResponse(generate()).stream_response(send) + + return messages + + assert asyncio.run(run()) == [] + + +def test_pd_master_stream_can_return_http_429(monkeypatch): + monkeypatch.setattr(api_stream_obj, "get_env_start_args", lambda: SimpleNamespace(run_mode="pd_master")) + app = FastAPI() + + @app.exception_handler(ServerBusyError) + async def handle_server_busy(_request: Request, _error: ServerBusyError): + return JSONResponse({"error": "server busy"}, status_code=429) + + @app.get("/") + async def generate_stream(): + async def generate(): + if False: + yield + raise ServerBusyError() + + return CustomStreamingResponse(generate()) + + async def run(): + transport = httpx.ASGITransport(app=app, raise_app_exceptions=False) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: + return await client.get("/") + + response = asyncio.run(run()) + assert response.status_code == 429 + assert response.json() == {"error": "server busy"} + + +def test_pd_master_anthropic_stream_preserves_error_envelope(monkeypatch): + metric_client = _MetricClient() + monkeypatch.setattr(api_http.g_objs, "metric_client", metric_client) + monkeypatch.setattr(api_http, "get_env_start_args", lambda: SimpleNamespace(run_mode="normal")) + monkeypatch.setattr(api_stream_obj, "get_env_start_args", lambda: SimpleNamespace(run_mode="pd_master")) + + async def anthropic_messages_impl(_request): + async def generate(): + if False: + yield + raise ServerBusyError() + + return CustomStreamingResponse(generate(), media_type="text/event-stream") + + monkeypatch.setattr(api_anthropic, "anthropic_messages_impl", anthropic_messages_impl) + + async def run(): + transport = httpx.ASGITransport(app=api_http.app, raise_app_exceptions=False) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: + return await client.post("/v1/messages", json={}) + + response = asyncio.run(run()) + assert response.status_code == 429 + assert response.json() == { + "type": "error", + "error": { + "type": "rate_limit_error", + "message": "Server is busy, please try again later (Status code: 429)", + }, + } + assert metric_client.counters == ["lightllm_request_failure"] + + +def test_safe_stream_reports_busy_error_after_first_chunk(): + async def run(): + async def generate(): + yield "first" + raise ServerBusyError() + + return [item async for item in api_openai._safe_stream_wrapper(generate())] + + chunks = asyncio.run(run()) + error = json.loads(chunks[1].removeprefix("data: ")) + + assert chunks[0] == "first" + assert error["error"]["type"] == "server_error" + assert error["error"]["code"] == "stream_error" + + +def test_safe_stream_propagates_busy_error_before_first_chunk(): + async def run(): + async def generate(): + if False: + yield + raise ServerBusyError() + + with pytest.raises(ServerBusyError): + async for _ in api_openai._safe_stream_wrapper(generate()): + pass + + asyncio.run(run())