Introduction
I have a telegram bot based on aiogram
framework. It does GPU intensive computations locally then user makes a request and outputs a text file after computations are done (which takes from 1 to 10 minutes).
I figured out that I should use a second thread for computations and leave the main thread for bot handlers, so my bot wouldn't freeze in this time.
Usually people recommend creating tasks in several threads in advance and starting them using asyncio.gather()
, but I don't see how it could be implemented in my case.
Instead, I create tasks as I need and await for them only then they are finished using a timer.
Code
Here I'm creating a function I want to compute and wrapping it in aiogram function which basically tells the bot there to send the output file. When user sends video to the bot, this handler is called. (end of the handler is relevant)
# src/bot/action_handlers.py
import logging
from functools import partial
from pathlib import Path
from aiogram.types import Message
from src.bot.actions import run_video, worker
from src.bot.bot_locale import BotReply
from src.config import Config
from src.database import Database
from src.setup_handler import get_handler
router = Router()
replies = BotReply()
database = Database()
defaults = Config('./configs/settings_defaults.yaml')
bot_settings = Config('./configs/bot_settings.yaml')
@router.message(F.video)
async def video_handler(message: Message, bot: Bot) -> None:
logger.info('Call: video_handler')
user_id = message.from_user.id
if await is_generating(message):
return
file_name = message.video.file_name
if file_name is None:
msg = 'Could not find file name'
raise ValueError(msg)
file_path = Path(f'./temp') / file_name
await bot.download(
file=message.video.file_id,
destination=file_path
)
await message.answer(
replies.answers(message.from_user.id, 'general')[
'got_video']
)
# Preparing data
with database as db:
settings = db.get_settings(user_id)
audio_model = bot_settings['audio_model']
text_model = bot_settings['text_model']
document_format = settings['document_format']
document_language = settings['document_language']
text_format = settings['text_format']
summary_func = partial(
run_video,
file_path,
text_model,
audio_model,
document_format,
text_format,
document_language,
)
wrapper = partial(message_wrapper, summary_func, message)
await worker(message, wrapper)
def message_wrapper(foo, message: Message):
out = foo()
return out, message
Now I move the function on the second thread and save the future object in the dict:
# src/bot/actions.py
import asyncio
import logging
from concurrent.futures import ThreadPoolExecutor
from aiogram.types import Message
from src.bot.bot_locale import BotReply
from src.bot.exceptions import TooManyTasksError
from src.config import Config
from src.database import user_tasks # Importing empty dict
bot_settings = Config('./configs/bot_settings.yaml')
executor = ThreadPoolExecutor(max_workers=1)
replies = BotReply()
async def worker(message: Message, foo):
len_t = len(user_tasks)
if len_t > bot_settings['max_tasks']:
msg = f"Tasks: {user_tasks} of len={len_t} > {bot_settings['max_tasks']}"
raise TooManyTasksError(msg)
loop = asyncio.get_running_loop()
task = asyncio.ensure_future(loop.run_in_executor(executor, foo))
# Store the future object to be awaited later
user_tasks[message.from_user.id] = task
The code awaits for the tasks to finish in the main function. It checks the user_tasks
dict every few seconds for finished tasks and awaits for them if there are any:
# src/bot/bot.py
import asyncio
import logging
from apscheduler.schedulers.asyncio import AsyncIOScheduler
from src.bot.action_handlers import router as action_router
from src.bot.bot_locale import BotReply
from src.config import Config
from src.database import user_tasks # Importing empty dict
database = Database()
bot_settings = Config('./configs/bot_settings.yaml')
async def check_tasks():
to_del = []
for user_id, task in user_tasks.items():
if task.done():
await task
out, message = task.result()
to_del.append(user_id)
user_id = message.from_user.id
path = await out
doc_file = FSInputFile(path)
caption = replies.answers(message.from_user.id, 'general')['document_caption']
await message.answer_document(
doc_file,
caption
)
Path(path).unlink(missing_ok=True)
if to_del:
for id in to_del:
del user_tasks[id]
async def main() -> None:
# ...
bot = Bot(token=TOKEN, parse_mode="MarkdownV2")
dp = Dispatcher()
# Setting up scheduler
look_after_tasks = AsyncIOScheduler()
look_after_tasks.add_job(
check_tasks,
trigger='interval',
seconds=4,
)
look_after_tasks.start()
# ...
The full code can be found in my repository https://github.com/dvarkless/video_summarizer_bot
Questions
- Is it the optimal way of solving this problem?
- Is storing future objects in the dict reliable? Can it be substituted with
asyncio
primitives likeSemaphore
?
Any other tips would be helpful. Thanks