mirror of
https://github.com/bigscience-workshop/petals
synced 2024-10-31 09:20:41 +00:00
f0c7383181
- finish renaming RemoteSequenceInfo -> RemoteSequenceManager (why: if it was an *Info, user would expect it to be similar - to a dataclass; whereas in actuality, the class is doing heavy network interactions on its own) - implement RemoteSequenceManager.make_sequence (from https://pastebin.com/uXgy2U8B ) - make RemoteSequentialInferenceSession use RemoteSequenceManager.make_sequence - make tests pass again - make it possible to create inference session without RemoteTransformerBlock - make a standalone test for RemoteSequential - rollback convert-model Co-authored-by: Tim Dettmers <tim.dettmers@gmail.com>
52 lines
1.7 KiB
Python
52 lines
1.7 KiB
Python
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()
|