experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Ensure distributed_file_utils.remove_temp_dirpath() can be safely called multiple times.

Browse Source 
PiperOrigin-RevId: 308716361
Change-Id: I16ea406b81e78f7560e7015d3c9e2dbbe19df267
This commit is contained in:
Rick Chao 2020-04-27 16:00:16 -07:00 committed by TensorFlower Gardener
parent 690afb0e68
commit 85e3642cfe
2 changed files with 24 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
tensorflow/python/distribute/distributed_file_utils.py
View File

@ -64,8 +64,9 @@ def _is_temp_dir(dirpath, strategy):
def _get_temp_dir(dirpath, strategy):
if _is_temp_dir(dirpath, strategy):
return dirpath
temp_dir = os.path.join(dirpath, _get_base_dirpath(strategy))
temp_dir = dirpath
else:
temp_dir = os.path.join(dirpath, _get_base_dirpath(strategy))
file_io.recursive_create_dir_v2(temp_dir)
return temp_dir

21
tensorflow/python/distribute/distributed_file_utils_test.py
View File

@ -106,6 +106,27 @@ class DistributedFileUtilsTest(test.TestCase):
self.assertTrue(os.path.exists(file_to_write))
distributed_file_utils.remove_temp_dirpath(temp_dir, strategy)
self.assertFalse(os.path.exists(file_to_write))
self.assertFalse(os.path.exists(os.path.dirname(file_to_write)))
def testMultipleRemoveOrigDirPathIsFine(self):
temp_dir = self.get_temp_dir()
strategy = DistributedFileUtilsTest.MockedWorkerStrategy()
dir_to_write = distributed_file_utils.write_dirpath(temp_dir, strategy)
file_to_write = os.path.join(dir_to_write, 'tmp')
self._write_dummy_file(file_to_write)
distributed_file_utils.remove_temp_dirpath(temp_dir, strategy)
distributed_file_utils.remove_temp_dirpath(temp_dir, strategy)
distributed_file_utils.remove_temp_dirpath(temp_dir, strategy)
def testMultipleRemoveDirToWritePathIsFine(self):
temp_dir = self.get_temp_dir()
strategy = DistributedFileUtilsTest.MockedWorkerStrategy()
dir_to_write = distributed_file_utils.write_dirpath(temp_dir, strategy)
file_to_write = os.path.join(dir_to_write, 'tmp')
self._write_dummy_file(file_to_write)
distributed_file_utils.remove_temp_dirpath(dir_to_write, strategy)
distributed_file_utils.remove_temp_dirpath(dir_to_write, strategy)
distributed_file_utils.remove_temp_dirpath(dir_to_write, strategy)
if __name__ == '__main__':