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.enums.celery_queue import CeleryQueue, CeleryTaskName 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.video_upscale_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.module_generation_v2_tasks", "app.tasks.user_oauth_tasks", "app.tasks.cleanup", "app.tasks.private_portrait_asset_tasks", ) RECOVERY_QUEUE = settings.CELERY_RECOVERY_QUEUE or CeleryQueue.GEN_RECOVERY.value 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}" def _beat_schedule() -> dict: schedule: dict = {} if bool(getattr(settings, "POLL_DUE_DISPATCH_ENABLED", True)): schedule["dispatch-due-poll-tasks-every-minute"] = { "task": CeleryTaskName.DISPATCH_DUE_POLL.value, "schedule": max(1, int(settings.POLL_DUE_DISPATCH_INTERVAL_SECONDS or 60)), "options": { "queue": RECOVERY_QUEUE, "priority": settings.DOWNLOAD_TASK_PRIORITY_RECOVER, }, } schedule["generation-create-recovery"] = { "task": CeleryTaskName.RECOVER_CREATE.value, "schedule": max(1, int(settings.GENERATION_CREATE_RECOVERY_INTERVAL_SECONDS or 60)), "options": { "queue": RECOVERY_QUEUE, "priority": settings.DOWNLOAD_TASK_PRIORITY_RECOVER, }, } schedule["module-async-recovery"] = { "task": CeleryTaskName.MODULE_ASYNC_RECOVERY.value, "schedule": max(1, int(settings.MODULE_ASYNC_RECOVERY_INTERVAL_SECONDS or 60)), "options": { "queue": RECOVERY_QUEUE, "priority": settings.DOWNLOAD_TASK_PRIORITY_RECOVER, }, } schedule["video-upscale-recovery-every-minute"] = { "task": CeleryTaskName.VIDEO_UPSCALE_RECOVER.value, "schedule": 60, "options": { "queue": RECOVERY_QUEUE, "priority": settings.DOWNLOAD_TASK_PRIORITY_RECOVER, }, } schedule["private-portrait-sync-due-assets-every-minute"] = { "task": CeleryTaskName.PRIVATE_PORTRAIT_SYNC_DUE_ASSETS.value, "schedule": 60, "options": {"queue": CeleryQueue.GEN_PRIVATE_PORTRAIT.value}, } schedule["private-portrait-recover-remote-deletes-every-5-minutes"] = { "task": CeleryTaskName.PRIVATE_PORTRAIT_RECOVER_REMOTE_DELETES.value, "schedule": 300, "options": {"queue": CeleryQueue.GEN_PRIVATE_PORTRAIT.value}, } return schedule 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", result_expires=max(60, int(settings.CELERY_RESULT_EXPIRES_SECONDS or 7200)), task_store_errors_even_if_ignored=True, 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, beat_schedule=_beat_schedule(), task_annotations={ # 生成链路任务以数据库状态为准，不依赖 Celery result backend。 # 这里忽略结果可避免任务误返回 ORM / 非 JSON 对象时触发结果序列化失败。 # "generation.chatapi_create_generation_task": {"ignore_result": True}, # "generation.poll_generation_task": {"ignore_result": True}, # "generation.download_generation_result_task": {"ignore_result": True}, "hot_opening.start_image_prompt_optimize": {"ignore_result": True}, "hot_opening.start_video_prompt_optimize": {"ignore_result": True}, "shot_replicate.analyze_original_video": {"ignore_result": True}, "shot_replicate.analyze_custom_segment_video": {"ignore_result": True}, "shot_replicate.split_one_segment": {"ignore_result": True}, "shot_replicate.start_image_prompt_optimize": {"ignore_result": True}, "shot_replicate.start_video_prompt_optimize": {"ignore_result": True}, "module_generation_v2.start_video_prompt_optimize": {"ignore_result": True}, CeleryTaskName.VIDEO_UPSCALE_EXECUTE_LOCAL.value: { "ignore_result": True, "soft_time_limit": max(60, int(settings.VIDEO_UPSCALE_LOCAL_TIMEOUT_SECONDS or 3600)) + 60, "time_limit": max(60, int(settings.VIDEO_UPSCALE_LOCAL_TIMEOUT_SECONDS or 3600)) + 300, }, CeleryTaskName.VIDEO_UPSCALE_SUBMIT_REMOTE.value: {"ignore_result": True}, CeleryTaskName.VIDEO_UPSCALE_POLL_REMOTE.value: {"ignore_result": True}, CeleryTaskName.VIDEO_UPSCALE_DOWNLOAD_REMOTE_RESULT.value: {"ignore_result": True}, CeleryTaskName.VIDEO_UPSCALE_FINALIZE.value: {"ignore_result": True}, CeleryTaskName.VIDEO_UPSCALE_RECOVER.value: {"ignore_result": True}, CeleryTaskName.RECOVER_CREATE.value: {"ignore_result": True}, CeleryTaskName.MODULE_ASYNC_RECOVERY.value: {"ignore_result": True}, }, worker_prefetch_multiplier=1, broker_transport_options={ "visibility_timeout": max( 3600, int(settings.VIDEO_UPSCALE_LOCAL_TIMEOUT_SECONDS or 3600) + 600, int(settings.VIDEO_UPSCALE_REMOTE_RESULT_DOWNLOAD_TIMEOUT_SECONDS or 600) + 600, ), "queue_order_strategy": "priority", "priority_steps": list(range(10)), "sep": ":", }, task_routes={ CeleryTaskName.CHATAPI_CREATE.value: {"queue": CeleryQueue.GEN_CHATAPI_CREATE.value}, CeleryTaskName.POLL_GENERATION.value: {"queue": CeleryQueue.GEN_PROVIDER_POLL.value}, CeleryTaskName.DOWNLOAD_GENERATION_RESULT.value: {"queue": CeleryQueue.GEN_RESULT_DOWNLOAD.value}, CeleryTaskName.VIDEO_UPSCALE_EXECUTE_LOCAL.value: { "queue": settings.VIDEO_UPSCALE_LOCAL_QUEUE or CeleryQueue.GEN_VIDEO_UPSCALE_LOCAL.value }, CeleryTaskName.VIDEO_UPSCALE_SUBMIT_REMOTE.value: { "queue": settings.VIDEO_UPSCALE_REMOTE_QUEUE or CeleryQueue.GEN_VIDEO_UPSCALE_REMOTE.value }, CeleryTaskName.VIDEO_UPSCALE_POLL_REMOTE.value: { "queue": settings.VIDEO_UPSCALE_REMOTE_QUEUE or CeleryQueue.GEN_VIDEO_UPSCALE_REMOTE.value }, CeleryTaskName.VIDEO_UPSCALE_DOWNLOAD_REMOTE_RESULT.value: { "queue": settings.VIDEO_UPSCALE_REMOTE_QUEUE or CeleryQueue.GEN_VIDEO_UPSCALE_REMOTE.value }, CeleryTaskName.VIDEO_UPSCALE_FINALIZE.value: { "queue": settings.VIDEO_UPSCALE_LOCAL_QUEUE or CeleryQueue.GEN_VIDEO_UPSCALE_LOCAL.value }, CeleryTaskName.VIDEO_UPSCALE_RECOVER.value: {"queue": RECOVERY_QUEUE}, CeleryTaskName.DISPATCH_DUE_POLL.value: {"queue": RECOVERY_QUEUE}, "hot_opening.start_image_prompt_optimize": {"queue": CeleryQueue.GEN_CHATAPI_CREATE.value}, "hot_opening.start_video_prompt_optimize": {"queue": CeleryQueue.GEN_CHATAPI_CREATE.value}, "shot_replicate.analyze_original_video": {"queue": CeleryQueue.GEN_CHATAPI_CREATE.value}, "shot_replicate.analyze_custom_segment_video": {"queue": CeleryQueue.GEN_CHATAPI_CREATE.value}, "shot_replicate.split_one_segment": {"queue": CeleryQueue.GEN_RESULT_DOWNLOAD.value}, "shot_replicate.start_image_prompt_optimize": {"queue": CeleryQueue.GEN_CHATAPI_CREATE.value}, "shot_replicate.start_video_prompt_optimize": {"queue": CeleryQueue.GEN_CHATAPI_CREATE.value}, "module_generation_v2.start_video_prompt_optimize": {"queue": CeleryQueue.GEN_CHATAPI_CREATE.value}, # 恢复扫描统一走独立队列，避免占用下载/轮询/创建业务 worker。 CeleryTaskName.STARTUP_RECOVERY.value: {"queue": RECOVERY_QUEUE}, CeleryTaskName.SHOT_SPLIT_RECOVERY.value: {"queue": RECOVERY_QUEUE}, CeleryTaskName.RECOVER_DOWNLOAD.value: {"queue": RECOVERY_QUEUE}, CeleryTaskName.RECOVER_GENERATION.value: {"queue": RECOVERY_QUEUE}, CeleryTaskName.RECOVER_CREATE.value: {"queue": RECOVERY_QUEUE}, CeleryTaskName.MODULE_ASYNC_RECOVERY.value: {"queue": RECOVERY_QUEUE}, "user_oauth.update_oauth_accounts": {"queue": CeleryQueue.DEFAULT.value}, "app.tasks.cleanup.*": {"queue": CeleryQueue.DEFAULT.value}, CeleryTaskName.PRIVATE_PORTRAIT_POLL_ASSET.value: {"queue": CeleryQueue.GEN_PRIVATE_PORTRAIT.value}, CeleryTaskName.PRIVATE_PORTRAIT_SYNC_DUE_ASSETS.value: {"queue": CeleryQueue.GEN_PRIVATE_PORTRAIT.value}, CeleryTaskName.PRIVATE_PORTRAIT_DELETE_ASSET.value: {"queue": CeleryQueue.GEN_PRIVATE_PORTRAIT.value}, CeleryTaskName.PRIVATE_PORTRAIT_DELETE_PROJECT.value: {"queue": CeleryQueue.GEN_PRIVATE_PORTRAIT.value}, CeleryTaskName.PRIVATE_PORTRAIT_RECOVER_REMOTE_DELETES.value: {"queue": CeleryQueue.GEN_PRIVATE_PORTRAIT.value}, }, ) 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 启动时做一次容灾恢复。 注意： - 启动容灾只投递一个 recovery.startup_recovery_once 协调任务。 - Celery Beat 只用于每分钟触发轻量 generation.dispatch_due_poll_tasks，不跑完整启动容灾。 - 协调任务走独立 gen_recovery 队列，串行扫描并把真实业务任务投回原队列。 - 所有 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 startup_recovery_once countdown = max(0, int(settings.DOWNLOAD_RECOVERY_STARTUP_DELAY_SECONDS or 0)) startup_recovery_once.apply_async( countdown=countdown, queue=RECOVERY_QUEUE, priority=settings.DOWNLOAD_TASK_PRIORITY_RECOVER, ) logger.info( "启动容灾恢复协调任务已投递。queue=%s countdown=%s", RECOVERY_QUEUE, 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()