← back to Exo
fix to creating engines
cd4d324a5e538ed387946dcbe78948e1dced3c74 · 2024-10-22 23:52:49 -0700 · josh
Files touched
M exo/main.pyM exo/networking/udp/udp_discovery.pyM exo/orchestration/standard_node.py
Diff
commit cd4d324a5e538ed387946dcbe78948e1dced3c74
Author: josh <eyasunigussie@Eyasus-MacBook-Air.local>
Date: Tue Oct 22 23:52:49 2024 -0700
fix to creating engines
---
exo/main.py | 3 ++-
exo/networking/udp/udp_discovery.py | 6 +----
exo/orchestration/standard_node.py | 52 +++++++++++++++++++++----------------
3 files changed, 33 insertions(+), 28 deletions(-)
diff --git a/exo/main.py b/exo/main.py
index e3171ad1..30a56473 100644
--- a/exo/main.py
+++ b/exo/main.py
@@ -95,7 +95,8 @@ node = StandardNode(
discovery,
partitioning_strategy=RingMemoryWeightedPartitioningStrategy(),
max_generate_tokens=args.max_generate_tokens,
- topology_viz=topology_viz
+ topology_viz=topology_viz,
+ shard_downloader=shard_downloader
)
server = GRPCServer(node, args.node_host, args.node_port)
node.server = server
diff --git a/exo/networking/udp/udp_discovery.py b/exo/networking/udp/udp_discovery.py
index eb2a90c3..adeb8316 100644
--- a/exo/networking/udp/udp_discovery.py
+++ b/exo/networking/udp/udp_discovery.py
@@ -159,11 +159,7 @@ 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)
- if message["type"] == "supported_inference_engines":
- logger.error(f'supported_inference_engines: {message}')
- 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),
local_addr=("0.0.0.0", self.listen_port))
diff --git a/exo/orchestration/standard_node.py b/exo/orchestration/standard_node.py
index f81cbfe1..0cf172d5 100644
--- a/exo/orchestration/standard_node.py
+++ b/exo/orchestration/standard_node.py
@@ -4,7 +4,7 @@ import asyncio
import uuid
import time
import traceback
-from typing import List, Dict, Optional, Tuple, Union
+from typing import List, Dict, Optional, Tuple, Union, Set
from exo.networking import Discovery, PeerHandle, Server
from exo.inference.inference_engine import InferenceEngine, Shard
from .node import Node
@@ -15,6 +15,8 @@ from exo import DEBUG
from exo.helpers import AsyncCallbackSystem
from exo.viz.topology_viz import TopologyViz
from exo.download.hf.hf_helpers import RepoProgressEvent
+from exo.inference.inference_engine import get_inference_engine, InferenceEngine
+from exo.download.hf.hf_shard_download import HFShardDownloader
import logging
logger = logging.getLogger(__name__)
@@ -29,6 +31,7 @@ class StandardNode(Node):
partitioning_strategy: PartitioningStrategy = None,
max_generate_tokens: int = 1024,
topology_viz: Optional[TopologyViz] = None,
+ shard_downloader: Optional[HFShardDownloader] = None,
):
self.id = _id
self.inference_engine = inference_engine
@@ -45,7 +48,8 @@ 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] = []
+ self.topology_inference_engines_pool: List[List[str]] = []
+ self.shard_downloader = shard_downloader
async def start(self, wait_for_peers: int = 0) -> None:
await self.server.start()
@@ -62,6 +66,14 @@ class StandardNode(Node):
def on_node_status(self, request_id, opaque_status):
try:
status_data = json.loads(opaque_status)
+ if status_data.get("type", "") == "supported_inference_engines":
+ logger.error(f'supported_inference_engines: {status_data}')
+ logger.error('inside on_status data')
+ node_id = status_data.get("node_id")
+ engines = status_data.get("engines", [])
+ logger.error(f'engines: {engines}')
+ self.topology_inference_engines_pool.append(engines)
+ logger.error(f'topology_inference_engines_pool: {self.topology_inference_engines_pool}')
if status_data.get("type", "") == "node_status":
if status_data.get("status", "").startswith("start_"):
self.current_topology.active_node_id = status_data.get("node_id")
@@ -82,22 +94,24 @@ class StandardNode(Node):
def get_supported_inference_engines(self):
supported_engine_names = []
if self.inference_engine.__class__.__name__ == 'MLXDynamicShardInferenceEngine':
- supported_engine_names.extend(['mlx', 'tinygrad'])
+ supported_engine_names.append('mlx')
+ supported_engine_names.append('tinygrad')
else:
- supported_engine_names.append('tinygrad')
+ supported_engine_names.append('tinygrad')
return supported_engine_names
- async def broadcast_supported_engines(self, supported_engines: List):
+ async def broadcast_supported_engines(self, supported_engines_names: List[str]):
status_message = json.dumps({
- "type": "supported_inference_engines",
- "node_id": self.id,
- "engines": supported_engines
+ "type": "supported_inference_engines",
+ "node_id": self.id,
+ "engines": supported_engines_names
})
logger.error(f'broadcast_supported_engines: {status_message}')
await self.broadcast_opaque_status("", status_message)
logger.error(f'broadcast_supported_engines: done')
- def get_topology_inference_engines(self) -> List[str]:
+ def get_topology_inference_engines(self) -> List[List[str]]:
+ logger.error(f'topology_inference_engines_pool: {self.topology_inference_engines_pool}')
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]:
@@ -365,20 +379,15 @@ class StandardNode(Node):
async def select_best_inference_engine(self):
supported_engines = self.get_supported_inference_engines()
await self.broadcast_supported_engines(supported_engines)
- logger.error('ABOVE and ALL')
- logger.error("Topology inference engines pool: %s", self.get_topology_inference_engines())
+ logger.error("Topology inference engines pool: %s", len(self.get_topology_inference_engines()))
logger.error(f'result:{self.get_topology_inference_engines()}')
- if self.get_topology_inference_engines():
- logger.info("Topology inference engines pool: %s", self.get_topology_inference_engines())
- topology_inference_engines_pool = self.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"
+ if len(self.get_topology_inference_engines()):
+ if any(len(engines) == 1 and "tinygrad" in engines for engines in self.get_topology_inference_engines()):
+ logger.info("Found node with only tinygrad, using tinygrad on all nodes")
+ self.inference_engine = get_inference_engine("tinygrad", self.shard_downloader)
else:
- raise ValueError("No compatible inference engine found across all nodes")
-
+ logger.info("all nodes can use mlx, using mlx for inference")
+ self.inference_engine = get_inference_engine("mlx", self.shard_downloader)
async def periodic_topology_collection(self, interval: int):
while True:
@@ -462,7 +471,6 @@ class StandardNode(Node):
async def send_status_to_peer(peer):
try:
- status_dict = json.loads(status)
await asyncio.wait_for(peer.send_opaque_status(request_id, status), timeout=15.0)
except asyncio.TimeoutError:
print(f"Timeout sending opaque status to {peer.id()}")
← da523573 fixed errors
·
back to Exo
·
removed logger 57c62c2c →