celery 容灾升级

video-gen-api/app/services/generation/ai/task_group_service.py

@@ -27,6 +27,7 @@ ACTIVE_STAGES = {
     ChatGenerationPipelineStage.QUEUED.value,
     ChatGenerationPipelineStage.PREPARING.value,
     ChatGenerationPipelineStage.CREATING_PROVIDER_TASK.value,
+    ChatGenerationPipelineStage.PROVIDER_RESULT_STAGED.value,
     ChatGenerationPipelineStage.WAITING_REMOTE.value,
     ChatGenerationPipelineStage.POLLING.value,
     ChatGenerationPipelineStage.RESULT_READY.value,
@@ -153,31 +154,7 @@ def _build_summary(children: list[ChatGenerationTask]) -> str | None:
     return f"{len(children)}项中" + "，".join(parts)
 
 
-async def aggregate_main_task_status(
-    db: AsyncSession,
-    *,
-    parent_task_id: str,
-) -> ChatGenerationTask | None:
-    result = await execute_with_lock_timeout(
-        db,
-
-        select(ChatGenerationTask)
-        .where(
-            ChatGenerationTask.id == parent_task_id,
-            ChatGenerationTask.generation_mode == GenerationMode.CHATAPI_MAIN.value,
-        )
-        .with_for_update()
-        .limit(1)
-    )
-    main = result.scalar_one_or_none()
-    if not main or main.deleted_at is not None:
-        return main
-
-    children_map = await load_children_map(db, [parent_task_id], include_deleted=True)
-    children = children_map.get(parent_task_id, [])
-    if not children:
-        return main
-
+def _apply_main_task_status(main: ChatGenerationTask, children: list[ChatGenerationTask]) -> dict[str, object]:
     previous_status = main.status
     previous_stage = main.pipeline_stage
     active_children = [child for child in children if is_task_active(child)]
@@ -217,7 +194,6 @@ async def aggregate_main_task_status(
         )
         main.error_message = _build_summary(children)
     else:
-        # 所有子任务真实生成结果均成功；资源是否软删除不改变生成历史终态。
         main.status = ChatGenerationTaskStatus.COMPLETED.value
         main.pipeline_stage = ChatGenerationPipelineStage.DONE.value
         main.generated_at = max(
@@ -232,29 +208,73 @@ async def aggregate_main_task_status(
         main.text_tokens_used = sum(int(child.text_tokens_used or 0) for child in children)
         main.image_tokens_used = sum(int(child.image_tokens_used or 0) for child in children)
         main.video_tokens_used = sum(int(child.video_tokens_used or 0) for child in children)
-        # main 的手动重试次数只代表 main 自身，不能累加 child 的轮询/重试次数。
         main.retry_count = int(main.manual_retry_count or 0)
         main.poll_count = sum(int(child.poll_count or 0) for child in children)
         main.poll_error_count = sum(int(child.poll_error_count or 0) for child in children)
 
-    await db.flush()
-    log_operation_event(
-        domain="generation_ai_batch",
-        event_type="MAIN_STATUS_AGGREGATED",
-        event_status="success",
-        source="service",
-        user_id=main.user_id,
-        group_id=main.id,
-        task_id=main.id,
-        detail={
-            "before_status": previous_status,
-            "before_stage": previous_stage,
-            "after_status": main.status,
-            "after_stage": main.pipeline_stage,
-            "summary": _build_summary(children),
-        },
+    return {
+        "before_status": previous_status,
+        "before_stage": previous_stage,
+        "after_status": main.status,
+        "after_stage": main.pipeline_stage,
+        "summary": _build_summary(children),
+    }
+
+
+async def aggregate_main_tasks_status_batch(
+    db: AsyncSession,
+    *,
+    parent_task_ids: Sequence[str] | Iterable[str],
+) -> dict[str, ChatGenerationTask]:
+    ids = list(dict.fromkeys(str(item) for item in parent_task_ids if item))
+    if not ids:
+        return {}
+    result = await execute_with_lock_timeout(
+        db,
+        select(ChatGenerationTask)
+        .where(
+            ChatGenerationTask.id.in_(ids),
+            ChatGenerationTask.generation_mode == GenerationMode.CHATAPI_MAIN.value,
+        )
+        .order_by(ChatGenerationTask.id.asc())
+        .with_for_update(),
     )
-    return main
+    mains = list(result.scalars().all())
+    children_map = await load_children_map(db, ids, include_deleted=True)
+    log_snapshots: list[tuple[str, str | None, dict[str, object]]] = []
+    main_map: dict[str, ChatGenerationTask] = {}
+    for main in mains:
+        main_id = str(main.id)
+        main_map[main_id] = main
+        if main.deleted_at is not None:
+            continue
+        children = children_map.get(main_id, [])
+        if not children:
+            continue
+        detail = _apply_main_task_status(main, children)
+        log_snapshots.append((main_id, str(main.user_id) if main.user_id else None, detail))
+    await db.flush()
+    for main_id, user_id, detail in log_snapshots:
+        log_operation_event(
+            domain="generation_ai_batch",
+            event_type="MAIN_STATUS_AGGREGATED",
+            event_status="success",
+            source="service",
+            user_id=user_id,
+            group_id=main_id,
+            task_id=main_id,
+            detail=detail,
+        )
+    return main_map
+
+
+async def aggregate_main_task_status(
+    db: AsyncSession,
+    *,
+    parent_task_id: str,
+) -> ChatGenerationTask | None:
+    main_map = await aggregate_main_tasks_status_batch(db, parent_task_ids=[parent_task_id])
+    return main_map.get(str(parent_task_id))
 
 
 async def aggregate_parent_for_child(db: AsyncSession, child: ChatGenerationTask | None) -> ChatGenerationTask | None: