diff --git a/src/exo/master/placement.py b/src/exo/master/placement.py index 68567b87..a6494905 100644 --- a/src/exo/master/placement.py +++ b/src/exo/master/placement.py @@ -30,6 +30,7 @@ from exo.shared.types.worker.instances import ( MlxJacclInstance, MlxRingInstance, ) +from exo.routing.connection_message import IpAddress def random_ephemeral_port() -> int: @@ -130,13 +131,13 @@ def place_instance( jaccl_coordinators=mlx_jaccl_coordinators, ) case InstanceMeta.MlxRing: - hosts: list[Host] = get_hosts_from_subgraph(cycle_digraph) + hosts: list[IpAddress] = get_hosts_from_subgraph(cycle_digraph) target_instances[instance_id] = MlxRingInstance( instance_id=instance_id, shard_assignments=shard_assignments, hosts=[ Host( - ip=host.ip, + ip=str(host), port=random_ephemeral_port(), ) for host in hosts diff --git a/src/exo/master/placement_utils.py b/src/exo/master/placement_utils.py index 559891b4..a62a36b0 100644 --- a/src/exo/master/placement_utils.py +++ b/src/exo/master/placement_utils.py @@ -5,7 +5,7 @@ from loguru import logger from pydantic import BaseModel from exo.shared.topology import Topology -from exo.shared.types.common import Host, NodeId +from exo.shared.types.common import NodeId from exo.shared.types.memory import Memory from exo.shared.types.models import ModelMetadata from exo.shared.types.profiling import NodePerformanceProfile @@ -17,6 +17,7 @@ from exo.shared.types.worker.shards import ( ShardMetadata, TensorShardMetadata, ) +from exo.routing.connection_message import IpAddress class NodeWithProfile(BaseModel): @@ -153,7 +154,7 @@ def get_shard_assignments( ) -def get_hosts_from_subgraph(cycle_digraph: Topology) -> list[Host]: +def get_hosts_from_subgraph(cycle_digraph: Topology) -> list[IpAddress]: # this function is wrong. cycles = cycle_digraph.get_cycles() expected_length = len(list(cycle_digraph.list_nodes())) @@ -172,7 +173,7 @@ def get_hosts_from_subgraph(cycle_digraph: Topology) -> list[Host]: logger.info(f"Using thunderbolt cycle: {get_thunderbolt}") cycle = cycles[0] - hosts: list[Host] = [] + hosts: list[IpAddress] = [] for i in range(len(cycle)): current_node = cycle[i] next_node = cycle[(i + 1) % len(cycle)] @@ -185,11 +186,7 @@ def get_hosts_from_subgraph(cycle_digraph: Topology) -> list[Host]: if get_thunderbolt and not connection.is_thunderbolt(): continue assert connection.sink_addr is not None - host = Host( - ip=str(connection.sink_addr.ip), - port=connection.sink_addr.port, - ) - hosts.append(host) + hosts.append(connection.sink_addr) break return hosts @@ -246,7 +243,7 @@ def _find_connection_ip( connection.source_id == node_i.node_id and connection.sink_id == node_j.node_id ): - yield str(connection.sink_addr.ip) + yield str(connection.sink_addr) def _find_interface_name_for_ip( diff --git a/src/exo/shared/topology.py b/src/exo/shared/topology.py index 2d1b345c..5dd11548 100644 --- a/src/exo/shared/topology.py +++ b/src/exo/shared/topology.py @@ -132,10 +132,7 @@ class Topology: return for connection in self.list_connections(): - if ( - connection.local_node_id == node_id - or connection.send_back_node_id == node_id - ): + if connection.source_id == node_id or connection.sink_id == node_id: self.remove_connection(connection) rx_idx = self._node_id_to_rx_id_map[node_id] diff --git a/src/exo/shared/types/topology.py b/src/exo/shared/types/topology.py index f632dd1b..4a8da37b 100644 --- a/src/exo/shared/types/topology.py +++ b/src/exo/shared/types/topology.py @@ -1,4 +1,4 @@ -from exo.routing.connection_message import SocketAddress +from exo.routing.connection_message import IpAddress from exo.shared.types.common import NodeId from exo.shared.types.profiling import ConnectionProfile, NodePerformanceProfile from exo.utils.pydantic_ext import CamelCaseModel @@ -12,7 +12,7 @@ class NodeInfo(CamelCaseModel): class Connection(CamelCaseModel): source_id: NodeId sink_id: NodeId - sink_addr: SocketAddress + sink_addr: IpAddress connection_profile: ConnectionProfile | None = None def __hash__(self) -> int: @@ -34,4 +34,4 @@ class Connection(CamelCaseModel): ) def is_thunderbolt(self) -> bool: - return str(self.sink_addr.ip).startswith("169.254") + return str(self.sink_addr).startswith("169.254") diff --git a/src/exo/worker/engines/mlx/__init__.py b/src/exo/worker/engines/mlx/__init__.py index 6c0a8323..bf8601f8 100644 --- a/src/exo/worker/engines/mlx/__init__.py +++ b/src/exo/worker/engines/mlx/__init__.py @@ -1,41 +1,40 @@ -from typing import TYPE_CHECKING, Any +from typing import Any -if TYPE_CHECKING: - import mlx.core as mx - import mlx.nn as nn - from mlx_lm.models.cache import KVCache +import mlx.core as mx +import mlx.nn as nn +from mlx_lm.models.cache import KVCache - # These are wrapper functions to fix the fact that mlx is not strongly typed in the same way that EXO is. - # For example - MLX has no guarantee of the interface that nn.Module will expose. But we need a guarantee that it has a __call__() function +# These are wrapper functions to fix the fact that mlx is not strongly typed in the same way that EXO is. +# For example - MLX has no guarantee of the interface that nn.Module will expose. But we need a guarantee that it has a __call__() function - class Model(nn.Module): - layers: list[nn.Module] +class Model(nn.Module): + layers: list[nn.Module] - def __call__( - self, - x: mx.array, - cache: list[KVCache] | None, - input_embeddings: mx.array | None = None, - ) -> mx.array: ... + def __call__( + self, + x: mx.array, + cache: list[KVCache] | None, + input_embeddings: mx.array | None = None, + ) -> mx.array: ... - class Detokenizer: - def reset(self) -> None: ... - def add_token(self, token: int) -> None: ... - def finalize(self) -> None: ... +class Detokenizer: + def reset(self) -> None: ... + def add_token(self, token: int) -> None: ... + def finalize(self) -> None: ... - @property - def last_segment(self) -> str: ... + @property + def last_segment(self) -> str: ... - class TokenizerWrapper: - bos_token: str | None - eos_token_ids: list[int] - detokenizer: Detokenizer +class TokenizerWrapper: + bos_token: str | None + eos_token_ids: list[int] + detokenizer: Detokenizer - def encode(self, text: str, add_special_tokens: bool = True) -> list[int]: ... + def encode(self, text: str, add_special_tokens: bool = True) -> list[int]: ... - def apply_chat_template( - self, - messages_dicts: list[dict[str, Any]], - tokenize: bool = False, - add_generation_prompt: bool = True, - ) -> str: ... + def apply_chat_template( + self, + messages_dicts: list[dict[str, Any]], + tokenize: bool = False, + add_generation_prompt: bool = True, + ) -> str: ... diff --git a/src/exo/worker/main.py b/src/exo/worker/main.py index ded9577f..1134ba61 100644 --- a/src/exo/worker/main.py +++ b/src/exo/worker/main.py @@ -20,9 +20,12 @@ from exo.shared.types.events import ( NodePerformanceMeasured, TaskCreated, TaskStatusUpdated, + TopologyEdgeCreated, + TopologyEdgeDeleted, ) from exo.shared.types.profiling import MemoryPerformanceProfile, NodePerformanceProfile from exo.shared.types.state import State +from exo.shared.types.topology import Connection from exo.shared.types.tasks import ( CreateRunner, DownloadModel, @@ -254,11 +257,12 @@ class Worker: async def _connection_message_event_writer(self): with self.connection_message_receiver as connection_messages: async for msg in connection_messages: + break + # TODO: use mdns for partial discovery for event in check_connections(self.node_id, msg, self.state): logger.info(f"Worker discovered connection {event}") await self.event_sender.send(event) - 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. @@ -385,30 +389,21 @@ class Worker: async def _poll_connection_updates(self): while True: - # TODO: EdgeDeleted - edges = set(self.state.topology.list_connections()) + edges = self.state.topology.out_edges(self.node_id) + pure_edges = set(edge for _, edge in edges) conns = await check_reachable(self.state.topology) + + for nid, conn in edges: + if nid in conns and conn.sink_addr in conns.get(nid, set()): + continue + + logger.debug(f"ping failed to discover {conn=}") + await self.event_sender.send(TopologyEdgeDeleted(edge=conn)) for nid in conns: for ip in conns[nid]: - edge = Connection( - local_node_id=self.node_id, - send_back_node_id=nid, - # nonsense multiaddr - send_back_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: + edge = Connection(sink_id=self.node_id, source_id=nid, sink_addr=ip) + if edge not in pure_edges: logger.debug(f"ping discovered {edge=}") await self.event_sender.send(TopologyEdgeCreated(edge=edge)) - for nid, conn in self.state.topology.out_edges(self.node_id): - if ( - nid not in conns - or conn.send_back_multiaddr.ip_address not in conns.get(nid, set()) - ): - logger.debug(f"ping failed to discover {conn=}") - await self.event_sender.send(TopologyEdgeDeleted(edge=conn)) - await anyio.sleep(10) diff --git a/src/exo/worker/utils/net_profile.py b/src/exo/worker/utils/net_profile.py index 9137d259..8563a6d0 100644 --- a/src/exo/worker/utils/net_profile.py +++ b/src/exo/worker/utils/net_profile.py @@ -1,19 +1,21 @@ import socket +from ipaddress import ip_address from anyio import create_task_group, to_thread from exo.shared.topology import Topology from exo.shared.types.common import NodeId +from exo.routing.connection_message import IpAddress # TODO: ref. api port async def check_reachability( - target_ip: str, target_node_id: NodeId, out: dict[NodeId, set[str]] + target_ip: IpAddress, target_node_id: NodeId, out: dict[NodeId, set[IpAddress]] ) -> 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, 52415)) + result = await to_thread.run_sync(sock.connect_ex, (str(target_ip), 52415)) except socket.gaierror: # seems to throw on ipv6 loopback. oh well # logger.warning(f"invalid {target_ip=}") @@ -24,18 +26,18 @@ async def check_reachability( if result == 0: if target_node_id not in out: out[target_node_id] = set() - out[target_node_id].add(target_ip) + out[target_node_id].add(ip_address(target_ip)) -async def check_reachable(topology: Topology) -> dict[NodeId, set[str]]: - reachable: dict[NodeId, set[str]] = {} +async def check_reachable(topology: Topology) -> dict[NodeId, set[IpAddress]]: + reachable: dict[NodeId, set[IpAddress]] = {} 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, str(iface.ip_address), node.node_id, reachable + check_reachability, iface.ip_address, node.node_id, reachable ) return reachable