Flush buffered frames before pushing the synchronization frame so downstream processors see the buffered frames first. Switch to a while-loop with pop(0) so frames added to the buffer during flush are also drained.
212 lines
8.2 KiB
Python
212 lines
8.2 KiB
Python
#
|
|
# Copyright (c) 2024-2026, Daily
|
|
#
|
|
# SPDX-License-Identifier: BSD 2-Clause License
|
|
#
|
|
|
|
"""Parallel pipeline implementation for concurrent frame processing.
|
|
|
|
This module provides a parallel pipeline that processes frames through multiple
|
|
sub-pipelines concurrently, with coordination for system frames and proper
|
|
handling of pipeline lifecycle events.
|
|
"""
|
|
|
|
from itertools import chain
|
|
from typing import Dict, List
|
|
|
|
from loguru import logger
|
|
|
|
from pipecat.frames.frames import CancelFrame, EndFrame, Frame, StartFrame
|
|
from pipecat.pipeline.base_pipeline import BasePipeline
|
|
from pipecat.pipeline.pipeline import Pipeline, PipelineSink, PipelineSource
|
|
from pipecat.processors.frame_processor import FrameDirection, FrameProcessor, FrameProcessorSetup
|
|
|
|
|
|
class ParallelPipeline(BasePipeline):
|
|
"""Pipeline that processes frames through multiple sub-pipelines concurrently.
|
|
|
|
Creates multiple parallel processing branches from the provided processor lists,
|
|
coordinating frame flow and ensuring proper synchronization of lifecycle events
|
|
like EndFrames. Each branch runs independently while system frames are handled
|
|
specially to maintain pipeline coordination.
|
|
"""
|
|
|
|
def __init__(self, *args):
|
|
"""Initialize the parallel pipeline with processor lists.
|
|
|
|
Args:
|
|
*args: Variable number of processor lists, each becoming a parallel branch.
|
|
|
|
Raises:
|
|
Exception: If no processor lists are provided.
|
|
TypeError: If any argument is not a list of processors.
|
|
"""
|
|
# We don't set it to direct mode because we use frame pausing and that
|
|
# requires queues.
|
|
super().__init__()
|
|
|
|
if len(args) == 0:
|
|
raise Exception(f"ParallelPipeline needs at least one argument")
|
|
|
|
self._pipelines = []
|
|
|
|
self._seen_ids = set()
|
|
self._frame_counter: Dict[int, int] = {}
|
|
self._synchronizing: bool = False
|
|
self._buffered_frames: list[tuple[Frame, FrameDirection]] = []
|
|
|
|
logger.debug(f"Creating {self} pipelines")
|
|
|
|
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
|
|
#
|
|
|
|
@property
|
|
def processors(self):
|
|
"""Return the list of sub-processors contained within this processor.
|
|
|
|
Only compound processors (e.g. pipelines and parallel pipelines) have
|
|
sub-processors. Non-compound processors will return an empty list.
|
|
|
|
Returns:
|
|
The list of sub-processors if this is a compound processor.
|
|
"""
|
|
return self._pipelines
|
|
|
|
@property
|
|
def entry_processors(self) -> List["FrameProcessor"]:
|
|
"""Return the list of entry processors for this processor.
|
|
|
|
Entry processors are the first processors in a compound processor
|
|
(e.g. pipelines, parallel pipelines). Note that pipelines can also be an
|
|
entry processor as pipelines are processors themselves. Non-compound
|
|
processors will simply return an empty list.
|
|
|
|
Returns:
|
|
The list of entry processors.
|
|
"""
|
|
return self._pipelines
|
|
|
|
def processors_with_metrics(self) -> List[FrameProcessor]:
|
|
"""Collect processors that can generate metrics from all parallel branches.
|
|
|
|
Returns:
|
|
List of frame processors that support metrics collection from all branches.
|
|
"""
|
|
return list(chain.from_iterable(p.processors_with_metrics() for p in self._pipelines))
|
|
|
|
async def setup(self, setup: FrameProcessorSetup):
|
|
"""Set up the parallel pipeline and all its branches.
|
|
|
|
Args:
|
|
setup: Configuration for frame processor setup.
|
|
|
|
Raises:
|
|
TypeError: If any processor list argument is not actually a list.
|
|
"""
|
|
await super().setup(setup)
|
|
for p in self._pipelines:
|
|
await p.setup(setup)
|
|
|
|
async def cleanup(self):
|
|
"""Clean up the parallel pipeline and all its branches."""
|
|
await super().cleanup()
|
|
for p in self._pipelines:
|
|
await p.cleanup()
|
|
|
|
async def process_frame(self, frame: Frame, direction: FrameDirection):
|
|
"""Process frames through all parallel branches with lifecycle coordination.
|
|
|
|
Args:
|
|
frame: The frame to process.
|
|
direction: The direction of frame flow.
|
|
"""
|
|
await super().process_frame(frame, direction)
|
|
|
|
# Parallel pipeline synchronized frames.
|
|
#
|
|
# - StartFrame: If a fast branch completes first, processors in
|
|
# other branches that haven't received StartFrame yet could
|
|
# receive other frames before it, causing errors.
|
|
#
|
|
# - EndFrame: If EndFrame escapes from a fast branch, downstream
|
|
# processors (e.g. output transport) begin shutting down while
|
|
# other branches still have frames to flush, causing lost output.
|
|
#
|
|
# - CancelFrame: PipelineTask waits for CancelFrame to reach the
|
|
# pipeline sink. If it escapes from a fast branch while slower
|
|
# branches are still running, the task considers cancellation
|
|
# complete prematurely.
|
|
if isinstance(frame, (StartFrame, EndFrame, CancelFrame)):
|
|
self._frame_counter[frame.id] = len(self._pipelines)
|
|
self._synchronizing = True
|
|
await self.pause_processing_system_frames()
|
|
await self.pause_processing_frames()
|
|
|
|
# Process frames in each of the sub-pipelines.
|
|
for p in self._pipelines:
|
|
await p.queue_frame(frame, direction)
|
|
|
|
async def _parallel_push_frame(self, frame: Frame, direction: FrameDirection):
|
|
"""Push frames while avoiding duplicates using frame ID tracking.
|
|
|
|
During lifecycle frame synchronization, non-lifecycle frames are buffered
|
|
to prevent them from escaping the parallel pipeline before all branches
|
|
have finished processing the lifecycle frame.
|
|
"""
|
|
if frame.id not in self._seen_ids:
|
|
self._seen_ids.add(frame.id)
|
|
if self._synchronizing:
|
|
self._buffered_frames.append((frame, direction))
|
|
else:
|
|
await self.push_frame(frame, direction)
|
|
|
|
async def _pipeline_sink_push_frame(self, frame: Frame, direction: FrameDirection):
|
|
# Parallel pipeline synchronized frames.
|
|
if isinstance(frame, (StartFrame, EndFrame, CancelFrame)):
|
|
# Decrement counter.
|
|
frame_counter = self._frame_counter.get(frame.id, 0)
|
|
if frame_counter > 0:
|
|
self._frame_counter[frame.id] -= 1
|
|
frame_counter = self._frame_counter[frame.id]
|
|
|
|
# Only push the frame when all pipelines have processed it.
|
|
if frame_counter == 0:
|
|
self._synchronizing = False
|
|
# StartFrame should always go before any other frame.
|
|
if isinstance(frame, StartFrame):
|
|
await self._parallel_push_frame(frame, direction)
|
|
await self._flush_buffered_frames()
|
|
else:
|
|
await self._flush_buffered_frames()
|
|
await self._parallel_push_frame(frame, direction)
|
|
await self.resume_processing_system_frames()
|
|
await self.resume_processing_frames()
|
|
else:
|
|
await self._parallel_push_frame(frame, direction)
|
|
|
|
async def _flush_buffered_frames(self):
|
|
"""Flush frames that were buffered during lifecycle frame synchronization."""
|
|
while len(self._buffered_frames) > 0:
|
|
frame, direction = self._buffered_frames.pop(0)
|
|
await self.push_frame(frame, direction)
|