使用argparse和mpi4py解析参数

Question

我想在多个并行进程中运行一个Python脚本，并且需要传递命令行参数。我在Python中使用了argparse模块，但有时候用起来有点麻烦。如果我没有指定正确的参数，所有进程都会抱怨，这样我就会收到很多重复的错误信息。

我试过让只有进程0来解析参数，然后把结果广播给其他进程，但如果解析失败，其他进程就会卡住，什么都不能广播。

我该如何解析命令行参数，并在解析失败时打印出易于理解的错误信息呢？

Answer 1

如果出现错误，通常最简单的办法就是让程序直接停止，而不是费劲去处理清理工作。在你的情况下，你可以让最开始的那个进程（也就是编号为0的进程）调用停止命令，这样其他所有进程也会跟着退出：

comm.abort()

这样一来，就不用担心大家去对比结果了，所有进程都会自动停止。

Answer 2

我需要的额外步骤是在第0个进程中把参数解析的部分放在一个try/finally结构里。在finally块里，向其他进程发送一个信息。如果参数解析失败，你就发送None，这样它们就可以安静地退出了。

from mpi4py import MPI
from time import sleep
import argparse

def parseOptions(comm):
    parser = argparse.ArgumentParser(
        description='Print some messages.')

    parser.add_argument('iteration_count', help='How many times', type=int)
    parser.add_argument('message',
                        help='What to say',
                        nargs=argparse.OPTIONAL,
                        default='Hello, World!')

    args = None
    try:
        if comm.Get_rank() == 0:
            args = parser.parse_args()
    finally:
        args = comm.bcast(args, root=0)

    if args is None:
        exit(0)
    return args

def main():
    comm = MPI.COMM_WORLD  # @UndefinedVariable
    rank = comm.Get_rank()
    size = comm.Get_size()

    args = parseOptions(comm)

    if rank == 0:
        print args.message

    for i in range(args.iteration_count):
        if i%size == rank:
            print '{} in rank {} started.'.format(i, rank)
            sleep(.5)
            print '...'
            sleep(.5)
            print '{} in rank {} ended.'.format(i, rank)

if __name__ == '__main__':
    main()

我用类似这样的命令来运行代码：

mpirun -np 4 python scratch.py 13