pd master health, tokens and server busy error (#895)

Author: shihaobai

lightllm/server/api_http.py

@@ -41,9 +41,10 @@
 from .multimodal_params import MultimodalParams
 from .httpserver.manager import HttpServerManager
 from .httpserver_for_pd_master.manager import HttpServerManagerForPDMaster
-from .api_lightllm import lightllm_get_score, lightllm_pd_generate_stream
+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.log_utils import init_logger
+from lightllm.utils.error_utils import ServerBusyError
 from lightllm.server.metrics.manager import MetricClient
 from lightllm.utils.envs_utils import get_unique_server_name
 from dataclasses import dataclass
@@ -136,6 +137,9 @@ def get_model_name():
 @app.get("/health", summary="Check server health")
 @app.head("/health", summary="Check server health")
 async def healthcheck(request: Request):
+    if g_objs.args.run_mode == "pd_master":
+        return JSONResponse({"message": "Ok"}, status_code=200)
+
     if os.environ.get("DEBUG_HEALTHCHECK_RETURN_FAIL") == "true":
         return JSONResponse({"message": "Error"}, status_code=503)
     from lightllm.utils.health_check import health_check, health_obj
@@ -175,6 +179,9 @@ async def token_load(request: Request):
 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))
     except Exception as e:
         logger.error("An error occurred: %s", str(e), exc_info=True)
         return create_error_response(HTTPStatus.EXPECTATION_FAILED, str(e))
@@ -184,15 +191,9 @@ async def generate(request: Request) -> Response:
 async def generate_stream(request: Request) -> Response:
     try:
         return await g_objs.g_generate_stream_func(request, g_objs.httpserver_manager)
-    except Exception as e:
-        logger.error("An error occurred: %s", str(e), exc_info=True)
-        return create_error_response(HTTPStatus.EXPECTATION_FAILED, str(e))
-
-
-@app.post("/pd_generate_stream")
-async def pd_generate_stream(request: Request) -> Response:
-    try:
-        return await lightllm_pd_generate_stream(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))
     except Exception as e:
         logger.error("An error occurred: %s", str(e), exc_info=True)
         return create_error_response(HTTPStatus.EXPECTATION_FAILED, str(e))

lightllm/server/api_lightllm.py

@@ -136,29 +136,3 @@ async def stream_results() -> AsyncGenerator[bytes, None]:
 
     background_tasks = BackgroundTasks()
     return StreamingResponse(stream_results(), media_type="text/event-stream", background=background_tasks)
-
-
-async def lightllm_pd_generate_stream(request: Request, httpserver_manager: HttpServerManager) -> Response:
-
-    request_dict = await request.json()
-    prompt = request_dict.pop("inputs")
-    sample_params_dict = request_dict["parameters"]
-    _ = sample_params_dict.pop("return_details", False)
-    sampling_params = SamplingParams()
-    sampling_params.init(tokenizer=httpserver_manager.tokenizer, **sample_params_dict)
-    sampling_params.verify()
-    if sampling_params.best_of != 1:
-        raise Exception("stream api only support best_of == 1")
-
-    multimodal_params_dict = request_dict.get("multimodal_params", {})
-    multimodal_params = MultimodalParams(**multimodal_params_dict)
-    results_generator = httpserver_manager.generate(prompt, sampling_params, multimodal_params, request=request)
-
-    # Streaming case
-    async def stream_results() -> AsyncGenerator[bytes, None]:
-        async for sub_req_id, request_output, metadata, finish_status in results_generator:
-            ret = [sub_req_id, request_output, metadata, finish_status.value]
-            yield ("data:" + json.dumps(ret, ensure_ascii=False) + "\n\n").encode("utf-8")
-
-    background_tasks = BackgroundTasks()
-    return StreamingResponse(stream_results(), media_type="text/event-stream", background=background_tasks)

lightllm/server/api_openai.py

@@ -26,7 +26,7 @@
 from .multimodal_params import MultimodalParams
 from .httpserver.manager import HttpServerManager
 from .httpserver_for_pd_master.manager import HttpServerManagerForPDMaster
-from .api_lightllm import lightllm_get_score, lightllm_pd_generate_stream
+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.log_utils import init_logger

lightllm/server/httpserver_for_pd_master/manager.py

@@ -24,6 +24,7 @@
 from lightllm.server.metrics.manager import MetricClient
 from lightllm.utils.statics_utils import MovingAverage
 from lightllm.server.httpserver.manager import AsyncQueue
+from lightllm.utils.error_utils import ServerBusyError
 
 logger = init_logger(__name__)
 
@@ -87,9 +88,22 @@ async def update_req_status(self, upkv_status: UpKVStatus):
             pass
         return
 
-    def tokens(self, prompt: str):
-        # to do
-        raise NotImplementedError("tokens is not implements")
+    def tokens(self, prompt, multimodal_params, samping_params: SamplingParams, kwargs=None):
+        kwargs = {} if kwargs is None else kwargs
+        prompt_ids = self.tokenizer.encode(prompt, None, **kwargs)
+        image_tokens = 0
+        img_count = 0
+        audio_tokens = 0
+        audio_count = 0
+        for img in multimodal_params.images:
+            img_count += 1
+            self.tokenizer.init_imageitem_extral_params(img, multimodal_params, samping_params)
+            image_tokens += self.tokenizer.get_image_token_length(img)
+        for audio in multimodal_params.audios:
+            audio_count += 1
+            self.tokenizer.init_audioitem_extral_params(audio, multimodal_params, samping_params)
+            audio_tokens += self.tokenizer.get_audio_token_length(audio)
+        return len(prompt_ids) + image_tokens + img_count + audio_tokens + audio_count
 
     async def select_p_d_node(
         self, prompt: Union[str, List[int]], sampling_params: SamplingParams, multimodal_params: MultimodalParams
@@ -219,8 +233,8 @@ async def fetch_stream(
         try:
             await asyncio.wait_for(up_status_event.wait(), timeout=60)
         except asyncio.TimeoutError:
-            logger.warning(f"group_request_id: {group_request_id} kv move time out err")
-            assert False, f"req_id {group_request_id} kv move time out, server is busy"
+            logger.warning(f"group_request_id: {group_request_id} kv move time out err, server is busy now.")
+            raise ServerBusyError()
 
         sampling_params.move_kv_to_decode_node.initialize(None)
         sampling_params.max_new_tokens = old_max_new_tokens - 1

lightllm/utils/error_utils.py

@@ -0,0 +1,18 @@
+class ServerBusyError(Exception):
+    """Custom exception for server busy/overload situations"""
+
+    def __init__(self, message="Server is busy, please try again later", status_code=503):
+        """
+        Initialize the ServerBusyError
+
+        Args:
+            message (str): Error message to display
+            status_code (int): HTTP status code (default 503 Service Unavailable)
+        """
+        super().__init__(message)
+        self.message = message
+        self.status_code = status_code  # HTTP 503 Service Unavailable
+
+    def __str__(self):
+        """String representation of the error"""
+        return f"{self.message} (Status code: {self.status_code})"