prism-engine 1.0.0__py3-none-any.whl

prism/__init__.py

@@ -0,0 +1,156 @@
+"""Prism Compute Engine — speculative-redundant-compute elimination for ML workloads.
+
+Public API:
+    from prism import PrismEngine, WorkloadClassifier, WasteReport, PrismConfig
+"""
+import contextlib
+
+from prism.hashing import hash_object, hash_call, hash_numpy, fingerprint_array
+from prism.config import PrismConfig, CacheConfig, HashingConfig, PredictorConfig, AdapterConfig, LoggingConfig, load_config
+
+__version__ = "1.0.0"
+
+
+# ─── Default engine + decorator API ──────────────────────────────────────────
+_default_engine = None
+
+
+def get_default_engine():
+    """Return (creating if needed) the module-global default PrismEngine.
+
+    The default engine is in ACTIVE mode with the default config. Use
+    ``set_default_engine`` to override it (e.g. with a disk-backed or observe
+    engine).
+    """
+    global _default_engine
+    if _default_engine is None:
+        from prism.core import PrismEngine, EngineMode
+        _default_engine = PrismEngine(mode=EngineMode.ACTIVE)
+    return _default_engine
+
+
+def set_default_engine(engine) -> None:
+    """Override the module-global default engine used by ``@prism.cache``."""
+    global _default_engine
+    _default_engine = engine
+
+
+def cached(arg=None):
+    """Decorator that routes a function's calls through a Prism engine.
+
+    Two usage forms::
+
+        @prism.cached                      # uses the default active engine
+        def train(X, y):
+            ...
+
+        @prism.cached(my_engine)           # use an explicit engine
+        def train(X, y):
+            ...
+
+    In ACTIVE mode the wrapped function only runs on a cache miss; in OBSERVE
+    mode it always runs and the would-have-been-hit is recorded.
+
+    (Named ``prism.cached`` rather than ``prism.cache`` to avoid clashing with
+    the ``prism.cache`` submodule that holds the cache backends.)
+    """
+    # `arg` is a plain function -> used as `@prism.cached` with no engine arg.
+    if arg is None or not hasattr(arg, "call"):
+        engine = get_default_engine()
+        if arg is None:
+            # `@prism.cached(engine)` form: return the decorator bound to engine.
+            return engine.cached
+        # `@prism.cached` direct form: arg is the function.
+        return engine.cached(arg)
+    # `arg` is an engine -> `@prism.cached(engine)` returns its cache decorator.
+    return arg.cached
+
+
+@contextlib.contextmanager
+def scope(engine, *adapter_classes):
+    """Context manager that wraps a set of adapters around an engine and
+    guarantees they are unwrapped on exit (even on exception).
+
+    Example::
+
+        from prism import PrismEngine, EngineMode, scope
+        from prism.adapters.sklearn_adapter import SklearnAdapter
+
+        engine = PrismEngine(mode=EngineMode.ACTIVE)
+        with scope(engine, SklearnAdapter):
+            run_pipeline()   # sklearn ops routed through Prism
+        # adapters restored on exit
+    """
+    adapters = []
+    try:
+        for cls in adapter_classes:
+            a = cls(engine)
+            a.wrap()
+            adapters.append(a)
+        yield engine
+    finally:
+        for a in adapters:
+            try:
+                a.unwrap()
+            except Exception:
+                pass
+
+
+def __getattr__(name: str):
+    """Lazy imports — modules are loaded on first attribute access."""
+    _lazy = {
+        "PrismEngine": ("prism.core", "PrismEngine"),
+        "EngineMode": ("prism.core", "EngineMode"),
+        "WorkloadClassifier": ("prism.classifier", "WorkloadClassifier"),
+        "ClassificationReport": ("prism.classifier", "ClassificationReport"),
+        "Observer": ("prism.observer", "Observer"),
+        "WasteReport": ("prism.report", "WasteReport"),
+        "StatisticalPredictor": ("prism.predictor", "StatisticalPredictor"),
+        "LSTMPredictor": ("prism.predictor_lstm", "LSTMPredictor"),
+        "MetricsExporter": ("prism.metrics", "MetricsExporter"),
+        "RedisCache": ("prism.redis_cache", "RedisCache"),
+        "Serializer": ("prism.serialization", "Serializer"),
+        "SemanticCache": ("prism.semantic_cache", "SemanticCache"),
+        "TelemetryHook": ("prism.telemetry", "TelemetryHook"),
+        "CostConfig": ("prism.config", "CostConfig"),
+    }
+    if name in _lazy:
+        import importlib
+
+        mod_name, attr = _lazy[name]
+        mod = importlib.import_module(mod_name)
+        return getattr(mod, attr)
+    raise AttributeError(f"module 'prism' has no attribute {name!r}")
+
+
+__all__ = [
+    "PrismEngine",
+    "EngineMode",
+    "WorkloadClassifier",
+    "ClassificationReport",
+    "Observer",
+    "WasteReport",
+    "StatisticalPredictor",
+    "LSTMPredictor",
+    "MetricsExporter",
+    "RedisCache",
+    "Serializer",
+    "SemanticCache",
+    "TelemetryHook",
+    "CostConfig",
+    "PrismConfig",
+    "CacheConfig",
+    "HashingConfig",
+    "PredictorConfig",
+    "AdapterConfig",
+    "LoggingConfig",
+    "load_config",
+    "hash_object",
+    "hash_call",
+    "hash_numpy",
+    "fingerprint_array",
+    "cached",
+    "scope",
+    "get_default_engine",
+    "set_default_engine",
+]

