Clean up response_factory when the final flag is set

Author: krishung5

src/pb_utils.h

@@ -225,6 +225,7 @@ struct ResponseSenderBase {
   bi::managed_external_buffer::handle_t error;
   intptr_t request_address;
   intptr_t response_factory_address;
+  bool is_response_factory_cleaned;
 };
 
 struct ResponseSendMessage : ResponseSenderBase {

src/python_be.cc

@@ -1310,6 +1310,7 @@ ModelInstanceState::ResponseSendDecoupled(
         TRITONBACKEND_ResponseFactory, backend::ResponseFactoryDeleter>
         response_factory(reinterpret_cast<TRITONBACKEND_ResponseFactory*>(
             send_message_payload->response_factory_address));
+    send_message_payload->is_response_factory_cleaned = true;
   }
 }
 

src/response_sender.cc

@@ -1,4 +1,4 @@
-// Copyright 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright 2022-2024, 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
@@ -41,17 +41,18 @@ ResponseSender::ResponseSender(
     const std::shared_ptr<PbCancel>& pb_cancel)
     : request_address_(request_address),
       response_factory_address_(response_factory_address), shm_pool_(shm_pool),
-      closed_(false), pb_cancel_(pb_cancel)
+      closed_(false), pb_cancel_(pb_cancel), is_response_factory_cleaned_(false)
 {
 }
 
 ResponseSender::~ResponseSender()
 {
-  // std::cerr << "===== ResponseSender::~ResponseSender() =====" << std::endl;
-  // std::unique_ptr<Stub>& stub = Stub::GetOrCreateInstance();
-  // stub->EnqueueCleanupId(
-  //     reinterpret_cast<void*>(response_factory_address_),
-  //     PYTHONSTUB_DecoupledResponseFactoryCleanup);
+  if (!is_response_factory_cleaned_) {
+    std::unique_ptr<Stub>& stub = Stub::GetOrCreateInstance();
+    stub->EnqueueCleanupId(
+        reinterpret_cast<void*>(response_factory_address_),
+        PYTHONSTUB_DecoupledResponseFactoryCleanup);
+  }
 }
 
 void
@@ -189,6 +190,9 @@ ResponseSender::Send(
     }
   }
 
+  is_response_factory_cleaned_ =
+      send_message_payload->is_response_factory_cleaned;
+
   if (send_message_payload->has_error) {
     if (send_message_payload->is_error_set) {
       std::unique_ptr<PbString> error = PbString::LoadFromSharedMemory(

src/response_sender.h

@@ -1,4 +1,4 @@
-// Copyright 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright 2022-2024, 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
@@ -46,7 +46,16 @@ class ResponseSender {
   intptr_t request_address_;
   intptr_t response_factory_address_;
   std::unique_ptr<SharedMemoryManager>& shm_pool_;
+  // The flag to indicate if the response sender is closed. It is set to true
+  // once the TRITONSERVER_RESPONSE_COMPLETE_FINAL flag is set, meaning that the
+  // response_sender should not be used anymore. This flag is separate from the
+  // `is_response_factory_cleaned_` flag because errors might occur after
+  // complete_final flag is set but before the response_factory gets cleaned up.
   bool closed_;
   std::shared_ptr<PbCancel> pb_cancel_;
+  // The flag to indicate if the response_factory is already cleaned up in the
+  // python backend side. If not, the response_factory will be cleaned up in the
+  // destructor.
+  bool is_response_factory_cleaned_;
 };
 }}}  // namespace triton::backend::python