feat: 添加队列监控功能,支持手动触发和监控间隔设置,增强任务处理能力
This commit is contained in:
@@ -6,4 +6,5 @@ aiohttp
|
||||
boto3
|
||||
aiosqlite
|
||||
sqlalchemy
|
||||
greenlet
|
||||
psutil
|
||||
@@ -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(
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)}")
|
||||
|
||||
Reference in New Issue
Block a user