fix checking nvidia gpu

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