Add PythonFutureTask

.github/CODEOWNERS

@@ -1,21 +1,22 @@
 #cpp code owners
-cpp/                              @rapidsai/ucxx-cpp-codeowners
+cpp/                                @rapidsai/ucxx-cpp-codeowners
 
 #python code owners
-python/                           @rapidsai/ucxx-python-codeowners
+python/                             @rapidsai/ucxx-python-codeowners
 
 #cmake code owners
-cpp/CMakeLists.txt                @rapidsai/ucxx-cmake-codeowners
-cpp/python/CMakeLists.txt         @rapidsai/ucxx-cmake-codeowners
-cpp/examples/CMakeLists.txt       @rapidsai/ucxx-cmake-codeowners
-cpp/benchmarks/CMakeLists.txt     @rapidsai/ucxx-cmake-codeowners
-cpp/tests/CMakeLists.txt          @rapidsai/ucxx-cmake-codeowners
-python/CMakeLists.txt             @rapidsai/ucxx-cmake-codeowners
-python/ucxx/_lib/CMakeLists.txt   @rapidsai/ucxx-cmake-codeowners
-fetch_rapids.cmake                @rapidsai/ucxx-cmake-codeowners
-**/cmake/                         @rapidsai/ucxx-cmake-codeowners
+cpp/CMakeLists.txt                  @rapidsai/ucxx-cmake-codeowners
+cpp/python/CMakeLists.txt           @rapidsai/ucxx-cmake-codeowners
+cpp/examples/CMakeLists.txt         @rapidsai/ucxx-cmake-codeowners
+cpp/benchmarks/CMakeLists.txt       @rapidsai/ucxx-cmake-codeowners
+cpp/tests/CMakeLists.txt            @rapidsai/ucxx-cmake-codeowners
+python/CMakeLists.txt               @rapidsai/ucxx-cmake-codeowners
+python/ucxx/examples/CMakeLists.txt @rapidsai/ucxx-cmake-codeowners
+python/ucxx/_lib/CMakeLists.txt     @rapidsai/ucxx-cmake-codeowners
+fetch_rapids.cmake                  @rapidsai/ucxx-cmake-codeowners
+**/cmake/                           @rapidsai/ucxx-cmake-codeowners
 
 #build/ops code owners
-.github/                          @rapidsai/ops-codeowners
-ci/                               @rapidsai/ops-codeowners
-conda/                            @rapidsai/ops-codeowners
+.github/                            @rapidsai/ops-codeowners
+ci/                                 @rapidsai/ops-codeowners
+conda/                              @rapidsai/ops-codeowners

ci/test_python.sh

@@ -132,3 +132,6 @@ run_distributed_ucxx_tests      thread          0                           0
 run_distributed_ucxx_tests      thread          0                           1
 run_distributed_ucxx_tests      thread          1                           0
 run_distributed_ucxx_tests      thread          1                           1
+
+rapids-logger "C++ future -> Python future notifier example"
+python -m ucxx.examples.python_future_task_example

cpp/python/include/ucxx/python/constructors.h

