concurrency_benchmark.py

"""Concurrency sweep: same request shape at increasing simultaneous in-flight counts.

Each level runs two waves against one model:

* prefill probe  - ``level`` simultaneous requests with max_tokens=1. Computed prompt
  tokens / probe wall time is the endpoint's *aggregate* prefill throughput at that
  concurrency (wall includes queueing, transport and one sampled token, so it slightly
  understates pure prefill).
* generation     - ``requests_per_level`` requests, never more than ``level`` in flight.
  Output tokens / wall is aggregate generation throughput; per-request decode speed is
  the per-stream client estimate (or vLLM per-request decode time when metrics are clean).

Server metrics are attributed only when the endpoint is idle at both boundaries, the
process/engine is unchanged, the completed-request count equals ours and token usage
matches exactly. Otherwise phase rates are reported unavailable with a reason; TTFT is
never substituted for prefill.
"""
import asyncio
import math
import random
import time

DEFAULT_LEVELS = [1, 2, 4, 8]
MAX_CONCURRENCY = 64
MAX_REQUESTS = 1024
MAX_LEVELS = 8

_WORDS = ('cache', 'node', 'shard', 'packet', 'ledger', 'socket', 'kernel', 'buffer', 'tensor', 'queue',
          'replica', 'stream', 'driver', 'vector', 'thread', 'signal', 'border', 'harbor', 'meadow', 'copper',
          'lantern', 'orbit', 'canyon', 'falcon', 'ember', 'glacier', 'violet', 'summit', 'river', 'pixel',
          'cobalt', 'marble', 'quartz', 'timber', 'saddle', 'beacon', 'anchor', 'circuit', 'mosaic', 'prism')


def _int(v, lo, hi):
    return isinstance(v, int) and not isinstance(v, bool) and lo <= v <= hi


def validate_concurrency(data):
    levels = data.get('concurrency_levels', DEFAULT_LEVELS)
    count = data.get('requests_per_level', 16)
    if (not isinstance(levels, list) or not 1 <= len(levels) <= MAX_LEVELS
            or any(not _int(x, 1, MAX_CONCURRENCY) for x in levels) or len(set(levels)) != len(levels)):
        raise ValueError(f'Invalid concurrency levels (1-{MAX_LEVELS} unique integers from 1 to {MAX_CONCURRENCY})')
    if not _int(count, 1, MAX_REQUESTS) or count < max(levels):
        raise ValueError(f'requests_per_level must be at least the largest concurrency level and at most {MAX_REQUESTS}')
    return {'concurrency_levels': sorted(levels), 'requests_per_level': count}


def planned_calls(levels, requests_per_level):
    return sum(level + requests_per_level for level in levels)


def build_messages(context_size, seed, run_tag, phase, level, index, context_tag=None):
    """Deterministic prompt of ~context_size tokens (4 chars/token estimate).

    The leading line is unique per run/phase/context/level/request so prefix caching
    cannot share KV between requests (or between the probe and generation waves).
    """
    rng = random.Random(f'{seed}:{run_tag}:{phase}:{context_tag}:{level}:{index}')
    band = f'-{context_tag}' if context_tag else ''
    head = f'Load request {run_tag}-{phase}{band}-c{level}-r{index:04d} nonce {rng.getrandbits(64):016x}.\n'
    task = ('\nTask: write a long, continuous technical maintenance narrative about the records above. '
            'Do not stop early; keep writing until you run out of space.')
    target = max(len(head) + len(task) + 16, int(context_size) * 4)
    parts = []; size = len(head) + len(task)
    while size < target:
        line = f'Record {len(parts):05d}: ' + ' '.join(rng.choice(_WORDS) for _ in range(10)) + '.\n'
        parts.append(line); size += len(line)
    body = ''.join(parts)[:max(0, target - len(head) - len(task))]
    return [{'role': 'user', 'content': head + body + task}]


