def testGetMetadataValueTypeFailed(self):
tfx_value = pipeline_pb2.Value()
text_format.Parse(
"""
runtime_parameter {
name: 'rp'
}""", tfx_value)
with self.assertRaisesRegex(RuntimeError,
'Expecting field_value but got'):
common_utils.get_metadata_value_type(tfx_value)
def prepare_execution(
metadata_handler: metadata.Metadata,
execution_type: metadata_store_pb2.ExecutionType,
state: metadata_store_pb2.Execution.State,
exec_properties: Optional[Mapping[Text, types.Property]] = None,
) -> metadata_store_pb2.Execution:
"""Creates an execution proto based on the information provided.
Args:
metadata_handler: A handler to access MLMD store.
execution_type: A metadata_pb2.ExecutionType message describing the type of
the execution.
state: The state of the execution.
exec_properties: Execution properties that need to be attached.
Returns:
A metadata_store_pb2.Execution message.
"""
execution = metadata_store_pb2.Execution()
execution.last_known_state = state
execution.type_id = common_utils.register_type_if_not_exist(
metadata_handler, execution_type).id
exec_properties = exec_properties or {}
# For every execution property, put it in execution.properties if its key is
# in execution type schema. Otherwise, put it in execution.custom_properties.
for k, v in exec_properties.items():
if (execution_type.properties.get(k) ==
common_utils.get_metadata_value_type(v)):
common_utils.set_metadata_value(execution.properties[k], v)
else:
common_utils.set_metadata_value(execution.custom_properties[k], v)
logging.debug('Prepared EXECUTION:\n %s', execution)
return execution
def testGetMetadataValueType(self):
tfx_value = pipeline_pb2.Value()
text_format.Parse(
"""
field_value {
int_value: 1
}""", tfx_value)
self.assertEqual(common_utils.get_metadata_value_type(tfx_value),
metadata_store_pb2.INT)
def _generate_context_proto(
metadata_handler: metadata.Metadata,
context_spec: pipeline_pb2.ContextSpec) -> metadata_store_pb2.Context:
"""Generates metadata_pb2.Context based on the ContextSpec message.
Args:
metadata_handler: A handler to access MLMD store.
context_spec: A pipeline_pb2.ContextSpec message that instructs registering
of a context.
Returns:
A metadata_store_pb2.Context message.
Raises:
RuntimeError: When actual property type does not match provided metadata
type schema.
"""
context_type = common_utils.register_type_if_not_exist(
metadata_handler, context_spec.type)
context_name = common_utils.get_value(context_spec.name)
assert isinstance(context_name, Text), 'context name should be string.'
context = metadata_store_pb2.Context(type_id=context_type.id,
name=context_name)
for k, v in context_spec.properties.items():
if k in context_type.properties:
actual_property_type = common_utils.get_metadata_value_type(v)
if context_type.properties.get(k) == actual_property_type:
common_utils.set_metadata_value(context.properties[k], v)
else:
raise RuntimeError(
'Property type %s different from provided metadata type property type %s for key %s'
%
(actual_property_type, context_type.properties.get(k), k))
else:
common_utils.set_metadata_value(context.custom_properties[k], v)
return context
def testGetMetadataValueTypePrimitiveValue(self):
self.assertEqual(common_utils.get_metadata_value_type(1),
metadata_store_pb2.INT)