Repository: getindata/kedro-kubeflow
Pull request #40 (2021-01-15 10:50:09)

Allow effective user change
#### Description

Adds ability to change the effective user id on the container to ensure the files are written/read with the proper rights

##### PR Checklist
- [x] Tests added
- [x] [Changelog](CHANGELOG.md) updated 



Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7574dcd..5aab205 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,7 @@
 ### Changed
 
 - Increased test coverage for the CLI handling
+- Ability to change the effective user id for steps if the ownership of the volume needs it
 
 ## [0.1.10] - 2021-01-11
 
diff --git a/kedro_kubeflow/config.py b/kedro_kubeflow/config.py
index bcf6079..649d17c 100644
--- a/kedro_kubeflow/config.py
+++ b/kedro_kubeflow/config.py
@@ -15,6 +15,7 @@ run_config:
     #size: 1Gi
     #access_modes: [ReadWriteOnce]
     #skip_init: False
+    #owner: 0
 """
 
 
@@ -57,6 +58,10 @@ class VolumeConfig(Config):
     def skip_init(self):
         return self._get_or_default("skip_init", False)
 
+    @property
+    def owner(self):
+        return self._get_or_default("owner", None)
+
     def _get_prefix(self):
         return "run_config.volume."
 
diff --git a/kedro_kubeflow/kfpclient.py b/kedro_kubeflow/kfpclient.py
index e5df753..5a721e9 100644
--- a/kedro_kubeflow/kfpclient.py
+++ b/kedro_kubeflow/kfpclient.py
@@ -9,7 +9,7 @@ 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 kubernetes.client import V1EnvVar, V1SecurityContext
 from tabulate import tabulate
 
 from .utils import is_mlflow_enabled
@@ -93,6 +93,14 @@ class KubeflowClient(object):
             return jwt_token
 
     def generate_pipeline(self, pipeline, image, image_pull_policy):
+        def _customize_op(op):
+            op.container.set_image_pull_policy(image_pull_policy)
+            if self.volume_meta and self.volume_meta.owner is not None:
+                op.container.set_security_context(
+                    V1SecurityContext(run_as_user=self.volume_meta.owner)
+                )
+            return op
+
         @dsl.pipeline(
             name=self.project_name,
             description="Kubeflow pipeline for Kedro project",
@@ -122,24 +130,25 @@ class KubeflowClient(object):
             if self.volume_meta.skip_init:
                 return {"/home/kedro/data": vop.volume}
             else:
-                volume_init = dsl.ContainerOp(
-                    name="data-volume-init",
-                    image=image,
-                    command=["sh", "-c"],
-                    arguments=[
-                        " ".join(
-                            [
-                                "cp",
-                                "--verbose",
-                                "-r",
-                                "/home/kedro/data/*",
-                                "/home/kedro/datavolume",
-                            ]
-                        )
-                    ],
-                    pvolumes={"/home/kedro/datavolume": vop.volume},
+                volume_init = _customize_op(
+                    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(
@@ -155,19 +164,18 @@ class KubeflowClient(object):
             ]
 
             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
+                kfp_ops["mlflow-start-run"] = _customize_op(
+                    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"},
+                    )
                 )
                 env.append(
                     V1EnvVar(
@@ -178,16 +186,15 @@ class KubeflowClient(object):
 
             for node in node_dependencies:
                 name = _clean_name(node.name)
-                kfp_ops[node.name] = dsl.ContainerOp(
-                    name=name,
-                    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] = _customize_op(
+                    dsl.ContainerOp(
+                        name=name,
+                        image=image,
+                        command=["kedro"],
+                        arguments=["run", "--node", node.name],
+                        pvolumes=node_volumes,
+                        container_kwargs={"env": env},
+                    )
                 )
 
             return kfp_ops


Test patch:
diff --git a/tests/test_mod.py b/tests/test_mod.py
index ba68d60..2b88a74 100644
--- a/tests/test_mod.py
+++ b/tests/test_mod.py
@@ -437,11 +437,15 @@ class TestKubeflowClient(unittest.TestCase):
         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.security_context is None
         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"
+            assert (
+                dsl_pipeline.ops[node_name].container.security_context is None
+            )
 
     def test_should_support_inter_steps_volume_with_given_spec(self):
         # given
@@ -490,6 +494,58 @@ class TestKubeflowClient(unittest.TestCase):
         assert volume_spec.access_modes == ["ReadWriteOnce"]
         assert volume_spec.storage_class_name == "nfs"
 
+    def test_should_change_effective_user_if_to_volume_owner(self):
+        # given
+        run_mock = unittest.mock.MagicMock()
+        self.kfp_client_mock.create_run_from_pipeline_func.return_value = (
+            run_mock
+        )
+        self.create_client(
+            {
+                "volume": {
+                    "storageclass": "nfs",
+                    "size": "1Mi",
+                    "access_modes": ["ReadWriteOnce"],
+                    "owner": 47,
+                }
+            }
+        )
+
+        # 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]()
+
+        volume_init_spec = dsl_pipeline.ops["data-volume-init"].container
+        assert volume_init_spec.security_context.run_as_user == 47
+        for node_name in ["data-volume-init", "node1", "node2"]:
+            assert (
+                dsl_pipeline.ops[
+                    node_name
+                ].container.security_context.run_as_user
+                == 47
+            )
+
     def test_should_add_mlflow_init_step_if_enabled(self):
         # given
         run_mock = unittest.mock.MagicMock()