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
+19 -14
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)
task.state = resp["state"]
if task.state == ComputeTaskState.ERROR:
task.error_str = resp["error"]["message"]
result.worker_id = self.node_id
result.result_str = resp["content"]
result.result_message = resp["message"]
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.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())