electron-rare/exo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fixup

Browse Source
This commit is contained in:
Evan
2025-12-24 19:54:48 +00:00
parent 780fd1da1a
commit a05b0b91ad
7 changed files with 68 additions and 77 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/exo/master/placement.py
+3 -2
View File
@@ -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
src/exo/master/placement_utils.py
+6 -9
View File
@@ -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(
src/exo/shared/topology.py
+1 -4
View File
@@ -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]
src/exo/shared/types/topology.py
+3 -3
View File
@@ -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")
src/exo/worker/engines/mlx/__init__.py
+31 -32
View File
@@ -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: ...
src/exo/worker/main.py
+16 -21
View File
@@ -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)
src/exo/worker/utils/net_profile.py
+8 -6
View File
@@ -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