custom timeout

hivemind/averaging/averager.py

@@ -109,6 +109,7 @@ class DecentralizedAverager(mp.Process, ServicerBase):
         averaging_expiration: float = None,
         min_matchmaking_time: float = 5.0,
         request_timeout: float = 3.0,
+        load_state_message_timeout: float = 10.0,
         averaging_alpha: float = 1.0,
         part_size_bytes: int = DEFAULT_PART_SIZE_BYTES,
         allreduce_timeout: Optional[float] = None,
@@ -164,6 +165,7 @@ class DecentralizedAverager(mp.Process, ServicerBase):
             tensor.share_memory_()
         self.total_size = sum(map(torch.Tensor.numel, self._averaged_tensors))
         self.schema_hash = compute_schema_hash(self._averaged_tensors)
+        self.load_state_message_timeout = load_state_message_timeout
         self.shutdown_timeout = shutdown_timeout
         self.bandwidth = bandwidth
 
@@ -643,7 +645,7 @@ class DecentralizedAverager(mp.Process, ServicerBase):
                         stream = await stub.rpc_download_state(averaging_pb2.DownloadRequest())
                         current_tensor_parts, tensors = [], []
 
-                        async for message in aiter_with_timeout(stream, timeout=self.request_timeout):
+                        async for message in aiter_with_timeout(stream, timeout=self.load_state_message_timeout):
                             if message.metadata:
                                 metadata = self.serializer.loads(message.metadata)
                             if message.tensor_part.dtype and current_tensor_parts: