diff --git a/src/pb_env.cc b/src/pb_env.cc index d9643a62..6c61da4e 100644 --- a/src/pb_env.cc +++ b/src/pb_env.cc @@ -240,95 +240,149 @@ EnvironmentManager::EnvironmentManager() } std::string -EnvironmentManager::ExtractIfNotExtracted(std::string env_path) +EnvironmentManager::ExtractIfNotExtracted(const 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; - // 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(env_path.c_str(), &info) != 0) { throw PythonBackendException( - std::string("stat() of : ") + canonical_env_path + " returned error."); + "stat() of : " + 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: ") + - canonical_env_path) + ("Returning path since EXECUTION_ENV_PATH does " + "not contain compressed path. Path: " + + env_path) .c_str()); - return canonical_env_path; + return env_path; } - const auto env_itr = env_map_.find(canonical_env_path); - if (env_itr != env_map_.end()) { + + // 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); + + auto env_itr = env_map_.find(env_path); + // Extract only if the env has not been extracted yet. + if (env_itr == env_map_.end()) { + LOG_MESSAGE( + TRITONSERVER_LOG_VERBOSE, + ("Extracting Python execution env " + env_path).c_str()); + + 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_path, env_path, dst_env_path, last_modified_time) + .first; + } else { + Environment& 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) { - env_extracted = true; - } else { + if (env.LastModifiedTime() != last_modified_time) { + LOG_MESSAGE( + TRITONSERVER_LOG_VERBOSE, + ("Re-extracting Python execution env " + env_path).c_str()); // 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; + env.Update(last_modified_time); } } - // Extract only if the env has not been extracted yet. - if (!env_extracted) { - LOG_MESSAGE( - TRITONSERVER_LOG_VERBOSE, - (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; - } else { - dst_env_path = - std::string(base_path_) + "/" + std::to_string(env_map_.size()); - } + Environment& env = env_itr->second; - std::string canonical_env_path_str(canonical_env_path); + // Reference counter must be incremented on each ExtractIfNotExtracted call. + env.IncrementRefCount(); - 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}}); + LOG_MESSAGE( + TRITONSERVER_LOG_VERBOSE, + ("Successfully extracted Python execution env " + env_path).c_str()); + + return env.Destination(); +} + +void +EnvironmentManager::DropEnvironment(const std::string& env_path) +{ + LOG_MESSAGE( + TRITONSERVER_LOG_VERBOSE, + ("Trying to drop Python execution env " + env_path).c_str()); + + std::lock_guard lk(mutex_); + + auto env_itr = env_map_.find(env_path); + if (env_itr != env_map_.end()) { + if (env_itr->second.DecrementRefCount() == 0) { + env_map_.erase(env_itr); + LOG_MESSAGE( + TRITONSERVER_LOG_VERBOSE, + ("Successfully dropped Python execution env " + env_path).c_str()); } - return dst_env_path; } else { - return env_map_.find(canonical_env_path)->second.first; + LOG_MESSAGE( + TRITONSERVER_LOG_VERBOSE, ("The environment with the key '" + env_path + + "' is not presented the env_map") + .c_str()); } } EnvironmentManager::~EnvironmentManager() { + // All environments must be deleted before EnvironmentManager is called + if (!env_map_.empty()) { + LOG_MESSAGE( + TRITONSERVER_LOG_ERROR, "Environment mananager map is not empty"); + } + RecursiveDirectoryDelete(base_path_); } + +EnvironmentManager::Environment::Environment( + const std::string& source, const std::string& destination, + const time_t& last_modified_time) + : source_(source), destination_(destination), + last_modified_time_(last_modified_time) +{ + Extract(); +} + +void +EnvironmentManager::Environment::Extract() +{ + int status = + mkdir(destination_.c_str(), S_IRWXU | S_IRWXG | S_IROTH | S_IXOTH); + if (status != 0) { + throw PythonBackendException( + "Failed to create environment directory for '" + destination_ + "'."); + } + ExtractTarFile(source_, destination_); +} + +void +EnvironmentManager::Environment::Update(const time_t& last_modified_time) +{ + Delete(); + Extract(); + last_modified_time_ = last_modified_time; +} + +void +EnvironmentManager::Environment::Delete() +{ + RecursiveDirectoryDelete(destination_.c_str()); +} + +EnvironmentManager::Environment::~Environment() +{ + Delete(); +} + #endif }}} // namespace triton::backend::python diff --git a/src/pb_env.h b/src/pb_env.h index 04e01fa3..0bbf7d4d 100644 --- a/src/pb_env.h +++ b/src/pb_env.h @@ -46,17 +46,51 @@ 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& destination, + const time_t& last_modified_time); + ~Environment(); + + void Update(const time_t& last_modified_time); + void IncrementRefCount() { ++ref_count_; } + size_t DecrementRefCount() { return --ref_count_; } + + const std::string& Source() const { return source_; } + const std::string& Destination() const { return destination_; } + const time_t& LastModifiedTime() const { return last_modified_time_; } + + private: + void Extract(); + void Delete(); + + std::string source_; + std::string destination_; + time_t last_modified_time_; + + size_t ref_count_ = 0; + }; + EnvironmentManager(); // 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 + std::string ExtractIfNotExtracted(const std::string& env_path); + ~EnvironmentManager(); + + // Decrement the refcount for the environment identified by + // env_path. If the refcount reaches zero, the environment is + // removed from the map. + void DropEnvironment(const std::string& env_path); + + private: + size_t env_path_counter_ = 0; + std::map env_map_; + char base_path_[PATH_MAX + 1]; + std::mutex mutex_; }; #endif diff --git a/src/python_be.cc b/src/python_be.cc index 4c9ca49e..78c7d622 100644 --- a/src/python_be.cc +++ b/src/python_be.cc @@ -2027,7 +2027,8 @@ ModelInstanceState::~ModelInstanceState() } // Terminate stub first to allow any last messages to be received by the back // end before deallocating the queue memory - Stub()->TerminateStub(); + ModelState* model_state = reinterpret_cast(Model()); + Stub()->TerminateStub(model_state); TerminateMonitor(); Stub()->ClearQueues(); Stub().reset(); @@ -2056,7 +2057,7 @@ ModelState::Create(TRITONBACKEND_Model* triton_model, ModelState** state) RETURN_IF_ERROR((*state)->SetModelConfig()); (*state)->Stub()->UpdateHealth(); - (*state)->Stub()->TerminateStub(); + (*state)->Stub()->TerminateStub(*state); (*state)->Stub()->ClearQueues(); (*state)->Stub().reset(); } @@ -2172,7 +2173,7 @@ ModelState::LaunchAutoCompleteStubProcess() } catch (const BackendModelException& ex) { Stub()->UpdateHealth(); - Stub()->TerminateStub(); + Stub()->TerminateStub(this); Stub()->ClearQueues(); Stub().reset(); RETURN_ERROR_IF_TRUE( diff --git a/src/stub_launcher.cc b/src/stub_launcher.cc index 6a1c8f2b..cc976777 100644 --- a/src/stub_launcher.cc +++ b/src/stub_launcher.cc @@ -106,6 +106,19 @@ StubLauncher::Initialize(ModelState* model_state) // are supported. if (python_execution_env_ != "") { #ifndef _WIN32 + // Resolve symlinks to avoid duplicate environment entries. + char canonical_env_path[PATH_MAX + 1]; + char* err = realpath(python_execution_env_.c_str(), canonical_env_path); + if (err == nullptr) { + return TRITONSERVER_ErrorNew( + TRITONSERVER_ERROR_INTERNAL, + ("Failed to get the canonical path for " + python_execution_env_ + + ".") + .c_str()); + } + + python_execution_env_ = canonical_env_path; + RETURN_IF_ERROR(GetPythonEnvironment(model_state)); #else return TRITONSERVER_ErrorNew( @@ -798,7 +811,7 @@ StubLauncher::UpdateHealth() } void -StubLauncher::TerminateStub() +StubLauncher::TerminateStub(ModelState* model_state) { if (is_initialized_) { bool force_kill = false; @@ -825,6 +838,12 @@ StubLauncher::TerminateStub() } } +// Nothing will happen if python_execution_env_ doesn't exist +#ifndef _WIN32 + model_state->StateForBackend()->env_manager->DropEnvironment( + python_execution_env_); +#endif + // First destroy the IPCControl. This makes sure that IPCControl is // destroyed before the shared memory manager goes out of scope. ipc_control_.reset(); diff --git a/src/stub_launcher.h b/src/stub_launcher.h index fba116df..cd0873ab 100644 --- a/src/stub_launcher.h +++ b/src/stub_launcher.h @@ -145,7 +145,7 @@ class StubLauncher { } // Destruct Stub process - void TerminateStub(); + void TerminateStub(ModelState* model_state); // Reset log queue and bls decoupled queue pointers void ClearQueues();