57 lines
2.2 KiB
Python
57 lines
2.2 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, 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 未安装)")
|