__init__.py 14 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317
  1. """
  2. This is a Distributed Hash Table optimized for rapidly accessing a lot of lightweight metadata.
  3. Hivemind DHT is based on Kademlia [1] with added support for improved bulk store/get operations and caching.
  4. The code is organized as follows:
  5. * **class DHT (__init__.py)** - high-level class for model training. Runs DHTNode in a background process.
  6. * **class DHTNode (node.py)** - an asyncio implementation of dht server, stores AND gets keys.
  7. * **class DHTProtocol (protocol.py)** - an RPC protocol to request data from dht nodes.
  8. * **async def traverse_dht (traverse.py)** - a search algorithm that crawls DHT peers.
  9. - [1] Maymounkov P., Mazieres D. (2002) Kademlia: A Peer-to-Peer Information System Based on the XOR Metric.
  10. - [2] https://github.com/bmuller/kademlia , Brian, if you're reading this: THANK YOU! you're awesome :)
  11. """
  12. from __future__ import annotations
  13. import asyncio
  14. import multiprocessing as mp
  15. import os
  16. from concurrent.futures import ThreadPoolExecutor
  17. from functools import partial
  18. from typing import Awaitable, Callable, Iterable, List, Optional, Sequence, TypeVar, Union
  19. from multiaddr import Multiaddr
  20. from hivemind.dht.node import DEFAULT_NUM_WORKERS, DHTNode
  21. from hivemind.dht.routing import DHTID, DHTKey, DHTValue, Subkey
  22. from hivemind.dht.validation import CompositeValidator, RecordValidatorBase
  23. from hivemind.p2p import P2P, PeerID
  24. from hivemind.utils import DHTExpiration, MPFuture, ValueWithExpiration, await_cancelled, get_logger, switch_to_uvloop
  25. logger = get_logger(__name__)
  26. ReturnType = TypeVar("ReturnType")
  27. class DHT(mp.Process):
  28. """
  29. A high-level interface to a hivemind DHT that runs a single DHT node in a background process.
  30. * hivemind servers periodically announce their experts via declare_experts (dht_handler.py)
  31. * trainers find most suitable experts via RemoteMixtureOfExperts (beam_search.py)
  32. :param initial_peers: multiaddrs of one or more active DHT peers (if you want to join an existing DHT)
  33. :param start: if True, automatically starts the background process on creation. Otherwise await manual start
  34. :param daemon: if True, the background process is marked as daemon and automatically terminated after main process
  35. :param num_workers: declare_experts and get_experts will use up to this many parallel workers
  36. (but no more than one per key)
  37. :param expiration: experts declared from this node expire after this many seconds (default = 5 minutes)
  38. :param record_validators: instances of RecordValidatorBase used for signing and validating stored records.
  39. The validators will be combined using the CompositeValidator class. It merges them when possible
  40. (according to their `.merge_with()` policies) and orders them according to the `.priority` properties.
  41. :param shutdown_timeout: when calling .shutdown, wait for up to this many seconds before terminating
  42. :param await_ready: if True, the constructor waits until the DHT process is ready to process incoming requests
  43. :param kwargs: any other params will be forwarded to DHTNode and hivemind.p2p.P2P upon creation
  44. """
  45. _node: DHTNode
  46. def __init__(
  47. self,
  48. initial_peers: Optional[Sequence[Union[Multiaddr, str]]] = None,
  49. *,
  50. start: bool,
  51. daemon: bool = True,
  52. num_workers: int = DEFAULT_NUM_WORKERS,
  53. record_validators: Iterable[RecordValidatorBase] = (),
  54. shutdown_timeout: float = 3,
  55. await_ready: bool = True,
  56. **kwargs,
  57. ):
  58. self._parent_pid = os.getpid()
  59. super().__init__()
  60. if not (
  61. initial_peers is None
  62. or (
  63. isinstance(initial_peers, Sequence)
  64. and all(isinstance(item, (Multiaddr, str)) for item in initial_peers)
  65. )
  66. ):
  67. raise TypeError("initial_peers should be of type Optional[Sequence[Union[Multiaddr, str]]]")
  68. self.initial_peers = initial_peers
  69. self.kwargs = kwargs
  70. self.num_workers = num_workers
  71. self._record_validator = CompositeValidator(record_validators)
  72. self._inner_pipe, self._outer_pipe = mp.Pipe(duplex=True)
  73. self.shutdown_timeout = shutdown_timeout
  74. self.ready = mp.Event()
  75. self.daemon = daemon
  76. # These values will be fetched from the child process when requested
  77. self._peer_id = None
  78. self._client_mode = None
  79. self._p2p_replica = None
  80. if start:
  81. self.run_in_background(await_ready=await_ready)
  82. def run(self) -> None:
  83. """Serve DHT forever. This function will not return until DHT node is shut down"""
  84. loop = switch_to_uvloop()
  85. with ThreadPoolExecutor(max_workers=1) as pipe_awaiter:
  86. async def _run():
  87. self._node = await DHTNode.create(
  88. initial_peers=self.initial_peers,
  89. num_workers=self.num_workers,
  90. record_validator=self._record_validator,
  91. **self.kwargs,
  92. )
  93. self.ready.set()
  94. while True:
  95. method, args, kwargs = await loop.run_in_executor(pipe_awaiter, self._inner_pipe.recv)
  96. task = asyncio.create_task(getattr(self, method)(*args, **kwargs))
  97. if method == "_shutdown":
  98. await task
  99. break
  100. coro = _run()
  101. loop.run_until_complete(coro)
  102. def run_in_background(self, await_ready=True, timeout=None):
  103. """
  104. Starts DHT in a background process. if await_ready, this method will wait until background dht
  105. is ready to process incoming requests or for :timeout: seconds max.
  106. """
  107. self.start()
  108. if await_ready and not self.ready.wait(timeout=timeout):
  109. raise TimeoutError(f"DHT didn't notify .ready in {timeout} seconds")
  110. def shutdown(self) -> None:
  111. """Shut down a running dht process"""
  112. if self.is_alive():
  113. self._outer_pipe.send(("_shutdown", [], {}))
  114. self.join(self.shutdown_timeout)
  115. if self.is_alive():
  116. logger.warning("DHT did not shut down within the grace period; terminating it the hard way.")
  117. self.terminate()
  118. async def _shutdown(self):
  119. await self._node.shutdown()
  120. def get(
  121. self, key: DHTKey, latest: bool = False, return_future: bool = False, **kwargs
  122. ) -> Union[Optional[ValueWithExpiration[DHTValue]], MPFuture]:
  123. """
  124. Search for a key across DHT and return either first or latest entry (if found).
  125. :param key: same key as in node.store(...)
  126. :param latest: if True, finds the latest value, otherwise finds any non-expired value (which is much faster)
  127. :param return_future: if False (default), return when finished. Otherwise return MPFuture and run in background.
  128. :param kwargs: parameters forwarded to DHTNode.get_many_by_id
  129. :returns: (value, expiration time); if value was not found, returns None
  130. """
  131. future = MPFuture()
  132. self._outer_pipe.send(("_get", [], dict(key=key, latest=latest, future=future, **kwargs)))
  133. return future if return_future else future.result()
  134. async def _get(self, key: DHTKey, latest: bool, future: MPFuture, **kwargs):
  135. try:
  136. result = await self._node.get(key, latest=latest, **kwargs)
  137. if not future.done():
  138. future.set_result(result)
  139. except BaseException as e:
  140. if not future.done():
  141. future.set_exception(e)
  142. raise
  143. def store(
  144. self,
  145. key: DHTKey,
  146. value: DHTValue,
  147. expiration_time: DHTExpiration,
  148. subkey: Optional[Subkey] = None,
  149. return_future: bool = False,
  150. **kwargs,
  151. ) -> Union[bool, MPFuture]:
  152. """
  153. Find num_replicas best nodes to store (key, value) and store it there until expiration time.
  154. :param key: msgpack-serializable key to be associated with value until expiration.
  155. :param value: msgpack-serializable value to be stored under a given key until expiration.
  156. :param expiration_time: absolute time when the entry should expire, based on hivemind.get_dht_time()
  157. :param subkey: if specified, add a value under that subkey instead of overwriting key (see DHTNode.store_many)
  158. :param return_future: if False (default), return when finished. Otherwise return MPFuture and run in background.
  159. :returns: True if store succeeds, False if it fails (due to no response or newer value)
  160. """
  161. future = MPFuture()
  162. self._outer_pipe.send(
  163. (
  164. "_store",
  165. [],
  166. dict(key=key, value=value, expiration_time=expiration_time, subkey=subkey, future=future, **kwargs),
  167. )
  168. )
  169. return future if return_future else future.result()
  170. async def _store(
  171. self,
  172. key: DHTKey,
  173. value: DHTValue,
  174. expiration_time: DHTExpiration,
  175. subkey: Optional[Subkey],
  176. future: MPFuture,
  177. **kwargs,
  178. ):
  179. try:
  180. result = await self._node.store(key, value, expiration_time, subkey=subkey, **kwargs)
  181. if not future.done():
  182. future.set_result(result)
  183. except BaseException as e:
  184. if not future.done():
  185. future.set_exception(e)
  186. raise
  187. def run_coroutine(
  188. self, coro: Callable[[DHT, DHTNode], Awaitable[ReturnType]], return_future: bool = False
  189. ) -> Union[ReturnType, MPFuture[ReturnType]]:
  190. """
  191. Execute an asynchronous function on a DHT participant and return results. This is meant as an interface
  192. for running custom functions DHT for special cases (e.g. declare experts, beam search)
  193. :param coro: async function to be executed. Receives 2 arguments: this DHT daemon and a running DHTNode
  194. :param return_future: if False (default), return when finished. Otherwise return MPFuture and run in background.
  195. :returns: coroutine outputs or MPFuture for these outputs
  196. :note: the coroutine will be executed inside the DHT process. As such, any changes to global variables or
  197. DHT fields made by this coroutine will not be accessible from the host process.
  198. :note: all time-consuming operations in coro should be asynchronous (e.g. asyncio.sleep instead of time.sleep)
  199. or use asyncio.get_event_loop().run_in_executor(...) to prevent coroutine from blocking background DHT tasks
  200. :note: when run_coroutine is called with wait=False, MPFuture can be cancelled to interrupt the task.
  201. """
  202. future = MPFuture()
  203. self._outer_pipe.send(("_run_coroutine", [], dict(coro=coro, future=future)))
  204. return future if return_future else future.result()
  205. async def _run_coroutine(
  206. self, coro: Callable[[DHT, DHTNode], Awaitable[ReturnType]], future: MPFuture[ReturnType]
  207. ):
  208. main_task = asyncio.create_task(coro(self, self._node))
  209. cancel_task = asyncio.create_task(await_cancelled(future))
  210. try:
  211. await asyncio.wait({main_task, cancel_task}, return_when=asyncio.FIRST_COMPLETED)
  212. if future.cancelled():
  213. main_task.cancel()
  214. else:
  215. future.set_result(await main_task)
  216. except BaseException as e:
  217. logger.exception(f"Caught an exception when running a coroutine: {e}")
  218. if not future.done():
  219. future.set_exception(e)
  220. def add_validators(self, record_validators: Iterable[RecordValidatorBase]) -> None:
  221. if not self.ready.is_set():
  222. raise RuntimeError(
  223. "Can't append new validators before the DHT process has started. "
  224. "Consider adding them to the initial list via DHT.__init__(record_validators=...)"
  225. )
  226. self.run_coroutine(partial(DHT._add_validators, record_validators=record_validators))
  227. @staticmethod
  228. async def _add_validators(_dht: DHT, node: DHTNode, record_validators: Iterable[RecordValidatorBase]) -> None:
  229. node.protocol.record_validator.extend(record_validators)
  230. @property
  231. def peer_id(self) -> PeerID:
  232. if self._peer_id is None:
  233. self._peer_id = self.run_coroutine(DHT._get_peer_id)
  234. return self._peer_id
  235. @staticmethod
  236. async def _get_peer_id(_dht: DHT, node: DHTNode) -> PeerID:
  237. return node.peer_id
  238. @property
  239. def client_mode(self) -> bool:
  240. if self._client_mode is None:
  241. self._client_mode = self.run_coroutine(DHT._get_client_mode)
  242. return self._client_mode
  243. @staticmethod
  244. async def _get_client_mode(_dht: DHT, node: DHTNode) -> bool:
  245. return node.protocol.client_mode
  246. def get_visible_maddrs(self, latest: bool = False) -> List[Multiaddr]:
  247. """
  248. Get multiaddrs of the current DHT node that should be accessible by other peers.
  249. :param latest: ask the P2P daemon to refresh the visible multiaddrs
  250. """
  251. return self.run_coroutine(partial(DHT._get_visible_maddrs, latest=latest))
  252. @staticmethod
  253. async def _get_visible_maddrs(_dht: DHT, node: DHTNode, latest: bool = False) -> List[Multiaddr]:
  254. return await node.get_visible_maddrs(latest=latest)
  255. async def replicate_p2p(self) -> P2P:
  256. """
  257. Get a replica of a P2P instance used in the DHT process internally.
  258. The replica uses the same P2P daemon as the DHT and only works while DHT is alive.
  259. """
  260. if self._p2p_replica is None:
  261. daemon_listen_maddr = self.run_coroutine(DHT._get_p2p_daemon_listen_maddr)
  262. self._p2p_replica = await P2P.replicate(daemon_listen_maddr)
  263. return self._p2p_replica
  264. @staticmethod
  265. async def _get_p2p_daemon_listen_maddr(_dht: DHT, node: DHTNode) -> Multiaddr:
  266. return node.p2p.daemon_listen_maddr
  267. def __del__(self):
  268. if self._parent_pid == os.getpid() and self.is_alive():
  269. self.shutdown()