use packages_to_install for kfp spec

Signed-off-by: Humair Khan <[email protected]>

Author: HumairAK

samples/v2/sample_test.py

@@ -48,11 +48,7 @@ def get_package_path(subdir: str) -> str:
 
 # Set the component configuration BEFORE importing any pipeline modules
 dsl.component = functools.partial(
-    dsl.component, kfp_package_path=get_kfp_package_path())
-
-# Ensure that we are installing pipeline-spec from source during pipeline execution.
-if 'KFP_PIPELINE_SPEC_PACKAGE_PATH' not in os.environ:
-    os.environ['KFP_PIPELINE_SPEC_PACKAGE_PATH'] = get_kfp_pipeline_spec_path()
+    dsl.component, kfp_package_path=get_kfp_package_path(), packages_to_install=[get_kfp_pipeline_spec_path()])
 
 # Now import the pipeline modules, this way we can leverage the kfp_package and pipeline
 # spec defined above

sdk/python/kfp/dsl/component_factory.py

@@ -14,7 +14,6 @@
 import dataclasses
 import inspect
 import itertools
-import os
 import pathlib
 import re
 import textwrap
@@ -144,8 +143,6 @@ def make_pip_install_command(
 . "$tmp/venv/bin/activate"
 '''
 
-_KFP_PIPELINE_SPEC_PACKAGE_PATH = "KFP_PIPELINE_SPEC_PACKAGE_PATH"
-
 def _get_packages_to_install_command(
     kfp_package_path: Optional[str] = None,
     pip_index_urls: Optional[List[str]] = None,
@@ -167,17 +164,17 @@ def _get_packages_to_install_command(
     index_url_options = make_index_url_options(pip_index_urls,
                                                pip_trusted_hosts)
 
-    # Setting pipeline spec path is not intended to be public api and
-    # is largely intended for testing and development purposes to
-    # allow pipeline runtime execution to install pipeline spec from
-    # development branches and pull request CIs.
-    pipeline_spec_path = os.getenv(_KFP_PIPELINE_SPEC_PACKAGE_PATH)
-    if pipeline_spec_path:
-        kfp_pipeline_spec_install_command = make_pip_install_command(
-            install_parts=[pipeline_spec_path],
-            index_url_options=index_url_options
+    # Install packages before KFP. This allows us to
+    # control where we source kfp-pipeline-spec.
+    # This is particularly useful for development and
+    # CI use-case when you want to install the spec
+    # from source.
+    if packages_to_install:
+        user_packages_pip_install_command = make_pip_install_command(
+            install_parts=packages_to_install,
+            index_url_options=index_url_options,
         )
-        pip_install_strings.append(kfp_pipeline_spec_install_command)
+        pip_install_strings.append(user_packages_pip_install_command)
         if inject_kfp_install:
             pip_install_strings.append(' && ')
 
@@ -200,16 +197,6 @@ def _get_packages_to_install_command(
             )
         pip_install_strings.append(kfp_pip_install_command)
 
-        if packages_to_install:
-            pip_install_strings.append(' && ')
-
-    if packages_to_install:
-        user_packages_pip_install_command = make_pip_install_command(
-            install_parts=packages_to_install,
-            index_url_options=index_url_options,
-        )
-        pip_install_strings.append(user_packages_pip_install_command)
-
     return [
         'sh', '-c',
         _install_python_packages_script_template.format(

sdk/python/kfp/dsl/component_factory_test.py

@@ -132,8 +132,18 @@ def test_with_user_packages_to_install_and_no_pip_index_url(self):
         self.assertEqual(
             strip_kfp_version(command),
             strip_kfp_version([
-                '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==2.1.3\' \'--no-deps\' \'typing-extensions>=3.7.4,<5; python_version<"3.9"\'  &&  python3 -m pip install --quiet --no-warn-script-location \'package1\' \'package2\' && "$0" "$@"\n'
+                'sh',
+                '-c',
+                '\n'
+                'if ! [ -x "$(command -v pip)" ]; then\n'
+                '    python3 -m ensurepip || python3 -m ensurepip --user || apt-get install '
+                'python3-pip\n'
+                'fi\n'
+                '\n'
+                'PIP_DISABLE_PIP_VERSION_CHECK=1 python3 -m pip install --quiet '
+                "--no-warn-script-location 'package1' 'package2'  &&  python3 -m pip install "
+                "--quiet --no-warn-script-location kfp '--no-deps' "
+                '\'typing-extensions>=3.7.4,<5; python_version<"3.9"\' && "$0" "$@"\n'
             ]))
 
     def test_with_packages_to_install_with_pip_index_url(self):
@@ -148,8 +158,21 @@ def test_with_packages_to_install_with_pip_index_url(self):
         self.assertEqual(
             strip_kfp_version(command),
             strip_kfp_version([
-                '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 --index-url https://myurl.org/simple --trusted-host https://myurl.org/simple \'kfp==2.1.3\' \'--no-deps\' \'typing-extensions>=3.7.4,<5; python_version<"3.9"\'  &&  python3 -m pip install --quiet --no-warn-script-location --index-url https://myurl.org/simple --trusted-host https://myurl.org/simple \'package1\' \'package2\' && "$0" "$@"\n'
+                'sh',
+                '-c',
+                '\n'
+                'if ! [ -x "$(command -v pip)" ]; then\n'
+                '    python3 -m ensurepip || python3 -m ensurepip --user || apt-get install '
+                'python3-pip\n'
+                'fi\n'
+                '\n'
+                'PIP_DISABLE_PIP_VERSION_CHECK=1 python3 -m pip install --quiet '
+                '--no-warn-script-location --index-url https://myurl.org/simple '
+                "--trusted-host https://myurl.org/simple 'package1' 'package2'  &&  python3 "
+                '-m pip install --quiet --no-warn-script-location --index-url '
+                'https://myurl.org/simple --trusted-host https://myurl.org/simple kfp '
+                '\'--no-deps\' \'typing-extensions>=3.7.4,<5; python_version<"3.9"\' && "$0" '
+                '"$@"\n'
             ]))
 
     def test_with_packages_to_install_with_pip_index_url_and_trusted_host(self):
@@ -166,8 +189,20 @@ def test_with_packages_to_install_with_pip_index_url_and_trusted_host(self):
         self.assertEqual(
             strip_kfp_version(command),
             strip_kfp_version([
-                '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 --index-url https://myurl.org/simple --trusted-host myurl.org \'kfp==2.1.3\' \'--no-deps\' \'typing-extensions>=3.7.4,<5; python_version<"3.9"\'  &&  python3 -m pip install --quiet --no-warn-script-location --index-url https://myurl.org/simple --trusted-host myurl.org \'package1\' \'package2\' && "$0" "$@"\n'
+                'sh',
+                '-c',
+                '\n'
+                'if ! [ -x "$(command -v pip)" ]; then\n'
+                '    python3 -m ensurepip || python3 -m ensurepip --user || apt-get install '
+                'python3-pip\n'
+                'fi\n'
+                '\n'
+                'PIP_DISABLE_PIP_VERSION_CHECK=1 python3 -m pip install --quiet '
+                '--no-warn-script-location --index-url https://myurl.org/simple '
+                "--trusted-host myurl.org 'package1' 'package2'  &&  python3 -m pip install "
+                '--quiet --no-warn-script-location --index-url https://myurl.org/simple '
+                "--trusted-host myurl.org kfp '--no-deps' 'typing-extensions>=3.7.4,<5; "
+                'python_version<"3.9"\' && "$0" "$@"\n'
             ]))
 
     def test_with_packages_to_install_with_pip_index_url_and_empty_trusted_host(
@@ -183,10 +218,23 @@ def test_with_packages_to_install_with_pip_index_url_and_empty_trusted_host(
 
         self.assertEqual(
             strip_kfp_version(command),
-            strip_kfp_version([
-                '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 --index-url https://myurl.org/simple \'kfp==2.1.3\' \'--no-deps\' \'typing-extensions>=3.7.4,<5; python_version<"3.9"\'  &&  python3 -m pip install --quiet --no-warn-script-location --index-url https://myurl.org/simple \'package1\' \'package2\' && "$0" "$@"\n'
-            ]))
+            strip_kfp_version(
+                [
+                    'sh',
+                    '-c',
+                    '\n'
+                    'if ! [ -x "$(command -v pip)" ]; then\n'
+                    '    python3 -m ensurepip || python3 -m ensurepip --user || apt-get install '
+                    'python3-pip\n'
+                    'fi\n'
+                    '\n'
+                    'PIP_DISABLE_PIP_VERSION_CHECK=1 python3 -m pip install --quiet '
+                    "--no-warn-script-location --index-url https://myurl.org/simple 'package1' "
+                    "'package2'  &&  python3 -m pip install --quiet --no-warn-script-location "
+                    "--index-url https://myurl.org/simple kfp '--no-deps' "
+                    '\'typing-extensions>=3.7.4,<5; python_version<"3.9"\' && "$0" "$@"\n'
+                ]
+            ))
 
 
 class TestInvalidParameterName(unittest.TestCase):

sdk/python/kfp/local/pipeline_orchestrator_test.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Tests for pipeline_orchestrator.py."""
-
+import functools
 import io as stdlib_io
 import os
 from typing import NamedTuple
@@ -31,16 +31,11 @@
 from kfp.local import testing_utilities
 
 ROOT_FOR_TESTING = './testing_root'
+CURRENT_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__)))))
+kfp_pipeline_spec_path = os.path.join(CURRENT_DIR, 'api', 'v2alpha1', 'python')
 
-@pytest.fixture(autouse=True)
-def set_env_for_test_classes(monkeypatch, request):
-    root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__)))))
-    # Only apply to specific classes
-    if request.cls.__name__ in {
-        "TestRunLocalPipeline",
-        "TestFstringContainerComponent",
-    }:
-        monkeypatch.setenv("KFP_PIPELINE_SPEC_PACKAGE_PATH", os.path.join(root, 'api', 'v2alpha1', 'python'))
+dsl.component = functools.partial(
+    dsl.component, packages_to_install=[kfp_pipeline_spec_path])
 
 
 class TestRunLocalPipeline(testing_utilities.LocalRunnerEnvironmentTestCase):
@@ -88,7 +83,7 @@ def my_pipeline():
     def test_no_io(self):
         local.init(local.SubprocessRunner(), pipeline_root=ROOT_FOR_TESTING)
 
-        @dsl.component
+        @dsl.component()
         def pass_op():
             pass
 
@@ -721,5 +716,6 @@ def my_pipeline(string: str = 'baz') -> str:
         task = my_pipeline()
         self.assertEqual(task.output, 'foo-bar-baz')
 
+
 if __name__ == '__main__':
     unittest.main()

sdk/python/kfp/local/subprocess_task_handler_test.py

@@ -13,13 +13,13 @@
 # limitations under the License.
 """Tests for subprocess_local_task_handler.py."""
 import contextlib
+import functools
 import io
 import os
 from typing import NamedTuple, Optional
 import unittest
 from unittest import mock
 
-import pytest
 from absl.testing import parameterized
 from kfp import dsl
 from kfp import local
@@ -36,18 +36,11 @@
 # impact of such an error we should not install into the main test process'
 # environment.
 
-@pytest.fixture(autouse=True)
-def set_env_for_test_classes(monkeypatch, request):
-    root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__)))))
-    # Only apply to specific classes
-    if request.cls.__name__ in {
-        "TestSubprocessRunner",
-        "TestRunLocalSubproces",
-        "TestUseCurrentPythonExecutable",
-        "TestUseVenv",
-        "TestLightweightPythonComponentLogic"
-    }:
-        monkeypatch.setenv("KFP_PIPELINE_SPEC_PACKAGE_PATH", os.path.join(root, 'api', 'v2alpha1', 'python'))
+CURRENT_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__)))))
+kfp_pipeline_spec_path = os.path.join(CURRENT_DIR, 'api', 'v2alpha1', 'python')
+
+dsl.component = functools.partial(
+    dsl.component, packages_to_install=[kfp_pipeline_spec_path])
 
 class TestSubprocessRunner(testing_utilities.LocalRunnerEnvironmentTestCase):
 
@@ -99,7 +92,7 @@ def comp():
     def test_cannot_run_containerized_python_component(self):
         local.init(runner=local.SubprocessRunner(use_venv=True))
 
-        @dsl.component(target_image='foo')
+        @dsl.component(target_image='foo', packages_to_install=[])
         def comp():
             pass
 
@@ -152,7 +145,7 @@ class TestUseVenv(testing_utilities.LocalRunnerEnvironmentTestCase):
     def test_use_venv_true(self, **kwargs):
         local.init(**kwargs)
 
-        @dsl.component(packages_to_install=['cloudpickle'])
+        @dsl.component(packages_to_install=[kfp_pipeline_spec_path, 'cloudpickle'])
         def installer_component():
             import cloudpickle
             print('Cloudpickle is installed:', cloudpickle)

sdk/python/kfp/local/task_dispatcher_test.py

@@ -19,13 +19,13 @@
 irrespective of the runner. While there will inevitably some overlap, we
 should seek to minimize it.
 """
+import functools
 import io
 import os
 import re
 import unittest
 from unittest import mock
 
-import pytest
 from absl.testing import parameterized
 from kfp import dsl
 from kfp import local
@@ -42,18 +42,11 @@
 # impact of such an error we should not install into the main test process'
 # environment.
 
-@pytest.fixture(autouse=True)
-def set_env_for_test_classes(monkeypatch, request):
-    root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__)))))
-    # Only apply to specific classes
-    if request.cls.__name__ in {
-        "TestLocalExecutionValidation",
-        "TestSupportOfComponentTypes",
-        "TestSupportOfComponentTypes",
-        "TestExceptionHandlingAndLogging",
-        "TestPipelineRootPaths"
-    }:
-        monkeypatch.setenv("KFP_PIPELINE_SPEC_PACKAGE_PATH", os.path.join(root, 'api', 'v2alpha1', 'python'))
+CURRENT_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__)))))
+kfp_pipeline_spec_path = os.path.join(CURRENT_DIR, 'api', 'v2alpha1', 'python')
+
+dsl.component = functools.partial(
+    dsl.component, packages_to_install=[kfp_pipeline_spec_path])
 
 class TestLocalExecutionValidation(
         testing_utilities.LocalRunnerEnvironmentTestCase):

sdk/python/test_data/components/component_with_pip_install.py

@@ -11,11 +11,17 @@
 # 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
+
 from kfp.dsl import component
 
 
+PACKAGES_TO_INSTALL=['yapf']
+if 'KFP_PIPELINE_SPEC_PACKAGE_PATH' in os.environ:
+    PACKAGES_TO_INSTALL.append(os.environ['KFP_PIPELINE_SPEC_PACKAGE_PATH'])
+
 @component(
-    pip_index_urls=['https://pypi.org/simple'], packages_to_install=['yapf'])
+    pip_index_urls=['https://pypi.org/simple'], packages_to_install=PACKAGES_TO_INSTALL)
 def component_with_pip_install():
     import yapf
     print(dir(yapf))

sdk/python/test_data/components/component_with_pip_install.yaml

@@ -18,10 +18,10 @@ 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 --index-url https://pypi.org/simple\
-          \ --trusted-host https://pypi.org/simple 'kfp==2.13.0' '--no-deps' 'typing-extensions>=3.7.4,<5;\
-          \ python_version<\"3.9\"'  &&  python3 -m pip install --quiet --no-warn-script-location\
-          \ --index-url https://pypi.org/simple --trusted-host https://pypi.org/simple\
-          \ 'yapf' && \"$0\" \"$@\"\n"
+          \ --trusted-host https://pypi.org/simple 'yapf'  &&  python3 -m pip install\
+          \ --quiet --no-warn-script-location --index-url https://pypi.org/simple\
+          \ --trusted-host https://pypi.org/simple 'kfp==2.14.0' '--no-deps' 'typing-extensions>=3.7.4,<5;\
+          \ python_version<\"3.9\"' && \"$0\" \"$@\"\n"
         - sh
         - -ec
         - 'program_path=$(mktemp -d)
@@ -50,4 +50,4 @@ root:
           name: component-with-pip-install
           taskName: component-with-pip-install
 schemaVersion: 2.1.0
-sdkVersion: kfp-2.13.0
+sdkVersion: kfp-2.14.0

sdk/python/test_data/components/component_with_pip_install_in_venv.py

@@ -11,12 +11,18 @@
 # 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
+
 from kfp.dsl import component
 
 
+PACKAGES_TO_INSTALL=['yapf']
+if 'KFP_PIPELINE_SPEC_PACKAGE_PATH' in os.environ:
+    PACKAGES_TO_INSTALL.append(os.environ['KFP_PIPELINE_SPEC_PACKAGE_PATH'])
+
 @component(
     pip_index_urls=["https://pypi.org/simple"],
-    packages_to_install=["yapf"],
+    packages_to_install=[PACKAGES_TO_INSTALL],
     use_venv=True,
 )
 def component_with_pip_install():

sdk/python/test_data/pipelines/component_with_pip_index_urls.py

@@ -11,14 +11,18 @@
 # 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
 
 from kfp import compiler
 from kfp import dsl
 from kfp.dsl import component
 
+PACKAGES_TO_INSTALL=['yapf']
+if 'KFP_PIPELINE_SPEC_PACKAGE_PATH' in os.environ:
+    PACKAGES_TO_INSTALL.append(os.environ['KFP_PIPELINE_SPEC_PACKAGE_PATH'])
 
 @component(
-    pip_index_urls=['https://pypi.org/simple'], packages_to_install=['yapf'])
+    pip_index_urls=['https://pypi.org/simple'], packages_to_install=PACKAGES_TO_INSTALL)
 def component_op():
     import yapf
     print(dir(yapf))