Fix response iterator memory leak (#302)

triton-inference-server · Sep 26, 2023 · b136bf3 · b136bf3
1 parent 238e0d0
commit b136bf3
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 5 deletions.
diff --git a/src/pb_response_iterator.cc b/src/pb_response_iterator.cc
@@ -100,7 +100,7 @@ ResponseIterator::Next()
   }
 }
 
-py::iterator
+void
 ResponseIterator::Iter()
 {
   if (is_finished_) {
@@ -111,8 +111,6 @@ ResponseIterator::Iter()
       idx_ = 0;
     }
   }
-
-  return py::cast(*this);
 }
 
 void

diff --git a/src/pb_response_iterator.h b/src/pb_response_iterator.h
@@ -38,7 +38,7 @@ class ResponseIterator {
   ~ResponseIterator();
 
   std::shared_ptr<InferResponse> Next();
-  py::iterator Iter();
+  void Iter();
   void EnqueueResponse(std::shared_ptr<InferResponse> infer_response);
   void* Id();
   void Clear();

diff --git a/src/pb_stub.cc b/src/pb_stub.cc
@@ -1544,7 +1544,12 @@ PYBIND11_EMBEDDED_MODULE(c_python_backend_utils, module)
   py::class_<ResponseIterator, std::shared_ptr<ResponseIterator>>(
       module, "ResponseIterator")
       .def(py::init<const std::shared_ptr<InferResponse>&>())
-      .def("__iter__", &ResponseIterator::Iter, py::keep_alive<0, 1>())
+      .def(
+          "__iter__",
+          [](ResponseIterator& it) -> ResponseIterator& {
+            it.Iter();
+            return it;
+          })
       .def("__next__", &ResponseIterator::Next);
 
   py::class_<Logger> logger(module, "Logger");