Repository: tensorflow/tfx

Issue:
Customizing `mlpipeline_ui_metadata.json`
**System information**

- TFX Version (you are using): 0.29
- Environment in which you plan to use the feature (e.g., Local
  (Linux/MacOS/Windows), Interactive Notebook, Google Cloud, etc..): Google Cloud
- Are you willing to contribute it (Yes/No): Yes

**Describe the feature and the current behavior/state.**

https://github.com/tensorflow/tfx/blob/79fa0d8adce177d04feaa737f1440689344521df/tfx/orchestration/kubeflow/container_entrypoint.py#L210-L296

For now, `_dump_ui_metadata` function in `tfx/orchestration/kubeflow/container_entrypoint.py` does not consider already existing `/mlpipeline-ui-metadata.json`. So if I write some metadata that I want to see in KubeFlow UI in a custom component, the file will be overwritten.

If `/mlpipeline-ui-metadata.json` already exists(maybe created in the component) and has a valid format, I think extending an existing file is better.

**Will this change the current API? How?**

No

**Who will benefit with this feature?**

**Do you have a workaround or are completely blocked by this?** :

**Name of your Organization (Optional)**

**Any Other info.**



Discussion:
@kennethyang404 is there any progress here?
Sorry for the late response. Your suggestion sounds great. If you can provide some PR, I'm happy to review and merge it. 

cc. @chensun @zijianjoy @Bobgy

Resolving patch:
diff --git a/tfx/orchestration/kubeflow/container_entrypoint.py b/tfx/orchestration/kubeflow/container_entrypoint.py
--- a/tfx/orchestration/kubeflow/container_entrypoint.py
+++ b/tfx/orchestration/kubeflow/container_entrypoint.py
@@ -242,6 +242,9 @@ def _dump_ui_metadata(
 
   For general components we just render a simple Markdown file for
     exec_properties/inputs/outputs.
+  If the file already exists and is a valid format(have a list of
+  dictionaries in outputs key), we append the existing UI metadata items
+  to our output json file.
 
   Args:
     node: associated TFX node.
@@ -347,6 +350,35 @@ def _dump_output_populated_artifacts(
       tensorboard_output = {'type': 'tensorboard', 'source': output_model.uri}
       outputs.append(tensorboard_output)
 
+  # Add existing KFP UI Metadata if the file exists and is a valid format
+  if os.path.isfile(ui_metadata_path):
+
+    def _read_validated_ui_metadata(
+        ui_metadata_path: str) -> List[Dict[str, str]]:
+      """Read validated existing KFP UI Metadata file.
+
+      Args:
+        ui_metadata_path: path for ui metadata
+
+      Returns:
+        A list of UI metadata if the file is valid. An empty list otherwise.
+      """
+      result = []
+      try:
+        with open(ui_metadata_path, 'r') as f:
+          metadata_dict = json.load(f)
+
+        if ('outputs' in metadata_dict and
+            isinstance(metadata_dict['outputs'], list)):
+          for ui_metadata in metadata_dict['outputs']:
+            if isinstance(ui_metadata, dict):
+              result.append(ui_metadata)
+      except json.JSONDecodeError:
+        pass
+      return result
+
+    outputs += _read_validated_ui_metadata(ui_metadata_path)
+
   metadata_dict = {'outputs': outputs}
 
   with open(ui_metadata_path, 'w') as f:


Test patch:
diff --git a/tfx/orchestration/kubeflow/container_entrypoint_test.py b/tfx/orchestration/kubeflow/container_entrypoint_test.py
--- a/tfx/orchestration/kubeflow/container_entrypoint_test.py
+++ b/tfx/orchestration/kubeflow/container_entrypoint_test.py
@@ -119,6 +119,59 @@ def testDumpUiMetadata(self):
       self.assertEqual('tensorboard', ui_metadata['outputs'][-1]['type'])
       self.assertEqual('model_run_uri', ui_metadata['outputs'][-1]['source'])
 
+  def testDumpUiMetadataWithPreExistingFile(self):
+    dummy_node = pipeline_pb2.PipelineNode()
+    dummy_node.node_info.type.name = 'class_path_for_dummy_node.DummyComponent'
+    exec_info = data_types.ExecutionInfo(
+        input_dict={}, output_dict={}, exec_properties={}, execution_id='id')
+
+    ui_metadata_path = os.path.join(
+        os.environ.get('TEST_UNDECLARED_OUTPUTS_DIR', self.get_temp_dir()),
+        self._testMethodName, 'json')
+    fileio.makedirs(os.path.dirname(ui_metadata_path))
+
+    # Check with valid file
+    example_ui_metadata_item = {
+        'type': 'table',
+        'storage': 'gcs',
+        'format': 'csv',
+        'header': ['example-header1', 'example-header2'],
+        'source': 'gs://example-data-source/example.csv',
+    }
+    with fileio.open(ui_metadata_path, 'w') as f:
+      f.write(json.dumps({'outputs': [example_ui_metadata_item]}))
+
+    container_entrypoint._dump_ui_metadata(dummy_node, exec_info,
+                                           ui_metadata_path)
+
+    with open(ui_metadata_path) as f:
+      ui_metadata = json.load(f)
+      self.assertLen(ui_metadata['outputs'], 2)
+      self.assertTrue(
+          any('markdown' == item['type'] for item in ui_metadata['outputs']))
+      self.assertTrue(
+          any('table' == item['type'] for item in ui_metadata['outputs']))
+
+    # Check with invalid file
+    invalid_contents = [
+        json.dumps({'wrong_key': [{
+            'foo': 1
+        }]}),
+        json.dumps({'outputs': [1]}),  # not a dictionary item
+        'not a json',
+    ]
+    for content in invalid_contents:
+      with fileio.open(ui_metadata_path, 'w') as f:
+        f.write(content)
+
+      container_entrypoint._dump_ui_metadata(dummy_node, exec_info,
+                                             ui_metadata_path)
+
+      with open(ui_metadata_path) as f:
+        ui_metadata = json.load(f)
+        self.assertLen(ui_metadata['outputs'], 1)
+        self.assertEqual('markdown', ui_metadata['outputs'][0]['type'])
+
   def testOverrideRegisterExecution(self):
     # Mock all real operations of driver / executor / MLMD accesses.
     mock_targets = (  # (cls, method, return_value)