From d077e6ecce8b102c15e066649d47fa3d67384ea6 Mon Sep 17 00:00:00 2001
From: "Yuan (Bob) Gong" <4957653+Bobgy@users.noreply.github.com>
Date: Thu, 19 Aug 2021 12:59:53 +0800
Subject: [PATCH] feat(frontend): integrate with TFX 1.2.0 metadata &
 visualization, no longer support previous versions. Part of #6138 (#6388)

---
 frontend/src/mlmd/MlmdUtils.test.ts |  4 ++--
 frontend/src/mlmd/MlmdUtils.ts      | 16 +++++-----------
 2 files changed, 7 insertions(+), 13 deletions(-)

diff --git a/frontend/src/mlmd/MlmdUtils.test.ts b/frontend/src/mlmd/MlmdUtils.test.ts
index ce44f0518cc..f940b7ca515 100644
--- a/frontend/src/mlmd/MlmdUtils.test.ts
+++ b/frontend/src/mlmd/MlmdUtils.test.ts
@@ -45,8 +45,8 @@ V2_CONTEXT.setName(RUN_ID);
 V2_CONTEXT.setType('system.PipelineRun');
 
 const TFX_CONTEXT = new Context();
-TFX_CONTEXT.setName('run.run-st448');
-TFX_CONTEXT.setType('run');
+TFX_CONTEXT.setName('run-st448');
+TFX_CONTEXT.setType('pipeline_run');
 
 const V1_CONTEXT = new Context();
 V1_CONTEXT.setName(WORKFLOW_NAME);
diff --git a/frontend/src/mlmd/MlmdUtils.ts b/frontend/src/mlmd/MlmdUtils.ts
index 1f60b4f06bc..af764625790 100644
--- a/frontend/src/mlmd/MlmdUtils.ts
+++ b/frontend/src/mlmd/MlmdUtils.ts
@@ -70,17 +70,9 @@ async function getContext({ type, name }: { type: string; name: string }): Promi
  * @throws error when network error, or not found
  */
 async function getTfxRunContext(argoWorkflowName: string): Promise<Context> {
-  // argoPodName has the general form "pipelineName-workflowId-executionId".
-  // All components of a pipeline within a single run will have the same
-  // "pipelineName-workflowId" prefix.
-  const pipelineName = argoWorkflowName
-    .split('-')
-    .slice(0, -1)
-    .join('_');
-  const runID = argoWorkflowName;
-  // An example run context name is parameterized_tfx_oss.parameterized-tfx-oss-4rq5v.
-  const tfxRunContextName = `${pipelineName}.${runID}`;
-  return await getContext({ name: tfxRunContextName, type: 'run' });
+  // Context: https://github.com/kubeflow/pipelines/issues/6138
+  // Require TFX version to be >= 1.2.0.
+  return await getContext({ name: argoWorkflowName, type: 'pipeline_run' });
 }
 
 /**
@@ -155,6 +147,8 @@ export const ExecutionHelpers = {
       getStringProperty(execution, ExecutionProperties.NAME) ||
       getStringProperty(execution, ExecutionProperties.COMPONENT_ID) ||
       getStringProperty(execution, ExecutionCustomProperties.TASK_ID, true) ||
+      // TFX 1.2.0 executions do not have any of the above, adding pod name as a fallback name
+      getStringProperty(execution, KfpExecutionProperties.KFP_POD_NAME, true) ||
       '(No name)'}`;
   },
   getState(execution: Execution): string | number | undefined {