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

digitalkin/modules/job_manager/single_job_manager.py

@@ -0,0 +1,218 @@
+"""Background module manager with single instance."""
+
+import asyncio
+import uuid
+from collections.abc import AsyncGenerator, AsyncIterator
+from contextlib import asynccontextmanager
+from typing import Any, Generic
+
+import grpc
+
+from digitalkin.logger import logger
+from digitalkin.models import ModuleStatus
+from digitalkin.models.module import InputModelT, OutputModelT, SetupModelT
+from digitalkin.modules._base_module import BaseModule
+from digitalkin.modules.job_manager.base_job_manager import BaseJobManager
+from digitalkin.services.services_models import ServicesMode
+
+
+class SingleJobManager(BaseJobManager, Generic[InputModelT, SetupModelT]):
+    """Manages a single instance of a module job.
+
+    This class ensures that only one instance of a module job is active at a time.
+    It provides functionality to create, stop, and monitor module jobs, as well as
+    to handle their output data.
+    """
+
+    modules: dict[str, BaseModule]
+    queue: dict[str, asyncio.Queue]
+
+    def __init__(
+        self,
+        module_class: type[BaseModule],
+        services_mode: ServicesMode,
+    ) -> None:
+        """Initialize the job manager.
+
+        Args:
+            module_class: The class of the module to be managed.
+            services_mode: The mode of operation for the services (e.g., ASYNC or SYNC).
+        """
+        super().__init__(module_class, services_mode)
+
+        self._lock = asyncio.Lock()
+        self.modules: dict[str, BaseModule] = {}
+        self.queues: dict[str, asyncio.Queue] = {}
+
+    async def add_to_queue(self, job_id: str, output_data: OutputModelT) -> None:  # type: ignore
+        """Add output data to the queue for a specific job.
+
+        This method is used as a callback to handle output data generated by a module job.
+
+        Args:
+            job_id: The unique identifier of the job.
+            output_data: The output data produced by the job.
+        """
+        await self.queues[job_id].put(output_data.model_dump())
+
+    @asynccontextmanager  # type: ignore
+    async def generate_stream_consumer(self, job_id: str) -> AsyncIterator[AsyncGenerator[dict[str, Any], None]]:  # type: ignore
+        """Generate a stream consumer for a module's output data.
+
+        This method creates an asynchronous generator that streams output data
+        from a specific module job. If the module does not exist, it generates
+        an error message.
+
+        Args:
+            job_id: The unique identifier of the job.
+
+        Yields:
+            AsyncGenerator: A stream of output data or error messages.
+        """
+        module = self.modules.get(job_id, None)
+
+        logger.debug("Module %s found: %s", job_id, module)
+
+        async def _stream() -> AsyncGenerator[dict[str, Any], Any]:
+            """Stream output data from the module.
+
+            Yields:
+                dict: Output data generated by the module.
+            """
+            if module is None:
+                yield {
+                    "error": {
+                        "error_message": f"Module {job_id} not found",
+                        "code": grpc.StatusCode.NOT_FOUND,
+                    }
+                }
+                return
+
+            try:
+                while module.status == ModuleStatus.RUNNING or (
+                    not self.queues[job_id].empty()
+                    and module.status
+                    in {
+                        ModuleStatus.STOPPED,
+                        ModuleStatus.STOPPING,
+                    }
+                ):
+                    logger.info(f"{job_id=}: {module.status=}")
+                    yield await self.queues[job_id].get()
+                logger.info(f"{job_id=}: {module.status=} | {self.queues[job_id].empty()}")
+
+            finally:
+                del self.queues[job_id]
+
+        yield _stream()
+
+    async def create_job(
+        self,
+        input_data: InputModelT,
+        setup_data: SetupModelT,
+        mission_id: str,
+        setup_version_id: str,
+    ) -> str:
+        """Create and start a new module job.
+
+        This method initializes a new module job, assigns it a unique job ID,
+        and starts it in the background.
+
+        Args:
+            input_data: The input data required to start the job.
+            setup_data: The setup configuration for the module.
+            mission_id: The mission ID associated with the job.
+            setup_version_id: The setup ID associated with the module.
+
+        Returns:
+            str: The unique identifier (job ID) of the created job.
+
+        Raises:
+            Exception: If the module fails to start.
+        """
+        job_id = str(uuid.uuid4())
+        # TODO: Ensure the job_id is unique.
+        module = self.module_class(job_id, mission_id=mission_id, setup_version_id=setup_version_id)
+        self.modules[job_id] = module
+        self.queues[job_id] = asyncio.Queue()
+
+        try:
+            await module.start(
+                input_data,
+                setup_data,
+                await self.job_specific_callback(self.add_to_queue, job_id),
+            )
+            logger.debug("Module %s (%s) started successfully", job_id, module.name)
+        except Exception:
+            # Remove the module from the manager in case of an error.
+            del self.modules[job_id]
+            logger.exception("Failed to start module %s: %s", job_id)
+            raise
+        else:
+            return job_id
+
+    async def stop_module(self, job_id: str) -> bool:
+        """Stop a running module job.
+
+        Args:
+            job_id: The unique identifier of the job to stop.
+
+        Returns:
+            bool: True if the module was successfully stopped, False if it does not exist.
+
+        Raises:
+            Exception: If an error occurs while stopping the module.
+        """
+        async with self._lock:
+            module = self.modules.get(job_id)
+            if not module:
+                logger.warning(f"Module {job_id} not found")
+                return False
+            try:
+                await module.stop()
+                # should maybe be added in finally
+                del self.queues[job_id]
+                del self.modules[job_id]
+                logger.debug(f"Module {job_id} ({module.name}) stopped successfully")
+            except Exception as e:
+                logger.error(f"Error while stopping module {job_id}: {e}")
+                raise
+            else:
+                return True
+
+    async def get_module_status(self, job_id: str) -> ModuleStatus | None:
+        """Retrieve the status of a module job.
+
+        Args:
+            job_id: The unique identifier of the job.
+
+        Returns:
+            ModuleStatus | None: The status of the module, or None if it does not exist.
+        """
+        module = self.modules.get(job_id)
+        return module.status if module else None
+
+    async def stop_all_modules(self) -> None:
+        """Stop all currently running module jobs.
+
+        This method ensures that all active jobs are gracefully terminated.
+        """
+        async with self._lock:
+            stop_tasks = [self.stop_module(job_id) for job_id in list(self.modules.keys())]
+            if stop_tasks:
+                await asyncio.gather(*stop_tasks, return_exceptions=True)
+
+    async def list_modules(self) -> dict[str, dict[str, Any]]:
+        """List all modules along with their statuses.
+
+        Returns:
+            dict[str, dict[str, Any]]: A dictionary containing information about all modules and their statuses.
+        """
+        return {
+            job_id: {
+                "name": module.name,
+                "status": module.status,
+                "class": module.__class__.__name__,
+            }
+            for job_id, module in self.modules.items()
+        }

digitalkin/modules/job_manager/taskiq_broker.py

@@ -0,0 +1,173 @@
+"""Taskiq broker & RSTREAM producer for the job manager."""
+
+import asyncio
+import json
+import logging
+import os
+import pickle  # noqa: S403
+
+from rstream import Producer
+from rstream.exceptions import PreconditionFailed
+from taskiq import Context, TaskiqDepends, TaskiqMessage
+from taskiq.abc.formatter import TaskiqFormatter
+from taskiq.compat import model_validate
+from taskiq.message import BrokerMessage
+from taskiq_aio_pika import AioPikaBroker
+
+from digitalkin.logger import logger
+from digitalkin.models.module.module_types import OutputModelT
+from digitalkin.modules._base_module import BaseModule
+from digitalkin.modules.job_manager.base_job_manager import BaseJobManager
+from digitalkin.modules.job_manager.job_manager_models import StreamCodeModel
+from digitalkin.services.services_config import ServicesConfig
+from digitalkin.services.services_models import ServicesMode
+
+logging.getLogger("taskiq").setLevel(logging.INFO)
+logging.getLogger("aiormq").setLevel(logging.INFO)
+logging.getLogger("aio_pika").setLevel(logging.INFO)
+logging.getLogger("rstream").setLevel(logging.INFO)
+
+
+class PickleFormatter(TaskiqFormatter):
+    """Formatter that pickles the JSON-dumped TaskiqMessage.
+
+    This lets you send arbitrary Python objects (classes, functions, etc.)
+    by first converting to JSON-safe primitives, then pickling that string.
+    """
+
+    def dumps(self, message: TaskiqMessage) -> BrokerMessage:  # noqa: PLR6301
+        """Dumps message from python complex object to JSON.
+
+        Args:
+            message: TaskIQ message
+
+        Returns:
+            BrokerMessage with mandatory information for TaskIQ
+        """
+        payload: bytes = pickle.dumps(message)
+
+        return BrokerMessage(
+            task_id=message.task_id,
+            task_name=message.task_name,
+            message=payload,
+            labels=message.labels,
+        )
+
+    def loads(self, message: bytes) -> TaskiqMessage:  # noqa: PLR6301
+        """Recreate Python object from bytes.
+
+        Args:
+            message: Broker message from bytes.
+
+        Returns:
+            message with TaskIQ format
+        """
+        json_str = pickle.loads(message)  # noqa: S301
+        return model_validate(TaskiqMessage, json_str)
+
+
+def define_producer() -> Producer:
+    """Get from the env the connection parameter to RabbitMQ.
+
+    Returns:
+        Producer
+    """
+    host: str = os.environ.get("RABBITMQ_RSTREAM_HOST", "localhost")
+    port: str = os.environ.get("RABBITMQ_RSTREAM_PORT", "5552")
+    username: str = os.environ.get("RABBITMQ_RSTREAM_USERNAME", "guest")
+    password: str = os.environ.get("RABBITMQ_RSTREAM_PASSWORD", "guest")
+
+    logger.info("Connection to RabbitMQ: %s:%s.", host, port)
+    return Producer(host=host, port=int(port), username=username, password=password)
+
+
+async def init_rstream() -> None:
+    """Init a stream for every tasks."""
+    try:
+        await RSTREAM_PRODUCER.create_stream(
+            STREAM,
+            exists_ok=True,
+            arguments={"max-length-bytes": STREAM_RETENTION},
+        )
+    except PreconditionFailed:
+        logger.warning("stream already exist")
+
+
+def define_broker() -> AioPikaBroker:
+    """Define broker with from env paramter.
+
+    Returns:
+        Broker: connected to RabbitMQ and with custom formatter.
+    """
+    host: str = os.environ.get("RABBITMQ_BROKER_HOST", "localhost")
+    port: str = os.environ.get("RABBITMQ_BROKER_PORT", "5672")
+    username: str = os.environ.get("RABBITMQ_BROKER_USERNAME", "guest")
+    password: str = os.environ.get("RABBITMQ_BROKER_PASSWORD", "guest")
+
+    broker = AioPikaBroker(
+        f"amqp://{username}:{password}@{host}:{port}",
+        startup=[init_rstream],
+    )
+    broker.formatter = PickleFormatter()
+    return broker
+
+
+STREAM = "taskiq_data"
+STREAM_RETENTION = 200_000
+RSTREAM_PRODUCER = define_producer()
+TASKIQ_BROKER = define_broker()
+
+
+async def send_message_to_stream(job_id: str, output_data: OutputModelT) -> None:  # type: ignore
+    """Callback define to add a message frame to the Rstream.
+
+    Args:
+        job_id: id of the job that sent the message
+        output_data: message body as a OutputModelT or error / stream_code
+    """
+    body = json.dumps({"job_id": job_id, "output_data": output_data.model_dump()}).encode("utf-8")
+    await RSTREAM_PRODUCER.send(stream=STREAM, message=body)
+
+
+@TASKIQ_BROKER.task
+async def run_task(
+    mission_id: str,
+    setup_version_id: str,
+    module_class: type[BaseModule],
+    services_mode: ServicesMode,
+    input_data: dict,
+    setup_data: dict,
+    context: Context = TaskiqDepends(),
+) -> None:
+    """TaskIQ task allowing a module to compute in the background asynchronously.
+
+    Args:
+        mission_id: str,
+        setup_version_id: The setup ID associated with the module.
+        module_class: type[BaseModule],
+        services_mode: ServicesMode,
+        input_data: dict,
+        setup_data: dict,
+        context: Allow TaskIQ context access
+    """
+    logger.warning("%s", services_mode)
+    services_config = ServicesConfig(
+        services_config_strategies=module_class.services_config_strategies,
+        services_config_params=module_class.services_config_params,
+        mode=services_mode,
+    )
+    setattr(module_class, "services_config", services_config)
+    logger.warning("%s | %s", services_config, module_class.services_config)
+
+    job_id = context.message.task_id
+    callback = await BaseJobManager.job_specific_callback(send_message_to_stream, job_id)
+    module = module_class(job_id, mission_id=mission_id, setup_version_id=setup_version_id)
+
+    await module.start(
+        input_data,
+        setup_data,
+        callback,
+        # ensure that the callback is called when the task is done + allow asyncio to run
+        # TODO: should define a BaseModel for stream code / error
+        done_callback=lambda _: asyncio.create_task(callback(StreamCodeModel(code="__END_OF_STREAM__"))),
+    )

