digitalkin 0.1.1__py3-none-any.whl

digitalkin/grpc/module_server.py

@@ -0,0 +1,233 @@
+"""Module gRPC server implementation for DigitalKin."""
+
+import logging
+import uuid
+
+import grpc
+
+from digitalkin.grpc._base_server import BaseServer
+from digitalkin.grpc.module_servicer import ModuleServicer
+from digitalkin.grpc.utils.exceptions import ServerError
+from digitalkin.grpc.utils.models import ModuleServerConfig, SecurityMode
+from digitalkin.modules._base_module import BaseModule
+
+from digitalkin_proto.digitalkin.module.v2 import module_service_pb2, module_service_pb2_grpc
+from digitalkin_proto.digitalkin.module_registry.v2 import (
+    metadata_pb2,
+    module_registry_service_pb2_grpc,
+    registration_pb2,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class ModuleServer(BaseServer):
+    """gRPC server for a DigitalKin module.
+
+    This server exposes the module's functionality through the ModuleService gRPC interface.
+    It can optionally register itself with a ModuleRegistry server.
+
+    Attributes:
+        module: The module instance being served.
+        config: Server configuration.
+        module_servicer: The gRPC servicer handling module requests.
+    """
+
+    def __init__(
+        self,
+        module_class: type[BaseModule],
+        config: ModuleServerConfig,
+    ) -> None:
+        """Initialize the module server.
+
+        Args:
+            module_class: The module instance to be served.
+            config: Server configuration including registry address if auto-registration is desired.
+        """
+        super().__init__(config)
+        self.module_class = module_class
+        self.config = config
+        self.module_servicer: ModuleServicer | None = None
+
+    def _register_servicers(self) -> None:
+        """Register the module servicer with the gRPC server.
+
+        Raises:
+            RuntimeError: No registered server
+        """
+        if self.server is None:
+            msg = "Server must be created before registering servicers"
+            raise RuntimeError(msg)
+
+        logger.info("Registering module servicer for %s", self.module_class.__name__)
+        self.module_servicer = ModuleServicer(self.module_class)
+        self.register_servicer(
+            self.module_servicer,
+            module_service_pb2_grpc.add_ModuleServiceServicer_to_server,
+            service_descriptor=module_service_pb2.DESCRIPTOR,
+        )
+        logger.info("Registered Async Greeter servicer")
+
+    def start(self) -> None:
+        """Start the module server and register with the registry if configured."""
+        logger.critical(self.config)
+        super().start()
+
+        logger.critical(self.config)
+        # If a registry address is provided, register the module
+        if self.config.registry_address:
+            try:
+                self._register_with_registry()
+            except Exception:
+                logger.exception("Failed to register with registry")
+
+    async def start_async(self) -> None:
+        """Start the module server and register with the registry if configured."""
+        logger.critical(self.config)
+        await super().start_async()
+
+        logger.critical(self.config)
+        # If a registry address is provided, register the module
+        if self.config.registry_address:
+            try:
+                self._register_with_registry()
+            except Exception:
+                logger.exception("Failed to register with registry")
+
+    def stop(self, grace: float | None = None) -> None:
+        """Stop the module server and deregister from the registry if needed."""
+        # If registered with a registry, deregister
+        if self.config.registry_address:
+            try:
+                self._deregister_from_registry()
+            except ServerError:
+                logger.exception("Failed to deregister from registry")
+
+        super().stop(grace)
+
+    def _register_with_registry(self) -> None:
+        """Register this module with the registry server.
+
+        Raises:
+            ServerError: If communication with the registry server fails.
+        """
+        logger.info("Registering module with registry at %s", self.config.registry_address)
+
+        # Create appropriate channel based on security mode
+        channel = self._create_registry_channel()
+
+        with channel:
+            # Create a stub (client)
+            stub = module_registry_service_pb2_grpc.ModuleRegistryServiceStub(channel)
+
+            # Determine module type
+            module_type = self._determine_module_type()
+
+            metadata = metadata_pb2.Metadata(
+                name=self.module_class.metadata["name"],
+                tags=[metadata_pb2.Tag(tag=tag) for tag in self.module_class.metadata["tags"]],
+                description=self.module_class.metadata["description"],
+            )
+
+            self.module_class.metadata["module_id"] = f"{self.module_class.metadata['name']}:{uuid.uuid4()}"
+            # Create registration request
+            request = registration_pb2.RegisterRequest(
+                module_id=self.module_class.metadata["module_id"],
+                version=self.module_class.metadata["version"],
+                module_type=module_type,
+                address=self.config.address,
+                metadata=metadata,
+            )
+
+            try:
+                # Call the register method
+                logger.info(
+                    "Request sent to registry for module: %s:%s",
+                    self.module_class.metadata["name"],
+                    self.module_class.metadata["module_id"],
+                )
+                response = stub.RegisterModule(request)
+
+                if response.success:
+                    logger.info("Module registered successfully")
+                else:
+                    logger.error("Module registration failed")
+            except grpc.RpcError:
+                logger.exception("RPC error during registration:")
+                raise ServerError
+
+    def _deregister_from_registry(self) -> None:
+        """Deregister this module from the registry server.
+
+        Raises:
+            ServerError: If communication with the registry server fails.
+        """
+        logger.info("Deregistering module from registry at %s", self.config.registry_address)
+
+        # Create appropriate channel based on security mode
+        channel = self._create_registry_channel()
+
+        with channel:
+            # Create a stub (client)
+            stub = module_registry_service_pb2_grpc.ModuleRegistryServiceStub(channel)
+
+            # Create deregistration request
+            request = registration_pb2.DeregisterRequest(
+                module_id=self.module_class.metadata["module_id"],
+            )
+        try:
+            # Call the deregister method
+            response = stub.DeregisterModule(request)
+
+            if response.success:
+                logger.info("Module deregistered successfull")
+            else:
+                logger.error("Module deregistration failed")
+        except grpc.RpcError:
+            logger.exception("RPC error during deregistration")
+            raise ServerError
+
+    def _create_registry_channel(self) -> 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 self.config.security == SecurityMode.SECURE and self.config.credentials:
+            # Secure channel
+            """TODO: use Path(self.config.credentials.server_cert_path).read_bytes()"""
+            with open(self.config.credentials.server_cert_path, "rb") as cert_file:  # noqa: FURB101
+                certificate_chain = cert_file.read()
+
+            root_certificates = None
+            if self.config.credentials.root_cert_path:
+                with open(self.config.credentials.root_cert_path, "rb") as root_cert_file:  # noqa: FURB101
+                    root_certificates = root_cert_file.read()
+
+            # Create channel credentials
+            channel_credentials = grpc.ssl_channel_credentials(root_certificates=root_certificates or certificate_chain)
+
+            return grpc.secure_channel(self.config.registry_address, channel_credentials)
+        # Insecure channel
+        return grpc.insecure_channel(self.config.registry_address)
+
+    def _determine_module_type(self) -> str:
+        """Determine the module type based on its class.
+
+        Returns:
+            A string representing the module type.
+        """
+        module_type = "UNKNOWN"
+        class_name = self.module_class.__name__
+
+        if class_name == "ToolModule":
+            module_type = "TOOL"
+        elif class_name == "TriggerModule":
+            module_type = "TRIGGER"
+        elif class_name == "ArchetypeModule":
+            module_type = "KIN"
+
+        return module_type

digitalkin/grpc/module_servicer.py

@@ -0,0 +1,304 @@
+"""Module servicer implementation for DigitalKin."""
+
+import asyncio
+import logging
+from collections.abc import AsyncGenerator
+from typing import Any
+
+import grpc
+from digitalkin_proto.digitalkin.module.v2 import (
+    information_pb2,
+    lifecycle_pb2,
+    module_service_pb2_grpc,
+    monitoring_pb2,
+)
+from google.protobuf import json_format, struct_pb2
+
+from digitalkin.models.module.module import ModuleStatus
+from digitalkin.modules._base_module import BaseModule
+from digitalkin.modules.job_manager import JobManager
+
+logger = logging.getLogger(__name__)
+
+
+class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer):
+    """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.
+    """
+
+    def __init__(self, module_class: type[BaseModule]) -> None:
+        """Initialize the module servicer.
+
+        Args:
+            module_class: The module type to serve.
+        """
+        super().__init__()
+        self.queue: asyncio.Queue = asyncio.Queue()
+        self.module_class = module_class
+        self.job_manager = JobManager(module_class)
+
+    async def add_to_queue(self, job_id: str, output_data: dict[str, Any]) -> None:
+        """Callback used to add the output data to the queue of messages."""
+        logger.info("JOB: %s added an output_data: %s", job_id, output_data)
+        await self.queue.put({job_id: output_data})
+
+    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.
+        """
+        logger.info("StartModule called for module: '%s'", self.module_class.__name__)
+        # Process the module input
+        input_data = dict(request.input.items())
+        setup_data = self.module_class.storage.get(
+            table="setups",
+            data={
+                "keys": [
+                    # remove prefix 'setups:'
+                    request.setup_id[7:],
+                ],
+            },
+        )[0]
+
+        # setup_id should be use to request a precise setup from the module
+        # Create a job for this execution
+        job_id, module = await self.job_manager.create_job(
+            input_data,
+            setup_data,
+            callback=self.add_to_queue,
+        )
+
+        while module.status == ModuleStatus.RUNNING or not self.queue.empty():
+            output_data = await self.queue.get()
+            if output_data.get("error", None) is not None:
+                context.set_code(output_data["error"]["code"])
+                context.set_details(output_data["error"]["error_message"])
+                yield lifecycle_pb2.StartModuleResponse(success=False)
+                return
+            else:
+                # TODO: add a check for the job_id / error handling
+                output_proto = {key: str(value) for key, value in output_data[job_id].items()}
+                yield lifecycle_pb2.StartModuleResponse(
+                    success=True,
+                    output=output_proto,
+                    job_id=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.info("StopModule called for module: '%s'", self.module_class.__name__)
+
+        job_id = request.job_id
+        if job_id not in self.job_manager.modules:
+            message = f"Job {job_id} not found"
+            logger.warning(message)
+            context.set_code(grpc.StatusCode.NOT_FOUND)
+            context.set_details(message)
+            return lifecycle_pb2.StopModuleResponse(success=False)
+
+        # Update the job status
+        await self.job_manager.modules[job_id].stop()
+
+        logger.info("Job %s stopped successfully", job_id)
+        return lifecycle_pb2.StopModuleResponse(success=True)
+
+    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.info("GetModuleStatus called for module: '%s'", self.module_class.__name__)
+
+        # If job_id is specified, get status for that job
+        if request.job_id:
+            if request.job_id not in self.job_manager.modules:
+                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()
+
+            status = self.job_manager.modules[request.job_id].status
+            logger.info("Job %s status: '%s'", request.job_id, status)
+            return monitoring_pb2.GetModuleStatusResponse(
+                success=True,
+                status=status.name,
+                job_id=request.job_id,
+            )
+
+        logger.info("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,
+        )
+
+    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.info("GetModuleJobs called for module: '%s'", self.module_class.__name__)
+
+        # 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 self.job_manager.modules.items()
+            ],
+        )
+
+    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.info("GetModuleInput called for module: '%s'", self.module_class.__name__)
+
+        # Get input schema if available
+        try:
+            # Convert schema to proto format
+            input_schema_proto = self.module_class.get_input_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(e)
+            return information_pb2.GetModuleInputResponse()
+
+        return information_pb2.GetModuleInputResponse(
+            success=True,
+            input_schema=input_format_struct,
+        )
+
+    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.info("GetModuleOutput called for module: '%s'", self.module_class.__name__)
+
+        # Get output schema if available
+        try:
+            # Convert schema to proto format
+            output_schema_proto = self.module_class.get_output_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(e)
+            return information_pb2.GetModuleOutputResponse()
+
+        return information_pb2.GetModuleOutputResponse(
+            success=True,
+            output_schema=output_format_struct,
+        )
+
+    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.info("GetModuleSetup called for module: '%s'", self.module_class.__name__)
+
+        # Get setup schema if available
+        try:
+            # Convert schema to proto format
+            setup_schema_proto = self.module_class.get_setup_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(e)
+            return information_pb2.GetModuleSetupResponse()
+
+        return information_pb2.GetModuleSetupResponse(
+            success=True,
+            setup_schema=setup_format_struct,
+        )

digitalkin/grpc/registry_server.py

@@ -0,0 +1,63 @@
+"""Registry gRPC server implementation for DigitalKin."""
+
+import logging
+
+from digitalkin_proto.digitalkin.module_registry.v2 import module_registry_service_pb2_grpc
+
+from digitalkin.grpc._base_server import BaseServer
+from digitalkin.grpc.registry_servicer import RegistryModule, RegistryServicer
+from digitalkin.grpc.utils.models import RegistryServerConfig
+
+logger = logging.getLogger(__name__)
+
+
+class RegistryServer(BaseServer):
+    """gRPC server for DigitalKin module registry.
+
+    This server implements the ModuleRegistryService which allows modules to register
+    themselves and be discovered by other components in the system.
+
+    Attributes:
+        config: Server configuration.
+        registry_servicer: The gRPC servicer handling registry requests.
+    """
+
+    def __init__(
+        self,
+        config: RegistryServerConfig,
+    ) -> None:
+        """Initialize the registry server.
+
+        Args:
+            config: Server configuration.
+        """
+        super().__init__(config)
+        self.config = config
+        self.registry_servicer: RegistryServicer | None = None
+
+    def _register_servicers(self) -> None:
+        """Register the registry servicer with the gRPC server.
+
+        Raises:
+            RuntimeError: If server is not registered during server creation
+        """
+        if self.server is None:
+            msg = "Server must be created before registering servicers"
+            raise RuntimeError(msg)
+
+        logger.info("Registering module registry servicer")
+        self.registry_servicer = RegistryServicer()
+        module_registry_service_pb2_grpc.add_ModuleRegistryServiceServicer_to_server(
+            self.registry_servicer, self.server
+        )
+        self._servicers.append(self.registry_servicer)
+
+    def get_registered_modules(self) -> list[RegistryModule]:
+        """Get a list of all registered modules.
+
+        Returns:
+            A list of module information objects.
+        """
+        if self.registry_servicer:
+            return list(self.registry_servicer.registered_modules.values())
+        return []