"""Worker 启动时的初始化任务 — 孤儿任务清理等.""" import logging from celery.signals import worker_ready from worker_app.db import SessionLocal logger = logging.getLogger(__name__) # 孤儿任务超时阈值:渲染任务超过此时间未更新则视为卡死 ORPHAN_TASK_TIMEOUT_MINUTES = 10 # Pending 任务超时阈值:pending 任务在队列中等待超过此时间则自动清理 PENDING_TASK_TIMEOUT_MINUTES = 30 def cleanup_orphan_tasks(timeout_minutes: int = ORPHAN_TASK_TIMEOUT_MINUTES) -> int: # pragma: no cover """清理数据库中超时未更新的 running GenerationTask(孤儿任务)。 worker 重启或崩溃后,之前处于 running 状态的任务会变成孤儿任务, 一直卡在 running 不动。通过 updated_at 超时判断并标记为 failed。 Args: timeout_minutes: 超时时间(分钟),默认 10 分钟 Returns: 清理的任务数量 """ from packages.adapters.sqlalchemy_impl.generation_task_repository import ( SQLAlchemyGenerationTaskRepository, ) try: session = SessionLocal() repo = SQLAlchemyGenerationTaskRepository(session) count = repo.cleanup_stale_running(timeout_minutes) session.close() if count > 0: logger.warning("清理了 %d 个超时的孤儿 GenerationTask(超过 %d 分钟未更新)", count, timeout_minutes) else: logger.info("无孤儿 GenerationTask 需要清理") return count except Exception as e: logger.error("清理孤儿 GenerationTask 失败: %s", e, exc_info=True) return 0 def cleanup_stale_jobs(timeout_minutes: int = ORPHAN_TASK_TIMEOUT_MINUTES) -> int: # pragma: no cover """清理数据库中超时未更新的 running Job(孤儿任务)。 与 cleanup_orphan_tasks 配合,同时清理 Job 表和 GenerationTask 表。 Returns: 清理的任务数量 """ from datetime import datetime, timedelta, timezone from packages.adapters.sqlalchemy_impl.models import JobModel from packages.domain.job import JobStatus try: session = SessionLocal() cutoff = datetime.now(timezone.utc) - timedelta(minutes=timeout_minutes) stale_jobs = ( session.query(JobModel) .filter( JobModel.status == JobStatus.RUNNING.value, JobModel.updated_at < cutoff, ) .all() ) count = 0 for model in stale_jobs: model.status = JobStatus.FAILED.value model.error_message = f"任务执行中断(超过 {timeout_minutes} 分钟未更新)" count += 1 if count > 0: session.commit() logger.warning("清理了 %d 个超时的孤儿 Job(超过 %d 分钟未更新)", count, timeout_minutes) else: logger.info("无孤儿 Job 需要清理") session.close() return count except Exception as e: logger.error("清理孤儿 Job 失败: %s", e, exc_info=True) return 0 def cleanup_stale_pending_tasks(timeout_minutes: int = PENDING_TASK_TIMEOUT_MINUTES) -> int: # pragma: no cover """清理数据库中卡在 pending 状态超时的 GenerationTask。 全局任务队列有 pending 数量上限,长期卡在 pending 的任务会占满队列, 导致新用户无法创建任务。通过 created_at 超时判断并标记为 failed。 Args: timeout_minutes: 超时时间(分钟),默认 PENDING_TASK_TIMEOUT_MINUTES Returns: 清理的任务数量 """ from packages.adapters.sqlalchemy_impl.generation_task_repository import ( SQLAlchemyGenerationTaskRepository, ) session = SessionLocal() try: repo = SQLAlchemyGenerationTaskRepository(session) count = repo.cleanup_stale_pending(timeout_minutes) if count > 0: logger.warning("清理了 %d 个超时的 pending GenerationTask(超过 %d 分钟未处理)", count, timeout_minutes) else: logger.info("无超时 pending GenerationTask 需要清理") return count except Exception as e: logger.error("清理超时 pending GenerationTask 失败: %s", e, exc_info=True) return 0 finally: session.close() def cleanup_all_stale_tasks(timeout_minutes: int = ORPHAN_TASK_TIMEOUT_MINUTES) -> dict: # pragma: no cover """统一清理所有超时的孤儿任务。 同时清理 GenerationTask 和 Job 两类表。 Returns: {"generation_tasks": int, "jobs": int} """ gen_count = cleanup_orphan_tasks(timeout_minutes) job_count = cleanup_stale_jobs(timeout_minutes) pending_count = cleanup_stale_pending_tasks(PENDING_TASK_TIMEOUT_MINUTES) total = gen_count + job_count + pending_count if total > 0: logger.warning( "任务清理完成: 孤儿 GenerationTask=%d, 孤儿 Job=%d, 超时 pending=%d, 总计=%d", gen_count, job_count, pending_count, total, ) return {"generation_tasks": gen_count, "jobs": job_count, "pending": pending_count} @worker_ready.connect def _on_worker_ready(sender, **kwargs): # pragma: no cover """Worker 启动完成后执行 — 清理孤儿任务。""" logger.info("Worker 启动完成,开始清理孤儿 running 任务...") result = cleanup_all_stale_tasks() total = result["generation_tasks"] + result["jobs"] logger.info("Worker 启动清理完成,共清理 %d 个孤儿任务", total)