manta-node 0.5b0__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",
+]