Build docker gpu image

soedinglab · Nov 16, 2024 · 8616910 · 8616910
1 parent 3740224
commit 8616910
Show file tree

Hide file tree

Showing 2 changed files with 108 additions and 25 deletions.
diff --git a/.github/workflows/docker.yml b/.github/workflows/docker.yml
@@ -11,9 +11,11 @@ on:
       tag:
         required: true
         type: string
+        description: "Docker tag"
       latest:
         default: false
         type: boolean
+        description: "Mark as latest"
 
 
 env:
@@ -77,4 +79,61 @@ jobs:
             ${{ steps.dispatch_tag.outputs.tag }}
             ${{ steps.dispatch_tag.outputs.latest }}
           labels: ${{ steps.meta.outputs.labels }}
+  build-and-push-gpu-image:
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+      packages: write
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v3
+
+      # - name: Set up QEMU
+      #   uses: docker/setup-qemu-action@v2
+
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v2
+
+      - name: Log in to the Container registry
+        uses: docker/login-action@f054a8b539a109f9f41c372932f1ae047eff08c9
+        with:
+          registry: ${{ env.REGISTRY }}
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Extract metadata (tags, labels) for Docker
+        id: meta
+        uses: docker/metadata-action@98669ae865ea3cffbcbaa878cf57c20bbf1c6c38
+        with:
+          images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}
 
+      - name: Tag for workflow_dispatch
+        id: dispatch_tag
+        run: |
+          if [ x"$TAG" != x"" ];then
+            echo "::set-output name=tag::${FULL_TAG}"
+          fi
+          if [ x"$LATEST" = x"true" ]; then
+            echo "::set-output name=latest::${LATEST_TAG}"
+          fi
+        env:
+          FULL_TAG: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:${{ github.event.inputs.tag }}-cuda12
+          LATEST_TAG: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:latest-cuda12
+          TAG: ${{ github.event.inputs.tag }}-cuda12
+          LATEST: ${{ github.event.inputs.latest }}-cuda12
+
+      - name: Build and push Docker image
+        uses: docker/build-push-action@ad44023a93711e3deb337508980b4b5e9bcdc5dc
+        with:
+          context: .
+          platforms: linux/amd64
+          push: true
+          cache-from: type=gha
+          cache-to: type=gha,mode=max
+          build-args: |
+            GPU=1
+          tags: |
+            ${{ steps.meta.outputs.tags }}
+            ${{ steps.dispatch_tag.outputs.tag }}
+            ${{ steps.dispatch_tag.outputs.latest }}
+          labels: ${{ steps.meta.outputs.labels }}
diff --git a/Dockerfile b/Dockerfile
@@ -1,15 +1,21 @@
 ARG APP=mmseqs
-FROM --platform=$BUILDPLATFORM debian:stable-slim as builder
+FROM --platform=$BUILDPLATFORM debian:bookworm-slim AS builder
 ARG TARGETARCH
 ARG APP
+ARG GPU
 
 RUN dpkg --add-architecture $TARGETARCH \
     && apt-get update \
     && apt-get install -y \
-      build-essential cmake xxd git \
+      build-essential cmake xxd git wget \
       zlib1g-dev libbz2-dev libatomic1 \
