From 94b6f25ba69bb4cf4b9468c7d253df18d99ad2ad Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonatan=20K=C5=82osko?= <jonatanklosko@gmail.com>
Date: Thu, 25 Jul 2024 22:48:18 +0200
Subject: [PATCH] Reduce CUDA images size (#2724)

---
 .dockerignore                      |  2 ++
 Dockerfile                         |  2 ++
 docker/base/elixir-cuda.dockerfile | 33 +++++++++++++-----------------
 docker/build_and_push.sh           |  2 +-
 versions                           |  2 +-
 5 files changed, 20 insertions(+), 21 deletions(-)

diff --git a/.dockerignore b/.dockerignore
index 64601c04d..938280e5c 100644
--- a/.dockerignore
+++ b/.dockerignore
@@ -12,3 +12,5 @@ npm-debug.log
 /assets/node_modules/
 /tmp/
 /livebook
+# Ignore app release files, including build artifacts
+/rel/app
diff --git a/Dockerfile b/Dockerfile
index ecfdbf89e..40366c496 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -51,6 +51,8 @@ RUN mix do compile, release livebook
 # Consequently the release doesn't include ERTS as we have it anyway.
 FROM ${BASE_IMAGE}
 
+ARG DEBIAN_FRONTEND=noninteractive
+
 RUN apt-get update && apt-get upgrade -y && \
     apt-get install --no-install-recommends -y \
         # Runtime dependencies
diff --git a/docker/base/elixir-cuda.dockerfile b/docker/base/elixir-cuda.dockerfile
index 85a467211..ac7c955b1 100644
--- a/docker/base/elixir-cuda.dockerfile
+++ b/docker/base/elixir-cuda.dockerfile
@@ -2,25 +2,20 @@ ARG ELIXIR_VERSION
 ARG ERLANG_VERSION
 ARG UBUNTU_VERSION
 
+FROM hexpm/elixir:${ELIXIR_VERSION}-erlang-${ERLANG_VERSION}-ubuntu-${UBUNTU_VERSION}
+
 ARG CUDA_VERSION
 
-FROM hexpm/elixir:${ELIXIR_VERSION}-erlang-${ERLANG_VERSION}-ubuntu-${UBUNTU_VERSION} AS elixir
+RUN distro="ubuntu$(. /etc/lsb-release; echo "$DISTRIB_RELEASE" | tr -d '.')" && \
+    # Official Docker images use the sbsa packages when targetting arm64.
+    # See https://gitlab.com/nvidia/container-images/cuda/-/blob/85f465ea3343a2d7f7753a0a838701999ed58a01/dist/12.5.1/ubuntu2204/base/Dockerfile#L12
+    arch="$(if [ "$(uname -m)" = "aarch64" ]; then echo "sbsa"; else echo "x86_64"; fi)" && \
+    apt update -q && apt install -y ca-certificates wget && \
+    wget -qO /tmp/cuda-keyring.deb https://developer.download.nvidia.com/compute/cuda/repos/$distro/$arch/cuda-keyring_1.1-1_all.deb && \
+    dpkg -i /tmp/cuda-keyring.deb && apt update -q
 
-FROM nvidia/cuda:${CUDA_VERSION}-cudnn8-devel-ubuntu20.04
-
-ENV DEBIAN_FRONTEND nonintaeractive
-ENV LANG=C.UTF-8
-
-# Erlang runtime dependencies, see https://github.com/hexpm/bob/blob/3b5721dccdfe9d59766f374e7b4fb7fb8a7c720e/priv/scripts/docker/erlang-ubuntu-focal.dockerfile#L41-L45
-RUN apt-get update && \
-  apt-get -y --no-install-recommends install \
-    libodbc1 \
-    libssl1.1 \
-    libsctp1
-
-# We copy the top-level directory first to preserve symlinks in /usr/local/bin
-COPY --from=elixir /usr/local /usr/ELIXIR_LOCAL
-
-RUN cp -r /usr/ELIXIR_LOCAL/lib/* /usr/local/lib && \
-  cp -r /usr/ELIXIR_LOCAL/bin/* /usr/local/bin && \
-  rm -rf /usr/ELIXIR_LOCAL
+# In order to minimize the image size, we install only a subset of
+# the CUDA toolkit that is required by Elixir numerical packages
+# (nvcc and runtime libraries). Note that we do not need to install
+# the driver, it is already provided by NVIDIA Container Toolkit.
+RUN apt install -y git cuda-nvcc-${CUDA_VERSION} cuda-libraries-${CUDA_VERSION} libcudnn8
diff --git a/docker/build_and_push.sh b/docker/build_and_push.sh
index a1d23cf43..872ef941c 100755
--- a/docker/build_and_push.sh
+++ b/docker/build_and_push.sh
@@ -18,7 +18,7 @@ if [ -z "$livebook_version" ]; then
 fi
 
 cuda_tag_list=("cuda11.8" "cuda12.1")
-cuda_version_list=("11.8.0" "12.1.0")
+cuda_version_list=("11-8" "12-1")
 
 for idx in "${!cuda_tag_list[@]}"; do
   cuda_tag="${cuda_tag_list[idx]}"
diff --git a/versions b/versions
index bc1aa28dd..8c2bbb299 100644
--- a/versions
+++ b/versions
@@ -3,4 +3,4 @@ otp="27.0"
 openssl="1.1.1s"
 rebar3="3.22.0"
 debian="bookworm-20240701-slim"
-ubuntu="focal-20240530"
+ubuntu="jammy-20240530"