life-echo/api/app/tasks/celery_app.py

"""
Celery 应用配置
配置从 app.core.config.settings 读取。
Worker 启动时需聚合注册所有 feature 的 model，否则 User 等 relationship("Order", ...) 解析时会报找不到 Order。
与 main.py / Alembic 一致：下方 import 仅用于注册 ORM model。
"""

from __future__ import annotations

from typing import Any

from app.core.logging import get_logger, setup_logging

# 与 app.main 一致：先配置 loguru + InterceptHandler，再加载会打日志的依赖
setup_logging()

from celery import Celery
from celery.signals import task_failure, task_postrun, task_prerun, task_success

from app.core.celery_log_context import clear_celery_log_extras, set_celery_log_extras
from app.core.config import settings
from app.core.log_events import celery_prerun_extras
from app.features.asset import models as _asset_models  # noqa: F401 - register Asset
from app.features.auth import models as _auth_models  # noqa: F401
from app.features.conversation import models as _conv_models  # noqa: F401
from app.features.evaluation import models as _eval_models  # noqa: F401
from app.features.memoir import models as _memoir_models  # noqa: F401
from app.features.memory import models as _memory_models  # noqa: F401
from app.features.payment import models as _payment_models  # noqa: F401
from app.features.story import models as _story_models  # noqa: F401
from app.features.user import models as _user_models  # noqa: F401

REDIS_URL = settings.redis_url

# 创建 Celery 应用
celery_app = Celery(
    "life_echo",
    broker=REDIS_URL,
    backend=REDIS_URL,
    include=[
        "app.tasks.memoir_tasks",
        "app.tasks.story_title_tasks",
        "app.tasks.story_image_tasks",
        "app.tasks.chapter_cover_tasks",
        "app.tasks.chapter_compose_tasks",
        "app.tasks.memory_compaction_tasks",
        "app.tasks.memory_enrichment_tasks",
        "app.tasks.memoir_quality_pass_tasks",
    ],
)

# Celery 配置
celery_app.conf.update(
    # 不劫持根 logger，便于与 loguru + InterceptHandler 统一格式与等级
    worker_hijack_root_logger=False,
    # 任务序列化
    task_serializer="json",
    accept_content=["json"],
    result_serializer="json",
    # 时区
    timezone="UTC",
    enable_utc=True,
    # 任务结果过期时间（1小时）
    result_expires=3600,
    # 任务执行设置
    task_soft_time_limit=300,  # 5分钟软超时
    task_time_limit=600,  # 10分钟硬超时
    # 并发设置
    worker_prefetch_multiplier=1,  # 每次只预取一个任务
    worker_concurrency=4,  # 并发 worker 数量
    # 任务重试设置
    task_acks_late=True,  # 任务完成后再确认
    task_reject_on_worker_lost=True,  # worker 丢失时拒绝任务
    task_routes={
        "app.tasks.memory_enrichment_tasks.enrich_memory_source": {
            "queue": settings.celery_memory_enrichment_queue,
        },
    },
)

celery_app.conf.task_annotations = {
    "app.tasks.memory_enrichment_tasks.enrich_memory_source": {
        "soft_time_limit": 660,
        "time_limit": 960,
    },
}

celery_app.conf.beat_schedule = {
    "memory-compaction-sweep": {
        "task": "app.tasks.memory_compaction_tasks.memory_compaction_sweep",
        "schedule": 6 * 3600.0,
    },
}

_celery_lifecycle_log = get_logger(__name__)


def _summarize_task_return(retval: object) -> str:
    if retval is None:
        return "None"
    if isinstance(retval, dict):
        keys = list(retval.keys())[:14]
        return "dict:" + ",".join(str(k) for k in keys)
    text = repr(retval)
    if len(text) > 180:
        return text[:180] + "..."
    return text


@task_prerun.connect
def _log_task_prerun(
    task_id: str | None = None,
    task: object | None = None,
    args: tuple[Any, ...] | None = None,
    kwargs: dict[str, Any] | None = None,
    **_: object,
) -> None:
    name = getattr(task, "name", None) or "?"
    extras = celery_prerun_extras(name, tuple(args or ()), dict(kwargs or {}))
    if task_id:
        extras["task_id"] = str(task_id).strip()
    set_celery_log_extras(extras if extras else None)
    _celery_lifecycle_log.info(
        "event=celery_task_start task={} task_id={} msg=Celery 任务已开始",
        name,
        task_id,
    )


@task_success.connect
def _log_task_success(
    sender: object | None = None, result: object | None = None, **_: object
) -> None:
    """仅成功路径；失败见 ``task_failure``（避免 ``task_postrun`` 在异常态仍触发）。"""
    name = getattr(sender, "name", None) if sender is not None else None
    name = name or "?"
    task_id: str | None = None
    if sender is not None:
        req = getattr(sender, "request", None)
        if req is not None:
            task_id = getattr(req, "id", None)
    _celery_lifecycle_log.info(
        "event=celery_task_ok task={} task_id={} result={} msg=Celery 任务已成功结束",
        name,
        task_id,
        _summarize_task_return(result),
    )


@task_failure.connect
def _log_task_failure(
    task_id: str | None = None,
    task: object | None = None,
    exception: BaseException | None = None,
    **kwargs: object,
) -> None:
    name = getattr(task, "name", None) or "?"
    et = type(exception).__name__ if isinstance(exception, BaseException) else "?"
    _celery_lifecycle_log.warning(
        "event=celery_task_failed task={} task_id={} exc_type={} exc={} msg=Celery 任务失败",
        name,
        task_id,
        et,
        exception,
    )


@task_postrun.connect
def _clear_worker_log_context(**kwargs: object) -> None:
    """任务体结束后清除 ContextVar，避免同一 worker 进程串上下文。"""
    clear_celery_log_extras()