def testBuildTwoStepPipeline(self):
my_builder = pipeline_builder.PipelineBuilder(
tfx_pipeline=test_utils.two_step_pipeline(),
default_image='gcr.io/my-tfx:latest')
actual_pipeline_spec = my_builder.build()
self.assertProtoEquals(
test_utils.get_proto_from_test_data('expected_two_step_pipeline.pbtxt',
pipeline_pb2.PipelineSpec()),
actual_pipeline_spec)
def testBuildRuntimeConfig(self):
my_builder = pipeline_builder.RuntimeConfigBuilder(
pipeline_info=test_utils.two_step_pipeline().pipeline_info,
parameter_values={
'string_param': 'test-string',
'int_param': 42,
'float_param': 3.14
})
actual_output_path_config = my_builder.build()
self.assertProtoEquals(test_utils.TEST_RUNTIME_CONFIG,
actual_output_path_config)
def testBuildTwoStepPipelineWithCacheEnabled(self):
pipeline = test_utils.two_step_pipeline()
pipeline.enable_cache = True
builder = pipeline_builder.PipelineBuilder(
tfx_pipeline=pipeline, default_image='gcr.io/my-tfx:latest')
pipeline_spec = builder.build()
self.assertProtoEquals(
test_utils.get_proto_from_test_data(
'expected_two_step_pipeline_with_cache_enabled.pbtxt',
pipeline_pb2.PipelineSpec()), pipeline_spec)
def testPipelineWithExitHandler(self):
pipeline = test_utils.two_step_pipeline()
# define exit handler
exit_handler = test_utils.dummy_exit_handler(
param1=decorators.FinalStatusStr())
builder = pipeline_builder.PipelineBuilder(
tfx_pipeline=pipeline,
default_image='gcr.io/my-tfx:latest',
exit_handler=exit_handler)
pipeline_spec = builder.build()
self.assertProtoEquals(
test_utils.get_proto_from_test_data(
'expected_two_step_pipeline_with_exit_handler.pbtxt',
pipeline_pb2.PipelineSpec()), pipeline_spec)
def testCompileTwoStepPipeline(self, fake_now, fake_sys_version):
fake_now.return_value = datetime.date(2020, 1, 1)
fake_sys_version.major = 3
fake_sys_version.minor = 7
runner = kubeflow_v2_dag_runner.KubeflowV2DagRunner(
output_dir=_TEST_DIR,
output_filename=_TEST_FILE_NAME,
config=kubeflow_v2_dag_runner.KubeflowV2DagRunnerConfig(
display_name='my-pipeline',
default_image='gcr.io/my-tfx:latest'))
self._compare_against_testdata(
runner=runner,
pipeline=test_utils.two_step_pipeline(),
golden_file='expected_two_step_pipeline_job.json')
