chore(sdk): extract DSL into `kfp-dsl` package (#9738)

* move dsl dir

* make kfp-dsl a package

* make kfp-dsl a package

* additional changes

* address review feedback

sdk/python/build.sh

@@ -21,7 +21,7 @@
 #   ./build.sh [output_file]
 
 
-target_archive_file=${1:-kfp.tar.gz}
+target_archive_file=$1
 
 pushd "$(dirname "$0")"
 dist_dir=$(mktemp -d)

sdk/python/install_from_source.sh

@@ -0,0 +1,4 @@
+#!/bin/bash
+
+pip3 install -e sdk/python/kfp-dsl
+pip3 install -e sdk/python

sdk/python/kfp-dsl/README.md

@@ -0,0 +1,25 @@
+## kfp-dsl package
+
+`kfp-dsl` is a subpackage of the KFP SDK that is released separately in order to provide a minimal dependency runtime package for Lightweight Python Components. **`kfp-dsl` should not be installed and used directly.**
+
+`kfp-dsl` enables the KFP runtime code and objects to be installed at Lightweight Python Component runtime without needing to install the full KFP SDK package.
+
+### Release
+`kfp-dsl` should be released immediately prior to each full `kfp` release. The version of `kfp-dsl` should match the version of `kfp` that depends on it.
+
+### Development
+To develop on `kfp` with a version of `kfp-dsl` built from source, run the following from the repository root:
+
+```sh
+source sdk/python/install_from_source.sh
+```
+
+**Note:** Modules in the `kfp-dsl` package are only permitted to have *top-level* imports from the Python standard library, the `typing-extensions` package, and the `kfp-dsl` package itself. Imports from other subpackages of the main `kfp` package or its transitive dependencies must be nested within functions to avoid runtime import errors when only `kfp-dsl` is installed.
+
+### Testing
+The `kfp-dsl` code is tested alongside the full KFP SDK in `sdk/python/kfp/dsl-test`. This is because many of the DSL tests require the full KFP SDK to be installed (e.g., requires creating and compiling a component/pipeline).
+
+There are also dedicated `kfp-dsl` tests `./sdk/python/kfp-dsl/runtime_tests/` which test the dedicated runtime code in `kfp-dsl` and should *not* be run with the full KFP SDK installed. Specifically, these tests ensure:
+* That KFP runtime logic is correct
+* That `kfp-dsl` specifies all of its dependencies (i.e., no module not found errors from missing `kfp-dsl` dependencies)
+* That `kfp-dsl` dependencies on the main `kfp` package have associated imports nested inside function calls (i.e., no module not found errors from missing `kfp` dependencies)

sdk/python/kfp-dsl/build.sh

@@ -0,0 +1,30 @@
+#!/bin/bash -ex
+#
+# Copyright 2018 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.
+
+
+# The scripts creates the Kubeflow Pipelines python SDK package.
+#
+# Usage:
+#   ./build.sh [output_file]
+
+
+target_archive_file=$1
+
+pushd "$(dirname "$0")"
+dist_dir=$(mktemp -d)
+python3 setup.py sdist --format=gztar --dist-dir "$dist_dir"
+cp "$dist_dir"/*.tar.gz "$target_archive_file"
+popd

sdk/python/kfp-dsl/kfp/dsl/__init__.py

@@ -50,6 +50,8 @@ __all__ = [
     'PipelineTask',
 ]
 
+_kfp_dsl_import_error_msg = 'It looks like only `kfp-dsl` is installed. Please install the full KFP SDK using `pip install kfp`.'
+
 try:
     from typing import Annotated
 except ImportError:

sdk/python/kfp-dsl/kfp/dsl/base_component.py

@@ -19,7 +19,6 @@ from typing import List
 from kfp.dsl import pipeline_task
 from kfp.dsl import structures
 from kfp.dsl.types import type_utils
-from kfp.pipeline_spec import pipeline_spec_pb2
 
 
 class BaseComponent(abc.ABC):
@@ -103,13 +102,13 @@ class BaseComponent(abc.ABC):
         )
 
     @property
-    def pipeline_spec(self) -> pipeline_spec_pb2.PipelineSpec:
+    def pipeline_spec(self) -> 'pipeline_spec_pb2.PipelineSpec':
         """Returns the pipeline spec of the component."""
         with BlockPipelineTaskRegistration():
             return self.component_spec.to_pipeline_spec()
 
     @property
-    def platform_spec(self) -> pipeline_spec_pb2.PlatformSpec:
+    def platform_spec(self) -> 'pipeline_spec_pb2.PlatformSpec':
         """Returns the PlatformSpec of the component.
 
         Useful when the component is a GraphComponent, else will be

sdk/python/kfp-dsl/kfp/dsl/component_factory.py

@@ -20,7 +20,7 @@ import textwrap
 from typing import Callable, List, Mapping, Optional, Tuple, Type, Union
 import warnings
 
-import docstring_parser
+from kfp import dsl
 from kfp.dsl import container_component_artifact_channel
 from kfp.dsl import container_component_class
 from kfp.dsl import graph_component
@@ -124,9 +124,9 @@ def _get_packages_to_install_command(
     return ['sh', '-c', install_python_packages_script]
 
 
-def _get_default_kfp_package_path() -> str:
+def _get_kfp_dsl_requirement() -> str:
     import kfp
-    return f'kfp=={kfp.__version__}'
+    return f'kfp-dsl=={kfp.__version__}'
 
 
 def _get_function_source_definition(func: Callable) -> str:
@@ -175,6 +175,12 @@ def extract_component_interface(
     parameters = list(signature.parameters.values())
 
     original_docstring = inspect.getdoc(func)
+
+    try:
+        import docstring_parser
+    except ImportError as e:
+        raise ImportError(dsl._kfp_dsl_import_error_msg) from e
+
     parsed_docstring = docstring_parser.parse(original_docstring)
 
     inputs = {}
@@ -475,7 +481,7 @@ def create_component_from_func(
 
     if install_kfp_package and target_image is None:
         if kfp_package_path is None:
-            kfp_package_path = _get_default_kfp_package_path()
+            kfp_package_path = _get_kfp_dsl_requirement()
         packages_to_install.append(kfp_package_path)
 
     packages_to_install_command = _get_packages_to_install_command(
@@ -622,7 +628,7 @@ def create_graph_component_from_func(
 
 def get_pipeline_description(
     decorator_description: Union[str, None],
-    docstring: docstring_parser.Docstring,
+    docstring: 'docstring_parser.Docstring',
 ) -> Union[str, None]:
     """Obtains the correct pipeline description from the pipeline decorator's
     description argument and the parsed docstring.

sdk/python/kfp-dsl/kfp/dsl/graph_component.py

@@ -17,12 +17,11 @@ import inspect
 from typing import Callable, Optional
 import uuid
 
-from kfp.compiler import pipeline_spec_builder as builder
+from kfp import dsl
 from kfp.dsl import base_component
 from kfp.dsl import pipeline_channel
 from kfp.dsl import pipeline_context
 from kfp.dsl import structures
-from kfp.pipeline_spec import pipeline_spec_pb2
 
 
 class GraphComponent(base_component.BaseComponent):
@@ -65,6 +64,11 @@ class GraphComponent(base_component.BaseComponent):
         pipeline_group = dsl_pipeline.groups[0]
         pipeline_group.name = uuid.uuid4().hex
 
+        try:
+            from kfp.compiler import pipeline_spec_builder as builder
+        except ImportError as e:
+            raise ImportError(dsl._kfp_dsl_import_error_msg) from e
+
         pipeline_spec, platform_spec = builder.create_pipeline_spec(
             pipeline=dsl_pipeline,
             component_spec=self.component_spec,
@@ -83,7 +87,7 @@ class GraphComponent(base_component.BaseComponent):
         self.component_spec.platform_spec = platform_spec
 
     @property
-    def pipeline_spec(self) -> pipeline_spec_pb2.PipelineSpec:
+    def pipeline_spec(self) -> 'pipeline_spec_pb2.PipelineSpec':
         """Returns the pipeline spec of the component."""
         return self.component_spec.implementation.graph
 

sdk/python/kfp-dsl/kfp/dsl/pipeline_task.py

@@ -20,13 +20,13 @@ import re
 from typing import Any, Dict, List, Mapping, Optional, Union
 import warnings
 
+import kfp
 from kfp.dsl import constants
 from kfp.dsl import pipeline_channel
 from kfp.dsl import placeholders
 from kfp.dsl import structures
 from kfp.dsl import utils
 from kfp.dsl.types import type_utils
-from kfp.pipeline_spec import pipeline_spec_pb2
 
 _register_task_handler = lambda task: utils.maybe_rename_for_k8s(
     task.component_spec.name)
@@ -89,6 +89,7 @@ class PipelineTask:
                 error_message_prefix=(
                     f'Incompatible argument passed to the input '
                     f'{input_name!r} of component {component_spec.name!r}: '),
+                raise_on_error=kfp.TYPE_CHECK,
             )
 
         self.component_spec = component_spec
@@ -149,7 +150,7 @@ class PipelineTask:
         ])
 
     @property
-    def platform_spec(self) -> pipeline_spec_pb2.PlatformSpec:
+    def platform_spec(self) -> 'pipeline_spec_pb2.PlatformSpec':
         """PlatformSpec for all tasks in the pipeline as task.
 
         Only for use on tasks created from GraphComponents.

sdk/python/kfp-dsl/kfp/dsl/structures.py

@@ -18,22 +18,18 @@ import collections
 import dataclasses
 import itertools
 import re
-from typing import Any, Dict, List, Mapping, Optional, Tuple, Union
+from typing import Any, Dict, List, Mapping, Optional, Union
 import uuid
 
-from google.protobuf import json_format
-import kfp
+from kfp import dsl
 from kfp.dsl import placeholders
 from kfp.dsl import utils
-from kfp.dsl import v1_components
 from kfp.dsl import v1_structures
 from kfp.dsl.container_component_artifact_channel import \
     ContainerComponentArtifactChannel
 from kfp.dsl.types import artifact_types
 from kfp.dsl.types import type_annotations
 from kfp.dsl.types import type_utils
-from kfp.pipeline_spec import pipeline_spec_pb2
-import yaml
 
 
 @dataclasses.dataclass
@@ -370,7 +366,7 @@ class RetryPolicy:
     backoff_factor: Optional[float] = None
     backoff_max_duration: Optional[str] = None
 
-    def to_proto(self) -> pipeline_spec_pb2.PipelineTaskSpec.RetryPolicy:
+    def to_proto(self) -> 'pipeline_spec_pb2.PipelineTaskSpec.RetryPolicy':
         # include defaults so that IR is more reflective of runtime behavior
         max_retry_count = self.max_retry_count or 0
         backoff_duration = self.backoff_duration or '0s'
@@ -381,6 +377,12 @@ class RetryPolicy:
         backoff_duration_seconds = f'{convert_duration_to_seconds(backoff_duration)}s'
         backoff_max_duration_seconds = f'{min(convert_duration_to_seconds(backoff_max_duration), 3600)}s'
 
+        try:
+            from google.protobuf import json_format
+            from kfp.pipeline_spec import pipeline_spec_pb2
+        except ImportError as e:
+            raise ImportError(dsl._kfp_dsl_import_error_msg) from e
+
         return json_format.ParseDict(
             {
                 'max_retry_count': max_retry_count,
@@ -480,6 +482,13 @@ class Implementation:
                 executor['container']) if executor else None
             return Implementation(container=container_spec)
         else:
+
+            try:
+                from google.protobuf import json_format
+                from kfp.pipeline_spec import pipeline_spec_pb2
+            except ImportError as e:
+                raise ImportError(dsl._kfp_dsl_import_error_msg) from e
+
             pipeline_spec = json_format.ParseDict(
                 pipeline_spec_dict, pipeline_spec_pb2.PipelineSpec())
             return Implementation(graph=pipeline_spec)
@@ -551,6 +560,14 @@ def check_placeholder_references_valid_io_name(
         raise TypeError(f'Unexpected argument "{arg}" of type {type(arg)}.')
 
 
+def _import_and_make_platform_spec() -> 'pipeline_spec_pb2.PlatformSpec':
+    try:
+        from kfp.pipeline_spec import pipeline_spec_pb2
+    except ImportError as e:
+        raise ImportError(dsl._kfp_dsl_import_error_msg) from e
+    return pipeline_spec_pb2.PlatformSpec()
+
+
 @dataclasses.dataclass
 class ComponentSpec:
     """The definition of a component.
@@ -568,8 +585,9 @@ class ComponentSpec:
     description: Optional[str] = None
     inputs: Optional[Dict[str, InputSpec]] = None
     outputs: Optional[Dict[str, OutputSpec]] = None
-    platform_spec: pipeline_spec_pb2.PlatformSpec = dataclasses.field(
-        default_factory=pipeline_spec_pb2.PlatformSpec)
+    platform_spec: Optional[
+        'pipeline_spec_pb2.PlatformSpec'] = dataclasses.field(
+            default_factory=_import_and_make_platform_spec)
 
     def __post_init__(self) -> None:
         self._transform_name()
@@ -652,6 +670,13 @@ class ComponentSpec:
         })
 
         inputs = {}
+
+        try:
+            from kfp.pipeline_spec import pipeline_spec_pb2
+        except ImportError as e:
+            raise ImportError(dsl._kfp_dsl_import_error_msg) from e
+        parameter_types_mapping = type_utils.get_parameter_types_mapping()
+
         for spec in component_dict.get('inputs', []):
             type_ = spec.get('type')
             optional = spec.get('optional', False) or 'default' in spec
@@ -667,9 +692,9 @@ class ComponentSpec:
                     type=type_, optional=True)
                 continue
 
-            elif isinstance(type_, str) and type_.lower(
-            ) in type_utils._PARAMETER_TYPES_MAPPING:
-                type_enum = type_utils._PARAMETER_TYPES_MAPPING[type_.lower()]
+            elif isinstance(type_,
+                            str) and type_.lower() in parameter_types_mapping:
+                type_enum = parameter_types_mapping[type_.lower()]
                 ir_parameter_type_name = pipeline_spec_pb2.ParameterType.ParameterTypeEnum.Name(
                     type_enum)
                 in_memory_parameter_type_name = type_utils.IR_TYPE_TO_IN_MEMORY_SPEC_TYPE[
@@ -720,9 +745,9 @@ class ComponentSpec:
             if isinstance(type_, str):
                 type_ = type_utils.get_canonical_name_for_outer_generic(type_)
 
-            if isinstance(type_, str) and type_.lower(
-            ) in type_utils._PARAMETER_TYPES_MAPPING:
-                type_enum = type_utils._PARAMETER_TYPES_MAPPING[type_.lower()]
+            if isinstance(type_,
+                          str) and type_.lower() in parameter_types_mapping:
+                type_enum = parameter_types_mapping[type_.lower()]
                 ir_parameter_type_name = pipeline_spec_pb2.ParameterType.ParameterTypeEnum.Name(
                     type_enum)
                 in_memory_parameter_type_name = type_utils.IR_TYPE_TO_IN_MEMORY_SPEC_TYPE[
@@ -824,6 +849,12 @@ class ComponentSpec:
             implementation.container.command or
             []) if implementation.container else None
 
+        try:
+            from google.protobuf import json_format
+            from kfp.pipeline_spec import pipeline_spec_pb2
+        except ImportError as e:
+            raise ImportError(dsl._kfp_dsl_import_error_msg) from e
+
         platform_spec = pipeline_spec_pb2.PlatformSpec()
         json_format.ParseDict(platform_spec_dict, platform_spec)
 
@@ -836,53 +867,6 @@ class ComponentSpec:
             platform_spec=platform_spec,
         )
 
-    @classmethod
-    def from_yaml_documents(cls, component_yaml: str) -> 'ComponentSpec':
-        """Loads V1 or V2 component YAML into a ComponentSpec.
-
-        Args:
-            component_yaml: PipelineSpec and optionally PlatformSpec YAML documents as a single string.
-
-        Returns:
-            ComponentSpec: The ComponentSpec object.
-        """
-
-        def extract_description(component_yaml: str) -> Union[str, None]:
-            heading = '# Description: '
-            multi_line_description_prefix = '#             '
-            index_of_heading = 2
-            if heading in component_yaml:
-                description = component_yaml.splitlines()[index_of_heading]
-
-                # Multi line
-                comments = component_yaml.splitlines()
-                index = index_of_heading + 1
-                while comments[index][:len(multi_line_description_prefix
-                                          )] == multi_line_description_prefix:
-                    description += '\n' + comments[index][
-                        len(multi_line_description_prefix) + 1:]
-                    index += 1
-
-                return description[len(heading):]
-            else:
-                return None
-
-        pipeline_spec_dict, platform_spec_dict = load_documents_from_yaml(
-            component_yaml)
-
-        is_v1 = 'implementation' in set(pipeline_spec_dict.keys())
-        if is_v1:
-            v1_component = v1_components._load_component_spec_from_component_text(
-                component_yaml)
-            return cls.from_v1_component_spec(v1_component)
-        else:
-            component_spec = ComponentSpec.from_ir_dicts(
-                pipeline_spec_dict, platform_spec_dict)
-            if not component_spec.description:
-                component_spec.description = extract_description(
-                    component_yaml=component_yaml)
-            return component_spec
-
     def save_to_component_yaml(self, output_file: str) -> None:
         """Saves ComponentSpec into IR YAML file.
 
@@ -892,6 +876,12 @@ class ComponentSpec:
         from kfp.compiler import pipeline_spec_builder as builder
 
         pipeline_spec = self.to_pipeline_spec()
+
+        try:
+            from kfp.pipeline_spec import pipeline_spec_pb2
+        except ImportError as e:
+            raise ImportError(dsl._kfp_dsl_import_error_msg) from e
+
         builder.write_pipeline_spec_to_file(
             pipeline_spec,
             None,
@@ -899,7 +889,7 @@ class ComponentSpec:
             output_file,
         )
 
-    def to_pipeline_spec(self) -> pipeline_spec_pb2.PipelineSpec:
+    def to_pipeline_spec(self) -> 'pipeline_spec_pb2.PipelineSpec':
         """Creates a pipeline instance and constructs the pipeline spec for a
         single component.
 
@@ -950,6 +940,12 @@ class ComponentSpec:
 
         utils.validate_pipeline_name(pipeline_name)
 
+        try:
+            import kfp
+            from kfp.pipeline_spec import pipeline_spec_pb2
+        except ImportError as e:
+            raise ImportError(dsl._kfp_dsl_import_error_msg) from e
+
         pipeline_spec = pipeline_spec_pb2.PipelineSpec()
         pipeline_spec.pipeline_info.name = pipeline_name
         pipeline_spec.sdk_version = f'kfp-{kfp.__version__}'
@@ -1052,24 +1048,3 @@ def convert_duration_to_seconds(duration: str) -> int:
         raise ValueError(
             f"Unsupported duration unit: '{duration[-1]}' for '{duration}'.")
     return int(duration[:-1]) * seconds_per_unit[duration[-1]]
-
-
-def load_documents_from_yaml(component_yaml: str) -> Tuple[dict, dict]:
-    """Loads up to two YAML documents from a YAML string.
-
-    First document must always be present. If second document is
-    present, it is returned as a dict, else an empty dict.
-    """
-    documents = list(yaml.safe_load_all(component_yaml))
-    num_docs = len(documents)
-    if num_docs == 1:
-        pipeline_spec_dict = documents[0]
-        platform_spec_dict = {}
-    elif num_docs == 2:
-        pipeline_spec_dict = documents[0]
-        platform_spec_dict = documents[1]
-    else:
-        raise ValueError(
-            f'Expected one or two YAML documents in the IR YAML file. Got: {num_docs}.'
-        )
-    return pipeline_spec_dict, platform_spec_dict

sdk/python/kfp-dsl/kfp/dsl/types/type_utils.py

@@ -19,12 +19,11 @@ import json
 from typing import Any, Callable, Dict, Optional, Type, Union
 import warnings
 
-import kfp
+from kfp import dsl
 from kfp.dsl import structures
 from kfp.dsl import task_final_status
 from kfp.dsl.types import artifact_types
 from kfp.dsl.types import type_annotations
-from kfp.pipeline_spec import pipeline_spec_pb2
 
 DEFAULT_ARTIFACT_SCHEMA_VERSION = '0.0.1'
 PARAMETER_TYPES = Union[str, int, float, bool, dict, list]
@@ -44,24 +43,32 @@ _ARTIFACT_CLASSES_MAPPING = {
 _GOOGLE_TYPES_PATTERN = r'^google.[A-Za-z]+$'
 _GOOGLE_TYPES_VERSION = DEFAULT_ARTIFACT_SCHEMA_VERSION
 
+
 # ComponentSpec I/O types to (IR) PipelineTaskSpec I/O types mapping.
 # The keys are normalized (lowercased). These are types viewed as Parameters.
 # The values are the corresponding IR parameter primitive types.
-_PARAMETER_TYPES_MAPPING = {
-    'integer': pipeline_spec_pb2.ParameterType.NUMBER_INTEGER,
-    'int': pipeline_spec_pb2.ParameterType.NUMBER_INTEGER,
-    'double': pipeline_spec_pb2.ParameterType.NUMBER_DOUBLE,
-    'float': pipeline_spec_pb2.ParameterType.NUMBER_DOUBLE,
-    'string': pipeline_spec_pb2.ParameterType.STRING,
-    'str': pipeline_spec_pb2.ParameterType.STRING,
-    'text': pipeline_spec_pb2.ParameterType.STRING,
-    'bool': pipeline_spec_pb2.ParameterType.BOOLEAN,
-    'boolean': pipeline_spec_pb2.ParameterType.BOOLEAN,
-    'dict': pipeline_spec_pb2.ParameterType.STRUCT,
-    'list': pipeline_spec_pb2.ParameterType.LIST,
-    'jsonobject': pipeline_spec_pb2.ParameterType.STRUCT,
-    'jsonarray': pipeline_spec_pb2.ParameterType.LIST,
-}
+def get_parameter_types_mapping():
+
+    try:
+        from kfp.pipeline_spec import pipeline_spec_pb2
+    except ImportError as e:
+        raise ImportError(dsl._kfp_dsl_import_error_msg) from e
+
+    return {
+        'integer': pipeline_spec_pb2.ParameterType.NUMBER_INTEGER,
+        'int': pipeline_spec_pb2.ParameterType.NUMBER_INTEGER,
+        'double': pipeline_spec_pb2.ParameterType.NUMBER_DOUBLE,
+        'float': pipeline_spec_pb2.ParameterType.NUMBER_DOUBLE,
+        'string': pipeline_spec_pb2.ParameterType.STRING,
+        'str': pipeline_spec_pb2.ParameterType.STRING,
+        'text': pipeline_spec_pb2.ParameterType.STRING,
+        'bool': pipeline_spec_pb2.ParameterType.BOOLEAN,
+        'boolean': pipeline_spec_pb2.ParameterType.BOOLEAN,
+        'dict': pipeline_spec_pb2.ParameterType.STRUCT,
+        'list': pipeline_spec_pb2.ParameterType.LIST,
+        'jsonobject': pipeline_spec_pb2.ParameterType.STRUCT,
+        'jsonarray': pipeline_spec_pb2.ParameterType.LIST,
+    }
 
 
 def bool_cast_fn(default: Union[str, bool]) -> bool:
@@ -109,18 +116,6 @@ def deserialize_v1_component_yaml_default(type_: str, default: Any) -> Any:
     return default
 
 
-# Mapping primitive types to their IR message field names.
-# This is used in constructing condition strings.
-_PARAMETER_TYPES_VALUE_REFERENCE_MAPPING = {
-    pipeline_spec_pb2.ParameterType.NUMBER_INTEGER: 'number_value',
-    pipeline_spec_pb2.ParameterType.NUMBER_DOUBLE: 'number_value',
-    pipeline_spec_pb2.ParameterType.STRING: 'string_value',
-    pipeline_spec_pb2.ParameterType.BOOLEAN: 'bool_value',
-    pipeline_spec_pb2.ParameterType.STRUCT: 'struct_value',
-    pipeline_spec_pb2.ParameterType.LIST: 'list_value',
-}
-
-
 def is_task_final_status_type(type_name: Optional[Union[str, dict]]) -> bool:
     """Check if a ComponentSpec I/O type is PipelineTaskFinalStatus.
 
@@ -150,15 +145,21 @@ def is_parameter_type(type_name: Optional[Union[str, dict]]) -> bool:
     else:
         return False
 
-    return type_name.lower(
-    ) in _PARAMETER_TYPES_MAPPING or is_task_final_status_type(type_name)
+    return type_name.lower() in get_parameter_types_mapping(
+    ) or is_task_final_status_type(type_name)
 
 
 def bundled_artifact_to_artifact_proto(
-        bundled_artifact_str: str) -> pipeline_spec_pb2.ArtifactTypeSchema:
+        bundled_artifact_str: str) -> 'pipeline_spec_pb2.ArtifactTypeSchema':
     """Gets the IR ArtifactTypeSchema proto for a bundled artifact in form
     `<namespace>.<Name>@x.x.x` (e.g., system.Artifact@0.0.1)."""
     bundled_artifact_str, schema_version = bundled_artifact_str.split('@')
+
+    try:
+        from kfp.pipeline_spec import pipeline_spec_pb2
+    except ImportError as e:
+        raise ImportError(dsl._kfp_dsl_import_error_msg) from e
+
     return pipeline_spec_pb2.ArtifactTypeSchema(
         schema_title=bundled_artifact_str,
         schema_version=schema_version,
@@ -167,7 +168,7 @@ def bundled_artifact_to_artifact_proto(
 
 def get_parameter_type(
     param_type: Optional[Union[Type, str, dict]]
-) -> pipeline_spec_pb2.ParameterType:
+) -> 'pipeline_spec_pb2.ParameterType':
     """Get the IR I/O parameter type for the given ComponentSpec I/O type.
 
     Args:
@@ -189,12 +190,17 @@ def get_parameter_type(
         type_name = list(param_type.keys())[0]
     else:
         type_name = type_annotations.get_short_type_name(str(param_type))
-    return _PARAMETER_TYPES_MAPPING.get(type_name.lower())
+    return get_parameter_types_mapping().get(type_name.lower())
 
 
 def get_parameter_type_name(
         param_type: Optional[Union[Type, str, dict]]) -> str:
     """Gets the parameter type name."""
+    try:
+        from kfp.pipeline_spec import pipeline_spec_pb2
+    except ImportError as e:
+        raise ImportError(dsl._kfp_dsl_import_error_msg) from e
+
     return pipeline_spec_pb2.ParameterType.ParameterTypeEnum.Name(
         get_parameter_type(param_type))
 
@@ -213,6 +219,21 @@ def get_parameter_type_field_name(type_name: Optional[str]) -> Optional[str]:
     Raises:
       AttributeError: if type_name is not a string type.
     """
+    try:
+        from kfp.pipeline_spec import pipeline_spec_pb2
+    except ImportError as e:
+        raise ImportError(dsl._kfp_dsl_import_error_msg) from e
+
+    # Mapping primitive types to their IR message field names.
+    # This is used in constructing condition strings.
+    _PARAMETER_TYPES_VALUE_REFERENCE_MAPPING = {
+        pipeline_spec_pb2.ParameterType.NUMBER_INTEGER: 'number_value',
+        pipeline_spec_pb2.ParameterType.NUMBER_DOUBLE: 'number_value',
+        pipeline_spec_pb2.ParameterType.STRING: 'string_value',
+        pipeline_spec_pb2.ParameterType.BOOLEAN: 'bool_value',
+        pipeline_spec_pb2.ParameterType.STRUCT: 'struct_value',
+        pipeline_spec_pb2.ParameterType.LIST: 'list_value',
+    }
     return _PARAMETER_TYPES_VALUE_REFERENCE_MAPPING.get(
         get_parameter_type(type_name))
 
@@ -251,6 +272,7 @@ def verify_type_compatibility(
     expected_spec: Union[structures.InputSpec, structures.OutputSpec],
     error_message_prefix: str,
     checks_input: bool = True,
+    raise_on_error: bool = True,
 ) -> bool:
     """Verifies the given argument type is compatible with the expected type.
 
@@ -259,12 +281,13 @@ def verify_type_compatibility(
         expected_spec: The InputSpec or OutputSpec that describes the expected type of given_value.
         error_message_prefix: The prefix for the error message.
         checks_input: True if checks an argument (given_value) against a component/pipeline input type (expected_spec). False if checks a component output (argument_value) against the pipeline output type (expected_spec).
+        raise_on_error: Whether to raise on type compatibility error. Should be passed kfp.TYPE_CHECK.
 
     Returns:
         True if types are compatible, and False if otherwise.
 
     Raises:
-        InconsistentTypeException if types are incompatible and TYPE_CHECK==True.
+        InconsistentTypeException if raise_on_error=True.
     """
     # extract and normalize types
     expected_type = expected_spec.type
@@ -307,7 +330,7 @@ def verify_type_compatibility(
         else:
             error_message_suffix = f'Output of type {given_type!r} cannot be surfaced as pipeline output type {expected_type!r}'
         error_text = error_message_prefix + error_message_suffix
-        if kfp.TYPE_CHECK:
+        if raise_on_error:
             raise InconsistentTypeException(error_text)
         else:
             warnings.warn(InconsistentTypeWarning(error_text))
@@ -424,11 +447,21 @@ class TypeCheckManager:
         Returns:
             TypeCheckManager: Returns itself.
         """
+        try:
+            import kfp
+        except ImportError as e:
+            raise ImportError(dsl._kfp_dsl_import_error_msg) from e
+
         self._prev = kfp.TYPE_CHECK
         kfp.TYPE_CHECK = self._enable
         return self
 
     def __exit__(self, *unused_args) -> None:
+
+        try:
+            import kfp
+        except ImportError as e:
+            raise ImportError(dsl._kfp_dsl_import_error_msg) from e
         """Restore type check mode to its previous state."""
         kfp.TYPE_CHECK = self._prev
 

sdk/python/kfp-dsl/kfp/dsl/v1_structures.py

@@ -16,7 +16,6 @@ from collections import OrderedDict
 from typing import Any, Dict, List, Mapping, Optional, Union
 
 from kfp.dsl.v1_modelbase import ModelBase
-import yaml
 
 PrimitiveTypes = Union[str, int, float, bool]
 PrimitiveTypesIncludingNone = Optional[PrimitiveTypes]
@@ -437,17 +436,6 @@ class ComponentSpec(ModelBase):
                                     f'Argument "{argument}" references non-existing input.'
                                 )
 
-    def save(self, file_path: str):
-        """Saves the component definition to file.
-
-        It can be shared online and later loaded using the
-        load_component function.
-        """
-
-        component_yaml = yaml.dump(self.to_dict(), sort_keys=True)
-        with open(file_path, 'w') as f:
-            f.write(component_yaml)
-
 
 class ComponentReference(ModelBase):
     """Component reference.

sdk/python/kfp-dsl/kfp/dsl/yaml_component.py

@@ -13,10 +13,9 @@
 # limitations under the License.
 """Component loaded from YAML."""
 
-from google.protobuf import json_format
+from kfp import dsl
 from kfp.dsl import base_component
 from kfp.dsl import structures
-from kfp.pipeline_spec import pipeline_spec_pb2
 
 
 class YamlComponent(base_component.BaseComponent):
@@ -38,14 +37,21 @@ class YamlComponent(base_component.BaseComponent):
         self.component_yaml = component_yaml
 
     @property
-    def pipeline_spec(self) -> pipeline_spec_pb2.PipelineSpec:
+    def pipeline_spec(self) -> 'pipeline_spec_pb2.PipelineSpec':
         """Returns the pipeline spec of the component."""
-        component_dict = structures.load_documents_from_yaml(
+        try:
+            from google.protobuf import json_format
+            from kfp.components import load_yaml_utilities
+            from kfp.pipeline_spec import pipeline_spec_pb2
+        except ImportError as e:
+            raise ImportError(dsl._kfp_dsl_import_error_msg) from e
+        component_dict = load_yaml_utilities._load_documents_from_yaml(
             self.component_yaml)[0]
         is_v1 = 'implementation' in set(component_dict.keys())
         if is_v1:
             return self.component_spec.to_pipeline_spec()
         else:
+
             return json_format.ParseDict(component_dict,
                                          pipeline_spec_pb2.PipelineSpec())
 

sdk/python/kfp-dsl/runtime_tests/execute_commands_args_test.py

@@ -0,0 +1,139 @@
+# Copyright 2023 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.
+import dataclasses
+import json
+import os
+import shutil
+import subprocess
+import tempfile
+from typing import Any, Dict
+
+from absl.testing import parameterized
+import yaml
+
+TEST_DATA_DIR = os.path.join(os.path.dirname(__file__), 'test_data')
+TMP_DIR = tempfile.mkdtemp()
+
+
+@dataclasses.dataclass
+class RuntimeTestConfig:
+    pipeline_file_relpath: str
+    executor_name: str
+    executor_input: Dict[str, Any]
+
+
+TEST_CONFIGS = [
+    RuntimeTestConfig(
+        pipeline_file_relpath=os.path.join(
+            TEST_DATA_DIR, 'pipeline_with_task_final_status.yaml'),
+        executor_name='exec-print-op',
+        executor_input={
+            'inputs': {
+                'parameterValues': {
+                    'message': 'Hello World!'
+                },
+                'parameters': {
+                    'message': {
+                        'stringValue': 'Hello World!'
+                    }
+                }
+            },
+            'outputs': {
+                'outputFile':
+                    '/gcs/cjmccarthy-kfp-default-bucket/271009669852/pipeline-with-task-final-status-07-14-2023-18-50-32/print-op_-9063136771365142528/executor_output.json'
+            }
+        },
+    ),
+    RuntimeTestConfig(
+        pipeline_file_relpath=os.path.join(
+            TEST_DATA_DIR, 'pipeline_with_task_final_status.yaml'),
+        executor_name='exec-exit-op',
+        executor_input={
+            'inputs': {
+                'parameterValues': {
+                    'status': {
+                        'error': {
+                            'code':
+                                9,
+                            'message':
+                                'The DAG failed because some tasks failed. The failed tasks are: [print-op, fail-op].'
+                        },
+                        'pipelineJobResourceName':
+                            'projects/271009669852/locations/us-central1/pipelineJobs/pipeline-with-task-final-status-07-14-2023-19-07-11',
+                        'pipelineTaskName':
+                            'my-pipeline',
+                        'state':
+                            'FAILED'
+                    },
+                    'user_input': 'Hello World!'
+                },
+                'parameters': {
+                    'status': {
+                        'stringValue':
+                            "{\"error\":{\"code\":9,\"message\":\"The DAG failed because some tasks failed. The failed tasks are: [print-op, fail-op].\"},\"pipelineJobResourceName\":\"projects/271009669852/locations/us-central1/pipelineJobs/pipeline-with-task-final-status-07-14-2023-19-07-11\",\"pipelineTaskName\":\"my-pipeline\",\"state\":\"FAILED\"}"
+                    },
+                    'user_input': {
+                        'stringValue': 'Hello World!'
+                    }
+                }
+            },
+            'outputs': {
+                'outputFile':
+                    '/gcs/cjmccarthy-kfp-default-bucket/271009669852/pipeline-with-task-final-status-07-14-2023-19-07-11/exit-op_-6100894116462198784/executor_output.json'
+            }
+        },
+    )
+]
+
+
+def run_commands_and_args(
+    config: RuntimeTestConfig,
+    temp_dir: str,
+) -> subprocess.CompletedProcess:
+    with open(config.pipeline_file_relpath) as f:
+        pipline_spec_dict = yaml.safe_load(f)
+        container = pipline_spec_dict['deploymentSpec']['executors'][
+            config.executor_name]['container']
+
+        command_and_args = container['command'] + container['args']
+        executor_input_json = json.dumps(config.executor_input).replace(
+            '/gcs/', TMP_DIR)
+        command_and_args = [
+            v.replace('{{$}}', executor_input_json) for v in command_and_args
+        ]
+
+        return subprocess.run(
+            command_and_args,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True)
+
+
+class TestRuntime(parameterized.TestCase):
+
+    @classmethod
+    def setUp(cls):
+        cls.temp_dir = tempfile.mkdtemp()
+
+    @classmethod
+    def tearDown(cls):
+        shutil.rmtree(cls.temp_dir)
+
+    @parameterized.parameters(TEST_CONFIGS)
+    def test(self, config: RuntimeTestConfig):
+        process = run_commands_and_args(
+            config=config,
+            temp_dir=self.temp_dir,
+        )
+        self.assertEqual(process.returncode, 0, process.stderr)

sdk/python/kfp-dsl/runtime_tests/import_objects_test.py

@@ -0,0 +1,21 @@
+# Copyright 2023 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.
+class TestImportObjects:
+
+    def test(self):
+        # from kfp.dsl import * only allowed at module level, so emulate behavior
+        from kfp import dsl
+        for obj_name in dir(dsl):
+            if not obj_name.startswith('_'):
+                getattr(dsl, obj_name)

sdk/python/kfp-dsl/runtime_tests/test_data/pipeline_with_task_final_status.py

@@ -0,0 +1,58 @@
+# Copyright 2022 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.
+"""Pipeline using ExitHandler with PipelineTaskFinalStatus."""
+
+from kfp import compiler
+from kfp import dsl
+from kfp.dsl import component
+from kfp.dsl import PipelineTaskFinalStatus
+
+
+@component
+def exit_op(user_input: str, status: PipelineTaskFinalStatus):
+    """Checks pipeline run status."""
+    print('Pipeline status: ', status.state)
+    print('Job resource name: ', status.pipeline_job_resource_name)
+    print('Pipeline task name: ', status.pipeline_task_name)
+    print('Error code: ', status.error_code)
+    print('Error message: ', status.error_message)
+
+
+@component
+def print_op(message: str):
+    """Prints a message."""
+    print(message)
+
+
+@component
+def fail_op(message: str):
+    """Fails."""
+    import sys
+    print(message)
+    sys.exit(1)
+
+
+@dsl.pipeline(name='pipeline-with-task-final-status')
+def my_pipeline(message: str = 'Hello World!'):
+    exit_task = exit_op(user_input=message)
+
+    with dsl.ExitHandler(exit_task, name='my-pipeline'):
+        print_op(message=message)
+        fail_op(message='Task failed.')
+
+
+if __name__ == '__main__':
+    compiler.Compiler().compile(
+        pipeline_func=my_pipeline,
+        package_path=__file__.replace('.py', '.yaml'))

sdk/python/kfp-dsl/runtime_tests/test_data/pipeline_with_task_final_status.yaml

@@ -0,0 +1,183 @@
+# PIPELINE DEFINITION
+# Name: pipeline-with-task-final-status
+# Inputs:
+#    message: str [Default: 'Hello World!']
+components:
+  comp-exit-handler-1:
+    dag:
+      tasks:
+        fail-op:
+          cachingOptions:
+            enableCache: true
+          componentRef:
+            name: comp-fail-op
+          inputs:
+            parameters:
+              message:
+                runtimeValue:
+                  constant: Task failed.
+          taskInfo:
+            name: fail-op
+        print-op:
+          cachingOptions:
+            enableCache: true
+          componentRef:
+            name: comp-print-op
+          inputs:
+            parameters:
+              message:
+                componentInputParameter: pipelinechannel--message
+          taskInfo:
+            name: print-op
+    inputDefinitions:
+      parameters:
+        pipelinechannel--message:
+          parameterType: STRING
+  comp-exit-op:
+    executorLabel: exec-exit-op
+    inputDefinitions:
+      parameters:
+        status:
+          isOptional: true
+          parameterType: TASK_FINAL_STATUS
+        user_input:
+          parameterType: STRING
+  comp-fail-op:
+    executorLabel: exec-fail-op
+    inputDefinitions:
+      parameters:
+        message:
+          parameterType: STRING
+  comp-print-op:
+    executorLabel: exec-print-op
+    inputDefinitions:
+      parameters:
+        message:
+          parameterType: STRING
+deploymentSpec:
+  executors:
+    exec-exit-op:
+      container:
+        args:
+        - --executor_input
+        - '{{$}}'
+        - --function_to_execute
+        - exit_op
+        command:
+        - sh
+        - -c
+        - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
+          \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
+          \ && \"$0\" \"$@\"\n"
+        - sh
+        - -ec
+        - 'program_path=$(mktemp -d)
+
+          printf "%s" "$0" > "$program_path/ephemeral_component.py"
+
+          python3 -m kfp.dsl.executor_main                         --component_module_path                         "$program_path/ephemeral_component.py"                         "$@"
+
+          '
+        - "\nimport kfp\nfrom kfp import dsl\nfrom kfp.dsl import *\nfrom typing import\
+          \ *\n\ndef exit_op(user_input: str, status: PipelineTaskFinalStatus):\n\
+          \    \"\"\"Checks pipeline run status.\"\"\"\n    print('Pipeline status:\
+          \ ', status.state)\n    print('Job resource name: ', status.pipeline_job_resource_name)\n\
+          \    print('Pipeline task name: ', status.pipeline_task_name)\n    print('Error\
+          \ code: ', status.error_code)\n    print('Error message: ', status.error_message)\n\
+          \n"
+        image: python:3.7
+    exec-fail-op:
+      container:
+        args:
+        - --executor_input
+        - '{{$}}'
+        - --function_to_execute
+        - fail_op
+        command:
+        - sh
+        - -c
+        - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
+          \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
+          \ && \"$0\" \"$@\"\n"
+        - sh
+        - -ec
+        - 'program_path=$(mktemp -d)
+
+          printf "%s" "$0" > "$program_path/ephemeral_component.py"
+
+          python3 -m kfp.dsl.executor_main                         --component_module_path                         "$program_path/ephemeral_component.py"                         "$@"
+
+          '
+        - "\nimport kfp\nfrom kfp import dsl\nfrom kfp.dsl import *\nfrom typing import\
+          \ *\n\ndef fail_op(message: str):\n    \"\"\"Fails.\"\"\"\n    import sys\n\
+          \    print(message)\n    sys.exit(1)\n\n"
+        image: python:3.7
+    exec-print-op:
+      container:
+        args:
+        - --executor_input
+        - '{{$}}'
+        - --function_to_execute
+        - print_op
+        command:
+        - sh
+        - -c
+        - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
+          \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
+          \ && \"$0\" \"$@\"\n"
+        - sh
+        - -ec
+        - 'program_path=$(mktemp -d)
+
+          printf "%s" "$0" > "$program_path/ephemeral_component.py"
+
+          python3 -m kfp.dsl.executor_main                         --component_module_path                         "$program_path/ephemeral_component.py"                         "$@"
+
+          '
+        - "\nimport kfp\nfrom kfp import dsl\nfrom kfp.dsl import *\nfrom typing import\
+          \ *\n\ndef print_op(message: str):\n    \"\"\"Prints a message.\"\"\"\n\
+          \    print(message)\n\n"
+        image: python:3.7
+pipelineInfo:
+  name: pipeline-with-task-final-status
+root:
+  dag:
+    tasks:
+      exit-handler-1:
+        componentRef:
+          name: comp-exit-handler-1
+        inputs:
+          parameters:
+            pipelinechannel--message:
+              componentInputParameter: message
+        taskInfo:
+          name: my-pipeline
+      exit-op:
+        cachingOptions:
+          enableCache: true
+        componentRef:
+          name: comp-exit-op
+        dependentTasks:
+        - exit-handler-1
+        inputs:
+          parameters:
+            status:
+              taskFinalStatus:
+                producerTask: exit-handler-1
+            user_input:
+              componentInputParameter: message
+        taskInfo:
+          name: exit-op
+        triggerPolicy:
+          strategy: ALL_UPSTREAM_TASKS_COMPLETED
+  inputDefinitions:
+    parameters:
+      message:
+        defaultValue: Hello World!
+        isOptional: true
+        parameterType: STRING
+schemaVersion: 2.1.0
+sdkVersion: kfp-2.0.1

sdk/python/kfp-dsl/setup.py

@@ -0,0 +1,52 @@
+# Copyright 2023 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.
+
+import os
+import re
+
+import setuptools
+
+
+def find_version(*file_path_parts: str) -> str:
+    """Get version from a file that defines a __version__ variable."""
+
+    file_path = os.path.join(os.path.dirname(__file__), *file_path_parts)
+    with open(file_path, 'r') as f:
+        version_file_text = f.read()
+
+    version_match = re.search(
+        r"^__version__ = ['\"]([^'\"]*)['\"]",
+        version_file_text,
+        re.M,
+    )
+    if version_match:
+        return version_match.group(1)
+
+    raise RuntimeError(f'Unable to find version string in file: {file_path}.')
+
+
+setuptools.setup(
+    name='kfp-dsl',
+    version=find_version(
+        os.path.dirname(os.path.dirname(__file__)), 'kfp', '__init__.py'),
+    description='A KFP SDK subpackage containing the DSL and runtime code.',
+    author='google',
+    author_email='kubeflow-pipelines@google.com',
+    url='https://github.com/kubeflow/pipelines',
+    packages=setuptools.find_namespace_packages(include=['kfp.*']),
+    python_requires='>=3.7.0',
+    install_requires=['typing-extensions>=3.7.4,<5; python_version<"3.9"'],
+    include_package_data=True,
+    license='Apache 2.0',
+)

sdk/python/kfp/compiler/pipeline_spec_builder.py

@@ -1748,6 +1748,7 @@ def _validate_dag_output_types(
             output_spec,
             error_message_prefix,
             checks_input=False,
+            raise_on_error=kfp.TYPE_CHECK,
         )
 
 

sdk/python/kfp/components/load_yaml_utilities.py

@@ -13,11 +13,15 @@
 # limitations under the License.
 """Functions for loading components from compiled YAML."""
 
-from typing import Optional, Tuple
+import hashlib
+from typing import Optional, Tuple, Union
+import warnings
 
 from kfp.dsl import structures
+from kfp.dsl import v1_structures
 from kfp.dsl import yaml_component
 import requests
+import yaml
 
 
 def load_component_from_text(text: str) -> yaml_component.YamlComponent:
@@ -30,7 +34,7 @@ def load_component_from_text(text: str) -> yaml_component.YamlComponent:
         Component loaded from YAML.
     """
     return yaml_component.YamlComponent(
-        component_spec=structures.ComponentSpec.from_yaml_documents(text),
+        component_spec=_load_component_spec_from_yaml_documents(text),
         component_yaml=text)
 
 
@@ -86,3 +90,97 @@ def load_component_from_url(
     resp.raise_for_status()
 
     return load_component_from_text(resp.content.decode('utf-8'))
+
+
+def _load_documents_from_yaml(component_yaml: str) -> Tuple[dict, dict]:
+    """Loads up to two YAML documents from a YAML string.
+
+    First document must always be present. If second document is
+    present, it is returned as a dict, else an empty dict.
+    """
+    documents = list(yaml.safe_load_all(component_yaml))
+    num_docs = len(documents)
+    if num_docs == 1:
+        pipeline_spec_dict = documents[0]
+        platform_spec_dict = {}
+    elif num_docs == 2:
+        pipeline_spec_dict = documents[0]
+        platform_spec_dict = documents[1]
+    else:
+        raise ValueError(
+            f'Expected one or two YAML documents in the IR YAML file. Got: {num_docs}.'
+        )
+    return pipeline_spec_dict, platform_spec_dict
+
+
+def _load_component_spec_from_yaml_documents(
+        component_yaml: str) -> structures.ComponentSpec:
+    """Loads V1 or V2 component YAML into a ComponentSpec.
+
+    Args:
+        component_yaml: PipelineSpec and optionally PlatformSpec YAML documents as a single string.
+
+    Returns:
+        ComponentSpec: The ComponentSpec object.
+    """
+
+    def extract_description(component_yaml: str) -> Union[str, None]:
+        heading = '# Description: '
+        multi_line_description_prefix = '#             '
+        index_of_heading = 2
+        if heading in component_yaml:
+            description = component_yaml.splitlines()[index_of_heading]
+
+            # Multi line
+            comments = component_yaml.splitlines()
+            index = index_of_heading + 1
+            while comments[index][:len(multi_line_description_prefix
+                                      )] == multi_line_description_prefix:
+                description += '\n' + comments[index][
+                    len(multi_line_description_prefix) + 1:]
+                index += 1
+
+            return description[len(heading):]
+        else:
+            return None
+
+    pipeline_spec_dict, platform_spec_dict = _load_documents_from_yaml(
+        component_yaml)
+
+    is_v1 = 'implementation' in set(pipeline_spec_dict.keys())
+    if is_v1:
+        v1_component = load_v1_component_spec_from_component_text(
+            component_yaml)
+        return structures.ComponentSpec.from_v1_component_spec(v1_component)
+    else:
+        component_spec = structures.ComponentSpec.from_ir_dicts(
+            pipeline_spec_dict, platform_spec_dict)
+        if not component_spec.description:
+            component_spec.description = extract_description(
+                component_yaml=component_yaml)
+        return component_spec
+
+
+def load_v1_component_spec_from_component_text(
+        text) -> v1_structures.ComponentSpec:
+    component_dict = yaml.safe_load(text)
+    component_spec = v1_structures.ComponentSpec.from_dict(component_dict)
+
+    if isinstance(component_spec.implementation,
+                  v1_structures.ContainerImplementation) and (
+                      component_spec.implementation.container.command is None):
+        warnings.warn(
+            'Container component must specify command to be compatible with KFP '
+            'v2 compatible mode and emissary executor, which will be the default'
+            ' executor for KFP v2.'
+            'https://www.kubeflow.org/docs/components/pipelines/installation/choose-executor/',
+            category=FutureWarning,
+        )
+
+    # Calculating hash digest for the component
+    data = text if isinstance(text, bytes) else text.encode('utf-8')
+    data = data.replace(b'\r\n', b'\n')  # Normalizing line endings
+    digest = hashlib.sha256(data).hexdigest()
+    component_spec._digest = digest
+
+    return component_spec

sdk/python/kfp/components/load_yaml_utilities_test.py

@@ -19,6 +19,7 @@ import textwrap
 import unittest
 
 from kfp import components
+from kfp.components import load_yaml_utilities
 from kfp.dsl import structures
 
 SAMPLE_YAML = textwrap.dedent("""\
@@ -124,5 +125,47 @@ class LoadYamlTests(unittest.TestCase):
             'python:3.7')
 
 
+class TestLoadDocumentsFromYAML(unittest.TestCase):
+
+    def test_no_documents(self):
+        with self.assertRaisesRegex(
+                ValueError,
+                r'Expected one or two YAML documents in the IR YAML file\. Got\: 0\.'
+        ):
+            load_yaml_utilities._load_documents_from_yaml('')
+
+    def test_one_document(self):
+        doc1, doc2 = load_yaml_utilities._load_documents_from_yaml(
+            textwrap.dedent("""\
+            key1: value1
+            """))
+        self.assertEqual(doc1, {'key1': 'value1'})
+        self.assertEqual(doc2, {})
+
+    def test_two_documents(self):
+        doc1, doc2 = load_yaml_utilities._load_documents_from_yaml(
+            textwrap.dedent("""\
+            key1: value1
+            ---
+            key2: value2
+            """))
+        self.assertEqual(doc1, {'key1': 'value1'})
+        self.assertEqual(doc2, {'key2': 'value2'})
+
+    def test_three_documents(self):
+        with self.assertRaisesRegex(
+                ValueError,
+                r'Expected one or two YAML documents in the IR YAML file\. Got\: 3\.'
+        ):
+            load_yaml_utilities._load_documents_from_yaml(
+                textwrap.dedent("""\
+                key3: value3
+                ---
+                key3: value3
+                ---
+                key3: value3
+                """))
+
+
 if __name__ == '__main__':
     unittest.main()

sdk/python/kfp/dsl-test/component_decorator_test.py

@@ -17,8 +17,8 @@ import tempfile
 from typing import Dict, List, NamedTuple
 import unittest
 
+from kfp.components import load_yaml_utilities
 from kfp.dsl import python_component
-from kfp.dsl import structures
 from kfp.dsl.component_decorator import component
 
 
@@ -104,7 +104,8 @@ class TestComponentDecorator(unittest.TestCase):
             with open(filepath, 'r') as f:
                 yaml_text = f.read()
 
-        component_spec = structures.ComponentSpec.from_yaml_documents(yaml_text)
+        component_spec = load_yaml_utilities._load_component_spec_from_yaml_documents(
+            yaml_text)
         self.assertEqual(component_spec.name, comp.component_spec.name)
 
     def test_output_named_tuple_with_dict(self):

sdk/python/kfp/dsl-test/pipeline_task_test.py

@@ -18,6 +18,7 @@ import unittest
 
 from absl.testing import parameterized
 from kfp import dsl
+from kfp.components import load_yaml_utilities
 from kfp.dsl import pipeline_task
 from kfp.dsl import placeholders
 from kfp.dsl import structures
@@ -112,8 +113,8 @@ class PipelineTaskTest(parameterized.TestCase):
         )
 
         task = pipeline_task.PipelineTask(
-            component_spec=structures.ComponentSpec.from_yaml_documents(
-                V2_YAML),
+            component_spec=load_yaml_utilities
+            ._load_component_spec_from_yaml_documents(V2_YAML),
             args={'input1': 'value'},
         )
         self.assertEqual(task._task_spec, expected_task_spec)
@@ -125,8 +126,8 @@ class PipelineTaskTest(parameterized.TestCase):
                 ValueError,
                 "Component 'component1' got an unexpected input: 'input0'."):
             task = pipeline_task.PipelineTask(
-                component_spec=structures.ComponentSpec.from_yaml_documents(
-                    V2_YAML),
+                component_spec=load_yaml_utilities
+                ._load_component_spec_from_yaml_documents(V2_YAML),
                 args={
                     'input1': 'value',
                     'input0': 'abc',
@@ -135,8 +136,8 @@ class PipelineTaskTest(parameterized.TestCase):
 
     def test_set_caching_options(self):
         task = pipeline_task.PipelineTask(
-            component_spec=structures.ComponentSpec.from_yaml_documents(
-                V2_YAML),
+            component_spec=load_yaml_utilities
+            ._load_component_spec_from_yaml_documents(V2_YAML),
             args={'input1': 'value'},
         )
         task.set_caching_options(False)
@@ -163,8 +164,8 @@ class PipelineTaskTest(parameterized.TestCase):
     def test_set_valid_cpu_request_limit(self, cpu: str,
                                          expected_cpu_number: float):
         task = pipeline_task.PipelineTask(
-            component_spec=structures.ComponentSpec.from_yaml_documents(
-                V2_YAML),
+            component_spec=load_yaml_utilities
+            ._load_component_spec_from_yaml_documents(V2_YAML),
             args={'input1': 'value'},
         )
         task.set_cpu_request(cpu)
@@ -182,8 +183,8 @@ class PipelineTaskTest(parameterized.TestCase):
     def test_set_valid_gpu_limit(self, gpu_limit: str,
                                  expected_gpu_number: int):
         task = pipeline_task.PipelineTask(
-            component_spec=structures.ComponentSpec.from_yaml_documents(
-                V2_YAML),
+            component_spec=load_yaml_utilities
+            ._load_component_spec_from_yaml_documents(V2_YAML),
             args={'input1': 'value'},
         )
         with self.assertWarnsRegex(
@@ -196,8 +197,8 @@ class PipelineTaskTest(parameterized.TestCase):
 
     def test_add_valid_node_selector_constraint(self):
         task = pipeline_task.PipelineTask(
-            component_spec=structures.ComponentSpec.from_yaml_documents(
-                V2_YAML),
+            component_spec=load_yaml_utilities
+            ._load_component_spec_from_yaml_documents(V2_YAML),
             args={'input1': 'value'},
         )
         with self.assertWarnsRegex(
@@ -220,8 +221,8 @@ class PipelineTaskTest(parameterized.TestCase):
     )
     def test_set_accelerator_limit(self, limit, expected):
         task = pipeline_task.PipelineTask(
-            component_spec=structures.ComponentSpec.from_yaml_documents(
-                V2_YAML),
+            component_spec=load_yaml_utilities
+            ._load_component_spec_from_yaml_documents(V2_YAML),
             args={'input1': 'value'},
         )
 
@@ -285,8 +286,8 @@ class PipelineTaskTest(parameterized.TestCase):
     )
     def test_set_memory_limit(self, memory: str, expected_memory_number: int):
         task = pipeline_task.PipelineTask(
-            component_spec=structures.ComponentSpec.from_yaml_documents(
-                V2_YAML),
+            component_spec=load_yaml_utilities
+            ._load_component_spec_from_yaml_documents(V2_YAML),
             args={'input1': 'value'},
         )
         task.set_memory_request(memory)
@@ -298,8 +299,8 @@ class PipelineTaskTest(parameterized.TestCase):
 
     def test_set_accelerator_type_with_type_only(self):
         task = pipeline_task.PipelineTask(
-            component_spec=structures.ComponentSpec.from_yaml_documents(
-                V2_YAML),
+            component_spec=load_yaml_utilities
+            ._load_component_spec_from_yaml_documents(V2_YAML),
             args={'input1': 'value'},
         )
         task.set_accelerator_type('NVIDIA_TESLA_K80')
@@ -310,8 +311,8 @@ class PipelineTaskTest(parameterized.TestCase):
 
     def test_set_accelerator_type_with_accelerator_count(self):
         task = pipeline_task.PipelineTask(
-            component_spec=structures.ComponentSpec.from_yaml_documents(
-                V2_YAML),
+            component_spec=load_yaml_utilities
+            ._load_component_spec_from_yaml_documents(V2_YAML),
             args={'input1': 'value'},
         )
         task.set_accelerator_limit('5').set_accelerator_type('TPU_V3')
@@ -322,8 +323,8 @@ class PipelineTaskTest(parameterized.TestCase):
 
     def test_set_env_variable(self):
         task = pipeline_task.PipelineTask(
-            component_spec=structures.ComponentSpec.from_yaml_documents(
-                V2_YAML),
+            component_spec=load_yaml_utilities
+            ._load_component_spec_from_yaml_documents(V2_YAML),
             args={'input1': 'value'},
         )
         task.set_env_variable('env_name', 'env_value')
@@ -331,8 +332,8 @@ class PipelineTaskTest(parameterized.TestCase):
 
     def test_set_display_name(self):
         task = pipeline_task.PipelineTask(
-            component_spec=structures.ComponentSpec.from_yaml_documents(
-                V2_YAML),
+            component_spec=load_yaml_utilities
+            ._load_component_spec_from_yaml_documents(V2_YAML),
             args={'input1': 'value'},
         )
         task.set_display_name('test_name')

sdk/python/kfp/dsl-test/structures_test.py

@@ -22,6 +22,7 @@ from absl.testing import parameterized
 from kfp import compiler
 from kfp import components
 from kfp import dsl
+from kfp.components import load_yaml_utilities
 from kfp.dsl import component_factory
 from kfp.dsl import placeholders
 from kfp.dsl import structures
@@ -263,7 +264,7 @@ class StructuresTest(parameterized.TestCase):
             # test that it can be read back correctly
             with open(output_path, 'r') as f:
                 contents = f.read()
-            new_component_spec = structures.ComponentSpec.from_yaml_documents(
+            new_component_spec = load_yaml_utilities._load_component_spec_from_yaml_documents(
                 contents)
 
         self.assertEqual(original_component_spec, new_component_spec)
@@ -318,7 +319,7 @@ schemaVersion: 2.1.0
 sdkVersion: kfp-2.0.0-alpha.2
         """)
 
-        generated_spec = structures.ComponentSpec.from_yaml_documents(
+        generated_spec = load_yaml_utilities._load_component_spec_from_yaml_documents(
             component_yaml_v2)
 
         expected_spec = structures.ComponentSpec(
@@ -359,7 +360,8 @@ sdkVersion: kfp-2.0.0-alpha.2
     )
     def test_component_spec_placeholder_load_from_v2_component_yaml(
             self, yaml, expected_component):
-        generated_spec = structures.ComponentSpec.from_yaml_documents(yaml)
+        generated_spec = load_yaml_utilities._load_component_spec_from_yaml_documents(
+            yaml)
         self.assertEqual(generated_spec, expected_component)
 
     def test_component_spec_load_from_v1_component_yaml(self):
@@ -388,7 +390,7 @@ sdkVersion: kfp-2.0.0-alpha.2
             - {outputPath: Output 2}
         """)
 
-        generated_spec = structures.ComponentSpec.from_yaml_documents(
+        generated_spec = load_yaml_utilities._load_component_spec_from_yaml_documents(
             component_yaml_v1)
 
         expected_spec = structures.ComponentSpec(
@@ -639,7 +641,7 @@ V1_YAML = textwrap.dedent("""\
 class TestReadInComponent(parameterized.TestCase):
 
     def test_read_v1(self):
-        component_spec = structures.ComponentSpec.from_yaml_documents(
+        component_spec = load_yaml_utilities._load_component_spec_from_yaml_documents(
             V1_YAML_IF_PLACEHOLDER)
         self.assertEqual(component_spec.name, 'component-if')
         self.assertEqual(component_spec.implementation.container.image,
@@ -694,7 +696,7 @@ root:
         parameterType: STRING
 schemaVersion: 2.1.0
 sdkVersion: kfp-2.0.0-alpha.2""")
-        loaded_component_spec = structures.ComponentSpec.from_yaml_documents(
+        loaded_component_spec = load_yaml_utilities._load_component_spec_from_yaml_documents(
             compiled_yaml)
         component_spec = structures.ComponentSpec(
             name='component1',
@@ -762,7 +764,7 @@ root:
         parameterType: STRING
 schemaVersion: 2.1.0
 sdkVersion: kfp-2.0.0-alpha.2""")
-        loaded_component_spec = structures.ComponentSpec.from_yaml_documents(
+        loaded_component_spec = load_yaml_utilities._load_component_spec_from_yaml_documents(
             compiled_yaml)
         component_spec = structures.ComponentSpec(
             name='if',
@@ -833,7 +835,7 @@ root:
         parameterType: STRING
 schemaVersion: 2.1.0
 sdkVersion: kfp-2.0.0-alpha.2""")
-        loaded_component_spec = structures.ComponentSpec.from_yaml_documents(
+        loaded_component_spec = load_yaml_utilities._load_component_spec_from_yaml_documents(
             compiled_yaml)
         component_spec = structures.ComponentSpec(
             name='concat',
@@ -1113,47 +1115,5 @@ implementation:
         self.assertEqual(outputs['output4'].type, 'Dict')
 
 
-class TestLoadDocumentsFromYAML(unittest.TestCase):
-
-    def test_no_documents(self):
-        with self.assertRaisesRegex(
-                ValueError,
-                r'Expected one or two YAML documents in the IR YAML file\. Got\: 0\.'
-        ):
-            structures.load_documents_from_yaml('')
-
-    def test_one_document(self):
-        doc1, doc2 = structures.load_documents_from_yaml(
-            textwrap.dedent("""\
-            key1: value1
-            """))
-        self.assertEqual(doc1, {'key1': 'value1'})
-        self.assertEqual(doc2, {})
-
-    def test_two_documents(self):
-        doc1, doc2 = structures.load_documents_from_yaml(
-            textwrap.dedent("""\
-            key1: value1
-            ---
-            key2: value2
-            """))
-        self.assertEqual(doc1, {'key1': 'value1'})
-        self.assertEqual(doc2, {'key2': 'value2'})
-
-    def test_three_documents(self):
-        with self.assertRaisesRegex(
-                ValueError,
-                r'Expected one or two YAML documents in the IR YAML file\. Got\: 3\.'
-        ):
-            structures.load_documents_from_yaml(
-                textwrap.dedent("""\
-                key3: value3
-                ---
-                key3: value3
-                ---
-                key3: value3
-                """))
-
-
 if __name__ == '__main__':
     unittest.main()

sdk/python/kfp/dsl-test/types/test_data/expected_bulk_loaded_confusion_matrix.json

@@ -7,4 +7,4 @@
     [2, 6, 0],
     [3, 5, 6],
     [5, 7, 8]]
-}
+}

sdk/python/kfp/dsl-test/types/test_data/expected_confusion_matrix.json

@@ -7,4 +7,4 @@
     [2, 6, 0],
     [3, 0, 0],
     [0, 0, 0]]
-}
+}

sdk/python/kfp/dsl-test/types/test_data/expected_io_types_bulk_load_classification_metrics.json

@@ -1,11 +1,11 @@
 {
-  "confidenceMetrics": 
+  "confidenceMetrics":
   [
     {
       "confidenceThreshold": 53.6,
       "recall": 52.6,
       "falsePositiveRate": 85.1
-    }, 
+    },
     {
       "confidenceThreshold": 53.6,
       "recall": 52.6,
@@ -16,20 +16,20 @@
       "recall": 52.6,
       "falsePositiveRate": 85.1
     }
-  ], 
-  "confusionMatrix": 
+  ],
+  "confusionMatrix":
   {
-    "annotationSpecs": 
+    "annotationSpecs":
     [
       {"displayName": "dog"},
       {"displayName": "cat"},
       {"displayName": "horses"}
     ],
-    "rows": 
+    "rows":
     [
       {"row": [2, 6, 0]},
       {"row": [3, 5, 6]},
       {"row" : [5, 7, 8]}
     ]
   }
-}
+}

sdk/python/kfp/dsl-test/types/test_data/expected_io_types_classification_metrics.json

@@ -5,22 +5,22 @@
       "confidenceThreshold": 0.1,
       "recall": 98.2,
       "falsePositiveRate": 96.2
-    }, 
+    },
     {
       "confidenceThreshold": 24.3,
       "recall": 24.5,
       "falsePositiveRate": 98.4
     }
   ],
-  "confusionMatrix": 
+  "confusionMatrix":
     {
-      "annotationSpecs": 
+      "annotationSpecs":
       [
         {"displayName": "dog"},
         {"displayName": "cat"},
         {"displayName": "horses"}
-      ], 
-      "rows": 
+      ],
+      "rows":
         [
           {"row" : [2, 6, 0]},
           {"row" : [3, 0, 0]},

sdk/python/kfp/dsl-test/types/type_utils_test.py

@@ -727,6 +727,7 @@ class TestTypeChecking(parameterized.TestCase):
                     given_value=argument_value,
                     expected_spec=parameter_input_spec,
                     error_message_prefix='',
+                    raise_on_error=kfp.TYPE_CHECK,
                 ))
         else:
             with self.assertRaises(InconsistentTypeException):
@@ -734,6 +735,7 @@ class TestTypeChecking(parameterized.TestCase):
                     given_value=argument_value,
                     expected_spec=parameter_input_spec,
                     error_message_prefix='',
+                    raise_on_error=kfp.TYPE_CHECK,
                 )
 
     def test_list_of_artifacts_across_compilation_valid(self):

sdk/python/kfp/dsl/v1_components.py

@@ -1,44 +0,0 @@
-# Copyright 2018-2022 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.
-
-import hashlib
-import warnings
-
-from kfp.dsl import v1_structures
-import yaml
-
-
-def _load_component_spec_from_component_text(
-        text) -> v1_structures.ComponentSpec:
-    component_dict = yaml.safe_load(text)
-    component_spec = v1_structures.ComponentSpec.from_dict(component_dict)
-
-    if isinstance(component_spec.implementation,
-                  v1_structures.ContainerImplementation) and (
-                      component_spec.implementation.container.command is None):
-        warnings.warn(
-            'Container component must specify command to be compatible with KFP '
-            'v2 compatible mode and emissary executor, which will be the default'
-            ' executor for KFP v2.'
-            'https://www.kubeflow.org/docs/components/pipelines/installation/choose-executor/',
-            category=FutureWarning,
-        )
-
-    # Calculating hash digest for the component
-    data = text if isinstance(text, bytes) else text.encode('utf-8')
-    data = data.replace(b'\r\n', b'\n')  # Normalizing line endings
-    digest = hashlib.sha256(data).hexdigest()
-    component_spec._digest = digest
-
-    return component_spec

sdk/python/requirements.in

@@ -10,6 +10,7 @@ google-api-core>=1.31.5,<3.0.0dev,!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.0
 google-auth>=1.6.1,<3
 # https://github.com/googleapis/python-storage/blob/main/CHANGELOG.md#221-2022-03-15
 google-cloud-storage>=2.2.1,<3
+kfp-dsl==2.0.1
 # pin kfp-pipeline-spec to an exact version, since this is the contract between a given KFP SDK version and the BE. we don't want old version of the SDK to write new fields and to have the BE reject the new unsupported field (even if the new field backward compatible from a proto perspective)
 kfp-pipeline-spec==0.2.2
 # Update the upper version whenever a new major version of the

sdk/python/test_data/components/add_numbers.yaml

@@ -32,7 +32,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec

sdk/python/test_data/components/component_with_metadata_fields.yaml

@@ -48,7 +48,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec

sdk/python/test_data/components/component_with_pip_install.yaml

@@ -19,7 +19,7 @@ deploymentSpec:
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
           \ python3 -m pip install --quiet     --no-warn-script-location --index-url\
           \ https://pypi.org/simple --trusted-host https://pypi.org/simple 'yapf'\
-          \ 'kfp==2.0.1' && \"$0\" \"$@\"\n"
+          \ 'kfp-dsl==2.0.1' && \"$0\" \"$@\"\n"
         - sh
         - -ec
         - 'program_path=$(mktemp -d)

sdk/python/test_data/components/component_with_task_final_status.yaml

@@ -24,7 +24,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec

sdk/python/test_data/components/concat_message.yaml

@@ -32,7 +32,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec

sdk/python/test_data/components/dict_input.yaml

@@ -23,7 +23,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec

sdk/python/test_data/components/identity.yaml

@@ -29,7 +29,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec

sdk/python/test_data/components/input_artifact.yaml

@@ -25,7 +25,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec

sdk/python/test_data/components/nested_return.yaml

@@ -23,7 +23,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec

sdk/python/test_data/components/output_metrics.yaml

@@ -27,7 +27,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec

sdk/python/test_data/components/preprocess.yaml

@@ -56,7 +56,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec

sdk/python/test_data/pipelines/component_with_optional_inputs.yaml

@@ -29,7 +29,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec

sdk/python/test_data/pipelines/component_with_pip_index_urls.yaml

@@ -19,7 +19,7 @@ deploymentSpec:
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
           \ python3 -m pip install --quiet     --no-warn-script-location --index-url\
           \ https://pypi.org/simple --trusted-host https://pypi.org/simple 'yapf'\
-          \ 'kfp==2.0.1' && \"$0\" \"$@\"\n"
+          \ 'kfp-dsl==2.0.1' && \"$0\" \"$@\"\n"
         - sh
         - -ec
         - 'program_path=$(mktemp -d)

sdk/python/test_data/pipelines/components_with_optional_artifacts.yaml

@@ -126,7 +126,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec
@@ -155,7 +155,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec

sdk/python/test_data/pipelines/lightweight_python_functions_pipeline.yaml

@@ -78,7 +78,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec
@@ -130,7 +130,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec

sdk/python/test_data/pipelines/lightweight_python_functions_with_outputs.yaml

@@ -81,7 +81,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec
@@ -108,7 +108,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec
@@ -135,7 +135,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec
@@ -162,7 +162,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec

sdk/python/test_data/pipelines/parallelfor_fan_in/artifacts_complex.yaml

@@ -285,7 +285,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec
@@ -315,7 +315,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec
@@ -345,7 +345,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec
@@ -375,7 +375,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec
@@ -403,7 +403,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec

sdk/python/test_data/pipelines/parallelfor_fan_in/artifacts_simple.yaml

@@ -90,7 +90,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec
@@ -136,7 +136,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec

sdk/python/test_data/pipelines/parallelfor_fan_in/conditional_producer_and_consumers.yaml

@@ -132,7 +132,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec
@@ -158,7 +158,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec

sdk/python/test_data/pipelines/parallelfor_fan_in/nested_with_parameters.yaml

@@ -150,7 +150,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec
@@ -177,7 +177,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec
@@ -203,7 +203,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec
@@ -229,7 +229,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec

sdk/python/test_data/pipelines/parallelfor_fan_in/parameters_complex.yaml

@@ -224,7 +224,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec
@@ -251,7 +251,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec
@@ -277,7 +277,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec
@@ -303,7 +303,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec
@@ -330,7 +330,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec
@@ -357,7 +357,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec
@@ -383,7 +383,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec

sdk/python/test_data/pipelines/parallelfor_fan_in/parameters_simple.yaml

@@ -75,7 +75,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec
@@ -111,7 +111,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec

sdk/python/test_data/pipelines/parallelfor_fan_in/pipeline_producer_consumer.yaml

@@ -206,7 +206,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec
@@ -233,7 +233,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec
@@ -259,7 +259,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec
@@ -286,7 +286,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec

sdk/python/test_data/pipelines/pipeline_as_exit_task.yaml

@@ -129,7 +129,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec
@@ -156,7 +156,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec
@@ -183,7 +183,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec
@@ -210,7 +210,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec

sdk/python/test_data/pipelines/pipeline_in_pipeline.yaml

@@ -74,7 +74,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec
@@ -101,7 +101,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec

sdk/python/test_data/pipelines/pipeline_in_pipeline_complex.yaml

@@ -161,7 +161,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec
@@ -188,7 +188,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec

sdk/python/test_data/pipelines/pipeline_in_pipeline_loaded_from_yaml.yaml

@@ -152,7 +152,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec
@@ -179,7 +179,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec
@@ -206,7 +206,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec

sdk/python/test_data/pipelines/pipeline_with_condition.yaml

@@ -88,7 +88,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec
@@ -116,7 +116,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec
@@ -144,7 +144,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec
@@ -171,7 +171,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec
@@ -198,7 +198,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec

sdk/python/test_data/pipelines/pipeline_with_dynamic_importer_metadata.yaml

@@ -94,7 +94,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec

sdk/python/test_data/pipelines/pipeline_with_env.yaml

@@ -41,7 +41,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec

sdk/python/test_data/pipelines/pipeline_with_exit_handler.yaml

@@ -65,7 +65,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec
@@ -92,7 +92,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec
@@ -119,7 +119,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.1'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp-dsl==2.0.1'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec

sdk/python/test_data/pipelines/pipeline_with_google_artifact_type.yaml

@@ -57,7 +57,7 @@ deploymentSpec:
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
           \ python3 -m pip install --quiet     --no-warn-script-location 'aiplatform'\
-          \ 'kfp==2.0.1' 'kfp==2.0.1' && \"$0\" \"$@\"\n"
+          \ 'kfp-dsl==2.0.1' 'kfp-dsl==2.0.1' && \"$0\" \"$@\"\n"
         - sh
         - -ec
         - 'program_path=$(mktemp -d)
@@ -90,7 +90,7 @@ deploymentSpec:
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
           \ python3 -m pip install --quiet     --no-warn-script-location 'aiplatform'\
-          \ 'kfp==2.0.1' && \"$0\" \"$@\"\n"
+          \ 'kfp-dsl==2.0.1' && \"$0\" \"$@\"\n"
         - sh
         - -ec
         - 'program_path=$(mktemp -d)