From 5fbbb719f42bf282501385eafd2127e053c302a6 Mon Sep 17 00:00:00 2001 From: iHeyTang Date: Fri, 15 Aug 2025 11:04:47 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E6=B7=BB=E5=8A=A0ComfyUI=E6=9C=8D?= =?UTF-8?q?=E5=8A=A1=E5=99=A8=E6=8C=81=E4=B9=85=E5=8C=96=E5=8A=9F=E8=83=BD?= =?UTF-8?q?=EF=BC=8C=E6=94=AF=E6=8C=81=E6=9C=8D=E5=8A=A1=E5=99=A8=E7=8A=B6?= =?UTF-8?q?=E6=80=81=E8=87=AA=E5=8A=A8=E4=BF=9D=E5=AD=98=E4=B8=8E=E6=81=A2?= =?UTF-8?q?=E5=A4=8D=EF=BC=8C=E5=A2=9E=E5=BC=BA=E7=B3=BB=E7=BB=9F=E7=9A=84?= =?UTF-8?q?=E5=8F=AF=E9=9D=A0=E6=80=A7=E5=92=8C=E5=8F=AF=E7=BB=B4=E6=8A=A4?= =?UTF-8?q?=E6=80=A7?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- README.md | 68 ++++++ workflow_service/comfy/comfy_server.py | 279 ++++++++++++++++++++----- workflow_service/database/__init__.py | 3 +- workflow_service/database/api.py | 2 +- workflow_service/database/models.py | 37 +++- 5 files changed, 333 insertions(+), 56 deletions(-) diff --git a/README.md b/README.md index 5d3fc3f..870a382 100644 --- a/README.md +++ b/README.md @@ -7,6 +7,7 @@ - 工作流上传 - 工作流版本控制 - 工作流加载 +- ComfyUI 服务器管理（支持持久化） - 需配置工作流服务器 (详见**WAAS(工作流即服务) Demo API 服务器**) # WAAS(工作流即服务) Demo API 服务器 @@ -42,6 +43,21 @@ version TEXT NOT NULL, workflow_json TEXT NOT NULL ) + + CREATE TABLE IF NOT EXISTS comfyui_servers ( + name TEXT PRIMARY KEY, + http_url TEXT NOT NULL, + ws_url TEXT NOT NULL, + status TEXT NOT NULL DEFAULT 'offline', + last_heartbeat DATETIME, + last_health_check DATETIME, + current_tasks INTEGER DEFAULT 0, + max_concurrent_tasks INTEGER DEFAULT 1, + capabilities TEXT, + server_metadata TEXT, + created_at DATETIME DEFAULT CURRENT_TIMESTAMP, + updated_at DATETIME DEFAULT CURRENT_TIMESTAMP + ) ``` - 路由 @@ -100,6 +116,58 @@ } ``` +# ComfyUI 服务器持久化功能 + +## 概述 + +ComfyUI 服务器管理器现在支持持久化功能，这意味着即使管理服务重启，也能从数据库中恢复已注册的服务器及其状态信息。 + +## 主要特性 + +- **自动持久化**: 服务器注册、状态变化、注销时自动保存到数据库 +- **自动恢复**: 服务启动时自动从数据库加载已注册的服务器 +- **健康检查**: 定期检查服务器健康状态并更新数据库 +- **心跳机制**: 支持服务器主动发送心跳，自动检测超时 + +## 使用方法 + +### 1. 数据库迁移 + +首次使用需要运行迁移脚本： + +```bash +python migrate_database.py +``` + +### 2. 测试功能 + +运行测试脚本验证功能： + +```bash +python test_persistence.py +``` + +### 3. API 使用 + +```python +from workflow_service.comfy.comfy_server import server_manager + +# 注册服务器 +await server_manager.register_server( + name="my_server", + http_url="http://localhost:8188", + ws_url="ws://localhost:8188", + max_concurrent_tasks=2 +) + +# 获取服务器状态 +servers = await server_manager.get_all_servers() +``` + +详细文档请参考 [PERSISTENCE_README.md](PERSISTENCE_README.md) + +--- + # JS 脚本使用指南 ### MacOS diff --git a/workflow_service/comfy/comfy_server.py b/workflow_service/comfy/comfy_server.py index 9f48cea..a7e7886 100644 --- a/workflow_service/comfy/comfy_server.py +++ b/workflow_service/comfy/comfy_server.py @@ -9,12 +9,19 @@ from enum import Enum import aiohttp from aiohttp import ClientTimeout +from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy import select, update, delete +from sqlalchemy.orm import selectinload + +from workflow_service.database.connection import AsyncSessionLocal +from workflow_service.database.models import ComfyUIServer as ComfyUIServerModel logger = logging.getLogger(__name__) class ServerStatus(Enum): """服务器状态枚举""" + ONLINE = "online" OFFLINE = "offline" BUSY = "busy" @@ -24,6 +31,7 @@ class ServerStatus(Enum): @dataclass class ComfyUIServerInfo: """ComfyUI服务器信息""" + name: str # 服务器名称，由注册方自行拟定 http_url: str ws_url: str @@ -34,26 +42,141 @@ class ComfyUIServerInfo: 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 = {} + @classmethod + def from_model(cls, model: ComfyUIServerModel) -> "ComfyUIServerInfo": + """从数据库模型创建实例""" + capabilities = json.loads(model.capabilities) if model.capabilities else {} + metadata = json.loads(model.server_metadata) if model.server_metadata else {} + + return cls( + name=model.name, + http_url=model.http_url, + ws_url=model.ws_url, + status=ServerStatus(model.status), + last_heartbeat=model.last_heartbeat, + last_health_check=model.last_health_check, + current_tasks=model.current_tasks, + max_concurrent_tasks=model.max_concurrent_tasks, + capabilities=capabilities, + metadata=metadata, + ) + + def to_model(self) -> ComfyUIServerModel: + """转换为数据库模型""" + return ComfyUIServerModel( + name=self.name, + http_url=self.http_url, + ws_url=self.ws_url, + status=self.status.value, + last_heartbeat=self.last_heartbeat, + last_health_check=self.last_health_check, + current_tasks=self.current_tasks, + max_concurrent_tasks=self.max_concurrent_tasks, + capabilities=json.dumps(self.capabilities), + server_metadata=json.dumps(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() - + self._initialized = False + + async def initialize(self): + """初始化管理器，从数据库加载服务器信息""" + if self._initialized: + return + + async with self.lock: + try: + await self._load_servers_from_db() + self._initialized = True + logger.info(f"从数据库加载了 {len(self.servers)} 个服务器") + except Exception as e: + logger.error(f"从数据库加载服务器失败: {e}") + self._initialized = True # 即使失败也标记为已初始化，避免重复尝试 + + async def _load_servers_from_db(self): + """从数据库加载服务器信息""" + async with AsyncSessionLocal() as session: + async with session.begin(): + result = await session.execute(select(ComfyUIServerModel)) + models = result.scalars().all() + + for model in models: + try: + server_info = ComfyUIServerInfo.from_model(model) + self.servers[server_info.name] = server_info + except Exception as e: + logger.error(f"加载服务器 {model.name} 失败: {e}") + + async def _save_server_to_db(self, server: ComfyUIServerInfo): + """保存服务器信息到数据库""" + try: + async with AsyncSessionLocal() as session: + async with session.begin(): + # 检查是否已存在 + existing = await session.execute( + select(ComfyUIServerModel).where( + ComfyUIServerModel.name == server.name + ) + ) + existing_model = existing.scalar_one_or_none() + + if existing_model: + # 更新现有记录 + await session.execute( + update(ComfyUIServerModel) + .where(ComfyUIServerModel.name == server.name) + .values( + http_url=server.http_url, + ws_url=server.ws_url, + status=server.status.value, + last_heartbeat=server.last_heartbeat, + last_health_check=server.last_health_check, + current_tasks=server.current_tasks, + max_concurrent_tasks=server.max_concurrent_tasks, + capabilities=json.dumps(server.capabilities), + server_metadata=json.dumps(server.metadata), + updated_at=datetime.utcnow(), + ) + ) + else: + # 创建新记录 + new_model = server.to_model() + session.add(new_model) + + await session.commit() + except Exception as e: + logger.error(f"保存服务器 {server.name} 到数据库失败: {e}") + + async def _delete_server_from_db(self, name: str): + """从数据库删除服务器""" + try: + async with AsyncSessionLocal() as session: + async with session.begin(): + await session.execute( + delete(ComfyUIServerModel).where( + ComfyUIServerModel.name == name + ) + ) + await session.commit() + except Exception as e: + logger.error(f"从数据库删除服务器 {name} 失败: {e}") + async def register_server( self, name: str, @@ -61,92 +184,116 @@ class ComfyUIServerManager: ws_url: str, max_concurrent_tasks: int = 1, capabilities: Optional[Dict[str, any]] = None, - metadata: Optional[Dict[str, any]] = None + metadata: Optional[Dict[str, any]] = None, ) -> bool: """注册新的ComfyUI服务器""" + # 确保已初始化 + await self.initialize() + # 确保健康检查任务已启动 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, # 注册时状态为离线，等待健康检查 + 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 {} + 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 + + # 保存到数据库 + await self._save_server_to_db(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] + # 从数据库删除 + await self._delete_server_from_db(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 + # 保存到数据库 + await self._save_server_to_db(self.servers[name]) return True return False - - async def get_available_server(self, required_capabilities: Optional[Dict[str, any]] = None) -> Optional[ComfyUIServerInfo]: + + async def get_available_server( + self, required_capabilities: Optional[Dict[str, any]] = None + ) -> Optional[ComfyUIServerInfo]: """获取可用的服务器（负载均衡）""" + # 确保已初始化 + await self.initialize() + # 确保健康检查任务已启动 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 ( + server.status == ServerStatus.ONLINE + and server.current_tasks < server.max_concurrent_tasks + ): + # 检查能力要求 if required_capabilities: - if not self._check_capabilities(server.capabilities, 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): + 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 + # 保存到数据库 + await self._save_server_to_db(server) return True return False - + async def release_server(self, name: str) -> bool: """释放服务器资源""" async with self.lock: @@ -154,22 +301,35 @@ class ComfyUIServerManager: 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: + if ( + server.status == ServerStatus.BUSY + and server.current_tasks < server.max_concurrent_tasks + ): server.status = ServerStatus.ONLINE + # 保存到数据库 + await self._save_server_to_db(server) return True return False - + async def get_server_status(self, name: str) -> Optional[ComfyUIServerInfo]: """获取服务器状态""" + # 确保已初始化 + await self.initialize() + async with self.lock: return self.servers.get(name) - + async def get_all_servers(self) -> List[ComfyUIServerInfo]: """获取所有服务器信息""" + # 确保已初始化 + await self.initialize() + async with self.lock: return list(self.servers.values()) - - def _check_capabilities(self, server_capabilities: Dict[str, any], required_capabilities: Dict[str, any]) -> bool: + + 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: @@ -180,7 +340,7 @@ class ComfyUIServerManager: elif server_capabilities[key] != value: return False return True - + async def _check_server_health(self, server: ComfyUIServerInfo) -> bool: """检查服务器健康状态""" try: @@ -191,9 +351,9 @@ class ComfyUIServerManager: return True except Exception as e: logger.debug(f"服务器 {server.name} 健康检查失败: {e}") - + return False - + async def _health_check_loop(self): """健康检查循环""" while True: @@ -201,35 +361,46 @@ class ComfyUIServerManager: 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)): + if ( + server.last_heartbeat + and current_time - server.last_heartbeat + > timedelta(seconds=self.heartbeat_timeout) + ): server.status = ServerStatus.OFFLINE logger.warning(f"服务器 {server.name} 心跳超时，标记为离线") - + # 保存到数据库 + await self._save_server_to_db(server) + # 定期健康检查 - if (not server.last_health_check or - current_time - server.last_health_check > timedelta(seconds=self.health_check_interval)): - + 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} 健康检查失败") - + + # 保存到数据库 + await self._save_server_to_db(server) + def _start_health_check(self): """启动健康检查任务""" try: @@ -240,12 +411,12 @@ class ComfyUIServerManager: 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(): @@ -266,6 +437,8 @@ 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]: +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/database/__init__.py b/workflow_service/database/__init__.py index 2378848..4524d35 100644 --- a/workflow_service/database/__init__.py +++ b/workflow_service/database/__init__.py @@ -3,7 +3,7 @@ 提供工作流和工作流运行的数据库操作接口 """ -from .models import Workflow, WorkflowRun, WorkflowRunNode +from .models import Workflow, WorkflowRun, WorkflowRunNode, ComfyUIServer from .api import ( init_db, save_workflow, @@ -27,6 +27,7 @@ __all__ = [ "Workflow", "WorkflowRun", "WorkflowRunNode", + "ComfyUIServer", # API函数 "init_db", diff --git a/workflow_service/database/api.py b/workflow_service/database/api.py index 917dda8..13c3520 100644 --- a/workflow_service/database/api.py +++ b/workflow_service/database/api.py @@ -11,7 +11,7 @@ import uuid from sqlalchemy.future import select from sqlalchemy import delete -from .models import Base, Workflow, WorkflowRun, WorkflowRunNode +from .models import Base, Workflow, WorkflowRun, WorkflowRunNode, ComfyUIServer from .connection import async_engine, AsyncSessionLocal diff --git a/workflow_service/database/models.py b/workflow_service/database/models.py index 404e61c..2665bcb 100644 --- a/workflow_service/database/models.py +++ b/workflow_service/database/models.py @@ -4,7 +4,7 @@ """ from datetime import datetime -from sqlalchemy import Column, String, Text, DateTime, ForeignKey +from sqlalchemy import Column, String, Text, DateTime, ForeignKey, Integer, Boolean from sqlalchemy.ext.declarative import declarative_base from sqlalchemy.orm import relationship @@ -109,3 +109,38 @@ class WorkflowRunNode(Base): "output_data": self.output_data, "error_message": self.error_message, } + + +class ComfyUIServer(Base): + """ComfyUI 服务器模型""" + + __tablename__ = "comfyui_servers" + + name = Column(String, primary_key=True) + http_url = Column(String, nullable=False) + ws_url = Column(String, nullable=False) + status = Column(String, nullable=False, default="offline") + last_heartbeat = Column(DateTime) + last_health_check = Column(DateTime) + current_tasks = Column(Integer, default=0) + max_concurrent_tasks = Column(Integer, default=1) + capabilities = Column(Text) # JSON 字符串 + server_metadata = Column(Text) # JSON 字符串 + created_at = Column(DateTime, default=datetime.utcnow) + updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow) + + def to_dict(self) -> dict: + return { + "name": self.name, + "http_url": self.http_url, + "ws_url": self.ws_url, + "status": self.status, + "last_heartbeat": self.last_heartbeat.isoformat() if self.last_heartbeat else None, + "last_health_check": self.last_health_check.isoformat() if self.last_health_check else None, + "current_tasks": self.current_tasks, + "max_concurrent_tasks": self.max_concurrent_tasks, + "capabilities": self.capabilities, + "server_metadata": self.server_metadata, + "created_at": self.created_at.isoformat() if self.created_at else None, + "updated_at": self.updated_at.isoformat() if self.updated_at else None, + }