Repository: alibaba/BladeDISC (language: python)

Issue:
TensorRT static linkage is not supported.



Discussion / hints:
tensorflow_blade still links to cuDNN shared libs.

Fix patch:
diff --git a/docker/dev/Dockerfile b/docker/dev/Dockerfile
index 6b6e3dbc7b4..84cf402f8c1 100644
--- a/docker/dev/Dockerfile
+++ b/docker/dev/Dockerfile
@@ -16,7 +16,7 @@ RUN bash /install/scripts/find-fastest-apt.sh && \
 
 ARG DEVICE=cu110
 RUN bash /install/scripts/install-tensorrt.sh "${DEVICE}"
-RUn bash /install/scripts/patch-cu102.sh "${DEVICE}"
+RUN bash /install/scripts/patch-cuda.sh "${DEVICE}"
 
 ENV PATH="/opt/cmake/bin:${PATH}"
 ENV LD_LIBRARY_PATH="/usr/local/TensorRT/lib/:${LD_LIBRARY_PATH}"
\ No newline at end of file
diff --git a/docker/dev/Dockerfile.torch-disc b/docker/dev/Dockerfile.torch-disc
index 9ca6ec5f795..c2e6f616f76 100644
--- a/docker/dev/Dockerfile.torch-disc
+++ b/docker/dev/Dockerfile.torch-disc
@@ -3,7 +3,7 @@ ARG PYTORCH_COMMIT=""
 ENV PYTORCH_COMMIT=${PYTORCH_COMMIT}
 
 COPY ./scripts /opt/scripts
-RUN apt-get update -y && \
+RUN bash /opt/scripts/find-fastest-apt.sh && \
     apt-get install -y clang-8 clang++-8 python3.8 python3.8-dev && \
     rm /usr/bin/python && \
     ln -s /usr/bin/python3.8 /usr/bin/python && \
diff --git a/docker/runtime/Dockerfile.pytorch b/docker/runtime/Dockerfile.pytorch
index 3b2215daa61..0e42593bd01 100644
--- a/docker/runtime/Dockerfile.pytorch
+++ b/docker/runtime/Dockerfile.pytorch
@@ -6,7 +6,7 @@ ADD ./docker/scripts /install/scripts
 RUN bash /install/scripts/find-fastest-apt.sh && \
         apt-get install -y -qq wget curl unzip
 RUN bash /install/scripts/install-tensorrt.sh "${DEVICE}"
-RUn bash /install/scripts/patch-cu102.sh "${DEVICE}"
+RUN bash /install/scripts/patch-cuda.sh "${DEVICE}"
 
 ADD ./build/torch_blade*.whl  /install/python/
 
diff --git a/docker/scripts/patch-cu102.sh b/docker/scripts/patch-cuda.sh
similarity index 55%
rename from docker/scripts/patch-cu102.sh
rename to docker/scripts/patch-cuda.sh
index 369fca0c1be..6281855ca6d 100644
--- a/docker/scripts/patch-cu102.sh
+++ b/docker/scripts/patch-cuda.sh
@@ -11,28 +11,34 @@
 # limitations under the License.
 
 # This script copies headers and libraries of cuDNN from /usr to /usr/local/cuda in
-# nvidia:cuda-10.2-xxx images. This is needed because tensorflow assumes that.
+# nvidia:cuda-10.x images. This is needed because tensorflow assumes that.
 
-[[ $1 != "cu102" ]] && { echo "Skip pathcing for non-cu102 images."; exit 0; }
+device=$1
+echo "Patching ${device} ..."
 
