[object Object]

← back to Exo

make configurable discovery timeout

63a05d5b4ff50efcd95f9914e01bf8f8b9495c0e · 2024-07-23 20:03:31 -0700 · Alex Cheema

Files touched

Diff

commit 63a05d5b4ff50efcd95f9914e01bf8f8b9495c0e
Author: Alex Cheema <alexcheema123@gmail.com>
Date:   Tue Jul 23 20:03:31 2024 -0700

    make configurable discovery timeout
---
 exo/networking/grpc/grpc_discovery.py | 6 +++---
 main.py                               | 1 +
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/exo/networking/grpc/grpc_discovery.py b/exo/networking/grpc/grpc_discovery.py
index 3064dff4..274a5105 100644
--- a/exo/networking/grpc/grpc_discovery.py
+++ b/exo/networking/grpc/grpc_discovery.py
@@ -23,7 +23,7 @@ class ListenProtocol(asyncio.DatagramProtocol):
 
 
 class GRPCDiscovery(Discovery):
-    def __init__(self, node_id: str, node_port: int, listen_port: int, broadcast_port: int = None, broadcast_interval: int = 1, device_capabilities: DeviceCapabilities = UNKNOWN_DEVICE_CAPABILITIES):
+    def __init__(self, node_id: str, node_port: int, listen_port: int, broadcast_port: int = None, broadcast_interval: int = 1, device_capabilities: DeviceCapabilities = UNKNOWN_DEVICE_CAPABILITIES, discovery_timeout: int = 30):
         self.node_id = node_id
         self.node_port = node_port
         self.device_capabilities = device_capabilities
@@ -34,6 +34,7 @@ class GRPCDiscovery(Discovery):
         self.broadcast_task = None
         self.listen_task = None
         self.cleanup_task = None
+        self.discovery_timeout = discovery_timeout
 
     async def start(self):
         self.device_capabilities = device_capabilities()
@@ -121,10 +122,9 @@ class GRPCDiscovery(Discovery):
         while True:
             try:
                 current_time = time.time()
-                timeout = 15 * self.broadcast_interval
                 peers_to_remove = [
                     peer_handle.id() for peer_handle, connected_at, last_seen in self.known_peers.values() if
-                    (not await peer_handle.is_connected() and current_time - connected_at > timeout) or current_time - last_seen > timeout
+                    (not await peer_handle.is_connected() and current_time - connected_at > self.discovery_timeout) or current_time - last_seen > self.discovery_timeout
                 ]
                 if DEBUG_DISCOVERY >= 2: print("Peer statuses:", {peer_handle.id(): f"is_connected={await peer_handle.is_connected()}, {connected_at=}, {last_seen=}" for peer_handle, connected_at, last_seen in self.known_peers.values()})
                 if DEBUG_DISCOVERY >= 2 and len(peers_to_remove) > 0: print(f"Cleaning up peers: {peers_to_remove}")
diff --git a/main.py b/main.py
index d0ae3f52..2cab3de4 100644
--- a/main.py
+++ b/main.py
@@ -18,6 +18,7 @@ parser.add_argument("--node-port", type=int, default=None, help="Node port")
 parser.add_argument("--listen-port", type=int, default=5678, help="Listening port for discovery")
 parser.add_argument("--prometheus-client-port", type=int, default=None, help="Prometheus client port")
 parser.add_argument("--broadcast-port", type=int, default=5678, help="Broadcast port for discovery")
+parser.add_argument("--discovery-timeout", type=int, default=30, help="Discovery timeout in seconds")
 parser.add_argument("--wait-for-peers", type=int, default=0, help="Number of peers to wait to connect to before starting")
 parser.add_argument("--chatgpt-api-port", type=int, default=8000, help="ChatGPT API port")
 parser.add_argument("--chatgpt-api-response-timeout-secs", type=int, default=90, help="ChatGPT API response timeout in seconds")

← 8d2bb819 add llama-3.1 notice to README  ·  back to Exo  ·  make configurable discovery timeout 9a373c2b →