Cache mechanism improvements:

theano/gof/cc.py

@@ -958,11 +958,30 @@ class CLinker(link.Linker):
 
     def compile_cmodule(self, location=None):
         """
-        This method is a callback for `ModuleCache.module_from_key`
+        Compile the module and return it.
+        """
+        # Go through all steps of the compilation process.
+        for step_result in self.compile_cmodule_by_step(location=location):
+            pass
+        # And return the output of the last step, which should be the module
+        # itself.
+        return step_result
+
+    def compile_cmodule_by_step(self, location=None):
+        """
+        This method is a callback for `ModuleCache.module_from_key`.
+
+        It is a generator (thus the 'by step'), so that:
+            - it first yields the module's C code
+            - it last yields the module itself
+            - it may yield other intermediate outputs in-between if needed
+              in the future (but this is not currently the case)
         """
         if location is None:
             location = cmodule.dlimport_workdir(config.compiledir)
         mod = self.build_dynamic_module()
+        src_code = mod.code()
+        yield src_code
         get_lock()
         try:
             debug("LOCATION", location)
@@ -970,7 +989,7 @@ class CLinker(link.Linker):
             libs = self.libraries()
             preargs = self.compile_args()
             if c_compiler.__name__=='nvcc_module_compile_str' and config.lib.amdlibm:
-                #this lib don't work correctly with nvcc in device code.
+                # This lib does not work correctly with nvcc in device code.
                 if '<amdlibm.h>' in mod.includes:
                     mod.includes.remove('<amdlibm.h>')
                 if '-DREPLACE_WITH_AMDLIBM' in preargs:
