"""Benchmark engine: load time, prompt tps, generation tps, TTFT, peak/avg memory, CPU/GPU, thermal.""" from __future__ import annotations import asyncio import json import statistics import time from .hardware import GB, sample_telemetry_async from .jobs import Job BENCH_PROMPT_SHORT = "Explain in three sentences why unified memory matters for running large language models on Apple Silicon." BENCH_PROMPT_LONG = ("You are a careful technical writer. Summarize the following passage, then list five key facts.\n\n" + ("Apple Silicon systems share a single pool of unified memory between the CPU and the GPU. " "This lets a model's weights be read by the GPU without copies, but it also means the operating system, " "the application and the model compete for the same bytes. ") * 40) async def run_benchmark(state, model_id: str, job: Job, params: dict | None = None) -> dict: manager = state.manager db = state.db jobs = state.jobs params = params or {} max_tokens = int(params.get("max_tokens", 256)) runs = int(params.get("runs", 2)) long_prompt = bool(params.get("long_prompt", True)) m = await manager.registry.get(model_id) if not m: raise ValueError("model not found") # 1. cold load timing (unload first if loaded) if manager.loaded.get(model_id): jobs.update(job, progress=0.02, stage="unloading for cold start") await manager.unload(model_id, reason="benchmark cold start") jobs.update(job, progress=0.05, stage="loading") t0 = time.time() lm = await manager.ensure_loaded(model_id, reason="benchmark") load_ms = (time.time() - t0) * 1000 jobs.update(job, progress=0.25, stage="loaded", load_ms=round(load_ms)) mem_samples: list[float] = [] cpu_samples: list[float] = [] gpu_samples: list[float] = [] stop = asyncio.Event() async def sampler(): while not stop.is_set(): tel = await sample_telemetry_async(state.settings.models_dir) mem_samples.append(lm.handle.memory_bytes() / GB) cpu_samples.append(tel.cpu_percent) if tel.gpu_percent is not None: gpu_samples.append(tel.gpu_percent) await asyncio.sleep(0.5) samp = asyncio.create_task(sampler()) results = [] try: prompts = [BENCH_PROMPT_SHORT] + ([BENCH_PROMPT_LONG] if long_prompt else []) total_steps = runs * len(prompts) step = 0 async with manager.use(lm): for r in range(runs): for p in prompts: if job.cancelled: break step += 1 jobs.update(job, progress=0.25 + 0.7 * step / total_steps, stage=f"run {step}/{total_steps}") if m["embedding"]: t1 = time.time() resp = await manager.client.post(f"{lm.handle.base_url}/v1/embeddings", json={"model": model_id, "input": [p] * 8}, timeout=600) resp.raise_for_status() d = resp.json() dt = time.time() - t1 ptoks = d.get("usage", {}).get("prompt_tokens", 0) results.append({"prompt_tokens": ptoks, "prompt_tps": round(ptoks / dt, 1) if dt else None, "generation_tokens": 0, "generation_tps": None, "ttft_ms": round(dt * 1000, 1)}) continue if m["reranker"]: t1 = time.time() resp = await manager.client.post(f"{lm.handle.base_url}/v1/rerank", json={"model": model_id, "query": "unified memory", "documents": [p] * 8}, timeout=600) resp.raise_for_status() dt = time.time() - t1 results.append({"prompt_tokens": None, "prompt_tps": None, "generation_tokens": 0, "generation_tps": None, "ttft_ms": round(dt * 1000, 1)}) continue body = {"model": model_id, "messages": [{"role": "user", "content": p}], "max_tokens": max_tokens, "temperature": 0.0, "stream": True, "stream_options": {"include_usage": True}} t1 = time.time() first = None n = 0 usage = {} timings = {} async with manager.client.stream("POST", f"{lm.handle.base_url}/v1/chat/completions", json=body, timeout=1800) as resp: async for line in resp.aiter_lines(): if not line.startswith("data: ") or line.strip() == "data: [DONE]": continue try: obj = json.loads(line[6:]) except json.JSONDecodeError: continue for ch in obj.get("choices") or []: if (ch.get("delta") or {}).get("content"): if first is None: first = time.time() n += 1 if obj.get("usage"): usage = obj["usage"] if obj.get("timings"): timings = obj["timings"] t2 = time.time() gen_tokens = usage.get("completion_tokens") or timings.get("predicted_n") or n ttft = (first - t1) * 1000 if first else (t2 - t1) * 1000 gen_s = (t2 - (first or t1)) results.append({ "prompt_tokens": usage.get("prompt_tokens") or timings.get("prompt_n"), "prompt_tps": timings.get("prompt_tps") or timings.get("prompt_per_second") or ( round((usage.get("prompt_tokens") or 0) / (ttft / 1000), 1) if ttft else None), "generation_tokens": gen_tokens, "generation_tps": timings.get("generation_tps") or timings.get("predicted_per_second") or ( round(gen_tokens / gen_s, 2) if gen_s > 0 else None), "ttft_ms": round(ttft, 1), "total_ms": round((t2 - t1) * 1000, 1), "peak_memory_gb": timings.get("peak_memory_gb"), }) finally: stop.set() samp.cancel() def avg(key): vals = [r[key] for r in results if r.get(key) is not None] return round(statistics.fmean(vals), 2) if vals else None tel = await sample_telemetry_async(state.settings.models_dir) row = { "model_id": model_id, "created_at": time.time(), "load_ms": round(load_ms), "prompt_tokens": avg("prompt_tokens"), "prompt_tps": avg("prompt_tps"), "generation_tokens": avg("generation_tokens"), "generation_tps": avg("generation_tps"), "ttft_ms": avg("ttft_ms"), "peak_memory_gb": round(max(mem_samples), 2) if mem_samples else None, "avg_memory_gb": round(statistics.fmean(mem_samples), 2) if mem_samples else None, "cpu_percent": round(statistics.fmean(cpu_samples), 1) if cpu_samples else None, "gpu_percent": round(statistics.fmean(gpu_samples), 1) if gpu_samples else None, "thermal_state": tel.thermal_state, "context": lm.handle.context, "runtime": m["runtime"], "params": json.dumps({"max_tokens": max_tokens, "runs": runs, "long_prompt": long_prompt}), "notes": json.dumps({"runs": results}), } cols = ", ".join(row.keys()) qs = ", ".join("?" for _ in row) bid = await db.execute(f"INSERT INTO model_benchmarks({cols}) VALUES({qs})", list(row.values())) if row["generation_tps"]: await manager.registry.update(model_id, avg_tps=row["generation_tps"]) if row["ttft_ms"]: await manager.registry.update(model_id, first_token_latency_ms=row["ttft_ms"]) await db.model_event(model_id, "benchmark", {k: row[k] for k in ("load_ms", "generation_tps", "prompt_tps", "ttft_ms", "peak_memory_gb")}) row["id"] = bid row["runs"] = results row.pop("notes", None) return row