386 lines
17 KiB
Python
386 lines
17 KiB
Python
"""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())
|