def testFullTaxiGcpPipeline(self):
pipeline_name = 'gcp-perf-test-full-e2e-test-{}'.format(
test_utils.random_id())
# Custom CAIP training job using a testing image.
ai_platform_training_args = {
'project': self._GCP_PROJECT_ID,
'region': self._GCP_REGION,
'scaleTier': 'CUSTOM',
'masterType': 'large_model',
'masterConfig': {
'imageUri': self._CONTAINER_IMAGE
},
'workerType': self._WORKER_TYPE,
'parameterServerType': 'standard',
'workerCount': self._WORKER_COUNT,
'parameterServerCount': self._PARAMETER_SERVER_COUNT
}
pipeline = taxi_pipeline_kubeflow_gcp.create_pipeline(
pipeline_name=pipeline_name,
pipeline_root=self._pipeline_root(pipeline_name),
module_file=self._MODULE_FILE,
ai_platform_training_args=ai_platform_training_args,
ai_platform_serving_args=self._AI_PLATFORM_SERVING_ARGS,
beam_pipeline_args=self._BEAM_PIPELINE_ARGS)
self._compile_and_run_pipeline(
pipeline=pipeline,
query_sample_rate=1,
# (1M * batch_size=200) / 200M records ~ 1 epoch
train_steps=1000000,
eval_steps=10000,
worker_count=20,
parameter_server_count=3,
)
def setUp(self):
super(BaseKubeflowTest, self).setUp()
self._test_dir = self.tmp_dir
self.enter_context(test_case_utils.change_working_dir(self.tmp_dir))
self._test_output_dir = 'gs://{}/test_output'.format(self._BUCKET_NAME)
test_id = test_utils.random_id()
self._testdata_root = 'gs://{}/test_data/{}'.format(
self._BUCKET_NAME, test_id)
subprocess.run(
['gsutil', 'cp', '-r', self._TEST_DATA_ROOT, self._testdata_root],
check=True,
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL,
)
self._data_root = os.path.join(self._testdata_root, 'external', 'csv')
self._transform_module = os.path.join(self._MODULE_ROOT,
'transform_module.py')
self._trainer_module = os.path.join(self._MODULE_ROOT,
'trainer_module.py')
self.addCleanup(self._delete_test_dir, test_id)
def testSimpleEnd2EndPipeline(self):
"""End-to-End test for a simple pipeline."""
pipeline_name = 'kubeflow-v2-bqeg-test-{}'.format(
test_utils.random_id())
components = kubeflow_v2_test_utils.create_pipeline_components(
pipeline_root=self._pipeline_root(pipeline_name),
transform_module=self._MODULE_FILE,
trainer_module=self._MODULE_FILE,
bigquery_query=_BIGQUERY_QUERY)
beam_pipeline_args = [
'--temp_location=' + os.path.join(
self._pipeline_root(pipeline_name), 'dataflow', 'temp'),
'--project={}'.format(self._GCP_PROJECT_ID),
# TODO(b/171733562): Remove `use_runner_v2` once it is the default for
# Dataflow.
'--experiments=use_runner_v2',
'--worker_harness_container_image=%s' % self.container_image,
]
pipeline = self._create_pipeline(pipeline_name, components,
beam_pipeline_args)
self._run_pipeline(pipeline)
def testSimpleEnd2EndPipeline(self, bigquery_query, csv_input_location,
use_custom_dataflow_image):
"""End-to-End test for a simple pipeline."""
pipeline_name = 'kubeflow-v2-e2e-test-{}'.format(test_utils.random_id())
components = kubeflow_v2_test_utils.create_pipeline_components(
pipeline_root=self._pipeline_root(pipeline_name),
transform_module=self._MODULE_FILE,
trainer_module=self._MODULE_FILE,
bigquery_query=bigquery_query,
csv_input_location=csv_input_location)
beam_pipeline_args = [
'--temp_location=' +
os.path.join(self._pipeline_root(pipeline_name), 'dataflow', 'temp'),
'--project={}'.format(self._GCP_PROJECT_ID)
]
if use_custom_dataflow_image:
beam_pipeline_args.extend([
# TODO(b/171733562): Remove `use_runner_v2` once it is the default for
# Dataflow.
'--experiments=use_runner_v2',
'--worker_harness_container_image=%s' % self.container_image,
])
pipeline = self._create_pipeline(pipeline_name, components,
beam_pipeline_args)
self._run_pipeline(pipeline, pipeline_name)
self._check_job_status(pipeline_name)
def testEndToEndPipelineRun(self):
"""End-to-end test for pipeline with RuntimeParameter."""
pipeline_name = 'kubeflow-e2e-test-parameter-{}'.format(
test_utils.random_id())
kubeflow_pipeline = penguin_pipeline_kubeflow.create_pipeline(
pipeline_name=pipeline_name,
pipeline_root=self._pipeline_root(pipeline_name),
data_root=self._penguin_data_root,
module_file=self._penguin_dependency_file,
enable_tuning=False,
user_provided_schema_path=self._penguin_schema_file,
ai_platform_training_args=penguin_pipeline_kubeflow
._ai_platform_training_args,
ai_platform_serving_args=penguin_pipeline_kubeflow
._ai_platform_serving_args,
beam_pipeline_args=penguin_pipeline_kubeflow
._beam_pipeline_args_by_runner['DirectRunner'],
use_aip_component=False,
serving_model_dir=self._serving_model_dir)
parameters = {
'train-args': '{"num_steps": 100}',
'eval-args': '{"num_steps": 50}',
}
self._compile_and_run_pipeline(
pipeline=kubeflow_pipeline, parameters=parameters)
self.assertTrue(fileio.exists(self._serving_model_dir))
def testSuccessfulExecution(self):
example_importer = importer.Importer(
artifact_type=simple_artifacts.File,
reimport=False,
source_uri=f'gs://{self._TEST_DATA_BUCKET}/ai-platform-training/mnist'
).with_id('examples')
train = ai_platform_training_component.create_ai_platform_training(
name='simple_aip_training',
project_id=self._GCP_PROJECT_ID,
region=self._GCP_REGION,
image_uri=self._TRAINING_IMAGE,
args=[
'--dataset',
placeholders.InputUriPlaceholder('examples'),
'--model-dir',
placeholders.OutputUriPlaceholder('model'),
'--lr',
placeholders.InputValuePlaceholder('learning_rate'),
],
scale_tier='BASIC',
inputs={'examples': example_importer.outputs['result']},
outputs={'model': standard_artifacts.Model},
parameters={'learning_rate': '0.001'})
pipeline_name = _PIPELINE_NAME_PREFIX.format(test_utils.random_id())
aip_training_pipeline = pipeline.Pipeline(
pipeline_name=pipeline_name,
pipeline_root=self._pipeline_root(pipeline_name),
components=[example_importer, train],
)
self._run_pipeline(aip_training_pipeline)
def testCreateContainerComponentEnd2EndPipeline(self):
"""End-to-End test for container components."""
pipeline_name = 'kubeflow-container-e2e-test-{}'.format(
test_utils.random_id())
text_url = (
'https://storage.googleapis.com/ml-pipeline-playground/hamlet.txt')
pattern = 'art thou'
component_instances = download_grep_print_pipeline.create_pipeline_component_instances(
text_url=text_url,
pattern=pattern,
)
# Test that the pipeline can be executed successfully.
pipeline = self._create_pipeline(pipeline_name, component_instances)
self._compile_and_run_pipeline(pipeline=pipeline,
workflow_name=pipeline_name)
# Test if the correct value has been passed.
artifacts = self._get_artifacts_with_type_and_pipeline(
type_name='ExternalArtifact', pipeline_name=pipeline_name)
# There should be exactly two artifacts.
self.assertEqual(len(artifacts), 2)
for artifact in artifacts:
# TODO(b/150515270) Remove the '/data' suffix when b/150515270 is fixed.
artifact_value = tf.io.gfile.GFile(artifact.uri + '/data',
'r').read()
self.assertGreater(len(artifact_value), 100)
def testEndToEndPipelineRun(self):
"""End-to-end test for pipeline with RuntimeParameter."""
pipeline_name = 'kubeflow-e2e-test-parameter-{}'.format(
test_utils.random_id())
pipeline = taxi_pipeline_runtime_parameter._create_parameterized_pipeline(
pipeline_name=pipeline_name,
pipeline_root=taxi_pipeline_runtime_parameter._pipeline_root,
enable_cache=True,
beam_pipeline_args=[
'--direct_running_mode=multi_processing',
# 0 means auto-detect based on on the number of CPUs available
# during execution time.
'--direct_num_workers=4',
])
parameters = {
'pipeline-root': self._pipeline_root(pipeline_name),
'transform-module': self._transform_module,
'trainer-module': self._trainer_module,
'data-root': self._data_root,
'train-steps': 10,
'eval-steps': 5,
'slicing-column': 'trip_start_hour',
}
self._compile_and_run_pipeline(pipeline=pipeline, parameters=parameters)
def testAIPlatformPusherPipeline(self):
"""Pusher-only test pipeline to AI Platform Prediction."""
pipeline_name_base = 'kubeflow-aip-pusher-test-{}'.format(
test_utils.random_id())
# AI Platform does not accept '-' in the model name.
model_name = ('%s_model' % pipeline_name_base).replace('-', '_')
self.addCleanup(kubeflow_test_utils.delete_ai_platform_model,
model_name)
def _pusher(model_importer, model_blessing_importer):
return Pusher(
custom_executor_spec=executor_spec.ExecutorClassSpec(
ai_platform_pusher_executor.Executor),
model=model_importer.outputs['result'],
model_blessing=model_blessing_importer.outputs['result'],
custom_config={
ai_platform_pusher_executor.SERVING_ARGS_KEY: {
'model_name': model_name,
'project_id': self._GCP_PROJECT_ID,
}
},
)
# Use default service_name / api_version.
service_name, api_version = runner.get_service_name_and_api_version({})
api = discovery.build(service_name, api_version)
# The model should be NotFound yet.
with self.assertRaisesRegex(googleapiclient_errors.HttpError,
'HttpError 404'):
self._sendDummyRequestToModel(api, self._GCP_PROJECT_ID,
model_name)
# Test creation of multiple versions under the same model_name.
pipeline_name_1 = '%s-1' % pipeline_name_base
pipeline_1 = self._create_pipeline(pipeline_name_1, [
self.model_1_importer,
self.model_blessing_importer,
_pusher(self.model_1_importer, self.model_blessing_importer),
])
self._compile_and_run_pipeline(pipeline_1)
self.assertEqual(
1,
self._getNumberOfVersionsForModel(api, self._GCP_PROJECT_ID,
model_name))
self._sendDummyRequestToModel(api, self._GCP_PROJECT_ID, model_name)
pipeline_name_2 = '%s-2' % pipeline_name_base
pipeline_2 = self._create_pipeline(pipeline_name_2, [
self.model_2_importer,
self.model_blessing_importer,
_pusher(self.model_2_importer, self.model_blessing_importer),
])
self._compile_and_run_pipeline(pipeline_2)
self.assertEqual(
2,
self._getNumberOfVersionsForModel(api, self._GCP_PROJECT_ID,
model_name))
self._sendDummyRequestToModel(api, self._GCP_PROJECT_ID, model_name)
def testCsvExampleGenOnDataflowRunner(self):
"""CsvExampleGen-only test pipeline on DataflowRunner invocation."""
pipeline_name = 'kubeflow-csv-example-gen-dataflow-test-{}'.format(
test_utils.random_id())
pipeline = self._create_dataflow_pipeline(pipeline_name, [
CsvExampleGen(input_base=self._data_root),
])
self._compile_and_run_pipeline(pipeline)
def testStatisticsGenOnDataflowRunner(self):
"""StatisticsGen-only test pipeline on DataflowRunner."""
pipeline_name = 'kubeflow-statistics-gen-dataflow-test-{}'.format(
test_utils.random_id())
pipeline = self._create_dataflow_pipeline(pipeline_name, [
self.raw_examples_importer,
StatisticsGen(examples=self.raw_examples_importer.outputs['result'])
])
self._compile_and_run_pipeline(pipeline)
def setUp(self):
super().setUp()
random_id = orchestration_test_utils.random_id()
self._target_container_image = 'gcr.io/{}/{}:{}'.format(
self._GCP_PROJECT_ID, 'taxi-template-kubeflow_v2-e2e-test',
random_id)
# Overriding the pipeline name to
self._pipeline_name = 'taxi_template_kubeflow_v2_e2e_test_{}'.format(
random_id)
self._prepare_skaffold()
def testTransformOnDataflowRunner(self):
"""Transform-only test pipeline on DataflowRunner."""
pipeline_name = 'kubeflow-transform-dataflow-test-{}'.format(
test_utils.random_id())
pipeline = self._create_dataflow_pipeline(pipeline_name, [
self.raw_examples_importer, self.schema_importer,
Transform(examples=self.raw_examples_importer.outputs['result'],
schema=self.schema_importer.outputs['result'],
module_file=self._transform_module)
])
self._compile_and_run_pipeline(pipeline)
def testSimpleEnd2EndPipeline(self):
"""End-to-End test for simple pipeline."""
pipeline_name = 'kubeflow-e2e-test-{}'.format(test_utils.random_id())
components = kubeflow_test_utils.create_e2e_components(
self._pipeline_root(pipeline_name),
self._data_root,
self._transform_module,
self._trainer_module,
)
pipeline = self._create_pipeline(pipeline_name, components)
self._compile_and_run_pipeline(pipeline)
def testArtifactValuePlaceholders(self):
component_instances = (
_tasks_for_pipeline_with_artifact_value_passing())
pipeline_name = 'kubeflow-v2-test-artifact-value-{}'.format(
test_utils.random_id())
pipeline = self._create_pipeline(
pipeline_name,
pipeline_components=component_instances,
)
self._run_pipeline(pipeline)
def setUpClass(cls):
super(BaseKubeflowTest, cls).setUpClass()
if ':' not in cls._BASE_CONTAINER_IMAGE:
# Generate base container image for the test if tag is not specified.
cls.container_image = '{}:{}'.format(cls._BASE_CONTAINER_IMAGE,
test_utils.random_id())
# Create a container image for use by test pipelines.
test_utils.build_and_push_docker_image(cls.container_image,
cls._REPO_BASE)
else: # Use the given image as a base image.
cls.container_image = cls._BASE_CONTAINER_IMAGE
def setUp(self):
super().setUp()
random_id = orchestration_test_utils.random_id()
if ':' not in self._BASE_CONTAINER_IMAGE:
self._base_container_image = '{}:{}'.format(
self._BASE_CONTAINER_IMAGE, random_id)
self._prepare_base_container_image()
else:
self._base_container_image = self._BASE_CONTAINER_IMAGE
self._target_container_image = 'gcr.io/{}/{}:{}'.format(
self._GCP_PROJECT_ID, 'taxi-template-vertex-e2e-test', random_id)
# Overriding the pipeline name to
self._pipeline_name = 'taxi-template-vertex-e2e-{}'.format(random_id)
def testEvaluatorOnDataflowRunner(self):
"""Evaluator-only test pipeline on DataflowRunner."""
pipeline_name = 'kubeflow-evaluator-dataflow-test-{}'.format(
test_utils.random_id())
pipeline = self._create_dataflow_pipeline(pipeline_name, [
self.raw_examples_importer, self.model_1_importer,
Evaluator(
examples=self.raw_examples_importer.outputs['result'],
model=self.model_1_importer.outputs['result'],
feature_slicing_spec=evaluator_pb2.FeatureSlicingSpec(specs=[
evaluator_pb2.SingleSlicingSpec(
column_for_slicing=['trip_start_hour'])
]))
])
self._compile_and_run_pipeline(pipeline)
def setUp(self):
super().setUp()
random_id = orchestration_test_utils.random_id()
self._pipeline_name = self._generate_pipeline_name(random_id)
logging.info('Pipeline: %s', self._pipeline_name)
if ':' not in self._BASE_CONTAINER_IMAGE:
self._base_container_image = '{}:{}'.format(
self._BASE_CONTAINER_IMAGE, random_id)
self._prepare_base_container_image()
else:
self._base_container_image = self._BASE_CONTAINER_IMAGE
self._target_container_image = 'gcr.io/{}/{}'.format(
self._GCP_PROJECT_ID, self._pipeline_name)
def setUp(self):
super().setUp()
random_id = orchestration_test_utils.random_id()
self._pipeline_name = 'taxi-template-kubeflow-e2e-test-' + random_id
logging.info('Pipeline: %s', self._pipeline_name)
self._endpoint = self._get_endpoint()
self._kfp_client = kfp.Client(host=self._endpoint)
logging.info('ENDPOINT: %s', self._endpoint)
self._base_container_image = '{}:{}'.format(self._BASE_CONTAINER_IMAGE,
random_id)
self._target_container_image = 'gcr.io/{}/{}:{}'.format(
self._GCP_PROJECT_ID, 'taxi-template-kubeflow-e2e-test', random_id)
self._prepare_base_container_image()
self._prepare_skaffold()
def testExitHandlerPipelineSuccess(self):
"""End-to-End test for a successful pipeline with exit handler."""
pipeline_name = 'kubeflow-v2-exit-handler-test-{}'.format(
orchestration_test_utils.random_id())
components = test_utils.simple_pipeline_components(_TEST_DATA_ROOT)
beam_pipeline_args = [
'--temp_location=' + os.path.join(
self._pipeline_root(pipeline_name), 'dataflow', 'temp'),
'--project={}'.format(self._GCP_PROJECT_ID)
]
pipeline = self._create_pipeline(pipeline_name, components,
beam_pipeline_args)
output_file_dir = os.path.join(self._pipeline_root(pipeline_name),
_success_file_name)
exit_handler = custom_exit_handler.test_exit_handler(
final_status=tfx.orchestration.experimental.FinalStatusStr(),
file_dir=output_file_dir)
self._run_pipeline(pipeline=pipeline, exit_handler=exit_handler)
# verify execution results
actual_final_status_str = io_utils.read_string_file(output_file_dir)
expected_successful_final_status_str = """
{
"state":"SUCCEEDED",
"error":{}
}
"""
expected_successful_final_status = (
pipeline_spec_pb2.PipelineTaskFinalStatus())
json_format.Parse(expected_successful_final_status_str,
expected_successful_final_status)
actual_final_status = pipeline_spec_pb2.PipelineTaskFinalStatus()
json_format.Parse(actual_final_status_str, actual_final_status)
self.assertProtoPartiallyEquals(
expected_successful_final_status,
actual_final_status,
ignored_fields=['pipeline_job_resource_name'])
def testEndToEndPipelineRun(self):
"""End-to-end test for pipeline with RuntimeParameter."""
pipeline_name = 'kubeflow-e2e-test-parameter-{}'.format(
test_utils.random_id())
pipeline = taxi_pipeline_runtime_parameter._create_parameterized_pipeline(
pipeline_name=pipeline_name, direct_num_workers=4)
parameters = {
'pipeline-root': self._pipeline_root(pipeline_name),
'transform-module': self._transform_module,
'trainer-module': self._trainer_module,
'data-root': self._data_root,
'train-steps': 10,
'eval-steps': 5,
'slicing-column': 'trip_start_hour',
}
self._compile_and_run_pipeline(pipeline=pipeline,
parameters=parameters)
def testSimpleEnd2EndPipeline(self):
"""End-to-End test for a simple pipeline."""
pipeline_name = 'kubeflow-v2-fbeg-test-{}'.format(
test_utils.random_id())
components = kubeflow_v2_test_utils.create_pipeline_components(
pipeline_root=self._pipeline_root(pipeline_name),
transform_module=self._MODULE_FILE,
trainer_module=self._MODULE_FILE,
csv_input_location=_TEST_DATA_ROOT)
beam_pipeline_args = [
'--temp_location=' + os.path.join(
self._pipeline_root(pipeline_name), 'dataflow', 'temp'),
'--project={}'.format(self._GCP_PROJECT_ID)
]
pipeline = self._create_pipeline(pipeline_name, components,
beam_pipeline_args)
self._run_pipeline(pipeline)
def testAIPlatformPusherPipeline(self):
"""Pusher-only test pipeline to AI Platform Prediction."""
pipeline_name_base = 'kubeflow-aip-pusher-test-{}'.format(
test_utils.random_id())
# AI Platform does not accept '-' in the model name.
model_name = ('%s_model' % pipeline_name_base).replace('-', '_')
self.addCleanup(kubeflow_test_utils.delete_ai_platform_model,
model_name)
def _pusher(model_importer, model_blessing_importer):
return Pusher(
custom_executor_spec=executor_spec.ExecutorClassSpec(
ai_platform_pusher_executor.Executor),
model=model_importer.outputs['result'],
model_blessing=model_blessing_importer.outputs['result'],
custom_config={
ai_platform_pusher_executor.SERVING_ARGS_KEY: {
'model_name': model_name,
'project_id': self._GCP_PROJECT_ID,
}
},
)
# Test creation of multiple versions under the same model_name.
pipeline_name_1 = '%s-1' % pipeline_name_base
pipeline_1 = self._create_pipeline(pipeline_name_1, [
self.model_1_importer,
self.model_blessing_importer,
_pusher(self.model_1_importer, self.model_blessing_importer),
])
self._compile_and_run_pipeline(pipeline_1)
pipeline_name_2 = '%s-2' % pipeline_name_base
pipeline_2 = self._create_pipeline(pipeline_name_2, [
self.model_2_importer,
self.model_blessing_importer,
_pusher(self.model_2_importer, self.model_blessing_importer),
])
self._compile_and_run_pipeline(pipeline_2)
def setUp(self):
super().setUp()
self._test_dir = self.tmp_dir
self.enter_context(test_case_utils.change_working_dir(self.tmp_dir))
self._test_output_dir = 'gs://{}/test_output'.format(self._BUCKET_NAME)
test_id = test_utils.random_id()
self._testdata_root = 'gs://{}/test_data/{}'.format(
self._BUCKET_NAME, test_id)
io_utils.copy_dir(self._TEST_DATA_ROOT, self._testdata_root)
self._data_root = os.path.join(self._testdata_root, 'external', 'csv')
self._transform_module = os.path.join(self._MODULE_ROOT,
'transform_module.py')
self._trainer_module = os.path.join(self._MODULE_ROOT,
'trainer_module.py')
self._serving_model_dir = os.path.join(self._testdata_root, 'output')
self.addCleanup(self._delete_test_dir, test_id)
def testFullTaxiGcpPipeline(self):
pipeline_name = 'gcp-perf-test-full-e2e-test-{}'.format(
test_utils.random_id())
# Custom CAIP training job using a testing image.
ai_platform_training_args = {
'project': self._GCP_PROJECT_ID,
'region': self._GCP_REGION,
'scaleTier': 'CUSTOM',
'masterType': 'large_model',
'masterConfig': {
'imageUri': self.container_image
},
'workerType': self._WORKER_TYPE,
'parameterServerType': 'standard',
'workerCount': self._WORKER_COUNT,
'parameterServerCount': self._PARAMETER_SERVER_COUNT
}
pipeline = taxi_pipeline_kubeflow_gcp.create_pipeline(
pipeline_name=pipeline_name,
pipeline_root=self._pipeline_root(pipeline_name),
module_file=self._MODULE_FILE,
ai_platform_training_args=ai_platform_training_args,
ai_platform_serving_args=self._AI_PLATFORM_SERVING_ARGS,
beam_pipeline_args=self._BEAM_PIPELINE_ARGS)
# TODO(b/162451308): Add this clean-up back after we re-enable AIP pusher
# when AIP prediction service supports TF>=2.3.
# self.addCleanup(kubeflow_test_utils.delete_ai_platform_model,
# self._MODEL_NAME)
self._compile_and_run_pipeline(
pipeline=pipeline,
query_sample_rate=1,
# (1M * batch_size=200) / 200M records ~ 1 epoch
train_steps=1000000,
eval_steps=10000,
worker_count=20,
parameter_server_count=3,
)
def testPrimitiveEnd2EndPipeline(self):
"""End-to-End test for primitive artifacts passing."""
pipeline_name = 'kubeflow-primitive-e2e-test-{}'.format(
test_utils.random_id())
components = kubeflow_test_utils.create_primitive_type_components(
pipeline_name)
# Test that the pipeline can be executed successfully.
pipeline = self._create_pipeline(pipeline_name, components)
self._compile_and_run_pipeline(pipeline=pipeline,
workflow_name=pipeline_name + '-run-1')
# Test if the correct value has been passed.
str_artifacts = self._get_artifacts_with_type_and_pipeline(
type_name='String', pipeline_name=pipeline_name)
# There should be exactly one string artifact.
self.assertEqual(1, len(str_artifacts))
self.assertEqual(self._get_value_of_string_artifact(str_artifacts[0]),
'hello %s\n' % pipeline_name)
# Test caching.
self._compile_and_run_pipeline(pipeline=pipeline,
workflow_name=pipeline_name + '-run-2')
cached_execution = self._get_executions_by_pipeline_name_and_state(
pipeline_name=pipeline_name, state=metadata.EXECUTION_STATE_CACHED)
self.assertEqual(2, len(cached_execution))
def testSimpleEnd2EndPipeline(self, bigquery_query, csv_input_location):
"""End-to-End test for a simple pipeline."""
pipeline_name = '-e2e-test-{}'.format(
test_utils.random_id())
components = kubeflow_v2_test_utils.create_pipeline_components(
pipeline_root=self._pipeline_root(pipeline_name),
transform_module=self._MODULE_FILE,
trainer_module=self._MODULE_FILE,
bigquery_query=bigquery_query,
csv_input_location=csv_input_location)
beam_pipeline_args = [
'--temp_location=' +
os.path.join(self._pipeline_root(pipeline_name), 'dataflow', 'temp'),
'--project={}'.format(self._GCP_PROJECT_ID)
]
pipeline = self._create_pipeline(pipeline_name, components,
beam_pipeline_args)
job_name = self._run_pipeline(pipeline)
self._check_job_status(job_name)
def testAIPlatformGenericTrainerPipeline(self):
"""Trainer-only pipeline on AI Platform Training with GenericTrainer."""
pipeline_name = 'kubeflow-aip-generic-trainer-test-{}'.format(
test_utils.random_id())
pipeline = self._create_pipeline(pipeline_name, [
self.schema_importer, self.transformed_examples_importer,
self.transform_graph_importer,
Trainer(custom_executor_spec=executor_spec.ExecutorClassSpec(
ai_platform_trainer_executor.GenericExecutor),
module_file=self._trainer_module,
transformed_examples=self.transformed_examples_importer.
outputs['result'],
schema=self.schema_importer.outputs['result'],
transform_graph=self.transform_graph_importer.
outputs['result'],
train_args=trainer_pb2.TrainArgs(num_steps=10),
eval_args=trainer_pb2.EvalArgs(num_steps=5),
custom_config={
ai_platform_trainer_executor.TRAINING_ARGS_KEY:
self._getCaipTrainingArgs(pipeline_name)
})
])
self._compile_and_run_pipeline(pipeline)
self._assertNumberOfTrainerOutputIsOne(pipeline_name)
def testAIPlatformDistributedTunerPipeline(self):
"""Tuner-only pipeline for distributed Tuner flock on AIP Training."""
pipeline_name = 'kubeflow-aip-dist-tuner-test-{}'.format(
test_utils.random_id())
pipeline = self._create_pipeline(
pipeline_name,
[
self.iris_examples_importer,
self.iris_schema_importer,
ai_platform_tuner_component.Tuner(
examples=self.iris_examples_importer.outputs['result'],
module_file=self._iris_tuner_module,
schema=self.iris_schema_importer.outputs['result'],
train_args=trainer_pb2.TrainArgs(num_steps=10),
eval_args=trainer_pb2.EvalArgs(num_steps=5),
# 3 worker parallel tuning.
tune_args=tuner_pb2.TuneArgs(num_parallel_trials=3),
custom_config={
ai_platform_trainer_executor.TRAINING_ARGS_KEY:
self._getCaipTrainingArgs(pipeline_name)
})
])
self._compile_and_run_pipeline(pipeline)
self._assertHyperparametersAreWritten(pipeline_name)