From b3ce493a16f588e8946a06d9d77dc0aa3a8d9ce6 Mon Sep 17 00:00:00 2001
From: Nick Mowen <nickmowen213@gmail.com>
Date: Mon, 17 Jul 2023 15:19:13 -0600
Subject: [PATCH] Add wheels target

---
 .github/workflows/ci.yml   |  5 ++---
 docker/main/Dockerfile     | 10 ++++++++++
 docker/tensorrt/Dockerfile | 15 ++++-----------
 docker/tensorrt/trt.hcl    |  9 ++++++++-
 docker/tensorrt/trt.mk     |  6 +++---
 5 files changed, 27 insertions(+), 18 deletions(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 223ff6229..94d644e42 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -5,7 +5,6 @@ on:
     branches:
       - dev
       - master
-      - community-boards
 
 # only run the latest commit to avoid cache overwrites
 concurrency:
@@ -65,7 +64,7 @@ jobs:
           files: docker/rpi/rpi.hcl
           set: |
             rpi.tags=ghcr.io/${{ steps.lowercaseRepo.outputs.lowercase }}:${{ github.ref_name }}-${{ env.SHORT_SHA }}-rpi
-            rpi.cache-from=type=gha
+            *.cache-from=type=gha
       - name: Build and push TensorRT
         uses: docker/bake-action@v3
         with:
@@ -74,4 +73,4 @@ jobs:
           files: docker/tensorrt/trt.hcl
           set: |
             tensorrt.tags=ghcr.io/${{ steps.lowercaseRepo.outputs.lowercase }}:${{ github.ref_name }}-${{ env.SHORT_SHA }}-tensorrt
-            tensorrt.cache-from=type=gha
+            *.cache-from=type=gha
diff --git a/docker/main/Dockerfile b/docker/main/Dockerfile
index 06c54f680..05a45c808 100644
--- a/docker/main/Dockerfile
+++ b/docker/main/Dockerfile
@@ -223,6 +223,16 @@ RUN --mount=type=bind,source=./requirements-dev.txt,target=/workspace/frigate/re
 
 CMD ["sleep", "infinity"]
 
+# Dev Container w/ TRT
+FROM devcontainer AS devcontainer-trt
+
+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 docker/tensorrt/detector/rootfs/ /
+COPY --from=trt-deps /usr/local/lib/libyolo_layer.so /usr/local/lib/libyolo_layer.so
+RUN --mount=type=bind,from=trt-wheels,source=/trt-wheels,target=/deps/trt-wheels \
+    pip3 install -U /deps/trt-wheels/*.whl
+
 
 # Frigate web build
 # This should be architecture agnostic, so speed up the build on multiarch by not using QEMU.
diff --git a/docker/tensorrt/Dockerfile b/docker/tensorrt/Dockerfile
index 8c77a701b..2e53a3d22 100644
--- a/docker/tensorrt/Dockerfile
+++ b/docker/tensorrt/Dockerfile
@@ -10,7 +10,7 @@ ARG TARGETARCH
 
 # Add TensorRT wheels to another folder
 COPY requirements-tensorrt.txt /requirements-tensorrt.txt
-RUN mkdir -p /trt-wheels && pip3 wheel --wheel-dir=/trt-wheels -r requirements-tensorrt.txt
+RUN mkdir -p /trt-wheels && pip3 wheel --wheel-dir=/trt-wheels -r /requirements-tensorrt.txt
 
 # Build TensorRT-specific library
 FROM nvcr.io/nvidia/tensorrt:23.03-py3 AS trt-deps
@@ -19,7 +19,7 @@ RUN --mount=type=bind,source=docker/tensorrt/detector/tensorrt_libyolo.sh,target
     /tensorrt_libyolo.sh
 
 # Frigate w/ TensorRT Support as separate image
-FROM frigate AS frigate-tensorrt
+FROM deps AS frigate-tensorrt
 
 #Disable S6 Global timeout
 ENV S6_CMD_WAIT_FOR_SERVICES_MAXTIME=0
@@ -35,12 +35,5 @@ RUN --mount=type=bind,from=trt-wheels,source=/trt-wheels,target=/deps/trt-wheels
     pip3 install -U /deps/trt-wheels/*.whl && \
     ldconfig
 
-# Dev Container w/ TRT
-FROM devcontainer AS devcontainer-trt
-
-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 docker/tensorrt/detector/rootfs/ /
-COPY --from=trt-deps /usr/local/lib/libyolo_layer.so /usr/local/lib/libyolo_layer.so
-RUN --mount=type=bind,from=trt-wheels,source=/trt-wheels,target=/deps/trt-wheels \
-    pip3 install -U /deps/trt-wheels/*.whl
+WORKDIR /opt/frigate/
+COPY --from=rootfs / /
diff --git a/docker/tensorrt/trt.hcl b/docker/tensorrt/trt.hcl
index db15ba971..409811c3b 100644
--- a/docker/tensorrt/trt.hcl
+++ b/docker/tensorrt/trt.hcl
@@ -1,4 +1,4 @@
-target "deps" {
+target deps {
   dockerfile = "docker/main/Dockerfile"
   platforms = ["linux/amd64"]
 }
@@ -8,11 +8,18 @@ target rootfs {
   platforms = ["linux/amd64"]
 }
 
+target wheels {
+  dockerfile = "docker/main/Dockerfile"
+  platforms = ["linux/amd64"]
+}
+
 target "tensorrt" {
   dockerfile = "docker/tensorrt/Dockerfile"
+  context = "."
   contexts = {
     deps = "target:deps",
     rootfs = "target:rootfs"
+    wheels = "target:wheels"
   }
   platforms = ["linux/amd64"]
 }
\ No newline at end of file
diff --git a/docker/tensorrt/trt.mk b/docker/tensorrt/trt.mk
index caf300511..df6b7e2fd 100644
--- a/docker/tensorrt/trt.mk
+++ b/docker/tensorrt/trt.mk
@@ -1,10 +1,10 @@
 BOARDS += trt
 
 local-trt: version
-	docker buildx bake --file=docker/rpi/bake.hcl --set rpi.tagsfrigate:latest-tensorrt tensorrt
+	docker buildx bake --file=docker/tensorrt/trt.hcl --set tensorrt.tags=frigate:latest-tensorrt tensorrt
 
 build-trt:
-	docker buildx bake --file=docker/rpi/bake.hcl --set rpi.tags=$(IMAGE_REPO):${GITHUB_REF_NAME}-$(COMMIT_HASH)-tensorrt tensorrt
+	docker buildx bake --file=docker/tensorrt/trt.hcl --set tensorrt.tags=$(IMAGE_REPO):${GITHUB_REF_NAME}-$(COMMIT_HASH)-tensorrt tensorrt
 
 push-trt: build-trt
-	docker buildx bake --push --file=docker/rpi/bake.hcl --set rpi.tags=$(IMAGE_REPO):${GITHUB_REF_NAME}-$(COMMIT_HASH)-tensorrt tensorrt
\ No newline at end of file
+	docker buildx bake --push --file=docker/tensorrt/trt.hcl --set tensorrt.tags=$(IMAGE_REPO):${GITHUB_REF_NAME}-$(COMMIT_HASH)-tensorrt tensorrt
\ No newline at end of file