Merge pull request #70 from streetycat/MVP

Add the llama nodes at local dynamically
2023-09-28 09:17:39 -07:00
parent 14a76476ef 861096fb0c
commit 8ccab139aa
5 changed files with 290 additions and 95 deletions
@@ -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")
@@ -1,11 +1,13 @@

+import json
 import logging
 import requests
 from typing import Optional, List
 from pydantic import BaseModel

-from .compute_task import ComputeTask, ComputeTaskState, ComputeTaskType
+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__)

@@ -14,81 +16,149 @@ This is a custom implementation, it should be redesigned.
 """

 class LocalLlama_ComputeNode(Queue_ComputeNode):
-    async def execute_task(self, task: ComputeTask) -> {
-        "content": str,
-        "message": str,
-        "state": ComputeTaskState,
-        "error": {
-            "code": int,
-            "message": str,
-        }
-    }:
-        class GenerateResponse(BaseModel):
-            error: Optional[int]
-            msg: Optional[str]
-            results: Optional[List[str]]
+    def __init__(self, url: str, model_name: str):
+        super().__init__()
+        self.url = url
+        self.model_name = model_name

-        try:
-            prompt_msgs = []
-            for prompt in task.params["prompts"]:
-                prompt_msgs.append(prompt["content"])
+    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 ({self.url}, {self.model_name}) {model_name} input: {input}")
+
+                self.embedding(input, result)
                
-            body = {
-                "prompts": prompt_msgs
-            }
-            
-            response = requests.post("http://aigc:7880/generate", json = body, verify=False, headers={"Content-Type": "application/json"})
-            response.close()
-
-            logger.info(f"LocalLlama_ComputeNode task responsed, request: {body}, status-code: {response.status_code}, headers: {response.headers}, content: {response.content}")
-
-            if response.status_code != 200:
-                return {
-                    "state": ComputeTaskState.ERROR,
-                    "error": {
-                        "code": response.status_code,
-                        "message": "http request failed: " + str(response.status_code)
-                    }
-                }
-            else:
-                resp = response.json()
-                if "error" in resp:
-                    return {
-                        "state": ComputeTaskState.ERROR,
-                        "error": {
-                            "code": resp["error"],
-                            "message": "local llama failed:" + resp["msg"]
-                        }
-                    }
+                if result.result_code == ComputeTaskResultCode.OK:
+                    task.state = ComputeTaskState.DONE
                else:
-                    return {
-                        "state": ComputeTaskState.DONE,
-                        "content": str(resp["results"]),
-                        "message": str(resp["results"])
-                    }
-        except Exception as err:
-            import traceback
-            logger.error(f"{traceback.format_exc()}, error: {err}")
-            
-            return {
-                "state": ComputeTaskState.ERROR,
-                "error": {
-                    "code": -1,
-                    "message": "unknown exception: " + str(err)
-                }
-            }
+                    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"]
+                    
+                logger.info(f"local-llama({self.url}, {self.model_name}) prompts: {prompts}")
+
+                self.completion(task, result)
+
+                if result.result_code == ComputeTaskResultCode.OK:
+                    task.state = ComputeTaskState.DONE
+                else:
+                    task.state = ComputeTaskState.ERROR
+                    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"LocalLlama_ComputeNode: {self.node_id}"
+        return f"local-llama: {self.node_id}"

    def get_capacity(self):
        pass

    def is_support(self, task: ComputeTask) -> bool:
-        return task.task_type == ComputeTaskType.LLM_COMPLETION and (not task.params["model_name"] or task.params["model_name"] == "llama")
+        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
+            
+        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"]
+
+                match status_code:
+                    case "function_call":
+                        task.state = ComputeTaskState.DONE
+                    case "stop":
+                        task.state = ComputeTaskState.DONE
+                    case _:
+                        task.state = ComputeTaskState.ERROR
+                        task.error_str = f"The status code was {status_code}."
+                        result.error_str = f"The status code was {status_code}."
+                        result.result_code = ComputeTaskResultCode.ERROR
+                        return None
+                    
+                result.result_code = ComputeTaskResultCode.OK
+                result.result_str = resp["choices"][0]["message"]["content"]
+                result.result_message = resp["choices"][0]["message"]
+                if token_usage:
+                    result.result_refers["token_usage"] = token_usage
+
+                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
+                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
@@ -16,15 +16,7 @@ class Queue_ComputeNode(ComputeNode):
        self.is_start = False

    @abstractmethod
-    async def execute_task(self, task: ComputeTask) -> {
-        "content": str,
-        "message": str,
-        "state": ComputeTaskState,
-        "error": {
-            "code": int,
-            "message": str,
-        }
-    }:
+    async def execute_task(self, task: ComputeTask, result: ComputeTaskResult):
        pass

    async def push_task(self, task: ComputeTask, proiority: int = 0):
@@ -36,24 +28,15 @@ class Queue_ComputeNode(ComputeNode):

    async def _run_task(self, task: ComputeTask):
        task.state = ComputeTaskState.RUNNING
-        
-        resp = await self.execute_task(task)
+

        result = ComputeTaskResult()
-
-        result.worker_id = self.node_id
-        task.state = resp["state"]
-        
-        if task.state == ComputeTaskState.ERROR:
-            result.result_code = ComputeTaskResultCode.ERROR
-            task.error_str = resp["error"]["message"]
-        else:
-            result.result_code = ComputeTaskResultCode.OK
-            result.result_str = resp["content"]
-            result.result_message = resp["message"]
-
+        result.result_code = ComputeTaskResultCode.ERROR
        result.set_from_task(task)
+        result.worker_id = self.node_id

+        await self.execute_task(task, result)
+            
        return result

    def start(self):