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

digitalkin/__version__.py

@@ -5,4 +5,4 @@ from importlib.metadata import PackageNotFoundError, version
 try:
     __version__ = version("digitalkin")
 except PackageNotFoundError:
-    __version__ = "0.2.12"
+    __version__ = "0.2.13"

digitalkin/grpc_servers/_base_server.py

@@ -2,7 +2,6 @@
 
 import abc
 import asyncio
-import logging
 from collections.abc import Callable
 from concurrent import futures
 from pathlib import Path
@@ -20,8 +19,7 @@ from digitalkin.grpc_servers.utils.exceptions import (
 )
 from digitalkin.grpc_servers.utils.models import SecurityMode, ServerConfig, ServerMode
 from digitalkin.grpc_servers.utils.types import GrpcServer, ServiceDescriptor, T
-
-logger = logging.getLogger(__name__)
+from digitalkin.logger import logger
 
 
 class BaseServer(abc.ABC):
@@ -132,7 +130,7 @@ class BaseServer(abc.ABC):
             # reflection queries with detailed service information
             reflection.enable_server_reflection(service_names, self.server)
 
-            logger.info("Added gRPC reflection service with services: %s", service_names)
+            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:
@@ -165,7 +163,7 @@ class BaseServer(abc.ABC):
                 if service_name not in self._service_names:
                     self._service_names.append(service_name)
 
-            logger.info("Added gRPC health checking service")
+            logger.debug("Added gRPC health checking service")
 
             # Set all services as SERVING
             for service_name in self._service_names:
@@ -255,7 +253,7 @@ class BaseServer(abc.ABC):
                 sync_server = cast("grpc.Server", server)
                 sync_server.add_secure_port(self.config.address, server_credentials)
 
-            logger.info("Added secure port %s", self.config.address)
+            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
@@ -277,7 +275,7 @@ class BaseServer(abc.ABC):
                 sync_server = cast("grpc.Server", server)
                 sync_server.add_insecure_port(self.config.address)
 
-            logger.info("Added insecure port %s", 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
@@ -301,7 +299,7 @@ class BaseServer(abc.ABC):
         self._add_reflection()
 
         # Start the server
-        logger.info("Starting gRPC server on %s", self.config.address)
+        logger.debug("Starting gRPC server on %s", self.config.address)
         try:
             if self.config.mode == ServerMode.ASYNC:
                 # For async server, use the event loop
@@ -314,7 +312,7 @@ class BaseServer(abc.ABC):
                 # For sync server, directly call start
                 sync_server = cast("grpc.Server", self.server)
                 sync_server.start()
-            logger.info("✅ gRPC server started on %s", self.config.address)
+            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}"
@@ -351,7 +349,7 @@ class BaseServer(abc.ABC):
         self._add_reflection()
 
         # Start the server
-        logger.info("Starting gRPC server on %s", self.config.address)
+        logger.debug("Starting gRPC server on %s", self.config.address)
         try:
             if self.config.mode == ServerMode.ASYNC:
                 await self._start_async()
@@ -359,7 +357,7 @@ class BaseServer(abc.ABC):
                 # For sync server in async context
                 sync_server = cast("grpc.Server", self.server)
                 sync_server.start()
-            logger.info("✅ gRPC server started on %s", self.config.address)
+            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}"
@@ -375,7 +373,7 @@ class BaseServer(abc.ABC):
             logger.warning("Attempted to stop server, but no server is running")
             return
 
-        logger.info("Stopping gRPC server...")
+        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:
@@ -392,13 +390,13 @@ class BaseServer(abc.ABC):
                     )
                     # Set server to None to avoid further operations
                     self.server = None
-                    logger.info("✅ gRPC server marked as stopped")
+                    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.info("Creating new event loop for shutdown")
+                logger.debug("Creating new event loop for shutdown")
                 try:
                     new_loop = asyncio.new_event_loop()
                     asyncio.set_event_loop(new_loop)
@@ -410,7 +408,7 @@ class BaseServer(abc.ABC):
             sync_server = cast("grpc.Server", self.server)
             sync_server.stop(grace=grace)
 
-        logger.info("✅ gRPC server stopped")
+        logger.debug("✅ gRPC server stopped")
         self.server = None
 
     async def _stop_async(self, grace: float | None = None) -> None:
