From 9cc498b1fa17e2571092084d5dc6d3762fbe0bfd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= Date: Mon, 5 May 2025 21:27:49 -0700 Subject: [PATCH 1/2] TaskManager: use a dictionary instead of a set to store tasks --- CHANGELOG.md | 2 ++ src/pipecat/utils/asyncio.py | 15 ++++++++------- 2 files changed, 10 insertions(+), 7 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 95350157f..ac4e2db11 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed +- Fixed a `TaskManager` that was causing dangling tasks to be reported. + - Fixed an issue that could cause data to be sent to the transports when they were still not ready. diff --git a/src/pipecat/utils/asyncio.py b/src/pipecat/utils/asyncio.py index acc4acec8..cea447329 100644 --- a/src/pipecat/utils/asyncio.py +++ b/src/pipecat/utils/asyncio.py @@ -6,7 +6,7 @@ import asyncio from abc import ABC, abstractmethod -from typing import Coroutine, Optional, Set +from typing import Coroutine, Dict, Optional, Sequence, Set from loguru import logger @@ -69,14 +69,14 @@ class BaseTaskManager(ABC): pass @abstractmethod - def current_tasks(self) -> Set[asyncio.Task]: + def current_tasks(self) -> Sequence[asyncio.Task]: """Returns the list of currently created/registered tasks.""" pass class TaskManager(BaseTaskManager): def __init__(self) -> None: - self._tasks: Set[asyncio.Task] = set() + self._tasks: Dict[str, asyncio.Task] = {} self._loop: Optional[asyncio.AbstractEventLoop] = None def set_event_loop(self, loop: asyncio.AbstractEventLoop): @@ -179,16 +179,17 @@ class TaskManager(BaseTaskManager): finally: self._remove_task(task) - def current_tasks(self) -> Set[asyncio.Task]: + def current_tasks(self) -> Sequence[asyncio.Task]: """Returns the list of currently created/registered tasks.""" - return self._tasks + return list(self._tasks.values()) def _add_task(self, task: asyncio.Task): - self._tasks.add(task) + name = task.get_name() + self._tasks[name] = task def _remove_task(self, task: asyncio.Task): name = task.get_name() try: - self._tasks.remove(task) + del self._tasks[name] except KeyError as e: logger.trace(f"{name}: unable to remove task (already removed?): {e}") From 45839053135f8f6dac67ea448fdd4f1909d17e95 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= Date: Mon, 5 May 2025 21:33:21 -0700 Subject: [PATCH 2/2] PipelineTask: cleanup if task is cancelled from outside Pipecat --- CHANGELOG.md | 3 +++ src/pipecat/pipeline/task.py | 31 ++++++++++++++++++++----------- 2 files changed, 23 insertions(+), 11 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ac4e2db11..d8be39cdb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed +- Fixed a `PipelineTask` issue that would cause tasks to not be cancelled if + task was cancelled from outside of Pipecat. + - Fixed a `TaskManager` that was causing dangling tasks to be reported. - Fixed an issue that could cause data to be sent to the transports when they diff --git a/src/pipecat/pipeline/task.py b/src/pipecat/pipeline/task.py index 8279373cb..c40173899 100644 --- a/src/pipecat/pipeline/task.py +++ b/src/pipecat/pipeline/task.py @@ -286,12 +286,7 @@ class PipelineTask(BaseTask): 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()) - # Only cancel the push task. Everything else will be cancelled in run(). - await self._task_manager.cancel_task(self._process_push_task) + await self._cancel() async def run(self): """Starts and manages the pipeline execution until completion or cancellation.""" @@ -309,11 +304,17 @@ class PipelineTask(BaseTask): # well, because you get a CancelledError in every place you are # awaiting a task. pass - await self._cancel_tasks() - await self._cleanup(cleanup_pipeline) - if self._check_dangling_tasks: - self._print_dangling_tasks() - self._finished = True + finally: + # It's possibe that we get an asyncio.CancelledError from the + # outside, if so we need to make sure everything gets cancelled + # properly. + if cleanup_pipeline: + await self._cancel() + await self._cancel_tasks() + await self._cleanup(cleanup_pipeline) + if self._check_dangling_tasks: + self._print_dangling_tasks() + self._finished = True async def queue_frame(self, frame: Frame): """Queue a single frame to be pushed down the pipeline. @@ -336,6 +337,14 @@ class PipelineTask(BaseTask): for frame in frames: await self.queue_frame(frame) + async def _cancel(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()) + # Only cancel the push task. Everything else will be cancelled in run(). + await self._task_manager.cancel_task(self._process_push_task) + async def _create_tasks(self): self._process_up_task = self._task_manager.create_task( self._process_up_queue(), f"{self}::_process_up_queue"