Check if a directory exists before deleting it. On some file systems like S3, an empty directory raises a NotFoundError.

PiperOrigin-RevId: 440493079

Author: varshaan

tensorflow_transform/beam/tft_beam_io/transform_fn_io.py

@@ -15,7 +15,6 @@
 
 import os
 
-from absl import logging
 import apache_beam as beam
 import tensorflow_transform as tft
 from tensorflow_transform import impl_helper
@@ -100,20 +99,16 @@ def expand(self, transform_fn):
     def publish_outputs(unused_element, metadata_source_path,
                         transform_fn_source_path):
       import tensorflow as tf  # pylint: disable=g-import-not-at-top
-      # TF 2.6 split support for filesystems such as Amazon S3 out to the
-      # `tensorflow_io` package. Hence, this import is needed wherever we touch
-      # the filesystem.
-      try:
-        import tensorflow_io as _  # pytype: disable=import-error # pylint: disable=g-import-not-at-top
-      except ModuleNotFoundError:
-        logging.info('tensorflow_io is not available.')
       if not tf.io.gfile.exists(self._path):
         tf.io.gfile.makedirs(self._path)
 
       tf.io.gfile.rename(metadata_source_path, metadata_path, overwrite=True)
       tf.io.gfile.rename(
           transform_fn_source_path, transform_fn_path, overwrite=True)
-      tf.io.gfile.rmtree(base_temp_dir)
+      # TODO(b/211615643): Remove the exists check once importing TFIO in S3
+      # addresses NotFoundError.
+      if tf.io.gfile.exists(base_temp_dir):
+        tf.io.gfile.rmtree(base_temp_dir)
 
     # TODO(KesterTong): Move this "must follows" logic into a tfx_bsl helper
     # function or into Beam.