From 327973657fc3fef35af0e0a317a8fb31d3968668 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= Date: Wed, 25 Jun 2025 11:22:05 -0700 Subject: [PATCH] TaskManager: remove wathcdog timer when main task is done --- src/pipecat/utils/asyncio.py | 23 +++++++++++++---------- 1 file changed, 13 insertions(+), 10 deletions(-) diff --git a/src/pipecat/utils/asyncio.py b/src/pipecat/utils/asyncio.py index a2c75ceb6..a4b286273 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 @@ -164,6 +164,7 @@ 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, @@ -269,14 +270,6 @@ class TaskManager(BaseTaskManager): async def _remove_task(self, task: asyncio.Task): name = task.get_name() try: - task_data = self._tasks[name] - if task_data.watchdog_task: - try: - task_data.watchdog_task.cancel() - await task_data.watchdog_task - except asyncio.CancelledError: - pass - task_data.watchdog_task = None del self._tasks[name] except KeyError as e: logger.trace(f"{name}: unable to remove task (already removed?): {e}") @@ -293,10 +286,20 @@ class TaskManager(BaseTaskManager): 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}") + 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 + except KeyError as e: + logger.trace(f"{name}: unable to find task (already removed?): {e}")