110 lines
3.8 KiB
Python
110 lines
3.8 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
from app.config import settings
|
|
from app.models.base import async_session
|
|
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_local(upscale_task_id: str) -> None:
|
|
async with async_session() as db:
|
|
await run_local_upscale(db, upscale_task_id)
|
|
|
|
|
|
async def _run_submit(upscale_task_id: str, *, count_attempt: bool = True) -> None:
|
|
async with async_session() as db:
|
|
await run_remote_submit(db, upscale_task_id, count_attempt=count_attempt)
|
|
|
|
|
|
async def _run_poll(upscale_task_id: str) -> None:
|
|
async with async_session() as db:
|
|
await run_remote_poll(db, upscale_task_id)
|
|
|
|
|
|
async def _run_download(upscale_task_id: str) -> None:
|
|
async with async_session() as db:
|
|
await run_remote_result_download(db, upscale_task_id)
|
|
|
|
|
|
async def _run_finalize(upscale_task_id: str) -> None:
|
|
async with async_session() as db:
|
|
await run_finalize_upscale(db, upscale_task_id)
|
|
|
|
|
|
async def _run_recovery() -> dict[str, Any]:
|
|
async with async_session() as db:
|
|
return await recover_video_upscale_tasks_once(db)
|
|
|
|
|
|
if celery_app:
|
|
|
|
@celery_app.task(name="video_upscale.execute_local", 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="video_upscale.submit_remote", 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="video_upscale.poll_remote", 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="video_upscale.download_remote_result", 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="video_upscale.finalize", 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="video_upscale.recover_once", bind=True)
|
|
def recover_once(self) -> dict[str, Any]:
|
|
return run_async(_run_recovery())
|
|
|
|
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()
|