From 31ff34776854a8ac0fb3805c8b454488229ebeba Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonatan=20K=C5=82osko?= <jonatanklosko@gmail.com>
Date: Tue, 17 Jun 2025 21:01:09 +0200
Subject: [PATCH] Bump Docker images to CUDA 12.8 and include NCCL (#3018)

---
 .github/workflows/release.yml | 3 ++-
 Dockerfile                    | 6 +++---
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
index d55836792..f0cc88fe2 100644
--- a/.github/workflows/release.yml
+++ b/.github/workflows/release.yml
@@ -75,7 +75,8 @@ jobs:
             tag_suffix: "-cuda12"
             build_args: |
               VARIANT=cuda
-              CUDA_VERSION=12-6
+              CUDA_VERSION_MAJOR=12
+              CUDA_VERSION_MINOR=8
     steps:
       - uses: actions/checkout@v4
       - run: |
diff --git a/Dockerfile b/Dockerfile
index 522e74656..a7520bf92 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -7,7 +7,8 @@ FROM ${BASE_IMAGE} AS base-default
 
 FROM ${BASE_IMAGE} AS base-cuda
 
-ARG CUDA_VERSION
+ARG CUDA_VERSION_MAJOR
+ARG CUDA_VERSION_MINOR
 
 RUN distro="ubuntu$(. /etc/lsb-release; echo "$DISTRIB_RELEASE" | tr -d '.')" && \
   # Official Docker images use the sbsa packages when targetting arm64.
@@ -20,8 +21,7 @@ RUN distro="ubuntu$(. /etc/lsb-release; echo "$DISTRIB_RELEASE" | tr -d '.')" &&
   # 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.
-  cuda_version="${CUDA_VERSION}" && cuda_major="${cuda_version%-*}" && \
-  apt-get install -y git cuda-nvcc-${CUDA_VERSION} cuda-libraries-${CUDA_VERSION} libcudnn9-cuda-$cuda_major && \
+  apt-get install -y git cuda-nvcc-${CUDA_VERSION_MAJOR}-${CUDA_VERSION_MINOR} cuda-libraries-${CUDA_VERSION_MAJOR}-${CUDA_VERSION_MINOR} libcudnn9-cuda-${CUDA_VERSION_MAJOR} libnccl2=*+cuda${CUDA_VERSION_MAJOR}.${CUDA_VERSION_MINOR} && \
   apt-get clean -y && rm -rf /var/lib/apt/lists/*
 
 ENV PATH="/usr/local/nvidia/bin:/usr/local/cuda/bin:$PATH"