def test_dht_single_node():
node = hivemind.DHT(start=True, expiration=999)
beam_search = MoEBeamSearcher(node, 'expert.')
assert all(
node.declare_experts(['expert.1', 'expert.2', 'expert.3'],
f"{hivemind.LOCALHOST}:1337").values())
assert len(node.declare_experts(["ffn.1", "ffn.2"],
endpoint="that_place")) == 4
assert len(
node.declare_experts(['e.1.2.3', 'e.1.2.5', 'e.2.0'],
f"{hivemind.LOCALHOST}:42")) == 7
for expert in node.get_experts(['expert.3', 'expert.2']):
assert expert.endpoint == f"{hivemind.LOCALHOST}:1337"
assert all(
node.declare_experts(['expert.5', 'expert.2'],
f"{hivemind.LOCALHOST}:1337").values())
found_experts = beam_search.find_best_experts(
[(0., 1., 2., 3., 4., 5., 6., 7., 8.)], beam_size=2)
assert len(found_experts) == 2 and [
expert.uid for expert in found_experts
] == ['expert.5', 'expert.3']
successors = beam_search.get_active_successors(
['e.1.2.', 'e.2.', 'e.4.5.'])
assert len(successors['e.1.2.']) == 2
assert successors['e.1.2.'][3] == UidEndpoint('e.1.2.3', f'{LOCALHOST}:42')
assert successors['e.1.2.'][5] == UidEndpoint('e.1.2.5', f'{LOCALHOST}:42')
assert len(
successors['e.2.']) == 1 and successors['e.2.'][0] == UidEndpoint(
'e.2.0', f'{LOCALHOST}:42')
assert successors['e.4.5.'] == {}
initial_beam = beam_search.get_initial_beam((3, 2, 1, 0, -1, -2, -3),
beam_size=3)
assert len(initial_beam) == 3
assert initial_beam[0][:2] == (2.0, 'expert.1.')
assert initial_beam[1][:2] == (1.0, 'expert.2.')
assert initial_beam[2][:2] == (0.0, 'expert.3.')
with pytest.raises(AssertionError):
beam_search = MoEBeamSearcher(node, 'expert.1.ffn')
with pytest.raises(AssertionError):
beam_search.get_active_successors(['e.1.2.', 'e.2', 'e.4.5.'])
async def test_negative_caching():
peers = []
for i in range(10):
neighbors_i = [
f'{LOCALHOST}:{node.port}'
for node in random.sample(peers, min(3, len(peers)))
]
peers.append(
hivemind.DHT(initial_peers=neighbors_i,
cache_locally=False,
start=True))
writer_peer = random.choice(peers)
assert all(
hivemind.declare_experts(writer_peer, ['ffn.1.2.3', 'ffn.3.4.5'],
'myaddr:1234').values())
neighbors_i = [
f'{LOCALHOST}:{node.port}'
for node in random.sample(peers, min(3, len(peers)))
]
neg_caching_peer = hivemind.DHT(initial_peers=neighbors_i,
cache_locally=False,
start=True)
beam_search = MoEBeamSearcher(neg_caching_peer,
uid_prefix='ffn.',
grid_size=(10, 10, 10),
negative_caching=True)
# get prefixes by the peer with negative caching. Cache "no data" entries for ffn.0.*, ffn.2.*, ffn.4.*, ffn.5.*
assert len(
beam_search.get_initial_beam(scores=[.1, .2, .3, .4, .5, .6],
beam_size=3)) == 2
node = await hivemind.DHTNode.create(initial_peers=neighbors_i)
fetched = await asyncio.gather(*(node.get(f'ffn.{i}.') for i in range(10)))
for i in range(6):
assert fetched[i] is not None, f"node should have cached ffn.{i}."
for i in range(6, len(fetched)):
assert fetched[i] is None, f"node shouldn't have cached ffn.{i}."