如何找到MPI(4PY)可用的核心数？

动机

我的Python程序分层次地生成MPI实例.

第一个spawn总是发生并创建4个实例.由于我的计算结构,增加这个数字是没有意义的,所以我硬编码了.

根据主程序的命令行选项,4个实例中的每个实例都会调用几乎线性扩展的外部Python软件.

我用这个外部软件叫

N=3

child=MPI.COMM_SELF.Spawn(sys.executable,args=[`external.py`],maxprocs=N)

目前,我使用N = 3,这样第一个产生的4个实例每个产生3个外部程序实例,总共产生12个实例,与我工作站上的核心数相匹配.

但是,为了便携,我想做

N_avail = #on my workstation: N_avail=12

N = N_avail/MPI.COMM_WORLD.Get_size() #on my workstation: N=12/4=3

这样就不需要对可用核心的数量进行硬编码.

这可能吗,它有意义吗？

笔记

我曾希望不指定maxprocs可以完成这项工作,就像带有out -np的mpirun产生与可用内核一样多的实例.但是,Spawn默认为maxprocs = 1.

外部库的调用是阻塞的,这就是为什么我不会(不会)从N_avail中减去第一个spawn的4个实例.

我不能只使用multiprocessing.cpu_count(),因为这只会给我当前节点上的核心(在集群设置中).我打算使用SLURM调度程序在集群上运行我的代码.

MPI provides an attribute on MPI_COMM_WORLD, MPI_UNIVERSE_SIZE, that allows the application to obtain this information in a portable manner. This attribute indicates the total number of processes that are expected. … An application typically subtracts the size of MPI_COMM_WORLD from MPI_UNIVERSE_SIZE to find out how many processes it should spawn. …

在mpi4py中,它可以打印为：

from mpi4py import MPI

version= MPI.Get_version()

print "mpi version is ",version

comm = MPI.COMM_WORLD

rank = comm.Get_rank()

size = comm.Get_size()

print "size is ",size

universe_size=comm.Get_attr(MPI.UNIVERSE_SIZE)

print "universe size is ",universe_size

mpirun -np 1 -H localhost,localhost,localhost python main.py

如果您的MPI版本高于或等于3,MPI_Info MPI_INFO_ENV可以帮助您.它有两个键可能提供一些信息：

maxprocs Maximum number of MPI processes to start.

soft Allowed values for number of processors.

要在mpi4py中使用它,您可以尝试：

soft=MPI.INFO_ENV.get("soft")

print soft

maxprocs=MPI.INFO_ENV.get("maxprocs")

print maxprocs