def testBuildTask(self):
query = 'SELECT * FROM TABLE'
bq_example_gen = big_query_example_gen_component.BigQueryExampleGen(
query=query)
deployment_config = pipeline_pb2.PipelineDeploymentConfig()
component_defs = {}
my_builder = step_builder.StepBuilder(
node=bq_example_gen,
image='gcr.io/tensorflow/tfx:latest',
deployment_config=deployment_config,
component_defs=component_defs,
enable_cache=True)
actual_step_spec = self._sole(my_builder.build())
actual_component_def = self._sole(component_defs)
self.assertProtoEquals(
test_utils.get_proto_from_test_data(
'expected_bq_example_gen_component.pbtxt',
pipeline_pb2.ComponentSpec()), actual_component_def)
self.assertProtoEquals(
test_utils.get_proto_from_test_data(
'expected_bq_example_gen_task.pbtxt',
pipeline_pb2.PipelineTaskSpec()), actual_step_spec)
self.assertProtoEquals(
test_utils.get_proto_from_test_data(
'expected_bq_example_gen_executor.pbtxt',
pipeline_pb2.PipelineDeploymentConfig()), deployment_config)
def testConstructWithInputConfig(self):
big_query_example_gen = component.BigQueryExampleGen(
input_config=example_gen_pb2.Input(splits=[
example_gen_pb2.Input.Split(name='train', pattern='query1'),
example_gen_pb2.Input.Split(name='eval', pattern='query2'),
]))
self.assertEqual(standard_artifacts.Examples.TYPE_NAME,
big_query_example_gen.outputs['examples'].type_name)
def testConstruct(self):
big_query_example_gen = component.BigQueryExampleGen(query='query')
self.assertEqual(standard_artifacts.Examples.TYPE_NAME,
big_query_example_gen.outputs['examples'].type_name)
artifact_collection = big_query_example_gen.outputs['examples'].get()
self.assertEqual(1, len(artifact_collection))
self.assertEqual(['train', 'eval'],
artifact_utils.decode_split_names(
artifact_collection[0].split_names))
def testConstructWithOutputConfig(self):
big_query_example_gen = component.BigQueryExampleGen(
query='query',
output_config=example_gen_pb2.Output(
split_config=example_gen_pb2.SplitConfig(splits=[
example_gen_pb2.SplitConfig.Split(name='train',
hash_buckets=2),
example_gen_pb2.SplitConfig.Split(name='eval',
hash_buckets=1),
])))
self.assertEqual(standard_artifacts.Examples.TYPE_NAME,
big_query_example_gen.outputs['examples'].type_name)
def _two_step_pipeline() -> tfx_pipeline.Pipeline:
table_name = data_types.RuntimeParameter(
name='table-name', ptype=Text, default='default-table')
example_gen = big_query_example_gen_component.BigQueryExampleGen(
query='SELECT * FROM %s' % str(table_name))
statistics_gen = statistics_gen_component.StatisticsGen(
examples=example_gen.outputs['examples'])
return tfx_pipeline.Pipeline(
pipeline_name='two_step_pipeline',
pipeline_root='pipeline_root',
metadata_connection_config=metadata_store_pb2.ConnectionConfig(),
components=[example_gen, statistics_gen],
)
def two_step_pipeline() -> tfx_pipeline.Pipeline:
"""Returns a simple 2-step pipeline under test."""
example_gen = big_query_example_gen_component.BigQueryExampleGen(
query='SELECT * FROM TABLE')
statistics_gen = components.StatisticsGen(
examples=example_gen.outputs['examples'])
return tfx_pipeline.Pipeline(
pipeline_name=_TEST_TWO_STEP_PIPELINE_NAME,
pipeline_root=_TEST_PIPELINE_ROOT,
components=[example_gen, statistics_gen],
# Needs to set GCP project because BQ is used.
beam_pipeline_args=[
'--project=my-gcp-project',
])
def testConstructWithInputConfig(self):
big_query_example_gen = component.BigQueryExampleGen(
input_config=example_gen_pb2.Input(splits=[
example_gen_pb2.Input.Split(name='train', pattern='query1'),
example_gen_pb2.Input.Split(name='eval', pattern='query2'),
example_gen_pb2.Input.Split(name='test', pattern='query3')
]))
self.assertEqual(standard_artifacts.Examples.TYPE_NAME,
big_query_example_gen.outputs['examples'].type_name)
artifact_collection = big_query_example_gen.outputs['examples'].get()
self.assertEqual(1, len(artifact_collection))
self.assertEqual(['train', 'eval', 'test'],
artifact_utils.decode_split_names(
artifact_collection[0].split_names))
def testConstructWithOutputConfig(self):
big_query_example_gen = component.BigQueryExampleGen(
query='query',
output_config=example_gen_pb2.Output(
split_config=example_gen_pb2.SplitConfig(splits=[
example_gen_pb2.SplitConfig.Split(name='train', hash_buckets=2),
example_gen_pb2.SplitConfig.Split(name='eval', hash_buckets=1),
example_gen_pb2.SplitConfig.Split(name='test', hash_buckets=1)
])))
self.assertEqual(standard_artifacts.Examples.TYPE_NAME,
big_query_example_gen.outputs['examples'].type_name)
artifact_collection = big_query_example_gen.outputs['examples'].get()
self.assertEqual(1, len(artifact_collection))
self.assertEqual(['train', 'eval', 'test'],
artifact_utils.decode_split_names(
artifact_collection[0].split_names))
def testConstructWithRangeConfig(self):
range_config = range_config_pb2.RangeConfig(
static_range=range_config_pb2.StaticRange(start_span_number=2,
end_span_number=2))
# @span_yyyymmdd_utc will replaced to '19700103' to query, span `2` will be
# recorded in output Example artifact.
big_query_example_gen = component.BigQueryExampleGen(
query='select * from table where date=@span_yyyymmdd_utc',
range_config=range_config)
self.assertEqual(
standard_artifacts.Examples.TYPE_NAME,
big_query_example_gen.outputs[
standard_component_specs.EXAMPLES_KEY].type_name)
stored_range_config = range_config_pb2.RangeConfig()
proto_utils.json_to_proto(
big_query_example_gen.exec_properties[
standard_component_specs.RANGE_CONFIG_KEY],
stored_range_config)
self.assertEqual(range_config, stored_range_config)
def _two_step_pipeline() -> tfx_pipeline.Pipeline:
default_input_config = json.dumps({
'splits': [{
'name': 'single_split',
'pattern': 'SELECT * FROM default-table'
}]
})
input_config = data_types.RuntimeParameter(name='input_config',
ptype=str,
default=default_input_config)
example_gen = big_query_example_gen_component.BigQueryExampleGen(
input_config=input_config, output_config=example_gen_pb2.Output())
statistics_gen = statistics_gen_component.StatisticsGen(
examples=example_gen.outputs['examples'])
return tfx_pipeline.Pipeline(
pipeline_name='two_step_pipeline',
pipeline_root='pipeline_root',
metadata_connection_config=metadata_store_pb2.ConnectionConfig(),
components=[example_gen, statistics_gen],
)
def create_pipeline(
pipeline_name: Text,
pipeline_root: Text,
data_path: Text,
# TODO(step 7): (Optional) Uncomment here to use BigQuery as a data source.
query: Text,
preprocessing_fn: Text,
run_fn: Text,
train_args: trainer_pb2.TrainArgs,
eval_args: trainer_pb2.EvalArgs,
eval_accuracy_threshold: float,
serving_model_dir: Text,
metadata_connection_config: Optional[
metadata_store_pb2.ConnectionConfig] = None,
beam_pipeline_args: Optional[List[Text]] = None,
ai_platform_training_args: Optional[Dict[Text, Text]] = None,
ai_platform_serving_args: Optional[Dict[Text, Any]] = None,
) -> pipeline.Pipeline:
"""Implements the chicago taxi pipeline with TFX."""
components = []
# Brings data into the pipeline or otherwise joins/converts training data.
# example_gen = CsvExampleGen(input=external_input(data_path))
# TODO(step 7): (Optional) Uncomment here to use BigQuery as a data source.
example_gen = big_query_example_gen_component.BigQueryExampleGen(
query=query)
components.append(example_gen)
# Computes statistics over data for visualization and example validation.
statistics_gen = StatisticsGen(examples=example_gen.outputs['examples'])
# TODO(step 5): Uncomment here to add StatisticsGen to the pipeline.
components.append(statistics_gen)
# Generates schema based on statistics files.
schema_gen = SchemaGen(statistics=statistics_gen.outputs['statistics'],
infer_feature_shape=True)
# TODO(step 5): Uncomment here to add SchemaGen to the pipeline.
components.append(schema_gen)
# Performs anomaly detection based on statistics and data schema.
example_validator = ExampleValidator( # pylint: disable=unused-variable
statistics=statistics_gen.outputs['statistics'],
schema=schema_gen.outputs['schema'])
# TODO(step 5): Uncomment here to add ExampleValidator to the pipeline.
components.append(example_validator)
# Performs transformations and feature engineering in training and serving.
transform = Transform(examples=example_gen.outputs['examples'],
schema=schema_gen.outputs['schema'],
preprocessing_fn=preprocessing_fn)
# TODO(step 6): Uncomment here to add Transform to the pipeline.
components.append(transform)
# Uses user-provided Python function that implements a model using TF-Learn.
trainer_args = {
'run_fn':
run_fn,
'transformed_examples':
transform.outputs['transformed_examples'],
'schema':
schema_gen.outputs['schema'],
'transform_graph':
transform.outputs['transform_graph'],
'train_args':
train_args,
'eval_args':
eval_args,
'custom_executor_spec':
executor_spec.ExecutorClassSpec(trainer_executor.GenericExecutor),
}
if ai_platform_training_args is not None:
trainer_args.update({
'custom_executor_spec':
executor_spec.ExecutorClassSpec(
ai_platform_trainer_executor.GenericExecutor),
'custom_config': {
ai_platform_trainer_executor.TRAINING_ARGS_KEY:
ai_platform_training_args,
}
})
trainer = Trainer(**trainer_args)
# TODO(step 6): Uncomment here to add Trainer to the pipeline.
components.append(trainer)
# Get the latest blessed model for model validation.
model_resolver = ResolverNode(
instance_name='latest_blessed_model_resolver',
resolver_class=latest_blessed_model_resolver.
LatestBlessedModelResolver,
model=Channel(type=Model),
model_blessing=Channel(type=ModelBlessing))
# TODO(step 6): Uncomment here to add ResolverNode to the pipeline.
components.append(model_resolver)
# Uses TFMA to compute a evaluation statistics over features of a model and
# perform quality validation of a candidate model (compared to a baseline).
eval_config = tfma.EvalConfig(
model_specs=[tfma.ModelSpec(label_key='big_tipper')],
slicing_specs=[tfma.SlicingSpec()],
metrics_specs=[
tfma.MetricsSpec(metrics=[
tfma.MetricConfig(
class_name='BinaryAccuracy',
threshold=tfma.MetricThreshold(
value_threshold=tfma.GenericValueThreshold(
lower_bound={'value': eval_accuracy_threshold}),
change_threshold=tfma.GenericChangeThreshold(
direction=tfma.MetricDirection.HIGHER_IS_BETTER,
absolute={'value': -1e-10})))
])
])
evaluator = Evaluator(
examples=example_gen.outputs['examples'],
model=trainer.outputs['model'],
baseline_model=model_resolver.outputs['model'],
# Change threshold will be ignored if there is no baseline (first run).
eval_config=eval_config)
# TODO(step 6): Uncomment here to add Evaluator to the pipeline.
components.append(evaluator)
# Checks whether the model passed the validation steps and pushes the model
# to a file destination if check passed.
pusher_args = {
'model':
trainer.outputs['model'],
'model_blessing':
evaluator.outputs['blessing'],
'push_destination':
pusher_pb2.PushDestination(
filesystem=pusher_pb2.PushDestination.Filesystem(
base_directory=serving_model_dir)),
}
if ai_platform_serving_args is not None:
pusher_args.update({
'custom_executor_spec':
executor_spec.ExecutorClassSpec(
ai_platform_pusher_executor.Executor),
'custom_config': {
ai_platform_pusher_executor.SERVING_ARGS_KEY:
ai_platform_serving_args
},
})
pusher = Pusher(**pusher_args) # pylint: disable=unused-variable
# TODO(step 6): Uncomment here to add Pusher to the pipeline.
components.append(pusher)
return pipeline.Pipeline(
pipeline_name=pipeline_name,
pipeline_root=pipeline_root,
components=components,
# Change this value to control caching of execution results. Default value
# is `False`.
# enable_cache=True,
metadata_connection_config=metadata_connection_config,
beam_pipeline_args=beam_pipeline_args,
)
def testConstruct(self):
big_query_example_gen = component.BigQueryExampleGen(query='query')
self.assertEqual(
standard_artifacts.Examples.TYPE_NAME,
big_query_example_gen.outputs[
standard_component_specs.EXAMPLES_KEY].type_name)
def create_pipeline(
pipeline_name: Text,
pipeline_root: Text,
module_file: Text,
ai_platform_training_args: Dict[Text, Text],
ai_platform_serving_args: Dict[Text, Text],
beam_pipeline_args: Optional[List[Text]] = None) -> pipeline.Pipeline:
"""Implements the chicago taxi pipeline with TFX and Kubeflow Pipelines.
Args:
pipeline_name: name of the TFX pipeline being created.
pipeline_root: root directory of the pipeline. Should be a valid GCS path.
module_file: uri of the module files used in Trainer and Transform
components.
ai_platform_training_args: Args of CAIP training job. Please refer to
https://cloud.google.com/ml-engine/reference/rest/v1/projects.jobs#Job
for detailed description.
ai_platform_serving_args: Args of CAIP model deployment. Please refer to
https://cloud.google.com/ml-engine/reference/rest/v1/projects.models
for detailed description.
beam_pipeline_args: Optional list of beam pipeline options. Please refer to
https://cloud.google.com/dataflow/docs/guides/specifying-exec-params#setting-other-cloud-dataflow-pipeline-options.
When this argument is not provided, the default is to use GCP
DataflowRunner with 50GB disk size as specified in this function. If an
empty list is passed in, default specified by Beam will be used, which can
be found at
https://cloud.google.com/dataflow/docs/guides/specifying-exec-params#setting-other-cloud-dataflow-pipeline-options
Returns:
A TFX pipeline object.
"""
# The rate at which to sample rows from the Taxi dataset using BigQuery.
# The full taxi dataset is > 200M record. In the interest of resource
# savings and time, we've set the default for this example to be much smaller.
# Feel free to crank it up and process the full dataset!
# By default it generates a 0.1% random sample.
query_sample_rate = data_types.RuntimeParameter(name='query_sample_rate',
ptype=float,
default=0.001)
# This is the upper bound of FARM_FINGERPRINT in Bigquery (ie the max value of
# signed int64).
max_int64 = '0x7FFFFFFFFFFFFFFF'
# The query that extracts the examples from BigQuery. The Chicago Taxi dataset
# used for this example is a public dataset available on Google AI Platform.
# https://console.cloud.google.com/marketplace/details/city-of-chicago-public-data/chicago-taxi-trips
query = """
SELECT
pickup_community_area,
fare,
EXTRACT(MONTH FROM trip_start_timestamp) AS trip_start_month,
EXTRACT(HOUR FROM trip_start_timestamp) AS trip_start_hour,
EXTRACT(DAYOFWEEK FROM trip_start_timestamp) AS trip_start_day,
UNIX_SECONDS(trip_start_timestamp) AS trip_start_timestamp,
pickup_latitude,
pickup_longitude,
dropoff_latitude,
dropoff_longitude,
trip_miles,
pickup_census_tract,
dropoff_census_tract,
payment_type,
company,
trip_seconds,
dropoff_community_area,
tips
FROM `bigquery-public-data.chicago_taxi_trips.taxi_trips`
WHERE (ABS(FARM_FINGERPRINT(unique_key)) / {max_int64})
< {query_sample_rate}""".format(
max_int64=max_int64, query_sample_rate=str(query_sample_rate))
# Beam args to run data processing on DataflowRunner.
#
# TODO(b/151114974): Remove `disk_size_gb` flag after default is increased.
# TODO(b/151116587): Remove `shuffle_mode` flag after default is changed.
# TODO(b/156874687): Remove `machine_type` after IP addresses are no longer a
# scaling bottleneck.
if beam_pipeline_args is None:
beam_pipeline_args = [
'--runner=DataflowRunner',
'--project=' + _project_id,
'--temp_location=' + os.path.join(_output_bucket, 'tmp'),
'--region=' + _gcp_region,
# Temporary overrides of defaults.
'--disk_size_gb=50',
'--experiments=shuffle_mode=auto',
'--machine_type=e2-standard-8',
]
# Number of epochs in training.
train_steps = data_types.RuntimeParameter(
name='train_steps',
default=10000,
ptype=int,
)
# Number of epochs in evaluation.
eval_steps = data_types.RuntimeParameter(
name='eval_steps',
default=5000,
ptype=int,
)
# Brings data into the pipeline or otherwise joins/converts training data.
example_gen = big_query_example_gen_component.BigQueryExampleGen(
query=query)
# Computes statistics over data for visualization and example validation.
statistics_gen = StatisticsGen(examples=example_gen.outputs['examples'])
# Generates schema based on statistics files.
schema_gen = SchemaGen(statistics=statistics_gen.outputs['statistics'],
infer_feature_shape=False)
# Performs anomaly detection based on statistics and data schema.
example_validator = ExampleValidator(
statistics=statistics_gen.outputs['statistics'],
schema=schema_gen.outputs['schema'])
# Performs transformations and feature engineering in training and serving.
transform = Transform(examples=example_gen.outputs['examples'],
schema=schema_gen.outputs['schema'],
module_file=module_file)
# Update ai_platform_training_args if distributed training was enabled.
# Number of worker machines used in distributed training.
worker_count = data_types.RuntimeParameter(
name='worker_count',
default=2,
ptype=int,
)
# Type of worker machines used in distributed training.
worker_type = data_types.RuntimeParameter(
name='worker_type',
default='standard',
ptype=str,
)
local_training_args = copy.deepcopy(ai_platform_training_args)
if FLAGS.distributed_training:
local_training_args.update({
# You can specify the machine types, the number of replicas for workers
# and parameter servers.
# https://cloud.google.com/ml-engine/reference/rest/v1/projects.jobs#ScaleTier
'scaleTier': 'CUSTOM',
'masterType': 'large_model',
'workerType': worker_type,
'parameterServerType': 'standard',
'workerCount': worker_count,
'parameterServerCount': 1
})
# Uses user-provided Python function that implements a model using TF-Learn
# to train a model on Google Cloud AI Platform.
trainer = Trainer(
custom_executor_spec=executor_spec.ExecutorClassSpec(
ai_platform_trainer_executor.Executor),
module_file=module_file,
transformed_examples=transform.outputs['transformed_examples'],
schema=schema_gen.outputs['schema'],
transform_graph=transform.outputs['transform_graph'],
train_args={'num_steps': train_steps},
eval_args={'num_steps': eval_steps},
custom_config={
ai_platform_trainer_executor.TRAINING_ARGS_KEY: local_training_args
})
# Get the latest blessed model for model validation.
model_resolver = ResolverNode(
instance_name='latest_blessed_model_resolver',
resolver_class=latest_blessed_model_resolver.
LatestBlessedModelResolver,
model=Channel(type=Model),
model_blessing=Channel(type=ModelBlessing))
# Uses TFMA to compute a evaluation statistics over features of a model and
# perform quality validation of a candidate model (compared to a baseline).
eval_config = tfma.EvalConfig(
model_specs=[tfma.ModelSpec(signature_name='eval')],
slicing_specs=[
tfma.SlicingSpec(),
tfma.SlicingSpec(feature_keys=['trip_start_hour'])
],
metrics_specs=[
tfma.MetricsSpec(
thresholds={
'accuracy':
tfma.config.MetricThreshold(
value_threshold=tfma.GenericValueThreshold(
lower_bound={'value': 0.6}),
change_threshold=tfma.GenericChangeThreshold(
direction=tfma.MetricDirection.HIGHER_IS_BETTER,
absolute={'value': -1e-10}))
})
])
evaluator = Evaluator(
examples=example_gen.outputs['examples'],
model=trainer.outputs['model'],
baseline_model=model_resolver.outputs['model'],
# Change threshold will be ignored if there is no baseline (first run).
eval_config=eval_config)
# Checks whether the model passed the validation steps and pushes the model
# to Google Cloud AI Platform if check passed.
# TODO(b/162451308): Add pusher back to components list once AIP Prediction
# Service supports TF>=2.3.
_ = Pusher(custom_executor_spec=executor_spec.ExecutorClassSpec(
ai_platform_pusher_executor.Executor),
model=trainer.outputs['model'],
model_blessing=evaluator.outputs['blessing'],
custom_config={
ai_platform_pusher_executor.SERVING_ARGS_KEY:
ai_platform_serving_args
})
return pipeline.Pipeline(
pipeline_name=pipeline_name,
pipeline_root=pipeline_root,
components=[
example_gen, statistics_gen, schema_gen, example_validator,
transform, trainer, model_resolver, evaluator
],
beam_pipeline_args=beam_pipeline_args,
)
def create_pipeline_components(
pipeline_root: Text,
transform_module: Text,
trainer_module: Text,
bigquery_query: Text = '',
csv_input_location: Text = '',
) -> List[base_node.BaseNode]:
"""Creates components for a simple Chicago Taxi TFX pipeline for testing.
Args:
pipeline_root: The root of the pipeline output.
transform_module: The location of the transform module file.
trainer_module: The location of the trainer module file.
bigquery_query: The query to get input data from BigQuery. If not empty,
BigQueryExampleGen will be used.
csv_input_location: The location of the input data directory.
Returns:
A list of TFX components that constitutes an end-to-end test pipeline.
"""
if bool(bigquery_query) == bool(csv_input_location):
raise ValueError(
'Exactly one example gen is expected. ',
'Please provide either bigquery_query or csv_input_location.')
if bigquery_query:
example_gen = big_query_example_gen_component.BigQueryExampleGen(
query=bigquery_query)
else:
example_gen = components.CsvExampleGen(input_base=csv_input_location)
statistics_gen = components.StatisticsGen(
examples=example_gen.outputs['examples'])
schema_gen = components.SchemaGen(
statistics=statistics_gen.outputs['statistics'],
infer_feature_shape=False)
example_validator = components.ExampleValidator(
statistics=statistics_gen.outputs['statistics'],
schema=schema_gen.outputs['schema'])
transform = components.Transform(examples=example_gen.outputs['examples'],
schema=schema_gen.outputs['schema'],
module_file=transform_module)
latest_model_resolver = resolver.Resolver(
strategy_class=latest_artifacts_resolver.LatestArtifactsResolver,
model=channel.Channel(type=standard_artifacts.Model)).with_id(
'Resolver.latest_model_resolver')
trainer = components.Trainer(
custom_executor_spec=executor_spec.ExecutorClassSpec(Executor),
transformed_examples=transform.outputs['transformed_examples'],
schema=schema_gen.outputs['schema'],
base_model=latest_model_resolver.outputs['model'],
transform_graph=transform.outputs['transform_graph'],
train_args=trainer_pb2.TrainArgs(num_steps=10),
eval_args=trainer_pb2.EvalArgs(num_steps=5),
module_file=trainer_module,
)
# Get the latest blessed model for model validation.
model_resolver = resolver.Resolver(
strategy_class=latest_blessed_model_resolver.
LatestBlessedModelResolver,
model=channel.Channel(type=standard_artifacts.Model),
model_blessing=channel.Channel(
type=standard_artifacts.ModelBlessing)).with_id(
'Resolver.latest_blessed_model_resolver')
# Set the TFMA config for Model Evaluation and Validation.
eval_config = tfma.EvalConfig(
model_specs=[tfma.ModelSpec(signature_name='eval')],
metrics_specs=[
tfma.MetricsSpec(
metrics=[tfma.MetricConfig(class_name='ExampleCount')],
thresholds={
'binary_accuracy':
tfma.MetricThreshold(
value_threshold=tfma.GenericValueThreshold(
lower_bound={'value': 0.5}),
change_threshold=tfma.GenericChangeThreshold(
direction=tfma.MetricDirection.HIGHER_IS_BETTER,
absolute={'value': -1e-10}))
})
],
slicing_specs=[
tfma.SlicingSpec(),
tfma.SlicingSpec(feature_keys=['trip_start_hour'])
])
evaluator = components.Evaluator(
examples=example_gen.outputs['examples'],
model=trainer.outputs['model'],
baseline_model=model_resolver.outputs['model'],
eval_config=eval_config)
pusher = components.Pusher(
model=trainer.outputs['model'],
model_blessing=evaluator.outputs['blessing'],
push_destination=pusher_pb2.PushDestination(
filesystem=pusher_pb2.PushDestination.Filesystem(
base_directory=os.path.join(pipeline_root, 'model_serving'))))
return [
example_gen, statistics_gen, schema_gen, example_validator, transform,
latest_model_resolver, trainer, model_resolver, evaluator, pusher
]
def _create_pipeline(
pipeline_name: Text, pipeline_root: Text, query: Text, module_file: Text,
beam_pipeline_args: List[Text], ai_platform_training_args: Dict[Text, Text],
bigquery_serving_args: Dict[Text, Text]) -> pipeline.Pipeline:
"""Implements the chicago taxi pipeline with TFX and Kubeflow Pipelines."""
# Brings data into the pipeline or otherwise joins/converts training data.
example_gen = big_query_example_gen_component.BigQueryExampleGen(query=query)
# Computes statistics over data for visualization and example validation.
statistics_gen = StatisticsGen(examples=example_gen.outputs['examples'])
# Generates schema based on statistics files.
schema_gen = SchemaGen(
statistics=statistics_gen.outputs['statistics'], infer_feature_shape=True)
# Performs anomaly detection based on statistics and data schema.
example_validator = ExampleValidator(
statistics=statistics_gen.outputs['statistics'],
schema=schema_gen.outputs['schema'])
# Performs transformations and feature engineering in training and serving.
transform = Transform(
examples=example_gen.outputs['examples'],
schema=schema_gen.outputs['schema'],
module_file=module_file)
# Uses user-provided Python function that implements a model using TF-Learn
# to train a model on Google Cloud AI Platform.
trainer = Trainer(
custom_executor_spec=executor_spec.ExecutorClassSpec(
ai_platform_trainer_executor.Executor),
module_file=module_file,
transformed_examples=transform.outputs['transformed_examples'],
schema=schema_gen.outputs['schema'],
transform_graph=transform.outputs['transform_graph'],
train_args=trainer_pb2.TrainArgs(num_steps=10000),
eval_args=trainer_pb2.EvalArgs(num_steps=5000),
custom_config={'ai_platform_training_args': ai_platform_training_args})
# Uses TFMA to compute a evaluation statistics over features of a model.
evaluator = Evaluator(
examples=example_gen.outputs['examples'],
model=trainer.outputs['model'],
feature_slicing_spec=evaluator_pb2.FeatureSlicingSpec(specs=[
evaluator_pb2.SingleSlicingSpec(
column_for_slicing=['trip_start_hour'])
]))
# Performs quality validation of a candidate model (compared to a baseline).
model_validator = ModelValidator(
examples=example_gen.outputs['examples'], model=trainer.outputs['model'])
# Checks whether the model passed the validation steps and pushes the model
# to Google Cloud BigQuery ML if check passed.
pusher = Pusher(
custom_executor_spec=executor_spec.ExecutorClassSpec(
bigquery_pusher_executor.Executor),
model=trainer.outputs['model'],
model_blessing=model_validator.outputs['blessing'],
custom_config={'bigquery_serving_args': bigquery_serving_args})
return pipeline.Pipeline(
pipeline_name=pipeline_name,
pipeline_root=pipeline_root,
components=[
example_gen, statistics_gen, schema_gen, example_validator, transform,
trainer, evaluator, model_validator, pusher
],
beam_pipeline_args=beam_pipeline_args,
)
def testConstruct(self):
big_query_example_gen = component.BigQueryExampleGen(query='query')
self.assertEqual(standard_artifacts.Examples.TYPE_NAME,
big_query_example_gen.outputs['examples'].type_name)
artifact_collection = big_query_example_gen.outputs['examples'].get()
self.assertEqual(1, len(artifact_collection))
