diff --git a/backend/metadata_writer/src/metadata_helpers.py b/backend/metadata_writer/src/metadata_helpers.py index e488bf177c44..048de9509a97 100644 --- 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 index 4d27ede2cee4..ed13ce580cce 100644 --- 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', [])