digitalkin 0.3.2.dev2__py3-none-any.whl

digitalkin/grpc_servers/_base_server.py

@@ -0,0 +1,486 @@
+"""Base gRPC server implementation for DigitalKin."""
+
+import abc
+import asyncio
+from collections.abc import Callable
+from concurrent import futures
+from pathlib import Path
+from typing import Any, cast
+
+import grpc
+from grpc import aio as grpc_aio
+
+from digitalkin.grpc_servers.utils.exceptions import (
+    ConfigurationError,
+    ReflectionError,
+    SecurityError,
+    ServerStateError,
+    ServicerError,
+)
+from digitalkin.logger import logger
+from digitalkin.models.grpc_servers.models import SecurityMode, ServerConfig, ServerMode
+from digitalkin.models.grpc_servers.types import GrpcServer, ServiceDescriptor, T
+
+
+class BaseServer(abc.ABC):
+    """Base class for gRPC servers in DigitalKin.
+
+    This class provides the foundation for both synchronous and asynchronous gRPC
+    servers used in the DigitalKin ecosystem. It supports both secure and insecure
+    communication modes.
+
+    Attributes:
+        config: The server configuration.
+        server: The gRPC server instance (either sync or async).
+        _servicers: List of registered servicers.
+        _service_names: List of service names for reflection.
+        _health_servicer: Optional health check servicer.
+    """
+
+    def __init__(self, config: ServerConfig) -> None:
+        """Initialize the base gRPC server.
+
+        Args:
+            config: The server configuration.
+        """
+        self.config = config
+        self.server: GrpcServer | None = None
+        self._servicers: list[Any] = []
+        self._service_names: list[str] = []  # Track service names for reflection
+        self._health_servicer: Any = None  # For health checking
+
+    def register_servicer(
+        self,
+        servicer: T,
+        add_to_server_fn: Callable[[T, GrpcServer], None],
+        service_descriptor: ServiceDescriptor | None = None,
+        service_names: list[str] | None = None,
+    ) -> None:
+        """Register a servicer with the gRPC server and track it for reflection.
+
+        Args:
+            servicer: The servicer implementation instance
+            add_to_server_fn: The function to add the servicer to the server
+            service_descriptor: Optional service descriptor (pb2 DESCRIPTOR)
+            service_names: Optional explicit list of service full names
+
+        Raises:
+            ServicerError: If the server is not created before calling
+        """
+        if self.server is None:
+            msg = "Server must be created before registering servicers"
+            raise ServicerError(msg)
+
+        # Register the servicer
+        try:
+            add_to_server_fn(servicer, self.server)
+            self._servicers.append(servicer)
+        except Exception as e:
+            msg = f"Failed to register servicer: {e}"
+            raise ServicerError(msg) from e
+
+        # Add service names from explicit list
+        if service_names:
+            for name in service_names:
+                if name not in self._service_names:
+                    self._service_names.append(name)
+                    logger.debug("Registered explicit service name for reflection: %s", name)
+
+        # If a descriptor is provided, extract service names
+        if service_descriptor and hasattr(service_descriptor, "services_by_name"):
+            for service_name in service_descriptor.services_by_name:
+                full_name = service_descriptor.services_by_name[service_name].full_name  # ignore: PLC0206
+                if full_name not in self._service_names:
+                    self._service_names.append(full_name)
+                    logger.debug("Registered service name from descriptor: %s", full_name)
+
+    @abc.abstractmethod
+    def _register_servicers(self) -> None:
+        """Register servicers with the gRPC server.
+
+        This method should be implemented by subclasses to register
+        the appropriate servicers for their specific functionality.
+
+        Raises:
+            ServicerError: If the server is not created before calling this method.
+        """
+
+    def _add_reflection(self) -> None:
+        """Add reflection service to the gRPC server if enabled.
+
+        Raises:
+            ReflectionError: If reflection initialization fails.
+        """
+        if not self.config.enable_reflection or self.server is None or not self._service_names:
+            return
+
+        try:
+            # Import here to avoid dependency if not used
+            from grpc_reflection.v1alpha import reflection  # noqa: PLC0415
+
+            # Get all registered service names
+            service_names = self._service_names.copy()
+
+            # Add the reflection service name
+            reflection_service = reflection.SERVICE_NAME
+            service_names.append(reflection_service)
+
+            # Register services with the reflection service
+            # This creates a dynamic file descriptor database that can respond to
+            # reflection queries with detailed service information
+            reflection.enable_server_reflection(service_names, self.server)
+
+            logger.debug("Added gRPC reflection service with services: %s", service_names)
+        except ImportError:
+            logger.warning("Could not enable reflection: grpcio-reflection package not installed")
+        except Exception as e:
+            error_msg = f"Failed to enable reflection: {e}"
+            logger.warning(error_msg)
+            raise ReflectionError(error_msg) from e
+
+    def _add_health_service(self) -> None:
+        """Add health checking service to the gRPC server.
+
+        The health service allows clients to check server status.
+        """
+        if self.server is None:
+            return
+
+        try:
+            # Import here to avoid dependency if not used
+            from grpc_health.v1 import health_pb2, health_pb2_grpc  # noqa: PLC0415
+            from grpc_health.v1.health import HealthServicer  # noqa: PLC0415
+
+            # Create health servicer
+            health_servicer = HealthServicer()
+
+            # Register health servicer
+            health_pb2_grpc.add_HealthServicer_to_server(health_servicer, self.server)
+
+            # Add service name to reflection list
+            if health_pb2.DESCRIPTOR.services_by_name:
+                service_name = health_pb2.DESCRIPTOR.services_by_name["Health"].full_name
+                if service_name not in self._service_names:
+                    self._service_names.append(service_name)
+
+            logger.debug("Added gRPC health checking service")
+
+            # Set all services as SERVING
+            for service_name in self._service_names:
+                health_servicer.set(service_name, health_pb2.HealthCheckResponse.SERVING)
+
+            # Set overall service status
+            health_servicer.set("", health_pb2.HealthCheckResponse.SERVING)
+
+            # Store reference to health servicer
+            self._health_servicer = health_servicer
+
+        except ImportError:
+            logger.warning("Could not enable health service: grpcio-health-checking package not installed")
+        except Exception as e:
+            logger.warning("Failed to enable health service: %s", e)
+
+    def _create_server(self) -> GrpcServer:
+        """Create a gRPC server instance based on the configuration.
+
+        Returns:
+            A configured gRPC server instance.
+
+        Raises:
+            ConfigurationError: If the server configuration is invalid.
+        """
+        try:
+            # Create the server based on mode
+            if self.config.mode == ServerMode.ASYNC:
+                server = grpc_aio.server(options=self.config.server_options)
+            else:
+                server = grpc.server(
+                    futures.ThreadPoolExecutor(max_workers=self.config.max_workers),
+                    options=self.config.server_options,
+                )
+
+            # Add the appropriate port
+            if self.config.security == SecurityMode.SECURE:
+                self._add_secure_port(server)
+            else:
+                self._add_insecure_port(server)
+
+        except Exception as e:
+            msg = f"Failed to create server: {e}"
+            raise ConfigurationError(msg) from e
+        else:
+            return server
+
+    def _add_secure_port(self, server: GrpcServer) -> None:
+        """Add a secure port to the server.
+
+        Args:
+            server: The gRPC server to add the port to.
+
+        Raises:
+            SecurityError: If credentials are not configured correctly.
+        """
+        if not self.config.credentials:
+            msg = "Credentials must be provided for secure server"
+            raise SecurityError(msg)
+
+        try:
+            # Read key and certificate files
+            private_key = Path(self.config.credentials.server_key_path).read_bytes()
+            certificate_chain = Path(self.config.credentials.server_cert_path).read_bytes()
+
+            # Read root certificate if provided
+            root_certificates = None
+            if self.config.credentials.root_cert_path:
+                root_certificates = Path(self.config.credentials.root_cert_path).read_bytes()
+        except OSError as e:
+            msg = f"Failed to read credential files: {e}"
+            raise SecurityError(msg) from e
+
+        try:
+            # Create server credentials
+            server_credentials = grpc.ssl_server_credentials(
+                [(private_key, certificate_chain)],
+                root_certificates=root_certificates,
+                require_client_auth=(root_certificates is not None),
+            )
+
+            # Add secure port to server
+            if self.config.mode == ServerMode.ASYNC:
+                async_server = cast("grpc_aio.Server", server)
+                async_server.add_secure_port(self.config.address, server_credentials)
+            else:
+                sync_server = cast("grpc.Server", server)
+                sync_server.add_secure_port(self.config.address, server_credentials)
+
+            logger.debug("Added secure port %s", self.config.address)
+        except Exception as e:
+            msg = f"Failed to configure secure port: {e}"
+            raise SecurityError(msg) from e
+
+    def _add_insecure_port(self, server: GrpcServer) -> None:
+        """Add an insecure port to the server.
+
+        Args:
+            server: The gRPC server to add the port to.
+
+        Raises:
+            ConfigurationError: If adding the insecure port fails.
+        """
+        try:
+            if self.config.mode == ServerMode.ASYNC:
+                async_server = cast("grpc_aio.Server", server)
+                async_server.add_insecure_port(self.config.address)
+            else:
+                sync_server = cast("grpc.Server", server)
+                sync_server.add_insecure_port(self.config.address)
+
+            logger.debug("Added insecure port %s", self.config.address)
+        except Exception as e:
+            msg = f"Failed to add insecure port: {e}"
+            raise ConfigurationError(msg) from e
+
+    def start(self) -> None:
+        """Start the gRPC server.
+
+        If using async mode, this will use the event loop to start the server.
+        If using sync mode, this will start the server in a non-blocking way.
+
+        Raises:
+            ServerStateError: If the server fails to start.
+        """
+        self.server = self._create_server()
+        self._register_servicers()
+
+        # Add health service
+        self._add_health_service()
+
+        # Add reflection if enabled
+        self._add_reflection()
+
+        # Start the server
+        logger.debug("Starting gRPC server on %s", self.config.address, extra={"config": self.config})
+        try:
+            if self.config.mode == ServerMode.ASYNC:
+                # For async server, use the event loop
+                loop = asyncio.get_event_loop()
+                if loop.is_closed():
+                    loop = asyncio.new_event_loop()
+                    asyncio.set_event_loop(loop)
+                loop.run_until_complete(self._start_async())
+            else:
+                # For sync server, directly call start
+                sync_server = cast("grpc.Server", self.server)
+                sync_server.start()
+            logger.debug("✅ gRPC server started on %s", self.config.address)
+        except Exception as e:
+            logger.exception("❎ Error starting server")
+            msg = f"Failed to start server: {e}"
+            raise ServerStateError(msg) from e
+
+    async def _start_async(self) -> None:
+        """Start the async gRPC server.
+
+        Raises:
+            ServerStateError: If the server is not created.
+        """
+        if self.server is None:
+            msg = "Server is not created"
+            raise ServerStateError(msg)
+
+        async_server = cast("grpc_aio.Server", self.server)
+        await async_server.start()
+
+    async def start_async(self) -> None:
+        """Start the gRPC server asynchronously.
+
+        This method should be used directly in an async context.
+
+        Raises:
+            ServerStateError: If the server fails to start.
+        """
+        self.server = self._create_server()
+        self._register_servicers()
+
+        # Add health service
+        self._add_health_service()
+
+        # Add reflection if enabled
+        self._add_reflection()
+
+        # Start the server
+        logger.debug("Starting gRPC server on %s", self.config.address)
+        try:
+            if self.config.mode == ServerMode.ASYNC:
+                await self._start_async()
+            else:
+                # For sync server in async context
+                sync_server = cast("grpc.Server", self.server)
+                sync_server.start()
+            logger.debug("✅ gRPC server started on %s", self.config.address)
+        except Exception as e:
+            logger.exception("❎ Error starting server")
+            msg = f"Failed to start server: {e}"
+            raise ServerStateError(msg) from e
+
+    def stop(self, grace: float | None = None) -> None:
+        """Stop the gRPC server.
+
+        Args:
+            grace: Optional grace period in seconds for existing RPCs to complete.
+        """
+        if self.server is None:
+            logger.warning("Attempted to stop server, but no server is running")
+            return
+
+        logger.debug("Stopping gRPC server...")
+        if self.config.mode == ServerMode.ASYNC:
+            # We'll use a different approach that works whether we're in a running event loop or not
+            try:
+                # Get the current event loop
+                loop = asyncio.get_event_loop()
+
+                if loop.is_running():
+                    # If we're in a running event loop, we can't run_until_complete
+                    # Just warn the user they should use stop_async
+                    logger.warning(
+                        "Called stop() on async server from a running event loop. "
+                        "This might not fully shut down the server. "
+                        "Use await stop_async() in async contexts instead."
+                    )
+                    # Set server to None to avoid further operations
+                    self.server = None
+                    logger.debug("✅ gRPC server marked as stopped")
+                    return
+                # If not in a running event loop, use run_until_complete
+                loop.run_until_complete(self._stop_async(grace))
+            except RuntimeError:
+                # Event loop issues - try with a new loop
+                logger.debug("Creating new event loop for shutdown")
+                try:
+                    new_loop = asyncio.new_event_loop()
+                    asyncio.set_event_loop(new_loop)
+                    new_loop.run_until_complete(self._stop_async(grace))
+                finally:
+                    new_loop.close()
+        else:
+            # For sync server, we can just call stop
+            sync_server = cast("grpc.Server", self.server)
+            sync_server.stop(grace=grace)
+
+        logger.debug("✅ gRPC server stopped")
+        self.server = None
+
+    async def _stop_async(self, grace: float | None = None) -> None:
+        """Stop the async gRPC server.
+
+        Args:
+            grace: Optional grace period in seconds for existing RPCs to complete.
+        """
+        if self.server is None:
+            return
+
+        async_server = cast("grpc_aio.Server", self.server)
+        await async_server.stop(grace=grace)
+
+    async def stop_async(self, grace: float | None = None) -> None:
+        """Stop the gRPC server asynchronously.
+
+        This method should be used in async contexts.
+
+        Args:
+            grace: Optional grace period in seconds for existing RPCs to complete.
+        """
+        if self.server is None:
+            logger.warning("Attempted to stop server, but no server is running")
+            return
+
+        logger.debug("Stopping gRPC server asynchronously...")
+        if self.config.mode == ServerMode.ASYNC:
+            await self._stop_async(grace)
+        else:
+            # For sync server, we can just call stop
+            sync_server = cast("grpc.Server", self.server)
+            sync_server.stop(grace=grace)
+
+        logger.debug("✅ gRPC server stopped")
+        self.server = None
+
+    def wait_for_termination(self) -> None:
+        """Wait for the server to terminate.
+
+        In synchronous mode, this blocks until the server is terminated.
+        In asynchronous mode, a warning is logged suggesting to use `await_termination`.
+        """
+        if self.server is None:
+            logger.warning("Attempted to wait for termination, but no server is running")
+            return
+
+        if self.config.mode == ServerMode.SYNC:
+            # For sync server
+            sync_server = cast("grpc.Server", self.server)
+            sync_server.wait_for_termination()
+        else:
+            # For async server, the caller should use await_termination instead
+            logger.warning(
+                "Called wait_for_termination on async server. Use await_termination instead for async servers.",
+            )
+
+    async def await_termination(self) -> None:
+        """Wait for the async server to terminate.
+
+        This method should only be used with async servers.
+        """
+        if self.config.mode == ServerMode.SYNC:
+            logger.warning(
+                "Called await_termination on sync server. Use wait_for_termination instead for sync servers.",
+            )
+            return
+
+        if self.server is None:
+            logger.warning("Attempted to await termination, but no server is running")
+            return
+
+        # For async server
+        async_server = cast("grpc_aio.Server", self.server)
+        await async_server.wait_for_termination()

digitalkin/grpc_servers/module_server.py

@@ -0,0 +1,208 @@
+"""Module gRPC server implementation for DigitalKin."""
+
+from typing import TYPE_CHECKING
+
+from agentic_mesh_protocol.module.v1 import (
+    module_service_pb2,
+    module_service_pb2_grpc,
+)
+
+from digitalkin.grpc_servers._base_server import BaseServer
+from digitalkin.grpc_servers.module_servicer import ModuleServicer
+from digitalkin.logger import logger
+from digitalkin.models.grpc_servers.models import (
+    ClientConfig,
+    ModuleServerConfig,
+    SecurityMode,
+)
+from digitalkin.modules._base_module import BaseModule
+from digitalkin.services.registry import GrpcRegistry
+
+if TYPE_CHECKING:
+    from digitalkin.services.registry import RegistryStrategy
+
+
+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 Registry server.
+
+    Attributes:
+        module: The module instance being served.
+        server_config: Server configuration.
+        client_config: Setup client configuration.
+        module_servicer: The gRPC servicer handling module requests.
+    """
+
+    def __init__(
+        self,
+        module_class: type[BaseModule],
+        server_config: ModuleServerConfig,
+        client_config: ClientConfig | None = None,
+    ) -> None:
+        """Initialize the module server.
+
+        Args:
+            module_class: The module instance to be served.
+            server_config: Server configuration including registry address if auto-registration is desired.
+            client_config: Client configuration used by services.
+        """
+        super().__init__(server_config)
+        self.module_class = module_class
+        self.server_config = server_config
+        self.client_config = client_config
+        self.module_servicer: ModuleServicer | None = None
+        self.registry: RegistryStrategy | None = None
+
+        self._registry_client_config: ClientConfig | None = None
+        if self.server_config.registry_address:
+            self._registry_client_config = self._build_registry_client_config()
+        self._prepare_registry_config()
+
+    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.debug("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.debug("Registered Module servicer")
+
+    def _prepare_registry_config(self) -> None:
+        """Prepare registry client config on module_class before server starts.
+
+        This ensures ServicesConfig created by JobManager will have registry config,
+        allowing spawned module instances to inherit the registry configuration.
+        """
+        if not self._registry_client_config:
+            return
+
+        self.module_class.services_config_params["registry"] = {"client_config": self._registry_client_config}
+
+    def _build_registry_client_config(self) -> ClientConfig:
+        """Build ClientConfig for registry from server_config.registry_address.
+
+        Returns:
+            ClientConfig configured for registry connection.
+        """
+        host, port = self.server_config.registry_address.rsplit(":", 1)
+        return ClientConfig(
+            host=host,
+            port=int(port),
+            mode=self.server_config.mode,
+            security=self.client_config.security if self.client_config else SecurityMode.INSECURE,
+            credentials=self.client_config.credentials if self.client_config else None,
+        )
+
+    def _init_registry(self) -> None:
+        """Initialize server-level registry client for registration.
+
+        Note: services_config_params["registry"] is already set in _prepare_registry_config()
+        which runs in __init__(). This method only creates the server-level client instance.
+        """
+        if not self._registry_client_config:
+            return
+
+        self.registry = GrpcRegistry("", "", "", self._registry_client_config)
+
+    def start(self) -> None:
+        """Start the module server and register with the registry if configured."""
+        logger.info("Starting module server", extra={"server_config": self.server_config})
+        super().start()
+
+        try:
+            self._init_registry()
+            self._register_with_registry()
+        except Exception:
+            logger.exception("Failed to register with registry")
+
+        if self.module_servicer is not None:
+            logger.debug("Setup post init started", extra={"client_config": self.client_config})
+            self.module_servicer.setup.__post_init__(self.client_config)
+
+    async def start_async(self) -> None:
+        """Start the module server and register with the registry if configured."""
+        logger.info("Starting module server", extra={"server_config": self.server_config})
+        await super().start_async()
+
+        try:
+            self._init_registry()
+            self._register_with_registry()
+        except Exception:
+            logger.exception("Failed to register with registry")
+
+        if self.module_servicer is not None:
+            logger.info("Setup post init started", extra={"client_config": self.client_config})
+            await self.module_servicer.job_manager.start()
+            self.module_servicer.setup.__post_init__(self.client_config)
+
+    def stop(self, grace: float | None = None) -> None:
+        """Stop the module server.
+
+        Modules become inactive when they stop sending heartbeats
+        """
+        super().stop(grace)
+
+    def _register_with_registry(self) -> None:
+        """Register this module with the registry server."""
+        if not self.registry:
+            logger.debug("No registry configured, skipping registration")
+            return
+
+        module_id = self.module_class.get_module_id()
+        version = self.module_class.metadata.get("version", "0.0.0")
+
+        if not module_id or module_id == "unknown":
+            logger.warning(
+                "Module has no valid module_id, skipping registration",
+                extra={"module_class": self.module_class.__name__},
+            )
+            return
+
+        logger.info(
+            "Attempting to register module with registry",
+            extra={
+                "module_id": module_id,
+                "address": self.server_config.address,
+                "port": self.server_config.port,
+                "version": version,
+                "registry_address": self.server_config.registry_address,
+            },
+        )
+
+        result = self.registry.register(
+            module_id=module_id,
+            address=self.server_config.address,
+            port=self.server_config.port,
+            version=version,
+        )
+
+        if result:
+            logger.info(
+                "Module registered successfully",
+                extra={
+                    "module_id": result.module_id,
+                    "address": self.server_config.address,
+                    "port": self.server_config.port,
+                    "registry_address": self.server_config.registry_address,
+                },
+            )
+        else:
+            logger.warning(
+                "Module registration returned None (module may not exist in registry)",
+                extra={
+                    "module_id": module_id,
+                    "registry_address": self.server_config.registry_address,
+                },
+            )