TRFC-1 Pipelines redesign

[s3rius]

Abstract

Our current pipeline solution doesn't fit many use cases. We aim to support as many cases as possible with minimal effort.

Current Pipeline Features:

• Filter
• Map
• Chain (call_next, call_after)

Missing Functionality:

• Groups
• Compensational tasks
• Conditional tasks

Additionally, every component should be composable—meaning you can nest groups within other groups.

---

Solution

First, we should update our PipelineAPI and introduce a generic class that the pipeline can append to. All components will be added using a single method: .then().

Example usage:

pipe = (
    Pipeline(task1)
    .then(Call(task2))
    .then(Map(task2))
    .then(Group().add(task3).add(task4).add(task5))
    .then(Conditional(condition, task6))
    .then(Compesable(on_error_task, task7))
)

This API allows us to mimic a bind-like operator from functional programming languages in Python. Here's the same pipeline using the bind (>>) approach:

class Pipeline:

    def add(self, step: Any) -> Pipeline:
        ...

    def __rshift__(self, step: Any) -> Pipeline:
        return self.add(step)

With this syntax, pipelines can be built like so:

pipe = (
    Pipeline(task1)
    >> Call(task2)
    >> Map(task2)
    >> Group().add(task3).add(task4).add(task5)
    >> Conditional(condition, task6)
    >> Compesable(on_error_task, task7)
)

Also, this approach allows us to make it easier to extens the whole pipeline engine within user's codebases.

Interface

In order to achieve desired API, we need to be able to serialize every component of the pipeline into a single intermediate representation.

Here’s the current generic representation of a pipeline step:

class DumpedStep(pydantic.BaseModel):
    """Dumped state model."""

    step_type: str
    step_data: Dict[str, Any]
    task_id: str

As you can see, this structure doesn't support groups and other nested features well. With nested pipelines or groups, we might want to track each step individually—or at least be aware of each nested step—at the time when users calls kiq to monitor progress effectively.

Other ideas:

We might want to introduce unique byte sequence that will mean that result of previous function should be used. It will allow us to make much cleaner code.

As an example:

# Defined as a constant at taskiq_pipelines
STEP_RESULT: Any = b"taskiq-taskiq-\x99\x00\x1337"

##### 

# Current solution
pipe = (
        Pipeline(broker, add_one)
        # Here param_name is used
        .call_next(repeat, param_name="value", reps=4)
)

# Proposed solution
# Imagine the following signature of a call_next method.
# call_next[Pspec](func: Callable[Pspec, Any], *args: Pspec.args, **kwargs: Pspec.kwargs) -> Pipeline: ...

pipe = (
        Pipeline(broker, add_one)
        # STEP_RESULT passed as a first argument
        .call_next(repeat, STEP_RESULT, reps=4)
        # Or alternatively as a kwargs
        .call_next(repeat, value=STEP_RESULT, reps=4)
)

The STEP_RESULT should be a unique byte sequence that is very unlikely to be met.

This will allow us to improve static type checking, because now it will be validating all the arguments of tasks during creation of pipeline, because of ParamSpec. Of course we need to verify that this idea can be implemented as proposed.

[wallseat]

First of all, thank you for your work, I really like it.

I want to add smth to this (T)RFC

1. As part of the additional functionality, it is also worth mentioning the star map and chunks
   https://docs.celeryq.dev/en/latest/userguide/canvas.html#map-starmap
   https://docs.celeryq.dev/en/latest/userguide/canvas.html#chunks
2. Can you describe in more detail how nested tasks will work? And what an alternative way to store information about pipeline state?
3. What about pass intial arguments to Group?

[s3rius]

1. This should be relatively easy to implement given the proposed architecture.

2. First, nesting should execute a number of tasks and wait for them to complete.
   For example, groups spawn one main task that waits for all declared tasks to finish.
   Groups with nested groups will wait for the main task of the nested group.

Consider this code:

pipe = Pipeline(broker) >> (
    Group()
    .add(
        # Group 2
        Group().add(task21).add(task22).add(task23)
    )
    .add(task11)
    .add(
        # Group 3
        Group().add(task31).add(task32).add(task33)
    )
)

Each nested group becomes a separate task that waits for the completion of all its nested tasks. Here's a visual example:

A task in Group1 will only wait for the completion of task11, Group2, and Group3. Each group waits only for its direct children to complete.

3. I believe you just need to define all arguments when adding tasks to the group.