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

digitalkin/__version__.py

@@ -5,4 +5,4 @@ from importlib.metadata import PackageNotFoundError, version
 try:
     __version__ = version("digitalkin")
 except PackageNotFoundError:
-    __version__ = "0.3.0-rc1"
+    __version__ = "0.3.1"

digitalkin/core/common/__init__.py

@@ -0,0 +1,9 @@
+"""Common utilities for the core module."""
+
+from digitalkin.core.common.factories import ConnectionFactory, ModuleFactory, QueueFactory
+
+__all__ = [
+    "ConnectionFactory",
+    "ModuleFactory",
+    "QueueFactory",
+]

digitalkin/core/common/factories.py

@@ -0,0 +1,156 @@
+"""Common factory functions for reducing code duplication in core module."""
+
+import asyncio
+import datetime
+
+from digitalkin.core.task_manager.surrealdb_repository import SurrealDBConnection
+from digitalkin.logger import logger
+from digitalkin.modules._base_module import BaseModule
+
+
+class ConnectionFactory:
+    """Factory for creating SurrealDB connections with consistent configuration."""
+
+    @staticmethod
+    async def create_surreal_connection(
+        database: str = "task_manager",
+        timeout: datetime.timedelta = datetime.timedelta(seconds=5),
+        *,
+        auto_init: bool = True,
+    ) -> SurrealDBConnection:
+        """Create and optionally initialize a SurrealDB connection.
+
+        This factory method centralizes the creation of SurrealDB connections
+        to ensure consistent configuration across the codebase.
+
+        Args:
+            database: Database name to connect to
+            timeout: Connection timeout
+            auto_init: Whether to automatically initialize the connection
+
+        Returns:
+            Initialized or uninitialized SurrealDBConnection instance
+
+        Example:
+            # Create and auto-initialize
+            conn = await ConnectionFactory.create_surreal_connection("taskiq_worker")
+
+            # Create without initialization
+            conn = await ConnectionFactory.create_surreal_connection(auto_init=False)
+            await conn.init_surreal_instance()
+        """
+        logger.debug(
+            "Creating SurrealDB connection for database: %s, timeout: %s",
+            database,
+            timeout,
+            extra={"database": database, "timeout": str(timeout)},
+        )
+
+        connection: SurrealDBConnection = SurrealDBConnection(database, timeout)
+
+        if auto_init:
+            await connection.init_surreal_instance()
+            logger.debug("SurrealDB connection initialized for database: %s", database)
+
+        return connection
+
+
+class ModuleFactory:
+    """Factory for creating module instances with consistent configuration."""
+
+    @staticmethod
+    def create_module_instance(
+        module_class: type[BaseModule],
+        job_id: str,
+        mission_id: str,
+        setup_id: str,
+        setup_version_id: str,
+    ) -> BaseModule:
+        """Create a module instance with standard parameters.
+
+        This factory method centralizes module instantiation to ensure
+        consistent parameter passing across the codebase.
+
+        Args:
+            module_class: The module class to instantiate
+            job_id: Unique job identifier
+            mission_id: Mission identifier
+            setup_id: Setup identifier
+            setup_version_id: Setup version identifier
+
+        Returns:
+            Instantiated module
+
+        Raises:
+            ValueError: If job_id or mission_id is empty
+
+        Example:
+            module = ModuleFactory.create_module_instance(
+                MyModule,
+                job_id="job_123",
+                mission_id="mission:test",
+                setup_id="setup:config",
+                setup_version_id="v1.0",
+            )
+        """
+        # Validate parameters
+        if not job_id:
+            msg = "job_id cannot be empty"
+            raise ValueError(msg)
+        if not mission_id:
+            msg = "mission_id cannot be empty"
+            raise ValueError(msg)
+
+        logger.debug(
+            "Creating module instance: %s for job: %s",
+            module_class.__name__,
+            job_id,
+            extra={
+                "module_class": module_class.__name__,
+                "job_id": job_id,
+                "mission_id": mission_id,
+                "setup_id": setup_id,
+                "setup_version_id": setup_version_id,
+            },
+        )
+
+        return module_class(
+            job_id=job_id,
+            mission_id=mission_id,
+            setup_id=setup_id,
+            setup_version_id=setup_version_id,
+        )
+
+
+class QueueFactory:
+    """Factory for creating asyncio queues with consistent configuration."""
+
+    # Default max queue size to prevent unbounded memory growth
+    DEFAULT_MAX_QUEUE_SIZE = 1000
+
+    @staticmethod
+    def create_bounded_queue(maxsize: int = DEFAULT_MAX_QUEUE_SIZE) -> asyncio.Queue:
+        """Create a bounded asyncio queue with standard configuration.
+
+        Args:
+            maxsize: Maximum queue size (default 1000, 0 means unlimited)
+
+        Returns:
+            Bounded asyncio.Queue instance
+
+        Raises:
+            ValueError: If maxsize is negative
+
+        Example:
+            queue = QueueFactory.create_bounded_queue()
+            # or with custom size
+            queue = QueueFactory.create_bounded_queue(maxsize=500)
+            # unlimited queue
+            queue = QueueFactory.create_bounded_queue(maxsize=0)
+        """
+        if maxsize < 0:
+            msg = "maxsize must be >= 0"
+            raise ValueError(msg)
+
+        logger.debug("Creating bounded queue with maxsize: %d", maxsize, extra={"maxsize": maxsize})
+        return asyncio.Queue(maxsize=maxsize)

