# Copyright 2019 Bytedance Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# =============================================================================

FROM ubuntu:16.04

ARG REGION

RUN rm -f /tmp/pip.conf &&\
    echo -e '[global]\nindex-url = https://pypi.douban.com/simple' >> /tmp/pip.conf

RUN if [ "x$REGION" = "xchina" ]; then mkdir -p ~/.pip && mv /tmp/pip.conf ~/.pip/; fi

ENV LD_LIBRARY_PATH /root/incubator-mxnet/lib/:/usr/local/lib:$LD_LIBRARY_PATH

# To enable RDMA, add `USE_RDMA=1` to `SERVER_BUILD_OPTS` below.
ENV SERVER_BUILD_OPTS "USE_BLAS=openblas USE_MKL=1 USE_DIST_KVSTORE=1"
ENV BYTEPS_SERVER_MXNET_PATH /root/incubator-mxnet
ENV MXNET_SERVER_LINK https://github.com/bytedance/incubator-mxnet

ENV BYTEPS_BASE_PATH /usr/local
ENV BYTEPS_PATH $BYTEPS_BASE_PATH/byteps
ENV BYTEPS_GIT_LINK https://github.com/bytedance/byteps

RUN apt-get update && apt-get install -y --allow-downgrades --allow-change-held-packages --no-install-recommends \
        build-essential \
        ca-certificates \
        git \
        curl \
        wget \
        vim \
        libopenblas-dev \
        liblapack-dev \
        libopencv-dev \
        python \
        python-dev \
        python-setuptools \
        libjemalloc-dev \
        graphviz \
        cmake \
        libjpeg-dev \
        libpng-dev \
        iftop \
        lsb-release

RUN apt-get update &&\
    apt-get -y install python-pip &&\
    pip install --upgrade pip

RUN pip --no-cache-dir install \
        matplotlib \
        numpy==1.15.2 \
        scipy \
        sklearn \
        pandas \
        graphviz==0.9.0 \
        mxboard \
        tensorboard==1.0.0a6

WORKDIR /root/

RUN git clone --single-branch --branch byteps --recurse-submodules $MXNET_SERVER_LINK

RUN cd $BYTEPS_SERVER_MXNET_PATH && \
    make clean_all && make -j16 $SERVER_BUILD_OPTS

RUN cd $BYTEPS_SERVER_MXNET_PATH && \
    cd python && \
    python setup.py build && \
    python setup.py install &&\
    python setup.py bdist_wheel

# fix git source code
RUN cd $BYTEPS_BASE_PATH
RUN git clone --recurse-submodules $BYTEPS_GIT_LINK
RUN cd $BYTEPS_PATH/3rdparty/ps-lite
RUN git checkout byteps
RUN git pull origin byteps
RUN git reset --hard 62bbf9a1a7c489c45b65379e2b02e8a620f8a4c1
RUN cd -


# 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
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

# ssh env
COPY sync_ssh_env.sh .
RUN bash ./sync_ssh_env.sh
RUN rm -rf ./sync_ssh_env.sh