ParallelPipeline: use PipelineSource/PipelineSink and remove tasks
This commit is contained in:
@@ -54,8 +54,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
|||||||
|
|
||||||
### Performance
|
### Performance
|
||||||
|
|
||||||
- Improve `ParallelPipeline` performance by using direct mode and by not
|
- Improve `ParallelPipeline` performance by using direct mode, by not
|
||||||
creating a task for each frame and every sub-pipeline.
|
creating a task for each frame and every sub-pipeline and also by removing
|
||||||
|
other unnecessary tasks.
|
||||||
|
|
||||||
- `Pipeline` performance improvements by using direct mode.
|
- `Pipeline` performance improvements by using direct mode.
|
||||||
|
|
||||||
|
|||||||
@@ -11,106 +11,15 @@ sub-pipelines concurrently, with coordination for system frames and proper
|
|||||||
handling of pipeline lifecycle events.
|
handling of pipeline lifecycle events.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import asyncio
|
|
||||||
from itertools import chain
|
from itertools import chain
|
||||||
from typing import Awaitable, Callable, Dict, List
|
from typing import Dict, List
|
||||||
|
|
||||||
from loguru import logger
|
from loguru import logger
|
||||||
|
|
||||||
from pipecat.frames.frames import (
|
from pipecat.frames.frames import EndFrame, Frame, StartFrame
|
||||||
CancelFrame,
|
|
||||||
EndFrame,
|
|
||||||
Frame,
|
|
||||||
StartFrame,
|
|
||||||
StartInterruptionFrame,
|
|
||||||
SystemFrame,
|
|
||||||
)
|
|
||||||
from pipecat.pipeline.base_pipeline import BasePipeline
|
from pipecat.pipeline.base_pipeline import BasePipeline
|
||||||
from pipecat.pipeline.pipeline import Pipeline
|
from pipecat.pipeline.pipeline import Pipeline, PipelineSink, PipelineSource
|
||||||
from pipecat.processors.frame_processor import FrameDirection, FrameProcessor, FrameProcessorSetup
|
from pipecat.processors.frame_processor import FrameDirection, FrameProcessor, FrameProcessorSetup
|
||||||
from pipecat.utils.asyncio.watchdog_queue import WatchdogQueue
|
|
||||||
|
|
||||||
|
|
||||||
class ParallelPipelineSource(FrameProcessor):
|
|
||||||
"""Source processor for parallel pipeline branches.
|
|
||||||
|
|
||||||
Handles frame routing for parallel pipeline inputs, directing system frames
|
|
||||||
to the parent push function and other upstream frames to a queue for processing.
|
|
||||||
"""
|
|
||||||
|
|
||||||
def __init__(
|
|
||||||
self,
|
|
||||||
upstream_queue: asyncio.Queue,
|
|
||||||
push_frame_func: Callable[[Frame, FrameDirection], Awaitable[None]],
|
|
||||||
):
|
|
||||||
"""Initialize the parallel pipeline source.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
upstream_queue: Queue for collecting upstream frames from this branch.
|
|
||||||
push_frame_func: Function to push frames to the parent parallel pipeline.
|
|
||||||
"""
|
|
||||||
super().__init__(enable_direct_mode=True)
|
|
||||||
self._up_queue = upstream_queue
|
|
||||||
self._push_frame_func = push_frame_func
|
|
||||||
|
|
||||||
async def process_frame(self, frame: Frame, direction: FrameDirection):
|
|
||||||
"""Process frames with special handling for system frames.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
frame: The frame to process.
|
|
||||||
direction: The direction of frame flow.
|
|
||||||
"""
|
|
||||||
await super().process_frame(frame, direction)
|
|
||||||
|
|
||||||
match direction:
|
|
||||||
case FrameDirection.UPSTREAM:
|
|
||||||
if isinstance(frame, SystemFrame):
|
|
||||||
await self._push_frame_func(frame, direction)
|
|
||||||
else:
|
|
||||||
await self._up_queue.put(frame)
|
|
||||||
case FrameDirection.DOWNSTREAM:
|
|
||||||
await self.push_frame(frame, direction)
|
|
||||||
|
|
||||||
|
|
||||||
class ParallelPipelineSink(FrameProcessor):
|
|
||||||
"""Sink processor for parallel pipeline branches.
|
|
||||||
|
|
||||||
Handles frame routing for parallel pipeline outputs, directing system frames
|
|
||||||
to the parent push function and other downstream frames to a queue for coordination.
|
|
||||||
"""
|
|
||||||
|
|
||||||
def __init__(
|
|
||||||
self,
|
|
||||||
downstream_queue: asyncio.Queue,
|
|
||||||
push_frame_func: Callable[[Frame, FrameDirection], Awaitable[None]],
|
|
||||||
):
|
|
||||||
"""Initialize the parallel pipeline sink.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
downstream_queue: Queue for collecting downstream frames from this branch.
|
|
||||||
push_frame_func: Function to push frames to the parent parallel pipeline.
|
|
||||||
"""
|
|
||||||
super().__init__(enable_direct_mode=True)
|
|
||||||
self._down_queue = downstream_queue
|
|
||||||
self._push_frame_func = push_frame_func
|
|
||||||
|
|
||||||
async def process_frame(self, frame: Frame, direction: FrameDirection):
|
|
||||||
"""Process frames with special handling for system frames.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
frame: The frame to process.
|
|
||||||
direction: The direction of frame flow.
|
|
||||||
"""
|
|
||||||
await super().process_frame(frame, direction)
|
|
||||||
|
|
||||||
match direction:
|
|
||||||
case FrameDirection.UPSTREAM:
|
|
||||||
await self.push_frame(frame, direction)
|
|
||||||
case FrameDirection.DOWNSTREAM:
|
|
||||||
if isinstance(frame, SystemFrame):
|
|
||||||
await self._push_frame_func(frame, direction)
|
|
||||||
else:
|
|
||||||
await self._down_queue.put(frame)
|
|
||||||
|
|
||||||
|
|
||||||
class ParallelPipeline(BasePipeline):
|
class ParallelPipeline(BasePipeline):
|
||||||
@@ -132,21 +41,38 @@ class ParallelPipeline(BasePipeline):
|
|||||||
Exception: If no processor lists are provided.
|
Exception: If no processor lists are provided.
|
||||||
TypeError: If any argument is not a list of processors.
|
TypeError: If any argument is not a list of processors.
|
||||||
"""
|
"""
|
||||||
super().__init__(enable_direct_mode=True)
|
# We don't set it to direct mode because we use frame pausing and that
|
||||||
|
# requires queues.
|
||||||
|
super().__init__()
|
||||||
|
|
||||||
if len(args) == 0:
|
if len(args) == 0:
|
||||||
raise Exception(f"ParallelPipeline needs at least one argument")
|
raise Exception(f"ParallelPipeline needs at least one argument")
|
||||||
|
|
||||||
self._args = args
|
|
||||||
self._pipelines = []
|
self._pipelines = []
|
||||||
|
|
||||||
self._seen_ids = set()
|
self._seen_ids = set()
|
||||||
self._endframe_counter: Dict[int, int] = {}
|
self._frame_counter: Dict[int, int] = {}
|
||||||
self._start_frame_counter: Dict[int, int] = {}
|
|
||||||
self._started = False
|
|
||||||
|
|
||||||
self._up_task = None
|
logger.debug(f"Creating {self} pipelines")
|
||||||
self._down_task = None
|
|
||||||
|
for processors in args:
|
||||||
|
if not isinstance(processors, list):
|
||||||
|
raise TypeError(f"ParallelPipeline argument {processors} is not a list")
|
||||||
|
|
||||||
|
num_pipelines = len(self._pipelines)
|
||||||
|
|
||||||
|
# We add a source before the pipeline and a sink after so we control
|
||||||
|
# the frames that are pushed upstream and downstream.
|
||||||
|
source = PipelineSource(
|
||||||
|
self._parallel_push_frame, name=f"{self}::Source{num_pipelines}"
|
||||||
|
)
|
||||||
|
sink = PipelineSink(self._pipeline_sink_push_frame, name=f"{self}::Sink{num_pipelines}")
|
||||||
|
|
||||||
|
# Create pipeline
|
||||||
|
pipeline = Pipeline(processors, source=source, sink=sink)
|
||||||
|
self._pipelines.append(pipeline)
|
||||||
|
|
||||||
|
logger.debug(f"Finished creating {self} pipelines")
|
||||||
|
|
||||||
#
|
#
|
||||||
# Frame processor
|
# Frame processor
|
||||||
@@ -196,32 +122,14 @@ class ParallelPipeline(BasePipeline):
|
|||||||
TypeError: If any processor list argument is not actually a list.
|
TypeError: If any processor list argument is not actually a list.
|
||||||
"""
|
"""
|
||||||
await super().setup(setup)
|
await super().setup(setup)
|
||||||
|
for p in self._pipelines:
|
||||||
self._up_queue = WatchdogQueue(setup.task_manager)
|
await p.setup(setup)
|
||||||
self._down_queue = WatchdogQueue(setup.task_manager)
|
|
||||||
|
|
||||||
logger.debug(f"Creating {self} pipelines")
|
|
||||||
for processors in self._args:
|
|
||||||
if not isinstance(processors, list):
|
|
||||||
raise TypeError(f"ParallelPipeline argument {processors} is not a list")
|
|
||||||
|
|
||||||
# We add a source before the pipeline and a sink after so we control
|
|
||||||
# the frames that are pushed upstream and downstream.
|
|
||||||
source = ParallelPipelineSource(self._up_queue, self._parallel_push_frame)
|
|
||||||
sink = ParallelPipelineSink(self._down_queue, self._pipeline_sink_push_frame)
|
|
||||||
|
|
||||||
# Create pipeline
|
|
||||||
pipeline = Pipeline([source] + processors + [sink])
|
|
||||||
self._pipelines.append(pipeline)
|
|
||||||
|
|
||||||
logger.debug(f"Finished creating {self} pipelines")
|
|
||||||
|
|
||||||
await asyncio.gather(*[p.setup(setup) for p in self._pipelines])
|
|
||||||
|
|
||||||
async def cleanup(self):
|
async def cleanup(self):
|
||||||
"""Clean up the parallel pipeline and all its branches."""
|
"""Clean up the parallel pipeline and all its branches."""
|
||||||
await super().cleanup()
|
await super().cleanup()
|
||||||
await asyncio.gather(*[p.cleanup() for p in self._pipelines])
|
for p in self._pipelines:
|
||||||
|
await p.cleanup()
|
||||||
|
|
||||||
async def process_frame(self, frame: Frame, direction: FrameDirection):
|
async def process_frame(self, frame: Frame, direction: FrameDirection):
|
||||||
"""Process frames through all parallel branches with lifecycle coordination.
|
"""Process frames through all parallel branches with lifecycle coordination.
|
||||||
@@ -232,77 +140,16 @@ class ParallelPipeline(BasePipeline):
|
|||||||
"""
|
"""
|
||||||
await super().process_frame(frame, direction)
|
await super().process_frame(frame, direction)
|
||||||
|
|
||||||
if isinstance(frame, StartFrame):
|
# Parallel pipeline synchronized frames.
|
||||||
self._start_frame_counter[frame.id] = len(self._pipelines)
|
if isinstance(frame, (StartFrame, EndFrame)):
|
||||||
elif isinstance(frame, EndFrame):
|
self._frame_counter[frame.id] = len(self._pipelines)
|
||||||
self._endframe_counter[frame.id] = len(self._pipelines)
|
await self.pause_processing_system_frames()
|
||||||
elif isinstance(frame, CancelFrame):
|
await self.pause_processing_frames()
|
||||||
await self._cancel()
|
|
||||||
|
|
||||||
# Process frames in each of the sub-pipelines.
|
# Process frames in each of the sub-pipelines.
|
||||||
for p in self._pipelines:
|
for p in self._pipelines:
|
||||||
await p.queue_frame(frame, direction)
|
await p.queue_frame(frame, direction)
|
||||||
|
|
||||||
# Handle interruptions after everything has been cancelled.
|
|
||||||
if isinstance(frame, StartInterruptionFrame):
|
|
||||||
await self._handle_interruption()
|
|
||||||
# Wait for tasks to finish.
|
|
||||||
elif isinstance(frame, EndFrame):
|
|
||||||
await self._stop()
|
|
||||||
|
|
||||||
async def _start(self, frame: StartFrame):
|
|
||||||
"""Start the parallel pipeline processing tasks."""
|
|
||||||
await self._create_tasks()
|
|
||||||
|
|
||||||
async def _stop(self):
|
|
||||||
"""Stop all parallel pipeline processing tasks."""
|
|
||||||
if self._up_task:
|
|
||||||
# The up task doesn't receive an EndFrame, so we just cancel it.
|
|
||||||
await self.cancel_task(self._up_task)
|
|
||||||
self._up_task = None
|
|
||||||
|
|
||||||
if self._down_task:
|
|
||||||
# The down tasks waits for the last EndFrame sent by the internal
|
|
||||||
# pipelines.
|
|
||||||
await self._down_task
|
|
||||||
self._down_task = None
|
|
||||||
|
|
||||||
async def _cancel(self):
|
|
||||||
"""Cancel all parallel pipeline processing tasks."""
|
|
||||||
if self._up_task:
|
|
||||||
self._up_queue.cancel()
|
|
||||||
await self.cancel_task(self._up_task)
|
|
||||||
self._up_task = None
|
|
||||||
if self._down_task:
|
|
||||||
self._down_queue.cancel()
|
|
||||||
await self.cancel_task(self._down_task)
|
|
||||||
self._down_task = None
|
|
||||||
|
|
||||||
async def _create_tasks(self):
|
|
||||||
"""Create upstream and downstream processing tasks if not already running."""
|
|
||||||
if not self._up_task:
|
|
||||||
self._up_task = self.create_task(self._process_up_queue())
|
|
||||||
if not self._down_task:
|
|
||||||
self._down_task = self.create_task(self._process_down_queue())
|
|
||||||
|
|
||||||
async def _drain_queue(self, queue: asyncio.Queue):
|
|
||||||
try:
|
|
||||||
while not queue.empty():
|
|
||||||
queue.get_nowait()
|
|
||||||
except asyncio.QueueEmpty:
|
|
||||||
logger.debug(f"Draining {self} queue already empty")
|
|
||||||
|
|
||||||
async def _drain_queues(self):
|
|
||||||
"""Drain all frames from upstream and downstream queues."""
|
|
||||||
await self._drain_queue(self._up_queue)
|
|
||||||
await self._drain_queue(self._down_queue)
|
|
||||||
|
|
||||||
async def _handle_interruption(self):
|
|
||||||
"""Handle interruption by cancelling tasks, draining queues, and restarting."""
|
|
||||||
await self._cancel()
|
|
||||||
await self._drain_queues()
|
|
||||||
await self._create_tasks()
|
|
||||||
|
|
||||||
async def _parallel_push_frame(self, frame: Frame, direction: FrameDirection):
|
async def _parallel_push_frame(self, frame: Frame, direction: FrameDirection):
|
||||||
"""Push frames while avoiding duplicates using frame ID tracking."""
|
"""Push frames while avoiding duplicates using frame ID tracking."""
|
||||||
if frame.id not in self._seen_ids:
|
if frame.id not in self._seen_ids:
|
||||||
@@ -310,52 +157,18 @@ class ParallelPipeline(BasePipeline):
|
|||||||
await self.push_frame(frame, direction)
|
await self.push_frame(frame, direction)
|
||||||
|
|
||||||
async def _pipeline_sink_push_frame(self, frame: Frame, direction: FrameDirection):
|
async def _pipeline_sink_push_frame(self, frame: Frame, direction: FrameDirection):
|
||||||
if isinstance(frame, StartFrame):
|
# Parallel pipeline synchronized frames.
|
||||||
# Decrement counter and check if all pipelines have processed the StartFrame
|
if isinstance(frame, (StartFrame, EndFrame)):
|
||||||
start_frame_counter = self._start_frame_counter.get(frame.id, 0)
|
# Decrement counter.
|
||||||
if start_frame_counter > 0:
|
frame_counter = self._frame_counter.get(frame.id, 0)
|
||||||
self._start_frame_counter[frame.id] -= 1
|
if frame_counter > 0:
|
||||||
start_frame_counter = self._start_frame_counter[frame.id]
|
self._frame_counter[frame.id] -= 1
|
||||||
|
frame_counter = self._frame_counter[frame.id]
|
||||||
|
|
||||||
# Only push the StartFrame when all pipelines have processed it
|
# Only push the frame when all pipelines have processed it.
|
||||||
if start_frame_counter == 0:
|
if frame_counter == 0:
|
||||||
self._started = True
|
|
||||||
await self._start(frame)
|
|
||||||
await self._parallel_push_frame(frame, direction)
|
await self._parallel_push_frame(frame, direction)
|
||||||
|
await self.resume_processing_system_frames()
|
||||||
|
await self.resume_processing_frames()
|
||||||
else:
|
else:
|
||||||
if self._started:
|
await self._parallel_push_frame(frame, direction)
|
||||||
await self._parallel_push_frame(frame, direction)
|
|
||||||
else:
|
|
||||||
await self._down_queue.put(frame)
|
|
||||||
|
|
||||||
async def _process_up_queue(self):
|
|
||||||
"""Process upstream frames from all parallel branches."""
|
|
||||||
while True:
|
|
||||||
frame = await self._up_queue.get()
|
|
||||||
await self._parallel_push_frame(frame, FrameDirection.UPSTREAM)
|
|
||||||
self._up_queue.task_done()
|
|
||||||
|
|
||||||
async def _process_down_queue(self):
|
|
||||||
"""Process downstream frames with EndFrame coordination.
|
|
||||||
|
|
||||||
Coordinates EndFrames to ensure they are only pushed upstream once
|
|
||||||
all parallel branches have completed processing them.
|
|
||||||
"""
|
|
||||||
running = True
|
|
||||||
while running:
|
|
||||||
frame = await self._down_queue.get()
|
|
||||||
|
|
||||||
endframe_counter = self._endframe_counter.get(frame.id, 0)
|
|
||||||
|
|
||||||
# If we have a counter, decrement it.
|
|
||||||
if endframe_counter > 0:
|
|
||||||
self._endframe_counter[frame.id] -= 1
|
|
||||||
endframe_counter = self._endframe_counter[frame.id]
|
|
||||||
|
|
||||||
# If we don't have a counter or we reached 0, push the frame.
|
|
||||||
if endframe_counter == 0:
|
|
||||||
await self._parallel_push_frame(frame, FrameDirection.DOWNSTREAM)
|
|
||||||
|
|
||||||
running = not (endframe_counter == 0 and isinstance(frame, EndFrame))
|
|
||||||
|
|
||||||
self._down_queue.task_done()
|
|
||||||
|
|||||||
Reference in New Issue
Block a user