1. 背景与动机1.1 PyTorch 张量共享的困境在使用 Python 标准库进行并行计算时PyTorch 开发者面临着严峻的挑战。让我们从实际问题说起。问题一标准库进程池的序列化开销当使用multiprocessing.Pool或concurrent.futures.ProcessPoolExecutor传递 PyTorch 张量时会发生完整的序列化 - 反序列化过程frommultiprocessingimportPoolimporttorchdefprocess_tensor(tensor):returntensor.sum()if__name____main__:large_tensortorch.randn(10000,1000)# 约 40MBwithPool(processes4)aspool:# 每次提交都会序列化整个张量futures[pool.apply_async(process_tensor,args(large_tensor,))for_inrange(10)]results[f.get()forfinfutures]严重性分析每次任务提交都需要将 40MB 张量序列化跨进程传输序列化数据子进程反序列化重建张量对于大规模模型训练这种开销可能是灾难性的问题二PyTorch 官方进程池的已知 BugPyTorch 官方提供了torch.multiprocessing.pool.Pool作为解决方案但实际使用时会直接报错fromtorch.multiprocessing.poolimportPool poolPool()错误信息TypeError: SimpleQueue.__init__() missing 1 required keyword-only argument: ctx这是 PyTorch 的一个已知 bug。该问题于 2019 年 2 月提出torch.multiprocessing.pool.Pool broken至今尚未解决。但当你尝试使用torch.multiprocessing.Pool注意不是刚才的torch.multiprocessing.pool.Pool时发现它是可用的但遗憾的是torch.multiprocessing.Pool仅仅是multiprocessing.Pool的直接链接依然不能实现 tensor 共享。fromtorch.multiprocessingimportPoolasPool1frommultiprocessingimportPoolasPool2print(Pool1isPool2)# 输出True这意味着torch.multiprocessing.Pool并没有解决 PyTorch 张量的跨进程共享问题它只是multiprocessing.Pool的一个别名。1.2 ONNX Runtime 推理的挑战除了 PyTorch 训练场景ONNX Runtime 推理同样面临并行化的难题。问题多线程共享 InferenceSession 的不安全性与 Provider 选择困境ONNX Runtime 的InferenceSession对象内部维护了大量状态如 CUDA 流、DML 命令列表等。在多线程环境下共享同一个InferenceSession实例并同时调用run()会导致数据竞争和内存损坏importonnxruntimeasortfromconcurrent.futuresimportThreadPoolExecutor sessort.InferenceSession(model.onnx)definfer(input_data):returnsess.run(None,{input:input_data})# 多线程共享同一个 session → 内存损坏withThreadPoolExecutor(max_workers4)aspool:resultslist(pool.map(infer,inputs))即使规避了线程安全问题传统方案还需要手动指定 ExecutionProvider# 必须人工指定 provider无法根据负载动态切换sessort.InferenceSession(model.onnx,providers[CUDAExecutionProvider])# 或sessort.InferenceSession(model.onnx,providers[DmlExecutionProvider])# 或sessort.InferenceSession(model.onnx,providers[CPUExecutionProvider])手动指定 provider 的问题选中 CUDA 但显存不足时 → OOM 崩溃选中 DML 但 GPU 繁忙时 → 任务排队等待CPU 空闲浪费选中 CPU 但 GPU 空闲时 → 硬件利用率低下无法根据系统实时负载动态切换 provider导致硬件饥饿或过载传统解决方案及其缺陷方案缺陷每个线程独立创建 Session每次创建加载模型权重显存/内存开销成倍增长全局锁保护 Session退化为串行推理失去多线程优势手动管理 Session 池代码复杂度高缺少资源感知和自动设备调度1.3 SmartPool真正的解决方案面对标准库的序列化开销和 PyTorch 官方的未解 bug以及 ONNX Runtime 多线程推理的陷阱开发者急需一个统一的解决方案。SmartPool 应运而生。针对 PyTorchSmartPool 通过use_torchTrue启用 PyTorch 支持实现真正的张量共享fromsmartpoolimportProcessPoolimporttorch# 这个函数应该写在其他模块不能写在主模块中defprocess_tensor(tensor):# 零序列化开销张量通过共享内存传递returntensor.sum()if__name____main__:large_tensortorch.randn(10000,1000)# 约 40MBlarge_tensor.share_memory_()withProcessPool(use_torchTrue)aspool:# 张量通过共享内存传递无序列化开销futures[pool.submit(process_tensor,args(large_tensor,))for_inrange(10)]results[f.result()forfinfutures]针对 ONNX RuntimeSmartPool 提供InferSessionPool自动管理线程安全的推理会话fromsmartpoolimportInferSessionPoolwithInferSessionPool()aspool:futurepool.submit(model.onnx,args(input_tensor,))resultfuture.result()这仅仅是 SmartPool 的开始。无论是 PyTorch 训练还是 ONNX Runtime 推理SmartPool 都提供了一系列智能资源管理特性。2. 超越张量共享智能资源管理2.1 动态资源感知调度传统进程池通过max_workers参数控制并发进程数从而间接控制内存使用fromconcurrent.futuresimportProcessPoolExecutor# 静态限制始终最多 4 个进程withProcessPoolExecutor(max_workers4)aspool:resultspool.map(memory_intensive_func,data_list)问题分析不同任务的内存使用情况不一为防止太多内存消耗大的任务同时运行导致 OOM设置processes4限制最多启动的进程数。但当内存消耗大的任务完成退出进程池后进程池处于内存使用低谷阶段如仅占用 500MB时无法启动更多任务剩余的 CPU 和内存被闲置。这种静态配置存在三个核心问题过于保守内存峰值过去后空闲资源无法被利用缺乏弹性无法根据实际内存使用情况动态调整并发度资源浪费在多 GPU 环境下静态配置无法适应动态变化的负载分布SmartPool 的动态调度方案SmartPool 允许任务声明资源需求系统根据实时资源状态动态调度fromsmartpoolimportProcessPool,DataSize,Resourcedefmemory_intensive_task(data):importnumpyasnp big_arraynp.zeros((10000,10000))# 约 800MB# 处理数据...returnresultif__name____main__:withProcessPool(use_torchTrue)aspool:# 声明资源需求futurepool.submit(memory_intensive_task,args(dataset,),cpu_mode_resResource(cpu_cores_in_python2,# 需要 2 个 CPU 核心cpu_mem4*DataSize.GB,# 需要 4GB 系统内存),gpu_mode_resResource(cpu_cores_in_python1,cpu_mem500*DataSize.MB,gpu_cores1024,# 需要 1024 个 CUDA 核心gpu_mem2*DataSize.GB,# 需要 2GB GPU 显存),)# 若当前资源不足任务自动进入等待队列# 待其他任务释放资源后自动启动执行resultfuture.result()与传统方案的对比特性传统进程池max_workersSmartPool动态调度并发控制固定进程数动态资源配额内存管理静态上限按需分配弹性伸缩资源利用率保守估计资源闲置充分利用并防止耗尽适应性无法应对负载波动实时监控动态调整配置复杂度需手动估算自动管理优势分析弹性伸缩当任务处于内存低谷期自动启动更多任务填充空闲资源防 OOM 机制通过资源声明和监控确保总消耗不超过系统容量细粒度控制可为每个任务单独指定 CPU、内存、GPU 需求2.2 异构负载均衡现代计算节点通常包含多核 CPU 和多块 GPU如何均匀分配任务成为关键问题典型场景4 核 CPU 2 块 GPU每块 GPU 有多个 CUDA 核心需要同时训练 20 个模型任务可能分配到 CPU 或任意一块 GPU传统方案的不足手动调度复杂需要显式管理每个任务的设备分配资源竞争多个任务可能竞争同一块 GPU导致显存溢出负载不均某些 GPU 过载而其他 GPU 空闲缺乏统一抽象CPU 和 GPU 任务需要不同的代码路径SmartPool 的统一调度方案SmartPool 统一管理 CPU 和 GPU 资源自动将任务分配到最优设备任务可通过best_device()函数获取到分配到的设备fromsmartpoolimportProcessPool,DataSize,Resource,best_devicedeftrain_model(model_config):devicebest_device()# 自动获取最优设备# 在最优设备上执行训练returntrain_on_device(model_config,device)if__name____main__:configs[...]# 20 个模型配置withProcessPool(use_torchTrue)aspool:futures[]forconfiginconfigs:futurepool.submit(train_model,args(config,),cpu_mode_resResource(cpu_cores_in_python1,# 需要 1 个 CPU 核心cpu_mem1*DataSize.GB,# 需要 1GB 系统内存),gpu_mode_resResource(cpu_cores_in_python1,cpu_mem500*DataSize.MB,gpu_cores1024,# 需要 1024 个 CUDA 核心gpu_mem1*DataSize.GB,# 需要 1GB GPU 显存),)# SmartPool 自动均衡分配到 CPU 和多块 GPUfutures.append(future)results[f.result()forfinfutures]2.3 训练热迁移在动态资源环境中GPU 可用性可能随时间变化场景描述初始状态所有 GPU 被占用部分训练任务在 CPU 上启动时间点 T1某个 GPU 上的任务快速完成释放 GPU问题仍在 CPU 上训练的任务继续占用 CPU新释放的 GPU 空闲结果硬件资源浪费整体吞吐量下降理想方案应能检测到 GPU 空闲并将 CPU 上的训练任务自动迁移到 GPU无需重启训练过程。SmartPool 的热迁移方案SmartPool 支持在训练过程中动态检测设备变化并自动迁移任务fromsmartpoolimportProcessPool,best_device,move_optimizer_toimporttorchdeftraining_task(model,optimizer,data_loader):# 初始化时获取设备devicebest_device()old_devicedevice model.to(device)forepochinrange(epochs):forbatch_idx,(data,target)inenumerate(data_loader):# 每个 batch 检查设备变化devicebest_device()data,targetdata.to(device),target.to(device)ifdevice!old_device:# 迁移模型和优化器model.to(device)move_optimizer_to(optimizer,device)old_devicedevice# 在新设备上继续训练outputmodel(data)losscriterion(output,target)# ... 反向传播returnmodelif__name____main__:withProcessPool(use_torchTrue)aspool:# 初始可能在 CPU 上训练futurepool.submit(training_task,args(model,optimizer,data))# 当 GPU 空闲时自动迁移到 GPUresultfuture.result()工作流程初始分配任务提交时若 GPU 繁忙先在 CPU 上启动持续监控每个训练步骤调用best_device()检查设备状态触发迁移当检测到更优设备如空闲 GPU时触发迁移逻辑状态保持迁移模型参数和优化器状态保持训练连续性无缝切换对训练逻辑透明无需重启或保存 checkpoint性能收益避免 GPU 空闲等待缩短整体训练时间提高硬件投资回报率2.4 模块亲和性优化在多进程环境中不同进程已加载的模块集合不同。如果任务所需的模块已在目标进程中存在可以避免重复内存分配示例进程 A 已加载numpy, torch, sklearn进程 B 已加载numpy, pandas新任务需要numpy, torch最优选择将任务提交给进程 A因为模块重合度更高无需额外加载 torch。传统进程池采用轮询或随机策略忽略了模块亲和性带来的内存优化机会。SmartPool 的亲和性提交方案SmartPool 在提交任务时优先选择模块重合度最高的工作进程fromsmartpoolimportProcessPool# 假设有两个工作进程# Worker A: 已加载 numpy, torch, sklearn# Worker B: 已加载 numpy, pandasdeftorch_task(data):importtorch# 需要 torchreturntorch_function(data)if__name____main__:withProcessPool()aspool:# 优先提交给 Worker A已有 torchfuturepool.submit(torch_task,args(data,))优化机制模块追踪记录每个工作进程已导入的模块集合重合度计算计算任务所需模块与进程已有模块的重合比例最优选择选择重合度最高的进程减少新模块加载内存节省避免重复加载相同模块降低总体内存占用2.5 弹性伸缩进程数在内存资源需求量少时SmartPool 会尽可能多的创建更多进程提过并行度在内存资源需求量大时如果内存不够SmartPool 会回收闲置的进程以释放进程中常驻内存。3. ONNX Runtime 推理支持除了 PyTorch 训练场景SmartPool 也为 ONNX Runtime 推理提供了专门的并行化方案。3.1 InferSessionPool线程安全的推理会话池InferSessionPool继承自ThreadPool在保证线程安全的前提下尽可能共享和复用InferenceSession实例兼顾安全性与内存效率。与传统方案最大的不同是无需人工指定 provider系统根据实时负载自动选择最优的 provider 以实现负载均衡推理。核心原理使用全局 LRU 缓存池管理 sessionkey 由模型路径 provider 名称 provider 选项组成对于CPUExecutionProvider和CUDAExecutionProvider线程安全所有线程共享同一个 session对于DmlExecutionProvider非线程安全每个线程独立持有自己的 session以线程 ID 为隔离粒度全局锁保护缓存访问确保增删操作的原子性自动 Provider 选择结合 SmartPool 的资源监控系统实时评估各设备负载为每个推理任务动态选择最优 providerfromsmartpoolimportInferSessionPool# 无需指定 provider系统自动选择最优设备withInferSessionPool()aspool:futurepool.submit(model.onnx,args(input_tensor,))resultfuture.result()内部工作流程提交 ONNX 模型路径pool 校验文件合法性SmartPool 资源监控系统评估当前 CPU、CUDA GPU、DML GPU 的负载情况自动选择负载最低、最合适的 provider并分配到对应工作线程线程调用_get_session()从缓存中获取或创建 sessionCPU/CUDA 场景下 session 跨线程共享DML 场景下按线程 ID 隔离3.2 多线程并发推理InferSessionPool支持批量提交推理任务充分利用多核 CPU 和多 GPU 的并行能力fromsmartpoolimportInferSessionPool,DataSize,Resource# 批量提交推理任务自动均衡分配到各线程withInferSessionPool()aspool:futures[pool.submit(model.onnx,args(image,),cpu_mode_resResource(cpu_cores_in_python1,cpu_mem500*DataSize.MB),gpu_mode_resResource(gpu_cores1024,gpu_mem1*DataSize.GB),)forimageinimage_batch]results[f.result()forfinfutures]与传统方案对比方案session 数量线程安全provider 选择资源感知代码复杂度单 session 全局锁1是串行人工指定不可切换无低每个线程独立创建N是人工指定不可切换无中手动 session 池N需自行保证人工指定不可切换无高InferSessionPoolN自动管理安全自动选择负载均衡支持低3.3 自动 Provider 选择与负载均衡不同 ExecutionProvider 的线程安全特性不同CPUExecutionProvider线程安全全局共享一个 sessionCUDAExecutionProvider线程安全onnxruntime-gpu全局共享一个 sessionDMLExecutionProvider非线程安全共享 session 调用run()会导致UnicodeDecodeError等内存损坏异常InferSessionPool针对不同 Provider 采取不同的策略CPU/CUDA 全局共享 session 以减少显存占用DML 按线程隔离 session 以保证安全。与传统方案的关键区别在于 Provider 的选择方式传统方案人工硬编码 provider → 某设备过载时无法切换 → 硬件饥饿或 OOMInferSessionPool资源监控系统实时评估各设备的 GPU 显存、CUDA 核心、DML 计算单元等负载 → 自动选择最优 provider → 实现异构设备间的负载均衡推理运行完整的 ONNX 推理示例python-msmartpool_examples.onnx_infer --max-workers4该示例会自动下载 YOLOv8 ONNX 模型和 COCO 验证集图片使用InferSessionPool并发执行推理显示下载进度条和推理进度条输出每张图片的 Top-5 分类结果4. 安装与快速开始4.1 安装方式# 基础安装pipinstallpysmartpool# 安装示例包pipinstallsmartpool-examples4.2 快速示例# test_smartpool.pyfromsmartpoolimportProcessPooldefcompute(x):returnx*2if__name____main__:withProcessPool()aspool:resultslist(pool.map(compute,range(10)))print(results)# [0, 2, 4, ..., 18]4.3 多个深度学习模型同时交叉验证运行如下命令即可启动使用 smartpool.ProcessPool 对 7 个深度学习模型同时进行 5 折交叉验证的例子运行后会给出源码目录可用 VSCode 打开源码目录查看。打开任务管理器可以看到充分利用的 GPU 和健康的内存利用率。随后得出多模型比较的柱状图。python-msmartpool_examples.cross_validation4.4 ONNX 模型并发推理运行如下命令即可启动使用 smartpool.InferSessionPool 使用 Yolo v8n 的 ONNX Runtime 模型对 20 张图片进行目标识别的示例运行后会给出源码目录可用 VSCode 打开源码目录查看。并能打印出每个推理任务使用的 provider 和 session如下图所示。python-msmartpool_examples.onnx_infer --max-workers45. 最佳实践为防止 CPU 超订应通过limit_num_single_thread设置数值计算库采用单线程重要limit_num_single_thread必须在导入 torch/numpy 之前调用。fromsmartpoolimportlimit_num_single_thread limit_num_single_thread()# 先调用importtorch# 然后导入importnumpyasnp6. 总结SmartPool 针对 PyTorch 训练和 ONNX Runtime 推理两大场景提供了完整解决方案PyTorch 训练支持真正的张量共享消除跨进程序列化开销相比标准库性能提升显著动态资源调度超越静态 max_workers 限制提高资源利用率异构负载均衡统一调度 CPU 和多 GPU简化代码复杂度训练热迁移最大化硬件利用率缩短整体训练时间模块亲和性优化内存使用效率减少重复加载ONNX Runtime 推理支持线程安全的会话管理CPU/CUDA 全局共享 sessionDML 按线程隔离兼顾安全与内存效率自动 Provider 选择无需人工指定 provider系统根据实时负载自动选择最优设备异构负载均衡在 CPU、CUDA、DML 之间动态调度推理任务避免硬件饥饿或过载声明式资源需求通过Resource指定 CPU/GPU 资源要求pool 自动调度这些特性使得 SmartPool 成为机器学习训练和推理的理想选择。项目资源github仓库: https://github.com/time-coder/smartpool/PyPI索引: https://pypi.org/project/pysmartpool/作者邮箱: binghui.wangfoxmail.comSmartPool 由本人开发欢迎大家提意见、提建议、提 issue。
SmartPool：智能资源感知的 PyTorch 并行计算框架
1. 背景与动机1.1 PyTorch 张量共享的困境在使用 Python 标准库进行并行计算时PyTorch 开发者面临着严峻的挑战。让我们从实际问题说起。问题一标准库进程池的序列化开销当使用multiprocessing.Pool或concurrent.futures.ProcessPoolExecutor传递 PyTorch 张量时会发生完整的序列化 - 反序列化过程frommultiprocessingimportPoolimporttorchdefprocess_tensor(tensor):returntensor.sum()if__name____main__:large_tensortorch.randn(10000,1000)# 约 40MBwithPool(processes4)aspool:# 每次提交都会序列化整个张量futures[pool.apply_async(process_tensor,args(large_tensor,))for_inrange(10)]results[f.get()forfinfutures]严重性分析每次任务提交都需要将 40MB 张量序列化跨进程传输序列化数据子进程反序列化重建张量对于大规模模型训练这种开销可能是灾难性的问题二PyTorch 官方进程池的已知 BugPyTorch 官方提供了torch.multiprocessing.pool.Pool作为解决方案但实际使用时会直接报错fromtorch.multiprocessing.poolimportPool poolPool()错误信息TypeError: SimpleQueue.__init__() missing 1 required keyword-only argument: ctx这是 PyTorch 的一个已知 bug。该问题于 2019 年 2 月提出torch.multiprocessing.pool.Pool broken至今尚未解决。但当你尝试使用torch.multiprocessing.Pool注意不是刚才的torch.multiprocessing.pool.Pool时发现它是可用的但遗憾的是torch.multiprocessing.Pool仅仅是multiprocessing.Pool的直接链接依然不能实现 tensor 共享。fromtorch.multiprocessingimportPoolasPool1frommultiprocessingimportPoolasPool2print(Pool1isPool2)# 输出True这意味着torch.multiprocessing.Pool并没有解决 PyTorch 张量的跨进程共享问题它只是multiprocessing.Pool的一个别名。1.2 ONNX Runtime 推理的挑战除了 PyTorch 训练场景ONNX Runtime 推理同样面临并行化的难题。问题多线程共享 InferenceSession 的不安全性与 Provider 选择困境ONNX Runtime 的InferenceSession对象内部维护了大量状态如 CUDA 流、DML 命令列表等。在多线程环境下共享同一个InferenceSession实例并同时调用run()会导致数据竞争和内存损坏importonnxruntimeasortfromconcurrent.futuresimportThreadPoolExecutor sessort.InferenceSession(model.onnx)definfer(input_data):returnsess.run(None,{input:input_data})# 多线程共享同一个 session → 内存损坏withThreadPoolExecutor(max_workers4)aspool:resultslist(pool.map(infer,inputs))即使规避了线程安全问题传统方案还需要手动指定 ExecutionProvider# 必须人工指定 provider无法根据负载动态切换sessort.InferenceSession(model.onnx,providers[CUDAExecutionProvider])# 或sessort.InferenceSession(model.onnx,providers[DmlExecutionProvider])# 或sessort.InferenceSession(model.onnx,providers[CPUExecutionProvider])手动指定 provider 的问题选中 CUDA 但显存不足时 → OOM 崩溃选中 DML 但 GPU 繁忙时 → 任务排队等待CPU 空闲浪费选中 CPU 但 GPU 空闲时 → 硬件利用率低下无法根据系统实时负载动态切换 provider导致硬件饥饿或过载传统解决方案及其缺陷方案缺陷每个线程独立创建 Session每次创建加载模型权重显存/内存开销成倍增长全局锁保护 Session退化为串行推理失去多线程优势手动管理 Session 池代码复杂度高缺少资源感知和自动设备调度1.3 SmartPool真正的解决方案面对标准库的序列化开销和 PyTorch 官方的未解 bug以及 ONNX Runtime 多线程推理的陷阱开发者急需一个统一的解决方案。SmartPool 应运而生。针对 PyTorchSmartPool 通过use_torchTrue启用 PyTorch 支持实现真正的张量共享fromsmartpoolimportProcessPoolimporttorch# 这个函数应该写在其他模块不能写在主模块中defprocess_tensor(tensor):# 零序列化开销张量通过共享内存传递returntensor.sum()if__name____main__:large_tensortorch.randn(10000,1000)# 约 40MBlarge_tensor.share_memory_()withProcessPool(use_torchTrue)aspool:# 张量通过共享内存传递无序列化开销futures[pool.submit(process_tensor,args(large_tensor,))for_inrange(10)]results[f.result()forfinfutures]针对 ONNX RuntimeSmartPool 提供InferSessionPool自动管理线程安全的推理会话fromsmartpoolimportInferSessionPoolwithInferSessionPool()aspool:futurepool.submit(model.onnx,args(input_tensor,))resultfuture.result()这仅仅是 SmartPool 的开始。无论是 PyTorch 训练还是 ONNX Runtime 推理SmartPool 都提供了一系列智能资源管理特性。2. 超越张量共享智能资源管理2.1 动态资源感知调度传统进程池通过max_workers参数控制并发进程数从而间接控制内存使用fromconcurrent.futuresimportProcessPoolExecutor# 静态限制始终最多 4 个进程withProcessPoolExecutor(max_workers4)aspool:resultspool.map(memory_intensive_func,data_list)问题分析不同任务的内存使用情况不一为防止太多内存消耗大的任务同时运行导致 OOM设置processes4限制最多启动的进程数。但当内存消耗大的任务完成退出进程池后进程池处于内存使用低谷阶段如仅占用 500MB时无法启动更多任务剩余的 CPU 和内存被闲置。这种静态配置存在三个核心问题过于保守内存峰值过去后空闲资源无法被利用缺乏弹性无法根据实际内存使用情况动态调整并发度资源浪费在多 GPU 环境下静态配置无法适应动态变化的负载分布SmartPool 的动态调度方案SmartPool 允许任务声明资源需求系统根据实时资源状态动态调度fromsmartpoolimportProcessPool,DataSize,Resourcedefmemory_intensive_task(data):importnumpyasnp big_arraynp.zeros((10000,10000))# 约 800MB# 处理数据...returnresultif__name____main__:withProcessPool(use_torchTrue)aspool:# 声明资源需求futurepool.submit(memory_intensive_task,args(dataset,),cpu_mode_resResource(cpu_cores_in_python2,# 需要 2 个 CPU 核心cpu_mem4*DataSize.GB,# 需要 4GB 系统内存),gpu_mode_resResource(cpu_cores_in_python1,cpu_mem500*DataSize.MB,gpu_cores1024,# 需要 1024 个 CUDA 核心gpu_mem2*DataSize.GB,# 需要 2GB GPU 显存),)# 若当前资源不足任务自动进入等待队列# 待其他任务释放资源后自动启动执行resultfuture.result()与传统方案的对比特性传统进程池max_workersSmartPool动态调度并发控制固定进程数动态资源配额内存管理静态上限按需分配弹性伸缩资源利用率保守估计资源闲置充分利用并防止耗尽适应性无法应对负载波动实时监控动态调整配置复杂度需手动估算自动管理优势分析弹性伸缩当任务处于内存低谷期自动启动更多任务填充空闲资源防 OOM 机制通过资源声明和监控确保总消耗不超过系统容量细粒度控制可为每个任务单独指定 CPU、内存、GPU 需求2.2 异构负载均衡现代计算节点通常包含多核 CPU 和多块 GPU如何均匀分配任务成为关键问题典型场景4 核 CPU 2 块 GPU每块 GPU 有多个 CUDA 核心需要同时训练 20 个模型任务可能分配到 CPU 或任意一块 GPU传统方案的不足手动调度复杂需要显式管理每个任务的设备分配资源竞争多个任务可能竞争同一块 GPU导致显存溢出负载不均某些 GPU 过载而其他 GPU 空闲缺乏统一抽象CPU 和 GPU 任务需要不同的代码路径SmartPool 的统一调度方案SmartPool 统一管理 CPU 和 GPU 资源自动将任务分配到最优设备任务可通过best_device()函数获取到分配到的设备fromsmartpoolimportProcessPool,DataSize,Resource,best_devicedeftrain_model(model_config):devicebest_device()# 自动获取最优设备# 在最优设备上执行训练returntrain_on_device(model_config,device)if__name____main__:configs[...]# 20 个模型配置withProcessPool(use_torchTrue)aspool:futures[]forconfiginconfigs:futurepool.submit(train_model,args(config,),cpu_mode_resResource(cpu_cores_in_python1,# 需要 1 个 CPU 核心cpu_mem1*DataSize.GB,# 需要 1GB 系统内存),gpu_mode_resResource(cpu_cores_in_python1,cpu_mem500*DataSize.MB,gpu_cores1024,# 需要 1024 个 CUDA 核心gpu_mem1*DataSize.GB,# 需要 1GB GPU 显存),)# SmartPool 自动均衡分配到 CPU 和多块 GPUfutures.append(future)results[f.result()forfinfutures]2.3 训练热迁移在动态资源环境中GPU 可用性可能随时间变化场景描述初始状态所有 GPU 被占用部分训练任务在 CPU 上启动时间点 T1某个 GPU 上的任务快速完成释放 GPU问题仍在 CPU 上训练的任务继续占用 CPU新释放的 GPU 空闲结果硬件资源浪费整体吞吐量下降理想方案应能检测到 GPU 空闲并将 CPU 上的训练任务自动迁移到 GPU无需重启训练过程。SmartPool 的热迁移方案SmartPool 支持在训练过程中动态检测设备变化并自动迁移任务fromsmartpoolimportProcessPool,best_device,move_optimizer_toimporttorchdeftraining_task(model,optimizer,data_loader):# 初始化时获取设备devicebest_device()old_devicedevice model.to(device)forepochinrange(epochs):forbatch_idx,(data,target)inenumerate(data_loader):# 每个 batch 检查设备变化devicebest_device()data,targetdata.to(device),target.to(device)ifdevice!old_device:# 迁移模型和优化器model.to(device)move_optimizer_to(optimizer,device)old_devicedevice# 在新设备上继续训练outputmodel(data)losscriterion(output,target)# ... 反向传播returnmodelif__name____main__:withProcessPool(use_torchTrue)aspool:# 初始可能在 CPU 上训练futurepool.submit(training_task,args(model,optimizer,data))# 当 GPU 空闲时自动迁移到 GPUresultfuture.result()工作流程初始分配任务提交时若 GPU 繁忙先在 CPU 上启动持续监控每个训练步骤调用best_device()检查设备状态触发迁移当检测到更优设备如空闲 GPU时触发迁移逻辑状态保持迁移模型参数和优化器状态保持训练连续性无缝切换对训练逻辑透明无需重启或保存 checkpoint性能收益避免 GPU 空闲等待缩短整体训练时间提高硬件投资回报率2.4 模块亲和性优化在多进程环境中不同进程已加载的模块集合不同。如果任务所需的模块已在目标进程中存在可以避免重复内存分配示例进程 A 已加载numpy, torch, sklearn进程 B 已加载numpy, pandas新任务需要numpy, torch最优选择将任务提交给进程 A因为模块重合度更高无需额外加载 torch。传统进程池采用轮询或随机策略忽略了模块亲和性带来的内存优化机会。SmartPool 的亲和性提交方案SmartPool 在提交任务时优先选择模块重合度最高的工作进程fromsmartpoolimportProcessPool# 假设有两个工作进程# Worker A: 已加载 numpy, torch, sklearn# Worker B: 已加载 numpy, pandasdeftorch_task(data):importtorch# 需要 torchreturntorch_function(data)if__name____main__:withProcessPool()aspool:# 优先提交给 Worker A已有 torchfuturepool.submit(torch_task,args(data,))优化机制模块追踪记录每个工作进程已导入的模块集合重合度计算计算任务所需模块与进程已有模块的重合比例最优选择选择重合度最高的进程减少新模块加载内存节省避免重复加载相同模块降低总体内存占用2.5 弹性伸缩进程数在内存资源需求量少时SmartPool 会尽可能多的创建更多进程提过并行度在内存资源需求量大时如果内存不够SmartPool 会回收闲置的进程以释放进程中常驻内存。3. ONNX Runtime 推理支持除了 PyTorch 训练场景SmartPool 也为 ONNX Runtime 推理提供了专门的并行化方案。3.1 InferSessionPool线程安全的推理会话池InferSessionPool继承自ThreadPool在保证线程安全的前提下尽可能共享和复用InferenceSession实例兼顾安全性与内存效率。与传统方案最大的不同是无需人工指定 provider系统根据实时负载自动选择最优的 provider 以实现负载均衡推理。核心原理使用全局 LRU 缓存池管理 sessionkey 由模型路径 provider 名称 provider 选项组成对于CPUExecutionProvider和CUDAExecutionProvider线程安全所有线程共享同一个 session对于DmlExecutionProvider非线程安全每个线程独立持有自己的 session以线程 ID 为隔离粒度全局锁保护缓存访问确保增删操作的原子性自动 Provider 选择结合 SmartPool 的资源监控系统实时评估各设备负载为每个推理任务动态选择最优 providerfromsmartpoolimportInferSessionPool# 无需指定 provider系统自动选择最优设备withInferSessionPool()aspool:futurepool.submit(model.onnx,args(input_tensor,))resultfuture.result()内部工作流程提交 ONNX 模型路径pool 校验文件合法性SmartPool 资源监控系统评估当前 CPU、CUDA GPU、DML GPU 的负载情况自动选择负载最低、最合适的 provider并分配到对应工作线程线程调用_get_session()从缓存中获取或创建 sessionCPU/CUDA 场景下 session 跨线程共享DML 场景下按线程 ID 隔离3.2 多线程并发推理InferSessionPool支持批量提交推理任务充分利用多核 CPU 和多 GPU 的并行能力fromsmartpoolimportInferSessionPool,DataSize,Resource# 批量提交推理任务自动均衡分配到各线程withInferSessionPool()aspool:futures[pool.submit(model.onnx,args(image,),cpu_mode_resResource(cpu_cores_in_python1,cpu_mem500*DataSize.MB),gpu_mode_resResource(gpu_cores1024,gpu_mem1*DataSize.GB),)forimageinimage_batch]results[f.result()forfinfutures]与传统方案对比方案session 数量线程安全provider 选择资源感知代码复杂度单 session 全局锁1是串行人工指定不可切换无低每个线程独立创建N是人工指定不可切换无中手动 session 池N需自行保证人工指定不可切换无高InferSessionPoolN自动管理安全自动选择负载均衡支持低3.3 自动 Provider 选择与负载均衡不同 ExecutionProvider 的线程安全特性不同CPUExecutionProvider线程安全全局共享一个 sessionCUDAExecutionProvider线程安全onnxruntime-gpu全局共享一个 sessionDMLExecutionProvider非线程安全共享 session 调用run()会导致UnicodeDecodeError等内存损坏异常InferSessionPool针对不同 Provider 采取不同的策略CPU/CUDA 全局共享 session 以减少显存占用DML 按线程隔离 session 以保证安全。与传统方案的关键区别在于 Provider 的选择方式传统方案人工硬编码 provider → 某设备过载时无法切换 → 硬件饥饿或 OOMInferSessionPool资源监控系统实时评估各设备的 GPU 显存、CUDA 核心、DML 计算单元等负载 → 自动选择最优 provider → 实现异构设备间的负载均衡推理运行完整的 ONNX 推理示例python-msmartpool_examples.onnx_infer --max-workers4该示例会自动下载 YOLOv8 ONNX 模型和 COCO 验证集图片使用InferSessionPool并发执行推理显示下载进度条和推理进度条输出每张图片的 Top-5 分类结果4. 安装与快速开始4.1 安装方式# 基础安装pipinstallpysmartpool# 安装示例包pipinstallsmartpool-examples4.2 快速示例# test_smartpool.pyfromsmartpoolimportProcessPooldefcompute(x):returnx*2if__name____main__:withProcessPool()aspool:resultslist(pool.map(compute,range(10)))print(results)# [0, 2, 4, ..., 18]4.3 多个深度学习模型同时交叉验证运行如下命令即可启动使用 smartpool.ProcessPool 对 7 个深度学习模型同时进行 5 折交叉验证的例子运行后会给出源码目录可用 VSCode 打开源码目录查看。打开任务管理器可以看到充分利用的 GPU 和健康的内存利用率。随后得出多模型比较的柱状图。python-msmartpool_examples.cross_validation4.4 ONNX 模型并发推理运行如下命令即可启动使用 smartpool.InferSessionPool 使用 Yolo v8n 的 ONNX Runtime 模型对 20 张图片进行目标识别的示例运行后会给出源码目录可用 VSCode 打开源码目录查看。并能打印出每个推理任务使用的 provider 和 session如下图所示。python-msmartpool_examples.onnx_infer --max-workers45. 最佳实践为防止 CPU 超订应通过limit_num_single_thread设置数值计算库采用单线程重要limit_num_single_thread必须在导入 torch/numpy 之前调用。fromsmartpoolimportlimit_num_single_thread limit_num_single_thread()# 先调用importtorch# 然后导入importnumpyasnp6. 总结SmartPool 针对 PyTorch 训练和 ONNX Runtime 推理两大场景提供了完整解决方案PyTorch 训练支持真正的张量共享消除跨进程序列化开销相比标准库性能提升显著动态资源调度超越静态 max_workers 限制提高资源利用率异构负载均衡统一调度 CPU 和多 GPU简化代码复杂度训练热迁移最大化硬件利用率缩短整体训练时间模块亲和性优化内存使用效率减少重复加载ONNX Runtime 推理支持线程安全的会话管理CPU/CUDA 全局共享 sessionDML 按线程隔离兼顾安全与内存效率自动 Provider 选择无需人工指定 provider系统根据实时负载自动选择最优设备异构负载均衡在 CPU、CUDA、DML 之间动态调度推理任务避免硬件饥饿或过载声明式资源需求通过Resource指定 CPU/GPU 资源要求pool 自动调度这些特性使得 SmartPool 成为机器学习训练和推理的理想选择。项目资源github仓库: https://github.com/time-coder/smartpool/PyPI索引: https://pypi.org/project/pysmartpool/作者邮箱: binghui.wangfoxmail.comSmartPool 由本人开发欢迎大家提意见、提建议、提 issue。
