diff --git a/ucp/__init__.py b/ucp/__init__.py index 2b338618a..f1d14fe7e 100644 --- a/ucp/__init__.py +++ b/ucp/__init__.py @@ -41,6 +41,14 @@ logger = get_ucxpy_logger() +try: + import uvloop + + uvloop.install() +except ImportError: + pass + + __version__ = _get_versions()["version"] __ucx_version__ = "%d.%d.%d" % get_ucx_version() diff --git a/ucp/continuous_ucx_progress.py b/ucp/continuous_ucx_progress.py index 2bf27c85a..d7068a274 100644 --- a/ucp/continuous_ucx_progress.py +++ b/ucp/continuous_ucx_progress.py @@ -82,7 +82,8 @@ def _fd_reader_callback(self): # 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() + if self.asyncio_task is None or self.asyncio_task.done(): + return self.asyncio_task = self.event_loop.create_task(self._arm_worker()) async def _arm_worker(self):