def export(self, estimator, export_path, checkpoint_path, eval_result,
is_the_final_export):
"""Exports the given `Estimator` to a specific format.
Performs the export as defined by the base_exporter and invokes all of the
specified rewriters.
Args:
estimator: the `Estimator` to export.
export_path: A string containing a directory where to write the export.
checkpoint_path: The checkpoint path to export.
eval_result: The output of `Estimator.evaluate` on this checkpoint.
is_the_final_export: This boolean is True when this is an export in the
end of training. It is False for the intermediate exports during the
training. When passing `Exporter` to `tf.estimator.train_and_evaluate`
`is_the_final_export` is always False if `TrainSpec.max_steps` is
`None`.
Returns:
The string path to the base exported directory or `None` if export is
skipped.
Raises:
RuntimeError: Unable to create a temporary rewrite directory.
"""
base_path = self._base_exporter.export(estimator, export_path,
checkpoint_path, eval_result,
is_the_final_export)
if not base_path:
return None
tmp_rewrite_folder = 'tmp-rewrite-' + str(int(time.time()))
tmp_rewrite_path = os.path.join(export_path, tmp_rewrite_folder)
if fileio.exists(tmp_rewrite_path):
raise RuntimeError(
'Unable to create a unique temporary rewrite path.')
fileio.makedirs(tmp_rewrite_path)
_invoke_rewriter(base_path, tmp_rewrite_path, self._rewriter_inst,
rewriter.ModelType.SAVED_MODEL,
rewriter.ModelType.ANY_MODEL)
fileio.rmtree(base_path)
fileio.rename(tmp_rewrite_path, base_path)
return base_path
def _check_pipeline_existence(self,
pipeline_name: Text,
required: bool = True) -> None:
"""Check if pipeline folder exists and if not, exit system.
Args:
pipeline_name: Name of the pipeline.
required: Set it as True if pipeline needs to exist else set it to False.
"""
handler_pipeline_path = os.path.join(self._handler_home_dir,
pipeline_name)
# Check if pipeline folder exists.
exists = fileio.exists(handler_pipeline_path)
if required and not exists:
# Check pipeline directory prior 0.25 and move files to the new location
# automatically.
old_handler_pipeline_path = os.path.join(
self._get_deprecated_handler_home(), pipeline_name)
if fileio.exists(old_handler_pipeline_path):
fileio.makedirs(os.path.dirname(handler_pipeline_path))
fileio.rename(old_handler_pipeline_path, handler_pipeline_path)
engine_flag = self.flags_dict[labels.ENGINE_FLAG]
handler_home_variable = engine_flag.upper() + '_HOME'
click.echo((
'[WARNING] Pipeline "{pipeline_name}" was found in "{old_path}", '
'but the location that TFX stores pipeline information was moved '
'since TFX 0.25.0.\n'
'[WARNING] Your files in "{old_path}" was automatically moved to '
'the new location, "{new_path}".\n'
'[WARNING] If you want to keep the files at the old location, set '
'`{handler_home}` environment variable to "{old_handler_home}".'
).format(pipeline_name=pipeline_name,
old_path=old_handler_pipeline_path,
new_path=handler_pipeline_path,
handler_home=handler_home_variable,
old_handler_home=self._get_deprecated_handler_home()),
err=True)
else:
sys.exit('Pipeline "{}" does not exist.'.format(pipeline_name))
elif not required and exists:
sys.exit('Pipeline "{}" already exists.'.format(pipeline_name))
def package_user_module_file(instance_name: Text, module_path: Text,
pipeline_root: Text) -> Tuple[Text, Text]:
"""Package the given user module file into a Python Wheel package.
Args:
instance_name: Name of the component instance, for creating a unique wheel
package name.
module_path: Path to the module file to be packaged.
pipeline_root: Text
Returns:
dist_file_path: Path to the generated wheel file.
user_module_path: Path for referencing the user module when stored
as the _MODULE_PATH_KEY execution property. Format should be treated
as opaque by the user.
Raises:
RuntimeError: When wheel building fails.
"""
module_path = os.path.abspath(io_utils.ensure_local(module_path))
if not module_path.endswith('.py'):
raise ValueError('Module path %r is not a ".py" file.' % module_path)
if not os.path.exists(module_path):
raise ValueError('Module path %r does not exist.' % module_path)
user_module_dir, module_file_name = os.path.split(module_path)
user_module_name = re.sub(r'\.py$', '', module_file_name)
source_files = []
# Discover all Python source files in this directory for inclusion.
for file_name in os.listdir(user_module_dir):
if file_name.endswith('.py'):
source_files.append(file_name)
module_names = []
for file_name in source_files:
if file_name in (_EPHEMERAL_SETUP_PY_FILE_NAME, '__init__.py'):
continue
module_name = re.sub(r'\.py$', '', file_name)
module_names.append(module_name)
# Set up build directory.
build_dir = tempfile.mkdtemp()
for source_file in source_files:
shutil.copyfile(os.path.join(user_module_dir, source_file),
os.path.join(build_dir, source_file))
# Generate an ephemeral wheel for this module.
logging.info(
'Generating ephemeral wheel package for %r (including modules: %s).',
module_path, module_names)
version_hash = _get_version_hash(user_module_dir, source_files)
logging.info('User module package has hash fingerprint version %s.',
version_hash)
setup_py_path = os.path.join(build_dir, _EPHEMERAL_SETUP_PY_FILE_NAME)
with open(setup_py_path, 'w') as f:
f.write(
_get_ephemeral_setup_py_contents(
'tfx-user-code-%s' % instance_name, '0.0+%s' % version_hash,
module_names))
temp_dir = tempfile.mkdtemp()
dist_dir = tempfile.mkdtemp()
bdist_command = [
sys.executable, setup_py_path, 'bdist_wheel', '--bdist-dir', temp_dir,
'--dist-dir', dist_dir
]
logging.info('Executing: %s', bdist_command)
try:
subprocess.check_call(bdist_command, cwd=build_dir)
except subprocess.CalledProcessError as e:
raise RuntimeError('Failed to build wheel.') from e
dist_files = os.listdir(dist_dir)
if len(dist_files) != 1:
raise RuntimeError(
'Unexpectedly found %d output files in wheel output directory %s.'
% (len(dist_files), dist_dir))
build_dist_file_path = os.path.join(dist_dir, dist_files[0])
# Copy wheel file atomically to wheel staging directory.
dist_wheel_directory = os.path.join(pipeline_root, '_wheels')
dist_file_path = os.path.join(dist_wheel_directory, dist_files[0])
temp_dist_file_path = dist_file_path + '.tmp'
fileio.makedirs(dist_wheel_directory)
fileio.copy(build_dist_file_path, temp_dist_file_path, overwrite=True)
fileio.rename(temp_dist_file_path, dist_file_path, overwrite=True)
logging.info(
('Successfully built user code wheel distribution at %r; target user '
'module is %r.'), dist_file_path, user_module_name)
# Encode the user module key as a specification of a user module name within
# a packaged wheel path.
assert '@' not in user_module_name, ('Unexpected invalid module name: %s' %
user_module_name)
user_module_path = '%s@%s' % (user_module_name, dist_file_path)
logging.info('Full user module path is %r', user_module_path)
return dist_file_path, user_module_path