mlx-gpt-oss-120b/check_hardware.py

#!/usr/bin/env python3
"""
Check available hardware for model inference
"""

import torch
import psutil
import logging
import subprocess
import platform

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def check_hardware():
    """Check available hardware resources"""
    logger.info("=" * 50)
    logger.info("🖥️  Hardware Assessment")
    logger.info("=" * 50)
    
    # System RAM
    ram_gb = psutil.virtual_memory().total / (1024 ** 3)
    ram_available_gb = psutil.virtual_memory().available / (1024 ** 3)
    logger.info(f"💾 System RAM: {ram_gb:.1f} GB total, {ram_available_gb:.1f} GB available")
    
    # GPU info (if available)
    if torch.cuda.is_available():
        gpu_count = torch.cuda.device_count()
        logger.info(f"🎮 CUDA GPUs: {gpu_count}")
        for i in range(gpu_count):
            gpu_name = torch.cuda.get_device_name(i)
            gpu_memory = torch.cuda.get_device_properties(i).total_memory / (1024 ** 3)
            logger.info(f"   GPU {i}: {gpu_name} ({gpu_memory:.1f} GB VRAM)")
    else:
        logger.info("🎮 No CUDA GPUs detected")
    
    # Apple Silicon GPU (M1/M2/M3)
    try:
        # Try to get Apple GPU info
        result = subprocess.run(['system_profiler', 'SPDisplaysDataType'], 
                              capture_output=True, text=True)
        if 'Chip' in result.stdout:
            lines = result.stdout.split('\n')
            for line in lines:
                if 'Chip' in line and 'Model' in line:
                    logger.info(f"🍎 Apple Silicon: {line.strip()}")
    except:
        pass
    
    # Disk space
    disk = psutil.disk_usage('/')
    disk_free_gb = disk.free / (1024 ** 3)
    logger.info(f"💿 Disk space: {disk_free_gb:.1f} GB free")
    
    # Model requirements
    logger.info("\n📋 GPT-OSS-120B Requirements:")
    logger.info("   Minimum: 64GB RAM (very slow)")
    logger.info("   Recommended: 128GB+ RAM with 80GB+ GPU VRAM")
    logger.info("   Ideal: Multiple high-end GPUs with 80GB+ VRAM each")
    
    # Recommendations
    logger.info("\n💡 Recommendations:")
    if ram_gb >= 128:
        logger.info("   ✅ You have enough RAM to attempt loading (will be slow)")
    elif ram_gb >= 64:
        logger.info("   ⚠️  Borderline RAM - loading may be very slow or fail")
    else:
        logger.info("   ❌ Insufficient RAM for this model")
    
    if disk_free_gb < 100:
        logger.info("   ⚠️  Low disk space - consider freeing up space")

if __name__ == "__main__":
    check_hardware()