diff --git a/sdk/python/kfp/compiler/_data_passing_rewriter.py b/sdk/python/kfp/compiler/_data_passing_rewriter.py index 9cc54eb47efc..74ff74fc549c 100644 --- 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 index e5289d1b27ff..98b47af23005 100644 --- 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 index f4fe041be1b4..3ca258e30e30 100644 --- 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 index cacbd01713f7..b6ae8ef41e5f 100644 --- 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,6 +144,9 @@ spec: volumeMounts: - mountPath: /data name: create-volume + inputs: + parameters: + - name: create-volume-name name: step4-output - dag: tasks: @@ -144,6 +154,8 @@ spec: 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 index cf5ea7f723e3..6c0d2d521b0e 100644 --- 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 index e2e45a27e06d..05a8ba2035ae 100644 --- 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 index d79ee0ae4df4..8a9341a75b00 100644 --- 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 index 63b2f3536b5b..f1cf8bbbad9d 100644 --- 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