← back to Exo
make all I/O ops in manual_discovery.py run inside a ThreadPoolExecutor
b066c944f3994f2a06ae43aba24755c1215af664 · 2024-11-15 10:20:42 +0700 · Ian Paul
Files touched
M exo/networking/manual/manual_discovery.py
Diff
commit b066c944f3994f2a06ae43aba24755c1215af664
Author: Ian Paul <iandouglaspaul@gmail.com>
Date: Fri Nov 15 10:20:42 2024 +0700
make all I/O ops in manual_discovery.py run inside a ThreadPoolExecutor
---
exo/networking/manual/manual_discovery.py | 60 ++++++++++++++++++++-----------
1 file changed, 40 insertions(+), 20 deletions(-)
diff --git a/exo/networking/manual/manual_discovery.py b/exo/networking/manual/manual_discovery.py
index e9351821..4b98514c 100644
--- a/exo/networking/manual/manual_discovery.py
+++ b/exo/networking/manual/manual_discovery.py
@@ -1,8 +1,9 @@
import os
import asyncio
-from exo.networking.discovery import Discovery
from typing import Dict, List, Callable, Optional
+from concurrent.futures import ThreadPoolExecutor
+from exo.networking.discovery import Discovery
from exo.topology.device_capabilities import DeviceCapabilities
from exo.networking.manual.network_topology_config import NetworkTopology, PeerConfig
from exo.helpers import DEBUG_DISCOVERY
@@ -26,6 +27,8 @@ class ManualDiscovery(Discovery):
self._cached_peers: Dict[str, PeerConfig] = {}
self._last_modified_time: Optional[float] = None
+ self._file_executor = ThreadPoolExecutor(max_workers=1)
+ self._lock = asyncio.Lock()
async def start(self) -> None:
self.listen_task = asyncio.create_task(self.task_find_peers_from_config())
@@ -34,6 +37,7 @@ class ManualDiscovery(Discovery):
async def stop(self) -> None:
if self.listen_task: self.listen_task.cancel()
if self.cleanup_task: self.cleanup_task.cancel()
+ self._file_executor.shutdown(wait=True)
async def discover_peers(self, wait_for_peers: int = 0) -> List[PeerHandle]:
if wait_for_peers > 0:
@@ -46,7 +50,8 @@ class ManualDiscovery(Discovery):
async def task_find_peers_from_config(self):
if DEBUG_DISCOVERY >= 2: print("Starting task to find peers from config...")
while True:
- for peer_id, peer_config in self._get_peers().items():
+ peers_from_config = await self._get_peers()
+ for peer_id, peer_config in peers_from_config.items():
try:
if DEBUG_DISCOVERY >= 2: print(f"Checking peer {peer_id=} at {peer_config.address}:{peer_config.port}")
peer = self.known_peers.get(peer_id)
@@ -72,7 +77,7 @@ class ManualDiscovery(Discovery):
async def task_clean_up_peers_from_config(self):
if DEBUG_DISCOVERY >= 2: print("Starting task to clean up peers from config...")
while True:
- peers_from_config = self._get_peers()
+ peers_from_config = await self._get_peers()
if peers_from_config:
peers_to_remove = [peer for peer in self.known_peers.keys() if peer not in peers_from_config]
@@ -83,30 +88,45 @@ class ManualDiscovery(Discovery):
await asyncio.sleep(1.0)
- def _get_peers(self):
+ async def _get_peers(self):
try:
- current_mtime = os.path.getmtime(self.network_config_path)
-
- if self._cached_peers is not None and self._last_modified_time is not None and current_mtime <= self._last_modified_time:
- return self._cached_peers
+ async with self._lock:
+ loop = asyncio.get_running_loop()
+ current_mtime = await loop.run_in_executor(
+ self._file_executor,
+ os.path.getmtime,
+ self.network_config_path
+ )
- topology = NetworkTopology.from_path(self.network_config_path)
+ if (self._cached_peers is not None and
+ self._last_modified_time is not None and
+ current_mtime <= self._last_modified_time):
+ return self._cached_peers
- if self.node_id not in topology.peers:
- raise ValueError(
- f"Node ID {self.node_id} not found in network config file "
- f"{self.network_config_path}. Please run with `node_id` set to "
- f"one of the keys in the config file: {[k for k, _ in topology.peers]}"
+ topology = await loop.run_in_executor(
+ self._file_executor,
+ NetworkTopology.from_path,
+ self.network_config_path
)
- peers_in_network: Dict[str, PeerConfig] = topology.peers
- peers_in_network.pop(self.node_id)
+ if self.node_id not in topology.peers:
+ raise ValueError(
+ f"Node ID {self.node_id} not found in network config file "
+ f"{self.network_config_path}. Please run with `node_id` set to "
+ f"one of the keys in the config file: {[k for k, _ in topology.peers]}"
+ )
+
+ peers_in_network: Dict[str, PeerConfig] = topology.peers
+ peers_in_network.pop(self.node_id)
- self._cached_peers = peers_in_network
- self._last_modified_time = current_mtime
+ self._cached_peers = peers_in_network
+ self._last_modified_time = current_mtime
- return peers_in_network
+ return peers_in_network
except Exception as e:
- if DEBUG_DISCOVERY >= 2: print(f"Error when loading network config file from {self.network_config_path}. Please update the config file in order to successfully discover peers. Exception: {e}")
+ if DEBUG_DISCOVERY >= 2:
+ print(f"Error when loading network config file from {self.network_config_path}. "
+ f"Please update the config file in order to successfully discover peers. "
+ f"Exception: {e}")
return self._cached_peers
← 0e34ce21 patch after rebasing to main
·
back to Exo
·
make popping from dict threadsafe 18acb97b →