From d270f3eae095285361f4b22381c246d29c26748f Mon Sep 17 00:00:00 2001
From: Josh Bailey <josh@vandervecken.com>
Date: Thu, 12 Oct 2023 07:19:51 +0000
Subject: [PATCH] Add yolo integration test.

---
 .github/workflows/docker-test.yml |  2 ++
 docker/Dockerfile.cuda-torchserve |  1 +
 docker/Dockerfile.torchserve      | 19 +------------------
 tests/test_torchserve.sh          | 18 ++++++++++++++++++
 4 files changed, 22 insertions(+), 18 deletions(-)
 create mode 100755 tests/test_torchserve.sh

diff --git a/.github/workflows/docker-test.yml b/.github/workflows/docker-test.yml
index ed3e42d9..fdf5d00f 100644
--- a/.github/workflows/docker-test.yml
+++ b/.github/workflows/docker-test.yml
@@ -9,6 +9,8 @@ jobs:
       run: |
         cd docker
         docker build -f Dockerfile.torchserve . -t iqtlabs/gamutrf-torchserve:latest
+        cd ..
+        ./tests/test_torchserve.sh
   test-gamutrf-extra-images:
     runs-on: ubuntu-latest
     steps:
diff --git a/docker/Dockerfile.cuda-torchserve b/docker/Dockerfile.cuda-torchserve
index e77162b2..8419ebb2 100644
--- a/docker/Dockerfile.cuda-torchserve
+++ b/docker/Dockerfile.cuda-torchserve
@@ -10,6 +10,7 @@ RUN pip config set global.no-cache-dir false && \
   cd serve && \
   python3 ./ts_scripts/install_dependencies.py --cuda cu118 --environment prod && \
   pip3 install . && \
+  pip3 install -r examples/object_detector/yolo/yolov8/requirements.txt && \
   cd .. && \
   rm -rf serve
 RUN /usr/local/bin/torchserve --help
diff --git a/docker/Dockerfile.torchserve b/docker/Dockerfile.torchserve
index 79e4400d..931a75a2 100644
--- a/docker/Dockerfile.torchserve
+++ b/docker/Dockerfile.torchserve
@@ -10,24 +10,7 @@ RUN pip config set global.no-cache-dir false && \
   cd serve && \
   python3 ./ts_scripts/install_dependencies.py --environment prod && \
   pip3 install . && \
+  pip3 install -r examples/object_detector/yolo/yolov8/requirements.txt && \
   cd .. && \
   rm -rf serve
 RUN /usr/local/bin/torchserve --help
-
-## example usage with YOLOv8 (https://github.com/pytorch/serve/tree/v0.8.2/examples/object_detector/yolo/yolov8):
-
-# docker build -f Dockerfile.torchserve . -ts iqtlabs/gamutrf-torchserve
-# docker run -ti iqtlabs/gamutrf-torchserve bash
-# apt-get update && apt-get install -y curl wget
-# pip3 install torch-model-archiver
-# cd /tmp
-# git clone https://github.com/pytorch/serve -b v0.8.2
-# cd serve/examples/object_detector/yolo/yolov8
-# pip install -r requirements.txt
-# wget https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8n.pt
-# torch-model-archiver --model-name yolov8n --version 1.0 --serialized-file yolov8n.pt --handler custom_handler.py
-# mkdir model_store
-# mv yolov8n.mar model_store/.
-# torchserve --start --model-store model_store --ncs
-# curl -X POST "localhost:8081/models?model_name=yolov8n&url=yolov8n.mar&initial_workers=4&batch_size=2"
-# curl http://127.0.0.1:8080/predictions/yolov8n -T persons.jpg  & curl http://127.0.0.1:8080/predictions/yolov8n -T bus.jpg
diff --git a/tests/test_torchserve.sh b/tests/test_torchserve.sh
new file mode 100755
index 00000000..9ee81358
--- /dev/null
+++ b/tests/test_torchserve.sh
@@ -0,0 +1,18 @@
+#!/bin/bash
+
+set -e
+TMPDIR=/tmp
+sudo apt-get update && sudo apt-get install -y curl wget
+sudo pip3 install torch-model-archiver
+cd $TMPDIR
+git clone https://github.com/pytorch/serve -b v0.8.2
+cd serve/examples/object_detector/yolo/yolov8
+# TODO: use gamutRF weights here.
+wget https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8n.pt
+torch-model-archiver --force --model-name yolov8n --version 1.0 --serialized-file yolov8n.pt --handler custom_handler.py
+rm -rf model_store && mkdir model_store
+mv yolov8n.mar model_store/
+docker run -v $(pwd)/model_store:/model_store --net host -d iqtlabs/gamutrf-torchserve timeout 60s torchserve --start --model-store /model_store --ncs --foreground
+sleep 5
+curl -X POST "localhost:8081/models?model_name=yolov8n&url=yolov8n.mar&initial_workers=4&batch_size=2"
+curl http://127.0.0.1:8080/predictions/yolov8n -T persons.jpg