Fix for deadlock in python callback

Author: sgonorov

src/python/py_utils.cpp

@@ -374,10 +374,22 @@ ov::Any py_object_to_any(const py::object& py_obj, std::string property_name) {
         return py::cast<std::shared_ptr<ov::genai::Generator>>(py_obj);
     } else if (py::isinstance<py::function>(py_obj) && property_name == "callback") {
         auto py_callback = py::cast<py::function>(py_obj);
+        auto shared_callback = std::shared_ptr<py::function>(
+            new py::function(py_callback),
+            [](py::function* f) {
+                if (Py_IsInitialized()) {
+                    py::gil_scoped_acquire acquire;
+                    delete f;
+                } else {
+                    delete f;
+                }
+            }
+        );
+
         return std::function<bool(size_t, size_t, ov::Tensor&)>(
-            [py_callback](size_t step, size_t num_steps, ov::Tensor& latent) -> bool {
+            [shared_callback](size_t step, size_t num_steps, ov::Tensor& latent) -> bool {
                 py::gil_scoped_acquire acquire;
-                return py_callback(step, num_steps, latent).cast<bool>();
+                return (*shared_callback)(step, num_steps, latent).cast<bool>();
             }
         );
     } else if ((py::isinstance<py::function>(py_obj) || py::isinstance<ov::genai::StreamerBase>(py_obj) || py::isinstance<std::monostate>(py_obj)) && property_name == "streamer") {
@@ -443,21 +455,40 @@ ov::genai::StreamerVariant pystreamer_to_streamer(const PyBindStreamerVariant& p
 
     std::visit(overloaded {
         [&streamer](const std::function<std::optional<uint16_t>(py::str)>& py_callback){
-            // Wrap python streamer with manual utf-8 decoding. Do not rely
-            // on pybind automatic decoding since it raises exceptions on incomplete strings.
-            auto callback_wrapped = [py_callback](std::string subword) -> ov::genai::StreamingStatus {
+            auto shared_callback = std::shared_ptr<std::function<std::optional<uint16_t>(py::str)>>(
+                new std::function<std::optional<uint16_t>(py::str)>(py_callback),
+                [](std::function<std::optional<uint16_t>(py::str)>* f) {
+                    if (Py_IsInitialized()) {
+                        py::gil_scoped_acquire acquire;
+                        delete f;
+                    } else {
+                        delete f;
+                    }
+                }
+            );
+
+            auto callback_wrapped = [shared_callback](std::string subword) -> ov::genai::StreamingStatus {
                 py::gil_scoped_acquire acquire;
-                auto py_str = PyUnicode_DecodeUTF8(subword.data(), subword.length(), "replace");
-                std::optional<uint16_t> callback_output = py_callback(py::reinterpret_borrow<py::str>(py_str));
+                PyObject* py_str = PyUnicode_DecodeUTF8(subword.data(), subword.length(), "replace");
+                if (!py_str) {
+                    PyErr_WriteUnraisable(nullptr);
+                    return StreamingStatus::RUNNING;
+                }
+                auto py_str_obj = py::reinterpret_steal<py::str>(py_str);
+                std::optional<uint16_t> callback_output;
+                try {
+                    callback_output = (*shared_callback)(py_str_obj);
+                } catch (const py::error_already_set&) {
+                    return StreamingStatus::RUNNING;
+                }
                 if (callback_output.has_value()) {
-                    if (*callback_output == (uint16_t)StreamingStatus::RUNNING)
+                    if (*callback_output == static_cast<uint16_t>(StreamingStatus::RUNNING))
                         return StreamingStatus::RUNNING;
-                    else if (*callback_output == (uint16_t)StreamingStatus::CANCEL)
+                    else if (*callback_output == static_cast<uint16_t>(StreamingStatus::CANCEL))
                         return StreamingStatus::CANCEL;
                     return StreamingStatus::STOP;
-                } else {
-                    return StreamingStatus::RUNNING;
                 }
+                return StreamingStatus::RUNNING;
             };
             streamer = callback_wrapped;
         },

src/python/py_utils.hpp

@@ -19,7 +19,7 @@ namespace ov::genai::pybind::utils {
 // When StreamerVariant is used utf-8 decoding is done by pybind and can lead to exception on incomplete texts.
 // Therefore strings decoding should be handled with PyUnicode_DecodeUTF8(..., "replace") to not throw errors.
 using PyBindStreamerVariant = std::variant<
-    std::function<std::optional<uint16_t>(std::string)>,
+    std::function<std::optional<uint16_t>(py::str)>,
     std::shared_ptr<StreamerBase>,
     std::monostate>;