Cleanup refresh to make the cleanup optional and avoid taking the lock

theano/gof/cmodule.py

@@ -627,8 +627,7 @@ class ModuleCache(object):
             self.stats[0] += 1
         return self.module_from_name[name]
 
-    def refresh(self, age_thresh_use=None, delete_if_problem=False,
-                no_clean_empty=False):
+    def refresh(self, age_thresh_use=None, delete_if_problem=False, cleanup=True):
         """Update cache data by walking the cache directory structure.
 
         Load key.pkl files that have not been loaded yet.
@@ -638,12 +637,15 @@ class ModuleCache(object):
         :param age_thresh_use: Do not use modules olther than this.
         Defaults to self.age_thresh_use.
 
-        :param delete_if_problem: If True, cache entries that meet one of those
-        two conditions are deleted:
-            - Those for which unpickling the KeyData file fails with an
-              unknown exception.
+        :param delete_if_problem: If True, cache entries that meet one
+        of those two conditions are deleted:
+            - Those for which unpickling the KeyData file fails with
+              an unknown exception.
             - Duplicated modules, regardless of their age.
 
+        :param cleanup: Do a cleanup of the cache removing expired and
+        broken modules.
+
         :returns: a list of modules of age higher than age_thresh_use.
         """
         if age_thresh_use is None:
@@ -651,8 +653,11 @@ class ModuleCache(object):
         start_time = time.time()
         too_old_to_use = []
 
-        compilelock.get_lock()
-        try:
+        to_delete = []
+        def rmtree(*args, **kwargs):
+            if cleanup:
+                to_delete.append((args, kwargs))
+
         # add entries that are not in the entry_from_key dictionary
         time_now = time.time()
         # Go through directories in alphabetical order to ensure consistent
@@ -666,8 +671,8 @@ class ModuleCache(object):
             if not os.path.isdir(root):
                 continue
             files = os.listdir(root)
