Env amends

src/pb_env.cc

@@ -36,6 +36,7 @@
 #include <cstdlib>
 #include <cstring>
 #include <iostream>
+#include <string>
 
 #include "pb_utils.h"
 
@@ -228,6 +229,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,50 +243,67 @@ EnvironmentManager::EnvironmentManager()
   strcpy(base_path_, tmp_dir_template);
 }
 
-std::string
-EnvironmentManager::ExtractIfNotExtracted(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* err = realpath(env_path.c_str(), canonical_env_path);
-  if (err == nullptr) {
-    throw PythonBackendException(
-        std::string("Failed to get the canonical path for ") + env_path + ".");
-  }
-
-  time_t last_modified_time;
-  LastModifiedTime(canonical_env_path, &last_modified_time);
 
-  bool env_extracted = false;
-  bool re_extraction = false;
+std::optional<EnvironmentManager::EnvironmentGuard>
+EnvironmentManager::ExtractIfNotExtracted(const std::string& env_path)
+{
+  std::string canonical_env_path = [&] {
+    char canonical_env_path[PATH_MAX + 1];
+    char* err = realpath(env_path.c_str(), canonical_env_path);
+    if (err == nullptr) {
+      throw PythonBackendException(
+          "Failed to get the canonical path for " + env_path + ".");
+    }
+    return std::string(canonical_env_path);
+  }();
 
   // If the path is not a conda-packed file, then bypass the extraction process
   struct stat info;
-  if (stat(canonical_env_path, &info) != 0) {
+  if (stat(canonical_env_path.c_str(), &info) != 0) {
     throw PythonBackendException(
-        std::string("stat() of : ") + canonical_env_path + " returned error.");
+        "stat() of : " + canonical_env_path + " returned error.");
   } else if (S_ISDIR(info.st_mode)) {
     LOG_MESSAGE(
         TRITONSERVER_LOG_VERBOSE,
-        (std::string("Returning canonical path since EXECUTION_ENV_PATH does "
-                     "not contain compressed path. Path: ") +
+        ("Returning canonical path since EXECUTION_ENV_PATH does "
+         "not contain compressed path. Path: " +
          canonical_env_path)
             .c_str());
-    return canonical_env_path;
+    return std::nullopt;
   }
-  const auto env_itr = env_map_.find(canonical_env_path);
+
+  auto& env = GetEnvironment(canonical_env_path);
+  return EnvironmentGuard(this, &env);
+}
+
+EnvironmentManager::Environment&
+EnvironmentManager::GetEnvironment(const std::string& env_path)
+{
+  // Lock the mutex. Only a single thread should modify the map.
+  std::lock_guard<std::mutex> lk(mutex_);
+
+  time_t last_modified_time;
+  LastModifiedTime(env_path, &last_modified_time);
+
+  bool env_extracted = false;
+  bool re_extraction = false;
+
+  const std::string& env_key = env_path;
+  auto env_itr = env_map_.find(env_key);
+  Environment* env = nullptr;
   if (env_itr != env_map_.end()) {
+    env = &env_itr->second;
+
     // 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 +312,114 @@ EnvironmentManager::ExtractIfNotExtracted(std::string env_path)
   if (!env_extracted) {
     LOG_MESSAGE(
         TRITONSERVER_LOG_VERBOSE,
-        (std::string("Extracting Python execution env ") + canonical_env_path)
-            .c_str());
-    std::string dst_env_path;
+        ("Extracting Python execution env " + env_path).c_str());
+
     if (re_extraction) {
-      dst_env_path = env_map_[canonical_env_path].first;
+      // Just replace with new environment (by updated source)
+      env->Update(last_modified_time);
     } else {
-      dst_env_path =
-          std::string(base_path_) + "/" + std::to_string(env_map_.size());
+      std::string dst_env_path =
+          std::string(base_path_) + "/" + std::to_string(env_path_counter_);
+      ++env_path_counter_;
+
+      // Add the environment to the list of environments
+      env_itr =
+          env_map_
+              .try_emplace(env_key, env_path, dst_env_path, last_modified_time)
+              .first;
+      env = &env_itr->second;
     }
+  }
 
-    std::string canonical_env_path_str(canonical_env_path);
+  env->AddOwner();
+  return *env;
+}
 
-    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;
-    } else {
-      // Add the path to the list of environments
-      env_map_.insert({canonical_env_path, {dst_env_path, last_modified_time}});
-    }
-    return dst_env_path;
-  } else {
-    return env_map_.find(canonical_env_path)->second.first;
+void
+EnvironmentManager::DropEnvironment(Environment& env)
+{
+  std::lock_guard<std::mutex> lk(mutex_);
+
+  size_t env_owners_counter = env.RemoveOwner();
+  if (env_owners_counter == 0) {
+    env_map_.erase(env.Source());
   }
 }
 
 EnvironmentManager::~EnvironmentManager()
 {
   RecursiveDirectoryDelete(base_path_);
 }
+
+EnvironmentManager::Environment::Environment(
+    const std::string& source, const std::string& path,
+    const time_t& last_modified_time)
+    : source_(source), path_(path), last_modified_time_(last_modified_time)
+{
+  Extract();
+}
+
+void
+EnvironmentManager::Environment::Extract()
+{
+  int status = mkdir(path_.c_str(), S_IRWXU | S_IRWXG | S_IROTH | S_IXOTH);
+  if (status != 0) {
+    throw PythonBackendException(
+        "Failed to create environment directory for '" + path_ + "'.");
+  }
+  ExtractTarFile(source_, path_);
+}
+
+void
+EnvironmentManager::Environment::Update(const time_t& last_modified_time)
+{
+  Delete();
+  Extract();
+  last_modified_time_ = last_modified_time;
+}
+
+void
+EnvironmentManager::Environment::Delete()
+{
+  RecursiveDirectoryDelete(path_.c_str());
+}
+
+EnvironmentManager::Environment::~Environment()
+{
+  Delete();
+}
+
+EnvironmentManager::EnvironmentGuard::EnvironmentGuard(
+    EnvironmentManager* manager, Environment* env)
+    : manager_(manager), environment_(env), environment_proxy_(env)
+{
+}
+
+EnvironmentManager::EnvironmentGuard::EnvironmentGuard(
+    EnvironmentGuard&& other_guard)
+    : manager_(other_guard.manager_), environment_(other_guard.environment_),
+      environment_proxy_(std::move(other_guard.environment_proxy_))
+{
+  other_guard.manager_ = nullptr;
+  other_guard.environment_ = nullptr;
+}
+
+EnvironmentManager::EnvironmentGuard&
+EnvironmentManager::EnvironmentGuard::operator=(EnvironmentGuard&& other_guard)
+{
+  EnvironmentGuard new_guard(std::move(other_guard));
+  std::swap(*this, new_guard);
+  return *this;
+}
+
+EnvironmentManager::EnvironmentGuard::~EnvironmentGuard()
+{
+  if (environment_ != nullptr && manager_ != nullptr) {
+    manager_->DropEnvironment(*environment_);
+  }
+}
+
+
 #endif
 
 }}}  // namespace triton::backend::python

