Repository: getindata/kedro-kubeflow
Pull request #9 (2021-01-06 09:25:17)

MLOPS-54 | Use volume for the inter-steps data

szczeles: @em-pe I decided to use `run_config` directly in the client class (instead of reading it in 4 functions in `plugin.py` and passing as params), maybe we can do same for other args: image, pull policy, etc. It would simplify the code - and anyway  - all these are required for most of the client functions.
szczeles: @em-pe Done, please recheck

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 9309bc5..40925d3 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,13 +2,15 @@
 
 ## [Unreleased]
 
+- Support for inter-steps volume: setup (one volume per pipeline run), initial load (the content of `data/` directory from the image and mount to all the steps for artifacts passing.
+
 ## [0.1.8] - 2021-01-05
 
 ### Added
 
 -   *Initial release* of kedro-kubeflow plugin
 -   Ability to run an anonymous pipeline once as within a specified experiment `kedro kubeflow run-once`.
--   Ability to upload pipeline `kedro kubeflow upload`.
+-   Ability to upload pipeline `kedro kubeflow upload-pipeline`.
 -   Method to schedule runs for most recent version of given pipeline `kedro kubeflow schedule` 
 -   Shortcut to open UI for pipelines using `kedro kubeflow ui` 
 
diff --git a/kedro_kubeflow/kfpclient.py b/kedro_kubeflow/kfpclient.py
index af84ad4..4d5bd2f 100644
--- a/kedro_kubeflow/kfpclient.py
+++ b/kedro_kubeflow/kfpclient.py
@@ -16,7 +16,8 @@ from kedro.pipeline.node import Node
 
 IAP_CLIENT_ID = "IAP_CLIENT_ID"
 
-WAIT_TIMEOUT = 24*60*60
+WAIT_TIMEOUT = 24 * 60 * 60
+
 
 class KubeflowClient(object):
 
@@ -24,22 +25,33 @@ 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")
 
     def list_pipelines(self):
         pipelines = self.client.list_pipelines(page_size=30).pipelines
-        return tabulate(map(lambda x: [x.name, x.id], pipelines), headers=["Name", "ID"])
+        return tabulate(
+            map(lambda x: [x.name, x.id], pipelines), headers=["Name", "ID"]
+        )
 
-    def run_once(self, pipeline, image, experiment_name, run_name, wait, image_pull_policy="IfNotPresent") -> None:
+    def run_once(
+        self,
+        pipeline,
+        image,
+        experiment_name,
+        run_name,
+        wait,
+        image_pull_policy="IfNotPresent",
+    ) -> None:
         run = self.client.create_run_from_pipeline_func(
             self.generate_pipeline(pipeline, image, image_pull_policy),
             arguments={},
             experiment_name=experiment_name,
-            run_name=run_name
+            run_name=run_name,
         )
 
         if wait:
@@ -55,7 +67,9 @@ class KubeflowClient(object):
         jwt_token = None
 
         if not client_id:
-            self.log.info("No IAP_CLIENT_ID provided, skipping custom IAP authentication")
+            self.log.info(
+                "No IAP_CLIENT_ID provided, skipping custom IAP authentication"
+            )
             return jwt_token
 
         try:
@@ -63,31 +77,76 @@ class KubeflowClient(object):
             jwt_token = id_token.fetch_id_token(Request(), client_id)
             self.log.info("Obtained JWT token for MLFLOW connectivity.")
         except DefaultCredentialsError as ex:
-            self.log.warning(str(ex) + (" Note that this authentication method does not work with default credentials"
-                                   " obtained via 'gcloud auth application-default login' command. Refer to"
-                                   " documentation on how to configure service account locally "
-                                   "(https://cloud.google.com/docs/authentication/production#manually)"))
+            self.log.warning(
+                str(ex)
+                + (
+                    " Note that this authentication method does not work with default credentials"
+                    " obtained via 'gcloud auth application-default login' command. Refer to"
+                    " documentation on how to configure service account locally "
+                    "(https://cloud.google.com/docs/authentication/production#manually)"
+                )
+            )
         except Exception as e:
             self.log.error("Failed to obtain IAP access token. " + str(e))
         finally:
             return jwt_token
 
     def generate_pipeline(self, pipeline, image, image_pull_policy):
-        @dsl.pipeline(name=self.project_name, description="Kubeflow pipeline for Kedro project")
+        @dsl.pipeline(
+            name=self.project_name,
+            description="Kubeflow pipeline for Kedro project",
+        )
         def convert_kedro_pipeline_to_kfp() -> None:
             """Convert from a Kedro pipeline into a kfp container graph."""
 
-            node_dependencies = self.context.pipelines.get(pipeline).node_dependencies
-            kfp_ops = _build_kfp_ops(node_dependencies)
+            node_volumes = (
+                _setup_volumes() if self.volume_meta is not None else {}
+            )
+            node_dependencies = self.context.pipelines.get(
+                pipeline
+            ).node_dependencies
+            kfp_ops = _build_kfp_ops(node_dependencies, node_volumes)
             for node, dependencies in node_dependencies.items():
                 for dependency in dependencies:
                     kfp_ops[node.name].after(kfp_ops[dependency.name])
 
-        def _build_kfp_ops(node_dependencies: Dict[Node, Set[Node]]) -> Dict[str, dsl.ContainerOp]:
+        def _setup_volumes():
+            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"),
+            )
+            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
+        ) -> Dict[str, dsl.ContainerOp]:
             """Build kfp container graph from Kedro node dependencies. """
             kfp_ops = {}
 
