gyt/opendan
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

rebase to main

Browse Source
This commit is contained in:
streetycat
2023-09-28 09:09:30 +00:00
parent 4f6b04fe48
commit 193c627bdb
4 changed files with 255 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/aios_kernel/compute_node_config.py
+87
View File
@@ -0,0 +1,87 @@
"""
Configuration for nodes:
```
├── nodes
│ └── llama
| └── 0
| | └── url
| | └── model_name
| └── 1
| └── url
| └── model_name
```
"""
import logging
from typing import List
import os
import toml
from .local_llama_compute_node import LocalLlama_ComputeNode
from .storage import AIStorage
# define singleton class knowledge pipline
class ComputeNodeConfig:
_instance = None
@classmethod
def get_instance(cls):
if cls._instance is None:
cls._instance = ComputeNodeConfig()
cls._instance.__singleton_init__()
return cls._instance
def initial(self) -> List[LocalLlama_ComputeNode]:
config_path = self.__config_path()
logging.info(f"initial nodes from {config_path}")
if os.path.exists(config_path):
self.config = toml.load(self.__config_path())
if self.config is None:
return []
nodes = []
llama_nodes_cfg = self.config["llama"]
if llama_nodes_cfg is not None:
for cfg in llama_nodes_cfg:
node = LocalLlama_ComputeNode(url=cfg["url"], model_name=cfg["model_name"])
nodes.append(node)
return nodes
return []
def save(self):
with open(self.__config_path(), "w") as f:
toml.dump(self.config, f)
def add_node(self, model_type: str, url: str, model_name: str):
if model_type == "llama":
llama_nodes_cfg = self.config.get("llama") or []
for cfg in llama_nodes_cfg:
if url == cfg["url"] and model_name == cfg["model_name"]:
return
llama_nodes_cfg.append({"url": url, "model_name": model_name})
self.config["llama"] = llama_nodes_cfg
def remove_node(self, model_type: str, url: str, model_name: str):
if model_type == "llama":
llama_nodes_cfg = self.config.get("llama") or []
for i in range(0, len(llama_nodes_cfg)):
cfg = llama_nodes_cfg[i]
if url == cfg["url"] and model_name == cfg["model_name"]:
llama_nodes_cfg.pop(i)
def list(self) -> str:
return toml.dumps(self.config)
def __singleton_init__(self):
self.config = {}
@classmethod
def __config_path(cls) -> str:
user_data_dir = AIStorage.get_instance().get_myai_dir()
return os.path.abspath(f"{user_data_dir}/etc/compute_nodes.cfg.toml")
src/aios_kernel/local_llama_compute_node.py
+106 -59
View File
@@ -4,10 +4,10 @@ import logging
import requests
from typing import Optional, List
from pydantic import BaseModel
from llama_cpp import Llama
from .compute_task import ComputeTask, ComputeTaskResult, ComputeTaskResultCode, ComputeTaskState, ComputeTaskType
from .queue_compute_node import Queue_ComputeNode
from .storage import AIStorage,UserConfig
logger = logging.getLogger(__name__)
@@ -16,59 +16,117 @@ This is a custom implementation, it should be redesigned.
"""
class LocalLlama_ComputeNode(Queue_ComputeNode):
def __init__(self, model_path: str, model_name: str):
def __init__(self, url: str, model_name: str):
super().__init__()
self.model_path = model_path
self.url = url
self.model_name = model_name
self.llm = Llama(model_path=model_path)
async def execute_task(self, task: ComputeTask, result: ComputeTaskResult) -> ComputeTaskResult:
async def execute_task(self, task: ComputeTask, result: ComputeTaskResult):
match task.task_type:
case ComputeTaskType.TEXT_EMBEDDING:
model_name = task.params["model_name"]
input = task.params["input"]
logger.info(f"call local-llama {model_name} input: {input}")
logger.info(f"call local-llama ({self.url}, {self.model_name}) {model_name} input: {input}")
try:
embedding = self.llm.embed(input=input)
logger.info(f"local-llama({self.model_path}) response: {embedding}")
except Exception as e:
logger.error(f"call local-llama {model_name} run TEXT_EMBEDDING task error: {e}")
task.state = ComputeTaskState.ERROR
task.error_str = str(e)
result.error_str = str(e)
return result
self.embedding(input, result)
logger.info(f"local-llama({self.model_path}) response: {embedding}")
task.state = ComputeTaskState.DONE
result.result_code = ComputeTaskResultCode.OK
result.result = embedding
if result.result_code == ComputeTaskResultCode.OK:
task.state = ComputeTaskState.DONE
else:
task.state = ComputeTaskState.ERROR
task.error_str = result.error_str
return result
case ComputeTaskType.LLM_COMPLETION:
mode_name = task.params["model_name"]
prompts = task.params["prompts"]
max_token_size = task.params.get("max_token_size")
llm_inner_functions = task.params.get("inner_functions")
if max_token_size is None:
max_token_size = 4000
logger.info(f"local-llama({self.model_path}) prompts: {prompts}")
logger.info(f"local-llama({self.url}, {self.model_name}) prompts: {prompts}")
try:
resp = self.llm.create_chat_completion(model=mode_name,
messages=prompts,
functions=llm_inner_functions, # function has not support?
max_tokens=max_token_size,
temperature=0.7) # TODO: add temperature to task params?
except Exception as e:
logger.error(f"local-llama({self.model_path}) run LLM_COMPLETION task error: {e}")
self.completion(task, result)
if result.result_code == ComputeTaskResultCode.OK:
task.state = ComputeTaskState.DONE
else:
task.state = ComputeTaskState.ERROR
task.error_str = str(e)
result.error_str = str(e)
return result
task.error_str = result.error_str
case _:
task.state = ComputeTaskState.ERROR
result.result_code = ComputeTaskResultCode.ERROR
task.error_str = f"ComputeTask's TaskType : {task.task_type} not support!"
result.error_str = f"ComputeTask's TaskType : {task.task_type} not support!"
return None
async def initial(self) -> bool:
return True
def display(self) -> str:
return f"local-llama: {self.node_id}"
def get_capacity(self):
pass
def is_support(self, task: ComputeTask) -> bool:
return (task.task_type == ComputeTaskType.TEXT_EMBEDDING or task.task_type == ComputeTaskType.LLM_COMPLETION) and (not task.params["model_name"] or task.params["model_name"] == self.model_name)
def is_local(self) -> bool:
return True
def embedding(self, input: str, result: ComputeTaskResult):
body = {
"input": input
}
try:
response = requests.post(self.url + "/v1/embeddings", json = body, verify=False, headers={"Content-Type": "application/json"})
response.close()
logger.info(f"local-llama({self.url}, {self.model_name}) task responsed, request: {body}, status-code: {response.status_code}, headers: {response.headers}, content: {response.content}")
if response.status_code == 200:
resp = response.json()
result.result = resp["data"][0]["embedding"]
elif response.status_code == 422:
resp = response.json()
result.result_code = ComputeTaskResultCode.ERROR
result.error_str = "http request failed: " + str(resp["detail"][0]["msg"])
else:
result.result_code = ComputeTaskResultCode.ERROR
result.error_str = "http request failed: " + str(response.status_code)
except Exception as e:
logger.error(f"call local-llama({self.url}, {self.model_name}) run TEXT_EMBEDDING task error: {e}")
result.result_code = ComputeTaskResultCode.ERROR
result.error_str = str(e)
return result
def completion(self, task: ComputeTask, result: ComputeTaskResult):
mode_name = task.params["model_name"]
prompts = task.params["prompts"]
max_token_size = task.params.get("max_token_size")
llm_inner_functions = task.params.get("inner_functions")
if max_token_size is None:
max_token_size = max_token_size
logger.info(f"local-llama({self.model_path}) response: {json.dumps(resp, indent=4)}")
body = {
"messages": [],
"max_tokens": 4000
}
for prompt in prompts:
body["messages"].append({
"role": prompt["role"],
"content": prompt["content"]
})
try:
response = requests.post(self.url + "/v1/chat/completions", json = body, verify=False, headers={"Content-Type": "application/json"})
response.close()
logger.info(f"local-llama({self.url}, {self.model_name}) task responsed, request: {body}, status-code: {response.status_code}, headers: {response.headers}, content: {response.content}")
if response.status_code == 200:
resp = response.json()
status_code = resp["choices"][0]["finish_reason"]
token_usage = resp["usage"]
@@ -91,27 +149,16 @@ class LocalLlama_ComputeNode(Queue_ComputeNode):
if token_usage:
result.result_refers["token_usage"] = token_usage
logger.info(f"local-llama({self.model_path}) success response: {result.result_str}")
return result
case _:
task.state = ComputeTaskState.ERROR
logger.info(f"local-llama({self.url}, {self.model_name}) success response: {result.result_str}")
elif response.status_code == 422:
resp = response.json()
result.result_code = ComputeTaskResultCode.ERROR
task.error_str = f"ComputeTask's TaskType : {task.task_type} not support!"
result.error_str = f"ComputeTask's TaskType : {task.task_type} not support!"
return None
async def initial(self) -> bool:
return True
def display(self) -> str:
return f"LocalLlama_ComputeNode: {self.node_id}"
def get_capacity(self):
pass
def is_support(self, task: ComputeTask) -> bool:
return (task.task_type == ComputeTaskType.TEXT_EMBEDDING or task.task_type == ComputeTaskType.LLM_COMPLETION) and (not task.params["model_name"] or task.params["model_name"] == self.model_name)
def is_local(self) -> bool:
return True
result.error_str = "http request failed: " + str(resp["detail"][0]["msg"])
else:
result.result_code = ComputeTaskResultCode.ERROR
result.error_str = "http request failed: " + str(response.status_code)
except Exception as e:
logger.error(f"call local-llama({self.url}, {self.model_name}) run LLM_COMPLETION task error: {e}")
result.result_code = ComputeTaskResultCode.ERROR
result.error_str = str(e)
return result