prism/adapters/__init__.py

@@ -0,0 +1,5 @@
+"""Prism adapters — framework-specific wrappers that route ML ops through PrismEngine.
+
+Each adapter patches framework methods to route through a PrismEngine, and
+provides unwrap() to restore original behavior.
+"""

prism/adapters/data_pipeline.py

@@ -0,0 +1,146 @@
+"""Data pipeline adapter — routes pandas/numpy data transforms through Prism.
+
+Validated workload: repeated transform() / fit_transform() on identical
+DataFrames or arrays — dedup the transform when input data is unchanged.
+
+Patches:
+  - numpy: np.sort, np.unique, np.argsort, np.partition, np.searchsorted,
+            np.histogram, np.bincount, np.digitize, np.percentile, np.quantile
+  - pandas: DataFrame.apply, DataFrame.groupby, DataFrame.merge,
+             DataFrame.pivot_table, DataFrame.resample, Series.map,
+             Series.value_counts
+"""
+from __future__ import annotations
+
+from typing import Any, Optional
+
+import numpy as np
+
+from prism.core import PrismEngine
+
+
+# numpy functions to patch
+_NUMPY_FUNCS = (
+    "sort", "unique", "argsort", "partition", "searchsorted",
+    "histogram", "bincount", "digitize", "percentile", "quantile",
+)
+
+# pandas DataFrame methods to patch
+_PANDAS_DF_METHODS = (
+    "apply", "groupby", "merge", "pivot_table", "resample",
+)
+
+# pandas Series methods to patch
+_PANDAS_SERIES_METHODS = (
+    "map", "value_counts",
+)
+
+
+class DataPipelineAdapter:
+    """Patch common data-transform functions to route through Prism.
+
+    Usage:
+        adapter = DataPipelineAdapter(engine)
+        adapter.wrap()
+        ...run your data pipeline...
+        adapter.unwrap()
+    """
+
+    def __init__(self, engine: PrismEngine) -> None:
+        self.engine = engine
+        self._originals: dict[str, Any] = {}
+
+    def wrap(self) -> "DataPipelineAdapter":
+        # numpy functions
+        for fn_name in _NUMPY_FUNCS:
+            key = f"np.{fn_name}"
+            if key in self._originals:
+                continue
+            original = getattr(np, fn_name, None)
+            if original is None:
+                continue
+            self._originals[key] = original
+            setattr(np, fn_name, self._make_np_wrapper(original, fn_name))
+
+        # pandas methods (optional)
+        try:
+            import pandas as pd
+
+            for method_name in _PANDAS_DF_METHODS:
+                key = f"DataFrame.{method_name}"
+                if key in self._originals:
+                    continue
+                if hasattr(pd.DataFrame, method_name):
+                    original = getattr(pd.DataFrame, method_name)
+                    self._originals[key] = original
+                    setattr(
+                        pd.DataFrame,
+                        method_name,
+                        self._make_method_wrapper(original, f"DataFrame.{method_name}"),
+                    )
+
+            for method_name in _PANDAS_SERIES_METHODS:
+                key = f"Series.{method_name}"
+                if key in self._originals:
+                    continue
+                if hasattr(pd.Series, method_name):
+                    original = getattr(pd.Series, method_name)
+                    self._originals[key] = original
+                    setattr(
+                        pd.Series,
+                        method_name,
+                        self._make_method_wrapper(original, f"Series.{method_name}"),
+                    )
+        except ImportError:
+            pass
+
+        return self
+
+    def unwrap(self) -> None:
+        # Restore numpy
+        for key, original in list(self._originals.items()):
+            if key.startswith("np."):
+                setattr(np, key[3:], original)
+                del self._originals[key]
+        # Restore pandas
+        try:
+            import pandas as pd
+
+            for key, original in list(self._originals.items()):
+                if key.startswith("DataFrame."):
+                    setattr(pd.DataFrame, key[len("DataFrame."):], original)
+                    del self._originals[key]
+                elif key.startswith("Series."):
+                    setattr(pd.Series, key[len("Series."):], original)
+                    del self._originals[key]
+        except ImportError:
+            pass
+        self._originals.clear()
+
+    # ─── Context-manager protocol ─────────────────────────────────────────
+    def __enter__(self) -> "DataPipelineAdapter":
+        return self.wrap()
+
+    def __exit__(self, *exc) -> bool:
+        self.unwrap()
+        return False
+
+    def _make_np_wrapper(self, original, name: str):
+        engine = self.engine
+
+        def wrapped(*args, **kwargs):
+            fn_name = f"np.{name}"
+            return engine.call(original, args, kwargs, fn_name=fn_name)
+
+        wrapped.__name__ = name
+        return wrapped
+
+    def _make_method_wrapper(self, original, name: str):
+        engine = self.engine
+
+        def wrapped(self_obj, *args, **kwargs):
+            fn_name = name
+            return engine.call(original, (self_obj,) + args, kwargs, fn_name=fn_name)
+
+        wrapped.__name__ = name.split(".")[-1]
+        return wrapped

