Cache (1/5): Add content-addressed local inference caches #121
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,341 @@ | ||
| """Content-addressed local SQLite caches.""" | ||
|
|
||
| from __future__ import annotations | ||
|
|
||
| import hashlib | ||
| import json | ||
| import os | ||
| import sqlite3 | ||
| import uuid | ||
| from datetime import UTC, datetime | ||
| from pathlib import Path | ||
| from typing import Any, Literal | ||
| from urllib.parse import quote | ||
|
|
||
| import pandas as pd | ||
|
|
||
| COMPLETION_DB_NAME = "completions.db" | ||
| JUDGEMENT_DB_NAME = "judgements.db" | ||
| DESCRIPTOR_FILENAME = "metadata.json" | ||
|
|
||
| CacheKind = Literal["completions", "judgements"] | ||
|
|
||
|
|
||
| def stable_json_dumps(value: Any) -> str: | ||
| return json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False) | ||
|
|
||
|
|
||
| def descriptor_hash(descriptor: dict[str, Any]) -> str: | ||
| return hashlib.sha256(stable_json_dumps(descriptor).encode()).hexdigest()[:16] | ||
|
|
||
|
|
||
| def input_hash(input_text: str) -> str: | ||
| return hashlib.sha256(input_text.encode()).hexdigest() | ||
|
|
||
|
|
||
| def cache_folder( | ||
| store_root: Path | str, | ||
| kind: CacheKind, | ||
| task: str, | ||
| model_spec: str, | ||
| descriptor: dict[str, Any], | ||
| ) -> Path: | ||
| return cache_model_folder(store_root, kind, task, model_spec) / descriptor_hash( | ||
| descriptor | ||
| ) | ||
|
|
||
|
|
||
| def cache_model_folder( | ||
| store_root: Path | str, | ||
| kind: CacheKind, | ||
| task: str, | ||
| model_spec: str, | ||
| ) -> Path: | ||
| provider, model = model_spec.split("/", 1) | ||
| return ( | ||
| Path(store_root) | ||
| / kind | ||
| / quote(task, safe="") | ||
| / quote(provider, safe="") | ||
| / quote(model, safe="") | ||
| ) | ||
|
|
||
|
|
||
| def write_descriptor(folder: Path, descriptor: dict[str, Any]) -> Path: | ||
| folder.mkdir(parents=True, exist_ok=True) | ||
| path = folder / DESCRIPTOR_FILENAME | ||
| if path.exists(): | ||
| if json.loads(path.read_text()) != descriptor: | ||
| raise ValueError(f"Descriptor does not match existing metadata at {path}.") | ||
| return path | ||
|
|
||
| path.write_text(json.dumps(descriptor, indent=2, sort_keys=True) + "\n") | ||
| return path | ||
|
|
||
|
|
||
| def read_descriptor(folder: Path) -> dict[str, Any]: | ||
| descriptor = json.loads((folder / DESCRIPTOR_FILENAME).read_text()) | ||
| if folder.name != descriptor_hash(descriptor): | ||
| raise ValueError(f"Descriptor hash does not match cache folder {folder}.") | ||
| return descriptor | ||
|
|
||
|
|
||
| class _SQLiteCache: | ||
| table: str | ||
| schema: str | ||
|
|
||
| def __init__(self, db_path: Path | str) -> None: | ||
| self.db_path = Path(db_path) | ||
| self._conn: sqlite3.Connection | None = None | ||
|
|
||
| def _connect(self) -> sqlite3.Connection: | ||
| if self._conn is None: | ||
| self.db_path.parent.mkdir(parents=True, exist_ok=True) | ||
| self._conn = sqlite3.connect(self.db_path) | ||
| self._conn.execute(self.schema) | ||
| self._conn.commit() | ||
| return self._conn | ||
|
|
||
| def _query( | ||
| self, | ||
| input_hashes: list[str] | None, | ||
| conditions: list[str], | ||
| params: list[Any], | ||
| ) -> pd.DataFrame: | ||
| if input_hashes is not None: | ||
| if not input_hashes: | ||
| return pd.read_sql( | ||
| f"SELECT * FROM {self.table} WHERE 0", self._connect() | ||
| ) | ||
| placeholders = ",".join("?" * len(input_hashes)) | ||
| conditions.append(f"input_hash IN ({placeholders})") | ||
| params.extend(input_hashes) | ||
|
|
||
| where = f" WHERE {' AND '.join(conditions)}" if conditions else "" | ||
| return pd.read_sql( | ||
| f"SELECT * FROM {self.table}{where} ORDER BY instruction_id", | ||
| self._connect(), | ||
| params=params, | ||
| ) | ||
|
|
||
| def _delete(self, conditions: list[str], params: list[Any]) -> int: | ||
| if not conditions: | ||
| raise ValueError("Delete requires at least one filter.") | ||
| where = f" WHERE {' AND '.join(conditions)}" if conditions else "" | ||
| with self._connect() as conn: | ||
| cursor = conn.execute(f"DELETE FROM {self.table}{where}", params) | ||
| return cursor.rowcount | ||
|
|
||
| def merge_from(self, other_db: Path) -> int: | ||
| """Atomically merge another cache database using pushed_at last-write-wins.""" | ||
| self.close() | ||
| frames = [] | ||
| for db_path in (other_db, self.db_path): | ||
| if db_path.exists(): | ||
| with sqlite3.connect(db_path) as conn: | ||
| frames.append(pd.read_sql(f"SELECT * FROM {self.table}", conn)) | ||
| rows = ( | ||
| pd.concat(frames, ignore_index=True) | ||
| .sort_values("pushed_at", kind="stable") | ||
| .drop_duplicates("input_hash", keep="last") | ||
| ) | ||
|
|
||
| self.db_path.parent.mkdir(parents=True, exist_ok=True) | ||
| temporary_db = self.db_path.with_name(f".{self.db_path.name}.{uuid.uuid4()}") | ||
| try: | ||
| with sqlite3.connect(temporary_db) as conn: | ||
| conn.execute(self.schema) | ||
| rows.to_sql(self.table, conn, if_exists="append", index=False) | ||
| os.replace(temporary_db, self.db_path) | ||
|
Member
There was a problem hiding this comment. We might need some handling and protection for the cache. Assuming we can start two jobs at the same time, while reading it from frames.append(pd.read_sql(f"SELECT * FROM {self.table}", conn))if another process appends something to this table while os.replace(temporary_db, self.db_path)executed, than information will be lost. Also pd.read_sql(...)
pd.concat(...)
.sort_values(...)
.drop_duplicates(...)
rows.to_sql(...)these operations load both databases and rewrites every row. If we have a huge cache this is expensive. We can use ATTACH DATABASE ? AS incoming;
INSERT INTO completions (...)
SELECT ...
FROM incoming.completions
ON CONFLICT(descriptor_id, input_hash) DO UPDATE ...;However this is not required and current is also enough (unless we use VERY large cache) |
||
| finally: | ||
| temporary_db.unlink(missing_ok=True) | ||
| return len(rows) | ||
|
|
||
| def close(self) -> None: | ||
| if self._conn is not None: | ||
| self._conn.close() | ||
| self._conn = None | ||
|
|
||
| def __enter__(self) -> _SQLiteCache: | ||
| return self | ||
|
|
||
| def __exit__(self, *_: object) -> None: | ||
| self.close() | ||
|
|
||
|
|
||
| class CompletionCache(_SQLiteCache): | ||
| """Completion rows keyed by the exact rendered model input.""" | ||
|
|
||
| table = "completions" | ||
| schema = """ | ||
| CREATE TABLE IF NOT EXISTS completions ( | ||
| input_hash TEXT PRIMARY KEY, | ||
| input_text TEXT NOT NULL, | ||
| completion TEXT NOT NULL, | ||
| benchmark TEXT NOT NULL, | ||
| instruction_id TEXT NOT NULL, | ||
| model TEXT NOT NULL, | ||
| pushed_by TEXT NOT NULL, | ||
| pushed_at TEXT NOT NULL, | ||
| run_id TEXT NOT NULL | ||
| ) | ||
| """ | ||
|
|
||
| def save( | ||
| self, | ||
| rows: pd.DataFrame, | ||
| *, | ||
| pushed_by: str, | ||
| run_id: str | None = None, | ||
| ) -> int: | ||
| now = datetime.now(UTC).isoformat() | ||
| resolved_run_id = run_id or str(uuid.uuid4()) | ||
| values = [ | ||
| ( | ||
| input_hash(str(row["input_text"])), | ||
| str(row["input_text"]), | ||
| str(row["completion"]), | ||
| str(row["benchmark"]), | ||
| str(row["instruction_id"]), | ||
| str(row["model"]), | ||
| pushed_by, | ||
| now, | ||
| resolved_run_id, | ||
| ) | ||
| for _, row in rows.iterrows() | ||
| ] | ||
| with self._connect() as conn: | ||
| conn.executemany( | ||
| "INSERT OR REPLACE INTO completions VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)", | ||
| values, | ||
| ) | ||
| return len(values) | ||
|
|
||
| def query( | ||
| self, | ||
| input_hashes: list[str] | None = None, | ||
| *, | ||
| instruction_id: str | None = None, | ||
| model: str | None = None, | ||
| ) -> pd.DataFrame: | ||
| conditions: list[str] = [] | ||
| params: list[Any] = [] | ||
| if instruction_id is not None: | ||
| conditions.append("instruction_id = ?") | ||
| params.append(str(instruction_id)) | ||
| if model is not None: | ||
| conditions.append("model = ?") | ||
| params.append(model) | ||
| return self._query(input_hashes, conditions, params) | ||
|
|
||
| def delete( | ||
| self, | ||
| *, | ||
| instruction_id: str | None = None, | ||
| model: str | None = None, | ||
| ) -> int: | ||
| conditions: list[str] = [] | ||
| params: list[Any] = [] | ||
| if instruction_id is not None: | ||
| conditions.append("instruction_id = ?") | ||
| params.append(str(instruction_id)) | ||
| if model is not None: | ||
| conditions.append("model = ?") | ||
| params.append(model) | ||
| return self._delete(conditions, params) | ||
|
|
||
|
|
||
| class JudgementCache(_SQLiteCache): | ||
| """Raw judge completions keyed by the exact rendered judge input.""" | ||
|
|
||
| table = "judgements" | ||
| schema = """ | ||
| CREATE TABLE IF NOT EXISTS judgements ( | ||
| input_hash TEXT PRIMARY KEY, | ||
| judge_input TEXT NOT NULL, | ||
| judge_completion TEXT NOT NULL, | ||
| benchmark TEXT NOT NULL, | ||
| instruction_id TEXT NOT NULL, | ||
| model_a TEXT NOT NULL, | ||
| model_b TEXT NOT NULL, | ||
| judge TEXT NOT NULL, | ||
| top_logprobs TEXT, | ||
| -- direct/reversed relative to the source model order, when applicable | ||
| orientation TEXT, | ||
| pushed_by TEXT NOT NULL, | ||
| pushed_at TEXT NOT NULL, | ||
| run_id TEXT NOT NULL | ||
| ) | ||
| """ | ||
|
|
||
| def save( | ||
| self, | ||
| rows: pd.DataFrame, | ||
| *, | ||
| pushed_by: str, | ||
| run_id: str | None = None, | ||
| ) -> int: | ||
| now = datetime.now(UTC).isoformat() | ||
| resolved_run_id = run_id or str(uuid.uuid4()) | ||
| values = [ | ||
| ( | ||
| input_hash(str(row["judge_input"])), | ||
| str(row["judge_input"]), | ||
| str(row["judge_completion"]), | ||
| str(row["benchmark"]), | ||
| str(row["instruction_id"]), | ||
| str(row["model_a"]), | ||
| str(row["model_b"]), | ||
|
Member
There was a problem hiding this comment. Note: For sample-wise this can be |
||
| str(row["judge"]), | ||
| ( | ||
| stable_json_dumps(row["top_logprobs"]) | ||
| if row.get("top_logprobs") is not None | ||
| else None | ||
| ), | ||
| row.get("orientation"), | ||
| pushed_by, | ||
| now, | ||
| resolved_run_id, | ||
| ) | ||
| for _, row in rows.iterrows() | ||
| ] | ||
| with self._connect() as conn: | ||
| conn.executemany( | ||
| "INSERT OR REPLACE INTO judgements " | ||
| "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", | ||
| values, | ||
| ) | ||
| return len(values) | ||
|
|
||
| def query( | ||
| self, | ||
| input_hashes: list[str] | None = None, | ||
| *, | ||
| instruction_id: str | None = None, | ||
| model: str | None = None, | ||
| ) -> pd.DataFrame: | ||
| conditions: list[str] = [] | ||
| params: list[Any] = [] | ||
| if instruction_id is not None: | ||
| conditions.append("instruction_id = ?") | ||
| params.append(str(instruction_id)) | ||
| if model is not None: | ||
| conditions.append("(model_a = ? OR model_b = ?)") | ||
| params.extend((model, model)) | ||
| return self._query(input_hashes, conditions, params) | ||
|
|
||
| def delete( | ||
| self, | ||
| *, | ||
| instruction_id: str | None = None, | ||
| model: str | None = None, | ||
| ) -> int: | ||
| conditions: list[str] = [] | ||
| params: list[Any] = [] | ||
| if instruction_id is not None: | ||
| conditions.append("instruction_id = ?") | ||
| params.append(str(instruction_id)) | ||
| if model is not None: | ||
| conditions.append("(model_a = ? OR model_b = ?)") | ||
| params.extend((model, model)) | ||
| return self._delete(conditions, params) | ||
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Why there are two checks here for
input_hashes?