celery 容灾升级

video-gen-api/app/tasks/shot_replicate_flow_tasks.py

@@ -10,6 +10,7 @@ from app.services.module_async_recovery_service import (
     TASK_SHOT_VIDEO_PROMPT,
     acquire_object_lock,
     cleanup_active_if_terminal,
+    ensure_object_lock_owned,
     mark_active_started,
     release_object_lock,
     register_module_step_task,
@@ -37,7 +38,16 @@ async def _run_image_prompt(project_id: str, step_id: str | None = None) -> None
         await mark_active_started(object_type=OBJECT_MODULE_STEP, object_id=step_id)
     try:
         async with async_session() as db:
-            await run_image_prompt_optimize(db, project_id=project_id, step_id=step_id)
+            await run_image_prompt_optimize(
+                db,
+                project_id=project_id,
+                step_id=step_id,
+                execution_guard=(
+                    (lambda: ensure_object_lock_owned(token=lock_token))
+                    if lock_token
+                    else None
+                ),
+            )
             await db.commit()
             if step_id:
                 await cleanup_active_if_terminal(db, object_type=OBJECT_MODULE_STEP, object_id=step_id)
@@ -63,7 +73,16 @@ async def _run_video_prompt(project_id: str, step_id: str | None = None) -> None
         await mark_active_started(object_type=OBJECT_MODULE_STEP, object_id=step_id)
     try:
         async with async_session() as db:
-            await run_video_prompt_optimize(db, project_id=project_id, step_id=step_id)
+            await run_video_prompt_optimize(
+                db,
+                project_id=project_id,
+                step_id=step_id,
+                execution_guard=(
+                    (lambda: ensure_object_lock_owned(token=lock_token))
+                    if lock_token
+                    else None
+                ),
+            )
             await db.commit()
             if step_id:
                 await cleanup_active_if_terminal(db, object_type=OBJECT_MODULE_STEP, object_id=step_id)