From b7fd5dcb4a7ab0c573c17aaad0712b4e4c80649d Mon Sep 17 00:00:00 2001 From: Mert <101130780+mertalev@users.noreply.github.com> Date: Fri, 1 Sep 2023 21:59:17 -0400 Subject: [PATCH] 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 | 2 +- machine-learning/load_test.sh | 24 ---------- machine-learning/locustfile.py | 84 +++++++++++++++++++++++++--------- machine-learning/start.sh | 2 +- 4 files changed, 64 insertions(+), 48 deletions(-) delete mode 100755 machine-learning/load_test.sh diff --git a/docker/docker-compose.dev.yml b/docker/docker-compose.dev.yml index 14c5238d0d..a77253cee9 100644 --- a/docker/docker-compose.dev.yml +++ b/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 diff --git a/machine-learning/load_test.sh b/machine-learning/load_test.sh deleted file mode 100755 index 2aaf6c4151..0000000000 --- a/machine-learning/load_test.sh +++ /dev/null @@ -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 \ No newline at end of file diff --git a/machine-learning/locustfile.py b/machine-learning/locustfile.py index c9fbae36dc..0bf2a37d7f 100644 --- a/machine-learning/locustfile.py +++ b/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) diff --git a/machine-learning/start.sh b/machine-learning/start.sh index b6b7616519..36c2b86259 100755 --- a/machine-learning/start.sh +++ b/machine-learning/start.sh @@ -10,4 +10,4 @@ gunicorn app.main:app \ -k uvicorn.workers.UvicornWorker \ -w $MACHINE_LEARNING_WORKERS \ -b $MACHINE_LEARNING_HOST:$MACHINE_LEARNING_PORT \ - --log-config-json log_conf.json \ No newline at end of file + --log-config-json log_conf.json