Add graceful shutdown in test_dht_experts

learning-at-home · Jun 9, 2024 · e955347 · e955347
1 parent a98c59d
commit e955347
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/tests/test_dht_experts.py b/tests/test_dht_experts.py
@@ -82,6 +82,10 @@ def test_beam_search(
         assert all(isinstance(e, hivemind.RemoteExpert) for experts in batch_experts for e in experts)
         assert all(len(experts) == beam_size for experts in batch_experts)
 
+    you.shutdown()
+    for dht_instance in dht_instances:
+        dht_instance.shutdown()
+
 
 @pytest.mark.forked
 def test_dht_single_node():
@@ -119,6 +123,8 @@ def test_dht_single_node():
     with pytest.raises(AssertionError):
         beam_search.get_active_successors(["e.1.2.", "e.2", "e.4.5."])
 
+    node.shutdown()
+
 
 def test_uid_patterns():
     valid_experts = [