From e94075ec223b246b1f9056f3da6036261a5593a4 Mon Sep 17 00:00:00 2001 From: Vinya Kestur Date: Thu, 16 Jul 2026 16:21:44 +0000 Subject: [PATCH 1/4] fix: ensure Python backend shm cleanup on forced shutdown Register parent-owned shm regions for atexit cleanup, remove regions explicitly in TerminateStub, and honor stub-timeout-seconds during stub teardown to avoid orphaned regions when server exit times out. --- src/shm_manager.cc | 50 +++++++++++++++++++++++++++++++++++++++++++ src/shm_manager.h | 3 +++ src/stub_launcher.cc | 51 +++++++++++++++++++++++++++++++++++++++----- src/stub_launcher.h | 4 ++++ 4 files changed, 103 insertions(+), 5 deletions(-) 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_; From 1852204b8380261536f43ed96dc828b18eb37846 Mon Sep 17 00:00:00 2001 From: Vinya Kestur Date: Thu, 23 Jul 2026 18:49:31 +0000 Subject: [PATCH 2/4] style: fix pre-commit (copyright bump + clang-format) - Bump copyright to 2021-2026 on src/shm_manager.{cc,h}. - Collapse wrapped stub_timeout_seconds_ assignment per clang-format. --- src/shm_manager.cc | 2 +- src/shm_manager.h | 2 +- src/stub_launcher.cc | 3 +-- 3 files changed, 3 insertions(+), 4 deletions(-) diff --git a/src/shm_manager.cc b/src/shm_manager.cc index 61fad9dc..ca548af2 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 diff --git a/src/shm_manager.h b/src/shm_manager.h index 49a61188..a8e1c7fc 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 diff --git a/src/stub_launcher.cc b/src/stub_launcher.cc index 56befbed..91b2bbc0 100644 --- a/src/stub_launcher.cc +++ b/src/stub_launcher.cc @@ -64,8 +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; - stub_timeout_seconds_ = - model_state->StateForBackend()->stub_timeout_seconds; + 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_); From aa7708f38f2332ce756ff82e91712456a80b8987 Mon Sep 17 00:00:00 2001 From: Vinya Kestur Date: Thu, 23 Jul 2026 19:00:34 +0000 Subject: [PATCH 3/4] =?UTF-8?q?fix:=20address=20Greptile=20review=20?= =?UTF-8?q?=E2=80=94=20bounded=20total=20shutdown=20budget?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Enforce stub_timeout_seconds_ as a single total budget across the finalize-wait and process-exit-wait phases; healthy teardown previously could take up to 2 * stub_timeout_seconds_. - Clamp the Pop timeout to INT_MAX to avoid narrowing when passing int64_t into MessageQueue::Pop(int const&). - Re-poll waitpid once after the final sleep window in WaitForStubProcessWithTimeout so a stub that exits during that second is not force-killed unnecessarily. --- src/stub_launcher.cc | 36 +++++++++++++++++++++++++++++++++--- 1 file changed, 33 insertions(+), 3 deletions(-) diff --git a/src/stub_launcher.cc b/src/stub_launcher.cc index 91b2bbc0..7c2dc013 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" @@ -802,9 +805,18 @@ void StubLauncher::TerminateStub() { if (is_initialized_) { + // Enforce stub_timeout_seconds_ as a single total budget across the + // finalize wait and the subsequent process-exit wait; without this the + // healthy teardown path could take up to 2 * stub_timeout_seconds_. bool force_kill = false; + int64_t remaining_seconds = stub_timeout_seconds_; if (is_healthy_) { - const int64_t timeout_ms = stub_timeout_seconds_ * 1000; + const int64_t total_timeout_ms = stub_timeout_seconds_ * 1000; + // MessageQueue::Pop takes `int const&`; clamp to INT_MAX so a large + // stub_timeout_seconds_ cannot silently narrow to a negative value. + 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 */); @@ -812,7 +824,14 @@ StubLauncher::TerminateStub() ipc_message->Command() = PYTHONSTUB_FinalizeRequest; stub_message_queue_->Push(ipc_message->ShmHandle()); bool success = false; - parent_message_queue_->Pop(timeout_ms, success); + 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; } @@ -826,7 +845,7 @@ StubLauncher::TerminateStub() if (force_kill) { KillStubProcess(); - } else if (!WaitForStubProcessWithTimeout(stub_timeout_seconds_)) { + } else if (!WaitForStubProcessWithTimeout(remaining_seconds)) { KillStubProcess(); } } @@ -964,6 +983,17 @@ StubLauncher::WaitForStubProcessWithTimeout(int64_t timeout_seconds) sleep(1); } + // The process may have exited during the final sleep window; re-check once + // more so we don't force-kill an already-dead process. + { + int status; + pid_t ret = waitpid(stub_pid_, &status, WNOHANG); + if (ret == stub_pid_ || ret == -1) { + stub_pid_ = 0; + return true; + } + } + return false; #endif } From 39096f6369c3a141e1f5312829f4e40e7ad9a13b Mon Sep 17 00:00:00 2001 From: Vinya Kestur Date: Thu, 23 Jul 2026 19:15:51 +0000 Subject: [PATCH 4/4] polish: log atexit registration failure; trim comments - Log a warning to stderr if std::atexit registration fails so the loss of the last-resort shm cleanup is visible; primary cleanup path is unaffected. - Tighten inline comments to explain only intent/constraints, not mechanics. --- src/shm_manager.cc | 6 +++++- src/shm_manager.h | 2 +- src/stub_launcher.cc | 10 +++------- 3 files changed, 9 insertions(+), 9 deletions(-) diff --git a/src/shm_manager.cc b/src/shm_manager.cc index ca548af2..b59a2ba4 100644 --- a/src/shm_manager.cc +++ b/src/shm_manager.cc @@ -61,7 +61,11 @@ RegisterParentShmRegion(const std::string& shm_region_name) parent_shm_regions.insert(shm_region_name); } if (!parent_shm_atexit_registered.exchange(true)) { - std::atexit(CleanupParentShmRegions); + if (std::atexit(CleanupParentShmRegions) != 0) { + std::cerr << "python_backend: failed to register atexit shm cleanup " + "handler; relying on TerminateStub for cleanup" + << std::endl; + } } } diff --git a/src/shm_manager.h b/src/shm_manager.h index a8e1c7fc..dc25f781 100644 --- a/src/shm_manager.h +++ b/src/shm_manager.h @@ -194,7 +194,7 @@ class SharedMemoryManager { void SetDeleteRegion(bool delete_region); - // Remove the parent-owned shared memory region from the filesystem. + // Idempotent: removes the parent-owned shm region and clears delete_region_. void RemoveShmRegion(); std::unique_ptr& GetCUDAMemoryPoolManager() diff --git a/src/stub_launcher.cc b/src/stub_launcher.cc index 7c2dc013..c676c5b4 100644 --- a/src/stub_launcher.cc +++ b/src/stub_launcher.cc @@ -805,15 +805,12 @@ void StubLauncher::TerminateStub() { if (is_initialized_) { - // Enforce stub_timeout_seconds_ as a single total budget across the - // finalize wait and the subsequent process-exit wait; without this the - // healthy teardown path could take up to 2 * stub_timeout_seconds_. + // 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; - // MessageQueue::Pop takes `int const&`; clamp to INT_MAX so a large - // stub_timeout_seconds_ cannot silently narrow to a negative value. + // 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()))); @@ -983,8 +980,7 @@ StubLauncher::WaitForStubProcessWithTimeout(int64_t timeout_seconds) sleep(1); } - // The process may have exited during the final sleep window; re-check once - // more so we don't force-kill an already-dead process. + // Stub may have exited during the last sleep(1); recheck before killing. { int status; pid_t ret = waitpid(stub_pid_, &status, WNOHANG);