neuron-tai/packages/node/meshnet_node/seam_telemetry.py

"""Bounded, in-process telemetry for distributed activation seams.

The generation path records one cheap counter update per activation.  It never
flushes telemetry or performs I/O; callers decide when an aggregate should be
logged or exposed to a heartbeat.
"""

from __future__ import annotations

from dataclasses import dataclass
import time


@dataclass
class _SeamAggregate:
    phase: str
    hop: int
    node: str
    count: int = 0
    latency_ms: float = 0.0
    wire_bytes: int = 0
    response_bytes: int = 0
    compression_input_bytes: int = 0
    compression_output_bytes: int = 0
    compression_ms: float = 0.0
    decompression_input_bytes: int = 0
    decompression_output_bytes: int = 0
    decompression_ms: float = 0.0
    reused_connections: int = 0
    last_activation_id: str = ""


class GenerationTelemetry:
    """Aggregate activation measurements for one stable Route Session."""

    def __init__(
        self,
        session_id: str,
        *,
        report_every: int = 32,
        report_interval: float = 5.0,
        now: float | None = None,
    ) -> None:
        self.session_id = session_id
        self.report_every = max(1, report_every)
        self.report_interval = max(0.0, report_interval)
        self.started = time.monotonic() if now is None else now
        self._last_report = self.started
        self._total_tokens = 0
        self._closed = False
        self._report_due = False
        self._seams: dict[tuple[str, int, str], _SeamAggregate] = {}

    def record_seam(
        self,
        *,
        activation_id: str,
        phase: str,
        hop: int,
        node: str,
        latency_seconds: float,
        wire_bytes: int,
        response_bytes: int,
        connection_reused: bool,
        now: float | None = None,
    ) -> bool:
        """Record one activation locally and say whether a summary is due."""
        if self._closed:
            return False
        observed = time.monotonic() if now is None else now
        key = (phase, hop, node)
        aggregate = self._seams.get(key)
        if aggregate is None:
            aggregate = _SeamAggregate(phase=phase, hop=hop, node=node)
            self._seams[key] = aggregate
        aggregate.count += 1
        aggregate.latency_ms += max(0.0, latency_seconds) * 1000.0
        aggregate.wire_bytes += max(0, wire_bytes)
        aggregate.response_bytes += max(0, response_bytes)
        aggregate.reused_connections += int(connection_reused)
        aggregate.last_activation_id = activation_id
        due = (
            aggregate.count == 1
            or aggregate.count % self.report_every == 0
            or observed - self._last_report >= self.report_interval
        )
        self._report_due = self._report_due or due
        return due

    @property
    def report_due(self) -> bool:
        return self._report_due

    def note_tokens(self, tokens: int) -> None:
        if not self._closed:
            self._total_tokens = max(0, tokens)

    def record_compression(
        self, *, phase: str, hop: int, node: str, input_bytes: int,
        output_bytes: int, elapsed_seconds: float, decompression: bool = False,
    ) -> None:
        """Attach compression work to the same bounded seam aggregate."""
        if self._closed:
            return
        key = (phase, hop, node)
        aggregate = self._seams.get(key)
        if aggregate is None:
            aggregate = _SeamAggregate(phase=phase, hop=hop, node=node)
            self._seams[key] = aggregate
        if decompression:
            aggregate.decompression_input_bytes += max(0, input_bytes)
            aggregate.decompression_output_bytes += max(0, output_bytes)
            aggregate.decompression_ms += max(0.0, elapsed_seconds) * 1000.0
        else:
            aggregate.compression_input_bytes += max(0, input_bytes)
            aggregate.compression_output_bytes += max(0, output_bytes)
            aggregate.compression_ms += max(0.0, elapsed_seconds) * 1000.0

    def snapshot(self, *, now: float | None = None) -> dict:
        observed = time.monotonic() if now is None else now
        elapsed = max(observed - self.started, 1e-6)
        seams = []
        for aggregate in self._seams.values():
            seams.append({
                "phase": aggregate.phase,
                "hop": aggregate.hop,
                "node": aggregate.node,
                "activations": aggregate.count,
                "latency_ms": round(aggregate.latency_ms, 3),
                "avg_latency_ms": round(aggregate.latency_ms / max(1, aggregate.count), 3),
                "wire_bytes": aggregate.wire_bytes,
                "response_bytes": aggregate.response_bytes,
                "compression_input_bytes": aggregate.compression_input_bytes,
                "compression_output_bytes": aggregate.compression_output_bytes,
                "compression_ms": round(aggregate.compression_ms, 3),
                "decompression_input_bytes": aggregate.decompression_input_bytes,
                "decompression_output_bytes": aggregate.decompression_output_bytes,
                "decompression_ms": round(aggregate.decompression_ms, 3),
                "connection_reuse": aggregate.reused_connections,
                "last_activation_id": aggregate.last_activation_id,
            })
        return {
            "session_id": self.session_id,
            "tokens_per_sec": round(self._total_tokens / elapsed, 2),
            "seams": seams,
        }

    def mark_reported(self, *, now: float | None = None) -> None:
        self._last_report = time.monotonic() if now is None else now
        self._report_due = False

    def close(self) -> None:
        self._closed = True
        self._seams.clear()
        self._report_due = False