def _attach_artifact_properties(spec: pipeline_pb2.OutputSpec.ArtifactSpec,
artifact: types.Artifact):
"""Attaches properties of an artifact using ArtifactSpec."""
for key, value in spec.additional_properties.items():
if not value.HasField('field_value'):
raise RuntimeError('Property value is not a field_value for %s' %
key)
setattr(artifact, key,
data_types_utils.get_metadata_value(value.field_value))
for key, value in spec.additional_custom_properties.items():
if not value.HasField('field_value'):
raise RuntimeError('Property value is not a field_value for %s' %
key)
value_type = value.field_value.WhichOneof('value')
if value_type == 'int_value':
artifact.set_int_custom_property(key, value.field_value.int_value)
elif value_type == 'string_value':
artifact.set_string_custom_property(key,
value.field_value.string_value)
elif value_type == 'double_value':
artifact.set_float_custom_property(key,
value.field_value.double_value)
else:
raise RuntimeError(f'Unexpected value_type: {value_type}')
def _MarkPushed(self, model_push: types.Artifact, pushed_destination: Text,
pushed_version: Optional[Text] = None) -> None:
model_push.set_int_custom_property('pushed', 1)
model_push.set_string_custom_property(
_PUSHED_DESTINATION_KEY, pushed_destination)
if pushed_version is not None:
model_push.set_string_custom_property(_PUSHED_VERSION_KEY, pushed_version)
def set_file_format(examples: types.Artifact, file_format: str):
"""Sets the file format custom property for `examples`.
Args:
examples: A standard_artifacts.Examples artifact.
file_format: One of the file format that tfx_bsl understands.
"""
assert examples.type_name == standard_artifacts.Examples.TYPE_NAME, (
'examples must be of type standard_artifacts.Examples')
examples.set_string_custom_property(
example_gen_utils.FILE_FORMAT_PROPERTY_NAME, file_format)
def set_payload_format(examples: types.Artifact, payload_format: int):
"""Sets the payload format custom property for `examples`.
Args:
examples: A standard_artifacts.Examples artifact.
payload_format: One of the enums in example_gen_pb2.PayloadFormat.
"""
assert examples.type is standard_artifacts.Examples, (
'examples must be of type standard_artifacts.Examples')
examples.set_string_custom_property(
example_gen_utils.PAYLOAD_FORMAT_PROPERTY_NAME,
example_gen_pb2.PayloadFormat.Name(payload_format))
def _set_artifact_properties(artifact: types.Artifact,
properties: Optional[Dict[str, Any]],
custom_properties: Optional[Dict[str, Any]]):
"""Sets properties and custom_properties to the given artifact."""
if properties is not None:
for key, value in properties.items():
setattr(artifact, key, value)
if custom_properties is not None:
for key, value in custom_properties.items():
if isinstance(value, int):
artifact.set_int_custom_property(key, value)
elif isinstance(value, (str, bytes)):
artifact.set_string_custom_property(key, value)
else:
raise NotImplementedError(
f'Unexpected custom_property value type:{type(value)}')
