From 6cb92d45c8f6718d416f42f469f831a7afe4ec5f Mon Sep 17 00:00:00 2001
From: Alexey Volkov <alexey.volkov@ark-kun.com>
Date: Sat, 25 Apr 2020 01:49:28 -0700
Subject: [PATCH] SDK - Compiler - Include the SDK version information in the
 compiled workflows (#3583)

* SDK - Compiler - Include the SDK version information in the compiled workflows

* Fixed the unit tests

* Removed the sdk_version annotation.
---
 sdk/python/kfp/compiler/compiler.py         | 14 ++++++++++--
 sdk/python/tests/compiler/compiler_tests.py | 24 ++++++++++-----------
 2 files changed, 24 insertions(+), 14 deletions(-)

diff --git a/sdk/python/kfp/compiler/compiler.py b/sdk/python/kfp/compiler/compiler.py
index f637364d94..e596cfc3d6 100644
--- a/sdk/python/kfp/compiler/compiler.py
+++ b/sdk/python/kfp/compiler/compiler.py
@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import datetime
 import json
 from collections import defaultdict
 from deprecated import deprecated
@@ -20,6 +21,7 @@ import uuid
 import zipfile
 from typing import Callable, Set, List, Text, Dict, Tuple, Any, Union, Optional
 
+import kfp
 from kfp.dsl import _for_loop
 
 from .. import dsl
@@ -831,8 +833,16 @@ class Compiler(object):
     from ._data_passing_rewriter import fix_big_data_passing
     workflow = fix_big_data_passing(workflow)
 
-    import json
-    workflow.setdefault('metadata', {}).setdefault('annotations', {})['pipelines.kubeflow.org/pipeline_spec'] = json.dumps(pipeline_meta.to_dict(), sort_keys=True)
+    metadata = workflow.setdefault('metadata', {})
+    annotations = metadata.setdefault('annotations', {})
+
+    annotations['pipelines.kubeflow.org/kfp_sdk_version'] = kfp.__version__
+    annotations['pipelines.kubeflow.org/pipeline_compilation_time'] = datetime.datetime.now().isoformat()
+    annotations['pipelines.kubeflow.org/pipeline_spec'] = json.dumps(pipeline_meta.to_dict(), sort_keys=True)
+
+    # Labels might be logged better than annotations so adding some information here as well
+    labels = metadata.setdefault('labels', {})
+    labels['pipelines.kubeflow.org/kfp_sdk_version'] = kfp.__version__
 
     return workflow
 
diff --git a/sdk/python/tests/compiler/compiler_tests.py b/sdk/python/tests/compiler/compiler_tests.py
index 6d3671fe34..a117f3752c 100644
--- a/sdk/python/tests/compiler/compiler_tests.py
+++ b/sdk/python/tests/compiler/compiler_tests.py
@@ -178,6 +178,9 @@ class TestCompiler(unittest.TestCase):
         golden = yaml.safe_load(f)
       compiled = self._get_yaml_from_zip(package_path)
 
+      for workflow in golden, compiled:
+        del workflow['metadata']
+
       self.maxDiff = None
       # Comment next line for generating golden yaml.
       self.assertEqual(golden, compiled)
@@ -205,8 +208,7 @@ class TestCompiler(unittest.TestCase):
         golden = yaml.safe_load(f)
 
       for workflow in golden, compiled_workflow:
-        annotations = workflow['metadata']['annotations']
-        del annotations['pipelines.kubeflow.org/pipeline_spec']
+        del workflow['metadata']
 
       self.assertEqual(golden, compiled_workflow)
     finally:
@@ -232,8 +234,7 @@ class TestCompiler(unittest.TestCase):
       compiled = self._get_yaml_from_zip(compose_package_path)
 
       for workflow in golden, compiled:
-        annotations = workflow['metadata']['annotations']
-        del annotations['pipelines.kubeflow.org/pipeline_spec']
+        del workflow['metadata']
 
       self.maxDiff = None
       # Comment next line for generating golden yaml.
@@ -263,8 +264,7 @@ class TestCompiler(unittest.TestCase):
       compiled = self._get_yaml_from_zip(target_zip)
 
       for workflow in golden, compiled:
-        annotations = workflow['metadata']['annotations']
-        del annotations['pipelines.kubeflow.org/pipeline_spec']
+        del workflow['metadata']
 
       self.maxDiff = None
       self.assertEqual(golden, compiled)
@@ -285,8 +285,7 @@ class TestCompiler(unittest.TestCase):
       compiled = self._get_yaml_from_zip(target_zip)
 
       for workflow in golden, compiled:
-        annotations = workflow['metadata']['annotations']
-        del annotations['pipelines.kubeflow.org/pipeline_spec']
+        del workflow['metadata']
 
       self.maxDiff = None
       self.assertEqual(golden, compiled)
@@ -306,8 +305,7 @@ class TestCompiler(unittest.TestCase):
       compiled = self._get_yaml_from_tar(target_tar)
 
       for workflow in golden, compiled:
-        annotations = workflow['metadata']['annotations']
-        del annotations['pipelines.kubeflow.org/pipeline_spec']
+        del workflow['metadata']
 
       self.maxDiff = None
       self.assertEqual(golden, compiled)
@@ -329,8 +327,7 @@ class TestCompiler(unittest.TestCase):
         compiled = yaml.safe_load(f)
       
       for workflow in golden, compiled:
-        annotations = workflow['metadata']['annotations']
-        del annotations['pipelines.kubeflow.org/pipeline_spec']
+        del workflow['metadata']
 
       self.maxDiff = None
       self.assertEqual(golden, compiled)
@@ -354,6 +351,9 @@ class TestCompiler(unittest.TestCase):
       with open(os.path.join(test_data_dir, target_yaml), 'r') as f:
         compiled = yaml.safe_load(f)
 
+      for workflow in golden, compiled:
+        del workflow['metadata']
+
       self.maxDiff = None
       self.assertEqual(golden, compiled)
     finally: