def setUp(self):
super().setUp()
self.artifact_struct_dict = {
'a1':
text_format.Parse(
"""
elements {
artifact {
artifact {
id: 123
}
type {
name: 't1'
}
}
}
""", metadata_store_service_pb2.ArtifactStructList()),
'a2':
text_format.Parse(
"""
elements {
artifact {
artifact {
id: 456
}
type {
name: 't2'
}
}
}
""", metadata_store_service_pb2.ArtifactStructList())
}
self.artifact_dict = {
'a1': [
artifact_utils.deserialize_artifact(
metadata_store_pb2.ArtifactType(name='t1'),
metadata_store_pb2.Artifact(id=123))
],
'a2': [
artifact_utils.deserialize_artifact(
metadata_store_pb2.ArtifactType(name='t2'),
metadata_store_pb2.Artifact(id=456))
]
}
self.metadata_value_dict = {
'p0': metadata_store_pb2.Value(int_value=0),
'p1': metadata_store_pb2.Value(int_value=1),
'p2': metadata_store_pb2.Value(string_value='hello'),
'p3': metadata_store_pb2.Value(string_value='')
}
self.value_dict = {'p0': 0, 'p1': 1, 'p2': 'hello', 'p3': ''}
def _build_proto_artifact_dict(artifact_dict):
"""Build PythonExecutorExecutionInfo input/output artifact dicts."""
proto_dict = {}
for k, v in artifact_dict.items():
artifact_list = metadata_store_service_pb2.ArtifactStructList()
for artifact in v:
artifact_struct = metadata_store_service_pb2.ArtifactStruct(
artifact=metadata_store_service_pb2.ArtifactAndType(
artifact=artifact.mlmd_artifact,
type=artifact.artifact_type))
artifact_list.elements.append(artifact_struct)
proto_dict[k] = artifact_list
return proto_dict
def _build_proto_artifact_dict(
artifact_dict: Mapping[str, Iterable[types.Artifact]]
) -> Dict[str, metadata_store_service_pb2.ArtifactStructList]:
"""Builds PythonExecutorExecutionInfo input/output artifact dicts."""
result = {}
if not artifact_dict:
return result
for k, v in artifact_dict.items():
artifact_list = metadata_store_service_pb2.ArtifactStructList()
for artifact in v:
artifact_struct = metadata_store_service_pb2.ArtifactStruct(
artifact=metadata_store_service_pb2.ArtifactAndType(
artifact=artifact.mlmd_artifact, type=artifact.artifact_type))
artifact_list.elements.append(artifact_struct)
result[k] = artifact_list
return result