本当は怖いHPC

HPC屋の趣味＆実益ブログ

MPIでRank順に出力する

mpi python

小ネタ。 MPIのプログラムから printf などを使ってデバッグ出力をする場合、全プロセスから一斉に同じ出力をしたときの順序は保証されていません。

例えば、下のようなプログラムを考えます。

#include <stdio.h>
#include <mpi.h>

int main(int argc, char **argv) {
    int rank;

    MPI_Init(&argc, &argv);
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    printf("I'm rank %d\n", rank);
    MPI_Finalize();
}

# 実行例
$ mpiexec -n 10 ./a.out
I'm rank 4
I'm rank 5
I'm rank 6
I'm rank 7
I'm rank 9
I'm rank 0
I'm rank 1
I'm rank 8
I'm rank 2
I'm rank 3

この順序を制御して、例えばRank順に出力したいということは頻繁にあります。MPIを使っている人なら、だいたい自前で書いてしまう処理で、以下のように簡単に書けます。

    for (i = 0; i < size; i++) {
        if (i == rank) {
            printf("I'm rank %d\n", rank);
        }
        MPI_Barrier(MPI_COMM_WORLD);
    }
# 実行例
$ mpiexec -n 10 ./a.out
I'm rank 0
I'm rank 1
I'm rank 2
I'm rank 3
I'm rank 4
I'm rank 5
I'm rank 6
I'm rank 7
I'm rank 8
I'm rank 9

次に、mpi4pyを使ったPythonスクリプトでも簡単に書きたいなーと思ったので書いてみました。 ループを書くのは面倒なので with 構文で書けないかなーと思い、やってみたら意外と簡単にできました。 with便利。

class RankOrdered(object):
    def __init__(self, comm):
        self._comm = comm

    def __enter__(self):
        for i in range(0, comm.rank):
            comm.barrier()
    
    def __exit__(self, exception_type, exception_value, traceback):
        for i in range(comm.size - comm.rank):
            comm.barrier()

# 利用例
from mpi4py import MPI
with RankOrdered(MPI.COMM_WORLD):
    print("rank {}".format(comm.rank))
【広告】