【Spinning Up】四、极简版并行网格搜索的脚本(python同时启动多个不同参数脚本)

kaixindelele

前言：

上篇文章说到，利用spinup的run_entrypoint.py可以实现依次批量调参，但这个仍然不能利用好，我们电脑上的其他CPU核心，每次只能等一个参数调整完毕之后，才能在当前CPU核心下，启动下一个进程。

这还是不够优雅~

因此我们利用mpi_tool.py中的mpi_fork函数，实现同时启动多个参数；

最新版代码（博客里的代码是两个月前的，可能有bug，功能不够完善）：

https://github.com/kaixindelele/DRLib/tree/main/tune_exps

mpi_fork()函数详解：

这个函数还是要花点时间去理解的，我的注释不一定能直接看懂；
最好多试试，这种破玩意儿，debug都太好使，但是也必须debug。

def mpi_fork(n, bind_to_core=False):
	# n是CPU核心数目，即要同时启动几个脚本的意思~
	# 如果n不大于1，直接就返回，不折腾了
    if n <= 1:
        return
    # 判断是否在MPI里面，如果不在的话，则执行下面的语句
    # 进去之后，把标志符变为True，防止套娃无限递归
    if os.getenv("IN_MPI") is None:
        env = os.environ.copy()
        # 更新环境信息
        env.update(
            MKL_NUM_THREADS="1",
            OMP_NUM_THREADS="1",
            IN_MPI="1"
        )
        # 准备好标准的mpi启动脚本，如果同时启动4个脚本，例子为：
        # mpirun -np 4 python demo.py args
        args = ["mpirun", "-np", str(n)]
        if bind_to_core:
            args += ["-bind-to", "core"]
        # sys.executable 即python解释器；
        # sys.argv的值为，当前启动当前脚本，命令行传来的参数
        # 里面有待执行的脚本文件，和参数，在我们的任务中，是打包好的thunk_plus，既有.py文件，又有参数
        args += [sys.executable] + sys.argv
        # 拿到这个脚本字符串列表，用subprocess.check_call(args)启动就行了。
        # 如果不加上面的in_mpi判断，那么就会无限套娃
        subprocess.check_call(args, env=env)

有了这个同时启动的函数，那么我们就可以想办法，同时启动多个超参数了；

接下来就是想办法，给每个进程的脚本，分配不同超参数~

假设我们有2组，一组5个超参数，一个6个。我们只有4个CPU核心；

该怎么设计？哈哈，花了一天的时间，我设计好了_{优雅简洁，但是可读性感觉也降低了不少。}

不管怎么说，我们每次能启动四个进程了；

tune_mpi_funcs极简例程：

极简例程涉及mpi，需要在tune_exps.py中用mpi_fork()函数启动多个脚本，其次还要在tune_func.py中，用mpi的pid来判断，该选列表中的哪个参数；
新建三个函数：

tune_exps.py: 主函数
tune_func.py: 待执行函数
run_entrypoint.py: 执行隔离后的函数
mpi_tool.py：利用其中的mpi_fork()函数，上面解释了

tune_func.py：待执行函数

这次的tune_func.py，要稍微复杂一些，就是传一个字典参数列表，根据当前pid来选择参数字典:

import sys
import os
from mpi4py import MPI
import numpy as np

def func(params_dict_list):
    proc_id = MPI.COMM_WORLD.Get_rank()
    if proc_id > len(params_dict_list)-1:
        print("proc_id:", proc_id)
        print("sys.exit()")
        sys.exit()
        print("sys.exit()")
    params_dict = params_dict_list[proc_id]
    print("proc_id:", proc_id)
    print("params_dict:", params_dict)
    print("-"*20)


if __name__=='__main__':
    params_dict = {
        'lr': [2, 3, 4, 5, 6, 7],
        "batch": [10, 20, 30, 40, 50,],
        "epoch": [100, 200, 300, 400, 500, 600],
    }
    import itertools
    
    params_list = [list(value) for value in itertools.product(*params_dict.values())]        
    params_dict_list = [{key: cur_param.pop(0) for key, value in params_dict.items()} for cur_param in params_list]
    for i in range(2):        
        func(params_dict_list=params_dict_list)

