Simplify handling for async execute future object

triton-inference-server · Apr 3, 2024 · 37c8524 · 37c8524
1 parent 42ed090
commit 37c8524
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 37 deletions.
diff --git a/src/pb_stub.cc b/src/pb_stub.cc
@@ -35,6 +35,7 @@
 #include <boost/interprocess/sync/scoped_lock.hpp>
 #include <boost/thread/thread_time.hpp>
 #include <cstdlib>
+#include <future>
 #include <iomanip>
 #include <iostream>
 #include <memory>
@@ -904,33 +905,22 @@ Stub::RunCoroutine(py::object coroutine)
   py::object py_future = py::module_::import("asyncio").attr(
       "run_coroutine_threadsafe")(coroutine, loop);
 
-  {
-    std::lock_guard<std::mutex> lock(async_event_futures_mu_);
-
-    std::shared_ptr<std::future<void>> shared_future(new std::future<void>());
-    std::future<void> c_future = std::async(
-        std::launch::async, [this, shared_future, py_future]() mutable {
-          {
-            py::gil_scoped_acquire gil_acquire;
-            try {
-              py_future.attr("result")();
-            }
-            catch (const PythonBackendException& pb_exception) {
-              LOG_ERROR << pb_exception.what();
-            }
-            catch (const py::error_already_set& error) {
-              LOG_ERROR << error.what();
-            }
-            py_future = py::none();
-          }
-          {
-            std::lock_guard<std::mutex> lock(async_event_futures_mu_);
-            async_event_futures_.erase(shared_future);
-          }
-        });
-    *shared_future = std::move(c_future);
-    async_event_futures_.emplace(std::move(shared_future));
-  }
+  std::shared_ptr<std::future<void>> shared_future(new std::future<void>());
+  std::future<void> c_future = std::async(
+      std::launch::async, [this, shared_future, py_future]() mutable {
+        py::gil_scoped_acquire gil_acquire;
+        try {
+          py_future.attr("result")();
+        }
+        catch (const PythonBackendException& pb_exception) {
+          LOG_ERROR << pb_exception.what();
+        }
+        catch (const py::error_already_set& error) {
+          LOG_ERROR << error.what();
+        }
+        py_future = py::none();
+      });
+  *shared_future = std::move(c_future);
 
   return py::none();
 }
@@ -948,10 +938,6 @@ Stub::Finalize()
   finalizing_ = true;
   // Stop async event loop if created.
   if (!py::isinstance<py::none>(async_event_loop_)) {
-    if (!async_event_futures_.empty()) {
-      LOG_ERROR << "Finalizing stub with " << async_event_futures_.size()
-                << " ongoing coroutines";
-    }
     async_event_loop_.attr("stop")();
   }
   // Call finalize if exists.
@@ -1016,7 +1002,6 @@ Stub::~Stub()
 
   {
     py::gil_scoped_acquire acquire;
-    async_event_futures_.clear();
     async_event_loop_ = py::none();
     model_instance_ = py::none();
   }

diff --git a/src/pb_stub.h b/src/pb_stub.h
@@ -31,9 +31,6 @@
 #include <pybind11/stl.h>
 
 #include <filesystem>
-#include <future>
-#include <memory>
-#include <unordered_set>
 
 #include "infer_request.h"
 #include "infer_response.h"
@@ -371,8 +368,6 @@ class Stub {
   py::object deserialize_bytes_;
   py::object serialize_bytes_;
   py::object async_event_loop_;
-  std::unordered_set<std::shared_ptr<std::future<void>>> async_event_futures_;
-  std::mutex async_event_futures_mu_;
   std::unique_ptr<MessageQueue<bi::managed_external_buffer::handle_t>>
       stub_message_queue_;
   std::unique_ptr<MessageQueue<bi::managed_external_buffer::handle_t>>