wehub-resource-sync
1046 lines
38 KiB
Python
1046 lines
38 KiB
Python
"""Benchmark runner functions for multi-turn LLM inference benchmarks.
|
|
|
|
Contains:
|
|
- Concurrency-based (closed-loop) runner
|
|
- Rate-based (constant QPS) runner
|
|
- Smoke test runner
|
|
- Direct entry point (run_direct)
|
|
"""
|
|
|
|
import argparse
|
|
import asyncio
|
|
import hashlib
|
|
import json
|
|
import logging
|
|
import math
|
|
import random
|
|
import sys
|
|
import time
|
|
from collections import Counter
|
|
from statistics import mean
|
|
from typing import Optional
|
|
|
|
import aiohttp
|
|
import numpy as np
|
|
from tqdm import tqdm
|
|
|
|
from ray.llm._internal.serve.benchmark.http_client import send_chat_completion
|
|
from ray.llm._internal.serve.benchmark.models import TurnMetric, WorkloadSpec
|
|
from ray.llm._internal.serve.benchmark.reporting import report_results
|
|
from ray.llm._internal.serve.benchmark.text_gen import (
|
|
Conversation,
|
|
TextGenerator,
|
|
conversation_factory,
|
|
)
|
|
from ray.llm._internal.serve.benchmark.turn import execute_single_turn
|
|
|
|
# ============================================================================
|
|
# LOGGING SETUP
|
|
# ============================================================================
|
|
|
|
|
|
class TqdmLoggingHandler(logging.Handler):
|
|
"""Custom logging handler that writes through tqdm to keep progress bar at bottom."""
|
|
|
|
def emit(self, record):
|
|
try:
|
|
msg = self.format(record)
|
|
tqdm.write(msg, file=sys.stderr)
|
|
self.flush()
|
|
except Exception:
|
|
self.handleError(record)
|
|
|
|
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format="%(asctime)s [%(levelname)s] %(message)s",
|
|
datefmt="%H:%M:%S",
|
|
handlers=[TqdmLoggingHandler()],
|
|
)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
# ============================================================================
|
|
# BENCHMARK STATE (concurrency mode)
|
|
# ============================================================================
|
|
|
|
|
|
class BenchmarkState:
|
|
"""Centralized state management for benchmark execution.
|
|
|
|
Consolidates all shared state and synchronization into a single class
|
|
with one lock, replacing the previous 7+ separate locks.
|
|
"""
|
|
|
|
def __init__(self, spec: WorkloadSpec):
|
|
self.spec = spec
|
|
|
|
# Metrics storage
|
|
self.warmup_metrics: list[TurnMetric] = []
|
|
self.measured_metrics: list[TurnMetric] = []
|
|
|
|
# Session tracking
|
|
self.conversations: dict[int, Conversation] = {}
|
|
self.session_is_warmup: dict[int, bool] = {}
|
|
self.next_session_id = 0
|
|
self.completed_sessions = 0
|
|
|
|
# Concurrency tracking
|
|
self.inflight = 0
|
|
self.max_inflight = 0
|
|
self.inflight_samples: list[int] = []
|
|
|
|
# Warm-up state
|
|
self.warmup_complete = False
|
|
self.active_turns: dict[str, int] = {} # session_id -> turn_idx
|
|
|
|
# Single lock for all state
|
|
self._lock = asyncio.Lock()
|
|
|
|
async def record_metric(self, metric: TurnMetric) -> None:
|
|
"""Thread-safe metric recording."""
|
|
async with self._lock:
|
|
if self.warmup_complete:
|
|
self.measured_metrics.append(metric)
|
|
else:
|
|
self.warmup_metrics.append(metric)
|
|
|
|
async def get_next_session(self) -> Optional[int]:
|
|
"""Get next available session ID, or None if exhausted."""
|
|
async with self._lock:
|
|
limit = self.spec.num_sessions
|
|
if limit is None or self.next_session_id < limit:
|
|
session_id = self.next_session_id
|
|
self.next_session_id += 1
|
|
return session_id
|
|
return None
|
|
|
|
async def has_remaining_sessions(self) -> bool:
|
|
"""True if more session indices can still be allocated (does not consume an ID)."""
|
|
async with self._lock:
|
|
limit = self.spec.num_sessions
|
|
if limit is None:
|
|
return True
|
|
return self.next_session_id < limit
|
|
|
|
async def get_inflight(self) -> int:
|
|
"""Current in-flight HTTP requests (concurrency benchmark)."""
|
|
async with self._lock:
|
|
return self.inflight
|
|
|
|
async def get_or_create_conversation(
|
|
self,
|
|
session_idx: int,
|
|
spec: WorkloadSpec,
|
|
shared_system_text: str,
|
|
text_gen: TextGenerator,
|
|
) -> tuple[Conversation, bool]:
|
|
"""Get or create a conversation. Returns (conversation, is_warmup)."""
|
|
async with self._lock:
|
|
if session_idx not in self.conversations:
|
|
conv = conversation_factory(
|
|
session_idx, spec, shared_system_text, text_gen
|
|
)
|
|
self.conversations[session_idx] = conv
|
|
# Mark as warmup if steady-state not yet reached
|
|
self.session_is_warmup[session_idx] = not self.warmup_complete
|
|
else:
|
|
conv = self.conversations[session_idx]
|
|
|
|
is_warmup = self.session_is_warmup[session_idx]
|
|
return conv, is_warmup
|
|
|
|
async def mark_session_active(self, session_id: str, turn_idx: int) -> None:
|
|
"""Mark a session as active at a specific turn."""
|
|
async with self._lock:
|
|
self.active_turns[session_id] = turn_idx
|
|
|
|
async def mark_session_complete(self, session_idx: int, session_id: str) -> None:
|
|
"""Mark a session as complete and clean up."""
|
|
async with self._lock:
|
|
self.active_turns.pop(session_id, None)
|
|
if session_idx in self.conversations:
|
|
del self.conversations[session_idx]
|
|
self.completed_sessions += 1
|
|
|
|
async def track_inflight_start(self) -> int:
|
|
"""Track start of inflight request. Returns current inflight count."""
|
|
async with self._lock:
|
|
self.inflight += 1
|
|
if self.inflight > self.max_inflight:
|
|
self.max_inflight = self.inflight
|
|
self.inflight_samples.append(self.inflight)
|
|
return self.inflight
|
|
|
|
async def track_inflight_end(self) -> int:
|
|
"""Track end of inflight request. Returns inflight count BEFORE decrement."""
|
|
async with self._lock:
|
|
current = self.inflight
|
|
self.inflight -= 1
|
|
return current
|
|
|
|
async def mark_warmup_complete(self) -> None:
|
|
"""Mark warm-up phase complete. All future sessions are measurement."""
|
|
async with self._lock:
|
|
self.warmup_complete = True
|
|
|
|
async def check_and_complete_warmup(self) -> bool:
|
|
"""Check if steady-state reached and complete warm-up if so.
|
|
|
|
Returns True if warm-up was just completed, False otherwise.
|
|
"""
|
|
if self.warmup_complete:
|
|
return False
|
|
|
|
async with self._lock:
|
|
# Double-check after acquiring lock
|
|
if self.warmup_complete:
|
|
return False
|
|
|
|
if not self.active_turns:
|
|
return False
|
|
|
|
# Calculate entropy. For num_turns==1, max_entropy=log2(1)=0 so the
|
|
# threshold is 0.0 — any non-empty active_turns distribution satisfies
|
|
# entropy >= 0, so warm-up ends on the first turn completion.
|
|
max_entropy = np.log2(self.spec.num_turns)
|
|
entropy_threshold = 0.50 * max_entropy
|
|
entropy = self._calculate_entropy()
|
|
|
|
if entropy >= entropy_threshold:
|
|
self.warmup_complete = True
|
|
pct_of_target = (
|
|
(entropy / entropy_threshold * 100)
|
|
if entropy_threshold > 0
|
|
else 100.0
|
|
)
|
|
logger.info(
|
|
"Steady-state reached! Entropy=%.3f (threshold=%.3f, %.0f%% of target)",
|
|
entropy,
|
|
entropy_threshold,
|
|
pct_of_target,
|
|
)
|
|
logger.info(
|
|
"Warm-up complete (%d sessions, %d requests discarded). "
|
|
"Now measuring steady-state...",
|
|
self.completed_sessions,
|
|
len(self.warmup_metrics),
|
|
)
|
|
return True
|
|
|
|
return False
|
|
|
|
def _calculate_entropy(self) -> float:
|
|
"""Calculate Shannon entropy of active turn distribution (must be called with lock held)."""
|
|
if not self.active_turns:
|
|
return 0.0
|
|
counts = Counter(self.active_turns.values())
|
|
total = sum(counts.values())
|
|
probs = [c / total for c in counts.values()]
|
|
return -sum(p * np.log2(p) for p in probs if p > 0)
|
|
|
|
async def get_stats(self) -> dict:
|
|
"""Get current benchmark stats (thread-safe)."""
|
|
async with self._lock:
|
|
return {
|
|
"warmup_metrics": len(self.warmup_metrics),
|
|
"measured_metrics": len(self.measured_metrics),
|
|
"completed_sessions": self.completed_sessions,
|
|
"max_inflight": self.max_inflight,
|
|
"avg_inflight": (
|
|
mean(self.inflight_samples) if self.inflight_samples else 0
|
|
),
|
|
"warmup_complete": self.warmup_complete,
|
|
}
|
|
|
|
|
|
# ============================================================================
|
|
# PROGRESS BAR
|
|
# ============================================================================
|
|
|
|
|
|
def create_progress_bar(total: int) -> tqdm:
|
|
"""Create a sticky progress bar for request tracking."""
|
|
return tqdm(
|
|
total=total,
|
|
desc="Requests",
|
|
unit="req",
|
|
bar_format=(
|
|
"{l_bar}{bar}| {n_fmt}/{total_fmt} "
|
|
"[{elapsed}<{remaining}, {rate_fmt}, inflight={postfix}]"
|
|
),
|
|
position=0,
|
|
leave=True,
|
|
file=sys.stderr,
|
|
)
|
|
|
|
|
|
# ============================================================================
|
|
# CONCURRENCY-BASED BENCHMARK
|
|
# ============================================================================
|
|
|
|
|
|
async def run_benchmark(
|
|
spec: WorkloadSpec,
|
|
base_url: str,
|
|
model: str,
|
|
shared_system_text: str,
|
|
text_gen: TextGenerator,
|
|
first_chunk_threshold: int = 16,
|
|
api_key: Optional[str] = None,
|
|
warmup_jitter_max_s: float = 10.0,
|
|
) -> list[TurnMetric]:
|
|
"""Run the full benchmark with rolling session pool and steady-state warm-up.
|
|
|
|
Maintains a pool of `concurrency` worker tasks that continuously pull turns
|
|
from a queue. This creates a natural turn distribution mix (some sessions at
|
|
turn 0, others at turn 5, etc.) after warm-up.
|
|
|
|
Warm-up continues until steady-state is reached: turn distribution entropy
|
|
>= 50% of maximum (indicating reasonable spread across multiple turns).
|
|
"""
|
|
if spec.concurrency is None:
|
|
raise ValueError("closed-loop benchmark requires concurrency to be set")
|
|
if spec.num_sessions is None:
|
|
raise ValueError("closed-loop (concurrency) benchmark requires --num-sessions")
|
|
|
|
state = BenchmarkState(spec)
|
|
request_sem = asyncio.Semaphore(spec.concurrency)
|
|
|
|
bench_start_ns = time.perf_counter_ns()
|
|
|
|
total_expected = spec.num_sessions * spec.num_turns
|
|
|
|
connector = aiohttp.TCPConnector(limit=spec.concurrency * 2)
|
|
async with aiohttp.ClientSession(connector=connector) as http_session:
|
|
pbar = create_progress_bar(total_expected)
|
|
pbar.set_postfix_str("0")
|
|
|
|
# Turn queue: (session_idx, turn_idx, delay_until)
|
|
turn_queue: asyncio.Queue[tuple[int, int, float]] = asyncio.Queue()
|
|
|
|
async def enqueue_new_session() -> bool:
|
|
"""Add a new session (turn 0) to the queue if any remain."""
|
|
session_idx = await state.get_next_session()
|
|
if session_idx is not None:
|
|
# Stagger initial session launches using configured ramp_interval
|
|
delay = session_idx * spec.ramp_interval
|
|
await turn_queue.put((session_idx, 0, time.perf_counter() + delay))
|
|
return True
|
|
return False
|
|
|
|
# Initialize: fill queue with just enough sessions to saturate concurrency
|
|
for _ in range(min(spec.concurrency, spec.num_sessions)):
|
|
await enqueue_new_session()
|
|
|
|
async def execute_turn(session_idx: int, turn_idx: int) -> None:
|
|
"""Execute a single turn for a session."""
|
|
conv, is_warmup = await state.get_or_create_conversation(
|
|
session_idx, spec, shared_system_text, text_gen
|
|
)
|
|
await state.mark_session_active(conv.session_id, turn_idx)
|
|
|
|
async with request_sem:
|
|
await state.track_inflight_start()
|
|
try:
|
|
outcome = await execute_single_turn(
|
|
http_session=http_session,
|
|
conv=conv,
|
|
turn_idx=turn_idx,
|
|
base_url=base_url,
|
|
model=model,
|
|
max_tokens=spec.osl,
|
|
bench_start_ns=bench_start_ns,
|
|
first_chunk_threshold=first_chunk_threshold,
|
|
api_key=api_key,
|
|
)
|
|
except Exception as e:
|
|
logger.error(
|
|
"Session %s turn %d failed: %s",
|
|
conv.session_id,
|
|
turn_idx,
|
|
e,
|
|
)
|
|
current_inflight = await state.track_inflight_end()
|
|
pbar.set_postfix_str(str(current_inflight))
|
|
pbar.update(1)
|
|
await state.mark_session_complete(session_idx, conv.session_id)
|
|
await enqueue_new_session()
|
|
return
|
|
current_inflight = await state.track_inflight_end()
|
|
|
|
await state.record_metric(outcome.metric)
|
|
pbar.set_postfix_str(str(current_inflight))
|
|
pbar.update(1)
|
|
|
|
if turn_idx < spec.num_turns - 1:
|
|
next_turn = turn_idx + 1
|
|
if is_warmup:
|
|
jitter = np.random.uniform(0.0, warmup_jitter_max_s)
|
|
delay_until = time.perf_counter() + jitter
|
|
else:
|
|
delay_until = time.perf_counter() + spec.think_time
|
|
await turn_queue.put((session_idx, next_turn, delay_until))
|
|
else:
|
|
await state.mark_session_complete(session_idx, conv.session_id)
|
|
await enqueue_new_session()
|
|
|
|
await state.check_and_complete_warmup()
|
|
|
|
async def worker() -> None:
|
|
"""Worker task: continuously pull turns from queue and execute them."""
|
|
while True:
|
|
try:
|
|
session_idx, turn_idx, delay_until = await asyncio.wait_for(
|
|
turn_queue.get(), timeout=1.0
|
|
)
|
|
except asyncio.TimeoutError:
|
|
# Do not call get_next_session() here — it allocates a session ID
|
|
# without enqueueing work. Peek completion instead.
|
|
if (
|
|
not await state.has_remaining_sessions()
|
|
and turn_queue.empty()
|
|
and await state.get_inflight() == 0
|
|
):
|
|
break
|
|
continue
|
|
|
|
# Wait until delay_until (for think time/jitter)
|
|
now = time.perf_counter()
|
|
wait_time = delay_until - now
|
|
if wait_time > 0:
|
|
await asyncio.sleep(wait_time)
|
|
|
|
await execute_turn(session_idx, turn_idx)
|
|
|
|
max_entropy = np.log2(spec.num_turns)
|
|
entropy_threshold = 0.50 * max_entropy
|
|
|
|
logger.info(
|
|
"Starting benchmark: %d sessions x %d turns = %d requests, concurrency=%d",
|
|
spec.num_sessions,
|
|
spec.num_turns,
|
|
total_expected,
|
|
spec.concurrency,
|
|
)
|
|
logger.info(
|
|
" Rolling pool: maintaining ~%d active sessions across all turns",
|
|
spec.concurrency,
|
|
)
|
|
logger.info(
|
|
" Warm-up: until steady-state (entropy >= %.3f, 50%% of max=%.3f)",
|
|
entropy_threshold,
|
|
max_entropy,
|
|
)
|
|
logger.info(
|
|
" Measurement: starts after warm-up, using think-time=%.1fs",
|
|
spec.think_time,
|
|
)
|
|
|
|
# Launch worker pool
|
|
workers = [asyncio.create_task(worker()) for _ in range(spec.concurrency)]
|
|
await asyncio.gather(*workers)
|
|
|
|
pbar.close()
|
|
|
|
bench_elapsed_s = (time.perf_counter_ns() - bench_start_ns) / 1e9
|
|
|
|
stats = await state.get_stats()
|
|
logger.info(
|
|
"Benchmark complete: %d measured requests in %.1fs (%.1f req/s), "
|
|
"%d warmup requests discarded",
|
|
stats["measured_metrics"],
|
|
bench_elapsed_s,
|
|
stats["measured_metrics"] / bench_elapsed_s if bench_elapsed_s > 0 else 0,
|
|
stats["warmup_metrics"],
|
|
)
|
|
logger.info(
|
|
"Concurrency: max=%d, avg=%.1f",
|
|
stats["max_inflight"],
|
|
stats["avg_inflight"],
|
|
)
|
|
|
|
return state.measured_metrics
|
|
|
|
|
|
# ============================================================================
|
|
# RATE-BASED BENCHMARK
|
|
# ============================================================================
|
|
|
|
|
|
async def run_rate_based_benchmark(
|
|
spec: WorkloadSpec,
|
|
base_url: str,
|
|
model: str,
|
|
shared_system_text: str,
|
|
text_gen: TextGenerator,
|
|
first_chunk_threshold: int = 16,
|
|
duration_s: float = 0.0,
|
|
warmup_s: float = 0.0,
|
|
api_key: Optional[str] = None,
|
|
) -> tuple[list[TurnMetric], int]:
|
|
"""Closed-loop constant-QPS benchmark with turn-order guarantees.
|
|
|
|
Uses separate asyncio locks for metrics/inflight here (time-based warm-up
|
|
filtering) while :class:`BenchmarkState` powers the concurrency path
|
|
(entropy-based warm-up). Unifying both under one state object is possible
|
|
but is left for a follow-up refactor.
|
|
|
|
Architecture
|
|
============
|
|
- **Ready queue**: sessions whose previous turn has completed push themselves
|
|
here so their next turn can be dispatched.
|
|
- **Pacer**: a single coroutine drains the ready queue at exactly
|
|
``spec.request_rate`` req/s using a token-bucket algorithm. It sleeps
|
|
until the next slot is due, pops one item from the ready queue, and fires
|
|
the turn as a background task.
|
|
- **Lazy session pool**: the pacer holds a pre-computed budget of session
|
|
indices (``ceil(rate * duration / num_turns)``). When the ready queue is
|
|
empty (all active sessions are waiting for server responses), it mints a
|
|
fresh conversation from the budget on the spot -- no upfront allocation.
|
|
If the budget is also exhausted (server extremely slow), it generates
|
|
extra conversations beyond the budget and logs a warning.
|
|
- **Duration-based**: the benchmark runs for ``duration_s`` seconds, then
|
|
the pacer stops and waits for all in-flight requests to complete.
|
|
- **Turn-order guarantee**: turn k+1 is only enqueued after the HTTP
|
|
response for turn k is fully received and injected into history.
|
|
|
|
Session budget
|
|
==============
|
|
Budget = ceil(request_rate * duration_s / num_turns).
|
|
This is the exact number of unique sessions needed if the server keeps up.
|
|
If the server is slower, some in-flight sessions are unavailable and the
|
|
pacer mints extras beyond the budget to maintain QPS.
|
|
"""
|
|
if spec.request_rate is None:
|
|
raise ValueError("rate-based benchmark requires request_rate to be set")
|
|
if warmup_s < 0:
|
|
raise ValueError("warmup_s must be >= 0")
|
|
|
|
bench_start_ns = time.perf_counter_ns()
|
|
|
|
# --------------------------------------------------------------------------
|
|
# Session budget: computed from rate * duration / num_turns
|
|
# --------------------------------------------------------------------------
|
|
if duration_s > 0:
|
|
session_budget = math.ceil(spec.request_rate * duration_s / spec.num_turns)
|
|
elif spec.num_sessions is not None:
|
|
session_budget = spec.num_sessions
|
|
else:
|
|
session_budget = 1000 # fallback; shouldn't reach here after CLI validation
|
|
|
|
# Monotonically increasing session index (never resets; unique per conversation)
|
|
next_session_idx = 0
|
|
sessions_beyond_budget = 0 # count of extra sessions created due to server lag
|
|
|
|
def _next_conv() -> Conversation:
|
|
"""Mint a fresh conversation, lazily, from the session budget (or beyond)."""
|
|
nonlocal next_session_idx, sessions_beyond_budget
|
|
idx = next_session_idx
|
|
next_session_idx += 1
|
|
if idx >= session_budget:
|
|
sessions_beyond_budget += 1
|
|
return conversation_factory(idx, spec, shared_system_text, text_gen)
|
|
|
|
# --------------------------------------------------------------------------
|
|
# Ready queue: (conv, turn_idx) -- only sessions whose prior turn completed
|
|
# --------------------------------------------------------------------------
|
|
ready_queue: asyncio.Queue[tuple[Conversation, int]] = asyncio.Queue()
|
|
|
|
# --------------------------------------------------------------------------
|
|
# Shared counters (only mutated by execute_turn background tasks)
|
|
# --------------------------------------------------------------------------
|
|
metrics: list[TurnMetric] = []
|
|
warmup_metrics: list[TurnMetric] = []
|
|
metrics_lock = asyncio.Lock()
|
|
|
|
inflight = 0
|
|
max_inflight = 0
|
|
inflight_samples: list[int] = []
|
|
inflight_lock = asyncio.Lock()
|
|
|
|
stop_dispatching = asyncio.Event()
|
|
all_done = asyncio.Event()
|
|
|
|
total_slots = int(spec.request_rate * duration_s) if duration_s > 0 else 0
|
|
|
|
# Two stacked progress bars: sent (top) chased by done (bottom).
|
|
_bar_fmt_sent = (
|
|
"{desc}: {bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]"
|
|
)
|
|
_bar_fmt_done = (
|
|
"{desc}: {bar}| {n_fmt}/{total_fmt} "
|
|
"[{elapsed}<{remaining}, {rate_fmt}] {postfix}"
|
|
)
|
|
_bar_fmt_no_total = "{desc}: {bar}| {n_fmt} [{elapsed}, {rate_fmt}]"
|
|
_bar_fmt_no_total_pf = "{desc}: {bar}| {n_fmt} [{elapsed}, {rate_fmt}] {postfix}"
|
|
|
|
pbar_sent = tqdm(
|
|
total=total_slots if total_slots > 0 else None,
|
|
desc=" Sent",
|
|
unit="req",
|
|
bar_format=_bar_fmt_sent if total_slots > 0 else _bar_fmt_no_total,
|
|
position=1,
|
|
leave=True,
|
|
file=sys.stderr,
|
|
colour="cyan",
|
|
)
|
|
pbar_done = tqdm(
|
|
total=total_slots if total_slots > 0 else None,
|
|
desc=" Done",
|
|
unit="req",
|
|
bar_format=_bar_fmt_done if total_slots > 0 else _bar_fmt_no_total_pf,
|
|
position=0,
|
|
leave=True,
|
|
file=sys.stderr,
|
|
colour="green",
|
|
)
|
|
pbar_done.set_postfix_str("inflight=0")
|
|
|
|
logger.info(
|
|
"Closed-loop QPS benchmark: target=%.2f req/s, duration=%.1fs, "
|
|
"num_turns=%d, session_budget=%d (auto from rate*duration/turns)",
|
|
spec.request_rate,
|
|
duration_s,
|
|
spec.num_turns,
|
|
session_budget,
|
|
)
|
|
logger.info(
|
|
"Stats warm-up: %.1fs (requests dispatched before this are excluded "
|
|
"from reported metrics)",
|
|
warmup_s,
|
|
)
|
|
|
|
warmup_cutoff_ns = bench_start_ns + int(warmup_s * 1e9)
|
|
|
|
# --------------------------------------------------------------------------
|
|
# Per-turn execution (background task, one per in-flight request)
|
|
# --------------------------------------------------------------------------
|
|
async def execute_turn(
|
|
conv: Conversation,
|
|
turn_idx: int,
|
|
http_session: aiohttp.ClientSession,
|
|
) -> None:
|
|
nonlocal inflight, max_inflight
|
|
|
|
async with inflight_lock:
|
|
inflight += 1
|
|
if inflight > max_inflight:
|
|
max_inflight = inflight
|
|
inflight_samples.append(inflight)
|
|
pbar_done.set_postfix_str(f"inflight={inflight}")
|
|
|
|
try:
|
|
outcome = await execute_single_turn(
|
|
http_session=http_session,
|
|
conv=conv,
|
|
turn_idx=turn_idx,
|
|
base_url=base_url,
|
|
model=model,
|
|
max_tokens=spec.osl,
|
|
bench_start_ns=bench_start_ns,
|
|
first_chunk_threshold=first_chunk_threshold,
|
|
api_key=api_key,
|
|
)
|
|
except Exception as e:
|
|
logger.error(
|
|
"Session %s turn %d failed: %s",
|
|
conv.session_id,
|
|
turn_idx,
|
|
e,
|
|
)
|
|
async with inflight_lock:
|
|
inflight -= 1
|
|
pbar_done.set_postfix_str(f"inflight={inflight}")
|
|
pbar_done.update(1)
|
|
if not stop_dispatching.is_set():
|
|
await ready_queue.put((_next_conv(), 0))
|
|
async with inflight_lock:
|
|
if stop_dispatching.is_set() and inflight == 0:
|
|
all_done.set()
|
|
return
|
|
|
|
async with inflight_lock:
|
|
inflight -= 1
|
|
pbar_done.set_postfix_str(f"inflight={inflight}")
|
|
|
|
async with metrics_lock:
|
|
req_start_ns = bench_start_ns + int(outcome.metric.start_time_ms * 1e6)
|
|
if req_start_ns >= warmup_cutoff_ns:
|
|
metrics.append(outcome.metric)
|
|
else:
|
|
warmup_metrics.append(outcome.metric)
|
|
pbar_done.update(1)
|
|
|
|
if not stop_dispatching.is_set():
|
|
next_turn = turn_idx + 1
|
|
if next_turn < spec.num_turns:
|
|
await ready_queue.put((conv, next_turn))
|
|
else:
|
|
await ready_queue.put((_next_conv(), 0))
|
|
|
|
async with inflight_lock:
|
|
if stop_dispatching.is_set() and inflight == 0:
|
|
all_done.set()
|
|
|
|
# --------------------------------------------------------------------------
|
|
# Pacer: token-bucket, single coroutine, no lock needed
|
|
# --------------------------------------------------------------------------
|
|
async def pacer(http_session: aiohttp.ClientSession) -> None:
|
|
"""Dispatch one request per slot at exactly request_rate req/s."""
|
|
nonlocal next_session_idx, sessions_beyond_budget
|
|
interval_s = 1.0 / spec.request_rate
|
|
next_dispatch = time.perf_counter()
|
|
|
|
# Throttle "server lagging" log to once per second
|
|
_last_lag_log: float = 0.0
|
|
_lag_count_since_log: int = 0
|
|
|
|
while not stop_dispatching.is_set():
|
|
now = time.perf_counter()
|
|
wait = next_dispatch - now
|
|
if wait > 0:
|
|
await asyncio.sleep(wait)
|
|
|
|
if stop_dispatching.is_set():
|
|
break
|
|
|
|
# Prefer a session whose prior turn already completed (ready queue).
|
|
# If empty the server is lagging -- mint a fresh session from the
|
|
# budget (or beyond) so QPS is maintained.
|
|
try:
|
|
conv, turn_idx = ready_queue.get_nowait()
|
|
# Ready queue had something -- server is keeping up at this slot
|
|
pbar_sent.set_description(" Sent")
|
|
except asyncio.QueueEmpty:
|
|
_lag_count_since_log += 1
|
|
now = time.perf_counter()
|
|
beyond_budget = next_session_idx >= session_budget
|
|
# Update sent bar description to show lag state continuously
|
|
if beyond_budget:
|
|
pbar_sent.set_description(
|
|
f" Sent [lag+{sessions_beyond_budget + _lag_count_since_log}]"
|
|
)
|
|
else:
|
|
pbar_sent.set_description(
|
|
f" Sent [lag {next_session_idx}/{session_budget}]"
|
|
)
|
|
if now - _last_lag_log >= 1.0:
|
|
elapsed = (time.perf_counter_ns() - bench_start_ns) / 1e9
|
|
if beyond_budget:
|
|
logger.warning(
|
|
"Server lagging at %.1fs: ready queue empty, budget exhausted "
|
|
"(budget=%d). Minted %d extra session(s) in last %.1fs to hold QPS.",
|
|
elapsed,
|
|
session_budget,
|
|
_lag_count_since_log,
|
|
now - _last_lag_log if _last_lag_log else elapsed,
|
|
)
|
|
else:
|
|
logger.info(
|
|
"Server lagging at %.1fs: ready queue empty, pulling from budget "
|
|
"(%d/%d sessions used). %d slot(s) filled from budget in last %.1fs.",
|
|
elapsed,
|
|
next_session_idx,
|
|
session_budget,
|
|
_lag_count_since_log,
|
|
now - _last_lag_log if _last_lag_log else elapsed,
|
|
)
|
|
_last_lag_log = now
|
|
_lag_count_since_log = 0
|
|
conv = _next_conv()
|
|
turn_idx = 0
|
|
|
|
asyncio.create_task(execute_turn(conv, turn_idx, http_session))
|
|
pbar_sent.update(1)
|
|
next_dispatch += interval_s
|
|
|
|
# Pacer finished -- trigger all_done if nothing is in-flight
|
|
async with inflight_lock:
|
|
if inflight == 0:
|
|
all_done.set()
|
|
|
|
# --------------------------------------------------------------------------
|
|
# Duration timer
|
|
# --------------------------------------------------------------------------
|
|
async def duration_timer() -> None:
|
|
await asyncio.sleep(duration_s)
|
|
logger.info("Duration %.1fs elapsed -- stopping new dispatches.", duration_s)
|
|
stop_dispatching.set()
|
|
|
|
# --------------------------------------------------------------------------
|
|
# Main execution
|
|
# --------------------------------------------------------------------------
|
|
connector = aiohttp.TCPConnector(limit=0)
|
|
async with aiohttp.ClientSession(connector=connector) as http_session:
|
|
# Seed the ready queue with enough turn-0 conversations to keep the
|
|
# pacer busy for the first second without generating them mid-sleep.
|
|
seed_count = max(1, int(spec.request_rate) + 1)
|
|
for _ in range(seed_count):
|
|
await ready_queue.put((_next_conv(), 0))
|
|
|
|
timer_task = asyncio.create_task(duration_timer())
|
|
pacer_task = asyncio.create_task(pacer(http_session))
|
|
|
|
await timer_task
|
|
await pacer_task
|
|
|
|
pbar_sent.set_description(" Sent")
|
|
logger.info("Waiting for %d in-flight requests to complete...", inflight)
|
|
await all_done.wait()
|
|
|
|
pbar_sent.close()
|
|
pbar_done.close()
|
|
|
|
bench_elapsed_s = (time.perf_counter_ns() - bench_start_ns) / 1e9
|
|
avg_inf = sum(inflight_samples) / len(inflight_samples) if inflight_samples else 0.0
|
|
logger.info(
|
|
"Benchmark complete: %d requests in %.1fs (%.2f req/s actual); "
|
|
"max_inflight=%d, avg_inflight=%.1f",
|
|
len(metrics),
|
|
bench_elapsed_s,
|
|
len(metrics) / bench_elapsed_s if bench_elapsed_s > 0 else 0,
|
|
max_inflight,
|
|
avg_inf,
|
|
)
|
|
if warmup_metrics:
|
|
logger.info(
|
|
"Warm-up filtered out %d requests (start_time < %.1fs)",
|
|
len(warmup_metrics),
|
|
warmup_s,
|
|
)
|
|
logger.info(
|
|
"Sessions: budget=%d, used=%d (%s)",
|
|
session_budget,
|
|
next_session_idx,
|
|
f"{sessions_beyond_budget} extra beyond budget -- server was lagging"
|
|
if sessions_beyond_budget > 0
|
|
else "within budget",
|
|
)
|
|
return metrics, len(warmup_metrics)
|
|
|
|
|
|
# ============================================================================
|
|
# SEED & ENTRY POINTS
|
|
# ============================================================================
|
|
|
|
|
|
def set_seed(args: argparse.Namespace) -> int:
|
|
"""Set random seed combining base seed with workload spec hash.
|
|
|
|
This ensures:
|
|
- Reproducibility: same configuration always generates same conversations
|
|
- Variation: different configurations (concurrency, rate, ISL, etc.)
|
|
generate different conversations
|
|
|
|
Args:
|
|
args: Parsed command-line arguments.
|
|
|
|
Returns:
|
|
The combined seed value used.
|
|
"""
|
|
# Generate random base seed if not provided
|
|
if args.seed is None:
|
|
base_seed = random.randint(0, 2**31 - 1)
|
|
logger.info("No seed provided, generated random base seed: %d", base_seed)
|
|
else:
|
|
base_seed = args.seed
|
|
|
|
# Get all args but exclude parameters that don't affect workload shape
|
|
# (e.g., output paths, server URLs, infrastructure settings)
|
|
exclude_params = {
|
|
"seed", # Base seed (used separately)
|
|
"save_result", # Output path doesn't affect workload
|
|
"base_url", # Server URL doesn't affect workload
|
|
"model", # Model name doesn't affect workload shape
|
|
"tokenizer", # Tokenizer path doesn't affect workload (token counts do)
|
|
"first_chunk_threshold", # Measurement detail, not workload shape
|
|
"warmup_jitter_max", # Warm-up pacing only; does not change generated text
|
|
}
|
|
|
|
# Include all workload-affecting parameters
|
|
workload_params = {k: v for k, v in vars(args).items() if k not in exclude_params}
|
|
|
|
# Create stable string representation and hash it
|
|
workload_str = json.dumps(workload_params, sort_keys=True, default=str)
|
|
workload_hash = int(hashlib.sha256(workload_str.encode()).hexdigest()[:8], 16)
|
|
|
|
# Combine base seed with workload hash; clamp to numpy's valid range [0, 2**32-1]
|
|
combined_seed = (base_seed + workload_hash) % (2**32)
|
|
|
|
# Set seeds
|
|
random.seed(combined_seed)
|
|
np.random.seed(combined_seed)
|
|
|
|
logger.info(
|
|
"Random seed: base=%d, workload_hash=%d, combined=%d",
|
|
base_seed,
|
|
workload_hash,
|
|
combined_seed,
|
|
)
|
|
|
|
return combined_seed
|
|
|
|
|
|
async def _run_smoke_async(args) -> dict:
|
|
"""Run a single smoke-test request and return metrics."""
|
|
messages = [{"role": "user", "content": "Say 'this is a test' and nothing else."}]
|
|
async with aiohttp.ClientSession() as session:
|
|
result = await send_chat_completion(
|
|
session=session,
|
|
base_url=args.base_url,
|
|
model=args.model,
|
|
messages=messages,
|
|
first_chunk_threshold=args.first_chunk_threshold,
|
|
api_key=getattr(args, "api_key", None),
|
|
)
|
|
return {
|
|
"ttft_ms": round(result.ttft_ms, 2),
|
|
"fc_ms": round(result.fc_ms, 2),
|
|
"itl_ms": round(result.itl_ms, 2),
|
|
"e2e_latency_ms": round(result.e2e_latency_ms, 2),
|
|
"input_tokens": result.input_tokens,
|
|
"output_tokens": result.output_tokens,
|
|
}
|
|
|
|
|
|
def run_smoke(args) -> int:
|
|
"""Run smoke test synchronously, return exit code."""
|
|
try:
|
|
result = asyncio.run(_run_smoke_async(args))
|
|
print(json.dumps(result, indent=2))
|
|
return 0
|
|
except Exception as e:
|
|
logger.error(f"Smoke test failed: {e}")
|
|
return 1
|
|
|
|
|
|
def run_direct(args) -> int:
|
|
"""Run a single benchmark (concurrency or rate-based) and report results."""
|
|
# Resolve tokenizer
|
|
tokenizer_name = args.tokenizer if args.tokenizer else args.model
|
|
try:
|
|
from transformers import AutoTokenizer
|
|
|
|
logger.info("Loading tokenizer: %s", tokenizer_name)
|
|
tokenizer = AutoTokenizer.from_pretrained(
|
|
tokenizer_name, trust_remote_code=True
|
|
)
|
|
except Exception as e:
|
|
logger.error(
|
|
"Failed to load tokenizer '%s'. If your --model is not a valid "
|
|
"HuggingFace repo, provide --tokenizer explicitly. Error: %s",
|
|
tokenizer_name,
|
|
e,
|
|
)
|
|
return 1
|
|
|
|
set_seed(args)
|
|
text_gen = TextGenerator(tokenizer)
|
|
|
|
# Validate
|
|
if args.concurrency is not None and args.num_sessions is None:
|
|
logger.error("--num-sessions is required when using --concurrency mode.")
|
|
return 1
|
|
if (
|
|
args.request_rate is not None
|
|
and args.num_sessions is None
|
|
and (args.duration is None or args.duration <= 0)
|
|
):
|
|
logger.error(
|
|
"For --request-rate mode without --num-sessions, specify --duration > 0."
|
|
)
|
|
return 1
|
|
|
|
# Build spec (simple mode only)
|
|
# Filter out None values so dataclass defaults (e.g. num_turns=1, osl=1)
|
|
# are preserved when the user omits those CLI flags.
|
|
spec_kwargs = dict(
|
|
num_sessions=args.num_sessions,
|
|
num_turns=args.num_turns,
|
|
osl=args.osl,
|
|
think_time=args.think_time,
|
|
concurrency=args.concurrency,
|
|
request_rate=args.request_rate,
|
|
ramp_interval=args.ramp_interval,
|
|
shared_system_prompt_ratio=args.shared_system_prompt_ratio,
|
|
isl=args.isl,
|
|
hit_rate=args.hit_rate,
|
|
duration_s=args.duration or 0.0,
|
|
)
|
|
spec = WorkloadSpec(**{k: v for k, v in spec_kwargs.items() if v is not None})
|
|
try:
|
|
spec.resolve()
|
|
except ValueError as e:
|
|
logger.error("Workload spec error: %s", e)
|
|
return 1
|
|
|
|
spec.print_summary()
|
|
|
|
shared_system_text = text_gen.generate(spec.shared_s)
|
|
|
|
bench_start = time.perf_counter()
|
|
discarded_warmup_requests = 0
|
|
report_elapsed = None
|
|
|
|
api_key = getattr(args, "api_key", None)
|
|
fc_thresh = args.first_chunk_threshold
|
|
warmup_jitter_max = getattr(args, "warmup_jitter_max", 10.0)
|
|
|
|
if spec.request_rate is not None:
|
|
if spec.duration_s > 0:
|
|
effective_duration = spec.duration_s
|
|
else:
|
|
# _validate() ensures num_sessions is set when duration_s <= 0
|
|
assert spec.num_sessions is not None
|
|
effective_duration = spec.num_sessions * spec.num_turns / spec.request_rate
|
|
if args.warm_up >= effective_duration:
|
|
logger.error(
|
|
"--warm-up (%.3fs) must be less than effective duration (%.3fs).",
|
|
args.warm_up,
|
|
effective_duration,
|
|
)
|
|
return 1
|
|
metrics, discarded_warmup_requests = asyncio.run(
|
|
run_rate_based_benchmark(
|
|
spec,
|
|
args.base_url,
|
|
args.model,
|
|
shared_system_text,
|
|
text_gen,
|
|
first_chunk_threshold=fc_thresh,
|
|
duration_s=effective_duration,
|
|
warmup_s=args.warm_up,
|
|
api_key=api_key,
|
|
)
|
|
)
|
|
report_elapsed = max(1e-9, effective_duration - args.warm_up)
|
|
else:
|
|
metrics = asyncio.run(
|
|
run_benchmark(
|
|
spec,
|
|
args.base_url,
|
|
args.model,
|
|
shared_system_text,
|
|
text_gen,
|
|
first_chunk_threshold=fc_thresh,
|
|
api_key=api_key,
|
|
warmup_jitter_max_s=warmup_jitter_max,
|
|
)
|
|
)
|
|
|
|
bench_elapsed = time.perf_counter() - bench_start
|
|
if report_elapsed is None:
|
|
report_elapsed = bench_elapsed
|
|
|
|
report_results(
|
|
metrics,
|
|
spec,
|
|
report_elapsed,
|
|
first_chunk_threshold=fc_thresh,
|
|
save_path=args.save_result,
|
|
warmup_s=args.warm_up if spec.request_rate is not None else 0.0,
|
|
discarded_warmup_requests=discarded_warmup_requests,
|
|
)
|
|
return 0
|