Env amends

src/pb_env.cc

@@ -228,6 +228,9 @@ RecursiveDirectoryDelete(const char* dir)
 
 EnvironmentManager::EnvironmentManager()
 {
+  LOG_MESSAGE(
+      TRITONSERVER_LOG_VERBOSE,
+      "EnvironmentManager constructor: initializing Python env manager");
   char tmp_dir_template[PATH_MAX + 1];
   strcpy(tmp_dir_template, "/tmp/python_env_XXXXXX");
 
@@ -239,13 +242,14 @@ EnvironmentManager::EnvironmentManager()
   strcpy(base_path_, tmp_dir_template);
 }
 
-std::string
-EnvironmentManager::ExtractIfNotExtracted(std::string env_path)
+std::shared_ptr<
+    Environment>  // TODO: write logic with shared and weak ptrs in this method
+EnvironmentManager::ExtractIfNotExtracted(const std::string& env_path)
 {
   // Lock the mutex. Only a single thread should modify the map.
   std::lock_guard<std::mutex> lk(mutex_);
-  char canonical_env_path[PATH_MAX + 1];
 
+  char canonical_env_path[PATH_MAX + 1];
   char* err = realpath(env_path.c_str(), canonical_env_path);
   if (err == nullptr) {
     throw PythonBackendException(
@@ -270,19 +274,24 @@ EnvironmentManager::ExtractIfNotExtracted(std::string env_path)
                      "not contain compressed path. Path: ") +
          canonical_env_path)
             .c_str());
-    return canonical_env_path;
+    return nullptr;
   }
-  const auto env_itr = env_map_.find(canonical_env_path);
+
+  std::string canonical_env_path_str(canonical_env_path);
+  std::string env_key = canonical_env_path_str;
+  const auto env_itr = env_map_[env_key];
+  std::shared_ptr<Environment> env;
   if (env_itr != env_map_.end()) {
+    env = env_itr->second.lock();
     // Check if the environment has been modified and would
-    // need to be extracted again.
-    if (env_itr->second.second == last_modified_time) {
+    // need to be extracted again (or the current environment has no owners
+    // anymore).
+    if (env->LastModifiedTime() == last_modified_time) {
       env_extracted = true;
     } else {
       // Environment file has been updated. Need to clear
       // the previously extracted environment and extract
       // the environment to the same destination directory.
-      RecursiveDirectoryDelete(env_itr->second.first.c_str());
       re_extraction = true;
     }
   }
@@ -291,44 +300,84 @@ EnvironmentManager::ExtractIfNotExtracted(std::string env_path)
   if (!env_extracted) {
     LOG_MESSAGE(
         TRITONSERVER_LOG_VERBOSE,
-        (std::string("Extracting Python execution env ") + canonical_env_path)
+        (std::string("Extracting Python execution env ") +
+         canonical_env_path)
             .c_str());
+
     std::string dst_env_path;
     if (re_extraction) {
-      dst_env_path = env_map_[canonical_env_path].first;
+      dst_env_path = env->Path();
     } else {
       dst_env_path =
-          std::string(base_path_) + "/" + std::to_string(env_map_.size());
+          std::string(base_path_) + "/" + std::to_string(env_path_counter_);
+      ++env_path_counter_;
     }
 
-    std::string canonical_env_path_str(canonical_env_path);
-
-    int status =
-        mkdir(dst_env_path.c_str(), S_IRWXU | S_IRWXG | S_IROTH | S_IXOTH);
-    if (status == 0) {
-      ExtractTarFile(canonical_env_path_str, dst_env_path);
-    } else {
-      throw PythonBackendException(
-          std::string("Failed to create environment directory for '") +
-          dst_env_path.c_str() + "'.");
-    }
-    if (re_extraction) {
-      // Just update the last modified timestamp
-      env_map_[canonical_env_path].second = last_modified_time;
+    if (re_extraction ) {
+      // Just replace with new environment (by updated source)
+      env->Update(last_modified_time);
     } else {
-      // Add the path to the list of environments
-      env_map_.insert({canonical_env_path, {dst_env_path, last_modified_time}});
+      // Add the environment to the list of environments
+      env = std::make_shared<Environment>(
+        canonical_env_path_str, dst_env_path, last_modified_time);
+      env->SetManager(this);
+      env_map_.insert({env_key, new_env});
     }
-    return dst_env_path;
-  } else {
-    return env_map_.find(canonical_env_path)->second.first;
   }
+
+  return env;
 }
 
 EnvironmentManager::~EnvironmentManager()
 {
   RecursiveDirectoryDelete(base_path_);
 }
+
+Environment::Environment(
+    const std::string& source, const std::string& path,
+    const time_t& last_modified_time)
+    : source_(source), path_(path),
+      last_modified_time_(std::to_string(last_modified_time))
+{
+  Extract();
+}
+
+void
+Environment::Extract()
+{
+  int status =
+      mkdir(dst_env_path.c_str(), S_IRWXU | S_IRWXG | S_IROTH | S_IXOTH);
+  if (status != 0) {
+    throw PythonBackendException(
+        std::string("Failed to create environment directory for '") +
+        dst_env_path.c_str() + "'.");
+  }
+  ExtractTarFile(source_, path_);
+}
+
+void
+Environment::Update(const time_t& last_modified_time)
+{
+  Delete();
+  Extract();
+  last_modified_time_ = last_modified_time;
+}
+
+void
+Environment::Delete()
+{
+  RecursiveDirectoryDelete(path_.c_str());
+}
+
+Environment::~Environment()
+{
+  if (manager_ != nullptr) {
+    manager_->env_map_.erase(source_);
+  }
+  Delete();
+}
+
+
 #endif
 
 }}}  // namespace triton::backend::python

