From 0d6bb6714a94b422674dcdc56ef3f7f2b7f83f2a Mon Sep 17 00:00:00 2001
From: Nate Meyer <Nate.Devel@gmail.com>
Date: Thu, 21 Sep 2023 06:23:51 -0400
Subject: [PATCH] Add support for selecting a specific GPU to use when
 converting TRT models (#7857)

---
 .../s6-overlay/s6-rc.d/trt-model-prepare/run   | 18 ++++++++++++++++++
 docs/docs/configuration/object_detectors.md    |  8 ++++++++
 2 files changed, 26 insertions(+)

diff --git a/docker/tensorrt/detector/rootfs/etc/s6-overlay/s6-rc.d/trt-model-prepare/run b/docker/tensorrt/detector/rootfs/etc/s6-overlay/s6-rc.d/trt-model-prepare/run
index 93b132a38..c39c7a0aa 100755
--- a/docker/tensorrt/detector/rootfs/etc/s6-overlay/s6-rc.d/trt-model-prepare/run
+++ b/docker/tensorrt/detector/rootfs/etc/s6-overlay/s6-rc.d/trt-model-prepare/run
@@ -43,6 +43,15 @@ if [[ -z ${MODEL_CONVERT} ]]; then
     exit 0
 fi
 
+# Setup ENV to select GPU for conversion
+if [ ! -z ${TRT_MODEL_PREP_DEVICE+x} ]; then
+    if [ ! -z ${CUDA_VISIBLE_DEVICES+x} ]; then
+        PREVIOUS_CVD="$CUDA_VISIBLE_DEVICES"
+        unset CUDA_VISIBLE_DEVICES
+    fi
+    export CUDA_VISIBLE_DEVICES="$TRT_MODEL_PREP_DEVICE"
+fi
+
 # On Jetpack 4.6, the nvidia container runtime will mount several host nvidia libraries into the
 # container which should not be present in the image - if they are, TRT model generation will
 # fail or produce invalid models. Thus we must request the user to install them on the host in
@@ -87,5 +96,14 @@ do
     echo "Generated ${model}.trt in $(($(date +%s)-start)) seconds"
 done
 
+# Restore ENV after conversion
+if [ ! -z ${TRT_MODEL_PREP_DEVICE+x} ]; then
+    unset CUDA_VISIBLE_DEVICES
+    if [ ! -z ${PREVIOUS_CVD+x} ]; then
+        export CUDA_VISIBLE_DEVICES="$PREVIOUS_CVD"
+    fi
+fi
+
+# Print which models exist in output folder
 echo "Available tensorrt models:"
 cd ${OUTPUT_FOLDER} && ls *.trt;
diff --git a/docs/docs/configuration/object_detectors.md b/docs/docs/configuration/object_detectors.md
index 2c827469b..e1847ead4 100644
--- a/docs/docs/configuration/object_detectors.md
+++ b/docs/docs/configuration/object_detectors.md
@@ -239,6 +239,14 @@ frigate:
     - USE_FP16=false
 ```
 
+If you have multiple GPUs passed through to Frigate, you can specify which one to use for the model conversion.  The conversion script will use the first visible GPU, however in systems with mixed GPU models you may not want to use the default index for object detection.  Add the `TRT_MODEL_PREP_DEVICE` environment variable to select a specific GPU.
+
+```yml
+frigate:
+  environment:
+    - TRT_MODEL_PREP_DEVICE=0 # Optionally, select which GPU is used for  model optimization
+```
+
 ### Configuration Parameters
 
 The TensorRT detector can be selected by specifying `tensorrt` as the model type. The GPU will need to be passed through to the docker container using the same methods described in the [Hardware Acceleration](hardware_acceleration.md#nvidia-gpu) section. If you pass through multiple GPUs, you can select which GPU is used for a detector with the `device` configuration parameter. The `device` parameter is an integer value of the GPU index, as shown by `nvidia-smi` within the container.