chore: add latest cuda versions

.github/workflows/build_and_push_to_docker.yml

@@ -107,9 +107,41 @@ jobs:
         include:
           # taken from tensorflow compatibility chart at https://www.tensorflow.org/install/source#gpu
           - CUDA_VERSION: "11.0"
-            BASE_IMAGE: renku/renkulab-py:python-3.8.8
+            PYTHON_VERSION: "3.8.8"
+            EXTRA_LIBRARIES: "libcusolver-11-2"
+            CUDA_CUDART_PACKAGE: "cuda-cudart-11-0=11.0.221-1"
+            CUDA_COMPAT_PACKAGE: "cuda-compat-11-0"
+            LIBCUDNN_PACKAGE: "libcudnn8=8.0.5.39-1+cuda11.0"
+          - CUDA_VERSION: "11.1"
+            PYTHON_VERSION: "3.8.8"
+            EXTRA_LIBRARIES: "libcusolver-11-2"
+            CUDA_CUDART_PACKAGE: "cuda-cudart-11-1=11.1.74-1"
+            CUDA_COMPAT_PACKAGE: "cuda-compat-11-1"
+            LIBCUDNN_PACKAGE: "libcudnn8=8.0.5.39-1+cuda11.1"
           - CUDA_VERSION: "11.2"
-            BASE_IMAGE: renku/renkulab-py:python-3.9.7
+            PYTHON_VERSION: "3.9.7"
+            EXTRA_LIBRARIES: ""
+            CUDA_CUDART_PACKAGE: "cuda-cudart-11-2=11.2.152-1"
+            CUDA_COMPAT_PACKAGE: "cuda-compat-11-2"
+            LIBCUDNN_PACKAGE: "libcudnn8=8.1.1.33-1+cuda11.2"
+          - CUDA_VERSION: "11.3"
+            PYTHON_VERSION: "3.9.7"
+            EXTRA_LIBRARIES: ""
+            CUDA_CUDART_PACKAGE: "cuda-cudart-11-3=11.3.109-1"
+            CUDA_COMPAT_PACKAGE: "cuda-compat-11-3"
+            LIBCUDNN_PACKAGE: "libcudnn8=8.2.1.32-1+cuda11.3"
+          - CUDA_VERSION: "11.4"
+            PYTHON_VERSION: "3.9.7"
+            EXTRA_LIBRARIES: ""
+            CUDA_CUDART_PACKAGE: "cuda-cudart-11-4=11.4.148-1"
+            CUDA_COMPAT_PACKAGE: "cuda-compat-11-4"
+            LIBCUDNN_PACKAGE: "libcudnn8=8.2.4.15-1+cuda11.4"
+          - CUDA_VERSION: "11.5"
+            PYTHON_VERSION: "3.9.7"
+            EXTRA_LIBRARIES: ""
+            CUDA_CUDART_PACKAGE: "cuda-cudart-11-5=11.5.117-1"
+            CUDA_COMPAT_PACKAGE: "cuda-compat-11-5"
+            LIBCUDNN_PACKAGE: "libcudnn8=8.3.2.44-1+cuda11.5"
 
     steps:
     - name: Docker Login
@@ -130,8 +162,13 @@ jobs:
           export LABEL=$(echo ${{ github.sha }} | cut -c 1-7)
         fi
 
