diff --git a/imperative/python/megengine/utils/persistent_cache.py b/imperative/python/megengine/utils/persistent_cache.py index b88e58d290cb6d93956165f23ee9ac1112f19e59..f33d46fbc37cfc8095e23a8485f3a824478d60cf 100644 --- a/imperative/python/megengine/utils/persistent_cache.py +++ b/imperative/python/megengine/utils/persistent_cache.py @@ -20,21 +20,27 @@ from ..version import __version__ class _FakeRedisConn: cache_file = None + _is_shelve = False + _dict = {} def __init__(self): - try: - from ..hub.hub import _get_megengine_home - - cache_dir = os.path.expanduser( - os.path.join(_get_megengine_home(), "persistent_cache") - ) - os.makedirs(cache_dir, exist_ok=True) - self.cache_file = os.path.join(cache_dir, "cache") - self._dict = shelve.open(self.cache_file) - self._is_shelve = True - except: + if os.getenv("MGE_FASTRUN_CACHE_TYPE") == "MEMORY": self._dict = {} self._is_shelve = False + else: + try: + from ..hub.hub import _get_megengine_home + + cache_dir = os.path.expanduser( + os.path.join(_get_megengine_home(), "persistent_cache") + ) + os.makedirs(cache_dir, exist_ok=True) + self.cache_file = os.path.join(cache_dir, "cache") + self._dict = shelve.open(self.cache_file) + self._is_shelve = True + except: + self._dict = {} + self._is_shelve = False def get(self, key): if self._is_shelve and isinstance(key, bytes): @@ -48,6 +54,10 @@ class _FakeRedisConn: self._dict[key] = val + def clear(self): + print("{} cache item in {} deleted".format(len(self._dict), self.cache_file)) + self._dict.clear() + def __del__(self): if self._is_shelve: self._dict.close() @@ -87,3 +97,8 @@ class PersistentCacheOnServer(_PersistentCache): key = self._make_key(category, key) self._prev_get_refkeep = conn.get(key) return self._prev_get_refkeep + + def clean(self): + conn = self._conn + if isinstance(conn, _FakeRedisConn): + conn.clear()