Yield from reachability checks (#1427)

## Motivation

check_reachable waits for all connection profile checks to be completed.
Since there are retries on failures, this can take around 20s to
resolve, preventing any instances from showing up. This feels very slow
for UX, and it slows down distributed testing.

## Changes

Made check_reachable an async generator.

## Test Plan

### Manual Testing
Works for me at least.

src/exo/utils/info_gatherer/net_profile.py

@@ -1,4 +1,5 @@
-from collections.abc import Mapping
+from collections import defaultdict
+from collections.abc import AsyncGenerator, Mapping
 
 import anyio
 import httpx
@@ -8,6 +9,7 @@ from loguru import logger
 from exo.shared.topology import Topology
 from exo.shared.types.common import NodeId
 from exo.shared.types.profiling import NodeNetworkInfo
+from exo.utils.channels import Sender, channel
 
 REACHABILITY_ATTEMPTS = 3
 
@@ -80,10 +82,10 @@ async def check_reachable(
     topology: Topology,
     self_node_id: NodeId,
     node_network: Mapping[NodeId, NodeNetworkInfo],
-) -> dict[NodeId, set[str]]:
-    """Check which nodes are reachable and return their IPs."""
+) -> AsyncGenerator[tuple[str, NodeId], None]:
+    """Yield (ip, node_id) pairs as reachability probes complete."""
 
-    reachable: dict[NodeId, set[str]] = {}
+    send, recv = channel[tuple[str, NodeId]]()
 
     # these are intentionally httpx's defaults so we can tune them later
     timeout = httpx.Timeout(timeout=5.0)
@@ -93,6 +95,18 @@ async def check_reachable(
         keepalive_expiry=5,
     )
 
+    async def _probe(
+        target_ip: str,
+        expected_node_id: NodeId,
+        client: httpx.AsyncClient,
+        send: Sender[tuple[str, NodeId]],
+    ) -> None:
+        async with send:
+            out: defaultdict[NodeId, set[str]] = defaultdict(set)
+            await check_reachability(target_ip, expected_node_id, out, client)
+            if expected_node_id in out:
+                await send.send((target_ip, expected_node_id))
+
     async with (
         httpx.AsyncClient(timeout=timeout, limits=limits) as client,
         create_task_group() as tg,
@@ -103,12 +117,9 @@ async def check_reachable(
             if node_id == self_node_id:
                 continue
             for iface in node_network[node_id].interfaces:
-                tg.start_soon(
-                    check_reachability,
-                    iface.ip_address,
-                    node_id,
-                    reachable,
-                    client,
-                )
+                tg.start_soon(_probe, iface.ip_address, node_id, client, send.clone())
+        send.close()
 
-    return reachable
+        with recv:
+            async for item in recv:
+                yield item

src/exo/worker/main.py

@@ -1,3 +1,4 @@
+from collections import defaultdict
 from datetime import datetime, timezone
 from random import random
 from typing import Iterator
@@ -345,29 +346,29 @@ class Worker:
             edges = set(
                 conn.edge for conn in self.state.topology.out_edges(self.node_id)
             )
-            conns = await check_reachable(
+            conns: defaultdict[NodeId, set[str]] = defaultdict(set)
+            async for ip, nid in check_reachable(
                 self.state.topology,
                 self.node_id,
                 self.state.node_network,
-            )
-            for nid in conns:
-                for ip in conns[nid]:
-                    edge = SocketConnection(
-                        # nonsense multiaddr
-                        sink_multiaddr=Multiaddr(address=f"/ip4/{ip}/tcp/52415")
-                        if "." in ip
-                        # nonsense multiaddr
-                        else Multiaddr(address=f"/ip6/{ip}/tcp/52415"),
-                    )
-                    if edge not in edges:
-                        logger.debug(f"ping discovered {edge=}")
-                        await self.event_sender.send(
-                            TopologyEdgeCreated(
-                                conn=Connection(
-                                    source=self.node_id, sink=nid, edge=edge
-                                )
-                            )
+            ):
+                if ip in conns[nid]:
+                    continue
+                conns[nid].add(ip)
+                edge = SocketConnection(
+                    # nonsense multiaddr
+                    sink_multiaddr=Multiaddr(address=f"/ip4/{ip}/tcp/52415")
+                    if "." in ip
+                    # nonsense multiaddr
+                    else Multiaddr(address=f"/ip6/{ip}/tcp/52415"),
+                )
+                if edge not in edges:
+                    logger.debug(f"ping discovered {edge=}")
+                    await self.event_sender.send(
+                        TopologyEdgeCreated(
+                            conn=Connection(source=self.node_id, sink=nid, edge=edge)
                         )
+                    )
 
             for conn in self.state.topology.out_edges(self.node_id):
                 if not isinstance(conn.edge, SocketConnection):
@@ -377,8 +378,7 @@ class Worker:
                     continue
                 if (
                     conn.sink not in conns
-                    or conn.edge.sink_multiaddr.ip_address
-                    not in conns.get(conn.sink, set())
+                    or conn.edge.sink_multiaddr.ip_address not in conns[conn.sink]
                 ):
                     logger.debug(f"ping failed to discover {conn=}")
                     await self.event_sender.send(TopologyEdgeDeleted(conn=conn))