[Benchmark] Fixing minor bugs in benchmark scripts (#192)

* Fixing minor bugs in benchmark scripts

Signed-off-by: YuhanLiu11 <[email protected]>

* Fixing minor bugs in benchmark scripts

Signed-off-by: YuhanLiu11 <[email protected]>

* Fixing minor bugs in benchmark scripts

Signed-off-by: YuhanLiu11 <[email protected]>

---------

Signed-off-by: YuhanLiu11 <[email protected]>

benchmarks/multi-round-qa/plot.py

@@ -0,0 +1,70 @@
+import os
+
+import numpy as np
+import pandas as pd
+from matplotlib import pyplot as plt
+
+stack_results = []
+qpses = []
+for qps in np.arange(0.3, 1.2, 0.2):
+    # round qps to 1 decimal
+    q = round(qps, 1)
+    file = f"stack_output_{q}.csv"
+    if not os.path.exists(file):
+        continue
+    qpses += [q]
+    # Read csv file
+    data = pd.read_csv(file)["ttft"].tolist()
+    stack_results.append(sum(data) / len(data))
+print(stack_results)
+
+plt.plot(
+    qpses,
+    stack_results,
+    label="Production Stack",
+    marker="x",
+    color="blue",
+    linewidth=2,
+    markersize=8,
+)
+aibrix_results = []
+qpses = []
+for qps in np.arange(0.3, 1.2, 0.2):
+    # round qps to 1 decimal
+    q = round(qps, 1)
+
+    file = f"aibrix_output_{q}.csv"
+    if not os.path.exists(file):
+        continue
+    qpses += [q]
+    # Read csv file
+    data = pd.read_csv(file)["ttft"].tolist()
+    aibrix_results.append(sum(data) / len(data))
+print(aibrix_results)
+plt.plot(
+    qpses,
+    aibrix_results,
+    label="Aibrix",
+    marker="o",
+    color="red",
+    linewidth=2,
+    markersize=8,
+)
+
+native_results = []
+for qps in np.arange(0.3, 1.2, 0.2):
+    # round qps to 1 decimal
+    q = round(qps, 1)
+    file = f"native_output_{q}.csv"
+    # Read csv file
+    data = pd.read_csv(file)["ttft"].tolist()
+    native_results.append(sum(data) / len(data))
+plt.plot(np.arange(0.3, 1.2, 0.2), native_results, label="Native K8S")
+
+
+plt.xlim(left=0)
+plt.ylim(top=10)
+plt.xlabel("QPS")
+plt.ylabel("Average Time to First Token (s)")
+plt.title("Average Time to First Token vs QPS")
+plt.savefig("figure.png")

benchmarks/multi-round-qa/run.sh

@@ -1,8 +1,8 @@
 #!/bin/bash
 
 # Ensure correct number of arguments
-if [[ $# -ne 2 ]]; then
-    echo "Usage: $0 <model> <base url>"
+if [[ $# -ne 3 ]]; then
+    echo "Usage: $0 <model> <base url> <output file key>"
     exit 1
 fi
 
@@ -36,16 +36,11 @@ run_benchmark() {
         --time 100
 }
 
-# Validate if a key argument is provided
-if [[ -z "$3" ]]; then
-    echo "Error: Missing key argument"
-    exit 1
-fi
 
 key=$3
 
 # Run benchmarks for different QPS values
-for qps in 0.3 0.5 0.7 0.9 1.1; do
+for qps in 0.1 0.3 0.5 0.7 0.9 1.1; do
     output_file="${key}_output_${qps}.csv"
     run_benchmark "$qps" "$output_file"
 done