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
def get_qualified_artifacts( self, contexts: List[metadata_store_pb2.Context], type_name: Text, producer_component_id: Optional[Text] = None, output_key: Optional[Text] = None, ) -> List[metadata_store_service_pb2.ArtifactAndType]: """Gets qualified artifacts that have the right producer info. Args: contexts: context constraints to filter artifacts type_name: type constraint to filter artifacts producer_component_id: producer constraint to filter artifacts output_key: output key constraint to filter artifacts Returns: A list of ArtifactAndType, containing qualified artifacts. """ def _match_producer_component_id(execution, component_id): if component_id: return execution.properties[ _EXECUTION_TYPE_KEY_COMPONENT_ID].string_value == component_id else: return True def _match_output_key(event, key): if key: assert len( event.path.steps) == 2, 'Event must have two path steps.' return (event.type == metadata_store_pb2.Event.OUTPUT and event.path.steps[0].key == key) else: return event.type == metadata_store_pb2.Event.OUTPUT try: artifact_type = self.store.get_artifact_type(type_name) if not artifact_type: raise tf.errors.NotFoundError( None, None, 'No artifact type found for %s.' % type_name) except tf.errors.NotFoundError: return [] # Gets the executions that are associated with all contexts. assert contexts, 'Must have at least one context.' executions_dict = {} for context in contexts: executions = self.store.get_executions_by_context(context.id) executions_dict.update(dict((e.id, e) for e in executions)) executions_within_context = executions_dict.values() # Filters the executions to match producer component id. qualified_producer_executions = [ e.id for e in executions_within_context if _match_producer_component_id(e, producer_component_id) ] # Gets the output events that have the matched output key. qualified_output_events = [ ev for ev in self.store.get_events_by_execution_ids( qualified_producer_executions) if _match_output_key(ev, output_key) ] # Gets the candidate artifacts from output events. candidate_artifacts = self.store.get_artifacts_by_id( list(set(ev.artifact_id for ev in qualified_output_events))) # Filters the artifacts that have the right artifact type and state. qualified_artifacts = [ a for a in candidate_artifacts if a.type_id == artifact_type.id and self._get_artifact_state(a) == ArtifactState.PUBLISHED ] return [ metadata_store_service_pb2.ArtifactAndType(artifact=a, type=artifact_type) for a in qualified_artifacts ]
def get_qualified_artifacts( self, context: metadata_store_pb2.Context, type_name: Text, producer_component_id: Optional[Text] = None, output_key: Optional[Text] = None, ) -> List[metadata_store_service_pb2.ArtifactAndType]: """Gets qualified artifacts that have the right producer info. Args: context: context constraint to filter artifacts type_name: type constraint to filter artifacts producer_component_id: producer constraint to filter artifacts output_key: output key constraint to filter artifacts Returns: A list of ArtifactAndType, containing qualified artifacts. """ def _match_producer_component_id(execution, component_id): if component_id: return execution.properties['component_id'].string_value == component_id else: return True def _match_output_key(event, key): if key: assert len(event.path.steps) == 2, 'Event must have two path steps.' return (event.type == metadata_store_pb2.Event.OUTPUT and event.path.steps[0].key == key) else: return event.type == metadata_store_pb2.Event.OUTPUT try: artifact_type = self.store.get_artifact_type(type_name) if not artifact_type: raise tf.errors.NotFoundError( None, None, 'No artifact type found for %s.' % type_name) except tf.errors.NotFoundError: return [] executions_within_context = self.store.get_executions_by_context(context.id) qualified_producer_executions = [ e.id for e in executions_within_context if _match_producer_component_id(e, producer_component_id) ] qualified_output_events = [ ev for ev in self.store.get_events_by_execution_ids( qualified_producer_executions) if _match_output_key(ev, output_key) ] candidate_artifacts = self.store.get_artifacts_by_id( list(set(ev.artifact_id for ev in qualified_output_events))) qualified_artifacts = [ a for a in candidate_artifacts if a.type_id == artifact_type.id and self._get_artifact_state(a) == ArtifactState.PUBLISHED ] return [ metadata_store_service_pb2.ArtifactAndType( artifact=a, type=artifact_type) for a in qualified_artifacts ]