benchmarks/system_info.py

"""
System Information Collection Module
收集系统基础信息：操作系统、CPU、内存、磁盘、GPU
"""

import platform
import os
import psutil
import cpuinfo
import time
from typing import Dict, Any, Optional, List


def get_os_info() -> Dict[str, str]:
    """获取操作系统信息"""
    try:
        # 尝试读取 /etc/os-release
        os_release = {}
        if os.path.exists('/etc/os-release'):
            with open('/etc/os-release', 'r') as f:
                for line in f:
                    if '=' in line:
                        key, value = line.strip().split('=', 1)
                        os_release[key] = value.strip('"')

        # 获取系统负载
        load_avg = [round(x, 2) for x in os.getloadavg()]
        
        # 获取启动时间
        boot_time = psutil.boot_time()
        uptime = time.time() - boot_time
        
        return {
            "system": platform.system(),
            "distro": os_release.get('PRETTY_NAME', platform.platform()),
            "kernel": platform.release(),
            "hostname": platform.node(),
            "architecture": platform.machine(),
            "load_average": load_avg,
            "uptime_seconds": int(uptime),
        }
    except Exception as e:
        return {"error": str(e)}


def get_cpu_info() -> Dict[str, Any]:
    """获取CPU详细信息"""
    try:
        info = cpuinfo.get_cpu_info()
        
        # 获取频率信息
        freq = psutil.cpu_freq()
        freq_info = {
            "current_mhz": round(freq.current, 2) if freq else None,
            "min_mhz": round(freq.min, 2) if freq and freq.min else None,
            "max_mhz": round(freq.max, 2) if freq and freq.max else None,
        }
        
        # 获取缓存信息
        cache_info = {}
        for key in ['l1_data_cache_size', 'l1_instruction_cache_size', 'l2_cache_size', 'l3_cache_size']:
            if key in info:
                cache_info[key] = info[key]
        
        return {
            "brand": info.get('brand_raw', 'Unknown'),
            "arch": info.get('arch', platform.machine()),
            "bits": info.get('bits', 64),
            "cores_physical": psutil.cpu_count(logical=False),
            "cores_logical": psutil.cpu_count(logical=True),
            "frequency": freq_info,
            "cache": cache_info,
            "flags": info.get('flags', [])[:20],  # 只取前20个特性
            "stats": {
                "ctx_switches": psutil.cpu_stats().ctx_switches,
                "interrupts": psutil.cpu_stats().interrupts,
                "soft_interrupts": psutil.cpu_stats().soft_interrupts,
            }
        }
    except Exception as e:
        return {"error": str(e)}


def get_memory_info() -> Dict[str, Any]:
    """获取内存信息"""
    try:
        mem = psutil.virtual_memory()
        swap = psutil.swap_memory()
        
        def format_bytes(b: int) -> str:
            """格式化字节数"""
            for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
                if b < 1024:
                    return f"{b:.2f} {unit}"
                b /= 1024
            return f"{b:.2f} PB"
        
        return {
            "total": mem.total,
            "total_formatted": format_bytes(mem.total),
            "available": mem.available,
            "available_formatted": format_bytes(mem.available),
            "used": mem.used,
            "used_formatted": format_bytes(mem.used),
            "percent": mem.percent,
            "swap_total": swap.total,
            "swap_total_formatted": format_bytes(swap.total),
            "swap_used": swap.used,
            "swap_percent": swap.percent,
        }
    except Exception as e:
        return {"error": str(e)}


def get_disk_info() -> List[Dict[str, Any]]:
    """获取磁盘信息"""
    try:
        disks = []
        
        def format_bytes(b: int) -> str:
            for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
                if b < 1024:
                    return f"{b:.2f} {unit}"
                b /= 1024
            return f"{b:.2f} PB"
        
        for partition in psutil.disk_partitions():
            try:
                usage = psutil.disk_usage(partition.mountpoint)
                disks.append({
                    "device": partition.device,
                    "mountpoint": partition.mountpoint,
                    "fstype": partition.fstype,
                    "total": usage.total,
                    "total_formatted": format_bytes(usage.total),
                    "used": usage.used,
                    "used_formatted": format_bytes(usage.used),
                    "free": usage.free,
                    "free_formatted": format_bytes(usage.free),
                    "percent": usage.percent,
                })
            except PermissionError:
                continue
        
        return disks
    except Exception as e:
        return [{"error": str(e)}]


def get_gpu_info() -> Optional[List[Dict[str, Any]]]:
    """获取GPU信息（如果有NVIDIA GPU）"""
    try:
        import subprocess
        
        # 尝试运行 nvidia-smi
        result = subprocess.run(
            ['nvidia-smi', '--query-gpu=name,memory.total,memory.free,memory.used,driver_version,compute_cap', 
             '--format=csv,noheader,nounits'],
            capture_output=True,
            text=True,
            timeout=10
        )
        
        if result.returncode != 0:
            return None
        
        gpus = []
        for i, line in enumerate(result.stdout.strip().split('\n')):
            if line:
                parts = [p.strip() for p in line.split(',')]
                if len(parts) >= 6:
                    gpus.append({
                        "index": i,
                        "name": parts[0],
                        "memory_total_mb": int(parts[1]),
                        "memory_free_mb": int(parts[2]),
                        "memory_used_mb": int(parts[3]),
                        "driver_version": parts[4],
                        "compute_capability": parts[5],
                    })
        
        return gpus if gpus else None
        
    except (FileNotFoundError, subprocess.TimeoutExpired, Exception):
        return None


def get_all_system_info() -> Dict[str, Any]:
    """获取所有系统信息"""
    return {
        "os": get_os_info(),
        "cpu": get_cpu_info(),
        "memory": get_memory_info(),
        "disk": get_disk_info(),
        "gpu": get_gpu_info(),
    }