diff --git a/src/shm_manager.cc b/src/shm_manager.cc index 134cee6f..61fad9dc 100644 --- a/src/shm_manager.cc +++ b/src/shm_manager.cc @@ -26,13 +26,54 @@ #include "shm_manager.h" +#include #include #include #include +#include #include +#include +#include namespace triton { namespace backend { namespace python { +namespace { + +std::mutex parent_shm_regions_mu; +std::unordered_set parent_shm_regions; +std::atomic parent_shm_atexit_registered{false}; + +void +CleanupParentShmRegions() +{ + std::lock_guard lock(parent_shm_regions_mu); + for (const auto& region : parent_shm_regions) { + bi::shared_memory_object::remove(region.c_str()); + } + parent_shm_regions.clear(); +} + +void +RegisterParentShmRegion(const std::string& shm_region_name) +{ + { + std::lock_guard lock(parent_shm_regions_mu); + parent_shm_regions.insert(shm_region_name); + } + if (!parent_shm_atexit_registered.exchange(true)) { + std::atexit(CleanupParentShmRegions); + } +} + +void +UnregisterParentShmRegion(const std::string& shm_region_name) +{ + std::lock_guard lock(parent_shm_regions_mu); + parent_shm_regions.erase(shm_region_name); +} + +} // namespace + void CUDAMemoryPoolManager::SetCUDAPoolAddress( const int32_t device_id, void* cuda_pool_address) @@ -139,6 +180,7 @@ SharedMemoryManager::SharedMemoryManager( if (create) { *total_size_ = current_capacity_; new (shm_mutex_) bi::interprocess_mutex; + RegisterParentShmRegion(shm_region_name_); } } @@ -226,9 +268,17 @@ SharedMemoryManager::FreeMemory() SharedMemoryManager::~SharedMemoryManager() noexcept(false) +{ + RemoveShmRegion(); +} + +void +SharedMemoryManager::RemoveShmRegion() { if (delete_region_) { bi::shared_memory_object::remove(shm_region_name_.c_str()); + UnregisterParentShmRegion(shm_region_name_); + delete_region_ = false; } } diff --git a/src/shm_manager.h b/src/shm_manager.h index 8517faf3..49a61188 100644 --- a/src/shm_manager.h +++ b/src/shm_manager.h @@ -194,6 +194,9 @@ class SharedMemoryManager { void SetDeleteRegion(bool delete_region); + // Remove the parent-owned shared memory region from the filesystem. + void RemoveShmRegion(); + std::unique_ptr& GetCUDAMemoryPoolManager() { return cuda_memory_pool_manager_; diff --git a/src/stub_launcher.cc b/src/stub_launcher.cc index 6a1c8f2b..56befbed 100644 --- a/src/stub_launcher.cc +++ b/src/stub_launcher.cc @@ -42,7 +42,7 @@ namespace triton { namespace backend { namespace python { StubLauncher::StubLauncher(const std::string stub_process_kind) : parent_pid_(0), is_initialized_(false), stub_process_kind_(stub_process_kind), model_instance_name_(""), - device_id_(0), kind_("") + device_id_(0), kind_(""), stub_timeout_seconds_(30) { } @@ -51,7 +51,7 @@ StubLauncher::StubLauncher( const int32_t device_id, const std::string kind) : is_initialized_(false), stub_process_kind_(stub_process_kind), model_instance_name_(model_instance_name), device_id_(device_id), - kind_(kind) + kind_(kind), stub_timeout_seconds_(30) { } @@ -64,6 +64,8 @@ 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; + stub_timeout_seconds_ = + model_state->StateForBackend()->stub_timeout_seconds; python_execution_env_ = model_state->PythonExecutionEnv(); python_lib_ = model_state->StateForBackend()->python_lib; model_state->ModelConfig().Write(&model_config_buffer_); @@ -803,13 +805,18 @@ StubLauncher::TerminateStub() if (is_initialized_) { bool force_kill = false; if (is_healthy_) { + const int64_t timeout_ms = stub_timeout_seconds_ * 1000; // Finalize command does not have any arguments. std::unique_ptr ipc_message = IPCMessage::Create(shm_pool_, false /* inline_response */); ipc_message->Command() = PYTHONSTUB_FinalizeRequest; stub_message_queue_->Push(ipc_message->ShmHandle()); - parent_message_queue_->Pop(); + bool success = false; + parent_message_queue_->Pop(timeout_ms, success); + if (!success) { + force_kill = true; + } stub_message_queue_.reset(); parent_message_queue_.reset(); @@ -820,11 +827,15 @@ StubLauncher::TerminateStub() if (force_kill) { KillStubProcess(); - } else { - WaitForStubProcess(); + } else if (!WaitForStubProcessWithTimeout(stub_timeout_seconds_)) { + KillStubProcess(); } } + if (shm_pool_ != nullptr) { + shm_pool_->RemoveShmRegion(); + } + // First destroy the IPCControl. This makes sure that IPCControl is // destroyed before the shared memory manager goes out of scope. ipc_control_.reset(); @@ -924,10 +935,40 @@ StubLauncher::WaitForStubProcess() // Added this check to ensure server doesn't hang waiting after stub // process has already be killed and cannot be waited on waitpid(stub_pid_, &status, 0); + stub_pid_ = 0; } #endif } +bool +StubLauncher::WaitForStubProcessWithTimeout(int64_t timeout_seconds) +{ +#ifdef _WIN32 + WaitForStubProcess(); + return true; +#else + if (stub_pid_ == 0) { + return true; + } + + for (int64_t elapsed = 0; elapsed < timeout_seconds; ++elapsed) { + int status; + pid_t ret = waitpid(stub_pid_, &status, WNOHANG); + if (ret == stub_pid_) { + stub_pid_ = 0; + return true; + } + if (ret == -1) { + stub_pid_ = 0; + return true; + } + sleep(1); + } + + return false; +#endif +} + #ifdef TRITON_ENABLE_GPU void StubLauncher::ShareCUDAMemoryPool( diff --git a/src/stub_launcher.h b/src/stub_launcher.h index fba116df..ed6abbe7 100644 --- a/src/stub_launcher.h +++ b/src/stub_launcher.h @@ -161,6 +161,9 @@ class StubLauncher { // Wait for stub process void WaitForStubProcess(); + // Wait for stub process with timeout. Returns true if the stub exited. + bool WaitForStubProcessWithTimeout(int64_t timeout_seconds); + #ifndef _WIN32 // FIXME [DLIS-5969]: Enable for Windows when custom execution environments // are supported. @@ -199,6 +202,7 @@ class StubLauncher { int64_t shm_default_byte_size_; int64_t shm_growth_byte_size_; int64_t shm_message_queue_size_; + int64_t stub_timeout_seconds_; // Path to python execution environment std::string path_to_libpython_;