docker (nvidia-docker) を使ってマルチノードで ChainerMN を実行する方法(仮)

docker (nvidia-docker) を使ってマルチノードで ChainerMN を実行する方法(仮)

シングルノード

シングルノードなら特に難しい所はない。docker コンテナ内に OpenMPI や ChainerMN をインストールして、nvidia-docker run で mpiexec すれば良い。

Dockerfile

FROM nvidia/cuda:9.0-cudnn7-devel-ubuntu16.04

RUN apt-get update \
  && apt-get install -y --no-install-recommends python3-dev python3-pip \
  && pip3 install setuptools && pip3 install --upgrade pip \
  && apt-get install -y --no-install-recommends git m4 autoconf automake libtool flex \
  && apt-get clean
RUN git clone --depth=1 -b v3.0.0 https://github.com/open-mpi/ompi.git /tmp/ompi \
  && cd /tmp/ompi \
  && ./autogen.pl \
  && ./configure --with-cuda \
  && make -j4 \
  && make install \
  && rm -rf /tmp/ompi

ADD nccl-repo-ubuntu1604-2.0.5-ga-cuda9.0_3-1_amd64.deb /tmp/
RUN dpkg -i /tmp/nccl-repo-ubuntu1604-2.0.5-ga-cuda9.0_3-1_amd64.deb \
  && apt-get update \
  && apt-get install -y libnccl2 libnccl-dev \
  && apt-get clean

ENV LD_LIBRARY_PATH=/usr/local/lib:$LD_LIBRARY_PATH
ENV CPATH=/usr/local/include:$CPATH
ENV CUDA_PATH=/usr/local/cuda
ENV PATH=$CUDA_PATH/bin:$PATH
ENV CPATH=$CUDA_PATH/include:$CPATH
ENV LD_LIBRARY_PATH=$CUDA_PATH/lib64:$CUDA_PATH/lib:$LD_LIBRARY_PATH

RUN pip3 install cupy==3.0.0a1 chainer==4.0.0a1
RUN pip3 install cython && pip3 install chainermn

nccl-repo-ubuntu1604-2.0.5-ga-cuda9.0_3-1_amd64.deb は NVIDIA のサイトから手動でダウンロードしてくる必要がある

ビルド

mkdir docker
vim docker/Dockerfile
cd docker
nvidia-docker build -t chainermn .

実行

docker コンテナ内で mpiexec する

$ git clone https://github.com/chainer/chainermn # examples が欲しいだけ
$ nvidia-docker run -v $(pwd):/mnt $HOME/.chainer:/root/.chainer chainermn \
  mpiexec --allow-run-as-root -n 8 \
  python3 /mnt/chainermn/examples/mnist/train_mnist.py --gpu

-n 8 はGPU数にあわせて調整

マルチノード

mpiexec すると、ホストそれぞれに ssh して、mpiexec の引数のコマンドを実行する。
ホストAのコンテナから、ホストBのコンテナにsshできるようにする必要がある。

Dockerfile

FROM nvidia/cuda:9.0-cudnn7-devel-ubuntu16.04

RUN apt-get update \
  && apt-get install -y --no-install-recommends python3-dev python3-pip \
  && pip3 install setuptools && pip3 install --upgrade pip \
  && apt-get install -y --no-install-recommends git m4 autoconf automake libtool flex \
  && apt-get clean
RUN git clone --depth=1 -b v3.0.0 https://github.com/open-mpi/ompi.git /tmp/ompi \
  && cd /tmp/ompi \
  && ./autogen.pl \
  && ./configure --with-cuda \
  && make -j4 \
  && make install \
  && rm -rf /tmp/ompi

ADD nccl-repo-ubuntu1604-2.0.5-ga-cuda9.0_3-1_amd64.deb /tmp/
RUN dpkg -i /tmp/nccl-repo-ubuntu1604-2.0.5-ga-cuda9.0_3-1_amd64.deb \
  && apt-get update \
  && apt-get install -y libnccl2 libnccl-dev \
  && apt-get clean

ENV LD_LIBRARY_PATH=/usr/local/lib:$LD_LIBRARY_PATH
ENV CPATH=/usr/local/include:$CPATH
ENV CUDA_PATH=/usr/local/cuda
ENV PATH=$CUDA_PATH/bin:$PATH
ENV CPATH=$CUDA_PATH/include:$CPATH
ENV LD_LIBRARY_PATH=$CUDA_PATH/lib64:$CUDA_PATH/lib:$LD_LIBRARY_PATH

RUN apt-get install -y --no-install-recommends openssh-server \
  && mkdir /var/run/sshd \
  && sed -i 's/Port 22/Port 10022/' /etc/ssh/sshd_config \
  && ssh-keygen -t rsa -N '' -f /root/.ssh/id_rsa \
  && cp /root/.ssh/id_rsa.pub /root/.ssh/authorized_keys \
  && chmod 600 /root/.ssh/authorized_keys \
  && echo 'Host *' > /root/.ssh/config \
  && echo '  Port 10022' >> /root/.ssh/config \
  && echo '  StrictHostKeyChecking no' >> /root/.ssh/config \
  && echo '  UserKnownHostsFile /dev/null' >> /root/.ssh/config

