Support additional backend dependency directory

src/backend_manager.cc

@@ -26,6 +26,8 @@
 
 #include "backend_manager.h"
 
+#include <algorithm>
+
 #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<TritonBackend>(
       new TritonBackend(name, dir, libpath, backend_config));
 
+  auto it = find_if(
+      backend_cmdline_config.begin(), backend_cmdline_config.end(),
+      [](const std::pair<std::string, std::string>& 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<SharedLibrary> 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 */,

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();
 

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

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;