Pythonlatency_benchmark.py
Copy
1
import time
2
import statistics
3
from openai import OpenAI
4
 
5
client = OpenAI()
6
 
7
def measure_latency(model: str, prompt: str, n_trials: int = 5) -> dict:
8
    ttfts, totals, tpss = [], [], []
9
 
10
    for _ in range(n_trials):
11
        start = time.monotonic()
12
        first_token_time = None
13
        output_tokens = 0
14
 
15
        stream = client.chat.completions.create(
16
            model=model,
17
            messages=[{"role": "user", "content": prompt}],
18
            stream=True,
19
        )
20
 
21
        for chunk in stream:
22
            delta = chunk.choices[0].delta.content
23
            if delta:
24
                if first_token_time is None:
25
                    first_token_time = time.monotonic()
26
                output_tokens += 1
27
 
28
        total = time.monotonic() - start
29
        ttft = (first_token_time or total) - start
30
        tps = output_tokens / total if total > 0 else 0
31
 
32
        ttfts.append(ttft)
33
        totals.append(total)
34
        tpss.append(tps)
35
 
36
    def pct(data: list[float], p: int) -> float:
37
        return round(statistics.quantiles(data, n=100)[p - 1], 3)
38
 
39
    return {
40
        "model": model,
41
        "ttft_p50": pct(ttfts, 50),
42
        "ttft_p95": pct(ttfts, 95),
43
        "total_p50": pct(totals, 50),
44
        "total_p95": pct(totals, 95),
45
        "tps_mean": round(statistics.mean(tpss), 1),
46
    }
47
 
48
SHORT = "What is 2 + 2?"
49
MEDIUM = "Explain gradient descent in detail."
50
 
51
for model in ["gpt-4o-mini", "gpt-4o"]:
52
    result = measure_latency(model, MEDIUM)
53
    print(result)
54