from abc import ABC, abstractmethod import random from typing import Optional import logging import asyncio import tiktoken from asyncio import Queue from knowledge import ObjectID from .agent_base import AgentPrompt from .compute_node import ComputeNode from .compute_task import ComputeTask, ComputeTaskState, ComputeTaskResult, ComputeTaskType,ComputeTaskResultCode 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, # is the CORE GOAL of the entire computing task schedule system (aios_kernel). class ComputeKernel: _instance = None @classmethod def get_instance(cls): if cls._instance is None: cls._instance = ComputeKernel() return cls._instance def __init__(self) -> None: self.is_start = False self.task_queue = Queue() self.is_start = False self.compute_nodes = {} 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") return # add task to working_queue self.task_queue.put_nowait(task) async 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: task = await self.task_queue.get() logger.info(f"compute_kernel get task: {task.display()}") c_node: ComputeNode = self._schedule(task) if c_node: await c_node.push_task(task) logger.warn("compute_kernel is stoped!") asyncio.create_task(_run_task_loop()) def _schedule(self, task) -> ComputeNode: # find all the node which supports this task support_nodes = [] total_weights = 0 for node in self.compute_nodes.values(): if node.is_support(task) is True: support_nodes.append({ "pos": total_weights, "node": node }) total_weights += node.weight() if len(support_nodes) < 1: logger.warning(f"task {task.display()} is not support by any compute node") return None # hit a random node with weight hit_pos = random.randint(0, total_weights - 1) for i in range(min(len(support_nodes) - 1, hit_pos), -1, -1): if support_nodes[i]["pos"] <= hit_pos: return support_nodes[i]["node"] logger.warning( f"task {task.display()} is not support by any compute node") return None 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") 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): 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: return True @staticmethod def llm_num_tokens_from_text(text:str,model:str) -> int: try: encoding = tiktoken.encoding_for_model(model) except KeyError: logger.debug("Warning: model not found. Using cl100k_base encoding.") encoding = tiktoken.get_encoding("cl100k_base") token_count = len(encoding.encode(text)) return token_count # friendly interface for use: def llm_completion(self, prompt: AgentPrompt, mode_name: Optional[str] = None, max_token: int = 0,inner_functions = None): # 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,inner_functions) self.run(task_req) return task_req async def _wait_task(self,task_req:ComputeTask)->ComputeTaskResult: async def check_timer(): check_times = 0 while True: if task_req.state == ComputeTaskState.DONE: break if task_req.state == ComputeTaskState.ERROR: break if check_times >= 120: task_req.state = ComputeTaskState.ERROR break await asyncio.sleep(0.5) check_times += 1 await asyncio.create_task(check_timer()) if task_req.result: return task_req.result else: time_out_result = ComputeTaskResult() time_out_result.result_code = ComputeTaskResultCode.TIMEOUT time_out_result.set_from_task(task_req) ## craete timeout task_result async def do_llm_completion(self, prompt: AgentPrompt, mode_name: Optional[str] = None, max_token: int = 0, inner_functions = None) -> str: task_req = self.llm_completion(prompt, mode_name, max_token,inner_functions) return await self._wait_task(task_req) def text_embedding(self,input:str,model_name:Optional[str] = None): task_req = ComputeTask() task_req.set_text_embedding_params(input,model_name) self.run(task_req) return task_req async def do_text_embedding(self,input:str,model_name:Optional[str] = None) -> [float]: task_req = self.text_embedding(input,model_name) task_result = await self._wait_task(task_req) if task_req.state == ComputeTaskState.DONE: return task_result.result.get("content") else: logging.warning(f"do_text_embedding error: {task_req.error_str},input: {input}") return None def image_embedding(self,input:ObjectID,model_name:Optional[str] = None): task_req = ComputeTask() task_req.set_image_embedding_params(input,model_name) self.run(task_req) return task_req async def do_image_embedding(self,input:ObjectID,model_name:Optional[str] = None) -> [float]: task_req = self.image_embedding(input,model_name) task_result = await self._wait_task(task_req) if task_req.state == ComputeTaskState.DONE: return task_result.result.get("content") return None async def do_text_to_speech(self, input:str, language_code:Optional[str] = None, gender: Optional[str] = None, age: Optional[str] = None, voice_name: Optional[str] = None, tone: Optional[str] = None): task_req = ComputeTask() task_req.params["text"] = input task_req.params["language_code"] = language_code task_req.params["gender"] = gender task_req.params["age"] = age task_req.params["voice_name"] = voice_name task_req.params["tone"] = tone task_req.task_type = ComputeTaskType.TEXT_2_VOICE self.run(task_req) task_result = await self._wait_task(task_req) if task_req.state == ComputeTaskState.DONE: return task_result.result def text_2_image(self, prompt:str, model_name:Optional[str] = None, negative_prompt = None): task = ComputeTask() task.set_text_2_image_params(prompt,model_name, negative_prompt) self.run(task) return task async def do_text_2_image(self, prompt:str, model_name:Optional[str] = None, negative_prompt = None) -> ComputeTaskResult: task = self.text_2_image(prompt,model_name, negative_prompt) task = await self._wait_task(task) return task.result # if task_req.state == ComputeTaskState.DONE: # return None, task_result # return task_req.error_str, None