@@ -437,7 +435,7 @@ class BaseServer(abc.ABC):
             logger.warning("Attempted to stop server, but no server is running")
             return
 
-        logger.info("Stopping gRPC server asynchronously...")
+        logger.debug("Stopping gRPC server asynchronously...")
         if self.config.mode == ServerMode.ASYNC:
             await self._stop_async(grace)
         else:
@@ -445,7 +443,7 @@ class BaseServer(abc.ABC):
             sync_server = cast("grpc.Server", self.server)
             sync_server.stop(grace=grace)
 
-        logger.info("✅ gRPC server stopped")
+        logger.debug("✅ gRPC server stopped")
         self.server = None
 
     def wait_for_termination(self) -> None:

digitalkin/grpc_servers/module_server.py

@@ -1,6 +1,5 @@
 """Module gRPC server implementation for DigitalKin."""
 
-import logging
 from pathlib import Path
 import uuid
 
@@ -25,8 +24,7 @@ from digitalkin_proto.digitalkin.module_registry.v2 import (
     module_registry_service_pb2_grpc,
     registration_pb2,
 )
-
-logger = logging.getLogger(__name__)
+from digitalkin.logger import logger
 
 
 class ModuleServer(BaseServer):
@@ -70,14 +68,14 @@ class ModuleServer(BaseServer):
             msg = "Server must be created before registering servicers"
             raise RuntimeError(msg)
 
-        logger.info("Registering module servicer for %s", self.module_class.__name__)
+        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.info("Registered Module servicer")
+        logger.debug("Registered Module servicer")
 
     def start(self) -> None:
         """Start the module server and register with the registry if configured."""
@@ -115,6 +113,7 @@ class ModuleServer(BaseServer):
             logger.critical(
                 "Setup post init started with config: %s", 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:
@@ -134,7 +133,7 @@ class ModuleServer(BaseServer):
         Raises:
             ServerError: If communication with the registry server fails.
         """
-        logger.info(
+        logger.debug(
             "Registering module with registry at %s",
             self.server_config.registry_address,
         )
@@ -172,7 +171,7 @@ class ModuleServer(BaseServer):
 
             try:
                 # Call the register method
-                logger.info(
+                logger.debug(
                     "Request sent to registry for module: %s:%s",
                     self.module_class.metadata["name"],
                     self.module_class.metadata["module_id"],
@@ -180,7 +179,7 @@ class ModuleServer(BaseServer):
                 response = stub.RegisterModule(request)
 
                 if response.success:
-                    logger.info("Module registered successfully")
+                    logger.debug("Module registered successfully")
                 else:
                     logger.error("Module registration failed")
             except grpc.RpcError:
@@ -193,7 +192,7 @@ class ModuleServer(BaseServer):
         Raises:
             ServerError: If communication with the registry server fails.
         """
-        logger.info(
+        logger.debug(
             "Deregistering module from registry at %s",
             self.server_config.registry_address,
         )
@@ -214,7 +213,7 @@ class ModuleServer(BaseServer):
             response = stub.DeregisterModule(request)
 
             if response.success:
-                logger.info("Module deregistered successfull")
+                logger.debug("Module deregistered successfull")
             else:
                 logger.error("Module deregistration failed")
         except grpc.RpcError:

digitalkin/grpc_servers/module_servicer.py

@@ -1,7 +1,6 @@
 """Module servicer implementation for DigitalKin."""
 
-import asyncio
-import logging
+from argparse import ArgumentParser, Namespace
 from collections.abc import AsyncGenerator
 from typing import Any
 
@@ -15,19 +14,20 @@ from digitalkin_proto.digitalkin.module.v2 import (
 from google.protobuf import json_format, struct_pb2
 
 from digitalkin.grpc_servers.utils.exceptions import ServicerError
-from digitalkin.models.module import OutputModelT
+from digitalkin.logger import logger
 from digitalkin.models.module.module import ModuleStatus
 from digitalkin.modules._base_module import BaseModule
-from digitalkin.modules.job_manager import JobManager
+from digitalkin.modules.job_manager.base_job_manager import BaseJobManager
+from digitalkin.modules.job_manager.job_manager_models import JobManagerMode
 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
 
-logger = logging.getLogger(__name__)
 
-
-class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer):
+class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
     """Implementation of the ModuleService.
 
     This servicer handles interactions with a DigitalKin module.
@@ -37,7 +37,31 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer):
         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.
@@ -46,15 +70,16 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer):
             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)
-        self.setup = GrpcSetup() if self.job_manager.args.services_mode == ServicesMode.REMOTE else DefaultSetup()
+        job_manager_class = self.args.job_manager_mode.get_manager_class()
+        self.job_manager = job_manager_class(module_class, self.args.services_mode)
 
-    async def add_to_queue(self, job_id: str, output_data: OutputModelT) -> None:  # type: ignore
-        """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})
+        logger.debug(
+            "ModuleServicer initialized with job manager: %s | %s",
+            self.args.job_manager_mode,
+            self.job_manager,
+        )
+        self.setup = GrpcSetup() if self.args.services_mode == ServicesMode.REMOTE else DefaultSetup()
 
     async def StartModule(  # noqa: N802
         self,
@@ -87,47 +112,45 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer):
         if not setup_data_class:
             msg = "No setup data returned."
             raise ServicerError(msg)