+# Copy cudnn headers and lib for cuda 
 headers=$(ls /usr/include/*.h | grep -E 'cublas|cudnn|nvblas')
 for hdr in ${headers[@]}; do
     echo "Copy ${hdr} to cuda home."
-	cp ${hdr} /usr/local/cuda/include/
+    cp ${hdr} /usr/local/cuda/include/
 done
 
 libs=$(ls /usr/lib/x86_64-linux-gnu/ | grep -E 'cublas|cudnn|nvblas' | grep -E '\.so|\.a')
 for lib in ${libs[@]}; do
     echo "Copy /usr/lib/x86_64-linux-gnu/${lib} to cuda home."
-	cp /usr/lib/x86_64-linux-gnu/${lib} /usr/local/cuda/lib64/
+    cp /usr/lib/x86_64-linux-gnu/${lib} /usr/local/cuda/lib64/
 done
 
-# patch 1
-curl -sL https://developer.download.nvidia.com/compute/cuda/10.2/Prod/patches/1/cuda_10.2.1_linux.run -o cuda_10.2.1_linux.run
-bash ./cuda_10.2.1_linux.run --silent --toolkit
-rm -f ./cuda_10.2.1_linux.run
 
-# patch 2
-curl -sL https://developer.download.nvidia.com/compute/cuda/10.2/Prod/patches/2/cuda_10.2.2_linux.run -o cuda_10.2.2_linux.run
-bash ./cuda_10.2.2_linux.run --silent --toolkit
-rm -f ./cuda_10.2.2_linux.run
\ No newline at end of file
+if [[ ${device} == "cu102" ]]; then 
+    echo "Installing cuda 10.2 patches ..."
+    # patch 1
+    curl -sL https://developer.download.nvidia.com/compute/cuda/10.2/Prod/patches/1/cuda_10.2.1_linux.run -o cuda_10.2.1_linux.run
+    bash ./cuda_10.2.1_linux.run --silent --toolkit
+    rm -f ./cuda_10.2.1_linux.run
+
+    # patch 2
+    curl -sL https://developer.download.nvidia.com/compute/cuda/10.2/Prod/patches/2/cuda_10.2.2_linux.run -o cuda_10.2.2_linux.run
+    bash ./cuda_10.2.2_linux.run --silent --toolkit
+    rm -f ./cuda_10.2.2_linux.run
+fi
\ No newline at end of file
diff --git a/tensorflow_blade/src/tensorrt/BUILD b/tensorflow_blade/src/tensorrt/BUILD
index 0eb60fd4c1c..1490b6ea7aa 100644
--- a/tensorflow_blade/src/tensorrt/BUILD
+++ b/tensorflow_blade/src/tensorrt/BUILD
@@ -1,5 +1,5 @@
 load("@local_config_tf//:build_defs.bzl", "tf_copts")
-load("@local_config_cuda_supplement//:build_defs.bzl", "if_has_cublaslt")
+load("@local_config_cuda_supplement//:build_defs.bzl", "if_has_cublaslt", "if_has_cudnn_static")
 load("@local_config_trt//:build_defs.bzl", "if_has_myelin")
 
 package(default_visibility = ["//visibility:public"])
@@ -16,17 +16,17 @@ cc_library(
     ],
     copts = tf_copts(),
     linkstatic = 1,  # this target only allow static linking, no .so file generated.
+    visibility = ["//visibility:private"],
     deps = [
-        "//src/util:tf_allocator_util",
         "//src/util:logging",
+        "//src/util:tf_allocator_util",
         "@local_config_cuda//cuda:cuda_headers",
-        "@local_config_trt//:nvinfer_static",
-        "@local_config_trt//:nvinferplugin_static",
         "@local_config_tf//:libtensorflow_framework",
         "@local_config_tf//:tf_header_lib",
+        "@local_config_trt//:nvinfer_static",
+        "@local_config_trt//:nvinferplugin_static",
     ],
     alwayslink = 1,  # targets depending on it should carry all symbols in its children.
-    visibility = ["//visibility:private"]
 )
 
 cc_library(
@@ -53,10 +53,22 @@ cc_library(
         ":tensorrt_tf_impl",
         "@local_config_trt//:nvonnxparser_static",
         "@local_config_cuda//cuda:cudart_static",
-        "@local_config_cuda//cuda:cudnn",
-        "@local_config_cuda//cuda:cublas",
+        "@local_config_cuda_supplement//:cublas_static",
+        "@local_config_cuda_supplement//:culibos_static",
         "//src/util:logging",
-    ] + if_has_cublaslt([
+    ] + if_has_cudnn_static(
+        [
+            "@local_config_cuda_supplement//:cudnn_static",
+        ],
+        if_false = [
+            # cuda 10.2 has no libcudnn_static.a
+            "@local_config_cuda_supplement//:cudnn_adv_infer_static",
+            "@local_config_cuda_supplement//:cudnn_cnn_infer_static",
+            "@local_config_cuda_supplement//:cudnn_cnn_train_static",
+            "@local_config_cuda_supplement//:cudnn_ops_infer_static",
+            "@local_config_cuda_supplement//:cudnn_ops_train_static",
+        ],
+    ) + if_has_cublaslt([
         "@local_config_cuda_supplement//:cublasLt_static",
     ]) + if_has_myelin([
         "@local_config_cuda_supplement//:nvrtc",


Test patch:
diff --git a/docker/scripts/find-fastest-apt.sh b/docker/scripts/find-fastest-apt.sh
index 2a4483ef209..e64070032f3 100755
--- a/docker/scripts/find-fastest-apt.sh
+++ b/docker/scripts/find-fastest-apt.sh
@@ -11,7 +11,6 @@
 # limitations under the License.
 
 
-
 # Returns the domain name of an URL.
 function get_domain_from_url() {
     # cut extract the domain, domain is the third field split by '/'
@@ -79,7 +78,14 @@ deb-src $fastest $codename-backports main restricted universe multiverse
 EOF
 }
 
+set -xe
+
+# Note(xiafei.qiuxf): Since we never install CUDA-related packages via apt-get, but
+#                     actually from .tar.gz file on OSS, just remove cuda's apt repo to
+#                     get free of its GPG key rotation issue.
+rm -f /etc/apt/sources.list.d/cuda.list /etc/apt/sources.list.d/nvidia-ml.list
+
 apt-get -qq update
 apt-get install -y iputils-ping bc
 find_fastest_apt_source > /etc/apt/sources.list
-apt-get -qq update
\ No newline at end of file
+apt-get -qq update