219 lines
9.5 KiB
Python
219 lines
9.5 KiB
Python
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.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",
|
|
"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["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",
|
|
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},
|
|
},
|
|
worker_prefetch_multiplier=1,
|
|
broker_transport_options={
|
|
"visibility_timeout": 3600,
|
|
"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.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},
|
|
# 恢复扫描统一走独立队列，避免占用下载/轮询/创建业务 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.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()
|