diff --git a/src/exo/worker/engines/mlx/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py index c0540a9d..59ad30a9 100644 --- a/src/exo/worker/engines/mlx/utils_mlx.py +++ b/src/exo/worker/engines/mlx/utils_mlx.py @@ -1,10 +1,9 @@ +import json import os import resource import time from pathlib import Path from typing import Any, Callable, cast -import json - from mlx_lm.models.cache import KVCache, QuantizedKVCache, RotatingKVCache from mlx_lm.models.deepseek_v3 import DeepseekV3Model diff --git a/src/exo/worker/main.py b/src/exo/worker/main.py index aa53ff23..8629ee55 100644 --- a/src/exo/worker/main.py +++ b/src/exo/worker/main.py @@ -50,6 +50,7 @@ from exo.worker.download.shard_downloader import RepoDownloadProgress, ShardDown from exo.worker.plan import plan from exo.worker.runner.runner_supervisor import RunnerSupervisor from exo.worker.utils import start_polling_memory_metrics, start_polling_node_metrics +from exo.worker.utils.net_profile import connect_all class Worker: @@ -122,6 +123,7 @@ class Worker: tg.start_soon(self._resend_out_for_delivery) tg.start_soon(self._event_applier) tg.start_soon(self._forward_events) + tg.start_soon(self._poll_connection_updates) # TODO: This is a little gross, but not too bad for msg in self._initial_connection_messages: await self.event_sender.send( @@ -394,6 +396,26 @@ class Worker: await self.local_event_sender.send(fe) self.out_for_delivery[event.event_id] = fe + async def _poll_connection_updates(self): + while True: + # TODO: EdgeDeleted + edges = set(self.state.topology.list_connections()) + conns = await connect_all(self.state.topology) + for nid in conns: + for ip in conns[nid]: + edge = Connection( + local_node_id=self.node_id, + send_back_node_id=nid, + send_back_multiaddr=Multiaddr(address=f"/ip4/{ip}/tcp/8000") + if "." in ip + else Multiaddr(address=f"/ip6/{ip}/tcp/8000"), + ) + if edge not in edges: + logger.debug(f"manually discovered {edge=}") + await self.event_sender.send(TopologyEdgeCreated(edge=edge)) + + await anyio.sleep(10) + def event_relevant_to_worker(event: Event, worker: Worker): # TODO diff --git a/src/exo/worker/utils/net_profile.py b/src/exo/worker/utils/net_profile.py new file mode 100644 index 00000000..923048b0 --- /dev/null +++ b/src/exo/worker/utils/net_profile.py @@ -0,0 +1,41 @@ +import socket + +from anyio import create_task_group, to_thread + +from exo.shared.topology import Topology +from exo.shared.types.common import NodeId + + +# TODO: ref. api port +async def check_reachability( + target_ip: str, target_node_id: NodeId, out: dict[NodeId, set[str]] +) -> None: + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.settimeout(1) # 1 second timeout + try: + result = await to_thread.run_sync(sock.connect_ex, (target_ip, 8000)) + except socket.gaierror: + # seems to throw on ipv6 loopback. oh well + # logger.warning(f"invalid {target_ip=}") + return + finally: + sock.close() + + if result == 0: + if target_node_id not in out: + out[target_node_id] = set() + out[target_node_id].add(target_ip) + + +async def connect_all(topology: Topology) -> dict[NodeId, set[str]]: + reachable: dict[NodeId, set[str]] = {} + async with create_task_group() as tg: + for node in topology.list_nodes(): + if not node.node_profile: + continue + for iface in node.node_profile.network_interfaces: + tg.start_soon( + check_reachability, iface.ip_address, node.node_id, reachable + ) + + return reachable diff --git a/src/exo/worker/utils/profile.py b/src/exo/worker/utils/profile.py index 9506428b..30aca08c 100644 --- a/src/exo/worker/utils/profile.py +++ b/src/exo/worker/utils/profile.py @@ -12,14 +12,15 @@ from exo.shared.types.profiling import ( NodePerformanceProfile, SystemPerformanceProfile, ) -from exo.worker.utils.macmon import ( + +from .macmon import ( MacMonError, Metrics, ) -from exo.worker.utils.macmon import ( +from .macmon import ( get_metrics_async as macmon_get_metrics_async, ) -from exo.worker.utils.system_info import ( +from .system_info import ( get_friendly_name, get_model_and_chip, get_network_interfaces, diff --git a/tmp/disable_bridge_enable_dhcp.sh b/tmp/disable_bridge_enable_dhcp.sh new file mode 100755 index 00000000..8bce9333 --- /dev/null +++ b/tmp/disable_bridge_enable_dhcp.sh @@ -0,0 +1,24 @@ +#!/usr/bin/env bash +set -euo pipefail + +networksetup -listallnetworkservices | grep -q '^Thunderbolt Bridge$' \ + && echo "Disabling bridge in networksetup" \ + && networksetup -setnetworkserviceenabled "Thunderbolt Bridge" off + +networksetup -listallnetworkservices | grep -q '^\*Thunderbolt Bridge$' \ + && echo "Bridge disabled in networksetup" + +ifconfig bridge0 &>/dev/null && { + ifconfig bridge0 | grep -q 'member' && echo "Removing bridge members in ifconfig" && { + ifconfig bridge0 | \ + awk '/member/ {print $2}' | \ + xargs -n1 sudo ifconfig bridge0 deletem + } + ifconfig bridge0 | grep -q 'status: active' && sudo ifconfig bridge0 down + ifconfig bridge0 | grep -q 'status: inactive' && echo "Bridge disabled in ifconfig" +} + +for iface in $(seq 2 7); do + sudo ipconfig set "en$iface" dhcp && echo "enabled dhcp on en$iface" || echo "failed to enable dhcp on en$iface" +done +