Fix issue where Python sources can deadlock when the subscriber is unsubscribed

python/mrc/_pymrc/include/pymrc/segment.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2021-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-FileCopyrightText: Copyright (c) 2021-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
@@ -252,5 +252,7 @@ class BuilderProxy
     static void init_module(mrc::segment::IBuilder& self, std::shared_ptr<mrc::modules::SegmentModule> module);
 };
 
+void iterator_thread(pybind11::iterator itr, pybind11::object queue, pybind11::object exception_queue);
+
 #pragma GCC visibility pop
 }  // namespace mrc::pymrc

python/mrc/_pymrc/src/segment.cpp

@@ -35,14 +35,18 @@
 #include <glog/logging.h>
 #include <pybind11/cast.h>
 #include <pybind11/gil.h>
+#include <pybind11/pybind11.h>
 #include <pybind11/pytypes.h>
+#include <pyerrors.h>  // for PyExc_SystemExit
 #include <rxcpp/rx.hpp>
 
+#include <chrono>
 #include <exception>
 #include <fstream>
 #include <functional>
 #include <iterator>
 #include <map>
+#include <memory>
 #include <stdexcept>
 #include <string>
 #include <type_traits>
@@ -52,6 +56,58 @@
 // IWYU thinks we need array for py::print
 // IWYU pragma: no_include <array>
 
+namespace {
+namespace py = pybind11;
+using namespace std::chrono_literals;
+constexpr std::chrono::seconds ThreadShutdownTimeoutSec = 10s;
+
+void stop_python_thread(py::object& thread)
+{
+    if (thread.attr("is_alive")().cast<bool>())
+    {
+        // The source generator thread needs to stop prior to the thread object going out of scope
+        auto thread_id = thread.attr("ident").cast<unsigned long>();
+        try
+        {
+            DVLOG(10) << "Attempting to kill python thread " << thread.attr("name").cast<std::string>()
+                      << " id: " << thread_id;
+
+            // Using PyExc_SystemExit since Python won't log it as an uncaught exception
+            auto num_threads_stopped = PyThreadState_SetAsyncExc(thread_id, PyExc_SystemExit);
+            if (num_threads_stopped != 1)
+            {
+                LOG(WARNING) << "Expected return value of 1 from PyThreadState_SetAsyncExc, received: "
+                             << num_threads_stopped;
+            }
+            else
+            {
+                // wait until the thread has actually stopped
+                bool is_alive       = thread.attr("is_alive")().cast<bool>();
+                const auto deadline = std::chrono::system_clock::now() + ThreadShutdownTimeoutSec;
+                while (is_alive && std::chrono::system_clock::now() < deadline)
+                {
+                    {
+                        pybind11::gil_scoped_release no_gil;
+                        boost::this_fiber::yield();
+                    }
+
+                    is_alive = thread.attr("is_alive")().cast<bool>();
+                }
+
+                if (is_alive)
+                {
+                    LOG(WARNING) << "Unable to stop thread: " << thread_id << " after "
+                                 << ThreadShutdownTimeoutSec.count() << " seconds";
+                }
+            }
+        } catch (const std::exception& e)
+        {
+            LOG(ERROR) << "Encountered error while attempting to stop thread: " << thread_id << ": " << e.what();
+        }
+    }
+}
+}  // namespace
+
 namespace mrc::pymrc {
 
 namespace py = pybind11;
@@ -152,6 +208,31 @@
     PyHolder m_value{};
 };
 