prism/adapters/jax_adapter.py

@@ -0,0 +1,77 @@
+"""JAX adapter — routes jit-compiled training steps through Prism.
+
+Validated workload: training sweep with repeated (params + data) — dedup
+jit-compiled steps. Optional dependency: skips gracefully when jax is absent.
+"""
+from __future__ import annotations
+
+from typing import Any, Optional
+
+from prism.core import PrismEngine
+
+
+class JAXAdapter:
+    """Patch jax.jit (and jax.pmap) to route through Prism.
+
+    Usage:
+        adapter = JAXAdapter(engine)
+        adapter.wrap()
+        ...run your jax workload...
+        adapter.unwrap()
+
+    If jax is not installed, wrap()/unwrap() are no-ops.
+    """
+
+    def __init__(self, engine: PrismEngine) -> None:
+        self.engine = engine
+        self._originals: dict[str, Any] = {}
+        self._jax_available = self._check_jax()
+
+    def _check_jax(self) -> bool:
+        try:
+            import jax  # noqa: F401
+
+            return True
+        except ImportError:
+            return False
+
+    def wrap(self) -> "JAXAdapter":
+        if not self._jax_available:
+            return self
+        import jax
+
+        if "jit" in self._originals:
+            return self
+        self._originals["jit"] = jax.jit
+        engine = self.engine
+        original_jit = jax.jit
+
+        def prism_jit(fun, *args, **kwargs):
+            # Wrap the jitted function so each call routes through Prism
+            jitted = original_jit(fun, *args, **kwargs)
+
+            def prism_call(*call_args, **call_kwargs):
+                fn_name = getattr(fun, "__qualname__", repr(fun))
+                return engine.call(jitted, call_args, call_kwargs, fn_name=fn_name)
+
+            return prism_call
+
+        jax.jit = prism_jit
+        return self
+
+    def unwrap(self) -> None:
+        if not self._jax_available or not self._originals:
+            return
+        import jax
+
+        for name, original in self._originals.items():
+            setattr(jax, name, original)
+        self._originals.clear()
+
+    # ─── Context-manager protocol ─────────────────────────────────────────
+    def __enter__(self) -> "JAXAdapter":
+        return self.wrap()
+
+    def __exit__(self, *exc) -> bool:
+        self.unwrap()
+        return False

