diff --git a/.env b/.env index 6a9eb05..7a332d3 100644 --- a/.env +++ b/.env @@ -22,7 +22,7 @@ # "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"}]' +# 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/js/publisher.js b/js/publisher.js index 25e1161..547c3ab 100644 --- a/js/publisher.js +++ b/js/publisher.js @@ -2,12 +2,39 @@ import { app } from "../../scripts/app.js"; import { api } from "../../scripts/api.js"; // 本地存储的 key -const SETTINGS_STORAGE_KEY = "comfyui_publisher_settings"; +const COMFY_SERVER_REGISTRATION_KEY = "comfyui_server_registration"; + +// ComfyUI服务器注册配置 +const DEFAULT_SERVER_CONFIG = { + name: "comfyui_node_1", + http_url: "http://127.0.0.1:8188", + workflow_service_host: "http://localhost:8000", +}; + +// 获取ComfyUI服务器注册配置 +function getComfyServerConfig() { + try { + const stored = localStorage.getItem(COMFY_SERVER_REGISTRATION_KEY); + return stored ? JSON.parse(stored) : { ...DEFAULT_SERVER_CONFIG }; + } catch (error) { + console.error("读取ComfyUI服务器配置失败:", error); + return { ...DEFAULT_SERVER_CONFIG }; + } +} + +// 保存ComfyUI服务器注册配置 +function saveComfyServerConfig(config) { + try { + localStorage.setItem(COMFY_SERVER_REGISTRATION_KEY, JSON.stringify(config)); + } catch (error) { + console.error("保存ComfyUI服务器配置失败:", error); + } +} // 获取本地设置 function getLocalSettings() { try { - const stored = localStorage.getItem(SETTINGS_STORAGE_KEY); + const stored = localStorage.getItem("workflow_publisher_settings"); return stored ? JSON.parse(stored) : { host: "" }; } catch (error) { console.error("读取本地设置失败:", error); @@ -15,32 +42,247 @@ function getLocalSettings() { } } -// 保存设置到本地 +// 保存本地设置 function saveLocalSettings(settings) { try { - localStorage.setItem(SETTINGS_STORAGE_KEY, JSON.stringify(settings)); + localStorage.setItem( + "workflow_publisher_settings", + JSON.stringify(settings) + ); } catch (error) { console.error("保存本地设置失败:", error); } } -// 自定义 fetchApi 函数，使用本地存储的 host +// ComfyUI服务器注册管理器 +class ComfyServerManager { + constructor() { + this.registrationInterval = null; + this.heartbeatInterval = null; + this.isRegistered = false; + this.serverConfig = getComfyServerConfig(); + } + + // 注册服务器到工作流服务 + async registerServer() { + if (!this.serverConfig.workflow_service_host) { + console.warn("未设置工作流服务地址，无法注册"); + return false; + } + + try { + const response = await fetch( + `${this.serverConfig.workflow_service_host}/api/comfy/register`, + { + method: "POST", + headers: { + "Content-Type": "application/json", + }, + body: JSON.stringify({ + name: this.serverConfig.name, + http_url: this.serverConfig.http_url, + ws_url: this.getWebSocketUrl(), + }), + } + ); + + if (response.ok) { + const result = await response.json(); + console.log("ComfyUI服务器注册成功:", result); + this.isRegistered = true; + this.startHeartbeat(); + return true; + } else { + const errorText = await response.text(); + console.error("ComfyUI服务器注册失败:", response.status, errorText); + return false; + } + } catch (error) { + console.error("注册ComfyUI服务器时发生错误:", error); + return false; + } + } + + // 注销服务器 + async unregisterServer() { + if (!this.serverConfig.workflow_service_host || !this.isRegistered) { + return false; + } + + try { + const response = await fetch( + `${this.serverConfig.workflow_service_host}/api/comfy/unregister`, + { + method: "POST", + headers: { + "Content-Type": "application/json", + }, + body: JSON.stringify({ name: this.serverConfig.name }), + } + ); + + if (response.ok) { + console.log("ComfyUI服务器注销成功"); + this.isRegistered = false; + this.stopHeartbeat(); + return true; + } else { + console.error("ComfyUI服务器注销失败:", response.status); + return false; + } + } catch (error) { + console.error("注销ComfyUI服务器时发生错误:", error); + return false; + } + } + + // 发送心跳 + async sendHeartbeat() { + if (!this.serverConfig.workflow_service_host || !this.isRegistered) { + return false; + } + + try { + const response = await fetch( + `${this.serverConfig.workflow_service_host}/api/comfy/heartbeat`, + { + method: "POST", + headers: { + "Content-Type": "application/json", + }, + body: JSON.stringify({ name: this.serverConfig.name }), + } + ); + + if (response.ok) { + console.debug("心跳发送成功"); + return true; + } else { + console.warn("心跳发送失败:", response.status); + // 如果心跳失败，尝试重新注册 + if (response.status === 404) { + console.log("服务器不存在，尝试重新注册..."); + this.isRegistered = false; + await this.registerServer(); + } + return false; + } + } catch (error) { + console.error("发送心跳时发生错误:", error); + // 如果网络错误，尝试重新注册 + this.isRegistered = false; + await this.registerServer(); + return false; + } + } + + // 启动心跳 + startHeartbeat() { + if (this.heartbeatInterval) { + clearInterval(this.heartbeatInterval); + } + + // 每30秒发送一次心跳 + this.heartbeatInterval = setInterval(() => { + this.sendHeartbeat(); + }, 30000); + + console.log("ComfyUI服务器心跳已启动"); + } + + // 停止心跳 + stopHeartbeat() { + if (this.heartbeatInterval) { + clearInterval(this.heartbeatInterval); + this.heartbeatInterval = null; + } + console.log("ComfyUI服务器心跳已停止"); + } + + // 更新服务器配置 + updateConfig(newConfig) { + this.serverConfig = { ...this.serverConfig, ...newConfig }; + saveComfyServerConfig(this.serverConfig); + + // 如果已注册，需要重新注册以更新配置 + if (this.isRegistered) { + this.unregisterServer().then(() => { + setTimeout(() => this.registerServer(), 1000); + }); + } + } + + // 设置工作流服务地址 + setWorkflowServiceHost(host) { + this.serverConfig.workflow_service_host = host; + saveComfyServerConfig(this.serverConfig); + } + + // 获取WebSocket URL（自动从HTTP URL生成） + getWebSocketUrl() { + if (!this.serverConfig.http_url) return ""; + const url = new URL(this.serverConfig.http_url); + // 根据协议自动选择ws或wss + const protocol = url.protocol === "https:" ? "wss:" : "ws:"; + return `${protocol}//${url.host}/ws`; + } + + // 获取服务器状态 + async getServerStatus() { + if (!this.serverConfig.workflow_service_host) { + return null; + } + + try { + const response = await fetch( + `${this.serverConfig.workflow_service_host}/api/comfy/status/${this.serverConfig.name}` + ); + if (response.ok) { + return await response.json(); + } + } catch (error) { + console.error("获取服务器状态失败:", error); + } + return null; + } + + // 获取系统健康状态 + async getSystemHealth() { + if (!this.serverConfig.workflow_service_host) { + return null; + } + + try { + const response = await fetch( + `${this.serverConfig.workflow_service_host}/api/comfy/health` + ); + if (response.ok) { + return await response.json(); + } + } catch (error) { + console.error("获取系统健康状态失败:", error); + } + return null; + } +} + +// 自定义 fetchApi 函数，使用ComfyUI服务器配置中的工作流服务地址 async function customFetchApi(endpoint, options = {}) { try { - // 从本地存储获取设置 - const settings = getLocalSettings(); + // 从ComfyUI服务器配置获取工作流服务地址 + const serverConfig = getComfyServerConfig(); - // 如果没有配置 host，回退到默认的 api.fetchApi - if (!settings.host) { - console.warn("未配置 host，使用默认的 ComfyUI API"); + // 如果没有配置工作流服务地址，回退到默认的 api.fetchApi + if (!serverConfig.workflow_service_host) { + console.warn("未配置工作流服务地址，使用默认的 ComfyUI API"); return await api.fetchApi(endpoint, options); } // 构建完整的 URL - const host = settings.host.replace(/\/$/, ""); // 移除末尾的斜杠 + const host = serverConfig.workflow_service_host.replace(/\/$/, ""); // 移除末尾的斜杠 const fullUrl = `${host}${endpoint}`; - console.log(`使用本地存储的 host 请求: ${fullUrl}`); + console.log(`使用工作流服务地址请求: ${fullUrl}`); // 使用 fetch 发送请求 const response = await fetch(fullUrl, { @@ -186,22 +428,18 @@ app.registerExtension({ const btnPublish = createButton(iconPublish, "发布当前工作流", { onClick: () => showPublishDialog(), }); - const btnSettings = createButton(`⚙️`, "发布设置", { - onClick: () => showSettingsDialog(), + + const btnComfyServer = createButton(`🖥️`, "ComfyUI服务器管理", { + onClick: () => showComfyServerDialog(), }); const separator = document.createElement("div"); separator.style.cssText = "border-left: 1px solid var(--border-color); height: 22px; margin: 0 4px;"; - container.append(btnGet, btnPublish, btnSettings, separator); + container.append(btnGet, btnPublish, btnComfyServer, separator); menu.insertBefore(container, menu.firstChild); }); - const settingsDialog = new Modal( - "workflow-publisher-settings-dialog", - "发布设置" - ); - const publishDialog = new Modal( "workflow-publisher-publish-dialog", "发布工作流" @@ -213,6 +451,12 @@ app.registerExtension({ true ); + const comfyServerDialog = new Modal( + "comfy-server-dialog", + "ComfyUI服务器管理", + true + ); + async function showWorkflowListDialog() { workflowListDialog.setContent( '

