Deprecate KFP v1 support

Author: nikelite

build/BUILD

@@ -25,7 +25,6 @@ sh_binary(
         "//tfx/extensions/experimental/kfp_compatibility/proto:kfp_component_spec_pb2.py",
         "//tfx/extensions/google_cloud_big_query/experimental/elwc_example_gen/proto:elwc_config_pb2.py",
         "//tfx/orchestration/experimental/core:component_generated_alert_pb2.py",
-        "//tfx/orchestration/kubeflow/proto:kubeflow_pb2.py",
         "//tfx/proto:bulk_inferrer_pb2.py",
         "//tfx/proto:distribution_validator_pb2.py",
         "//tfx/proto:evaluator_pb2.py",

tfx/dependencies.py

@@ -69,9 +69,8 @@ def make_pipeline_sdk_required_install_packages():
       # TODO(b/176812386): Deprecate usage of jinja2 for placeholders.
       'jinja2>=2.7.3,<4',
       # typing-extensions allows consistent & future-proof interface for typing.
-      # Since kfp<2 uses typing-extensions<4, lower bound is the latest 3.x, and
-      # upper bound is <5 as the semver started from 4.0 according to their doc.
-      'typing-extensions>=3.10.0.2,<5',
+      # Upper bound is <5 as the semver started from 4.0 according to their doc.
+      'typing-extensions<5',
   ]
 
 
@@ -87,7 +86,7 @@ def make_required_install_packages():
       'google-cloud-bigquery>=3,<4',
       'grpcio>=1.28.1,<2',
       'keras-tuner>=1.0.4,<2,!=1.4.0,!=1.4.1',
-      'kubernetes>=10.0.1,<13',
+      'kubernetes>=10.0.1,<27',
       'numpy>=1.16,<2',
       'pyarrow>=10,<11',
       # TODO: b/358471141 - Orjson 3.10.7 breaks TFX OSS tests.
@@ -146,27 +145,29 @@ def make_extra_packages_airflow():
 def make_extra_packages_kfp():
   """Prepare extra packages needed for Kubeflow Pipelines orchestrator."""
   return [
-      # TODO(b/304892416): Migrate from KFP SDK v1 to v2.
-      'kfp>=1.8.14,<2',
-      'kfp-pipeline-spec>0.1.13,<0.2',
+      'kfp>=2',
+      'kfp-pipeline-spec>=0.2.2',
   ]
 
 
 def make_extra_packages_test():
   """Prepare extra packages needed for running unit tests."""
   # Note: It is okay to pin packages to exact versions in this list to minimize
   # conflicts.