-            env = V1EnvVar(name=IAP_CLIENT_ID, value=os.environ.get(IAP_CLIENT_ID, ""))
+            env = V1EnvVar(
+                name=IAP_CLIENT_ID, value=os.environ.get(IAP_CLIENT_ID, "")
+            )
 
             for node in node_dependencies:
                 name = _clean_name(node.name)
@@ -96,25 +155,33 @@ class KubeflowClient(object):
                     image=image,
                     command=["kedro"],
                     arguments=["run", "--node", node.name],
+                    pvolumes=node_volumes,
                 )
                 kfp_ops[node.name].container.add_env_variable(env)
-                kfp_ops[node.name].container.set_image_pull_policy(image_pull_policy)
+                kfp_ops[node.name].container.set_image_pull_policy(
+                    image_pull_policy
+                )
 
             return kfp_ops
 
         return convert_kedro_pipeline_to_kfp
 
-
-    def compile(self, pipeline, image, output, image_pull_policy='IfNotPresent'):
-        Compiler().compile(self.generate_pipeline(pipeline, image, image_pull_policy), output)
+    def compile(
+        self, pipeline, image, output, image_pull_policy="IfNotPresent"
+    ):
+        Compiler().compile(
+            self.generate_pipeline(pipeline, image, image_pull_policy), output
+        )
         self.log.info("Generated pipeline definition was saved to %s" % output)
 
-    def upload(self, pipeline, image, image_pull_policy='IfNotPresent'):
+    def upload(self, pipeline, image, image_pull_policy="IfNotPresent"):
         pipeline = self.generate_pipeline(pipeline, image, image_pull_policy)
 
         if self._pipeline_exists(self.project_name):
             pipeline_id = self._get_pipeline_id(self.project_name)
-            version_id = self._upload_pipeline_version(pipeline, pipeline_id, self.project_name)
+            version_id = self._upload_pipeline_version(
+                pipeline, pipeline_id, self.project_name
+            )
             self.log.info("New version of pipeline created: %s", version_id)
         else:
             (pipeline_id, version_id) = self._upload_pipeline(
@@ -141,7 +208,11 @@ class KubeflowClient(object):
                 filter=json.dumps(
                     {
                         "predicates": [
-                            {"key": "name", "op": 1, "string_value": pipeline_name}
+                            {
+                                "key": "name",
+                                "op": 1,
+                                "string_value": pipeline_name,
+                            }
                         ]
                     }
                 )
@@ -150,7 +221,9 @@ class KubeflowClient(object):
             .id
         )
 
-    def _upload_pipeline_version(self, pipeline_func, pipeline_id, pipeline_name):
+    def _upload_pipeline_version(
+        self, pipeline_func, pipeline_id, pipeline_name
+    ):
         version_name = f"{_clean_name(pipeline_name)}-{uuid.uuid4()}"[:100]
         with NamedTemporaryFile(suffix=".yaml") as f:
             Compiler().compile(pipeline_func, f.name)
@@ -173,21 +246,18 @@ class KubeflowClient(object):
             )
             self.log.info(f"Existing experiment found: {experiment.id}")
         except:
-            experiment = self.client.create_experiment(
-                experiment_name
-            )
+            experiment = self.client.create_experiment(experiment_name)
             self.log.info(f"New experiment created: {experiment.id}")
 
         return experiment.id
 
