From 3079bbf494f574d843a9bdb9b0fc1b7b2f507b2a Mon Sep 17 00:00:00 2001 From: fluffy314 Date: Thu, 16 Jul 2026 21:36:34 +0800 Subject: [PATCH] feat(bench): add three-phase prefill fleet benchmark Provide one terminal task for remote compute, Primary hot hit, and allens cold restore; persist privacy-safe timing and hit metrics and render live/history/detail reports on the network dashboard. Co-authored-by: Cursor --- docs/ops/distributed-prefill-kv-network.md | 27 ++ .../bench/prefill_fleet_report.py | 99 +++++++ inference_engine/network/api.py | 62 ++++ inference_engine/network/dashboard.py | 11 +- inference_engine/network/state.py | 104 ++++++- scripts/benchmark_prefill_architecture.py | 269 ++++++++++++++++++ scripts/run_prefill_architecture_benchmark.sh | 11 + .../bench/test_prefill_fleet_report.py | 67 +++++ .../test_prefill_architecture_benchmark.py | 46 +++ .../network/test_network_api.py | 58 ++++ .../network/test_network_state.py | 45 +++ 11 files changed, 794 insertions(+), 5 deletions(-) create mode 100644 inference_engine/bench/prefill_fleet_report.py create mode 100644 scripts/benchmark_prefill_architecture.py create mode 100644 scripts/run_prefill_architecture_benchmark.sh create mode 100644 tests/inference_engine/bench/test_prefill_fleet_report.py create mode 100644 tests/inference_engine/bridge/test_prefill_architecture_benchmark.py diff --git a/docs/ops/distributed-prefill-kv-network.md b/docs/ops/distributed-prefill-kv-network.md index 8d098254..5b7e014f 100644 --- a/docs/ops/distributed-prefill-kv-network.md +++ b/docs/ops/distributed-prefill-kv-network.md @@ -275,6 +275,33 @@ publish failures, fallbacks, or `/tmp/kakeya-cache-fill.stop`, and never expects resident fleet usage to exceed the configured 1+8 GiB ceiling. Churn is accepted when `bytes_evicted` increases while resident bytes remain bounded. +## Three-phase architecture benchmark + +Run from Primary; services are started only if missing and remain running: + +```bash +bash scripts/run_prefill_architecture_benchmark.sh \ + --output-tokens 32 \ + --report /tmp/kakeya-prefill-benchmark.json +``` + +The task runs `remote_compute`, `primary_hot_hit`, and +`allens_cold_restore`, recording client-side append latency, TTFT, decode +latency/tokens-per-second, E2E throughput, and server-side hit/promotion deltas. +Reports never persist prompts, token IDs, cache keys, raw addresses, or user +paths. + +Benchmark APIs are public reads and API-key writes: + +```bash +curl -fsS https://kakeya.ai/v1/network/benchmarks +curl -fsS https://kakeya.ai/v1/network/benchmarks/live +curl -fsS https://kakeya.ai/v1/network/benchmarks/ +``` + +The `Benchmarks` dashboard tab shows live progress, phase comparison, history, +and complete redacted stage details. + ## Rollback The cache is an optimization; inference correctness does not depend on it. diff --git a/inference_engine/bench/prefill_fleet_report.py b/inference_engine/bench/prefill_fleet_report.py new file mode 100644 index 00000000..f4b69898 --- /dev/null +++ b/inference_engine/bench/prefill_fleet_report.py @@ -0,0 +1,99 @@ +"""Canonical report schema and aggregation for the two-Mac prefill benchmark.""" +from __future__ import annotations + +import statistics +from typing import Any, Sequence + +PHASES = ("remote_compute", "primary_hot_hit", "allens_cold_restore") +HIT_SOURCES = ("remote_worker", "primary_hot", "allens_offload", "unknown") +_PRIVATE_KEYS = { + "prompt", + "token_ids", + "cache_key", + "block_hash", + "payload_sha256", + "peer_address", + "source_path", +} + + +def normalize_stage(stage: dict[str, Any]) -> dict[str, Any]: + name = stage.get("name") + if name not in PHASES: + raise ValueError(f"unknown benchmark phase {name!r}") + hit_source = stage.get("hit_source", "unknown") + if hit_source not in HIT_SOURCES: + raise ValueError(f"unknown hit_source {hit_source!r}") + output_tokens = int(stage.get("output_tokens", 0)) + prefix_tokens = int(stage.get("prefix_tokens", 0)) + append_s = float(stage.get("append_s", 0.0)) + decode_s = float(stage.get("decode_s", 0.0)) + e2e_s = float(stage.get("e2e_s", 0.0)) + if min(output_tokens, prefix_tokens) < 0 or min(append_s, decode_s, e2e_s) < 0: + raise ValueError("benchmark token and duration values must be non-negative") + normalized = dict(stage) + normalized.update({ + "name": name, + "hit_source": hit_source, + "prefix_tokens": prefix_tokens, + "output_tokens": output_tokens, + "append_s": append_s, + "ttft_s": float(stage.get("ttft_s", 0.0)), + "decode_s": decode_s, + "e2e_s": e2e_s, + "prefill_or_restore_tok_s": ( + prefix_tokens / append_s if append_s > 0 else 0.0 + ), + "decode_tok_s": output_tokens / decode_s if decode_s > 0 else 0.0, + "generation_latency_ms_per_token": ( + decode_s / output_tokens * 1000.0 if output_tokens > 0 else 0.0 + ), + "e2e_tok_s": output_tokens / e2e_s if e2e_s > 0 else 0.0, + }) + assert_public_safe(normalized) + return normalized + + +def summarize_stages(stages: Sequence[dict[str, Any]]) -> dict[str, Any]: + normalized = [normalize_stage(stage) for stage in stages] + sources = {source: 0 for source in HIT_SOURCES} + for stage in normalized: + sources[stage["hit_source"]] += 1 + decode = [stage["decode_tok_s"] for stage in normalized] + return { + "stages_total": len(normalized), + "stages_failed": sum(not stage.get("ok", False) for stage in normalized), + "ttft_p50_s": _median(stage["ttft_s"] for stage in normalized), + "prefill_tok_s_p50": _median( + stage["prefill_or_restore_tok_s"] for stage in normalized + ), + "decode_tok_s_p50": statistics.median(decode) if decode else 0.0, + "e2e_tok_s_p50": _median(stage["e2e_tok_s"] for stage in normalized), + "generation_latency_ms_p50": _median( + stage["generation_latency_ms_per_token"] for stage in normalized + ), + "bytes_received": sum( + int(stage.get("delta", {}).get("bytes_received", 0)) + for stage in normalized + ), + "hit_source_counts": sources, + } + + +def assert_public_safe(value: Any) -> None: + if isinstance(value, dict): + for key, child in value.items(): + if key in _PRIVATE_KEYS: + raise ValueError(f"private benchmark field {key!r} is forbidden") + assert_public_safe(child) + elif isinstance(value, (list, tuple)): + for child in value: + assert_public_safe(child) + elif isinstance(value, str): + if "/Users/" in value or "169.254." in value: + raise ValueError("benchmark report contains a private path or address") + + +def _median(values) -> float: + items = list(values) + return float(statistics.median(items)) if items else 0.0 diff --git a/inference_engine/network/api.py b/inference_engine/network/api.py index 78e2173a..3899ca34 100644 --- a/inference_engine/network/api.py +++ b/inference_engine/network/api.py @@ -40,6 +40,18 @@ class DrainCaptureRequest(BaseModel): max_items: int = Field(default=8, ge=1, le=64) +class BenchmarkCreateRequest(BaseModel): + kind: str = Field(default="distributed_prefill_fleet_benchmark", max_length=100) + config: dict = Field(default_factory=dict) + started_at: Optional[float] = None + + +class BenchmarkUpdateRequest(BaseModel): + stages: list[dict] = Field(default_factory=list) + status: Optional[str] = None + finished_at: Optional[float] = None + + def create_network_app( state: NetworkState, *, @@ -124,6 +136,56 @@ def tokens(): def prefill(): return state.prefill_stats() + @app.get("/v1/network/benchmarks") + def benchmarks(limit: int = 20, status: Optional[str] = None): + try: + return state.list_benchmarks(limit=limit, status=status) + except ValueError as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + + @app.get("/v1/network/benchmarks/live") + def benchmark_live(): + return state.live_benchmark() + + @app.get("/v1/network/benchmarks/{run_id}") + def benchmark_detail(run_id: str): + try: + return state.get_benchmark(run_id) + except KeyError as exc: + raise HTTPException(status_code=404, detail="benchmark not found") from exc + + @app.get("/v1/network/benchmarks/{run_id}/stages") + def benchmark_stages(run_id: str, offset: int = 0, limit: int = 50): + if offset < 0 or not 1 <= limit <= 200: + raise HTTPException(status_code=400, detail="invalid stage pagination") + try: + stages = state.get_benchmark(run_id)["stages"] + except KeyError as exc: + raise HTTPException(status_code=404, detail="benchmark not found") from exc + return {"items": stages[offset:offset + limit], "total": len(stages)} + + @app.post( + "/v1/network/benchmarks", + dependencies=[Depends(require_key)], + ) + def create_benchmark(request: BenchmarkCreateRequest): + try: + return state.create_benchmark(**request.model_dump()) + except ValueError as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + + @app.patch( + "/v1/network/benchmarks/{run_id}", + dependencies=[Depends(require_key)], + ) + def update_benchmark(run_id: str, request: BenchmarkUpdateRequest): + try: + return state.update_benchmark(run_id, **request.model_dump()) + except KeyError as exc: + raise HTTPException(status_code=404, detail="benchmark not found") from exc + except ValueError as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + @app.get( "/v1/network/maintenance/capture", dependencies=[Depends(require_maintenance_key)], diff --git a/inference_engine/network/dashboard.py b/inference_engine/network/dashboard.py index f3d89237..bfcb79f6 100644 --- a/inference_engine/network/dashboard.py +++ b/inference_engine/network/dashboard.py @@ -30,7 +30,7 @@ def dashboard_html() -> str:

