From 3686735bfa7e6840275a87418ad350ad4ed254fe Mon Sep 17 00:00:00 2001 From: fpetrini15 Date: Thu, 10 Oct 2024 10:16:00 -0700 Subject: [PATCH 1/6] Support additional backend dependency directory --- src/backend_manager.cc | 28 +++++++++++++-- src/backend_manager.h | 2 +- src/shared_library.cc | 77 ++++++++++++++++++++++++++++++++++++++++++ src/shared_library.h | 8 +++++ 4 files changed, 112 insertions(+), 3 deletions(-) diff --git a/src/backend_manager.cc b/src/backend_manager.cc index 23f765c8a..fdb212158 100644 --- a/src/backend_manager.cc +++ b/src/backend_manager.cc @@ -26,6 +26,8 @@ #include "backend_manager.h" +#include + #include "backend_memory_manager.h" #include "server_message.h" #include "shared_library.h" @@ -73,8 +75,19 @@ TritonBackend::Create( auto local_backend = std::shared_ptr( new TritonBackend(name, dir, libpath, backend_config)); + auto it = find_if( + backend_cmdline_config.begin(), backend_cmdline_config.end(), + [](const std::pair& config_pair) { + return config_pair.first == "additional-dependency-dir"; + }); + std::string additional_dependency_dir_path; + if (it != backend_cmdline_config.end()) { + additional_dependency_dir_path = it->second; + } + // Load the library and initialize all the entrypoints - RETURN_IF_ERROR(local_backend->LoadBackendLibrary()); + RETURN_IF_ERROR( + local_backend->LoadBackendLibrary(additional_dependency_dir_path)); // Backend initialization is optional... The TRITONBACKEND_Backend // object is this TritonBackend object. We must set set shared @@ -163,7 +176,8 @@ TritonBackend::ClearHandles() } Status -TritonBackend::LoadBackendLibrary() +TritonBackend::LoadBackendLibrary( + const std::string& additional_dependency_dir_path) { TritonBackendInitFn_t bifn; TritonBackendFiniFn_t bffn; @@ -178,8 +192,18 @@ TritonBackend::LoadBackendLibrary() std::unique_ptr slib; RETURN_IF_ERROR(SharedLibrary::Acquire(&slib)); + std::wstring original_path; + if (!additional_dependency_dir_path.empty()) { + RETURN_IF_ERROR(slib->AddAdditionalDependencyDir( + additional_dependency_dir_path, original_path)); + } + RETURN_IF_ERROR(slib->OpenLibraryHandle(libpath_, &dlhandle_)); + if (!original_path.empty()) { + RETURN_IF_ERROR(slib->RemoveAdditionalDependencyDir(original_path)); + } + // Backend initialize and finalize functions, optional RETURN_IF_ERROR(slib->GetEntrypoint( dlhandle_, "TRITONBACKEND_Initialize", true /* optional */, diff --git a/src/backend_manager.h b/src/backend_manager.h index 904004579..bfa2ea79a 100644 --- a/src/backend_manager.h +++ b/src/backend_manager.h @@ -127,7 +127,7 @@ class TritonBackend { const std::string& libpath, const TritonServerMessage& backend_config); void ClearHandles(); - Status LoadBackendLibrary(); + Status LoadBackendLibrary(const std::string& additional_dependency_dir_path); Status UpdateAttributes(); diff --git a/src/shared_library.cc b/src/shared_library.cc index f7530c3a7..c1fb0d85e 100644 --- a/src/shared_library.cc +++ b/src/shared_library.cc @@ -242,4 +242,81 @@ SharedLibrary::GetEntrypoint( return Status::Success; } +Status +SharedLibrary::AddAdditionalDependencyDir( + const std::string& additional_path, std::wstring& original_path) +{ +#ifdef _WIN32 + const std::wstring PATH(L"Path"); + + DWORD len = GetEnvironmentVariableW(PATH.c_str(), NULL, 0); + if (len > 0) { + original_path.resize(len); + GetEnvironmentVariableW(PATH.c_str(), &original_path[0], len); + } else { + return Status(Status::Code::INTERNAL, "PATH variable is empty"); + } + std::wcout << "Before Add: " << original_path << std::endl; + + std::wstring updated_path_value = + std::wstring(additional_path.begin(), additional_path.end()); + updated_path_value += original_path; + + if (!SetEnvironmentVariableW(PATH.c_str(), updated_path_value.c_str())) { + LPSTR err_buffer = nullptr; + size_t size = FormatMessageA( + FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM | + FORMAT_MESSAGE_IGNORE_INSERTS, + NULL, GetLastError(), MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), + (LPSTR)&err_buffer, 0, NULL); + std::string errstr(err_buffer, size); + LocalFree(err_buffer); + return Status( + Status::Code::INTERNAL, + "failed to append user-provided directory to PATH " + errstr); + } + + // TODO: Delete -- just for sanity purposes + std::wstring path_after; + len = GetEnvironmentVariableW(PATH.c_str(), NULL, 0); + if (len > 0) { + path_after.resize(len); + GetEnvironmentVariableW(PATH.c_str(), &path_after[0], len); + } + std::wcout << "After Add: " << path_after << std::endl; +#endif + return Status::Success; +} + +Status +SharedLibrary::RemoveAdditionalDependencyDir(std::wstring& original_path) +{ +#ifdef _WIN32 + const std::wstring PATH(L"Path"); + if (!SetEnvironmentVariableW(PATH.c_str(), original_path.c_str())) { + LPSTR err_buffer = nullptr; + size_t size = FormatMessageA( + FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM | + FORMAT_MESSAGE_IGNORE_INSERTS, + NULL, GetLastError(), MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), + (LPSTR)&err_buffer, 0, NULL); + std::string errstr(err_buffer, size); + LocalFree(err_buffer); + return Status( + Status::Code::INTERNAL, + "failed to restore PATH to its original configuration " + errstr); + } + + // TODO: Delete -- just for sanity purposes + std::wstring path_after; + DWORD len = GetEnvironmentVariableW(PATH.c_str(), NULL, 0); + if (len > 0) { + path_after.resize(len); + GetEnvironmentVariableW(PATH.c_str(), &path_after[0], len); + } + std::wcout << "After Restore: " << path_after << std::endl; +#endif + return Status::Success; +} + }} // namespace triton::core diff --git a/src/shared_library.h b/src/shared_library.h index 39599c004..4c767936e 100644 --- a/src/shared_library.h +++ b/src/shared_library.h @@ -65,6 +65,14 @@ class SharedLibrary { Status GetEntrypoint( void* handle, const std::string& name, const bool optional, void** befn); + // Add an additional dependency directory to PATH (Windows-only). + Status AddAdditionalDependencyDir( + const std::string& additional_path, std::wstring& original_path); + + // Restore PATH to its original configuration. Should be used in + // conjunction with AddAdditionalDependencyDir (Windows-only). + Status RemoveAdditionalDependencyDir(std::wstring& original_path); + private: DISALLOW_COPY_AND_ASSIGN(SharedLibrary); explicit SharedLibrary() = default; From caaabce0182020c13e5222956d75b330bb904b38 Mon Sep 17 00:00:00 2001 From: fpetrini15 Date: Mon, 14 Oct 2024 16:24:26 -0700 Subject: [PATCH 2/6] Verbose logging and tweaks --- src/backend_manager.cc | 2 +- src/shared_library.cc | 41 ++++++++++++++++++++++++++--------------- 2 files changed, 27 insertions(+), 16 deletions(-) diff --git a/src/backend_manager.cc b/src/backend_manager.cc index fdb212158..b27940039 100644 --- a/src/backend_manager.cc +++ b/src/backend_manager.cc @@ -78,7 +78,7 @@ TritonBackend::Create( auto it = find_if( backend_cmdline_config.begin(), backend_cmdline_config.end(), [](const std::pair& config_pair) { - return config_pair.first == "additional-dependency-dir"; + return config_pair.first == "additional-dependency-dirs"; }); std::string additional_dependency_dir_path; if (it != backend_cmdline_config.end()) { diff --git a/src/shared_library.cc b/src/shared_library.cc index c1fb0d85e..2c9cb64b3 100644 --- a/src/shared_library.cc +++ b/src/shared_library.cc @@ -256,7 +256,9 @@ SharedLibrary::AddAdditionalDependencyDir( } else { return Status(Status::Code::INTERNAL, "PATH variable is empty"); } - std::wcout << "Before Add: " << original_path << std::endl; + + LOG_VERBOSE(1) << "Environment before extending PATH: " + << std::string(original_path.begin(), original_path.end()); std::wstring updated_path_value = std::wstring(additional_path.begin(), additional_path.end()); @@ -276,14 +278,21 @@ SharedLibrary::AddAdditionalDependencyDir( "failed to append user-provided directory to PATH " + errstr); } - // TODO: Delete -- just for sanity purposes - std::wstring path_after; - len = GetEnvironmentVariableW(PATH.c_str(), NULL, 0); - if (len > 0) { - path_after.resize(len); - GetEnvironmentVariableW(PATH.c_str(), &path_after[0], len); + if (LOG_VERBOSE_IS_ON(1)) { + std::wstring path_after; + len = GetEnvironmentVariableW(PATH.c_str(), NULL, 0); + if (len > 0) { + path_after.resize(len); + GetEnvironmentVariableW(PATH.c_str(), &path_after[0], len); + } + LOG_VERBOSE(1) << "Environment after extending PATH: " + << std::string(path_after.begin(), path_after.end()); } - std::wcout << "After Add: " << path_after << std::endl; +#else + LOG_WARNING + << "The parameter \"additional-dependency-dirs\" has been specified but " + "is not supported for Linux. It is currently a Windows-only feature. " + "No change to the environment will take effect."; #endif return Status::Success; } @@ -307,14 +316,16 @@ SharedLibrary::RemoveAdditionalDependencyDir(std::wstring& original_path) "failed to restore PATH to its original configuration " + errstr); } - // TODO: Delete -- just for sanity purposes - std::wstring path_after; - DWORD len = GetEnvironmentVariableW(PATH.c_str(), NULL, 0); - if (len > 0) { - path_after.resize(len); - GetEnvironmentVariableW(PATH.c_str(), &path_after[0], len); + if (LOG_VERBOSE_IS_ON(1)) { + std::wstring path_after; + DWORD len = GetEnvironmentVariableW(PATH.c_str(), NULL, 0); + if (len > 0) { + path_after.resize(len); + GetEnvironmentVariableW(PATH.c_str(), &path_after[0], len); + } + LOG_VERBOSE(1) << "Environment after restoring PATH: " + << std::string(path_after.begin(), path_after.end()); } - std::wcout << "After Restore: " << path_after << std::endl; #endif return Status::Success; } From a18758cf7ad1591697e39f71732a893f466dcb9f Mon Sep 17 00:00:00 2001 From: fpetrini15 Date: Tue, 15 Oct 2024 16:07:49 -0700 Subject: [PATCH 3/6] Update copyrights --- src/backend_manager.cc | 2 +- src/backend_manager.h | 2 +- src/shared_library.cc | 2 +- src/shared_library.h | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/backend_manager.cc b/src/backend_manager.cc index b27940039..4fd869eff 100644 --- a/src/backend_manager.cc +++ b/src/backend_manager.cc @@ -1,4 +1,4 @@ -// Copyright 2020-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// Copyright 2020-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved. // // Redistribution and use in source and binary forms, with or without // modification, are permitted provided that the following conditions diff --git a/src/backend_manager.h b/src/backend_manager.h index bfa2ea79a..209253a80 100644 --- a/src/backend_manager.h +++ b/src/backend_manager.h @@ -1,4 +1,4 @@ -// Copyright 2020-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// Copyright 2020-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved. // // Redistribution and use in source and binary forms, with or without // modification, are permitted provided that the following conditions diff --git a/src/shared_library.cc b/src/shared_library.cc index 2c9cb64b3..c28406487 100644 --- a/src/shared_library.cc +++ b/src/shared_library.cc @@ -1,4 +1,4 @@ -// Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +// Copyright (c) 2021-2024, NVIDIA CORPORATION. All rights reserved. // // Redistribution and use in source and binary forms, with or without // modification, are permitted provided that the following conditions diff --git a/src/shared_library.h b/src/shared_library.h index 4c767936e..b114b4a7c 100644 --- a/src/shared_library.h +++ b/src/shared_library.h @@ -1,4 +1,4 @@ -// Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +// Copyright (c) 2021-2024, NVIDIA CORPORATION. All rights reserved. // // Redistribution and use in source and binary forms, with or without // modification, are permitted provided that the following conditions From 0da79f4503c6db027ed98a15bbe33bdaff718f51 Mon Sep 17 00:00:00 2001 From: fpetrini15 Date: Fri, 18 Oct 2024 12:09:41 -0700 Subject: [PATCH 4/6] Review comments --- src/backend_manager.cc | 2 +- src/shared_library.cc | 4 ++-- src/shared_library.h | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/backend_manager.cc b/src/backend_manager.cc index 4fd869eff..c4ff0d137 100644 --- a/src/backend_manager.cc +++ b/src/backend_manager.cc @@ -200,7 +200,7 @@ TritonBackend::LoadBackendLibrary( RETURN_IF_ERROR(slib->OpenLibraryHandle(libpath_, &dlhandle_)); - if (!original_path.empty()) { + if (!additional_dependency_dir_path.empty()) { RETURN_IF_ERROR(slib->RemoveAdditionalDependencyDir(original_path)); } diff --git a/src/shared_library.cc b/src/shared_library.cc index c28406487..4df9794bf 100644 --- a/src/shared_library.cc +++ b/src/shared_library.cc @@ -254,7 +254,7 @@ SharedLibrary::AddAdditionalDependencyDir( original_path.resize(len); GetEnvironmentVariableW(PATH.c_str(), &original_path[0], len); } else { - return Status(Status::Code::INTERNAL, "PATH variable is empty"); + original_path = L""; } LOG_VERBOSE(1) << "Environment before extending PATH: " @@ -298,7 +298,7 @@ SharedLibrary::AddAdditionalDependencyDir( } Status -SharedLibrary::RemoveAdditionalDependencyDir(std::wstring& original_path) +SharedLibrary::RemoveAdditionalDependencyDir(const std::wstring& original_path) { #ifdef _WIN32 const std::wstring PATH(L"Path"); diff --git a/src/shared_library.h b/src/shared_library.h index b114b4a7c..fc69f306d 100644 --- a/src/shared_library.h +++ b/src/shared_library.h @@ -71,7 +71,7 @@ class SharedLibrary { // Restore PATH to its original configuration. Should be used in // conjunction with AddAdditionalDependencyDir (Windows-only). - Status RemoveAdditionalDependencyDir(std::wstring& original_path); + Status RemoveAdditionalDependencyDir(const std::wstring& original_path); private: DISALLOW_COPY_AND_ASSIGN(SharedLibrary); From 45f3f84d73640c1be787a433df05ed7afac51deb Mon Sep 17 00:00:00 2001 From: fpetrini15 Date: Fri, 18 Oct 2024 15:21:28 -0700 Subject: [PATCH 5/6] Add new error case --- src/shared_library.cc | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/src/shared_library.cc b/src/shared_library.cc index 4df9794bf..6fc237f6a 100644 --- a/src/shared_library.cc +++ b/src/shared_library.cc @@ -249,6 +249,13 @@ SharedLibrary::AddAdditionalDependencyDir( #ifdef _WIN32 const std::wstring PATH(L"Path"); + if (additional_path.back() != ';') { + return Status( + Status::Code::INVALID_ARG, + "backend config parameter \"additional-dependency-dirs\" is malformed. " + "Each additional path provided should terminate with a ';'."); + } + DWORD len = GetEnvironmentVariableW(PATH.c_str(), NULL, 0); if (len > 0) { original_path.resize(len); From 6468f3b1cb3e8040a24039a6117e26afdc04a2fe Mon Sep 17 00:00:00 2001 From: fpetrini15 Date: Mon, 21 Oct 2024 17:40:32 -0700 Subject: [PATCH 6/6] Review comments --- src/shared_library.h | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/shared_library.h b/src/shared_library.h index fc69f306d..74cc8f1e2 100644 --- a/src/shared_library.h +++ b/src/shared_library.h @@ -65,12 +65,12 @@ class SharedLibrary { Status GetEntrypoint( void* handle, const std::string& name, const bool optional, void** befn); - // Add an additional dependency directory to PATH (Windows-only). + // Add an additional dependency directory to PATH. Status AddAdditionalDependencyDir( const std::string& additional_path, std::wstring& original_path); // Restore PATH to its original configuration. Should be used in - // conjunction with AddAdditionalDependencyDir (Windows-only). + // conjunction with AddAdditionalDependencyDir. Status RemoveAdditionalDependencyDir(const std::wstring& original_path); private: