Getting all running tasks

[kevinhikaruevans]

I'd like to get all the tasks that are running. I'm currently using Redis as my broker and result backend.

Is that possible to do out-of-box? I'm still pretty new to taskiq, but it looks like this wouldn't be possible because the task is removed from the Redis queue (taskiq) during its execution and only accessible when the task is finished.

Do I need to create a middleware to temporarily store running tasks?

Thanks!

[kevinhikaruevans]

Here's a quick update. I'm not sure if this is the best way of doing this, but I wrote a middleware to track the task_ids for Redis:

from taskiq_redis.redis_broker import Redis, BaseRedisBroker

class RunningTasksMiddleware(TaskiqMiddleware):

    def __init__(self, list_name: str = 'taskiq_running') -> None:
        super().__init__()
        
        self.list_name = list_name

    async def startup(self) -> Coroutine[Any, Any, None] | None:
        broker: BaseRedisBroker = self.broker

        if broker.is_worker_process:
            # clear on startup
            async with Redis(connection_pool=broker.connection_pool) as redis_conn:
                await redis_conn.delete(self.list_name)
            
        return super().startup()
    
    async def get_running_task_ids(self) -> list[str]:
        """gets a list of running task IDs

        Returns:
            list[str]: list of task IDs as strings
        """
        broker: BaseRedisBroker = self.broker

        async with Redis(connection_pool=broker.connection_pool) as redis_conn:
            return await redis_conn.lrange(self.list_name, 0, -1)

    async def pre_send(self, message: TaskiqMessage) -> TaskiqMessage | Coroutine[Any, Any, TaskiqMessage]:
        broker: BaseRedisBroker = self.broker
        
        async with Redis(connection_pool=broker.connection_pool) as redis_conn:
            await redis_conn.lpush(self.list_name, message.task_id)
    
        return super().pre_send(message)
        
    
    async def post_save(self, message: TaskiqMessage, result: TaskiqResult[Any]) -> Coroutine[Any, Any, None] | None:
        broker: BaseRedisBroker = self.broker
        
        async with Redis(connection_pool=broker.connection_pool) as redis_conn:
            await redis_conn.lrem(self.list_name, 0, message.task_id)

Although I might end up serializing the TaskiqMessages and sticking those in Redis instead of just their task_ids, but that might be a mistake. 🤔

[kevinhikaruevans]

Here's one that serializes the messages:

from taskiq_redis.redis_broker import Redis, BaseRedisBroker

class RunningTasksMiddleware(TaskiqMiddleware):

    def __init__(self, list_name: str = 'taskiq_running') -> None:
        super().__init__()
        
        self.list_name = list_name

    async def startup(self) -> Coroutine[Any, Any, None] | None:
        broker: BaseRedisBroker = self.broker

        if broker.is_worker_process:
            # clear on startup
            async with Redis(connection_pool=broker.connection_pool) as redis_conn:
                await redis_conn.delete(self.list_name)
            
        return super().startup()
    
    async def get_running_tasks(self) -> list[TaskiqMessage]:
        """gets a list of running task IDs

        Returns:
            list[str]: list of task IDs as strings
        """
        broker: BaseRedisBroker = self.broker
        
        async with Redis(connection_pool=broker.connection_pool) as redis_conn:
            return list(map(lambda item: TaskiqMessage.model_validate_json(item), await redis_conn.lrange(self.list_name, 0, -1)))

    async def pre_send(self, message: TaskiqMessage) -> TaskiqMessage | Coroutine[Any, Any, TaskiqMessage]:
        broker: BaseRedisBroker = self.broker
        
        async with Redis(connection_pool=broker.connection_pool) as redis_conn:
            await redis_conn.lpush(self.list_name, message.model_dump_json())
    
        return super().pre_send(message)
        
    
    async def post_save(self, message: TaskiqMessage, result: TaskiqResult[Any]) -> Coroutine[Any, Any, None] | None:
        broker: BaseRedisBroker = self.broker
        
        async with Redis(connection_pool=broker.connection_pool) as redis_conn:
            await redis_conn.lrem(self.list_name, 0, message.model_dump_json())

Not sure if this is a terrible practice or not...

[s3rius]

Actually, not bad implementation. But there's a small flaw in the design. You perform N operations every time you delete a task. It's fine for not highload systems, but may result in poor performance.

I'd suggest you to use scan for keys and batch operations, like mget.

from typing import Any
from taskiq import TaskiqMessage, TaskiqMiddleware, TaskiqResult
from taskiq_redis import ListQueueBroker, RedisAsyncResultBackend
from taskiq_redis.redis_broker import BaseRedisBroker
from redis.asyncio import Redis


class RunningTasksMiddleware(TaskiqMiddleware):
    def __init__(self, prefix: str = "running:") -> None:
        super().__init__()

        self.prefix = prefix

    async def startup(self) -> None:
        broker: BaseRedisBroker = self.broker

        if broker.is_worker_process:
            # clear on startup
            async with Redis(connection_pool=broker.connection_pool) as redis_conn:
                cursor = "0"
                while cursor != 0:
                    cursor, keys = await redis_conn.scan(
                        cursor=cursor, match="running:*"
                    )
                    if not keys:
                        continue
                    await redis_conn.delete(*keys)

        return super().startup()

    async def get_running_tasks(self) -> list[TaskiqMessage]:
        """gets a list of running task IDs

        Returns:
            list[str]: list of task IDs as strings
        """
        broker: BaseRedisBroker = self.broker
        messages = []
        async with Redis(connection_pool=broker.connection_pool) as redis_conn:
            cursor = "0"
            while cursor != 0:
                cursor, keys = await redis_conn.scan(cursor=cursor, match="running:*")
                values = await redis_conn.mget(keys)
                for value in values:
                    messages.append(TaskiqMessage.model_validate_json(value))
        return messages

    async def pre_send(self, message: TaskiqMessage) -> TaskiqMessage:
        broker: BaseRedisBroker = self.broker

        async with Redis(connection_pool=broker.connection_pool) as redis_conn:
            await redis_conn.set(
                f"{self.prefix}{message.task_id}",
                message.model_dump_json(),
            )

        return super().pre_send(message)

    async def post_save(self, message: TaskiqMessage, _: TaskiqResult[Any]) -> None:
        broker: BaseRedisBroker = self.broker

        async with Redis(connection_pool=broker.connection_pool) as redis_conn:
            await redis_conn.delete(f"{self.prefix}{message.task_id}")

This operation performs more actions when you call get_running_tasks, but it set and delete results much faster which won't affect your worker runtime. Also it makes possible to store large amount of tasks in redis.

The idea was totally correct, I just made it more reliable to bigger workloads.

[pushpak1300]

Is there way to get all the failed taks ? @s3rius