digitalkin 0.3.2.dev2__py3-none-any.whl

digitalkin/core/task_manager/remote_task_manager.py

@@ -0,0 +1,87 @@
+"""Remote task manager for distributed execution."""
+
+import datetime
+from collections.abc import Coroutine
+from typing import Any
+
+from digitalkin.core.task_manager.base_task_manager import BaseTaskManager
+from digitalkin.logger import logger
+from digitalkin.modules._base_module import BaseModule
+
+
+class RemoteTaskManager(BaseTaskManager):
+    """Task manager for distributed/remote execution.
+
+    Only manages task metadata and signals - actual execution happens in remote workers.
+    Suitable for horizontally scaled deployments with Taskiq/Celery workers.
+    """
+
+    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:
+        """Register task for remote execution (metadata only).
+
+        Creates TaskSession for signal handling and monitoring, but doesn't execute the coroutine.
+        The coroutine will be recreated and executed by a remote worker.
+
+        Args:
+            task_id: Unique identifier for the task
+            mission_id: Mission identifier
+            module: Module instance for metadata (not executed here)
+            coro: Coroutine (will be closed - execution happens in worker)
+            heartbeat_interval: Interval between heartbeats
+            connection_timeout: Connection timeout for SurrealDB
+
+        Raises:
+            ValueError: If task_id duplicated
+            RuntimeError: If task overload
+        """
+        # Validation
+        await self._validate_task_creation(task_id, mission_id, coro)
+
+        logger.info(
+            "Registering remote task: '%s'",
+            task_id,
+            extra={
+                "mission_id": mission_id,
+                "task_id": task_id,
+                "heartbeat_interval": heartbeat_interval,
+                "connection_timeout": connection_timeout,
+            },
+        )
+
+        try:
+            # Create session for metadata and signal handling
+            _channel, _session = await self._create_session(
+                task_id, mission_id, module, heartbeat_interval, connection_timeout
+            )
+
+            # Close coroutine - worker will recreate and execute it
+            coro.close()
+
+            logger.info(
+                "Remote task registered: '%s'",
+                task_id,
+                extra={
+                    "mission_id": mission_id,
+                    "task_id": task_id,
+                    "total_sessions": len(self.tasks_sessions),
+                },
+            )
+
+        except Exception as e:
+            logger.error(
+                "Failed to register remote 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

digitalkin/core/task_manager/surrealdb_repository.py

@@ -0,0 +1,266 @@
+"""SurrealDB connection management."""
+
+import asyncio
+import datetime
+import os
+from collections.abc import AsyncGenerator
+from typing import Any, Generic, TypeVar, cast
+from uuid import UUID
+
+from surrealdb import AsyncHttpSurrealConnection, AsyncSurreal, AsyncWsSurrealConnection, RecordID
+
+from digitalkin.logger import logger
+
+TSurreal = TypeVar("TSurreal", bound=AsyncHttpSurrealConnection | AsyncWsSurrealConnection)
+
+
+class SurrealDBSetupBadIDError(Exception):
+    """Exception raised when an invalid ID is encountered during the setup process in the SurrealDB repository.
+
+    This error is used to indicate that the provided ID does not meet the
+    expected format or criteria.
+    """
+
+
+class SurrealDBSetupVersionBadIDError(Exception):
+    """Exception raised when an invalid ID is encountered during the setup of a SurrealDB version.
+
+    This error is intended to signal that the provided ID does not meet
+    the expected format or criteria for a valid SurrealDB setup version ID.
+    """
+
+
+class SurrealDBConnection(Generic[TSurreal]):
+    """Base repository for database operations.
+
+    This class provides common database operations that can be used by
+    specific table repositories.
+    """
+
+    db: TSurreal
+    timeout: datetime.timedelta
+    _live_queries: set[UUID]  # Track active live queries for cleanup
+
+    @staticmethod
+    def _valid_id(raw_id: str, table_name: str) -> RecordID:
+        """Validate and parse a raw ID string into a RecordID.
+
+        Args:
+            raw_id: The raw ID string to validate
+            table_name: table name to enforce
+
+        Raises:
+            SurrealDBSetupBadIDError: If the raw ID string is not valid
+
+        Returns:
+            RecordID: Parsed RecordID object if valid, None otherwise
+        """
+        try:
+            split_id = raw_id.split(":")
+            if split_id[0] != table_name:
+                msg = f"Invalid table name for ID: {raw_id}"
+                raise SurrealDBSetupBadIDError(msg)
+            return RecordID(split_id[0], split_id[1])
+        except IndexError:
+            raise SurrealDBSetupBadIDError
+
+    def __init__(
+        self,
+        database: str | None = None,
+        timeout: datetime.timedelta = datetime.timedelta(seconds=5),
+    ) -> None:
+        """Initialize the repository.
+
+        Args:
+            database: AsyncSurrealDB connection to a specific database
+            timeout: Timeout for database operations
+        """
+        self.timeout = timeout
+        base_url = os.getenv("SURREALDB_URL", "ws://localhost").strip()
+        port = (os.getenv("SURREALDB_PORT") or "").strip()
+        self.url = f"{base_url}{f':{port}' if port else ''}/rpc"
+
+        self.username = os.getenv("SURREALDB_USERNAME", "root")
+        self.password = os.getenv("SURREALDB_PASSWORD", "root")
+        self.namespace = os.getenv("SURREALDB_NAMESPACE", "test")
+        self.database = database or os.getenv("SURREALDB_DATABASE", "task_manager")
+        self._live_queries = set()  # Initialize live queries tracker
+
+    async def init_surreal_instance(self) -> None:
+        """Init a SurrealDB connection instance."""
+        logger.debug("Connecting to SurrealDB at %s", self.url)
+        self.db = AsyncSurreal(self.url)  # type: ignore
+        await self.db.signin({"username": self.username, "password": self.password})
+        await self.db.use(self.namespace, self.database)  # type: ignore[arg-type]
+        logger.debug("Successfully connected to SurrealDB")
+
+    async def close(self) -> None:
+        """Close the SurrealDB connection if it exists.
+
+        This will also kill all active live queries to prevent memory leaks.
+        """
+        # Kill all tracked live queries before closing connection
+        if self._live_queries:
+            logger.debug("Killing %d active live queries before closing", len(self._live_queries))
+            live_query_ids = list(self._live_queries)
+
+            # Kill all queries concurrently, capturing any exceptions
+            results = await asyncio.gather(
+                *[self.db.kill(live_id) for live_id in live_query_ids], return_exceptions=True
+            )
+
+            # Process results and track failures
+            failed_queries = []
+            for live_id, result in zip(live_query_ids, results):
+                if isinstance(result, ConnectionError | TimeoutError | Exception):
+                    failed_queries.append((live_id, str(result)))
+                else:
+                    self._live_queries.discard(live_id)
+
+            # Log aggregated failures once instead of per-query
+            if failed_queries:
+                logger.warning(
+                    "Failed to kill %d live queries: %s",
+                    len(failed_queries),
+                    failed_queries[:5],  # Only log first 5 to avoid log spam
+                    extra={"total_failed": len(failed_queries)},
+                )
+
+        logger.debug("Closing SurrealDB connection")
+        await self.db.close()
+
+    async def create(
+        self,
+        table_name: str,
+        data: dict[str, Any],
+    ) -> list[dict[str, Any]] | dict[str, Any]:
+        """Create a new record.
+
+        Args:
+            table_name: Name of the table to insert into
+            data: Data to insert
+
+        Returns:
+            Dict[str, Any]: The created record as returned by the database
+        """
+        logger.debug("Creating record in %s with data: %s", table_name, data)
+        result = await self.db.create(table_name, data)
+        logger.debug("create result: %s", result)
+        return cast("list[dict[str, Any]] | dict[str, Any]", result)
+
+    async def merge(
+        self,
+        table_name: str,
+        record_id: str | RecordID,
+        data: dict[str, Any],
+    ) -> list[dict[str, Any]] | dict[str, Any]:
+        """Update an existing record.
+
+        Args:
+            table_name: Name of the table to insert into
+            record_id: record ID to update
+            data: Data to insert
+
+        Returns:
+            Dict[str, Any]: The created record as returned by the database
+        """
+        if isinstance(record_id, str):
+            # validate surrealDB id if raw str
+            record_id = self._valid_id(record_id, table_name)
+        logger.debug("Updating record in %s with data: %s", record_id, data)
+        result = await self.db.merge(record_id, data)
+        logger.debug("update result: %s", result)
+        return cast("list[dict[str, Any]] | dict[str, Any]", result)
+
+    async def update(
+        self,
+        table_name: str,
+        record_id: str | RecordID,
+        data: dict[str, Any],
+    ) -> list[dict[str, Any]] | dict[str, Any]:
+        """Update an existing record.
+
+        Args:
+            table_name: Name of the table to insert into
+            record_id: record ID to update
+            data: Data to insert
+
+        Returns:
+            Dict[str, Any]: The created record as returned by the database
+        """
+        if isinstance(record_id, str):
+            # validate surrealDB id if raw str
+            record_id = self._valid_id(record_id, table_name)
+        logger.debug("Updating record in %s with data: %s", record_id, data)
+        result = await self.db.update(record_id, data)
+        logger.debug("update result: %s", result)
+        return cast("list[dict[str, Any]] | dict[str, Any]", result)
+
+    async def execute_query(self, query: str, params: dict[str, Any] | None = None) -> list[dict[str, Any]]:
+        """Execute a custom SurrealQL query.
+
+        Args:
+            query: SurrealQL query
+            params: Query parameters
+
+        Returns:
+            List[Dict[str, Any]]: Query results
+        """
+        logger.debug("execute_query: %s with params: %s", query, params)
+        result = await self.db.query(query, params or {})
+        logger.debug("execute_query result: %s", result)
+        return cast("list[dict[str, Any]]", [result] if isinstance(result, dict) else result)
+
+    async def select_by_task_id(self, table: str, value: str) -> dict[str, Any]:
+        """Fetch a record from a table by a unique field.
+
+        Args:
+            table: Table name
+            value: Field value to match
+
+        Raises:
+            ValueError: If no records are found
+
+        Returns:
+            Dict with record data if found, else None
+        """
+        query = "SELECT * FROM type::table($table) WHERE task_id = $value;"
+        params = {"table": table, "value": value}
+
+        result = await self.execute_query(query, params)
+        if not result:
+            msg = f"No records found in table '{table}' with task_id '{value}'"
+            logger.error(msg)
+            raise ValueError(msg)
+
+        return result[0]
+
+    async def start_live(
+        self,
+        table_name: str,
+    ) -> tuple[UUID, AsyncGenerator[dict[str, Any], None]]:
+        """Create and subscribe to a live SurrealQL query.
+
+        The live query ID is tracked to ensure proper cleanup on connection close.
+
+        Args:
+            table_name: Name of the table to insert into
+
+        Returns:
+            tuple[UUID, AsyncGenerator]: Live query ID and subscription generator
+        """
+        live_id = await self.db.live(table_name, diff=False)
+        self._live_queries.add(live_id)  # Track for cleanup
+        logger.debug("Started live query %s for table %s (total: %d)", live_id, table_name, len(self._live_queries))
+        return live_id, await self.db.subscribe_live(live_id)
+
+    async def stop_live(self, live_id: UUID) -> None:
+        """Kill a live SurrealQL query.
+
+        Args:
+            live_id: Live query ID to kill
+        """
+        logger.debug("Killing live query: %s", live_id)
+        await self.db.kill(live_id)
+        self._live_queries.discard(live_id)  # Remove from tracker
+        logger.debug("Stopped live query %s (remaining: %d)", live_id, len(self._live_queries))

digitalkin/core/task_manager/task_executor.py

@@ -0,0 +1,249 @@
+"""Task executor for running tasks with full lifecycle management."""
+
+import asyncio
+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 (
+    CancellationReason,
+    SignalMessage,
+    SignalType,
+    TaskStatus,
+)
+
+
+class TaskExecutor:
+    """Executes tasks with the supervisor pattern (main + heartbeat + signal listener).
+
+    Pure execution logic - no task registry or orchestration.
+    Used by workers to run distributed tasks or by TaskManager for local execution.
+    """
+
+    @staticmethod
+    async def execute_task(  # noqa: C901, PLR0915
+        task_id: str,
+        mission_id: str,
+        coro: Coroutine[Any, Any, None],
+        session: TaskSession,
+        channel: SurrealDBConnection,
+    ) -> asyncio.Task[None]:
+        """Execute a task using the supervisor pattern.
+
+        Runs three concurrent sub-tasks:
+        - Main coroutine (the actual work)
+        - Heartbeat generator (sends heartbeats to SurrealDB)
+        - Signal listener (watches for stop/pause/resume signals)
+
+        The first task to complete determines the outcome.
+
+        Args:
+            task_id: Unique identifier for the task
+            mission_id: Mission identifier for the task
+            coro: The coroutine to execute (module.start(...))
+            session: TaskSession for state management
+            channel: SurrealDB connection for signals
+
+        Returns:
+            asyncio.Task: The supervisor task managing the lifecycle
+        """
+
+        async def signal_wrapper() -> None:
+            """Create initial signal record and listen for signals."""
+            try:
+                await 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 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:
+            """Generate heartbeats for task health monitoring."""
+            try:
+                await session.generate_heartbeats()
+            except asyncio.CancelledError:
+                logger.debug("Heartbeat 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:  # noqa: C901, PLR0912, PLR0915
+            """Supervise the three concurrent tasks and handle outcomes.
+
+            Raises:
+                RuntimeError: If the heartbeat task stops unexpectedly.
+                asyncio.CancelledError: If the supervisor task is cancelled.
+            """
+            session.started_at = datetime.datetime.now(datetime.timezone.utc)
+            session.status = TaskStatus.RUNNING
+
+            # Create tasks with proper exception handling
+            main_task = None
+            hb_task = None
+            sig_task = None
+            cleanup_reason = CancellationReason.UNKNOWN
+
+            try:
+                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")
+                done, pending = await asyncio.wait(
+                    [main_task, sig_task, hb_task],
+                    return_when=asyncio.FIRST_COMPLETED,
+                )
+
+                # Determine cleanup reason based on which task completed first
+                completed = next(iter(done))
+
+                if completed is main_task:
+                    # Main task finished - cleanup is due to success
+                    cleanup_reason = CancellationReason.SUCCESS_CLEANUP
+                elif completed is sig_task or (completed is hb_task and sig_task.done()):
+                    # Signal task finished - external cancellation
+                    cleanup_reason = CancellationReason.SIGNAL
+                elif completed is hb_task:
+                    # Heartbeat stopped - failure cleanup
+                    cleanup_reason = CancellationReason.FAILURE_CLEANUP
+
+                # Cancel pending tasks with proper reason logging
+                if pending:
+                    pending_names = [t.get_name() for t in pending]
+                    logger.debug(
+                        "Cancelling pending tasks: %s, reason: %s",
+                        pending_names,
+                        cleanup_reason.value,
+                        extra={
+                            "mission_id": mission_id,
+                            "task_id": task_id,
+                            "pending_tasks": pending_names,
+                            "cancellation_reason": cleanup_reason.value,
+                        },
+                    )
+                    for t in pending:
+                        t.cancel()
+
+                # Propagate exception/result from the finished task
+                await completed
+
+                # Determine final status based on which task completed
+                if completed is main_task:
+                    session.status = TaskStatus.COMPLETED
+                    logger.info(
+                        "Main task completed successfully",
+                        extra={"mission_id": mission_id, "task_id": task_id},
+                    )
+                elif completed is sig_task or (completed is hb_task and sig_task.done()):
+                    session.status = TaskStatus.CANCELLED
+                    session.cancellation_reason = CancellationReason.SIGNAL
+                    logger.info(
+                        "Task cancelled via external signal",
+                        extra={
+                            "mission_id": mission_id,
+                            "task_id": task_id,
+                            "cancellation_reason": CancellationReason.SIGNAL.value,
+                        },
+                    )
+                elif completed is hb_task:
+                    session.status = TaskStatus.FAILED
+                    session.cancellation_reason = CancellationReason.HEARTBEAT_FAILURE
+                    logger.error(
+                        "Heartbeat stopped unexpectedly for task: '%s'",
+                        task_id,
+                        extra={
+                            "mission_id": mission_id,
+                            "task_id": task_id,
+                            "cancellation_reason": CancellationReason.HEARTBEAT_FAILURE.value,
+                        },
+                    )
+                    msg = f"Heartbeat stopped for {task_id}"
+                    raise RuntimeError(msg)  # noqa: TRY301
+
+            except asyncio.CancelledError:
+                session.status = TaskStatus.CANCELLED
+                # Only set reason if not already set (preserve original reason)
+                logger.info(
+                    "Task cancelled externally: '%s', reason: %s",
+                    task_id,
+                    session.cancellation_reason.value,
+                    extra={
+                        "mission_id": mission_id,
+                        "task_id": task_id,
+                        "cancellation_reason": session.cancellation_reason.value,
+                    },
+                )
+                cleanup_reason = CancellationReason.FAILURE_CLEANUP
+                raise
+            except Exception:
+                session.status = TaskStatus.FAILED
+                cleanup_reason = CancellationReason.FAILURE_CLEANUP
+                logger.exception(
+                    "Task failed with exception: '%s'",
+                    task_id,
+                    extra={"mission_id": mission_id, "task_id": task_id},
+                )
+                raise
+            finally:
+                session.completed_at = datetime.datetime.now(datetime.timezone.utc)
+                # Ensure all tasks are cleaned up with proper reason
+                tasks_to_cleanup = [t for t in [main_task, hb_task, sig_task] if t is not None and not t.done()]
+                if tasks_to_cleanup:
+                    cleanup_names = [t.get_name() for t in tasks_to_cleanup]
+                    logger.debug(
+                        "Final cleanup of %d remaining tasks: %s, reason: %s",
+                        len(tasks_to_cleanup),
+                        cleanup_names,
+                        cleanup_reason.value,
+                        extra={
+                            "mission_id": mission_id,
+                            "task_id": task_id,
+                            "cleanup_count": len(tasks_to_cleanup),
+                            "cleanup_tasks": cleanup_names,
+                            "cancellation_reason": cleanup_reason.value,
+                        },
+                    )
+                    for t in tasks_to_cleanup:
+                        t.cancel()
+                    await asyncio.gather(*tasks_to_cleanup, return_exceptions=True)
+
+                duration = (
+                    (session.completed_at - session.started_at).total_seconds()
+                    if session.started_at and session.completed_at
+                    else None
+                )
+                logger.info(
+                    "Task execution completed: '%s', status: %s, reason: %s, duration: %.2fs",
+                    task_id,
+                    session.status.value,
+                    session.cancellation_reason.value if session.status == TaskStatus.CANCELLED else "n/a",
+                    duration or 0,
+                    extra={
+                        "mission_id": mission_id,
+                        "task_id": task_id,
+                        "status": session.status.value,
+                        "cancellation_reason": session.cancellation_reason.value,
+                        "duration": duration,
+                    },
+                )
+
+        # Return the supervisor task to be awaited by caller
+        return asyncio.create_task(supervisor(), name=f"{task_id}_supervisor")