Match packages in devel and non-devel GPU with TF dockerfile.

PiperOrigin-RevId: 454022961
tensorflow · Jun 9, 2022 · a8ffec7 · a8ffec7
1 parent 3cda321
commit a8ffec7
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 28 deletions.
diff --git a/tensorflow_serving/tools/docker/Dockerfile.devel-gpu b/tensorflow_serving/tools/docker/Dockerfile.devel-gpu
@@ -60,28 +60,25 @@ RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/
         zip \
         zlib1g-dev \
         python3-distutils \
-        python-distutils-extra \
-        && \
-    apt-get clean && \
-    rm -rf /var/lib/apt/lists/* && \
+        python-distutils-extra && \
     find /usr/local/cuda-11.2/lib64/ -type f -name 'lib*_static.a' -not -name 'libcudart_static.a' -delete && \
     rm /usr/lib/x86_64-linux-gnu/libcudnn_static_v8.a
 
 # NOTE: libnvinfer uses cuda11.1 versions
-RUN apt-get update && \
-        apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1804/x86_64/7fa2af80.pub && \
-        echo "deb https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1804/x86_64 /"  > /etc/apt/sources.list.d/tensorRT.list && \
-        apt-get install -y --no-install-recommends libnvinfer7=${TF_TENSORRT_VERSION}-1+cuda11.0 \
-        libnvinfer-dev=${TF_TENSORRT_VERSION}-1+cuda11.0 \
-        libnvinfer-plugin-dev=${TF_TENSORRT_VERSION}-1+cuda11.0 \
-        libnvinfer-plugin7=${TF_TENSORRT_VERSION}-1+cuda11.0 \
-        # TODO: need to verify
-        libnvonnxparsers7=${TF_TENSORRT_VERSION}-1+cuda11.0 \
-        libnvparsers7=${TF_TENSORRT_VERSION}-1+cuda11.0\
-        libnvonnxparsers-dev=${TF_TENSORRT_VERSION}-1+cuda11.0 \
-        libnvparsers-dev=${TF_TENSORRT_VERSION}-1+cuda11.0 \
-        && apt-get clean \
-        && rm -rf /var/lib/apt/lists/*;
+RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1804/x86_64/7fa2af80.pub && \
+    echo "deb https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1804/x86_64 /"  > /etc/apt/sources.list.d/tensorRT.list && \
+    apt-get update && \
+    apt-get install -y --no-install-recommends libnvinfer7=${TF_TENSORRT_VERSION}-1+cuda11.0 \
+      libnvinfer-dev=${TF_TENSORRT_VERSION}-1+cuda11.0 \
+      libnvinfer-plugin-dev=${TF_TENSORRT_VERSION}-1+cuda11.0 \
+      libnvinfer-plugin7=${TF_TENSORRT_VERSION}-1+cuda11.0 \
+      # TODO: need to verify
+      libnvonnxparsers7=${TF_TENSORRT_VERSION}-1+cuda11.0 \
+      libnvparsers7=${TF_TENSORRT_VERSION}-1+cuda11.0\
+      libnvonnxparsers-dev=${TF_TENSORRT_VERSION}-1+cuda11.0 \
+      libnvparsers-dev=${TF_TENSORRT_VERSION}-1+cuda11.0 && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*;
 
 # Install python 3.7.
 RUN add-apt-repository ppa:deadsnakes/ppa && \

diff --git a/tensorflow_serving/tools/docker/Dockerfile.gpu b/tensorflow_serving/tools/docker/Dockerfile.gpu
@@ -34,6 +34,7 @@ RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/
     apt-get update && apt-get install -y --no-install-recommends \
         ca-certificates \
         cuda-command-line-tools-11-2 \
+        cuda-nvrtc-${CUDA/./-} \
         libcublas-11-2 \
         libcublas-dev-11-2 \
         libcufft-11-2 \
@@ -42,18 +43,21 @@ RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/
         libcusparse-11-2 \
         libcudnn8=${CUDNN_VERSION}-1+cuda${CUDA} \
         libgomp1 \
-        && \
-    apt-get clean && \
-    rm -rf /var/lib/apt/lists/*
+        build-essential \
+        curl \
+        libfreetype6-dev \
+        pkg-config \
+        software-properties-common \
+        unzip
 
 # We don't install libnvinfer-dev since we don't need to build against TensorRT
-RUN apt-get update && \
-        apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1804/x86_64/7fa2af80.pub && \
-        echo "deb https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1804/x86_64 /"  > /etc/apt/sources.list.d/tensorRT.list && \
-        apt-get install -y --no-install-recommends libnvinfer7=${TF_TENSORRT_VERSION}-1+cuda11.0 \
-        libnvinfer-plugin7=${TF_TENSORRT_VERSION}-1+cuda11.0 \
-        && apt-get clean \
-        && rm -rf /var/lib/apt/lists/*;
+RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1804/x86_64/7fa2af80.pub && \
+    echo "deb https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1804/x86_64 /"  > /etc/apt/sources.list.d/tensorRT.list && \
+    apt-get update && \
+    apt-get install -y --no-install-recommends libnvinfer7=${TF_TENSORRT_VERSION}-1+cuda11.0 \
+      libnvinfer-plugin7=${TF_TENSORRT_VERSION}-1+cuda11.0 \
+    && apt-get clean && \
+    rm -rf /var/lib/apt/lists/*;
 
 # Install TF Serving GPU pkg
 COPY --from=build_image /usr/local/bin/tensorflow_model_server /usr/bin/tensorflow_model_server