digitalkin 0.3.2.dev2__py3-none-any.whl

digitalkin/grpc_servers/module_servicer.py

@@ -0,0 +1,516 @@
+"""Module servicer implementation for DigitalKin."""
+
+from argparse import ArgumentParser, Namespace
+from collections.abc import AsyncGenerator
+from typing import Any
+
+import grpc
+from agentic_mesh_protocol.module.v1 import (
+    information_pb2,
+    lifecycle_pb2,
+    module_service_pb2_grpc,
+    monitoring_pb2,
+)
+from google.protobuf import json_format, struct_pb2
+
+from digitalkin.core.job_manager.base_job_manager import BaseJobManager
+from digitalkin.grpc_servers.utils.exceptions import ServicerError
+from digitalkin.logger import logger
+from digitalkin.models.core.job_manager_models import JobManagerMode
+from digitalkin.models.module.module import ModuleStatus
+from digitalkin.modules._base_module import BaseModule
+from digitalkin.services.services_models import ServicesMode
+from digitalkin.services.setup.default_setup import DefaultSetup
+from digitalkin.services.setup.grpc_setup import GrpcSetup
+from digitalkin.services.setup.setup_strategy import SetupStrategy
+from digitalkin.utils.arg_parser import ArgParser
+from digitalkin.utils.development_mode_action import DevelopmentModeMappingAction
+
+
+class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
+    """Implementation of the ModuleService.
+
+    This servicer handles interactions with a DigitalKin module.
+
+    Attributes:
+        module: The module instance being served.
+        active_jobs: Dictionary tracking active module jobs.
+    """
+
+    args: Namespace
+    setup: SetupStrategy
+    job_manager: BaseJobManager
+
+    def _add_parser_args(self, parser: ArgumentParser) -> None:
+        super()._add_parser_args(parser)
+        parser.add_argument(
+            "-d",
+            "--dev-mode",
+            env_var="SERVICE_MODE",
+            choices=ServicesMode.__members__,
+            default="local",
+            action=DevelopmentModeMappingAction,
+            dest="services_mode",
+            help="Define Module Service configurations for endpoints",
+        )
+        parser.add_argument(
+            "-jm",
+            "--job-manager",
+            type=JobManagerMode,
+            choices=list(JobManagerMode),
+            default=JobManagerMode.SINGLE,
+            dest="job_manager_mode",
+            help="Define Module job manager configurations for load balancing",
+        )
+
+    def __init__(self, module_class: type[BaseModule]) -> None:
+        """Initialize the module servicer.
+
+        Args:
+            module_class: The module type to serve.
+        """
+        super().__init__()
+        module_class.discover()
+        self.module_class = module_class
+        job_manager_class = self.args.job_manager_mode.get_manager_class()
+        self.job_manager = job_manager_class(module_class, self.args.services_mode)
+
+        logger.debug(
+            "ModuleServicer initialized with job manager: %s",
+            self.args.job_manager_mode,
+            extra={"job_manager": self.job_manager},
+        )
+        self.setup = GrpcSetup() if self.args.services_mode == ServicesMode.REMOTE else DefaultSetup()
+
+    async def ConfigSetupModule(  # noqa: N802
+        self,
+        request: lifecycle_pb2.ConfigSetupModuleRequest,
+        context: grpc.aio.ServicerContext,
+    ) -> lifecycle_pb2.ConfigSetupModuleResponse:
+        """Configure the module setup.
+
+        Args:
+            request: The configuration request.
+            context: The gRPC context.
+
+        Returns:
+            A response indicating success or failure.
+
+        Raises:
+            ServicerError: if the setup data is not returned or job creation fails.
+        """
+        logger.info(
+            "ConfigSetupVersion called for module: '%s'",
+            self.module_class.__name__,
+            extra={
+                "module_class": self.module_class,
+                "setup_version": request.setup_version,
+                "mission_id": request.mission_id,
+            },
+        )
+        # Process the module input
+        # TODO: Secret should be used here as well
+        setup_version = request.setup_version
+        config_setup_data = self.module_class.create_config_setup_model(json_format.MessageToDict(request.content))
+        setup_version_data = await self.module_class.create_setup_model(
+            json_format.MessageToDict(request.setup_version.content),
+            config_fields=True,
+        )
+
+        if not setup_version_data:
+            msg = "No setup data returned."
+            raise ServicerError(msg)
+
+        if not config_setup_data:
+            msg = "No config setup data returned."
+            raise ServicerError(msg)
+
+        # create a task to run the module in background
+        job_id = await self.job_manager.create_config_setup_instance_job(
+            config_setup_data,
+            request.mission_id,
+            setup_version.setup_id,
+            setup_version.id,
+        )
+
+        if job_id is None:
+            context.set_code(grpc.StatusCode.NOT_FOUND)
+            context.set_details("Failed to create module instance")
+            return lifecycle_pb2.ConfigSetupModuleResponse(success=False)
+
+        updated_setup_data = await self.job_manager.generate_config_setup_module_response(job_id)
+        logger.info("Setup updated")
+        logger.debug(f"Updated setup data: {updated_setup_data=}")
+        setup_version.content = json_format.ParseDict(
+            updated_setup_data,
+            struct_pb2.Struct(),
+            ignore_unknown_fields=True,
+        )
+        return lifecycle_pb2.ConfigSetupModuleResponse(success=True, setup_version=setup_version)
+
+    async def StartModule(  # noqa: N802
+        self,
+        request: lifecycle_pb2.StartModuleRequest,
+        context: grpc.aio.ServicerContext,
+    ) -> AsyncGenerator[lifecycle_pb2.StartModuleResponse, Any]:
+        """Start a module execution.
+
+        Args:
+            request: Iterator of start module requests.
+            context: The gRPC context.
+
+        Yields:
+            Responses during module execution.
+
+        Raises:
+            ServicerError: the necessary query didn't work.
+        """
+        logger.info(
+            "StartModule called for module: '%s'",
+            self.module_class.__name__,
+            extra={"module_class": self.module_class, "setup_id": request.setup_id, "mission_id": request.mission_id},
+        )
+        # Process the module input
+        # TODO: Check failure of input data format
+        input_data = self.module_class.create_input_model(json_format.MessageToDict(request.input))
+
+        setup_data_class = self.setup.get_setup(
+            setup_dict={
+                "setup_id": request.setup_id,
+                "mission_id": request.mission_id,
+            }
+        )
+
+        if not setup_data_class:
+            msg = "No setup data returned."
+            raise ServicerError(msg)
+
+        setup_data = await self.module_class.create_setup_model(setup_data_class.current_setup_version.content)
+
+        # create a task to run the module in background
+        job_id = await self.job_manager.create_module_instance_job(
+            input_data,
+            setup_data,
+            mission_id=request.mission_id,
+            setup_id=setup_data_class.current_setup_version.setup_id,
+            setup_version_id=setup_data_class.current_setup_version.id,
+        )
+
+        if job_id is None:
+            context.set_code(grpc.StatusCode.NOT_FOUND)
+            context.set_details("Failed to create module instance")
+            yield lifecycle_pb2.StartModuleResponse(success=False)
+            return
+
+        try:
+            async with self.job_manager.generate_stream_consumer(job_id) as stream:  # type: ignore
+                async for message in stream:
+                    if message.get("error", None) is not None:
+                        logger.error("Error in output_data", extra={"message": message})
+                        context.set_code(message["error"]["code"])
+                        context.set_details(message["error"]["error_message"])
+                        yield lifecycle_pb2.StartModuleResponse(success=False, job_id=job_id)
+                        break
+
+                    if message.get("exception", None) is not None:
+                        logger.error("Exception in output_data", extra={"message": message})
+                        context.set_code(message["short_description"])
+                        context.set_details(message["exception"])
+                        yield lifecycle_pb2.StartModuleResponse(success=False, job_id=job_id)
+                        break
+
+                    logger.info("Yielding message from job %s: %s", job_id, message)
+                    proto = json_format.ParseDict(message, struct_pb2.Struct(), ignore_unknown_fields=True)
+                    yield lifecycle_pb2.StartModuleResponse(success=True, output=proto, job_id=job_id)
+
+                    if message.get("root", {}).get("protocol") == "end_of_stream":
+                        logger.info(
+                            "End of stream signal received",
+                            extra={"job_id": job_id, "mission_id": request.mission_id},
+                        )
+                        break
+        finally:
+            await self.job_manager.wait_for_completion(job_id)
+            await self.job_manager.clean_session(job_id, mission_id=request.mission_id)
+
+        logger.info("Job %s finished", job_id)
+
+    async def StopModule(  # noqa: N802
+        self,
+        request: lifecycle_pb2.StopModuleRequest,
+        context: grpc.ServicerContext,
+    ) -> lifecycle_pb2.StopModuleResponse:
+        """Stop a running module execution.
+
+        Args:
+            request: The stop module request.
+            context: The gRPC context.
+
+        Returns:
+            A response indicating success or failure.
+        """
+        logger.debug("StopModule called for module: '%s'", self.module_class.__name__)
+
+        response: bool = await self.job_manager.stop_module(request.job_id)
+        if not response:
+            message = f"Job {request.job_id} not found"
+            logger.warning(message)
+            context.set_code(grpc.StatusCode.NOT_FOUND)
+            context.set_details(message)
+            return lifecycle_pb2.StopModuleResponse(success=False)
+
+        logger.debug("Job %s stopped successfully", request.job_id, extra={"job_id": request.job_id})
+        return lifecycle_pb2.StopModuleResponse(success=True)
+
+    async def GetModuleStatus(  # noqa: N802
+        self,
+        request: monitoring_pb2.GetModuleStatusRequest,
+        context: grpc.ServicerContext,
+    ) -> monitoring_pb2.GetModuleStatusResponse:
+        """Get the status of a module.
+
+        Args:
+            request: The get module status request.
+            context: The gRPC context.
+
+        Returns:
+            A response with the module status.
+        """
+        logger.debug("GetModuleStatus called for module: '%s'", self.module_class.__name__)
+
+        if not request.job_id:
+            logger.debug("Job %s status: '%s'", request.job_id, ModuleStatus.NOT_FOUND)
+            return monitoring_pb2.GetModuleStatusResponse(
+                success=False,
+                status=ModuleStatus.NOT_FOUND.name,
+                job_id=request.job_id,
+            )
+
+        status = await self.job_manager.get_module_status(request.job_id)
+
+        if status is None:
+            message = f"Job {request.job_id} not found"
+            logger.warning(message)
+            context.set_code(grpc.StatusCode.NOT_FOUND)
+            context.set_details(message)
+            return monitoring_pb2.GetModuleStatusResponse()
+
+        logger.debug("Job %s status: '%s'", request.job_id, status)
+        return monitoring_pb2.GetModuleStatusResponse(
+            success=True,
+            status=status.name,
+            job_id=request.job_id,
+        )
+
+    async def GetModuleJobs(  # noqa: N802
+        self,
+        request: monitoring_pb2.GetModuleJobsRequest,  # noqa: ARG002
+        context: grpc.ServicerContext,  # noqa: ARG002
+    ) -> monitoring_pb2.GetModuleJobsResponse:
+        """Get information about the module's jobs.
+
+        Args:
+            request: The get module jobs request.
+            context: The gRPC context.
+
+        Returns:
+            A response with information about active jobs.
+        """
+        logger.debug("GetModuleJobs called for module: '%s'", self.module_class.__name__)
+
+        modules = await self.job_manager.list_modules()
+
+        # Create job info objects for each active job
+        return monitoring_pb2.GetModuleJobsResponse(
+            jobs=[
+                monitoring_pb2.JobInfo(
+                    job_id=job_id,
+                    job_status=job_data["status"].name,
+                )
+                for job_id, job_data in modules.items()
+            ],
+        )
+
+    async def GetModuleInput(  # noqa: N802
+        self,
+        request: information_pb2.GetModuleInputRequest,
+        context: grpc.ServicerContext,
+    ) -> information_pb2.GetModuleInputResponse:
+        """Get information about the module's expected input.
+
+        Args:
+            request: The get module input request.
+            context: The gRPC context.
+
+        Returns:
+            A response with the module's input schema.
+        """
+        logger.debug("GetModuleInput called for module: '%s'", self.module_class.__name__)
+
+        # Get input schema if available
+        try:
+            # Convert schema to proto format
+            input_schema_proto = await self.module_class.get_input_format(
+                llm_format=request.llm_format,
+            )
+            input_format_struct = json_format.Parse(
+                text=input_schema_proto,
+                message=struct_pb2.Struct(),  # pylint: disable=no-member
+                ignore_unknown_fields=True,
+            )
+        except NotImplementedError as e:
+            logger.warning(e)
+            context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+            context.set_details(str(e))
+            return information_pb2.GetModuleInputResponse()
+
+        return information_pb2.GetModuleInputResponse(
+            success=True,
+            input_schema=input_format_struct,
+        )
+
+    async def GetModuleOutput(  # noqa: N802
+        self,
+        request: information_pb2.GetModuleOutputRequest,
+        context: grpc.ServicerContext,
+    ) -> information_pb2.GetModuleOutputResponse:
+        """Get information about the module's expected output.
+
+        Args:
+            request: The get module output request.
+            context: The gRPC context.
+
+        Returns:
+            A response with the module's output schema.
+        """
+        logger.debug("GetModuleOutput called for module: '%s'", self.module_class.__name__)
+
+        # Get output schema if available
+        try:
+            # Convert schema to proto format
+            output_schema_proto = await self.module_class.get_output_format(
+                llm_format=request.llm_format,
+            )
+            output_format_struct = json_format.Parse(
+                text=output_schema_proto,
+                message=struct_pb2.Struct(),  # pylint: disable=no-member
+                ignore_unknown_fields=True,
+            )
+        except NotImplementedError as e:
+            logger.warning(e)
+            context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+            context.set_details(str(e))
+            return information_pb2.GetModuleOutputResponse()
+
+        return information_pb2.GetModuleOutputResponse(
+            success=True,
+            output_schema=output_format_struct,
+        )
+
+    async def GetModuleSetup(  # noqa: N802
+        self,
+        request: information_pb2.GetModuleSetupRequest,
+        context: grpc.ServicerContext,
+    ) -> information_pb2.GetModuleSetupResponse:
+        """Get information about the module's setup and configuration.
+
+        Args:
+            request: The get module setup request.
+            context: The gRPC context.
+
+        Returns:
+            A response with the module's setup information.
+        """
+        logger.debug("GetModuleSetup called for module: '%s'", self.module_class.__name__)
+
+        # Get setup schema if available
+        try:
+            # Convert schema to proto format
+            setup_schema_proto = await self.module_class.get_setup_format(llm_format=request.llm_format)
+            setup_format_struct = json_format.Parse(
+                text=setup_schema_proto,
+                message=struct_pb2.Struct(),  # pylint: disable=no-member
+                ignore_unknown_fields=True,
+            )
+        except NotImplementedError as e:
+            logger.warning(e)
+            context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+            context.set_details(str(e))
+            return information_pb2.GetModuleSetupResponse()
+
+        return information_pb2.GetModuleSetupResponse(
+            success=True,
+            setup_schema=setup_format_struct,
+        )
+
+    async def GetModuleSecret(  # noqa: N802
+        self,
+        request: information_pb2.GetModuleSecretRequest,
+        context: grpc.ServicerContext,
+    ) -> information_pb2.GetModuleSecretResponse:
+        """Get information about the module's secrets.
+
+        Args:
+            request: The get module secret request.
+            context: The gRPC context.
+
+        Returns:
+            A response with the module's secret schema.
+        """
+        logger.info("GetModuleSecret called for module: '%s'", self.module_class.__name__)
+
+        # Get secret schema if available
+        try:
+            # Convert schema to proto format
+            secret_schema_proto = await self.module_class.get_secret_format(llm_format=request.llm_format)
+            secret_format_struct = json_format.Parse(
+                text=secret_schema_proto,
+                message=struct_pb2.Struct(),  # pylint: disable=no-member
+                ignore_unknown_fields=True,
+            )
+        except NotImplementedError as e:
+            logger.warning(e)
+            context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+            context.set_details(str(e))
+            return information_pb2.GetModuleSecretResponse()
+
+        return information_pb2.GetModuleSecretResponse(
+            success=True,
+            secret_schema=secret_format_struct,
+        )
+
+    async def GetConfigSetupModule(  # noqa: N802
+        self,
+        request: information_pb2.GetConfigSetupModuleRequest,
+        context: grpc.ServicerContext,
+    ) -> information_pb2.GetConfigSetupModuleResponse:
+        """Get information about the module's setup and configuration.
+
+        Args:
+            request: The get module setup request.
+            context: The gRPC context.
+
+        Returns:
+            A response with the module's setup information.
+        """
+        logger.debug("GetConfigSetupModule called for module: '%s'", self.module_class.__name__)
+
+        # Get setup schema if available
+        try:
+            # Convert schema to proto format
+            config_setup_schema_proto = await self.module_class.get_config_setup_format(llm_format=request.llm_format)
+            config_setup_format_struct = json_format.Parse(
+                text=config_setup_schema_proto,
+                message=struct_pb2.Struct(),  # pylint: disable=no-member
+                ignore_unknown_fields=True,
+            )
+        except NotImplementedError as e:
+            logger.warning(e)
+            context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+            context.set_details(str(e))
+            return information_pb2.GetConfigSetupModuleResponse()
+
+        return information_pb2.GetConfigSetupModuleResponse(
+            success=True,
+            config_setup_schema=config_setup_format_struct,
+        )

