diff --git a/src/backend_manager.cc b/src/backend_manager.cc index 23f765c8a..c4ff0d137 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 @@ -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-dirs"; + }); + 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 (!additional_dependency_dir_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..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 @@ -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..6fc237f6a 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 @@ -242,4 +242,99 @@ 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"); + + 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); + GetEnvironmentVariableW(PATH.c_str(), &original_path[0], len); + } else { + original_path = L""; + } + + 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()); + 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); + } + + 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()); + } +#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; +} + +Status +SharedLibrary::RemoveAdditionalDependencyDir(const 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); + } + + 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()); + } +#endif + return Status::Success; +} + }} // namespace triton::core diff --git a/src/shared_library.h b/src/shared_library.h index 39599c004..74cc8f1e2 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 @@ -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. + Status AddAdditionalDependencyDir( + const std::string& additional_path, std::wstring& original_path); + + // Restore PATH to its original configuration. Should be used in + // conjunction with AddAdditionalDependencyDir. + Status RemoveAdditionalDependencyDir(const std::wstring& original_path); + private: DISALLOW_COPY_AND_ASSIGN(SharedLibrary); explicit SharedLibrary() = default;