-                if (not files and not no_clean_empty) or 'delete.me' in files:
-                    _rmtree(root, ignore_nocleanup=True,
+            if not files or 'delete.me' in files:
+                rmtree(root, ignore_nocleanup=True,
                        msg="delete.me found in dir")
                 continue
             elif 'key.pkl' in files:
@@ -681,7 +686,7 @@ class ModuleCache(object):
                         _logger.warning("ModuleCache.refresh() Found key "
                                         "without dll in cache, deleting it. %s",
                                         key_pkl)
-                        _rmtree(root, ignore_nocleanup=True,
+                    rmtree(root, ignore_nocleanup=True,
                            msg="missing module file", level=logging.INFO)
                 continue
                 if (time_now - last_access_time(entry)) < age_thresh_use:
@@ -698,7 +703,7 @@ class ModuleCache(object):
                         # Happened once... not sure why (would be worth
                         # investigating if it ever happens again).
                         unpickle_failure()
-                            _rmtree(root, ignore_nocleanup=True,
+                        rmtree(root, ignore_nocleanup=True,
                                msg='broken cache directory [EOF]',
                                level=logging.WARNING)
                         continue
@@ -711,7 +716,7 @@ class ModuleCache(object):
                     except Exception:
                         unpickle_failure()
                         if delete_if_problem:
-                                _rmtree(root, ignore_nocleanup=True,
+                            rmtree(root, ignore_nocleanup=True,
                                    msg='broken cache directory',
                                    level=logging.INFO)
                         else:
@@ -731,7 +736,7 @@ class ModuleCache(object):
                         # do not know the config options that were used.
                         # As a result, we delete it instead (which is also
                         # simpler to implement).
-                            _rmtree(root, ignore_nocleanup=True,
+                        rmtree(root, ignore_nocleanup=True,
                                msg=(
                                 'invalid cache entry format -- this '
                                 'should not happen unless your cache '
@@ -754,7 +759,7 @@ class ModuleCache(object):
                             key_data.key_pkl = key_pkl
                         else:
                             # This is suspicious. Better get rid of it.
-                                _rmtree(root, ignore_nocleanup=True,
+                            rmtree(root, ignore_nocleanup=True,
                                    msg='module file path mismatch',
                                    level=logging.INFO)
                         continue
@@ -773,7 +778,7 @@ class ModuleCache(object):
                                 'Found a mix of unversioned and '
                                 'versioned keys for the same '
                                 'module %s', key_pkl)
-                            _rmtree(root, ignore_nocleanup=True,
+                        rmtree(root, ignore_nocleanup=True,
                                msg="unversioned key(s) in cache",
                                level=logging.INFO)
                         continue
@@ -788,9 +793,10 @@ class ModuleCache(object):
                         # Note that it is important to walk through
                         # directories in alphabetical order so as to make
                         # sure all new processes only use the first one.
+                        if cleanup:
                             age = time.time() - last_access_time(entry)
                             if delete_if_problem or age > self.age_thresh_del:
-                                _rmtree(root, ignore_nocleanup=True,
+                                rmtree(root, ignore_nocleanup=True,
                                        msg='duplicated module',
                                        level=logging.DEBUG)
                             else:
@@ -881,8 +887,10 @@ class ModuleCache(object):
                                         pkl_file_to_remove)
                     self.loaded_key_pkl.remove(pkl_file_to_remove)
 
-        finally:
-            compilelock.release_lock()
+        if to_delete:
+            with compilelock.lock_ctx():
+                for a, kw in to_delete:
+                    _rmtree(*a, **kw)
 
         _logger.debug('Time needed to refresh cache: %s',
                       (time.time() - start_time))
@@ -919,18 +927,13 @@ class ModuleCache(object):
                           "previous one")
             key_data = self.module_hash_to_key_data[module_hash]
             module = self._get_from_key(None, key_data)
-            lock_taken = False
+            with compilelock.lock_ctx(keep_lock=keep_lock):
                 try:
-                compilelock.get_lock()
-                lock_taken = True
                     key_data.add_key(key, save_pkl=bool(key[0]))
                     key_broken = False
                 except cPickle.PicklingError:
                     key_data.remove_key(key)
                     key_broken = True
-            finally:
-                if lock_taken and not keep_lock:
-                    compilelock.release_lock()
             if (key[0] and not key_broken and
                 self.check_for_broken_eq):
                 self.check_key(key, key_data.key_pkl)
@@ -1026,16 +1029,11 @@ class ModuleCache(object):
         if module is not None:
             return module
 
-        try:
-            # Compile the module since it's not cached
-            compilelock.get_lock()
-            lock_taken = True
+        with compilelock.lock_ctx(keep_lock=keep_lock):
             # Maybe somebody else compiled it for us while we
             # where waiting for the lock. Try to load it again
+            self.refresh(cleanup=False)
 
-            # Need no_clean_empty otherwise it deletes the workdir we
-            # created above.
-            self.refresh()
             module = self._get_from_key(key)
             if module is not None:
                 return module
@@ -1073,10 +1071,6 @@ class ModuleCache(object):
 
             key_data = self._add_to_cache(module, key, module_hash)
             self.module_hash_to_key_data[module_hash] = key_data
-        finally:
-            # Release lock if needed.
-            if not keep_lock and lock_taken:
-                compilelock.release_lock()
 
         self.stats[2] += 1
         return module
@@ -1159,8 +1153,7 @@ class ModuleCache(object):
         else:
             age_thresh_use = None
 
-        compilelock.get_lock()
-        try:
+        with compilelock.lock_ctx():
             # Update the age of modules that have been accessed by other
             # processes and get all module that are too old to use
             # (not loaded in self.entry_from_key).
@@ -1179,9 +1172,6 @@ class ModuleCache(object):
                 _rmtree(parent, msg='old cache directory', level=logging.INFO,
                         ignore_nocleanup=True)
 
-        finally:
-            compilelock.release_lock()
-
     def clear(self, unversioned_min_age=None, clear_base_files=False,
               delete_if_problem=False):
         """
@@ -1198,16 +1188,13 @@ class ModuleCache(object):
 
         :param delete_if_problem: See help of refresh() method.
         """
-        compilelock.get_lock()
-        try:
+        with compilelock.lock_ctx():
             self.clear_old(
                     age_thresh_del=-1.0,
                     delete_if_problem=delete_if_problem)
             self.clear_unversioned(min_age=unversioned_min_age)
             if clear_base_files:
                 self.clear_base_files()
-        finally:
-            compilelock.release_lock()
 
     def clear_base_files(self):
         """
@@ -1219,8 +1206,7 @@ class ModuleCache(object):
         rename them with the '.delete.me' extension, to mark them to be deleted
         next time we clear the cache.
         """
-        compilelock.get_lock()
-        try:
+        with compilelock.lock_ctx()
             for base_dir in ('cuda_ndarray', 'cutils_ext', 'lazylinker_ext',
                              'scan_perform'):
                 to_delete = os.path.join(self.dirname, base_dir + '.delete.me')
@@ -1238,8 +1224,6 @@ class ModuleCache(object):
                     except Exception:
                         _logger.warning('Could not move %s to %s',
                                         to_rename, to_delete)
-        finally:
-            compilelock.release_lock()
 
     def clear_unversioned(self, min_age=None):
         """
@@ -1254,9 +1238,8 @@ class ModuleCache(object):
         if min_age is None:
             min_age = self.age_thresh_del_unversioned
 
-        compilelock.get_lock()
+        with compilelock.lock_ctx():
             all_key_datas = self.module_hash_to_key_data.values()
-        try:
             for key_data in all_key_datas:
                 if not key_data.keys:
                     # May happen for broken versioned keys.
@@ -1329,17 +1312,12 @@ class ModuleCache(object):
                             _rmtree(os.path.join(self.dirname, filename),
                                     msg='old unversioned', level=logging.INFO,
                                     ignore_nocleanup=True)
-        finally:
-            compilelock.release_lock()
 
     def _on_atexit(self):
         # Note: no need to call refresh() since it is called by clear_old().
-        compilelock.get_lock()
-        try:
+        with compilelock.lock_ctx()
             self.clear_old()
             self.clear_unversioned()
-        finally:
-            compilelock.release_lock()
         _logger.debug('Time spent checking keys: %s',
                       self.time_spent_in_check_key)
 

theano/gof/compilelock.py

@@ -7,6 +7,7 @@ import random
 import socket  # only used for gethostname()
 import time
 import logging
+import contextlib
 
 from theano import config
 from theano.configparser import AddConfigVar, IntParam
@@ -44,6 +45,14 @@ def force_unlock():
     release_lock()
 
 
+@contextmanager
+def lock_ctx(lock_dir=None, keep_lock=False, **kw):
+    get_lock(lock_dir=lock_dir, **kw)
+    yield
+    if not keep_lock:
+        release_lock()
+
+
 def get_lock(lock_dir=None, **kw):
     """
     Obtain lock on compilation directory.