introduce heartbeat frames

This commit is contained in:
Aleix Conchillo Flaqué
2025-01-16 18:56:01 -08:00
parent 3e178fd46f
commit ff8aa68942
3 changed files with 165 additions and 36 deletions

View File

@@ -14,6 +14,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
pipeline. This can be useful, for example, to implement frame loggers or
debuggers among other things.
- Introduced heartbeat frames. The pipeline task can now push periodic
heartbeats down the pipeline when `enable_heartbeats=True`. Heartbeats are
system frames that are supposed to make it all the way to the end of the
pipeline. When a heartbeat frame is received the traversing time (i.e. the time
it took to go through the whole pipeline) will be displayed (with TRACE
logging) otherwise a warning will be shown.
- Added `30-observer.py` to show how to add an Observer to a pipeline for
debugging.

View File

@@ -424,6 +424,16 @@ class FatalErrorFrame(ErrorFrame):
fatal: bool = field(default=True, init=False)
@dataclass
class HeartbeatFrame(SystemFrame):
"""This frame is used by the pipeline task as a mechanism to know if the
pipeline is running properly.
"""
timestamp: int
@dataclass
class EndTaskFrame(SystemFrame):
"""This is used to notify the pipeline task that the pipeline should be

View File

@@ -19,6 +19,7 @@ from pipecat.frames.frames import (
EndTaskFrame,
ErrorFrame,
Frame,
HeartbeatFrame,
MetricsFrame,
StartFrame,
StopTaskFrame,
@@ -29,11 +30,14 @@ from pipecat.pipeline.base_pipeline import BasePipeline
from pipecat.processors.frame_processor import FrameDirection, FrameProcessor
from pipecat.utils.utils import obj_count, obj_id
HEARTBEAT_SECONDS = 1.0
class PipelineParams(BaseModel):
model_config = ConfigDict(arbitrary_types_allowed=True)
allow_interruptions: bool = False
enable_heartbeats: bool = False
enable_metrics: bool = False
enable_usage_metrics: bool = False
send_initial_empty_metrics: bool = True
@@ -58,25 +62,10 @@ class Source(FrameProcessor):
match direction:
case FrameDirection.UPSTREAM:
await self._handle_upstream_frame(frame)
await self._up_queue.put(frame)
case FrameDirection.DOWNSTREAM:
await self.push_frame(frame, direction)
async def _handle_upstream_frame(self, frame: Frame):
if isinstance(frame, EndTaskFrame):
# Tell the task we should end nicely.
await self._up_queue.put(EndTaskFrame())
elif isinstance(frame, CancelTaskFrame):
# Tell the task we should end right away.
await self._up_queue.put(CancelTaskFrame())
elif isinstance(frame, ErrorFrame):
logger.error(f"Error running app: {frame}")
if frame.fatal:
# Cancel all tasks downstream.
await self.push_frame(CancelFrame())
# Tell the task we should stop.
await self._up_queue.put(StopTaskFrame())
class Sink(FrameProcessor):
"""This is the sink processor that is linked at the end of the pipeline
@@ -91,10 +80,7 @@ class Sink(FrameProcessor):
async def process_frame(self, frame: Frame, direction: FrameDirection):
await super().process_frame(frame, direction)
# We really just want to know when the EndFrame reached the sink.
if isinstance(frame, EndFrame):
await self._down_queue.put(frame)
await self._down_queue.put(frame)
class Observer(BaseObserver):
@@ -135,9 +121,18 @@ class PipelineTask:
self._params = params
self._finished = False
# This queue receives frames coming from the pipeline upstream.
self._up_queue = asyncio.Queue()
# This queue receives frames coming from the pipeline downstream.
self._down_queue = asyncio.Queue()
# This queue is the queue used to push frames to the pipeline.
self._push_queue = asyncio.Queue()
# 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()
self._source = Source(self._up_queue)
self._source.link(pipeline)
@@ -148,33 +143,49 @@ class PipelineTask:
self._observer = Observer(params.observers)
def has_finished(self):
"""Indicates whether the tasks has finished. That is, all processors
have stopped.
"""
return self._finished
async def stop_when_done(self):
"""This is a helper function that sends an EndFrame to the pipeline in
order to stop the task after everything in it has been processed.
"""
logger.debug(f"Task {self} scheduled to stop when done")
await self.queue_frame(EndFrame())
async def cancel(self):
"""
Stops the running pipeline immediately.
"""
logger.debug(f"Canceling pipeline task {self}")
# Make sure everything is cleaned up downstream. This is sent
# out-of-band from the main streaming task which is what we want since
# we want to cancel right away.
await self._source.push_frame(CancelFrame())
self._process_push_task.cancel()
self._process_up_task.cancel()
await self._process_push_task
await self._process_up_task
await self._cancel_tasks(True)
async def run(self):
self._process_up_task = asyncio.create_task(self._process_up_queue())
self._process_push_task = asyncio.create_task(self._process_push_queue())
await asyncio.gather(self._process_up_task, self._process_push_task)
"""
Starts running the given pipeline.
"""
tasks = self._create_tasks()
await asyncio.gather(*tasks)
self._finished = True
async def queue_frame(self, frame: Frame):
"""
Queue a frame to be pushed down the pipeline.
"""
await self._push_queue.put(frame)
async def queue_frames(self, frames: Iterable[Frame] | AsyncIterable[Frame]):
"""
Queues multiple frames to be pushed down the pipeline.
"""
if isinstance(frames, AsyncIterable):
async for frame in frames:
await self.queue_frame(frame)
@@ -182,6 +193,39 @@ class PipelineTask:
for frame in frames:
await self.queue_frame(frame)
def _create_tasks(self):
tasks = []
self._process_up_task = asyncio.create_task(self._process_up_queue())
self._process_down_task = asyncio.create_task(self._process_down_queue())
self._process_push_task = asyncio.create_task(self._process_push_queue())
tasks = [self._process_up_task, self._process_down_task, self._process_push_task]
if self._params.enable_heartbeats:
self._heartbeat_push_task = asyncio.create_task(self._heartbeat_push_handler())
self._heartbeat_monitor_task = asyncio.create_task(self._heartbeat_monitor_handler())
tasks.append(self._heartbeat_push_task)
tasks.append(self._heartbeat_monitor_task)
return tasks
async def _cancel_tasks(self, cancel_push: bool):
if cancel_push:
self._process_push_task.cancel()
await self._process_push_task
self._process_up_task.cancel()
await self._process_up_task
self._process_down_task.cancel()
await self._process_down_task
if self._params.enable_heartbeats:
self._heartbeat_push_task.cancel()
await self._heartbeat_push_task
self._heartbeat_monitor_task.cancel()
await self._heartbeat_monitor_task
def _initial_metrics_frame(self) -> MetricsFrame:
processors = self._pipeline.processors_with_metrics()
data = []
@@ -190,7 +234,16 @@ class PipelineTask:
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 _process_push_queue(self):
"""This is the task that runs the pipeline for the first time by sending
a StartFrame and by pushing any other frames queued by the user. It runs
until the tasks is canceled or stopped (e.g. with an EndFrame).
"""
self._clock.start()
start_frame = StartFrame(
@@ -224,29 +277,88 @@ class PipelineTask:
await self._source.cleanup()
await self._pipeline.cleanup()
await self._sink.cleanup()
# We just enqueue None to terminate the task gracefully.
self._process_up_task.cancel()
await self._process_up_task
async def _wait_for_endframe(self):
# NOTE(aleix): the Sink element just pushes EndFrames to the down queue,
# so just wait for it. In the future we might do something else here,
# but for now this is fine.
await self._down_queue.get()
# Finally, cancel internal tasks. We don't cancel the push tasks because
# that's us.
await self._cancel_tasks(False)
async def _process_up_queue(self):
"""This is the task that processes frames coming upstream from the
pipeline. These frames might indicate, for example, that we want the
pipeline to be stopped (e.g. EndTaskFrame) in which case we would send
an EndFrame down the pipeline.
"""
while True:
try:
frame = await self._up_queue.get()
if isinstance(frame, EndTaskFrame):
# Tell the task we should end nicely.
await self.queue_frame(EndFrame())
elif isinstance(frame, CancelTaskFrame):
# Tell the task we should end right away.
await self.queue_frame(CancelFrame())
elif isinstance(frame, StopTaskFrame):
await self.queue_frame(StopTaskFrame())
elif isinstance(frame, ErrorFrame):
logger.error(f"Error running app: {frame}")
if frame.fatal:
# Cancel all tasks downstream.
await self.queue_frame(CancelFrame())
# Tell the task we should stop.
await self.queue_frame(StopTaskFrame())
self._up_queue.task_done()
except asyncio.CancelledError:
break
async def _process_down_queue(self):
"""This tasks process frames coming downstream from the pipeline. For
example, heartbeat frames or an EndFrame which would indicate all
processors have handled the EndFrame and therefore we can exit the task
cleanly.
"""
while True:
try:
frame = await self._down_queue.get()
if isinstance(frame, EndFrame):
self._endframe_event.set()
elif isinstance(frame, HeartbeatFrame):
await self._heartbeat_queue.put(frame)
self._down_queue.task_done()
except asyncio.CancelledError:
break
async def _heartbeat_push_handler(self):
"""
This tasks pushes a heartbeat frame every HEARTBEAT_SECONDS.
"""
while True:
try:
await self.queue_frame(HeartbeatFrame(timestamp=self._clock.get_time()))
await asyncio.sleep(HEARTBEAT_SECONDS)
except asyncio.CancelledError:
break
async def _heartbeat_monitor_handler(self):
"""This tasks monitors heartbeat frames. If a heartbeat frame has not
been received for a long period a warning will be logged. It also logs
the time that a heartbeat frame takes to processes, that is how long it
takes for the heartbeat frame to traverse all the pipeline.
"""
wait_time = HEARTBEAT_SECONDS * 2
while True:
try:
frame = await asyncio.wait_for(self._heartbeat_queue.get(), timeout=wait_time)
process_time = (self._clock.get_time() - frame.timestamp) / 1_000_000_000
logger.trace(f"{self}: heartbeat frame processed in {process_time} seconds")
self._heartbeat_queue.task_done()
except asyncio.TimeoutError:
logger.warning(
f"{self}: heartbeat frame not received for more than {wait_time} seconds"
)
except asyncio.CancelledError:
break
def __str__(self):
return self.name