responsible-request 0.4.0__py3-none-any.whl

responsible_request/sources.py

@@ -0,0 +1,207 @@
+"""Read request records back from the log while the client runs (cache lookups, spent cost)."""
+
+from __future__ import annotations
+
+import contextlib
+import json
+import os
+import sqlite3
+from pathlib import Path
+from typing import IO, Any, Protocol
+
+from .logfiles import compression_of, open_binary, rotated_segments
+
+
+class _Target(Protocol):
+    @property
+    def sqlite(self) -> str | Path | None: ...
+    @property
+    def jsonl(self) -> str | Path | None: ...
+
+
+def _cacheable(row: dict[str, Any]) -> bool:
+    return (
+        bool(row.get("cache_key"))
+        and row.get("status_code") == 200
+        and not row.get("error")
+        and not row.get("stream")
+        and isinstance(row.get("response_body"), dict)
+    )
+
+
+class SQLiteSource:
+    """Looks up responses in the ``requests`` table written by the SQLite log sink."""
+
+    _QUERY = (
+        "SELECT response_body FROM requests WHERE cache_key = ? AND status_code = 200 "
+        "AND error IS NULL AND NOT stream AND response_body IS NOT NULL ORDER BY rowid LIMIT 1"
+    )
+    _COST = "SELECT SUM(cost_usd) FROM requests WHERE NOT cache_hit"
+
+    def __init__(self, path: str | Path) -> None:
+        self.path = Path(path)
+        self._conn: sqlite3.Connection | None = None
+
+    def _query(self, sql: str, *args: Any) -> Any:
+        if self._conn is None:
+            if not self.path.exists():  # the log sink creates it with the first record
+                return None
+            self._conn = sqlite3.connect(self.path, check_same_thread=False)
+        try:
+            return self._conn.execute(sql, args).fetchone()
+        except sqlite3.OperationalError:  # table or column not created yet
+            return None
+
+    def get(self, key: str) -> dict[str, Any] | None:
+        row = self._query(self._QUERY, key)
+        if row is None:
+            return None
+        with contextlib.suppress(ValueError):
+            body = json.loads(row[0])
+            if isinstance(body, dict):
+                return body
+        return None
+
+    def total_cost(self) -> float:
+        """Sum of ``cost_usd`` over all records that were not cache hits."""
+        row = self._query(self._COST)
+        return float(row[0]) if row is not None and row[0] is not None else 0.0
+
+    def close(self) -> None:
+        if self._conn is not None:
+            self._conn.close()
+            self._conn = None
+
+
+class _Rotated(Exception):
+    """The active file was replaced since it was last indexed."""
+
+
+class JSONLSource:
+    """Indexes a JSONL log by cache key and reads matching lines on demand.
+
+    The segments loguru rotated away from the log are indexed too. Responses from compressed
+    segments are kept in memory (only the first one per cache key); everything else is read back
+    by byte offset. New lines in the active file are indexed incrementally on each lookup, and a
+    rotation while the client runs triggers a full re-index.
+    """
+
+    def __init__(self, path: str | Path) -> None:
+        self.path = Path(path)
+        self._live = compression_of(self.path) is None  # a compressed file does not grow
+        self._reset()
+
+    def _reset(self) -> None:
+        # cache key -> (plain file, byte offset of the first matching line) or response JSON
+        self._index: dict[str, tuple[Path, int] | str] = {}
+        self._cost = 0.0  # cost_usd of the indexed records that were not cache hits
+        self._loaded = False  # rotated or compressed segments indexed
+        self._file_id: tuple[int, int] | None = None  # (device, inode) of the active file
+        self._offset = 0  # bytes of the active file indexed so far
+
+    def get(self, key: str) -> dict[str, Any] | None:
+        if key not in self._index:
+            self._scan()
+        for _ in range(2):
+            entry = self._index.get(key)
+            if entry is None:
+                return None
+            body = self._read(key, entry)
+            if body is not None:
+                return body
+            # a rotation renamed, compressed or replaced the file since it was indexed
+            self._reset()
+            self._scan()
+        return None
+
+    @staticmethod
+    def _read(key: str, entry: tuple[Path, int] | str) -> dict[str, Any] | None:
+        if isinstance(entry, str):
+            body: dict[str, Any] = json.loads(entry)
+            return body
+        path, pos = entry
+        try:
+            with path.open("rb") as fh:
+                fh.seek(pos)
+                row = json.loads(fh.readline())
+        except (FileNotFoundError, ValueError):
+            return None
+        if not isinstance(row, dict) or row.get("cache_key") != key:
+            return None
+        body = row["response_body"]
+        return body
+
+    def _scan(self) -> None:
+        for _ in range(5):
+            try:
+                if not self._loaded:
+                    self._load_segments()
+                if self._live:
+                    self._scan_active()
+                return
+            except (FileNotFoundError, _Rotated):  # loguru renamed or compressed a file under us
+                self._reset()
+
+    def _load_segments(self) -> None:
+        segments = rotated_segments(self.path) if self._live else [self.path]
+        for segment in segments:
+            with open_binary(segment) as fh:
+                self._index_lines(fh, segment, 0, in_memory=compression_of(segment) is not None)
+        self._loaded = True
+        if self._live:
+            self._scan_active()
+            if rotated_segments(self.path) != segments:  # rotated while we were reading
+                raise _Rotated
+
+    def _scan_active(self) -> None:
+        try:
+            fh = self.path.open("rb")
+        except FileNotFoundError:  # nothing written yet, or in the middle of a rotation
+            return
+        with fh:
+            stat = os.fstat(fh.fileno())
+            file_id = (stat.st_dev, stat.st_ino)
+            if self._file_id is None:
+                self._file_id = file_id
+            elif file_id != self._file_id or stat.st_size < self._offset:
+                raise _Rotated
+            if stat.st_size > self._offset:
+                fh.seek(self._offset)
+                self._offset = self._index_lines(fh, self.path, self._offset, in_memory=False)
+
+    def _index_lines(self, fh: IO[bytes], path: Path, pos: int, *, in_memory: bool) -> int:
+        for line in fh:
+            if not line.endswith(b"\n"):  # still being written
+                break
+            with contextlib.suppress(ValueError):
+                row = json.loads(line)
+                if not isinstance(row, dict):
+                    continue
+                if _cacheable(row) and row["cache_key"] not in self._index:
+                    self._index[row["cache_key"]] = (
+                        json.dumps(row["response_body"]) if in_memory else (path, pos)
+                    )
+                if not row.get("cache_hit"):
+                    self._cost += float(row.get("cost_usd") or 0.0)
+            pos += len(line)
+        return pos
+
+    def total_cost(self) -> float:
+        """Sum of ``cost_usd`` over all records that were not cache hits."""
+        self._scan()
+        return self._cost
+
+    def close(self) -> None:
+        pass
+
+
+RecordSource = SQLiteSource | JSONLSource
+
+
+def open_source(target: _Target | None) -> RecordSource | None:
+    """A reader for the SQLite database of ``target``, else its JSONL file, else None."""
+    if target is not None and target.sqlite is not None:
+        return SQLiteSource(target.sqlite)
+    if target is not None and target.jsonl is not None:
+        return JSONLSource(target.jsonl)
+    return None

