Repository: getindata/kedro-kubeflow
Pull request #81 (2021-12-23 07:59:48)

Passing all KEDRO_CONFIG_ env to the nodes
#### Description

In order to pass some additional runtime context to the pipeline (for example the commit id or branch name), this change extends the environment of the nodes by all `KEDRO_CONFIG_` prefixed values

##### PR Checklist
- [x] Tests added
- [x] [Changelog](CHANGELOG.md) updated 



Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 28f2215..ccddf5b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,8 @@
 
 ## [Unreleased]
 
+-   Passing all `KEDRO_CONFIG_` environment variables to the pipeline nodes
+
 ## [0.4.5] - 2021-12-22
 
 -   Add `node_merge_strategy` alongside with `full` option to run a whole pipeline in one pod
diff --git a/kedro_kubeflow/generators/one_pod_pipeline_generator.py b/kedro_kubeflow/generators/one_pod_pipeline_generator.py
index 7077fdd..2f574a1 100644
--- a/kedro_kubeflow/generators/one_pod_pipeline_generator.py
+++ b/kedro_kubeflow/generators/one_pod_pipeline_generator.py
@@ -1,12 +1,14 @@
 import logging
-import os
 
 import kubernetes.client as k8s
 from kfp import dsl
 
-from ..auth import IAP_CLIENT_ID
 from ..utils import clean_name
-from .utils import create_params, maybe_add_params
+from .utils import (
+    create_container_environment,
+    create_params,
+    maybe_add_params,
+)
 
 
 class OnePodPipelineGenerator(object):
@@ -37,11 +39,7 @@ class OnePodPipelineGenerator(object):
         image_pull_policy,
     ) -> dsl.ContainerOp:
         kwargs = {
-            "env": [
-                k8s.V1EnvVar(
-                    name=IAP_CLIENT_ID, value=os.environ.get(IAP_CLIENT_ID, "")
-                )
-            ],
+            "env": create_container_environment(),
             "image_pull_policy": image_pull_policy,
         }
         default_resources = self.run_config.resources.get_for("__default__")
diff --git a/kedro_kubeflow/generators/pod_per_node_pipeline_generator.py b/kedro_kubeflow/generators/pod_per_node_pipeline_generator.py
index a80eac9..db0a3d0 100644
--- a/kedro_kubeflow/generators/pod_per_node_pipeline_generator.py
+++ b/kedro_kubeflow/generators/pod_per_node_pipeline_generator.py
@@ -1,6 +1,5 @@
 import contextlib
 import logging
-import os
 from typing import Dict, Set
 
 import kubernetes.client as k8s
@@ -8,9 +7,12 @@ from kedro.pipeline.node import Node
 from kfp import dsl
 from kfp.compiler._k8s_helper import sanitize_k8s_name
 
-from ..auth import IAP_CLIENT_ID
 from ..utils import clean_name, is_mlflow_enabled
-from .utils import create_params, maybe_add_params
+from .utils import (
+    create_container_environment,
+    create_params,
+    maybe_add_params,
+)
 
 
 class PodPerNodePipelineGenerator(object):
@@ -106,10 +108,7 @@ class PodPerNodePipelineGenerator(object):
             else {}
         )
 
-        iap_env_var = k8s.V1EnvVar(
-            name=IAP_CLIENT_ID, value=os.environ.get(IAP_CLIENT_ID, "")
-        )
-        nodes_env = [iap_env_var]
+        nodes_env = create_container_environment()
 
         if is_mlflow_enabled():
             kfp_ops["mlflow-start-run"] = self._customize_op(
@@ -124,7 +123,7 @@ class PodPerNodePipelineGenerator(object):
                         "mlflow-start",
                         dsl.RUN_ID_PLACEHOLDER,
                     ],
-                    container_kwargs={"env": [iap_env_var]},
+                    container_kwargs={"env": nodes_env},
                     file_outputs={"mlflow_run_id": "/tmp/mlflow_run_id"},
                 ),
                 image_pull_policy,
diff --git a/kedro_kubeflow/generators/utils.py b/kedro_kubeflow/generators/utils.py
index 296f674..a880e60 100644
--- a/kedro_kubeflow/generators/utils.py
+++ b/kedro_kubeflow/generators/utils.py
@@ -1,9 +1,13 @@
+import os
 from functools import wraps
 from inspect import Parameter, signature
 from typing import Iterable
 
