Metadata Writer - Fixed pod name property setting (kubeflow#3563)

RedbackThomson · Apr 28, 2020 · c019b6c · c019b6c
1 parent 6f94d26
commit c019b6c
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 6 deletions.
diff --git a/backend/metadata_writer/src/metadata_helpers.py b/backend/metadata_writer/src/metadata_helpers.py
@@ -238,6 +238,8 @@ def create_new_execution_in_existing_context(
 EXECUTION_RUN_ID_PROPERTY_NAME = "run_id"
 CONTEXT_RUN_ID_PROPERTY_NAME = "run_id"
 
+KFP_POD_NAME_EXECUTION_PROPERTY_NAME = 'kfp_pod_name'
+
 
 def get_or_create_run_context(
     store,
@@ -263,6 +265,7 @@ def create_new_execution_in_existing_run_context(
     store,
     execution_type_name: str,
     context_id: int,
+    pod_name: str,
     # TODO: Remove when UX stops relying on thsese properties
     pipeline_name: str = None,
     run_id: str = None,
@@ -287,7 +290,9 @@ def create_new_execution_in_existing_run_context(
             EXECUTION_RUN_ID_PROPERTY_NAME: metadata_store_pb2.Value(string_value=run_id),
             EXECUTION_COMPONENT_ID_PROPERTY_NAME: metadata_store_pb2.Value(string_value=instance_id), # should set to task ID, not component ID
         },
-        custom_properties=custom_properties,
+        custom_properties={
+            KFP_POD_NAME_EXECUTION_PROPERTY_NAME: metadata_store_pb2.Value(string_value=pod_name),
+        },
     )
 
 

diff --git a/backend/metadata_writer/src/metadata_writer.py b/backend/metadata_writer/src/metadata_writer.py
@@ -76,8 +76,6 @@ def patch_pod_metadata(
 ARGO_COMPLETED_LABEL_KEY = 'workflows.argoproj.io/completed'
 METADATA_WRITTEN_LABEL_KEY = 'pipelines.kubeflow.org/metadata_written'
 
-KFP_POD_NAME_EXECUTION_PROPERTY_NAME = 'kfp_pod_name'
-
 
 def output_name_to_argo(name: str) -> str:
     import re
@@ -176,12 +174,10 @@ def is_tfx_pod(pod) -> bool:
                     store=mlmd_store,
                     context_id=run_context.id,
                     execution_type_name=KFP_EXECUTION_TYPE_NAME_PREFIX + component_version,
+                    pod_name=pod_name,
                     pipeline_name=argo_workflow_name,
                     run_id=argo_workflow_name,
                     instance_id=component_name,
-                    custom_properties={
-                        KFP_POD_NAME_EXECUTION_PROPERTY_NAME: pod_name,
-                    }
                 )
 
                 argo_input_artifacts = argo_template.get('inputs', {}).get('artifacts', [])