def get_pipeline_outputs( metadata_connection_config: Optional[ metadata_store_pb2.ConnectionConfig], pipeline_name: str ) -> Dict[Text, Dict[Text, Dict[int, types.Artifact]]]: """Returns a dictionary of pipeline output artifacts for every component. Args: metadata_connection_config: connection configuration to MLMD. pipeline_name: Name of the pipeline. Returns: a dictionary of holding list of artifacts for a component id. """ output_map = {} with metadata.Metadata(metadata_connection_config) as m: executions = pipeline_recorder_utils.get_latest_executions( m, pipeline_name) for execution in executions: component_id = pipeline_recorder_utils.get_component_id_from_execution( m, execution) output_dict = {} for event in m.store.get_events_by_execution_ids([execution.id]): if event.type == metadata_store_pb2.Event.OUTPUT: artifacts = m.store.get_artifacts_by_id( [event.artifact_id]) steps = event.path.steps if not steps or not steps[0].HasField('key'): raise ValueError( 'Artifact key is not recorded in the MLMD.') key = steps[0].key artifacts = m.store.get_artifacts_by_id( [event.artifact_id]) if key not in output_dict: output_dict[key] = {} for pb_artifact in artifacts: if len(steps) < 2 or not steps[1].HasField('index'): raise ValueError( 'Artifact index is not recorded in the MLMD.') artifact_index = steps[1].index if artifact_index in output_dict[key]: raise ValueError('Artifact already in output_dict') [artifact_type] = m.store.get_artifact_types_by_id( [pb_artifact.type_id]) artifact = artifact_utils.deserialize_artifact( artifact_type, pb_artifact) output_dict[key][artifact_index] = artifact output_map[component_id] = output_dict return output_map
def testStubbedImdbPipelineBeam(self): pipeline_ir = compiler.Compiler().compile(self.imdb_pipeline) pipeline_mock.replace_executor_with_stub(pipeline_ir, self._recorded_output_dir, []) BeamDagRunner().run(pipeline_ir) self.assertTrue(fileio.exists(self._metadata_path)) metadata_config = metadata.sqlite_metadata_connection_config( self._metadata_path) # Verify that recorded files are successfully copied to the output uris. with metadata.Metadata(metadata_config) as m: for execution in m.store.get_executions(): component_id = pipeline_recorder_utils.get_component_id_from_execution( m, execution) if component_id.startswith('ResolverNode'): continue eid = [execution.id] events = m.store.get_events_by_execution_ids(eid) output_events = [ x for x in events if x.type == metadata_store_pb2.Event.OUTPUT ] for event in output_events: steps = event.path.steps assert steps[0].HasField('key') name = steps[0].key artifacts = m.store.get_artifacts_by_id([event.artifact_id]) for idx, artifact in enumerate(artifacts): self.assertDirectoryEqual( artifact.uri, os.path.join(self._recorded_output_dir, component_id, name, str(idx))) # Calls verifier for pipeline output artifacts, excluding the resolver node. BeamDagRunner().run(self.imdb_pipeline) pipeline_outputs = executor_verifier_utils.get_pipeline_outputs( self.imdb_pipeline.metadata_connection_config, self._pipeline_name) verifier_map = { 'model': self._verify_model, 'model_run': self._verify_model, 'examples': self._verify_examples, 'schema': self._verify_schema, 'anomalies': self._verify_anomalies, 'evaluation': self._verify_evaluation, # A subdirectory of updated_analyzer_cache has changing name. 'updated_analyzer_cache': self._veryify_root_dir, } # List of components to verify. ResolverNode is ignored because it # doesn't have an executor. verify_component_ids = [ component.id for component in self.imdb_pipeline.components if not component.id.startswith('ResolverNode') ] for component_id in verify_component_ids: for key, artifact_dict in pipeline_outputs[component_id].items(): for idx, artifact in artifact_dict.items(): logging.info('Verifying %s', component_id) recorded_uri = os.path.join(self._recorded_output_dir, component_id, key, str(idx)) verifier_map.get(key, self._verify_file_path)(artifact.uri, recorded_uri)
def testStubbedTaxiPipelineBeam(self): pipeline_ir = compiler.Compiler().compile(self.taxi_pipeline) logging.info('Replacing with test_data_dir:%s', self._recorded_output_dir) pipeline_mock.replace_executor_with_stub(pipeline_ir, self._recorded_output_dir, []) BeamDagRunner().run_with_ir(pipeline_ir) self.assertTrue(fileio.exists(self._metadata_path)) metadata_config = metadata.sqlite_metadata_connection_config( self._metadata_path) # Verify that recorded files are successfully copied to the output uris. with metadata.Metadata(metadata_config) as m: artifacts = m.store.get_artifacts() artifact_count = len(artifacts) executions = m.store.get_executions() execution_count = len(executions) # Artifact count is greater by 7 due to extra artifacts produced by # Evaluator(blessing and evaluation), Trainer(model and model_run) and # Transform(example, graph, cache, pre_transform_statistics, # pre_transform_schema, post_transform_statistics, post_transform_schema, # post_transform_anomalies) minus Resolver which doesn't generate # new artifact. self.assertEqual(artifact_count, execution_count + 7) self.assertLen(self.taxi_pipeline.components, execution_count) for execution in executions: component_id = pipeline_recorder_utils.get_component_id_from_execution( m, execution) if component_id.startswith('Resolver'): continue eid = [execution.id] events = m.store.get_events_by_execution_ids(eid) output_events = [ x for x in events if x.type == metadata_store_pb2.Event.OUTPUT ] for event in output_events: steps = event.path.steps self.assertTrue(steps[0].HasField('key')) name = steps[0].key artifacts = m.store.get_artifacts_by_id( [event.artifact_id]) for idx, artifact in enumerate(artifacts): self.assertDirectoryEqual( artifact.uri, os.path.join(self._recorded_output_dir, component_id, name, str(idx))) # Calls verifier for pipeline output artifacts, excluding the resolver node. BeamDagRunner().run(self.taxi_pipeline) pipeline_outputs = executor_verifier_utils.get_pipeline_outputs( self.taxi_pipeline.metadata_connection_config, self._pipeline_name) verifier_map = { 'model': self._verify_model, 'model_run': self._verify_model, 'examples': self._verify_examples, 'schema': self._verify_schema, 'anomalies': self._verify_anomalies, 'evaluation': self._verify_evaluation, # A subdirectory of updated_analyzer_cache has changing name. 'updated_analyzer_cache': self._veryify_root_dir, } # List of components to verify. Resolver is ignored because it # doesn't have an executor. verify_component_ids = [ component.id for component in self.taxi_pipeline.components if not component.id.startswith('Resolver') ] for component_id in verify_component_ids: logging.info('Verifying %s', component_id) for key, artifact_dict in pipeline_outputs[component_id].items(): for idx, artifact in artifact_dict.items(): recorded_uri = os.path.join(self._recorded_output_dir, component_id, key, str(idx)) verifier_map.get(key, self._verify_file_path)(artifact.uri, recorded_uri)