Repository: getindata/kedro-kubeflow (python)

Issue:
Add support for additional custom exit handler
In order to send a notification when the pipeline ends we need ability to run the custom code as an exit handler.

Ideally, there should be new config param "exit_handler_pipeline" and when it's set, an exit handler calls the requested kedro pipeline, passing `status` and `failures` as params (see: https://github.com/argoproj/argo-workflows/blob/master/docs/variables.md#exit-handler)

Pull request:
Support custom exit handler
#### Description

Provides ability to run custom Kedro pipeline on run success/failure.
Resolves #91 

##### PR Checklist
- [x] Tests added
- [x] [Changelog](CHANGELOG.md) updated 


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2d48427..fd5e387 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,8 @@
 
 ## [Unreleased]
 
+-   Possibility to run custom Kedro pipeline as on-exit-handler
+
 ## [0.5.0] - 2022-01-27
 
 -   Kedro paramters of complex types (lists and dicts) are now supported
diff --git a/kedro_kubeflow/cli.py b/kedro_kubeflow/cli.py
index 8a8ec26..1c0d961 100644
--- a/kedro_kubeflow/cli.py
+++ b/kedro_kubeflow/cli.py
@@ -300,3 +300,21 @@ def mlflow_start(ctx, kubeflow_run_id: str, output: str):
     with open(output, "w") as f:
         f.write(run.info.run_id)
     click.echo(f"Started run: {run.info.run_id}")
+
+
+@kubeflow_group.command(hidden=True)
+@click.argument("pvc_name", type=str)
+def delete_pipeline_volume(pvc_name: str):
+    import kubernetes.client
+    import kubernetes.config
+
+    kubernetes.config.load_incluster_config()
+    current_namespace = open(
+        "/var/run/secrets/kubernetes.io/serviceaccount/namespace"
+    ).read()
+
+    kubernetes.client.CoreV1Api().delete_namespaced_persistent_volume_claim(
+        pvc_name,
+        current_namespace,
+    )
+    click.echo(f"Volume removed: {pvc_name}")
diff --git a/kedro_kubeflow/config.py b/kedro_kubeflow/config.py
index ac5bcf3..85811bd 100644
--- a/kedro_kubeflow/config.py
+++ b/kedro_kubeflow/config.py
@@ -38,6 +38,10 @@ run_config:
   # volume after pipeline finishes) [in seconds]. Default: 1 week
   ttl: 604800
 
+  # What Kedro pipeline should be run as the last step regardless of the
+  # pipeline status. Used to send notifications or raise the alerts
+  # on_exit_pipeline: notify_via_slack
+
   # This sets the caching option for pipeline using
   # execution_options.caching_strategy.max_cache_staleness
   # See https://en.wikipedia.org/wiki/ISO_8601 in section 'Duration'
@@ -248,6 +252,10 @@ class RunConfig(Config):
     def ttl(self):
         return int(self._get_or_default("ttl", 3600 * 24 * 7))
 
