Repository: getindata/kedro-kubeflow
Pull request #98 (2022-02-02 12:33:29)

Support exit pipeline in full merge strategy
#### Description

Exit pipeline handler is added also in case of "full" merge strategy

##### PR Checklist
- [x] Tests added
- [x] [Changelog](CHANGELOG.md) updated 


em-pe: > @em-pe / @mariusz89016 Code climate strikes again, this time even with false-positive about `nullcontext`. What would you say to disable this integration?

I would stick to code climate and just intentionally ignore the failures in such case and merge anyway. You should be able to merge with your permission level - let me know if that's not the case. 


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 80f8f83..92cca43 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -3,6 +3,7 @@
 ## [Unreleased]
 
 -   Kedro pipeline name is now added into Kubeflow pipeline name during upload
+-   Added missing on-exit-handler for `node_merge_strategy: full`
 
 ## [0.5.1] - 2022-01-28
 
diff --git a/kedro_kubeflow/generators/one_pod_pipeline_generator.py b/kedro_kubeflow/generators/one_pod_pipeline_generator.py
index 0800f45..bc7a126 100644
--- a/kedro_kubeflow/generators/one_pod_pipeline_generator.py
+++ b/kedro_kubeflow/generators/one_pod_pipeline_generator.py
@@ -8,6 +8,7 @@ from .utils import (
     create_arguments_from_parameters,
     create_command_using_params_dumper,
     create_container_environment,
+    create_pipeline_exit_handler,
     maybe_add_params,
 )
 
@@ -29,7 +30,14 @@ class OnePodPipelineGenerator(object):
             dsl.get_pipeline_conf().set_ttl_seconds_after_finished(
                 self.run_config.ttl
             )
-            self._build_kfp_op(pipeline, image, image_pull_policy)
+            with create_pipeline_exit_handler(
+                pipeline,
+                image,
+                image_pull_policy,
+                self.run_config,
+                self.context,
+            ):
+                self._build_kfp_op(pipeline, image, image_pull_policy)
 
         return convert_kedro_pipeline_to_kfp
 
diff --git a/kedro_kubeflow/generators/pod_per_node_pipeline_generator.py b/kedro_kubeflow/generators/pod_per_node_pipeline_generator.py
index 0cc75e8..4bdc910 100644
--- a/kedro_kubeflow/generators/pod_per_node_pipeline_generator.py
+++ b/kedro_kubeflow/generators/pod_per_node_pipeline_generator.py
@@ -1,17 +1,17 @@
-import contextlib
 import logging
 from typing import Dict, Set
 
 import kubernetes.client as k8s
 from kedro.pipeline.node import Node
 from kfp import dsl
-from kfp.compiler._k8s_helper import sanitize_k8s_name
 
 from ..utils import clean_name, is_mlflow_enabled
 from .utils import (
     create_arguments_from_parameters,
     create_command_using_params_dumper,
     create_container_environment,
+    create_pipeline_exit_handler,
+    customize_op,
     maybe_add_params,
 )
 
@@ -47,8 +47,12 @@ class PodPerNodePipelineGenerator(object):
             node_dependencies = self.context.pipelines.get(
                 pipeline
             ).node_dependencies
