feat(sdk): add iterateParamPassStyle and itemPassStyle params to loop (kubeflow#1059)

* add iterateParamPassStyle and itemPassStyle params to loop

* update new fields to store in extra configs

* update dsl to take extra_fields as new arguments

* make extra field values configurable

* update error message

Author: Tomcli

sdk/python/kfp_tekton/compiler/compiler.py

@@ -573,6 +573,23 @@ def process_parameter(parameter):
       if sub_group.parallelism is not None and sub_group.parallelism > 0:
         self.loops_pipeline[group_name]['spec']['parallelism'] = sub_group.parallelism
 
+      def insert_extra_config_field(config_name, config_object, extra_field_name):
+        # Default allowed values
+        config_value_list = ['inline', 'file']
+        config_value = config_object.lower()
+        # Update the list of allowed values if exist
+        if hasattr(sub_group, 'config_value_list'):
+          config_value_list = sub_group.config_value_list.get(extra_field_name, config_value_list)
+        if config_value in config_value_list:
+          self.loops_pipeline[group_name]['spec'][config_name] = config_value
+        else:
+          raise ValueError("%s value in loop %s must be one of [%s], not %s" %
+                           (config_name, group_name, ",".join(config_value_list), config_value))
+      if hasattr(sub_group, 'iterate_param_pass_style') and sub_group.iterate_param_pass_style is not None:
+        insert_extra_config_field('iterateParamPassStyle', sub_group.iterate_param_pass_style, 'iterate_param_pass_style')
+      if hasattr(sub_group, 'item_pass_style') and sub_group.item_pass_style is not None:
+        insert_extra_config_field('itemPassStyle', sub_group.item_pass_style, 'item_pass_style')
+
     return template
 
   def _create_dag_templates(self, pipeline, op_transformers=None, params=None, op_to_templates_handler=None):

sdk/python/kfp_tekton/tekton.py

@@ -239,23 +239,41 @@ class Loop(dsl.ParallelFor):
 
   @classmethod
   def sequential(cls,
-                 loop_args: _for_loop.ItemList):
-    return cls(loop_args=loop_args, parallelism=1)
+                 loop_args: _for_loop.ItemList,
+                 extra_fields: Optional[dict] = None,
+                 valid_extra_field_values: Optional[dict] = None):
+    return cls(loop_args=loop_args,
+               parallelism=1,
+               extra_fields=extra_fields,
+               valid_extra_field_values=valid_extra_field_values)
 
   @classmethod
   def from_string(cls,
                   loop_args: Union[str, _pipeline_param.PipelineParam],
                   separator: Optional[Union[str, _pipeline_param.PipelineParam]] = None,
-                  parallelism: Optional[int] = None):
-    return cls(loop_args=loop_args, separator=separator, parallelism=parallelism)
+                  parallelism: Optional[int] = None,
+                  extra_fields: Optional[dict] = None,
+                  valid_extra_field_values: Optional[dict] = None):
+    return cls(loop_args=loop_args,
+               separator=separator,
+               parallelism=parallelism,
+               extra_fields=extra_fields,
+               valid_extra_field_values=valid_extra_field_values)
 
   @classmethod
   def range(cls,
             start: Union[_Num, PipelineParam],
             end: Union[_Num, PipelineParam],
             step: Optional[Union[_Num, PipelineParam]] = None,
-            parallelism: Optional[int] = None):
-    return cls(start=start, step=step, end=end, parallelism=parallelism)
+            parallelism: Optional[int] = None,
+            extra_fields: Optional[dict] = None,
+            valid_extra_field_values: Optional[dict] = None):
+    return cls(start=start,
+               step=step,
+               end=end,
+               parallelism=parallelism,
+               extra_fields=extra_fields,
+               valid_extra_field_values=valid_extra_field_values)
 
   def add_pod_annotation(self, name: str, value: str):
     """Adds a pod's metadata annotation.
@@ -291,14 +309,31 @@ def __init__(self,
                end: Union[_Num, PipelineParam, None] = None,
                step: Union[_Num, PipelineParam, None] = None,
                separator: Optional[Union[str, _pipeline_param.PipelineParam]] = None,
-               parallelism: Optional[int] = None):
+               parallelism: Optional[int] = None,
+               extra_fields: Optional[dict] = None,
+               valid_extra_field_values: Optional[dict] = None):
     self.start = None
     self.end = None
     self.step = None
     self.call_enumerate = False
     self.iteration_number = None
     self.pod_annotations = {}
     self.pod_labels = {}
+    self.iterate_param_pass_style = None
+    self.item_pass_style = None
+    self.config_value_list = {"iterate_param_pass_style": ["inline", "file"],
+                              "item_pass_style": ["inline", "file"]}
+    if extra_fields:
+        if extra_fields.get('iterate_param_pass_style'):
+            self.iterate_param_pass_style = extra_fields['iterate_param_pass_style']
+        if extra_fields.get('item_pass_style'):
+            self.item_pass_style = extra_fields['item_pass_style']
+
+    # Update allowed values in the extra fields
+    if valid_extra_field_values:
+        for k, v in valid_extra_field_values.items():
+            self.config_value_list[k] = v
+
     if start and end:
         super().__init__(loop_args=["iteration"], parallelism=parallelism)
         self.start = start

sdk/python/tests/compiler/testdata/loop_with_numeric.py

@@ -47,10 +47,12 @@ def empty(self):
 def pipeline(my_pipe_param: int = 10, start: int = 1, end: int = 2):
     start_2 = 1
     end_2 = 2
-    with Loop.range(start=start, end=end) as item:
+    iterate_param_pass_style_field = {'iterate_param_pass_style': 'inline'}
+    item_pass_style_field = {'item_pass_style': 'file'}
+    with Loop.range(start=start, end=end, extra_fields=iterate_param_pass_style_field) as item:
         op1_template = components.load_component_from_text(op1_yaml)
         op1_template(item, my_pipe_param)
-        with Loop.range(start=start_2, end=end_2) as item2:
+        with Loop.range(start=start_2, end=end_2, extra_fields=item_pass_style_field) as item2:
             op1_template(item2, my_pipe_param)
 
 

sdk/python/tests/compiler/testdata/loop_with_numeric.yaml

@@ -58,6 +58,7 @@ spec:
         value: $(params.my_pipe_param)
       - name: to
         value: $(params.end)
+      iterateParamPassStyle: inline
       taskSpec:
         apiVersion: custom.tekton.dev/v1alpha1
         kind: PipelineLoop
@@ -107,6 +108,7 @@ spec:
                 value: $(params.my_pipe_param)
               - name: to
                 value: '2'
+              itemPassStyle: file
               taskSpec:
                 apiVersion: custom.tekton.dev/v1alpha1
                 kind: PipelineLoop

sdk/python/tests/compiler/testdata/loop_with_numeric_noninlined.yaml

@@ -42,16 +42,17 @@ metadata:
       "loop-item-param-1", "type": "string"}, {"name": "my_pipe_param", "type": "string"}],
       "steps": [{"args": ["set -e\necho op1 \"$0\" \"$1\"\n", "$(inputs.params.loop-item-param-1)",
       "$(inputs.params.my_pipe_param)"], "command": ["sh", "-c"], "image": "library/bash:4.4.23",
-      "name": "main"}]}}, {"name": "my-pipeline-for-loop-4", "params": [{"name": "from",
-      "value": "1"}, {"name": "my_pipe_param", "value": "$(params.my_pipe_param)"},
-      {"name": "to", "value": "2"}], "taskRef": {"apiVersion": "custom.tekton.dev/v1alpha1",
-      "kind": "PipelineLoop", "name": "my-pipeline-for-loop-4"}}]}}}, {"apiVersion":
-      "custom.tekton.dev/v1alpha1", "kind": "PipelineLoop", "metadata": {"name": "my-pipeline-for-loop-4"},
-      "spec": {"iterateNumeric": "loop-item-param-3", "pipelineSpec": {"params": [{"name":
-      "loop-item-param-3", "type": "string"}, {"name": "my_pipe_param", "type": "string"}],
-      "tasks": [{"name": "my-in-coop1-2", "params": [{"name": "loop-item-param-3",
-      "value": "$(params.loop-item-param-3)"}, {"name": "my_pipe_param", "value":
-      "$(params.my_pipe_param)"}], "taskSpec": {"metadata": {"annotations": {"pipelines.kubeflow.org/component_spec_digest":
+      "name": "main"}]}}, {"itemPassStyle": "file", "name": "my-pipeline-for-loop-4",
+      "params": [{"name": "from", "value": "1"}, {"name": "my_pipe_param", "value":
+      "$(params.my_pipe_param)"}, {"name": "to", "value": "2"}], "taskRef": {"apiVersion":
+      "custom.tekton.dev/v1alpha1", "kind": "PipelineLoop", "name": "my-pipeline-for-loop-4"}}]}}},
+      {"apiVersion": "custom.tekton.dev/v1alpha1", "kind": "PipelineLoop", "metadata":
+      {"name": "my-pipeline-for-loop-4"}, "spec": {"iterateNumeric": "loop-item-param-3",
+      "pipelineSpec": {"params": [{"name": "loop-item-param-3", "type": "string"},
+      {"name": "my_pipe_param", "type": "string"}], "tasks": [{"name": "my-in-coop1-2",
+      "params": [{"name": "loop-item-param-3", "value": "$(params.loop-item-param-3)"},
+      {"name": "my_pipe_param", "value": "$(params.my_pipe_param)"}], "taskSpec":
+      {"metadata": {"annotations": {"pipelines.kubeflow.org/component_spec_digest":
       "{\"name\": \"my-in-coop1\", \"outputs\": [], \"version\": \"my-in-coop1@sha256=42f0c6a60c7e435441b8afaeb382a771a9741fe3aabb203748fdbd72b25f1628\"}"},
       "labels": {"pipelines.kubeflow.org/cache_enabled": "true"}}, "params": [{"name":
       "loop-item-param-3", "type": "string"}, {"name": "my_pipe_param", "type": "string"}],
@@ -90,3 +91,4 @@ spec:
         value: $(params.my_pipe_param)
       - name: to
         value: $(params.end)
+      iterateParamPassStyle: inline