diff --git a/workflow_service/comfy/comfy_queue.py b/workflow_service/comfy/comfy_queue.py index 34fcdbc..099be22 100644 --- a/workflow_service/comfy/comfy_queue.py +++ b/workflow_service/comfy/comfy_queue.py @@ -1,27 +1,13 @@ import asyncio -import json import logging -import os -import random -import uuid -import websockets -from datetime import datetime -from typing import Dict, Any, Optional +from typing import Any, Optional import aiohttp -from aiohttp import ClientTimeout from workflow_service.comfy.comfy_workflow import ComfyWorkflow +from workflow_service.comfy.comfy_run import ComfyRun from workflow_service.config import Settings from workflow_service.comfy.comfy_server import server_manager, ComfyUIServerInfo -from workflow_service.database.api import ( - create_workflow_run, - update_workflow_run_status, - create_workflow_run_nodes, - update_workflow_run_node_status, - get_workflow_run, - get_workflow_run_nodes, -) settings = Settings() @@ -29,22 +15,6 @@ logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) -API_INPUT_PREFIX = "INPUT_" -API_OUTPUT_PREFIX = "OUTPUT_" - - -class ComfyUIExecutionError(Exception): - def __init__(self, error_data: dict): - self.error_data = error_data - # 创建一个对开发者友好的异常消息 - message = ( - f"ComfyUI节点执行失败。节点ID: {error_data.get('node_id')}, " - f"节点类型: {error_data.get('node_type')}. " - f"错误: {error_data.get('exception_message', 'N/A')}" - ) - super().__init__(message) - - # 全局任务队列管理器 class WorkflowQueueManager: def __init__(self, monitor_interval: int = 5): @@ -71,33 +41,60 @@ class WorkflowQueueManager: pass logger.info("队列监控任务已停止") - async def set_monitor_interval(self, interval: int): - """设置监控间隔""" - if interval < 1: - raise ValueError("监控间隔不能小于1秒") - self._monitor_interval = interval - logger.info(f"队列监控间隔已设置为 {interval} 秒") - - # 如果监控任务正在运行，重启它以应用新间隔 - if self._queue_monitor_task and not self._queue_monitor_task.done(): - await self.stop_queue_monitor() - await self.start_queue_monitor() - - async def trigger_queue_processing(self): - """手动触发队列处理（用于测试或紧急情况）""" - logger.info("手动触发队列处理") - asyncio.create_task(self._process_queue()) - - async def get_queue_status(self) -> dict: - """获取队列状态信息""" + async def add_task( + self, + workflow: ComfyWorkflow, + request_data: dict, + ): + """添加新任务到队列""" async with self.lock: - return { - "pending_tasks_count": len(self.pending_tasks), - "running_tasks_count": len(self.running_tasks), - "monitor_active": self._queue_monitor_task is not None - and not self._queue_monitor_task.done(), - "monitor_interval": self._monitor_interval, - } + # 创建ComfyRun实例并保存到数据库 + comfy_run = await ComfyRun.create(workflow, request_data) + + # 添加到待处理队列 + self.pending_tasks.append(comfy_run.run_id) + logger.info( + f"任务 {comfy_run.run_id} 已添加到队列，当前队列长度: {len(self.pending_tasks)}" + ) + + return comfy_run.run_id + + async def get_server_status( + self, server: ComfyUIServerInfo, session: aiohttp.ClientSession + ) -> dict[str, Any]: + """ + 检查单个ComfyUI服务器的详细状态。 + 返回一个包含可达性、队列状态和详细队列内容的字典。 + """ + # [BUG修复] 确保初始字典结构与成功时的结构一致，以满足Pydantic模型 + status_info = { + "is_reachable": False, + "is_free": False, + "queue_details": {"running_count": 0, "pending_count": 0}, + } + try: + queue_url = f"{server.http_url}/queue" + async with session.get(queue_url, timeout=60) as response: + response.raise_for_status() + queue_data = await response.json() + + status_info["is_reachable"] = True + + running_count = len(queue_data.get("queue_running", [])) + pending_count = len(queue_data.get("queue_pending", [])) + + status_info["queue_details"] = { + "running_count": running_count, + "pending_count": pending_count, + } + + status_info["is_free"] = running_count == 0 and pending_count == 0 + + except Exception as e: + # 当请求失败时，将返回上面定义的、结构正确的初始 status_info + logger.warning(f"无法检查服务器 {server.http_url} 的队列状态: {e}") + + return status_info async def _monitor_queue(self): """定期监控队列，当有可用服务器时自动处理""" @@ -107,11 +104,12 @@ class WorkflowQueueManager: # 每次定时检测触发时，打印当前状态 async with self.lock: + from datetime import datetime + pending_count = len(self.pending_tasks) running_count = len(self.running_tasks) current_time = datetime.now() timestamp = current_time.strftime("%Y-%m-%d %H:%M:%S") - time_only = current_time.strftime("%H:%M:%S") # 计算距离上次处理队列的时间 time_since_last_processing = "" @@ -161,44 +159,12 @@ class WorkflowQueueManager: # 出错时稍微延长等待时间，避免频繁重试 await asyncio.sleep(self._monitor_interval * 2) - async def add_task( - self, - workflow: ComfyWorkflow, - request_data: dict, - ): - """添加新任务到队列""" - workflow_run_id = str(uuid.uuid4()) - async with self.lock: - # 创建任务记录 - await create_workflow_run( - workflow_run_id=workflow_run_id, - workflow_name=workflow.workflow_name, - workflow_json=json.dumps(workflow.workflow_data.model_dump()), - api_spec=json.dumps(workflow.get_api_spec().model_dump()), - request_data=json.dumps(request_data), - ) - - # 创建工作流节点记录 - nodes_data = [] - for node in workflow.workflow_data.nodes: - nodes_data.append({"id": str(node.id), "type": node.type}) - await create_workflow_run_nodes(workflow_run_id, nodes_data) - - # 添加到待处理队列 - self.pending_tasks.append(workflow_run_id) - logger.info( - f"任务 {workflow_run_id} 已添加到队列，当前队列长度: {len(self.pending_tasks)}" - ) - - # 注意：不再手动触发队列处理，由定时监控自动处理 - # 这样可以避免在没有可用服务器时的不必要尝试 - - return workflow_run_id - async def _process_queue(self): """处理队列中的任务""" async with self.lock: # 更新上次处理队列的时间 + from datetime import datetime + self._last_processing_time = datetime.now() if not self.pending_tasks: @@ -221,13 +187,7 @@ class WorkflowQueueManager: f"开始处理任务 {workflow_run_id}，使用服务器 {server.name} ({server.http_url})" ) - # 分配服务器资源 - await server_manager.allocate_server(server.name) - - # 标记任务为运行中 - await update_workflow_run_status( - workflow_run_id, "running", server.http_url - ) + # 记录运行中任务 self.running_tasks[server.http_url] = { "workflow_run_id": workflow_run_id, "started_at": datetime.now(), @@ -251,307 +211,25 @@ class WorkflowQueueManager: async def _execute_task(self, workflow_run_id: str, server: ComfyUIServerInfo): """执行任务""" - cleanup_paths = [] try: - # 获取工作流数据 - workflow_run = await get_workflow_run(workflow_run_id) - if not workflow_run: + # 从run_id创建ComfyRun实例并执行 + comfy_run = await ComfyRun.from_run_id(workflow_run_id) + if not comfy_run: raise Exception(f"找不到工作流运行记录: {workflow_run_id}") - workflow_data = json.loads(workflow_run.workflow_json) - request_data = json.loads(workflow_run.request_data) - workflow = ComfyWorkflow(workflow_run.workflow_name, workflow_data) - # 执行工作流 - result = await _execute_prompt_on_server( - workflow, request_data, server, workflow_run_id - ) - - # 保存处理后的结果到数据库 - await update_workflow_run_status( - workflow_run_id, - "completed", - result=json.dumps(result, ensure_ascii=False), - ) + # 执行工作流（ComfyRun内部处理所有逻辑包括资源管理） + await comfy_run.execute(server) except Exception as e: logger.error(f"执行任务 {workflow_run_id} 时出错: {e}") - await update_workflow_run_status( - workflow_run_id, "failed", error_message=str(e) - ) - finally: - # 清理临时文件 - if cleanup_paths: - logger.info(f"正在清理 {len(cleanup_paths)} 个临时文件...") - for path in cleanup_paths: - try: - if os.path.exists(path): - os.remove(path) - logger.info(f" - 已删除: {path}") - except Exception as e: - logger.warning(f" - 删除 {path} 时出错: {e}") + finally: # 清理运行状态 async with self.lock: if server.http_url in self.running_tasks: del self.running_tasks[server.http_url] - # 释放服务器资源 - await server_manager.release_server(server.name) - - async def get_task_status(self, workflow_run_id: str) -> dict: - """获取任务状态""" - workflow_run = await get_workflow_run(workflow_run_id) - if not workflow_run: - return {"error": "任务不存在"} - nodes = await get_workflow_run_nodes(workflow_run_id) - - result = { - "id": workflow_run_id, - "status": workflow_run.status, - "created_at": workflow_run.created_at, - "started_at": workflow_run.started_at, - "completed_at": workflow_run.completed_at, - "server_url": workflow_run.server_url, - "error_message": workflow_run.error_message, - "nodes": nodes, - } - - # 如果任务完成，从数据库获取结果 - if workflow_run.status == "completed" and workflow_run.result: - try: - result["result"] = json.loads(workflow_run.result) - except (json.JSONDecodeError, TypeError): - result["result"] = None - - return result - - async def get_server_status( - self, server: ComfyUIServerInfo, session: aiohttp.ClientSession - ) -> dict[str, Any]: - """ - 检查单个ComfyUI服务器的详细状态。 - 返回一个包含可达性、队列状态和详细队列内容的字典。 - """ - # [BUG修复] 确保初始字典结构与成功时的结构一致，以满足Pydantic模型 - status_info = { - "is_reachable": False, - "is_free": False, - "queue_details": {"running_count": 0, "pending_count": 0}, - } - try: - queue_url = f"{server.http_url}/queue" - async with session.get(queue_url, timeout=60) as response: - response.raise_for_status() - queue_data = await response.json() - - status_info["is_reachable"] = True - - running_count = len(queue_data.get("queue_running", [])) - pending_count = len(queue_data.get("queue_pending", [])) - - status_info["queue_details"] = { - "running_count": running_count, - "pending_count": pending_count, - } - - status_info["is_free"] = running_count == 0 and pending_count == 0 - - except Exception as e: - # 当请求失败时，将返回上面定义的、结构正确的初始 status_info - logger.warning(f"无法检查服务器 {server.http_url} 的队列状态: {e}") - - return status_info - # 全局队列管理器实例 # 从配置中读取监控间隔，默认为5秒 queue_manager = WorkflowQueueManager(monitor_interval=5) - - -async def _execute_prompt_on_server( - workflow: ComfyWorkflow, - request_data: dict, - server: ComfyUIServerInfo, - workflow_run_id: str, -) -> dict: - """ - 在指定的服务器上执行一个准备好的prompt。 - 现在支持节点级别的状态跟踪。 - """ - client_id = str(uuid.uuid4()) - - # 构建prompt - prompt = await workflow.build_prompt(server, request_data) - - # 更新工作流运行状态，记录prompt_id和client_id - await update_workflow_run_status( - workflow_run_id, - "running", - server.http_url, - None, # prompt_id将在_queue_prompt后更新 - client_id, - ) - - # 提交到ComfyUI - prompt_id = await _queue_prompt(workflow, prompt, client_id, server.http_url) - - # 更新prompt_id - await update_workflow_run_status( - workflow_run_id, "running", server.http_url, prompt_id, client_id - ) - - logger.info( - f"工作流 {workflow_run_id} 已在 {server.http_url} 上入队，Prompt ID: {prompt_id}" - ) - - # 获取执行结果，现在支持节点级别的状态跟踪 - results = await _get_execution_results( - workflow, prompt_id, client_id, server.ws_url, workflow_run_id - ) - return results - - -async def _queue_prompt( - workflow: ComfyWorkflow, prompt: dict, client_id: str, http_url: str -) -> str: - """通过HTTP POST将工作流任务提交到指定的ComfyUI服务器。""" - for node_id in prompt: - prompt[node_id]["inputs"][f"cache_buster_{uuid.uuid4().hex}"] = random.random() - payload = { - "prompt": prompt, - "client_id": client_id, - "extra_data": { - "api_key_comfy_org": "", - "extra_pnginfo": {"workflow": workflow.workflow_data.model_dump()}, - }, - } - async with aiohttp.ClientSession(timeout=ClientTimeout(total=90)) as session: - prompt_url = f"{http_url}/prompt" - try: - async with session.post(prompt_url, json=payload) as response: - logger.info(f"ComfyUI /prompt 端点返回的响应: {response}") - response.raise_for_status() - result = await response.json() - logger.info(f"ComfyUI /prompt 端点返回的响应: {result}") - if "prompt_id" not in result: - raise Exception(f"从 ComfyUI /prompt 端点返回的响应无效: {result}") - return result["prompt_id"] - except Exception as e: - logger.error(f"提交到 ComfyUI /prompt 端点时发生错误: {e}") - raise e - - -async def _get_execution_results( - workflow: ComfyWorkflow, - prompt_id: str, - client_id: str, - ws_url: str, - workflow_run_id: str, -) -> dict: - """ - 通过WebSocket连接到指定的ComfyUI服务器，聚合执行结果。 - 现在支持节点级别的状态跟踪。 - """ - full_ws_url = f"{ws_url}?clientId={client_id}" - aggregated_outputs = {} - - try: - async with websockets.connect(full_ws_url) as websocket: - while True: - try: - out = await websocket.recv() - if not isinstance(out, str): - continue - - message = json.loads(out) - msg_type = message.get("type") - data = message.get("data") - - if not (data and data.get("prompt_id") == prompt_id): - continue - - # 捕获并处理执行错误 - if msg_type == "execution_error": - error_data = data - logger.error( - f"ComfyUI执行错误 (Prompt ID: {prompt_id}): {error_data}" - ) - - # 更新节点状态为失败 - node_id = error_data.get("node_id") - if node_id: - await update_workflow_run_node_status( - workflow_run_id, - node_id, - "failed", - error_message=error_data.get( - "exception_message", "Unknown error" - ), - ) - - # 抛出自定义异常，将错误详情传递出去 - raise ComfyUIExecutionError(error_data) - - # 处理节点开始执行 - if msg_type == "executing" and data.get("node"): - node_id = data.get("node") - logger.info(f"节点 {node_id} 开始执行 (Prompt ID: {prompt_id})") - - # 更新节点状态为运行中 - await update_workflow_run_node_status( - workflow_run_id, node_id, "running" - ) - - # 处理节点执行完成 - if msg_type == "executed": - node_id = data.get("node") - output_data = data.get("output") - if node_id and output_data: - node = next( - ( - x - for x in workflow.workflow_data.nodes - if str(x.id) == node_id - ), - None, - ) - if ( - node - and node.title - and node.title.startswith(API_OUTPUT_PREFIX) - ): - title = node.title.replace(API_OUTPUT_PREFIX, "") - aggregated_outputs[title] = output_data - logger.info( - f"收到节点 {node_id} 的输出 (Prompt ID: {prompt_id})" - ) - - # 更新节点状态为完成 - await update_workflow_run_node_status( - workflow_run_id, - node_id, - "completed", - output_data=json.dumps(output_data), - ) - - # 处理整个工作流执行完成 - elif msg_type == "executing" and data.get("node") is None: - logger.info(f"Prompt ID: {prompt_id} 执行完成。") - return aggregated_outputs - - except websockets.exceptions.ConnectionClosed as e: - logger.warning( - f"WebSocket 连接已关闭 (Prompt ID: {prompt_id})。错误: {e}" - ) - return aggregated_outputs - except Exception as e: - # 重新抛出我们自己的异常，或者处理其他意外错误 - if not isinstance(e, ComfyUIExecutionError): - logger.error(f"处理 prompt {prompt_id} 时发生意外错误: {e}") - raise e - - except websockets.exceptions.InvalidURI as e: - logger.error( - f"错误: 尝试连接的WebSocket URI无效: '{full_ws_url}'. 原始URL: '{ws_url}'. 错误: {e}" - ) - raise e diff --git a/workflow_service/comfy/comfy_run.py b/workflow_service/comfy/comfy_run.py new file mode 100644 index 0000000..8d28002 --- /dev/null +++ b/workflow_service/comfy/comfy_run.py @@ -0,0 +1,345 @@ +import json +import logging +import random +import uuid +from typing import Optional + +import aiohttp +import websockets +from aiohttp import ClientTimeout + +from workflow_service.comfy.comfy_workflow import ComfyWorkflow, API_OUTPUT_PREFIX +from workflow_service.comfy.comfy_server import ComfyUIServerInfo, server_manager +from workflow_service.database.api import ( + create_workflow_run, + create_workflow_run_nodes, + get_workflow_run, + update_workflow_run_status, + update_workflow_run_node_status, +) + + +class ComfyUIExecutionError(Exception): + """ComfyUI执行错误""" + + def __init__(self, error_data): + self.error_data = error_data + super().__init__(f"ComfyUI execution error: {error_data}") + + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + + +class ComfyRun: + """ + ComfyUI工作流运行实例，封装单个工作流运行的所有操作 + """ + + def __init__(self, workflow: ComfyWorkflow, run_id: str, request_data: dict): + """ + 初始化工作流运行实例 + + Args: + workflow: ComfyWorkflow实例 + run_id: 运行ID + request_data: 请求数据 + """ + self.workflow = workflow + self.run_id = run_id + self.request_data = request_data + self.client_id = str(uuid.uuid4()) + self.prompt_id: Optional[str] = None + self.server: Optional[ComfyUIServerInfo] = None + + @classmethod + async def create(cls, workflow: ComfyWorkflow, request_data: dict) -> "ComfyRun": + """ + 创建新的ComfyRun实例并保存到数据库 + + Args: + workflow: ComfyWorkflow实例 + request_data: 请求数据 + + Returns: + ComfyRun实例 + """ + run_id = str(uuid.uuid4()) + + # 创建任务记录 + await create_workflow_run( + workflow_run_id=run_id, + workflow_name=workflow.workflow_name, + workflow_json=json.dumps(workflow.workflow_data.model_dump()), + api_spec=json.dumps(workflow.get_api_spec().model_dump()), + request_data=json.dumps(request_data), + ) + + # 创建工作流节点记录 + nodes_data = [] + for node in workflow.workflow_data.nodes: + nodes_data.append({"id": str(node.id), "type": node.type}) + await create_workflow_run_nodes(run_id, nodes_data) + + return cls(workflow, run_id, request_data) + + @classmethod + async def from_run_id(cls, run_id: str) -> Optional["ComfyRun"]: + """ + 从run_id创建ComfyRun实例 + + Args: + run_id: 运行ID + + Returns: + ComfyRun实例，如果找不到则返回None + """ + workflow_run = await get_workflow_run(run_id) + if not workflow_run: + return None + + workflow_data = json.loads(workflow_run.workflow_json) + workflow = ComfyWorkflow(workflow_run.workflow_name, workflow_data) + request_data = json.loads(workflow_run.request_data) + + return cls(workflow, run_id, request_data) + + async def execute(self, server: ComfyUIServerInfo) -> dict: + """ + 在指定的服务器上执行工作流 + + Args: + server: ComfyUI服务器信息 + + Returns: + 执行结果 + """ + self.server = server + + try: + # 分配服务器资源 + await server_manager.allocate_server(server.name) + + # 构建prompt + prompt = await self.workflow.build_prompt(server, self.request_data) + + # 更新运行状态为running + await self._update_status( + "running", server.http_url, client_id=self.client_id + ) + + # 提交到ComfyUI + self.prompt_id = await self._queue_prompt(prompt, server.http_url) + + # 更新prompt_id + await self._update_status( + "running", server.http_url, self.prompt_id, self.client_id + ) + + logger.info( + f"工作流 {self.run_id} 已在 {server.http_url} 上入队，Prompt ID: {self.prompt_id}" + ) + + # 获取执行结果 + results = await self._get_execution_results(server.ws_url) + + # 标记完成 + await self._update_status( + "completed", result=json.dumps(results, ensure_ascii=False) + ) + + return results + + except Exception as e: + # 标记失败 + await self._update_status("failed", error_message=str(e)) + raise + + finally: + # 释放服务器资源 + try: + await server_manager.release_server(server.name) + except Exception as e: + logger.error(f"释放服务器资源时出错: {e}") + + async def _update_status( + self, + status: str, + server_url: Optional[str] = None, + prompt_id: Optional[str] = None, + client_id: Optional[str] = None, + result: Optional[str] = None, + error_message: Optional[str] = None, + ): + """更新工作流运行状态""" + await update_workflow_run_status( + self.run_id, status, server_url, prompt_id, client_id, error_message, result + ) + + async def _queue_prompt(self, prompt: dict, http_url: str) -> str: + """提交工作流到ComfyUI服务器""" + # 添加随机缓存破坏器避免缓存 + for node_id in prompt: + prompt[node_id]["inputs"][ + f"cache_buster_{uuid.uuid4().hex}" + ] = random.random() + + payload = { + "prompt": prompt, + "client_id": self.client_id, + "extra_data": { + "api_key_comfy_org": "", + "extra_pnginfo": {"workflow": self.workflow.workflow_data.model_dump()}, + }, + } + + async with aiohttp.ClientSession(timeout=ClientTimeout(total=90)) as session: + prompt_url = f"{http_url}/prompt" + try: + async with session.post(prompt_url, json=payload) as response: + logger.info(f"ComfyUI /prompt 端点返回的响应: {response}") + response.raise_for_status() + result = await response.json() + logger.info(f"ComfyUI /prompt 端点返回的响应: {result}") + + if "prompt_id" not in result: + raise Exception( + f"从 ComfyUI /prompt 端点返回的响应无效: {result}" + ) + + return result["prompt_id"] + + except Exception as e: + logger.error(f"提交到 ComfyUI /prompt 端点时发生错误: {e}") + raise + + async def _get_execution_results(self, ws_url: str) -> dict: + """ + 通过WebSocket连接获取执行结果，支持节点级别的状态跟踪 + """ + aggregated_outputs = {} + full_ws_url = f"{ws_url}?clientId={self.client_id}" + + try: + async with websockets.connect(full_ws_url) as websocket: + logger.info(f"已连接到WebSocket: {full_ws_url}") + + while True: + try: + out = await websocket.recv() + if not isinstance(out, str): + continue + + message = json.loads(out) + msg_type = message.get("type") + data = message.get("data") + + if not (data and data.get("prompt_id") == self.prompt_id): + continue + + # 捕获并处理执行错误 + if msg_type == "execution_error": + error_data = data + logger.error( + f"ComfyUI执行错误 (Prompt ID: {self.prompt_id}): {error_data}" + ) + + # 更新节点状态为失败 + node_id = error_data.get("node_id") + if node_id: + await update_workflow_run_node_status( + self.run_id, + node_id, + "failed", + error_message=error_data.get( + "exception_message", "Unknown error" + ), + ) + + # 抛出自定义异常，将错误详情传递出去 + raise ComfyUIExecutionError(error_data) + + # 处理节点开始执行 + if msg_type == "executing" and data.get("node"): + node_id = data.get("node") + logger.info( + f"节点 {node_id} 开始执行 (Prompt ID: {self.prompt_id})" + ) + + # 更新节点状态为运行中 + await update_workflow_run_node_status( + self.run_id, node_id, "running" + ) + + # 处理节点执行完成 + elif msg_type == "executed": + await self._handle_node_executed(data, aggregated_outputs) + + # 处理整个工作流执行完成 + elif msg_type == "executing" and data.get("node") is None: + logger.info(f"Prompt ID: {self.prompt_id} 执行完成。") + return aggregated_outputs + + except websockets.exceptions.ConnectionClosed as e: + logger.warning( + f"WebSocket 连接已关闭 (Prompt ID: {self.prompt_id})。错误: {e}" + ) + return aggregated_outputs + except Exception as e: + # 重新抛出我们自己的异常，或者处理其他意外错误 + if not isinstance(e, ComfyUIExecutionError): + logger.error( + f"处理 prompt {self.prompt_id} 时发生意外错误: {e}" + ) + raise e + + except websockets.exceptions.InvalidURI as e: + logger.error( + f"错误: 尝试连接的WebSocket URI无效: '{full_ws_url}'. 原始URL: '{ws_url}'. 错误: {e}" + ) + raise e + except Exception as e: + logger.error(f"WebSocket连接出错: {e}") + raise + + async def _handle_node_executed(self, data: dict, aggregated_outputs: dict): + """处理节点执行完成事件""" + node_id = data.get("node") + output_data = data.get("output") + + if not node_id or not output_data: + return + + # 查找对应的节点 + node = next( + (x for x in self.workflow.workflow_data.nodes if str(x.id) == node_id), + None, + ) + + if not node: + return + + # 如果是输出节点，收集结果 + if node.title and node.title.startswith(API_OUTPUT_PREFIX): + title = node.title.replace(API_OUTPUT_PREFIX, "") + aggregated_outputs[title] = output_data + + logger.info(f"收到节点 {node_id} 的输出 (Prompt ID: {self.prompt_id})") + + # 更新节点状态为完成 + await update_workflow_run_node_status( + self.run_id, + node_id, + "completed", + output_data=json.dumps(output_data), + ) + + def get_status_info(self) -> dict: + """获取运行状态信息""" + return { + "run_id": self.run_id, + "workflow_name": self.workflow.workflow_name, + "client_id": self.client_id, + "prompt_id": self.prompt_id, + "server_url": self.server.http_url if self.server else None, + } diff --git a/workflow_service/routes/run.py b/workflow_service/routes/run.py index 01df260..f96df76 100644 --- a/workflow_service/routes/run.py +++ b/workflow_service/routes/run.py @@ -7,7 +7,12 @@ from fastapi.responses import JSONResponse from workflow_service.comfy.comfy_queue import queue_manager from workflow_service.comfy.comfy_workflow import ComfyWorkflow -from workflow_service.database.api import get_workflow, get_workflow_runs_recent +from workflow_service.database.api import ( + get_workflow, + get_workflow_run, + get_workflow_run_nodes, + get_workflow_runs_recent, +) run_router = APIRouter( prefix="/api/run", @@ -133,7 +138,30 @@ async def get_run_status(workflow_run_id: str): 获取工作流执行状态。 """ try: - status = await queue_manager.get_task_status(workflow_run_id) - return status + """获取任务状态""" + workflow_run = await get_workflow_run(workflow_run_id) + if not workflow_run: + return {"error": "任务不存在"} + nodes = await get_workflow_run_nodes(workflow_run_id) + + result = { + "id": workflow_run_id, + "status": workflow_run.status, + "created_at": workflow_run.created_at, + "started_at": workflow_run.started_at, + "completed_at": workflow_run.completed_at, + "server_url": workflow_run.server_url, + "error_message": workflow_run.error_message, + "nodes": nodes, + } + + # 如果任务完成，从数据库获取结果 + if workflow_run.status == "completed" and workflow_run.result: + try: + result["result"] = json.loads(workflow_run.result) + except (json.JSONDecodeError, TypeError): + result["result"] = None + + return result except Exception as e: raise HTTPException(status_code=500, detail=f"获取状态失败: {str(e)}")