← back to Exo
formatting and fixing tests after rebasing
b003292b89d1294b38c0d0dee47b7bf14fe93ce8 · 2024-12-28 12:31:15 +0700 · Ian Paul
Files touched
M exo/networking/manual/manual_discovery.pyM exo/networking/manual/test_manual_discovery.py
Diff
commit b003292b89d1294b38c0d0dee47b7bf14fe93ce8
Author: Ian Paul <iandouglaspaul@gmail.com>
Date: Sat Dec 28 12:31:15 2024 +0700
formatting and fixing tests after rebasing
---
exo/networking/manual/manual_discovery.py | 4 +-
exo/networking/manual/test_manual_discovery.py | 329 ++++++++++++-------------
2 files changed, 158 insertions(+), 175 deletions(-)
diff --git a/exo/networking/manual/manual_discovery.py b/exo/networking/manual/manual_discovery.py
index 69c6cb21..1c232bb6 100644
--- a/exo/networking/manual/manual_discovery.py
+++ b/exo/networking/manual/manual_discovery.py
@@ -15,7 +15,7 @@ class ManualDiscovery(Discovery):
self,
network_config_path: str,
node_id: str,
- create_peer_handle: Callable[[str, str, DeviceCapabilities], PeerHandle],
+ create_peer_handle: Callable[[str, str, str, DeviceCapabilities], PeerHandle],
):
self.network_config_path = network_config_path
self.node_id = node_id
@@ -54,7 +54,7 @@ class ManualDiscovery(Discovery):
peer = self.known_peers.get(peer_id)
if not peer:
if DEBUG_DISCOVERY >= 2: print(f"{peer_id=} not found in known peers. Adding.")
- peer = self.create_peer_handle(peer_id, f"{peer_config.address}:{peer_config.port}", peer_config.device_capabilities)
+ peer = self.create_peer_handle(peer_id, f"{peer_config.address}:{peer_config.port}", "MAN", peer_config.device_capabilities)
is_healthy = await peer.health_check()
if is_healthy:
if DEBUG_DISCOVERY >= 2: print(f"{peer_id=} at {peer_config.address}:{peer_config.port} is healthy.")
diff --git a/exo/networking/manual/test_manual_discovery.py b/exo/networking/manual/test_manual_discovery.py
index 90a435c0..317fba9d 100644
--- a/exo/networking/manual/test_manual_discovery.py
+++ b/exo/networking/manual/test_manual_discovery.py
@@ -12,187 +12,170 @@ root_path = "./exo/networking/manual/test_data/test_config.json"
class TestSingleNodeManualDiscovery(unittest.IsolatedAsyncioTestCase):
- async def asyncSetUp(self):
- self.peer1 = mock.AsyncMock(spec=Node)
- self.peer1.connect = mock.AsyncMock()
- self.server1 = GRPCServer(self.peer1, "localhost", 8000)
- await self.server1.start()
- self.discovery1 = ManualDiscovery(
- root_path,
- "node1",
- create_peer_handle=lambda peer_id, address, device_capabilities: GRPCPeerHandle(
- peer_id, address, device_capabilities
- ),
- )
- await self.discovery1.start()
-
- async def asyncTearDown(self):
- await self.discovery1.stop()
- await self.server1.stop()
-
- async def test_discovery(self):
- peers1 = await self.discovery1.discover_peers(wait_for_peers=0)
- self.assertEqual(len(peers1), 0)
-
- self.peer1.connect.assert_not_called()
+ async def asyncSetUp(self):
+ self.peer1 = mock.AsyncMock()
+ self.peer1.connect = mock.AsyncMock()
+ self.discovery1 = ManualDiscovery(
+ root_path,
+ "node1",
+ create_peer_handle=lambda peer_id, address, description, device_capabilities: self.peer1,
+ )
+ await self.discovery1.start()
+
+ async def asyncTearDown(self):
+ await self.discovery1.stop()
+
+ async def test_discovery(self):
+ peers1 = await self.discovery1.discover_peers(wait_for_peers=0)
+ assert len(peers1) == 0
+
+ self.peer1.connect.assert_not_called()
class TestManualDiscovery(unittest.IsolatedAsyncioTestCase):
- async def asyncSetUp(self):
- self.peer1 = mock.AsyncMock()
- self.peer2 = mock.AsyncMock()
-
- self.peer1.id = mock.MagicMock(return_value="node2")
- self.peer2.id = mock.MagicMock(return_value="node1")
-
- self.peer1.connect = mock.AsyncMock()
- self.peer2.connect = mock.AsyncMock()
- self.discovery1 = ManualDiscovery(
- root_path,
- "node1",
- create_peer_handle=lambda peer_id, address, description, device_capabilities: self.peer1,
- )
- self.discovery2 = ManualDiscovery(
- root_path,
- "node2",
- create_peer_handle=lambda peer_id, address, description, device_capabilities: self.peer2,
- )
- await self.discovery1.start()
- await self.discovery2.start()
-
- async def asyncTearDown(self):
- await self.discovery1.stop()
- await self.discovery2.stop()
-
- async def test_discovery(self):
- peers1 = await self.discovery1.discover_peers(wait_for_peers=1)
- self.assertEqual(len(peers1), 1)
- peers2 = await self.discovery2.discover_peers(wait_for_peers=1)
- self.assertEqual(len(peers2), 1)
-
- # connect has to be explicitly called after discovery
- self.peer1.connect.assert_not_called()
- self.peer2.connect.assert_not_called()
+ async def asyncSetUp(self):
+ self.peer1 = mock.AsyncMock()
+ self.peer2 = mock.AsyncMock()
+ self.peer1.connect = mock.AsyncMock()
+ self.peer2.connect = mock.AsyncMock()
+ self.discovery1 = ManualDiscovery(
+ root_path,
+ "node1",
+ create_peer_handle=lambda peer_id, address, description, device_capabilities: self.peer1,
+ )
+ self.discovery2 = ManualDiscovery(
+ root_path,
+ "node2",
+ create_peer_handle=lambda peer_id, address, description, device_capabilities: self.peer2,
+ )
+ await self.discovery1.start()
+ await self.discovery2.start()
+
+ async def asyncTearDown(self):
+ await self.discovery1.stop()
+ await self.discovery2.stop()
+
+ async def test_discovery(self):
+ peers1 = await self.discovery1.discover_peers(wait_for_peers=1)
+ assert len(peers1) == 1
+ peers2 = await self.discovery2.discover_peers(wait_for_peers=1)
+ assert len(peers2) == 1
+
+ # connect has to be explicitly called after discovery
+ self.peer1.connect.assert_not_called()
+ self.peer2.connect.assert_not_called()
class TestManualDiscoveryWithGRPCPeerHandle(unittest.IsolatedAsyncioTestCase):
- async def asyncSetUp(self):
- config = NetworkTopology.from_path(root_path)
-
- self.node1 = mock.AsyncMock(spec=Node)
- self.node2 = mock.AsyncMock(spec=Node)
- self.server1 = GRPCServer(
- self.node1, config.peers["node1"].address, config.peers["node1"].port
- )
- self.server2 = GRPCServer(
- self.node2, config.peers["node2"].address, config.peers["node2"].port
- )
- await self.server1.start()
- await self.server2.start()
- self.discovery1 = ManualDiscovery(
- root_path,
- "node1",
- create_peer_handle=lambda peer_id, address, description, device_capabilities: GRPCPeerHandle(
- peer_id, address, description, device_capabilities
- ),
- )
- self.discovery2 = ManualDiscovery(
- root_path,
- "node2",
- create_peer_handle=lambda peer_id, address, description, device_capabilities: GRPCPeerHandle(
- peer_id, address, description, device_capabilities
- ),
- )
- await self.discovery1.start()
- await self.discovery2.start()
-
- async def asyncTearDown(self):
- await self.discovery1.stop()
- await self.discovery2.stop()
- await self.server1.stop()
- await self.server2.stop()
-
- async def test_grpc_discovery(self):
- peers1 = await self.discovery1.discover_peers(wait_for_peers=1)
- self.assertEqual(len(peers1), 1)
- peers2 = await self.discovery2.discover_peers(wait_for_peers=1)
- self.assertEqual(len(peers2), 1)
-
- # Connect
- await peers1[0].connect()
- await peers2[0].connect()
- self.assertTrue(await peers1[0].is_connected())
- self.assertTrue(await peers2[0].is_connected())
-
- # Kill server1
- await self.server1.stop()
-
- self.assertTrue(await peers1[0].is_connected())
- self.assertFalse(await peers2[0].is_connected())
-
- # Kill server2
- await self.server2.stop()
-
- self.assertFalse(await peers1[0].is_connected())
- self.assertFalse(await peers2[0].is_connected())
-
- async def test_dynamic_config_update(self):
- initial_peers = await self.discovery1.discover_peers(wait_for_peers=1)
- self.assertEqual(len(initial_peers), 1)
-
- # Save original config for cleanup
- with open(root_path, "r") as f:
- original_config = json.load(f)
-
- try:
- updated_config = {
- "peers": {
- **original_config["peers"],
- "node3": {
- "address": "localhost",
- "port": 50053,
- "device_capabilities": {
- "model": "Unknown Model",
- "chip": "Unknown Chip",
- "memory": 0,
- "flops": {"fp32": 0, "fp16": 0, "int8": 0},
- },
- },
- }
- }
-
- with open(root_path, "w") as f:
- json.dump(updated_config, f, indent=2)
-
- node3 = mock.AsyncMock(spec=Node)
- server3 = GRPCServer(node3, "localhost", 50053)
- await server3.start()
-
- try:
- # Wait for the config to be reloaded
- await asyncio.sleep(1.5)
-
- updated_peers = await self.discovery1.discover_peers(wait_for_peers=2)
- self.assertEqual(len(updated_peers), 2)
-
- for peer in updated_peers:
- await peer.connect()
- self.assertTrue(await peer.is_connected())
-
- finally:
- await server3.stop()
-
- finally:
- # Restore the original config file
- with open(root_path, "w") as f:
- json.dump(original_config, f, indent=2)
-
- # Wait for the config to be reloaded again
+ async def asyncSetUp(self):
+ config = NetworkTopology.from_path(root_path)
+
+ self.node1 = mock.AsyncMock(spec=Node)
+ self.node2 = mock.AsyncMock(spec=Node)
+ self.server1 = GRPCServer(self.node1, config.peers["node1"].address, config.peers["node1"].port)
+ self.server2 = GRPCServer(self.node2, config.peers["node2"].address, config.peers["node2"].port)
+ await self.server1.start()
+ await self.server2.start()
+ self.discovery1 = ManualDiscovery(
+ root_path,
+ "node1",
+ create_peer_handle=lambda peer_id, address, description, device_capabilities: GRPCPeerHandle(peer_id, address, description, device_capabilities),
+ )
+ self.discovery2 = ManualDiscovery(
+ root_path,
+ "node2",
+ create_peer_handle=lambda peer_id, address, description, device_capabilities: GRPCPeerHandle(peer_id, address, description, device_capabilities),
+ )
+ await self.discovery1.start()
+ await self.discovery2.start()
+
+ async def asyncTearDown(self):
+ await self.discovery1.stop()
+ await self.discovery2.stop()
+ await self.server1.stop()
+ await self.server2.stop()
+
+ async def test_grpc_discovery(self):
+ peers1 = await self.discovery1.discover_peers(wait_for_peers=1)
+ assert len(peers1) == 1
+ peers2 = await self.discovery2.discover_peers(wait_for_peers=1)
+ assert len(peers2) == 1
+
+ # Connect
+ await peers1[0].connect()
+ await peers2[0].connect()
+ self.assertTrue(await peers1[0].is_connected())
+ self.assertTrue(await peers2[0].is_connected())
+
+ # Kill server1
+ await self.server1.stop()
+
+ self.assertTrue(await peers1[0].is_connected())
+ self.assertFalse(await peers2[0].is_connected())
+
+ # Kill server2
+ await self.server2.stop()
+
+ self.assertFalse(await peers1[0].is_connected())
+ self.assertFalse(await peers2[0].is_connected())
+
+ async def test_dynamic_config_update(self):
+ initial_peers = await self.discovery1.discover_peers(wait_for_peers=1)
+ self.assertEqual(len(initial_peers), 1)
+
+ # Save original config for cleanup
+ with open(root_path, "r") as f:
+ original_config = json.load(f)
+
+ try:
+ updated_config = {
+ "peers": {
+ **original_config["peers"],
+ "node3": {
+ "address": "localhost",
+ "port": 50053,
+ "device_capabilities": {
+ "model": "Unknown Model",
+ "chip": "Unknown Chip",
+ "memory": 0,
+ "flops": {"fp32": 0, "fp16": 0, "int8": 0},
+ },
+ },
+ }
+ }
+
+ with open(root_path, "w") as f:
+ json.dump(updated_config, f, indent=2)
+
+ node3 = mock.AsyncMock(spec=Node)
+ server3 = GRPCServer(node3, "localhost", 50053)
+ await server3.start()
+
+ try:
+ # Wait for the config to be reloaded
await asyncio.sleep(1.5)
- updated_peers = await self.discovery1.discover_peers(wait_for_peers=1)
- self.assertEqual(len(updated_peers), 1)
+ updated_peers = await self.discovery1.discover_peers(wait_for_peers=2)
+ self.assertEqual(len(updated_peers), 2)
+
+ for peer in updated_peers:
+ await peer.connect()
+ self.assertTrue(await peer.is_connected())
+
+ finally:
+ await server3.stop()
+
+ finally:
+ # Restore the original config file
+ with open(root_path, "w") as f:
+ json.dump(original_config, f, indent=2)
+
+ # Wait for the config to be reloaded again
+ await asyncio.sleep(1.5)
+
+ updated_peers = await self.discovery1.discover_peers(wait_for_peers=1)
+ self.assertEqual(len(updated_peers), 1)
if __name__ == "__main__":
- asyncio.run(unittest.main())
+ asyncio.run(unittest.main())
← 1dfd058c rm unecessary lock
·
back to Exo
·
fix inference engine for inference state fff8a1a6 →