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

* fixed dev docker compose

* updated locustfile

* deleted old script, moved comments to locustfile

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):
-        self.client.post(
-            "/image-classifier/tag-image", data=self.data, headers=self.headers
-        )
+    def classify(self) -> None:
+        data = [
+            ("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):
-        self.client.post(
-            "/sentence-transformer/encode-image",
-            data=self.data,
-            headers=self.headers,
-        )
+    def encode_text(self) -> None:
+        data = [
+            ("modelName", self.environment.parsed_options.clip_model),
+            ("modelType", "clip"),
+            ("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):
-        self.client.post(
-            "/facial-recognition/detect-faces",
-            data=self.data,
-            headers=self.headers,
-        )
+    def encode_image(self) -> None:
+        data = [
+            ("modelName", self.environment.parsed_options.clip_model),
+            ("modelType", "clip"),
+            ("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)