Prune non requested outputs from decoupled models

triton-inference-server · Jun 11, 2024 · 37c199e · 37c199e
1 parent e63594f
commit 37c199e
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 24 deletions.
diff --git a/src/response_sender.cc b/src/response_sender.cc
@@ -74,23 +74,17 @@ ResponseSender::~ResponseSender()
       PYTHONSTUB_DecoupledResponseFactoryCleanup);
 }
 
-bool
-ResponseSender::IsDecoupled() const
+void
+ResponseSender::UpdateStateAndCounters(
+    const std::shared_ptr<InferResponse>& response, const uint32_t flags)
 {
   if (is_decoupled_ == nullptr) {
     // TODO: Can a model access the response sender on a BLS infer request?
     throw PythonBackendException(
         "Unable to send response. Response sender has no reference to the "
         "decoupled state of the model.");
   }
-  return *is_decoupled_;
-}
-
-void
-ResponseSender::UpdateStateAndCounters(
-    const std::shared_ptr<InferResponse>& response, const uint32_t flags)
-{
-  bool is_decoupled = IsDecoupled();
+  bool is_decoupled = *is_decoupled_;
 
   std::lock_guard<std::mutex> lk(mu_);
 
@@ -119,16 +113,6 @@ ResponseSender::UpdateStateAndCounters(
   number_of_response_sent_++;
 }
 
-void
-ResponseSender::PruneNonRequestedOutputs(
-    const std::shared_ptr<InferResponse>& infer_response) const
-{
-  // TODO: should this be limited to non decoupled only?
-  if (!IsDecoupled() && infer_response) {
-    infer_response->PruneOutputTensors(requested_output_names_);
-  }
-}
-
 void
 ResponseSender::Send(
     std::shared_ptr<InferResponse> infer_response, const uint32_t flags)
@@ -142,7 +126,9 @@ ResponseSender::Send(
 
   CheckResponseSenderArguments(infer_response, flags);
   UpdateStateAndCounters(infer_response, flags);
-  PruneNonRequestedOutputs(infer_response);
+  if (infer_response) {
+    infer_response->PruneOutputTensors(requested_output_names_);
+  }
 
   std::unique_ptr<Stub>& stub = Stub::GetOrCreateInstance();
 

diff --git a/src/response_sender.h b/src/response_sender.h
@@ -50,11 +50,8 @@ class ResponseSender {
   void Close();
 
  private:
-  bool IsDecoupled() const;
   void UpdateStateAndCounters(
       const std::shared_ptr<InferResponse>& response, const uint32_t flags);
-  void PruneNonRequestedOutputs(
-      const std::shared_ptr<InferResponse>& infer_response) const;
 
   intptr_t request_address_;
   intptr_t response_factory_address_;