[FSTORE-690] Setup Online Feature Store Benchmark with Locust (logicalclocks#948)

Author: moritzmeister

locust_benchmark/Dockerfile

@@ -0,0 +1,11 @@
+FROM locustio/locust:2.15.1
+
+USER root
+
+RUN apt-get update -y && apt-get -y install gcc python3-dev librdkafka-dev git
+
+USER locust
+WORKDIR /home/locust
+
+COPY requirements.txt .
+RUN pip3 install -r requirements.txt

locust_benchmark/README.md

@@ -0,0 +1,178 @@
+## Online Feature Store Benchmark Suite
+
+This directory provides the tools to quickly bootstrap the process of benchmarking the Hopsworks Online Feature Store.
+
+The benchmark suite is built around [locust](https://locust.io/) which allows you to quickly start so-called user classes to perform requests against a system.
+
+In the feature store context, a user class represents a serving application performing requests against the online feature store, looking up feature vectors.
+
+Locust takes care of the hard part of collecting the results and aggregating them.
+
+### Quickstart: Single process locust
+
+By default, each user class gets spawned in its own Python thread, which means, when you run locust in a single Python process, all user classes will share the same core. This is due to the Python Global Interpreter lock.
+
+This section illustrates how to quickly run locust in a single process, see the next section on how to scale locust with a master and multiple worker processes.
+
+#### Prerequisites
+
+1. Hopsworks cluster
+   1. You can follow the documentation on how to setup a cluster on [managed.hopsworks.ai](https://docs.hopsworks.ai/3.1/setup_installation/aws/cluster_creation/)
+   2. Ideally you create an extra node in the same VPC as the Hopsworks cluster, which you can use to run the benchmark suite from. The easiest way to do this is by using a dedicated [API node](https://docs.hopsworks.ai/3.1/setup_installation/common/rondb/#api-nodes).
+   3. Once the cluster is created you need the following information:
+      - Hopsworks domain name: **[UUID].cloud.hopsworks.ai** which you can find in the *Details* tab of your managed.hopsworks.ai control plane
+      - A project and its name on the hopsworks cluster. For that log into your cluster and create a project.
+      - An API key for authentication. Head to your project and [create an API key for your account](https://docs.hopsworks.ai/3.1/user_guides/projects/api_key/create_api_key/).
+2. Python environment
+   In order to run the benchmark suite, you need a Python environment.
+   For that we will SSH into the created API node. Make sure that [SSH is enabled](https://docs.hopsworks.ai/3.1/setup_installation/common/services/) on your cluster.
+
+   1. On AWS all nodes have public IPs, so you can head to your cloud console to find the IP of the API node and SSH directly:
+      `ssh ubuntu@[ip-of-api-node] -i your-ssh-key.pem`
+      On Azure, you will have to proxy through the Hopsworks head node, you can either use the public IP or the Hopsworks domain name as noted above.
+      On GCP, we don't attach SSH keys, but you can use the GCP CLI to create an ssh session.
+
+   2. API nodes come with Anaconda pre-installed, which we can use to create a python environment and install our dependencies:
+
+      ```bash
+      conda init bash
+      ```
+
+      After this you need to create a new SSH session.
+
+      ```bash
+      conda create -n locust-benchmark python=3.10
+      conda activate locust-benchmark
+      ```
+
+      Clone the feature-store-api repository to retrieve the benchmark suite and install the requirements.
+      Make sure to checkout the release branch matching the Hopsworks version of your cluster.
+
+      ```bash
+      git clone https://github.com/logicalclocks/feature-store-api.git
+      cd feature-store-api
+      git checkout master
+      cd locust_benchmark
+      pip install -r requirements.txt
+      ```
+
+#### Creating feature group for lookups
+
+1. Save the previously created API key in a file named `.api_key`
+
+```bash
+echo "[YOUR KEY]" > .api_key
+```
+
+2. Now we need to configure the test, for that modify the `hopsworks_config.json` template:
+
+```json
+{
+    "host": "[UUID].cloud.hopsworks.ai",
+    "port": 443,
+    "project": "test",
+    "external": false,
+    "rows": 100000,
+    "schema_repetitions": 1,
+    "recreate_feature_group": true,
+    "batch_size": 100
+}
+```
+
+- `host`: Domain name of your Hopsworks cluster.
+- `port`: optional, for managed.hopsworks.ai it should be `443`.
+- `project`: the project to run the benchmark in.
+- `external`: if you are not running the benchmark suite from the same VPC as Hopsworks, set this to `true`.
+- `rows`: Number of rows/unique primary key values in the feature group used for lookup benchmarking.
+- `schema_repetitions`: This controls the number of features for the lookup. One schema repetition will result in 10 features plus primary key. Five repetitions will result in 50 features plus primary key.
+- `recreate_feature_group`: This controls if the previous feature group should be dropped and recreated. Set this to true when rerunning the benchmark with different size of rows or schema repetitions.
+- `batch_size`: This is relevant for the actual benchmark and controls how many feature vectors are looked up in the batch benchmark.
+
+3. Create the feature group
+
+```bash
+python create_feature_group.py
+```
+
+Note, the `recreate_feature_group` only matters if you rerun the `create_feature_group.py` script.
+
+#### Run one process locust benchmark
+
+You are now ready to run the load test:
+```bash
+locust -f locustfile.py --headless -u 4 -r 1 -t 30 -s 1 --html=result.html
+```
+
+Options:
+- `u`: number of users sharing the python process
+- `r`: spawn rate of the users in seconds, it will launch one user every `r` seconds until there are `u` users
+- `t`: total time to run the test for
+- `s`: shutdown timeout, no need to be changed
+- `html`: path for the output file
+
+You can also run only single feature vector or batch lookups by running only the respective user class:
+
+```bash
+locust -f locustfile.py FeatureVectorLookup --headless -u 4 -r 1 -t 30 -s 1 --html=result.html
+
+locust -f locustfile.py FeatureVectorBatchLookup --headless -u 4 -r 1 -t 30 -s 1 --html=result.html
+```
+
+### Distributed Locust Benchmark using Docker Compose
+
+As you will see already with 4 users, the single core tends to be saturated and locust will print a warning.
+
+
+
+```
+[2023-03-09 12:02:11,304] ip-10-0-0-187/WARNING/root: CPU usage above 90%! This may constrain your throughput and may even give inconsistent response time measurements! See https://docs.locust.io/en/stable/running-distributed.html for how to distribute the load over multiple CPU cores or machines
+```
+
+First we need to install docker, for which we will use [the provided convenience script](https://docs.docker.com/engine/install/ubuntu/#install-using-the-convenience-script), however, you can use any other method too:
+```bash
+curl -fsSL https://get.docker.com -o get-docker.sh
+sudo sh get-docker.sh
+```
+
+#### Load the docker image
+
+```bash
+wget https://repo.hops.works/dev/moritz/locust_benchmark/locust-2.15.1-hsfs-3.3.0.dev1-amd64.tgz
+sudo docker load < locust-2.15.0-hsfs-3.2.0.dev1-amd64.tgz
+```
+
+#### Create feature group and configure test
+
+See the above section to setup the `hopsworks_config.json` and create the feature group to perform lookups against.
+
+#### Run multiple locust processes using docker compose
+
+Using docker compose we can now start multiple containers, one dedicated master process, which is responsible for collecting metrics and orchestrating the benchmark, and a variable number of worker processes.
+
+```bash
+sudo docker compose up
+```
+
+Similarly to the configuration of a single locust process, you have the possibility to configure the test using docker compose in the `docker-compose.yml` file.
+
+Locust command: modify the locust command in the same way as the parameters are described above.
+
+```yml
+   command: -f /home/locust/locustfile.py --master --headless --expect-workers 4 -u 16 -r 1 -t 30 -s 1 --html=/home/locust/result.html
+```
+
+- `--expect-workers`: this is the only option that's new with this setup and let's you control for how many worker containers locust will wait to be up and running before starting the load test
+- `u`: the number of users, is the total number of users which will be distributed evenly among the worker processes. So with 4 workers and 16 users, each worker will run 4 users.
+
+By default, docker compose will launch 4 worker instances, however, you can scale number of workers with
+
+```bash
+sudo docker compose up --scale worker=6
+```
+
+However, note that the number of workes shouldn't be lower than the `expect-workers` parameter.
+As a rule of thumb, you should be running one worker per core of your host.
+
+#### Collect the results
+
+By default, the directory with the test configuration will be mounted in the containers, and locust will create the `result.html` in the mounted directory, so you will be able to access it after the containers are shut down and the test concluded.

locust_benchmark/common/__init__.py

@@ -0,0 +1,99 @@
+import datetime
+import random
+import string
+import json
+
+import numpy as np
+import pandas as pd
+
+from locust.runners import MasterRunner, LocalRunner
+import hsfs
+
+from hsfs import client
+from hsfs.client.exceptions import RestAPIError
+
+
+class HopsworksClient:
+    def __init__(self, environment=None):
+        with open("hopsworks_config.json") as json_file:
+            self.hopsworks_config = json.load(json_file)
+        if environment is None or isinstance(
+            environment.runner, (MasterRunner, LocalRunner)
+        ):
+            print(self.hopsworks_config)
+        self.connection = hsfs.connection(
+            project=self.hopsworks_config.get("project", "test"),
+            host=self.hopsworks_config.get("host", "localhost"),
+            port=self.hopsworks_config.get("port", 443),
+            api_key_file=".api_key",
+            secrets_store="local",
+        )
+        self.fs = self.connection.get_feature_store()
+
+        # test settings
+        self.external = self.hopsworks_config.get("external", False)
+        self.rows = self.hopsworks_config.get("rows", 1_000_000)
+        self.schema_repetitions = self.hopsworks_config.get("schema_repetitions", 1)
+        self.recreate_feature_group = self.hopsworks_config.get(
+            "recreate_feature_group", False
+        )
+        self.batch_size = self.hopsworks_config.get("batch_size", 100)
+
+    def get_or_create_fg(self):
+        locust_fg = self.fs.get_or_create_feature_group(
+            name="locust_fg",
+            version=1,
+            primary_key=["ip"],
+            online_enabled=True,
+            stream=True,
+        )
+        return locust_fg
+
+    def insert_data(self, locust_fg):
+        if locust_fg.id is not None and self.recreate_feature_group:
+            locust_fg.delete()
+            locust_fg = self.get_or_create_fg()
+        if locust_fg.id is None:
+            df = self.generate_insert_df(self.rows, self.schema_repetitions)
+            locust_fg.insert(df, write_options={"internal_kafka": not self.external})
+        return locust_fg
+
+    def get_or_create_fv(self, fg=None):
+        try:
+            return self.fs.get_feature_view("locust_fv", version=1)
+        except RestAPIError:
+            return self.fs.create_feature_view(
+                name="locust_fv",
+                query=fg.select_all(),
+                version=1,
+            )
+
+    def close(self):
+        if client._client is not None:
+            self.connection.close()
+
+    def generate_insert_df(self, rows, schema_repetitions):
+        data = {"ip": range(0, rows)}
+        df = pd.DataFrame.from_dict(data)
+
+        for i in range(0, schema_repetitions):
+            df["rand_ts_1_" + str(i)] = datetime.datetime.now()
+            df["rand_ts_2_" + str(i)] = datetime.datetime.now()
+            df["rand_int_1" + str(i)] = np.random.randint(0, 100000)
+            df["rand_int_2" + str(i)] = np.random.randint(0, 100000)
+            df["rand_float_1" + str(i)] = np.random.uniform(low=0.0, high=1.0)
+            df["rand_float_2" + str(i)] = np.random.uniform(low=0.0, high=1.0)
+            df["rand_string_1" + str(i)] = "".join(
+                random.choices(string.ascii_lowercase, k=5)
+            )
+            df["rand_string_2" + str(i)] = "".join(
+                random.choices(string.ascii_lowercase, k=5)
+            )
+            df["rand_string_3" + str(i)] = "".join(
+                random.choices(string.ascii_lowercase, k=5)
+            )
+            df["rand_string_4" + str(i)] = "".join(
+                random.choices(string.ascii_lowercase, k=5)
+            )
+
+        return df

locust_benchmark/common/hopsworks_client.py

@@ -0,0 +1,34 @@
+import inspect
+import time
+from locust import events
+
+
+def stopwatch(func):
+    def wrapper(*args, **kwargs):
+        # get task's function name
+        previous_frame = inspect.currentframe().f_back
+        _, _, task_name, _, _ = inspect.getframeinfo(previous_frame)
+
+        start = time.time()
+        result = None
+        try:
+            result = func(*args, **kwargs)
+        except Exception as e:
+            total = int((time.time() - start) * 1000)
+            events.request.fire(
+                request_type=task_name[:3].upper(),
+                name=task_name,
+                response_time=total,
+                exception=e,
+            )
+        else:
+            total = int((time.time() - start) * 1000)
+            events.request.fire(
+                request_type=task_name[:3].upper(),
+                name=task_name,
+                response_time=total,
+                response_length=0,
+            )
+        return result
+
+    return wrapper

locust_benchmark/common/stop_watch.py

@@ -0,0 +1,8 @@
+from common.hopsworks_client import HopsworksClient
+
+if __name__ == "__main__":
+
+    hopsworks_client = HopsworksClient()
+    fg = hopsworks_client.get_or_create_fg()
+    hopsworks_client.insert_data(fg)
+    hopsworks_client.close()

locust_benchmark/create_feature_group.py

@@ -0,0 +1,21 @@
+version: '3'
+
+services:
+  master:
+    extra_hosts:
+      - "host.docker.internal:host-gateway"
+    image: locust-hsfs:master
+    ports:
+     - "8089:8089"
+    volumes:
+      - ./:/home/locust
+    command: -f /home/locust/locustfile.py --master --headless --expect-workers 4 -u 16 -r 1 -t 30 -s 1 --html=/home/locust/result.html
+
+  worker:
+    extra_hosts:
+      - "host.docker.internal:host-gateway"
+    image: locust-hsfs:master
+    volumes:
+      - ./:/home/locust
+    command: -f /home/locust/locustfile.py --worker --master-host master
+    scale: 4

locust_benchmark/docker-compose.yml

@@ -0,0 +1,10 @@
+{
+    "host": "localhost",
+    "port": 443,
+    "project": "test",
+    "external": true,
+    "rows": 100000,
+    "schema_repetitions": 1,
+    "recreate_feature_group": true,
+    "batch_size": 100
+}