+import kubernetes.client as k8s
 from kfp import dsl
 
+from ..auth import IAP_CLIENT_ID
+
 
 def maybe_add_params(kedro_parameters):
     def decorator(f):
@@ -26,3 +30,16 @@ def create_params(param_keys: Iterable[str]) -> str:
     return ",".join(
         [f"{param}:{dsl.PipelineParam(param)}" for param in param_keys]
     )
+
+
+def create_container_environment():
+    env_vars = [
+        k8s.V1EnvVar(
+            name=IAP_CLIENT_ID, value=os.environ.get(IAP_CLIENT_ID, "")
+        )
+    ]
+    for key in os.environ.keys():
+        if key.startswith("KEDRO_CONFIG_"):
+            env_vars.append(k8s.V1EnvVar(name=key, value=os.environ[key]))
+
+    return env_vars


Test patch:
diff --git a/tests/test_generator.py b/tests/test_generator.py
index d9ddcca..92f1ed0 100644
--- a/tests/test_generator.py
+++ b/tests/test_generator.py
@@ -1,5 +1,6 @@
 """Test generator"""
 
+import os
 import unittest
 from inspect import signature
 from unittest.mock import MagicMock
@@ -325,6 +326,32 @@ class TestGenerator(unittest.TestCase):
         # then
         assert "schedule-volume-termination" not in dsl_pipeline.ops
 
+    def test_should_pass_kedro_config_env_to_nodes(self):
+        # given
+        self.create_generator()
+        os.environ["KEDRO_CONFIG_MY_KEY"] = "42"
+        os.environ["SOME_VALUE"] = "100"
+
+        try:
+            # when
+            with kfp.dsl.Pipeline(None) as dsl_pipeline:
+                self.generator_under_test.generate_pipeline(
+                    "pipeline", "unittest-image", "Always"
+                )()
+
+            # then
+            for node_name in ["node1", "node2"]:
+                env_values = {
+                    e.name: e.value
+                    for e in dsl_pipeline.ops[node_name].container.env
+                }
+                assert "KEDRO_CONFIG_MY_KEY" in env_values
+                assert env_values["KEDRO_CONFIG_MY_KEY"] == "42"
+                assert "SOME_VALUE" not in env_values
+        finally:
+            del os.environ["KEDRO_CONFIG_MY_KEY"]
+            del os.environ["SOME_VALUE"]
+
     def create_generator(self, config={}, params={}, catalog={}):
         project_name = "my-awesome-project"
         config_loader = MagicMock()
diff --git a/tests/test_one_pod_pipeline_generator.py b/tests/test_one_pod_pipeline_generator.py
index 96e6159..5ce9de7 100644
--- a/tests/test_one_pod_pipeline_generator.py
+++ b/tests/test_one_pod_pipeline_generator.py
@@ -1,5 +1,6 @@
 """Test generator"""
 
+import os
 import unittest
 from inspect import signature
 from unittest.mock import MagicMock
@@ -152,6 +153,31 @@ class TestGenerator(unittest.TestCase):
         # then
         assert dsl_pipeline.ops["pipeline"].file_outputs == {}
 
+    def test_should_pass_kedro_config_env_to_nodes(self):
+        # given
+        self.create_generator(params={"param1": 0.3, "param2": 42})
+        os.environ["KEDRO_CONFIG_MY_KEY"] = "42"
+        os.environ["SOME_VALUE"] = "100"
+
+        try:
+            # when
+            with kfp.dsl.Pipeline(None) as dsl_pipeline:
+                self.generator_under_test.generate_pipeline(
+                    "pipeline", "unittest-image", "Always"
+                )()
+
+            # then
+            env_values = {
+                e.name: e.value
+                for e in dsl_pipeline.ops["pipeline"].container.env
+            }
+            assert "KEDRO_CONFIG_MY_KEY" in env_values
+            assert env_values["KEDRO_CONFIG_MY_KEY"] == "42"
+            assert "SOME_VALUE" not in env_values
+        finally:
+            del os.environ["KEDRO_CONFIG_MY_KEY"]
+            del os.environ["SOME_VALUE"]
+
     def create_generator(self, config=None, params=None, catalog=None):
         if config is None:
             config = {}