adelevett's picture
Redeploy ELBench-live
68670f0 verified
Raw History Blame Contribute Delete
16.8 kB
"""Run a visitor selection through BenchmarkRunner, checkpoint every sample, and persist to LeaderboardStore."""
from __future__ import annotations
import asyncio
import json
import os
import shutil
import threading
from contextlib import contextmanager
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
from uuid import uuid4
from elbench.dashboard.catalog import Catalog
from elbench.dashboard.english import (
SUMMARY_GROUP_ENGLISH_LABELS,
english_module,
english_preflight_message,
english_summary_name,
)
from elbench.dashboard.leaderboard import (
LeaderboardStore,
basic_education_valid_means,
compute_normalized_scores,
)
from elbench.dashboard.plan import Plan
from elbench.dashboard.protocol import VISITOR_KEY_ENV
from elbench.execution.preflight import PreflightError, PreflightRunner
from elbench.execution.runner import BenchmarkRunner, RunOptions
from elbench.persistence.writers import OutputPaths
from elbench.summary import build_summary
from elbench.utils.spend import SpendLimits, spend_limits
_ENV_LOCK = threading.Lock()
@contextmanager
def visitor_api_key(api_key: str):
"""Hold the only visitor key slot for the whole critical section."""
secret = (api_key or "").strip()
if not secret:
raise ValueError("OpenRouter API key is empty.")
with _ENV_LOCK:
previous = os.environ.get(VISITOR_KEY_ENV)
os.environ[VISITOR_KEY_ENV] = secret
try:
yield secret
finally:
if previous is None:
os.environ.pop(VISITOR_KEY_ENV, None)
else:
os.environ[VISITOR_KEY_ENV] = previous
def make_run_id(config_hash: str | None = None) -> str:
stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S")
suffix = (config_hash or uuid4().hex)[:8]
return f"live-{stamp}-{suffix}"
def execute_run(
catalog: Catalog,
plan: Plan,
*,
max_concurrency: int,
accept_unpriced: bool,
run_id: str | None = None,
resume_previous: bool = True,
leaderboard_store: LeaderboardStore | None = None,
persist_to_leaderboard: bool = True,
total_budget_usd: float | None = None,
role_budgets_usd: dict[str, float | None] | None = None,
) -> dict[str, object]:
"""Execute ``plan`` with immediate per-sample checkpointing and fault-tolerant recovery."""
if plan.blocked:
raise PreflightError(" | ".join(plan.error_messages()) or "Selection is blocked.")
if plan.needs_price_ack and not accept_unpriced:
raise PreflightError(
"OpenRouter returned no token price for this model. Confirm the unknown price before running."
)
config = catalog.config.model_copy(deep=True)
# Flush every single sample immediately so long-running tasks never lose progress
config.app.default_run.checkpoint_interval = 1
model = plan.protocol.model_config
config.models[model.model_id] = model
if plan.judge_model_config is not None:
config.models[plan.judge_model_config.model_id] = plan.judge_model_config
config.judges.default_judge_model_id = plan.judge_model_config.model_id
config.judges.basic_education_evaluator_model_id = plan.judge_model_config.model_id
config.judges.basic_education_test_endpoint_model_id = plan.judge_model_config.model_id
if plan.multiturn_model_config is not None:
config.models[plan.multiturn_model_config.model_id] = plan.multiturn_model_config
config.judges.basic_education_test_endpoint_model_id = plan.multiturn_model_config.model_id
resumed_from_existing = False
resolved_run_id = (run_id or "").strip()
if leaderboard_store is not None and resume_previous:
existing = leaderboard_store.find_resumable_run(
config_hash=plan.config_hash,
explicit_run_id=resolved_run_id if resolved_run_id else None,
)
if existing is not None and existing.get("run_id"):
resolved_run_id = str(existing["run_id"])
resumed_from_existing = True
if not resolved_run_id:
resolved_run_id = make_run_id(plan.config_hash)
output_root = _hosted_output_root(config.project_root, resolved_run_id)
config.app.output_root = output_root
paths = OutputPaths.build(output_root, resolved_run_id, model.model_id)
if leaderboard_store is not None and (resumed_from_existing or resume_previous):
restored = leaderboard_store.prepare_resume_directory(
run_id=resolved_run_id,
model_id=model.model_id,
output_root=output_root,
)
if restored:
resumed_from_existing = True
report = PreflightRunner(config).run(
model_id=model.model_id,
modules=set(plan.selection.modules),
subsets=None,
source_files=set(plan.selection.source_files),
dimensions=set(plan.selection.dimensions) or None,
max_samples=plan.selection.max_samples,
judge_enabled=True,
require_real_judges=True,
)
if not report.ok:
messages = [english_preflight_message(issue) for issue in report.issues if issue.level == "error"]
raise PreflightError(" | ".join(messages))
secret = os.environ.get(VISITOR_KEY_ENV, "")
basic_root = config.project_root / "outputs" / "basic_education" / resolved_run_id
interrupted_error: BaseException | None = None
result: dict[str, Any] = {}
limits = build_spend_limits(plan, total_budget_usd, role_budgets_usd)
ledger_path = output_root / "spend.jsonl"
spent: dict[str, float] = {}
try:
with spend_limits(limits, ledger_path) as ledger:
result = _run_blocking(
BenchmarkRunner(config).run(
RunOptions(
model_id=model.model_id,
run_id=resolved_run_id,
modules=set(plan.selection.modules),
subsets=None,
source_files=set(plan.selection.source_files),
dimensions=set(plan.selection.dimensions) or None,
max_samples=plan.selection.max_samples,
max_samples_per_file=plan.selection.max_samples_per_file,
max_concurrency=max_concurrency,
resume=bool(resume_previous),
judge_enabled=True,
progress_enabled=True,
)
)
)
spent = ledger.spent()
except BaseException as exc: # noqa: BLE001
interrupted_error = exc
if ledger_path.exists():
spent = _ledger_totals(ledger_path, limits)
finally:
scrub_tree(output_root, secret)
scrub_tree(basic_root, secret)
if basic_root.exists():
shutil.rmtree(basic_root, ignore_errors=True)
has_partial_records = (
(paths.judged_path.exists() and paths.judged_path.stat().st_size > 0)
or (paths.failures_path.exists() and paths.failures_path.stat().st_size > 0)
)
if interrupted_error is not None and not has_partial_records:
raise interrupted_error
if not paths.summary_path.exists() or interrupted_error is not None:
summary = build_summary(paths.judged_path, paths.failures_path)
paths.summary_path.write_text(
json.dumps(summary, ensure_ascii=False, indent=2),
encoding="utf-8",
)
else:
summary = json.loads(paths.summary_path.read_text(encoding="utf-8"))
total_judged = int(summary.get("total_judged") or 0)
total_failures = int(summary.get("total_failures") or 0)
quota_exhausted = bool(result.get("quota_exhausted", False))
spent_total = round(sum(spent.values()), 6)
budget_stopped = quota_exhausted and _limit_reached(limits, spent)
if interrupted_error is not None:
run_status = "interrupted"
elif quota_exhausted or total_failures > 0:
run_status = "partial"
else:
run_status = "completed"
normalized_scores = compute_normalized_scores(
summary,
basic_valid_means=basic_education_valid_means(paths.judged_path),
)
leaderboard_record = {
"run_id": resolved_run_id,
"config_hash": plan.config_hash,
"target_model": plan.protocol.openrouter_model,
"judge_model": (
plan.judge_model_config.model_name
if plan.judge_model_config is not None
else (plan.judge_model_id or "default")
),
"use_system_one_judge": bool(plan.use_system_one_judge),
"use_multiturn_chat_model": bool(plan.use_multiturn_chat_model),
"multiturn_model": (
plan.multiturn_model_config.model_name
if plan.multiturn_model_config is not None
else None
),
"modules": [english_module(m) for m in sorted(plan.selection.modules)],
"source_files": sorted(plan.selection.source_files),
"max_samples": plan.selection.max_samples,
"max_samples_per_file": plan.selection.max_samples_per_file,
"dimensions": sorted(plan.selection.dimensions),
"is_full_benchmark": bool(plan.is_full_benchmark),
"covers_core": bool(plan.covers_core),
"routing": {role: route.record() for role, route in plan.routes.items()},
# False when every answer was checked against the answer key, so the run scores the same on every judge's board.
"judge_used": bool(plan.judge_calls or plan.basic_tasks),
"status": run_status,
"resumed": resumed_from_existing,
"total_judged": total_judged,
"total_failures": total_failures,
"quota_exhausted": quota_exhausted,
"interrupted_error": (
public_error(interrupted_error, secret) if interrupted_error is not None else None
),
"prompt_usd_per_token": plan.prompt_usd_per_token,
"completion_usd_per_token": plan.completion_usd_per_token,
"low_usd": plan.low_usd,
"high_usd": plan.high_usd,
"spent_usd": spent_total,
"spent_usd_by_role": {role: round(value, 6) for role, value in spent.items()},
"budget_usd": total_budget_usd,
"role_budgets_usd": {k: v for k, v in (role_budgets_usd or {}).items() if v is not None},
"stopped_by_budget": budget_stopped,
"normalized_scores": normalized_scores,
}
if persist_to_leaderboard and leaderboard_store is not None and (total_judged > 0 or total_failures > 0):
leaderboard_store.save_run(
leaderboard_record,
output_Paths=paths,
push_to_hub=True,
)
download_paths = [
str(p)
for p in (paths.judged_path, paths.summary_path, paths.failures_path)
if p.exists() and p.stat().st_size > 0
]
return {
**result,
"run_id": resolved_run_id,
"status": run_status,
"resumed": resumed_from_existing,
"config_hash": plan.config_hash,
"spent_usd": spent_total,
"spent_usd_by_role": leaderboard_record["spent_usd_by_role"],
"stopped_by_budget": budget_stopped,
"interrupted_error": leaderboard_record["interrupted_error"],
"summary_path": str(paths.summary_path),
"judged_path": str(paths.judged_path),
"failures_path": str(paths.failures_path),
"summary": summary,
"normalized_scores": normalized_scores,
"leaderboard_record": leaderboard_record,
"summary_rows": summary_rows(summary),
"download_paths": download_paths,
"output_root": str(output_root),
}
def build_spend_limits(
plan: Plan,
total_budget_usd: float | None,
role_budgets_usd: dict[str, float | None] | None,
) -> SpendLimits:
"""Attribute each request to the model under test, the judge, or the simulated student."""
roles = {plan.protocol.model_config.model_id: "target"}
prices: dict[str, tuple[float, float]] = {}
if plan.prompt_usd_per_token is not None and plan.completion_usd_per_token is not None:
prices["target"] = (plan.prompt_usd_per_token, plan.completion_usd_per_token)
if plan.judge_model_config is not None:
roles[plan.judge_model_config.model_id] = "judge"
if plan.judge_prices is not None:
prices["judge"] = plan.judge_prices
if plan.multiturn_model_config is not None:
roles[plan.multiturn_model_config.model_id] = "student"
if plan.student_prices is not None:
prices["student"] = plan.student_prices
elif plan.judge_prices is not None:
# Without a separate student model, the judge model also plays the student.
prices["student"] = plan.judge_prices
return SpendLimits(
total_usd=_positive_or_none(total_budget_usd),
per_role_usd={
role: float(value)
for role, value in (role_budgets_usd or {}).items()
if _positive_or_none(value) is not None
},
prices=prices,
roles_by_model_id=roles,
)
def _positive_or_none(value: float | None) -> float | None:
if value is None:
return None
try:
number = float(value)
except (TypeError, ValueError):
return None
return number if number > 0 else None
def _limit_reached(limits: SpendLimits, spent: dict[str, float]) -> bool:
if limits.total_usd is not None and sum(spent.values()) >= limits.total_usd:
return True
return any(spent.get(role, 0.0) >= cap for role, cap in limits.per_role_usd.items())
def _ledger_totals(path: Path, limits: SpendLimits) -> dict[str, float]:
from elbench.utils.spend import SpendLedger
return SpendLedger(path, limits).spent()
def judged_progress_path(project_root: Path, run_id: str, model_id: str) -> Path:
output_root = _hosted_output_root(project_root, run_id)
paths = OutputPaths.build(output_root, run_id, model_id)
return paths.judged_path
def failures_progress_path(project_root: Path, run_id: str, model_id: str) -> Path:
output_root = _hosted_output_root(project_root, run_id)
paths = OutputPaths.build(output_root, run_id, model_id)
return paths.failures_path
def summary_rows(summary: dict) -> list[list[object]]:
rows: list[list[object]] = []
for section in ("by_module", "by_subset", "by_task", "by_source_file", "by_dimension", "by_scene"):
block = summary.get(section) or {}
if not isinstance(block, dict):
continue
group_label = SUMMARY_GROUP_ENGLISH_LABELS.get(section, section)
for name, stats in block.items():
if not isinstance(stats, dict):
continue
avg_score = stats.get("avg_score")
pass_rate = stats.get("pass_rate")
rows.append(
[
group_label,
english_summary_name(section, str(name)),
stats.get("count"),
round(float(avg_score), 4) if isinstance(avg_score, (int, float)) else avg_score,
round(float(pass_rate), 4) if isinstance(pass_rate, (int, float)) else pass_rate,
]
)
return rows
def scrub_tree(root: Path, secret: str) -> None:
if not secret or len(secret) < 8 or not root.exists():
return
token = secret.encode("utf-8")
replacement = b"[redacted]"
for path in root.rglob("*"):
if not path.is_file():
continue
try:
size = path.stat().st_size
except OSError:
continue
if size > 50_000_000 or size == 0:
continue
try:
data = path.read_bytes()
except OSError:
continue
if token not in data:
continue
path.write_bytes(data.replace(token, replacement))
def public_error(exc: BaseException, secret: str) -> str:
text = f"{type(exc).__name__}: {exc}"
if secret:
text = text.replace(secret, "[redacted]")
return text
def _hosted_output_root(project_root: Path, run_id: str) -> Path:
path = project_root / "outputs" / "hosted" / run_id
path.mkdir(parents=True, mode=0o700, exist_ok=True)
path.chmod(0o700)
return path
def _run_blocking(coro):
try:
asyncio.get_running_loop()
except RuntimeError:
return asyncio.run(coro)
box: dict[str, object] = {}
def _target() -> None:
try:
box["value"] = asyncio.run(coro)
except BaseException as exc: # noqa: BLE001
box["error"] = exc
thread = threading.Thread(target=_target)
thread.start()
thread.join()
if "error" in box:
raise box["error"]
return box["value"]