398 lines
16 KiB
Python
398 lines
16 KiB
Python
"""Data models for the multi-turn benchmark."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from dataclasses import dataclass, field
|
|
from typing import List, Optional
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class TurnResult:
|
|
"""Result of a single turn's HTTP request."""
|
|
|
|
ttft_ms: float # time to first token
|
|
fc_ms: float # first-chunk latency (time to N-th content chunk)
|
|
itl_ms: float # mean inter-token latency across output tokens
|
|
e2e_latency_ms: float # total request latency
|
|
input_tokens: int # reported by server (usage.prompt_tokens)
|
|
output_tokens: int # reported by server (usage.completion_tokens)
|
|
generated_text: str # generated text
|
|
itl_ms_list: List[float] = field(default_factory=list) # per-token ITL values
|
|
|
|
|
|
@dataclass
|
|
class TurnMetric:
|
|
"""Metrics for a single turn."""
|
|
|
|
session_id: str
|
|
turn: int # 0-indexed
|
|
ttft_ms: float
|
|
fc_ms: float # first-chunk latency
|
|
itl_ms: float # mean inter-token latency
|
|
e2e_latency_ms: float
|
|
input_tokens: int
|
|
output_tokens: int
|
|
start_time_ms: float # relative to benchmark start
|
|
itl_ms_list: List[float] = field(default_factory=list) # per-token ITL values
|
|
|
|
|
|
@dataclass
|
|
class WorkloadSpec:
|
|
"""Workload specification for multi-turn session benchmarks.
|
|
|
|
Supports simple mode: specify isl + hit_rate, derive user_tokens and sys_tokens.
|
|
All parameters are scalar (fixed) values -- no distributions.
|
|
"""
|
|
|
|
# Core parameters
|
|
num_sessions: Optional[int] = None # total unique sessions (None = duration-based)
|
|
num_turns: int = 1 # turns per session
|
|
osl: int = 1 # output sequence length per turn
|
|
think_time: float = 0.0 # seconds between turns within a session
|
|
|
|
# Traffic (use either concurrency or request_rate, not both)
|
|
concurrency: Optional[int] = None # max concurrent in-flight requests
|
|
request_rate: Optional[float] = None # requests per second (constant rate mode)
|
|
ramp_interval: float = -1.0 # seconds between session launches (-1 = auto)
|
|
|
|
# Duration-based mode (used with request_rate)
|
|
duration_s: float = 0.0 # seconds to run benchmark (0 = use num_sessions)
|
|
|
|
# Fraction of system prompt shared across all sessions
|
|
# 1.0 = identical system prompt, 0.0 = all unique
|
|
shared_system_prompt_ratio: float = 1.0
|
|
|
|
# Simple mode inputs (derive user_tokens, sys_tokens)
|
|
isl: Optional[int] = None
|
|
hit_rate: Optional[float] = None
|
|
|
|
# Resolved values (computed by resolve())
|
|
_user_tokens: int = field(default=0, init=False, repr=False)
|
|
_sys_tokens: int = field(default=0, init=False, repr=False)
|
|
|
|
def resolve(self) -> "WorkloadSpec":
|
|
"""Resolve the spec: derive user_tokens and sys_tokens from inputs. Call after init."""
|
|
if self.isl is None or self.hit_rate is None:
|
|
raise ValueError("Simple mode requires both --isl and --hit-rate.")
|
|
self._validate()
|
|
self._derive_from_simple()
|
|
return self
|
|
|
|
def _derive_from_simple(self) -> None:
|
|
"""Derive user_tokens and sys_tokens from (ISL, hit_rate, num_turns, OSL, shared_system_prompt_ratio).
|
|
|
|
Two equations, two unknowns (u = user_tokens, s = sys_tokens):
|
|
|
|
(1) ISL = s + (n+1)/2 · u + (n-1)/2 · a [average input length]
|
|
(2) (1-h)·ISL = (1-f)·s/n + u [average new-token fraction]
|
|
|
|
where n = num_turns, a = osl, f = shared_system_prompt_ratio, h = hit_rate.
|
|
|
|
Substituting s from (1) into (2) and solving for u:
|
|
|
|
u = [ (1-h)·ISL - (1-f)/n · (ISL - (n-1)·a/2) ]
|
|
/ [ 1 - (1-f)·(n+1)/(2n) ]
|
|
|
|
Then s = ISL - (n+1)/2 · u - (n-1)/2 · a.
|
|
|
|
Special case: when n=1 and f=0, equations (1) and (2) collapse to
|
|
s + u = ISL with h = s/(s+u), giving s = h·ISL and u = (1-h)·ISL.
|
|
"""
|
|
isl = self.isl
|
|
h = self.hit_rate
|
|
n = self.num_turns
|
|
a = self.osl
|
|
f = self.shared_system_prompt_ratio
|
|
|
|
denom = 1 - (1 - f) * (n + 1) / (2 * n)
|
|
if abs(denom) < 1e-9:
|
|
# n=1, f=0, h=0 (validated earlier): s=0, u=ISL.
|
|
sys_tokens = 0.0
|
|
user_tokens = float(isl)
|
|
else:
|
|
numer = (1 - h) * isl - (1 - f) / n * (isl - (n - 1) * a / 2)
|
|
user_tokens = numer / denom
|
|
sys_tokens = isl - (n + 1) / 2 * user_tokens - (n - 1) / 2 * a
|
|
|
|
if user_tokens < 0.5 or sys_tokens < -0.5:
|
|
suggestions = self._feasibility_suggestions()
|
|
which = "user_tokens" if user_tokens < 0.5 else "sys_tokens"
|
|
val = user_tokens if user_tokens < 0.5 else sys_tokens
|
|
raise ValueError(
|
|
f"Derived {which} = {val:.1f} is infeasible with "
|
|
f"(ISL={isl}, hit_rate={h}, num_turns={n}, "
|
|
f"OSL={a}, shared_system_prompt_ratio={f}).\n"
|
|
f"To fix, try one of:\n{suggestions}"
|
|
)
|
|
|
|
self._user_tokens = max(1, int(round(user_tokens)))
|
|
self._sys_tokens = max(0, int(round(sys_tokens)))
|
|
|
|
def _feasibility_suggestions(self) -> str:
|
|
"""Compute feasible boundary values for each parameter and return suggestions.
|
|
|
|
For each workload parameter, search for a boundary value that makes
|
|
the solver yield user_tokens >= 0.5 and sys_tokens >= -0.5 (the
|
|
minimum values that round to physically meaningful token counts:
|
|
at least 1 user token and non-negative system tokens).
|
|
"""
|
|
isl = self.isl
|
|
hit_rate = self.hit_rate
|
|
num_turns = self.num_turns
|
|
osl = self.osl
|
|
sharing = self.shared_system_prompt_ratio
|
|
lines = []
|
|
|
|
def _try_solve(isl_, hit_rate_, num_turns_, osl_, sharing_):
|
|
"""Solve for (user_tokens, sys_tokens) or return None if degenerate."""
|
|
denom = 1 - (1 - sharing_) * (num_turns_ + 1) / (2 * num_turns_)
|
|
if abs(denom) < 1e-9:
|
|
if hit_rate_ > 1e-9:
|
|
return None
|
|
return (float(isl_), 0.0)
|
|
numer = (1 - hit_rate_) * isl_ - (1 - sharing_) / num_turns_ * (
|
|
isl_ - (num_turns_ - 1) * osl_ / 2
|
|
)
|
|
user_tokens = numer / denom
|
|
sys_tokens = (
|
|
isl_ - (num_turns_ + 1) / 2 * user_tokens - (num_turns_ - 1) / 2 * osl_
|
|
)
|
|
return (user_tokens, sys_tokens)
|
|
|
|
def _feasible(isl_, hit_rate_, num_turns_, osl_, sharing_):
|
|
result = _try_solve(isl_, hit_rate_, num_turns_, osl_, sharing_)
|
|
# user_tokens >= 0.5 rounds to at least 1 token per turn;
|
|
# sys_tokens >= -0.5 rounds to at least 0 system prompt tokens.
|
|
return result is not None and result[0] >= 0.5 and result[1] >= -0.5
|
|
|
|
# Min ISL (binary search)
|
|
lo, hi = isl, isl * 20
|
|
if _feasible(hi, hit_rate, num_turns, osl, sharing):
|
|
while hi - lo > 1:
|
|
mid = (lo + hi) // 2
|
|
if _feasible(mid, hit_rate, num_turns, osl, sharing):
|
|
hi = mid
|
|
else:
|
|
lo = mid
|
|
lines.append(f" - ISL >= {hi} (with current params)")
|
|
|
|
# Max OSL
|
|
lo, hi = 1, osl
|
|
if _feasible(isl, hit_rate, num_turns, lo, sharing):
|
|
while hi - lo > 1:
|
|
mid = (lo + hi) // 2
|
|
if _feasible(isl, hit_rate, num_turns, mid, sharing):
|
|
lo = mid
|
|
else:
|
|
hi = mid
|
|
lines.append(f" - OSL <= {lo} (with current ISL={isl})")
|
|
|
|
# Min hit_rate / max hit_rate (search in 0.01 steps)
|
|
for h_try in range(0, 100):
|
|
h_val = h_try / 100.0
|
|
if _feasible(isl, h_val, num_turns, osl, sharing):
|
|
if h_val != hit_rate:
|
|
if h_val > hit_rate:
|
|
lines.append(
|
|
f" - hit_rate >= {h_val:.2f} (with current ISL/OSL)"
|
|
)
|
|
else:
|
|
lines.append(
|
|
f" - hit_rate <= {h_val:.2f} (with current ISL/OSL)"
|
|
)
|
|
break
|
|
|
|
# Max num_turns
|
|
for n_try in range(num_turns, 0, -1):
|
|
if _feasible(isl, hit_rate, n_try, osl, sharing):
|
|
if n_try != num_turns:
|
|
lines.append(f" - num_turns <= {n_try} (with current ISL/OSL)")
|
|
break
|
|
|
|
# Min shared_system_prompt_ratio
|
|
if sharing < 1.0:
|
|
for f_try in range(int(sharing * 100), 101):
|
|
f_val = f_try / 100.0
|
|
if _feasible(isl, hit_rate, num_turns, osl, f_val):
|
|
if f_val != sharing:
|
|
lines.append(f" - shared_system_prompt_ratio >= {f_val:.2f}")
|
|
break
|
|
|
|
return "\n".join(lines) if lines else " (no single-parameter fix found)"
|
|
|
|
def _validate(self) -> None:
|
|
"""Validate resolved parameters."""
|
|
if self.num_turns < 1:
|
|
raise ValueError("num_turns must be >= 1.")
|
|
if self.osl < 1:
|
|
raise ValueError("osl must be >= 1.")
|
|
if self.num_sessions is not None and self.num_sessions < 1:
|
|
raise ValueError("num_sessions must be >= 1.")
|
|
if self.num_sessions is None and self.duration_s <= 0:
|
|
raise ValueError(
|
|
"Must specify either --num-sessions or --duration (> 0) for rate-based mode."
|
|
)
|
|
if not (0 <= self.shared_system_prompt_ratio <= 1):
|
|
raise ValueError("shared_system_prompt_ratio must be in [0, 1].")
|
|
if self.think_time < 0:
|
|
raise ValueError("think_time must be >= 0.")
|
|
if (
|
|
self.num_turns == 1
|
|
and self.shared_system_prompt_ratio == 0
|
|
and self.hit_rate is not None
|
|
and self.hit_rate > 1e-9
|
|
):
|
|
raise ValueError(
|
|
f"Cannot achieve hit_rate={self.hit_rate} with num_turns=1 and "
|
|
f"shared_system_prompt_ratio=0. There is no caching source "
|
|
f"(no multi-turn history, no shared prefix). "
|
|
f"Set shared_system_prompt_ratio > 0 to enable cross-session "
|
|
f"prefix caching, or use num_turns > 1 for multi-turn caching."
|
|
)
|
|
|
|
if self.concurrency is None and self.request_rate is None:
|
|
raise ValueError("Must specify either --concurrency or --request-rate.")
|
|
if self.concurrency is not None and self.request_rate is not None:
|
|
raise ValueError("Cannot specify both --concurrency and --request-rate.")
|
|
if self.concurrency is not None and self.concurrency < 1:
|
|
raise ValueError("concurrency must be >= 1.")
|
|
if self.request_rate is not None and self.request_rate <= 0:
|
|
raise ValueError("request_rate must be > 0.")
|
|
|
|
if self.ramp_interval < 0:
|
|
if self.concurrency is not None:
|
|
if self.think_time > 0:
|
|
self.ramp_interval = self.think_time / self.concurrency
|
|
else:
|
|
self.ramp_interval = 0.0
|
|
else:
|
|
self.ramp_interval = 0.0
|
|
|
|
if (
|
|
self.concurrency is not None
|
|
and self.think_time > 0
|
|
and self.num_sessions is not None
|
|
and self.num_sessions < self.concurrency * 2
|
|
):
|
|
logger.warning(
|
|
"num_sessions=%d may be too low to sustain concurrency=%d "
|
|
"with think_time=%.1f. Consider increasing num_sessions.",
|
|
self.num_sessions,
|
|
self.concurrency,
|
|
self.think_time,
|
|
)
|
|
|
|
@property
|
|
def user_tokens(self) -> int:
|
|
return self._user_tokens
|
|
|
|
@property
|
|
def sys_tokens(self) -> int:
|
|
return self._sys_tokens
|
|
|
|
@property
|
|
def shared_s(self) -> int:
|
|
return int(round(self._sys_tokens * self.shared_system_prompt_ratio))
|
|
|
|
@property
|
|
def unique_s(self) -> int:
|
|
return self._sys_tokens - self.shared_s
|
|
|
|
def turn_input_tokens(self, k: int) -> int:
|
|
"""Total input tokens at turn k (1-indexed)."""
|
|
return self._sys_tokens + k * self._user_tokens + (k - 1) * self.osl
|
|
|
|
@property
|
|
def effective_isl(self) -> float:
|
|
n = self.num_turns
|
|
return (
|
|
self._sys_tokens + self._user_tokens * (n + 1) / 2 + self.osl * (n - 1) / 2
|
|
)
|
|
|
|
@property
|
|
def effective_h(self) -> float:
|
|
f = self.shared_system_prompt_ratio
|
|
n = self.num_turns
|
|
avg_new = (1 - f) * self._sys_tokens / n + self._user_tokens
|
|
isl = self.effective_isl
|
|
return 1.0 - avg_new / isl if isl > 0 else 0.0
|
|
|
|
def summary(self) -> dict:
|
|
per_turn = []
|
|
for k in range(1, self.num_turns + 1):
|
|
total = self.turn_input_tokens(k)
|
|
if k == 1:
|
|
cached = int(round(self._sys_tokens * self.shared_system_prompt_ratio))
|
|
else:
|
|
cached = (
|
|
self._sys_tokens + (k - 1) * self._user_tokens + (k - 1) * self.osl
|
|
)
|
|
new = total - cached
|
|
h_k = cached / total if total > 0 else 0.0
|
|
per_turn.append(
|
|
{
|
|
"turn": k,
|
|
"total": total,
|
|
"cached": cached,
|
|
"new": new,
|
|
"hit_rate": round(h_k, 4),
|
|
}
|
|
)
|
|
|
|
return {
|
|
"num_sessions": self.num_sessions,
|
|
"duration_s": self.duration_s,
|
|
"num_turns": self.num_turns,
|
|
"osl": self.osl,
|
|
"think_time": self.think_time,
|
|
"concurrency": self.concurrency,
|
|
"request_rate": self.request_rate,
|
|
"shared_system_prompt_ratio": self.shared_system_prompt_ratio,
|
|
"user_tokens_per_turn": self._user_tokens,
|
|
"system_prompt_tokens": self._sys_tokens,
|
|
"shared_system_prompt": self.shared_s,
|
|
"unique_system_prompt": self.unique_s,
|
|
"effective_isl": round(self.effective_isl, 1),
|
|
"effective_hit_rate": round(self.effective_h, 4),
|
|
"per_turn": per_turn,
|
|
}
|
|
|
|
def print_summary(self) -> None:
|
|
s = self.summary()
|
|
print("=" * 70)
|
|
print("Workload Spec (resolved)")
|
|
print("=" * 70)
|
|
if s["num_sessions"] is not None:
|
|
print(f" Sessions (N_s): {s['num_sessions']}")
|
|
else:
|
|
print(" Sessions (N_s): unlimited (duration-based)")
|
|
if s["duration_s"] > 0:
|
|
print(f" Duration: {s['duration_s']}s")
|
|
print(f" Turns per session (N_t): {s['num_turns']}")
|
|
print(f" User tokens/turn (u): {s['user_tokens_per_turn']}")
|
|
print(
|
|
f" System prompt (s): {s['system_prompt_tokens']} "
|
|
f"(shared={s['shared_system_prompt']}, unique={s['unique_system_prompt']})"
|
|
)
|
|
print(f" Output tokens (o): {s['osl']}")
|
|
print(f" Think time: {s['think_time']}s")
|
|
if self.concurrency is not None:
|
|
print(f" Concurrency (C): {self.concurrency}")
|
|
print(f" Ramp interval: {self.ramp_interval:.3f}s")
|
|
if self.request_rate is not None:
|
|
print(f" Request rate (QPS): {self.request_rate}")
|
|
print(f" Shared sys prompt ratio: {s['shared_system_prompt_ratio']}")
|
|
print(f" Effective avg ISL: {s['effective_isl']}")
|
|
print(f" Effective avg hit rate: {s['effective_hit_rate']:.1%}")
|
|
print("-" * 70)
|
|
print(f" {'Turn':<6} {'Total':<8} {'Cached':<8} {'New':<8} {'Hit Rate':<10}")
|
|
for t in s["per_turn"]:
|
|
print(
|
|
f" {t['turn']:<6} {t['total']:<8} {t['cached']:<8} "
|
|
f"{t['new']:<8} {t['hit_rate']:.1%}"
|
|
)
|
|
print("=" * 70)
|