From 24664aae217eac2ae736b0e0c77b6928f6771727 Mon Sep 17 00:00:00 2001
From: zacharya19 <zacharya19@gmail.com>
Date: Thu, 30 May 2024 17:00:13 +0300
Subject: [PATCH] chor(benchmark): create basic Python control

---
 tools/benchmark/runner.py | 99 +++++++++++++++++++++++++++++++++++++++
 tools/requirements.txt    |  2 +
 2 files changed, 101 insertions(+)
 create mode 100644 tools/benchmark/runner.py

diff --git a/tools/benchmark/runner.py b/tools/benchmark/runner.py
new file mode 100644
index 000000000000..b0855a284b55
--- /dev/null
+++ b/tools/benchmark/runner.py
@@ -0,0 +1,99 @@
+#!/usr/bin/env python3
+import os
+from typing import List
+
+import attr
+from kubernetes import client, config, watch
+from kubernetes.client import models
+
+
+@attr.define
+class TestCase:
+    name: str
+    timeout: int = 6000
+
+
+TEST_CASES: List[TestCase] = [TestCase(name="default")]
+
+config.load_kube_config()
+v1 = client.BatchV1Api()
+
+
+def wait_for_job_completion(namespace: str, job_name: str, timeout_seconds: int):
+    w = watch.Watch()
+    try:
+        for event in w.stream(
+            v1.list_namespaced_job, namespace=namespace, timeout_seconds=timeout_seconds
+        ):
+            job = event["object"]  # type: ignore
+
+            if job.metadata.name != job_name:  # type: ignore
+                continue
+
+            conditions = job.status.conditions  # type: ignore
+            if not conditions:
+                continue
+
+            for condition in conditions:
+                if condition.status != "True":
+                    continue
+
+                if condition.type == "Complete":
+                    print(f"Job {job_name} completed successfully.")
+                    return
+
+                if condition.type == "Failed":
+                    print(f"Job {job_name} failed.")
+                    return
+    finally:
+        w.stop()
+
+    print(
+        f"Timeout or job {job_name} did not reach a complete/failed state within the specified timeout."
+    )
+
+
+def main():
+    namespace = os.environ["NAMESPACE"]
+
+    for case in TEST_CASES:
+        job_name = f"memtier-benchmark-{case.name}"
+
+        print(f"Running case {job_name}")
+        v1.create_namespaced_job(
+            namespace=namespace,
+            body=models.V1Job(
+                metadata={"name": f"memtier-benchmark-{case.name}"},
+                spec=models.V1JobSpec(
+                    backoff_limit=0,
+                    template=models.V1JobTemplateSpec(
+                        spec={
+                            "restartPolicy": "Never",
+                            "containers": [
+                                models.V1Container(
+                                    name="memtier",
+                                    image="redislabs/memtier_benchmark:latest",
+                                    args=[
+                                        "memtier_benchmark --pipeline=30 --key-maximum=10000 -c 10 -t 2 --requests=500000 --expiry-range=10-100 --reconnect-interval=10000 --distinct-client-seed --hide-histogram -s dragonfly-sample"
+                                    ],
+                                    command=[
+                                        # This is important! without it memtier cannot DIG the dragonfly SVC domain
+                                        "sh",
+                                        "-c",
+                                    ],
+                                    resources={
+                                        "requests": {"cpu": "2", "memory": "500Mi"},
+                                        "limits": {"cpu": "2", "memory": "500Mi"},
+                                    },
+                                )
+                            ],
+                        }
+                    ),
+                ),
+            ),
+        )
+        wait_for_job_completion(namespace, job_name, case.timeout)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tools/requirements.txt b/tools/requirements.txt
index 25f05b2be1bb..7ccde0b937a9 100644
--- a/tools/requirements.txt
+++ b/tools/requirements.txt
@@ -8,3 +8,5 @@ requests==2.28.1
 aiocsv==1.2.3
 aiofiles==22.1.0
 numpy==1.24.1
+kubernetes==26.1.0
+attrs==22.2.0