752 lines
29 KiB
Python
752 lines
29 KiB
Python
"""Real runtime drivers for the recipe benchmark.
|
|
|
|
This module is the only place that imports torch, transformers, or spawns a
|
|
llama.cpp server, and :mod:`meshnet_node.recipe_benchmark` imports it lazily.
|
|
That keeps the default test suite deterministic, GPU-free and download-free
|
|
while the real evidence runs through exactly the same measurement core.
|
|
|
|
Fairness is the whole point of a baseline, so both drivers are held to the same
|
|
rules:
|
|
|
|
* They are handed a **pre-formatted prompt string**. Neither applies a chat
|
|
template, because a template applied twice — or differently — by two runtimes
|
|
would show up as a speed and drift difference that has nothing to do with the
|
|
runtime.
|
|
* They are given the **same CPU thread budget**, so the comparison measures
|
|
kernels rather than how many cores each runtime felt entitled to take.
|
|
* They report the runtime's **own prefill/decode split** where it has one, and
|
|
say so honestly where it does not.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import hashlib
|
|
import hmac
|
|
import json
|
|
import os
|
|
import platform
|
|
import re
|
|
import socket
|
|
import stat
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
import urllib.error
|
|
import urllib.request
|
|
import uuid
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
from typing import Any, Mapping
|
|
|
|
from cryptography.hazmat.primitives import serialization
|
|
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
|
|
|
|
from .performance_contract import (
|
|
PROVENANCE_SCHEMA_VERSION,
|
|
REAL_REPORT_PRODUCER,
|
|
_canonical_sha256,
|
|
report_signing_payload,
|
|
)
|
|
from .recipe_benchmark import (
|
|
BenchmarkError,
|
|
BenchmarkPlan,
|
|
GenerationSample,
|
|
Lane,
|
|
LoadStats,
|
|
PromptSpec,
|
|
RecipeSpec,
|
|
SamplingPolicy,
|
|
build_report,
|
|
measure_recipe,
|
|
)
|
|
|
|
REAL_INFERENCE_ENV = "MESHNET_ENABLE_REAL_INFERENCE_TESTS"
|
|
EVIDENCE_SIGNING_KEY_ENV = "MESHNET_EVIDENCE_SIGNING_KEY"
|
|
|
|
|
|
def real_inference_enabled() -> bool:
|
|
"""Real runtimes stay off unless the operator opts in explicitly."""
|
|
return os.environ.get(REAL_INFERENCE_ENV) == "1"
|
|
|
|
|
|
def require_real_inference() -> None:
|
|
if not real_inference_enabled():
|
|
raise BenchmarkError(
|
|
f"real model execution is opt-in: set {REAL_INFERENCE_ENV}=1 to run this benchmark"
|
|
)
|
|
|
|
|
|
def _load_evidence_signing_key() -> Ed25519PrivateKey:
|
|
raw_path = os.environ.get(EVIDENCE_SIGNING_KEY_ENV)
|
|
if not raw_path:
|
|
raise BenchmarkError(
|
|
f"real evidence requires {EVIDENCE_SIGNING_KEY_ENV} to name an Ed25519 private key"
|
|
)
|
|
path = Path(raw_path).expanduser().resolve(strict=True)
|
|
if os.name != "nt" and stat.S_IMODE(path.stat().st_mode) != 0o600:
|
|
raise BenchmarkError("evidence signing key must have mode 0600")
|
|
key = serialization.load_pem_private_key(path.read_bytes(), password=None)
|
|
if not isinstance(key, Ed25519PrivateKey):
|
|
raise BenchmarkError("evidence signing key must be Ed25519")
|
|
return key
|
|
|
|
|
|
def _utc_now() -> str:
|
|
return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
|
|
|
|
|
|
def _sign_report(report: dict[str, Any], key: Ed25519PrivateKey) -> None:
|
|
public_key = key.public_key().public_bytes(
|
|
serialization.Encoding.Raw, serialization.PublicFormat.Raw
|
|
)
|
|
report["provenance"]["signer_public_key_sha256"] = hashlib.sha256(
|
|
public_key
|
|
).hexdigest()
|
|
report["provenance"]["signature"] = base64.b64encode(
|
|
key.sign(report_signing_payload(report))
|
|
).decode("ascii")
|
|
|
|
|
|
def _process_rss(pid: int | None = None) -> int:
|
|
"""Resident bytes for a process and its children, or 0 when unobservable."""
|
|
try:
|
|
import psutil
|
|
except ImportError:
|
|
return 0
|
|
try:
|
|
process = psutil.Process(pid) if pid else psutil.Process()
|
|
total = process.memory_info().rss
|
|
for child in process.children(recursive=True):
|
|
try:
|
|
total += child.memory_info().rss
|
|
except psutil.Error:
|
|
continue
|
|
return int(total)
|
|
except Exception:
|
|
return 0
|
|
|
|
|
|
def _directory_bytes(path: Path) -> int:
|
|
if path.is_file():
|
|
return path.stat().st_size
|
|
return sum(entry.stat().st_size for entry in path.rglob("*") if entry.is_file())
|
|
|
|
|
|
def _artifact_sha256(path: Path) -> str:
|
|
"""Hash an artifact file or a deterministic directory content manifest.
|
|
|
|
A file uses the ordinary SHA-256 digest. A directory hashes each sorted
|
|
relative path, resolved file size, and file bytes, so tokenizer/config drift
|
|
cannot hide behind a weight-only digest.
|
|
"""
|
|
digest = hashlib.sha256()
|
|
if path.is_file():
|
|
entries = [(None, path)]
|
|
else:
|
|
entries = [
|
|
(entry.relative_to(path).as_posix(), entry)
|
|
for entry in sorted(path.rglob("*"))
|
|
if entry.is_file()
|
|
]
|
|
if not entries:
|
|
raise BenchmarkError(f"artifact directory is empty: {path}")
|
|
|
|
for relative, entry in entries:
|
|
if relative is not None:
|
|
encoded = relative.encode("utf-8")
|
|
digest.update(len(encoded).to_bytes(8, "big"))
|
|
digest.update(encoded)
|
|
digest.update(entry.stat().st_size.to_bytes(8, "big"))
|
|
with entry.open("rb") as stream:
|
|
while chunk := stream.read(8 * 1024 * 1024):
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def _host_manifest(config: Mapping[str, Any] | None = None) -> dict[str, Any]:
|
|
"""Capture non-secret host facts with the report rather than trusting prose."""
|
|
manifest: dict[str, Any] = {
|
|
"hostname": socket.gethostname(),
|
|
"platform": platform.platform(),
|
|
"python": sys.version.split()[0],
|
|
"cpu_count": os.cpu_count(),
|
|
}
|
|
try:
|
|
import torch
|
|
import transformers
|
|
|
|
manifest["torch_version"] = torch.__version__
|
|
manifest["transformers_version"] = transformers.__version__
|
|
manifest["cuda_available"] = bool(torch.cuda.is_available())
|
|
if torch.cuda.is_available():
|
|
manifest["accelerator_name"] = torch.cuda.get_device_name(0)
|
|
manifest["accelerator_runtime"] = getattr(torch.version, "cuda", None) or getattr(
|
|
torch.version, "hip", None
|
|
)
|
|
except ImportError:
|
|
manifest["torch_version"] = None
|
|
|
|
llama_identities: dict[str, dict[str, str]] = {}
|
|
for spec in (config or {}).get("recipes", ()):
|
|
driver = spec.get("driver", {})
|
|
if driver.get("type") != "llama-cpp-server":
|
|
continue
|
|
binary = Path(driver["binary"]).resolve(strict=True)
|
|
key = str(binary)
|
|
if key in llama_identities:
|
|
continue
|
|
version_result = subprocess.run(
|
|
[str(binary), "--version"],
|
|
check=True,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.STDOUT,
|
|
text=True,
|
|
timeout=10,
|
|
)
|
|
llama_identities[key] = {
|
|
"sha256": _artifact_sha256(binary),
|
|
"version": " | ".join(version_result.stdout.strip().splitlines()),
|
|
}
|
|
if llama_identities:
|
|
manifest["llama_server_identities"] = llama_identities
|
|
return manifest
|
|
|
|
|
|
def _validate_config(config: Mapping[str, Any]) -> None:
|
|
"""Reject comparisons that mix models, artifacts, devices, or budgets."""
|
|
try:
|
|
plan = config["plan"]
|
|
root = Path(config["artifact_storage_root"]).resolve(strict=True)
|
|
recipes = config["recipes"]
|
|
except (KeyError, TypeError, OSError) as exc:
|
|
raise BenchmarkError(
|
|
"benchmark config needs an existing artifact_storage_root, plan, and recipes"
|
|
) from exc
|
|
if not root.is_absolute() or root == Path("/home") or Path("/home") in root.parents:
|
|
raise BenchmarkError("model artifacts must use configured mounted-drive storage, never /home")
|
|
if not isinstance(recipes, list) or not recipes:
|
|
raise BenchmarkError("benchmark config needs at least one recipe")
|
|
|
|
sampling = plan.get("sampling", {})
|
|
if (
|
|
float(sampling.get("temperature", 0.0)) != 0.0
|
|
or int(sampling.get("top_k", 1)) != 1
|
|
or float(sampling.get("top_p", 1.0)) != 1.0
|
|
):
|
|
raise BenchmarkError("the quality comparison requires greedy sampling")
|
|
|
|
if len(plan.get("prompts", ())) < 3 or int(plan.get("repeats", 0)) < 3:
|
|
raise BenchmarkError("contract-grade evidence requires at least 3 prompts and 3 repeats")
|
|
if int(plan.get("warmup_requests", 0)) < 1:
|
|
raise BenchmarkError("contract-grade evidence requires at least one warmup")
|
|
if int(sampling.get("max_output_tokens", 0)) < 32:
|
|
raise BenchmarkError("contract-grade evidence requires at least 32 output tokens")
|
|
|
|
thread_budgets: set[int] = set()
|
|
max_concurrency = max(int(level) for level in plan.get("concurrency_levels", (1, 4)))
|
|
for spec in recipes:
|
|
if spec.get("source_model_id") != plan.get("model_id"):
|
|
raise BenchmarkError("every recipe must declare the plan's exact source_model_id")
|
|
if spec.get("source_model_revision") != plan.get("model_revision"):
|
|
raise BenchmarkError("every recipe must declare the plan's exact source_model_revision")
|
|
|
|
digest = spec.get("artifact_sha256", "")
|
|
if not isinstance(digest, str) or re.fullmatch(r"[0-9a-f]{64}", digest) is None:
|
|
raise BenchmarkError("every recipe must declare a lowercase SHA-256 artifact digest")
|
|
artifact = Path(spec.get("artifact_path", "")).resolve(strict=True)
|
|
if artifact != root and root not in artifact.parents:
|
|
raise BenchmarkError("every model artifact must be beneath artifact_storage_root")
|
|
actual_digest = _artifact_sha256(artifact)
|
|
if not hmac.compare_digest(digest, actual_digest):
|
|
raise BenchmarkError(
|
|
f"artifact digest mismatch for {spec.get('id', '<unknown>')}: "
|
|
f"declared {digest}, measured {actual_digest}"
|
|
)
|
|
|
|
driver = spec.get("driver")
|
|
if not isinstance(driver, Mapping):
|
|
raise BenchmarkError("every recipe needs a driver object")
|
|
kind = driver.get("type")
|
|
if kind == "transformers":
|
|
driver_artifact = Path(driver.get("model_path", "")).resolve(strict=True)
|
|
elif kind == "llama-cpp-server":
|
|
driver_artifact = Path(driver.get("gguf_path", "")).resolve(strict=True)
|
|
binary = Path(driver.get("binary", "")).resolve(strict=True)
|
|
binary_digest = driver.get("binary_sha256", "")
|
|
if (
|
|
not isinstance(binary_digest, str)
|
|
or re.fullmatch(r"[0-9a-f]{64}", binary_digest) is None
|
|
or not hmac.compare_digest(binary_digest, _artifact_sha256(binary))
|
|
):
|
|
raise BenchmarkError("llama.cpp binary SHA-256 mismatch")
|
|
if int(driver.get("n_parallel", max_concurrency)) < max_concurrency:
|
|
raise BenchmarkError("llama.cpp parallel slots must cover maximum concurrency")
|
|
if driver.get("device", "cpu") != "cpu" or int(driver.get("n_gpu_layers", 0)) != 0:
|
|
raise BenchmarkError(
|
|
"v1 benchmark supports CPU-only llama.cpp until process VRAM is measurable"
|
|
)
|
|
else:
|
|
raise BenchmarkError(f"unknown driver type {kind!r}")
|
|
if driver_artifact != artifact:
|
|
raise BenchmarkError("driver artifact path must match the hashed recipe artifact")
|
|
if driver.get("device", "cpu") != spec.get("device"):
|
|
raise BenchmarkError("recipe and driver must declare the same device")
|
|
thread_budgets.add(int(driver.get("threads", 8)))
|
|
|
|
if len(thread_budgets) != 1:
|
|
raise BenchmarkError("every recipe must use the same CPU thread budget")
|
|
|
|
|
|
class TransformersDriver:
|
|
"""The current Transformers/safetensors recipe: the correctness reference.
|
|
|
|
Generation is a hand-written prefill-then-decode loop rather than
|
|
``model.generate`` because the benchmark needs the two phases separated: one
|
|
forward over the prompt gives an exact prefill time and TTFT, and the cached
|
|
single-token steps that follow give an exact decode rate. ``generate`` would
|
|
hand back one blended number.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
model_path: str,
|
|
*,
|
|
device: str = "cpu",
|
|
dtype: str = "bfloat16",
|
|
threads: int = 8,
|
|
) -> None:
|
|
self.model_path = Path(model_path)
|
|
self.device = device
|
|
self.dtype = dtype
|
|
self.threads = threads
|
|
self._model: Any = None
|
|
self._tokenizer: Any = None
|
|
self._torch: Any = None
|
|
self._rss_baseline = 0
|
|
|
|
def load(self) -> LoadStats:
|
|
self._rss_baseline = _process_rss()
|
|
import torch
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
|
|
self._torch = torch
|
|
torch.set_num_threads(self.threads)
|
|
torch.manual_seed(0)
|
|
|
|
started = time.monotonic()
|
|
self._tokenizer = AutoTokenizer.from_pretrained(
|
|
str(self.model_path), local_files_only=True
|
|
)
|
|
self._model = AutoModelForCausalLM.from_pretrained(
|
|
str(self.model_path),
|
|
dtype=getattr(torch, self.dtype),
|
|
local_files_only=True,
|
|
)
|
|
self._model.to(self.device)
|
|
self._model.eval()
|
|
load_ms = (time.monotonic() - started) * 1000
|
|
|
|
return LoadStats(
|
|
artifact_bytes=_directory_bytes(self.model_path),
|
|
load_ms=round(load_ms, 4),
|
|
rss_bytes=max(0, _process_rss() - self._rss_baseline),
|
|
vram_bytes=self._vram_bytes(),
|
|
backend_detail=(
|
|
f"torch {torch.__version__}; dtype {self.dtype}; "
|
|
f"device {self.device}; intra-op threads {self.threads}"
|
|
),
|
|
)
|
|
|
|
def _vram_bytes(self) -> int:
|
|
torch = self._torch
|
|
if torch is None or self.device == "cpu":
|
|
return 0
|
|
try:
|
|
if torch.cuda.is_available():
|
|
return int(torch.cuda.max_memory_allocated())
|
|
except Exception:
|
|
return 0
|
|
return 0
|
|
|
|
def generate(self, prompt: str, sampling: SamplingPolicy) -> GenerationSample:
|
|
if self._model is None:
|
|
raise BenchmarkError("TransformersDriver.generate called before load()")
|
|
torch = self._torch
|
|
|
|
# add_special_tokens=False: the plan owns the prompt format, and the
|
|
# llama.cpp recipe is given the identical string.
|
|
encoded = self._tokenizer(prompt, return_tensors="pt", add_special_tokens=False)
|
|
input_ids = encoded["input_ids"].to(self.device)
|
|
prompt_tokens = int(input_ids.shape[-1])
|
|
eos_ids = {self._tokenizer.eos_token_id} | set(
|
|
getattr(self._model.generation_config, "eos_token_id", None) or []
|
|
if isinstance(getattr(self._model.generation_config, "eos_token_id", None), list)
|
|
else []
|
|
)
|
|
eos_ids.discard(None)
|
|
|
|
started = time.monotonic()
|
|
with torch.inference_mode():
|
|
outputs = self._model(input_ids=input_ids, use_cache=True)
|
|
past = outputs.past_key_values
|
|
next_id = self._select(outputs.logits[:, -1, :], sampling)
|
|
ttft_ms = (time.monotonic() - started) * 1000
|
|
|
|
token_ids = [int(next_id.item())]
|
|
decode_started = time.monotonic()
|
|
while len(token_ids) < sampling.max_output_tokens and token_ids[-1] not in eos_ids:
|
|
outputs = self._model(
|
|
input_ids=next_id.view(1, 1), past_key_values=past, use_cache=True
|
|
)
|
|
past = outputs.past_key_values
|
|
next_id = self._select(outputs.logits[:, -1, :], sampling)
|
|
token_ids.append(int(next_id.item()))
|
|
decode_ms = (time.monotonic() - decode_started) * 1000
|
|
|
|
total_ms = (time.monotonic() - started) * 1000
|
|
emitted = [token for token in token_ids if token not in eos_ids]
|
|
return GenerationSample(
|
|
text=self._tokenizer.decode(emitted, skip_special_tokens=True),
|
|
prompt_tokens=prompt_tokens,
|
|
# The first token is produced by the prefill forward, so the decode
|
|
# rate must not be credited with it.
|
|
decode_tokens=max(0, len(token_ids) - 1),
|
|
ttft_ms=ttft_ms,
|
|
prefill_ms=ttft_ms,
|
|
decode_ms=decode_ms,
|
|
total_ms=total_ms,
|
|
)
|
|
|
|
def _select(self, logits: Any, sampling: SamplingPolicy) -> Any:
|
|
if sampling.temperature > 0:
|
|
raise BenchmarkError(
|
|
"this benchmark is greedy-only: sampling noise is indistinguishable from "
|
|
"quantization drift, which is precisely what the quality lane must isolate"
|
|
)
|
|
return logits.argmax(dim=-1)
|
|
|
|
def memory_probe(self) -> tuple[int, int]:
|
|
return max(0, _process_rss() - self._rss_baseline), self._vram_bytes()
|
|
|
|
def close(self) -> None:
|
|
self._model = None
|
|
self._tokenizer = None
|
|
if self._torch is not None:
|
|
import gc
|
|
|
|
gc.collect()
|
|
|
|
|
|
def _free_port() -> int:
|
|
with socket.socket() as probe:
|
|
probe.bind(("127.0.0.1", 0))
|
|
return int(probe.getsockname()[1])
|
|
|
|
|
|
class LlamaCppServerDriver:
|
|
"""The whole-model llama.cpp/GGUF recipe, driven through ``llama-server``.
|
|
|
|
``llama-server`` is used rather than an in-process binding because it is the
|
|
shape llama.cpp is actually deployed in and the only one that offers
|
|
continuous batching across parallel slots — which is the runtime property
|
|
this project cares about most. It also reports its own prefill/decode
|
|
timings per request, so the decode rate is the runtime's own number and not
|
|
an inference drawn from a client-side stopwatch.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
binary: str,
|
|
gguf_path: str,
|
|
*,
|
|
binary_sha256: str,
|
|
device: str = "cpu",
|
|
threads: int = 8,
|
|
n_parallel: int = 4,
|
|
context_per_slot: int = 1024,
|
|
n_gpu_layers: int = 0,
|
|
startup_timeout_s: float = 120.0,
|
|
) -> None:
|
|
self.binary = Path(binary)
|
|
self.binary_sha256 = binary_sha256
|
|
self.gguf_path = Path(gguf_path)
|
|
self.device = device
|
|
self.threads = threads
|
|
self.n_parallel = n_parallel
|
|
self.context_per_slot = context_per_slot
|
|
self.n_gpu_layers = n_gpu_layers
|
|
self.startup_timeout_s = startup_timeout_s
|
|
self._process: subprocess.Popen | None = None
|
|
self._port = 0
|
|
self._log: Any = None
|
|
|
|
@property
|
|
def _url(self) -> str:
|
|
return f"http://127.0.0.1:{self._port}"
|
|
|
|
def _log_excerpt(self) -> str:
|
|
if self._log is None:
|
|
return ""
|
|
try:
|
|
self._log.flush()
|
|
self._log.seek(0)
|
|
return self._log.read()[-4096:].decode("utf-8", errors="replace").strip()
|
|
except Exception:
|
|
return ""
|
|
|
|
def load(self) -> LoadStats:
|
|
if not self.binary.exists():
|
|
raise BenchmarkError(f"llama-server binary not found at {self.binary}")
|
|
if not self.gguf_path.exists():
|
|
raise BenchmarkError(f"GGUF artifact not found at {self.gguf_path}")
|
|
measured_binary_sha256 = _artifact_sha256(self.binary)
|
|
if not hmac.compare_digest(self.binary_sha256, measured_binary_sha256):
|
|
raise BenchmarkError("llama-server binary changed after config validation")
|
|
version = " | ".join(
|
|
subprocess.run(
|
|
[str(self.binary), "--version"],
|
|
check=True,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.STDOUT,
|
|
text=True,
|
|
timeout=10,
|
|
).stdout.strip().splitlines()
|
|
)
|
|
|
|
self._port = _free_port()
|
|
command = [
|
|
str(self.binary),
|
|
"--model", str(self.gguf_path),
|
|
"--host", "127.0.0.1",
|
|
"--port", str(self._port),
|
|
"--threads", str(self.threads),
|
|
"--parallel", str(self.n_parallel),
|
|
# Every slot must hold a whole request, so the pool is sized for the
|
|
# worst case rather than letting llama.cpp silently truncate context.
|
|
"--ctx-size", str(self.context_per_slot * self.n_parallel),
|
|
"--n-gpu-layers", str(self.n_gpu_layers),
|
|
"--no-webui",
|
|
]
|
|
started = time.monotonic()
|
|
self._log = tempfile.TemporaryFile(mode="w+b")
|
|
self._process = subprocess.Popen(
|
|
command, stdout=self._log, stderr=subprocess.STDOUT
|
|
)
|
|
self._await_health(started)
|
|
load_ms = (time.monotonic() - started) * 1000
|
|
|
|
return LoadStats(
|
|
artifact_bytes=self.gguf_path.stat().st_size,
|
|
load_ms=round(load_ms, 4),
|
|
rss_bytes=_process_rss(self._process.pid),
|
|
vram_bytes=0,
|
|
backend_detail=(
|
|
f"{version}; binary sha256 {measured_binary_sha256}; "
|
|
f"threads {self.threads}; parallel slots {self.n_parallel}; "
|
|
f"ctx/slot {self.context_per_slot}; gpu layers {self.n_gpu_layers}"
|
|
),
|
|
)
|
|
|
|
def _await_health(self, started: float) -> None:
|
|
while time.monotonic() - started < self.startup_timeout_s:
|
|
if self._process is not None and self._process.poll() is not None:
|
|
raise BenchmarkError(
|
|
f"llama-server exited with code {self._process.returncode} during startup; "
|
|
f"log tail: {self._log_excerpt()}"
|
|
)
|
|
try:
|
|
with urllib.request.urlopen(f"{self._url}/health", timeout=2) as response:
|
|
if response.status == 200:
|
|
return
|
|
except (urllib.error.URLError, OSError):
|
|
time.sleep(0.25)
|
|
raise BenchmarkError(
|
|
f"llama-server did not become healthy within {self.startup_timeout_s:.0f}s; "
|
|
f"log tail: {self._log_excerpt()}"
|
|
)
|
|
|
|
def generate(self, prompt: str, sampling: SamplingPolicy) -> GenerationSample:
|
|
if self._process is None:
|
|
raise BenchmarkError("LlamaCppServerDriver.generate called before load()")
|
|
if sampling.temperature > 0:
|
|
raise BenchmarkError("this benchmark is greedy-only; see TransformersDriver._select")
|
|
|
|
body = json.dumps({
|
|
"prompt": prompt,
|
|
"n_predict": sampling.max_output_tokens,
|
|
"temperature": 0.0,
|
|
"top_k": 1,
|
|
"top_p": 1.0,
|
|
"seed": sampling.seed,
|
|
# Prompt cache reuse across repeats would measure the cache, not the
|
|
# prefill, and the safetensors recipe has no equivalent.
|
|
"cache_prompt": False,
|
|
"stream": True,
|
|
}).encode()
|
|
request = urllib.request.Request(
|
|
f"{self._url}/completion", data=body,
|
|
headers={"Content-Type": "application/json"}, method="POST",
|
|
)
|
|
|
|
started = time.monotonic()
|
|
chunks: list[str] = []
|
|
timings: Mapping[str, Any] = {}
|
|
with urllib.request.urlopen(request, timeout=600) as response:
|
|
for raw in response:
|
|
line = raw.decode("utf-8").strip()
|
|
if not line.startswith("data:"):
|
|
continue
|
|
payload = json.loads(line[len("data:"):].strip())
|
|
content = payload.get("content", "")
|
|
chunks.append(content)
|
|
if payload.get("stop"):
|
|
timings = payload.get("timings") or {}
|
|
total_ms = (time.monotonic() - started) * 1000
|
|
|
|
if not timings:
|
|
raise BenchmarkError("llama-server returned no timings; cannot report an honest split")
|
|
|
|
prefill_ms = float(timings.get("prompt_ms", 0.0))
|
|
decode_ms = float(timings.get("predicted_ms", 0.0))
|
|
return GenerationSample(
|
|
text="".join(chunks),
|
|
prompt_tokens=int(timings.get("prompt_n", 0)),
|
|
# llama.cpp starts predicted_ms after sampling the first token while
|
|
# predicted_n includes it. Exclude that token to match the
|
|
# Transformers inter-token decode metric.
|
|
decode_tokens=max(0, int(timings.get("predicted_n", 0)) - 1),
|
|
# Use the runtime's prompt/first-token timing, matching the
|
|
# in-process Transformers boundary. HTTP/SSE and slot delay remain
|
|
# represented by total latency and queue_wait_ms.
|
|
ttft_ms=prefill_ms,
|
|
prefill_ms=prefill_ms,
|
|
decode_ms=decode_ms,
|
|
total_ms=total_ms,
|
|
# Whatever the wall clock saw but the runtime did not attribute to
|
|
# compute is time this request spent waiting for a slot.
|
|
queue_wait_ms=max(0.0, total_ms - prefill_ms - decode_ms),
|
|
)
|
|
|
|
def memory_probe(self) -> tuple[int, int]:
|
|
if self._process is None:
|
|
return 0, 0
|
|
return _process_rss(self._process.pid), 0
|
|
|
|
def close(self) -> None:
|
|
if self._process is not None:
|
|
if self._process.poll() is None:
|
|
self._process.terminate()
|
|
try:
|
|
self._process.wait(timeout=20)
|
|
except subprocess.TimeoutExpired:
|
|
self._process.kill()
|
|
self._process.wait(timeout=10)
|
|
self._process = None
|
|
if self._log is not None:
|
|
self._log.close()
|
|
self._log = None
|
|
|
|
|
|
def build_driver(spec: Mapping[str, Any], plan: BenchmarkPlan) -> RecipeDriverBundle:
|
|
"""Construct the driver named by a recipe's ``driver`` block."""
|
|
driver_spec = dict(spec["driver"])
|
|
kind = driver_spec.pop("type")
|
|
if kind == "transformers":
|
|
return TransformersDriver(**driver_spec)
|
|
if kind == "llama-cpp-server":
|
|
driver_spec.setdefault("n_parallel", max(plan.concurrency_levels))
|
|
return LlamaCppServerDriver(**driver_spec)
|
|
raise BenchmarkError(f"unknown driver type {kind!r}")
|
|
|
|
|
|
RecipeDriverBundle = Any # a RecipeDriver; named for readability at the call site
|
|
|
|
|
|
def _plan_from_config(config: Mapping[str, Any]) -> BenchmarkPlan:
|
|
raw = config["plan"]
|
|
return BenchmarkPlan(
|
|
plan_id=raw["plan_id"],
|
|
model_id=raw["model_id"],
|
|
model_revision=raw["model_revision"],
|
|
prompts=tuple(PromptSpec(**prompt) for prompt in raw["prompts"]),
|
|
sampling=SamplingPolicy(**raw.get("sampling", {})),
|
|
concurrency_levels=tuple(raw.get("concurrency_levels", (1, 4))),
|
|
repeats=int(raw.get("repeats", 1)),
|
|
warmup_requests=int(raw.get("warmup_requests", 1)),
|
|
)
|
|
|
|
|
|
def _recipe_from_config(spec: Mapping[str, Any]) -> RecipeSpec:
|
|
return RecipeSpec(
|
|
id=spec["id"],
|
|
runtime=spec["runtime"],
|
|
weight_format=spec["weight_format"],
|
|
weight_quantization=spec["weight_quantization"],
|
|
lane=Lane(spec["lane"]),
|
|
device=spec["device"],
|
|
artifact_path=spec.get("artifact_path", ""),
|
|
source_model_id=spec.get("source_model_id", ""),
|
|
source_model_revision=spec.get("source_model_revision", ""),
|
|
artifact_sha256=spec.get("artifact_sha256", ""),
|
|
is_reference=bool(spec.get("is_reference", False)),
|
|
notes=spec.get("notes", ""),
|
|
)
|
|
|
|
|
|
def run_configured_benchmark(config: Mapping[str, Any]) -> dict:
|
|
"""Run every recipe in ``config`` against one shared plan and return the report.
|
|
|
|
A recipe whose runtime cannot start is recorded as unavailable with the real
|
|
reason rather than dropped: a report that silently omits the recipe that
|
|
crashed would read as a clean result.
|
|
"""
|
|
require_real_inference()
|
|
_validate_config(config)
|
|
evidence_class = config.get("evidence_class", "local-real")
|
|
if evidence_class not in {"local-real", "multi-machine-real"}:
|
|
raise BenchmarkError("canonical real runner cannot emit synthetic evidence")
|
|
signing_key = _load_evidence_signing_key()
|
|
started_at = _utc_now()
|
|
run_id = str(uuid.uuid4())
|
|
config_sha256 = _canonical_sha256(config)
|
|
plan = _plan_from_config(config)
|
|
|
|
from .recipe_benchmark import RecipeMeasurement # local import keeps the seam obvious
|
|
|
|
measurements = []
|
|
for spec in config["recipes"]:
|
|
recipe = _recipe_from_config(spec)
|
|
driver = None
|
|
try:
|
|
driver = build_driver(spec, plan)
|
|
measurements.append(measure_recipe(driver, recipe, plan))
|
|
except Exception as exc:
|
|
measurements.append(RecipeMeasurement(
|
|
recipe=recipe,
|
|
load=LoadStats(artifact_bytes=0, load_ms=0.0),
|
|
unavailable_reason=f"{type(exc).__name__}: {exc}",
|
|
))
|
|
finally:
|
|
if driver is not None:
|
|
driver.close()
|
|
|
|
report = build_report(
|
|
plan,
|
|
measurements,
|
|
host={**dict(config.get("host", {})), **_host_manifest(config)},
|
|
evidence_class=evidence_class,
|
|
provenance={
|
|
"schema_version": PROVENANCE_SCHEMA_VERSION,
|
|
"producer": REAL_REPORT_PRODUCER,
|
|
"run_id": run_id,
|
|
"started_at": started_at,
|
|
"completed_at": _utc_now(),
|
|
"config_sha256": config_sha256,
|
|
"signature_algorithm": "ed25519",
|
|
},
|
|
)
|
|
_sign_report(report, signing_key)
|
|
return report
|