""" MinerU Tianshu - Task Scheduler (Optional) 天枢任务调度器（可选） 在 Worker 自动循环模式下，调度器主要用于： 1. 监控队列状态（默认5分钟一次） 2. 健康检查（默认15分钟一次） 3. 统计信息收集 4. 故障恢复（重置超时任务） 注意： - 如果 workers 启用了自动循环模式（默认），则不需要调度器来触发任务处理 - Worker 已经主动工作，调度器只是偶尔检查系统状态 - 较长的间隔可以最小化系统开销，同时保持必要的监控能力 - 5分钟监控、15分钟健康检查对于自动运行的系统来说已经足够及时 """ import asyncio import aiohttp from loguru import logger from task_db import TaskDB import signal class TaskScheduler: """ 任务调度器（可选） 职责（在 Worker 自动循环模式下）： 1. 监控 SQLite 任务队列状态 2. 健康检查 Workers 3. 故障恢复（重置超时任务） 4. 收集和展示统计信息 职责（在传统模式下）： 1. 触发 Workers 拉取任务 """ def __init__( self, litserve_url='http://localhost:9000/predict', monitor_interval=300, health_check_interval=900, stale_task_timeout=60, cleanup_old_files_days=7, cleanup_old_records_days=0, worker_auto_mode=True ): """ 初始化调度器 Args: litserve_url: LitServe Worker 的 URL monitor_interval: 监控间隔（秒，默认300秒=5分钟） health_check_interval: 健康检查间隔（秒，默认900秒=15分钟） stale_task_timeout: 超时任务重置时间（分钟） cleanup_old_files_days: 清理多少天前的结果文件（0=禁用，默认7天） cleanup_old_records_days: 清理多少天前的数据库记录（0=禁用，不推荐删除） worker_auto_mode: Worker 是否启用自动循环模式 """ self.litserve_url = litserve_url self.monitor_interval = monitor_interval self.health_check_interval = health_check_interval self.stale_task_timeout = stale_task_timeout self.cleanup_old_files_days = cleanup_old_files_days self.cleanup_old_records_days = cleanup_old_records_days self.worker_auto_mode = worker_auto_mode self.db = TaskDB() self.running = True async def check_worker_health(self, session: aiohttp.ClientSession): """ 检查 worker 健康状态 """ try: async with session.post( self.litserve_url, json={'action': 'health'}, timeout=aiohttp.ClientTimeout(total=10) ) as resp: if resp.status == 200: result = await resp.json() return result else: logger.error(f"Health check failed with status {resp.status}") return None except asyncio.TimeoutError: logger.warning("Health check timeout") return None except Exception as e: logger.error(f"Health check error: {e}") return None async def schedule_loop(self): """ 主监控循环 """ logger.info("🔄 Task scheduler started") logger.info(f" LitServe URL: {self.litserve_url}") logger.info(f" Worker Mode: {'Auto-Loop' if self.worker_auto_mode else 'Scheduler-Driven'}") logger.info(f" Monitor Interval: {self.monitor_interval}s") logger.info(f" Health Check Interval: {self.health_check_interval}s") logger.info(f" Stale Task Timeout: {self.stale_task_timeout}m") if self.cleanup_old_files_days > 0: logger.info(f" Cleanup Old Files: {self.cleanup_old_files_days} days") else: logger.info(f" Cleanup Old Files: Disabled") if self.cleanup_old_records_days > 0: logger.info(f" Cleanup Old Records: {self.cleanup_old_records_days} days (Not Recommended)") else: logger.info(f" Cleanup Old Records: Disabled (Keep Forever)") health_check_counter = 0 stale_task_counter = 0 cleanup_counter = 0 async with aiohttp.ClientSession() as session: while self.running: try: # 1. 监控队列状态 stats = self.db.get_queue_stats() pending_count = stats.get('pending', 0) processing_count = stats.get('processing', 0) completed_count = stats.get('completed', 0) failed_count = stats.get('failed', 0) if pending_count > 0 or processing_count > 0: logger.info( f"📊 Queue: {pending_count} pending, {processing_count} processing, " f"{completed_count} completed, {failed_count} failed" ) # 2. 定期健康检查 health_check_counter += 1 if health_check_counter * self.monitor_interval >= self.health_check_interval: health_check_counter = 0 logger.info("🏥 Performing health check...") health_result = await self.check_worker_health(session) if health_result: logger.info(f"✅ Workers healthy: {health_result}") else: logger.warning("⚠️ Workers health check failed") # 3. 定期重置超时任务 stale_task_counter += 1 if stale_task_counter * self.monitor_interval >= self.stale_task_timeout * 60: stale_task_counter = 0 reset_count = self.db.reset_stale_tasks(self.stale_task_timeout) if reset_count > 0: logger.warning(f"⚠️ Reset {reset_count} stale tasks (timeout: {self.stale_task_timeout}m)") # 4. 定期清理旧任务文件和记录 cleanup_counter += 1 # 每24小时清理一次（基于当前监控间隔计算） cleanup_interval_cycles = (24 * 3600) / self.monitor_interval if cleanup_counter >= cleanup_interval_cycles: cleanup_counter = 0 # 清理旧结果文件（保留数据库记录） if self.cleanup_old_files_days > 0: logger.info(f"🧹 Cleaning up result files older than {self.cleanup_old_files_days} days...") file_count = self.db.cleanup_old_task_files(days=self.cleanup_old_files_days) if file_count > 0: logger.info(f"✅ Cleaned up {file_count} result directories (DB records kept)") # 清理极旧的数据库记录（可选，默认不启用） if self.cleanup_old_records_days > 0: logger.warning( f"🗑️ Cleaning up database records older than {self.cleanup_old_records_days} days..." ) record_count = self.db.cleanup_old_task_records(days=self.cleanup_old_records_days) if record_count > 0: logger.warning(f"⚠️ Deleted {record_count} task records permanently") # 等待下一次监控 await asyncio.sleep(self.monitor_interval) except Exception as e: logger.error(f"Scheduler loop error: {e}") await asyncio.sleep(self.monitor_interval) logger.info("⏹️ Task scheduler stopped") def start(self): """启动调度器""" logger.info("🚀 Starting MinerU Tianshu Task Scheduler...") # 设置信号处理 def signal_handler(sig, frame): logger.info("\n🛑 Received stop signal, shutting down...") self.running = False signal.signal(signal.SIGINT, signal_handler) signal.signal(signal.SIGTERM, signal_handler) # 运行调度循环 asyncio.run(self.schedule_loop()) def stop(self): """停止调度器""" self.running = False async def health_check(litserve_url: str) -> bool: """ 健康检查：验证 LitServe Worker 是否可用 """ try: async with aiohttp.ClientSession() as session: async with session.get( litserve_url.replace('/predict', '/health'), timeout=aiohttp.ClientTimeout(total=5) ) as resp: return resp.status == 200 except: return False if __name__ == '__main__': import argparse parser = argparse.ArgumentParser(description='MinerU Tianshu Task Scheduler (Optional)') parser.add_argument('--litserve-url', type=str, default='http://localhost:9000/predict', help='LitServe worker URL') parser.add_argument('--monitor-interval', type=int, default=300, help='Monitor interval in seconds (default: 300s = 5 minutes)') parser.add_argument('--health-check-interval', type=int, default=900, help='Health check interval in seconds (default: 900s = 15 minutes)') parser.add_argument('--stale-task-timeout', type=int, default=60, help='Timeout for stale tasks in minutes (default: 60)') parser.add_argument('--cleanup-old-files-days', type=int, default=7, help='Delete result files older than N days (0=disable, default: 7)') parser.add_argument('--cleanup-old-records-days', type=int, default=0, help='Delete DB records older than N days (0=disable, NOT recommended)') parser.add_argument('--wait-for-workers', action='store_true', help='Wait for workers to be ready before starting') parser.add_argument('--no-worker-auto-mode', action='store_true', help='Disable worker auto-loop mode assumption') args = parser.parse_args() # 等待 workers 就绪（可选） if args.wait_for_workers: logger.info("⏳ Waiting for LitServe workers to be ready...") import time max_retries = 30 for i in range(max_retries): if asyncio.run(health_check(args.litserve_url)): logger.info("✅ LitServe workers are ready!") break time.sleep(2) if i == max_retries - 1: logger.error("❌ LitServe workers not responding, starting anyway...") # 创建并启动调度器 scheduler = TaskScheduler( litserve_url=args.litserve_url, monitor_interval=args.monitor_interval, health_check_interval=args.health_check_interval, stale_task_timeout=args.stale_task_timeout, cleanup_old_files_days=args.cleanup_old_files_days, cleanup_old_records_days=args.cleanup_old_records_days, worker_auto_mode=not args.no_worker_auto_mode ) try: scheduler.start() except KeyboardInterrupt: logger.info("👋 Scheduler interrupted by user")