From 863ac616b792590e3f0f82630adda2b59b927440 Mon Sep 17 00:00:00 2001
From: danielle-tfh <danielle.nagar@toolsforhumanity.com>
Date: Wed, 22 Jan 2025 18:51:03 +0100
Subject: [PATCH] test smaller image size

---
 .github/workflows/build-and-push-base.yaml    |  2 +-
 .../workflows/temp-branch-build-and-push.yaml |  2 +-
 Dockerfile                                    |  2 +-
 Dockerfile.base                               | 28 ++++++++++++++++++-
 4 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/.github/workflows/build-and-push-base.yaml b/.github/workflows/build-and-push-base.yaml
index e35c72b60..7afc262f1 100644
--- a/.github/workflows/build-and-push-base.yaml
+++ b/.github/workflows/build-and-push-base.yaml
@@ -41,7 +41,7 @@ jobs:
           context: .
           push: true
           tags: |
-            ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:cuda${{ env.CUDA_VERSION }}-nccl${{ env.NCCL_VERSION }}
+            ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:cuda-libraries${{ env.CUDA_VERSION }}-nccl${{ env.NCCL_VERSION }}
           platforms: linux/amd64
           cache-from: type=gha
           cache-to: type=gha,mode=max
diff --git a/.github/workflows/temp-branch-build-and-push.yaml b/.github/workflows/temp-branch-build-and-push.yaml
index fc150d0b2..7dc9fe770 100644
--- a/.github/workflows/temp-branch-build-and-push.yaml
+++ b/.github/workflows/temp-branch-build-and-push.yaml
@@ -3,7 +3,7 @@ name: Branch - Build and push docker image
 on:
   push:
     branches:
-      - "add-graceful-shutdown-during-s3-sync"
+      - "reduce-size-docker-image"
 
 concurrency:
   group: '${{ github.workflow }} @ ${{ github.event.pull_request.head.label || github.head_ref || github.ref }}'
diff --git a/Dockerfile b/Dockerfile
index 147996daf..37a26b567 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -30,7 +30,7 @@ WORKDIR /src/gpu-iris-mpc
 COPY . .
 RUN cargo build --release --target x86_64-unknown-linux-gnu --bin nccl --bin server --bin client --bin key-manager --bin upgrade-server --bin upgrade-client --bin upgrade-checker --bin reshare-server --bin reshare-client
 
-FROM --platform=linux/amd64 ghcr.io/worldcoin/iris-mpc-base:cuda12_2-nccl2_22_3_1
+FROM --platform=linux/amd64 ghcr.io/worldcoin/iris-mpc-base:cuda-libraries12_2-nccl2_22_3_1
 ENV DEBIAN_FRONTEND=noninteractive
 
 # Include client, server and key-manager, upgrade-client and upgrade-server binaries
diff --git a/Dockerfile.base b/Dockerfile.base
index 5ddd4c845..7bdf1728a 100644
--- a/Dockerfile.base
+++ b/Dockerfile.base
@@ -1,4 +1,4 @@
-FROM --platform=linux/amd64 ubuntu:22.04 as build-image
+FROM --platform=linux/amd64 ubuntu:22.04 as base-image
 
 RUN apt-get update && apt-get install -y pkg-config wget libssl-dev ca-certificates protobuf-compiler \
     && rm -rf /var/lib/apt/lists/*
@@ -46,3 +46,29 @@ RUN cd /tmp \
     --with-nccl=/tmp/nccl/build \
     --with-mpi=/opt/amazon/openmpi/ \
     && make && make install
+
+
+# Runtime Image
+FROM --platform=linux/amd64 ubuntu:22.04 as build-image
+
+RUN apt-get update && apt-get install -y pkg-config wget
+
+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-libraries-12-2 libnccl2=2.22.3-1+cuda12.2 libnccl-dev=2.22.3-1+cuda12.2 \
+    && rm -f cuda-keyring_1.1-1_all.deb
+
+# Set environment variables for runtime
+ENV LD_LIBRARY_PATH=/usr/local/cuda/extras/CUPTI/lib64:/opt/aws-ofi-nccl/install/lib:/usr/local/cuda/lib64:$LD_LIBRARY_PATH
+ENV PATH=/opt/aws-ofi-nccl/install/bin:/usr/local/cuda/bin:$PATH
+
+# Copy necessary files from build stage
+COPY --from=base-image /opt/gdrcopy /opt/gdrcopy
+COPY --from=base-image /opt/aws-ofi-nccl/install /opt/aws-ofi-nccl/install
+COPY --from=base-image /opt/amazon/efa /opt/amazon/efa
+
+ENV LD_LIBRARY_PATH=/opt/gdrcopy/lib:/usr/local/cuda/compat:$LD_LIBRARY_PATH
+ENV LIBRARY_PATH=/opt/gdrcopy/lib:/usr/local/cuda/compat/:$LIBRARY_PATH
+ENV CPATH=/opt/gdrcopy/include:$CPATH
+ENV PATH=/opt/gdrcopy/bin:$PATH