diff --git a/CHANGELOG.md b/CHANGELOG.md index 95350157f..d8be39cdb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,11 @@ 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 were still not ready. 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" 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}")