change: add queryLineageResult visualizer load test & integ test

Author: Yi-Ting Lee

tests/integ/sagemaker/lineage/helpers.py

@@ -12,16 +12,10 @@
 # language governing permissions and limitations under the License.
 """This module contains helper methods for tests of SageMaker Lineage"""
 from __future__ import absolute_import
-from urllib import response
 
 import uuid
 from datetime import datetime
 import time
-import boto3
-from botocore.config import Config
-from sagemaker.lineage import association
-from sagemaker.lineage.artifact import Artifact
-from sagemaker.lineage.association import Association
 
 
 def name():
@@ -92,7 +86,6 @@ def __init__(self, sagemaker_session):
         self.artifacts = []
         self.actions = []
         self.contexts = []
-        self.trialComponents = []
         self.associations = []
 
     def create_artifact(self, artifact_name, artifact_type="Dataset"):
@@ -115,11 +108,10 @@ def create_action(self, action_name, action_type="ModelDeployment"):
             ActionName=action_name,
             Source={
                 "SourceUri": "Test-action-" + action_name,
-                "SourceTypes": [
-                    {"SourceIdType": "S3ETag", "Value": "Test-action-sourceId-value"},
-                ],
+                "SourceType": "S3ETag",
+                "SourceId": "Test-action-sourceId-value",
             },
-            ActionType=action_type
+            ActionType=action_type,
         )
         self.actions.append(response["ActionArn"])
 
@@ -130,22 +122,15 @@ def create_context(self, context_name, context_type="Endpoint"):
             ContextName=context_name,
             Source={
                 "SourceUri": "Test-context-" + context_name,
-                "SourceTypes": [
-                    {"SourceIdType": "S3ETag", "Value": "Test-context-sourceId-value"},
-                ],
+                "SourceType": "S3ETag",
+                "SourceId": "Test-context-sourceId-value",
             },
-            ContextType=context_type
+            ContextType=context_type,
         )
         self.contexts.append(response["ContextArn"])
 
         return response["ContextArn"]
 
