Improve error message when BLS is used in 'initialize' or 'finalize' function (triton-inference-server#211)

krishung5 · web-flow · commit f007255d977d · 2023-03-10T13:10:55.000-08:00
* Improve error message when BLS is used in 'initialize' or 'finalize' function

* Address comment
diff --git a/src/infer_request.cc b/src/infer_request.cc
@@ -373,9 +373,15 @@ InferRequest::GetResponseSender()
 std::shared_ptr<InferResponse>
 InferRequest::Exec(const bool is_decoupled)
 {
+  // BLS should not be used in "initialize" or "finalize" function.
+  std::unique_ptr<Stub>& stub = Stub::GetOrCreateInstance();
+  if (!stub->IsInitialized() || stub->IsFinalizing()) {
+    throw PythonBackendException(
+        "BLS is only supported during the 'execute' function.");
+  }
+
   ResponseBatch* response_batch = nullptr;
   bool responses_is_set = false;
-  std::unique_ptr<Stub>& stub = Stub::GetOrCreateInstance();
   std::unique_ptr<SharedMemoryManager>& shm_pool = stub->SharedMemory();
   bi::managed_external_buffer::handle_t* response_handle = nullptr;
 
@@ -529,8 +535,7 @@ InferRequest::Exec(const bool is_decoupled)
 
     for (auto& output_tensor : error_response->OutputTensors()) {
       if (!output_tensor->IsCPU()) {
-        uint64_t memory_release_id =
-            output_tensor->Memory()->MemoryReleaseId();
+        uint64_t memory_release_id = output_tensor->Memory()->MemoryReleaseId();
         output_tensor->Memory()->SetMemoryReleaseCallback(
             [&memory_manager_message_queue, memory_release_id]() {
               memory_manager_message_queue->Push(memory_release_id);
diff --git a/src/pb_stub.cc b/src/pb_stub.cc
@@ -84,6 +84,7 @@ Stub::Instantiate(
   name_ = name;
   health_mutex_ = nullptr;
   initialized_ = false;
+  finalizing_ = false;
   stub_to_parent_thread_ = false;
   parent_to_stub_thread_ = false;
 
@@ -799,6 +800,7 @@ Stub::UpdateHealth()
 void
 Stub::Finalize()
 {
+  finalizing_ = true;
   // Call finalize if exists.
   if (initialized_ && py::hasattr(model_instance_, "finalize")) {
     try {
@@ -1120,6 +1122,18 @@ Stub::SaveResponseIterator(std::shared_ptr<ResponseIterator> response_iterator)
           response_iterator->Id(), response_iterator));
 }
 
+bool
+Stub::IsInitialized()
+{
+  return initialized_;
+}
+
+bool
+Stub::IsFinalizing()
+{
+  return finalizing_;
+}
+
 std::unique_ptr<Logger> Logger::log_instance_;
 
 std::unique_ptr<Logger>&
diff --git a/src/pb_stub.h b/src/pb_stub.h
@@ -256,6 +256,12 @@ class Stub {
   /// Add cleanup id to queue
   void EnqueueCleanupId(void* id);
 
+  /// Is the stub initialized
+  bool IsInitialized();
+
+  /// Is the stub in the finalize stage
+  bool IsFinalizing();
+
 
  private:
   bi::interprocess_mutex* stub_mutex_;
@@ -282,6 +288,7 @@ class Stub {
       parent_to_stub_mq_;
   std::unique_ptr<MessageQueue<uint64_t>> memory_manager_message_queue_;
   bool initialized_;
+  bool finalizing_;
   static std::unique_ptr<Stub> stub_instance_;
   std::vector<std::shared_ptr<PbTensor>> gpu_tensors_;
   std::queue<std::unique_ptr<PbLog>> log_request_buffer_;
