scripts/gpu_detect.py

#!/usr/bin/env python3
"""
GPU Detection and Validation Module
- Detects NVIDIA GPU
- Checks VRAM availability
- Checks CUDA version compatibility
- Returns validation result (JSON)
"""
import json
import os
import re
import subprocess
import sys


def get_cuda_version():
    """
    Detect the installed CUDA version on the system.

    Attempts to determine the CUDA version by checking multiple sources in order:
    1. nvcc compiler version (--version flag)
    2. nvidia-smi CUDA version output

    Both methods parse version strings using regex to extract major and minor version numbers.

    Returns:
        tuple: (major_version, minor_version) as integers if CUDA is detected.
               (None, None) if CUDA is not found or detection fails.

    Raises:
        None: Silently handles subprocess timeouts and file not found errors.

    Example:
        >>> major, minor = get_cuda_version()
        >>> if major is not None:
        ...     print(f"CUDA {major}.{minor} detected")
    """
    # 1. Try nvcc
    try:
        result = subprocess.run(
            ["nvcc", "--version"],
            capture_output=True,
            text=True,
            timeout=5
        )
        combined_output = (result.stdout or "") + "\n" + (result.stderr or "")
        match = re.search(r'release\s+(\d+)\.(\d+)', combined_output)
        if match:
            return int(match.group(1)), int(match.group(2))
    except (subprocess.TimeoutExpired, FileNotFoundError):
        pass

    # 2. Fallback to nvidia-smi CUDA version
    try:
        result = subprocess.run(
            ["nvidia-smi"],
            capture_output=True,
            text=True,
            timeout=5
        )
        combined_output = (result.stdout or "") + "\n" + (result.stderr or "")
        match = re.search(r'CUDA Version:\s*(\d+)\.(\d+)', combined_output)
        if match:
            return int(match.group(1)), int(match.group(2))
    except (subprocess.TimeoutExpired, FileNotFoundError):
        pass

    return None, None


def get_gpu_info():
    """
    Retrieve information about available NVIDIA GPUs.

    Uses nvidia-smi to query GPU information including name, total VRAM, and used VRAM.
    The command executed is: nvidia-smi --query-gpu=name,memory.total,memory.used
    --format=csv,noheader,nounits

    Returns:
        list[dict]: A list of dictionaries, each containing:
            - name (str): GPU model name
            - total_vram_mb (int): Total VRAM in megabytes
            - used_vram_mb (int): Currently used VRAM in megabytes
                   Returns an empty list if no GPUs are detected or if nvidia-smi fails.

    Raises:
        None: Silently handles subprocess timeouts, file not found errors, and value errors.

    Example:
        >>> gpus = get_gpu_info()
        >>> for gpu in gpus:
        ...     print(f"{gpu['name']}: {gpu['total_vram_mb']}MB total, {gpu['used_vram_mb']}MB used")
    """
    try:
        result = subprocess.run(
            ["nvidia-smi", "--query-gpu=name,memory.total,memory.used",
             "--format=csv,noheader,nounits"],
            capture_output=True,
            text=True,
            timeout=10
        )
        gpus = []
        for line in result.stdout.strip().split('\n'):
            if not line.strip():
                continue
            parts = [p.strip() for p in line.split(',')]
            if len(parts) >= 3:
                gpus.append({
                    'name': parts[0],
                    'total_vram_mb': int(float(parts[1])),
                    'used_vram_mb': int(float(parts[2]))
                })
        return gpus
    except (subprocess.TimeoutExpired, FileNotFoundError, ValueError):
        return []


