diff --git a/CHANGELOG.md b/CHANGELOG.md index f153dda13..33064fa80 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,22 +12,24 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Added logging and improved error handling to help diagnose and prevent potential Pipeline freezes. +- Added `WatchdogQueue`, `WatchdogPriorityQueue`, `WatchdogEvent` and + `WatchdogAsyncIterator`. These helper utilities reset watchdog timers + appropriately before they expire. When watchdog timers are disabled, the + utilities behave as standard counterparts without side effects. + - Introduce task watchdog timers. Watchdog timers are used to detect if a - Pipecat task is taking longer than expected (by default 5 seconds). It is + Pipecat task is taking longer than expected (by default 5 seconds). Watchdog + timers are disabled by default and can be enabled globally by passing + `enable_watchdog_timers` argument to `PipelineTask` constructor. It is possible to change the default watchdog timer timeout by using the - `watchdog_timeout` constructor argument when creating a `PipelineTask`. With - watchdog timers it is also possible to log how long each processing step is - taking (e.g. processing an element from a queue inside a task). This is done - with the `enable_watchdog_logging` constructor argument when creating a - `PipelineTask.` It is also possible to control these two values per each frame - processor. That is, you can set set `enable_watchdog_logging` and + `watchdog_timeout` argument. You can also log how long it takes to reset the + watchdog timers which is done with the `enable_watchdog_logging`. You can + control all these settings per each frame processor or even per task. That is, + you can set `enable_watchdog_timers`, `enable_watchdog_logging` and `watchdog_timeout` when creating any frame processor through their constructor - arguments. Finally, you can also set these values per task. So, if you are - writing a frame processor that creates multiple tasks and you only want to - enable logging for one of them, you can do so by passing the same argument - names to the `FrameProcessor.create_task()` function. Note that watchdog - timers only work with Pipecat tasks but not if you use `asycio.create_task()` - or similar. + arguments or when you create a task with `FrameProcessor.create_task()`. Note + that watchdog timers only work with Pipecat tasks and will not work if you use + `asycio.create_task()` or similar. - Added `lexicon_names` parameter to `AWSPollyTTSService.InputParams`. diff --git a/src/pipecat/frames/frames.py b/src/pipecat/frames/frames.py index 4b2d934ab..3a602a3e2 100644 --- a/src/pipecat/frames/frames.py +++ b/src/pipecat/frames/frames.py @@ -453,8 +453,8 @@ class StartFrame(SystemFrame): allow_interruptions: bool = False enable_metrics: bool = False enable_usage_metrics: bool = False - report_only_initial_ttfb: bool = False interruption_strategies: List[BaseInterruptionStrategy] = field(default_factory=list) + report_only_initial_ttfb: bool = False @dataclass diff --git a/src/pipecat/pipeline/parallel_pipeline.py b/src/pipecat/pipeline/parallel_pipeline.py index 3a08f44ed..f6ac78827 100644 --- a/src/pipecat/pipeline/parallel_pipeline.py +++ b/src/pipecat/pipeline/parallel_pipeline.py @@ -21,6 +21,7 @@ from pipecat.frames.frames import ( from pipecat.pipeline.base_pipeline import BasePipeline from pipecat.pipeline.pipeline import Pipeline from pipecat.processors.frame_processor import FrameDirection, FrameProcessor, FrameProcessorSetup +from pipecat.utils.watchdog_queue import WatchdogQueue class ParallelPipelineSource(FrameProcessor): @@ -76,20 +77,36 @@ class ParallelPipeline(BasePipeline): if len(args) == 0: raise Exception(f"ParallelPipeline needs at least one argument") + self._args = args self._sources = [] self._sinks = [] + self._pipelines = [] + self._seen_ids = set() self._endframe_counter: Dict[int, int] = {} self._up_task = None self._down_task = None - self._up_queue = asyncio.Queue() - self._down_queue = asyncio.Queue() - self._pipelines = [] + # + # BasePipeline + # + + def processors_with_metrics(self) -> List[FrameProcessor]: + return list(chain.from_iterable(p.processors_with_metrics() for p in self._pipelines)) + + # + # Frame processor + # + + async def setup(self, setup: FrameProcessorSetup): + await super().setup(setup) + + self._up_queue = WatchdogQueue(self, watchdog_enabled=setup.watchdog_timers_enabled) + self._down_queue = WatchdogQueue(self, watchdog_enabled=setup.watchdog_timers_enabled) logger.debug(f"Creating {self} pipelines") - for processors in args: + for processors in self._args: if not isinstance(processors, list): raise TypeError(f"ParallelPipeline argument {processors} is not a list") @@ -107,19 +124,6 @@ class ParallelPipeline(BasePipeline): logger.debug(f"Finished creating {self} pipelines") - # - # BasePipeline - # - - def processors_with_metrics(self) -> List[FrameProcessor]: - return list(chain.from_iterable(p.processors_with_metrics() for p in self._pipelines)) - - # - # Frame processor - # - - async def setup(self, setup: FrameProcessorSetup): - await super().setup(setup) await asyncio.gather(*[s.setup(setup) for s in self._sources]) await asyncio.gather(*[p.setup(setup) for p in self._pipelines]) await asyncio.gather(*[s.setup(setup) for s in self._sinks]) @@ -134,7 +138,7 @@ class ParallelPipeline(BasePipeline): await super().process_frame(frame, direction) if isinstance(frame, StartFrame): - await self._start() + await self._start(frame) elif isinstance(frame, EndFrame): self._endframe_counter[frame.id] = len(self._pipelines) elif isinstance(frame, CancelFrame): @@ -154,7 +158,7 @@ class ParallelPipeline(BasePipeline): elif isinstance(frame, EndFrame): await self._stop() - async def _start(self): + async def _start(self, frame: StartFrame): await self._create_tasks() async def _stop(self): @@ -202,18 +206,14 @@ class ParallelPipeline(BasePipeline): async def _process_up_queue(self): while True: frame = await self._up_queue.get() - self.start_watchdog() await self._parallel_push_frame(frame, FrameDirection.UPSTREAM) self._up_queue.task_done() - self.reset_watchdog() async def _process_down_queue(self): running = True while running: frame = await self._down_queue.get() - self.start_watchdog() - endframe_counter = self._endframe_counter.get(frame.id, 0) # If we have a counter, decrement it. @@ -228,5 +228,3 @@ class ParallelPipeline(BasePipeline): running = not (endframe_counter == 0 and isinstance(frame, EndFrame)) self._down_queue.task_done() - - self.reset_watchdog() diff --git a/src/pipecat/pipeline/sync_parallel_pipeline.py b/src/pipecat/pipeline/sync_parallel_pipeline.py index 4cf9f5033..f78ca0de3 100644 --- a/src/pipecat/pipeline/sync_parallel_pipeline.py +++ b/src/pipecat/pipeline/sync_parallel_pipeline.py @@ -15,6 +15,7 @@ from pipecat.frames.frames import ControlFrame, EndFrame, Frame, SystemFrame from pipecat.pipeline.base_pipeline import BasePipeline from pipecat.pipeline.pipeline import Pipeline from pipecat.processors.frame_processor import FrameDirection, FrameProcessor, FrameProcessorSetup +from pipecat.utils.watchdog_queue import WatchdogQueue @dataclass @@ -61,15 +62,30 @@ class SyncParallelPipeline(BasePipeline): if len(args) == 0: raise Exception(f"SyncParallelPipeline needs at least one argument") + self._args = args self._sinks = [] self._sources = [] self._pipelines = [] - self._up_queue = asyncio.Queue() - self._down_queue = asyncio.Queue() + # + # BasePipeline + # + + def processors_with_metrics(self) -> List[FrameProcessor]: + return list(chain.from_iterable(p.processors_with_metrics() for p in self._pipelines)) + + # + # Frame processor + # + + async def setup(self, setup: FrameProcessorSetup): + await super().setup(setup) + + self._up_queue = WatchdogQueue(self, watchdog_enabled=setup.watchdog_timers_enabled) + self._down_queue = WatchdogQueue(self, watchdog_enabled=setup.watchdog_timers_enabled) logger.debug(f"Creating {self} pipelines") - for processors in args: + for processors in self._args: if not isinstance(processors, list): raise TypeError(f"SyncParallelPipeline argument {processors} is not a list") @@ -92,19 +108,6 @@ class SyncParallelPipeline(BasePipeline): logger.debug(f"Finished creating {self} pipelines") - # - # BasePipeline - # - - def processors_with_metrics(self) -> List[FrameProcessor]: - return list(chain.from_iterable(p.processors_with_metrics() for p in self._pipelines)) - - # - # Frame processor - # - - async def setup(self, setup: FrameProcessorSetup): - await super().setup(setup) await asyncio.gather(*[s["processor"].setup(setup) for s in self._sources]) await asyncio.gather(*[p.setup(setup) for p in self._pipelines]) await asyncio.gather(*[s["processor"].setup(setup) for s in self._sinks]) diff --git a/src/pipecat/pipeline/task.py b/src/pipecat/pipeline/task.py index 88a47189b..68b76401e 100644 --- a/src/pipecat/pipeline/task.py +++ b/src/pipecat/pipeline/task.py @@ -41,6 +41,8 @@ from pipecat.processors.frame_processor import FrameDirection, FrameProcessor, F from pipecat.utils.asyncio import WATCHDOG_TIMEOUT, BaseTaskManager, TaskManager, TaskManagerParams from pipecat.utils.tracing.setup import is_tracing_available from pipecat.utils.tracing.turn_trace_observer import TurnTraceObserver +from pipecat.utils.watchdog_queue import WatchdogQueue +from pipecat.utils.watchdog_reseter import WatchdogReseter HEARTBEAT_SECONDS = 1.0 HEARTBEAT_MONITOR_SECONDS = HEARTBEAT_SECONDS * 10 @@ -131,7 +133,7 @@ class PipelineTaskSink(FrameProcessor): await self._down_queue.put(frame) -class PipelineTask(BasePipelineTask): +class PipelineTask(WatchdogReseter, BasePipelineTask): """Manages the execution of a pipeline, handling frame processing and task lifecycle. It has a couple of event handlers `on_frame_reached_upstream` and @@ -188,6 +190,7 @@ class PipelineTask(BasePipelineTask): enable_tracing: Whether to enable tracing. enable_turn_tracking: Whether to enable turn tracking. enable_watchdog_logging: Whether to print task processing times. + enable_watchdog_timers: Whether to enable task watchdog timers. idle_timeout_frames: A tuple with the frames that should trigger an idle timeout if not received withing `idle_timeout_seconds`. idle_timeout_secs: Timeout (in seconds) to consider pipeline idle or @@ -211,6 +214,7 @@ class PipelineTask(BasePipelineTask): enable_tracing: bool = False, enable_turn_tracking: bool = True, enable_watchdog_logging: bool = False, + enable_watchdog_timers: bool = False, idle_timeout_frames: Tuple[Type[Frame], ...] = ( BotSpeakingFrame, LLMFullResponseEndFrame, @@ -231,6 +235,7 @@ class PipelineTask(BasePipelineTask): self._enable_tracing = enable_tracing and is_tracing_available() self._enable_turn_tracking = enable_turn_tracking self._enable_watchdog_logging = enable_watchdog_logging + self._enable_watchdog_timers = enable_watchdog_timers self._idle_timeout_frames = idle_timeout_frames self._idle_timeout_secs = idle_timeout_secs self._watchdog_timeout_secs = watchdog_timeout_secs @@ -261,18 +266,24 @@ class PipelineTask(BasePipelineTask): self._cancelled = False # This queue receives frames coming from the pipeline upstream. - self._up_queue = asyncio.Queue() + self._up_queue = WatchdogQueue(self, watchdog_enabled=enable_watchdog_timers) + self._process_up_task: Optional[asyncio.Task] = None # This queue receives frames coming from the pipeline downstream. - self._down_queue = asyncio.Queue() + self._down_queue = WatchdogQueue(self, watchdog_enabled=enable_watchdog_timers) + self._process_down_task: Optional[asyncio.Task] = None # This queue is the queue used to push frames to the pipeline. - self._push_queue = asyncio.Queue() + self._push_queue = WatchdogQueue(self, watchdog_enabled=enable_watchdog_timers) + self._process_push_task: Optional[asyncio.Task] = None # 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() + self._heartbeat_queue = WatchdogQueue(self, watchdog_enabled=enable_watchdog_timers) + self._heartbeat_push_task: Optional[asyncio.Task] = None + self._heartbeat_monitor_task: Optional[asyncio.Task] = None # This is the idle queue. When frames are received downstream they are # put in the queue. If no frame is received the pipeline is considered # idle. - self._idle_queue = asyncio.Queue() + self._idle_queue = WatchdogQueue(self, watchdog_enabled=enable_watchdog_timers) + self._idle_monitor_task: Optional[asyncio.Task] = None # 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() @@ -424,6 +435,9 @@ class PipelineTask(BasePipelineTask): for frame in frames: await self.queue_frame(frame) + def reset_watchdog(self): + self._task_manager.reset_watchdog(asyncio.current_task()) + async def _cancel(self): if not self._cancelled: logger.debug(f"Canceling pipeline task {self}") @@ -433,7 +447,9 @@ class PipelineTask(BasePipelineTask): # we want to cancel right away. await self._source.push_frame(CancelFrame()) # Only cancel the push task. Everything else will be cancelled in run(). - await self._task_manager.cancel_task(self._process_push_task) + if self._process_push_task: + await self._task_manager.cancel_task(self._process_push_task) + self._process_push_task = None async def _create_tasks(self): self._process_up_task = self._task_manager.create_task( @@ -446,7 +462,7 @@ class PipelineTask(BasePipelineTask): self._process_push_queue(), f"{self}::_process_push_queue" ) - await self._observer.start() + await self._observer.start(self._enable_watchdog_timers) return self._process_push_task @@ -468,20 +484,33 @@ class PipelineTask(BasePipelineTask): async def _cancel_tasks(self): await self._observer.stop() - await self._task_manager.cancel_task(self._process_up_task) - await self._task_manager.cancel_task(self._process_down_task) + if self._process_up_task: + await self._task_manager.cancel_task(self._process_up_task) + self._process_up_task = None + + if self._process_down_task: + await self._task_manager.cancel_task(self._process_down_task) + self._process_down_task = None await self._maybe_cancel_heartbeat_tasks() await self._maybe_cancel_idle_task() async def _maybe_cancel_heartbeat_tasks(self): - if self._params.enable_heartbeats: + if not self._params.enable_heartbeats: + return + + if self._heartbeat_push_task: await self._task_manager.cancel_task(self._heartbeat_push_task) + self._heartbeat_push_task = None + + if self._heartbeat_monitor_task: await self._task_manager.cancel_task(self._heartbeat_monitor_task) + self._heartbeat_monitor_task = None async def _maybe_cancel_idle_task(self): - if self._idle_timeout_secs: + if self._idle_timeout_secs and self._idle_monitor_task: await self._task_manager.cancel_task(self._idle_monitor_task) + self._idle_monitor_task = None def _initial_metrics_frame(self) -> MetricsFrame: processors = self._pipeline.processors_with_metrics() @@ -499,6 +528,7 @@ class PipelineTask(BasePipelineTask): mgr_params = TaskManagerParams( loop=params.loop, enable_watchdog_logging=self._enable_watchdog_logging, + enable_watchdog_timers=self._enable_watchdog_timers, watchdog_timeout=self._watchdog_timeout_secs, ) self._task_manager.setup(mgr_params) @@ -507,6 +537,7 @@ class PipelineTask(BasePipelineTask): clock=self._clock, task_manager=self._task_manager, observer=self._observer, + watchdog_timers_enabled=self._enable_watchdog_timers, ) await self._source.setup(setup) await self._pipeline.setup(setup) @@ -526,8 +557,6 @@ class PipelineTask(BasePipelineTask): await self._pipeline.cleanup() await self._sink.cleanup() - await self._task_manager.cleanup() - 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 diff --git a/src/pipecat/pipeline/task_observer.py b/src/pipecat/pipeline/task_observer.py index 950ddc43b..40d4ef3ed 100644 --- a/src/pipecat/pipeline/task_observer.py +++ b/src/pipecat/pipeline/task_observer.py @@ -12,6 +12,8 @@ from attr import dataclass from pipecat.observers.base_observer import BaseObserver, FramePushed from pipecat.utils.asyncio import BaseTaskManager +from pipecat.utils.watchdog_queue import WatchdogQueue +from pipecat.utils.watchdog_reseter import WatchdogReseter @dataclass @@ -26,7 +28,7 @@ class Proxy: observer: BaseObserver -class TaskObserver(BaseObserver): +class TaskObserver(WatchdogReseter, BaseObserver): """This is a pipeline frame observer that is meant to be used as a proxy to the user provided observers. That is, this is the observer that should be passed to the frame processors. Then, every time a frame is pushed this @@ -52,6 +54,7 @@ class TaskObserver(BaseObserver): self._proxies: Optional[Dict[BaseObserver, Proxy]] = ( None # Becomes a dict after start() is called ) + self._watchdog_timers_enabled = False def add_observer(self, observer: BaseObserver): # Add the observer to the list. @@ -76,12 +79,16 @@ class TaskObserver(BaseObserver): if observer in self._observers: self._observers.remove(observer) - async def start(self): + async def start(self, watchdog_timers_enabled: bool = False): """Starts all proxy observer tasks.""" + self._watchdog_timers_enabled = watchdog_timers_enabled self._proxies = self._create_proxies(self._observers) async def stop(self): """Stops all proxy observer tasks.""" + if not self._proxies: + return + for proxy in self._proxies.values(): await self._task_manager.cancel_task(proxy.task) @@ -89,11 +96,14 @@ class TaskObserver(BaseObserver): for proxy in self._proxies.values(): await proxy.queue.put(data) + def reset_watchdog(self): + self._task_manager.reset_watchdog(asyncio.current_task()) + def _started(self) -> bool: return self._proxies is not None def _create_proxy(self, observer: BaseObserver) -> Proxy: - queue = asyncio.Queue() + queue = WatchdogQueue(self, watchdog_enabled=self._watchdog_timers_enabled) task = self._task_manager.create_task( self._proxy_task_handler(queue, observer), f"TaskObserver::{observer}::_proxy_task_handler", diff --git a/src/pipecat/processors/aggregators/dtmf_aggregator.py b/src/pipecat/processors/aggregators/dtmf_aggregator.py index cc6218a1f..3008fb398 100644 --- a/src/pipecat/processors/aggregators/dtmf_aggregator.py +++ b/src/pipecat/processors/aggregators/dtmf_aggregator.py @@ -119,6 +119,7 @@ class DTMFAggregator(FrameProcessor): await asyncio.wait_for(self._digit_event.wait(), timeout=self._idle_timeout) self._digit_event.clear() except asyncio.TimeoutError: + self.reset_watchdog() if self._aggregation: await self._flush_aggregation() diff --git a/src/pipecat/processors/aggregators/llm_response.py b/src/pipecat/processors/aggregators/llm_response.py index 1984e3cf8..40016aaa0 100644 --- a/src/pipecat/processors/aggregators/llm_response.py +++ b/src/pipecat/processors/aggregators/llm_response.py @@ -470,6 +470,7 @@ class LLMUserContextAggregator(LLMContextResponseAggregator): ) self._emulating_vad = False finally: + self.reset_watchdog() self._aggregation_event.clear() async def _maybe_emulate_user_speaking(self): diff --git a/src/pipecat/processors/consumer_processor.py b/src/pipecat/processors/consumer_processor.py index 121dd7712..10cae11a3 100644 --- a/src/pipecat/processors/consumer_processor.py +++ b/src/pipecat/processors/consumer_processor.py @@ -10,6 +10,7 @@ from typing import Awaitable, Callable, Optional from pipecat.frames.frames import CancelFrame, EndFrame, Frame, StartFrame from pipecat.processors.frame_processor import FrameDirection, FrameProcessor from pipecat.processors.producer_processor import ProducerProcessor, identity_transformer +from pipecat.utils.watchdog_queue import WatchdogQueue class ConsumerProcessor(FrameProcessor): @@ -31,7 +32,7 @@ class ConsumerProcessor(FrameProcessor): super().__init__(**kwargs) self._transformer = transformer self._direction = direction - self._queue: asyncio.Queue = producer.add_consumer() + self._queue: WatchdogQueue = producer.add_consumer(self) self._consumer_task: Optional[asyncio.Task] = None async def process_frame(self, frame: Frame, direction: FrameDirection): @@ -61,7 +62,5 @@ class ConsumerProcessor(FrameProcessor): async def _consumer_task_handler(self): while True: frame = await self._queue.get() - self.start_watchdog() new_frame = await self._transformer(frame) await self.push_frame(new_frame, self._direction) - self.reset_watchdog() diff --git a/src/pipecat/processors/frame_processor.py b/src/pipecat/processors/frame_processor.py index bee5ce91c..134a69da7 100644 --- a/src/pipecat/processors/frame_processor.py +++ b/src/pipecat/processors/frame_processor.py @@ -31,6 +31,9 @@ from pipecat.observers.base_observer import BaseObserver, FramePushed from pipecat.processors.metrics.frame_processor_metrics import FrameProcessorMetrics from pipecat.utils.asyncio import BaseTaskManager from pipecat.utils.base_object import BaseObject +from pipecat.utils.watchdog_event import WatchdogEvent +from pipecat.utils.watchdog_queue import WatchdogQueue +from pipecat.utils.watchdog_reseter import WatchdogReseter class FrameDirection(Enum): @@ -43,15 +46,17 @@ class FrameProcessorSetup: clock: BaseClock task_manager: BaseTaskManager observer: Optional[BaseObserver] = None + watchdog_timers_enabled: bool = False -class FrameProcessor(BaseObject): +class FrameProcessor(WatchdogReseter, BaseObject): def __init__( self, *, name: Optional[str] = None, - metrics: Optional[FrameProcessorMetrics] = None, enable_watchdog_logging: Optional[bool] = None, + enable_watchdog_timers: Optional[bool] = None, + metrics: Optional[FrameProcessorMetrics] = None, watchdog_timeout_secs: Optional[float] = None, **kwargs, ): @@ -60,11 +65,14 @@ class FrameProcessor(BaseObject): self._prev: Optional["FrameProcessor"] = None self._next: Optional["FrameProcessor"] = None + # Enable watchdog timers for all tasks created by this frame processor. + self._enable_watchdog_timers = enable_watchdog_timers + # Enable watchdog logging for all tasks created by this frame processor. self._enable_watchdog_logging = enable_watchdog_logging # Allow this frame processor to control their tasks timeout. - self._watchdog_timeout = watchdog_timeout_secs + self._watchdog_timeout_secs = watchdog_timeout_secs # Clock self._clock: Optional[BaseClock] = None @@ -81,6 +89,7 @@ class FrameProcessor(BaseObject): self._enable_usage_metrics = False self._report_only_initial_ttfb = False self._interruption_strategies: List[BaseInterruptionStrategy] = [] + self._watchdog_timers_enabled = False # Indicates whether we have received the StartFrame. self.__started = False @@ -101,7 +110,7 @@ class FrameProcessor(BaseObject): # is called. To resume processing frames we need to call # `resume_processing_frames()` which will wake up the event. self.__should_block_frames = False - self.__input_event = asyncio.Event() + self.__input_event = None self.__input_frame_task: Optional[asyncio.Task] = None # Every processor in Pipecat should only output frames from a single @@ -137,6 +146,10 @@ class FrameProcessor(BaseObject): def interruption_strategies(self) -> Sequence[BaseInterruptionStrategy]: return self._interruption_strategies + @property + def watchdog_timers_enabled(self): + return self._watchdog_timers_enabled + def can_generate_metrics(self) -> bool: return False @@ -185,6 +198,7 @@ class FrameProcessor(BaseObject): name: Optional[str] = None, *, enable_watchdog_logging: Optional[bool] = None, + enable_watchdog_timers: Optional[bool] = None, watchdog_timeout_secs: Optional[float] = None, ) -> asyncio.Task: if name: @@ -199,8 +213,11 @@ class FrameProcessor(BaseObject): if enable_watchdog_logging else self._enable_watchdog_logging ), + enable_watchdog_timers=( + enable_watchdog_timers if enable_watchdog_timers else self.watchdog_timers_enabled + ), watchdog_timeout=( - watchdog_timeout_secs if watchdog_timeout_secs else self._watchdog_timeout + watchdog_timeout_secs if watchdog_timeout_secs else self._watchdog_timeout_secs ), ) @@ -210,9 +227,6 @@ class FrameProcessor(BaseObject): async def wait_for_task(self, task: asyncio.Task, timeout: Optional[float] = None): await self.get_task_manager().wait_for_task(task, timeout) - def start_watchdog(self): - self.get_task_manager().start_watchdog(asyncio.current_task()) - def reset_watchdog(self): self.get_task_manager().reset_watchdog(asyncio.current_task()) @@ -220,8 +234,13 @@ class FrameProcessor(BaseObject): self._clock = setup.clock self._task_manager = setup.task_manager self._observer = setup.observer + self._watchdog_timers_enabled = ( + self._enable_watchdog_timers + if self._enable_watchdog_timers + else setup.watchdog_timers_enabled + ) if self._metrics is not None: - await self._metrics.setup(self._task_manager) + await self._metrics.setup(self._task_manager, self._watchdog_timers_enabled) async def cleanup(self): await super().cleanup() @@ -279,7 +298,8 @@ class FrameProcessor(BaseObject): async def resume_processing_frames(self): logger.trace(f"{self}: resuming frame processing") - self.__input_event.set() + if self.__input_event: + self.__input_event.set() async def process_frame(self, frame: Frame, direction: FrameDirection): if isinstance(frame, StartFrame): @@ -313,8 +333,8 @@ class FrameProcessor(BaseObject): self._allow_interruptions = frame.allow_interruptions self._enable_metrics = frame.enable_metrics self._enable_usage_metrics = frame.enable_usage_metrics - self._report_only_initial_ttfb = frame.report_only_initial_ttfb self._interruption_strategies = frame.interruption_strategies + self._report_only_initial_ttfb = frame.report_only_initial_ttfb self.__create_input_task() self.__create_push_task() @@ -396,8 +416,12 @@ class FrameProcessor(BaseObject): def __create_input_task(self): if not self.__input_frame_task: self.__should_block_frames = False + if not self.__input_event: + self.__input_event = WatchdogEvent( + self, watchdog_enabled=self.watchdog_timers_enabled + ) self.__input_event.clear() - self.__input_queue = asyncio.Queue() + self.__input_queue = WatchdogQueue(self, watchdog_enabled=self.watchdog_timers_enabled) self.__input_frame_task = self.create_task(self.__input_frame_task_handler()) async def __cancel_input_task(self): @@ -407,7 +431,7 @@ class FrameProcessor(BaseObject): async def __input_frame_task_handler(self): while True: - if self.__should_block_frames: + if self.__should_block_frames and self.__input_event: logger.trace(f"{self}: frame processing paused") await self.__input_event.wait() self.__input_event.clear() @@ -416,7 +440,6 @@ class FrameProcessor(BaseObject): (frame, direction, callback) = await self.__input_queue.get() try: - self.start_watchdog() # Process the frame. await self.process_frame(frame, direction) # If this frame has an associated callback, call it now. @@ -427,11 +450,10 @@ class FrameProcessor(BaseObject): await self.push_error(ErrorFrame(str(e))) finally: self.__input_queue.task_done() - self.reset_watchdog() def __create_push_task(self): if not self.__push_frame_task: - self.__push_queue = asyncio.Queue() + self.__push_queue = WatchdogQueue(self, watchdog_enabled=self.watchdog_timers_enabled) self.__push_frame_task = self.create_task(self.__push_frame_task_handler()) async def __cancel_push_task(self): @@ -442,7 +464,5 @@ class FrameProcessor(BaseObject): async def __push_frame_task_handler(self): while True: (frame, direction) = await self.__push_queue.get() - self.start_watchdog() await self.__internal_push_frame(frame, direction) self.__push_queue.task_done() - self.reset_watchdog() diff --git a/src/pipecat/processors/frameworks/rtvi.py b/src/pipecat/processors/frameworks/rtvi.py index 29e85e5d7..1646a9fa6 100644 --- a/src/pipecat/processors/frameworks/rtvi.py +++ b/src/pipecat/processors/frameworks/rtvi.py @@ -68,6 +68,7 @@ from pipecat.transports.base_input import BaseInputTransport from pipecat.transports.base_output import BaseOutputTransport from pipecat.transports.base_transport import BaseTransport from pipecat.utils.string import match_endofsentence +from pipecat.utils.watchdog_queue import WatchdogQueue RTVI_PROTOCOL_VERSION = "0.3.0" @@ -650,11 +651,9 @@ class RTVIProcessor(FrameProcessor): self._registered_services: Dict[str, RTVIService] = {} # A task to process incoming action frames. - self._action_queue = asyncio.Queue() self._action_task: Optional[asyncio.Task] = None # A task to process incoming transport messages. - self._message_queue = asyncio.Queue() self._message_task: Optional[asyncio.Task] = None self._register_event_handler("on_bot_started") @@ -756,8 +755,10 @@ class RTVIProcessor(FrameProcessor): async def _start(self, frame: StartFrame): if not self._action_task: + self._action_queue = WatchdogQueue(self, watchdog_enabled=self.watchdog_timers_enabled) self._action_task = self.create_task(self._action_task_handler()) if not self._message_task: + self._message_queue = WatchdogQueue(self, watchdog_enabled=self.watchdog_timers_enabled) self._message_task = self.create_task(self._message_task_handler()) await self._call_event_handler("on_bot_started") @@ -783,18 +784,14 @@ class RTVIProcessor(FrameProcessor): async def _action_task_handler(self): while True: frame = await self._action_queue.get() - self.start_watchdog() await self._handle_action(frame.message_id, frame.rtvi_action_run) self._action_queue.task_done() - self.reset_watchdog() async def _message_task_handler(self): while True: message = await self._message_queue.get() - self.start_watchdog() await self._handle_message(message) self._message_queue.task_done() - self.reset_watchdog() async def _handle_transport_message(self, frame: TransportMessageUrgentFrame): try: diff --git a/src/pipecat/processors/metrics/frame_processor_metrics.py b/src/pipecat/processors/metrics/frame_processor_metrics.py index 40a83fa38..9ee1ccdd3 100644 --- a/src/pipecat/processors/metrics/frame_processor_metrics.py +++ b/src/pipecat/processors/metrics/frame_processor_metrics.py @@ -31,7 +31,7 @@ class FrameProcessorMetrics(BaseObject): self._last_ttfb_time = 0 self._should_report_ttfb = True - async def setup(self, task_manager: TaskManager): + async def setup(self, task_manager: TaskManager, watchdog_timers_enabled: bool = False): self._task_manager = task_manager async def cleanup(self): diff --git a/src/pipecat/processors/metrics/sentry.py b/src/pipecat/processors/metrics/sentry.py index 20d854ffb..083ff621b 100644 --- a/src/pipecat/processors/metrics/sentry.py +++ b/src/pipecat/processors/metrics/sentry.py @@ -9,6 +9,8 @@ import asyncio from loguru import logger from pipecat.utils.asyncio import TaskManager +from pipecat.utils.watchdog_queue import WatchdogQueue +from pipecat.utils.watchdog_reseter import WatchdogReseter try: import sentry_sdk @@ -20,7 +22,7 @@ except ModuleNotFoundError as e: from pipecat.processors.metrics.frame_processor_metrics import FrameProcessorMetrics -class SentryMetrics(FrameProcessorMetrics): +class SentryMetrics(WatchdogReseter, FrameProcessorMetrics): def __init__(self): super().__init__() self._ttfb_metrics_tx = None @@ -28,13 +30,12 @@ class SentryMetrics(FrameProcessorMetrics): self._sentry_available = sentry_sdk.is_initialized() if not self._sentry_available: logger.warning("Sentry SDK not initialized. Sentry features will be disabled.") - self._sentry_queue = asyncio.Queue() self._sentry_task = None - async def setup(self, task_manager: TaskManager): - await super().setup(task_manager) + async def setup(self, task_manager: TaskManager, watchdog_timers_enabled: bool = False): + await super().setup(task_manager, watchdog_timers_enabled) if self._sentry_available: - self._sentry_queue = asyncio.Queue() + self._sentry_queue = WatchdogQueue(self, watchdog_enabled=watchdog_timers_enabled) self._sentry_task = self.task_manager.create_task( self._sentry_task_handler(), name=f"{self}::_sentry_task_handler" ) @@ -48,6 +49,10 @@ class SentryMetrics(FrameProcessorMetrics): logger.trace(f"{self} Flushing Sentry metrics") sentry_sdk.flush(timeout=5.0) + def reset_watchdog(self): + if self._task_manager: + self._task_manager.reset_watchdog(asyncio.current_task()) + async def start_ttfb_metrics(self, report_only_initial_ttfb): await super().start_ttfb_metrics(report_only_initial_ttfb) @@ -93,3 +98,4 @@ class SentryMetrics(FrameProcessorMetrics): if tx: await self.task_manager.get_event_loop().run_in_executor(None, tx.finish) running = tx is not None + self._sentry_queue.task_done() diff --git a/src/pipecat/processors/producer_processor.py b/src/pipecat/processors/producer_processor.py index 6ada2ed83..ad08802e2 100644 --- a/src/pipecat/processors/producer_processor.py +++ b/src/pipecat/processors/producer_processor.py @@ -9,6 +9,7 @@ from typing import Awaitable, Callable, List from pipecat.frames.frames import Frame from pipecat.processors.frame_processor import FrameDirection, FrameProcessor +from pipecat.utils.watchdog_queue import WatchdogQueue async def identity_transformer(frame: Frame): @@ -36,14 +37,14 @@ class ProducerProcessor(FrameProcessor): self._passthrough = passthrough self._consumers: List[asyncio.Queue] = [] - def add_consumer(self): + def add_consumer(self, consumer: FrameProcessor): """ Adds a new consumer and returns its associated queue. Returns: asyncio.Queue: The queue for the newly added consumer. """ - queue = asyncio.Queue() + queue = WatchdogQueue(consumer, watchdog_enabled=self.watchdog_timers_enabled) self._consumers.append(queue) return queue diff --git a/src/pipecat/services/anthropic/llm.py b/src/pipecat/services/anthropic/llm.py index 236f269fa..b5334c383 100644 --- a/src/pipecat/services/anthropic/llm.py +++ b/src/pipecat/services/anthropic/llm.py @@ -47,6 +47,7 @@ from pipecat.processors.aggregators.openai_llm_context import ( from pipecat.processors.frame_processor import FrameDirection from pipecat.services.llm_service import FunctionCallFromLLM, LLMService from pipecat.utils.tracing.service_decorators import traced_llm +from pipecat.utils.watchdog_async_iterator import WatchdogAsyncIterator try: from anthropic import NOT_GIVEN, AsyncAnthropic, NotGiven @@ -203,7 +204,9 @@ class AnthropicLLMService(LLMService): json_accumulator = "" function_calls = [] - async for event in response: + async for event in WatchdogAsyncIterator( + response, reseter=self, watchdog_enabled=self.watchdog_timers_enabled + ): # Aggregate streaming content, create frames, trigger events if event.type == "content_block_delta": diff --git a/src/pipecat/services/assemblyai/stt.py b/src/pipecat/services/assemblyai/stt.py index 09aaa4d25..452d4cfb6 100644 --- a/src/pipecat/services/assemblyai/stt.py +++ b/src/pipecat/services/assemblyai/stt.py @@ -189,17 +189,16 @@ class AssemblyAISTTService(STTService): try: while self._connected: try: - message = await self._websocket.recv() - self.start_watchdog() + message = await asyncio.wait_for(self._websocket.recv(), timeout=1.0) data = json.loads(message) await self._handle_message(data) + except asyncio.TimeoutError: + self.reset_watchdog() except websockets.exceptions.ConnectionClosedOK: break except Exception as e: logger.error(f"Error processing WebSocket message: {e}") break - finally: - self.reset_watchdog() except Exception as e: logger.error(f"Fatal error in receive handler: {e}") diff --git a/src/pipecat/services/aws/llm.py b/src/pipecat/services/aws/llm.py index dcec91463..d3217e7a1 100644 --- a/src/pipecat/services/aws/llm.py +++ b/src/pipecat/services/aws/llm.py @@ -711,6 +711,8 @@ class AWSBedrockLLMService(LLMService): function_calls = [] for event in response["stream"]: + self.reset_watchdog() + # Handle text content if "contentBlockDelta" in event: delta = event["contentBlockDelta"]["delta"] @@ -762,6 +764,7 @@ class AWSBedrockLLMService(LLMService): completion_tokens += usage.get("outputTokens", 0) cache_read_input_tokens += usage.get("cacheReadInputTokens", 0) cache_creation_input_tokens += usage.get("cacheWriteInputTokens", 0) + await self.run_function_calls(function_calls) except asyncio.CancelledError: # If we're interrupted, we won't get a complete usage report. So set our flag to use the diff --git a/src/pipecat/services/aws/stt.py b/src/pipecat/services/aws/stt.py index c9c1b32d1..c4170ebad 100644 --- a/src/pipecat/services/aws/stt.py +++ b/src/pipecat/services/aws/stt.py @@ -284,9 +284,7 @@ class AWSTranscribeSTTService(STTService): break try: - response = await self._ws_client.recv() - - self.start_watchdog() + response = await asyncio.wait_for(self._ws_client.recv(), timeout=1.0) headers, payload = decode_event(response) @@ -337,6 +335,8 @@ class AWSTranscribeSTTService(STTService): else: logger.debug(f"{self} Other message type received: {headers}") logger.debug(f"{self} Payload: {payload}") + except asyncio.TimeoutError: + self.reset_watchdog() except websockets.exceptions.ConnectionClosed as e: logger.error( f"{self} WebSocket connection closed in receive loop with code {e.code}: {e.reason}" @@ -345,5 +345,3 @@ class AWSTranscribeSTTService(STTService): except Exception as e: logger.error(f"{self} Unexpected error in receive loop: {e}") break - finally: - self.reset_watchdog() diff --git a/src/pipecat/services/aws_nova_sonic/aws.py b/src/pipecat/services/aws_nova_sonic/aws.py index 718e3dc50..93eb77e90 100644 --- a/src/pipecat/services/aws_nova_sonic/aws.py +++ b/src/pipecat/services/aws_nova_sonic/aws.py @@ -697,9 +697,9 @@ class AWSNovaSonicLLMService(LLMService): try: while self._stream and not self._disconnecting: output = await self._stream.await_output() - result = await output[1].receive() + result = await asyncio.wait_for(output[1].receive(), timeout=1.0) - self.start_watchdog() + self.reset_watchdog() if result.value and result.value.bytes_: response_data = result.value.bytes_.decode("utf-8") @@ -728,13 +728,12 @@ class AWSNovaSonicLLMService(LLMService): elif "completionEnd" in event_json: # Handle the LLM completion ending await self._handle_completion_end_event(event_json) - + except asyncio.TimeoutError: + self.reset_watchdog() except Exception as e: logger.error(f"{self} error processing responses: {e}") if self._wants_connection: await self.reset_conversation() - finally: - self.reset_watchdog() async def _handle_completion_start_event(self, event_json): pass diff --git a/src/pipecat/services/cartesia/tts.py b/src/pipecat/services/cartesia/tts.py index 22493f229..8ac997f27 100644 --- a/src/pipecat/services/cartesia/tts.py +++ b/src/pipecat/services/cartesia/tts.py @@ -30,6 +30,7 @@ from pipecat.transcriptions.language import Language from pipecat.utils.text.base_text_aggregator import BaseTextAggregator from pipecat.utils.text.skip_tags_aggregator import SkipTagsAggregator from pipecat.utils.tracing.service_decorators import traced_tts +from pipecat.utils.watchdog_async_iterator import WatchdogAsyncIterator # See .env.example for Cartesia configuration needed try: @@ -255,7 +256,9 @@ class CartesiaTTSService(AudioContextWordTTSService): self._context_id = None async def _receive_messages(self): - async for message in self._get_websocket(): + async for message in WatchdogAsyncIterator( + self._get_websocket(), reseter=self, watchdog_enabled=self.watchdog_timers_enabled + ): msg = json.loads(message) if not msg or not self.audio_context_available(msg["context_id"]): continue diff --git a/src/pipecat/services/elevenlabs/tts.py b/src/pipecat/services/elevenlabs/tts.py index ccd9b5b3f..7665632fc 100644 --- a/src/pipecat/services/elevenlabs/tts.py +++ b/src/pipecat/services/elevenlabs/tts.py @@ -33,6 +33,7 @@ from pipecat.services.tts_service import ( ) from pipecat.transcriptions.language import Language from pipecat.utils.tracing.service_decorators import traced_tts +from pipecat.utils.watchdog_async_iterator import WatchdogAsyncIterator # See .env.example for ElevenLabs configuration needed try: @@ -394,7 +395,9 @@ class ElevenLabsTTSService(AudioContextWordTTSService): self._started = False async def _receive_messages(self): - async for message in self._get_websocket(): + async for message in WatchdogAsyncIterator( + self._get_websocket(), reseter=self, watchdog_enabled=self.watchdog_timers_enabled + ): msg = json.loads(message) received_ctx_id = msg.get("contextId") @@ -425,8 +428,10 @@ class ElevenLabsTTSService(AudioContextWordTTSService): self._cumulative_time = word_times[-1][1] async def _keepalive_task_handler(self): + KEEPALIVE_SLEEP = 10 if self.watchdog_timers_enabled else 3 while True: - await asyncio.sleep(10) + self.reset_watchdog() + await asyncio.sleep(KEEPALIVE_SLEEP) try: if self._websocket and self._websocket.open: if self._context_id: diff --git a/src/pipecat/services/gemini_multimodal_live/gemini.py b/src/pipecat/services/gemini_multimodal_live/gemini.py index d8a90fada..c713c3cab 100644 --- a/src/pipecat/services/gemini_multimodal_live/gemini.py +++ b/src/pipecat/services/gemini_multimodal_live/gemini.py @@ -60,7 +60,8 @@ from pipecat.services.openai.llm import ( from pipecat.transcriptions.language import Language from pipecat.utils.string import match_endofsentence from pipecat.utils.time import time_now_iso8601 -from pipecat.utils.tracing.service_decorators import traced_gemini_live, traced_stt, traced_tts +from pipecat.utils.tracing.service_decorators import traced_gemini_live, traced_stt +from pipecat.utils.watchdog_async_iterator import WatchdogAsyncIterator from . import events @@ -686,9 +687,9 @@ class GeminiMultimodalLiveLLMService(LLMService): # async def _receive_task_handler(self): - async for message in self._websocket: - self.start_watchdog() - + async for message in WatchdogAsyncIterator( + self._websocket, reseter=self, watchdog_enabled=self.watchdog_timers_enabled + ): evt = events.parse_server_event(message) # logger.debug(f"Received event: {message[:500]}") # logger.debug(f"Received event: {evt}") @@ -711,8 +712,6 @@ class GeminiMultimodalLiveLLMService(LLMService): # errors are fatal, so exit the receive loop return - self.reset_watchdog() - # # # diff --git a/src/pipecat/services/gladia/stt.py b/src/pipecat/services/gladia/stt.py index 885fe8dc2..75e8bbee3 100644 --- a/src/pipecat/services/gladia/stt.py +++ b/src/pipecat/services/gladia/stt.py @@ -27,6 +27,7 @@ from pipecat.services.stt_service import STTService from pipecat.transcriptions.language import Language from pipecat.utils.time import time_now_iso8601 from pipecat.utils.tracing.service_decorators import traced_stt +from pipecat.utils.watchdog_async_iterator import WatchdogAsyncIterator try: import websockets @@ -391,8 +392,8 @@ class GladiaSTTService(STTService): await self._send_buffered_audio() # Start tasks - self._receive_task = asyncio.create_task(self._receive_task_handler()) - self._keepalive_task = asyncio.create_task(self._keepalive_task_handler()) + self._receive_task = self.create_task(self._receive_task_handler()) + self._keepalive_task = self.create_task(self._keepalive_task_handler()) # Wait for tasks to complete await asyncio.gather(self._receive_task, self._keepalive_task) @@ -403,9 +404,9 @@ class GladiaSTTService(STTService): # Clean up tasks if self._receive_task: - self._receive_task.cancel() + await self.cancel_task(self._receive_task) if self._keepalive_task: - self._keepalive_task.cancel() + await self.cancel_task(self._keepalive_task) # Attempt reconnect using helper if not await self._maybe_reconnect(): @@ -484,9 +485,11 @@ class GladiaSTTService(STTService): async def _keepalive_task_handler(self): """Send periodic empty audio chunks to keep the connection alive.""" try: + KEEPALIVE_SLEEP = 20 if self.watchdog_timers_enabled else 3 while self._connection_active: - # Send keepalive every 20 seconds (Gladia times out after 30 seconds) - await asyncio.sleep(20) + self.reset_watchdog() + # Send keepalive (Gladia times out after 30 seconds) + await asyncio.sleep(KEEPALIVE_SLEEP) if self._websocket and not self._websocket.closed: # Send an empty audio chunk as keepalive empty_audio = b"" @@ -501,9 +504,9 @@ class GladiaSTTService(STTService): async def _receive_task_handler(self): try: - async for message in self._websocket: - self.start_watchdog() - + async for message in WatchdogAsyncIterator( + self._websocket, reseter=self, watchdog_enabled=self.watchdog_timers_enabled + ): content = json.loads(message) # Handle audio chunk acknowledgments @@ -568,8 +571,6 @@ class GladiaSTTService(STTService): pass except Exception as e: logger.error(f"Error in Gladia WebSocket handler: {e}") - finally: - self.reset_watchdog() async def _maybe_reconnect(self) -> bool: """Handle exponential backoff reconnection logic.""" diff --git a/src/pipecat/services/google/llm.py b/src/pipecat/services/google/llm.py index f983b7342..5fe005fbd 100644 --- a/src/pipecat/services/google/llm.py +++ b/src/pipecat/services/google/llm.py @@ -48,6 +48,7 @@ from pipecat.services.openai.llm import ( OpenAIUserContextAggregator, ) from pipecat.utils.tracing.service_decorators import traced_llm +from pipecat.utils.watchdog_async_iterator import WatchdogAsyncIterator # Suppress gRPC fork warnings os.environ["GRPC_ENABLE_FORK_SUPPORT"] = "false" @@ -557,7 +558,9 @@ class GoogleLLMService(LLMService): ) function_calls = [] - async for chunk in response: + async for chunk in WatchdogAsyncIterator( + response, reseter=self, watchdog_enabled=self.watchdog_timers_enabled + ): # Stop TTFB metrics after the first chunk await self.stop_ttfb_metrics() if chunk.usage_metadata: diff --git a/src/pipecat/services/google/llm_openai.py b/src/pipecat/services/google/llm_openai.py index a497cb229..e76ac1886 100644 --- a/src/pipecat/services/google/llm_openai.py +++ b/src/pipecat/services/google/llm_openai.py @@ -11,6 +11,7 @@ from openai import AsyncStream from openai.types.chat import ChatCompletionChunk from pipecat.services.llm_service import FunctionCallFromLLM +from pipecat.utils.watchdog_async_iterator import WatchdogAsyncIterator # Suppress gRPC fork warnings os.environ["GRPC_ENABLE_FORK_SUPPORT"] = "false" @@ -53,7 +54,9 @@ class GoogleLLMOpenAIBetaService(OpenAILLMService): context ) - async for chunk in chunk_stream: + async for chunk in WatchdogAsyncIterator( + chunk_stream, reseter=self, watchdog_enabled=self.watchdog_timers_enabled + ): if chunk.usage: tokens = LLMTokenUsage( prompt_tokens=chunk.usage.prompt_tokens, diff --git a/src/pipecat/services/google/stt.py b/src/pipecat/services/google/stt.py index 26186e6bf..80f061b44 100644 --- a/src/pipecat/services/google/stt.py +++ b/src/pipecat/services/google/stt.py @@ -10,6 +10,7 @@ import os import time from pipecat.utils.tracing.service_decorators import traced_stt +from pipecat.utils.watchdog_async_iterator import WatchdogAsyncIterator # Suppress gRPC fork warnings os.environ["GRPC_ENABLE_FORK_SUPPORT"] = "false" @@ -436,7 +437,6 @@ class GoogleSTTService(STTService): self._location = location self._stream = None self._config = None - self._request_queue = asyncio.Queue() self._streaming_task = None # Used for keep-alive logic @@ -683,23 +683,15 @@ class GoogleSTTService(STTService): ), ) + self._request_queue = asyncio.Queue() self._streaming_task = self.create_task(self._stream_audio()) async def _disconnect(self): """Clean up streaming recognition resources.""" if self._streaming_task: logger.debug("Disconnecting from Google Speech-to-Text") - # Send sentinel value to stop request generator - await self._request_queue.put(None) await self.cancel_task(self._streaming_task) self._streaming_task = None - # Clear any remaining items in the queue - while not self._request_queue.empty(): - try: - self._request_queue.get_nowait() - self._request_queue.task_done() - except asyncio.QueueEmpty: - break async def _request_generator(self): """Generates requests for the streaming recognize method.""" @@ -714,29 +706,23 @@ class GoogleSTTService(STTService): ) while True: - try: - audio_data = await self._request_queue.get() - if audio_data is None: # Sentinel value to stop - break + audio_data = await self._request_queue.get() - # Check streaming limit - if (int(time.time() * 1000) - self._stream_start_time) > self.STREAMING_LIMIT: - logger.debug("Streaming limit reached, initiating graceful reconnection") - # Instead of immediate reconnection, we'll break and let the stream close naturally - self._last_audio_input = self._audio_input - self._audio_input = [] - self._restart_counter += 1 - # Put the current audio chunk back in the queue - await self._request_queue.put(audio_data) - break + self._request_queue.task_done() - self._audio_input.append(audio_data) - yield cloud_speech.StreamingRecognizeRequest(audio=audio_data) - - except asyncio.CancelledError: + # Check streaming limit + if (int(time.time() * 1000) - self._stream_start_time) > self.STREAMING_LIMIT: + logger.debug("Streaming limit reached, initiating graceful reconnection") + # Instead of immediate reconnection, we'll break and let the stream close naturally + self._last_audio_input = self._audio_input + self._audio_input = [] + self._restart_counter += 1 + # Put the current audio chunk back in the queue + await self._request_queue.put(audio_data) break - finally: - self._request_queue.task_done() + + self._audio_input.append(audio_data) + yield cloud_speech.StreamingRecognizeRequest(audio=audio_data) except Exception as e: logger.error(f"Error in request generator: {e}") @@ -747,8 +733,6 @@ class GoogleSTTService(STTService): try: while True: try: - self.start_watchdog() - if self._request_queue.empty(): # wait for 10ms in case we don't have audio await asyncio.sleep(0.01) @@ -763,8 +747,6 @@ class GoogleSTTService(STTService): # Process responses await self._process_responses(streaming_recognize) - self.reset_watchdog() - # If we're here, check if we need to reconnect if (int(time.time() * 1000) - self._stream_start_time) > self.STREAMING_LIMIT: logger.debug("Reconnecting stream after timeout") @@ -779,8 +761,6 @@ class GoogleSTTService(STTService): await asyncio.sleep(1) # Brief delay before reconnecting self._stream_start_time = int(time.time() * 1000) - finally: - self.reset_watchdog() except Exception as e: logger.error(f"Error in streaming task: {e}") @@ -804,17 +784,15 @@ class GoogleSTTService(STTService): async def _process_responses(self, streaming_recognize): """Process streaming recognition responses.""" try: - async for response in streaming_recognize: - self.start_watchdog() - + async for response in WatchdogAsyncIterator( + streaming_recognize, reseter=self, watchdog_enabled=self.watchdog_timers_enabled + ): # Check streaming limit if (int(time.time() * 1000) - self._stream_start_time) > self.STREAMING_LIMIT: logger.debug("Stream timeout reached in response processing") - self.reset_watchdog() break if not response.results: - self.reset_watchdog() continue for result in response.results: @@ -856,11 +834,8 @@ class GoogleSTTService(STTService): result=result, ) ) - - self.reset_watchdog() except Exception as e: logger.error(f"Error processing Google STT responses: {e}") - self.reset_watchdog() # Re-raise the exception to let it propagate (e.g. in the case of a # timeout, propagate to _stream_audio to reconnect) raise diff --git a/src/pipecat/services/neuphonic/tts.py b/src/pipecat/services/neuphonic/tts.py index bfceca50b..85bd9d0cc 100644 --- a/src/pipecat/services/neuphonic/tts.py +++ b/src/pipecat/services/neuphonic/tts.py @@ -30,6 +30,7 @@ from pipecat.processors.frame_processor import FrameDirection from pipecat.services.tts_service import InterruptibleTTSService, TTSService from pipecat.transcriptions.language import Language from pipecat.utils.tracing.service_decorators import traced_tts +from pipecat.utils.watchdog_async_iterator import WatchdogAsyncIterator try: import websockets @@ -221,7 +222,9 @@ class NeuphonicTTSService(InterruptibleTTSService): self._websocket = None async def _receive_messages(self): - async for message in self._websocket: + async for message in WatchdogAsyncIterator( + self._websocket, reseter=self, watchdog_enabled=self.watchdog_timers_enabled + ): if isinstance(message, str): msg = json.loads(message) if msg.get("data", {}).get("audio") is not None: @@ -232,8 +235,10 @@ class NeuphonicTTSService(InterruptibleTTSService): await self.push_frame(frame) async def _keepalive_task_handler(self): + KEEPALIVE_SLEEP = 10 if self.watchdog_timers_enabled else 3 while True: - await asyncio.sleep(10) + self.reset_watchdog() + await asyncio.sleep(KEEPALIVE_SLEEP) await self._send_text("") async def _send_text(self, text: str): diff --git a/src/pipecat/services/openai/base_llm.py b/src/pipecat/services/openai/base_llm.py index 2badfed96..f0029ad22 100644 --- a/src/pipecat/services/openai/base_llm.py +++ b/src/pipecat/services/openai/base_llm.py @@ -36,6 +36,7 @@ from pipecat.processors.aggregators.openai_llm_context import ( from pipecat.processors.frame_processor import FrameDirection from pipecat.services.llm_service import FunctionCallFromLLM, LLMService from pipecat.utils.tracing.service_decorators import traced_llm +from pipecat.utils.watchdog_async_iterator import WatchdogAsyncIterator class BaseOpenAILLMService(LLMService): @@ -192,7 +193,9 @@ class BaseOpenAILLMService(LLMService): context ) - async for chunk in chunk_stream: + async for chunk in WatchdogAsyncIterator( + chunk_stream, reseter=self, watchdog_enabled=self.watchdog_timers_enabled + ): if chunk.usage: tokens = LLMTokenUsage( prompt_tokens=chunk.usage.prompt_tokens, diff --git a/src/pipecat/services/openai_realtime_beta/openai.py b/src/pipecat/services/openai_realtime_beta/openai.py index 7f459000a..8d5168c70 100644 --- a/src/pipecat/services/openai_realtime_beta/openai.py +++ b/src/pipecat/services/openai_realtime_beta/openai.py @@ -52,7 +52,8 @@ from pipecat.services.llm_service import FunctionCallFromLLM, LLMService from pipecat.services.openai.llm import OpenAIContextAggregatorPair from pipecat.transcriptions.language import Language from pipecat.utils.time import time_now_iso8601 -from pipecat.utils.tracing.service_decorators import traced_openai_realtime, traced_stt, traced_tts +from pipecat.utils.tracing.service_decorators import traced_openai_realtime, traced_stt +from pipecat.utils.watchdog_async_iterator import WatchdogAsyncIterator from . import events from .context import ( @@ -369,8 +370,9 @@ class OpenAIRealtimeBetaLLMService(LLMService): # async def _receive_task_handler(self): - async for message in self._websocket: - self.start_watchdog() + async for message in WatchdogAsyncIterator( + self._websocket, reseter=self, watchdog_enabled=self.watchdog_timers_enabled + ): evt = events.parse_server_event(message) if evt.type == "session.created": await self._handle_evt_session_created(evt) @@ -401,7 +403,6 @@ class OpenAIRealtimeBetaLLMService(LLMService): await self._handle_evt_error(evt) # errors are fatal, so exit the receive loop return - self.reset_watchdog() @traced_openai_realtime(operation="llm_setup") async def _handle_evt_session_created(self, evt): diff --git a/src/pipecat/services/riva/stt.py b/src/pipecat/services/riva/stt.py index 91c207c8b..284252adb 100644 --- a/src/pipecat/services/riva/stt.py +++ b/src/pipecat/services/riva/stt.py @@ -23,6 +23,7 @@ from pipecat.services.stt_service import SegmentedSTTService, STTService from pipecat.transcriptions.language import Language from pipecat.utils.time import time_now_iso8601 from pipecat.utils.tracing.service_decorators import traced_stt +from pipecat.utils.watchdog_queue import WatchdogQueue try: import riva.client @@ -198,7 +199,9 @@ class RivaSTTService(STTService): self._thread_task = self.create_task(self._thread_task_handler()) if not self._response_task: - self._response_queue = asyncio.Queue() + self._response_queue = WatchdogQueue( + self, watchdog_enabled=self.watchdog_timers_enabled + ) self._response_task = self.create_task(self._response_task_handler()) async def stop(self, frame: EndFrame): @@ -224,13 +227,12 @@ class RivaSTTService(STTService): streaming_config=self._config, ) for response in responses: - self.start_watchdog() + self.reset_watchdog() if not response.results: continue asyncio.run_coroutine_threadsafe( self._response_queue.put(response), self.get_event_loop() ) - self.reset_watchdog() async def _thread_task_handler(self): try: @@ -285,9 +287,8 @@ class RivaSTTService(STTService): async def _response_task_handler(self): while True: response = await self._response_queue.get() - self.start_watchdog() await self._handle_response(response) - self.reset_watchdog() + self._response_queue.task_done() async def run_stt(self, audio: bytes) -> AsyncGenerator[Frame, None]: await self.start_ttfb_metrics() diff --git a/src/pipecat/services/sambanova/llm.py b/src/pipecat/services/sambanova/llm.py index 01a8d294c..3ca2ee5be 100644 --- a/src/pipecat/services/sambanova/llm.py +++ b/src/pipecat/services/sambanova/llm.py @@ -19,6 +19,7 @@ from pipecat.processors.aggregators.openai_llm_context import OpenAILLMContext from pipecat.services.llm_service import FunctionCallFromLLM from pipecat.services.openai.llm import OpenAILLMService from pipecat.utils.tracing.service_decorators import traced_llm +from pipecat.utils.watchdog_async_iterator import WatchdogAsyncIterator class SambaNovaLLMService(OpenAILLMService): # type: ignore @@ -94,7 +95,9 @@ class SambaNovaLLMService(OpenAILLMService): # type: ignore context ) - async for chunk in chunk_stream: + async for chunk in WatchdogAsyncIterator( + chunk_stream, reseter=self, watchdog_enabled=self.watchdog_timers_enabled + ): if chunk.usage: tokens = LLMTokenUsage( prompt_tokens=chunk.usage.prompt_tokens, diff --git a/src/pipecat/services/simli/video.py b/src/pipecat/services/simli/video.py index 76381b9c6..2bca697cb 100644 --- a/src/pipecat/services/simli/video.py +++ b/src/pipecat/services/simli/video.py @@ -18,6 +18,7 @@ from pipecat.frames.frames import ( TTSAudioRawFrame, ) from pipecat.processors.frame_processor import FrameDirection, FrameProcessor, StartFrame +from pipecat.utils.watchdog_async_iterator import WatchdogAsyncIterator try: from av.audio.frame import AudioFrame @@ -61,8 +62,10 @@ class SimliVideoService(FrameProcessor): async def _consume_and_process_audio(self): await self._pipecat_resampler_event.wait() - async for audio_frame in self._simli_client.getAudioStreamIterator(): - self.start_watchdog() + audio_iterator = self._simli_client.getAudioStreamIterator() + async for audio_frame in WatchdogAsyncIterator( + audio_iterator, reseter=self, watchdog_enabled=self.watchdog_timers_enabled + ): resampled_frames = self._pipecat_resampler.resample(audio_frame) for resampled_frame in resampled_frames: audio_array = resampled_frame.to_ndarray() @@ -75,12 +78,13 @@ class SimliVideoService(FrameProcessor): num_channels=1, ), ) - self.reset_watchdog() async def _consume_and_process_video(self): await self._pipecat_resampler_event.wait() - async for video_frame in self._simli_client.getVideoStreamIterator(targetFormat="rgb24"): - self.start_watchdog() + video_iterator = self._simli_client.getVideoStreamIterator(targetFormat="rgb24") + async for video_frame in WatchdogAsyncIterator( + video_iterator, reseter=self, watchdog_enabled=self.watchdog_timers_enabled + ): # Process the video frame convertedFrame: OutputImageRawFrame = OutputImageRawFrame( image=video_frame.to_rgb().to_image().tobytes(), @@ -89,7 +93,6 @@ class SimliVideoService(FrameProcessor): ) convertedFrame.pts = video_frame.pts await self.push_frame(convertedFrame) - self.reset_watchdog() async def process_frame(self, frame: Frame, direction: FrameDirection): await super().process_frame(frame, direction) diff --git a/src/pipecat/services/tavus/video.py b/src/pipecat/services/tavus/video.py index 4ec744c51..9688aaebc 100644 --- a/src/pipecat/services/tavus/video.py +++ b/src/pipecat/services/tavus/video.py @@ -27,6 +27,7 @@ from pipecat.frames.frames import ( from pipecat.processors.frame_processor import FrameDirection, FrameProcessorSetup from pipecat.services.ai_service import AIService from pipecat.transports.services.tavus import TavusCallbacks, TavusParams, TavusTransportClient +from pipecat.utils.watchdog_queue import WatchdogQueue class TavusVideoService(AIService): @@ -71,7 +72,6 @@ class TavusVideoService(AIService): self._resampler = create_default_resampler() self._audio_buffer = bytearray() - self._queue = asyncio.Queue() self._send_task: Optional[asyncio.Task] = None # This is the custom track destination expected by Tavus self._transport_destination: Optional[str] = "stream" @@ -188,7 +188,7 @@ class TavusVideoService(AIService): async def _create_send_task(self): if not self._send_task: - self._queue = asyncio.Queue() + self._queue = WatchdogQueue(self, watchdog_enabled=self.watchdog_timers_enabled) self._send_task = self.create_task(self._send_task_handler()) async def _cancel_send_task(self): @@ -217,7 +217,6 @@ class TavusVideoService(AIService): async def _send_task_handler(self): while True: frame = await self._queue.get() - self.start_watchdog() if isinstance(frame, OutputAudioRawFrame) and self._client: await self._client.write_audio_frame(frame) - self.reset_watchdog() + self._queue.task_done() diff --git a/src/pipecat/services/tts_service.py b/src/pipecat/services/tts_service.py index 904f603a9..4fe7f1905 100644 --- a/src/pipecat/services/tts_service.py +++ b/src/pipecat/services/tts_service.py @@ -39,6 +39,7 @@ from pipecat.utils.text.base_text_aggregator import BaseTextAggregator from pipecat.utils.text.base_text_filter import BaseTextFilter from pipecat.utils.text.simple_text_aggregator import SimpleTextAggregator from pipecat.utils.time import seconds_to_nanoseconds +from pipecat.utils.watchdog_queue import WatchdogQueue class TTSService(AIService): @@ -315,6 +316,8 @@ class TTSService(AIService): if has_started: await self.push_frame(TTSStoppedFrame()) has_started = False + finally: + self.reset_watchdog() class WordTTSService(TTSService): @@ -327,7 +330,6 @@ class WordTTSService(TTSService): def __init__(self, **kwargs): super().__init__(**kwargs) self._initial_word_timestamp = -1 - self._words_queue = asyncio.Queue() self._words_task = None self._llm_response_started: bool = False @@ -369,6 +371,7 @@ class WordTTSService(TTSService): def _create_words_task(self): if not self._words_task: + self._words_queue = WatchdogQueue(self, watchdog_enabled=self.watchdog_timers_enabled) self._words_task = self.create_task(self._words_task_handler()) async def _stop_words_task(self): @@ -521,7 +524,6 @@ class AudioContextWordTTSService(WebsocketWordTTSService): def __init__(self, **kwargs): super().__init__(**kwargs) - self._contexts_queue = asyncio.Queue() self._contexts: Dict[str, asyncio.Queue] = {} self._audio_context_task = None @@ -578,7 +580,9 @@ class AudioContextWordTTSService(WebsocketWordTTSService): def _create_audio_context_task(self): if not self._audio_context_task: - self._contexts_queue = asyncio.Queue() + self._contexts_queue = WatchdogQueue( + self, watchdog_enabled=self.watchdog_timers_enabled + ) self._contexts: Dict[str, asyncio.Queue] = {} self._audio_context_task = self.create_task(self._audio_context_task_handler()) @@ -620,10 +624,12 @@ class AudioContextWordTTSService(WebsocketWordTTSService): while running: try: frame = await asyncio.wait_for(queue.get(), timeout=AUDIO_CONTEXT_TIMEOUT) + self.reset_watchdog() if frame: await self.push_frame(frame) running = frame is not None except asyncio.TimeoutError: + self.reset_watchdog() # We didn't get audio, so let's consider this context finished. logger.trace(f"{self} time out on audio context {context_id}") break diff --git a/src/pipecat/transports/base_input.py b/src/pipecat/transports/base_input.py index d707c5969..93cd90825 100644 --- a/src/pipecat/transports/base_input.py +++ b/src/pipecat/transports/base_input.py @@ -368,8 +368,6 @@ class BaseInputTransport(FrameProcessor): self._audio_in_queue.get(), timeout=AUDIO_INPUT_TIMEOUT_SECS ) - self.start_watchdog() - # If an audio filter is available, run it before VAD. if self._params.audio_in_filter: frame.audio = await self._params.audio_in_filter.filter(frame.audio) diff --git a/src/pipecat/transports/base_output.py b/src/pipecat/transports/base_output.py index 386f223d7..f477be447 100644 --- a/src/pipecat/transports/base_output.py +++ b/src/pipecat/transports/base_output.py @@ -40,6 +40,7 @@ from pipecat.frames.frames import ( from pipecat.processors.frame_processor import FrameDirection, FrameProcessor from pipecat.transports.base_transport import TransportParams from pipecat.utils.time import nanoseconds_to_seconds +from pipecat.utils.watchdog_priority_queue import WatchdogPriorityQueue BOT_VAD_STOP_SECS = 0.35 @@ -441,8 +442,10 @@ class BaseOutputTransport(FrameProcessor): frame = await asyncio.wait_for( self._audio_queue.get(), timeout=vad_stop_secs ) + self._transport.reset_watchdog() yield frame except asyncio.TimeoutError: + self._transport.reset_watchdog() # Notify the bot stopped speaking upstream if necessary. await self._bot_stopped_speaking() @@ -452,11 +455,13 @@ class BaseOutputTransport(FrameProcessor): while True: try: frame = self._audio_queue.get_nowait() + self._transport.reset_watchdog() if isinstance(frame, OutputAudioRawFrame): frame.audio = await self._mixer.mix(frame.audio) last_frame_time = time.time() yield frame except asyncio.QueueEmpty: + self._transport.reset_watchdog() # Notify the bot stopped speaking upstream if necessary. diff_time = time.time() - last_frame_time if diff_time > vad_stop_secs: @@ -597,7 +602,9 @@ class BaseOutputTransport(FrameProcessor): def _create_clock_task(self): if not self._clock_task: - self._clock_queue = asyncio.PriorityQueue() + self._clock_queue = WatchdogPriorityQueue( + self._transport, watchdog_enabled=self._transport.watchdog_timers_enabled + ) self._clock_task = self._transport.create_task(self._clock_task_handler()) async def _cancel_clock_task(self): diff --git a/src/pipecat/transports/network/fastapi_websocket.py b/src/pipecat/transports/network/fastapi_websocket.py index 7f6256b83..3d19cb05f 100644 --- a/src/pipecat/transports/network/fastapi_websocket.py +++ b/src/pipecat/transports/network/fastapi_websocket.py @@ -26,11 +26,12 @@ from pipecat.frames.frames import ( TransportMessageFrame, TransportMessageUrgentFrame, ) -from pipecat.processors.frame_processor import FrameDirection, FrameProcessorSetup +from pipecat.processors.frame_processor import FrameDirection from pipecat.serializers.base_serializer import FrameSerializer, FrameSerializerType from pipecat.transports.base_input import BaseInputTransport from pipecat.transports.base_output import BaseOutputTransport from pipecat.transports.base_transport import BaseTransport, TransportParams +from pipecat.utils.watchdog_async_iterator import WatchdogAsyncIterator try: from fastapi import WebSocket @@ -178,12 +179,12 @@ class FastAPIWebsocketInputTransport(BaseInputTransport): async def _receive_messages(self): try: - async for message in self._client.receive(): + async for message in WatchdogAsyncIterator( + self._client.receive(), reseter=self, watchdog_enabled=self.watchdog_timers_enabled + ): if not self._params.serializer: continue - self.start_watchdog() - frame = await self._params.serializer.deserialize(message) if not frame: @@ -193,13 +194,9 @@ class FastAPIWebsocketInputTransport(BaseInputTransport): await self.push_audio_frame(frame) else: await self.push_frame(frame) - - self.reset_watchdog() except Exception as e: logger.error(f"{self} exception receiving data: {e.__class__.__name__} ({e})") - self.reset_watchdog() - await self._client.trigger_client_disconnected() async def _monitor_websocket(self): diff --git a/src/pipecat/transports/network/small_webrtc.py b/src/pipecat/transports/network/small_webrtc.py index 5b36fec37..086aa383c 100644 --- a/src/pipecat/transports/network/small_webrtc.py +++ b/src/pipecat/transports/network/small_webrtc.py @@ -33,6 +33,7 @@ from pipecat.transports.base_input import BaseInputTransport from pipecat.transports.base_output import BaseOutputTransport from pipecat.transports.base_transport import BaseTransport, TransportParams from pipecat.transports.network.webrtc_connection import SmallWebRTCConnection +from pipecat.utils.watchdog_async_iterator import WatchdogAsyncIterator try: import cv2 @@ -422,19 +423,22 @@ class SmallWebRTCInputTransport(BaseInputTransport): async def _receive_audio(self): try: - async for audio_frame in self._client.read_audio_frame(): - self.start_watchdog() + audio_iterator = self._client.read_audio_frame() + async for audio_frame in WatchdogAsyncIterator( + audio_iterator, reseter=self, watchdog_enabled=self.watchdog_timers_enabled + ): if audio_frame: await self.push_audio_frame(audio_frame) - self.reset_watchdog() except Exception as e: logger.error(f"{self} exception receiving data: {e.__class__.__name__} ({e})") async def _receive_video(self): try: - async for video_frame in self._client.read_video_frame(): - self.start_watchdog() + video_iterator = self._client.read_video_frame() + async for video_frame in WatchdogAsyncIterator( + video_iterator, reseter=self, watchdog_enabled=self.watchdog_timers_enabled + ): if video_frame: await self.push_video_frame(video_frame) @@ -453,7 +457,6 @@ class SmallWebRTCInputTransport(BaseInputTransport): await self.push_video_frame(image_frame) # Remove from pending requests del self._image_requests[req_id] - self.reset_watchdog() except Exception as e: logger.error(f"{self} exception receiving data: {e.__class__.__name__} ({e})") diff --git a/src/pipecat/transports/services/daily.py b/src/pipecat/transports/services/daily.py index 776da1693..f92d632ca 100644 --- a/src/pipecat/transports/services/daily.py +++ b/src/pipecat/transports/services/daily.py @@ -40,6 +40,8 @@ from pipecat.transports.base_input import BaseInputTransport from pipecat.transports.base_output import BaseOutputTransport from pipecat.transports.base_transport import BaseTransport, TransportParams from pipecat.utils.asyncio import BaseTaskManager +from pipecat.utils.watchdog_queue import WatchdogQueue +from pipecat.utils.watchdog_reseter import WatchdogReseter try: from daily import ( @@ -250,7 +252,7 @@ class DailyAudioTrack: track: CustomAudioTrack -class DailyTransportClient(EventHandler): +class DailyTransportClient(WatchdogReseter, EventHandler): """Core client for interacting with Daily's API. Manages the connection to Daily rooms and handles all low-level API interactions. @@ -304,6 +306,7 @@ class DailyTransportClient(EventHandler): self._leave_counter = 0 self._task_manager: Optional[BaseTaskManager] = None + self._watchdog_timers_enabled = False # We use the executor to cleanup the client. We just do it from one # place, so only one thread is really needed. @@ -320,9 +323,6 @@ class DailyTransportClient(EventHandler): # waits for it to finish using completions (and a future) we will # deadlock because completions use event handlers (which are holding the # GIL). - self._event_queue = asyncio.Queue() - self._audio_queue = asyncio.Queue() - self._video_queue = asyncio.Queue() self._event_task = None self._audio_task = None self._video_task = None @@ -395,11 +395,18 @@ class DailyTransportClient(EventHandler): if not frame.transport_destination and self._camera: self._camera.write_frame(frame.image) + def reset_watchdog(self): + if self._task_manager: + self._task_manager.reset_watchdog(asyncio.current_task()) + async def setup(self, setup: FrameProcessorSetup): if self._task_manager: return self._task_manager = setup.task_manager + self._watchdog_timers_enabled = setup.watchdog_timers_enabled + + self._event_queue = WatchdogQueue(self, watchdog_enabled=self._watchdog_timers_enabled) self._event_task = self._task_manager.create_task( self._callback_task_handler(self._event_queue), f"{self}::event_callback_task", @@ -424,12 +431,14 @@ class DailyTransportClient(EventHandler): self._out_sample_rate = self._params.audio_out_sample_rate or frame.audio_out_sample_rate if self._params.audio_in_enabled and not self._audio_task and self._task_manager: + self._audio_queue = WatchdogQueue(self, watchdog_enabled=self._watchdog_timers_enabled) self._audio_task = self._task_manager.create_task( self._callback_task_handler(self._audio_queue), f"{self}::audio_callback_task", ) if self._params.video_in_enabled and not self._video_task and self._task_manager: + self._video_queue = WatchdogQueue(self, watchdog_enabled=self._watchdog_timers_enabled) self._video_task = self._task_manager.create_task( self._callback_task_handler(self._video_queue), f"{self}::video_callback_task", @@ -934,6 +943,7 @@ class DailyTransportClient(EventHandler): await self._joined_event.wait() (callback, *args) = await queue.get() await callback(*args) + queue.task_done() def _get_event_loop(self) -> asyncio.AbstractEventLoop: if not self._task_manager: diff --git a/src/pipecat/transports/services/livekit.py b/src/pipecat/transports/services/livekit.py index ecb2718c8..67ea6b32a 100644 --- a/src/pipecat/transports/services/livekit.py +++ b/src/pipecat/transports/services/livekit.py @@ -28,6 +28,7 @@ from pipecat.transports.base_input import BaseInputTransport from pipecat.transports.base_output import BaseOutputTransport from pipecat.transports.base_transport import BaseTransport, TransportParams from pipecat.utils.asyncio import BaseTaskManager +from pipecat.utils.watchdog_async_iterator import WatchdogAsyncIterator try: from livekit import rtc @@ -341,8 +342,9 @@ class LiveKitTransportClient: logger.warning(f"Received unexpected event type: {type(event)}") async def get_next_audio_frame(self): - frame, participant_id = await self._audio_queue.get() - return frame, participant_id + while True: + frame, participant_id = await self._audio_queue.get() + yield frame, participant_id def __str__(self): return f"{self._transport_name}::LiveKitTransportClient" @@ -413,9 +415,10 @@ class LiveKitInputTransport(BaseInputTransport): async def _audio_in_task_handler(self): logger.info("Audio input task started") - while True: - audio_data = await self._client.get_next_audio_frame() - self.start_watchdog() + audio_iterator = self._client.get_next_audio_frame() + async for audio_data in WatchdogAsyncIterator( + audio_iterator, reseter=self, watchdog_enabled=self.watchdog_timers_enabled + ): if audio_data: audio_frame_event, participant_id = audio_data pipecat_audio_frame = await self._convert_livekit_audio_to_pipecat( @@ -428,7 +431,6 @@ class LiveKitInputTransport(BaseInputTransport): num_channels=pipecat_audio_frame.num_channels, ) await self.push_audio_frame(input_audio_frame) - self.reset_watchdog() async def _convert_livekit_audio_to_pipecat( self, audio_frame_event: rtc.AudioFrameEvent diff --git a/src/pipecat/utils/asyncio.py b/src/pipecat/utils/asyncio.py index 36479edd4..1b8a3221f 100644 --- a/src/pipecat/utils/asyncio.py +++ b/src/pipecat/utils/asyncio.py @@ -8,7 +8,7 @@ import asyncio import time from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import Coroutine, Dict, List, Optional, Sequence +from typing import Coroutine, Dict, Optional, Sequence from loguru import logger @@ -18,6 +18,7 @@ WATCHDOG_TIMEOUT = 5.0 @dataclass class TaskManagerParams: loop: asyncio.AbstractEventLoop + enable_watchdog_timers: bool = False enable_watchdog_logging: bool = False watchdog_timeout: float = WATCHDOG_TIMEOUT @@ -27,10 +28,6 @@ class BaseTaskManager(ABC): def setup(self, params: TaskManagerParams): pass - @abstractmethod - async def cleanup(self): - pass - @abstractmethod def get_event_loop(self) -> asyncio.AbstractEventLoop: pass @@ -42,6 +39,7 @@ class BaseTaskManager(ABC): name: str, *, enable_watchdog_logging: Optional[bool] = None, + enable_watchdog_timers: Optional[bool] = None, watchdog_timeout: Optional[float] = None, ) -> asyncio.Task: """ @@ -54,6 +52,7 @@ class BaseTaskManager(ABC): coroutine (Coroutine): The coroutine to be executed within the task. name (str): The name to assign to the task for identification. enable_watchdog_logging(bool): whether this task should log watchdog processing times. + enable_watchdog_timers(bool): whether this task should have a watchdog timer. watchdog_timeout(float): watchdog timer timeout for this task. Returns: @@ -97,14 +96,6 @@ class BaseTaskManager(ABC): """Returns the list of currently created/registered tasks.""" pass - @abstractmethod - def start_watchdog(self, task: asyncio.Task): - """Starts the given task watchdog timer. If not reset, a warning will be - logged indicating the task is stalling. - - """ - pass - @abstractmethod def reset_watchdog(self, task: asyncio.Task): """Resets the given task watchdog timer. If not reset, a warning will be @@ -117,31 +108,22 @@ class BaseTaskManager(ABC): @dataclass class TaskData: task: asyncio.Task - watchdog_start: asyncio.Event watchdog_timer: asyncio.Event enable_watchdog_logging: bool + enable_watchdog_timers: bool watchdog_timeout: float + watchdog_task: Optional[asyncio.Task] class TaskManager(BaseTaskManager): def __init__(self) -> None: self._tasks: Dict[str, TaskData] = {} self._params: Optional[TaskManagerParams] = None - self._watchdog_tasks: List[asyncio.Task] = [] def setup(self, params: TaskManagerParams): if not self._params: self._params = params - async def cleanup(self): - for task in self._watchdog_tasks: - try: - task.cancel() - await task - except asyncio.CancelledError: - # This is expected, no need to re-raise. - pass - def get_event_loop(self) -> asyncio.AbstractEventLoop: if not self._params: raise Exception("TaskManager is not setup: unable to get event loop") @@ -153,6 +135,7 @@ class TaskManager(BaseTaskManager): name: str, *, enable_watchdog_logging: Optional[bool] = None, + enable_watchdog_timers: Optional[bool] = None, watchdog_timeout: Optional[float] = None, ) -> asyncio.Task: """ @@ -165,6 +148,7 @@ class TaskManager(BaseTaskManager): coroutine (Coroutine): The coroutine to be executed within the task. name (str): The name to assign to the task for identification. enable_watchdog_logging(bool): whether this task should log watchdog processing time. + enable_watchdog_timers(bool): whether this task should have a watchdog timer. watchdog_timeout(float): watchdog timer timeout for this task. Returns: @@ -186,20 +170,26 @@ class TaskManager(BaseTaskManager): task = self._params.loop.create_task(run_coroutine()) task.set_name(name) + task.add_done_callback(self._task_done_handler) self._add_task( TaskData( task=task, - watchdog_start=asyncio.Event(), watchdog_timer=asyncio.Event(), enable_watchdog_logging=( enable_watchdog_logging if enable_watchdog_logging else self._params.enable_watchdog_logging ), + enable_watchdog_timers=( + enable_watchdog_timers + if enable_watchdog_timers + else self._params.enable_watchdog_timers + ), watchdog_timeout=( watchdog_timeout if watchdog_timeout else self._params.watchdog_timeout ), - ) + watchdog_task=None, + ), ) logger.trace(f"{name}: task created") return task @@ -230,8 +220,6 @@ class TaskManager(BaseTaskManager): raise except Exception as e: logger.exception(f"{name}: unexpected exception while stopping task: {e}") - finally: - self._remove_task(task) async def cancel_task(self, task: asyncio.Task, timeout: Optional[float] = None): """Cancels the given asyncio Task and awaits its completion with an @@ -264,82 +252,59 @@ class TaskManager(BaseTaskManager): except BaseException as e: logger.critical(f"{name}: fatal base exception while cancelling task: {e}") raise - finally: - self._remove_task(task) def current_tasks(self) -> Sequence[asyncio.Task]: """Returns the list of currently created/registered tasks.""" return [data.task for data in self._tasks.values()] - def start_watchdog(self, task: asyncio.Task): - """Starts the given task watchdog timer. If not reset, a warning will be - logged indicating the task is stalling. If the timer was already started - a warning will be logged. - - """ - name = task.get_name() - if name in self._tasks: - if self._tasks[name].watchdog_start.is_set(): - logger.warning(f"Watchdog timer for task {name} already started") - else: - self._tasks[name].watchdog_timer.clear() - self._tasks[name].watchdog_start.set() - else: - logger.warning(f"Unable to start watchdog timer: task {name} does not exist") - def reset_watchdog(self, task: asyncio.Task): - """Resets the given task watchdog timer. If not reset, a warning will be - logged indicating the task is stalling. + """Resets the given task watchdog timer. If not reset on time, a warning + will be logged indicating the task is stalling. """ name = task.get_name() if name in self._tasks: - self._tasks[name].watchdog_start.clear() - self._tasks[name].watchdog_timer.set() + if self._tasks[name].enable_watchdog_timers: + self._tasks[name].watchdog_timer.set() else: logger.warning(f"Unable to reset watchdog timer: task {name} does not exist") def _add_task(self, task_data: TaskData): name = task_data.task.get_name() self._tasks[name] = task_data - watchdog_task = self.get_event_loop().create_task( - self._watchdog_task_handler(self._tasks[name]) - ) - self._watchdog_tasks.append(watchdog_task) - - def _remove_task(self, task: asyncio.Task): - name = task.get_name() - try: - del self._tasks[name] - except KeyError as e: - logger.trace(f"{name}: unable to remove task (already removed?): {e}") + if self._params and task_data.enable_watchdog_timers: + watchdog_task = self.get_event_loop().create_task( + self._watchdog_task_handler(task_data) + ) + task_data.watchdog_task = watchdog_task async def _watchdog_task_handler(self, task_data: TaskData): name = task_data.task.get_name() - start = task_data.watchdog_start timer = task_data.watchdog_timer enable_watchdog_logging = task_data.enable_watchdog_logging watchdog_timeout = task_data.watchdog_timeout - async def wait_for_reset(): - waiting = True - while waiting: - try: - start_time = time.time() - await asyncio.wait_for(timer.wait(), timeout=watchdog_timeout) - total_time = time.time() - start_time - if enable_watchdog_logging: - logger.debug(f"{name} task processing time: {total_time:.20f}") - waiting = False - except asyncio.TimeoutError: - logger.warning( - f"{name}: task is taking too long {WATCHDOG_TIMEOUT} second(s) (forgot to reset watchdog?)" - ) - finally: - timer.clear() - while True: - # Wait for the user to start the watchdog timer. - await start.wait() - # Now, waiting for the task to finish. - await wait_for_reset() + try: + start_time = time.time() + await asyncio.wait_for(timer.wait(), timeout=watchdog_timeout) + total_time = time.time() - start_time + if enable_watchdog_logging: + logger.debug(f"{name} time between watchdog timer resets: {total_time:.20f}") + except asyncio.TimeoutError: + logger.warning( + f"{name}: task is taking too long {WATCHDOG_TIMEOUT} second(s) (forgot to reset watchdog?)" + ) + finally: + timer.clear() + + def _task_done_handler(self, task: asyncio.Task): + name = task.get_name() + try: + task_data = self._tasks[name] + if task_data.watchdog_task: + task_data.watchdog_task.cancel() + task_data.watchdog_task = None + del self._tasks[name] + except KeyError as e: + logger.trace(f"{name}: unable to remove task data (already removed?): {e}") diff --git a/src/pipecat/utils/watchdog_async_iterator.py b/src/pipecat/utils/watchdog_async_iterator.py new file mode 100644 index 000000000..62b6d1a23 --- /dev/null +++ b/src/pipecat/utils/watchdog_async_iterator.py @@ -0,0 +1,74 @@ +# +# Copyright (c) 2024–2025, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +import asyncio +from typing import AsyncIterator, Optional + +from pipecat.utils.watchdog_reseter import WatchdogReseter + + +class WatchdogAsyncIterator: + """An asynchronous iterator that monitors activity and resets the current + task watchdog timer. This is necessary to avoid task watchdog timers to + expire while we are waiting to get an item from the iterator. + + """ + + def __init__( + self, + async_iterable, + *, + reseter: WatchdogReseter, + timeout: float = 2.0, + watchdog_enabled: bool = False, + ): + self._async_iterable = async_iterable + self._reseter = reseter + self._timeout = timeout + self._iter: Optional[AsyncIterator] = None + self._current_anext_task: Optional[asyncio.Task] = None + self._watchdog_enabled = watchdog_enabled + + def __aiter__(self): + return self + + async def __anext__(self): + if not self._iter: + self._iter = await self._ensure_async_iterator(self._async_iterable) + + if self._watchdog_enabled: + return await self._watchdog_anext() + else: + return await self._iter.__anext__() + + async def _watchdog_anext(self): + while True: + try: + if not self._current_anext_task: + self._current_anext_task = asyncio.create_task(self._iter.__anext__()) + + item = await asyncio.wait_for( + asyncio.shield(self._current_anext_task), + timeout=self._timeout, + ) + + self._reseter.reset_watchdog() + + # The task has finish, so we will create a new one for th next item. + self._current_anext_task = None + + return item + except asyncio.TimeoutError: + self._reseter.reset_watchdog() + except StopAsyncIteration: + self._current_anext_task = None + raise + + async def _ensure_async_iterator(self, obj) -> AsyncIterator: + aiter = obj.__aiter__() + if asyncio.iscoroutine(aiter): + aiter = await aiter + return aiter diff --git a/src/pipecat/utils/watchdog_event.py b/src/pipecat/utils/watchdog_event.py new file mode 100644 index 000000000..001a0cf26 --- /dev/null +++ b/src/pipecat/utils/watchdog_event.py @@ -0,0 +1,44 @@ +# +# Copyright (c) 2024–2025, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +import asyncio + +from pipecat.utils.watchdog_reseter import WatchdogReseter + + +class WatchdogEvent(asyncio.Event): + """An asynchronous event that resets the current task watchdog timer. This + is necessary to avoid task watchdog timers to expire while we are waiting on + the event. + + """ + + def __init__( + self, + reseter: WatchdogReseter, + *, + timeout: float = 2.0, + watchdog_enabled: bool = False, + ) -> None: + super().__init__() + self._reseter = reseter + self._timeout = timeout + self._watchdog_enabled = watchdog_enabled + + async def wait(self): + if self._watchdog_enabled: + return await self._watchdog_wait() + else: + return await super().wait() + + async def _watchdog_wait(self): + while True: + try: + await asyncio.wait_for(super().wait(), timeout=self._timeout) + self._reseter.reset_watchdog() + return True + except asyncio.TimeoutError: + self._reseter.reset_watchdog() diff --git a/src/pipecat/utils/watchdog_priority_queue.py b/src/pipecat/utils/watchdog_priority_queue.py new file mode 100644 index 000000000..a3635667c --- /dev/null +++ b/src/pipecat/utils/watchdog_priority_queue.py @@ -0,0 +1,50 @@ +# +# Copyright (c) 2024–2025, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +import asyncio + +from pipecat.utils.watchdog_reseter import WatchdogReseter + + +class WatchdogPriorityQueue(asyncio.PriorityQueue): + """An asynchronous priority queue that resets the current task watchdog + timer. This is necessary to avoid task watchdog timers to expire while we + are waiting to get an item from the queue. + + """ + + def __init__( + self, + reseter: WatchdogReseter, + *, + maxsize: int = 0, + timeout: float = 2.0, + watchdog_enabled: bool = False, + ) -> None: + super().__init__(maxsize) + self._reseter = reseter + self._timeout = timeout + self._watchdog_enabled = watchdog_enabled + + async def get(self): + if self._watchdog_enabled: + return await self._watchdog_get() + else: + return await super().get() + + def task_done(self): + if self._watchdog_enabled: + self._reseter.reset_watchdog() + super().task_done() + + async def _watchdog_get(self): + while True: + try: + item = await asyncio.wait_for(super().get(), timeout=self._timeout) + self._reseter.reset_watchdog() + return item + except asyncio.TimeoutError: + self._reseter.reset_watchdog() diff --git a/src/pipecat/utils/watchdog_queue.py b/src/pipecat/utils/watchdog_queue.py new file mode 100644 index 000000000..6eb9dab10 --- /dev/null +++ b/src/pipecat/utils/watchdog_queue.py @@ -0,0 +1,50 @@ +# +# Copyright (c) 2024–2025, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +import asyncio + +from pipecat.utils.watchdog_reseter import WatchdogReseter + + +class WatchdogQueue(asyncio.Queue): + """An asynchronous queue that resets the current task watchdog timer. This + is necessary to avoid task watchdog timers to expire while we are waiting to + get an item from the queue. + + """ + + def __init__( + self, + reseter: WatchdogReseter, + *, + maxsize: int = 0, + timeout: float = 2.0, + watchdog_enabled: bool = False, + ) -> None: + super().__init__(maxsize) + self._reseter = reseter + self._timeout = timeout + self._watchdog_enabled = watchdog_enabled + + async def get(self): + if self._watchdog_enabled: + return await self._watchdog_get() + else: + return await super().get() + + def task_done(self): + if self._watchdog_enabled: + self._reseter.reset_watchdog() + super().task_done() + + async def _watchdog_get(self): + while True: + try: + item = await asyncio.wait_for(super().get(), timeout=self._timeout) + self._reseter.reset_watchdog() + return item + except asyncio.TimeoutError: + self._reseter.reset_watchdog() diff --git a/src/pipecat/utils/watchdog_reseter.py b/src/pipecat/utils/watchdog_reseter.py new file mode 100644 index 000000000..ee70207b3 --- /dev/null +++ b/src/pipecat/utils/watchdog_reseter.py @@ -0,0 +1,13 @@ +# +# Copyright (c) 2024–2025, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +from abc import ABC, abstractmethod + + +class WatchdogReseter(ABC): + @abstractmethod + def reset_watchdog(self): + pass