def testGenerateEvent(self):
self.assertProtoEquals(
"""
type: INPUT
path {
steps {
key: 'key'
}
steps {
index: 1
}
}
artifact_id: 2
execution_id: 3
""",
event_lib.generate_event(event_type=metadata_store_pb2.Event.INPUT,
key='key',
index=1,
artifact_id=2,
execution_id=3))
def _create_artifact_and_event_pairs(
metadata_handler: metadata.Metadata,
artifact_dict: MutableMapping[Text, Sequence[types.Artifact]],
event_type: metadata_store_pb2.Event.Type,
) -> List[Tuple[metadata_store_pb2.Artifact, metadata_store_pb2.Event]]:
"""Creates a list of [Artifact, Event] tuples.
The result of this function will be used in a MLMD put_execution() call.
Args:
metadata_handler: A handler to access MLMD store.
artifact_dict: The source of artifacts to work on. For each artifact in the
dict, creates a tuple for that. Note that all artifacts of the same key in
the artifact_dict are expected to share the same artifact type.
event_type: The event type of the event to be attached to the artifact
Returns:
A list of [Artifact, Event] tuples
"""
result = []
for key, artifact_list in artifact_dict.items():
artifact_type = None
for index, artifact in enumerate(artifact_list):
# TODO(b/153904840): If artifact id is present, skip putting the artifact
# into the pair when MLMD API is ready.
event = event_lib.generate_event(event_type=event_type,
key=key,
index=index)
# Reuses already registered type in the same list whenever possible as
# the artifacts in the same list share the same artifact type.
if artifact_type:
assert artifact_type.name == artifact.artifact_type.name, (
'Artifacts under the same key should share the same artifact type.'
)
artifact_type = common_utils.register_type_if_not_exist(
metadata_handler, artifact.artifact_type)
artifact.set_mlmd_artifact_type(artifact_type)
result.append((artifact.mlmd_artifact, event))
return result