From 3eef24dcaaab167b398442bf0a4acfe80a2e9c8d Mon Sep 17 00:00:00 2001
From: James <namnh0122@gmail.com>
Date: Tue, 5 Nov 2024 11:29:50 +0700
Subject: [PATCH] add dylib for windows

---
 engine/services/engine_service.cc | 20 +++++++++++++++++++-
 engine/services/engine_service.h  |  1 +
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/engine/services/engine_service.cc b/engine/services/engine_service.cc
index 0780e65ec..07433b771 100644
--- a/engine/services/engine_service.cc
+++ b/engine/services/engine_service.cc
@@ -467,7 +467,6 @@ cpp::result<bool, std::string> EngineService::DownloadCuda(
     auto engine_path = file_manager_utils::GetCudaToolkitPath(engine);
     archive_utils::ExtractArchive(finishedTask.items[0].localPath.string(),
                                   engine_path.string());
-
     try {
       std::filesystem::remove(finishedTask.items[0].localPath);
     } catch (std::exception& e) {
@@ -758,6 +757,15 @@ cpp::result<void, std::string> EngineService::LoadEngine(
       } else {
         LOG_WARN << "Could not add dll directory: " << p;
       }
+
+      auto cuda_path = file_manager_utils::GetCudaToolkitPath(e_type);
+      if (auto cuda_cookie = AddDllDirectory(cuda_path.string());
+          cuda_cookie != 0) {
+        LOG_INFO << "Added cuda dll directory: " << p;
+        engines_[e_type].cuda_cookie = cuda_cookie;
+      } else {
+        LOG_WARN << "Could not add cuda dll directory: " << p;
+      }
     };
 
     if (bool should_use_dll_search_path = !(getenv("ENGINE_PATH"));
@@ -770,6 +778,11 @@ cpp::result<void, std::string> EngineService::LoadEngine(
         } else {
           LOG_INFO << "Removed dll directory: " << kLlamaRepo;
         }
+        if (!RemoveDllDirectory(engines_[kLlamaRepo].cuda_cookie)) {
+          LOG_WARN << "Could not remove cuda dll directory: " << kLlamaRepo;
+        } else {
+          LOG_INFO << "Removed cuda dll directory: " << kLlamaRepo;
+        }
 
         add_dll(ne, engine_dir_path.string());
       } else if (IsEngineLoaded(kTrtLlmRepo) && ne == kLlamaRepo) {
@@ -821,6 +834,11 @@ cpp::result<void, std::string> EngineService::UnloadEngine(
   } else {
     LOG_INFO << "Removed dll directory: " << ne;
   }
+  if (!RemoveDllDirectory(engines_[ne].cuda_cookie)) {
+    LOG_WARN << "Could not remove cuda dll directory: " << ne;
+  } else {
+    LOG_INFO << "Removed cuda dll directory: " << ne;
+  }
 #endif
   engines_.erase(ne);
   LOG_INFO << "Unloaded engine " + ne;
diff --git a/engine/services/engine_service.h b/engine/services/engine_service.h
index 5b99e31dd..4e58fccfd 100644
--- a/engine/services/engine_service.h
+++ b/engine/services/engine_service.h
@@ -76,6 +76,7 @@ class EngineService {
     EngineV engine;
 #if defined(_WIN32)
     DLL_DIRECTORY_COOKIE cookie;
+    DLL_DIRECTORY_COOKIE cuda_cookie;
 #endif
   };