feat: Incorporate "additional-dependency-dirs"

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 <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-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<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 (!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 */,

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

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

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;