272 lines
11 KiB
Python
272 lines
11 KiB
Python
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)
|