Created separate Dockerfiles for CPU and GPU devices, bumped TensorFl…

…ow==2.0.0 and PyTorch==1.3.0 (#1460)

Signed-off-by: Travis Addair <taddair@uber.com>

Dockerfile.cpu

@@ -0,0 +1,80 @@
+FROM ubuntu:18.04
+
+ENV TENSORFLOW_VERSION=2.0.0
+ENV PYTORCH_VERSION=1.3.0
+ENV TORCHVISION_VERSION=0.4.1
+ENV MXNET_VERSION=1.5.0
+
+# Python 2.7 or 3.6 is supported by Ubuntu Bionic out of the box
+ARG python=2.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 \
+        librdmacm1 \
+        libibverbs1 \
+        ibverbs-providers
+
+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
+
+# Install TensorFlow, Keras, PyTorch and MXNet
+RUN pip install future typing
+RUN pip install numpy \
+        tensorflow==${TENSORFLOW_VERSION} \
+        keras \
+        h5py
+RUN pip install torch==${PYTORCH_VERSION} torchvision==${TORCHVISION_VERSION}
+RUN pip install mxnet==${MXNET_VERSION}
+
+# 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
+
+# Install Horovod
+RUN HOROVOD_WITH_TENSORFLOW=1 HOROVOD_WITH_PYTORCH=1 HOROVOD_WITH_MXNET=1 \
+    pip install --no-cache-dir horovod
+
+# 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
+
+# Download examples
+RUN apt-get install -y --no-install-recommends subversion && \
+    svn checkout https://github.com/horovod/horovod/trunk/examples && \
+    rm -rf /examples/.svn
+
+WORKDIR "/examples"

Dockerfile → Dockerfile.gpu

@@ -1,9 +1,9 @@
 FROM nvidia/cuda:10.0-devel-ubuntu18.04
 
 # TensorFlow version is tightly coupled to CUDA and cuDNN so it should be selected carefully
-ENV TENSORFLOW_VERSION=1.14.0
-ENV PYTORCH_VERSION=1.2.0
-ENV TORCHVISION_VERSION=0.4.0
+ENV TENSORFLOW_VERSION=2.0.0
+ENV PYTORCH_VERSION=1.3.0
+ENV TORCHVISION_VERSION=0.4.1
 ENV CUDNN_VERSION=7.6.0.64-1+cuda10.0
 ENV NCCL_VERSION=2.4.7-1+cuda10.0
 ENV MXNET_VERSION=1.5.0
@@ -50,8 +50,9 @@ RUN pip install numpy \
         tensorflow-gpu==${TENSORFLOW_VERSION} \
         keras \
         h5py
-RUN pip install https://download.pytorch.org/whl/cu100/torch-${PYTORCH_VERSION}-$(python -c "import wheel.pep425tags as w; print('-'.join(w.get_supported()[0][:-1]))")-manylinux1_x86_64.whl \
-        https://download.pytorch.org/whl/cu100/torchvision-${TORCHVISION_VERSION}-$(python -c "import wheel.pep425tags as w; print('-'.join(w.get_supported()[0][:-1]))")-manylinux1_x86_64.whl
+
+RUN pip install https://download.pytorch.org/whl/cu100/torch-${PYTORCH_VERSION}%2Bcu100-$(python -c "import wheel.pep425tags as w; print('-'.join(w.get_supported()[0][:-1]))")-linux_x86_64.whl \
+        https://download.pytorch.org/whl/cu100/torchvision-${TORCHVISION_VERSION}%2Bcu100-$(python -c "import wheel.pep425tags as w; print('-'.join(w.get_supported()[0][:-1]))")-linux_x86_64.whl
 RUN pip install mxnet-cu100==${MXNET_VERSION}
 
 # Install Open MPI

build-docker-images.sh

@@ -6,23 +6,33 @@ set -x
 function build_one()
 {
     py=$1
-    tag=horovod-build-py${py}:$(date +%Y%m%d-%H%M%S)
-    docker build -t ${tag} --build-arg python=${py} --no-cache .
+    device=$2
+
+    tensorflow_pkg='tensorflow'
+    if [[ ${device} == 'gpu' ]]; then
+        tensorflow_pkg='tensorflow-gpu'
+    fi
+
+    tag=horovod-build-py${py}-${device}:$(date +%Y%m%d-%H%M%S)
+    docker build -f Dockerfile.${device} -t ${tag} --build-arg python=${py} --no-cache .
     horovod_version=$(docker run ${tag} pip freeze | grep ^horovod= | awk -F== '{print $2}')
-    tensorflow_version=$(docker run ${tag} pip freeze | grep ^tensorflow-gpu= | awk -F== '{print $2}')
+    tensorflow_version=$(docker run ${tag} pip freeze | grep ^${tensorflow_pkg}= | awk -F== '{print $2}')
     pytorch_version=$(docker run ${tag} pip freeze | grep ^torch= | awk -F== '{print $2}' | awk -F+ '{print $1}')
     mxnet_version=$(docker run ${tag} pip freeze | grep ^mxnet | awk -F== '{print $2}')
-    final_tag=horovod/horovod:${horovod_version}-tf${tensorflow_version}-torch${pytorch_version}-mxnet${mxnet_version}-py${py}
+    final_tag=horovod/horovod:${horovod_version}-tf${tensorflow_version}-torch${pytorch_version}-mxnet${mxnet_version}-py${py}-${device}
     docker tag ${tag} ${final_tag}
     docker rmi ${tag}
 }
 
-# clear upstream image, ok to fail if image does not exist
-docker rmi $(cat Dockerfile | grep FROM | awk '{print $2}') || true
+# clear upstream images, ok to fail if images do not exist
+docker rmi $(cat Dockerfile.gpu | grep FROM | awk '{print $2}') || true
+docker rmi $(cat Dockerfile.cpu | grep FROM | awk '{print $2}') || true
 
-# build for py2 and py3
-build_one 2.7
-build_one 3.6
+# build for py2 and py3, cpu and gpu
+build_one 2.7 gpu
+build_one 3.6 gpu
+build_one 2.7 cpu
+build_one 3.6 cpu
 
 # print recent images
 docker images horovod/horovod

docs/contributors.rst

@@ -244,13 +244,15 @@ From a clean copy of the ``horovod`` repository on a Linux machine:
 If you have trouble connecting to external URLs, try changing ``docker build ...`` to
 ``docker build --network host ...`` in ``build-docker-images.sh``.
 
-Upload artifacts for both Python 2.7 and Python 3.6:
+Upload artifacts for Python 2.7 and Python 3.6, CPU and GPU:
 
 .. code-block:: bash
 
     $ docker login
-    $ docker push horovod/horovod:0.18.1-tf1.14.0-torch1.2.0-mxnet1.5.0-py2.7
-    $ docker push horovod/horovod:0.18.1-tf1.14.0-torch1.2.0-mxnet1.5.0-py3.6
+    $ docker push horovod/horovod:0.18.1-tf1.14.0-torch1.2.0-mxnet1.5.0-py2.7-gpu
+    $ docker push horovod/horovod:0.18.1-tf1.14.0-torch1.2.0-mxnet1.5.0-py3.6-gpu
+    $ docker push horovod/horovod:0.18.1-tf1.14.0-torch1.2.0-mxnet1.5.0-py2.7-cpu
+    $ docker push horovod/horovod:0.18.1-tf1.14.0-torch1.2.0-mxnet1.5.0-py3.6-cpu
 
 Check the horovod `Docker Hub project <https://cloud.docker.com/u/horovod/repository/docker/horovod/horovod>`_
 to verify that the image artifacts were successfully uploaded.