def test_resolve_payload_format_and_data_view_uri(
self,
payload_formats,
data_view_uris=None,
data_view_ids=None,
expected_payload_format=None,
expected_data_view_uri=None,
expected_error_type=None,
expected_error_msg_regex=None):
examples = []
if data_view_uris is None:
data_view_uris = [None] * len(payload_formats)
if data_view_ids is None:
data_view_ids = [None] * len(payload_formats)
for payload_format, data_view_uri, data_view_id in zip(
payload_formats, data_view_uris, data_view_ids):
artifact = standard_artifacts.Examples()
examples_utils.set_payload_format(artifact, payload_format)
if data_view_uri is not None:
artifact.set_string_custom_property(
constants.DATA_VIEW_URI_PROPERTY_KEY, data_view_uri)
if data_view_id is not None:
artifact.set_int_custom_property(
constants.DATA_VIEW_ID_PROPERTY_KEY, data_view_id)
examples.append(artifact)
if expected_error_type is None:
payload_format, data_view_uri = (
tfxio_utils.resolve_payload_format_and_data_view_uri(examples))
self.assertEqual(payload_format, expected_payload_format)
self.assertEqual(data_view_uri, expected_data_view_uri)
else:
with self.assertRaisesRegex(expected_error_type,
expected_error_msg_regex):
_ = tfxio_utils.resolve_payload_format_and_data_view_uri(
examples)
def _run_model_inference(
self,
data_spec: bulk_inferrer_pb2.DataSpec,
output_example_spec: bulk_inferrer_pb2.OutputExampleSpec,
examples: List[types.Artifact],
output_examples: Optional[types.Artifact],
inference_result: Optional[types.Artifact],
inference_endpoint: model_spec_pb2.InferenceSpecType,
) -> None:
"""Runs model inference on given examples data.
Args:
data_spec: bulk_inferrer_pb2.DataSpec instance.
output_example_spec: bulk_inferrer_pb2.OutputExampleSpec instance.
examples: List of `standard_artifacts.Examples` artifacts.
output_examples: Optional output `standard_artifacts.Examples` artifact.
inference_result: Optional output `standard_artifacts.InferenceResult`
artifact.
inference_endpoint: Model inference endpoint.
"""
example_uris = {}
for example_artifact in examples:
for split in artifact_utils.decode_split_names(
example_artifact.split_names):
if data_spec.example_splits:
if split in data_spec.example_splits:
example_uris[split] = artifact_utils.get_split_uri(
[example_artifact], split)
else:
example_uris[split] = artifact_utils.get_split_uri(
[example_artifact], split)
payload_format, _ = tfxio_utils.resolve_payload_format_and_data_view_uri(
examples)
tfxio_factory = tfxio_utils.get_tfxio_factory_from_artifact(
examples,
_TELEMETRY_DESCRIPTORS,
schema=None,
read_as_raw_records=True,
# We have to specify this parameter in order to create a RawRecord TFXIO
# but we won't use the RecordBatches so the column name of the raw
# records does not matter.
raw_record_column_name='unused')
if output_examples:
output_examples.split_names = artifact_utils.encode_split_names(
sorted(example_uris.keys()))
with self._make_beam_pipeline() as pipeline:
data_list = []
for split, example_uri in example_uris.items():
tfxio = tfxio_factory(
[io_utils.all_files_pattern(example_uri)])
assert isinstance(
tfxio, record_based_tfxio.RecordBasedTFXIO
), ('Unable to use TFXIO {} as it does not support reading raw records.'
.format(type(tfxio)))
# pylint: disable=no-value-for-parameter
data = (pipeline
| 'ReadData[{}]'.format(split) >>
tfxio.RawRecordBeamSource()
| 'RunInference[{}]'.format(split) >> _RunInference(
payload_format, inference_endpoint))
if output_examples:
output_examples_split_uri = artifact_utils.get_split_uri(
[output_examples], split)
logging.info('Path of output examples split `%s` is %s.',
split, output_examples_split_uri)
_ = (data
| 'WriteExamples[{}]'.format(split) >> _WriteExamples(
output_example_spec, output_examples_split_uri))
# pylint: enable=no-value-for-parameter
data_list.append(data)
if inference_result:
_ = (
data_list
|
'FlattenInferenceResult' >> beam.Flatten(pipeline=pipeline)
| 'WritePredictionLogs' >> beam.io.WriteToTFRecord(
os.path.join(inference_result.uri,
_PREDICTION_LOGS_FILE_NAME),
file_name_suffix='.gz',
coder=beam.coders.ProtoCoder(
prediction_log_pb2.PredictionLog)))
if output_examples:
logging.info('Output examples written to %s.', output_examples.uri)
if inference_result:
logging.info('Inference result written to %s.',
inference_result.uri)
def Do(self, input_dict: Dict[Text, List[types.Artifact]],
output_dict: Dict[Text, List[types.Artifact]],
exec_properties: Dict[Text, Any]) -> None:
"""TensorFlow Transform executor entrypoint.
This implements BaseExecutor.Do() and is invoked by orchestration systems.
This is not inteded for manual usage or further customization. Please use
the Transform() function which takes an input format with no artifact
dependency.
Args:
input_dict: Input dict from input key to a list of artifacts, including:
- input_data: A list of type `standard_artifacts.Examples` which
should contain two splits 'train' and 'eval'.
- schema: A list of type `standard_artifacts.Schema` which should
contain a single schema artifact.
output_dict: Output dict from key to a list of artifacts, including:
- transform_output: Output of 'tf.Transform', which includes an exported
Tensorflow graph suitable for both training and serving;
- transformed_examples: Materialized transformed examples, which
includes both 'train' and 'eval' splits.
exec_properties: A dict of execution properties, including either one of:
- module_file: The file path to a python module file, from which the
'preprocessing_fn' function will be loaded.
- preprocessing_fn: The module path to a python function that
implements 'preprocessing_fn'.
Returns:
None
"""
self._log_startup(input_dict, output_dict, exec_properties)
train_data_uri = artifact_utils.get_split_uri(input_dict[EXAMPLES_KEY],
'train')
eval_data_uri = artifact_utils.get_split_uri(input_dict[EXAMPLES_KEY],
'eval')
payload_format, data_view_uri = (
tfxio_utils.resolve_payload_format_and_data_view_uri(
input_dict[EXAMPLES_KEY]))
schema_file = io_utils.get_only_uri_in_dir(
artifact_utils.get_single_uri(input_dict[SCHEMA_KEY]))
transform_graph_uri = artifact_utils.get_single_uri(
input_dict[TRANSFORM_GRAPH_KEY])
transform_output = artifact_utils.get_single_uri(
output_dict[TRANSFORM_OUTPUT_KEY])
temp_path = os.path.join(transform_output,
_TEMP_DIR_IN_TRANSFORM_OUTPUT)
absl.logging.debug('Using temp path %s for tft.beam', temp_path)
materialize_output_paths = []
if output_dict.get(TRANSFORMED_EXAMPLES_KEY) is not None:
transformed_example_artifact = artifact_utils.get_single_instance(
output_dict[TRANSFORMED_EXAMPLES_KEY])
# TODO(b/161490287): move the split_names setting to executor for all
# components.
transformed_example_artifact.split_names = (
artifact_utils.encode_split_names(
artifact.DEFAULT_EXAMPLE_SPLITS))
transformed_train_output = artifact_utils.get_split_uri(
output_dict[TRANSFORMED_EXAMPLES_KEY], 'train')
transformed_eval_output = artifact_utils.get_split_uri(
output_dict[TRANSFORMED_EXAMPLES_KEY], 'eval')
materialize_output_paths = [
os.path.join(transformed_train_output,
_DEFAULT_TRANSFORMED_EXAMPLES_PREFIX),
os.path.join(transformed_eval_output,
_DEFAULT_TRANSFORMED_EXAMPLES_PREFIX)
]
def _GetCachePath(label, params_dict):
if label not in params_dict:
return None
else:
return artifact_utils.get_single_uri(params_dict[label])
label_inputs = {
'transform_graph_uri':
transform_graph_uri,
labels.COMPUTE_STATISTICS_LABEL:
False,
labels.SCHEMA_PATH_LABEL:
schema_file,
labels.EXAMPLES_DATA_FORMAT_LABEL:
payload_format,
labels.DATA_VIEW_LABEL:
data_view_uri,
labels.ANALYZE_DATA_PATHS_LABEL:
io_utils.all_files_pattern(train_data_uri),
labels.ANALYZE_PATHS_FILE_FORMATS_LABEL:
labels.FORMAT_TFRECORD,
labels.TRANSFORM_DATA_PATHS_LABEL: [
io_utils.all_files_pattern(train_data_uri),
io_utils.all_files_pattern(eval_data_uri)
],
labels.TRANSFORM_PATHS_FILE_FORMATS_LABEL:
[labels.FORMAT_TFRECORD, labels.FORMAT_TFRECORD],
labels.CUSTOM_CONFIG:
exec_properties.get('custom_config', None),
}
cache_input = _GetCachePath('cache_input_path', input_dict)
if cache_input is not None:
label_inputs[labels.CACHE_INPUT_PATH_LABEL] = cache_input
label_outputs = {
labels.TRANSFORM_METADATA_OUTPUT_PATH_LABEL: transform_output,
labels.TRANSFORM_MATERIALIZE_OUTPUT_PATHS_LABEL:
materialize_output_paths,
labels.TEMP_OUTPUT_LABEL: str(temp_path),
}
cache_output = _GetCachePath('cache_output_path', output_dict)
if cache_output is not None:
label_outputs[labels.CACHE_OUTPUT_PATH_LABEL] = cache_output
status_file = 'status_file' # Unused
self.Transform(label_inputs, label_outputs, status_file)
absl.logging.debug('Cleaning up temp path %s on executor success',
temp_path)
io_utils.delete_dir(temp_path)