Add vector db benchmark

redis_benchmarks_specification/__builder__/builder.py

@@ -381,6 +381,8 @@ def builder_process_stream(
                 #     build_vars_str,
                 # )
                 build_command = "sh -c 'make -j'"
+                if "build_command" in build_config:
+                    build_command = build_config["build_command"]
                 if b"build_command" in testDetails:
                     build_command = testDetails[b"build_command"].decode()
                 server_name = "redis"
@@ -649,7 +651,13 @@ def generate_benchmark_stream_request(
     prefix = f"github_org={github_org}/github_repo={github_repo}/git_branch={str(git_branch)}/git_version={str(git_version)}/git_hash={str(git_hash)}"
     for artifact in build_artifacts:
         bin_key = f"zipped:artifacts:{prefix}:{id}:{artifact}.zip"
-        bin_artifact = open(f"{redis_temporary_dir}src/{artifact}", "rb").read()
+        if artifact == "redisearch.so":
+            bin_artifact = open(
+                f"{redis_temporary_dir}modules/redisearch/src/bin/linux-x64-release/search-community/{artifact}",
+                "rb",
+            ).read()
+        else:
+            bin_artifact = open(f"{redis_temporary_dir}src/{artifact}", "rb").read()
         bin_artifact_len = len(bytes(bin_artifact))
         assert bin_artifact_len > 0
         conn.set(bin_key, bytes(bin_artifact), ex=REDIS_BINS_EXPIRE_SECS)

redis_benchmarks_specification/__self_contained_coordinator__/clients.py

@@ -22,3 +22,27 @@ def prepare_memtier_benchmark_parameters(
         benchmark_command_str = benchmark_command_str + " " + clientconfig["arguments"]
 
     return None, benchmark_command_str
+
+
+def prepare_vector_db_benchmark_parameters(
+    clientconfig, full_benchmark_path, port, server, password, client_mnt_point
+):
+    benchmark_command = []
+    # if port is not None:
+    #     benchmark_command.extend(["REDIS_PORT={}".format(port)])
+    # if password is not None:
+    #     benchmark_command.extend(["REDIS_AUTH={}".format(password)])
+    benchmark_command.extend(
+        [
+            full_benchmark_path,
+            "--host",
+            f"{server}",
+        ]
+    )
+    benchmark_command.extend(["--engines", clientconfig.get("engines", "redis-test")])
+    benchmark_command.extend(
+        ["--datasets", clientconfig.get("datasets", "glove-100-angular")]
+    )
+    benchmark_command_str = " ".join(benchmark_command)
+    benchmark_command_str = f"bash -c 'ITERATIONS=1 {benchmark_command_str} && mv /code/results {client_mnt_point}.'"
+    return None, benchmark_command_str

redis_benchmarks_specification/__self_contained_coordinator__/post_processing.py

@@ -0,0 +1,19 @@
+import os
+import json
+
+
+def post_process_vector_db(temporary_dir):
+    results_dir = os.path.join(temporary_dir, "results")
+    results = {}
+    for file in os.listdir(results_dir):
+        if "upload" in file:
+            with open(os.path.join(results_dir, file), "r") as f:
+                upload_results = json.load(f)
+                results["upload_time"] = upload_results["results"]["upload_time"]
+        else:
+            with open(os.path.join(results_dir, file), "r") as f:
+                query_results = json.load(f)
+                results["rps"] = query_results["results"]["rps"]
+                results["precision"] = query_results["results"]["mean_precisions"]
+                results["total_time"] = query_results["results"]["total_time"]
+    return results

redis_benchmarks_specification/__self_contained_coordinator__/runners.py

@@ -56,6 +56,7 @@
 )
 from redis_benchmarks_specification.__self_contained_coordinator__.clients import (
     prepare_memtier_benchmark_parameters,
+    prepare_vector_db_benchmark_parameters,
 )
 from redis_benchmarks_specification.__self_contained_coordinator__.cpuset import (
     extract_db_cpu_limit,
@@ -347,9 +348,12 @@ def process_self_contained_coordinator_stream(
                             # backwards compatible
                             if benchmark_tool is None:
                                 benchmark_tool = "redis-benchmark"
-                            full_benchmark_path = "/usr/local/bin/{}".format(
-                                benchmark_tool
-                            )
+                            if benchmark_tool == "vector_db_benchmark":
+                                full_benchmark_path = "python /code/run.py"
+                            else:
+                                full_benchmark_path = "/usr/local/bin/{}".format(
+                                    benchmark_tool
+                                )
 
                             # setup the benchmark
                             (
@@ -370,32 +374,42 @@ def process_self_contained_coordinator_stream(
                                     local_benchmark_output_filename
                                 )
                             )
-                            if "memtier_benchmark" not in benchmark_tool:
-                                # prepare the benchmark command
+                            if "memtier_benchmark" in benchmark_tool:
                                 (
-                                    benchmark_command,
+                                    _,
                                     benchmark_command_str,
-                                ) = prepare_benchmark_parameters(
-                                    benchmark_config,
+                                ) = prepare_memtier_benchmark_parameters(
+                                    benchmark_config["clientconfig"],
                                     full_benchmark_path,
                                     redis_proc_start_port,
                                     "localhost",
                                     local_benchmark_output_filename,
-                                    False,
                                     benchmark_tool_workdir,
-                                    False,
                                 )
-                            else:
+                            elif "vector_db_benchmark" in benchmark_tool:
                                 (
                                     _,
                                     benchmark_command_str,
-                                ) = prepare_memtier_benchmark_parameters(
+                                ) = prepare_vector_db_benchmark_parameters(
                                     benchmark_config["clientconfig"],
                                     full_benchmark_path,
                                     redis_proc_start_port,
                                     "localhost",
+                                )
+                            else:
+                                # prepare the benchmark command
+                                (
+                                    benchmark_command,
+                                    benchmark_command_str,
+                                ) = prepare_benchmark_parameters(
+                                    benchmark_config,
+                                    full_benchmark_path,
+                                    redis_proc_start_port,
+                                    "localhost",
                                     local_benchmark_output_filename,
+                                    False,
                                     benchmark_tool_workdir,
+                                    False,
                                 )
 
                             client_container_image = extract_client_container_image(