celery 容灾升级

video-gen-api/app/services/hot_opening_replicate_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 String, cast, func, or_, select
@@ -874,7 +874,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)
@@ -966,6 +972,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,
@@ -1029,6 +1037,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,
@@ -1227,7 +1237,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)
@@ -1336,6 +1352,8 @@ async def run_video_prompt_optimize(db: AsyncSession, *, project_id: str, step_i
             step_id=step_id_value,
             trace_id=f"hot-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,
@@ -1402,6 +1420,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,