run_entrypoint.py 入口函数

这个函数用的spinup自带的，也比较简单，就是一个接收命令行传参，一个反序列化过程，一个执行恢复后的函数的过程

import zlib
import pickle
import base64

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('encoded_thunk')
    args = parser.parse_args()
    thunk = pickle.loads(zlib.decompress(base64.b64decode(args.encoded_thunk)))
    thunk()

tune_exps.py 调用主函数：

这里东西比较多，相比spinup原本的run_utils.py，我已经删到不能再删的地步了。
另外，这次不能用lambda来封装函数了，必须要用thunk_plus()函数封装了。
整个流程：

导入待执行函数
将批量参数放在列表中，如果一个不够，就多层循环套娃一下；
将函数和参数传入callexp()中；
用lambda语句，将函数和参数打包在一起；
将打包后的序列化成二进制字符串encoded_thunk ；
封装成命令行字符串列表：cmd = ['python', 'run_entrypoint.py', encoded_thunk]
利用subprocess.call_check(cmd)

完了~

import base64
from copy import deepcopy
import cloudpickle
import numpy as np
import os
import os.path as osp
import psutil
import string
import subprocess
from subprocess import CalledProcessError
import sys
from textwrap import dedent
import time
import zlib

# 导入待执行的函数
from spinup_utils.mpi_tools import mpi_fork
from spinup_utils.tune_func import func


DIV_LINE_WIDTH = 80


def call_experiment(thunk, params_dict_list, **kwargs):
    """
        :params_dict thunk:待启动的函数
        :params_dict params_dict:批量参数名
        :params kwargs: 其他的一些没考虑到的参数~用处不大，没事儿最好别写这个,容易造成混乱~    
        正常的函数，传入参数之后，就会直接执行。
        但是通过这个神奇的lambda，就可以即把参数传进去，又不执行。返回出一个函数
        再次调用的时候，只需要将返回值，加上括号，即当一个无参数传入的函数执行就可以了。
    """
    def thunk_plus():
        # Fork into multiple processes
        mpi_fork(4)
        # Run thunk
        thunk(params_dict_list)
    # lambda封装会让tune_func.py中导入MPI模块报初始化错误。
    # thunk_plus = lambda: thunk(params_dict)
    # mpi_fork(len(params_dict))
    pickled_thunk = cloudpickle.dumps(thunk_plus)
    encoded_thunk = base64.b64encode(zlib.compress(pickled_thunk)).decode('utf-8')
    # 当前脚本和entry_point.py的路径要在一起，要不然下面的语句要改。
    entrypoint = osp.join(osp.abspath(osp.dirname(__file__)), 'run_entrypoint.py')
    # subprocess的输入就是一个字符串列表，正常在命令行，该怎么输入，这个就该怎么写。
    cmd = [sys.executable if sys.executable else 'python', entrypoint, encoded_thunk]
    print("tune_exps_pid:", os.getpid())
    try:
        subprocess.check_call(cmd, env=os.environ)
    except CalledProcessError:
        err_msg = '\n'*3 + '='*DIV_LINE_WIDTH + '\n' + dedent("""
            Check the traceback above to see what actually went wrong. 
            """) + '='*DIV_LINE_WIDTH + '\n'*3
        print(err_msg)
        raise


if __name__ == '__main__':    
    cpu_num = 5
    params_dict = {
        'lr': [2, 3, 4],
        "batch": [10, 20, 30],
        "epoch": [9, 8, 7],
    }

    import itertools
    # 将字典变为排列组合列表
    params_list = [list(value) for value in itertools.product(*params_dict.values())]        
    # 将列表列表变为单个文件的字典列表
    params_dict_list = [{key: cur_param.pop(0) for key, value in params_dict.items()} for cur_param in params_list]
    print(params_dict_list)
    # 每次传入cpu_num数个字典。
    for i in range(0, len(params_dict_list), cpu_num):
        cur_params_dict_list = params_dict_list[i:i+cpu_num]
        print("cur_params_dict_list:", cur_params_dict_list)
        call_experiment(thunk=func, params_dict_list=cur_params_dict_list)

打印结果：