-        docker build -f Dockerfile-${{ matrix.CUDA_VERSION }} . \
-          --build-arg BASE_IMAGE="${{ matrix.BASE_IMAGE }}-$LABEL" \
+        docker build . \
+          --build-arg BASE_IMAGE="renku/renkulab-py:python-${{ matrix.PYTHON_VERSION }}-$LABEL" \
+          --build-arg CUDA_VERSION="${{ matrix.CUDA_VERSION }}" \
+          --build-arg EXTRA_LIBRARIES="${{ matrix.EXTRA_LIBRARIES }}" \
+          --build-arg CUDA_CUDART_PACKAGE="${{ matrix.CUDA_CUDART_PACKAGE }}" \
+          --build-arg CUDA_COMPAT_PACKAGE="${{ matrix.CUDA_COMPAT_PACKAGE }}" \
+          --build-arg LIBCUDNN_PACKAGE="${{ matrix.LIBCUDNN_PACKAGE }}" \
           --tag $DOCKER_NAME-cuda:${{ matrix.CUDA_VERSION }}-$LABEL
         echo "::set-output name=IMAGE_NAME::$DOCKER_NAME-cuda:${{ matrix.CUDA_VERSION }}-$LABEL"
     - name: Image Acceptance Tests
@@ -147,7 +184,7 @@ jobs:
         docker push ${{ steps.build-image.outputs.IMAGE_NAME }}
 
         # on master push latest image
-        if [ "$REF" == "refs/heads/master" ] && [ "${{ matrix.CUDA_VERSION }}" == "11.2" ]; then
+        if [ "$REF" == "refs/heads/master" ] && [ "${{ matrix.CUDA_VERSION }}" == "11.5" ]; then
           docker tag ${{ steps.build-image.outputs.IMAGE_NAME }} $DOCKER_NAME-cuda:latest
           docker push $DOCKER_NAME-cuda:latest
         fi

docker/cuda/Dockerfile-11.0 → docker/cuda/Dockerfile

@@ -13,13 +13,17 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
     apt-get purge --autoremove -y curl \
     && rm -rf /var/lib/apt/lists/*
 
-ENV CUDA_VERSION 11.0.3
+ARG CUDA_VERSION=11.2
+ENV CUDA_VERSION $CUDA_VERSION
+ARG CUDA_CUDART_PACKAGE=cuda-cudart-11-2=11.2.152-1
+ARG CUDA_COMPAT_PACKAGE=cuda-compat-11-2
+ARG EXTRA_LIBRARIES=""
 
 # For libraries in the cuda-compat-* package: https://docs.nvidia.com/cuda/eula/index.html#attachment-a
 RUN apt-get update && apt-get install -y --no-install-recommends \
-    cuda-cudart-11-0=11.0.221-1 \
-    cuda-compat-11-0 \
-    && ln -s cuda-11.0 /usr/local/cuda && \
+    $CUDA_CUDART_PACKAGE \
+    $CUDA_COMPAT_PACKAGE \
+    && ln -s cuda-${CUDA_VERSION} /usr/local/cuda && \
     rm -rf /var/lib/apt/lists/*
 
 # Required for nvidia-docker v1
@@ -38,10 +42,9 @@ ENV NVIDIA_REQUIRE_CUDA "cuda>=11.0 brand=tesla,driver>=418,driver<419 brand=tes
 # Install TensorFlow (from https://github.com/tensorflow/tensorflow/blob/master/tensorflow/tools/dockerfiles/dockerfiles/gpu.Dockerfile)
 #
 
-ARG CUDA=11.0
-ARG CUDNN=8.0.5.39
-ARG CUDNN_MAJOR_VERSION=8
+ARG CUDA=${CUDA_VERSION}
 ARG LIB_DIR_PREFIX=x86_64
+ARG LIBCUDNN_PACKAGE="libcudnn8=8.1.1.33-1+cuda11.2"
 
 # Needed for string substitution
 SHELL ["/bin/bash", "-c"]
@@ -54,10 +57,10 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
     libcufft-${CUDA/./-} \
     libcurand-${CUDA/./-} \
     libcusolver-${CUDA/./-} \
-    libcusolver-11-2 \
     libcusparse-${CUDA/./-} \
+    $EXTRA_LIBRARIES \
     curl \
-    libcudnn8=${CUDNN}-1+cuda${CUDA} \
+    ${LIBCUDNN_PACKAGE} \
     libfreetype6-dev \
     libhdf5-serial-dev \
     libzmq3-dev \