[object Object]

← back to Exo

move everything under exo module

5bbde22a23706db6022c745007e7f7a6764987a5 · 2024-07-14 21:00:37 -0700 · Alex Cheema

Files touched

Diff

commit 5bbde22a23706db6022c745007e7f7a6764987a5
Author: Alex Cheema <alexcheema123@gmail.com>
Date:   Sun Jul 14 21:00:37 2024 -0700

    move everything under exo module
---
 example_user.py                                              | 10 +++++-----
 example_user_2.py                                            | 10 +++++-----
 {inference => exo}/__init__.py                               |  0
 {inference/mlx => exo/inference}/__init__.py                 |  0
 {inference => exo/inference}/inference_engine.py             |  0
 {inference/mlx/models => exo/inference/mlx}/__init__.py      |  0
 {networking/grpc => exo/inference/mlx/models}/__init__.py    |  0
 {inference => exo/inference}/mlx/models/sharded_llama.py     |  0
 {inference => exo/inference}/mlx/sharded_inference_engine.py |  0
 {inference => exo/inference}/mlx/sharded_model.py            |  0
 {inference => exo/inference}/mlx/sharded_utils.py            |  0
 {inference => exo/inference}/mlx/test_sharded_llama.py       |  6 +++---
 {inference => exo/inference}/mlx/test_sharded_model.py       |  4 ++--
 {inference => exo/inference}/shard.py                        |  0
 {inference => exo/inference}/test_inference_engine.py        |  8 ++++----
 {inference => exo/inference}/tinygrad/inference.py           |  6 +++---
 {inference => exo/inference}/tinygrad/models/llama.py        |  0
 {networking => exo/networking}/__init__.py                   |  0
 {networking => exo/networking}/discovery.py                  |  0
 {topology => exo/networking/grpc}/__init__.py                |  0
 {networking => exo/networking}/grpc/grpc_discovery.py        |  2 +-
 {networking => exo/networking}/grpc/grpc_peer_handle.py      |  6 +++---
 {networking => exo/networking}/grpc/grpc_server.py           |  4 ++--
 {networking => exo/networking}/grpc/node_service.proto       |  0
 {networking => exo/networking}/grpc/node_service_pb2.py      |  0
 {networking => exo/networking}/grpc/node_service_pb2_grpc.py |  0
 {networking => exo/networking}/grpc/test_grpc_discovery.py   |  0
 {networking => exo/networking}/peer_handle.py                |  6 +++---
 {networking => exo/networking}/server.py                     |  0
 {orchestration => exo/orchestration}/__init__.py             |  0
 {orchestration => exo/orchestration}/node.py                 |  4 ++--
 {orchestration => exo/orchestration}/standard_node.py        | 12 ++++++------
 {orchestration => exo/orchestration}/test_node.py            |  2 +-
 exo/topology/__init__.py                                     |  0
 {topology => exo/topology}/device_capabilities.py            |  0
 {topology => exo/topology}/partitioning_strategy.py          |  4 ++--
 .../topology}/ring_memory_weighted_partitioning_strategy.py  |  2 +-
 {topology => exo/topology}/test_device_capabilities.py       |  2 +-
 .../test_ring_memory_weighted_partitioning_strategy.py       |  0
 {topology => exo/topology}/topology.py                       |  0
 main.py                                                      | 12 ++++++------
 main_dynamic.py                                              | 12 ++++++------
 42 files changed, 56 insertions(+), 56 deletions(-)

diff --git a/example_user.py b/example_user.py
index 3700c1f1..61284964 100644
--- a/example_user.py
+++ b/example_user.py
@@ -2,11 +2,11 @@
 # 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 exo.inference.mlx.sharded_utils import get_model_path, load_tokenizer
+from exo.inference.shard import Shard
+from exo.networking.peer_handle import PeerHandle
+from exo.networking.grpc.grpc_peer_handle import GRPCPeerHandle
+from exo.topology.device_capabilities import DeviceCapabilities
 from typing import List
 import asyncio
 import argparse
diff --git a/example_user_2.py b/example_user_2.py
index 01431a78..b54233e3 100644
--- a/example_user_2.py
+++ b/example_user_2.py
@@ -2,11 +2,11 @@
 # 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 exo.inference.mlx.sharded_utils import get_model_path, load_tokenizer
+from exo.inference.shard import Shard
+from exo.networking.peer_handle import PeerHandle
+from exo.networking.grpc.grpc_peer_handle import GRPCPeerHandle
+from exo.topology.device_capabilities import DeviceCapabilities
 from typing import List
 import asyncio
 import argparse
