From 8b5228a105ddf452d5ff220217403598d76c4bd2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= Date: Sun, 26 Jan 2025 18:49:55 -0800 Subject: [PATCH] utils: move task functions to asyncio module --- src/pipecat/pipeline/runner.py | 3 +- src/pipecat/pipeline/task.py | 3 +- src/pipecat/pipeline/task_observer.py | 3 +- src/pipecat/processors/frame_processor.py | 3 +- .../services/gemini_multimodal_live/gemini.py | 1 - src/pipecat/transports/base_output.py | 2 +- src/pipecat/transports/services/daily.py | 2 +- src/pipecat/transports/services/livekit.py | 2 +- src/pipecat/utils/asyncio.py | 114 ++++++++++++++++++ src/pipecat/utils/utils.py | 107 ---------------- 10 files changed, 125 insertions(+), 115 deletions(-) create mode 100644 src/pipecat/utils/asyncio.py diff --git a/src/pipecat/pipeline/runner.py b/src/pipecat/pipeline/runner.py index 8cdfc3eaa..9963454b6 100644 --- a/src/pipecat/pipeline/runner.py +++ b/src/pipecat/pipeline/runner.py @@ -10,7 +10,8 @@ import signal from loguru import logger from pipecat.pipeline.task import PipelineTask -from pipecat.utils.utils import current_tasks, obj_count, obj_id +from pipecat.utils.asyncio import current_tasks +from pipecat.utils.utils import obj_count, obj_id class PipelineRunner: diff --git a/src/pipecat/pipeline/task.py b/src/pipecat/pipeline/task.py index 58935ff06..45596980f 100644 --- a/src/pipecat/pipeline/task.py +++ b/src/pipecat/pipeline/task.py @@ -30,7 +30,8 @@ from pipecat.pipeline.base_pipeline import BasePipeline from pipecat.pipeline.base_task import BaseTask from pipecat.pipeline.task_observer import TaskObserver from pipecat.processors.frame_processor import FrameDirection, FrameProcessor -from pipecat.utils.utils import cancel_task, create_task, obj_count, obj_id, wait_for_task +from pipecat.utils.asyncio import cancel_task, create_task, wait_for_task +from pipecat.utils.utils import obj_count, obj_id HEARTBEAT_SECONDS = 1.0 HEARTBEAT_MONITOR_SECONDS = HEARTBEAT_SECONDS * 5 diff --git a/src/pipecat/pipeline/task_observer.py b/src/pipecat/pipeline/task_observer.py index da5a6818f..c9816f7d3 100644 --- a/src/pipecat/pipeline/task_observer.py +++ b/src/pipecat/pipeline/task_observer.py @@ -12,7 +12,8 @@ from attr import dataclass from pipecat.frames.frames import Frame from pipecat.observers.base_observer import BaseObserver from pipecat.processors.frame_processor import FrameDirection, FrameProcessor -from pipecat.utils.utils import cancel_task, create_task, obj_count, obj_id +from pipecat.utils.asyncio import cancel_task, create_task +from pipecat.utils.utils import obj_count, obj_id @dataclass diff --git a/src/pipecat/processors/frame_processor.py b/src/pipecat/processors/frame_processor.py index 6a8252677..c986d7025 100644 --- a/src/pipecat/processors/frame_processor.py +++ b/src/pipecat/processors/frame_processor.py @@ -24,7 +24,8 @@ from pipecat.frames.frames import ( ) from pipecat.metrics.metrics import LLMTokenUsage, MetricsData from pipecat.processors.metrics.frame_processor_metrics import FrameProcessorMetrics -from pipecat.utils.utils import cancel_task, create_task, obj_count, obj_id +from pipecat.utils.asyncio import cancel_task, create_task +from pipecat.utils.utils import obj_count, obj_id class FrameDirection(Enum): diff --git a/src/pipecat/services/gemini_multimodal_live/gemini.py b/src/pipecat/services/gemini_multimodal_live/gemini.py index c2fc103b7..770b64711 100644 --- a/src/pipecat/services/gemini_multimodal_live/gemini.py +++ b/src/pipecat/services/gemini_multimodal_live/gemini.py @@ -51,7 +51,6 @@ from pipecat.services.openai import ( OpenAIUserContextAggregator, ) from pipecat.utils.time import time_now_iso8601 -from pipecat.utils.utils import wait_for_task from . import events from .audio_transcriber import AudioTranscriber diff --git a/src/pipecat/transports/base_output.py b/src/pipecat/transports/base_output.py index 1ec5ea7fa..5182fc790 100644 --- a/src/pipecat/transports/base_output.py +++ b/src/pipecat/transports/base_output.py @@ -35,8 +35,8 @@ from pipecat.frames.frames import ( ) from pipecat.processors.frame_processor import FrameDirection, FrameProcessor from pipecat.transports.base_transport import TransportParams +from pipecat.utils.asyncio import wait_for_task from pipecat.utils.time import nanoseconds_to_seconds -from pipecat.utils.utils import wait_for_task class BaseOutputTransport(FrameProcessor): diff --git a/src/pipecat/transports/services/daily.py b/src/pipecat/transports/services/daily.py index 70df2f05a..6d32786d9 100644 --- a/src/pipecat/transports/services/daily.py +++ b/src/pipecat/transports/services/daily.py @@ -46,7 +46,7 @@ from pipecat.transcriptions.language import Language from pipecat.transports.base_input import BaseInputTransport from pipecat.transports.base_output import BaseOutputTransport from pipecat.transports.base_transport import BaseTransport, TransportParams -from pipecat.utils.utils import cancel_task, create_task +from pipecat.utils.asyncio import cancel_task, create_task try: from daily import CallClient, Daily, EventHandler diff --git a/src/pipecat/transports/services/livekit.py b/src/pipecat/transports/services/livekit.py index f0a2add0e..eff892ac0 100644 --- a/src/pipecat/transports/services/livekit.py +++ b/src/pipecat/transports/services/livekit.py @@ -28,7 +28,7 @@ from pipecat.processors.frame_processor import FrameDirection from pipecat.transports.base_input import BaseInputTransport from pipecat.transports.base_output import BaseOutputTransport from pipecat.transports.base_transport import BaseTransport, TransportParams -from pipecat.utils.utils import create_task +from pipecat.utils.asyncio import create_task try: from livekit import rtc diff --git a/src/pipecat/utils/asyncio.py b/src/pipecat/utils/asyncio.py new file mode 100644 index 000000000..d10895758 --- /dev/null +++ b/src/pipecat/utils/asyncio.py @@ -0,0 +1,114 @@ +# +# Copyright (c) 2024–2025, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +import asyncio +from typing import Coroutine, Optional, Set + +from loguru import logger + +_TASKS: Set[asyncio.Task] = set() + + +def create_task(loop: asyncio.AbstractEventLoop, coroutine: Coroutine, name: str) -> asyncio.Task: + """ + Creates and schedules a new asyncio Task that runs the given coroutine. + + The task is added to a global set of created tasks. + + Args: + loop (asyncio.AbstractEventLoop): The event loop to use for creating the task. + coroutine (Coroutine): The coroutine to be executed within the task. + name (str): The name to assign to the task for identification. + + Returns: + asyncio.Task: The created task object. + """ + + async def run_coroutine(): + try: + await coroutine + except asyncio.CancelledError: + logger.trace(f"{name}: task cancelled") + # Re-raise the exception to ensure the task is cancelled. + raise + except Exception as e: + logger.exception(f"{name}: unexpected exception: {e}") + + task = loop.create_task(run_coroutine()) + task.set_name(name) + _TASKS.add(task) + logger.trace(f"{name}: task created") + return task + + +async def wait_for_task(task: asyncio.Task, timeout: Optional[float] = None): + """Wait for an asyncio.Task to complete with optional timeout handling. + + This function awaits the specified asyncio.Task and handles scenarios for + timeouts, cancellations, and other exceptions. It also ensures that the task + is removed from the set of registered tasks upon completion or failure. + + Args: + task (asyncio.Task): The asyncio Task to wait for. + timeout (Optional[float], optional): The maximum number of seconds + to wait for the task to complete. If None, waits indefinitely. + Defaults to None. + """ + name = task.get_name() + try: + if timeout: + await asyncio.wait_for(task, timeout=timeout) + else: + await task + except asyncio.TimeoutError: + logger.warning(f"{name}: timed out waiting for task to finish") + except asyncio.CancelledError: + logger.trace(f"{name}: unexpected task cancellation (maybe Ctrl-C?)") + except Exception as e: + logger.exception(f"{name}: unexpected exception while stopping task: {e}") + finally: + try: + _TASKS.remove(task) + except KeyError as e: + logger.trace(f"{name}: unable to remove task (already removed?): {e}") + + +async def cancel_task(task: asyncio.Task, timeout: Optional[float] = None): + """Cancels the given asyncio Task and awaits its completion with an + optional timeout. + + This function removes the task from the set of registered tasks upon + completion or failure. + + Args: + task (asyncio.Task): The task to be cancelled. + timeout (Optional[float]): The optional timeout in seconds to wait for the task to cancel. + + """ + name = task.get_name() + task.cancel() + try: + if timeout: + await asyncio.wait_for(task, timeout=timeout) + else: + await task + except asyncio.TimeoutError: + logger.warning(f"{name}: timed out waiting for task to cancel") + except asyncio.CancelledError: + # Here are sure the task is cancelled properly. + pass + except Exception as e: + logger.exception(f"{name}: unexpected exception while cancelling task: {e}") + finally: + try: + _TASKS.remove(task) + except KeyError as e: + logger.trace(f"{name}: unable to remove task (already removed?): {e}") + + +def current_tasks() -> Set[asyncio.Task]: + """Returns the list of currently created/registered tasks.""" + return _TASKS diff --git a/src/pipecat/utils/utils.py b/src/pipecat/utils/utils.py index 435131684..4cdb73050 100644 --- a/src/pipecat/utils/utils.py +++ b/src/pipecat/utils/utils.py @@ -4,16 +4,11 @@ # SPDX-License-Identifier: BSD 2-Clause License # -import asyncio import collections import itertools -from typing import Coroutine, Optional, Set - -from loguru import logger _COUNTS = collections.defaultdict(itertools.count) _ID = itertools.count() -_TASKS: Set[asyncio.Task] = set() def obj_id() -> int: @@ -41,105 +36,3 @@ def obj_count(obj) -> int: 0 """ return next(_COUNTS[obj.__class__.__name__]) - - -def create_task(loop: asyncio.AbstractEventLoop, coroutine: Coroutine, name: str) -> asyncio.Task: - """ - Creates and schedules a new asyncio Task that runs the given coroutine. - - The task is added to a global set of created tasks. - - Args: - loop (asyncio.AbstractEventLoop): The event loop to use for creating the task. - coroutine (Coroutine): The coroutine to be executed within the task. - name (str): The name to assign to the task for identification. - - Returns: - asyncio.Task: The created task object. - """ - - async def run_coroutine(): - try: - await coroutine - except asyncio.CancelledError: - logger.trace(f"{name}: task cancelled") - # Re-raise the exception to ensure the task is cancelled. - raise - except Exception as e: - logger.exception(f"{name}: unexpected exception: {e}") - - task = loop.create_task(run_coroutine()) - task.set_name(name) - _TASKS.add(task) - logger.trace(f"{name}: task created") - return task - - -async def wait_for_task(task: asyncio.Task, timeout: Optional[float] = None): - """Wait for an asyncio.Task to complete with optional timeout handling. - - This function awaits the specified asyncio.Task and handles scenarios for - timeouts, cancellations, and other exceptions. It also ensures that the task - is removed from the set of registered tasks upon completion or failure. - - Args: - task (asyncio.Task): The asyncio Task to wait for. - timeout (Optional[float], optional): The maximum number of seconds - to wait for the task to complete. If None, waits indefinitely. - Defaults to None. - """ - name = task.get_name() - try: - if timeout: - await asyncio.wait_for(task, timeout=timeout) - else: - await task - except asyncio.TimeoutError: - logger.warning(f"{name}: timed out waiting for task to finish") - except asyncio.CancelledError: - logger.trace(f"{name}: unexpected task cancellation (maybe Ctrl-C?)") - except Exception as e: - logger.exception(f"{name}: unexpected exception while stopping task: {e}") - finally: - try: - _TASKS.remove(task) - except KeyError as e: - logger.trace(f"{name}: unable to remove task (already removed?): {e}") - - -async def cancel_task(task: asyncio.Task, timeout: Optional[float] = None): - """Cancels the given asyncio Task and awaits its completion with an - optional timeout. - - This function removes the task from the set of registered tasks upon - completion or failure. - - Args: - task (asyncio.Task): The task to be cancelled. - timeout (Optional[float]): The optional timeout in seconds to wait for the task to cancel. - - """ - name = task.get_name() - task.cancel() - try: - if timeout: - await asyncio.wait_for(task, timeout=timeout) - else: - await task - except asyncio.TimeoutError: - logger.warning(f"{name}: timed out waiting for task to cancel") - except asyncio.CancelledError: - # Here are sure the task is cancelled properly. - pass - except Exception as e: - logger.exception(f"{name}: unexpected exception while cancelling task: {e}") - finally: - try: - _TASKS.remove(task) - except KeyError as e: - logger.trace(f"{name}: unable to remove task (already removed?): {e}") - - -def current_tasks() -> Set[asyncio.Task]: - """Returns the list of currently created/registered tasks.""" - return _TASKS