+void iterator_thread(py::iterator itr, py::object queue, py::object exception_queue)
+{
+    try
+    {
+        py::gil_scoped_acquire gil;
+        PyIteratorIterator wrapped_iter(std::move(itr));
+        PyIteratorIterator sentinel;
+        while (wrapped_iter != sentinel)
+        {
+            {
+                // Copy the current value into the queue.
+                queue.attr("put")((*wrapped_iter).copy_obj());
+            }
+
+            ++wrapped_iter;
+        }
+    } catch (py::error_already_set py_except)
+    {
+        exception_queue.attr("put")(py::str(py_except.what()));
+    }
+}
+
+class EmptyQueue : public std::exception
+{};
+
 class PyIteratorWrapper
 {
   public:
@@ -160,7 +241,7 @@
     // NOLINTEND(readability-identifier-naming)
 
     // Create from an iterator
-    PyIteratorWrapper(py::iterator source_iterator) :
+    PyIteratorWrapper(py::iterator source_iterator, PyObjectSubscriber* subscriber = nullptr) :
       m_iter_factory([iterator = PyObjectHolder(std::move(source_iterator))]() mutable {
           // Check if the iterator has been started already
           if (!iterator)
@@ -195,58 +276,138 @@
     // Create directly
     PyIteratorWrapper(std::function<py::iterator()> iter_factory) : m_iter_factory(std::move(iter_factory)) {}
 
-    iterator begin()
+    py::object get_next_value()
     {
-        // Grab the GIL and create the iterator from the factory
+        using namespace py::literals;
+        if (!m_initialized)
+        {
+            init();
+        }
+
         AcquireGIL gil;
+        try
+        {
+            auto value = m_queue.attr("get_nowait")();
+            return value;
+        } catch (py::error_already_set py_except)
+        {
+            if (py_except.matches(m_empty_exception))
+            {
+                // Check to see if we got an exception
+                if (!m_exception_queue.attr("empty")().cast<bool>())
+                {
+                    auto py_err_str = m_exception_queue.attr("get")("block"_a = true, "timeout"_a = 0.5);
+                    throw std::runtime_error(py_err_str.cast<std::string>());
+                }
 
-        auto iter = m_iter_factory();
+                if (!m_thread.attr("is_alive")().cast<bool>())
+                {
+                    throw pybind11::stop_iteration();
+                }
+
+                // if the thread is alive and we don't have an exception, just signal that we don't have any data
+                throw EmptyQueue();
+            }
 
-        return iterator{std::move(iter)};
+            throw;
+        }
     }
 
-    iterator end()  // NOLINT(readability-convert-member-functions-to-static)
+    ~PyIteratorWrapper()
     {
-        // Do we need the GIL here?
-        return iterator{};
+        {
+            pybind11::gil_scoped_acquire gil;
+
+            if (m_initialized)
+            {
+                stop_python_thread(m_thread);
+            }
+
+            m_exception_queue = py::object();
+            m_queue           = py::object();
+            m_empty_exception = py::object();
+            m_thread          = py::object();
+        }
     }
 
   private:
+    void init()
+    {
+        using namespace py::literals;
+
+        AcquireGIL gil;
+        auto iter = m_iter_factory();
+
+        // Use Python to create/track the thread to allow the interpreter to shutdown safely
+        auto functools_mod = py::module_::import("functools");
+        auto queue_mod     = py::module_::import("queue");
+        auto threading_mod = py::module_::import("threading");
+
+        // We want the python bound version of iterator_thread
+        auto segment_mod = py::module_::import("mrc.core.segment");
+
+        m_empty_exception         = queue_mod.attr("Empty");
+        m_queue                   = queue_mod.attr("Queue")();
+        m_exception_queue         = queue_mod.attr("Queue")();
+        auto iter_thread_fn       = segment_mod.attr("_iterator_thread");
+        auto bound_iter_thread_fn = functools_mod.attr(
+            "partial")(iter_thread_fn, "itr"_a = iter, "queue"_a = m_queue, "exception_queue"_a = m_exception_queue);
+
+        m_thread = threading_mod.attr(
+            "Thread")("target"_a = bound_iter_thread_fn, "name"_a = "py_gen_src", "daemon"_a = true);
+        m_thread.attr("start")();
+
+        m_initialized = true;
+    }
+
+    bool m_initialized = false;
+    py::object m_thread{};
+    py::object m_exception_queue{};  // replace with exception queue
+    py::object m_queue{};
+    py::object m_empty_exception{};
     std::function<py::iterator()> m_iter_factory;
 };
 
 std::shared_ptr<mrc::segment::ObjectProperties> build_source(mrc::segment::IBuilder& self,
                                                              const std::string& name,
                                                              PyIteratorWrapper iter_wrapper)
 {
-    auto wrapper = [iter_wrapper = std::move(iter_wrapper)](PyObjectSubscriber& subscriber) mutable {
+    auto wrapper = [src_iter_wrapper = std::move(iter_wrapper)](PyObjectSubscriber& subscriber) mutable {
         auto& ctx = runnable::Context::get_runtime_context();
 
-        try
-        {
-            DVLOG(10) << ctx.info() << " Starting source";
+        DVLOG(10) << ctx.info() << " Starting source";
 
-            for (auto next_val : iter_wrapper)
+        // Taking a copy, if the source has pe_count>1 or engines_per_pe>1 we will need an indepdenent copy of the
+        // iterator wrapper
+        PyIteratorWrapper iter_wrapper = src_iter_wrapper;
+        bool received_stop_iteration   = false;
+        while (!received_stop_iteration && subscriber.is_subscribed())
+        {
+            try
             {
-                //  Only send if its subscribed. Very important to ensure the object has been moved!
+                auto next_val = iter_wrapper.get_next_value();
                 if (subscriber.is_subscribed())
                 {
                     subscriber.on_next(std::move(next_val));
                 }
-                else
+
+            } catch (EmptyQueue)
+            {
+                // No value
+                if (subscriber.is_subscribed())
                 {
-                    DVLOG(10) << ctx.info() << " Source unsubscribed. Stopping";
-                    break;
+                    boost::this_fiber::yield();
                 }
+            } catch (pybind11::stop_iteration)
+            {
+                DVLOG(10) << ctx.info() << " Received stop_iteration";
+                received_stop_iteration = true;
+            } catch (const std::exception& e)
+            {
+                LOG(ERROR) << ctx.info() << " Error occurred in source. Error msg: " << e.what();
+                subscriber.on_error(std::current_exception());
+                return;
             }
-
-        } catch (const std::exception& e)
-        {
-            LOG(ERROR) << ctx.info() << "Error occurred in source. Error msg: " << e.what();
-
-            // gil.release();
-            subscriber.on_error(std::current_exception());
-            return;
         }
 
         subscriber.on_completed();
@@ -266,22 +427,30 @@
         // // the GIL now to avoid dropping and reacquiring multiple times
         // AcquireGIL gil;
 
-        if (!current.is_initialized())
-        {
-            // On the first pass, initialize the iterator. This will generate the first value
-            current = iter_wrapper.begin();
-        }
-        else
+        auto& ctx = runnable::Context::get_runtime_context();
+
+        bool received_value          = false;
+        bool received_stop_iteration = false;
+        while (!received_value && !received_stop_iteration)
         {
-            // Otherwise pre-increment the iterator
-            ++current;
-        }
+            try
+            {
+                data           = iter_wrapper.get_next_value();
+                received_value = true;
 
-        // Copy the current value into data. No need for the GIL here due to PyHolder
-        data = *current;
+            } catch (EmptyQueue)
+            {
+                // No value
+                boost::this_fiber::yield();
+            } catch (pybind11::stop_iteration)
+            {
+                DVLOG(10) << ctx.info() << "Received stop_iteration";
+                received_stop_iteration = true;
+            }
+        }
 
         // Return closed if the current iterator is complete
-        return current == PyIteratorIterator() ? channel::Status::closed : channel::Status::success;
+        return received_stop_iteration ? channel::Status::closed : channel::Status::success;
     };
 
     return self.construct_object<PythonSourceComponent<PyHolder>>(name, std::move(get_next));

python/mrc/_pymrc/src/utilities/acquire_gil.cpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2021-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-FileCopyrightText: Copyright (c) 2021-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
@@ -17,16 +17,13 @@
 
 #include "pymrc/utilities/acquire_gil.hpp"
 
-#include <nlohmann/json_fwd.hpp>
 #include <pybind11/gil.h>
 #include <pybind11/pybind11.h>
 
 namespace mrc::pymrc {
 
 namespace py = pybind11;
 
-using nlohmann::json;
-
 AcquireGIL::AcquireGIL() : m_gil(std::make_unique<py::gil_scoped_acquire>()) {}
 
 AcquireGIL::~AcquireGIL() = default;

python/mrc/_pymrc/tests/test_executor.cpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2021-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-FileCopyrightText: Copyright (c) 2021-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  *
  * Licensed under the Apache License, Version 2.0 (the "License");

python/mrc/core/segment.cpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2021-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-FileCopyrightText: Copyright (c) 2021-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
@@ -251,6 +251,8 @@ PYBIND11_MODULE(segment, py_mod)
 
     Builder.def("make_node_full", &BuilderProxy::make_node_full, py::return_value_policy::reference_internal);
 
+    py_mod.def("_iterator_thread", &iterator_thread, py::arg("itr"), py::arg("queue"), py::arg("exception_queue"));
+
     py_mod.attr("__version__") = MRC_CONCAT_STR(mrc_VERSION_MAJOR << "." << mrc_VERSION_MINOR << "."
                                                                   << mrc_VERSION_PATCH);
 }