@@ -980,7 +999,7 @@ class CLinker(link.Linker):
             try:
                 module = c_compiler(
                     module_name=mod.name,
-                    src_code = mod.code(),
+                    src_code=src_code,
                     location=location,
                     include_dirs=self.header_dirs(),
                     lib_dirs=self.lib_dirs(),
@@ -992,8 +1011,7 @@ class CLinker(link.Linker):
         finally:
             release_lock()
 
-        return module
-
+        yield module
 
     def build_dynamic_module(self):
         """Return a cmodule.DynamicModule instance full of the code for our env.
@@ -1056,10 +1074,10 @@ class CLinker(link.Linker):
         except KeyError:
             key = None
         if key is None:
-            #if we can't get a key, then forget the cache mechanism
+            # If we can't get a key, then forget the cache mechanism.
             module = self.compile_cmodule()
         else:
-            module = get_module_cache().module_from_key(key=key, fn=self.compile_cmodule, keep_lock=keep_lock)
+            module = get_module_cache().module_from_key(key=key, fn=self.compile_cmodule_by_step, keep_lock=keep_lock)
 
         vars = self.inputs + self.outputs + self.orphans
         # List of indices that should be ignored when passing the arguments

theano/gof/cmodule.py

@@ -205,26 +205,22 @@ def module_name_from_dir(dirname):
     return os.path.join(dirname, name)
 
 
-def get_module_hash(module_file, key):
+def get_module_hash(src_code, key):
 
     """
     Return an MD5 hash that uniquely identifies a module.
 
     This hash takes into account:
-        1. The 'mod.cpp' or 'mod.cu' file used to compile `module_file`.
+        1. The C source code of the module (`src_code`).
         2. The version part of the key.
         3. The compiler options defined in `key` (command line parameters and
            libraries to link against).
     """
-    source_code = os.path.join(os.path.dirname(module_file), 'mod.cpp')
-    if not os.path.exists(source_code):
-        source_code = os.path.join(os.path.dirname(module_file), 'mod.cu')
-        assert os.path.exists(source_code)
     # `to_hash` will contain any element such that we know for sure that if
     # it changes, then the module hash should be different.
     # We start with the source code itself (stripping blanks might avoid
     # recompiling after a basic indentation fix for instance).
-    to_hash = map(str.strip, open(source_code).readlines())
+    to_hash = map(str.strip, src_code.split('\n'))
     # Get the version part of the key.
     to_hash += map(str, key[0])
     c_link_key = key[1]
@@ -477,8 +473,8 @@ 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).
-                            debug('Deleting deprecated cache entry', key_pkl)
-                            _rmtree(root, ignore_nocleanup=True)
+                            _rmtree(root, ignore_nocleanup=True,
+                                    msg='deprecated cache entry')
                             continue
 
                         # Find unversioned keys.
@@ -571,9 +567,12 @@ class ModuleCache(object):
 
     def module_from_key(self, key, fn=None, keep_lock=False):
         """
-        :param fn: a callable object that will return a module for the key (it is called only if the key isn't in
-        the cache).  This function will be called with a single keyword argument "location"
-        that is a path on the filesystem wherein the function should write the module.
+        :param fn: A callable object that will return an iterable object when
+        called, such that the first element in this iterable object is the
+        source code of the module, and the last element is the module itself.
+        `fn` is called only if the key is not already in the cache, with
+        a single keyword argument `location` that is the path to the directory
+        where the module should be compiled.
         """
         rval = None
         try:
@@ -606,15 +605,46 @@ class ModuleCache(object):
             try:
                 location = dlimport_workdir(self.dirname)
                 #debug("LOCATION*", location)
+
+                compile_steps = fn(location=location).__iter__()
+
+                # Check if we already know a module with the same hash. If we
+                # do, then there is no need to even compile it.
+                duplicated_module = False
+                # The first compilation step is to yield the source code.
+                src_code = compile_steps.next()
+                module_hash = get_module_hash(src_code, key)
+                if module_hash in self.module_hash_to_key_data:
+                    debug("Duplicated module! Will re-use the previous one")
+                    duplicated_module = True
+                    # Load the already existing module.
+                    key_data = self.module_hash_to_key_data[module_hash]
+                    # Note that we do not pass the `fn` argument, since it
+                    # should not be used considering that the module should
+                    # already be compiled.
+                    module = self.module_from_key(
+                            key=key_data.keys.__iter__().next())
+                    name = module.__file__
+                    # Add current key to the set of keys associated to the same
+                    # module.
+                    key_data.add_key(key)
+                    # We can delete the work directory.
+                    _rmtree(location, ignore_nocleanup=True)
+                else:
                     try:
-                    module = fn(location=location)  # WILL FAIL FOR BAD C CODE
+                        # Will fail if there is an error compiling the C code.
+                        while True:
+                            try:
+                                # The module should be returned by the last
+                                # step of the compilation.
+                                module = compile_steps.next()
+                            except StopIteration:
+                                break
                     except Exception, e:
                         _rmtree(location)
-                    #try:
-                    #except Exception, ee:
-                        #error('failed to cleanup location', location, ee)
                         raise
 
+                    # Obtain path to the '.so' module file.
                     name = module.__file__
                     
                     debug("Adding module to cache", key, name)
@@ -626,30 +656,12 @@ class ModuleCache(object):
                     assert hash(key) == hash_key
                     assert key not in self.entry_from_key
 
-                # Check if we already know a module with the same hash.
-                duplicated_module = False
-                module_hash = get_module_hash(name, key)
-                if module_hash in self.module_hash_to_key_data:
-                    debug("Duplicated module! Will re-use the previous one")
-                    duplicated_module = True
-                    # Load the already existing module.
-                    key_data = self.module_hash_to_key_data[module_hash]
-                    module = self.module_from_key(
-                            key=key_data.keys.__iter__().next(),
-                            keep_lock=True)
-                    # Add current key to the set of keys associated to the same
-                    # module.
-                    key_data.add_key(key)
-                    # We can delete this module.
-                    debug("Deleting: ", os.path.dirname(name))
-                    shutil.rmtree(os.path.dirname(name))
-                    name = module.__file__
-
-                if not duplicated_module and _version: # save the key
+                    if _version: # save the key
                         key_pkl = os.path.join(location, 'key.pkl')
+                        assert not os.path.exists(key_pkl)
                         key_data = KeyData(
                                 keys=set([key]),
-                            module_hash=get_module_hash(name, key),
+                                module_hash=module_hash,
                                 key_pkl=key_pkl)
                         try:
                             key_data.save_pkl()
@@ -697,6 +709,11 @@ class ModuleCache(object):
                 else:
                     self.entry_from_key[k] = name
 
+            if name in self.module_from_name:
+                # May happen if we are re-using an existing module.
+                assert duplicated_module
+                assert self.module_from_name[name] is module
+            else:
                 self.module_from_name[name] = module
 
             self.stats[2] += 1
@@ -709,21 +726,17 @@ class ModuleCache(object):
 
     """The default age threshold for `clear_old` (in seconds)
     """
-    def clear_old(self, age_thresh_del=None, get_lock=True):
+    def clear_old(self, age_thresh_del=None):
         """
         Delete entries from the filesystem for cache entries that are too old.
 
         :param age_thresh_del: Dynamic modules whose last access time is more
         than ``age_thresh_del`` seconds ago will be erased. Defaults to 31-day
         age if not provided.
-
-        :param get_lock: If True, then this function acquires and releases the
-        lock on the compile dir.
         """
         if age_thresh_del is None:
             age_thresh_del = self.age_thresh_del
 
-        if get_lock:
         compilelock.get_lock()
         try:
             # update the age of modules that have been accessed by other processes
@@ -755,11 +768,9 @@ class ModuleCache(object):
                             del self.entry_from_key[key]
                     parent = os.path.dirname(entry)
                     assert parent.startswith(os.path.join(self.dirname, 'tmp'))
-                    info("clear_old removing cache dir", parent)
-                    _rmtree(parent)
+                    _rmtree(parent, msg='old cache directory', level='info')
 
         finally:
-            if get_lock:
             compilelock.release_lock()
 
     def clear(self, unversioned_min_age=None, clear_base_files=False):
@@ -776,21 +787,17 @@ class ModuleCache(object):
         """
         compilelock.get_lock()
         try:
-            self.clear_old(-1.0, get_lock=False)
-            self.clear_unversioned(min_age=unversioned_min_age, get_lock=False)
+            self.clear_old(-1.0)
+            self.clear_unversioned(min_age=unversioned_min_age)
             if clear_base_files:
-                self.clear_base_files(get_lock=False)
+                self.clear_base_files()
         finally:
             compilelock.release_lock()
 
-    def clear_base_files(self, get_lock=True):
+    def clear_base_files(self):
         """
         Delete base directories 'cuda_ndarray' and 'cutils_ext' if present.
-
-        :param get_lock: If True, then this function acquires then releases the
-        lock on the compile dir.
         """
-        if get_lock:
         compilelock.get_lock()
         try:
             for base_dir in ('cuda_ndarray', 'cutils_ext'):
@@ -802,10 +809,9 @@ class ModuleCache(object):
                     except:
                         warning('Could not delete %s' % to_delete)
         finally:
-            if get_lock:
             compilelock.release_lock()
 
-    def clear_unversioned(self, min_age=None, get_lock=True):
+    def clear_unversioned(self, min_age=None):
         """
         Delete unversioned dynamic modules.
         
@@ -814,15 +820,11 @@ class ModuleCache(object):
 
         :param min_age: Minimum age to be deleted, in seconds. Defaults to
         7-day age if not provided.
-
-        :param get_lock: If True, then this function acquires and releases the
-        lock on the compile dir.
         """
         if min_age is None:
             min_age = self.age_thresh_del_unversioned
         items_copy = list(self.entry_from_key.iteritems())
 
-        if get_lock:
         compilelock.get_lock()
 
         try:
@@ -839,8 +841,7 @@ class ModuleCache(object):
 
                     parent = os.path.dirname(entry)
                     assert parent.startswith(os.path.join(self.dirname, 'tmp'))
-                    info("clear_unversioned removing cache dir", parent)
-                    _rmtree(parent)
+                    _rmtree(parent, msg='unversioned', level='info')
 
             time_now = time.time()
             for filename in os.listdir(self.dirname):
@@ -860,23 +861,27 @@ class ModuleCache(object):
                         # take care of the clean-up.
                         if age > min_age:
                             info("clear_unversioned removing cache dir", filename)
-                            _rmtree(os.path.join(self.dirname, filename))
+                            _rmtree(os.path.join(self.dirname, filename),
+                                    msg='unversioned', level='info')
         finally:
-            if get_lock:
             compilelock.release_lock()
 
     def _on_atexit(self):
         # Note: no need to call refresh() since it is called by clear_old().
         compilelock.get_lock()
         try:
-            self.clear_old(get_lock=False)
-            self.clear_unversioned(get_lock=False)
+            self.clear_old()
+            self.clear_unversioned()
         finally:
             compilelock.release_lock()
 
-def _rmtree(parent, ignore_nocleanup=False):
+def _rmtree(parent, ignore_nocleanup=False, msg='', level='debug'):
     try:
         if ignore_nocleanup or not config.nocleanup:
+            log_msg = 'Deleting'
+            if msg:
+                log_msg += ' (%s)'
+            eval(level)('%s: %s' % (log_msg, parent))
             shutil.rmtree(parent)
     except Exception, e:
         # If parent still exists, mark it for deletion by a future refresh()