Repository: kubeflow/pipelines

Issue:
[v2compat] re-evaluate execution custom properties schema
* [ ] rename task execution `task_name` to `display_name`?


Discussion:

Resolving patch:
diff --git a/sdk/python/kfp/compiler/_default_transformers.py b/sdk/python/kfp/compiler/_default_transformers.py
--- a/sdk/python/kfp/compiler/_default_transformers.py
+++ b/sdk/python/kfp/compiler/_default_transformers.py
@@ -41,6 +41,11 @@ def add_kfp_pod_env(op: BaseOp) -> BaseOp:
                           value_from=k8s_client.V1EnvVarSource(
                               field_ref=k8s_client.V1ObjectFieldSelector(
                                   field_path='metadata.name')))
+  ).add_env_variable(
+      k8s_client.V1EnvVar(name='KFP_POD_UID',
+                          value_from=k8s_client.V1EnvVarSource(
+                              field_ref=k8s_client.V1ObjectFieldSelector(
+                                  field_path='metadata.uid')))
   ).add_env_variable(
       k8s_client.V1EnvVar(name='KFP_NAMESPACE',
                           value_from=k8s_client.V1EnvVarSource(
@@ -52,6 +57,12 @@ def add_kfp_pod_env(op: BaseOp) -> BaseOp:
           value_from=k8s_client.
           V1EnvVarSource(field_ref=k8s_client.V1ObjectFieldSelector(
               field_path="metadata.labels['workflows.argoproj.io/workflow']")))
+  ).add_env_variable(
+      k8s_client.V1EnvVar(
+          name='KFP_RUN_ID',
+          value_from=k8s_client.
+          V1EnvVarSource(field_ref=k8s_client.V1ObjectFieldSelector(
+              field_path="metadata.labels['pipeline/runid']")))
   ).add_env_variable(
       k8s_client.V1EnvVar(
           name='ENABLE_CACHING',
diff --git a/sdk/python/kfp/compiler/v2_compat.py b/sdk/python/kfp/compiler/v2_compat.py
--- a/sdk/python/kfp/compiler/v2_compat.py
+++ b/sdk/python/kfp/compiler/v2_compat.py
@@ -100,10 +100,16 @@ def update_op(op: dsl.ContainerOp,
       op.name,
       "--pipeline_name",
       pipeline_name,
-      "--pipeline_run_id",
-      "$(WORKFLOW_ID)",
-      "--pipeline_task_id",
+      "--run_id",
+      "$(KFP_RUN_ID)",
+      "--run_resource",
+      "workflows.argoproj.io/$(WORKFLOW_ID)",
+      "--namespace",
+      "$(KFP_NAMESPACE)",
+      "--pod_name",
       "$(KFP_POD_NAME)",
+      "--pod_uid",
+      "$(KFP_POD_UID)",
       "--pipeline_root",
       pipeline_root,
       "--enable_caching",


Test patch:
diff --git a/frontend/src/mlmd/MlmdUtils.test.ts b/frontend/src/mlmd/MlmdUtils.test.ts
--- a/frontend/src/mlmd/MlmdUtils.test.ts
+++ b/frontend/src/mlmd/MlmdUtils.test.ts
@@ -30,6 +30,7 @@ import { Workflow, WorkflowSpec, WorkflowStatus } from 'third_party/argo-ui/argo
 testBestPractices();
 
 const WORKFLOW_NAME = 'run-st448';
+const RUN_ID = 'abcdefghijk';
 const WORKFLOW_EMPTY: Workflow = {
   metadata: {
     name: WORKFLOW_NAME,
@@ -40,8 +41,8 @@ const WORKFLOW_EMPTY: Workflow = {
 };
 
 const V2_CONTEXT = new Context();
-V2_CONTEXT.setName(WORKFLOW_NAME);
-V2_CONTEXT.setType('kfp.PipelineRun');
+V2_CONTEXT.setName(RUN_ID);
+V2_CONTEXT.setType('system.PipelineRun');
 
 const TFX_CONTEXT = new Context();
 TFX_CONTEXT.setName('run.run-st448');
@@ -55,26 +56,29 @@ describe('MlmdUtils', () => {
   describe('getRunContext', () => {
     it('gets KFP v2 context', async () => {
       mockGetContextByTypeAndName([V2_CONTEXT]);
-      const context = await getRunContext({
-        ...WORKFLOW_EMPTY,
-        metadata: {
-          ...WORKFLOW_EMPTY.metadata,
-          annotations: { 'pipelines.kubeflow.org/v2_pipeline': 'true' },
+      const context = await getRunContext(
+        {
+          ...WORKFLOW_EMPTY,
+          metadata: {
+            ...WORKFLOW_EMPTY.metadata,
+            annotations: { 'pipelines.kubeflow.org/v2_pipeline': 'true' },
+          },
         },
-      });
+        RUN_ID,
+      );
       expect(context).toEqual(V2_CONTEXT);
     });
 
     it('gets TFX context', async () => {
       mockGetContextByTypeAndName([TFX_CONTEXT, V1_CONTEXT]);
-      const context = await getRunContext(WORKFLOW_EMPTY);
+      const context = await getRunContext(WORKFLOW_EMPTY, RUN_ID);
       expect(context).toEqual(TFX_CONTEXT);
     });
 
     it('gets KFP v1 context', async () => {
       const verify = expectWarnings();
       mockGetContextByTypeAndName([V1_CONTEXT]);
-      const context = await getRunContext(WORKFLOW_EMPTY);
+      const context = await getRunContext(WORKFLOW_EMPTY, RUN_ID);
       expect(context).toEqual(V1_CONTEXT);
       verify();
     });
@@ -82,7 +86,7 @@ describe('MlmdUtils', () => {
     it('throws error when not found', async () => {
       const verify = expectWarnings();
       mockGetContextByTypeAndName([]);
-      await expect(getRunContext(WORKFLOW_EMPTY)).rejects.toThrow();
+      await expect(getRunContext(WORKFLOW_EMPTY, RUN_ID)).rejects.toThrow();
       verify();
     });
   });
diff --git a/frontend/src/mlmd/__snapshots__/LineageActionBar.test.tsx.snap b/frontend/src/mlmd/__snapshots__/LineageActionBar.test.tsx.snap
--- a/frontend/src/mlmd/__snapshots__/LineageActionBar.test.tsx.snap
+++ b/frontend/src/mlmd/__snapshots__/LineageActionBar.test.tsx.snap
@@ -261,7 +261,7 @@ exports[`LineageActionBar Adds the artifact to the history state and DOM when pu
         key="breadcrumb-1"
         onClick={[Function]}
       >
-        (unnamed)
+        (No name)
       </button>
     </div>
     <div>
diff --git a/samples/core/exit_handler/exit_handler_test.py b/samples/core/exit_handler/exit_handler_test.py
--- a/samples/core/exit_handler/exit_handler_test.py
+++ b/samples/core/exit_handler/exit_handler_test.py
@@ -23,10 +23,8 @@
 from .exit_handler import pipeline_exit_handler
 from ...test.util import run_pipeline_func, TestCase, KfpMlmdClient
 
-def verify(
-    argo_workflow_name: str, mlmd_connection_config, run: kfp_server_api.ApiRun,
-    **kwargs
-):
+
+def verify(mlmd_connection_config, run: kfp_server_api.ApiRun, **kwargs):
     t = unittest.TestCase()
     t.maxDiff = None  # we always want to see full diff
 
@@ -34,19 +32,20 @@ def verify(
 
     # Verify MLMD state
     client = KfpMlmdClient(mlmd_connection_config=mlmd_connection_config)
-    tasks = client.get_tasks(argo_workflow_name=argo_workflow_name)
+    tasks = client.get_tasks(run_id=run.id)
     task_names = [*tasks.keys()]
     t.assertEqual(task_names, ['echo-msg', 'print-file', 'download-from-gcs'])
 
     for task in task_names:
         pprint(f'======= {task} =======')
         pprint(tasks.get(task).get_dict())
-    
+
     t.assertEqual(
         tasks.get('echo-msg').inputs.parameters.get('msg'),
         'exit!',
     )
 
+
 # %%
 
 if __name__ == '__main__':
diff --git a/samples/core/visualization/confusion_matrix_test.py b/samples/core/visualization/confusion_matrix_test.py
--- a/samples/core/visualization/confusion_matrix_test.py
+++ b/samples/core/visualization/confusion_matrix_test.py
@@ -21,32 +21,36 @@
 import kfp
 
 
-def verify(
-    run: kfp_server_api.ApiRun, mlmd_connection_config, argo_workflow_name: str,
-    **kwargs
-):
+def verify(run: kfp_server_api.ApiRun, mlmd_connection_config, **kwargs):
     t = unittest.TestCase()
     t.maxDiff = None  # we always want to see full diff
     t.assertEqual(run.status, 'Succeeded')
     client = KfpMlmdClient(mlmd_connection_config=mlmd_connection_config)
-    tasks = client.get_tasks(argo_workflow_name=argo_workflow_name)
+    tasks = client.get_tasks(run_id=run.id)
     pprint(tasks)
 
     confusion_visualization = tasks['confusion-visualization']
     output = [
-        a for a in confusion_visualization.outputs.artifacts if a.name == 'mlpipeline_ui_metadata'
+        a for a in confusion_visualization.outputs.artifacts
+        if a.name == 'mlpipeline_ui_metadata'
     ][0]
     pprint(output)
 
-    t.assertEqual(confusion_visualization.get_dict()['outputs']['artifacts'][0]['name'],
-                  'mlpipeline_ui_metadata')
-
-
-run_pipeline_func([TestCase(pipeline_func=confusion_matrix_pipeline,
-                            verify_func=verify,
-                            mode=kfp.dsl.PipelineExecutionMode.V2_COMPATIBLE,
-                            arguments={
-                            }),
-                   TestCase(pipeline_func=confusion_matrix_pipeline,
-                            mode=kfp.dsl.PipelineExecutionMode.V1_LEGACY
-                            )])
+    t.assertEqual(
+        confusion_visualization.get_dict()['outputs']['artifacts'][0]['name'],
+        'mlpipeline_ui_metadata'
+    )
+
+
+run_pipeline_func([
+    TestCase(
+        pipeline_func=confusion_matrix_pipeline,
+        verify_func=verify,
+        mode=kfp.dsl.PipelineExecutionMode.V2_COMPATIBLE,
+        arguments={}
+    ),
+    TestCase(
+        pipeline_func=confusion_matrix_pipeline,
+        mode=kfp.dsl.PipelineExecutionMode.V1_LEGACY
+    )
+])
diff --git a/samples/core/visualization/roc_test.py b/samples/core/visualization/roc_test.py
--- a/samples/core/visualization/roc_test.py
+++ b/samples/core/visualization/roc_test.py
@@ -21,32 +21,36 @@
 import kfp
 
 
-def verify(
-    run: kfp_server_api.ApiRun, mlmd_connection_config, argo_workflow_name: str,
-    **kwargs
-):
+def verify(run: kfp_server_api.ApiRun, mlmd_connection_config, **kwargs):
     t = unittest.TestCase()
     t.maxDiff = None  # we always want to see full diff
     t.assertEqual(run.status, 'Succeeded')
     client = KfpMlmdClient(mlmd_connection_config=mlmd_connection_config)
-    tasks = client.get_tasks(argo_workflow_name=argo_workflow_name)
+    tasks = client.get_tasks(run_id=run.id)
     pprint(tasks)
 
     roc_visualization = tasks['roc-visualization']
     output = [
-        a for a in roc_visualization.outputs.artifacts if a.name == 'mlpipeline_ui_metadata'
+        a for a in roc_visualization.outputs.artifacts
+        if a.name == 'mlpipeline_ui_metadata'
     ][0]
     pprint(output)
 
-    t.assertEqual(roc_visualization.get_dict()['outputs']['artifacts'][0]['name'],
-                  'mlpipeline_ui_metadata')
-
-
-run_pipeline_func([TestCase(pipeline_func=roc_curve_pipeline,
-                            verify_func=verify,
-                            mode=kfp.dsl.PipelineExecutionMode.V2_COMPATIBLE,
-                            arguments={
-                            }),
-                   TestCase(pipeline_func=roc_curve_pipeline,
-                            mode=kfp.dsl.PipelineExecutionMode.V1_LEGACY
-                            )])
+    t.assertEqual(
+        roc_visualization.get_dict()['outputs']['artifacts'][0]['name'],
+        'mlpipeline_ui_metadata'
+    )
+
+
+run_pipeline_func([
+    TestCase(
+        pipeline_func=roc_curve_pipeline,
+        verify_func=verify,
+        mode=kfp.dsl.PipelineExecutionMode.V2_COMPATIBLE,
+        arguments={}
+    ),
+    TestCase(
+        pipeline_func=roc_curve_pipeline,
+        mode=kfp.dsl.PipelineExecutionMode.V1_LEGACY
+    )
+])
diff --git a/samples/core/visualization/tensorboard_minio_test.py b/samples/core/visualization/tensorboard_minio_test.py
--- a/samples/core/visualization/tensorboard_minio_test.py
+++ b/samples/core/visualization/tensorboard_minio_test.py
@@ -23,16 +23,13 @@
 from ...test.util import run_pipeline_func, TestCase, KfpMlmdClient
 
 
-def verify(
-    run: kfp_server_api.ApiRun, mlmd_connection_config, argo_workflow_name: str,
-    **kwargs
-):
+def verify(run: kfp_server_api.ApiRun, mlmd_connection_config, **kwargs):
     t = unittest.TestCase()
     t.maxDiff = None  # we always want to see full diff
     t.assertEqual(run.status, 'Succeeded')
 
     client = KfpMlmdClient(mlmd_connection_config=mlmd_connection_config)
-    tasks = client.get_tasks(argo_workflow_name=argo_workflow_name)
+    tasks = client.get_tasks(run_id=run.id)
     # uncomment to debug
     # pprint(tasks)
     vis = tasks['create-tensorboard-visualization']
diff --git a/samples/test/cache_v2_test.py b/samples/test/cache_v2_test.py
--- a/samples/test/cache_v2_test.py
+++ b/samples/test/cache_v2_test.py
@@ -21,7 +21,6 @@
 import string
 import unittest
 import functools
-import time
 
 import kfp
 import kfp_server_api
@@ -31,14 +30,10 @@
 from ml_metadata.proto import Execution
 
 
-def get_tasks(mlmd_connection_config, argo_workflow_name: str):
-    # Verify MLMD state
-    client = KfpMlmdClient(mlmd_connection_config=mlmd_connection_config)
-    tasks = client.get_tasks(argo_workflow_name=argo_workflow_name)
-    return tasks
-
-
-def verify_tasks(t: unittest.TestCase, tasks: dict[str, KfpTask], task_state, uri: str, some_int:int):
+def verify_tasks(
+    t: unittest.TestCase, tasks: dict[str, KfpTask], task_state, uri: str,
+    some_int: int
+):
     task_names = [*tasks.keys()]
     t.assertEqual(task_names, ['train-op', 'preprocess'], 'task names')
 
@@ -56,22 +51,26 @@ def verify_tasks(t: unittest.TestCase, tasks: dict[str, KfpTask], task_state, ur
         },
         'outputs': {
             'artifacts': [{
-                'metadata': {},
+                'metadata': {
+                    'display_name': 'output_dataset_one',
+                },
                 'name': 'output_dataset_one',
                 'type': 'system.Dataset'
             }],
             'parameters': {
-                'output_parameter_one':some_int
+                'output_parameter_one': some_int
             }
         },
-        'type': 'kfp.ContainerExecution',
+        'type': 'system.ContainerExecution',
         'state': task_state,
     }, preprocess.get_dict())
     t.assertEqual({
         'name': 'train-op',
         'inputs': {
             'artifacts': [{
-                'metadata': {},
+                'metadata': {
+                    'display_name': 'output_dataset_one',
+                },
                 'name': 'dataset',
                 'type': 'system.Dataset',
             }],
@@ -81,28 +80,31 @@ def verify_tasks(t: unittest.TestCase, tasks: dict[str, KfpTask], task_state, ur
         },
         'outputs': {
             'artifacts': [{
-                'metadata': {},
+                'metadata': {
+                    'display_name': 'model',
+                },
                 'name': 'model',
                 'type': 'system.Model',
             }],
             'parameters': {}
         },
-        'type': 'kfp.ContainerExecution',
+        'type': 'system.ContainerExecution',
         'state': task_state,
     }, train.get_dict())
 
 
 def verify(
-        run: kfp_server_api.ApiRun, mlmd_connection_config, argo_workflow_name: str,
-        uri: str, some_int, state: int,
-        **kwargs
+    run: kfp_server_api.ApiRun, mlmd_connection_config, uri: str, some_int,
+    state: int, **kwargs
 ):
     t = unittest.TestCase()
     t.maxDiff = None  # we always want to see full diff
     t.assertEqual(run.status, 'Succeeded')
-    tasks = get_tasks(mlmd_connection_config, argo_workflow_name)
+    client = KfpMlmdClient(mlmd_connection_config=mlmd_connection_config)
+    tasks = client.get_tasks(run_id=run.id)
     verify_tasks(t, tasks, state, uri, some_int)
 
+
 if __name__ == '__main__':
     letters = string.ascii_lowercase
     random_uri = 'http://' + ''.join(random.choice(letters) for i in range(5))
@@ -110,8 +112,16 @@ def verify(
     run_pipeline_func([
         TestCase(
             pipeline_func=two_step_pipeline,
-            arguments={'uri': f'{random_uri}', 'some_int': f'{random_int}'},
-            verify_func=functools.partial(verify, uri=random_uri, some_int=random_int, state = Execution.State.COMPLETE,),
+            arguments={
+                'uri': f'{random_uri}',
+                'some_int': f'{random_int}'
+            },
+            verify_func=functools.partial(
+                verify,
+                uri=random_uri,
+                some_int=random_int,
+                state=Execution.State.COMPLETE,
+            ),
             mode=kfp.dsl.PipelineExecutionMode.V2_COMPATIBLE,
             enable_caching=True
         ),
@@ -119,8 +129,16 @@ def verify(
     run_pipeline_func([
         TestCase(
             pipeline_func=two_step_pipeline,
-            arguments={'uri': f'{random_uri}', 'some_int': f'{random_int}'},
-            verify_func=functools.partial(verify, uri=random_uri, some_int=random_int, state = Execution.State.CACHED),
+            arguments={
+                'uri': f'{random_uri}',
+                'some_int': f'{random_int}'
+            },
+            verify_func=functools.partial(
+                verify,
+                uri=random_uri,
+                some_int=random_int,
+                state=Execution.State.CACHED
+            ),
             mode=kfp.dsl.PipelineExecutionMode.V2_COMPATIBLE,
             enable_caching=True
         ),
diff --git a/samples/test/lightweight_python_functions_v2_pipeline_test.py b/samples/test/lightweight_python_functions_v2_pipeline_test.py
--- a/samples/test/lightweight_python_functions_v2_pipeline_test.py
+++ b/samples/test/lightweight_python_functions_v2_pipeline_test.py
@@ -22,15 +22,12 @@
 from ml_metadata.proto import Execution
 
 
-def verify(
-    run: kfp_server_api.ApiRun, mlmd_connection_config, argo_workflow_name: str,
-    **kwargs
-):
+def verify(run: kfp_server_api.ApiRun, mlmd_connection_config, **kwargs):
     t = unittest.TestCase()
     t.maxDiff = None  # we always want to see full diff
     t.assertEqual(run.status, 'Succeeded')
     client = KfpMlmdClient(mlmd_connection_config=mlmd_connection_config)
-    tasks = client.get_tasks(argo_workflow_name=argo_workflow_name)
+    tasks = client.get_tasks(run_id=run.id)
 
     task_names = [*tasks.keys()]
     t.assertCountEqual(task_names, ['preprocess', 'train'], 'task names')
@@ -51,11 +48,15 @@ def verify(
             'name': 'preprocess',
             'outputs': {
                 'artifacts': [{
-                    'metadata': {},
+                    'metadata': {
+                        'display_name': 'output_dataset_one'
+                    },
                     'name': 'output_dataset_one',
                     'type': 'system.Dataset'
                 }, {
-                    'metadata': {},
+                    'metadata': {
+                        'display_name': 'output_dataset_two'
+                    },
                     'name': 'output_dataset_two',
                     'type': 'system.Dataset'
                 }],
@@ -66,7 +67,7 @@ def verify(
                     'output_parameter': 'message'
                 }
             },
-            'type': 'kfp.ContainerExecution',
+            'type': 'system.ContainerExecution',
             'state': Execution.State.COMPLETE,
         },
         preprocess.get_dict(),
@@ -75,11 +76,15 @@ def verify(
         {
             'inputs': {
                 'artifacts': [{
-                    'metadata': {},
+                    'metadata': {
+                        'display_name': 'output_dataset_one'
+                    },
                     'name': 'dataset_one',
                     'type': 'system.Dataset'
                 }, {
-                    'metadata': {},
+                    'metadata': {
+                        'display_name': 'output_dataset_two'
+                    },
                     'name': 'dataset_two',
                     'type': 'system.Dataset'
                 }],
@@ -95,16 +100,15 @@ def verify(
             'outputs': {
                 'artifacts': [{
                     'metadata': {
-                        'accuracy': {
-                            'doubleValue': 0.9
-                        }
+                        'display_name': 'model',
+                        'accuracy': 0.9,
                     },
                     'name': 'model',
                     'type': 'system.Model'
                 }],
                 'parameters': {}
             },
-            'type': 'kfp.ContainerExecution',
+            'type': 'system.ContainerExecution',
             'state': Execution.State.COMPLETE,
         },
         train.get_dict(),
diff --git a/samples/test/lightweight_python_functions_v2_with_outputs_test.py b/samples/test/lightweight_python_functions_v2_with_outputs_test.py
--- a/samples/test/lightweight_python_functions_v2_with_outputs_test.py
+++ b/samples/test/lightweight_python_functions_v2_with_outputs_test.py
@@ -23,15 +23,12 @@
 from .util import KfpMlmdClient, run_pipeline_func, TestCase
 
 
-def verify(
-    run: kfp_server_api.ApiRun, mlmd_connection_config, argo_workflow_name: str,
-    **kwargs
-):
+def verify(run: kfp_server_api.ApiRun, mlmd_connection_config, **kwargs):
     t = unittest.TestCase()
     t.maxDiff = None  # we always want to see full diff
     t.assertEqual(run.status, 'Succeeded')
     client = KfpMlmdClient(mlmd_connection_config=mlmd_connection_config)
-    tasks = client.get_tasks(argo_workflow_name=argo_workflow_name)
+    tasks = client.get_tasks(run_id=run.id)
     pprint(tasks)
 
     output_artifact = tasks['output-artifact']
diff --git a/samples/test/metrics_visualization_v2_test.py b/samples/test/metrics_visualization_v2_test.py
--- a/samples/test/metrics_visualization_v2_test.py
+++ b/samples/test/metrics_visualization_v2_test.py
@@ -23,15 +23,12 @@
 from ml_metadata.proto import Execution
 
 
-def verify(
-    run: kfp_server_api.ApiRun, mlmd_connection_config, argo_workflow_name: str,
-    **kwargs
-):
+def verify(run: kfp_server_api.ApiRun, mlmd_connection_config, **kwargs):
     t = unittest.TestCase()
     t.maxDiff = None  # we always want to see full diff
     t.assertEqual(run.status, 'Succeeded')
     client = KfpMlmdClient(mlmd_connection_config=mlmd_connection_config)
-    tasks = client.get_tasks(argo_workflow_name=argo_workflow_name)
+    tasks = client.get_tasks(run_id=run.id)
 
     task_names = [*tasks.keys()]
     t.assertCountEqual(
@@ -60,54 +57,53 @@ def verify(
         'outputs': {
             'artifacts': [{
                 'metadata': {
+                    'display_name': 'metrics',
                     'confidenceMetrics': {
-                        'structValue': {
-                            'list': [{
-                                'confidenceThreshold': 2.0,
-                                'falsePositiveRate': 0.0,
-                                'recall': 0.0
-                            }, {
-                                'confidenceThreshold': 1.0,
-                                'falsePositiveRate': 0.0,
-                                'recall': 0.33962264150943394
-                            }, {
-                                'confidenceThreshold': 0.9,
-                                'falsePositiveRate': 0.0,
-                                'recall': 0.6037735849056604
-                            }, {
-                                'confidenceThreshold': 0.8,
-                                'falsePositiveRate': 0.0,
-                                'recall': 0.8490566037735849
-                            }, {
-                                'confidenceThreshold': 0.6,
-                                'falsePositiveRate': 0.0,
-                                'recall': 0.8867924528301887
-                            }, {
-                                'confidenceThreshold': 0.5,
-                                'falsePositiveRate': 0.0125,
-                                'recall': 0.9245283018867925
-                            }, {
-                                'confidenceThreshold': 0.4,
-                                'falsePositiveRate': 0.075,
-                                'recall': 0.9622641509433962
-                            }, {
-                                'confidenceThreshold': 0.3,
-                                'falsePositiveRate': 0.0875,
-                                'recall': 1.0
-                            }, {
-                                'confidenceThreshold': 0.2,
-                                'falsePositiveRate': 0.2375,
-                                'recall': 1.0
-                            }, {
-                                'confidenceThreshold': 0.1,
-                                'falsePositiveRate': 0.475,
-                                'recall': 1.0
-                            }, {
-                                'confidenceThreshold': 0.0,
-                                'falsePositiveRate': 1.0,
-                                'recall': 1.0
-                            }]
-                        }
+                        'list': [{
+                            'confidenceThreshold': 2.0,
+                            'falsePositiveRate': 0.0,
+                            'recall': 0.0
+                        }, {
+                            'confidenceThreshold': 1.0,
+                            'falsePositiveRate': 0.0,
+                            'recall': 0.33962264150943394
+                        }, {
+                            'confidenceThreshold': 0.9,
+                            'falsePositiveRate': 0.0,
+                            'recall': 0.6037735849056604
+                        }, {
+                            'confidenceThreshold': 0.8,
+                            'falsePositiveRate': 0.0,
+                            'recall': 0.8490566037735849
+                        }, {
+                            'confidenceThreshold': 0.6,
+                            'falsePositiveRate': 0.0,
+                            'recall': 0.8867924528301887
+                        }, {
+                            'confidenceThreshold': 0.5,
+                            'falsePositiveRate': 0.0125,
+                            'recall': 0.9245283018867925
+                        }, {
+                            'confidenceThreshold': 0.4,
+                            'falsePositiveRate': 0.075,
+                            'recall': 0.9622641509433962
+                        }, {
+                            'confidenceThreshold': 0.3,
+                            'falsePositiveRate': 0.0875,
+                            'recall': 1.0
+                        }, {
+                            'confidenceThreshold': 0.2,
+                            'falsePositiveRate': 0.2375,
+                            'recall': 1.0
+                        }, {
+                            'confidenceThreshold': 0.1,
+                            'falsePositiveRate': 0.475,
+                            'recall': 1.0
+                        }, {
+                            'confidenceThreshold': 0.0,
+                            'falsePositiveRate': 1.0,
+                            'recall': 1.0
+                        }]
                     }
                 },
                 'name': 'metrics',
@@ -115,7 +111,7 @@ def verify(
             }],
             'parameters': {}
         },
-        'type': 'kfp.ContainerExecution',
+        'type': 'system.ContainerExecution',
         'state': Execution.State.COMPLETE,
     }, wine_classification.get_dict())
     t.assertEqual(
@@ -130,7 +126,8 @@ def verify(
             'outputs': {
                 'artifacts': [{
                     'metadata': {
-                        'confusionMatrix': {'structValue': {'struct': {
+                        'display_name': 'metrics',
+                        'confusionMatrix': {'struct': {
                             'annotationSpecs': [{
                                 'displayName': 'Setosa'
                             }, {
@@ -145,19 +142,19 @@ def verify(
                             }, {
                                 'row': [mock.ANY, mock.ANY, mock.ANY]
                             }]
-                        }}}
+                        }}
                     },
                     'name': 'metrics',
                     'type': 'system.ClassificationMetrics'
                 }],
                 'parameters': {}
             },
-            'type': 'kfp.ContainerExecution',
+            'type': 'system.ContainerExecution',
             'state': Execution.State.COMPLETE,
         }, iris_sgdclassifier.get_dict()
     )
     rows = iris_sgdclassifier.get_dict()['outputs']['artifacts'][0]['metadata'][
-        'confusionMatrix']['structValue']['struct']['rows']
+        'confusionMatrix']['struct']['rows']
     for i, row in enumerate(rows):
         for j, item in enumerate(row['row']):
             t.assertIsInstance(
@@ -174,16 +171,15 @@ def verify(
         'outputs': {
             'artifacts': [{
                 'metadata': {
-                    'accuracy': {
-                        'doubleValue': 92.0
-                    },
+                    'display_name': 'metrics',
+                    'accuracy': 92.0,
                 },
                 'name': 'metrics',
                 'type': 'system.Metrics'
             }],
             'parameters': {}
         },
-        'type': 'kfp.ContainerExecution',
+        'type': 'system.ContainerExecution',
         'state': Execution.State.COMPLETE,
     }, digit_classification.get_dict())
 
diff --git a/samples/test/reused_component_test.py b/samples/test/reused_component_test.py
--- a/samples/test/reused_component_test.py
+++ b/samples/test/reused_component_test.py
@@ -19,14 +19,14 @@
 from .util import run_pipeline_func, TestCase, KfpMlmdClient
 
 
-def verify(run, argo_workflow_name: str, mlmd_connection_config, **kwargs):
+def verify(run, mlmd_connection_config, **kwargs):
     t = unittest.TestCase()
     t.maxDiff = None  # we always want to see full diff
     t.assertEqual(run.status, 'Succeeded')
 
     client = KfpMlmdClient(mlmd_connection_config=mlmd_connection_config)
 
-    tasks = client.get_tasks(argo_workflow_name=argo_workflow_name)
+    tasks = client.get_tasks(run_id=run.id)
     t.assertEqual(
         17, tasks['add-3'].outputs.parameters['sum'], 'add result should be 17'
     )
diff --git a/samples/test/two_step_test.py b/samples/test/two_step_test.py
--- a/samples/test/two_step_test.py
+++ b/samples/test/two_step_test.py
@@ -27,13 +27,6 @@
 from ml_metadata.proto import Execution
 
 
-def get_tasks(mlmd_connection_config, argo_workflow_name: str):
-    # Verify MLMD state
-    client = KfpMlmdClient(mlmd_connection_config=mlmd_connection_config)
-    tasks = client.get_tasks(argo_workflow_name=argo_workflow_name)
-    return tasks
-
-
 def verify_tasks(t: unittest.TestCase, tasks: dict[str, KfpTask]):
     task_names = [*tasks.keys()]
     t.assertEqual(task_names, ['train-op', 'preprocess'], 'task names')
@@ -58,7 +51,9 @@ def verify_tasks(t: unittest.TestCase, tasks: dict[str, KfpTask]):
         },
         'outputs': {
             'artifacts': [{
-                'metadata': {},
+                'metadata': {
+                    'display_name': 'output_dataset_one',
+                },
                 'name': 'output_dataset_one',
                 'type': 'system.Dataset'
             }],
@@ -66,14 +61,16 @@ def verify_tasks(t: unittest.TestCase, tasks: dict[str, KfpTask]):
                 'output_parameter_one': 1234
             }
         },
-        'type': 'kfp.ContainerExecution',
+        'type': 'system.ContainerExecution',
         'state': Execution.State.COMPLETE,
     }, preprocess.get_dict())
     t.assertEqual({
         'name': 'train-op',
         'inputs': {
             'artifacts': [{
-                'metadata': {},
+                'metadata': {
+                    'display_name': 'output_dataset_one',
+                },
                 'name': 'dataset',
                 'type': 'system.Dataset',
             }],
@@ -83,13 +80,15 @@ def verify_tasks(t: unittest.TestCase, tasks: dict[str, KfpTask]):
         },
         'outputs': {
             'artifacts': [{
-                'metadata': {},
+                'metadata': {
+                    'display_name': 'model',
+                },
                 'name': 'model',
                 'type': 'system.Model',
             }],
             'parameters': {}
         },
-        'type': 'kfp.ContainerExecution',
+        'type': 'system.ContainerExecution',
         'state': Execution.State.COMPLETE,
     }, train.get_dict())
 
@@ -100,37 +99,35 @@ def verify_artifacts(t: unittest.TestCase, tasks: dict, artifact_uri_prefix):
             t.assertTrue(artifact.uri.startswith(artifact_uri_prefix))
 
 
-def verify(
-    run: kfp_server_api.ApiRun, mlmd_connection_config, argo_workflow_name: str,
-    **kwargs
-):
+def verify(run: kfp_server_api.ApiRun, mlmd_connection_config, **kwargs):
     t = unittest.TestCase()
     t.maxDiff = None  # we always want to see full diff
     t.assertEqual(run.status, 'Succeeded')
-    tasks = get_tasks(mlmd_connection_config, argo_workflow_name)
+    client = KfpMlmdClient(mlmd_connection_config=mlmd_connection_config)
+    tasks = client.get_tasks(run_id=run.id)
     verify_tasks(t, tasks)
 
 
 def verify_with_default_pipeline_root(
-    run: kfp_server_api.ApiRun, mlmd_connection_config, argo_workflow_name: str,
-    **kwargs
+    run: kfp_server_api.ApiRun, mlmd_connection_config, **kwargs
 ):
     t = unittest.TestCase()
     t.maxDiff = None  # we always want to see full diff
     t.assertEqual(run.status, 'Succeeded')
-    tasks = get_tasks(mlmd_connection_config, argo_workflow_name)
+    client = KfpMlmdClient(mlmd_connection_config=mlmd_connection_config)
+    tasks = client.get_tasks(run_id=run.id)
     verify_tasks(t, tasks)
     verify_artifacts(t, tasks, 'minio://mlpipeline/v2/artifacts')
 
 
 def verify_with_specific_pipeline_root(
-    run: kfp_server_api.ApiRun, mlmd_connection_config, argo_workflow_name: str,
-    **kwargs
+    run: kfp_server_api.ApiRun, mlmd_connection_config, **kwargs
 ):
     t = unittest.TestCase()
     t.maxDiff = None  # we always want to see full diff
     t.assertEqual(run.status, 'Succeeded')
-    tasks = get_tasks(mlmd_connection_config, argo_workflow_name)
+    client = KfpMlmdClient(mlmd_connection_config=mlmd_connection_config)
+    tasks = client.get_tasks(run_id=run.id)
     verify_tasks(t, tasks)
     verify_artifacts(t, tasks, 'minio://mlpipeline/override/artifacts')
 
diff --git a/samples/test/util.py b/samples/test/util.py
--- a/samples/test/util.py
+++ b/samples/test/util.py
@@ -66,8 +66,9 @@ def run_pipeline_func(test_cases: List[TestCase]):
     """
 
     def test_wrapper(
-        run_pipeline: Callable[[Callable, kfp.dsl.PipelineExecutionMode, bool, dict],
-                               kfp_server_api.ApiRunDetail],
+        run_pipeline: Callable[
+            [Callable, kfp.dsl.PipelineExecutionMode, bool, dict],
+            kfp_server_api.ApiRunDetail],
         mlmd_connection_config: metadata_store_pb2.MetadataStoreClientConfig,
     ):
         for case in test_cases:
@@ -222,6 +223,20 @@ def _create_run():
     fire.Fire(main)
 
 
+def simplify_proto_struct(data: dict) -> dict:
+    res = {}
+    for key, value in data.items():
+        if value.get('stringValue') is not None:
+            res[key] = value['stringValue']
+        elif value.get('doubleValue') is not None:
+            res[key] = value['doubleValue']
+        elif value.get('structValue') is not None:
+            res[key] = value['structValue']
+        else:
+            res[key] = value
+    return res
+
+
 @dataclass
 class KfpArtifact:
     name: str
@@ -241,7 +256,9 @@ def new(
         artifact_name = mlmd_event.path.steps[0].key
         # The original field is custom_properties, but MessageToDict converts it
         # to customProperties.
-        metadata = MessageToDict(mlmd_artifact).get('customProperties', {})
+        metadata = simplify_proto_struct(
+            MessageToDict(mlmd_artifact).get('customProperties', {})
+        )
         return cls(
             name=artifact_name,
             type=mlmd_artifact_type.name,
@@ -353,14 +370,14 @@ def __init__(
             )
         self.mlmd_store = metadata_store.MetadataStore(mlmd_connection_config)
 
-    def get_tasks(self, argo_workflow_name: str):
+    def get_tasks(self, run_id: str):
         run_context = self.mlmd_store.get_context_by_type_and_name(
-            type_name='kfp.PipelineRun',
-            context_name=argo_workflow_name,
+            type_name='system.PipelineRun',
+            context_name=run_id,
         )
         if not run_context:
             raise Exception(
-                f'Cannot find kfp.PipelineRun context "{argo_workflow_name}"'
+                f'Cannot find system.PipelineRun context "{run_id}"'
             )
         logging.info(
             f'run_context: name={run_context.name} id={run_context.id}'
diff --git a/sdk/python/tests/compiler/testdata/add_pod_env.yaml b/sdk/python/tests/compiler/testdata/add_pod_env.yaml
--- a/sdk/python/tests/compiler/testdata/add_pod_env.yaml
+++ b/sdk/python/tests/compiler/testdata/add_pod_env.yaml
@@ -21,6 +21,10 @@ spec:
         valueFrom:
           fieldRef:
             fieldPath: metadata.name
+      - name: KFP_POD_UID
+        valueFrom:
+          fieldRef:
+            fieldPath: metadata.uid
       - name: KFP_NAMESPACE
         valueFrom:
           fieldRef:
@@ -29,6 +33,10 @@ spec:
         valueFrom:
           fieldRef:
             fieldPath: metadata.labels['workflows.argoproj.io/workflow']
+      - name: KFP_RUN_ID
+        valueFrom:
+          fieldRef:
+            fieldPath: metadata.labels['pipeline/runid']
       - name: ENABLE_CACHING
         valueFrom:
           fieldRef:
diff --git a/v2/metadata/client_test.go b/v2/metadata/client_test.go
--- a/v2/metadata/client_test.go
+++ b/v2/metadata/client_test.go
@@ -35,6 +35,8 @@ import (
 const (
 	testMlmdServerAddress = "localhost"
 	testMlmdServerPort    = "8080"
+	namespace             = "kubeflow"
+	runResource           = "workflows.argoproj.io/hello-world-abcd"
 )
 
 func Test_schemaToArtifactType(t *testing.T) {
@@ -84,9 +86,9 @@ func Test_GetPipeline(t *testing.T) {
 	mlmdClient, err := NewTestMlmdClient()
 	fatalIf(err)
 
-	_, err = client.GetPipeline(ctx, "get-pipeline-test", runId)
+	_, err = client.GetPipeline(ctx, "get-pipeline-test", runId, namespace, runResource)
 	fatalIf(err)
-	runCtxType := "kfp.PipelineRun"
+	runCtxType := "system.PipelineRun"
 	pipeline := "get-pipeline-test"
 
 	res, err := mlmdClient.GetContextByTypeAndName(ctx, &pb.GetContextByTypeAndNameRequest{
@@ -130,7 +132,7 @@ func Test_GetPipelineConcurrently(t *testing.T) {
 		wg.Add(1)
 		go func() {
 			defer wg.Done()
-			_, err := client.GetPipeline(ctx, fmt.Sprintf("get-pipeline-concurrently-test-%s", runIdText), runIdText)
+			_, err := client.GetPipeline(ctx, fmt.Sprintf("get-pipeline-concurrently-test-%s", runIdText), runIdText, namespace, "workflows.argoproj.io/hello-world-"+runIdText)
 			if err != nil {
 				t.Error(err)
 			}
@@ -142,7 +144,7 @@ func Test_GetPipelineConcurrently(t *testing.T) {
 		wg.Add(1)
 		go func() {
 			defer wg.Done()
-			_, err := client.GetPipeline(ctx, fmt.Sprintf("get-pipeline-concurrently-test-%s", runIdText), runIdText)
+			_, err := client.GetPipeline(ctx, fmt.Sprintf("get-pipeline-concurrently-test-%s", runIdText), runIdText, namespace, "workflows.argoproj.io/hello-world-"+runIdText)
 			if err != nil {
 				t.Error(err)
 			}