You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
57 lines
1.9 KiB
Python
57 lines
1.9 KiB
Python
import threading
|
|
import time
|
|
|
|
import pytest
|
|
import torch
|
|
from hivemind import DHT, get_logger
|
|
|
|
from petals import AutoDistributedConfig
|
|
from petals.client import RemoteSequenceManager, RemoteSequential
|
|
from petals.data_structures import UID_DELIMITER
|
|
from test_utils import *
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
@pytest.mark.forked
|
|
@pytest.mark.parametrize("mode", ["max_throughput", "min_latency"])
|
|
def test_sequence_manager_basics(mode: str):
|
|
config = AutoDistributedConfig.from_pretrained(MODEL_NAME, initial_peers=INITIAL_PEERS)
|
|
dht = DHT(initial_peers=config.initial_peers, client_mode=True, start=True)
|
|
sequential = RemoteSequential(config, dht=dht)
|
|
shutdown_evt = threading.Event()
|
|
|
|
# test RemoteSequential with lossy compression
|
|
block_uids = [f"{config.dht_prefix}{UID_DELIMITER}{i}" for i in range(config.num_hidden_layers)]
|
|
sequential = RemoteSequential(
|
|
config,
|
|
sequence_manager=RemoteSequenceManagerWithChecks(config, block_uids, dht=dht, _was_shut_down=shutdown_evt),
|
|
)
|
|
|
|
sequence = sequential.sequence_manager.make_sequence(mode=mode)
|
|
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 RemoteSequenceManagerWithChecks(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()
|