diff --git a/inference/__init__.py b/exo/__init__.py
similarity index 100%
rename from inference/__init__.py
rename to exo/__init__.py
diff --git a/inference/mlx/__init__.py b/exo/inference/__init__.py
similarity index 100%
rename from inference/mlx/__init__.py
rename to exo/inference/__init__.py
diff --git a/inference/inference_engine.py b/exo/inference/inference_engine.py
similarity index 100%
rename from inference/inference_engine.py
rename to exo/inference/inference_engine.py
diff --git a/inference/mlx/models/__init__.py b/exo/inference/mlx/__init__.py
similarity index 100%
rename from inference/mlx/models/__init__.py
rename to exo/inference/mlx/__init__.py
diff --git a/networking/grpc/__init__.py b/exo/inference/mlx/models/__init__.py
similarity index 100%
rename from networking/grpc/__init__.py
rename to exo/inference/mlx/models/__init__.py
diff --git a/inference/mlx/models/sharded_llama.py b/exo/inference/mlx/models/sharded_llama.py
similarity index 100%
rename from inference/mlx/models/sharded_llama.py
rename to exo/inference/mlx/models/sharded_llama.py
diff --git a/inference/mlx/sharded_inference_engine.py b/exo/inference/mlx/sharded_inference_engine.py
similarity index 100%
rename from inference/mlx/sharded_inference_engine.py
rename to exo/inference/mlx/sharded_inference_engine.py
diff --git a/inference/mlx/sharded_model.py b/exo/inference/mlx/sharded_model.py
similarity index 100%
rename from inference/mlx/sharded_model.py
rename to exo/inference/mlx/sharded_model.py
diff --git a/inference/mlx/sharded_utils.py b/exo/inference/mlx/sharded_utils.py
similarity index 100%
rename from inference/mlx/sharded_utils.py
rename to exo/inference/mlx/sharded_utils.py
diff --git a/inference/mlx/test_sharded_llama.py b/exo/inference/mlx/test_sharded_llama.py
similarity index 89%
rename from inference/mlx/test_sharded_llama.py
rename to exo/inference/mlx/test_sharded_llama.py
index b744c182..50f42318 100644
--- a/inference/mlx/test_sharded_llama.py
+++ b/exo/inference/mlx/test_sharded_llama.py
@@ -1,7 +1,7 @@
 import mlx.core as mx
-from inference.mlx.sharded_model import StatefulShardedModel
-from inference.mlx.sharded_utils import load_shard
-from inference.shard import Shard
+from exo.inference.mlx.sharded_model import StatefulShardedModel
+from exo.inference.mlx.sharded_utils import load_shard
+from exo.inference.shard import Shard
 
 shard_full = Shard("llama", 0, 31, 32)
 shard1 = Shard("llama", 0, 12, 32)
diff --git a/inference/mlx/test_sharded_model.py b/exo/inference/mlx/test_sharded_model.py
similarity index 93%
rename from inference/mlx/test_sharded_model.py
rename to exo/inference/mlx/test_sharded_model.py
index 9129424a..8b7b8970 100644
--- a/inference/mlx/test_sharded_model.py
+++ b/exo/inference/mlx/test_sharded_model.py
@@ -1,5 +1,5 @@
-from inference.shard import Shard
-from inference.mlx.sharded_model import StatefulShardedModel
+from exo.inference.shard import Shard
+from exo.inference.mlx.sharded_model import StatefulShardedModel
 import mlx.core as mx
 import mlx.nn as nn
 from typing import Optional
diff --git a/inference/shard.py b/exo/inference/shard.py
similarity index 100%
rename from inference/shard.py
rename to exo/inference/shard.py
diff --git a/inference/test_inference_engine.py b/exo/inference/test_inference_engine.py
similarity index 82%
rename from inference/test_inference_engine.py
rename to exo/inference/test_inference_engine.py
index 123a06de..191fccc5 100644
--- a/inference/test_inference_engine.py
+++ b/exo/inference/test_inference_engine.py
@@ -1,7 +1,7 @@
-from inference.mlx.sharded_inference_engine import MLXDynamicShardInferenceEngine
-from inference.inference_engine import InferenceEngine
-from inference.shard import Shard
-from inference.tinygrad.inference import TinygradDynamicShardInferenceEngine
+from exo.inference.mlx.sharded_inference_engine import MLXDynamicShardInferenceEngine
+from exo.inference.inference_engine import InferenceEngine
+from exo.inference.shard import Shard
+from exo.inference.tinygrad.inference import TinygradDynamicShardInferenceEngine
 import numpy as np
 
 # An inference engine should work the same for any number of Shards, as long as the Shards are continuous.
