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

Merge pull request #45 from glen0125/MVP

Browse Source 
Add stability node (Text2Img)
This commit is contained in:
Liu Zhicong
2023-09-09 10:59:20 -07:00
committed by GitHub
parent 1bbdf2e722 27ed3e20e6
commit 15dd600fdb
6 changed files with 240 additions and 92 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/aios_kernel/compute_kernel.py
+33 -29
View File
@@ -6,15 +6,18 @@ from asyncio import Queue
from .agent import AgentPrompt from .agent import AgentPrompt
from .compute_node import ComputeNode from .compute_node import ComputeNode
from .compute_task import ComputeTask,ComputeTaskState,ComputeTaskResult from .compute_task import ComputeTask, ComputeTaskState, ComputeTaskResult
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
# How to dispatch different computing tasks (some tasks may contain a large amount of state for correct execution) # 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). # is the CORE GOAL of the entire computing task schedule system (aios_kernel).
class ComputeKernel: class ComputeKernel:
_instance = None _instance = None
def __new__(cls): def __new__(cls):
if cls._instance is None: if cls._instance is None:
cls._instance = super().__new__(cls) cls._instance = super().__new__(cls)
@@ -25,7 +28,7 @@ class ComputeKernel:
def __init__(self) -> None: def __init__(self) -> None:
if self.is_start is True: if self.is_start is True:
return return
self.is_start = True self.is_start = True
self.task_queue = Queue() self.task_queue = Queue()
self.is_start = False self.is_start = False
@@ -33,70 +36,72 @@ class ComputeKernel:
self.start() self.start()
def run(self,task:ComputeTask) -> None: def run(self, task: ComputeTask) -> None:
# check there is compute node can support this task # check there is compute node can support this task
if self.is_task_support(task) is False: 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 return
# add task to working_queue # add task to working_queue
self.task_queue.put_nowait(task) self.task_queue.put_nowait(task)
def start(self): def start(self):
if self.is_start is True: if self.is_start is True:
logger.warn("compute_kernel is already start") logger.warn("compute_kernel is already start")
return return
self.is_start = True self.is_start = True
async def _run_task_loop(): async def _run_task_loop():
while True: while True:
logger.info("compute_kernel is waiting for task...") logger.info("compute_kernel is waiting for task...")
task = await self.task_queue.get() task = await self.task_queue.get()
logger.info(f"compute_kernel get task: {task.display()}") 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) 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(): 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 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 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: 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 return
self.compute_nodes[node.node_id] = node self.compute_nodes[node.node_id] = node
logger.info(f"add compute_node {node.display()} to compute_kernel") 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) node = self.compute_nodes.get(node_id)
if node is None: if node is None:
logger.warn(f"compute_node {node_id} not in compute_kernel") logger.warn(f"compute_node {node_id} not in compute_kernel")
return return
node.enable = False node.enable = False
def is_task_support(self,task:ComputeTask) -> bool: def is_task_support(self, task: ComputeTask) -> bool:
return True return True
# friendly interface for use: # 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 # 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) # then task_schedule would run this task.(might schedule some work_task to another host)
task_req = ComputeTask() 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) self.run(task_req)
return task_req return task_req
async def do_llm_completion(self,prompt:AgentPrompt,mode_name:Optional[str] = None,max_token:int = 0) -> str: 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) task_req = self.llm_completion(prompt, mode_name, max_token)
async def check_timer(): async def check_timer():
check_times = 0 check_times = 0
while True: while True:
@@ -106,16 +111,15 @@ class ComputeKernel:
if task_req.state == ComputeTaskState.ERROR: if task_req.state == ComputeTaskState.ERROR:
break break
if check_times >= 20: if check_times >= 20:
task_req.state = ComputeTaskState.ERROR task_req.state = ComputeTaskState.ERROR
break break
await asyncio.sleep(0.5) await asyncio.sleep(0.5)
check_times += 1 check_times += 1
await asyncio.create_task(check_timer()) await asyncio.create_task(check_timer())
if task_req.state == ComputeTaskState.DONE: if task_req.state == ComputeTaskState.DONE:
return task_req.result.result_str return task_req.result.result_str
return "error!" return "error!"
src/aios_kernel/compute_node.py
+11 -14
View File
@@ -1,5 +1,5 @@
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
from .compute_task import ComputeTask from .compute_task import ComputeTask, ComputeTaskType
class ComputeNode(ABC): class ComputeNode(ABC):
@@ -8,15 +8,15 @@ class ComputeNode(ABC):
self.enable = True self.enable = True
@abstractmethod @abstractmethod
async def push_task(self,task:ComputeTask,proiority:int = 0): async def push_task(self, task: ComputeTask, proiority: int = 0):
pass
@abstractmethod
async def remove_task(self,task_id:str):
pass pass
@abstractmethod @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 pass
@abstractmethod @abstractmethod
@@ -28,7 +28,7 @@ class ComputeNode(ABC):
pass pass
@abstractmethod @abstractmethod
def is_support(self,task_type:str) -> bool: def is_support(self, task_type: ComputeTaskType) -> bool:
pass pass
@abstractmethod @abstractmethod
@@ -37,17 +37,14 @@ class ComputeNode(ABC):
def is_trusted(self) -> bool: def is_trusted(self) -> bool:
return True return True
def get_fee_type(self) -> str: def get_fee_type(self) -> str:
return "free" return "free"
class LocalComputeNode(ComputeNode): class LocalComputeNode(ComputeNode):
def display(self) -> str: def display(self) -> str:
return super().display() return super().display()
def is_local(self) -> bool: def is_local(self) -> bool:
return True return True
src/aios_kernel/compute_task.py
+28 -18
View File
@@ -3,6 +3,7 @@ from enum import Enum
import uuid import uuid
import time import time
class ComputeTaskState(Enum): class ComputeTaskState(Enum):
DONE = 0 DONE = 0
INIT = 1 INIT = 1
@@ -11,30 +12,39 @@ class ComputeTaskState(Enum):
PENDING = 4 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: class ComputeTask:
def __init__(self) -> None: def __init__(self) -> None:
self.task_type = "llm_completion" self.task_type = ComputeTaskType.NONE
self.create_time = None self.create_time = None
self.task_id:str = None self.task_id: str = None
self.callchain_id:str = None self.callchain_id: str = None
self.params:dict = {} self.params: dict = {}
self.refers:dict = None self.refers: dict = None
self.pading_data:bytearray = None self.pading_data: bytearray = None
self.state = ComputeTaskState.INIT self.state = ComputeTaskState.INIT
self.result = None self.result = None
self.error_str = 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.task_type = ComputeTaskType.LLM_COMPLETION
self.create_time = time.time() self.create_time = time.time()
self.task_id = uuid.uuid4().hex self.task_id = uuid.uuid4().hex
self.callchain_id = callchain_id self.callchain_id = callchain_id
self.params["prompts"] = prompts.messages self.params["prompts"] = prompts.messages
if model_name is not None: if model_name is not None:
self.params["model_name"] = model_name self.params["model_name"] = model_name
else: else:
self.params["model_name"] = "gpt-4-0613" self.params["model_name"] = "gpt-4-0613"
self.params["max_token_size"] = max_token_size self.params["max_token_size"] = max_token_size
@@ -45,16 +55,16 @@ class ComputeTask:
class ComputeTaskResult: class ComputeTaskResult:
def __init__(self) -> None: def __init__(self) -> None:
self.create_time = None self.create_time = None
self.task_id:str = None self.task_id: str = None
self.callchain_id:str = None self.callchain_id: str = None
self.worker_id:str = None self.worker_id: str = None
self.result_code:int = 0 self.result_code: int = 0
self.result_str:str = None self.result_str: str = None
self.result:dict = {} self.result: dict = {}
self.result_refers:dict = None self.result_refers: dict = None
self.pading_data:bytearray = 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.task_id = task.task_id
self.callchain_id = task.callchain_id self.callchain_id = task.callchain_id
src/aios_kernel/open_ai_node.py
+24 -31
View File
@@ -5,47 +5,49 @@ import asyncio
from asyncio import Queue from asyncio import Queue
import logging import logging
from .compute_task import ComputeTask,ComputeTaskResult,ComputeTaskState from .compute_task import ComputeTask, ComputeTaskResult, ComputeTaskState, ComputeTaskType
from .compute_node import ComputeNode from .compute_node import ComputeNode
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
class OpenAI_ComputeNode(ComputeNode): class OpenAI_ComputeNode(ComputeNode):
_instance = None _instance = None
def __new__(cls): def __new__(cls):
if cls._instance is None: if cls._instance is None:
cls._instance = super(OpenAI_ComputeNode, cls).__new__(cls) cls._instance = super(OpenAI_ComputeNode, cls).__new__(cls)
cls._instance.is_start = False cls._instance.is_start = False
return cls._instance return cls._instance
def __init__(self) -> None: def __init__(self) -> None:
super().__init__() super().__init__()
if self.is_start is True: if self.is_start is True:
logger.warn("OpenAI_ComputeNode is already start") logger.warn("OpenAI_ComputeNode is already start")
return return
self.is_start = True self.is_start = True
#openai.organization = "org-AoKrOtF2myemvfiFfnsSU8rF" #buckycloud # openai.organization = "org-AoKrOtF2myemvfiFfnsSU8rF" #buckycloud
self.openai_api_key = "" self.openai_api_key = ""
self.node_id = "openai_node" self.node_id = "openai_node"
self.task_queue = Queue() 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") openai.api_key = os.getenv("OPENAI_API_KEY")
else: else:
openai.api_key = self.openai_api_key openai.api_key = self.openai_api_key
self.start() 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()}") logger.info(f"openai_node push task: {task.display()}")
self.task_queue.put_nowait(task) self.task_queue.put_nowait(task)
async def remove_task(self,task_id:str): async def remove_task(self, task_id: str):
pass pass
def _run_task(self,task:ComputeTask): def _run_task(self, task: ComputeTask):
task.state = ComputeTaskState.RUNNING task.state = ComputeTaskState.RUNNING
mode_name = task.params["model_name"] mode_name = task.params["model_name"]
# max_token_size = task.params["max_token_size"] # max_token_size = task.params["max_token_size"]
@@ -57,19 +59,19 @@ class OpenAI_ComputeNode(ComputeNode):
max_tokens=4000, max_tokens=4000,
temperature=1.2) temperature=1.2)
logger.info(f"openai response: {resp}") logger.info(f"openai response: {resp}")
status_code = resp["choices"][0]["finish_reason"] status_code = resp["choices"][0]["finish_reason"]
if status_code != "stop": if status_code != "stop":
task.state = ComputeTaskState.ERROR 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 return None
result = ComputeTaskResult() result = ComputeTaskResult()
result.set_from_task(task) result.set_from_task(task)
result.worker_id = self.node_id result.worker_id = self.node_id
result.result_str = resp["choices"][0]["message"]["content"] result.result_str = resp["choices"][0]["message"]["content"]
result.result = resp["choices"][0]["message"] result.result = resp["choices"][0]["message"]
return result return result
def start(self): def start(self):
@@ -82,29 +84,20 @@ class OpenAI_ComputeNode(ComputeNode):
if result is not None: if result is not None:
task.state = ComputeTaskState.DONE task.state = ComputeTaskState.DONE
task.result = result task.result = result
asyncio.create_task(_run_task_loop()) asyncio.create_task(_run_task_loop())
def display(self) -> str: def display(self) -> str:
return f"OpenAI_ComputeNode: {self.node_id}" 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): def get_capacity(self):
pass pass
def is_support(self, task_type: ComputeTaskType) -> bool:
def is_support(self,task_type:str) -> bool: return task_type == ComputeTaskType.LLM_COMPLETION
return True
def is_local(self) -> bool: def is_local(self) -> bool:
return False return False
src/aios_kernel/stability_node.py
+136
View File
@@ -0,0 +1,136 @@
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 = ""
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")
# 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")
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
src/requirements.txt
+8
View File
@@ -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