1
\$\begingroup\$

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

  1. Is it the optimal way of solving this problem?
  2. Is storing future objects in the dict reliable? Can it be substituted with asyncio primitives like Semaphore?

Any other tips would be helpful. Thanks

\$\endgroup\$

0

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge that you have read and understand our privacy policy and code of conduct.

Browse other questions tagged or ask your own question.