Skip to content

Commit 40adfeb

Browse files
varshaantfx-copybara
authored andcommitted
Check if a directory exists before deleting it. On some file systems like S3, an empty directory raises a NotFoundError.
PiperOrigin-RevId: 440493079
1 parent 34ae2dd commit 40adfeb

File tree

1 file changed

+4
-9
lines changed

1 file changed

+4
-9
lines changed

tensorflow_transform/beam/tft_beam_io/transform_fn_io.py

Lines changed: 4 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@
1515

1616
import os
1717

18-
from absl import logging
1918
import apache_beam as beam
2019
import tensorflow_transform as tft
2120
from tensorflow_transform import impl_helper
@@ -100,20 +99,16 @@ def expand(self, transform_fn):
10099
def publish_outputs(unused_element, metadata_source_path,
101100
transform_fn_source_path):
102101
import tensorflow as tf # pylint: disable=g-import-not-at-top
103-
# TF 2.6 split support for filesystems such as Amazon S3 out to the
104-
# `tensorflow_io` package. Hence, this import is needed wherever we touch
105-
# the filesystem.
106-
try:
107-
import tensorflow_io as _ # pytype: disable=import-error # pylint: disable=g-import-not-at-top
108-
except ModuleNotFoundError:
109-
logging.info('tensorflow_io is not available.')
110102
if not tf.io.gfile.exists(self._path):
111103
tf.io.gfile.makedirs(self._path)
112104

113105
tf.io.gfile.rename(metadata_source_path, metadata_path, overwrite=True)
114106
tf.io.gfile.rename(
115107
transform_fn_source_path, transform_fn_path, overwrite=True)
116-
tf.io.gfile.rmtree(base_temp_dir)
108+
# TODO(b/211615643): Remove the exists check once importing TFIO in S3
109+
# addresses NotFoundError.
110+
if tf.io.gfile.exists(base_temp_dir):
111+
tf.io.gfile.rmtree(base_temp_dir)
117112

118113
# TODO(KesterTong): Move this "must follows" logic into a tfx_bsl helper
119114
# function or into Beam.

0 commit comments

Comments
 (0)