From dd3187fc8aba98bd69b5ff1fa48a0f41732f2688 Mon Sep 17 00:00:00 2001 From: zhangzhen Date: Thu, 14 Sep 2023 15:22:38 +0800 Subject: [PATCH] Provide the local llama service --- src/aios_kernel/compute_kernel.py | 2 +- src/aios_kernel/compute_node.py | 5 +- src/aios_kernel/local_llama_compute_node.py | 86 +++++++++++++++++++++ src/aios_kernel/open_ai_node.py | 4 +- src/aios_kernel/queue_compute_node.py | 69 +++++++++++++++++ src/aios_kernel/stability_node.py | 4 +- 6 files changed, 162 insertions(+), 8 deletions(-) create mode 100644 src/aios_kernel/local_llama_compute_node.py create mode 100644 src/aios_kernel/queue_compute_node.py diff --git a/src/aios_kernel/compute_kernel.py b/src/aios_kernel/compute_kernel.py index 5fea958..381dda9 100644 --- a/src/aios_kernel/compute_kernel.py +++ b/src/aios_kernel/compute_kernel.py @@ -65,7 +65,7 @@ class ComputeKernel: def _schedule(self, task) -> ComputeNode: for node in self.compute_nodes.values(): - if node.is_support(task.task_type) is True: + if node.is_support(task) is True: return node logger.warning( f"task {task.display()} is not support by any compute node") diff --git a/src/aios_kernel/compute_node.py b/src/aios_kernel/compute_node.py index 7e159f7..873fecd 100644 --- a/src/aios_kernel/compute_node.py +++ b/src/aios_kernel/compute_node.py @@ -28,7 +28,7 @@ class ComputeNode(ABC): pass @abstractmethod - def is_support(self, task_type: ComputeTaskType) -> bool: + def is_support(self, task: ComputeTask) -> bool: pass @abstractmethod @@ -41,10 +41,9 @@ class ComputeNode(ABC): def get_fee_type(self) -> str: return "free" - class LocalComputeNode(ComputeNode): def display(self) -> str: return super().display() def is_local(self) -> bool: - return True + return True \ No newline at end of file diff --git a/src/aios_kernel/local_llama_compute_node.py b/src/aios_kernel/local_llama_compute_node.py new file mode 100644 index 0000000..0931f1f --- /dev/null +++ b/src/aios_kernel/local_llama_compute_node.py @@ -0,0 +1,86 @@ + +import logging +import requests +from typing import Optional +from pydantic import BaseModel + +from .compute_task import ComputeTask, ComputeTaskState, ComputeTaskType +from .queue_compute_node import Queue_ComputeNode + +logger = logging.getLogger(__name__) + +""" +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[str] + + try: + body = { + "prompts": task.params["prompts"] + } + + response = requests.post("http://aigc:7880/generate", data = body, verify=False) + 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: " + response.status_code + } + } + else: + resp = GenerateResponse.parse_raw(response.content.decode("utf-8")) + if resp.error: + return { + "state": ComputeTaskState.ERROR, + "error": { + "code": resp.error, + "message": "local llama failed:" + resp.msg + } + } + else: + return { + "content": str(resp.results), + "message": {} + } + 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) + } + } + + 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.LLM_COMPLETION and (not task.params["model_name"] or task.params["model_name"] == "llama") + + def is_local(self) -> bool: + return True diff --git a/src/aios_kernel/open_ai_node.py b/src/aios_kernel/open_ai_node.py index 1be47fd..e2ef098 100644 --- a/src/aios_kernel/open_ai_node.py +++ b/src/aios_kernel/open_ai_node.py @@ -103,8 +103,8 @@ class OpenAI_ComputeNode(ComputeNode): def get_capacity(self): pass - def is_support(self, task_type: ComputeTaskType) -> bool: - return task_type == ComputeTaskType.LLM_COMPLETION + 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"] == "gpt-4-0613") def is_local(self) -> bool: return False diff --git a/src/aios_kernel/queue_compute_node.py b/src/aios_kernel/queue_compute_node.py new file mode 100644 index 0000000..e975cd8 --- /dev/null +++ b/src/aios_kernel/queue_compute_node.py @@ -0,0 +1,69 @@ + +import asyncio +from asyncio import Queue +import logging +from abc import abstractmethod + +from .compute_task import ComputeTask, ComputeTaskResult, ComputeTaskState, ComputeTaskType +from .compute_node import ComputeNode + +logger = logging.getLogger(__name__) + +class Queue_ComputeNode(ComputeNode): + def __init__(self): + super().__init__() + self.task_queue = Queue() + + @abstractmethod + async def execute_task(self, task: ComputeTask) -> { + "content": str, + "message": str, + "state": ComputeTaskState, + "error": { + "code": int, + "message": str, + } + }: + pass + + async def push_task(self, task: ComputeTask, proiority: int = 0): + logger.info(f"{self.display()} push task: {task.display()}") + self.task_queue.put_nowait(task) + + async def remove_task(self, task_id: str): + pass + + async def _run_task(self, task: ComputeTask): + task.state = ComputeTaskState.RUNNING + + resp = await self.execute_task(task) + + result = ComputeTaskResult() + result.set_from_task(task) + + task.state = resp["state"] + + if task.state == ComputeTaskState.ERROR: + task.error_str = resp["error"]["message"] + + + result.worker_id = self.node_id + result.result_str = resp["content"] + result.result_message = resp["message"] + + return result + + def start(self): + async def _run_task_loop(): + while True: + task = await self.task_queue.get() + logger.info(f"{self.display()} get task: {task.display()}") + result = await self._run_task(task) + if result is not None: + task.result = result + + asyncio.create_task(_run_task_loop()) + + + def get_task_state(self, task_id: str): + pass diff --git a/src/aios_kernel/stability_node.py b/src/aios_kernel/stability_node.py index 8d40878..b1f0da9 100644 --- a/src/aios_kernel/stability_node.py +++ b/src/aios_kernel/stability_node.py @@ -129,8 +129,8 @@ class Stability_ComputeNode(ComputeNode): def get_capacity(self): pass - def is_support(self, task_type: ComputeTaskType) -> bool: - return task_type == ComputeTaskType.TEXT_2_IMAGE + def is_support(self, task: ComputeTask) -> bool: + return task.task_type == ComputeTaskType.TEXT_2_IMAGE def is_local(self) -> bool: return False