thoxmesh_node/telemetry.py

"""Rolling performance telemetry for ranking a node inside the mesh.

``meshstack_model_endpoints_v2`` has **no** latency, success-rate or load
columns. It has ``capabilities jsonb``, ``resource_requirements jsonb`` and a
``status`` enum. Rather than propose new columns (which would fork the Mesh
backend lane's schema mid-flight), this module packs telemetry into the jsonb
fields the contract already defines, under a namespaced ``thox_telemetry`` key
so it cannot collide with capability flags that ThoxRoute matches on.

That choice has a real consequence worth stating: jsonb values are not indexed
for range queries here, so ranking on latency happens in the router, not in
Postgres. At this tier's endpoint counts that is the right trade; if the mesh
ever holds thousands of endpoints, promoting these to real columns becomes
worthwhile and this module is the single place that changes.

The ``status`` field is the one piece of telemetry the schema *does* model, so
health degradation is mapped onto it: a node whose recent success rate collapses
reports ``degraded`` instead of silently continuing to advertise ``healthy``.
"""

from __future__ import annotations

import threading
import time
from collections import deque
from dataclasses import dataclass, field
from typing import Any, Deque

#: Endpoint status values permitted by the v2 schema's CHECK constraint.
STATUS_HEALTHY = "healthy"
STATUS_DEGRADED = "degraded"
STATUS_OFFLINE = "offline"

#: Success rate below which the node self-reports ``degraded``.
DEGRADED_SUCCESS_RATE = 0.80

#: Minimum completed requests before success rate is trusted for status.
MIN_SAMPLES_FOR_STATUS = 5


@dataclass
class _Sample:
    """One completed inference request."""

    latency_ms: float
    ok: bool
    completion_tokens: int
    at: float = field(default_factory=time.time)


class TelemetryRecorder:
    """Thread-safe rolling window of request outcomes.

    A bounded deque is used rather than cumulative counters so the reported
    figures track *current* behaviour. On a Colab node that matters: a session
    that was fast for an hour and then hits a throttled GPU should start
    advertising the degraded reality within a window, not average it away.
    """

    def __init__(self, window: int = 100, *, clock=time.time) -> None:
        if window < 1:
            raise ValueError("window must be >= 1")
        self._samples: Deque[_Sample] = deque(maxlen=window)
        self._lock = threading.Lock()
        self._clock = clock
        self._in_flight = 0
        self._started_at = clock()
        self._total = 0
        self._total_failed = 0

    # ── Recording ────────────────────────────────────────────────────────

    def request_started(self) -> None:
        with self._lock:
            self._in_flight += 1

    def request_finished(self, *, latency_ms: float, ok: bool, completion_tokens: int = 0) -> None:
        with self._lock:
            self._in_flight = max(0, self._in_flight - 1)
            self._total += 1
            if not ok:
                self._total_failed += 1
            self._samples.append(
                _Sample(
                    latency_ms=float(latency_ms),
                    ok=bool(ok),
                    completion_tokens=int(completion_tokens),
                    at=self._clock(),
                )
            )

    # ── Derived metrics ──────────────────────────────────────────────────

    def snapshot(self) -> dict[str, Any]:
        """Current metrics. Safe to call from the heartbeat thread."""
        with self._lock:
            samples = list(self._samples)
            in_flight = self._in_flight
            total = self._total
            total_failed = self._total_failed
            uptime = max(0.0, self._clock() - self._started_at)

        if not samples:
            return {
                "samples": 0,
                "in_flight": in_flight,
                "total_requests": total,
                "total_failed": total_failed,
                "success_rate": None,
                "latency_ms_p50": None,
                "latency_ms_p95": None,
                "tokens_per_second": None,
                "uptime_seconds": round(uptime, 1),
            }

        latencies = sorted(s.latency_ms for s in samples)
        ok_count = sum(1 for s in samples if s.ok)
        elapsed = max(1e-6, samples[-1].at - samples[0].at) if len(samples) > 1 else None
        tokens = sum(s.completion_tokens for s in samples)
        return {
            "samples": len(samples),
            "in_flight": in_flight,
            "total_requests": total,
            "total_failed": total_failed,
            "success_rate": round(ok_count / len(samples), 4),
            "latency_ms_p50": round(_percentile(latencies, 0.50), 1),
            "latency_ms_p95": round(_percentile(latencies, 0.95), 1),
            "tokens_per_second": round(tokens / elapsed, 2) if elapsed and tokens else None,
            "uptime_seconds": round(uptime, 1),
        }

    def status(self) -> str:
        """Map telemetry onto the schema's ``status`` enum.

        Only demotes once there is enough evidence; a single early failure on a
        cold node should not pull the endpoint out of the candidate set.
        """
        snap = self.snapshot()
        rate = snap["success_rate"]
        if rate is None or snap["samples"] < MIN_SAMPLES_FOR_STATUS:
            return STATUS_HEALTHY
        return STATUS_HEALTHY if rate >= DEGRADED_SUCCESS_RATE else STATUS_DEGRADED

    def load_factor(self) -> float:
        """Concurrent requests in flight, as a plain gauge for the router."""
        with self._lock:
            return float(self._in_flight)


