Passing the annotations and labels to the ContainerOp

Currently the annotations and labels are not passed from component to the ContainerOp. This PR fixes that.

Fixes kubeflow#1013

sdk/python/kfp/components/_dsl_bridge.py

@@ -13,8 +13,9 @@
 # limitations under the License.
 
 from collections import OrderedDict
+import json
 from typing import Mapping
-from ._structures import ConcatPlaceholder, IfPlaceholder, InputValuePlaceholder, InputPathPlaceholder, IsPresentPlaceholder, OutputPathPlaceholder, TaskSpec
+from ._structures import ConcatPlaceholder, IfPlaceholder, InputValuePlaceholder, InputPathPlaceholder, IsPresentPlaceholder, OutputPathPlaceholder, ComponentSpec, TaskSpec
 from ._components import _generate_output_file_name, _default_component_name
 from kfp.dsl._metadata import ComponentMeta, ParameterMeta, TypeMeta, _annotation_to_typemeta
 
@@ -126,12 +127,13 @@ def expand_argument_list(argument_list):
         output_paths=output_paths,
         env=container_spec.env,
         component_spec=component_spec,
+        task_spec=task_spec,
     )
 
 
 _dummy_pipeline=None
 
-def _create_container_op_from_resolved_task(name:str, container_image:str, command=None, arguments=None, output_paths=None, env : Mapping[str, str]=None, component_spec=None):
+def _create_container_op_from_resolved_task(name:str, container_image:str, command=None, arguments=None, output_paths=None, env : Mapping[str, str]=None, component_spec: ComponentSpec=None, task_spec: TaskSpec=None):
     from .. import dsl
     global _dummy_pipeline
     need_dummy = dsl.Pipeline._default_pipeline is None
@@ -170,7 +172,15 @@ def _create_container_op_from_resolved_task(name:str, container_image:str, comma
         from kubernetes import client as k8s_client
         for name, value in env.items():
             task.container.add_env_variable(k8s_client.V1EnvVar(name=name, value=value))
-
+
+    if component_spec.metadata:
+        for key, value in (component_spec.metadata.annotations or {}).items():
+            task.add_pod_annotation(key, value)
+        for key, value in (component_spec.metadata.labels or {}).items():
+            task.add_pod_label(key, value)
+    task.add_pod_annotation('kubeflow.org/pipelines/component_spec', json.dumps(component_spec.to_struct()))
+    task.add_pod_annotation('kubeflow.org/pipelines/task_spec', json.dumps(task_spec.to_struct()))
+
     if need_dummy:
         _dummy_pipeline.__exit__()
 

sdk/python/tests/components/test_components.py

@@ -547,6 +547,23 @@ def test_handle_default_values_in_task_factory(self):
         task2 = task_factory1('456')
         self.assertEqual(task2.arguments, ['456'])
 
+    def test_passing_component_metadata_to_container_op(self):
+        component_text = '''\
+metadata:
+  annotations:
+    key1: value1
+  labels:
+    key1: value1
+implementation:
+  container:
+    image: busybox
+'''
+        task_factory1 = comp.load_component_from_text(text=component_text)
+
+        task1 = task_factory1()
+        self.assertEqual(task1.pod_annotations['key1'], 'value1')
+        self.assertEqual(task1.pod_labels['key1'], 'value1')
+
     def test_type_compatibility_check_for_simple_types(self):
         component_a = '''\
 outputs: