change: lineage query visualization experience enhancement

Author: Yi-Ting Lee

src/sagemaker/lineage/query.py

@@ -15,7 +15,6 @@
 
 from datetime import datetime
 from enum import Enum
-from platform import node
 from typing import Optional, Union, List, Dict
 import re
 
@@ -263,7 +262,7 @@ def __init__(self, graph_styles):
             self.Network,
             self.Options,
             self.IFrame,
-            self.BeautifulSoup
+            self.BeautifulSoup,
         ) = self._import_visual_modules()
 
         self.graph_styles = graph_styles
@@ -316,9 +315,53 @@ def _node_color(self, entity):
         return self.graph_styles[entity]["style"]["background-color"]
 
     def _add_legend(self, path):
-        f = open(path, "r+")
-        soup = self.BeautifulSoup(f, 'html.parser')
-        print(soup.prettify())
+        """Embed legend to html file generated by pyvis."""
+        f = open(path, "r")
+        content = self.BeautifulSoup(f, "html.parser")
+
+        legend = """
+            <div style="display: inline-block; font-size: 1vw; font-family: verdana;
+                vertical-align: top; padding: 1vw;">
+                <div>
+                    <div style="background-color: rgb(246, 207, 97); width: 1.6vw; height: 1.6vw;
+                        display: inline-block; font-size: 1.5vw; vertical-align: -0.2em;"></div>
+                    <div style="width: 0.3vw; height: 1.5vw; display: inline-block;"></div>
+                    <div style="display: inline-block; font-size: 1.5vw;">Trial Component</div>
+                </div>
+                <div>
+                    <div style="background-color: rgb(255, 153, 0); width: 1.6vw; height: 1.6vw;
+                        display: inline-block; font-size: 1.5vw; vertical-align: -0.2em;"></div>
+                    <div style="width: 0.3vw; height: 1.5vw; display: inline-block;"></div>
+                    <div style="display: inline-block; font-size: 1.5vw;">Context</div>
+                </div>
+                <div>
+                    <div style="background-color: rgb(136, 195, 150); width: 1.6vw; height: 1.6vw;
+                        display: inline-block; font-size: 1.5vw; vertical-align: -0.2em;"></div>
+                    <div style="width: 0.3vw; height: 1.5vw; display: inline-block;"></div>
+                    <div style="display: inline-block; font-size: 1.5vw;">Action</div>
+                </div>
+                <div>
+                    <div style="background-color: rgb(20, 110, 180); width: 1.6vw; height: 1.6vw;
+                        display: inline-block; font-size: 1.5vw; vertical-align: -0.2em;"></div>
+                    <div style="width: 0.3vw; height: 1.5vw; display: inline-block;"></div>
+                    <div style="display: inline-block; font-size: 1.5vw;">Artifact</div>
+                </div>
+                <div>
+                    <div style="background-color: rgb(255, 255, 255); width: 1.6vw; height: 1.6vw;
+                        display: inline-block; font-size: 0.9vw;">star</div>
+                    <div style="width: 0.3vw; height: 1.5vw; display: inline-block;"></div>
+                    <div style="display: inline-block; font-size: 1.5vw;">StartArn</div>
+                </div>
+            </div>
+        """
+        legend_div = self.BeautifulSoup(legend, "html.parser")
+
+        content.div.insert_after(legend_div)
+
+        html = content.prettify()
+
+        with open(path, "w", encoding="utf8") as file:
+            file.write(html)
 
     def render(self, elements, path="pyvisExample.html"):
         """Render graph for lineage query result.
@@ -338,19 +381,42 @@ def render(self, elements, path="pyvisExample.html"):
             display graph: The interactive visualization is presented as a static HTML file.
 
         """
-        net = self.Network(height="500px", width="100%", notebook=True, directed=True)
+        net = self.Network(height="600px", width="82%", notebook=True, directed=True)
         net.set_options(self._options)
 
         # add nodes to graph
         for arn, source, entity, is_start_arn in elements["nodes"]:
+            entity_text = re.sub(r"(\w)([A-Z])", r"\1 \2", entity)
             source = re.sub(r"(\w)([A-Z])", r"\1 \2", source)
-            node_info = "Entity: " + entity + "\n" + "Type: " + source + "\n" + "Name: " + arn
+            account_id = re.search(r":\d{12}:", arn)
+            name = re.search(r"\/.*", arn)
+            node_info = (
+                "Entity: "
+                + entity_text
+                + "\nType: "
+                + source
+                + "\nAccount ID: "
+                + str(account_id.group()[1:-1])
+                + "\nName: "
+                + str(name.group()[1:])
+            )
             if is_start_arn:  # startarn
                 net.add_node(
-                    arn, label=source, title=node_info, color=self._node_color(entity), shape="star", borderWidth=3
+                    arn,
+                    label=source,
+                    title=node_info,
+                    color=self._node_color(entity),
+                    shape="star",
+                    borderWidth=3,
                 )
             else:
-                net.add_node(arn, label=source, title=node_info, color=self._node_color(entity), borderWidth=3)
+                net.add_node(
+                    arn,
+                    label=source,
+                    title=node_info,
+                    color=self._node_color(entity),
+                    borderWidth=3,
+                )
 
         # add edges to graph
         for src, dest, asso_type in elements["edges"]:
@@ -359,7 +425,7 @@ def render(self, elements, path="pyvisExample.html"):
         net.write_html(path)
         self._add_legend(path)
 
-        return self.IFrame(path, width="100%", height="500px")
+        return self.IFrame(path, width="100%", height="600px")
 
 
 class LineageQueryResult(object):

tests/integ/sagemaker/lineage/test_lineage_visualize.py

@@ -14,6 +14,7 @@
 from __future__ import absolute_import
 import time
 import os
+import re
 
 import pytest
 
@@ -160,31 +161,56 @@ def test_graph_visualize(sagemaker_session, extract_data_from_html):
                 "color": "#146eb4",
                 "label": "Model",
                 "shape": "star",
-                "title": "Artifact",
+                "title": "Entity: Artifact"
+                + "\nType: Model"
+                + "\nAccount ID: "
+                + str(re.search(r":\d{12}:", graph_startarn).group()[1:-1])
+                + "\nName: "
+                + str(re.search(r"\/.*", graph_startarn).group()[1:]),
             },
             image_artifact: {
                 "color": "#146eb4",
                 "label": "Image",
                 "shape": "dot",
-                "title": "Artifact",
+                "title": "Entity: Artifact"
+                + "\nType: Image"
+                + "\nAccount ID: "
+                + str(re.search(r":\d{12}:", image_artifact).group()[1:-1])
+                + "\nName: "
+                + str(re.search(r"\/.*", image_artifact).group()[1:]),
             },
             dataset_artifact: {
                 "color": "#146eb4",
-                "label": "DataSet",
+                "label": "Data Set",
                 "shape": "dot",
-                "title": "Artifact",
+                "title": "Entity: Artifact"
+                + "\nType: Data Set"
+                + "\nAccount ID: "
+                + str(re.search(r":\d{12}:", dataset_artifact).group()[1:-1])
+                + "\nName: "
+                + str(re.search(r"\/.*", dataset_artifact).group()[1:]),
             },
             modeldeploy_action: {
                 "color": "#88c396",
-                "label": "ModelDeploy",
+                "label": "Model Deploy",
                 "shape": "dot",
-                "title": "Action",
+                "title": "Entity: Action"
+                + "\nType: Model Deploy"
+                + "\nAccount ID: "
+                + str(re.search(r":\d{12}:", modeldeploy_action).group()[1:-1])
+                + "\nName: "
+                + str(re.search(r"\/.*", modeldeploy_action).group()[1:]),
             },
             endpoint_context: {
                 "color": "#ff9900",
                 "label": "Endpoint",
                 "shape": "dot",
-                "title": "Context",
+                "title": "Entity: Context"
+                + "\nType: Endpoint"
+                + "\nAccount ID: "
+                + str(re.search(r":\d{12}:", endpoint_context).group()[1:-1])
+                + "\nName: "
+                + str(re.search(r"\/.*", endpoint_context).group()[1:]),
             },
         }
 

tests/unit/sagemaker/lineage/test_query.py

@@ -580,7 +580,7 @@ def test_query_lineage_result_str(sagemaker_session):
 
     assert (
         response_str
-        == "{\n'edges': [\n\t{'source_arn': 'arn1', 'destination_arn': 'arn2', 'association_type': 'Produced'}],"
+        == "{'edges': [\n\t{'source_arn': 'arn1', 'destination_arn': 'arn2', 'association_type': 'Produced'}],"
         + "\n\n'vertices': [\n\t{'arn': 'arn1', 'lineage_entity': 'Artifact', 'lineage_source': 'Endpoint', "
         + "'_session': <Mock id=''>}, \n\t{'arn': 'arn2', 'lineage_entity': 'Context', 'lineage_source': "
         + "'Model', '_session': <Mock id=''>}],\n\n'startarn': "