remove mdns discovered peers from appearing in state

## Motivation

Running RDMA from source is not well documented as is. Several
surprising things that took time to debug internally too.

App should be updated to detect MacOS versions in future.

README.md

@@ -107,6 +107,10 @@ uv run exo
 
 This starts the exo dashboard and API at http://localhost:52415/
 
+
+*Please view the section on RDMA to enable this feature on MacOS >=26.2!*
+
+
 ### Run from Source (Linux)
 
 **Prerequisites:**
@@ -230,7 +234,7 @@ This removes:
 
 RDMA is a new capability added to macOS 26.2. It works on any Mac with Thunderbolt 5 (M4 Pro Mac Mini, M4 Max Mac Studio, M4 Max MacBook Pro, M3 Ultra Mac Studio).
 
-Note that on Mac Studio, you cannot use the Thunderbolt 5 port next to the Ethernet port.
+Please refer to the caveats for immediate troubleshooting.
 
 To enable RDMA on macOS, follow these steps:
 
@@ -247,6 +251,14 @@ To enable RDMA on macOS, follow these steps:
 
 After that, RDMA will be enabled in macOS and exo will take care of the rest.
 
+**Important Caveats**
+
+1. Devices that wish to be part of an RDMA cluster must be connected to all other devices in the cluster.
+2. The cables must support TB5.
+3. On a Mac Studio, you cannot use the Thunderbolt 5 port next to the Ethernet port.
+4. If running from source, please use the script found at `tmp/set_rdma_network_config.sh`, which will disable Thunderbolt Bridge and set dhcp on each RDMA port.
+5. RDMA ports may be unable to discover each other on different versions of MacOS. Please ensure that OS versions match exactly (even beta version numbers) on all devices.
+
 ---
 
 ### Using the API

pyproject.toml

@@ -6,6 +6,8 @@ readme = "README.md"
 requires-python = ">=3.13"
 dependencies = [
     "aiofiles>=24.1.0",
+    "aiohttp>=3.12.14",
+    "types-aiofiles>=24.1.0.20250708",
     "pydantic>=2.11.7",
     "fastapi>=0.116.1",
     "filelock>=3.18.0",

src/exo/download/coordinator.py

@@ -121,7 +121,6 @@ class DownloadCoordinator:
     def _start_download_task(
         self, shard: ShardMetadata, initial_progress: RepoDownloadProgress
     ) -> None:
-        logger.warning("starting download for {shard}")
         model_id = shard.model_card.model_id
 
         # Emit ongoing status

src/exo/download/download_utils.py

@@ -8,13 +8,13 @@ import traceback
 from collections.abc import Awaitable
 from datetime import timedelta
 from pathlib import Path
-from typing import Callable, Literal, cast
+from typing import Callable, Literal
 from urllib.parse import urljoin
 
 import aiofiles
 import aiofiles.os as aios
+import aiohttp
 import certifi
-import httpx
 from huggingface_hub import (
     snapshot_download,  # pyright: ignore[reportUnknownVariableType]
 )
@@ -176,7 +176,7 @@ async def fetch_file_list_with_cache(
         # Fetch failed - try cache fallback
         if await aios.path.exists(cache_file):
             logger.warning(
-                f"{type(e).__name__}: Failed to fetch file list for {model_id}, using cached data"
+                f"Failed to fetch file list for {model_id}, using cached data: {e}"
             )
             async with aiofiles.open(cache_file, "r") as f:
                 return TypeAdapter(list[FileListEntry]).validate_json(await f.read())
@@ -196,7 +196,7 @@ async def fetch_file_list_with_retry(
         except Exception as e:
             if attempt == n_attempts - 1:
                 raise e
-            await asyncio.sleep(min(16, 0.5 * float(2.0 ** int(attempt))))
+            await asyncio.sleep(min(8, 0.1 * float(2.0 ** int(attempt))))
     raise Exception(
         f"Failed to fetch file list for {model_id=} {revision=} {path=} {recursive=}"
     )
@@ -211,25 +211,26 @@ async def _fetch_file_list(
     headers = await get_download_headers()
     async with (
         create_http_session(timeout_profile="short") as session,
+        session.get(url, headers=headers) as response,
     ):
-        response = await session.get(url, headers=headers)
-        if response.status_code in [401, 403]:
-            msg = await _build_auth_error_message(response.status_code, model_id)
+        if response.status in [401, 403]:
+            msg = await _build_auth_error_message(response.status, model_id)
             raise HuggingFaceAuthenticationError(msg)
-        if response.status_code != 200:
-            raise Exception(f"Failed to fetch file list: {response.status_code}")
-
-        data = TypeAdapter(list[FileListEntry]).validate_json(response.text)
-        files: list[FileListEntry] = []
-        for item in data:
-            if item.type == "file":
-                files.append(FileListEntry.model_validate(item))
-            elif item.type == "directory" and recursive:
-                subfiles = await _fetch_file_list(
-                    model_id, revision, item.path, recursive
-                )
-                files.extend(subfiles)
-        return files
+        if response.status == 200:
+            data_json = await response.text()
+            data = TypeAdapter(list[FileListEntry]).validate_json(data_json)
+            files: list[FileListEntry] = []
+            for item in data:
+                if item.type == "file":
+                    files.append(FileListEntry.model_validate(item))
+                elif item.type == "directory" and recursive:
+                    subfiles = await _fetch_file_list(
+                        model_id, revision, item.path, recursive
+                    )
+                    files.extend(subfiles)
+            return files
+        else:
+            raise Exception(f"Failed to fetch file list: {response.status}")
 
 
 async def get_download_headers() -> dict[str, str]:
@@ -237,29 +238,34 @@ async def get_download_headers() -> dict[str, str]:
 
 
 def create_http_session(
+    auto_decompress: bool = False,
     timeout_profile: Literal["short", "long"] = "long",
-) -> httpx.AsyncClient:
+) -> aiohttp.ClientSession:
     if timeout_profile == "short":
         total_timeout = 30
         connect_timeout = 10
-        read_timeout = 30
+        sock_read_timeout = 30
+        sock_connect_timeout = 10
     else:
         total_timeout = 1800
         connect_timeout = 60
-        read_timeout = 1800
+        sock_read_timeout = 1800
+        sock_connect_timeout = 60
 
     ssl_context = ssl.create_default_context(
         cafile=os.getenv("SSL_CERT_FILE") or certifi.where()
     )
+    connector = aiohttp.TCPConnector(ssl=ssl_context)
 
-    # default here is to load env vars
-    return httpx.AsyncClient(
-        verify=ssl_context,
-        timeout=httpx.Timeout(
+    return aiohttp.ClientSession(
+        auto_decompress=auto_decompress,
+        connector=connector,
+        proxy=os.getenv("HTTPS_PROXY") or os.getenv("HTTP_PROXY") or None,
+        timeout=aiohttp.ClientTimeout(
+            total=total_timeout,
             connect=connect_timeout,
-            read=read_timeout,
-            write=total_timeout,
-            pool=total_timeout,
+            sock_read=sock_read_timeout,
+            sock_connect=sock_connect_timeout,
         ),
     )
 
@@ -286,28 +292,26 @@ async def file_meta(
     headers = await get_download_headers()
     async with (
         create_http_session(timeout_profile="short") as session,
-        session.stream("HEAD", url, headers=headers) as r,
+        session.head(url, headers=headers) as r,
     ):
-        if r.status_code == 307:
+        if r.status == 307:
             # On redirect, only trust Hugging Face's x-linked-* headers.
-            x_linked_size = cast(str | None, r.headers.get("x-linked-size"))
-            x_linked_etag = cast(str | None, r.headers.get("x-linked-etag"))
+            x_linked_size = r.headers.get("x-linked-size")
+            x_linked_etag = r.headers.get("x-linked-etag")
             if x_linked_size and x_linked_etag:
                 content_length = int(x_linked_size)
                 etag = trim_etag(x_linked_etag)
                 return content_length, etag
             # Otherwise, follow the redirect to get authoritative size/hash
-            redirected_location = cast(str | None, r.headers.get("location"))
+            redirected_location = r.headers.get("location")
             return await file_meta(model_id, revision, path, redirected_location)
-        if r.status_code in [401, 403]:
-            msg = await _build_auth_error_message(r.status_code, model_id)
+        if r.status in [401, 403]:
+            msg = await _build_auth_error_message(r.status, model_id)
             raise HuggingFaceAuthenticationError(msg)
-        content_length = cast(
-            str | None,
-            r.headers.get("x-linked-size") or r.headers.get("content-length"),
+        content_length = int(
+            r.headers.get("x-linked-size") or r.headers.get("content-length") or 0
         )
-        content_length = 0 if content_length is None else int(content_length)
-        etag = cast(str | None, r.headers.get("x-linked-etag") or r.headers.get("etag"))
+        etag = r.headers.get("x-linked-etag") or r.headers.get("etag")
         assert content_length > 0, f"No content length for {url}"
         assert etag is not None, f"No remote hash for {url}"
         etag = trim_etag(etag)
@@ -336,7 +340,7 @@ async def download_file_with_retry(
                 f"Download error on attempt {attempt}/{n_attempts} for {model_id=} {revision=} {path=} {target_dir=}"
             )
             logger.error(traceback.format_exc())
-            await asyncio.sleep(min(16, 0.5 * (2.0**attempt)))
+            await asyncio.sleep(min(8, 0.1 * (2.0**attempt)))
     raise Exception(
         f"Failed to download file {model_id=} {revision=} {path=} {target_dir=}"
     )
@@ -349,7 +353,6 @@ async def _download_file(
     target_dir: Path,
     on_progress: Callable[[int, int, bool], None] = lambda _, __, ___: None,
 ) -> Path:
-    logger.warning(f"downloading {path} from {model_id} to {target_dir}")
     target_path = target_dir / path
 
     if await aios.path.exists(target_path):
@@ -389,20 +392,20 @@ async def _download_file(
         n_read = resume_byte_pos or 0
         async with (
             create_http_session(timeout_profile="long") as session,
-            session.stream("GET", url, headers=headers, follow_redirects=True) as r,
+            session.get(url, headers=headers) as r,
         ):
-            if r.status_code == 404:
+            if r.status == 404:
                 raise FileNotFoundError(f"File not found: {url}")
-            if r.status_code in [401, 403]:
-                msg = await _build_auth_error_message(r.status_code, model_id)
+            if r.status in [401, 403]:
+                msg = await _build_auth_error_message(r.status, model_id)
                 raise HuggingFaceAuthenticationError(msg)
-            assert r.status_code in [200, 206], (
-                f"Failed to download {path} from {url}: {r.status_code}"
+            assert r.status in [200, 206], (
+                f"Failed to download {path} from {url}: {r.status}"
             )
             async with aiofiles.open(
                 partial_path, "ab" if resume_byte_pos else "wb"
             ) as f:
-                async for chunk in r.aiter_bytes(8 * 1024 * 1024):
+                while chunk := await r.content.read(8 * 1024 * 1024):
                     n_read = n_read + (await f.write(chunk))
                     on_progress(n_read, length, False)
 

src/exo/download/impl_shard_downloader.py

@@ -168,8 +168,7 @@ class ResumableShardDownloader(ShardDownloader):
                 yield await task
             # TODO: except Exception
             except Exception as e:
-                task.cancel()
-                logger.opt(exception=e).error("Error downloading shard")
+                logger.error("Error downloading shard:", e)
 
     async def get_shard_download_status_for_shard(
         self, shard: ShardMetadata

src/exo/main.py

@@ -90,7 +90,6 @@ class Node:
             worker = Worker(
                 node_id,
                 session_id,
-                connection_message_receiver=router.receiver(topics.CONNECTION_MESSAGES),
                 global_event_receiver=router.receiver(topics.GLOBAL_EVENTS),
                 local_event_sender=router.sender(topics.LOCAL_EVENTS),
                 command_sender=router.sender(topics.COMMANDS),
@@ -227,9 +226,6 @@ class Node:
                         self.worker = Worker(
                             self.node_id,
                             result.session_id,
-                            connection_message_receiver=self.router.receiver(
-                                topics.CONNECTION_MESSAGES
-                            ),
                             global_event_receiver=self.router.receiver(
                                 topics.GLOBAL_EVENTS
                             ),

src/exo/utils/keyed_backoff.py

@@ -1,8 +1,10 @@
 import time
-from collections.abc import Hashable
+from typing import Generic, TypeVar
+
+K = TypeVar("K")
 
 
-class KeyedBackoff[K: Hashable]:
+class KeyedBackoff(Generic[K]):
     """Tracks exponential backoff state per key."""
 
     def __init__(self, base: float = 0.5, cap: float = 10.0):

src/exo/worker/main.py

@@ -7,7 +7,6 @@ from anyio import CancelScope, create_task_group, fail_after
 from anyio.abc import TaskGroup
 from loguru import logger
 
-from exo.routing.connection_message import ConnectionMessage, ConnectionMessageType
 from exo.shared.apply import apply
 from exo.shared.models.model_cards import ModelId
 from exo.shared.types.api import ImageEditsInternalParams
@@ -57,7 +56,6 @@ class Worker:
         node_id: NodeId,
         session_id: SessionId,
         *,
-        connection_message_receiver: Receiver[ConnectionMessage],
         global_event_receiver: Receiver[ForwarderEvent],
         local_event_sender: Sender[ForwarderEvent],
         # This is for requesting updates. It doesn't need to be a general command sender right now,
@@ -74,7 +72,6 @@ class Worker:
         self.event_index_counter = event_index_counter
         self.command_sender = command_sender
         self.download_command_sender = download_command_sender
-        self.connection_message_receiver = connection_message_receiver
         self.event_buffer = OrderedBuffer[Event]()
         self.out_for_delivery: dict[EventId, ForwarderEvent] = {}
 
@@ -105,7 +102,6 @@ class Worker:
             tg.start_soon(info_gatherer.run)
             tg.start_soon(self._forward_info, info_recv)
             tg.start_soon(self.plan_step)
-            tg.start_soon(self._connection_message_event_writer)
             tg.start_soon(self._resend_out_for_delivery)
             tg.start_soon(self._event_applier)
             tg.start_soon(self._forward_events)
@@ -279,41 +275,6 @@ class Worker:
         instance = self.state.instances[task.instance_id]
         return instance.shard_assignments.node_to_runner[self.node_id]
 
-    async def _connection_message_event_writer(self):
-        with self.connection_message_receiver as connection_messages:
-            async for msg in connection_messages:
-                await self.event_sender.send(
-                    self._convert_connection_message_to_event(msg)
-                )
-
-    def _convert_connection_message_to_event(self, msg: ConnectionMessage):
-        match msg.connection_type:
-            case ConnectionMessageType.Connected:
-                return TopologyEdgeCreated(
-                    conn=Connection(
-                        source=self.node_id,
-                        sink=msg.node_id,
-                        edge=SocketConnection(
-                            sink_multiaddr=Multiaddr(
-                                address=f"/ip4/{msg.remote_ipv4}/tcp/{msg.remote_tcp_port}"
-                            ),
-                        ),
-                    ),
-                )
-
-            case ConnectionMessageType.Disconnected:
-                return TopologyEdgeDeleted(
-                    conn=Connection(
-                        source=self.node_id,
-                        sink=msg.node_id,
-                        edge=SocketConnection(
-                            sink_multiaddr=Multiaddr(
-                                address=f"/ip4/{msg.remote_ipv4}/tcp/{msg.remote_tcp_port}"
-                            ),
-                        ),
-                    ),
-                )
-
     async def _nack_request(self, since_idx: int) -> None:
         # We request all events after (and including) the missing index.
         # This function is started whenever we receive an event that is out of sequence.

tmp/set_rdma_network_config.sh

@@ -0,0 +1,49 @@
+#!/usr/bin/env bash
+
+set -euo pipefail
+
+PREFS="/Library/Preferences/SystemConfiguration/preferences.plist"
+
+# Remove bridge0 interface
+ifconfig bridge0 &>/dev/null && {
+  ifconfig bridge0 | grep -q 'member' && {
+    ifconfig bridge0 | awk '/member/ {print $2}' | xargs -n1 ifconfig bridge0 deletem 2>/dev/null || true
+  }
+  ifconfig bridge0 destroy 2>/dev/null || true
+}
+
+# Remove Thunderbolt Bridge from VirtualNetworkInterfaces in preferences.plist
+/usr/libexec/PlistBuddy -c "Delete :VirtualNetworkInterfaces:Bridge:bridge0" "$PREFS" 2>/dev/null || true
+
+networksetup -listlocations | grep -q exo || {
+  networksetup -createlocation exo
+}
+
+networksetup -switchtolocation exo
+networksetup -listallhardwareports \
+  | awk -F': ' '/Hardware Port: / {print $2}' \
+  | while IFS=":" read -r name; do
+      case "$name" in
+        "Ethernet Adapter"*)
+                ;;
+        "Thunderbolt Bridge")
+                ;;
+        "Thunderbolt "*)
+          networksetup -listallnetworkservices \
+            | grep -q "EXO $name" \
+              || networksetup -createnetworkservice "EXO $name" "$name" 2>/dev/null \
+              || continue
+          networksetup -setdhcp "EXO $name"
+                ;;
+        *)
+          networksetup -listallnetworkservices \
+            | grep -q "$name" \
+              || networksetup -createnetworkservice "$name" "$name" 2>/dev/null \
+              || continue
+                ;;
+      esac
+    done
+
+networksetup -listnetworkservices | grep -q "Thunderbolt Bridge" && {
+  networksetup -setnetworkserviceenabled "Thunderbolt Bridge" off
+} || true

uv.lock

@@ -366,6 +366,7 @@ version = "0.3.0"
 source = { editable = "." }
 dependencies = [
     { name = "aiofiles", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
+    { name = "aiohttp", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "anyio", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "exo-pyo3-bindings", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "fastapi", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -386,6 +387,7 @@ dependencies = [
     { name = "rustworkx", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "tiktoken", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "tomlkit", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
+    { name = "types-aiofiles", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
 ]
 
 [package.dev-dependencies]
@@ -401,6 +403,7 @@ dev = [
 [package.metadata]
 requires-dist = [
     { name = "aiofiles", specifier = ">=24.1.0" },
+    { name = "aiohttp", specifier = ">=3.12.14" },
     { name = "anyio", specifier = "==4.11.0" },
     { name = "exo-pyo3-bindings", editable = "rust/exo_pyo3_bindings" },
     { name = "fastapi", specifier = ">=0.116.1" },
@@ -421,6 +424,7 @@ requires-dist = [
     { name = "rustworkx", specifier = ">=0.17.1" },
     { name = "tiktoken", specifier = ">=0.12.0" },
     { name = "tomlkit", specifier = ">=0.14.0" },
+    { name = "types-aiofiles", specifier = ">=24.1.0.20250708" },
 ]
 
 [package.metadata.requires-dev]