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

Author: peterschmidt85

docker/base/base-efa.Dockerfile

@@ -3,13 +3,15 @@
 # Build stage
 FROM nvidia/cuda:12.1.1-base-ubuntu20.04 AS builder
 
+ARG NCCL_VERSION=2.26.2-1
+ARG EFA_VERSION=1.38.1
+ARG OFI_VERSION=1.14.0
+
 ENV NCCL_HOME=/opt/nccl
-ENV CUDA_PATH=/usr/local/cuda
+ENV CUDA_HOME=/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}"
 
 # Install build dependencies
 RUN export DEBIAN_FRONTEND=noninteractive \
@@ -33,40 +35,18 @@ RUN export DEBIAN_FRONTEND=noninteractive \
         python3 \
         build-essential
 
-# EFA
-
-ARG EFA_VERSION=1.38.1
-
-RUN cd $HOME \
+RUN cd /tmp \
     && 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
-
-# Build NCCL tests
-RUN git clone https://github.com/NVIDIA/nccl-tests ${NCCL_TESTS_HOME} \
-    && cd ${NCCL_TESTS_HOME} \
-    && make -j$(nproc) \
-        MPI=1 \
-        MPI_HOME=${OPEN_MPI_PATH} \
-        CUDA_HOME=${CUDA_HOME} \
-        NCCL_HOME=${NCCL_HOME}
-
-
-# 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-nccl=${NCCL_HOME} \
@@ -75,15 +55,19 @@ RUN cd $HOME \
     && make -j$(numproc) \
     && make install
 
-# NCCL Tests
+# Build NCCL
+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}
 
-RUN cd $HOME \
-    && git clone https://github.com/NVIDIA/nccl-tests \
-    && cd nccl-tests \
-    && make -j$(numproc) \
+# Build NCCL tests
+RUN git clone https://github.com/NVIDIA/nccl-tests ${NCCL_TESTS_HOME} \
+    && cd ${NCCL_TESTS_HOME} \
+    && make -j$(nproc) \
         MPI=1 \
         MPI_HOME=${OPEN_MPI_PATH} \
-        CUDA_HOME=${CUDA_PATH} \
+        CUDA_HOME=${CUDA_HOME} \
         NCCL_HOME=${NCCL_HOME}
 
 # Final stage