FROM nvidia/cuda:10.0-devel-ubuntu16.04

# Install OpenSSH
RUN apt-get update \
    && apt-get install -y --no-install-recommends openssh-client openssh-server \
    && apt-get clean \
    && rm -rf /var/lib/apt/lists/*

# Install OpenMPI
RUN apt-get update \
    && apt-get install -y --no-install-recommends wget\
    && apt-get clean \
    && rm -rf /var/lib/apt/lists/*
RUN mkdir -p /tmp/openmpi \
    && cd /tmp/openmpi \
    && wget https://www.open-mpi.org/software/ompi/v4.0/downloads/openmpi-4.0.0.tar.gz \
    && tar zxf openmpi-4.0.0.tar.gz \
    && cd openmpi-4.0.0 \
    && ./configure --enable-orterun-prefix-by-default \
    && make -j $(nproc) all \
    && make install \
    && ldconfig \
    && rm -rf /tmp/openmpi

# Install Python3. Ubuntu 16.04 has Python 3.5
# python3-dev and cmake is for building Horovod
RUN apt-get update \
    && apt-get install -y --no-install-recommends \
         python3 python3-pip python3-setuptools python3-dev cmake \
    && apt-get clean \
    && rm -rf /var/lib/apt/lists/*

# Install PyTorch
# Only numpy < 1.19 supports Python 3.5
RUN pip3 install -f https://download.pytorch.org/whl/torch_stable.html torch "numpy<1.19" \
    && rm -rf ~/.cache/pip

# Install Horovod with PyTorch
# Only horovod < 0.20 supports Python 3.5
RUN HOROVOD_WITH_PYTORCH=1 pip3 install "horovod[pytorch]<0.20" \
    && rm -rf ~/.cache/pip

WORKDIR /workspace