fix checking nvidia gpu

This commit is contained in:
Alex Bezdieniezhnykh
2025-04-30 22:19:42 +03:00
parent b3db108f59
commit ae83bc8542
+19 -7
View File
@@ -1,6 +1,7 @@
import re import re
import subprocess import subprocess
import psutil import psutil
import pynvml
cdef class HardwareInfo: cdef class HardwareInfo:
def __init__(self, str cpu, str gpu, str memory, str mac_address): def __init__(self, str cpu, str gpu, str memory, str mac_address):
@@ -46,14 +47,25 @@ cdef class HardwareService:
@staticmethod @staticmethod
cdef has_nvidia_gpu(): cdef has_nvidia_gpu():
try: try:
output = subprocess.check_output(['nvidia-smi']).decode() pynvml.nvmlInit()
match = re.search(r'CUDA Version:\s*([\d.]+)', output) device_count = pynvml.nvmlDeviceGetCount()
if match:
return float(match.group(1)) > 11 if device_count > 0:
return False print(f"Found NVIDIA GPU(s).")
except Exception as e: return True
print(e) else:
print("No NVIDIA GPUs found by NVML.")
return False
except pynvml.NVMLError as error:
print(f"Failed to find NVIDIA GPU")
return False return False
finally:
try:
pynvml.nvmlShutdown()
except:
print('Failed to shutdown pynvml cause probably no NVidia GPU')
pass
cdef HardwareInfo get_hardware_info(self): cdef HardwareInfo get_hardware_info(self):
if self.is_windows: if self.is_windows: