Fix response factory cleanup

triton-inference-server · Sep 24, 2024 · 697dfc0 · 697dfc0
1 parent c42afe1
commit 697dfc0
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 7 deletions.
diff --git a/src/infer_request.h b/src/infer_request.h
@@ -96,6 +96,7 @@ class InferRequest {
   InferenceTrace& GetTrace();
   uint32_t ReleaseFlags();
   void SetReleaseFlags(const uint32_t& flags);
+  intptr_t GetResponseFactoryAddress() { return response_factory_address_; }
 
 #ifdef TRITON_PB_STUB
   std::shared_ptr<InferResponse> Exec(const bool is_decoupled);

diff --git a/src/python_be.cc b/src/python_be.cc
@@ -1089,6 +1089,17 @@ ModelInstanceState::ResponseSendDecoupled(
   ResponseSendMessage* send_message_payload =
       reinterpret_cast<ResponseSendMessage*>(send_message.data_.get());
   std::unique_ptr<PbString> error_message;
+  ScopedDefer response_factory_deleter([send_message_payload] {
+    if (send_message_payload->flags == TRITONSERVER_RESPONSE_COMPLETE_FINAL) {
+      TRITONBACKEND_ResponseFactory* response_factory =
+          reinterpret_cast<TRITONBACKEND_ResponseFactory*>(
+              send_message_payload->response_factory_address);
+      std::unique_ptr<
+          TRITONBACKEND_ResponseFactory, backend::ResponseFactoryDeleter>
+          lresponse_factory(reinterpret_cast<TRITONBACKEND_ResponseFactory*>(
+              response_factory));
+    }
+  });
   ScopedDefer _([send_message_payload] {
     {
       bi::scoped_lock<bi::interprocess_mutex> guard{send_message_payload->mu};
@@ -1214,13 +1225,6 @@ ModelInstanceState::ResponseSendDecoupled(
     SetErrorForResponseSendMessage(
         send_message_payload, WrapTritonErrorInSharedPtr(error), error_message);
   }
-
-  if (send_message_payload->flags == TRITONSERVER_RESPONSE_COMPLETE_FINAL) {
-    std::unique_ptr<
-        TRITONBACKEND_ResponseFactory, backend::ResponseFactoryDeleter>
-        lresponse_factory(
-            reinterpret_cast<TRITONBACKEND_ResponseFactory*>(response_factory));
-  }
 }
 
 TRITONSERVER_Error*
@@ -1357,6 +1361,16 @@ ModelInstanceState::ProcessRequests(
           (*responses)[r] = nullptr;
           continue;
         }
+        {
+          TRITONBACKEND_ResponseFactory* response_factory =
+              reinterpret_cast<TRITONBACKEND_ResponseFactory*>(
+                  pb_infer_requests[r]->GetResponseFactoryAddress());
+          std::unique_ptr<
+              TRITONBACKEND_ResponseFactory, backend::ResponseFactoryDeleter>
+              lresponse_factory(
+                  reinterpret_cast<TRITONBACKEND_ResponseFactory*>(
+                      response_factory));
+        }
         infer_response = InferResponse::LoadFromSharedMemory(
             Stub()->ShmPool(), response_shm_handle[r],
             false /* open_cuda_handle */);

diff --git a/src/response_sender.h b/src/response_sender.h
@@ -43,6 +43,7 @@ class ResponseSender {
       const std::set<std::string>& requested_output_names,
       std::unique_ptr<SharedMemoryManager>& shm_pool,
       const std::shared_ptr<PbCancel>& pb_cancel);
+  intptr_t ResponseFactory() { return response_factory_address_; }
   ~ResponseSender();
   void Send(std::shared_ptr<InferResponse> response, const uint32_t flags);
   bool IsCancelled();