Merge pull request #170 from zillow/yunw/system-tag-in-kfp-plugin

Add service / k8s tags as system tags

Author: cloudw

metaflow/metadata/metadata.py

@@ -511,6 +511,15 @@ def _tags(self):
             tags.append("metaflow_r_version:" + env["metaflow_r_version"])
         if "r_version_code" in env:
             tags.append("r_version:" + env["r_version_code"])
+        # KFP plugin tags
+        kfp_tags = {
+            "k8s_namespace": os.environ.get("MF_POD_NAMESPACE"),
+            "zodiac_service": os.environ.get("ZODIAC_SERVICE"),
+            "zodiac_team": os.environ.get("ZODIAC_TEAM"),
+        }
+        for key in ["k8s_namespace", "zodiac_service", "zodiac_team"]:
+            if kfp_tags[key]:
+                tags.append(f"{key}:{kfp_tags[key]}")
         return tags
 
     def _register_code_package_metadata(self, run_id, step_name, task_id, attempt):

metaflow/plugins/kfp/kfp_decorator.py

@@ -126,7 +126,12 @@ def task_pre_step(
         self.task_datastore = task_datastore
 
         # TODO: any other KFP environment variables to get and register to Metadata service?
-        meta = {"kfp-execution": run_id}
+        meta = {
+            "kfp-execution": run_id,
+            "pod_name": os.environ.get("MF_POD_NAME"),
+            "argo_workflow": os.environ.get("MF_ARGO_WORKFLOW_NAME"),
+        }
+
         entries = [
             MetaDatum(field=k, value=v, type=k, tags=[]) for k, v in meta.items()
         ]

metaflow/plugins/kfp/kfp_exit_handler.py

@@ -58,14 +58,14 @@ def email_notify(send_to):
             kfp_run_id,
         )
 
-        pod_namespace = get_env("POD_NAMESPACE", "")
+        k8s_namespace = get_env("POD_NAMESPACE", "")
         argo_workflow_name = get_env("MF_ARGO_WORKFLOW_NAME", "")
         email_body = get_env("METAFLOW_NOTIFY_EMAIL_BODY", "")
         body = (
             f"status = {status} <br/>\n"
             f"{kfp_run_url} <br/>\n"
             f"Metaflow RunId = kfp-{kfp_run_id} <br/>\n"
-            f"argo -n {pod_namespace} get {argo_workflow_name} <br/>"
+            f"argo -n {k8s_namespace} get {argo_workflow_name} <br/>"
             "<br/>"
             f"{email_body}"
         )

metaflow/plugins/kfp/kfp_metaflow_step.py

@@ -56,14 +56,7 @@ def _step_cli(
 
     input_paths = None
 
-    tags_extended: List[str] = [
-        f"--tag argo_workflow:{workflow_name}",
-        "--tag pod_name:$MF_POD_NAME",
-        "--tag pod_namespace:$MF_POD_NAMESPACE",
-        # TODO(talebz): A Metaflow plugin framework to customize tags, labels, etc.
-        "--tag zodiac_service:$ZODIAC_SERVICE",
-        "--tag zodiac_team:$ZODIAC_TEAM",
-    ]
+    tags_extended: List[str] = []
     if tags:
         tags_extended.extend("--tag %s" % tag for tag in tags)
 

metaflow/plugins/kfp/tests/flows/metadata_flow.py

@@ -37,8 +37,7 @@ def end(self):
         assert "metaflow_test" in start_step_tags
         assert f"zodiac_service:{os.environ['ZODIAC_SERVICE']}" in start_step_tags
         assert f"zodiac_team:{os.environ['ZODIAC_TEAM']}" in start_step_tags
-        assert f"argo_workflow:{os.environ['MF_ARGO_WORKFLOW_NAME']}" in start_step_tags
-        assert f"pod_namespace:{os.environ['MF_POD_NAMESPACE']}" in start_step_tags
+        assert f"k8s_namespace:{os.environ['MF_POD_NAMESPACE']}" in start_step_tags
 
         print("MetadataFlow is all done.")