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

Author: peterschmidt85

.github/workflows/docker.yml

@@ -51,8 +51,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        # flavor: ["base", "devel", "devel-efa"]
-        flavor: ["base-efa"]
+        flavor: ["base", "devel", "devel-efa"]
     steps:
       - name: Checkout repository
         uses: actions/checkout@v4
@@ -71,8 +70,6 @@ jobs:
             FILE="base/Dockerfile"
           elif [ "${{ matrix.flavor }}" = "devel" ]; then
             FILE="base/Dockerfile"
-          elif [ "${{ matrix.flavor }}" = "base-efa" ]; then
-            FILE="base/base-efa.Dockerfile"
           else
             FILE="base/efa.Dockerfile"
           fi

docker/base/Dockerfile

@@ -55,7 +55,7 @@ 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_TESTS_HOME}/build ${NCCL_TESTS_HOME}/build
 
 ARG FLAVOR
 

docker/base/base-efa.Dockerfile

@@ -2,14 +2,15 @@
 
 INCLUDE+ base/Dockerfile.common
 
+ENV NCCL_HOME=/opt/nccl
 ENV PREFIX=/usr/local
-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 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 \
@@ -25,47 +26,48 @@ RUN cuda_version=$(echo ${CUDA_VERSION} | awk -F . '{ print $1"-"$2 }') \
 
 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
+    && ./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=${PREFIX} \
+    && 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} \
         --disable-tests \
         --prefix=${PREFIX} \
-    && make -j$(numproc) \
-    && make install
+    && make -j$(nproc) \
+    && make install \
+    && rm -rf /tmp/aws-ofi-nccl
 
 # NCCL Tests
 
-RUN cd $HOME \
+RUN cd $NCCL_HOME \
     && 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} \
+        CUDA_HOME=${CUDA_HOME} \
         NCCL_HOME=${PREFIX}