PipelineTask: improve task cancellation

This commit is contained in:
Aleix Conchillo Flaqué
2025-08-05 11:25:29 -07:00
parent 61cb45d61b
commit 49a5a1e375
3 changed files with 25 additions and 14 deletions

View File

@@ -459,15 +459,20 @@ class PipelineTask(BasePipelineTask):
# awaiting a task. # awaiting a task.
pass pass
finally: finally:
# It's possibe that we get an asyncio.CancelledError from the # We only cancel things cleanly if we know we are the ones
# outside, if so we need to make sure everything gets cancelled # cancelling. It's possibe that we get an asyncio.CancelledError
# properly. # from the outside, in which case it is very likely other tasks have
if cleanup_pipeline: # been already cancelled (e.g. when python is shutting down) so we
await self._cancel() # can't assume things are being cancelled nicely.
await self._cancel_tasks() if self._cancelled:
await self._cleanup(cleanup_pipeline) await self._cancel_tasks()
if self._check_dangling_tasks: await self._cleanup(cleanup_pipeline)
self._print_dangling_tasks() if self._check_dangling_tasks:
self._print_dangling_tasks()
else:
logger.warning(
f"Pipeline task {self} is not being cancelled properly (use cancel() method)"
)
self._finished = True self._finished = True
async def queue_frame(self, frame: Frame): async def queue_frame(self, frame: Frame):
@@ -494,7 +499,7 @@ class PipelineTask(BasePipelineTask):
async def _cancel(self): async def _cancel(self):
"""Internal cancellation logic for the pipeline task.""" """Internal cancellation logic for the pipeline task."""
if not self._cancelled: if not self._cancelled:
logger.debug(f"Canceling pipeline task {self}") logger.debug(f"Cancelling pipeline task {self}")
self._cancelled = True self._cancelled = True
# Make sure everything is cleaned up downstream. This is sent # Make sure everything is cleaned up downstream. This is sent
# out-of-band from the main streaming task which is what we want since # out-of-band from the main streaming task which is what we want since
@@ -502,7 +507,9 @@ class PipelineTask(BasePipelineTask):
await self._source.push_frame(CancelFrame()) await self._source.push_frame(CancelFrame())
# Wait for CancelFrame to make it throught the pipeline. # Wait for CancelFrame to make it throught the pipeline.
await self._wait_for_pipeline_end() await self._wait_for_pipeline_end()
# Only cancel the push task. Everything else will be cancelled in run(). # 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: if self._process_push_task:
await self._task_manager.cancel_task(self._process_push_task) await self._task_manager.cancel_task(self._process_push_task)
self._process_push_task = None self._process_push_task = None
@@ -544,6 +551,10 @@ class PipelineTask(BasePipelineTask):
"""Cancel all running pipeline tasks.""" """Cancel all running pipeline tasks."""
await self._observer.stop() await self._observer.stop()
if self._process_push_task:
await self._task_manager.cancel_task(self._process_push_task)
self._process_push_task = None
if self._process_up_task: if self._process_up_task:
await self._task_manager.cancel_task(self._process_up_task) await self._task_manager.cancel_task(self._process_up_task)
self._process_up_task = None self._process_up_task = None

View File

@@ -556,7 +556,7 @@ class FrameProcessor(BaseObject):
# #
async def _start_interruption(self): async def _start_interruption(self):
"""Start handling an interruption by canceling current tasks.""" """Start handling an interruption by cancelling current tasks."""
try: try:
# Cancel the input task. This will stop processing queued frames. # Cancel the input task. This will stop processing queued frames.
await self.__cancel_input_task() await self.__cancel_input_task()

View File

@@ -183,7 +183,7 @@ class HeyGenVideoService(AIService):
async def stop(self, frame: EndFrame): async def stop(self, frame: EndFrame):
"""Stop the HeyGen video service gracefully. """Stop the HeyGen video service gracefully.
Performs cleanup by ending the conversation and canceling ongoing tasks Performs cleanup by ending the conversation and cancelling ongoing tasks
in a controlled manner. in a controlled manner.
Args: Args:
@@ -241,7 +241,7 @@ class HeyGenVideoService(AIService):
Manages the interruption flow by: Manages the interruption flow by:
1. Setting the interruption flag 1. Setting the interruption flag
2. Signaling the client to interrupt current speech 2. Signaling the client to interrupt current speech
3. Canceling ongoing audio sending tasks 3. Cancelling ongoing audio sending tasks
4. Creating a new send task 4. Creating a new send task
5. Activating the avatar's listening animation 5. Activating the avatar's listening animation
""" """