Simplify tensorrt (#16835)

* Remove unneccessary trt wheels build * Cleanup * Try without local cuda * Keep specific cuda libs only * Cleanup * Add newer libcufft * remove target * Include more
blakeblackshear · Feb 26, 2025 · 4f855f8 · 4f855f8
1 parent d0e9bcb
commit 4f855f8
Show file tree

Hide file tree

Showing 5 changed files with 15 additions and 16 deletions.
diff --git a/docker/tensorrt/Dockerfile.amd64 b/docker/tensorrt/Dockerfile.amd64
@@ -3,22 +3,16 @@
 # https://askubuntu.com/questions/972516/debian-frontend-environment-variable
 ARG DEBIAN_FRONTEND=noninteractive
 
-# Make this a separate target so it can be built/cached optionally
-FROM wheels as trt-wheels
-ARG DEBIAN_FRONTEND
-ARG TARGETARCH
-RUN python3 -m pip config set global.break-system-packages true
-
-# Add TensorRT wheels to another folder
-COPY docker/tensorrt/requirements-amd64.txt /requirements-tensorrt.txt
-RUN mkdir -p /trt-wheels && pip3 wheel --wheel-dir=/trt-wheels -r /requirements-tensorrt.txt
+# Globally set pip break-system-packages option to avoid having to specify it every time
+ARG PIP_BREAK_SYSTEM_PACKAGES=1
 
 FROM tensorrt-base AS frigate-tensorrt
+ARG PIP_BREAK_SYSTEM_PACKAGES
 ENV TRT_VER=8.6.1
-RUN python3 -m pip config set global.break-system-packages true
-RUN --mount=type=bind,from=trt-wheels,source=/trt-wheels,target=/deps/trt-wheels \
-    pip3 install -U /deps/trt-wheels/*.whl && \
-    ldconfig
+
+# Install TensorRT wheels
+COPY docker/tensorrt/requirements-amd64.txt /requirements-tensorrt.txt
+RUN pip3 install -U -r /requirements-tensorrt.txt && ldconfig
 
 WORKDIR /opt/frigate/
 COPY --from=rootfs / /

diff --git a/docker/tensorrt/Dockerfile.base b/docker/tensorrt/Dockerfile.base
@@ -22,9 +22,14 @@ FROM deps AS tensorrt-base
 #Disable S6 Global timeout
 ENV S6_CMD_WAIT_FOR_SERVICES_MAXTIME=0
 
+# COPY TensorRT Model Generation Deps
 COPY --from=trt-deps /usr/local/lib/libyolo_layer.so /usr/local/lib/libyolo_layer.so
 COPY --from=trt-deps /usr/local/src/tensorrt_demos /usr/local/src/tensorrt_demos
-COPY --from=trt-deps /usr/local/cuda-12.* /usr/local/cuda
+
+# COPY Individual CUDA deps
+COPY --from=trt-deps /usr/local/cuda-12.3/targets/x86_64-linux/lib/libcurand.so.* /usr/local/cuda/
+COPY --from=trt-deps /usr/local/cuda-12.3/targets/x86_64-linux/lib/libnvrtc.so.* /usr/local/cuda/
+
 COPY docker/tensorrt/detector/rootfs/ /
 ENV YOLO_MODELS=""
 

diff --git a/docker/tensorrt/detector/rootfs/etc/ld.so.conf.d/cuda_tensorrt.conf b/docker/tensorrt/detector/rootfs/etc/ld.so.conf.d/cuda_tensorrt.conf
@@ -1,5 +1,5 @@
 /usr/local/lib
-/usr/local/cuda/lib64
+/usr/local/cuda
 /usr/local/lib/python3.11/dist-packages/nvidia/cudnn/lib
 /usr/local/lib/python3.11/dist-packages/nvidia/cuda_runtime/lib
 /usr/local/lib/python3.11/dist-packages/nvidia/cublas/lib

diff --git a/docker/tensorrt/requirements-amd64.txt b/docker/tensorrt/requirements-amd64.txt
@@ -11,6 +11,7 @@ nvidia-cublas-cu11 == 11.11.3.6; platform_machine == 'x86_64'
 nvidia-cudnn-cu11 == 8.6.0.*; platform_machine == 'x86_64'
 nvidia-cudnn-cu12 == 9.5.0.*; platform_machine == 'x86_64'
 nvidia-cufft-cu11==10.*; platform_machine == 'x86_64'
+nvidia-cufft-cu12==11.*; platform_machine == 'x86_64'
 onnx==1.16.*; platform_machine == 'x86_64'
 onnxruntime-gpu==1.20.*; platform_machine == 'x86_64'
 protobuf==3.20.3; platform_machine == 'x86_64'
diff --git a/docker/tensorrt/trt.hcl b/docker/tensorrt/trt.hcl
@@ -95,7 +95,6 @@ target "tensorrt" {
     wget = "target:wget",
     tensorrt-base = "target:tensorrt-base",
     rootfs = "target:rootfs"
-    wheels = "target:wheels"
   }
   target = "frigate-tensorrt"
   inherits = ["_build_args"]