#!/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)