fix: ensure Python backend shm cleanup on forced shutdown (#450)

* 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.

* 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.

* fix: address Greptile review — bounded total shutdown budget

- 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.

* 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.

* refactor: inline shm cleanup callback into atexit registration

CleanupParentShmRegions was only referenced as the std::atexit callback
inside RegisterParentShmRegion. Inline it as a capture-less lambda so the
teardown logic lives next to the registration and there's one fewer
top-level helper in the anonymous namespace.

No behavior change.

Author: Vinya567

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,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
+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)) {
+    if (std::atexit([]() {
+          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();
+        }) != 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<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

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

src/stub_launcher.cc

@@ -26,7 +26,10 @@
 
 #include "stub_launcher.h"
 
+#include <algorithm>
+#include <chrono>
 #include <filesystem>
+#include <limits>
 
 #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<int>(std::min<int64_t>(
+          total_timeout_ms,
+          static_cast<int64_t>(std::numeric_limits<int>::max())));
       // 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;
+      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::seconds>(
+              std::chrono::steady_clock::now() - pop_start)
+              .count();
+      remaining_seconds =
+          std::max<int64_t>(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(

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