Follow up #501 and #511 with minor fixes (#513)

- In `hivemind.Server`, use the graceful shutdown for `ConnectionHandler`
- In `hivemind.P2P`, if we are the first peer, skip checking if the provided identity is free

(cherry picked from commit 13cdd135284d1571240013d0697210c03a965cd7)

hivemind/moe/server/server.py

@@ -247,10 +247,8 @@ class Server(threading.Thread):
         if self.checkpoint_saver is not None:
             self.checkpoint_saver.start()
 
-        for process in self.conn_handlers:
-            if not process.is_alive():
-                process.start()
-            process.ready.result()
+        for handler in self.conn_handlers:
+            handler.run_in_background()
 
         try:
             self.runtime.run()
@@ -287,9 +285,8 @@ class Server(threading.Thread):
         """
         self.ready.clear()
 
-        for process in self.conn_handlers:
-            process.terminate()
-            process.join()
+        for handler in self.conn_handlers:
+            handler.shutdown()
         logger.debug("Connection handlers terminated")
 
         if self.module_backends:
@@ -301,11 +298,10 @@ class Server(threading.Thread):
             self.checkpoint_saver.join()
 
         self.dht.shutdown()
-        self.dht.join()
 
         logger.debug(f"Shutting down runtime")
-
         self.runtime.shutdown()
+
         logger.info("Server shutdown succesfully")
 
 

hivemind/p2p/p2p_daemon.py

@@ -131,7 +131,8 @@ class P2P:
         :param quic: Deprecated, has no effect since libp2p 0.17.0
         :param use_relay_hop: Deprecated, has no effect since libp2p 0.17.0
         :param use_relay_discovery: Deprecated, has no effect since libp2p 0.17.0
-        :param check_if_identity_free: If enabled (default) and ``identity_path`` is provided,
+        :param check_if_identity_free: If enabled (default), ``identity_path`` is provided,
+                                       and we are connecting to an existing swarm,
                                        ensure that this identity is not used by other peers already.
                                        This slows down ``P2P.create()`` but protects from unintuitive libp2p errors
                                        appearing in case of the identity collision.
@@ -176,7 +177,7 @@ class P2P:
 
         if identity_path is not None:
             if os.path.isfile(identity_path):
-                if check_if_identity_free:
+                if check_if_identity_free and need_bootstrap:
                     logger.info(f"Checking that identity from `{identity_path}` is not used by other peers")
                     if await cls.is_identity_taken(
                         identity_path,

tests/test_start_server.py

@@ -52,9 +52,6 @@ def test_cli_run_server_identity_path():
             encoding="utf-8",
         )
 
-        line = server_2_proc.stderr.readline()
-        assert re.search(r"Checking that identity.+is not used by other peers", line) is not None
-
         line = server_2_proc.stderr.readline()
         addrs_pattern_result = re.search(pattern, line)
         assert addrs_pattern_result is not None, line