changes

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(
@@ -580,6 +567,18 @@ def create_component_from_func(
         pip_trusted_hosts=pip_trusted_hosts,
         use_venv=use_venv,
     )
+    print("hihihi")
+    print("hihihi")
+    print("hihihi")
+    print("hihihi")
+    print("hihihi")
+    print("hihihi")
+    print("hihihi")
+    print("hihihi")
+    print(packages_to_install)
+
+    print(packages_to_install_command)
+
 
     command = []
     args = []

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):
 

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_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))

sdk/python/test_data/pipelines/pipeline_with_google_artifact_type.py

@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import os
 import sys
 import tempfile
 
@@ -76,6 +77,8 @@ def path(self) -> str:
     from kfp.dsl import Output
 
     PACKAGES_TO_INSTALL = ['aiplatform']
+    if 'KFP_PIPELINE_SPEC_PACKAGE_PATH' in os.environ:
+        PACKAGES_TO_INSTALL.append(os.environ['KFP_PIPELINE_SPEC_PACKAGE_PATH'])
 
     @dsl.component(packages_to_install=PACKAGES_TO_INSTALL)
     def model_producer(model: Output[aiplatform.VertexModel]):

sdk/python/test_data/pipelines/pythonic_artifact_with_single_return.py

@@ -11,13 +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 import dsl
 from kfp.dsl import Dataset
 from kfp.dsl import Model
 
+PACKAGES_TO_INSTALL = ['dill==0.3.7']
+if 'KFP_PIPELINE_SPEC_PACKAGE_PATH' in os.environ:
+    PACKAGES_TO_INSTALL.append(os.environ['KFP_PIPELINE_SPEC_PACKAGE_PATH'])
 
-@dsl.component(packages_to_install=['dill==0.3.7'])
+@dsl.component(packages_to_install=PACKAGES_TO_INSTALL)
 def make_language_model(text_dataset: Dataset) -> Model:
     # dill allows pickling objects belonging to a function's local namespace
     import dill