212 lines
8.3 KiB
Python
212 lines
8.3 KiB
Python
from __future__ import annotations
|
|
|
|
import uuid
|
|
from collections.abc import Awaitable, Callable
|
|
from typing import Any
|
|
|
|
from app.config import settings
|
|
from app.enums.celery_queue import CeleryQueue, CeleryTaskName
|
|
from app.enums.celery_runtime import CeleryRuntimeDomain
|
|
from app.models.base import async_session
|
|
from app.services.celery_runtime.runtime_service import CeleryRuntimeLease, RuntimeIdentity
|
|
from app.services.redis_registry_service import RedisExecutionLockLease
|
|
from app.services.video_upscale.task_service import (
|
|
recover_video_upscale_tasks_once,
|
|
run_finalize_upscale,
|
|
run_local_upscale,
|
|
run_remote_poll,
|
|
run_remote_result_download,
|
|
run_remote_submit,
|
|
)
|
|
from app.tasks.async_runner import run_async
|
|
from app.tasks.celery_app import celery_app
|
|
|
|
|
|
async def _run_with_execution_lock(
|
|
upscale_task_id: str,
|
|
callback: Callable[[str, str, Callable[[], Awaitable[None]]], Awaitable[None]],
|
|
*,
|
|
task_name: str,
|
|
queue: str,
|
|
pipeline_stage: str,
|
|
) -> None:
|
|
lock_key = f"{settings.VIDEO_UPSCALE_EXECUTION_LOCK_KEY_PREFIX}:{upscale_task_id}"
|
|
token = uuid.uuid4().hex
|
|
lease = await CeleryRuntimeLease.acquire(
|
|
identity=RuntimeIdentity(
|
|
domain=CeleryRuntimeDomain.VIDEO_UPSCALE.value,
|
|
owner_type="video_upscale_task",
|
|
owner_id=upscale_task_id,
|
|
attempt_no=1,
|
|
task_name=task_name,
|
|
queue=queue,
|
|
registry_item_id=f"video_upscale:{upscale_task_id}",
|
|
),
|
|
lock_key=lock_key,
|
|
hash_key=settings.VIDEO_UPSCALE_ACTIVE_REDIS_HASH_KEY,
|
|
zset_key=settings.VIDEO_UPSCALE_ACTIVE_REDIS_ZSET_KEY,
|
|
token=token,
|
|
ttl_seconds=max(30, int(settings.VIDEO_UPSCALE_EXECUTION_LOCK_TTL_SECONDS or 900)),
|
|
heartbeat_interval_seconds=max(1, int(settings.REDIS_EXECUTION_LOCK_RENEW_INTERVAL_SECONDS or 20)),
|
|
pipeline_stage=pipeline_stage,
|
|
)
|
|
if lease is None:
|
|
# 重复消息已有其他 Worker 推进，不属于业务失败。
|
|
return
|
|
try:
|
|
await callback(upscale_task_id, lease.token, lease.ensure_owned)
|
|
await lease.ensure_owned()
|
|
finally:
|
|
await lease.close()
|
|
|
|
|
|
async def _run_local(upscale_task_id: str) -> None:
|
|
async def _execute(task_id: str, token: str, guard: Callable[[], Awaitable[None]]) -> None:
|
|
async with async_session() as db:
|
|
await run_local_upscale(db, task_id, execution_token=token, execution_guard=guard)
|
|
await _run_with_execution_lock(
|
|
upscale_task_id,
|
|
_execute,
|
|
task_name=CeleryTaskName.VIDEO_UPSCALE_EXECUTE_LOCAL.value,
|
|
queue=CeleryQueue.GEN_VIDEO_UPSCALE_LOCAL.value,
|
|
pipeline_stage="execute_local",
|
|
)
|
|
|
|
|
|
async def _run_submit(upscale_task_id: str, *, count_attempt: bool = True) -> None:
|
|
async def _execute(task_id: str, token: str, guard: Callable[[], Awaitable[None]]) -> None:
|
|
async with async_session() as db:
|
|
await run_remote_submit(
|
|
db, task_id, count_attempt=count_attempt, execution_token=token, execution_guard=guard
|
|
)
|
|
await _run_with_execution_lock(
|
|
upscale_task_id,
|
|
_execute,
|
|
task_name=CeleryTaskName.VIDEO_UPSCALE_SUBMIT_REMOTE.value,
|
|
queue=CeleryQueue.GEN_VIDEO_UPSCALE_REMOTE.value,
|
|
pipeline_stage="submit_remote",
|
|
)
|
|
|
|
|
|
async def _run_poll(upscale_task_id: str) -> None:
|
|
async def _execute(task_id: str, token: str, guard: Callable[[], Awaitable[None]]) -> None:
|
|
async with async_session() as db:
|
|
await run_remote_poll(db, task_id, execution_token=token, execution_guard=guard)
|
|
await _run_with_execution_lock(
|
|
upscale_task_id,
|
|
_execute,
|
|
task_name=CeleryTaskName.VIDEO_UPSCALE_POLL_REMOTE.value,
|
|
queue=CeleryQueue.GEN_VIDEO_UPSCALE_REMOTE.value,
|
|
pipeline_stage="poll_remote",
|
|
)
|
|
|
|
|
|
async def _run_download(upscale_task_id: str) -> None:
|
|
async def _execute(task_id: str, token: str, guard: Callable[[], Awaitable[None]]) -> None:
|
|
async with async_session() as db:
|
|
await run_remote_result_download(db, task_id, execution_token=token, execution_guard=guard)
|
|
await _run_with_execution_lock(
|
|
upscale_task_id,
|
|
_execute,
|
|
task_name=CeleryTaskName.VIDEO_UPSCALE_DOWNLOAD_REMOTE_RESULT.value,
|
|
queue=CeleryQueue.GEN_VIDEO_UPSCALE_REMOTE.value,
|
|
pipeline_stage="download_remote_result",
|
|
)
|
|
|
|
|
|
async def _run_finalize(upscale_task_id: str) -> None:
|
|
async def _execute(task_id: str, token: str, guard: Callable[[], Awaitable[None]]) -> None:
|
|
async with async_session() as db:
|
|
await run_finalize_upscale(db, task_id, execution_token=token, execution_guard=guard)
|
|
await _run_with_execution_lock(
|
|
upscale_task_id,
|
|
_execute,
|
|
task_name=CeleryTaskName.VIDEO_UPSCALE_FINALIZE.value,
|
|
queue=CeleryQueue.GEN_VIDEO_UPSCALE_LOCAL.value,
|
|
pipeline_stage="finalize",
|
|
)
|
|
|
|
|
|
async def _run_recovery() -> dict[str, Any]:
|
|
lease = await RedisExecutionLockLease.acquire(
|
|
lock_key=settings.VIDEO_UPSCALE_RECOVERY_LOCK_KEY,
|
|
ttl_seconds=max(30, int(settings.VIDEO_UPSCALE_EXECUTION_LOCK_TTL_SECONDS or 900)),
|
|
log_context="video_upscale_recovery",
|
|
renew_interval_seconds=max(1, int(settings.REDIS_EXECUTION_LOCK_RENEW_INTERVAL_SECONDS or 20)),
|
|
)
|
|
if lease is None:
|
|
return {"checked": 0, "results": {"lock_busy": 1}}
|
|
try:
|
|
async with async_session() as db:
|
|
result = await recover_video_upscale_tasks_once(db)
|
|
await lease.ensure_owned()
|
|
return result
|
|
finally:
|
|
await lease.close()
|
|
|
|
|
|
if celery_app:
|
|
|
|
@celery_app.task(name=CeleryTaskName.VIDEO_UPSCALE_EXECUTE_LOCAL.value, bind=True, max_retries=2)
|
|
def execute_local(self, upscale_task_id: str) -> None:
|
|
try:
|
|
return run_async(_run_local(upscale_task_id))
|
|
except Exception as exc:
|
|
raise self.retry(exc=exc, countdown=max(5, int(settings.VIDEO_UPSCALE_RETRY_BACKOFF_SECONDS or 60)))
|
|
|
|
|
|
@celery_app.task(name=CeleryTaskName.VIDEO_UPSCALE_SUBMIT_REMOTE.value, bind=True, max_retries=2)
|
|
def submit_remote(self, upscale_task_id: str, count_attempt: bool = True) -> None:
|
|
try:
|
|
return run_async(_run_submit(upscale_task_id, count_attempt=count_attempt))
|
|
except Exception as exc:
|
|
raise self.retry(exc=exc, countdown=max(5, int(settings.VIDEO_UPSCALE_RETRY_BACKOFF_SECONDS or 60)))
|
|
|
|
|
|
@celery_app.task(name=CeleryTaskName.VIDEO_UPSCALE_POLL_REMOTE.value, bind=True, max_retries=2)
|
|
def poll_remote(self, upscale_task_id: str) -> None:
|
|
try:
|
|
return run_async(_run_poll(upscale_task_id))
|
|
except Exception as exc:
|
|
raise self.retry(exc=exc, countdown=max(5, int(settings.VIDEO_UPSCALE_RETRY_BACKOFF_SECONDS or 60)))
|
|
|
|
|
|
@celery_app.task(name=CeleryTaskName.VIDEO_UPSCALE_DOWNLOAD_REMOTE_RESULT.value, bind=True, max_retries=2)
|
|
def download_remote_result(self, upscale_task_id: str) -> None:
|
|
try:
|
|
return run_async(_run_download(upscale_task_id))
|
|
except Exception as exc:
|
|
raise self.retry(exc=exc, countdown=max(5, int(settings.VIDEO_UPSCALE_RETRY_BACKOFF_SECONDS or 60)))
|
|
|
|
|
|
@celery_app.task(name=CeleryTaskName.VIDEO_UPSCALE_FINALIZE.value, bind=True, max_retries=2)
|
|
def finalize(self, upscale_task_id: str) -> None:
|
|
try:
|
|
return run_async(_run_finalize(upscale_task_id))
|
|
except Exception as exc:
|
|
raise self.retry(exc=exc, countdown=max(5, int(settings.VIDEO_UPSCALE_RETRY_BACKOFF_SECONDS or 60)))
|
|
|
|
|
|
@celery_app.task(name=CeleryTaskName.VIDEO_UPSCALE_RECOVER.value, bind=True, max_retries=2)
|
|
def recover_once(self) -> dict[str, Any]:
|
|
try:
|
|
return run_async(_run_recovery())
|
|
except Exception as exc:
|
|
raise self.retry(exc=exc, countdown=max(5, int(settings.VIDEO_UPSCALE_RETRY_BACKOFF_SECONDS or 60)))
|
|
|
|
else:
|
|
|
|
class _DisabledTask:
|
|
def delay(self, *args: Any, **kwargs: Any) -> None:
|
|
raise RuntimeError("Celery is disabled")
|
|
|
|
def apply_async(self, *args: Any, **kwargs: Any) -> None:
|
|
raise RuntimeError("Celery is disabled")
|
|
|
|
execute_local = _DisabledTask()
|
|
submit_remote = _DisabledTask()
|
|
poll_remote = _DisabledTask()
|
|
download_remote_result = _DisabledTask()
|
|
finalize = _DisabledTask()
|
|
recover_once = _DisabledTask()
|