Merge branch 'main' of github.com:risingsunomi/exo-nvidia

Author: risingsunomi

.circleci/config.yml

@@ -17,11 +17,11 @@ commands:
             source env/bin/activate
 
             # Start first instance
-            HF_HOME="$(pwd)/.hf_cache_node1" DEBUG_DISCOVERY=7 DEBUG=7 python3 main.py --inference-engine <<parameters.inference_engine>> --node-id "node1" --listen-port 5678 --broadcast-port 5679 --chatgpt-api-port 8000 --chatgpt-api-response-timeout 900 2>&1 | tee output1.log &
+            HF_HOME="$(pwd)/.hf_cache_node1" DEBUG_DISCOVERY=7 DEBUG=7 exo --inference-engine <<parameters.inference_engine>> --node-id "node1" --listen-port 5678 --broadcast-port 5679 --chatgpt-api-port 8000 --chatgpt-api-response-timeout 900 2>&1 | tee output1.log &
             PID1=$!
 
             # Start second instance
-            HF_HOME="$(pwd)/.hf_cache_node2" DEBUG_DISCOVERY=7 DEBUG=7 python3 main.py --inference-engine <<parameters.inference_engine>> --node-id "node2" --listen-port 5679 --broadcast-port 5678 --chatgpt-api-port 8001 --chatgpt-api-response-timeout 900 2>&1 | tee output2.log &
+            HF_HOME="$(pwd)/.hf_cache_node2" DEBUG_DISCOVERY=7 DEBUG=7 exo --inference-engine <<parameters.inference_engine>> --node-id "node2" --listen-port 5679 --broadcast-port 5678 --chatgpt-api-port 8001 --chatgpt-api-response-timeout 900 2>&1 | tee output2.log &
             PID2=$!
 
             # Wait for discovery
@@ -138,9 +138,9 @@ jobs:
           name: Run discovery integration test
           command: |
             source env/bin/activate
-            DEBUG_DISCOVERY=7 DEBUG=7 python3 main.py --node-id "node1" --listen-port 5678 --broadcast-port 5679 --chatgpt-api-port 8000 > output1.log 2>&1 &
+            DEBUG_DISCOVERY=7 DEBUG=7 exo --node-id "node1" --listen-port 5678 --broadcast-port 5679 --chatgpt-api-port 8000 > output1.log 2>&1 &
             PID1=$!
-            DEBUG_DISCOVERY=7 DEBUG=7 python3 main.py --node-id "node2" --listen-port 5679 --broadcast-port 5678 --chatgpt-api-port 8001 > output2.log 2>&1 &
+            DEBUG_DISCOVERY=7 DEBUG=7 exo --node-id "node2" --listen-port 5679 --broadcast-port 5678 --chatgpt-api-port 8001 > output2.log 2>&1 &
             PID2=$!
             sleep 10
             kill $PID1 $PID2

README.md

@@ -124,19 +124,7 @@ def remap_messages(messages: List[Message]) -> List[Message]:
 
 def build_prompt(tokenizer, _messages: List[Message]):
   messages = remap_messages(_messages)
-  if DEBUG >= 3:
-    print(f"messages: {messages}")
-  prompt = tokenizer.apply_chat_template(
-    messages, 
-    tokenize=False, 
-    add_generation_prompt=True
-  )
-
-  if DEBUG >= 3:
-    print(f"prompt: {str(prompt)}")
-    for msg in messages:
-      print(f"chat role: {msg.role}\ncontent: {msg.content}")
-
+  prompt = tokenizer.apply_chat_template([m.to_dict() for m in messages], tokenize=False, add_generation_prompt=True)
   image_str = None
   for message in messages:
     if not isinstance(message.content, list):
