Update Dockerfile

Dockerfile

@@ -1,18 +1,44 @@
-FROM ubuntu:20.04 as base
+ARG IMAGE_NAME
+FROM ${IMAGE_NAME}:11.8.0-runtime-ubuntu20.04 as base
+
+ENV NV_CUDA_LIB_VERSION "11.8.0-1"
 
 FROM base as base-amd64
 
-ENV NVARCH x86_64
+ENV NV_CUDA_CUDART_DEV_VERSION 11.8.89-1
+ENV NV_NVML_DEV_VERSION 11.8.86-1
+ENV NV_LIBCUSPARSE_DEV_VERSION 11.7.5.86-1
+ENV NV_LIBNPP_DEV_VERSION 11.8.0.86-1
+ENV NV_LIBNPP_DEV_PACKAGE libnpp-dev-11-8=${NV_LIBNPP_DEV_VERSION}
+
+ENV NV_LIBCUBLAS_DEV_VERSION 11.11.3.6-1
+ENV NV_LIBCUBLAS_DEV_PACKAGE_NAME libcublas-dev-11-8
+ENV NV_LIBCUBLAS_DEV_PACKAGE ${NV_LIBCUBLAS_DEV_PACKAGE_NAME}=${NV_LIBCUBLAS_DEV_VERSION}
 
-ENV NVIDIA_REQUIRE_CUDA "cuda>=11.8 brand=tesla,driver>=450,driver<451 brand=tesla,driver>=470,driver<471 brand=unknown,driver>=470,driver<471 brand=nvidia,driver>=470,driver<471 brand=nvidiartx,driver>=470,driver<471 brand=geforce,driver>=470,driver<471 brand=geforcertx,driver>=470,driver<471 brand=quadro,driver>=470,driver<471 brand=quadrortx,driver>=470,driver<471 brand=titan,driver>=470,driver<471 brand=titanrtx,driver>=470,driver<471 brand=tesla,driver>=510,driver<511 brand=unknown,driver>=510,driver<511 brand=nvidia,driver>=510,driver<511 brand=nvidiartx,driver>=510,driver<511 brand=geforce,driver>=510,driver<511 brand=geforcertx,driver>=510,driver<511 brand=quadro,driver>=510,driver<511 brand=quadrortx,driver>=510,driver<511 brand=titan,driver>=510,driver<511 brand=titanrtx,driver>=510,driver<511 brand=tesla,driver>=515,driver<516 brand=unknown,driver>=515,driver<516 brand=nvidia,driver>=515,driver<516 brand=nvidiartx,driver>=515,driver<516 brand=geforce,driver>=515,driver<516 brand=geforcertx,driver>=515,driver<516 brand=quadro,driver>=515,driver<516 brand=quadrortx,driver>=515,driver<516 brand=titan,driver>=515,driver<516 brand=titanrtx,driver>=515,driver<516"
-ENV NV_CUDA_CUDART_VERSION 11.8.89-1
-ENV NV_CUDA_COMPAT_PACKAGE cuda-compat-11-8
+ENV NV_NVPROF_VERSION 11.8.87-1
+ENV NV_NVPROF_DEV_PACKAGE cuda-nvprof-11-8=${NV_NVPROF_VERSION}
 
+ENV NV_LIBNCCL_DEV_PACKAGE_NAME libnccl-dev
+ENV NV_LIBNCCL_DEV_PACKAGE_VERSION 2.16.2-1
+ENV NCCL_VERSION 2.16.2-1
+ENV NV_LIBNCCL_DEV_PACKAGE ${NV_LIBNCCL_DEV_PACKAGE_NAME}=${NV_LIBNCCL_DEV_PACKAGE_VERSION}+cuda11.8
 FROM base as base-arm64
 
