474 lines
17 KiB
Python
474 lines
17 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 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"),
|
|
)
|