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

Author: peterschmidt85

.github/workflows/docker.yml

@@ -69,15 +69,12 @@ jobs:
         run: |
           if [ "${{ matrix.flavor }}" = "base" ]; then
             TAG_SUFFIX=""
-            FLAVOR="base"
             FILE="base/Dockerfile"
           elif [ "${{ matrix.flavor }}" = "efa" ]; then
             TAG_SUFFIX="-efa"
-            FLAVOR="devel"
             FILE="base/efa/Dockerfile"
           else
             TAG_SUFFIX="-devel"
-            FLAVOR="devel"
             FILE="base/devel/Dockerfile"
           fi
           docker buildx build \

docker/base/Dockerfile

@@ -1,5 +1,4 @@
-ARG FLAVOR
-FROM nvidia/cuda:12.1.1-${FLAVOR}-ubuntu20.04
+FROM nvidia/cuda:12.1.1-base-ubuntu20.04
 
 ARG PYTHON
 ARG _UV_HOME="/opt/uv"

docker/base/devel/Dockerfile

@@ -9,20 +9,30 @@ ARG NCCL_VERSION=2.26.2-1
 ENV NCCL_HOME=/usr/local
 ENV CUDA_PATH=/usr/local/cuda
 ENV OPEN_MPI_PATH=/usr/lib/x86_64-linux-gnu/openmpi
+ENV NCCL_TESTS_PATH=/opt/nccl-tests
+ENV PATH="${LIBFABRIC_PATH}/bin:${OPEN_MPI_PATH}/bin:${NCCL_TESTS_PATH}:${PATH}"
 ENV LD_LIBRARY_PATH="${OPEN_MPI_PATH}/lib:${NCCL_HOME}/lib:${LD_LIBRARY_PATH}"
-ENV PATH="${OPEN_MPI_PATH}/bin:${HOME}/nccl-tests/build:${PATH}"
 
-RUN apt-get install -y --no-install-recommends \
-        libopenmpi-dev \
-    && cd $HOME \
+RUN cuda_version=$(echo ${CUDA_VERSION} | awk -F . '{ print $1"-"$2 }') \
+    && apt-get install -y --no-install-recommends \
+        cuda-libraries-dev-${cuda_version} \
+        cuda-nvcc-${cuda_version} \
+        libhwloc-dev \
+        autoconf \
+        automake \
+        libtool \
+        libopenmpi-dev
+
+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 \
-    && git clone https://github.com/NVIDIA/nccl-tests \
-    && cd nccl-tests \
+    && git clone https://github.com/NVIDIA/nccl-tests ${NCCL_TESTS_PATH} \
+    && cd ${NCCL_TESTS_PATH} \
     && make -j$(nproc) \
         MPI=1 \
         MPI_HOME=${OPEN_MPI_PATH} \
         CUDA_HOME=${CUDA_PATH} \
-        NCCL_HOME=${NCCL_HOME}
+        NCCL_HOME=${NCCL_HOME} \
+    && echo "${NCCL_HOME}/lib" >> /etc/ld.so.conf.d/nccl.conf \
+    && ldconfig

docker/base/efa/Dockerfile

@@ -6,13 +6,16 @@ ENV NCCL_HOME=/usr/local
 ENV CUDA_PATH=/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:${HOME}/nccl-tests/build:${PATH}"
-ENV LD_LIBRARY_PATH="${OPEN_MPI_PATH}/lib:${LD_LIBRARY_PATH}"
+ENV NCCL_TESTS_PATH=/opt/nccl-tests
+ENV PATH="${LIBFABRIC_PATH}/bin:${OPEN_MPI_PATH}/bin:${NCCL_TESTS_PATH}:${PATH}"
+ENV LD_LIBRARY_PATH="${OPEN_MPI_PATH}/lib:${NCCL_HOME}/lib:${LD_LIBRARY_PATH}"
 
 # Prerequisites
 
-RUN apt-get update \
+RUN cuda_version=$(echo ${CUDA_VERSION} | awk -F . '{ print $1"-"$2 }') \
     && apt-get install -y --no-install-recommends \
+        cuda-libraries-dev-${cuda_version} \
+        cuda-nvcc-${cuda_version} \
         libhwloc-dev \
         autoconf \
         automake \
@@ -58,11 +61,13 @@ RUN cd $HOME \
 
 # NCCL Tests
 
-RUN cd $HOME \
-    && git clone https://github.com/NVIDIA/nccl-tests \
-    && cd nccl-tests \
+RUN git clone https://github.com/NVIDIA/nccl-tests ${NCCL_TESTS_PATH} \
+    && cd ${NCCL_TESTS_PATH} \
     && make -j$(numproc) \
         MPI=1 \
         MPI_HOME=${OPEN_MPI_PATH} \
         CUDA_HOME=${CUDA_PATH} \
-        NCCL_HOME=${NCCL_HOME}
+        NCCL_HOME=${NCCL_HOME} \
+    && echo "${NCCL_HOME}/lib" >> /etc/ld.so.conf.d/nccl.conf \
+    && ldconfig
+