omdev 0.0.0.dev28__py3-none-any.whl → 0.0.0.dev30__py3-none-any.whl

omdev/cache/compute/cache.py

@@ -0,0 +1,120 @@
+"""
+TODO:
+ - filesystem OPTIONAL
+ - also postgres + (s3?) blobstore
+ - locking
+
+==
+
+TODO (old):
+ - are pickles stable?
+ - ttl
+ - np mmap
+  - https://numpy.org/doc/stable/reference/generated/numpy.ndarray.flags.html freeze ndarray w writable=False
+ - compress?
+ - overlap w/ jobs/dags/batches/whatever
+ - joblib
+ - ** INPUTS **
+  - if underlying impl changes, bust
+  - kinda reacty/reffy/signally
+ - decorator unwrapping and shit
+ - proactive deep invalidate
+ ----
+ - version can be anything - hashes, etc
+ - version generators - one for ast
+ - configurable serde - marshal vs pickle? marshal w/ override for ndarray to write to file?
+ - ok:
+  - @fn - version, passive=False, deps=[Objectable, …]
+   - it no version use ast - specifically {'ast': <md5>}
+   - but if present just use literal they gave, probably int
+   - idiom: Version can be a frozendict, conventionally of str -> ta.Hashable
+  - auto deps - fn can get containing Packages
+  - Module, Resource, …
+  - hrm.. LiteralVersion, MapVersion? + custom Marshal? need to deser as frozendict
+ - storage
+  - object table? w/ versions? strictly one row per object, evict objects with diff versions than those encountered
+  - nah Cache iface, SimpleCache, SqlCache
+  - dir structure: __package__/__qualname__/... ?
+ - next: Versions get squashed into VersionHash, store whole version in db but only pass and cmp md5
+  - thus VersionHashMap
+
+manifest stuff
+ - serialization_version
+ - lib_version
+ - lib_revision
+
+See:
+ - https://jax.readthedocs.io/en/latest/autodidax.html
+ - https://github.com/tinygrad/tinygrad/blob/78699d9924feb96dc0bac88c3646b5d4f9ecad23/tinygrad/engine/jit.py
+ - https://github.com/SeaOfNodes/Simple/tree/c7445ad142aeaece5b2b1059c193735ba7e509d9 (gvn)
+ - https://github.com/joblib/joblib/tree/bca1f4216a38cff82a85371c45dde79bed977d0e/joblib
+ - https://docs.python.org/3/library/pickle.html#pickle.Pickler.dispatch_table
+
+Don't see:
+ - https://github.com/amakelov/mandala
+"""
+import copy
+import typing as ta
+
+from omlish import collections as col
+
+from .storage import Storage
+from .types import CacheEntry
+from .types import CacheKey
+from .types import CacheResult
+from .types import CacheStats
+from .types import Name
+from .types import ObjectResolver
+from .types import VersionMap
+
+
+class Cache:
+    def __init__(
+            self,
+            resolver: ObjectResolver,
+            storage: Storage,
+    ) -> None:
+        super().__init__()
+
+        self._resolver = resolver
+        self._storage = storage
+
+        self._stats = CacheStats()
+
+    @property
+    def stats(self) -> CacheStats:
+        return copy.deepcopy(self._stats)
+
+    def _build_version_map(self, names: ta.Iterable[Name]) -> VersionMap:
+        dct = {}
+        for n in names:
+            c = self._resolver.resolve(n)
+            dct[n] = c.version
+        return col.frozendict(dct)
+
+    def get(self, key: CacheKey) -> CacheResult | None:
+        entry = self._storage.get(key)
+        if entry is None:
+            self._stats.num_misses += 1
+            return None
+
+        new_versions = self._build_version_map(entry.versions)
+        if entry.versions != new_versions:
+            self._storage.delete(key)
+            self._stats.num_invalidates += 1
+            return None
+
+        self._stats.num_hits += 1
+        return CacheResult(
+            True,
+            entry.versions,
+            entry.value,
+        )
+
+    def put(self, key: CacheKey, versions: VersionMap, val: ta.Any) -> None:
+        self._storage.put(CacheEntry(
+            key,
+            versions,
+            val,
+        ))
+        self._stats.num_puts += 1

omdev/cache/compute/contexts.py

