Repository: getindata/kedro-kubeflow
Pull request #66 (2021-08-18 06:28:34)

Small improvements to Vertex pipeline generation
#### Description

Small improvements to Vertex pipeline generation
- passing env in kedro commands
- setting proper artifact type based on layer in catalog

##### PR Checklist
- [x] Tests added
- [x] [Changelog](CHANGELOG.md) updated 



Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index a9fd9ae..91f80c1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,10 @@
 
 ## [Unreleased]
 
+-   Passing Kedro environment and pipeline name in Vertex nodes
+-   Setting artifact type based on catalog layer in Vertex pipeline
+-   Added `pipeline` param to `schedule` in Vertex
+
 ## [0.4.0] - 2021-08-11
 
 -   Support for kedro-mlflow>=0.7
diff --git a/kedro_kubeflow/cli.py b/kedro_kubeflow/cli.py
index fd9eff3..74211d4 100644
--- a/kedro_kubeflow/cli.py
+++ b/kedro_kubeflow/cli.py
@@ -158,6 +158,14 @@ def upload_pipeline(ctx, image, pipeline) -> None:
 
 
 @kubeflow_group.command()
+@click.option(
+    "-p",
+    "--pipeline",
+    "pipeline",
+    type=str,
+    help="Name of pipeline to run",
+    default="__default__",
+)
 @click.option(
     "-c",
     "--cron-expression",
@@ -183,7 +191,11 @@ def upload_pipeline(ctx, image, pipeline) -> None:
 )
 @click.pass_context
 def schedule(
-    ctx, experiment_namespace: str, experiment_name: str, cron_expression: str
+    ctx,
+    pipeline: str,
+    experiment_namespace: str,
+    experiment_name: str,
+    cron_expression: str,
 ):
     """Schedules recurring execution of latest version of the pipeline"""
     context_helper = ctx.obj["context_helper"]