Kakeya Inference Network

P2P Prefill KV sharing across trusted inference nodes
-
+
0Online nodes
0Inference groups
0Completed tokens
0%KV-assisted tokens
0 GBShared cache online
0Remote prefill jobs
0Remote KV imports
0Tokens reused
0LRU evictions
0Publish failures
@@ -38,21 +38,24 @@ def dashboard_html() -> str:
+

Exact IPs, raw prompt hashes and cache keys are administrator-only. Region is operator-selected and coarse.

""" diff --git a/inference_engine/network/state.py b/inference_engine/network/state.py index 43a97594..d4fc74b4 100644 --- a/inference_engine/network/state.py +++ b/inference_engine/network/state.py @@ -13,6 +13,11 @@ from inference_engine.distributed.capability import CapabilityRegistry from inference_engine.distributed.kv_namespace import VirtualKVNamespace from inference_engine.distributed.prefill_cache import PrefixCacheStore +from inference_engine.bench.prefill_fleet_report import ( + assert_public_safe, + normalize_stage, + summarize_stages, +) class NetworkState: @@ -94,6 +99,89 @@ def record_tokens( counters["kv_assisted"] += int(kv_assisted) self._save() + def create_benchmark( + self, + *, + kind: str, + config: dict[str, Any], + started_at: float | None = None, + ) -> dict[str, Any]: + assert_public_safe(config) + run = { + "id": "br_" + secrets.token_hex(8), + "schema_version": 1, + "kind": kind, + "status": "running", + "started_at": float(started_at or time.time()), + "finished_at": None, + "config": dict(config), + "stages": [], + "summary": {}, + } + with self._lock: + self._data["benchmark_runs"].append(run) + self._data["benchmark_runs"] = self._data["benchmark_runs"][-200:] + self._data["benchmark_live"] = run["id"] + self._save() + return dict(run) + + def update_benchmark( + self, + run_id: str, + *, + stages: list[dict[str, Any]] | None = None, + status: str | None = None, + finished_at: float | None = None, + ) -> dict[str, Any]: + if status not in (None, "running", "completed", "failed"): + raise ValueError("invalid benchmark status") + with self._lock: + run = self._benchmark_locked(run_id) + if stages: + run["stages"].extend(normalize_stage(stage) for stage in stages) + if status is not None: + run["status"] = status + if finished_at is not None: + run["finished_at"] = float(finished_at) + run["summary"] = summarize_stages(run["stages"]) + if run["status"] != "running" and self._data["benchmark_live"] == run_id: + self._data["benchmark_live"] = None + self._save() + return json.loads(json.dumps(run)) + + def list_benchmarks( + self, + *, + limit: int = 20, + status: str | None = None, + ) -> list[dict[str, Any]]: + if not 1 <= limit <= 200: + raise ValueError("benchmark limit must be in [1, 200]") + with self._lock: + runs = [ + run for run in self._data["benchmark_runs"] + if status is None or run["status"] == status + ][-limit:] + return [ + { + key: run[key] + for key in ( + "id", "schema_version", "kind", "status", + "started_at", "finished_at", "config", "summary", + ) + } + for run in reversed(runs) + ] + + def get_benchmark(self, run_id: str) -> dict[str, Any]: + with self._lock: + return json.loads(json.dumps(self._benchmark_locked(run_id))) + + def live_benchmark(self) -> dict[str, Any] | None: + with self._lock: + run_id = self._data.get("benchmark_live") + return self.get_benchmark(run_id) if run_id else None + def nodes(self) -> list[dict[str, Any]]: registrations = { item["alias"]: item @@ -272,10 +360,24 @@ def _load(self) -> dict[str, Any]: data.setdefault("registrations", []) data.setdefault("groups", []) data.setdefault("tokens", {}) + data.setdefault("benchmark_runs", []) + data.setdefault("benchmark_live", None) return data except (OSError, ValueError): pass - return {"registrations": [], "groups": [], "tokens": {}} + return { + "registrations": [], + "groups": [], + "tokens": {}, + "benchmark_runs": [], + "benchmark_live": None, + } + + def _benchmark_locked(self, run_id: str) -> dict[str, Any]: + for run in self._data["benchmark_runs"]: + if run["id"] == run_id: + return run + raise KeyError(run_id) def _save(self) -> None: self.state_path.parent.mkdir(parents=True, exist_ok=True) diff --git a/scripts/benchmark_prefill_architecture.py b/scripts/benchmark_prefill_architecture.py new file mode 100644 index 00000000..55942d57 --- /dev/null +++ b/scripts/benchmark_prefill_architecture.py @@ -0,0 +1,269 @@ +#!/usr/bin/env python3 +"""One-command three-phase benchmark for Primary decode + allens prefill.""" +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import socket +import subprocess +import time +import urllib.request +import uuid +from pathlib import Path + +PHASE_KEYS = ( + "local_hits", "remote_hits", "remote_jobs", "tokens_reused", + "tokens_computed", "bytes_received", "hot_promotions", + "hot_promotion_bytes", "fallbacks", "remote_job_failures", +) + + +def _json_request(url: str, *, api_key: str = "", method: str = "GET", body=None): + data = None if body is None else json.dumps(body).encode() + headers = {"Content-Type": "application/json"} + if api_key: + headers["X-API-Key"] = api_key + request = urllib.request.Request(url, data=data, headers=headers, method=method) + with urllib.request.urlopen(request, timeout=10) as response: + return json.load(response) + + +def _delta(before: dict, after: dict) -> dict: + return { + key: int(after.get(key, 0)) - int(before.get(key, 0)) + for key in PHASE_KEYS + } + + +def _port_ready(host: str, port: int) -> bool: + try: + with socket.create_connection((host, port), timeout=2): + return True + except OSError: + return False + + +def _wait_ready(host: str, port: int, *, timeout: float = 120.0) -> None: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if _port_ready(host, port): + return + time.sleep(2) + raise RuntimeError(f"service did not become ready on {host}:{port}") + + +def _ensure_services(worker_ssh: str) -> None: + uid = os.getuid() + primary = f"gui/{uid}/ai.kakeya.grpc-runtime-prefill" + if subprocess.run( + ["launchctl", "print", primary], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ).returncode: + plist = Path.home() / "Library/LaunchAgents/ai.kakeya.grpc-runtime-prefill.plist" + subprocess.run(["launchctl", "bootstrap", f"gui/{uid}", str(plist)], check=True) + remote = ( + 'DOMAIN="gui/$(id -u)"; ' + 'launchctl print "$DOMAIN/ai.kakeya.prefill-worker" >/dev/null 2>&1 || ' + 'launchctl bootstrap "$DOMAIN" ' + '"$HOME/Library/LaunchAgents/ai.kakeya.prefill-worker.plist"' + ) + subprocess.run(["ssh", worker_ssh, remote], check=True) + _wait_ready("127.0.0.1", 51051) + _wait_ready("169.254.27.104", 53051) + + +def _restart_primary() -> None: + label = f"gui/{os.getuid()}/ai.kakeya.grpc-runtime-prefill" + subprocess.run(["launchctl", "kickstart", "-k", label], check=True) + _wait_ready("127.0.0.1", 51051, timeout=180) + _wait_ready("127.0.0.1", 8090, timeout=180) + time.sleep(35) + + +def _run_stage(client, eos_ids, token_ids, output_tokens: int, get_stats) -> dict: + before = get_stats() + e2e_start = time.perf_counter() + with client.create_session(eos_token_ids=eos_ids, client_label="fleet-benchmark") as s: + append_start = time.perf_counter() + s.append(token_ids) + append_end = time.perf_counter() + first_token_at = None + count = 0 + for _token in s.generate(max_tokens=output_tokens): + count += 1 + if first_token_at is None: + first_token_at = time.perf_counter() + done_at = time.perf_counter() + after = get_stats() + first_token_at = first_token_at or done_at + return { + "prefix_tokens": len(token_ids), + "output_tokens": count, + "append_s": append_end - append_start, + "ttft_s": first_token_at - e2e_start, + "decode_s": max(0.0, done_at - append_end), + "e2e_s": done_at - e2e_start, + "delta": _delta(before, after), + } + + +def _gate(name: str, stage: dict) -> tuple[bool, str]: + delta = stage["delta"] + if delta["fallbacks"] or delta["remote_job_failures"]: + return False, "fallback_or_remote_failure" + if delta["tokens_computed"] != 0: + return False, "primary_computed_prefill" + if name == "remote_compute": + ok = ( + delta["remote_jobs"] >= 1 + and delta["remote_hits"] >= 1 + and delta["hot_promotions"] >= 1 + ) + return ok, "remote_worker" + if name == "primary_hot_hit": + ok = ( + delta["local_hits"] >= 1 + and delta["remote_hits"] == 0 + and delta["remote_jobs"] == 0 + ) + return ok, "primary_hot" + ok = ( + delta["remote_hits"] >= 1 + and delta["remote_jobs"] == 0 + and delta["hot_promotions"] >= 1 + ) + return ok, "allens_offload" + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--worker-ssh", default="allens") + parser.add_argument("--address", default="127.0.0.1:51051") + parser.add_argument("--dashboard", default="http://127.0.0.1:8090") + parser.add_argument("--api-key-file", default="~/.kakeya/network_api_key") + parser.add_argument("--tokenizer-id", required=True) + parser.add_argument("--minimum-prefix-tokens", type=int, default=128) + parser.add_argument("--output-tokens", type=int, default=32) + parser.add_argument("--report", default="/tmp/kakeya-prefill-benchmark.json") + parser.add_argument("--skip-ensure", action="store_true") + args = parser.parse_args() + + from kakeya import Client + from transformers import AutoTokenizer + from scripts.chat_grpc import _resolve_eos_token_ids + + if not args.skip_ensure: + _ensure_services(args.worker_ssh) + api_key = Path(args.api_key_file).expanduser().read_text().strip() + nodes = _json_request(f"{args.dashboard}/v1/network/nodes") + worker = next( + (node for node in nodes if node["id"] == "allens-mini"), + None, + ) + if not worker or not worker.get("prefill_worker"): + raise SystemExit("allens prefill worker capability is not online") + + tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_id) + nonce = uuid.uuid4().hex + sentence = f"Distributed prefill benchmark run {nonce}. " + text = sentence + token_ids = tokenizer.encode(text, add_special_tokens=True) + while len(token_ids) < args.minimum_prefix_tokens: + text += sentence + token_ids = tokenizer.encode(text, add_special_tokens=True) + prefix_id = hashlib.sha256( + os.urandom(32) + + b"".join(int(t).to_bytes(4, "little") for t in token_ids) + ).hexdigest() + + run = _json_request( + f"{args.dashboard}/v1/network/benchmarks", + api_key=api_key, + method="POST", + body={ + "kind": "distributed_prefill_fleet_benchmark", + "config": { + "model_id": "gemma-4-26B-A4B-it-mlx-4bit", + "topology": "primary-decode-allens-prefill", + "prefill_policy": "remote-required", + "prefix_tokens": len(token_ids), + "output_tokens": args.output_tokens, + "prefix_id": prefix_id, + }, + }, + ) + run_id = run["id"] + stages = [] + + def get_stats(): + return _json_request(f"{args.dashboard}/v1/network/prefill") + + try: + with Client(args.address) as client: + for name in ("remote_compute", "primary_hot_hit"): + stage = _run_stage( + client, + _resolve_eos_token_ids(tokenizer), + token_ids, + args.output_tokens, + get_stats, + ) + stage["name"] = name + stage["ok"], stage["hit_source"] = _gate(name, stage) + stages.append(stage) + _json_request( + f"{args.dashboard}/v1/network/benchmarks/{run_id}", + api_key=api_key, + method="PATCH", + body={"stages": [stage]}, + ) + if not stage["ok"]: + raise RuntimeError(f"phase failed: {name}") + _restart_primary() + with Client(args.address) as client: + stage = _run_stage( + client, + _resolve_eos_token_ids(tokenizer), + token_ids, + args.output_tokens, + get_stats, + ) + stage["name"] = "allens_cold_restore" + stage["ok"], stage["hit_source"] = _gate(stage["name"], stage) + stages.append(stage) + if not stage["ok"]: + raise RuntimeError("phase failed: allens_cold_restore") + completed = _json_request( + f"{args.dashboard}/v1/network/benchmarks/{run_id}", + api_key=api_key, + method="PATCH", + body={ + "stages": [stage], + "status": "completed", + "finished_at": time.time(), + }, + ) + except Exception: + _json_request( + f"{args.dashboard}/v1/network/benchmarks/{run_id}", + api_key=api_key, + method="PATCH", + body={"status": "failed", "finished_at": time.time()}, + ) + raise + + Path(args.report).write_text(json.dumps(completed, indent=2)) + print(json.dumps({ + "ok": True, + "run_id": run_id, + "report": args.report, + "summary": completed["summary"], + }, indent=2)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/run_prefill_architecture_benchmark.sh b/scripts/run_prefill_architecture_benchmark.sh new file mode 100644 index 00000000..30f73c11 --- /dev/null +++ b/scripts/run_prefill_architecture_benchmark.sh @@ -0,0 +1,11 @@ +#!/usr/bin/env bash +set -euo pipefail + +REPO_ROOT="$(cd "$(dirname "$0")/.." && pwd)" +PYTHON="${KAKEYA_BENCH_PYTHON:-$HOME/.venv-distwan/bin/python}" +MODEL="${KAKEYA_BENCH_MODEL:-$HOME/kakeya-models/gemma-4-26B-A4B-it-mlx-4bit}" + +exec env PYTHONPATH="$REPO_ROOT:$REPO_ROOT/sdks/python" \ + "$PYTHON" "$REPO_ROOT/scripts/benchmark_prefill_architecture.py" \ + --tokenizer-id "$MODEL" \ + "$@" diff --git a/tests/inference_engine/bench/test_prefill_fleet_report.py b/tests/inference_engine/bench/test_prefill_fleet_report.py new file mode 100644 index 00000000..dff0a5af --- /dev/null +++ b/tests/inference_engine/bench/test_prefill_fleet_report.py @@ -0,0 +1,67 @@ +import pytest + +from inference_engine.bench.prefill_fleet_report import ( + assert_public_safe, + normalize_stage, + summarize_stages, +) + + +def _stage(name="remote_compute", source="remote_worker"): + return { + "name": name, + "hit_source": source, + "ok": True, + "prefix_tokens": 100, + "output_tokens": 10, + "append_s": 5.0, + "ttft_s": 5.2, + "decode_s": 2.0, + "e2e_s": 7.0, + "delta": {"bytes_received": 1000}, + } + + +def test_normalize_stage_derives_throughput_and_latency(): + stage = normalize_stage(_stage()) + assert stage["prefill_or_restore_tok_s"] == 20.0 + assert stage["decode_tok_s"] == 5.0 + assert stage["generation_latency_ms_per_token"] == 200.0 + assert stage["e2e_tok_s"] == 10 / 7 + + +def test_summary_aggregates_sources_and_medians(): + summary = summarize_stages([ + _stage(), + _stage("primary_hot_hit", "primary_hot"), + {**_stage("allens_cold_restore", "allens_offload"), "ok": False}, + ]) + assert summary["stages_total"] == 3 + assert summary["stages_failed"] == 1 + assert summary["hit_source_counts"]["remote_worker"] == 1 + assert summary["hit_source_counts"]["primary_hot"] == 1 + assert summary["bytes_received"] == 3000 + assert summarize_stages([])["decode_tok_s_p50"] == 0 + + +def test_schema_rejects_unknown_and_private_fields(): + with pytest.raises(ValueError, match="unknown benchmark phase"): + normalize_stage(_stage("bad")) + with pytest.raises(ValueError, match="unknown hit_source"): + normalize_stage(_stage(source="peer:1")) + with pytest.raises(ValueError, match="non-negative"): + normalize_stage({**_stage(), "append_s": -1}) + with pytest.raises(ValueError, match="private benchmark field"): + assert_public_safe({"prompt": "secret"}) + with pytest.raises(ValueError, match="private path"): + assert_public_safe({"value": "/Users/private/model"}) + with pytest.raises(ValueError, match="private path"): + assert_public_safe(["169.254.27.104"]) + assert normalize_stage({ + **_stage(), + "output_tokens": 0, + "prefix_tokens": 0, + "append_s": 0, + "decode_s": 0, + "e2e_s": 0, + })["decode_tok_s"] == 0 diff --git a/tests/inference_engine/bridge/test_prefill_architecture_benchmark.py b/tests/inference_engine/bridge/test_prefill_architecture_benchmark.py new file mode 100644 index 00000000..538bb178 --- /dev/null +++ b/tests/inference_engine/bridge/test_prefill_architecture_benchmark.py @@ -0,0 +1,46 @@ +from scripts.benchmark_prefill_architecture import _delta, _gate + + +def _stage(**delta): + values = { + "remote_jobs": 0, + "remote_hits": 0, + "local_hits": 0, + "tokens_reused": 0, + "tokens_computed": 0, + "hot_promotions": 0, + "fallbacks": 0, + "remote_job_failures": 0, + } + values.update(delta) + return {"delta": values} + + +def test_phase_gates_cover_worker_hot_and_offload_paths(): + assert _gate("remote_compute", _stage( + remote_jobs=1, remote_hits=1, hot_promotions=1, + )) == (True, "remote_worker") + assert _gate("primary_hot_hit", _stage(local_hits=1)) == ( + True, "primary_hot", + ) + assert _gate("allens_cold_restore", _stage( + remote_hits=1, hot_promotions=1, + )) == (True, "allens_offload") + + +def test_phase_gates_reject_primary_prefill_and_failures(): + assert _gate("remote_compute", _stage(tokens_computed=1))[0] is False + assert _gate("remote_compute", _stage(fallbacks=1))[1] == ( + "fallback_or_remote_failure" + ) + assert _gate("primary_hot_hit", _stage(remote_hits=1))[0] is False + assert _gate("allens_cold_restore", _stage(remote_jobs=1))[0] is False + + +def test_metric_delta_uses_known_phase_keys(): + before = {"local_hits": 1, "remote_hits": 2} + after = {"local_hits": 3, "remote_hits": 2} + result = _delta(before, after) + assert result["local_hits"] == 2 + assert result["remote_hits"] == 0 + assert result["fallbacks"] == 0 diff --git a/tests/inference_engine/network/test_network_api.py b/tests/inference_engine/network/test_network_api.py index 5faead90..e58d5aa9 100644 --- a/tests/inference_engine/network/test_network_api.py +++ b/tests/inference_engine/network/test_network_api.py @@ -43,6 +43,7 @@ def test_dashboard_health_and_read_apis(tmp_path): assert "Kakeya Inference Network" in dashboard assert "Remote prefill jobs" in dashboard assert "LRU evictions" in dashboard + assert "Benchmarks" in dashboard assert client.get("/healthz").json()["status"] == "ok" assert client.get("/v1/network/summary").json()["online_nodes"] == 1 assert len(client.get("/v1/network/nodes").json()) == 1 @@ -153,3 +154,60 @@ def test_disabled_capture_and_missing_maintenance_key(tmp_path): cache_fill_capture=CacheFillCapture(), )) assert without_key.get("/v1/network/maintenance/capture").status_code == 401 + + +def test_benchmark_api_create_update_list_detail_and_pagination(tmp_path): + client = _client(tmp_path) + assert client.post( + "/v1/network/benchmarks", + json={"kind": "test", "config": {}}, + ).status_code == 401 + created = client.post( + "/v1/network/benchmarks", + json={"kind": "distributed_prefill_fleet_benchmark", "config": {}}, + headers={"X-API-Key": "secret"}, + ).json() + run_id = created["id"] + assert client.get("/v1/network/benchmarks/live").json()["id"] == run_id + stage = { + "name": "remote_compute", + "hit_source": "remote_worker", + "ok": True, + "prefix_tokens": 10, + "output_tokens": 2, + "append_s": 1, + "ttft_s": 1.1, + "decode_s": 0.5, + "e2e_s": 1.5, + "delta": {}, + } + updated = client.patch( + f"/v1/network/benchmarks/{run_id}", + json={"stages": [stage], "status": "completed", "finished_at": 2}, + headers={"X-API-Key": "secret"}, + ) + assert updated.status_code == 200 + assert client.get("/v1/network/benchmarks/live").json() is None + assert client.get("/v1/network/benchmarks").json()[0]["id"] == run_id + assert client.get(f"/v1/network/benchmarks/{run_id}").json()["stages"][0]["ok"] + page = client.get(f"/v1/network/benchmarks/{run_id}/stages?offset=0&limit=1") + assert page.json()["total"] == 1 + assert client.get("/v1/network/benchmarks?limit=0").status_code == 400 + assert client.get(f"/v1/network/benchmarks/{run_id}/stages?offset=-1").status_code == 400 + assert client.get("/v1/network/benchmarks/missing").status_code == 404 + assert client.get("/v1/network/benchmarks/missing/stages").status_code == 404 + assert client.patch( + "/v1/network/benchmarks/missing", + json={}, + headers={"X-API-Key": "secret"}, + ).status_code == 404 + assert client.patch( + f"/v1/network/benchmarks/{run_id}", + json={"status": "invalid"}, + headers={"X-API-Key": "secret"}, + ).status_code == 400 + assert client.post( + "/v1/network/benchmarks", + json={"kind": "bad", "config": {"prompt": "secret"}}, + headers={"X-API-Key": "secret"}, + ).status_code == 400 diff --git a/tests/inference_engine/network/test_network_state.py b/tests/inference_engine/network/test_network_state.py index 8e6c1649..70ee5c4b 100644 --- a/tests/inference_engine/network/test_network_state.py +++ b/tests/inference_engine/network/test_network_state.py @@ -127,3 +127,48 @@ def test_prefill_stats_serializes_runtime_dataclass(tmp_path): tokens_reused=256, ) assert state.prefill_stats()["remote_jobs"] == 4 + + +def test_benchmark_lifecycle_persistence_and_retention(tmp_path): + state = _state(tmp_path) + run = state.create_benchmark( + kind="distributed_prefill_fleet_benchmark", + config={"model_id": "gemma"}, + started_at=10, + ) + assert state.live_benchmark()["id"] == run["id"] + stage = { + "name": "remote_compute", + "hit_source": "remote_worker", + "ok": True, + "prefix_tokens": 100, + "output_tokens": 10, + "append_s": 5, + "ttft_s": 5.1, + "decode_s": 2, + "e2e_s": 7, + "delta": {}, + } + completed = state.update_benchmark( + run["id"], + stages=[stage], + status="completed", + finished_at=20, + ) + assert completed["summary"]["decode_tok_s_p50"] == 5 + assert state.live_benchmark() is None + assert state.list_benchmarks(limit=1)[0]["id"] == run["id"] + assert state.list_benchmarks(status="completed")[0]["status"] == "completed" + assert _state(tmp_path).get_benchmark(run["id"])["finished_at"] == 20 + with __import__("pytest").raises(ValueError): + state.list_benchmarks(limit=0) + with __import__("pytest").raises(ValueError): + state.update_benchmark(run["id"], status="invalid") + with __import__("pytest").raises(KeyError): + state.get_benchmark("missing") + with __import__("pytest").raises(ValueError): + state.create_benchmark(kind="x", config={"prompt": "private"}) + + for index in range(205): + state.create_benchmark(kind="retention", config={"index": index}) + assert len(state._data["benchmark_runs"]) == 200