use cuda libraries only for runtime

.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

.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 }}'

Dockerfile

@@ -1,4 +1,4 @@
 FROM --platform=linux/amd64 ubuntu:22.04 as build-image
 
 WORKDIR /src
 RUN apt-get update && apt-get install -y \
@@ -16,21 +16,21 @@
     wget
 
 RUN curl https://sh.rustup.rs -sSf | sh -s -- -y
 ENV PATH "/root/.cargo/bin:${PATH}"
 ENV RUSTUP_HOME "/root/.rustup"
 ENV CARGO_HOME "/root/.cargo"
 RUN rustup toolchain install nightly-2024-07-10
 RUN rustup default nightly-2024-07-10
 RUN rustup component add cargo
 RUN cargo install cargo-build-deps \
     && cargo install cargo-edit
 
 FROM --platform=linux/amd64 build-image as build-app
 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

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/*
@@ -12,7 +12,7 @@
 ARG NCCL_INSTALL_VERSION=v2.22.3-1
 ARG AWS_OFI_NCCL_VERSION=aws
 
 ENV LD_LIBRARY_PATH=/usr/local/cuda/extras/CUPTI/lib64:/opt/amazon/openmpi/lib:/opt/nccl/build/lib:/opt/amazon/efa/lib:/opt/aws-ofi-nccl/install/lib:/usr/local/lib:$LD_LIBRARY_PATH
 ENV PATH=/opt/amazon/openmpi/bin/:/opt/amazon/efa/bin:$PATH
 
 RUN git clone -b ${GDRCOPY_VERSION} https://github.com/NVIDIA/gdrcopy.git /tmp/gdrcopy \
@@ -20,8 +20,8 @@
     && make prefix=/opt/gdrcopy install
 
 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
 
 # Install EFA installer
@@ -46,3 +46,29 @@
     --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