digitalkin 0.2.25rc0__py3-none-any.whl → 0.3.2.dev14__py3-none-any.whl

digitalkin/grpc_servers/utils/factory.py

@@ -1,180 +0,0 @@
-"""Factory functions for creating gRPC servers."""
-
-from pathlib import Path
-from typing import Any
-
-from digitalkin.grpc_servers.module_server import ModuleServer
-from digitalkin.grpc_servers.registry_server import RegistryServer
-from digitalkin.grpc_servers.utils.models import (
-    ModuleServerConfig,
-    RegistryServerConfig,
-    SecurityMode,
-    ServerCredentials,
-    ServerMode,
-)
-from digitalkin.modules._base_module import BaseModule
-
-
-def create_module_server(
-    module: type[BaseModule],
-    host: str = "0.0.0.0",  # noqa: S104
-    port: int = 50051,
-    max_workers: int = 10,
-    mode: str = "sync",
-    security: str = "insecure",
-    registry_address: str | None = None,
-    server_key_path: str | None = None,
-    server_cert_path: str | None = None,
-    root_cert_path: str | None = None,
-    server_options: list[tuple[str, Any]] | None = None,
-) -> ModuleServer:
-    """Create a new module server.
-
-    Args:
-        module: The module to serve.
-        host: The host address to bind to.
-        port: The port to listen on.
-        max_workers: Maximum number of workers for the thread pool (sync mode only).
-        mode: Server mode ("sync" or "async").
-        security: Security mode ("secure" or "insecure").
-        registry_address: Optional address of a registry server for auto-registration.
-        server_key_path: Path to server private key (required for secure mode).
-        server_cert_path: Path to server certificate (required for secure mode).
-        root_cert_path: Optional path to root certificate.
-        server_options: Additional server options.
-
-    Returns:
-        A configured ModuleServer instance.
-
-    Raises:
-        ValueError: If secure mode is requested but credentials are missing.
-    """
-    # Create configuration with credentials if needed
-    config = _create_server_config(
-        ModuleServerConfig,
-        host=host,
-        port=port,
-        max_workers=max_workers,
-        mode=mode,
-        security=security,
-        server_key_path=server_key_path,
-        server_cert_path=server_cert_path,
-        root_cert_path=root_cert_path,
-        server_options=server_options,
-        registry_address=registry_address,
-    )
-
-    # Create and return the server
-    return ModuleServer(module, config)
-
-
-def create_registry_server(
-    host: str = "0.0.0.0",  # noqa: S104
-    port: int = 50052,
-    max_workers: int = 10,
-    mode: str = "sync",
-    security: str = "insecure",
-    database_url: str | None = None,
-    server_key_path: str | None = None,
-    server_cert_path: str | None = None,
-    root_cert_path: str | None = None,
-    server_options: list[tuple[str, Any]] | None = None,
-) -> RegistryServer:
-    """Create a new registry server.
-
-    Args:
-        host: The host address to bind to.
-        port: The port to listen on.
-        max_workers: Maximum number of workers for the thread pool (sync mode only).
-        mode: Server mode ("sync" or "async").
-        security: Security mode ("secure" or "insecure").
-        database_url: Optional database URL for registry data storage.
-        server_key_path: Path to server private key (required for secure mode).
-        server_cert_path: Path to server certificate (required for secure mode).
-        root_cert_path: Optional path to root certificate.
-        server_options: Additional server options.
-
-    Returns:
-        A configured RegistryServer instance.
-
-    Raises:
-        ValueError: If secure mode is requested but credentials are missing.
-    """
-    # Create configuration with credentials if needed
-    config = _create_server_config(
-        RegistryServerConfig,
-        host=host,
-        port=port,
-        max_workers=max_workers,
-        mode=mode,
-        security=security,
-        server_key_path=server_key_path,
-        server_cert_path=server_cert_path,
-        root_cert_path=root_cert_path,
-        server_options=server_options,
-        database_url=database_url,
-    )
-
-    # Create and return the server
-    return RegistryServer(config)
-
-
-def _create_server_config(
-    config_class: Any,
-    host: str,
-    port: int,
-    max_workers: int,
-    mode: str,
-    security: str,
-    server_key_path: str | None,
-    server_cert_path: str | None,
-    root_cert_path: str | None,
-    server_options: list[tuple[str, Any]] | None,
-    **kwargs,
-) -> ModuleServerConfig | RegistryServerConfig:
-    """Create a server configuration with appropriate settings.
-
-    Args:
-        config_class: The configuration class to instantiate.
-        host: The host address.
-        port: The port number.
-        max_workers: Maximum number of workers.
-        mode: Server mode.
-        security: Security mode.
-        server_key_path: Path to server key.
-        server_cert_path: Path to server certificate.
-        root_cert_path: Path to root certificate.
-        server_options: Additional server options.
-        **kwargs: Additional configuration parameters.
-
-    Returns:
-        A server configuration instance.
-
-    Raises:
-        ValueError: If secure mode is requested but credentials are missing.
-    """
-    # Create basic config
-    config_params = {
-        "host": host,
-        "port": port,
-        "max_workers": max_workers,
-        "mode": ServerMode(mode),
-        "security": SecurityMode(security),
-        "server_options": server_options or [],
-        **kwargs,
-    }
-
-    # Add credentials if secure mode
-    if security == "secure":
-        if not server_key_path or not server_cert_path:
-            raise ValueError(
-                "Server key and certificate paths are required for secure mode"
-            )
-
-        config_params["credentials"] = ServerCredentials(
-            server_key_path=Path(server_key_path),
-            server_cert_path=Path(server_cert_path),
-            root_cert_path=Path(root_cert_path) if root_cert_path else None,
-        )
-
-    return config_class(**config_params)

digitalkin/modules/job_manager/single_job_manager.py

@@ -1,294 +0,0 @@
-"""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 InputModelT, OutputModelT, SetupModelT
-from digitalkin.modules._base_module import BaseModule
-from digitalkin.modules.job_manager.base_job_manager import BaseJobManager
-from digitalkin.modules.job_manager.job_manager_models import StreamCodeModel
-from digitalkin.services.services_models import ServicesMode
-
-
-class SingleJobManager(BaseJobManager, Generic[InputModelT, 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.
-    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: SetupModelT,
-        mission_id: str,
-        setup_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_id: The setup ID associated with the module.
-            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_id=setup_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,
-                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()
-
-            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_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_id: The setup ID associated with the module.
-            setup_version_id: The setup Version 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_id=setup_id,
-            setup_version_id=setup_version_id,
-        )
-        self.modules[job_id] = module
-        self.queues[job_id] = asyncio.Queue()
-        callback = await self.job_specific_callback(self.add_to_queue, job_id)
-
-        try:
-            await module.start(
-                input_data,
-                setup_data,
-                callback,
-                done_callback=lambda _: asyncio.create_task(callback(StreamCodeModel(code="__END_OF_STREAM__"))),
-            )
-            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()
-        }