digitalkin/grpc_servers/utils/__init__.py

@@ -0,0 +1 @@
+"""gRPC servers utilities package."""

digitalkin/grpc_servers/utils/exceptions.py

@@ -0,0 +1,29 @@
+"""Exceptions for the DigitalKin gRPC package."""
+
+
+class DigitalKinError(Exception):
+    """Base exception for all DigitalKin errors."""
+
+
+class ServerError(DigitalKinError):
+    """Base class for server-related errors."""
+
+
+class ConfigurationError(ServerError):
+    """Error related to server configuration."""
+
+
+class ServicerError(ServerError):
+    """Error related to servicer operations."""
+
+
+class SecurityError(ServerError):
+    """Error related to security configuration."""
+
+
+class ServerStateError(ServerError):
+    """Error related to server state (e.g., already started, not started)."""
+
+
+class ReflectionError(ServerError):
+    """Error related to gRPC reflection service."""

digitalkin/grpc_servers/utils/grpc_client_wrapper.py

@@ -0,0 +1,88 @@
+"""Client wrapper to ease channel creation with specific ServerConfig."""
+
+from pathlib import Path
+from typing import Any
+
+import grpc
+
+from digitalkin.grpc_servers.utils.exceptions import ServerError
+from digitalkin.logger import logger
+from digitalkin.models.grpc_servers.models import ClientConfig, SecurityMode
+
+
+class GrpcClientWrapper:
+    """gRPC client shared by the different services."""
+
+    stub: Any
+
+    @staticmethod
+    def _init_channel(config: ClientConfig) -> grpc.Channel:
+        """Create an appropriate channel to the registry server.
+
+        Returns:
+            A gRPC channel for communication with the registry.
+
+        Raises:
+            ValueError: If credentials are required but not provided.
+        """
+        if config.security == SecurityMode.SECURE and config.credentials is not None:
+            # Secure channel
+            root_certificates = Path(config.credentials.root_cert_path).read_bytes()
+
+            # mTLS channel
+            private_key = None
+            certificate_chain = None
+            if config.credentials.client_cert_path is not None and config.credentials.client_key_path is not None:
+                private_key = Path(config.credentials.client_key_path).read_bytes()
+                certificate_chain = Path(config.credentials.client_cert_path).read_bytes()
+
+            # Create channel credentials
+            channel_credentials = grpc.ssl_channel_credentials(
+                root_certificates=root_certificates,
+                certificate_chain=certificate_chain,
+                private_key=private_key,
+            )
+
+            return grpc.secure_channel(config.address, channel_credentials, options=config.channel_options)
+        # Insecure channel
+        return grpc.insecure_channel(config.address, options=config.channel_options)
+
+    def exec_grpc_query(self, query_endpoint: str, request: Any) -> Any:  # noqa: ANN401
+        """Execute a gRPC query with from the query's rpc endpoint name.
+
+        Arguments:
+            query_endpoint: rpc query name
+            request: gRPC object to match the rpc query
+
+        Returns:
+            corresponding gRPC reponse.
+
+        Raises:
+            ServerError: gRPC error catching with status code and details
+        """
+        service_name = getattr(self, "service_name", "unknown")
+        try:
+            logger.debug(
+                "Sending gRPC request to %s",
+                query_endpoint,
+                extra={"request": str(request), "service": service_name},
+            )
+            response = getattr(self.stub, query_endpoint)(request)
+            logger.debug(
+                "Received gRPC response from %s",
+                query_endpoint,
+                extra={"response": str(response), "service": service_name},
+            )
+        except grpc.RpcError as e:
+            status_code = e.code().name if hasattr(e, "code") else "UNKNOWN"
+            details = e.details() if hasattr(e, "details") else str(e)
+            msg = f"[{status_code}] {details}"
+            logger.error(
+                "gRPC %s failed: %s",
+                query_endpoint,
+                msg,
+                extra={"service": service_name},
+            )
+            raise ServerError(msg) from e
+        else:
+            return response