可以看到主进程ID一直不变，子进程的ID一直更新。参数传输正常，一切OK。

[{'lr': 2, 'batch': 10, 'epoch': 9}, {'lr': 2, 'batch': 10, 'epoch': 8}, {'lr': 2, 'batch': 10, 'epoch': 7}, {'lr': 2, 'batch': 20, 'epoch': 9}, {'lr': 2, 'batch': 20, 'epoch': 8}, {'lr': 2, 'batch': 20, 'epoch': 7}, {'lr': 2, 'batch': 30, 'epoch': 9}, {'lr': 2, 'batch': 30, 'epoch': 8}, {'lr': 2, 'batch': 30, 'epoch': 7}, {'lr': 3, 'batch': 10, 'epoch': 9}, {'lr': 3, 'batch': 10, 'epoch': 8}, {'lr': 3, 'batch': 10, 'epoch': 7}, {'lr': 3, 'batch': 20, 'epoch': 9}, {'lr': 3, 'batch': 20, 'epoch': 8}, {'lr': 3, 'batch': 20, 'epoch': 7}, {'lr': 3, 'batch': 30, 'epoch': 9}, {'lr': 3, 'batch': 30, 'epoch': 8}, {'lr': 3, 'batch': 30, 'epoch': 7}, {'lr': 4, 'batch': 10, 'epoch': 9}, {'lr': 4, 'batch': 10, 'epoch': 8}, {'lr': 4, 'batch': 10, 'epoch': 7}, {'lr': 4, 'batch': 20, 'epoch': 9}, {'lr': 4, 'batch': 20, 'epoch': 8}, {'lr': 4, 'batch': 20, 'epoch': 7}, {'lr': 4, 'batch': 30, 'epoch': 9}, {'lr': 4, 'batch': 30, 'epoch': 8}, {'lr': 4, 'batch': 30, 'epoch': 7}]
cur_params_dict_list: [{'lr': 2, 'batch': 10, 'epoch': 9}, {'lr': 2, 'batch': 10, 'epoch': 8}, {'lr': 2, 'batch': 10, 'epoch': 7}, {'lr': 2, 'batch': 20, 'epoch': 9}, {'lr': 2, 'batch': 20, 'epoch': 8}]
tune_exps_pid: 15411
proc_id: 1
params_dict: {'lr': 2, 'batch': 10, 'epoch': 8}
--------------------
proc_id: 0
params_dict: {'lr': 2, 'batch': 10, 'epoch': 9}
--------------------
proc_id: 2
params_dict: {'lr': 2, 'batch': 10, 'epoch': 7}
--------------------
proc_id: 3
params_dict: {'lr': 2, 'batch': 20, 'epoch': 9}
--------------------
proc_id: 0
params_dict: {'lr': 2, 'batch': 10, 'epoch': 9}
--------------------
cur_params_dict_list: [{'lr': 2, 'batch': 20, 'epoch': 7}, {'lr': 2, 'batch': 30, 'epoch': 9}, {'lr': 2, 'batch': 30, 'epoch': 8}, {'lr': 2, 'batch': 30, 'epoch': 7}, {'lr': 3, 'batch': 10, 'epoch': 9}]
tune_exps_pid: 15411
proc_id: 3
params_dict: {'lr': 2, 'batch': 30, 'epoch': 7}
--------------------
proc_id: 2
params_dict: {'lr': 2, 'batch': 30, 'epoch': 8}
--------------------
proc_id: 0
params_dict: {'lr': 2, 'batch': 20, 'epoch': 7}
--------------------
proc_id: 1
params_dict: {'lr': 2, 'batch': 30, 'epoch': 9}
--------------------
proc_id: 0
params_dict: {'lr': 2, 'batch': 20, 'epoch': 7}
--------------------
cur_params_dict_list: [{'lr': 3, 'batch': 10, 'epoch': 8}, {'lr': 3, 'batch': 10, 'epoch': 7}, {'lr': 3, 'batch': 20, 'epoch': 9}, {'lr': 3, 'batch': 20, 'epoch': 8}, {'lr': 3, 'batch': 20, 'epoch': 7}]
tune_exps_pid: 15411
proc_id: 2
params_dict: {'lr': 3, 'batch': 20, 'epoch': 9}
--------------------
proc_id: 3
params_dict: {'lr': 3, 'batch': 20, 'epoch': 8}
--------------------
proc_id: 0
params_dict: {'lr': 3, 'batch': 10, 'epoch': 8}
--------------------
proc_id: 1
params_dict: {'lr': 3, 'batch': 10, 'epoch': 7}
--------------------
proc_id: 0
params_dict: {'lr': 3, 'batch': 10, 'epoch': 8}
--------------------
cur_params_dict_list: [{'lr': 3, 'batch': 30, 'epoch': 9}, {'lr': 3, 'batch': 30, 'epoch': 8}, {'lr': 3, 'batch': 30, 'epoch': 7}, {'lr': 4, 'batch': 10, 'epoch': 9}, {'lr': 4, 'batch': 10, 'epoch': 8}]
tune_exps_pid: 15411
proc_id: 1
params_dict: {'lr': 3, 'batch': 30, 'epoch': 8}
--------------------
proc_id: 0
params_dict: {'lr': 3, 'batch': 30, 'epoch': 9}
--------------------
proc_id: 2
params_dict: {'lr': 3, 'batch': 30, 'epoch': 7}
--------------------
proc_id: 3
params_dict: {'lr': 4, 'batch': 10, 'epoch': 9}
--------------------
proc_id: 0
params_dict: {'lr': 3, 'batch': 30, 'epoch': 9}
--------------------
cur_params_dict_list: [{'lr': 4, 'batch': 10, 'epoch': 7}, {'lr': 4, 'batch': 20, 'epoch': 9}, {'lr': 4, 'batch': 20, 'epoch': 8}, {'lr': 4, 'batch': 20, 'epoch': 7}, {'lr': 4, 'batch': 30, 'epoch': 9}]
tune_exps_pid: 15411
proc_id: 1
params_dict: {'lr': 4, 'batch': 20, 'epoch': 9}
--------------------
proc_id: 3
params_dict: {'lr': 4, 'batch': 20, 'epoch': 7}
--------------------
proc_id: 0
params_dict: {'lr': 4, 'batch': 10, 'epoch': 7}
--------------------
proc_id: 2
params_dict: {'lr': 4, 'batch': 20, 'epoch': 8}
--------------------
proc_id: 0
params_dict: {'lr': 4, 'batch': 10, 'epoch': 7}
--------------------
cur_params_dict_list: [{'lr': 4, 'batch': 30, 'epoch': 8}, {'lr': 4, 'batch': 30, 'epoch': 7}]
tune_exps_pid: 15411
proc_id: 1
params_dict: {'lr': 4, 'batch': 30, 'epoch': 7}
--------------------
proc_id: 0
params_dict: {'lr': 4, 'batch': 30, 'epoch': 8}
--------------------
proc_id: 2
sys.exit()
proc_id: 3
sys.exit()
proc_id: 0
params_dict: {'lr': 4, 'batch': 30, 'epoch': 8}
--------------------

诡异bug：

今天遇到了一个极度诡异的bug；
在虚拟环境1中，我的MPI.COMM_WORLD.Get_rank()
拿到的值一直都是0！
不管mpi_fork()几次，都是0，极度离谱；
debug什么的都无法解决，搞得我以为我电脑坏了，因为上午还是好好的；
后来我想起来，可能是我的虚拟环境换了？
换回了另一个虚拟环境2，终于跑成功了，这样的话，就明白了，是我的虚拟环境1坏掉了，我也找不到bug在哪儿，他也不报错，艹，不报错的bug是最恶心的；
我只是看了一下mpi4py的版本，虚拟环境1 的是3.0.3，虚拟环境2的是3.0.2，应该不是这个的原因吧？

联系方式：

ps: 欢迎做强化的同学加群一起学习：

深度强化学习-DRL：799378128

欢迎关注知乎帐号：未入门的炼丹学徒

CSDN帐号：https://blog.csdn.net/hehedadaq

极简spinup+HER+PER代码实现，两小时之内配置完毕：https://github.com/kaixindelele/DRLib

Document