feat(sdk): Enable setting OwnerReference on ResourceOps. Fixes kubefl…

…ow#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] kubeflow#4498 [3] kubeflow#3537 [4] kubeflow#1779 Signed-off-by: Ilias Katsakioris <elikatsis@arrikto.com>
arrikto · Nov 26, 2020 · 7119f59 · 7119f59
1 parent 4fe4a30
commit 7119f59
Show file tree

Hide file tree

Showing 9 changed files with 20 additions and 1 deletion.
diff --git a/sdk/python/kfp/dsl/_resource_op.py b/sdk/python/kfp/dsl/_resource_op.py
@@ -32,20 +32,23 @@ class Resource(object):
         "merge_strategy": "str",
         "success_condition": "str",
         "failure_condition": "str",
+        "set_owner_reference": "bool",
         "manifest": "str"
     }
     openapi_types = {
         "action": "str",
         "merge_strategy": "str",
         "success_condition": "str",
         "failure_condition": "str",
+        "set_owner_reference": "bool",
         "manifest": "str"
     }
     attribute_map = {
         "action": "action",
         "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)

diff --git a/sdk/python/kfp/dsl/_volume_op.py b/sdk/python/kfp/dsl/_volume_op.py
@@ -75,6 +75,9 @@ def __init__(self,
         # Add size to attribute outputs
         self.attribute_outputs = {"size": "{.status.capacity.storage}"}
 
+        if "set_owner_reference" not in kwargs:
+            kwargs["set_owner_reference"] = True
+
         if "k8s_resource" in kwargs:
             if resource_name or size or storage_class or modes or annotations:
                 raise ValueError("You cannot provide k8s_resource along with "

diff --git a/sdk/python/tests/compiler/testdata/param_substitutions.yaml b/sdk/python/tests/compiler/testdata/param_substitutions.yaml
@@ -44,6 +44,7 @@ spec:
       manifest: "apiVersion: v1\nkind: PersistentVolumeClaim\nmetadata:\n  name: '{{workflow.name}}-data'\n\
         spec:\n  accessModes:\n  - ReadWriteMany\n  resources:\n    requests:\n  \
         \    storage: 1Gi\n"
+      setOwnerReference: true
   - dag:
       tasks:
       - arguments:

diff --git a/sdk/python/tests/compiler/testdata/volume_snapshotop_rokurl.yaml b/sdk/python/tests/compiler/testdata/volume_snapshotop_rokurl.yaml
@@ -75,6 +75,7 @@ spec:
         \    rok/origin: '{{inputs.parameters.rok_url}}'\n  name: '{{workflow.name}}-vol1'\n\
         spec:\n  accessModes:\n  - ReadWriteMany\n  resources:\n    requests:\n  \
         \    storage: 1Gi\n"
+      setOwnerReference: true
   - inputs:
       parameters:
       - name: create-snapshot-1-name
@@ -97,6 +98,7 @@ spec:
         spec:\n  accessModes:\n  - ReadWriteMany\n  dataSource:\n    apiGroup: snapshot.storage.k8s.io\n\
         \    kind: VolumeSnapshot\n    name: '{{inputs.parameters.create-snapshot-1-name}}'\n\
         \  resources:\n    requests:\n      storage: '{{inputs.parameters.create-snapshot-1-size}}'\n"
+      setOwnerReference: true
   - inputs:
       parameters:
       - name: create-snapshot-2-name
@@ -119,6 +121,7 @@ spec:
         spec:\n  accessModes:\n  - ReadWriteMany\n  dataSource:\n    apiGroup: snapshot.storage.k8s.io\n\
         \    kind: VolumeSnapshot\n    name: '{{inputs.parameters.create-snapshot-2-name}}'\n\
         \  resources:\n    requests:\n      storage: '{{inputs.parameters.create-snapshot-2-size}}'\n"
+      setOwnerReference: true
   - container:
       args:
       - cat /data/file*| gzip -c >/data/full.gz

diff --git a/sdk/python/tests/compiler/testdata/volume_snapshotop_sequential.yaml b/sdk/python/tests/compiler/testdata/volume_snapshotop_sequential.yaml
@@ -29,6 +29,7 @@ spec:
       manifest: "apiVersion: v1\nkind: PersistentVolumeClaim\nmetadata:\n  name: '{{workflow.name}}-vol1'\n\
         spec:\n  accessModes:\n  - ReadWriteMany\n  resources:\n    requests:\n  \
         \    storage: 1Gi\n"
+      setOwnerReference: true
   - container:
       args:
       - mkdir /data/step1 && gsutil cat {{inputs.parameters.url}} | gzip -c >/data/step1/file1.gz

diff --git a/sdk/python/tests/compiler/testdata/volumeop_basic.yaml b/sdk/python/tests/compiler/testdata/volumeop_basic.yaml
@@ -50,6 +50,7 @@ spec:
       manifest: "apiVersion: v1\nkind: PersistentVolumeClaim\nmetadata:\n  name: '{{workflow.name}}-my-pvc'\n\
         spec:\n  accessModes:\n  - ReadWriteMany\n  resources:\n    requests:\n  \
         \    storage: '{{inputs.parameters.size}}'\n"
+      setOwnerReference: true
   - dag:
       tasks:
       - arguments:

diff --git a/sdk/python/tests/compiler/testdata/volumeop_dag.yaml b/sdk/python/tests/compiler/testdata/volumeop_dag.yaml
@@ -28,6 +28,7 @@ spec:
       manifest: "apiVersion: v1\nkind: PersistentVolumeClaim\nmetadata:\n  name: '{{workflow.name}}-my-pvc'\n\
         spec:\n  accessModes:\n  - ReadWriteMany\n  resources:\n    requests:\n  \
         \    storage: 10Gi\n"
+      setOwnerReference: true
   - container:
       args:
       - echo 1 | tee /mnt/file1

diff --git a/sdk/python/tests/compiler/testdata/volumeop_parallel.yaml b/sdk/python/tests/compiler/testdata/volumeop_parallel.yaml
@@ -28,6 +28,7 @@ spec:
       manifest: "apiVersion: v1\nkind: PersistentVolumeClaim\nmetadata:\n  name: '{{workflow.name}}-my-pvc'\n\
         spec:\n  accessModes:\n  - ReadWriteMany\n  resources:\n    requests:\n  \
         \    storage: 10Gi\n"
+      setOwnerReference: true
   - container:
       args:
       - echo 1 | tee /mnt/file1

diff --git a/sdk/python/tests/compiler/testdata/volumeop_sequential.yaml b/sdk/python/tests/compiler/testdata/volumeop_sequential.yaml
@@ -28,6 +28,7 @@ spec:
       manifest: "apiVersion: v1\nkind: PersistentVolumeClaim\nmetadata:\n  name: '{{workflow.name}}-newpvc'\n\
         spec:\n  accessModes:\n  - ReadWriteMany\n  resources:\n    requests:\n  \
         \    storage: 10Gi\n"
+      setOwnerReference: true
   - container:
       args:
       - echo 1|tee /data/file1