tnfr 4.5.2__py3-none-any.whl → 8.5.0__py3-none-any.whl

tnfr/utils/cache.pyi

@@ -0,0 +1,470 @@
+from __future__ import annotations
+
+import logging
+import threading
+from collections import defaultdict
+from collections.abc import (
+    Callable,
+    Hashable,
+    Iterable,
+    Iterator,
+    Mapping,
+    MutableMapping,
+)
+from dataclasses import dataclass
+from typing import Any, ClassVar, ContextManager, Generic, TypeVar
+
+import networkx as nx
+from cachetools import LRUCache
+
+from ..types import GraphLike, NodeId, TimingContext, TNFRGraph
+
+K = TypeVar("K", bound=Hashable)
+V = TypeVar("V")
+T = TypeVar("T")
+
+class SecurityError(RuntimeError):
+    """Raised when a cache payload fails hardened validation."""
+
+    ...
+
+@dataclass(frozen=True)
+class CacheCapacityConfig:
+    default_capacity: int | None
+    overrides: dict[str, int | None]
+
+@dataclass(frozen=True)
+class CacheStatistics:
+    hits: int = ...
+    misses: int = ...
+    evictions: int = ...
+    total_time: float = ...
+    timings: int = ...
+
+    def merge(self, other: CacheStatistics) -> CacheStatistics: ...
+
+class CacheLayer:
+    def load(self, name: str) -> Any: ...
+    def store(self, name: str, value: Any) -> None: ...
+    def delete(self, name: str) -> None: ...
+    def clear(self) -> None: ...
+    def close(self) -> None: ...
+
+class MappingCacheLayer(CacheLayer):
+    def __init__(self, storage: MutableMapping[str, Any] | None = ...) -> None: ...
+
+class ShelveCacheLayer(CacheLayer):
+    def __init__(
+        self,
+        path: str,
+        *,
+        flag: str = ...,
+        protocol: int | None = ...,
+        writeback: bool = ...,
+    ) -> None: ...
+
+class RedisCacheLayer(CacheLayer):
+    def __init__(self, client: Any | None = ..., *, namespace: str = ...) -> None: ...
+
+class CacheManager:
+    _MISSING: ClassVar[object]
+
+    def __init__(
+        self,
+        storage: MutableMapping[str, Any] | None = ...,
+        *,
+        default_capacity: int | None = ...,
+        overrides: Mapping[str, int | None] | None = ...,
+        layers: Iterable[CacheLayer] | None = ...,
+    ) -> None: ...
+    @staticmethod
+    def _normalise_capacity(value: int | None) -> int | None: ...
+    def register(
+        self,
+        name: str,
+        factory: Callable[[], Any],
+        *,
+        lock_factory: Callable[[], threading.Lock | threading.RLock] | None = ...,
+        reset: Callable[[Any], Any] | None = ...,
+        create: bool = ...,
+        encoder: Callable[[Any], Any] | None = ...,
+        decoder: Callable[[Any], Any] | None = ...,
+    ) -> None: ...
+    def configure(
+        self,
+        *,
+        default_capacity: int | None | object = ...,
+        overrides: Mapping[str, int | None] | None = ...,
+        replace_overrides: bool = ...,
+    ) -> None: ...
+    def configure_from_mapping(self, config: Mapping[str, Any]) -> None: ...
+    def export_config(self) -> CacheCapacityConfig: ...
+    def get_capacity(
+        self,
+        name: str,
+        *,
+        requested: int | None = ...,
+        fallback: int | None = ...,
+        use_default: bool = ...,
+    ) -> int | None: ...
+    def has_override(self, name: str) -> bool: ...
+    def get_lock(self, name: str) -> threading.Lock | threading.RLock: ...
+    def names(self) -> Iterator[str]: ...
+    def get(self, name: str, *, create: bool = ...) -> Any: ...
+    def peek(self, name: str) -> Any: ...
+    def store(self, name: str, value: Any) -> None: ...
+    def update(
+        self,
+        name: str,
+        updater: Callable[[Any], Any],
+        *,
+        create: bool = ...,
+    ) -> Any: ...
+    def clear(self, name: str | None = ...) -> None: ...
+    def increment_hit(
+        self,
+        name: str,
+        *,
+        amount: int = ...,
+        duration: float | None = ...,
+    ) -> None: ...
+    def increment_miss(
+        self,
+        name: str,
+        *,
+        amount: int = ...,
+        duration: float | None = ...,
+    ) -> None: ...
+    def increment_eviction(self, name: str, *, amount: int = ...) -> None: ...
+    def record_timing(self, name: str, duration: float) -> None: ...
+    def timer(self, name: str) -> TimingContext: ...
+    def get_metrics(self, name: str) -> CacheStatistics: ...
+    def iter_metrics(self) -> Iterator[tuple[str, CacheStatistics]]: ...
+    def aggregate_metrics(self) -> CacheStatistics: ...
+    def register_metrics_publisher(
+        self, publisher: Callable[[str, CacheStatistics], None]
+    ) -> None: ...
+    def publish_metrics(
+        self,
+        *,
+        publisher: Callable[[str, CacheStatistics], None] | None = ...,
+    ) -> None: ...
+    def log_metrics(self, logger: logging.Logger, *, level: int = ...) -> None: ...
+
+class InstrumentedLRUCache(MutableMapping[K, V], Generic[K, V]):
+    _MISSING: ClassVar[object]
+
+    def __init__(
+        self,
+        maxsize: int,
+        *,
+        manager: CacheManager | None = ...,
+        metrics_key: str | None = ...,
+        telemetry_callbacks: (
+            Iterable[Callable[[K, V], None]] | Callable[[K, V], None] | None
+        ) = ...,
+        eviction_callbacks: (
+            Iterable[Callable[[K, V], None]] | Callable[[K, V], None] | None
+        ) = ...,
+        locks: MutableMapping[K, Any] | None = ...,
+        getsizeof: Callable[[V], int] | None = ...,
+        count_overwrite_hit: bool = ...,
+    ) -> None: ...
+    @property
+    def telemetry_callbacks(self) -> tuple[Callable[[K, V], None], ...]: ...
+    @property
+    def eviction_callbacks(self) -> tuple[Callable[[K, V], None], ...]: ...
+    def set_telemetry_callbacks(
+        self,
+        callbacks: Iterable[Callable[[K, V], None]] | Callable[[K, V], None] | None,
+        *,
+        append: bool = ...,
+    ) -> None: ...
+    def set_eviction_callbacks(
+        self,
+        callbacks: Iterable[Callable[[K, V], None]] | Callable[[K, V], None] | None,
+        *,
+        append: bool = ...,
+    ) -> None: ...
+    def pop(self, key: K, default: Any = ...) -> V: ...
+    def popitem(self) -> tuple[K, V]: ...
+    def clear(self) -> None: ...
+    @property
+    def maxsize(self) -> int: ...
+    @property
+    def currsize(self) -> int: ...
+    def get(self, key: K, default: V | None = ...) -> V | None: ...
+
+class ManagedLRUCache(LRUCache[K, V], Generic[K, V]):
+    def __init__(
+        self,
+        maxsize: int,
+        *,
+        manager: CacheManager | None = ...,
+        metrics_key: str | None = ...,
+        eviction_callbacks: (
+            Iterable[Callable[[K, V], None]] | Callable[[K, V], None] | None
+        ) = ...,
+        telemetry_callbacks: (
+            Iterable[Callable[[K, V], None]] | Callable[[K, V], None] | None
+        ) = ...,
+        locks: MutableMapping[K, Any] | None = ...,
+    ) -> None: ...
+    def popitem(self) -> tuple[K, V]: ...
+
+def prune_lock_mapping(
+    cache: Mapping[K, Any] | MutableMapping[K, Any] | None,
+    locks: MutableMapping[K, Any] | None,
+) -> None: ...
+
+__all__ = (
+    "CacheLayer",
+    "CacheManager",
+    "CacheCapacityConfig",
+    "CacheStatistics",
+    "InstrumentedLRUCache",
+    "ManagedLRUCache",
+    "MappingCacheLayer",
+    "RedisCacheLayer",
+    "ShelveCacheLayer",
+    "prune_lock_mapping",
+    "EdgeCacheManager",
+    "NODE_SET_CHECKSUM_KEY",
+    "cached_node_list",
+    "cached_nodes_and_A",
+    "clear_node_repr_cache",
+    "edge_version_cache",
+    "edge_version_update",
+    "ensure_node_index_map",
+    "ensure_node_offset_map",
+    "get_graph_version",
+    "increment_edge_version",
+    "increment_graph_version",
+    "node_set_checksum",
+    "stable_json",
+    "configure_graph_cache_limits",
+    "DNFR_PREP_STATE_KEY",
+    "DnfrCache",
+    "DnfrPrepState",
+    "new_dnfr_cache",
+    "build_cache_manager",
+    "configure_global_cache_layers",
+    "reset_global_cache_manager",
+    "_GRAPH_CACHE_LAYERS_KEY",
+    "_SeedHashCache",
+    "ScopedCounterCache",
+)
+
+NODE_SET_CHECKSUM_KEY: str
+_GRAPH_CACHE_LAYERS_KEY: str
+DNFR_PREP_STATE_KEY: str
+
+class DnfrCache:
+    idx: dict[Any, int]
+    theta: list[float]
+    epi: list[float]
+    vf: list[float]
+    cos_theta: list[float]
+    sin_theta: list[float]
+    neighbor_x: list[float]
+    neighbor_y: list[float]
+    neighbor_epi_sum: list[float]
+    neighbor_vf_sum: list[float]
+    neighbor_count: list[float]
+    neighbor_deg_sum: list[float] | None
+    th_bar: list[float] | None
+    epi_bar: list[float] | None
+    vf_bar: list[float] | None
+    deg_bar: list[float] | None
+    degs: dict[Any, float] | None
+    deg_list: list[float] | None
+    theta_np: Any | None
+    epi_np: Any | None
+    vf_np: Any | None
+    cos_theta_np: Any | None
+    sin_theta_np: Any | None
+    deg_array: Any | None
+    edge_src: Any | None
+    edge_dst: Any | None
+    checksum: Any | None
+    neighbor_x_np: Any | None
+    neighbor_y_np: Any | None
+    neighbor_epi_sum_np: Any | None
+    neighbor_vf_sum_np: Any | None
+    neighbor_count_np: Any | None
+    neighbor_deg_sum_np: Any | None
+    th_bar_np: Any | None
+    epi_bar_np: Any | None
+    vf_bar_np: Any | None
+    deg_bar_np: Any | None
+    grad_phase_np: Any | None
+    grad_epi_np: Any | None
+    grad_vf_np: Any | None
+    grad_topo_np: Any | None
+    grad_total_np: Any | None
+    dense_components_np: Any | None
+    dense_accum_np: Any | None
+    dense_degree_np: Any | None
+    neighbor_accum_np: Any | None
+    neighbor_inv_count_np: Any | None
+    neighbor_cos_avg_np: Any | None
+    neighbor_sin_avg_np: Any | None
+    neighbor_mean_tmp_np: Any | None
+    neighbor_mean_length_np: Any | None
+    edge_signature: Any | None
+    neighbor_accum_signature: Any | None
+    neighbor_edge_values_np: Any | None
+
+class EdgeCacheState:
+    cache: MutableMapping[Hashable, Any]
+    locks: defaultdict[Hashable, threading.RLock]
+    max_entries: int | None
+    dirty: bool
+
+def new_dnfr_cache() -> DnfrCache: ...
+
+class DnfrPrepState:
+    cache: DnfrCache
+    cache_lock: threading.RLock
+    vector_lock: threading.RLock
+
+class EdgeCacheManager:
+    _STATE_KEY: str
+
+    def __init__(self, graph: MutableMapping[str, Any]) -> None: ...
+    def record_hit(self) -> None: ...
+    def record_miss(self, *, track_metrics: bool = ...) -> None: ...
+    def record_eviction(self, *, track_metrics: bool = ...) -> None: ...
+    def timer(self) -> TimingContext: ...
+    def _default_state(self) -> EdgeCacheState: ...
+    def resolve_max_entries(self, max_entries: int | None | object) -> int | None: ...
+    def _build_state(self, max_entries: int | None) -> EdgeCacheState: ...
+    def _ensure_state(
+        self, state: EdgeCacheState | None, max_entries: int | None | object
+    ) -> EdgeCacheState: ...
+    def _reset_state(self, state: EdgeCacheState | None) -> EdgeCacheState: ...
+    def get_cache(
+        self,
+        max_entries: int | None | object,
+        *,
+        create: bool = ...,
+    ) -> EdgeCacheState | None: ...
+    def flush_state(self, state: EdgeCacheState) -> None: ...
+    def clear(self) -> None: ...
+
+def get_graph_version(graph: Any, key: str, default: int = ...) -> int: ...
+def increment_graph_version(graph: Any, key: str) -> int: ...
+def stable_json(obj: Any) -> str: ...
+def clear_node_repr_cache() -> None: ...
+def configure_global_cache_layers(
+    *,
+    shelve: Mapping[str, Any] | None = ...,
+    redis: Mapping[str, Any] | None = ...,
+    replace: bool = ...,
+) -> None: ...
+def node_set_checksum(
+    G: nx.Graph,
+    nodes: Iterable[Any] | None = ...,
+    *,
+    presorted: bool = ...,
+    store: bool = ...,
+) -> str: ...
+def reset_global_cache_manager() -> None: ...
+def build_cache_manager(
+    *,
+    graph: MutableMapping[str, Any] | None = ...,
+    storage: MutableMapping[str, Any] | None = ...,
+    default_capacity: int | None = ...,
+    overrides: Mapping[str, int | None] | None = ...,
+) -> CacheManager: ...
+def cached_node_list(G: nx.Graph) -> tuple[Any, ...]: ...
+def ensure_node_index_map(G: TNFRGraph) -> dict[NodeId, int]: ...
+def ensure_node_offset_map(G: TNFRGraph) -> dict[NodeId, int]: ...
+def configure_graph_cache_limits(
+    G: GraphLike | TNFRGraph | MutableMapping[str, Any],
+    *,
+    default_capacity: int | None | object = CacheManager._MISSING,
+    overrides: Mapping[str, int | None] | None = ...,
+    replace_overrides: bool = ...,
+) -> CacheCapacityConfig: ...
+def increment_edge_version(G: Any) -> None: ...
+def edge_version_cache(
+    G: Any,
+    key: Hashable,
+    builder: Callable[[], T],
+    *,
+    max_entries: int | None | object = CacheManager._MISSING,
+) -> T: ...
+def cached_nodes_and_A(
+    G: nx.Graph,
+    *,
+    cache_size: int | None = ...,
+    require_numpy: bool = ...,
+    prefer_sparse: bool = ...,
+    nodes: tuple[Any, ...] | None = ...,
+) -> tuple[tuple[Any, ...], Any]: ...
+def edge_version_update(G: TNFRGraph) -> ContextManager[None]: ...
+
+class _SeedCacheState:
+    cache: InstrumentedLRUCache[tuple[int, int], int] | None
+    maxsize: int
+
+class _CounterState(Generic[K]):
+    cache: InstrumentedLRUCache[K, int]
+    locks: dict[K, threading.RLock]
+    max_entries: int
+
+class _SeedHashCache(MutableMapping[tuple[int, int], int]):
+    _state_key: str
+
+    def __init__(
+        self,
+        *,
+        manager: CacheManager | None = ...,
+        state_key: str = ...,
+        default_maxsize: int = ...,
+    ) -> None: ...
+    def configure(self, maxsize: int) -> None: ...
+    def __getitem__(self, key: tuple[int, int]) -> int: ...
+    def __setitem__(self, key: tuple[int, int], value: int) -> None: ...
+    def __delitem__(self, key: tuple[int, int]) -> None: ...
+    def __iter__(self) -> Iterator[tuple[int, int]]: ...
+    def __len__(self) -> int: ...
+    def clear(self) -> None: ...
+    @property
+    def maxsize(self) -> int: ...
+    @property
+    def enabled(self) -> bool: ...
+    @property
+    def data(self) -> InstrumentedLRUCache[tuple[int, int], int] | None: ...
+
+class ScopedCounterCache(Generic[K]):
+    _state_key: str
+
+    def __init__(
+        self,
+        name: str,
+        max_entries: int | None = ...,
+        *,
+        manager: CacheManager | None = ...,
+        default_max_entries: int = ...,
+    ) -> None: ...
+    def configure(
+        self, *, force: bool = ..., max_entries: int | None = ...
+    ) -> None: ...
+    def clear(self) -> None: ...
+    def bump(self, key: K) -> int: ...
+    def __len__(self) -> int: ...
+    @property
+    def lock(self) -> threading.Lock | threading.RLock: ...
+    @property
+    def max_entries(self) -> int: ...
+    @property
+    def cache(self) -> InstrumentedLRUCache[K, int]: ...
+    @property
+    def locks(self) -> dict[K, threading.RLock]: ...
+
+# Internal symbols used by utils.__init__.py
+_GRAPH_CACHE_MANAGER_KEY: str
+
+def _graph_cache_manager(graph: MutableMapping[str, Any]) -> CacheManager: ...

