火山引擎SMS API|celery容灾优化|生成模型引擎积分列表API

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

@@ -5,7 +5,11 @@ custom named tasks are registered when workers start.
 """
 
 try:
-    from app.tasks import generation_create_tasks, generation_poll_tasks, generation_download_tasks  # noqa: F401
+    from app.tasks import (  # noqa: F401
+        generation_create_tasks,
+        generation_poll_tasks,
+        generation_download_tasks,
+        generation_recovery_tasks,
+    )
 except Exception:
-    # Keep application importable even when optional Celery dependencies/config are absent.
     pass

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

@@ -1,16 +1,20 @@
+import logging
+
 from celery import Celery
+from celery.signals import worker_process_init, worker_process_shutdown, worker_ready
+
 from app.config import settings
-
-from celery.signals import worker_process_init, worker_process_shutdown
-
-from app.tasks.async_runner import run_async, close_loop
 from app.models.base import engine
+from app.tasks.async_runner import close_loop, run_async
+
+logger = logging.getLogger("video_gen")
 
 
 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}"
@@ -37,11 +41,16 @@ if broker_url:
         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"},
+            "generation.recover_download_tasks_once": {"queue": "gen_result_download"},
+            "generation.recover_generation_tasks_once": {"queue": "gen_result_download"},
             "app.tasks.cleanup.*": {"queue": "default"},
         },
     )
@@ -50,15 +59,47 @@ else:
     celery_app = None
 
 
+@worker_ready.connect
+def on_worker_ready(sender=None, **kwargs):
+    """Celery worker 启动时做一次容灾恢复。
+
+    注意：
+    - 不启用 Celery beat。
+    - 不要求新增第四条启动命令。
+    - 只让 gen_result_download worker 投递恢复任务，避免三个 worker 同时重复扫描。
+    """
+    if celery_app is None:
+        return
+
+    hostname = str(getattr(sender, "hostname", "") or "")
+    if "gen_result_download" not in hostname:
+        return
+
+    try:
+        from app.tasks.generation_recovery_tasks import (
+            recover_download_tasks_once,
+            recover_generation_tasks_once,
+        )
+
+        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,
+        )
+    except Exception:
+        logger.exception("启动容灾恢复任务投递失败")
+
+
 @worker_process_init.connect
 def on_worker_process_init(**kwargs):
-    """
-    Linux prefork 子进程启动后执行。
-
-    目的：
-    1. 丢弃 fork 前可能继承的连接池状态。
-    2. 后续任务会在当前子进程自己的长期 event loop 上重新建连接池。
-    """
+    """Linux prefork 子进程启动后丢弃 fork 前可能继承的连接池状态。"""
     try:
         run_async(engine.dispose())
     except Exception:
@@ -67,12 +108,17 @@ def on_worker_process_init(**kwargs):
 
 @worker_process_shutdown.connect
 def on_worker_process_shutdown(**kwargs):
-    """
-    子进程退出前关闭连接池和 event loop。
-    """
+    """子进程退出前关闭连接池和 event loop。"""
     try:
         run_async(engine.dispose())
     except Exception:
         pass
+
+    try:
+        from app.services.celery_download_recovery_service import close_registry_redis
+
+        run_async(close_registry_redis())
+    except Exception:
+        pass
     finally:
-        close_loop()
+        close_loop()

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

@@ -172,9 +172,9 @@ async def _run(task_id: str):
                 task.pipeline_stage = "result_ready"
                 await db.commit()
 
-                from app.tasks.generation_download_tasks import download_generation_result_task
+                from app.tasks.generation_download_tasks import enqueue_download_task
 
-                download_generation_result_task.delay(task.id)
+                await enqueue_download_task(db, task, reason="create_remote_result_ready")
                 return
 
             else:
@@ -223,9 +223,9 @@ async def _run(task_id: str):
                 )
 
             if task.pipeline_stage == "result_ready":
-                from app.tasks.generation_download_tasks import download_generation_result_task
+                from app.tasks.generation_download_tasks import enqueue_download_task
 
-                download_generation_result_task.delay(task.id)
+                await enqueue_download_task(db, task, reason="create_result_ready")
             else:
                 from app.tasks.generation_poll_tasks import poll_generation_task
 

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

@@ -1,10 +1,19 @@
 from app.tasks.async_runner import run_async
 from datetime import datetime, timezone, timedelta
+import uuid
 
 from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
 
+from app.config import settings
 from app.models.base import async_session
 from app.models.chat_generation_task import ChatGenerationTask
+from app.services.celery_download_recovery_service import (
+    build_download_active_payload,
+    ensure_aware_utc,
+    remove_download_active,
+    upsert_download_active,
+)
 from app.services.error_codes import extract_error_message
 from app.services.generation_download_service import download_generation_result
 from app.services.generation_log_service import log_task_event
@@ -12,52 +21,132 @@ from app.services.generation_refund_service import mark_chat_generation_task_fai
 from app.services.resource_accounting_service import record_chat_task_generated_resource
 from app.tasks.celery_app import celery_app
 
-
-# downloading 卡住多久后允许自动恢复。
-# 说明：
-# - worker 在 pipeline_stage 改成 downloading 后，如果被 kill，任务可能永远停在 downloading。
-# - 这里允许超过该时间的 downloading 任务重新进入下载流程。
-# - 如果你的视频文件特别大，可以把这个时间调大，比如 20 * 60。
-DOWNLOAD_STUCK_SECONDS = 10 * 60
+DOWNLOAD_QUEUE = "gen_result_download"
+DOWNLOAD_STAGE_QUEUED = "download_queued"
+DOWNLOAD_STAGE_DOWNLOADING = "downloading"
+DOWNLOAD_STAGE_RETRY_WAITING = "retry_waiting"
+DOWNLOAD_STAGE_DONE = "done"
+DOWNLOAD_STAGE_FAILED = "download_failed"
 
 
-def _to_aware_utc(dt):
-    """
-    把 datetime 统一转成 timezone-aware UTC，避免 offset-naive 和 offset-aware 比较报错。
-    PostgreSQL / SQLite / 不同驱动返回的 updated_at 可能有时区，也可能没有。
-    """
-    if not dt:
+def _now() -> datetime:
+    return datetime.now(timezone.utc)
+
+
+def _task_created_date_dir(task: ChatGenerationTask) -> str:
+    created_at = ensure_aware_utc(getattr(task, "created_at", None)) or _now()
+    return created_at.strftime("%Y/%m/%d")
+
+
+def _queue_timeout_at(now: datetime | None = None) -> datetime:
+    now = now or _now()
+    return now + timedelta(seconds=int(settings.DOWNLOAD_TASK_QUEUE_TIMEOUT_SECONDS or 300))
+
+
+def _lease_until(now: datetime | None = None) -> datetime:
+    now = now or _now()
+    return now + timedelta(seconds=int(settings.DOWNLOAD_TASK_LEASE_SECONDS or 600))
+
+
+def _retry_at(attempt: int, now: datetime | None = None) -> datetime:
+    now = now or _now()
+    base = int(settings.DOWNLOAD_TASK_RETRY_BACKOFF_SECONDS or 30)
+    return now + timedelta(seconds=max(1, base * max(1, attempt)))
+
+
+def _is_expired(value: datetime | None, now: datetime | None = None) -> bool:
+    value = ensure_aware_utc(value)
+    if value is None:
+        return True
+    return value <= (now or _now())
+
+
+def _is_already_completed(task: ChatGenerationTask) -> bool:
+    if task.status == "completed" or task.pipeline_stage == DOWNLOAD_STAGE_DONE:
+        if task.gen_type == "image" and task.image_url:
+            return True
+        if task.gen_type == "video" and task.video_url:
+            return True
+    return False
+
+
+def _build_celery_task_id(task_id: str, attempt: int | None = None, reason: str | None = None) -> str:
+    safe_reason = (reason or "download").replace(" ", "_")[:32]
+    return f"download:{task_id}:{int(attempt or 0)}:{safe_reason}:{uuid.uuid4().hex[:12]}"
+
+
+async def _register_active_from_task(
+    task: ChatGenerationTask,
+    *,
+    check_at: datetime,
+    priority: int,
+    reason: str | None = None,
+) -> None:
+    payload = build_download_active_payload(
+        record_id=task.id,
+        celery_task_id=task.download_celery_task_id,
+        stage=task.pipeline_stage or "",
+        attempt=task.download_attempt_count or 0,
+        queue=DOWNLOAD_QUEUE,
+        priority=priority,
+        enqueue_at=task.download_enqueued_at,
+        started_at=task.download_started_at,
+        lease_until=task.download_lease_until,
+        next_retry_at=task.download_next_retry_at,
+        check_at=check_at,
+        reason=reason,
+    )
+    await upsert_download_active(record_id=task.id, payload=payload, check_at=check_at)
+
+
+async def enqueue_download_task(
+    db: AsyncSession,
+    task: ChatGenerationTask,
+    *,
+    recover: bool = False,
+    reason: str | None = None,
+    countdown: int | None = None,
+) -> str | None:
+    """统一投递图片/视频下载任务，并同步 DB + Redis active 注册表。"""
+    if not task or task.generation_mode != "chatapi_async":
         return None
-    if dt.tzinfo is None:
-        return dt.replace(tzinfo=timezone.utc)
-    return dt.astimezone(timezone.utc)
+    if task.status != "generating":
+        return None
+    if not task.remote_result_url:
+        return None
+
+    now = _now()
+    priority = settings.DOWNLOAD_TASK_PRIORITY_RECOVER if recover else settings.DOWNLOAD_TASK_PRIORITY_NORMAL
+    celery_task_id = _build_celery_task_id(
+        task.id,
+        attempt=task.download_attempt_count or task.retry_count or 0,
+        reason=reason or ("recover" if recover else "normal"),
+    )
+
+    task.pipeline_stage = DOWNLOAD_STAGE_QUEUED
+    task.download_celery_task_id = celery_task_id
+    task.download_enqueued_at = now
+    task.download_next_retry_at = None
+    if not task.download_storage_date_dir:
+        task.download_storage_date_dir = _task_created_date_dir(task)
+
+    await db.commit()
+
+    check_at = _queue_timeout_at(now)
+    await _register_active_from_task(task, check_at=check_at, priority=priority, reason=reason)
+
+    if celery_app:
+        download_generation_result_task.apply_async(
+            args=[task.id],
+            queue=DOWNLOAD_QUEUE,
+            priority=priority,
+            countdown=countdown,
+            task_id=celery_task_id,
+        )
+    return celery_task_id
 
 
-def _is_recent_downloading(task: ChatGenerationTask) -> bool:
-    """
-    判断 downloading 是否仍然是较新的下载任务。
-
-    返回 True：
-    - 说明可能有另一个 worker 刚进入下载，不要重复下载。
-
-    返回 False：
-    - 说明 downloading 已经超过 DOWNLOAD_STUCK_SECONDS，认为可能卡死，可以恢复。
-    """
-    updated_at = _to_aware_utc(getattr(task, "updated_at", None))
-    if not updated_at:
-        return False
-
-    return datetime.now(timezone.utc) - updated_at < timedelta(seconds=DOWNLOAD_STUCK_SECONDS)
-
-
-async def _reload_task(db, task_id: str) -> ChatGenerationTask | None:
-    """
-    rollback 后重新查询任务对象。
-
-    说明：
-    - SQLAlchemy rollback 后，当前 ORM 对象可能过期。
-    - 继续访问旧 task 有概率触发异步懒加载异常。
-    """
+async def _reload_task(db: AsyncSession, task_id: str) -> ChatGenerationTask | None:
     result = await db.execute(
         select(ChatGenerationTask).where(
             ChatGenerationTask.id == task_id,
@@ -67,6 +156,110 @@ async def _reload_task(db, task_id: str) -> ChatGenerationTask | None:
     return result.scalar_one_or_none()
 
 
+async def _claim_download_lease(db: AsyncSession, task: ChatGenerationTask) -> bool:
+    now = _now()
+
+    if not task or task.generation_mode != "chatapi_async":
+        return False
+    if task.status != "generating":
+        return False
+    if _is_already_completed(task):
+        return False
+    if not task.remote_result_url:
+        return False
+
+    stage = task.pipeline_stage
+
+    if stage == DOWNLOAD_STAGE_DOWNLOADING:
+        if not _is_expired(task.download_lease_until, now):
+            return False
+        await log_task_event(
+            task,
+            event_type="DOWNLOAD_STUCK_RECOVER",
+            message=f"downloading lease 已过期，重新抢占下载。lease_until={task.download_lease_until}",
+        )
+    elif stage == DOWNLOAD_STAGE_RETRY_WAITING:
+        if not _is_expired(task.download_next_retry_at, now):
+            return False
+    elif stage in (DOWNLOAD_STAGE_QUEUED, "result_ready"):
+        pass
+    else:
+        return False
+
+    old_stage = stage
+    task.pipeline_stage = DOWNLOAD_STAGE_DOWNLOADING
+    task.download_started_at = now
+    task.download_lease_until = _lease_until(now)
+    task.download_next_retry_at = None
+    task.download_attempt_count = int(task.download_attempt_count or 0) + 1
+    task.retry_count = task.download_attempt_count
+    if not task.download_storage_date_dir:
+        task.download_storage_date_dir = _task_created_date_dir(task)
+
+    await db.commit()
+
+    await _register_active_from_task(
+        task,
+        check_at=task.download_lease_until,
+        priority=settings.DOWNLOAD_TASK_PRIORITY_NORMAL,
+        reason="claim_download_lease",
+    )
+
+    await log_task_event(
+        task,
+        event_type="DOWNLOAD_START",
+        from_stage=old_stage,
+        to_stage=DOWNLOAD_STAGE_DOWNLOADING,
+        detail={
+            "attempt": task.download_attempt_count,
+            "lease_until": task.download_lease_until,
+            "download_celery_task_id": task.download_celery_task_id,
+        },
+    )
+    return True
+
+
+async def _mark_retry_waiting(db: AsyncSession, task: ChatGenerationTask, exc: Exception) -> datetime:
+    now = _now()
+    attempt = int(task.download_attempt_count or task.retry_count or 0)
+    next_retry_at = _retry_at(attempt, now)
+    error_message = extract_error_message(exc, "下载") if callable(extract_error_message) else str(exc)
+
+    retry_celery_task_id = _build_celery_task_id(task.id, attempt=attempt, reason="retry_waiting")
+
+    task.pipeline_stage = DOWNLOAD_STAGE_RETRY_WAITING
+    task.download_celery_task_id = retry_celery_task_id
+    task.download_next_retry_at = next_retry_at
+    task.download_lease_until = None
+    task.download_last_error = error_message
+    task.retry_count = attempt
+    await db.commit()
+
+    await _register_active_from_task(
+        task,
+        check_at=next_retry_at,
+        priority=settings.DOWNLOAD_TASK_PRIORITY_RECOVER,
+        reason="download_retry_waiting",
+    )
+
+    await log_task_event(
+        task,
+        event_type="DOWNLOAD_RETRY_WAITING",
+        message=error_message,
+        to_stage=DOWNLOAD_STAGE_RETRY_WAITING,
+        detail={
+            "attempt": attempt,
+            "next_retry_at": next_retry_at,
+            "download_celery_task_id": retry_celery_task_id,
+        },
+    )
+    return next_retry_at
+
+
+def _should_final_fail(task: ChatGenerationTask) -> bool:
+    return int(task.download_attempt_count or task.retry_count or 0) >= int(settings.DOWNLOAD_TASK_MAX_ATTEMPTS or 3)
+
+
 async def _run(task_id: str):
     async with async_session() as db:
         result = await db.execute(
@@ -76,46 +269,20 @@ async def _run(task_id: str):
             ).with_for_update().limit(1)
         )
         task = result.scalar_one_or_none()
-        if not task or task.generation_mode != "chatapi_async":
+        if not task:
             return
 
-        if task.status != "generating":
-            return
-
-        # 关键修改 3：
-        # 原来只允许 result_ready 进入下载。
-        # 现在允许 downloading 恢复，但只有“卡住超过 DOWNLOAD_STUCK_SECONDS”的 downloading 才继续。
-        if task.pipeline_stage == "downloading":
-            if _is_recent_downloading(task):
-                # downloading 很新，说明可能有 worker 正在下载，直接跳过，避免并发重复下载。
-                return
-
-            # downloading 已经很久没更新，认为 worker 可能挂了，允许恢复下载。
-            await log_task_event(
-                task,
-                event_type="DOWNLOAD_STUCK_RECOVER",
-                message=f"downloading 超过 {DOWNLOAD_STUCK_SECONDS} 秒，重新进入下载流程",
-            )
-
-        elif task.pipeline_stage != "result_ready":
+        claimed = await _claim_download_lease(db, task)
+        if not claimed:
             return
 
         try:
-            old_stage = task.pipeline_stage
-
-            # 无论从 result_ready 进入，还是从 stuck downloading 恢复，都重新标记为 downloading。
-            task.pipeline_stage = "downloading"
-            await db.commit()
-
-            await log_task_event(
-                task,
-                event_type="DOWNLOAD_START",
-                from_stage=old_stage,
-                to_stage="downloading",
-            )
-
             downloaded = await download_generation_result(task)
 
+            task = await _reload_task(db, task_id)
+            if not task:
+                return
+
             if task.gen_type == "image":
                 task.image_url = downloaded.url
             else:
@@ -123,9 +290,12 @@ async def _run(task_id: str):
                 task.video_cover_url = downloaded.cover_url
 
             task.status = "completed"
-            task.pipeline_stage = "done"
-            task.generated_at = datetime.now(timezone.utc)
+            task.pipeline_stage = DOWNLOAD_STAGE_DONE
+            task.generated_at = _now()
             task.retry_count = 0
+            task.download_lease_until = None
+            task.download_next_retry_at = None
+            task.download_last_error = None
 
             await record_chat_task_generated_resource(
                 db,
@@ -138,22 +308,22 @@ async def _run(task_id: str):
             )
 
             await db.commit()
+            await remove_download_active(task.id)
 
             await log_task_event(
                 task,
                 event_type="DOWNLOAD_SUCCESS",
                 to_status="completed",
-                to_stage="done",
+                to_stage=DOWNLOAD_STAGE_DONE,
                 detail={
                     "resource_url": downloaded.url,
                     "video_cover_url": downloaded.cover_url,
                     "file_size_bytes": downloaded.file_size_bytes,
+                    "download_attempt_count": task.download_attempt_count,
                 },
             )
 
         except Exception as exc:
-            # 关键修改 2：
-            # 异常后先 rollback，再重新查询 task，不继续使用 rollback 前的旧 ORM 对象。
             try:
                 await db.rollback()
             except Exception:
@@ -163,33 +333,42 @@ async def _run(task_id: str):
             if not task:
                 return
 
-            task.retry_count = (task.retry_count or 0) + 1
-
-            if task.retry_count > 3:
+            if _should_final_fail(task):
                 error_message = extract_error_message(exc, "下载") if callable(extract_error_message) else str(exc)
                 await mark_chat_generation_task_failed_and_refund_once(
                     db,
                     task=task,
                     error_message=error_message,
-                    pipeline_stage="download_failed",
+                    pipeline_stage=DOWNLOAD_STAGE_FAILED,
                 )
+                task.download_last_error = error_message
+                task.download_lease_until = None
+                task.download_next_retry_at = None
                 await db.commit()
 
+                await remove_download_active(task.id)
+
                 await log_task_event(
                     task,
                     event_type="DOWNLOAD_FAILED",
                     message=task.error_message,
+                    detail={
+                        "download_attempt_count": task.download_attempt_count,
+                        "max_attempts": settings.DOWNLOAD_TASK_MAX_ATTEMPTS,
+                    },
                 )
             else:
-                # 下载失败但未超过重试次数，改回 result_ready，等待下一次下载。
-                # 这样不会卡死在 downloading。
-                task.pipeline_stage = "result_ready"
-                await db.commit()
+                next_retry_at = await _mark_retry_waiting(db, task, exc)
 
-                download_generation_result_task.apply_async(
-                    args=[task.id],
-                    countdown=30 * task.retry_count,
-                )
+                if celery_app:
+                    delay_seconds = max(1, int((next_retry_at - _now()).total_seconds()))
+                    download_generation_result_task.apply_async(
+                        args=[task.id],
+                        queue=DOWNLOAD_QUEUE,
+                        priority=settings.DOWNLOAD_TASK_PRIORITY_RECOVER,
+                        countdown=delay_seconds,
+                        task_id=task.download_celery_task_id,
+                    )
 
 
 if celery_app:

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

@@ -140,8 +140,9 @@ async def _run(task_id: str):
 
                 await log_task_event(task, event_type="POLL_SUCCESS", to_stage="result_ready")
 
-                from app.tasks.generation_download_tasks import download_generation_result_task
-                download_generation_result_task.delay(task.id)
+                from app.tasks.generation_download_tasks import enqueue_download_task
+
+                await enqueue_download_task(db, task, reason="poll_success_result_ready")
                 return
 
             if _is_failed(status):

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

@@ -0,0 +1,46 @@
+# app/tasks/generation_recovery_tasks.py
+from __future__ import annotations
+
+from typing import Any, Dict
+
+from app.models.base import async_session
+from app.tasks.async_runner import run_async
+from app.tasks.celery_app import celery_app
+
+
+async def _run_download_once() -> Dict[str, Any]:
+    from app.services.generation_recovery_service import recover_download_tasks_once
+
+    async with async_session() as db:
+        return await recover_download_tasks_once(db)
+
+
+async def _run_generation_once() -> Dict[str, Any]:
+    from app.services.generation_recovery_service import recover_generation_tasks_once
+
+    async with async_session() as db:
+        return await recover_generation_tasks_once(db)
+
+
+if celery_app:
+
+    @celery_app.task(name="generation.recover_download_tasks_once")
+    def recover_download_tasks_once() -> Dict[str, Any]:
+        return run_async(_run_download_once())
+
+
+    @celery_app.task(name="generation.recover_generation_tasks_once")
+    def recover_generation_tasks_once() -> Dict[str, Any]:
+        return run_async(_run_generation_once())
+
+else:
+
+    class _DisabledTask:
+        def delay(self, *args: Any, **kwargs: Any) -> None:
+            raise RuntimeError("Celery is disabled")
+
+        def apply_async(self, *args: Any, **kwargs: Any) -> None:
+            raise RuntimeError("Celery is disabled")
+
+    recover_download_tasks_once = _DisabledTask()
+    recover_generation_tasks_once = _DisabledTask()