From f54e1b600d575f9b7560559a57156f2983899f29 Mon Sep 17 00:00:00 2001 From: Yi-Ting Lee Date: Tue, 9 Aug 2022 11:56:56 -0700 Subject: [PATCH 1/6] graph style fine-tune --- src/sagemaker/lineage/query.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/sagemaker/lineage/query.py b/src/sagemaker/lineage/query.py index f6966d9fa3..bcd4c26a85 100644 --- a/src/sagemaker/lineage/query.py +++ b/src/sagemaker/lineage/query.py @@ -332,14 +332,14 @@ def render(self, elements, path="pyvisExample.html"): for arn, source, entity, is_start_arn in elements["nodes"]: if is_start_arn: # startarn net.add_node( - arn, label=source, title=entity, color=self._node_color(entity), shape="star" + arn, label=source, title=entity, color=self._node_color(entity), shape="star", borderWidth=3 ) else: - net.add_node(arn, label=source, title=entity, color=self._node_color(entity)) + net.add_node(arn, label=source, title=entity, color=self._node_color(entity), borderWidth=3) # add edges to graph for src, dest, asso_type in elements["edges"]: - net.add_edge(src, dest, title=asso_type) + net.add_edge(src, dest, title=asso_type, width=2) return net.show(path) From ff136c10471c3d7c90597abd9de42fbc07949a70 Mon Sep 17 00:00:00 2001 From: Yi-Ting Lee Date: Tue, 9 Aug 2022 16:03:02 -0700 Subject: [PATCH 2/6] query visualize more info on hover node --- src/sagemaker/lineage/query.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/sagemaker/lineage/query.py b/src/sagemaker/lineage/query.py index bcd4c26a85..9e257d7733 100644 --- a/src/sagemaker/lineage/query.py +++ b/src/sagemaker/lineage/query.py @@ -332,10 +332,10 @@ def render(self, elements, path="pyvisExample.html"): for arn, source, entity, is_start_arn in elements["nodes"]: if is_start_arn: # startarn net.add_node( - arn, label=source, title=entity, color=self._node_color(entity), shape="star", borderWidth=3 + arn, label=source, title=entity+"\n"+arn, color=self._node_color(entity), shape="star", borderWidth=3 ) else: - net.add_node(arn, label=source, title=entity, color=self._node_color(entity), borderWidth=3) + net.add_node(arn, label=source, title=entity+"\n"+arn, color=self._node_color(entity), borderWidth=3) # add edges to graph for src, dest, asso_type in elements["edges"]: From 1a3b5f3b64783fef36bc846f9af0de6c02ad91ec Mon Sep 17 00:00:00 2001 From: Yi-Ting Lee Date: Wed, 10 Aug 2022 14:25:02 -0700 Subject: [PATCH 3/6] info on hover --- src/sagemaker/lineage/query.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/src/sagemaker/lineage/query.py b/src/sagemaker/lineage/query.py index 9e257d7733..348bd46697 100644 --- a/src/sagemaker/lineage/query.py +++ b/src/sagemaker/lineage/query.py @@ -15,7 +15,9 @@ from datetime import datetime from enum import Enum +from platform import node from typing import Optional, Union, List, Dict +import re from sagemaker.lineage._utils import get_resource_name_from_arn, get_module @@ -330,12 +332,14 @@ def render(self, elements, path="pyvisExample.html"): # add nodes to graph for arn, source, entity, is_start_arn in elements["nodes"]: + source = re.sub(r"(\w)([A-Z])", r"\1 \2", source) + node_info = "Entity: " + entity + "\n" + "Type: " + source + "\n" + "Name: " + arn if is_start_arn: # startarn net.add_node( - arn, label=source, title=entity+"\n"+arn, 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=entity+"\n"+arn, 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"]: @@ -391,7 +395,7 @@ def __str__(self): """ return ( - "{\n" + "{" + "\n\n".join("'{}': {},".format(key, val) for key, val in self.__dict__.items()) + "\n}" ) From a9e51141ccffb37f85b7559e84f3802dd61130c2 Mon Sep 17 00:00:00 2001 From: Yi-Ting Lee Date: Thu, 11 Aug 2022 14:46:52 -0700 Subject: [PATCH 4/6] split generate html file & display --- src/sagemaker/lineage/query.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/src/sagemaker/lineage/query.py b/src/sagemaker/lineage/query.py index 348bd46697..8258be0273 100644 --- a/src/sagemaker/lineage/query.py +++ b/src/sagemaker/lineage/query.py @@ -262,6 +262,8 @@ def __init__(self, graph_styles): ( self.Network, self.Options, + self.IFrame, + self.BeautifulSoup ) = self._import_visual_modules() self.graph_styles = graph_styles @@ -302,13 +304,22 @@ def _import_visual_modules(self): get_module("pyvis") from pyvis.network import Network from pyvis.options import Options + from IPython.display import IFrame - return Network, Options + get_module("bs4") + from bs4 import BeautifulSoup + + return Network, Options, IFrame, BeautifulSoup def _node_color(self, entity): """Return node color by background-color specified in graph styles.""" 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()) + def render(self, elements, path="pyvisExample.html"): """Render graph for lineage query result. @@ -345,7 +356,10 @@ def render(self, elements, path="pyvisExample.html"): for src, dest, asso_type in elements["edges"]: net.add_edge(src, dest, title=asso_type, width=2) - return net.show(path) + net.write_html(path) + self._add_legend(path) + + return self.IFrame(path, width="100%", height="500px") class LineageQueryResult(object): From c1be76563914857aebea563cd4d9665c34331e76 Mon Sep 17 00:00:00 2001 From: Yi-Ting Lee Date: Fri, 12 Aug 2022 11:35:51 -0700 Subject: [PATCH 5/6] change: lineage query visualization experience enhancement --- src/sagemaker/lineage/query.py | 86 ++++++++++++++++--- .../lineage/test_lineage_visualize.py | 40 +++++++-- tests/unit/sagemaker/lineage/test_query.py | 2 +- 3 files changed, 110 insertions(+), 18 deletions(-) diff --git a/src/sagemaker/lineage/query.py b/src/sagemaker/lineage/query.py index 8258be0273..31885d3e7f 100644 --- a/src/sagemaker/lineage/query.py +++ b/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 = """ +
+
+
+
+
Trial Component
+
+
+
+
+
Context
+
+
+
+
+
Action
+
+
+
+
+
Artifact
+
+
+
star
+
+
StartArn
+
+
+ """ + 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): diff --git a/tests/integ/sagemaker/lineage/test_lineage_visualize.py b/tests/integ/sagemaker/lineage/test_lineage_visualize.py index 555b98452e..4b9e816623 100644 --- a/tests/integ/sagemaker/lineage/test_lineage_visualize.py +++ b/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:]), }, } diff --git a/tests/unit/sagemaker/lineage/test_query.py b/tests/unit/sagemaker/lineage/test_query.py index 0a357eb1fc..bac5cb6cdb 100644 --- a/tests/unit/sagemaker/lineage/test_query.py +++ b/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': }, \n\t{'arn': 'arn2', 'lineage_entity': 'Context', 'lineage_source': " + "'Model', '_session': }],\n\n'startarn': " From e1190092b79544ed3f978533a5939225f536ce03 Mon Sep 17 00:00:00 2001 From: Yi-Ting Lee Date: Fri, 12 Aug 2022 16:33:34 -0700 Subject: [PATCH 6/6] generate legend divs programmatically --- src/sagemaker/lineage/query.py | 56 +++++++++++++++------------------- 1 file changed, 25 insertions(+), 31 deletions(-) diff --git a/src/sagemaker/lineage/query.py b/src/sagemaker/lineage/query.py index 31885d3e7f..659f88a59c 100644 --- a/src/sagemaker/lineage/query.py +++ b/src/sagemaker/lineage/query.py @@ -314,6 +314,25 @@ def _node_color(self, entity): """Return node color by background-color specified in graph styles.""" return self.graph_styles[entity]["style"]["background-color"] + def _get_legend_line(self, component_name): + """Generate lengend div line for each graph component in graph_styles.""" + if self.graph_styles[component_name]["isShape"] == "False": + return '
\ +
\ +
{name}
'.format( + color=self.graph_styles[component_name]["style"]["background-color"], + name=self.graph_styles[component_name]["name"], + ) + else: + return '
{shape}
\ +
\ +
{name}
'.format( + shape=self.graph_styles[component_name]["style"]["shape"], + name=self.graph_styles[component_name]["name"], + ) + def _add_legend(self, path): """Embed legend to html file generated by pyvis.""" f = open(path, "r") @@ -322,38 +341,13 @@ def _add_legend(self, path): legend = """
-
-
-
-
Trial Component
-
-
-
-
-
Context
-
-
-
-
-
Action
-
-
-
-
-
Artifact
-
-
-
star
-
-
StartArn
-
-
""" + # iterate through graph styles to get legend + for component in self.graph_styles.keys(): + legend += self._get_legend_line(component_name=component) + + legend += "" + legend_div = self.BeautifulSoup(legend, "html.parser") content.div.insert_after(legend_div)