ONNXRUNTIME-GPU多进程多GPU二次循环阻塞问题求助

1. ORT CUDA资源未彻底释放：OnnxRuntime的CUDA会话（Session）与GPU上下文、显存资源绑定紧密，子进程退出时，ORT内部的底层资源（如CUDA上下文锁、线程池状态）可能未被彻底清理。第一次循环结束后，残留的资源锁会导致第二次循环启动的新进程无法正常获取GPU资源，进而阻塞。
2. 序列化传递Session无效：ORT的Session包含底层C++资源句柄，无法通过pickle正确序列化和反序列化。即使强行传递，也会导致资源状态异常，无法解决阻塞问题。

1. 子进程内显式清理ORT资源 #

在子进程任务结束后，主动销毁Session实例并触发垃圾回收，确保GPU资源被释放：

def impl(model_path, ctx, num_gpus):
    device_id = ctx % num_gpus
    model = Onnx(model_path, device_id)
    # 执行模型推理等业务逻辑
    
    # 显式清理资源
    del model.session
    del model
    import gc
    gc.collect()

2. 改用进程池复用进程 #

避免每次循环重复创建销毁进程，用multiprocessing.Pool复用进程，既减少资源泄漏风险，又提升性能：

from multiprocessing import Pool

# 每个进程仅初始化一次模型，全局变量存储实例
_model = None

def init_process(model_path, device_id):
    global _model
    options = ort.SessionOptions()
    options.intra_op_num_threads = 1
    options.inter_op_num_threads = 1
    options.enable_mem_pattern = False  # 禁用内存池复用
    providers = [('CUDAExecutionProvider', {'device_id': device_id}), 'CPUExecutionProvider']
    _model = ort.InferenceSession(model_path, options, providers=providers)

def impl(_):
    # 复用已初始化的模型执行推理
    # outputs = _model.run(None, input_dict)
    pass

def main():
    num_workers = 4
    num_gpus = 2
    n = 2
    # 为每个进程分配固定GPU设备
    init_args_list = [(model_path, i % num_gpus) for i in range(num_workers)]
    
    with Pool(num_workers) as pool:
        # 批量初始化每个进程的模型
        pool.starmap(init_process, init_args_list)
        
        for i in range(n):
            print("START LOOP {}".format(i))
            # 提交任务到复用的进程
            pool.map(impl, range(num_workers))
            print("END LOOP {}".format(i))

3. 优化ORT Session配置 #

创建Session时增加资源限制选项，减少资源锁定风险：

class Onnx:
    def __init__(self, model_path, device_id):
        options = ort.SessionOptions()
        options.intra_op_num_threads = 1
        options.inter_op_num_threads = 1
        # 禁用内存池和CPU内存 arena，避免跨进程资源冲突
        options.enable_mem_pattern = False
        options.enable_cpu_mem_arena = False
        # 禁用CUDA图，防止资源锁定
        providers = [('CUDAExecutionProvider', {
            'device_id': device_id,
            'enable_cuda_graph': False
        }), 'CPUExecutionProvider']
        self.session = ort.InferenceSession(model_path, options, providers=providers)

4. 排查GPU资源泄漏 #

在循环间隙执行nvidia-smi命令，检查显存是否完全释放。若第一次循环后显存未归零，说明存在资源泄漏，需进一步检查资源清理逻辑。

内容的提问来源于stack exchange，提问作者walstruzz