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 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 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 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 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 };