-ENV NVARCH sbsa
-ENV NVIDIA_REQUIRE_CUDA "cuda>=11.8"
-ENV NV_CUDA_CUDART_VERSION 11.8.89-1
+ENV NV_CUDA_CUDART_DEV_VERSION 11.8.89-1
+ENV NV_NVML_DEV_VERSION 11.8.86-1
+ENV NV_LIBCUSPARSE_DEV_VERSION 11.7.5.86-1
+ENV NV_LIBNPP_DEV_VERSION 11.8.0.86-1
+ENV NV_LIBNPP_DEV_PACKAGE libnpp-dev-11-8=${NV_LIBNPP_DEV_VERSION}
+
+ENV NV_LIBCUBLAS_DEV_PACKAGE_NAME libcublas-dev-11-8
+ENV NV_LIBCUBLAS_DEV_VERSION 11.11.3.6-1
+ENV NV_LIBCUBLAS_DEV_PACKAGE ${NV_LIBCUBLAS_DEV_PACKAGE_NAME}=${NV_LIBCUBLAS_DEV_VERSION}
+
+ENV NV_LIBNCCL_DEV_PACKAGE_NAME libnccl-dev
+ENV NV_LIBNCCL_DEV_PACKAGE_VERSION 2.16.2-1
+ENV NCCL_VERSION 2.16.2-1
+ENV NV_LIBNCCL_DEV_PACKAGE ${NV_LIBNCCL_DEV_PACKAGE_NAME}=${NV_LIBNCCL_DEV_PACKAGE_VERSION}+cuda11.8
+
 
 FROM base-${TARGETARCH}
 
@@ -21,30 +47,22 @@ ARG TARGETARCH
 LABEL maintainer "NVIDIA CORPORATION <[email protected]>"
 
 RUN apt-get update && apt-get install -y --no-install-recommends \
-    gnupg2 curl ca-certificates && \
-    curl -fsSL https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/${NVARCH}/3bf863cc.pub | apt-key add - && \
-    echo "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/${NVARCH} /" > /etc/apt/sources.list.d/cuda.list && \
-    apt-get purge --autoremove -y curl \
+    libtinfo5 libncursesw5 \
+    cuda-cudart-dev-11-8=${NV_CUDA_CUDART_DEV_VERSION} \
+    cuda-command-line-tools-11-8=${NV_CUDA_LIB_VERSION} \
+    cuda-minimal-build-11-8=${NV_CUDA_LIB_VERSION} \
+    cuda-libraries-dev-11-8=${NV_CUDA_LIB_VERSION} \
+    cuda-nvml-dev-11-8=${NV_NVML_DEV_VERSION} \
+    ${NV_NVPROF_DEV_PACKAGE} \
+    ${NV_LIBNPP_DEV_PACKAGE} \
+    libcusparse-dev-11-8=${NV_LIBCUSPARSE_DEV_VERSION} \
+    ${NV_LIBCUBLAS_DEV_PACKAGE} \
+    ${NV_LIBNCCL_DEV_PACKAGE} \
     && rm -rf /var/lib/apt/lists/*
 
-ENV CUDA_VERSION 11.8.0
-
-# 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-8=${NV_CUDA_CUDART_VERSION} \
-    ${NV_CUDA_COMPAT_PACKAGE} \
-    && rm -rf /var/lib/apt/lists/*
-
-# Required for nvidia-docker v1
-RUN echo "/usr/local/nvidia/lib" >> /etc/ld.so.conf.d/nvidia.conf \
-    && echo "/usr/local/nvidia/lib64" >> /etc/ld.so.conf.d/nvidia.conf
-
-ENV PATH /usr/local/nvidia/bin:/usr/local/cuda/bin:${PATH}
-ENV LD_LIBRARY_PATH /usr/local/nvidia/lib:/usr/local/nvidia/lib64
-
-# nvidia-container-runtime
-ENV NVIDIA_VISIBLE_DEVICES all
-ENV NVIDIA_DRIVER_CAPABILITIES compute,utility
+# Keep apt from auto upgrading the cublas and nccl packages. See https://gitlab.com/nvidia/container-images/cuda/-/issues/88
+RUN apt-mark hold ${NV_LIBCUBLAS_DEV_PACKAGE_NAME} ${NV_LIBNCCL_DEV_PACKAGE_NAME}
+ENV LIBRARY_PATH /usr/local/cuda/lib64/stubs
 
 
 WORKDIR /app