+    @property
+    def on_exit_pipeline(self):
+        return self._get_or_default("on_exit_pipeline", None)
+
     @property
     def vertex_ai_networking(self):
         return VertexAiNetworkingConfig(
diff --git a/kedro_kubeflow/generators/pod_per_node_pipeline_generator.py b/kedro_kubeflow/generators/pod_per_node_pipeline_generator.py
index 9099fa9..0cc75e8 100644
--- a/kedro_kubeflow/generators/pod_per_node_pipeline_generator.py
+++ b/kedro_kubeflow/generators/pod_per_node_pipeline_generator.py
@@ -47,7 +47,9 @@ class PodPerNodePipelineGenerator(object):
             node_dependencies = self.context.pipelines.get(
                 pipeline
             ).node_dependencies
-            with self._create_pipeline_exit_handler(pipeline):
+            with self._create_pipeline_exit_handler(
+                pipeline, image, image_pull_policy
+            ):
                 kfp_ops = self._build_kfp_ops(
                     pipeline, node_dependencies, image, image_pull_policy
                 )
@@ -59,29 +61,48 @@ class PodPerNodePipelineGenerator(object):
 
         return convert_kedro_pipeline_to_kfp
 
-    def _create_pipeline_exit_handler(self, pipeline):
+    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:
+        if not enable_volume_cleaning and not self.run_config.on_exit_pipeline:
             return contextlib.nullcontext()
 
-        exit_container_op = dsl.ContainerOp(
-            name="schedule-volume-termination",
-            image="gcr.io/cloud-builders/kubectl",
-            command=[
-                "kubectl",
-                "delete",
-                "pvc",
+        commands = []
+
+        if enable_volume_cleaning:
+            commands.append(
+                "kedro kubeflow delete-pipeline-volume "
                 "{{workflow.name}}-"
-                + sanitize_k8s_name(f"{pipeline}-data-volume"),
-                "--wait=false",
-                "--ignore-not-found",
-                "--output",
-                "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, ""]:
@@ -89,7 +110,9 @@ class PodPerNodePipelineGenerator(object):
                 self.run_config.max_cache_staleness
             )
 
-        return dsl.ExitHandler(exit_container_op)
+        return dsl.ExitHandler(
+            self._customize_op(exit_container_op, image_pull_policy)
+        )
 
     def _build_kfp_ops(
         self,


Test patch:
diff --git a/tests/test_cli.py b/tests/test_cli.py
index b7c89a5..32ca2b9 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -1,5 +1,6 @@
 import os
 import unittest
+import unittest.mock as um
 from collections import namedtuple
 from pathlib import Path
 from tempfile import TemporaryDirectory
@@ -9,6 +10,7 @@ from click.testing import CliRunner
 
 from kedro_kubeflow.cli import (
     compile,
+    delete_pipeline_volume,
     init,
     list_pipelines,
     mlflow_start,
@@ -237,3 +239,20 @@ class TestPluginCLI(unittest.TestCase):
                 assert f.read() == "MLFLOW_RUN_ID"
 
         set_tag_mock.assert_called_with("kubeflow_run_id", "KUBEFLOW_RUN_ID")
+
+    @patch("kubernetes.client")
+    @patch("kubernetes.config")
+    def test_delete_pipeline_volume(self, k8s_config_mock, k8s_client_mock):
+        with um.patch(
+            "builtins.open", um.mock_open(read_data="unittest-namespace")
+        ):
+            runner = CliRunner()
+            result = runner.invoke(
+                delete_pipeline_volume,
+                ["workflow-name"],
+            )
+            assert result.exit_code == 0
+            core_api = k8s_client_mock.CoreV1Api()
+            core_api.delete_namespaced_persistent_volume_claim.assert_called_with(
+                "workflow-name", "unittest-namespace"
+            )
diff --git a/tests/test_generator.py b/tests/test_generator.py
index fd64959..c716678 100644
--- a/tests/test_generator.py
+++ b/tests/test_generator.py
@@ -47,7 +47,15 @@ class TestGenerator(unittest.TestCase):
 
         # then
         assert len(dsl_pipeline.ops) == 5
-        assert "schedule-volume-termination" in dsl_pipeline.ops
+        assert "on-exit" in dsl_pipeline.ops
+        assert (
+            dsl_pipeline.ops["on-exit"]
+            .container.command[-1]
+            .endswith(
+                "kedro kubeflow delete-pipeline-volume "
+                "{{workflow.name}}-pipeline-data-volume"
+            )
+        )
         volume_spec = dsl_pipeline.ops["data-volume-create"].k8s_resource.spec
         assert volume_spec.resources.requests["storage"] == "1Gi"
         assert volume_spec.access_modes == ["ReadWriteOnce"]
@@ -68,6 +76,54 @@ class TestGenerator(unittest.TestCase):
                 == 0
             )
 
+    def test_should_generate_on_exit_pipeline_run(self):
+        # given
+        self.create_generator(config={"on_exit_pipeline": "notify_via_slack"})
+
+        # when
+        pipeline = self.generator_under_test.generate_pipeline(
+            "pipeline", "unittest-image", "IfNotPresent"
+        )
+        with kfp.dsl.Pipeline(None) as dsl_pipeline:
+            pipeline()
+
+        # then
+        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 test_should_generate_volume_removal_and_on_exit_pipeline_run(self):
+        # given
+        self.create_generator(
+            config={"volume": {}, "on_exit_pipeline": "notify_via_slack"}
+        )
+
+        # when
+        pipeline = self.generator_under_test.generate_pipeline(
+            "pipeline", "unittest-image", "IfNotPresent"
+        )
+        with kfp.dsl.Pipeline(None) as dsl_pipeline:
+            pipeline()
+
+        # then
+        assert "on-exit" in dsl_pipeline.ops
+        assert (
+            dsl_pipeline.ops["on-exit"]
+            .container.command[-1]
+            .endswith(
+                "kedro kubeflow delete-pipeline-volume "
+                "{{workflow.name}}-pipeline-data-volume;"
+                "kedro run --config config.yaml "
+                "--env unittests --pipeline notify_via_slack"
+            )
+        )
+
     def test_should_support_inter_steps_volume_with_given_spec(self):
         # given
         self.create_generator(
@@ -89,7 +145,7 @@ class TestGenerator(unittest.TestCase):
 
         # then
         assert len(dsl_pipeline.ops) == 5
-        assert "schedule-volume-termination" in dsl_pipeline.ops
+        assert "on-exit" in dsl_pipeline.ops
         volume_spec = dsl_pipeline.ops["data-volume-create"].k8s_resource.spec
         assert volume_spec.resources.requests["storage"] == "1Mi"
         assert volume_spec.access_modes == ["ReadWriteOnce"]
@@ -174,7 +230,7 @@ class TestGenerator(unittest.TestCase):
         # then
         assert len(dsl_pipeline.ops) == 4
         assert "data-volume-create" in dsl_pipeline.ops
-        assert "schedule-volume-termination" in dsl_pipeline.ops
+        assert "on-exit" in dsl_pipeline.ops
         assert "data-volume-init" not in dsl_pipeline.ops
         for node_name in ["node1", "node2"]:
             volumes = dsl_pipeline.ops[node_name].container.volume_mounts