Limit Python based backend search to backend directory

triton-inference-server · Nov 28, 2023 · 0d89719 · 0d89719
1 parent 1c2c618
commit 0d89719
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 6 deletions.
diff --git a/src/backend_model.cc b/src/backend_model.cc
@@ -312,8 +312,8 @@ TritonModel::GetBackendLibraryProperties(
   std::string backend_libname = model_config->runtime();
   if (backend_libname.empty()) {
     RETURN_IF_ERROR(GetBackendRuntimeLibraryName(
-        backend_name, *search_paths, &backend_libname, backend_libdir,
-        backend_libpath, is_python_based_backend));
+        backend_dir, backend_name, *search_paths, &backend_libname,
+        backend_libdir, backend_libpath, is_python_based_backend));
     if (!*is_python_based_backend) {
       // All variables are correctly set for C++ backends on initial search.
       return Status::Success;
@@ -376,7 +376,7 @@ TritonModel::GetBackendLibraryProperties(
 
 Status
 TritonModel::GetBackendRuntimeLibraryName(
-    const std::string& backend_name,
+    const std::string& backend_dir, const std::string& backend_name,
     const std::vector<std::string>& search_paths, std::string* backend_libname,
     std::string* backend_libdir, std::string* backend_libpath,
     bool* is_python_based_backend)
@@ -392,7 +392,7 @@ TritonModel::GetBackendRuntimeLibraryName(
   // Try Python runtime
   *backend_libname = kPythonFilename;
   RETURN_IF_ERROR(FindBackendLibraryPath(
-      search_paths, *backend_libname, backend_libdir, backend_libpath));
+      {backend_dir}, *backend_libname, backend_libdir, backend_libpath));
   if (!backend_libpath->empty()) {
     *is_python_based_backend = true;
     return Status::Success;

diff --git a/src/backend_model.h b/src/backend_model.h
@@ -212,9 +212,10 @@ class TritonModel : public Model {
 
   // Get 'backend_libname', 'backend_libdir', 'backend_libpath' and
   // 'is_python_based_backend' by searching for different possible backend
-  // library names on 'search_paths'
+  // library names on 'search_paths'. Searching for Python based backend
+  // runtime is limited to 'backend_dir'.
   static Status GetBackendRuntimeLibraryName(
-      const std::string& backend_name,
+      const std::string& backend_dir, const std::string& backend_name,
       const std::vector<std::string>& search_paths,
       std::string* backend_libname, std::string* backend_libdir,
       std::string* backend_libpath, bool* is_python_based_backend);