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