def setUp(self):
self._output_dict = {'output_name': [standard_artifacts.Examples()]}
self._pipeline_properties = base_component.PipelineProperties(
output_dir='output_dir',
log_root='log_root',
)
with dsl.Pipeline('test_pipeline'):
self.component = base_component.BaseComponent(
component_name='TFXComponent',
input_dict=collections.OrderedDict([
('input_data', 'input-data-contents'),
('train_steps', 300),
('accuracy_threshold', 0.3),
]),
output_dict=self._output_dict,
exec_properties=collections.OrderedDict([
('module_file', '/path/to/module.py')
]),
executor_class_path='some.executor.Class',
pipeline_properties=self._pipeline_properties,
)
def _construct_pipeline_graph(self, pipeline):
"""Constructs a Kubeflow Pipeline graph.
Args:
pipeline: The logical TFX pipeline to base the construction on.
"""
output_dir = os.path.join(pipeline.pipeline_args['pipeline_root'],
pipeline.pipeline_args['pipeline_name'])
beam_pipeline_args = []
tfx_image = None
if 'additional_pipeline_args' in pipeline.pipeline_args:
additional_pipeline_args = pipeline.pipeline_args[
'additional_pipeline_args']
beam_pipeline_args = additional_pipeline_args.get(
'beam_pipeline_args', [])
tfx_image = additional_pipeline_args.get('tfx_image')
pipeline_properties = base_component.PipelineProperties(
output_dir=output_dir,
log_root=pipeline.pipeline_args['log_root'],
beam_pipeline_args=beam_pipeline_args,
tfx_image=tfx_image,
)
# producers is a map from an output Channel, to a Kubeflow component that
# is responsible for the named output represented by the Channel.
# Assumption: Channels are unique in a pipeline.
producers = {}
# Assumption: There is a partial ordering of components in the list, i.e.,
# if component A depends on component B and C, then A appears after B and C
# in the list.
for component in pipeline.components:
input_dict = {}
for input_name, input_channel in component.input_dict.items():
if input_channel in producers:
output = getattr(
producers[input_channel]['component'].outputs,
producers[input_channel]['channel_name'])
if not isinstance(output, dsl.PipelineParam):
raise ValueError(
'Component outputs should be of type dsl.PipelineParam.'
' Got type {} for output {}'.format(
type(output), output))
input_dict[input_name] = output
else:
input_dict[input_name] = json.dumps(
[x.json_dict() for x in input_channel.get()])
executor_class_path = '.'.join(
[component.executor.__module__, component.executor.__name__])
kfp_component = base_component.BaseComponent(
component_name=component.component_name,
input_dict=input_dict,
output_dict=self._prepare_output_dict(component.outputs),
exec_properties=component.exec_properties,
executor_class_path=executor_class_path,
pipeline_properties=pipeline_properties)
for channel_name, channel in component.outputs.get_all().items():
producers[channel] = {}
producers[channel]['component'] = kfp_component
producers[channel]['channel_name'] = channel_name