Merge pull request #1748 from pipecat-ai/aleix/task-manager-dictionary

task manager dictionary and cleanup PipelineTask

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.
 

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"

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}")