diff --git a/Azaion.Inference/hardware_service.pyx b/Azaion.Inference/hardware_service.pyx index 7a1e97c..9782b21 100644 --- a/Azaion.Inference/hardware_service.pyx +++ b/Azaion.Inference/hardware_service.pyx @@ -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 - return False - except Exception as e: - print(e) + 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 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: