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

Author: peterschmidt85

docker/base/efa.Dockerfile

@@ -2,57 +2,70 @@
 
 INCLUDE+ base/Dockerfile.common
 
-ENV NCCL_HOME=/usr/local
-ENV CUDA_HOME=/usr/local/cuda
+ENV PREFIX=/usr/local
+ENV CUDA_PATH=/usr/local/cuda
 ENV LIBFABRIC_PATH=/opt/amazon/efa
 ENV OPEN_MPI_PATH=/opt/amazon/openmpi
-ENV NCCL_TESTS_HOME=/opt/nccl-tests
 ENV PATH="${LIBFABRIC_PATH}/bin:${OPEN_MPI_PATH}/bin:${PATH}"
+ENV LD_LIBRARY_PATH="${OPEN_MPI_PATH}/lib:${LD_LIBRARY_PATH}"
 
-ARG EFA_VERSION=1.38.1
-ARG NCCL_VERSION=2.26.2-1
-ARG OFI_VERSION=1.14.0
-ARG FLAVOR
+# prerequisites
 
-RUN apt-get update \
-    && cuda_version=$(echo ${CUDA_VERSION} | awk -F . '{ print $1"-"$2 }') \
+RUN cuda_version=$(echo ${CUDA_VERSION} | awk -F . '{ print $1"-"$2 }') \
+    && apt-get update \
     && apt-get install -y --no-install-recommends \
         cuda-libraries-dev-${cuda_version} \
         cuda-nvcc-${cuda_version} \
         libhwloc-dev \
         autoconf \
         automake \
-        libtool \
-    && cd $HOME \
+        libtool
+
+# EFA
+
+ARG EFA_VERSION=1.38.1
+
+RUN 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 \
-    && cd $HOME \
+    && ./efa_installer.sh -y --skip-kmod -g
+
+# NCCL
+
+ARG NCCL_VERSION=2.26.2-1
+
+RUN cd $HOME \
     && git clone https://github.com/NVIDIA/nccl.git -b v${NCCL_VERSION} \
     && cd nccl \
-    && make -j$(nproc) src.build BUILDDIR=${NCCL_HOME} \
-    && cd $HOME \
+    && make -j$(nproc) src.build BUILDDIR=${PREFIX}
+
+# AWS OFI NCCL
+
+ARG OFI_VERSION=1.14.0
+
+RUN cd $HOME \
     && git clone https://github.com/aws/aws-ofi-nccl.git -b v${OFI_VERSION} \
     && cd aws-ofi-nccl \
     && ./autogen.sh \
     && ./configure \
-        --with-cuda=${CUDA_HOME} \
+        --with-cuda=${CUDA_PATH} \
         --with-libfabric=${LIBFABRIC_PATH} \
         --with-mpi=${OPEN_MPI_PATH} \
-        --with-nccl=${NCCL_HOME} \
+        --with-cuda=${CUDA_PATH} \
+        --with-nccl=${PREFIX} \
         --disable-tests \
-        --prefix=${NCCL_HOME} \
+        --prefix=${PREFIX} \
     && make -j$(numproc) \
-    && make install \
-    && git clone https://github.com/NVIDIA/nccl-tests ${HOME}/nccl-tests \
-    && cd ${HOME}/nccl-tests \
+    && make install
+
+# NCCL Tests
+
+RUN cd $HOME \
+    && git clone https://github.com/NVIDIA/nccl-tests \
+    && cd nccl-tests \
     && make -j$(numproc) \
         MPI=1 \
         MPI_HOME=${OPEN_MPI_PATH} \
-        CUDA_HOME=${CUDA_HOME} \
-        NCCL_HOME=${NCCL_HOME} \
-    && ln -s ${HOME}/nccl-tests/build ${NCCL_TESTS_HOME} \
-    && echo "${OPEN_MPI_PATH}/lib" >> /etc/ld.so.conf.d/openmpi.conf \
-    && echo "${NCCL_HOME}/lib" >> /etc/ld.so.conf.d/nccl.conf \
-    && ldconfig
+        CUDA_HOME=${CUDA_PATH} \
+        NCCL_HOME=${PREFIX}