Skip to content

P
pytensor
提交 753ad528 authored 作者: Arnaud Bergeron's avatar Arnaud Bergeron
浏览文件
操作

Change the ModuleCache to only take the lock for actual compilation.

Also, once we acquire the lock do a refresh of the cache and try to pick up modules compiled by others.
上级 34a15a3f
隐藏空白字符变更
内嵌 并排
正在显示 包含 187 行增加 和 220 行删除
theano/gof/cmodule.py
浏览文件 @ 753ad528
...@@ -617,6 +617,16 @@ class ModuleCache(object): ...@@ -617,6 +617,16 @@ class ModuleCache(object):
Older modules will be deleted in ``clear_old``. Older modules will be deleted in ``clear_old``.
""" """
def _get_module(self, name):
if name not in self.module_from_name:
_logger.debug('loading name %s', name)
self.module_from_name[name] = dlimport(name)
self.stats[1] += 1
else:
_logger.debug('returning compiled module from cache %s', name)
self.stats[0] += 1
return self.module_from_name[name]
def refresh(self, age_thresh_use=None, delete_if_problem=False): def refresh(self, age_thresh_use=None, delete_if_problem=False):
"""Update cache data by walking the cache directory structure. """Update cache data by walking the cache directory structure.
...@@ -877,243 +887,200 @@ class ModuleCache(object): ...@@ -877,243 +887,200 @@ class ModuleCache(object):
return too_old_to_use return too_old_to_use
def module_from_key(self, key, fn=None, keep_lock=False, key_data=None): def _get_from_key(self, key, key_data=None):
""" """
:param fn: A callable object that will return an iterable object when Returns a module if the passed-in key is found in the cache
called, such that the first element in this iterable object is the and None otherwise.
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.
:param key_data: If not None, it should be a KeyData object and the May raise ValueError if the key is malformed.
key parameter should be None. In this case, we use the info from the
KeyData object to recover the module, rather than the key itself. Note
that this implies the module already exists (and may or may not have
already been loaded).
""" """
# We should only use one of the two ways to get a module. name = None
assert key_data is None or key is None
rval = None
if key is not None: if key is not None:
assert key_data is None
try: try:
_version, _rest = key _version, _rest = key
except (TypeError, ValueError): except (TypeError, ValueError):
raise ValueError( raise ValueError(
"Invalid key. key must have form (version, rest)", key) "Invalid key. key must have form (version, rest)", key)
name = None if name in self.entry_from_key:
if key is not None and key in self.entry_from_key: name = self.entry_from_key[key]
# We have seen this key either in this process or previously.
name = self.entry_from_key[key]
elif key_data is not None:
name = key_data.get_entry()
if name is not None:
# This is an existing module we can recover.
if name not in self.module_from_name:
_logger.debug('loading name %s', name)
self.module_from_name[name] = dlimport(name)
self.stats[1] += 1
else:
self.stats[0] += 1
_logger.debug('returning compiled module from cache %s', name)
rval = self.module_from_name[name]
else: else:
hash_key = hash(key) assert key_data is not None
key_data = None name = key_data.get_entry()
# We have never seen this key before. if name is None:
return None
# We acquire the lock later only if we were able to return self._get_module(name)
# generate C code. Otherwise, we would take the lock for ops
# that have only a perform(). def _get_from_hash(self, module_hash, key, keep_lock=False):
if module_hash in self.module_hash_to_key_data:
_logger.debug("Duplicated module! Will re-use the "
"previous one")
key_data = self.module_hash_to_key_data[module_hash]
module = self._get_from_key(None, key_data)
lock_taken = False lock_taken = False
# This try/finally block ensures that the lock is released once we
# are done writing in the cache file or after raising an exception.
try: try:
# Embedding two try statements for Python 2.4 compatibility compilelock.get_lock()
# (cannot do try / except / finally). lock_taken = True
try: key_data.add_key(key, save_pkl=bool(key[0]))
location = dlimport_workdir(self.dirname) key_broken = False
except OSError, e: except cPickle.PicklingError:
_logger.error(e) key_data.remove_key(key)
if e.errno == 31: key_broken = True
_logger.error('There are %i files in %s', finally:
len(os.listdir(config.compiledir)), if lock_taken and not keep_lock:
config.compiledir) compilelock.release_lock()
raise if (key[0] and not key_broken and
self.check_for_broken_eq):
self.check_key(key, key_data.key_pkl)
self._update_mappings(key, key_data, module.__file__)
return module
else:
return None
def _update_mappings(self, key, key_data, name):
all_keys = key_data.keys
if not all_keys:
all_keys = [key]
assert key in all_keys
for k in all_keys:
if k in self.entry_from_key:
assert self.entry_from_key[k] == name
else:
self.entry_from_key[k] = name
if key[0]:
self.similar_keys.setdefault(get_safe_part(k),
[]).append(key)
def _compile_code(self, compile_steps):
"""
Compiles the passed-in source code.
This expects that the compile lock is held during the call.
"""
location = None
try:
location = dlimport_workdir(self.dirname)
except OSError, e:
_logger.error(e)
if e.errno == 31:
_logger.error('There are %i files in %s',
len(os.listdir(config.compiledir)),
config.compiledir)
raise
try:
while True:
try: try:
compile_steps = fn(location=location).__iter__() # The module should be returned by the last
# step of the compilation.
# Check if we already know a module with the same hash. module = next(compile_steps)
# If we do, then there is no need to even compile it. except StopIteration:
duplicated_module = False break
# The first compilation step is to yield the source code. name = module.__file__
src_code = next(compile_steps) assert name.startswith(location)
module_hash = get_module_hash(src_code, key) assert name not in self.module_from_name
return module
# The op has c_code, so take the lock. except Exception:
compilelock.get_lock() _rmtree(location, ignore_if_missing=True,
lock_taken = True msg='exception during compilation')
raise
if not os.path.exists(location):
# Temporary fix, we should make sure it don't
# get deleted by the clear*() fct.
os.makedirs(location)
if module_hash in self.module_hash_to_key_data:
_logger.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=None,
key_data=key_data)
name = module.__file__
# Add current key to the set of keys associated to the
# same module. We only save the KeyData object of
# versioned modules.
try:
key_data.add_key(key, save_pkl=bool(_version))
key_broken = False
except cPickle.PicklingError:
# This should only happen if we tried to save the
# pickled file.
assert _version
# The key we are trying to add is broken: we will
# not add it after all.
key_data.remove_key(key)
key_broken = True
if (_version and not key_broken and
self.check_for_broken_eq):
self.check_key(key, key_data.key_pkl)
# We can delete the work directory.
_rmtree(location, ignore_nocleanup=True,
msg='temporary workdir of duplicated module')
else: def _add_to_cache(self, module, key, module_hash):
# Will fail if there is an error compiling the C code. """
# The exception will be caught and the work dir will be This function expects the compile lock to be held.
# deleted. """
while True: name = module.__file__
try: _logger.debug("Adding module to cache %s %s",
# The module should be returned by the last key, name)
# step of the compilation. # Changing the hash of the key is not allowed during
module = next(compile_steps) # compilation. That is the only cause found that makes
except StopIteration: # the following assert fail.
break assert key not in self.entry_from_key
# Obtain path to the '.so' module file. key_pkl = os.path.join(location, 'key.pkl')
name = module.__file__ assert not os.path.exists(key_pkl)
key_data = KeyData(
_logger.debug("Adding module to cache %s %s", keys=set([key]),
key, name) module_hash=module_hash,
assert name.startswith(location) key_pkl=key_pkl,
assert name not in self.module_from_name entry=name)
# Changing the hash of the key is not allowed during
# compilation. That is the only cause found that makes if key[0]:
# the following assert fail. try:
assert hash(key) == hash_key key_data.save_pkl()
assert key not in self.entry_from_key key_broken = False
except cPickle.PicklingError:
key_pkl = os.path.join(location, 'key.pkl') key_broken = True
assert not os.path.exists(key_pkl) key_data.remove_key(key)
key_data = KeyData( key_data.save_pkl()
keys=set([key]), if not key_broken and self.check_for_broken_eq:
module_hash=module_hash, self.check_key(key, key_pkl)
key_pkl=key_pkl, self.loaded_key_pkl.add(key_pkl)
entry=name) elif config.cmodule.warn_no_version:
key_flat = flatten(key)
# Note that we only save KeyData objects associated to ops = [k for k in key_flat if isinstance(k, theano.Op)]
# versioned modules. So for unversioned key, the _logger.warning("not all the"
# `key_pkl` field of the KeyData object will be a " following op(s) implement"
# non-existing file (which does not matter since it " c_code_cache_version(). This makes them"
# will not be accessed). " recompiled for each process." + str(ops))
if _version: self._update_mappings(key, key_data, module)
try: return key_data
key_data.save_pkl()
key_broken = False def module_from_key(self, key, fn=None, keep_lock=False):
except cPickle.PicklingError: """
key_broken = True :param fn: A callable object that will return an iterable object when
# Remove key from the KeyData object, to make called, such that the first element in this iterable object is the
# sure we never try to save it again. source code of the module, and the last element is the module itself.
# We still keep the KeyData object and save it `fn` is called only if the key is not already in the cache, with
# so that the module can be re-used in the a single keyword argument `location` that is the path to the directory
# future. where the module should be compiled.
key_data.keys = set() """
key_data.save_pkl() # Is the module in the cache?
module = self._get_from_key(key)
if not key_broken and self.check_for_broken_eq: if module is not None:
self.check_key(key, key_pkl) return module
# Adding the KeyData file to this set means it is a # Is the source code already in the cache?
# versioned module. compile_steps = fn(location=location).__iter__()
self.loaded_key_pkl.add(key_pkl) src_code = next(compile_steps)
elif config.cmodule.warn_no_version: module_hash = get_module_hash(src_code, key)
key_flat = flatten(key) module = self._get_from_hash(module_hash, key, keep_lock=keep_lock)
ops = [k for k in key_flat if module is not None:
if isinstance(k, theano.Op)] return module
_logger.warning("not all the"
" following op(s) implement" # Compile the module since it's not cached
" c_code_cache_version(). This makes them" try:
" recompiled for each process." + str(ops)) # The op has c_code, so take the lock.
compilelock.get_lock()
# Map the new module to its KeyData object. Note that lock_taken = True
# we need to do it regardless of whether the key is # Maybe somebody else compiled it for us while we
# versioned or not if we want to be able to re-use this # where waiting for the lock. Try to load it again
# module inside the same process. self.refresh()
self.module_hash_to_key_data[module_hash] = key_data module = self._get_from_key(key)
if module is not None:
return module
except Exception: module = self._get_from_hash(module_hash, key, keep_lock=keep_lock)
# This may happen e.g. when an Op has no C implementation. if module is not None:
# In any case, we do not want to keep around the temporary return module
# work directory, as it may cause trouble if we create too
# many of these. The 'ignore_if_missing' flag is set just
# in case this directory would have already been deleted.
_rmtree(location, ignore_if_missing=True,
msg=('exception -- '
'typically means no C implementation'))
raise
finally: hash_key = hash(key)
# Release lock if needed.
if not keep_lock and lock_taken:
compilelock.release_lock()
# Update map from key to module name for all keys associated to module = self._compile_module(compile_steps)
# this same module.
all_keys = key_data.keys
if not all_keys:
# Should only happen for broken keys.
assert key_broken
all_keys = [key]
else:
assert key in key_data.keys
for k in all_keys:
if k in self.entry_from_key:
# If we had already seen this key, then it should be
# associated to the same module.
assert self.entry_from_key[k] == name
else:
self.entry_from_key[k] = name
if _version:
self.similar_keys.setdefault(get_safe_part(k),
[]).append(key)
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 # Changing the hash of the key is not allowed during
rval = module # compilation.
#_logger.debug('stats %s %i', self.stats, sum(self.stats)) assert hash(key) == hash_key
return rval
key_data = self._add_to_cache(module, key)
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
def check_key(self, key, key_pkl): def check_key(self, key, key_pkl):
""" """
......
Markdown 格式
0% 或
您添加了 0 人 到此讨论。请谨慎行事。
请先完成此评论的编辑！
请 注册 或者 后发表评论