RUN echo 'export CUDA_PATH="/usr/local/cuda"\n\
export PATH="$CUDA_PATH/bin:$PATH"\n\
export CPATH="$CUDA_PATH/include:$CPATH"\n\
export LD_LIBRARY_PATH="$CUDA_PATH/lib64:$CUDA_PATH/lib:$LD_LIBRARY_PATH"\n\
export PATH="/usr/local/nvidia/bin:$PATH"\n\
export LD_LIBRARY_PATH="/usr/local/nvidia/lib:/usr/local/nvidia/lib64:$LD_LIBRARY_PATH"\n\
export NCCL_SOCKET_IFNAME=^docker0' > /etc/profile.d/chainermn.sh

RUN pip3 install cupy==3.0.0a1 chainer==4.0.0a1
RUN pip3 install cython && pip3 install chainermn

EXPOSE 10022

シングルノード用の Dockerfile に対して、sshd のセットアップと、ssh したときに必要な環境変数を設定できるように /etc/profile.d/chainermn.sh ファイルの作成を追加している。

ホストに向けたMPI通信をコンテナに簡易にフォワードするため、またネットワーク速度向上のため、docker run 時には --net=host を指定する。ホストの sshd 22 番ポートと被らないように、コンテナの sshd は 10022 番ポートで起動するようにしている。

再掲するが、nccl-repo-ubuntu1604-2.0.5-ga-cuda9.0_3-1_amd64.deb は NVIDIA のサイトから手動でダウンロードしてくる必要がある

ビルド

mkdir docker
vim docker/Dockerfile
cd docker
nvidia-docker build -t chainermn .

※ ここでビルドしたイメージを全ホストで利用すること。今回の Dockerfile では、それぞれのホストで docker build すると、別の ssh 鍵ができてコンテナ間 ssh できなくなってしまう。

実行

従属ホスト(ホストB)

コンテナで sshd を起動だけしておく

$ git clone https://github.com/chainer/chainermn # examples が欲しいだけ
$ nvidia-docker run --rm -it --net=host \
  -v $(pwd):/mnt -v $HOME/.chainer:/root/.chainer chainermn \
  /bin/bash -lc '/etc/init.d/ssh start && /bin/bash'

マスターホスト(ホストA)

hostfile を用意。cpu=N に並列実行数を指定する。

hostfile
ホストA cpu=8
ホストB cpu=8

実行

$ git clone https://github.com/chainer/chainermn # examples が欲しいだけ
$ nvidia-docker run --rm -it --net=host \
  -v $(pwd):/mnt -v $HOME/.chainer:/root/.chainer chainermn \
  /bin/bash -lc '/etc/init.d/ssh start && \
  mpiexec --allow-run-as-root --mca btl_tcp_if_exclude docker0,lo --hostfile /mnt/hostfile \
  /bin/bash -lc "python3 /mnt/chainermn/examples/mnist/train_mnist.py --gpu"'

今回は、実行ファイル(chainermn の examples)を各ホストに配置しているが、Docker イメージに含める方法もあるとは思う。chainermn のルールとして、全ホストに実行ファイルがないとダメ。

便利スクリプト

全従属ホストでコンテナを立ち上げて、マスターホストで mpiexec を実行する便利スクリプト。
git clone https://github.com/chainer/chainermn は実行したいスクリプトに合わせてご編集ください。
python3 /mnt/chainermn/examples/mnist/train_mnist.py --gpu も実行したいスクリプトに合わせてご編集ください。

run_chainermn.sh
#!/bin/bash

PREPARE_SCRIPT=/tmp/run_chainermn_prepare.sh
SLAVE_SCRIPT=/tmp/run_chainermn_slave.sh
MASTER_SCRIPT=/tmp/run_chainermn_master.sh

cat <<'EOF' > $PREPARE_SCRIPT
git clone https://github.com/chainer/chainermn chainermn
EOF

cp $PREPARE_SCRIPT $SLAVE_SCRIPT
cat <<'EOF' >> $SLAVE_SCRIPT
nvidia-docker run --rm -it --net=host \
  -v $(pwd):/mnt -v $HOME/.chainer:/root/.chainer chainermn \
  /usr/sbin/sshd -D
EOF

cp $PREPARE_SCRIPT $MASTER_SCRIPT
cat <<'EOF' >> $MASTER_SCRIPT
nvidia-docker run --rm -it --net=host \
  -v $(pwd):/mnt -v $HOME/.chainer:/root/.chainer chainermn \
  /bin/bash -lc '/etc/init.d/ssh start && \
  mpiexec --allow-run-as-root --mca btl_tcp_if_exclude docker0,lo --hostfile /mnt/hostfile \
  /bin/bash -lc "python3 /mnt/chainermn/examples/mnist/train_mnist.py --gpu"'
EOF

master_host=$(hostname)
slave_hosts=$(grep -v ${master_host} hostfile | awk '{print $1}')
hosts=$(cat hostfile | awk '{print $1}')
n_hosts=$(cat hostfile | wc -l)

