manta-node 0.5b0__py3-none-any.whl

manta_node/infrastructure/grpc/world_servicer.py

@@ -0,0 +1,284 @@
+from logging import getLogger
+from typing import AsyncIterator
+
+from manta_common.build.common.results import (
+    ChunkResultValues,
+    GlobalTag,
+    GlobalUpdate,
+    GlobalValue,
+    Result,
+    ResultQuery,
+)
+from manta_common.build.common.system import EnvId, Ok
+from manta_common.build.common.tasks import TaskScheduling, TaskUpdate
+from manta_common.build.node.light_service import (
+    GetGlobalRequest,
+    LightResult,
+    LightResultQuery,
+    SetGlobalRequest,
+    WorldBase,
+)
+from manta_common.build.node.node_service import StopSwarmProcess
+from manta_common.conversions import ID
+from manta_common.decorators import catch_batch_errors, catch_streaming_errors
+from manta_common.traces import Tracer
+from ...tasks import TaskProgressions
+from ..grpc.client import NodeClient
+
+__all__ = ["WorldServicer"]
+
+
+class WorldServicer(WorldBase):
+    """
+    World Servicer allows containers to access data stored in the Manager
+
+    Parameters
+    ----------
+
+        Node ID
+    node_client : NodeClient
+        Node client
+    tasks : TaskProgressions
+        Task Progressions from Task Manager
+    """
+
+    __slots__ = ["tracer", "node_client", "tasks"]
+
+    def __init__(self, node_client: NodeClient, tasks: TaskProgressions):
+        self.tracer = Tracer(getLogger(__name__))
+        self.node_client = node_client
+        self.tasks = tasks
+        super().__init__()
+
+    def task_update(self, swarm_id: ID, task_id: ID) -> TaskUpdate:
+        """
+        Convenient method to access task update from tasks
+        given a task ID
+
+        Parameters
+        ----------
+        swarm_id : PairId
+            Swarm ID
+        task_id : PairId
+            Task ID
+
+        Returns
+        -------
+        TaskUpdate
+            Task Update
+        """
+        return self.tasks[(swarm_id.tid, task_id.tid)].task_update
+
+    @catch_streaming_errors
+    async def get_task_result(
+        self, light_request: LightResultQuery
+    ) -> AsyncIterator[ChunkResultValues]:
+        """
+        Get the swarm results
+
+        Parameters
+        ----------
+        request : LightResultQuery
+            Result query from the container
+
+        Returns
+        -------
+        AsyncIterator[ChunkResultValues]
+            Result values
+        """
+        task_id = ID(light_request.task_id)
+        swarm_id = ID(light_request.swarm_id)
+        task_update = self.task_update(swarm_id, task_id)
+        iteration = task_update.iteration
+        circular = task_update.circular
+        self.tracer.set_metadata(
+            task_id=task_id,
+            swarm_id=swarm_id,
+            iteration=iteration,
+            circular=circular,
+        )
+
+        async for chunk in self.node_client.get_task_result(
+            ResultQuery(
+                swarm_id=swarm_id.oid,
+                task_id=task_id.oid,
+                node_id=self.tracer.metadata.node_id,
+                iteration=iteration,
+                tag=light_request.tag,
+                size=light_request.size,
+                method=light_request.method,
+            )
+        ):
+            self.tracer.set_metadata(
+                task_id=task_id,
+                swarm_id=swarm_id,
+                iteration=iteration,
+                circular=circular,
+            )
+            yield chunk
+
+    @catch_batch_errors
+    async def add_task_result(self, stream: AsyncIterator[LightResult]) -> Ok:
+        """
+        Set a swarm result
+
+        Parameters
+        ----------
+        request : LightResult
+            Result from the container
+
+        Returns
+        -------
+        Ok
+            Response
+        """
+
+        async def stream_generator(
+            stream: AsyncIterator[LightResult],
+        ) -> AsyncIterator[Result]:
+            async for light_request in stream:
+                task_id = ID(light_request.task_id)
+                swarm_id = ID(light_request.swarm_id)
+                task_update = self.task_update(swarm_id, task_id)
+                self.tracer.set_metadata(
+                    task_id=task_id,
+                    swarm_id=swarm_id,
+                    iteration=task_update.iteration,
+                    circular=task_update.circular,
+                )
+
+                yield Result(
+                    swarm_id=swarm_id.oid,
+                    task_id=task_id.oid,
+                    node_id=self.tracer.metadata.node_id,
+                    iteration=task_update.iteration,
+                    tag=light_request.tag,
+                    data=light_request.data,
+                )
+
+        return await self.node_client.add_task_result(stream_generator(stream))
+
+    @catch_streaming_errors
+    async def get_global(self, request: GetGlobalRequest) -> AsyncIterator[GlobalValue]:
+        """
+        Set a swarm global value
+
+        Parameters
+        ----------
+        request : GlobalTag
+            Task ID and global tag
+
+        Returns
+        -------
+        GlobalValue
+            Global value
+        """
+        task_id = ID(request.task_id)
+        swarm_id = ID(request.swarm_id)
+        task_update = self.task_update(swarm_id, task_id)
+        circular = task_update.circular
+        iteration = task_update.iteration
+
+        async for chunk in self.node_client.get_global(
+            GlobalTag(swarm_id=swarm_id.oid, tag=request.tag)
+        ):
+            self.tracer.set_metadata(
+                task_id=task_id,
+                swarm_id=swarm_id,
+                iteration=iteration,
+                circular=circular,
+            )
+            self.tracer.info(f"Received chunk for global of size {len(chunk.data)}")
+            yield chunk
+
+    @catch_batch_errors
+    async def set_global(self, stream: AsyncIterator[SetGlobalRequest]) -> Ok:
+        """
+        Set a swarm result
+
+        Parameters
+        ----------
+        request : SetGlobalRequest
+            Task ID, global tag and value
+
+        Returns
+        -------
+        Ok
+            Response
+        """
+
+        async def stream_generator(
+            stream: AsyncIterator[SetGlobalRequest],
+        ) -> AsyncIterator[GlobalUpdate]:
+            async for request in stream:
+                swarm_id = ID(request.swarm_id)
+                task_id = ID(request.task_id)
+                task_update = self.task_update(swarm_id, task_id)
+                self.tracer.set_metadata(
+                    task_id=task_id,
+                    swarm_id=swarm_id,
+                    iteration=task_update.iteration,
+                    circular=task_update.circular,
+                )
+
+                yield GlobalUpdate(
+                    swarm_id=swarm_id.oid,
+                    tag=request.tag,
+                    data=request.data,
+                )
+
+        return await self.node_client.set_global(stream_generator(stream))
+
+    @catch_batch_errors
+    async def stop_swarm(self, request: EnvId) -> Ok:
+        """
+        Stop the swarm
+
+        Parameters
+        ----------
+        request : EnvId
+            Task ID
+
+        Returns
+        -------
+        Ok
+            Response
+        """
+        swarm_id = ID(request.swarm_id)
+        task_id = ID(request.task_id)
+        task_update = self.task_update(swarm_id, task_id)
+        self.tracer.set_metadata(
+            task_id=task_id,
+            swarm_id=swarm_id,
+            iteration=task_update.iteration,
+            circular=task_update.circular,
+        )
+        return await self.node_client.stop_swarm(
+            StopSwarmProcess(swarm_id=task_update.swarm_id)
+        )
+
+    @catch_batch_errors
+    async def schedule_task(self, light_request: TaskScheduling) -> Ok:
+        """
+        Schedule a task
+
+        Parameters
+        ----------
+        request : TaskScheduling
+            Task Scheduling from container
+
+        Returns
+        -------
+        Ok
+            Response
+        """
+        swarm_id = ID(light_request.swarm_id)
+        task_id = ID(light_request.task_id)
+        task_update = self.task_update(swarm_id, task_id)
+        self.tracer.set_metadata(
+            task_id=task_id,
+            swarm_id=swarm_id,
+            iteration=task_update.iteration,
+            circular=task_update.circular,
+        )
+        return await self.node_client.schedule_task(light_request)

