Added intial group support.

Author: s3rius

.pre-commit-config.yaml

@@ -31,7 +31,8 @@ repos:
         args:
           - "check"
           - "--fix"
-          - "."
+          - "taskiq_pipelines"
+          - "tests"
 
       - id: mypy
         name: Validate types with MyPy

poetry.lock

@@ -3,9 +3,11 @@
 from taskiq_pipelines.exceptions import AbortPipeline, PipelineError
 from taskiq_pipelines.middleware import PipelineMiddleware
 from taskiq_pipelines.pipeliner import Pipeline
+from taskiq_pipelines.task_group import Group
 
 __all__ = [
     "AbortPipeline",
+    "Group",
     "Pipeline",
     "PipelineError",
     "PipelineMiddleware",

taskiq_pipelines/__init__.py

@@ -2,5 +2,6 @@
 
 CURRENT_STEP = "_pipe_current_step"
 PIPELINE_DATA = "_pipe_data"
+PARENT_TASK_ID = "_parent_task_id"
 
 EMPTY_PARAM_NAME: Literal[-1] = -1

taskiq_pipelines/constants.py

@@ -20,6 +20,8 @@
 
 from taskiq_pipelines.constants import CURRENT_STEP, EMPTY_PARAM_NAME, PIPELINE_DATA
 from taskiq_pipelines.steps import FilterStep, MapperStep, SequentialStep, parse_step
+from taskiq_pipelines.steps.group import GroupStep
+from taskiq_pipelines.task_group import Group
 
 _ReturnType = TypeVar("_ReturnType")
 _FuncParams = ParamSpec("_FuncParams")
@@ -325,6 +327,31 @@ def filter(
         )
         return self
 
+    def group(
+        self: "Pipeline[_FuncParams, _ReturnType]",
+        group: Group[_T2],
+    ) -> "Pipeline[_FuncParams, _T2]":
+        """
+        Add group task execution step.
+
+        This step will run all tasks in parallel
+        and will wait for all of them to finish.
+
+        Results of all tasks will be returned as an iterable
+        where each item is a result of the task in the group
+        with the same order.
+
+        :param group: group to execute.
+        """
+        self.steps.append(
+            DumpedStep(
+                step_type=GroupStep._step_name,
+                step_data=group.to_step().model_dump(),
+                task_id="",
+            ),
+        )
+        return self  # type: ignore
+
     def dumpb(self) -> bytes:
         """
         Dumps current pipeline as string.

taskiq_pipelines/pipeliner.py

@@ -5,6 +5,7 @@
 
 from taskiq_pipelines.abc import AbstractStep
 from taskiq_pipelines.steps.filter import FilterStep
+from taskiq_pipelines.steps.group import GroupStep
 from taskiq_pipelines.steps.mapper import MapperStep
 from taskiq_pipelines.steps.sequential import SequentialStep
 
@@ -21,6 +22,7 @@ def parse_step(step_type: str, step_data: Dict[str, Any]) -> AbstractStep:
 
 __all__ = [
     "FilterStep",
+    "GroupStep",
     "MapperStep",
     "SequentialStep",
 ]

taskiq_pipelines/steps/__init__.py

@@ -0,0 +1,115 @@
+from typing import Any, Dict, List, Optional
+
+import pydantic
+from taskiq import (
+    AsyncBroker,
+    Context,
+    TaskiqDepends,
+    TaskiqMessage,
+    TaskiqResult,
+    async_shared_broker,
+)
+
+from taskiq_pipelines.abc import AbstractStep
+from taskiq_pipelines.constants import CURRENT_STEP, PIPELINE_DATA
+from taskiq_pipelines.steps.mapper import wait_tasks
+
+
+@async_shared_broker.task(task_name="taskiq_pipelines.shared.wait_group_tasks")
+async def wait_group_tasks(
+    task_ids: List[str],
+    check_interval: float,
+    skip_errors: bool = True,
+    context: Context = TaskiqDepends(),
+) -> tuple[Any, ...]:
+    """Waits for subtasks to complete."""
+    res = await wait_tasks(
+        task_ids,
+        check_interval=check_interval,
+        skip_errors=False,
+        none_if_errors=skip_errors,
+        context=context,
+    )
+    return tuple(res)
+
+
+class GroupStepItem(pydantic.BaseModel):
+    """Item of a group step."""
+
+    task_name: str
+    labels: Dict[str, Any]
+    labels_types: Optional[Dict[str, int]] = None
+    args: List[Any]
+    kwargs: Dict[str, Any]
+
+    def from_message(self, message: TaskiqMessage) -> None:
+        """
+        Parse labels and kwargs from message.
+
+        :param message: message to parse.
+        """
+        self.labels = message.labels
+        self.labels_types = message.labels_types
+        self.args = message.args
+        self.kwargs = message.kwargs
+
+    def to_message(self, task_id: str) -> TaskiqMessage:
+        """
+        Convert this item to message.
+
+        :return: message
+        """
+        return TaskiqMessage(
+            task_id=task_id,
+            task_name=self.task_name,
+            labels=self.labels,
+            labels_types=self.labels_types,
+            args=self.args,
+            kwargs=self.kwargs,
+        )
+
+
+class GroupStep(pydantic.BaseModel, AbstractStep, step_name="group"):
+    """Step that maps iterables."""
+
+    tasks: list[GroupStepItem]
+    skip_errors: bool
+    check_interval: float
+
+    async def act(
+        self,
+        broker: AsyncBroker,
+        step_number: int,
+        parent_task_id: str,
+        task_id: str,
+        pipe_data: str,
+        result: "TaskiqResult[Any]",
+    ) -> None:
+        """
+        Execute group action.
+
+        This steps creates many small tasks
+        and one waiter task.
+
+        The waiter task awaits for all small tasks to complete,
+        and then assembles the final result.
+        """
+        ids: List[str] = []
+        for task in self.tasks:
+            subtask_id = broker.id_generator()
+            ids.append(subtask_id)
+            await broker.kick(broker.formatter.dumps(task.to_message(subtask_id)))
+
+        await (
+            wait_group_tasks.kicker()
+            .with_broker(broker)
+            .with_task_id(task_id)
+            .with_labels(
+                **{CURRENT_STEP: step_number, PIPELINE_DATA: pipe_data},  # type: ignore
+            )
+            .kiq(
+                task_ids=ids,
+                skip_errors=True,
+                check_interval=self.check_interval,
+            )
+        )

taskiq_pipelines/steps/group.py

@@ -1,4 +1,5 @@
 import asyncio
+from logging import getLogger
 from typing import Any, Dict, Iterable, List, Optional, Union
 
 import pydantic
@@ -16,12 +17,15 @@
 from taskiq_pipelines.constants import CURRENT_STEP, PIPELINE_DATA
 from taskiq_pipelines.exceptions import AbortPipeline, MappingError
 
+logger = getLogger("taskiq_pipelines")
+
 
 @async_shared_broker.task(task_name="taskiq_pipelines.shared.wait_tasks")
-async def wait_tasks(
+async def wait_tasks(  # noqa: C901
     task_ids: List[str],
     check_interval: float,
     skip_errors: bool = True,
+    none_if_errors: bool = False,
     context: Context = TaskiqDepends(),
 ) -> List[Any]:
     """
@@ -53,12 +57,16 @@ async def wait_tasks(
         if tasks_set:
             await asyncio.sleep(check_interval)
 
-    results = []
+    results: List[Any] = []
     for task_id in ordered_ids:
         result = await context.broker.result_backend.get_result(task_id)
+        logger.warning("Found error: %s", result.error)
         if result.is_err:
             if skip_errors:
                 continue
+            if none_if_errors:
+                results.append(None)
+                continue
             err_cause = None
             if isinstance(result.error, BaseException):
                 err_cause = result.error
@@ -137,6 +145,7 @@ async def act(
                 sub_task_ids,
                 check_interval=self.check_interval,
                 skip_errors=self.skip_errors,
+                none_if_errors=False,
             )
         )
 

taskiq_pipelines/steps/mapper.py

@@ -0,0 +1,86 @@
+from types import CoroutineType
+from typing import Any, Coroutine, Generic, Tuple, Union, overload
+
+from taskiq import AsyncTaskiqDecoratedTask
+from taskiq.kicker import AsyncKicker
+from typing_extensions import ParamSpec, TypeVar, TypeVarTuple, Unpack
+
+from taskiq_pipelines.steps.group import GroupStep, GroupStepItem
+
+_Tups = TypeVarTuple("_Tups")
+_T = TypeVar("_T")
+_TVal = TypeVar("_TVal")
+_Params = ParamSpec("_Params")
+
+
+class Group(Generic[_T]):
+    """
+    Group of tasks.
+
+    This class gathers multiple tasks together.
+    They will run in parallel
+
+    :param skip_errors: If True, errors in one task will not affect others.
+    """
+
+    def __init__(
+        self: "Group[Tuple[()]]",
+        skip_errors: bool = False,
+        check_interval: float = 0.1,
+    ) -> None:
+        self.tasks: Tuple[GroupStepItem, ...] = ()
+        self.skip_errors = skip_errors
+        self.check_interval = check_interval
+
+    @overload
+    def add(
+        self: "Group[Tuple[Unpack[_Tups]]]",
+        task: Union[
+            AsyncKicker[_Params, Coroutine[Any, Any, _TVal]],
+            AsyncKicker[_Params, "CoroutineType[Any, Any, _TVal]"],
+            AsyncTaskiqDecoratedTask[_Params, Coroutine[Any, Any, _TVal]],
+            AsyncTaskiqDecoratedTask[_Params, "CoroutineType[Any, Any, _TVal]"],
+        ],
+        *args: _Params.args,
+        **kwargs: _Params.kwargs,
+    ) -> "Group[Tuple[Unpack[_Tups], _TVal]]": ...
+
+    @overload
+    def add(
+        self: "Group[Tuple[Unpack[_Tups]]]",
+        task: Union[
+            AsyncKicker[_Params, _TVal],
+            AsyncTaskiqDecoratedTask[_Params, _TVal],
+        ],
+        *args: _Params.args,
+        **kwargs: _Params.kwargs,
+    ) -> "Group[Tuple[Unpack[_Tups], _TVal]]": ...
+
+    def add(
+        self: "Group[Any]",
+        task: Union[AsyncKicker[_Params, Any], AsyncTaskiqDecoratedTask[_Params, Any]],
+        *args: _Params.args,
+        **kwargs: _Params.kwargs,
+    ) -> "Any":
+        """Add task to a group."""
+        kicker = task.kicker() if isinstance(task, AsyncTaskiqDecoratedTask) else task
+        message = kicker._prepare_message(*args, **kwargs)
+        self.tasks = (
+            *self.tasks,
+            GroupStepItem(
+                task_name=message.task_name,
+                labels=message.labels,
+                labels_types=message.labels_types,
+                args=message.args,
+                kwargs=message.kwargs,
+            ),
+        )
+        return self
+
+    def to_step(self) -> GroupStep:
+        """Convert group definition to a step."""
+        return GroupStep(
+            tasks=list(self.tasks),
+            skip_errors=self.skip_errors,
+            check_interval=self.check_interval,
+        )