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)
if task_req.state == ComputeTaskState.DONE:
return task_result.result
return task_result.result_str
return "error!"
+2 -2
View File
@@ -36,7 +36,7 @@ class LocalSentenceTransformer_Text_ComputeNode(Queue_ComputeNode):
except Exception as err:
logger.error(f"load model {self.model} failed: {err}")
return False
self.start()
return True
async def execute_task(
@@ -126,7 +126,7 @@ class LocalSentenceTransformer_Image_ComputeNode(Queue_ComputeNode):
except Exception as err:
logger.error(f"load model {self.model} failed: {err}")
return False
self.start()
return True
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
result.result_code = ComputeTaskResultCode.OK
result.worker_id = self.node_id
result.result = resp["data"][0]["embedding"]
result.result_str = resp["data"][0]["embedding"]
return result
case ComputeTaskType.LLM_COMPLETION:
+15 -10
View File
@@ -4,7 +4,7 @@ from asyncio import Queue
import logging
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
logger = logging.getLogger(__name__)
@@ -13,6 +13,7 @@ class Queue_ComputeNode(ComputeNode):
def __init__(self):
super().__init__()
self.task_queue = Queue()
self.is_start = False
@abstractmethod
async def execute_task(self, task: ComputeTask) -> {
@@ -39,28 +40,32 @@ class Queue_ComputeNode(ComputeNode):
resp = await self.execute_task(task)
result = ComputeTaskResult()
result.set_from_task(task)
result.worker_id = self.node_id
task.state = resp["state"]
if task.state == ComputeTaskState.ERROR:
result.result_code = ComputeTaskResultCode.ERROR
task.error_str = resp["error"]["message"]
else:
result.result_code = ComputeTaskResultCode.OK
result.result_str = resp["content"]
result.result_message = resp["message"]
result.worker_id = self.node_id
result.result_str = resp["content"]
result.result_message = resp["message"]
result.set_from_task(task)
return result
async def start(self):
if self.is_start is True:
return
self.is_start = True
async def _run_task_loop():
while True:
task = await self.task_queue.get()
logger.info(f"{self.display()} get task: {task.display()}")
result = await self._run_task(task)
if result is not None:
task.result = result
logger.info(f"openai_node get task: {task.display()}")
self._run_task(task)
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:
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]