修复信号量 未释放bug

Author: shell-nlp

gpt_server/model_worker/base/model_worker_base.py

@@ -203,7 +203,7 @@ def get_worker(
         controller_addr: str = "http://localhost:21001",
         worker_id: str = str(uuid.uuid4())[:8],
         model_names: List[str] = [""],
-        limit_worker_concurrency: int = 10000,
+        limit_worker_concurrency: int = 512,
         conv_template: str = None,  # type: ignore
     ):
         worker = cls(
@@ -318,14 +318,9 @@ def acquire_worker_semaphore():
 
 
 def create_background_tasks(request_id):
-    async def abort_request() -> None:
-        await worker.backend.engine.abort(request_id)
-
     background_tasks = BackgroundTasks()
     background_tasks.add_task(release_worker_semaphore)
-    #
-    if os.getenv("backend") == "vllm":
-        background_tasks.add_task(abort_request)
+
     return background_tasks
 
 
@@ -348,8 +343,8 @@ async def api_generate_stream(request: Request):
     params.pop("prompt")
     logger.debug(f"params {params}")
     generator = worker.generate_stream_gate(params)
-    # background_tasks = create_background_tasks(request_id)
-    return StreamingResponse(generator, background=None)
+    background_tasks = create_background_tasks(request_id)
+    return StreamingResponse(generator, background=background_tasks)
 
 
 @app.post("/worker_generate_voice_stream")
@@ -361,7 +356,7 @@ async def api_generate_stream(request: Request):
     params["request"] = request
     logger.debug(f"params {params}")
     generator = worker.generate_voice_stream(params)
-    # background_tasks = create_background_tasks(request_id)
+    background_tasks = create_background_tasks(request_id)
     response_format = params["response_format"]
     content_type = {
         "mp3": "audio/mpeg",
@@ -373,7 +368,7 @@ async def api_generate_stream(request: Request):
     }.get(response_format, f"audio/{response_format}")
     return StreamingResponse(
         generator,
-        background=None,
+        background=background_tasks,
         media_type=content_type,
         headers={
             "Content-Disposition": f"attachment; filename=speech.{response_format}",
@@ -395,8 +390,7 @@ async def api_generate(request: Request):
     logger.debug(f"params {params}")
     output = await worker.generate_gate(params)
     release_worker_semaphore()
-    # if os.getenv("backend") == "vllm":
-    #     await worker.backend.engine.abort(request_id)
+
     return JSONResponse(output)
 
 

gpt_server/serving/openai_api_server.py

@@ -508,7 +508,8 @@ async def chat_completion_stream_generator(
             try:
                 error_code = content["error_code"]
             except Exception as e:
-                print(content)
+                logger.exception(f"发生异常 content：{content}")
+                content["error_code"] = ErrorCode.INTERNAL_ERROR
             if content["error_code"] != 0:
                 yield f"data: {json.dumps(content, ensure_ascii=False)}\n\n"
                 yield "data: [DONE]\n\n"
@@ -678,16 +679,14 @@ async def generate_completion_stream_generator(
 
 
 async def generate_completion_stream(payload: Dict[str, Any], worker_addr: str):
-    async with httpx.AsyncClient(
-        limits=httpx.Limits(max_connections=1000, max_keepalive_connections=100)
-    ) as client:
+    async with httpx.AsyncClient() as client:
         delimiter = b"\0"
         async with client.stream(
             "POST",
             worker_addr + "/worker_generate_stream",
             headers=headers,
             json=payload,
-            timeout=WORKER_API_TIMEOUT,
+            timeout=30,
         ) as response:
             # content = await response.aread()
             buffer = b""