@@ -197,7 +185,7 @@ def __init__(self, node: Node, inference_engine_classname: str, response_timeout
     cors.add(self.app.router.add_post("/chat/completions", self.handle_post_chat_completions), {"*": cors_options})
     cors.add(self.app.router.add_post("/v1/chat/completions", self.handle_post_chat_completions), {"*": cors_options})
 
-    self.static_dir = Path(__file__).parent.parent.parent/"tinychat/examples/tinychat"
+    self.static_dir = Path(__file__).parent.parent/"tinychat"
     self.app.router.add_get("/", self.handle_root)
     self.app.router.add_static("/", self.static_dir, name="static")
 

exo/api/chatgpt_api.py

@@ -17,6 +17,7 @@
 
 T = TypeVar("T")
 
+
 async def get_local_snapshot_dir(repo_id: str, revision: str = "main") -> Optional[Path]:
   refs_dir = get_repo_root(repo_id)/"refs"
   refs_file = refs_dir/revision
@@ -69,6 +70,8 @@ def _add_wildcard_to_directories(pattern: str) -> str:
     return pattern + "*"
   return pattern
 
+def get_hf_endpoint() -> str:
+    return os.environ.get('HF_ENDPOINT', "https://huggingface.co")
 
 def get_hf_home() -> Path:
   """Get the Hugging Face home directory."""
@@ -99,7 +102,7 @@ def get_repo_root(repo_id: str) -> Path:
 
 
 async def fetch_file_list(session, repo_id, revision, path=""):
-  api_url = f"https://huggingface.co/api/models/{repo_id}/tree/{revision}"
+  api_url = f"{get_hf_endpoint()}/api/models/{repo_id}/tree/{revision}"
   url = f"{api_url}/{path}" if path else api_url
 
   headers = await get_auth_headers()
@@ -124,7 +127,7 @@ async def fetch_file_list(session, repo_id, revision, path=""):
 async def download_file(
   session: aiohttp.ClientSession, repo_id: str, revision: str, file_path: str, save_directory: str, progress_callback: Optional[RepoFileProgressCallback] = None, use_range_request: bool = True
 ):
-  base_url = f"https://huggingface.co/{repo_id}/resolve/{revision}/"
+  base_url = f"{get_hf_endpoint()}/{repo_id}/resolve/{revision}/"
   url = urljoin(base_url, file_path)
   local_path = os.path.join(save_directory, file_path)
 
@@ -214,7 +217,7 @@ async def resolve_revision_to_commit_hash(repo_id: str, revision: str) -> str:
 
   # Fetch the commit hash for the given revision
   async with aiohttp.ClientSession() as session:
-    api_url = f"https://huggingface.co/api/models/{repo_id}/revision/{revision}"
+    api_url = f"{get_hf_endpoint()}/api/models/{repo_id}/revision/{revision}"
     headers = await get_auth_headers()
     async with session.get(api_url, headers=headers) as response:
       if response.status != 200:

exo/download/hf/hf_helpers.py

@@ -169,7 +169,7 @@ def is_valid_uuid(val):
 
 
 def get_or_create_node_id():
-  NODE_ID_FILE = Path(os.path.dirname(os.path.abspath(__file__)))/".exo_node_id"
+  NODE_ID_FILE = Path(tempfile.gettempdir()) / ".exo_node_id"
   try:
     if NODE_ID_FILE.is_file():
       with open(NODE_ID_FILE, "r") as f:

exo/helpers.py

@@ -4,18 +4,15 @@
 from exo.inference.tinygrad.models.llama import Transformer, convert_from_huggingface, fix_bf16
 from exo.inference.shard import Shard
 from exo.inference.tokenizers import resolve_tokenizer
-from tinygrad.nn.state import safe_load, torch_load, load_state_dict
-from tinygrad import Tensor, dtypes, nn, Context
-from transformers import AutoTokenizer
+from tinygrad.nn.state import load_state_dict
+from tinygrad import Tensor, nn, Context
 from exo.inference.inference_engine import InferenceEngine
 from typing import Optional, Tuple
 import numpy as np
 from exo.inference.tinygrad.tinygrad_helpers import concat_weights, load
 from exo.download.shard_download import ShardDownloader
 from concurrent.futures import ThreadPoolExecutor
 import asyncio
-import threading
-from functools import partial
 
 Tensor.no_grad = True
 # default settings

exo/inference/tinygrad/__init__.py

@@ -5,10 +5,11 @@
 import time
 import traceback
 import uuid
+import sys
 from exo.orchestration.standard_node import StandardNode
 from exo.networking.grpc.grpc_server import GRPCServer
-from exo.networking.udp_discovery import UDPDiscovery
-from exo.networking.tailscale_discovery import TailscaleDiscovery
+from exo.networking.udp.udp_discovery import UDPDiscovery
+from exo.networking.tailscale.tailscale_discovery import TailscaleDiscovery
 from exo.networking.grpc.grpc_peer_handle import GRPCPeerHandle
 from exo.topology.ring_memory_weighted_partitioning_strategy import RingMemoryWeightedPartitioningStrategy
 from exo.api import ChatGPTAPI
@@ -24,6 +25,8 @@
 
 # parse args
 parser = argparse.ArgumentParser(description="Initialize GRPC Discovery")
+parser.add_argument("command", nargs="?", choices=["run"], help="Command to run")
+parser.add_argument("model_name", nargs="?", help="Model name to run")
 parser.add_argument("--node-id", type=str, default=None, help="Node ID")
 parser.add_argument("--node-host", type=str, default="0.0.0.0", help="Node host")
 parser.add_argument("--node-port", type=int, default=None, help="Node port")
@@ -180,14 +183,18 @@ def handle_exit():
 
   await node.start(wait_for_peers=args.wait_for_peers)
 
-  if args.run_model:
-    await run_model_cli(node, inference_engine, args.run_model, args.prompt)
+  if args.command == "run" or args.run_model:
+    model_name = args.model_name or args.run_model
+    if not model_name:
+      print("Error: Model name is required when using 'run' command or --run-model")
+      return
+    await run_model_cli(node, inference_engine, model_name, args.prompt)
   else:
     asyncio.create_task(api.run(port=args.chatgpt_api_port))  # Start the API server as a non-blocking task
     await asyncio.Event().wait()
 
 
-if __name__ == "__main__":
+def run():
   loop = asyncio.new_event_loop()
   asyncio.set_event_loop(loop)
   try:
@@ -197,3 +204,6 @@ def handle_exit():
   finally:
     loop.run_until_complete(shutdown(signal.SIGTERM, loop))
     loop.close()
+
+if __name__ == "__main__":
+  run()

exo/inference/tinygrad/inference.py

@@ -53,6 +53,12 @@
   },
 
   ### qwen
