[add]support req-disconnected

Author: sangchengmeng

lightllm/server/api_http.py

@@ -368,7 +368,7 @@ async def tokens(request: Request):
 
         multimodal_params_dict = request_dict.get("multimodal_params", {})
         multimodal_params = MultimodalParams(**multimodal_params_dict)
-        await multimodal_params.verify_and_preload()
+        await multimodal_params.verify_and_preload(request)
         return JSONResponse(
             {
                 "ntokens": g_objs.httpserver_manager.tokens(

lightllm/server/httpserver/manager.py

@@ -228,7 +228,7 @@ async def generate(
                 original_multimodal_params = copy.deepcopy(multimodal_params)
 
             if self.pd_mode.is_P_or_NORMAL():
-                await multimodal_params.verify_and_preload()
+                await multimodal_params.verify_and_preload(request)
 
             # 记录请求到达的相关信息
             await self._log_req_header(request_headers, group_request_id)

lightllm/server/multimodal_params.py

@@ -6,6 +6,7 @@
 from PIL import Image
 from lightllm.utils.image_utils import fetch_image
 import base64
+from fastapi import Request
 
 
 class ImageItem:
@@ -24,11 +25,11 @@ def __init__(self, **kwargs):
         self._preload_data = None
         self.extra_params = {}
 
-    async def preload(self):
+    async def preload(self, request: Request):
         try:
             if self._type == "url":
                 timeout = int(os.getenv("REQUEST_TIMEOUT", "5"))
-                img_data = await fetch_image(self._data, timeout=timeout)
+                img_data = await fetch_image(self._data, request, timeout=timeout)
             elif self._type == "base64":
                 img_data = base64.b64decode(self._data)
             elif self._type == "image_size":
@@ -81,9 +82,9 @@ def __init__(
         self.images = [ImageItem(**i) for i in images]
         return
 
-    async def verify_and_preload(self):
+    async def verify_and_preload(self, request: Request):
         for image in self.images:
-            await image.preload()
+            await image.preload(request)
         return
 
     def to_dict(self):

lightllm/utils/image_utils.py

@@ -1,7 +1,13 @@
+import time
 import base64
 import httpx
+import logging
 from PIL import Image
 from io import BytesIO
+from fastapi import Request
+from lightllm.utils.log_utils import init_logger
+
+logger = init_logger(__name__)
 
 
 def image2base64(img_str: str):
@@ -13,17 +19,23 @@ def image2base64(img_str: str):
     return base64.b64encode(buffer.getvalue()).decode("utf-8")
 
 
-async def fetch_image(url, timeout):
+async def fetch_image(url, request: Request, timeout):
+    logger.info(f"Begin to download image from url: {url}")
+    start_time = time.time()
     async with httpx.AsyncClient() as client:
         async with client.stream("GET", url, timeout=timeout) as response:
             response.raise_for_status()
             ans_bytes = []
-
             async for chunk in response.aiter_bytes(chunk_size=1024 * 1024):
+                if await request.is_disconnected():
+                    await response.aclose()
+                    raise Exception("Request disconnected. User cancelled download.")
                 ans_bytes.append(chunk)
                 # 接收的数据不能大于128M
                 if len(ans_bytes) > 128:
-                    raise Exception("image data is too big")
+                    raise Exception("Image data is too big")
 
             content = b"".join(ans_bytes)
-            return content
+    end_time = time.time()
+    logger.info("Download image time: {:.2f} seconds".format(end_time - start_time))
+    return content