tnfr/{callback_utils.py → utils/callbacks.py}

@@ -9,28 +9,35 @@ modify shared state.
 
 from __future__ import annotations
 
-
-from typing import Any, TypedDict
-from enum import Enum
+import threading
+import traceback
 from collections import defaultdict, deque
-from collections.abc import Callable, Mapping, Iterable
+from collections.abc import Callable, Iterable, Mapping
+from enum import Enum
+from typing import Any, NamedTuple
+
+import networkx as nx
+
+from ..constants import DEFAULTS
+from ..locking import get_lock
+from .init import get_logger
+from .data import is_non_string_sequence
+from ..types import CallbackError
 
-import traceback
-import threading
-from .logging_utils import get_logger
-from .constants import DEFAULTS
-from .locking import get_lock
 
-from .trace import CallbackSpec
-from .collections_utils import is_non_string_sequence
+class CallbackSpec(NamedTuple):
+    """Specification for a registered callback."""
+
+    name: str | None
+    func: Callable[..., Any]
 
-import networkx as nx  # type: ignore[import-untyped]
 
 __all__ = (
     "CallbackEvent",
     "CallbackManager",
     "callback_manager",
     "CallbackError",
+    "CallbackSpec",
 )
 
 logger = get_logger(__name__)
@@ -42,6 +49,7 @@ class CallbackEvent(str, Enum):
     BEFORE_STEP = "before_step"
     AFTER_STEP = "after_step"
     ON_REMESH = "on_remesh"
