Skip to content

feature: query lineage visualizer advanced styling & interactive component handling #4

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 17 commits into from
Jul 26, 2022
Merged

feature: query lineage visualizer advanced styling & interactive component handling #4

Changes from 16 commits
Commits
Show all changes
17 commits
Select commit Hold shift + click to select a range
7359b3d
feature: query lineage visualizer for general case
Jul 8, 2022
d524416
startarn added to lineageQueryResult
Jul 14, 2022
acf64f4
color node by lineage entity
Jul 14, 2022
e07e83c
identify startarn node by shape
Jul 14, 2022
687d6b6
Add code comments
Jul 14, 2022
e3a4c9d
Double sided arrows handled
Jul 15, 2022
1808f87
legend added
Jul 18, 2022
60904d5
try except raise
Jul 19, 2022
9125eb9
startarn added
Jul 19, 2022
3db7fa1
add get element function
Jul 19, 2022
c135ab9
add DashVisualizer class
Jul 19, 2022
daa5cc4
style check
Jul 19, 2022
0adc4db
Merge branch 'master' into advanced-style
Jul 19, 2022
28a9eeb
feature: query lineage visualizer advanced styling & interactive comp…
Jul 19, 2022
67346a7
Merge pull request #1 from Snoyark/master
ytlee93 Jul 20, 2022
7c2b0c3
add functions that generate html components and style selectors
Jul 20, 2022
41d2453
inject graph data to DashVisualizer task
Jul 21, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
177 changes: 149 additions & 28 deletions src/sagemaker/lineage/query.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -207,29 +207,84 @@ class DashVisualizer(object):
def __init__(self):
"""Init for DashVisualizer."""
# import visualization packages
self.cyto, self.JupyterDash, self.html = self._import_visual_modules()
(
self.cyto,
self.JupyterDash,
self.html,
self.Input,
self.Output,
) = self._import_visual_modules()

self.entity_color = {
"TrialComponent": "#f6cf61",
"Context": "#ff9900",
"Action": "#88c396",
"Artifact": "#146eb4",
}

def _import_visual_modules(self):
"""Import modules needed for visualization."""
try:
import dash_cytoscape as cyto
except ImportError as e:
print(e)
print("try pip install dash-cytoscape")
print("Try: pip install dash-cytoscape")
raise

try:
from jupyter_dash import JupyterDash
except ImportError as e:
print(e)
print("try pip install jupyter-dash")
print("Try: pip install jupyter-dash")
raise

try:
from dash import html
except ImportError as e:
print(e)
print("try pip install dash")
print("Try: pip install dash")
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, text, color, colorText=""):
"""Create legend component div."""
return self.html.Div(
[
self.html.Div(
colorText,
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"},
),
]
)

return cyto, JupyterDash, html
def _create_entity_selector(self, entity_name, color):
"""Create selector for each lineage entity."""
return {"selector": "." + entity_name, "style": {"background-color": color}}

def _get_app(self, elements):
"""Create JupyterDash app for interactivity on Jupyter notebook."""
Expand All @@ -238,10 +293,17 @@ def _get_app(self, elements):

app.layout = self.html.Div(
[
# graph section
self.cyto.Cytoscape(
id="cytoscape-layout-1",
id="cytoscape-graph",
elements=elements,
style={"width": "100%", "height": "350px"},
style={
"width": "84%",
"height": "350px",
"display": "inline-block",
"border-width": "1vw",
"border-color": "#232f3e",
},
layout={"name": "klay"},
stylesheet=[
{
Expand All @@ -251,6 +313,10 @@ def _get_app(self, elements):
"font-size": "3.5vw",
"height": "10vw",
"width": "10vw",
"border-width": "0.8",
"border-opacity": "0",
"border-color": "#232f3e",
"font-family": "verdana",
},
},
{
Expand All @@ -259,23 +325,63 @@ def _get_app(self, elements):
"label": "data(label)",
"color": "gray",
"text-halign": "left",
"text-margin-y": "3px",
"text-margin-x": "-2px",
"font-size": "3%",
"width": "1%",
"curve-style": "taxi",
"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": ".startarn", "style": {"shape": "star"}},
{"selector": ".select", "style": {"border-opacity": "0.7"}},
]
+ [self._create_entity_selector(k, v) for k, v in self.entity_color.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(k, v) for k, v in self.entity_color.items()]
Copy link
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Won't this result in the string TrialComponent being shown instead of Trial Component? Could either string split on capitalization or just include the title in the config as well as the selector id

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

resolved

+ [self._create_legend_component("StartArn", "#ffffff", "★")],
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):
Expand All @@ -292,6 +398,7 @@ def __init__(
self,
edges: List[Edge] = None,
vertices: List[Vertex] = None,
startarn: List[str] = None,
):
"""Init for LineageQueryResult.

Expand All @@ -301,63 +408,75 @@ def __init__(
"""
self.edges = []
self.vertices = []
self.startarn = []

if edges is not None:
self.edges = edges

if vertices is not None:
self.vertices = vertices

if startarn is not None:
self.startarn = startarn

def __str__(self):
"""Define string representation of ``LineageQueryResult``.

Format:
{
'edges':[
{
"{
'source_arn': 'string', 'destination_arn': 'string',
'association_type': 'string'
},
}",
...
]
],
'vertices':[
{
"{
'arn': 'string', 'lineage_entity': 'string',
'lineage_source': 'string',
'_session': <sagemaker.session.Session object>
},
}",
...
],
'startarn':[
'string',
...
]
}

"""
result_dict = vars(self)
return str({k: [vars(val) for val in v] for k, v in result_dict.items()})
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:
verts.append((vert.arn, vert.lineage_source))
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 = []
# get edge info in the form of (source, target, label)
for edge in self.edges:
edges.append((edge.source_arn, edge.destination_arn, edge.association_type))
return edges

def _get_visualization_elements(self):
"""Get elements for visualization."""
# get vertices and edges info for graph
verts = self._covert_vertices_to_tuples()
edges = self._covert_edges_to_tuples()

nodes = [
{
"data": {"id": id, "label": label},
}
for id, label in verts
{"data": {"id": id, "label": label}, "classes": classes} for id, label, classes in verts
]

edges = [
Expand All @@ -373,6 +492,7 @@ def visualize(self):
"""Visualize lineage query result."""
elements = self._get_visualization_elements()

# initialize DashVisualizer instance to render graph & interactive components
dash_vis = DashVisualizer()

dash_server = dash_vis.render(elements=elements, mode="inline")
Expand Down Expand Up @@ -453,9 +573,8 @@ def _get_vertex(self, vertex):
sagemaker_session=self._session,
)

def _convert_api_response(self, response) -> LineageQueryResult:
def _convert_api_response(self, response, converted) -> LineageQueryResult:
"""Convert the lineage query API response to its Python representation."""
converted = LineageQueryResult()
converted.edges = [self._get_edge(edge) for edge in response["Edges"]]
converted.vertices = [self._get_vertex(vertex) for vertex in response["Vertices"]]

Expand Down Expand Up @@ -538,7 +657,9 @@ def query(
Filters=query_filter._to_request_dict() if query_filter else {},
MaxDepth=max_depth,
)
query_response = self._convert_api_response(query_response)
# create query result for startarn info
query_result = LineageQueryResult(startarn=start_arns)
query_response = self._convert_api_response(query_response, query_result)
query_response = self._collapse_cross_account_artifacts(query_response)

return query_response