-      crossbuild-essential-$TARGETARCH zlib1g-dev:$TARGETARCH libbz2-dev:$TARGETARCH \
-    && rm -rf /var/lib/apt/lists/*
+      crossbuild-essential-$TARGETARCH zlib1g-dev:$TARGETARCH libbz2-dev:$TARGETARCH; \
+    if [ "$GPU" = "1" ]; then \
+      wget https://developer.download.nvidia.com/compute/cuda/repos/debian12/x86_64/cuda-keyring_1.1-1_all.deb; \
+      dpkg -i cuda-keyring_1.1-1_all.deb; \
+      apt-get update && apt-get install -y cuda-nvcc-12-6 ninja-build; \
+    fi; \
+    rm -rf /var/lib/apt/lists/*;
 
 WORKDIR /opt/build
 ADD . .
@@ -22,33 +28,51 @@ RUN if [ "$TARGETARCH" = "arm64" ]; then \
       mv src/${APP} /opt/build/${APP}_arch; \
       touch /opt/build/${APP}_sse2 /opt/build/${APP}_sse41 /opt/build/${APP}_avx2; \
     else \
-      mkdir -p build_sse2/src && mkdir -p build_sse41/src && mkdir -p build_avx2/src; \
-      cd /opt/build/build_sse2; \
-      cmake -DHAVE_SSE2=1 -DHAVE_MPI=0 -DHAVE_TESTS=0 -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=. ..; \
-      make -j $(nproc --all); \
-      mv src/${APP} /opt/build/${APP}_sse2; \
-      cd /opt/build/build_sse41; \
-      cmake -DHAVE_SSE4_1=1 -DHAVE_MPI=0 -DHAVE_TESTS=0 -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=. ..; \
-      make -j $(nproc --all); \
-      mv src/${APP} /opt/build/${APP}_sse41; \
-      cd /opt/build/build_avx2; \
-      cmake -DHAVE_AVX2=1 -DHAVE_MPI=0 -DHAVE_TESTS=0 -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=. ..; \
-      make -j $(nproc --all); \
-      mv src/${APP} /opt/build/${APP}_avx2; \
-      touch /opt/build/${APP}_arch; \
+      if [ "$GPU" = "1" ]; then \
+        export CUDACXX=/usr/local/cuda/bin/nvcc; \
+        mkdir -p build_avx2/src; \
+        cd /opt/build/build_avx2; \
+        cmake -GNinja -DHAVE_AVX2=1 -DHAVE_MPI=0 -DHAVE_TESTS=0 \
+        -DPREFER_STATIC=1 -DBUILD_SHARED_LIBS=OFF \
+        -DCMAKE_EXE_LINKER_FLAGS="-static-libgcc -static-libstdc++" \
+        -DENABLE_CUDA=1 -DCMAKE_CUDA_ARCHITECTURES="75-real;80-real;86-real;89-real;90" \
+        -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=. ..; \
+        cmake --build . -j$(nproc --all); \
+        mv src/${APP} /opt/build/${APP}_avx2; \
+        touch /opt/build/${APP}_arch /opt/build/${APP}_sse41 /opt/build/${APP}_sse2; \
+      else \
+        mkdir -p build_sse2/src && mkdir -p build_sse41/src && mkdir -p build_avx2/src; \
+        cd /opt/build/build_sse2; \
+        cmake -DHAVE_SSE2=1 -DHAVE_MPI=0 -DHAVE_TESTS=0 -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=. ..; \
+        make -j $(nproc --all); \
+        mv src/${APP} /opt/build/${APP}_sse2; \
+        cd /opt/build/build_sse41; \
+        cmake -DHAVE_SSE4_1=1 -DHAVE_MPI=0 -DHAVE_TESTS=0 -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=. ..; \
+        make -j $(nproc --all); \
+        mv src/${APP} /opt/build/${APP}_sse41; \
+        cd /opt/build/build_avx2; \
+        cmake -DHAVE_AVX2=1 -DHAVE_MPI=0 -DHAVE_TESTS=0 -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=. ..; \
+        make -j $(nproc --all); \
+        mv src/${APP} /opt/build/${APP}_avx2; \
+        touch /opt/build/${APP}_arch; \
+      fi; \
     fi
 
-FROM debian:stable-slim
+FROM debian:bookworm-slim
 ARG TARGETARCH
 ARG APP
-
-RUN apt-get update && apt-get install -y \
-      gawk bash grep libstdc++6 libgomp1 libatomic1 zlib1g libbz2-1.0 wget tar \
-    && rm -rf /var/lib/apt/lists/*
+ARG GPU
 
 COPY --from=builder /opt/build/${APP}_arch /opt/build/${APP}_sse2 /opt/build/${APP}_sse41 /opt/build/${APP}_avx2 /usr/local/bin/
 ADD util/${APP}_wrapper.sh /usr/local/bin/entrypoint
-RUN if [ "$TARGETARCH" = "arm64" ]; then rm -f /usr/local/bin/entrypoint; ln -s /usr/local/bin/${APP}_arch /usr/local/bin/entrypoint; fi
 
-ENTRYPOINT ["/usr/local/bin/entrypoint"]
+RUN apt-get update && apt-get install -y \
+      gawk bash grep libstdc++6 libgomp1 libatomic1 zlib1g libbz2-1.0 wget tar aria2 \
+    && rm -rf /var/lib/apt/lists/*; \
+    if [ "$TARGETARCH" = "arm64" ]; then \
+      rm -f /usr/local/bin/entrypoint; ln -s /usr/local/bin/${APP}_arch /usr/local/bin/entrypoint; \
+    elif [ "$GPU" = "1" ]; then \
+      rm -f /usr/local/bin/entrypoint; ln -s /usr/local/bin/${APP}_avx2 /usr/local/bin/entrypoint; \
+    fi
 
+ENTRYPOINT ["/usr/local/bin/entrypoint"]