#!/usr/bin/env python3 # -*- coding: utf-8 -*- # cache.py """ Memoization and caching utilities. """ # pylint: disable=redefined-builtin,redefined-outer-name,missing-docstring # pylint: disable=no-self-use,arguments-differ # pylint: disable=dangerous-default-value,redefined-builtin # pylint: disable=abstract-method import os import pickle from functools import namedtuple, update_wrapper, wraps import psutil import redis from . import config, constants _CacheInfo = namedtuple("CacheInfo", ["hits", "misses", "currsize"]) def memory_full(): """Check if the memory is too full for further caching.""" current_process = psutil.Process(os.getpid()) return (current_process.memory_percent() > config.MAXIMUM_CACHE_MEMORY_PERCENTAGE) class _HashedSeq(list): """This class guarantees that ``hash()`` will be called no more than once per element. This is important because the ``lru_cache()`` will hash the key multiple times on a cache miss. """ __slots__ = ('hashvalue',) def __init__(self, tup, hash=hash): super().__init__() self[:] = tup self.hashvalue = hash(tup) def __hash__(self): return self.hashvalue def _make_key(args, kwds, typed, kwd_mark=(object(),), fasttypes={int, str, frozenset, type(None)}, sorted=sorted, tuple=tuple, type=type, len=len): """Make a cache key from optionally typed positional and keyword arguments. The key is constructed in a way that is flat as possible rather than as a nested structure that would take more memory. If there is only a single argument and its data type is known to cache its hash value, then that argument is returned without a wrapper. This saves space and improves lookup speed. """ key = args if kwds: sorted_items = sorted(kwds.items()) key += kwd_mark for item in sorted_items: key += item if typed: key += tuple(type(v) for v in args) if kwds: key += tuple(type(v) for k, v in sorted_items) elif len(key) == 1 and type(key[0]) in fasttypes: return key[0] return _HashedSeq(key) def cache(cache={}, maxmem=config.MAXIMUM_CACHE_MEMORY_PERCENTAGE, typed=False): """Memory-limited cache decorator. ``maxmem`` is a float between 0 and 100, inclusive, specifying the maximum percentage of physical memory that the cache can use. If ``typed`` is ``True``, arguments of different types will be cached separately. For example, f(3.0) and f(3) will be treated as distinct calls with distinct results. Arguments to the cached function must be hashable. View the cache statistics named tuple (hits, misses, currsize) with f.cache_info(). Clear the cache and statistics with f.cache_clear(). Access the underlying function with f.__wrapped__. """ # Constants shared by all lru cache instances: # Unique object used to signal cache misses. sentinel = object() # Build a key from the function arguments. make_key = _make_key def decorating_function(user_function, hits=0, misses=0): full = False # Bound method to look up a key or return None. cache_get = cache.get if not maxmem: def wrapper(*args, **kwds): # Simple caching without memory limit. nonlocal hits, misses key = make_key(args, kwds, typed) result = cache_get(key, sentinel) if result is not sentinel: hits += 1 return result result = user_function(*args, **kwds) cache[key] = result misses += 1 return result else: def wrapper(*args, **kwds): # Memory-limited caching. nonlocal hits, misses, full key = make_key(args, kwds, typed) result = cache_get(key) if result is not None: hits += 1 return result result = user_function(*args, **kwds) if not full: cache[key] = result # Cache is full if the total recursive usage is greater # than the maximum allowed percentage. current_process = psutil.Process(os.getpid()) full = current_process.memory_percent() > maxmem misses += 1 return result def cache_info(): """Report cache statistics.""" return _CacheInfo(hits, misses, len(cache)) def cache_clear(): """Clear the cache and cache statistics.""" nonlocal hits, misses, full cache.clear() hits = misses = 0 full = False wrapper.cache_info = cache_info wrapper.cache_clear = cache_clear return update_wrapper(wrapper, user_function) return decorating_function class DictCache: """A generic dictionary-based cache. Intended to be used as an object-level cache of method results. """ def __init__(self): self.cache = {} self.hits = 0 self.misses = 0 def clear(self): self.cache = {} self.hits = 0 self.misses = 0 def size(self): """Number of items in cache""" return len(self.cache) def info(self): """Return info about cache hits, misses, and size""" return _CacheInfo(self.hits, self.misses, self.size()) def get(self, key): """Get a value out of the cache. Returns None if the key is not in the cache. Updates cache statistics. """ if key in self.cache: self.hits += 1 return self.cache[key] self.misses += 1 return None def set(self, key, value): """Set a value in the cache""" self.cache[key] = value # TODO: handle **kwarg keys if needed # See joblib.func_inspect.filter_args def key(self, *args, _prefix=None, **kwargs): """Get the cache key for the given function args. Kwargs: prefix: A constant to prefix to the key. """ if kwargs: raise NotImplementedError( 'kwarg cache keys not implemented') return (_prefix,) + tuple(args) def redis_init(db): return redis.StrictRedis(host=config.REDIS_CONFIG['host'], port=config.REDIS_CONFIG['port'], db=db) # Expose the StrictRedis API, maintaining one connection pool # The connection pool is multi-process safe, and is reinitialized when the # client detects a fork. See: # https://github.com/andymccurdy/redis-py/blob/5109cb4f/redis/connection.py#L950 # # TODO: rebuild connection after config changes? redis_conn = redis_init(config.REDIS_CONFIG['db']) def redis_available(): """Check if the Redis server is connected.""" try: return redis_conn.ping() except redis.exceptions.ConnectionError: return False # TODO: use a cache prefix? # TODO: key schema for easy access/queries class RedisCache: def clear(self): """Flush the cache.""" redis_conn.flushdb() redis_conn.config_resetstat() @staticmethod def size(): """Size of the Redis cache. .. note:: This is the size of the entire Redis database. """ return redis_conn.dbsize() def info(self): """Return cache information. .. note:: This is not the cache info for the entire Redis key space. """ info = redis_conn.info() return _CacheInfo(info['keyspace_hits'], info['keyspace_misses'], self.size()) def get(self, key): """Get a value from the cache. Returns None if the key is not in the cache. """ value = redis_conn.get(key) if value is not None: value = pickle.loads(value) return value def set(self, key, value): """Set a value in the cache.""" value = pickle.dumps(value, protocol=constants.PICKLE_PROTOCOL) redis_conn.set(key, value) def key(self): """Delegate to subclasses.""" raise NotImplementedError def validate_parent_cache(parent_cache): # TODO: also validate that subsystem is a cut version of # parent_cache.subsystem? Do we need to check this at all? if parent_cache.subsystem.is_cut: raise ValueError("parent_cache must be from an uncut subsystem") class RedisMICECache(RedisCache): """A Redis-backed cache for |Subsystem.find_mice()|. See |MICECache| for more info. """ def __init__(self, subsystem, parent_cache=None): super().__init__() self.subsystem = subsystem self.subsystem_hash = hash(subsystem) if parent_cache is not None: validate_parent_cache(parent_cache) # Store the hash of the parent subsystem. We don't want to store # the parent subsystem explicitly so that it does not need to be # passed between processes. self.parent_subsystem_hash = parent_cache.subsystem_hash else: self.parent_subsystem_hash = None # TODO: if the value is found in the parent cache, store it in this # cache so we don't have to call `damaged_by_cut` over and over? def get(self, key): """Get a value from the cache. If the |MICE| cannot be found in this cache, try and find it in the parent cache. """ mice = super().get(key) if mice is not None: # Hit return mice # Try and get the key from the parent cache. if self.parent_subsystem_hash: parent_key = key.replace(str(self.subsystem_hash), str(self.parent_subsystem_hash), 1) mice = super().get(parent_key) if mice is not None and not mice.damaged_by_cut(self.subsystem): return mice return None def set(self, key, value): """Only need to set if the subsystem is uncut. Caches are only inherited from uncut subsystems. """ if not self.subsystem.is_cut: super().set(key, value) def key(self, direction, mechanism, purviews=False, _prefix=None): """Cache key. This is the call signature of |Subsystem.find_mice()|.""" return "subsys:{}:{}:{}:{}:{}".format( self.subsystem_hash, _prefix, direction, mechanism, purviews) class DictMICECache(DictCache): """A subsystem-local cache for |MICE| objects. See |MICECache| for more info. """ def __init__(self, subsystem, parent_cache=None): super().__init__() self.subsystem = subsystem if parent_cache is not None: validate_parent_cache(parent_cache) self._build(parent_cache) def _build(self, parent_cache): """Build the initial cache from the parent. Only include the |MICE| which are unaffected by the subsystem cut. A |MICE| is affected if either the cut splits the mechanism or splits the connections between the purview and mechanism """ for key, mice in parent_cache.cache.items(): if not mice.damaged_by_cut(self.subsystem): self.cache[key] = mice def set(self, key, mice): """Set a value in the cache. Only cache if: - The subsystem is uncut (caches are only inherited from uncut subsystems so there is no reason to cache on cut subsystems.) - |small_phi| > 0. Ideally we would cache all mice, but the size of the cache grows way too large, making parallel computations incredibly inefficient because the caches have to be passed between process. This will be changed once global caches are implemented. - Memory is not too full. """ if (not self.subsystem.is_cut and mice.phi > 0 and not memory_full()): self.cache[key] = mice def key(self, direction, mechanism, purviews=False, _prefix=None): """Cache key. This is the call signature of |Subsystem.find_mice()|.""" return (_prefix, direction, mechanism, purviews) def MICECache(subsystem, parent_cache=None): """Construct a |MICE| cache. Uses either a Redis-backed cache or a local dict cache on the object. Args: subsystem (Subsystem): The subsystem that this is a cache for. Kwargs: parent_cache (MICECache): The cache generated by the uncut version of ``subsystem``. Any cached |MICE| which are unaffected by the cut are reused in this cache. If None, the cache is initialized empty. """ if config.REDIS_CACHE: cls = RedisMICECache else: cls = DictMICECache return cls(subsystem, parent_cache=parent_cache) class PurviewCache(DictCache): """A network-level cache for possible purviews.""" def set(self, key, value): """Only set if purview caching is enabled""" if config.CACHE_POTENTIAL_PURVIEWS: self.cache[key] = value def method(cache_name, key_prefix=None): """Caching decorator for object-level method caches. Cache key generation is delegated to the cache. Args: cache_name (str): The name of the (already-instantiated) cache on the decorated object which should be used to store results of this method. *key_prefix: A constant to use as part of the cache key in addition to the method arguments. """ def decorator(func): if (func.__name__ in ['cause_repertoire', 'effect_repertoire'] and not config.CACHE_REPERTOIRES): return func @wraps(func) def wrapper(obj, *args, **kwargs): cache = getattr(obj, cache_name) # Delegate key generation key = cache.key(*args, _prefix=key_prefix, **kwargs) # Get cached value, or compute value = cache.get(key) if value is None: # miss value = func(obj, *args, **kwargs) cache.set(key, value) return value return wrapper return decorator