def _wait_for_lock_to_disappear(handle, lock_file, lock_file_timeout_sec):
"""Waits for the lock file to disappear.
The lock file was created by another process that is performing a download
into its own temporary directory. The name of this temp directory is
sha1(<module>).<uuid>.tmp where <uuid> comes from the lock file.
Args:
handle: The location from where a module is being download.
lock_file: Lock file created by another process downloading this module.
lock_file_timeout_sec: The amount of time to wait (in seconds) before we
can declare that the other downloaded has been
abandoned. The download is declared abandoned if
there is no file size change in the temporary
directory within the last 'lock_file_timeout_sec'.
"""
locked_tmp_dir_size = 0
locked_tmp_dir_size_check_time = time.time()
lock_file_content = None
while tf.compat.v1.gfile.Exists(lock_file):
try:
logging.log_every_n(
logging.INFO,
"Module '%s' already being downloaded by '%s'. Waiting.", 10,
handle, tf_utils.read_file_to_string(lock_file))
if (time.time() - locked_tmp_dir_size_check_time >
lock_file_timeout_sec):
# Check whether the holder of the current lock downloaded anything
# in its temporary directory in the last 'lock_file_timeout_sec'.
cur_locked_tmp_dir_size = _locked_tmp_dir_size(lock_file)
cur_lock_file_content = tf_utils.read_file_to_string(lock_file)
if (cur_locked_tmp_dir_size == locked_tmp_dir_size
and cur_lock_file_content == lock_file_content):
# There is was no data downloaded in the past
# 'lock_file_timeout_sec'. Steal the lock and proceed with the
# local download.
logging.warning("Deleting lock file %s due to inactivity.",
lock_file)
tf.compat.v1.gfile.Remove(lock_file)
break
locked_tmp_dir_size = cur_locked_tmp_dir_size
locked_tmp_dir_size_check_time = time.time()
lock_file_content = cur_lock_file_content
except tf.errors.NotFoundError:
# Lock file or temp directory were deleted during check. Continue
# to check whether download succeeded or we need to start our own
# download.
pass
finally:
time.sleep(5)
def _wait_for_lock_to_disappear(handle, lock_file, lock_file_timeout_sec):
"""Waits for the lock file to disappear.
The lock file was created by another process that is performing a download
into its own temporary directory. The name of this temp directory is
sha1(<module>).<uuid>.tmp where <uuid> comes from the lock file.
Args:
handle: The location from where a module is being download.
lock_file: Lock file created by another process downloading this module.
lock_file_timeout_sec: The amount of time to wait (in seconds) before we
can declare that the other downloaded has been
abandoned. The download is declared abandoned if
there is no file size change in the temporary
directory within the last 'lock_file_timeout_sec'.
"""
locked_tmp_dir_size = 0
locked_tmp_dir_size_check_time = time.time()
lock_file_content = None
while tf.gfile.Exists(lock_file):
try:
tf.logging.log_every_n(
tf.logging.INFO,
"Module '%s' already being downloaded by '%s'. Waiting.", 10,
handle, tf_utils.read_file_to_string(lock_file))
if (time.time() - locked_tmp_dir_size_check_time >
lock_file_timeout_sec):
# Check whether the holder of the current lock downloaded anything
# in its temporary directory in the last 'lock_file_timeout_sec'.
cur_locked_tmp_dir_size = _locked_tmp_dir_size(lock_file)
cur_lock_file_content = tf_utils.read_file_to_string(lock_file)
if (cur_locked_tmp_dir_size == locked_tmp_dir_size and
cur_lock_file_content == lock_file_content):
# There is was no data downloaded in the past
# 'lock_file_timeout_sec'. Steal the lock and proceed with the
# local download.
tf.logging.warning("Deleting lock file %s due to inactivity." %
lock_file)
tf.gfile.Remove(lock_file)
break
locked_tmp_dir_size = cur_locked_tmp_dir_size
locked_tmp_dir_size_check_time = time.time()
lock_file_content = cur_lock_file_content
except tf.errors.NotFoundError:
# Lock file or temp directory were deleted during check. Continue
# to check whether download succeeded or we need to start our own
# download.
pass
finally:
time.sleep(5)
def testModuleAlreadyDownloaded(self):
# Simulate the case when a rogue process finishes downloading a module
# right before the current process can perform a rename of a temp directory
# to a permanent module directory.
module_dir = os.path.join(self.get_temp_dir(), "module")
def fake_download_fn_with_rogue_behavior(handle, tmp_dir):
del handle, tmp_dir
# Create module directory
tf_v1.gfile.MakeDirs(module_dir)
tf_utils.atomic_write_string_to_file(
os.path.join(module_dir, "file"), "content", False)
self.assertEqual(
module_dir,
resolver.atomic_download("module", fake_download_fn_with_rogue_behavior,
module_dir))
self.assertEqual(tf_v1.gfile.ListDirectory(module_dir), ["file"])
self.assertFalse(tf_v1.gfile.Exists(resolver._lock_filename(module_dir)))
parent_dir = os.path.abspath(os.path.join(module_dir, ".."))
self.assertEqual(
sorted(tf_v1.gfile.ListDirectory(parent_dir)),
["module", "module.descriptor.txt"])
self.assertRegexpMatches(
tf_utils.read_file_to_string(
resolver._module_descriptor_file(module_dir)),
"Module: module\n"
"Download Time: .*\n"
"Downloader Hostname: %s .PID:%d." % (re.escape(socket.gethostname()),
os.getpid()))
def testWaitForLockToDisappear_DownloadOngoing(self):
module_dir = os.path.join(self.get_temp_dir(), "module")
task_uid = uuid.uuid4().hex
lock_filename = resolver._lock_filename(module_dir)
lock_file_content = resolver._lock_file_contents(task_uid)
tf_utils.atomic_write_string_to_file(
lock_filename, lock_file_content, overwrite=False)
lock_expiration_wait_time_secs = 10
thread = threading.Thread(
target=resolver._wait_for_lock_to_disappear,
args=(
"module",
lock_filename,
lock_expiration_wait_time_secs,
))
thread.start()
# Simulate download by writing a file every 1 sec. While writes are happing
# the lock file remains in place.
tmp_dir = resolver._temp_download_dir(self.get_temp_dir(), task_uid)
tf_v1.gfile.MakeDirs(tmp_dir)
for x in range(2 * lock_expiration_wait_time_secs):
tf_utils.atomic_write_string_to_file(
os.path.join(tmp_dir, "file_%d" % x), "test", overwrite=False)
# While writes are happening the original lock file is in place.
self.assertEqual(lock_file_content,
tf_utils.read_file_to_string(lock_filename))
time.sleep(1)
thread.join(lock_expiration_wait_time_secs)
def testModuleDownloadedWhenEmptyFolderExists(self):
# Simulate the case when a module is cached in /tmp/module_dir but module
# files inside the folder are deleted. In this case, the download should
# still be conducted.
module_dir = os.path.join(self.get_temp_dir(), "module")
def fake_download_fn(handle, tmp_dir):
del handle, tmp_dir
tf.compat.v1.gfile.MakeDirs(module_dir)
tf_utils.atomic_write_string_to_file(
os.path.join(module_dir, "file"), "content", False)
# Create an empty folder before downloading.
self.assertFalse(tf.compat.v1.gfile.Exists(module_dir))
tf.compat.v1.gfile.MakeDirs(module_dir)
self.assertEqual(
module_dir,
resolver.atomic_download("module", fake_download_fn, module_dir))
self.assertEqual(tf.compat.v1.gfile.ListDirectory(module_dir),
["file"])
self.assertFalse(
tf.compat.v1.gfile.Exists(resolver._lock_filename(module_dir)))
parent_dir = os.path.abspath(os.path.join(module_dir, ".."))
self.assertEqual(sorted(tf.compat.v1.gfile.ListDirectory(parent_dir)),
["module", "module.descriptor.txt"])
self.assertRegexpMatches(
tf_utils.read_file_to_string(
resolver._module_descriptor_file(module_dir)),
"Module: module\n"
"Download Time: .*\n"
"Downloader Hostname: %s .PID:%d." %
(re.escape(socket.gethostname()), os.getpid()))
def testModuleAlreadyDownloaded(self):
# Simulate the case when a rogue process finishes downloading a module
# right before the current process can perform a rename of a temp directory
# to a permanent module directory.
module_dir = os.path.join(self.get_temp_dir(), "module")
def fake_download_fn_with_rogue_behavior(handle, tmp_dir):
del handle, tmp_dir
# Create module directory
tf.gfile.MakeDirs(module_dir)
tf_utils.atomic_write_string_to_file(
os.path.join(module_dir, "file"), "content", False)
self.assertEqual(
module_dir,
resolver.atomic_download("module", fake_download_fn_with_rogue_behavior,
module_dir))
self.assertEqual(tf.gfile.ListDirectory(module_dir), ["file"])
self.assertFalse(tf.gfile.Exists(resolver._lock_filename(module_dir)))
parent_dir = os.path.abspath(os.path.join(module_dir, ".."))
self.assertEqual(
sorted(tf.gfile.ListDirectory(parent_dir)),
["module", "module.descriptor.txt"])
self.assertRegexpMatches(
tf_utils.read_file_to_string(
resolver._module_descriptor_file(module_dir)),
"Module: module\n"
"Download Time: .*\n"
"Downloader Hostname: %s .PID:%d." % (re.escape(socket.gethostname()),
os.getpid()))
def testWaitForLockToDisappear_DownloadOngoing(self):
module_dir = os.path.join(self.get_temp_dir(), "module")
task_uid = uuid.uuid4().hex
lock_filename = resolver._lock_filename(module_dir)
lock_file_content = resolver._lock_file_contents(task_uid)
tf_utils.atomic_write_string_to_file(
lock_filename, lock_file_content, overwrite=False)
lock_expiration_wait_time_secs = 10
thread = threading.Thread(
target=resolver._wait_for_lock_to_disappear,
args=(
"module",
lock_filename,
lock_expiration_wait_time_secs,
))
thread.start()
# Simulate download by writing a file every 1 sec. While writes are happing
# the lock file remains in place.
tmp_dir = resolver._temp_download_dir(self.get_temp_dir(), task_uid)
tf.gfile.MakeDirs(tmp_dir)
for x in range(2 * lock_expiration_wait_time_secs):
tf_utils.atomic_write_string_to_file(
os.path.join(tmp_dir, "file_%d" % x), "test", overwrite=False)
# While writes are happening the original lock file is in place.
self.assertEqual(lock_file_content,
tf_utils.read_file_to_string(lock_filename))
time.sleep(1)
thread.join(lock_expiration_wait_time_secs)
def testModuleDescriptor(self):
FLAGS.tfhub_cache_dir = os.path.join(self.get_temp_dir(), "cache_dir")
http_resolver = compressed_module_resolver.HttpCompressedFileResolver()
path = http_resolver(self.module_handle)
desc = tf_utils.read_file_to_string(resolver._module_descriptor_file(path))
self.assertRegexpMatches(desc, "Module: %s\n"
"Download Time: .*\n"
"Downloader Hostname: %s .PID:%d." %
(re.escape(self.module_handle),
re.escape(socket.gethostname()), os.getpid()))
def testModuleDescriptor(self):
FLAGS.tfhub_cache_dir = os.path.join(self.get_temp_dir(), "cache_dir")
http_resolver = compressed_module_resolver.HttpCompressedFileResolver()
path = http_resolver(self.module_handle)
desc = tf_utils.read_file_to_string(resolver._module_descriptor_file(path))
self.assertRegexpMatches(desc, "Module: %s\n"
"Download Time: .*\n"
"Downloader Hostname: %s .PID:%d." %
(re.escape(self.module_handle),
re.escape(socket.gethostname()), os.getpid()))
def testModuleAlreadyDownloaded(self):
# Simulate the case when a rogue process finishes downloading a module
# right before the current process can perform a rename of a temp directory
# to a permanent module directory.
module_dir = os.path.join(self.get_temp_dir(), "module")
def fake_download_fn_with_rogue_behavior(handle, tmp_dir):
del handle, tmp_dir
# Create module directory
tf.compat.v1.gfile.MakeDirs(module_dir)
tf_utils.atomic_write_string_to_file(
os.path.join(module_dir, "file"), "content", False)
self.assertEqual(
module_dir,
resolver.atomic_download("module",
fake_download_fn_with_rogue_behavior,
module_dir))
self.assertEqual(tf.compat.v1.gfile.ListDirectory(module_dir),
["file"])
self.assertFalse(
tf.compat.v1.gfile.Exists(resolver._lock_filename(module_dir)))
parent_dir = os.path.abspath(os.path.join(module_dir, ".."))
self.assertEqual(sorted(tf.compat.v1.gfile.ListDirectory(parent_dir)),
["module", "module.descriptor.txt"])
self.assertRegexpMatches(
tf_utils.read_file_to_string(
resolver._module_descriptor_file(module_dir)),
"Module: module\n"
"Download Time: .*\n"
"Downloader Hostname: %s .PID:%d." %
(re.escape(socket.gethostname()), os.getpid()))
# Try downloading the model again. Mock
# tf_utils.atomic_write_string_to_file() to throw an exception. Since the
# model is already downloaded, the function will never get called and the
# download succeeds.
with mock.patch.object(
tf_utils,
"atomic_write_string_to_file",
side_effect=ValueError("This error should never be raised!")):
self.assertEqual(
module_dir,
resolver.atomic_download("module",
fake_download_fn_with_rogue_behavior,
module_dir))
self.assertEqual(tf.compat.v1.gfile.ListDirectory(module_dir),
["file"])
self.assertFalse(
tf.compat.v1.gfile.Exists(resolver._lock_filename(module_dir)))
def atomic_download(handle,
download_fn,
module_dir,
lock_file_timeout_sec=10 * 60):
"""Returns the path to a Module directory for a given TF-Hub Module handle.
Args:
handle: (string) Location of a TF-Hub Module.
download_fn: Callback function that actually performs download. The callback
receives two arguments, handle and the location of a temporary
directory to download the content into.
module_dir: Directory where to download the module files to.
lock_file_timeout_sec: The amount of time we give the current holder of
the lock to make progress in downloading a module.
If no progress is made, the lock is revoked.
Returns:
A string containing the path to a TF-Hub Module directory.
Raises:
ValueError: if the Module is not found.
tf.errors.OpError: file I/O failures raise the appropriate subtype.
"""
lock_file = _lock_filename(module_dir)
task_uid = uuid.uuid4().hex
lock_contents = _lock_file_contents(task_uid)
tmp_dir = _temp_download_dir(module_dir, task_uid)
# Function to check whether model has already been downloaded.
check_module_exists = lambda: (tf.compat.v1.gfile.Exists(module_dir) and tf
.compat.v1.gfile.ListDirectory(module_dir))
# Check whether the model has already been downloaded before locking
# the destination path.
if check_module_exists():
return module_dir
# Attempt to protect against cases of processes being cancelled with
# KeyboardInterrupt by using a try/finally clause to remove the lock
# and tmp_dir.
try:
while True:
try:
tf_utils.atomic_write_string_to_file(lock_file,
lock_contents,
overwrite=False)
# Must test condition again, since another process could have created
# the module and deleted the old lock file since last test.
if check_module_exists():
# Lock file will be deleted in the finally-clause.
return module_dir
if tf.compat.v1.gfile.Exists(module_dir):
tf.compat.v1.gfile.DeleteRecursively(module_dir)
break # Proceed to downloading the module.
# These errors are believed to be permanent problems with the
# module_dir that justify failing the download.
except (tf.errors.NotFoundError, tf.errors.PermissionDeniedError,
tf.errors.UnauthenticatedError,
tf.errors.ResourceExhaustedError, tf.errors.InternalError,
tf.errors.InvalidArgumentError,
tf.errors.UnimplementedError):
raise
# All other errors are retried.
# TODO(b/144424849): Retrying an AlreadyExistsError from the atomic write
# should be good enough, but see discussion about misc filesystem types.
# TODO(b/144475403): How atomic is the overwrite=False check?
except tf.errors.OpError:
pass
# Wait for lock file to disappear.
_wait_for_lock_to_disappear(handle, lock_file,
lock_file_timeout_sec)
# At this point we either deleted a lock or a lock got removed by the
# owner or another process. Perform one more iteration of the while-loop,
# we would either terminate due tf.compat.v1.gfile.Exists(module_dir) or
# because we would obtain a lock ourselves, or wait again for the lock to
# disappear.
# Lock file acquired.
logging.info("Downloading TF-Hub Module '%s'.", handle)
tf.compat.v1.gfile.MakeDirs(tmp_dir)
download_fn(handle, tmp_dir)
# Write module descriptor to capture information about which module was
# downloaded by whom and when. The file stored at the same level as a
# directory in order to keep the content of the 'model_dir' exactly as it
# was define by the module publisher.
#
# Note: The descriptor is written purely to help the end-user to identify
# which directory belongs to which module. The descriptor is not part of the
# module caching protocol and no code in the TF-Hub library reads its
# content.
_write_module_descriptor_file(handle, module_dir)
try:
tf.compat.v1.gfile.Rename(tmp_dir, module_dir)
logging.info("Downloaded TF-Hub Module '%s'.", handle)
except tf.errors.AlreadyExistsError:
logging.warning("Module already exists in %s", module_dir)
finally:
try:
# Temp directory is owned by the current process, remove it.
tf.compat.v1.gfile.DeleteRecursively(tmp_dir)
except tf.errors.NotFoundError:
pass
try:
contents = tf_utils.read_file_to_string(lock_file)
except tf.errors.NotFoundError:
contents = ""
if contents == lock_contents:
# Lock file exists and is owned by this process.
try:
tf.compat.v1.gfile.Remove(lock_file)
except tf.errors.NotFoundError:
pass
return module_dir
def _task_uid_from_lock_file(lock_filename):
"""Returns task UID of the task that created a given lock file."""
lock = tf_utils.read_file_to_string(lock_filename)
return lock.split(".")[-1]
def atomic_download(handle,
download_fn,
module_dir,
lock_file_timeout_sec=10 * 60):
"""Returns the path to a Module directory for a given TF-Hub Module handle.
Args:
handle: (string) Location of a TF-Hub Module.
download_fn: Callback function that actually performs download. The callback
receives two arguments, handle and the location of a temporary
directory to download the content into.
module_dir: Directory where to download the module files to.
lock_file_timeout_sec: The amount of time we give the current holder of
the lock to make progress in downloading a module.
If no progress is made, the lock is revoked.
Returns:
A string containing the path to a TF-Hub Module directory.
Raises:
ValueError: if the Module is not found.
"""
lock_file = _lock_filename(module_dir)
task_uid = uuid.uuid4().hex
lock_contents = _lock_file_contents(task_uid)
tmp_dir = _temp_download_dir(module_dir, task_uid)
# Attempt to protect against cases of processes being cancelled with
# KeyboardInterrupt by using a try/finally clause to remove the lock
# and tmp_dir.
try:
while True:
try:
tf_utils.atomic_write_string_to_file(lock_file,
lock_contents,
overwrite=False)
# Must test condition again, since another process could have created
# the module and deleted the old lock file since last test.
if tf.gfile.Exists(module_dir):
# Lock file will be deleted in the finally-clause.
return module_dir
break # Proceed to downloading the module.
except tf.errors.OpError:
pass
# Wait for lock file to disappear.
_wait_for_lock_to_disappear(handle, lock_file,
lock_file_timeout_sec)
# At this point we either deleted a lock or a lock got removed by the
# owner or another process. Perform one more iteration of the while-loop,
# we would either terminate due tf.gfile.Exists(module_dir) or because we
# would obtain a lock ourselves, or wait again for the lock to disappear.
# Lock file acquired.
tf.logging.info("Downloading TF-Hub Module '%s'.", handle)
tf.gfile.MakeDirs(tmp_dir)
download_fn(handle, tmp_dir)
# Write module descriptor to capture information about which module was
# downloaded by whom and when. The file stored at the same level as a
# directory in order to keep the content of the 'model_dir' exactly as it
# was define by the module publisher.
#
# Note: The descriptor is written purely to help the end-user to identify
# which directory belongs to which module. The descriptor is not part of the
# module caching protocol and no code in the TF-Hub library reads its
# content.
_write_module_descriptor_file(handle, module_dir)
try:
tf.gfile.Rename(tmp_dir, module_dir)
tf.logging.info("Downloaded TF-Hub Module '%s'.", handle)
except tf.errors.AlreadyExistsError:
tf.logging.warning("Module already exists in %s" % module_dir)
finally:
try:
# Temp directory is owned by the current process, remove it.
tf.gfile.DeleteRecursively(tmp_dir)
except tf.errors.NotFoundError:
pass
try:
contents = tf_utils.read_file_to_string(lock_file)
except tf.errors.NotFoundError:
contents = ""
if contents == lock_contents:
# Lock file exists and is owned by this process.
try:
tf.gfile.Remove(lock_file)
except tf.errors.NotFoundError:
pass
return module_dir
def atomic_download(handle,
download_fn,
module_dir,
lock_file_timeout_sec=10 * 60):
"""Returns the path to a Module directory for a given TF-Hub Module handle.
Args:
handle: (string) Location of a TF-Hub Module.
download_fn: Callback function that actually performs download. The callback
receives two arguments, handle and the location of a temporary
directory to download the content into.
module_dir: Directory where to download the module files to.
lock_file_timeout_sec: The amount of time we give the current holder of
the lock to make progress in downloading a module.
If no progress is made, the lock is revoked.
Returns:
A string containing the path to a TF-Hub Module directory.
Raises:
ValueError: if the Module is not found.
"""
lock_file = _lock_filename(module_dir)
task_uid = uuid.uuid4().hex
lock_contents = _lock_file_contents(task_uid)
tmp_dir = _temp_download_dir(module_dir, task_uid)
# Attempt to protect against cases of processes being cancelled with
# KeyboardInterrupt by using a try/finally clause to remove the lock
# and tmp_dir.
try:
while True:
try:
tf_utils.atomic_write_string_to_file(lock_file, lock_contents,
overwrite=False)
# Must test condition again, since another process could have created
# the module and deleted the old lock file since last test.
if tf.gfile.Exists(module_dir):
# Lock file will be deleted in the finally-clause.
return module_dir
break # Proceed to downloading the module.
except tf.errors.OpError:
pass
# Wait for lock file to disappear.
_wait_for_lock_to_disappear(handle, lock_file, lock_file_timeout_sec)
# At this point we either deleted a lock or a lock got removed by the
# owner or another process. Perform one more iteration of the while-loop,
# we would either terminate due tf.gfile.Exists(module_dir) or because we
# would obtain a lock ourselves, or wait again for the lock to disappear.
# Lock file acquired.
tf.logging.info("Downloading TF-Hub Module '%s'.", handle)
tf.gfile.MakeDirs(tmp_dir)
download_fn(handle, tmp_dir)
# Write module descriptor to capture information about which module was
# downloaded by whom and when. The file stored at the same level as a
# directory in order to keep the content of the 'model_dir' exactly as it
# was define by the module publisher.
#
# Note: The descriptor is written purely to help the end-user to identify
# which directory belongs to which module. The descriptor is not part of the
# module caching protocol and no code in the TF-Hub library reads its
# content.
_write_module_descriptor_file(handle, module_dir)
try:
tf.gfile.Rename(tmp_dir, module_dir)
tf.logging.info("Downloaded TF-Hub Module '%s'.", handle)
except tf.errors.AlreadyExistsError:
tf.logging.warning("Module already exists in %s" % module_dir)
finally:
try:
# Temp directory is owned by the current process, remove it.
tf.gfile.DeleteRecursively(tmp_dir)
except tf.errors.NotFoundError:
pass
try:
contents = tf_utils.read_file_to_string(lock_file)
except tf.errors.NotFoundError:
contents = ""
if contents == lock_contents:
# Lock file exists and is owned by this process.
try:
tf.gfile.Remove(lock_file)
except tf.errors.NotFoundError:
pass
return module_dir
def _task_uid_from_lock_file(lock_filename):
"""Returns task UID of the task that created a given lock file."""
lock = tf_utils.read_file_to_string(lock_filename)
return lock.split(".")[-1]
