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