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

Author: peterschmidt85

.github/workflows/docker.yml

@@ -78,7 +78,7 @@ jobs:
           fi
           docker buildx build \
             --platform linux/amd64 \
-            --tag dstackai/${{ env.BUILD_DOCKER_REPO }}:${{ matrix.flavor }}-${{ inputs.image_version }} \
+            --tag dstackai/${{ env.BUILD_DOCKER_REPO }}:${{ inputs.image_version }}-${{ matrix.flavor }} \
             --build-arg FLAVOR=${{ matrix.flavor }} \
             --provenance=false \
             --push \

docker/base/Dockerfile

@@ -68,5 +68,6 @@ RUN apt-get install -y --no-install-recommends openmpi-bin \
             cuda-nvcc-${cuda_version} \
             libhwloc-dev; \
     fi \
+    && rm -rf /var/lib/apt/lists/* \
     && echo "${NCCL_HOME}/lib" >> /etc/ld.so.conf.d/nccl.conf \
     && ldconfig

docker/base/Dockerfile.common

@@ -8,17 +8,21 @@ ENV LANG=C.UTF-8 LC_ALL=C.UTF-8
 
 ENV PATH="${UV_INSTALL_DIR}:${PATH}"
 
-RUN export DEBIAN_FRONTEND=noninteractive && \
-    apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/3bf863cc.pub && \
-    apt-get update --fix-missing && \
-    apt-get upgrade -y && \
-    ln -fs /usr/share/zoneinfo/America/New_York /etc/localtime && \
-    apt-get install -y tzdata && \
-    dpkg-reconfigure --frontend noninteractive tzdata && \
-    apt-get install -y bzip2 ca-certificates curl build-essential git libglib2.0-0 libsm6 libxext6 libxrender1 mercurial openssh-server subversion wget \
-    libibverbs1 ibverbs-providers ibverbs-utils libibverbs-dev infiniband-diags && \
-    sed -i "s/.*PasswordAuthentication.*/PasswordAuthentication no/g" /etc/ssh/sshd_config && mkdir /run/sshd && \
-    mkdir ~/.ssh && chmod 700 ~/.ssh && touch ~/.ssh/authorized_keys && chmod 600 ~/.ssh/authorized_keys && rm /etc/ssh/ssh_host_*
+RUN export DEBIAN_FRONTEND=noninteractive \
+    && apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/3bf863cc.pub \
+    && apt-get update --fix-missing \
+    && apt-get upgrade -y \
+    && ln -fs /usr/share/zoneinfo/America/New_York /etc/localtime \
+    && apt-get install -y tzdata \
+    && dpkg-reconfigure --frontend noninteractive tzdata \
+    && apt-get install -y bzip2 ca-certificates curl build-essential git libglib2.0-0 libsm6 libxext6 libxrender1 mercurial openssh-server subversion wget \
+        libibverbs1 ibverbs-providers ibverbs-utils libibverbs-dev infiniband-diags \
+    && rm -rf /var/lib/apt/lists/* \
+    && sed -i "s/.*PasswordAuthentication.*/PasswordAuthentication no/g" /etc/ssh/sshd_config \
+    && mkdir /run/sshd \
+    && mkdir ~/.ssh && chmod 700 ~/.ssh && touch ~/.ssh/authorized_keys \
+    && chmod 600 ~/.ssh/authorized_keys \
+    && rm /etc/ssh/ssh_host_*
 
-RUN curl -LsSf https://astral.sh/uv/install.sh | INSTALLER_NO_MODIFY_PATH=1 sh && \
-    uv python install --preview --default
+RUN curl -LsSf https://astral.sh/uv/install.sh | INSTALLER_NO_MODIFY_PATH=1 sh \
+    && uv python install --preview --default

docker/base/efa.Dockerfile

@@ -5,10 +5,9 @@ INCLUDE+ base/Dockerfile.common
 ENV NCCL_HOME=/usr/local
 ENV CUDA_HOME=/usr/local/cuda
 ENV LIBFABRIC_PATH=/opt/amazon/efa
-ENV MPI_HOME=/opt/amazon/openmpi
+ENV OPEN_MPI_PATH=/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}"
+ENV PATH="${LIBFABRIC_PATH}/bin:${OPEN_MPI_PATH}/bin:${PATH}"
 
 ARG EFA_VERSION=1.38.1
 ARG NCCL_VERSION=2.26.2-1
@@ -39,18 +38,20 @@ RUN cuda_version=$(echo ${CUDA_VERSION} | awk -F . '{ print $1"-"$2 }') \
     && ./configure \
         --with-cuda=${CUDA_HOME} \
         --with-libfabric=${LIBFABRIC_PATH} \
-        --with-mpi=${MPI_HOME} \
+        --with-mpi=${OPEN_MPI_PATH} \
         --with-nccl=${NCCL_HOME} \
         --disable-tests \
         --prefix=${NCCL_HOME} \
     && make -j$(numproc) \
     && make install \
-    && git clone https://github.com/NVIDIA/nccl-tests ${NCCL_TESTS_HOME} \
-    && cd ${NCCL_TESTS_HOME} \
+    && git clone https://github.com/NVIDIA/nccl-tests ${HOME}/nccl-tests \
+    && cd ${HOME}/nccl-tests \
     && make -j$(numproc) \
         MPI=1 \
-        MPI_HOME=${MPI_HOME} \
+        MPI_HOME=${OPEN_MPI_PATH} \
         CUDA_HOME=${CUDA_HOME} \
         NCCL_HOME=${NCCL_HOME} \
+    && ln -s ${HOME}/nccl-tests/build ${NCCL_TESTS_HOME} \
+    && echo "${OPEN_MPI_PATH}/lib" >> /etc/ld.so.conf.d/openmpi.conf \
     && echo "${NCCL_HOME}/lib" >> /etc/ld.so.conf.d/nccl.conf \
     && ldconfig