diff --git a/apps/worker/worker_app/celery_app.py b/apps/worker/worker_app/celery_app.py index 6e51e61ff..2cf626277 100755 --- a/apps/worker/worker_app/celery_app.py +++ b/apps/worker/worker_app/celery_app.py @@ -19,4 +19,14 @@ celery_app.conf.imports = ( "worker_app.tasks.batch_download", "worker_app.tasks._startup", "apps.worker.video_processing.dedup", + "worker_app.tasks.cleanup", ) + +# Celery Beat 定时任务调度 +celery_app.conf.beat_schedule = { + "cleanup-stale-pending-tasks": { + "task": "worker.cleanup_stale_pending_tasks", + "schedule": 600.0, # 每 10 分钟（秒） + "options": {"expires": 300}, # 5 分钟过期，避免堆积 + }, +} diff --git a/apps/worker/worker_app/tasks/_startup.py b/apps/worker/worker_app/tasks/_startup.py index ff776c516..c841f4406 100644 --- a/apps/worker/worker_app/tasks/_startup.py +++ b/apps/worker/worker_app/tasks/_startup.py @@ -10,6 +10,9 @@ 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（孤儿任务）。 @@ -83,6 +86,38 @@ def cleanup_stale_jobs(timeout_minutes: int = ORPHAN_TASK_TIMEOUT_MINUTES) -> in 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 """统一清理所有超时的孤儿任务。 @@ -93,15 +128,17 @@ def cleanup_all_stale_tasks(timeout_minutes: int = ORPHAN_TASK_TIMEOUT_MINUTES) """ gen_count = cleanup_orphan_tasks(timeout_minutes) job_count = cleanup_stale_jobs(timeout_minutes) - total = gen_count + job_count + 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, 总计=%d", + "任务清理完成: 孤儿 GenerationTask=%d, 孤儿 Job=%d, 超时 pending=%d, 总计=%d", gen_count, job_count, + pending_count, total, ) - return {"generation_tasks": gen_count, "jobs": job_count} + return {"generation_tasks": gen_count, "jobs": job_count, "pending": pending_count} @worker_ready.connect diff --git a/apps/worker/worker_app/tasks/cleanup.py b/apps/worker/worker_app/tasks/cleanup.py new file mode 100644 index 000000000..28d5e7e22 --- /dev/null +++ b/apps/worker/worker_app/tasks/cleanup.py @@ -0,0 +1,35 @@ +"""定期清理任务 — Celery Beat 调度。 + +包含： +- cleanup_stale_pending_tasks: 定期清理卡在 pending 超时的 generation_tasks +""" + +import logging + +from celery import shared_task +from worker_app.tasks._startup import ( + PENDING_TASK_TIMEOUT_MINUTES, + cleanup_stale_pending_tasks, +) + +logger = logging.getLogger(__name__) + + +@shared_task(name="worker.cleanup_stale_pending_tasks") +def scheduled_cleanup_stale_pending(timeout_minutes: int = PENDING_TASK_TIMEOUT_MINUTES) -> dict: + """Celery Beat 调度的定期任务：清理超时的 pending 任务。 + + 每 10 分钟执行一次（由 celery_app.py 的 beat_schedule 配置）， + 查找所有 status='pending' 且 created_at < NOW() - timeout_minutes + 的 generation_tasks，批量更新为 failed。 + + Args: + timeout_minutes: 超时时间（分钟），默认 30 分钟 + + Returns: + {"cleaned": int} + """ + count = cleanup_stale_pending_tasks(timeout_minutes) + if count > 0: + logger.info("[Beat] 清理了 %d 个超时 pending 任务（超时阈值 %d 分钟）", count, timeout_minutes) + return {"cleaned": count} diff --git a/infra/docker/entrypoint-worker.sh b/infra/docker/entrypoint-worker.sh index d03eb4bd1..f2e208d96 100755 --- a/infra/docker/entrypoint-worker.sh +++ b/infra/docker/entrypoint-worker.sh @@ -6,8 +6,13 @@ set -e CONCURRENCY="${WORKER_CONCURRENCY:-2}" +# ⚠️ 部署约束：此 Worker 必须且只能运行单实例（replicas=1） +# -B 标志嵌入 celery beat，beat 负责定期触发 pending 超时清理等定时任务 +# 多实例部署会导致每个 Worker 独立运行 Beat，造成定时任务重复执行 +# 若需横向扩展 Worker，必须将 Beat 拆分为独立服务（celery beat -A worker_app.celery_app） exec celery \ -A worker_app.celery_app \ worker \ --loglevel=info \ + "-B" \ "--concurrency=${CONCURRENCY}" diff --git a/packages/adapters/sqlalchemy_impl/generation_task_repository.py b/packages/adapters/sqlalchemy_impl/generation_task_repository.py index b0ed071af..75f394edc 100755 --- a/packages/adapters/sqlalchemy_impl/generation_task_repository.py +++ b/packages/adapters/sqlalchemy_impl/generation_task_repository.py @@ -310,3 +310,42 @@ class SQLAlchemyGenerationTaskRepository: model.completed_at = datetime.now(timezone.utc) self.session.commit() return len(models) + + def cleanup_stale_pending(self, timeout_minutes: int = 30) -> int: + """清理超时的 pending 任务（未被 Worker 拉取的任务）。 + + 全局任务队列有 pending 数量上限，长期卡在 pending 的任务会占满队列， + 导致新用户无法创建任务。将超时的 pending 任务标记为 failed。 + + Args: + timeout_minutes: 超时时间（分钟），默认 30 分钟 + + Returns: + 清理的任务数量 + """ + from datetime import timedelta + + cutoff = datetime.now(timezone.utc) - timedelta(minutes=timeout_minutes) + error_info = { + "error_type": "PendingTimeout", + "message": f"任务在 pending 状态停留超过 {timeout_minutes} 分钟，自动清理", + "failed_at": datetime.now(timezone.utc).isoformat(), + } + count = ( + self.session.query(GenerationTaskModel) + .filter( + GenerationTaskModel.status == GenerationTaskStatus.PENDING.value, + GenerationTaskModel.created_at < cutoff, + ) + .update( + { + GenerationTaskModel.status: GenerationTaskStatus.FAILED.value, + GenerationTaskModel.error_message: "pending timeout: auto cleanup", + GenerationTaskModel.error_info: error_info, + GenerationTaskModel.completed_at: datetime.now(timezone.utc), + }, + synchronize_session=False, + ) + ) + self.session.commit() + return count diff --git a/tests/unit/test_generation_task_pending_cleanup.py b/tests/unit/test_generation_task_pending_cleanup.py new file mode 100644 index 000000000..252f1be19 --- /dev/null +++ b/tests/unit/test_generation_task_pending_cleanup.py @@ -0,0 +1,151 @@ +"""GenerationTaskRepository - cleanup_stale_pending 超时 pending 清理单元测试。""" + +import sys +from datetime import datetime, timedelta, timezone +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "apps" / "api")) + +from sqlalchemy import create_engine, text +from sqlalchemy.orm import sessionmaker + +from packages.adapters.sqlalchemy_impl.generation_task_repository import ( + SQLAlchemyGenerationTaskRepository, +) +from packages.adapters.sqlalchemy_impl.models import Base +from packages.domain import GenerationTask, GenerationTaskStatus + + +def _repository(): + engine = create_engine("sqlite:///:memory:") + Base.metadata.create_all(engine) + session = sessionmaker(bind=engine)() + return SQLAlchemyGenerationTaskRepository(session), session, engine + + +def _make_task(**kwargs) -> GenerationTask: + defaults = dict( + project_id="proj-1", + asset_library_id="lib-1", + created_by_user_id="user-1", + ) + defaults.update(kwargs) + return GenerationTask.create(**defaults) + + +# --------------------------------------------------------------------------- +# cleanup_stale_pending 基本测试 +# --------------------------------------------------------------------------- + + +def test_cleanup_stale_pending_no_tasks_returns_zero(): + """没有任务时返回 0。""" + repo, _, _ = _repository() + count = repo.cleanup_stale_pending(timeout_minutes=30) + assert count == 0 + + +def test_cleanup_stale_pending_recent_pending_not_cleaned(): + """30 分钟内的 pending 任务不被清理。""" + repo, _, _ = _repository() + task = _make_task() + repo.create(task) + # 刚创建的 pending 任务不应被清理 + count = repo.cleanup_stale_pending(timeout_minutes=30) + assert count == 0 + assert repo.get(task.id).status == GenerationTaskStatus.PENDING + + +def test_cleanup_stale_pending_old_pending_marked_failed(): + """超过 30 分钟的 pending 任务被标记为 failed。""" + repo, _, engine = _repository() + task = _make_task() + repo.create(task) + + # 手动把 created_at 改到 1 小时前 + with engine.connect() as conn: + conn.execute( + text("UPDATE generation_tasks SET created_at = :ts WHERE id = :id"), + {"ts": datetime.now(timezone.utc) - timedelta(hours=1), "id": task.id}, + ) + conn.commit() + + count = repo.cleanup_stale_pending(timeout_minutes=30) + assert count == 1 + + saved = repo.get(task.id) + assert saved.status == GenerationTaskStatus.FAILED + assert saved.error_message == "pending timeout: auto cleanup" + assert saved.error_info.get("error_type") == "PendingTimeout" + assert "30" in saved.error_info["message"] + assert "failed_at" in saved.error_info + assert saved.completed_at is not None + + +def test_cleanup_stale_pending_running_not_touched(): + """running 任务不受影响，只清理 pending。""" + repo, _, engine = _repository() + task = _make_task() + repo.create(task) + task.mark_processing() + repo.update(task) + + # 回写 created_at 到 1 小时前 + with engine.connect() as conn: + conn.execute( + text("UPDATE generation_tasks SET created_at = :ts WHERE id = :id"), + {"ts": datetime.now(timezone.utc) - timedelta(hours=1), "id": task.id}, + ) + conn.commit() + + count = repo.cleanup_stale_pending(timeout_minutes=30) + assert count == 0 + assert repo.get(task.id).status == GenerationTaskStatus.RUNNING + + +def test_cleanup_stale_pending_custom_timeout(): + """自定义超时时间生效。""" + repo, _, engine = _repository() + task = _make_task() + repo.create(task) + + # 回写 created_at 到 20 分钟前 + with engine.connect() as conn: + conn.execute( + text("UPDATE generation_tasks SET created_at = :ts WHERE id = :id"), + {"ts": datetime.now(timezone.utc) - timedelta(minutes=20), "id": task.id}, + ) + conn.commit() + + # 30 分钟超时：不清理 + count_30 = repo.cleanup_stale_pending(timeout_minutes=30) + assert count_30 == 0 + # 15 分钟超时：清理 + count_15 = repo.cleanup_stale_pending(timeout_minutes=15) + assert count_15 == 1 + assert repo.get(task.id).status == GenerationTaskStatus.FAILED + + +def test_cleanup_stale_pending_multiple(): + """批量清理多个超时的 pending 任务。""" + repo, _, engine = _repository() + + tasks = [] + for i in range(5): + t = _make_task(project_id=f"proj-{i}") + repo.create(t) + tasks.append(t) + + # 全部回写 created_at 到 2 小时前 + with engine.connect() as conn: + for t in tasks: + conn.execute( + text("UPDATE generation_tasks SET created_at = :ts WHERE id = :id"), + {"ts": datetime.now(timezone.utc) - timedelta(hours=2), "id": t.id}, + ) + conn.commit() + + count = repo.cleanup_stale_pending(timeout_minutes=30) + assert count == 5 + for t in tasks: + assert repo.get(t.id).status == GenerationTaskStatus.FAILED