Spaces:
Sleeping
Sleeping
File size: 4,674 Bytes
8a169a0 cfac7e4 8a169a0 cfac7e4 8a169a0 2db5fdd 8a169a0 2db5fdd 8a169a0 2db5fdd 8a169a0 2db5fdd 8a169a0 2db5fdd cfac7e4 2db5fdd 8a169a0 0bbc8ce 8a169a0 | 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 | """SQLite-backed cache for simulation sweep results, keyed on SHA256 of params."""
from __future__ import annotations
import hashlib
import json
import logging
import sqlite3
from typing import Any, Dict, Optional, Tuple
logger = logging.getLogger(__name__)
SCHEMA_VERSION = 3
class SweepCache:
"""Stores simulation scalar results in a local SQLite database.
Keys are derived from a SHA256 hash of the canonicalised parameter dict,
so identical runs are served from cache instead of re-simulated.
Schema is versioned via ``PRAGMA user_version``. On open, if the stored
version differs from ``SCHEMA_VERSION``, all rows are dropped and the
version is written.
"""
def __init__(self, db_path: str = "sweep_cache.db") -> None:
# check_same_thread=False: Gradio serves requests from a worker pool,
# so the same SweepCache instance is used across threads. Access is
# serialized at the Python level by SQLite's own locking, so cross-
# thread use is safe in practice.
self._conn = sqlite3.connect(db_path, check_same_thread=False)
self._conn.execute("PRAGMA journal_mode=WAL")
self._conn.execute(
"CREATE TABLE IF NOT EXISTS cache ("
" key TEXT PRIMARY KEY,"
" scalars TEXT,"
" wall_time REAL,"
" ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP"
")"
)
self._migrate_if_needed()
self._conn.commit()
self._session_hits = 0
self._session_misses = 0
def _migrate_if_needed(self) -> None:
"""If stored schema version != SCHEMA_VERSION, drop all rows."""
cur = self._conn.execute("PRAGMA user_version")
current = cur.fetchone()[0]
if current != SCHEMA_VERSION:
logger.warning(
"sweep_cache schema mismatch (stored=%s, expected=%s); "
"wiping cache",
current, SCHEMA_VERSION,
)
self._conn.execute("DELETE FROM cache")
self._conn.execute(f"PRAGMA user_version = {SCHEMA_VERSION}")
# ------------------------------------------------------------------
# Internal helpers
# ------------------------------------------------------------------
@staticmethod
def _make_key(params: Dict[str, Any]) -> str:
"""Canonical SHA256 key: round floats to 8 dp, sort keys, hash."""
cleaned: Dict[str, Any] = {}
for k in sorted(params.keys()):
v = params[k]
cleaned[k] = round(v, 8) if isinstance(v, float) else v
blob = json.dumps(cleaned, sort_keys=True).encode()
return hashlib.sha256(blob).hexdigest()[:16]
# ------------------------------------------------------------------
# Public API
# ------------------------------------------------------------------
def get(self, params: Dict[str, Any]) -> Optional[Tuple[Dict[str, Any], float]]:
"""Return ``(scalars_dict, wall_time)`` or ``None`` on cache miss."""
key = self._make_key(params)
row = self._conn.execute(
"SELECT scalars, wall_time FROM cache WHERE key = ?", (key,)
).fetchone()
if row is None:
self._session_misses += 1
return None
self._session_hits += 1
return json.loads(row[0]), row[1]
def put(
self,
params: Dict[str, Any],
scalars: Dict[str, Any],
wall_time: float,
) -> None:
"""Insert (or replace) a cache entry."""
key = self._make_key(params)
self._conn.execute(
"INSERT OR REPLACE INTO cache (key, scalars, wall_time) VALUES (?, ?, ?)",
(key, json.dumps(scalars), wall_time),
)
self._conn.commit()
def clear(self) -> int:
"""Delete all cached entries. Returns the number of rows deleted."""
count = self._conn.execute("SELECT COUNT(*) FROM cache").fetchone()[0]
self._conn.execute("DELETE FROM cache")
self._conn.commit()
self._session_hits = 0
self._session_misses = 0
return count
@property
def stats(self) -> Dict[str, Any]:
"""Return cache statistics for the current session."""
total = self._conn.execute("SELECT COUNT(*) FROM cache").fetchone()[0]
total_lookups = self._session_hits + self._session_misses
return {
"total_cached": total,
"session_hits": self._session_hits,
"session_misses": self._session_misses,
"hit_rate": (
self._session_hits / total_lookups if total_lookups > 0 else 0.0
),
}
|