Introduce DALI_PRELOAD_PLUGINS

Signed-off-by: Joaquin Anton <janton@nvidia.com>

dali/c_api/c_api.cc

@@ -717,7 +717,15 @@ void daliDeletePipeline(daliPipelineHandle_t pipe_handle) {
 }
 
 void daliLoadLibrary(const char* lib_path) {
-    dali::PluginManager::LoadLibrary(lib_path);
+  dali::PluginManager::LoadLibrary(lib_path);
+}
+
+void daliLoadPluginDirectory(const char* plugin_dir) {
+  dali::PluginManager::LoadDirectory(plugin_dir);
+}
+
+void daliLoadDefaultPlugins() {
+  dali::PluginManager::LoadDefaultPlugins();
 }
 
 void daliGetReaderMetadata(daliPipelineHandle_t pipe_handle, const char *reader_name,

dali/operators/operators.cc

@@ -12,17 +12,20 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "dali/core/api_helper.h"
 #include "dali/operators.h"
-#include "dali/npp/npp.h"
+#include "dali/core/api_helper.h"
 #include "dali/core/cuda_stream_pool.h"
+#include "dali/npp/npp.h"
+#include "dali/plugin/plugin_manager.h"
 
 #if DALI_USE_NVJPEG
-  #include "dali/operators/decoder/nvjpeg/nvjpeg_helper.h"
+#include "dali/operators/decoder/nvjpeg/nvjpeg_helper.h"
 #endif
 
+#include <dlfcn.h>
 #include <nvimgcodec.h>
 
+
 /*
  * The point of these functions is to force the linker to link against dali_operators lib
  * and not optimize-out symbols from dali_operators
@@ -35,8 +38,10 @@ namespace dali {
 
 DLL_PUBLIC void InitOperatorsLib() {
   (void)CUDAStreamPool::instance();
+  dali::PluginManager::LoadDefaultPlugins();
 }
 
+
 DLL_PUBLIC int GetNppVersion() {
   return NPPGetVersion();
 }

dali/plugin/plugin_manager.cc

@@ -12,18 +12,95 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include <dlfcn.h>
 #include "dali/plugin/plugin_manager.h"
+#include <dlfcn.h>
+#include <filesystem>
+#include <utility>
 #include "dali/core/error_handling.h"
 
+namespace fs = std::filesystem;
+
 namespace dali {
 
 void PluginManager::LoadLibrary(const std::string& lib_path, bool global_symbols) {
-    // dlopen is thread safe
-    int flags = global_symbols ? RTLD_GLOBAL : RTLD_LOCAL;
-    flags |= RTLD_LAZY;
-    auto handle = dlopen(lib_path.c_str(), flags);
-    DALI_ENFORCE(handle != nullptr, "Failed to load library: " + std::string(dlerror()));
+  // dlopen is thread safe
+  int flags = global_symbols ? RTLD_GLOBAL : RTLD_LOCAL;
+  flags |= RTLD_LAZY;
+  auto handle = dlopen(lib_path.c_str(), flags);
+  DALI_ENFORCE(handle != nullptr, "Failed to load library: " + std::string(dlerror()));
+}
+
+inline const std::string& DefaultPluginPath() {
+  static const std::string path = [&]() -> std::string {
+    Dl_info info;
+    if (dladdr((const void*)DefaultPluginPath, &info)) {
+      fs::path path(info.dli_fname);
+      // use the directory of the current shared-object file as starting point to autodiscover the
+      // plugin default directory
+      // ~/.local/lib/python3.8/site-packages/nvidia/dali/libdali.so ->
+      //     ~/.local/lib/python3.8/site-packages/nvidia/dali/plugins/{plugin_name}/libdali_{plugin_name}.so
+      path = path.parent_path();
+      path /= "plugin";
+      return path.string();
+    }
+    return {};
+  }();
+  return path;
+}
+
+inline void PluginManager::LoadDirectory(const std::string& path, bool global_symbols) {
+  std::vector<std::string> plugin_paths;
+  if (!fs::is_directory(path)) {
+    LOG_LINE << path << " is not a directory. Nothing to load\n";
+    return;
+  }
+  for (const auto& fpath : fs::recursive_directory_iterator(path)) {
+    // pos=0 limits the search to the prefix
+    if (fpath.path().stem().string().rfind("libdali_", 0) == 0 &&
+        fpath.path().extension() == ".so") {
+      // filename starts with libdali_ and ends with .so
+      auto p = fpath.path().string();
+      std::cout << "Loading " << p << "\n";
+      PluginManager::LoadLibrary(std::move(p), global_symbols);
+    }
+  }
+}
+
+inline void PreloadPluginList(const std::string& dali_preload_plugins) {
+  const char delimiter = ':';
+  size_t previous = 0;
+  std::string preload(dali_preload_plugins);
+  size_t index = dali_preload_plugins.find(delimiter);
+  std::vector<std::string> plugins;
+  while (index != string::npos) {
+    auto plugin_path = dali_preload_plugins.substr(previous, index - previous);
+    std::cout << "Loading " << plugin_path << "\n";
+    PluginManager::LoadLibrary(plugin_path);
+    previous = index + 1;
+    index = dali_preload_plugins.find(delimiter, previous);
+  }
+  auto plugin_path = dali_preload_plugins.substr(previous);
+  std::cout << "Loading " << plugin_path << "\n";
+  PluginManager::LoadLibrary(plugin_path);
+  plugins.push_back(dali_preload_plugins.substr(previous));
+}
+
+void PluginManager::LoadDefaultPlugins() {
+  static bool run_once = []() {
+    std::string preload_plugins_str = "default";
+    const char* dali_preload_plugins = std::getenv("DALI_PRELOAD_PLUGINS");
+    if (dali_preload_plugins)
+      preload_plugins_str = dali_preload_plugins;
+    if (preload_plugins_str == "default") {
+      PluginManager::LoadDirectory(DefaultPluginPath());
+    } else if (fs::is_directory(preload_plugins_str)) {
+      PluginManager::LoadDirectory(preload_plugins_str);
+    } else {
+      PreloadPluginList(dali_preload_plugins);
+    }
+    return true;
+  }();
+  (void) run_once;
 }
 
 }  // namespace dali

dali/plugin/plugin_manager.h

@@ -16,21 +16,44 @@
 #define DALI_PLUGIN_PLUGIN_MANAGER_H_
 
 #include <string>
+#include <vector>
 #include "dali/core/common.h"
 
 namespace dali {
 
 class DLL_PUBLIC PluginManager {
  public:
-    /**
-     * @brief Load plugin library
-     * @remarks Will invoke dlopen()
-     * @param [in] lib_path path to the plugin library, e.g. "/usr/lib/libcustomplugin.so"
-     * @param [in] global_symbols if true, the library is loaded with RTLD_GLOBAL flag or equivalent
-*                                 otherwise, RTLD_LOCAL is used
-     * @throws std::runtime_error if the library could not be loaded
-     */
-    static DLL_PUBLIC void LoadLibrary(const std::string& lib_path, bool global_symbols = false);
+  /**
+   * @brief Load plugin library
+   * @remarks Will invoke dlopen()
+   * @param [in] lib_path path to the plugin library, e.g. "/usr/lib/libcustomplugin.so"
+   * @param [in] global_symbols if true, the library is loaded with RTLD_GLOBAL flag or equivalent
+   *                            otherwise, RTLD_LOCAL is used
+   * @throws std::runtime_error if the library could not be loaded
+   */
+  static DLL_PUBLIC void LoadLibrary(const std::string& lib_path, bool global_symbols = false);
+
+  /**
+   * @brief Load plugin directory. The plugin paths will have the following pattern:
+   *        {lib_path}/{sub_path}/lib_dali{plugin_name}.so
+   * @param [in] lib_path path to the root directory where the plugins are located
+   * @param [in] global_symbols if true, the library is loaded with RTLD_GLOBAL flag or equivalent
+   *                            otherwise, RTLD_LOCAL is used
+   * @throws std::runtime_error if the library could not be loaded
+   */
+  static DLL_PUBLIC void LoadDirectory(const std::string& lib_path, bool global_symbols = false);
+
+  /**
+   * @brief Load default plugin library
+   * @remarks DALI_PRELOAD_PLUGINS are environment variables that can be used to control what
+   * plugins are loaded. If the variable is set, it can be one of the following:
+   *          - The string "default", in which case the default plugin directory is scanned (a
+   * subdir plugins under the DALI installation directory)
+   *          - A path to a directory which will be used for the plugin search
+   *          - A string containing colon (:) separated paths to plugin files, which will be
+   * preloaded If the variable is unset, it behaves as "default"
+   */
+  static DLL_PUBLIC void LoadDefaultPlugins();
 };
 
 }  // namespace dali

