1"""Benchmark engine: load time, prompt tps, generation tps, TTFT, peak/avg memory, CPU/GPU, thermal."""23from __future__ import annotations45import asyncio6import json7import statistics8import time910from .hardware import GB, sample_telemetry_async11from .jobs import Job1213BENCH_PROMPT_SHORT = "Explain in three sentences why unified memory matters for running large language models on Apple Silicon."14BENCH_PROMPT_LONG = ("You are a careful technical writer. Summarize the following passage, then list five key facts.\n\n" +15 ("Apple Silicon systems share a single pool of unified memory between the CPU and the GPU. "16 "This lets a model's weights be read by the GPU without copies, but it also means the operating system, "17 "the application and the model compete for the same bytes. ") * 40)181920async def run_benchmark(state, model_id: str, job: Job, params: dict | None = None) -> dict:21 manager = state.manager22 db = state.db23 jobs = state.jobs24 params = params or {}25 max_tokens = int(params.get("max_tokens", 256))26 runs = int(params.get("runs", 2))27 long_prompt = bool(params.get("long_prompt", True))28 m = await manager.registry.get(model_id)29 if not m:30 raise ValueError("model not found")3132 # 1. cold load timing (unload first if loaded)33 if manager.loaded.get(model_id):34 jobs.update(job, progress=0.02, stage="unloading for cold start")35 await manager.unload(model_id, reason="benchmark cold start")36 jobs.update(job, progress=0.05, stage="loading")37 t0 = time.time()38 lm = await manager.ensure_loaded(model_id, reason="benchmark")39 load_ms = (time.time() - t0) * 100040 jobs.update(job, progress=0.25, stage="loaded", load_ms=round(load_ms))4142 mem_samples: list[float] = []43 cpu_samples: list[float] = []44 gpu_samples: list[float] = []45 stop = asyncio.Event()4647 async def sampler():48 while not stop.is_set():49 tel = await sample_telemetry_async(state.settings.models_dir)50 mem_samples.append(lm.handle.memory_bytes() / GB)51 cpu_samples.append(tel.cpu_percent)52 if tel.gpu_percent is not None:53 gpu_samples.append(tel.gpu_percent)54 await asyncio.sleep(0.5)5556 samp = asyncio.create_task(sampler())57 results = []58 try:59 prompts = [BENCH_PROMPT_SHORT] + ([BENCH_PROMPT_LONG] if long_prompt else [])60 total_steps = runs * len(prompts)61 step = 062 async with manager.use(lm):63 for r in range(runs):64 for p in prompts:65 if job.cancelled:66 break67 step += 168 jobs.update(job, progress=0.25 + 0.7 * step / total_steps, stage=f"run {step}/{total_steps}")69 if m["embedding"]:70 t1 = time.time()71 resp = await manager.client.post(f"{lm.handle.base_url}/v1/embeddings",72 json={"model": model_id, "input": [p] * 8}, timeout=600)73 resp.raise_for_status()74 d = resp.json()75 dt = time.time() - t176 ptoks = d.get("usage", {}).get("prompt_tokens", 0)77 results.append({"prompt_tokens": ptoks, "prompt_tps": round(ptoks / dt, 1) if dt else None,78 "generation_tokens": 0, "generation_tps": None, "ttft_ms": round(dt * 1000, 1)})79 continue80 if m["reranker"]:81 t1 = time.time()82 resp = await manager.client.post(f"{lm.handle.base_url}/v1/rerank",83 json={"model": model_id, "query": "unified memory", "documents": [p] * 8}, timeout=600)84 resp.raise_for_status()85 dt = time.time() - t186 results.append({"prompt_tokens": None, "prompt_tps": None, "generation_tokens": 0,87 "generation_tps": None, "ttft_ms": round(dt * 1000, 1)})88 continue89 body = {"model": model_id, "messages": [{"role": "user", "content": p}], "max_tokens": max_tokens,90 "temperature": 0.0, "stream": True, "stream_options": {"include_usage": True}}91 t1 = time.time()92 first = None93 n = 094 usage = {}95 timings = {}96 async with manager.client.stream("POST", f"{lm.handle.base_url}/v1/chat/completions", json=body, timeout=1800) as resp:97 async for line in resp.aiter_lines():98 if not line.startswith("data: ") or line.strip() == "data: [DONE]":99 continue100 try:101 obj = json.loads(line[6:])102 except json.JSONDecodeError:103 continue104 for ch in obj.get("choices") or []:105 if (ch.get("delta") or {}).get("content"):106 if first is None:107 first = time.time()108 n += 1109 if obj.get("usage"):110 usage = obj["usage"]111 if obj.get("timings"):112 timings = obj["timings"]113 t2 = time.time()114 gen_tokens = usage.get("completion_tokens") or timings.get("predicted_n") or n115 ttft = (first - t1) * 1000 if first else (t2 - t1) * 1000116 gen_s = (t2 - (first or t1))117 results.append({118 "prompt_tokens": usage.get("prompt_tokens") or timings.get("prompt_n"),119 "prompt_tps": timings.get("prompt_tps") or timings.get("prompt_per_second") or (120 round((usage.get("prompt_tokens") or 0) / (ttft / 1000), 1) if ttft else None),121 "generation_tokens": gen_tokens,122 "generation_tps": timings.get("generation_tps") or timings.get("predicted_per_second") or (123 round(gen_tokens / gen_s, 2) if gen_s > 0 else None),124 "ttft_ms": round(ttft, 1), "total_ms": round((t2 - t1) * 1000, 1),125 "peak_memory_gb": timings.get("peak_memory_gb"),126 })127 finally:128 stop.set()129 samp.cancel()130131 def avg(key):132 vals = [r[key] for r in results if r.get(key) is not None]133 return round(statistics.fmean(vals), 2) if vals else None134135 tel = await sample_telemetry_async(state.settings.models_dir)136 row = {137 "model_id": model_id, "created_at": time.time(), "load_ms": round(load_ms), "prompt_tokens": avg("prompt_tokens"),138 "prompt_tps": avg("prompt_tps"), "generation_tokens": avg("generation_tokens"), "generation_tps": avg("generation_tps"),139 "ttft_ms": avg("ttft_ms"), "peak_memory_gb": round(max(mem_samples), 2) if mem_samples else None,140 "avg_memory_gb": round(statistics.fmean(mem_samples), 2) if mem_samples else None,141 "cpu_percent": round(statistics.fmean(cpu_samples), 1) if cpu_samples else None,142 "gpu_percent": round(statistics.fmean(gpu_samples), 1) if gpu_samples else None,143 "thermal_state": tel.thermal_state, "context": lm.handle.context, "runtime": m["runtime"],144 "params": json.dumps({"max_tokens": max_tokens, "runs": runs, "long_prompt": long_prompt}),145 "notes": json.dumps({"runs": results}),146 }147 cols = ", ".join(row.keys())148 qs = ", ".join("?" for _ in row)149 bid = await db.execute(f"INSERT INTO model_benchmarks({cols}) VALUES({qs})", list(row.values()))150 if row["generation_tps"]:151 await manager.registry.update(model_id, avg_tps=row["generation_tps"])152 if row["ttft_ms"]:153 await manager.registry.update(model_id, first_token_latency_ms=row["ttft_ms"])154 await db.model_event(model_id, "benchmark", {k: row[k] for k in ("load_ms", "generation_tps", "prompt_tps", "ttft_ms", "peak_memory_gb")})155 row["id"] = bid156 row["runs"] = results157 row.pop("notes", None)158 return row159