From 6ae325cf2707ccbc76a54c8a7a9a1100501eb972 Mon Sep 17 00:00:00 2001
From: wojciechsromek <wojciech.sromek@toolsforhumanity.com>
Date: Thu, 8 Aug 2024 09:06:45 +0200
Subject: [PATCH] chore: Speeding up Docker build

---
 .github/workflows/build-and-push-base.yaml | 48 ++++++++++++++++++++++
 Dockerfile                                 | 13 ++----
 Dockerfile.base                            |  8 ++++
 3 files changed, 59 insertions(+), 10 deletions(-)
 create mode 100644 .github/workflows/build-and-push-base.yaml
 create mode 100644 Dockerfile.base

diff --git a/.github/workflows/build-and-push-base.yaml b/.github/workflows/build-and-push-base.yaml
new file mode 100644
index 000000000..e35c72b60
--- /dev/null
+++ b/.github/workflows/build-and-push-base.yaml
@@ -0,0 +1,48 @@
+name: Build and push base docker image (CUDA, NCCL)
+
+on:
+  push:
+    paths:
+      - Dockerfile.base
+      - .github/workflows/build-and-push-base.yaml
+
+  workflow_dispatch:
+
+env:
+  REGISTRY: ghcr.io
+  IMAGE_NAME: ${{ github.repository }}-base
+  CUDA_VERSION: 12_2
+  NCCL_VERSION: 2_22_3_1
+jobs:
+  docker:
+    runs-on:
+      labels: ubuntu-22.04-64core
+    permissions:
+      packages: write
+      contents: read
+      attestations: write
+      id-token: write
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+      - name: Set up QEMU
+        uses: docker/setup-qemu-action@v3
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
+      - name: Log in to the Container registry
+        uses: docker/login-action@v3
+        with:
+          registry: ${{ env.REGISTRY }}
+          username: ${{ github.repository_owner }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+      - name: Build and Push
+        uses: docker/build-push-action@v6
+        with:
+          context: .
+          push: true
+          tags: |
+            ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:cuda${{ env.CUDA_VERSION }}-nccl${{ env.NCCL_VERSION }}
+          platforms: linux/amd64
+          cache-from: type=gha
+          cache-to: type=gha,mode=max
+          file: Dockerfile.base
diff --git a/Dockerfile b/Dockerfile
index cfb3a7fac..c98177cab 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -29,20 +29,13 @@ WORKDIR /src/gpu-iris-mpc
 COPY . .
 RUN cargo build --release --target x86_64-unknown-linux-gnu --bin server --bin client --bin key-manager
 
-FROM --platform=linux/amd64 ubuntu:22.04
+FROM --platform=linux/amd64 ghcr.io/worldcoin/gpu-iris-mpc-base:cuda12_2-nccl2_23_2_1
 ENV DEBIAN_FRONTEND=noninteractive
+
+# Include client, server and key-manager
 COPY --from=build-app /src/gpu-iris-mpc/target/x86_64-unknown-linux-gnu/release/server /bin/server
-# include client for testing
 COPY --from=build-app /src/gpu-iris-mpc/target/x86_64-unknown-linux-gnu/release/client /bin/client
-# include key-manager for testing
 COPY --from=build-app /src/gpu-iris-mpc/target/x86_64-unknown-linux-gnu/release/key-manager /bin/key-manager
 
-RUN apt-get update && apt-get install -y pkg-config wget libssl-dev ca-certificates \
-    && rm -rf /var/lib/apt/lists/*
-RUN wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/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-toolkit-12-2 libnccl2=2.22.3-1+cuda12.2 libnccl-dev=2.22.3-1+cuda12.2
-
 USER 65534
 ENTRYPOINT ["/bin/server"]
diff --git a/Dockerfile.base b/Dockerfile.base
new file mode 100644
index 000000000..782f40900
--- /dev/null
+++ b/Dockerfile.base
@@ -0,0 +1,8 @@
+FROM --platform=linux/amd64 ubuntu:22.04 as build-image
+
+RUN apt-get update && apt-get install -y pkg-config wget libssl-dev ca-certificates \
+    && rm -rf /var/lib/apt/lists/*
+RUN wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/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-toolkit-12-2 libnccl2=2.22.3-1+cuda12.2 libnccl-dev=2.22.3-1+cuda12.2