x

src/moldinsight/services/task_dispatcher.py

@@ -0,0 +1,51 @@
+# 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, file_path: 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, file_path, stp_file_id, process_params
+        )
+
+
+def dispatch_processing(task_id: str, file_path: str, stp_file_id: int, process_params: dict):
+    """调度 STP 处理任务：优先 Celery（进程隔离），否则 API 进程内 asyncio 后台执行。"""
+    if _use_celery:
+        process_stp_task.delay(task_id, file_path, stp_file_id, process_params)
+        logger.info(f"[DISPATCH] Celery 任务已调度: task_id={task_id}")
+        return
+
+    task = asyncio.create_task(
+        _run_with_limit(task_id, file_path, 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 未安装)")