Gabriel20xx/VisionFlowAutomator
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Refactor GPU memory info retrieval to improve module imports and error handling; remove timeout mechanism in favor of caching

Browse Source
This commit is contained in:
Gabriel20xx
2025-07-31 14:26:58 +02:00
parent 7489658551
commit dafe588601
1 changed files with 8 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files
main.py
+8 -34
View File
@@ -51,6 +51,12 @@ def get_gpu_memory_info():
Returns dict with GPU memory info or basic info if no GPU found. Returns dict with GPU memory info or basic info if no GPU found.
Uses caching to prevent repeated expensive system calls. Uses caching to prevent repeated expensive system calls.
""" """
# Import modules at function level to avoid scope issues
import json
import subprocess
import platform
import os
# Cache GPU info for 30 seconds to prevent UI hanging # Cache GPU info for 30 seconds to prevent UI hanging
current_time = time.time() current_time = time.time()
cache_duration = 30 # seconds cache_duration = 30 # seconds
@@ -158,9 +164,6 @@ def get_gpu_memory_info():
# Try Windows-specific methods # Try Windows-specific methods
try: try:
import subprocess
import re
# Try nvidia-smi command for NVIDIA GPUs # Try nvidia-smi command for NVIDIA GPUs
try: try:
result = subprocess.run(['nvidia-smi', '--query-gpu=name,memory.total,memory.used,memory.free', '--format=csv,noheader,nounits'], result = subprocess.run(['nvidia-smi', '--query-gpu=name,memory.total,memory.used,memory.free', '--format=csv,noheader,nounits'],
@@ -200,7 +203,6 @@ def get_gpu_memory_info():
], capture_output=True, text=True, timeout=3) # Reduced timeout ], capture_output=True, text=True, timeout=3) # Reduced timeout
if wmi_result.returncode == 0 and wmi_result.stdout.strip(): if wmi_result.returncode == 0 and wmi_result.stdout.strip():
import json
wmi_data = json.loads(wmi_result.stdout.strip()) wmi_data = json.loads(wmi_result.stdout.strip())
# Handle both single GPU and multiple GPUs # Handle both single GPU and multiple GPUs
@@ -236,7 +238,6 @@ def get_gpu_memory_info():
], capture_output=True, text=True, timeout=3) # Reduced timeout ], capture_output=True, text=True, timeout=3) # Reduced timeout
if wmi_result.returncode == 0 and wmi_result.stdout.strip(): if wmi_result.returncode == 0 and wmi_result.stdout.strip():
import json
wmi_data = json.loads(wmi_result.stdout.strip()) wmi_data = json.loads(wmi_result.stdout.strip())
# Handle both single GPU and multiple GPUs # Handle both single GPU and multiple GPUs
@@ -271,7 +272,6 @@ def get_gpu_memory_info():
], capture_output=True, text=True, timeout=3) # Reduced timeout ], capture_output=True, text=True, timeout=3) # Reduced timeout
if wmi_result.returncode == 0 and wmi_result.stdout.strip(): if wmi_result.returncode == 0 and wmi_result.stdout.strip():
import json
wmi_data = json.loads(wmi_result.stdout.strip()) wmi_data = json.loads(wmi_result.stdout.strip())
# Handle both single GPU and multiple GPUs # Handle both single GPU and multiple GPUs
@@ -306,8 +306,6 @@ def get_gpu_memory_info():
# Try DirectX DXGI detection (Windows 10+) # Try DirectX DXGI detection (Windows 10+)
try: try:
import subprocess
# Use PowerShell to query DirectX information # Use PowerShell to query DirectX information
dxdiag_result = subprocess.run([ dxdiag_result = subprocess.run([
'powershell', '-Command', 'powershell', '-Command',
@@ -319,7 +317,6 @@ def get_gpu_memory_info():
], capture_output=True, text=True, timeout=3) # Reduced timeout ], capture_output=True, text=True, timeout=3) # Reduced timeout
if dxdiag_result.returncode == 0 and dxdiag_result.stdout.strip(): if dxdiag_result.returncode == 0 and dxdiag_result.stdout.strip():
import json
dx_data = json.loads(dxdiag_result.stdout.strip()) dx_data = json.loads(dxdiag_result.stdout.strip())
if isinstance(dx_data, dict): if isinstance(dx_data, dict):
@@ -346,9 +343,6 @@ def get_gpu_memory_info():
# Try Linux/macOS methods for integrated GPUs # Try Linux/macOS methods for integrated GPUs
try: try:
import subprocess
import platform
system = platform.system().lower() system = platform.system().lower()
if system == 'linux': if system == 'linux':
@@ -390,7 +384,6 @@ def get_gpu_memory_info():
# Try checking /sys/class/drm for Intel iGPU on Linux # Try checking /sys/class/drm for Intel iGPU on Linux
try: try:
import os
drm_path = '/sys/class/drm' drm_path = '/sys/class/drm'
if os.path.exists(drm_path): if os.path.exists(drm_path):
drm_devices = os.listdir(drm_path) drm_devices = os.listdir(drm_path)
@@ -419,7 +412,6 @@ def get_gpu_memory_info():
], capture_output=True, text=True, timeout=15) ], capture_output=True, text=True, timeout=15)
if profiler_result.returncode == 0: if profiler_result.returncode == 0:
import json
display_data = json.loads(profiler_result.stdout) display_data = json.loads(profiler_result.stdout)
displays = display_data.get('SPDisplaysDataType', []) displays = display_data.get('SPDisplaysDataType', [])
@@ -1025,26 +1017,8 @@ class MainWindow(QtWidgets.QMainWindow):
# Get system info # Get system info
system_memory = psutil.virtual_memory() system_memory = psutil.virtual_memory()
# Get GPU information with timeout protection # Get GPU information with error handling (no timeout needed due to caching)
try: try:
# Use a simple timeout mechanism for GPU detection
import signal
def timeout_handler(signum, frame):
raise TimeoutError("GPU detection timeout")
# Set up timeout for GPU detection (3 seconds max)
old_handler = signal.signal(signal.SIGALRM, timeout_handler)
signal.alarm(3)
try:
gpu_info = get_gpu_memory_info()
finally:
signal.alarm(0) # Cancel the alarm
signal.signal(signal.SIGALRM, old_handler)
except (TimeoutError, AttributeError):
# Fallback if timeout or signal not available (Windows)
gpu_info = get_gpu_memory_info() gpu_info = get_gpu_memory_info()
except Exception as e: except Exception as e:
logger.debug(f"GPU detection failed: {e}") logger.debug(f"GPU detection failed: {e}")
@@ -1069,7 +1043,7 @@ class MainWindow(QtWidgets.QMainWindow):
'gpu_method': gpu_info.get('method', 'none') 'gpu_method': gpu_info.get('method', 'none')
} }
except ImportError: except ImportError:
# psutil not available, get GPU info anyway with timeout protection # psutil not available, get GPU info anyway with error handling
try: try:
gpu_info = get_gpu_memory_info() gpu_info = get_gpu_memory_info()
except Exception as e: except Exception as e: