collect global topology with local peer visibility, ring memory weighted partitioning strategy

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

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

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

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)

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

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}")

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

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)

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<string, DeviceCapabilities> nodes = 1;
+  map<string, Peers> peer_graph = 2;
+}
+
+message Peers {
+    repeated string peer_ids = 1;
+}
+
+message DeviceCapabilities {
+  string model = 1;
+  string chip = 2;
+  int32 memory = 3;
+}
+
 message Empty {}

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)

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)

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

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

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}")
-
-        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")
+        result = await self.inference_engine.infer_tensor(shard, tensor)
+        print(f"Got result from tensor: {len(tensor)}. Result: {result}")
 
-            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

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

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

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

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

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}}})"