52 lines
2.1 KiB
Python
52 lines
2.1 KiB
Python
|
# 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 未安装)")
|