local_swarm/src/hardware/detector.py

"""Hardware detection module for Local Swarm."""

from dataclasses import dataclass
from typing import Optional, List
import os
import platform
import psutil


@dataclass
class GPUInfo:
    """Information about a GPU."""
    name: str
    vram_gb: float
    driver_version: Optional[str] = None
    device_id: Optional[int] = None
    is_apple_silicon: bool = False
    is_nvidia: bool = False
    is_amd: bool = False
    is_mobile: bool = False
    compute_capability: Optional[str] = None  # CUDA compute capability
    device_count: int = 1  # Number of GPUs available


@dataclass
class HardwareProfile:
    """Complete hardware profile of the system."""
    os: str  # 'windows', 'darwin', 'linux'
    cpu_cores: int
    ram_gb: float
    ram_available_gb: float
    gpu: Optional[GPUInfo]
    is_apple_silicon: bool

    @property
    def has_gpu(self) -> bool:
        return self.gpu is not None
    
    @property
    def has_dedicated_gpu(self) -> bool:
        """Check if system has a dedicated GPU (not Apple Silicon integrated)."""
        return self.gpu is not None and not self.is_apple_silicon

    @property
    def available_memory_gb(self) -> float:
        """Calculate maximum available memory for LLM instances (hard limit)."""
        if self.gpu and not self.is_apple_silicon:
            # External GPU: use 100% of VRAM minus 10% buffer
            return self.gpu.vram_gb * 0.9
        elif self.is_apple_silicon:
            # Apple Silicon: allow up to RAM - 4GB safety buffer (like CPU-only)
            return max(self.ram_gb - 4.0, 4.0)
        else:
            # CPU only: use system RAM minus 4GB safety buffer
            return max(self.ram_gb - 4.0, 4.0)
    
    @property
    def recommended_memory_gb(self) -> float:
        """Get recommended memory usage (for display purposes)."""
        if self.gpu and not self.is_apple_silicon:
            # External GPU: use 100% of VRAM minus 10% buffer
            return self.gpu.vram_gb * 0.9
        elif self.is_apple_silicon:
            # Apple Silicon: recommend 50% of unified RAM
            return self.ram_gb * 0.5
        else:
            # CPU only: recommend 50% of system RAM
            return self.ram_gb * 0.5
    
    @property
    def max_memory_gb(self) -> float:
        """Get maximum safe memory usage."""
        return self.available_memory_gb


def is_android() -> bool:
    """Check if running on Android (beyond just Termux)."""
    # Check multiple Android indicators
    
    # 1. Check for Android-specific environment variables
    android_env_vars = [
        "ANDROID_ROOT",
        "ANDROID_DATA",
        "ANDROID_ART_ROOT",
        "ANDROID_I18N_ROOT",
        "ANDROID_TZDATA_ROOT",
    ]
    if any(os.environ.get(var) for var in android_env_vars):
        return True
    
    # 2. Check for Android-specific paths
    android_paths = [
        "/system/build.prop",
        "/system/bin/app_process",
        "/data/data",
    ]
    if any(os.path.exists(path) for path in android_paths):
        return True
    
    # 3. Check for Termux (which runs on Android)
    if _is_android_or_termux():
        return True
    
    # 4. Check /proc/sys/kernel/osrelease for Android
    try:
        if os.path.exists("/proc/sys/kernel/osrelease"):
            with open("/proc/sys/kernel/osrelease", "r") as f:
                release = f.read().lower()
                if "android" in release:
                    return True
    except Exception:
        pass
    
    return False


def detect_os() -> str:
    """Detect the operating system."""
    system = platform.system().lower()
    
    # Check for Android first (reports as Linux)
    if system == "linux" and is_android():
        return "android"
    elif system == "darwin":
        return "darwin"
    elif system == "windows":
        return "windows"
    elif system == "linux":
        return "linux"
    else:
        return "unknown"


def detect_cpu() -> int:
    """Detect number of CPU cores."""
    return psutil.cpu_count(logical=True) or 4


def detect_ram() -> tuple[float, float]:
    """Detect total and available RAM in GB."""
    mem = psutil.virtual_memory()
    total_gb = mem.total / (1024 ** 3)
    available_gb = mem.available / (1024 ** 3)
    return total_gb, available_gb


def is_apple_silicon() -> bool:
    """Check if running on Apple Silicon."""
    if platform.system().lower() != "darwin":
        return False
    return platform.machine().lower() == "arm64"