diff --git a/inference/tinygrad/inference.py b/exo/inference/tinygrad/inference.py
similarity index 97%
rename from inference/tinygrad/inference.py
rename to exo/inference/tinygrad/inference.py
index c1ff10e0..ae9712e5 100644
--- a/inference/tinygrad/inference.py
+++ b/exo/inference/tinygrad/inference.py
@@ -4,12 +4,12 @@ from typing import List
 import json, argparse, random, time
 import tiktoken
 from tiktoken.load import load_tiktoken_bpe
-from inference.tinygrad.models.llama import Transformer, convert_from_huggingface, fix_bf16
+from exo.inference.tinygrad.models.llama import Transformer, convert_from_huggingface, fix_bf16
 from tinygrad.nn.state import safe_load, torch_load, load_state_dict, get_parameters
 from tinygrad import Tensor, dtypes, nn, Context, Device, GlobalCounters
 from tinygrad.helpers import Profiling, Timing, DEBUG, colored, fetch, tqdm
-from inference.shard import Shard
-from inference.inference_engine import InferenceEngine
+from exo.inference.shard import Shard
+from exo.inference.inference_engine import InferenceEngine
 import numpy as np
 
 MODEL_PARAMS = {
diff --git a/inference/tinygrad/models/llama.py b/exo/inference/tinygrad/models/llama.py
similarity index 100%
rename from inference/tinygrad/models/llama.py
rename to exo/inference/tinygrad/models/llama.py
diff --git a/networking/__init__.py b/exo/networking/__init__.py
similarity index 100%
rename from networking/__init__.py
rename to exo/networking/__init__.py
diff --git a/networking/discovery.py b/exo/networking/discovery.py
similarity index 100%
rename from networking/discovery.py
rename to exo/networking/discovery.py
diff --git a/topology/__init__.py b/exo/networking/grpc/__init__.py
similarity index 100%
rename from topology/__init__.py
rename to exo/networking/grpc/__init__.py
diff --git a/networking/grpc/grpc_discovery.py b/exo/networking/grpc/grpc_discovery.py
similarity index 98%
rename from networking/grpc/grpc_discovery.py
rename to exo/networking/grpc/grpc_discovery.py
index 926a7097..d19c9e23 100644
--- a/networking/grpc/grpc_discovery.py
+++ b/exo/networking/grpc/grpc_discovery.py
@@ -6,7 +6,7 @@ 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, device_capabilities
+from exo.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, device_capabilities=None):
diff --git a/networking/grpc/grpc_peer_handle.py b/exo/networking/grpc/grpc_peer_handle.py
similarity index 96%
rename from networking/grpc/grpc_peer_handle.py
rename to exo/networking/grpc/grpc_peer_handle.py
index 3a56604c..3ad44a95 100644
--- a/networking/grpc/grpc_peer_handle.py
+++ b/exo/networking/grpc/grpc_peer_handle.py
@@ -7,9 +7,9 @@ from . import node_service_pb2
 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
+from exo.inference.shard import Shard
+from exo.topology.topology import Topology
+from exo.topology.device_capabilities import DeviceCapabilities
 
 class GRPCPeerHandle(PeerHandle):
     def __init__(self, id: str, address: str, device_capabilities: DeviceCapabilities):
diff --git a/networking/grpc/grpc_server.py b/exo/networking/grpc/grpc_server.py
similarity index 98%
rename from networking/grpc/grpc_server.py
rename to exo/networking/grpc/grpc_server.py
index 9647c217..9fdc4785 100644
--- a/networking/grpc/grpc_server.py
+++ b/exo/networking/grpc/grpc_server.py
@@ -4,9 +4,9 @@ import numpy as np
 
 from . import node_service_pb2
 from . import node_service_pb2_grpc
-from inference.shard import Shard
+from exo.inference.shard import Shard
 
-from orchestration import Node
+from exo.orchestration import Node
 
 import uuid
 
