← back to Exo
collect global topology with local peer visibility, ring memory weighted partitioning strategy
36b8456798d2ed9cb80ce662a72d0ce3be77902a · 2024-06-25 12:32:03 +0100 · Alex Cheema
Files touched
M example_user.pyA example_user_2.pyM inference/inference_engine.pyM inference/mlx/sharded_inference_engine.pyM main.pyM networking/grpc/grpc_discovery.pyM networking/grpc/grpc_peer_handle.pyM networking/grpc/grpc_server.pyM networking/grpc/node_service.protoM networking/grpc/node_service_pb2.pyM networking/grpc/node_service_pb2_grpc.pyM networking/peer_handle.pyM orchestration/node.pyM orchestration/standard_node.pyM topology/device_capabilities.pyM topology/partitioning_strategy.pyM topology/ring_memory_weighted_partitioning_strategy.pyA topology/test_ring_memory_weighted_partitioning_strategy.pyM topology/topology.py
Diff
commit 36b8456798d2ed9cb80ce662a72d0ce3be77902a
Author: Alex Cheema <alexcheema123@gmail.com>
Date: Tue Jun 25 12:32:03 2024 +0100
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 ++++-
.../ring_memory_weighted_partitioning_strategy.py | 31 ++++------
...t_ring_memory_weighted_partitioning_strategy.py | 31 ++++++++++
topology/topology.py | 43 +++++++++++--
19 files changed, 384 insertions(+), 68 deletions(-)
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<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 {}
\ 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}")
-
- 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
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}}})"
← 3a66a0a4 add requirements.txt
·
back to Exo
·
dynamically assign shards to nodes deterministically weighte 445eda15 →