Merge pull request #2872 from pipecat-ai/aleix/pipeline-task-cancellation-fixes

PipelineTask: fix task cancellation issues

CHANGELOG.md

@@ -27,12 +27,18 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - `CartesiaSTTService` now inherits from `WebsocketSTTService`.
 
-- Package upgrades:
+# Package upgrades:
   - `openai` upgraded to support up to 2.x.x.
   - `openpipe` upgraded to support up to 5.x.x.
 
 ### Fixed
 
+- Fixed multiple pipeline task cancellation issues. `asyncio.CancelledError` is
+  now handled properly in `PipelineTask` making it possible to cancel an asyncio
+  task that it's executing a `PipelineRunner` cleanly. Also,
+  `PipelineTask.cancel()` does not block anymore waiting for the `CancelFrame`
+  to reach the end of the pipeline (going back to the behavior in < 0.0.83).
+
 - Fixed an issue in `ElevenLabsTTSService` and `ElevenLabsHttpTTSService` where
   the Flash models would split words, resulting in a space being inserted
   between words.

src/pipecat/pipeline/runner.py

@@ -70,11 +70,15 @@ class PipelineRunner(BaseObject):
         """
         logger.debug(f"Runner {self} started running {task}")
         self._tasks[task.name] = task
-        params = PipelineTaskParams(loop=self._loop)
+
+        # PipelineTask handles asyncio.CancelledError to shutdown the pipeline
+        # properly and re-raises it in case there's more cleanup to do.
         try:
+            params = PipelineTaskParams(loop=self._loop)
             await task.run(params)
         except asyncio.CancelledError:
-            await self._cancel()
+            pass
+
         del self._tasks[task.name]
 
         # Cleanup base object.

src/pipecat/pipeline/task.py

@@ -269,6 +269,9 @@ class PipelineTask(BasePipelineTask):
         # StopFrame) has been received at the end of the pipeline.
         self._pipeline_end_event = asyncio.Event()
 
+        # This event is set when the pipeline truly finishes.
+        self._pipeline_finished_event = asyncio.Event()
+
         # This is the final pipeline. It is composed of a source processor,
         # followed by the user pipeline, and ending with a sink processor. The
         # source allows us to receive and react to upstream frames, and the sink
@@ -401,11 +404,7 @@ class PipelineTask(BasePipelineTask):
         await self.queue_frame(EndFrame())
 
     async def cancel(self):
-        """Immediately stop the running pipeline.
+        """Request the running pipeline to cancel."""
-
-        Cancels all running tasks and stops frame processing without
-        waiting for completion.
-        """
         if not self._finished:
             await self._cancel()
 
@@ -417,51 +416,38 @@ class PipelineTask(BasePipelineTask):
         """
         if self.has_finished():
             return
-        cleanup_pipeline = True
+
+        # Setup processors.
+        await self._setup(params)
+
+        # Create all main tasks and wait for the main push task. This is the
+        # task that pushes frames to the very beginning of our pipeline (i.e. to
+        # our controlled source processor).
+        await self._create_tasks()
+
         try:
-            # Setup processors.
+            # Wait for pipeline to finish.
-            await self._setup(params)
+            await self._wait_for_pipeline_finished()
-
-            # Create all main tasks and wait of the main push task. This is the
-            # task that pushes frames to the very beginning of our pipeline (our
-            # controlled source processor).
-            push_task = await self._create_tasks()
-            await push_task
-
-            # We have already cleaned up the pipeline inside the task.
-            cleanup_pipeline = False
-
-            # Pipeline has finished nicely.
-            self._finished = True
         except asyncio.CancelledError:
-            # Raise exception back to the pipeline runner so it can cancel this
+            logger.debug(f"Pipeline task {self} got cancelled from outside...")
-            # task properly.
+            # We have been cancelled from outside, let's just cancel everything.
+            await self._cancel()
+            # Wait again for pipeline to finish. This time we have really
+            # cancelled, so it should really finish.
+            await self._wait_for_pipeline_finished()
+            # Re-raise in case there's more cleanup to do.
             raise
         finally:
             # We can reach this point for different reasons:
             #