@@ -0,0 +1,137 @@
+"""
+TODO:
+ - weak children if not in debug?
+"""
+import abc
+import typing as ta
+
+from omlish import check
+from omlish import collections as col
+from omlish import lang
+
+from .types import CacheKey
+from .types import CacheResult
+from .types import Object
+from .types import VersionMap
+from .types import merge_version_maps
+
+
+##
+
+
+class Context(lang.Abstract, lang.Sealed):
+    def __init__(
+            self,
+            obj: Object,
+            *,
+            dependencies: VersionMap = col.frozendict(),
+            parent: ta.Optional['Context'] = None,
+    ) -> None:
+        super().__init__()
+
+        self._obj = obj
+        self._dependencies = dependencies
+        self._parent = parent
+
+        check.equal(set(dependencies), set(obj.dependencies))
+
+        self._result: CacheResult | None = None
+        self._children: list[Context] = []
+
+        if parent is not None:
+            check.state(not parent.done)
+            parent._children.append(self)  # noqa
+
+    #
+
+    @property
+    def object(self) -> Object:
+        return self._obj
+
+    @property
+    def dependencies(self) -> VersionMap:
+        return self._dependencies
+
+    @property
+    def parent(self) -> ta.Optional['Context']:
+        return self._parent
+
+    @property
+    def children(self) -> ta.Sequence['Context']:
+        return self._children
+
+    @property
+    @abc.abstractmethod
+    def done(self) -> bool:
+        raise NotImplementedError
+
+    @lang.cached_function
+    @ta.final
+    def versions(self) -> VersionMap:
+        check.state(self.done)
+        return merge_version_maps(
+            self._obj.as_version_map,
+            self._dependencies,
+            self._impl_versions(),
+            *[c.versions() for c in self._children],
+        )
+
+    @abc.abstractmethod
+    def _impl_versions(self) -> VersionMap:
+        raise NotImplementedError
+
+
+class ActiveContext(Context, lang.Final):
+    def __init__(self, obj: Object, key: CacheKey, **kwargs: ta.Any) -> None:
+        check.arg(not obj.passive)
+
+        super().__init__(obj, **kwargs)
+
+        self._key = key
+
+        self._result: CacheResult | None = None
+
+    @property
+    def key(self) -> CacheKey:
+        return self._key
+
+    @property
+    def done(self) -> bool:
+        return self._result is not None
+
+    def set_hit(self, result: CacheResult) -> None:
+        check.state(result.hit)
+        self._result = check.replacing_none(self._result, result)
+        self.versions()
+
+    def set_miss(self, val: ta.Any) -> None:
+        self._result = check.replacing_none(self._result, CacheResult(
+            False,
+            VersionMap(),
+            val,
+        ))
+        self.versions()
+
+    def _impl_versions(self) -> VersionMap:
+        return check.not_none(self._result).versions
+
+
+class PassiveContext(Context, lang.Final):
+    def __init__(self, obj: Object, **kwargs: ta.Any) -> None:
+        check.arg(obj.passive)
+
+        super().__init__(obj, **kwargs)
+
+        self._done = False
+
+    @property
+    def done(self) -> bool:
+        return self._done
+
+    def finish(self) -> None:
+        check.state(not self._done)
+        self._done = True
+
+    @lang.cached_function
+    def _impl_versions(self) -> VersionMap:
+        return col.frozendict()

omdev/cache/compute/currents.py

@@ -0,0 +1,78 @@
+"""
+TODO:
+ - threadlocal / contextvar / whatever
+"""
+import contextlib
+import typing as ta
+
+from omlish import check
+
+from .cache import Cache
+from .contexts import ActiveContext
+from .contexts import Context
+from .contexts import PassiveContext
+from .types import CacheKey
+from .types import Object
+
+
+CacheT = ta.TypeVar('CacheT', bound='Cache')
+
+
+##
+
+
+_CURRENT_CACHE: Cache | None = None
+
+
+@contextlib.contextmanager
+def setting_current_cache(cache: CacheT) -> ta.Iterator[CacheT]:
+    global _CURRENT_CACHE
+    prev = _CURRENT_CACHE
+
+    try:
+        _CURRENT_CACHE = cache
+        yield cache
+
+    finally:
+        check.is_(_CURRENT_CACHE, cache)
+        _CURRENT_CACHE = prev
+
+
+def get_current_cache() -> Cache | None:
+    return _CURRENT_CACHE
+
+
+#
+
+
+_CURRENT_CONTEXT: Context | None = None
+
+
+@contextlib.contextmanager
+def setting_current_context(
+        obj: Object,
+        key: CacheKey | None = None,
+        **kwargs: ta.Any,
+) -> ta.Iterator[Context]:
+    global _CURRENT_CONTEXT
+    prev = _CURRENT_CONTEXT
+
+    ctx_kw = dict(
+        parent=prev,
+        **kwargs,
+    )
+
+    ctx: Context
+    if obj.passive:
+        check.none(key)
+        ctx = PassiveContext(obj, **ctx_kw)
+    else:
+        ctx = ActiveContext(obj, check.isinstance(key, CacheKey), **ctx_kw)
+
+    try:
+        _CURRENT_CONTEXT = ctx
+        yield ctx
+
+    finally:
+        check.is_(_CURRENT_CONTEXT, ctx)
+        _CURRENT_CONTEXT = prev