正在从服务器加载工作流...

' @@ -318,76 +562,8 @@ app.registerExtension({ } } - async function showSettingsDialog() { - try { - // 首先尝试从本地存储获取设置 - const settings = getLocalSettings(); - - settingsDialog.setContent(` - - - `); - settingsDialog.addButtons([ - { label: "取消", callback: () => settingsDialog.hide() }, - { label: "保存", callback: saveSettings }, - ]); - settingsDialog.show(); - } catch (error) { - console.error("加载设置失败:", error); - alert(`加载设置失败: ${error.message}`); - } - } - - async function saveSettings() { - const host = document.getElementById("publisher-host").value; - - // 构建新的设置对象 - const newSettings = { host: host }; - - try { - // 保存到本地存储 - saveLocalSettings(newSettings); - - alert("设置已保存到本地！"); - settingsDialog.hide(); - } catch (error) { - alert(`保存设置失败: ${error}`); - } - } - async function showPublishDialog() { try { - // 从本地存储获取设置 - let settings = getLocalSettings(); - - // 如果本地没有设置，尝试从服务器获取 - if (!settings.host) { - try { - const settingsResp = await customFetchApi("/publisher/settings", { - cache: "no-store", - }); - if (settingsResp.ok) { - const text = await settingsResp.text(); - try { - const serverSettings = text ? JSON.parse(text) : {}; - settings = { ...settings, ...serverSettings }; - // 保存到本地存储 - saveLocalSettings(settings); - } catch (parseError) { - console.error("JSON解析失败:", parseError, "原始响应:", text); - } - } - } catch (error) { - console.error("从服务器获取设置失败:", error); - } - } - - // 确保设置对象有必需的字段 - settings = { - host: settings.host || "", - ...settings, - }; - publishDialog.setContent( `