-            # 1. The task has finished properly (e.g. `EndFrame`).
+            # 1. The pipeline task has finished (try case).
-            # 2. By calling `PipelineTask.cancel()`.
+            # 2. By an asyncio task cancellation (except case).
-            # 3. By asyncio task cancellation.
+            logger.debug(f"Pipeline task {self} is finishing...")
-            #
+            await self._cancel_tasks()
-            # Case (1) will execute the code below without issues because
+            if self._check_dangling_tasks:
-            # `self._finished` is true.
+                self._print_dangling_tasks()
-            #
+            self._finished = True
-            # Case (2) will execute the code below without issues because
+            logger.debug(f"Pipeline task {self} has finished")
-            # `self._cancelled` is true.
-            #
-            # Case (3) will raise the exception above (because we are cancelling
-            # the asyncio task). This will be then captured by the
-            # `PipelineRunner` which will call `PipelineTask.cancel()` and
-            # therefore becoming case (2).
-            if self._finished or self._cancelled:
-                logger.debug(f"Pipeline task {self} is finishing cleanup...")
-                await self._cancel_tasks()
-                await self._cleanup(cleanup_pipeline)
-                if self._check_dangling_tasks:
-                    self._print_dangling_tasks()
-                self._finished = True
-                logger.debug(f"Pipeline task {self} has finished")
 
     async def queue_frame(self, frame: Frame):
         """Queue a single frame to be pushed down the pipeline.
@@ -489,19 +475,7 @@ class PipelineTask(BasePipelineTask):
         if not self._cancelled:
             logger.debug(f"Cancelling pipeline task {self}")
             self._cancelled = True
-            cancel_frame = CancelFrame()
+            await self.queue_frame(CancelFrame())
-            # 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._pipeline.queue_frame(cancel_frame)
-            # Wait for CancelFrame to make it through the pipeline.
-            await self._wait_for_pipeline_end(cancel_frame)
-            # Only cancel the push task, we don't want to be able to process any
-            # other frame after cancel. Everything else will be cancelled in
-            # run().
-            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):
         """Create and start all pipeline processing tasks."""
@@ -603,6 +577,17 @@ class PipelineTask(BasePipelineTask):
 
         self._pipeline_end_event.clear()
 
+        # We are really done.
+        self._pipeline_finished_event.set()
+
+    async def _wait_for_pipeline_finished(self):
+        await self._pipeline_finished_event.wait()
+        self._pipeline_finished_event.clear()
+        # Make sure we wait for the main task to complete.
+        if self._process_push_task:
+            await self._process_push_task
+            self._process_push_task = None
+
     async def _setup(self, params: PipelineTaskParams):
         """Set up the pipeline task and all processors."""
         mgr_params = TaskManagerParams(loop=params.loop)

tests/test_pipeline.py

@@ -254,7 +254,7 @@ class TestPipelineTask(unittest.IsolatedAsyncioTestCase):
 
         try:
             await asyncio.wait_for(
-                asyncio.shield(task.run(PipelineTaskParams(loop=asyncio.get_event_loop()))),
+                task.run(PipelineTaskParams(loop=asyncio.get_event_loop())),
                 timeout=1.0,
             )
         except asyncio.TimeoutError:
@@ -290,7 +290,7 @@ class TestPipelineTask(unittest.IsolatedAsyncioTestCase):
         await task.queue_frame(TextFrame(text="Hello!"))
         try:
             await asyncio.wait_for(
-                asyncio.shield(task.run(PipelineTaskParams(loop=asyncio.get_event_loop()))),
+                task.run(PipelineTaskParams(loop=asyncio.get_event_loop())),
                 timeout=1.0,
             )
         except asyncio.TimeoutError:
@@ -301,11 +301,8 @@ class TestPipelineTask(unittest.IsolatedAsyncioTestCase):
         identity = IdentityFilter()
         pipeline = Pipeline([identity])
         task = PipelineTask(pipeline, idle_timeout_secs=0.2)
-        try:
+        # This shouldn't freeze, so nothing to check really.
-            await task.run(PipelineTaskParams(loop=asyncio.get_event_loop()))
+        await task.run(PipelineTaskParams(loop=asyncio.get_event_loop()))
-            assert False
-        except asyncio.CancelledError:
-            assert True
 
     async def test_no_idle_task(self):
         identity = IdentityFilter()
@@ -313,7 +310,7 @@ class TestPipelineTask(unittest.IsolatedAsyncioTestCase):
         task = PipelineTask(pipeline, idle_timeout_secs=0.2, cancel_on_idle_timeout=False)
         try:
             await asyncio.wait_for(
-                asyncio.shield(task.run(PipelineTaskParams(loop=asyncio.get_event_loop()))),
+                task.run(PipelineTaskParams(loop=asyncio.get_event_loop())),
                 timeout=0.3,
             )
         except asyncio.TimeoutError:
@@ -332,11 +329,7 @@ class TestPipelineTask(unittest.IsolatedAsyncioTestCase):
             ),
             idle_timeout_secs=0.3,
         )
-        try:
+        await task.run(PipelineTaskParams(loop=asyncio.get_event_loop()))
-            await task.run(PipelineTaskParams(loop=asyncio.get_event_loop()))
-            assert False
-        except asyncio.CancelledError:
-            assert True
 
     async def test_idle_task_event_handler_no_frames(self):
         identity = IdentityFilter()
@@ -351,11 +344,8 @@ class TestPipelineTask(unittest.IsolatedAsyncioTestCase):
             idle_timeout = True
             await task.cancel()
 
-        try:
+        await task.run(PipelineTaskParams(loop=asyncio.get_event_loop()))
-            await task.run(PipelineTaskParams(loop=asyncio.get_event_loop()))
+        assert idle_timeout
-            assert False
-        except asyncio.CancelledError:
-            assert idle_timeout
 
     async def test_idle_task_event_handler_quiet_user(self):
         identity = IdentityFilter()
@@ -416,12 +406,15 @@ class TestPipelineTask(unittest.IsolatedAsyncioTestCase):
             asyncio.create_task(delayed_frames()),
         ]
 
-        await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
+        _, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
 
         diff_time = time.time() - start_time
 
         self.assertGreater(diff_time, sleep_time_secs * 3)
 
+        # Wait for the pending tasks to complete.
+        await asyncio.gather(*pending)
+
     async def test_task_cancel_timeout(self):
         class CancelFilter(FrameProcessor):
             def __init__(self, **kwargs):