mrleo1nidandClaude Sonnet 5.5
f49ee90477
Includes earlier uncommitted work: Tomie and Unity adapters, benchmark, structured output, call metrics. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
492 lines
27 KiB
Python
492 lines
27 KiB
Python
"""Repeatable end-to-end translation benchmark against local LM Studio."""
|
||
import argparse
|
||
from contextlib import contextmanager
|
||
from datetime import datetime, timezone
|
||
import csv
|
||
import hashlib
|
||
import html
|
||
import json
|
||
import os
|
||
from pathlib import Path
|
||
import platform
|
||
import shutil
|
||
import statistics
|
||
import subprocess
|
||
import threading
|
||
import time
|
||
from urllib.parse import urlsplit
|
||
from uuid import uuid4
|
||
|
||
import httpx
|
||
from pydantic import Field
|
||
import uvicorn
|
||
|
||
from .api import create_app
|
||
from .config import ProviderConfig, Settings
|
||
from .models import Project, StrictModel, TextItem
|
||
from .storage import Storage
|
||
from .provider import Provider
|
||
|
||
|
||
class ModelSpec(StrictModel):
|
||
model: str = Field(min_length=1)
|
||
label: str | None = None
|
||
prompt_style: str | None = None
|
||
temperature: float | None = None
|
||
top_p: float | None = None
|
||
top_k: int | None = None
|
||
repetition_penalty: float | None = None
|
||
structured_output: bool = True
|
||
|
||
|
||
class BenchmarkConfig(StrictModel):
|
||
game: str | None = None
|
||
input: str | None = None
|
||
project_file: str = "examples/tomie_project.json"
|
||
models: list[ModelSpec] = Field(default_factory=list)
|
||
batch_sizes: list[int] = Field(default=[1, 5, 10], min_length=1, max_length=10)
|
||
repeats: int = Field(default=3, ge=1, le=20)
|
||
limit: int = Field(default=30, ge=1, le=500)
|
||
warmup_lines: int = Field(default=2, ge=1, le=20)
|
||
context_length: int = Field(default=8192, ge=512, le=262144)
|
||
gpu: float = Field(default=1, ge=0, le=1)
|
||
reuse_loaded: bool = False
|
||
port: int = Field(default=8788, ge=1024, le=65535)
|
||
job_timeout: float = Field(default=600, gt=0, le=7200)
|
||
retries: int = Field(default=2, ge=0, le=5)
|
||
provider_url: str = "http://127.0.0.1:1234/v1"
|
||
api_key_env: str = "LM_STUDIO_API_KEY"
|
||
|
||
|
||
def write_json(path, value):
|
||
path = Path(path)
|
||
path.parent.mkdir(parents=True, exist_ok=True)
|
||
temporary = path.with_suffix(path.suffix + ".tmp")
|
||
temporary.write_text(json.dumps(value, ensure_ascii=False, indent=2), encoding="utf-8")
|
||
temporary.replace(path)
|
||
|
||
|
||
class LMS:
|
||
def __init__(self, executable=None):
|
||
self.executable = executable or shutil.which("lms")
|
||
if not self.executable:
|
||
raise RuntimeError("lms CLI not found in PATH. Install the LM Studio CLI first.")
|
||
|
||
def run(self, *arguments):
|
||
result = subprocess.run([self.executable, *arguments], capture_output=True, text=True,
|
||
encoding="utf-8", errors="replace", stdin=subprocess.DEVNULL, timeout=180,
|
||
creationflags=0x08000000 if os.name == "nt" else 0)
|
||
if result.returncode:
|
||
raise RuntimeError(f"lms {' '.join(arguments[:2])} failed: {(result.stderr or result.stdout)[-1500:]}")
|
||
return result.stdout
|
||
|
||
def json(self, command):
|
||
value = json.loads(self.run(command, "--json"))
|
||
if not isinstance(value, list):
|
||
raise RuntimeError(f"Unexpected lms {command} --json response")
|
||
return value
|
||
|
||
|
||
class ModelManager:
|
||
"""Switch only idle LLMs and restore the original model IDs on exit."""
|
||
def __init__(self, lms, config):
|
||
self.lms, self.config = lms, config
|
||
self.originals = self.loaded()
|
||
self.owned = set()
|
||
self.restore_errors = []
|
||
self.assert_idle(self.originals)
|
||
|
||
def loaded(self):
|
||
return [item for item in self.lms.json("ps") if item.get("type") == "llm"]
|
||
|
||
@staticmethod
|
||
def assert_idle(models):
|
||
if any(item.get("status") != "idle" or item.get("queued", 0) for item in models):
|
||
raise RuntimeError("An LM Studio model is busy. Stop other generation jobs before benchmarking.")
|
||
|
||
def activate(self, metadata):
|
||
loaded = self.loaded()
|
||
self.assert_idle(loaded)
|
||
selected = next((item for item in loaded if self.config.reuse_loaded and item.get("path") == metadata.get("path")
|
||
and item.get("contextLength") == self.config.context_length), None)
|
||
for item in loaded:
|
||
if selected and item["identifier"] == selected["identifier"]:
|
||
continue
|
||
print(f"Temporarily unloading idle model: {item['identifier']}", flush=True)
|
||
self.lms.run("unload", item["identifier"])
|
||
if selected:
|
||
return selected, 0.0
|
||
identifier = "lt-bench-" + uuid4().hex[:10]
|
||
# Register before loading so a partial load is also cleaned up on error.
|
||
self.owned.add(identifier)
|
||
started = time.perf_counter()
|
||
key = metadata["modelKey"]
|
||
self.lms.run("load", key, "--identifier", identifier, "--context-length", str(self.config.context_length), "--gpu", str(self.config.gpu), "--yes")
|
||
elapsed = time.perf_counter() - started
|
||
selected = next((item for item in self.loaded() if item["identifier"] == identifier), None)
|
||
if not selected:
|
||
raise RuntimeError("Model was not registered after lms load")
|
||
if selected.get("path") != metadata.get("path"):
|
||
raise RuntimeError("LM Studio loaded a different model file; refusing a mislabeled benchmark")
|
||
return selected, elapsed
|
||
|
||
def restore(self):
|
||
try:
|
||
current = self.loaded()
|
||
self.assert_idle(current)
|
||
for item in current:
|
||
if item["identifier"] in self.owned:
|
||
self.lms.run("unload", item["identifier"])
|
||
identifiers = {item["identifier"] for item in self.loaded()}
|
||
for item in self.originals:
|
||
if item["identifier"] not in identifiers:
|
||
key = item["modelKey"]
|
||
self.lms.run("load", key, "--identifier", item["identifier"], "--context-length", str(item["contextLength"]), "--yes")
|
||
except Exception as exc:
|
||
self.restore_errors.append(str(exc))
|
||
|
||
|
||
class CapturingProvider(Provider):
|
||
def __init__(self, config, directory):
|
||
super().__init__(config)
|
||
self.directory = directory
|
||
|
||
async def translate(self, *args, **kwargs):
|
||
try:
|
||
return await super().translate(*args, **kwargs)
|
||
finally:
|
||
if self.last_response is not None:
|
||
write_json(self.directory / (uuid4().hex + ".json"), self.last_response)
|
||
|
||
|
||
@contextmanager
|
||
def isolated_api(settings, provider_factory=Provider):
|
||
# A dedicated port and DB avoid old queues, cache and project mutations.
|
||
import socket
|
||
listener = socket.socket()
|
||
try:
|
||
listener.bind(("127.0.0.1", settings.port))
|
||
except OSError as exc:
|
||
raise RuntimeError(f"Benchmark port {settings.port} is already in use") from exc
|
||
finally:
|
||
listener.close()
|
||
server = uvicorn.Server(uvicorn.Config(create_app(settings, provider_factory), host="127.0.0.1", port=settings.port,
|
||
log_level="warning", access_log=False))
|
||
thread = threading.Thread(target=server.run, daemon=True)
|
||
thread.start()
|
||
try:
|
||
deadline = time.monotonic() + 20
|
||
while not server.started:
|
||
if not thread.is_alive() or time.monotonic() >= deadline:
|
||
raise RuntimeError("Benchmark API failed to start")
|
||
time.sleep(0.05)
|
||
with httpx.Client(base_url=f"http://127.0.0.1:{settings.port}", timeout=30, trust_env=False) as client:
|
||
yield client
|
||
finally:
|
||
server.should_exit = True
|
||
thread.join(timeout=30)
|
||
if thread.is_alive():
|
||
raise RuntimeError("Benchmark API did not stop; do not start another benchmark on this port")
|
||
|
||
|
||
def request(client, method, path, payload=None):
|
||
response = client.request(method, path, json=payload)
|
||
if not response.is_success:
|
||
raise RuntimeError(f"Translation API HTTP {response.status_code}: {response.text[:1000]}")
|
||
return response.json()
|
||
|
||
|
||
def timed_job(client, project_id, items, batch_size, config, on_created=None):
|
||
started = time.perf_counter()
|
||
job = request(client, "POST", f"/projects/{project_id}/jobs", {
|
||
"items": items, "batch_size": batch_size, "use_cache": False, "retries": config.retries,
|
||
})
|
||
if on_created:
|
||
on_created(job["id"])
|
||
deadline = time.monotonic() + config.job_timeout
|
||
while job["status"] in {"queued", "running"}:
|
||
if time.monotonic() >= deadline:
|
||
# Never continue with another model while this job is still running.
|
||
raise TimeoutError(f"Job {job['id']} timed out; aborting benchmark. Results remain in benchmark DB.")
|
||
time.sleep(0.05)
|
||
job = request(client, "GET", "/jobs/" + job["id"])
|
||
return job, time.perf_counter() - started
|
||
|
||
|
||
def summarize_run(job, elapsed):
|
||
rows = job["items"]
|
||
successes = [row for row in rows if row["status"] == "completed"]
|
||
inferred = [row for row in successes if row["attempts"] > 0 and not row["cached"]]
|
||
metrics = job.get("metrics", {})
|
||
provider_seconds = metrics.get("provider_seconds", 0)
|
||
complete_usage = metrics.get("provider_calls", 0) > 0 and metrics.get("calls_with_completion_usage") == metrics.get("provider_calls")
|
||
tokens = metrics.get("completion_tokens") if complete_usage else None
|
||
return {
|
||
"job_id": job["id"], "status": job["status"], "seconds": elapsed,
|
||
"completed": len(successes), "failed": len(rows) - len(successes),
|
||
"model_lines": len(inferred), "cached_lines": sum(bool(row["cached"]) for row in rows),
|
||
"local_lines": sum(row["attempts"] == 0 for row in successes),
|
||
"retry_attempts": sum(max(0, row["attempts"] - 1) for row in rows),
|
||
"source_chars": sum(len(row["source"]) for row in inferred),
|
||
"lines_per_second": len(inferred) / elapsed if elapsed > 0 else 0,
|
||
"chars_per_second": sum(len(row["source"]) for row in inferred) / elapsed if elapsed > 0 else 0,
|
||
"unchanged_ids": [row["id"] for row in inferred if row["source"].strip() == row["translation"].strip()],
|
||
"completion_tokens_per_request_second": tokens / provider_seconds if tokens is not None and provider_seconds > 0 else None,
|
||
"error": job.get("error"), **metrics,
|
||
}
|
||
|
||
|
||
def aggregate(runs):
|
||
result = []
|
||
keys = dict.fromkeys((run["model"], run["batch_size"]) for run in runs)
|
||
for model, batch in keys:
|
||
group = [run for run in runs if run["model"] == model and run["batch_size"] == batch]
|
||
successful = [run for run in group if run["status"] == "completed" and run["failed"] == 0]
|
||
result.append({
|
||
"model": model, "batch_size": batch, "runs": len(group), "successful_runs": len(successful),
|
||
"median_seconds": statistics.median(run["seconds"] for run in successful) if successful else None,
|
||
"median_lines_per_second": statistics.median(run["lines_per_second"] for run in successful) if successful else None,
|
||
"median_chars_per_second": statistics.median(run["chars_per_second"] for run in successful) if successful else None,
|
||
"min_seconds": min((run["seconds"] for run in successful), default=None),
|
||
"max_seconds": max((run["seconds"] for run in successful), default=None),
|
||
"failed_items": sum(run["failed"] for run in group),
|
||
"retry_attempts": sum(run["retry_attempts"] for run in group),
|
||
"unchanged_items": sum(len(run["unchanged_ids"]) for run in group),
|
||
})
|
||
return result
|
||
|
||
|
||
def save_report(directory, report):
|
||
report["summary"] = aggregate(report["runs"])
|
||
write_json(directory / "report.json", report)
|
||
fields = ["model", "batch_size", "repeat", "status", "seconds", "completed", "failed", "model_lines", "cached_lines",
|
||
"local_lines", "retry_attempts", "lines_per_second", "chars_per_second", "provider_calls", "provider_seconds",
|
||
"prompt_tokens", "completion_tokens", "completion_tokens_per_request_second", "job_id"]
|
||
with (directory / "runs.csv").open("w", encoding="utf-8-sig", newline="") as stream:
|
||
writer = csv.DictWriter(stream, fieldnames=fields, extrasaction="ignore")
|
||
writer.writeheader()
|
||
writer.writerows(report["runs"])
|
||
escape = lambda value: html.escape(str(value))
|
||
number = lambda value: f"{value:.2f}" if value is not None else "—"
|
||
rows = "".join(f"<tr><td>{escape(row['model'])}</td><td>{row['batch_size']}</td>"
|
||
f"<td>{row['successful_runs']}/{row['runs']}</td><td>{number(row['median_seconds'])}</td>"
|
||
f"<td>{number(row['median_lines_per_second'])}</td><td>{number(row['median_chars_per_second'])}</td>"
|
||
f"<td>{row['retry_attempts']}</td><td>{row['unchanged_items']}</td></tr>" for row in report["summary"])
|
||
errors = "".join(f"<li>{escape(error)}</li>" for error in report["errors"] + report.get("restore_errors", []))
|
||
source = escape(report["dataset"]["source"])
|
||
document = f"""<!doctype html><html lang="ru"><meta charset="utf-8"><title>Translation benchmark</title>
|
||
<style>body{{font:16px system-ui;max-width:1200px;margin:40px auto;padding:0 24px;color:#20242a;background:#fafafa}}
|
||
table{{border-collapse:collapse;width:100%;background:white}}td,th{{padding:12px;text-align:left;border-bottom:1px solid #ddd}}
|
||
h1{{font-size:26px}}.note{{color:#555;line-height:1.6}}a{{color:#245b8c}}code{{overflow-wrap:anywhere}}</style>
|
||
<h1>Скорость перевода игры через Local Translator</h1>
|
||
<p class="note">Источник: {source} · {report['dataset']['items']} строк · {report['dataset']['source_chars']} символов · {escape(report['started_at'])}</p>
|
||
<p>Показаны медианы успешных повторов. Загрузка модели и прогрев исключены; HTTP, очередь, валидация, запись SQLite и повторы включены.
|
||
Кэш отключён. Полные совпадения со словарём удалены из тестового словаря. Неизменённые строки требуют ручной проверки и не считаются автоматически ошибками.</p>
|
||
<div style="overflow:auto"><table><thead><tr><th>Модель</th><th>Пачка</th><th>Успешные прогоны</th><th>Медиана, с</th><th>Строк/с</th><th>Символов/с</th><th>Повторы строк</th><th>Без изменений</th></tr></thead><tbody>{rows}</tbody></table></div>
|
||
<p class="note">Это проверка скорости и структуры, а не оценка литературного качества или отсутствия цензуры.
|
||
Токены/с в JSON/CSV делятся на полное время HTTP-запросов модели, включая обработку промпта; это не чистая скорость декодирования.
|
||
Нативный prompt cache LM Studio не сбрасывается. GPU, другие процессы и порядок прогонов влияют на результат.</p>
|
||
<ul>{errors}</ul><p><a href="report.json">Полный JSON</a> · <a href="runs.csv">CSV</a> · Переводы каждого прогона находятся в папке jobs.</p></html>"""
|
||
(directory / "report.html").write_text(document, encoding="utf-8")
|
||
|
||
|
||
def profile(spec, config):
|
||
hy = "hy-mt" in spec.model.lower() or "hy_mt" in spec.model.lower()
|
||
values = {"base_url": config.provider_url, "api_key_env": config.api_key_env, "timeout": min(config.job_timeout, 180),
|
||
"model": spec.model, "prompt_style": "hy_mt" if hy else "generic", "temperature": 0.7 if hy else 0.2,
|
||
"structured_output": spec.structured_output}
|
||
if hy:
|
||
values.update(top_p=0.6, top_k=20, repetition_penalty=1.05)
|
||
values.update({key: value for key, value in spec.model_dump().items() if key not in {"model", "label"} and value is not None})
|
||
return ProviderConfig.model_validate(values)
|
||
|
||
|
||
def resolve_metadata(spec, downloaded):
|
||
matches = [item for item in downloaded if spec.model in {item.get("modelKey"), item.get("path"), item.get("indexedModelIdentifier")}]
|
||
if len(matches) != 1:
|
||
raise ValueError(f"Expected one downloaded model for {spec.model}, found {len(matches)}. Use its exact indexedModelIdentifier from lms ls --json.")
|
||
return matches[0]
|
||
|
||
|
||
def load_dataset(config, root, directory):
|
||
if config.input:
|
||
source = (root / config.input).resolve()
|
||
elif config.game:
|
||
# The optional extractor runs independently and never writes into the game.
|
||
source = directory / "input.json"
|
||
import sys
|
||
result = subprocess.run([sys.executable, "-X", "utf8", str(Path(__file__).resolve().parents[1] / "examples/extract_tomie_ui.py"),
|
||
str((root / config.game).resolve()), "--output", str(source), "--limit", str(config.limit)],
|
||
capture_output=True, text=True, encoding="utf-8", errors="replace", timeout=120,
|
||
creationflags=0x08000000 if os.name == "nt" else 0)
|
||
if result.returncode:
|
||
raise RuntimeError("Game extraction failed. Install optional dependency with pip install -e .[unity].\n" + result.stderr[-2000:])
|
||
else:
|
||
raise ValueError("Set input or game in benchmark configuration")
|
||
raw = json.loads(source.read_text(encoding="utf-8-sig"))
|
||
if not isinstance(raw, list):
|
||
raise ValueError("Dataset must be an array of {id, text} items")
|
||
items = [TextItem.model_validate(item).model_dump() for item in raw][:config.limit]
|
||
items = [item for item in items if item["text"].strip()]
|
||
if not items or len({item["id"] for item in items}) != len(items):
|
||
raise ValueError("Dataset must contain nonempty text and unique IDs")
|
||
if sum(len(item["text"]) for item in items) > 1000000:
|
||
raise ValueError("Dataset exceeds API job limit of 1,000,000 characters")
|
||
write_json(directory / "dataset.json", items)
|
||
return items, {"source": str(source), "items": len(items), "source_chars": sum(len(item["text"]) for item in items),
|
||
"sha256": hashlib.sha256(json.dumps(items, ensure_ascii=False, sort_keys=True).encode()).hexdigest()}
|
||
|
||
|
||
def run(config, root, directory):
|
||
directory.mkdir(parents=True, exist_ok=False)
|
||
items, dataset = load_dataset(config, root, directory)
|
||
base_project = Project.model_validate(json.loads((root / config.project_file).read_text(encoding="utf-8-sig")))
|
||
source_strings = {item["text"] for item in items}
|
||
removed_keys = [key for key in base_project.glossary if key in source_strings]
|
||
base_project.glossary = {key: value for key, value in base_project.glossary.items() if key not in source_strings}
|
||
url = urlsplit(config.provider_url)
|
||
if url.scheme != "http" or url.hostname not in {"localhost", "127.0.0.1"} or url.username or url.password or url.query or url.fragment:
|
||
raise ValueError("Benchmark model switching requires a local HTTP LM Studio endpoint")
|
||
lms = LMS()
|
||
downloaded = [item for item in lms.json("ls") if item.get("type") == "llm"]
|
||
specs = config.models or [ModelSpec(model=item.get("indexedModelIdentifier") or item["modelKey"]) for item in downloaded]
|
||
if not specs:
|
||
raise RuntimeError("No local language models found. Download a model in LM Studio first.")
|
||
if len({spec.model for spec in specs}) != len(specs):
|
||
raise ValueError("Model list contains duplicates")
|
||
for batch_size in config.batch_sizes:
|
||
if not 1 <= batch_size <= 50:
|
||
raise ValueError("Batch sizes must be between 1 and 50")
|
||
if len(set(config.batch_sizes)) != len(config.batch_sizes):
|
||
raise ValueError("Batch sizes must be unique")
|
||
# Validate every profile before changing loaded models.
|
||
for spec in specs:
|
||
profile(spec, config)
|
||
source_root = Path(__file__).parent
|
||
implementation_hash = hashlib.sha256(b"".join(file.read_bytes() for file in sorted(source_root.glob("*.py")))).hexdigest()
|
||
gpu_info = None
|
||
if shutil.which("nvidia-smi"):
|
||
try:
|
||
gpu_info = subprocess.run(["nvidia-smi", "--query-gpu=name,memory.total,driver_version", "--format=csv,noheader"],
|
||
capture_output=True, text=True, timeout=5, creationflags=0x08000000 if os.name == "nt" else 0).stdout.strip() or None
|
||
except (OSError, subprocess.SubprocessError):
|
||
pass
|
||
report = {"started_at": datetime.now(timezone.utc).isoformat(), "state": "running", "config": config.model_dump(),
|
||
"implementation_sha256": implementation_hash,
|
||
"dataset": dataset, "removed_exact_glossary_keys": removed_keys, "host": {"platform": platform.platform(), "python": platform.python_version(), "nvidia_gpu_info": gpu_info},
|
||
"models": [], "runs": [], "errors": [], "restore_errors": []}
|
||
save_report(directory, report)
|
||
headers = {"Authorization": "Bearer " + os.environ[config.api_key_env]} if os.environ.get(config.api_key_env) else {}
|
||
with httpx.Client(timeout=3, trust_env=False, headers=headers) as probe:
|
||
try:
|
||
response = probe.get(config.provider_url.rstrip("/") + "/models")
|
||
response.raise_for_status()
|
||
except httpx.HTTPError:
|
||
lms.run("daemon", "up")
|
||
lms.run("server", "start", "--port", str(url.port or 1234), "--bind", "127.0.0.1")
|
||
manager = ModelManager(lms, config)
|
||
write_json(directory / "original-models.json", manager.originals)
|
||
try:
|
||
for index, spec in enumerate(specs):
|
||
try:
|
||
metadata = resolve_metadata(spec, downloaded)
|
||
except ValueError as exc:
|
||
report["errors"].append(str(exc))
|
||
save_report(directory, report)
|
||
continue
|
||
label = spec.label or f"{metadata.get('displayName', spec.model)} {metadata.get('quantization', {}).get('name', '')}".strip()
|
||
if any(model["label"] == label for model in report["models"]):
|
||
label += " (" + spec.model + ")"
|
||
model_record = {"label": label, "key": spec.model, "metadata": metadata}
|
||
report["models"].append(model_record)
|
||
print(f"\nModel {index + 1}/{len(specs)}: {label}", flush=True)
|
||
try:
|
||
loaded, load_seconds = manager.activate(metadata)
|
||
provider = profile(spec, config).model_copy(update={"model": loaded["identifier"]})
|
||
model_record.update(loaded=loaded, load_seconds=load_seconds, provider=provider.model_dump())
|
||
settings = Settings(database=str(directory / "translations.sqlite3"), port=config.port, providers={"lmstudio": provider})
|
||
factory = lambda provider_config: CapturingProvider(provider_config, directory / "raw-responses" / f"model-{index}")
|
||
with isolated_api(settings, factory) as client:
|
||
project = base_project.model_copy(update={"name": "Benchmark: " + label, "provider": "lmstudio", "model": loaded["identifier"]})
|
||
project_data = request(client, "POST", "/projects", project.model_dump())
|
||
warmup, warm_seconds = timed_job(client, project_data["id"], items[:config.warmup_lines], min(10, config.warmup_lines), config,
|
||
lambda identity: write_json(directory / "active-job.json", {"job_id": identity, "phase": "warmup", "model": label}))
|
||
write_json(directory / "jobs" / f"model-{index}-warmup.json", warmup)
|
||
model_record["warmup_seconds"] = warm_seconds
|
||
if warmup["status"] != "completed":
|
||
raise RuntimeError(f"Warmup failed: {warmup.get('error') or warmup['status']}; see jobs/model-{index}-warmup.json")
|
||
# Rotate batch order between repeats to reduce fixed-order bias.
|
||
for repeat in range(config.repeats):
|
||
order = config.batch_sizes[repeat % len(config.batch_sizes):] + config.batch_sizes[:repeat % len(config.batch_sizes)]
|
||
for batch_size in order:
|
||
print(f" repeat {repeat + 1}/{config.repeats}, batch {batch_size} ...", flush=True)
|
||
job, elapsed = timed_job(client, project_data["id"], items, batch_size, config,
|
||
lambda identity: write_json(directory / "active-job.json", {"job_id": identity, "phase": "measured", "model": label}))
|
||
filename = f"model-{index}-batch-{batch_size}-repeat-{repeat + 1}.json"
|
||
write_json(directory / "jobs" / filename, job)
|
||
row = {"model": label, "model_index": index, "batch_size": batch_size, "repeat": repeat + 1,
|
||
"result_file": "jobs/" + filename, **summarize_run(job, elapsed)}
|
||
report["runs"].append(row)
|
||
save_report(directory, report)
|
||
print(f" {row['seconds']:.2f}s, {row['lines_per_second']:.2f} lines/s, failed={row['failed']}, retries={row['retry_attempts']}", flush=True)
|
||
except (TimeoutError, KeyboardInterrupt):
|
||
raise
|
||
except Exception as exc:
|
||
report["errors"].append(f"{label}: {exc}")
|
||
print(f" FAILED: {exc}", flush=True)
|
||
save_report(directory, report)
|
||
except BaseException as exc:
|
||
report["errors"].append(f"Benchmark interrupted: {type(exc).__name__}: {exc}")
|
||
report["state"] = "interrupted"
|
||
raise
|
||
finally:
|
||
manager.restore()
|
||
report["restore_errors"] = manager.restore_errors
|
||
if report["state"] != "interrupted":
|
||
report["state"] = "completed_with_errors" if report["errors"] or manager.restore_errors or any(row["status"] != "completed" for row in report["runs"]) else "completed"
|
||
report["finished_at"] = datetime.now(timezone.utc).isoformat()
|
||
save_report(directory, report)
|
||
return report
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description=__doc__)
|
||
parser.add_argument("--config", type=Path, default=Path("benchmark.json"))
|
||
parser.add_argument("--models", nargs="+", help="Downloaded model keys or paths from lms ls --json")
|
||
parser.add_argument("--input", type=Path)
|
||
parser.add_argument("--limit", type=int)
|
||
parser.add_argument("--repeats", type=int)
|
||
parser.add_argument("--batches", type=int, nargs="+")
|
||
parser.add_argument("--output", type=Path, help="New output directory")
|
||
args = parser.parse_args()
|
||
config_path = args.config.resolve()
|
||
values = json.loads(config_path.read_text(encoding="utf-8-sig"))
|
||
root = config_path.parent
|
||
if args.models:
|
||
values["models"] = [{"model": model} for model in args.models]
|
||
for name in ("limit", "repeats"):
|
||
if getattr(args, name) is not None:
|
||
values[name] = getattr(args, name)
|
||
if args.batches:
|
||
values["batch_sizes"] = args.batches
|
||
if args.input:
|
||
values["input"] = str(args.input.resolve())
|
||
config = BenchmarkConfig.model_validate(values)
|
||
directory = args.output.resolve() if args.output else root / "data" / "benchmarks" / (datetime.now().strftime("%Y%m%d-%H%M%S") + "-" + uuid4().hex[:6])
|
||
lock = Storage(str(root / "data" / "benchmarks" / "runner.sqlite3"))
|
||
with lock.exclusive_worker():
|
||
report = run(config, root, directory)
|
||
print(f"\nReport: {directory / 'report.html'}")
|
||
print(f"State: {report['state']}")
|
||
if report["state"] != "completed":
|
||
raise SystemExit(1)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
try:
|
||
main()
|
||
except KeyboardInterrupt:
|
||
raise SystemExit("Benchmark interrupted; partial results preserved.")
|
||
except (RuntimeError, ValueError, OSError, httpx.HTTPError, subprocess.SubprocessError) as exc:
|
||
raise SystemExit(str(exc)) from exc
|