omdev/cache/compute/fns.py

@@ -0,0 +1,157 @@
+"""
+TODO:
+ - user specified key construction (skip args, default, transform, etc)
+ - option to not cache if too fast
+ - auto parent package chain per-module/package-ish CACHE_VERSION convention
+ - meditate on decos, descriptors, unwrapping, etc
+ - auto metadata:
+  - source
+  - qualname
+  - location
+  - ast? ast hash?
+  - keep src anyway, but just for warn
+   - strip comments?
+ - global tracking?
+  - *** or, at least, global constants? *** explicit Const obj, on access attach? hm...
+"""
+import functools
+import importlib
+import typing as ta
+
+from omlish import check
+from omlish import collections as col
+from omlish import dataclasses as dc
+from omlish import lang
+
+from .contexts import ActiveContext
+from .contexts import PassiveContext
+from .currents import get_current_cache
+from .currents import setting_current_context
+from .types import CacheKey
+from .types import Metadata
+from .types import Name
+from .types import Object
+from .types import ObjectResolver
+from .types import Version
+
+
+P = ta.ParamSpec('P')
+T = ta.TypeVar('T')
+
+
+##
+
+
+@dc.dataclass(frozen=True)
+class FnName(Name, lang.Final):
+    module: str
+    qualname: str
+
+
+@dc.dataclass(frozen=True)
+class FnObject(Object, lang.Final, ta.Generic[P, T]):
+    name: Name = dc.xfield(override=True)
+    fn: ta.Callable[P, T]  # type: ignore
+    version: Version = dc.xfield(override=True)
+
+    _: dc.KW_ONLY
+
+    dependencies: ta.AbstractSet[Name] = dc.xfield(default=frozenset(), override=True)
+    passive: bool = dc.xfield(default=False, override=True)
+    metadata: Metadata = dc.xfield(default=col.frozendict(), override=True)
+
+    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> T:
+        return self.fn(*args, **kwargs)
+
+
+class FnObjectResolver(ObjectResolver):
+    def resolve(self, name: Name) -> FnObject:
+        fname = check.isinstance(name, FnName)
+
+        mod = importlib.import_module(fname.module)
+        obj = mod
+        for a in fname.qualname.split('.'):
+            obj = getattr(obj, a)
+
+        check.callable(obj)
+        fc = check.isinstance(obj.__cacheable__, FnObject)
+
+        return fc
+
+
+@dc.dataclass(frozen=True)
+class FnCacheKey(CacheKey[FnName], lang.Final):
+    args: tuple
+    kwargs: col.frozendict[str, ta.Any]
+
+    @dc.validate
+    def _check_fn_types(self) -> bool:
+        return (
+            isinstance(self.name, FnName) and
+            isinstance(self.args, tuple) and
+            isinstance(self.kwargs, col.frozendict)
+        )
+
+
+##
+
+
+def fn(
+        version: Version,
+        *,
+        passive: bool = False,
+        metadata: Metadata = col.frozendict(),
+) -> ta.Callable[[T], T]:
+    def outer(fn):
+        @functools.wraps(fn)
+        def inner(*args, **kwargs):
+            # NOTE: just for testing :x allows updating
+            # TODO: proper wrapper obj probably (enforce name resolution)
+            obj = inner.__cacheable__  # type: ignore
+
+            if (cache := get_current_cache()) is None:
+                return fn(*args, **kwargs)
+
+            if obj.passive:
+                with setting_current_context(obj) as ctx:
+                    pctx = check.isinstance(ctx, PassiveContext)
+                    val = fn(*args, **kwargs)
+                    pctx.finish()
+                    return val
+
+            key = FnCacheKey(
+                obj.name,
+                args,
+                col.frozendict(kwargs),
+            )
+
+            with setting_current_context(obj, key) as ctx:
+                actx = check.isinstance(ctx, ActiveContext)
+
+                if (hit := cache.get(key)) is not None:
+                    actx.set_hit(hit)
+                    return hit.value
+
+                val = fn(*args, **kwargs)
+                actx.set_miss(val)
+                cache.put(
+                    key,
+                    actx.versions(),
+                    val,
+                )
+                return val
+
+        inner.__cacheable__ = FnObject(  # type: ignore
+            FnName(
+                fn.__module__,
+                fn.__qualname__,
+            ),
+            fn,
+            version,
+            passive=passive,
+            metadata=metadata,
+        )
+
+        return inner
+
+    return outer  # noqa

omdev/cache/compute/resolvers.py