+    CACHE_METRICS = "cache_metrics"
 
 
 class CallbackManager:
@@ -131,9 +139,7 @@ class CallbackManager:
             cb_name = name or getattr(func, "__name__", None)
             spec = CallbackSpec(cb_name, func)
             existing_map = cbs[event]
-            strict = bool(
-                G.graph.get("CALLBACKS_STRICT", DEFAULTS["CALLBACKS_STRICT"])
-            )
+            strict = bool(G.graph.get("CALLBACKS_STRICT", DEFAULTS["CALLBACKS_STRICT"]))
             key = _reconcile_callback(event, existing_map, spec, strict)
 
             existing_map[key] = spec
@@ -152,9 +158,7 @@ class CallbackManager:
         event = _normalize_event(event)
         with self._lock:
             cbs = dict(self._ensure_callbacks_nolock(G).get(event, {}))
-            strict = bool(
-                G.graph.get("CALLBACKS_STRICT", DEFAULTS["CALLBACKS_STRICT"])
-            )
+            strict = bool(G.graph.get("CALLBACKS_STRICT", DEFAULTS["CALLBACKS_STRICT"]))
         if ctx is None:
             ctx = {}
         for spec in cbs.values():
@@ -185,17 +189,6 @@ Callback = Callable[["nx.Graph", dict[str, Any]], None]
 CallbackRegistry = dict[str, dict[str, "CallbackSpec"]]
 
 
