Fixed handling of volumes

kubeflow · Sep 20, 2019 · 29602d6 · 29602d6
1 parent 0595ac5
commit 29602d6
Show file tree

Hide file tree

Showing 8 changed files with 154 additions and 17 deletions.
diff --git a/sdk/python/kfp/compiler/_data_passing_rewriter.py b/sdk/python/kfp/compiler/_data_passing_rewriter.py
@@ -28,6 +28,7 @@ def fix_big_data_passing(workflow: dict) -> dict:
 
     workflow = copy.deepcopy(workflow)
     templates = workflow['spec']['templates']
+    volume_map = {volume['name']: volume for volume in workflow['spec'].get('volumes', [])}
 
     container_templates = [template for template in workflow['spec']['templates'] if 'container' in template]
     dag_templates = [template for template in workflow['spec']['templates'] if 'dag' in template]
@@ -156,6 +157,20 @@ def fix_big_data_passing(workflow: dict) -> dict:
 
     # Searching for parameter input consumers in container and resource templates
     for template in container_templates + resource_templates:
+        # Hack for volumes: Prior to Argo 2.3.0, it was only possible to add volumes globally on the workflow level instead of at template level.
+        # When VolumeOps were added, the generated structures became hacky since the workflow-level volumes cound now contain template-level input references.
+        # Here we virtually add volumes back to the template to analyze the input references.
+        # To properly fix this the compiler should move volumes to templates.
+        # TODO: Fix the compiler and remove this hack
+        if 'container' in template:
+            template = copy.deepcopy(template)
+            template_volumes = template.setdefault('volumes', [])
+            for volume_mount in template['container'].get('volumeMounts', []):
+                volume_name = volume_mount['name']
+                if volume_name in volume_map:
+                    template_volumes.append(volume_map[volume_name])
+        # End hack
+
         template_name = template['name']
         placeholders = extract_all_placeholders(template)
         for placeholder in placeholders:

diff --git a/sdk/python/tests/compiler/testdata/resourceop_basic.yaml b/sdk/python/tests/compiler/testdata/resourceop_basic.yaml
@@ -24,6 +24,9 @@ spec:
       volumeMounts:
       - mountPath: /etc/secret-volume
         name: my-secret
+    inputs:
+      parameters:
+      - name: create-my-secret-name
     name: cop
   - inputs:
       parameters:
@@ -45,7 +48,11 @@ spec:
         \ generateName: my-secret-\ntype: Opaque\n"
   - dag:
       tasks:
-      - dependencies:
+      - arguments:
+          parameters:
+          - name: create-my-secret-name
+            value: '{{tasks.create-my-secret.outputs.parameters.create-my-secret-name}}'
+        dependencies:
         - create-my-secret
         name: cop
         template: cop

diff --git a/sdk/python/tests/compiler/testdata/volume_snapshotop_rokurl.yaml b/sdk/python/tests/compiler/testdata/volume_snapshotop_rokurl.yaml
@@ -129,6 +129,9 @@ spec:
       volumeMounts:
       - mountPath: /data
         name: create-volume-1
+    inputs:
+      parameters:
+      - name: create-volume-1-name
     name: step1-concat
   - container:
       command:
@@ -139,6 +142,9 @@ spec:
       volumeMounts:
       - mountPath: /data
         name: create-volume-2
+    inputs:
+      parameters:
+      - name: create-volume-2-name
     name: step2-gunzip
   - container:
       command:
@@ -148,6 +154,9 @@ spec:
       volumeMounts:
       - mountPath: /data
         name: create-volume-3
+    inputs:
+      parameters:
+      - name: create-volume-3-name
     name: step3-output
   - dag:
       tasks:
@@ -195,15 +204,27 @@ spec:
         - create-snapshot-2
         name: create-volume-3
         template: create-volume-3
-      - dependencies:
+      - arguments:
+          parameters:
+          - name: create-volume-1-name
+            value: '{{tasks.create-volume-1.outputs.parameters.create-volume-1-name}}'
+        dependencies:
         - create-volume-1
         name: step1-concat
         template: step1-concat
-      - dependencies:
+      - arguments:
+          parameters:
+          - name: create-volume-2-name
+            value: '{{tasks.create-volume-2.outputs.parameters.create-volume-2-name}}'
+        dependencies:
         - create-volume-2
         name: step2-gunzip
         template: step2-gunzip
