From 97d69e43f76219ead5accfadceb547b7c2fe6ba2 Mon Sep 17 00:00:00 2001 From: streetycat <305190374@qq.com> Date: Thu, 28 Sep 2023 05:03:06 +0000 Subject: [PATCH] fix bugs --- src/aios_kernel/local_llama_compute_node.py | 60 +++++++++++++++------ src/aios_kernel/queue_compute_node.py | 16 +++--- 2 files changed, 51 insertions(+), 25 deletions(-) diff --git a/src/aios_kernel/local_llama_compute_node.py b/src/aios_kernel/local_llama_compute_node.py index bca0c22..a16c3a3 100644 --- a/src/aios_kernel/local_llama_compute_node.py +++ b/src/aios_kernel/local_llama_compute_node.py @@ -6,7 +6,7 @@ from typing import Optional, List from pydantic import BaseModel from llama_cpp import Llama -from .compute_task import ComputeTask, ComputeTaskResult, ComputeTaskState, ComputeTaskType +from .compute_task import ComputeTask, ComputeTaskResult, ComputeTaskResultCode, ComputeTaskState, ComputeTaskType from .queue_compute_node import Queue_ComputeNode logger = logging.getLogger(__name__) @@ -22,19 +22,26 @@ class LocalLlama_ComputeNode(Queue_ComputeNode): self.model_name = model_name self.llm = Llama(model_path=model_path) - async def execute_task(self, task: ComputeTask) -> ComputeTaskResult: + async def execute_task(self, task: ComputeTask, result: ComputeTaskResult) -> ComputeTaskResult: match task.task_type: case ComputeTaskType.TEXT_EMBEDDING: model_name = task.params["model_name"] input = task.params["input"] - logger.info(f"call openai {model_name} input: {input}") + logger.info(f"call local-llama {model_name} input: {input}") - embedding = self.llm.embed(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 - logger.info(f"local-llama({self.model_path}) response: {resp}") - - result = ComputeTaskResult() - result.set_from_task(task) + logger.info(f"local-llama({self.model_path}) response: {embedding}") + task.state = ComputeTaskState.DONE + result.result_code = ComputeTaskResultCode.OK result.result = embedding return result @@ -48,19 +55,24 @@ class LocalLlama_ComputeNode(Queue_ComputeNode): logger.info(f"local-llama({self.model_path}) prompts: {prompts}") - 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? + 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}") + task.state = ComputeTaskState.ERROR + task.error_str = str(e) + result.error_str = str(e) + return result logger.info(f"local-llama({self.model_path}) response: {json.dumps(resp, indent=4)}") - result = ComputeTaskResult() - result.set_from_task(task) - status_code = resp["choices"][0]["finish_reason"] + token_usage = resp["usage"] match status_code: case "function_call": task.state = ComputeTaskState.DONE @@ -69,11 +81,25 @@ class LocalLlama_ComputeNode(Queue_ComputeNode): 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.model_path}) success response: {result.result_str}") + return result + 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 diff --git a/src/aios_kernel/queue_compute_node.py b/src/aios_kernel/queue_compute_node.py index 6b3469b..261aaa9 100644 --- a/src/aios_kernel/queue_compute_node.py +++ b/src/aios_kernel/queue_compute_node.py @@ -16,7 +16,7 @@ class Queue_ComputeNode(ComputeNode): self.is_start = False @abstractmethod - async def execute_task(self, task: ComputeTask) -> ComputeTaskResult: + async def execute_task(self, task: ComputeTask, result: ComputeTaskResult): pass async def push_task(self, task: ComputeTask, proiority: int = 0): @@ -28,13 +28,13 @@ class Queue_ComputeNode(ComputeNode): async def _run_task(self, task: ComputeTask): task.state = ComputeTaskState.RUNNING - - result = await self.execute_task(task) - if result is not None: - result.set_from_task(task) - result.worker_id = self.node_id - else: - task.state = ComputeTaskState.ERROR + + result = ComputeTaskResult() + result.result_code = ComputeTaskResultCode.ERROR + result.set_from_task(task) + result.worker_id = self.node_id + + await self.execute_task(task, result) return result