-            with self._create_pipeline_exit_handler(
-                pipeline, image, image_pull_policy
+            with create_pipeline_exit_handler(
+                pipeline,
+                image,
+                image_pull_policy,
+                self.run_config,
+                self.context,
             ):
                 kfp_ops = self._build_kfp_ops(
                     pipeline, node_dependencies, image, image_pull_policy
@@ -61,59 +65,6 @@ class PodPerNodePipelineGenerator(object):
 
         return convert_kedro_pipeline_to_kfp
 
-    def _create_pipeline_exit_handler(
-        self, pipeline, image, image_pull_policy
-    ):
-        enable_volume_cleaning = (
-            self.run_config.volume is not None
-            and not self.run_config.volume.keep
-        )
-
-        if not enable_volume_cleaning and not self.run_config.on_exit_pipeline:
-            return contextlib.nullcontext()
-
-        commands = []
-
-        if enable_volume_cleaning:
-            commands.append(
-                "kedro kubeflow delete-pipeline-volume "
-                "{{workflow.name}}-"
-                + sanitize_k8s_name(f"{pipeline}-data-volume")
-            )
-
-        if self.run_config.on_exit_pipeline:
-            commands.append(
-                "kedro run "
-                "--config config.yaml "
-                f"--env {self.context.env} "
-                f"--pipeline {self.run_config.on_exit_pipeline}"
-            )
-
-        exit_container_op = dsl.ContainerOp(
-            name="on-exit",
-            image=image,
-            command=create_command_using_params_dumper(";".join(commands)),
-            arguments=create_arguments_from_parameters(
-                self.context.params.keys()
-            )
-            + [
-                "status",
-                "{{workflow.status}}",
-                "failures",
-                "{{workflow.failures}}",
-            ],
-            container_kwargs={"env": create_container_environment()},
-        )
-
-        if self.run_config.max_cache_staleness not in [None, ""]:
-            exit_container_op.execution_options.caching_strategy.max_cache_staleness = (
-                self.run_config.max_cache_staleness
-            )
-
-        return dsl.ExitHandler(
-            self._customize_op(exit_container_op, image_pull_policy)
-        )
-
     def _build_kfp_ops(
         self,
         pipeline,
@@ -135,7 +86,7 @@ class PodPerNodePipelineGenerator(object):
         nodes_env = create_container_environment()
 
         if is_mlflow_enabled():
-            kfp_ops["mlflow-start-run"] = self._customize_op(
+            kfp_ops["mlflow-start-run"] = customize_op(
                 dsl.ContainerOp(
                     name="mlflow-start-run",
                     image=image,
@@ -151,6 +102,7 @@ class PodPerNodePipelineGenerator(object):
                     file_outputs={"mlflow_run_id": "/tmp/mlflow_run_id"},
                 ),
                 image_pull_policy,
+                self.run_config,
             )
 
             nodes_env.append(
@@ -169,7 +121,7 @@ class PodPerNodePipelineGenerator(object):
                     requests=self.run_config.resources.get_for(node.name),
                 )
 
-            kfp_ops[node.name] = self._customize_op(
+            kfp_ops[node.name] = customize_op(
                 dsl.ContainerOp(
                     name=name,
                     image=image,
@@ -196,18 +148,11 @@ class PodPerNodePipelineGenerator(object):
                     },
                 ),
                 image_pull_policy,
+                self.run_config,
             )
 
         return kfp_ops
 
-    def _customize_op(self, op, image_pull_policy):
-        op.container.set_image_pull_policy(image_pull_policy)
-        if self.run_config.volume and self.run_config.volume.owner is not None:
-            op.container.set_security_context(
-                k8s.V1SecurityContext(run_as_user=self.run_config.volume.owner)
-            )
-        return op
-
     def _setup_volumes(self, volume_name, image, image_pull_policy):
         vop = dsl.VolumeOp(
             name="data-volume-create",
@@ -226,7 +171,7 @@ class PodPerNodePipelineGenerator(object):
         if self.run_config.volume.skip_init:
             return {"/home/kedro/data": vop.volume}
         else:
-            volume_init = self._customize_op(
+            volume_init = customize_op(
                 dsl.ContainerOp(
                     name="data-volume-init",
                     image=image,
@@ -245,5 +190,6 @@ class PodPerNodePipelineGenerator(object):
                     pvolumes={"/home/kedro/datavolume": vop.volume},
                 ),
                 image_pull_policy,
+                self.run_config,
             )
             return {"/home/kedro/data": volume_init.pvolume}
diff --git a/kedro_kubeflow/generators/utils.py b/kedro_kubeflow/generators/utils.py
index 289ed95..54b1e7f 100644
--- a/kedro_kubeflow/generators/utils.py
+++ b/kedro_kubeflow/generators/utils.py
@@ -1,3 +1,4 @@
+import contextlib
 import itertools
 import os
 from functools import wraps
@@ -5,6 +6,7 @@ from inspect import Parameter, signature
 
 import kubernetes.client as k8s
 from kfp import dsl
+from kfp.compiler._k8s_helper import sanitize_k8s_name
 
 from ..auth import IAP_CLIENT_ID
 
@@ -58,3 +60,62 @@ def create_arguments_from_parameters(paramter_names):
             *[[param, dsl.PipelineParam(param)] for param in paramter_names]
         )
     )
+
+
+def create_pipeline_exit_handler(
+    pipeline, image, image_pull_policy, run_config, context
+):
+    enable_volume_cleaning = (
+        run_config.volume is not None and not run_config.volume.keep
+    )
+
+    if not enable_volume_cleaning and not run_config.on_exit_pipeline:
+        return contextlib.nullcontext()
+
+    commands = []
+
+    if enable_volume_cleaning:
+        commands.append(
+            "kedro kubeflow delete-pipeline-volume "
+            "{{workflow.name}}-" + sanitize_k8s_name(f"{pipeline}-data-volume")
+        )
+
+    if run_config.on_exit_pipeline:
+        commands.append(
+            "kedro run "
+            "--config config.yaml "
+            f"--env {context.env} "
+            f"--pipeline {run_config.on_exit_pipeline}"
+        )
+
+    exit_container_op = dsl.ContainerOp(
+        name="on-exit",
+        image=image,
+        command=create_command_using_params_dumper(";".join(commands)),
+        arguments=create_arguments_from_parameters(context.params.keys())
+        + [
+            "status",
+            "{{workflow.status}}",
+            "failures",
+            "{{workflow.failures}}",
+        ],
+        container_kwargs={"env": create_container_environment()},
+    )
+
+    if run_config.max_cache_staleness not in [None, ""]:
+        exit_container_op.execution_options.caching_strategy.max_cache_staleness = (
+            run_config.max_cache_staleness
+        )
+
+    return dsl.ExitHandler(
+        customize_op(exit_container_op, image_pull_policy, run_config)
+    )
+
+
+def customize_op(op, image_pull_policy, run_config):
+    op.container.set_image_pull_policy(image_pull_policy)
+    if run_config.volume and run_config.volume.owner is not None:
+        op.container.set_security_context(
+            k8s.V1SecurityContext(run_as_user=run_config.volume.owner)
+        )
+    return op


Test patch:
diff --git a/tests/test_one_pod_pipeline_generator.py b/tests/test_one_pod_pipeline_generator.py
index af8ae78..8e5d456 100644
--- a/tests/test_one_pod_pipeline_generator.py
+++ b/tests/test_one_pod_pipeline_generator.py
@@ -191,6 +191,29 @@ class TestGenerator(unittest.TestCase):
         assert "KUBEFLOW_RUN_ID" in env_values
         assert env_values["KUBEFLOW_RUN_ID"] == "{{workflow.uid}}"
 
+    def test_should_generate_exit_handler_if_requested(self):
+        # given
+        self.create_generator(config={"on_exit_pipeline": "notify_via_slack"})
+
+        # when
+        with kfp.dsl.Pipeline(None) as dsl_pipeline:
+            pipeline = self.generator_under_test.generate_pipeline(
+                "pipeline", "unittest-image", "Always"
+            )
+            pipeline()
+
+        # then
+        assert len(dsl_pipeline.ops) == 2
+        assert "on-exit" in dsl_pipeline.ops
+        assert (
+            dsl_pipeline.ops["on-exit"]
+            .container.command[-1]
+            .endswith(
+                "kedro run --config config.yaml "
+                "--env unittests --pipeline notify_via_slack"
+            )
+        )
+
     def create_generator(self, config=None, params=None, catalog=None):
         if config is None:
             config = {}