From ab24832a88efeb581aa29370c65ebda1fb18991c Mon Sep 17 00:00:00 2001
From: vansangpfiev <sang@jan.ai>
Date: Tue, 11 Feb 2025 13:54:18 +0700
Subject: [PATCH] fix: map nvidia and vulkan uuid

---
 engine/services/hardware_service.cc           |  2 +-
 engine/utils/hardware/gpu/vulkan/vulkan_gpu.h | 22 +++++++++----------
 engine/utils/hardware/gpu_info.h              |  9 ++++++++
 3 files changed, 20 insertions(+), 13 deletions(-)

diff --git a/engine/services/hardware_service.cc b/engine/services/hardware_service.cc
index b4300d828..2ace50172 100644
--- a/engine/services/hardware_service.cc
+++ b/engine/services/hardware_service.cc
@@ -491,7 +491,7 @@ std::vector<int> HardwareService::GetCudaConfig() {
   // Map uuid back to nvidia id
   for (auto const& uuid : uuids) {
     for (auto const& ngpu : nvidia_gpus) {
-      if (uuid == ngpu.uuid) {
+      if (ngpu.uuid.find(uuid) != std::string::npos) {
         res.push_back(std::stoi(ngpu.id));
       }
     }
diff --git a/engine/utils/hardware/gpu/vulkan/vulkan_gpu.h b/engine/utils/hardware/gpu/vulkan/vulkan_gpu.h
index b45d4d2dc..1b04029f0 100644
--- a/engine/utils/hardware/gpu/vulkan/vulkan_gpu.h
+++ b/engine/utils/hardware/gpu/vulkan/vulkan_gpu.h
@@ -421,18 +421,16 @@ inline cpp::result<std::vector<cortex::hw::GPU>, std::string> GetGpuInfoList() {
 #endif
     int free_vram_MiB =
         total_vram_MiB > used_vram_MiB ? total_vram_MiB - used_vram_MiB : 0;
-    if (total_vram_MiB > 0) {
-      gpus.emplace_back(cortex::hw::GPU{
-          .id = std::to_string(id),
-          .device_id = device_properties.deviceID,
-          .name = device_properties.deviceName,
-          .version = std::to_string(device_properties.driverVersion),
-          .add_info = cortex::hw::AmdAddInfo{},
-          .free_vram = free_vram_MiB,
-          .total_vram = total_vram_MiB,
-          .uuid = uuid_to_string(device_id_properties.deviceUUID),
-          .vendor = GetVendorStr(device_properties.vendorID)});
-    }
+    gpus.emplace_back(cortex::hw::GPU{
+        .id = std::to_string(id),
+        .device_id = device_properties.deviceID,
+        .name = device_properties.deviceName,
+        .version = std::to_string(device_properties.driverVersion),
+        .add_info = cortex::hw::AmdAddInfo{},
+        .free_vram = free_vram_MiB,
+        .total_vram = total_vram_MiB,
+        .uuid = uuid_to_string(device_id_properties.deviceUUID),
+        .vendor = GetVendorStr(device_properties.vendorID)});
     id++;
   }
 
diff --git a/engine/utils/hardware/gpu_info.h b/engine/utils/hardware/gpu_info.h
index 288259570..b9f7173fe 100644
--- a/engine/utils/hardware/gpu_info.h
+++ b/engine/utils/hardware/gpu_info.h
@@ -29,6 +29,15 @@ inline std::vector<GPU> GetGPUInfo() {
     }
   }
 
+  // Erase invalid GPUs
+  for (std::vector<cortex::hw::GPU>::iterator it = vulkan_gpus.begin();
+       it != vulkan_gpus.end();) {
+    if ((*it).total_vram <= 0)
+      it = vulkan_gpus.erase(it);
+    else
+      ++it;
+  }
+
   if (use_vulkan_info) {
     return vulkan_gpus;
   } else {