core-api/src/domains/tools/system/service.py

"""
System stats service for collecting host system metrics
"""
import subprocess
import socket
from datetime import datetime, timezone

import psutil

from src.shared.logging import get_logger
from typing import List

from src.domains.tools.system.schemas import (
    SystemStatsResponse,
    CpuStats,
    MemoryStats,
    DiskStats,
    NetworkStats,
    GpuStats,
)

# Filesystem types to exclude (virtual/system filesystems)
EXCLUDED_FSTYPES = {
    "tmpfs", "devtmpfs", "devfs", "squashfs", "overlay",
    "aufs", "proc", "sysfs", "cgroup", "cgroup2",
    "debugfs", "tracefs", "securityfs", "pstore",
    "hugetlbfs", "mqueue", "binfmt_misc", "autofs",
    "fuse.lxcfs", "nsfs", "efivarfs",
}

logger = get_logger(__name__)


class SystemStatsService:
    """Service for collecting host system statistics"""

    async def get_stats(self) -> SystemStatsResponse:
        """
        Collect current system statistics.

        Returns:
            SystemStatsResponse with CPU, memory, disks, network, and GPU stats
        """
        cpu = self._get_cpu_stats()
        memory = self._get_memory_stats()
        disks = self._get_all_disk_stats()
        network = self._get_network_stats()
        gpu = self._get_gpu_stats()

        return SystemStatsResponse(
            cpu=cpu,
            memory=memory,
            disks=disks,
            network=network,
            gpu=gpu,
            hostname=socket.gethostname(),
            queried_at=datetime.now(timezone.utc),
        )

    def _get_cpu_stats(self) -> CpuStats:
        """Get CPU usage statistics"""
        # Get CPU percentage (blocking call with interval for accuracy)
        cpu_percent = psutil.cpu_percent(interval=0.1)
        cpu_count = psutil.cpu_count()

        # Get load averages (Unix only)
        try:
            load_1, load_5, load_15 = psutil.getloadavg()
        except (AttributeError, OSError):
            load_1 = load_5 = load_15 = None

        return CpuStats(
            usage_percent=cpu_percent,
            cores=cpu_count or 1,
            load_1m=load_1,
            load_5m=load_5,
            load_15m=load_15,
        )

    def _get_memory_stats(self) -> MemoryStats:
        """Get memory usage statistics"""
        mem = psutil.virtual_memory()

        return MemoryStats(
            usage_percent=mem.percent,
            total_bytes=mem.total,
            used_bytes=mem.used,
            available_bytes=mem.available,
        )

    def _get_all_disk_stats(self) -> List[DiskStats]:
        """Get disk usage statistics for all mounted real filesystems"""
        disks = []
        seen_devices = set()

        for partition in psutil.disk_partitions(all=False):
            # Skip excluded filesystem types
            if partition.fstype.lower() in EXCLUDED_FSTYPES:
                continue

            # Skip duplicate devices (same device mounted multiple times)
            if partition.device in seen_devices:
                continue
            seen_devices.add(partition.device)

            # Skip Docker/container overlays
            if partition.mountpoint.startswith("/var/lib/docker"):
                continue

            try:
                usage = psutil.disk_usage(partition.mountpoint)
                disks.append(DiskStats(
                    mount_point=partition.mountpoint,
                    device=partition.device,
                    fstype=partition.fstype,
                    usage_percent=usage.percent,
                    total_bytes=usage.total,
                    used_bytes=usage.used,
                    free_bytes=usage.free,
                ))
            except (PermissionError, OSError) as e:
                logger.debug(f"Skipping {partition.mountpoint}: {e}")
                continue

        # Sort by mount point for consistent ordering
        disks.sort(key=lambda d: d.mount_point)

        return disks

    def _get_network_stats(self) -> NetworkStats:
        """Get network I/O statistics"""
        net_io = psutil.net_io_counters()

        return NetworkStats(
            bytes_sent=net_io.bytes_sent,
            bytes_recv=net_io.bytes_recv,
            bytes_total=net_io.bytes_sent + net_io.bytes_recv,
        )

    def _get_gpu_stats(self) -> GpuStats:
        """Get GPU/VRAM statistics using nvidia-smi"""
        try:
            # Query nvidia-smi for GPU memory info
            result = subprocess.run(
                [
                    "nvidia-smi",
                    "--query-gpu=name,memory.total,memory.used,memory.free",
                    "--format=csv,noheader,nounits",
                ],
                capture_output=True,
                text=True,
                timeout=5,
            )

            if result.returncode != 0:
                logger.debug("nvidia-smi not available or failed")
                return GpuStats(available=False)

            # Parse output: "NVIDIA GeForce RTX 3080, 10240, 2048, 8192"
            line = result.stdout.strip().split("\n")[0]  # First GPU
            parts = [p.strip() for p in line.split(",")]

            if len(parts) >= 4:
                name = parts[0]
                total_mb = int(parts[1])
                used_mb = int(parts[2])
                free_mb = int(parts[3])

                total_bytes = total_mb * 1024 * 1024
                used_bytes = used_mb * 1024 * 1024
                free_bytes = free_mb * 1024 * 1024
                usage_percent = (used_mb / total_mb * 100) if total_mb > 0 else 0

                return GpuStats(
                    available=True,
                    name=name,
                    usage_percent=round(usage_percent, 1),
                    total_bytes=total_bytes,
                    used_bytes=used_bytes,
                    free_bytes=free_bytes,
                )

        except FileNotFoundError:
            logger.debug("nvidia-smi not found - no NVIDIA GPU available")
        except subprocess.TimeoutExpired:
            logger.warning("nvidia-smi timed out")
        except Exception as e:
            logger.warning(f"Failed to get GPU stats: {e}")

        return GpuStats(available=False)