Repository: getindata/kedro-kubeflow
Pull request #80 (2021-12-22 19:38:19)

Release 0.4.5
Bump version and CHANGELOG for next release.


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7f1be5d..28f2215 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,10 @@
 
 ## [Unreleased]
 
+## [0.4.5] - 2021-12-22
+
+-   Add `node_merge_strategy` alongside with `full` option to run a whole pipeline in one pod
+
 ## [0.4.4] - 2021-09-29
 
 -   Custom networking setup for Vertex AI pipelines run
@@ -88,7 +92,9 @@
 -   Method to schedule runs for most recent version of given pipeline `kedro kubeflow schedule` 
 -   Shortcut to open UI for pipelines using `kedro kubeflow ui` 
 
-[Unreleased]: https://github.com/getindata/kedro-kubeflow/compare/0.4.4...HEAD
+[Unreleased]: https://github.com/getindata/kedro-kubeflow/compare/0.4.5...HEAD
+
+[0.4.5]: https://github.com/getindata/kedro-kubeflow/compare/0.4.4...0.4.5
 
 [0.4.4]: https://github.com/getindata/kedro-kubeflow/compare/0.4.3...0.4.4
 
diff --git a/kedro_kubeflow/__init__.py b/kedro_kubeflow/__init__.py
index a6f375b..fe52fc1 100644
--- a/kedro_kubeflow/__init__.py
+++ b/kedro_kubeflow/__init__.py
@@ -1,3 +1,3 @@
 """kedro_kubeflow."""
 
-version = "0.4.4"
+version = "0.4.5"
diff --git a/kedro_kubeflow/config.py b/kedro_kubeflow/config.py
index a0418df..dd09369 100644
--- a/kedro_kubeflow/config.py
+++ b/kedro_kubeflow/config.py
@@ -45,6 +45,20 @@ run_config:
   # intermediate results in the MLMD
   #store_kedro_outputs_as_kfp_artifacts: True
 
+  # Strategy used to generate Kubeflow pipeline nodes from Kedro nodes
+  # Available strategies:
+  #  * none (default) - nodes in Kedro pipeline are mapped to separate nodes
+  #                     in Kubeflow pipelines. This strategy allows to inspect
+  #                     a whole processing graph in Kubeflow UI and override
+  #                     resources for each node (because they are run in separate pods)
+  #                     Although, performance may not be optimal due to potential
+  #                     sharing of intermediate datasets through disk.
+  #  * full - nodes in Kedro pipeline are mapped to one node in Kubeflow pipelines.
+  #           This strategy mitigate potential performance issues with `none` strategy
+  #           but at the cost of degraded user experience within Kubeflow UI: a graph
+  #           is collapsed to one node.
+  #node_merge_strategy: none
+
   # Optional volume specification (only for non vertex-ai)
   volume:
 
@@ -231,6 +245,16 @@ class RunConfig(Config):
             self._get_or_default("vertex_ai_networking", {})
         )
 
+    @property
+    def node_merge_strategy(self):
+        strategy = str(self._get_or_default("node_merge_strategy", "none"))
+        if strategy not in ["none", "full"]:
+            raise ValueError(
+                f"Invalid {self._get_prefix()}node_merge_strategy: {strategy}"
+            )
+        else:
+            return strategy
+
     def _get_prefix(self):
         return "run_config."
 
diff --git a/kedro_kubeflow/context_helper.py b/kedro_kubeflow/context_helper.py
index c010836..a1fbf32 100644
--- a/kedro_kubeflow/context_helper.py
+++ b/kedro_kubeflow/context_helper.py
@@ -11,11 +11,7 @@ class ContextHelper(object):
 
     CONFIG_FILE_PATTERN = "kubeflow*"
 
-    def __init__(
-        self,
-        metadata,
-        env,
-    ):
+    def __init__(self, metadata, env):
         self._metadata = metadata
         self._env = env
 
@@ -59,21 +55,12 @@ class ContextHelper(object):
             )
 
     @staticmethod
-    def init(
-        metadata,
-        env,
-    ):
+    def init(metadata, env):
         version = VersionInfo.parse(kedro_version)
         if version.match(">=0.17.0"):
-            return ContextHelper(
-                metadata,
-                env,
-            )
+            return ContextHelper(metadata, env)
         else:
-            return ContextHelper16(
-                metadata,
-                env,
-            )
+            return ContextHelper16(metadata, env)
 
 
 class ContextHelper16(ContextHelper):
