Repository: getindata/kedro-kubeflow
Pull request #22 (2021-01-08 10:48:49)

Release 0.1.9
Bump version and CHANGELOG for next release.


Patch:
diff --git a/.codeclimate.yml b/.codeclimate.yml
new file mode 100644
index 0000000..dbf9a37
--- /dev/null
+++ b/.codeclimate.yml
@@ -0,0 +1,3 @@
+plugins:
+  pylint:
+    enabled: true
diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS
new file mode 100644
index 0000000..95045a2
--- /dev/null
+++ b/.github/CODEOWNERS
@@ -0,0 +1,2 @@
+# Define global code owners
+* @empe @szczeles
diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md
new file mode 100644
index 0000000..f5f00af
--- /dev/null
+++ b/.github/pull_request_template.md
@@ -0,0 +1,9 @@
+#### Description
+
+`describe the purpose of the change here` 
+
+Resolves `<issue nr here>`
+
+##### PR Checklist
+- [ ] Tests added
+- [ ] [Changelog](CHANGELOG.md) updated 
diff --git a/.github/workflows/python-package.yml b/.github/workflows/python-package.yml
index 43abbb5..6273b9c 100644
--- a/.github/workflows/python-package.yml
+++ b/.github/workflows/python-package.yml
@@ -4,6 +4,7 @@ on:
   push:
     branches: 
       - master
+      - develop
   pull_request:
 
 jobs:
@@ -16,6 +17,7 @@ jobs:
       uses: actions/setup-python@v2.2.1
       with:
         python-version: 3.7
+
     - name: Setup virtualenv
       run: |
         python -V
@@ -23,8 +25,20 @@ jobs:
         virtualenv venv
         source venv/bin/activate
    
+    - name: Check pre-commit status
+      run: |
+        pip install -r requirements-dev.txt
+        pre-commit run --all-files
+
     - name: Test with tox
       run: |
-        python -m pip install flake8 tox
-        if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
         tox -e py37
+
+    - name: Report coverage
+      uses: paambaati/codeclimate-action@v2.7.5
+      env:
+        CC_TEST_REPORTER_ID: ${{ secrets.CODE_CLIMATE }}
+      with:
+        coverageCommand: coverage xml
+        debug: true
+        coverageLocations: coverage.xml:coverage.py
diff --git a/.isort.cfg b/.isort.cfg
new file mode 100644
index 0000000..cdab8f2
--- /dev/null
+++ b/.isort.cfg
@@ -0,0 +1,2 @@
+[settings]
+known_third_party = click,google,kedro,kfp,kubernetes,semver,setuptools,tabulate
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
new file mode 100644
index 0000000..ca0370f
--- /dev/null
+++ b/.pre-commit-config.yaml
@@ -0,0 +1,13 @@
+repos:
+- repo: https://github.com/pre-commit/mirrors-isort
+  rev: v4.3.21
+  hooks:
+  - id: isort
+- repo: https://github.com/psf/black
+  rev: stable
+  hooks:
+  - id: black
+- repo: https://github.com/pre-commit/pre-commit-hooks
+  rev: v2.3.0
+  hooks:
+  - id: flake8
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 9309bc5..0a281af 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,16 +2,24 @@
 
 ## [Unreleased]
 
+## [0.1.9] - 2021-01-08
+
+-   Support for MLFlow - if the package is installed then additional step is added with parent run init. Then all separate nodes runs register under this run.
+-   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.
+-   `kubeflow init` command added to generate sample configuration file.
+
 ## [0.1.8] - 2021-01-05
 
 ### Added
 
--   *Initial release* of kedro-kubeflow plugin
+-   _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` 
 
-[Unreleased]: https://github.com/getindata/kedro-kubeflow/compare/0.1.8...HEAD
+[Unreleased]: https://github.com/getindata/kedro-kubeflow/compare/0.1.9...HEAD
+
+[0.1.9]: https://github.com/getindata/kedro-kubeflow/compare/0.1.8...0.1.9
 
 [0.1.8]: https://github.com/getindata/kedro-kubeflow/compare/ea219ae5f70e726b7afc9d0864da4b6649e4deaf...0.1.8
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index 8ce2c91..22e88d6 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -1,15 +1,25 @@
-# Release workflow
+
+## PR Guidelines
+1. Fork branch from `develop`.
+1. Ensure to provide unit tests for new functionality.
+1. Install dev requirements: `pip install -r requirements-dev.txt` and setup a hook: `pre-commit install`
+1. Update documentation accordingly.
+1. Update [changelog](CHANGELOG.md) according to ["Keep a changelog"](https://keepachangelog.com/en/1.0.0/) guidelines.
+1. Squash changes with a single commit as much as possible and ensure verbose PR name.
+1. Open a PR against `develop`
+
+*We reserve the right to take over and modify or abandon PRs that do not match the workflow or are abandoned.* 
+
+## Release workflow
 
 1. Create the release candidate:
-    - Go to the `create-release-candidate action`
+    - Go to the [Prepare release](https://github.com/getindata/kedro-kubeflow/actions?query=workflow%3A%22Prepare+release%22) action.
     - Click "Run workflow"
-    - Enter the part of the version to bump (one of `<major>.<minor>.<patch>`)
+    - Enter the part of the version to bump (one of `<major>.<minor>.<patch>`). Minor (x.**x**.x) is a default. 
 2. If the workflow has run sucessfully:
     - Go to the newly openened PR named `Release candidate `<version>`