digitalkin/core/job_manager/base_job_manager.py

@@ -5,7 +5,8 @@ from collections.abc import AsyncGenerator, AsyncIterator, Callable, Coroutine
 from contextlib import asynccontextmanager
 from typing import Any, Generic
 
-from digitalkin.core.task_manager.task_manager import TaskManager
+from digitalkin.core.task_manager.base_task_manager import BaseTaskManager
+from digitalkin.core.task_manager.task_session import TaskSession
 from digitalkin.models.core.task_monitor import TaskStatus
 from digitalkin.models.module import InputModelT, OutputModelT, SetupModelT
 from digitalkin.models.module.module import ModuleCodeModel
@@ -14,9 +15,115 @@ from digitalkin.services.services_config import ServicesConfig
 from digitalkin.services.services_models import ServicesMode
 
 
-class BaseJobManager(abc.ABC, TaskManager, Generic[InputModelT, SetupModelT, OutputModelT]):
-    """Abstract base class for managing background module jobs."""
+class BaseJobManager(abc.ABC, Generic[InputModelT, OutputModelT, SetupModelT]):
+    """Abstract base class for managing background module jobs.
 
+    Uses composition to delegate task lifecycle management to a TaskManager.
+    """
+
+    module_class: type[BaseModule]
+    services_mode: ServicesMode
+    _task_manager: BaseTaskManager
+
+    def __init__(
+        self,
+        module_class: type[BaseModule],
+        services_mode: ServicesMode,
+        task_manager: BaseTaskManager,
+    ) -> None:
+        """Initialize the job manager.
+
+        Args:
+            module_class: The class of the module to be managed.
+            services_mode: The mode of operation for the services (e.g., ASYNC or SYNC).
+            task_manager: The task manager instance to use for task lifecycle management.
+        """
+        self.module_class = module_class
+        self.services_mode = services_mode
+        self._task_manager = task_manager
+
+        services_config = ServicesConfig(
+            services_config_strategies=self.module_class.services_config_strategies,
+            services_config_params=self.module_class.services_config_params,
+            mode=services_mode,
+        )
+        setattr(self.module_class, "services_config", services_config)
+
+    # Properties to expose task manager attributes
+    @property
+    def tasks_sessions(self) -> dict[str, TaskSession]:
+        """Get task sessions from the task manager."""
+        return self._task_manager.tasks_sessions
+
+    @property
+    def tasks(self) -> dict[str, Any]:
+        """Get tasks from the task manager."""
+        return self._task_manager.tasks
+
+    # Delegate task lifecycle methods to task manager
+    async def create_task(
+        self,
+        task_id: str,
+        mission_id: str,
+        module: BaseModule,
+        coro: Coroutine[Any, Any, None],
+        **kwargs: Any,  # noqa: ANN401
+    ) -> None:
+        """Create a task using the task manager.
+
+        Args:
+            task_id: Unique identifier for the task
+            mission_id: Mission identifier
+            module: Module instance
+            coro: Coroutine to execute
+            **kwargs: Additional arguments for task creation
+        """
+        await self._task_manager.create_task(task_id, mission_id, module, coro, **kwargs)
+
+    async def clean_session(self, task_id: str, mission_id: str) -> bool:
+        """Clean a task's session.
+
+        Args:
+            task_id: Unique identifier for the task.
+            mission_id: Mission identifier.
+
+        Returns:
+            bool: True if the task was successfully cancelled, False otherwise.
+        """
+        return await self._task_manager.clean_session(task_id, mission_id)
+
+    async def cancel_task(self, task_id: str, mission_id: str, timeout: float | None = None) -> bool:
+        """Cancel a task.
+
+        Args:
+            task_id: Unique identifier for the task.
+            mission_id: Mission identifier.
+            timeout: Optional timeout in seconds to wait for the cancellation to complete.
+
+        Returns:
+            bool: True if the task was successfully cancelled, False otherwise.
+        """
+        return await self._task_manager.cancel_task(task_id, mission_id, timeout)
+
+    async def send_signal(self, task_id: str, mission_id: str, signal_type: str, payload: dict) -> bool:
+        """Send signal to a task.
+
+        Args:
+            task_id: Unique identifier for the task.
+            mission_id: Mission identifier.
+            signal_type: Type of signal to send.
+            payload: Payload data for the signal.
+
+        Returns:
+            bool: True if the signal was successfully sent, False otherwise.
+        """
+        return await self._task_manager.send_signal(task_id, mission_id, signal_type, payload)
+
+    async def shutdown(self, mission_id: str, timeout: float = 30.0) -> None:
+        """Shutdown all tasks."""
+        await self._task_manager.shutdown(mission_id, timeout)
+
+    @abc.abstractmethod
     async def start(self) -> None:
         """Start the job manager.
 
@@ -52,29 +159,6 @@ class BaseJobManager(abc.ABC, TaskManager, Generic[InputModelT, SetupModelT, Out
 
         return callback_wrapper
 
-    def __init__(
-        self,
-        module_class: type[BaseModule],
-        services_mode: ServicesMode,
-        **kwargs,  # noqa: ANN003
-    ) -> None:
-        """Initialize the job manager.
-
-        Args:
-            module_class: The class of the module to be managed.
-            services_mode: The mode of operation for the services (e.g., ASYNC or SYNC).
-            **kwargs: Additional keyword arguments for the job manager.
-        """
-        self.module_class = module_class
-
-        services_config = ServicesConfig(
-            services_config_strategies=self.module_class.services_config_strategies,
-            services_config_params=self.module_class.services_config_params,
-            mode=services_mode,
-        )
-        setattr(self.module_class, "services_config", services_config)
-        super().__init__(**kwargs)
-
     @abc.abstractmethod  # type: ignore
     @asynccontextmanager  # type: ignore
     async def generate_stream_consumer(self, job_id: str) -> AsyncIterator[AsyncGenerator[dict[str, Any], None]]:
@@ -110,7 +194,7 @@ class BaseJobManager(abc.ABC, TaskManager, Generic[InputModelT, SetupModelT, Out
         """
 
     @abc.abstractmethod
