From 1ebd4e269f74a8a9e7aa288a0454da7c6fb5f082 Mon Sep 17 00:00:00 2001 From: Christopher Harris Date: Wed, 25 Oct 2023 14:47:18 -0500 Subject: [PATCH] Move Pycoro from Morpheus to MRC (#409) Moves pycoro from Morpheus to MRC and incorperates tests from https://github.com/nv-morpheus/Morpheus/pull/1286 Closes https://github.com/nv-morpheus/Morpheus/issues/1268 Authors: - Christopher Harris (https://github.com/cwharris) Approvers: - Devin Robison (https://github.com/drobison00) URL: https://github.com/nv-morpheus/MRC/pull/409 --- ci/conda/environments/dev_env.yml | 1 + python/mrc/_pymrc/CMakeLists.txt | 1 + python/mrc/_pymrc/include/pymrc/coro.hpp | 433 ++++++++++++++++++++ python/mrc/_pymrc/src/coro.cpp | 26 ++ python/mrc/_pymrc/tests/CMakeLists.txt | 2 + python/mrc/_pymrc/tests/coro/CMakeLists.txt | 29 ++ python/mrc/_pymrc/tests/coro/module.cpp | 70 ++++ python/mrc/core/CMakeLists.txt | 1 + python/mrc/core/coro.cpp | 68 +++ python/tests/test_coro.py | 152 +++++++ 10 files changed, 783 insertions(+) create mode 100644 python/mrc/_pymrc/include/pymrc/coro.hpp create mode 100644 python/mrc/_pymrc/src/coro.cpp create mode 100644 python/mrc/_pymrc/tests/coro/CMakeLists.txt create mode 100644 python/mrc/_pymrc/tests/coro/module.cpp create mode 100644 python/mrc/core/coro.cpp create mode 100644 python/tests/test_coro.py diff --git a/ci/conda/environments/dev_env.yml b/ci/conda/environments/dev_env.yml index 5af8a91c9..08d31e0c0 100644 --- a/ci/conda/environments/dev_env.yml +++ b/ci/conda/environments/dev_env.yml @@ -59,6 +59,7 @@ dependencies: - pybind11-stubgen=0.10 - pytest - pytest-timeout + - pytest-asyncio - python=3.10 - scikit-build>=0.17 - sysroot_linux-64=2.17 diff --git a/python/mrc/_pymrc/CMakeLists.txt b/python/mrc/_pymrc/CMakeLists.txt index 2e81eac88..ed385504f 100644 --- a/python/mrc/_pymrc/CMakeLists.txt +++ b/python/mrc/_pymrc/CMakeLists.txt @@ -18,6 +18,7 @@ find_package(prometheus-cpp REQUIRED) # Keep all source files sorted!!! add_library(pymrc + src/coro.cpp src/executor.cpp src/logging.cpp src/module_registry.cpp diff --git a/python/mrc/_pymrc/include/pymrc/coro.hpp b/python/mrc/_pymrc/include/pymrc/coro.hpp new file mode 100644 index 000000000..5c80398cc --- /dev/null +++ b/python/mrc/_pymrc/include/pymrc/coro.hpp @@ -0,0 +1,433 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +// Dont directly include python headers +// IWYU pragma: no_include + +namespace mrc::pymrc::coro { + +class PYBIND11_EXPORT StopIteration : public pybind11::stop_iteration +{ + public: + StopIteration(pybind11::object&& result) : stop_iteration("--"), m_result(std::move(result)){}; + ~StopIteration() override; + + void set_error() const override + { + PyErr_SetObject(PyExc_StopIteration, this->m_result.ptr()); + } + + private: + pybind11::object m_result; +}; + +class PYBIND11_EXPORT CppToPyAwaitable : public std::enable_shared_from_this +{ + public: + CppToPyAwaitable() = default; + + template + CppToPyAwaitable(mrc::coroutines::Task&& task) + { + auto converter = [](mrc::coroutines::Task incoming_task) -> mrc::coroutines::Task { + DCHECK_EQ(PyGILState_Check(), 0) << "Should not have the GIL when resuming a C++ coroutine"; + + mrc::pymrc::PyHolder holder; + + if constexpr (std::is_same_v) + { + co_await incoming_task; + + // Need the GIL to make the return object + pybind11::gil_scoped_acquire gil; + + holder = pybind11::none(); + } + else + { + auto result = co_await incoming_task; + + // Need the GIL to cast the return object + pybind11::gil_scoped_acquire gil; + + holder = pybind11::cast(std::move(result)); + } + + co_return holder; + }; + + m_task = converter(std::move(task)); + } + + CppToPyAwaitable(mrc::coroutines::Task&& task) : m_task(std::move(task)) {} + + std::shared_ptr iter() + { + return this->shared_from_this(); + } + + std::shared_ptr await() + { + return this->shared_from_this(); + } + + void next() + { + // Need to release the GIL before waiting + pybind11::gil_scoped_release nogil; + + // Run the tick function which will resume the coroutine + this->tick(); + + if (m_task.is_ready()) + { + pybind11::gil_scoped_acquire gil; + + // job done -> throw + auto exception = StopIteration(std::move(m_task.promise().result())); + + // Destroy the task now that we have the value + m_task.destroy(); + + throw exception; + } + } + + protected: + virtual void tick() + { + if (!m_has_resumed) + { + m_has_resumed = true; + + m_task.resume(); + } + } + + bool m_has_resumed{false}; + mrc::coroutines::Task m_task; +}; + +/** + * @brief Similar to CppToPyAwaitable but will yield to other fibers when waiting for the coroutine to finish. Use this + * once per loop at the main entry point for the asyncio loop + * + */ +class PYBIND11_EXPORT BoostFibersMainPyAwaitable : public CppToPyAwaitable +{ + public: + using CppToPyAwaitable::CppToPyAwaitable; + + protected: + void tick() override + { + // Call the base class and then see if any fibers need processing by calling yield + CppToPyAwaitable::tick(); + + bool has_fibers = boost::fibers::has_ready_fibers(); + + if (has_fibers) + { + // Yield to other fibers + boost::this_fiber::yield(); + } + } +}; + +class PYBIND11_EXPORT PyTaskToCppAwaitable +{ + public: + PyTaskToCppAwaitable() = default; + PyTaskToCppAwaitable(mrc::pymrc::PyObjectHolder&& task) : m_task(std::move(task)) + { + pybind11::gil_scoped_acquire acquire; + if (pybind11::module_::import("inspect").attr("iscoroutine")(m_task).cast()) + { + m_task = pybind11::module_::import("asyncio").attr("create_task")(m_task); + } + } + + static bool await_ready() noexcept // NOLINT(readability-convert-member-functions-to-static) + { + // Always suspend + return false; + } + + void await_suspend(std::coroutine_handle<> caller) noexcept + { + pybind11::gil_scoped_acquire gil; + + auto done_callback = pybind11::cpp_function([this, caller](pybind11::object future) { + try + { + // Save the result value + m_result = future.attr("result")(); + } catch (pybind11::error_already_set) + { + m_exception_ptr = std::current_exception(); + } + + pybind11::gil_scoped_release nogil; + + // Resume the coroutine + caller.resume(); + }); + + m_task.attr("add_done_callback")(done_callback); + } + + mrc::pymrc::PyHolder await_resume() + { + if (m_exception_ptr) + { + std::rethrow_exception(m_exception_ptr); + } + + return std::move(m_result); + } + + private: + mrc::pymrc::PyObjectHolder m_task; + mrc::pymrc::PyHolder m_result; + std::exception_ptr m_exception_ptr; +}; + +// ====== HELPER MACROS ====== + +#define MRC_PYBIND11_FAIL_ABSTRACT(cname, fnname) \ + pybind11::pybind11_fail(MRC_CONCAT_STR("Tried to call pure virtual function \"" << PYBIND11_STRINGIFY(cname) \ + << "::" << fnname << "\"")); + +// ====== OVERRIDE PURE TEMPLATE ====== +#define MRC_PYBIND11_OVERRIDE_PURE_TEMPLATE_NAME(ret_type, abstract_cname, cname, name, fn, ...) \ + do \ + { \ + PYBIND11_OVERRIDE_IMPL(PYBIND11_TYPE(ret_type), PYBIND11_TYPE(cname), name, __VA_ARGS__); \ + if constexpr (std::is_same_v) \ + { \ + MRC_PYBIND11_FAIL_ABSTRACT(PYBIND11_TYPE(abstract_cname), name); \ + } \ + else \ + { \ + return cname::fn(__VA_ARGS__); \ + } \ + } while (false) + +#define MRC_PYBIND11_OVERRIDE_PURE_TEMPLATE(ret_type, abstract_cname, cname, fn, ...) \ + MRC_PYBIND11_OVERRIDE_PURE_TEMPLATE_NAME(PYBIND11_TYPE(ret_type), \ + PYBIND11_TYPE(abstract_cname), \ + PYBIND11_TYPE(cname), \ + #fn, \ + fn, \ + __VA_ARGS__) +// ====== OVERRIDE PURE TEMPLATE ====== + +// ====== OVERRIDE COROUTINE IMPL ====== +#define MRC_PYBIND11_OVERRIDE_CORO_IMPL(ret_type, cname, name, ...) \ + do \ + { \ + DCHECK_EQ(PyGILState_Check(), 0) << "Should not have the GIL when resuming a C++ coroutine"; \ + pybind11::gil_scoped_acquire gil; \ + pybind11::function override = pybind11::get_override(static_cast(this), name); \ + if (override) \ + { \ + auto o_coro = override(__VA_ARGS__); \ + auto asyncio_module = pybind11::module::import("asyncio"); \ + /* Return type must be a coroutine to allow calling asyncio.create_task() */ \ + if (!asyncio_module.attr("iscoroutine")(o_coro).cast()) \ + { \ + pybind11::pybind11_fail(MRC_CONCAT_STR("Return value from overriden async function " \ + << PYBIND11_STRINGIFY(cname) << "::" << name \ + << " did not return a coroutine. Returned: " \ + << pybind11::str(o_coro).cast())); \ + } \ + auto o_task = asyncio_module.attr("create_task")(o_coro); \ + mrc::pymrc::PyHolder o_result; \ + { \ + pybind11::gil_scoped_release nogil; \ + o_result = co_await mrc::pymrc::coro::PyTaskToCppAwaitable(std::move(o_task)); \ + DCHECK_EQ(PyGILState_Check(), 0) << "Should not have the GIL after returning from co_await"; \ + } \ + if (pybind11::detail::cast_is_temporary_value_reference::value) \ + { \ + static pybind11::detail::override_caster_t caster; \ + co_return pybind11::detail::cast_ref(std::move(o_result), caster); \ + } \ + co_return pybind11::detail::cast_safe(std::move(o_result)); \ + } \ + } while (false) +// ====== OVERRIDE COROUTINE IMPL====== + +// ====== OVERRIDE COROUTINE ====== +#define MRC_PYBIND11_OVERRIDE_CORO_NAME(ret_type, cname, name, fn, ...) \ + do \ + { \ + MRC_PYBIND11_OVERRIDE_CORO_IMPL(PYBIND11_TYPE(ret_type), PYBIND11_TYPE(cname), name, __VA_ARGS__); \ + return cname::fn(__VA_ARGS__); \ + } while (false) + +#define MRC_PYBIND11_OVERRIDE_CORO(ret_type, cname, fn, ...) \ + MRC_PYBIND11_OVERRIDE_CORO_NAME(PYBIND11_TYPE(ret_type), PYBIND11_TYPE(cname), #fn, fn, __VA_ARGS__) +// ====== OVERRIDE COROUTINE ====== + +// ====== OVERRIDE COROUTINE PURE====== +#define MRC_PYBIND11_OVERRIDE_CORO_PURE_NAME(ret_type, cname, name, fn, ...) \ + do \ + { \ + MRC_PYBIND11_OVERRIDE_CORO_IMPL(PYBIND11_TYPE(ret_type), PYBIND11_TYPE(cname), name, __VA_ARGS__); \ + MRC_PYBIND11_FAIL_ABSTRACT(PYBIND11_TYPE(cname), name); \ + } while (false) + +#define MRC_PYBIND11_OVERRIDE_CORO_PURE(ret_type, cname, fn, ...) \ + MRC_PYBIND11_OVERRIDE_CORO_PURE_NAME(PYBIND11_TYPE(ret_type), PYBIND11_TYPE(cname), #fn, fn, __VA_ARGS__) +// ====== OVERRIDE COROUTINE PURE====== + +// ====== OVERRIDE COROUTINE PURE TEMPLATE====== +#define MRC_PYBIND11_OVERRIDE_CORO_PURE_TEMPLATE_NAME(ret_type, abstract_cname, cname, name, fn, ...) \ + do \ + { \ + MRC_PYBIND11_OVERRIDE_CORO_IMPL(PYBIND11_TYPE(ret_type), PYBIND11_TYPE(cname), name, __VA_ARGS__); \ + if constexpr (std::is_same_v) \ + { \ + MRC_PYBIND11_FAIL_ABSTRACT(PYBIND11_TYPE(abstract_cname), name); \ + } \ + else \ + { \ + co_return co_await cname::fn(__VA_ARGS__); \ + } \ + } while (false) + +#define MRC_PYBIND11_OVERRIDE_CORO_PURE_TEMPLATE(ret_type, abstract_cname, cname, fn, ...) \ + MRC_PYBIND11_OVERRIDE_CORO_PURE_TEMPLATE_NAME(PYBIND11_TYPE(ret_type), \ + PYBIND11_TYPE(abstract_cname), \ + PYBIND11_TYPE(cname), \ + #fn, \ + fn, \ + __VA_ARGS__) +// ====== OVERRIDE COROUTINE PURE TEMPLATE====== + +} // namespace mrc::pymrc::coro + +// NOLINTNEXTLINE(modernize-concat-nested-namespaces) +namespace PYBIND11_NAMESPACE { +namespace detail { + +/** + * @brief Provides a type caster for converting a C++ coroutine to a python awaitable. Include this file in any pybind11 + * module to automatically convert the types. Allows for converting arguments and return values. + * + * @tparam ReturnT The return type of the coroutine + */ +template +struct type_caster> +{ + public: + /** + * This macro establishes the name 'inty' in + * function signatures and declares a local variable + * 'value' of type inty + */ + PYBIND11_TYPE_CASTER(mrc::coroutines::Task, _("typing.Awaitable[") + make_caster::name + _("]")); + + /** + * Conversion part 1 (Python->C++): convert a PyObject into a inty + * instance or return false upon failure. The second argument + * indicates whether implicit conversions should be applied. + */ + bool load(handle src, bool convert) + { + if (!src || src.is_none()) + { + return false; + } + + if (!PyCoro_CheckExact(src.ptr())) + { + return false; + } + + auto cpp_coro = [](mrc::pymrc::PyHolder py_task) -> mrc::coroutines::Task { + DCHECK_EQ(PyGILState_Check(), 0) << "Should not have the GIL when resuming a C++ coroutine"; + + // Always assume we are resuming without the GIL + pybind11::gil_scoped_acquire gil; + + auto asyncio_task = pybind11::module_::import("asyncio").attr("create_task")(py_task); + + mrc::pymrc::PyHolder py_result; + { + // Release the GIL before awaiting + pybind11::gil_scoped_release nogil; + + py_result = co_await mrc::pymrc::coro::PyTaskToCppAwaitable(std::move(asyncio_task)); + } + + // Now cast back to the C++ type + if (pybind11::detail::cast_is_temporary_value_reference::value) + { + static pybind11::detail::override_caster_t caster; + co_return pybind11::detail::cast_ref(std::move(py_result), caster); + } + co_return pybind11::detail::cast_safe(std::move(py_result)); + }; + + value = cpp_coro(pybind11::reinterpret_borrow(std::move(src))); + + return true; + } + + /** + * Conversion part 2 (C++ -> Python): convert an inty instance into + * a Python object. The second and third arguments are used to + * indicate the return value policy and parent object (for + * ``return_value_policy::reference_internal``) and are generally + * ignored by implicit casters. + */ + static handle cast(mrc::coroutines::Task src, return_value_policy policy, handle parent) + { + // Wrap the object in a CppToPyAwaitable + std::shared_ptr awaitable = + std::make_shared(std::move(src)); + + // Convert the object to a python object + auto py_awaitable = pybind11::cast(std::move(awaitable)); + + return py_awaitable.release(); + } +}; + +} // namespace detail +} // namespace PYBIND11_NAMESPACE diff --git a/python/mrc/_pymrc/src/coro.cpp b/python/mrc/_pymrc/src/coro.cpp new file mode 100644 index 000000000..8bb57cb84 --- /dev/null +++ b/python/mrc/_pymrc/src/coro.cpp @@ -0,0 +1,26 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "pymrc/coro.hpp" + +namespace mrc::pymrc::coro { + +namespace py = pybind11; + +StopIteration::~StopIteration() = default; + +} // namespace mrc::pymrc::coro diff --git a/python/mrc/_pymrc/tests/CMakeLists.txt b/python/mrc/_pymrc/tests/CMakeLists.txt index 4ac354a78..f40e20d72 100644 --- a/python/mrc/_pymrc/tests/CMakeLists.txt +++ b/python/mrc/_pymrc/tests/CMakeLists.txt @@ -17,6 +17,8 @@ list(APPEND CMAKE_MESSAGE_CONTEXT "tests") find_package(pybind11 REQUIRED) +add_subdirectory(coro) + # Keep all source files sorted!!! add_executable(test_pymrc test_codable_pyobject.cpp diff --git a/python/mrc/_pymrc/tests/coro/CMakeLists.txt b/python/mrc/_pymrc/tests/coro/CMakeLists.txt new file mode 100644 index 000000000..788d04832 --- /dev/null +++ b/python/mrc/_pymrc/tests/coro/CMakeLists.txt @@ -0,0 +1,29 @@ +# ============================================================================= +# Copyright (c) 2022-2023, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except +# in compliance with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software distributed under the License +# is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +# or implied. See the License for the specific language governing permissions and limitations under +# the License. +# ============================================================================= + +list(APPEND CMAKE_MESSAGE_CONTEXT "coro") + +set(MODULE_SOURCE_FILES) + +# Add the module file +list(APPEND MODULE_SOURCE_FILES module.cpp) + +# Create the python module +mrc_add_pybind11_module(coro + INCLUDE_DIRS ${CMAKE_CURRENT_SOURCE_DIR}/include + SOURCE_FILES ${MODULE_SOURCE_FILES} + LINK_TARGETS mrc::pymrc +) + +list(POP_BACK CMAKE_MESSAGE_CONTEXT) diff --git a/python/mrc/_pymrc/tests/coro/module.cpp b/python/mrc/_pymrc/tests/coro/module.cpp new file mode 100644 index 000000000..c5332c78c --- /dev/null +++ b/python/mrc/_pymrc/tests/coro/module.cpp @@ -0,0 +1,70 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +#include +#include + +#include +#include + +mrc::coroutines::Task subtract(int a, int b) +{ + co_return a - b; +} + +mrc::coroutines::Task call_fib_async(mrc::pymrc::PyHolder fib, int value, int minus) +{ + auto result = co_await subtract(value, minus); + co_return co_await mrc::pymrc::coro::PyTaskToCppAwaitable([](auto fib, auto result) { + pybind11::gil_scoped_acquire acquire; + return fib(result); + }(fib, result)); +} + +mrc::coroutines::Task raise_at_depth_async(mrc::pymrc::PyHolder fn, int depth) +{ + if (depth <= 0) + { + throw std::runtime_error("depth reached zero in c++"); + } + + co_return co_await mrc::pymrc::coro::PyTaskToCppAwaitable([](auto fn, auto depth) { + pybind11::gil_scoped_acquire acquire; + return fn(depth - 1); + }(fn, depth)); +} + +mrc::coroutines::Task call_async(mrc::pymrc::PyHolder fn) +{ + co_return co_await mrc::pymrc::coro::PyTaskToCppAwaitable([](auto fn) { + pybind11::gil_scoped_acquire acquire; + return fn(); + }(fn)); +} + +PYBIND11_MODULE(coro, _module) +{ + pybind11::module_::import("mrc.core.coro"); // satisfies automatic type conversions for tasks + + _module.def("call_fib_async", &call_fib_async); + _module.def("raise_at_depth_async", &raise_at_depth_async); + _module.def("call_async", &call_async); +} diff --git a/python/mrc/core/CMakeLists.txt b/python/mrc/core/CMakeLists.txt index d635e071f..f04b17f1f 100644 --- a/python/mrc/core/CMakeLists.txt +++ b/python/mrc/core/CMakeLists.txt @@ -16,6 +16,7 @@ list(APPEND CMAKE_MESSAGE_CONTEXT "core") mrc_add_pybind11_module(common SOURCE_FILES common.cpp) +mrc_add_pybind11_module(coro SOURCE_FILES coro.cpp) mrc_add_pybind11_module(executor SOURCE_FILES executor.cpp) mrc_add_pybind11_module(logging SOURCE_FILES logging.cpp) mrc_add_pybind11_module(node SOURCE_FILES node.cpp) diff --git a/python/mrc/core/coro.cpp b/python/mrc/core/coro.cpp new file mode 100644 index 000000000..8139ce9ec --- /dev/null +++ b/python/mrc/core/coro.cpp @@ -0,0 +1,68 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "pymrc/coro.hpp" + +#include +#include +#include +#include +#include +#include // IWYU pragma: keep + +#include +#include +#include +#include +#include + +namespace mrc::pymrc::coro { + +namespace py = pybind11; + +PYBIND11_MODULE(coro, _module) +{ + _module.doc() = R"pbdoc( + ----------------------- + .. currentmodule:: morpheus.llm + .. autosummary:: + :toctree: _generate + + )pbdoc"; + + py::class_>(_module, "CppToPyAwaitable") + .def(py::init<>()) + .def("__iter__", &CppToPyAwaitable::iter) + .def("__await__", &CppToPyAwaitable::await) + .def("__next__", &CppToPyAwaitable::next); + + py::class_>( // + _module, + "BoostFibersMainPyAwaitable") + .def(py::init<>()); + + _module.def("wrap_coroutine", [](coroutines::Task> fn) -> coroutines::Task { + DCHECK_EQ(PyGILState_Check(), 0) << "Should not have the GIL when resuming a C++ coroutine"; + + auto strings = co_await fn; + + co_return strings[0]; + }); + + // _module.attr("__version__") = + // MRC_CONCAT_STR(morpheus_VERSION_MAJOR << "." << morpheus_VERSION_MINOR << "." << morpheus_VERSION_PATCH); +} +} // namespace mrc::pymrc::coro diff --git a/python/tests/test_coro.py b/python/tests/test_coro.py new file mode 100644 index 000000000..940160f18 --- /dev/null +++ b/python/tests/test_coro.py @@ -0,0 +1,152 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import asyncio + +import pytest + +from mrc._pymrc.tests.coro.coro import call_async +from mrc._pymrc.tests.coro.coro import call_fib_async +from mrc._pymrc.tests.coro.coro import raise_at_depth_async +from mrc.core import coro + + +@pytest.mark.asyncio +async def test_coro(): + + # hit_inside = False + + async def inner(): + + # nonlocal hit_inside + + result = await coro.wrap_coroutine(asyncio.sleep(1, result=['a', 'b', 'c'])) + + # hit_inside = True + + return [result] + + returned_val = await coro.wrap_coroutine(inner()) + + assert returned_val == 'a' + # assert hit_inside + + +@pytest.mark.asyncio +async def test_coro_many(): + + expected_count = 1000 + hit_count = 0 + + start_time = asyncio.get_running_loop().time() + + async def inner(): + + nonlocal hit_count + + await asyncio.sleep(0.1) + + hit_count += 1 + + return ['a', 'b', 'c'] + + coros = [coro.wrap_coroutine(inner()) for _ in range(expected_count)] + + returned_vals = await asyncio.gather(*coros) + + end_time = asyncio.get_running_loop().time() + + assert returned_vals == ['a'] * expected_count + assert hit_count == expected_count + assert (end_time - start_time) < 1.5 + + +@pytest.mark.asyncio +async def test_python_cpp_async_interleave(): + + def fib(n): + if n < 0: + raise ValueError() + + if n < 2: + return 1 + + return fib(n - 1) + fib(n - 2) + + async def fib_async(n): + if n < 0: + raise ValueError() + + if n < 2: + return 1 + + task_a = call_fib_async(fib_async, n, 1) + task_b = call_fib_async(fib_async, n, 2) + + [a, b] = await asyncio.gather(task_a, task_b) + + return a + b + + assert fib(15) == await fib_async(15) + + +@pytest.mark.asyncio +async def test_python_cpp_async_exception(): + + async def py_raise_at_depth_async(n: int): + if n <= 0: + raise RuntimeError("depth reached zero in python") + + await raise_at_depth_async(py_raise_at_depth_async, n - 1) + + depth = 100 + + with pytest.raises(RuntimeError) as ex: + await raise_at_depth_async(py_raise_at_depth_async, depth + 1) + assert "python" in str(ex.value) + + with pytest.raises(RuntimeError) as ex: + await raise_at_depth_async(py_raise_at_depth_async, depth) + assert "c++" in str(ex.value) + + +@pytest.mark.asyncio +async def test_can_cancel_coroutine_from_python(): + + counter = 0 + + async def increment_recursively(): + nonlocal counter + await asyncio.sleep(0) + counter += 1 + await call_async(increment_recursively) + + task = asyncio.ensure_future(call_async(increment_recursively)) + + await asyncio.sleep(0) + assert counter == 0 + await asyncio.sleep(0) + await asyncio.sleep(0) + assert counter == 1 + await asyncio.sleep(0) + await asyncio.sleep(0) + assert counter == 2 + + task.cancel() + + with pytest.raises(asyncio.exceptions.CancelledError): + await task + + assert counter == 3