← back to Exo
chatgpt api integration
acc94b50c73c7f65658461a6c6b3f44994e0fab7 · 2024-07-28 16:12:21 +0530 · Varshith
Files touched
M exo/api/chatgpt_api.pyM exo/inference/mlx/sharded_inference_engine.pyM exo/inference/mlx/sharded_utils.pyM exo/inference/mlx/test_sharded_llava.pyM exo/networking/grpc/grpc_peer_handle.pyM exo/networking/grpc/grpc_server.pyM exo/networking/grpc/node_service.protoM exo/networking/grpc/node_service_pb2.pyM exo/networking/peer_handle.pyM exo/orchestration/node.pyM exo/orchestration/standard_node.pyM setup.py
Diff
commit acc94b50c73c7f65658461a6c6b3f44994e0fab7
Author: Varshith <varshith.bathini@sprinklr.com>
Date: Sun Jul 28 16:12:21 2024 +0530
chatgpt api integration
---
exo/api/chatgpt_api.py | 43 ++++++++++++++----
exo/inference/mlx/sharded_inference_engine.py | 13 ++++--
exo/inference/mlx/sharded_utils.py | 14 ++++++
exo/inference/mlx/test_sharded_llava.py | 8 ++--
exo/networking/grpc/grpc_peer_handle.py | 3 +-
exo/networking/grpc/grpc_server.py | 5 ++-
exo/networking/grpc/node_service.proto | 5 ++-
exo/networking/grpc/node_service_pb2.py | 64 +++++++++++++--------------
exo/networking/peer_handle.py | 2 +-
exo/orchestration/node.py | 2 +-
exo/orchestration/standard_node.py | 21 +++++----
setup.py | 2 +-
12 files changed, 119 insertions(+), 63 deletions(-)
diff --git a/exo/api/chatgpt_api.py b/exo/api/chatgpt_api.py
index 9795fa1a..fff4a63d 100644
--- a/exo/api/chatgpt_api.py
+++ b/exo/api/chatgpt_api.py
@@ -3,7 +3,7 @@ import time
import asyncio
import json
from pathlib import Path
-from transformers import AutoTokenizer
+from transformers import AutoTokenizer, AutoProcessor
from typing import List, Literal, Union, Dict
from aiohttp import web
import aiohttp_cors
@@ -42,11 +42,15 @@ shard_mappings = {
"deepseek-coder-v2-lite": {
"MLXDynamicShardInferenceEngine": Shard(model_id="mlx-community/DeepSeek-Coder-V2-Lite-Instruct-4bit-mlx", start_layer=0, end_layer=0, n_layers=27),
},
+ ### llava
+ "llava-1.5-7b-hf": {
+ "MLXDynamicShardInferenceEngine": Shard(model_id="llava-hf/llava-1.5-7b-hf", start_layer=0, end_layer=0, n_layers=32),
+ },
}
class Message:
- def __init__(self, role: str, content: str):
+ def __init__(self, role: str, content: Union[str, list]):
self.role = role
self.content = content
@@ -68,6 +72,18 @@ def resolve_tinygrad_tokenizer(model_id: str):
async def resolve_tokenizer(model_id: str):
+ try:
+ if DEBUG >= 2: print(f"Trying to AutoProcessor for {model_id}")
+ processor = AutoProcessor.from_pretrained(model_id)
+ processor.eos_token_id = processor.tokenizer.eos_token_id
+ processor.encode = processor.tokenizer.encode
+ return processor
+ except Exception as e:
+ if DEBUG >= 2: print(f"Failed to load processor for {model_id}. Error: {e}")
+ import traceback
+
+ if DEBUG >= 2: print(traceback.format_exc())
+
try:
if DEBUG >= 2: print(f"Trying AutoTokenizer for {model_id}")
return AutoTokenizer.from_pretrained(model_id)
@@ -138,7 +154,18 @@ def generate_completion(
def build_prompt(tokenizer, messages: List[Message]):
- return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+ prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+ image_str = None
+ for message in messages:
+ if not isinstance(message.content, list):
+ continue
+
+ for content in message.content:
+ if content.get("type", None) == "image":
+ image_str = content.get("image", None)
+ break
+
+ return prompt, image_str
def parse_message(data: dict):
@@ -195,7 +222,7 @@ class ChatGPTAPI:
shard = shard_mappings.get(data.get("model", "llama-3.1-8b"), {}).get(self.inference_engine_classname)
messages = [parse_message(msg) for msg in data.get("messages", [])]
tokenizer = await resolve_tokenizer(shard.model_id)
- return web.json_response({"length": len(build_prompt(tokenizer, messages))})
+ return web.json_response({"length": len(build_prompt(tokenizer, messages)[0])})
async def handle_post_chat_completions(self, request):
data = await request.json()
@@ -219,13 +246,13 @@ class ChatGPTAPI:
tokenizer = await resolve_tokenizer(shard.model_id)
if DEBUG >= 4: print(f"Resolved tokenizer: {tokenizer}")
- prompt = build_prompt(tokenizer, chat_request.messages)
+ prompt, image_str = build_prompt(tokenizer, chat_request.messages)
callback_id = f"chatgpt-api-wait-response-{request_id}"
callback = self.node.on_token.register(callback_id)
- if DEBUG >= 2: print(f"Sending prompt from ChatGPT api {request_id=} {shard=} {prompt=}")
+ if DEBUG >= 2: print(f"Sending prompt from ChatGPT api {request_id=} {shard=} {prompt=} {image_str=}")
try:
- await self.node.process_prompt(shard, prompt, request_id=request_id)
+ await self.node.process_prompt(shard, prompt, image_str, request_id=request_id)
except Exception as e:
if DEBUG >= 2:
import traceback
@@ -294,7 +321,7 @@ class ChatGPTAPI:
)
finish_reason = "length"
- eos_token_id = tokenizer.special_tokens_map.get("eos_token_id") if isinstance(tokenizer._tokenizer, AutoTokenizer) else tokenizer.eos_token_id
+ eos_token_id = tokenizer.special_tokens_map.get("eos_token_id") if isinstance(getattr(tokenizer, "_tokenizer", None), AutoTokenizer) else tokenizer.eos_token_id
if DEBUG >= 2: print(f"Checking if end of tokens result {tokens[-1]=} is {eos_token_id=}")
if tokens[-1] == eos_token_id:
tokens = tokens[:-1]
diff --git a/exo/inference/mlx/sharded_inference_engine.py b/exo/inference/mlx/sharded_inference_engine.py
index b5104c72..9726123b 100644
--- a/exo/inference/mlx/sharded_inference_engine.py
+++ b/exo/inference/mlx/sharded_inference_engine.py
@@ -2,7 +2,7 @@ import numpy as np
import mlx.core as mx
from ..inference_engine import InferenceEngine
from .sharded_model import StatefulShardedModel
-from .sharded_utils import load_shard
+from .sharded_utils import load_shard, get_image_from_str
from ..shard import Shard
from typing import Optional
@@ -11,9 +11,16 @@ class MLXDynamicShardInferenceEngine(InferenceEngine):
def __init__(self):
self.shard = None
- async def infer_prompt(self, request_id: str, shard: Shard, prompt: str, inference_state: Optional[str] = None) -> (np.ndarray, str, bool):
+ async def infer_prompt(self, request_id: str, shard: Shard, prompt: str, image_str: Optional[str] = None, inference_state: Optional[str] = None) -> (np.ndarray, str, bool):
await self.ensure_shard(shard)
- output_data: np.ndarray = np.array(self.stateful_sharded_model.step(request_id, mx.array(self.tokenizer.encode(prompt))))
+ if image_str:
+ image = get_image_from_str(image_str)
+ inputs = self.tokenizer(prompt, image, return_tensors="np")
+ pixel_values = mx.array(inputs["pixel_values"])
+ input_ids = mx.array(inputs["input_ids"])
+ output_data: np.ndarray = np.array(self.stateful_sharded_model.step(request_id, input_ids, pixel_values))
+ else:
+ output_data: np.ndarray = np.array(self.stateful_sharded_model.step(request_id, mx.array(self.tokenizer.encode(prompt))))
return output_data, "", output_data.size == 1 and output_data.item() == self.tokenizer.eos_token_id
async def infer_tensor(self, request_id: str, shard: Shard, input_data: np.ndarray, inference_state: Optional[str] = None) -> (np.ndarray, str, bool):
diff --git a/exo/inference/mlx/sharded_utils.py b/exo/inference/mlx/sharded_utils.py
index e5429c64..296d552f 100644
--- a/exo/inference/mlx/sharded_utils.py
+++ b/exo/inference/mlx/sharded_utils.py
@@ -8,6 +8,9 @@ import asyncio
from functools import partial
from pathlib import Path
from typing import Optional, Tuple
+import requests
+from PIL import Image
+from io import BytesIO
import mlx.core as mx
import mlx.nn as nn
@@ -222,7 +225,18 @@ async def load_shard(
# TODO: figure out a generic solution
if model.model_type == "llava":
processor = AutoProcessor.from_pretrained(model_path)
+ processor.eos_token_id = processor.tokenizer.eos_token_id
+ processor.encode = processor.tokenizer.encode
return model, processor
else:
tokenizer = load_tokenizer(model_path, tokenizer_config)
return model, tokenizer
+
+def get_image_from_str(image_str: str):
+ if image_str.startswith("http"):
+ response = requests.get(image_str, timeout=10)
+ image = Image.open(BytesIO(response.content)).convert("RGB")
+ else:
+ imgdata = base64.b64decode(image_str)
+ image = Image.open(io.BytesIO(imgdata))
+ return image
diff --git a/exo/inference/mlx/test_sharded_llava.py b/exo/inference/mlx/test_sharded_llava.py
index 5e2e13ad..f2e75c48 100644
--- a/exo/inference/mlx/test_sharded_llava.py
+++ b/exo/inference/mlx/test_sharded_llava.py
@@ -15,9 +15,11 @@ shard_full = Shard("llava", 0, 31, 32)
shard1 = Shard("llava", 0, 12, 32)
shard2 = Shard("llava", 13, 31, 32)
-full_model_shard, full_processor = asyncio.run(load_shard("llava-hf/llava-1.5-7b-hf", shard=shard_full))
-model_shard1, processor1 = asyncio.run(load_shard("llava-hf/llava-1.5-7b-hf", shard=shard1))
-model_shard2, processor2 = asyncio.run(load_shard("llava-hf/llava-1.5-7b-hf", shard=shard2))
+model_path = "llava-hf/llava-1.5-7b-hf"
+
+full_model_shard, full_processor = asyncio.run(load_shard(model_path, shard=shard_full))
+model_shard1, processor1 = asyncio.run(load_shard(model_path, shard=shard1))
+model_shard2, processor2 = asyncio.run(load_shard(model_path, shard=shard2))
full = StatefulShardedModel(shard_full, full_model_shard)
m1 = StatefulShardedModel(shard1, model_shard1)
diff --git a/exo/networking/grpc/grpc_peer_handle.py b/exo/networking/grpc/grpc_peer_handle.py
index 78009d49..e3e98dce 100644
--- a/exo/networking/grpc/grpc_peer_handle.py
+++ b/exo/networking/grpc/grpc_peer_handle.py
@@ -39,9 +39,10 @@ class GRPCPeerHandle(PeerHandle):
self.channel = None
self.stub = None
- async def send_prompt(self, shard: Shard, prompt: str, request_id: Optional[str] = None, inference_state: Optional[str] = None) -> Optional[np.array]:
+ async def send_prompt(self, shard: Shard, prompt: str, image_str: Optional[str] = None, request_id: Optional[str] = None, inference_state: Optional[str] = None) -> Optional[np.array]:
request = node_service_pb2.PromptRequest(
prompt=prompt,
+ image_str=image_str,
shard=node_service_pb2.Shard(
model_id=shard.model_id,
start_layer=shard.start_layer,
diff --git a/exo/networking/grpc/grpc_server.py b/exo/networking/grpc/grpc_server.py
index 13d4bf6f..5992ce1f 100644
--- a/exo/networking/grpc/grpc_server.py
+++ b/exo/networking/grpc/grpc_server.py
@@ -45,9 +45,10 @@ class GRPCServer(node_service_pb2_grpc.NodeServiceServicer):
n_layers=request.shard.n_layers,
)
prompt = request.prompt
+ image_str = request.image_str
request_id = request.request_id
- result = await self.node.process_prompt(shard, prompt, request_id)
- if DEBUG >= 2: print(f"SendPrompt {shard=} {prompt=} {request_id=} result: {result}")
+ result = await self.node.process_prompt(shard, prompt, image_str, request_id)
+ if DEBUG >= 2: print(f"SendPrompt {shard=} {prompt=} {image=} {request_id=} result: {result}")
tensor_data = result.tobytes() if result is not None else None
return node_service_pb2.Tensor(tensor_data=tensor_data, shape=result.shape, dtype=str(result.dtype)) if result is not None else node_service_pb2.Tensor()
diff --git a/exo/networking/grpc/node_service.proto b/exo/networking/grpc/node_service.proto
index d76430ca..6fcee351 100644
--- a/exo/networking/grpc/node_service.proto
+++ b/exo/networking/grpc/node_service.proto
@@ -21,8 +21,9 @@ message Shard {
message PromptRequest {
Shard shard = 1;
string prompt = 2;
- optional string request_id = 3;
- optional string inference_state = 4;
+ optional string image_str = 3;
+ optional string request_id = 4;
+ optional string inference_state = 5;
}
message TensorRequest {
diff --git a/exo/networking/grpc/node_service_pb2.py b/exo/networking/grpc/node_service_pb2.py
index 66e516c7..62ffb1bb 100644
--- a/exo/networking/grpc/node_service_pb2.py
+++ b/exo/networking/grpc/node_service_pb2.py
@@ -14,7 +14,7 @@ _sym_db = _symbol_database.Default()
-DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x12node_service.proto\x12\x0cnode_service\"S\n\x05Shard\x12\x10\n\x08model_id\x18\x01 \x01(\t\x12\x13\n\x0bstart_layer\x18\x02 \x01(\x05\x12\x11\n\tend_layer\x18\x03 \x01(\x05\x12\x10\n\x08n_layers\x18\x04 \x01(\x05\"\x9d\x01\n\rPromptRequest\x12\"\n\x05shard\x18\x01 \x01(\x0b\x32\x13.node_service.Shard\x12\x0e\n\x06prompt\x18\x02 \x01(\t\x12\x17\n\nrequest_id\x18\x03 \x01(\tH\x00\x88\x01\x01\x12\x1c\n\x0finference_state\x18\x04 \x01(\tH\x01\x88\x01\x01\x42\r\n\x0b_request_idB\x12\n\x10_inference_state\"\xb3\x01\n\rTensorRequest\x12\"\n\x05shard\x18\x01 \x01(\x0b\x32\x13.node_service.Shard\x12$\n\x06tensor\x18\x02 \x01(\x0b\x32\x14.node_service.Tensor\x12\x17\n\nrequest_id\x18\x03 \x01(\tH\x00\x88\x01\x01\x12\x1c\n\x0finference_state\x18\x04 \x01(\tH\x01\x88\x01\x01\x42\r\n\x0b_request_idB\x12\n\x10_inference_state\"/\n\x19GetInferenceResultRequest\x12\x12\n\nrequest_id\x18\x01 \x01(\t\"\\\n\x0fInferenceResult\x12)\n\x06tensor\x18\x01 \x01(\x0b\x32\x14.node_service.TensorH\x00\x88\x01\x01\x12\x13\n\x0bis_finished\x18\x02 \x01(\x08\x42\t\n\x07_tensor\";\n\x06Tensor\x12\x13\n\x0btensor_data\x18\x01 \x01(\x0c\x12\r\n\x05shape\x18\x02 \x03(\x05\x12\r\n\x05\x64type\x18\x03 \x01(\t\"<\n\x16\x43ollectTopologyRequest\x12\x0f\n\x07visited\x18\x01 \x03(\t\x12\x11\n\tmax_depth\x18\x02 \x01(\x05\"\x8e\x02\n\x08Topology\x12\x30\n\x05nodes\x18\x01 \x03(\x0b\x32!.node_service.Topology.NodesEntry\x12\x39\n\npeer_graph\x18\x02 \x03(\x0b\x32%.node_service.Topology.PeerGraphEntry\x1aN\n\nNodesEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12/\n\x05value\x18\x02 \x01(\x0b\x32 .node_service.DeviceCapabilities:\x02\x38\x01\x1a\x45\n\x0ePeerGraphEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\"\n\x05value\x18\x02 \x01(\x0b\x32\x13.node_service.Peers:\x02\x38\x01\"\x19\n\x05Peers\x12\x10\n\x08peer_ids\x18\x01 \x03(\t\"7\n\x0b\x44\x65viceFlops\x12\x0c\n\x04\x66p32\x18\x01 \x01(\x02\x12\x0c\n\x04\x66p16\x18\x02 \x01(\x02\x12\x0c\n\x04int8\x18\x03 \x01(\x02\"k\n\x12\x44\x65viceCapabilities\x12\r\n\x05model\x18\x01 \x01(\t\x12\x0c\n\x04\x63hip\x18\x02 \x01(\t\x12\x0e\n\x06memory\x18\x03 \x01(\x05\x12(\n\x05\x66lops\x18\x04 \x01(\x0b\x32\x19.node_service.DeviceFlops\"L\n\x11SendResultRequest\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0e\n\x06result\x18\x02 \x03(\x05\x12\x13\n\x0bis_finished\x18\x03 \x01(\x08\"=\n\x17SendOpaqueStatusRequest\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0e\n\x06status\x18\x02 \x01(\t\"\x07\n\x05\x45mpty2\xde\x03\n\x0bNodeService\x12\x41\n\nSendPrompt\x12\x1b.node_service.PromptRequest\x1a\x14.node_service.Tensor\"\x00\x12\x41\n\nSendTensor\x12\x1b.node_service.TensorRequest\x1a\x14.node_service.Tensor\"\x00\x12^\n\x12GetInferenceResult\x12\'.node_service.GetInferenceResultRequest\x1a\x1d.node_service.InferenceResult\"\x00\x12Q\n\x0f\x43ollectTopology\x12$.node_service.CollectTopologyRequest\x1a\x16.node_service.Topology\"\x00\x12\x44\n\nSendResult\x12\x1f.node_service.SendResultRequest\x1a\x13.node_service.Empty\"\x00\x12P\n\x10SendOpaqueStatus\x12%.node_service.SendOpaqueStatusRequest\x1a\x13.node_service.Empty\"\x00\x62\x06proto3')
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x12node_service.proto\x12\x0cnode_service\"S\n\x05Shard\x12\x10\n\x08model_id\x18\x01 \x01(\t\x12\x13\n\x0bstart_layer\x18\x02 \x01(\x05\x12\x11\n\tend_layer\x18\x03 \x01(\x05\x12\x10\n\x08n_layers\x18\x04 \x01(\x05\"\xc3\x01\n\rPromptRequest\x12\"\n\x05shard\x18\x01 \x01(\x0b\x32\x13.node_service.Shard\x12\x0e\n\x06prompt\x18\x02 \x01(\t\x12\x16\n\timage_str\x18\x03 \x01(\tH\x00\x88\x01\x01\x12\x17\n\nrequest_id\x18\x04 \x01(\tH\x01\x88\x01\x01\x12\x1c\n\x0finference_state\x18\x05 \x01(\tH\x02\x88\x01\x01\x42\x0c\n\n_image_strB\r\n\x0b_request_idB\x12\n\x10_inference_state\"\xb3\x01\n\rTensorRequest\x12\"\n\x05shard\x18\x01 \x01(\x0b\x32\x13.node_service.Shard\x12$\n\x06tensor\x18\x02 \x01(\x0b\x32\x14.node_service.Tensor\x12\x17\n\nrequest_id\x18\x03 \x01(\tH\x00\x88\x01\x01\x12\x1c\n\x0finference_state\x18\x04 \x01(\tH\x01\x88\x01\x01\x42\r\n\x0b_request_idB\x12\n\x10_inference_state\"/\n\x19GetInferenceResultRequest\x12\x12\n\nrequest_id\x18\x01 \x01(\t\"\\\n\x0fInferenceResult\x12)\n\x06tensor\x18\x01 \x01(\x0b\x32\x14.node_service.TensorH\x00\x88\x01\x01\x12\x13\n\x0bis_finished\x18\x02 \x01(\x08\x42\t\n\x07_tensor\";\n\x06Tensor\x12\x13\n\x0btensor_data\x18\x01 \x01(\x0c\x12\r\n\x05shape\x18\x02 \x03(\x05\x12\r\n\x05\x64type\x18\x03 \x01(\t\"<\n\x16\x43ollectTopologyRequest\x12\x0f\n\x07visited\x18\x01 \x03(\t\x12\x11\n\tmax_depth\x18\x02 \x01(\x05\"\x8e\x02\n\x08Topology\x12\x30\n\x05nodes\x18\x01 \x03(\x0b\x32!.node_service.Topology.NodesEntry\x12\x39\n\npeer_graph\x18\x02 \x03(\x0b\x32%.node_service.Topology.PeerGraphEntry\x1aN\n\nNodesEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12/\n\x05value\x18\x02 \x01(\x0b\x32 .node_service.DeviceCapabilities:\x02\x38\x01\x1a\x45\n\x0ePeerGraphEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\"\n\x05value\x18\x02 \x01(\x0b\x32\x13.node_service.Peers:\x02\x38\x01\"\x19\n\x05Peers\x12\x10\n\x08peer_ids\x18\x01 \x03(\t\"7\n\x0b\x44\x65viceFlops\x12\x0c\n\x04\x66p32\x18\x01 \x01(\x02\x12\x0c\n\x04\x66p16\x18\x02 \x01(\x02\x12\x0c\n\x04int8\x18\x03 \x01(\x02\"k\n\x12\x44\x65viceCapabilities\x12\r\n\x05model\x18\x01 \x01(\t\x12\x0c\n\x04\x63hip\x18\x02 \x01(\t\x12\x0e\n\x06memory\x18\x03 \x01(\x05\x12(\n\x05\x66lops\x18\x04 \x01(\x0b\x32\x19.node_service.DeviceFlops\"L\n\x11SendResultRequest\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0e\n\x06result\x18\x02 \x03(\x05\x12\x13\n\x0bis_finished\x18\x03 \x01(\x08\"=\n\x17SendOpaqueStatusRequest\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0e\n\x06status\x18\x02 \x01(\t\"\x07\n\x05\x45mpty2\xde\x03\n\x0bNodeService\x12\x41\n\nSendPrompt\x12\x1b.node_service.PromptRequest\x1a\x14.node_service.Tensor\"\x00\x12\x41\n\nSendTensor\x12\x1b.node_service.TensorRequest\x1a\x14.node_service.Tensor\"\x00\x12^\n\x12GetInferenceResult\x12\'.node_service.GetInferenceResultRequest\x1a\x1d.node_service.InferenceResult\"\x00\x12Q\n\x0f\x43ollectTopology\x12$.node_service.CollectTopologyRequest\x1a\x16.node_service.Topology\"\x00\x12\x44\n\nSendResult\x12\x1f.node_service.SendResultRequest\x1a\x13.node_service.Empty\"\x00\x12P\n\x10SendOpaqueStatus\x12%.node_service.SendOpaqueStatusRequest\x1a\x13.node_service.Empty\"\x00\x62\x06proto3')
_globals = globals()
_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
@@ -28,35 +28,35 @@ if not _descriptor._USE_C_DESCRIPTORS:
_globals['_SHARD']._serialized_start=36
_globals['_SHARD']._serialized_end=119
_globals['_PROMPTREQUEST']._serialized_start=122
- _globals['_PROMPTREQUEST']._serialized_end=279
- _globals['_TENSORREQUEST']._serialized_start=282
- _globals['_TENSORREQUEST']._serialized_end=461
- _globals['_GETINFERENCERESULTREQUEST']._serialized_start=463
- _globals['_GETINFERENCERESULTREQUEST']._serialized_end=510
- _globals['_INFERENCERESULT']._serialized_start=512
- _globals['_INFERENCERESULT']._serialized_end=604
- _globals['_TENSOR']._serialized_start=606
- _globals['_TENSOR']._serialized_end=665
- _globals['_COLLECTTOPOLOGYREQUEST']._serialized_start=667
- _globals['_COLLECTTOPOLOGYREQUEST']._serialized_end=727
- _globals['_TOPOLOGY']._serialized_start=730
- _globals['_TOPOLOGY']._serialized_end=1000
- _globals['_TOPOLOGY_NODESENTRY']._serialized_start=851
- _globals['_TOPOLOGY_NODESENTRY']._serialized_end=929
- _globals['_TOPOLOGY_PEERGRAPHENTRY']._serialized_start=931
- _globals['_TOPOLOGY_PEERGRAPHENTRY']._serialized_end=1000
- _globals['_PEERS']._serialized_start=1002
- _globals['_PEERS']._serialized_end=1027
- _globals['_DEVICEFLOPS']._serialized_start=1029
- _globals['_DEVICEFLOPS']._serialized_end=1084
- _globals['_DEVICECAPABILITIES']._serialized_start=1086
- _globals['_DEVICECAPABILITIES']._serialized_end=1193
- _globals['_SENDRESULTREQUEST']._serialized_start=1195
- _globals['_SENDRESULTREQUEST']._serialized_end=1271
- _globals['_SENDOPAQUESTATUSREQUEST']._serialized_start=1273
- _globals['_SENDOPAQUESTATUSREQUEST']._serialized_end=1334
- _globals['_EMPTY']._serialized_start=1336
- _globals['_EMPTY']._serialized_end=1343
- _globals['_NODESERVICE']._serialized_start=1346
- _globals['_NODESERVICE']._serialized_end=1824
+ _globals['_PROMPTREQUEST']._serialized_end=317
+ _globals['_TENSORREQUEST']._serialized_start=320
+ _globals['_TENSORREQUEST']._serialized_end=499
+ _globals['_GETINFERENCERESULTREQUEST']._serialized_start=501
+ _globals['_GETINFERENCERESULTREQUEST']._serialized_end=548
+ _globals['_INFERENCERESULT']._serialized_start=550
+ _globals['_INFERENCERESULT']._serialized_end=642
+ _globals['_TENSOR']._serialized_start=644
+ _globals['_TENSOR']._serialized_end=703
+ _globals['_COLLECTTOPOLOGYREQUEST']._serialized_start=705
+ _globals['_COLLECTTOPOLOGYREQUEST']._serialized_end=765
+ _globals['_TOPOLOGY']._serialized_start=768
+ _globals['_TOPOLOGY']._serialized_end=1038
+ _globals['_TOPOLOGY_NODESENTRY']._serialized_start=889
+ _globals['_TOPOLOGY_NODESENTRY']._serialized_end=967
+ _globals['_TOPOLOGY_PEERGRAPHENTRY']._serialized_start=969
+ _globals['_TOPOLOGY_PEERGRAPHENTRY']._serialized_end=1038
+ _globals['_PEERS']._serialized_start=1040
+ _globals['_PEERS']._serialized_end=1065
+ _globals['_DEVICEFLOPS']._serialized_start=1067
+ _globals['_DEVICEFLOPS']._serialized_end=1122
+ _globals['_DEVICECAPABILITIES']._serialized_start=1124
+ _globals['_DEVICECAPABILITIES']._serialized_end=1231
+ _globals['_SENDRESULTREQUEST']._serialized_start=1233
+ _globals['_SENDRESULTREQUEST']._serialized_end=1309
+ _globals['_SENDOPAQUESTATUSREQUEST']._serialized_start=1311
+ _globals['_SENDOPAQUESTATUSREQUEST']._serialized_end=1372
+ _globals['_EMPTY']._serialized_start=1374
+ _globals['_EMPTY']._serialized_end=1381
+ _globals['_NODESERVICE']._serialized_start=1384
+ _globals['_NODESERVICE']._serialized_end=1862
# @@protoc_insertion_point(module_scope)
diff --git a/exo/networking/peer_handle.py b/exo/networking/peer_handle.py
index 27ff4201..cf232d00 100644
--- a/exo/networking/peer_handle.py
+++ b/exo/networking/peer_handle.py
@@ -28,7 +28,7 @@ class PeerHandle(ABC):
pass
@abstractmethod
- async def send_prompt(self, shard: Shard, prompt: str, request_id: Optional[str] = None, inference_state: Optional[str] = None) -> Optional[np.array]:
+ async def send_prompt(self, shard: Shard, prompt: str, image_str: Optional[str] = None, request_id: Optional[str] = None, inference_state: Optional[str] = None) -> Optional[np.array]:
pass
@abstractmethod
diff --git a/exo/orchestration/node.py b/exo/orchestration/node.py
index 93b777d4..60b72974 100644
--- a/exo/orchestration/node.py
+++ b/exo/orchestration/node.py
@@ -16,7 +16,7 @@ class Node(ABC):
pass
@abstractmethod
- async def process_prompt(self, shard: Shard, prompt: str, request_id: Optional[str] = None, inference_state: Optional[str] = None) -> Optional[np.ndarray]:
+ async def process_prompt(self, shard: Shard, prompt: str, image_str: Optional[str] = None, request_id: Optional[str] = None, inference_state: Optional[str] = None) -> Optional[np.ndarray]:
pass
@abstractmethod
diff --git a/exo/orchestration/standard_node.py b/exo/orchestration/standard_node.py
index 23e3e032..55efdebc 100644
--- a/exo/orchestration/standard_node.py
+++ b/exo/orchestration/standard_node.py
@@ -69,7 +69,7 @@ class StandardNode(Node):
await self.discovery.stop()
await self.server.stop()
- async def process_prompt(self, base_shard: Shard, prompt: str, request_id: Optional[str] = None, inference_state: Optional[str] = None) -> Optional[np.ndarray]:
+ async def process_prompt(self, base_shard: Shard, prompt: str, image_str: Optional[str] = None, request_id: Optional[str] = None, inference_state: Optional[str] = None) -> Optional[np.ndarray]:
shard = self.get_current_shard(base_shard)
asyncio.create_task(
self.broadcast_opaque_status(
@@ -82,6 +82,7 @@ class StandardNode(Node):
"base_shard": base_shard.to_dict(),
"shard": shard.to_dict(),
"prompt": prompt,
+ "image_str": image_str,
"inference_state": inference_state,
"request_id": request_id,
}
@@ -89,7 +90,7 @@ class StandardNode(Node):
)
)
start_time = time.perf_counter_ns()
- resp = await self._process_prompt(base_shard, prompt, request_id, inference_state)
+ resp = await self._process_prompt(base_shard, prompt, image_str, request_id, inference_state)
end_time = time.perf_counter_ns()
elapsed_time_ns = end_time - start_time
asyncio.create_task(
@@ -103,6 +104,7 @@ class StandardNode(Node):
"base_shard": base_shard.to_dict(),
"shard": shard.to_dict(),
"prompt": prompt,
+ "image_str": image_str,
"inference_state": inference_state,
"request_id": request_id,
"elapsed_time_ns": elapsed_time_ns,
@@ -113,20 +115,20 @@ class StandardNode(Node):
)
return resp
- async def _process_prompt(self, base_shard: Shard, prompt: str, request_id: Optional[str] = None, inference_state: Optional[str] = None) -> Optional[np.ndarray]:
+ async def _process_prompt(self, base_shard: Shard, prompt: str, image_str: Optional[str] = None, request_id: Optional[str] = None, inference_state: Optional[str] = None) -> Optional[np.ndarray]:
if request_id is None:
request_id = str(uuid.uuid4())
if request_id not in self.buffered_token_output:
self.buffered_token_output[request_id] = ([], False)
shard = self.get_current_shard(base_shard)
- if DEBUG >= 2: print(f"[{request_id}] process prompt: {base_shard=} {shard=} {prompt=}")
+ if DEBUG >= 2: print(f"[{request_id}] process prompt: {base_shard=} {shard=} {prompt=} {image_str=}")
if shard.start_layer != 0:
- if DEBUG >= 2: print(f"[{request_id}] forwarding to next shard: {base_shard=} {shard=} {prompt=}")
- await self.forward_to_next_shard(shard, prompt, request_id)
+ if DEBUG >= 2: print(f"[{request_id}] forwarding to next shard: {base_shard=} {shard=} {prompt=} {image_str=}")
+ await self.forward_to_next_shard(shard, prompt, request_id, image_str)
return
- result, inference_state, is_finished = await self.inference_engine.infer_prompt(request_id, shard, prompt, inference_state=inference_state)
+ result, inference_state, is_finished = await self.inference_engine.infer_prompt(request_id, shard, prompt, image_str, inference_state=inference_state)
is_finished = is_finished or len(self.buffered_token_output[request_id][0]) >= self.max_generate_tokens
if is_finished:
self.buffered_token_output[request_id] = (self.buffered_token_output[request_id][0], True)
@@ -234,6 +236,7 @@ class StandardNode(Node):
base_shard: Shard,
tensor_or_prompt: Union[np.ndarray, str],
request_id: str,
+ image_str: Optional[str] = None,
inference_state: Optional[str] = None,
) -> None:
if not self.partitioning_strategy:
@@ -255,7 +258,7 @@ class StandardNode(Node):
if isinstance(tensor_or_prompt, np.ndarray):
await self.process_tensor(shard, tensor_or_prompt, request_id, inference_state=inference_state)
else:
- await self.process_prompt(shard, tensor_or_prompt, request_id, inference_state=inference_state)
+ await self.process_prompt(shard, tensor_or_prompt, image_str, request_id, inference_state=inference_state)
return
target_peer = next((p for p in self.peers if p.id() == next_partition.node_id), None)
@@ -267,7 +270,7 @@ class StandardNode(Node):
if isinstance(tensor_or_prompt, np.ndarray):
await target_peer.send_tensor(next_shard, tensor_or_prompt, request_id=request_id, inference_state=inference_state)
else:
- await target_peer.send_prompt(next_shard, tensor_or_prompt, request_id=request_id, inference_state=inference_state)
+ await target_peer.send_prompt(next_shard, tensor_or_prompt, image_str=image_str, request_id=request_id, inference_state=inference_state)
def get_current_shard(self, base_shard: Shard) -> Shard:
partitions = self.partitioning_strategy.partition(self.topology)
diff --git a/setup.py b/setup.py
index 79fef2e8..53f84e19 100644
--- a/setup.py
+++ b/setup.py
@@ -22,7 +22,7 @@ install_requires = [
"tiktoken==0.7.0",
"tokenizers==0.19.1",
"tqdm==4.66.4",
- "transformers==4.41.2",
+ "transformers==4.43.3",
"uuid==1.30",
"tinygrad @ git+https://github.com/tinygrad/tinygrad.git@a9f5a764dc640a5e5cbaaeeee21df7c8ca37da38",
]
← 33cbacf5 fix llava sanitize
·
back to Exo
·
update readme 8d3d3df1 →