Updated node refs

763fbf848612f795f305af8c34207c62709b83bc · 2024-12-11 02:54:01 -0800 · Nel Nibcord
Files touched

M README.md
M exo/main.py
M exo/orchestration/__init__.py
M exo/orchestration/test_node.py
Diff

commit 763fbf848612f795f305af8c34207c62709b83bc
Author: Nel Nibcord <blindcrone@tuta.io>
Date:   Wed Dec 11 02:54:01 2024 -0800

    Updated node refs
---
 README.md                      | 4 ++--
 exo/main.py                    | 5 ++---
 exo/orchestration/__init__.py  | 3 +--
 exo/orchestration/test_node.py | 4 ++--
 4 files changed, 7 insertions(+), 9 deletions(-)

diff --git a/README.md b/README.md
index ff4a1bad..9881097a 100644
--- a/README.md
+++ b/README.md
@@ -46,7 +46,7 @@ exo [optimally splits up models](exo/topology/ring_memory_weighted_partitioning_
 
 ### Automatic Device Discovery
 
-exo will [automatically discover](https://github.com/exo-explore/exo/blob/945f90f676182a751d2ad7bcf20987ab7fe0181e/exo/orchestration/standard_node.py#L154) other devices using the best method available. Zero manual configuration.
+exo will [automatically discover](https://github.com/exo-explore/exo/blob/945f90f676182a751d2ad7bcf20987ab7fe0181e/exo/orchestration/node.py#L154) other devices using the best method available. Zero manual configuration.
 
 ### ChatGPT-compatible API
 
@@ -54,7 +54,7 @@ exo provides a [ChatGPT-compatible API](exo/api/chatgpt_api.py) for running mode
 
 ### Device Equality
 
-Unlike other distributed inference frameworks, exo does not use a master-worker architecture. Instead, exo devices [connect p2p](https://github.com/exo-explore/exo/blob/945f90f676182a751d2ad7bcf20987ab7fe0181e/exo/orchestration/standard_node.py#L161). As long as a device is connected somewhere in the network, it can be used to run models.
+Unlike other distributed inference frameworks, exo does not use a master-worker architecture. Instead, exo devices [connect p2p](https://github.com/exo-explore/exo/blob/945f90f676182a751d2ad7bcf20987ab7fe0181e/exo/orchestration/node.py#L161). As long as a device is connected somewhere in the network, it can be used to run models.
 
 Exo supports different [partitioning strategies](exo/topology/partitioning_strategy.py) to split up a model across devices. The default partitioning strategy is [ring memory weighted partitioning](exo/topology/ring_memory_weighted_partitioning_strategy.py). This runs an inference in a ring where each device runs a number of model layers proportional to the memory of the device.
 
diff --git a/exo/main.py b/exo/main.py
index c8bb72e5..3a01f08c 100644
--- a/exo/main.py
+++ b/exo/main.py
@@ -17,7 +17,7 @@ from tqdm.asyncio import tqdm_asyncio
 from exo.train.dataset import load_dataset, iterate_batches, compose
 from exo.networking.manual.manual_discovery import ManualDiscovery
 from exo.networking.manual.network_topology_config import NetworkTopology
-from exo.orchestration.standard_node import StandardNode
+from exo.orchestration.node import Node
 from exo.networking.grpc.grpc_server import GRPCServer
 from exo.networking.udp.udp_discovery import UDPDiscovery
 from exo.networking.tailscale.tailscale_discovery import TailscaleDiscovery
@@ -30,7 +30,6 @@ from exo.helpers import print_yellow_exo, find_available_port, DEBUG, get_system
 from exo.inference.shard import Shard
 from exo.inference.inference_engine import get_inference_engine, InferenceEngine
 from exo.inference.tokenizers import resolve_tokenizer
-from exo.orchestration.node import Node
 from exo.models import build_base_shard, get_repo
 from exo.viz.topology_viz import TopologyViz
 from exo.download.hf.hf_helpers import has_hf_home_read_access, has_hf_home_write_access, get_hf_home, move_models_to_hf
@@ -129,7 +128,7 @@ elif args.discovery_module == "manual":
     raise ValueError(f"--discovery-config-path is required when using manual discovery. Please provide a path to a config json file.")
   discovery = ManualDiscovery(args.discovery_config_path, args.node_id, create_peer_handle=lambda peer_id, address, description, device_capabilities: GRPCPeerHandle(peer_id, address, description, device_capabilities))
 topology_viz = TopologyViz(chatgpt_api_endpoints=chatgpt_api_endpoints, web_chat_urls=web_chat_urls) if not args.disable_tui else None
-node = StandardNode(
+node = Node(
   args.node_id,
   None,
   inference_engine,
diff --git a/exo/orchestration/__init__.py b/exo/orchestration/__init__.py
index 478af537..022a9160 100644
--- a/exo/orchestration/__init__.py
+++ b/exo/orchestration/__init__.py
@@ -1,4 +1,3 @@
 from .node import Node
-from .standard_node import StandardNode
 
-__all__ = ["Node", "StandardNode"]
+__all__ = ["Node"]
diff --git a/exo/orchestration/test_node.py b/exo/orchestration/test_node.py
index 230ef0cf..5ea0ab84 100644
--- a/exo/orchestration/test_node.py
+++ b/exo/orchestration/test_node.py
@@ -2,7 +2,7 @@ import unittest
 from unittest.mock import Mock, AsyncMock
 import numpy as np
 
-from .standard_node import StandardNode
+from .node import Node
 from exo.networking.peer_handle import PeerHandle
 
 
@@ -21,7 +21,7 @@ class TestNode(unittest.IsolatedAsyncioTestCase):
     mock_peer2.id.return_value = "peer2"
     self.mock_discovery.discover_peers = AsyncMock(return_value=[mock_peer1, mock_peer2])
 
-    self.node = StandardNode("test_node", self.mock_server, self.mock_inference_engine, "localhost", 50051, self.mock_discovery)
+    self.node = Node("test_node", self.mock_server, self.mock_inference_engine, "localhost", 50051, self.mock_discovery)
 
   async def asyncSetUp(self):
     await self.node.start()
← 59af2dd5 Do we need casting here? · back to Exo · Moved nodes around c2332e24 →