如何在 HPC 集群上将所有已分配的节点与 python 一起使用

Question

我有一个安装了 SLURM 的 HPC 集群。我可以为自己正确分配节点和核心。我希望能够使用所有分配的内核，而不管它们位于哪个节点。正如我在这个线程Using the multiprocessing module for cluster computing 中看到的那样，multiprocessing 无法实现这一点。

我的脚本看起来像这样（过于简化的版本）：

def func(input_data):
    #lots of computing
    return data

parallel_pool = multiprocessing.Pool(processes=300)
returned_data_list = []
for i in parallel_pool.imap_unordered(func, lots_of_input_data)
    returned_data_list.append(i)
# Do additional computing with the returned_data
....

这个脚本工作得很好，但是正如我提到的，多处理对我来说不是一个好工具，因为即使 SLURM 为我分配了 3 个节点，多处理也只能使用一个。据我了解，这是多处理的限制。

我可以使用 SLURM 的srun 协议，但是我们必须执行相同的脚本 N 次，并且我需要对并行进程的输出进行额外的计算。我当然可以将输出存储在某个地方，然后将它们重新放入，但必须有一些更优雅的解决方案。

在提到的线程中有类似jug 的建议，但是当我阅读它时，我还没有找到适合自己的解决方案。

也许py4mpi 可以成为我的解决方案？那个教程看起来很混乱，我也没有找到解决我的问题的具体解决方案。 （与 mpi 并行运行一个函数，然后继续执行脚本）。

我尝试了subprocess 调用，但它们的工作方式似乎与multiprocess 调用相同，因此它们只在一个节点上运行。我还没有找到任何确认，所以这只是我的试错猜测。

我该如何克服这个问题？目前我可以使用超过 300 个核心，但一个节点只有 32 个，所以如果我能找到解决方案，那么我的项目运行速度将快近 10 倍。

谢谢

Answer 1

经过很多麻烦scoop 是解决我问题的库。