-      - dependencies:
+      - arguments:
+          parameters:
+          - name: create-volume-3-name
+            value: '{{tasks.create-volume-3.outputs.parameters.create-volume-3-name}}'
+        dependencies:
         - create-volume-3
         name: step3-output
         template: step3-output

diff --git a/sdk/python/tests/compiler/testdata/volume_snapshotop_sequential.yaml b/sdk/python/tests/compiler/testdata/volume_snapshotop_sequential.yaml
@@ -41,6 +41,7 @@ spec:
         name: create-volume
     inputs:
       parameters:
+      - name: create-volume-name
       - name: url
     name: step1-ingest
   - inputs:
@@ -74,6 +75,9 @@ spec:
       volumeMounts:
       - mountPath: /data
         name: create-volume
+    inputs:
+      parameters:
+      - name: create-volume-name
     name: step2-gunzip
   - inputs:
       parameters:
@@ -106,6 +110,9 @@ spec:
       volumeMounts:
       - mountPath: /data
         name: create-volume
+    inputs:
+      parameters:
+      - name: create-volume-name
     name: step3-copy
   - inputs:
       parameters:
@@ -137,13 +144,18 @@ spec:
       volumeMounts:
       - mountPath: /data
         name: create-volume
+    inputs:
+      parameters:
+      - name: create-volume-name
     name: step4-output
   - dag:
       tasks:
       - name: create-volume
         template: create-volume
       - arguments:
           parameters:
+          - name: create-volume-name
+            value: '{{tasks.create-volume.outputs.parameters.create-volume-name}}'
           - name: url
             value: '{{inputs.parameters.url}}'
         dependencies:
@@ -159,7 +171,11 @@ spec:
         - step1-ingest
         name: step1-snap
         template: step1-snap
-      - dependencies:
+      - arguments:
+          parameters:
+          - name: create-volume-name
+            value: '{{tasks.create-volume.outputs.parameters.create-volume-name}}'
+        dependencies:
         - create-volume
         - step1-ingest
         name: step2-gunzip
@@ -173,7 +189,11 @@ spec:
         - step2-gunzip
         name: step2-snap
         template: step2-snap
-      - dependencies:
+      - arguments:
+          parameters:
+          - name: create-volume-name
+            value: '{{tasks.create-volume.outputs.parameters.create-volume-name}}'
+        dependencies:
         - create-volume
         - step2-gunzip
         name: step3-copy
@@ -187,7 +207,11 @@ spec:
         - step3-copy
         name: step3-snap
         template: step3-snap
-      - dependencies:
+      - arguments:
+          parameters:
+          - name: create-volume-name
+            value: '{{tasks.create-volume.outputs.parameters.create-volume-name}}'
+        dependencies:
         - create-volume
         - step3-copy
         name: step4-output

diff --git a/sdk/python/tests/compiler/testdata/volumeop_basic.yaml b/sdk/python/tests/compiler/testdata/volumeop_basic.yaml
@@ -22,6 +22,9 @@ spec:
       volumeMounts:
       - mountPath: /mnt
         name: create-pvc
+    inputs:
+      parameters:
+      - name: create-pvc-name
     name: cop
   - inputs:
       parameters:
@@ -45,7 +48,11 @@ spec:
         \    storage: '{{inputs.parameters.size}}'\n"
   - dag:
       tasks:
-      - dependencies:
+      - arguments:
+          parameters:
+          - name: create-pvc-name
+            value: '{{tasks.create-pvc.outputs.parameters.create-pvc-name}}'
+        dependencies:
         - create-pvc
         name: cop
         template: cop

diff --git a/sdk/python/tests/compiler/testdata/volumeop_dag.yaml b/sdk/python/tests/compiler/testdata/volumeop_dag.yaml
@@ -38,6 +38,9 @@ spec:
       volumeMounts:
       - mountPath: /mnt
         name: create-pvc
+    inputs:
+      parameters:
+      - name: create-pvc-name
     name: step1
   - container:
       args:
@@ -49,6 +52,9 @@ spec:
       volumeMounts:
       - mountPath: /mnt2
         name: create-pvc
+    inputs:
+      parameters:
+      - name: create-pvc-name
     name: step2
   - container:
       args:
@@ -60,20 +66,35 @@ spec:
       volumeMounts:
       - mountPath: /mnt
         name: create-pvc
+    inputs:
+      parameters:
+      - name: create-pvc-name
     name: step3
   - dag:
       tasks:
       - name: create-pvc
         template: create-pvc
-      - dependencies:
+      - arguments:
+          parameters:
+          - name: create-pvc-name
+            value: '{{tasks.create-pvc.outputs.parameters.create-pvc-name}}'
+        dependencies:
         - create-pvc
         name: step1
         template: step1
-      - dependencies:
+      - arguments:
+          parameters:
+          - name: create-pvc-name
+            value: '{{tasks.create-pvc.outputs.parameters.create-pvc-name}}'
+        dependencies:
         - create-pvc
         name: step2
         template: step2
-      - dependencies:
+      - arguments:
+          parameters:
+          - name: create-pvc-name
+            value: '{{tasks.create-pvc.outputs.parameters.create-pvc-name}}'
+        dependencies:
         - create-pvc
         - step1
         - step2

diff --git a/sdk/python/tests/compiler/testdata/volumeop_parallel.yaml b/sdk/python/tests/compiler/testdata/volumeop_parallel.yaml
@@ -38,6 +38,9 @@ spec:
       volumeMounts:
       - mountPath: /mnt
         name: create-pvc
+    inputs:
+      parameters:
+      - name: create-pvc-name
     name: step1
   - container:
       args:
@@ -49,6 +52,9 @@ spec:
       volumeMounts:
       - mountPath: /common
         name: create-pvc
+    inputs:
+      parameters:
+      - name: create-pvc-name
     name: step2
   - container:
       args:
@@ -60,20 +66,35 @@ spec:
       volumeMounts:
       - mountPath: /mnt3
         name: create-pvc
+    inputs:
+      parameters:
+      - name: create-pvc-name
     name: step3
   - dag:
       tasks:
       - name: create-pvc
         template: create-pvc
-      - dependencies:
+      - arguments:
+          parameters:
+          - name: create-pvc-name
+            value: '{{tasks.create-pvc.outputs.parameters.create-pvc-name}}'
+        dependencies:
         - create-pvc
         name: step1
         template: step1
-      - dependencies:
+      - arguments:
+          parameters:
+          - name: create-pvc-name
+            value: '{{tasks.create-pvc.outputs.parameters.create-pvc-name}}'
+        dependencies:
         - create-pvc
         name: step2
         template: step2
-      - dependencies:
+      - arguments:
+          parameters:
+          - name: create-pvc-name
+            value: '{{tasks.create-pvc.outputs.parameters.create-pvc-name}}'
+        dependencies:
         - create-pvc
         name: step3
         template: step3

diff --git a/sdk/python/tests/compiler/testdata/volumeop_sequential.yaml b/sdk/python/tests/compiler/testdata/volumeop_sequential.yaml
@@ -38,6 +38,9 @@ spec:
       volumeMounts:
       - mountPath: /data
         name: mypvc
+    inputs:
+      parameters:
+      - name: mypvc-name
     name: step1
   - container:
       args:
@@ -49,6 +52,9 @@ spec:
       volumeMounts:
       - mountPath: /data
         name: mypvc
+    inputs:
+      parameters:
+      - name: mypvc-name
     name: step2
   - container:
       command:
@@ -59,21 +65,36 @@ spec:
       volumeMounts:
       - mountPath: /mnt
         name: mypvc
+    inputs:
+      parameters:
+      - name: mypvc-name
     name: step3
   - dag:
       tasks:
       - name: mypvc
         template: mypvc
-      - dependencies:
+      - arguments:
+          parameters:
+          - name: mypvc-name
+            value: '{{tasks.mypvc.outputs.parameters.mypvc-name}}'
+        dependencies:
         - mypvc
         name: step1
         template: step1
-      - dependencies:
+      - arguments:
+          parameters:
+          - name: mypvc-name
+            value: '{{tasks.mypvc.outputs.parameters.mypvc-name}}'
+        dependencies:
         - mypvc
         - step1
         name: step2
         template: step2
-      - dependencies:
+      - arguments:
+          parameters:
+          - name: mypvc-name
+            value: '{{tasks.mypvc.outputs.parameters.mypvc-name}}'
+        dependencies:
         - mypvc
         - step2
         name: step3