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