From 1af6b425731e0559b161eddf6d4fdfe024de2c7a Mon Sep 17 00:00:00 2001 From: iHeyTang Date: Wed, 13 Aug 2025 15:16:02 +0800 Subject: [PATCH] fix --- .env | 51 ++-- run_service.py | 12 + workflow_service/comfyui_client.py | 460 +++++++++++++++++++++++++--- workflow_service/database.py | 180 ++++++++++- workflow_service/main.py | 464 +++++++++++++++++------------ 5 files changed, 919 insertions(+), 248 deletions(-) diff --git a/.env b/.env index 50e540a..6a9eb05 100644 --- a/.env +++ b/.env @@ -1,29 +1,28 @@ -COMFYUI_SERVERS_JSON='[ - { - "http_url": "https://bowongai-dev--comfyui-for-waas-ui-1.modal.run", - "ws_url": "wss://bowongai-dev--comfyui-for-waas-ui-1.modal.run/ws", - "input_dir": "waas/comfyui_1/input", - "output_dir": "waas/comfyui_1/output" - }, - { - "http_url": "https://bowongai-dev--comfyui-for-waas-ui-2.modal.run", - "ws_url": "wss://bowongai-dev--comfyui-for-waas-ui-2.modal.run/ws", - "input_dir": "waas/comfyui_2/input", - "output_dir": "waas/comfyui_2/output" - }, - { - "http_url": "https://bowongai-dev--comfyui-for-waas-ui-3.modal.run", - "ws_url": "wss://bowongai-dev--comfyui-for-waas-ui-3.modal.run/ws", - "input_dir": "waas/comfyui_3/input", - "output_dir": "waas/comfyui_3/output" - }, - { - "http_url": "https://bowongai-dev--comfyui-for-waas-ui-4.modal.run", - "ws_url": "wss://bowongai-dev--comfyui-for-waas-ui-4.modal.run/ws", - "input_dir": "waas/comfyui_4/input", - "output_dir": "waas/comfyui_4/output" - } -]' + # { + # "http_url": "https://bowongai-dev--comfyui-for-waas-ui-1.modal.run", + # "ws_url": "wss://bowongai-dev--comfyui-for-waas-ui-1.modal.run/ws", + # "input_dir": "waas/comfyui_1/input", + # "output_dir": "waas/comfyui_1/output" + # }, + # { + # "http_url": "https://bowongai-dev--comfyui-for-waas-ui-2.modal.run", + # "ws_url": "wss://bowongai-dev--comfyui-for-waas-ui-2.modal.run/ws", + # "input_dir": "waas/comfyui_2/input", + # "output_dir": "waas/comfyui_2/output" + # }, + # { + # "http_url": "https://bowongai-dev--comfyui-for-waas-ui-3.modal.run", + # "ws_url": "wss://bowongai-dev--comfyui-for-waas-ui-3.modal.run/ws", + # "input_dir": "waas/comfyui_3/input", + # "output_dir": "waas/comfyui_3/output" + # }, + # { + # "http_url": "https://bowongai-dev--comfyui-for-waas-ui-4.modal.run", + # "ws_url": "wss://bowongai-dev--comfyui-for-waas-ui-4.modal.run/ws", + # "input_dir": "waas/comfyui_4/input", + # "output_dir": "waas/comfyui_4/output" + # } +COMFYUI_SERVERS_JSON='[{"http_url":"http://127.0.0.1:8000","ws_url":"ws://127.0.0.1:8000/ws","input_dir":"waas/comfyui_4/input","output_dir":"waas/comfyui_4/output"}]' DB_FILE="workflows_service.sqlite" diff --git a/run_service.py b/run_service.py index e567aa9..22c4b3d 100644 --- a/run_service.py +++ b/run_service.py @@ -4,6 +4,18 @@ """ import sys import os +from pathlib import Path + +# 加载 .env 文件 +try: + from dotenv import load_dotenv + env_path = Path(__file__).parent / ".env" + load_dotenv(env_path) + print(f"已加载环境变量文件: {env_path}") +except ImportError: + print("警告: 未安装 python-dotenv，无法加载 .env 文件") +except Exception as e: + print(f"警告: 加载 .env 文件失败: {e}") # 将项目根目录添加到Python路径 project_root = os.path.dirname(os.path.abspath(__file__)) diff --git a/workflow_service/comfyui_client.py b/workflow_service/comfyui_client.py index 5f0ffac..95e27e6 100644 --- a/workflow_service/comfyui_client.py +++ b/workflow_service/comfyui_client.py @@ -1,24 +1,211 @@ import asyncio -from collections import defaultdict -from ctypes import Union import json +import logging +import os import random import uuid -from typing import Dict, Any +import websockets +from collections import defaultdict +from datetime import datetime +from typing import Dict, Any, Optional, List, Set import aiohttp -import websockets from aiohttp import ClientTimeout from workflow_service.config import Settings, ComfyUIServer +from workflow_service.database import ( + create_workflow_run, + update_workflow_run_status, + create_workflow_run_nodes, + update_workflow_run_node_status, + get_workflow_run, + get_pending_workflow_runs, + get_running_workflow_runs, + get_workflow_run_nodes, +) settings = Settings() API_INPUT_PREFIX = "INPUT_" API_OUTPUT_PREFIX = "OUTPUT_" +# 配置日志 +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) -# [新增] 定义一个自定义异常，用于封装来自ComfyUI的执行错误 + +# 全局任务队列管理器 +class WorkflowQueueManager: + def __init__(self): + self.running_tasks = {} # server_url -> task_info + self.pending_tasks = [] # 等待队列 + self.lock = asyncio.Lock() + + async def add_task( + self, + workflow_run_id: str, + workflow_name: str, + workflow_data: dict, + api_spec: dict, + request_data: dict, + ): + """添加新任务到队列""" + async with self.lock: + # 创建任务记录 + await create_workflow_run( + workflow_run_id=workflow_run_id, + workflow_name=workflow_name, + workflow_json=json.dumps(workflow_data), + api_spec=json.dumps(api_spec), + request_data=json.dumps(request_data), + ) + + # 创建工作流节点记录 + nodes_data = [] + for node in workflow_data.get("nodes", []): + nodes_data.append( + {"id": str(node["id"]), "type": node.get("type", "unknown")} + ) + 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)}" + ) + + # 尝试处理队列 + asyncio.create_task(self._process_queue()) + + return workflow_run_id + + async def _process_queue(self): + """处理队列中的任务""" + async with self.lock: + if not self.pending_tasks: + return + + # 检查是否有空闲的服务器 + available_servers = await self._get_available_servers() + if not available_servers: + logger.info("没有可用的服务器，等待中...") + return + + # 获取一个待处理任务 + workflow_run_id = self.pending_tasks.pop(0) + server = available_servers[0] + + # 标记任务为运行中 + 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(), + } + + # 启动任务执行 + asyncio.create_task(self._execute_task(workflow_run_id, server)) + + async def _get_available_servers(self) -> list[ComfyUIServer]: + """获取可用的服务器""" + available_servers = [] + for server in settings.SERVERS: + if server.http_url not in self.running_tasks: + # 检查服务器状态 + try: + async with aiohttp.ClientSession() as session: + status = await get_server_status(server, session) + if status["is_reachable"] and status["is_free"]: + available_servers.append(server) + except Exception as e: + logger.warning(f"检查服务器 {server.http_url} 状态时出错: {e}") + + return available_servers + + async def _execute_task(self, workflow_run_id: str, server: ComfyUIServer): + """执行任务""" + cleanup_paths = [] + try: + # 获取工作流数据 + workflow_run = await get_workflow_run(workflow_run_id) + if not workflow_run: + raise Exception(f"找不到工作流运行记录: {workflow_run_id}") + + workflow_data = json.loads(workflow_run["workflow_json"]) + api_spec = json.loads(workflow_run["api_spec"]) + request_data = json.loads(workflow_run["request_data"]) + + # 执行工作流 + result = await execute_prompt_on_server( + workflow_data, api_spec, request_data, server, workflow_run_id + ) + + # 保存处理后的结果到数据库 + await update_workflow_run_status( + workflow_run_id, + "completed", + result=json.dumps(result, ensure_ascii=False), + ) + + 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}") + + # 清理运行状态 + async with self.lock: + if server.http_url in self.running_tasks: + del self.running_tasks[server.http_url] + + # 继续处理队列 + asyncio.create_task(self._process_queue()) + + 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.get("result"): + try: + result["result"] = json.loads(workflow_run["result"]) + except (json.JSONDecodeError, TypeError): + result["result"] = None + + return result + + +# 全局队列管理器实例 +queue_manager = WorkflowQueueManager() + + +# 定义一个自定义异常，用于封装来自ComfyUI的执行错误 class ComfyUIExecutionError(Exception): def __init__(self, error_data: dict): self.error_data = error_data @@ -64,7 +251,7 @@ async def get_server_status( except Exception as e: # 当请求失败时，将返回上面定义的、结构正确的初始 status_info - print(f"警告: 无法检查服务器 {server.http_url} 的队列状态: {e}") + logger.warning(f"无法检查服务器 {server.http_url} 的队列状态: {e}") return status_info @@ -88,7 +275,7 @@ async def select_server_for_execution() -> ComfyUIServer: if free_servers: selected_server = random.choice(free_servers) - print( + logger.info( f"发现 {len(free_servers)} 个空闲服务器。已选择: {selected_server.http_url}" ) return selected_server @@ -99,7 +286,7 @@ async def select_server_for_execution() -> ComfyUIServer: ] if reachable_servers: selected_server = random.choice(reachable_servers) - print( + logger.info( f"所有服务器当前都在忙。从可达服务器中随机选择: {selected_server.http_url}" ) return selected_server @@ -108,17 +295,70 @@ async def select_server_for_execution() -> ComfyUIServer: raise ConnectionError("所有配置的ComfyUI服务器都不可达。") -async def execute_prompt_on_server(prompt: Dict, server: ComfyUIServer) -> Dict: +async def execute_prompt_on_server( + workflow_data: Dict, + api_spec: Dict, + request_data: Dict, + server: ComfyUIServer, + workflow_run_id: str, +) -> Dict: """ 在指定的服务器上执行一个准备好的prompt。 + 现在支持节点级别的状态跟踪。 """ client_id = str(uuid.uuid4()) - prompt_id = await _queue_prompt(prompt, client_id, server.http_url) - print(f"工作流已在 {server.http_url} 上入队，Prompt ID: {prompt_id}") - results = await _get_execution_results(prompt_id, client_id, server.ws_url) + + # 应用请求数据到工作流 + patched_workflow = patch_workflow(workflow_data, api_spec, request_data) + + # 转换为prompt格式 + prompt = convert_workflow_to_prompt_api_format(patched_workflow) + + # 更新工作流运行状态，记录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_data, 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_data, prompt_id, client_id, server.ws_url, workflow_run_id + ) return results +async def submit_workflow_to_queue( + workflow_name: str, workflow_data: Dict, api_spec: Dict, request_data: Dict +) -> str: + """ + 提交工作流到队列，立即返回任务ID。 + 这是新的异步接口，调用者可以通过任务ID查询状态。 + """ + workflow_run_id = str(uuid.uuid4()) + + # 添加到队列管理器 + await queue_manager.add_task( + workflow_run_id, workflow_name, workflow_data, api_spec, request_data + ) + + return workflow_run_id + + def parse_api_spec(workflow_data: dict) -> Dict[str, Dict[str, Any]]: """ 解析工作流，并根据规范 '{基础名}_{属性名}_{可选计数}' 生成API参数名。 @@ -198,9 +438,13 @@ def parse_api_spec(workflow_data: dict) -> Dict[str, Dict[str, Any]]: return spec -def patch_workflow(workflow_data: dict, api_spec: dict, request_data: dict) -> dict: +def patch_workflow( + workflow_data: dict, + api_spec: dict[str, dict[str, Any]], + request_data: dict[str, Any], +) -> dict: """ - [BUG修复] 根据API请求数据，正确地修改工作流JSON中的 `widgets_values`。 + 将request_data中的参数值，patch到workflow_data中。并返回修改后的workflow_data。 """ if "nodes" not in workflow_data: raise ValueError("无效的工作流格式") @@ -232,8 +476,8 @@ def patch_workflow(workflow_data: dict, api_spec: dict, request_data: dict) -> d break if target_widget_index == -1: - print( - f"警告: 在节点 {node_id} 中未找到名为 '{widget_name_to_patch}' 的 widget。跳过此参数。" + logger.warning( + f"在节点 {node_id} 中未找到名为 '{widget_name_to_patch}' 的 widget。跳过此参数。" ) continue @@ -259,8 +503,8 @@ def patch_workflow(workflow_data: dict, api_spec: dict, request_data: dict) -> d try: target_node["widgets_values"][target_widget_index] = target_type(value) except (ValueError, TypeError) as e: - print( - f"警告: 无法将参数 '{param_name}' 的值 '{value}' 转换为类型 '{spec['type']}'。错误: {e}" + logger.warning( + f"无法将参数 '{param_name}' 的值 '{value}' 转换为类型 '{spec['type']}'。错误: {e}" ) continue @@ -322,25 +566,47 @@ def convert_workflow_to_prompt_api_format(workflow_data: dict) -> dict: return prompt_api_format -async def _queue_prompt(prompt: dict, client_id: str, http_url: str) -> str: +async def _queue_prompt( + workflow: dict, 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} + payload = { + "prompt": prompt, + "client_id": client_id, + "extra_data": { + "api_key_comfy_org": "", + "extra_pnginfo": {"workflow": workflow}, + }, + } + logger.info(f"提交到 ComfyUI /prompt 端点的payload: {json.dumps(payload)}") async with aiohttp.ClientSession(timeout=ClientTimeout(total=90)) as session: prompt_url = f"{http_url}/prompt" - async with session.post(prompt_url, json=payload) as response: - response.raise_for_status() - result = await response.json() - if "prompt_id" not in result: - raise Exception(f"从 ComfyUI /prompt 端点返回的响应无效: {result}") - return result["prompt_id"] + 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(prompt_id: str, client_id: str, ws_url: str) -> dict: +async def _get_execution_results( + workflow_data: Dict, + prompt_id: str, + client_id: str, + ws_url: str, + workflow_run_id: str, +) -> dict: """ 通过WebSocket连接到指定的ComfyUI服务器，聚合执行结果。 - [核心改动] 新增对 'execution_error' 消息的处理。 + 现在支持节点级别的状态跟踪。 """ full_ws_url = f"{ws_url}?clientId={client_id}" aggregated_outputs = {} @@ -360,34 +626,154 @@ async def _get_execution_results(prompt_id: str, client_id: str, ws_url: str) -> if not (data and data.get("prompt_id") == prompt_id): continue - # [核心改动] 捕获并处理执行错误 + # 捕获并处理执行错误 if msg_type == "execution_error": - print(f"ComfyUI执行错误 (Prompt ID: {prompt_id}): {data}") - # 抛出自定义异常，将错误详情传递出去 - raise ComfyUIExecutionError(data) + 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_data["nodes"] + if str(x["id"]) == node_id + ), + None, + ) + if ( + node + and node.get("title", "") + and node["title"].startswith(API_OUTPUT_PREFIX) + ): + aggregated_outputs[node["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 + return aggregated_outputs + + +async def _get_execution_results_legacy( + prompt_id: str, client_id: str, ws_url: str +) -> dict: + """ + 简化版的执行结果获取函数，不包含数据库状态跟踪。 + """ + 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}" + ) + # 抛出自定义异常，将错误详情传递出去 + raise ComfyUIExecutionError(error_data) + + # 处理节点执行完成 if msg_type == "executed": node_id = data.get("node") output_data = data.get("output") if node_id and output_data: aggregated_outputs[node_id] = output_data - print(f"收到节点 {node_id} 的输出 (Prompt ID: {prompt_id})") + logger.info( + f"收到节点 {node_id} 的输出 (Prompt ID: {prompt_id})" + ) + # 处理整个工作流执行完成 elif msg_type == "executing" and data.get("node") is None: - print(f"Prompt ID: {prompt_id} 执行完成。") + logger.info(f"Prompt ID: {prompt_id} 执行完成。") return aggregated_outputs except websockets.exceptions.ConnectionClosed as e: - print(f"WebSocket 连接已关闭 (Prompt ID: {prompt_id})。错误: {e}") + logger.warning( + f"WebSocket 连接已关闭 (Prompt ID: {prompt_id})。错误: {e}" + ) return aggregated_outputs except Exception as e: # 重新抛出我们自己的异常，或者处理其他意外错误 if not isinstance(e, ComfyUIExecutionError): - print(f"处理 prompt {prompt_id} 时发生意外错误: {e}") + logger.error(f"处理 prompt {prompt_id} 时发生意外错误: {e}") raise e except websockets.exceptions.InvalidURI as e: - print( + logger.error( f"错误: 尝试连接的WebSocket URI无效: '{full_ws_url}'. 原始URL: '{ws_url}'. 错误: {e}" ) raise e diff --git a/workflow_service/database.py b/workflow_service/database.py index 93c2c57..ded3de0 100644 --- a/workflow_service/database.py +++ b/workflow_service/database.py @@ -3,6 +3,7 @@ from typing import Optional import aiosqlite import json import re +import uuid from workflow_service.config import Settings @@ -18,6 +19,43 @@ async def init_db(): workflow_json TEXT NOT NULL ) """) + + # 新增 workflow_run 表 + await db.execute(""" + CREATE TABLE IF NOT EXISTS workflow_run ( + id TEXT PRIMARY KEY, + workflow_name TEXT NOT NULL, + prompt_id TEXT, + client_id TEXT, + status TEXT NOT NULL DEFAULT 'pending', + server_url TEXT, + created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + started_at TIMESTAMP, + completed_at TIMESTAMP, + error_message TEXT, + workflow_json TEXT NOT NULL, + api_spec TEXT NOT NULL, + request_data TEXT NOT NULL, + result TEXT + ) + """) + + # 新增 workflow_run_nodes 表，记录每个节点的运行状态 + await db.execute(""" + CREATE TABLE IF NOT EXISTS workflow_run_nodes ( + id TEXT PRIMARY KEY, + workflow_run_id TEXT NOT NULL, + node_id TEXT NOT NULL, + node_type TEXT NOT NULL, + status TEXT NOT NULL DEFAULT 'pending', + started_at TIMESTAMP, + completed_at TIMESTAMP, + output_data TEXT, + error_message TEXT, + FOREIGN KEY (workflow_run_id) REFERENCES workflow_run (id) + ) + """) + await db.commit() print(f"数据库 '{DATABASE_FILE}' 已准备就绪。") @@ -72,4 +110,144 @@ async def delete_workflow(name: str) -> bool: async with aiosqlite.connect(DATABASE_FILE) as db: cursor = await db.execute("DELETE FROM workflows WHERE name = ?", (name,)) await db.commit() - return cursor.rowcount > 0 \ No newline at end of file + return cursor.rowcount > 0 + +# 新增 workflow_run 相关函数 +async def create_workflow_run( + workflow_run_id: str, + workflow_name: str, + workflow_json: str, + api_spec: str, + request_data: str +) -> str: + """创建新的工作流运行记录""" + async with aiosqlite.connect(DATABASE_FILE) as db: + await db.execute(""" + INSERT INTO workflow_run (id, workflow_name, workflow_json, api_spec, request_data) + VALUES (?, ?, ?, ?, ?) + """, (workflow_run_id, workflow_name, workflow_json, api_spec, request_data)) + await db.commit() + return workflow_run_id + +async def update_workflow_run_status( + workflow_run_id: str, + status: str, + server_url: str = None, + prompt_id: str = None, + client_id: str = None, + error_message: str = None, + result: str = None +): + """更新工作流运行状态""" + async with aiosqlite.connect(DATABASE_FILE) as db: + if status == 'running': + await db.execute(""" + UPDATE workflow_run + SET status = ?, server_url = ?, prompt_id = ?, client_id = ?, started_at = CURRENT_TIMESTAMP + WHERE id = ? + """, (status, server_url, prompt_id, client_id, workflow_run_id)) + elif status == 'completed': + await db.execute(""" + UPDATE workflow_run + SET status = ?, completed_at = CURRENT_TIMESTAMP, result = ? + WHERE id = ? + """, (status, result, workflow_run_id)) + elif status == 'failed': + await db.execute(""" + UPDATE workflow_run + SET status = ?, error_message = ?, completed_at = CURRENT_TIMESTAMP + WHERE id = ? + """, (status, error_message, workflow_run_id)) + else: + await db.execute(""" + UPDATE workflow_run + SET status = ? + WHERE id = ? + """, (status, workflow_run_id)) + await db.commit() + +async def create_workflow_run_nodes(workflow_run_id: str, nodes_data: list): + """创建工作流运行节点记录""" + async with aiosqlite.connect(DATABASE_FILE) as db: + for node in nodes_data: + node_id = str(uuid.uuid4()) + await db.execute(""" + INSERT INTO workflow_run_nodes (id, workflow_run_id, node_id, node_type) + VALUES (?, ?, ?, ?) + """, (node_id, workflow_run_id, node["id"], node["type"])) + await db.commit() + +async def update_workflow_run_node_status( + workflow_run_id: str, + node_id: str, + status: str, + output_data: str = None, + error_message: str = None +): + """更新工作流运行节点状态""" + async with aiosqlite.connect(DATABASE_FILE) as db: + if status == 'running': + await db.execute(""" + UPDATE workflow_run_nodes + SET status = ?, started_at = CURRENT_TIMESTAMP + WHERE workflow_run_id = ? AND node_id = ? + """, (status, workflow_run_id, node_id)) + elif status == 'completed': + await db.execute(""" + UPDATE workflow_run_nodes + SET status = ?, output_data = ?, completed_at = CURRENT_TIMESTAMP + WHERE workflow_run_id = ? AND node_id = ? + """, (status, output_data, workflow_run_id, node_id)) + elif status == 'failed': + await db.execute(""" + UPDATE workflow_run_nodes + SET status = ?, error_message = ?, completed_at = CURRENT_TIMESTAMP + WHERE workflow_run_id = ? AND node_id = ? + """, (status, error_message, workflow_run_id, node_id)) + else: + await db.execute(""" + UPDATE workflow_run_nodes + SET status = ? + WHERE workflow_run_id = ? AND node_id = ? + """, (status, workflow_run_id, node_id)) + await db.commit() + +async def get_workflow_run(workflow_run_id: str) -> dict | None: + """获取工作流运行记录""" + async with aiosqlite.connect(DATABASE_FILE) as db: + db.row_factory = aiosqlite.Row + cursor = await db.execute(""" + SELECT * FROM workflow_run WHERE id = ? + """, (workflow_run_id,)) + row = await cursor.fetchone() + return dict(row) if row else None + +async def get_workflow_run_nodes(workflow_run_id: str) -> list[dict]: + """获取工作流运行节点记录""" + async with aiosqlite.connect(DATABASE_FILE) as db: + db.row_factory = aiosqlite.Row + cursor = await db.execute(""" + SELECT * FROM workflow_run_nodes WHERE workflow_run_id = ? + """, (workflow_run_id,)) + rows = await cursor.fetchall() + return [dict(row) for row in rows] + +async def get_pending_workflow_runs() -> list[dict]: + """获取所有待处理的工作流运行记录""" + async with aiosqlite.connect(DATABASE_FILE) as db: + db.row_factory = aiosqlite.Row + cursor = await db.execute(""" + SELECT * FROM workflow_run WHERE status = 'pending' ORDER BY created_at ASC + """) + rows = await cursor.fetchall() + return [dict(row) for row in rows] + +async def get_running_workflow_runs() -> list[dict]: + """获取所有正在运行的工作流运行记录""" + async with aiosqlite.connect(DATABASE_FILE) as db: + db.row_factory = aiosqlite.Row + cursor = await db.execute(""" + SELECT * FROM workflow_run WHERE status = 'running' + """) + rows = await cursor.fetchall() + return [dict(row) for row in rows] \ No newline at end of file diff --git a/workflow_service/main.py b/workflow_service/main.py index 9baa004..fe76cb1 100644 --- a/workflow_service/main.py +++ b/workflow_service/main.py @@ -15,7 +15,11 @@ from pydantic import BaseModel from workflow_service import comfyui_client from workflow_service import database from workflow_service.utils.s3_client import upload_file_to_s3 -from workflow_service.comfyui_client import ComfyUIExecutionError +from workflow_service.comfyui_client import ( + ComfyUIExecutionError, + queue_manager, + submit_workflow_to_queue, +) from workflow_service.config import Settings settings = Settings() @@ -81,7 +85,6 @@ async def startup_event(): print(f"错误: 无法在启动时初始化服务器目录: {e}") -# --- Section 1: 工作流管理API --- @web_app.post("/api/workflow", status_code=201) async def publish_workflow_endpoint(request: Request): """ @@ -117,7 +120,6 @@ async def get_all_workflows_endpoint(): raise HTTPException(status_code=500, detail=f"Failed to get workflows: {e}") -# [BUG修复] 修正API路径，使其与其他管理端点保持一致 @web_app.delete(f"/api/workflow/{{workflow_name:path}}") async def delete_workflow_endpoint( workflow_name: str = Path( @@ -140,179 +142,8 @@ async def delete_workflow_endpoint( raise HTTPException(status_code=500, detail=f"Failed to delete workflow: {e}") -# --- Section 2: 工作流执行API (核心改动) --- - - -def get_files_in_dir(directory: str) -> Set[str]: - file_set = set() - for root, _, files in os.walk(directory): - for filename in files: - if not filename.startswith("."): - file_set.add(os.path.join(root, filename)) - return file_set - - -async def download_file_from_url( - session: aiohttp.ClientSession, url: str, save_path: str -): - async with session.get(url) as response: - response.raise_for_status() - with open(save_path, "wb") as f: - while True: - chunk = await response.content.read(8192) - if not chunk: - break - f.write(chunk) - - -async def handle_file_upload(file_path: str, base_name: str) -> str: - s3_object_name = f"outputs/{base_name}/{uuid.uuid4()}_{os.path.basename(file_path)}" - return await upload_file_to_s3(file_path, settings.S3_BUCKET_NAME, s3_object_name) - - -@web_app.post("/api/run/") -async def execute_workflow_endpoint( - base_name: str, request_data_raw: Dict[str, Any], version: Optional[str] = None -): - """ - 执行工作流 - """ - cleanup_paths = [] - try: - # 1. 获取工作流定义 - workflow_data = await database.get_workflow(base_name, version) - - if not workflow_data: - detail = ( - f"工作流 '{base_name}'" - + (f" 带版本 '{version}'" if version else " (最新版)") - + " 未找到。" - ) - raise HTTPException(status_code=404, detail=detail) - - workflow = json.loads(workflow_data["workflow_json"]) - api_spec = comfyui_client.parse_api_spec(workflow) - request_data = {k.lower(): v for k, v in request_data_raw.items()} - - # 2. [核心改动] 使用智能调度选择服务器 - try: - selected_server = await comfyui_client.select_server_for_execution() - except Exception as e: - raise HTTPException( - status_code=503, detail=f"无法选择ComfyUI服务器执行任务: {e}" - ) - - server_input_dir = selected_server.input_dir - server_output_dir = selected_server.output_dir - - # 3. 下载文件到选定服务器的输入目录 - async with aiohttp.ClientSession() as session: - for param_name, param_def in api_spec["inputs"].items(): - if param_def["type"] == "UploadFile" and param_name in request_data: - image_url = request_data[param_name] - if not isinstance(image_url, str) or not image_url.startswith( - "http" - ): - raise HTTPException( - status_code=400, - detail=f"参数 '{param_name}' 必须是一个有效的URL。", - ) - - filename = f"api_download_{uuid.uuid4().hex}{os.path.splitext(image_url.split('/')[-1])[1] or '.dat'}" - save_path = os.path.join(server_input_dir, filename) - try: - await download_file_from_url(session, image_url, save_path) - # 直接更新 request_data，以便后续的 patch_workflow 使用 - request_data[param_name] = filename - cleanup_paths.append(save_path) - except Exception as e: - raise HTTPException( - status_code=500, - detail=f"为 '{param_name}' 从 {image_url} 下载文件失败。错误: {e}", - ) - - # 4. Patch工作流，生成最终的API Prompt - patched_workflow = comfyui_client.patch_workflow( - workflow, api_spec, request_data - ) - prompt_to_run = comfyui_client.convert_workflow_to_prompt_api_format( - patched_workflow - ) - - # 5. 执行前快照，并在选定服务器上执行工作流 - files_before = get_files_in_dir(server_output_dir) - output_nodes = await comfyui_client.execute_prompt_on_server( - prompt_to_run, selected_server - ) - - # 6. 处理输出（与之前逻辑相同） - files_after = get_files_in_dir(server_output_dir) - new_files = files_after - files_before - - output_response = {} - if new_files: - s3_urls = [] - for file_path in new_files: - cleanup_paths.append(file_path) - try: - s3_urls.append(await handle_file_upload(file_path, base_name)) - except Exception as e: - print(f"上传文件 {file_path} 到S3时出错: {e}") - if s3_urls: - output_response["output_files"] = s3_urls - - for final_param_name, param_def in api_spec["outputs"].items(): - node_id = param_def["node_id"] - if node_id in output_nodes: - node_output = output_nodes[node_id] - original_output_name = param_def["output_name"] - - output_value = None - if original_output_name in node_output: - output_value = node_output[original_output_name] - elif "text" in node_output: # 备选，例如对于'ShowText'节点 - output_value = node_output["text"] - - if output_value is None: - continue - - if isinstance(output_value, list): - output_value = output_value[0] if output_value else None - - output_response[final_param_name] = output_value - - return JSONResponse(content=output_response) - # [核心改动] 捕获来自ComfyUI的执行失败异常 - except ComfyUIExecutionError as e: - print(f"捕获到ComfyUI执行错误: {e.error_data}") - # 返回 502 Bad Gateway 状态码，表示上游服务器出错 - # detail 中包含结构化的错误信息，方便客户端处理 - raise HTTPException( - status_code=500, - detail={ - "message": "工作流在上游ComfyUI节点中执行失败。", - "error_details": e.error_data, - }, - ) - except Exception as e: - # 捕获其他所有异常，作为通用的服务器内部错误 - print(f"执行工作流时发生未知错误: {e}") - raise HTTPException(status_code=500, detail=str(e)) - finally: - if cleanup_paths: - print(f"正在清理 {len(cleanup_paths)} 个临时文件...") - for path in cleanup_paths: - try: - if os.path.exists(path): - os.remove(path) - print(f" - 已删除: {path}") - except Exception as e: - print(f" - 删除 {path} 时出错: {e}") - - -# --- Section 3: 工作流元数据/规范API --- -@web_app.get("/api/spec/") -async def get_workflow_spec_endpoint(base_name: str, version: Optional[str] = None): +@web_app.get("/api/workflow/{base_name:path}") +async def get_one_workflow_endpoint(base_name: str, version: Optional[str] = None): """ 获取工作流规范 """ @@ -328,14 +159,231 @@ async def get_workflow_spec_endpoint(base_name: str, version: Optional[str] = No try: workflow = json.loads(workflow_data["workflow_json"]) - return comfyui_client.parse_api_spec(workflow) + return { + "workflow": workflow, + "api_spec": comfyui_client.parse_api_spec(workflow), + } except Exception as e: raise HTTPException( status_code=500, detail=f"Failed to parse workflow specification: {e}" ) -# --- NEW Section: 服务器监控API --- +@web_app.post("/api/run") +async def run_workflow( + request: Request, workflow_name: str, workflow_version: Optional[str] = None +): + """ + 异步执行工作流。 + 立即返回任务ID，调用者可以通过任务ID查询执行状态。 + """ + try: + data = await request.json() + + if not workflow_name: + raise HTTPException(status_code=400, detail="`workflow_name` 字段是必需的") + + # 获取工作流定义 + workflow_data = await database.get_workflow(workflow_name, workflow_version) + if not workflow_data: + detail = ( + f"工作流 '{workflow_name}'" + + (f" 带版本 '{workflow_version}'" if workflow_version else " (最新版)") + + " 未找到。" + ) + raise HTTPException(status_code=404, detail=detail) + + workflow = json.loads(workflow_data["workflow_json"]) + api_spec = comfyui_client.parse_api_spec(workflow) + + # 提交到队列 + workflow_run_id = await submit_workflow_to_queue( + workflow_name=workflow_name, workflow_data=workflow, api_spec=api_spec, request_data=data + ) + + return JSONResponse( + content={ + "workflow_run_id": workflow_run_id, + "status": "queued", + "message": "工作流已提交到队列，正在等待执行", + }, + status_code=202, + ) + + except Exception as e: + raise HTTPException(status_code=500, detail=f"提交工作流失败: {str(e)}") + + +@web_app.get("/api/run/{workflow_run_id}") +async def get_run_status(workflow_run_id: str): + """ + 获取工作流执行状态。 + """ + try: + status = await queue_manager.get_task_status(workflow_run_id) + return status + except Exception as e: + raise HTTPException(status_code=500, detail=f"获取状态失败: {str(e)}") + + +@web_app.get("/api/metrics") +async def get_metrics(): + """ + 获取队列状态概览。 + """ + try: + pending_count = len(queue_manager.pending_tasks) + running_count = len(queue_manager.running_tasks) + + return { + "pending_tasks": pending_count, + "running_tasks": running_count, + "total_servers": len(queue_manager.running_tasks), + "queue_manager_status": "active", + } + except Exception as e: + raise HTTPException(status_code=500, detail=f"获取队列状态失败: {str(e)}") + + +# 同步工作流执行API（带S3上传功能） +# @web_app.post("/api/run_sync/") +# async def execute_workflow_sync_endpoint( +# base_name: str, request_data_raw: Dict[str, Any], version: Optional[str] = None +# ): +# """ +# 同步执行工作流（支持S3文件上传） +# """ +# cleanup_paths = [] +# try: +# # 1. 获取工作流定义 +# workflow_data = await database.get_workflow(base_name, version) + +# if not workflow_data: +# detail = ( +# f"工作流 '{base_name}'" +# + (f" 带版本 '{version}'" if version else " (最新版)") +# + " 未找到。" +# ) +# raise HTTPException(status_code=404, detail=detail) + +# workflow = json.loads(workflow_data["workflow_json"]) +# api_spec = comfyui_client.parse_api_spec(workflow) +# request_data = {k.lower(): v for k, v in request_data_raw.items()} + +# # 2. [核心改动] 使用智能调度选择服务器 +# try: +# selected_server = await comfyui_client.select_server_for_execution() +# except Exception as e: +# raise HTTPException( +# status_code=503, detail=f"无法选择ComfyUI服务器执行任务: {e}" +# ) + +# server_input_dir = selected_server.input_dir +# server_output_dir = selected_server.output_dir + +# # 3. 下载文件到选定服务器的输入目录 +# async with aiohttp.ClientSession() as session: +# for param_name, param_def in api_spec["inputs"].items(): +# if param_def["type"] == "UploadFile" and param_name in request_data: +# image_url = request_data[param_name] +# if not isinstance(image_url, str) or not image_url.startswith( +# "http" +# ): +# raise HTTPException( +# status_code=400, +# detail=f"参数 '{param_name}' 必须是一个有效的URL。", +# ) + +# filename = f"api_download_{uuid.uuid4().hex}{os.path.splitext(image_url.split('/')[-1])[1] or '.dat'}" +# save_path = os.path.join(server_input_dir, filename) +# try: +# await download_file_from_url(session, image_url, save_path) +# # 直接更新 request_data，以便后续的 patch_workflow 使用 +# request_data[param_name] = filename +# cleanup_paths.append(save_path) +# except Exception as e: +# raise HTTPException( +# status_code=500, +# detail=f"为 '{param_name}' 从 {image_url} 下载文件失败。错误: {e}", +# ) + +# # 4. Patch工作流，生成最终的API Prompt +# patched_workflow = comfyui_client.patch_workflow( +# workflow, api_spec, request_data +# ) +# prompt_to_run = comfyui_client.convert_workflow_to_prompt_api_format( +# patched_workflow +# ) + +# # 5. 执行前快照，并在选定服务器上执行工作流 +# files_before = get_files_in_dir(server_output_dir) +# output_nodes = await comfyui_client.execute_prompt_on_server_legacy( +# prompt_to_run, selected_server +# ) + +# # 6. 处理输出（与之前逻辑相同） +# files_after = get_files_in_dir(server_output_dir) +# new_files = files_after - files_before + +# output_response = {} +# if new_files: +# s3_urls = [] +# for file_path in new_files: +# cleanup_paths.append(file_path) +# try: +# s3_urls.append(await handle_file_upload(file_path, base_name)) +# except Exception as e: +# print(f"上传文件 {file_path} 到S3时出错: {e}") +# if s3_urls: +# output_response["output_files"] = s3_urls + +# for final_param_name, param_def in api_spec["outputs"].items(): +# node_id = param_def["node_id"] +# if node_id in output_nodes: +# node_output = output_nodes[node_id] +# original_output_name = param_def["output_name"] + +# output_value = None +# if original_output_name in node_output: +# output_value = node_output[original_output_name] +# elif "text" in node_output: # 备选，例如对于'ShowText'节点 +# output_value = node_output["text"] + +# if output_value is None: +# continue + +# if isinstance(output_value, list): +# output_value = output_value[0] if output_value else None + +# output_response[final_param_name] = output_value + +# return JSONResponse(content=output_response) +# # [核心改动] 捕获来自ComfyUI的执行失败异常 +# except ComfyUIExecutionError as e: +# print(f"捕获到ComfyUI执行错误: {e.error_data}") +# # 返回 502 Bad Gateway 状态码，表示上游服务器出错 +# # detail 中包含结构化的错误信息，方便客户端处理 +# raise HTTPException( +# status_code=500, +# detail={ +# "message": "工作流在上游ComfyUI节点中执行失败。", +# "error_details": e.error_data, +# }, +# ) +# except Exception as e: +# # 捕获其他所有异常，作为通用的服务器内部错误 +# print(f"执行工作流时发生未知错误: {e}") +# raise HTTPException(status_code=500, detail=str(e)) +# finally: +# if cleanup_paths: +# print(f"正在清理 {len(cleanup_paths)} 个临时文件...") +# for path in cleanup_paths: +# try: +# if os.path.exists(path): +# os.remove(path) +# print(f" - 已删除: {path}") +# except Exception as e: +# print(f" - 删除 {path} 时出错: {e}") @web_app.get( @@ -449,6 +497,27 @@ def read_root(): }, { "step": 2, + "action": "异步执行工作流 (Execute Workflow Async with S3 Upload)", + "description": "提交工作流到异步队列，立即返回任务ID，支持长时间运行的工作流。执行完成后自动上传输出文件到S3。", + "endpoint": "/api/run", + "example_curl": 'curl -X POST "http://127.0.0.1:18000/api/run?workflow_name=my_t2i&workflow_version=20250801120000" -H \'Content-Type: application/json\' -d \'{\n "prompt_prompt": "a beautiful cat sitting on a roof",\n "sampler_seed": 12345\n}\'', + }, + { + "step": 3, + "action": "查询异步任务状态 (Check Async Task Status)", + "description": "通过任务ID查询异步工作流的执行状态。", + "endpoint": "/api/run/{workflow_run_id}", + "example_curl": "curl -X GET 'http://127.0.0.1:18000/api/run/your-task-id-here'", + }, + { + "step": 4, + "action": "查看队列状态 (Check Queue Status)", + "description": "查看当前异步任务队列的状态。", + "endpoint": "/api/metrics", + "example_curl": "curl -X GET 'http://127.0.0.1:18000/api/metrics'", + }, + { + "step": 5, "action": "查看工作流的API规范 (Inspect Workflow API Spec)", "description": "一旦发布，你可以查询工作流的API规范，以了解需要提供哪些输入参数以及可以期望哪些输出。", "endpoint": "/api/spec/", @@ -465,10 +534,10 @@ def read_root(): "example_curl": "curl -X GET 'http://127.0.0.1:18000/api/spec/?base_name=my_t2i'", }, { - "step": 3, - "action": "执行工作流 (Execute a Workflow)", - "description": "使用获取到的API规范，通过HTTP POST请求来执行工作流。对于文件输入，请提供可公开访问的URL。输出文件将被上传到S3并返回URL。", - "endpoint": "/api/run/", + "step": 6, + "action": "执行工作流 (Execute a Workflow - Sync with S3 Upload)", + "description": "使用获取到的API规范，通过HTTP POST请求来同步执行工作流。对于文件输入，请提供可公开访问的URL。输出文件将被自动上传到S3并返回URL。", + "endpoint": "/api/run_sync/", "parameters": [ {"name": "base_name", "description": "工作流的基础名称。"}, { @@ -476,10 +545,10 @@ def read_root(): "description": "可选。要执行的特定版本。如果省略，则执行最新版本。", }, ], - "example_curl": "curl -X POST 'http://127.0.0.1:18000/api/run/?base_name=my_t2i' -H 'Content-Type: application/json' -d '{\n \"prompt_prompt\": \"a beautiful cat sitting on a roof\",\n \"sampler_seed\": 12345\n}'", + "example_curl": "curl -X POST 'http://127.0.0.1:18000/api/run_sync/?base_name=my_t2i' -H 'Content-Type: application/json' -d '{\n \"prompt_prompt\": \"a beautiful cat sitting on a roof\",\n \"sampler_seed\": 12345\n}'", }, { - "step": 4, + "step": 7, "action": "管理工作流 (Manage Workflows)", "description": "你也可以列出所有已发布的工作流或删除不再需要的工作流。", "endpoints": [ @@ -503,6 +572,33 @@ def read_root(): return JSONResponse(content=guide) +def get_files_in_dir(directory: str) -> Set[str]: + file_set = set() + for root, _, files in os.walk(directory): + for filename in files: + if not filename.startswith("."): + file_set.add(os.path.join(root, filename)) + return file_set + + +async def download_file_from_url( + session: aiohttp.ClientSession, url: str, save_path: str +): + async with session.get(url) as response: + response.raise_for_status() + with open(save_path, "wb") as f: + while True: + chunk = await response.content.read(8192) + if not chunk: + break + f.write(chunk) + + +async def handle_file_upload(file_path: str, base_name: str) -> str: + s3_object_name = f"outputs/{base_name}/{uuid.uuid4()}_{os.path.basename(file_path)}" + return await upload_file_to_s3(file_path, settings.S3_BUCKET_NAME, s3_object_name) + + if __name__ == "__main__": uvicorn.run( "workflow_service.main:web_app", host="0.0.0.0", port=18000, reload=True