feat: AsyncPipeline that can schedule components to run concurrently

haystack/__init__.py

@@ -6,7 +6,7 @@
 import haystack.tracing
 from haystack.core.component import component
 from haystack.core.errors import ComponentError, DeserializationError
-from haystack.core.pipeline import Pipeline, PredefinedPipeline
+from haystack.core.pipeline import AsyncPipeline, Pipeline, PredefinedPipeline
 from haystack.core.serialization import default_from_dict, default_to_dict
 from haystack.dataclasses import Answer, Document, ExtractedAnswer, GeneratedAnswer
 
@@ -18,15 +18,16 @@
 haystack.tracing.auto_enable_tracing()
 
 __all__ = [
-    "component",
-    "default_from_dict",
-    "default_to_dict",
-    "DeserializationError",
+    "Answer",
+    "AsyncPipeline",
     "ComponentError",
-    "Pipeline",
-    "PredefinedPipeline",
+    "DeserializationError",
     "Document",
-    "Answer",
-    "GeneratedAnswer",
     "ExtractedAnswer",
+    "GeneratedAnswer",
+    "Pipeline",
+    "PredefinedPipeline",
+    "component",
+    "default_from_dict",
+    "default_to_dict",
 ]

haystack/core/pipeline/__init__.py

@@ -2,7 +2,8 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+from .async_pipeline import AsyncPipeline
 from .pipeline import Pipeline
 from .template import PredefinedPipeline
 
-__all__ = ["Pipeline", "PredefinedPipeline"]
+__all__ = ["AsyncPipeline", "Pipeline", "PredefinedPipeline"]

haystack/telemetry/_telemetry.py

@@ -8,7 +8,7 @@
 import uuid
 from collections import defaultdict
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
 
 import posthog
 import yaml
@@ -18,7 +18,7 @@
 from haystack.telemetry._environment import collect_system_specs
 
 if TYPE_CHECKING:
-    from haystack.core.pipeline import Pipeline
+    from haystack.core.pipeline import AsyncPipeline, Pipeline
 
 
 HAYSTACK_TELEMETRY_ENABLED = "HAYSTACK_TELEMETRY_ENABLED"
@@ -135,7 +135,7 @@ def send_telemetry_wrapper(*args, **kwargs):
 
 
 @send_telemetry
-def pipeline_running(pipeline: "Pipeline") -> Optional[Tuple[str, Dict[str, Any]]]:
+def pipeline_running(pipeline: Union["Pipeline", "AsyncPipeline"]) -> Optional[Tuple[str, Dict[str, Any]]]:
     """
     Collects telemetry data for a pipeline run and sends it to Posthog.
 
@@ -170,6 +170,7 @@ def pipeline_running(pipeline: "Pipeline") -> Optional[Tuple[str, Dict[str, Any]
     # Data sent to Posthog
     return "Pipeline run (2.x)", {
         "pipeline_id": str(id(pipeline)),
+        "pipeline_type": generate_qualified_class_name(type(pipeline)),
         "runs": pipeline._telemetry_runs,
         "components": components,
     }

releasenotes/notes/feat-async-pipeline-338856a142e1318c.yaml

@@ -0,0 +1,8 @@
+---
+highlights: >
+    We are introducing the `AsyncPipeline`: Supports running pipelines asynchronously. Schedules components concurrently
+    whenever possible. Leads to major speed improvements for any pipelines that may run workloads in parallel.
+features:
+  - |
+    Added a new `AsyncPipeline` implementation that allows pipelines to be executed from async code,
+    supporting concurrent scheduling of pipeline components for faster processing.

test/core/pipeline/features/conftest.py

@@ -2,14 +2,25 @@
 from typing import Tuple, List, Dict, Any, Set, Union
 from pathlib import Path
 import re
+import pytest
+import asyncio
 
 from pytest_bdd import when, then, parsers
 
-from haystack import Pipeline
+from haystack import Pipeline, AsyncPipeline
 
 PIPELINE_NAME_REGEX = re.compile(r"\[(.*)\]")
 
 
+@pytest.fixture(params=[AsyncPipeline, Pipeline])
+def pipeline_class(request):
+    """
+    A parametrized fixture that will yield AsyncPipeline for one test run
+    and Pipeline for the next test run.
+    """
+    return request.param
+
+
 @dataclass
 class PipelineRunData:
     """
@@ -34,6 +45,54 @@ class _PipelineResult:
 
 @when("I run the Pipeline", target_fixture="pipeline_result")
 def run_pipeline(
+    pipeline_data: Tuple[Union[AsyncPipeline, Pipeline], List[PipelineRunData]], spying_tracer
+) -> Union[List[Tuple[_PipelineResult, PipelineRunData]], Exception]:
+    if isinstance(pipeline_data[0], AsyncPipeline):
+        return run_async_pipeline(pipeline_data, spying_tracer)
+    else:
+        return run_sync_pipeline(pipeline_data, spying_tracer)
+
+
+def run_async_pipeline(
+    pipeline_data: Tuple[Union[AsyncPipeline], List[PipelineRunData]], spying_tracer
+) -> Union[List[Tuple[_PipelineResult, PipelineRunData]], Exception]:
+    """
+    Attempts to run a pipeline with the given inputs.
+    `pipeline_data` is a tuple that must contain:
+    * A Pipeline instance
+    * The data to run the pipeline with
+
+    If successful returns a tuple of the run outputs and the expected outputs.
+    In case an exceptions is raised returns that.
+    """
+    pipeline, pipeline_run_data = pipeline_data[0], pipeline_data[1]
+
+    results: List[_PipelineResult] = []
+
+    async def run_inner(data, include_outputs_from):
+        """Wrapper function to call pipeline.run_async method with required params."""
+        return await pipeline.run_async(data=data.inputs, include_outputs_from=include_outputs_from)
+
+    for data in pipeline_run_data:
+        try:
+            outputs = asyncio.run(run_inner(data, data.include_outputs_from))
+
+            component_calls = {
+                (span.tags["haystack.component.name"], span.tags["haystack.component.visits"]): span.tags[
+                    "haystack.component.input"
+                ]
+                for span in spying_tracer.spans
+                if "haystack.component.name" in span.tags and "haystack.component.visits" in span.tags
+            }
+            results.append(_PipelineResult(outputs=outputs, component_calls=component_calls))
+            spying_tracer.spans.clear()
+        except Exception as e:
+            return e
+
+    return [e for e in zip(results, pipeline_run_data)]
+
+
+def run_sync_pipeline(
     pipeline_data: Tuple[Pipeline, List[PipelineRunData]], spying_tracer
 ) -> Union[List[Tuple[_PipelineResult, PipelineRunData]], Exception]:
     """
@@ -61,7 +120,6 @@ def run_pipeline(
                 if "haystack.component.name" in span.tags and "haystack.component.visits" in span.tags
             }
             results.append(_PipelineResult(outputs=outputs, component_calls=component_calls))
-
             spying_tracer.spans.clear()
         except Exception as e:
             return e

test/core/pipeline/features/pipeline_run.feature

@@ -3,8 +3,8 @@ Feature: Pipeline running
     Scenario Outline: Running a correct Pipeline
         Given a pipeline <kind>
         When I run the Pipeline
-        Then it should return the expected result
-        And components are called with the expected inputs
+        Then components are called with the expected inputs
+        And it should return the expected result
 
         Examples:
         | kind |