Compare commits

..

3 commits

Author SHA1 Message Date
monosans ab3a019e0e
Merge 0b08c5b712 into 719f29fb29 2023-10-07 15:07:40 +00:00
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 asyncio
import starlette import starlette
from typing import Set from typing import Any, Coroutine, Set
from rich import print from rich import print
from dotenv import load_dotenv from dotenv import load_dotenv
@ -24,7 +24,18 @@ CRITICAL_API_ERRORS = ['invalid_api_key', 'account_deactivated']
keymanager = providerkeys.manager 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( async def respond(
path: str='/v1/chat/completions', path: str='/v1/chat/completions',
@ -169,7 +180,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)
task = asyncio.create_task( create_background_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,
@ -181,5 +192,3 @@ async def respond(
model=model, model=model,
) )
) )
background_tasks.add(task)
task.add_done_callback(background_tasks.discard)