多线程安全推理与并发设计

多线程安全推理与并发设计

学习目标

读完本章后，你将能够：

• 理解 MinerU 线程安全优化的设计动机
• 掌握模型单例的线程安全实现
• 分析 API 服务的并发处理流程
• 设计多线程并发推理方案

前置知识

• 三引擎后端的源码实现与调度
• 异步 API 与任务编排

---

1. 线程安全优化的背景

在 v3.0 之前，MinerU 的模型推理存在线程安全问题：

# 线程不安全的模型调用
class UnsafeModel:
    def predict(self, image):
        # 多线程同时调用可能导致状态混乱
        self._internal_state = preprocess(image)
        return self._model(self._internal_state)

问题场景：

• FastAPI 默认使用异步事件循环处理请求
• 多个并发请求可能同时调用模型
• PyTorch 模型本身线程安全，但 MinerU 的预处理/后处理逻辑不是

2. 线程安全实现

2.1 RLock 保护模型访问

mineru/backend/pipeline/pipeline_analyze.py

class ModelSingleton:
    _instance = None
    _models = {}
    _lock = threading.RLock()  # 可重入锁

    def get_model(self, lang=None, formula_enable=None, table_enable=None):
        key = (lang, formula_enable, table_enable)
        with self._lock:
            if key not in self._models:
                self._models[key] = custom_model_init(...)
        return self._models[key]

为什么用 RLock 而不是 Lock？

• RLock（可重入锁）允许同一线程多次获取锁
• 模型初始化过程中可能嵌套调用，Lock 会导致死锁

2.2 推理阶段的并发控制

# VLM 引擎的执行保护
def predictor_execution_guard(func):
    """装饰器：保护 VLM 推理的线程安全"""
    _inference_lock = threading.Lock()

    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        with _inference_lock:
            return func(*args, **kwargs)
    return wrapper

# 异步版本
async def aio_predictor_execution_guard(func, *args, **kwargs):
    _aio_lock = asyncio.Lock()
    async with _aio_lock:
        return await func(*args, **kwargs)

2.3 PDFium 的线程安全

PDF 渲染（pypdfium2）需要特殊处理：

mineru/utils/pdfium_guard.py

def open_pdfium_document(pdf_bytes) -> pdfium.PdfDocument:
    """线程安全地打开 PDF 文档"""
    # 每个线程使用独立的 PDF 文档实例
    # 避免共享文档导致的状态冲突
    return pdfium.PdfDocument(pdf_bytes)

def close_pdfium_document(doc):
    """安全关闭 PDF 文档"""
    try:
        doc.close()
    except Exception:
        pass  # 忽略关闭时的异常

3. 并发处理流程

3.1 FastAPI 并发模型

关键设计：

• FastAPI 异步接收请求
• GPU 推理串行化（通过 Lock）
• CPU 预处理/后处理可并行

3.2 并发数配置

mineru/cli/api_protocol.py

DEFAULT_MAX_CONCURRENT_REQUESTS = 3  # 默认最大并发数

通过配置文件调整：

{
  "max_concurrent_requests": 5
}

配置建议：

• GPU 推理瓶颈时：设为 1-3
• CPU 充足时：可适当增加

4. mineru-router 多卡部署

mineru-router 通过多进程实现多 GPU 负载均衡：

mineru/cli/router.py

def main():
    workers = args.workers  # GPU 数量

    # 启动多个 Worker 进程
    processes = []
    for i in range(workers):
        # 每个 Worker 绑定到不同的 GPU
        env = os.environ.copy()
        env["CUDA_VISIBLE_DEVICES"] = str(i)

        proc = subprocess.Popen(
            ["mineru-api", "--port", str(8888 + i)],
            env=env,
        )
        processes.append(proc)

    # 启动负载均衡代理
    start_load_balancer(
        workers=[f"127.0.0.1:{8888 + i}" for i in range(workers)],
    )

负载均衡策略：

• 轮询：按顺序分配请求到 Worker
• 最少连接：分配给当前任务最少的 Worker
• 接口兼容：与 mineru-api 完全一致的 API

5. 陷阱与对策

陷阱	对策
多线程同时初始化模型	使用 RLock 保护 get_model 方法
PyTorch CUDA 并发冲突	通过 Lock 序列化推理调用
PDFium 文档对象共享	每个线程使用独立文档实例
并发数过高导致 GPU OOM	调整 max_concurrent_requests 配置
后台任务清理不及时	使用 BackgroundTasks 确保清理

6. 设计权衡

序列化推理 vs 并行推理

方案	优势	劣势
序列化推理（Lock）	简单可靠，无并发冲突	GPU 利用率低
并行推理	GPU 利用率高	需要模型支持批量推理

取舍：MinerU 默认序列化推理（安全优先），但通过 batch 推理和 mineru-router 实现并发加速。

多进程 vs 多线程

方案	优势	劣势
多线程	共享内存，通信简单	GIL 限制，Python 并发效率低
多进程	真正的并行，无 GIL	进程间通信开销

取舍：mineru-router 使用多进程方案，每个 Worker 进程独占一个 GPU，避免 GPU 资源竞争。

---

参考来源

• MinerU 源码：mineru/backend/pipeline/pipeline_analyze.py、mineru/backend/vlm/vlm_analyze.py、mineru/cli/router.py
• Python threading 文档：https://docs.python.org/3/library/threading.html
• MinerU CHANGELOG v3.0：线程安全优化