[Docker] Add dstackai/efa image (#2422)

Author: un-def

.github/workflows/docker-efa.yml

@@ -0,0 +1,46 @@
+name: Build EFA Docker image
+
+on:
+  workflow_dispatch:
+    inputs:
+      image_name:
+        description: "Docker image name"
+        required: true
+        default: "dstackai/efa"
+      dstack_revision:
+        description: "Docker image revision"
+        required: true
+        default: 0
+
+jobs:
+  build-efa:
+    defaults:
+      run:
+        working-directory: docker/efa
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
+      - name: Login to DockerHub
+        uses: docker/login-action@v3
+        with:
+          username: ${{ secrets.DOCKERHUB_USERNAME }}
+          password: ${{ secrets.DOCKERHUB_TOKEN }}
+      - name: Build and upload to DockerHub
+        run: |
+          IMAGE_NAME=${{ inputs.image_name }}
+          BUILD_DATE=$(date --utc --iso-8601=seconds)Z
+          docker buildx build . \
+            --load \
+            --provenance=false \
+            --platform linux/amd64 \
+            --build-arg IMAGE_NAME=${IMAGE_NAME} \
+            --build-arg DSTACK_REVISION=${{ inputs.dstack_revision }} \
+            --build-arg BUILD_DATE=${BUILD_DATE} \
+            --tag ${IMAGE_NAME}:latest
+          VERSION=$(docker inspect --format '{{ index .Config.Labels "org.opencontainers.image.version" }}' ${IMAGE_NAME})
+          docker tag ${IMAGE_NAME}:latest ${IMAGE_NAME}:${VERSION}
+          docker push ${IMAGE_NAME}:${VERSION}
+          docker push ${IMAGE_NAME}:latest

docker/efa/Dockerfile

@@ -0,0 +1,79 @@
+ARG BASE_IMAGE=dstackai/base:py3.12-0.7-cuda-12.1
+
+FROM ${BASE_IMAGE}
+
+ENV PREFIX=/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:${PATH}"
+ENV LD_LIBRARY_PATH="${OPEN_MPI_PATH}/lib:${LD_LIBRARY_PATH}"
+
+# prerequisites
+
+RUN cuda_version=$(echo ${CUDA_VERSION} | awk -F . '{ print $1"-"$2 }') \
+    && apt-get update \
+    && apt-get install -y --no-install-recommends \
+        cuda-libraries-dev-${cuda_version} \
+        cuda-nvcc-${cuda_version} \
+        libhwloc-dev \
+        autoconf \
+        automake \
+        libtool
+
+# EFA
+
+ARG EFA_VERSION=1.38.1
+
+RUN cd $HOME \
+    && 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
+
+# NCCL
+
+ARG NCCL_VERSION=2.26.2-1
+
+RUN cd $HOME \
+    && git clone https://github.com/NVIDIA/nccl.git -b v${NCCL_VERSION} \
+    && cd nccl \
+    && make -j$(nproc) src.build BUILDDIR=${PREFIX}
+
+# AWS OFI NCCL
+
+ARG OFI_VERSION=1.14.0
+
+RUN cd $HOME \
+    && 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-libfabric=${LIBFABRIC_PATH} \
+        --with-mpi=${OPEN_MPI_PATH} \
+        --with-cuda=${CUDA_PATH} \
+        --with-nccl=${PREFIX} \
+        --disable-tests \
+        --prefix=${PREFIX} \
+    && make -j$(numproc) \
+    && make install
+
+# NCCL Tests
+
+RUN cd $HOME \
+    && git clone https://github.com/NVIDIA/nccl-tests \
+    && cd nccl-tests \
+    && make -j$(numproc) \
+        MPI=1 \
+        MPI_HOME=${OPEN_MPI_PATH} \
+        CUDA_HOME=${CUDA_PATH} \
+        NCCL_HOME=${PREFIX}
+
+ARG BUILD_DATE
+ARG IMAGE_NAME
+ARG DSTACK_REVISION
+
+LABEL org.opencontainers.image.title="${IMAGE_NAME}"
+LABEL org.opencontainers.image.version="${EFA_VERSION}-${DSTACK_REVISION}"
+LABEL org.opencontainers.image.created="${BUILD_DATE}"

docker/efa/README.md

@@ -0,0 +1,9 @@
+# dstack AWS EFA
+
+This image has the following installed:
+
+* CUDA 12.1
+* AWS EFA Installer 1.38.1 (Libfabric + Open MPI 4 + Open MPI 5)
+* NCCL 2.26.2-1
+* AWS OFI NCCL 1.14.0
+* NCCL Tests