Spaces:
Sleeping
Sleeping
| """Optional spending limits for one run, shared with the basic-education runtime subprocess. | |
| Nothing here is active unless a caller installs limits with ``spend_limits``. | |
| The CLI never does, so a normal benchmark run behaves exactly as before. | |
| Each paid request appends one line to a JSONL ledger, so the main process and the | |
| runtime subprocess (which inherits the two environment variables below) count | |
| against the same totals. Requests are checked before they are sent, so a limit | |
| can be overshot by the requests already in flight. | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import os | |
| import threading | |
| from contextlib import contextmanager | |
| from dataclasses import dataclass, field | |
| from pathlib import Path | |
| from typing import Any, Iterator | |
| LEDGER_ENV = "ELBENCH_SPEND_LEDGER" | |
| LIMITS_ENV = "ELBENCH_SPEND_LIMITS" | |
| ROLE_LABELS = { | |
| "target": "model under test", | |
| "judge": "judge", | |
| "student": "simulated student", | |
| } | |
| class BudgetExceeded(RuntimeError): | |
| """Raised instead of sending a request once a spending limit is reached. | |
| ``retryable = False`` keeps ``BenchmarkRunner`` from retrying it, and the message | |
| says "quota exceeded" so the runner stops the way it does when a key runs out of | |
| credit: finished items are kept and the run can be resumed with a higher limit. | |
| """ | |
| retryable = False | |
| class SpendLimits: | |
| total_usd: float | None = None | |
| per_role_usd: dict[str, float] = field(default_factory=dict) | |
| # USD per token, (prompt, completion), used when a response carries no cost. | |
| prices: dict[str, tuple[float, float]] = field(default_factory=dict) | |
| roles_by_model_id: dict[str, str] = field(default_factory=dict) | |
| def to_json(self) -> str: | |
| return json.dumps( | |
| { | |
| "total_usd": self.total_usd, | |
| "per_role_usd": self.per_role_usd, | |
| "prices": {role: list(pair) for role, pair in self.prices.items()}, | |
| "roles_by_model_id": self.roles_by_model_id, | |
| } | |
| ) | |
| def from_json(cls, raw: str) -> "SpendLimits": | |
| data = json.loads(raw) | |
| return cls( | |
| total_usd=data.get("total_usd"), | |
| per_role_usd={k: float(v) for k, v in (data.get("per_role_usd") or {}).items()}, | |
| prices={k: (float(v[0]), float(v[1])) for k, v in (data.get("prices") or {}).items()}, | |
| roles_by_model_id=dict(data.get("roles_by_model_id") or {}), | |
| ) | |
| class SpendLedger: | |
| def __init__(self, path: Path, limits: SpendLimits) -> None: | |
| self.path = path | |
| self.limits = limits | |
| self._lock = threading.Lock() | |
| self._offset = 0 | |
| self._spent: dict[str, float] = {} | |
| def spent(self) -> dict[str, float]: | |
| with self._lock: | |
| self._read_new_lines() | |
| return dict(self._spent) | |
| def check(self, role: str) -> None: | |
| spent = self.spent() | |
| total = sum(spent.values()) | |
| limit = self.limits.total_usd | |
| if limit is not None and total >= limit: | |
| raise BudgetExceeded( | |
| f"Spending limit reached (quota exceeded): ${total:.4f} spent of the ${limit:.2f} total limit." | |
| ) | |
| role_limit = self.limits.per_role_usd.get(role) | |
| if role_limit is not None and spent.get(role, 0.0) >= role_limit: | |
| raise BudgetExceeded( | |
| f"Spending limit reached (quota exceeded): ${spent.get(role, 0.0):.4f} spent on the " | |
| f"{ROLE_LABELS.get(role, role)} of its ${role_limit:.2f} limit." | |
| ) | |
| def record(self, role: str, usage: Any) -> float: | |
| cost = usage_cost(usage, self.limits.prices.get(role)) | |
| if cost <= 0: | |
| return 0.0 | |
| line = json.dumps({"role": role, "usd": cost}) + "\n" | |
| with self._lock: | |
| with self.path.open("a", encoding="utf-8") as handle: | |
| handle.write(line) | |
| return cost | |
| def _read_new_lines(self) -> None: | |
| if not self.path.exists(): | |
| return | |
| with self.path.open("rb") as handle: | |
| handle.seek(self._offset) | |
| chunk = handle.read() | |
| end = chunk.rfind(b"\n") | |
| if end < 0: | |
| return | |
| self._offset += end + 1 | |
| for raw in chunk[: end + 1].splitlines(): | |
| try: | |
| entry = json.loads(raw) | |
| except json.JSONDecodeError: | |
| continue | |
| role = str(entry.get("role") or "") | |
| self._spent[role] = self._spent.get(role, 0.0) + float(entry.get("usd") or 0.0) | |
| def usage_cost(usage: Any, price: tuple[float, float] | None) -> float: | |
| """USD for one response: OpenRouter's reported ``cost`` when present, else tokens x price.""" | |
| if not isinstance(usage, dict): | |
| return 0.0 | |
| cost = usage.get("cost") | |
| if isinstance(cost, (int, float)) and not isinstance(cost, bool) and cost >= 0: | |
| return float(cost) | |
| if price is None: | |
| return 0.0 | |
| prompt_tokens = usage.get("prompt_tokens", usage.get("input_tokens")) or 0 | |
| completion_tokens = usage.get("completion_tokens", usage.get("output_tokens")) or 0 | |
| try: | |
| return float(prompt_tokens) * price[0] + float(completion_tokens) * price[1] | |
| except (TypeError, ValueError): | |
| return 0.0 | |
| _ACTIVE: SpendLedger | None = None | |
| _ACTIVE_LOCK = threading.Lock() | |
| def active_ledger() -> SpendLedger | None: | |
| """The installed ledger, or one rebuilt from the environment inside the runtime subprocess.""" | |
| global _ACTIVE | |
| if _ACTIVE is not None: | |
| return _ACTIVE | |
| path = os.environ.get(LEDGER_ENV) | |
| raw = os.environ.get(LIMITS_ENV) | |
| if not path or not raw: | |
| return None | |
| with _ACTIVE_LOCK: | |
| if _ACTIVE is None: | |
| _ACTIVE = SpendLedger(Path(path), SpendLimits.from_json(raw)) | |
| return _ACTIVE | |
| def role_for_model(model_id: str | None) -> str | None: | |
| ledger = active_ledger() | |
| if ledger is None or not model_id: | |
| return None | |
| return ledger.limits.roles_by_model_id.get(model_id) | |
| def check(role: str | None) -> None: | |
| ledger = active_ledger() | |
| if ledger is not None and role: | |
| ledger.check(role) | |
| def record(role: str | None, usage: Any) -> None: | |
| ledger = active_ledger() | |
| if ledger is not None and role: | |
| ledger.record(role, usage) | |
| def spend_limits(limits: SpendLimits, ledger_path: Path) -> Iterator[SpendLedger]: | |
| """Install ``limits`` for this process and any subprocess it starts.""" | |
| global _ACTIVE | |
| ledger_path.parent.mkdir(parents=True, exist_ok=True) | |
| ledger_path.touch(exist_ok=True) | |
| ledger = SpendLedger(ledger_path, limits) | |
| previous = {name: os.environ.get(name) for name in (LEDGER_ENV, LIMITS_ENV)} | |
| os.environ[LEDGER_ENV] = str(ledger_path) | |
| os.environ[LIMITS_ENV] = limits.to_json() | |
| with _ACTIVE_LOCK: | |
| _ACTIVE = ledger | |
| try: | |
| yield ledger | |
| finally: | |
| with _ACTIVE_LOCK: | |
| _ACTIVE = None | |
| for name, value in previous.items(): | |
| if value is None: | |
| os.environ.pop(name, None) | |
| else: | |
| os.environ[name] = value | |