add dylib for windows

janhq · Nov 5, 2024 · 3eef24d · 3eef24d
1 parent eb96ed4
commit 3eef24d
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 1 deletion.
diff --git 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
@@ -76,6 +76,7 @@ class EngineService {
     EngineV engine;
 #if defined(_WIN32)
     DLL_DIRECTORY_COOKIE cookie;
+    DLL_DIRECTORY_COOKIE cuda_cookie;
 #endif
   };