digitalkin 0.2.25rc0__py3-none-any.whl → 0.3.2.dev14__py3-none-any.whl

digitalkin/modules/job_manager/taskiq_job_manager.py

@@ -1,290 +0,0 @@
-"""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
-
-    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: SetupModelT,
-        mission_id: str,
-        setup_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_id: The setup ID associated with the module.
-            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_id,
-            setup_version_id,
-            self.module_class,
-            self.services_mode,
-            config_setup_data.model_dump(),  # type: ignore
-        )
-
-        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_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_id: The setup ID associated with 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_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-0.2.25rc0.dist-info/RECORD

@@ -1,89 +0,0 @@
-base_server/__init__.py,sha256=gs8t9Dg0dNVHRdYYbEQY8bn8tzEUv2zE6eBhKNPG3kU,88
-base_server/server_async_insecure.py,sha256=Rvj5Xj8tDMglAj6aOGlQdMeD-FL97_OBEhHYHxxQrVo,3887
-base_server/server_async_secure.py,sha256=SV_CqEmD6YNHHP2Y369T2PLPJ-9JG3bvXytchX_Ensk,4684
-base_server/server_sync_insecure.py,sha256=VgSH8YghagK3fiwHhc7d6__zjN6lj_FgBUponFbOoxM,3127
-base_server/server_sync_secure.py,sha256=MMkq67vAZNiDLJSySjoXtHIEIK1pAlRz57n75Egnndk,3939
-base_server/mock/__init__.py,sha256=YZFT-F1l_TpvJYuIPX-7kTeE1CfOjhx9YmNRXVoi-jQ,143
-base_server/mock/mock_pb2.py,sha256=sETakcS3PAAm4E-hTCV1jIVaQTPEAIoVVHupB8Z_k7Y,1843
-base_server/mock/mock_pb2_grpc.py,sha256=BbOT70H6q3laKgkHfOx1QdfmCS_HxCY4wCOX84YAdG4,3180
-digitalkin/__init__.py,sha256=7LLBAba0th-3SGqcpqFO-lopWdUkVLKzLZiMtB-mW3M,162
-digitalkin/__version__.py,sha256=LHqr5g2PwW2SCnat77bHfB45UQ6a8RBjYaxJyzdJzPo,194
-digitalkin/logger.py,sha256=wjJnRzA9p6xFGBtY1A32hlHixRBhMzz-tAx5Hoz5lRs,3347
-digitalkin/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-digitalkin/grpc_servers/__init__.py,sha256=0cJBlwipSmFdXkyH3T0i6OJ1WpAtNsZgYX7JaSnkbtg,804
-digitalkin/grpc_servers/_base_server.py,sha256=NXnnZPjJqUDWoumhEbb7EOEWB7d8XYwpQs-l97NTe4k,18647
-digitalkin/grpc_servers/module_server.py,sha256=81l1i32F7XGQDgBXjQ7Dpukh-KTpmHzM5PKka8LjviE,10386
-digitalkin/grpc_servers/module_servicer.py,sha256=hhUAH6K5OcvkQEbXQUeUTAHJxkus-KlVwJnG1bGHTo4,18963
-digitalkin/grpc_servers/registry_server.py,sha256=StY18DKYoPKQIU1SIzgito6D4_QA1aMVddZ8O2WGlHY,2223
-digitalkin/grpc_servers/registry_servicer.py,sha256=dqsKGHZ0LnaIvGt4ipaAuigd37sbJBndT4MAT029GsY,16471
-digitalkin/grpc_servers/utils/exceptions.py,sha256=SyOgvjggaUECYmSiqy8KJLHwHVt5IClSTxslHM-IZzI,931
-digitalkin/grpc_servers/utils/factory.py,sha256=jm6rFjiqmtSv7BIHNAOxsG9xXtSvWpx9TfzSQiX97MQ,5899
-digitalkin/grpc_servers/utils/grpc_client_wrapper.py,sha256=iPHwVAaGeuV0ZW7YhTc7E3SWPyWX244Cy0iFf2DSf_Q,2685
-digitalkin/grpc_servers/utils/models.py,sha256=hiwiGHy3sEBBFdOUBCoC1n_7GVeJhjRmi0sl4tCEJy0,8965
-digitalkin/grpc_servers/utils/types.py,sha256=rQ78s4nAet2jy-NIDj_PUWriT0kuGHr_w6ELjmjgBao,539
-digitalkin/models/__init__.py,sha256=hDHtUfswaNh8wo4NZaBItg9JqC0uNSRqXArNWSrGynY,163
-digitalkin/models/module/__init__.py,sha256=fgTVbsNmLZgM43Vy-Ea5-g0EG3pncmAmJl9nk-OQdzo,561
-digitalkin/models/module/module.py,sha256=NWOwCeD_NH0NL89AF_NJ0SwE7G1n-HTcRalodLCDmpM,749
-digitalkin/models/module/module_context.py,sha256=AL6bzwjYfYPLHb1Hu5kKBjaIOxCpBCPMvxEnapdyrX4,1003
-digitalkin/models/module/module_types.py,sha256=IygT4MnS-ihI3s0oaPf4DibWcKlq7hvQCjgPc9JAH6g,3511
-digitalkin/models/services/__init__.py,sha256=HsW7MUGFPvH7Ri28WN4BHHBfEQk5dzU_9FOWAc-0lSE,56
-digitalkin/models/services/cost.py,sha256=QTEuFD6xz62nob0z4ksE-INJWcZ-iFiuNW5mvXhpFes,1599
-digitalkin/models/services/storage.py,sha256=cYTVIriGKiprF9OerhSxmc_jM6fUTVwmeon1yQCinkE,143
-digitalkin/modules/__init__.py,sha256=VwVbKok81NGyPIBZgEj_SR-59G8tTlSb4eBJI9W6Vx4,281
-digitalkin/modules/_base_module.py,sha256=uZ19kOgF-z5iEXjtfp0mBlex5XnifvuxFc7_Gj810qE,17365
-digitalkin/modules/archetype_module.py,sha256=lOe3yYufwfylZR_VGy1w-zqdqVaMI_JANfKkbH9eODE,471
-digitalkin/modules/tool_module.py,sha256=h9oo2vrFJdiBLW2qL_m79a9hEzfXV0L6bZd4KD5OSQs,422
-digitalkin/modules/trigger_handler.py,sha256=TpxFc_xkYia9B9rAvHkj02e818W08JOfIqssQqbsCpM,1785
-digitalkin/modules/job_manager/base_job_manager.py,sha256=4qrCw68ZAleZp8vnpd1qYH320aQO0ApH2cZrKM87Grk,6121
-digitalkin/modules/job_manager/job_manager_models.py,sha256=onHy-DfInLZQveniMIWIKwTKSQjojz500JvHB54x93c,1129
-digitalkin/modules/job_manager/single_job_manager.py,sha256=nMC1BX6G62Ehmam7aVSAlKW9ShJDB1ntuMq7KbD_sSI,10564
-digitalkin/modules/job_manager/taskiq_broker.py,sha256=4q3U03SNlb1ypup0VOa73KHLKJ9peuZUPCR5eFLMNAk,7498
-digitalkin/modules/job_manager/taskiq_job_manager.py,sha256=JxvNS95J_kEuRP08PvC3pQIO3DL8V1GRv7u6WUG59xg,10254
-digitalkin/services/__init__.py,sha256=LqGk_5DJy8Bzz62ajIq9jCeYNKQUIgtSCpafZk15FLc,910
-digitalkin/services/base_strategy.py,sha256=yA9KUJGRKuuaxA6l3GcMv8zKfWoIsW03UxJT80Yea2I,766
-digitalkin/services/services_config.py,sha256=JnyzZcG7OYBelwgn-wdVgY2n3yFTEkwLPHRZB8Tjw10,7468
-digitalkin/services/services_models.py,sha256=5zXkWcfKnXGwQi9sN4OAL3XrgqOcmsTl8ai5Mi4RPsw,1668
-digitalkin/services/agent/__init__.py,sha256=vJc8JN0pdtA8ecypLBeHrwAUIW6H2C8NyW-dk24rTpk,244
-digitalkin/services/agent/agent_strategy.py,sha256=42Q9RciHX6tg3CgDQkbrlIx4h_TX0WIuSpLmCjitVmA,492
-digitalkin/services/agent/default_agent.py,sha256=4N_E_eQxJGOx1KVUUg5jNOje-3ncMxF3ePB-uDuGrJc,345
-digitalkin/services/cost/__init__.py,sha256=Wi9ZB4LSXFsUYgkX-V1UJQkVXYDNDpp8q2dXccR2uRM,303
-digitalkin/services/cost/cost_strategy.py,sha256=MpPX33P_S5b2by6F4zT-rcyeRuh2V4NYPZe05VpDOGQ,2649
-digitalkin/services/cost/default_cost.py,sha256=XE7kNFde8NmbulU9m1lc3mi-vHFkbaJf0XHUc0D4UHE,3945
-digitalkin/services/cost/grpc_cost.py,sha256=cGtb0atPXSEEOrNIWee-o3ScfNRSAFXJGDu0vcWT6zg,6295
-digitalkin/services/filesystem/__init__.py,sha256=BhwMl_BUvM0d65fmglkp0SVwn3RfYiUOKJgIMnOCaGM,381
-digitalkin/services/filesystem/default_filesystem.py,sha256=qfC4jorfo86fBBnth-4ft13VuaGdvIHgWOVurCykXIk,15182
-digitalkin/services/filesystem/filesystem_strategy.py,sha256=zibVLvX_IBQ-kgh-KYzHdszDeiHFPEAZszu_k99x1GQ,9487
-digitalkin/services/filesystem/grpc_filesystem.py,sha256=ilxTFjelmf8_bVSs3xLlsyWFHVlSJf27c4j8H7R1Rkw,12987
-digitalkin/services/identity/__init__.py,sha256=InkeyLgFYYwItx8mePA8HpfacOMWZwwuc0G4pWtKq9s,270
-digitalkin/services/identity/default_identity.py,sha256=Y2auZHrGSZTIN5D8HyjLvLcNbYFM1CNUE23x7p5VIGw,386
-digitalkin/services/identity/identity_strategy.py,sha256=skappBbds1_qa0Gr24FGrNX1N0_OYhYT1Lh7dUaAirE,429
-digitalkin/services/registry/__init__.py,sha256=Zl4QAkCe9tOmmKGBWVuLQVFepdZiL0ec3VDj27IeyYM,270
-digitalkin/services/registry/default_registry.py,sha256=VnWkF6nHpFxUKuUbZLPqzXqdA6oXmyV_ySpeuOCf_ko,277
-digitalkin/services/registry/registry_strategy.py,sha256=uBXgZIv25jeXbeVO8vWvlNPxxNYu7_KiCw2PoE6AWr8,423
-digitalkin/services/setup/__init__.py,sha256=t6xcvEWqTbcRZstBFK9cESEqaZKvpW14VtYygxIqfYQ,65
-digitalkin/services/setup/default_setup.py,sha256=9VM3KwsuQcFQQ08RoOHWOE_-9BsRW0YGRtDWYTbQGdA,8246
-digitalkin/services/setup/grpc_setup.py,sha256=qjdED3HYZtpgeYYpiIwBWwPgB1NtoAuKmB5unNNRnVQ,12480
-digitalkin/services/setup/setup_strategy.py,sha256=ZnJ_HwWCkHCPrqKekSD5L9y3p8wMwfjQ8sj2hLZq6go,4004
-digitalkin/services/snapshot/__init__.py,sha256=Uzlnzo0CYlSpVsdiI37hW7xQk8hu3YA1fOI6O6MSzB0,270
-digitalkin/services/snapshot/default_snapshot.py,sha256=Mb8QwWRsHh9I_tN0ln_ZiFa1QCZxOVWmuVLemQOTWpc,1058
-digitalkin/services/snapshot/snapshot_strategy.py,sha256=B1TU3V_k9A-OdqBkdyc41-ihnrW5Btcwd1KyQdHT46A,898
-digitalkin/services/storage/__init__.py,sha256=T-ocYLLphudkQgzvG47jBOm5GQsRFRIGA88y7Ur4akg,341
-digitalkin/services/storage/default_storage.py,sha256=D8e-UYUkb2GvDEHMWcN3EkcIKXWA8DrsaQsXVjoXAYQ,7975
-digitalkin/services/storage/grpc_storage.py,sha256=3ZHGq3wnbF8dZjJcYN1EGTEUXA4rur0XEGtTBwAcGB4,7206
-digitalkin/services/storage/storage_strategy.py,sha256=sERF5tIJnzpb1iNqTXic9xRkGaXMifo6kb709ubB-Yo,8848
-digitalkin/utils/__init__.py,sha256=sJnY-ZUgsjMfojAjONC1VN14mhgIDnzyOlGkw21rRnM,28
-digitalkin/utils/arg_parser.py,sha256=nvjI1pKDY1HfS0oGcMQPtdTQcggXLtpxXMbnMxNEKRU,3109
-digitalkin/utils/development_mode_action.py,sha256=TqRuAF_A7bDD4twRB4PnZcRoNeaiAnEdxM5kvy4aoaA,1511
-digitalkin/utils/llm_ready_schema.py,sha256=JjMug_lrQllqFoanaC091VgOqwAd-_YzcpqFlS7p778,2375
-digitalkin/utils/package_discover.py,sha256=3e9-6Vf3yAAv2VkHHVK4QVqHJBxQqg3d8uuDTsXph24,13471
-digitalkin-0.2.25rc0.dist-info/licenses/LICENSE,sha256=Ies4HFv2r2hzDRakJYxk3Y60uDFLiG-orIgeTpstnIo,20327
-modules/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-modules/cpu_intensive_module.py,sha256=ejB9XPnFfA0uCuFUQbM3fy5UYfqqAlF36rv_P5Ri8ho,8363
-modules/minimal_llm_module.py,sha256=Ijld__ZnhzfLwpXD1XVkLZ7jyKZKyOFZczOpiPttJZc,11216
-modules/text_transform_module.py,sha256=bwPSnEUthZQyfLwcTLo52iAxItAoknkLh8Y3m5aywaY,7251
-services/filesystem_module.py,sha256=71Mcja8jCQqiqFHPdsIXplFIHTvgkxRhp0TRXuCfgkk,7430
-services/storage_module.py,sha256=ybTMqmvGaTrR8PqJ4FU0cwxaDjT36TskVrGoetTGmno,6955
-digitalkin-0.2.25rc0.dist-info/METADATA,sha256=gLkfxQPu5ATtTX7MPvBshfvEdh8hnh9ZRIIsc5YIsZA,30582
-digitalkin-0.2.25rc0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-digitalkin-0.2.25rc0.dist-info/top_level.txt,sha256=gcjqlyrZuLjIyxrOIavCQM_olpr6ND5kPKkZd2j0xGo,40
-digitalkin-0.2.25rc0.dist-info/RECORD,,