fix: ensure Python backend shm cleanup on forced shutdown

src/shm_manager.cc

@@ -26,13 +26,54 @@
 
 #include "shm_manager.h"
 
+#include <atomic>
 #include <boost/interprocess/managed_external_buffer.hpp>
 #include <boost/interprocess/mapped_region.hpp>
 #include <boost/interprocess/shared_memory_object.hpp>
+#include <cstdlib>
 #include <iostream>
+#include <mutex>
+#include <unordered_set>
 
 namespace triton { namespace backend { namespace python {
 
+namespace {
+
+std::mutex parent_shm_regions_mu;
+std::unordered_set<std::string> parent_shm_regions;
+std::atomic<bool> parent_shm_atexit_registered{false};
+
+void
+CleanupParentShmRegions()
+{
+  std::lock_guard<std::mutex> 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<std::mutex> 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<std::mutex> 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;
   }
 }
 

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<CUDAMemoryPoolManager>& GetCUDAMemoryPoolManager()
   {
     return cuda_memory_pool_manager_;

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<IPCMessage> 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(

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_;