diff --git a/src/pipecat/pipeline/task.py b/src/pipecat/pipeline/task.py index 873fc60d6..a2aad1e93 100644 --- a/src/pipecat/pipeline/task.py +++ b/src/pipecat/pipeline/task.py @@ -22,6 +22,7 @@ from pipecat.frames.frames import ( HeartbeatFrame, MetricsFrame, StartFrame, + StopFrame, StopTaskFrame, ) from pipecat.metrics.metrics import ProcessingMetricsData, TTFBMetricsData @@ -163,9 +164,9 @@ class PipelineTask(BaseTask): # This is the heartbeat queue. When a heartbeat frame is received in the # down queue we add it to the heartbeat queue for processing. self._heartbeat_queue = asyncio.Queue() - # This event is used to indicate an EndFrame has been received in the - # down queue. - self._endframe_event = asyncio.Event() + # This event is used to indicate a finalize frame (e.g. EndFrame, + # StopFrame) has been received in the down queue. + self._pipeline_end_event = asyncio.Event() self._source = PipelineTaskSource(self._up_queue) self._source.link(pipeline) @@ -224,9 +225,12 @@ class PipelineTask(BaseTask): """Starts and manages the pipeline execution until completion or cancellation.""" if self.has_finished(): return + cleanup_pipeline = True try: push_task = await self._create_tasks() await self._task_manager.wait_for_task(push_task) + # We have already cleaned up the pipeline inside the task. + cleanup_pipeline = False except asyncio.CancelledError: # We are awaiting on the push task and it might be cancelled # (e.g. Ctrl-C). This means we will get a CancelledError here as @@ -234,7 +238,7 @@ class PipelineTask(BaseTask): # awaiting a task. pass await self._cancel_tasks() - await self._cleanup() + await self._cleanup(cleanup_pipeline) if self._check_dangling_tasks: self._print_dangling_tasks() self._finished = True @@ -305,13 +309,14 @@ class PipelineTask(BaseTask): data.append(ProcessingMetricsData(processor=p.name, value=0.0)) return MetricsFrame(data=data) - async def _wait_for_endframe(self): - await self._endframe_event.wait() - self._endframe_event.clear() + async def _wait_for_pipeline_end(self): + await self._pipeline_end_event.wait() + self._pipeline_end_event.clear() - async def _cleanup(self): + async def _cleanup(self, cleanup_pipeline: bool): await self._source.cleanup() - await self._pipeline.cleanup() + if cleanup_pipeline: + await self._pipeline.cleanup() await self._sink.cleanup() async def _process_push_queue(self): @@ -342,18 +347,16 @@ class PipelineTask(BaseTask): await self._source.queue_frame(self._initial_metrics_frame(), FrameDirection.DOWNSTREAM) running = True - should_cleanup = True + cleanup_pipeline = True while running: frame = await self._push_queue.get() await self._source.queue_frame(frame, FrameDirection.DOWNSTREAM) - if isinstance(frame, EndFrame): - await self._wait_for_endframe() - running = not isinstance(frame, (CancelFrame, EndFrame, StopTaskFrame)) - should_cleanup = not isinstance(frame, StopTaskFrame) + if isinstance(frame, (EndFrame, StopFrame)): + await self._wait_for_pipeline_end() + running = not isinstance(frame, (CancelFrame, EndFrame, StopFrame)) + cleanup_pipeline = not isinstance(frame, StopFrame) self._push_queue.task_done() - # Cleanup only if we need to. - if should_cleanup: - await self._cleanup() + await self._cleanup(cleanup_pipeline) async def _process_up_queue(self): """This is the task that processes frames coming upstream from the @@ -371,7 +374,8 @@ class PipelineTask(BaseTask): # Tell the task we should end right away. await self.queue_frame(CancelFrame()) elif isinstance(frame, StopTaskFrame): - await self.queue_frame(StopTaskFrame()) + # Tell the task we should stop nicely. + await self.queue_frame(StopFrame()) elif isinstance(frame, ErrorFrame): logger.error(f"Error running app: {frame}") if frame.fatal: @@ -390,8 +394,8 @@ class PipelineTask(BaseTask): """ while True: frame = await self._down_queue.get() - if isinstance(frame, EndFrame): - self._endframe_event.set() + if isinstance(frame, (EndFrame, StopFrame)): + self._pipeline_end_event.set() elif isinstance(frame, HeartbeatFrame): await self._heartbeat_queue.put(frame) self._down_queue.task_done()