Compare commits

..

2 commits

Author SHA1 Message Date
monosans 0b08c5b712
Add docstring to create_background_task 2023-10-07 18:07:35 +03:00
monosans 266aa3e0fe
Fix dangling asyncio tasks 2023-10-07 16:09:55 +03:00

View file

@ -7,7 +7,7 @@ import aiohttp
import asyncio
import starlette
from typing import Set
from typing import Any, Coroutine, Set
from rich import print
from dotenv import load_dotenv
@ -24,7 +24,18 @@ CRITICAL_API_ERRORS = ['invalid_api_key', 'account_deactivated']
keymanager = providerkeys.manager
background_tasks: Set[asyncio.Task] = set()
background_tasks: Set[asyncio.Task[Any]] = 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(
path: str='/v1/chat/completions',
@ -169,7 +180,7 @@ async def respond(
if (not is_stream) and server_json_response:
yield json.dumps(server_json_response)
task = asyncio.create_task(
create_background_task(
after_request.after_request(
incoming_request=incoming_request,
target_request=target_request,
@ -181,5 +192,3 @@ async def respond(
model=model,
)
)
background_tasks.add(task)
task.add_done_callback(background_tasks.discard)