video-gen/video-gen-api/app/tasks/celery_app.py

import logging

from celery import Celery
from celery.signals import worker_process_init, worker_process_shutdown, worker_ready

from app.config import settings
from app.models.base import engine
from app.tasks.async_runner import close_loop, run_async

logger = logging.getLogger("video_gen")


# 显式注册所有 Celery 任务模块，避免新增任务文件后 worker 启动时未注册任务。
# 不再依赖 app.tasks.__init__ 内部 import，也不再依赖 autodiscover_tasks。
CELERY_TASK_IMPORTS = (
    "app.tasks.generation_create_tasks",
    "app.tasks.generation_poll_tasks",
    "app.tasks.generation_download_tasks",
    "app.tasks.generation_recovery_tasks",
    "app.tasks.hot_opening_replicate_tasks",
    "app.tasks.shot_replicate_tasks",
    "app.tasks.shot_replicate_flow_tasks",
    "app.tasks.module_async_recovery_tasks",
    "app.tasks.user_oauth_tasks",
    "app.tasks.cleanup",
)


def _derive_redis_db(url: str, db_no: int) -> str:
    if not url:
        return url
    import re

    if re.search(r"/\d+$", url):
        return re.sub(r"/\d+$", f"/{db_no}", url)
    return url.rstrip("/") + f"/{db_no}"


broker_url = settings.CELERY_BROKER_URL or (_derive_redis_db(settings.REDIS_URL, 1) if settings.REDIS_URL else "")
backend_url = settings.CELERY_RESULT_BACKEND or (_derive_redis_db(settings.REDIS_URL, 2) if settings.REDIS_URL else "")

if broker_url:
    celery_app = Celery("videogen", include=CELERY_TASK_IMPORTS)
    celery_app.conf.update(
        broker_url=broker_url,
        result_backend=backend_url or broker_url,
        imports=CELERY_TASK_IMPORTS,
        task_serializer="json",
        accept_content=["json"],
        result_serializer="json",
        timezone="Asia/Shanghai",
        enable_utc=True,
        task_soft_time_limit=600,
        task_time_limit=900,
        task_acks_late=True,
        task_reject_on_worker_lost=True,
        task_track_started=True,
        worker_prefetch_multiplier=1,
        broker_transport_options={
            "visibility_timeout": 3600,
            "queue_order_strategy": "priority",
            "priority_steps": list(range(10)),
            "sep": ":",
        },
        task_routes={
            "generation.chatapi_create_generation_task": {"queue": "gen_chatapi_create"},
            "generation.poll_generation_task": {"queue": "gen_provider_poll"},
            "generation.download_generation_result_task": {"queue": "gen_result_download"},
            "hot_opening.start_image_prompt_optimize": {"queue": "gen_chatapi_create"},
            "hot_opening.start_video_prompt_optimize": {"queue": "gen_chatapi_create"},
            "shot_replicate.analyze_original_video": {"queue": "gen_chatapi_create"},
            "shot_replicate.analyze_custom_segment_video": {"queue": "gen_chatapi_create"},
            "shot_replicate.split_one_segment": {"queue": "gen_result_download"},
            "shot_replicate.start_image_prompt_optimize": {"queue": "gen_chatapi_create"},
            "shot_replicate.start_video_prompt_optimize": {"queue": "gen_chatapi_create"},
            "shot_replicate.recover_split_tasks_once": {"queue": "gen_result_download"},
            "generation.recover_download_tasks_once": {"queue": "gen_result_download"},
            "generation.recover_generation_tasks_once": {"queue": "gen_result_download"},
            "module_async.recover_module_async_tasks_once": {"queue": "gen_result_download"},
            "user_oauth.update_oauth_accounts": {"queue": "default"},
            "app.tasks.cleanup.*": {"queue": "default"},
        },
    )
else:
    celery_app = None


async def _try_acquire_startup_recovery_lock() -> bool:
    """任意 worker 启动时都可尝试抢恢复锁，避免依赖 hostname 命名。"""
    from app.services.redis_registry_service import redis_acquire_lock

    token = await redis_acquire_lock(
        lock_key=settings.CELERY_STARTUP_RECOVERY_LOCK_KEY,
        ttl_seconds=int(settings.CELERY_STARTUP_RECOVERY_LOCK_TTL_SECONDS or 120),
        log_context="celery_startup_recovery",
    )
    return bool(token)


@worker_ready.connect
def on_worker_ready(sender=None, **kwargs):
    """Celery worker 启动时做一次容灾恢复。

    注意：
    - 不启用 Celery beat。
    - 不要求新增第四条启动命令。
    - 不再依赖 worker hostname 是否包含 gen_result_download。
    - 所有 worker 都尝试抢 Redis 锁，只有抢到锁的 worker 投递恢复任务。
    """
    if celery_app is None:
        return
    if not bool(getattr(settings, "CELERY_STARTUP_RECOVERY_ENABLED", True)):
        logger.info("启动容灾恢复已关闭。CELERY_STARTUP_RECOVERY_ENABLED=false")
        return

    try:
        if not run_async(_try_acquire_startup_recovery_lock()):
            return
    except Exception:
        # Redis 不可用时不阻塞 worker 启动，避免影响稳定生成链路。
        logger.exception("启动容灾恢复锁获取失败，已跳过本次自动恢复投递")
        return

    try:
        from app.tasks.generation_recovery_tasks import (
            recover_download_tasks_once,
            recover_generation_tasks_once,
        )
        from app.tasks.shot_replicate_tasks import recover_split_tasks_once
        from app.tasks.module_async_recovery_tasks import recover_module_async_tasks_once_task

        countdown = max(0, int(settings.DOWNLOAD_RECOVERY_STARTUP_DELAY_SECONDS or 0))

        recover_generation_tasks_once.apply_async(
            countdown=countdown,
            queue="gen_result_download",
            priority=settings.DOWNLOAD_TASK_PRIORITY_RECOVER,
        )
        recover_download_tasks_once.apply_async(
            countdown=countdown + 5,
            queue="gen_result_download",
            priority=settings.DOWNLOAD_TASK_PRIORITY_RECOVER,
        )
        recover_split_tasks_once.apply_async(
            countdown=countdown + 10,
            queue="gen_result_download",
            priority=settings.DOWNLOAD_TASK_PRIORITY_RECOVER,
        )
        recover_module_async_tasks_once_task.apply_async(
            countdown=countdown + 15,
            queue="gen_result_download",
            priority=settings.DOWNLOAD_TASK_PRIORITY_RECOVER,
        )

        logger.info("启动容灾恢复任务已投递。countdown=%s", countdown)
    except Exception:
        logger.exception("启动容灾恢复任务投递失败")


@worker_process_init.connect
def on_worker_process_init(**kwargs):
    """Linux prefork 子进程启动后丢弃 fork 前可能继承的连接池状态。"""
    try:
        run_async(engine.dispose())
    except Exception:
        pass


@worker_process_shutdown.connect
def on_worker_process_shutdown(**kwargs):
    """子进程退出前关闭连接池、Redis 注册表连接和 event loop。"""
    try:
        run_async(engine.dispose())
    except Exception:
        pass

    try:
        from app.services.redis_registry_service import close_registry_redis

        run_async(close_registry_redis())
    except Exception:
        pass
    finally:
        close_loop()