Make sure that WriteTransformFn.publish_outputs is idempotent.

PiperOrigin-RevId: 700817208
tensorflow · Nov 27, 2024 · a14023c · a14023c
1 parent 9ab3ebc
commit a14023c
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/tensorflow_transform/beam/tft_beam_io/transform_fn_io.py b/tensorflow_transform/beam/tft_beam_io/transform_fn_io.py
@@ -102,9 +102,15 @@ def publish_outputs(unused_element, metadata_source_path,
       if not tf.io.gfile.exists(self._path):
         tf.io.gfile.makedirs(self._path)
 
+      if tf.io.gfile.exists(metadata_path):
+        tf.io.gfile.rmtree(metadata_path)
       tf.io.gfile.rename(metadata_source_path, metadata_path, overwrite=True)
+
+      if tf.io.gfile.exists(transform_fn_path):
+        tf.io.gfile.rmtree(transform_fn_path)
       tf.io.gfile.rename(
           transform_fn_source_path, transform_fn_path, overwrite=True)
+
       # TODO(b/211615643): Remove the exists check once importing TFIO in S3
       # addresses NotFoundError.
       if tf.io.gfile.exists(base_temp_dir):