Skip to content

Commit

Permalink
formatting and fixing tests after rebasing
Browse files Browse the repository at this point in the history
  • Loading branch information
@ianpaul10
ianpaul10 committed Dec 28, 2024
1 parent 1dfd058 commit b003292
Show file tree
Hide file tree
Showing 2 changed files with 158 additions and 175 deletions.
4 changes: 2 additions & 2 deletions exo/networking/manual/manual_discovery.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ def __init__(
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
Expand Down Expand Up @@ -54,7 +54,7 @@ async def task_find_peers_from_config(self):
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.")
Expand Down
329 changes: 156 additions & 173 deletions exo/networking/manual/test_manual_discovery.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,187 +12,170 @@


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())

0 comments on commit b003292

Please sign in to comment.