celery 容灾升级

video-gen-api/app/services/module_generation_v2/flow_service.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import json
 from dataclasses import dataclass
 from types import SimpleNamespace
-from typing import Any
+from typing import Any, Awaitable, Callable
 
 from fastapi import HTTPException
 from sqlalchemy import select
@@ -730,6 +730,7 @@ async def run_video_prompt_optimize_v2(
     *,
     project_id: str,
     step_id: str,
+    execution_guard: Callable[[], Awaitable[None]] | None = None,
 ) -> ModuleGenerationStep | None:
     try:
         meta_result = await execute_with_lock_timeout(
@@ -800,6 +801,8 @@ async def run_video_prompt_optimize_v2(
             step_id=project_snapshot["step_id"],
         )
 
+        if execution_guard is not None:
+            await execution_guard()
         locked = await execute_with_lock_timeout(
             db,
             select(ModuleGenerationProject, ModuleGenerationStep)
@@ -885,6 +888,8 @@ async def run_video_prompt_optimize_v2(
     except Exception as exc:
         await db.rollback()
         try:
+            if execution_guard is not None:
+                await execution_guard()
             result = await execute_with_lock_timeout(
                 db,
                 select(ModuleGenerationProject, ModuleGenerationStep)