Merge branch 'triton-inference-server:main' into main

Author: aleksn7

src/infer_request.cc

@@ -1,4 +1,4 @@
-// Copyright 2021-2024, 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
@@ -418,7 +418,7 @@ InferRequest::Exec(const bool is_decoupled)
   py::gil_scoped_release release;
 
   // BLS should not be used in "initialize" or "finalize" function.
-  std::unique_ptr<Stub>& stub = Stub::GetOrCreateInstance();
+  auto stub = Stub::GetOrCreateInstance();
   if (!stub->IsInitialized() || stub->IsFinalizing()) {
     throw PythonBackendException(
         "BLS is only supported during the 'execute' function.");

src/metric.cc

@@ -1,4 +1,4 @@
-// Copyright 2023-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright 2023-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
@@ -147,7 +147,7 @@ void
 Metric::SendCreateMetricRequest()
 {
   // Send the request to create the Metric to the parent process
-  std::unique_ptr<Stub>& stub = Stub::GetOrCreateInstance();
+  auto stub = Stub::GetOrCreateInstance();
   SaveToSharedMemory(stub->ShmPool());
   CustomMetricsMessage* custom_metrics_msg = nullptr;
   AllocatedSharedMemory<CustomMetricsMessage> custom_metrics_shm;
@@ -170,7 +170,7 @@ Metric::SendIncrementRequest(const double& value)
   py::gil_scoped_release release;
   try {
     CheckIfCleared();
-    std::unique_ptr<Stub>& stub = Stub::GetOrCreateInstance();
+    auto stub = Stub::GetOrCreateInstance();
     operation_value_ = value;
     SaveToSharedMemory(stub->ShmPool());
     AllocatedSharedMemory<CustomMetricsMessage> custom_metrics_shm;
@@ -189,7 +189,7 @@ Metric::SendSetValueRequest(const double& value)
 {
   try {
     CheckIfCleared();
-    std::unique_ptr<Stub>& stub = Stub::GetOrCreateInstance();
+    auto stub = Stub::GetOrCreateInstance();
     operation_value_ = value;
     SaveToSharedMemory(stub->ShmPool());
     AllocatedSharedMemory<CustomMetricsMessage> custom_metrics_shm;
@@ -208,7 +208,7 @@ Metric::SendObserveRequest(const double& value)
   py::gil_scoped_release release;
   try {
     CheckIfCleared();
-    std::unique_ptr<Stub>& stub = Stub::GetOrCreateInstance();
+    auto stub = Stub::GetOrCreateInstance();
     operation_value_ = value;
     SaveToSharedMemory(stub->ShmPool());
     AllocatedSharedMemory<CustomMetricsMessage> custom_metrics_shm;
@@ -228,7 +228,7 @@ Metric::SendGetValueRequest()
   AllocatedSharedMemory<CustomMetricsMessage> custom_metrics_shm;
   try {
     CheckIfCleared();
-    std::unique_ptr<Stub>& stub = Stub::GetOrCreateInstance();
+    auto stub = Stub::GetOrCreateInstance();
     SaveToSharedMemory(stub->ShmPool());
     stub->SendMessage<CustomMetricsMessage>(
         custom_metrics_shm, PYTHONSTUB_MetricRequestValue, shm_handle_);
@@ -251,7 +251,7 @@ Metric::Clear()
   // scope/being deleted.
   if (!is_cleared_) {
     is_cleared_ = true;
-    std::unique_ptr<Stub>& stub = Stub::GetOrCreateInstance();
+    auto stub = Stub::GetOrCreateInstance();
     SaveToSharedMemory(stub->ShmPool());
     AllocatedSharedMemory<CustomMetricsMessage> custom_metrics_shm;
     try {

src/metric_family.cc

@@ -1,4 +1,4 @@
-// Copyright 2023-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright 2023-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
@@ -55,7 +55,7 @@ MetricFamily::~MetricFamily()
   }
 
   // Send the request to delete the MetricFamily to the parent process
-  std::unique_ptr<Stub>& stub = Stub::GetOrCreateInstance();
+  auto stub = Stub::GetOrCreateInstance();
   SaveToSharedMemory(stub->ShmPool());
   AllocatedSharedMemory<CustomMetricsMessage> custom_metrics_shm;
   try {
@@ -147,7 +147,7 @@ MetricFamily::CreateMetricFamily(
 void
 MetricFamily::SendCreateMetricFamilyRequest()
 {
-  std::unique_ptr<Stub>& stub = Stub::GetOrCreateInstance();
+  auto stub = Stub::GetOrCreateInstance();
   SaveToSharedMemory(stub->ShmPool());
   CustomMetricsMessage* custom_metrics_msg = nullptr;
   AllocatedSharedMemory<CustomMetricsMessage> custom_metrics_shm;

src/model_loader.cc

@@ -1,4 +1,4 @@
-// Copyright 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright 2023-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
@@ -101,7 +101,7 @@ ModelLoader::ModelLoader(
 void
 ModelLoader::SendLoadModelRequest()
 {
-  std::unique_ptr<Stub>& stub = Stub::GetOrCreateInstance();
+  auto stub = Stub::GetOrCreateInstance();
   SaveToSharedMemory(stub->ShmPool());
   AllocatedSharedMemory<ModelLoaderMessage> model_loader_msg_shm;
 
@@ -118,7 +118,7 @@ ModelLoader::SendLoadModelRequest()
 void
 ModelLoader::SendUnloadModelRequest()
 {
-  std::unique_ptr<Stub>& stub = Stub::GetOrCreateInstance();
+  auto stub = Stub::GetOrCreateInstance();
   SaveToSharedMemory(stub->ShmPool());
   AllocatedSharedMemory<ModelLoaderMessage> model_loader_msg_shm;
   try {
@@ -134,7 +134,7 @@ ModelLoader::SendUnloadModelRequest()
 bool
 ModelLoader::SendModelReadinessRequest()
 {
-  std::unique_ptr<Stub>& stub = Stub::GetOrCreateInstance();
+  auto stub = Stub::GetOrCreateInstance();
   SaveToSharedMemory(stub->ShmPool());
   ModelLoaderMessage* model_loader_msg = nullptr;
   AllocatedSharedMemory<ModelLoaderMessage> model_loader_msg_shm;

src/pb_bls_cancel.cc

@@ -1,4 +1,4 @@
-// Copyright 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright 2025-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
@@ -67,7 +67,7 @@ PbBLSCancel::Cancel()
     return;
   }
   if (!updating_) {
-    std::unique_ptr<Stub>& stub = Stub::GetOrCreateInstance();
+    auto stub = Stub::GetOrCreateInstance();
     if (!stub->StubToParentServiceActive()) {
       LOG_ERROR << "Cannot communicate with parent service";
       return;

src/pb_cancel.cc

@@ -1,4 +1,4 @@
-// Copyright 2023-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright 2023-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
@@ -68,7 +68,7 @@ PbCancel::IsCancelled()
     return is_cancelled_;
   }
   if (!updating_) {
-    std::unique_ptr<Stub>& stub = Stub::GetOrCreateInstance();
+    auto stub = Stub::GetOrCreateInstance();
     if (!stub->StubToParentServiceActive()) {
       LOG_ERROR << "Cannot communicate with parent service";
       return false;

src/pb_response_iterator.cc

@@ -1,4 +1,4 @@
-// Copyright 2023-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright 2023-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
@@ -133,7 +133,7 @@ ResponseIterator::Id()
 void
 ResponseIterator::Clear()
 {
-  std::unique_ptr<Stub>& stub = Stub::GetOrCreateInstance();
+  auto stub = Stub::GetOrCreateInstance();
   stub->EnqueueCleanupId(id_, PYTHONSTUB_BLSDecoupledInferPayloadCleanup);
   {
     std::lock_guard<std::mutex> lock{mu_};

src/pb_stub.cc

@@ -130,7 +130,7 @@ PyParametersToJSON(const py::dict& parameters)
 void
 AsyncEventFutureDoneCallback(const py::object& py_future)
 {
-  std::unique_ptr<Stub>& stub = Stub::GetOrCreateInstance();
+  auto stub = Stub::GetOrCreateInstance();
   stub->BackgroundFutureDone(py_future);
 }
 
@@ -514,7 +514,7 @@ Stub::AutoCompleteModelConfig(
   python_backend_utils.def(
       "get_model_dir",
       []() {
-        std::unique_ptr<Stub>& stub = Stub::GetOrCreateInstance();
+        auto stub = Stub::GetOrCreateInstance();
         return stub->GetModelDir();
       },
       py::return_value_policy::reference);
@@ -568,7 +568,7 @@ Stub::Initialize(bi::managed_external_buffer::handle_t map_handle)
   python_backend_utils.def(
       "get_model_dir",
       []() {
-        std::unique_ptr<Stub>& stub = Stub::GetOrCreateInstance();
+        auto stub = Stub::GetOrCreateInstance();
         return stub->GetModelDir();
       },
       py::return_value_policy::reference);
@@ -1073,16 +1073,22 @@ Stub::~Stub()
   memory_manager_message_queue_.reset();
 }
 
-std::unique_ptr<Stub> Stub::stub_instance_;
+static std::shared_ptr<triton::backend::python::Stub> stub_instance{nullptr};
 
-std::unique_ptr<Stub>&
+std::shared_ptr<triton::backend::python::Stub>
 Stub::GetOrCreateInstance()
 {
-  if (Stub::stub_instance_.get() == nullptr) {
-    Stub::stub_instance_ = std::make_unique<Stub>();
+  if (!stub_instance) {
+    stub_instance.reset(new triton::backend::python::Stub());
   }
 
-  return Stub::stub_instance_;
+  return stub_instance;
+}
+
+void
+Stub::DestroyInstance()
+{
+  stub_instance.reset();
 }
 
 void
@@ -1834,7 +1840,7 @@ PYBIND11_EMBEDDED_MODULE(c_python_backend_utils, module)
           "exec",
           [](std::shared_ptr<InferRequest>& infer_request,
              const bool decoupled) {
-            std::unique_ptr<Stub>& stub = Stub::GetOrCreateInstance();
+            auto stub = Stub::GetOrCreateInstance();
             std::shared_ptr<InferResponse> response =
                 infer_request->Exec(decoupled);
             py::object response_object;
@@ -1852,7 +1858,7 @@ PYBIND11_EMBEDDED_MODULE(c_python_backend_utils, module)
           "async_exec",
           [](std::shared_ptr<InferRequest>& infer_request,
              const bool decoupled) {
-            std::unique_ptr<Stub>& stub = Stub::GetOrCreateInstance();
+            auto stub = Stub::GetOrCreateInstance();
             py::object loop =
                 py::module_::import("asyncio").attr("get_running_loop")();
             py::cpp_function callback = [&stub, infer_request, decoupled]() {
@@ -2137,7 +2143,7 @@ main(int argc, char** argv)
   std::string name = argv[8];
   std::string runtime_modeldir = argv[9];
 
-  std::unique_ptr<Stub>& stub = Stub::GetOrCreateInstance();
+  auto stub = Stub::GetOrCreateInstance();
   try {
     stub->Instantiate(
         shm_growth_size, shm_default_size, shm_region_name, model_path,
@@ -2147,7 +2153,7 @@ main(int argc, char** argv)
   catch (const PythonBackendException& pb_exception) {
     LOG_INFO << "Failed to preinitialize Python stub: " << pb_exception.what();
     logger.reset();
-    stub.reset();
+    Stub::DestroyInstance();
     exit(1);
   }
 
@@ -2160,7 +2166,7 @@ main(int argc, char** argv)
 #endif
   std::atomic<bool> background_thread_running = {true};
   std::thread background_thread =
-      std::thread([&parent_pid, &background_thread_running, &stub, &logger] {
+      std::thread([stub, &parent_pid, &background_thread_running, &logger] {
         // Send a dummy message after the stub process is launched to notify the
         // parent process that the health thread has started.
         std::unique_ptr<IPCMessage> ipc_message = IPCMessage::Create(
@@ -2193,7 +2199,7 @@ main(int argc, char** argv)
             // Destroy stub and exit.
             stub->DestroyPythonObjects();
             logger.reset();
-            stub.reset();
+            Stub::DestroyInstance();
             exit(1);
           }
         }
@@ -2227,7 +2233,7 @@ main(int argc, char** argv)
   // will result in segfault.
   stub->DestroyPythonObjects();
   logger.reset();
-  stub.reset();
+  Stub::DestroyInstance();
 
   return 0;
 }

src/pb_stub.h

@@ -96,7 +96,9 @@ struct UtilsMessagePayload {
 class Stub {
  public:
   Stub() : stub_to_parent_thread_(false), parent_to_stub_thread_(false){};
-  static std::unique_ptr<Stub>& GetOrCreateInstance();
+  static std::shared_ptr<Stub> GetOrCreateInstance();
+
+  static void DestroyInstance();
 
   /// Instantiate a new Python backend Stub.
   void Instantiate(
@@ -301,7 +303,6 @@ class Stub {
   std::unique_ptr<MessageQueue<uint64_t>> memory_manager_message_queue_;
   bool initialized_;
   bool finalizing_;
-  static std::unique_ptr<Stub> stub_instance_;
   std::vector<std::shared_ptr<PbTensor>> gpu_tensors_;
   std::queue<std::unique_ptr<UtilsMessagePayload>> stub_to_parent_buffer_;
   std::thread stub_to_parent_queue_monitor_;

src/pb_stub_log.cc

@@ -1,4 +1,4 @@
-// Copyright 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright 2025-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
@@ -51,7 +51,7 @@ Logger::GetOrCreateInstance()
 void
 Logger::Log(const std::string& message, LogLevel level)
 {
-  std::unique_ptr<Stub>& stub = Stub::GetOrCreateInstance();
+  auto stub = Stub::GetOrCreateInstance();
   py::object frame = py::module_::import("inspect").attr("currentframe");
   py::object caller_frame = frame();
   py::object info = py::module_::import("inspect").attr("getframeinfo");
@@ -110,7 +110,7 @@ Logger::Log(
 #endif
   } else {
     // Ensure we do not create a stub instance before it has initialized
-    std::unique_ptr<Stub>& stub = Stub::GetOrCreateInstance();
+    auto stub = Stub::GetOrCreateInstance();
     std::unique_ptr<PbLog> log_msg(new PbLog(filename, lineno, message, level));
     stub->EnqueueLogRequest(log_msg);
   }