def detect_nvidia_gpu() -> Optional[GPUInfo]:
    """Detect NVIDIA GPU using pynvml."""
    try:
        import pynvml
        pynvml.nvmlInit()

        try:
            device_count = pynvml.nvmlDeviceGetCount()
            if device_count == 0:
                return None

            # Get first GPU for now
            handle = pynvml.nvmlDeviceGetHandleByIndex(0)
            name = pynvml.nvmlDeviceGetName(handle)
            if isinstance(name, bytes):
                name = name.decode('utf-8')

            # Get memory info
            mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
            vram_gb = mem_info.total / (1024 ** 3)

            # Get driver version
            try:
                driver = pynvml.nvmlSystemGetDriverVersion()
                if isinstance(driver, bytes):
                    driver = driver.decode('utf-8')
            except Exception:
                driver = None

            # Get compute capability
            compute_capability = None
            try:
                major, minor = pynvml.nvmlDeviceGetCudaComputeCapability(handle)
                compute_capability = f"{major}.{minor}"
            except Exception:
                pass

            return GPUInfo(
                name=name,
                vram_gb=vram_gb,
                driver_version=driver,
                device_id=0,
                is_nvidia=True,
                is_apple_silicon=False,
                is_amd=False,
                compute_capability=compute_capability,
                device_count=device_count
            )
        finally:
            pynvml.nvmlShutdown()

    except ImportError:
        return None
    except Exception:
        return None


def detect_apple_gpu() -> Optional[GPUInfo]:
    """Detect Apple Silicon GPU."""
    if not is_apple_silicon():
        return None

    # On Apple Silicon, GPU shares unified memory with CPU
    total_ram, _ = detect_ram()

    return GPUInfo(
        name="Apple Silicon GPU",
        vram_gb=total_ram,  # Unified memory
        is_apple_silicon=True,
        is_nvidia=False,
        is_amd=False
    )


def _is_android_or_termux() -> bool:
    """Check if running on Android/Termux (where platform.system() returns 'Linux')."""
    try:
        from hardware.qualcomm import is_termux
        return is_termux()
    except Exception:
        return False


def detect_gpu() -> Optional[GPUInfo]:
    """Detect GPU based on platform."""
    os_name = detect_os()

    if os_name == "darwin":
        return detect_apple_gpu()
    elif os_name in ("linux", "windows"):
        # Check for Android/Termux first (reports as Linux)
        if _is_android_or_termux():
            try:
                from hardware.qualcomm import detect_qualcomm_gpu
                return detect_qualcomm_gpu()
            except ImportError:
                return None

        # Priority: NVIDIA > AMD > Intel
        gpu = detect_nvidia_gpu()
        if gpu:
            return gpu

        # Try AMD
        try:
            from hardware.amd import detect_amd_gpu
            gpu = detect_amd_gpu()
            if gpu:
                return gpu
        except ImportError:
            pass

        # Try Intel
        try:
            from hardware.intel import detect_intel_gpu
            gpu = detect_intel_gpu()
            if gpu:
                return gpu
        except ImportError:
            pass

        return None

    return None


def calculate_gpu_layers(gpu: Optional[GPUInfo]) -> int:
    """Calculate optimal number of GPU layers to offload.
    
    Args:
        gpu: GPU information (None if no GPU)
        
    Returns:
        Number of layers to offload (-1 = all, 0 = CPU only)
    """
    if gpu is None:
        return 0
    
    if gpu.is_apple_silicon:
        # Apple Silicon: offload all layers (unified memory)
        return -1
    
    if gpu.is_nvidia:
        # NVIDIA: Check compute capability for compatibility
        if gpu.compute_capability:
            major, _ = gpu.compute_capability.split('.')
            if int(major) < 5:
                # Very old GPUs (Kepler and earlier) may have issues
                return 0
        
        # Multi-GPU support: use device_count to determine layers
        # For now, offload all layers if we have any NVIDIA GPU
        return -1
    
    if gpu.is_amd:
        # AMD: ROCm support varies, be conservative
        return -1
    
    # Unknown GPU type: use CPU
    return 0


def validate_gpu_layers(requested_layers: int, gpu: Optional[GPUInfo]) -> int:
    """Validate and adjust requested GPU layers.
    
    Args:
        requested_layers: Requested number of layers (-1 = all)
        gpu: GPU information
        
    Returns:
        Validated layer count
    """
    if requested_layers == 0:
        return 0
    
    if gpu is None:
        if requested_layers != 0:
            raise ValueError(
                f"Requested {requested_layers} GPU layers but no GPU detected. "
                "Use n_gpu_layers=0 for CPU-only mode."
            )
        return 0
    
    if gpu.is_apple_silicon:
        # Apple Silicon always uses all layers
        return -1
    
    if gpu.is_nvidia and gpu.compute_capability:
        major, _ = gpu.compute_capability.split('.')
        if int(major) < 5:
            raise ValueError(
                f"NVIDIA GPU {gpu.name} has compute capability {gpu.compute_capability}. "
                f"Minimum required is 5.0. Use n_gpu_layers=0 for CPU mode."
            )
    
    return requested_layers


def detect_hardware() -> HardwareProfile:
    """Detect complete hardware profile."""
    os_name = detect_os()
    cpu_cores = detect_cpu()
    ram_gb, ram_available = detect_ram()
    gpu = detect_gpu()
    apple_silicon = is_apple_silicon()

    return HardwareProfile(
        os=os_name,
        cpu_cores=cpu_cores,
        ram_gb=ram_gb,
        ram_available_gb=ram_available,
        gpu=gpu,
        is_apple_silicon=apple_silicon
    )