Multiprocessing.Pool指定4进程却占用全部CPU的解决问询

问题描述 #

在40核的共享Ubuntu 20.04服务器上运行540个强化学习实验，为避免占用全部资源，使用multiprocessing.Pool(processes=4)控制同时运行的进程数，但通过htop观察发现所有40核均被占满，每个子进程占用大量CPU核心。尝试调整chunksize参数及map/map_async等Pool方法后，问题未解决。

使用的核心代码如下：

from stable_baselines3.common.callbacks import EvalCallback
from stable_baselines3.common.logger import configure
from stable_baselines3 import PPO
from vrpEnv import VRPEnv
import multiprocessing
import time  # 原代码遗漏time模块导入

expList = [*range(1, 541)]

ITERATIONS = 500
TIMESTEPS = 2048*10

def runExperiment(exp):
    models_dir = "models/" + str(exp)  # 原代码exp为整数，需转字符串避免报错
    log_dir = "logs/"

    rank = multiprocessing.current_process()._identity[0]
    print(f"Processor {rank} working on case {exp}")

    env = VRPEnv()

    model = PPO("MultiInputPolicy", env, verbose=1, tensorboard_log=log_dir, device='cpu')
   
    eval_callback = EvalCallback(env, best_model_save_path=models_dir,
                            eval_freq=40000, deterministic=True, render=False)

    new_logger = configure(log_dir, format_strings = ["log", "csv", "tensorboard"])  # 原代码缺少log_dir参数
    nombreExp = f"exp_{exp}"  # 原代码未定义nombreExp变量

    start_time = time.time()

    for _ in range(1, ITERATIONS+1):
        model.set_logger(new_logger)
        model.learn(total_timesteps = TIMESTEPS, reset_num_timesteps = False, tb_log_name = nombreExp, callback=eval_callback)

    print("---%s: %s minutos ---" % (exp, round((time.time() - start_time)/60, 2)))

    env.close()

if __name__ == '__main__':
    pool = multiprocessing.Pool(processes = 4)
    result = pool.imap(runExperiment, expList)
    pool.close()
    pool.join()
    print('END')

问题原因 #

multiprocessing.Pool(processes=4)仅限制了同时运行的进程数，但Stable Baselines3的PPO模型默认会启用内部多线程加速（依赖PyTorch、NumPy等库的并行计算能力），每个子进程会自动占用多个CPU核心，最终导致所有40核被占满。

解决方案 #

1. 限制PyTorch线程数 #

在每个子进程的初始化阶段，强制PyTorch仅使用1个线程，避免单个进程占用多核心。在runExperiment函数开头添加：

import torch
torch.set_num_threads(1)
torch.set_num_interop_threads(1)

2. 限制OpenMP相关库的线程数 #

NumPy、MKL等依赖OpenMP的库也会自动启用多线程，需通过环境变量限制其线程数。可以在代码开头全局设置：

import os
os.environ["OMP_NUM_THREADS"] = "1"
os.environ["MKL_NUM_THREADS"] = "1"
os.environ["NUMEXPR_NUM_THREADS"] = "1"

或者在终端启动脚本时直接设置环境变量：

OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 NUMEXPR_NUM_THREADS=1 python your_script.py

3. 修正原代码中的潜在问题 #

原代码存在几处语法/逻辑错误，需修正后才能正常运行：

• 遗漏time模块导入
• models_dir = "models/" + exp中exp是整数，需转为字符串
• configure函数缺少log_dir参数
• 未定义nombreExp变量

验证效果 #

应用上述设置后，每个子进程仅会占用1个CPU核心，4个并行进程最多占用4核，不会再出现全核心占用的情况。

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