From ae83bc8542cf1b9c712a4aba8b904d723c4f8803 Mon Sep 17 00:00:00 2001
From: Alex Bezdieniezhnykh <zxsanny@gmail.com>
Date: Wed, 30 Apr 2025 22:19:42 +0300
Subject: [PATCH] fix checking nvidia gpu

---
 Azaion.Inference/hardware_service.pyx | 26 +++++++++++++++++++-------
 1 file changed, 19 insertions(+), 7 deletions(-)

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: