fix queue compute node
This commit is contained in:
@@ -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!"
|
||||||
|
|
||||||
|
|||||||
@@ -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]:
|
||||||
|
|||||||
@@ -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:
|
||||||
|
|||||||
@@ -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())
|
||||||
|
|
||||||
|
|||||||
@@ -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]
|
||||||
|
|||||||
Reference in New Issue
Block a user