geMoldInsight/src/moldinsight/services/task_dispatcher.py

# services/task_dispatcher.py
"""后台处理任务分派器 - 统一 upload/batch 路由的 Celery/asyncio 分派逻辑。

修复两个问题：
1. fire-and-forget：asyncio.create_task 返回值未持有引用，任务可能被 GC 中途回收，
   异常也无从浮现（python 官方文档明确警告的模式）；
2. 复制粘贴：upload_router 与 batch_router 各自维护一份相同的分派代码，易漂移。
"""
import asyncio

from shared.utils.logger import get_logger

logger = get_logger(__name__)

try:
    from celery_tasks import process_stp_task
    _use_celery = True
except ImportError:
    process_stp_task = None
    _use_celery = False

# 持有后台任务强引用，防止被 GC 回收；完成后自动移出
_background_tasks: set = set()

# API 进程内并发处理上限（celery 路径由 worker 并发数控制，不走这里）。
# asyncio.Semaphore 自 3.10 起惰性绑定事件循环，模块级创建安全；
# 本模块仅在 API 进程（单一事件循环）导入使用。
_dispatch_semaphore = asyncio.Semaphore(2)


async def _run_with_limit(task_id: str, stp_file_id: int, process_params: dict):
    async with _dispatch_semaphore:
        from moldinsight.services.processing_service import processing_service
        await processing_service.process_file_with_storage(
            task_id, stp_file_id, process_params
        )


def dispatch_processing(task_id: str, stp_file_id: int, process_params: dict):
    """调度 STP 处理任务：优先 Celery（进程隔离），否则 API 进程内 asyncio 后台执行。

    入参只传 stp_file_id（D6）：源文件由处理方按 PG 元数据从 RustFS 获取，
    不再跨进程传节点本地路径——API 与 Celery worker 容器文件系统不互通，
    传路径在容器化部署下必然失败。
    """
    if _use_celery:
        process_stp_task.delay(task_id, stp_file_id, process_params)
        logger.info(f"[DISPATCH] Celery 任务已调度: task_id={task_id}")
        return

    task = asyncio.create_task(
        _run_with_limit(task_id, stp_file_id, process_params)
    )
    _background_tasks.add(task)
    task.add_done_callback(_background_tasks.discard)
    logger.info(f"[DISPATCH] 进程内后台处理: task_id={task_id} (celery 未安装)")