mirror of
https://github.com/bigscience-workshop/petals
synced 2024-11-19 21:25:38 +00:00
58 lines
1.9 KiB
Python
58 lines
1.9 KiB
Python
import threading
|
|
import time
|
|
|
|
import pytest
|
|
import torch
|
|
from hivemind import DHT, get_logger, use_hivemind_log_handler
|
|
from test_utils import *
|
|
|
|
from petals.client import RemoteSequenceManager, RemoteSequential
|
|
from petals.client.remote_model import DistributedBloomConfig
|
|
from petals.data_structures import UID_DELIMITER
|
|
|
|
use_hivemind_log_handler("in_root_logger")
|
|
logger = get_logger(__file__)
|
|
|
|
|
|
@pytest.mark.forked
|
|
def test_sequence_manager_shutdown():
|
|
config = DistributedBloomConfig.from_pretrained(MODEL_NAME, initial_peers=INITIAL_PEERS)
|
|
dht = DHT(initial_peers=config.initial_peers, client_mode=True, start=True)
|
|
sequential = RemoteSequential(config, dht)
|
|
shutdown_evt = threading.Event()
|
|
|
|
# test RemoteSequential with lossy compression
|
|
block_uids = [f"{config.dht_prefix}{UID_DELIMITER}{i}" for i in range(config.n_layer)]
|
|
sequential = RemoteSequential(
|
|
config,
|
|
dht,
|
|
sequence_manager=TestSequenceManager(dht, block_uids, sequential.p2p, _was_shut_down=shutdown_evt, start=True),
|
|
)
|
|
|
|
sequence = sequential.sequence_manager.make_sequence()
|
|
assert all(sequence[i].peer_id != sequence[i + 1].peer_id for i in range(len(sequence) - 1))
|
|
|
|
assert sequential.sequence_manager.is_alive()
|
|
assert sequential.sequence_manager._thread.ready.is_set()
|
|
assert not shutdown_evt.is_set()
|
|
sequential(torch.randn(1, 2, config.hidden_size))
|
|
|
|
sequential.sequence_manager.shutdown()
|
|
del sequential
|
|
time.sleep(1)
|
|
|
|
assert shutdown_evt.is_set()
|
|
|
|
|
|
class TestSequenceManager(RemoteSequenceManager):
|
|
"""A sequence manager that signals if it was shut down"""
|
|
|
|
def __init__(self, *args, _was_shut_down: threading.Event, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self._was_shut_down = _was_shut_down
|
|
|
|
def shutdown(self):
|
|
super().shutdown()
|
|
assert not self.is_alive()
|
|
self._was_shut_down.set()
|