celery 容灾升级

2026-07-22 14:48:29 +08:00
parent 3f1c4063b0
commit 69e7dec807
67 changed files with 6161 additions and 1958 deletions
@@ -8,7 +8,8 @@ from typing import Any
 from fastapi import HTTPException

 from app.config import settings
-from sqlalchemy import String, cast, func, or_, select
+from app.enums.celery_queue import CeleryQueue
+from sqlalchemy import String, case, cast, func, or_, select
 from sqlalchemy.ext.asyncio import AsyncSession

 from app.enums.shot_replicate import (
@@ -341,67 +342,96 @@ async def _next_segment_index(db: AsyncSession, task_set_id: str) -> int:
    return int(result.scalar() or 0) + 1


-async def refresh_task_set_split_summary(db: AsyncSession, task_set_id: str) -> None:
-    task_set_result = await db.execute(select(ShotReplicateTaskSet).where(ShotReplicateTaskSet.id == task_set_id).with_for_update().limit(1))
-    task_set = task_set_result.scalar_one_or_none()
-    if not task_set:
+async def refresh_task_set_split_summaries(db: AsyncSession, task_set_ids: set[str] | list[str]) -> None:
+    ids = sorted({str(item) for item in task_set_ids if item})
+    if not ids:
        return

-    result = await db.execute(
-        select(ShotReplicateSegment).where(
-            ShotReplicateSegment.task_set_id == task_set_id,
+    task_set_result = await db.execute(
+        select(ShotReplicateTaskSet)
+        .where(ShotReplicateTaskSet.id.in_(ids))
+        .order_by(ShotReplicateTaskSet.id.asc())
+        .with_for_update()
+    )
+    task_sets = list(task_set_result.scalars().all())
+    if not task_sets:
+        return
+
+    count_result = await db.execute(
+        select(
+            ShotReplicateSegment.task_set_id,
+            func.count(ShotReplicateSegment.id).label("total"),
+            func.sum(
+                case(
+                    (ShotReplicateSegment.split_status == ShotSplitStatusEnum.COMPLETED.value, 1),
+                    else_=0,
+                )
+            ).label("completed"),
+            func.sum(
+                case(
+                    (ShotReplicateSegment.split_status == ShotSplitStatusEnum.FAILED.value, 1),
+                    else_=0,
+                )
+            ).label("failed"),
+        )
+        .where(
+            ShotReplicateSegment.task_set_id.in_(ids),
            ShotReplicateSegment.deleted_at.is_(None),
        )
+        .group_by(ShotReplicateSegment.task_set_id)
    )
-    segments = list(result.scalars().all())
-    total = len(segments)
-    completed = len([s for s in segments if s.split_status == ShotSplitStatusEnum.COMPLETED.value])
-    failed = len([s for s in segments if s.split_status == ShotSplitStatusEnum.FAILED.value])
+    count_map = {
+        str(row.task_set_id): (int(row.total or 0), int(row.completed or 0), int(row.failed or 0))
+        for row in count_result.all()
+    }

-    task_set.segment_count = total
-    task_set.completed_segment_count = completed
-    task_set.failed_segment_count = failed
+    for task_set in task_sets:
+        total, completed, failed = count_map.get(str(task_set.id), (0, 0, 0))
+        task_set.segment_count = total
+        task_set.completed_segment_count = completed
+        task_set.failed_segment_count = failed

-    if total <= 0:
-        task_set.split_status = ShotSplitStatusEnum.NONE.value
-        if task_set.analysis_status == ShotAnalysisStatusEnum.COMPLETED.value:
-            task_set.status = ShotTaskSetStatusEnum.ANALYSIS_COMPLETED.value
-        return
+        old_status = task_set.status
+        old_split_status = task_set.split_status
+        if total <= 0:
+            task_set.split_status = ShotSplitStatusEnum.NONE.value
+            if task_set.analysis_status == ShotAnalysisStatusEnum.COMPLETED.value:
+                task_set.status = ShotTaskSetStatusEnum.ANALYSIS_COMPLETED.value
+        elif completed == total:
+            task_set.split_status = ShotSplitStatusEnum.COMPLETED.value
+            task_set.status = ShotTaskSetStatusEnum.SPLIT_COMPLETED.value
+        elif failed == total:
+            task_set.split_status = ShotSplitStatusEnum.FAILED.value
+            task_set.status = ShotTaskSetStatusEnum.FAILED.value
+        elif failed > 0:
+            task_set.split_status = ShotSplitStatusEnum.FAILED.value
+            task_set.status = ShotTaskSetStatusEnum.PARTIAL_FAILED.value
+        else:
+            task_set.split_status = ShotSplitStatusEnum.PROCESSING.value
+            task_set.status = ShotTaskSetStatusEnum.SPLITTING.value

-    old_status = task_set.status
-    old_split_status = task_set.split_status
+        if old_status != task_set.status or old_split_status != task_set.split_status:
+            log_module_event_file(
+                module=MODULE,
+                event_type="SHOT_SPLIT_STATUS_CHANGED",
+                project_id=task_set.id,
+                user_id=task_set.user_id,
+                message="拆镜总任务集拆分状态变更",
+                detail={
+                    "task_set_id": task_set.id,
+                    "from_status": old_status,
+                    "to_status": task_set.status,
+                    "from_split_status": old_split_status,
+                    "to_split_status": task_set.split_status,
+                    "segment_count": total,
+                    "completed_segment_count": completed,
+                    "failed_segment_count": failed,
+                },
+            )

-    if completed == total:
-        task_set.split_status = ShotSplitStatusEnum.COMPLETED.value
-        task_set.status = ShotTaskSetStatusEnum.SPLIT_COMPLETED.value
-    elif failed == total:
-        task_set.split_status = ShotSplitStatusEnum.FAILED.value
-        task_set.status = ShotTaskSetStatusEnum.FAILED.value
-    elif failed > 0:
-        task_set.split_status = ShotSplitStatusEnum.FAILED.value
-        task_set.status = ShotTaskSetStatusEnum.PARTIAL_FAILED.value
-    else:
-        task_set.split_status = ShotSplitStatusEnum.PROCESSING.value
-        task_set.status = ShotTaskSetStatusEnum.SPLITTING.value

-    if old_status != task_set.status or old_split_status != task_set.split_status:
-        log_module_event_file(
-            module=MODULE,
-            event_type="SHOT_SPLIT_STATUS_CHANGED",
-            project_id=task_set.id,
-            user_id=task_set.user_id,
-            message="拆镜总任务集拆分状态变更",
-            detail={
-                "task_set_id": task_set.id,
-                "from_status": old_status,
-                "to_status": task_set.status,
-                "from_split_status": old_split_status,
-                "to_split_status": task_set.split_status,
-                "segment_count": total,
-                "completed_segment_count": completed,
-                "failed_segment_count": failed,
-            },
-        )
+async def refresh_task_set_split_summary(db: AsyncSession, task_set_id: str) -> None:
+    await refresh_task_set_split_summaries(db, {task_set_id})


 async def create_segments_by_ai(
@@ -582,7 +612,7 @@ async def prepare_retry_split_segment(
            "reason": reason,
            "source_path": task_set.video_path,
            "celery_task_name": "shot_replicate.split_one_segment",
-            "queue": "gen_result_download",
+            "queue": CeleryQueue.GEN_SHOT_SPLIT.value,
            "status": "pending",
        },
    )
@@ -658,7 +688,7 @@ async def enqueue_segment_split(segment_id: str, *, countdown: int | None = None

    split_one_segment.apply_async(
        args=[segment_id],
-        queue="gen_result_download",
+        queue=CeleryQueue.GEN_SHOT_SPLIT.value,
        countdown=countdown,
        priority=settings.DOWNLOAD_TASK_PRIORITY_RECOVER if recover else settings.DOWNLOAD_TASK_PRIORITY_NORMAL,
    )
@@ -1019,6 +1049,10 @@ async def prepare_reanalyze_task_set(

    task_set.status = ShotTaskSetStatusEnum.PENDING_ANALYSIS.value
    task_set.analysis_status = ShotAnalysisStatusEnum.PENDING.value
+    task_set.analysis_attempt_no = max(1, int(task_set.analysis_attempt_no or 1)) + 1
+    task_set.analysis_claim_token = None
+    task_set.analysis_started_at = None
+    task_set.analysis_lease_until = None
    task_set.analysis_error_message = None
    task_set.original_video_content = None
    task_set.original_video_category = None
@@ -1076,6 +1110,10 @@ async def prepare_reanalyze_segment(
        raise HTTPException(status_code=409, detail="AI 建议片段默认无需单独分析，如确需重跑请传 force=true")

    segment.analysis_status = ShotSegmentAnalysisStatusEnum.PENDING.value
+    segment.analysis_attempt_no = max(1, int(segment.analysis_attempt_no or 1)) + 1
+    segment.analysis_claim_token = None
+    segment.analysis_started_at = None
+    segment.analysis_lease_until = None
    segment.analysis_error_message = None
    segment.analysis_json = None
    segment.original_video_content = None