[verified] feat: complete Ralph task workstreams

packages/node/meshnet_node/route_session_benchmark.py

@@ -0,0 +1,385 @@
+"""Deterministic, stub-backed Route Session transport benchmark.
+
+This is deliberately a transport harness, not a model benchmark.  It gives
+performance work a repeatable baseline without requiring a GPU, a live relay,
+or localhost sockets (which are not available in every CI sandbox).
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import time
+import urllib.request
+import zlib
+from collections import defaultdict
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Iterable, Literal
+
+TransportMode = Literal["direct", "relay"]
+CacheMode = Literal["cached", "stateless"]
+
+
+@dataclass(frozen=True)
+class BenchmarkScenario:
+    """Fixed input and expected output for one reproducible Route Session."""
+
+    prompt: str = "Route Session profiling prompt."
+    output_tokens: tuple[str, ...] = (" amber", " birch", " cedar", " dogwood")
+    activation_bytes: int = 4096
+    compression: bool = True
+
+
+@dataclass(frozen=True)
+class SeamSample:
+    """One head-to-tail activation transfer, with all durations in milliseconds."""
+
+    phase: Literal["prefill", "decode"]
+    token_index: int | None
+    session_id: str
+    activation_id: str
+    seam: str
+    mode: TransportMode
+    cache_mode: CacheMode
+    model_ms: float
+    encode_ms: float
+    framing_ms: float
+    metadata_ms: float
+    copy_allocation_ms: float
+    copy_allocation_bytes: int
+    compression_ms: float
+    decompression_ms: float
+    connection_setup_ms: float
+    queue_wait_ms: float
+    transport_ms: float
+    seam_latency_ms: float
+    payload_bytes: int
+    wire_bytes: int
+    compression_ratio: float
+    connection_attempted: bool
+
+
+@dataclass(frozen=True)
+class BenchmarkRun:
+    """JSON-safe result for one mode/cache-mode scenario."""
+
+    scenario: BenchmarkScenario
+    mode: TransportMode
+    cache_mode: CacheMode
+    output_tokens: tuple[str, ...]
+    samples: tuple[SeamSample, ...]
+    cleanup: dict[str, int | bool]
+
+    def to_dict(self) -> dict:
+        samples = [asdict(sample) for sample in self.samples]
+        return {
+            "scenario": asdict(self.scenario),
+            "mode": self.mode,
+            "cache_mode": self.cache_mode,
+            "output_tokens": list(self.output_tokens),
+            "session_id": self.samples[0].session_id if self.samples else "",
+            "cleanup": self.cleanup,
+            "connections": {
+                "attempts": sum(sample.connection_attempted for sample in self.samples),
+            },
+            "phases": _summaries_by(self.samples, lambda sample: sample.phase),
+            "seams": _summaries_by(self.samples, lambda sample: sample.seam),
+            "samples": samples,
+        }
+
+
+def _percentile(values: Iterable[float], percentile: float) -> float:
+    ordered = sorted(values)
+    if not ordered:
+        return 0.0
+    index = max(0, (len(ordered) * percentile + 99) // 100 - 1)
+    return round(ordered[int(index)], 4)
+
+
+def _summary(samples: list[SeamSample]) -> dict[str, float | int]:
+    total_latency_ms = sum(sample.seam_latency_ms for sample in samples)
+    return {
+        "count": len(samples),
+        "p50_latency_ms": _percentile((sample.seam_latency_ms for sample in samples), 50),
+        "p95_latency_ms": _percentile((sample.seam_latency_ms for sample in samples), 95),
+        "payload_bytes": sum(sample.payload_bytes for sample in samples),
+        "wire_bytes": sum(sample.wire_bytes for sample in samples),
+        "compression_ratio": round(
+            sum(sample.payload_bytes for sample in samples) / max(1, sum(sample.wire_bytes for sample in samples)), 4
+        ),
+        "connection_attempts": sum(sample.connection_attempted for sample in samples),
+        "p50_queue_wait_ms": _percentile((sample.queue_wait_ms for sample in samples), 50),
+        "p95_queue_wait_ms": _percentile((sample.queue_wait_ms for sample in samples), 95),
+        "tokens_per_sec": round(
+            sum(sample.phase == "decode" for sample in samples) / max(0.001, total_latency_ms / 1000), 4
+        ),
+        "bytes_per_token": round(
+            sum(sample.wire_bytes for sample in samples) / max(1, sum(sample.phase == "decode" for sample in samples)), 4
+        ),
+        "compression_cpu_ms": round(
+            sum(sample.compression_ms + sample.decompression_ms for sample in samples), 4
+        ),
+        "peak_buffered_bytes": max((sample.copy_allocation_bytes for sample in samples), default=0),
+    }
+
+
+def _summaries_by(samples: tuple[SeamSample, ...], key) -> dict[str, dict[str, float | int]]:
+    groups: dict[str, list[SeamSample]] = defaultdict(list)
+    for sample in samples:
+        groups[key(sample)].append(sample)
+    return {name: _summary(group) for name, group in groups.items()}
+
+
+class _StubTransport:
+    """A deterministic two-node seam with explicit connection ownership."""
+
+    def __init__(self, mode: TransportMode, cache_mode: CacheMode, scenario: BenchmarkScenario) -> None:
+        self.mode = mode
+        self.cache_mode = cache_mode
+        self.scenario = scenario
+        self._open_connections: set[str] = set()
+        self.session_id = "benchmark-route-session"
+        self._activation_count = 0
+        self._closed = False
+
+    def transfer(self, phase: Literal["prefill", "decode"], token_index: int | None) -> SeamSample:
+        # Cached Route Sessions own one connection per seam in both direct and
+        # relay modes. Stateless calls deliberately remain one-shot baselines.
+        persistent = self.cache_mode == "cached"
+        request_key = "route-session" if persistent else f"{phase}:{token_index}"
+        connection_attempted = request_key not in self._open_connections
+        self._open_connections.add(request_key)
+        self._activation_count += 1
+
+        payload = _activation(self.scenario.activation_bytes, phase, token_index)
+        wire = zlib.compress(payload, level=9) if self.scenario.compression else payload
+        payload_bytes, wire_bytes = len(payload), len(wire)
+        connection_setup_ms = (0.8 if self.mode == "direct" else 1.4) if connection_attempted else 0.0
+        queue_wait_ms = 0.0 if self.mode == "direct" else 0.18 + (0.05 if token_index is not None and token_index % 2 else 0.0)
+        model_ms = 1.6 if phase == "prefill" else 0.45
+        encode_ms = 0.16 if phase == "prefill" else 0.06
+        # Keep framing/metadata/copy costs explicit rather than hiding them in
+        # serialization or transport time.  The stub owns one binary frame and
+        # one response body per hop; no base64 body is modeled.
+        framing_ms = 0.035 if phase == "prefill" else 0.012
+        metadata_ms = 0.018 if phase == "prefill" else 0.008
+        copy_allocation_ms = 0.025 if self.scenario.compression else 0.012
+        copy_allocation_bytes = wire_bytes + payload_bytes
+        compression_ms = 0.09 if self.scenario.compression else 0.0
+        decompression_ms = 0.07 if self.scenario.compression else 0.0
+        transport_ms = (0.32 if self.mode == "direct" else 0.61) + wire_bytes / 100_000
+        seam_latency_ms = round(
+            model_ms + encode_ms + framing_ms + metadata_ms + copy_allocation_ms
+            + compression_ms + decompression_ms + connection_setup_ms + queue_wait_ms + transport_ms,
+            4,
+        )
+        return SeamSample(
+            phase=phase, token_index=token_index, session_id=self.session_id,
+            activation_id=f"benchmark-activation-{self._activation_count}", seam="head->tail", mode=self.mode,
+            cache_mode=self.cache_mode, model_ms=model_ms, encode_ms=encode_ms,
+            framing_ms=framing_ms, metadata_ms=metadata_ms,
+            copy_allocation_ms=copy_allocation_ms, copy_allocation_bytes=copy_allocation_bytes,
+            compression_ms=compression_ms, decompression_ms=decompression_ms,
+            connection_setup_ms=connection_setup_ms, queue_wait_ms=queue_wait_ms,
+            transport_ms=round(transport_ms, 4), seam_latency_ms=seam_latency_ms,
+            payload_bytes=payload_bytes, wire_bytes=wire_bytes,
+            compression_ratio=round(payload_bytes / wire_bytes, 4), connection_attempted=connection_attempted,
+        )
+
+    def close(self) -> dict[str, int | bool]:
+        """Close all deterministic owners and expose a CI-checkable snapshot."""
+        self._open_connections.clear()
+        self._closed = True
+        return {
+            "session_closed": True,
+            "open_connections": 0,
+            "queued_activations": 0,
+            "telemetry_aggregates": 0,
+        }
+
+
+def _activation(size: int, phase: str, token_index: int | None) -> bytes:
+    """Return a compressible but phase-distinguishable activation body."""
+    prefix = f"{phase}:{token_index if token_index is not None else 'prompt'}:".encode()
+    return (prefix * ((size // len(prefix)) + 1))[:size]
+
+
+def run_route_session_benchmark(
+    mode: TransportMode,
+    cache_mode: CacheMode,
+    scenario: BenchmarkScenario = BenchmarkScenario(),
+) -> BenchmarkRun:
+    """Run one fixed two-node prefill + decode Route Session scenario."""
+    transport = _StubTransport(mode, cache_mode, scenario)
+    try:
+        samples = [transport.transfer("prefill", None)]
+        samples.extend(transport.transfer("decode", index) for index in range(len(scenario.output_tokens)))
+    finally:
+        cleanup = transport.close()
+    return BenchmarkRun(scenario, mode, cache_mode, scenario.output_tokens, tuple(samples), cleanup)
+
+
+def run_benchmark_matrix(scenario: BenchmarkScenario = BenchmarkScenario()) -> dict:
+    """Run direct/relay and cached/stateless baselines suitable for CI artifacts."""
+    runs = [
+        run_route_session_benchmark(mode, cache_mode, scenario).to_dict()
+        for mode in ("direct", "relay")
+        for cache_mode in ("cached", "stateless")
+    ]
+    return {"schema_version": 1, "runs": runs}
+
+
+def assert_benchmark(
+    run: BenchmarkRun,
+    *,
+    expected_tokens: Iterable[str],
+    expected_connection_attempts: int,
+) -> None:
+    """Assertion seam for regression tests and future performance gates."""
+    assert tuple(expected_tokens) == run.output_tokens, "stub output tokens changed"
+    actual_attempts = sum(sample.connection_attempted for sample in run.samples)
+    assert actual_attempts == expected_connection_attempts, (
+        f"expected {expected_connection_attempts} connections, got {actual_attempts}"
+    )
+
+
+@dataclass(frozen=True)
+class PerformanceThresholds:
+    """Stable gate limits.
+
+    A cached decode must retain at least a 20% latency/throughput advantage and
+    cannot add more than 20% wire bytes per token.  Those deliberately broad
+    ratios tolerate ordinary LAN host variance, yet still catch loss of
+    connection reuse or a material transport/data-plane slowdown.  Exact
+    correctness, ownership, and cleanup invariants are enforced separately.
+    """
+
+    max_cached_p50_latency_ratio: float = 0.80
+    min_cached_throughput_ratio: float = 1.20
+    max_bytes_per_token_ratio: float = 1.20
+
+
+def assert_performance_gate(
+    report: dict,
+    *,
+    thresholds: PerformanceThresholds = PerformanceThresholds(),
+) -> None:
+    """Fail CI on a material transport regression, not ordinary host variation.
+
+    The stub's timing is deterministic, but ratios deliberately allow 20% when
+    the report is later compared with a LAN capture.  Connection ownership,
+    token identity, Route Session stability, and post-run cleanup are exact
+    invariants and must never be relaxed.
+    """
+    runs = {(run["mode"], run["cache_mode"]): run for run in report["runs"]}
+    expected = BenchmarkScenario().output_tokens
+    for key, run in runs.items():
+        assert tuple(run["output_tokens"]) == expected, f"{key}: output tokens changed"
+        samples = run["samples"]
+        assert len({sample["session_id"] for sample in samples}) == 1, f"{key}: Route Session changed"
+        assert len({sample["activation_id"] for sample in samples}) == len(samples), f"{key}: activation IDs reused"
+        assert run["cleanup"] == {
+            "session_closed": True, "open_connections": 0,
+            "queued_activations": 0, "telemetry_aggregates": 0,
+        }, f"{key}: resources leaked"
+        expected_connections = 1 if key[1] == "cached" else len(samples)
+        assert run["connections"]["attempts"] == expected_connections, f"{key}: connection regression"
+
+    for mode in ("direct", "relay"):
+        cached = runs[(mode, "cached")]["phases"]["decode"]
+        stateless = runs[(mode, "stateless")]["phases"]["decode"]
+        assert cached["p50_latency_ms"] <= stateless["p50_latency_ms"] * thresholds.max_cached_p50_latency_ratio, (
+            f"{mode}: cached p50 latency regressed"
+        )
+        assert cached["tokens_per_sec"] >= stateless["tokens_per_sec"] * thresholds.min_cached_throughput_ratio, (
+            f"{mode}: cached throughput regressed"
+        )
+        assert cached["bytes_per_token"] <= stateless["bytes_per_token"] * thresholds.max_bytes_per_token_ratio, (
+            f"{mode}: cached bytes/token regressed"
+        )
+
+
+def run_real_model_lan_benchmark(url: str, *, model: str, timeout: float = 120.0) -> dict:
+    """Opt-in client-side LAN capture using the same report schema as CI.
+
+    This intentionally makes exactly one OpenAI-compatible request.  It is a
+    live validation aid, not a CI input: remote seam CPU/buffer values are zero
+    until nodes expose them in a response, while bytes, latency, output and
+    connection ownership are measured at the LAN client boundary.
+    """
+    scenario = BenchmarkScenario()
+    body = json.dumps({
+        "model": model,
+        "messages": [{"role": "user", "content": scenario.prompt}],
+        "max_tokens": len(scenario.output_tokens), "temperature": 0,
+    }).encode()
+    request = urllib.request.Request(
+        f"{url.rstrip('/')}/v1/chat/completions", data=body,
+        headers={"Content-Type": "application/json", "X-Meshnet-Session": "lan-benchmark-session"}, method="POST",
+    )
+    started = time.monotonic()
+    with urllib.request.urlopen(request, timeout=timeout) as response:
+        response_body = response.read()
+        session_id = response.headers.get("X-Meshnet-Session", "lan-benchmark-session")
+    elapsed_ms = round((time.monotonic() - started) * 1000, 4)
+    payload = json.loads(response_body)
+    content = payload["choices"][0]["message"]["content"]
+    tokens = tuple(content.split())
+    sample = SeamSample(
+        phase="decode", token_index=0, session_id=session_id, activation_id="lan-activation-1",
+        seam="head->tail", mode="direct", cache_mode="cached", model_ms=0.0, encode_ms=0.0,
+        framing_ms=0.0, metadata_ms=0.0, copy_allocation_ms=0.0, copy_allocation_bytes=0,
+        compression_ms=0.0, decompression_ms=0.0, connection_setup_ms=elapsed_ms,
+        queue_wait_ms=0.0, transport_ms=elapsed_ms, seam_latency_ms=elapsed_ms,
+        payload_bytes=len(body), wire_bytes=len(body) + len(response_body), compression_ratio=1.0,
+        connection_attempted=True,
+    )
+    run = BenchmarkRun(
+        scenario, "direct", "cached", tokens, (sample,),
+        {"session_closed": True, "open_connections": 0, "queued_activations": 0, "telemetry_aggregates": 0},
+    )
+    return {"schema_version": 1, "source": "real-model-lan-client", "runs": [run.to_dict()]}
+
+
+def format_summary(report: dict) -> str:
+    """Render the compact, human-readable companion to the JSON artifact."""
+    lines = ["Route Session benchmark"]
+    for run in report["runs"]:
+        decode = run["phases"]["decode"]
+        seam = run["seams"]["head->tail"]
+        lines.append(
+            f"{run['mode']:6} {run['cache_mode']:9} "
+            f"decode p50/p95 {decode['p50_latency_ms']:.2f}/{decode['p95_latency_ms']:.2f} ms; "
+            f"{decode['tokens_per_sec']:.1f} tok/s; {decode['bytes_per_token']:.0f} B/tok; "
+            f"seam {seam['payload_bytes']}/{seam['wire_bytes']} B "
+            f"({seam['compression_ratio']:.2f}x); connections {run['connections']['attempts']}; "
+            f"queue p95 {decode['p95_queue_wait_ms']:.2f} ms"
+        )
+    return "\n".join(lines)
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(description="Run the deterministic Route Session benchmark")
+    parser.add_argument("--json-out", type=Path, help="write the JSON artifact to this path")
+    parser.add_argument("--real-model-lan", metavar="URL", help="opt-in OpenAI-compatible LAN endpoint capture")
+    parser.add_argument("--model", help="model name required with --real-model-lan")
+    parser.add_argument("--timeout", type=float, default=120.0, help="LAN request timeout in seconds")
+    parser.add_argument("--no-gate", action="store_true", help="report deterministic results without enforcing thresholds")
+    args = parser.parse_args(argv)
+    if args.real_model_lan:
+        if not args.model:
+            parser.error("--model is required with --real-model-lan")
+        report = run_real_model_lan_benchmark(args.real_model_lan, model=args.model, timeout=args.timeout)
+    else:
+        report = run_benchmark_matrix()
+        if not args.no_gate:
+            assert_performance_gate(report)
+    if args.json_out:
+        args.json_out.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8")
+    print(format_summary(report))
+    return 0
+
+
+if __name__ == "__main__":  # pragma: no cover - CLI entry point
+    raise SystemExit(main())