[UX] Pre-build a EFA version of the default Docker image #2793

Author: peterschmidt85

docker/base/efa/Dockerfile

@@ -10,7 +10,10 @@ ENV NCCL_TESTS_HOME=/opt/nccl-tests
 ENV PATH="${LIBFABRIC_PATH}/bin:${MPI_HOME}/bin:${NCCL_TESTS_HOME}/build:${PATH}"
 ENV LD_LIBRARY_PATH="${MPI_HOME}/lib:${NCCL_HOME}/lib:${LD_LIBRARY_PATH}"
 
-# Prerequisites
+ARG EFA_VERSION=1.38.1
+ARG NCCL_VERSION=2.26.2-1
+ARG OFI_VERSION=1.14.0
+ARG FLAVOR
 
 RUN cuda_version=$(echo ${CUDA_VERSION} | awk -F . '{ print $1"-"$2 }') \
     && apt-get install -y --no-install-recommends \
@@ -19,32 +22,17 @@ RUN cuda_version=$(echo ${CUDA_VERSION} | awk -F . '{ print $1"-"$2 }') \
         libhwloc-dev \
         autoconf \
         automake \
-        libtool
-
-# EFA
-
-ARG EFA_VERSION=1.38.1
-
-RUN cd $HOME \
+        libtool \
+    && cd $HOME \
     && curl -O https://s3-us-west-2.amazonaws.com/aws-efa-installer/aws-efa-installer-${EFA_VERSION}.tar.gz \
     && tar -xf aws-efa-installer-${EFA_VERSION}.tar.gz \
     && cd aws-efa-installer \
-    && ./efa_installer.sh -y --skip-kmod -g
-
-# NCCL
-
-ARG NCCL_VERSION=2.26.2-1
-
-RUN cd $HOME \
+    && ./efa_installer.sh -y --skip-kmod -g \
+    && cd $HOME \
     && git clone https://github.com/NVIDIA/nccl.git -b v${NCCL_VERSION} \
     && cd nccl \
-    && make -j$(nproc) src.build BUILDDIR=${NCCL_HOME}
-
-# AWS OFI NCCL
-
-ARG OFI_VERSION=1.14.0
-
-RUN cd $HOME \
+    && make -j$(nproc) src.build BUILDDIR=${NCCL_HOME} \
+    && cd $HOME \
     && git clone https://github.com/aws/aws-ofi-nccl.git -b v${OFI_VERSION} \
     && cd aws-ofi-nccl \
     && ./autogen.sh \
@@ -56,33 +44,24 @@ RUN cd $HOME \
         --disable-tests \
         --prefix=${NCCL_HOME} \
     && make -j$(numproc) \
-    && make install
-
-# NCCL Tests
-
-RUN git clone https://github.com/NVIDIA/nccl-tests ${NCCL_TESTS_HOME} \
+    && make install \
+    && git clone https://github.com/NVIDIA/nccl-tests ${NCCL_TESTS_HOME} \
     && cd ${NCCL_TESTS_HOME} \
     && make -j$(numproc) \
         MPI=1 \
         MPI_HOME=${MPI_HOME} \
         CUDA_HOME=${CUDA_HOME} \
         NCCL_HOME=${NCCL_HOME} \
     && echo "${NCCL_HOME}/lib" >> /etc/ld.so.conf.d/nccl.conf \
-    && ldconfig
-    
-ARG FLAVOR
-ENV FLAVOR=${FLAVOR}    
-
-# If FLAVOR is base, uninstall development packages to reduce image size
-RUN if [ "$FLAVOR" = "base" ]; then \
-    cuda_version=$(echo ${CUDA_VERSION} | awk -F . '{ print $1"-"$2 }') \
-    && apt-get remove -y \
-        cuda-nvcc-${cuda_version} \
-        libhwloc-dev \
-        autoconf \
-        automake \
-        libtool \
-    && apt-get autoremove -y \
-    && apt-get clean \
-    && rm -rf /var/lib/apt/lists/*; \
-fi
+    && ldconfig \
+    && if [ "$FLAVOR" = "base" ]; then \
+        apt-get remove -y \
+            cuda-nvcc-${cuda_version} \
+            libhwloc-dev \
+            autoconf \
+            automake \
+            libtool \
+        && apt-get autoremove -y \
+        && apt-get clean \
+        && rm -rf /var/lib/apt/lists/*; \
+    fi

docker/base/nvidia/Dockerfile

@@ -5,6 +5,7 @@ INCLUDE+ base/Dockerfile.common
 # NCCL & NCCL tests
 
 ARG NCCL_VERSION=2.26.2-1
+ARG FLAVOR
 
 ENV NCCL_HOME=/usr/local
 ENV CUDA_HOME=/usr/local/cuda
@@ -21,9 +22,8 @@ RUN cuda_version=$(echo ${CUDA_VERSION} | awk -F . '{ print $1"-"$2 }') \
         autoconf \
         automake \
         libtool \
-        libopenmpi-dev
-
-RUN cd $HOME \
+        libopenmpi-dev \
+    && cd $HOME \
     && git clone https://github.com/NVIDIA/nccl.git -b v${NCCL_VERSION} \
     && cd nccl \
     && make -j$(nproc) src.build BUILDDIR=${NCCL_HOME} \
@@ -35,22 +35,15 @@ RUN cd $HOME \
         CUDA_HOME=${CUDA_HOME} \
         NCCL_HOME=${NCCL_HOME} \
     && echo "${NCCL_HOME}/lib" >> /etc/ld.so.conf.d/nccl.conf \
-    && ldconfig
-
-ARG FLAVOR
-ENV FLAVOR=${FLAVOR}    
-
-# If FLAVOR is base, uninstall development packages to reduce image size
-RUN if [ "$FLAVOR" = "base" ]; then \
-    cuda_version=$(echo ${CUDA_VERSION} | awk -F . '{ print $1"-"$2 }') \
-    && apt-get remove -y \
-        cuda-nvcc-${cuda_version} \
-        libhwloc-dev \
-        autoconf \
-        automake \
-        libtool \
-        libopenmpi-dev \
-    && apt-get autoremove -y \
-    && apt-get clean \
-    && rm -rf /var/lib/apt/lists/*; \
-fi
+    && ldconfig \
+    && if [ "$FLAVOR" = "base" ]; then \
+        apt-get remove -y \
+            cuda-nvcc-${cuda_version} \
+            libhwloc-dev \
+            autoconf \
+            automake \
+            libtool \
+        && apt-get autoremove -y \
+        && apt-get clean \
+        && rm -rf /var/lib/apt/lists/*; \
+    fi