在 Slurm 中使用 Python 多处理，以及我需要哪种 ntasks 或 ncpus 组合

Question

我正在尝试在 slurm 集群上运行 python 脚本，并且我正在使用 python 的内置multiprocessing模块。

我使用了一个非常简单的设置，出于测试目的，示例是：

len(arg_list)
Out[2]: 5

threads = multiprocessing.Pool(5)
output = threads.map(func, arg_list)

Sofunc中的 5 个参数并行应用 5 次arg_list。我想知道的是如何在 slurm 中分配正确数量的 cpu/任务以使其按预期工作。这是我的 slurm 批处理脚本的相关部分的样子：

#!/bin/bash

# Runtime and memory
#SBATCH --time=90:00:00
#SBATCH --mem-per-cpu=2G

# For parallel jobs
#SBATCH --cpus-per-task=10
##SBATCH --nodes=2
#SBATCH --ntasks=1
##SBATCH --ntasks-per-node=4  

#### Your shell commands below this line ####

srun ./script_wrapper.py 'test'

如您所见，目前我有ntasks=1和cpus-per-task=10。请注意，func 的主要部分包含一个倾向于在两个内核上运行的 scipy 例程（即使用 200% 的 cpu 使用率，这就是我想要 10 个 cpu 而不是 5 个的原因）。

这是为我的目的分配资源的正确方法吗，因为目前这项工作需要比预期更长的时间（更像是它在单个线程中运行）。

我需要设置ntasks=5吗？因为我从在线文档中得到的印象是ntasks=5调用srun ./script_wrapper.py 'test'5 次而不是，这不是我想要的。我的假设正确吗？

另外，有没有一种方法可以轻松检查诸如 CPU 使用率和 multiprocessing.Pool 调用的 Python 任务的所有进程 ID 之类的东西？目前我正在尝试使用sacct -u <user> --format=JobID,JobName,MaxRSS,Elapsed,AveCPU，但是由于某种原因，AveCPU和MaxRSS字段总是为空（？），虽然我将第一个脚本视为一个进程，但我没有看到其他 5 个应该由多处理调用的脚本。例子：

       JobID    JobName     MaxRSS    Elapsed     AveCPU 
------------ ---------- ---------- ---------- ---------- 
16260892             GP              00:13:07            
16260892.0   script_wr+              00:13:07            

Answer 1

我认为你的 Slurm 任务分配是正确的。Python 的多处理只能在一台机器上运行，在我看来，您在一个节点上正确分配了 10 个 CPU。导致该问题的原因可能是多重处理Pool.map默认对输入列表的“块”起作用，而不是一次对一个元素起作用。这样做是为了在任务较短时最大限度地减少开销。要强制多重处理一次处理列表中的一个元素，请将映射的 chunksize 参数设置为 1，例如

threads.map(func, arglist, 1)

有关更多信息，请参阅多处理文档。

因为您说您正在使用 SciPy 的多线程版本，所以您可能还想检查底层库的相关线程级别。例如，如果您的 SciPy 是针对 Intel Math Kernel Library 构建的，请尝试设置OMP_NUM_THREADS和MKL_NUM_THREADS 环境变量，以确保每个进程使用不超过 2 个线程，并充分利用（而不是过度使用）分配的 SLURM 资源。

编辑： sacct 只会为您提供由 srun 直接启动的任何进程的运行时间，而不是任何子进程的运行时间。因此，在您的情况下，您只能从单个 srun 命令获得一个进程。要监控子流程，您可能需要研究在系统级别运行的监控工具，而不是通过 Slurm。