Add incomplete_files method for when multiple files with the same incomplete prefix are written

PiperOrigin-RevId: 653384142

Author: The TensorFlow Datasets Authors

tensorflow_datasets/core/utils/__init__.py

@@ -42,6 +42,7 @@
 from tensorflow_datasets.core.utils.py_utils import get_class_url
 from tensorflow_datasets.core.utils.py_utils import has_sufficient_disk_space
 from tensorflow_datasets.core.utils.py_utils import incomplete_file
+from tensorflow_datasets.core.utils.py_utils import incomplete_files
 from tensorflow_datasets.core.utils.py_utils import indent
 from tensorflow_datasets.core.utils.py_utils import is_incomplete_file
 from tensorflow_datasets.core.utils.py_utils import is_notebook

tensorflow_datasets/core/utils/py_utils.py

@@ -310,12 +310,13 @@ def nullcontext(enter_result: T = None) -> Iterator[T]:
   yield enter_result
 
 
+def _tmp_file_prefix() -> str:
+  return f'{constants.INCOMPLETE_PREFIX}{uuid.uuid4().hex}'
+
+
 def _tmp_file_name(path: epath.PathLike) -> epath.Path:
   path = epath.Path(path)
-  return (
-      path.parent
-      / f'{constants.INCOMPLETE_PREFIX}{uuid.uuid4().hex}.{path.name}'
-  )
+  return path.parent / f'{_tmp_file_prefix()}.{path.name}'
 
 
 @contextlib.contextmanager
@@ -332,6 +333,25 @@ def incomplete_file(
     tmp_path.unlink(missing_ok=True)
 
 
+@contextlib.contextmanager
+def incomplete_files(
+    path: epath.Path,
+) -> Iterator[epath.Path]:
+  """Writes to path atomically, by writing to temp file and renaming it."""
+  tmp_file_prefix = _tmp_file_prefix()
+  tmp_path = path.parent / f'{tmp_file_prefix}.{path.name}'
+  try:
+    yield tmp_path
+    # Rename all tmp files to their final name.
+    for tmp_file in path.parent.glob(f'{tmp_file_prefix}.*'):
+      file_name = tmp_file.name.removeprefix(tmp_file_prefix + '.')
+      tmp_file.replace(path.parent / file_name)
+  finally:
+    # Eventually delete the tmp_path if exception was raised
+    for tmp_file in path.parent.glob(f'{tmp_file_prefix}.*'):
+      tmp_file.unlink(missing_ok=True)
+
+
 def is_incomplete_file(path: epath.Path) -> bool:
   """Returns whether the given filename suggests that it's incomplete."""
   return bool(

tensorflow_datasets/core/writer.py

@@ -483,7 +483,7 @@ def _write_final_shard(
     shard_path = self._filename_template.sharded_filepath(
         shard_index=shard_id, num_shards=len(non_empty_shard_ids)
     )
-    with utils.incomplete_file(epath.Path(shard_path)) as tmp_path:
+    with utils.incomplete_files(epath.Path(shard_path)) as tmp_path:
       logging.info(
           "Writing %d examples to %s.", len(example_by_key), os.fspath(tmp_path)
       )