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 testAbandondedLockFile(self):
# Tests that the caching procedure is resilient to an abandonded lock
# file.
FLAGS.tfhub_cache_dir = os.path.join(self.get_temp_dir(), "cache_dir")
# Create an "abandoned" lock file, i.e. a lock file with no process actively
# downloading anymore.
module_dir = compressed_module_resolver._module_dir(self.module_handle)
task_uid = uuid.uuid4().hex
lock_filename = resolver._lock_filename(module_dir)
tf_utils.atomic_write_string_to_file(
lock_filename,
resolver._lock_file_contents(task_uid),
overwrite=False)
with mock.patch.object(
compressed_module_resolver.HttpCompressedFileResolver,
"_lock_file_timeout_sec",
return_value=10):
http_resolver = compressed_module_resolver.HttpCompressedFileResolver(
)
handle = "http://localhost:%d/mock_module.tar.gz" % self.server_port
# After seeing the lock file is abandoned, this resolver will download the
# module and return a path to the extracted contents.
path = http_resolver(handle)
files = os.listdir(path)
self.assertListEqual(sorted(files), ["file1", "file2", "file3"])
self.assertFalse(tf.compat.v1.gfile.Exists(lock_filename))
def testDirSize(self):
fake_task_uid = 1234
# Create a directory with some files and sub-directory and check its size.
test_dir = resolver._temp_download_dir(self.get_temp_dir(), fake_task_uid)
tf.gfile.MakeDirs(test_dir)
tf_utils.atomic_write_string_to_file(
os.path.join(test_dir, "file1"), "content1", False)
tf_utils.atomic_write_string_to_file(
os.path.join(test_dir, "file2"), "content2", False)
test_sub_dir = os.path.join(test_dir, "sub_dir")
tf.gfile.MakeDirs(test_sub_dir)
tf_utils.atomic_write_string_to_file(
os.path.join(test_sub_dir, "file3"), "content3", False)
self.assertEqual(3 * 8, resolver._dir_size(test_dir))
self.assertEqual(8, resolver._dir_size(test_sub_dir))
# Treat the directory as a temporary directory used by a module download by
# referring to that directory from the lock file.
fake_lock_filename = resolver._lock_filename(self.get_temp_dir())
tf_utils.atomic_write_string_to_file(
fake_lock_filename, resolver._lock_file_contents(fake_task_uid), False)
self.assertEqual(3 * 8, resolver._locked_tmp_dir_size(fake_lock_filename))
# Check that if temp directory doesn't exist, 0 is returned.
tf.gfile.DeleteRecursively(test_dir)
self.assertEqual(0, resolver._locked_tmp_dir_size(fake_lock_filename))
def testDirSize(self):
fake_task_uid = 1234
# Create a directory with some files and sub-directory and check its size.
test_dir = resolver._temp_download_dir(self.get_temp_dir(), fake_task_uid)
tf_v1.gfile.MakeDirs(test_dir)
tf_utils.atomic_write_string_to_file(
os.path.join(test_dir, "file1"), "content1", False)
tf_utils.atomic_write_string_to_file(
os.path.join(test_dir, "file2"), "content2", False)
test_sub_dir = os.path.join(test_dir, "sub_dir")
tf_v1.gfile.MakeDirs(test_sub_dir)
tf_utils.atomic_write_string_to_file(
os.path.join(test_sub_dir, "file3"), "content3", False)
self.assertEqual(3 * 8, resolver._dir_size(test_dir))
self.assertEqual(8, resolver._dir_size(test_sub_dir))
# Treat the directory as a temporary directory used by a module download by
# referring to that directory from the lock file.
fake_lock_filename = resolver._lock_filename(self.get_temp_dir())
tf_utils.atomic_write_string_to_file(
fake_lock_filename, resolver._lock_file_contents(fake_task_uid), False)
self.assertEqual(3 * 8, resolver._locked_tmp_dir_size(fake_lock_filename))
# Check that if temp directory doesn't exist, 0 is returned.
tf_v1.gfile.DeleteRecursively(test_dir)
self.assertEqual(0, resolver._locked_tmp_dir_size(fake_lock_filename))
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 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 testReadTaskUidFromLockFile(self):
module_dir = os.path.join(self.get_temp_dir(), "module")
task_uid = uuid.uuid4().hex
lock_filename = resolver._lock_filename(module_dir)
tf_utils.atomic_write_string_to_file(
lock_filename, resolver._lock_file_contents(task_uid), overwrite=False)
self.assertEqual(task_uid, resolver._task_uid_from_lock_file(lock_filename))
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 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 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 testAbandondedLockFile(self):
# Tests that the caching procedure is resilient to an abandonded lock
# file.
FLAGS.tfhub_cache_dir = os.path.join(self.get_temp_dir(), "cache_dir")
# Create an "abandoned" lock file, i.e. a lock file with no process actively
# downloading anymore.
module_dir = compressed_module_resolver._module_dir(self.module_handle)
task_uid = uuid.uuid4().hex
lock_filename = resolver._lock_filename(module_dir)
tf_utils.atomic_write_string_to_file(lock_filename,
resolver._lock_file_contents(task_uid),
overwrite=False)
with mock.patch.object(
compressed_module_resolver.HttpCompressedFileResolver,
"_lock_file_timeout_sec",
return_value=10):
http_resolver = compressed_module_resolver.HttpCompressedFileResolver()
handle = "http://localhost:%d/mock_module.tar.gz" % self.server_port
# After seeing the lock file is abandoned, this resolver will download the
# module and return a path to the extracted contents.
path = http_resolver(handle)
files = os.listdir(path)
self.assertListEqual(sorted(files), ["file1", "file2", "file3"])
self.assertFalse(tf.gfile.Exists(lock_filename))
def testReadTaskUidFromLockFile(self):
module_dir = os.path.join(self.get_temp_dir(), "module")
task_uid = uuid.uuid4().hex
lock_filename = resolver._lock_filename(module_dir)
tf_utils.atomic_write_string_to_file(lock_filename,
resolver._lock_file_contents(task_uid),
overwrite=False)
self.assertEqual(task_uid, resolver._task_uid_from_lock_file(lock_filename))
def testWaitForLockToDisappear_DownloadCompletes(self):
module_dir = os.path.join(self.get_temp_dir(), "module")
task_uid = uuid.uuid4().hex
lock_filename = resolver._lock_filename(module_dir)
# Write lock file
tf_utils.atomic_write_string_to_file(lock_filename,
resolver._lock_file_contents(task_uid),
overwrite=False)
# Wait for the lock file to disappear (in a separate thread)
thread = threading.Thread(target=resolver._wait_for_lock_to_disappear,
args=("module", lock_filename, 600,))
thread.start()
# Delete the lock file.
tf.gfile.Remove(lock_filename)
thread.join(10)
def testWaitForLockToDisappear_DownloadCompletes(self):
module_dir = os.path.join(self.get_temp_dir(), "module")
task_uid = uuid.uuid4().hex
lock_filename = resolver._lock_filename(module_dir)
# Write lock file
tf_utils.atomic_write_string_to_file(lock_filename,
resolver._lock_file_contents(task_uid),
overwrite=False)
# Wait for the lock file to disappear (in a separate thread)
thread = threading.Thread(target=resolver._wait_for_lock_to_disappear,
args=("module", lock_filename, 600,))
thread.start()
# Delete the lock file.
tf_v1.gfile.Remove(lock_filename)
thread.join(10)
def testWaitForLockToDisappear_DownloadAborted(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)
tmp_dir = resolver._temp_download_dir(self.get_temp_dir(), task_uid)
tf.gfile.MakeDirs(tmp_dir)
thread = threading.Thread(target=resolver._wait_for_lock_to_disappear,
args=("module", lock_filename, 10,))
thread.start()
thread.join(30)
# Because nobody was writing to tmp_dir, the lock file got reclaimed by
# resolver._wait_for_lock_to_disappear.
self.assertFalse(tf.gfile.Exists(lock_filename))
def testWaitForLockToDisappear_DownloadAborted(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)
tmp_dir = resolver._temp_download_dir(self.get_temp_dir(), task_uid)
tf_v1.gfile.MakeDirs(tmp_dir)
thread = threading.Thread(target=resolver._wait_for_lock_to_disappear,
args=("module", lock_filename, 10,))
thread.start()
thread.join(30)
# Because nobody was writing to tmp_dir, the lock file got reclaimed by
# resolver._wait_for_lock_to_disappear.
self.assertFalse(tf_v1.gfile.Exists(lock_filename))
def kill_download(handle, tmp_dir): del handle, tmp_dir # Simulate lock loss by removing the lock. tf_v1.gfile.Remove(resolver._lock_filename(module_dir)) # Throw an error to simulate aborted download. raise OSError(download_aborted_msg)
def kill_download(handle, tmp_dir): del handle, tmp_dir # Simulate lock loss by removing the lock. tf.gfile.Remove(resolver._lock_filename(module_dir)) # Throw an error to simulate aborted download. raise OSError(download_aborted_msg)
def testLockFileName(self):
self.assertEquals("/a/b/c.lock", resolver._lock_filename("/a/b/c/"))
def testLockFileName(self):
self.assertEquals("/a/b/c.lock", resolver._lock_filename("/a/b/c/"))
