← back to Exo
wip state
03ba31c020d965c79ec0ed53c5ff2a87b8752160 · 2024-07-16 10:13:13 -0700 · Alex Cheema
Files touched
M exo/inference/inference_engine.pyM exo/inference/mlx/sharded_inference_engine.pyM exo/inference/test_inference_engine.pyM exo/inference/tinygrad/inference.pyM exo/networking/grpc/grpc_server.pyM exo/networking/grpc/node_service.protoM exo/networking/grpc/node_service_pb2.pyM exo/orchestration/node.pyM exo/orchestration/standard_node.py
Diff
commit 03ba31c020d965c79ec0ed53c5ff2a87b8752160
Author: Alex Cheema <alexcheema123@gmail.com>
Date: Tue Jul 16 10:13:13 2024 -0700
wip state
---
exo/inference/inference_engine.py | 7 ++-
exo/inference/mlx/sharded_inference_engine.py | 7 +--
exo/inference/test_inference_engine.py | 8 ++--
exo/inference/tinygrad/inference.py | 22 +++++++---
exo/networking/grpc/grpc_server.py | 3 +-
exo/networking/grpc/node_service.proto | 6 ++-
exo/networking/grpc/node_service_pb2.py | 62 +++++++++++++--------------
exo/orchestration/node.py | 4 +-
exo/orchestration/standard_node.py | 18 ++++----
9 files changed, 75 insertions(+), 62 deletions(-)
diff --git a/exo/inference/inference_engine.py b/exo/inference/inference_engine.py
index 2b5035b1..2fb4d243 100644
--- a/exo/inference/inference_engine.py
+++ b/exo/inference/inference_engine.py
@@ -1,17 +1,16 @@
import numpy as np
-import mlx.nn as nn
-from typing import Tuple
+from typing import Tuple, Optional
from abc import ABC, abstractmethod
from .shard import Shard
class InferenceEngine(ABC):
@abstractmethod
- async def infer_tensor(self, shard: Shard, input_data: np.ndarray) -> Tuple[np.ndarray, bool]:
+ async def infer_tensor(self, shard: Shard, input_data: np.ndarray, inference_state: Optional[str] = None) -> Tuple[np.ndarray, str, bool]:
pass
@abstractmethod
- async def infer_prompt(self, shard: Shard, prompt: str) -> Tuple[np.ndarray, bool]:
+ async def infer_prompt(self, shard: Shard, prompt: str, inference_state: Optional[str] = None) -> Tuple[np.ndarray, str, bool]:
pass
@abstractmethod
diff --git a/exo/inference/mlx/sharded_inference_engine.py b/exo/inference/mlx/sharded_inference_engine.py
index 133aef8e..72487949 100644
--- a/exo/inference/mlx/sharded_inference_engine.py
+++ b/exo/inference/mlx/sharded_inference_engine.py
@@ -4,6 +4,7 @@ from ..inference_engine import InferenceEngine
from .sharded_model import StatefulShardedModel
from .sharded_utils import load_shard
from ..shard import Shard
+from typing import Optional
class MLXFixedShardInferenceEngine(InferenceEngine):
def __init__(self, model_path: str, shard: Shard):
@@ -12,7 +13,7 @@ class MLXFixedShardInferenceEngine(InferenceEngine):
model_shard, self.tokenizer = load_shard(model_path, shard)
self.stateful_sharded_model = StatefulShardedModel(shard, model_shard)
- async def infer_prompt(self, shard: Shard, prompt: str) -> (np.ndarray, bool):
+ async def infer_prompt(self, shard: Shard, prompt: str, inference_state: Optional[str] = None) -> (np.ndarray, bool):
if shard != self.shard:
raise ValueError(f"Shard mismatch: {shard} != {self.shard}")
@@ -38,12 +39,12 @@ class MLXDynamicShardInferenceEngine(InferenceEngine):
def __init__(self):
self.shard = None
- async def infer_prompt(self, shard: Shard, prompt: str) -> (np.ndarray, bool):
+ async def infer_prompt(self, shard: Shard, prompt: str, inference_state: Optional[str] = None) -> (np.ndarray, bool):
await self.ensure_shard(shard)
output_data: np.ndarray = np.array(self.stateful_sharded_model.step(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, shard: Shard, input_data: np.ndarray) -> (np.ndarray, bool):
+ async def infer_tensor(self, shard: Shard, input_data: np.ndarray, inference_state: Optional[str] = None) -> (np.ndarray, bool):
await self.ensure_shard(shard)
output_data: np.ndarray = np.array(self.stateful_sharded_model.step(mx.array(input_data)))
return output_data, output_data.size == 1 and output_data.item() == self.tokenizer.eos_token_id
diff --git a/exo/inference/test_inference_engine.py b/exo/inference/test_inference_engine.py
index 191fccc5..e8df20ca 100644
--- a/exo/inference/test_inference_engine.py
+++ b/exo/inference/test_inference_engine.py
@@ -7,15 +7,17 @@ import numpy as np
# An inference engine should work the same for any number of Shards, as long as the Shards are continuous.
async def test_inference_engine(inference_engine: InferenceEngine, model_id: str, input_data: np.array):
# inference_engine.reset_shard(Shard("", 0,0,0))
- resp_full, _ = await inference_engine.infer_prompt(shard=Shard(model_id=model_id, start_layer=0, end_layer=1, n_layers=2), prompt="In one word, what is the capital of USA? ")
+ prompt = "In a single word only, what is the capital of Japan? "
+ resp_full, _, _ = await inference_engine.infer_prompt(shard=Shard(model_id=model_id, start_layer=0, end_layer=1, n_layers=2), prompt=prompt)
print("resp_full", resp_full)
print("decoded", inference_engine.tokenizer.decode(resp_full))
# inference_engine.reset_shard(Shard("", 0,0,0))
- # resp1, _ = await inference_engine.infer_tensor(shard=Shard(model_id=model_id, start_layer=0, end_layer=0, n_layers=2), input_data=input_data)
- # resp2, _ = await inference_engine.infer_tensor(shard=Shard(model_id=model_id, start_layer=1, end_layer=1, n_layers=2), input_data=resp1)
+ resp1, inference_state, _ = await inference_engine.infer_tensor(shard=Shard(model_id=model_id, start_layer=0, end_layer=0, n_layers=2), input_data=input_data)
+ print(f"Intermediate {inference_state=}")
+ resp2, _, _ = await inference_engine.infer_tensor(shard=Shard(model_id=model_id, start_layer=1, end_layer=1, n_layers=2), input_data=resp1, inference_state=inference_state)
# assert np.array_equal(resp_full, resp2)
diff --git a/exo/inference/tinygrad/inference.py b/exo/inference/tinygrad/inference.py
index 1c009a96..db3a4ed1 100644
--- a/exo/inference/tinygrad/inference.py
+++ b/exo/inference/tinygrad/inference.py
@@ -1,6 +1,6 @@
from pathlib import Path
-from typing import List
+from typing import List, Optional
import json, argparse, random, time
import tiktoken
from tiktoken.load import load_tiktoken_bpe
@@ -147,7 +147,7 @@ class TinygradDynamicShardInferenceEngine(InferenceEngine):
def __init__(self):
self.shard = None
- async def infer_prompt(self, shard: Shard, prompt: str) -> (np.ndarray, bool):
+ async def infer_prompt(self, shard: Shard, prompt: str, inference_state: Optional[str] = None) -> (np.ndarray, str, bool):
def encode_role(role: str):
return [self.tokenizer.special_tokens["<|start_header_id|>"]] + self.tokenizer.encode(role) + [self.tokenizer.special_tokens["<|end_header_id|>"]] + self.tokenizer.encode("\n\n")
def encode_message(role: str, content: str):
@@ -161,14 +161,22 @@ class TinygradDynamicShardInferenceEngine(InferenceEngine):
last_tok = toks[-1]
output_data = np.array(self.model(Tensor([[last_tok]]), start_pos, TEMPERATURE, TOP_K, TOP_P, ALPHA_F, ALPHA_P).tolist())
- start_pos += 1
+ print(f"{output_data.size=}")
+ if output_data.size == 1:
+ start_pos += 1
- return output_data, output_data.size == 1 and output_data.item() in self.tokenizer.stop_tokens
+ return output_data, json.dumps({"start_pos": start_pos}), output_data.size == 1 and output_data.item() in self.tokenizer.stop_tokens
- async def infer_tensor(self, shard: Shard, input_data: np.ndarray) -> (np.ndarray, bool):
+ async def infer_tensor(self, shard: Shard, input_data: np.ndarray, inference_state: Optional[str] = None) -> (np.ndarray, str, bool):
await self.ensure_shard(shard)
- output_data: np.ndarray = np.array(self.model(Tensor([input_data]), 0, TEMPERATURE, TOP_K, TOP_P, ALPHA_F, ALPHA_P).tolist())
- return output_data, output_data.size == 1 and output_data.item() in self.tokenizer.stop_tokens
+
+ start_pos = json.loads(inference_state)["start_pos"] if inference_state else 0
+ output_data: np.ndarray = np.array(self.model(Tensor([input_data]), start_pos, TEMPERATURE, TOP_K, TOP_P, ALPHA_F, ALPHA_P).tolist())
+ print(f"{output_data.size=}")
+ if output_data.size == 1:
+ start_pos += 1
+
+ return output_data, json.dumps({"start_pos": start_pos}), output_data.size == 1 and output_data.item() in self.tokenizer.stop_tokens
async def reset_shard(self, shard: Shard):
await self.ensure_shard(shard)
diff --git a/exo/networking/grpc/grpc_server.py b/exo/networking/grpc/grpc_server.py
index 63267839..62bf7936 100644
--- a/exo/networking/grpc/grpc_server.py
+++ b/exo/networking/grpc/grpc_server.py
@@ -44,8 +44,9 @@ class GRPCServer(node_service_pb2_grpc.NodeServiceServicer):
shard = Shard(model_id=request.shard.model_id, start_layer=request.shard.start_layer, end_layer=request.shard.end_layer, n_layers=request.shard.n_layers)
tensor = np.frombuffer(request.tensor.tensor_data, dtype=np.dtype(request.tensor.dtype)).reshape(request.tensor.shape)
request_id = request.request_id
+ inference_state = request.inference_state
- result = await self.node.process_tensor(shard, tensor, request_id)
+ result = await self.node.process_tensor(shard, tensor, request_id, inference_state)
if DEBUG >= 2: print(f"SendTensor tensor {shard=} {tensor=} {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 80295de2..31f10164 100644
--- a/exo/networking/grpc/node_service.proto
+++ b/exo/networking/grpc/node_service.proto
@@ -21,13 +21,15 @@ message Shard {
message PromptRequest {
Shard shard = 1;
string prompt = 2;
- optional string request_id = 3;
+ optional string inference_state = 3;
+ optional string request_id = 4;
}
message TensorRequest {
Shard shard = 1;
Tensor tensor = 2;
- optional string request_id = 3;
+ optional string inference_state = 3;
+ optional string request_id = 4;
}
message GetInferenceResultRequest {
diff --git a/exo/networking/grpc/node_service_pb2.py b/exo/networking/grpc/node_service_pb2.py
index bad9a291..f596e8fd 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\"k\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\x42\r\n\x0b_request_id\"\x81\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\x42\r\n\x0b_request_id\"/\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\"7\n\x11ResetShardRequest\x12\"\n\x05shard\x18\x01 \x01(\x0b\x32\x13.node_service.Shard\"<\n\x16\x43ollectTopologyRequest\x12\x0f\n\x07visited\x18\x01 \x03(\t\x12\x11\n\tmax_depth\x18\x02 \x01(\x05\"a\n\x12GlobalResetRequest\x12\'\n\nbase_shard\x18\x01 \x01(\x0b\x32\x13.node_service.Shard\x12\x0f\n\x07visited\x18\x02 \x03(\t\x12\x11\n\tmax_depth\x18\x03 \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\"A\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\"\x07\n\x05\x45mpty2\xd4\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\x44\n\nResetShard\x12\x1f.node_service.ResetShardRequest\x1a\x13.node_service.Empty\"\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\x46\n\x0bGlobalReset\x12 .node_service.GlobalResetRequest\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\"\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\x1c\n\x0finference_state\x18\x03 \x01(\tH\x00\x88\x01\x01\x12\x17\n\nrequest_id\x18\x04 \x01(\tH\x01\x88\x01\x01\x42\x12\n\x10_inference_stateB\r\n\x0b_request_id\"\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\x1c\n\x0finference_state\x18\x03 \x01(\tH\x00\x88\x01\x01\x12\x17\n\nrequest_id\x18\x04 \x01(\tH\x01\x88\x01\x01\x42\x12\n\x10_inference_stateB\r\n\x0b_request_id\"/\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\"7\n\x11ResetShardRequest\x12\"\n\x05shard\x18\x01 \x01(\x0b\x32\x13.node_service.Shard\"<\n\x16\x43ollectTopologyRequest\x12\x0f\n\x07visited\x18\x01 \x03(\t\x12\x11\n\tmax_depth\x18\x02 \x01(\x05\"a\n\x12GlobalResetRequest\x12\'\n\nbase_shard\x18\x01 \x01(\x0b\x32\x13.node_service.Shard\x12\x0f\n\x07visited\x18\x02 \x03(\t\x12\x11\n\tmax_depth\x18\x03 \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\"A\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\"\x07\n\x05\x45mpty2\xd4\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\x44\n\nResetShard\x12\x1f.node_service.ResetShardRequest\x1a\x13.node_service.Empty\"\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\x46\n\x0bGlobalReset\x12 .node_service.GlobalResetRequest\x1a\x13.node_service.Empty\"\x00\x62\x06proto3')
_globals = globals()
_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
@@ -27,34 +27,34 @@ if not _descriptor._USE_C_DESCRIPTORS:
_globals['_TOPOLOGY_PEERGRAPHENTRY']._serialized_options = b'8\001'
_globals['_SHARD']._serialized_start=36
_globals['_SHARD']._serialized_end=119
- _globals['_PROMPTREQUEST']._serialized_start=121
- _globals['_PROMPTREQUEST']._serialized_end=228
- _globals['_TENSORREQUEST']._serialized_start=231
- _globals['_TENSORREQUEST']._serialized_end=360
- _globals['_GETINFERENCERESULTREQUEST']._serialized_start=362
- _globals['_GETINFERENCERESULTREQUEST']._serialized_end=409
- _globals['_INFERENCERESULT']._serialized_start=411
- _globals['_INFERENCERESULT']._serialized_end=503
- _globals['_TENSOR']._serialized_start=505
- _globals['_TENSOR']._serialized_end=564
- _globals['_RESETSHARDREQUEST']._serialized_start=566
- _globals['_RESETSHARDREQUEST']._serialized_end=621
- _globals['_COLLECTTOPOLOGYREQUEST']._serialized_start=623
- _globals['_COLLECTTOPOLOGYREQUEST']._serialized_end=683
- _globals['_GLOBALRESETREQUEST']._serialized_start=685
- _globals['_GLOBALRESETREQUEST']._serialized_end=782
- _globals['_TOPOLOGY']._serialized_start=785
- _globals['_TOPOLOGY']._serialized_end=1055
- _globals['_TOPOLOGY_NODESENTRY']._serialized_start=906
- _globals['_TOPOLOGY_NODESENTRY']._serialized_end=984
- _globals['_TOPOLOGY_PEERGRAPHENTRY']._serialized_start=986
- _globals['_TOPOLOGY_PEERGRAPHENTRY']._serialized_end=1055
- _globals['_PEERS']._serialized_start=1057
- _globals['_PEERS']._serialized_end=1082
- _globals['_DEVICECAPABILITIES']._serialized_start=1084
- _globals['_DEVICECAPABILITIES']._serialized_end=1149
- _globals['_EMPTY']._serialized_start=1151
- _globals['_EMPTY']._serialized_end=1158
- _globals['_NODESERVICE']._serialized_start=1161
- _globals['_NODESERVICE']._serialized_end=1629
+ _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['_RESETSHARDREQUEST']._serialized_start=667
+ _globals['_RESETSHARDREQUEST']._serialized_end=722
+ _globals['_COLLECTTOPOLOGYREQUEST']._serialized_start=724
+ _globals['_COLLECTTOPOLOGYREQUEST']._serialized_end=784
+ _globals['_GLOBALRESETREQUEST']._serialized_start=786
+ _globals['_GLOBALRESETREQUEST']._serialized_end=883
+ _globals['_TOPOLOGY']._serialized_start=886
+ _globals['_TOPOLOGY']._serialized_end=1156
+ _globals['_TOPOLOGY_NODESENTRY']._serialized_start=1007
+ _globals['_TOPOLOGY_NODESENTRY']._serialized_end=1085
+ _globals['_TOPOLOGY_PEERGRAPHENTRY']._serialized_start=1087
+ _globals['_TOPOLOGY_PEERGRAPHENTRY']._serialized_end=1156
+ _globals['_PEERS']._serialized_start=1158
+ _globals['_PEERS']._serialized_end=1183
+ _globals['_DEVICECAPABILITIES']._serialized_start=1185
+ _globals['_DEVICECAPABILITIES']._serialized_end=1250
+ _globals['_EMPTY']._serialized_start=1252
+ _globals['_EMPTY']._serialized_end=1259
+ _globals['_NODESERVICE']._serialized_start=1262
+ _globals['_NODESERVICE']._serialized_end=1730
# @@protoc_insertion_point(module_scope)
diff --git a/exo/orchestration/node.py b/exo/orchestration/node.py
index 8dac1085..13e8d785 100644
--- a/exo/orchestration/node.py
+++ b/exo/orchestration/node.py
@@ -14,11 +14,11 @@ class Node(ABC):
pass
@abstractmethod
- async def process_prompt(self, shard: Shard, prompt: str, request_id: Optional[str] = None) -> Optional[np.ndarray]:
+ async def process_prompt(self, shard: Shard, prompt: str, request_id: Optional[str] = None, inference_state: Optional[str] = None) -> Optional[np.ndarray]:
pass
@abstractmethod
- async def process_tensor(self, shard: Shard, tensor: np.ndarray, request_id: Optional[str] = None) -> Optional[np.ndarray]:
+ async def process_tensor(self, shard: Shard, tensor: np.ndarray, 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 5fd194d5..c6ede1e4 100644
--- a/exo/orchestration/standard_node.py
+++ b/exo/orchestration/standard_node.py
@@ -37,7 +37,7 @@ class StandardNode(Node):
await self.discovery.stop()
await self.server.stop()
- async def process_prompt(self, shard: Shard, prompt: str, request_id: Optional[str] = None) -> Optional[np.ndarray]:
+ async def process_prompt(self, shard: Shard, prompt: str, 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:
@@ -49,7 +49,7 @@ class StandardNode(Node):
await self.forward_to_next_shard(shard, prompt, request_id)
return
- result, is_finished = await self.inference_engine.infer_prompt(self.get_current_shard(shard), prompt)
+ result, inference_state, is_finished = await self.inference_engine.infer_prompt(self.get_current_shard(shard), prompt, 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)
@@ -61,11 +61,11 @@ class StandardNode(Node):
if DEBUG >= 2: print(f"[{request_id}] result size: {result.size}, is finished: {is_finished}, buffered tokens: {len(self.buffered_token_output[request_id])}")
if not is_finished:
- asyncio.create_task(self.forward_to_next_shard(shard, result, request_id))
+ asyncio.create_task(self.forward_to_next_shard(shard, result, request_id, inference_state=inference_state))
return np.array(self.buffered_token_output[request_id]) if len(self.buffered_token_output[request_id]) > 0 else None
- async def process_tensor(self, shard: Shard, tensor: np.ndarray, request_id: Optional[str] = None) -> Optional[np.ndarray]:
+ async def process_tensor(self, shard: Shard, tensor: np.ndarray, 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:
@@ -73,7 +73,7 @@ class StandardNode(Node):
try:
if DEBUG >= 1: print(f"[{request_id}] process_tensor: {tensor.size=} {tensor.shape=}")
- result, is_finished = await self.inference_engine.infer_tensor(self.get_current_shard(shard), tensor)
+ result, inference_state, is_finished = await self.inference_engine.infer_tensor(self.get_current_shard(shard), tensor, 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)
@@ -84,7 +84,7 @@ class StandardNode(Node):
if DEBUG >= 2: print(f"[{request_id}] result size: {result.size}, is finished: {is_finished}, buffered tokens: {len(self.buffered_token_output[request_id])}")
if not is_finished:
- asyncio.create_task(self.forward_to_next_shard(shard, result, request_id))
+ asyncio.create_task(self.forward_to_next_shard(shard, result, request_id, inference_state=inference_state))
return np.array(self.buffered_token_output[request_id][0]) if len(self.buffered_token_output[request_id][0]) > 0 else None
except Exception as e:
@@ -93,7 +93,7 @@ class StandardNode(Node):
traceback.print_exc()
return None
- async def forward_to_next_shard(self, shard: Shard, tensor_or_prompt: Union[np.ndarray, str], request_id: str) -> None:
+ async def forward_to_next_shard(self, shard: Shard, tensor_or_prompt: Union[np.ndarray, str], request_id: str, inference_state: Optional[str] = None) -> None:
if not self.partitioning_strategy:
if DEBUG >= 1: print("No partitioning strategy found. Skipping forward.")
return
@@ -109,9 +109,9 @@ class StandardNode(Node):
if next_partition:
if next_partition.node_id == self.id:
if isinstance(tensor_or_prompt, np.ndarray):
- await self.process_tensor(shard, tensor_or_prompt, request_id)
+ 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)
+ await self.process_prompt(shard, tensor_or_prompt, request_id, inference_state=inference_state)
return
target_peer = next((p for p in self.peers if p.id() == next_partition.node_id), None)
← ed7672e3 Update README.md
·
back to Exo
·
add an opaque inference_state that inference engines can use dd8d1812 →