From 784fc32bfa27526eb83b5561225933949abebac2 Mon Sep 17 00:00:00 2001
From: Yu Yang <yuyang18@baidu.com>
Date: Sat, 21 Oct 2017 06:00:34 +0800
Subject: [PATCH] Add nccl to docker image

---
 Dockerfile                     | 2 +-
 paddle/scripts/docker/build.sh | 7 +++++++
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/Dockerfile b/Dockerfile
index 136db772c..150344a81 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -22,7 +22,7 @@ COPY ./paddle/scripts/docker/root/ /root/
 
 RUN apt-get update && \
     apt-get install -y \
-    git python-pip python-dev openssh-server bison  \
+    git python-pip python-dev openssh-server bison libnccl-dev \
     wget unzip unrar tar xz-utils bzip2 gzip coreutils ntp \
     curl sed grep graphviz libjpeg-dev zlib1g-dev  \
     python-matplotlib gcc-4.8 g++-4.8 \
diff --git a/paddle/scripts/docker/build.sh b/paddle/scripts/docker/build.sh
index 2ac455d77..a08716c5a 100644
--- a/paddle/scripts/docker/build.sh
+++ b/paddle/scripts/docker/build.sh
@@ -141,10 +141,17 @@ RUN sed -i '${APT_MIRROR}' /etc/apt/sources.list
 EOF
 fi
 
+if [[ ${WITH_GPU} == "ON"  ]]; then
+  NCCL_DEPS="apt-get install -y libnccl-dev &&"
+else
+  NCCL_DEPS="" 
+fi
+
 cat >> /paddle/build/Dockerfile <<EOF
 ADD python/dist/*.whl /
 # run paddle version to install python packages first
 RUN apt-get update &&\
+    ${NCCL_DEPS}\
     apt-get install -y wget python-pip && pip install -U pip && \
     pip install /*.whl; apt-get install -f -y && \
     apt-get clean -y && \
-- 
GitLab