Dockerfile-openmpi

gistfile1.txt

ROM nvidia/cuda:9.2-cudnn7-devel

ARG OPENMPI_VERSION="3.1.3"
ARG NCCL_VERSION="v2.3.7-1"
ARG CHAINER_VERSION="5.0.0"

# Install basic dependencies and locales
RUN apt-get update && apt-get install -yq --no-install-recommends --allow-change-held-packages \
      locales wget sudo ca-certificates ssh build-essential devscripts debhelper git \
      openssh-client openssh-server \
      libnccl2 libnccl-dev \
      python3-dev \
      python3-pip \
      python3-wheel \
      python3-setuptools && \
    rm -rf /var/lib/apt/lists/* /var/cache/apt/archives/* && \
    echo "en_US.UTF-8 UTF-8" > /etc/locale.gen && locale-gen

# Install OpenMPI with cuda
RUN cd /tmp && \
  wget -q https://www.open-mpi.org/software/ompi/v${OPENMPI_VERSION%\.*}/downloads/openmpi-$OPENMPI_VERSION.tar.bz2 && \
  tar -xjf openmpi-$OPENMPI_VERSION.tar.bz2 && \
  cd /tmp/openmpi-$OPENMPI_VERSION && \
  ./configure --prefix=/usr --with-cuda && make -j2 && make install && rm -r /tmp/openmpi-$OPENMPI_VERSION* && \
  ompi_info --parsable --all | grep -q "mpi_built_with_cuda_support:value:true"

# Install ChainerMN
RUN pip3 install --no-cache-dir mpi4py

# Set default NCCL parameters
RUN echo NCCL_DEBUG=INFO >> /etc/nccl.conf

# Install OpenSSH for MPI to communicate between containers
RUN mkdir -p /var/run/sshd

# Allow OpenSSH to talk to containers without asking for confirmation
RUN cat /etc/ssh/ssh_config | grep -v StrictHostKeyChecking > /etc/ssh/ssh_config.new && \
    echo "    StrictHostKeyChecking no" >> /etc/ssh/ssh_config.new && \
    mv /etc/ssh/ssh_config.new /etc/ssh/ssh_config
    
RUN pip3 install cupy==5.1.0 chainer==5.1.0