-    async def generate_config_setup_module_response(self, job_id: str) -> SetupModelT:
+    async def generate_config_setup_module_response(self, job_id: str) -> SetupModelT | ModuleCodeModel:
         """Generate a stream consumer for a module's output data.
 
         This method creates an asynchronous generator that streams output data
@@ -121,7 +205,7 @@ class BaseJobManager(abc.ABC, TaskManager, Generic[InputModelT, SetupModelT, Out
             job_id: The unique identifier of the job.
 
         Returns:
-            SetupModelT: the SetupModelT object fully processed.
+            SetupModelT | ModuleCodeModel: the SetupModelT object fully processed, or an error code.
         """
 
     @abc.abstractmethod
@@ -172,6 +256,22 @@ class BaseJobManager(abc.ABC, TaskManager, Generic[InputModelT, SetupModelT, Out
             ModuleStatu: The status of the job.
         """
 
+    @abc.abstractmethod
+    async def wait_for_completion(self, job_id: str) -> None:
+        """Wait for a task to complete.
+
+        This method blocks until the specified job has reached a terminal state.
+        The implementation varies by job manager type:
+        - SingleJobManager: Awaits the asyncio.Task directly
+        - TaskiqJobManager: Polls task status from SurrealDB
+
+        Args:
+            job_id: The unique identifier of the job to wait for.
+
+        Raises:
+            KeyError: If the job_id is not found.
+        """
+
     @abc.abstractmethod
     async def stop_all_modules(self) -> None:
         """Stop all currently running module jobs.

digitalkin/core/job_manager/single_job_manager.py

@@ -5,12 +5,13 @@ import datetime
 import uuid
 from collections.abc import AsyncGenerator, AsyncIterator
 from contextlib import asynccontextmanager
-from typing import Any, Generic
+from typing import Any
 
 import grpc
 
+from digitalkin.core.common import ConnectionFactory, ModuleFactory
 from digitalkin.core.job_manager.base_job_manager import BaseJobManager
-from digitalkin.core.task_manager.surrealdb_repository import SurrealDBConnection
+from digitalkin.core.task_manager.local_task_manager import LocalTaskManager
 from digitalkin.core.task_manager.task_session import TaskSession
 from digitalkin.logger import logger
 from digitalkin.models.core.task_monitor import TaskStatus
@@ -20,7 +21,7 @@ from digitalkin.modules._base_module import BaseModule
 from digitalkin.services.services_models import ServicesMode
 
 
-class SingleJobManager(BaseJobManager, Generic[InputModelT, OutputModelT, SetupModelT]):
+class SingleJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
     """Manages a single instance of a module job.
 
     This class ensures that only one instance of a module job is active at a time.
@@ -30,21 +31,29 @@ class SingleJobManager(BaseJobManager, Generic[InputModelT, OutputModelT, SetupM
 
     async def start(self) -> None:
         """Start manager."""
-        self.channel: SurrealDBConnection = SurrealDBConnection("task_manager", datetime.timedelta(seconds=5))
-        await self.channel.init_surreal_instance()
+        self.channel = await ConnectionFactory.create_surreal_connection("task_manager", datetime.timedelta(seconds=5))
 
     def __init__(
         self,
         module_class: type[BaseModule],
         services_mode: ServicesMode,
+        default_timeout: float = 10.0,
+        max_concurrent_tasks: int = 100,
     ) -> None:
         """Initialize the job manager.
 
         Args:
             module_class: The class of the module to be managed.
             services_mode: The mode of operation for the services (e.g., ASYNC or SYNC).
+            default_timeout: Default timeout for task operations
+            max_concurrent_tasks: Maximum number of concurrent tasks
         """
-        super().__init__(module_class, services_mode)
+        # Create local task manager for same-process execution
+        task_manager = LocalTaskManager(default_timeout, max_concurrent_tasks)
+
+        # Initialize base job manager with task manager
+        super().__init__(module_class, services_mode, task_manager)
+
         self._lock = asyncio.Lock()
 
     async def generate_config_setup_module_response(self, job_id: str) -> SetupModelT | ModuleCodeModel:
@@ -68,7 +77,14 @@ class SingleJobManager(BaseJobManager, Generic[InputModelT, OutputModelT, SetupM
 
         logger.debug("Module %s found: %s", job_id, session.module)
         try:
-            return await session.queue.get()
+            # Add timeout to prevent indefinite blocking
+            return await asyncio.wait_for(session.queue.get(), timeout=30.0)
+        except asyncio.TimeoutError:
+            logger.error("Timeout waiting for config setup response from module %s", job_id)
+            return ModuleCodeModel(
+                code=str(grpc.StatusCode.DEADLINE_EXCEEDED),
+                message=f"Module {job_id} did not respond within 30 seconds",
+            )
         finally:
             logger.info(f"{job_id=}: {session.queue.empty()}")
 
@@ -98,7 +114,7 @@ class SingleJobManager(BaseJobManager, Generic[InputModelT, OutputModelT, SetupM
         """
         job_id = str(uuid.uuid4())
         # TODO: Ensure the job_id is unique.
-        module = self.module_class(job_id, mission_id=mission_id, setup_id=setup_id, setup_version_id=setup_version_id)
+        module = ModuleFactory.create_module_instance(self.module_class, job_id, mission_id, setup_id, setup_version_id)
         self.tasks_sessions[job_id] = TaskSession(job_id, mission_id, self.channel, module)
 
         try:
@@ -161,26 +177,45 @@ class SingleJobManager(BaseJobManager, Generic[InputModelT, OutputModelT, SetupM
         logger.debug("Session: %s with Module %s", job_id, session.module)
 
         async def _stream() -> AsyncGenerator[dict[str, Any], Any]:
-            """Stream output data from the module.
+            """Stream output data from the module with simple blocking pattern.
+
+            This implementation uses a simple one-item-at-a-time pattern optimized
+            for local execution where we have direct access to session status:
+            1. Block waiting for each item
+            2. Check termination conditions after each item
+            3. Clean shutdown when task completes
+
+            This pattern provides:
+            - Immediate termination when task completes
+            - Direct session status monitoring
+            - Simple, predictable behavior for local tasks
 
             Yields:
                 dict: Output data generated by the module.
             """
             while True:
-                # if queue is empty but producer not finished yet, block on get()
+                # Block for next item - if queue is empty but producer not finished yet
                 msg = await session.queue.get()
                 try:
                     yield msg
                 finally:
+                    # Always mark task as done, even if consumer raises exception
                     session.queue.task_done()
 
-                # If the producer marked finished and no more items, break soon:
+                # Check termination conditions after each message
+                # This allows immediate shutdown when the task completes
                 if (
                     session.is_cancelled.is_set()
                     or (session.status is TaskStatus.COMPLETED and session.queue.empty())
                     or session.status is TaskStatus.FAILED
                 ):
-                    # and session.queue.empty():
+                    logger.debug(
+                        "Stream ending for job %s: cancelled=%s, status=%s, queue_empty=%s",
+                        job_id,
+                        session.is_cancelled.is_set(),
+                        session.status,
+                        session.queue.empty(),
+                    )
                     break
 
         yield _stream()
@@ -212,12 +247,7 @@ class SingleJobManager(BaseJobManager, Generic[InputModelT, OutputModelT, SetupM
             Exception: If the module fails to start.
         """
         job_id = str(uuid.uuid4())
-        module = self.module_class(
-            job_id,
-            mission_id=mission_id,
-            setup_id=setup_id,
-            setup_version_id=setup_version_id,
-        )
+        module = ModuleFactory.create_module_instance(self.module_class, job_id, mission_id, setup_id, setup_version_id)
         callback = await self.job_specific_callback(self.add_to_queue, job_id)
 
         await self.create_task(
@@ -251,9 +281,7 @@ class SingleJobManager(BaseJobManager, Generic[InputModelT, OutputModelT, SetupM
                 return False
             try:
                 await session.module.stop()
-
-                if job_id in self.tasks:
-                    await self.cancel_task(job_id, session.mission_id)
+                await self.cancel_task(job_id, session.mission_id)
                 logger.debug(f"session {job_id} ({session.module.name}) stopped successfully")
             except Exception as e:
                 logger.error(f"Error while stopping module {job_id}: {e}")
@@ -273,15 +301,42 @@ class SingleJobManager(BaseJobManager, Generic[InputModelT, OutputModelT, SetupM
         session = self.tasks_sessions.get(job_id, None)
         return session.status if session is not None else TaskStatus.FAILED
 
+    async def wait_for_completion(self, job_id: str) -> None:
+        """Wait for a task to complete by awaiting its asyncio.Task.
+
+        Args:
+            job_id: The unique identifier of the job to wait for.
+
+        Raises:
+            KeyError: If the job_id is not found in tasks.
+        """
+        if job_id not in self._task_manager.tasks:
+            msg = f"Job {job_id} not found"
+            raise KeyError(msg)
+        await self._task_manager.tasks[job_id]
+
     async def stop_all_modules(self) -> None:
         """Stop all currently running module jobs.
 
-        This method ensures that all active jobs are gracefully terminated.
+        This method ensures that all active jobs are gracefully terminated
+        and closes the SurrealDB connection.
         """
+        # Snapshot job IDs while holding lock
         async with self._lock:
-            stop_tasks = [self.stop_module(job_id) for job_id in list(self.tasks_sessions.keys())]
-            if stop_tasks:
-                await asyncio.gather(*stop_tasks, return_exceptions=True)
+            job_ids = list(self.tasks_sessions.keys())
+
+        # Release lock before calling stop_module (which has its own lock)
+        if job_ids:
+            stop_tasks = [self.stop_module(job_id) for job_id in job_ids]
+            await asyncio.gather(*stop_tasks, return_exceptions=True)
+
+        # Close SurrealDB connection after stopping all modules
+        if hasattr(self, "channel"):
+            try:
+                await self.channel.close()
+                logger.info("SingleJobManager: SurrealDB connection closed")
+            except Exception as e:
+                logger.warning("Failed to close SurrealDB connection: %s", e)
 
     async def list_modules(self) -> dict[str, dict[str, Any]]:
         """List all modules along with their statuses.