diff --git a/kedro_kubeflow/generators/__init__.py b/kedro_kubeflow/generators/__init__.py
new file mode 100644
index 0000000..23ac6ff
--- /dev/null
+++ b/kedro_kubeflow/generators/__init__.py
@@ -0,0 +1,1 @@
+"""kedro_kubeflow.generators"""
diff --git a/kedro_kubeflow/generators/one_pod_pipeline_generator.py b/kedro_kubeflow/generators/one_pod_pipeline_generator.py
new file mode 100644
index 0000000..7077fdd
--- /dev/null
+++ b/kedro_kubeflow/generators/one_pod_pipeline_generator.py
@@ -0,0 +1,78 @@
+import logging
+import os
+
+import kubernetes.client as k8s
+from kfp import dsl
+
+from ..auth import IAP_CLIENT_ID
+from ..utils import clean_name
+from .utils import create_params, maybe_add_params
+
+
+class OnePodPipelineGenerator(object):
+    log = logging.getLogger(__name__)
+
+    def __init__(self, config, project_name, context):
+        self.project_name = project_name
+        self.context = context
+        dsl.ContainerOp._DISABLE_REUSABLE_COMPONENT_WARNING = True
+        self.run_config = config.run_config
+        self.catalog = context.config_loader.get("catalog*")
+
+    def generate_pipeline(self, pipeline, image, image_pull_policy):
+        @dsl.pipeline(self.project_name, self.run_config.description)
+        @maybe_add_params(self.context.params)
+        def convert_kedro_pipeline_to_kfp() -> None:
+            dsl.get_pipeline_conf().set_ttl_seconds_after_finished(
+                self.run_config.ttl
+            )
+            self._build_kfp_op(pipeline, image, image_pull_policy)
+
+        return convert_kedro_pipeline_to_kfp
+
+    def _build_kfp_op(
+        self,
+        pipeline,
+        image,
+        image_pull_policy,
+    ) -> dsl.ContainerOp:
+        kwargs = {
+            "env": [
+                k8s.V1EnvVar(
+                    name=IAP_CLIENT_ID, value=os.environ.get(IAP_CLIENT_ID, "")
+                )
+            ],
+            "image_pull_policy": image_pull_policy,
+        }
+        default_resources = self.run_config.resources.get_for("__default__")
+        if default_resources:
+            kwargs["resources"] = k8s.V1ResourceRequirements(
+                limits=default_resources, requests=default_resources
+            )
+        container_op = dsl.ContainerOp(
+            name=clean_name(pipeline),
+            image=image,
+            command=["kedro"],
+            arguments=[
+                "run",
+                "--env",
+                self.context.env,
+                "--params",
+                create_params(self.context.params.keys()),
+                "--pipeline",
+                pipeline,
+            ],
+            container_kwargs=kwargs,
+            file_outputs={
+                output: f"/home/kedro/{self.catalog[output]['filepath']}"
+                for output in self.catalog
+                if "filepath" in self.catalog[output]
+                and self.run_config.store_kedro_outputs_as_kfp_artifacts
+            },
+        )
+
+        container_op.execution_options.caching_strategy.max_cache_staleness = (
+            self.run_config.max_cache_staleness
+        )
+
+        return container_op
diff --git a/kedro_kubeflow/generator.py b/kedro_kubeflow/generators/pod_per_node_pipeline_generator.py
similarity index 90%
rename from kedro_kubeflow/generator.py
rename to kedro_kubeflow/generators/pod_per_node_pipeline_generator.py
index dc52862..a80eac9 100644
--- a/kedro_kubeflow/generator.py
+++ b/kedro_kubeflow/generators/pod_per_node_pipeline_generator.py
@@ -1,8 +1,6 @@
 import contextlib
 import logging
 import os
-from functools import wraps
-from inspect import Parameter, signature
 from typing import Dict, Set
 
 import kubernetes.client as k8s
@@ -10,29 +8,12 @@ from kedro.pipeline.node import Node
 from kfp import dsl
 from kfp.compiler._k8s_helper import sanitize_k8s_name
 
-from .auth import IAP_CLIENT_ID
-from .utils import clean_name, is_mlflow_enabled
+from ..auth import IAP_CLIENT_ID
+from ..utils import clean_name, is_mlflow_enabled
+from .utils import create_params, maybe_add_params
 
 
-def maybe_add_params(kedro_parameters):
-    def decorator(f):
-        @wraps(f)
-        def wrapper(*args, **kwargs):
-            return f()
-
-        sig = signature(f)
-        new_params = (
-            Parameter(name, Parameter.KEYWORD_ONLY, default=default)
-            for name, default in kedro_parameters.items()
-        )
-        wrapper.__signature__ = sig.replace(parameters=new_params)
-        return wrapper
-
-    return decorator
-
-
-class PipelineGenerator(object):
-
+class PodPerNodePipelineGenerator(object):
     log = logging.getLogger(__name__)
 
     def __init__(self, config, project_name, context):