def validate_environment(min_vram_mb=15000, min_cuda_major=11, min_cuda_minor=8):
    """
    Validate the GPU environment for running Large Language Models.

    Performs comprehensive validation checks to ensure the environment meets
    requirements for running LLM models like Qwen2.5-14B. Checks include:
    1. NVIDIA GPU detection
    2. Total VRAM availability across all GPUs
    3. CUDA version compatibility

    Also generates warnings for suboptimal configurations (e.g., CUDA 11.x
    when CUDA 12.x is recommended).

    Args:
        min_vram_mb (int): Minimum required total VRAM in megabytes.
                           Default: 15000 (15GB).
        min_cuda_major (int): Minimum required CUDA major version.
                              Default: 11.
        min_cuda_minor (int): Minimum required CUDA minor version.
                              Default: 8.

    Returns:
        tuple: (is_valid, result_dict) where:
            - is_valid (bool): True if environment passes all checks, False otherwise.
            - result_dict (dict): Detailed validation result containing:
                - gpu_detected (bool): Whether GPUs were found
                - gpus (list[dict]): List of GPU info dictionaries
                - cuda_detected (bool): Whether CUDA was found
                - cuda_version (str or None): Detected CUDA version string
                - total_vram_mb (int): Total VRAM across all GPUs
                - errors (list[str]): List of validation error messages
                - warnings (list[str]): List of validation warning messages
                - status (str): "valid", "error", or other status
                - selected_model (str): Model name from MODEL_NAME env var or default

    Raises:
        None: Returns validation results via the tuple instead of raising exceptions.

    Example:
        >>> is_valid, result = validate_environment(min_vram_mb=20000)
        >>> if is_valid:
        ...     print("Environment is ready for LLM inference")
        ... else:
        ...     for error in result['errors']:
        ...         print(f"Error: {error}")
    """
    gpu_info = get_gpu_info()
    cuda_major, cuda_minor = get_cuda_version()

    result = {
        'gpu_detected': len(gpu_info) > 0,
        'gpus': gpu_info,
        'cuda_detected': cuda_major is not None,
        'cuda_version': f"{cuda_major}.{cuda_minor}" if cuda_major is not None else None,
        'errors': [],
        'warnings': []
    }

    # Check 1: GPU detected
    if not result['gpu_detected']:
        result['errors'].append("No NVIDIA GPU detected. GPU acceleration is required.")
        result['status'] = "error"
        return False, result

    # Check 2: Total VRAM
    total_vram = sum(gpu['total_vram_mb'] for gpu in gpu_info)
    result['total_vram_mb'] = total_vram
    if total_vram < min_vram_mb:
        result['errors'].append(
            f"Insufficient VRAM: {total_vram} MB detected, "
            f"{min_vram_mb} MB minimum required."
        )
        result['status'] = "error"
        return False, result

    # Check 3: CUDA version
    if not result['cuda_detected']:
        result['errors'].append(
            f"CUDA not detected. CUDA {min_cuda_major}.{min_cuda_minor}+ required."
        )
        result['status'] = "error"
        return False, result

    if cuda_major < min_cuda_major or (cuda_major == min_cuda_major and cuda_minor < min_cuda_minor):
        result['errors'].append(
            f"CUDA {cuda_major}.{cuda_minor} too old. "
            f"Minimum required: CUDA {min_cuda_major}.{min_cuda_minor}."
        )
        result['status'] = "error"
        return False, result

    # Check 4: Warning for CUDA < 12.0
    if cuda_major == 11:
        result['warnings'].append(
            f"CUDA {cuda_major}.{cuda_minor} detected. "
            f"Recommended: CUDA 12.x for optimal performance."
        )

    result['selected_model'] = os.getenv('MODEL_NAME', 'qwen2.5:14b')
    result['status'] = "valid"
    return True, result


if __name__ == "__main__":
    min_vram = int(os.getenv('MIN_VRAM_MB', '15000'))
    min_cuda_major = int(os.getenv('MIN_CUDA_MAJOR', '11'))
    min_cuda_minor = int(os.getenv('MIN_CUDA_MINOR', '8'))

    is_valid, result = validate_environment(min_vram, min_cuda_major, min_cuda_minor)

    # Save to /tmp/gpu_info.json for /gpu-info endpoint if possible
    try:
        os.makedirs('/tmp', exist_ok=True)
        with open('/tmp/gpu_info.json', 'w', encoding='utf-8') as f:
            json.dump(result, f, indent=2)
    except Exception:
        pass

    # Output JSON to stdout
    print(json.dumps(result, indent=2))

    if not is_valid:
        sys.exit(1)