src/pb_env.h

@@ -27,6 +27,7 @@
 #pragma once
 #include <climits>
 #include <map>
+#include <memory>
 #include <mutex>
 #include <string>
 
@@ -46,17 +47,46 @@ bool FileExists(std::string& path);
 //
 #ifndef _WIN32
 class EnvironmentManager {
-  std::map<std::string, std::pair<std::string, time_t>> env_map_;
-  char base_path_[PATH_MAX + 1];
-  std::mutex mutex_;
-
  public:
+  class Environment {
+   public:
+    friend class EnvironmentManager;
+    Environment(
+        const std::string& source, const std::string& path,
+        const time_t& last_modified_time);
+    void SetManager(EnvironmentManager* manager) { manager_ = manager; }
+
+    void Update(const time_t& last_modified_time);
+
+    const std::string& Source() const { return source_; }
+    const std::string& Path() const { return path_; }
+    const time_t& LastModifiedTime() const { return last_modified_time_; }
+    explicit operator std::string() const { return Path(); }
+
+   private:
+    void Extract();
+    void Delete();
+
+    std::string source_;
+    std::string path_;
+    time_t last_modified_time_;
+
+    EnvironmentManager* manager_ = nullptr;
+  };
+
   EnvironmentManager();
 
   // Extracts the tar.gz file in the 'env_path' if it has not been
   // already extracted.
-  std::string ExtractIfNotExtracted(std::string env_path);
+  std::shared_ptr<Environment> ExtractIfNotExtracted(
+      const std::string& env_path);
   ~EnvironmentManager();
+
+ private:
+  size_t env_path_counter_ = 0;
+  std::map<std::string, std::weak_ptr<Environment>> env_map_;
+  char base_path_[PATH_MAX + 1];
+  std::mutex mutex_;
 };
 #endif
 

src/stub_launcher.cc

@@ -64,7 +64,7 @@ StubLauncher::Initialize(ModelState* model_state)
   shm_growth_byte_size_ = model_state->StateForBackend()->shm_growth_byte_size;
   shm_message_queue_size_ =
       model_state->StateForBackend()->shm_message_queue_size;
-  python_execution_env_ = model_state->PythonExecutionEnv();
+  python_execution_env_source_ = model_state->PythonExecutionEnv();
   python_lib_ = model_state->StateForBackend()->python_lib;
   model_state->ModelConfig().Write(&model_config_buffer_);
   is_decoupled_ = model_state->IsDecoupled();
@@ -104,7 +104,7 @@ StubLauncher::Initialize(ModelState* model_state)
 
   // FIXME [DLIS-5969]: Enable for Windows when custom execution environments
   // are supported.
-  if (python_execution_env_ != "") {
+  if (python_execution_env_source_ != "") {
 #ifndef _WIN32
     RETURN_IF_ERROR(GetPythonEnvironment(model_state));
 #else
@@ -405,7 +405,7 @@ StubLauncher::Launch()
   // executables and libraries.
   ipc_control_->uses_env = false;
 
-  if (python_execution_env_ != "") {
+  if (python_execution_env_source_ != "") {
     ipc_control_->uses_env = true;
 
     // Parse environment variables from activation script
@@ -592,20 +592,20 @@ StubLauncher::Launch()
 TRITONSERVER_Error*
 StubLauncher::GetPythonEnvironment(ModelState* model_state)
 {
-  std::string python_execution_env = "";
   try {
-    python_execution_env =
-        model_state->StateForBackend()->env_manager->ExtractIfNotExtracted(
-            python_execution_env_);
+    python_execution_env_ =
+        model_state->StateForBackend()->env_manager->GetEnvironment(model_state);
   }
   catch (PythonBackendException& pb_exception) {
     return TRITONSERVER_ErrorNew(
         TRITONSERVER_ERROR_INTERNAL, pb_exception.what());
   }
 
-  path_to_activate_ = python_execution_env + "/bin/activate";
-  path_to_libpython_ = python_execution_env + "/lib";
-  if (python_execution_env.length() > 0 && !FileExists(path_to_activate_)) {
+  std::string python_execution_env_path = python_execution_env_->Path();
+  path_to_activate_ = python_execution_env_path + "/bin/activate";
+  path_to_libpython_ = python_execution_env_path + "/lib";
+  if (python_execution_env_path.length() > 0 &&
+      !FileExists(path_to_activate_)) {
     return TRITONSERVER_ErrorNew(
         TRITONSERVER_ERROR_INTERNAL,
         ("Path " + path_to_activate_ +

src/stub_launcher.h

@@ -50,6 +50,8 @@
 #include "triton/core/tritonbackend.h"
 #include "triton/core/tritonserver.h"
 
+#include "pb_env.h"
+
 namespace triton { namespace backend { namespace python {
 
 class ModelState;
@@ -203,7 +205,8 @@ class StubLauncher {
   // Path to python execution environment
   std::string path_to_libpython_;
   std::string path_to_activate_;
-  std::string python_execution_env_;
+  std::shared_ptr<EnvironmentManager::Environment> python_execution_env_;
+  std::string python_execution_env_source_;
 
   common::TritonJson::WriteBuffer model_config_buffer_;
   common::TritonJson::Value auto_complete_config_;