From 3eac598c972c7a8c3a1ecc01cae76202118f898d Mon Sep 17 00:00:00 2001 From: Song Date: Thu, 7 Sep 2023 12:50:13 +0800 Subject: [PATCH 1/3] Add stability node (Text2Img) --- src/aios_kernel/compute_node.py | 25 +++--- src/aios_kernel/compute_task.py | 42 +++++---- src/aios_kernel/open_ai_node.py | 55 +++++------- src/aios_kernel/stability_node.py | 140 ++++++++++++++++++++++++++++++ src/requirements.txt | 8 ++ 5 files changed, 209 insertions(+), 61 deletions(-) create mode 100644 src/aios_kernel/stability_node.py create mode 100644 src/requirements.txt diff --git a/src/aios_kernel/compute_node.py b/src/aios_kernel/compute_node.py index 6fe0059..7e159f7 100644 --- a/src/aios_kernel/compute_node.py +++ b/src/aios_kernel/compute_node.py @@ -1,5 +1,5 @@ from abc import ABC, abstractmethod -from .compute_task import ComputeTask +from .compute_task import ComputeTask, ComputeTaskType class ComputeNode(ABC): @@ -8,15 +8,15 @@ class ComputeNode(ABC): self.enable = True @abstractmethod - async def push_task(self,task:ComputeTask,proiority:int = 0): - pass - - @abstractmethod - async def remove_task(self,task_id:str): + async def push_task(self, task: ComputeTask, proiority: int = 0): pass @abstractmethod - def get_task_state(self,task_id:str): + async def remove_task(self, task_id: str): + pass + + @abstractmethod + def get_task_state(self, task_id: str): pass @abstractmethod @@ -28,7 +28,7 @@ class ComputeNode(ABC): pass @abstractmethod - def is_support(self,task_type:str) -> bool: + def is_support(self, task_type: ComputeTaskType) -> bool: pass @abstractmethod @@ -37,17 +37,14 @@ class ComputeNode(ABC): def is_trusted(self) -> bool: return True - + 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 - - diff --git a/src/aios_kernel/compute_task.py b/src/aios_kernel/compute_task.py index f2253e0..7895283 100644 --- a/src/aios_kernel/compute_task.py +++ b/src/aios_kernel/compute_task.py @@ -3,6 +3,7 @@ from enum import Enum import uuid import time + class ComputeTaskState(Enum): DONE = 0 INIT = 1 @@ -11,22 +12,31 @@ class ComputeTaskState(Enum): PENDING = 4 +class ComputeTaskType(Enum): + NONE = -1 + LLM_COMPLETION = 0 + TEXT_2_IMAGE = 1 + IMAGE_2_IMAGE = 2 + VOICE_2_TEXT = 3 + TEXT_2_VOICE = 4 + + class ComputeTask: def __init__(self) -> None: self.task_type = "llm_completion" self.create_time = None - self.task_id:str = None - self.callchain_id:str = None - self.params:dict = {} - self.refers:dict = None - self.pading_data:bytearray = None + self.task_id: str = None + self.callchain_id: str = None + self.params: dict = {} + self.refers: dict = None + self.pading_data: bytearray = None self.state = ComputeTaskState.INIT self.result = None self.error_str = None - def set_llm_params(self,prompts,model_name,max_token_size,callchain_id = None): + def set_llm_params(self, prompts, model_name, max_token_size, callchain_id=None): self.task_type = "llm_completion" self.create_time = time.time() self.task_id = uuid.uuid4().hex @@ -34,7 +44,7 @@ class ComputeTask: self.params["prompts"] = prompts.messages if model_name is not None: self.params["model_name"] = model_name - else: + else: self.params["model_name"] = "gpt-4-0613" self.params["max_token_size"] = max_token_size @@ -45,16 +55,16 @@ class ComputeTask: class ComputeTaskResult: def __init__(self) -> None: self.create_time = None - self.task_id:str = None - self.callchain_id:str = None - self.worker_id:str = None - self.result_code:int = 0 - self.result_str:str = None + self.task_id: str = None + self.callchain_id: str = None + self.worker_id: str = None + self.result_code: int = 0 + self.result_str: str = None - self.result:dict = {} - self.result_refers:dict = None - self.pading_data:bytearray = None + self.result: dict = {} + self.result_refers: dict = None + self.pading_data: bytearray = None - def set_from_task(self,task:ComputeTask): + def set_from_task(self, task: ComputeTask): self.task_id = task.task_id self.callchain_id = task.callchain_id diff --git a/src/aios_kernel/open_ai_node.py b/src/aios_kernel/open_ai_node.py index 7f5c8e6..06e0f09 100644 --- a/src/aios_kernel/open_ai_node.py +++ b/src/aios_kernel/open_ai_node.py @@ -5,47 +5,49 @@ import asyncio from asyncio import Queue import logging -from .compute_task import ComputeTask,ComputeTaskResult,ComputeTaskState +from .compute_task import ComputeTask, ComputeTaskResult, ComputeTaskState, ComputeTaskType from .compute_node import ComputeNode logger = logging.getLogger(__name__) + class OpenAI_ComputeNode(ComputeNode): _instance = None + def __new__(cls): if cls._instance is None: cls._instance = super(OpenAI_ComputeNode, cls).__new__(cls) cls._instance.is_start = False return cls._instance - + def __init__(self) -> None: super().__init__() if self.is_start is True: logger.warn("OpenAI_ComputeNode is already start") return - + self.is_start = True - #openai.organization = "org-AoKrOtF2myemvfiFfnsSU8rF" #buckycloud + # openai.organization = "org-AoKrOtF2myemvfiFfnsSU8rF" #buckycloud self.openai_api_key = "" self.node_id = "openai_node" self.task_queue = Queue() - if os.getenv("OPENAI_API_KEY") is not None: + if os.getenv("OPENAI_API_KEY") is not None: openai.api_key = os.getenv("OPENAI_API_KEY") else: openai.api_key = self.openai_api_key - + self.start() - - async def push_task(self,task:ComputeTask,proiority:int = 0): + + async def push_task(self, task: ComputeTask, proiority: int = 0): logger.info(f"openai_node push task: {task.display()}") self.task_queue.put_nowait(task) - - async def remove_task(self,task_id:str): + + async def remove_task(self, task_id: str): pass - - def _run_task(self,task:ComputeTask): + + def _run_task(self, task: ComputeTask): task.state = ComputeTaskState.RUNNING mode_name = task.params["model_name"] # max_token_size = task.params["max_token_size"] @@ -57,19 +59,19 @@ class OpenAI_ComputeNode(ComputeNode): max_tokens=4000, temperature=1.2) logger.info(f"openai response: {resp}") - + status_code = resp["choices"][0]["finish_reason"] if status_code != "stop": task.state = ComputeTaskState.ERROR - task.error_str =f"The status code was {status_code}." + task.error_str = f"The status code was {status_code}." return None - - result = ComputeTaskResult() + + result = ComputeTaskResult() result.set_from_task(task) result.worker_id = self.node_id result.result_str = resp["choices"][0]["message"]["content"] result.result = resp["choices"][0]["message"] - + return result def start(self): @@ -82,29 +84,20 @@ class OpenAI_ComputeNode(ComputeNode): if result is not None: task.state = ComputeTaskState.DONE task.result = result - + asyncio.create_task(_run_task_loop()) def display(self) -> str: return f"OpenAI_ComputeNode: {self.node_id}" - - def get_task_state(self,task_id:str): - pass + def get_task_state(self, task_id: str): + pass def get_capacity(self): pass - - def is_support(self,task_type:str) -> bool: - return True - + def is_support(self, task_type: ComputeTaskType) -> bool: + return task_type == ComputeTaskType.LLM_COMPLETION def is_local(self) -> bool: return False - - - - - - \ No newline at end of file diff --git a/src/aios_kernel/stability_node.py b/src/aios_kernel/stability_node.py new file mode 100644 index 0000000..a5ed4e0 --- /dev/null +++ b/src/aios_kernel/stability_node.py @@ -0,0 +1,140 @@ +import os +import io +import asyncio +from asyncio import Queue +import logging + +from PIL import Image +from stability_sdk import client +import stability_sdk.interfaces.gooseai.generation.generation_pb2 as generation + +from .compute_task import ComputeTask, ComputeTaskResult, ComputeTaskState, ComputeTaskType +from .compute_node import ComputeNode + +logger = logging.getLogger(__name__) + + +class Stability_ComputeNode(ComputeNode): + _instanace = None + + def __new__(cls): + if cls._instanace is None: + cls._instanace = super(Stability_ComputeNode, cls).__new__(cls) + cls._instanace.is_start = False + return cls._instanace + + def __init__(self) -> None: + super().__init__() + if self.is_start is True: + logger.warn("Stability_ComputeNode is already start") + return + + self.is_start = True + self.node_id = "stability_node" + self.api_key = "" # "sk-RQDlJtBFQg6I3IueeGCGZTPhWPYAl3bgRdvFDMkcEXsKbUc0" + self.engine = "" # stable-diffusion-512-v2-0 + + self.task_queue = Queue() + + if os.getenv("STABILITY_API_KEY") is not None: + self.api_key = os.getenv("STABILITY_API_KEY") + else: + self.api_key = "sk-RQDlJtBFQg6I3IueeGCGZTPhWPYAl3bgRdvFDMkcEXsKbUc0" + + # Check out the following link for a list of available engines: https://platform.stability.ai/docs/features/api-parameters#engine + if os.getenv("STABILITY_ENGINE") is not None: + self.engine = os.getenv("STABILITY_ENGINE") + else: + self.engine = "stable-diffusion-512-v2-1" + + self.client = client.StabilityInference( + key=self.api_key, + verbose=True, # Print debug messages. + engine=self.engine, + ) + + self.start() + + async def push_task(self, task: ComputeTask, proiority: int = 0): + logger.info(f"stability_node push task: {task.display()}") + self.task_queue.put_nowait(task) + + async def remove_task(self, task_id: str): + pass + + def _run_task(self, task: ComputeTask): + task.state = ComputeTaskState.RUNNING + # model_name && max_token_size not used here + prompts = task.params["prompts"] + + logging.info(f"call stability {self.engine} prompts: {prompts}") + answers = self.client.generate( + prompt=prompts, + # If a seed is provided, the resulting generated image will be deterministic. + seed=0, + # What this means is that as long as all generation parameters remain the same, you can always recall the same image simply by generating it again. + # Note: This isn't quite the case for Clip Guided generations, which we'll tackle in a future example notebook. + # Amount of inference steps performed on image generation. Defaults to 30. + steps=30, + # Influences how strongly your generation is guided to match your prompt. + cfg_scale=7.0, + # Setting this value higher increases the strength in which it tries to match your prompt. + # Defaults to 7.0 if not specified. + width=512, # Generation width, defaults to 512 if not included. + height=512, # Generation height, defaults to 512 if not included. + # Number of images to generate, defaults to 1 if not included. + samples=1, + # Choose which sampler we want to denoise our generation with. + sampler=generation.SAMPLER_K_DPMPP_2M + # Defaults to k_dpmpp_2m if not specified. Clip Guidance only supports ancestral samplers. + # (Available Samplers: ddim, plms, k_euler, k_euler_ancestral, k_heun, k_dpm_2, k_dpm_2_ancestral, k_dpmpp_2s_ancestral, k_lms, k_dpmpp_2m, k_dpmpp_sde) + ) + + for resp in answers: + for artifact in resp.artifacts: + logger.info("artifact:", artifact.id, + artifact.type, artifact.finish_reason) + if artifact.finish_reason == generation.FILTER: + logging.warn("request activated the API's safety filters") + if artifact.type == generation.ARTIFACT_IMAGE: + img = Image.open(io.BytesIO(artifact.binary)) + # Save our generated images with the task_id as the filename. + file_name = task.task_id + ".png" # which dir to save? + img.save(file_name) + + result = ComputeTaskResult() + result.set_from_task(task) + result.worker_id = self.node_id + result.result = {"file": file_name} + + return result + + return None + + def start(self): + async def _run_task_loop(): + while True: + logger.info("stability_node is waiting for task...") + task = await self.task_queue.get() + logger.info(f"stability_node get task: {task.display()}") + result = self._run_task(task) + if result is not None: + task.state = ComputeTaskState.DONE + task.result = result + + asyncio.create_task(_run_task_loop()) + + def display(self) -> str: + return f"Stability_ComputeNode: {self.node_id}" + + def get_task_state(self, task_id: str): + pass + + def get_capacity(self): + pass + + def is_support(self, task_type: ComputeTaskType) -> bool: + return task_type == ComputeTaskType.TEXT_2_IMAGE + + def is_local(self) -> bool: + return False diff --git a/src/requirements.txt b/src/requirements.txt new file mode 100644 index 0000000..605dc72 --- /dev/null +++ b/src/requirements.txt @@ -0,0 +1,8 @@ +aiofiles==23.2.1 +aiohttp==3.8.5 +openai==0.28.0 +Pillow==10.0.0 +Pillow==10.0.0 +prompt_toolkit==3.0.39 +stability_sdk==0.8.4 +toml==0.10.2 From 0d793bd12d4fd5e24b13d9a829a81ad91adac3b4 Mon Sep 17 00:00:00 2001 From: Song Date: Sat, 9 Sep 2023 22:07:31 +0800 Subject: [PATCH 2/3] Remove some test information --- src/aios_kernel/compute_task.py | 4 ++-- src/aios_kernel/stability_node.py | 8 ++------ 2 files changed, 4 insertions(+), 8 deletions(-) diff --git a/src/aios_kernel/compute_task.py b/src/aios_kernel/compute_task.py index 7895283..086d946 100644 --- a/src/aios_kernel/compute_task.py +++ b/src/aios_kernel/compute_task.py @@ -23,7 +23,7 @@ class ComputeTaskType(Enum): class ComputeTask: def __init__(self) -> None: - self.task_type = "llm_completion" + self.task_type = ComputeTaskType.NONE self.create_time = None self.task_id: str = None @@ -37,7 +37,7 @@ class ComputeTask: self.error_str = None def set_llm_params(self, prompts, model_name, max_token_size, callchain_id=None): - self.task_type = "llm_completion" + self.task_type = ComputeTaskType.LLM_COMPLETION self.create_time = time.time() self.task_id = uuid.uuid4().hex self.callchain_id = callchain_id diff --git a/src/aios_kernel/stability_node.py b/src/aios_kernel/stability_node.py index a5ed4e0..8d40878 100644 --- a/src/aios_kernel/stability_node.py +++ b/src/aios_kernel/stability_node.py @@ -31,21 +31,17 @@ class Stability_ComputeNode(ComputeNode): self.is_start = True self.node_id = "stability_node" - self.api_key = "" # "sk-RQDlJtBFQg6I3IueeGCGZTPhWPYAl3bgRdvFDMkcEXsKbUc0" - self.engine = "" # stable-diffusion-512-v2-0 + self.api_key = "" + self.engine = "stable-diffusion-512-v2-1" self.task_queue = Queue() if os.getenv("STABILITY_API_KEY") is not None: self.api_key = os.getenv("STABILITY_API_KEY") - else: - self.api_key = "sk-RQDlJtBFQg6I3IueeGCGZTPhWPYAl3bgRdvFDMkcEXsKbUc0" # Check out the following link for a list of available engines: https://platform.stability.ai/docs/features/api-parameters#engine if os.getenv("STABILITY_ENGINE") is not None: self.engine = os.getenv("STABILITY_ENGINE") - else: - self.engine = "stable-diffusion-512-v2-1" self.client = client.StabilityInference( key=self.api_key, From 27ed3e20e6cf448739a61da1ef97283be2c8b557 Mon Sep 17 00:00:00 2001 From: Song Date: Sat, 9 Sep 2023 22:08:39 +0800 Subject: [PATCH 3/3] The argument for is_support should be task_type --- src/aios_kernel/compute_kernel.py | 62 ++++++++++++++++--------------- 1 file changed, 33 insertions(+), 29 deletions(-) diff --git a/src/aios_kernel/compute_kernel.py b/src/aios_kernel/compute_kernel.py index 1d05afd..2501ad6 100644 --- a/src/aios_kernel/compute_kernel.py +++ b/src/aios_kernel/compute_kernel.py @@ -6,15 +6,18 @@ from asyncio import Queue from .agent import AgentPrompt from .compute_node import ComputeNode -from .compute_task import ComputeTask,ComputeTaskState,ComputeTaskResult +from .compute_task import ComputeTask, ComputeTaskState, ComputeTaskResult logger = logging.getLogger(__name__) # How to dispatch different computing tasks (some tasks may contain a large amount of state for correct execution) -# to suitable computing nodes, achieving a balance of speed, cost, and power consumption, +# to suitable computing nodes, achieving a balance of speed, cost, and power consumption, # is the CORE GOAL of the entire computing task schedule system (aios_kernel). + + class ComputeKernel: _instance = None + def __new__(cls): if cls._instance is None: cls._instance = super().__new__(cls) @@ -25,7 +28,7 @@ class ComputeKernel: def __init__(self) -> None: if self.is_start is True: return - + self.is_start = True self.task_queue = Queue() self.is_start = False @@ -33,70 +36,72 @@ class ComputeKernel: self.start() - def run(self,task:ComputeTask) -> None: + def run(self, task: ComputeTask) -> None: # check there is compute node can support this task if self.is_task_support(task) is False: - logger.error(f"task {task.display()} is not support by any compute node") + logger.error( + f"task {task.display()} is not support by any compute node") return # add task to working_queue self.task_queue.put_nowait(task) - def start(self): if self.is_start is True: logger.warn("compute_kernel is already start") return - + self.is_start = True + async def _run_task_loop(): while True: logger.info("compute_kernel is waiting for task...") task = await self.task_queue.get() logger.info(f"compute_kernel get task: {task.display()}") - c_node:ComputeNode = self._schedule(task) + c_node: ComputeNode = self._schedule(task) await c_node.push_task(task) - - logger.warn("compute_kernel is stoped!") - - asyncio.create_task(_run_task_loop()) - - def _schedule(self,task) -> ComputeNode: + logger.warn("compute_kernel is stoped!") + + asyncio.create_task(_run_task_loop()) + + def _schedule(self, task) -> ComputeNode: for node in self.compute_nodes.values(): - if node.is_support(task) is True: + if node.is_support(task.task_type) is True: return node - logger.warning(f"task {task.display()} is not support by any compute node") + logger.warning( + f"task {task.display()} is not support by any compute node") return None - def add_compute_node(self,node:ComputeNode): + def add_compute_node(self, node: ComputeNode): if self.compute_nodes.get(node.node_id) is not None: - logger.warn(f"compute_node {node.display()} already in compute_kernel") + logger.warn( + f"compute_node {node.display()} already in compute_kernel") return self.compute_nodes[node.node_id] = node logger.info(f"add compute_node {node.display()} to compute_kernel") - def disable_compute_node(self,node_id:str): + def disable_compute_node(self, node_id: str): node = self.compute_nodes.get(node_id) if node is None: logger.warn(f"compute_node {node_id} not in compute_kernel") return node.enable = False - def is_task_support(self,task:ComputeTask) -> bool: + def is_task_support(self, task: ComputeTask) -> bool: return True - # friendly interface for use: - def llm_completion(self,prompt:AgentPrompt,mode_name:Optional[str] = None,max_token:int = 0): + def llm_completion(self, prompt: AgentPrompt, mode_name: Optional[str] = None, max_token: int = 0): # craete a llm_work_task ,push on queue's end # then task_schedule would run this task.(might schedule some work_task to another host) task_req = ComputeTask() - task_req.set_llm_params(prompt,mode_name,max_token) + task_req.set_llm_params(prompt, mode_name, max_token) self.run(task_req) return task_req - async def do_llm_completion(self,prompt:AgentPrompt,mode_name:Optional[str] = None,max_token:int = 0) -> str: - task_req = self.llm_completion(prompt,mode_name,max_token) + async def do_llm_completion(self, prompt: AgentPrompt, mode_name: Optional[str] = None, max_token: int = 0) -> str: + task_req = self.llm_completion(prompt, mode_name, max_token) + async def check_timer(): check_times = 0 while True: @@ -106,16 +111,15 @@ class ComputeKernel: if task_req.state == ComputeTaskState.ERROR: break - if check_times >= 20: + if check_times >= 20: task_req.state = ComputeTaskState.ERROR break await asyncio.sleep(0.5) check_times += 1 - + await asyncio.create_task(check_timer()) if task_req.state == ComputeTaskState.DONE: return task_req.result.result_str - + return "error!" - \ No newline at end of file