digitalkin 0.3.0rc2__py3-none-any.whl → 0.3.1__py3-none-any.whl

digitalkin/core/task_manager/task_manager.py

@@ -1,442 +0,0 @@
-"""Task manager with comprehensive lifecycle management."""
-
-import asyncio
-import contextlib
-import datetime
-from collections.abc import Coroutine
-from typing import Any
-
-from digitalkin.core.task_manager.surrealdb_repository import SurrealDBConnection
-from digitalkin.core.task_manager.task_session import TaskSession
-from digitalkin.logger import logger
-from digitalkin.models.core.task_monitor import SignalMessage, SignalType, TaskStatus
-from digitalkin.modules._base_module import BaseModule
-
-
-class TaskManager:
-    """Task manager with comprehensive lifecycle management.
-
-    Handle the tasks creation, execution, monitoring, signaling, and cancellation.
-    """
-
-    tasks: dict[str, asyncio.Task]
-    tasks_sessions: dict[str, TaskSession]
-    channel: SurrealDBConnection
-    default_timeout: float
-    max_concurrent_tasks: int
-    _shutdown_event: asyncio.Event
-
-    def __init__(self, default_timeout: float = 10.0, max_concurrent_tasks: int = 1000) -> None:
-        """Defining task manager properties."""
-        self.tasks = {}
-        self.tasks_sessions = {}
-        self.default_timeout = default_timeout
-        self.max_concurrent_tasks = max_concurrent_tasks
-        self._shutdown_event = asyncio.Event()
-
-        logger.info(
-            "TaskManager initialized with max_concurrent_tasks: %d, default_timeout: %.1f",
-            max_concurrent_tasks,
-            default_timeout,
-            extra={
-                "max_concurrent_tasks": max_concurrent_tasks,
-                "default_timeout": default_timeout,
-            },
-        )
-
-    @property
-    def task_count(self) -> int:
-        """Number of managed tasks."""
-        return len(self.tasks_sessions)
-
-    @property
-    def running_tasks(self) -> set[str]:
-        """Get IDs of currently running tasks."""
-        return {task_id for task_id, task in self.tasks.items() if not task.done()}
-
-    async def _cleanup_task(self, task_id: str, mission_id: str) -> None:
-        """Clean up task resources.
-
-        Args:
-            task_id (str): The ID of the task to clean up.
-            mission_id (str): The ID of the mission associated with the task.
-        """
-        logger.debug(
-            "Cleaning up resources for task: '%s'", task_id, extra={"mission_id": mission_id, "task_id": task_id}
-        )
-        if task_id in self.tasks_sessions:
-            await self.tasks_sessions[task_id].db.close()
-        # Remove from collections
-        self.tasks.pop(task_id, None)
-        self.tasks_sessions.pop(task_id, None)
-
-    async def _task_wrapper(  # noqa: C901, PLR0915
-        self,
-        task_id: str,
-        mission_id: str,
-        coro: Coroutine[Any, Any, None],
-        session: TaskSession,
-    ) -> asyncio.Task[None]:
-        """Task wrapper that runs main, heartbeat, and listener concurrently.
-
-        The first to finish determines the outcome. Returns a Task that the
-        caller can await externally.
-
-        Returns:
-            asyncio.Task[None]: The supervisor task managing the lifecycle.
-        """
-
-        async def signal_wrapper() -> None:
-            try:
-                await self.channel.create(
-                    "tasks",
-                    SignalMessage(
-                        task_id=task_id,
-                        mission_id=mission_id,
-                        status=session.status,
-                        action=SignalType.START,
-                    ).model_dump(),
-                )
-                await session.listen_signals()
-            except asyncio.CancelledError:
-                logger.debug("Signal listener cancelled", extra={"mission_id": mission_id, "task_id": task_id})
-            finally:
-                await self.channel.create(
-                    "tasks",
-                    SignalMessage(
-                        task_id=task_id,
-                        mission_id=mission_id,
-                        status=session.status,
-                        action=SignalType.STOP,
-                    ).model_dump(),
-                )
-                logger.info("Signal listener ended", extra={"mission_id": mission_id, "task_id": task_id})
-
-        async def heartbeat_wrapper() -> None:
-            try:
-                await session.generate_heartbeats()
-            except asyncio.CancelledError:
-                logger.debug("Signal listener cancelled", extra={"mission_id": mission_id, "task_id": task_id})
-            finally:
-                logger.info("Heartbeat task ended", extra={"mission_id": mission_id, "task_id": task_id})
-
-        async def supervisor() -> None:
-            session.started_at = datetime.datetime.now(datetime.timezone.utc)
-            session.status = TaskStatus.RUNNING
-
-            main_task = asyncio.create_task(coro, name=f"{task_id}_main")
-            hb_task = asyncio.create_task(heartbeat_wrapper(), name=f"{task_id}_heartbeat")
-            sig_task = asyncio.create_task(signal_wrapper(), name=f"{task_id}_listener")
-
-            try:
-                done, pending = await asyncio.wait(
-                    [main_task, sig_task, hb_task],
-                    return_when=asyncio.FIRST_COMPLETED,
-                )
-
-                # One task completed -> cancel the others
-                for t in pending:
-                    t.cancel()
-
-                # Propagate exception/result from the finished task
-                completed = next(iter(done))
-                await completed
-
-                if completed is main_task:
-                    session.status = TaskStatus.COMPLETED
-                elif completed is sig_task or (completed is hb_task and sig_task.done()):
-                    logger.debug(f"Task cancelled due to signal {sig_task=}")
-                    session.status = TaskStatus.CANCELLED
-                elif completed is hb_task:
-                    session.status = TaskStatus.FAILED
-                    msg = f"Heartbeat stopped for {task_id}"
-                    raise RuntimeError(msg)  # noqa: TRY301
-
-            except asyncio.CancelledError:
-                session.status = TaskStatus.CANCELLED
-                raise
-            except Exception:
-                session.status = TaskStatus.FAILED
-                raise
-            finally:
-                session.completed_at = datetime.datetime.now(datetime.timezone.utc)
-                # Ensure all tasks are cleaned up
-                for t in [main_task, hb_task, sig_task]:
-                    if not t.done():
-                        t.cancel()
-                await asyncio.gather(main_task, hb_task, sig_task, return_exceptions=True)
-
-        # Return the supervisor task to be awaited outside
-        return asyncio.create_task(supervisor(), name=f"{task_id}_supervisor")
-
-    async def create_task(
-        self,
-        task_id: str,
-        mission_id: str,
-        module: BaseModule,
-        coro: Coroutine[Any, Any, None],
-        heartbeat_interval: datetime.timedelta = datetime.timedelta(seconds=2),
-        connection_timeout: datetime.timedelta = datetime.timedelta(seconds=5),
-    ) -> None:
-        """Create and start a new managed task.
-
-        Raises:
-            ValueError: task_id duplicated
-            RuntimeError: task overload
-        """
-        if task_id in self.tasks:
-            # close Coroutine during runtime
-            coro.close()
-            logger.warning(
-                "Task creation failed - task already exists: '%s'",
-                task_id,
-                extra={"mission_id": mission_id, "task_id": task_id},
-            )
-            msg = f"Task {task_id} already exists"
-            raise ValueError(msg)
-
-        if len(self.tasks) >= self.max_concurrent_tasks:
-            coro.close()
-            logger.error(
-                "Task creation failed - max concurrent tasks reached: %d",
-                self.max_concurrent_tasks,
-                extra={
-                    "mission_id": mission_id,
-                    "task_id": task_id,
-                    "current_count": len(self.tasks),
-                    "max_concurrent": self.max_concurrent_tasks,
-                },
-            )
-            msg = f"Maximum concurrent tasks ({self.max_concurrent_tasks}) reached"
-            raise RuntimeError(msg)
-
-        logger.info(
-            "Creating new task: '%s'",
-            task_id,
-            extra={
-                "mission_id": mission_id,
-                "task_id": task_id,
-                "heartbeat_interval": heartbeat_interval,
-                "connection_timeout": connection_timeout,
-            },
-        )
-
-        try:
-            # Initialize components
-            channel: SurrealDBConnection = SurrealDBConnection("task_manager", connection_timeout)
-            await channel.init_surreal_instance()
-            session = TaskSession(task_id, mission_id, channel, module, heartbeat_interval)
-
-            self.tasks_sessions[task_id] = session
-
-            # Create wrapper task
-            self.tasks[task_id] = asyncio.create_task(
-                self._task_wrapper(
-                    task_id,
-                    mission_id,
-                    coro,
-                    session,
-                ),
-                name=task_id,
-            )
-
-            logger.info(
-                "Task created successfully: '%s'",
-                task_id,
-                extra={
-                    "mission_id": mission_id,
-                    "task_id": task_id,
-                    "total_tasks": len(self.tasks),
-                },
-            )
-
-        except Exception as e:
-            logger.error(
-                "Failed to create task: '%s'",
-                task_id,
-                extra={"mission_id": mission_id, "task_id": task_id, "error": str(e)},
-                exc_info=True,
-            )
-            # Cleanup on failure
-            await self._cleanup_task(task_id, mission_id=mission_id)
-            raise
-
-    async def send_signal(self, task_id: str, mission_id: str, signal_type: str, payload: dict) -> bool:
-        """Send signal to a specific task.
-
-        Returns:
-            bool: True if the task sent successfully the given signal, False otherwise.
-        """
-        if task_id not in self.tasks_sessions:
-            logger.warning(
-                "Cannot send signal - task not found: '%s'",
-                task_id,
-                extra={"mission_id": mission_id, "task_id": task_id, "signal_type": signal_type},
-            )
-            return False
-
-        logger.info(
-            "Sending signal '%s' to task: '%s'",
-            signal_type,
-            task_id,
-            extra={"mission_id": mission_id, "task_id": task_id, "signal_type": signal_type, "payload": payload},
-        )
-
-        await self.channel.update("tasks", signal_type, payload)
-        return True
-
-    async def cancel_task(self, task_id: str, mission_id: str, timeout: float | None = None) -> bool:
-        """Cancel a task with graceful shutdown and fallback.
-
-        Returns:
-            bool: True if the task was cancelled successfully, False otherwise.
-        """
-        if task_id not in self.tasks:
-            logger.warning(
-                "Cannot cancel - task not found: '%s'", task_id, extra={"mission_id": mission_id, "task_id": task_id}
-            )
-            return True
-
-        timeout = timeout or self.default_timeout
-        task = self.tasks[task_id]
-
-        logger.info(
-            "Initiating task cancellation: '%s', timeout: %.1fs",
-            task_id,
-            timeout,
-            extra={"mission_id": mission_id, "task_id": task_id, "timeout": timeout},
-        )
-
-        try:
-            # Phase 1: Cooperative cancellation
-            # await self.send_signal(task_id, mission_id, "cancel")  # noqa: ERA001
-
-            # Wait for graceful shutdown
-            await asyncio.wait_for(task, timeout=timeout)
-
-            logger.info(
-                "Task cancelled gracefully: '%s'", task_id, extra={"mission_id": mission_id, "task_id": task_id}
-            )
-
-        except asyncio.TimeoutError:
-            logger.warning(
-                "Graceful cancellation timed out for task: '%s', forcing cancellation",
-                task_id,
-                extra={"mission_id": mission_id, "task_id": task_id, "timeout": timeout},
-            )
-
-            # Phase 2: Force cancellation
-            task.cancel()
-            with contextlib.suppress(asyncio.CancelledError):
-                await task
-
-            logger.warning("Task force-cancelled: '%s'", task_id, extra={"mission_id": mission_id, "task_id": task_id})
-            return True
-
-        except Exception as e:
-            logger.error(
-                "Error during task cancellation: '%s'",
-                task_id,
-                extra={"mission_id": mission_id, "task_id": task_id, "error": str(e)},
-                exc_info=True,
-            )
-            return False
-        return True
-
-    async def clean_session(self, task_id: str, mission_id: str) -> bool:
-        """Clean up task session without cancelling the task.
-
-        Returns:
-            bool: True if the task was cleaned successfully, False otherwise.
-        """
-        if task_id not in self.tasks_sessions:
-            logger.warning(
-                "Cannot clean session - task not found: '%s'",
-                task_id,
-                extra={"mission_id": mission_id, "task_id": task_id},
-            )
-            return False
-
-        await self.tasks_sessions[task_id].module.stop()
-        await self.cancel_task(task_id, mission_id)
-
-        logger.info("Cleaning up session for task: '%s'", task_id, extra={"mission_id": mission_id, "task_id": task_id})
-        self.tasks_sessions.pop(task_id, None)
-        self.tasks.pop(task_id, None)
-        return True
-
-    async def pause_task(self, task_id: str, mission_id: str) -> bool:
-        """Pause a running task.
-
-        Returns:
-            bool: True if the task was paused successfully, False otherwise.
-        """
-        return await self.send_signal(task_id, mission_id, "pause", {})
-
-    async def resume_task(self, task_id: str, mission_id: str) -> bool:
-        """Resume a paused task.
-
-        Returns:
-            bool: True if the task was paused successfully, False otherwise.
-        """
-        return await self.send_signal(task_id, mission_id, "resume", {})
-
-    async def get_task_status(self, task_id: str, mission_id: str) -> bool:
-        """Request status from a task.
-
-        Returns:
-            bool: True if the task was paused successfully, False otherwise.
-        """
-        return await self.send_signal(task_id, mission_id, "status", {})
-
-    async def cancel_all_tasks(self, mission_id: str, timeout: float | None = None) -> dict[str, bool]:
-        """Cancel all running tasks.
-
-        Returns:
-            dict[str: bool]: True if the tasks were paused successfully, False otherwise.
-        """
-        timeout = timeout or self.default_timeout
-        task_ids = list(self.running_tasks)
-
-        logger.info(
-            "Cancelling all tasks: %d tasks",
-            len(task_ids),
-            extra={"mission_id": mission_id, "task_count": len(task_ids), "timeout": timeout},
-        )
-
-        results = {}
-        for task_id in task_ids:
-            results[task_id] = await self.cancel_task(task_id, mission_id, timeout)
-
-        return results
-
-    async def shutdown(self, mission_id: str, timeout: float = 30.0) -> None:
-        """Graceful shutdown of all tasks."""
-        logger.info(
-            "TaskManager shutdown initiated, timeout: %.1fs",
-            timeout,
-            extra={"mission_id": mission_id, "timeout": timeout, "active_tasks": len(self.running_tasks)},
-        )
-
-        self._shutdown_event.set()
-        results = await self.cancel_all_tasks(mission_id, timeout)
-
-        failed_tasks = [task_id for task_id, success in results.items() if not success]
-        if failed_tasks:
-            logger.error(
-                "Failed to cancel %d tasks during shutdown: %s",
-                len(failed_tasks),
-                failed_tasks,
-                extra={"mission_id": mission_id, "failed_tasks": failed_tasks, "failed_count": len(failed_tasks)},
-            )
-
-        logger.info(
-            "TaskManager shutdown completed, cancelled: %d, failed: %d",
-            len(results) - len(failed_tasks),
-            len(failed_tasks),
-            extra={
-                "mission_id": mission_id,
-                "cancelled_count": len(results) - len(failed_tasks),
-                "failed_count": len(failed_tasks),
-            },
-        )