【python】在SLURM上使用multiprocessing并行

示例：VRP并行计算

• 主函数：main-process.py
  • complex_func为需要并行的函数，模拟退火解VRP问题，一堆参数
  • itertools：生成不同参数组合成的元组，传入complex_func，根据结果选最佳参数
  • 工作进程池Pool
  • apply_async()：单次启动一个任务，启动后不等这个进程结束又开始执行新任务。 func 只会在一个进程池中的一个工作进程中执行
  • r.get()顺序输出complex_func返回的结果：不用在并行过程读写文件就不用Lock；没有r.get()的话并行的complex_func不会执行
  • n_jobs = int(os.getenv('SLURM_CPUS_PER_TASK'))：提交到超算，直接读取sbatch的设置，防止py代码并行核数和超算设置不同
  • 其他函数print信息也提交到标准错误输出，标准输出只保留结果便于处理：print(xxx, file=sys.stderr)或者sys.stderr.write()，例如

print('%s is running %s' % (multiprocessing.current_process().name, \
        yourfunc.__name__), file=sys.stderr, flush=True)

• 参考：
  • https://blog.csdn.net/BobYuan888/article/details/109266020
  • https://zhuanlan.zhihu.com/p/194349143
  • https://docs.python.org/zh-cn/3/library/multiprocessing.html#examples

# -*- coding: UTF-8 -*-
import os
import time
import multiprocessing
import itertools
import sys
from VNSSA import VNSSA
from insert import insert
from Data import Data

def complex_func(files, T0, cool_rate, init_capacity, N_weight, iter_num):

    data = Data()
    data.readSolomonData(files)
    # print('当前文件:',files)
    mode = 1
    start = time.time()
    oldcapacity = data.capacity
    data.capacity = int(data.capacity * init_capacity)
    p = insert(data)
    data.capacity = oldcapacity
    best_p = VNSSA(p, mode, data, T0, cool_rate, N_weight)

    end = time.time()
    obj = best_p.cal_totalObj(mode, data)
    run_time = end-start

    vehicle_num = len(best_p.routes)
    return [files, obj, vehicle_num, run_time, T0, cool_rate, init_capacity, \
        N_weight, iter_num, str(best_p.routes)]


def main(n_jobs=10):

    file_path = 'solomon_dataset/'
    file_list = ['r108.txt','r109.txt','r110.txt','r111.txt']
    # file_list = ['r201.txt','r202.txt','r203.txt','r204.txt']
    # file_list = ['r206.txt','r207.txt','r208.txt','r209.txt',\
    #     'r210.txt','r211.txt']
    T0 = [10, 50, 100, 150]
    cool_rate = [0.997, 0.9997, 0.99975]
    init_capacity = [0.3, 0.5, 0.7, 1]
    N_weight = [250, 450]
    iter_num = [i for i in range(3)]

    # test cases
    # file_list = ['r205.txt', 'r206.txt']
    # T0 = [10]
    # cool_rate = [0.997, 0.9997]
    # init_capacity = [0.3, 1]
    # N_weight = [250, 450]
    # iter_num = [i for i in range(1)]

    files = []
    for file in file_list:
        files.append(file_path + file)

    params = list(itertools.product(files, T0, cool_rate, init_capacity, \
                                    N_weight, iter_num))
    for param in params:
        print(param)
    print('num of tasks: {}'.format(len(params)))

    print('Creating pool with %d processes\n' % n_jobs)
    with multiprocessing.Pool(n_jobs) as pool:
        results = [pool.apply_async(complex_func, param) for param in params]
        for r in results:
            while 1:
                sys.stdout.flush()
                try:
                    sys.stdout.write('\n%s' % r.get())
                    break
                except multiprocessing.TimeoutError:
                    sys.stdout.write('.')


if __name__ == "__main__":
    n_jobs = int(os.getenv('SLURM_CPUS_PER_TASK'))
    # start_time = time.time()
    main(n_jobs=n_jobs)
    # end_time = time.time()
    # print(n_jobs, end_time - start_time)

• run.sh
  • 想要实现的效果：我要测试多种参数组合，彼此之间互不干扰，循环可以并行，每个核跑一组参数组合
  • multiprocessing是基于进程的并行，而非线程，本地跑没啥问题，在SLURM系统并行核数应设置--cpus-per-task而不是--ntasks，若-n 56 -c 1则会所有的核都跑同一参数组合，非常地蠢
  • 登录计算结点，用top查看CPU利用率，对-n 1 -c 56的设置基本都100%，对-n 56 -c 1的设置CPU利用率就3%左右
  • multiprocessing只能单结点并行，python命令前加srun
  • 定向标准错误输出--error，和标准输出分开，便于后续数据处理
  • 线程 vs. 进程：傻傻分不清楚
• 参考：
  • https://researchcomputing.princeton.edu/support/knowledge-base/python#multiprocessing
  • https://sulis-hpc.github.io/gettingstarted/batchq/singlenode.html#python-multiprocessing

#!/bin/bash
#SBATCH --job-name=r108-111
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=56
#SBATCH --error=error.log

module purge
module load anaconda/3
srun python main-process.py