Compare commits

..

2 commits

Author SHA1 Message Date
monosans f5be895f48
Merge 03f8c9ad0a into 719f29fb29 2023-10-07 06:14:57 +00:00
monosans 03f8c9ad0a
Fix dangling asyncio tasks 2023-10-07 09:14:51 +03:00

View file

@ -7,7 +7,7 @@ import aiohttp
import asyncio import asyncio
import starlette import starlette
from typing import Any, Coroutine, Set from typing import Set
from rich import print from rich import print
from dotenv import load_dotenv from dotenv import load_dotenv
@ -24,18 +24,7 @@ CRITICAL_API_ERRORS = ['invalid_api_key', 'account_deactivated']
keymanager = providerkeys.manager keymanager = providerkeys.manager
background_tasks: Set[asyncio.Task[Any]] = set() background_tasks: Set[asyncio.Task] = set()
def create_background_task(coro: Coroutine[Any, Any, Any]) -> None:
"""asyncio.create_task, which prevents the task from being garbage collected.
https://docs.python.org/3/library/asyncio-task.html#asyncio.create_task
"""
task = asyncio.create_task(coro)
background_tasks.add(task)
task.add_done_callback(background_tasks.discard)
async def respond( async def respond(
path: str='/v1/chat/completions', path: str='/v1/chat/completions',
@ -180,7 +169,7 @@ async def respond(
if (not is_stream) and server_json_response: if (not is_stream) and server_json_response:
yield json.dumps(server_json_response) yield json.dumps(server_json_response)
create_background_task( task = asyncio.create_task(
after_request.after_request( after_request.after_request(
incoming_request=incoming_request, incoming_request=incoming_request,
target_request=target_request, target_request=target_request,
@ -192,3 +181,5 @@ async def respond(
model=model, model=model,
) )
) )
background_tasks.add(task)
task.add_done_callback(background_tasks.discard)