+  "qwen-2.5-coder-1.5b": {
+    "MLXDynamicShardInferenceEngine": Shard(model_id="mlx-community/Qwen2.5-Coder-1.5B-Instruct-4bit", start_layer=0, end_layer=0, n_layers=28),
+  },
+  "qwen-2.5-coder-7b": {
+    "MLXDynamicShardInferenceEngine": Shard(model_id="mlx-community/Qwen2.5-Coder-7B-Instruct-4bit", start_layer=0, end_layer=0, n_layers=28),
+  },
   "qwen-2.5-7b": {
     "MLXDynamicShardInferenceEngine": Shard(model_id="mlx-community/Qwen2.5-7B-Instruct-4bit", start_layer=0, end_layer=0, n_layers=28),
   },

exo/inference/tinygrad/models/__init__.py

@@ -3,8 +3,8 @@
 import traceback
 from typing import List, Dict, Callable, Tuple
 from tailscale import Tailscale, Device
-from .discovery import Discovery
-from .peer_handle import PeerHandle
+from exo.networking.discovery import Discovery
+from exo.networking.peer_handle import PeerHandle
 from exo.topology.device_capabilities import DeviceCapabilities, device_capabilities, UNKNOWN_DEVICE_CAPABILITIES
 from exo.helpers import DEBUG, DEBUG_DISCOVERY
 from .tailscale_helpers import get_device_id, update_device_attributes, get_device_attributes, update_device_attributes

main.py renamed to exo/main.py

@@ -2,7 +2,7 @@
 import asyncio
 import unittest
 from unittest import mock
-from exo.networking.tailscale_discovery import TailscaleDiscovery
+from exo.networking.tailscale.tailscale_discovery import TailscaleDiscovery
 from exo.networking.peer_handle import PeerHandle
 
 class TestTailscaleDiscovery(unittest.IsolatedAsyncioTestCase):

exo/models.py

@@ -1,7 +1,7 @@
 import asyncio
 import unittest
 from unittest import mock
-from exo.networking.udp_discovery import UDPDiscovery
+from exo.networking.udp.udp_discovery import UDPDiscovery
 from exo.networking.grpc.grpc_peer_handle import GRPCPeerHandle
 from exo.networking.grpc.grpc_server import GRPCServer
 from exo.orchestration.node import Node

exo/networking/tailscale/__init__.py

@@ -4,8 +4,8 @@
 import time
 import traceback
 from typing import List, Dict, Callable, Tuple, Coroutine
-from .discovery import Discovery
-from .peer_handle import PeerHandle
+from exo.networking.discovery import Discovery
+from exo.networking.peer_handle import PeerHandle
 from exo.topology.device_capabilities import DeviceCapabilities, device_capabilities, UNKNOWN_DEVICE_CAPABILITIES
 from exo.helpers import DEBUG, DEBUG_DISCOVERY, get_all_ip_addresses
 

exo/networking/tailscale_discovery.py renamed to exo/networking/tailscale/tailscale_discovery.py

@@ -43,6 +43,8 @@
 <option value="deepseek-coder-v2-lite">Deepseek Coder V2 Lite</option>
 <option value="deepseek-coder-v2.5">Deepseek Coder V2.5</option>
 <option value="llava-1.5-7b-hf">LLaVa 1.5 7B (Vision Model)</option>
+<option value="qwen-2.5-coder-1.5b">Qwen 2.5 Coder 1.5B</option>
+<option value="qwen-2.5-coder-7b">Qwen 2.5 Coder 7B</option>
 <option value="qwen-2.5-7b">Qwen 2.5 7B</option>
 <option value="qwen-2.5-math-7b">Qwen 2.5 7B (Math)</option>
 <option value="qwen-2.5-14b">Qwen 2.5 14B</option>

exo/networking/tailscale_helpers.py renamed to exo/networking/tailscale/tailscale_helpers.py

@@ -2,4 +2,4 @@
 
 python3 -m venv .venv
 source .venv/bin/activate
-pip install .
+pip install -e .