diff --git a/src/shm_manager.cc b/src/shm_manager.cc index 134cee6f..b59a2ba4 100644 --- a/src/shm_manager.cc +++ b/src/shm_manager.cc @@ -1,4 +1,4 @@ -// Copyright 2021-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// Copyright 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. // // Redistribution and use in source and binary forms, with or without // modification, are permitted provided that the following conditions @@ -26,13 +26,58 @@ #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)) { + if (std::atexit(CleanupParentShmRegions) != 0) { + std::cerr << "python_backend: failed to register atexit shm cleanup " + "handler; relying on TerminateStub for cleanup" + << std::endl; + } + } +} + +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 +184,7 @@ SharedMemoryManager::SharedMemoryManager( if (create) { *total_size_ = current_capacity_; new (shm_mutex_) bi::interprocess_mutex; + RegisterParentShmRegion(shm_region_name_); } } @@ -226,9 +272,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..dc25f781 100644 --- a/src/shm_manager.h +++ b/src/shm_manager.h @@ -1,4 +1,4 @@ -// Copyright 2021-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// Copyright 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. // // Redistribution and use in source and binary forms, with or without // modification, are permitted provided that the following conditions @@ -194,6 +194,9 @@ class SharedMemoryManager { void SetDeleteRegion(bool delete_region); + // Idempotent: removes the parent-owned shm region and clears delete_region_. + 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..c676c5b4 100644 --- a/src/stub_launcher.cc +++ b/src/stub_launcher.cc @@ -26,7 +26,10 @@ #include "stub_launcher.h" +#include +#include #include +#include #include "pb_utils.h" #include "python_be.h" @@ -42,7 +45,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 +54,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 +67,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; + 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_); @@ -801,15 +805,33 @@ void StubLauncher::TerminateStub() { if (is_initialized_) { + // Single teardown budget: finalize + wait share stub_timeout_seconds_. bool force_kill = false; + int64_t remaining_seconds = stub_timeout_seconds_; if (is_healthy_) { + const int64_t total_timeout_ms = stub_timeout_seconds_ * 1000; + // Clamp to INT_MAX; MessageQueue::Pop takes int. + const int pop_timeout_ms = static_cast(std::min( + total_timeout_ms, + static_cast(std::numeric_limits::max()))); // 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; + const auto pop_start = std::chrono::steady_clock::now(); + parent_message_queue_->Pop(pop_timeout_ms, success); + const int64_t pop_elapsed_s = + std::chrono::duration_cast( + std::chrono::steady_clock::now() - pop_start) + .count(); + remaining_seconds = + std::max(0, stub_timeout_seconds_ - pop_elapsed_s); + if (!success) { + force_kill = true; + } stub_message_queue_.reset(); parent_message_queue_.reset(); @@ -820,11 +842,15 @@ StubLauncher::TerminateStub() if (force_kill) { KillStubProcess(); - } else { - WaitForStubProcess(); + } else if (!WaitForStubProcessWithTimeout(remaining_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 +950,50 @@ 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); + } + + // Stub may have exited during the last sleep(1); recheck before killing. + { + int status; + pid_t ret = waitpid(stub_pid_, &status, WNOHANG); + if (ret == stub_pid_ || ret == -1) { + stub_pid_ = 0; + return true; + } + } + + 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_;