Fix segfault for decoupled models (#327) (#328)

* Set release flags and clean up response factory map before returning error * Address comments * Move the cleanup function to the outside scope * Delete response factory when response sender goes out of scope
triton-inference-server · Dec 20, 2023 · 599a75e · 599a75e
1 parent f32a352
commit 599a75e
Show file tree

Hide file tree

Showing 5 changed files with 19 additions and 10 deletions.
diff --git a/src/ipc_message.h b/src/ipc_message.h
@@ -55,7 +55,7 @@ typedef enum PYTHONSTUB_commandtype_enum {
   PYTHONSTUB_AutoCompleteResponse,
   PYTHONSTUB_LogRequest,
   PYTHONSTUB_BLSDecoupledInferPayloadCleanup,
-  PYTHONSTUB_DecoupledResponseFactoryCleanup,
+  PYTHONSTUB_BLSDecoupledResponseFactoryCleanup,
   PYTHONSTUB_MetricFamilyRequestNew,
   PYTHONSTUB_MetricFamilyRequestDelete,
   PYTHONSTUB_MetricRequestNew,

diff --git a/src/pb_stub.cc b/src/pb_stub.cc
@@ -997,7 +997,7 @@ Stub::ServiceStubToParentRequests()
           (utils_msg_payload->command_type ==
            PYTHONSTUB_BLSDecoupledInferPayloadCleanup) ||
           (utils_msg_payload->command_type ==
-           PYTHONSTUB_DecoupledResponseFactoryCleanup)) {
+           PYTHONSTUB_BLSDecoupledResponseFactoryCleanup)) {
         SendCleanupId(utils_msg_payload, utils_msg_payload->command_type);
       } else if (
           utils_msg_payload->command_type == PYTHONSTUB_IsRequestCancelled) {

diff --git a/src/python_be.cc b/src/python_be.cc
@@ -830,8 +830,8 @@ ModelInstanceState::StubToParentMQMonitor()
         break;
       }
       case PYTHONSTUB_BLSDecoupledInferPayloadCleanup:
-      case PYTHONSTUB_DecoupledResponseFactoryCleanup: {
-        ProcessCleanupRequest(message);
+      case PYTHONSTUB_BLSDecoupledResponseFactoryCleanup: {
+        ProcessBLSCleanupRequest(message);
         break;
       }
       case PYTHONSTUB_IsRequestCancelled: {
@@ -928,9 +928,17 @@ ModelInstanceState::ProcessCleanupRequest(
       Stub()->ShmPool()->Load<char>(message->Args());
   CleanupMessage* cleanup_message_ptr =
       reinterpret_cast<CleanupMessage*>(cleanup_request_message.data_.get());
-
-  void* id = cleanup_message_ptr->id;
-  infer_payload_.erase(reinterpret_cast<intptr_t>(id));
+  intptr_t id = reinterpret_cast<intptr_t>(cleanup_message_ptr->id);
+  if (message->Command() == PYTHONSTUB_BLSDecoupledInferPayloadCleanup) {
+    // Remove the InferPayload object from the map.
+    infer_payload_.erase(id);
+  } else if (
+      message->Command() == PYTHONSTUB_BLSDecoupledResponseFactoryCleanup) {
+    // Delete response factory
+    std::unique_ptr<
+        TRITONBACKEND_ResponseFactory, backend::ResponseFactoryDeleter>
+        response_factory(reinterpret_cast<TRITONBACKEND_ResponseFactory*>(id));
+  }
 
   {
     bi::scoped_lock<bi::interprocess_mutex> lock{*(message->ResponseMutex())};

diff --git a/src/python_be.h b/src/python_be.h
@@ -400,8 +400,9 @@ class ModelInstanceState : public BackendModelInstance {
       std::unique_ptr<InferResponse>* infer_response,
       bi::managed_external_buffer::handle_t* response_handle);
 
-  // Process the decoupled cleanup request for InferPayload and ResponseFactory
-  void ProcessCleanupRequest(const std::unique_ptr<IPCMessage>& message);
+  // Process the bls decoupled cleanup request for InferPayload and
+  // ResponseFactory
+  void ProcessBLSCleanupRequest(const std::unique_ptr<IPCMessage>& message);
 
   // Process request cancellation query
   void ProcessIsRequestCancelled(const std::unique_ptr<IPCMessage>& message);

diff --git a/src/response_sender.cc b/src/response_sender.cc
@@ -50,7 +50,7 @@ ResponseSender::~ResponseSender()
   std::unique_ptr<Stub>& stub = Stub::GetOrCreateInstance();
   stub->EnqueueCleanupId(
       reinterpret_cast<void*>(response_factory_address_),
-      PYTHONSTUB_DecoupledResponseFactoryCleanup);
+      PYTHONSTUB_BLSDecoupledResponseFactoryCleanup);
 }
 
 void