-    - Check that changelog and version have been properly updated.
-    - *(If everything is normal, skip this step)* Eventually pull the branch and make changes if necessary
+    - Check that changelog and version have been properly updated. If not pull the branch and apply manual changes if necessary.
     - Merge the PR to master
-3. Checkout the `publish workflow`to see if:
-    - The package has been uploaded on PyPI sucessfully
+3. Checkout the [Publish](https://github.com/getindata/kedro-kubeflow/actions?query=workflow%3APublish) workflow to see if:
+    - The package has been uploaded on PyPI successfully
     - The changes have been merged back to develop
-4. If the pipeline has failed, please raise an issue to correct the CI, and ensure merge on develop "by hand"".
diff --git a/README.md b/README.md
index ebedf55..cdedd74 100644
--- a/README.md
+++ b/README.md
@@ -1,9 +1,13 @@
 # Kedro Kubeflow Plugin
 
-[![Python Version](https://img.shields.io/badge/python-3.6%20%7C%203.7%20%7C%203.8-blue.svg)](https://github.com/getindata/kedro-kubeflow)
+[![Python Version](https://img.shields.io/badge/python-3.7%20%7C%203.8-blue.svg)](https://github.com/getindata/kedro-kubeflow)
 [![License](https://img.shields.io/badge/license-Apache%202.0-blue.svg)](https://opensource.org/licenses/Apache-2.0) 
 [![SemVer](https://img.shields.io/badge/semver-2.0.0-green)](https://semver.org/)
+[![PyPI version](https://badge.fury.io/py/kedro-kubeflow.svg)](https://pypi.org/project/kedro-kubeflow/)
+[![Downloads](https://pepy.tech/badge/kedro-kubeflow)](https://pepy.tech/project/kedro-kubeflow) 
 
+[![Maintainability](https://api.codeclimate.com/v1/badges/fff07cbd2e5012a045a3/maintainability)](https://codeclimate.com/github/getindata/kedro-kubeflow/maintainability) 
+[![Test Coverage](https://api.codeclimate.com/v1/badges/fff07cbd2e5012a045a3/test_coverage)](https://codeclimate.com/github/getindata/kedro-kubeflow/test_coverage)
 ## About
 
 The main purpose of this plugin is to enable running kedro pipeline on Kubeflow Pipelines. It supports translation from 
@@ -24,9 +28,29 @@ Usage: kedro kubeflow [OPTIONS] COMMAND [ARGS]...
  
  Commands:
    compile          Translates Kedro pipeline into YAML file with Kubeflow pipeline definition
+   init             Initializes configuration for the plugin
    list-pipelines   List deployed pipeline definitions
    run-once         Deploy pipeline as a single run within given experiment.
    schedule         Schedules recurring execution of latest version of the pipeline
    ui               Open Kubeflow Pipelines UI in new browser tab
    upload-pipeline  Uploads pipeline to Kubeflow server
 ```
+
+## Configuration file
+
+`kedro init` generates configuration file for the plugin, but users may want
+to adjust it to the requirements of the environment:
+
+```
+host: http://10.43.77.224
+
+run_config:
+  image: new-kedro-project
+  experiment_name: New Kedro Project
+  run_name: New Kedro Project
+  wait_for_completion: False
+  volume:
+    storageclass: # default
+    size: 1Gi
+    access_modes: [ReadWriteOnce]
+```
diff --git a/kedro_kubeflow/__init__.py b/kedro_kubeflow/__init__.py
index b5bad32..35777a6 100644
--- a/kedro_kubeflow/__init__.py
+++ b/kedro_kubeflow/__init__.py
@@ -1,4 +1,3 @@
 """kedro_kubeflow."""
 
-version = "0.1.8"
-
+version = "0.1.9"
diff --git a/kedro_kubeflow/hooks.py b/kedro_kubeflow/hooks.py
new file mode 100644
index 0000000..0e0ef6f
--- /dev/null
+++ b/kedro_kubeflow/hooks.py
@@ -0,0 +1,38 @@
+import os
+from typing import Any, Dict
+
+from kedro.framework.context import load_context
+from kedro.framework.hooks import hook_impl
+from kedro.io import DataCatalog
+from kedro.pipeline import Pipeline
+
+from .utils import is_mlflow_enabled
+
+
+class MLFlowActivateParentHook:
+    @hook_impl
+    def before_pipeline_run(
+        self,
+        run_params: Dict[str, Any],
+        pipeline: Pipeline,
+        catalog: DataCatalog,
+    ) -> None:
+
+        if not is_mlflow_enabled():
+            return
+
+        import mlflow
+        from kedro_mlflow.framework.context import get_mlflow_config
+
+        context = load_context(
+            project_path=run_params["project_path"],
+            env=run_params["env"],
+            extra_params=run_params["extra_params"],
+        )
+        mlflow_conf = get_mlflow_config(context)
+        mlflow_conf.setup(context)
+
+        mlflow.start_run(run_id=os.getenv("MLFLOW_PARENT_ID"))
+
+
+mlflow_activate_parent_run = MLFlowActivateParentHook()
diff --git a/kedro_kubeflow/kfpclient.py b/kedro_kubeflow/kfpclient.py
index af84ad4..36401e5 100644
--- a/kedro_kubeflow/kfpclient.py
+++ b/kedro_kubeflow/kfpclient.py
@@ -1,22 +1,23 @@
-import re
-import os
-import logging
 import json
+import logging
+import os
+import re
 import uuid
-from pathlib import Path
 from tempfile import NamedTemporaryFile
+from typing import Dict, Set
 
+from kedro.pipeline.node import Node
 from kfp import Client, dsl
 from kfp.compiler import Compiler
 from kubernetes.client import V1EnvVar
 from tabulate import tabulate
-from typing import Dict, Set
-from kedro.pipeline.node import Node
 
+from .utils import is_mlflow_enabled
 
 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,101 @@ 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, "")
+                )
+            ]
+
+            if is_mlflow_enabled():
+                kfp_ops["mlflow-start-run"] = dsl.ContainerOp(
+                    name="mlflow-start-run",
+                    image=image,
+                    command=["kedro"],
+                    arguments=[
+                        "kubeflow",
+                        "mlflow-start",
+                        dsl.RUN_ID_PLACEHOLDER,
+                    ],
+                    file_outputs={"mlflow_run_id": "/tmp/mlflow_run_id"},
+                )
+                kfp_ops["mlflow-start-run"].container.set_image_pull_policy(
+                    image_pull_policy
+                )
+                env.append(
+                    V1EnvVar(
+                        name="MLFLOW_PARENT_ID",
+                        value=kfp_ops["mlflow-start-run"].output,
+                    )
+                )
 
             for node in node_dependencies:
                 name = _clean_name(node.name)
@@ -96,25 +180,33 @@ class KubeflowClient(object):
                     image=image,
                     command=["kedro"],
                     arguments=["run", "--node", node.name],
+                    pvolumes=node_volumes,
+                    container_kwargs={"env": env},
+                )
+                kfp_ops[node.name].container.set_image_pull_policy(
+                    image_pull_policy
                 )
-                kfp_ops[node.name].container.add_env_variable(env)
-                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(
@@ -129,28 +221,29 @@ class KubeflowClient(object):
         )
 
     def _pipeline_exists(self, pipeline_name):
-        try:
-            self._get_pipeline_id(pipeline_name)
-            return True
-        except:
-            return False
+        return self._get_pipeline_id(pipeline_name) is not None
 
     def _get_pipeline_id(self, pipeline_name):
-        return (
-            self.client.pipelines.list_pipelines(
-                filter=json.dumps(
-                    {
-                        "predicates": [
-                            {"key": "name", "op": 1, "string_value": pipeline_name}
-                        ]
-                    }
-                )
+        pipelines = self.client.pipelines.list_pipelines(
+            filter=json.dumps(
+                {
+                    "predicates": [
+                        {
+                            "key": "name",
+                            "op": 1,
+                            "string_value": pipeline_name,
+                        }
+                    ]
+                }
             )
-            .pipelines[0]
-            .id
-        )
+        ).pipelines
+
+        if pipelines:
+            return pipelines[0].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)
@@ -172,22 +265,22 @@ class KubeflowClient(object):
                 experiment_name=experiment_name
             )
             self.log.info(f"Existing experiment found: {experiment.id}")
-        except:
-            experiment = self.client.create_experiment(
-                experiment_name
-            )
+        except ValueError as e:
+            if not str(e).startswith("No experiment is found"):
+                raise
+
+            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(
+        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 +289,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 +301,3 @@ class KubeflowClient(object):
 
 def _clean_name(name: str) -> str:
     return re.sub(r"[\W_]+", "-", name).strip("-")
-
-
diff --git a/kedro_kubeflow/plugin.py b/kedro_kubeflow/plugin.py
index c77c95b..057aed8 100644
--- a/kedro_kubeflow/plugin.py
+++ b/kedro_kubeflow/plugin.py
@@ -1,9 +1,12 @@
-import click
+from pathlib import Path
 
-from semver import VersionInfo
+import click
 from kedro import __version__ as kedro_version
 from kedro.config import MissingConfigException
+from semver import VersionInfo
+
 from .kfpclient import KubeflowClient
+from .utils import strip_margin
 
 KEDRO_VERSION = VersionInfo.parse(kedro_version)
 
@@ -11,13 +14,16 @@ if KEDRO_VERSION.match(">=0.17.0"):
     from kedro.framework.session import KedroSession
 
     def get_context(metadata, env):
-        return KedroSession.create(metadata.package_name, env=env).load_context()
+        return KedroSession.create(
+            metadata.package_name, env=env
+        ).load_context()
 
     def get_project_name(metadata, env):
         return metadata.project_name
-else: # <0.17
+
+
+else:  # <0.17
     from kedro.framework.context import load_context
-    from pathlib import Path
 
     def get_context(metadata, env):
         return load_context(Path.cwd(), env=env)
@@ -28,99 +34,263 @@ else: # <0.17
 
 CONFIG_FILE_PATTERN = "kubeflow*"
 
+
 @click.group("Kubeflow")
 def commands():
     """Kedro plugin adding support for Kubeflow Pipelines"""
     pass
 
-@commands.group(name="kubeflow", context_settings=dict(help_option_names=["-h", "--help"]))
-@click.option("-e", "--env", "env", type=str, default="base", help="Environment to use.")
+
+@commands.group(
+    name="kubeflow", context_settings=dict(help_option_names=["-h", "--help"])
+)
+@click.option(
+    "-e", "--env", "env", type=str, default="base", help="Environment to use."
+)
 @click.pass_obj
 @click.pass_context
 def kubeflow_group(ctx, metadata, env):
     """Interact with Kubeflow Pipelines"""
     ctx.ensure_object(dict)
-    ctx.obj['kedro_ctx'] = get_context(metadata, env)
-    ctx.obj['config'] = ctx.obj['kedro_ctx'].config_loader.get(CONFIG_FILE_PATTERN)
-    if 'host' not in ctx.obj['config'].keys():
+    ctx.obj["kedro_ctx"] = get_context(metadata, env)
+    ctx.obj["project_name"] = get_project_name(metadata, env)
+
+
+def setup_context(ctx):
+    ctx.obj["config"] = ctx.obj["kedro_ctx"].config_loader.get(
+        CONFIG_FILE_PATTERN
+    )
+    if "host" not in ctx.obj["config"].keys():
         raise MissingConfigException("No 'host' defined in kubeflow.yml")
 
-    ctx.obj['kfp_client'] = KubeflowClient(ctx.obj['config'], get_project_name(metadata, env), ctx.obj['kedro_ctx'])
+    ctx.obj["kfp_client"] = KubeflowClient(
+        ctx.obj["config"], ctx.obj["project_name"], ctx.obj["kedro_ctx"]
+    )
 
 
 @kubeflow_group.command()
 @click.pass_context
 def list_pipelines(ctx):
     """List deployed pipeline definitions"""
-    click.echo(ctx.obj['kfp_client'].list_pipelines())
+    setup_context(ctx)
+    click.echo(ctx.obj["kfp_client"].list_pipelines())
 
 
 @kubeflow_group.command()
-@click.option("-i", "--image", type=str, help="Docker image to use for pipeline execution.")
-@click.option("-p", "--pipeline", "pipeline", type=str, help="Name of pipeline to run", default="__default__")
-@click.option("-x", "--experiment-name", "experiment_name", type=str, help="Name of experiment associated with this run.")
-@click.option("-r", "--run-name", "run_name", type=str, help="Name for this run.")
+@click.option(
+    "-i",
+    "--image",
+    type=str,
+    help="Docker image to use for pipeline execution.",
+)
+@click.option(
+    "-p",
+    "--pipeline",
+    "pipeline",
+    type=str,
+    help="Name of pipeline to run",
+    default="__default__",
+)
+@click.option(
+    "-x",
+    "--experiment-name",
+    "experiment_name",
+    type=str,
+    help="Name of experiment associated with this run.",
+)
+@click.option(
+    "-r", "--run-name", "run_name", type=str, help="Name for this run."
+)
 @click.option("-w", "--wait", "wait", type=bool, help="Wait for completion.")
-@click.option("--image-pull-policy", type=str, default="IfNotPresent", help="Image pull policy.")
+@click.option(
+    "--image-pull-policy",
+    type=str,
+    default="IfNotPresent",
+    help="Image pull policy.",
+)
 @click.pass_context
-def run_once(ctx, image: str, pipeline: str, experiment_name: str, run_name: str, wait: bool, image_pull_policy: str):
-    """Deploy pipeline as a single run within given experiment. Config can be specified in kubeflow.yml as well."""
-    conf = ctx.obj['config']
+def run_once(
+    ctx,
+    image: str,
+    pipeline: str,
+    experiment_name: str,
+    run_name: str,
+    wait: bool,
+    image_pull_policy: str,
+):
+    """Deploy pipeline as a single run within given experiment.
+    Config can be specified in kubeflow.yml as well."""
+    setup_context(ctx)
+    conf = ctx.obj["config"]
     run_conf = conf.get("run_config", {})
-    image = image if image else run_conf['image']
-    experiment_name = experiment_name if experiment_name else run_conf['experiment_name']
-    run_name = run_name if run_name else run_conf['run_name']
+    image = image if image else run_conf["image"]
+    experiment_name = (
+        experiment_name if experiment_name else run_conf["experiment_name"]
+    )
+    run_name = run_name if run_name else run_conf["run_name"]
     wait = wait if wait is not None else bool(run_conf["wait_for_completion"])
-    image_pull_policy = run_conf.get('image_pull_policy', image_pull_policy)
+    image_pull_policy = run_conf.get("image_pull_policy", image_pull_policy)
 
-    ctx.obj['kfp_client'].run_once(pipeline, image, experiment_name, run_name, wait, image_pull_policy)
+    ctx.obj["kfp_client"].run_once(
+        pipeline, image, experiment_name, run_name, wait, image_pull_policy
+    )
 
 
 @kubeflow_group.command()
 @click.pass_context
 def ui(ctx) -> None:
     """Open Kubeflow Pipelines UI in new browser tab"""
+    setup_context(ctx)
     import webbrowser
-    host = ctx.obj['config']['host']
+
+    host = ctx.obj["config"]["host"]
     webbrowser.open_new_tab(host)
 
 
 @kubeflow_group.command()
-@click.option("-i", "--image", type=str, help="Docker image to use for pipeline execution.")
-@click.option("-p", "--pipeline", "pipeline", type=str, help="Name of pipeline to run", default="__default__")
-@click.option("-o", "--output", type=str, default="pipeline.yml", help="Pipeline YAML definition file.")
-@click.option("--image-pull-policy", type=str, default="IfNotPresent", help="Image pull policy.")
+@click.option(
+    "-i",
+    "--image",
+    type=str,
+    help="Docker image to use for pipeline execution.",
+)
+@click.option(
+    "-p",
+    "--pipeline",
+    "pipeline",
+    type=str,
+    help="Name of pipeline to run",
+    default="__default__",
+)
+@click.option(
+    "-o",
+    "--output",
+    type=str,
+    default="pipeline.yml",
+    help="Pipeline YAML definition file.",
+)
+@click.option(
+    "--image-pull-policy",
+    type=str,
+    default="IfNotPresent",
+    help="Image pull policy.",
+)
 @click.pass_context
 def compile(ctx, image, pipeline, output, image_pull_policy) -> None:
     """Translates Kedro pipeline into YAML file with Kubeflow Pipeline definition"""
-    conf = ctx.obj['config']
+    setup_context(ctx)
+    conf = ctx.obj["config"]
     run_conf = conf.get("run_config", {})
-    image = image if image else run_conf['image']
-    ctx.obj['kfp_client'].compile(pipeline, image, output, image_pull_policy)
+    image = image if image else run_conf["image"]
+    ctx.obj["kfp_client"].compile(pipeline, image, output, image_pull_policy)
 
 
 @kubeflow_group.command()
-@click.option("-i", "--image", type=str, help="Docker image to use for pipeline execution.")
-@click.option("-p", "--pipeline", "pipeline", type=str, help="Name of pipeline to upload", default="__default__")
-@click.option("--image-pull-policy", type=str, default="IfNotPresent", help="Image pull policy.")
+@click.option(
+    "-i",
+    "--image",
+    type=str,
+    help="Docker image to use for pipeline execution.",
+)
+@click.option(
+    "-p",
+    "--pipeline",
+    "pipeline",
+    type=str,
+    help="Name of pipeline to upload",
+    default="__default__",
+)
+@click.option(
+    "--image-pull-policy",
+    type=str,
+    default="IfNotPresent",
+    help="Image pull policy.",
+)
 @click.pass_context
 def upload_pipeline(ctx, image, pipeline, image_pull_policy) -> None:
     """Uploads pipeline to Kubeflow server"""
-    conf = ctx.obj['config']
+    setup_context(ctx)
+    conf = ctx.obj["config"]
     run_conf = conf.get("run_config", {})
-    image = image if image else run_conf['image']
-    ctx.obj['kfp_client'].upload(pipeline, image, image_pull_policy)
+    image = image if image else run_conf["image"]
+    ctx.obj["kfp_client"].upload(pipeline, image, image_pull_policy)
 
 
 @kubeflow_group.command()
-@click.option("-c", "--cron-expression", type=str, help="Cron expression for recurring run", required=True)
-@click.option("-x", "--experiment-name", "experiment_name", type=str, help="Name of experiment associated with this run.")
+@click.option(
+    "-c",
+    "--cron-expression",
+    type=str,
+    help="Cron expression for recurring run",
+    required=True,
+)
+@click.option(
+    "-x",
+    "--experiment-name",
+    "experiment_name",
+    type=str,
+    help="Name of experiment associated with this run.",
+)
 @click.pass_context
 def schedule(ctx, experiment_name: str, cron_expression: str):
     """Schedules recurring execution of latest version of the pipeline"""
-    conf = ctx.obj['config']
+    setup_context(ctx)
+    conf = ctx.obj["config"]
     run_conf = conf.get("run_config", {})
-    experiment_name = experiment_name if experiment_name else run_conf['experiment_name']
+    experiment_name = (
+        experiment_name if experiment_name else run_conf["experiment_name"]
+    )
+
+    ctx.obj["kfp_client"].schedule(experiment_name, cron_expression)
+
+
+@kubeflow_group.command()
+@click.argument("kfp_url", type=str)
+@click.option(
+    "-x",
+    "--experiment-name",
+    "experiment_name",
+    type=str,
+    help="Name of experiment associated with this run.",
+    default="default",
+)
+@click.pass_context
+def init(ctx, kfp_url: str, experiment_name: str):
+    """Initializes configuration for the plugin"""
+    image = ctx.obj["kedro_ctx"].project_path.name  # default from kedro-docker
+    config = f"""
+    |host: {kfp_url}
+    |
+    |run_config:
+    |  image: {image}
+    |  experiment_name: {ctx.obj['project_name']}
+    |  run_name: {ctx.obj['project_name']}
+    |  wait_for_completion: False
+    |  volume:
+    |    storageclass: # default
+    |    #size: 1Gi
+    |    #access_modes: [ReadWriteOnce]
+    """
+    config_path = Path.cwd().joinpath("conf/base/kubeflow.yaml")
+    with open(config_path, "w") as f:
+        f.write(strip_margin(config))
+
+    click.echo(f"Configuration generated in {config_path}")
+
 
-    ctx.obj['kfp_client'].schedule(experiment_name, cron_expression)
+@kubeflow_group.command()
+@click.argument("kubeflow_run_id", type=str)
+@click.pass_context
+def mlflow_start(ctx, kubeflow_run_id: str):
+    from kedro_mlflow.framework.context import get_mlflow_config
+    import mlflow
 
+    mlflow_conf = get_mlflow_config(ctx.obj["kedro_ctx"])
+    mlflow_conf.setup(ctx.obj["kedro_ctx"])
+    run = mlflow.start_run(
+        experiment_id=mlflow_conf.experiment.experiment_id, nested=False
+    )
+    mlflow.set_tag("kubeflow_run_id", kubeflow_run_id)
+    with open("/tmp/mlflow_run_id", "w") as f:
+        f.write(run.info.run_id)
+    click.echo(f"Started run: {run.info.run_id}")
diff --git a/kedro_kubeflow/utils.py b/kedro_kubeflow/utils.py
new file mode 100644
index 0000000..f660a94
--- /dev/null
+++ b/kedro_kubeflow/utils.py
@@ -0,0 +1,15 @@
+import re
+
+
+def strip_margin(text: str) -> str:
+    return re.sub("\n[ \t]*\\|", "\n", text).strip()
+
+
+def is_mlflow_enabled() -> bool:
+    try:
+        from kedro_mlflow.framework.context import get_mlflow_config  # NOQA
+        import mlflow  # NOQA
+
+        return True
+    except ImportError:
+        return False
diff --git a/pyproject.toml b/pyproject.toml
new file mode 100644
index 0000000..a8f43fe
--- /dev/null
+++ b/pyproject.toml
@@ -0,0 +1,2 @@
+[tool.black]
+line-length = 79
diff --git a/requirements-dev.txt b/requirements-dev.txt
index 9955dec..103ce28 100644
--- a/requirements-dev.txt
+++ b/requirements-dev.txt
@@ -1,2 +1,3 @@
-pytest
-pytest-cov
+pre-commit
+coverage
+tox
diff --git a/requirements.txt b/requirements.txt
index b9db2a8..97f1530 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -2,3 +2,4 @@ kedro
 click
 kfp
 tabulate
+kedro-mlflow
diff --git a/setup.cfg b/setup.cfg
index e21e1b8..9f804bd 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 0.1.8
+current_version = 0.1.9
 
 [bumpversion:file:setup.py]
 
diff --git a/setup.py b/setup.py
index 2a16a7e..9c0d01f 100644
--- a/setup.py
+++ b/setup.py
@@ -1,43 +1,49 @@
 """kedro_kubeflow module."""
 
 from setuptools import find_packages, setup
-import os
+
 with open("README.md") as f:
-    readme = f.read()
+    README = f.read()
 
 # Runtime Requirements.
-install_requires = ["kedro>=0.16,<=0.18", "click", "kfp", "tabulate"]
+INSTALL_REQUIRES = ["kedro>=0.16,<=0.18", "click", "kfp", "tabulate", "semver"]
 
 # Dev Requirements
-extra_require = {
+EXTRA_REQUIRE = {
+    "mlflow": ["kedro-mlflow"],
     "test": ["pytest", "pytest-cov"],
     "dev": ["pytest", "pytest-cov", "pre-commit"],
 }
 
 setup(
     name="kedro-kubeflow",
-    version="0.1.8",
+    version="0.1.9",
     description="Kedro plugin with Kubeflow support",
-    long_description=readme,
+    long_description=README,
     long_description_content_type="text/markdown",
     license="Apache Software License (Apache 2.0)",
     python_requires=">=3",
     classifiers=[
-        "Programming Language :: Python :: 3.6",
+        "Development Status :: 4 - Beta",
         "Programming Language :: Python :: 3.7",
         "Programming Language :: Python :: 3.8",
     ],
     keywords="kedro kubeflow plugin",
     author=u"Mateusz Pytel, Mariusz Strzelecki",
-    author_email="mateusz@getindata.com, mariusz.strzelecki@getindata.com",
+    author_email="mateusz@getindata.com",
     url="https://github.com/getindata/kedro-kubeflow/",
     packages=find_packages(exclude=["ez_setup", "examples", "tests", "docs"]),
     include_package_data=True,
     zip_safe=False,
-    install_requires=install_requires,
-    extras_require=extra_require,
+    install_requires=INSTALL_REQUIRES,
+    extras_require=EXTRA_REQUIRE,
     entry_points={
-        "kedro.project_commands": ["kubeflow = kedro_kubeflow.plugin:commands"],
-        "kedro.global_commands": ["kubeflow = kedro_kubeflow.plugin:commands"]
+        "kedro.project_commands": [
+            "kubeflow = kedro_kubeflow.plugin:commands"
+        ],
+        "kedro.global_commands": ["kubeflow = kedro_kubeflow.plugin:commands"],
+        "kedro.hooks": [
+            "kubeflow_mlflow_hook = kedro_kubeflow.hooks:mlflow_activate_parent_run",
+        ],
     },
 )
diff --git a/tox.ini b/tox.ini
index 262ffa1..4535d80 100644
--- a/tox.ini
+++ b/tox.ini
@@ -2,7 +2,9 @@
 envlist = py37
 
 [testenv]
-extras = test
+extras = 
+    test
+    mlflow
 commands=
     python -m pytest --cov kedro_kubeflow --cov-report xml --cov-report term-missing --ignore=venv
 


Test patch:
diff --git a/tests/test_function.py b/tests/test_function.py
index 409db61..d5c89d2 100644
--- a/tests/test_function.py
+++ b/tests/test_function.py
@@ -1,4 +1,5 @@
 """Test kedro_kubeflow functions."""
 
+
 def test_app():
     pass
diff --git a/tests/test_mod.py b/tests/test_mod.py
index f1b02e0..99a2e0e 100644
--- a/tests/test_mod.py
+++ b/tests/test_mod.py
@@ -1,16 +1,16 @@
 """Test kedro_kubeflow module."""
 
 import os
-import re
 import unittest
-from unittest.mock import patch
 from tempfile import NamedTemporaryFile
+from unittest.mock import patch
 
 import kfp
 from google.auth.exceptions import DefaultCredentialsError
+from kedro.pipeline import Pipeline, node
 
 from kedro_kubeflow.kfpclient import KubeflowClient
-from kedro.pipeline import Pipeline, node
+from kedro_kubeflow.utils import strip_margin
 
 
 def identity(input1: str):
@@ -18,10 +18,6 @@ def identity(input1: str):
 
 
 class TestKubeflowClient(unittest.TestCase):
-    @staticmethod
-    def _strip_margin(text: str) -> str:
-        return re.sub("\n[ \t]*\\|", "\n", text).strip()
-
     def create_pipeline(self):
         return Pipeline(
             [
@@ -33,13 +29,22 @@ class TestKubeflowClient(unittest.TestCase):
     def create_experiment(self, id="123"):
         return type("obj", (object,), {"id": id})
 
+    def create_empty_pipelines_list(self):
+        return type(
+            "obj",
+            (object,),
+            {"pipelines": None},
+        )
+
     def create_pipelines_list(self):
         return type(
             "obj",
             (object,),
             {
                 "pipelines": [
-                    type("obj", (object,), {"name": "somename", "id": "someid"})
+                    type(
+                        "obj", (object,), {"name": "somename", "id": "someid"}
+                    )
                 ]
             },
         )
@@ -66,7 +71,9 @@ class TestKubeflowClient(unittest.TestCase):
 
     def test_should_list_pipelines_tabularized(self):
         # given
-        self.kfp_client_mock.list_pipelines.return_value = self.create_pipelines_list()
+        self.kfp_client_mock.list_pipelines.return_value = (
+            self.create_pipelines_list()
+        )
 
         # when
         output = self.client_under_test.list_pipelines()
@@ -76,12 +83,14 @@ class TestKubeflowClient(unittest.TestCase):
         |Name      ID
         |--------  ------
         |somename  someid"""
-        self.assertEqual(output, TestKubeflowClient._strip_margin(expected_output))
+        self.assertEqual(output, strip_margin(expected_output))
 
     def test_should_run_pipeline_without_waiting(self):
         # given
         run_mock = unittest.mock.MagicMock()
-        self.kfp_client_mock.create_run_from_pipeline_func.return_value = run_mock
+        self.kfp_client_mock.create_run_from_pipeline_func.return_value = (
+            run_mock
+        )
 
         # when
         self.client_under_test.run_once(
@@ -95,7 +104,10 @@ class TestKubeflowClient(unittest.TestCase):
         # 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
+        (
+            args,
+            kwargs,
+        ) = self.kfp_client_mock.create_run_from_pipeline_func.call_args
         assert kwargs == {
             "arguments": {},
             "experiment_name": "experiment",
@@ -105,13 +117,19 @@ 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"
+        assert (
+            dsl_pipeline.ops["node1"].container.image_pull_policy
+            == "IfNotPresent"
+        )
 
     def test_should_run_pipeline_and_wait(self):
         # given
         run_mock = unittest.mock.MagicMock()
-        self.kfp_client_mock.create_run_from_pipeline_func.return_value = run_mock
+        self.kfp_client_mock.create_run_from_pipeline_func.return_value = (
+            run_mock
+        )
 
         # when
         self.client_under_test.run_once(
@@ -147,7 +165,9 @@ class TestKubeflowClient(unittest.TestCase):
         fetch_id_token_mock.return_value = "unittest-token"
 
         # when
-        self.client_under_test = KubeflowClient({"host": "http://unittest"}, None, None)
+        self.client_under_test = KubeflowClient(
+            {"host": "http://unittest"}, None, None
+        )
 
         # then
         kfp_client_mock.assert_called_with(
@@ -163,7 +183,9 @@ class TestKubeflowClient(unittest.TestCase):
         os.environ["IAP_CLIENT_ID"] = "unittest-client-id"
         fetch_id_token_mock.side_effect = DefaultCredentialsError()
 
-        with self.assertLogs("kedro_kubeflow.kfpclient", level="WARNING") as cm:
+        with self.assertLogs(
+            "kedro_kubeflow.kfpclient", level="WARNING"
+        ) as cm:
             # when
             self.client_under_test = KubeflowClient(
                 {"host": "http://unittest"}, None, None
@@ -175,11 +197,15 @@ class TestKubeflowClient(unittest.TestCase):
             )
 
         # then
-        kfp_client_mock.assert_called_with("http://unittest", existing_token=None)
+        kfp_client_mock.assert_called_with(
+            "http://unittest", existing_token=None
+        )
 
     @patch("google.oauth2.id_token.fetch_id_token")
     @patch("kedro_kubeflow.kfpclient.Client")
-    def test_should_error_on_invalid_creds(self, kfp_client_mock, fetch_id_token_mock):
+    def test_should_error_on_invalid_creds(
+        self, kfp_client_mock, fetch_id_token_mock
+    ):
         # given
         os.environ["IAP_CLIENT_ID"] = "unittest-client-id"
         fetch_id_token_mock.side_effect = Exception()
@@ -193,12 +219,16 @@ class TestKubeflowClient(unittest.TestCase):
             assert "Failed to obtain IAP access token" in cm.output[0]
 
         # then
-        kfp_client_mock.assert_called_with("http://unittest", existing_token=None)
+        kfp_client_mock.assert_called_with(
+            "http://unittest", existing_token=None
+        )
 
     def test_should_modify_pull_policy_in_run(self):
         # given
         run_mock = unittest.mock.MagicMock()
-        self.kfp_client_mock.create_run_from_pipeline_func.return_value = run_mock
+        self.kfp_client_mock.create_run_from_pipeline_func.return_value = (
+            run_mock
+        )
 
         # when
         self.client_under_test.run_once(
@@ -211,7 +241,10 @@ class TestKubeflowClient(unittest.TestCase):
         )
 
         # then
-        args, kwargs = self.kfp_client_mock.create_run_from_pipeline_func.call_args
+        (
+            args,
+            kwargs,
+        ) = self.kfp_client_mock.create_run_from_pipeline_func.call_args
         with kfp.dsl.Pipeline(None) as dsl_pipeline:
             args[0]()
 
@@ -236,7 +269,9 @@ class TestKubeflowClient(unittest.TestCase):
 
     def test_should_schedule_pipeline(self):
         # given
-        self.kfp_client_mock.get_experiment.return_value = self.create_experiment()
+        self.kfp_client_mock.get_experiment.return_value = (
+            self.create_experiment()
+        )
         self.kfp_client_mock.pipelines = unittest.mock.MagicMock()
         self.kfp_client_mock.pipelines.list_pipelines.return_value = (
             self.create_pipelines_list()
@@ -244,7 +279,8 @@ class TestKubeflowClient(unittest.TestCase):
 
         # when
         self.client_under_test.schedule(
-            experiment_name="EXPERIMENT", cron_expression="0 * * * * *",
+            experiment_name="EXPERIMENT",
+            cron_expression="0 * * * * *",
         )
 
         # then
@@ -259,8 +295,12 @@ class TestKubeflowClient(unittest.TestCase):
 
     def test_should_schedule_pipeline_and_create_experiment_if_needed(self):
         # given
-        self.kfp_client_mock.get_experiment.side_effect = Exception()
-        self.kfp_client_mock.create_experiment.return_value = self.create_experiment()
+        self.kfp_client_mock.get_experiment.side_effect = ValueError(
+            "No experiment is found with name ...."
+        )
+        self.kfp_client_mock.create_experiment.return_value = (
+            self.create_experiment()
+        )
         self.kfp_client_mock.pipelines = unittest.mock.MagicMock()
         self.kfp_client_mock.pipelines.list_pipelines.return_value = (
             self.create_pipelines_list()
@@ -268,7 +308,8 @@ class TestKubeflowClient(unittest.TestCase):
 
         # when
         self.client_under_test.schedule(
-            experiment_name="EXPERIMENT", cron_expression="0 * * * * *",
+            experiment_name="EXPERIMENT",
+            cron_expression="0 * * * * *",
         )
 
         # then
@@ -283,18 +324,21 @@ class TestKubeflowClient(unittest.TestCase):
 
     def test_should_disable_old_runs_before_schedule(self):
         # given
-        self.kfp_client_mock.get_experiment.return_value = self.create_experiment()
+        self.kfp_client_mock.get_experiment.return_value = (
+            self.create_experiment()
+        )
         self.kfp_client_mock.pipelines = unittest.mock.MagicMock()
         self.kfp_client_mock.pipelines.list_pipelines.return_value = (
             self.create_pipelines_list()
         )
-        self.kfp_client_mock.list_recurring_runs.return_value = self.create_recurring_jobs_list(
-            "someid"
+        self.kfp_client_mock.list_recurring_runs.return_value = (
+            self.create_recurring_jobs_list("someid")
         )
 
         # when
         self.client_under_test.schedule(
-            experiment_name="EXPERIMENT", cron_expression="0 * * * * *",
+            experiment_name="EXPERIMENT",
+            cron_expression="0 * * * * *",
         )
 
         # then
@@ -311,11 +355,15 @@ class TestKubeflowClient(unittest.TestCase):
     def test_should_upload_new_pipeline(self):
         # given
         self.kfp_client_mock.pipelines = unittest.mock.MagicMock()
-        self.kfp_client_mock.pipelines.list_pipelines.side_effect = Exception()
+        self.kfp_client_mock.pipelines.list_pipelines.return_value = (
+            self.create_empty_pipelines_list()
+        )
 
         # when
         self.client_under_test.upload(
-            pipeline="pipeline", image="unittest-image", image_pull_policy="Always",
+            pipeline="pipeline",
+            image="unittest-image",
+            image_pull_policy="Always",
         )
 
         # then
@@ -332,23 +380,181 @@ class TestKubeflowClient(unittest.TestCase):
 
         # when
         self.client_under_test.upload(
-            pipeline="pipeline", image="unittest-image", image_pull_policy="Always",
+            pipeline="pipeline",
+            image="unittest-image",
+            image_pull_policy="Always",
         )
 
         # then
         self.kfp_client_mock.pipeline_uploads.upload_pipeline.assert_not_called()
         self.kfp_client_mock.pipeline_uploads.upload_pipeline_version.assert_called()
 
-    def tearDown(self):
-        os.environ["IAP_CLIENT_ID"] = ""
+    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_name in ["data-volume-init", "node1", "node2"]:
+            volumes = dsl_pipeline.ops[node_name].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 test_should_add_mlflow_init_step_if_enabled(self):
+        # given
+        run_mock = unittest.mock.MagicMock()
+        self.kfp_client_mock.create_run_from_pipeline_func.return_value = (
+            run_mock
+        )
+        self.mock_mlflow(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
+
+        with kfp.dsl.Pipeline(None) as dsl_pipeline:
+            args[0]()
+
+        assert len(dsl_pipeline.ops) == 3
+        init_step = dsl_pipeline.ops["mlflow-start-run"].container
+        assert init_step.image == "unittest-image"
+        assert init_step.args == [
+            "kubeflow",
+            "mlflow-start",
+            "{{workflow.uid}}",
+        ]
+        for node_name in ["node1", "node2"]:
+            env = dsl_pipeline.ops[node_name].container.env
+            assert env[1].name == "MLFLOW_PARENT_ID"
+            assert (
+                env[1].value
+                == "{{pipelineparam:op=mlflow-start-run;name=mlflow_run_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()}},
+            "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 mock_mlflow(self, enabled=False):
+        def fakeimport(name, *args, **kw):
+            if not enabled and name == "mlflow":
+                raise ImportError
+            return self.realimport(name, *args, **kw)
+
+        __builtins__["__import__"] = fakeimport
+
+    def setUp(self):
+        self.realimport = __builtins__["__import__"]
+        self.mock_mlflow(False)
+        self.create_client({})
+
+    def tearDown(self):
+        os.environ["IAP_CLIENT_ID"] = ""