"""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"]