Record a bit more metrics

tests/colosseum/controller_load_test.py

@@ -1,7 +1,10 @@
 import os
+import csv
+import json
 import time
 import random
 import itertools
+from statistics import quantiles
 import multiprocessing as mp
 
 import tyro
@@ -70,42 +73,89 @@ PROMPTS = [
 ] * 2
 
 
-def request(prompt: str) -> tuple[str, str, str, str, float]:
+def request(prompt: str) -> tuple[str, float, float, float]:
     time.sleep(random.random() * 5)
-    client = ControllerClient(CONTROLLER_ADDR, timeout=30)
+    client = ControllerClient(CONTROLLER_ADDR, timeout=60)
     
     response_a, response_b = "", ""
+    first_token_latency = -1.0
+    num_tokens = 0
     start_time = time.monotonic()
-    for resp_a, resp_b in itertools.zip_longest(
+    for i, (resp_a, resp_b) in enumerate(itertools.zip_longest(
         client.prompt(prompt, index=0),
         client.prompt(prompt, index=1),
-    ):
+    )):
+        if i == 0:
+            first_token_latency = time.monotonic() - start_time
         if resp_a is not None:
+            num_tokens += 1
             response_a += resp_a
         if resp_b is not None:
+            num_tokens += 1
             response_b += resp_b
 
     latency = time.monotonic() - start_time
-    return client.request_id, prompt, response_a, response_b, latency
+    tokens_per_second = num_tokens / latency
+    return client.request_id, latency, first_token_latency, tokens_per_second
 
 
-def main(concurrency: int = len(PROMPTS), logfile: str = "load_test_results.csv"):
-    latencies = []
+def main(
+    concurrencies: list[int] = [10],
+    result_csv: str = "load_test_results.csv",
+    ftl_json: str = "ftl_dist.json",
+):
+    data = []
+    ftl_dist = {}
 
-    start_time = time.monotonic()
-    with mp.Pool(processes=concurrency) as pool:
-        for request_id, prompt, response_a, response_b, latency in pool.imap_unordered(request, PROMPTS):
-            latencies.append(latency)
-            print(f"Request ID {request_id} finished, {latency=:.2f}s")
-
-    total_time = time.monotonic() - start_time
-    average_latency = sum(latencies) / len(latencies)
-    requests_per_second = len(latencies) / total_time
-    print(f"Total time: {total_time:.2f}s")
-    print(f"Average latency: {average_latency:.2f}s")
-    print(f"Requests per second: {requests_per_second:.2f}")
-    with open(logfile, "a") as f:
-        f.write(f"{concurrency},{total_time},{average_latency},{requests_per_second}\n")
+    for concurrency in concurrencies:
+        latencies = []
+        first_token_latencies = []
+        tps = []
+
+        start_time = time.monotonic()
+        with mp.Pool(processes=concurrency) as pool:
+            for request_id, latency, first_token_latency, tokens_per_second in pool.imap_unordered(request, PROMPTS):
+                latencies.append(latency)
+                first_token_latencies.append(first_token_latency)
+                tps.append(tokens_per_second)
+                print(f"Request ID {request_id} finished, {latency=:.2f}s, {first_token_latency=:.2f}s, {tokens_per_second=:.2f} tokens/s")
+
+        total_time = time.monotonic() - start_time
+        average_latency = sum(latencies) / len(latencies)
+        average_first_token_latency = sum(first_token_latencies) / len(first_token_latencies)
+        first_token_latency_quartiles = quantiles(first_token_latencies, n=10)
+        ftl_dist[concurrency] = first_token_latencies
+        average_tokens_per_second = sum(tps) / len(tps)
+        requests_per_second = len(latencies) / total_time
+        print(f"Total time: {total_time:.2f}s")
+        print(f"Average latency: {average_latency:.2f}s")
+        print(f"Average first token latency: {average_first_token_latency:.2f}s")
+        print(f"Average tokens per second: {average_tokens_per_second:.2f}")
+        print(f"Requests per second: {requests_per_second:.2f}")
+        print(f"First token latency quartiles: {first_token_latency_quartiles}")
+        data.append((
+            concurrency,
+            total_time,
+            average_latency,
+            average_first_token_latency,
+            average_tokens_per_second,
+            requests_per_second,
+        ))
+
+    with open(result_csv, "w") as f:
+        writer = csv.writer(f)
+        writer.writerow((
+            "concurrency",
+            "total_time",
+            "average_latency",
+            "average_first_token_latency",
+            "average_tokens_per_second",
+            "requests_per_second",
+        ))
+        writer.writerows(data)
+
+    with open(ftl_json, "w") as f:
+        json.dump(ftl_dist, f)
 
 
 if __name__ == "__main__":