prism/adapters/pytorch_adapter.py

@@ -0,0 +1,167 @@
+"""PyTorch adapter — routes forward passes and eval loops through Prism.
+
+Validated workload: eval loop over a model on the same val set across epochs —
+forward passes on unchanged (model + input) are deduped.
+
+Patching strategy:
+  - `nn.Module.__call__` is patched to intercept all forward passes.
+  - A re-entrancy guard prevents intercepting nested submodule calls.
+  - In training mode (`model.training == True`), the model's parameters change
+    between calls, so we include `model.training` in the hash key and also
+    hash the model's parameter state — this means training-step forward passes
+    are only deduped if the parameters haven't changed (rare in training, common
+    in eval).
+  - In eval mode (`model.training == False`), parameters are frozen, so
+    repeated forward passes on the same input are always deduped.
+"""
+from __future__ import annotations
+
+import threading
+from typing import Any, Optional
+
+from prism.core import PrismEngine
+from prism.hashing import hash_object
+
+
+class PyTorchAdapter:
+    """Patch torch.nn.Module.__call__ to route through Prism.
+
+    Usage:
+        adapter = PyTorchAdapter(engine)
+        adapter.wrap()
+        ...run your torch workload...
+        adapter.unwrap()
+
+    Options:
+        hash_params: If True (default), include model parameter state in the
+            cache key. This is essential for correctness in training mode —
+            without it, a forward pass after a parameter update would return
+            the stale cached result. In eval-only workloads, set to False for
+            faster hashing.
+        dedup_training: If False (default), don't dedup forward passes when
+            model.training == True (training steps almost always change params,
+            so dedup is pointless and the param hash is expensive). If True,
+            attempt dedup in training mode too (uses param hashing).
+    """
+
+    def __init__(
+        self,
+        engine: PrismEngine,
+        hash_params: bool = True,
+        dedup_training: bool = False,
+    ) -> None:
+        self.engine = engine
+        self.hash_params = hash_params
+        self.dedup_training = dedup_training
+        self._original_call = None
+        self._original_state_dict = None
+        self._torch_available = self._check_torch()
+        self._in_prism_call = threading.local()
+
+    def _check_torch(self) -> bool:
+        try:
+            import torch  # noqa: F401
+
+            return True
+        except ImportError:
+            return False
+
+    def wrap(self) -> "PyTorchAdapter":
+        if not self._torch_available:
+            return self
+        import torch.nn as nn
+
+        if self._original_call is not None:
+            return self
+        self._original_call = nn.Module.__call__
+        engine = self.engine
+        original = self._original_call
+        hash_params = self.hash_params
+        dedup_training = self.dedup_training
+
+        def patched_call(self_mod, *args, **kwargs):
+            # Re-entrancy guard
+            if getattr(self._in_prism_call, "active", False):
+                return original(self_mod, *args, **kwargs)
+
+            # In training mode, skip dedup unless explicitly enabled
+            if self_mod.training and not dedup_training:
+                return original(self_mod, *args, **kwargs)
+
+            self._in_prism_call.active = True
+            try:
+                fn_name = f"{type(self_mod).__qualname__}.forward"
+
+                if hash_params and hasattr(self_mod, "state_dict"):
+                    try:
+                        param_hash = _hash_model_params(self_mod)
+                    except Exception:
+                        param_hash = None
+                    if param_hash is not None:
+                        # Pass param_hash as a kwarg to the engine so it's
+                        # included in the cache key. The wrapper strips it
+                        # before calling the original function.
+                        _orig = original
+
+                        def _param_aware_wrapper(*a, **kw):
+                            kw.pop("__prism_param_hash__", None)
+                            return _orig(*a, **kw)
+
+                        return engine.call(
+                            _param_aware_wrapper,
+                            (self_mod,) + args,
+                            {**kwargs, "__prism_param_hash__": param_hash},
+                            fn_name=fn_name,
+                        )
+
+                return engine.call(
+                    original,
+                    (self_mod,) + args,
+                    kwargs,
+                    fn_name=fn_name,
+                )
+            finally:
+                self._in_prism_call.active = False
+
+        nn.Module.__call__ = patched_call
+        return self
+
+    def unwrap(self) -> None:
+        if not self._torch_available or self._original_call is None:
+            return
+        import torch.nn as nn
+
+        nn.Module.__call__ = self._original_call
+        self._original_call = None
+
+    # ─── Context-manager protocol ─────────────────────────────────────────
+    def __enter__(self) -> "PyTorchAdapter":
+        return self.wrap()
+
+    def __exit__(self, *exc) -> bool:
+        self.unwrap()
+        return False
+
+
+def _hash_model_params(model: Any) -> str:
+    """Hash a model's parameters — used to detect parameter changes between calls.
+
+    Sums all parameter tensors and hashes the sum + count. This is O(n) in
+    parameters but much cheaper than hashing every tensor's full content.
+    For large models, this could be replaced with a sampled fingerprint.
+    """
+    import torch
+
+    parts = []
+    total_params = 0
+    param_sum = torch.tensor(0.0)
+    for name, param in model.named_parameters():
+        total_params += param.numel()
+        parts.append(f"{name}:{param.shape}")
+        param_sum = param_sum + param.detach().float().sum()
+    # Include buffers (e.g. BatchNorm running stats) — they change in eval too
+    for name, buf in model.named_buffers():
+        total_params += buf.numel()
+        parts.append(f"buf:{name}:{buf.shape}")
+        param_sum = param_sum + buf.detach().float().sum()
+    return f"params={total_params}|sum={float(param_sum):.10f}|shapes={','.join(parts)}"

