Add blocking progress mode to Python async

ci/test_python.sh

@@ -109,24 +109,28 @@ rapids-logger "Python Async Tests"
 # run_tests_async PROGRESS_MODE   ENABLE_DELAYED_SUBMISSION ENABLE_PYTHON_FUTURE SKIP
 run_tests_async   thread          0                         0                    0
 run_tests_async   thread          1                         1                    0
+run_tests_async   blocking        0                         0                    0
 
 rapids-logger "Python Benchmarks"
 # run_py_benchmark  BACKEND   PROGRESS_MODE   ASYNCIO_WAIT  ENABLE_DELAYED_SUBMISSION ENABLE_PYTHON_FUTURE NBUFFERS SLOW
 run_py_benchmark    ucxx-core thread          0             0                         0                    1        0
 run_py_benchmark    ucxx-core thread          1             0                         0                    1        0
 
-for nbuf in 1 8; do
-  if [[ ! $RAPIDS_CUDA_VERSION =~ 11.2.* ]]; then
-    # run_py_benchmark  BACKEND     PROGRESS_MODE   ASYNCIO_WAIT  ENABLE_DELAYED_SUBMISSION ENABLE_PYTHON_FUTURE NBUFFERS SLOW
-    run_py_benchmark    ucxx-async  thread          0             0                         0                    ${nbuf}  0
-    run_py_benchmark    ucxx-async  thread          0             0                         1                    ${nbuf}  0
-    run_py_benchmark    ucxx-async  thread          0             1                         0                    ${nbuf}  0
-    run_py_benchmark    ucxx-async  thread          0             1                         1                    ${nbuf}  0
-  fi
+for progress_mode in "blocking" "thread"; do
+  for nbuf in 1 8; do
+    if [[ ! $RAPIDS_CUDA_VERSION =~ 11.2.* ]]; then
+      # run_py_benchmark  BACKEND     PROGRESS_MODE     ASYNCIO_WAIT  ENABLE_DELAYED_SUBMISSION ENABLE_PYTHON_FUTURE NBUFFERS SLOW
+      run_py_benchmark    ucxx-async  ${progress_mode}  0             0                         0                    ${nbuf}  0
+      run_py_benchmark    ucxx-async  ${progress_mode}  0             0                         1                    ${nbuf}  0
+      run_py_benchmark    ucxx-async  ${progress_mode}  0             1                         0                    ${nbuf}  0
+      run_py_benchmark    ucxx-async  ${progress_mode}  0             1                         1                    ${nbuf}  0
+    fi
+  done
 done
 
 rapids-logger "Distributed Tests"
 # run_distributed_ucxx_tests    PROGRESS_MODE   ENABLE_DELAYED_SUBMISSION   ENABLE_PYTHON_FUTURE
+run_distributed_ucxx_tests      blocking        0                           0
 run_distributed_ucxx_tests      polling         0                           0
 run_distributed_ucxx_tests      thread          0                           0
 run_distributed_ucxx_tests      thread          0                           1

cpp/include/ucxx/worker.h