` ); @@ -397,8 +573,176 @@ app.registerExtension({ ]); publishDialog.show(); } catch (error) { - console.error("加载设置失败:", error); - alert(`加载设置失败: ${error.message}`); + console.error("显示发布对话框失败:", error); + alert(`显示对话框失败: ${error.message}`); + } + } + + // ComfyUI 服务器管理对话框 + async function showComfyServerDialog() { + try { + const serverConfig = getComfyServerConfig(); + const settings = getLocalSettings(); + + // 创建服务器配置表单 + const configForm = ` +
+
+ + +
+ +
+ + + WebSocket地址将自动生成为: ${serverConfig.http_url ? serverConfig.http_url.replace(/^https?:\/\//, 'ws://') + '/ws' : ''} +
+ + + +
+ + +
+ +
+ 服务器状态: 未连接 +
+ +
+ + + +
+
+ `; + + comfyServerDialog.setContent(configForm); + comfyServerDialog.addButtons([ + { label: "关闭", callback: () => comfyServerDialog.hide() }, + ]); + comfyServerDialog.show(); + + // 初始化 ComfyServerManager + const serverManager = new ComfyServerManager(); + if (settings.host) { + serverManager.setWorkflowServiceHost(settings.host); + } + + // 绑定事件 + const btnRegister = document.getElementById("btn-register"); + const btnUnregister = document.getElementById("btn-unregister"); + const btnRefreshStatus = document.getElementById("btn-refresh-status"); + const statusText = document.getElementById("status-text"); + + // 添加HTTP地址输入框的实时预览功能 + const httpUrlInput = document.getElementById("server-http-url"); + const wsUrlPreview = document.getElementById("ws-url-preview"); + + httpUrlInput.addEventListener("input", function() { + const httpUrl = this.value; + if (httpUrl) { + try { + const url = new URL(httpUrl); + // 根据协议自动选择ws或wss + const protocol = url.protocol === "https:" ? "wss:" : "ws:"; + const wsUrl = `${protocol}//${url.host}/ws`; + wsUrlPreview.textContent = wsUrl; + } catch (error) { + wsUrlPreview.textContent = "无效的URL格式"; + } + } else { + wsUrlPreview.textContent = ""; + } + }); + + // 注册服务器 + btnRegister.onclick = async () => { + const newConfig = { + name: document.getElementById("server-name").value, + http_url: document.getElementById("server-http-url").value, + }; + + const workflowHost = document.getElementById( + "workflow-service-host" + ).value; + if (workflowHost) { + serverManager.setWorkflowServiceHost(workflowHost); + // 保存到本地设置 + saveLocalSettings({ host: workflowHost }); + } + + serverManager.updateConfig(newConfig); + + btnRegister.disabled = true; + btnRegister.textContent = "注册中..."; + + const success = await serverManager.registerServer(); + if (success) { + statusText.textContent = "已注册"; + statusText.style.color = "var(--success-color)"; + btnRegister.textContent = "已注册"; + btnRegister.disabled = true; + btnUnregister.disabled = false; + } else { + statusText.textContent = "注册失败"; + statusText.style.color = "var(--error-color)"; + btnRegister.textContent = "注册服务器"; + btnRegister.disabled = false; + } + }; + + // 注销服务器 + btnUnregister.onclick = async () => { + btnUnregister.disabled = true; + btnUnregister.textContent = "注销中..."; + + const success = await serverManager.unregisterServer(); + if (success) { + statusText.textContent = "已注销"; + statusText.style.color = "orange"; + btnRegister.disabled = false; + btnRegister.textContent = "注册服务器"; + btnUnregister.disabled = true; + btnUnregister.textContent = "注销服务器"; + } else { + statusText.textContent = "注销失败"; + statusText.style.color = "var(--error-color)"; + btnUnregister.disabled = false; + btnUnregister.textContent = "注销服务器"; + } + }; + + // 刷新状态 + btnRefreshStatus.onclick = async () => { + if (settings.host) { + serverManager.setWorkflowServiceHost(settings.host); + const status = await serverManager.getServerStatus(); + if (status) { + statusText.textContent = `在线 - 最后心跳: ${new Date( + status.last_heartbeat + ).toLocaleString()}`; + statusText.style.color = "var(--success-color)"; + btnRegister.disabled = true; + btnRegister.textContent = "已注册"; + btnUnregister.disabled = false; + } else { + statusText.textContent = "未注册"; + statusText.style.color = "orange"; + btnRegister.disabled = false; + btnRegister.textContent = "注册服务器"; + btnUnregister.disabled = true; + } + } + }; + + // 初始状态检查 + if (settings.host) { + btnRefreshStatus.click(); + } + } catch (error) { + console.error("显示ComfyUI服务器管理对话框失败:", error); + alert(`显示对话框失败: ${error.message}`); } } }, diff --git a/workflow_service/comfy/comfy_queue.py b/workflow_service/comfy/comfy_queue.py index 71ee2c8..731c5ed 100644 --- a/workflow_service/comfy/comfy_queue.py +++ b/workflow_service/comfy/comfy_queue.py @@ -12,7 +12,8 @@ import aiohttp from aiohttp import ClientTimeout from workflow_service.comfy.comfy_workflow import build_prompt -from workflow_service.config import Settings, ComfyUIServer +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, @@ -105,6 +106,9 @@ class WorkflowQueueManager: workflow_run_id = self.pending_tasks.pop(0) server = available_servers[0] + # 分配服务器资源 + await server_manager.allocate_server(server.name) + # 标记任务为运行中 await update_workflow_run_status( workflow_run_id, "running", server.http_url @@ -112,28 +116,21 @@ class WorkflowQueueManager: self.running_tasks[server.http_url] = { "workflow_run_id": workflow_run_id, "started_at": datetime.now(), + "server_name": server.name, } # 启动任务执行 asyncio.create_task(self._execute_task(workflow_run_id, server)) - async def _get_available_servers(self) -> list[ComfyUIServer]: + async def _get_available_servers(self) -> list[ComfyUIServerInfo]: """获取可用的服务器""" - 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 self.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}") + # 使用新的服务器管理器获取可用服务器 + available_server = await server_manager.get_available_server() + if available_server: + return [available_server] + return [] - return available_servers - - async def _execute_task(self, workflow_run_id: str, server: ComfyUIServer): + async def _execute_task(self, workflow_run_id: str, server: ComfyUIServerInfo): """执行任务""" cleanup_paths = [] try: @@ -179,6 +176,9 @@ class WorkflowQueueManager: 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) # 继续处理队列 asyncio.create_task(self._process_queue()) @@ -211,7 +211,7 @@ class WorkflowQueueManager: return result async def get_server_status( - self, server: ComfyUIServer, session: aiohttp.ClientSession + self, server: ComfyUIServerInfo, session: aiohttp.ClientSession ) -> dict[str, Any]: """ 检查单个ComfyUI服务器的详细状态。 @@ -256,7 +256,7 @@ async def _execute_prompt_on_server( workflow_data: dict, api_spec: dict, request_data: dict, - server: ComfyUIServer, + server: ComfyUIServerInfo, workflow_run_id: str, ) -> dict: """ diff --git a/workflow_service/comfy/comfy_server.py b/workflow_service/comfy/comfy_server.py new file mode 100644 index 0000000..9f48cea --- /dev/null +++ b/workflow_service/comfy/comfy_server.py @@ -0,0 +1,271 @@ +import asyncio +import json +import logging +import time +from datetime import datetime, timedelta +from typing import Dict, List, Optional, Set +from dataclasses import dataclass, asdict +from enum import Enum + +import aiohttp +from aiohttp import ClientTimeout + +logger = logging.getLogger(__name__) + + +class ServerStatus(Enum): + """服务器状态枚举""" + ONLINE = "online" + OFFLINE = "offline" + BUSY = "busy" + ERROR = "error" + + +@dataclass +class ComfyUIServerInfo: + """ComfyUI服务器信息""" + name: str # 服务器名称，由注册方自行拟定 + http_url: str + ws_url: str + status: ServerStatus = ServerStatus.OFFLINE + last_heartbeat: Optional[datetime] = None + last_health_check: Optional[datetime] = None + current_tasks: int = 0 + max_concurrent_tasks: int = 1 + capabilities: Dict[str, any] = None # 服务器能力信息，如支持的模型等 + metadata: Dict[str, any] = None # 其他元数据 + + def __post_init__(self): + if self.capabilities is None: + self.capabilities = {} + if self.metadata is None: + self.metadata = {} + + +class ComfyUIServerManager: + """ComfyUI服务器管理器""" + + def __init__(self): + self.servers: Dict[str, ComfyUIServerInfo] = {} # name -> server_info + self.lock = asyncio.Lock() + self.health_check_interval = 30 # 健康检查间隔（秒） + self.heartbeat_timeout = 60 # 心跳超时时间（秒） + self._health_check_task: Optional[asyncio.Task] = None + # 延迟启动健康检查，避免在模块导入时启动异步任务 + # self._start_health_check() + + async def register_server( + self, + name: str, + http_url: str, + ws_url: str, + max_concurrent_tasks: int = 1, + capabilities: Optional[Dict[str, any]] = None, + metadata: Optional[Dict[str, any]] = None + ) -> bool: + """注册新的ComfyUI服务器""" + # 确保健康检查任务已启动 + await self._ensure_health_check_started() + + async with self.lock: + if name in self.servers: + logger.warning(f"服务器名称 {name} 已存在，将更新配置") + + server_info = ComfyUIServerInfo( + name=name, + http_url=http_url.rstrip('/'), + ws_url=ws_url.rstrip('/'), + status=ServerStatus.OFFLINE, # 注册时状态为离线，等待健康检查 + max_concurrent_tasks=max_concurrent_tasks, + capabilities=capabilities or {}, + metadata=metadata or {} + ) + + # 立即进行健康检查 + if await self._check_server_health(server_info): + server_info.status = ServerStatus.ONLINE + server_info.last_health_check = datetime.now() + server_info.last_heartbeat = datetime.now() + + self.servers[name] = server_info + logger.info(f"服务器 {name} 注册成功: {http_url}") + return True + + async def unregister_server(self, name: str) -> bool: + """注销ComfyUI服务器""" + async with self.lock: + if name in self.servers: + del self.servers[name] + logger.info(f"服务器 {name} 已注销") + return True + return False + + async def update_server_heartbeat(self, name: str) -> bool: + """更新服务器心跳""" + async with self.lock: + if name in self.servers: + self.servers[name].last_heartbeat = datetime.now() + self.servers[name].status = ServerStatus.ONLINE + return True + return False + + async def get_available_server(self, required_capabilities: Optional[Dict[str, any]] = None) -> Optional[ComfyUIServerInfo]: + """获取可用的服务器（负载均衡）""" + # 确保健康检查任务已启动 + await self._ensure_health_check_started() + + async with self.lock: + available_servers = [] + + for server in self.servers.values(): + if (server.status == ServerStatus.ONLINE and + server.current_tasks < server.max_concurrent_tasks): + + # 检查能力要求 + if required_capabilities: + if not self._check_capabilities(server.capabilities, required_capabilities): + continue + + available_servers.append(server) + + if not available_servers: + return None + + # 简单的负载均衡：选择当前任务数最少的服务器 + return min(available_servers, key=lambda s: s.current_tasks) + + async def allocate_server(self, name: str) -> bool: + """分配服务器资源""" + async with self.lock: + if name in self.servers: + server = self.servers[name] + if (server.status == ServerStatus.ONLINE and + server.current_tasks < server.max_concurrent_tasks): + server.current_tasks += 1 + if server.current_tasks >= server.max_concurrent_tasks: + server.status = ServerStatus.BUSY + return True + return False + + async def release_server(self, name: str) -> bool: + """释放服务器资源""" + async with self.lock: + if name in self.servers: + server = self.servers[name] + if server.current_tasks > 0: + server.current_tasks -= 1 + if server.status == ServerStatus.BUSY and server.current_tasks < server.max_concurrent_tasks: + server.status = ServerStatus.ONLINE + return True + return False + + async def get_server_status(self, name: str) -> Optional[ComfyUIServerInfo]: + """获取服务器状态""" + async with self.lock: + return self.servers.get(name) + + async def get_all_servers(self) -> List[ComfyUIServerInfo]: + """获取所有服务器信息""" + async with self.lock: + return list(self.servers.values()) + + def _check_capabilities(self, server_capabilities: Dict[str, any], required_capabilities: Dict[str, any]) -> bool: + """检查服务器是否满足能力要求""" + for key, value in required_capabilities.items(): + if key not in server_capabilities: + return False + if isinstance(value, dict) and isinstance(server_capabilities[key], dict): + if not self._check_capabilities(server_capabilities[key], value): + return False + elif server_capabilities[key] != value: + return False + return True + + async def _check_server_health(self, server: ComfyUIServerInfo) -> bool: + """检查服务器健康状态""" + try: + timeout = ClientTimeout(total=10) + async with aiohttp.ClientSession(timeout=timeout) as session: + async with session.get(f"{server.http_url}/system_stats") as response: + if response.status == 200: + return True + except Exception as e: + logger.debug(f"服务器 {server.name} 健康检查失败: {e}") + + return False + + async def _health_check_loop(self): + """健康检查循环""" + while True: + try: + await self._perform_health_checks() + except Exception as e: + logger.error(f"健康检查过程中发生错误: {e}") + + await asyncio.sleep(self.health_check_interval) + + async def _perform_health_checks(self): + """执行健康检查""" + current_time = datetime.now() + + async with self.lock: + for server in self.servers.values(): + # 检查心跳超时 + if (server.last_heartbeat and + current_time - server.last_heartbeat > timedelta(seconds=self.heartbeat_timeout)): + server.status = ServerStatus.OFFLINE + logger.warning(f"服务器 {server.name} 心跳超时，标记为离线") + + # 定期健康检查 + if (not server.last_health_check or + current_time - server.last_health_check > timedelta(seconds=self.health_check_interval)): + + is_healthy = await self._check_server_health(server) + server.last_health_check = current_time + + if is_healthy and server.status != ServerStatus.ONLINE: + server.status = ServerStatus.ONLINE + logger.info(f"服务器 {server.name} 恢复在线") + elif not is_healthy and server.status == ServerStatus.ONLINE: + server.status = ServerStatus.ERROR + logger.warning(f"服务器 {server.name} 健康检查失败") + + def _start_health_check(self): + """启动健康检查任务""" + try: + # 检查是否有运行中的事件循环 + loop = asyncio.get_running_loop() + if self._health_check_task is None or self._health_check_task.done(): + self._health_check_task = asyncio.create_task(self._health_check_loop()) + except RuntimeError: + # 没有运行中的事件循环，延迟启动 + logger.debug("没有运行中的事件循环，健康检查将在事件循环启动后自动开始") + + async def _ensure_health_check_started(self): + """确保健康检查任务已启动""" + if self._health_check_task is None or self._health_check_task.done(): + self._start_health_check() + + async def shutdown(self): + """关闭管理器""" + if self._health_check_task and not self._health_check_task.done(): + self._health_check_task.cancel() + try: + await self._health_check_task + except asyncio.CancelledError: + pass + + +# 全局服务器管理器实例 +server_manager = ComfyUIServerManager() + + +# 兼容性函数，用于替代原有的静态配置 +async def get_available_servers() -> List[ComfyUIServerInfo]: + """获取可用的服务器列表（兼容性函数）""" + return await server_manager.get_all_servers() + + +async def get_server_for_task(required_capabilities: Optional[Dict[str, any]] = None) -> Optional[ComfyUIServerInfo]: + """为任务获取合适的服务器（兼容性函数）""" + return await server_manager.get_available_server(required_capabilities) diff --git a/workflow_service/comfy/comfy_workflow.py b/workflow_service/comfy/comfy_workflow.py index 31810c7..0d4d664 100644 --- a/workflow_service/comfy/comfy_workflow.py +++ b/workflow_service/comfy/comfy_workflow.py @@ -2,7 +2,7 @@ import logging from typing import Any, Dict import aiohttp -from workflow_service.config import ComfyUIServer +from workflow_service.comfy.comfy_server import ComfyUIServerInfo logging.basicConfig(level=logging.INFO) @@ -17,7 +17,7 @@ async def build_prompt( workflow_data: dict, api_spec: dict[str, dict[str, Any]], request_data: dict[str, Any], - server: ComfyUIServer, + server: ComfyUIServerInfo, ): """ 构建prompt @@ -190,7 +190,7 @@ async def _patch_workflow( workflow_data: dict, api_spec: dict[str, dict[str, Any]], request_data: dict[str, Any], - server: ComfyUIServer, + server: ComfyUIServerInfo, ) -> dict: """ 将request_data中的参数值，patch到workflow_data中。并返回修改后的workflow_data。 @@ -326,7 +326,7 @@ def _convert_workflow_to_prompt_api_format(workflow_data: dict) -> dict: return prompt_api_format -async def _upload_image_to_comfy(file_path: str, server: ComfyUIServer) -> str: +async def _upload_image_to_comfy(file_path: str, server: ComfyUIServerInfo) -> str: """ 上传文件到服务器。 file 是一个http链接 diff --git a/workflow_service/config.py b/workflow_service/config.py index e01b425..ce66577 100644 --- a/workflow_service/config.py +++ b/workflow_service/config.py @@ -20,11 +20,6 @@ class Settings(BaseSettings): model_config = SettingsConfigDict(env_file=".env", env_file_encoding="utf-8") - # [核心改动] 使用JSON字符串来定义一个服务器列表 - COMFYUI_SERVERS_JSON: str = ( - '[{"http_url": "http://127.0.0.1:8188", "ws_url": "ws://127.0.0.1:8188/ws", "input_dir": "/comfyui/input", "output_dir": "/comfyui/output"}]' - ) - DB_FILE: str = "/db/workflows_service.sqlite" # S3 和 AWS 配置保持不变 S3_BUCKET_NAME: str @@ -32,21 +27,5 @@ class Settings(BaseSettings): AWS_SECRET_ACCESS_KEY: str AWS_REGION_NAME: str - @property - def SERVERS(self) -> List[ComfyUIServer]: - """ - 解析JSON配置字符串，返回一个经过验证的服务器配置对象列表。 - """ - try: - server_list = json.loads(self.COMFYUI_SERVERS_JSON) - # 使用Pydantic模型进行验证和类型转换 - return [ComfyUIServer(**server) for server in server_list] - except json.JSONDecodeError: - raise ValueError("COMFYUI_SERVERS_JSON 格式无效，必须是有效的JSON数组。") - except ValidationError as e: - raise ValueError(f"COMFYUI_SERVERS_JSON 中的服务器配置无效: {e}") - except Exception as e: - raise e - settings = Settings() diff --git a/workflow_service/main.py b/workflow_service/main.py index 87c8ee2..546c37a 100644 --- a/workflow_service/main.py +++ b/workflow_service/main.py @@ -4,9 +4,10 @@ import uvicorn from fastapi import FastAPI from fastapi.middleware.cors import CORSMiddleware +from workflow_service.comfy.comfy_server import server_manager from workflow_service.database.api import init_db from workflow_service.config import Settings -from workflow_service.routes import service, workflow, run, runx, monitor +from workflow_service.routes import service, workflow, run, runx, monitor, comfy_server settings = Settings() @@ -26,23 +27,20 @@ web_app.include_router(workflow.workflow_router) web_app.include_router(run.run_router) web_app.include_router(runx.runx_router) web_app.include_router(monitor.monitor_router) +web_app.include_router(comfy_server.router) @web_app.on_event("startup") async def startup_event(): - """服务启动时，初始化数据库并为所有配置的服务器创建输入/输出目录。""" + """服务启动时，初始化数据库。""" await init_db() try: - servers = settings.SERVERS + servers = await server_manager.get_all_servers() print(f"检测到 {len(servers)} 个 ComfyUI 服务器配置。") for server in servers: - print(f" - 正在为服务器 {server.http_url} 准备目录...") - os.makedirs(server.input_dir, exist_ok=True) - print(f" - 输入目录: {os.path.abspath(server.input_dir)}") - os.makedirs(server.output_dir, exist_ok=True) - print(f" - 输出目录: {os.path.abspath(server.output_dir)}") + print(f" - 服务器: {server.http_url}") except ValueError as e: - print(f"错误: 无法在启动时初始化服务器目录: {e}") + print(f"错误: 无法在启动时获取服务器信息: {e}") if __name__ == "__main__": diff --git a/workflow_service/routes/comfy_server.py b/workflow_service/routes/comfy_server.py new file mode 100644 index 0000000..2dbfc37 --- /dev/null +++ b/workflow_service/routes/comfy_server.py @@ -0,0 +1,224 @@ +from fastapi import APIRouter, HTTPException, BackgroundTasks +from pydantic import BaseModel, HttpUrl +from typing import Dict, List, Optional, Any +import logging + +from workflow_service.comfy.comfy_server import ( + server_manager, + ComfyUIServerInfo, + ServerStatus, +) + +logger = logging.getLogger(__name__) + +router = APIRouter(prefix="/api/comfy", tags=["ComfyUI Server Management"]) + + +class ServerRegistrationRequest(BaseModel): + """服务器注册请求模型""" + + name: str + http_url: str + ws_url: str + + +class ServerHeartbeatRequest(BaseModel): + """服务器心跳请求模型""" + + name: str + + +class ServerUnregisterRequest(BaseModel): + """服务器注销请求模型""" + + name: str + + +class ServerStatusResponse(BaseModel): + """服务器状态响应模型""" + + name: str + http_url: str + ws_url: str + status: str + last_heartbeat: Optional[str] = None + last_health_check: Optional[str] = None + current_tasks: int + max_concurrent_tasks: int + capabilities: Dict[str, Any] + metadata: Dict[str, Any] + + +@router.post("/register", response_model=Dict[str, str]) +async def register_server(request: ServerRegistrationRequest): + """注册ComfyUI服务器""" + try: + print(request) + success = await server_manager.register_server( + name=request.name, + http_url=request.http_url, + ws_url=request.ws_url, + max_concurrent_tasks=1, + capabilities={}, + metadata={}, + ) + + if success: + logger.info(f"服务器 {request.name} 注册成功") + return {"message": f"服务器 {request.name} 注册成功", "status": "success"} + else: + raise HTTPException(status_code=400, detail="服务器注册失败") + + except Exception as e: + logger.error(f"注册服务器 {request.name} 时发生错误: {e}") + raise HTTPException(status_code=500, detail=f"注册失败: {str(e)}") + + +@router.post("/heartbeat", response_model=Dict[str, str]) +async def update_heartbeat(request: ServerHeartbeatRequest): + """更新服务器心跳""" + try: + success = await server_manager.update_server_heartbeat(request.name) + + if success: + return { + "message": f"服务器 {request.name} 心跳更新成功", + "status": "success", + } + else: + raise HTTPException(status_code=404, detail=f"服务器 {request.name} 不存在") + + except Exception as e: + logger.error(f"更新服务器 {request.name} 心跳时发生错误: {e}") + raise HTTPException(status_code=500, detail=f"心跳更新失败: {str(e)}") + + +@router.post("/unregister", response_model=Dict[str, str]) +async def unregister_server(request: ServerUnregisterRequest): + """注销ComfyUI服务器""" + try: + success = await server_manager.unregister_server(request.name) + + if success: + logger.info(f"服务器 {request.name} 注销成功") + return {"message": f"服务器 {request.name} 注销成功", "status": "success"} + else: + raise HTTPException(status_code=404, detail=f"服务器 {request.name} 不存在") + + except Exception as e: + logger.error(f"注销服务器 {request.name} 时发生错误: {e}") + raise HTTPException(status_code=500, detail=f"注销失败: {str(e)}") + + +@router.get("/status/{server_name}", response_model=ServerStatusResponse) +async def get_server_status(server_name: str): + """获取指定服务器状态""" + try: + server_info = await server_manager.get_server_status(server_name) + + if server_info is None: + raise HTTPException(status_code=404, detail=f"服务器 {server_name} 不存在") + + return ServerStatusResponse( + name=server_info.name, + http_url=server_info.http_url, + ws_url=server_info.ws_url, + status=server_info.status.value, + last_heartbeat=( + server_info.last_heartbeat.isoformat() + if server_info.last_heartbeat + else None + ), + last_health_check=( + server_info.last_health_check.isoformat() + if server_info.last_health_check + else None + ), + current_tasks=server_info.current_tasks, + max_concurrent_tasks=server_info.max_concurrent_tasks, + capabilities=server_info.capabilities, + metadata=server_info.metadata, + ) + + except Exception as e: + logger.error(f"获取服务器 {server_name} 状态时发生错误: {e}") + raise HTTPException(status_code=500, detail=f"获取状态失败: {str(e)}") + + +@router.get("/list", response_model=List[ServerStatusResponse]) +async def list_all_servers(): + """获取所有服务器列表""" + try: + servers = await server_manager.get_all_servers() + + return [ + ServerStatusResponse( + name=server.name, + http_url=server.http_url, + ws_url=server.ws_url, + status=server.status.value, + last_heartbeat=( + server.last_heartbeat.isoformat() if server.last_heartbeat else None + ), + last_health_check=( + server.last_health_check.isoformat() + if server.last_health_check + else None + ), + current_tasks=server.current_tasks, + max_concurrent_tasks=server.max_concurrent_tasks, + capabilities=server.capabilities, + metadata=server.metadata, + ) + for server in servers + ] + + except Exception as e: + logger.error(f"获取服务器列表时发生错误: {e}") + raise HTTPException(status_code=500, detail=f"获取列表失败: {str(e)}") + + +@router.get("/health", response_model=Dict[str, Any]) +async def get_system_health(): + """获取系统整体健康状态""" + try: + servers = await server_manager.get_all_servers() + + total_servers = len(servers) + online_servers = len([s for s in servers if s.status == ServerStatus.ONLINE]) + busy_servers = len([s for s in servers if s.status == ServerStatus.BUSY]) + offline_servers = len([s for s in servers if s.status == ServerStatus.OFFLINE]) + error_servers = len([s for s in servers if s.status == ServerStatus.ERROR]) + + total_tasks = sum(s.current_tasks for s in servers) + total_capacity = sum(s.max_concurrent_tasks for s in servers) + + return { + "total_servers": total_servers, + "online_servers": online_servers, + "busy_servers": busy_servers, + "offline_servers": offline_servers, + "error_servers": error_servers, + "current_tasks": total_tasks, + "total_capacity": total_capacity, + "utilization_rate": ( + (total_tasks / total_capacity * 100) if total_capacity > 0 else 0 + ), + } + + except Exception as e: + logger.error(f"获取系统健康状态时发生错误: {e}") + raise HTTPException(status_code=500, detail=f"获取健康状态失败: {str(e)}") + + +@router.post("/force_health_check", response_model=Dict[str, str]) +async def force_health_check(background_tasks: BackgroundTasks): + """强制执行健康检查""" + try: + # 在后台执行健康检查 + background_tasks.add_task(server_manager._perform_health_checks) + return {"message": "健康检查已启动", "status": "success"} + + except Exception as e: + logger.error(f"启动健康检查时发生错误: {e}") + raise HTTPException(status_code=500, detail=f"启动健康检查失败: {str(e)}") diff --git a/workflow_service/routes/monitor.py b/workflow_service/routes/monitor.py index 81e291c..833feec 100644 --- a/workflow_service/routes/monitor.py +++ b/workflow_service/routes/monitor.py @@ -2,13 +2,18 @@ from fastapi import APIRouter, HTTPException from fastapi.responses import HTMLResponse import psutil import asyncio +import logging from datetime import datetime, timedelta from typing import Dict, Any, List import json +# 设置日志 +logger = logging.getLogger(__name__) + from workflow_service.config import settings from workflow_service.database.api import get_workflow_runs_recent, get_workflow_run_nodes from workflow_service.comfy.comfy_queue import WorkflowQueueManager +from workflow_service.comfy.comfy_server import server_manager monitor_router = APIRouter(prefix="/monitor", tags=["监控"]) @@ -223,12 +228,24 @@ async def monitor_dashboard(): const statusClass = server.status === 'online' ? 'online' : 'offline'; const statusColor = server.status === 'online' ? 'status-online' : 'status-offline'; + // 格式化时间 + const formatTime = (timeStr) => { + if (!timeStr) return 'N/A'; + try { + return new Date(timeStr).toLocaleString(); + } catch { + return timeStr; + } + }; + container.innerHTML += `

${server.http_url}

${server.status}

-

输入目录: ${server.input_dir}

-

输出目录: ${server.output_dir}

+

任务状态: ${server.current_tasks || 0}/${server.max_concurrent_tasks || 1}

+

最后心跳: ${formatTime(server.last_heartbeat)}

+

最后检查: ${formatTime(server.last_health_check)}

+ ${server.error ? `

错误: ${server.error}

` : ''}
`; }); @@ -324,6 +341,13 @@ async def get_server_status() -> List[Dict[str, Any]]: import aiohttp import asyncio + # 从 ComfyUIServerManager 获取所有注册的服务器 + all_servers = await server_manager.get_all_servers() + + if not all_servers: + logger.info("当前没有动态注册的服务器") + return [] + async def check_server_status(server): try: timeout = aiohttp.ClientTimeout(total=5) @@ -333,32 +357,38 @@ async def get_server_status() -> List[Dict[str, Any]]: return { "http_url": server.http_url, "ws_url": server.ws_url, - "input_dir": server.input_dir, - "output_dir": server.output_dir, "status": "online", - "response_time": response.headers.get('X-Response-Time', 'N/A') + "response_time": response.headers.get('X-Response-Time', 'N/A'), + "current_tasks": getattr(server, 'current_tasks', 0), + "max_concurrent_tasks": getattr(server, 'max_concurrent_tasks', 1), + "last_heartbeat": getattr(server, 'last_heartbeat', None), + "last_health_check": getattr(server, 'last_health_check', None) } else: return { "http_url": server.http_url, "ws_url": server.ws_url, - "input_dir": server.input_dir, - "output_dir": server.output_dir, "status": "offline", - "error": f"HTTP {response.status}" + "error": f"HTTP {response.status}", + "current_tasks": getattr(server, 'current_tasks', 0), + "max_concurrent_tasks": getattr(server, 'max_concurrent_tasks', 1), + "last_heartbeat": getattr(server, 'last_heartbeat', None), + "last_health_check": getattr(server, 'last_health_check', None) } except Exception as e: return { "http_url": server.http_url, "ws_url": server.ws_url, - "input_dir": server.input_dir, - "output_dir": server.output_dir, "status": "offline", - "error": str(e) + "error": str(e), + "current_tasks": getattr(server, 'current_tasks', 0), + "max_concurrent_tasks": getattr(server, 'max_concurrent_tasks', 1), + "last_heartbeat": getattr(server, 'last_heartbeat', None), + "last_health_check": getattr(server, 'last_health_check', None) } # 并发检查所有服务器状态 - tasks = [check_server_status(server) for server in settings.SERVERS] + tasks = [check_server_status(server) for server in all_servers] results = await asyncio.gather(*tasks, return_exceptions=True) # 过滤掉异常结果 @@ -409,15 +439,24 @@ async def health_check() -> Dict[str, Any]: from workflow_service.database.connection import get_db db = get_db() - # 检查配置 - servers_count = len(settings.SERVERS) + # 从 ComfyUIServerManager 获取服务器信息 + all_servers = await server_manager.get_all_servers() + dynamic_servers_count = len(all_servers) + + # 检查服务器健康状态 + online_servers = [s for s in all_servers if s.status == 'online'] + offline_servers = [s for s in all_servers if s.status != 'online'] return { "status": "healthy", "timestamp": datetime.now().isoformat(), "version": "1.0.0", "database": "connected", - "servers_configured": servers_count, + "servers": { + "total_dynamic": dynamic_servers_count, + "online": len(online_servers), + "offline": len(offline_servers) + }, "uptime": "N/A" # 可以添加启动时间跟踪 } except Exception as e: diff --git a/workflow_service/routes/service.py b/workflow_service/routes/service.py index 4284058..5ac8fd8 100644 --- a/workflow_service/routes/service.py +++ b/workflow_service/routes/service.py @@ -8,6 +8,7 @@ from fastapi import APIRouter, HTTPException, Path from pydantic import BaseModel from workflow_service.comfy.comfy_queue import queue_manager +from workflow_service.comfy.comfy_server import server_manager from workflow_service.config import Settings settings = Settings() @@ -28,8 +29,6 @@ class ServerStatus(BaseModel): server_index: int http_url: str ws_url: str - input_dir: str - output_dir: str is_reachable: bool is_free: bool queue_details: ServerQueueDetails @@ -72,7 +71,7 @@ async def get_servers_status(): """ 获取所有已配置的ComfyUI服务器的配置信息和实时状态。 """ - servers = settings.SERVERS + servers = await server_manager.get_all_servers() if not servers: return [] @@ -90,8 +89,6 @@ async def get_servers_status(): server_index=i, http_url=server_config.http_url, ws_url=server_config.ws_url, - input_dir=server_config.input_dir, - output_dir=server_config.output_dir, is_reachable=status_data["is_reachable"], is_free=status_data["is_free"], queue_details=status_data["queue_details"], @@ -106,8 +103,9 @@ async def list_server_files( ): """ 获取指定ComfyUI服务器的输入和输出文件夹中的文件列表。 + 注意：由于节点管理不再维护目录，此接口返回空列表。 """ - servers = settings.SERVERS + servers = await server_manager.get_all_servers() if server_index >= len(servers): raise HTTPException( status_code=404, @@ -116,10 +114,8 @@ async def list_server_files( server_config = servers[server_index] - input_files, output_files = await asyncio.gather( - _get_folder_contents(server_config.input_dir), - _get_folder_contents(server_config.output_dir), - ) + # 由于节点管理不再维护目录，返回空列表 + input_files, output_files = [], [] return ServerFiles( server_index=server_index,