@@ -19,6 +19,9 @@ namespace python {
 
 std::shared_ptr<::ucxx::Future> createFuture(std::shared_ptr<::ucxx::Notifier> notifier);
 
+std::shared_ptr<::ucxx::Future> createFutureWithEventLoop(
+  PyObject* asyncioEventLoop, std::shared_ptr<::ucxx::Notifier> notifier);
+
 std::shared_ptr<::ucxx::Notifier> createNotifier();
 
 std::shared_ptr<::ucxx::Worker> createWorker(std::shared_ptr<ucxx::Context> context,

cpp/python/include/ucxx/python/future.h

@@ -55,6 +55,60 @@ PyObject* future_set_result(PyObject* future, PyObject* value);
  */
 PyObject* future_set_exception(PyObject* future, PyObject* exception, const char* message);
 
+/**
+ * @brief Create a Python asyncio future with associated event loop.
+ *
+ * Create Python asyncio future associated with the event loop passed via the `event_loop`
+ * argument, effectively equal to calling `loop.create_future()` directly in Python.
+ *
+ * Note that this call will take the Python GIL and requires that the current thread have
+ * an asynchronous event loop set.
+ *
+ * @param[in] event_loop  the Python asyncio event loop to which the future will belong to.
+ *
+ * @returns The Python asyncio future object.
+ */
+PyObject* create_python_future_with_event_loop(PyObject* event_loop);
+
+/**
+ * @brief Set the result of a Python future with associated event loop.
+ *
+ * Schedule setting the result of a Python future in the given event loop using the
+ * threadsafe method `event_loop.call_soon_threadsafe`. The event loop given must be the
+ * same specified when creating the future object with `create_python_future`.
+ *
+ * Note that this may be called from any thread and will take the Python GIL to run.
+ *
+ * @param[in] future  Python object containing the `_asyncio.Future` object.
+ * @param[in] value   Python object containing an arbitrary value to set the future result
+ *                    to.
+ *
+ * @returns The result of the call to `_asyncio.Future.set_result()`.
+ */
+PyObject* future_set_result_with_event_loop(PyObject* event_loop,
+                                            PyObject* future,
+                                            PyObject* value);
+
+/**
+ * @brief Set the exception of a Python future with associated event loop.
+ *
+ * Schedule setting an exception of a Python future in the given event loop using the
+ * threadsafe method `event_loop.call_soon_threadsafe`. The event loop given must be the
+ * same specified when creating the future object with `create_python_future`.
+ *
+ * Note that this may be called from any thread and will take the Python GIL to run.
+ *
+ * @param[in] future    Python object containing the `_asyncio.Future` object.
+ * @param[in] exception a Python exception derived of the `Exception` class.
+ * @param[in] message   human-readable error message for the exception.
+ *
+ * @returns The result of the call to `_asyncio.Future.set_result()`.
+ */
+PyObject* future_set_exception_with_event_loop(PyObject* event_loop,
+                                               PyObject* future,
+                                               PyObject* exception,
+                                               const char* message);
+
 }  // namespace python
 
 }  // namespace ucxx

cpp/python/include/ucxx/python/python_future.h