def _percentile(sorted_values: list[float], q: float) -> float:
    """Nearest-rank percentile over a pre-sorted list."""
    if not sorted_values:
        raise ValueError("cannot take a percentile of an empty list")
    if len(sorted_values) == 1:
        return sorted_values[0]
    index = min(len(sorted_values) - 1, max(0, int(round(q * (len(sorted_values) - 1)))))
    return sorted_values[index]


def build_capabilities(
    *,
    base: dict[str, Any],
    recorder: TelemetryRecorder,
    node_kind: str,
    ephemeral: bool,
    heartbeat_seconds: int,
    ttl_seconds: int,
) -> dict[str, Any]:
    """Compose the ``capabilities`` jsonb sent on register and heartbeat.

    ``base`` holds the classification/capability flags ThoxRoute matches on
    (``chat``, ``code``, ``reasoning``, ...). ``resolve_route`` in the edge
    function does a truthiness lookup ``capabilities[requested_capability]``, so
    those flags must stay at the top level as plain booleans — telemetry is
    namespaced beneath ``thox_telemetry`` precisely to keep that lookup clean.

    ``thox_lease`` carries the ephemerality contract. The backend does not reap
    stale endpoints today, so the node publishes its own TTL and the router
    treats an endpoint whose ``expires_at`` has passed as ineligible. This is a
    client-side lease: it degrades safely if the backend later grows a reaper.
    """
    now = time.time()
    capabilities = dict(base)
    capabilities["thox_telemetry"] = recorder.snapshot()
    capabilities["thox_node"] = {
        "kind": node_kind,
        "ephemeral": ephemeral,
        "agent": "thoxmesh-node",
    }
    capabilities["thox_lease"] = {
        "heartbeat_seconds": heartbeat_seconds,
        "ttl_seconds": ttl_seconds,
        "renewed_at": _iso(now),
        "expires_at": _iso(now + ttl_seconds),
    }
    return capabilities


def is_lease_expired(capabilities: dict[str, Any] | None, *, now: float | None = None) -> bool:
    """True when an endpoint's self-published lease has lapsed.

    Used by the router to skip endpoints left behind by a Colab session that
    died without deregistering — the dominant failure mode in this tier.
    Endpoints with no lease block are treated as live, so endpoints registered
    by other (non-ephemeral) device types are unaffected.
    """
    if not capabilities:
        return False
    lease = capabilities.get("thox_lease")
    if not isinstance(lease, dict):
        return False
    expires_at = lease.get("expires_at")
    if not isinstance(expires_at, str):
        return False
    parsed = _parse_iso(expires_at)
    if parsed is None:
        return False
    return (now if now is not None else time.time()) > parsed


def _iso(epoch_seconds: float) -> str:
    """UTC ISO-8601 with a trailing Z, matching the backend's timestamps."""
    return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime(epoch_seconds))


def _parse_iso(value: str) -> float | None:
    """Parse the subset of ISO-8601 this module emits. ``None`` when unparseable."""
    text = value.strip().replace("Z", "+0000")
    for fmt in ("%Y-%m-%dT%H:%M:%S%z", "%Y-%m-%dT%H:%M:%S.%f%z"):
        try:
            import datetime as _dt

            return _dt.datetime.strptime(text, fmt).timestamp()
        except ValueError:
            continue
    return None