From b3382cc3bffabf7d7fa8e92a1bd5e488ec21419e Mon Sep 17 00:00:00 2001
From: Ilias Katsakioris <elikatsis@arrikto.com>
Date: Mon, 23 Nov 2020 19:49:00 +0200
Subject: [PATCH] feat(sdk): Enable setting OwnerReference on ResourceOps.
 Fixes #1779

Argo supports a field in the ResourceTemplate that makes the controller
add an owner reference of the workflow to the created resource since
v2.4.0 [1].
With the upgrade of Argo client [2] and deployment [3] we are now able
to exploit it.

We set it to 'false' by default on all ResourceOps (actually, leave it
empty) and 'true' for VolumeOps. This allows the garbage collection of
PVCs upon workflow cleanup without any further change [4].

[1] https://github.com/argoproj/argo/blob/v2.4.0/pkg/apis/workflow/v1alpha1/workflow_types.go#L1044-L1045
[2] https://github.com/kubeflow/pipelines/pull/4498
[3] https://github.com/kubeflow/pipelines/pull/3537
[4] https://github.com/kubeflow/pipelines/issues/1779

Signed-off-by: Ilias Katsakioris <elikatsis@arrikto.com>
---
 sdk/python/kfp/dsl/_resource_op.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/sdk/python/kfp/dsl/_resource_op.py b/sdk/python/kfp/dsl/_resource_op.py
index b91878aae6c2..7ef76831ebb9 100644
--- a/sdk/python/kfp/dsl/_resource_op.py
+++ b/sdk/python/kfp/dsl/_resource_op.py
@@ -32,6 +32,7 @@ class Resource(object):
         "merge_strategy": "str",
         "success_condition": "str",
         "failure_condition": "str",
+        "set_owner_reference": "bool",
         "manifest": "str"
     }
     openapi_types = {
@@ -39,6 +40,7 @@ class Resource(object):
         "merge_strategy": "str",
         "success_condition": "str",
         "failure_condition": "str",
+        "set_owner_reference": "bool",
         "manifest": "str"
     }
     attribute_map = {
@@ -46,6 +48,7 @@ class Resource(object):
         "merge_strategy": "mergeStrategy",
         "success_condition": "successCondition",
         "failure_condition": "failureCondition",
+        "set_owner_reference": "setOwnerReference",
         "manifest": "manifest"
     }
 
@@ -54,12 +57,14 @@ def __init__(self,
                  merge_strategy: str = None,
                  success_condition: str = None,
                  failure_condition: str = None,
+                 set_owner_reference: bool = None,
                  manifest: str = None):
         """Create a new instance of Resource"""
         self.action = action
         self.merge_strategy = merge_strategy
         self.success_condition = success_condition
         self.failure_condition = failure_condition
+        self.set_owner_reference = set_owner_reference
         self.manifest = manifest
 
 
@@ -95,6 +100,7 @@ def __init__(self,
                  merge_strategy: str = None,
                  success_condition: str = None,
                  failure_condition: str = None,
+                 set_owner_reference: bool = None,
                  attribute_outputs: Dict[str, str] = None,
                  **kwargs):
 
@@ -121,7 +127,8 @@ def __init__(self,
             "action": action,
             "merge_strategy": merge_strategy,
             "success_condition": success_condition,
-            "failure_condition": failure_condition
+            "failure_condition": failure_condition,
+            "set_owner_reference": set_owner_reference,
         }
         # `resource` prop in `io.argoproj.workflow.v1alpha1.Template`
         self._resource = Resource(**init_resource)