diff --git a/flake.lock b/flake.lock index 309def17..869ba848 100644 --- a/flake.lock +++ b/flake.lock @@ -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, diff --git a/src/exo/master/tests/conftest.py b/src/exo/master/tests/conftest.py index dec1a113..e931d1a7 100644 --- a/src/exo/master/tests/conftest.py +++ b/src/exo/master/tests/conftest.py @@ -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 + ), + ) diff --git a/src/exo/master/tests/test_placement.py b/src/exo/master/tests/test_placement.py index c551c0b0..a8b95329 100644 --- a/src/exo/master/tests/test_placement.py +++ b/src/exo/master/tests/test_placement.py @@ -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() diff --git a/src/exo/master/tests/test_placement_utils.py b/src/exo/master/tests/test_placement_utils.py index 7a0d33be..e490a77c 100644 --- a/src/exo/master/tests/test_placement_utils.py +++ b/src/exo/master/tests/test_placement_utils.py @@ -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" + )