echo "$slave_hosts" | xargs -I{} -P${n_hosts} echo scp $SLAVE_SCRIPT {}:/tmp/chainermn.sh
echo "$slave_hosts" | xargs -I{} -P${n_hosts} scp $SLAVE_SCRIPT {}:/tmp/chainermn.sh
echo cp $MASTER_SCRIPT /tmp/chainermn.sh
cp $MASTER_SCRIPT /tmp/chainermn.sh

function cleanup() {
  echo "$hosts" | xargs -I{} -P${n_hosts} echo ssh {} docker stop chainermn
  echo "$hosts" | xargs -I{} -P${n_hosts} ssh {} docker stop chainermn
}
trap cleanup EXIT

echo "$hosts" | xargs -I{} -P${n_hosts} echo ssh -tt {} /bin/bash /tmp/chainermn.sh
echo "$hosts" | xargs -I{} -P${n_hosts} ssh -tt {} /bin/bash /tmp/chainermn.sh

hostfile を用意して、実行

bash run_chainermn.sh

Ctrl-C するとリモートホストの docker コンテナが stop される (ちょっと時間かかるのでお待ちください)

トラブルシューティング

See also AWS GPU インスタンスにおける ChainerMN の分散効率に対する評価#トラブルシューティング

Unexpected end of /proc/mounts line `overlay

コンテナ内で mpiexec すると以下のような warning が出る

Unexpected end of /proc/mounts line `overlay / overlay rw,relatime,lowerdir=/var/lib/docker/overlay2/l/ETGN4JWRCAAOR5SDTEX4LJN7Z6:/var/lib/docker/overlay2/l/6GKXZXZ25OYNPLCCTBNRKHDQG6:/var/lib/docker/overlay2/l/OUCL4BE45M72TBGXLVGOUGULMH:/var/lib/docker/overlay2/l/DHW4TE7ZW55P5UOZBX547LDZ77:/var/lib/docker/overlay2/l/YSIOMV4IYP3T2YVN77UFZFQ5XI:/var/lib/docker/overlay2/l/F5AUK2J5LITCBCVWGBCJD3LETY:/var/lib/docker/overlay2/l/ZPR6AAOCLV6SOQA7I6DIDMRY7G:/var/lib/docker/overlay2/l/UE6YIXVHWKAVHG2UJAMUWEUE5Y:/var/lib/docker/overlay2/l/UY4QMTUE5XDUS'
Unexpected end of /proc/mounts line `INR6Y7546V2HV:/var/lib/docker/overlay2/l/F6JQBSMJYZOOQ3QIHYUG3VVHMW:/var/lib/docker/overlay2/l/VIP7WZGLNKZMRDKH6KBINRIJSW:/var/lib/docker/overlay2/l/TDLKNGLFEVYF4SI73USXIW4A4A:/var/lib/docker/overlay2/l/VZ6RLPGCUQFQ4K5HHWOSZXXK75:/var/lib/docker/overlay2/l/UVHFPFXQRL3XJS32BZRUMP6QEP:/var/lib/docker/overlay2/l/STHMTEHKNZWTCGWFDSYT6J3ND4:/var/lib/docker/overlay2/l/L7ANIC2YQAOSWMBVRR6QZS5SZ2:/var/lib/docker/overlay2/l/DC7BAJG5WLKXZVJJM73DGCNIKE:/var/lib/docker/overlay2/l/4W423XWKPJV3WINYECSJQ7JBY3:/var/lib/do'
Unexpected end of /proc/mounts line `cker/overlay2/l/FJF6DA25TZWDVGOMCZ7J7PRT7Q:/var/lib/docker/overlay2/l/4PHONTLPJTZAVDXCRN4W2BCQPW:/var/lib/docker/overlay2/l/RU7ZODNEXSEKWG5J4UU53AYGDB:/var/lib/docker/overlay2/l/6VF7XXRRSZC2HH24Z3JJRVZBVA:/var/lib/docker/overlay2/l/M3WQNRNV5IA3KXFBHV5GY3KD6U:/var/lib/docker/overlay2/l/ML2ZZUHFRV5HELDLE5O5NDWPLX:/var/lib/docker/overlay2/l/T535AZCEOGAX5BT2XGDEBQ3RH3:/var/lib/docker/overlay2/l/4SZPXUFR4LU2WRUTITUPSOPIJE,upperdir=/var/lib/docker/overlay2/333f260a32ed9645c209c166c9e2d6ce807ec87b0ee351a86dc365fa'

https://stackoverflow.com/questions/46138549/docker-openmpi-and-unexpected-end-of-proc-mounts-line この問題らしいが、無視して動作したので無視した。

今後の発展

マルチノードでの docker コンテナの管理を、docker swarn や kubernetes のようなコンテナクラスタ管理ツールを使って行えるようにすると立ち上げるのが簡単になりそう。

一方で、うまく MPI や GPU を扱えるようにするには設定が難しそう (少なくも GKE のようなマネージドサービスを使うのは難しそう)。https://github.com/NLKNguyen/alpine-mpich が参考になりそう (docker swarn を使っている)。