@@ -21,7 +21,8 @@ namespace python {
 
 class Future : public ::ucxx::Future {
  private:
-  PyObject* _handle{create_python_future()};  ///< The handle to the Python future
+  PyObject* _asyncioEventLoop{nullptr};  ///< The asyncio event loop the Python future belongs to.
+  PyObject* _handle{nullptr};            ///< The handle to the Python future
 
   /**
    * @brief Construct a future that may be notified from a notifier thread.
@@ -32,9 +33,12 @@ class Future : public ::ucxx::Future {
    * This class may also be used to set the result or exception from any thread, but that
    * currently requires explicitly taking the GIL before calling `set()`.
    *
+   * @param[in] asyncioEventLoop pointer to a valid Python object containing the event loop
+   *                             that the application is using, to which the future will
+   *                             belong to.
    * @param[in] notifier  notifier object running on a separate thread.
    */
-  explicit Future(std::shared_ptr<::ucxx::Notifier> notifier);
+  explicit Future(PyObject* asyncioEventLoop, std::shared_ptr<::ucxx::Notifier> notifier);
 
  public:
   Future()                         = delete;
@@ -56,6 +60,23 @@ class Future : public ::ucxx::Future {
    */
   friend std::shared_ptr<::ucxx::Future> createFuture(std::shared_ptr<::ucxx::Notifier> notifier);
 
+  /**
+   * @brief Constructor of `shared_ptr<ucxx::python::Future>`.
+   *
+   * The constructor for a `shared_ptr<ucxx::python::Future>` object. The default
+   * constructor is made private to ensure all UCXX objects are shared pointers and correct
+   * lifetime management.
+   *
+   * @param[in] asyncioEventLoop  pointer to a valid Python object containing the event loop
+   *                              that the application is using, to which the future will
+   *                              belong to.
+   * @param[in] notifier          notifier object running on a separate thread.
+   *
+   * @returns The `shared_ptr<ucxx::python::Worker>` object
+   */
+  friend std::shared_ptr<::ucxx::Future> createFutureWithEventLoop(
+    PyObject* asyncioEventLoop, std::shared_ptr<::ucxx::Notifier> notifier);
+
   /**
    * @brief Virtual destructor.
    *

cpp/python/include/ucxx/python/python_future_task.h

@@ -0,0 +1,213 @@
+/**
+ * SPDX-FileCopyrightText: Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+#pragma once
+
+#include <chrono>
+#include <future>
+#include <memory>
+#include <queue>
+#include <sstream>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include <iostream>
+
+#include <Python.h>
+
+#include <ucxx/log.h>
+#include <ucxx/python/future.h>
+
+namespace ucxx {
+
+namespace python {
+
+template <typename ReturnType, typename... TaskArgs>
+class PythonFutureTask : public std::enable_shared_from_this<PythonFutureTask<ReturnType>> {
+ private:
+  std::packaged_task<ReturnType(TaskArgs...)> _task{};  ///< The user-defined C++ task to run
+  std::function<PyObject*(ReturnType)>
+    _pythonConvert{};                 ///< Function to convert the C++ result into Python value
+  PyObject* _asyncioEventLoop{};      ///< The handle to the Python asyncio event loop
+  PyObject* _handle{};                ///< The handle to the Python future
+  std::future<ReturnType> _future{};  ///< The C++ future containing the task result
+
+  /**
+   * @brief Set the result value of the Python future.
+   *
+   * Set the result value of the underlying Python future using the `pythonConvert` function
+   * specified in the constructor to convert the C++ result into the `PyObject*`.
+   *
+   * This function will take the GIL to convert the C++ result into the `PyObject*`.
+   *
+   * @param[in] result the C++ value that will be converted and set as the result of the
+   *                   Python future.
+   */
+  void setResult(const ReturnType result)
+  {
+    // PyLong_FromSize_t requires the GIL
+    if (_handle == nullptr) throw std::runtime_error("Invalid object or already released");
+    PyGILState_STATE state = PyGILState_Ensure();
+    ucxx::python::future_set_result_with_event_loop(
+      _asyncioEventLoop, _handle, _pythonConvert(result));
+    PyGILState_Release(state);
+  }
+
+  /**
+   * @brief Set the exception of the Python future.
+   *
+   * Set the exception of the underlying Python future. Currently any exceptions that the
+   * task may raise must be derived from `std::exception`.
+   *
+   * @param[in] pythonException the Python exception type to raise.
+   * @param[in] message the message of the exception.
+   */
+  void setPythonException(PyObject* pythonException, const std::string& message)
+  {
+    if (_handle == nullptr) throw std::runtime_error("Invalid object or already released");
+    ucxx::python::future_set_exception_with_event_loop(
+      _asyncioEventLoop, _handle, pythonException, message.c_str());
+  }
+
+  /**
+   * @brief Parse C++ exception as a Python exception and set the Python future exception.
+   *
+   * Parse a C++ exception as a Python exception and set the Python future exception.
+   * Currently any exceptions that the task may raise must be derived from `std::exception`.
+   *
+   * @param[in] exception the C++ exception that was raised by the user-defined task.
+   */
+  void setException(const std::exception& exception)
+  {
+    try {
+      throw exception;
+    } catch (const std::bad_alloc& e) {
+      setPythonException(PyExc_MemoryError, e.what());
+    } catch (const std::bad_cast& e) {
+      setPythonException(PyExc_TypeError, e.what());
+    } catch (const std::bad_typeid& e) {
+      setPythonException(PyExc_TypeError, e.what());
+    } catch (const std::domain_error& e) {
+      setPythonException(PyExc_ValueError, e.what());
+    } catch (const std::invalid_argument& e) {
+      setPythonException(PyExc_ValueError, e.what());
+    } catch (const std::ios_base::failure& e) {
+      setPythonException(PyExc_IOError, e.what());
+    } catch (const std::out_of_range& e) {
+      setPythonException(PyExc_IndexError, e.what());
+    } catch (const std::overflow_error& e) {
+      setPythonException(PyExc_OverflowError, e.what());
+    } catch (const std::range_error& e) {
+      setPythonException(PyExc_ArithmeticError, e.what());
+    } catch (const std::underflow_error& e) {
+      setPythonException(PyExc_ArithmeticError, e.what());
+    } catch (const std::exception& e) {
+      setPythonException(PyExc_RuntimeError, e.what());
+    } catch (...) {
+      setPythonException(PyExc_RuntimeError, "Unknown exception");
+    }
+  }
+
+ public:
+  /**
+   * @brief Construct a Python future backed by C++ `std::packaged_task`.
+   *
+   * Construct a future object that receives a user-defined C++ `std::packaged_task` which
+   * runs asynchronously using an internal `std::async` that ultimately notifies a Python
+   * future that can be awaited in Python code.
+   *
+   * @param[in] task the user-defined C++ task.
+   * @param[in] pythonConvert C-Python function to convert a C object into a `PyObject*`
+   *                          representing the result of the task.
+   * @param[in] asyncioEventLoop pointer to a valid Python object containing the event loop
+   *                             that the application is using, to which the Python future
+   *                             will belong to.
+   * @param[in] launchPolicy launch policy for the async C++ task.
+   */
+  explicit PythonFutureTask(std::packaged_task<ReturnType(TaskArgs...)> task,
+                            std::function<PyObject*(ReturnType)> pythonConvert,
+                            PyObject* asyncioEventLoop,
+                            std::launch launchPolicy = std::launch::async)
+    : _task{std::move(task)},
+      _pythonConvert(pythonConvert),
+      _asyncioEventLoop(asyncioEventLoop),
+      _handle{ucxx::python::create_python_future_with_event_loop(asyncioEventLoop)}
+  {
+    _future = std::async(launchPolicy, [this]() {
+      std::future<ReturnType> result = this->_task.get_future();
+      this->_task();
+      try {
+        const ReturnType r = result.get();
+        this->setResult(r);
+        return r;
+      } catch (std::exception& e) {
+        this->setException(e);
+      }
+    });
+  }
+  PythonFutureTask(const PythonFutureTask&)            = delete;
+  PythonFutureTask& operator=(PythonFutureTask const&) = delete;
+  PythonFutureTask(PythonFutureTask&& o)               = delete;
+  PythonFutureTask& operator=(PythonFutureTask&& o)    = delete;
+
+  /**
+   * @brief Python future destructor.
+   *
+   * Decrement the reference count of the underlying Python future.
+   */
+  ~PythonFutureTask() { Py_XDECREF(_handle); }
+
+  /**
+   * @brief Get the C++ future.
+   *
+   * Get the underlying C++ future that can be awaited and have its result read.
+   *
+   * @returns The underlying C++ future.
+   */
+  std::future<ReturnType>& getFuture() { return _future; }
+
+  /**
+   * @brief Get the underlying future `PyObject*` handle but does not release ownership.
+   *
+   * Get the underlying `PyObject*` handle without releasing ownership. This can be useful
+   * for example for logging, where we want to see the address of the pointer but do not
+   * want to transfer ownership.
+   *
+   * @warning The destructor will also destroy the Python future, a pointer taken via this
+   * method will cause the object to become invalid.
+   *
+   * @throws std::runtime_error if the object is invalid or has been already released.
+   *
+   * @returns The underlying `PyObject*` handle.
+   */
+  PyObject* getHandle()
+  {
+    if (_handle == nullptr) throw std::runtime_error("Invalid object or already released");
+
+    return _handle;
+  }
+
+  /**
+   * @brief Get the underlying future `PyObject*` handle and release ownership.
+   *
+   * Get the underlying `PyObject*` handle releasing ownership. This should be used when
+   * the future needs to be permanently transferred to Python code. After calling this
+   * method the object becomes invalid for any other uses.
+   *
+   * @throws std::runtime_error if the object is invalid or has been already released.
+   *
+   * @returns The underlying `PyObject*` handle.
+   */
+  PyObject* release()
+  {
+    if (_handle == nullptr) throw std::runtime_error("Invalid object or already released");
+
+    return std::exchange(_handle, nullptr);
+  }
+};
+
+}  // namespace python
+
+}  // namespace ucxx