88 lines
3.0 KiB
Python
88 lines
3.0 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
from sqlalchemy import select
|
|
|
|
from app.models.base import async_session
|
|
from app.models.module_generation_project import ModuleGenerationProject
|
|
from app.services.module_async_recovery_service import (
|
|
OBJECT_MODULE_STEP,
|
|
TASK_MODULE_V2_VIDEO_PROMPT,
|
|
acquire_object_lock,
|
|
cleanup_active_if_terminal,
|
|
ensure_object_lock_owned,
|
|
mark_active_started,
|
|
register_module_step_task,
|
|
release_object_lock,
|
|
)
|
|
from app.services.module_generation_v2.config import VIDEO_PROMPT_OPTIMIZE
|
|
from app.services.module_generation_v2.flow_service import run_video_prompt_optimize_v2
|
|
from app.tasks.async_runner import run_async
|
|
from app.tasks.celery_app import celery_app
|
|
|
|
|
|
async def _run_video_prompt(project_id: str, step_id: str) -> None:
|
|
lock_token = await acquire_object_lock(object_type=OBJECT_MODULE_STEP, object_id=step_id)
|
|
if not lock_token:
|
|
return None
|
|
try:
|
|
async with async_session() as db:
|
|
result = await db.execute(
|
|
select(ModuleGenerationProject.module)
|
|
.where(
|
|
ModuleGenerationProject.id == project_id,
|
|
ModuleGenerationProject.deleted_at.is_(None),
|
|
)
|
|
.limit(1)
|
|
)
|
|
module = result.scalar_one_or_none()
|
|
if not module:
|
|
return None
|
|
await register_module_step_task(
|
|
module=str(module),
|
|
project_id=project_id,
|
|
step_id=step_id,
|
|
step_code=VIDEO_PROMPT_OPTIMIZE,
|
|
task_name=TASK_MODULE_V2_VIDEO_PROMPT,
|
|
)
|
|
await mark_active_started(object_type=OBJECT_MODULE_STEP, object_id=step_id)
|
|
await run_video_prompt_optimize_v2(
|
|
db,
|
|
project_id=project_id,
|
|
step_id=step_id,
|
|
execution_guard=lambda: ensure_object_lock_owned(token=lock_token),
|
|
)
|
|
await cleanup_active_if_terminal(db, object_type=OBJECT_MODULE_STEP, object_id=step_id)
|
|
finally:
|
|
await release_object_lock(
|
|
object_type=OBJECT_MODULE_STEP,
|
|
object_id=step_id,
|
|
token=lock_token,
|
|
)
|
|
|
|
|
|
if celery_app:
|
|
@celery_app.task(
|
|
name="module_generation_v2.start_video_prompt_optimize",
|
|
bind=True,
|
|
max_retries=3,
|
|
default_retry_delay=30,
|
|
ignore_result=True,
|
|
)
|
|
def start_video_prompt_optimize_v2(self, project_id: str, step_id: str):
|
|
try:
|
|
run_async(_run_video_prompt(project_id, step_id))
|
|
return None
|
|
except Exception as exc:
|
|
raise self.retry(exc=exc) from exc
|
|
else:
|
|
class _DisabledTask:
|
|
def delay(self, *args: Any, **kwargs: Any):
|
|
raise RuntimeError("Celery is disabled")
|
|
|
|
def apply_async(self, *args: Any, **kwargs: Any):
|
|
raise RuntimeError("Celery is disabled")
|
|
|
|
start_video_prompt_optimize_v2 = _DisabledTask()
|