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

Author: peterschmidt85

docker/base/base-efa.Dockerfile

@@ -8,6 +8,7 @@ ARG EFA_VERSION=1.38.1
 ARG OFI_VERSION=1.14.0
 
 ENV NCCL_HOME=/opt/nccl
+ENV OFI_NCCL_HOME=/opt/amazon/ofi-nccl
 ENV CUDA_HOME=/usr/local/cuda
 ENV LIBFABRIC_PATH=/opt/amazon/efa
 ENV OPEN_MPI_PATH=/opt/amazon/openmpi
@@ -76,19 +77,24 @@ ENV NCCL_HOME=/opt/nccl
 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}"
-# TODO: Unsure if this is required, updating ` /etc/ld.so.conf.d` should be enough
-ENV LD_LIBRARY_PATH="${OPEN_MPI_PATH}/lib:${LD_LIBRARY_PATH}"
 
 COPY --from=builder ${NCCL_HOME} ${NCCL_HOME}
-COPY --from=builder ${LIBFABRIC_PATH} ${LIBFABRIC_PATH}
-COPY --from=builder ${OPEN_MPI_PATH} ${OPEN_MPI_PATH}
+COPY --from=builder ${OFI_NCCL_HOME} ${OFI_NCCL_HOME}
+COPY --from=builder /etc/ld.so.conf.d/100_ofinccl.conf /etc/ld.so.conf.d/100_ofinccl.conf
 COPY --from=builder ${NCCL_TESTS_HOME}/build ${NCCL_TESTS_HOME}
-COPY --from=builder /etc/ld.so.conf.d/000_efa.conf /etc/ld.so.conf.d/000_efa.conf
-COPY --from=builder /etc/profile.d/zippy_efa.sh /etc/profile.d/zippy_efa.sh
 
-RUN echo "${NCCL_HOME}/lib" >> /etc/ld.so.conf.d/nccl.conf \
+RUN apt-get update \
+    && apt-get install -y --no-install-recommends \
+        libevent-dev \
+        libhwloc-dev \
+    && 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
+    && rm -rf /tmp/aws-efa-installer /var/lib/apt/lists/* \
+    && echo "${NCCL_HOME}/lib" >> /etc/ld.so.conf.d/nccl.conf \
     && echo "${OPEN_MPI_PATH}/lib" >> /etc/ld.so.conf.d/openmpi.conf \
     && echo "${LIBFABRIC_PATH}/lib" >> /etc/ld.so.conf.d/efa.conf \
     && ldconfig