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

digitalkin/modules/job_manager/taskiq_broker.py

@@ -0,0 +1,214 @@
+"""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_start_module(
+    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__"))),
+    )
+
+
+@TASKIQ_BROKER.task
+async def run_config_module(
+    mission_id: str,
+    setup_version_id: str,
+    module_class: type[BaseModule],
+    services_mode: ServicesMode,
+    config_setup_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,
+        config_setup_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_config_setup(
+        module_class.create_config_setup_model(config_setup_data),
+        module_class.create_setup_model(setup_data),
+        callback,
+    )

digitalkin/modules/job_manager/taskiq_job_manager.py

@@ -0,0 +1,286 @@
+"""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 ConfigSetupModelT, 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, ConfigSetupModelT]):
+    """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
+
+    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 generate_config_setup_module_response(self, job_id: str) -> SetupModelT:
+        """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.
+
+        Returns:
+            SetupModelT: the SetupModelT object fully processed.
+        """
+        queue: asyncio.Queue = asyncio.Queue(maxsize=self.max_queue_size)
+        self.job_queues[job_id] = queue
+
+        try:
+            item = await queue.get()
+            queue.task_done()
+            return item
+        finally:
+            logger.info(f"generate_config_setup_module_response: {job_id=}: {self.job_queues[job_id].empty()}")
+            self.job_queues.pop(job_id, None)
+
+    async def create_config_setup_instance_job(
+        self,
+        config_setup_data: ConfigSetupModelT,
+        setup_data: SetupModelT,
+        mission_id: str,
+        setup_version_id: str,
+    ) -> str:
+        """Create and start a new module setup configuration job.
+
+        This method initializes a new module job, assigns it a unique job ID,
+        and starts the config setup it in the background.
+
+        Args:
+            config_setup_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.
+
+        Returns:
+            str: The unique identifier (job ID) of the created job.
+
+        Raises:
+            TypeError: If the function is called with bad data type.
+            ValueError: If the module fails to start.
+        """
+        task = TASKIQ_BROKER.find_task("digitalkin.modules.job_manager.taskiq_broker:run_config_module")
+
+        if task is None:
+            msg = "Task not found"
+            raise ValueError(msg)
+
+        if config_setup_data is None:
+            msg = "config_setup_data must be a valid model with model_dump method"
+            raise TypeError(msg)
+
+        running_task: AsyncTaskiqTask[Any] = await task.kiq(
+            mission_id,
+            setup_version_id,
+            self.module_class,
+            self.services_mode,
+            config_setup_data.model_dump(),  # type: ignore
+            setup_data.model_dump(),
+        )
+
+        job_id = running_task.task_id
+        result = await running_task.wait_result(timeout=10)
+        logger.info("Job %s with data %s", job_id, result)
+        return job_id
+
+    @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)
+
+    async def create_module_instance_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_start_module")
+
+        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/modules/tool_module.py

@@ -3,8 +3,9 @@
 from abc import ABC
 
 from digitalkin.models.module import InputModelT, OutputModelT, SecretModelT, SetupModelT
+from digitalkin.models.module.module_types import ConfigSetupModelT
 from digitalkin.modules._base_module import BaseModule # type: ignore
 
 
-class ToolModule(BaseModule[InputModelT, OutputModelT, SetupModelT, SecretModelT], ABC):
+class ToolModule(BaseModule[InputModelT, OutputModelT, SetupModelT, SecretModelT,ConfigSetupModelT,], ABC):
     """ToolModule extends BaseModule to implement specific module types."""

digitalkin/modules/trigger_module.py

@@ -2,8 +2,10 @@
 
 from abc import ABC
 
+from digitalkin.models.module.module_types import ConfigSetupModelT
 from digitalkin.modules._base_module import BaseModule, InputModelT, OutputModelT, SecretModelT, SetupModelT
 
 
-class TriggerModule(BaseModule[InputModelT, OutputModelT, SetupModelT, SecretModelT], ABC):
+class TriggerModule(BaseModule[InputModelT, OutputModelT, SetupModelT, SecretModelT,
+        ConfigSetupModelT,], ABC):
     """TriggerModule extends BaseModule to implement specific module types."""

digitalkin/services/cost/default_cost.py

@@ -1,11 +1,15 @@
 """Default cost."""
 
-import logging
 from typing import Literal
 
-from digitalkin.services.cost.cost_strategy import CostConfig, CostData, CostServiceError, CostStrategy, CostType
-
-logger = logging.getLogger(__name__)
+from digitalkin.logger import logger
+from digitalkin.services.cost.cost_strategy import (
+    CostConfig,
+    CostData,
+    CostServiceError,
+    CostStrategy,
+    CostType,
+)
 
 
 class DefaultCost(CostStrategy):

digitalkin/services/cost/grpc_cost.py

@@ -1,6 +1,5 @@
-"""This module implements the default Cost strategy."""
+"""This module implements the gRPC Cost strategy."""
 
-import logging
 from collections.abc import Generator
 from contextlib import contextmanager
 from typing import Any, Literal
@@ -11,9 +10,14 @@ from google.protobuf import json_format
 from digitalkin.grpc_servers.utils.exceptions import ServerError
 from digitalkin.grpc_servers.utils.grpc_client_wrapper import GrpcClientWrapper
 from digitalkin.grpc_servers.utils.models import ClientConfig
-from digitalkin.services.cost.cost_strategy import CostConfig, CostData, CostServiceError, CostStrategy, CostType
-
-logger = logging.getLogger(__name__)
+from digitalkin.logger import logger
+from digitalkin.services.cost.cost_strategy import (
+    CostConfig,
+    CostData,
+    CostServiceError,
+    CostStrategy,
+    CostType,
+)
 
 
 class GrpcCost(CostStrategy, GrpcClientWrapper):
@@ -51,13 +55,17 @@ class GrpcCost(CostStrategy, GrpcClientWrapper):
             raise CostServiceError(msg) from e
 
     def __init__(
-        self, mission_id: str, setup_version_id: str, config: dict[str, CostConfig], client_config: ClientConfig
+        self,
+        mission_id: str,
+        setup_version_id: str,
+        config: dict[str, CostConfig],
+        client_config: ClientConfig,
     ) -> None:
         """Initialize the cost."""
         super().__init__(mission_id=mission_id, setup_version_id=setup_version_id, config=config)
         channel = self._init_channel(client_config)
         self.stub = cost_service_pb2_grpc.CostServiceStub(channel)
-        logger.info("Channel client 'Cost' initialized succesfully")
+        logger.debug("Channel client 'Cost' initialized succesfully")
 
     def add(
         self,

digitalkin/services/filesystem/default_filesystem.py

@@ -1,10 +1,10 @@
 """Default filesystem."""
 
-import logging
 import os
 import tempfile
 from pathlib import Path
 
+from digitalkin.logger import logger
 from digitalkin.services.filesystem.filesystem_strategy import (
     FilesystemData,
     FilesystemServiceError,
@@ -12,8 +12,6 @@ from digitalkin.services.filesystem.filesystem_strategy import (
     FileType,
 )
 
-logger = logging.getLogger(__name__)
-
 
 class DefaultFilesystem(FilesystemStrategy):
     """Default state filesystem strategy."""
@@ -177,7 +175,7 @@ class DefaultFilesystem(FilesystemStrategy):
         try:
             os.remove(file_path)
             del self.db[name]
-            logger.info("File %s successfully deleted.", name)
+            logger.debug("File %s successfully deleted.", name)
 
         except OSError:
             msg = f"Error deleting file {name} from filesystem"

digitalkin/services/filesystem/grpc_filesystem.py

@@ -1,16 +1,21 @@
 """Grpc filesystem."""
 
-import logging
 from collections.abc import Generator
 from contextlib import contextmanager
 from typing import Any
 
-from digitalkin_proto.digitalkin.filesystem.v2 import filesystem_pb2, filesystem_service_pb2_grpc
-from digitalkin_proto.digitalkin.filesystem.v2.filesystem_pb2 import FileType as FileTypeProto
+from digitalkin_proto.digitalkin.filesystem.v2 import (
+    filesystem_pb2,
+    filesystem_service_pb2_grpc,
+)
+from digitalkin_proto.digitalkin.filesystem.v2.filesystem_pb2 import (
+    FileType as FileTypeProto,
+)
 
 from digitalkin.grpc_servers.utils.exceptions import ServerError
 from digitalkin.grpc_servers.utils.grpc_client_wrapper import GrpcClientWrapper
 from digitalkin.grpc_servers.utils.models import ClientConfig
+from digitalkin.logger import logger
 from digitalkin.services.filesystem.filesystem_strategy import (
     FilesystemData,
     FilesystemServiceError,
@@ -18,8 +23,6 @@ from digitalkin.services.filesystem.filesystem_strategy import (
     FileType,
 )
 
-logger = logging.getLogger(__name__)
-
 
 class GrpcFilesystem(FilesystemStrategy, GrpcClientWrapper):
     """Default state filesystem strategy."""

digitalkin/services/setup/__init__.py

@@ -0,0 +1 @@
+"""This module is responsible for handling the setup service."""