diff --git a/src/pb_env.cc b/src/pb_env.cc index d9643a62..b67b455f 100644 --- a/src/pb_env.cc +++ b/src/pb_env.cc @@ -36,6 +36,7 @@ #include #include #include +#include #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 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::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 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,37 +312,37 @@ 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 lk(mutex_); + + size_t env_owners_counter = env.RemoveOwner(); + if (env_owners_counter == 0) { + env_map_.erase(env.Source()); } } @@ -329,6 +350,76 @@ 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 diff --git a/src/pb_env.h b/src/pb_env.h index 04e01fa3..0c0f20d0 100644 --- a/src/pb_env.h +++ b/src/pb_env.h @@ -27,7 +27,9 @@ #pragma once #include #include +#include #include +#include #include #ifdef WIN32 @@ -46,17 +48,92 @@ bool FileExists(std::string& path); // #ifndef _WIN32 class EnvironmentManager { - std::map> 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 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 env_map_; + char base_path_[PATH_MAX + 1]; + std::mutex mutex_; }; #endif diff --git a/src/stub_launcher.cc b/src/stub_launcher.cc index 6a1c8f2b..1978afe8 100644 --- a/src/stub_launcher.cc +++ b/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,26 +592,31 @@ StubLauncher::Launch() TRITONSERVER_Error* StubLauncher::GetPythonEnvironment(ModelState* model_state) { - std::string python_execution_env = ""; + auto python_execution_env_source = model_state->PythonExecutionEnv(); try { - python_execution_env = + python_execution_env_ = model_state->StateForBackend()->env_manager->ExtractIfNotExtracted( - python_execution_env_); + python_execution_env_source); } 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_source; + if (python_execution_env_.has_value()) { + 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_ + " does not exist. The Python environment should contain an " - "'activate' script.") - .c_str()); + "'activate' script.").c_str()); } return nullptr; } diff --git a/src/stub_launcher.h b/src/stub_launcher.h index fba116df..72906798 100644 --- a/src/stub_launcher.h +++ b/src/stub_launcher.h @@ -36,6 +36,7 @@ #include #include #include +#include #include #include #include @@ -50,6 +51,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 +206,8 @@ class StubLauncher { // Path to python execution environment std::string path_to_libpython_; std::string path_to_activate_; - std::string python_execution_env_; + std::optional python_execution_env_; + std::string python_execution_env_source_; common::TritonJson::WriteBuffer model_config_buffer_; common::TritonJson::Value auto_complete_config_;