digitalkin/modules/job_manager/taskiq_job_manager.py

@@ -0,0 +1,213 @@
+"""Taskiq job manager module."""
+
+try:
+    import taskiq  # noqa: F401
+
+except ImportError:
+    msg = "Install digitalkin[taskiq] to use this functionality\n$ uv pip install digitalkin[taskiq]."
+    raise ImportError(msg)
+
+import asyncio
+import contextlib
+import json
+import os
+from collections.abc import AsyncGenerator, AsyncIterator
+from contextlib import asynccontextmanager
+from typing import TYPE_CHECKING, Any, Generic
+
+from rstream import Consumer, ConsumerOffsetSpecification, MessageContext, OffsetType
+
+from digitalkin.logger import logger
+from digitalkin.models.module import InputModelT, SetupModelT
+from digitalkin.models.module.module import ModuleStatus
+from digitalkin.modules._base_module import BaseModule
+from digitalkin.modules.job_manager.base_job_manager import BaseJobManager
+from digitalkin.modules.job_manager.taskiq_broker import STREAM, STREAM_RETENTION, TASKIQ_BROKER
+from digitalkin.services.services_models import ServicesMode
+
+if TYPE_CHECKING:
+    from taskiq.task import AsyncTaskiqTask
+
+
+class TaskiqJobManager(BaseJobManager, Generic[InputModelT, SetupModelT]):
+    """Taskiq job manager for running modules in Taskiq tasks."""
+
+    services_mode: ServicesMode
+
+    @staticmethod
+    def _define_consumer() -> Consumer:
+        """Get from the env the connection parameter to RabbitMQ.
+
+        Returns:
+            Consumer
+        """
+        host: str = os.environ.get("RABBITMQ_RSTREAM_HOST", "localhost")
+        port: str = os.environ.get("RABBITMQ_RSTREAM_PORT", "5552")
+        username: str = os.environ.get("RABBITMQ_RSTREAM_USERNAME", "guest")
+        password: str = os.environ.get("RABBITMQ_RSTREAM_PASSWORD", "guest")
+
+        logger.info("Connection to RabbitMQ: %s:%s.", host, port)
+        return Consumer(host=host, port=int(port), username=username, password=password)
+
+    async def _on_message(self, message: bytes, message_context: MessageContext) -> None:  # noqa: ARG002
+        """Internal callback: parse JSON and route to the correct job queue."""
+        try:
+            data = json.loads(message.decode("utf-8"))
+        except json.JSONDecodeError:
+            return
+        job_id = data.get("job_id")
+        if not job_id:
+            return
+        queue = self.job_queues.get(job_id)
+        if queue:
+            await queue.put(data.get("output_data"))
+
+    async def _start(self) -> None:
+        await TASKIQ_BROKER.startup()
+
+        self.stream_consumer = self._define_consumer()
+
+        await self.stream_consumer.create_stream(
+            STREAM,
+            exists_ok=True,
+            arguments={"max-length-bytes": STREAM_RETENTION},
+        )
+        await self.stream_consumer.start()
+
+        start_spec = ConsumerOffsetSpecification(OffsetType.LAST)
+        # on_message use bytes instead of AMQPMessage
+        await self.stream_consumer.subscribe(
+            stream=STREAM,
+            subscriber_name=f"""subscriber_{os.environ.get("SERVER_NAME", "module_servicer")}""",
+            callback=self._on_message,  # type: ignore
+            offset_specification=start_spec,
+        )
+        self.stream_consumer_task = asyncio.create_task(
+            self.stream_consumer.run(),
+            name="stream_consumer_task",
+        )
+
+    async def _stop(self) -> None:
+        # Signal the consumer to stop
+        await self.stream_consumer.close()
+        # Cancel the background task
+        self.stream_consumer_task.cancel()
+        with contextlib.suppress(asyncio.CancelledError):
+            await self.stream_consumer_task
+
+    @asynccontextmanager  # type: ignore
+    async def generate_stream_consumer(self, job_id: str) -> AsyncIterator[AsyncGenerator[dict[str, Any], None]]:  # type: ignore
+        """Generate a stream consumer for the RStream stream.
+
+        Args:
+            job_id: The job ID to filter messages.
+
+        Yields:
+            messages: The stream messages from the associated module.
+        """
+        queue: asyncio.Queue = asyncio.Queue(maxsize=self.max_queue_size)
+        self.job_queues[job_id] = queue
+
+        async def _stream() -> AsyncGenerator[dict[str, Any], Any]:
+            """Generate the stream allowing flowless communication.
+
+            Yields:
+                dict: generated object from the module
+            """
+            while True:
+                item = await queue.get()
+                queue.task_done()
+                yield item
+
+                while True:
+                    try:
+                        item = queue.get_nowait()
+                    except asyncio.QueueEmpty:
+                        break
+                    queue.task_done()
+                    yield item
+
+        try:
+            yield _stream()
+        finally:
+            self.job_queues.pop(job_id, None)
+
+    def __init__(
+        self,
+        module_class: type[BaseModule],
+        services_mode: ServicesMode,
+    ) -> None:
+        """Initialize the Taskiq job manager."""
+        super().__init__(module_class, services_mode)
+
+        logger.warning("TaskiqJobManager initialized with app: %s", TASKIQ_BROKER)
+        self.services_mode = services_mode
+        self.job_queues: dict[str, asyncio.Queue] = {}
+        self.max_queue_size = 1000
+
+    async def create_job(
+        self,
+        input_data: InputModelT,
+        setup_data: SetupModelT,
+        mission_id: str,
+        setup_version_id: str,
+    ) -> str:
+        """Launches the module_task in Taskiq, returns the Taskiq task id as job_id.
+
+        Args:
+            input_data: Input data for the module
+            setup_data: Setup data for the module
+            mission_id: Mission ID for the module
+            setup_version_id: The setup ID associated with the module.
+
+        Returns:
+            job_id: The Taskiq task id.
+
+        Raises:
+            ValueError: If the task is not found.
+        """
+        task = TASKIQ_BROKER.find_task("digitalkin.modules.job_manager.taskiq_broker:run_task")
+
+        if task is None:
+            msg = "Task not found"
+            raise ValueError(msg)
+
+        running_task: AsyncTaskiqTask[Any] = await task.kiq(
+            mission_id,
+            setup_version_id,
+            self.module_class,
+            self.services_mode,
+            input_data.model_dump(),
+            setup_data.model_dump(),
+        )
+        job_id = running_task.task_id
+        result = await running_task.wait_result(timeout=10)
+        logger.debug("Job %s with data %s", job_id, result)
+        return job_id
+
+    async def stop_module(self, job_id: str) -> bool:
+        """Revoke (terminate) the Taskiq task with id.
+
+        Args:
+            job_id: The Taskiq task id to stop.
+
+        Raises:
+            bool: True if the task was successfully revoked, False otherwise.
+        """
+        msg = "stop_module not implemented in TaskiqJobManager"
+        raise NotImplementedError(msg)
+
+    async def stop_all_modules(self) -> None:
+        """Stop all running modules."""
+        msg = "stop_all_modules not implemented in TaskiqJobManager"
+        raise NotImplementedError(msg)
+
+    async def get_module_status(self, job_id: str) -> ModuleStatus | None:
+        """Query a module status."""
+        msg = "get_module_status not implemented in TaskiqJobManager"
+        raise NotImplementedError(msg)
+
+    async def list_modules(self) -> dict[str, dict[str, Any]]:
+        """List all modules."""
+        msg = "list_modules not implemented in TaskiqJobManager"
+        raise NotImplementedError(msg)

digitalkin/services/base_strategy.py

@@ -9,10 +9,12 @@ class BaseStrategy(ABC):
     This class defines the interface for all strategies.
     """
 
-    def __init__(self, mission_id: str) -> None:
+    def __init__(self, mission_id: str, setup_version_id: str) -> None:
         """Initialize the strategy.
 
         Args:
             mission_id: The ID of the mission this strategy is associated with
+            setup_version_id: The ID of the setup version this strategy is associated with
         """
         self.mission_id: str = mission_id
+        self.setup_version_id: str = setup_version_id