dev(ml): fixed docker-compose.dev.yml, updated locust (#3951)

* fixed dev docker compose

* updated locustfile

* deleted old script, moved comments to locustfile

docker/docker-compose.dev.yml

@@ -34,7 +34,7 @@ services:
     ports:
       - 3003:3003
     volumes:
-      - ../machine-learning/app:/usr/src/app
+      - ../machine-learning:/usr/src/app
       - model-cache:/cache
     env_file:
       - .env

machine-learning/load_test.sh

@@ -1,24 +0,0 @@
-export MACHINE_LEARNING_CACHE_FOLDER=/tmp/model_cache
-export MACHINE_LEARNING_MIN_FACE_SCORE=0.034 # returns 1 face per request; setting this to 0 blows up the number of faces to the thousands
-export MACHINE_LEARNING_MIN_TAG_SCORE=0.0
-export PID_FILE=/tmp/locust_pid
-export LOG_FILE=/tmp/gunicorn.log
-export HEADLESS=false
-export HOST=127.0.0.1:3003
-export CONCURRENCY=4
-export NUM_ENDPOINTS=3
-export PYTHONPATH=app
-
-gunicorn app.main:app --worker-class uvicorn.workers.UvicornWorker \
-    --bind $HOST --daemon --error-logfile $LOG_FILE --pid $PID_FILE
-while true ; do
-    echo "Loading models..."
-    sleep 5
-    if cat $LOG_FILE | grep -q -E "startup complete"; then break; fi
-done
-
-# "users" are assigned only one task, so multiply concurrency by the number of tasks
-locust --host http://$HOST --web-host 127.0.0.1 \
-    --run-time 120s --users $(($CONCURRENCY * $NUM_ENDPOINTS)) $(if $HEADLESS; then echo "--headless"; fi)
-
-if [[ -e $PID_FILE ]]; then kill $(cat $PID_FILE); fi

machine-learning/locustfile.py

@@ -1,13 +1,32 @@
 from io import BytesIO
+import json
+from typing import Any
 
 from locust import HttpUser, events, task
+from locust.env import Environment
 from PIL import Image
+from argparse import ArgumentParser
+byte_image = BytesIO()
+
+
+@events.init_command_line_parser.add_listener
+def _(parser: ArgumentParser) -> None:
+    parser.add_argument("--tag-model", type=str, default="microsoft/resnet-50")
+    parser.add_argument("--clip-model", type=str, default="ViT-B-32::openai")
+    parser.add_argument("--face-model", type=str, default="buffalo_l")
+    parser.add_argument("--tag-min-score", type=int, default=0.0, 
+                        help="Returns all tags at or above this score. The default returns all tags.")
+    parser.add_argument("--face-min-score", type=int, default=0.034, 
+                        help=("Returns all faces at or above this score. The default returns 1 face per request; "
+                              "setting this to 0 blows up the number of faces to the thousands."))
+    parser.add_argument("--image-size", type=int, default=1000)
 
 
 @events.test_start.add_listener
-def on_test_start(environment, **kwargs):
+def on_test_start(environment: Environment, **kwargs: Any) -> None:
     global byte_image
-    image = Image.new("RGB", (1000, 1000))
+    assert environment.parsed_options is not None
+    image = Image.new("RGB", (environment.parsed_options.image_size, environment.parsed_options.image_size))
     byte_image = BytesIO()
     image.save(byte_image, format="jpeg")
 
@@ -19,34 +38,55 @@ class InferenceLoadTest(HttpUser):
     headers: dict[str, str] = {"Content-Type": "image/jpg"}
 
     # re-use the image across all instances in a process
-    def on_start(self):
+    def on_start(self) -> None:
         global byte_image
         self.data = byte_image.getvalue()
 
 
-class ClassificationLoadTest(InferenceLoadTest):
+class ClassificationFormDataLoadTest(InferenceLoadTest):
     @task
-    def classify(self):
+    def classify(self) -> None:
-        self.client.post(
+        data = [
-            "/image-classifier/tag-image", data=self.data, headers=self.headers
+            ("modelName", self.environment.parsed_options.clip_model),
-        )
+            ("modelType", "clip"),
+            ("options", json.dumps({"minScore": self.environment.parsed_options.tag_min_score})),
+        ]
+        files = {"image": self.data}
+        self.client.post("/predict", data=data, files=files)
 
 
-class CLIPLoadTest(InferenceLoadTest):
+class CLIPTextFormDataLoadTest(InferenceLoadTest):
     @task
-    def encode_image(self):
+    def encode_text(self) -> None:
-        self.client.post(
+        data = [
-            "/sentence-transformer/encode-image",
+            ("modelName", self.environment.parsed_options.clip_model),
-            data=self.data,
+            ("modelType", "clip"),
-            headers=self.headers,
+            ("options", json.dumps({"mode": "text"})),
-        )
+            ("text", "test search query")
+        ]
+        self.client.post("/predict", data=data)
 
 
-class RecognitionLoadTest(InferenceLoadTest):
+class CLIPVisionFormDataLoadTest(InferenceLoadTest):
     @task
-    def recognize(self):
+    def encode_image(self) -> None:
-        self.client.post(
+        data = [
-            "/facial-recognition/detect-faces",
+            ("modelName", self.environment.parsed_options.clip_model),
-            data=self.data,
+            ("modelType", "clip"),
-            headers=self.headers,
+            ("options", json.dumps({"mode": "vision"})),
-        )
+        ]
+        files = {"image": self.data}
+        self.client.post("/predict", data=data, files=files)
+
+
+class RecognitionFormDataLoadTest(InferenceLoadTest):
+    @task
+    def recognize(self) -> None:
+        data = [
+            ("modelName", self.environment.parsed_options.face_model),
+            ("modelType", "facial-recognition"),
+            ("options", json.dumps({"minScore": self.environment.parsed_options.face_min_score})),
+        ]
+        files = {"image": self.data}
+            
+        self.client.post("/predict", data=data, files=files)