main rename

32f2e36fd33fd49a02ed04f09bd00728930d12e5 · 2024-07-14 21:01:28 -0700 · Alex Cheema
Files touched

M main.py
R077 main_dynamic.py main_static.py
Diff

commit 32f2e36fd33fd49a02ed04f09bd00728930d12e5
Author: Alex Cheema <alexcheema123@gmail.com>
Date:   Sun Jul 14 21:01:28 2024 -0700

    main rename
---
 main.py                           | 20 ++++++++++++--------
 main_dynamic.py => main_static.py | 20 ++++++++------------
 2 files changed, 20 insertions(+), 20 deletions(-)

diff --git a/main.py b/main.py
index f78fcc60..5e337d4e 100644
--- a/main.py
+++ b/main.py
@@ -3,9 +3,10 @@ import asyncio
 import signal
 import mlx.core as mx
 import mlx.nn as nn
+from typing import List
 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.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
@@ -17,16 +18,16 @@ parser.add_argument("--node-host", type=str, default="0.0.0.0", help="Node host"
 parser.add_argument("--node-port", type=int, default=8080, help="Node port")
 parser.add_argument("--listen-port", type=int, default=5678, help="Listening port for discovery")
 parser.add_argument("--broadcast-port", type=int, default=5678, help="Broadcast port for discovery")
-parser.add_argument("--model-id", type=str, default="mlx-community/Meta-Llama-3-8B-Instruct-4bit", help="Path to the model")
-parser.add_argument("--n-layers", type=int, default=32, help="Number of layers in the model")
-parser.add_argument("--start-layer", type=int, default=0, help="Start layer index")
-parser.add_argument("--end-layer", type=int, default=31, help="End layer index")
 parser.add_argument("--wait-for-peers", type=int, default=0, help="Number of peers to wait to connect to before starting")
 args = parser.parse_args()
 
-inference_engine = MLXFixedShardInferenceEngine(args.model_id, shard=Shard(model_id=args.model_id, n_layers=args.n_layers, start_layer=args.start_layer, end_layer=args.end_layer))
+
+inference_engine = MLXDynamicShardInferenceEngine()
+def on_token(tokens: List[int]):
+    if inference_engine.tokenizer:
+        print(inference_engine.tokenizer.decode(tokens))
 discovery = GRPCDiscovery(args.node_id, args.node_port, args.listen_port, args.broadcast_port)
-node = StandardNode(args.node_id, None, inference_engine, discovery, partitioning_strategy=RingMemoryWeightedPartitioningStrategy())
+node = StandardNode(args.node_id, None, inference_engine, discovery, partitioning_strategy=RingMemoryWeightedPartitioningStrategy(), on_token=on_token)
 server = GRPCServer(node, args.node_host, args.node_port)
 node.server = server
 
@@ -38,7 +39,7 @@ async def shutdown(signal, loop):
     [task.cancel() for task in server_tasks]
     print(f"Cancelling {len(server_tasks)} outstanding tasks")
     await asyncio.gather(*server_tasks, return_exceptions=True)
-    await server.shutdown()
+    await server.stop()
     loop.stop()
 
 async def main():
@@ -60,5 +61,8 @@ if __name__ == "__main__":
     asyncio.set_event_loop(loop)
     try:
         loop.run_until_complete(main())
+    except KeyboardInterrupt:
+        print("Received keyboard interrupt. Shutting down...")
     finally:
+        loop.run_until_complete(shutdown(signal.SIGTERM, loop))
         loop.close()
diff --git a/main_dynamic.py b/main_static.py
similarity index 77%
rename from main_dynamic.py
rename to main_static.py
index 5e337d4e..f78fcc60 100644
--- a/main_dynamic.py
+++ b/main_static.py
@@ -3,10 +3,9 @@ import asyncio
 import signal
 import mlx.core as mx
 import mlx.nn as nn
-from typing import List
 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.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
@@ -18,16 +17,16 @@ parser.add_argument("--node-host", type=str, default="0.0.0.0", help="Node host"
 parser.add_argument("--node-port", type=int, default=8080, help="Node port")
 parser.add_argument("--listen-port", type=int, default=5678, help="Listening port for discovery")
 parser.add_argument("--broadcast-port", type=int, default=5678, help="Broadcast port for discovery")
+parser.add_argument("--model-id", type=str, default="mlx-community/Meta-Llama-3-8B-Instruct-4bit", help="Path to the model")
+parser.add_argument("--n-layers", type=int, default=32, help="Number of layers in the model")
+parser.add_argument("--start-layer", type=int, default=0, help="Start layer index")
+parser.add_argument("--end-layer", type=int, default=31, help="End layer index")
 parser.add_argument("--wait-for-peers", type=int, default=0, help="Number of peers to wait to connect to before starting")
 args = parser.parse_args()
 
-
-inference_engine = MLXDynamicShardInferenceEngine()
-def on_token(tokens: List[int]):
-    if inference_engine.tokenizer:
-        print(inference_engine.tokenizer.decode(tokens))
+inference_engine = MLXFixedShardInferenceEngine(args.model_id, shard=Shard(model_id=args.model_id, n_layers=args.n_layers, start_layer=args.start_layer, end_layer=args.end_layer))
 discovery = GRPCDiscovery(args.node_id, args.node_port, args.listen_port, args.broadcast_port)
-node = StandardNode(args.node_id, None, inference_engine, discovery, partitioning_strategy=RingMemoryWeightedPartitioningStrategy(), on_token=on_token)
+node = StandardNode(args.node_id, None, inference_engine, discovery, partitioning_strategy=RingMemoryWeightedPartitioningStrategy())
 server = GRPCServer(node, args.node_host, args.node_port)
 node.server = server
 
@@ -39,7 +38,7 @@ async def shutdown(signal, loop):
     [task.cancel() for task in server_tasks]
     print(f"Cancelling {len(server_tasks)} outstanding tasks")
     await asyncio.gather(*server_tasks, return_exceptions=True)
-    await server.stop()
+    await server.shutdown()
     loop.stop()
 
 async def main():
@@ -61,8 +60,5 @@ if __name__ == "__main__":
     asyncio.set_event_loop(loop)
     try:
         loop.run_until_complete(main())
-    except KeyboardInterrupt:
-        print("Received keyboard interrupt. Shutting down...")
     finally:
-        loop.run_until_complete(shutdown(signal.SIGTERM, loop))
         loop.close()
← 5bbde22a move everything under exo module · back to Exo · slightly nicer readme 8ff3e263 →