Fix: Vertex ML pipeline test failures (#7727)

nikelite · web-flow · commit cfab8a4c487e · 2024-12-05T07:40:05.000+09:00
* Update wheel build to include data files for Vertex tests
* Fix Keras Tuner compatibility issue with the example Keras model
* Update the taxi template model to avoid using deprecated Keras APIs
* Restore exit handlers that were removed with KFP v1 support
diff --git a/package_build/initialize.sh b/package_build/initialize.sh
@@ -27,6 +27,7 @@ do
   ln -sf $BASEDIR/setup.py $BASEDIR/package_build/$CONFIG_NAME/
   ln -sf $BASEDIR/dist $BASEDIR/package_build/$CONFIG_NAME/
   ln -sf $BASEDIR/tfx $BASEDIR/package_build/$CONFIG_NAME/
+  ln -sf $BASEDIR/MANIFEST.in $BASEDIR/package_build/$CONFIG_NAME/
   ln -sf $BASEDIR/README*.md $BASEDIR/package_build/$CONFIG_NAME/
   ln -sf $BASEDIR/LICENSE $BASEDIR/package_build/$CONFIG_NAME/
 
diff --git a/tfx/examples/chicago_taxi_pipeline/taxi_utils.py b/tfx/examples/chicago_taxi_pipeline/taxi_utils.py
@@ -246,7 +246,6 @@ def _build_keras_model(
   output = tf.keras.layers.Dense(1, activation='sigmoid')(
       tf.keras.layers.concatenate([deep, wide])
   )
-  output = tf.squeeze(output, -1)
 
   model = tf.keras.Model(input_layers, output)
   model.compile(
@@ -371,4 +370,4 @@ def run_fn(fn_args: fn_args_utils.FnArgs):
           model, tf_transform_output
       ),
   }
-  model.save(fn_args.serving_model_dir, save_format='tf', signatures=signatures)
+  tf.saved_model.save(model, fn_args.serving_model_dir, signatures=signatures)
diff --git a/tfx/experimental/templates/taxi/models/keras_model/model.py b/tfx/experimental/templates/taxi/models/keras_model/model.py
@@ -215,4 +215,4 @@ def run_fn(fn_args):
       'transform_features':
           _get_transform_features_signature(model, tf_transform_output),
   }
-  model.save(fn_args.serving_model_dir, save_format='tf', signatures=signatures)
+  tf.saved_model.save(model, fn_args.serving_model_dir, signatures=signatures)
diff --git a/tfx/orchestration/kubeflow/v2/test_utils.py b/tfx/orchestration/kubeflow/v2/test_utils.py
@@ -234,25 +234,28 @@ def create_pipeline_components(
       model_blessing=tfx.dsl.Channel(
           type=tfx.types.standard_artifacts.ModelBlessing)).with_id(
               'Resolver.latest_blessed_model_resolver')
-  # Set the TFMA config for Model Evaluation and Validation.
+  # 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')],
-      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}))
-              })
+      model_specs=[
+          tfma.ModelSpec(
+              signature_name='serving_default', label_key='tips_xf',
+              preprocessing_function_names=['transform_features'])
       ],
