/
caching.py
170 lines (137 loc) · 4.66 KB
/
caching.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
import gc
import weakref
import sympy
from sympy import cache
__all__ = ['Cached', '_SymbolCache', 'CacheManager']
_SymbolCache = {}
"""The symbol cache."""
class AugmentedWeakRef(weakref.ref):
def __new__(cls, obj, meta):
obj = super().__new__(cls, obj)
obj.nbytes = meta.get('nbytes', 0)
return obj
class Cached(object):
"""
Mixin class for cached symbolic objects.
"""
@classmethod
def _cache_key(cls, *args, **kwargs):
"""
A unique, deterministic key from the input arguments.
Notes
-----
To be implemented by subclasses.
Returns
-------
The cache key. It must be hashable.
"""
raise NotImplementedError("Subclass must implement _cache_key")
@classmethod
def _cache_get(cls, key):
"""
Look up the cache for a given key.
Parameters
----------
key : object
The cache key. It must be hashable.
Returns
-------
The object if in the cache and alive, otherwise None.
"""
if key in _SymbolCache:
# There is indeed an object mapped to `key`. But is it still alive?
obj = _SymbolCache[key]()
if obj is None:
# Cleanup _SymbolCache (though practically unnecessary)
del _SymbolCache[key]
return None
else:
return obj
else:
return None
def __init__(self, key, *aliases):
"""
Store `self` in the symbol cache.
Parameters
----------
key : object
The cache key. It must be hashable.
*aliases
Additional keys to which self is mapped.
"""
# Precompute hash. This uniquely depends on the cache key
self._cache_key_hash = hash(key)
# Add ourselves to the symbol cache
awr = AugmentedWeakRef(self, self._cache_meta())
for i in (key,) + aliases:
_SymbolCache[i] = awr
def __init_cached__(self, key):
"""
Initialise `self` with a cached object state.
Parameters
----------
key : object
The cache key of the object whose state is used to initialize `self`.
It must be hashable.
"""
self.__dict__ = _SymbolCache[key]().__dict__.copy()
def __hash__(self):
"""
The hash value of a cached object is the hash of its cache key.
"""
return self._cache_key_hash
def _cache_meta(self):
"""
Metadata attached when ``self`` is added to the symbol cache.
Notes
-----
This should be specialized by the individual subclasses. This is useful
to implement callbacks to be executed upon eviction.
"""
return {}
class CacheManager(object):
"""
Drop unreferenced objects from the SymPy and Devito caches. The associated
data is lost (and thus memory is freed).
"""
gc_ths = 3*10**8
"""
The `clear` function will trigger garbage collection if at least one weak
reference points to an unreachable object whose size in bytes is greated
than the `gc_ths` value. Garbage collection is an expensive operation, so
we do it judiciously.
"""
force_ths = 100
"""
After `force_ths` *consecutive* calls ``clear(force=False)``, the flag
``force`` is ignored, and thus ``clear(force=True)`` is executed.
``
"""
ncalls_w_force_false = 0
@classmethod
def clear(cls, force=True):
# Wipe out the "true" SymPy cache
cache.clear_cache()
# Wipe out the hidden module-private SymPy caches
sympy.polys.rootoftools.ComplexRootOf.clear_cache()
sympy.polys.rings._ring_cache.clear()
sympy.polys.fields._field_cache.clear()
sympy.polys.domains.modularinteger._modular_integer_cache.clear()
# Maybe trigger garbage collection
if force is False:
if cls.ncalls_w_force_false + 1 == cls.force_ths:
# Case 1: too long since we called gc.collect, let's do it now
gc.collect()
cls.ncalls_w_force_false = 0
elif any(i.nbytes > cls.gc_ths for i in _SymbolCache.values()):
# Case 2: we got big objects in cache, we try to reclaim memory
gc.collect()
cls.ncalls_w_force_false = 0
else:
# We won't call gc.collect() this time
cls.ncalls_w_force_false += 1
else:
gc.collect()
for key, obj in list(_SymbolCache.items()):
if obj() is None:
del _SymbolCache[key]