[object Object]

← back to Exo

Yield from reachability checks (#1427)

2204f651c8f7d59061f9e3eb26ef2e9b1feabcd3 · 2026-02-10 12:18:45 +0000 · rltakashige

## 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.

Files touched

Diff

commit 2204f651c8f7d59061f9e3eb26ef2e9b1feabcd3
Author: rltakashige <rl.takashige@gmail.com>
Date:   Tue Feb 10 12:18:45 2026 +0000

    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 | 37 +++++++++++++++++---------
 src/exo/worker/main.py                     | 42 +++++++++++++++---------------
 2 files changed, 45 insertions(+), 34 deletions(-)

diff --git a/src/exo/utils/info_gatherer/net_profile.py b/src/exo/utils/info_gatherer/net_profile.py
index 1c8fa01b..85986ead 100644
--- a/src/exo/utils/info_gatherer/net_profile.py
+++ b/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,
-                )
-
-    return reachable
+                tg.start_soon(_probe, iface.ip_address, node_id, client, send.clone())
+        send.close()
+
+        with recv:
+            async for item in recv:
+                yield item
diff --git a/src/exo/worker/main.py b/src/exo/worker/main.py
index 35834360..4f059bfb 100644
--- a/src/exo/worker/main.py
+++ b/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))

← 4abdaaf7 Address GPU timeouts (#1429)  ·  back to Exo  ·  Add error handling to info gatherer monitor loops (#1422) 163cf183 →