diff --git a/networking/grpc/node_service.proto b/exo/networking/grpc/node_service.proto
similarity index 100%
rename from networking/grpc/node_service.proto
rename to exo/networking/grpc/node_service.proto
diff --git a/networking/grpc/node_service_pb2.py b/exo/networking/grpc/node_service_pb2.py
similarity index 100%
rename from networking/grpc/node_service_pb2.py
rename to exo/networking/grpc/node_service_pb2.py
diff --git a/networking/grpc/node_service_pb2_grpc.py b/exo/networking/grpc/node_service_pb2_grpc.py
similarity index 100%
rename from networking/grpc/node_service_pb2_grpc.py
rename to exo/networking/grpc/node_service_pb2_grpc.py
diff --git a/networking/grpc/test_grpc_discovery.py b/exo/networking/grpc/test_grpc_discovery.py
similarity index 100%
rename from networking/grpc/test_grpc_discovery.py
rename to exo/networking/grpc/test_grpc_discovery.py
diff --git a/networking/peer_handle.py b/exo/networking/peer_handle.py
similarity index 88%
rename from networking/peer_handle.py
rename to exo/networking/peer_handle.py
index 78a2e86f..c3c2c6c5 100644
--- a/networking/peer_handle.py
+++ b/exo/networking/peer_handle.py
@@ -1,9 +1,9 @@
 from abc import ABC, abstractmethod
 from typing import Optional, Tuple
 import numpy as np
-from inference.shard import Shard
-from topology.device_capabilities import DeviceCapabilities
-from topology.topology import Topology
+from exo.inference.shard import Shard
+from exo.topology.device_capabilities import DeviceCapabilities
+from exo.topology.topology import Topology
 
 class PeerHandle(ABC):
     @abstractmethod
diff --git a/networking/server.py b/exo/networking/server.py
similarity index 100%
rename from networking/server.py
rename to exo/networking/server.py
diff --git a/orchestration/__init__.py b/exo/orchestration/__init__.py
similarity index 100%
rename from orchestration/__init__.py
rename to exo/orchestration/__init__.py
diff --git a/orchestration/node.py b/exo/orchestration/node.py
similarity index 91%
rename from orchestration/node.py
rename to exo/orchestration/node.py
index 007321a0..98415923 100644
--- a/orchestration/node.py
+++ b/exo/orchestration/node.py
@@ -1,8 +1,8 @@
 from typing import Optional, Tuple
 import numpy as np
 from abc import ABC, abstractmethod
-from inference.shard import Shard
-from topology.topology import Topology
+from exo.inference.shard import Shard
+from exo.topology.topology import Topology
 
 class Node(ABC):
     @abstractmethod
diff --git a/orchestration/standard_node.py b/exo/orchestration/standard_node.py
similarity index 96%
rename from orchestration/standard_node.py
rename to exo/orchestration/standard_node.py
index 2e402ea8..ae10ed2e 100644
--- a/orchestration/standard_node.py
+++ b/exo/orchestration/standard_node.py
@@ -1,12 +1,12 @@
 from typing import List, Dict, Optional, Callable, Tuple
 import numpy as np
-from networking import Discovery, PeerHandle, Server
-from inference.inference_engine import InferenceEngine, Shard
+from exo.networking import Discovery, PeerHandle, Server
+from exo.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
+from exo.topology.topology import Topology
+from exo.topology.device_capabilities import device_capabilities
+from exo.topology.partitioning_strategy import PartitioningStrategy
+from exo.topology.partitioning_strategy import Partition
 import asyncio
 import uuid
 
diff --git a/orchestration/test_node.py b/exo/orchestration/test_node.py
similarity index 97%
rename from orchestration/test_node.py
rename to exo/orchestration/test_node.py
index 5713b3ab..c13a108c 100644
--- a/orchestration/test_node.py
+++ b/exo/orchestration/test_node.py
@@ -3,7 +3,7 @@ from unittest.mock import Mock, AsyncMock
 import numpy as np
 
 from .standard_node import StandardNode
-from networking.peer_handle import PeerHandle
+from exo.networking.peer_handle import PeerHandle
 
 class TestNode(unittest.IsolatedAsyncioTestCase):
     def setUp(self):
diff --git a/exo/topology/__init__.py b/exo/topology/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/topology/device_capabilities.py b/exo/topology/device_capabilities.py
similarity index 100%
rename from topology/device_capabilities.py
rename to exo/topology/device_capabilities.py
diff --git a/topology/partitioning_strategy.py b/exo/topology/partitioning_strategy.py
similarity index 86%
rename from topology/partitioning_strategy.py
rename to exo/topology/partitioning_strategy.py
index a988bbfd..52140198 100644
--- a/topology/partitioning_strategy.py
+++ b/exo/topology/partitioning_strategy.py
@@ -1,8 +1,8 @@
 from abc import ABC, abstractmethod
 from typing import List, Optional
 from dataclasses import dataclass
