← back to Exo
added error
a03f3a2a5cf924ec997f90eca63497ec261ff6c2 · 2024-10-22 18:59:59 -0700 · josh
Files touched
M exo/main.pyM exo/networking/udp/udp_discovery.pyM exo/orchestration/standard_node.py
Diff
commit a03f3a2a5cf924ec997f90eca63497ec261ff6c2
Author: josh <eyasunigussie@Eyasus-MacBook-Air.local>
Date: Tue Oct 22 18:59:59 2024 -0700
added error
---
exo/main.py | 55 +++++++++++++++----------------------
exo/networking/udp/udp_discovery.py | 4 +++
exo/orchestration/standard_node.py | 21 ++++++++------
3 files changed, 39 insertions(+), 41 deletions(-)
diff --git a/exo/main.py b/exo/main.py
index 7faf63c1..b3761b8a 100644
--- a/exo/main.py
+++ b/exo/main.py
@@ -2,6 +2,8 @@ import argparse
import asyncio
import signal
import json
+import logging
+import seqlog
import time
import traceback
import uuid
@@ -51,6 +53,13 @@ args = parser.parse_args()
print_yellow_exo()
+seqlog.log_to_seq(
+ server_url = 'http://localhost:5341',
+ level = logging.DEBUG
+)
+logger = logging.getLogger('exo_logger')
+logger.setLevel(logging.DEBUG)
+
system_info = get_system_info()
print(f"Detected system: {system_info}")
@@ -145,37 +154,18 @@ async def shutdown(signal, loop):
loop.stop()
async def select_best_inference_engine(node: StandardNode):
- supported_engines = set(node.get_supported_inference_engines())
+ supported_engines = node.get_supported_inference_engines()
await node.broadcast_supported_engines(supported_engines)
- num_peers = len(node.peers)
- all_peers_responded = asyncio.Event()
- def check_all_responses():
- if len(node.received_opaque_statuses) >= num_peers:
- all_peers_responded.set()
- node.on_opaque_status.register("engine_selection").on_next(lambda *args: check_all_responses())
- try:
- await asyncio.wait_for(all_peers_responded.wait(), timeout=10.0)
- except asyncio.TimeoutError:
- print("Timed out waiting for peer nodes to respond.")
- node.on_opaque_status.unregister("engine_selection")
- all_supported_engines = [supported_engines]
- for peer_id, status in node.received_opaque_statuses:
- try:
- status_data = json.loads(status)
- if status_data.get("type") == "supported_inference_engines":
- all_supported_engines.append(set(status_data.get("engines", [])))
- except json.JSONDecodeError:
- continue
- if any("tinygrad" in engines and len(engines) == 1 for engines in all_supported_engines):
- return "tinygrad"
- common_engine_across_peers = set.intersection(*all_supported_engines)
- if "mlx" in common_engine_across_peers:
- print('mlx')
- return "mlx"
- elif "tinygrad" in common_engine_across_peers:
- return "tinygrad"
- else:
- raise ValueError("No compatible inference engine found across all nodes")
+
+ if node.get_topology_inference_engines():
+ topology_inference_engines_pool = node.get_topology_inference_engines()
+ if any("tinygrad" in engines and len(engines) == 1 for engines in topology_inference_engines_pool):
+ return "tinygrad"
+ common_engine_across_peers = set.intersection(*topology_inference_engines_pool)
+ if "mlx" in common_engine_across_peers:
+ return "mlx"
+ else:
+ raise ValueError("No compatible inference engine found across all nodes")
async def run_model_cli(node: Node, inference_engine: InferenceEngine, model_name: str, prompt: str):
@@ -217,9 +207,8 @@ async def main():
loop.add_signal_handler(s, handle_exit)
await node.start(wait_for_peers=args.wait_for_peers)
- if len(node.peers) > 1:
- compatible_engine = await select_best_inference_engine(node)
- node.inference_engine = get_inference_engine(compatible_engine, shard_downloader)
+ await select_best_inference_engine(node)
+
if args.command == "run" or args.run_model:
model_name = args.model_name or args.run_model
diff --git a/exo/networking/udp/udp_discovery.py b/exo/networking/udp/udp_discovery.py
index adeb8316..00197228 100644
--- a/exo/networking/udp/udp_discovery.py
+++ b/exo/networking/udp/udp_discovery.py
@@ -159,6 +159,10 @@ class UDPDiscovery(Discovery):
if peer_id in self.known_peers: del self.known_peers[peer_id]
return
if peer_id in self.known_peers: self.known_peers[peer_id] = (self.known_peers[peer_id][0], self.known_peers[peer_id][1], time.time(), peer_prio)
+ elif message["type"] == "supported_inference_engines":
+ peer_id = message["node_id"]
+ engines = message["engines"]
+ if peer_id in self.known_peers: self.known_peers[peer_id][0].topology_inference_engines_pool.append(engines)
async def task_listen_for_peers(self):
await asyncio.get_event_loop().create_datagram_endpoint(lambda: ListenProtocol(self.on_listen_message),
diff --git a/exo/orchestration/standard_node.py b/exo/orchestration/standard_node.py
index 54fba341..aef5bc7e 100644
--- a/exo/orchestration/standard_node.py
+++ b/exo/orchestration/standard_node.py
@@ -43,6 +43,7 @@ class StandardNode(Node):
self._on_opaque_status = AsyncCallbackSystem[str, Tuple[str, str]]()
self._on_opaque_status.register("node_status").on_next(self.on_node_status)
self.node_download_progress: Dict[str, RepoProgressEvent] = {}
+ self.topology_inference_engines_pool: List[str] = []
async def start(self, wait_for_peers: int = 0) -> None:
await self.server.start()
@@ -77,19 +78,23 @@ class StandardNode(Node):
if DEBUG >= 1: traceback.print_exc()
def get_supported_inference_engines(self):
- supported_engines = []
- if self.inferenceEngine == 'mlx':
- supported_engines.extend('mlx', 'tinygrad')
+ supported_engine_names = []
+ if self.inference_engine.__class__.__name__ == 'MLXDynamicShardInferenceEngine':
+ supported_engine_names.extend(['mlx', 'tinygrad'])
else:
- supported_engines.append('tinygrad')
- return supported_engines
+ supported_engine_names.append('tinygrad')
+ return supported_engine_names
async def broadcast_supported_engines(self, supported_engines: List):
- await self.broadcast_opaque_status("", json.dumps({
+ status_message = json.dumps({
"type": "supported_inference_engines",
- "node_id": self.id,
+ "node_id": self.id,
"engines": supported_engines
- }))
+ })
+ await self.broadcast_opaque_status("", status_message)
+
+ def get_topology_inference_engines(self) -> List[str]:
+ return self.topology_inference_engines_pool
async def process_prompt(self, base_shard: Shard, prompt: str, image_str: Optional[str] = None, request_id: Optional[str] = None, inference_state: Optional[str] = None) -> Optional[np.ndarray]:
shard = self.get_current_shard(base_shard)
← ae47fe11 Moving conditonal apple silicon logic to setup.py
·
back to Exo
·
added logger f4a5562c →