Spaces:
Sleeping
Sleeping
File size: 8,415 Bytes
91ab82e 497602c 91ab82e 497602c 91ab82e 97f856b 91ab82e 497602c 91ab82e 97f856b 91ab82e 97f856b 91ab82e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 | #!/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)
|