@@ -158,12 +139,6 @@ class PipelineGenerator(object):
 
         for node in node_dependencies:
             name = clean_name(node.name)
-            params = ",".join(
-                [
-                    f"{param}:{dsl.PipelineParam(param)}"
-                    for param in self.context.params.keys()
-                ]
-            )
             kwargs = {"env": nodes_env}
             if self.run_config.resources.is_set_for(node.name):
                 kwargs["resources"] = k8s.V1ResourceRequirements(
@@ -181,7 +156,7 @@ class PipelineGenerator(object):
                         "--env",
                         self.context.env,
                         "--params",
-                        params,
+                        create_params(self.context.params.keys()),
                         "--pipeline",
                         pipeline,
                         "--node",
diff --git a/kedro_kubeflow/generators/utils.py b/kedro_kubeflow/generators/utils.py
new file mode 100644
index 0000000..296f674
--- /dev/null
+++ b/kedro_kubeflow/generators/utils.py
@@ -0,0 +1,28 @@
+from functools import wraps
+from inspect import Parameter, signature
+from typing import Iterable
+
+from kfp import dsl
+
+
+def maybe_add_params(kedro_parameters):
+    def decorator(f):
+        @wraps(f)
+        def wrapper(*args, **kwargs):
+            return f()
+
+        sig = signature(f)
+        new_params = (
+            Parameter(name, Parameter.KEYWORD_ONLY, default=default)
+            for name, default in kedro_parameters.items()
+        )
+        wrapper.__signature__ = sig.replace(parameters=new_params)
+        return wrapper
+
+    return decorator
+
+
+def create_params(param_keys: Iterable[str]) -> str:
+    return ",".join(
+        [f"{param}:{dsl.PipelineParam(param)}" for param in param_keys]
+    )
diff --git a/kedro_kubeflow/kfpclient.py b/kedro_kubeflow/kfpclient.py
index 0481831..231123a 100644
--- a/kedro_kubeflow/kfpclient.py
+++ b/kedro_kubeflow/kfpclient.py
@@ -7,8 +7,14 @@ from kfp import Client
 from kfp.compiler import Compiler
 from tabulate import tabulate
 
+from kedro_kubeflow.generators.one_pod_pipeline_generator import (
+    OnePodPipelineGenerator,
+)
+from kedro_kubeflow.generators.pod_per_node_pipeline_generator import (
+    PodPerNodePipelineGenerator,
+)
+
 from .auth import AuthHandler
-from .generator import PipelineGenerator
 from .utils import clean_name
 
 WAIT_TIMEOUT = 24 * 60 * 60
@@ -35,7 +41,18 @@ class KubeflowClient(object):
 
         self.project_name = project_name
         self.pipeline_description = config.run_config.description
-        self.generator = PipelineGenerator(config, project_name, context)
+        if config.run_config.node_merge_strategy == "none":
+            self.generator = PodPerNodePipelineGenerator(
+                config, project_name, context
+            )
+        elif config.run_config.node_merge_strategy == "full":
+            self.generator = OnePodPipelineGenerator(
+                config, project_name, context
+            )
+        else:
+            raise Exception(
+                f"Invalid `node_merge_strategy`: {config.run_config.node_merge_strategy}"
+            )
 
     def list_pipelines(self):
         pipelines = self.client.list_pipelines(page_size=30).pipelines
diff --git a/setup.cfg b/setup.cfg
index 97b3e8e..4806d8f 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 0.4.4
+current_version = 0.4.5
 
 [bumpversion:file:setup.py]
 
diff --git a/setup.py b/setup.py
index b879989..aab0c38 100644
--- a/setup.py
+++ b/setup.py
@@ -37,7 +37,7 @@ EXTRA_REQUIRE = {
 
 setup(
     name="kedro-kubeflow",
-    version="0.4.4",
+    version="0.4.5",
     description="Kedro plugin with Kubeflow support",
     long_description=README,
     long_description_content_type="text/markdown",


Test patch:
diff --git a/tests/test_generator.py b/tests/test_generator.py
index 92cf41a..d9ddcca 100644
--- a/tests/test_generator.py
+++ b/tests/test_generator.py
@@ -8,7 +8,9 @@ import kfp
 from kedro.pipeline import Pipeline, node
 
 from kedro_kubeflow.config import PluginConfig
-from kedro_kubeflow.generator import PipelineGenerator
+from kedro_kubeflow.generators.pod_per_node_pipeline_generator import (
+    PodPerNodePipelineGenerator,
+)
 
 
 def identity(input1: str):
@@ -16,14 +18,6 @@ def identity(input1: str):
 
 
 class TestGenerator(unittest.TestCase):
-    def create_pipeline(self):
-        return Pipeline(
-            [
-                node(identity, "A", "B", name="node1"),
-                node(identity, "B", "C", name="node2"),
-            ]
-        )
-
     def test_support_modification_of_pull_policy(self):
         # given
         self.create_generator()
@@ -352,7 +346,7 @@ class TestGenerator(unittest.TestCase):
                 },
             },
         )
