Skip to content

Commit

Permalink
Add downloading time to averager (#618)
Browse files Browse the repository at this point in the history
(cherry picked from commit 14ff472)
  • Loading branch information
samsja authored and mryab committed Oct 13, 2024
1 parent 1e7e31a commit 17f61c5
Showing 1 changed file with 4 additions and 1 deletion.
5 changes: 4 additions & 1 deletion hivemind/averaging/averager.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
import random
import signal
import threading
import time
import weakref
from dataclasses import asdict
from typing import Any, AsyncIterator, Dict, Optional, Sequence, Tuple, Union
Expand Down Expand Up @@ -700,6 +701,7 @@ async def _load_state_from_peers(self, future: MPFuture, timeout: Optional[float
metadata = None
for peer in sorted(peer_priority.keys(), key=peer_priority.get, reverse=True):
if peer != self.peer_id:
t0 = time.monotonic()
logger.info(f"Downloading parameters from peer {peer}")
try:
stub = self.get_stub(self._p2p, peer, namespace=self.prefix)
Expand All @@ -722,7 +724,8 @@ async def _load_state_from_peers(self, future: MPFuture, timeout: Optional[float
logger.debug(f"Peer {peer} did not send its state")
continue

logger.info(f"Finished downloading state from {peer}")
t1 = time.monotonic()
logger.info(f"Finished downloading state in {t1 - t0:.3f}s from {peer}")
future.set_result((metadata, tensors))
return
except Exception as e:
Expand Down

0 comments on commit 17f61c5

Please sign in to comment.