-        # TODO: Check failure of setup data format
+
         setup_data = self.module_class.create_setup_model(setup_data_class.current_setup_version.content)
 
-        # setup_id should be use to request a precise setup from the module
-        # Create a job for this execution
-        result: tuple[str, BaseModule] = await self.job_manager.create_job(
+        # create a task to run the module in background
+        job_id = await self.job_manager.create_job(
             input_data,
             setup_data,
             mission_id=request.mission_id,
             setup_version_id=setup_data_class.current_setup_version.id,
-            callback=self.add_to_queue,
         )
-        job_id, module = result
-
-        while module.status == ModuleStatus.RUNNING or not self.queue.empty():
-            output_data: dict = await self.queue.get()
-
-            if job_id not in output_data or 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)
-                yield lifecycle_pb2.StartModuleResponse(success=False)
-                return
-
-            if output_data[job_id].get("error", None) is not None:
-                context.set_code(output_data[job_id]["error"]["code"])
-                context.set_details(output_data[job_id]["error"]["error_message"])
-                yield lifecycle_pb2.StartModuleResponse(success=False)
-                return
-
-            output_proto = json_format.ParseDict(
-                output_data[job_id],
-                struct_pb2.Struct(),
-                ignore_unknown_fields=True,
-            )
-            yield lifecycle_pb2.StartModuleResponse(
-                success=True,
-                output=output_proto,
-                job_id=job_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
+
+        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:
+                    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("Error in output_data")
+                    context.set_code(message["short_description"])
+                    context.set_details(message["exception"])
+                    yield lifecycle_pb2.StartModuleResponse(success=False, job_id=job_id)
+                    break
+
+                if message.get("code", None) is not None and message.get("code") == "__END_OF_STREAM__":
+                    yield lifecycle_pb2.StartModuleResponse(success=True, job_id=job_id)
+                    break
+
+                proto = json_format.ParseDict(message, struct_pb2.Struct(), ignore_unknown_fields=True)
+                yield lifecycle_pb2.StartModuleResponse(success=True, output=proto, job_id=job_id)
+        logger.info("Job %s finished", job_id)
 
     async def StopModule(  # noqa: N802
         self,
@@ -143,23 +166,20 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer):
         Returns:
             A response indicating success or failure.
         """
-        logger.info("StopModule called for module: '%s'", self.module_class.__name__)
+        logger.debug("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"
+        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)
 
-        # Update the job status
-        await self.job_manager.modules[job_id].stop()
-
-        logger.info("Job %s stopped successfully", job_id)
+        logger.debug("Job %s stopped successfully", request.job_id)
         return lifecycle_pb2.StopModuleResponse(success=True)
 
-    def GetModuleStatus(  # noqa: N802
+    async def GetModuleStatus(  # noqa: N802
         self,
         request: monitoring_pb2.GetModuleStatusRequest,
         context: grpc.ServicerContext,
@@ -173,33 +193,33 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer):
         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)
+        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=True,
-                status=status.name,
+                success=False,
+                status=ModuleStatus.NOT_FOUND.name,
                 job_id=request.job_id,
             )
 
-        logger.info("Job %s status: '%s'", request.job_id, ModuleStatus.NOT_FOUND)
+        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=False,
-            status=ModuleStatus.NOT_FOUND.name,
+            success=True,
+            status=status.name,
             job_id=request.job_id,
         )
 
-    def GetModuleJobs(  # noqa: N802
+    async def GetModuleJobs(  # noqa: N802
         self,
         request: monitoring_pb2.GetModuleJobsRequest,  # noqa: ARG002
         context: grpc.ServicerContext,  # noqa: ARG002
@@ -213,20 +233,22 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer):
         Returns:
             A response with information about active jobs.
         """
-        logger.info("GetModuleJobs called for module: '%s'", self.module_class.__name__)
+        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,
+                    job_status=job_data["status"].name,
                 )
-                for job_id, job_data in self.job_manager.modules.items()
+                for job_id, job_data in modules.items()
             ],
         )
 
-    def GetModuleInput(  # noqa: N802
+    async def GetModuleInput(  # noqa: N802
         self,
         request: information_pb2.GetModuleInputRequest,
         context: grpc.ServicerContext,
@@ -240,7 +262,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer):
         Returns:
             A response with the module's input schema.
         """
-        logger.info("GetModuleInput called for module: '%s'", self.module_class.__name__)
+        logger.debug("GetModuleInput called for module: '%s'", self.module_class.__name__)
 
         # Get input schema if available
         try:
@@ -262,7 +284,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer):
             input_schema=input_format_struct,
         )
 
-    def GetModuleOutput(  # noqa: N802
+    async def GetModuleOutput(  # noqa: N802
         self,
         request: information_pb2.GetModuleOutputRequest,
         context: grpc.ServicerContext,
@@ -276,7 +298,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer):
         Returns:
             A response with the module's output schema.
         """
-        logger.info("GetModuleOutput called for module: '%s'", self.module_class.__name__)
+        logger.debug("GetModuleOutput called for module: '%s'", self.module_class.__name__)
 
         # Get output schema if available
         try:
@@ -298,7 +320,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer):
             output_schema=output_format_struct,
         )
 
-    def GetModuleSetup(  # noqa: N802
+    async def GetModuleSetup(  # noqa: N802
         self,
         request: information_pb2.GetModuleSetupRequest,
         context: grpc.ServicerContext,
@@ -312,7 +334,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer):
         Returns:
             A response with the module's setup information.
         """
-        logger.info("GetModuleSetup called for module: '%s'", self.module_class.__name__)
+        logger.debug("GetModuleSetup called for module: '%s'", self.module_class.__name__)
 
         # Get setup schema if available
         try:

digitalkin/grpc_servers/registry_server.py

@@ -1,7 +1,5 @@
 """Registry gRPC server implementation for DigitalKin."""
 
-import logging
-
 from digitalkin_proto.digitalkin.module_registry.v2 import (
     module_registry_service_pb2,
     module_registry_service_pb2_grpc,
@@ -10,8 +8,7 @@ from digitalkin_proto.digitalkin.module_registry.v2 import (
 from digitalkin.grpc_servers._base_server import BaseServer
 from digitalkin.grpc_servers.registry_servicer import RegistryModule, RegistryServicer
 from digitalkin.grpc_servers.utils.models import RegistryServerConfig
-
-logger = logging.getLogger(__name__)
+from digitalkin.logger import logger
 
 
 class RegistryServer(BaseServer):
@@ -48,14 +45,14 @@ class RegistryServer(BaseServer):
             msg = "Server must be created before registering servicers"
             raise RuntimeError(msg)
 
-        logger.info("Registering registry servicer")
+        logger.debug("Registering registry servicer")
         self.registry_servicer = RegistryServicer()
         self.register_servicer(
             self.registry_servicer,
             module_registry_service_pb2_grpc.add_ModuleRegistryServiceServicer_to_server,
             service_descriptor=module_registry_service_pb2.DESCRIPTOR,
         )
-        logger.info("Registered registry servicer")
+        logger.debug("Registered registry servicer")
 
     def get_registered_modules(self) -> list[RegistryModule]:
         """Get a list of all registered modules.