Convert DHT to libp2p backend

benchmarks/benchmark_averaging.py

@@ -34,7 +34,9 @@ def sample_tensors(hid_size, num_layers):
 def benchmark_averaging(num_peers: int, target_group_size: int, num_rounds: int,
                         averaging_expiration: float, request_timeout: float, round_timeout: float,
                         hid_size: int, num_layers: int, spawn_dtime: float):
-    dht_root = hivemind.DHT(listen_on=f'{LOCALHOST}:*', start=True)
+    dht_root = hivemind.DHT(start=True)
+    initial_peers = dht_root.get_visible_maddrs()
+
     num_groups = 2 ** int(round(math.log2(num_peers / target_group_size)))
     nbits = int(round(math.log2(num_groups)))
     peer_tensors = [sample_tensors(hid_size, num_layers)
@@ -45,9 +47,7 @@ def benchmark_averaging(num_peers: int, target_group_size: int, num_rounds: int,
 
     def run_averager(index):
         nonlocal successful_steps, total_steps, lock_stats
-        dht = hivemind.DHT(listen_on=f'{LOCALHOST}:*',
-                           initial_peers=[f"{LOCALHOST}:{dht_root.port}"],
-                           start=True)
+        dht = hivemind.DHT(initial_peers=initial_peers, start=True)
         initial_bits = bin(index % num_groups)[2:].rjust(nbits, '0')
         averager = hivemind.averaging.DecentralizedAverager(
             peer_tensors[i], dht, prefix='my_tensor', initial_group_bits=initial_bits, listen_on=f"{LOCALHOST}:*",

benchmarks/benchmark_dht.py

@@ -23,9 +23,9 @@ def benchmark_dht(num_peers: int, initial_peers: int, num_experts: int, expert_b
     logger.info("Creating peers...")
     peers = []
     for _ in trange(num_peers):
-        neighbors = [f'0.0.0.0:{node.port}' for node in random.sample(peers, min(initial_peers, len(peers)))]
-        peer = hivemind.DHT(initial_peers=neighbors, start=True, wait_timeout=wait_timeout,
-                            listen_on=f'0.0.0.0:*')
+        neighbors = sum([peer.get_visible_maddrs()
+                         for peer in random.sample(peers, min(initial_peers, len(peers)))], [])
+        peer = hivemind.DHT(initial_peers=neighbors, start=True, wait_timeout=wait_timeout)
         peers.append(peer)
 
     store_peer, get_peer = peers[-2:]

examples/albert/arguments.py

@@ -1,5 +1,5 @@
-from typing import Optional, List
 from dataclasses import dataclass, field
+from typing import Optional, List
 
 from transformers import TrainingArguments
 
@@ -11,11 +11,26 @@ class BaseTrainingArguments:
     )
     initial_peers: List[str] = field(
         default_factory=list,
-        metadata={"help": "One or more peers (comma-separated) that will welcome you into the collaboration"}
+        metadata={"help":
+            "Multiaddrs of the peers that will welcome you into the existing collaboration. "
+            "Example: /ip4/203.0.113.1/tcp/31337/p2p/XXXX /ip4/203.0.113.2/udp/7777/quic/p2p/YYYY"}
     )
-    dht_listen_on: str = field(
-        default="[::]:*",
-        metadata={"help": "Network interface used for incoming DHT communication. Default: all ipv6"}
+    use_ipfs: bool = field(
+        default=False,
+        metadata={"help":
+            "Use IPFS to find initial_peers. If enabled, you only need to provide /p2p/XXXX part of the multiaddrs "
+            "for the initial_peers (no need to specify a particular IPv4/IPv6 host and port)"}
+    )
+    host_maddrs: List[str] = field(
+        default_factory=lambda: ['/ip4/0.0.0.0/tcp/0', '/ip4/0.0.0.0/udp/0/quic'],
+        metadata={"help":
+            "Multiaddrs to listen for external connections from other p2p instances. "
+            "Defaults to all IPv4 interfaces with TCP and QUIC (over UDP) protocols: "
+            "/ip4/0.0.0.0/tcp/0 /ip4/0.0.0.0/udp/0/quic"}
+    )
+    announce_maddrs: List[str] = field(
+        default_factory=list,
+        metadata={"help": "Visible multiaddrs the host announces for external connections from other p2p instances"}
     )
 
 
@@ -97,10 +112,6 @@ class CollaborationArguments(AveragerArguments, CollaborativeOptimizerArguments,
         default=600,
         metadata={"help": "Statistics will be removed if not updated in this many seconds"}
     )
-    endpoint: Optional[str] = field(
-        default=None,
-        metadata={"help": "This node's IP for inbound connections, used when running from behind a proxy"}
-    )
 
 
 @dataclass

examples/albert/run_trainer.py

@@ -18,8 +18,8 @@
 from torch_optimizer import Lamb
 
 import hivemind
+import utils
 from arguments import CollaborationArguments, DatasetArguments, AlbertTrainingArguments
-import metrics_utils
 
 
 logger = logging.getLogger(__name__)
@@ -130,7 +130,7 @@ def on_step_end(self, args: TrainingArguments, state: transformers.TrainerState,
                 self.last_reported_collaboration_step = self.collaborative_optimizer.local_step
                 self.total_samples_processed += self.samples
                 samples_per_second = self.collaborative_optimizer.performance_ema.samples_per_second
-                statistics = metrics_utils.LocalMetrics(
+                statistics = utils.LocalMetrics(
                     step=self.collaborative_optimizer.local_step,
                     samples_per_second=samples_per_second,
                     samples_accumulated=self.samples,
@@ -219,13 +219,16 @@ def main():
 
     opt, scheduler = get_optimizer_and_scheduler(training_args, model)
 
-    validators, local_public_key = metrics_utils.make_validators(
+    validators, local_public_key = utils.make_validators(
         collaboration_args_dict['experiment_prefix'])
-    dht = hivemind.DHT(
-        start=True, initial_peers=collaboration_args_dict.pop('initial_peers'),
-        listen=not collaboration_args_dict['client_mode'],
-        listen_on=collaboration_args_dict.pop('dht_listen_on'),
-        endpoint=collaboration_args_dict.pop('endpoint'), record_validators=validators)
+    dht = hivemind.DHT(start=True,
+                       initial_peers=collaboration_args_dict.pop('initial_peers'),
+                       listen=not collaboration_args_dict['client_mode'],
+                       record_validators=validators,
+                       use_ipfs=collaboration_args_dict.pop('use_ipfs'),
+                       host_maddrs=collaboration_args_dict.pop('host_maddrs'),
+                       announce_maddrs=collaboration_args_dict.pop('announce_maddrs'))
+    utils.log_visible_maddrs(dht.get_visible_maddrs())
 
     total_batch_size_per_step = training_args.per_device_train_batch_size * training_args.gradient_accumulation_steps
     statistics_expiration = collaboration_args_dict.pop('statistics_expiration')

examples/albert/run_training_monitor.py

@@ -1,21 +1,22 @@
 #!/usr/bin/env python
 
-from dataclasses import dataclass, field, asdict
 import subprocess
 import time
+from dataclasses import asdict, dataclass, field
+from ipaddress import ip_address
 from typing import Optional
 
 import torch
+import wandb
 from torch_optimizer import Lamb
 from transformers import AlbertForPreTraining, AlbertConfig, HfArgumentParser
-import wandb
-from whatsmyip.providers import GoogleDnsProvider
 from whatsmyip.ip import get_ip
+from whatsmyip.providers import GoogleDnsProvider
 
-from arguments import BaseTrainingArguments, CollaborativeOptimizerArguments, AveragerArguments
 import hivemind
+import utils
+from arguments import BaseTrainingArguments, CollaborativeOptimizerArguments, AveragerArguments
 from hivemind.utils.logging import get_logger
-import metrics_utils
 
 logger = get_logger(__name__)
 
@@ -27,10 +28,9 @@ class CoordinatorArguments(BaseTrainingArguments):
     new workers still can join the collaboration via alive initial peers' addresses.
     Specify initial_peers argument for that purpose
     """
-    address: Optional[str] = field(
-        default=None,
-        metadata={"help": "This machine's network address. Use public IP for global experiments, "
-                          "local address for private runs"}
+    use_google_dns: bool = field(
+        default=False,
+        metadata={"help": "Use Google DNS to determine our public IP address (and add it to --announce_maddrs)"}
     )
     refresh_period: float = field(
         default=30,
@@ -139,17 +139,21 @@ def upload_checkpoint(self, current_loss):
     parser = HfArgumentParser((CoordinatorArguments, CollaborativeOptimizerArguments, AveragerArguments))
     coordinator_args, collab_optimizer_args, averager_args = parser.parse_args_into_dataclasses()
 
-    if coordinator_args.address is None:
-        logger.warning("No address specified. Attempting to infer address from DNS.")
-        coordinator_args.address = get_ip(GoogleDnsProvider)
+    if coordinator_args.use_google_dns:
+        address = get_ip(GoogleDnsProvider)
+        logger.info(f"Google DNS responds that our IP address is {address}")
+        version = ip_address(address).version
+        coordinator_args.announce_maddrs += [f'/ip{version}/{address}/tcp/0', f'/ip{version}/{address}/udp/0/quic']
 
     experiment_prefix = coordinator_args.experiment_prefix
-    validators, local_public_key = metrics_utils.make_validators(experiment_prefix)
-    dht = hivemind.DHT(start=True, listen_on=coordinator_args.dht_listen_on,
-                       endpoint=f"{coordinator_args.address}:*", initial_peers=coordinator_args.initial_peers,
-                       record_validators=validators)
-
-    logger.info(f"Running DHT root at {coordinator_args.address}:{dht.port}")
+    validators, local_public_key = utils.make_validators(experiment_prefix)
+    dht = hivemind.DHT(start=True,
+                       initial_peers=coordinator_args.initial_peers,
+                       record_validators=validators
+                       use_ipfs=coordinator_args.use_ipfs,
+                       host_maddrs=coordinator_args.host_maddrs,
+                       announce_maddrs=coordinator_args.announce_maddrs)
+    utils.log_visible_maddrs(dht.get_visible_maddrs())
 
     if coordinator_args.wandb_project is not None:
         wandb.init(project=coordinator_args.wandb_project)
@@ -162,7 +166,7 @@ def upload_checkpoint(self, current_loss):
         metrics_dict = dht.get(experiment_prefix + '_metrics', latest=True)
         if metrics_dict is not None:
             metrics_dict = metrics_dict.value
-            metrics = [metrics_utils.LocalMetrics.parse_obj(metrics_dict[peer].value)
+            metrics = [utils.LocalMetrics.parse_obj(metrics_dict[peer].value)
                        for peer in metrics_dict]
             latest_step = max(item.step for item in metrics)
             if latest_step != current_step:
@@ -184,6 +188,7 @@ def upload_checkpoint(self, current_loss):
                     num_samples += item.samples_accumulated
                     sum_mini_steps += item.mini_steps
                 current_loss = sum_loss / sum_mini_steps
+                logger.info(f"Step #{current_step}\tloss = {current_loss:.5f}")
 
                 if coordinator_args.wandb_project is not None:
                     wandb.log({
@@ -198,6 +203,5 @@ def upload_checkpoint(self, current_loss):
                         checkpoint_handler.save_state(current_step)
                         if checkpoint_handler.is_time_to_upload():
                             checkpoint_handler.upload_checkpoint(current_loss)
-                    logger.info(f"Step #{current_step}\tloss = {current_loss:.5f}")
         logger.debug("Peer is still alive...")
         time.sleep(coordinator_args.refresh_period)

examples/albert/metrics_utils.py → examples/albert/utils.py

@@ -1,9 +1,15 @@
 from typing import Dict, List, Tuple
 
+from multiaddr import Multiaddr
+from pydantic import BaseModel, StrictFloat, confloat, conint
+
 from hivemind.dht.crypto import RSASignatureValidator
 from hivemind.dht.schema import BytesWithPublicKey, SchemaValidator
 from hivemind.dht.validation import RecordValidatorBase
-from pydantic import BaseModel, StrictFloat, confloat, conint
+from hivemind.utils.logging import get_logger
+
+
+logger = get_logger(__name__)
 
 
 class LocalMetrics(BaseModel):
@@ -23,3 +29,15 @@ def make_validators(experiment_prefix: str) -> Tuple[List[RecordValidatorBase],
     validators = [SchemaValidator(MetricSchema, prefix=experiment_prefix),
                   signature_validator]
     return validators, signature_validator.local_public_key
+
+
+class TextStyle:
+    BOLD = '\033[1m'
+    BLUE = '\033[34m'
+    RESET = '\033[0m'
+
+
+def log_visible_maddrs(visible_maddrs: List[Multiaddr]) -> None:
+    initial_peers_str = ' '.join(str(addr) for addr in visible_maddrs)
+    logger.info(f"Running a DHT node. To connect, supply "
+                f"{TextStyle.BOLD}{TextStyle.BLUE}--initial_peers {initial_peers_str}{TextStyle.RESET}")

hivemind/averaging/averager.py

@@ -12,6 +12,7 @@
 import weakref
 from concurrent.futures.thread import ThreadPoolExecutor
 from dataclasses import asdict
+from ipaddress import ip_address
 from typing import Sequence, Optional, Tuple, Any, Union, Dict, AsyncIterator
 
 import grpc
@@ -30,6 +31,7 @@
 from hivemind.utils.asyncio import anext, achain, aiter, switch_to_uvloop
 from hivemind.utils.compression import serialize_torch_tensor, deserialize_torch_tensor
 from hivemind.utils.grpc import ChannelCache, GRPC_KEEPALIVE_OPTIONS, split_for_streaming, combine_from_streaming
+from hivemind.utils.networking import choose_ip_address, strip_port
 from hivemind.utils.serializer import MSGPackSerializer, SerializerBase
 from hivemind.utils.timed_storage import get_dht_time, ValueWithExpiration, DHTExpiration
 
@@ -68,6 +70,8 @@ class DecentralizedAverager(mp.Process, averaging_pb2_grpc.DecentralizedAveragin
     :param listen: if True (default), this averager will accept incoming requests from other peers and perform allreduce
             if False, the averager will register as a freeloader and attempt to fetch vectors from other averagers
     :param listen_on: network interface, e.g. "0.0.0.0:1337" or "localhost:*" (* means pick any port) or "[::]:7654"
+    :param announced_host: visible IP address the averager will announce for external connections from other peers.
+          If None, the address will be chosen from p2p.get_visible_maddrs() (global IPv4 addresses are preferred)
     :param channel_options: options for grpc.aio.insecure_channel, e.g. [('grpc.enable_retries', 0)]
           see https://grpc.github.io/grpc/core/group__grpc__arg__keys.html for a list of all options
     :param kwargs: extra parameters forwarded to grpc.aio.server
@@ -102,7 +106,8 @@ def __init__(self, averaged_tensors: Sequence[torch.Tensor], dht: DHT, *, start:
                  throughput: Optional[float] = None, min_vector_size: int = 0,
                  auxiliary: bool = False, allow_state_sharing: Optional[bool] = None,
                  listen: bool = True, listen_on: Endpoint = '0.0.0.0:*', daemon: bool = True,
-                 channel_options: Optional[Sequence[Tuple[str, Any]]] = None,
+                 announced_host: Optional[str] = None,
+                 channel_options: Sequence[Tuple[str, Any]] = (),
                  shutdown_timeout: float = 5, **kwargs):
         assert '.' not in prefix, "group prefix must be a string without trailing '.'"
         assert throughput is None or (throughput >= 0 and np.isfinite(np.float32(throughput))), \
@@ -122,6 +127,9 @@ def __init__(self, averaged_tensors: Sequence[torch.Tensor], dht: DHT, *, start:
         else:
             self.mode = AveragingMode.NODE
 
+        if announced_host is None:
+            announced_host = self._choose_announced_host()
+        self.announced_host = announced_host
         self.channel_options = channel_options
         self.daemon = daemon
 
@@ -163,6 +171,17 @@ def __init__(self, averaged_tensors: Sequence[torch.Tensor], dht: DHT, *, start:
         if start:
             self.run_in_background(await_ready=True)
 
+    def _choose_announced_host(self) -> Hostname:
+        announced_host = strip_port(self.listen_on).strip('[]')  # Stripping square brackets for IPv6
+        if ip_address(announced_host) not in [ip_address('0.0.0.0'), ip_address('::')]:
+            return announced_host
+
+        maddrs = self.dht.get_visible_maddrs()
+        announced_host = choose_ip_address(maddrs)
+        logger.info(f'Choosing IP {announced_host} as endpoint for DecentralizedAverager '
+                    f'from visible multiaddrs {maddrs}')
+        return announced_host
+
     @property
     def port(self) -> Optional[Port]:
         return self._port.value if self._port.value != 0 else None
@@ -183,7 +202,7 @@ def allow_state_sharing(self, value: bool):
     def endpoint(self) -> Optional[Endpoint]:
         if self.listen and self._averager_endpoint is None:
             assert self.port is not None, "Averager is not running yet"
-            self._averager_endpoint = f"{self.dht.get_visible_address()}:{self.port}"
+            self._averager_endpoint = f"{self.announced_host}:{self.port}"
             logger.debug(f"Assuming averager endpoint to be {self._averager_endpoint}")
         return self._averager_endpoint
 
@@ -499,7 +518,8 @@ async def _load_state_from_peers(self, future: MPFuture):
                     logger.info(f"Downloading parameters from peer {peer}")
                     stream = None
                     try:
-                        stub = ChannelCache.get_stub(peer, averaging_pb2_grpc.DecentralizedAveragingStub, aio=True)
+                        stub = ChannelCache.get_stub(peer, averaging_pb2_grpc.DecentralizedAveragingStub, aio=True,
+                                                     options=self.channel_options)
                         stream = stub.rpc_download_state(averaging_pb2.DownloadRequest())
                         current_tensor_parts, tensors = [], []
                         async for message in stream: