From 5609e429b00ad529bceafc4fef2b942399600815 Mon Sep 17 00:00:00 2001 From: Yi-Ting Lee Date: Fri, 29 Jul 2022 10:00:29 -0700 Subject: [PATCH 1/5] test_get_visualization_elements added --- src/sagemaker/lineage/query.py | 4 ++-- tests/integ/sagemaker/lineage/helpers.py | 4 +++- tests/unit/sagemaker/lineage/test_query.py | 28 ++++++++++++++++++++++ 3 files changed, 33 insertions(+), 3 deletions(-) diff --git a/src/sagemaker/lineage/query.py b/src/sagemaker/lineage/query.py index 5f3164bdc0..a46e88867a 100644 --- a/src/sagemaker/lineage/query.py +++ b/src/sagemaker/lineage/query.py @@ -364,7 +364,7 @@ def _get_visualization_elements(self): elements = {"nodes": verts, "edges": edges} return elements - def visualize(self): + def visualize(self, path="pyvisExample.html"): """Visualize lineage query result.""" lineage_graph = { # nodes can have shape / color @@ -398,7 +398,7 @@ def visualize(self): pyvis_vis = PyvisVisualizer(lineage_graph) elements = self._get_visualization_elements() - return pyvis_vis.render(elements=elements) + return pyvis_vis.render(elements=elements, path=path) class LineageFilter(object): diff --git a/tests/integ/sagemaker/lineage/helpers.py b/tests/integ/sagemaker/lineage/helpers.py index fb71d1d88c..4ed78127bb 100644 --- a/tests/integ/sagemaker/lineage/helpers.py +++ b/tests/integ/sagemaker/lineage/helpers.py @@ -16,7 +16,8 @@ import uuid from datetime import datetime import time - +import boto3 +from botocore.config import Config def name(): return "lineage-integ-{}-{}".format( @@ -78,3 +79,4 @@ def visit(arn, visited: set): ret = [] return visit(start_arn, set()) + diff --git a/tests/unit/sagemaker/lineage/test_query.py b/tests/unit/sagemaker/lineage/test_query.py index ae76fd199c..b5b809138d 100644 --- a/tests/unit/sagemaker/lineage/test_query.py +++ b/tests/unit/sagemaker/lineage/test_query.py @@ -524,3 +524,31 @@ def test_vertex_to_object_unconvertable(sagemaker_session): with pytest.raises(ValueError): vertex.to_lineage_object() + + +def test_get_visualization_elements(sagemaker_session): + lineage_query = LineageQuery(sagemaker_session) + sagemaker_session.sagemaker_client.query_lineage.return_value = { + "Vertices": [ + {"Arn": "arn1", "Type": "Endpoint", "LineageType": "Artifact"}, + {"Arn": "arn2", "Type": "Model", "LineageType": "Context"}, + ], + "Edges": [{"SourceArn": "arn1", "DestinationArn": "arn2", "AssociationType": "Produced"}], + } + + query_response = lineage_query.query( + start_arns=["arn:aws:sagemaker:us-west-2:0123456789012:context/mycontext"] + ) + + print(query_response) + + elements = query_response._get_visualization_elements() + + print(elements) + + assert elements["nodes"][0] == ("arn1", "Endpoint", "Artifact", False) + assert elements["nodes"][1] == ("arn2", "Model", "Context", False) + assert elements["edges"][0] == ("arn1", "arn2", "Produced") + + + From 628ba0ff3e4c327fb9eba2532090a2d7bec60a62 Mon Sep 17 00:00:00 2001 From: Yi-Ting Lee Date: Tue, 2 Aug 2022 09:55:08 -0700 Subject: [PATCH 2/5] __str__ function update --- src/sagemaker/lineage/query.py | 51 ++++++++++++++++++++++------------ 1 file changed, 33 insertions(+), 18 deletions(-) diff --git a/src/sagemaker/lineage/query.py b/src/sagemaker/lineage/query.py index a46e88867a..5277a2d806 100644 --- a/src/sagemaker/lineage/query.py +++ b/src/sagemaker/lineage/query.py @@ -104,6 +104,18 @@ def __str__(self): """ return str(self.__dict__) + def __repr__(self): + """Define string representation of ``Edge``. + + Format: + { + 'source_arn': 'string', 'destination_arn': 'string', + 'association_type': 'string' + } + + """ + return "\n\t" + str(self.__dict__) + class Vertex: """A vertex for a lineage graph.""" @@ -155,6 +167,19 @@ def __str__(self): """ return str(self.__dict__) + def __repr__(self): + """Define string representation of ``Vertex``. + + Format: + { + 'arn': 'string', 'lineage_entity': 'string', + 'lineage_source': 'string', + '_session': + } + + """ + return "\n\t" + str(self.__dict__) + def to_lineage_object(self): """Convert the ``Vertex`` object to its corresponding lineage object. @@ -312,29 +337,19 @@ def __str__(self): Format: { 'edges':[ - "{ - 'source_arn': 'string', 'destination_arn': 'string', - 'association_type': 'string' - }", - ... - ], + {'source_arn': 'string', 'destination_arn': 'string', 'association_type': 'string'}, + ...], + 'vertices':[ - "{ - 'arn': 'string', 'lineage_entity': 'string', - 'lineage_source': 'string', - '_session': - }", - ... - ], - 'startarn':[ - 'string', - ... - ] + {'arn': 'string', 'lineage_entity': 'string', 'lineage_source': 'string', '_session': }, + ...], + + 'startarn':['string', ...] } """ result_dict = vars(self) - return str({k: [str(val) for val in v] for k, v in result_dict.items()}) + return '{\n' + '\n\n'.join('\'{}\': {},'.format(key, val) for key, val in self.__dict__.items()) + '\n}' def _covert_edges_to_tuples(self): """Convert edges to tuple format for visualizer.""" From f9d0ae126dcb8ef740b560e5c8af4892ac9f20a1 Mon Sep 17 00:00:00 2001 From: Yi-Ting Lee Date: Tue, 2 Aug 2022 12:37:20 -0700 Subject: [PATCH 3/5] query lineage result str function test added --- src/sagemaker/lineage/query.py | 14 +++++---- tests/integ/sagemaker/lineage/helpers.py | 3 +- tests/unit/sagemaker/lineage/test_query.py | 33 ++++++++++++++++++---- 3 files changed, 38 insertions(+), 12 deletions(-) diff --git a/src/sagemaker/lineage/query.py b/src/sagemaker/lineage/query.py index 5277a2d806..0acdccd698 100644 --- a/src/sagemaker/lineage/query.py +++ b/src/sagemaker/lineage/query.py @@ -339,17 +339,21 @@ def __str__(self): 'edges':[ {'source_arn': 'string', 'destination_arn': 'string', 'association_type': 'string'}, ...], - + 'vertices':[ - {'arn': 'string', 'lineage_entity': 'string', 'lineage_source': 'string', '_session': }, + {'arn': 'string', 'lineage_entity': 'string', 'lineage_source': 'string', + '_session': }, ...], - + 'startarn':['string', ...] } """ - result_dict = vars(self) - return '{\n' + '\n\n'.join('\'{}\': {},'.format(key, val) for key, val in self.__dict__.items()) + '\n}' + return ( + "{\n" + + "\n\n".join("'{}': {},".format(key, val) for key, val in self.__dict__.items()) + + "\n}" + ) def _covert_edges_to_tuples(self): """Convert edges to tuple format for visualizer.""" diff --git a/tests/integ/sagemaker/lineage/helpers.py b/tests/integ/sagemaker/lineage/helpers.py index d2ef2c067d..0c40bbac91 100644 --- a/tests/integ/sagemaker/lineage/helpers.py +++ b/tests/integ/sagemaker/lineage/helpers.py @@ -16,8 +16,7 @@ import uuid from datetime import datetime import time -import boto3 -from botocore.config import Config + def name(): return "lineage-integ-{}-{}".format( diff --git a/tests/unit/sagemaker/lineage/test_query.py b/tests/unit/sagemaker/lineage/test_query.py index b5b809138d..2573314d3b 100644 --- a/tests/unit/sagemaker/lineage/test_query.py +++ b/tests/unit/sagemaker/lineage/test_query.py @@ -18,6 +18,7 @@ from sagemaker.lineage.lineage_trial_component import LineageTrialComponent from sagemaker.lineage.query import LineageEntityEnum, LineageSourceEnum, Vertex, LineageQuery import pytest +import re def test_lineage_query(sagemaker_session): @@ -540,15 +541,37 @@ def test_get_visualization_elements(sagemaker_session): start_arns=["arn:aws:sagemaker:us-west-2:0123456789012:context/mycontext"] ) - print(query_response) - elements = query_response._get_visualization_elements() - print(elements) - assert elements["nodes"][0] == ("arn1", "Endpoint", "Artifact", False) assert elements["nodes"][1] == ("arn2", "Model", "Context", False) assert elements["edges"][0] == ("arn1", "arn2", "Produced") - +def test_query_lineage_result_str(sagemaker_session): + lineage_query = LineageQuery(sagemaker_session) + sagemaker_session.sagemaker_client.query_lineage.return_value = { + "Vertices": [ + {"Arn": "arn1", "Type": "Endpoint", "LineageType": "Artifact"}, + {"Arn": "arn2", "Type": "Model", "LineageType": "Context"}, + ], + "Edges": [{"SourceArn": "arn1", "DestinationArn": "arn2", "AssociationType": "Produced"}], + } + + query_response = lineage_query.query( + start_arns=["arn:aws:sagemaker:us-west-2:0123456789012:context/mycontext"] + ) + + response_str = query_response.__str__() + pattern = "Mock id='\d*'" + replace = "Mock id=''" + response_str = re.sub(pattern, replace, response_str) + + assert ( + response_str + == "{\n'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': " + + "['arn:aws:sagemaker:us-west-2:0123456789012:context/mycontext'],\n}" + ) From 09c9fc166fb090b0526567b7e611fd357be12806 Mon Sep 17 00:00:00 2001 From: Yi-Ting Lee Date: Wed, 3 Aug 2022 15:46:18 -0700 Subject: [PATCH 4/5] change: add queryLineageResult visualizer unit test --- tests/integ/sagemaker/lineage/test_lineage_visualize.py | 4 ---- tests/unit/sagemaker/lineage/test_query.py | 4 ++-- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/tests/integ/sagemaker/lineage/test_lineage_visualize.py b/tests/integ/sagemaker/lineage/test_lineage_visualize.py index d9b3f879a8..555b98452e 100644 --- a/tests/integ/sagemaker/lineage/test_lineage_visualize.py +++ b/tests/integ/sagemaker/lineage/test_lineage_visualize.py @@ -57,10 +57,6 @@ def test_wide_graph_visualize(sagemaker_session): lq_result = lq.query(start_arns=[wide_graph_root_arn]) lq_result.visualize(path="wideGraph.html") - print("vertex len = ") - print(len(lq_result.vertices)) - assert False - except Exception as e: print(e) assert False diff --git a/tests/unit/sagemaker/lineage/test_query.py b/tests/unit/sagemaker/lineage/test_query.py index 2573314d3b..ada905008b 100644 --- a/tests/unit/sagemaker/lineage/test_query.py +++ b/tests/unit/sagemaker/lineage/test_query.py @@ -563,8 +563,8 @@ def test_query_lineage_result_str(sagemaker_session): ) response_str = query_response.__str__() - pattern = "Mock id='\d*'" - replace = "Mock id=''" + pattern = r"Mock id='\d*'" + replace = r"Mock id=''" response_str = re.sub(pattern, replace, response_str) assert ( From 684d45d556b874eb854a982e0ed857a6f93dc6fa Mon Sep 17 00:00:00 2001 From: Yi-Ting Lee Date: Thu, 4 Aug 2022 13:47:34 -0700 Subject: [PATCH 5/5] startarn added to lineage return value --- tests/unit/sagemaker/lineage/test_query.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/tests/unit/sagemaker/lineage/test_query.py b/tests/unit/sagemaker/lineage/test_query.py index ada905008b..0a357eb1fc 100644 --- a/tests/unit/sagemaker/lineage/test_query.py +++ b/tests/unit/sagemaker/lineage/test_query.py @@ -533,6 +533,11 @@ def test_get_visualization_elements(sagemaker_session): "Vertices": [ {"Arn": "arn1", "Type": "Endpoint", "LineageType": "Artifact"}, {"Arn": "arn2", "Type": "Model", "LineageType": "Context"}, + { + "Arn": "arn:aws:sagemaker:us-west-2:0123456789012:context/mycontext", + "Type": "Model", + "LineageType": "Context", + }, ], "Edges": [{"SourceArn": "arn1", "DestinationArn": "arn2", "AssociationType": "Produced"}], } @@ -545,6 +550,12 @@ def test_get_visualization_elements(sagemaker_session): assert elements["nodes"][0] == ("arn1", "Endpoint", "Artifact", False) assert elements["nodes"][1] == ("arn2", "Model", "Context", False) + assert elements["nodes"][2] == ( + "arn:aws:sagemaker:us-west-2:0123456789012:context/mycontext", + "Model", + "Context", + True, + ) assert elements["edges"][0] == ("arn1", "arn2", "Produced")