-    def create_trialComponent(self, trialComponent_name, trialComponent_type="TrainingJob"):
-        response = self.client.create_trial_component(
-            TrialComponentName=trialComponent_name,
-            
-        )
-
     def create_association(self, source_arn, dest_arn, association_type="AssociatedWith"):
         response = self.client.add_association(
             SourceArn=source_arn, DestinationArn=dest_arn, AssociationType=association_type
@@ -177,3 +162,10 @@ def clean_all(self):
                 time.sleep(0.5)
             except Exception as e:
                 print("skipped " + str(e))
+
+        for context_arn in self.contexts:
+            try:
+                self.client.delete_context(ContextArn=context_arn)
+                time.sleep(0.5)
+            except Exception as e:
+                print("skipped " + str(e))

tests/integ/sagemaker/lineage/test_lineage_visualize.py

@@ -11,16 +11,15 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 """This module contains code to test SageMaker ``LineageQueryResult.visualize()``"""
-
-import datetime
-import logging
+from __future__ import absolute_import
 import time
+import json
 
 import pytest
 
 import sagemaker.lineage.query
 from sagemaker.lineage.query import LineageQueryDirectionEnum
-from tests.integ.sagemaker.lineage.helpers import name, names, retry, LineageResourceHelper
+from tests.integ.sagemaker.lineage.helpers import name, LineageResourceHelper
 
 
 def test_LineageResourceHelper(sagemaker_session):
@@ -35,6 +34,7 @@ def test_LineageResourceHelper(sagemaker_session):
         print(e)
         assert False
 
+
 @pytest.mark.skip("visualizer load test")
 def test_wide_graph_visualize(sagemaker_session):
     lineage_resource_helper = LineageResourceHelper(sagemaker_session=sagemaker_session)
@@ -46,9 +46,11 @@ def test_wide_graph_visualize(sagemaker_session):
     #         \ \--> Artifact
     #          \--->  ...
     try:
-        for i in range(500):
+        for i in range(10):
             artifact_arn = lineage_resource_helper.create_artifact(artifact_name=name())
-            lineage_resource_helper.create_association(source_arn=wide_graph_root_arn, dest_arn=artifact_arn)
+            lineage_resource_helper.create_association(
+                source_arn=wide_graph_root_arn, dest_arn=artifact_arn
+            )
     except Exception as e:
         print(e)
         lineage_resource_helper.clean_all()
@@ -65,6 +67,7 @@ def test_wide_graph_visualize(sagemaker_session):
 
     lineage_resource_helper.clean_all()
 
+
 @pytest.mark.skip("visualizer load test")
 def test_long_graph_visualize(sagemaker_session):
     lineage_resource_helper = LineageResourceHelper(sagemaker_session=sagemaker_session)
@@ -74,9 +77,11 @@ def test_long_graph_visualize(sagemaker_session):
     # create long graph
     # Artifact -> Artifact -> ... -> Artifact
     try:
-        for i in range(20):
+        for i in range(10):
             new_artifact_arn = lineage_resource_helper.create_artifact(artifact_name=name())
-            lineage_resource_helper.create_association(source_arn=last_arn, dest_arn=new_artifact_arn)
+            lineage_resource_helper.create_association(
+                source_arn=last_arn, dest_arn=new_artifact_arn
+            )
             last_arn = new_artifact_arn
     except Exception as e:
         print(e)
@@ -85,12 +90,181 @@ def test_long_graph_visualize(sagemaker_session):
 
     try:
         lq = sagemaker.lineage.query.LineageQuery(sagemaker_session)
-        lq_result = lq.query(start_arns=[long_graph_root_arn], direction=LineageQueryDirectionEnum.DESCENDANTS)
+        lq_result = lq.query(
+            start_arns=[long_graph_root_arn], direction=LineageQueryDirectionEnum.DESCENDANTS
+        )
         # max depth = 10 -> graph rendered only has length of ten (in DESCENDANTS direction)
         lq_result.visualize(path="longGraph.html")
     except Exception as e:
         print(e)
         lineage_resource_helper.clean_all()
         assert False
 
-    lineage_resource_helper.clean_all()
+    lineage_resource_helper.clean_all()
+
+
+def test_graph_visualize(sagemaker_session):
+    lineage_resource_helper = LineageResourceHelper(sagemaker_session=sagemaker_session)
+
+    # create lineage data
+    # image artifact ------> model artifact(startarn) -> model deploy action -> endpoint context
+    #                    /->
+    # dataset artifact -/
+    try:
+        graph_startarn = lineage_resource_helper.create_artifact(
+            artifact_name=name(), artifact_type="Model"
+        )
+        image_artifact = lineage_resource_helper.create_artifact(
+            artifact_name=name(), artifact_type="Image"
+        )
+        lineage_resource_helper.create_association(
+            source_arn=image_artifact, dest_arn=graph_startarn, association_type="ContributedTo"
+        )
+        dataset_artifact = lineage_resource_helper.create_artifact(
+            artifact_name=name(), artifact_type="DataSet"
+        )
+        lineage_resource_helper.create_association(
+            source_arn=dataset_artifact, dest_arn=graph_startarn, association_type="AssociatedWith"
+        )
+        modeldeploy_action = lineage_resource_helper.create_action(
+            action_name=name(), action_type="ModelDeploy"
+        )
+        lineage_resource_helper.create_association(
+            source_arn=graph_startarn, dest_arn=modeldeploy_action, association_type="ContributedTo"
+        )
+        endpoint_context = lineage_resource_helper.create_context(
+            context_name=name(), context_type="Endpoint"
+        )
+        lineage_resource_helper.create_association(
+            source_arn=modeldeploy_action,
+            dest_arn=endpoint_context,
+            association_type="AssociatedWith",
+        )
+        time.sleep(1)
+    except Exception as e:
+        print(e)
+        lineage_resource_helper.clean_all()
+        assert False
+
+    # visualize
+    try:
+        lq = sagemaker.lineage.query.LineageQuery(sagemaker_session)
+        lq_result = lq.query(start_arns=[graph_startarn])
+        lq_result.visualize(path="testGraph.html")
+    except Exception as e:
+        print(e)
+        lineage_resource_helper.clean_all()
+        assert False
+
+    # check generated graph info
+    try:
+        fo = open("testGraph.html", "r")
+        lines = fo.readlines()
+        for line in lines:
+            if "nodes = " in line:
+                node = line
+            if "edges = " in line:
+                edge = line
+
+        # extract node data
+        start = node.find("[")
+        end = node.find("]")
+        res = node[start + 1 : end].split("}, ")
+        res = [i + "}" for i in res]
+        res[-1] = res[-1][:-1]
+        node_dict = [json.loads(i) for i in res]
+
+        # extract edge data
+        start = edge.find("[")
+        end = edge.find("]")
+        res = edge[start + 1 : end].split("}, ")
+        res = [i + "}" for i in res]
+        res[-1] = res[-1][:-1]
+        edge_dict = [json.loads(i) for i in res]
+
+        # check node number
+        assert len(node_dict) == 5
+
+        # check startarn
+        found_value = next(
+            dictionary for dictionary in node_dict if dictionary["id"] == graph_startarn
+        )
+        assert found_value["color"] == "#146eb4"
+        assert found_value["label"] == "Model"
+        assert found_value["shape"] == "star"
+        assert found_value["title"] == "Artifact"
+
+        # check image artifact
+        found_value = next(
+            dictionary for dictionary in node_dict if dictionary["id"] == image_artifact
+        )
+        assert found_value["color"] == "#146eb4"
+        assert found_value["label"] == "Image"
+        assert found_value["shape"] == "dot"
+        assert found_value["title"] == "Artifact"
+
+        # check dataset artifact
+        found_value = next(
+            dictionary for dictionary in node_dict if dictionary["id"] == dataset_artifact
+        )
+        assert found_value["color"] == "#146eb4"
+        assert found_value["label"] == "DataSet"
+        assert found_value["shape"] == "dot"
+        assert found_value["title"] == "Artifact"
+
+        # check modeldeploy action
+        found_value = next(
+            dictionary for dictionary in node_dict if dictionary["id"] == modeldeploy_action
+        )
+        assert found_value["color"] == "#88c396"
+        assert found_value["label"] == "ModelDeploy"
+        assert found_value["shape"] == "dot"
+        assert found_value["title"] == "Action"
+
+        # check endpoint context
+        found_value = next(
+            dictionary for dictionary in node_dict if dictionary["id"] == endpoint_context
+        )
+        assert found_value["color"] == "#ff9900"
+        assert found_value["label"] == "Endpoint"
+        assert found_value["shape"] == "dot"
+        assert found_value["title"] == "Context"
+
+        # check edge number
+        assert len(edge_dict) == 4
+
+        # check image_artifact ->  model_artifact(startarn) edge
+        found_value = next(
+            dictionary for dictionary in edge_dict if dictionary["from"] == image_artifact
+        )
+        assert found_value["to"] == graph_startarn
+        assert found_value["title"] == "ContributedTo"
+
+        # check dataset_artifact ->  model_artifact(startarn) edge
+        found_value = next(
+            dictionary for dictionary in edge_dict if dictionary["from"] == dataset_artifact
+        )
+        assert found_value["to"] == graph_startarn
+        assert found_value["title"] == "AssociatedWith"
+
+        # check model_artifact(startarn) ->  modeldeploy_action edge
+        found_value = next(
+            dictionary for dictionary in edge_dict if dictionary["from"] == graph_startarn
+        )
+        assert found_value["to"] == modeldeploy_action
+        assert found_value["title"] == "ContributedTo"
+
+        # check modeldeploy_action ->  endpoint_context edge
+        found_value = next(
+            dictionary for dictionary in edge_dict if dictionary["from"] == modeldeploy_action
+        )
+        assert found_value["to"] == endpoint_context
+        assert found_value["title"] == "AssociatedWith"
+
+    except Exception as e:
+        print(e)
+        lineage_resource_helper.clean_all()
+        assert False
+
+    # clean lineage data
+    lineage_resource_helper.clean_all()