responsible_request/throttle.py

@@ -0,0 +1,229 @@
+"""Per-model throttling state: limiter + load estimator + controller."""
+
+from __future__ import annotations
+
+import random
+import threading
+import time
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any
+
+from loguru import logger
+
+from .config import ThrottleConfig
+from .controller import State, ThrottleController, Transition
+from .estimator import LatencyEstimator, LoadEstimator
+from .limiter import AdaptiveLimiter
+from .records import RequestRecord
+
+
+@dataclass
+class LaneStats:
+    requests: int = 0
+    errors: int = 0
+    prompt_tokens: int = 0
+    completion_tokens: int = 0
+    cost_usd: float = 0.0
+
+    def add(self, record: RequestRecord) -> None:
+        self.requests += 1
+        if record.error or (record.status_code or 0) >= 400:
+            self.errors += 1
+        self.prompt_tokens += record.prompt_tokens or 0
+        self.completion_tokens += record.completion_tokens or 0
+        self.cost_usd += record.cost_usd or 0.0
+
+
+class Lane:
+    """Throttling state for one (endpoint, model) pair.
+
+    Endpoints usually serve every model from its own backend, so load is tracked per model.
+    Lanes for endpoints that are not in ``observe_paths`` are paced and react to server errors,
+    but they do not adapt to latency.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        config: ThrottleConfig,
+        *,
+        observed: bool,
+        clock: Callable[[], float] = time.monotonic,
+        rng: random.Random | None = None,
+    ) -> None:
+        self.name = name
+        self.config = config
+        self.observed = observed
+        self.limiter = AdaptiveLimiter(config.initial_rpm, config.max_concurrency)
+        self.estimator: LoadEstimator = (
+            config.estimator_factory(config)
+            if config.estimator_factory is not None
+            else LatencyEstimator(config, clock=clock)
+        )
+        self.controller = ThrottleController(config, clock=clock)
+        self.totals = LaneStats()
+        self.since_summary = LaneStats()
+        self.last_record: RequestRecord | None = None
+        self._rng = rng or random.Random()
+        self._backoff_level = 0  # consecutive 429s without Retry-After (fixed mode)
+        self._lock = threading.Lock()
+
+    @property
+    def rpm(self) -> float:
+        return self.controller.rpm
+
+    @property
+    def state(self) -> State:
+        return self.controller.state
+
+    def observe(
+        self, record: RequestRecord, *, eligible: bool, retry_after: float | None = None
+    ) -> Transition | None:
+        """Update load estimate and rate after a request finished.
+
+        ``retry_after`` is the server's ``Retry-After`` in seconds, if it sent one.
+        """
+        with self._lock:
+            if not self.config.adaptive:
+                self._handle_fixed(record, retry_after)
+            return self._observe(record, eligible=eligible)
+
+    def _handle_fixed(self, record: RequestRecord, retry_after: float | None) -> None:
+        """Fixed mode: pause this lane after a 429 instead of changing the rate."""
+        status = record.status_code
+        if status is not None and status < 400:
+            self._backoff_level = 0
+        if status != 429:
+            return
+        cfg = self.config
+        if retry_after is not None:
+            delay, reason = retry_after, "Retry-After"
+        else:
+            # 429s of requests sent before the current pause began do not raise the level.
+            if self.limiter.paused_s == 0 or self._backoff_level == 0:
+                self._backoff_level += 1
+            exponent = min(self._backoff_level - 1, 30)
+            delay = min(cfg.backoff_max_s, cfg.backoff_initial_s * 2**exponent)
+            reason = f"backoff #{self._backoff_level}"
+        delay += self._rng.uniform(0, cfg.backoff_jitter_s)
+        self.limiter.pause(delay)
+        logger.warning("{}: HTTP 429, pausing {:.1f} s ({})", self.name, delay, reason)
+
+    def _observe(self, record: RequestRecord, *, eligible: bool) -> Transition | None:
+        self.last_record = record
+        self.totals.add(record)
+        self.since_summary.add(record)
+        congested = is_congestion(record)
+        if not self.observed:
+            ratio: float | None = 1.0  # non-latency lanes only react to server errors
+        else:
+            if eligible and not congested and record.error is None:
+                self.estimator.observe(record, throttled=self.state is State.THROTTLED)
+            ratio = self.estimator.load_ratio()
+        transition = self.controller.update(ratio, congested=congested)
+        if transition is not None:
+            self.limiter.set_rpm(self.controller.rpm)
+            _log_transition(self.name, transition)
+        return transition
+
+    def snapshot(self) -> dict[str, Any]:
+        with self._lock:
+            return self._snapshot()
+
+    def _snapshot(self) -> dict[str, Any]:
+        return {
+            "state": self.state.value,
+            "rpm": round(self.rpm, 2),
+            "baseline": self.estimator.baseline,
+            "load_ratio": self.estimator.load_ratio() if self.observed else None,
+            "in_flight": self.limiter.in_flight,
+            "paused_s": round(self.limiter.paused_s, 2),
+            "requests": self.totals.requests,
+            "errors": self.totals.errors,
+            "prompt_tokens": self.totals.prompt_tokens,
+            "completion_tokens": self.totals.completion_tokens,
+            "cost_usd": round(self.totals.cost_usd, 6),
+        }
+
+
+def is_congestion(record: RequestRecord) -> bool:
+    """429, 5xx and timeouts mean the server is struggling."""
+    status = record.status_code
+    if status is not None and (status == 429 or status >= 500):
+        return True
+    return record.error is not None and "Timeout" in record.error
+
+
+def _log_transition(name: str, t: Transition) -> None:
+    if t.new_state is State.THROTTLED and t.old_state is not State.THROTTLED:
+        logger.warning(
+            "{}: {} -> throttling {:.0f} -> {:.0f} RPM", name, t.reason, t.old_rpm, t.new_rpm
+        )
+    elif t.old_rpm != t.new_rpm:
+        logger.info(
+            "{}: {} ({:.0f} -> {:.0f} RPM, {})",
+            name,
+            t.reason,
+            t.old_rpm,
+            t.new_rpm,
+            t.new_state.value,
+        )
+    else:
+        logger.info("{}: {} ({})", name, t.reason, t.new_state.value)
+
+
+class Throttle:
+    """All lanes of one client. Available as ``client.throttle`` on :class:`rr.AsyncOpenAI`.
+
+    To enforce one rate across several clients, e.g. one per thread with its own
+    ``asyncio.run``, create one ``Throttle`` and pass it as ``throttle=`` to each client.
+    """
+
+    def __init__(
+        self, config: ThrottleConfig | None = None, clock: Callable[[], float] = time.monotonic
+    ) -> None:
+        self.config = config or ThrottleConfig()
+        self._clock = clock
+        self._rng = random.Random()  # seeded from OS entropy: jitter differs between clients
+        self._lanes: dict[tuple[str, str | None], Lane] = {}
+        self._lock = threading.Lock()
+
+    def lane(self, path: str, model: str | None) -> Lane:
+        key = (path, model)
+        with self._lock:
+            lane = self._lanes.get(key)
+            if lane is None:
+                observed = self.config.adaptive and any(
+                    path.endswith(p) for p in self.config.observe_paths
+                )
+                name = model or path
+                lane = Lane(name, self.config, observed=observed, clock=self._clock, rng=self._rng)
+                self._lanes[key] = lane
+            return lane
+
+    def lanes(self) -> list[Lane]:
+        with self._lock:
+            return list(self._lanes.values())
+
+    def find(self, model: str) -> list[Lane]:
+        with self._lock:
+            return [lane for (_, m), lane in self._lanes.items() if m == model]
+
+    def stats(self) -> dict[str, dict[str, Any]]:
+        """Current state, rate, load and totals per lane (keyed by model, or path)."""
+        with self._lock:
+            items = list(self._lanes.items())
+        out: dict[str, dict[str, Any]] = {}
+        for (path, model), lane in items:
+            key = model if model is not None else path
+            if key in out:
+                key = f"{model} {path}"
+            out[key] = lane.snapshot()
+        return out
+
+    def reset_baseline(self, model: str | None = None) -> None:
+        """Forget learned baselines (of one model, or all)."""
+        for lane in self.lanes():
+            if model is None or lane.name == model:
+                lane.estimator.reset()