mpi4pyを試してみる

Pythonの並列処理のためのライブラリであるmpi4pyを試してみました．

公式のチュートリアルがわかりやすかったので，かなり参考になります．

Tutorial — MPI for Python 3.0.0 documentation

僕はいつも通り，Docker imageから作っていきます．

FROM python

# install via apt
USER root
RUN apt-get update
RUN pip install -U pip

# mpi4py
RUN apt-get install -y libopenmpi-dev
RUN pip install mpi4py

WORKDIR /root/work

ビルド済みイメージをDockerHubにあげておきました．

https://hub.docker.com/r/denden047/mpi4py/

なので，こんな感じで実行できます．

$ docker run --rm denden047/mpi4py /bin/bash

Hello Worldのプログラムを書いてみました．
これをmain.pyという名前で保存します．

from mpi4py import MPI

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

    print(f"Hello, world! from rank {rank} out of {size}")


if __name__ == "__main__":
    main()

っでこれを実行してみますが， 並列実行させるために，mpiexecコマンドを利用します．
具体的には以下のコマンドを実行します．

$ docker run --rm denden047/mpi4py /bin/bash -c "mpiexec --allow-run-as-root -np 4 python main.py"
Hello, world! from rank 1 out of 4
Hello, world! from rank 3 out of 4
Hello, world! from rank 0 out of 4
Hello, world! from rank 2 out of 4

Dockerコンテナで動かしているので，rootでも実行できるように， --allow-run-as-rootオプションを付けています．
また，-np 4としているので，4つのプロセスで並列に実行されています．

ここまでで，大体の感覚は掴めるので，公式のチュートリアルをやっていきましょう！