"""Run-local, exact-key cache for expensive thermodynamic calculations. The cache is deliberately bound to one simulation through ``ContextVar``. That keeps concurrent runs isolated and releases all cached states when the run finishes. Keys use the original Python values with no rounding or tolerance-based reuse that could flatten numerical residuals seen by ODE and nonlinear solvers. """ from __future__ import annotations from collections.abc import Callable, Generator from contextlib import contextmanager from contextvars import ContextVar from dataclasses import dataclass from functools import lru_cache, wraps import os from typing import ParamSpec, TypeVar from app.simulation.performance import PROFILE_MODE, record_property_cache _P = ParamSpec("_P") _R = TypeVar("_R") DEFAULT_PROPERTY_CACHE_MAX_ENTRIES = 8192 def _read_cache_enabled() -> bool: raw_value = os.getenv("SIMULATIONAPP_PROPERTY_CACHE", "on").strip().lower() if raw_value in {"", "1", "true", "yes", "on"}: return True if raw_value in {"0", "false", "no", "off"}: return False raise ValueError( "SIMULATIONAPP_PROPERTY_CACHE must be one of: on, off, true, false, 1, 0." ) PROPERTY_CACHE_ENABLED = _read_cache_enabled() @dataclass(frozen=True) class PropertyCacheInfo: hits: int misses: int max_entries_per_cache: int cache_count: int current_entries: int evictions: int class SimulationPropertyCache: """Bounded C-level LRUs owned by one simulation run.""" def __init__(self, max_entries: int = DEFAULT_PROPERTY_CACHE_MAX_ENTRIES) -> None: if max_entries <= 0: raise ValueError("Property cache max_entries must be positive.") self.max_entries_per_cache = int(max_entries) self._functions: dict[ tuple[str, int, Callable[..., object]], Callable[..., object], ] = {} self._failed_misses: dict[ tuple[str, int, Callable[..., object]], int, ] = {} self._owners: dict[int, tuple[object, int]] = {} self._next_owner_token = 0 def owner_token(self, owner: object) -> int: """Return a stable identity token and retain its owner for this run.""" identity = id(owner) existing = self._owners.get(identity) if existing is not None and existing[0] is owner: return existing[1] self._next_owner_token += 1 self._owners[identity] = (owner, self._next_owner_token) return self._next_owner_token def get_or_compute( self, operation: str, owner: object, function: Callable[..., _R], args: tuple[object, ...], kwargs: dict[str, object], ) -> _R: cache_key = (operation, self.owner_token(owner), function) cached_function = self._functions.get(cache_key) if cached_function is None: @lru_cache(maxsize=self.max_entries_per_cache, typed=True) def invoke(*cached_args: object, **cached_kwargs: object) -> _R: return function(owner, *cached_args, **cached_kwargs) cached_function = invoke self._functions[cache_key] = cached_function if PROFILE_MODE != "audit": try: return cached_function(*args, **kwargs) except Exception: self._failed_misses[cache_key] = ( self._failed_misses.get(cache_key, 0) + 1 ) raise before = cached_function.cache_info() # type: ignore[attr-defined] try: value = cached_function(*args, **kwargs) except Exception: self._failed_misses[cache_key] = ( self._failed_misses.get(cache_key, 0) + 1 ) raise finally: after = cached_function.cache_info() # type: ignore[attr-defined] hit = after.hits > before.hits record_property_cache(operation, hit=hit) return value def info(self) -> PropertyCacheInfo: cache_infos = { key: cached.cache_info() # type: ignore[attr-defined] for key, cached in self._functions.items() } return PropertyCacheInfo( hits=sum(info.hits for info in cache_infos.values()), misses=sum(info.misses for info in cache_infos.values()), max_entries_per_cache=self.max_entries_per_cache, cache_count=len(self._functions), current_entries=sum(info.currsize for info in cache_infos.values()), evictions=sum( max( 0, info.misses - self._failed_misses.get(key, 0) - info.currsize, ) for key, info in cache_infos.items() ), ) _CURRENT_PROPERTY_CACHE: ContextVar[SimulationPropertyCache | None] = ContextVar( "simulation_property_cache", default=None, ) def current_property_cache() -> SimulationPropertyCache | None: return _CURRENT_PROPERTY_CACHE.get() @contextmanager def property_cache_run( *, max_entries: int = DEFAULT_PROPERTY_CACHE_MAX_ENTRIES, ) -> Generator[SimulationPropertyCache | None, None, None]: """Bind a fresh cache to one top-level simulation run. Nested uses reuse the existing cache so lower-level simulation helpers can safely opt in without replacing the cache created by the API entry point. """ existing = _CURRENT_PROPERTY_CACHE.get() if existing is not None: yield existing return if not PROPERTY_CACHE_ENABLED: yield None return cache = SimulationPropertyCache(max_entries=max_entries) token = _CURRENT_PROPERTY_CACHE.set(cache) try: yield cache finally: _CURRENT_PROPERTY_CACHE.reset(token) def cache_property_calculation( operation: str, ) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]: """Cache one pure property calculation with hashable arguments per run.""" def decorate(function: Callable[_P, _R]) -> Callable[_P, _R]: if not PROPERTY_CACHE_ENABLED: return function @wraps(function) def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R: cache = _CURRENT_PROPERTY_CACHE.get() if cache is None: return function(*args, **kwargs) owner = args[0] if args else function return cache.get_or_compute( operation, owner, function, tuple(args[1:] if args else ()), dict(kwargs), ) return wrapper return decorate def with_property_cache(function: Callable[_P, _R]) -> Callable[_P, _R]: """Ensure a simulation entry point has a run-local cache.""" if not PROPERTY_CACHE_ENABLED: return function @wraps(function) def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R: with property_cache_run(): return function(*args, **kwargs) return wrapper __all__ = [ "DEFAULT_PROPERTY_CACHE_MAX_ENTRIES", "PROPERTY_CACHE_ENABLED", "PropertyCacheInfo", "SimulationPropertyCache", "cache_property_calculation", "current_property_cache", "property_cache_run", "with_property_cache", ]