-
     def schedule(self, experiment_name, cron_expression):
         experiment_id = self._ensure_experiment_exists(experiment_name)
         pipeline_id = self._get_pipeline_id(self.project_name)
         self._disable_runs(experiment_id, pipeline_id)
         job = self.client.create_recurring_run(
             experiment_id,
-            f'{self.project_name} on {cron_expression}',
+            f"{self.project_name} on {cron_expression}",
             cron_expression=cron_expression,
             pipeline_id=pipeline_id,
         )
@@ -196,7 +266,11 @@ class KubeflowClient(object):
     def _disable_runs(self, experiment_id, pipeline_id):
         runs = self.client.list_recurring_runs(experiment_id=experiment_id)
         if runs.jobs is not None:
-            my_runs = [job for job in runs.jobs if job.pipeline_spec.pipeline_id == pipeline_id]
+            my_runs = [
+                job
+                for job in runs.jobs
+                if job.pipeline_spec.pipeline_id == pipeline_id
+            ]
             for job in my_runs:
                 self.client.jobs.delete_job(job.id)
                 self.log.info(f"Previous schedule deleted {job.id}")
@@ -204,5 +278,3 @@ class KubeflowClient(object):
 
 def _clean_name(name: str) -> str:
     return re.sub(r"[\W_]+", "-", name).strip("-")
-
-


Test patch:
diff --git a/tests/test_mod.py b/tests/test_mod.py
index f1b02e0..8a8d7dd 100644
--- a/tests/test_mod.py
+++ b/tests/test_mod.py
@@ -105,6 +105,7 @@ class TestKubeflowClient(unittest.TestCase):
         with kfp.dsl.Pipeline(None) as dsl_pipeline:
             args[0]()
 
+        assert len(dsl_pipeline.ops) == 2
         assert dsl_pipeline.ops["node1"].container.image == "unittest-image"
         assert dsl_pipeline.ops["node1"].container.image_pull_policy == "IfNotPresent"
 
@@ -339,16 +340,104 @@ 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):
+        # given
+        run_mock = unittest.mock.MagicMock()
+        self.kfp_client_mock.create_run_from_pipeline_func.return_value = run_mock
+        self.create_client({"volume": {}})
+
+        # 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) == 4
+        volume_spec = dsl_pipeline.ops["data-volume-create"].k8s_resource.spec
+        assert volume_spec.resources.requests["storage"] == "1Gi"
+        assert volume_spec.access_modes == ["ReadWriteMany"]
+        assert volume_spec.storage_class_name is None
+        volume_init_spec = dsl_pipeline.ops["data-volume-init"].container
+        assert volume_init_spec.image == "unittest-image"
+        assert volume_init_spec.image_pull_policy == "IfNotPresent"
+        assert volume_init_spec.args[0].startswith("cp --verbose -r")
+        for node in ["data-volume-init", "node1", "node2"]:
+            volumes = dsl_pipeline.ops[node].container.volume_mounts
+            assert len(volumes) == 1
+            assert volumes[0].name == "data-volume-create"
+
+    def test_should_support_inter_steps_volume_with_given_spec(self):
+        # given
+        run_mock = unittest.mock.MagicMock()
+        self.kfp_client_mock.create_run_from_pipeline_func.return_value = run_mock
+        self.create_client(
+            {
+                "volume": {
+                    "storage_class": "nfs",
+                    "size": "1Mi",
+                    "access_modes": ["ReadWriteOnce"],
+                }
+            }
+        )
+
+        # 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) == 4
+        volume_spec = dsl_pipeline.ops["data-volume-create"].k8s_resource.spec
+        assert volume_spec.resources.requests["storage"] == "1Mi"
+        assert volume_spec.access_modes == ["ReadWriteOnce"]
+        assert volume_spec.storage_class_name == "nfs"
+
     def tearDown(self):
         os.environ["IAP_CLIENT_ID"] = ""
 
     @patch("kedro_kubeflow.kfpclient.Client")
-    def setUp(self, kfp_client_mock):
+    def create_client(self, config, kfp_client_mock):
         project_name = "my-awesome-project"
         context = type(
             "obj", (object,), {"pipelines": {"pipeline": self.create_pipeline()}},
         )
         self.client_under_test = KubeflowClient(
-            {"host": "http://unittest"}, project_name, context
+            {"host": "http://unittest", "run_config": config}, project_name, context
         )
+        self.client_under_test.client = kfp_client_mock
         self.kfp_client_mock = self.client_under_test.client
+
+    def setUp(self):
+        self.create_client({})