poolprobe 1.0.0__py3-none-any.whl

poolprobe/bench.py

@@ -0,0 +1,368 @@
+"""Latency and throughput, with the measurement error stated instead of hidden.
+
+A benchmark that derives tokens/second from SSE chunk counts is lying: chunk
+boundaries are a transport artifact, not a token count. One live route (a
+Claude-class model behind a self-hosted gateway) produced five chunks for a
+complete answer and was reported at 2.2 tok/s while answering perfectly.
+
+So every result records *where its token number came from* and refuses to
+extrapolate:
+
+``usage``
+    ``usage.completion_tokens`` from the reply. Exact. The only source that may
+    be turned into a tok/s figure without a caveat.
+``chunks``
+    Streaming with no ``usage`` block. The count is a hard lower bound and is
+    flagged as such everywhere it is displayed.
+``none``
+    No evidence. ``tok_s`` stays ``0.0`` rather than being guessed, because a
+    plausible-looking wrong number gets copied into someone's model ranking.
+"""
+
+from __future__ import annotations
+
+import json
+import statistics
+import time
+import urllib.error
+import urllib.request
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass
+
+from . import taxonomy
+from .probe import (
+    DEFAULT_TIMEOUT,
+    Gateway,
+    _flatten,
+    parallel_map,
+)
+
+#: A deterministic amount of work. "Say hello" measures the gateway's queue, not
+#: the model; a fixed enumeration gives every route the same job and produces a
+#: token count large enough that timing noise stays in the second decimal.
+DEFAULT_BENCH_PROMPT = "Count from 1 to 60, one number per line."
+
+DEFAULT_BENCH_MAX_TOKENS = 256
+
+#: Deliberately lower than the sweep default. Parallel requests contend inside
+#: the gateway, so a wide bench measures its own congestion; 4 keeps queueing
+#: noise small while still finishing a 60-model pool in reasonable time.
+DEFAULT_THREADS = 4
+
+
+@dataclass
+class BenchResult:
+    """One measured model. ``token_source`` says how much to trust ``tok_s``."""
+
+    model: str
+    provider: str = ""
+    ok: bool = False
+    status: int = 0
+    latency_s: float = 0.0
+    ttft_s: float = 0.0
+    tokens: int = 0
+    tok_s: float = 0.0
+    token_source: str = "none"  # usage | chunks | none
+    verdict: str = ""
+    kind: str = ""
+    error: str = ""
+
+    @property
+    def tok_s_exact(self) -> bool:
+        """True only when the token count came from the gateway's own usage."""
+        return self.token_source == "usage"
+
+    def as_dict(self) -> dict[str, object]:
+        return {
+            "model": self.model,
+            "provider": self.provider,
+            "ok": self.ok,
+            "status": self.status,
+            "latency_s": round(self.latency_s, 2),
+            "ttft_s": round(self.ttft_s, 2),
+            "tokens": self.tokens,
+            "tok_s": round(self.tok_s, 2),
+            "token_source": self.token_source,
+            "tok_s_exact": self.tok_s_exact,
+            "verdict": self.verdict,
+            "kind": self.kind,
+            "error": self.error,
+        }
+
+
+def bench_one(
+    gateway: Gateway,
+    model: str,
+    *,
+    prompt: str = DEFAULT_BENCH_PROMPT,
+    max_tokens: int = DEFAULT_BENCH_MAX_TOKENS,
+    timeout: float = DEFAULT_TIMEOUT,
+    stream: bool = False,
+    provider: str = "",
+    repeat: int = 1,
+) -> BenchResult:
+    """Measure one model, optionally ``repeat`` times, keeping the best truth.
+
+    With ``repeat > 1`` the reported latency is the **minimum** and the reported
+    throughput the **median** of the successful samples: the minimum is the
+    closest observation to an unloaded route, the median resists one lucky sample
+    inflating the ranking. Failed samples are never averaged into a success — a
+    dead attempt is not a slow one.
+    """
+    repeat = max(1, int(repeat))
+    samples = [
+        _sample(
+            gateway,
+            model,
+            prompt=prompt,
+            max_tokens=max_tokens,
+            timeout=timeout,
+            stream=stream,
+            provider=provider,
+        )
+        for _ in range(repeat)
+    ]
+    good = [sample for sample in samples if sample.ok]
+    if not good:
+        # A dead attempt is not a slow one: report the failure as measured,
+        # never averaged into a latency.
+        return samples[-1]
+
+    # Exact beats approximate: if any sample carried usage, the route does emit
+    # it, and that is the only source worth ranking on.
+    source = (
+        "usage"
+        if any(sample.token_source == "usage" for sample in good)
+        else "chunks"
+        if any(sample.token_source == "chunks" for sample in good)
+        else "none"
+    )
+    # Median over samples sharing the winning source. Folding chunk-counted and
+    # usage-counted throughputs into one median yields a number that describes
+    # neither measurement.
+    comparable = [sample for sample in good if sample.token_source == source] or good
+    throughputs = [sample.tok_s for sample in comparable if sample.tok_s > 0]
+    firsts = [sample.ttft_s for sample in good if sample.ttft_s > 0]
+
+    winner = min(good, key=lambda sample: sample.latency_s)
+    winner.latency_s = min(sample.latency_s for sample in good)
+    winner.ttft_s = min(firsts) if firsts else 0.0
+    winner.tokens = max(sample.tokens for sample in comparable)
+    winner.token_source = source
+    # Not winner.tokens / winner.latency_s: both of those are extremes taken
+    # across samples, so their ratio would be a number no sample ever produced.
+    winner.tok_s = _median(throughputs)
+    return winner
+
+
+def _sample(
+    gateway: Gateway,
+    model: str,
+    *,
+    prompt: str,
+    max_tokens: int,
+    timeout: float,
+    stream: bool,
+    provider: str,
+) -> BenchResult:
+    """A single request. Own transport: :func:`probe.probe` discards ``usage``."""
+    body: dict[str, object] = {
+        "model": model,
+        "messages": [{"role": "user", "content": prompt}],
+        "max_tokens": max_tokens,
+        "stream": stream,
+    }
+    if stream:
+        # Ask for the final usage chunk. Gateways that support it make the
+        # measurement exact; the ones that ignore it fall back to chunk counting.
+        body["stream_options"] = {"include_usage": True}
+    request = urllib.request.Request(
+        gateway.chat_url,
+        data=json.dumps(body).encode(),
+        headers=gateway.headers(json_body=True),
+    )
+
+    result = BenchResult(model=model, provider=provider)
+    started = time.monotonic()
+    try:
+        with urllib.request.urlopen(request, timeout=timeout) as response:
+            result.status = getattr(response, "status", 200) or 200
+            if stream:
+                _read_stream(response, result, started)
+            else:
+                _read_json(response.read(), result)
+    except urllib.error.HTTPError as exc:
+        result.status = exc.code
+        try:
+            detail = exc.read().decode("utf-8", "replace")
+        except Exception:
+            detail = str(exc.reason or exc)
+        result.error = _flatten(detail)
+        result.verdict = taxonomy.ERR
+    except Exception as exc:
+        result.error = _flatten(f"{type(exc).__name__}: {exc}")
+        result.verdict = (
+            taxonomy.TIMEOUT if taxonomy._is_timeout(result.error) else taxonomy.ERR
+        )
+    finally:
+        result.latency_s = time.monotonic() - started
+
+    # Throughput is derived here, once, from the final latency. Computing it
+    # inside a reader and then overwriting latency_s in the finally block leaves
+    # tok/s describing a duration the result no longer reports.
+    result.tok_s = (
+        result.tokens / result.latency_s
+        if result.tokens and result.latency_s > 0
+        else 0.0
+    )
+    result.kind = taxonomy.classify_death(result.verdict, result.error, result.status)
+    result.ok = result.verdict not in taxonomy.DEAD
+    return result
+
+
+def _read_json(raw: bytes, result: BenchResult) -> None:
+    """Non-streaming path: tokens come from ``usage`` or not at all."""
+    try:
+        body = json.loads(raw.decode("utf-8", "replace"))
+    except Exception as exc:
+        result.error = _flatten(f"non-JSON body: {type(exc).__name__}: {raw[:80]!r}")
+        result.verdict = taxonomy.ERR
+        return
+    if not isinstance(body, dict):
+        result.error = _flatten(f"reply is a {type(body).__name__}, not an object")
+        result.verdict = taxonomy.ERR
+        return
+
+    verdict, _note = taxonomy.classify(result.status, body)
+    result.verdict = verdict
+    usage = body.get("usage") if isinstance(body.get("usage"), dict) else {}
+    completion = usage.get("completion_tokens")
+    if isinstance(completion, int) and completion > 0:
+        result.tokens = completion
+        result.token_source = "usage"
+    else:
+        # No usage block. Falling back to len(content) would report characters
+        # as tokens, so the honest answer is "not measured".
+        result.token_source = "none"
+
+
+def _read_stream(response: object, result: BenchResult, started: float) -> None:
+    """Streaming path: time-to-first-token is exact, token count usually is not.
+
+    ``ttft`` is measured on the first chunk that actually carries text — not the
+    first byte, which for many gateways is an immediate role-only delta and
+    would report a time-to-first-token of a few milliseconds for a route that
+    then thinks for ten seconds.
+    """
+    chunks = 0
+    usage_tokens = 0
+    text_seen = False
+    try:
+        for raw_line in response:  # type: ignore[union-attr]
+            line = raw_line.decode("utf-8", "replace").strip()
+            if not line or line.startswith(":"):
+                continue  # SSE comment / keepalive
+            if not line.startswith("data:"):
+                continue
+            payload = line[len("data:") :].strip()
+            if payload == "[DONE]":
+                break
+            try:
+                chunk = json.loads(payload)
+            except Exception:
+                continue  # a half-flushed frame is not a failure of the model
+            if not isinstance(chunk, dict):
+                continue
+
+            usage = chunk.get("usage")
+            if isinstance(usage, dict):
+                completion = usage.get("completion_tokens")
+                if isinstance(completion, int) and completion > 0:
+                    usage_tokens = completion
+
+            try:
+                delta = (chunk.get("choices") or [{}])[0].get("delta") or {}
+                piece = (delta.get("content") or delta.get("reasoning_content") or "")
+            except Exception:
+                piece = ""
+            if isinstance(piece, str) and piece.strip():
+                chunks += 1
+                if not text_seen:
+                    text_seen = True
+                    result.ttft_s = time.monotonic() - started
+    except Exception as exc:
+        # Mid-stream disconnect after some text: the route works, the measurement
+        # is partial. Keep what was counted and record the interruption.
+        result.error = _flatten(f"stream interrupted: {type(exc).__name__}: {exc}")
+
+    if usage_tokens:
+        result.tokens = usage_tokens
+        result.token_source = "usage"
+    elif chunks:
+        result.tokens = chunks
+        result.token_source = "chunks"
+    else:
+        result.token_source = "none"
+
+    result.verdict = taxonomy.OK if text_seen or usage_tokens else taxonomy.EMPTY
+
+
+def bench_many(
+    gateway: Gateway,
+    models: Sequence[str],
+    *,
+    threads: int = DEFAULT_THREADS,
+    prompt: str = DEFAULT_BENCH_PROMPT,
+    max_tokens: int = DEFAULT_BENCH_MAX_TOKENS,
+    timeout: float = DEFAULT_TIMEOUT,
+    stream: bool = False,
+    repeat: int = 1,
+    providers: dict[str, str] | None = None,
+    on_result: Callable[[BenchResult], None] | None = None,
+) -> list[BenchResult]:
+    """Benchmark a pool. ``on_result`` fires from worker threads — lock it."""
+    lookup = providers or {}
+
+    def work(model: str) -> BenchResult:
+        measured = bench_one(
+            gateway,
+            model,
+            prompt=prompt,
+            max_tokens=max_tokens,
+            timeout=timeout,
+            stream=stream,
+            provider=lookup.get(model, ""),
+            repeat=repeat,
+        )
+        if on_result is not None:
+            on_result(measured)
+        return measured
+
+    return parallel_map(models, work, threads=threads)
+
+
+def summarize(results: Sequence[BenchResult]) -> dict[str, object]:
+    """Roll a bench run up into the numbers a report header needs."""
+    alive = [item for item in results if item.ok]
+    exact = [item for item in alive if item.tok_s_exact]
+    latencies = sorted(item.latency_s for item in alive)
+
+    # Ranked by latency, which is always measured exactly. Ranking by tok/s
+    # would put chunk-counted routes next to usage-counted ones and imply the
+    # two are comparable — they are not.
+    fastest = sorted(alive, key=lambda item: item.latency_s)[:5]
+    slowest = sorted(alive, key=lambda item: item.latency_s, reverse=True)[:5]
+
+    return {
+        "total": len(results),
+        "alive": len(alive),
+        "dead": len(results) - len(alive),
+        "with_exact_tokens": len(exact),
+        "median_latency_s": round(_median(latencies), 2),
+        "fastest": [item.as_dict() for item in fastest],
+        "slowest": [item.as_dict() for item in slowest],
+    }
+
+
+def _median(values: Sequence[float]) -> float:
+    return float(statistics.median(values)) if values else 0.0