Python多进程调用线程过度生成问题排查与优化咨询

背景与问题 #

正在开展一项分析工作，需为10-60个独立数据源分别拟合基于PyTorch实现的模型。为了将所有输出统一保存到同一文件，且主要在SLURM管理的高性能计算集群运行，选择使用torch.multiprocessing而非SLURM批量数组调用。

近期一批任务因生成过多线程导致CPU负载过高被集群取消，核心代码如下：

torch.set_num_threads(1) 

import torch.multiprocessing as mp
with mp.Pool(processes=20) as pool:
  output_to_save = pool.map(myModelFit, sourcesN)
  pool.close()

按照HPC管理员要求设置了20个进程（集群多数计算节点为48核），预期同时运行的线程数不超过20，但实际生成了数百个线程，导致CPU占用率异常偏高。本地服务器运行时也出现相同问题，因此排除SLURM参数（如--tasks-per-node 20）的影响。

尝试相关优化方案后，本地和集群的CPU使用率被限制在100%，但仍存在以下疑问：

• 这种CPU使用率被限制在100%的情况，是否属于合理高效的利用方式？
• 为何提前设置的torch.set_num_threads(1)单进程单线程配置未生效？
• 为何使用pool.map多进程调用时，线程数会远超进程数；而仅处理单个数据源（无多进程调用）时，只生成1个线程？

已知问题可能与myModelFit内部的PyTorch和NumPy调用有关，但不确定是否也受到mp.Pool调用机制的影响。

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