sdk(feat): add metadata support for loop (kubeflow#1028)

Author: Tomcli

sdk/python/kfp_tekton/compiler/compiler.py

@@ -461,6 +461,15 @@ def process_pipelineparam(s):
         "kind": "PipelineLoop",
         "name": group_name
       }
+      # Handle sub-pipeline metadata
+      if hasattr(sub_group, 'pod_annotations') and sub_group.pod_annotations:
+        self.loops_pipeline[group_name]['spec']['taskRef']['metadata'] = \
+          self.loops_pipeline[group_name]['spec']['taskRef'].setdefault('metadata', {'annotations': {}})
+        self.loops_pipeline[group_name]['spec']['taskRef']['metadata']['annotations'] = sub_group.pod_annotations
+      if hasattr(sub_group, 'pod_labels') and sub_group.pod_annotations:
+        self.loops_pipeline[group_name]['spec']['taskRef']['metadata'] = \
+          self.loops_pipeline[group_name]['spec']['taskRef'].setdefault('metadata', {'labels': {}})
+        self.loops_pipeline[group_name]['spec']['taskRef']['metadata']['labels'] = sub_group.pod_labels
       if sub_group.items_is_pipeline_param:
         # these loop args are a 'dynamic param' rather than 'static param'.
         # i.e., rather than a static list, they are either the output of another task or were input

sdk/python/kfp_tekton/tekton.py

@@ -257,6 +257,26 @@ def range(cls,
             parallelism: Optional[int] = None):
     return cls(start=start, step=step, end=end, parallelism=parallelism)
 
+  def add_pod_annotation(self, name: str, value: str):
+    """Adds a pod's metadata annotation.
+    Args:
+        name: The name of the annotation.
+        value: The value of the annotation.
+    """
+
+    self.pod_annotations[name] = value
+    return self
+
+  def add_pod_label(self, name: str, value: str):
+    """Adds a pod's metadata label.
+    Args:
+        name: The name of the label.
+        value: The value of the label.
+    """
+
+    self.pod_labels[name] = value
+    return self
+
   def _next_id(self):
     return str(_pipeline.Pipeline.get_default_pipeline().get_next_group_id())
 
@@ -274,6 +294,8 @@ def __init__(self,
     self.step = None
     self.call_enumerate = False
     self.iteration_number = None
+    self.pod_annotations = {}
+    self.pod_labels = {}
     if start and end:
         super().__init__(loop_args=["iteration"], parallelism=parallelism)
         self.start = start

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

@@ -47,7 +47,9 @@ def empty(self):
 def pipeline(my_pipe_param: int = 10, start: int = 1, end: int = 5, step: int = 3):
     start_2 = 1
     end_2 = 2
-    with Loop.range(start=start, end=end, step=step) as item:
+    range_object = Loop.range(start=start, end=end, step=step)\
+                    .add_pod_annotation('test_anno', 'test').add_pod_label('test_label', 'test')
+    with range_object 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:

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

@@ -89,6 +89,11 @@ spec:
         apiVersion: custom.tekton.dev/v1alpha1
         kind: PipelineLoop
         name: my-pipeline-for-loop-2
+        metadata:
+          annotations:
+            test_anno: test
+          labels:
+            test_label: test
       params:
       - name: from
         value: $(params.start)