simplify manual discovery

exo/main.py

@@ -85,7 +85,7 @@ elif args.discovery_module == "tailscale":
 elif args.discovery_module == "manual":
   if not args.discovery_config_path:
     raise ValueError(f"--discovery-config-path is required when using manual discovery. Please provide a path to a config json file.")
-  discovery = ManualDiscovery(args.discovery_config_path, args.node_id, create_peer_handle=lambda peer_id, address, device_capabilities: GRPCPeerHandle(peer_id, address, device_capabilities), discovery_timeout=args.discovery_timeout)
+  discovery = ManualDiscovery(args.discovery_config_path, args.node_id, create_peer_handle=lambda peer_id, address, device_capabilities: GRPCPeerHandle(peer_id, address, device_capabilities))
 topology_viz = TopologyViz(chatgpt_api_endpoints=chatgpt_api_endpoints, web_chat_urls=web_chat_urls) if not args.disable_tui else None
 node = StandardNode(
   args.node_id,

exo/networking/manual/manual_discovery.py

@@ -14,27 +14,18 @@ class ManualDiscovery(Discovery):
     network_config_path: str,
     node_id: str,
     create_peer_handle: Callable[[str, str, DeviceCapabilities], PeerHandle],
-    discovery_timeout: int = 30,
   ):
     self.topology = NetworkTopology.from_path(network_config_path)
-    self.node_id = node_id
     self.create_peer_handle = create_peer_handle
-    self.discovery_timeout = discovery_timeout
 
-    try:
-      self.node = self.topology.peers[node_id]
-    except KeyError as e:
-      print(f"Node ID {node_id} not found in network config file {network_config_path}. Please run with `node_id` set to one of the keys in the config file: {[k for k, _ in self.topology.peers]}")
-      raise e
-
-    self.node_port = self.node.port
+    if node_id not in self.topology.peers:
+      raise ValueError(f"Node ID {node_id} not found in network config file {network_config_path}. Please run with `node_id` set to one of the keys in the config file: {[k for k, _ in self.topology.peers]}")
 
     self.listen_task = None
-    self.cleanup_task = None
 
     self.known_peers: Dict[str, PeerHandle] = {}
     self.peers_in_network: Dict[str, PeerConfig] = self.topology.peers
-    self.node_config = self.peers_in_network.pop(node_id)
+    self.peers_in_network.pop(node_id)
 
   async def start(self) -> None:
     self.listen_task = asyncio.create_task(self.task_find_peers_from_config())
@@ -44,11 +35,11 @@ class ManualDiscovery(Discovery):
       self.listen_task.cancel()
 
   async def discover_peers(self, wait_for_peers: int = 0) -> List[PeerHandle]:
-    if DEBUG_DISCOVERY >= 2: print("Starting discovery...")
     if wait_for_peers > 0:
       while len(self.known_peers) < wait_for_peers:
         if DEBUG_DISCOVERY >= 2: print(f"Current peers: {len(self.known_peers)}/{wait_for_peers}. Waiting for more peers...")
         await asyncio.sleep(0.1)
+    if DEBUG_DISCOVERY >= 2: print(f"Discovered peers: {[peer.id() for peer in self.known_peers.values()]}")
     return list(self.known_peers.values())
 
 
@@ -68,14 +59,11 @@ class ManualDiscovery(Discovery):
             self.known_peers[peer_id] = peer
           else:
             if DEBUG_DISCOVERY >= 2: print(f"{peer_id=} at {peer_config.address}:{peer_config.port} is not healthy.")
-            try:
-              del self.known_peers[peer_id]
-            except KeyError:
-              pass  # peer was never added, so nothing to delete
+            try: del self.known_peers[peer_id]
+            except KeyError: pass
         except Exception as e:
-            if DEBUG_DISCOVERY >=2: print(f"Exception occured when attempting to add {peer_id=}: {e}")
-        finally:
-          await asyncio.sleep(self.discovery_timeout)
+            if DEBUG_DISCOVERY >= 2: print(f"Exception occured when attempting to add {peer_id=}: {e}")
+      await asyncio.sleep(1.0)
 
       if DEBUG_DISCOVERY >= 2: print(f"Current known peers: {[peer.id() for peer in self.known_peers.values()]}")