-        self.generator_under_test = PipelineGenerator(
+        self.generator_under_test = PodPerNodePipelineGenerator(
             PluginConfig({"host": "http://unittest", "run_config": config}),
             project_name,
             context,
diff --git a/tests/test_kfpclient.py b/tests/test_kfpclient.py
index 4890831..cdb9db8 100644
--- a/tests/test_kfpclient.py
+++ b/tests/test_kfpclient.py
@@ -136,7 +136,7 @@ class TestKubeflowClient(unittest.TestCase):
                 assert "generateName: my-awesome-project-" in yamlfile.read()
 
     @patch("kedro_kubeflow.kfpclient.AuthHandler")
-    @patch("kedro_kubeflow.kfpclient.PipelineGenerator")
+    @patch("kedro_kubeflow.kfpclient.PodPerNodePipelineGenerator")
     @patch("kedro_kubeflow.kfpclient.Client")
     def test_should_use_jwt_token_in_kfp_client(
         self, kfp_client_mock, pipeline_generator_mock, auth_handler_mock
@@ -163,7 +163,7 @@ class TestKubeflowClient(unittest.TestCase):
         )
 
     @patch("kedro_kubeflow.kfpclient.AuthHandler")
-    @patch("kedro_kubeflow.kfpclient.PipelineGenerator")
+    @patch("kedro_kubeflow.kfpclient.PodPerNodePipelineGenerator")
     @patch("kedro_kubeflow.kfpclient.Client")
     def test_should_use_dex_session_in_kfp_client(
         self, kfp_client_mock, pipeline_generator_mock, auth_handler_mock
@@ -321,7 +321,23 @@ class TestKubeflowClient(unittest.TestCase):
         self.kfp_client_mock.pipeline_uploads.upload_pipeline.assert_not_called()
         self.kfp_client_mock.pipeline_uploads.upload_pipeline_version.assert_called()
 
-    @patch("kedro_kubeflow.kfpclient.PipelineGenerator")
+    @patch("kedro_kubeflow.kfpclient.Client")
+    def test_should_raise_error_if_invalid_node_merge_strategy(
+        self, kfp_client_mock
+    ):
+        with self.assertRaises(ValueError):
+            KubeflowClient(
+                PluginConfig(
+                    {
+                        "host": "http://unittest",
+                        "run_config": {"node_merge_strategy": "other"},
+                    }
+                ),
+                None,
+                None,
+            )
+
+    @patch("kedro_kubeflow.kfpclient.PodPerNodePipelineGenerator")
     @patch("kedro_kubeflow.kfpclient.Client")
     def create_client(self, config, kfp_client_mock, pipeline_generator_mock):
         project_name = "my-awesome-project"
diff --git a/tests/test_one_pod_pipeline_generator.py b/tests/test_one_pod_pipeline_generator.py
new file mode 100644
index 0000000..96e6159
--- /dev/null
+++ b/tests/test_one_pod_pipeline_generator.py
@@ -0,0 +1,187 @@
+"""Test generator"""
+
+import unittest
+from inspect import signature
+from unittest.mock import MagicMock
+
+import kfp
+from kedro.pipeline import Pipeline, node
+
+from kedro_kubeflow.config import PluginConfig
+from kedro_kubeflow.generators.one_pod_pipeline_generator import (
+    OnePodPipelineGenerator,
+)
+
+
+def identity(input1: str):
+    return input1  # pragma: no cover
+
+
+class TestGenerator(unittest.TestCase):
+    def test_support_modification_of_pull_policy(self):
+        # given
+        self.create_generator()
+
+        # when
+        with kfp.dsl.Pipeline(None) as dsl_pipeline:
+            self.generator_under_test.generate_pipeline(
+                "pipeline", "unittest-image", "Never"
+            )()
+
+        # then
+        assert len(dsl_pipeline.ops) == 1
+        assert dsl_pipeline.ops["pipeline"].container.image == "unittest-image"
+        assert (
+            dsl_pipeline.ops["pipeline"].container.image_pull_policy == "Never"
+        )
+
+    def test_should_support_params_and_inject_them_to_the_node(self):
+        # given
+        self.create_generator(params={"param1": 0.3, "param2": 42})
+
+        # when
+        with kfp.dsl.Pipeline(None) as dsl_pipeline:
+            pipeline = self.generator_under_test.generate_pipeline(
+                "pipeline", "unittest-image", "Always"
+            )
+            default_params = signature(pipeline).parameters
+            pipeline()
+
+        # then
+        assert len(default_params) == 2
+        assert default_params["param1"].default == 0.3
+        assert default_params["param2"].default == 42
+        assert dsl_pipeline.ops["pipeline"].container.args == [
+            "run",
+            "--env",
+            "unittests",
+            "--params",
+            "param1:{{pipelineparam:op=;name=param1}},"
+            "param2:{{pipelineparam:op=;name=param2}}",
+            "--pipeline",
+            "pipeline",
+        ]
+
+    def test_should_not_add_resources_spec_if_not_requested(self):
+        # given
+        self.create_generator(config={})
+
+        # when
+        with kfp.dsl.Pipeline(None) as dsl_pipeline:
+            self.generator_under_test.generate_pipeline(
+                "pipeline", "unittest-image", "Always"
+            )()
+
+        # then
+        assert dsl_pipeline.ops["pipeline"].container.resources is None
+
+    def test_should_add_resources_spec(self):
+        # given
+        self.create_generator(
+            config={
+                "resources": {
+                    "__default__": {"cpu": "100m", "memory": "8Gi"},
+                    "node1": {"cpu": "400m", "memory": "64Gi"},
+                }
+            }
+        )
+
+        # when
+        with kfp.dsl.Pipeline(None) as dsl_pipeline:
+            self.generator_under_test.generate_pipeline(
+                "pipeline", "unittest-image", "Always"
+            )()
+
+        # then
+        resources = dsl_pipeline.ops["pipeline"].container.resources
+        assert resources.limits == {"cpu": "100m", "memory": "8Gi"}
+        assert resources.requests == {"cpu": "100m", "memory": "8Gi"}
+
+    def test_should_set_description(self):
+        # given
+        self.create_generator(config={"description": "DESC"})
+
+        # when
+        pipeline = self.generator_under_test.generate_pipeline(
+            "pipeline", "unittest-image", "Never"
+        )
+
+        # then
+        assert pipeline._component_description == "DESC"
+
+    def test_artifact_registration(self):
+        # given
+        self.create_generator(
+            catalog={
+                "B": {
+                    "type": "pandas.CSVDataSet",
+                    "filepath": "data/02_intermediate/b.csv",
+                }
+            }
+        )
+
+        # when
+        with kfp.dsl.Pipeline(None) as dsl_pipeline:
+            self.generator_under_test.generate_pipeline(
+                "pipeline", "unittest-image", "Always"
+            )()
+
+        # then
+        assert dsl_pipeline.ops["pipeline"].file_outputs == {
+            "B": "/home/kedro/data/02_intermediate/b.csv"
+        }
+
+    def test_should_skip_artifact_registration_if_requested(self):
+        # given
+        self.create_generator(
+            catalog={
+                "B": {
+                    "type": "pandas.CSVDataSet",
+                    "filepath": "data/02_intermediate/b.csv",
+                }
+            },
+            config={"store_kedro_outputs_as_kfp_artifacts": False},
+        )
+
+        # when
+        with kfp.dsl.Pipeline(None) as dsl_pipeline:
+            self.generator_under_test.generate_pipeline(
+                "pipeline", "unittest-image", "Always"
+            )()
+
+        # then
+        assert dsl_pipeline.ops["pipeline"].file_outputs == {}
+
+    def create_generator(self, config=None, params=None, catalog=None):
+        if config is None:
+            config = {}
+        if params is None:
+            params = {}
+        if catalog is None:
+            catalog = {}
+        config_loader = MagicMock()
+        config_loader.get.return_value = catalog
+        context = type(
+            "obj",
+            (object,),
+            {
+                "env": "unittests",
+                "params": params,
+                "config_loader": config_loader,
+                "pipelines": {
+                    "pipeline": Pipeline(
+                        [
+                            node(identity, "A", "B", name="node1"),
+                            node(identity, "B", "C", name="node2"),
+                        ]
+                    )
+                },
+            },
+        )
+        self.generator_under_test = OnePodPipelineGenerator(
+            config=PluginConfig(
+                {"host": "http://unittest", "run_config": config}
+            ),
+            project_name="my-awesome-project",
+            context=context,
+        )