[BugFix] Prevent the task of _force_log from being garbage collected (#4567)

This commit is contained in:
Yang, Bo 2024-05-02 18:35:18 -07:00 committed by GitHub
parent 344a5d0c33
commit 808632d3b4
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -4,6 +4,7 @@ import inspect
import re import re
from contextlib import asynccontextmanager from contextlib import asynccontextmanager
from http import HTTPStatus from http import HTTPStatus
from typing import Any, Set
import fastapi import fastapi
import uvicorn import uvicorn
@ -33,6 +34,8 @@ openai_serving_chat: OpenAIServingChat
openai_serving_completion: OpenAIServingCompletion openai_serving_completion: OpenAIServingCompletion
logger = init_logger(__name__) logger = init_logger(__name__)
_running_tasks: Set[asyncio.Task[Any]] = set()
@asynccontextmanager @asynccontextmanager
async def lifespan(app: fastapi.FastAPI): async def lifespan(app: fastapi.FastAPI):
@ -43,7 +46,9 @@ async def lifespan(app: fastapi.FastAPI):
await engine.do_log_stats() await engine.do_log_stats()
if not engine_args.disable_log_stats: if not engine_args.disable_log_stats:
asyncio.create_task(_force_log()) task = asyncio.create_task(_force_log())
_running_tasks.add(task)
task.add_done_callback(_running_tasks.remove)
yield yield