async def run_wave(call, total, concurrency):
    """Run ``call(i)`` for i in range(total) with at most ``concurrency`` in flight.

    Returns per-request outcomes (exceptions captured, not raised; CancelledError propagates),
    the observed maximum in-flight count and wall seconds from first start to last finish.
    """
    sem = asyncio.Semaphore(concurrency)
    state = {'in_flight': 0, 'max': 0}
    results = [None] * total

    async def one(i):
        async with sem:
            state['in_flight'] += 1; state['max'] = max(state['max'], state['in_flight'])
            started = time.monotonic()
            try:
                value = await call(i)
                results[i] = {'index': i, 'value': value, 'error': None, 'error_detail': None,
                              'started': started, 'finished': time.monotonic()}
            except asyncio.CancelledError:
                raise
            except Exception as e:  # recorded per request; batch continues
                results[i] = {'index': i, 'value': None, 'error': type(e).__name__, 'error_detail': str(e)[:500],
                              'started': started, 'finished': time.monotonic()}
            finally:
                state['in_flight'] -= 1

    t0 = time.monotonic()
    await asyncio.gather(*(one(i) for i in range(total)))
    return {'results': results, 'max_in_flight': state['max'], 'wall_s': time.monotonic() - t0}


_DELTA_KEYS = ('request_success_total', 'request_prompt_tokens_sum', 'request_generation_tokens_sum',
               'request_prefill_time_seconds_sum', 'request_decode_time_seconds_sum',
               'request_prefill_kv_computed_tokens_sum', 'request_prompt_tokens_count',
               'request_generation_tokens_count', 'request_prefill_time_seconds_count',
               'request_decode_time_seconds_count', 'request_prefill_kv_computed_tokens_count')


def batch_measurement(before, after, count, prompt_tokens, completion_tokens):
    """Attribute vLLM counter deltas to exactly this batch, or explain why not."""
    if not before or not after:
        return None, 'server metrics unavailable (endpoint exposes no parseable vLLM /metrics)'
    if before['engine'] != after['engine'] or before['process_start'] != after['process_start']:
        return None, 'engine restart between snapshots'
    if any(s['values'][k] != 0 for s in (before, after) for k in ('num_requests_running', 'num_requests_waiting')):
        return None, 'endpoint not idle at a batch boundary (other traffic)'
    delta = {k: after['values'][k] - before['values'][k] for k in _DELTA_KEYS}
    if any(v < 0 for v in delta.values()):
        return None, 'counter reset between snapshots'
    if delta['request_success_total'] != count or any(delta[k] != count for k in _DELTA_KEYS if k.endswith('_count')):
        return None, f"completed-request count mismatch (server {delta['request_success_total']:g}, ours {count}); foreign traffic or failed requests"
    if not isinstance(prompt_tokens, int) or not isinstance(completion_tokens, int):
        return None, 'client token usage missing'
    if delta['request_prompt_tokens_sum'] != prompt_tokens or delta['request_generation_tokens_sum'] != completion_tokens:
        return None, 'server token totals do not match client usage'
    computed = delta['request_prefill_kv_computed_tokens_sum']; pre_s = delta['request_prefill_time_seconds_sum']
    dec_s = delta['request_decode_time_seconds_sum']; decode_tokens = completion_tokens - count
    if not 0 < computed <= prompt_tokens or pre_s <= 0:
        return None, 'no computed prefill tokens or prefill time recorded'
    return {'prefill_tps': computed / pre_s, 'prefill_seconds_sum': pre_s,
            'decode_tps': decode_tokens / dec_s if dec_s > 0 and decode_tokens > 0 else None,
            'decode_seconds_sum': dec_s, 'computed_prompt_tokens': int(computed),
            'cached_tokens': int(prompt_tokens - computed), 'source': 'vllm_batch_metrics',
            'raw_metrics_delta': delta}, None


def _pct(values, p):
    if not values:
        return None
    vals = sorted(values)
    return vals[min(len(vals) - 1, max(0, math.ceil(p * len(vals)) - 1))]


def _num(v):
    return isinstance(v, (int, float)) and not isinstance(v, bool)


def _total(items, key):
    vals = [x.get(key) for x in items]
    return sum(vals) if vals and all(_num(v) for v in vals) else None


