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 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 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)))
