feat: 添加队列监控功能,支持手动触发和监控间隔设置,增强任务处理能力

This commit is contained in:
iHeyTang
2025-08-18 10:59:01 +08:00
parent 3fa1c895be
commit ab4cbf7188
4 changed files with 189 additions and 11 deletions

View File

@@ -6,4 +6,5 @@ aiohttp
boto3
aiosqlite
sqlalchemy
greenlet
psutil

View File

@@ -6,7 +6,7 @@ import random
import uuid
import websockets
from datetime import datetime
from typing import Dict, Any
from typing import Dict, Any, Optional
import aiohttp
from aiohttp import ClientTimeout
@@ -47,10 +47,117 @@ class ComfyUIExecutionError(Exception):
# 全局任务队列管理器
class WorkflowQueueManager:
def __init__(self):
def __init__(self, monitor_interval: int = 5):
self.running_tasks = {} # server_url -> task_info
self.pending_tasks = [] # 等待队列
self.lock = asyncio.Lock()
self._queue_monitor_task: Optional[asyncio.Task] = None
self._monitor_interval = monitor_interval # 监控间隔(秒)
self._last_processing_time = None # 上次处理队列的时间
async def start_queue_monitor(self):
"""启动队列监控任务"""
if self._queue_monitor_task is None or self._queue_monitor_task.done():
self._queue_monitor_task = asyncio.create_task(self._monitor_queue())
logger.info(f"队列监控任务已启动,监控间隔: {self._monitor_interval}")
async def stop_queue_monitor(self):
"""停止队列监控任务"""
if self._queue_monitor_task and not self._queue_monitor_task.done():
self._queue_monitor_task.cancel()
try:
await self._queue_monitor_task
except asyncio.CancelledError:
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 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,
}
async def _monitor_queue(self):
"""定期监控队列,当有可用服务器时自动处理"""
while True:
try:
await asyncio.sleep(self._monitor_interval)
# 每次定时检测触发时,打印当前状态
async with self.lock:
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 = ""
if self._last_processing_time:
time_diff = current_time - self._last_processing_time
minutes = int(time_diff.total_seconds() // 60)
seconds = int(time_diff.total_seconds() % 60)
time_since_last_processing = f"上次处理: {minutes}{seconds}秒前"
else:
time_since_last_processing = "上次处理: 从未"
# 格式化状态信息,提高可读性
status_info = (
f"⏰ 定时检测触发 [{timestamp}]\n"
f" 📋 待处理任务: {pending_count}\n"
f" 🚀 运行中任务: {running_count}\n"
f" ⏱️ 监控间隔: {self._monitor_interval}\n"
f" 🕐 {time_since_last_processing}"
)
logger.info(status_info)
if self.pending_tasks:
available_servers = await self._get_available_servers()
if available_servers:
logger.info(
f"监控发现 {len(available_servers)} 个可用服务器,开始处理队列"
)
# 使用create_task避免阻塞监控循环
asyncio.create_task(self._process_queue())
else:
# 只在有任务等待时记录日志,避免日志噪音
if len(self.pending_tasks) > 0:
logger.debug(
f"队列中有 {len(self.pending_tasks)} 个待处理任务,但无可用服务器"
)
else:
# 队列为空时,减少日志输出
logger.debug("队列为空,继续监控中...")
except asyncio.CancelledError:
logger.info("队列监控任务被取消")
break
except Exception as e:
logger.error(f"队列监控任务出错: {e}")
# 出错时稍微延长等待时间,避免频繁重试
await asyncio.sleep(self._monitor_interval * 2)
async def add_task(
self,
@@ -85,30 +192,40 @@ class WorkflowQueueManager:
f"任务 {workflow_run_id} 已添加到队列,当前队列长度: {len(self.pending_tasks)}"
)
# 尝试处理队列
asyncio.create_task(self._process_queue())
# 注意:不再手动触发队列处理,由定时监控自动处理
# 这样可以避免在没有可用服务器时的不必要尝试
return workflow_run_id
async def _process_queue(self):
"""处理队列中的任务"""
async with self.lock:
# 更新上次处理队列的时间
self._last_processing_time = datetime.now()
if not self.pending_tasks:
logger.debug("队列为空,无需处理")
return
# 检查是否有空闲的服务器
available_servers = await self._get_available_servers()
if not available_servers:
logger.info("没有可用的服务器,等待中...")
logger.info(
f"队列中有 {len(self.pending_tasks)} 个待处理任务,但没有可用的服务器"
)
return
# 获取一个待处理任务
workflow_run_id = self.pending_tasks.pop(0)
server = available_servers[0]
logger.info(
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
@@ -122,6 +239,10 @@ class WorkflowQueueManager:
# 启动任务执行
asyncio.create_task(self._execute_task(workflow_run_id, server))
logger.info(
f"任务 {workflow_run_id} 已分配到服务器 {server.name},当前运行中任务数: {len(self.running_tasks)}"
)
async def _get_available_servers(self) -> list[ComfyUIServerInfo]:
"""获取可用的服务器"""
# 使用新的服务器管理器获取可用服务器
@@ -176,13 +297,10 @@ 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())
async def get_task_status(self, workflow_run_id: str) -> dict:
"""获取任务状态"""
workflow_run = await get_workflow_run(workflow_run_id)
@@ -249,7 +367,8 @@ class WorkflowQueueManager:
# 全局队列管理器实例
queue_manager = WorkflowQueueManager()
# 从配置中读取监控间隔默认为5秒
queue_manager = WorkflowQueueManager(monitor_interval=5)
async def _execute_prompt_on_server(

View File

@@ -8,6 +8,7 @@ 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, comfy_server
from workflow_service.comfy.comfy_queue import queue_manager
settings = Settings()
@@ -39,10 +40,22 @@ async def startup_event():
print(f"检测到 {len(servers)} 个 ComfyUI 服务器配置。")
for server in servers:
print(f" - 服务器: {server.http_url}")
# 启动队列监控任务
await queue_manager.start_queue_monitor()
print("队列监控任务已启动")
except ValueError as e:
print(f"错误: 无法在启动时获取服务器信息: {e}")
@web_app.on_event("shutdown")
async def shutdown_event():
"""服务关闭时,停止队列监控任务。"""
await queue_manager.stop_queue_monitor()
print("队列监控任务已停止")
if __name__ == "__main__":
uvicorn.run(
"workflow_service.main:web_app", host="0.0.0.0", port=18000, reload=True

View File

@@ -465,3 +465,48 @@ async def health_check() -> Dict[str, Any]:
"timestamp": datetime.now().isoformat(),
"error": str(e)
}
@monitor_router.get("/queue/status")
async def get_queue_status():
"""获取队列状态"""
try:
status = await queue_manager.get_queue_status()
return {
"success": True,
"data": status,
"timestamp": datetime.now().isoformat()
}
except Exception as e:
logger.error(f"获取队列状态失败: {e}")
raise HTTPException(status_code=500, detail=f"获取队列状态失败: {str(e)}")
@monitor_router.post("/queue/trigger")
async def trigger_queue_processing():
"""手动触发队列处理"""
try:
await queue_manager.trigger_queue_processing()
return {
"success": True,
"message": "队列处理已手动触发",
"timestamp": datetime.now().isoformat()
}
except Exception as e:
logger.error(f"手动触发队列处理失败: {e}")
raise HTTPException(status_code=500, detail=f"手动触发队列处理失败: {str(e)}")
@monitor_router.post("/queue/interval")
async def set_monitor_interval(interval: int):
"""设置队列监控间隔"""
try:
if interval < 1:
raise HTTPException(status_code=400, detail="监控间隔不能小于1秒")
await queue_manager.set_monitor_interval(interval)
return {
"success": True,
"message": f"队列监控间隔已设置为 {interval}",
"timestamp": datetime.now().isoformat()
}
except Exception as e:
logger.error(f"设置监控间隔失败: {e}")
raise HTTPException(status_code=500, detail=f"设置监控间隔失败: {str(e)}")