Update cuda libraries to match TF

PiperOrigin-RevId: 616183821
tensorflow · Mar 15, 2024 · 45446cf · 45446cf
1 parent 7599531
commit 45446cf
Showing 1 changed file with 22 additions and 16 deletions.
diff --git a/tensorflow_serving/tools/docker/Dockerfile.devel-gpu b/tensorflow_serving/tools/docker/Dockerfile.devel-gpu
@@ -35,18 +35,31 @@ RUN curl -sSL --retry 5 https://github.com/tensorflow/serving/tarball/${TF_SERVI
 RUN /tensorflow-serving/tensorflow_serving/tools/docker/setup.sources.sh
 
 RUN apt-get install -y --no-install-recommends \
-        automake \
-        build-essential \
-        ca-certificates \
         cuda-command-line-tools-12-2 \
-        libcublas-12-2 \
-        libcublas-dev-12-2 \
-        cuda-nvrtc-12-2 \
-        cuda-nvrtc-dev-12-2 \
-        cuda-nvprune-12-2 \
+        cuda-cudart-dev-12-2 \
         cuda-nvcc-12-2 \
+        cuda-cupti-12-2 \
+        cuda-nvprune-12-2 \
+        cuda-libraries-12-2 \
+        cuda-libraries-dev-12-2 \
         cuda-nvml-dev-12-2 \
-        cuda-cudart-dev-12-2 \
+        libcufft-12-2 \
+        libcurand-12-2 \
+        libcusolver-dev-12-2 \
+        libcusparse-dev-12-2 \
+        libcublas-12-2 \
+        libcublas-dev-12-2 \
+        libnccl-dev=2.18.5-1+cuda12.2 \
+        libnccl2=2.18.5-1+cuda12.2 \
+        libcudnn8-dev=8.9.4.25-1+cuda12.2 \
+        libcudnn8=8.9.4.25-1+cuda12.2 \
+        libnvinfer-plugin8=8.6.1.6-1+cuda12.0 \
+        libnvinfer8=8.6.1.6-1+cuda12.0 \
+        libnvinfer-dev=8.6.1.6-1+cuda12.0 \
+        libnvinfer-plugin-dev=8.6.1.6-1+cuda12.0 \
+        automake \
+        build-essential \
+        ca-certificates \
         libcufft-dev-12-2 \
         libcurand-dev-12-2 \
         libcusolver-dev-12-2 \
@@ -57,8 +70,6 @@ RUN apt-get install -y --no-install-recommends \
         git \
         libfreetype6-dev \
         libtool \
-        libcudnn8=${CUDNN_VERSION}-1+cuda${CUDA} \
-        libcudnn8-dev=${CUDNN_VERSION}-1+cuda${CUDA} \
         libcurl3-dev \
         libzmq3-dev \
         mlocate \
@@ -76,11 +87,6 @@ RUN apt-get install -y --no-install-recommends \
         python-distutils-extra && \
     find /usr/local/cuda-12.2/lib64/ -type f -name 'lib*_static.a' -not -name 'libcudart_static.a' -delete
 
-RUN apt-get install -y --no-install-recommends \
-        libnvinfer-plugin8=8.6.1.6-1+cuda12.0 \
-        libnvinfer8=8.6.1.6-1+cuda12.0 \
-        libnvinfer-dev=8.6.1.6-1+cuda12.0 \
-        libnvinfer-plugin-dev=8.6.1.6-1+cuda12.0
 
 # Install python 3.9.
 RUN add-apt-repository ppa:deadsnakes/ppa && \