Support libp2p relays for NAT traversal

setup.cfg

@@ -37,7 +37,7 @@ install_requires =
     huggingface-hub==0.11.1
     transformers==4.25.1
     speedtest-cli==2.1.3
-    hivemind==1.1.3
+    hivemind==1.1.5
     tensor_parallel==1.0.23
     humanfriendly
     async-timeout>=4.0.2

src/petals/cli/run_server.py

@@ -127,7 +127,10 @@ def main():
     parser.add_argument("--mean_balance_check_period", type=float, default=60,
                         help="Check the swarm's balance every N seconds (and rebalance it if necessary)")
 
+    parser.add_argument("--auto_relay", action='store_true', help="Enabling relay for NAT traversal")
+    parser.add_argument('--no-auto_relay', dest='auto_relay', action='store_false')
     parser.add_argument("--use_auth_token", action='store_true', help="auth token for from_pretrained")
+
     parser.add_argument('--load_in_8bit', type=str, default=None,
                         help="Convert the loaded transformer blocks into mixed-8bit quantized model. "
                              "Default: True if GPU is available. Use `--load_in_8bit False` to disable this")
@@ -140,7 +143,7 @@ def main():
                         help="Skip checking this server's reachability via health.petals.ml "
                              "when connecting to the public swarm. If you connect to a private swarm, "
                              "the check is skipped by default. Use this option only if you know what you are doing")
-
+    parser.set_defaults(auto_relay=True)
     # fmt:on
     args = vars(parser.parse_args())
     args.pop("config", None)
@@ -158,6 +161,8 @@ def main():
 
     announce_maddrs = args.pop("announce_maddrs")
     public_ip = args.pop("public_ip")
+    use_auto_relay = args.pop("auto_relay")
+
     if public_ip is not None:
         assert announce_maddrs is None, "You can't use --public_ip and --announce_maddrs at the same time"
         assert port != 0, "Please specify a fixed non-zero --port when you use --public_ip (e.g., --port 31337)"
@@ -197,6 +202,7 @@ def main():
         compression=compression,
         max_disk_space=max_disk_space,
         attn_cache_size=attn_cache_size,
+        use_auto_relay=use_auto_relay,
     )
     try:
         server.run()

src/petals/client/remote_model.py

@@ -107,6 +107,8 @@ def __init__(self, config: DistributedBloomConfig):
                 num_workers=n_layer,
                 startup_timeout=config.daemon_startup_timeout,
                 start=True,
+                use_relay=True,
+                use_auto_relay=True,
             )
         )
         assert isinstance(dht, hivemind.DHT) and dht.is_alive(), "dht must be a running hivemind.DHT instance"

src/petals/server/server.py

@@ -78,6 +78,8 @@ def __init__(
         load_in_8bit: Optional[bool] = None,
         tensor_parallel_devices: Optional[Sequence[torch.device]] = None,
         skip_reachability_check: bool = False,
+        use_relay: bool = True,
+        use_auto_relay: bool = True,
         **kwargs,
     ):
         """Create a server with one or more bloom blocks. See run_server.py for documentation."""
@@ -117,7 +119,14 @@ def __init__(
         )
         self.module_uids = [f"{self.prefix}.{block_index}" for block_index in range(self.block_config.n_layer)]
 
-        self.dht = DHT(initial_peers=initial_peers, start=True, num_workers=self.block_config.n_layer, **kwargs)
+        self.dht = DHT(
+            initial_peers=initial_peers,
+            start=True,
+            num_workers=self.block_config.n_layer,
+            use_relay=use_relay,
+            use_auto_relay=use_auto_relay,
+            **kwargs,
+        )
         visible_maddrs_str = [str(a) for a in self.dht.get_visible_maddrs()]
         if initial_peers == PUBLIC_INITIAL_PEERS:
             logger.info(f"Connecting to the public swarm, peer_id = {self.dht.peer_id}")