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

The argument for is_support should be task_type

Browse Source
This commit is contained in:
Song
2023-09-09 22:08:39 +08:00
parent 0d793bd12d
commit 27ed3e20e6
1 changed files with 33 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/aios_kernel/compute_kernel.py
+23 -19
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)
@@ -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!") logger.warn("compute_kernel is stoped!")
asyncio.create_task(_run_task_loop()) asyncio.create_task(_run_task_loop())
def _schedule(self, task) -> ComputeNode:
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:
@@ -118,4 +123,3 @@ class ComputeKernel:
return task_req.result.result_str return task_req.result.result_str
return "error!" return "error!"