@@ -245,6 +245,23 @@ class Worker : public Component {
    */
   void initBlockingProgressMode();
 
+  /**
+   * @brief Get the epoll file descriptor associated with the worker.
+   *
+   * Get the epoll file descriptor associated with the worker when running in blocking mode.
+   * The worker only has an associated epoll file descriptor after
+   * `initBlockingProgressMode()` is executed.
+   *
+   * The file descriptor is destroyed as part of the `ucxx::Worker` destructor, thus any
+   * reference to it shall not be used after that.
+   *
+   * @throws std::runtime_error if `initBlockingProgressMode()` was not executed to run the
+   *                            worker in blocking progress mode.
+   *
+   * @returns the file descriptor.
+   */
+  int getEpollFileDescriptor();
+
   /**
    * @brief Arm the UCP worker.
    *

cpp/src/worker.cpp

@@ -207,6 +207,14 @@ void Worker::initBlockingProgressMode()
   if (err != 0) throw std::ios_base::failure(std::string("epoll_ctl() returned " + err));
 }
 
+int Worker::getEpollFileDescriptor()
+{
+  if (_epollFileDescriptor == 0)
+    throw std::runtime_error("Worker not running in blocking progress mode");
+
+  return _epollFileDescriptor;
+}
+
 bool Worker::arm()
 {
   ucs_status_t status = ucp_worker_arm(_handle);

python/ucxx/_lib/libucxx.pyx

@@ -504,6 +504,23 @@ cdef class UCXWorker():
         with nogil:
             self._worker.get().initBlockingProgressMode()
 
+    def arm(self):
+        cdef bint armed
+
+        with nogil:
+            armed = self._worker.get().arm()
+
+        return armed
+
+    @property
+    def epoll_file_descriptor(self):
+        cdef int epoll_file_descriptor = 0
+
+        with nogil:
+            epoll_file_descriptor = self._worker.get().getEpollFileDescriptor()
+
+        return epoll_file_descriptor
+
     def progress(self):
         with nogil:
             self._worker.get().progress()

python/ucxx/_lib/ucxx_api.pxd

@@ -228,6 +228,8 @@ cdef extern from "<ucxx/api.h>" namespace "ucxx" nogil:
             uint16_t port, ucp_listener_conn_callback_t callback, void *callback_args
         ) except +raise_py_error
         void initBlockingProgressMode() except +raise_py_error
+        int getEpollFileDescriptor()
+        bint arm() except +raise_py_error
         void progress()
         bint progressOnce()
         void progressWorkerEvent(int epoll_timeout)

python/ucxx/_lib_async/application_context.py

@@ -11,7 +11,7 @@
 from ucxx._lib.arr import Array
 from ucxx.exceptions import UCXMessageTruncatedError
 
-from .continuous_ucx_progress import PollingMode, ThreadMode
+from .continuous_ucx_progress import BlockingMode, PollingMode, ThreadMode
 from .endpoint import Endpoint
 from .exchange_peer_info import exchange_peer_info
 from .listener import ActiveClients, Listener, _listener_handler
@@ -60,6 +60,9 @@ def __init__(
             enable_python_future=enable_python_future,
         )
 
+        if self.progress_mode == "blocking":
+            self.worker.init_blocking_progress_mode()
+
         self.start_notifier_thread()
 
         weakref.finalize(self, self.progress_tasks.clear)
@@ -77,7 +80,7 @@ def _check_progress_mode(progress_mode):
             else:
                 progress_mode = "thread"
 
-        valid_progress_modes = ["polling", "thread", "thread-polling"]
+        valid_progress_modes = ["blocking", "polling", "thread", "thread-polling"]
         if not isinstance(progress_mode, str) or not any(
             progress_mode == m for m in valid_progress_modes
         ):
@@ -107,7 +110,7 @@ def _check_enable_delayed_submission(enable_delayed_submission, progress_mode):
             and explicit_enable_delayed_submission
         ):
             raise ValueError(
-                f"Delayed submission requested, but {progress_mode} does not "
+                f"Delayed submission requested, but '{progress_mode}' mode does not "
                 "support it, 'thread' or 'thread-polling' progress mode required."
             )
 
@@ -400,6 +403,8 @@ def continuous_ucx_progress(self, event_loop=None):
             task = ThreadMode(self.worker, loop, polling_mode=True)
         elif self.progress_mode == "polling":
             task = PollingMode(self.worker, loop)
+        elif self.progress_mode == "blocking":
+            task = BlockingMode(self.worker, loop, self.worker.epoll_file_descriptor)
 
         self.progress_tasks.append(task)
 

python/ucxx/_lib_async/continuous_ucx_progress.py

@@ -3,6 +3,8 @@
 
 
 import asyncio
+import socket
+import weakref
 
 
 class ProgressTask(object):
@@ -70,3 +72,50 @@ async def _progress_task(self):
             worker.progress()
             # Give other co-routines a chance to run.
             await asyncio.sleep(0)
+
+
+class BlockingMode(ProgressTask):
+    def __init__(self, worker, event_loop, epoll_fd):
+        super().__init__(worker, event_loop)
+
+        # Creating a job that is ready straightaway but with low priority.
+        # Calling `await self.event_loop.sock_recv(self.rsock, 1)` will
+        # return when all non-IO tasks are finished.
+        # See <https://stackoverflow.com/a/48491563>.
+        self.rsock, wsock = socket.socketpair()
+        self.rsock.setblocking(0)
+        wsock.setblocking(0)
+        wsock.close()
+
+        # Bind an asyncio reader to a UCX epoll file descripter
+        event_loop.add_reader(epoll_fd, self._fd_reader_callback)
+
+        # Remove the reader and close socket on finalization
+        weakref.finalize(self, event_loop.remove_reader, epoll_fd)
+        weakref.finalize(self, self.rsock.close)
+
+    def _fd_reader_callback(self):
+        self.worker.progress()
+
+        # Notice, we can safely overwrite `self.dangling_arm_task`
+        # since previous arm task is finished by now.
+        assert self.asyncio_task is None or self.asyncio_task.done()
+        self.asyncio_task = self.event_loop.create_task(self._arm_worker())
+
+    async def _arm_worker(self):
+        # When arming the worker, the following must be true:
+        #  - No more progress in UCX (see doc of ucp_worker_arm())
+        #  - All asyncio tasks that isn't waiting on UCX must be executed
+        #    so that the asyncio's next state is epoll wait.
+        #    See <https://github.com/rapidsai/ucx-py/issues/413>
+        while True:
+            self.worker.progress()
+
+            # This IO task returns when all non-IO tasks are finished.
+            # Notice, we do NOT hold a reference to `worker` while waiting.
+            await self.event_loop.sock_recv(self.rsock, 1)
+
+            if self.worker.arm():
+                # At this point we know that asyncio's next state is
+                # epoll wait.
+                break