@@ -191,7 +203,7 @@ def schedule(
     experiment = experiment_name if experiment_name else config.experiment_name
 
     context_helper.kfp_client.schedule(
-        experiment, experiment_namespace, cron_expression
+        pipeline, experiment, experiment_namespace, cron_expression
     )
 
 
diff --git a/kedro_kubeflow/kfpclient.py b/kedro_kubeflow/kfpclient.py
index d845ca1..0481831 100644
--- a/kedro_kubeflow/kfpclient.py
+++ b/kedro_kubeflow/kfpclient.py
@@ -157,7 +157,9 @@ class KubeflowClient(object):
 
         return experiment.id
 
-    def schedule(self, experiment_name, experiment_namespace, cron_expression):
+    def schedule(
+        self, pipeline, experiment_name, experiment_namespace, cron_expression
+    ):
         experiment_id = self._ensure_experiment_exists(
             experiment_name, experiment_namespace
         )
diff --git a/kedro_kubeflow/vertex_ai/client.py b/kedro_kubeflow/vertex_ai/client.py
index 83a361b..9bebe35 100644
--- a/kedro_kubeflow/vertex_ai/client.py
+++ b/kedro_kubeflow/vertex_ai/client.py
@@ -34,7 +34,7 @@ class VertexAIPipelinesClient:
         """
         pipelines = self.api_client.list_jobs()["pipelineJobs"]
         return tabulate(
-            map(lambda x: [x["displayName"], x["name"]], pipelines),
+            map(lambda x: [x.get("displayName"), x["name"]], pipelines),
             headers=["Name", "ID"],
         )
 
@@ -46,6 +46,7 @@ class VertexAIPipelinesClient:
         run_name,
         wait=False,
         image_pull_policy="IfNotPresent",
+        experiment_namespace=None,
     ):
         """
         Runs the pipeline in Vertex AI Pipelines
@@ -118,14 +119,16 @@ class VertexAIPipelinesClient:
     def schedule(
         self,
         pipeline,
-        image,
+        experiment_name,
+        experiment_namespace,
         cron_expression,
         image_pull_policy="IfNotPresent",
     ):
         """
         Schedule pipeline to Vertex AI with given cron expression
         :param pipeline:
-        :param image:
+        :param experiment_name:
+        :param experiment_namespace:
         :param cron_expression:
         :param image_pull_policy:
         :return:
@@ -135,7 +138,7 @@ class VertexAIPipelinesClient:
         ) as spec_output:
             self.compile(
                 pipeline,
-                image,
+                self.run_config.image,
                 output=spec_output.name,
                 image_pull_policy=image_pull_policy,
             )
diff --git a/kedro_kubeflow/vertex_ai/generator.py b/kedro_kubeflow/vertex_ai/generator.py
index c488574..2490adf 100644
--- a/kedro_kubeflow/vertex_ai/generator.py
+++ b/kedro_kubeflow/vertex_ai/generator.py
@@ -65,7 +65,9 @@ class PipelineGenerator:
             node_dependencies = self.context.pipelines.get(
                 pipeline
             ).node_dependencies
-            kfp_ops = self._build_kfp_ops(node_dependencies, image, token)
+            kfp_ops = self._build_kfp_ops(
+                node_dependencies, image, pipeline, token
+            )
             for node, dependencies in node_dependencies.items():
                 set_dependencies(node, dependencies, kfp_ops)
 
@@ -93,11 +95,12 @@ class PipelineGenerator:
                 "`dirname {{$.outputs.parameters['output'].output_file}}`",
                 "&&",
                 "MLFLOW_TRACKING_TOKEN={{$.inputs.parameters['mlflow_tracking_token']}} "
-                "kedro kubeflow mlflow-start "
+                f"kedro kubeflow -e {self.context.env} mlflow-start "
                 "--output {{$.outputs.parameters['output'].output_file}} "
                 + self.run_config.run_name,
             ]
         )
+
         spec = ComponentSpec(
             name="mlflow-start-run",
             inputs=[InputSpec("mlflow_tracking_token", "String")],
@@ -132,6 +135,7 @@ class PipelineGenerator:
         self,
         node_dependencies: Dict[Node, Set[Node]],
         image,
+        pipeline,
         tracking_token=None,
     ) -> Dict[str, dsl.ContainerOp]:
         """Build kfp container graph from Kedro node dependencies."""
@@ -162,8 +166,13 @@ class PipelineGenerator:
                 else []
             )
 
-            kedro_command = (
-                f'kedro run {params_parameter} --node "{node.name}"'
+            kedro_command = " ".join(
+                [
+                    f"kedro run -e {self.context.env}",
+                    f"--pipeline {pipeline}",
+                    f"{params_parameter}",
+                    f'--node "{node.name}"',
+                ]
             )
             node_command = " ".join(
                 [
@@ -229,6 +238,12 @@ class PipelineGenerator:
             f"{self.run_config.experiment_name}/{self.run_config.run_name}/data"
         )
 
+    def _get_mlruns_path(self):
+        return (
+            f"{self.run_config.root}/"
+            f"{self.run_config.experiment_name}/{self.run_config.run_name}/mlruns"
+        )
+
     def _setup_volume_op(self, image):
         command = " ".join(
             [
diff --git a/kedro_kubeflow/vertex_ai/io.py b/kedro_kubeflow/vertex_ai/io.py
index 7606988..e3795ac 100644
--- a/kedro_kubeflow/vertex_ai/io.py
+++ b/kedro_kubeflow/vertex_ai/io.py
@@ -54,6 +54,15 @@ def generate_inputs(
     return input_params, input_specs
 
 
+def get_output_type(output, catalog):
+    """
+    Returns Vertex output type based on the layer in Kedro catalog
+    """
+    if catalog[output].get("layer") == "models":
+        return "Model"
+    return "Dataset"
+
+
 def generate_outputs(node: Node, catalog):
     """
     Generates outputs for a particular kedro node
@@ -66,7 +75,8 @@ def generate_outputs(node: Node, catalog):
         and ":/" not in catalog[o]["filepath"]
     }
     output_specs = [
-        structures.OutputSpec(o, "Dataset") for o in data_mapping.keys()
+        structures.OutputSpec(o, get_output_type(o, catalog))
+        for o in data_mapping.keys()
     ]
     output_copy_commands = " ".join(
         [
diff --git a/setup.py b/setup.py
index 5cddc9b..d542a0a 100644
--- a/setup.py
+++ b/setup.py
@@ -16,7 +16,7 @@ INSTALL_REQUIRES = [
 
 # Dev Requirements
 EXTRA_REQUIRE = {
-    "mlflow": ["kedro-mlflow>=0.4.1,<0.5"],
+    "mlflow": ["kedro-mlflow>=0.4.1"],
     "tests": [
         "pytest>=5.4.0, <7.0.0",
         "pytest-cov>=2.8.0, <3.0.0",


Test patch:
diff --git a/tests/test_cli.py b/tests/test_cli.py
index 070eb1a..6962e79 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -132,12 +132,14 @@ class TestPluginCLI(unittest.TestCase):
         runner = CliRunner()
 
         result = runner.invoke(
-            schedule, ["-c", "* * *", "-x", "test_experiment"], obj=config
+            schedule,
+            ["-c", "* * *", "-x", "test_experiment", "-p", "my-pipeline"],
+            obj=config,
         )
 
         assert result.exit_code == 0
         context_helper.kfp_client.schedule.assert_called_with(
-            "test_experiment", None, "* * *"
+            "my-pipeline", "test_experiment", None, "* * *"
         )
 
     @patch.object(Path, "cwd")
diff --git a/tests/test_kfpclient.py b/tests/test_kfpclient.py
index 521ecf9..4890831 100644
--- a/tests/test_kfpclient.py
+++ b/tests/test_kfpclient.py
@@ -198,6 +198,7 @@ class TestKubeflowClient(unittest.TestCase):
 
         # when
         self.client_under_test.schedule(
+            pipeline=None,
             experiment_name="EXPERIMENT",
             cron_expression="0 * * * * *",
             experiment_namespace=None,
@@ -228,6 +229,7 @@ class TestKubeflowClient(unittest.TestCase):
 
         # when
         self.client_under_test.schedule(
+            pipeline=None,
             experiment_name="EXPERIMENT",
             cron_expression="0 * * * * *",
             experiment_namespace=None,
@@ -258,6 +260,7 @@ class TestKubeflowClient(unittest.TestCase):
 
         # when
         self.client_under_test.schedule(
+            pipeline=None,
             experiment_name="EXPERIMENT",
             cron_expression="0 * * * * *",
             experiment_namespace=None,
@@ -352,4 +355,5 @@ class TestKubeflowClient(unittest.TestCase):
         self.create_client({})
 
     def tearDown(self):
+        __builtins__["__import__"] = self.realimport
         os.environ["IAP_CLIENT_ID"] = ""
diff --git a/tests/test_vertex_ai_client.py b/tests/test_vertex_ai_client.py
index 47e1764..b159897 100644
--- a/tests/test_vertex_ai_client.py
+++ b/tests/test_vertex_ai_client.py
@@ -71,6 +71,10 @@ class TestKubeflowClient(unittest.TestCase):
                         "name": "projects/29350373243/locations/"
                         "europe-west4/pipelineJobs/run2",
                     },
+                    {
+                        "name": "projects/123/locations/"
+                        "europe-west4/pipelineJobs/no-display-name",
+                    },
                 ]
             }
 
@@ -79,9 +83,10 @@ class TestKubeflowClient(unittest.TestCase):
 
             expected_output = """
             |Name    ID
-            |------  -------------------------------------------------------------
+            |------  ----------------------------------------------------------------
             |run1    projects/29350373243/locations/europe-west4/pipelineJobs/run1
-            |run2    projects/29350373243/locations/europe-west4/pipelineJobs/run2"""
+            |run2    projects/29350373243/locations/europe-west4/pipelineJobs/run2
+            |        projects/123/locations/europe-west4/pipelineJobs/no-display-name"""
             assert tabulation == strip_margin(expected_output)
 
     def test_should_schedule_pipeline(self):
diff --git a/tests/test_vertex_ai_generator.py b/tests/test_vertex_ai_generator.py
index 1f104ae..e52ab2a 100644
--- a/tests/test_vertex_ai_generator.py
+++ b/tests/test_vertex_ai_generator.py
@@ -117,7 +117,12 @@ class TestGenerator(unittest.TestCase):
                 "B": {
                     "type": "pandas.CSVDataSet",
                     "filepath": "data/02_intermediate/b.csv",
-                }
+                },
+                "C": {
+                    "type": "pickle.PickleDataSet",
+                    "filepath": "data/06_models/model.pkl",
+                    "layer": "models",
+                },
             }
         )
 
@@ -136,7 +141,7 @@ class TestGenerator(unittest.TestCase):
             name="B", op_name="node1", param_type="Dataset"
         )
         outputs2 = dsl_pipeline.ops["node2"].outputs
-        assert len(outputs2) == 0  # output "C" is missing in the catalog
+        assert outputs2["C"].param_type == "Model"
 
     def test_should_skip_volume_removal_if_requested(self):
         # given
@@ -152,6 +157,43 @@ class TestGenerator(unittest.TestCase):
         # then
         assert "schedule-volume-termination" not in dsl_pipeline.ops
 
+    def test_should_add_env_and_pipeline_in_the_invocations(self):
+        # given
+        self.create_generator()
+        self.mock_mlflow(True)
+
+        # when
+        pipeline = self.generator_under_test.generate_pipeline(
+            "pipeline", "unittest-image", "Never", "MLFLOW_TRACKING_TOKEN"
+        )
+        with kfp.dsl.Pipeline(None) as dsl_pipeline:
+            pipeline()
+
+        # then
+        assert (
+            "kedro kubeflow -e unittests mlflow-start"
+            in dsl_pipeline.ops["mlflow-start-run"].container.args[0]
+        )
+        assert (
+            'kedro run -e unittests --pipeline pipeline  --node "node1"'
+            in dsl_pipeline.ops["node1"].container.args[0]
+        )
+
+    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)
+
+    def tearDown(self):
+        __builtins__["__import__"] = self.realimport
+
     def create_generator(self, config={}, params={}, catalog={}):
         project_name = "my-awesome-project"
         config_loader = MagicMock()
@@ -160,6 +202,7 @@ class TestGenerator(unittest.TestCase):
             "obj",
             (object,),
             {
+                "env": "unittests",
                 "params": params,
                 "config_loader": config_loader,
                 "pipelines": {