def create_pipeline(self, update: bool = False) -> None:
"""Creates or updates a pipeline to use in Vertex Pipelines.
Args:
update: set as true to update pipeline.
"""
if self.flags_dict.get(labels.BUILD_IMAGE):
build_image_fn = functools.partial(
kubeflow_handler.create_container_image,
base_image=self.flags_dict.get(labels.BASE_IMAGE))
else:
build_image_fn = None
patcher = kubeflow_v2_dag_runner_patcher.KubeflowV2DagRunnerPatcher(
call_real_run=True,
build_image_fn=build_image_fn,
prepare_dir_fn=functools.partial(self._prepare_pipeline_dir,
required=update))
context = self.execute_dsl(patcher)
pipeline_name = context[patcher.PIPELINE_NAME]
if update:
click.echo(
'Pipeline "{}" updated successfully.'.format(pipeline_name))
else:
click.echo(
'Pipeline "{}" created successfully.'.format(pipeline_name))
def compile_pipeline(self) -> None:
"""Compiles pipeline into Kubeflow V2 Pipelines spec."""
patcher = kubeflow_v2_dag_runner_patcher.KubeflowV2DagRunnerPatcher(
call_real_run=True)
context = self.execute_dsl(patcher)
click.echo(f'Pipeline {context[patcher.PIPELINE_NAME]} compiled '
'successfully.')
def testPatcherSavePipelineFn(self):
pipeline_name = 'dummy'
pipeline_dir = '/foo/pipeline'
mock_prepare_dir_fn = mock.MagicMock(return_value=pipeline_dir)
patcher = kubeflow_v2_dag_runner_patcher.KubeflowV2DagRunnerPatcher(
call_real_run=False, prepare_dir_fn=mock_prepare_dir_fn)
runner = kubeflow_v2_dag_runner.KubeflowV2DagRunner(
config=kubeflow_v2_dag_runner.KubeflowV2DagRunnerConfig())
pipeline = tfx_pipeline.Pipeline(pipeline_name, 'dummy_root')
with patcher.patch() as context:
runner.run(pipeline)
mock_prepare_dir_fn.assert_called_once_with(pipeline_name)
self.assertEqual(
context[patcher.OUTPUT_FILE_PATH],
os.path.join(pipeline_dir,
kubeflow_v2_dag_runner_patcher.OUTPUT_FILENAME))
def testPatcherBuildImageFn(self):
given_image_name = 'foo/bar'
built_image_name = 'foo/bar@sha256:1234567890'
mock_build_image_fn = mock.MagicMock(return_value=built_image_name)
patcher = kubeflow_v2_dag_runner_patcher.KubeflowV2DagRunnerPatcher(
call_real_run=True, build_image_fn=mock_build_image_fn)
runner_config = kubeflow_v2_dag_runner.KubeflowV2DagRunnerConfig(
default_image=given_image_name)
runner = kubeflow_v2_dag_runner.KubeflowV2DagRunner(
config=runner_config)
pipeline = tfx_pipeline.Pipeline('dummy', 'dummy_root')
with patcher.patch() as context:
runner.run(pipeline)
self.assertIn(patcher.OUTPUT_FILE_PATH, context)
mock_build_image_fn.assert_called_once_with(given_image_name)
self.assertEqual(runner_config.default_image, built_image_name)