File size: 7,953 Bytes
d8ebb7a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 | """Budget pacing, rate limits, device cooldown and the kill switch (spec section 10).
State is in memory and resets on restart. Provider prepaid credit is the hard cap; these
counters only pace. One threading.Lock guards all counters: every critical section is short and
never awaits, so it is safe from async generation code and from sync admin callbacks.
"""
from __future__ import annotations
import asyncio
import hashlib
import logging
import re
import secrets
import statistics
import threading
import time
from collections import defaultdict, deque
from collections.abc import AsyncIterator, Callable
from contextlib import asynccontextmanager
from dataclasses import dataclass
from typing import Literal
from src.config import LimitsSection
from src.errors import BudgetReachedError, CooldownError, PausedError, RateLimitedError
logger = logging.getLogger(__name__)
MINUTE = 60.0
HOUR = 3600.0
STATS_WINDOW = 600.0
MAX_EVENTS = 5000
MAX_TRACKED_DEVICES = 5000
UNKNOWN_DEVICE = "-"
_DEVICE_ID = re.compile(r"[A-Za-z0-9_-]{16,64}")
_SALT = secrets.token_hex(8)
Kind = Literal["ok", "error", "refusal", "limited"]
def new_device_id() -> str:
return secrets.token_urlsafe(16)
def clean_device_id(raw: object) -> str | None:
"""The stored id if it looks like one we issued. The browser controls it, so never trust it."""
return raw if isinstance(raw, str) and _DEVICE_ID.fullmatch(raw) else None
def hash_device_id(raw: str) -> str:
"""Short salted hash for limits and logs. The salt changes on every restart."""
return hashlib.sha256(f"{_SALT}:{raw}".encode()).hexdigest()[:12]
def device_identity(raw: object) -> tuple[str, str]:
"""(id to keep in the browser, hash to use for limits and logs)."""
device_id = clean_device_id(raw) or new_device_id()
return device_id, hash_device_id(device_id)
@dataclass
class Reservation:
count: int
cost: float
epoch: int = 0
settled: bool = False
@dataclass(frozen=True)
class Snapshot:
images_used: int
max_images: int
spend_usd: float
max_cost_usd: float
in_flight: int
paused: bool
recent: dict[Kind, int]
p50_seconds: dict[str, float]
class LimitService:
def __init__(self, limits: LimitsSection, clock: Callable[[], float] = time.monotonic) -> None:
self._limits = limits
self._clock = clock
self._lock = threading.Lock()
self._slots = asyncio.Semaphore(limits.max_concurrent_generations)
self._paused = False
self._epoch = 0
self._images = 0
self._spend = 0.0
self._in_flight = 0
self._starts: deque[float] = deque()
self._device_last: dict[str, float] = {}
self._device_starts: dict[str, deque[float]] = {}
self._events: deque[tuple[float, Kind, str, float]] = deque(maxlen=MAX_EVENTS)
def reserve(self, *, count: int, est_cost: float, device: str | None) -> Reservation:
"""Claim budget for `count` images before calling a provider, or raise a typed error."""
limits = self._limits
key = device or UNKNOWN_DEVICE
now = self._clock()
with self._lock:
if self._paused:
raise PausedError
last = self._device_last.get(key)
if last is not None and now - last < limits.device_cooldown_seconds:
raise CooldownError
hour = self._device_starts.get(key) or deque()
_drop_older(hour, now - HOUR)
if limits.device_max_images_per_hour and (
len(hour) + count > limits.device_max_images_per_hour
):
raise RateLimitedError
_drop_older(self._starts, now - MINUTE)
if len(self._starts) + count > limits.max_generations_per_minute:
raise RateLimitedError
if (
self._images + count > limits.max_images
or self._spend + est_cost > limits.max_cost_usd
):
raise BudgetReachedError
self._device_last[key] = now
self._device_starts[key] = hour
hour.extend([now] * count)
self._starts.extend([now] * count)
self._images += count
self._spend += est_cost
self._forget_idle_devices(now)
return Reservation(count=count, cost=est_cost, epoch=self._epoch)
def reconcile(self, reservation: Reservation, actual_cost: float) -> None:
with self._lock:
if reservation.settled:
return
reservation.settled = True
if reservation.epoch != self._epoch:
return
self._spend = max(0.0, self._spend + actual_cost - reservation.cost)
def release(self, reservation: Reservation) -> None:
"""Return the budget of a call that produced nothing. The cooldown still counts."""
with self._lock:
if reservation.settled:
return
reservation.settled = True
if reservation.epoch != self._epoch:
return
self._images = max(0, self._images - reservation.count)
self._spend = max(0.0, self._spend - reservation.cost)
@asynccontextmanager
async def slot(self) -> AsyncIterator[None]:
"""Hold one of the global concurrent-generation slots."""
async with self._slots:
with self._lock:
self._in_flight += 1
try:
yield
finally:
with self._lock:
self._in_flight -= 1
def record(self, kind: Kind, model: str, seconds: float) -> None:
with self._lock:
self._events.append((self._clock(), kind, model, seconds))
@property
def paused(self) -> bool:
with self._lock:
return self._paused
def set_paused(self, paused: bool) -> None:
with self._lock:
self._paused = paused
logger.info("admin action: paused=%s", paused)
def reset_counters(self) -> None:
"""Zero the counters. Reservations still in flight no longer count against the new ones."""
with self._lock:
self._epoch += 1
self._images = 0
self._spend = 0.0
self._starts.clear()
self._events.clear()
logger.info("admin action: counters reset")
def snapshot(self) -> Snapshot:
now = self._clock()
with self._lock:
recent = [e for e in self._events if now - e[0] <= STATS_WINDOW]
counts: dict[Kind, int] = {"ok": 0, "error": 0, "refusal": 0, "limited": 0}
latencies: defaultdict[str, list[float]] = defaultdict(list)
for _, kind, model, seconds in recent:
counts[kind] += 1
if kind == "ok":
latencies[model].append(seconds)
return Snapshot(
images_used=self._images,
max_images=self._limits.max_images,
spend_usd=self._spend,
max_cost_usd=self._limits.max_cost_usd,
in_flight=self._in_flight,
paused=self._paused,
recent=counts,
p50_seconds={m: statistics.median(v) for m, v in latencies.items()},
)
def _forget_idle_devices(self, now: float) -> None:
"""Bound memory if many device ids are forged. Caller holds the lock."""
if len(self._device_last) <= MAX_TRACKED_DEVICES:
return
stale = [k for k, t in self._device_last.items() if now - t > HOUR]
for key in stale:
self._device_last.pop(key, None)
for key in [k for k, v in self._device_starts.items() if not v or now - v[-1] > HOUR]:
del self._device_starts[key]
def _drop_older(stamps: deque[float], cutoff: float) -> None:
while stamps and stamps[0] < cutoff:
stamps.popleft()
|