dali/plugin/plugin_manager_test.cc

@@ -1,4 +1,4 @@
-// Copyright (c) 2018-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright (c) 2018-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -20,12 +20,12 @@ const char kNonExistingLibName[] = "not_a_dali_plugin.so";
 
 static const std::string& DummyPluginLibPath() {
   static const std::string plugin_lib =
-      dali::test::CurrentExecutableDir() + "/libcustomdummyplugin.so";
+      dali::test::CurrentExecutableDir() + "/libdali_customdummyplugin.so";
   return plugin_lib;
 }
 
 static const std::string& DummyPluginLibPathGlobal() {
-  static const std::string plugin_lib = "libcustomdummyplugin.so";
+  static const std::string plugin_lib = "libdali_customdummyplugin.so";
   return plugin_lib;
 }
 

dali/python/backend_impl.cc

@@ -1664,6 +1664,12 @@ PYBIND11_MODULE(backend_impl, m) {
     py::arg("lib_path"),
     py::arg("global_symbols") = false);
 
+  m.def("LoadDirectory", &PluginManager::LoadDirectory,
+    py::arg("dir_path"),
+    py::arg("global_symbols") = false);
+
+  m.def("LoadDefaultPlugins", &PluginManager::LoadDefaultPlugins);
+
   m.def("GetCxx11AbiFlag", &GetCxx11AbiFlag);
 
   m.def("IsDriverInitialized", [] {

dali/python/nvidia/dali/backend.py

@@ -15,7 +15,6 @@
 from nvidia.dali.backend_impl import (
     Init,
     OpSpec,
-    LoadLibrary,
     GetCudaVersion,
     GetCufftVersion,
     GetNppVersion,
@@ -31,11 +30,6 @@
 import os
 import sys
 
-# Note: If we ever need to add more complex functionality
-# for importing the DALI c++ extensions, we can do it here
-
-default_plugins = []
-
 
 def deprecation_warning(what):
     # show only this warning
@@ -78,9 +72,6 @@ def deprecation_warning(what):
             "Please update your environment to CUDA version 11 or newer."
         )
 
-    for lib in default_plugins:
-        LoadLibrary(os.path.join(os.path.dirname(__file__), lib))
-
 cuda_checked = False
 
 

dali/python/nvidia/dali/plugin_manager.py

@@ -33,3 +33,22 @@ def load_library(library_path: str, global_symbols: bool = False):
     """
     b.LoadLibrary(library_path, global_symbols)
     ops.Reload()
+
+def load_directory(plugin_dir_path: str, global_symbols: bool = False):
+    """Loads a DALI plugin directory, containing one or more DALI plugins, following the pattern:
+    {plugin_dir_path}/{sub_path}/libdali_{plugin_name}.so
+
+    Args:
+        plugin_dir_path: Path to the directory to search for plugins
+        global_symbols: If ``True``, the library is loaded with ``RTLD_GLOBAL`` flag or equivalent;
+            otherwise ``RTLD_LOCAL`` is used. Some libraries (for example Halide) require being
+            loaded with ``RTLD_GLOBAL`` - use this setting if your plugin uses any such library.
+
+    Returns:
+        None.
+
+    Raises:
+        RuntimeError: when unable to load the library.
+    """
+    b.LoadDirectory(plugin_dir_path, global_symbols)
+    ops.Reload()