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

Author: peterschmidt85

docker/base/Dockerfile

@@ -3,14 +3,12 @@
 # Build stage
 FROM nvidia/cuda:12.1.1-base-ubuntu20.04 AS builder
 
-ARG NCCL_VERSION=2.26.2-1
-
 ENV NCCL_HOME=/opt/nccl
 ENV CUDA_HOME=/usr/local/cuda
 ENV OPEN_MPI_PATH=/usr/lib/x86_64-linux-gnu/openmpi
-ENV NCCL_TESTS_HOME=/opt/nccl-tests
 
-# Install build dependencies
+# Prerequisites
+
 RUN export DEBIAN_FRONTEND=noninteractive \
     && apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/3bf863cc.pub \
     && apt-get update --fix-missing \
@@ -32,36 +30,42 @@ RUN export DEBIAN_FRONTEND=noninteractive \
         python3 \
         build-essential
 
-# Build NCCL
+# NCCL
+
+ARG NCCL_VERSION=2.26.2-1
+
 RUN cd /tmp \
     && git clone https://github.com/NVIDIA/nccl.git -b v${NCCL_VERSION} \
     && cd nccl \
     && make -j$(nproc) src.build BUILDDIR=${NCCL_HOME}
 
-# Build NCCL tests
-RUN git clone https://github.com/NVIDIA/nccl-tests ${NCCL_TESTS_HOME} \
-    && cd ${NCCL_TESTS_HOME} \
+# NCCL tests
+
+RUN cd /opt \
+    && git clone https://github.com/NVIDIA/nccl-tests \
+    && cd nccl-tests \
     && make -j$(nproc) \
         MPI=1 \
         MPI_HOME=${OPEN_MPI_PATH} \
         CUDA_HOME=${CUDA_HOME} \
         NCCL_HOME=${NCCL_HOME}
 
 # Final stage
+
 INCLUDE+ base/Dockerfile.common
 
 ENV NCCL_HOME=/opt/nccl
-ENV NCCL_TESTS_HOME=/opt/nccl-tests
 
-COPY --from=builder ${NCCL_HOME}/lib ${NCCL_HOME}/lib
-COPY --from=builder ${NCCL_HOME}/include ${NCCL_HOME}/include
-COPY --from=builder ${NCCL_TESTS_HOME}/build ${NCCL_TESTS_HOME}
+COPY --from=builder ${NCCL_HOME} ${NCCL_HOME}
+COPY --from=builder /opt/nccl-tests/build /opt/nccl-tests/build
 
 ARG FLAVOR
 
-# Configure library paths
+# MPI, NVCC, and /etc/ld.so.conf.d
+
 RUN apt-get update \
-    && apt-get install -y --no-install-recommends openmpi-bin \
+    && apt-get install -y --no-install-recommends \
+        openmpi-bin \
     && if [ "$FLAVOR" = "devel" ]; then \
         cuda_version=$(echo ${CUDA_VERSION} | awk -F . '{ print $1"-"$2 }') \
         && apt-get install -y --no-install-recommends \

docker/base/efa.Dockerfile

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