← back to Exo
exo topology visualisation that shows the topology of the network, device capabilities and the currently active node using opaque statuses. fixes #36. ready for #33
4b592f9d456567d74c1320261468fa284c5cc438 · 2024-07-18 22:50:18 -0700 · Alex Cheema
Files touched
M exo/inference/shard.pyM exo/networking/grpc/grpc_discovery.pyM exo/networking/grpc/grpc_peer_handle.pyM exo/networking/grpc/grpc_server.pyM exo/networking/grpc/node_service.protoM exo/networking/grpc/node_service_pb2.pyM exo/networking/grpc/node_service_pb2_grpc.pyM exo/orchestration/node.pyM exo/orchestration/standard_node.pyM exo/topology/topology.pyA exo/viz/test_topology_viz.pyA exo/viz/topology_viz.pyM main.pyM setup.py
Diff
commit 4b592f9d456567d74c1320261468fa284c5cc438
Author: Alex Cheema <alexcheema123@gmail.com>
Date: Thu Jul 18 22:50:18 2024 -0700
exo topology visualisation that shows the topology of the network, device capabilities and the currently active node using opaque statuses. fixes #36. ready for #33
---
exo/inference/shard.py | 8 ++
exo/networking/grpc/grpc_discovery.py | 15 ++--
exo/networking/grpc/grpc_peer_handle.py | 4 +
exo/networking/grpc/grpc_server.py | 7 ++
exo/networking/grpc/node_service.proto | 6 ++
exo/networking/grpc/node_service_pb2.py | 12 +--
exo/networking/grpc/node_service_pb2_grpc.py | 43 +++++++++++
exo/orchestration/node.py | 10 +++
exo/orchestration/standard_node.py | 61 +++++++++++++--
exo/topology/topology.py | 3 +-
exo/viz/test_topology_viz.py | 43 +++++++++++
exo/viz/topology_viz.py | 110 +++++++++++++++++++++++++++
main.py | 4 +
setup.py | 1 +
14 files changed, 310 insertions(+), 17 deletions(-)
diff --git a/exo/inference/shard.py b/exo/inference/shard.py
index d4c60408..79c8005e 100644
--- a/exo/inference/shard.py
+++ b/exo/inference/shard.py
@@ -12,3 +12,11 @@ class Shard:
def is_last_layer(self) -> bool:
return self.end_layer == self.n_layers - 1
+
+ def to_dict(self) -> dict:
+ return {
+ "model_id": self.model_id,
+ "start_layer": self.start_layer,
+ "end_layer": self.end_layer,
+ "n_layers": self.n_layers
+ }
diff --git a/exo/networking/grpc/grpc_discovery.py b/exo/networking/grpc/grpc_discovery.py
index f01f3886..3064dff4 100644
--- a/exo/networking/grpc/grpc_discovery.py
+++ b/exo/networking/grpc/grpc_discovery.py
@@ -30,7 +30,7 @@ class GRPCDiscovery(Discovery):
self.listen_port = listen_port
self.broadcast_port = broadcast_port if broadcast_port is not None else listen_port
self.broadcast_interval = broadcast_interval
- self.known_peers: Dict[str, Tuple[GRPCPeerHandle, float]] = {}
+ self.known_peers: Dict[str, Tuple[GRPCPeerHandle, float, float]] = {}
self.broadcast_task = None
self.listen_task = None
self.cleanup_task = None
@@ -73,7 +73,7 @@ class GRPCDiscovery(Discovery):
if DEBUG_DISCOVERY >= 2: print("No new peers discovered in the last grace period. Ending discovery process.")
break # No new peers found in the grace period, we are done
- return [peer_handle for peer_handle, _ in self.known_peers.values()]
+ return [peer_handle for peer_handle, _, _ in self.known_peers.values()]
async def task_broadcast_presence(self):
transport, _ = await asyncio.get_event_loop().create_datagram_endpoint(
@@ -109,9 +109,9 @@ class GRPCDiscovery(Discovery):
peer_port = message['grpc_port']
device_capabilities = DeviceCapabilities(**message['device_capabilities'])
if peer_id not in self.known_peers:
- self.known_peers[peer_id] = (GRPCPeerHandle(peer_id, f"{peer_host}:{peer_port}", device_capabilities), time.time())
+ self.known_peers[peer_id] = (GRPCPeerHandle(peer_id, f"{peer_host}:{peer_port}", device_capabilities), time.time(), time.time())
if DEBUG_DISCOVERY >= 2: print(f"Discovered new peer {peer_id} at {peer_host}:{peer_port}")
- self.known_peers[peer_id] = (self.known_peers[peer_id][0], time.time())
+ self.known_peers[peer_id] = (self.known_peers[peer_id][0], self.known_peers[peer_id][1], time.time())
async def task_listen_for_peers(self):
await asyncio.get_event_loop().create_datagram_endpoint(lambda: ListenProtocol(self.on_listen_message), local_addr=('0.0.0.0', self.listen_port))
@@ -122,8 +122,11 @@ class GRPCDiscovery(Discovery):
try:
current_time = time.time()
timeout = 15 * self.broadcast_interval
- peers_to_remove = [peer_handle.id() for peer_handle, last_seen in self.known_peers.values() if not await peer_handle.is_connected() or current_time - last_seen > timeout]
- if DEBUG_DISCOVERY >= 2: print("Peer statuses:", {peer_handle.id(): f"is_connected={await peer_handle.is_connected()}, last_seen={last_seen}" for peer_handle, last_seen in self.known_peers.values()})
+ peers_to_remove = [
+ peer_handle.id() for peer_handle, connected_at, last_seen in self.known_peers.values() if
+ (not await peer_handle.is_connected() and current_time - connected_at > timeout) or current_time - last_seen > timeout
+ ]
+ if DEBUG_DISCOVERY >= 2: print("Peer statuses:", {peer_handle.id(): f"is_connected={await peer_handle.is_connected()}, {connected_at=}, {last_seen=}" for peer_handle, connected_at, last_seen in self.known_peers.values()})
if DEBUG_DISCOVERY >= 2 and len(peers_to_remove) > 0: print(f"Cleaning up peers: {peers_to_remove}")
for peer_id in peers_to_remove:
if peer_id in self.known_peers: del self.known_peers[peer_id]
diff --git a/exo/networking/grpc/grpc_peer_handle.py b/exo/networking/grpc/grpc_peer_handle.py
index 962750ad..685df5f4 100644
--- a/exo/networking/grpc/grpc_peer_handle.py
+++ b/exo/networking/grpc/grpc_peer_handle.py
@@ -96,3 +96,7 @@ class GRPCPeerHandle(PeerHandle):
async def send_result(self, request_id: str, result: List[int], is_finished: bool) -> None:
request = node_service_pb2.SendResultRequest(request_id=request_id, result=result, is_finished=is_finished)
await self.stub.SendResult(request)
+
+ async def send_opaque_status(self, request_id: str, status: str) -> None:
+ request = node_service_pb2.SendOpaqueStatusRequest(request_id=request_id, status=status)
+ await self.stub.SendOpaqueStatus(request)
\ No newline at end of file
diff --git a/exo/networking/grpc/grpc_server.py b/exo/networking/grpc/grpc_server.py
index e956ddd5..41904708 100644
--- a/exo/networking/grpc/grpc_server.py
+++ b/exo/networking/grpc/grpc_server.py
@@ -90,3 +90,10 @@ class GRPCServer(node_service_pb2_grpc.NodeServiceServicer):
if DEBUG >= 2: print(f"Received SendResult request: {request_id=} {result=} {is_finished=}")
self.node.on_token.trigger_all(request_id, result, is_finished)
return node_service_pb2.Empty()
+
+ async def SendOpaqueStatus(self, request, context):
+ request_id = request.request_id
+ status = request.status
+ if DEBUG >= 2: print(f"Received SendOpaqueStatus request: {request_id=} {status=}")
+ self.node.on_opaque_status.trigger_all(request_id, status)
+ return node_service_pb2.Empty()
diff --git a/exo/networking/grpc/node_service.proto b/exo/networking/grpc/node_service.proto
index d6fb8f0e..a9defa99 100644
--- a/exo/networking/grpc/node_service.proto
+++ b/exo/networking/grpc/node_service.proto
@@ -10,6 +10,7 @@ service NodeService {
rpc CollectTopology (CollectTopologyRequest) returns (Topology) {}
rpc GlobalReset (GlobalResetRequest) returns (Empty) {}
rpc SendResult (SendResultRequest) returns (Empty) {}
+ rpc SendOpaqueStatus (SendOpaqueStatusRequest) returns (Empty) {}
}
message Shard {
@@ -91,4 +92,9 @@ message SendResultRequest {
bool is_finished = 3;
}
+message SendOpaqueStatusRequest {
+ string request_id = 1;
+ string status = 2;
+}
+
message Empty {}
\ No newline at end of file
diff --git a/exo/networking/grpc/node_service_pb2.py b/exo/networking/grpc/node_service_pb2.py
index 4909992d..b7a019a3 100644
--- a/exo/networking/grpc/node_service_pb2.py
+++ b/exo/networking/grpc/node_service_pb2.py
@@ -14,7 +14,7 @@ _sym_db = _symbol_database.Default()
-DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x12node_service.proto\x12\x0cnode_service\"S\n\x05Shard\x12\x10\n\x08model_id\x18\x01 \x01(\t\x12\x13\n\x0bstart_layer\x18\x02 \x01(\x05\x12\x11\n\tend_layer\x18\x03 \x01(\x05\x12\x10\n\x08n_layers\x18\x04 \x01(\x05\"\x9d\x01\n\rPromptRequest\x12\"\n\x05shard\x18\x01 \x01(\x0b\x32\x13.node_service.Shard\x12\x0e\n\x06prompt\x18\x02 \x01(\t\x12\x1c\n\x0finference_state\x18\x03 \x01(\tH\x00\x88\x01\x01\x12\x17\n\nrequest_id\x18\x04 \x01(\tH\x01\x88\x01\x01\x42\x12\n\x10_inference_stateB\r\n\x0b_request_id\"\xb3\x01\n\rTensorRequest\x12\"\n\x05shard\x18\x01 \x01(\x0b\x32\x13.node_service.Shard\x12$\n\x06tensor\x18\x02 \x01(\x0b\x32\x14.node_service.Tensor\x12\x1c\n\x0finference_state\x18\x03 \x01(\tH\x00\x88\x01\x01\x12\x17\n\nrequest_id\x18\x04 \x01(\tH\x01\x88\x01\x01\x42\x12\n\x10_inference_stateB\r\n\x0b_request_id\"/\n\x19GetInferenceResultRequest\x12\x12\n\nrequest_id\x18\x01 \x01(\t\"\\\n\x0fInferenceResult\x12)\n\x06tensor\x18\x01 \x01(\x0b\x32\x14.node_service.TensorH\x00\x88\x01\x01\x12\x13\n\x0bis_finished\x18\x02 \x01(\x08\x42\t\n\x07_tensor\";\n\x06Tensor\x12\x13\n\x0btensor_data\x18\x01 \x01(\x0c\x12\r\n\x05shape\x18\x02 \x03(\x05\x12\r\n\x05\x64type\x18\x03 \x01(\t\"7\n\x11ResetShardRequest\x12\"\n\x05shard\x18\x01 \x01(\x0b\x32\x13.node_service.Shard\"<\n\x16\x43ollectTopologyRequest\x12\x0f\n\x07visited\x18\x01 \x03(\t\x12\x11\n\tmax_depth\x18\x02 \x01(\x05\"a\n\x12GlobalResetRequest\x12\'\n\nbase_shard\x18\x01 \x01(\x0b\x32\x13.node_service.Shard\x12\x0f\n\x07visited\x18\x02 \x03(\t\x12\x11\n\tmax_depth\x18\x03 \x01(\x05\"\x8e\x02\n\x08Topology\x12\x30\n\x05nodes\x18\x01 \x03(\x0b\x32!.node_service.Topology.NodesEntry\x12\x39\n\npeer_graph\x18\x02 \x03(\x0b\x32%.node_service.Topology.PeerGraphEntry\x1aN\n\nNodesEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12/\n\x05value\x18\x02 \x01(\x0b\x32 .node_service.DeviceCapabilities:\x02\x38\x01\x1a\x45\n\x0ePeerGraphEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\"\n\x05value\x18\x02 \x01(\x0b\x32\x13.node_service.Peers:\x02\x38\x01\"\x19\n\x05Peers\x12\x10\n\x08peer_ids\x18\x01 \x03(\t\"7\n\x0b\x44\x65viceFlops\x12\x0c\n\x04\x66p32\x18\x01 \x01(\x02\x12\x0c\n\x04\x66p16\x18\x02 \x01(\x02\x12\x0c\n\x04int8\x18\x03 \x01(\x02\"k\n\x12\x44\x65viceCapabilities\x12\r\n\x05model\x18\x01 \x01(\t\x12\x0c\n\x04\x63hip\x18\x02 \x01(\t\x12\x0e\n\x06memory\x18\x03 \x01(\x05\x12(\n\x05\x66lops\x18\x04 \x01(\x0b\x32\x19.node_service.DeviceFlops\"L\n\x11SendResultRequest\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0e\n\x06result\x18\x02 \x03(\x05\x12\x13\n\x0bis_finished\x18\x03 \x01(\x08\"\x07\n\x05\x45mpty2\x9a\x04\n\x0bNodeService\x12\x41\n\nSendPrompt\x12\x1b.node_service.PromptRequest\x1a\x14.node_service.Tensor\"\x00\x12\x41\n\nSendTensor\x12\x1b.node_service.TensorRequest\x1a\x14.node_service.Tensor\"\x00\x12\x44\n\nResetShard\x12\x1f.node_service.ResetShardRequest\x1a\x13.node_service.Empty\"\x00\x12^\n\x12GetInferenceResult\x12\'.node_service.GetInferenceResultRequest\x1a\x1d.node_service.InferenceResult\"\x00\x12Q\n\x0f\x43ollectTopology\x12$.node_service.CollectTopologyRequest\x1a\x16.node_service.Topology\"\x00\x12\x46\n\x0bGlobalReset\x12 .node_service.GlobalResetRequest\x1a\x13.node_service.Empty\"\x00\x12\x44\n\nSendResult\x12\x1f.node_service.SendResultRequest\x1a\x13.node_service.Empty\"\x00\x62\x06proto3')
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x12node_service.proto\x12\x0cnode_service\"S\n\x05Shard\x12\x10\n\x08model_id\x18\x01 \x01(\t\x12\x13\n\x0bstart_layer\x18\x02 \x01(\x05\x12\x11\n\tend_layer\x18\x03 \x01(\x05\x12\x10\n\x08n_layers\x18\x04 \x01(\x05\"\x9d\x01\n\rPromptRequest\x12\"\n\x05shard\x18\x01 \x01(\x0b\x32\x13.node_service.Shard\x12\x0e\n\x06prompt\x18\x02 \x01(\t\x12\x1c\n\x0finference_state\x18\x03 \x01(\tH\x00\x88\x01\x01\x12\x17\n\nrequest_id\x18\x04 \x01(\tH\x01\x88\x01\x01\x42\x12\n\x10_inference_stateB\r\n\x0b_request_id\"\xb3\x01\n\rTensorRequest\x12\"\n\x05shard\x18\x01 \x01(\x0b\x32\x13.node_service.Shard\x12$\n\x06tensor\x18\x02 \x01(\x0b\x32\x14.node_service.Tensor\x12\x1c\n\x0finference_state\x18\x03 \x01(\tH\x00\x88\x01\x01\x12\x17\n\nrequest_id\x18\x04 \x01(\tH\x01\x88\x01\x01\x42\x12\n\x10_inference_stateB\r\n\x0b_request_id\"/\n\x19GetInferenceResultRequest\x12\x12\n\nrequest_id\x18\x01 \x01(\t\"\\\n\x0fInferenceResult\x12)\n\x06tensor\x18\x01 \x01(\x0b\x32\x14.node_service.TensorH\x00\x88\x01\x01\x12\x13\n\x0bis_finished\x18\x02 \x01(\x08\x42\t\n\x07_tensor\";\n\x06Tensor\x12\x13\n\x0btensor_data\x18\x01 \x01(\x0c\x12\r\n\x05shape\x18\x02 \x03(\x05\x12\r\n\x05\x64type\x18\x03 \x01(\t\"7\n\x11ResetShardRequest\x12\"\n\x05shard\x18\x01 \x01(\x0b\x32\x13.node_service.Shard\"<\n\x16\x43ollectTopologyRequest\x12\x0f\n\x07visited\x18\x01 \x03(\t\x12\x11\n\tmax_depth\x18\x02 \x01(\x05\"a\n\x12GlobalResetRequest\x12\'\n\nbase_shard\x18\x01 \x01(\x0b\x32\x13.node_service.Shard\x12\x0f\n\x07visited\x18\x02 \x03(\t\x12\x11\n\tmax_depth\x18\x03 \x01(\x05\"\x8e\x02\n\x08Topology\x12\x30\n\x05nodes\x18\x01 \x03(\x0b\x32!.node_service.Topology.NodesEntry\x12\x39\n\npeer_graph\x18\x02 \x03(\x0b\x32%.node_service.Topology.PeerGraphEntry\x1aN\n\nNodesEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12/\n\x05value\x18\x02 \x01(\x0b\x32 .node_service.DeviceCapabilities:\x02\x38\x01\x1a\x45\n\x0ePeerGraphEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\"\n\x05value\x18\x02 \x01(\x0b\x32\x13.node_service.Peers:\x02\x38\x01\"\x19\n\x05Peers\x12\x10\n\x08peer_ids\x18\x01 \x03(\t\"7\n\x0b\x44\x65viceFlops\x12\x0c\n\x04\x66p32\x18\x01 \x01(\x02\x12\x0c\n\x04\x66p16\x18\x02 \x01(\x02\x12\x0c\n\x04int8\x18\x03 \x01(\x02\"k\n\x12\x44\x65viceCapabilities\x12\r\n\x05model\x18\x01 \x01(\t\x12\x0c\n\x04\x63hip\x18\x02 \x01(\t\x12\x0e\n\x06memory\x18\x03 \x01(\x05\x12(\n\x05\x66lops\x18\x04 \x01(\x0b\x32\x19.node_service.DeviceFlops\"L\n\x11SendResultRequest\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0e\n\x06result\x18\x02 \x03(\x05\x12\x13\n\x0bis_finished\x18\x03 \x01(\x08\"=\n\x17SendOpaqueStatusRequest\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0e\n\x06status\x18\x02 \x01(\t\"\x07\n\x05\x45mpty2\xec\x04\n\x0bNodeService\x12\x41\n\nSendPrompt\x12\x1b.node_service.PromptRequest\x1a\x14.node_service.Tensor\"\x00\x12\x41\n\nSendTensor\x12\x1b.node_service.TensorRequest\x1a\x14.node_service.Tensor\"\x00\x12\x44\n\nResetShard\x12\x1f.node_service.ResetShardRequest\x1a\x13.node_service.Empty\"\x00\x12^\n\x12GetInferenceResult\x12\'.node_service.GetInferenceResultRequest\x1a\x1d.node_service.InferenceResult\"\x00\x12Q\n\x0f\x43ollectTopology\x12$.node_service.CollectTopologyRequest\x1a\x16.node_service.Topology\"\x00\x12\x46\n\x0bGlobalReset\x12 .node_service.GlobalResetRequest\x1a\x13.node_service.Empty\"\x00\x12\x44\n\nSendResult\x12\x1f.node_service.SendResultRequest\x1a\x13.node_service.Empty\"\x00\x12P\n\x10SendOpaqueStatus\x12%.node_service.SendOpaqueStatusRequest\x1a\x13.node_service.Empty\"\x00\x62\x06proto3')
_globals = globals()
_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
@@ -57,8 +57,10 @@ if not _descriptor._USE_C_DESCRIPTORS:
_globals['_DEVICECAPABILITIES']._serialized_end=1349
_globals['_SENDRESULTREQUEST']._serialized_start=1351
_globals['_SENDRESULTREQUEST']._serialized_end=1427
- _globals['_EMPTY']._serialized_start=1429
- _globals['_EMPTY']._serialized_end=1436
- _globals['_NODESERVICE']._serialized_start=1439
- _globals['_NODESERVICE']._serialized_end=1977
+ _globals['_SENDOPAQUESTATUSREQUEST']._serialized_start=1429
+ _globals['_SENDOPAQUESTATUSREQUEST']._serialized_end=1490
+ _globals['_EMPTY']._serialized_start=1492
+ _globals['_EMPTY']._serialized_end=1499
+ _globals['_NODESERVICE']._serialized_start=1502
+ _globals['_NODESERVICE']._serialized_end=2122
# @@protoc_insertion_point(module_scope)
diff --git a/exo/networking/grpc/node_service_pb2_grpc.py b/exo/networking/grpc/node_service_pb2_grpc.py
index 6bb772aa..920a2a3e 100644
--- a/exo/networking/grpc/node_service_pb2_grpc.py
+++ b/exo/networking/grpc/node_service_pb2_grpc.py
@@ -74,6 +74,11 @@ class NodeServiceStub(object):
request_serializer=node__service__pb2.SendResultRequest.SerializeToString,
response_deserializer=node__service__pb2.Empty.FromString,
_registered_method=True)
+ self.SendOpaqueStatus = channel.unary_unary(
+ '/node_service.NodeService/SendOpaqueStatus',
+ request_serializer=node__service__pb2.SendOpaqueStatusRequest.SerializeToString,
+ response_deserializer=node__service__pb2.Empty.FromString,
+ _registered_method=True)
class NodeServiceServicer(object):
@@ -121,6 +126,12 @@ class NodeServiceServicer(object):
context.set_details('Method not implemented!')
raise NotImplementedError('Method not implemented!')
+ def SendOpaqueStatus(self, request, context):
+ """Missing associated documentation comment in .proto file."""
+ context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+ context.set_details('Method not implemented!')
+ raise NotImplementedError('Method not implemented!')
+
def add_NodeServiceServicer_to_server(servicer, server):
rpc_method_handlers = {
@@ -159,6 +170,11 @@ def add_NodeServiceServicer_to_server(servicer, server):
request_deserializer=node__service__pb2.SendResultRequest.FromString,
response_serializer=node__service__pb2.Empty.SerializeToString,
),
+ 'SendOpaqueStatus': grpc.unary_unary_rpc_method_handler(
+ servicer.SendOpaqueStatus,
+ request_deserializer=node__service__pb2.SendOpaqueStatusRequest.FromString,
+ response_serializer=node__service__pb2.Empty.SerializeToString,
+ ),
}
generic_handler = grpc.method_handlers_generic_handler(
'node_service.NodeService', rpc_method_handlers)
@@ -358,3 +374,30 @@ class NodeService(object):
timeout,
metadata,
_registered_method=True)
+
+ @staticmethod
+ def SendOpaqueStatus(request,
+ target,
+ options=(),
+ channel_credentials=None,
+ call_credentials=None,
+ insecure=False,
+ compression=None,
+ wait_for_ready=None,
+ timeout=None,
+ metadata=None):
+ return grpc.experimental.unary_unary(
+ request,
+ target,
+ '/node_service.NodeService/SendOpaqueStatus',
+ node__service__pb2.SendOpaqueStatusRequest.SerializeToString,
+ node__service__pb2.Empty.FromString,
+ options,
+ channel_credentials,
+ insecure,
+ call_credentials,
+ compression,
+ wait_for_ready,
+ timeout,
+ metadata,
+ _registered_method=True)
diff --git a/exo/orchestration/node.py b/exo/orchestration/node.py
index da6f320f..1caf6618 100644
--- a/exo/orchestration/node.py
+++ b/exo/orchestration/node.py
@@ -38,7 +38,17 @@ class Node(ABC):
async def global_reset(self, base_shard: Shard, visited: set[str] = set(), max_depth: int = 2) -> None:
pass
+ @property
+ @abstractmethod
+ def current_topology(self) -> Topology:
+ pass
+
@property
@abstractmethod
def on_token(self) -> AsyncCallbackSystem[str, Tuple[str, List[int], bool]]:
pass
+
+ @property
+ @abstractmethod
+ def on_opaque_status(self) -> AsyncCallbackSystem[str, Tuple[str, str]]:
+ pass
diff --git a/exo/orchestration/standard_node.py b/exo/orchestration/standard_node.py
index 268663ea..0e393a62 100644
--- a/exo/orchestration/standard_node.py
+++ b/exo/orchestration/standard_node.py
@@ -1,5 +1,9 @@
-from typing import List, Dict, Optional, Callable, Tuple, Union
import numpy as np
+import json
+import asyncio
+import uuid
+import time
+from typing import List, Dict, Optional, Callable, Tuple, Union
from exo.networking import Discovery, PeerHandle, Server
from exo.inference.inference_engine import InferenceEngine, Shard
from .node import Node
@@ -9,8 +13,7 @@ from exo.topology.partitioning_strategy import PartitioningStrategy
from exo.topology.partitioning_strategy import Partition
from exo import DEBUG
from exo.helpers import AsyncCallback, AsyncCallbackSystem
-import asyncio
-import uuid
+from exo.viz.topology_viz import TopologyViz
class StandardNode(Node):
def __init__(self, id: str, server: Server, inference_engine: InferenceEngine, discovery: Discovery, partitioning_strategy: PartitioningStrategy = None, max_generate_tokens: int = 256):
@@ -23,8 +26,24 @@ class StandardNode(Node):
self.topology: Topology = Topology()
self.device_capabilities = device_capabilities()
self.buffered_token_output: Dict[str, Tuple[List[int], bool]] = {}
- self._on_token = AsyncCallbackSystem[str, Tuple[str, List[int], bool]]()
+ self.topology_viz = TopologyViz()
self.max_generate_tokens = max_generate_tokens
+ self._on_token = AsyncCallbackSystem[str, Tuple[str, List[int], bool]]()
+ self._on_opaque_status = AsyncCallbackSystem[str, str]()
+ self._on_opaque_status.register("node_status").on_next(self.on_node_status)
+
+ def on_node_status(self, request_id, opaque_status):
+ try:
+ status_data = json.loads(opaque_status)
+ if status_data.get("type", "") == "node_status":
+ if status_data.get("status", "").startswith("start_"):
+ self.current_topology.active_node_id = status_data.get("node_id")
+ elif status_data.get("status", "").startswith("end_"):
+ if status_data.get("node_id") == self.current_topology.active_node_id:
+ self.current_topology.active_node_id = None
+ self.topology_viz.update_visualization(self.current_topology, self.partitioning_strategy.partition(self.current_topology))
+ except json.JSONDecodeError:
+ pass
async def start(self, wait_for_peers: int = 0) -> None:
await self.server.start()
@@ -39,6 +58,15 @@ class StandardNode(Node):
await self.server.stop()
async def process_prompt(self, shard: Shard, prompt: str, request_id: Optional[str] = None, inference_state: Optional[str] = None) -> Optional[np.ndarray]:
+ asyncio.create_task(self.broadcast_opaque_status(request_id, json.dumps({"type": "node_status", "node_id": self.id, "status": "start_process_prompt", "shard": shard.to_dict(), "prompt": prompt, "inference_state": inference_state, "request_id": request_id})))
+ start_time = time.perf_counter_ns()
+ resp = await self._process_prompt(shard, prompt, request_id, inference_state)
+ end_time = time.perf_counter_ns()
+ elapsed_time_ns = end_time - start_time
+ asyncio.create_task(self.broadcast_opaque_status(request_id, json.dumps({"type": "node_status", "node_id": self.id, "status": "end_process_prompt", "shard": shard.to_dict(), "prompt": prompt, "inference_state": inference_state, "request_id": request_id, "elapsed_time_ns": elapsed_time_ns, "result_size": resp.size if resp is not None else 0})))
+ return resp
+
+ async def _process_prompt(self, shard: Shard, prompt: str, request_id: Optional[str] = None, inference_state: Optional[str] = None) -> Optional[np.ndarray]:
if request_id is None:
request_id = str(uuid.uuid4())
if request_id not in self.buffered_token_output:
@@ -68,6 +96,15 @@ class StandardNode(Node):
return np.array(self.buffered_token_output[request_id][0]) if len(self.buffered_token_output[request_id][0]) > 0 else None
async def process_tensor(self, shard: Shard, tensor: np.ndarray, request_id: Optional[str] = None, inference_state: Optional[str] = None) -> Optional[np.ndarray]:
+ asyncio.create_task(self.broadcast_opaque_status(request_id, json.dumps({"type": "node_status", "node_id": self.id, "status": "start_process_tensor", "shard": shard.to_dict(), "tensor_size": tensor.size, "tensor_shape": tensor.shape, "request_id": request_id, "inference_state": inference_state})))
+ start_time = time.perf_counter_ns()
+ resp = await self._process_tensor(shard, tensor, request_id, inference_state)
+ end_time = time.perf_counter_ns()
+ elapsed_time_ns = end_time - start_time
+ asyncio.create_task(self.broadcast_opaque_status(request_id, json.dumps({"type": "node_status", "node_id": self.id, "status": "end_process_tensor", "shard": shard.to_dict(), "request_id": request_id, "elapsed_time_ns": elapsed_time_ns, "result_size": resp.size if resp is not None else 0})))
+ return resp
+
+ async def _process_tensor(self, shard: Shard, tensor: np.ndarray, request_id: Optional[str] = None, inference_state: Optional[str] = None) -> Optional[np.ndarray]:
if request_id is None:
request_id = str(uuid.uuid4())
if request_id not in self.buffered_token_output:
@@ -206,7 +243,9 @@ class StandardNode(Node):
except Exception as e:
print(f"Error collecting topology from {peer.id()}: {e}")
+ next_topology.active_node_id = self.topology.active_node_id # this is not so clean.
self.topology = next_topology
+ self.topology_viz.update_visualization(self.current_topology, self.partitioning_strategy.partition(self.current_topology))
return next_topology
# TODO: unify this and collect_topology as global actions
@@ -237,6 +276,10 @@ class StandardNode(Node):
def on_token(self) -> AsyncCallbackSystem[str, Tuple[str, List[int], bool]]:
return self._on_token
+ @property
+ def on_opaque_status(self) -> AsyncCallbackSystem[str, str]:
+ return self._on_opaque_status
+
def trigger_on_token_callbacks(self, request_id: str, tokens: List[int], is_finished: bool) -> None:
if DEBUG >= 2: print(f"Triggering all on_token callbacks with {request_id=} num_tokens={len(tokens)} {is_finished=}")
self.on_token.trigger_all(request_id, tokens, is_finished)
@@ -253,4 +296,12 @@ class StandardNode(Node):
traceback.print_exc()
print(f"Broadcast result: {request_id=} {result=} {is_finished=}")
- await asyncio.gather(*[send_result_to_peer(peer) for peer in self.peers], return_exceptions=True)
\ No newline at end of file
+ await asyncio.gather(*[send_result_to_peer(peer) for peer in self.peers], return_exceptions=True)
+
+ async def broadcast_opaque_status(self, request_id: str, status: str) -> None:
+ for peer in self.peers:
+ await peer.send_opaque_status(request_id, status)
+
+ @property
+ def current_topology(self) -> Topology:
+ return self.topology
\ No newline at end of file
diff --git a/exo/topology/topology.py b/exo/topology/topology.py
index 38661e0f..1a250191 100644
--- a/exo/topology/topology.py
+++ b/exo/topology/topology.py
@@ -1,10 +1,11 @@
from .device_capabilities import DeviceCapabilities
-from typing import Dict, Set
+from typing import Dict, Set, Optional
class Topology:
def __init__(self):
self.nodes: Dict[str, DeviceCapabilities] = {} # Maps node IDs to DeviceCapabilities
self.peer_graph: Dict[str, Set[str]] = {} # Adjacency list representing the graph
+ self.active_node_id: Optional[str] = None
def update_node(self, node_id: str, device_capabilities: DeviceCapabilities):
self.nodes[node_id] = device_capabilities
diff --git a/exo/viz/test_topology_viz.py b/exo/viz/test_topology_viz.py
new file mode 100644
index 00000000..095eb6ef
--- /dev/null
+++ b/exo/viz/test_topology_viz.py
@@ -0,0 +1,43 @@
+import asyncio
+import unittest
+from exo.viz.topology_viz import TopologyViz
+from exo.topology.topology import Topology
+from exo.topology.device_capabilities import DeviceCapabilities, DeviceFlops
+from exo.topology.partitioning_strategy import Partition
+from exo.helpers import AsyncCallbackSystem
+
+class TestNodeViz(unittest.IsolatedAsyncioTestCase):
+ async def asyncSetUp(self):
+ self.topology = Topology()
+ self.topology.update_node("node1", DeviceCapabilities(model="ModelA", chip="ChipA", memory=8*1024, flops=DeviceFlops(fp32=1.0,fp16=2.0,int8=4.0)))
+ self.topology.update_node("node2", DeviceCapabilities(model="ModelB", chip="ChipB", memory=16*1024, flops=DeviceFlops(fp32=2.0,fp16=4.0,int8=8.0)))
+ self.topology.update_node("node3", DeviceCapabilities(model="ModelC", chip="ChipC", memory=32*1024, flops=DeviceFlops(fp32=4.0,fp16=8.0,int8=16.0)))
+ self.topology.update_node("node4", DeviceCapabilities(model="ModelD", chip="ChipD", memory=64*1024, flops=DeviceFlops(fp32=8.0,fp16=16.0,int8=32.0)))
+
+ self.top_viz = TopologyViz()
+ await asyncio.sleep(2) # Simulate running for a short time
+
+ async def test_layout_generation(self):
+ self.top_viz._generate_layout()
+ self.top_viz.refresh()
+ import time
+ time.sleep(2)
+ self.top_viz.update_visualization(self.topology, [
+ Partition("node1", 0, 0.2),
+ Partition("node4", 0.2, 0.4),
+ Partition("node2", 0.4, 0.8),
+ Partition("node3", 0.8, 1),
+ ])
+ time.sleep(2)
+ self.topology.active_node_id = "node3"
+ self.top_viz.update_visualization(self.topology, [
+ Partition("node1", 0, 0.3),
+ Partition("node2", 0.3, 0.7),
+ Partition("node4", 0.7, 0.9),
+ Partition("node3", 0.9, 1),
+ ])
+ time.sleep(2)
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/exo/viz/topology_viz.py b/exo/viz/topology_viz.py
new file mode 100644
index 00000000..62d7a42f
--- /dev/null
+++ b/exo/viz/topology_viz.py
@@ -0,0 +1,110 @@
+import math
+from typing import Dict, List
+from exo.helpers import exo_text
+from exo.orchestration.node import Node
+from exo.topology.topology import Topology
+from exo.topology.partitioning_strategy import Partition
+from rich.console import Console
+from rich.panel import Panel
+from rich.text import Text
+from rich.live import Live
+from rich.style import Style
+from exo.topology.device_capabilities import DeviceCapabilities, UNKNOWN_DEVICE_CAPABILITIES
+
+class TopologyViz:
+ def __init__(self):
+ self.console = Console()
+ self.topology = Topology()
+ self.partitions: List[Partition] = []
+ self.panel = Panel(self._generate_layout(), title=f"Exo Cluster (0 nodes)", border_style="bright_yellow")
+ self.live_panel = Live(self.panel, auto_refresh=False, console=self.console)
+ self.live_panel.start()
+
+ def update_visualization(self, topology: Topology, partitions: List[Partition]):
+ self.topology = topology
+ self.partitions = partitions
+ self.refresh()
+
+ def refresh(self):
+ self.panel.renderable = self._generate_layout()
+ # Update the panel title with the number of nodes and partitions
+ node_count = len(self.topology.nodes)
+ self.panel.title = f"Exo Cluster ({node_count} node{'s' if node_count != 1 else ''})"
+ self.live_panel.update(self.panel, refresh=True)
+
+ def _generate_layout(self) -> str:
+ # Calculate visualization parameters
+ num_partitions = len(self.partitions)
+ radius = 12 # Reduced radius
+ center_x, center_y = 45, 25 # Adjusted center_x to center the visualization
+
+ # Generate visualization
+ visualization = [[' ' for _ in range(90)] for _ in range(45)] # Increased width to 90
+
+ # Add exo_text at the top in bright yellow
+ exo_lines = exo_text.split('\n')
+ yellow_style = Style(color="bright_yellow")
+ max_line_length = max(len(line) for line in exo_lines)
+ for i, line in enumerate(exo_lines):
+ centered_line = line.center(max_line_length)
+ start_x = (90 - max_line_length) // 2 # Calculate starting x position to center the text
+ colored_line = Text(centered_line, style=yellow_style)
+ for j, char in enumerate(str(colored_line)):
+ if 0 <= start_x + j < 90 and i < len(visualization): # Ensure we don't exceed the width and height
+ visualization[i][start_x + j] = char
+
+ for i, partition in enumerate(self.partitions):
+ device_capabilities = self.topology.nodes.get(partition.node_id, UNKNOWN_DEVICE_CAPABILITIES)
+
+ angle = 2 * math.pi * i / num_partitions
+ x = int(center_x + radius * math.cos(angle))
+ y = int(center_y + radius * math.sin(angle))
+
+ # Place node with different color for active node
+ if partition.node_id == self.topology.active_node_id:
+ visualization[y][x] = '🔴' # Red circle for active node
+ else:
+ visualization[y][x] = '🔵' # Blue circle for inactive nodes
+
+ # Place node info (ID, start_layer, end_layer)
+ node_info = [
+ f"Model: {device_capabilities.model}",
+ f"Mem: {device_capabilities.memory // 1024}GB",
+ f"FLOPS: {device_capabilities.flops.fp16}T",
+ f"Part: {partition.start:.2f}-{partition.end:.2f}"
+ ]
+
+ # Calculate info position based on angle
+ info_distance = radius + 3 # Reduced distance
+ info_x = int(center_x + info_distance * math.cos(angle))
+ info_y = int(center_y + info_distance * math.sin(angle))
+
+ # Adjust text position to avoid overwriting the node icon
+ if info_x < x: # Text is to the left of the node
+ info_x = max(0, x - len(max(node_info, key=len)) - 1)
+ elif info_x > x: # Text is to the right of the node
+ info_x = min(89 - len(max(node_info, key=len)), info_x)
+
+ for j, line in enumerate(node_info):
+ for k, char in enumerate(line):
+ if 0 <= info_y + j < 45 and 0 <= info_x + k < 90: # Updated width check
+ # Ensure we're not overwriting the node icon
+ if info_y + j != y or info_x + k != x:
+ visualization[info_y + j][info_x + k] = char
+
+ # Draw line to next node
+ next_i = (i + 1) % num_partitions
+ next_angle = 2 * math.pi * next_i / num_partitions
+ next_x = int(center_x + radius * math.cos(next_angle))
+ next_y = int(center_y + radius * math.sin(next_angle))
+
+ # Simple line drawing
+ steps = max(abs(next_x - x), abs(next_y - y))
+ for step in range(1, steps):
+ line_x = int(x + (next_x - x) * step / steps)
+ line_y = int(y + (next_y - y) * step / steps)
+ if 0 <= line_y < 45 and 0 <= line_x < 90: # Updated width check
+ visualization[line_y][line_x] = '-'
+
+ # Convert to string
+ return '\n'.join(''.join(str(char) for char in row) for row in visualization)
diff --git a/main.py b/main.py
index b79fee74..f492ead1 100644
--- a/main.py
+++ b/main.py
@@ -5,8 +5,10 @@ import uuid
import platform
import psutil
import os
+import json
from typing import List
from exo.orchestration.standard_node import StandardNode
+from exo.viz.topology_viz import TopologyViz
from exo.networking.grpc.grpc_server import GRPCServer
from exo.networking.grpc.grpc_discovery import GRPCDiscovery
from exo.topology.ring_memory_weighted_partitioning_strategy import RingMemoryWeightedPartitioningStrategy
@@ -59,7 +61,9 @@ server = GRPCServer(node, args.node_host, args.node_port)
node.server = server
api = ChatGPTAPI(node, inference_engine.__class__.__name__)
+topology_viz = TopologyViz()
node.on_token.register("main_log").on_next(lambda _, tokens , __: print(inference_engine.tokenizer.decode(tokens) if hasattr(inference_engine, "tokenizer") else tokens))
+node.on_opaque_status.register("main_log").on_next(lambda request_id, status: print(f"!!! [{request_id}] Opaque Status: {status}"))
async def shutdown(signal, loop):
"""Gracefully shutdown the server and close the asyncio loop."""
diff --git a/setup.py b/setup.py
index 3cc53d65..4b1db878 100644
--- a/setup.py
+++ b/setup.py
@@ -15,6 +15,7 @@ install_requires = [
"psutil==6.0.0",
"pynvml==11.5.3",
"requests==2.32.3",
+ "rich==13.7.1",
"safetensors==0.4.3",
"tiktoken==0.7.0",
"tokenizers==0.19.1",
← 35177690 by default find an ephemeral node port fixes #35, more robus
·
back to Exo
·
fix issues with chatgpt api where it would generate too long dd09c597 →