Repository: getindata/kedro-kubeflow (python)

Issue:
Allow to skip volume init
In case if there is no raw data in the image, `data-volume-init` step is not needed.

To be controlled via config: `volume.skip_init = true`

Pull request:
KFPClient adjusted to new config, ability to skip volume init
#### Description

* KFP Client was adjusted to the new object-like configuration
* Ability to skip volume init using `skip-init` flag in the config file (in case there is no raw data in the image)

Resolves #24

##### PR Checklist
- [x] Tests added
- [x] [Changelog](CHANGELOG.md) updated 


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index bee54ce..192d76e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@
 - Added sample support for [TemplatedConfigLoader](https://kedro.readthedocs.io/en/latest/kedro.config.TemplatedConfigLoader.html)
 - MLFlow support updated to not use nested runs.
 - Simple configuration validation added.
+- Volume init step is now optional (useful, if there is raw data in the image)
 
 ## [0.1.9] - 2021-01-08
 
diff --git a/README.md b/README.md
index cdedd74..a3e26f9 100644
--- a/README.md
+++ b/README.md
@@ -53,4 +53,5 @@ run_config:
     storageclass: # default
     size: 1Gi
     access_modes: [ReadWriteOnce]
+    skip_init: False
 ```
diff --git a/kedro_kubeflow/config.py b/kedro_kubeflow/config.py
index a66a49a..354dce6 100644
--- a/kedro_kubeflow/config.py
+++ b/kedro_kubeflow/config.py
@@ -14,6 +14,7 @@ run_config:
     storageclass: # default
     #size: 1Gi
     #access_modes: [ReadWriteOnce]
+    #skip_init: False
 """
 
 
@@ -22,7 +23,7 @@ class Config(object):
         self._raw = raw
 
     def _get_or_default(self, prop, default):
-        return self._raw[prop] if prop in self._raw.keys() else default
+        return self._raw.get(prop, default)
 
     def _get_or_fail(self, prop):
         if prop in self._raw.keys():
@@ -42,7 +43,7 @@ class Config(object):
 class VolumeConfig(Config):
     @property
     def storageclass(self):
-        return self._get_or_fail("storageclass")
+        return self._get_or_default("storageclass", None)
 
     @property
     def size(self):
@@ -50,7 +51,11 @@ class VolumeConfig(Config):
 
     @property
     def access_modes(self):
-        return self._get_or_default("access_modes", "[ReadWriteMany]")
+        return self._get_or_default("access_modes", ["ReadWriteMany"])
+
+    @property
+    def skip_init(self):
+        return self._get_or_default("skip_init", False)
 
     def _get_prefix(self):
         return "run_config.volume."
@@ -75,8 +80,11 @@ class RunConfig(Config):
 
     @property
     def volume(self):
-        cfg = self._get_or_fail("volume")
-        return VolumeConfig(cfg)
+        if "volume" in self._raw.keys():
+            cfg = self._get_or_fail("volume")
+            return VolumeConfig(cfg)
+        else:
+            return None
 
     @property
     def wait_for_completion(self):
diff --git a/kedro_kubeflow/context_helper.py b/kedro_kubeflow/context_helper.py
index 122a1e0..df59ba6 100644
--- a/kedro_kubeflow/context_helper.py
+++ b/kedro_kubeflow/context_helper.py
@@ -53,7 +53,7 @@ class ContextHelper16(ContextHelper):
 
     @property
     def project_name(self):
-        return self.get_context(self._metadata, self._env).project_name
+        return self.context.project_name
 
     @property
     def context(self):
diff --git a/kedro_kubeflow/kfpclient.py b/kedro_kubeflow/kfpclient.py
index 9175169..66884ea 100644
--- a/kedro_kubeflow/kfpclient.py
+++ b/kedro_kubeflow/kfpclient.py
@@ -25,12 +25,12 @@ class KubeflowClient(object):
 
     def __init__(self, config, project_name, context):
         token = self.obtain_id_token()
-        self.host = config["host"]
+        self.host = config.host
         self.client = Client(self.host, existing_token=token)
         self.project_name = project_name
         self.context = context
         dsl.ContainerOp._DISABLE_REUSABLE_COMPONENT_WARNING = True
-        self.volume_meta = config.get("run_config", {}).get("volume")
+        self.volume_meta = config.run_config.volume
 
     def list_pipelines(self):
         pipelines = self.client.list_pipelines(page_size=30).pipelines
@@ -115,29 +115,32 @@ class KubeflowClient(object):
             vop = dsl.VolumeOp(
                 name="data-volume-create",
                 resource_name="data-volume",
-                size=self.volume_meta.get("size", "1Gi"),
-                modes=self.volume_meta.get("access_modes", ["ReadWriteMany"]),
-                storage_class=self.volume_meta.get("storage_class"),
+                size=self.volume_meta.size,
+                modes=self.volume_meta.access_modes,
+                storage_class=self.volume_meta.storageclass,
             )
-            volume_init = dsl.ContainerOp(
-                name="data-volume-init",
-                image=image,
-                command=["sh", "-c"],
-                arguments=[
-                    " ".join(
-                        [
-                            "cp",
-                            "--verbose",
-                            "-r",
-                            "/home/kedro/data/*",
-                            "/home/kedro/datavolume",
-                        ]
-                    )
-                ],
-                pvolumes={"/home/kedro/datavolume": vop.volume},
-            )
-            volume_init.container.set_image_pull_policy(image_pull_policy)
-            return {"/home/kedro/data": volume_init.pvolume}
+            if self.volume_meta.skip_init:
+                return {"/home/kedro/data": vop.volume}
+            else:
+                volume_init = dsl.ContainerOp(
+                    name="data-volume-init",
+                    image=image,
+                    command=["sh", "-c"],
+                    arguments=[
+                        " ".join(
+                            [
+                                "cp",
+                                "--verbose",
+                                "-r",
+                                "/home/kedro/data/*",
+                                "/home/kedro/datavolume",
+                            ]
+                        )
+                    ],
+                    pvolumes={"/home/kedro/datavolume": vop.volume},
+                )
+                volume_init.container.set_image_pull_policy(image_pull_policy)
+                return {"/home/kedro/data": volume_init.pvolume}
 
         def _build_kfp_ops(
             node_dependencies: Dict[Node, Set[Node]], node_volumes: Dict
diff --git a/kedro_kubeflow/plugin.py b/kedro_kubeflow/plugin.py
index 9bcebb9..3245c82 100644
--- a/kedro_kubeflow/plugin.py
+++ b/kedro_kubeflow/plugin.py
@@ -175,7 +175,6 @@ def init(ctx, kfp_url: str):
     sample_config = PluginConfig.sample_config(
         url=kfp_url, image=image, project_name=context_helper.project_name
     )
-
     config_path = Path.cwd().joinpath("conf/base/kubeflow.yaml")
     with open(config_path, "w") as f:
         f.write(sample_config)


Test patch:
diff --git a/tests/test_mod.py b/tests/test_mod.py
index 48be56f..ba68d60 100644
--- a/tests/test_mod.py
+++ b/tests/test_mod.py
@@ -9,6 +9,7 @@ import kfp
 from google.auth.exceptions import DefaultCredentialsError
 from kedro.pipeline import Pipeline, node
 
+from kedro_kubeflow.config import PluginConfig
 from kedro_kubeflow.kfpclient import KubeflowClient
 from kedro_kubeflow.utils import strip_margin
 
@@ -166,7 +167,9 @@ class TestKubeflowClient(unittest.TestCase):
 
         # when
         self.client_under_test = KubeflowClient(
-            {"host": "http://unittest"}, None, None
+            PluginConfig({"host": "http://unittest", "run_config": {}}),
+            None,
+            None,
         )
 
         # then
@@ -188,7 +191,9 @@ class TestKubeflowClient(unittest.TestCase):
         ) as cm:
             # when
             self.client_under_test = KubeflowClient(
-                {"host": "http://unittest"}, None, None
+                PluginConfig({"host": "http://unittest", "run_config": {}}),
+                None,
+                None,
             )
             # then
             assert (
@@ -213,7 +218,9 @@ class TestKubeflowClient(unittest.TestCase):
         with self.assertLogs("kedro_kubeflow.kfpclient", level="ERROR") as cm:
             # when
             self.client_under_test = KubeflowClient(
-                {"host": "http://unittest"}, None, None
+                PluginConfig({"host": "http://unittest", "run_config": {}}),
+                None,
+                None,
             )
             # then
             assert "Failed to obtain IAP access token" in cm.output[0]
@@ -389,7 +396,7 @@ 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()
 
-    def test_should_support_inter_steps_volume_with_defauls(self):
+    def test_should_support_inter_steps_volume_with_defaults(self):
         # given
         run_mock = unittest.mock.MagicMock()
         self.kfp_client_mock.create_run_from_pipeline_func.return_value = (
@@ -445,7 +452,7 @@ class TestKubeflowClient(unittest.TestCase):
         self.create_client(
             {
                 "volume": {
-                    "storage_class": "nfs",
+                    "storageclass": "nfs",
                     "size": "1Mi",
                     "access_modes": ["ReadWriteOnce"],
                 }
@@ -527,6 +534,46 @@ class TestKubeflowClient(unittest.TestCase):
                 == "{{pipelineparam:op=mlflow-start-run;name=mlflow_run_id}}"
             )
 
+    def test_should_skip_volume_init_if_requested(self):
+        # given
+        run_mock = unittest.mock.MagicMock()
+        self.kfp_client_mock.create_run_from_pipeline_func.return_value = (
+            run_mock
+        )
+        self.create_client({"volume": {"skip_init": True}})
+
+        # when
+        self.client_under_test.run_once(
+            run_name="unittest",
+            pipeline="pipeline",
+            image="unittest-image",
+            experiment_name="experiment",
+            wait=False,
+        )
+
+        # then
+        self.kfp_client_mock.create_run_from_pipeline_func.assert_called()
+        run_mock.wait_for_run_completion.assert_not_called()
+        (
+            args,
+            kwargs,
+        ) = self.kfp_client_mock.create_run_from_pipeline_func.call_args
+        assert kwargs == {
+            "arguments": {},
+            "experiment_name": "experiment",
+            "run_name": "unittest",
+        }
+
+        with kfp.dsl.Pipeline(None) as dsl_pipeline:
+            args[0]()
+
+        assert len(dsl_pipeline.ops) == 3
+        assert "data-volume-init" not in dsl_pipeline.ops
+        for node_name in ["node1", "node2"]:
+            volumes = dsl_pipeline.ops[node_name].container.volume_mounts
+            assert len(volumes) == 1
+            assert volumes[0].name == "data-volume-create"
+
     @patch("kedro_kubeflow.kfpclient.Client")
     def create_client(self, config, kfp_client_mock):
         project_name = "my-awesome-project"
@@ -536,7 +583,7 @@ class TestKubeflowClient(unittest.TestCase):
             {"pipelines": {"pipeline": self.create_pipeline()}},
         )
         self.client_under_test = KubeflowClient(
-            {"host": "http://unittest", "run_config": config},
+            PluginConfig({"host": "http://unittest", "run_config": config}),
             project_name,
             context,
         )