def summarize_level(concurrency, gen_results, gen_wall_s, gen_max_in_flight, gen_metrics, gen_reason,
                    probe_results, probe_wall_s, probe_max_in_flight, probe_metrics, probe_reason):
    done = [r for r in gen_results if r.get('status') == 'completed']
    calls = [c for r in done for c in r.get('calls', [])]
    output, prompt = _total(done, 'completion_tokens'), _total(done, 'prompt_tokens')
    usage_cached = _total(calls, 'cached_tokens')
    ttfts = [c['ttft_s'] for c in calls if _num(c.get('ttft_s'))]
    lats = [r['latency_s'] for r in done if _num(r.get('latency_s'))]
    dec = [c for c in calls if c.get('decode_source') and _num(c.get('decode_seconds')) and c['decode_seconds'] > 0 and _num(c.get('completion_tokens'))]
    dec_tokens = sum(c['completion_tokens'] - (1 if c['decode_source'] == 'client_observed_estimate' else 0) for c in dec)
    decode = dec_tokens / sum(c['decode_seconds'] for c in dec) if dec else None
    dec_sources = sorted({c['decode_source'] for c in dec})

    p_done = [r for r in probe_results if r.get('status') == 'completed']
    p_calls = [c for r in p_done for c in r.get('calls', [])]
    p_prompt = _total(p_done, 'prompt_tokens'); p_cached = _total(p_calls, 'cached_tokens')
    p_ttft = [c['ttft_s'] for c in p_calls if _num(c.get('ttft_s'))]
    agg, agg_src, agg_reason = None, None, None
    if not probe_wall_s or probe_wall_s <= 0 or len(p_done) != len(probe_results) or not probe_results:
        agg_reason = 'prefill probe incomplete (errors or no completed requests)'
    elif probe_metrics:
        agg, agg_src = probe_metrics['computed_prompt_tokens'] / probe_wall_s, 'vllm_batch_metrics'
    elif p_prompt is not None and p_cached is not None:
        agg, agg_src = (p_prompt - p_cached) / probe_wall_s, 'client_wall_usage'
    elif p_prompt is not None:
        agg, agg_src = p_prompt / probe_wall_s, 'client_wall_prompt_tokens_cache_unknown'
    else:
        agg_reason = 'probe token usage missing'

    wall_ok = _num(gen_wall_s) and gen_wall_s > 0
    return {
        'concurrency': concurrency, 'requests': len(gen_results), 'completed': len(done),
        'errors': len(gen_results) - len(done), 'wall_s': gen_wall_s,
        'observed_max_in_flight': gen_max_in_flight,
        'output_tokens': output, 'prompt_tokens': prompt,
        'cached_tokens': gen_metrics['cached_tokens'] if gen_metrics else usage_cached,
        'output_tps': output / gen_wall_s if output is not None and wall_ok else None,
        'input_tps': prompt / gen_wall_s if prompt is not None and wall_ok else None,
        'decode_tps': decode, 'decode_source': '+'.join(dec_sources) if dec_sources else None,
        'measured_decode_requests': len(dec),
        'server_decode_tps': gen_metrics['decode_tps'] if gen_metrics else None,
        'prefill_tps': gen_metrics['prefill_tps'] if gen_metrics else None,
        'prefill_source': gen_metrics['source'] if gen_metrics else None,
        'prefill_unavailable_reason': None if gen_metrics else gen_reason,
        'ttft_p50_s': _pct(ttfts, .5), 'ttft_p95_s': _pct(ttfts, .95),
        'latency_p50_s': _pct(lats, .5), 'latency_p95_s': _pct(lats, .95),
        'aggregate_prefill_tps': agg, 'aggregate_prefill_source': agg_src,
        'aggregate_prefill_unavailable_reason': agg_reason,
        'probe': {'requests': len(probe_results), 'completed': len(p_done), 'wall_s': probe_wall_s,
                  'observed_max_in_flight': probe_max_in_flight, 'prompt_tokens': p_prompt,
                  'cached_tokens': probe_metrics['cached_tokens'] if probe_metrics else p_cached,
                  'computed_prompt_tokens': probe_metrics['computed_prompt_tokens'] if probe_metrics else None,
                  'per_request_prefill_tps': probe_metrics['prefill_tps'] if probe_metrics else None,
                  'ttft_p50_s': _pct(p_ttft, .5), 'ttft_p95_s': _pct(p_ttft, .95),
                  'metrics_unavailable_reason': None if probe_metrics else probe_reason,
                  'raw_metrics_delta': probe_metrics['raw_metrics_delta'] if probe_metrics else None},
        'raw_metrics_delta': gen_metrics['raw_metrics_delta'] if gen_metrics else None,
    }