From ab4cbf718885a8f452cc3e2c8e212cd9bcaa86d3 Mon Sep 17 00:00:00 2001 From: iHeyTang Date: Mon, 18 Aug 2025 10:59:01 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E6=B7=BB=E5=8A=A0=E9=98=9F=E5=88=97?= =?UTF-8?q?=E7=9B=91=E6=8E=A7=E5=8A=9F=E8=83=BD=EF=BC=8C=E6=94=AF=E6=8C=81?= =?UTF-8?q?=E6=89=8B=E5=8A=A8=E8=A7=A6=E5=8F=91=E5=92=8C=E7=9B=91=E6=8E=A7?= =?UTF-8?q?=E9=97=B4=E9=9A=94=E8=AE=BE=E7=BD=AE=EF=BC=8C=E5=A2=9E=E5=BC=BA?= =?UTF-8?q?=E4=BB=BB=E5=8A=A1=E5=A4=84=E7=90=86=E8=83=BD=E5=8A=9B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- requirements.txt | 1 + workflow_service/comfy/comfy_queue.py | 141 ++++++++++++++++++++++++-- workflow_service/main.py | 13 +++ workflow_service/routes/monitor.py | 45 ++++++++ 4 files changed, 189 insertions(+), 11 deletions(-) diff --git a/requirements.txt b/requirements.txt index b07a582..0b894e4 100644 --- a/requirements.txt +++ b/requirements.txt @@ -6,4 +6,5 @@ aiohttp boto3 aiosqlite sqlalchemy +greenlet psutil \ No newline at end of file diff --git a/workflow_service/comfy/comfy_queue.py b/workflow_service/comfy/comfy_queue.py index 731c5ed..fc6a1c5 100644 --- a/workflow_service/comfy/comfy_queue.py +++ b/workflow_service/comfy/comfy_queue.py @@ -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( diff --git a/workflow_service/main.py b/workflow_service/main.py index 546c37a..a1e6f2d 100644 --- a/workflow_service/main.py +++ b/workflow_service/main.py @@ -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 diff --git a/workflow_service/routes/monitor.py b/workflow_service/routes/monitor.py index 833feec..da47209 100644 --- a/workflow_service/routes/monitor.py +++ b/workflow_service/routes/monitor.py @@ -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)}")