chore: cuda 12.1.1 for GPU docker images

jolibrain · Nov 3, 2024 · ee7ad32 · ee7ad32
1 parent 372e106
commit ee7ad32
Showing 1 changed file with 11 additions and 2 deletions.
diff --git a/docker/gpu.Dockerfile b/docker/gpu.Dockerfile
@@ -1,5 +1,14 @@
 # syntax = docker/dockerfile:1.0-experimental
-FROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04 AS build
+
+ARG DD_UBUNTU_VERSION=22.04
+ARG DD_CUDA_VERSION=12.1.1
+ARG DD_CUDNN_VERSION=8
+FROM nvidia/cuda:${DD_CUDA_VERSION}-cudnn${DD_CUDNN_VERSION}-devel-ubuntu${DD_UBUNTU_VERSION} AS build
+
+
+RUN echo UBUNTU_VERSION=${DD_UBUNTU_VERSION} >> /image-info
+RUN echo CUDA_VERSION=${DD_CUDA_VERSION} >> /image-info
+RUN echo CUDNN_VERSION=${DD_CUDNN_VERSION} >> /image-info
 
 ARG DEEPDETECT_RELEASE=OFF
 ARG DEEPDETECT_ARCH=gpu
@@ -106,7 +115,7 @@ RUN --mount=type=cache,target=/ccache/ mkdir build && cd build && ../build.sh
 RUN ./docker/get_libs.sh
 
 # Build final Docker image
-FROM nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04 AS runtime
+FROM nvidia/cuda:${DD_CUDA_VERSION}-cudnn${DD_CUDNN_VERSION}-runtime-ubuntu${DD_UBUNTU_VERSION} AS runtime
 
 ARG DEEPDETECT_ARCH=gpu