digitalkin/grpc_servers/utils/grpc_error_handler.py

@@ -0,0 +1,53 @@
+"""Shared error handling utilities for gRPC services."""
+
+from collections.abc import Generator
+from contextlib import contextmanager
+from typing import Any
+
+from digitalkin.grpc_servers.utils.exceptions import ServerError
+from digitalkin.logger import logger
+
+
+class GrpcErrorHandlerMixin:
+    """Mixin class providing common gRPC error handling functionality."""
+
+    @contextmanager
+    def handle_grpc_errors(  # noqa: PLR6301
+        self,
+        operation: str,
+        service_error_class: type[Exception] | None = None,
+    ) -> Generator[Any, Any, Any]:
+        """Handle gRPC errors for the given operation.
+
+        Args:
+            operation: Name of the operation being performed.
+            service_error_class: Optional specific service exception class to raise.
+                                If not provided, uses the generic ServerError.
+
+        Yields:
+            Context for the operation.
+
+        Raises:
+            ServerError: For gRPC-related errors.
+            service_error_class: For service-specific errors if provided.
+        """
+        if service_error_class is None:
+            service_error_class = ServerError
+
+        try:
+            yield
+        except service_error_class as e:
+            # Re-raise service-specific errors as-is
+            msg = f"{service_error_class.__name__} in {operation}: {e}"
+            logger.exception(msg)
+            raise service_error_class(msg) from e
+        except ServerError as e:
+            # Handle gRPC server errors
+            msg = f"gRPC {operation} failed: {e}"
+            logger.exception(msg)
+            raise ServerError(msg) from e
+        except Exception as e:
+            # Handle unexpected errors
+            msg = f"Unexpected error in {operation}: {e}"
+            logger.exception(msg)
+            raise service_error_class(msg) from e