poolprobe 1.0.0__py3-none-any.whl

poolprobe/probe.py

@@ -0,0 +1,352 @@
+"""Transport layer: one chat completion per probe, plus burst confirmation.
+
+Stdlib only (``urllib``). The request shape is boring on purpose — the
+interesting part is what comes back and how little of it can be trusted.
+"""
+
+from __future__ import annotations
+
+import json
+import time
+import urllib.error
+import urllib.request
+from collections.abc import Callable, Iterable, Sequence
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass, field
+from typing import TypeVar
+
+from . import taxonomy
+
+#: A short, cheap prompt. Long prompts make the sweep cost real tokens on
+#: metered upstreams and change the latency you are measuring.
+DEFAULT_PROMPT = "Reply with the single word: alive"
+
+#: Reasoning models spend tokens on thinking before content. At 32 tokens they
+#: return an empty body with ``finish_reason=length`` and look dead. 1024 is the
+#: floor that stopped the false positives; do not lower it to make sweeps faster.
+DEFAULT_MAX_TOKENS = 1024
+
+DEFAULT_TIMEOUT = 60.0
+
+#: Generic over the worker's return type so ``bench`` can reuse the same bounded
+#: pool without a second scheduler.
+T = TypeVar("T")
+R = TypeVar("R")
+
+
+class GatewayError(RuntimeError):
+    """The gateway itself is unreachable or refused us (not a model problem)."""
+
+
+@dataclass(frozen=True)
+class Gateway:
+    """An OpenAI-compatible endpoint.
+
+    ``base_url`` is the versioned root (``http://host:port/v1``); a trailing
+    slash is tolerated. ``api_key`` may be empty for keyless local gateways.
+    """
+
+    base_url: str
+    api_key: str = ""
+    name: str = "gateway"
+    extra_headers: dict[str, str] = field(default_factory=dict)
+
+    @property
+    def root(self) -> str:
+        return self.base_url.rstrip("/")
+
+    @property
+    def chat_url(self) -> str:
+        return f"{self.root}/chat/completions"
+
+    @property
+    def models_url(self) -> str:
+        return f"{self.root}/models"
+
+    def headers(self, json_body: bool = False) -> dict[str, str]:
+        out = dict(self.extra_headers)
+        if self.api_key:
+            out["Authorization"] = f"Bearer {self.api_key}"
+        if json_body:
+            out["Content-Type"] = "application/json"
+        return out
+
+
+@dataclass
+class ProbeResult:
+    """One observation about one model."""
+
+    model: str
+    verdict: str
+    note: str
+    kind: str
+    status: int = 0
+    latency_s: float = 0.0
+    finish_reason: str = "?"
+    attempts: int = 1
+    provider: str = ""
+    history: list[tuple[str, str]] = field(default_factory=list)
+
+    @property
+    def dead(self) -> bool:
+        return self.verdict in taxonomy.DEAD
+
+    @property
+    def advice(self) -> str:
+        return taxonomy.advice(self.kind)
+
+    def as_dict(self) -> dict[str, object]:
+        return {
+            "model": self.model,
+            "verdict": self.verdict,
+            "kind": self.kind,
+            "note": self.note,
+            "status": self.status,
+            "latency_s": round(self.latency_s, 2),
+            "finish_reason": self.finish_reason,
+            "attempts": self.attempts,
+            "provider": self.provider,
+            "advice": self.advice,
+            "history": [{"verdict": v, "note": n} for v, n in self.history],
+        }
+
+
+def probe(
+    gateway: Gateway,
+    model: str,
+    timeout: float = DEFAULT_TIMEOUT,
+    max_tokens: int = DEFAULT_MAX_TOKENS,
+    prompt: str = DEFAULT_PROMPT,
+) -> ProbeResult:
+    """POST one chat completion and classify whatever comes back.
+
+    ``"stream": false`` is load-bearing. Some gateway routes default to SSE
+    even for a non-streaming request; ``json.loads`` then reads the first chunk
+    as ``Expecting value: line 1 column 1`` and the model is reported dead while
+    it is answering perfectly. That false red repeated every 30 minutes on a
+    live pool until the flag was pinned.
+    """
+    payload = json.dumps(
+        {
+            "model": model,
+            "messages": [{"role": "user", "content": prompt}],
+            "max_tokens": max_tokens,
+            "stream": False,
+        }
+    ).encode()
+    request = urllib.request.Request(
+        gateway.chat_url, data=payload, headers=gateway.headers(json_body=True)
+    )
+
+    started = time.monotonic()
+    status = 0
+    # Defaults for every path that never reaches a parseable body. Without them
+    # an HTTPError or a timeout would raise UnboundLocalError while building the
+    # result, and one bad upstream would take the whole sweep down with it.
+    finish = "?"
+    note = ""
+    verdict = taxonomy.ERR
+    try:
+        with urllib.request.urlopen(request, timeout=timeout) as response:
+            status = getattr(response, "status", 200) or 200
+            raw = response.read()
+    except urllib.error.HTTPError as exc:
+        status = exc.code
+        # urlopen puts the response body on the exception; that body is the only
+        # evidence of *why* the upstream refused, so it becomes the note.
+        detail = ""
+        try:
+            detail = exc.read().decode("utf-8", "replace")
+        except Exception:
+            detail = str(exc.reason or exc)
+        note = _flatten(detail)
+        verdict = taxonomy.ERR
+    except Exception as exc:  # transport, DNS, TLS, socket timeout, proxy
+        note = _flatten(f"{type(exc).__name__}: {exc}")
+        verdict = taxonomy.TIMEOUT if taxonomy._is_timeout(note) else taxonomy.ERR
+    else:
+        try:
+            body = json.loads(raw.decode("utf-8", "replace"))
+        except Exception as exc:
+            body = None
+            note = _flatten(f"non-JSON body: {type(exc).__name__}: {raw[:80]!r}")
+            verdict = taxonomy.ERR
+        else:
+            verdict, note = taxonomy.classify(status, body)
+            finish = _finish_of(body)
+
+    elapsed = time.monotonic() - started
+    kind = taxonomy.classify_death(verdict, note, status)
+    return ProbeResult(
+        model=model,
+        verdict=verdict,
+        note=note,
+        kind=kind,
+        status=status,
+        latency_s=elapsed,
+        finish_reason=finish,
+        attempts=1,
+        history=[(verdict, note)],
+    )
+
+
+def burst(
+    gateway: Gateway,
+    model: str,
+    attempts: int = 5,
+    gap_s: float = 0.0,
+    timeout: float = DEFAULT_TIMEOUT,
+    max_tokens: int = DEFAULT_MAX_TOKENS,
+    prompt: str = DEFAULT_PROMPT,
+    early_exit_ok: bool = True,
+) -> ProbeResult:
+    """Probe the same model up to ``attempts`` times and keep the best truth.
+
+    One observation is not a verdict. On a measured pool, 39% of scheduled
+    liveness runs went red and every single red was a transient — including a
+    model that failed a probe at 17:47 and answered 5/5 two minutes later.
+
+    The returned result carries ``attempts`` actually made and the full
+    ``history``, so a report can show "recovered on attempt 3" instead of
+    hiding the flakiness.
+    """
+    attempts = max(1, int(attempts))
+    last: ProbeResult | None = None
+    best: ProbeResult | None = None
+    history: list[tuple[str, str]] = []
+    for index in range(attempts):
+        if index and gap_s:
+            time.sleep(gap_s)
+        result = probe(
+            gateway,
+            model,
+            timeout=timeout,
+            max_tokens=max_tokens,
+            prompt=prompt,
+        )
+        history.append((result.verdict, result.note))
+        last = result
+        if not result.dead:
+            best = result
+            if early_exit_ok:
+                break
+        elif result.kind not in taxonomy.PRUNABLE_KINDS and early_exit_ok:
+            # A transient failure is already the answer we needed: this model is
+            # not a prune candidate, so stop spending requests on it.
+            break
+    assert last is not None
+    # "Best truth wins": one good answer inside the burst outranks any number of
+    # bad ones, otherwise a model that recovers on attempt 3 and flakes on 5
+    # would be reported dead on the strength of its last observation.
+    winner = best or last
+    winner.attempts = len(history)
+    winner.history = history
+    return winner
+
+
+def list_models(gateway: Gateway, timeout: float = 20.0) -> list[str]:
+    """Model ids published by ``GET /models``.
+
+    Raises :class:`GatewayError` when the catalog cannot be read: guessing an
+    empty pool would silently turn every downstream gate green.
+    """
+    request = urllib.request.Request(gateway.models_url, headers=gateway.headers())
+    try:
+        with urllib.request.urlopen(request, timeout=timeout) as response:
+            payload = json.loads(response.read().decode("utf-8", "replace"))
+    except Exception as exc:
+        raise GatewayError(
+            f"cannot read {gateway.models_url}: {type(exc).__name__}: {exc}"
+        ) from exc
+
+    entries = payload.get("data") if isinstance(payload, dict) else None
+    if not isinstance(entries, list):
+        raise GatewayError(f"{gateway.models_url} returned no data[] list")
+    ids = [entry.get("id") for entry in entries if isinstance(entry, dict)]
+    return sorted({str(i) for i in ids if i})
+
+
+def ping(gateway: Gateway, timeout: float = 10.0) -> str:
+    """Cheap reachability check used by ``--selfcheck``-adjacent smoke runs."""
+    try:
+        count = len(list_models(gateway, timeout=timeout))
+    except GatewayError as exc:
+        return str(exc)
+    return f"{count} models published"
+
+
+def parallel_map(
+    items: Iterable[T],
+    worker: Callable[[T], R],
+    threads: int = 8,
+) -> list[R]:
+    """Run ``worker`` over ``items``, bounded and order-preserving.
+
+    Threads, not processes: one blocking socket read per model makes the sweep
+    I/O bound. Order is preserved so two runs can be diffed line by line, which
+    is how a slow drift in the pool stays visible across weeks of cron.
+    """
+    todo = list(items)
+    if not todo:
+        return []
+    threads = max(1, min(int(threads), len(todo)))
+    if threads == 1:
+        return [worker(model) for model in todo]
+    with ThreadPoolExecutor(max_workers=threads) as pool:
+        return list(pool.map(worker, todo))
+
+
+def sweep(
+    gateway: Gateway,
+    models: Sequence[str],
+    *,
+    threads: int = 8,
+    timeout: float = DEFAULT_TIMEOUT,
+    max_tokens: int = DEFAULT_MAX_TOKENS,
+    prompt: str = DEFAULT_PROMPT,
+    attempts: int = 1,
+    gap_s: float = 0.0,
+    on_result: Callable[[ProbeResult], None] | None = None,
+) -> list[ProbeResult]:
+    """Probe every model once, or burst-confirm each when ``attempts > 1``.
+
+    ``on_result`` fires from worker threads as each model lands, so a caller can
+    stream progress; it must be thread-safe (the CLI wraps its printer in a
+    lock). Results are returned in the order of ``models`` regardless.
+    """
+
+    def work(model: str) -> ProbeResult:
+        if attempts > 1:
+            result = burst(
+                gateway,
+                model,
+                attempts=attempts,
+                gap_s=gap_s,
+                timeout=timeout,
+                max_tokens=max_tokens,
+                prompt=prompt,
+            )
+        else:
+            result = probe(
+                gateway, model, timeout=timeout, max_tokens=max_tokens, prompt=prompt
+            )
+        if on_result is not None:
+            on_result(result)
+        return result
+
+    return parallel_map(models, work, threads=threads)
+
+
+def _flatten(text: str) -> str:
+    """Collapse a body into one printable line, capped."""
+    cleaned = " ".join((text or "").split())
+    return (cleaned or "empty")[:160]
+
+
+def _finish_of(body: object) -> str:
+    """``finish_reason`` of the first choice, or ``?`` when the shape is odd."""
+    try:
+        choices = (body or {}).get("choices") or [{}]  # type: ignore[union-attr]
+        return str(choices[0].get("finish_reason") or "?")
+    except Exception:
+        return "?"