Convert averager to libp2p backend

hivemind/averaging/allreduce.py

@@ -7,7 +7,7 @@
 from hivemind.averaging.partition import TensorPartContainer, TensorPartReducer, AllreduceException
 from hivemind.p2p import P2P, P2PContext, PeerID as Endpoint, ServicerBase, StubBase
 from hivemind.utils import get_logger
-from hivemind.utils.asyncio import anext, achain, aiter, aenumerate, amap_in_executor
+from hivemind.utils.asyncio import anext, achain, aiter, aenumerate, amap_in_executor, asingle
 from hivemind.utils.compression import serialize_torch_tensor, deserialize_torch_tensor
 from hivemind.proto import averaging_pb2
 
@@ -231,9 +231,9 @@ async def _accumulate_parts_streaming(self, stream: AsyncIterator[averaging_pb2.
             yield averaging_pb2.AveragingData(code=averaging_pb2.AVERAGED_PART, tensor_part=serialized_delta)
 
     async def _send_error_to_peer(self, peer_endpoint: Endpoint, code: averaging_pb2.MessageCode):
+        # In case of reporting the error, we expect the response stream to contain exactly one item
         error = averaging_pb2.AveragingData(group_id=self.group_id, endpoint=self.endpoint.to_base58(), code=code)
-        async for _ in self._get_peer_stub(peer_endpoint).rpc_aggregate_part(aiter(error)):
-            pass
+        await asingle(self._get_peer_stub(peer_endpoint).rpc_aggregate_part(aiter(error)))
 
     def finalize(self, *, cancel: bool = False, exception: Optional[BaseException] = None):
         """finish or terminate AllReduceRunner, propagate any errors / cancellations to peers."""

hivemind/p2p/p2p_daemon.py

@@ -15,7 +15,7 @@
 from hivemind.p2p.p2p_daemon_bindings.datastructures import PeerID, PeerInfo, StreamInfo
 from hivemind.p2p.p2p_daemon_bindings.utils import ControlFailure
 from hivemind.proto.p2pd_pb2 import RPCError
-from hivemind.utils.asyncio import aiter
+from hivemind.utils.asyncio import aiter, asingle
 from hivemind.utils.logging import get_logger
 
 logger = get_logger(__name__)
@@ -401,15 +401,7 @@ async def add_protobuf_handler(
         """
 
         async def _stream_handler(requests: P2P.TInputStream, context: P2PContext) -> P2P.TOutputStream:
-            if stream_input:
-                input = requests
-            else:
-                count = 0
-                async for input in requests:
-                    count += 1
-                if count != 1:
-                    raise ValueError(f"Got {count} requests for handler {name} instead of one")
-
+            input = requests if stream_input else await asingle(requests)
             output = handler(input, context)
 
             if isinstance(output, AsyncIterableABC):
@@ -429,13 +421,7 @@ async def call_protobuf_handler(
     ) -> Awaitable[TOutputProtobuf]:
         requests = input if isinstance(input, AsyncIterableABC) else aiter(input)
         responses = self._iterate_protobuf_stream_handler(peer_id, name, requests, output_protobuf_type)
-
-        count = 0
-        async for response in responses:
-            count += 1
-        if count != 1:
-            raise ValueError(f"Got {count} responses from handler {name} instead of one")
-        return response
+        return await asingle(responses)
 
     def iterate_protobuf_handler(
         self,

hivemind/utils/asyncio.py

@@ -59,6 +59,16 @@ async def aenumerate(aiterable: AsyncIterable[T]) -> AsyncIterable[Tuple[int, T]
         index += 1
 
 
+async def asingle(aiter: AsyncIterable[T]) -> T:
+    """If ``aiter`` has exactly one item, returns this item. Otherwise, raises `ValueError`."""
+    count = 0
+    async for item in aiter:
+        count += 1
+    if count != 1:
+        raise ValueError(f"Iterable has {count} items instead of one (as expected)")
+    return item
+
+
 async def await_cancelled(awaitable: Awaitable) -> bool:
     try:
         await awaitable