diff --git a/mineru/backend/pipeline/model_init.py b/mineru/backend/pipeline/model_init.py index 2006402a..f696dc00 100644 --- a/mineru/backend/pipeline/model_init.py +++ b/mineru/backend/pipeline/model_init.py @@ -1,6 +1,7 @@ # Copyright (c) Opendatalab. All rights reserved. import os import threading +from contextlib import contextmanager import torch from loguru import logger @@ -24,12 +25,39 @@ PIPELINE_LAYOUT_INFERENCE_LOCK = threading.RLock() PIPELINE_MFR_INFERENCE_LOCK = threading.RLock() PIPELINE_OCR_DET_INFERENCE_LOCK = threading.RLock() PIPELINE_OCR_REC_INFERENCE_LOCK = threading.RLock() +# 推理准入门闩:清理线程等待静默区时,阻止新的推理调用继续插队进入阶段锁。 +PIPELINE_INFERENCE_ADMISSION_LOCK = threading.RLock() +PIPELINE_INFERENCE_STAGE_LOCKS = ( + PIPELINE_LAYOUT_INFERENCE_LOCK, + PIPELINE_MFR_INFERENCE_LOCK, + PIPELINE_OCR_DET_INFERENCE_LOCK, + PIPELINE_OCR_REC_INFERENCE_LOCK, +) def _run_with_inference_lock(inference_lock, inference_callable, *args, **kwargs): """在指定推理锁内执行真实 native 模型调用。""" - with inference_lock: + with PIPELINE_INFERENCE_ADMISSION_LOCK: + inference_lock.acquire() + try: return inference_callable(*args, **kwargs) + finally: + inference_lock.release() + + +@contextmanager +def pipeline_inference_quiet_zone(): + """进入 pipeline/hybrid 共享模型推理静默区,等待四类阶段推理全部退出。""" + acquired_locks = [] + with PIPELINE_INFERENCE_ADMISSION_LOCK: + try: + for inference_lock in PIPELINE_INFERENCE_STAGE_LOCKS: + inference_lock.acquire() + acquired_locks.append(inference_lock) + yield + finally: + for inference_lock in reversed(acquired_locks): + inference_lock.release() def run_layout_inference(inference_callable, *args, **kwargs): diff --git a/mineru/utils/model_utils.py b/mineru/utils/model_utils.py index 6c1f486a..1c2e4f68 100644 --- a/mineru/utils/model_utils.py +++ b/mineru/utils/model_utils.py @@ -180,7 +180,15 @@ def get_res_list_from_layout_res(layout_res, overlap_threshold=0.8): return ocr_res_list, table_res_list, single_page_mfdetrec_res -def clean_memory(device='cuda'): +def _pipeline_inference_quiet_zone(): + """懒加载推理静默区,避免 utils 模块导入时和 pipeline backend 形成循环依赖。""" + from mineru.backend.pipeline.model_init import pipeline_inference_quiet_zone + + return pipeline_inference_quiet_zone() + + +def _clean_memory_without_inference(device='cuda'): + """执行实际显存/内存清理动作,调用方负责进入推理静默区。""" if str(device).startswith("cuda"): if torch.cuda.is_available(): torch.cuda.empty_cache() @@ -205,6 +213,11 @@ def clean_memory(device='cuda'): gc.collect() +def clean_memory(device='cuda'): + with _pipeline_inference_quiet_zone(): + _clean_memory_without_inference(device) + + def clean_vram(device, vram_threshold=8): total_memory = get_vram(device) if total_memory and total_memory <= vram_threshold: