Repository: kubeflow/pipelines

Issue:
[feature] dsl-compile supports v2 compatible mode
### Feature Area

<!-- Uncomment the labels below which are relevant to this feature: -->
<!-- /area frontend -->
<!-- /area backend -->
/area sdk
<!-- /area samples -->
<!-- /area components -->


### What feature would you like to see?

Add a flag to enable v2 compatible mode when using dsl-compile.
Ideally, v2 compatible mode can be enabled using an env var, so that we can switch to v2 compatible mode and do not need an explicit flag each time.

---

<!-- Don't delete message below to encourage users to support your feature request! -->
Love this idea? Give it a 👍. We prioritize fulfilling features with the most 👍.



Discussion:

Resolving patch:
diff --git a/sdk/python/kfp/compiler/compiler.py b/sdk/python/kfp/compiler/compiler.py
--- a/sdk/python/kfp/compiler/compiler.py
+++ b/sdk/python/kfp/compiler/compiler.py
@@ -59,12 +59,12 @@ def my_pipeline(a: int = 1, b: str = "default value"):
 
   def __init__(
       self,
-      mode: dsl.PipelineExecutionMode = dsl.PipelineExecutionMode.V1_LEGACY,
+      mode: dsl.PipelineExecutionMode = kfp.dsl.PipelineExecutionMode.V1_LEGACY,
       launcher_image: Optional[str] = None):
     """Creates a KFP compiler for compiling pipeline functions for execution.
 
     Args:
-      mode: The pipeline execution mode to use.
+      mode: The pipeline execution mode to use, defaults to kfp.dsl.PipelineExecutionMode.V1_LEGACY.
       launcher_image: Configurable image for KFP launcher to use. Only applies
         when `mode == dsl.PipelineExecutionMode.V2_COMPATIBLE`. Should only be
         needed for tests or custom deployments right now.
diff --git a/sdk/python/kfp/compiler/main.py b/sdk/python/kfp/compiler/main.py
--- a/sdk/python/kfp/compiler/main.py
+++ b/sdk/python/kfp/compiler/main.py
@@ -14,15 +14,15 @@
 
 
 import argparse
+from typing import Optional
 from kfp import dsl
 import kfp.compiler
 import os
-import shutil
-import subprocess
 import sys
-import tempfile
 from deprecated.sphinx import deprecated
 
+_KF_PIPELINES_COMPILER_MODE_ENV = 'KF_PIPELINES_COMPILER_MODE'
+
 
 def parse_arguments():
   """Parse command line arguments."""
@@ -44,12 +44,15 @@ def parse_arguments():
   parser.add_argument('--disable-type-check',
                       action='store_true',
                       help='disable the type check, default is enabled.')
+  parser.add_argument('--mode',
+                      type=str,
+                      help='compiler mode, defaults to V1, can also be V2_COMPATIBLE. You can override the default using env var KF_PIPELINES_COMPILER_MODE.')
 
   args = parser.parse_args()
   return args
 
 
-def _compile_pipeline_function(pipeline_funcs, function_name, output_path, type_check):
+def _compile_pipeline_function(pipeline_funcs, function_name, output_path, type_check, mode: Optional[dsl.PipelineExecutionMode]):
   if len(pipeline_funcs) == 0:
     raise ValueError('A function with @dsl.pipeline decorator is required in the py file.')
 
@@ -65,7 +68,7 @@ def _compile_pipeline_function(pipeline_funcs, function_name, output_path, type_
   else:
     pipeline_func = pipeline_funcs[0]
 
-  kfp.compiler.Compiler().compile(pipeline_func, output_path, type_check)
+  kfp.compiler.Compiler(mode=mode).compile(pipeline_func, output_path, type_check)
 
 
 class PipelineCollectorContext():
@@ -82,24 +85,36 @@ def __exit__(self, *args):
     dsl._pipeline._pipeline_decorator_handler = self.old_handler
 
 
-def compile_pyfile(pyfile, function_name, output_path, type_check):
+def compile_pyfile(pyfile, function_name, output_path, type_check, mode: Optional[dsl.PipelineExecutionMode]):
   sys.path.insert(0, os.path.dirname(pyfile))
   try:
     filename = os.path.basename(pyfile)
     with PipelineCollectorContext() as pipeline_funcs:
       __import__(os.path.splitext(filename)[0])
-    _compile_pipeline_function(pipeline_funcs, function_name, output_path, type_check)
+    _compile_pipeline_function(pipeline_funcs, function_name, output_path, type_check, mode)
   finally:
     del sys.path[0]
 
-
 def main():
   args = parse_arguments()
   if args.py is None:
     raise ValueError('The --py option must be specified.')
+  mode_str = args.mode
+  if not mode_str:
+    mode_str = os.environ.get(_KF_PIPELINES_COMPILER_MODE_ENV, 'V1')
+  mode = None
+  if mode_str == 'V1_LEGACY' or mode_str == 'V1':
+    mode = kfp.dsl.PipelineExecutionMode.V1_LEGACY
+  elif mode_str == 'V2_COMPATIBLE':
+    mode = kfp.dsl.PipelineExecutionMode.V2_COMPATIBLE
+  elif mode_str == 'V2_ENGINE':
+    mode = kfp.dsl.PipelineExecutionMode.V2_ENGINE
+  else:
+    raise ValueError(f'Got unexpected --mode option "{mode_str}", it must be one of V1, V2_COMPATIBLE or V2_ENGINE')
   compile_pyfile(
       args.py,
       args.function,
       args.output,
       not args.disable_type_check,
+      mode,
   )


Test patch:
diff --git a/sdk/python/tests/compiler/compiler_tests.py b/sdk/python/tests/compiler/compiler_tests.py
--- a/sdk/python/tests/compiler/compiler_tests.py
+++ b/sdk/python/tests/compiler/compiler_tests.py
@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import List
+from typing import Optional
 
 import kfp
 import kfp.compiler as compiler
@@ -24,9 +24,10 @@
 import zipfile
 import tarfile
 import tempfile
-import unittest
+import mock
 import yaml
 
+from absl.testing import parameterized
 from kfp.compiler import Compiler
 from kfp.dsl._component import component
 from kfp.dsl import ContainerOp, pipeline, PipelineParam
@@ -43,7 +44,7 @@ def some_op():
   )
 
 
-class TestCompiler(unittest.TestCase):
+class TestCompiler(parameterized.TestCase):
   # Define the places of samples covered by unit tests.
   core_sample_path = os.path.join(os.path.dirname(__file__), '..', '..', '..',
                                   '..', 'samples', 'core',)
@@ -383,6 +384,59 @@ def test_py_volume(self):
     """Test a pipeline with a volume and volume mount."""
     self._test_py_compile_yaml('volume')
 
+  @parameterized.parameters(
+    {'mode': 'V2_COMPATIBLE', 'is_v2': True},
+    {'mode': 'V1', 'is_v2': False},
+    {'mode': 'V1_LEGACY', 'is_v2': False},
+    {'mode': None, 'is_v2': False},
+    {'mode': 'V2_COMPATIBLE', 'env': 'V1', 'is_v2': True},
+    {'mode': None, 'env': 'V1', 'is_v2': False},
+    {'mode': None, 'env': 'V2_COMPATIBLE', 'is_v2': True},
+    {'mode': None, 'env': 'V1_LEGACY', 'is_v2': False},
+    {'mode': 'INVALID', 'error': True},
+    {'mode': None, 'env': 'INVALID', 'error': True},
+  )
+  def test_dsl_compile_mode(self, mode: Optional[str] = None, is_v2: Optional[bool] = None, env: Optional[str] = None, error: Optional[bool] = None):
+    with mock.patch.dict(os.environ, env and {'KF_PIPELINES_COMPILER_MODE': env} or {}):
+      file_base_name = 'two_step'
+      test_data_dir = os.path.join(os.path.dirname(__file__), 'testdata')
+      py_file = os.path.join(test_data_dir, f'{file_base_name}.py')
+      tmpdir = tempfile.mkdtemp()
+      try:
+        target_yaml = os.path.join(tmpdir, f'{file_base_name}.yaml')
+        args = ['dsl-compile', '--py', py_file, '--output', target_yaml]
+        if mode:
+          args = args + ['--mode', mode]
+        got_error = None
+        compiled = None
+        try:
+          subprocess.check_output(args)
+          with open(target_yaml, 'r') as f:
+            compiled = yaml.safe_load(f)
+        except subprocess.CalledProcessError as err:
+          got_error = err
+        if error:
+          if not got_error:
+            self.fail(f'expected error, but succeeded')
+        else:
+          if got_error:
+            self.fail(f'expected success, but got {got_error}')
+          v2_pipeline_annotation = compiled['metadata']['annotations'].get('pipelines.kubeflow.org/v2_pipeline')
+          if is_v2:
+            self.assertEqual(
+              'true',
+              v2_pipeline_annotation,
+              f'expected to compile in v2_compatible mode'
+            )
+          else:
+            self.assertEqual(
+              None,
+              v2_pipeline_annotation,
+              f'expected to compile in v1 mode'
+            )
+      finally:
+        shutil.rmtree(tmpdir)
+
   def test_py_retry_policy(self):
       """Test retry policy is set."""
 
diff --git a/sdk/python/tests/compiler/testdata/two_step.py b/sdk/python/tests/compiler/testdata/two_step.py
new file mode 100644
--- /dev/null
+++ b/sdk/python/tests/compiler/testdata/two_step.py
@@ -0,0 +1,59 @@
+# Copyright 2021 The Kubeflow Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Two step v2-compatible pipeline."""
+
+from kfp import components, dsl
+from kfp.components import InputPath, OutputPath
+
+
+def preprocess(
+    uri: str, some_int: int, output_parameter_one: OutputPath(int),
+    output_dataset_one: OutputPath('Dataset')
+):
+    '''Dummy Preprocess Step.'''
+    with open(output_dataset_one, 'w') as f:
+        f.write('Output dataset')
+    with open(output_parameter_one, 'w') as f:
+        f.write("{}".format(1234))
+
+
+preprocess_op = components.create_component_from_func(
+    preprocess, base_image='python:3.9'
+)
+
+
+@components.create_component_from_func
+def train_op(
+    dataset: InputPath('Dataset'),
+    model: OutputPath('Model'),
+    num_steps: int = 100
+):
+    '''Dummy Training Step.'''
+
+    with open(dataset, 'r') as input_file:
+        input_string = input_file.read()
+        with open(model, 'w') as output_file:
+            for i in range(num_steps):
+                output_file.write(
+                    "Step {}\n{}\n=====\n".format(i, input_string)
+                )
+
+
+@dsl.pipeline(name='two_step_pipeline')
+def two_step_pipeline():
+    preprocess_task = preprocess_op(uri='uri-to-import', some_int=12)
+    train_task = train_op(
+        num_steps=preprocess_task.outputs['output_parameter_one'],
+        dataset=preprocess_task.outputs['output_dataset_one']
+    )