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")