celery 容灾升级

video-gen-api/app/services/shot_replicate_flow_service.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import asyncio
 import json
 from datetime import datetime
-from typing import Any
+from typing import Any, Awaitable, Callable
 
 from fastapi import HTTPException
 from sqlalchemy import func, select
@@ -824,7 +824,13 @@ async def submit_image_prompt_optimize(
     return project, step
 
 
-async def run_image_prompt_optimize(db: AsyncSession, *, project_id: str, step_id: str | None = None) -> ModuleGenerationStep | None:
+async def run_image_prompt_optimize(
+    db: AsyncSession,
+    *,
+    project_id: str,
+    step_id: str | None = None,
+    execution_guard: Callable[[], Awaitable[None]] | None = None,
+) -> ModuleGenerationStep | None:
     await apply_short_lock_timeout(db)
     project_result = await db.execute(
         select(ModuleGenerationProject)
@@ -916,6 +922,8 @@ async def run_image_prompt_optimize(db: AsyncSession, *, project_id: str, step_i
             references=references,
             gen_type="image",
         )
+        if execution_guard is not None:
+            await execution_guard()
         project, step = await _reload_prompt_context_for_update(
             db,
             project_id=project_id_value,
@@ -979,6 +987,8 @@ async def run_image_prompt_optimize(db: AsyncSession, *, project_id: str, step_i
         raise
     except Exception as exc:
         await db.rollback()
+        if execution_guard is not None:
+            await execution_guard()
         project, step = await _reload_prompt_context_for_update(
             db,
             project_id=project_id_value,
@@ -1187,7 +1197,13 @@ async def submit_video_prompt_optimize(
     return project, step
 
 
-async def run_video_prompt_optimize(db: AsyncSession, *, project_id: str, step_id: str | None = None) -> ModuleGenerationStep | None:
+async def run_video_prompt_optimize(
+    db: AsyncSession,
+    *,
+    project_id: str,
+    step_id: str | None = None,
+    execution_guard: Callable[[], Awaitable[None]] | None = None,
+) -> ModuleGenerationStep | None:
     await apply_short_lock_timeout(db)
     project_result = await db.execute(
         select(ModuleGenerationProject)
@@ -1296,6 +1312,8 @@ async def run_video_prompt_optimize(db: AsyncSession, *, project_id: str, step_i
             step_id=step_id_value,
             trace_id=f"shot-video-prompt:{step_id_value}",
         )
+        if execution_guard is not None:
+            await execution_guard()
         project, step = await _reload_prompt_context_for_update(
             db,
             project_id=project_id_value,
@@ -1362,6 +1380,8 @@ async def run_video_prompt_optimize(db: AsyncSession, *, project_id: str, step_i
         raise
     except Exception as exc:
         await db.rollback()
+        if execution_guard is not None:
+            await execution_guard()
         project, step = await _reload_prompt_context_for_update(
             db,
             project_id=project_id_value,