def _run_publisher(self, output_dict: Dict[Text, List[types.Artifact]]) -> None: """Publish execution result to ml metadata.""" with self._metadata_connection as m: p = publisher.Publisher(metadata_handler=m) p.publish_execution( component_info=self._component_info, output_artifacts=output_dict)
def testPrepareExecutionComplete(self): p = publisher.Publisher(metadata_handler=self._mock_metadata) p.publish_execution(component_info=self._component_info, output_artifacts=self._output_dict, exec_properties=self._exec_properties) self._mock_metadata.publish_execution.assert_called_with( component_info=self._component_info, output_artifacts=self._output_dict, exec_properties=self._exec_properties)
def _run_publisher(self, use_cached_results: bool, execution_id: int, input_dict: Dict[Text, List[types.Artifact]], output_dict: Dict[Text, List[types.Artifact]]) -> None: """Publish execution result to ml metadata.""" with metadata.Metadata(self._metadata_connection_config) as m: p = publisher.Publisher(metadata_handler=m) p.publish_execution(execution_id=execution_id, input_dict=input_dict, output_dict=output_dict, use_cached_results=use_cached_results)
def testPrepareExecutionCached(self): input_dict = copy.deepcopy(self._input_dict) output_dict = copy.deepcopy(self._output_dict) p = publisher.Publisher(metadata_handler=self._mock_metadata) p.publish_execution( self._execution_id, input_dict, output_dict, use_cached_results=True) self._mock_metadata.publish_execution.assert_called_with( execution_id=self._execution_id, input_dict=input_dict, output_dict=output_dict, state=metadata.EXECUTION_STATE_CACHED)
def testPrepareExecutionComplete(self): p = publisher.Publisher(metadata_handler=self._mock_metadata) p.publish_execution(component_info=self._component_info, output_artifacts=self._output_dict, exec_properties=self._exec_properties) self._mock_metadata.publish_execution.assert_called_with( component_info=self._component_info, output_artifacts=self._output_dict, exec_properties=self._exec_properties) self.assertEqual( self._output_dict['output_data'][0].get_string_custom_property( 'tfx_version'), version.__version__)
def test_prepare_execution_complete(self): input_dict = copy.deepcopy(self._input_dict) output_dict = copy.deepcopy(self._output_dict) p = publisher.Publisher(metadata_handler=self._mock_metadata) p.publish_execution(self._execution_id, input_dict, output_dict, use_cached_results=False) self._mock_metadata.publish_execution.assert_called_with( execution_id=self._execution_id, input_dict=input_dict, output_dict=output_dict, state=metadata.EXECUTION_STATE_COMPLETE)
def _run_publisher( self, use_cached_results: bool, execution_id: int, input_dict: Dict[Text, List[types.TfxArtifact]], output_dict: Dict[Text, List[types.TfxArtifact]]) -> None: """Publish execution result to ml metadata.""" tf.logging.info('Run publisher for %s', self._component_info.component_id) with metadata.Metadata(self._metadata_connection_config) as m: p = publisher.Publisher(metadata_handler=m) p.publish_execution(execution_id=execution_id, input_dict=input_dict, output_dict=output_dict, use_cached_results=use_cached_results)
def test_prepare_execution_cached(self): input_dict = copy.deepcopy(self._input_dict) output_dict = copy.deepcopy(self._output_dict) p = publisher.Publisher(metadata_handler=self._mock_metadata) p.publish_execution(self._execution_id, input_dict, output_dict, use_cached_results=True) self._mock_metadata.publish_execution.assert_called_with( execution_id=self._execution_id, input_dict=input_dict, output_dict=output_dict, state='skipped')