-class CallbackError(TypedDict):
-    """Metadata for a failed callback invocation."""
-
-    event: str
-    step: int | None
-    error: str
-    traceback: str
-    fn: str
-    name: str | None
-
-
 def _func_id(fn: Callable[..., Any]) -> str:
     """Return a deterministic identifier for ``fn``.
 
@@ -211,9 +204,7 @@ def _func_id(fn: Callable[..., Any]) -> str:
     return f"{module}.{qualname}"
 
 
-def _validate_registry(
-    G: "nx.Graph", cbs: Any, dirty: set[str]
-) -> CallbackRegistry:
+def _validate_registry(G: "nx.Graph", cbs: Any, dirty: set[str]) -> CallbackRegistry:
     """Validate and normalise the callback registry.
 
     ``cbs`` is coerced to a ``defaultdict(dict)`` and any events listed in
@@ -246,13 +237,13 @@ def _validate_registry(
     return cbs
 
 
-
-
 def _normalize_callbacks(entries: Any) -> dict[str, CallbackSpec]:
     """Return ``entries`` normalised into a callback mapping."""
     if isinstance(entries, Mapping):
         entries_iter = entries.values()
-    elif isinstance(entries, Iterable) and not isinstance(entries, (str, bytes, bytearray)):
+    elif isinstance(entries, Iterable) and not isinstance(
+        entries, (str, bytes, bytearray)
+    ):
         entries_iter = entries
     else:
         return {}
@@ -328,7 +319,6 @@ def _normalize_callback_entry(entry: Any) -> "CallbackSpec | None":
         return None
 
 
-
 def _reconcile_callback(
     event: str,
     existing_map: dict[str, CallbackSpec],
@@ -383,6 +373,3 @@ def _reconcile_callback(
 # ---------------------------------------------------------------------------
 
 callback_manager = CallbackManager()
-
-
-

tnfr/utils/callbacks.pyi

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+import networkx as nx
+from ..types import CallbackError as CallbackError
+from _typeshed import Incomplete
+from collections.abc import Callable
+from enum import Enum
+from typing import Any, NamedTuple
+
+__all__ = [
+    "CallbackEvent",
+    "CallbackManager",
+    "callback_manager",
+    "CallbackError",
+    "CallbackSpec",
+]
+
+Callback = Callable[[nx.Graph, dict[str, Any]], None]
+
+class CallbackSpec(NamedTuple):
+    name: str | None
+    func: Callable[..., Any]
+
+class CallbackEvent(str, Enum):
+    BEFORE_STEP = "before_step"
+    AFTER_STEP = "after_step"
+    ON_REMESH = "on_remesh"
+    CACHE_METRICS = "cache_metrics"
+
+class CallbackManager:
+    def __init__(self) -> None: ...
+    def get_callback_error_limit(self) -> int: ...
+    def set_callback_error_limit(self, limit: int) -> int: ...
+    def register_callback(
+        self,
+        G: nx.Graph,
+        event: CallbackEvent | str,
+        func: Callback,
+        *,
+        name: str | None = None,
+    ) -> Callback: ...
+    def invoke_callbacks(
+        self, G: nx.Graph, event: CallbackEvent | str, ctx: dict[str, Any] | None = None
+    ) -> None: ...
+
+callback_manager: Incomplete
+
+def _normalize_callbacks(entries: Any) -> dict[str, CallbackSpec]: ...
+def _normalize_callback_entry(entry: Any) -> CallbackSpec | None: ...