SDK - Annotate pods with component_ref (#3727)

* SDK - Annotate pods with component_ref

This preserves the information about the digest of the component and the location from which the component was loaded.

* Fixed compiler tests

sdk/python/kfp/compiler/_op_to_template.py

@@ -277,6 +277,9 @@ def _op_to_template(op: BaseOp):
     if isinstance(op, dsl.ContainerOp) and op._metadata:
         template.setdefault('metadata', {}).setdefault('annotations', {})['pipelines.kubeflow.org/component_spec'] = json.dumps(op._metadata.to_dict(), sort_keys=True)
 
+    if hasattr(op, '_component_ref'):
+        template.setdefault('metadata', {}).setdefault('annotations', {})['pipelines.kubeflow.org/component_ref'] = json.dumps(op._component_ref.to_dict(), sort_keys=True)
+
     if isinstance(op, dsl.ContainerOp) and op.execution_options:
         if op.execution_options.caching_strategy.max_cache_staleness:
             template.setdefault('metadata', {}).setdefault('annotations', {})['pipelines.kubeflow.org/max_cache_staleness'] = str(op.execution_options.caching_strategy.max_cache_staleness)

sdk/python/kfp/dsl/_component_bridge.py

@@ -72,6 +72,7 @@ def _create_container_op_from_component_and_arguments(
     component_meta = copy.copy(component_spec)
     component_meta.implementation = None
     task._set_metadata(component_meta)
+    task._component_ref = component_ref
 
     if container_spec.env:
         from kubernetes import client as k8s_client

sdk/python/tests/compiler/compiler_tests.py

@@ -180,6 +180,8 @@ def test_basic_workflow(self):
 
       for workflow in golden, compiled:
         del workflow['metadata']
+        for template in workflow['spec']['templates']:
+          template.pop('metadata', None)
 
       self.maxDiff = None
       # Comment next line for generating golden yaml.
@@ -235,6 +237,8 @@ def test_composing_workflow(self):
 
       for workflow in golden, compiled:
         del workflow['metadata']
+        for template in workflow['spec']['templates']:
+          template.pop('metadata', None)
 
       self.maxDiff = None
       # Comment next line for generating golden yaml.
@@ -265,6 +269,8 @@ def test_package_compile(self):
 
       for workflow in golden, compiled:
         del workflow['metadata']
+        for template in workflow['spec']['templates']:
+          template.pop('metadata', None)
 
       self.maxDiff = None
       self.assertEqual(golden, compiled)
@@ -286,6 +292,8 @@ def _test_py_compile_zip(self, file_base_name):
 
       for workflow in golden, compiled:
         del workflow['metadata']
+        for template in workflow['spec']['templates']:
+          template.pop('metadata', None)
 
       self.maxDiff = None
       self.assertEqual(golden, compiled)
@@ -306,6 +314,8 @@ def _test_py_compile_targz(self, file_base_name):
 
       for workflow in golden, compiled:
         del workflow['metadata']
+        for template in workflow['spec']['templates']:
+          template.pop('metadata', None)
 
       self.maxDiff = None
       self.assertEqual(golden, compiled)
@@ -328,6 +338,8 @@ def _test_py_compile_yaml(self, file_base_name):
 
       for workflow in golden, compiled:
         del workflow['metadata']
+        for template in workflow['spec']['templates']:
+          template.pop('metadata', None)
 
       self.maxDiff = None
       self.assertEqual(golden, compiled)
@@ -353,6 +365,8 @@ def _test_sample_py_compile_yaml(self, file_base_name):
 
       for workflow in golden, compiled:
         del workflow['metadata']
+        for template in workflow['spec']['templates']:
+          template.pop('metadata', None)
 
       self.maxDiff = None
       self.assertEqual(golden, compiled)