wx44wx/pipecat
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

task: handle StopFrame and StopTaskFrame gracefully

Browse Source
This commit is contained in:
Aleix Conchillo Flaqué
2025-02-24 21:32:58 -08:00
parent 68789dfcf0
commit 376d969a77
1 changed files with 24 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

44
src/pipecat/pipeline/task.py
View File

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