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

from __future__ import annotations

import logging
from typing import Any

from app.config import settings
from app.enums.celery_queue import CeleryTaskName
from app.enums.celery_runtime import CeleryRuntimeEvent
from app.services.celery_runtime.recovery_service import (
    garbage_collect_registry_pair,
    garbage_collect_worker_registry,
    guard_periodic_recovery,
)
from app.services.operation_log_service import log_operation_event
from app.services.redis_registry_service import RedisExecutionLockLease
from app.tasks.async_runner import run_async
from app.tasks.celery_app import celery_app

logger = logging.getLogger("video_gen")


async def _run_reconcile_once() -> dict[str, Any]:
    guarded = await guard_periodic_recovery(check_global_lock=False)
    if guarded:
        return guarded
    lease = await RedisExecutionLockLease.acquire(
        lock_key=settings.CELERY_RUNTIME_GLOBAL_RECOVERY_LOCK_KEY,
        ttl_seconds=int(settings.CELERY_RECOVERY_TASK_LOCK_TTL_SECONDS or 600),
        log_context="celery_runtime_reconcile",
        renew_interval_seconds=int(settings.REDIS_EXECUTION_LOCK_RENEW_INTERVAL_SECONDS or 30),
    )
    if lease is None:
        return {"skipped": "global_recovery_lock_held"}
    try:
        from app.tasks.generation_recovery_tasks import (
            _run_create_once,
            _run_download_once,
            _run_module_async_once,
            _run_shot_analysis_once,
            _run_shot_split_once,
        )

        results: dict[str, Any] = {}
        for name, runner in (
            ("create", _run_create_once),
            ("download", _run_download_once),
            ("module_async", _run_module_async_once),
            ("shot_analysis", _run_shot_analysis_once),
            ("shot_split", _run_shot_split_once),
        ):
            try:
                results[name] = await runner()
            except Exception as exc:
                logger.exception("Celery runtime reconcile step failed. step=%s", name)
                results[name] = {"error": str(exc)}
        await lease.ensure_owned()
        log_operation_event(
            domain="celery_runtime",
            event_type=CeleryRuntimeEvent.REGISTRY_RECONCILE_DONE.value,
            event_status="success",
            source="recovery",
            detail={"steps": results},
        )
        return {"steps": results}
    finally:
        await lease.close()


async def _run_registry_gc_once() -> dict[str, Any]:
    guarded = await guard_periodic_recovery()
    if guarded:
        return guarded
    pairs = {
        "generation_create": (
            settings.GENERATION_CREATE_ACTIVE_REDIS_HASH_KEY,
            settings.GENERATION_CREATE_ACTIVE_REDIS_ZSET_KEY,
        ),
        "poll": (settings.POLL_ACTIVE_REDIS_HASH_KEY, settings.POLL_ACTIVE_REDIS_ZSET_KEY),
        "download": (settings.DOWNLOAD_ACTIVE_REDIS_HASH_KEY, settings.DOWNLOAD_ACTIVE_REDIS_ZSET_KEY),
        "module_async": (
            settings.MODULE_ASYNC_ACTIVE_REDIS_HASH_KEY,
            settings.MODULE_ASYNC_ACTIVE_REDIS_ZSET_KEY,
        ),
        "shot_analysis": (
            settings.SHOT_ANALYSIS_ACTIVE_REDIS_HASH_KEY,
            settings.SHOT_ANALYSIS_ACTIVE_REDIS_ZSET_KEY,
        ),
        "shot_split": (
            settings.SHOT_SPLIT_ACTIVE_REDIS_HASH_KEY,
            settings.SHOT_SPLIT_ACTIVE_REDIS_ZSET_KEY,
        ),
        "video_upscale": (
            settings.VIDEO_UPSCALE_ACTIVE_REDIS_HASH_KEY,
            settings.VIDEO_UPSCALE_ACTIVE_REDIS_ZSET_KEY,
        ),
        "private_portrait_poll": (
            settings.PRIVATE_PORTRAIT_POLL_ACTIVE_REDIS_HASH_KEY,
            settings.PRIVATE_PORTRAIT_POLL_ACTIVE_REDIS_ZSET_KEY,
        ),
        "private_portrait_delete": (
            settings.PRIVATE_PORTRAIT_DELETE_ACTIVE_REDIS_HASH_KEY,
            settings.PRIVATE_PORTRAIT_DELETE_ACTIVE_REDIS_ZSET_KEY,
        ),
    }
    results = {}
    for name, (hash_key, zset_key) in pairs.items():
        results[name] = await garbage_collect_registry_pair(
            hash_key=hash_key,
            zset_key=zset_key,
            limit=int(settings.CELERY_RUNTIME_GC_BATCH_SIZE or 500),
        )
    results["worker_registry_v2"] = await garbage_collect_worker_registry(
        limit=int(settings.CELERY_RUNTIME_GC_BATCH_SIZE or 500),
    )
    log_operation_event(
        domain="celery_runtime",
        event_type=CeleryRuntimeEvent.REGISTRY_GC_DONE.value,
        event_status="success",
        source="recovery",
        detail={"results": results},
    )
    return {"results": results}


if celery_app:

    @celery_app.task(
        name=CeleryTaskName.CELERY_RUNTIME_RECONCILE.value,
        bind=True,
        soft_time_limit=settings.CELERY_RECOVERY_SOFT_TIME_LIMIT_SECONDS,
        time_limit=settings.CELERY_RECOVERY_TIME_LIMIT_SECONDS,
        ignore_result=True,
    )
    def reconcile_once(self) -> dict[str, Any]:
        return run_async(_run_reconcile_once())


    @celery_app.task(
        name=CeleryTaskName.CELERY_RUNTIME_GC.value,
        bind=True,
        soft_time_limit=settings.CELERY_RECOVERY_SOFT_TIME_LIMIT_SECONDS,
        time_limit=settings.CELERY_RECOVERY_TIME_LIMIT_SECONDS,
        ignore_result=True,
    )
    def registry_gc_once(self) -> dict[str, Any]:
        return run_async(_run_registry_gc_once())
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")

    reconcile_once = _DisabledTask()
    registry_gc_once = _DisabledTask()