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

fix merge

Browse Source
This commit is contained in:
Evan
2025-12-24 19:54:48 +00:00
parent 427e0af1e4
commit 922632b48d
4 changed files with 109 additions and 93 deletions
Download Patch File Download Diff File Expand all files Collapse all files
flake.lock
Generated
+39
View File
@@ -1,5 +1,26 @@
{
"nodes": {
"fenix": {
"inputs": {
"nixpkgs": [
"nixpkgs"
],
"rust-analyzer-src": "rust-analyzer-src"
},
"locked": {
"lastModified": 1761893049,
"narHash": "sha256-1TtFDPhC+ZsrOOtBnry1EZC+WipTTvsOVjIEVugqji8=",
"owner": "nix-community",
"repo": "fenix",
"rev": "c2ac9a5c0d6d16630c3b225b874bd14528d1abe6",
"type": "github"
},
"original": {
"owner": "nix-community",
"repo": "fenix",
"type": "github"
}
},
"flake-utils": {
"inputs": {
"systems": "systems"
@@ -36,11 +57,29 @@
},
"root": {
"inputs": {
"fenix": "fenix",
"flake-utils": "flake-utils",
"nixpkgs": "nixpkgs",
"treefmt-nix": "treefmt-nix"
}
},
"rust-analyzer-src": {
"flake": false,
"locked": {
"lastModified": 1761849405,
"narHash": "sha256-igXdvC+WCUN+3gnfk+ptT7rMmxQuY6WbIg1rXMUN1DM=",
"owner": "rust-lang",
"repo": "rust-analyzer",
"rev": "f7de8ae045a5fe80f1203c5a1c3015b05f7c3550",
"type": "github"
},
"original": {
"owner": "rust-lang",
"ref": "nightly",
"repo": "rust-analyzer",
"type": "github"
}
},
"systems": {
"locked": {
"lastModified": 1681028828,
src/exo/master/tests/conftest.py
+45 -54
View File
@@ -1,8 +1,7 @@
from typing import Callable
from ipaddress import ip_address
from itertools import count
import pytest
from exo.routing.connection_message import SocketAddress
from exo.shared.types.common import NodeId
from exo.shared.types.profiling import (
MemoryPerformanceProfile,
@@ -10,61 +9,53 @@ from exo.shared.types.profiling import (
SystemPerformanceProfile,
)
from exo.shared.types.topology import Connection, ConnectionProfile, NodeInfo
from exo.routing.connection_message import SocketAddress
ip_octet_iter = count()
port_iter = count(5000)
@pytest.fixture
def create_node():
def _create_node(memory: int, node_id: NodeId | None = None) -> NodeInfo:
if node_id is None:
node_id = NodeId()
return NodeInfo(
node_id=node_id,
node_profile=NodePerformanceProfile(
model_id="test",
chip_id="test",
friendly_name="test",
memory=MemoryPerformanceProfile.from_bytes(
ram_total=1000,
ram_available=memory,
swap_total=1000,
swap_available=1000,
),
network_interfaces=[],
system=SystemPerformanceProfile(),
def create_node(memory: int, node_id: NodeId | None = None) -> NodeInfo:
if node_id is None:
node_id = NodeId()
return NodeInfo(
node_id=node_id,
node_profile=NodePerformanceProfile(
model_id="test",
chip_id="test",
friendly_name="test",
memory=MemoryPerformanceProfile.from_bytes(
ram_total=1000,
ram_available=memory,
swap_total=1000,
swap_available=1000,
),
)
return _create_node
network_interfaces=[],
system=SystemPerformanceProfile(),
),
)
# TODO: this is a hack to get the port for the send_back_multiaddr
@pytest.fixture
def create_connection() -> Callable[[NodeId, NodeId, int | None], Connection]:
port_counter = 1235
ip_counter = 1
def create_connection(
source_node_id: NodeId,
sink_node_id: NodeId,
*,
port: int | None = None,
ip_octet: int | None = None,
) -> Connection:
global ip_octet_iter
global port_iter
def _create_connection(
source_node_id: NodeId, sink_node_id: NodeId, send_back_port: int | None = None
) -> Connection:
nonlocal port_counter
nonlocal ip_counter
# assign unique ips
ip_counter += 1
if send_back_port is None:
send_back_port = port_counter
port_counter += 1
return Connection(
source_id=source_node_id,
sink_id=sink_node_id,
sink_addr=SocketAddress(
ip=ip_address("169.254.0.{ip_counter}"),
port=send_back_port,
zone_id=None,
return Connection(
source_id=source_node_id,
sink_id=sink_node_id,
sink_addr=SocketAddress(
ip=ip_address(
f"169.254.0.{ip_octet if ip_octet is not None else next(ip_octet_iter)}"
),
connection_profile=ConnectionProfile(
throughput=1000, latency=1000, jitter=1000
),
)
return _create_connection
port=port if port is not None else next(port_iter),
zone_id=None,
),
connection_profile=ConnectionProfile(
throughput=1000, latency=1000, jitter=1000
),
)
src/exo/master/tests/test_placement.py
+1
View File
@@ -26,6 +26,7 @@ from exo.shared.types.worker.shards import Sharding
from exo.master.tests.conftest import create_node, create_connection
@pytest.fixture
def topology() -> Topology:
return Topology()
src/exo/master/tests/test_placement_utils.py
+24 -39
View File
@@ -1,5 +1,3 @@
from typing import Callable
import pytest
from exo.master.placement_utils import (
@@ -14,9 +12,10 @@ from exo.shared.types.common import Host, NodeId
from exo.shared.types.memory import Memory
from exo.shared.types.models import ModelId, ModelMetadata
from exo.shared.types.profiling import NetworkInterfaceInfo, NodePerformanceProfile
from exo.shared.types.topology import Connection, NodeInfo
from exo.shared.types.worker.shards import Sharding
from exo.master.tests.conftest import create_connection, create_node
@pytest.fixture
def topology() -> Topology:
@@ -26,8 +25,6 @@ def topology() -> Topology:
def test_filter_cycles_by_memory(
topology: Topology,
create_node: Callable[[int, NodeId | None], NodeInfo],
create_connection: Callable[[NodeId, NodeId], Connection],
):
# arrange
node1_id = NodeId()
@@ -60,8 +57,6 @@ def test_filter_cycles_by_memory(
def test_filter_cycles_by_insufficient_memory(
topology: Topology,
create_node: Callable[[int, NodeId | None], NodeInfo],
create_connection: Callable[[NodeId, NodeId], Connection],
):
# arrange
node1_id = NodeId()
@@ -90,8 +85,6 @@ def test_filter_cycles_by_insufficient_memory(
def test_filter_multiple_cycles_by_memory(
topology: Topology,
create_node: Callable[[int, NodeId | None], NodeInfo],
create_connection: Callable[[NodeId, NodeId], Connection],
):
# arrange
node_a_id = NodeId()
@@ -129,8 +122,6 @@ def test_filter_multiple_cycles_by_memory(
def test_get_smallest_cycles(
topology: Topology,
create_node: Callable[[int, NodeId | None], NodeInfo],
create_connection: Callable[[NodeId, NodeId], Connection],
):
# arrange
node_a_id = NodeId()
@@ -169,8 +160,6 @@ def test_get_smallest_cycles(
)
def test_get_shard_assignments(
topology: Topology,
create_node: Callable[[int, NodeId | None], NodeInfo],
create_connection: Callable[[NodeId, NodeId], Connection],
available_memory: tuple[int, int, int],
total_layers: int,
expected_layers: tuple[int, int, int],
@@ -230,8 +219,6 @@ def test_get_shard_assignments(
def test_get_hosts_from_subgraph(
topology: Topology,
create_node: Callable[[int, NodeId | None], NodeInfo],
create_connection: Callable[[NodeId, NodeId, int | None], Connection],
):
# arrange
node_a_id = NodeId()
@@ -246,10 +233,10 @@ def test_get_hosts_from_subgraph(
topology.add_node(node_b)
topology.add_node(node_c)
topology.add_connection(create_connection(node_a_id, node_b_id, 5001))
topology.add_connection(create_connection(node_b_id, node_c_id, 5002))
topology.add_connection(create_connection(node_c_id, node_a_id, 5003))
topology.add_connection(create_connection(node_b_id, node_a_id, 5004))
topology.add_connection(create_connection(node_a_id, node_b_id))
topology.add_connection(create_connection(node_b_id, node_c_id))
topology.add_connection(create_connection(node_c_id, node_a_id))
topology.add_connection(create_connection(node_b_id, node_a_id))
# act
hosts = get_hosts_from_subgraph(topology)
@@ -267,8 +254,6 @@ def test_get_hosts_from_subgraph(
def test_get_mlx_jaccl_coordinators(
topology: Topology,
create_node: Callable[[int, NodeId | None], NodeInfo],
create_connection: Callable[[NodeId, NodeId, int | None], Connection],
):
# arrange
node_a_id = NodeId()
@@ -279,12 +264,12 @@ def test_get_mlx_jaccl_coordinators(
node_b = create_node(500 * 1024, node_b_id)
node_c = create_node(1000 * 1024, node_c_id)
conn_a_b = create_connection(node_a_id, node_b_id, 5001)
conn_b_a = create_connection(node_b_id, node_a_id, 5002)
conn_b_c = create_connection(node_b_id, node_c_id, 5003)
conn_c_b = create_connection(node_c_id, node_b_id, 5004)
conn_c_a = create_connection(node_c_id, node_a_id, 5005)
conn_a_c = create_connection(node_a_id, node_c_id, 5006)
conn_a_b = create_connection(node_a_id, node_b_id)
conn_b_a = create_connection(node_b_id, node_a_id)
conn_b_c = create_connection(node_b_id, node_c_id)
conn_c_b = create_connection(node_c_id, node_b_id)
conn_c_a = create_connection(node_c_id, node_a_id)
conn_a_c = create_connection(node_a_id, node_c_id)
# Update node profiles with network interfaces before adding to topology
assert node_a.node_profile is not None
@@ -299,11 +284,11 @@ def test_get_mlx_jaccl_coordinators(
network_interfaces=[
NetworkInterfaceInfo(
name="en3",
ip_address=conn_a_b.send_back_multiaddr.ip_address,
ip_address=conn_a_b.sink_addr.ip,
),
NetworkInterfaceInfo(
name="en4",
ip_address=conn_a_c.send_back_multiaddr.ip_address,
ip_address=conn_a_c.sink_addr.ip,
),
],
system=node_a.node_profile.system,
@@ -316,11 +301,11 @@ def test_get_mlx_jaccl_coordinators(
network_interfaces=[
NetworkInterfaceInfo(
name="en3",
ip_address=conn_b_a.send_back_multiaddr.ip_address,
ip_address=conn_b_a.sink_addr.ip,
),
NetworkInterfaceInfo(
name="en4",
ip_address=conn_b_c.send_back_multiaddr.ip_address,
ip_address=conn_b_c.sink_addr.ip,
),
],
system=node_b.node_profile.system,
@@ -333,11 +318,11 @@ def test_get_mlx_jaccl_coordinators(
network_interfaces=[
NetworkInterfaceInfo(
name="en3",
ip_address=conn_c_b.send_back_multiaddr.ip_address,
ip_address=conn_c_b.sink_addr.ip,
),
NetworkInterfaceInfo(
name="en4",
ip_address=conn_c_a.send_back_multiaddr.ip_address,
ip_address=conn_c_a.sink_addr.ip,
),
],
system=node_c.node_profile.system,
@@ -387,11 +372,11 @@ def test_get_mlx_jaccl_coordinators(
# Non-rank-0 nodes should use the specific IP from their connection to rank 0
# node_b uses the IP from conn_b_a (node_b -> node_a)
assert coordinators[node_b_id] == (
f"{conn_b_a.send_back_multiaddr.ip_address}:5000"
), "node_b should use the IP from conn_b_a"
assert coordinators[node_b_id] == (f"{conn_b_a.sink_addr.ip}:5000"), (
"node_b should use the IP from conn_b_a"
)
# node_c uses the IP from conn_c_a (node_c -> node_a)
assert coordinators[node_c_id] == (
f"{conn_c_a.send_back_multiaddr.ip_address}:5000"
), "node_c should use the IP from conn_c_a"
assert coordinators[node_c_id] == (f"{conn_c_a.sink_addr.ip}:5000"), (
"node_c should use the IP from conn_c_a"
)