def testGetValueFailed(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_value(tfx_value)
def get_executions(
metadata_handler: metadata.Metadata,
node: pipeline_pb2.PipelineNode) -> List[metadata_store_pb2.Execution]:
"""Returns all executions for the given pipeline node.
This finds all executions having the same set of contexts as the pipeline
node.
Args:
metadata_handler: A handler to access MLMD db.
node: The pipeline node for which to obtain executions.
Returns:
List of executions for the given node in MLMD db.
"""
# Get all the contexts associated with the node.
contexts = []
for context_spec in node.contexts.contexts:
context = metadata_handler.store.get_context_by_type_and_name(
context_spec.type.name, common_utils.get_value(context_spec.name))
if context is None:
# If no context is registered, it's certain that there is no
# associated execution for the node.
return []
contexts.append(context)
return execution_lib.get_executions_associated_with_all_contexts(
metadata_handler, contexts)
def _extract_properties(
execution: metadata_store_pb2.Execution) -> Dict[Text, types.Property]:
result = {}
for key, prop in itertools.chain(execution.properties.items(),
execution.custom_properties.items()):
result[key] = common_utils.get_value(prop)
return result
def _register_context_if_not_exist(
metadata_handler: metadata.Metadata,
context_spec: pipeline_pb2.ContextSpec,
) -> metadata_store_pb2.Context:
"""Registers a context if not exist, otherwise returns the existing one.
Args:
metadata_handler: A handler to access MLMD store.
context_spec: A pipeline_pb2.ContextSpec message that instructs registering
of a context.
Returns:
An MLMD context.
"""
context = _generate_context_proto(metadata_handler=metadata_handler,
context_spec=context_spec)
try:
[context_id] = metadata_handler.store.put_contexts([context])
context.id = context_id
except mlmd.errors.AlreadyExistsError:
context_name = common_utils.get_value(context_spec.name)
logging.debug('Context %s already exists.', context_name)
context = metadata_handler.store.get_context_by_type_and_name(
type_name=context_spec.type.name, context_name=context_name)
assert context is not None, (
'Context is missing for %s while put_contexts '
'reports that it existed.') % (context_spec.name)
logging.debug('ID of context %s is %s.', context_spec, context.id)
return context
def testGetValue(self):
tfx_value = pipeline_pb2.Value()
text_format.Parse(
"""
field_value {
int_value: 1
}""", tfx_value)
self.assertEqual(common_utils.get_value(tfx_value), 1)
def _resolve_single_channel(
metadata_handler: metadata.Metadata,
channel: pipeline_pb2.InputSpec.Channel) -> List[types.Artifact]:
"""Resolves input artifacts from a single channel."""
artifact_type = channel.artifact_query.type
output_key = channel.output_key or None
contexts = filter(None, [
metadata_handler.store.get_context_by_type_and_name(
context_query.type.name, common_utils.get_value(
context_query.name))
for context_query in channel.context_queries
])
return get_qualified_artifacts(metadata_handler=metadata_handler,
contexts=contexts,
artifact_type=artifact_type,
output_key=output_key)
def _resolve_single_channel(
metadata_handler: metadata.Metadata,
channel: pipeline_pb2.InputSpec.Channel) -> List[types.Artifact]:
"""Resolves input artifacts from a single channel."""
artifact_type = channel.artifact_query.type
output_key = channel.output_key or None
# 1. filter(None, list) filters "false" value out from the list
# 2. even if the filter() result is empty, its result is considered as "true"
# so turning it into a list explicitly.
contexts = list(filter(None, [
metadata_handler.store.get_context_by_type_and_name(
context_query.type.name, common_utils.get_value(context_query.name))
for context_query in channel.context_queries
]))
return get_qualified_artifacts(
metadata_handler=metadata_handler,
contexts=contexts,
artifact_type=artifact_type,
output_key=output_key)
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