-  return make_extra_packages_airflow() + make_extra_packages_kfp() + [
-      'pytest>=5,<7',
-  ]
+  return (
+      make_extra_packages_airflow()
+      + make_extra_packages_kfp()
+      + [
+          'pytest>=5,<7',
+      ]
+  )
 
 
 def make_extra_packages_docker_image():
   # Packages needed for tfx docker image.
   return [
-      # TODO(b/304892416): Migrate from KFP SDK v1 to v2.
-      'kfp>=1.8.14,<2',
-      'kfp-pipeline-spec>0.1.13,<0.2',
+      'kfp>=2',
+      'kfp-pipeline-spec>=0.3.0',
       'mmh>=2.2,<3',
       'python-snappy>=0.5,<0.6',
       # Required for tfx/examples/penguin/penguin_utils_cloud_tuner.py
@@ -194,10 +195,12 @@ def make_extra_packages_tf_ranking():
   # Packages needed for tf-ranking which is used in tfx/examples/ranking.
   return [
       'tensorflow-ranking>=0.5,<0.6',
-      'struct2tensor' + select_constraint(
+      'struct2tensor'
+      + select_constraint(
           default='>=0.46.0,<0.47.0',
           nightly='>=0.47.0.dev',
-          git_master='@git+https://github.com/google/struct2tensor@master'),
+          git_master='@git+https://github.com/google/struct2tensor@master',
+      ),
   ]
 
 

tfx/examples/penguin/penguin_pipeline_kubeflow.py

@@ -501,33 +501,28 @@ def main():
   else:
     beam_pipeline_args = _beam_pipeline_args_by_runner['DirectRunner']
 
-  if use_vertex:
-    dag_runner = tfx.orchestration.experimental.KubeflowV2DagRunner(
-        config=tfx.orchestration.experimental.KubeflowV2DagRunnerConfig(),
-        output_filename=_pipeline_definition_file)
-  else:
-    dag_runner = tfx.orchestration.experimental.KubeflowDagRunner(
-        config=tfx.orchestration.experimental.KubeflowDagRunnerConfig(
-            kubeflow_metadata_config=tfx.orchestration.experimental
-            .get_default_kubeflow_metadata_config()))
-
-    dag_runner.run(
-        create_pipeline(
-            pipeline_name=_pipeline_name,
-            pipeline_root=_pipeline_root,
-            data_root=_data_root,
-            module_file=_module_file,
-            enable_tuning=False,
-            enable_cache=True,
-            user_provided_schema_path=_user_provided_schema,
-            ai_platform_training_args=_ai_platform_training_args,
-            ai_platform_serving_args=_ai_platform_serving_args,
-            beam_pipeline_args=beam_pipeline_args,
-            use_cloud_component=use_cloud_component,
-            use_aip=use_aip,
-            use_vertex=use_vertex,
-            serving_model_dir=_serving_model_dir,
-        ))
+  dag_runner = tfx.orchestration.experimental.KubeflowV2DagRunner(
+      config=tfx.orchestration.experimental.KubeflowV2DagRunnerConfig(),
+      output_filename=_pipeline_definition_file,
+  )
+  dag_runner.run(
+      create_pipeline(
+          pipeline_name=_pipeline_name,
+          pipeline_root=_pipeline_root,
+          data_root=_data_root,
+          module_file=_module_file,
+          enable_tuning=False,
+          enable_cache=True,
+          user_provided_schema_path=_user_provided_schema,
+          ai_platform_training_args=_ai_platform_training_args,
+          ai_platform_serving_args=_ai_platform_serving_args,
+          beam_pipeline_args=beam_pipeline_args,
+          use_cloud_component=use_cloud_component,
+          use_aip=use_aip,
+          use_vertex=use_vertex,
+          serving_model_dir=_serving_model_dir,
+      )
+  )
 
 
 # To compile the pipeline:

tfx/examples/penguin/penguin_pipeline_kubeflow_test.py

@@ -64,24 +64,16 @@ def testPenguinPipelineConstructionAndDefinitionFileExists(
         serving_model_dir=penguin_pipeline_kubeflow._serving_model_dir)
     self.assertLen(kubeflow_pipeline.components, 9)
 
-    if use_vertex:
-      v2_dag_runner = orchestration.experimental.KubeflowV2DagRunner(
-          config=orchestration.experimental.KubeflowV2DagRunnerConfig(),
-          output_dir=self.tmp_dir,
-          output_filename=penguin_pipeline_kubeflow._pipeline_definition_file)
-      v2_dag_runner.run(kubeflow_pipeline)
-      file_path = os.path.join(
-          self.tmp_dir, penguin_pipeline_kubeflow._pipeline_definition_file)
-      self.assertTrue(fileio.exists(file_path))
-    else:
-      v1_dag_runner = orchestration.experimental.KubeflowDagRunner(
-          config=orchestration.experimental.KubeflowDagRunnerConfig(
-              kubeflow_metadata_config=orchestration.experimental
-              .get_default_kubeflow_metadata_config()))
-      v1_dag_runner.run(kubeflow_pipeline)
-      file_path = os.path.join(self.tmp_dir, 'penguin-kubeflow.tar.gz')
-      self.assertTrue(fileio.exists(file_path))
-
+    v2_dag_runner = orchestration.experimental.KubeflowV2DagRunner(
+        config=orchestration.experimental.KubeflowV2DagRunnerConfig(),
+        output_dir=self.tmp_dir,
+        output_filename=penguin_pipeline_kubeflow._pipeline_definition_file,
+    )
+    v2_dag_runner.run(kubeflow_pipeline)
+    file_path = os.path.join(
+        self.tmp_dir, penguin_pipeline_kubeflow._pipeline_definition_file
+    )
+    self.assertTrue(fileio.exists(file_path))
 
 if __name__ == '__main__':
   tf.test.main()

tfx/experimental/templates/taxi/kubeflow_runner.py

@@ -145,7 +145,7 @@ def component_run_context_name(self) -> str:
 class RuntimeParameter(json_utils.Jsonable):
   """Runtime parameter.
 
-  Currently only supported on KubeflowDagRunner.
+  Currently only supported on KubeflowV2DagRunner.
 
   For protos, use text type RuntimeParameter, which holds the proto json
   string, e.g., `'{"num_steps": 5}'` for TrainArgs proto.