feat: DGR-001 - Lock the safetensors-versus-GGUF performance contract
This commit is contained in:
Dobromir Popov
473
packages/node/meshnet_node/recipe_drivers.py
Normal file
473
packages/node/meshnet_node/recipe_drivers.py
Normal file
@@ -0,0 +1,473 @@
|
||||
"""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 json
|
||||
import os
|
||||
import socket
|
||||
import subprocess
|
||||
import time
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from pathlib import Path
|
||||
from typing import Any, Mapping
|
||||
|
||||
from .recipe_benchmark import (
|
||||
BenchmarkError,
|
||||
BenchmarkPlan,
|
||||
GenerationSample,
|
||||
Lane,
|
||||
LoadStats,
|
||||
PromptSpec,
|
||||
RecipeSpec,
|
||||
SamplingPolicy,
|
||||
build_report,
|
||||
measure_recipe,
|
||||
)
|
||||
|
||||
REAL_INFERENCE_ENV = "MESHNET_ENABLE_REAL_INFERENCE_TESTS"
|
||||
|
||||
|
||||
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 _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())
|
||||
|
||||
|
||||
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
|
||||
|
||||
def load(self) -> LoadStats:
|
||||
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=_process_rss(),
|
||||
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 _process_rss(), 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,
|
||||
*,
|
||||
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.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 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}")
|
||||
|
||||
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 = subprocess.PIPE
|
||||
self._process = subprocess.Popen(
|
||||
command, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
|
||||
)
|
||||
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"llama-server; 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"
|
||||
)
|
||||
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"
|
||||
)
|
||||
|
||||
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()
|
||||
ttft_ms = 0.0
|
||||
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", "")
|
||||
if content and not ttft_ms:
|
||||
ttft_ms = (time.monotonic() - started) * 1000
|
||||
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)),
|
||||
decode_tokens=int(timings.get("predicted_n", 0)),
|
||||
ttft_ms=ttft_ms or total_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 None:
|
||||
return
|
||||
self._process.terminate()
|
||||
try:
|
||||
self._process.wait(timeout=20)
|
||||
except subprocess.TimeoutExpired:
|
||||
self._process.kill()
|
||||
self._process.wait(timeout=10)
|
||||
self._process = 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", ""),
|
||||
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()
|
||||
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)
|
||||
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}",
|
||||
))
|
||||
|
||||
return build_report(
|
||||
plan,
|
||||
measurements,
|
||||
host=dict(config.get("host", {})),
|
||||
evidence_class=config.get("evidence_class", "local-real"),
|
||||
)
|
||||
Reference in New Issue
Block a user