digitalkin 0.2.25rc1__py3-none-any.whl → 0.3.0__py3-none-any.whl

digitalkin/modules/job_manager/task_manager.py

@@ -0,0 +1,391 @@
+"""Task manager with comprehensive lifecycle management."""
+
+import asyncio
+import contextlib
+import datetime
+from collections.abc import Coroutine
+from typing import Any
+
+from digitalkin.logger import logger
+from digitalkin.models.module.task_monitor import SignalMessage, SignalType, TaskStatus
+from digitalkin.modules._base_module import BaseModule
+from digitalkin.modules.job_manager.task_session import SurrealDBConnection, TaskSession
+
+
+class TaskManager:
+    """Task manager with comprehensive lifecycle management."""
+
+    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 = 100) -> None:
+        """."""
+        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:
+        """."""
+        return len(self.tasks_sessions)
+
+    @property
+    def running_tasks(self) -> set[str]:
+        """."""
+        return {task_id for task_id, task in self.tasks.items() if not task.done()}
+
+    async def _cleanup_task(self, task_id: str) -> None:
+        """Clean up task resources."""
+        logger.debug("Cleaning up resources for task: '%s'", task_id, extra={"task_id": task_id})
+        if task_id in self.tasks_sessions:
+            await self.tasks_sessions[task_id].db.close()
+        # Remove from collections
+
+    async def _task_wrapper(  # noqa: C901, PLR0915
+        self,
+        task_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,
+                        status=session.status,
+                        action=SignalType.START,
+                    ).model_dump(),
+                )
+                await session.listen_signals()
+            except asyncio.CancelledError:
+                logger.debug("Signal listener cancelled", extra={"task_id": task_id})
+            finally:
+                await self.channel.create(
+                    "tasks",
+                    SignalMessage(
+                        task_id=task_id,
+                        status=session.status,
+                        action=SignalType.STOP,
+                    ).model_dump(),
+                )
+                logger.info("Signal listener ended", extra={"task_id": task_id})
+
+        async def heartbeat_wrapper() -> None:
+            try:
+                await session.generate_heartbeats()
+            except asyncio.CancelledError:
+                logger.debug("Signal listener cancelled", extra={"task_id": task_id})
+            finally:
+                logger.info("Heartbeat task ended", extra={"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
+
+                logger.critical(f"{completed=} | {main_task=} | {hb_task=} | {sig_task=}")
+
+                if completed is main_task:
+                    session.status = TaskStatus.COMPLETED
+                elif completed is sig_task or (completed is hb_task and sig_task.done()):
+                    logger.critical(f"{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,
+        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={"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={
+                    "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={
+                "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, 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, coro, session), name=task_id)
+
+            logger.info(
+                "Task created successfully: '%s'",
+                task_id,
+                extra={
+                    "task_id": task_id,
+                    "total_tasks": len(self.tasks),
+                },
+            )
+
+        except Exception as e:
+            logger.error(
+                "Failed to create task: '%s'", task_id, extra={"task_id": task_id, "error": str(e)}, exc_info=True
+            )
+            # Cleanup on failure
+            await self._cleanup_task(task_id)
+            raise
+
+    async def send_signal(self, task_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={"task_id": task_id, "signal_type": signal_type},
+            )
+            return False
+
+        logger.info(
+            "Sending signal '%s' to task: '%s'",
+            signal_type,
+            task_id,
+            extra={"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, 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={"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={"task_id": task_id, "timeout": timeout},
+        )
+
+        try:
+            # Phase 1: Cooperative cancellation
+            # await self.send_signal(task_id, "cancel")  # noqa: ERA001
+
+            # Wait for graceful shutdown
+            await asyncio.wait_for(task, timeout=timeout)
+
+            logger.info("Task cancelled gracefully: '%s'", task_id, extra={"task_id": task_id})
+
+        except asyncio.TimeoutError:
+            logger.warning(
+                "Graceful cancellation timed out for task: '%s', forcing cancellation",
+                task_id,
+                extra={"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={"task_id": task_id})
+            return True
+
+        except Exception as e:
+            logger.error(
+                "Error during task cancellation: '%s'",
+                task_id,
+                extra={"task_id": task_id, "error": str(e)},
+                exc_info=True,
+            )
+            return False
+        return True
+
+    async def clean_session(self, task_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={"task_id": task_id})
+            return False
+
+        await self.tasks_sessions[task_id].module.stop()
+        await self.cancel_task(task_id)
+
+        logger.info("Cleaning up session for task: '%s'", task_id, extra={"task_id": task_id})
+        self.tasks_sessions.pop(task_id, None)
+        return True
+
+    async def pause_task(self, task_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, "pause", {})
+
+    async def resume_task(self, task_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, "resume", {})
+
+    async def get_task_status(self, task_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, "status", {})
+
+    async def cancel_all_tasks(self, 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={"task_count": len(task_ids), "timeout": timeout}
+        )
+
+        results = {}
+        for task_id in task_ids:
+            results[task_id] = await self.cancel_task(task_id, timeout)
+
+        return results
+
+    async def shutdown(self, timeout: float = 30.0) -> None:
+        """Graceful shutdown of all tasks."""
+        logger.info(
+            "TaskManager shutdown initiated, timeout: %.1fs",
+            timeout,
+            extra={"timeout": timeout, "active_tasks": len(self.running_tasks)},
+        )
+
+        self._shutdown_event.set()
+        results = await self.cancel_all_tasks(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={"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={"cancelled_count": len(results) - len(failed_tasks), "failed_count": len(failed_tasks)},
+        )

digitalkin/modules/job_manager/task_session.py

@@ -0,0 +1,276 @@
+"""."""
+
+import asyncio
+import datetime
+from collections.abc import AsyncGenerator
+
+from digitalkin.logger import logger
+from digitalkin.models.module.task_monitor import HeartbeatMessage, SignalMessage, SignalType, TaskStatus
+from digitalkin.modules._base_module import BaseModule
+from digitalkin.modules.job_manager.surrealdb_repository import SurrealDBConnection
+
+
+class TaskSession:
+    """Task Session with lifecycle management.
+
+    The Session defined the whole lifecycle of a task as an epheneral context.
+    """
+
+    db: SurrealDBConnection
+    module: BaseModule
+
+    status: TaskStatus
+    signal_queue: AsyncGenerator | None
+
+    task_id: str
+    signal_record_id: str | None
+    heartbeat_record_id: str | None
+
+    started_at: datetime.datetime | None
+    completed_at: datetime.datetime | None
+
+    is_cancelled: asyncio.Event
+    _paused: asyncio.Event
+    _heartbeat_interval: datetime.timedelta
+    _last_heartbeat: datetime.datetime
+
+    def __init__(
+        self,
+        task_id: str,
+        db: SurrealDBConnection,
+        module: BaseModule,
+        heartbeat_interval: datetime.timedelta = datetime.timedelta(seconds=2),
+    ) -> None:
+        """."""
+        self.db = db
+        self.module = module
+
+        self.status = TaskStatus.PENDING
+        self.queue: asyncio.Queue = asyncio.Queue()
+
+        self.task_id = task_id
+        self.heartbeat = None
+        self.started_at = None
+        self.completed_at = None
+
+        self.signal_record_id = None
+        self.heartbeat_record_id = None
+
+        self.is_cancelled = asyncio.Event()
+        self._paused = asyncio.Event()
+        self._heartbeat_interval = heartbeat_interval
+
+        logger.info(
+            "TaskContext initialized for task: '%s'",
+            task_id,
+            extra={"task_id": task_id, "heartbeat_interval": heartbeat_interval},
+        )
+
+    @property
+    def cancelled(self) -> bool:
+        """."""
+        return self.is_cancelled.is_set()
+
+    @property
+    def paused(self) -> bool:
+        """."""
+        return self._paused.is_set()
+
+    async def send_heartbeat(self) -> bool:
+        """Rate-limited heartbeat with connection resilience.
+
+        Returns:
+            bool: True if heartbeat was successful, False otherwise
+        """
+        heartbeat = HeartbeatMessage(
+            task_id=self.task_id,
+            timestamp=datetime.datetime.now(datetime.timezone.utc),
+        )
+
+        if self.heartbeat_record_id is None:
+            try:
+                success = await self.db.create("heartbeats", heartbeat.model_dump())
+                logger.critical(f"{success=} | {'code' not in success}")
+                if "code" not in success:
+                    self.heartbeat_record_id = success.get("id")  # type: ignore
+                    self._last_heartbeat = heartbeat.timestamp
+                    return True
+            except Exception as e:
+                logger.error(
+                    "Heartbeat exception for task: '%s'",
+                    self.task_id,
+                    extra={"task_id": self.task_id, "error": str(e)},
+                    exc_info=True,
+                )
+            logger.error(
+                "Initial heartbeat failed for task: '%s'",
+                self.task_id,
+                extra={"task_id": self.task_id},
+            )
+            return False
+
+        if (heartbeat.timestamp - self._last_heartbeat) < self._heartbeat_interval:
+            logger.debug(
+                "Heartbeat skipped due to rate limiting for task: '%s' | delta=%s",
+                self.task_id,
+                heartbeat.timestamp - self._last_heartbeat,
+            )
+            return True
+
+        try:
+            success = await self.db.merge("heartbeats", self.heartbeat_record_id, heartbeat.model_dump())
+            if "code" not in success:
+                self._last_heartbeat = heartbeat.timestamp
+                return True
+        except Exception as e:
+            logger.error(
+                "Heartbeat exception for task: '%s'",
+                self.task_id,
+                extra={"task_id": self.task_id, "error": str(e)},
+                exc_info=True,
+            )
+        logger.warning(
+            "Heartbeat failed for task: '%s'",
+            self.task_id,
+            extra={"task_id": self.task_id},
+        )
+        return False
+
+    async def generate_heartbeats(self) -> None:
+        """Periodic heartbeat generator with cancellation support."""
+        logger.critical("Heartbeat started")
+        while not self.cancelled:
+            logger.debug(f"Heartbeat tick for task: '{self.task_id}' | {self.cancelled=}")
+            success = await self.send_heartbeat()
+            if not success:
+                logger.error("Heartbeat failed, cancelling task: '%s'", self.task_id, extra={"task_id": self.task_id})
+                await self._handle_cancel()
+                break
+            await asyncio.sleep(self._heartbeat_interval.total_seconds())
+
+    async def wait_if_paused(self) -> None:
+        """Block execution if task is paused."""
+        if self._paused.is_set():
+            logger.info("Task paused, waiting for resume: '%s'", self.task_id, extra={"task_id": self.task_id})
+            await self._paused.wait()
+
+    async def listen_signals(self) -> None:  # noqa: C901
+        """Enhanced signal listener with comprehensive handling.
+
+        Raises:
+            CancelledError: Asyncio when task cancelling
+        """
+        logger.info("Signal listener started for task: '%s'", self.task_id, extra={"task_id": self.task_id})
+        if self.signal_record_id is None:
+            self.signal_record_id = (await self.db.select_by_task_id("tasks", self.task_id)).get("id")
+
+        live_id, live_signals = await self.db.start_live("tasks")
+        try:
+            async for signal in live_signals:
+                logger.critical("Signal received for task '%s': %s", self.task_id, signal)
+                if self.cancelled:
+                    break
+
+                if signal is None or signal["id"] == self.signal_record_id or "payload" not in signal:
+                    continue
+
+                if signal["action"] == "cancel":
+                    await self._handle_cancel()
+                elif signal["action"] == "pause":
+                    await self._handle_pause()
+                elif signal["action"] == "resume":
+                    await self._handle_resume()
+                elif signal["action"] == "status":
+                    await self._handle_status_request()
+
+        except asyncio.CancelledError:
+            logger.debug("Signal listener cancelled for task: '%s'", self.task_id, extra={"task_id": self.task_id})
+            raise
+        except Exception as e:
+            logger.error(
+                "Signal listener fatal error for task: '%s'",
+                self.task_id,
+                extra={"task_id": self.task_id, "error": str(e)},
+                exc_info=True,
+            )
+        finally:
+            await self.db.stop_live(live_id)
+            logger.info("Signal listener stopped for task: '%s'", self.task_id, extra={"task_id": self.task_id})
+
+    async def _handle_cancel(self) -> None:
+        """Idempotent cancellation with acknowledgment."""
+        logger.critical("Handle cancel called")
+        if self.is_cancelled.is_set():
+            logger.debug(
+                "Cancel signal ignored - task already cancelled: '%s'", self.task_id, extra={"task_id": self.task_id}
+            )
+            return
+
+        logger.info("Cancelling task: '%s'", self.task_id, extra={"task_id": self.task_id})
+
+        self.status = TaskStatus.CANCELLED
+        self.is_cancelled.set()
+
+        # Resume if paused so cancellation can proceed
+        if self._paused.is_set():
+            self._paused.set()
+
+        await self.db.update(
+            "tasks",
+            self.signal_record_id,  # type: ignore
+            SignalMessage(
+                task_id=self.task_id,
+                action=SignalType.ACK_CANCEL,
+                status=self.status,
+            ).model_dump(),
+        )
+
+    async def _handle_pause(self) -> None:
+        """Pause task execution."""
+        if not self._paused.is_set():
+            logger.info("Pausing task: '%s'", self.task_id, extra={"task_id": self.task_id})
+            self._paused.set()
+
+        await self.db.update(
+            "tasks",
+            self.signal_record_id,  # type: ignore
+            SignalMessage(
+                task_id=self.task_id,
+                action=SignalType.ACK_PAUSE,
+                status=self.status,
+            ).model_dump(),
+        )
+
+    async def _handle_resume(self) -> None:
+        """Resume paused task."""
+        if self._paused.is_set():
+            logger.info("Resuming task: '%s'", self.task_id, extra={"task_id": self.task_id})
+            self._paused.clear()
+
+        await self.db.update(
+            "tasks",
+            self.signal_record_id,  # type: ignore
+            SignalMessage(
+                task_id=self.task_id,
+                action=SignalType.ACK_RESUME,
+                status=self.status,
+            ).model_dump(),
+        )
+
+    async def _handle_status_request(self) -> None:
+        """Send current task status."""
+        await self.db.update(
+            "tasks",
+            self.signal_record_id,  # type: ignore
+            SignalMessage(
+                action=SignalType.ACK_STATUS,
+                task_id=self.task_id,
+                status=self.status,
+            ).model_dump(),
+        )
+
+        logger.debug(
+            "Status report sent for task: '%s'",
+            self.task_id,
+            extra={"task_id": self.task_id},
+        )

digitalkin/modules/job_manager/taskiq_job_manager.py

@@ -19,7 +19,7 @@ from rstream import Consumer, ConsumerOffsetSpecification, MessageContext, Offse
 
 from digitalkin.logger import logger
 from digitalkin.models.module import InputModelT, SetupModelT
-from digitalkin.models.module.module import ModuleStatus
+from digitalkin.models.module.task_monitor import TaskStatus
 from digitalkin.modules._base_module import BaseModule
 from digitalkin.modules.job_manager.base_job_manager import BaseJobManager
 from digitalkin.modules.job_manager.taskiq_broker import STREAM, STREAM_RETENTION, TASKIQ_BROKER
@@ -279,7 +279,7 @@ class TaskiqJobManager(BaseJobManager, Generic[InputModelT, SetupModelT]):
         msg = "stop_all_modules not implemented in TaskiqJobManager"
         raise NotImplementedError(msg)
 
-    async def get_module_status(self, job_id: str) -> ModuleStatus | None:
+    async def get_module_status(self, job_id: str) -> TaskStatus:
         """Query a module status."""
         msg = "get_module_status not implemented in TaskiqJobManager"
         raise NotImplementedError(msg)