-      slicing_specs=[
-          tfma.SlicingSpec(),
-          tfma.SlicingSpec(feature_keys=['trip_start_hour'])
+      slicing_specs=[tfma.SlicingSpec()],
+      metrics_specs=[
+          tfma.MetricsSpec(metrics=[
+              tfma.MetricConfig(
+                  class_name='BinaryAccuracy',
+                  threshold=tfma.MetricThreshold(
+                      value_threshold=tfma.GenericValueThreshold(
+                          lower_bound={'value': 0.6}),
+                      # Change threshold will be ignored if there is no
+                      # baseline model resolved from MLMD (first run).
+                      change_threshold=tfma.GenericChangeThreshold(
+                          direction=tfma.MetricDirection.HIGHER_IS_BETTER,
+                          absolute={'value': -1e-10})))
+          ])
       ])
   evaluator = tfx.components.Evaluator(
       examples=example_gen.outputs['examples'],
diff --git a/tfx/orchestration/kubeflow/v2/testdata/expected_full_taxi_pipeline_job.json b/tfx/orchestration/kubeflow/v2/testdata/expected_full_taxi_pipeline_job.json
@@ -706,7 +706,7 @@
               "parameters": {
                 "eval_config": {
                   "runtimeValue": {
-                    "constant": "{\n  \"metrics_specs\": [\n    {\n      \"metrics\": [\n        {\n          \"class_name\": \"ExampleCount\"\n        }\n      ],\n      \"thresholds\": {\n        \"binary_accuracy\": {\n          \"change_threshold\": {\n            \"absolute\": -1e-10,\n            \"direction\": \"HIGHER_IS_BETTER\"\n          },\n          \"value_threshold\": {\n            \"lower_bound\": 0.5\n          }\n        }\n      }\n    }\n  ],\n  \"model_specs\": [\n    {\n      \"signature_name\": \"eval\"\n    }\n  ],\n  \"slicing_specs\": [\n    {},\n    {\n      \"feature_keys\": [\n        \"trip_start_hour\"\n      ]\n    }\n  ]\n}"
+                    "constant": "{\n  \"metrics_specs\": [\n    {\n      \"metrics\": [\n        {\n          \"class_name\": \"BinaryAccuracy\",\n          \"threshold\": {\n            \"change_threshold\": {\n              \"absolute\": -1e-10,\n              \"direction\": \"HIGHER_IS_BETTER\"\n            },\n            \"value_threshold\": {\n              \"lower_bound\": 0.6\n            }\n          }\n        }\n      ]\n    }\n  ],\n  \"model_specs\": [\n    {\n      \"label_key\": \"tips_xf\",\n      \"preprocessing_function_names\": [\n        \"transform_features\"\n      ],\n      \"signature_name\": \"serving_default\"\n    }\n  ],\n  \"slicing_specs\": [\n    {}\n  ]\n}"
                   }
                 },
                 "example_splits": {
diff --git a/tfx/orchestration/kubeflow/v2/testdata/legacy/expected_full_taxi_pipeline_job.json b/tfx/orchestration/kubeflow/v2/testdata/legacy/expected_full_taxi_pipeline_job.json
@@ -698,7 +698,7 @@
                 "eval_config": {
                   "runtimeValue": {
                     "constantValue": {
-                      "stringValue": "{\n  \"metrics_specs\": [\n    {\n      \"metrics\": [\n        {\n          \"class_name\": \"ExampleCount\"\n        }\n      ],\n      \"thresholds\": {\n        \"binary_accuracy\": {\n          \"change_threshold\": {\n            \"absolute\": -1e-10,\n            \"direction\": \"HIGHER_IS_BETTER\"\n          },\n          \"value_threshold\": {\n            \"lower_bound\": 0.5\n          }\n        }\n      }\n    }\n  ],\n  \"model_specs\": [\n    {\n      \"signature_name\": \"eval\"\n    }\n  ],\n  \"slicing_specs\": [\n    {},\n    {\n      \"feature_keys\": [\n        \"trip_start_hour\"\n      ]\n    }\n  ]\n}"
+                      "stringValue": "{\n  \"metrics_specs\": [\n    {\n      \"metrics\": [\n        {\n          \"class_name\": \"BinaryAccuracy\",\n          \"threshold\": {\n            \"change_threshold\": {\n              \"absolute\": -1e-10,\n              \"direction\": \"HIGHER_IS_BETTER\"\n            },\n            \"value_threshold\": {\n              \"lower_bound\": 0.6\n            }\n          }\n        }\n      ]\n    }\n  ],\n  \"model_specs\": [\n    {\n      \"label_key\": \"tips_xf\",\n      \"preprocessing_function_names\": [\n        \"transform_features\"\n      ],\n      \"signature_name\": \"serving_default\"\n    }\n  ],\n  \"slicing_specs\": [\n    {}\n  ]\n}"
                     }
                   }
                 },
diff --git a/tfx/tools/docker/Dockerfile b/tfx/tools/docker/Dockerfile
@@ -27,8 +27,7 @@ WORKDIR ${TFX_DIR}
 ARG TFX_DEPENDENCY_SELECTOR
 ENV TFX_DEPENDENCY_SELECTOR=${TFX_DEPENDENCY_SELECTOR}
 
-RUN python -m pip install --upgrade pip wheel setuptools
-RUN python -m pip install tomli
+RUN python -m pip install --upgrade pip wheel setuptools tomli
 
 # TODO(b/175089240): clean up conditional checks on whether ml-pipelines-sdk is
 # built after TFX versions <= 0.25 are no longer eligible for cherry-picks.
diff --git a/tfx/v1/orchestration/experimental/__init__.py b/tfx/v1/orchestration/experimental/__init__.py
@@ -14,6 +14,9 @@
 """TFX orchestration.experimental module."""
 
 try:
+    from tfx.orchestration.kubeflow.decorators import exit_handler  # pylint: disable=g-import-not-at-top
+    from tfx.orchestration.kubeflow.decorators import FinalStatusStr  # pylint: disable=g-import-not-at-top
+
     from tfx.orchestration.kubeflow.v2.kubeflow_v2_dag_runner import (
         KubeflowV2DagRunner,
         KubeflowV2DagRunnerConfig,
@@ -24,11 +27,8 @@
 
 __all__ = [
     "FinalStatusStr",
-    "KubeflowDagRunner",
-    "KubeflowDagRunnerConfig",
     "KubeflowV2DagRunner",
     "KubeflowV2DagRunnerConfig",
-    "LABEL_KFP_SDK_ENV",
     "exit_handler",
     "get_default_kubeflow_metadata_config",
 ]

Original file line number	Diff line number	Diff line change
`@@ -246,7 +246,6 @@ def _build_keras_model(`
`246`	`246`	`output = tf.keras.layers.Dense(1, activation='sigmoid')(`
`247`	`247`	`tf.keras.layers.concatenate([deep, wide])`
`248`	`248`	`)`
`249`		`- output = tf.squeeze(output, -1)`
`250`	`249`
`251`	`250`	`model = tf.keras.Model(input_layers, output)`
`252`	`251`	`model.compile(`
`@@ -371,4 +370,4 @@ def run_fn(fn_args: fn_args_utils.FnArgs):`
`371`	`370`	`model, tf_transform_output`
`372`	`371`	`),`
`373`	`372`	`}`
`374`		`- model.save(fn_args.serving_model_dir, save_format='tf', signatures=signatures)`
	`373`	`+ tf.saved_model.save(model, fn_args.serving_model_dir, signatures=signatures)`
Original file line number	Diff line number	Diff line change
`@@ -215,4 +215,4 @@ def run_fn(fn_args):`
`215`	`215`	`'transform_features':`
`216`	`216`	`_get_transform_features_signature(model, tf_transform_output),`
`217`	`217`	`}`
`218`		`- model.save(fn_args.serving_model_dir, save_format='tf', signatures=signatures)`
	`218`	`+ tf.saved_model.save(model, fn_args.serving_model_dir, signatures=signatures)`
Original file line number	Diff line number	Diff line change
`@@ -706,7 +706,7 @@`
`706`	`706`	`"parameters": {`
`707`	`707`	`"eval_config": {`
`708`	`708`	`"runtimeValue": {`
`709`		- "constant": "{\n \"metrics_specs\": [\n {\n \"metrics\": [\n {\n \"class_name\": \"ExampleCount\"\n }\n ],\n \"thresholds\": {\n \"binary_accuracy\": {\n \"change_threshold\": {\n \"absolute\": -1e-10,\n \"direction\": \"HIGHER_IS_BETTER\"\n },\n \"value_threshold\": {\n \"lower_bound\": 0.5\n }\n }\n }\n }\n ],\n \"model_specs\": [\n {\n \"signature_name\": \"eval\"\n }\n ],\n \"slicing_specs\": [\n {},\n {\n \"feature_keys\": [\n \"trip_start_hour\"\n ]\n }\n ]\n}"
	`709`	+ "constant": "{\n \"metrics_specs\": [\n {\n \"metrics\": [\n {\n \"class_name\": \"BinaryAccuracy\",\n \"threshold\": {\n \"change_threshold\": {\n \"absolute\": -1e-10,\n \"direction\": \"HIGHER_IS_BETTER\"\n },\n \"value_threshold\": {\n \"lower_bound\": 0.6\n }\n }\n }\n ]\n }\n ],\n \"model_specs\": [\n {\n \"label_key\": \"tips_xf\",\n \"preprocessing_function_names\": [\n \"transform_features\"\n ],\n \"signature_name\": \"serving_default\"\n }\n ],\n \"slicing_specs\": [\n {}\n ]\n}"
`710`	`710`	`}`
`711`	`711`	`},`
`712`	`712`	`"example_splits": {`
Original file line number	Diff line number	Diff line change
`@@ -698,7 +698,7 @@`
`698`	`698`	`"eval_config": {`
`699`	`699`	`"runtimeValue": {`
`700`	`700`	`"constantValue": {`
`701`		- "stringValue": "{\n \"metrics_specs\": [\n {\n \"metrics\": [\n {\n \"class_name\": \"ExampleCount\"\n }\n ],\n \"thresholds\": {\n \"binary_accuracy\": {\n \"change_threshold\": {\n \"absolute\": -1e-10,\n \"direction\": \"HIGHER_IS_BETTER\"\n },\n \"value_threshold\": {\n \"lower_bound\": 0.5\n }\n }\n }\n }\n ],\n \"model_specs\": [\n {\n \"signature_name\": \"eval\"\n }\n ],\n \"slicing_specs\": [\n {},\n {\n \"feature_keys\": [\n \"trip_start_hour\"\n ]\n }\n ]\n}"
	`701`	+ "stringValue": "{\n \"metrics_specs\": [\n {\n \"metrics\": [\n {\n \"class_name\": \"BinaryAccuracy\",\n \"threshold\": {\n \"change_threshold\": {\n \"absolute\": -1e-10,\n \"direction\": \"HIGHER_IS_BETTER\"\n },\n \"value_threshold\": {\n \"lower_bound\": 0.6\n }\n }\n }\n ]\n }\n ],\n \"model_specs\": [\n {\n \"label_key\": \"tips_xf\",\n \"preprocessing_function_names\": [\n \"transform_features\"\n ],\n \"signature_name\": \"serving_default\"\n }\n ],\n \"slicing_specs\": [\n {}\n ]\n}"
`702`	`702`	`}`
`703`	`703`	`}`
`704`	`704`	`},`