fix queue compute node

This commit is contained in:
tsukasa
2023-09-27 17:24:00 +08:00
parent a5e5839460
commit 01f05285c7
5 changed files with 24 additions and 19 deletions
+1 -1
View File
@@ -148,7 +148,7 @@ class ComputeKernel:
task_result = await self._send_task(task_req) task_result = await self._send_task(task_req)
if task_req.state == ComputeTaskState.DONE: if task_req.state == ComputeTaskState.DONE:
return task_result.result return task_result.result_str
return "error!" return "error!"
+2 -2
View File
@@ -36,7 +36,7 @@ class LocalSentenceTransformer_Text_ComputeNode(Queue_ComputeNode):
except Exception as err: except Exception as err:
logger.error(f"load model {self.model} failed: {err}") logger.error(f"load model {self.model} failed: {err}")
return False return False
self.start()
return True return True
async def execute_task( async def execute_task(
@@ -126,7 +126,7 @@ class LocalSentenceTransformer_Image_ComputeNode(Queue_ComputeNode):
except Exception as err: except Exception as err:
logger.error(f"load model {self.model} failed: {err}") logger.error(f"load model {self.model} failed: {err}")
return False return False
self.start()
return True return True
def _load_image(self, source: Union[ObjectID, bytes]) -> Optional[Image]: def _load_image(self, source: Union[ObjectID, bytes]) -> Optional[Image]:
+1 -1
View File
@@ -98,7 +98,7 @@ class OpenAI_ComputeNode(ComputeNode):
task.state = ComputeTaskState.DONE task.state = ComputeTaskState.DONE
result.result_code = ComputeTaskResultCode.OK result.result_code = ComputeTaskResultCode.OK
result.worker_id = self.node_id result.worker_id = self.node_id
result.result = resp["data"][0]["embedding"] result.result_str = resp["data"][0]["embedding"]
return result return result
case ComputeTaskType.LLM_COMPLETION: case ComputeTaskType.LLM_COMPLETION:
+14 -9
View File
@@ -4,7 +4,7 @@ from asyncio import Queue
import logging import logging
from abc import abstractmethod from abc import abstractmethod
from .compute_task import ComputeTask, ComputeTaskResult, ComputeTaskState, ComputeTaskType from .compute_task import ComputeTask, ComputeTaskResult, ComputeTaskResultCode, ComputeTaskState, ComputeTaskType
from .compute_node import ComputeNode from .compute_node import ComputeNode
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -13,6 +13,7 @@ class Queue_ComputeNode(ComputeNode):
def __init__(self): def __init__(self):
super().__init__() super().__init__()
self.task_queue = Queue() self.task_queue = Queue()
self.is_start = False
@abstractmethod @abstractmethod
async def execute_task(self, task: ComputeTask) -> { async def execute_task(self, task: ComputeTask) -> {
@@ -39,28 +40,32 @@ class Queue_ComputeNode(ComputeNode):
resp = await self.execute_task(task) resp = await self.execute_task(task)
result = ComputeTaskResult() result = ComputeTaskResult()
result.set_from_task(task)
result.worker_id = self.node_id
task.state = resp["state"] task.state = resp["state"]
if task.state == ComputeTaskState.ERROR: if task.state == ComputeTaskState.ERROR:
result.result_code = ComputeTaskResultCode.ERROR
task.error_str = resp["error"]["message"] task.error_str = resp["error"]["message"]
else:
result.result_code = ComputeTaskResultCode.OK
result.worker_id = self.node_id
result.result_str = resp["content"] result.result_str = resp["content"]
result.result_message = resp["message"] result.result_message = resp["message"]
result.set_from_task(task)
return result return result
async def start(self): async def start(self):
if self.is_start is True:
return
self.is_start = True
async def _run_task_loop(): async def _run_task_loop():
while True: while True:
task = await self.task_queue.get() task = await self.task_queue.get()
logger.info(f"{self.display()} get task: {task.display()}") logger.info(f"openai_node get task: {task.display()}")
result = await self._run_task(task) self._run_task(task)
if result is not None:
task.result = result
asyncio.create_task(_run_task_loop()) asyncio.create_task(_run_task_loop())
+1 -1
View File
@@ -62,5 +62,5 @@ class KnowledgeStore:
def get_vector_store(self, model_name: str) -> VectorBase: def get_vector_store(self, model_name: str) -> VectorBase:
if model_name not in self.vector_store: if model_name not in self.vector_store:
self.vector_store[model_name] = ChromaVectorStore(model_name) self.vector_store[model_name] = ChromaVectorStore(self.root, model_name)
return self.vector_store[model_name] return self.vector_store[model_name]