src/pb_env.h

@@ -27,7 +27,9 @@
 #pragma once
 #include <climits>
 #include <map>
+#include <memory>
 #include <mutex>
+#include <optional>
 #include <string>
 
 #ifdef WIN32
@@ -46,17 +48,92 @@ 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:
+    Environment(
+        const std::string& source, const std::string& path,
+        const time_t& last_modified_time);
+    ~Environment();
+
+    void Update(const time_t& last_modified_time);
+    void AddOwner() { ++owners_counter_; }
+    size_t RemoveOwner() { return --owners_counter_; }
+
+    const std::string& Source() const { return source_; }
+    const std::string& Path() const { return path_; }
+    const time_t& LastModifiedTime() const { return last_modified_time_; }
+
+   private:
+    void Extract();
+    void Delete();
+
+    std::string source_;
+    std::string path_;
+    time_t last_modified_time_;
+
+    size_t owners_counter_ = 0;
+  };
+
+  class EnvironmentProxy {
+   public:
+    EnvironmentProxy(const Environment* env) : env_(env) {}
+
+    EnvironmentProxy(EnvironmentProxy&& other_proxy) : env_(other_proxy.env_)
+    {
+      other_proxy.env_ = nullptr;
+    }
+
+    ~EnvironmentProxy() = default;
+
+    const std::string& Source() const & { return env_->Source(); }
+    const std::string& Path() const & { return env_->Path(); }
+    const time_t& LastModifiedTime() const & { return env_->LastModifiedTime(); }
+
+   private:
+    const Environment* env_;
+  };
+
+  class EnvironmentGuard {
+   public:
+    EnvironmentGuard(EnvironmentManager* manager, Environment* environment);
+
+    EnvironmentGuard(const EnvironmentGuard&) = delete;
+    EnvironmentGuard(EnvironmentGuard&&);
+
+    EnvironmentGuard& operator=(const EnvironmentGuard&) = delete;
+    EnvironmentGuard& operator=(EnvironmentGuard&&);
+
+    const EnvironmentProxy* operator->() const { return &environment_proxy_; }
+    const EnvironmentProxy& operator*() const { return environment_proxy_; }
+
+    ~EnvironmentGuard();
+
+   private:
+    EnvironmentManager* manager_;
+    Environment* environment_;
+    EnvironmentProxy environment_proxy_;
+  };
+
   EnvironmentManager();
+  friend class EnvironmentGuard;
 
   // Extracts the tar.gz file in the 'env_path' if it has not been
-  // already extracted.
-  std::string ExtractIfNotExtracted(std::string env_path);
+  // already extracted. Returns nullopt when env_path is an uncompressed
+  // directory (caller uses that path directly).
+  std::optional<EnvironmentGuard> ExtractIfNotExtracted(
+      const std::string& env_path);
+
   ~EnvironmentManager();
+
+ private:
+  void DropEnvironment(Environment& environment);
+  Environment& GetEnvironment(const std::string& env_path);
+
+  size_t env_path_counter_ = 0;
+  std::map<std::string, Environment> env_map_;
+  char base_path_[PATH_MAX + 1];
+  std::mutex mutex_;
 };
 #endif