manta-node 0.5b0.dev9__py3-none-any.whl

manta_node/infrastructure/metrics/metrics_collector.py

@@ -0,0 +1,351 @@
+import asyncio
+import platform
+import time
+from abc import ABC, abstractmethod
+from datetime import datetime, timezone
+from logging import getLogger
+from typing import Optional
+
+import psutil
+
+from manta_common.build.common.informations import (
+    AdditionalMetrics,
+    CpuMetrics,
+    DiskUsage,
+    GpuDeviceList,
+    MemoryMetrics,
+    MetricsSnapshot,
+    NetworkUsage,
+    PlatformInfo,
+)
+from manta_common.traces import Tracer
+
+PSUTIL_AVAILABLE = True
+
+
+class MetricsCollectorBase(ABC):
+    """
+    Abstract base class for metrics collection.
+
+    This class defines the interface for metrics collectors that gather
+    system metrics like CPU, memory, and other hardware usage.
+
+    Different implementations can be created for specific hardware platforms
+    (Linux, Jetson, Mac, etc.) by extending this class.
+
+    Parameters
+    ----------
+
+        The ID of the node where metrics are being collected
+    collection_interval : float
+        Time in seconds between metrics collection (default: 30.0)
+    """
+
+    def __init__(self, collection_interval: float = 30.0):
+        self.collection_interval = collection_interval
+        self.tracer = Tracer(getLogger(__name__), None)
+        self._latest_metrics: Optional[MetricsSnapshot] = None
+        self._running = False
+        self._collection_task: Optional[asyncio.Task] = None
+        self._lock = asyncio.Lock()
+
+    @abstractmethod
+    async def collect_cpu_metrics(self) -> CpuMetrics:
+        """
+        Collect CPU usage metrics.
+
+        Returns
+        -------
+        CpuMetrics
+            Protobuf message with CPU metrics
+        """
+        pass
+
+    @abstractmethod
+    async def collect_memory_metrics(self) -> MemoryMetrics:
+        """
+        Collect memory usage metrics.
+
+        Returns
+        -------
+        MemoryMetrics
+            Protobuf message with memory metrics
+        """
+        pass
+
+    @abstractmethod
+    async def collect_gpu_metrics(self) -> GpuDeviceList:
+        """
+        Collect GPU usage metrics if available.
+
+        Returns
+        -------
+        GpuDeviceList
+            Protobuf message with GPU metrics, or empty if no GPU or metrics unavailable
+        """
+        pass
+
+    @abstractmethod
+    async def collect_additional_metrics(self) -> AdditionalMetrics:
+        """
+        Collect disk and network metrics.
+
+        Returns
+        -------
+        AdditionalMetrics
+            Protobuf message with disk and network metrics
+        """
+        pass
+
+    @classmethod
+    def get_platform_info(cls) -> PlatformInfo:
+        """
+        Get platform information (system, version, etc.).
+
+        Returns
+        -------
+        PlatformInfo
+            Protobuf message with platform information
+        """
+        return PlatformInfo(
+            system=platform.system(),
+            release=platform.release(),
+            version=platform.version(),
+            machine=platform.machine(),
+            processor=platform.processor(),
+            hostname=platform.node(),
+        )
+
+    async def collect_all_metrics(self) -> MetricsSnapshot:
+        """
+        Collect all available metrics.
+
+        Returns
+        -------
+        MetricsSnapshot
+            Protobuf message with all collected metrics
+        """
+        cpu_metrics = await self.collect_cpu_metrics()
+        memory_metrics = await self.collect_memory_metrics()
+        # gpu_metrics = await self.collect_gpu_metrics()
+        # additional_metrics = await self.collect_additional_metrics()
+        # platform_info = self.get_platform_info()
+
+        # Combine all metrics into a snapshot
+        metrics = MetricsSnapshot(
+            timestamp=datetime.now(timezone.utc),
+            cpu=cpu_metrics,
+            memory=memory_metrics,
+            # gpu=gpu_metrics,
+            # additional=additional_metrics,
+            # platform=platform_info,
+        )
+
+        # Update latest metrics with thread safety
+        async with self._lock:
+            self._latest_metrics = metrics
+
+        return metrics
+
+    async def start_collection(self):
+        """
+        Start the metrics collection background task.
+        """
+        if self._running:
+            self.tracer.warning("Metrics collection already running")
+            return
+
+        self._running = True
+        self._collection_task = asyncio.create_task(self._collection_loop())
+        self.tracer.info(
+            f"Started metrics collection (interval: {self.collection_interval}s)"
+        )
+
+    async def stop_collection(self):
+        """
+        Stop the metrics collection background task.
+        """
+        if not self._running:
+            return
+
+        self._running = False
+        if self._collection_task:
+            self._collection_task.cancel()
+            try:
+                await self._collection_task
+            except asyncio.CancelledError:
+                pass
+            self._collection_task = None
+
+        self.tracer.info("Stopped metrics collection")
+
+    async def _collection_loop(self):
+        """
+        Background loop that periodically collects metrics.
+        """
+        while self._running:
+            try:
+                await self.collect_all_metrics()
+                await asyncio.sleep(self.collection_interval)
+            except asyncio.CancelledError:
+                break
+            except Exception as e:
+                self.tracer.error(f"Error in metrics collection loop: {e}")
+                await asyncio.sleep(self.collection_interval)
+
+    async def get_latest_metrics(self) -> Optional[MetricsSnapshot]:
+        """
+        Get the most recently collected metrics.
+
+        Returns
+        -------
+        Optional[MetricsSnapshot]
+            The latest collected metrics, or None if no metrics have been collected
+        """
+        async with self._lock:
+            return self._latest_metrics
+
+
+class GenericMetricsCollector(MetricsCollectorBase):
+    """
+    Generic implementation of the metrics collector using psutil.
+
+    This implementation works on most platforms (Linux, macOS, Windows)
+    where psutil is available.
+
+    Parameters
+    ----------
+    node_id : Any
+        The ID of the node where metrics are being collected
+    collection_interval : float
+        Time in seconds between metrics collection (default: 30.0)
+    """
+
+    def __init__(self, collection_interval: float = 30.0):
+        super().__init__(collection_interval)
+
+        # psutil is now always available
+
+    async def collect_cpu_metrics(self) -> CpuMetrics:
+        """
+        Collect CPU usage metrics using psutil.
+
+        Returns
+        -------
+        CpuMetrics
+            Protobuf message with CPU metrics
+        """
+        # psutil is now always available
+
+        # Get CPU usage percentage (blocking call, so run in executor)
+        cpu_percent = await asyncio.to_thread(psutil.cpu_percent, interval=0.1)
+
+        # Get per-core CPU usage
+        per_cpu = await asyncio.to_thread(psutil.cpu_percent, interval=0.1, percpu=True)
+
+        # Get CPU times
+        cpu_times = await asyncio.to_thread(psutil.cpu_times_percent, interval=0.1)
+
+        return CpuMetrics(
+            usage_percent=cpu_percent,
+            per_cpu_percent=per_cpu,
+            user_percent=float(cpu_times.user),
+            system_percent=float(cpu_times.system),
+            idle_percent=float(cpu_times.idle),
+            count=psutil.cpu_count(logical=True) or 0,
+            physical_count=psutil.cpu_count(logical=False) or 0,
+        )
+
+    async def collect_memory_metrics(self) -> MemoryMetrics:
+        """
+        Collect memory usage metrics using psutil.
+
+        Returns
+        -------
+        MemoryMetrics
+            Protobuf message with memory metrics
+        """
+        # psutil is now always available
+
+        # Get virtual memory stats
+        memory = await asyncio.to_thread(psutil.virtual_memory)
+
+        # Get swap memory stats
+        swap = await asyncio.to_thread(psutil.swap_memory)
+
+        return MemoryMetrics(
+            total_bytes=memory.total,
+            available_bytes=memory.available,
+            used_bytes=memory.used,
+            free_bytes=memory.free,
+            percent_used=memory.percent,
+            swap_total_bytes=swap.total,
+            swap_used_bytes=swap.used,
+            swap_free_bytes=swap.free,
+            swap_percent_used=swap.percent,
+        )
+
+    async def collect_gpu_metrics(self) -> GpuDeviceList:
+        """
+        Collect GPU metrics if available.
+
+        This generic implementation doesn't collect GPU metrics.
+        Subclasses for specific platforms will override this method.
+
+        Returns
+        -------
+        GpuDeviceList
+            Empty build message (no GPU metrics in generic implementation)
+        """
+        # Generic implementation doesn't collect GPU metrics
+        # This will be overridden in platform-specific implementations
+        return GpuDeviceList(devices=[])
+
+    async def collect_additional_metrics(self) -> AdditionalMetrics:
+        """
+        Collect disk and network metrics using psutil.
+
+        Returns
+        -------
+        AdditionalMetrics
+            Protobuf message with disk and network metrics
+        """
+        # psutil is now always available
+
+        # Get disk usage
+        disk_usage = []
+        for partition in psutil.disk_partitions(all=False):
+            try:
+                usage = psutil.disk_usage(partition.mountpoint)
+                disk_usage.append(
+                    DiskUsage(
+                        mount_point=partition.mountpoint,
+                        total_bytes=usage.total,
+                        used_bytes=usage.used,
+                        free_bytes=usage.free,
+                        percent_used=usage.percent,
+                    )
+                )
+            except (PermissionError, OSError):
+                # Some mountpoints may not be accessible
+                pass
+
+        # Get basic network info
+        net_io = psutil.net_io_counters()
+        network = NetworkUsage(
+            bytes_sent=net_io.bytes_sent,
+            bytes_recv=net_io.bytes_recv,
+            packets_sent=net_io.packets_sent,
+            packets_recv=net_io.packets_recv,
+        )
+
+        # Get system uptime
+        boot_time = psutil.boot_time()
+        uptime_seconds = time.time() - boot_time
+
+        # Combine all additional metrics
+        return AdditionalMetrics(
+            disk_usage=disk_usage,
+            network=network,
+            uptime_seconds=int(uptime_seconds),
+        )