@@ -0,0 +1,23 @@
+from .types import Name
+from .types import Object
+from .types import ObjectResolver
+
+
+class CachingObjectResolver(ObjectResolver):
+    def __init__(self, child: ObjectResolver) -> None:
+        super().__init__()
+
+        self._child = child
+        self._dct: dict[Name, Object] = {}
+
+    def clear(self) -> None:
+        self._dct.clear()
+
+    def resolve(self, name: Name) -> Object:
+        try:
+            return self._dct[name]
+        except KeyError:
+            pass
+        ret = self._child.resolve(name)
+        self._dct[name] = ret
+        return ret

omdev/cache/compute/storage.py

@@ -0,0 +1,39 @@
+import abc
+
+from .types import CacheEntry
+from .types import CacheKey
+
+
+class Storage(abc.ABC):
+    @abc.abstractmethod
+    def get(self, key: CacheKey) -> CacheEntry | None:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def put(self, entry: CacheEntry) -> None:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def delete(self, key: CacheKey) -> None:
+        raise NotImplementedError
+
+
+class DictStorage(Storage):
+    def __init__(self) -> None:
+        super().__init__()
+
+        self._dct: dict[CacheKey, CacheEntry] = {}
+
+    def get(self, key: CacheKey) -> CacheEntry | None:
+        try:
+            return self._dct[key]
+        except KeyError:
+            return None
+
+    def put(self, entry: CacheEntry) -> None:
+        if entry.key in self._dct:
+            raise KeyError(entry.key)
+        self._dct[entry.key] = entry
+
+    def delete(self, key: CacheKey) -> None:
+        del self._dct[key]

omdev/cache/compute/types.py

@@ -0,0 +1,144 @@
+import abc
+import typing as ta
+
+from omlish import cached
+from omlish import collections as col
+from omlish import dataclasses as dc
+from omlish import lang
+
+
+T = ta.TypeVar('T')
+
+
+##
+
+
+NameT = ta.TypeVar('NameT', bound='Name')
+
+
+class Name(lang.Abstract):
+    pass
+
+
+##
+
+
+Version: ta.TypeAlias = ta.Hashable
+
+
+def version(**kwargs: Version) -> Version:
+    return col.frozendict(**kwargs)
+
+
+VersionMap: ta.TypeAlias = col.frozendict[Name, Version]
+
+
+def merge_version_maps(
+        *dcts: ta.Mapping[Name, Version],
+) -> VersionMap:
+    out: dict[Name, Version] = {}
+    for dct in dcts:
+        for name, version in dct.items():
+            try:
+                ex = out[name]
+            except KeyError:
+                out[name] = version
+            else:
+                if ex != version:
+                    raise Exception(f'Version mismatch: {ex} {version}')
+    return col.frozendict(out)
+
+
+##
+
+
+Metadata: ta.TypeAlias = ta.Mapping[str, ta.Any]  # *not* hashed - advisory
+
+
+class Object(lang.Abstract):
+    @property
+    @abc.abstractmethod
+    def name(self) -> Name:
+        raise NotImplementedError
+
+    @property
+    @abc.abstractmethod
+    def version(self) -> Version:
+        raise NotImplementedError
+
+    @property
+    @abc.abstractmethod
+    def dependencies(self) -> ta.AbstractSet[Name]:
+        raise NotImplementedError
+
+    @property
+    @abc.abstractmethod
+    def passive(self) -> bool:
+        raise NotImplementedError
+
+    @property
+    @abc.abstractmethod
+    def metadata(self) -> Metadata:
+        raise NotImplementedError
+
+    @cached.property
+    def as_version_map(self) -> VersionMap:
+        return col.frozendict({self.name: self.version})
+
+
+##
+
+
+@dc.dataclass(frozen=True)
+@dc.extra_params(cache_hash=True)
+class CacheKey(lang.Abstract, ta.Generic[NameT]):
+    name: NameT
+
+    @dc.validate
+    def _check_types(self) -> bool:
+        hash(self)
+        return isinstance(self.name, Name)
+
+
+@dc.dataclass(frozen=True)
+class CacheResult(ta.Generic[T], lang.Final):
+    hit: bool
+    versions: VersionMap
+    value: T
+
+
+##
+
+
+class ObjectResolver(lang.Abstract):
+    @abc.abstractmethod
+    def resolve(self, name: Name) -> Object:
+        raise NotImplementedError
+
+
+##
+
+
+@dc.dataclass()
+class CacheStats:
+    num_hits: int = 0
+    num_misses: int = 0
+    num_invalidates: int = 0
+    num_puts: int = 0
+
+
+##
+
+
+@dc.dataclass(frozen=True)
+class CacheEntry:
+    key: CacheKey
+    versions: VersionMap
+    value: ta.Any
+
+    @dc.validate
+    def _check_types(self) -> bool:
+        return (
+            isinstance(self.key, CacheKey) and
+            isinstance(self.versions, col.frozendict)
+        )