1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556 |
- 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()
|