Fix L0_backend_python timeout issue (triton-inference-server#218)

* Fix up missing mutex

* Rename variable

* Fix for the case where returning the first decoupled response is slower than the following responses

* Address comment

Author: krishung5

src/infer_request.cc

@@ -553,11 +553,11 @@ InferRequest::Exec(const bool is_decoupled)
 
   if (responses_is_set) {
     auto& memory_manager_message_queue = stub->MemoryManagerQueue();
-    std::unique_ptr<InferResponse> error_response =
+    std::unique_ptr<InferResponse> return_response =
         InferResponse::LoadFromSharedMemory(
             shm_pool, *response_handle, true /* open cuda handle */);
 
-    for (auto& output_tensor : error_response->OutputTensors()) {
+    for (auto& output_tensor : return_response->OutputTensors()) {
       if (!output_tensor->IsCPU()) {
         uint64_t memory_release_id = output_tensor->Memory()->MemoryReleaseId();
         output_tensor->Memory()->SetMemoryReleaseCallback(
@@ -567,7 +567,7 @@ InferRequest::Exec(const bool is_decoupled)
       }
     }
 
-    return error_response;
+    return return_response;
   } else {
     auto error_response = std::make_unique<InferResponse>(
         std::vector<std::shared_ptr<PbTensor>>{},

src/pb_response_iterator.cc

@@ -114,11 +114,11 @@ ResponseIterator::Iter()
 }
 
 void
-ResponseIterator::EnqueueResponse(std::unique_ptr<InferResponse> infer_response)
+ResponseIterator::EnqueueResponse(std::shared_ptr<InferResponse> infer_response)
 {
   {
     std::lock_guard<std::mutex> lock{mu_};
-    response_buffer_.push(std::move(infer_response));
+    response_buffer_.push(infer_response);
   }
   cv_.notify_one();
 }
@@ -144,4 +144,19 @@ ResponseIterator::Clear()
   is_cleared_ = true;
 }
 
+std::vector<std::shared_ptr<InferResponse>>
+ResponseIterator::GetExistingResponses()
+{
+  std::vector<std::shared_ptr<InferResponse>> responses;
+  std::unique_lock<std::mutex> lock{mu_};
+  while (!response_buffer_.empty()) {
+    responses.push_back(response_buffer_.front());
+    response_buffer_.pop();
+  }
+  is_finished_ = true;
+  is_cleared_ = true;
+
+  return responses;
+}
+
 }}}  // namespace triton::backend::python

src/pb_response_iterator.h

@@ -38,9 +38,10 @@ class ResponseIterator {
 
   std::shared_ptr<InferResponse> Next();
   py::iterator Iter();
-  void EnqueueResponse(std::unique_ptr<InferResponse> infer_response);
+  void EnqueueResponse(std::shared_ptr<InferResponse> infer_response);
   void* Id();
   void Clear();
+  std::vector<std::shared_ptr<InferResponse>> GetExistingResponses();
 
  private:
   std::vector<std::shared_ptr<InferResponse>> responses_;

src/pb_stub.cc

@@ -918,8 +918,11 @@ Stub::ServiceStubToParentRequests()
         break;
       } else {
         bls_response_cleanup_buffer_.pop();
+        {
+          std::lock_guard<std::mutex> lock(response_iterator_map_mu_);
+          response_iterator_map_.erase(id);
+        }
         SendCleanupId(id);
-        response_iterator_map_.erase(id);
       }
     }
   }
@@ -1093,7 +1096,11 @@ Stub::ParentToStubMQMonitor()
         response_iterator_map_[infer_response->Id()]->EnqueueResponse(
             std::move(infer_response));
       } else {
-        LOG_INFO << "Failed to enqueue the response to its response iterator.";
+        auto response_iterator =
+            std::make_shared<ResponseIterator>(std::move(infer_response));
+        response_iterator_map_.insert(
+            std::pair<void*, std::shared_ptr<ResponseIterator>>(
+                response_iterator->Id(), response_iterator));
       }
     }
 
@@ -1115,13 +1122,31 @@ Stub::ParentToStubServiceActive()
   return parent_to_stub_thread_;
 }
 
-void
-Stub::SaveResponseIterator(std::shared_ptr<ResponseIterator> response_iterator)
+std::shared_ptr<ResponseIterator>
+Stub::GetResponseIterator(std::shared_ptr<InferResponse> infer_response)
 {
   std::lock_guard<std::mutex> lock(response_iterator_map_mu_);
-  response_iterator_map_.insert(
-      std::pair<void*, std::shared_ptr<ResponseIterator>>(
-          response_iterator->Id(), response_iterator));
+  if (response_iterator_map_.find(infer_response->Id()) !=
+      response_iterator_map_.end()) {
+    // Need to re-construct the 'ResponseIterator' and update the
+    // 'response_iterator_map_' to make sure the 'ResponseIterator' object has
+    // the correct first response.
+    auto response_iterator = std::make_shared<ResponseIterator>(infer_response);
+    std::vector<std::shared_ptr<InferResponse>> existing_responses =
+        response_iterator_map_[infer_response->Id()]->GetExistingResponses();
+    for (auto& response : existing_responses) {
+      response_iterator->EnqueueResponse(response);
+    }
+
+    response_iterator_map_[infer_response->Id()] = response_iterator;
+  } else {
+    auto response_iterator = std::make_shared<ResponseIterator>(infer_response);
+    response_iterator_map_.insert(
+        std::pair<void*, std::shared_ptr<ResponseIterator>>(
+            response_iterator->Id(), response_iterator));
+  }
+
+  return response_iterator_map_[infer_response->Id()];
 }
 
 bool
@@ -1304,12 +1329,8 @@ PYBIND11_EMBEDDED_MODULE(c_python_backend_utils, module)
                 infer_request->Exec(decoupled);
             py::object response_object;
             if (decoupled) {
-              auto response_iterator =
-                  std::make_shared<ResponseIterator>(response);
+              auto response_iterator = stub->GetResponseIterator(response);
               response_object = py::cast(response_iterator);
-              if (response_iterator->Id() != nullptr) {
-                stub->SaveResponseIterator(response_iterator);
-              }
             } else {
               response_object = py::cast(response);
             }
@@ -1334,12 +1355,8 @@ PYBIND11_EMBEDDED_MODULE(c_python_backend_utils, module)
                   infer_request->Exec(decoupled);
               py::object response_object;
               if (decoupled) {
-                auto response_iterator =
-                    std::make_shared<ResponseIterator>(response);
+                auto response_iterator = stub->GetResponseIterator(response);
                 response_object = py::cast(response_iterator);
-                if (response_iterator->Id() != nullptr) {
-                  stub->SaveResponseIterator(response_iterator);
-                }
               } else {
                 response_object = py::cast(response);
               }

src/pb_stub.h

@@ -246,9 +246,9 @@ class Stub {
   /// Thread process
   void ParentToStubMQMonitor();
 
-  /// Keep track of the ResponseIterator object
-  void SaveResponseIterator(
-      std::shared_ptr<ResponseIterator> response_iterator);
+  /// Get the ResponseIterator object associated with the infer response
+  std::shared_ptr<ResponseIterator> GetResponseIterator(
+      std::shared_ptr<InferResponse> infer_response);
 
   /// Send the id to the python backend for object cleanup
   void SendCleanupId(void* id);