FROM ubuntu:18.04

# Python 2.7 or 3.6 is supported by Ubuntu Bionic out of the box
ARG python=3.7
ENV PYTHON_VERSION=${python}

# Set default shell to /bin/bash
SHELL ["/bin/bash", "-cu"]

RUN apt-get update && apt-get install -y --allow-downgrades --allow-change-held-packages --no-install-recommends \
        build-essential \
        cmake \
        g++-4.8 \
        git \
        curl \
        vim \
        wget \
        ca-certificates \
        libjpeg-dev \
        libpng-dev \
        python${PYTHON_VERSION} \
        python${PYTHON_VERSION}-dev \
        python3-distutils \
        librdmacm1 \
        libibverbs1 \
        ibverbs-providers

# Install OpenSSH for MPI to communicate between containers
RUN apt-get install -y --no-install-recommends openssh-client openssh-server && \
    mkdir -p /var/run/sshd
# Allow OpenSSH to talk to containers without asking for confirmation
# by disabling StrictHostKeyChecking.
# mpi-operator mounts the .ssh folder from a Secret. For that to work, we need
# to disable UserKnownHostsFile to avoid write permissions.
# Disabling StrictModes avoids directory and files read permission checks.
RUN sed -i 's/[ #]\(.*StrictHostKeyChecking \).*/ \1no/g' /etc/ssh/ssh_config && \
    echo "    UserKnownHostsFile /dev/null" >> /etc/ssh/ssh_config && \
    sed -i 's/#\(StrictModes \).*/\1no/g' /etc/ssh/sshd_config

# Install Open MPI
RUN mkdir /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

# Download imagenet tiny for data
RUN apt-get -q update && apt-get -q install -y wget unzip sudo
RUN wget -q http://cs231n.stanford.edu/tiny-imagenet-200.zip && mkdir -p /workspace/data && unzip -q tiny-imagenet-200.zip -d /workspace/data && rm tiny-imagenet-200.zip

RUN if [[ "${PYTHON_VERSION}" == "3.6" ]]; then \
        apt-get install -y python${PYTHON_VERSION}-distutils; \
    fi
RUN ln -s /usr/bin/python${PYTHON_VERSION} /usr/bin/python

RUN curl -O https://bootstrap.pypa.io/get-pip.py && \
    python get-pip.py && \
    rm get-pip.py

ENV TENSORFLOW_VERSION=1.14.0
# Install TensorFlow, Keras, PyTorch and MXNet
RUN pip install future typing
RUN pip install numpy \
        tensorflow==${TENSORFLOW_VERSION} \
        keras \
        h5py

ENV PYTORCH_VERSION=1.4.0
ENV TORCHVISION_VERSION=0.5.0
RUN pip install torch==${PYTORCH_VERSION} torchvision==${TORCHVISION_VERSION}

#ENV MXNET_VERSION=1.6.0
#RUN pip install mxnet==${MXNET_VERSION}

ENV HOROVOD_VERSION=0.27.0
# Install Horovod
#RUN HOROVOD_WITH_TENSORFLOW=0 HOROVOD_WITH_PYTORCH=1 HOROVOD_WITH_MXNET=0 \
#    pip install --no-cache-dir horovod==${HOROVOD_VERSION}


RUN python -m pip install numpy --upgrade
RUN python -m pip install torch --upgrade
RUN python -m pip install torchvision --upgrade
RUN python -m pip install python-etcd --upgrade

RUN mkdir /app
WORKDIR /workspace

ADD *.py /app/