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

Author: peterschmidt85

.github/workflows/docker.yml

@@ -52,7 +52,8 @@ jobs:
     strategy:
       matrix:
         python: ["3.9", "3.10", "3.11", "3.12", "3.13"]
-        flavor: ["base", "devel", "efa"]
+        version: ["nvidia", "efa"]
+        flavor: ["base", "devel"]
     steps:
       - name: Checkout repository
         uses: actions/checkout@v4
@@ -67,24 +68,14 @@ jobs:
         uses: docker/setup-qemu-action@v3
       - name: Build and upload to DockerHub
         run: |
-          if [ "${{ matrix.flavor }}" = "base" ]; then
-            TAG_SUFFIX=""
-            FILE="base/Dockerfile"
-          elif [ "${{ matrix.flavor }}" = "efa" ]; then
-            TAG_SUFFIX="-efa"
-            FILE="base/efa/Dockerfile"
-          else
-            TAG_SUFFIX="-devel"
-            FILE="base/devel/Dockerfile"
-          fi
           docker buildx build \
             --platform linux/amd64 \
-            --tag dstackai/${{ env.BUILD_DOCKER_REPO }}:py${{ matrix.python }}-${{ inputs.image_version }}-cuda-12.1${TAG_SUFFIX} \
-            --build-arg FLAVOR=$FLAVOR \
+            --tag dstackai/${{ env.BUILD_DOCKER_REPO }}:py${{ matrix.python }}-${{ inputs.image_version }}-${{ matrix.version }}-${{ matrix.flavor }} \
+            --build-arg FLAVOR=${{ matrix.flavor }} \
             --build-arg PYTHON=${{ matrix.python }} \
             --provenance=false \
             --push \
-            -f $FILE .
+            -f default/${{ matrix.version }}/Dockerfile .
 
   build-aws-images:
     needs: build-docker

docker/base/Dockerfile docker/default/Dockerfile.commondocker/base/Dockerfile renamed to docker/default/Dockerfile.common

@@ -3,12 +3,12 @@
 INCLUDE+ base/Dockerfile
 
 ENV NCCL_HOME=/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 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 MPI_HOME=/opt/amazon/openmpi
+ENV NCCL_TESTS_HOME=/opt/nccl-tests
+ENV PATH="${LIBFABRIC_PATH}/bin:${MPI_HOME}/bin:${NCCL_TESTS_HOME}/build:${PATH}"
+ENV LD_LIBRARY_PATH="${MPI_HOME}/lib:${NCCL_HOME}/lib:${LD_LIBRARY_PATH}"
 
 # Prerequisites
 
@@ -49,10 +49,9 @@ RUN cd $HOME \
     && 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-mpi=${MPI_HOME} \
         --with-nccl=${NCCL_HOME} \
         --disable-tests \
         --prefix=${NCCL_HOME} \
@@ -61,13 +60,29 @@ RUN cd $HOME \
 
 # NCCL Tests
 
-RUN git clone https://github.com/NVIDIA/nccl-tests ${NCCL_TESTS_PATH} \
-    && cd ${NCCL_TESTS_PATH} \
+RUN git clone https://github.com/NVIDIA/nccl-tests ${NCCL_TESTS_HOME} \
+    && cd ${NCCL_TESTS_HOME} \
     && make -j$(numproc) \
         MPI=1 \
-        MPI_HOME=${OPEN_MPI_PATH} \
-        CUDA_HOME=${CUDA_PATH} \
+        MPI_HOME=${MPI_HOME} \
+        CUDA_HOME=${CUDA_HOME} \
         NCCL_HOME=${NCCL_HOME} \
     && echo "${NCCL_HOME}/lib" >> /etc/ld.so.conf.d/nccl.conf \
     && ldconfig
 
+ARG FLAVOR
+ENV FLAVOR=${FLAVOR}    
+
+# If FLAVOR is base, uninstall development packages to reduce image size
+RUN if [ "$FLAVOR" = "base" ]; then \
+    cuda_version=$(echo ${CUDA_VERSION} | awk -F . '{ print $1"-"$2 }') \
+    && apt-get remove -y \
+        cuda-nvcc-${cuda_version} \
+        libhwloc-dev \
+        autoconf \
+        automake \
+        libtool \
+    && apt-get autoremove -y \
+    && apt-get clean \
+    && rm -rf /var/lib/apt/lists/*; \
+fi

docker/base/README.md docker/default/README.mddocker/base/README.md renamed to docker/default/README.md

@@ -1,16 +1,16 @@
 # syntax = edrevo/dockerfile-plus
 
-INCLUDE+ base/Dockerfile
+INCLUDE+ default/Dockerfile.common
 
 # NCCL & NCCL tests
 
 ARG NCCL_VERSION=2.26.2-1
 
 ENV NCCL_HOME=/usr/local
-ENV CUDA_PATH=/usr/local/cuda
+ENV CUDA_HOME=/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 NCCL_TESTS_HOME=/opt/nccl-tests
+ENV PATH="${NCCL_TESTS_HOME}/build:${PATH}"
 ENV LD_LIBRARY_PATH="${OPEN_MPI_PATH}/lib:${NCCL_HOME}/lib:${LD_LIBRARY_PATH}"
 
 RUN cuda_version=$(echo ${CUDA_VERSION} | awk -F . '{ print $1"-"$2 }') \
@@ -27,12 +27,30 @@ 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} \
-    && git clone https://github.com/NVIDIA/nccl-tests ${NCCL_TESTS_PATH} \
-    && cd ${NCCL_TESTS_PATH} \
+    && 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} \
     && echo "${NCCL_HOME}/lib" >> /etc/ld.so.conf.d/nccl.conf \
     && ldconfig
+
+ARG FLAVOR
+ENV FLAVOR=${FLAVOR}    
+
+# If FLAVOR is base, uninstall development packages to reduce image size
+RUN if [ "$FLAVOR" = "base" ]; then \
+    cuda_version=$(echo ${CUDA_VERSION} | awk -F . '{ print $1"-"$2 }') \
+    && apt-get remove -y \
+        cuda-nvcc-${cuda_version} \
+        libhwloc-dev \
+        autoconf \
+        automake \
+        libtool \
+        libopenmpi-dev \
+    && apt-get autoremove -y \
+    && apt-get clean \
+    && rm -rf /var/lib/apt/lists/*; \
+fi