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')
