|
15 | 15 |
|
16 | 16 | import os
|
17 | 17 |
|
18 |
| -from absl import logging |
19 | 18 | import apache_beam as beam
|
20 | 19 | import tensorflow_transform as tft
|
21 | 20 | from tensorflow_transform import impl_helper
|
@@ -100,20 +99,16 @@ def expand(self, transform_fn):
|
100 | 99 | def publish_outputs(unused_element, metadata_source_path,
|
101 | 100 | transform_fn_source_path):
|
102 | 101 | 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.') |
110 | 102 | if not tf.io.gfile.exists(self._path):
|
111 | 103 | tf.io.gfile.makedirs(self._path)
|
112 | 104 |
|
113 | 105 | tf.io.gfile.rename(metadata_source_path, metadata_path, overwrite=True)
|
114 | 106 | tf.io.gfile.rename(
|
115 | 107 | 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) |
117 | 112 |
|
118 | 113 | # TODO(KesterTong): Move this "must follows" logic into a tfx_bsl helper
|
119 | 114 | # function or into Beam.
|
|
0 commit comments