Pārlūkot izejas kodu

use conftest to clean up child processes

justheuristic 3 gadi atpakaļ
vecāks
revīzija
b02d09dd31
2 mainītis faili ar 52 papildinājumiem un 0 dzēšanām
  1. 51 0
      tests/conftest.py
  2. 1 0
      tests/test_full_model.py

+ 51 - 0
tests/conftest.py

@@ -0,0 +1,51 @@
+import asyncio
+import gc
+from contextlib import suppress
+
+import psutil
+import pytest
+from hivemind.utils.crypto import RSAPrivateKey
+from hivemind.utils.logging import get_logger, use_hivemind_log_handler
+from hivemind.utils.mpfuture import MPFuture
+
+use_hivemind_log_handler("in_root_logger")
+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
+
+    with RSAPrivateKey._process_wide_key_lock:
+        RSAPrivateKey._process_wide_key = None
+
+    gc.collect()  # Call .__del__() for removed objects
+
+    children = psutil.Process().children(recursive=True)
+    if children:
+        logger.info(f"Cleaning up {len(children)} leftover child processes")
+        for child in children:
+            with suppress(psutil.NoSuchProcess):
+                child.terminate()
+        psutil.wait_procs(children, timeout=1)
+        for child in children:
+            with suppress(psutil.NoSuchProcess):
+                child.kill()
+
+    MPFuture.reset_backend()

+ 1 - 0
tests/test_full_model.py

@@ -52,3 +52,4 @@ def test_full_model_exact_match(atol_forward=1e-3, atol_inference=1e-3):
             del ref_model, ref_outputs, dummy_mask
         else:
             logger.warning("Did not test exact match with local model: REF_NAME environment variable is not set")
+            assert False