manta_node/infrastructure/mqtt/__init__.py

@@ -0,0 +1,7 @@
+"""MQTT infrastructure - Message broker communication."""
+
+from .command_handler import CommandHandler
+
+__all__ = [
+    "CommandHandler",
+]

manta_node/infrastructure/mqtt/command_handler.py

@@ -0,0 +1,315 @@
+import asyncio
+import io
+from datetime import datetime
+from logging import getLogger
+from pathlib import Path
+from typing import AsyncIterable, Awaitable, Callable, List, Optional, cast
+
+import aiomqtt
+
+from manta_common.base_mqtt import MqttBase
+from manta_common.build.common.mqtt import MqMetricsSnapshot, MqTaskUpdate
+from manta_common.build.common.system import Logs
+from manta_common.build.common.tasks import TaskStatus, TaskUpdate
+from manta_common.const import CHUNK_SIZE
+from manta_common.conversions import ID
+from manta_common.errors import (
+    MantaError,
+    MantaGRPCError,
+    MantaMQTTError,
+    MantaNodeError,
+)
+from manta_common.traces import Tracer
+from ...task_manager import TaskManager
+from ..filesystem.dataset_manager import DatasetManager
+from ..grpc.client import NodeClient
+from ..metrics.collector import Collector
+from ..metrics.metrics_collector import MetricsCollectorBase
+
+TIMEOUT = 1  # s
+
+
+class CommandHandler(TaskManager, Collector, MqttBase):
+    """
+    CommandHandler class to handle commands from the Manager
+
+    Parameters
+    ----------
+    node_id: ID
+        Node ID
+    mqtt_client: MQTTClient
+        MQTT client
+    node_client: NodeClient
+        Node client
+    mqtt_broker_host : str
+        MQTT broker host
+    mqtt_broker_port : int
+        MQTT broker port
+    topics : List[str]
+        List of topics
+    light_service_port: int
+        Light service port
+    dataset_manager: DatasetManager
+        Dataset manager
+    metrics_collector: MetricsCollectorBase
+        Metrics collector
+    stop_callback: Callable[[], Awaitable[None]]
+        Stop callback
+    ssl_cert_folder : Optional[Path]
+        SSL Certification folder
+    """
+
+    def __init__(
+        self,
+        node_client: NodeClient,
+        mqtt_broker_host: str,
+        mqtt_broker_port: int,
+        light_service_port: int,
+        topics: List[str],
+        dataset_manager: DatasetManager,
+        metrics_collector: MetricsCollectorBase,
+        stop_callback: Callable[[], Awaitable[None]],
+        ssl_cert_folder: Optional[Path] = None,
+        ping_interval: int = TIMEOUT,
+        max_concurrent_tasks: int = 2,
+        task_timeout: int = 3600,
+    ) -> None:
+        self.tracer = Tracer(getLogger(__name__))
+        self.node_client = node_client
+        self.dataset_manager = dataset_manager
+        self.metrics_collector = metrics_collector
+        self.stop_callback = stop_callback
+        self.ping_interval = ping_interval
+        self.max_concurrent_tasks = max_concurrent_tasks
+        self.task_timeout = task_timeout
+
+        MqttBase.__init__(
+            self,
+            mqtt_broker_host=mqtt_broker_host,
+            mqtt_broker_port=mqtt_broker_port,
+            topics=topics,
+            client_id=ID(node_client.metadata.get("authorization", "node_client")),
+            ssl_cert_folder=ssl_cert_folder,
+        )
+
+        self.tracer.debug("Trying to connect to docker daemon...")
+        TaskManager.__init__(self, light_service_port)
+        self.tracer.info("Connected to docker daemon !")
+
+    async def clean(self):  # TODO : use __del__
+        """Delete the temporary folder and stop all containers"""
+        self.tracer.debug("Cleaning up the temporary folder")
+        self.temp_folder.cleanup()
+
+        self.tracer.info("Stopping all tasks")
+        for task_progression in self.tasks:
+            try:
+                task_id = task_progression.task_id
+                swarm_id = task_progression.swarm_id
+
+                task_progression.unlink()
+                await self.stop_container(swarm_id, task_id)
+                await self.remove_container(swarm_id, task_id)
+            except Exception:
+                self.tracer.exception(
+                    f"Error when stopping task {task_progression.task_id}"
+                )
+
+        self.tracer.info("Closing Docker client")
+        try:
+            await self.async_client.close()
+        except Exception:
+            self.tracer.exception("Error when closing Docker client")
+
+    async def loop_tasks(self):
+        """
+        Processes messages or check tasks to ping
+        """
+        while True:
+            # Wait until we either (1) receive a message or (2) publish a message
+            try:
+                subscribe_task = asyncio.create_task(self.client.messages.__anext__())
+                publish_task = asyncio.create_task(self.queue.get())
+                ping_task = asyncio.create_task(self._ping())
+                done, _ = await asyncio.wait(
+                    (subscribe_task, publish_task, ping_task),
+                    return_when=asyncio.FIRST_COMPLETED,
+                )
+            # If the asyncio.wait is cancelled, we must also cancel the queue task
+            except asyncio.CancelledError:
+                subscribe_task.cancel()
+                publish_task.cancel()
+                ping_task.cancel()
+                raise
+
+            # When we receive a message, process it
+            if subscribe_task in done and publish_task in done:
+                await self._subscribe_from_task(subscribe_task)
+                await self._publish_from_task(publish_task)
+                ping_task.cancel()
+            elif subscribe_task in done:
+                await self._subscribe_from_task(subscribe_task)
+                publish_task.cancel()
+                ping_task.cancel()
+            # Publish the message to its topic with specific QoS
+            elif publish_task in done:
+                await self._publish_from_task(publish_task)
+                subscribe_task.cancel()
+                ping_task.cancel()
+            # If we disconnect from the broker, stop the loop with an exception
+            elif ping_task in done:
+                ping_task.cancel()
+                publish_task.cancel()
+                subscribe_task.cancel()
+            else:
+                self.tracer.warning("Disconnected during message iteration")
+                publish_task.cancel()
+                subscribe_task.cancel()
+                ping_task.cancel()
+                raise aiomqtt.MqttError("Disconnected during message iteration")
+
+    async def _ping(self):
+        """
+        Ping the Manager
+        """
+        await asyncio.sleep(self.ping_interval)
+        updated = await self.check_tasks()
+        if len(updated) > 0:
+            for task_progression in updated:
+                if task_progression.task_status == TaskStatus.STOPPED:
+                    try:
+                        date, processed_logs = await self.get_task_logs(
+                            task_progression
+                        )
+                        self.tracer.info(
+                            f"Sending logs to the Manager for task {task_progression.task_id}"
+                        )
+                        await self.send_task_logs(date, processed_logs)
+                    except MantaGRPCError as e:
+                        self.tracer.exception(
+                            f"gRPCError when updating finished tasks: {e}"
+                        )
+                        task_progression.task_status = TaskStatus.FAILED
+                    except MantaError as manta_error:
+                        self.tracer.exception("Error updating finished tasks.")
+                        manta_error.metadata = self.tracer.metadata
+                        task_progression.task_status = TaskStatus.FAILED
+                        await self.node_client.send_error(manta_error)
+                    except Exception as exc:
+                        message = f"Error updating finished tasks: {repr(exc)}"
+                        self.tracer.exception(message)
+                        task_progression.task_status = TaskStatus.FAILED
+                        await self.node_client.send_error(
+                            MantaNodeError(message, metadata=self.tracer.metadata)
+                        )
+                else:
+                    try:
+                        await self.publish_task_update(task_progression.task_update)
+                    except Exception as exc:
+                        message = f"Error updating active tasks: {repr(exc)}"
+                        self.tracer.exception("Error updating active tasks.")
+                        await self.node_client.send_error(
+                            MantaNodeError(message, metadata=self.tracer.metadata)
+                        )
+        else:
+            await self.ping()
+
+    async def ping(self):
+        """
+        Ping the Manager
+        """
+        try:
+            metrics = await self.metrics_collector.collect_all_metrics()
+            await self.publish_message(
+                "manager/set_system_summary",
+                MqMetricsSnapshot(
+                    token=self.node_client.secured_token, snapshot=metrics
+                ),
+            )
+        except MantaGRPCError as e:
+            self.tracer.exception(f"Error when pinging the Manager: {e}")
+        except MantaError as e:
+            self.tracer.exception(f"Error when pinging the Manager: {e}")
+            await self.node_client.send_error(e)
+
+    async def process_message(self, message: aiomqtt.Message):
+        """
+        Handle messages from MQTT
+
+        Parameters
+        ----------
+        message : aiomqtt.Message
+            Message
+        """
+        try:
+            self.tracer.clean_metadata()
+            key = str(message.topic).split("/")[-1]
+            payload = cast(bytes, message.payload)
+
+            if key == "deploy_task":
+                await self.deploy_task(payload)
+            elif key == "stop_task":
+                task_update = await self.stop_task(payload)
+                if task_update is not None:
+                    await self.publish_task_update(task_update)
+            elif key == "stop_swarm":
+                await self.stop_swarm(payload)
+            elif key == "collect_information":
+                await self.collect_information(payload)
+            elif key == "stop_node":
+                await self.stop_node(payload)
+            else:
+                self.tracer.error(f"Not allowed method: {key}")
+                raise MantaMQTTError(
+                    f"Not allowed method: {key}", metadata=self.tracer.metadata
+                )
+        except MantaGRPCError as e:
+            self.tracer.exception(f"Error sending message to the Manager: {e}")
+        except MantaError as manta_error:
+            self.tracer.exception(f"Error when handling message: {manta_error}")
+            await self.node_client.send_error(manta_error)
+
+    async def publish_task_update(self, task_update: TaskUpdate):
+        """
+        Publish a task update to the Manager through MQTT
+        Message broker
+
+        Parameters
+        ----------
+        task_update : TaskUpdate
+            Task Update
+        """
+        await self.publish_message(
+            "manager/update_task_status",
+            MqTaskUpdate(token=self.node_client.secured_token, update=task_update),
+        )
+
+    async def send_task_logs(self, date: datetime, processed_logs: bytes):
+        """
+        Send task logs to the Manager
+
+        Parameters
+        ----------
+        date : datetime
+            Date when logs where created
+        processed_logs : bytes
+            Logs prepared from tracer
+        """
+
+        async def chunk_logs(content: bytes) -> AsyncIterable[Logs]:
+            data_stream = io.BytesIO(content)
+            while chunk := data_stream.read(CHUNK_SIZE):
+                yield Logs(content=chunk, timestamp=date)
+
+        await self.node_client.add_task_logs(chunk_logs(processed_logs))
+
+    async def stop_node(self, _payload: bytes):
+        """
+        Stop the node and prepare for graceful shutdown
+
+        Parameters
+        ----------
+        payload : bytes
+            Node ID to stop
+        """
+        await self.stop_callback()