refactor: 重构WorkflowQueueManager和相关API，移除冗余代码，简化任务处理逻辑，增强服务器状态检查功能
This commit is contained in:
iHeyTang
@@ -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
|
||||
|
||||
345
workflow_service/comfy/comfy_run.py
Normal file
345
workflow_service/comfy/comfy_run.py
Normal file
@@ -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,
|
||||
}
|
||||
@@ -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)}")
|
||||
|
||||
Reference in New Issue
Block a user