From 36b8456798d2ed9cb80ce662a72d0ce3be77902a Mon Sep 17 00:00:00 2001 From: Alex Cheema Date: Tue, 25 Jun 2024 12:32:03 +0100 Subject: [PATCH] collect global topology with local peer visibility, ring memory weighted partitioning strategy --- example_user.py | 11 +-- example_user_2.py | 72 +++++++++++++++++++ inference/inference_engine.py | 2 +- inference/mlx/sharded_inference_engine.py | 4 +- main.py | 6 +- networking/grpc/grpc_discovery.py | 18 +++-- networking/grpc/grpc_peer_handle.py | 20 +++++- networking/grpc/grpc_server.py | 7 ++ networking/grpc/node_service.proto | 20 ++++++ networking/grpc/node_service_pb2.py | 26 +++++-- networking/grpc/node_service_pb2_grpc.py | 45 +++++++++++- networking/peer_handle.py | 4 ++ orchestration/node.py | 14 ++-- orchestration/standard_node.py | 70 +++++++++++++----- topology/device_capabilities.py | 12 ++++ topology/partitioning_strategy.py | 16 ++++- ...g_memory_weighted_partitioning_strategy.py | 31 +++----- ...g_memory_weighted_partitioning_strategy.py | 31 ++++++++ topology/topology.py | 43 +++++++++-- 19 files changed, 384 insertions(+), 68 deletions(-) create mode 100644 example_user_2.py create mode 100644 topology/test_ring_memory_weighted_partitioning_strategy.py diff --git a/example_user.py b/example_user.py index 7be62e11..3700c1f1 100644 --- a/example_user.py +++ b/example_user.py @@ -6,6 +6,7 @@ from inference.mlx.sharded_utils import get_model_path, load_tokenizer from inference.shard import Shard from networking.peer_handle import PeerHandle from networking.grpc.grpc_peer_handle import GRPCPeerHandle +from topology.device_capabilities import DeviceCapabilities from typing import List import asyncio import argparse @@ -19,17 +20,19 @@ peers: List[PeerHandle] = [ GRPCPeerHandle( "node1", "localhost:8080", + DeviceCapabilities(model="test1", chip="test1", memory=10000) ), GRPCPeerHandle( "node2", "localhost:8081", + DeviceCapabilities(model="test2", chip="test2", memory=20000) ) ] shards: List[Shard] = [ - # Shard(model_id=path_or_hf_repo, start_layer=0, end_layer=15, n_layers=32), - # Shard(model_id=path_or_hf_repo, start_layer=16, end_layer=31, n_layers=32), - Shard(model_id=path_or_hf_repo, start_layer=0, end_layer=30, n_layers=32), - Shard(model_id=path_or_hf_repo, start_layer=31, end_layer=31, n_layers=32), + Shard(model_id=path_or_hf_repo, start_layer=0, end_layer=15, n_layers=32), + Shard(model_id=path_or_hf_repo, start_layer=16, end_layer=31, n_layers=32), + # Shard(model_id=path_or_hf_repo, start_layer=0, end_layer=30, n_layers=32), + # Shard(model_id=path_or_hf_repo, start_layer=31, end_layer=31, n_layers=32), ] async def run_prompt(prompt: str): diff --git a/example_user_2.py b/example_user_2.py new file mode 100644 index 00000000..d9ceb13d --- /dev/null +++ b/example_user_2.py @@ -0,0 +1,72 @@ +# In this example, a user is running a home cluster with 3 shards. +# They are prompting the cluster to generate a response to a question. +# The cluster is given the question, and the user is given the response. + +from inference.mlx.sharded_utils import get_model_path, load_tokenizer +from inference.shard import Shard +from networking.peer_handle import PeerHandle +from networking.grpc.grpc_peer_handle import GRPCPeerHandle +from topology.device_capabilities import DeviceCapabilities +from typing import List +import asyncio +import argparse + +path_or_hf_repo = "mlx-community/Meta-Llama-3-8B-Instruct-4bit" +model_path = get_model_path(path_or_hf_repo) +tokenizer_config = {} +tokenizer = load_tokenizer(model_path, tokenizer_config) + +peers: List[PeerHandle] = [ + GRPCPeerHandle( + "node1", + "localhost:8080", + DeviceCapabilities(model="test1", chip="test1", memory=10000) + ), +] +shards: List[Shard] = [ + Shard(model_id=path_or_hf_repo, start_layer=0, end_layer=15, n_layers=32), + # Shard(model_id=path_or_hf_repo, start_layer=0, end_layer=30, n_layers=32), + # Shard(model_id=path_or_hf_repo, start_layer=31, end_layer=31, n_layers=32), +] + +async def run_prompt(prompt: str): + if tokenizer.chat_template is None: + tokenizer.chat_template = tokenizer.default_chat_template + if ( + hasattr(tokenizer, "apply_chat_template") + and tokenizer.chat_template is not None + ): + messages = [{"role": "user", "content": prompt}] + prompt = tokenizer.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True + ) + + for peer, shard in zip(peers, shards): + await peer.connect() + await peer.reset_shard(shard) + + tokens = [] + last_output = prompt + + for _ in range(20): + for peer, shard in zip(peers, shards): + if isinstance(last_output, str): + last_output = await peer.send_prompt(shard, last_output) + print("prompt output:", last_output) + else: + last_output = await peer.send_tensor(shard, last_output) + print("tensor output:", last_output) + + if not last_output: + break + + tokens.append(last_output.item()) + + print(tokenizer.decode(tokens)) + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Run prompt") + parser.add_argument("--prompt", type=str, help="The prompt to run") + args = parser.parse_args() + + asyncio.run(run_prompt(args.prompt)) diff --git a/inference/inference_engine.py b/inference/inference_engine.py index c333a532..dc245827 100644 --- a/inference/inference_engine.py +++ b/inference/inference_engine.py @@ -6,7 +6,7 @@ from .shard import Shard class InferenceEngine(ABC): @abstractmethod - async def infer_shard(self, shard: Shard, input_data: np.ndarray) -> np.ndarray: + async def infer_tensor(self, shard: Shard, input_data: np.ndarray) -> np.ndarray: pass @abstractmethod diff --git a/inference/mlx/sharded_inference_engine.py b/inference/mlx/sharded_inference_engine.py index 3050e047..99e8a0b9 100644 --- a/inference/mlx/sharded_inference_engine.py +++ b/inference/mlx/sharded_inference_engine.py @@ -19,11 +19,11 @@ class MLXFixedShardInferenceEngine(InferenceEngine): output_data = self.stateful_sharded_model.step(mx.array(self.tokenizer.encode(prompt))) return np.array(output_data) - async def infer_shard(self, shard: Shard, input_data: np.ndarray) -> np.ndarray: + async def infer_tensor(self, shard: Shard, input_data: np.ndarray) -> np.ndarray: if shard != self.shard: raise ValueError(f"Shard mismatch: {shard} != {self.shard}") - print("infer_shard", shard, input_data) + print("infer_tensor", shard, input_data) output_data = self.stateful_sharded_model.step(mx.array(input_data)) return np.array(output_data) diff --git a/main.py b/main.py index 5c703eb9..dfeff3cd 100644 --- a/main.py +++ b/main.py @@ -8,6 +8,7 @@ from networking.grpc.grpc_server import GRPCServer from inference.mlx.sharded_inference_engine import MLXFixedShardInferenceEngine from inference.shard import Shard from networking.grpc.grpc_discovery import GRPCDiscovery +from topology.ring_memory_weighted_partitioning_strategy import RingMemoryWeightedPartitioningStrategy # parse args parser = argparse.ArgumentParser(description="Initialize GRPC Discovery") @@ -20,11 +21,12 @@ parser.add_argument("--model-id", type=str, default="mlx-community/Meta-Llama-3- parser.add_argument("--n-layers", type=int, default=32, help="Number of layers in the model") parser.add_argument("--start-layer", type=int, default=0, help="Start layer index") parser.add_argument("--end-layer", type=int, default=31, help="End layer index") +parser.add_argument("--wait-for-peers", type=int, default=0, help="Number of peers to wait to connect to before starting") args = parser.parse_args() inference_engine = MLXFixedShardInferenceEngine(args.model_id, shard=Shard(model_id=args.model_id, n_layers=args.n_layers, start_layer=args.start_layer, end_layer=args.end_layer)) discovery = GRPCDiscovery(args.node_id, args.node_port, args.listen_port, args.broadcast_port) -node = StandardNode(args.node_id, None, inference_engine, discovery) +node = StandardNode(args.node_id, None, inference_engine, discovery, partitioning_strategy=RingMemoryWeightedPartitioningStrategy()) server = GRPCServer(node, args.node_host, args.node_port) node.server = server @@ -49,7 +51,7 @@ async def main(): for s in [signal.SIGINT, signal.SIGTERM]: loop.add_signal_handler(s, handle_exit) - await node.start() + await node.start(wait_for_peers=args.wait_for_peers) await asyncio.Event().wait() diff --git a/networking/grpc/grpc_discovery.py b/networking/grpc/grpc_discovery.py index c61deb51..97537007 100644 --- a/networking/grpc/grpc_discovery.py +++ b/networking/grpc/grpc_discovery.py @@ -6,12 +6,13 @@ from typing import List, Dict from ..discovery import Discovery from ..peer_handle import PeerHandle from .grpc_peer_handle import GRPCPeerHandle -from topology.device_capabilities import DeviceCapabilities, mac_device_capabilities +from topology.device_capabilities import DeviceCapabilities, device_capabilities class GRPCDiscovery(Discovery): - def __init__(self, node_id: str, node_port: int, listen_port: int, broadcast_port: int = None, broadcast_interval: int = 1): + def __init__(self, node_id: str, node_port: int, listen_port: int, broadcast_port: int = None, broadcast_interval: int = 1, device_capabilities=None): self.node_id = node_id self.node_port = node_port + self.device_capabilities = device_capabilities self.listen_port = listen_port self.broadcast_port = broadcast_port if broadcast_port is not None else listen_port self.broadcast_interval = broadcast_interval @@ -62,7 +63,9 @@ class GRPCDiscovery(Discovery): return list(self.known_peers.values()) async def _broadcast_presence(self): - self.device_capabilities: DeviceCapabilities = mac_device_capabilities() + if not self.device_capabilities: + self.device_capabilities = device_capabilities() + sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM, socket.IPPROTO_UDP) sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1) sock.settimeout(0.5) @@ -70,7 +73,11 @@ class GRPCDiscovery(Discovery): "type": "discovery", "node_id": self.node_id, "grpc_port": self.node_port, - "device_capabilities": self.device_capabilities.to_dict() + "device_capabilities": { + "model": self.device_capabilities.model, + "chip": self.device_capabilities.chip, + "memory": self.device_capabilities.memory + } }).encode('utf-8') while True: @@ -90,7 +97,8 @@ class GRPCDiscovery(Discovery): peer_id = message['node_id'] peer_host = addr[0] peer_port = message['grpc_port'] - self.known_peers[peer_id] = GRPCPeerHandle(peer_id, f"{peer_host}:{peer_port}") + device_capabilities = DeviceCapabilities(**message['device_capabilities']) + self.known_peers[peer_id] = GRPCPeerHandle(peer_id, f"{peer_host}:{peer_port}", device_capabilities) self.peer_last_seen[peer_id] = time.time() except Exception as e: print(f"Error in peer discovery: {e}") diff --git a/networking/grpc/grpc_peer_handle.py b/networking/grpc/grpc_peer_handle.py index fd652251..fbd83818 100644 --- a/networking/grpc/grpc_peer_handle.py +++ b/networking/grpc/grpc_peer_handle.py @@ -8,15 +8,21 @@ from . import node_service_pb2_grpc from ..peer_handle import PeerHandle from inference.shard import Shard +from topology.topology import Topology +from topology.device_capabilities import DeviceCapabilities class GRPCPeerHandle(PeerHandle): - def __init__(self, id: str, address: str): + def __init__(self, id: str, address: str, device_capabilities: DeviceCapabilities): self._id = id self.address = address + self._device_capabilities = device_capabilities def id(self) -> str: return self._id + def device_capabilities(self) -> DeviceCapabilities: + return self._device_capabilities + async def connect(self): self.channel = grpc.aio.insecure_channel(self.address) self.stub = node_service_pb2_grpc.NodeServiceStub(self.channel) @@ -54,3 +60,15 @@ class GRPCPeerHandle(PeerHandle): request = node_service_pb2.ResetShardRequest(shard=node_service_pb2.Shard(model_id=shard.model_id, start_layer=shard.start_layer, end_layer=shard.end_layer, n_layers=shard.n_layers)) await self.stub.ResetShard(request) print(f"Reset shard {shard} on {self.address}") + + async def collect_topology(self, max_depth: int) -> Topology: + request = node_service_pb2.CollectTopologyRequest(max_depth=max_depth) + response = await self.stub.CollectTopology(request) + topology = Topology() + for node_id, capabilities in response.nodes.items(): + device_capabilities = DeviceCapabilities(model=capabilities.model, chip=capabilities.chip, memory=capabilities.memory) + topology.update_node(node_id, device_capabilities) + for node_id, peers in response.peer_graph.items(): + for peer_id in peers.peer_ids: + topology.add_edge(node_id, peer_id) + return topology diff --git a/networking/grpc/grpc_server.py b/networking/grpc/grpc_server.py index 69fe0520..c5909a05 100644 --- a/networking/grpc/grpc_server.py +++ b/networking/grpc/grpc_server.py @@ -48,3 +48,10 @@ class GRPCServer(node_service_pb2_grpc.NodeServiceServicer): print(f"Received ResetShard request: {shard}") await self.node.reset_shard(shard) return node_service_pb2.Empty() + + async def CollectTopology(self, request, context): + max_depth = request.max_depth + topology = await self.node.collect_topology(max_depth) + nodes = {node_id: node_service_pb2.DeviceCapabilities(model=cap.model, chip=cap.chip, memory=cap.memory) for node_id, cap in topology.nodes.items()} + peer_graph = {node_id: node_service_pb2.Peers(peer_ids=peers) for node_id, peers in topology.peer_graph.items()} + return node_service_pb2.Topology(nodes=nodes, peer_graph=peer_graph) diff --git a/networking/grpc/node_service.proto b/networking/grpc/node_service.proto index 55127a30..c0cfc20e 100644 --- a/networking/grpc/node_service.proto +++ b/networking/grpc/node_service.proto @@ -6,6 +6,7 @@ service NodeService { rpc SendPrompt (PromptRequest) returns (Tensor) {} rpc SendTensor (TensorRequest) returns (Tensor) {} rpc ResetShard (ResetShardRequest) returns (Empty) {} + rpc CollectTopology (CollectTopologyRequest) returns (Topology) {} } message Shard { @@ -35,4 +36,23 @@ message ResetShardRequest { Shard shard = 1; } +message CollectTopologyRequest { + int32 max_depth = 1; +} + +message Topology { + map nodes = 1; + map peer_graph = 2; +} + +message Peers { + repeated string peer_ids = 1; +} + +message DeviceCapabilities { + string model = 1; + string chip = 2; + int32 memory = 3; +} + message Empty {} \ No newline at end of file diff --git a/networking/grpc/node_service_pb2.py b/networking/grpc/node_service_pb2.py index 343b2aac..77c15ca9 100644 --- a/networking/grpc/node_service_pb2.py +++ b/networking/grpc/node_service_pb2.py @@ -14,13 +14,17 @@ _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\"C\n\rPromptRequest\x12\"\n\x05shard\x18\x01 \x01(\x0b\x32\x13.node_service.Shard\x12\x0e\n\x06prompt\x18\x02 \x01(\t\"Y\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\";\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\"\x07\n\x05\x45mpty2\xd9\x01\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\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\"C\n\rPromptRequest\x12\"\n\x05shard\x18\x01 \x01(\x0b\x32\x13.node_service.Shard\x12\x0e\n\x06prompt\x18\x02 \x01(\t\"Y\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\";\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\x11\n\tmax_depth\x18\x01 \x01(\x05\"\x19\n\x05Peers\x12\x10\n\x08peer_ids\x18\x01 \x03(\t\"\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\"A\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\"\x07\n\x05\x45mpty2\xac\x02\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\x12Q\n\x0f\x43ollectTopology\x12$.node_service.CollectTopologyRequest\x1a\x16.node_service.Topology\"\x00\x62\x06proto3') _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) _builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'node_service_pb2', _globals) if not _descriptor._USE_C_DESCRIPTORS: DESCRIPTOR._loaded_options = None + _globals['_TOPOLOGY_NODESENTRY']._loaded_options = None + _globals['_TOPOLOGY_NODESENTRY']._serialized_options = b'8\001' + _globals['_TOPOLOGY_PEERGRAPHENTRY']._loaded_options = None + _globals['_TOPOLOGY_PEERGRAPHENTRY']._serialized_options = b'8\001' _globals['_SHARD']._serialized_start=36 _globals['_SHARD']._serialized_end=119 _globals['_PROMPTREQUEST']._serialized_start=121 @@ -31,8 +35,20 @@ if not _descriptor._USE_C_DESCRIPTORS: _globals['_TENSOR']._serialized_end=340 _globals['_RESETSHARDREQUEST']._serialized_start=342 _globals['_RESETSHARDREQUEST']._serialized_end=397 - _globals['_EMPTY']._serialized_start=399 - _globals['_EMPTY']._serialized_end=406 - _globals['_NODESERVICE']._serialized_start=409 - _globals['_NODESERVICE']._serialized_end=626 + _globals['_COLLECTTOPOLOGYREQUEST']._serialized_start=399 + _globals['_COLLECTTOPOLOGYREQUEST']._serialized_end=442 + _globals['_PEERS']._serialized_start=444 + _globals['_PEERS']._serialized_end=469 + _globals['_TOPOLOGY']._serialized_start=472 + _globals['_TOPOLOGY']._serialized_end=742 + _globals['_TOPOLOGY_NODESENTRY']._serialized_start=593 + _globals['_TOPOLOGY_NODESENTRY']._serialized_end=671 + _globals['_TOPOLOGY_PEERGRAPHENTRY']._serialized_start=673 + _globals['_TOPOLOGY_PEERGRAPHENTRY']._serialized_end=742 + _globals['_DEVICECAPABILITIES']._serialized_start=744 + _globals['_DEVICECAPABILITIES']._serialized_end=809 + _globals['_EMPTY']._serialized_start=811 + _globals['_EMPTY']._serialized_end=818 + _globals['_NODESERVICE']._serialized_start=821 + _globals['_NODESERVICE']._serialized_end=1121 # @@protoc_insertion_point(module_scope) diff --git a/networking/grpc/node_service_pb2_grpc.py b/networking/grpc/node_service_pb2_grpc.py index 8529f3e7..b3752130 100644 --- a/networking/grpc/node_service_pb2_grpc.py +++ b/networking/grpc/node_service_pb2_grpc.py @@ -3,7 +3,7 @@ import grpc import warnings -import node_service_pb2 as node__service__pb2 +from . import node_service_pb2 as node__service__pb2 GRPC_GENERATED_VERSION = '1.64.1' GRPC_VERSION = grpc.__version__ @@ -54,6 +54,11 @@ class NodeServiceStub(object): request_serializer=node__service__pb2.ResetShardRequest.SerializeToString, response_deserializer=node__service__pb2.Empty.FromString, _registered_method=True) + self.CollectTopology = channel.unary_unary( + '/node_service.NodeService/CollectTopology', + request_serializer=node__service__pb2.CollectTopologyRequest.SerializeToString, + response_deserializer=node__service__pb2.Topology.FromString, + _registered_method=True) class NodeServiceServicer(object): @@ -77,6 +82,12 @@ class NodeServiceServicer(object): context.set_details('Method not implemented!') raise NotImplementedError('Method not implemented!') + def CollectTopology(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 = { @@ -95,6 +106,11 @@ def add_NodeServiceServicer_to_server(servicer, server): request_deserializer=node__service__pb2.ResetShardRequest.FromString, response_serializer=node__service__pb2.Empty.SerializeToString, ), + 'CollectTopology': grpc.unary_unary_rpc_method_handler( + servicer.CollectTopology, + request_deserializer=node__service__pb2.CollectTopologyRequest.FromString, + response_serializer=node__service__pb2.Topology.SerializeToString, + ), } generic_handler = grpc.method_handlers_generic_handler( 'node_service.NodeService', rpc_method_handlers) @@ -186,3 +202,30 @@ class NodeService(object): timeout, metadata, _registered_method=True) + + @staticmethod + def CollectTopology(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/CollectTopology', + node__service__pb2.CollectTopologyRequest.SerializeToString, + node__service__pb2.Topology.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) diff --git a/networking/peer_handle.py b/networking/peer_handle.py index 74b7ba57..b418d1d1 100644 --- a/networking/peer_handle.py +++ b/networking/peer_handle.py @@ -3,6 +3,7 @@ from typing import Optional import numpy as np from inference.shard import Shard from topology.device_capabilities import DeviceCapabilities +from topology.topology import Topology class PeerHandle(ABC): @abstractmethod @@ -32,3 +33,6 @@ class PeerHandle(ABC): @abstractmethod async def reset_shard(self, shard: Shard) -> None: pass + + async def collect_topology(self, max_depth: int) -> Topology: + pass diff --git a/orchestration/node.py b/orchestration/node.py index 973898e3..65a1606d 100644 --- a/orchestration/node.py +++ b/orchestration/node.py @@ -2,24 +2,28 @@ from typing import Optional import numpy as np from abc import ABC, abstractmethod from inference.shard import Shard +from topology.topology import Topology class Node(ABC): @abstractmethod - def start(self, wait_for_peers: int = 0) -> None: + async def start(self, wait_for_peers: int = 0) -> None: pass @abstractmethod - def stop(self) -> None: + async def stop(self) -> None: pass @abstractmethod - def process_tensor(self, shard: Shard, tensor: np.ndarray) -> None: + async def process_tensor(self, shard: Shard, tensor: np.ndarray) -> None: pass @abstractmethod - def process_prompt(self, shard: Shard, prompt: str) -> None: + async def process_prompt(self, shard: Shard, prompt: str) -> None: pass @abstractmethod - def reset_shard(self, shard: Shard) -> None: + async def reset_shard(self, shard: Shard) -> None: + pass + + async def collect_topology(self, max_depth: int = 2) -> Topology: pass diff --git a/orchestration/standard_node.py b/orchestration/standard_node.py index 01785446..5f9dd310 100644 --- a/orchestration/standard_node.py +++ b/orchestration/standard_node.py @@ -4,16 +4,20 @@ from networking import Discovery, PeerHandle, Server from inference.inference_engine import InferenceEngine, Shard from .node import Node from topology.topology import Topology +from topology.device_capabilities import device_capabilities +from topology.partitioning_strategy import PartitioningStrategy +from topology.partitioning_strategy import Partition class StandardNode(Node): - def __init__(self, id: str, server: Server, inference_engine: InferenceEngine, discovery: Discovery): + def __init__(self, id: str, server: Server, inference_engine: InferenceEngine, discovery: Discovery, partitioning_strategy: PartitioningStrategy = None): self.id = id self.inference_engine = inference_engine self.server = server self.discovery = discovery + self.partitioning_strategy = partitioning_strategy self.peers: List[PeerHandle] = {} self.topology: Topology = Topology() - self.successor: Optional[PeerHandle] = None + self.device_capabilities = device_capabilities() async def start(self, wait_for_peers: int = 0) -> None: await self.server.start() @@ -24,6 +28,8 @@ class StandardNode(Node): for peer in self.peers: await peer.connect() print(f"Connected to {peer.id()}") + await self.collect_topology() + print(f"Collected topology: {self.topology}") async def stop(self) -> None: await self.discovery.stop() @@ -32,30 +38,62 @@ class StandardNode(Node): async def process_prompt(self, shard: Shard, prompt: str) -> Optional[np.array]: print("Process prompt", shard, prompt) result = await self.inference_engine.infer_prompt(shard, prompt) - # Implement prompt processing logic print(f"Got result from prompt: {prompt}. Result: {result}") - # You might want to initiate inference here - if self.successor: - await self.succesor.send_tensor() + + await self.forward_tensor_to_next_shard(shard, result) return result - async def process_tensor(self, shard: Shard, tensor: np.ndarray, target: Optional[str] = None) -> None: + async def process_tensor(self, shard: Shard, tensor: np.ndarray) -> None: print("Process tensor", shard, tensor) - result = await self.inference_engine.infer_shard(shard, tensor) - # Implement prompt processing logic - print(f"Got result from prompt: {len(tensor)}. Result: {result}") + result = await self.inference_engine.infer_tensor(shard, tensor) + print(f"Got result from tensor: {len(tensor)}. Result: {result}") - if target: - target_peer = next((p for p in self.peers if p.id() == target), None) - if not target_peer: - raise ValueError(f"Peer {target} not found") - - await target_peer.send_tensor(result) + await self.forward_tensor_to_next_shard(shard, result) return result + async def forward_tensor_to_next_shard(self, shard: Shard, tensor: np.ndarray) -> None: + if not self.partitioning_strategy: + print("No partitioning strategy found. Skipping forward.") + return + + partitions = self.partitioning_strategy.partition(self.topology) + current_partition_index = next((i for i, p in enumerate(partitions) if p.node_id == self.id), None) + print(f"Current partition index: {current_partition_index}") + if current_partition_index is not None: + next_partition_index = (current_partition_index + 1) % len(partitions) + next_partition: Partition = partitions[next_partition_index] + print(f"Computed next from: {shard}, {self.topology}. Next partition: {next_partition}") + + if next_partition: + target_peer = next((p for p in self.peers if p.id() == next_partition.node_id), None) + if not target_peer: + raise ValueError(f"Peer for {next_partition} not found") + + start_layer = int(next_partition.start * shard.n_layers) + end_layer = int(next_partition.end * shard.n_layers) - 1 + next_shard = Shard(shard.model_id, start_layer, end_layer, shard.n_layers) + + print(f"Sending tensor to {target_peer.id()} for shard: {next_shard}") + + await target_peer.send_tensor(next_shard, tensor) + async def reset_shard(self, shard: Shard) -> None: # Implement shard reset logic print(f"Resetting shard: {shard}") await self.inference_engine.reset_shard(shard) + + async def collect_topology(self, max_depth: int = 4) -> Topology: + self.topology.update_node(self.id, self.device_capabilities) + + for peer in self.peers: + self.topology.update_node(peer.id(), peer.device_capabilities()) + self.topology.add_edge(self.id, peer.id()) + + if max_depth > 0: + other_topology = await peer.collect_topology(max_depth = max_depth - 1) + print(f"Collected topology from: {peer.id()}: {other_topology}") + self.topology.merge(other_topology) + + return self.topology diff --git a/topology/device_capabilities.py b/topology/device_capabilities.py index 6d71dd5a..c4e34986 100644 --- a/topology/device_capabilities.py +++ b/topology/device_capabilities.py @@ -1,5 +1,6 @@ from dataclasses import dataclass import subprocess +import platform @dataclass class DeviceCapabilities: @@ -7,6 +8,17 @@ class DeviceCapabilities: chip: str memory: int +def device_capabilities() -> DeviceCapabilities: + system = platform.system() + if system == 'Darwin': + return mac_device_capabilities() + # elif system == 'Linux': + # return linux_device_capabilities() + # elif system == 'Windows': + # return windows_device_capabilities() + else: + return DeviceCapabilities(model="Unknown Model", chip="Unknown Chip", memory=0) + def mac_device_capabilities() -> DeviceCapabilities: # Fetch the model of the Mac using system_profiler model = subprocess.check_output(['system_profiler', 'SPHardwareDataType']).decode('utf-8') diff --git a/topology/partitioning_strategy.py b/topology/partitioning_strategy.py index 329c4ab7..a988bbfd 100644 --- a/topology/partitioning_strategy.py +++ b/topology/partitioning_strategy.py @@ -1,10 +1,22 @@ from abc import ABC, abstractmethod -from typing import List +from typing import List, Optional +from dataclasses import dataclass from inference.shard import Shard from networking.peer_handle import PeerHandle from .topology import Topology +# Partitions shard-space into pieces of contiguous shards, represented by floating point range [start, end) between 0 and 1 +@dataclass +class Partition: + node_id: str + start: float + end: float + +class PartitioningStrategy(ABC): + def node_id(self) -> str: + pass + class PartitioningStrategy(ABC): @abstractmethod - def next_shard(self, current_shard: Shard, topology: Topology, node_stats: dict) -> Shard: + def partition(self, topology: Topology) -> List[Partition]: pass diff --git a/topology/ring_memory_weighted_partitioning_strategy.py b/topology/ring_memory_weighted_partitioning_strategy.py index 6f421098..ad768b24 100644 --- a/topology/ring_memory_weighted_partitioning_strategy.py +++ b/topology/ring_memory_weighted_partitioning_strategy.py @@ -1,27 +1,18 @@ +from typing import List from .partitioning_strategy import PartitioningStrategy from inference.shard import Shard from .topology import Topology +from .partitioning_strategy import Partition class RingMemoryWeightedPartitioningStrategy(PartitioningStrategy): - def next_shard(self, current_shard: Shard, topology: Topology, node_stats: dict) -> Shard: - # Get all nodes from the topology and include the current node + def partition(self, topology: Topology) -> List[Partition]: nodes = list(topology.all_nodes()) - nodes.append((self.id, None, node_stats)) - - # Sort nodes by their IDs nodes.sort(key=lambda x: x[0]) - - # Calculate the total memory of all nodes - total_memory = sum(node[2]['memory'] for node in nodes) - - # Calculate the number of layers to assign to each node proportional to its memory - layers_per_node = {node[0]: (node[2]['memory'] / total_memory) * current_shard.n_layers for node in nodes} - - # Find the successor node - node_ids = [node[0] for node in nodes] - current_index = node_ids.index(self.id) - successor_index = (current_index + 1) % len(node_ids) - successor_id = node_ids[successor_index] - - # Return the Shard calculated for the successor - return Shard(successor_id, layers_per_node[successor_id]) + total_memory = sum(node[1].memory for node in nodes) + partitions = [] + start = 0 + for node in nodes: + end = start + (node[1].memory / total_memory) + partitions.append(Partition(node[0], start, end)) + start = end + return partitions diff --git a/topology/test_ring_memory_weighted_partitioning_strategy.py b/topology/test_ring_memory_weighted_partitioning_strategy.py new file mode 100644 index 00000000..460e050f --- /dev/null +++ b/topology/test_ring_memory_weighted_partitioning_strategy.py @@ -0,0 +1,31 @@ +import unittest +from unittest.mock import MagicMock +from .ring_memory_weighted_partitioning_strategy import RingMemoryWeightedPartitioningStrategy +from .topology import Topology, DeviceCapabilities +from .partitioning_strategy import Partition + +class TestRingMemoryWeightedPartitioningStrategy(unittest.TestCase): + def test_partition(self): + # triangle + # node1 -> node2 -> node3 -> node1 + topology = Topology() + topology.update_node('node1', DeviceCapabilities(model="test1", chip="test1", memory=100)) + topology.update_node('node2', DeviceCapabilities(model="test2", chip="test2", memory=300)) + topology.update_node('node3', DeviceCapabilities(model="test3", chip="test3", memory=600)) + topology.add_edge('node1', 'node2') + topology.add_edge('node2', 'node3') + topology.add_edge('node3', 'node1') + topology.add_edge('node1', 'node3') + + strategy = RingMemoryWeightedPartitioningStrategy() + partitions = strategy.partition(topology) + + self.assertEqual(len(partitions), 3) + self.assertEqual(partitions, [ + Partition('node1', 0.0, 0.1), + Partition('node2', 0.1, 0.4), + Partition('node3', 0.4, 1.0) + ]) + +if __name__ == '__main__': + unittest.main() diff --git a/topology/topology.py b/topology/topology.py index 31b91711..38661e0f 100644 --- a/topology/topology.py +++ b/topology/topology.py @@ -1,12 +1,47 @@ +from .device_capabilities import DeviceCapabilities +from typing import Dict, Set + class Topology: def __init__(self): - self.nodes = {} # Maps node IDs to a tuple of (host, port, stats) + self.nodes: Dict[str, DeviceCapabilities] = {} # Maps node IDs to DeviceCapabilities + self.peer_graph: Dict[str, Set[str]] = {} # Adjacency list representing the graph - def update_node(self, node_id, stats): - self.nodes[node_id] = stats + def update_node(self, node_id: str, device_capabilities: DeviceCapabilities): + self.nodes[node_id] = device_capabilities - def get_node(self, node_id): + def get_node(self, node_id: str) -> DeviceCapabilities: return self.nodes.get(node_id) def all_nodes(self): return self.nodes.items() + + def add_edge(self, node1_id: str, node2_id: str): + if node1_id not in self.peer_graph: + self.peer_graph[node1_id] = set() + if node2_id not in self.peer_graph: + self.peer_graph[node2_id] = set() + self.peer_graph[node1_id].add(node2_id) + self.peer_graph[node2_id].add(node1_id) + + def get_neighbors(self, node_id: str) -> Set[str]: + return self.peer_graph.get(node_id, set()) + + def all_edges(self): + edges = [] + for node, neighbors in self.peer_graph.items(): + for neighbor in neighbors: + if (neighbor, node) not in edges: # Avoid duplicate edges + edges.append((node, neighbor)) + return edges + + def merge(self, other: 'Topology'): + for node_id, capabilities in other.nodes.items(): + self.update_node(node_id, capabilities) + for node_id, neighbors in other.peer_graph.items(): + for neighbor in neighbors: + self.add_edge(node_id, neighbor) + + def __str__(self): + nodes_str = ', '.join(f"{node_id}: {cap}" for node_id, cap in self.nodes.items()) + edges_str = ', '.join(f"{node}: {neighbors}" for node, neighbors in self.peer_graph.items()) + return f"Topology(Nodes: {{{nodes_str}}}, Edges: {{{edges_str}}})"