manta_node/infrastructure/metrics/__init__.py

@@ -0,0 +1,10 @@
+"""Metrics infrastructure - System monitoring and data collection."""
+
+from .collector import Collector
+from .metrics_collector import GenericMetricsCollector, MetricsCollectorBase
+
+__all__ = [
+    "Collector",
+    "GenericMetricsCollector",
+    "MetricsCollectorBase",
+]

manta_node/infrastructure/metrics/collector.py

@@ -0,0 +1,94 @@
+from manta_common.build.common.informations import NodeStatus, NodeStatusEnum
+from manta_common.errors import (
+    MantaError,
+    MantaGRPCError,
+    MantaMQTTError,
+    MantaNodeError,
+)
+from manta_common.traces import Tracer
+from ..filesystem.dataset_manager import DatasetManager
+from ..grpc.client import NodeClient
+from ..metrics.metrics_collector import MetricsCollectorBase
+
+
+class Collector:
+    """
+    Collector class to collect information from the node
+
+    Parameters
+    ----------
+    tracer : Tracer
+        Tracer for the collector
+    """
+
+    tracer: Tracer
+    node_client: NodeClient
+    metrics_collector: MetricsCollectorBase
+    dataset_manager: DatasetManager
+
+    async def collect_information(self, payload: bytes):
+        """
+        Send the node information to the Manager
+
+        Parameters
+        ----------
+        payload : bytes
+            Payload
+
+        Raises
+        ------
+        TaskManagerException
+            Exception when the method requested is not allowed
+        """
+        try:
+            method = payload.decode("utf-8")
+            if method == "set_node_status":
+                await self.set_node_status()
+            elif method == "set_system_summary":
+                await self.set_system_summary()
+            elif method == "set_overviews":
+                await self.set_overviews()
+            else:
+                raise MantaMQTTError(
+                    f"Not allowed method: {method}", metadata=self.tracer.metadata
+                )
+        except MantaGRPCError as e:
+            self.tracer.exception(f"Error when collecting information: {e}")
+        except MantaError as manta_error:
+            self.tracer.exception(repr(manta_error))
+            await self.node_client.send_error(manta_error)
+        except Exception as exc:
+            message = f"Error for collecting information: {repr(exc)}"
+            self.tracer.exception(message)
+            await self.node_client.send_error(
+                MantaNodeError(message, metadata=self.tracer.metadata)
+            )
+
+    async def set_node_status(self):  # DEAD CODE
+        """
+        Send the node information to the Manager
+        """
+        self.tracer.info("Sending NodeInfo to the Manager")
+        await self.node_client.set_node_status(
+            NodeStatus(status=NodeStatusEnum.CONNECTED)
+        )
+
+    async def set_system_summary(self):
+        """
+        Send the system information to the Manager
+        """
+        self.tracer.info("Sending SystemInfo to the Manager")
+        metrics = await self.metrics_collector.collect_all_metrics()
+        await self.node_client.set_system_summary(metrics)
+
+    async def set_overviews(self):
+        """
+        Send the data information to the Manager
+        """
+        self.tracer.info("Sending DataInfo to the Manager")
+
+        # Refresh the dataset metadata before sending
+        self.dataset_manager.refresh()
+
+        # Use the dataset_manager to get the overview information
+        await self.node_client.set_overviews(self.dataset_manager.to_proto())