Load Testing ML Model using Locust

Locust is a load testing tool to check the behavior of the system. The aim of load testing is to know how the system responds given the number of loads.

Photo by James Wainscoat on Unsplash

In this experiment, I built a simple sentiment classification model using Random Forest and Locust as the load testing tool.

I used the google play review from here and build the model using scikit-learn. The extract the text using TF-IDF and fit the data into RF model.

Classification report from the model using test data

After successfully build the model and dump into joblib. Now, we need to create an API using Flask.

@app.route("/predict", methods=["GET", "POST"])
def predict():
    if request.method == "GET":
        return jsonify({"message": "Please use POST method to predict sentiment from a review"})
    else:
        content = request.json
        logging.debug("REVIEW: " + content["review"])
        # label = sentclf(content["review"])

        review_mat = tfidf.transform([content['review']])
        labels = classifier.predict(review_mat)
        label = labels[0]
        
        return jsonify(
            {
                "label": label
            }
        )

Using the lines of code above, we are able to use the API to get the inference of the model.

Next, we need to setup locust script for load testing.

Import the modules needed for this testing.

import random
import time
from locust import HttpUser, task, tag, constant_throughput

let’s define the class and constant_throughput. throughput means that in every second, we only send a request once. In addition, we create predict() function for testing the url. decorator tag ‘model’ here is to differentiate with other functions, if you look at full script, it also test another url. Therefore, in this testing, we only focus on /predict .

We then set reviews for the script to randomly choose which review wants to predict and later send request post. the objective of time.sleep(2) is used for after the request send, it will wait for 2 seconds for sending new request. The time I was writing this test is to take a closer look on response time, so it is not suddently raise the RPS.

class HelloWorldUser(HttpUser):

    wait_time = constant_throughput(1)

    @tag("model")
    @task
    def predict(self):

        reviews = [
            "Perbanyak lagi item penjualannya, soale masih kalah jauh dengan olshop tetangga",
            "sebenarnya sangat mudah pengoperasiannya, tp untuk pemula mungkin agak bingung",
            "kerjasama promo dgn vendornya lbh variatif lagi dong, intip2 ecommerce sebelah",
            "ok bgt.. gratis ongkir.. cuma barang blm begitu banyak, blom begitu bervariasi",
            "aplikasi shoping trhancur yng pernah sya install... verifikasi hp bloon banget",
        ]

        self.client.post("/predict", json={"review": random.choice(reviews)}, timeout=1)

        time.sleep(2)

Finally, we run the app using gunicorn

gunicorn --workers=4 server:app

[2022-12-05 16:26:05 +0700] [1027] [INFO] Starting gunicorn 20.1.0
[2022-12-05 16:26:05 +0700] [1027] [INFO] Listening at: http://127.0.0.1:8000 (1027)
[2022-12-05 16:26:05 +0700] [1027] [INFO] Using worker: sync
[2022-12-05 16:26:05 +0700] [1028] [INFO] Booting worker with pid: 1028
[2022-12-05 16:26:05 +0700] [1029] [INFO] Booting worker with pid: 1029
[2022-12-05 16:26:05 +0700] [1030] [INFO] Booting worker with pid: 1030
[2022-12-05 16:26:05 +0700] [1036] [INFO] Booting worker with pid: 1036
^C[2022-12-05 16:26:07 +0700] [1027] [INFO] Handling signal: int
[2022-12-05 16:26:07 +0700] [1036] [INFO] Worker exiting (pid: 1036)
[2022-12-05 16:26:07 +0700] [1029] [INFO] Worker exiting (pid: 1029)
[2022-12-05 16:26:07 +0700] [1030] [INFO] Worker exiting (pid: 1030)
[2022-12-05 16:26:07 +0700] [1028] [INFO] Worker exiting (pid: 1028)

Then, run locust. tags model is to ensure that when we run testing, we only test the function that have tags ‘model’

locust --tags model

Then after we run several tests, the result will be something like this.

I used fly.io to deploy the model, you can look up the configurations on the branch flyio

Conclusion

To look into the details, you can clone this repository. If you have any question, reach out to JAIR discord [Invitation Link]