prism/adapters/sklearn_adapter.py

@@ -0,0 +1,140 @@
+"""sklearn adapter — routes fit / fit_transform / transform / predict through Prism.
+
+Validated workload: GridSearchCV with repeated param combos — fit on identical
+(estimator + params + data) is deduped. Also covers StandardScaler, PCA, etc.
+
+Patching strategy: sklearn estimators define `fit` on their own concrete class
+(shadowing BaseEstimator), so we walk all loaded BaseEstimator subclasses and
+patch each class that defines the target method in its own __dict__.
+
+fit/fit_transform special handling: sklearn's `fit` mutates `self` in place AND
+returns `self`. On a cache hit, the engine returns a cached deep copy of a
+previously-fitted estimator — but the target estimator is never mutated. So
+the wrapper must copy fitted attributes (those ending in `_`) from the cached
+result onto the target estimator on a hit.
+"""
+from __future__ import annotations
+
+import copy
+from typing import Any
+
+from prism.core import PrismEngine
+from prism.hashing import hash_call
+
+
+class SklearnAdapter:
+    _PATCHED_METHODS = ("fit", "fit_transform", "transform", "predict", "score")
+
+    def __init__(self, engine: PrismEngine) -> None:
+        self.engine = engine
+        self._originals: dict[tuple[type, str], Any] = {}
+
+    def wrap(self) -> "SklearnAdapter":
+        from sklearn.base import BaseEstimator
+
+        classes = self._all_subclasses(BaseEstimator)
+        classes.add(BaseEstimator)
+        for cls in classes:
+            for method_name in self._PATCHED_METHODS:
+                key = (cls, method_name)
+                if key in self._originals:
+                    continue
+                if method_name not in cls.__dict__:
+                    continue
+                original = cls.__dict__[method_name]
+                if getattr(original, "_prism_wrapped", False):
+                    continue
+                self._originals[key] = original
+                if method_name in ("fit", "fit_transform"):
+                    wrapped = self._make_fit_wrapper(original, method_name)
+                else:
+                    wrapped = self._make_simple_wrapper(original, method_name)
+                wrapped._prism_wrapped = True
+                setattr(cls, method_name, wrapped)
+        return self
+
+    def unwrap(self) -> None:
+        for (cls, method_name), original in self._originals.items():
+            setattr(cls, method_name, original)
+        self._originals.clear()
+
+    # ─── Context-manager protocol ─────────────────────────────────────────
+    def __enter__(self) -> "SklearnAdapter":
+        return self.wrap()
+
+    def __exit__(self, *exc) -> bool:
+        self.unwrap()
+        return False
+
+    def _all_subclasses(self, cls: type) -> set[type]:
+        result = set()
+        for sub in cls.__subclasses__():
+            result.add(sub)
+            result.update(self._all_subclasses(sub))
+        return result
+
+    def _make_simple_wrapper(self, original, method_name: str):
+        """For methods that just return a value (transform, predict, score)."""
+        engine = self.engine
+
+        def wrapped(self_est, *args, **kwargs):
+            fn_name = f"{type(self_est).__qualname__}.{method_name}"
+            return engine.call(original, (self_est,) + args, kwargs, fn_name=fn_name)
+
+        wrapped.__name__ = method_name
+        wrapped.__qualname__ = f"PrismAdapter.{method_name}"
+        return wrapped
+
+    def _make_fit_wrapper(self, original, method_name: str):
+        """For fit / fit_transform — handles in-place mutation + state restoration."""
+        engine = self.engine
+
+        def wrapped(self_est, *args, **kwargs):
+            fn_name = f"{type(self_est).__qualname__}.{method_name}"
+            # Use the engine, but we need to know if it was a hit to restore state.
+            # Strategy: check engine.stats.misses before and after the call.
+            misses_before = engine.stats.misses
+            result = engine.call(original, (self_est,) + args, kwargs, fn_name=fn_name)
+            was_hit = engine.stats.misses == misses_before  # no new miss = hit
+
+            if was_hit and result is not self_est:
+                # Cache hit — `result` is a deep copy of a previously-fitted
+                # estimator. Copy its fitted attributes onto `self_est`.
+                _restore_fitted_state(result, self_est)
+                # For fit_transform, also return the cached transformed data
+                # (which is `result`'s transformed output, already returned by
+                # engine.call). For fit, return self_est (the now-restored estimator).
+                if method_name == "fit":
+                    return self_est
+                # fit_transform: result is the cached transformed array — return it
+                return result
+            # Miss (real fit ran, self_est is mutated) — return as normal
+            return result
+
+        wrapped.__name__ = method_name
+        wrapped.__qualname__ = f"PrismAdapter.{method_name}"
+        return wrapped
+
+
+def _restore_fitted_state(source: Any, target: Any) -> None:
+    """Copy fitted attributes (ending in `_`) from source estimator onto target.
+
+    This is the key to correct fit caching: on a cache hit, the target estimator
+    hasn't been mutated by a real fit, so we must transfer the fitted state from
+    the cached copy.
+    """
+    # Copy all attributes ending in `_` (sklearn convention for fitted attrs)
+    # plus any attribute that exists on source but not target.
+    for attr in dir(source):
+        if attr.startswith("_"):
+            continue
+        if not attr.endswith("_"):
+            continue
+        try:
+            val = getattr(source, attr)
+        except AttributeError:
+            continue
+        try:
+            setattr(target, attr, copy.deepcopy(val))
+        except (AttributeError, TypeError):
+            pass