如何在多个节点上分配多进程CPU使用率？

Question

我正在尝试使用多处理在 HPC 上运行作业。每个进程的峰值内存使用量约为 44GB。我可以使用的作业类允许使用1-16个节点，每个节点有32个CPU和124GB内存。因此，如果我想尽快运行代码（并且在最大挂起时间限制内），我应该能够在每个节点上运行 2 个 CPU，在所有 16 个节点上最多运行 32 个 CPU。但是，当我指定 mp.Pool(32) 时，作业很快就超出了内存限制，我认为这是因为节点上使用了两个以上的 CPU。

我的本能反应是在运行 python 脚本的 pbs 脚本中指定 2 个 CPU 作为最大数量，但是系统上不允许这种配置。我真的很感激任何见解，今天的大部分时间我都在为这个问题摸不着头脑——而且过去也曾面临并解决过类似的问题，但没有解决这里起作用的基本原理。

以下两个脚本的简化版本：

#!/bin/sh 

#PBS -l select=16:ncpus=32:mem=124gb
#PBS -l walltime=24:00:00

module load anaconda3/personal
source activate py_env

python directory/script.py

#!/usr/bin/env python
# -*- coding: utf-8 -*- 

import numpy as np
import pandas as pd
import multiprocessing as mp

def df_function(df, arr1, arr2):
    df['col3'] = some_algorithm(df, arr1, arr2)
    return df

def parallelize_dataframe(df, func, num_cores):
    df_split = np.array_split(df, num_cores)
    with mp.Pool(num_cores, maxtasksperchild = 10 ** 3) as pool:
        df = pd.concat(pool.map(func, df_split))
    return df

def main():
    # Loading input data    
    direc = '/home/dir1/dir2/'
    file = 'input_data.csv'
    a_file = 'array_a.npy'
    b_file = 'array_b.npy'
    df = pd.read_csv(direc + file)
    a = np.load(direc + a_file)
    b = np.load(direc + b_file)

    # Globally defining function with keyword defaults
    global f
    def f(df):
        return df_function(df, arr1 = a, arr2 = b)

    num_cores = 32 # i.e. 2 per node if evenly distributed.
    # Running the function as a multiprocess:
    df = parallelize_dataframe(df, f, num_cores)
    # Saving:
    df.to_csv(direc + 'outfile.csv', index = False)

if __name__ == '__main__':
    main()

Answer 1

要按原样运行作业，您只需请求 ncpu=32，然后在 Python 脚本中设置 num_cores = 2。显然，这会让您支付 32 个核心的费用，然后让其中 30 个核心闲置，这是一种浪费。

这里真正的问题是你当前的算法是内存限制的，而不是 CPU 限制的。您应该竭尽全力将文件块读入内存，对这些块进行操作，然后将结果块写入磁盘以便稍后组织。

幸运的是 Dask 就是为了做这种事情而构建的。第一步，您可以取出 parallelize_dataframe 函数，并直接使用 dask.dataframe 和 dask.array：

#!/usr/bin/env python
# -*- coding: utf-8 -*- 

import dask.dataframe as dd
import dask.array as da


def main():
    # Loading input data    
    direc = '/home/dir1/dir2/'
    file = 'input_data.csv'
    a_file = 'array_a.npy'
    b_file = 'array_b.npy'
    df = dd.read_csv(direc + file, blocksize=25e6)
    a_and_b = da.from_np_stack(direc)

    df['col3'] = df.apply(some_algorithm, args=(a_and_b,))
    
    # dask is lazy, this is the only line that does any work    
    # Saving:
    df.to_csv(
        direc + 'outfile.csv',
        index = False,
        compute_kwargs={"scheduler": "threads"},  # also "processes", but try threads first
    )

if __name__ == '__main__':
    main()

这需要对 some_algorithm 和 to_csv 和 from_np_stack 工作方式略有不同，但是您将能够在自己的笔记本电脑上合理地运行这个东西，并且它将扩展到您的集群硬件。您可以使用分布式调度程序甚至将其直接部署到您的集群 dask-jobqueue。