From b6b5ee444d8cff2d3ccd651f846ab75d26d2f5bd Mon Sep 17 00:00:00 2001 From: Yi-Ting Lee Date: Mon, 25 Jul 2022 10:40:00 -0700 Subject: [PATCH 1/4] PyvisVisualizer added --- src/sagemaker/lineage/query.py | 104 ++++++++++++++++++++++++++++++++- 1 file changed, 101 insertions(+), 3 deletions(-) diff --git a/src/sagemaker/lineage/query.py b/src/sagemaker/lineage/query.py index 1ea7baecb6..9b65b77646 100644 --- a/src/sagemaker/lineage/query.py +++ b/src/sagemaker/lineage/query.py @@ -390,6 +390,78 @@ def render(self, elements, mode): return app.run_server(mode=mode) +class PyvisVisualizer(object): + """Create object used for visualizing graph using Pyvis library.""" + + def __init__(self, graph_styles): + """Init for PyvisVisualizer.""" + # import visualization packages + ( + self.pyvis, + self.Network, + self.Options, + ) = self._import_visual_modules() + + self.graph_styles = graph_styles + + def _import_visual_modules(self): + import pyvis + from pyvis.network import Network + from pyvis.options import Options + + return pyvis, Network, Options + + def _get_options(self): + options = """ + var options = { + "configure":{ + "enabled": true + }, + "layout": { + "hierarchical": { + "enabled": true, + "blockShifting": false, + "direction": "LR", + "sortMethod": "directed", + "shakeTowards": "roots" + } + }, + "interaction": { + "multiselect": true, + "navigationButtons": true + }, + "physics": { + "enabled": false, + "hierarchicalRepulsion": { + "centralGravity": 0, + "avoidOverlap": null + }, + "minVelocity": 0.75, + "solver": "hierarchicalRepulsion" + } + } + """ + return options + + def _node_color(self, n): + return self.graph_styles[n[2]]["style"]["background-color"] + + def render(self, elements): + net = self.Network(height='500px', width='100%', notebook = True, directed = True) + options = self._get_options() + net.set_options(options) + + for n in elements["nodes"]: + if(n[3]==True): # startarn + net.add_node(n[0], label=n[1], title=n[1], color=self._node_color(n), shape="star") + else: + net.add_node(n[0], label=n[1], title=n[1], color=self._node_color(n)) + + for e in elements["edges"]: + print(e) + net.add_edge(e[0], e[1], title=e[2]) + + return net.show('pyvisExample.html') class LineageQueryResult(object): """A wrapper around the results of a lineage query.""" @@ -469,6 +541,18 @@ def _covert_edges_to_tuples(self): edges.append((edge.source_arn, edge.destination_arn, edge.association_type)) return edges + def _pyvis_covert_vertices_to_tuples(self): + """Convert vertices to tuple format for visualizer.""" + verts = [] + # get vertex info in the form of (id, label, class) + for vert in self.vertices: + if vert.arn in self.startarn: + # add "startarn" class to node if arn is a startarn + verts.append((vert.arn, vert.lineage_source, vert.lineage_entity, True)) + else: + verts.append((vert.arn, vert.lineage_source, vert.lineage_entity, False)) + return verts + def _get_visualization_elements(self): """Get elements for visualization.""" # get vertices and edges info for graph @@ -488,6 +572,16 @@ def _get_visualization_elements(self): return elements + def _get_pyvis_visualization_elements(self): + verts = self._pyvis_covert_vertices_to_tuples() + edges = self._covert_edges_to_tuples() + + elements = { + "nodes": verts, + "edges": edges + } + return elements + def visualize(self): """Visualize lineage query result.""" elements = self._get_visualization_elements() @@ -523,11 +617,15 @@ def visualize(self): } # initialize DashVisualizer instance to render graph & interactive components - dash_vis = DashVisualizer(lineage_graph) + # dash_vis = DashVisualizer(lineage_graph) + + # dash_server = dash_vis.render(elements=elements, mode="inline") - dash_server = dash_vis.render(elements=elements, mode="inline") + # return dash_server - return dash_server + pyvis_vis = PyvisVisualizer(lineage_graph) + elements = self._get_pyvis_visualization_elements() + return pyvis_vis.render(elements=elements) class LineageFilter(object): From 794fb578bde195e81199d322e8973e66c32d9430 Mon Sep 17 00:00:00 2001 From: Yi-Ting Lee Date: Mon, 25 Jul 2022 14:06:00 -0700 Subject: [PATCH 2/4] modify options --- src/sagemaker/lineage/query.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/sagemaker/lineage/query.py b/src/sagemaker/lineage/query.py index 9b65b77646..69132e35c9 100644 --- a/src/sagemaker/lineage/query.py +++ b/src/sagemaker/lineage/query.py @@ -408,6 +408,7 @@ def _import_visual_modules(self): import pyvis from pyvis.network import Network from pyvis.options import Options + # No module named 'pyvis' return pyvis, Network, Options @@ -415,15 +416,15 @@ def _get_options(self): options = """ var options = { "configure":{ - "enabled": true + "enabled": false }, "layout": { "hierarchical": { "enabled": true, - "blockShifting": false, + "blockShifting": true, "direction": "LR", "sortMethod": "directed", - "shakeTowards": "roots" + "shakeTowards": "leaves" } }, "interaction": { From baeaace2292aa3a67ac3de859f233cc96f5fcb96 Mon Sep 17 00:00:00 2001 From: Yi-Ting Lee Date: Wed, 27 Jul 2022 10:01:23 -0700 Subject: [PATCH 3/4] change: change visualization to using pyvis library --- src/sagemaker/lineage/query.py | 268 ++++----------------------------- 1 file changed, 30 insertions(+), 238 deletions(-) diff --git a/src/sagemaker/lineage/query.py b/src/sagemaker/lineage/query.py index 69132e35c9..b96880209b 100644 --- a/src/sagemaker/lineage/query.py +++ b/src/sagemaker/lineage/query.py @@ -201,18 +201,16 @@ def _artifact_to_lineage_object(self): return Artifact.load(artifact_arn=self.arn, sagemaker_session=self._session) -class DashVisualizer(object): - """Create object used for visualizing graph using Dash library.""" +class PyvisVisualizer(object): + """Create object used for visualizing graph using Pyvis library.""" def __init__(self, graph_styles): - """Init for DashVisualizer.""" + """Init for PyvisVisualizer.""" # import visualization packages ( - self.cyto, - self.JupyterDash, - self.html, - self.Input, - self.Output, + self.pyvis, + self.Network, + self.Options, ) = self._import_visual_modules() self.graph_styles = graph_styles @@ -220,199 +218,30 @@ def __init__(self, graph_styles): def _import_visual_modules(self): """Import modules needed for visualization.""" try: - import dash_cytoscape as cyto + import pyvis except ImportError as e: print(e) - print("Try: pip install dash-cytoscape") + print("Try: pip install pyvis") raise try: - from jupyter_dash import JupyterDash + from pyvis.network import Network except ImportError as e: print(e) - print("Try: pip install jupyter-dash") + print("Try: pip install pyvis") raise try: - from dash import html + from pyvis.options import Options except ImportError as e: print(e) - print("Try: pip install dash") + print("Try: pip install pyvis") raise - try: - from dash.dependencies import Input, Output - except ImportError as e: - print(e) - print("Try: pip install dash") - raise - - return cyto, JupyterDash, html, Input, Output - - def _create_legend_component(self, style): - """Create legend component div.""" - text = style["name"] - symbol = "" - color = "#ffffff" - if style["isShape"] == "False": - color = style["style"]["background-color"] - else: - symbol = style["symbol"] - return self.html.Div( - [ - self.html.Div( - symbol, - style={ - "background-color": color, - "width": "1.5vw", - "height": "1.5vw", - "display": "inline-block", - "font-size": "1.5vw", - }, - ), - self.html.Div( - style={ - "width": "0.5vw", - "height": "1.5vw", - "display": "inline-block", - } - ), - self.html.Div( - text, - style={"display": "inline-block", "font-size": "1.5vw"}, - ), - ] - ) - - def _create_entity_selector(self, entity_name, style): - """Create selector for each lineage entity.""" - return {"selector": "." + entity_name, "style": style["style"]} - - def _get_app(self, elements): - """Create JupyterDash app for interactivity on Jupyter notebook.""" - app = self.JupyterDash(__name__) - self.cyto.load_extra_layouts() - - app.layout = self.html.Div( - [ - # graph section - self.cyto.Cytoscape( - id="cytoscape-graph", - elements=elements, - style={ - "width": "84%", - "height": "350px", - "display": "inline-block", - "border-width": "1vw", - "border-color": "#232f3e", - }, - layout={"name": "klay"}, - stylesheet=[ - { - "selector": "node", - "style": { - "label": "data(label)", - "font-size": "3.5vw", - "height": "10vw", - "width": "10vw", - "border-width": "0.8", - "border-opacity": "0", - "border-color": "#232f3e", - "font-family": "verdana", - }, - }, - { - "selector": "edge", - "style": { - "label": "data(label)", - "color": "gray", - "text-halign": "left", - "text-margin-y": "2.5", - "font-size": "3", - "width": "1", - "curve-style": "bezier", - "control-point-step-size": "15", - "target-arrow-color": "gray", - "target-arrow-shape": "triangle", - "line-color": "gray", - "arrow-scale": "0.5", - "font-family": "verdana", - }, - }, - {"selector": ".select", "style": {"border-opacity": "0.7"}}, - ] - + [self._create_entity_selector(k, v) for k, v in self.graph_styles.items()], - responsive=True, - ), - self.html.Div( - style={ - "width": "0.5%", - "display": "inline-block", - "font-size": "1vw", - "font-family": "verdana", - "vertical-align": "top", - }, - ), - # legend section - self.html.Div( - [self._create_legend_component(v) for k, v in self.graph_styles.items()], - style={ - "display": "inline-block", - "font-size": "1vw", - "font-family": "verdana", - "vertical-align": "top", - }, - ), - ] - ) - - @app.callback( - self.Output("cytoscape-graph", "elements"), - self.Input("cytoscape-graph", "tapNodeData"), - self.Input("cytoscape-graph", "elements"), - ) - def selectNode(tapData, elements): - for n in elements: - if tapData is not None and n["data"]["id"] == tapData["id"]: - # if is tapped node, add "select" class to node - n["classes"] += " select" - elif "classes" in n: - # remove "select" class in "classes" if node not selected - n["classes"] = n["classes"].replace("select", "") - - return elements - - return app - - def render(self, elements, mode): - """Render graph for lineage query result.""" - app = self._get_app(elements) - - return app.run_server(mode=mode) - -class PyvisVisualizer(object): - """Create object used for visualizing graph using Pyvis library.""" - - def __init__(self, graph_styles): - """Init for PyvisVisualizer.""" - # import visualization packages - ( - self.pyvis, - self.Network, - self.Options, - ) = self._import_visual_modules() - - self.graph_styles = graph_styles - - def _import_visual_modules(self): - import pyvis - from pyvis.network import Network - from pyvis.options import Options - # No module named 'pyvis' - return pyvis, Network, Options def _get_options(self): + """Get pyvis graph options.""" options = """ var options = { "configure":{ @@ -445,24 +274,29 @@ def _get_options(self): return options def _node_color(self, n): + """Return node color by background-color specified in graph styles.""" return self.graph_styles[n[2]]["style"]["background-color"] - def render(self, elements): - net = self.Network(height='500px', width='100%', notebook = True, directed = True) + def render(self, elements, path="pyvisExample.html"): + """Render graph for lineage query result.""" + net = self.Network(height="500px", width="100%", notebook=True, directed=True) options = self._get_options() - net.set_options(options) + net.set_options(options) + # add nodes to graph for n in elements["nodes"]: - if(n[3]==True): # startarn - net.add_node(n[0], label=n[1], title=n[1], color=self._node_color(n), shape="star") + if n[3]: # startarn + net.add_node(n[0], label=n[1], title=n[2], color=self._node_color(n), shape="star") else: - net.add_node(n[0], label=n[1], title=n[1], color=self._node_color(n)) + net.add_node(n[0], label=n[1], title=n[2], color=self._node_color(n)) + # add edges to graph for e in elements["edges"]: print(e) - net.add_edge(e[0], e[1], title=e[2]) + net.add_edge(e[0], e[1], title=e[2]) + + return net.show(path) - return net.show('pyvisExample.html') class LineageQueryResult(object): """A wrapper around the results of a lineage query.""" @@ -522,18 +356,6 @@ def __str__(self): result_dict = vars(self) return str({k: [str(val) for val in v] for k, v in result_dict.items()}) - def _covert_vertices_to_tuples(self): - """Convert vertices to tuple format for visualizer.""" - verts = [] - # get vertex info in the form of (id, label, class) - for vert in self.vertices: - if vert.arn in self.startarn: - # add "startarn" class to node if arn is a startarn - verts.append((vert.arn, vert.lineage_source, vert.lineage_entity + " startarn")) - else: - verts.append((vert.arn, vert.lineage_source, vert.lineage_entity)) - return verts - def _covert_edges_to_tuples(self): """Convert edges to tuple format for visualizer.""" edges = [] @@ -542,7 +364,7 @@ def _covert_edges_to_tuples(self): edges.append((edge.source_arn, edge.destination_arn, edge.association_type)) return edges - def _pyvis_covert_vertices_to_tuples(self): + def _covert_vertices_to_tuples(self): """Convert vertices to tuple format for visualizer.""" verts = [] # get vertex info in the form of (id, label, class) @@ -555,38 +377,15 @@ def _pyvis_covert_vertices_to_tuples(self): return verts def _get_visualization_elements(self): - """Get elements for visualization.""" - # get vertices and edges info for graph + """Get elements(nodes+edges) for visualization.""" verts = self._covert_vertices_to_tuples() edges = self._covert_edges_to_tuples() - nodes = [ - {"data": {"id": id, "label": label}, "classes": classes} for id, label, classes in verts - ] - - edges = [ - {"data": {"source": source, "target": target, "label": label}} - for source, target, label in edges - ] - - elements = nodes + edges - - return elements - - def _get_pyvis_visualization_elements(self): - verts = self._pyvis_covert_vertices_to_tuples() - edges = self._covert_edges_to_tuples() - - elements = { - "nodes": verts, - "edges": edges - } + elements = {"nodes": verts, "edges": edges} return elements def visualize(self): """Visualize lineage query result.""" - elements = self._get_visualization_elements() - lineage_graph = { # nodes can have shape / color "TrialComponent": { @@ -617,15 +416,8 @@ def visualize(self): }, } - # initialize DashVisualizer instance to render graph & interactive components - # dash_vis = DashVisualizer(lineage_graph) - - # dash_server = dash_vis.render(elements=elements, mode="inline") - - # return dash_server - pyvis_vis = PyvisVisualizer(lineage_graph) - elements = self._get_pyvis_visualization_elements() + elements = self._get_visualization_elements() return pyvis_vis.render(elements=elements) From 3bab76a2a2854d6f73ea71e55e15a9cf280a7a40 Mon Sep 17 00:00:00 2001 From: Yi-Ting Lee Date: Thu, 28 Jul 2022 10:25:37 -0700 Subject: [PATCH 4/4] used get_module() method for import & code style change --- .gitignore | 3 +- src/sagemaker/lineage/query.py | 64 ++++++++++++---------------------- 2 files changed, 24 insertions(+), 43 deletions(-) diff --git a/.gitignore b/.gitignore index 5b496055e9..7a61658c20 100644 --- a/.gitignore +++ b/.gitignore @@ -29,4 +29,5 @@ venv/ env/ .vscode/ **/tmp -.python-version \ No newline at end of file +.python-version +*.html \ No newline at end of file diff --git a/src/sagemaker/lineage/query.py b/src/sagemaker/lineage/query.py index b96880209b..5f3164bdc0 100644 --- a/src/sagemaker/lineage/query.py +++ b/src/sagemaker/lineage/query.py @@ -17,7 +17,7 @@ from enum import Enum from typing import Optional, Union, List, Dict -from sagemaker.lineage._utils import get_resource_name_from_arn +from sagemaker.lineage._utils import get_resource_name_from_arn, get_module class LineageEntityEnum(Enum): @@ -208,41 +208,14 @@ def __init__(self, graph_styles): """Init for PyvisVisualizer.""" # import visualization packages ( - self.pyvis, self.Network, self.Options, ) = self._import_visual_modules() self.graph_styles = graph_styles - def _import_visual_modules(self): - """Import modules needed for visualization.""" - try: - import pyvis - except ImportError as e: - print(e) - print("Try: pip install pyvis") - raise - - try: - from pyvis.network import Network - except ImportError as e: - print(e) - print("Try: pip install pyvis") - raise - - try: - from pyvis.options import Options - except ImportError as e: - print(e) - print("Try: pip install pyvis") - raise - - return pyvis, Network, Options - - def _get_options(self): - """Get pyvis graph options.""" - options = """ + # pyvis graph options + self._options = """ var options = { "configure":{ "enabled": false @@ -271,29 +244,36 @@ def _get_options(self): } } """ - return options - def _node_color(self, n): + def _import_visual_modules(self): + """Import modules needed for visualization.""" + get_module("pyvis") + from pyvis.network import Network + from pyvis.options import Options + + return Network, Options + + def _node_color(self, entity): """Return node color by background-color specified in graph styles.""" - return self.graph_styles[n[2]]["style"]["background-color"] + return self.graph_styles[entity]["style"]["background-color"] def render(self, elements, path="pyvisExample.html"): """Render graph for lineage query result.""" net = self.Network(height="500px", width="100%", notebook=True, directed=True) - options = self._get_options() - net.set_options(options) + net.set_options(self._options) # add nodes to graph - for n in elements["nodes"]: - if n[3]: # startarn - net.add_node(n[0], label=n[1], title=n[2], color=self._node_color(n), shape="star") + 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" + ) else: - net.add_node(n[0], label=n[1], title=n[2], color=self._node_color(n)) + net.add_node(arn, label=source, title=entity, color=self._node_color(entity)) # add edges to graph - for e in elements["edges"]: - print(e) - net.add_edge(e[0], e[1], title=e[2]) + for src, dest, asso_type in elements["edges"]: + net.add_edge(src, dest, title=asso_type) return net.show(path)