From 12cdbddd141215fa628038087115d9c2c74f12e7 Mon Sep 17 00:00:00 2001
From: Megvii Engine Team <megengine@megvii.com>
Date: Fri, 20 Aug 2021 18:14:23 +0800
Subject: [PATCH] fix(ci): clean fastrun cache in windows and macos ci

GitOrigin-RevId: d1a010287f5b9441f6813d0eebd5088e4da88ea3
---
 .../megengine/utils/persistent_cache.py       | 37 +++++++++++++------
 1 file changed, 26 insertions(+), 11 deletions(-)

diff --git a/imperative/python/megengine/utils/persistent_cache.py b/imperative/python/megengine/utils/persistent_cache.py
index b88e58d29..f33d46fbc 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()
-- 
GitLab