-from inference.shard import Shard
-from networking.peer_handle import PeerHandle
+from exo.inference.shard import Shard
+from exo.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
diff --git a/topology/ring_memory_weighted_partitioning_strategy.py b/exo/topology/ring_memory_weighted_partitioning_strategy.py
similarity index 94%
rename from topology/ring_memory_weighted_partitioning_strategy.py
rename to exo/topology/ring_memory_weighted_partitioning_strategy.py
index ad768b24..117e3a0d 100644
--- a/topology/ring_memory_weighted_partitioning_strategy.py
+++ b/exo/topology/ring_memory_weighted_partitioning_strategy.py
@@ -1,6 +1,6 @@
 from typing import List
 from .partitioning_strategy import PartitioningStrategy
-from inference.shard import Shard
+from exo.inference.shard import Shard
 from .topology import Topology
 from .partitioning_strategy import Partition
 
diff --git a/topology/test_device_capabilities.py b/exo/topology/test_device_capabilities.py
similarity index 96%
rename from topology/test_device_capabilities.py
rename to exo/topology/test_device_capabilities.py
index 381a379a..06cb3f53 100644
--- a/topology/test_device_capabilities.py
+++ b/exo/topology/test_device_capabilities.py
@@ -1,6 +1,6 @@
 import unittest
 from unittest.mock import patch
-from topology.device_capabilities import mac_device_capabilities, DeviceCapabilities
+from exo.topology.device_capabilities import mac_device_capabilities, DeviceCapabilities
 
 class TestMacDeviceCapabilities(unittest.TestCase):
     @patch('subprocess.check_output')
diff --git a/topology/test_ring_memory_weighted_partitioning_strategy.py b/exo/topology/test_ring_memory_weighted_partitioning_strategy.py
similarity index 100%
rename from topology/test_ring_memory_weighted_partitioning_strategy.py
rename to exo/topology/test_ring_memory_weighted_partitioning_strategy.py
diff --git a/topology/topology.py b/exo/topology/topology.py
similarity index 100%
rename from topology/topology.py
rename to exo/topology/topology.py
diff --git a/main.py b/main.py
index dfeff3cd..f78fcc60 100644
--- a/main.py
+++ b/main.py
@@ -3,12 +3,12 @@ import asyncio
 import signal
 import mlx.core as mx
 import mlx.nn as nn
-from orchestration.standard_node import StandardNode
-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
+from exo.orchestration.standard_node import StandardNode
+from exo.networking.grpc.grpc_server import GRPCServer
+from exo.inference.mlx.sharded_inference_engine import MLXFixedShardInferenceEngine
+from exo.inference.shard import Shard
+from exo.networking.grpc.grpc_discovery import GRPCDiscovery
+from exo.topology.ring_memory_weighted_partitioning_strategy import RingMemoryWeightedPartitioningStrategy
 
 # parse args
 parser = argparse.ArgumentParser(description="Initialize GRPC Discovery")
diff --git a/main_dynamic.py b/main_dynamic.py
index c5bf7bf0..5e337d4e 100644
--- a/main_dynamic.py
+++ b/main_dynamic.py
@@ -4,12 +4,12 @@ import signal
 import mlx.core as mx
 import mlx.nn as nn
 from typing import List
-from orchestration.standard_node import StandardNode
-from networking.grpc.grpc_server import GRPCServer
-from inference.mlx.sharded_inference_engine import MLXDynamicShardInferenceEngine
-from inference.shard import Shard
-from networking.grpc.grpc_discovery import GRPCDiscovery
-from topology.ring_memory_weighted_partitioning_strategy import RingMemoryWeightedPartitioningStrategy
+from exo.orchestration.standard_node import StandardNode
+from exo.networking.grpc.grpc_server import GRPCServer
+from exo.inference.mlx.sharded_inference_engine import MLXDynamicShardInferenceEngine
+from exo.inference.shard import Shard
+from exo.networking.grpc.grpc_discovery import GRPCDiscovery
+from exo.topology.ring_memory_weighted_partitioning_strategy import RingMemoryWeightedPartitioningStrategy
 
 # parse args
 parser = argparse.ArgumentParser(description="Initialize GRPC Discovery")

← c851644a update requirements, specify exact versions  ·  back to Exo  ·  main rename 32f2e36f →