Upgrade pytest-asyncio to 0.26.0 (#652)

* Upgrade pytest-asyncio to 0.26.0

* Add newline to pyproject.toml

* Remove the fixture override

* Cancel tasks for ControlClient only if the event loop is running

* Use cancel_task_if_running more broadly

Author: mryab

hivemind/p2p/p2p_daemon.py

@@ -20,7 +20,7 @@
 from hivemind.p2p.p2p_daemon_bindings.utils import ControlFailure
 from hivemind.proto import crypto_pb2
 from hivemind.proto.p2pd_pb2 import RPCError
-from hivemind.utils.asyncio import as_aiter, asingle
+from hivemind.utils.asyncio import as_aiter, asingle, cancel_task_if_running
 from hivemind.utils.crypto import RSAPrivateKey
 from hivemind.utils.logging import get_logger, golog_level_to_python, loglevel, python_level_to_golog
 from hivemind.utils.multiaddr import Multiaddr
@@ -647,9 +647,9 @@ def _terminate(self) -> None:
         if self._client is not None:
             self._client.close()
         if self._listen_task is not None:
-            self._listen_task.cancel()
+            cancel_task_if_running(self._listen_task)
         if self._reader_task is not None:
-            self._reader_task.cancel()
+            cancel_task_if_running(self._reader_task)
 
         self._alive = False
         if self._child is not None and self._child.returncode is None:

hivemind/p2p/p2p_daemon_bindings/control.py

@@ -19,6 +19,7 @@
     write_pbmsg,
 )
 from hivemind.proto import p2pd_pb2 as p2pd_pb
+from hivemind.utils.asyncio import cancel_task_if_running
 from hivemind.utils.logging import get_logger
 from hivemind.utils.multiaddr import Multiaddr, protocols
 
@@ -134,10 +135,8 @@ async def create(
         return control
 
     def close(self) -> None:
-        if self._read_task is not None:
-            self._read_task.cancel()
-        if self._write_task is not None:
-            self._write_task.cancel()
+        cancel_task_if_running(self._read_task)
+        cancel_task_if_running(self._write_task)
 
     def __del__(self):
         self.close()
@@ -194,7 +193,7 @@ async def _read_from_persistent_conn(self, reader: asyncio.StreamReader):
                 self._handler_tasks[call_id] = handler_task
 
             elif call_id in self._handler_tasks and resp.HasField("cancel"):
-                self._handler_tasks[call_id].cancel()
+                cancel_task_if_running(self._handler_tasks[call_id])
 
             elif call_id in self._pending_calls and resp.HasField("daemonError"):
                 daemon_exc = P2PDaemonError(resp.daemonError.message)

hivemind/utils/asyncio.py

@@ -195,3 +195,16 @@ async def enter_asynchronously(context: AbstractContextManager):
     """Wrap a non-async context so that it can be entered asynchronously"""
     async with _AsyncContextWrapper(context) as ret_value:
         yield ret_value
+
+
+def cancel_task_if_running(task: Optional[asyncio.Task]) -> None:
+    """Safely cancel a task if it's still running and the event loop is available."""
+    if task is not None and not task.done():
+        try:
+            loop = asyncio.get_event_loop()
+            if loop.is_running():
+                task.cancel()
+        except RuntimeError as e:
+            # Only ignore event loop closure errors
+            if "Event loop is closed" not in str(e):
+                raise

pyproject.toml

@@ -21,3 +21,6 @@ dummy-variable-rgx = "^_$"
 
 [tool.ruff.lint.isort]
 known-local-folder = ["arguments", "test_utils", "tests", "utils"]
+
+[tool.pytest.ini_options]
+asyncio_default_fixture_loop_scope = "function"

requirements-dev.txt

@@ -1,6 +1,6 @@
 pytest==8.3.5
 pytest-forked
-pytest-asyncio==0.16.0
+pytest-asyncio==0.26.0
 pytest-cov
 pytest-timeout
 coverage

tests/conftest.py

@@ -1,4 +1,3 @@
-import asyncio
 import gc
 
 import psutil
@@ -12,22 +11,6 @@
 logger = get_logger(__name__)
 
 
-@pytest.fixture
-def event_loop():
-    """
-    This overrides the ``event_loop`` fixture from pytest-asyncio
-    (e.g. to make it compatible with ``asyncio.subprocess``).
-
-    This fixture is identical to the original one but does not call ``loop.close()`` in the end.
-    Indeed, at this point, the loop is already stopped (i.e. next tests are free to create new loops).
-    However, finalizers of objects created in the current test may reference the current loop and fail if it is closed.
-    For example, this happens while using ``asyncio.subprocess`` (the ``asyncio.subprocess.Process`` finalizer
-    fails if the loop is closed, but works if the loop is only stopped).
-    """
-
-    yield asyncio.get_event_loop()
-
-
 @pytest.fixture(autouse=True, scope="session")
 def cleanup_children():
     yield

tests/test_connection_handler.py

@@ -5,6 +5,7 @@
 from typing import Any, Dict
 
 import pytest
+import pytest_asyncio
 import torch
 
 from hivemind.compression import deserialize_tensor_stream, deserialize_torch_tensor, serialize_torch_tensor
@@ -20,7 +21,7 @@
 from hivemind.utils.tensor_descr import BatchTensorDescriptor
 
 
-@pytest.fixture
+@pytest_asyncio.fixture
 async def client_stub():
     handler_dht = DHT(start=True)
     module_backends = {"expert1": DummyModuleBackend("expert1", k=1), "expert2": DummyModuleBackend("expert2", k=2)}

tests/test_dht_schema.py

@@ -2,6 +2,7 @@
 from typing import Dict
 
 import pytest
+import pytest_asyncio
 from pydantic.v1 import BaseModel, StrictInt, conint
 
 import hivemind
@@ -17,7 +18,7 @@ class SampleSchema(BaseModel):
     signed_data: Dict[BytesWithPublicKey, bytes]
 
 
-@pytest.fixture
+@pytest_asyncio.fixture
 async def dht_nodes_with_schema():
     validator = SchemaValidator(SampleSchema)
 

tests/test_p2p_daemon_bindings.py

@@ -3,6 +3,7 @@
 from contextlib import AsyncExitStack
 
 import pytest
+import pytest_asyncio
 from google.protobuf.message import EncodeError
 
 from hivemind.p2p.p2p_daemon_bindings.control import ControlClient, DaemonConnector, parse_conn_protocol
@@ -372,7 +373,7 @@ async def test_write_pbmsg_missing_fields(pb_msg):
         await write_pbmsg(MockReaderWriter(), pb_msg)
 
 
-@pytest.fixture
+@pytest_asyncio.fixture
 async def p2pcs():
     # TODO: Change back to gather style
     async with AsyncExitStack() as stack:

tests/test_p2p_servicer.py

@@ -2,13 +2,14 @@
 from typing import AsyncIterator
 
 import pytest
+import pytest_asyncio
 
 from hivemind.p2p import P2P, P2PContext, P2PDaemonError, ServicerBase
 from hivemind.proto import test_pb2
 from hivemind.utils.asyncio import anext
 
 
-@pytest.fixture
+@pytest_asyncio.fixture
 async def server_client():
     server = await P2P.create()
     client = await P2P.create(initial_peers=await server.get_visible_maddrs())