digitalkin 0.2.12__py3-none-any.whl → 0.2.14__py3-none-any.whl

digitalkin/modules/job_manager/base_job_manager.py

@@ -0,0 +1,181 @@
+"""Background module manager."""
+
+import abc
+from collections.abc import AsyncGenerator, AsyncIterator, Callable, Coroutine
+from contextlib import asynccontextmanager
+from typing import Any, Generic
+
+from digitalkin.models import ModuleStatus
+from digitalkin.models.module import InputModelT, OutputModelT, SetupModelT
+from digitalkin.models.module.module_types import ConfigSetupModelT
+from digitalkin.modules._base_module import BaseModule
+from digitalkin.services.services_config import ServicesConfig
+from digitalkin.services.services_models import ServicesMode
+
+
+class BaseJobManager(abc.ABC, Generic[InputModelT, SetupModelT, ConfigSetupModelT]):
+    """Abstract base class for managing background module jobs."""
+
+    async def _start(self) -> None:
+        """Start the job manager.
+
+        This method initializes any necessary resources or configurations
+        required for the job manager to function.
+        """
+
+    @staticmethod
+    async def job_specific_callback(
+        callback: Callable[[str, OutputModelT], Coroutine[Any, Any, None]], job_id: str
+    ) -> Callable[[OutputModelT], Coroutine[Any, Any, None]]:
+        """Generate a job-specific callback function.
+
+        Args:
+            callback: The callback function to be executed when the job completes.
+            job_id: The unique identifier of the job.
+
+        Returns:
+            Callable: A wrapped callback function that includes the job ID.
+        """
+
+        def callback_wrapper(output_data: OutputModelT) -> Coroutine[Any, Any, None]:
+            """Wrapper for the callback function.
+
+            Args:
+                output_data: The output data produced by the job.
+
+            Returns:
+                Coroutine: The wrapped callback function.
+            """
+            return callback(job_id, output_data)
+
+        return callback_wrapper
+
+    def __init__(
+        self,
+        module_class: type[BaseModule],
+        services_mode: ServicesMode,
+    ) -> 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).
+        """
+        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)
+
+    @abc.abstractmethod  # type: ignore
+    @asynccontextmanager  # type: ignore
+    async def generate_stream_consumer(self, job_id: str) -> AsyncIterator[AsyncGenerator[dict[str, Any], None]]:
+        """Generate a stream consumer for the job's message stream.
+
+        Args:
+            job_id: The unique identifier of the job to filter messages for.
+
+        Yields:
+            dict[str, Any]: The messages from the associated module's stream.
+        """
+
+    @abc.abstractmethod
+    async def create_module_instance_job(
+        self,
+        input_data: InputModelT,
+        setup_data: SetupModelT,
+        mission_id: str,
+        setup_version_id: str,
+    ) -> str:
+        """Create and start a new job for the module's instance.
+
+        Args:
+            input_data: The input data required to start the job.
+            setup_data: The setup configuration for the module.
+            mission_id: The mission ID associated with the job.
+            setup_version_id: The setup ID associated with the module.
+
+        Returns:
+            str: The unique identifier (job ID) of the created job.
+        """
+
+    @abc.abstractmethod
+    async def generate_config_setup_module_response(self, job_id: str) -> SetupModelT:
+        """Generate a stream consumer for a module's output data.
+
+        This method creates an asynchronous generator that streams output data
+        from a specific module job. If the module does not exist, it generates
+        an error message.
+
+        Args:
+            job_id: The unique identifier of the job.
+
+        Returns:
+            SetupModelT: the SetupModelT object fully processed.
+        """
+
+    @abc.abstractmethod
+    async def create_config_setup_instance_job(
+        self,
+        config_setup_data: ConfigSetupModelT,
+        setup_data: SetupModelT,
+        mission_id: str,
+        setup_version_id: str,
+    ) -> str:
+        """Create and start a new module job.
+
+        This method initializes a new module job, assigns it a unique job ID,
+        and starts it in the background.
+
+        Args:
+            config_setup_data: The input data required to start the job.
+            setup_data: The setup configuration for the module.
+            mission_id: The mission ID associated with the job.
+            setup_version_id: The setup ID.
+
+        Returns:
+            str: The unique identifier (job ID) of the created job.
+
+        Raises:
+            Exception: If the module fails to start.
+        """
+
+    @abc.abstractmethod
+    async def stop_module(self, job_id: str) -> bool:
+        """Stop a running module job.
+
+        Args:
+            job_id: The unique identifier of the job to stop.
+
+        Returns:
+            bool: True if the job was successfully stopped, False if it does not exist.
+        """
+
+    @abc.abstractmethod
+    async def get_module_status(self, job_id: str) -> ModuleStatus | None:
+        """Retrieve the status of a module job.
+
+        Args:
+            job_id: The unique identifier of the job.
+
+        Returns:
+            ModuleStatus | None: The status of the job, or None if the job does not exist.
+        """
+
+    @abc.abstractmethod
+    async def stop_all_modules(self) -> None:
+        """Stop all currently running module jobs.
+
+        This method ensures that all active jobs are gracefully terminated.
+        """
+
+    @abc.abstractmethod
+    async def list_modules(self) -> dict[str, dict[str, Any]]:
+        """List all modules along with their statuses.
+
+        Returns:
+            dict[str, dict[str, Any]]: A dictionary containing information about all modules and their statuses.
+        """

digitalkin/modules/job_manager/job_manager_models.py

@@ -0,0 +1,44 @@
+"""Job manager models."""
+
+from enum import Enum
+
+from pydantic import BaseModel
+
+from digitalkin.modules.job_manager.base_job_manager import BaseJobManager
+
+
+class StreamCodeModel(BaseModel):
+    """Typed error/code model."""
+
+    code: str
+
+
+class JobManagerMode(Enum):
+    """Job manager mode."""
+
+    SINGLE = "single"
+    TASKIQ = "taskiq"
+
+    def __str__(self) -> str:
+        """Get the string representation of the job manager mode.
+
+        Returns:
+            str: job manager mode name.
+        """
+        return self.value
+
+    def get_manager_class(self) -> type[BaseJobManager]:
+        """Get the job manager class based on the mode.
+
+        Returns:
+            type: The job manager class.
+        """
+        match self:
+            case JobManagerMode.SINGLE:
+                from digitalkin.modules.job_manager.single_job_manager import SingleJobManager  # noqa: PLC0415
+
+                return SingleJobManager
+            case JobManagerMode.TASKIQ:
+                from digitalkin.modules.job_manager.taskiq_job_manager import TaskiqJobManager  # noqa: PLC0415
+
+                return TaskiqJobManager

digitalkin/modules/job_manager/single_job_manager.py

@@ -0,0 +1,285 @@
+"""Background module manager with single instance."""
+
+import asyncio
+import uuid
+from collections.abc import AsyncGenerator, AsyncIterator
+from contextlib import asynccontextmanager
+from typing import Any, Generic
+
+import grpc
+
+from digitalkin.logger import logger
+from digitalkin.models import ModuleStatus
+from digitalkin.models.module import ConfigSetupModelT, InputModelT, OutputModelT, SetupModelT
+from digitalkin.modules._base_module import BaseModule
+from digitalkin.modules.job_manager.base_job_manager import BaseJobManager
+from digitalkin.services.services_models import ServicesMode
+
+
+class SingleJobManager(BaseJobManager, Generic[InputModelT, SetupModelT, ConfigSetupModelT]):
+    """Manages a single instance of a module job.
+
+    This class ensures that only one instance of a module job is active at a time.
+    It provides functionality to create, stop, and monitor module jobs, as well as
+    to handle their output data.
+    """
+
+    modules: dict[str, BaseModule]
+    queue: dict[str, asyncio.Queue]
+
+    def __init__(
+        self,
+        module_class: type[BaseModule],
+        services_mode: ServicesMode,
+    ) -> 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).
+        """
+        super().__init__(module_class, services_mode)
+
+        self._lock = asyncio.Lock()
+        self.modules: dict[str, BaseModule] = {}
+        self.queues: dict[str, asyncio.Queue] = {}
+
+    async def generate_config_setup_module_response(self, job_id: str) -> SetupModelT:
+        """Generate a stream consumer for a module's output data.
+
+        This method creates an asynchronous generator that streams output data
+        from a specific module job. If the module does not exist, it generates
+        an error message.
+
+        Args:
+            job_id: The unique identifier of the job.
+
+        Returns:
+            SetupModelT: the SetupModelT object fully processed.
+        """
+        module = self.modules.get(job_id, None)
+        logger.debug("Module %s found: %s", job_id, module)
+
+        try:
+            return await self.queues[job_id].get()
+        finally:
+            logger.info(f"{job_id=}: {self.queues[job_id].empty()}")
+            del self.queues[job_id]
+
+    async def create_config_setup_instance_job(
+        self,
+        config_setup_data: ConfigSetupModelT,
+        setup_data: SetupModelT,
+        mission_id: str,
+        setup_version_id: str,
+    ) -> str:
+        """Create and start a new module setup configuration job.
+
+        This method initializes a new module job, assigns it a unique job ID,
+        and starts the config setup it in the background.
+
+        Args:
+            config_setup_data: The input data required to start the job.
+            setup_data: The setup configuration for the module.
+            mission_id: The mission ID associated with the job.
+            setup_version_id: The setup ID.
+
+        Returns:
+            str: The unique identifier (job ID) of the created job.
+
+        Raises:
+            Exception: If the module fails to start.
+        """
+        job_id = str(uuid.uuid4())
+        # TODO: Ensure the job_id is unique.
+        module = self.module_class(job_id, mission_id=mission_id, setup_version_id=setup_version_id)
+        self.modules[job_id] = module
+        self.queues[job_id] = asyncio.Queue()
+
+        try:
+            await module.start_config_setup(
+                config_setup_data,
+                setup_data,
+                await self.job_specific_callback(self.add_to_queue, job_id),
+            )
+            logger.debug("Module %s (%s) started successfully", job_id, module.name)
+        except Exception:
+            # Remove the module from the manager in case of an error.
+            del self.modules[job_id]
+            logger.exception("Failed to start module %s: %s", job_id)
+            raise
+        else:
+            return job_id
+
+    async def add_to_queue(self, job_id: str, output_data: OutputModelT) -> None:  # type: ignore
+        """Add output data to the queue for a specific job.
+
+        This method is used as a callback to handle output data generated by a module job.
+
+        Args:
+            job_id: The unique identifier of the job.
+            output_data: The output data produced by the job.
+        """
+        await self.queues[job_id].put(output_data.model_dump())
+
+    @asynccontextmanager  # type: ignore
+    async def generate_stream_consumer(self, job_id: str) -> AsyncIterator[AsyncGenerator[dict[str, Any], None]]:  # type: ignore
+        """Generate a stream consumer for a module's output data.
+
+        This method creates an asynchronous generator that streams output data
+        from a specific module job. If the module does not exist, it generates
+        an error message.
+
+        Args:
+            job_id: The unique identifier of the job.
+
+        Yields:
+            AsyncGenerator: A stream of output data or error messages.
+        """
+        module = self.modules.get(job_id, None)
+
+        logger.debug("Module %s found: %s", job_id, module)
+
+        async def _stream() -> AsyncGenerator[dict[str, Any], Any]:
+            """Stream output data from the module.
+
+            Yields:
+                dict: Output data generated by the module.
+            """
+            if module is None:
+                yield {
+                    "error": {
+                        "error_message": f"Module {job_id} not found",
+                        "code": grpc.StatusCode.NOT_FOUND,
+                    }
+                }
+                return
+
+            try:
+                while module.status == ModuleStatus.RUNNING or (
+                    not self.queues[job_id].empty()
+                    and module.status
+                    in {
+                        ModuleStatus.STOPPED,
+                        ModuleStatus.STOPPING,
+                    }
+                ):
+                    logger.info(f"{job_id=}: {module.status=}")
+                    yield await self.queues[job_id].get()
+                logger.info(f"{job_id=}: {module.status=} | {self.queues[job_id].empty()}")
+
+            finally:
+                del self.queues[job_id]
+
+        yield _stream()
+
+    async def create_module_instance_job(
+        self,
+        input_data: InputModelT,
+        setup_data: SetupModelT,
+        mission_id: str,
+        setup_version_id: str,
+    ) -> str:
+        """Create and start a new module job.
+
+        This method initializes a new module job, assigns it a unique job ID,
+        and starts it in the background.
+
+        Args:
+            input_data: The input data required to start the job.
+            setup_data: The setup configuration for the module.
+            mission_id: The mission ID associated with the job.
+            setup_version_id: The setup ID associated with the module.
+
+        Returns:
+            str: The unique identifier (job ID) of the created job.
+
+        Raises:
+            Exception: If the module fails to start.
+        """
+        job_id = str(uuid.uuid4())
+        # TODO: Ensure the job_id is unique.
+        module = self.module_class(job_id, mission_id=mission_id, setup_version_id=setup_version_id)
+        self.modules[job_id] = module
+        self.queues[job_id] = asyncio.Queue()
+
+        try:
+            await module.start(
+                input_data,
+                setup_data,
+                await self.job_specific_callback(self.add_to_queue, job_id),
+            )
+            logger.debug("Module %s (%s) started successfully", job_id, module.name)
+        except Exception:
+            # Remove the module from the manager in case of an error.
+            del self.modules[job_id]
+            logger.exception("Failed to start module %s: %s", job_id)
+            raise
+        else:
+            return job_id
+
+    async def stop_module(self, job_id: str) -> bool:
+        """Stop a running module job.
+
+        Args:
+            job_id: The unique identifier of the job to stop.
+
+        Returns:
+            bool: True if the module was successfully stopped, False if it does not exist.
+
+        Raises:
+            Exception: If an error occurs while stopping the module.
+        """
+        async with self._lock:
+            module = self.modules.get(job_id)
+            if not module:
+                logger.warning(f"Module {job_id} not found")
+                return False
+            try:
+                await module.stop()
+                # should maybe be added in finally
+                del self.queues[job_id]
+                del self.modules[job_id]
+                logger.debug(f"Module {job_id} ({module.name}) stopped successfully")
+            except Exception as e:
+                logger.error(f"Error while stopping module {job_id}: {e}")
+                raise
+            else:
+                return True
+
+    async def get_module_status(self, job_id: str) -> ModuleStatus | None:
+        """Retrieve the status of a module job.
+
+        Args:
+            job_id: The unique identifier of the job.
+
+        Returns:
+            ModuleStatus | None: The status of the module, or None if it does not exist.
+        """
+        module = self.modules.get(job_id)
+        return module.status if module else None
+
+    async def stop_all_modules(self) -> None:
+        """Stop all currently running module jobs.
+
+        This method ensures that all active jobs are gracefully terminated.
+        """
+        async with self._lock:
+            stop_tasks = [self.stop_module(job_id) for job_id in list(self.modules.keys())]
+            if stop_tasks:
+                await asyncio.gather(*stop_tasks, return_exceptions=True)
+
+    async def list_modules(self) -> dict[str, dict[str, Any]]:
+        """List all modules along with their statuses.
+
+        Returns:
+            dict[str, dict[str, Any]]: A dictionary containing information about all modules and their statuses.
+        """
+        return {
+            job_id: {
+                "name": module.name,
+                "status": module.status,
+                "class": module.__class__.__name__,
+            }
+            for job_id, module in self.modules.items()
+        }